Bunch of changes

Author: slinkydeveloper

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

@@ -362,8 +362,7 @@ public CompletableFuture<Void> writeOutput(TerminalException throwable) {
 
   @Override
   public void pollAsyncResult(AsyncResultInternal<?> asyncResult) {
-    // We use the separate function for the recursion,
-    // as there's no need to jump back and forth between threads again.
+    this.stateMachine.pumpOutput();
     this.pollAsyncResultInner(asyncResult);
   }
 
@@ -400,10 +399,12 @@ private void pollAsyncResultInner(AsyncResultInternal<?> asyncResult) {
       if (response instanceof StateMachine.DoProgressResponse.AnyCompleted) {
         // Let it loop now
       } else if (response instanceof StateMachine.DoProgressResponse.WaitExternalProgress) {
+        this.stateMachine.pumpOutput();
         this.stateMachine.onExternalProgress(() -> this.pollAsyncResultInner(asyncResult));
         return;
       } else if (response instanceof StateMachine.DoProgressResponse.CancelSignalReceived) {
         asyncResult.tryCancel();
+        return;
       } else if (response instanceof StateMachine.DoProgressResponse.ExecuteRun) {
         triggerScheduledRun(((StateMachine.DoProgressResponse.ExecuteRun) response).handle());
         // Let it loop now

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

@@ -45,6 +45,10 @@ static StateMachine init(
 
   void onExternalProgress(Runnable runnable);
 
+  // --- Output pump (mirrors TS OutputPump.awaitNextProgress — drains ONE chunk)
+
+  void pumpOutput();
+
   // --- Async results
 
   sealed interface DoProgressResponse {

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

@@ -108,17 +108,16 @@ public void cancel() {
         });
   }
 
-  private void pumpOutput() {
+  @Override
+  public void pumpOutput() {
     if (outputSubscriber == null) {
       return;
     }
-    while (true) {
-      Optional<byte[]> chunk = vm.takeOutput();
-      if (chunk.isEmpty()) {
-        return;
-      }
-      outputSubscriber.onNext(Slice.wrap(chunk.get()));
+    byte[] chunk = vm.takeOutput();
+    if (chunk.length == 0) {
+      return;
     }
+    outputSubscriber.onNext(Slice.wrap(chunk));
   }
 
   private void checkReadyToExecute() {
@@ -171,14 +170,12 @@ public DoProgressResponse doProgress(List<Integer> handles) {
     if (result instanceof SharedCoreVM.DoProgressResult.AnyCompleted) {
       return DoProgressResponse.AnyCompleted.INSTANCE;
     } else if (result instanceof SharedCoreVM.DoProgressResult.WaitExternalProgress) {
-      pumpOutput();
       return DoProgressResponse.WaitExternalProgress.INSTANCE;
     } else if (result instanceof SharedCoreVM.DoProgressResult.CancelSignalReceived) {
       return DoProgressResponse.CancelSignalReceived.INSTANCE;
     } else if (result instanceof SharedCoreVM.DoProgressResult.ExecuteRun r) {
       return new DoProgressResponse.ExecuteRun(r.handle());
     } else if (result instanceof SharedCoreVM.DoProgressResult.Suspended) {
-      pumpOutput();
       ExceptionUtils.sneakyThrow(AbortedExecutionException.INSTANCE);
     }
     throw new IllegalStateException("Unknown DoProgressResult: " + result);
@@ -419,7 +416,7 @@ public void writeOutput(TerminalException exception) {
   public void end() {
     try {
       vm.sysEnd();
-      pumpOutput();
+      this.pumpOutput();
     } finally {
       if (outputSubscriber != null) {
         outputSubscriber.onComplete();

sdk-core/src/main/java/dev/restate/sdk/core/sharedcore/SharedCoreVM.java

@@ -114,16 +114,14 @@ public void notifyError(
         new VmNotifyError(message, stacktrace, delayOverrideMillis));
   }
 
-  public Optional<byte[]> takeOutput() {
+  public byte[] takeOutput() {
     if (closed) {
-      return Optional.empty();
+      return new byte[0];
     }
     LOG.trace("[vm=0x{}] takeOutput()", Integer.toHexString(vmPtr));
 
-    var ret =
-        instance.callCborVmFunction(exports -> exports.vmTakeOutput(vmPtr), TakeOutputReturn.class);
-    if (ret instanceof TakeOutputReturn.Eof) return Optional.empty();
-    return Optional.of(((TakeOutputReturn.Buffer) ret).bytes());
+    var ptr = instance.getExports().vmTakeOutput(vmPtr);
+    return instance.readAndFree(ptr);
   }
 
   public String getResponseContentType() {
@@ -583,17 +581,6 @@ record Ok(int handle) implements HandleReturn {}
     record Failure(int code, String message) implements HandleReturn {}
   }
 
-  @JsonTypeInfo(use = Id.NAME, include = As.PROPERTY, property = "type")
-  @JsonSubTypes({
-    @JsonSubTypes.Type(value = TakeOutputReturn.Buffer.class, name = "buffer"),
-    @JsonSubTypes.Type(value = TakeOutputReturn.Eof.class, name = "eof"),
-  })
-  public sealed interface TakeOutputReturn {
-    record Buffer(byte[] bytes) implements TakeOutputReturn {}
-
-    record Eof() implements TakeOutputReturn {}
-  }
-
   @JsonTypeInfo(use = Id.NAME, include = As.PROPERTY, property = "type")
   @JsonSubTypes({
     @JsonSubTypes.Type(value = IsReadyReturn.Ok.class, name = "ok"),

sdk-core/src/main/rust/src/lib.rs

@@ -252,8 +252,11 @@ fn vm_notify_error(rc_vm: &Rc<RefCell<WasmVM>>, input: VmNotifyError) {
 #[export_name = "vm_take_output"]
 pub unsafe extern "C" fn _vm_take_output(vm_pointer: *const RefCell<WasmVM>) -> u64 {
     let rc_vm = vm_ptr_to_rc(vm_pointer);
-    let res: TakeOutputReturn = VM::take_output(&mut rc_vm.borrow_mut().vm).into();
-    output_to_ptr(res)
+    let res: Vec<u8> = match VM::take_output(&mut rc_vm.borrow_mut().vm) {
+        TakeOutputResult::Buffer(b) => b.to_vec(),
+        TakeOutputResult::EOF => Vec::default(),
+    };
+    vec_to_ptr(res)
 }
 
 #[export_name = "vm_is_ready_to_execute"]
@@ -1224,16 +1227,6 @@ enum HandleReturn {
     Failure { code: u32, message: String },
 }
 
-#[derive(Serialize)]
-#[serde(tag = "type", rename_all = "camelCase", rename_all_fields = "camelCase")]
-enum TakeOutputReturn {
-    Buffer {
-        #[serde(with = "serde_bytes")]
-        bytes: Vec<u8>,
-    },
-    Eof,
-}
-
 #[derive(Serialize)]
 #[serde(tag = "type", rename_all = "camelCase", rename_all_fields = "camelCase")]
 enum IsReadyReturn {
@@ -1418,15 +1411,6 @@ impl From<Result<NotificationHandle, Error>> for HandleReturn {
     }
 }
 
-impl From<TakeOutputResult> for TakeOutputReturn {
-    fn from(value: TakeOutputResult) -> Self {
-        match value {
-            TakeOutputResult::Buffer(b) => TakeOutputReturn::Buffer { bytes: b.to_vec() },
-            TakeOutputResult::EOF => TakeOutputReturn::Eof,
-        }
-    }
-}
-
 impl From<ResponseHead> for ResponseHeadReturn {
     fn from(head: ResponseHead) -> Self {
         ResponseHeadReturn {