refactor: handle all events in BaseState

Author: bevzzz

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

@@ -641,40 +641,11 @@ public void onError(Throwable t) {
     }
   }
 
-  private final State AWAIT_STARTED = new BaseState("AWAIT_STARTED", BaseState.Action.PREPARE_NEXT) {
-    @Override
-    public void onEvent(Event event) {
-      if (event == Event.STARTED) {
-        setState(ACTIVE);
-      } else {
-        super.onEvent(event);
-      }
-    }
-  };
+  private final State AWAIT_STARTED = new BaseState("AWAIT_STARTED", BaseState.Action.PREPARE_NEXT);
   private final State ACTIVE = new BaseState("ACTIVE", BaseState.Action.PREPARE_NEXT, BaseState.Action.SEND);
-  private final State IN_FLIGHT = new BaseState("IN_FLIGHT") {
-    @Override
-    public void onEvent(Event event) {
-      if (event instanceof Event.Acks acks) {
-        Collection<String> removed = batch.clear();
-        if (!acks.acked().containsAll(removed)) {
-          throw ProtocolViolationException.incompleteAcks(List.copyOf(removed));
-        }
-        acks.acked().forEach(id -> {
-          TaskHandle task = wip.get(id);
-          if (task != null) {
-            task.setAcked();
-          }
-        });
-        setState(ACTIVE);
-      } else if (event instanceof Event.Oom oom) {
-        setState(new Oom(oom.delaySeconds()));
-      } else {
-        super.onEvent(event);
-      }
-    }
-  };
+  private final State IN_FLIGHT = new BaseState("IN_FLIGHT");
 
