From 2f2164106f5f39b447f949e7236144992d8fe181 Mon Sep 17 00:00:00 2001 From: Ayush Agrawal Date: Thu, 29 Oct 2026 02:27:03 -0700 Subject: [PATCH] fix: Close the HTTP response before the reader in ResponseStream.close() fixes #1372 FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/googleapis/java-genai/pull/1371 from googleapis:release-please--branches--main dea37a49b73d4d3c57eace0419246c310cdcbc1b PiperOrigin-RevId: 1009532954 --- .../java/com/google/genai/ResponseStream.java | 9 +- .../com/google/genai/ResponseStreamTest.java | 123 ++++++++++++++++++ 2 files changed, 128 insertions(+), 4 deletions(-) diff --git a/src/main/java/com/google/genai/ResponseStream.java b/src/main/java/com/google/genai/ResponseStream.java index 2436ea54f67..97c31408359 100644 --- a/src/main/java/com/google/genai/ResponseStream.java +++ b/src/main/java/com/google/genai/ResponseStream.java @@ -360,7 +360,12 @@ public void close() { } private void closeOwnResponse() { + // A blocked readLine() holds the reader's lock; closing the response ends it. try { + if (response != null) { + response.close(); + } + } finally { if (reader != null) { try { reader.close(); @@ -368,10 +373,6 @@ private void closeOwnResponse() { throw new GenAiIOException("Failed to close the response stream.", e); } } - } finally { - if (response != null) { - response.close(); - } } } diff --git a/src/test/java/com/google/genai/ResponseStreamTest.java b/src/test/java/com/google/genai/ResponseStreamTest.java index bd3968e648e..82046512c36 100644 --- a/src/test/java/com/google/genai/ResponseStreamTest.java +++ b/src/test/java/com/google/genai/ResponseStreamTest.java @@ -16,8 +16,11 @@ package com.google.genai; +import static java.util.concurrent.TimeUnit.SECONDS; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; import static org.junit.jupiter.api.Assertions.assertTrue; import com.fasterxml.jackson.databind.JsonNode; @@ -25,13 +28,26 @@ import com.google.common.collect.ImmutableList; import com.google.genai.types.Candidate; import com.google.genai.types.GenerateContentResponse; +import java.io.IOException; import java.nio.charset.StandardCharsets; +import java.time.Duration; import java.util.ArrayList; import java.util.Iterator; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import okhttp3.Headers; import okhttp3.MediaType; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; import okhttp3.ResponseBody; +import okio.Buffer; +import okio.Okio; +import okio.Source; +import okio.Timeout; import org.junit.jupiter.api.Test; public final class ResponseStreamTest { @@ -208,6 +224,113 @@ public ResponseStream next() { assertTrue(first.isConsumed()); } + /** + * Pins that close() from another thread releases a read that is blocked waiting for bytes, by + * closing the HTTP response before it waits for the reader's lock. + */ + @Test + public void testCloseClosesResponseWhileAnotherThreadIsReading() throws Exception { + BlockingSource source = new BlockingSource(); + ResponseBody body = + ResponseBody.create(Okio.buffer(source), MediaType.parse("text/event-stream"), -1); + HttpApiResponse response = + new HttpApiResponse( + new Response.Builder() + .request(new Request.Builder().url("https://example.com/stream").build()) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(body) + .build()); + ResponseStream stream = + new ResponseStream<>( + GenerateContentResponse.class, response, new DummyConverter(), "convert"); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future reading = executor.submit(() -> stream.iterator().hasNext()); + // While the source blocks, BufferedReader.readLine() holds the reader lock. + assertTrue(source.readStarted.await(5, SECONDS), "The stream read did not start"); + + Future closing = executor.submit(stream::close); + + assertTrue( + source.closed.await(5, SECONDS), + "close() did not close the HTTP response while another thread was reading"); + closing.get(5, SECONDS); + assertFalse(reading.get(5, SECONDS)); + } finally { + // Release the read so a failing run does not leave threads blocked. + source.releaseRead.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, SECONDS)); + } + } + + /** Pins that close() still closes the reader when closing the HTTP response throws. */ + @Test + public void testCloseClosesReaderWhenResponseCloseThrows() { + BlockingSource source = new BlockingSource(); + ResponseBody body = + ResponseBody.create(Okio.buffer(source), MediaType.parse("text/event-stream"), -1); + ApiResponse response = + new ApiResponse() { + @Override + public ResponseBody getBody() { + return body; + } + + @Override + public Headers getHeaders() { + return Headers.of(); + } + + @Override + public void close() { + throw new IllegalStateException("response close failed"); + } + }; + ResponseStream stream = + new ResponseStream<>( + GenerateContentResponse.class, response, new DummyConverter(), "convert"); + + IllegalStateException e = + assertTimeoutPreemptively( + Duration.ofSeconds(5), () -> assertThrows(IllegalStateException.class, stream::close)); + + assertEquals("response close failed", e.getMessage()); + assertEquals(0, source.closed.getCount(), "close() did not close the reader"); + } + + /** A response body source whose read blocks until the source is closed. */ + private static final class BlockingSource implements Source { + final CountDownLatch readStarted = new CountDownLatch(1); + final CountDownLatch releaseRead = new CountDownLatch(1); + final CountDownLatch closed = new CountDownLatch(1); + + @Override + public long read(Buffer sink, long byteCount) throws IOException { + readStarted.countDown(); + try { + releaseRead.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException(e); + } + return -1; + } + + @Override + public Timeout timeout() { + return Timeout.NONE; + } + + @Override + public void close() { + closed.countDown(); + releaseRead.countDown(); + } + } + private static ResponseStream textStream(String... texts) { StringBuilder sseData = new StringBuilder(); for (String text : texts) {