caohy1988
diff --git a/‎core/src/main/java/com/google/adk/plugins/agentanalytics/BatchProcessor.java‎
Lines changed: 192 additions & 14 deletions b/‎core/src/main/java/com/google/adk/plugins/agentanalytics/BatchProcessor.java‎
Lines changed: 192 additions & 14 deletions
@@ -24,6 +24,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.api.core.ApiFuture;
 import com.google.cloud.bigquery.storage.v1.AppendRowsResponse;
 import com.google.cloud.bigquery.storage.v1.Exceptions.AppendSerializationError;
 import com.google.cloud.bigquery.storage.v1.StreamWriter;
@@ -37,8 +38,12 @@
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Consumer;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.apache.arrow.memory.BufferAllocator;
@@ -55,6 +60,7 @@
 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
+import org.jspecify.annotations.Nullable;
 
 /** Handles asynchronous batching and writing of events to BigQuery. */
 class BatchProcessor implements AutoCloseable {
@@ -63,27 +69,53 @@ class BatchProcessor implements AutoCloseable {
   private final StreamWriter writer;
   private final int batchSize;
   private final Duration flushInterval;
+  private final Duration shutdownTimeout;
   @VisibleForTesting final BlockingQueue<Map<String, Object>> queue;
   private final ScheduledExecutorService executor;
   @VisibleForTesting final BufferAllocator allocator;
-  final AtomicBoolean flushLock = new AtomicBoolean(false);
+  // Mutual exclusion for flush; a ReentrantLock (not a CAS flag) so close() can WAIT, bounded,
+  // for an in-flight flush instead of guessing from queue emptiness.
+  private final ReentrantLock flushMutex = new ReentrantLock();
+  private final AtomicBoolean closed = new AtomicBoolean(false);
+  // Set by close() before it waits for the flush mutex: whichever party releases the mutex last
+  // (close, or an in-flight flush that outlived close's deadline) performs the actual teardown.
+  private final AtomicBoolean teardownRequested = new AtomicBoolean(false);
+  private final AtomicBoolean tornDown = new AtomicBoolean(false);
+  // While closing, bounds every drain/in-flight append to the remaining close budget.
+  private volatile @Nullable Instant closeDeadline;
+  // Delivered the FINAL drop-stat snapshot when teardown actually completes; see closeAndFold.
+  private volatile @Nullable Consumer<ImmutableMap<String, Long>> onFinalStats;
+  // Owner-preserving detached close for the StreamWriter (see teardownOnce). PluginState provides
+  // an implementation that is bounded, never blocks, and guarantees every writer's close
+  // eventually runs (a StreamWriter owns an internal client and a NON-DAEMON append thread that
+  // only ConnectionWorker.close() stops — abandoning one leaks process-level resources).
+  private final Consumer<StreamWriter> writerCloser;
+  // The periodic flush task; stored so per-invocation close() can cancel it instead of leaving a
+  // scheduled task retaining this (closed) processor until plugin-wide shutdown.
+  private volatile @Nullable ScheduledFuture<?> flushTask;
   private final Schema arrowSchema;
   private final VectorSchemaRoot root;
 
   // Drop accounting so hosts can programmatically detect lost analytics rows.
   private final AtomicLong droppedQueueFull = new AtomicLong();
   private final AtomicLong droppedAppendError = new AtomicLong();
   private final AtomicLong droppedSerializationError = new AtomicLong();
+  private final AtomicLong droppedAfterClose = new AtomicLong();
+  private final AtomicLong droppedShutdownTimeout = new AtomicLong();
 
   public BatchProcessor(
       StreamWriter writer,
       int batchSize,
       Duration flushInterval,
       int queueMaxSize,
-      ScheduledExecutorService executor) {
+      ScheduledExecutorService executor,
+      Duration shutdownTimeout,
+      Consumer<StreamWriter> writerCloser) {
     this.writer = writer;
+    this.writerCloser = writerCloser;
     this.batchSize = batchSize;
     this.flushInterval = flushInterval;
+    this.shutdownTimeout = shutdownTimeout;
     this.queue = new LinkedBlockingQueue<>(queueMaxSize);
     this.executor = executor;
     // It's safe to use Long.MAX_VALUE here as this is a top-level RootAllocator,
@@ -95,8 +127,7 @@ public BatchProcessor(
   }
 
   public void start() {
-    @SuppressWarnings("unused")
-    var unused =
+    this.flushTask =
         executor.scheduleWithFixedDelay(
             () -> {
               try {
@@ -111,19 +142,26 @@ public void start() {
   }
 
   public void append(Map<String, Object> row) {
+    if (closed.get()) {
+      // The owning invocation has already been finalized; accept-and-drop with accounting rather
+      // than silently enqueueing into a processor whose final drain has already run.
+      droppedAfterClose.incrementAndGet();
+      logger.warning("BatchProcessor is closed, dropping late event.");
+      return;
+    }
     if (!queue.offer(row)) {
       droppedQueueFull.incrementAndGet();
       logger.warning("BigQuery event queue is full, dropping event.");
       return;
     }
-    if (queue.size() >= batchSize && !flushLock.get()) {
+    if (queue.size() >= batchSize && !flushMutex.isLocked()) {
       executor.execute(this::flush);
     }
   }
 
   public void flush() {
-    // Acquire the flushLock. If another flush is already in progress, return immediately.
-    if (!flushLock.compareAndSet(false, true)) {
+    // Acquire the flush mutex. If another flush is already in progress, return immediately.
+    if (!flushMutex.tryLock()) {
       return;
     }
     try {
@@ -145,7 +183,16 @@ public void flush() {
         }
         root.setRowCount(batch.size());
         try (ArrowRecordBatch recordBatch = new VectorUnloader(root).getRecordBatch()) {
-          AppendRowsResponse result = writer.append(recordBatch).get();
+          // Bound the append so one stuck Storage Write RPC cannot block the flush path (and,
+          // during close(), the final drain) indefinitely.
+          ApiFuture<AppendRowsResponse> appendFuture = writer.append(recordBatch);
+          AppendRowsResponse result;
+          try {
+            result = appendFuture.get(appendTimeoutMillis(), MILLISECONDS);
+          } catch (TimeoutException e) {
+            appendFuture.cancel(true);
+            throw e;
+          }
           if (result.hasError()) {
             droppedAppendError.addAndGet(batch.size());
             logger.severe("BigQuery append error: " + result.getError().getMessage());
@@ -185,13 +232,32 @@ public void flush() {
         root.clear();
       }
     } finally {
-      flushLock.set(false);
-      if (queue.size() >= batchSize && !flushLock.get()) {
+      flushMutex.unlock();
+      // Deferred teardown: close() timed out waiting for this flush, transferring ownership of
+      // the final resource teardown (and drop-stat delivery) to us.
+      if (teardownRequested.get()) {
+        teardownOnce();
+      }
+      if (queue.size() >= batchSize && !flushMutex.isLocked()) {
         executor.execute(this::flush);
       }
     }
   }
 
+  /**
+   * Per-append deadline: normally {@code shutdownTimeout}; once close() has started, capped to the
+   * remaining close budget so the final drain cannot exceed the caller's bound.
+   */
+  private long appendTimeoutMillis() {
+    long timeoutMillis = shutdownTimeout.toMillis();
+    Instant deadline = this.closeDeadline;
+    if (deadline != null) {
+      long remaining = Duration.between(Instant.now(), deadline).toMillis();
+      timeoutMillis = Math.max(1, Math.min(timeoutMillis, remaining));
+    }
+    return timeoutMillis;
+  }
+
   private void populateVector(FieldVector vector, int index, Object value) {
     if (value == null || (value instanceof JsonNode jsonNode && jsonNode.isNull())) {
       vector.setNull(index);
@@ -266,13 +332,113 @@ ImmutableMap<String, Long> getDropStats() {
     return ImmutableMap.of(
         "queue_full", droppedQueueFull.get(),
         "append_error", droppedAppendError.get(),
-        "serialization_error", droppedSerializationError.get());
+        "serialization_error", droppedSerializationError.get(),
+        "after_close", droppedAfterClose.get(),
+        "shutdown_timeout", droppedShutdownTimeout.get());
+  }
+
+  /**
+   * Closes the processor and delivers the FINAL drop-stat snapshot to {@code statsConsumer} when
+   * teardown actually completes — which may be after this call returns, if an in-flight flush still
+   * owns the resources when the shutdownTimeout deadline expires (ownership of the teardown then
+   * transfers to that flush). This guarantees counters recorded by that last flush (e.g. its append
+   * failure) are included in the delivered snapshot exactly once.
+   */
+  void closeAndFold(Consumer<ImmutableMap<String, Long>> statsConsumer) {
+    closeAndFold(statsConsumer, Instant.now().plus(shutdownTimeout));
+  }
+
+  /**
+   * Deadline-accepting variant: the caller passes ONE absolute deadline shared across a larger
+   * shutdown operation (pending-task waits, sibling processors, executor termination), so this
+   * processor's drain consumes only the remaining budget instead of restarting a fresh
+   * shutdownTimeout.
+   */
+  void closeAndFold(Consumer<ImmutableMap<String, Long>> statsConsumer, Instant deadline) {
+    this.onFinalStats = statsConsumer;
+    close(deadline);
   }
 
   @Override
   public void close() {
-    while (this.queue != null && !this.queue.isEmpty()) {
+    close(Instant.now().plus(shutdownTimeout));
+  }
+
+  private void close(Instant drainDeadline) {
+    // Idempotent: ensureInvocationCompleted and plugin-wide shutdown may both close a processor.
+    if (!closed.compareAndSet(false, true)) {
+      return;
+    }
+    // Cancel the periodic flush task so a completed invocation does not leave a scheduled task
+    // retaining this processor (and its writer) until plugin-wide shutdown.
+    ScheduledFuture<?> task = this.flushTask;
+    if (task != null) {
+      task.cancel(false);
+    }
+    // Final drain, bounded by the caller's absolute deadline rather than looping until empty.
+    // Publishing the deadline caps every drain/in-flight append to the remaining close budget.
+    this.closeDeadline = drainDeadline;
+    while (!this.queue.isEmpty() && Instant.now().isBefore(drainDeadline)) {
       this.flush();
+      if (!this.queue.isEmpty()) {
+        // Another thread may hold the flush mutex; back off briefly instead of spinning.
+        try {
+          Thread.sleep(10);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          break;
+        }
+      }
+    }
+    int remaining = this.queue.size();
+    if (remaining > 0) {
+      droppedShutdownTimeout.addAndGet(remaining);
+      this.queue.clear();
+      logger.severe(
+          "Dropping " + remaining + " rows: final drain did not complete within shutdownTimeout.");
+    }
+    // Teardown ownership: request it, then try to acquire the flush mutex within the remaining
+    // budget. If acquired, no flush is active and we tear down here. If the wait expires, the
+    // in-flight flush performs the teardown (and final stats delivery) when it releases the mutex
+    // — resources are never destroyed underneath an active flush, and counters that flush records
+    // are still included in the final snapshot.
+    teardownRequested.set(true);
+    boolean acquired = false;
+    long waitMillis = Math.max(1, Duration.between(Instant.now(), drainDeadline).toMillis());
+    try {
+      acquired = flushMutex.tryLock(waitMillis, MILLISECONDS);
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
+    if (acquired) {
+      try {
+        teardownOnce();
+      } finally {
+        flushMutex.unlock();
+      }
+    } else {
+      logger.severe(
+          "Deferring resource teardown to the in-flight flush: it did not release the flush mutex"
+              + " within shutdownTimeout.");
+      // The flush may have released the mutex between the timed wait expiring and the request
+      // flag becoming visible to it; re-check so the teardown is never lost.
+      if (flushMutex.tryLock()) {
+        try {
+          teardownOnce();
+        } finally {
+          flushMutex.unlock();
+        }
+      }
+    }
+  }
+
+  /**
+   * Tears down Arrow and writer resources and delivers the final drop-stat snapshot, exactly once,
+   * regardless of whether close() or a deferred in-flight flush gets here first.
+   */
+  private void teardownOnce() {
+    if (!tornDown.compareAndSet(false, true)) {
+      return;
     }
     if (this.allocator != null) {
       try {
@@ -289,10 +455,22 @@ public void close() {
       }
     }
     if (this.writer != null) {
+      // StreamWriter.close() can block far beyond any shutdownTimeout (it joins the writer's
+      // internal non-daemon append thread, then may wait minutes on its internal client and
+      // callback pools). Delegate to the plugin-owned closer, which detaches the close without
+      // ever abandoning the writer: cleanup ownership is guaranteed by writer admission permits
+      // acquired before construction. No further work touches the writer here: teardown only
+      // runs after appends have stopped (closed gate + flush-mutex ownership), and drop counters
+      // are final below.
+      writerCloser.accept(this.writer);
+    }
+    // Deliver the final snapshot only now, when no flush can mutate the counters anymore.
+    Consumer<ImmutableMap<String, Long>> statsConsumer = this.onFinalStats;
+    if (statsConsumer != null) {
       try {
-        this.writer.close();
+        statsConsumer.accept(getDropStats());
       } catch (RuntimeException e) {
-        logger.log(Level.SEVERE, "Failed to close BigQuery writer", e);
+        logger.log(Level.WARNING, "Failed to deliver final drop stats", e);
       }
     }
   }