test(batch): test ClientError flow triggered by exhausted reconect retries

Author: bevzzz

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

@@ -195,7 +195,6 @@ public TaskHandle add(BatchReference reference) throws InterruptedException {
   }
 
   void start() {
-    System.out.println("RESET COUNTDOWN LATCH to [2]");
     workers = new CountDownLatch(2);
 
     messages = streamFactory.createStream(new Recv());
@@ -216,7 +215,6 @@ void start() {
    * is reached.
    */
   void reconnect() throws InterruptedException, ExecutionException {
-    System.out.println("RESET COUNTDOWN LATCH to [1]");
     workers = new CountDownLatch(2);
 
     messages = streamFactory.createStream(new Recv());
@@ -296,6 +294,10 @@ private void shutdown() {
   }
 
   private void shutdownNow(Exception ex) {
+    // Now report this error to the server and close the stream.
+    closing.completeExceptionally(ex);
+    messages.onError(Status.INTERNAL.withCause(ex).asRuntimeException());
+
     // Terminate the "send" routine and wait for it to exit.
     // Since we're already in the error state we do not care
     // much if it throws or not.
@@ -305,11 +307,7 @@ private void shutdownNow(Exception ex) {
     } catch (Exception e) {
     }
 
-    // Now report this error to the server and close the stream.
-    closing.completeExceptionally(ex);
-    messages.onError(Status.INTERNAL.withCause(ex).asRuntimeException());
-
-    // Since shutdownNow is never triggerred by the "main" thread,
+    // Since shutdownNow is never triggered by the "main" thread,
     // it may be blocked on trying to add to the queue. While batch
     // context is active, we own this thread and may interrupt it.
     parent.interrupt();
@@ -434,7 +432,10 @@ private void trySend() {
           if (task == TaskHandle.POISON) {
             System.out.println("took POISON");
             drain();
-            break;
+
+            messages.onNext(Message.stop());
+            messages.onCompleted();
+            return;
           }
 
           Data data = task.data();
@@ -447,11 +448,7 @@ private void trySend() {
         Thread.currentThread().interrupt();
       } catch (Exception e) {
         onEvent(new Event.ClientError(e));
-        return;
       }
-
-      messages.onNext(Message.stop());
-      messages.onCompleted();
     }
 
     /**
@@ -817,8 +814,6 @@ public void onEnter(State prev) {
         return;
       }
 
-      // send.cancel(true);
-
       if (!ServerShuttingDown.class.isAssignableFrom(prev.getClass())) {
         // This is NOT an orderly shutdown, we're reconnecting after a stream hangup.
         // Assume all WIP items have been lost and re-submit everything.

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

@@ -2,6 +2,7 @@
 
 import static java.util.Objects.requireNonNull;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -57,7 +58,7 @@ public class BatchContextTest {
   private static final int MAX_SIZE_BYTES = 2 * 1024;
   private static final int BATCH_SIZE = 10;
   private static final int QUEUE_SIZE = 1;
-  private static final int MAX_RECONNECT_RETRIES = 2;
+  private static final int MAX_RECONNECT_RETRIES = 1;
 
   private CompletableStreamFactory factory;
   private MockServer server;
@@ -390,6 +391,41 @@ public void test_closeAfterStreamHangup() throws Exception {
     server.hangup();
   }
 
+  @Test
+  public void test_maxReconnectRetries() throws Exception {
+    server.expectMessage(WeaviateProtoBatch.BatchStreamRequest.MessageCase.START);
+
+    // Drop the connection several times until the client exhausts its reconnect attempts.
+    int retries = 0;
+    while (retries < MAX_RECONNECT_RETRIES) {
+      server.hangup();
+      server.expectMessage(WeaviateProtoBatch.BatchStreamRequest.MessageCase.START);
+      retries++;
+    }
+
+    Future<?> fatalHangup = server.hangup();
+    try {
+      fatalHangup.get();
+    } catch (InterruptedException ignored) {
+      // BatchContext#shutdownNow might interrupt the parent thread.
+    }
+
+    Assertions.assertThatThrownBy(() -> context.close())
+        .isInstanceOf(IOException.class)
+        .hasMessageContaining("Server unavailable");
+
+    // Cleanup: unset the context and factory to prevent test teardown code
+    // from tripping on it while trying to close the context.
+    context = null;
+
+    try {
+      factory.close();
+    } catch (InterruptedException ignored) {
+    } finally {
+    factory = null;
+    }
+  }
+
   @Test(expected = IllegalStateException.class)
   public void test_add_closed() throws Exception {
     server.expectMessage(WeaviateProtoBatch.BatchStreamRequest.MessageCase.START);
@@ -465,12 +501,7 @@ public MockServer(
         BlockingQueue<WeaviateProtoBatch.BatchStreamRequest> requestQueue) {
       this.eventStream = requireNonNull(eventStream, "eventStream is null");
       this.requestQueue = requireNonNull(requestQueue, "requestQueue is null");
-      this.eventExecutor = eventExecutor;
-    }
-
-    /** Wait until the next request arrives. */
-    WeaviateProtoBatch.BatchStreamRequest recv() throws InterruptedException {
-      return requestQueue.take();
+      this.eventExecutor = requireNonNull(eventExecutor, "eventExecutor is null");
     }
 
     Future<?> emitEvent(Event event) {
@@ -492,8 +523,8 @@ Future<?> emitEvent(Event event) {
     }
 
     /** Terminate the server-side of the stream abruptly. */
-    void hangup() {
-      emitEvent(new Event.StreamHangup(new RuntimeException("whaam!")));
+    Future<?> hangup() {
+      return emitEvent(new Event.StreamHangup(new RuntimeException("whaam!")));
     }
 
     /**