refactor(batch): make TaskHandle retriable

This commit introduces a RetriableTask, which, at its core,
manages a lifecycle of a CompletableFuture such that it is
only completed after retrying this task is not possible.

Author: bevzzz

src/main/java/io/weaviate/client6/v1/api/collections/batch/BatchContext.java

@@ -223,6 +223,12 @@ public final class BatchContext<PropertiesT> implements Closeable {
    */
   private volatile CompletableFuture<?> recv;
 
+  /**
+   * Retry policy controls if and how many times
+   * a {@link RetriableTask} can be retried.
+   */
+  private final RetryPolicy retryPolicy;
+
   /**
    * Maximum number of times the client will attempt to re-open the stream
    * before terminating the context.
@@ -243,12 +249,14 @@ public final class BatchContext<PropertiesT> implements Closeable {
       int maxSizeBytes,
       CollectionDescriptor<PropertiesT> collectionDescriptor,
       CollectionHandleDefaults collectionHandleDefaults,
+      RetryPolicy retryPolicy,
       int batchSize,
       int queueSize,
       int maxReconnectRetries) {
-    this.streamFactory = requireNonNull(streamFactory, "streamFactory is null");
     this.collectionDescriptor = requireNonNull(collectionDescriptor, "collectionDescriptor is null");
     this.collectionHandleDefaults = requireNonNull(collectionHandleDefaults, "collectionHandleDefaults is null");
+    this.retryPolicy = requireNonNull(retryPolicy, "retryPolicy is null");
+    this.streamFactory = requireNonNull(streamFactory, "streamFactory is null");
 
     this.queue = new ArrayBlockingQueue<>(queueSize);
     this.batch = new Batch(batchSize, maxSizeBytes);
@@ -263,6 +271,7 @@ private BatchContext(Builder<PropertiesT> builder) {
         builder.maxSizeBytes,
         builder.collectionDescriptor,
         builder.collectionHandleDefaults,
+        builder.retryPolicy,
         builder.batchSize,
         builder.queueSize,
         builder.maxReconnectRetries);
@@ -272,18 +281,35 @@ private BatchContext(Builder<PropertiesT> builder) {
   public TaskHandle add(WeaviateObject<PropertiesT> object) throws InterruptedException {
     TaskHandle handle = new TaskHandle(
         object,
-        InsertManyRequest.buildObject(object, collectionDescriptor, collectionHandleDefaults));
+        InsertManyRequest.buildObject(object, collectionDescriptor, collectionHandleDefaults),
+        retryPolicy, this::retry);
     return add(handle);
   }
 
   /** Add {@link BatchReference} to the batch. */
   public TaskHandle add(BatchReference reference) throws InterruptedException {
     TaskHandle handle = new TaskHandle(
         reference,
-        InsertManyRequest.buildReference(reference, collectionHandleDefaults.tenant()));
+        InsertManyRequest.buildReference(reference, collectionHandleDefaults.tenant()),
+        retryPolicy, this::retry);
     return add(handle);
   }
 
