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

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();
   }
 

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();
   }
 

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,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-";
@@ -534,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) {
@@ -931,7 +934,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(
@@ -1004,6 +1007,8 @@ Future<?> populateArrowBufferedQueue(
                     String streamName = readSession.getStreams(0).getName();
 
                     while (true) {
+                      ServerStream<ReadRowsResponse> stream = null;
+                      boolean streamCompleted = false;
                       try {
                         ReadRowsRequest readRowsRequest =
                             ReadRowsRequest.newBuilder()
@@ -1013,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<ReadRowsResponse> 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) {
@@ -1048,29 +1057,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,
@@ -1337,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(),
@@ -1347,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<TableResult, Boolean> 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
@@ -1424,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;
@@ -1435,17 +1464,26 @@ 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"
                               + Thread.currentThread().getName()
                               + " Interrupted @ populateBufferAsync",
                           ex);
                       enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex);
-                      Thread.currentThread().interrupt();
                     } else {
                       LOG.log(
                           Level.WARNING,
@@ -1512,6 +1550,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);
@@ -1928,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<BigQueryArrowBatchWrapper> 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<BigQueryArrowBatchWrapper> 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<BigQueryFieldValueListWrapper> 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<BigQueryFieldValueListWrapper> 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)

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<BigQueryArrowBatchWrapper> 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())