Skip to content

Commit ad83e20

Browse files
caohy1988claude
caohy1988
and
claude
committed
fix(bigquery): address round-6 review — bounded writer-close service
Resolves the round-6 P1 on PR google#1344: detached StreamWriter closes no longer spawn one raw daemon thread per completed invocation. Teardown submits the close to a plugin-owned bounded closer service (2 daemon threads, 256-task backlog, core timeout) shared by all BatchProcessors, so a Storage outage converts invocation throughput into a bounded backlog instead of native-thread exhaustion. Saturation policy: never block teardown and never spawn unbounded threads — a rejected close is counted ("writer_close_saturated", unit: writers) and the StreamWriter is left for BigQueryWriteClient shutdown to reclaim at the transport level. Plugin-wide close shuts the closer service down within the same shared absolute deadline, interrupting closes still blocked past it. Tests: 25 completed invocations with indefinitely blocked writer closes grow closer threads by at most the pool bound (delta-based, immune to sibling-test instances), leave no retained processors, and plugin shutdown still completes within its bound with the close backlog pending. Focused BQAA suite 171 passed (previous: 170); full core 1,627 passed / 0 failed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 544690e commit ad83e20

4 files changed

Lines changed: 206 additions & 30 deletions

File tree

‎core/src/main/java/com/google/adk/plugins/agentanalytics/BatchProcessor.java‎

Lines changed: 38 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,9 @@
3636
import java.util.List;
3737
import java.util.Map;
3838
import java.util.concurrent.BlockingQueue;
39+
import java.util.concurrent.ExecutorService;
3940
import java.util.concurrent.LinkedBlockingQueue;
41+
import java.util.concurrent.RejectedExecutionException;
4042
import java.util.concurrent.ScheduledExecutorService;
4143
import java.util.concurrent.ScheduledFuture;
4244
import java.util.concurrent.TimeoutException;
@@ -85,6 +87,13 @@ class BatchProcessor implements AutoCloseable {
8587
private volatile @Nullable Instant closeDeadline;
8688
// Delivered the FINAL drop-stat snapshot when teardown actually completes; see closeAndFold.
8789
private volatile @Nullable Consumer<ImmutableMap<String, Long>> onFinalStats;
90+
// Plugin-owned BOUNDED service for detached StreamWriter closes (see teardownOnce): a raw
91+
// thread per teardown would grow without bound during a Storage outage (one blocked closer per
92+
// completed invocation) until native-thread exhaustion.
93+
private final ExecutorService writerCloseExecutor;
94+
// Writers whose detached close could not be admitted (closer service saturated); they are left
95+
// for BigQueryWriteClient shutdown to reclaim at the transport level. Unit: writers, not rows.
96+
private final AtomicLong writerCloseSaturated = new AtomicLong();
8897
// The periodic flush task; stored so per-invocation close() can cancel it instead of leaving a
8998
// scheduled task retaining this (closed) processor until plugin-wide shutdown.
9099
private volatile @Nullable ScheduledFuture<?> flushTask;
@@ -104,8 +113,10 @@ public BatchProcessor(
104113
Duration flushInterval,
105114
int queueMaxSize,
106115
ScheduledExecutorService executor,
107-
Duration shutdownTimeout) {
116+
Duration shutdownTimeout,
117+
ExecutorService writerCloseExecutor) {
108118
this.writer = writer;
119+
this.writerCloseExecutor = writerCloseExecutor;
109120
this.batchSize = batchSize;
110121
this.flushInterval = flushInterval;
111122
this.shutdownTimeout = shutdownTimeout;
@@ -327,7 +338,8 @@ ImmutableMap<String, Long> getDropStats() {
327338
"append_error", droppedAppendError.get(),
328339
"serialization_error", droppedSerializationError.get(),
329340
"after_close", droppedAfterClose.get(),
330-
"shutdown_timeout", droppedShutdownTimeout.get());
341+
"shutdown_timeout", droppedShutdownTimeout.get(),
342+
"writer_close_saturated", writerCloseSaturated.get());
331343
}
332344

