diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDaemonPollingTask.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDaemonPollingTask.java index d3e21f79ad17..d1782e5fa747 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDaemonPollingTask.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDaemonPollingTask.java @@ -17,14 +17,15 @@ 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 { @@ -32,43 +33,32 @@ class BigQueryDaemonPollingTask extends Thread { private static final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(BigQueryDaemonPollingTask.class.getName()); - static ReferenceQueue referenceQueueArrowRs; - static ReferenceQueue referenceQueueJsonRs; - static List arrowRsFinalizers; - static List 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> 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 arrowRsFinalizers, - ReferenceQueue referenceQueueArrowRs) { - super("BigQuery-GC-Daemon-Arrow"); - BigQueryDaemonPollingTask.referenceQueueArrowRs = referenceQueueArrowRs; - BigQueryDaemonPollingTask.arrowRsFinalizers = arrowRsFinalizers; - setDaemon(true); - } - - private BigQueryDaemonPollingTask( - ReferenceQueue referenceQueueJsonRs, - List jsonRsFinalizers) { - super("BigQuery-GC-Daemon-Json"); - BigQueryDaemonPollingTask.referenceQueueJsonRs = referenceQueueJsonRs; - BigQueryDaemonPollingTask.jsonRsFinalizers = jsonRsFinalizers; + @VisibleForTesting + BigQueryDaemonPollingTask( + String threadName, + ReferenceQueue referenceQueue, + List> 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 referenceQueueArrowRs, @@ -77,45 +67,42 @@ public static boolean startGcDaemonTask( List 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"); diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryResultSetFinalizers.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryResultSetFinalizers.java index 57f764979f53..6749f4a6c9b2 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryResultSetFinalizers.java +++ b/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 { + static class ArrowResultSetFinalizer extends PhantomReference + implements ResultSetFinalizer { Future ownedTask; public ArrowResultSetFinalizer( @@ -48,7 +54,8 @@ public void finalizeResources() { } @InternalApi - static class JsonResultSetFinalizer extends PhantomReference { + static class JsonResultSetFinalizer extends PhantomReference + implements ResultSetFinalizer { Future[] ownedTasks; public JsonResultSetFinalizer( 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 d3e312d79feb..b070356cb221 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 @@ -152,9 +152,9 @@ public class BigQueryStatement extends BigQueryNoOpsStatement { static ReferenceQueue referenceQueueArrowRs = new ReferenceQueue<>(); static ReferenceQueue referenceQueueJsonRs = new ReferenceQueue<>(); static List arrowResultSetFinalizers = - new ArrayList<>(); + Collections.synchronizedList(new ArrayList<>()); static List jsonResultSetFinalizers = - new ArrayList<>(); + Collections.synchronizedList(new ArrayList<>()); private static final ThreadFactory JDBC_THREAD_FACTORY = new BigQueryThreadFactory("BigQuery-Thread-"); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDaemonPollingTaskTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDaemonPollingTaskTest.java index 29941152e7c5..66c1c0b50c7e 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDaemonPollingTaskTest.java +++ b/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(); + } }