feat(batch): make sure to read all Results before closing the context

Author: bevzzz

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

@@ -23,7 +23,6 @@
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
-import java.util.function.BiConsumer;
 
 import javax.annotation.concurrent.GuardedBy;
 
@@ -276,7 +275,7 @@ void start() {
       throw new IllegalStateException("context is closed");
     }
 
-    messages = streamFactory.createStream(recv = new Recv());
+    messages = streamFactory.createStream(recv = new Recv(this));
     messages.onNext(Message.start(collectionHandleDefaults.consistencyLevel()));
 
     // "send" routine must start after the nextState has been set.
@@ -301,7 +300,7 @@ void reconnect() throws InterruptedException, ExecutionException {
     // after the server half of the stream is closed (EOF or hangup).
     recv.get();
 
-    messages = streamFactory.createStream(recv = new Recv());
+    messages = streamFactory.createStream(recv = new Recv(this));
     messages.onNext(Message.start(collectionHandleDefaults.consistencyLevel()));
   }
 
@@ -345,7 +344,7 @@ public void close() throws IOException {
       }
       throw new IOException(e.getCause());
     } finally {
-      shutdownExecutionServices();
+      shutdownExecutors();
     }
   }
 
@@ -393,22 +392,10 @@ private void shutdownNow(Exception ex) {
     }
   }
 
-  private void shutdownExecutionServices() {
-    BiConsumer<String, List<Runnable>> assertEmpty = (name, pending) -> {
-      assert pending.isEmpty() : "'%s' service had %d tasks awaiting execution"
-          .formatted(pending.size(), name);
-    };
-
-    List<Runnable> pending;
-
-    pending = sendService.shutdownNow();
-    assertEmpty.accept("send", pending);
-
-    pending = scheduledService.shutdownNow();
-    assertEmpty.accept("oom", pending);
-
-    pending = shutdownService.shutdownNow();
-    assertEmpty.accept("shutdown", pending);
+  private void shutdownExecutors() {
+    sendService.shutdown();
+    scheduledService.shutdown();
+    shutdownService.shutdown();
   }
 
   /** Set the new state and notify awaiting threads. */
@@ -523,13 +510,30 @@ private void trySend() {
             System.out.println("took POISON");
             drain();
 
-            // FIXME(dyma): we should wait until the WIP is empty
-            // and only then exit, in case the server restarts
-            // after ack'ing the last batch.
-
+            // Close our end of the stream and exit.
             messages.onNext(Message.stop());
             messages.onCompleted();
-            return;
+
+            // The SSB protocol requires the client to continue reading the stream
+            // until EOF. In the happy case, the server will close its half having
+            // processed all previous requests; the WIP buffer is empty in that case.
+            //
+            // It is possible that the server will be restarted or the stream will be
+            // hungup before client receives all Results, in which case we might need
+            // to re-submit the items remaining in the WIP buffer.
+            recv.get();
+
+            assert closed : "queue poisoned when context not closed";
+            if (wip.isEmpty()) {
+              return;
+            }
+
+            // Server closed the stream before reporting all expected Results.
+            // Return the poison to the queue and go another round -- if the
+            // client tried to reconnect the batch may've been filled again.
+            assert queue.isEmpty() : "queue must be empty after poison";
+            queue.add(task);
+            continue;
           }
 
           Data data = task.data();
@@ -618,11 +622,20 @@ private void awaitCanPrepareNext() throws InterruptedException {
     }
   }
 
