WIP trying to figure out where the RST stream issue even is...

Author: slinkydeveloper

gradle/libs.versions.toml

@@ -230,5 +230,5 @@
   restate = '2.5.0-SNAPSHOT'
   schema-kenerator = '2.1.2'
   spring-boot = '3.4.9'
-  vertx = '4.5.18'
+  vertx = '4.5.22'
   victools-json-schema = '4.38.0'

sdk-core/src/main/java/dev/restate/sdk/core/HandlerContextImpl.java

@@ -349,6 +349,10 @@ public void pollAsyncResult(AsyncResultInternal<?> asyncResult) {
 
   private void pollAsyncResultInner(AsyncResultInternal<?> asyncResult) {
     while (true) {
+      if (this.stateMachine.state() == InvocationState.CLOSED) {
+        asyncResult.publicFuture().completeExceptionally(AbortedExecutionException.INSTANCE);
+        return;
+      }
       if (asyncResult.isDone()) {
         return;
       }
@@ -412,9 +416,12 @@ private void pollAsyncResultInner(AsyncResultInternal<?> asyncResult) {
         // Let it loop now
       } else if (response instanceof StateMachine.DoProgressResponse.ReadFromInput
           || response instanceof StateMachine.DoProgressResponse.WaitingPendingRun) {
-        CompletableFuture.anyOf(
-                this.waitNextProcessedRun(), this.stateMachine.waitNextInputSignal())
-            .thenAccept(v -> this.pollAsyncResultInner(asyncResult));
+//        LOG.info("Gonna need to wait here {}", response);
+        this.stateMachine.onNextEvent(
+            () -> {
+              LOG.info("Triggered after wait {}", response);
+              this.pollAsyncResultInner(asyncResult);
+            });
         return;
       } else if (response instanceof StateMachine.DoProgressResponse.ExecuteRun) {
         triggerScheduledRun(((StateMachine.DoProgressResponse.ExecuteRun) response).handle());
@@ -430,7 +437,6 @@ public void proposeRunSuccess(int runHandle, Slice toWrite) {
     } catch (Exception e) {
       this.failWithoutContextSwitch(e);
     }
-    triggerNextProcessedRun();
   }
 
   @Override
@@ -444,15 +450,6 @@ public void proposeRunFailure(
     } catch (Exception e) {
       this.failWithoutContextSwitch(e);
     }
-    triggerNextProcessedRun();
-  }
-
-  private void triggerNextProcessedRun() {
-    if (this.nextProcessedRun != null) {
-      var fut = this.nextProcessedRun;
-      this.nextProcessedRun = null;
-      fut.complete(null);
-    }
   }
 
   private void triggerScheduledRun(int handle) {
@@ -475,13 +472,6 @@ public void proposeFailure(Throwable toWrite, @Nullable RetryPolicy retryPolicy)
         });
   }
 
-  private CompletableFuture<Void> waitNextProcessedRun() {
-    if (this.nextProcessedRun == null) {
-      this.nextProcessedRun = new CompletableFuture<>();
-    }
-    return this.nextProcessedRun;
-  }
-
   @Override
   public void close() {
     this.stateMachine.end();

sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StateMachine.java

@@ -42,9 +42,9 @@ static StateMachine init(
 
   CompletableFuture<Void> waitForReady();
 
-  // --- Await next input
+  // --- Await next event
 
-  CompletableFuture<Void> waitNextInputSignal();
+  void onNextEvent(Runnable runnable);
 
   // --- Async results
 

sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StateMachineImpl.java

@@ -28,6 +28,7 @@
 import java.util.function.Consumer;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.jspecify.annotations.NonNull;
 import org.jspecify.annotations.Nullable;
 
 class StateMachineImpl implements StateMachine {
@@ -38,7 +39,7 @@ class StateMachineImpl implements StateMachine {
 
   // Callbacks
   private final CompletableFuture<Void> waitForReadyFuture = new CompletableFuture<>();
-  private CompletableFuture<Void> waitNextProcessedInput;
+  private @NonNull Runnable nextEventListener = () -> {};
 
   // Java Flow and message handling
   private final MessageDecoder messageDecoder = new MessageDecoder();
@@ -72,22 +73,20 @@ public CompletableFuture<Void> waitForReady() {
   }
 
   @Override
-  public CompletableFuture<Void> waitNextInputSignal() {
-    if (this.stateContext.isInputClosed()) {
-      return CompletableFuture.completedFuture(null);
-    }
-    if (waitNextProcessedInput == null) {
-      this.waitNextProcessedInput = new CompletableFuture<>();
-    }
-    return this.waitNextProcessedInput;
+  public void onNextEvent(Runnable runnable) {
+//    LOG.info("Registering next event signal");
+    this.nextEventListener =
+        () -> {
+          this.nextEventListener.run();
+          runnable.run();
+        };
   }
 
-  private void triggerWaitNextInputSignal() {
-    if (this.waitNextProcessedInput != null) {
-      CompletableFuture<Void> fut = this.waitNextProcessedInput;
-      this.waitNextProcessedInput = null;
-      fut.complete(null);
-    }
+  private void triggerNextEventSignal() {
+//    LOG.info("Triggering next event signal");
+    Runnable listener = this.nextEventListener;
+    this.nextEventListener = () -> {};
+    listener.run();
   }
 
   // -- IO
@@ -142,7 +141,7 @@ public void onNext(Slice slice) {
       }
 
       if (shouldTriggerInputListener) {
-        this.triggerWaitNextInputSignal();
+        this.triggerNextEventSignal();
       }
 
     } catch (Throwable e) {
@@ -153,6 +152,7 @@ public void onNext(Slice slice) {
   @Override
   public void onError(Throwable throwable) {
     this.stateContext.getCurrentState().hitError(throwable, null, null, this.stateContext);
+    this.triggerNextEventSignal();
     cancelInputSubscription();
   }
 
@@ -164,7 +164,7 @@ public void onComplete() {
     } catch (Throwable e) {
       this.onError(e);
     }
-    this.triggerWaitNextInputSignal();
+    this.triggerNextEventSignal();
     this.cancelInputSubscription();
   }
 
@@ -548,7 +548,11 @@ public int run(String name) {
   @Override
   public void proposeRunCompletion(int handle, Slice value) {
     LOG.debug("Executing 'Run completed with success'");
-    this.stateContext.getCurrentState().proposeRunCompletion(handle, value, this.stateContext);
+    try {
+      this.stateContext.getCurrentState().proposeRunCompletion(handle, value, this.stateContext);
+    } finally {
+      this.triggerNextEventSignal();
+    }
   }
 
   @Override
@@ -558,9 +562,13 @@ public void proposeRunCompletion(
       Duration attemptDuration,
       @Nullable RetryPolicy retryPolicy) {
     LOG.debug("Executing 'Run completed with failure'");
-    this.stateContext
-        .getCurrentState()
-        .proposeRunCompletion(handle, exception, attemptDuration, retryPolicy, this.stateContext);
+    try {
+      this.stateContext
+          .getCurrentState()
+          .proposeRunCompletion(handle, exception, attemptDuration, retryPolicy, this.stateContext);
+    } finally {
+      this.triggerNextEventSignal();
+    }
   }
 
   @Override