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..607b76bec92e 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(); + } } } @@ -279,6 +283,7 @@ public boolean next() throws SQLException { // Advance the cursor. Potentially blocking operation. try (Scope scope = makeOriginalContextCurrent()) { BigQueryArrowBatchWrapper batchWrapper = this.buffer.take(); + checkClosed(); if (batchWrapper.getException() != null) { throw backgroundFetchError(batchWrapper.getException()); } @@ -542,13 +547,29 @@ 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); + try { + if (ownedTask != null) { + // cancel the producer task when result set is closed + ownedTask.cancel(true); + } + if (this.buffer != null) { + this.buffer.clear(); + // Unblock any consumer thread currently waiting on buffer.take() in next() + this.buffer.offer(BigQueryArrowBatchWrapper.of(null, true)); + } + if (this.arrowDeserializer != null) { + this.arrowDeserializer.close(); + } else if (this.allocator != null) { + this.allocator.close(); + } + } finally { + super.close(); } - super.close(); } @Override 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..46b6c06e8fc1 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 @@ -227,6 +227,7 @@ public boolean next() throws SQLException { try { // Advance the cursor,Potentially blocking operation this.cursor = this.buffer.take(); + checkClosed(); if (this.cursor.getException() != null) { throw backgroundFetchError(this.cursor.getException()); } @@ -382,16 +383,28 @@ 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) { - for (Future ownedTask : ownedTasks) { - if (ownedTask != null) { - ownedTask.cancel(true); + try { + if (ownedTasks != null) { + for (Future ownedTask : ownedTasks) { + if (ownedTask != null) { + ownedTask.cancel(true); + } } } + if (this.buffer != null) { + this.buffer.clear(); + // Unblock any consumer thread currently waiting on buffer.take() in next() + this.buffer.offer(BigQueryFieldValueListWrapper.ofEndOfStream(null)); + } + this.cursor = null; + } finally { + super.close(); } - super.close(); } @Override 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..060c60c607b2 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; @@ -104,6 +104,8 @@ 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; + // Bounds in-memory ArrowRecordBatch buffering per ResultSet (~15-20 MB of in-flight batches), + // keeping the gRPC ReadRows stream saturated while preventing unbounded heap accumulation. 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_"; @@ -535,6 +537,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 +1009,8 @@ Future populateArrowBufferedQueue( String streamName = readSession.getStreams(0).getName(); while (true) { + ServerStream stream = null; + boolean streamCompleted = false; try { ReadRowsRequest readRowsRequest = ReadRowsRequest.newBuilder() @@ -1014,19 +1020,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()) { - break; + if (isWorkerShuttingDown()) { + throw new InterruptedException( + "Arrow stream reading was interrupted or shut down before completion."); } ArrowRecordBatch currentBatch = response.getArrowRecordBatch(); - Uninterruptibles.putUninterruptibly( - arrowBatchWrapperBlockingQueue, + arrowBatchWrapperBlockingQueue.put( BigQueryArrowBatchWrapper.of(currentBatch)); rowsRead += response.getRowCount(); } + if (isWorkerShuttingDown()) { + throw new InterruptedException( + "Arrow stream reading was interrupted or shut down before completion."); + } + streamCompleted = true; break; } catch (ApiException e) { if (e.getStatusCode().getCode() == StatusCode.Code.NOT_FOUND) { @@ -1049,10 +1059,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 +1076,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 +1088,6 @@ Future populateArrowBufferedQueue( + " Interrupted @ arrowStreamProcessor", e); enqueueError(arrowBatchWrapperBlockingQueue, e); - Thread.currentThread().interrupt(); } else { LOG.log( Level.WARNING, @@ -1312,11 +1328,9 @@ Future runNextPageTaskAsync( try { while (currentPageToken != null) { // do not process further pages and shutdown - if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { - LOG.warning( - "%s Interrupted @ runNextPageTaskAsync", - Thread.currentThread().getName()); - break; + if (isWorkerShuttingDown()) { + throw new InterruptedException( + "Pagination was interrupted or shut down before completion."); } SpanBuilder spanBuilder = @@ -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,22 @@ Future runNextPageTaskAsync( paginationSpan.end(); } } + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + LOG.warning( + "%s Interrupted @ runNextPageTaskAsync: %s", + Thread.currentThread().getName(), ex.getMessage()); + enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); } 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)); + enqueueMarker(rpcResponseQueue, Tuple.of(null, false)); } // We cannot do queryTaskExecutor.shutdownNow() here as populate buffer method may // not @@ -1390,29 +1411,19 @@ Future parseAndPopulateRpcDataAsync( // as we have to process the first page boolean hasRows = true; while (hasRows) { - try { - Tuple nextPageTuple = rpcResponseQueue.take(); - if (nextPageTuple.x() != null) { - fieldValueLists = nextPageTuple.x().getValues(); - } else { - fieldValueLists = null; - } - hasRows = nextPageTuple.y(); - - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.log( - Level.WARNING, - "\n" + Thread.currentThread().getName() + " Interrupted", - e); - // Thread might get interrupted while calling the Cancel method, which is - // expected, so logging this instead of throwing the exception back - break; + Tuple nextPageTuple = rpcResponseQueue.take(); + if (nextPageTuple.x() != null) { + fieldValueLists = nextPageTuple.x().getValues(); + } else { + fieldValueLists = null; } + hasRows = nextPageTuple.y(); - if (Thread.currentThread().isInterrupted() - || executor.isShutdown() - || fieldValueLists == null) { + if (isWorkerShuttingDown()) { + throw new InterruptedException( + "Buffer population was interrupted or shut down before completion."); + } + if (fieldValueLists == null) { // do not process further pages and shutdown (outerloop) break; } @@ -1421,12 +1432,11 @@ Future parseAndPopulateRpcDataAsync( long results = 0; for (FieldValueList fieldValueList : fieldValueLists) { - if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { - // do not process further pages and shutdown (inner loop) - break; + if (isWorkerShuttingDown()) { + throw new InterruptedException( + "Buffer population was interrupted or shut down before completion."); } - Uninterruptibles.putUninterruptibly( - bigQueryFieldValueListWrapperBlockingQueue, + bigQueryFieldValueListWrapperBlockingQueue.put( BigQueryFieldValueListWrapper.of( schema.getFields(), fieldValueList, isComplexColumn)); results += 1; @@ -1436,9 +1446,21 @@ Future parseAndPopulateRpcDataAsync( results, (int) ((System.nanoTime() - startTime) / 1000000)); } + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + bigQueryFieldValueListWrapperBlockingQueue.clear(); + 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(); + bigQueryFieldValueListWrapperBlockingQueue.clear(); LOG.log( Level.WARNING, "\n" @@ -1446,7 +1468,6 @@ Future parseAndPopulateRpcDataAsync( + " Interrupted @ populateBufferAsync", ex); enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); - Thread.currentThread().interrupt(); } else { LOG.log( Level.WARNING, @@ -1933,20 +1954,44 @@ 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 enqueueMarker(BlockingQueue queue, T marker) { + if (isWorkerShuttingDown()) { + queue.offer(marker); + return; + } + try { + queue.put(marker); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } + } + private void enqueueError(BlockingQueue queue, Exception e) { - Uninterruptibles.putUninterruptibly(queue, BigQueryArrowBatchWrapper.ofError(e)); + enqueueMarker(queue, BigQueryArrowBatchWrapper.ofError(e)); } private void enqueueEndOfStream(BlockingQueue queue) { - Uninterruptibles.putUninterruptibly(queue, BigQueryArrowBatchWrapper.of(null, true)); + enqueueMarker(queue, BigQueryArrowBatchWrapper.of(null, true)); } private void enqueueBufferError(BlockingQueue queue, Exception e) { - Uninterruptibles.putUninterruptibly(queue, BigQueryFieldValueListWrapper.ofError(e)); + enqueueMarker(queue, BigQueryFieldValueListWrapper.ofError(e)); } private void enqueueBufferEndOfStream(BlockingQueue queue) { - Uninterruptibles.putUninterruptibly(queue, BigQueryFieldValueListWrapper.ofEndOfStream(null)); + enqueueMarker(queue, BigQueryFieldValueListWrapper.ofEndOfStream(null)); } 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..6f8a249513b1 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,36 @@ 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).hasSize(1); + assertThat(queue.peek().getCurrentArrowBatch()).isNull(); + assertThat(queue.peek().isLast()).isTrue(); + 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..41f3731c2806 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,19 @@ 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.ExecutionException; 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 +1358,206 @@ 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()"); + } + + @Test + public void testCancelStatementUnblocksConcurrentArrowResultSetNext() throws Exception { + BlockingQueue emptyQueue = new LinkedBlockingDeque<>(4); + Future producerFuture = mock(Future.class); + BigQueryArrowResultSet rs = + BigQueryArrowResultSet.of( + Schema.of(fieldList), + arrowSchema, + 100L, + bigQueryStatement, + emptyQueue, + producerFuture, + bigquery); + bigQueryStatement.currentResultSet = rs; + + CountDownLatch readerEnteredNext = new CountDownLatch(1); + Future readerFuture = + testExecutorService.submit( + () -> { + readerEnteredNext.countDown(); + rs.next(); + return null; + }); + + assertTrue(readerEnteredNext.await(5, TimeUnit.SECONDS)); + // Give the reader thread a brief moment to block inside emptyQueue.take() + Thread.sleep(50); + + bigQueryStatement.cancel(); + + ExecutionException ex = + assertThrows(ExecutionException.class, () -> readerFuture.get(5, TimeUnit.SECONDS)); + assertTrue(ex.getCause() instanceof SQLException); + } + + @Test + public void testCancelStatementUnblocksConcurrentJsonResultSetNext() throws Exception { + BlockingQueue emptyQueue = new LinkedBlockingDeque<>(4); + BigQueryJsonResultSet rs = + BigQueryJsonResultSet.of( + Schema.of(fieldList), 100L, emptyQueue, bigQueryStatement, null, bigquery); + bigQueryStatement.currentResultSet = rs; + + CountDownLatch readerEnteredNext = new CountDownLatch(1); + Future readerFuture = + testExecutorService.submit( + () -> { + readerEnteredNext.countDown(); + rs.next(); + return null; + }); + + assertTrue(readerEnteredNext.await(5, TimeUnit.SECONDS)); + Thread.sleep(50); + + bigQueryStatement.cancel(); + + ExecutionException ex = + assertThrows(ExecutionException.class, () -> readerFuture.get(5, TimeUnit.SECONDS)); + assertTrue(ex.getCause() instanceof SQLException); + } + + @Test + public void testPopulateArrowBufferedQueueEnqueuesErrorWhenWorkerShuttingDown() throws Exception { + BlockingQueue boundedQueue = new LinkedBlockingDeque<>(4); + 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(); + + @SuppressWarnings("unchecked") + ServerStream serverStream = mock(ServerStream.class); + Iterator responseIterator = + new Iterator() { + @Override + public boolean hasNext() { + // Simulate executor shutdown mid-stream before the next batch is processed + testExecutorService.shutdownNow(); + return true; + } + + @Override + public ReadRowsResponse next() { + 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); + producerFuture.get(5, TimeUnit.SECONDS); + + BigQueryArrowBatchWrapper marker = boundedQueue.poll(5, TimeUnit.SECONDS); + assertNotNull(marker); + assertNotNull( + marker.getException(), + "Stopping early due to worker shutdown must enqueue an error marker rather than normal end-of-stream"); + } }