fix(bigquery): address round-5 review — completion-ordered cleanup, deadline-safe Storage Writer entry and teardown

Resolves both round-5 P1 findings on PR google#1344:

1. Finalization is completion-ordered: RxJava's doFinally notifies the
   downstream BEFORE running its action, so blockingAwait()/subscribers
   could observe success while lifecycle marking, processor draining,
   and teardown were still running. Both ensureInvocationCompleted and
   plugin-wide close now run an idempotent cleanup action via
   onErrorComplete().andThen(Completable.fromAction(cleanup)) — cleanup
   completes before the returned Completable does — with the same
   action in a trailing doFinally to cover disposal.
2. Storage Writer calls obey the public bound. Admission: StreamWriter's
   default LimitExceededBehavior.Block parks append() up to five minutes
   on inflight quota; the writer is now built with ThrowException, so
   quota saturation surfaces as an accounted append failure instead of
   an unbounded block. Teardown: StreamWriter.close() can block minutes
   (thread join + client/callback-pool waits); it is detached to a
   daemon thread so teardownOnce — and the public close()/deferred-flush
   path — honors the absolute deadline. No further work touches the
   writer at that point, and drop counters are final before delivery.

Tests: new completion-ordering regression (async pending-task completion
with latch-blocked cleanup; the returned Completable must not complete
until cleanup releases); blocking StreamWriter.close() must not block
the public close() while the writer is still closed eventually; the
multi-processor close bound is tightened to <900ms so the old
per-processor-deadline behavior fails; writer.close() verifications use
timeout() now that teardown is detached. Focused BQAA suite 170 passed
(previous: 168); full core 1,626 passed / 0 failed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: caohy1988

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

@@ -448,11 +448,24 @@ private void teardownOnce() {
       }
     }
     if (this.writer != null) {
-      try {
-        this.writer.close();
-      } catch (RuntimeException e) {
-        logger.log(Level.SEVERE, "Failed to close BigQuery writer", e);
-      }
+      // StreamWriter.close() can block far beyond any shutdownTimeout (it joins the append
+      // thread, then may wait minutes on the client and callback pools). Detach it to a daemon
+      // thread so teardown — and therefore the public close()/deferred-flush path — honors the
+      // absolute deadline. No further work touches the writer: teardown only runs after appends
+      // have stopped (closed gate + flush-mutex ownership), and drop counters are final below.
+      StreamWriter writerToClose = this.writer;
+      Thread writerCloser =
+          new Thread(
+              () -> {
+                try {
+                  writerToClose.close();
+                } catch (RuntimeException e) {
+                  logger.log(Level.SEVERE, "Failed to close BigQuery writer", e);
+                }
+              },
+              "bq-analytics-writer-close");
+      writerCloser.setDaemon(true);
+      writerCloser.start();
     }
     // Deliver the final snapshot only now, when no flush can mutate the counters anymore.
     Consumer<ImmutableMap<String, Long>> statsConsumer = this.onFinalStats;

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

@@ -21,6 +21,7 @@
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
+import com.google.api.gax.batching.FlowController;
 import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.FixedHeaderProvider;
@@ -34,6 +35,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import io.reactivex.rxjava3.core.Completable;
+import io.reactivex.rxjava3.functions.Action;
 import java.io.IOException;
 import java.util.Collection;
 import java.util.Map;
