Skip to content

fix(bigquery-jdbc): stop producer and release resources on ResultSet.close #14712

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
logachev merged 6 commits into from
Oct 21, 2026
Merged

fix(bigquery-jdbc): stop producer and release resources on ResultSet.close #14712

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
7ce4107
fix(bigquery-jdbc): bound Storage Read API Arrow batch queue to 32 ba…
logachev Oct 20, 2026
02f9592
fix(bigquery-jdbc): stop producer and release resources on ResultSet.…
logachev Oct 20, 2026
3d3ba9b
chore(bigquery-jdbc): merge main into kirl/htapi_cleanup_on_close
logachev Oct 21, 2026
2465b08
fix(bigquery-jdbc): address PR #14712 review comments
logachev Oct 21, 2026
51cba45
refactor(bigquery-jdbc): move ResultSet buffer cleanup to BigQueryBas…
logachev Oct 21, 2026
ef9b74d
Revert "refactor(bigquery-jdbc): move ResultSet buffer cleanup to Big…
logachev Oct 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 27 additions & 6 deletions java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
}

Expand Down Expand Up @@ -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());
}
Expand Down Expand Up @@ -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();
}
Comment thread
logachev marked this conversation as resolved.

@Override
Expand Down
23 changes: 18 additions & 5 deletions java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
Expand Down Expand Up @@ -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
Expand Down
147 changes: 96 additions & 51 deletions java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;

@keshavdandeva keshavdandeva Oct 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess making default 32 makes sense but I wonder if this info should be added/explained somewhere or also if this should be a configurable conn prop?

@Neenu1995 Neenu1995 Oct 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the rationale behind 32 as the hardcoded number?

@logachev logachev Oct 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's an arbitrary constant. If consumer apps is reading data faster than we're fetching, buffer will be ~empty most of the time, if it is slower, we need some limit (similar to how we store 2 pages of results for Json).

We can make it configurable, but I'm not sure there is a lot of value in tweaking it.. Buffer is likely to be either empty or full in either case.

private static final String DEFAULT_DATASET_NAME = "_google_jdbc";
private static final String DEFAULT_TABLE_NAME = "temp_table_";
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -1005,6 +1009,8 @@ Future<?> populateArrowBufferedQueue(
String streamName = readSession.getStreams(0).getName();

while (true) {
ServerStream<ReadRowsResponse> stream = null;
boolean streamCompleted = false;
try {
ReadRowsRequest readRowsRequest =
ReadRowsRequest.newBuilder()
Expand All @@ -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<ReadRowsResponse> 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) {
Expand All @@ -1049,29 +1059,35 @@ 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"
+ Thread.currentThread().getName()
+ " 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"
+ Thread.currentThread().getName()
+ " Interrupted @ arrowStreamProcessor",
e);
enqueueError(arrowBatchWrapperBlockingQueue, e);
Thread.currentThread().interrupt();
} else {
LOG.log(
Level.WARNING,
Expand Down Expand Up @@ -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 =
Expand All @@ -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(),
Expand All @@ -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
Expand Down Expand Up @@ -1390,29 +1411,19 @@ Future<?> parseAndPopulateRpcDataAsync(
// as we have to process the first page
boolean hasRows = true;
while (hasRows) {
try {
Tuple<TableResult, Boolean> 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<TableResult, Boolean> 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;
}
Expand All @@ -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;
Expand All @@ -1436,17 +1446,28 @@ 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"
+ Thread.currentThread().getName()
+ " Interrupted @ populateBufferAsync",
ex);
enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex);
Thread.currentThread().interrupt();
} else {
LOG.log(
Level.WARNING,
Expand Down Expand Up @@ -1933,20 +1954,44 @@ enum QueryDialectType {
BIG_QUERY
}

private boolean isWorkerShuttingDown() {
Comment thread
logachev marked this conversation as resolved.
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 <T> void enqueueMarker(BlockingQueue<T> queue, T marker) {
if (isWorkerShuttingDown()) {
queue.offer(marker);
return;
}
try {
queue.put(marker);
} catch (InterruptedException ie) {
Thread.currentThread().interrupt();
}
}

private void enqueueError(BlockingQueue<BigQueryArrowBatchWrapper> queue, Exception e) {
Uninterruptibles.putUninterruptibly(queue, BigQueryArrowBatchWrapper.ofError(e));
enqueueMarker(queue, BigQueryArrowBatchWrapper.ofError(e));
}

private void enqueueEndOfStream(BlockingQueue<BigQueryArrowBatchWrapper> queue) {
Uninterruptibles.putUninterruptibly(queue, BigQueryArrowBatchWrapper.of(null, true));
enqueueMarker(queue, BigQueryArrowBatchWrapper.of(null, true));
}

private void enqueueBufferError(BlockingQueue<BigQueryFieldValueListWrapper> queue, Exception e) {
Uninterruptibles.putUninterruptibly(queue, BigQueryFieldValueListWrapper.ofError(e));
enqueueMarker(queue, BigQueryFieldValueListWrapper.ofError(e));
}

private void enqueueBufferEndOfStream(BlockingQueue<BigQueryFieldValueListWrapper> queue) {
Uninterruptibles.putUninterruptibly(queue, BigQueryFieldValueListWrapper.ofEndOfStream(null));
enqueueMarker(queue, BigQueryFieldValueListWrapper.ofEndOfStream(null));
}

Comment thread
logachev marked this conversation as resolved.
QueryStatistics describePositionalParameterQuery(String query)
Expand Down
Loading
Loading