-  private final class Recv extends CompletableFuture<Void> implements StreamObserver<Event> {
+  private static final class Recv extends CompletableFuture<Void> implements StreamObserver<Event> {
+    private final BatchContext<?> context;
+
+    private Recv(BatchContext<?> context) {
+      this.context = context;
+    }
 
     @Override
     public void onNext(Event event) {
-      onEvent(event);
+      try {
+        context.onEvent(event);
+      } catch (Exception e) {
+        context.onEvent(Event.StreamHangup.fromThrowable(e));
+      }
     }
 
     /**
@@ -633,7 +646,7 @@ public void onNext(Event event) {
     @Override
     public void onCompleted() {
       try {
-        onEvent(Event.EOF);
+        context.onEvent(Event.EOF);
       } finally {
         complete(null);
       }
@@ -643,7 +656,7 @@ public void onCompleted() {
     @Override
     public void onError(Throwable t) {
       try {
-        onEvent(Event.StreamHangup.fromThrowable(t));
+        context.onEvent(Event.StreamHangup.fromThrowable(t));
       } finally {
         complete(null);
       }

src/test/java/io/weaviate/client6/v1/api/collections/batch/BatchContextTest.java

@@ -2,7 +2,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Collection;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -150,13 +150,12 @@ public void test_sendOneBatch() throws Exception {
     Assertions.assertThat(tasks)
         .extracting(TaskHandle::isAcked).allMatch(CompletableFuture::isDone);
 
-    Future<?> results = out.emitEvent(new Event.Results(received, Collections.emptyMap()));
+    out.beforeEof(new Event.Results(received, Collections.emptyMap()));
 
     // Since MockServer runs in the same thread as this test,
     // the context will be updated before the last emitEvent returns.
     context.close();
 
-    results.get(); // Wait until the Results event has been processed.
     Assertions.assertThat(tasks).extracting(TaskHandle::result)
         .allMatch(CompletableFuture::isDone)
         .extracting(CompletableFuture::get).extracting(TaskHandle.Result::error)
@@ -179,18 +178,19 @@ public void test_drainOnClose() throws Exception {
     BACKGROUND.submit(() -> {
       try {
         List<String> received = recvDataAndAck();
-        Assertions.assertThat(tasks)
-            .extracting(TaskHandle::id).containsExactlyInAnyOrderElementsOf(received);
-        Assertions.assertThat(tasks)
-            .extracting(TaskHandle::isAcked).allMatch(CompletableFuture::isDone);
-        out.emitEvent(new Event.Results(received, Collections.emptyMap()));
+        Assertions.assertThat(tasks).extracting(TaskHandle::id)
+            .containsExactlyInAnyOrderElementsOf(received);
+        Assertions.assertThat(tasks).extracting(TaskHandle::isAcked)
+            .allMatch(CompletableFuture::isDone);
       } catch (Exception e) {
         throw new RuntimeException(e);
       }
     });
 
+    List<String> submitted = tasks.stream().map(TaskHandle::id).toList();
+    out.beforeEof(new Event.Results(submitted, Collections.emptyMap()));
+
     context.close();
-    awaitResults(tasks);
 
     Assertions.assertThat(tasks).extracting(TaskHandle::result)
         .allMatch(CompletableFuture::isDone)
@@ -220,15 +220,15 @@ public void test_backoff() throws Exception {
     // set by the Backoff message, i.e. BATCH_SIZE / 2.
     List<String> received = recvDataAndAck();
     Assertions.assertThat(received).hasSize(BATCH_SIZE / 2);
-    out.emitEvent(new Event.Results(received, Collections.emptyMap()));
+    out.beforeEof(new Event.Results(received, Collections.emptyMap()));
 
     backgroundAdd.get(); // Finish populating batch context.
 
     // Since testUser will try and add BATCH_SIZE no. objects,
     // we should expect there to be exactly 2 batches.
     received = recvDataAndAck();
     Assertions.assertThat(received).hasSize(BATCH_SIZE / 2);
-    out.emitEvent(new Event.Results(received, Collections.emptyMap()));
+    out.beforeEof(new Event.Results(received, Collections.emptyMap()));
 
     context.close();
 
@@ -262,20 +262,13 @@ public void test_backoffBacklog() throws Exception {
     tasks.add(context.add(WeaviateObject.of()));
 
     Assertions.assertThat(received = recvDataAndAck()).hasSize(batchSizeNew);
-    out.emitEvent(new Event.Results(received, Collections.emptyMap()));
+    out.beforeEof(new Event.Results(received, Collections.emptyMap()));
 
     Assertions.assertThat(received = recvDataAndAck()).hasSize(batchSizeNew);
-    out.emitEvent(new Event.Results(received, Collections.emptyMap()));
+    out.beforeEof(new Event.Results(received, Collections.emptyMap()));
 
     context.close();
 
-    // Wait until the Results event's been processed to guarantee
-    // that the tasks' futures are completed before asserting.
-    CompletableFuture.allOf(
-        tasks.stream().map(TaskHandle::result)
-            .toArray(CompletableFuture[]::new))
-        .get(100, TimeUnit.MILLISECONDS);
-
     Assertions.assertThat(tasks).extracting(TaskHandle::result)
         .allMatch(CompletableFuture::isDone)
         .extracting(CompletableFuture::get).extracting(TaskHandle.Result::error)
@@ -301,9 +294,11 @@ public void test_reconnect_onOom() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
+    List<TaskHandle> tasks = new ArrayList<>();
+
     // OOM is the opposite of Ack -- trigger a flush first.
     for (int i = 0; i < BATCH_SIZE; i++) {
-      context.add(WeaviateObject.of());
+      tasks.add(context.add(WeaviateObject.of()));
     }
 
     // Respond with OOM and wait for the client to close its end of the stream.
@@ -317,16 +312,20 @@ public void test_reconnect_onOom() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
     recvDataAndAck();
+
+    List<String> submitted = tasks.stream().map(TaskHandle::id).toList();
+    out.beforeEof(new Event.Results(submitted, Collections.emptyMap()));
   }
 
   @Test
   public void test_reconnect_onStreamHangup() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
+    List<TaskHandle> tasks = new ArrayList<>();
     // Trigger a flush.
     for (int i = 0; i < BATCH_SIZE; i++) {
-      context.add(WeaviateObject.of());
+      tasks.add(context.add(WeaviateObject.of()));
     }
 
     // Expect a new batch to arrive. Hangup the stream before sending the Acks.
@@ -348,14 +347,17 @@ public void test_reconnect_onStreamHangup() throws Exception {
     out.hangup();
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
-    context.add(WeaviateObject.of());
+    tasks.add(context.add(WeaviateObject.of()));
     recvDataAndAck();
 
     // Now fill up the rest of the batch to trigger a flush. Ack the incoming batch.
     for (int i = 0; i < BATCH_SIZE - 1; i++) {
-      context.add(WeaviateObject.of());
+      tasks.add(context.add(WeaviateObject.of()));
     }
     recvDataAndAck();
+
+    List<String> submitted = tasks.stream().map(TaskHandle::id).toList();
+    out.beforeEof(new Event.Results(submitted, Collections.emptyMap()));
   }
 
   @Test
@@ -391,20 +393,17 @@ public void test_reconnect_DrainAfterStreamHangup() throws Exception {
     out.emitEvent(Event.STARTED);
     Future<?> backgroundAcks = BACKGROUND.submit(() -> {
       try {
-        List<String> ids = recvDataAndAck();
-        out.emitEvent(new Event.Results(ids, Collections.emptyMap()));
-
-        ids = recvDataAndAck();
-        out.emitEvent(new Event.Results(ids, Collections.emptyMap()));
-
-        ids = recvDataAndAck();
-        Future<?> lastEvent = out.emitEvent(new Event.Results(ids, Collections.emptyMap()));
-        lastEvent.get();
+        recvDataAndAck();
+        recvDataAndAck();
+        recvDataAndAck();
       } catch (Exception e) {
         throw new RuntimeException(e);
       }
     });
 
+    List<String> submitted = tasks.stream().map(TaskHandle::id).toList();
+    out.beforeEof(new Event.Results(submitted, Collections.emptyMap()));
+
     context.close();
     backgroundAcks.get();
 
@@ -480,22 +479,14 @@ private List<String> recvData() throws InterruptedException {
         .toList();
   }
 
-  private void awaitResults(Collection<TaskHandle> tasks) throws Exception {
-    // Wait until the Results event's been processed to guarantee
-    // that the tasks' futures are completed before asserting.
-    CompletableFuture.allOf(
-        tasks.stream().map(TaskHandle::result)
-            .toArray(CompletableFuture[]::new))
-        .get(100, TimeUnit.MILLISECONDS);
-  }
-
   static String getBeacon(WeaviateProtoBatch.BatchReference reference) {
     return "weaviate://localhost/" + reference.getToCollection() + "/" + reference.getToUuid();
   }
 
   /** OutboundStream is a mock which dispatches server-side events. */
   private static final class OutboundStream {
     private final StreamObserver<Event> stream;
+    private List<Event> beforeEof = new ArrayList<>();
 
     OutboundStream(StreamObserver<Event> stream) {
       this.stream = stream;
@@ -508,13 +499,22 @@ Future<?> emitEvent(Event event) {
       return EVENT_THREAD.submit(() -> stream.onNext(event));
     }
 
-    /** Terminate the server-side of the stream abruptly. */
+    /** Terminate the server half of the stream abruptly. */
+
     Future<?> hangup() {
       return EVENT_THREAD.submit(() -> stream.onError(new RuntimeException("whaam!")));
     }
 
-    void eof() {
+    /** Emit events before closing the server half of the stream. */
+    void beforeEof(Event... events) {
+      this.beforeEof.addAll(Arrays.asList(events));
+    }
+
+    void eof(boolean ok) {
       assert Thread.currentThread() != TEST_THREAD : "test MUST NOT close/terminate the the stream";
+      if (ok) {
+        beforeEof.forEach(this::emitEvent);
+      }
       EVENT_THREAD.submit(stream::onCompleted);
     }
   }
@@ -548,13 +548,13 @@ WeaviateProtoBatch.BatchStreamRequest expectMessage(
     @Override
     public void onCompleted() {
       done.complete(null);
-      outbound.eof();
+      outbound.eof(true);
     }
 
     @Override
     public void onError(Throwable t) {
       done.completeExceptionally(t);
-      outbound.eof();
+      outbound.eof(false);
     }
 
     @Override