fix(bigquery-jdbc): poll the JSON reference queue and release processed finalizers

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

@@ -17,58 +17,48 @@
 package com.google.cloud.bigquery.jdbc;
 
 import com.google.api.core.InternalApi;
-import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
+import com.google.cloud.bigquery.jdbc.BigQueryResultSetFinalizers.ResultSetFinalizer;
+import com.google.common.annotations.VisibleForTesting;
 import java.lang.ref.Reference;
 import java.lang.ref.ReferenceQueue;
 import java.util.List;
 
 /**
- * This class is used to add polling threads required for polling the reference queues associated
- * with the two ResultSets
+ * Polls a {@link ReferenceQueue} of result-set phantom references and releases the resources they
+ * own. One instance (thread) is started per result-set type.
  */
 @InternalApi
 class BigQueryDaemonPollingTask extends Thread {
 
   private static final BigQueryJdbcCustomLogger LOG =
       new BigQueryJdbcCustomLogger(BigQueryDaemonPollingTask.class.getName());
 
-  static ReferenceQueue<BigQueryArrowResultSet> referenceQueueArrowRs;
-  static ReferenceQueue<BigQueryJsonResultSet> referenceQueueJsonRs;
-  static List<BigQueryResultSetFinalizers.ArrowResultSetFinalizer> arrowRsFinalizers;
-  static List<BigQueryResultSetFinalizers.JsonResultSetFinalizer> jsonRsFinalizers;
+  // Per-thread state. Each daemon owns exactly one queue and the list that keeps its
+  // references strongly reachable until they have been processed.
+  private final ReferenceQueue<?> referenceQueue;
+  private final List<? extends Reference<?>> finalizers;
 
+  // Singleton guard: the two daemons are started at most once per JVM.
   private static BigQueryDaemonPollingTask arrowDaemon;
   private static BigQueryDaemonPollingTask jsonDaemon;
-
   private static final Object LOCK = new Object();
 
-  private BigQueryDaemonPollingTask(
-      List<BigQueryResultSetFinalizers.ArrowResultSetFinalizer> arrowRsFinalizers,
-      ReferenceQueue<BigQueryArrowResultSet> referenceQueueArrowRs) {
-    super("BigQuery-GC-Daemon-Arrow");
-    BigQueryDaemonPollingTask.referenceQueueArrowRs = referenceQueueArrowRs;
-    BigQueryDaemonPollingTask.arrowRsFinalizers = arrowRsFinalizers;
-    setDaemon(true);
-  }
-
-  private BigQueryDaemonPollingTask(
-      ReferenceQueue<BigQueryJsonResultSet> referenceQueueJsonRs,
-      List<BigQueryResultSetFinalizers.JsonResultSetFinalizer> jsonRsFinalizers) {
-    super("BigQuery-GC-Daemon-Json");
-    BigQueryDaemonPollingTask.referenceQueueJsonRs = referenceQueueJsonRs;
-    BigQueryDaemonPollingTask.jsonRsFinalizers = jsonRsFinalizers;
+  @VisibleForTesting
+  BigQueryDaemonPollingTask(
+      String threadName,
+      ReferenceQueue<?> referenceQueue,
+      List<? extends Reference<?>> finalizers) {
+    super(threadName);
+    this.referenceQueue = referenceQueue;
+    this.finalizers = finalizers;
     setDaemon(true);
   }
 
   /**
-   * ThreadSafe method which creates two instances of polling task, one each for each type of
-   * ResultSet
+   * Thread-safe. Starts one polling thread per result-set type. Subsequent calls are no-ops, even
+   * if different queues are passed.
    *
-   * @param referenceQueueArrowRs ReferenceQueue for ArrowResultSet
-   * @param referenceQueueJsonRs ReferenceQueue for JsonResultSet
-   * @param arrowRsFinalizers Finalizer for ArrowResultSet
-   * @param jsonRsFinalizers Finalizer for JsonResultSet
-   * @return true if the tasks were created
+   * @return true if the daemons were started by this call
    */
   public static boolean startGcDaemonTask(
       ReferenceQueue<BigQueryArrowResultSet> referenceQueueArrowRs,
@@ -77,45 +67,42 @@ public static boolean startGcDaemonTask(
       List<BigQueryResultSetFinalizers.JsonResultSetFinalizer> jsonRsFinalizers) {
     LOG.finest("++enter++");
     synchronized (LOCK) {
-      // 2 Background threads will be required to monitor the respective queues
       if (arrowDaemon == null && jsonDaemon == null) {
-        arrowDaemon = new BigQueryDaemonPollingTask(arrowRsFinalizers, referenceQueueArrowRs);
+        arrowDaemon =
+            new BigQueryDaemonPollingTask(
+                "BigQuery-GC-Daemon-Arrow", referenceQueueArrowRs, arrowRsFinalizers);
         arrowDaemon.start();
 
-        jsonDaemon = new BigQueryDaemonPollingTask(referenceQueueJsonRs, jsonRsFinalizers);
+        jsonDaemon =
+            new BigQueryDaemonPollingTask(
+                "BigQuery-GC-Daemon-Json", referenceQueueJsonRs, jsonRsFinalizers);
         jsonDaemon.start();
-
         return true;
       }
     }
-    // Task(s) are already initialised
     return false;
   }
 
   @Override
   public void run() {
-
-    Reference<?> reference;
+    LOG.finest("++enter++");
     try {
-      LOG.finest("++enter++");
-      // poll for Arrow ResultSets
-      if (referenceQueueArrowRs != null) {
-
-        while ((reference = referenceQueueArrowRs.remove()) != null) {
-          LOG.fine("Clearing Arrow ResultSet reference " + referenceQueueArrowRs);
-          ((BigQueryResultSetFinalizers.ArrowResultSetFinalizer) reference).finalizeResources();
-          reference.clear();
-        }
-      }
-      // poll for JSON ResultSets
-      else if (referenceQueueJsonRs != null) {
-        while ((reference = referenceQueueJsonRs.remove()) != null) {
-          LOG.fine("Clearing Json ResultSet reference " + referenceQueueJsonRs);
-          ((BigQueryResultSetFinalizers.JsonResultSetFinalizer) reference).finalizeResources();
+      while (true) {
+        Reference<?> reference =
+            referenceQueue.remove(); // InterruptedException escapes to outer catch
+        try {
+          LOG.fine("Clearing ResultSet reference " + reference);
+          if (reference instanceof ResultSetFinalizer) {
+            ((ResultSetFinalizer) reference).finalizeResources();
+          } else {
+            LOG.warning("Unexpected reference type in GC queue: " + reference.getClass());
+          }
+        } catch (RuntimeException ex) {
+          LOG.warning(ex, "Failed to finalize ResultSet resources; continuing");
+        } finally {
           reference.clear();
+          finalizers.remove(reference);
         }
-      } else {
-        throw new BigQueryJdbcRuntimeException("Null Reference Queue");
       }
     } catch (InterruptedException ex) {
       LOG.warning(ex, "Interrupted in GC daemon task");

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

@@ -26,8 +26,14 @@ class BigQueryResultSetFinalizers {
   private static final BigQueryJdbcResultSetLogger LOG =
       BigQueryJdbcResultSetLogger.getLogger(BigQueryResultSetFinalizers.class);
 
+  /** Implemented by every PhantomReference the GC daemon knows how to clean up. */
+  interface ResultSetFinalizer {
+    void finalizeResources();
+  }
+
   @InternalApi
-  static class ArrowResultSetFinalizer extends PhantomReference<BigQueryArrowResultSet> {
+  static class ArrowResultSetFinalizer extends PhantomReference<BigQueryArrowResultSet>
+      implements ResultSetFinalizer {
     Future<?> ownedTask;
 
     public ArrowResultSetFinalizer(
@@ -48,7 +54,8 @@ public void finalizeResources() {
   }
 
   @InternalApi
-  static class JsonResultSetFinalizer extends PhantomReference<BigQueryJsonResultSet> {
+  static class JsonResultSetFinalizer extends PhantomReference<BigQueryJsonResultSet>
+      implements ResultSetFinalizer {
     Future<?>[] ownedTasks;
 
     public JsonResultSetFinalizer(

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

@@ -152,9 +152,9 @@ public class BigQueryStatement extends BigQueryNoOpsStatement {
   static ReferenceQueue<BigQueryArrowResultSet> referenceQueueArrowRs = new ReferenceQueue<>();
   static ReferenceQueue<BigQueryJsonResultSet> referenceQueueJsonRs = new ReferenceQueue<>();
   static List<BigQueryResultSetFinalizers.ArrowResultSetFinalizer> arrowResultSetFinalizers =
-      new ArrayList<>();
+      Collections.synchronizedList(new ArrayList<>());
   static List<BigQueryResultSetFinalizers.JsonResultSetFinalizer> jsonResultSetFinalizers =
-      new ArrayList<>();
+      Collections.synchronizedList(new ArrayList<>());
 
   private static final ThreadFactory JDBC_THREAD_FACTORY =
       new BigQueryThreadFactory("BigQuery-Thread-");

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDaemonPollingTaskTest.java

@@ -17,11 +17,17 @@
 package com.google.cloud.bigquery.jdbc;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.verify;
 
 import com.google.cloud.bigquery.jdbc.BigQueryResultSetFinalizers.ArrowResultSetFinalizer;
+import com.google.cloud.bigquery.jdbc.BigQueryResultSetFinalizers.JsonResultSetFinalizer;
 import java.lang.ref.ReferenceQueue;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.Future;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -36,8 +42,8 @@ public class BigQueryDaemonPollingTaskTest {
   public void setUp() {
     referenceQueueArrowRs = new ReferenceQueue<>();
     referenceQueueJsonRs = new ReferenceQueue<>();
-    arrowResultSetFinalizers = new ArrayList<>();
-    jsonResultSetFinalizers = new ArrayList<>();
+    arrowResultSetFinalizers = Collections.synchronizedList(new ArrayList<>());
+    jsonResultSetFinalizers = Collections.synchronizedList(new ArrayList<>());
   }
 
   @Test
@@ -60,4 +66,83 @@ public void testStartGcDaemonTask() {
                 jsonResultSetFinalizers))
         .isFalse();
   }
+
+  @Test
+  public void testJsonDaemonFinalizesAndReleasesReference() throws Exception {
+    Future<?> worker1 = mock(Future.class);
+    Future<?> worker2 = mock(Future.class);
+    JsonResultSetFinalizer finalizer =
+        new JsonResultSetFinalizer(null, referenceQueueJsonRs, new Future<?>[] {worker1, worker2});
+    jsonResultSetFinalizers.add(finalizer);
+    BigQueryDaemonPollingTask daemon =
+        new BigQueryDaemonPollingTask(
+            "test-json-daemon", referenceQueueJsonRs, jsonResultSetFinalizers);
+    daemon.start();
+    try {
+      finalizer.enqueue(); // simulate the GC enqueuing the phantom reference
+      verify(worker1, timeout(5000)).cancel(true);
+      verify(worker2, timeout(5000)).cancel(true);
+      awaitEmpty(jsonResultSetFinalizers);
+    } finally {
+      daemon.interrupt();
+    }
+  }
+
+  @Test
+  public void testArrowDaemonFinalizesAndReleasesReference() throws Exception {
+    Future<?> worker = mock(Future.class);
+    ArrowResultSetFinalizer finalizer =
+        new ArrowResultSetFinalizer(null, referenceQueueArrowRs, worker);
+    arrowResultSetFinalizers.add(finalizer);
+    BigQueryDaemonPollingTask daemon =
+        new BigQueryDaemonPollingTask(
+            "test-arrow-daemon", referenceQueueArrowRs, arrowResultSetFinalizers);
+    daemon.start();
+    try {
+      finalizer.enqueue();
+      verify(worker, timeout(5000)).cancel(true);
+      awaitEmpty(arrowResultSetFinalizers);
+    } finally {
+      daemon.interrupt();
+    }
+  }
+
+  @Test
+  public void testEachDaemonPollsOnlyItsOwnQueue() throws Exception {
+    Future<?> jsonWorker = mock(Future.class);
+    Future<?> arrowWorker = mock(Future.class);
+    JsonResultSetFinalizer jsonFinalizer =
+        new JsonResultSetFinalizer(null, referenceQueueJsonRs, new Future<?>[] {jsonWorker});
+    ArrowResultSetFinalizer arrowFinalizer =
+        new ArrowResultSetFinalizer(null, referenceQueueArrowRs, arrowWorker);
+    jsonResultSetFinalizers.add(jsonFinalizer);
+    arrowResultSetFinalizers.add(arrowFinalizer);
+    BigQueryDaemonPollingTask jsonDaemon =
+        new BigQueryDaemonPollingTask("test-json", referenceQueueJsonRs, jsonResultSetFinalizers);
+    BigQueryDaemonPollingTask arrowDaemon =
+        new BigQueryDaemonPollingTask(
+            "test-arrow", referenceQueueArrowRs, arrowResultSetFinalizers);
+    jsonDaemon.start();
+    arrowDaemon.start();
+    try {
+      jsonFinalizer.enqueue();
+      arrowFinalizer.enqueue();
+      verify(jsonWorker, timeout(5000)).cancel(true);
+      verify(arrowWorker, timeout(5000)).cancel(true);
+      awaitEmpty(jsonResultSetFinalizers);
+      awaitEmpty(arrowResultSetFinalizers);
+    } finally {
+      jsonDaemon.interrupt();
+      arrowDaemon.interrupt();
+    }
+  }
+
+  /** Small poll loop; the list is drained by the daemon thread shortly after cancel() runs. */
+  private static void awaitEmpty(List<?> list) throws InterruptedException {
+    long deadline = System.currentTimeMillis() + 5000;
+    while (!list.isEmpty() && System.currentTimeMillis() < deadline) {
+      Thread.sleep(20);
+    }
+    assertThat(list).isEmpty();
+  }
 }