From 7ce41072aec71b7cabed19b70f0f3739a5b7580d Mon Sep 17 00:00:00 2001 From: kirl Date: Tue, 20 Oct 2026 23:36:06 +0000 Subject: [PATCH 1/2] fix(bigquery-jdbc): bound Storage Read API Arrow batch queue to 32 batches --- .../bigquery/jdbc/BigQueryStatement.java | 7 ++++- .../bigquery/jdbc/BigQueryStatementTest.java | 29 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index af50bf8b8df3..0654d6036f0f 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -104,6 +104,7 @@ public class BigQueryStatement extends BigQueryNoOpsStatement { private final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(this.toString()); public static final int DEFAULT_BUFFER_SIZE = BigQuerySettings.DEFAULT_NUM_BUFFERED_ROWS * 2; + static final int DEFAULT_ARROW_BUFFER_SIZE = 32; private static final String DEFAULT_DATASET_NAME = "_google_jdbc"; private static final String DEFAULT_TABLE_NAME = "temp_table_"; private static final String JDBC_JOB_PREFIX = "google-jdbc-"; @@ -931,7 +932,7 @@ ResultSet processArrowResultSet(TableResult results, Job job) throws SQLExceptio .setMaxStreamCount(1); ReadSession readSession = getReadSession(builder.build()); - this.arrowBatchWrapperBlockingQueue = new LinkedBlockingDeque<>(getBufferSize()); + this.arrowBatchWrapperBlockingQueue = new LinkedBlockingDeque<>(getArrowBufferSize()); // deserialize and populate the buffer async, so that the client isn't blocked populateBufferWorker = populateArrowBufferedQueue( @@ -1512,6 +1513,10 @@ private int getBufferSize() { : Math.min(this.querySettings.getNumBufferedRows() * 2, 100000)); } + int getArrowBufferSize() { + return DEFAULT_ARROW_BUFFER_SIZE; + } + /** Returns the destinationTable from jobId by calling `jobs.get` API */ TableId getDestinationTable(JobId jobId) { Job job = this.bigQuery.getJob(jobId); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index 9245f627c147..a442efd61ed4 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -1318,4 +1318,33 @@ public void testJoblessQueryPopulatesQueryStatistics() throws Exception { verify(bigquery, Mockito.never()).getJob(any(JobId.class)); } + + @Test + public void testProcessArrowResultSetUsesBoundedQueueCapacity() throws SQLException { + BigQueryStatement defaultStmtSpy = Mockito.spy(bigQueryStatement); + BigQueryReadClient bigQueryReadClient = mock(BigQueryReadClient.class); + Schema schema = Schema.of(fieldList); + ReadSession readSession = ReadSession.getDefaultInstance(); + doReturn(bigQueryReadClient).when(defaultStmtSpy).getBigQueryReadClient(); + doReturn(readSession).when(defaultStmtSpy).getReadSession(any(CreateReadSessionRequest.class)); + Future mockWorker = mock(Future.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> defaultQueueCaptor = + ArgumentCaptor.forClass(BlockingQueue.class); + doReturn(mockWorker) + .when(defaultStmtSpy) + .populateArrowBufferedQueue( + any(ReadSession.class), defaultQueueCaptor.capture(), any(BigQueryReadClient.class)); + doReturn(arrowSchema).when(defaultStmtSpy).getArrowSchema(any(ReadSession.class)); + + TableResult result = mock(TableResult.class); + doReturn(schema).when(result).getSchema(); + doReturn(10L).when(result).getTotalRows(); + doReturn(TABLE_ID).when(defaultStmtSpy).getDestinationTable(any()); + + defaultStmtSpy.processArrowResultSet(result, null); + assertEquals( + BigQueryStatement.DEFAULT_ARROW_BUFFER_SIZE, + defaultQueueCaptor.getValue().remainingCapacity()); + } } From 02f9592015eb88e11f5de8ae5788c1955bf0a497 Mon Sep 17 00:00:00 2001 From: kirl Date: Tue, 20 Oct 2026 23:38:50 +0000 Subject: [PATCH 2/2] fix(bigquery-jdbc): stop producer and release resources on ResultSet.close() --- .../bigquery/jdbc/BigQueryArrowResultSet.java | 19 ++- .../bigquery/jdbc/BigQueryJsonResultSet.java | 7 + .../bigquery/jdbc/BigQueryStatement.java | 125 +++++++++++++++--- .../jdbc/BigQueryArrowResultSetTest.java | 30 +++++ .../bigquery/jdbc/BigQueryStatementTest.java | 99 ++++++++++++++ 5 files changed, 258 insertions(+), 22 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java index b6cdc393173f..94573278a69c 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java @@ -246,8 +246,12 @@ private void deserializeArrowBatch(ArrowRecordBatch batch) throws SQLException { @Override public void close() { LOG.fineTrace("close", () -> String.format("Closing BigQueryArrowResultSet %s.", this)); - vectorSchemaRoot.close(); - allocator.close(); + if (vectorSchemaRoot != null) { + vectorSchemaRoot.close(); + } + if (allocator != null) { + allocator.close(); + } } } @@ -542,12 +546,23 @@ private String formatRangeElement(Object element, StandardSQLTypeName elementTyp @Override public void close() { + if (this.isClosed) { + return; + } LOG.fineTrace("close", () -> String.format("Closing BigqueryArrowResultSet %s.", this)); this.isClosed = true; if (ownedTask != null) { // cancel the producer task when result set is closed ownedTask.cancel(true); } + if (this.buffer != null) { + this.buffer.clear(); + } + if (this.arrowDeserializer != null) { + this.arrowDeserializer.close(); + } else if (this.allocator != null) { + this.allocator.close(); + } super.close(); } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java index c8a2451aeeaa..afa54bff4fb9 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java @@ -382,6 +382,9 @@ private FieldValue getObjectInternal(int columnIndex) throws SQLException { @Override public void close() { + if (this.isClosed) { + return; + } LOG.fineTrace("close", () -> String.format("Closing BigqueryJsonResultSet %s.", this)); this.isClosed = true; if (ownedTasks != null) { @@ -391,6 +394,10 @@ public void close() { } } } + if (this.buffer != null) { + this.buffer.clear(); + } + this.cursor = null; super.close(); } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 0654d6036f0f..44eb3557f662 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -19,6 +19,7 @@ import com.google.api.core.InternalApi; import com.google.api.gax.paging.Page; import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.StatusCode; import com.google.cloud.Tuple; import com.google.cloud.bigquery.BigQuery; @@ -62,7 +63,6 @@ import com.google.cloud.bigquery.storage.v1.ReadSession; import com.google.cloud.bigquery.storage.v1.ReadSession.TableReadOptions; import com.google.common.annotations.VisibleForTesting; -import com.google.common.util.concurrent.Uninterruptibles; import io.grpc.Status; import io.grpc.StatusRuntimeException; import io.opentelemetry.api.common.AttributeKey; @@ -535,6 +535,8 @@ private void closeStatementResources() throws SQLException { tmp.close(); } this.batchQueries.clear(); + this.arrowBatchWrapperBlockingQueue = null; + this.bigQueryFieldValueListWrapperBlockingQueue = null; this.currentUpdateCount = -1; this.currentJobIdIndex = -1; if (this.connection != null) { @@ -1005,6 +1007,8 @@ Future populateArrowBufferedQueue( String streamName = readSession.getStreams(0).getName(); while (true) { + ServerStream stream = null; + boolean streamCompleted = false; try { ReadRowsRequest readRowsRequest = ReadRowsRequest.newBuilder() @@ -1014,19 +1018,23 @@ Future populateArrowBufferedQueue( // Process each block of rows as they arrive and decode using our simple row // reader. - com.google.api.gax.rpc.ServerStream stream = - bqReadClient.readRowsCallable().call(readRowsRequest); + stream = bqReadClient.readRowsCallable().call(readRowsRequest); for (ReadRowsResponse response : stream) { if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { + arrowBatchWrapperBlockingQueue.clear(); break; } ArrowRecordBatch currentBatch = response.getArrowRecordBatch(); - Uninterruptibles.putUninterruptibly( - arrowBatchWrapperBlockingQueue, + arrowBatchWrapperBlockingQueue.put( BigQueryArrowBatchWrapper.of(currentBatch)); rowsRead += response.getRowCount(); } + if (!Thread.currentThread().isInterrupted() && !executor.isShutdown()) { + streamCompleted = true; + } else { + arrowBatchWrapperBlockingQueue.clear(); + } break; } catch (ApiException e) { if (e.getStatusCode().getCode() == StatusCode.Code.NOT_FOUND) { @@ -1049,10 +1057,16 @@ Future populateArrowBufferedQueue( "Connection interrupted during arrow stream read, retrying. attempt: %d", retryCount); Thread.sleep(RETRY_DELAY_MS); + } finally { + if (stream != null && !streamCompleted) { + stream.cancel(); + } } } } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + arrowBatchWrapperBlockingQueue.clear(); LOG.log( Level.WARNING, "\n" @@ -1060,10 +1074,11 @@ Future populateArrowBufferedQueue( + " Interrupted @ arrowStreamProcessor", e); enqueueError(arrowBatchWrapperBlockingQueue, e); - Thread.currentThread().interrupt(); } catch (Exception e) { if (e.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) { + Thread.currentThread().interrupt(); + arrowBatchWrapperBlockingQueue.clear(); LOG.log( Level.WARNING, "\n" @@ -1071,7 +1086,6 @@ Future populateArrowBufferedQueue( + " Interrupted @ arrowStreamProcessor", e); enqueueError(arrowBatchWrapperBlockingQueue, e); - Thread.currentThread().interrupt(); } else { LOG.log( Level.WARNING, @@ -1338,8 +1352,7 @@ Future runNextPageTaskAsync( currentPageToken = currentResults.getNextPageToken(); // this will be parsed asynchronously without blocking the current // thread - Uninterruptibles.putUninterruptibly( - rpcResponseQueue, Tuple.of(currentResults, true)); + rpcResponseQueue.put(Tuple.of(currentResults, true)); LOG.fine( "Fetched %d results from the server in %d ms.", querySettings.getMaxResultPerPage(), @@ -1348,14 +1361,30 @@ Future runNextPageTaskAsync( paginationSpan.end(); } } + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + LOG.warning( + "%s Interrupted @ runNextPageTaskAsync: %s", + Thread.currentThread().getName(), ex.getMessage()); } catch (Exception ex) { - Uninterruptibles.putUninterruptibly( - bigQueryFieldValueListWrapperBlockingQueue, - BigQueryFieldValueListWrapper.ofError(ex)); + if (ex.getCause() instanceof InterruptedException + || Thread.currentThread().isInterrupted()) { + Thread.currentThread().interrupt(); + } + enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); } finally { // this will stop the parseDataTask as well when the pagination // completes - Uninterruptibles.putUninterruptibly(rpcResponseQueue, Tuple.of(null, false)); + Tuple endOfRpcTuple = Tuple.of(null, false); + if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { + rpcResponseQueue.offer(endOfRpcTuple); + } else { + try { + rpcResponseQueue.put(endOfRpcTuple); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } + } } // We cannot do queryTaskExecutor.shutdownNow() here as populate buffer method may // not @@ -1425,8 +1454,7 @@ Future parseAndPopulateRpcDataAsync( // do not process further pages and shutdown (inner loop) break; } - Uninterruptibles.putUninterruptibly( - bigQueryFieldValueListWrapperBlockingQueue, + bigQueryFieldValueListWrapperBlockingQueue.put( BigQueryFieldValueListWrapper.of( schema.getFields(), fieldValueList, isComplexColumn)); results += 1; @@ -1436,9 +1464,19 @@ Future parseAndPopulateRpcDataAsync( results, (int) ((System.nanoTime() - startTime) / 1000000)); } + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + LOG.log( + Level.WARNING, + "\n" + + Thread.currentThread().getName() + + " Interrupted @ populateBufferAsync", + ex); + enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); } catch (Exception ex) { if (ex.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) { + Thread.currentThread().interrupt(); LOG.log( Level.WARNING, "\n" @@ -1446,7 +1484,6 @@ Future parseAndPopulateRpcDataAsync( + " Interrupted @ populateBufferAsync", ex); enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); - Thread.currentThread().interrupt(); } else { LOG.log( Level.WARNING, @@ -1933,20 +1970,68 @@ enum QueryDialectType { BIG_QUERY } + private boolean isWorkerShuttingDown() { + if (Thread.currentThread().isInterrupted()) { + return true; + } + BigQueryConnection conn = this.connection; + if (conn == null) { + return true; + } + ExecutorService executor = conn.getExecutorService(); + return executor != null && executor.isShutdown(); + } + private void enqueueError(BlockingQueue queue, Exception e) { - Uninterruptibles.putUninterruptibly(queue, BigQueryArrowBatchWrapper.ofError(e)); + BigQueryArrowBatchWrapper errorWrapper = BigQueryArrowBatchWrapper.ofError(e); + if (isWorkerShuttingDown()) { + queue.offer(errorWrapper); + return; + } + try { + queue.put(errorWrapper); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } } private void enqueueEndOfStream(BlockingQueue queue) { - Uninterruptibles.putUninterruptibly(queue, BigQueryArrowBatchWrapper.of(null, true)); + BigQueryArrowBatchWrapper endWrapper = BigQueryArrowBatchWrapper.of(null, true); + if (isWorkerShuttingDown()) { + queue.offer(endWrapper); + return; + } + try { + queue.put(endWrapper); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } } private void enqueueBufferError(BlockingQueue queue, Exception e) { - Uninterruptibles.putUninterruptibly(queue, BigQueryFieldValueListWrapper.ofError(e)); + BigQueryFieldValueListWrapper errorWrapper = BigQueryFieldValueListWrapper.ofError(e); + if (isWorkerShuttingDown()) { + queue.offer(errorWrapper); + return; + } + try { + queue.put(errorWrapper); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } } private void enqueueBufferEndOfStream(BlockingQueue queue) { - Uninterruptibles.putUninterruptibly(queue, BigQueryFieldValueListWrapper.ofEndOfStream(null)); + BigQueryFieldValueListWrapper endWrapper = BigQueryFieldValueListWrapper.ofEndOfStream(null); + if (isWorkerShuttingDown()) { + queue.offer(endWrapper); + return; + } + try { + queue.put(endWrapper); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } } QueryStatistics describePositionalParameterQuery(String query) diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java index 8681fd9559a5..0797c829194d 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java @@ -24,6 +24,8 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import com.google.cloud.bigquery.Field; import com.google.cloud.bigquery.Field.Mode; @@ -339,6 +341,34 @@ public void testIsClosed() { assertThat(bigQueryArrowResultSet.isClosed()).isFalse(); } + @Test + public void testCloseCancelsProducerClearsBufferAndIsIdempotent() + throws SQLException, IOException { + ArrowSchema arrowSchema = + ArrowSchema.newBuilder() + .setSerializedSchema(serializeSchema(vectorSchemaRoot.getSchema())) + .build(); + Future workerTask = mock(Future.class); + BlockingQueue queue = new LinkedBlockingDeque<>(4); + queue.add(arrowBatchWrapper); + queue.add(arrowBatchWrapperLast); + + BigQueryArrowResultSet rs = + BigQueryArrowResultSet.of(QUERY_SCHEMA, arrowSchema, 1, statement, queue, workerTask, null); + + assertThat(rs.isClosed()).isFalse(); + assertThat(queue).hasSize(2); + + rs.close(); + assertThat(rs.isClosed()).isTrue(); + assertThat(queue).isEmpty(); + verify(workerTask, times(1)).cancel(true); + + // Second close() call should be a no-op + rs.close(); + verify(workerTask, times(1)).cancel(true); + } + @Test public void testResultSetHoldability() throws SQLException { assertThat(bigQueryArrowResultSet.getHoldability()) diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index a442efd61ed4..2fda81227e71 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -33,6 +33,8 @@ import static org.mockito.Mockito.withSettings; import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.ServerStream; +import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.gax.rpc.StatusCode; import com.google.cloud.ServiceOptions; import com.google.cloud.Tuple; @@ -62,11 +64,15 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.jdbc.BigQueryStatement.JobIdWrapper; import com.google.cloud.bigquery.spi.BigQueryRpcFactory; +import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch; import com.google.cloud.bigquery.storage.v1.ArrowSchema; import com.google.cloud.bigquery.storage.v1.ArrowSerializationOptions; import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest; +import com.google.cloud.bigquery.storage.v1.ReadRowsRequest; +import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; import com.google.cloud.bigquery.storage.v1.ReadSession; +import com.google.cloud.bigquery.storage.v1.ReadStream; import com.google.common.collect.ImmutableList; import com.google.common.collect.Maps; import io.opentelemetry.api.baggage.Baggage; @@ -83,14 +89,18 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.UUID; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Stream; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BitVector; @@ -1347,4 +1357,93 @@ public void testProcessArrowResultSetUsesBoundedQueueCapacity() throws SQLExcept BigQueryStatement.DEFAULT_ARROW_BUFFER_SIZE, defaultQueueCaptor.getValue().remainingCapacity()); } + + @Test + public void testPopulateArrowBufferedQueueStopsAndCancelsStreamWhenResultSetClosedWhileQueueFull() + throws Exception { + BlockingQueue boundedQueue = new LinkedBlockingDeque<>(1); + ReadSession readSession = + ReadSession.newBuilder() + .addStreams( + ReadStream.newBuilder().setName("projects/p/locations/l/sessions/s/streams/0")) + .build(); + + ReadRowsResponse batchResponse = + ReadRowsResponse.newBuilder() + .setArrowRecordBatch(ArrowRecordBatch.getDefaultInstance()) + .build(); + + CountDownLatch secondBatchReturnedLatch = new CountDownLatch(1); + AtomicBoolean streamCancelled = new AtomicBoolean(false); + + @SuppressWarnings("unchecked") + ServerStream serverStream = mock(ServerStream.class); + Mockito.doAnswer( + invocation -> { + streamCancelled.set(true); + return null; + }) + .when(serverStream) + .cancel(); + + Iterator responseIterator = + new Iterator() { + private int count = 0; + + @Override + public boolean hasNext() { + return true; + } + + @Override + public ReadRowsResponse next() { + count++; + if (count == 2) { + secondBatchReturnedLatch.countDown(); + } + return batchResponse; + } + }; + doReturn(responseIterator).when(serverStream).iterator(); + + @SuppressWarnings("unchecked") + ServerStreamingCallable callable = + mock(ServerStreamingCallable.class); + doReturn(serverStream).when(callable).call(any(ReadRowsRequest.class)); + + BigQueryReadClient readClient = mock(BigQueryReadClient.class); + doReturn(callable).when(readClient).readRowsCallable(); + + Future producerFuture = + bigQueryStatement.populateArrowBufferedQueue(readSession, boundedQueue, readClient); + + // Wait until the first batch is in the queue and the second batch has been yielded by the + // stream iterator (so the producer thread is blocked or about to block on boundedQueue.put()). + assertTrue(secondBatchReturnedLatch.await(5, TimeUnit.SECONDS)); + assertEquals(1, boundedQueue.size()); + + BigQueryArrowResultSet rs = + BigQueryArrowResultSet.of( + Schema.of(fieldList), + arrowSchema, + 100L, + bigQueryStatement, + boundedQueue, + producerFuture, + bigquery); + + // Closing the ResultSet while the queue is full must cancel the producer task, + // unblock the producer thread, cancel the underlying ServerStream, and clear the buffer. + rs.close(); + + assertTrue(producerFuture.isCancelled()); + testExecutorService.shutdown(); + assertTrue( + testExecutorService.awaitTermination(5, TimeUnit.SECONDS), + "Producer thread should terminate promptly after ResultSet.close()"); + assertTrue(streamCancelled.get(), "ServerStream.cancel() should be called on early close"); + assertTrue( + boundedQueue.stream().allMatch(w -> w.getCurrentArrowBatch() == null), + "No ArrowRecordBatch should remain buffered after ResultSet.close()"); + } }