+  private TaskHandle add(final TaskHandle taskHandle) throws InterruptedException {
+    if (closed) {
+      throw new IllegalStateException("context is closed");
+    }
+    requireNonNull(taskHandle, "taskHandle is null");
+
+    TaskHandle existing = wip.get(taskHandle.id());
+    if (existing != null) {
+      throw new DuplicateTaskException(taskHandle, existing);
+    }
+
+    queue.put(taskHandle);
+    return taskHandle;
+  }
+
   void start() {
     if (closed) {
       throw new IllegalStateException("context is closed");
@@ -316,12 +342,28 @@ void reconnect() throws InterruptedException, ExecutionException {
    *
    * <p>
    * BatchContext does not impose any limit on the number of times a task can
-   * be retried -- it is up to the user to implement an appropriate retry policy.
+   * be retried -- it is up to the user to select an appropriate retry policy.
    *
    * @see TaskHandle#timesRetried
+   * @see RetryPolicy
    */
-  public TaskHandle retry(TaskHandle taskHandle) throws InterruptedException {
-    return add(taskHandle.retry());
+  private void retry(String id) {
+    try {
+      requireNonNull(id, "id is null");
+
+      TaskHandle taskHandle = wip.get(id);
+      assert taskHandle != null : taskHandle + " is not wip";
+
+      // Put the handle back on the queue directly, circumventing
+      // the checks closed- and duplicate items checks we do for
+      // public methods. The retried task is guaranteed to be present
+      // in the WIP list and may be retried well after the context
+      // is closed to the user.
+      queue.put(taskHandle);
+    } catch (InterruptedException e) {
+      // Preserve interrupted state without throwing the exception.
+      Thread.currentThread().interrupt();
+    }
   }
 
   /**
@@ -562,21 +604,6 @@ private void onEvent(Event event) {
     }
   }
 
-  private TaskHandle add(final TaskHandle taskHandle) throws InterruptedException {
-    if (closed) {
-      throw new IllegalStateException("context is closed");
-    }
-    requireNonNull(taskHandle, "taskHandle is null");
-
-    TaskHandle existing = wip.get(taskHandle.id());
-    if (existing != null) {
-      throw new DuplicateTaskException(taskHandle, existing);
-    }
-
-    queue.put(taskHandle);
-    return taskHandle;
-  }
-
   private final class Send implements Runnable {
 
     @Override
@@ -856,10 +883,6 @@ private void onAcks(Event.Acks acks) {
       if (!acks.acked().containsAll(removed)) {
         throwInternal(ProtocolViolationException.incompleteAcks(List.copyOf(removed)));
       }
-      acks.acked().stream()
-          .map(wip::get).filter(Objects::nonNull)
-          .forEach(TaskHandle::setAcked);
-
       setState(ACTIVE);
     }
 
@@ -870,13 +893,18 @@ 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::remove).filter(Objects::nonNull)
-          .forEach(taskHandle -> taskHandle.setError(results.errors().get(taskHandle.id())));
+          .map(wip::get).filter(Objects::nonNull)
+          .forEach(taskHandle -> taskHandle.setError(
+              new ServerException(results.errors().get(taskHandle.id()))));
+
+      // TODO(dyma): notify receivedResults
     }
 
     private void onBackoff(Event.Backoff backoff) {
@@ -1194,10 +1222,16 @@ public static class Builder<PropertiesT> implements ObjectBuilder<BatchContext<P
       this.collectionHandleDefaults = collectionHandleDefaults;
     }
 
+    private RetryPolicy retryPolicy = RetryPolicy.never();
     private int batchSize = 1_000;
     private int queueSize = 1_000;
     private int maxReconnectRetries = 5;
 
+    public Builder<PropertiesT> retryPolicy(RetryPolicy retryPolicy) {
+      this.retryPolicy = retryPolicy;
+      return this;
+    }
+
     public Builder<PropertiesT> batchSize(int batchSize) {
       this.batchSize = batchSize;
       return this;

src/main/java/io/weaviate/client6/v1/api/collections/batch/RetriableTask.java

@@ -0,0 +1,80 @@
+package io.weaviate.client6.v1.api.collections.batch;
+
+import static java.util.Objects.requireNonNull;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
+
+abstract class RetriableTask {
+  private final String id;
+  private final CompletableFuture<Void> root;
+
+  private volatile CompletableFuture<Void> current = new CompletableFuture<>();
+  private volatile int retries = 0;
+
+  protected RetriableTask(String id, RetryPolicy retryPolicy, Consumer<String> onRetry) {
+    this.id = requireNonNull(id, "id is null");
+    this.root = retry(current, retryPolicy, onRetry);
+  }
+
+  private final CompletableFuture<Void> retry(
+      CompletableFuture<Void> f,
+      RetryPolicy retryPolicy,
+      Consumer<String> onRetry) {
+    requireNonNull(f, "f is null");
+    requireNonNull(retryPolicy, "retryPolicy is null");
+    requireNonNull(onRetry, "onRetry is null");
+
+    return f.exceptionallyCompose(t -> {
+      if (!retryPolicy.canRetry(this, t)) {
+        return CompletableFuture.failedFuture(t);
+      }
+      retries++;
+      current = new CompletableFuture<>();
+      onRetry.accept(id);
+      return retry(current, retryPolicy, onRetry);
+    });
+  }
+
+  /** Number of times this task has been retried. */
+  public final int timesRetried() {
+    return retries;
+  }
+
+  /** Retrieve the ID of this task. */
+  public final String id() {
+    return id;
+  }
+
+  /**
+   * Mark the task successful. This status cannot be changed, so calling
+   * {@link #setError} afterwards will have no effect.
+   */
+  public final boolean setSuccess() {
+    return current.complete(null);
+  }
+
+  /**
+   * 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.
+   */
+  public final boolean setError(Throwable t) {
+    return current.completeExceptionally(t);
+  }
+
+  /**
+   * Track completion of this task.
+   *
+   * @return A future which completes when the server reports success
+   *         for this tasks or the applied {@link RetryPolicy}
+   *         no longer permits retrying the task.
+   */
+  public final CompletableFuture<Void> done() {
+    return root;
+  }
+}

src/main/java/io/weaviate/client6/v1/api/collections/batch/RetryPolicy.java

@@ -0,0 +1,43 @@
+package io.weaviate.client6.v1.api.collections.batch;
+
+import static java.util.Objects.requireNonNull;
+
+import java.util.function.Predicate;
+
+public class RetryPolicy {
+  /** Create a retry policy that never permits retrying a task. */
+  static RetryPolicy never() {
+    return new RetryPolicy(__ -> false);
+  }
+
+  private final Predicate<RetriableTask> retry;
+
+  /**
+   * Construct a simple RetryPolicy that retries up to a certain number of times.
+   *
+   * @param maxRetries Maximum number of retries.
+   */
+  public RetryPolicy(int maxRetries) {
+    this(task -> task.timesRetried() < maxRetries);
+  }
+
+  /**
+   * Construct a RetryPolicy with a custom predicate.
+   *
+   * @param retry Predicate that returns true if the task should be retried.
+   */
+  public RetryPolicy(Predicate<RetriableTask> retry) {
+    this.retry = requireNonNull(retry, "retry is null");
+  }
+
+  /**
+   * Override this method to control which exceptions are considered retriable.
+   */
+  protected boolean canRetryThrowable(Throwable t) {
+    return t instanceof ServerException;
+  }
+
+  boolean canRetry(RetriableTask task, Throwable t) {
+    return canRetryThrowable(t) && retry.test(task);
+  }
+}

src/main/java/io/weaviate/client6/v1/api/collections/batch/ServerException.java

@@ -0,0 +1,13 @@
+package io.weaviate.client6.v1.api.collections.batch;
+
+import io.weaviate.client6.v1.api.WeaviateException;
+
+/**
+ * ServerException carries an error message the server returns in
+ * {@link Event.Results}. This is a retriable exception.
+ */
+public class ServerException extends WeaviateException {
+  ServerException(String message) {
+    super(message);
+  }
+}