@@ -209,6 +211,11 @@ protected StreamWriter createWriter() {
     try {
       return StreamWriter.newBuilder(streamName, writeClient)
           .setTraceId(BigQueryUtils.getVersionHeaderValue() + ":" + UUID.randomUUID())
+          // Nonblocking admission: the default LimitExceededBehavior.Block parks append() for up
+          // to five minutes waiting on inflight quota, escaping every shutdownTimeout bound this
+          // plugin enforces. ThrowException surfaces quota saturation as an append failure, which
+          // the flush path catches and accounts as dropped rows.
+          .setLimitExceededBehavior(FlowController.LimitExceededBehavior.ThrowException)
           .setRetrySettings(retrySettings)
           .setWriterSchema(BigQuerySchema.getArrowSchema())
           // Route Storage Write append RPCs to the dataset's region. Without this, appends to any
@@ -422,20 +429,13 @@ private Completable finalizeInvocation(String invocationId, java.time.Instant fi
               CompletableFuture.allOf(tasks.toArray(new CompletableFuture<?>[0])));
     }
     logger.fine("Waiting for pending tasks to complete for invocation ID: " + invocationId);
-    return tasksState
-        .timeout(config.shutdownTimeout().toMillis(), MILLISECONDS)
-        .doOnError(
-            e -> {
-              if (e instanceof TimeoutException) {
-                logger.log(
-                    Level.WARNING,
-                    "Timeout while waiting for pending tasks to complete for invocation ID: "
-                        + invocationId,
-                    e);
-              }
-            })
-        .onErrorComplete()
-        .doFinally(
+    // Idempotent cleanup shared by the completion-ordered andThen (normal path) and doFinally
+    // (disposal path): RxJava's doFinally notifies the downstream FIRST and runs its action
+    // afterwards, so relying on it alone would let blockingAwait()/subscribers observe success
+    // while finalization is still running. andThen(fromAction) runs the cleanup BEFORE the
+    // returned Completable completes.
+    Action cleanup =
+        runOnce(
             () -> {
               // Mark the durable lifecycle token FIRST (before removing the processor), so any
               // continuation that later finds no mapping is guaranteed to observe the terminal
@@ -464,6 +464,32 @@ private Completable finalizeInvocation(String invocationId, java.time.Instant fi
               logger.fine("Removing pending tasks for invocation ID: " + invocationId);
               pendingTasks.remove(invocationId);
             });
+    return tasksState
+        .timeout(config.shutdownTimeout().toMillis(), MILLISECONDS)
+        .doOnError(
+            e -> {
+              if (e instanceof TimeoutException) {
+                logger.log(
+                    Level.WARNING,
+                    "Timeout while waiting for pending tasks to complete for invocation ID: "
+                        + invocationId,
+                    e);
+              }
+            })
+        .onErrorComplete()
+        .andThen(Completable.fromAction(cleanup))
+        .doFinally(cleanup);
+  }
+
+  /** Wraps an action so repeated invocations (completion path + disposal path) run it once. */
+  private static Action runOnce(Action delegate) {
+    java.util.concurrent.atomic.AtomicBoolean ran =
+        new java.util.concurrent.atomic.AtomicBoolean(false);
+    return () -> {
+      if (ran.compareAndSet(false, true)) {
+        delegate.run();
+      }
+    };
   }
 
   private void foldStats(ImmutableMap<String, Long> stats) {
@@ -523,17 +549,9 @@ private Completable closeInternal(java.time.Instant closeDeadline) {
           Completable.fromCompletionStage(
               CompletableFuture.allOf(tasks.toArray(new CompletableFuture<?>[0])));
     }
-    return tasksState
-        .timeout(config.shutdownTimeout().toMillis(), MILLISECONDS)
-        .doOnError(
-            e -> {
-              if (e instanceof TimeoutException) {
-                logger.log(
-                    Level.WARNING, "Timeout while waiting for pending tasks to complete.", e);
-              }
-            })
-        .onErrorComplete()
-        .doFinally(
+    // Completion-ordered cleanup shared with the disposal path; see finalizeInvocation.
+    Action cleanup =
+        runOnce(
             () -> {
               for (InvocationLifecycle lifecycle : lifecycles.values()) {
                 lifecycle.markFinalized();
@@ -589,5 +607,17 @@ private Completable closeInternal(java.time.Instant closeDeadline) {
                 logger.log(Level.WARNING, "Failed to close GCS offloader", e);
               }
             });
+    return tasksState
+        .timeout(config.shutdownTimeout().toMillis(), MILLISECONDS)
+        .doOnError(
+            e -> {
+              if (e instanceof TimeoutException) {
+                logger.log(
+                    Level.WARNING, "Timeout while waiting for pending tasks to complete.", e);
+              }
+            })
+        .onErrorComplete()
+        .andThen(Completable.fromAction(cleanup))
+        .doFinally(cleanup);
   }
 }

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

@@ -24,7 +24,6 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -377,7 +376,7 @@ public void close_flushesAndClosesResources() throws Exception {
     }
 
     verify(mockWriter).append(any(ArrowRecordBatch.class));
-    verify(mockWriter).close();
+    verify(mockWriter, org.mockito.Mockito.timeout(2000)).close();
   }
 
   @Test
@@ -408,7 +407,7 @@ public void close_isIdempotent() {
     bp.close();
     bp.close();
 
-    verify(mockWriter, times(1)).close();
+    verify(mockWriter, org.mockito.Mockito.timeout(2000).times(1)).close();
   }
 
   @Test
@@ -481,7 +480,7 @@ public void close_waitsForInFlightFlushBeforeTeardown() throws Exception {
         "close should have waited for the in-flight flush, took " + elapsedMs + "ms",
         elapsedMs >= 100);
     assertTrue("close should be bounded, took " + elapsedMs + "ms", elapsedMs < 5_000);
-    verify(mockWriter).close();
+    verify(mockWriter, org.mockito.Mockito.timeout(2000)).close();
   }
 
   @Test
@@ -530,6 +529,35 @@ public void close_deadlineExpired_defersTeardownAndStatsToInFlightFlush() throws
         statsDelivered.await(5, java.util.concurrent.TimeUnit.SECONDS));
     inFlight.join(2000);
     assertEquals(1L, (long) finalStats.get().get("append_error"));
