feat(bigquery-jdbc): migrate statement execution thread tracking to connection-scoped executor

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java

@@ -38,6 +38,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.Future;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.FieldVector;
@@ -80,8 +81,8 @@ class BigQueryArrowResultSet extends BigQueryBaseResultSet {
   // Decoder object will be reused to avoid re-allocation and too much garbage collection.
   private VectorSchemaRoot vectorSchemaRoot;
   private VectorLoader vectorLoader;
-  // producer thread's reference
-  private final Thread ownedThread;
+  // producer task's reference
+  private final Future<?> ownedTask;
 
   private BigQueryArrowResultSet(
       Schema schema,
@@ -93,7 +94,7 @@ private BigQueryArrowResultSet(
       boolean isNested,
       int fromIndex,
       int toIndexExclusive,
-      Thread ownedThread,
+      Future<?> ownedTask,
       BigQuery bigQuery,
       Job job)
       throws SQLException {
@@ -105,7 +106,7 @@ private BigQueryArrowResultSet(
     this.fromIndex = fromIndex;
     this.toIndexExclusive = toIndexExclusive;
     this.nestedRowIndex = fromIndex - 1;
-    this.ownedThread = ownedThread;
+    this.ownedTask = ownedTask;
     if (!isNested && arrowSchema != null) {
       try {
         this.arrowDeserializer = new ArrowDeserializer(arrowSchema);
@@ -127,10 +128,10 @@ static BigQueryArrowResultSet of(
       long totalRows,
       BigQueryStatement statement,
       BlockingQueue<BigQueryArrowBatchWrapper> buffer,
-      Thread ownedThread,
+      Future<?> ownedTask,
       BigQuery bigQuery)
       throws SQLException {
-    return of(schema, arrowSchema, totalRows, statement, buffer, ownedThread, bigQuery, null);
+    return of(schema, arrowSchema, totalRows, statement, buffer, ownedTask, bigQuery, null);
   }
 
   static BigQueryArrowResultSet of(
@@ -139,7 +140,7 @@ static BigQueryArrowResultSet of(
       long totalRows,
       BigQueryStatement statement,
       BlockingQueue<BigQueryArrowBatchWrapper> buffer,
-      Thread ownedThread,
+      Future<?> ownedTask,
       BigQuery bigQuery,
       Job job)
       throws SQLException {
@@ -153,7 +154,7 @@ static BigQueryArrowResultSet of(
         false,
         -1,
         -1,
-        ownedThread,
+        ownedTask,
         bigQuery,
         job);
   }
@@ -165,7 +166,7 @@ static BigQueryArrowResultSet of(
     this.currentNestedBatch = null;
     this.fromIndex = 0;
     this.toIndexExclusive = 0;
-    this.ownedThread = null;
+    this.ownedTask = null;
     this.arrowDeserializer = null;
     this.vectorSchemaRoot = null;
     this.vectorLoader = null;
@@ -484,9 +485,9 @@ private String formatRangeElement(Object element, StandardSQLTypeName elementTyp
   public void close() {
     LOG.fineTrace("close", () -> String.format("Closing BigqueryArrowResultSet %s.", this));
     this.isClosed = true;
-    if (ownedThread != null && !ownedThread.isInterrupted()) {
-      // interrupt the producer thread when result set is closed
-      ownedThread.interrupt();
+    if (ownedTask != null) {
+      // cancel the producer task when result set is closed
+      ownedTask.cancel(true);
     }
     super.close();
   }

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java

@@ -348,7 +348,11 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
       this.headerProvider = createHeaderProvider();
       this.bigQuery = getBigQueryConnection();
       this.metadataExecutor = BigQueryJdbcMdc.newFixedThreadPool(metadataFetchThreadCount);
-      this.queryExecutor = BigQueryJdbcMdc.newCachedThreadPool();
+      // Use a bounded cached thread pool to prevent unbounded thread creation (and OOMs)
+      // under heavy load, while ensuring a limit (e.g., 100) high enough to prevent deadlocks
+      // between interdependent producer/consumer tasks (like nextPageWorker and
+      // populateBufferWorker).
+      this.queryExecutor = BigQueryJdbcMdc.newBoundedCachedThreadPool(100);
     }
   }
 

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java

@@ -72,6 +72,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -944,7 +945,7 @@ public ResultSet getProcedures(
 
     Thread fetcherThread = new Thread(procedureFetcher, "getProcedures-fetcher-" + catalog);
     BigQueryJsonResultSet resultSet =
-        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
+        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, wrapThread(fetcherThread));
 
     fetcherThread.start();
     LOG.info("Started background thread for getProcedures");
@@ -1206,7 +1207,7 @@ public ResultSet getProcedureColumns(
     Thread fetcherThread =
         new Thread(procedureColumnFetcher, "getProcedureColumns-fetcher-" + catalog);
     BigQueryJsonResultSet resultSet =
-        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
+        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, wrapThread(fetcherThread));
 
     fetcherThread.start();
     LOG.info("Started background thread for getProcedureColumns for catalog: " + catalog);
@@ -1877,7 +1878,7 @@ public ResultSet getTables(
 
     Thread fetcherThread = new Thread(tableFetcher, "getTables-fetcher-" + effectiveCatalog);
     BigQueryJsonResultSet resultSet =
-        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
+        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, wrapThread(fetcherThread));
 
     fetcherThread.start();
     LOG.info("Started background thread for getTables");
@@ -2017,7 +2018,8 @@ public ResultSet getCatalogs() {
     populateQueue(catalogRows, queue, schemaFields);
     signalEndOfData(queue, schemaFields);
 
-    return BigQueryJsonResultSet.of(catalogsSchema, catalogRows.size(), queue, null, new Thread[0]);
+    return BigQueryJsonResultSet.of(
+        catalogsSchema, catalogRows.size(), queue, null, new Future<?>[0]);
   }
 
   Schema defineGetCatalogsSchema() {
@@ -2049,7 +2051,7 @@ public ResultSet getTableTypes() {
     signalEndOfData(queue, tableTypesSchema.getFields());
 
     return BigQueryJsonResultSet.of(
-        tableTypesSchema, tableTypeRows.size(), queue, null, new Thread[0]);
+        tableTypesSchema, tableTypeRows.size(), queue, null, new Future<?>[0]);
   }
 
   static Schema defineGetTableTypesSchema() {
@@ -2203,7 +2205,7 @@ public ResultSet getColumns(
 
     Thread fetcherThread = new Thread(columnFetcher, "getColumns-fetcher-" + effectiveCatalog);
     BigQueryJsonResultSet resultSet =
-        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
+        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, wrapThread(fetcherThread));
 
     fetcherThread.start();
     LOG.info("Started background thread for getColumns");
@@ -2718,7 +2720,7 @@ public ResultSet getTypeInfo() {
     populateQueue(typeInfoRows, queue, schemaFields);
     signalEndOfData(queue, schemaFields);
     return BigQueryJsonResultSet.of(
-        typeInfoSchema, typeInfoRows.size(), queue, null, new Thread[0]);
+        typeInfoSchema, typeInfoRows.size(), queue, null, new Future<?>[0]);
   }
 
   Schema defineGetTypeInfoSchema() {
@@ -3713,7 +3715,7 @@ public ResultSet getSchemas(String catalog, String schemaPattern) {
 
     Thread fetcherThread = new Thread(schemaFetcher, "getSchemas-fetcher-" + catalog);
     BigQueryJsonResultSet resultSet =
-        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
+        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, wrapThread(fetcherThread));
 
     fetcherThread.start();
     LOG.info("Started background thread for getSchemas");
@@ -3832,7 +3834,7 @@ public ResultSet getClientInfoProperties() {
       signalEndOfData(queue, resultSchemaFields);
     }
     return BigQueryJsonResultSet.of(
-        resultSchema, collectedResults.size(), queue, null, new Thread[0]);
+        resultSchema, collectedResults.size(), queue, null, new Future<?>[0]);
   }
 
   Schema defineGetClientInfoPropertiesSchema() {
@@ -4007,7 +4009,7 @@ public ResultSet getFunctions(String catalog, String schemaPattern, String funct
 
     Thread fetcherThread = new Thread(functionFetcher, "getFunctions-fetcher-" + catalog);
     BigQueryJsonResultSet resultSet =
-        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
+        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, wrapThread(fetcherThread));
 
     fetcherThread.start();
     LOG.info("Started background thread for getFunctions");
@@ -4261,7 +4263,7 @@ public ResultSet getFunctionColumns(
     Thread fetcherThread =
         new Thread(functionColumnFetcher, "getFunctionColumns-fetcher-" + catalog);
     BigQueryJsonResultSet resultSet =
-        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
+        BigQueryJsonResultSet.of(resultSchema, -1, queue, null, wrapThread(fetcherThread));
 
     fetcherThread.start();
     LOG.info("Started background thread for getFunctionColumns for catalog: " + catalog);
@@ -5264,4 +5266,88 @@ private void loadDriverVersionProperties() {
       throw ex;
     }
   }
+
+  // TODO(keshav): This is a temporary compatibility bridge to wrap raw Threads into Futures.
+  // This should be removed when BigQueryDatabaseMetaData is refactored to use the ExecutorService
+  // directly.
+  static Future<?>[] wrapThread(final Thread thread) {
+    if (thread == null) {
+      return null;
+    }
+    return new Future<?>[] {
+      new Future<Object>() {
+        private volatile boolean cancelled = false;
+
+        @Override
+        public synchronized boolean cancel(boolean mayInterruptIfRunning) {
+          if (cancelled || thread.getState() == Thread.State.TERMINATED) {
+            return false;
+          }
+          cancelled = true;
+          if (mayInterruptIfRunning) {
+            thread.interrupt();
+          }
+          return true;
+        }
+
+        @Override
+        public boolean isCancelled() {
+          return cancelled;
+        }
+
+        @Override
+        public boolean isDone() {
+          return cancelled || thread.getState() == Thread.State.TERMINATED;
+        }
+
+        @Override
+        public Object get() throws InterruptedException, CancellationException {
+          if (isCancelled()) {
+            throw new CancellationException();
+          }
+          while (thread.getState() != Thread.State.TERMINATED) {
+            if (isCancelled()) {
+              throw new CancellationException();
+            }
+            if (thread.getState() == Thread.State.NEW) {
+              Thread.sleep(50);
+            } else {
+              thread.join(50);
+            }
+          }
+          return null;
+        }
+
+        @Override
+        public Object get(long timeout, TimeUnit unit)
+            throws InterruptedException, CancellationException, TimeoutException {
+          if (isCancelled()) {
+            throw new CancellationException();
+          }
+          long remainingNanos = unit.toNanos(timeout);
+          long deadline = System.nanoTime() + remainingNanos;
+          while (thread.getState() != Thread.State.TERMINATED) {
+            if (isCancelled()) {
+              throw new CancellationException();
+            }
+            if (remainingNanos <= 0) {
+              throw new TimeoutException();
+            }
+            long remainingMillis = TimeUnit.NANOSECONDS.toMillis(remainingNanos);
+            if (remainingMillis <= 0) {
+              remainingMillis = 1;
+            }
+            long delay = Math.min(remainingMillis, 50);
+            if (thread.getState() == Thread.State.NEW) {
+              Thread.sleep(delay);
+            } else {
+              thread.join(delay);
+            }
+            remainingNanos = deadline - System.nanoTime();
+          }
+          return null;
+        }
+      }
+    };
+  }
 }

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java

@@ -81,25 +81,25 @@ static ExecutorService newFixedThreadPool(int nThreads) {
   }
 
   /**
-   * Creates a new cached thread pool ExecutorService that automatically propagates MDC connection
-   * context from the submitting thread to the executing thread.
+   * Creates a new bounded cached thread pool ExecutorService that automatically propagates MDC
+   * connection context from the submitting thread to the executing thread.
    */
-  static ExecutorService newCachedThreadPool(ThreadFactory threadFactory) {
+  static ExecutorService newBoundedCachedThreadPool(int maxThreads, ThreadFactory threadFactory) {
     return new MdcThreadPoolExecutor(
         0,
-        Integer.MAX_VALUE,
+        maxThreads,
         60L,
         TimeUnit.SECONDS,
         new java.util.concurrent.SynchronousQueue<>(),
         new MdcThreadFactory(threadFactory));
   }
 
   /**
-   * Creates a new cached thread pool ExecutorService that automatically propagates MDC connection
-   * context from the submitting thread to the executing thread.
+   * Creates a new bounded cached thread pool ExecutorService that automatically propagates MDC
+   * connection context from the submitting thread to the executing thread.
    */
-  static ExecutorService newCachedThreadPool() {
-    return newCachedThreadPool(Executors.defaultThreadFactory());
+  static ExecutorService newBoundedCachedThreadPool(int maxThreads) {
+    return newBoundedCachedThreadPool(maxThreads, Executors.defaultThreadFactory());
   }
 
   private static class MdcThreadFactory implements ThreadFactory {