Skip to content

Commit 1b7d00d

Browse files
Neenu1995Neenu1995
authored
fix(bigquery-jdbc): poll the JSON reference queue and release processed finalizers (#14719)
b/524270586
1 parent b27385e commit 1b7d00d

4 files changed

Lines changed: 139 additions & 60 deletions

File tree

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

Lines changed: 41 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -17,58 +17,48 @@
1717
package com.google.cloud.bigquery.jdbc;
1818

1919
import com.google.api.core.InternalApi;
20-
import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
20+
import com.google.cloud.bigquery.jdbc.BigQueryResultSetFinalizers.ResultSetFinalizer;
21+
import com.google.common.annotations.VisibleForTesting;
2122
import java.lang.ref.Reference;
2223
import java.lang.ref.ReferenceQueue;
2324
import java.util.List;
2425

2526
/**
26-
* This class is used to add polling threads required for polling the reference queues associated
27-
* with the two ResultSets
27+
* Polls a {@link ReferenceQueue} of result-set phantom references and releases the resources they
28+
* own. One instance (thread) is started per result-set type.
2829
*/
2930
@InternalApi
3031
class BigQueryDaemonPollingTask extends Thread {
3132

3233
private static final BigQueryJdbcCustomLogger LOG =
3334
new BigQueryJdbcCustomLogger(BigQueryDaemonPollingTask.class.getName());
3435

35-
static ReferenceQueue<BigQueryArrowResultSet> referenceQueueArrowRs;
36-
static ReferenceQueue<BigQueryJsonResultSet> referenceQueueJsonRs;
37-
static List<BigQueryResultSetFinalizers.ArrowResultSetFinalizer> arrowRsFinalizers;
38-
static List<BigQueryResultSetFinalizers.JsonResultSetFinalizer> jsonRsFinalizers;
36+
// Per-thread state. Each daemon owns exactly one queue and the list that keeps its
37+
// references strongly reachable until they have been processed.
38+
private final ReferenceQueue<?> referenceQueue;
39+
private final List<? extends Reference<?>> finalizers;
3940

41+
// Singleton guard: the two daemons are started at most once per JVM.
4042
private static BigQueryDaemonPollingTask arrowDaemon;
4143
private static BigQueryDaemonPollingTask jsonDaemon;
42-
4344
private static final Object LOCK = new Object();
4445

45-
private BigQueryDaemonPollingTask(
46-
List<BigQueryResultSetFinalizers.ArrowResultSetFinalizer> arrowRsFinalizers,
47-
ReferenceQueue<BigQueryArrowResultSet> referenceQueueArrowRs) {
48-
super("BigQuery-GC-Daemon-Arrow");
49-
BigQueryDaemonPollingTask.referenceQueueArrowRs = referenceQueueArrowRs;
50-
BigQueryDaemonPollingTask.arrowRsFinalizers = arrowRsFinalizers;
51-
setDaemon(true);
52-
}
53-
54-
private BigQueryDaemonPollingTask(
55-
ReferenceQueue<BigQueryJsonResultSet> referenceQueueJsonRs,
56-
List<BigQueryResultSetFinalizers.JsonResultSetFinalizer> jsonRsFinalizers) {
57-
super("BigQuery-GC-Daemon-Json");
58-
BigQueryDaemonPollingTask.referenceQueueJsonRs = referenceQueueJsonRs;
59-
BigQueryDaemonPollingTask.jsonRsFinalizers = jsonRsFinalizers;
46+
@VisibleForTesting
47+
BigQueryDaemonPollingTask(
48+
String threadName,
49+
ReferenceQueue<?> referenceQueue,
50+
List<? extends Reference<?>> finalizers) {
51+
super(threadName);
52+
this.referenceQueue = referenceQueue;
53+
this.finalizers = finalizers;
6054
setDaemon(true);
6155
}
6256

6357
/**
64-
* ThreadSafe method which creates two instances of polling task, one each for each type of
65-
* ResultSet
58+
* Thread-safe. Starts one polling thread per result-set type. Subsequent calls are no-ops, even
59+
* if different queues are passed.
6660
*
67-
* @param referenceQueueArrowRs ReferenceQueue for ArrowResultSet
68-
* @param referenceQueueJsonRs ReferenceQueue for JsonResultSet
69-
* @param arrowRsFinalizers Finalizer for ArrowResultSet
70-
* @param jsonRsFinalizers Finalizer for JsonResultSet
71-
* @return true if the tasks were created
61+
* @return true if the daemons were started by this call
7262
*/
7363
public static boolean startGcDaemonTask(
7464
ReferenceQueue<BigQueryArrowResultSet> referenceQueueArrowRs,
@@ -77,45 +67,42 @@ public static boolean startGcDaemonTask(
7767
List<BigQueryResultSetFinalizers.JsonResultSetFinalizer> jsonRsFinalizers) {
7868
LOG.finest("++enter++");
7969
synchronized (LOCK) {
80-
// 2 Background threads will be required to monitor the respective queues
8170
if (arrowDaemon == null && jsonDaemon == null) {
82-
arrowDaemon = new BigQueryDaemonPollingTask(arrowRsFinalizers, referenceQueueArrowRs);
71+
arrowDaemon =
72+
new BigQueryDaemonPollingTask(
73+
"BigQuery-GC-Daemon-Arrow", referenceQueueArrowRs, arrowRsFinalizers);
8374
arrowDaemon.start();
8475

85-
jsonDaemon = new BigQueryDaemonPollingTask(referenceQueueJsonRs, jsonRsFinalizers);
76+
jsonDaemon =
77+
new BigQueryDaemonPollingTask(
78+
"BigQuery-GC-Daemon-Json", referenceQueueJsonRs, jsonRsFinalizers);
8679
jsonDaemon.start();
87-
8880
return true;
8981
}
9082
}
91-
// Task(s) are already initialised
9283
return false;
9384
}
9485

9586
@Override
9687
public void run() {
97-
98-
Reference<?> reference;
88+
LOG.finest("++enter++");
9989
try {
100-
LOG.finest("++enter++");
101-
// poll for Arrow ResultSets
102-
if (referenceQueueArrowRs != null) {
103-
104-
while ((reference = referenceQueueArrowRs.remove()) != null) {
105-
LOG.fine("Clearing Arrow ResultSet reference " + referenceQueueArrowRs);
106-
((BigQueryResultSetFinalizers.ArrowResultSetFinalizer) reference).finalizeResources();
107-
reference.clear();
108-
}
109-
}
110-
// poll for JSON ResultSets
111-
else if (referenceQueueJsonRs != null) {
112-
while ((reference = referenceQueueJsonRs.remove()) != null) {
113-
LOG.fine("Clearing Json ResultSet reference " + referenceQueueJsonRs);
114-
((BigQueryResultSetFinalizers.JsonResultSetFinalizer) reference).finalizeResources();
90+
while (true) {
91+
Reference<?> reference =
92+
referenceQueue.remove(); // InterruptedException escapes to outer catch
93+
try {
94+
LOG.fine("Clearing ResultSet reference " + reference);
95+
if (reference instanceof ResultSetFinalizer) {
96+
((ResultSetFinalizer) reference).finalizeResources();
97+
} else {
98+
LOG.warning("Unexpected reference type in GC queue: " + reference.getClass());
99+
}
100+
} catch (RuntimeException ex) {
101+
LOG.warning(ex, "Failed to finalize ResultSet resources; continuing");
102+
} finally {
115103
reference.clear();
104+
finalizers.remove(reference);
116105
}
117-
} else {
118-
throw new BigQueryJdbcRuntimeException("Null Reference Queue");
119106
}
120107
} catch (InterruptedException ex) {
121108
LOG.warning(ex, "Interrupted in GC daemon task");

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,14 @@ class BigQueryResultSetFinalizers {
2626
private static final BigQueryJdbcResultSetLogger LOG =
2727
BigQueryJdbcResultSetLogger.getLogger(BigQueryResultSetFinalizers.class);
2828

29+
/** Implemented by every PhantomReference the GC daemon knows how to clean up. */
30+
interface ResultSetFinalizer {
31+
void finalizeResources();
32+
}
33+
2934
@InternalApi
30-
static class ArrowResultSetFinalizer extends PhantomReference<BigQueryArrowResultSet> {
35+
static class ArrowResultSetFinalizer extends PhantomReference<BigQueryArrowResultSet>
36+
implements ResultSetFinalizer {
3137
Future<?> ownedTask;
3238

3339
public ArrowResultSetFinalizer(
@@ -48,7 +54,8 @@ public void finalizeResources() {
4854
}
4955

5056
@InternalApi
51-
static class JsonResultSetFinalizer extends PhantomReference<BigQueryJsonResultSet> {
57+
static class JsonResultSetFinalizer extends PhantomReference<BigQueryJsonResultSet>
58+
implements ResultSetFinalizer {
5259
Future<?>[] ownedTasks;
5360

5461
public JsonResultSetFinalizer(

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -152,9 +152,9 @@ public class BigQueryStatement extends BigQueryNoOpsStatement {
152152
static ReferenceQueue<BigQueryArrowResultSet> referenceQueueArrowRs = new ReferenceQueue<>();
153153
static ReferenceQueue<BigQueryJsonResultSet> referenceQueueJsonRs = new ReferenceQueue<>();
154154
static List<BigQueryResultSetFinalizers.ArrowResultSetFinalizer> arrowResultSetFinalizers =
155-
new ArrayList<>();
155+
Collections.synchronizedList(new ArrayList<>());
156156
static List<BigQueryResultSetFinalizers.JsonResultSetFinalizer> jsonResultSetFinalizers =
157-
new ArrayList<>();
157+
Collections.synchronizedList(new ArrayList<>());
158158

159159
private static final ThreadFactory JDBC_THREAD_FACTORY =
160160
new BigQueryThreadFactory("BigQuery-Thread-");

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

Lines changed: 87 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,17 @@
1717
package com.google.cloud.bigquery.jdbc;
1818

1919
import static com.google.common.truth.Truth.assertThat;
20+
import static org.mockito.Mockito.mock;
21+
import static org.mockito.Mockito.timeout;
22+
import static org.mockito.Mockito.verify;
2023

2124
import com.google.cloud.bigquery.jdbc.BigQueryResultSetFinalizers.ArrowResultSetFinalizer;
25+
import com.google.cloud.bigquery.jdbc.BigQueryResultSetFinalizers.JsonResultSetFinalizer;
2226
import java.lang.ref.ReferenceQueue;
2327
import java.util.ArrayList;
28+
import java.util.Collections;
2429
import java.util.List;
30+
import java.util.concurrent.Future;
2531
import org.junit.jupiter.api.BeforeEach;
2632
import org.junit.jupiter.api.Test;
2733

@@ -36,8 +42,8 @@ public class BigQueryDaemonPollingTaskTest {
3642
public void setUp() {
3743
referenceQueueArrowRs = new ReferenceQueue<>();
3844
referenceQueueJsonRs = new ReferenceQueue<>();
39-
arrowResultSetFinalizers = new ArrayList<>();
40-
jsonResultSetFinalizers = new ArrayList<>();
45+
arrowResultSetFinalizers = Collections.synchronizedList(new ArrayList<>());
46+
jsonResultSetFinalizers = Collections.synchronizedList(new ArrayList<>());
4147
}
4248

4349
@Test
@@ -60,4 +66,83 @@ public void testStartGcDaemonTask() {
6066
jsonResultSetFinalizers))
6167
.isFalse();
6268
}
69+
70+
@Test
71+
public void testJsonDaemonFinalizesAndReleasesReference() throws Exception {
72+
Future<?> worker1 = mock(Future.class);
73+
Future<?> worker2 = mock(Future.class);
74+
JsonResultSetFinalizer finalizer =
75+
new JsonResultSetFinalizer(null, referenceQueueJsonRs, new Future<?>[] {worker1, worker2});
76+
jsonResultSetFinalizers.add(finalizer);
77+
BigQueryDaemonPollingTask daemon =
78+
new BigQueryDaemonPollingTask(
79+
"test-json-daemon", referenceQueueJsonRs, jsonResultSetFinalizers);
80+
daemon.start();
81+
try {
82+
finalizer.enqueue(); // simulate the GC enqueuing the phantom reference
83+
verify(worker1, timeout(5000)).cancel(true);
84+
verify(worker2, timeout(5000)).cancel(true);
85+
awaitEmpty(jsonResultSetFinalizers);
86+
} finally {
87+
daemon.interrupt();
88+
}
89+
}
90+
91+
@Test
92+
public void testArrowDaemonFinalizesAndReleasesReference() throws Exception {
93+
Future<?> worker = mock(Future.class);
94+
ArrowResultSetFinalizer finalizer =
95+
new ArrowResultSetFinalizer(null, referenceQueueArrowRs, worker);
96+
arrowResultSetFinalizers.add(finalizer);
97+
BigQueryDaemonPollingTask daemon =
98+
new BigQueryDaemonPollingTask(
99+
"test-arrow-daemon", referenceQueueArrowRs, arrowResultSetFinalizers);
100+
daemon.start();
101+
try {
102+
finalizer.enqueue();
103+
verify(worker, timeout(5000)).cancel(true);
104+
awaitEmpty(arrowResultSetFinalizers);
105+
} finally {
106+
daemon.interrupt();
107+
}
108+
}
109+
110+
@Test
111+
public void testEachDaemonPollsOnlyItsOwnQueue() throws Exception {
112+
Future<?> jsonWorker = mock(Future.class);
113+
Future<?> arrowWorker = mock(Future.class);
114+
JsonResultSetFinalizer jsonFinalizer =
115+
new JsonResultSetFinalizer(null, referenceQueueJsonRs, new Future<?>[] {jsonWorker});
116+
ArrowResultSetFinalizer arrowFinalizer =
117+
new ArrowResultSetFinalizer(null, referenceQueueArrowRs, arrowWorker);
118+
jsonResultSetFinalizers.add(jsonFinalizer);
119+
arrowResultSetFinalizers.add(arrowFinalizer);
120+
BigQueryDaemonPollingTask jsonDaemon =
121+
new BigQueryDaemonPollingTask("test-json", referenceQueueJsonRs, jsonResultSetFinalizers);
122+
BigQueryDaemonPollingTask arrowDaemon =
123+
new BigQueryDaemonPollingTask(
124+
"test-arrow", referenceQueueArrowRs, arrowResultSetFinalizers);
125+
jsonDaemon.start();
126+
arrowDaemon.start();
127+
try {
128+
jsonFinalizer.enqueue();
129+
arrowFinalizer.enqueue();
130+
verify(jsonWorker, timeout(5000)).cancel(true);
131+
verify(arrowWorker, timeout(5000)).cancel(true);
132+
awaitEmpty(jsonResultSetFinalizers);
133+
awaitEmpty(arrowResultSetFinalizers);
134+
} finally {
135+
jsonDaemon.interrupt();
136+
arrowDaemon.interrupt();
137+
}
138+
}
139+
140+
/** Small poll loop; the list is drained by the daemon thread shortly after cancel() runs. */
141+
private static void awaitEmpty(List<?> list) throws InterruptedException {
142+
long deadline = System.currentTimeMillis() + 5000;
143+
while (!list.isEmpty() && System.currentTimeMillis() < deadline) {
144+
Thread.sleep(20);
145+
}
146+
assertThat(list).isEmpty();
147+
}
63148
}

0 commit comments

Comments
 (0)