333345
/**
@@ -449,23 +461,31 @@ private void teardownOnce() {
449461
}
450462
if (this.writer != null) {
451463
// StreamWriter.close() can block far beyond any shutdownTimeout (it joins the append
452-
// thread, then may wait minutes on the client and callback pools). Detach it to a daemon
453-
// thread so teardown — and therefore the public close()/deferred-flush path — honors the
454-
// absolute deadline. No further work touches the writer: teardown only runs after appends
455-
// have stopped (closed gate + flush-mutex ownership), and drop counters are final below.
464+
// thread, then may wait minutes on the client and callback pools). Detach it to the
465+
// plugin-owned BOUNDED closer service so teardown — and therefore the public
466+
// close()/deferred-flush path — honors the absolute deadline without growing one raw
467+
// thread per completed invocation during a Storage outage. No further work touches the
468+
// writer: teardown only runs after appends have stopped (closed gate + flush-mutex
469+
// ownership), and drop counters are final below.
456470
StreamWriter writerToClose = this.writer;
457-
Thread writerCloser =
458-
new Thread(
459-
() -> {
460-
try {
461-
writerToClose.close();
462-
} catch (RuntimeException e) {
463-
logger.log(Level.SEVERE, "Failed to close BigQuery writer", e);
464-
}
465-
},
466-
"bq-analytics-writer-close");
467-
writerCloser.setDaemon(true);
468-
writerCloser.start();
471+
try {
472+
writerCloseExecutor.execute(
473+
() -> {
474+
try {
475+
writerToClose.close();
476+
} catch (RuntimeException e) {
477+
logger.log(Level.SEVERE, "Failed to close BigQuery writer", e);
478+
}
479+
});
480+
} catch (RejectedExecutionException e) {
481+
// Saturation policy: never block teardown and never spawn unbounded threads. The writer
482+
// is left for BigQueryWriteClient shutdown to reclaim; account it so operators can see
483+
// sustained close backlog.
484+
writerCloseSaturated.incrementAndGet();
485+
logger.severe(
486+
"Writer-close service saturated; leaving StreamWriter to be reclaimed by client"
487+
+ " shutdown.");
488+
}
469489
}
470490
// Deliver the final snapshot only now, when no flush can mutate the counters anymore.
471491
Consumer<ImmutableMap<String, Long>> statsConsumer = this.onFinalStats;

‎core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,9 +66,15 @@ class PluginState {
6666
// Idle time before threads are terminated.
6767
private static final int GCS_OFFLOAD_IDLE_TIME_SECONDS = 30;
6868

69+
// Bounded detached-close service shared by all BatchProcessors: at most this many concurrent
70+
// StreamWriter closes, with a bounded backlog; saturation is accounted per processor.
71+
private static final int WRITER_CLOSE_MAX_THREADS = 2;
72+
private static final int WRITER_CLOSE_QUEUE_SIZE = 256;
73+
6974
private final BigQueryLoggerConfig config;
7075
private final ScheduledExecutorService executor;
7176
private final ExecutorService offloadExecutor;
77+
private final ThreadPoolExecutor writerCloseExecutor;
7278
private final BigQueryWriteClient writeClient;
7379
private static final AtomicLong threadCounter = new AtomicLong(0);
7480
// Map of invocation ID to BatchProcessor.
@@ -135,13 +141,30 @@ synchronized boolean isFinalized() {
135141
// continuation completed after the invocation was already finalized.
136142
private final AtomicLong droppedWriterCreateError = new AtomicLong();
137143
private final AtomicLong droppedAfterFinalize = new AtomicLong();
144+
// Writers whose detached close was rejected by the saturated closer service (unit: writers).
145+
private final AtomicLong writerCloseSaturated = new AtomicLong();
138146

139147
PluginState(BigQueryLoggerConfig config) throws IOException {
140148
this.config = config;
141149
this.executor =
142150
Executors.newScheduledThreadPool(
143151
2, r -> new Thread(r, "bq-analytics-plugin-" + threadCounter.getAndIncrement()));
144152
this.offloadExecutor = createGcsOffloadThreadPool();
153+
this.writerCloseExecutor =
154+
new ThreadPoolExecutor(
155+
WRITER_CLOSE_MAX_THREADS,
156+
WRITER_CLOSE_MAX_THREADS,
157+
30,
158+
SECONDS,
159+
new ArrayBlockingQueue<>(WRITER_CLOSE_QUEUE_SIZE),
160+
r -> {
161+
Thread t =
162+
new Thread(r, "bq-analytics-writer-close-" + threadCounter.getAndIncrement());
163+
t.setDaemon(true);
164+
return t;
165+
},
166+
new ThreadPoolExecutor.AbortPolicy());
167+
this.writerCloseExecutor.allowCoreThreadTimeOut(true);
145168
// One write client per plugin instance, shared by all invocations.
146169
this.writeClient = createWriteClient(config);
147170
this.processedInvocations =
@@ -259,7 +282,8 @@ BatchProcessor getBatchProcessor(String invocationId) {
259282
config.batchFlushInterval(),
260283
config.queueMaxSize(),
261284
executor,
262-
config.shutdownTimeout());
285+
config.shutdownTimeout(),
286+
writerCloseExecutor);
263287
p.start();
264288
return p;
265289
});
@@ -345,7 +369,8 @@ void appendRow(InvocationLifecycle lifecycle, String invocationId, Map<String, O
345369
config.batchFlushInterval(),
346370
config.queueMaxSize(),
347371
executor,
348-
config.shutdownTimeout());
372+
config.shutdownTimeout(),
373+
writerCloseExecutor);
349374
p.start();
350375
return p;
351376
});
@@ -498,6 +523,7 @@ private void foldStats(ImmutableMap<String, Long> stats) {
498523
droppedSerializationError.addAndGet(stats.getOrDefault("serialization_error", 0L));
499524
droppedAfterClose.addAndGet(stats.getOrDefault("after_close", 0L));
500525
droppedShutdownTimeout.addAndGet(stats.getOrDefault("shutdown_timeout", 0L));
526+
writerCloseSaturated.addAndGet(stats.getOrDefault("writer_close_saturated", 0L));
501527
}
502528

503529
/**
@@ -512,20 +538,23 @@ ImmutableMap<String, Long> getDropStats() {
512538
long serializationError = droppedSerializationError.get();
513539
long afterClose = droppedAfterClose.get();
514540
long shutdownTimeout = droppedShutdownTimeout.get();
541+
long closeSaturated = writerCloseSaturated.get();
515542
for (BatchProcessor processor : getBatchProcessors()) {
516543
ImmutableMap<String, Long> stats = processor.getDropStats();
517544
queueFull += stats.getOrDefault("queue_full", 0L);
518545
appendError += stats.getOrDefault("append_error", 0L);
519546
serializationError += stats.getOrDefault("serialization_error", 0L);
520547
afterClose += stats.getOrDefault("after_close", 0L);
521548
shutdownTimeout += stats.getOrDefault("shutdown_timeout", 0L);
549+
closeSaturated += stats.getOrDefault("writer_close_saturated", 0L);
522550
}
523551
return ImmutableMap.<String, Long>builder()
524552
.put("queue_full", queueFull)
525553
.put("append_error", appendError)
526554
.put("serialization_error", serializationError)
527555
.put("after_close", afterClose)
528556
.put("shutdown_timeout", shutdownTimeout)
557+
.put("writer_close_saturated", closeSaturated)
529558
.put("writer_create_error", droppedWriterCreateError.get())
530559
.put("late_after_finalize", droppedAfterFinalize.get())
531560
.buildOrThrow();
@@ -593,9 +622,19 @@ private Completable closeInternal(java.time.Instant closeDeadline) {
593622
} else {
594623
offloadExecutor.shutdownNow();
595624
}
625+
// Detached writer closes are best-effort by design; interrupt any still blocked
626+
// past the deadline (writeClient.close() below reclaims at the transport level).
627+
writerCloseExecutor.shutdown();
628+
remainingMillis =
629+
java.time.Duration.between(java.time.Instant.now(), closeDeadline).toMillis();
630+
if (remainingMillis <= 0
631+
|| !writerCloseExecutor.awaitTermination(remainingMillis, MILLISECONDS)) {
632+
var unusedPending = writerCloseExecutor.shutdownNow();
633+
}
596634
} catch (InterruptedException e) {
597635
executor.shutdownNow();
598636
offloadExecutor.shutdownNow();
637+
var unusedPending = writerCloseExecutor.shutdownNow();
599638
Thread.currentThread().interrupt();
600639
}
601640

‎core/src/test/java/com/google/adk/plugins/agentanalytics/BatchProcessorTest.java‎

Lines changed: 74 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
import java.util.HashMap;
4040
import java.util.List;
4141
import java.util.Map;
42+
import java.util.concurrent.ExecutorService;
4243
import java.util.concurrent.Executors;
4344
import java.util.concurrent.ScheduledExecutorService;
4445
import java.util.concurrent.ScheduledFuture;
@@ -78,12 +79,21 @@ public class BatchProcessorTest {
7879
private Schema schema;
7980
private Handler mockHandler;
8081

82+
private ExecutorService writerCloseExecutor;
83+
8184
@Before
8285
public void setUp() {
8386
executor = Executors.newScheduledThreadPool(1);
87+
writerCloseExecutor = Executors.newCachedThreadPool();
8488
batchProcessor =
8589
new BatchProcessor(
86-
mockWriter, 10, Duration.ofMinutes(1), 100, executor, Duration.ofSeconds(10));
90+
mockWriter,
91+
10,
92+
Duration.ofMinutes(1),
93+
100,
94+
executor,
95+
Duration.ofSeconds(10),
96+
writerCloseExecutor);
8797
schema = BigQuerySchema.getArrowSchema();
8898

8999
when(mockWriter.append(any(ArrowRecordBatch.class)))
@@ -295,7 +305,13 @@ public void append_triggersFlushWhenBatchSizeReached() {
295305
ScheduledExecutorService mockExecutor = mock(ScheduledExecutorService.class);
296306
BatchProcessor bp =
297307
new BatchProcessor(
298-
mockWriter, 2, Duration.ofMinutes(1), 10, mockExecutor, Duration.ofSeconds(10));
308+
mockWriter,
309+
2,
310+
Duration.ofMinutes(1),
311+
10,
312+
mockExecutor,
313+
Duration.ofSeconds(10),
314+
writerCloseExecutor);
299315

300316
Map<String, Object> row = new HashMap<>();
301317
bp.append(row);
@@ -369,7 +385,13 @@ public void flush_handlesAllocationFailure() throws Exception {
369385
public void close_flushesAndClosesResources() throws Exception {
370386
try (BatchProcessor bp =
371387
new BatchProcessor(
372-
mockWriter, 10, Duration.ofMinutes(1), 100, executor, Duration.ofSeconds(10))) {
388+
mockWriter,
389+
10,
390+
Duration.ofMinutes(1),
391+
100,
392+
executor,
393+
Duration.ofSeconds(10),
394+
writerCloseExecutor)) {
373395
Map<String, Object> row = new HashMap<>();
374396
row.put("event_type", "CLOSE_EVENT");
375397
bp.append(row);
@@ -388,7 +410,13 @@ public void close_cancelsPeriodicFlushTask() {
388410
.scheduleWithFixedDelay(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class));
389411
BatchProcessor bp =
390412
new BatchProcessor(
391-
mockWriter, 10, Duration.ofMinutes(1), 10, mockExecutor, Duration.ofSeconds(1));
413+
mockWriter,
414+
10,
415+
Duration.ofMinutes(1),
416+
10,
417+
mockExecutor,
418+
Duration.ofSeconds(1),
419+
writerCloseExecutor);
392420
bp.start();
393421

394422
bp.close();
@@ -402,7 +430,13 @@ public void close_cancelsPeriodicFlushTask() {
402430
public void close_isIdempotent() {
403431
BatchProcessor bp =
404432
new BatchProcessor(
405-
mockWriter, 10, Duration.ofMinutes(1), 10, executor, Duration.ofSeconds(1));
433+
mockWriter,
434+
10,
435+
Duration.ofMinutes(1),
436+
10,
437+
executor,
438+
Duration.ofSeconds(1),
439+
writerCloseExecutor);
406440

407441
bp.close();
408442
bp.close();
@@ -414,7 +448,13 @@ public void close_isIdempotent() {
414448
public void append_afterClose_dropsWithAccounting() {
415449
BatchProcessor bp =
416450
new BatchProcessor(
417-
mockWriter, 10, Duration.ofMinutes(1), 10, executor, Duration.ofSeconds(1));
451+
mockWriter,
452+
10,
453+
Duration.ofMinutes(1),
454+
10,
455+
executor,
456+
Duration.ofSeconds(1),
457+
writerCloseExecutor);
418458
bp.close();
419459

420460
Map<String, Object> row = new HashMap<>();
@@ -430,7 +470,13 @@ public void flush_neverCompletingAppend_isBoundedByShutdownTimeout() {
430470
when(mockWriter.append(any(ArrowRecordBatch.class))).thenReturn(SettableApiFuture.create());
431471
BatchProcessor bp =
432472
new BatchProcessor(
433-
mockWriter, 1, Duration.ofMinutes(1), 10, executor, Duration.ofMillis(200));
473+
mockWriter,
474+
1,
475+
Duration.ofMinutes(1),
476+
10,
477+
executor,
478+
Duration.ofMillis(200),
479+
writerCloseExecutor);
434480
Map<String, Object> row = new HashMap<>();
435481
row.put("event_type", "STUCK_EVENT");
436482
bp.queue.offer(row);
@@ -460,7 +506,13 @@ public void close_waitsForInFlightFlushBeforeTeardown() throws Exception {
460506
.append(any(ArrowRecordBatch.class));
461507
BatchProcessor bp =
462508
new BatchProcessor(
463-
mockWriter, 1, Duration.ofMinutes(1), 10, executor, Duration.ofSeconds(5));
509+
mockWriter,
510+
1,
511+
Duration.ofMinutes(1),
512+
10,
513+
executor,
514+
Duration.ofSeconds(5),
515+
writerCloseExecutor);
464516

465517
Map<String, Object> row = new HashMap<>();
466518
row.put("event_type", "IN_FLIGHT_EVENT");
@@ -499,7 +551,13 @@ public void close_deadlineExpired_defersTeardownAndStatsToInFlightFlush() throws
499551
.append(any(ArrowRecordBatch.class));
500552
BatchProcessor bp =
501553
new BatchProcessor(
502-
mockWriter, 1, Duration.ofMinutes(1), 10, executor, Duration.ofMillis(300));
554+
mockWriter,
555+
1,
556+
Duration.ofMinutes(1),
557+
10,
558+
executor,
559+
Duration.ofMillis(300),
560+
writerCloseExecutor);
503561

504562
Map<String, Object> row = new HashMap<>();
505563
row.put("event_type", "STUCK_EVENT");
@@ -547,7 +605,13 @@ public void close_blockingWriterClose_doesNotBlockTeardown() throws Exception {
547605
.close();
548606
BatchProcessor bp =
549607
new BatchProcessor(
550-
mockWriter, 10, Duration.ofMinutes(1), 10, executor, Duration.ofMillis(300));
608+
mockWriter,
609+
10,
610+
Duration.ofMinutes(1),
611+
10,
612+
executor,
613+
Duration.ofMillis(300),
614+
writerCloseExecutor);
551615

552616
long start = System.nanoTime();
553617
bp.close();

0 commit comments

Comments
 (0)