fix(batch): drain retried items

bevzzz · bevzzz · commit 1a190792b569 · 2026-03-27T14:50:52.000+01:00
diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/batch/BatchContext.java b/src/main/java/io/weaviate/client6/v1/api/collections/batch/BatchContext.java
@@ -21,7 +21,6 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
-import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
@@ -81,7 +80,7 @@
  * <li>{@code null} -- context hasn't been {@link #start}ed yet. The context
  * SHOULD NOT be used in this state, as it will likely result in an NPE.
  * <li>AwaitStarted -- client's opened the stream, sent Start,
- * and is now awaiting for the server to respond with Started.
+ * and is now awaiting the server to respond with Started.
  * <li>Active -- the server is ready to accept the next Data message.
  * <li>InFlight -- the latest batch has been sent, awaiting Acks.
  * <li>OOM -- server has OOM'ed and will not accept any more data.
@@ -145,8 +144,8 @@ public final class BatchContext<PropertiesT> implements Closeable {
    * <p>
    * In the event of abrupt stream termination ({@link Recv#onError} is called),
    * the "recv" thread MAY shutdown this service in order to interrupt the "send"
-   * thread; the latter may be blocked on {@link Send#awaitCanSend} or
-   * {@link Send#awaitCanPrepareNext}.
+   * thread; the latter may be blocked on {@link State#awaitCanSend} or
+   * {@link State#awaitCanPrepareNext}.
    */
   private final ExecutorService sendService = Executors.newSingleThreadExecutor();
 
@@ -158,6 +157,9 @@ public final class BatchContext<PropertiesT> implements Closeable {
    */
   private final ScheduledExecutorService scheduledService = Executors.newScheduledThreadPool(1);
 
+  /** Executor for processing {@link Event.Results} and enqueing retried items. */
+  private final ExecutorService retryService = Executors.newSingleThreadExecutor();
+
   /** The thread that created the context. */
   private final Thread parent = Thread.currentThread();
 
@@ -208,20 +210,6 @@ public final class BatchContext<PropertiesT> implements Closeable {
   /** stateChanged notifies threads about a state transition. */
   private final Condition stateChanged = lock.newCondition();
 
-  /**
-   * Releasing a permit notifies the "sender" about an incoming
-   * {@link Event.Results} batch. Acquire a permit to await the next batch.
-   *
-   * <p>
-   * A semaphore provides signal semantics, similar to a {@link Condition},
-   * but without being associated with a predicate. This comes handy when
-   * {@link wip} is being drained after the context is closed, and the "sender"
-   * needs to be notified about incoming {@link Event.Results}; a separate
-   * condition is not necessary, as we can simply probe the {@link queue} to
-   * find out if any new items have been added to it.
-   */
-  private final Semaphore awaitResults = new Semaphore(0);
-
   /**
    * Client-side part of the current stream, created on {@link #start}.
    * Other threads MAY use stream but MUST NOT update this field on their own.
@@ -323,6 +311,11 @@ private TaskHandle add(final TaskHandle taskHandle) throws InterruptedException
       throw new DuplicateTaskException(taskHandle, existing);
     }
 
+    // Remove the task from the WIP list as soon as it completes,
+    // successfully or otherwise. Note, that TaskHandle::done future
+    // only completes exceptionally after all retries have been exhausted.
+    taskHandle.done().whenComplete((__, t) -> wip.remove(taskHandle.id()));
+
     queue.put(taskHandle);
     return taskHandle;
   }
@@ -507,6 +500,7 @@ private void shutdownNow(Exception e) {
   private void shutdownExecutors() {
     sendService.shutdownNow();
     scheduledService.shutdownNow();
+    retryService.shutdownNow();
     shutdownService.shutdownNow();
     scheduledReconnectService.shutdownNow();
   }
@@ -647,7 +641,11 @@ private void trySend() {
 
           TaskHandle task = queue.take();
 
-          if (task == TaskHandle.POISON) {
+          if (task == TaskHandle.END_RESULTS) {
+            // This marker is only relevant when sent after POISON,
+            // as we process retried items in drainWip.
+            continue;
+          } else if (task == TaskHandle.POISON) {
             assert closed : "queue poisoned before the context is closed";
 
             log.debug("Took poison");
@@ -673,10 +671,9 @@ private void trySend() {
             return;
           }
 
-          Data data = task.data();
-          batch.add(data);
+          batch.add(task.data());
 
-          // Retred tasks already exist in the WIP list, replacing them is redundant.
+          // Retried tasks already exist in the WIP list, replacing them is redundant.
           wip.putIfAbsent(task.id(), task);
         }
       } catch (InterruptedException ignored) {
@@ -695,14 +692,14 @@ private void trySend() {
      * in the while-loop is that nothing is sent <i>unless</i> the batch is full.
      */
     private void send() throws InterruptedException {
-      log.atInfo()
-          .addKeyValue("batch_size_total_bytes", batch::sizeBytes)
-          .log("Send next batch");
-
       // Continue flushing until we get the batch to not a "not full" state.
       // This is to account for the backlog, which might re-fill the batch
       // after .clear().
       while (batch.isFull()) {
+        log.atInfo()
+            .addKeyValue("batch_size_items", batch::size)
+            .addKeyValue("batch_size_total_bytes", batch::sizeBytes)
+            .log("Send next batch");
         flush();
       }
       assert !batch.isFull() : "batch is full after send";
@@ -729,24 +726,22 @@ private void drainWip() throws InterruptedException {
             .addKeyValue("batch_size_total_items", batch::size)
             .addKeyValue("wip_tasks", wip::size)
             .log("Await Results");
-        awaitResults.acquire();
 
         TaskHandle task;
-        while ((task = queue.poll()) != null) {
-          Data data = task.data();
-          batch.add(data);
+        while ((task = queue.take()) != TaskHandle.END_RESULTS) {
+          batch.add(task.data());
         }
 
         assert batch.size() <= wip.size() : "batch has more items than wip";
 
         if (batch.size() == wip.size()) {
           // This means the batch already contains all items in WIP,
           // and no more tasks will be added to the queue until the
-          // current ones are send.
+          // current ones are sent.
           drain();
         } else {
           // Only sends if the batch is full. If the batch is not full,
-          // the we can keep accumulating items from the failed tasks.
+          // then we can keep accumulating items from the failed tasks.
           send();
         }
       }
@@ -773,7 +768,7 @@ private void drain() throws InterruptedException {
     }
 
     /**
-     * Block until the current state allows {@link State#canSend},
+     * Block until the current state allows {@link State#canSendNext},
      * then prepare the batch, send it, and set InFlight state.
      * Block until the current state allows {@link State#canPrepareNext}.
      *
@@ -953,24 +948,28 @@ private void onResults(Event.Results results) {
           .addKeyValue("wip_tasks", wip::size)
           .log("Received Results");
 
-      // Remove successfully completed tasks from the WIP list and mark them done.
-      results.successful().stream()
-          .map(wip::remove).filter(Objects::nonNull)
-          .forEach(TaskHandle::setSuccess);
-
-      // Report errors for failed tasks. Do NOT remove them from the WIP list.
-      results.errors().keySet().stream()
-          .map(wip::get).filter(Objects::nonNull)
-          .forEach(taskHandle -> taskHandle.setError(
-              new ServerException(results.errors().get(taskHandle.id()))));
-
-      log.atDebug()
-          .addKeyValue("count_success", results.errors().size())
-          .addKeyValue("count_errors", results.successful().size())
-          .log("Received results");
+      // Scheduling items to be retried may block temporarily if the queue is not
+      // adequately sized. Offload the next step to a different thread to avoid
+      // blocking the internal gRPC thread on which this callback is running.
+      retryService.execute(() -> {
+        // Mark successfully completed tasks in the WIP list.
+        // The whenComplete hook added in .add() will remove them from WIP.
+        results.successful().stream()
+            .map(wip::get).filter(Objects::nonNull)
+            .forEach(TaskHandle::setSuccess);
+
+        // Report errors for failed tasks. Do NOT remove them from the WIP list.
+        results.errors().keySet().stream()
+            .map(wip::get).filter(Objects::nonNull)
+            .forEach(taskHandle -> taskHandle.setError(
+                new ServerException(results.errors().get(taskHandle.id()))));
 
-      awaitResults.release();
-      assert awaitResults.availablePermits() == 1;
+        try {
+          queue.put(TaskHandle.END_RESULTS);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+        }
+      });
     }
 
     private void onBackoff(Event.Backoff backoff) {
diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/batch/RetriableTask.java b/src/main/java/io/weaviate/client6/v1/api/collections/batch/RetriableTask.java
@@ -58,10 +58,10 @@ public final boolean setSuccess() {
    * Mark the task failed. This status cannot be changed, so calling
    * {@link #setSuccess} afterwards will have no effect.
    *
-   * @param error Error message. Null values are tolerated, but are only expected
-   *              to occur due to a server's mistake.
-   *              Do not use {@code setError(null)} if the server reports success
-   *              status for the task; prefer {@link #setSuccess} in that case.
+   * @param t Error message. Null values are tolerated, but are only expected
+   *          to occur due to a server's mistake.
+   *          Do not use {@code setError(null)} if the server reports success
+   *          status for the task; prefer {@link #setSuccess} in that case.
    */
   public final boolean setError(Throwable t) {
     return current.completeExceptionally(t);
diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/batch/TaskHandle.java b/src/main/java/io/weaviate/client6/v1/api/collections/batch/TaskHandle.java
@@ -16,7 +16,17 @@
 @ThreadSafe
 @SuppressWarnings("deprecation") // protoc uses GeneratedMessageV3
 public final class TaskHandle extends RetriableTask {
-  static final TaskHandle POISON = new TaskHandle();
+  /**
+   * Poison pill informs the "sender" that the batch is closed
+   * and no more user-supplied items will be incoming.
+   */
+  static final TaskHandle POISON = new TaskHandle("POISON");
+  /**
+   * Break pill informs the "sender" that the event handler has finished
+   * processing an incoming {@link Event.Results} message and no more retry items
+   * will be added to the queue until the next message arrives.
+   */
+  static final TaskHandle END_RESULTS = new TaskHandle("END_RESULTS");
 
   /**
    * Input value as passed by the user.
@@ -56,16 +66,16 @@ private TaskHandle(Data data, RetryPolicy retryPolicy, Consumer<String> onRetry)
   }
 
   /**
-   * Poison pill constructor.
+   * Poison / Break pill constructor.
    *
    * <p>
    * A handle created with this constructor should not be
    * used for anything other that direct comparison using {@code ==} operator;
    * calling any method on a poison pill is likely to result in a
    * {@link NullPointerException} being thrown.
    */
-  private TaskHandle() {
-    super("POISON", RetryPolicy.never(), __ -> {
+  private TaskHandle(String name) {
+    super(name, RetryPolicy.never(), __ -> {
     });
     this.data = null;
   }
@@ -76,8 +86,8 @@ Data data() {
 
   @Override
   public String toString() {
-    if (this == POISON) {
-      return "TaskHandle<POISON>";
+    if (this == POISON || this == END_RESULTS) {
+      return "TaskHandle<%s>".formatted(id());
     }
     return "TaskHandle<id=%s, retried=%d, created=%s>".formatted(id(), timesRetried(), createdAt);
   }
diff --git a/src/test/java/io/weaviate/client6/v1/api/collections/batch/BatchContextTest.java b/src/test/java/io/weaviate/client6/v1/api/collections/batch/BatchContextTest.java