-    verify(mockWriter).close();
+    verify(mockWriter, org.mockito.Mockito.timeout(2000)).close();
+  }
+
+  @Test
+  public void close_blockingWriterClose_doesNotBlockTeardown() throws Exception {
+    // The real StreamWriter.close() can block for minutes (thread join + client/pool waits); the
+    // public close() must not inherit that, while the writer still gets closed eventually.
+    java.util.concurrent.CountDownLatch writerClosed = new java.util.concurrent.CountDownLatch(1);
+    org.mockito.Mockito.doAnswer(
+            invocation -> {
+              Thread.sleep(1500);
+              writerClosed.countDown();
+              return null;
+            })
+        .when(mockWriter)
+        .close();
+    BatchProcessor bp =
+        new BatchProcessor(
+            mockWriter, 10, Duration.ofMinutes(1), 10, executor, Duration.ofMillis(300));
+
+    long start = System.nanoTime();
+    bp.close();
+    long elapsedMs = (System.nanoTime() - start) / 1_000_000;
+
+    assertTrue(
+        "close() must not block on StreamWriter.close(), took " + elapsedMs + "ms",
+        elapsedMs < 1_000);
+    assertTrue(
+        "the writer must still be closed eventually",
+        writerClosed.await(5, java.util.concurrent.TimeUnit.SECONDS));
   }
 }

core/src/test/java/com/google/adk/plugins/agentanalytics/PluginStateTest.java

@@ -535,8 +535,63 @@ protected StreamWriter createWriter() {
     slowState.close().blockingAwait();
     long elapsedMs = (System.nanoTime() - start) / 1_000_000;
 
+    // Old behavior was ~one deadline PER processor (>=1000ms for two 500ms drains) before
+    // executor waits; the shared absolute deadline finishes in ~one timeout plus slack.
     assertTrue(
         "multi-processor close must share one shutdownTimeout, took " + elapsedMs + "ms",
-        elapsedMs < 1_100);
+        elapsedMs < 900);
+  }
+
+  @Test
+  public void ensureInvocationCompleted_doesNotCompleteBeforeCleanupFinishes() throws Exception {
+    // RxJava's doFinally notifies the downstream BEFORE running its action; finalization must be
+    // completion-ordered so callers cannot observe success while cleanup is still running.
+    String invocationId = "inv-ordered";
+    java.util.concurrent.CountDownLatch cleanupStarted = new java.util.concurrent.CountDownLatch(1);
+    java.util.concurrent.CountDownLatch cleanupRelease = new java.util.concurrent.CountDownLatch(1);
+    BatchProcessor blockingProcessor = mock(BatchProcessor.class);
+    org.mockito.Mockito.doAnswer(
+            invocation -> {
+              cleanupStarted.countDown();
+              cleanupRelease.await(5, TimeUnit.SECONDS);
+              return null;
+            })
+        .when(blockingProcessor)
+        .closeAndFold(org.mockito.ArgumentMatchers.any(), org.mockito.ArgumentMatchers.any());
+    TestPluginState orderedState =
+        new TestPluginState(config) {
+          @Override
+          protected BatchProcessor removeProcessor(String id) {
+            return id.equals(invocationId) ? blockingProcessor : super.removeProcessor(id);
+          }
+        };
+    CompletableFuture<Void> pending = new CompletableFuture<>();
+    orderedState.addPendingTask(invocationId, pending);
+
+    java.util.concurrent.atomic.AtomicBoolean observedComplete =
+        new java.util.concurrent.atomic.AtomicBoolean(false);
+    var unused =
+        orderedState
+            .ensureInvocationCompleted(invocationId)
+            .subscribe(() -> observedComplete.set(true));
+
+    // Complete the pending task ASYNCHRONOUSLY so the chain advances on another thread and
+    // blocks inside the (latched) cleanup.
+    Thread completer = new Thread(() -> pending.complete(null));
+    completer.start();
+    assertTrue(cleanupStarted.await(2, TimeUnit.SECONDS));
+
+    // Cleanup is running but blocked: the returned Completable must NOT have completed.
+    Thread.sleep(100);
+    assertTrue(
+        "completion must not be observable before cleanup finishes", !observedComplete.get());
+
+    cleanupRelease.countDown();
+    completer.join(2000);
+    long deadline = Instant.now().plusMillis(2000).toEpochMilli();
+    while (!observedComplete.get() && Instant.now().toEpochMilli() < deadline) {
+      Thread.sleep(10);
+    }
+    assertTrue("completion must be observable after cleanup finishes", observedComplete.get());
   }
 }