+  /** BaseState implements default handlers for all {@link Event} subclasses. */
   private class BaseState implements State {
     /** State's display name for logging. */
     private final String name;
@@ -717,24 +688,15 @@ public boolean canPrepareNext() {
       return permitted.contains(Action.PREPARE_NEXT);
     }
 
-    /**
-     * Handle events which may arrive at any moment without violating the protocol.
-     *
-     * <ul>
-     * <li>{@link Event.Results} -- update tasks in {@link #wip} and remove them.
-     * <li>{@link Event.Backoff} -- adjust batch size.
-     * <li>{@link Event#SHUTTING_DOWN} -- transition into
-     * {@link ServerShuttingDown}.
-     * <li>{@link Event.StreamHangup -- transition into {@link Reconnecting} state.
-     * <li>{@link Event.ClientError -- shutdown the service immediately.
-     * </ul>
-     *
-     * @throws ProtocolViolationException If event cannot be handled in this state.
-     * @see BatchContext#shutdownNow
-     */
     @Override
     public void onEvent(Event event) {
-      if (event instanceof Event.Results results) {
+      if (event == Event.STARTED) {
+        onStarted();
+      } else if (event instanceof Event.Acks acks) {
+        onAcks(acks);
+      } else if (event instanceof Event.Oom oom) {
+        onOom(oom);
+      } else if (event instanceof Event.Results results) {
         onResults(results);
       } else if (event instanceof Event.Backoff backoff) {
         onBackoff(backoff);
@@ -745,10 +707,28 @@ public void onEvent(Event event) {
       } else if (event instanceof Event.ClientError error) {
         onClientError(error);
       } else {
-        throw ProtocolViolationException.illegalStateTransition(this, event);
+        throw new AssertionError("unreachable with event " + event);
       }
     }
 
+    private void onStarted() {
+      setState(ACTIVE);
+    }
+
+    private void onAcks(Event.Acks acks) {
+      Collection<String> removed = batch.clear();
+      if (!acks.acked().containsAll(removed)) {
+        throw ProtocolViolationException.incompleteAcks(List.copyOf(removed));
+      }
+      acks.acked().forEach(id -> {
+        TaskHandle task = wip.get(id);
+        if (task != null) {
+          task.setAcked();
+        }
+      });
+      setState(ACTIVE);
+    }
+
     private void onResults(Event.Results results) {
       results.successful().forEach(id -> wip.remove(id).setSuccess());
       results.errors().forEach((id, error) -> wip.remove(id).setError(error));
@@ -758,6 +738,10 @@ private void onBackoff(Event.Backoff backoff) {
       batch.setMaxSize(backoff.maxSize());
     }
 
+    private void onOom(Event.Oom oom) {
+      setState(new Oom(oom.delaySeconds()));
+    }
+
     private void onShuttingDown() {
       setState(new ServerShuttingDown(this));
     }

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

@@ -119,7 +119,7 @@ public synchronized void reset() throws Exception {
   }
 
   @AfterClass
-  public static void closeExecutor() throws Exception {
+  public static void shutdownExecutors() throws Exception {
     boolean terminated;
 
     BACKGROUND.shutdown();
@@ -136,7 +136,7 @@ public static void closeExecutor() throws Exception {
   private static final WeaviateProtoBatch.BatchStreamRequest.MessageCase DATA = WeaviateProtoBatch.BatchStreamRequest.MessageCase.DATA;
 
   @Test
-  public void test_sendOneBatch() throws Exception {
+  public synchronized void test_sendOneBatch() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
@@ -166,7 +166,7 @@ public void test_sendOneBatch() throws Exception {
   }
 
   @Test
-  public void test_drainOnClose() throws Exception {
+  public synchronized void test_drainOnClose() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
@@ -202,7 +202,7 @@ public void test_drainOnClose() throws Exception {
   }
 
   @Test
-  public void test_backoff() throws Exception {
+  public synchronized void test_backoff() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
@@ -242,7 +242,7 @@ public void test_backoff() throws Exception {
   }
 
   @Test
-  public void test_backoffBacklog() throws Exception {
+  public synchronized void test_backoffBacklog() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
@@ -279,7 +279,7 @@ public void test_backoffBacklog() throws Exception {
   }
 
   @Test
-  public void test_reconnect_onShutdown() throws Exception {
+  public synchronized void test_reconnect_onShutdown() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
@@ -293,7 +293,7 @@ public void test_reconnect_onShutdown() throws Exception {
   }
 
   @Test
-  public void test_reconnect_onOom() throws Exception {
+  public synchronized void test_reconnect_onOom() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
@@ -321,7 +321,7 @@ public void test_reconnect_onOom() throws Exception {
   }
 
   @Test
-  public void test_reconnect_onStreamHangup() throws Exception {
+  public synchronized void test_reconnect_onStreamHangup() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
@@ -364,7 +364,7 @@ public void test_reconnect_onStreamHangup() throws Exception {
   }
 
   @Test
-  public void test_reconnect_DrainAfterStreamHangup() throws Exception {
+  public synchronized void test_reconnect_DrainAfterStreamHangup() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
@@ -417,15 +417,15 @@ public void test_reconnect_DrainAfterStreamHangup() throws Exception {
   }
 
   @Test
-  public void test_closeAfterStreamHangup() throws Exception {
+  public synchronized void test_closeAfterStreamHangup() throws Exception {
     in.expectMessage(START);
     out.emitEvent(Event.STARTED);
 
     out.hangup();
   }
 
   @Test
-  public void test_maxReconnectRetries() throws Exception {
+  public synchronized void test_maxReconnectRetries() throws Exception {
     in.expectMessage(START);
 
     // Drop the connection several times until the client
@@ -450,14 +450,14 @@ public void test_maxReconnectRetries() throws Exception {
   }
 
   @Test(expected = IllegalStateException.class)
-  public void test_add_closed() throws Exception {
+  public synchronized void test_add_closed() throws Exception {
     in.expectMessage(START);
     context.close();
     context.add(WeaviateObject.of(o -> o.properties(Map.of())));
   }
 
   @Test(expected = IllegalStateException.class)
-  public void test_startAfterClose() throws Exception {
+  public synchronized void test_startAfterClose() throws Exception {
     in.expectMessage(START);
     context.close();
     context.start();