Small things

Author: slinkydeveloper

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

@@ -50,7 +50,11 @@ non-sealed interface AsyncResultInternal<T> extends AsyncResult<T>, UnresolvedFu
 
     void tryCancel();
 
-    void tryComplete(NotificationReader reader);
+    /**
+     * Attempts to complete this result from available notifications; returns whether it is now
+     * done.
+     */
+    boolean tryComplete(NotificationReader reader);
 
     CompletableFuture<T> publicFuture();
 
@@ -121,19 +125,21 @@ public void tryCancel() {
     }
 
     @Override
-    public void tryComplete(NotificationReader reader) {
-      if (this.isDone()) return;
-      reader
-          .take(handle)
-          .ifPresent(
-              value -> {
-                try {
-                  completer.complete(value, publicFuture);
-                } catch (Throwable e) {
-                  contextInternal.fail(e);
-                  publicFuture.completeExceptionally(AbortedExecutionException.INSTANCE);
-                }
-              });
+    public boolean tryComplete(NotificationReader reader) {
+      if (this.isDone()) {
+        return true;
+      }
+      Optional<NotificationValue> value = reader.take(handle);
+      if (value.isEmpty()) {
+        return false;
+      }
+      try {
+        completer.complete(value.get(), publicFuture);
+      } catch (Throwable e) {
+        contextInternal.fail(e);
+        publicFuture.completeExceptionally(AbortedExecutionException.INSTANCE);
+      }
+      return true;
     }
 
     @Override
@@ -174,9 +180,12 @@ public void tryCancel() {
     }
 
     @Override
-    public void tryComplete(NotificationReader reader) {
-      if (this.isDone()) return;
+    public boolean tryComplete(NotificationReader reader) {
+      if (this.isDone()) return true;
       asyncResult.tryComplete(reader);
+      // This is defensive, because the async result might not compelte immediately if there's a
+      // mapper?!
+      return this.isDone();
     }
 
     @Override
@@ -293,15 +302,16 @@ public void tryCancel() {
     }
 
     @Override
-    public void tryComplete(NotificationReader reader) {
-      if (this.isDone()) return;
+    public boolean tryComplete(NotificationReader reader) {
+      if (this.isDone()) return true;
       asyncResults.forEach(ar -> ar.tryComplete(reader));
       for (int i = 0; i < asyncResults.size(); i++) {
         if (asyncResults.get(i).isDone()) {
           publicFuture.complete(i);
-          return;
+          return true;
         }
       }
+      return this.isDone();
     }
 
     @Override
@@ -343,8 +353,8 @@ public void tryCancel() {
     }
 
     @Override
-    public void tryComplete(NotificationReader reader) {
-      if (this.isDone()) return;
+    public boolean tryComplete(NotificationReader reader) {
+      if (this.isDone()) return true;
       asyncResults.forEach(ar -> ar.tryComplete(reader));
       asyncResults.stream()
           .filter(ar -> ar.publicFuture().isCompletedExceptionally())
@@ -357,6 +367,7 @@ public void tryComplete(NotificationReader reader) {
                   this.publicFuture.completeExceptionally(e.getCause());
                 }
               });
+      return this.isDone();
     }
 
     @Override

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

@@ -424,17 +424,19 @@ private void pollAsyncResultInner(AsyncResultInternal<?> asyncResult) {
         return;
       }
 
-      // Let's start by trying to complete it
+      // Try to complete it; tryComplete returns whether it is now done, fusing the completeness
+      // check so we don't walk the async-result tree twice.
+      boolean done;
       try {
-        asyncResult.tryComplete(this::takeNotification);
+        done = asyncResult.tryComplete(this::takeNotification);
       } catch (Throwable e) {
         // This can happen if the state machine was closed in the meantime.
         failWithoutContextSwitch(e);
         asyncResult.publicFuture().completeExceptionally(AbortedExecutionException.INSTANCE);
         return;
       }
       // If done, nothing more to do here.
-      if (asyncResult.isDone()) {
+      if (done) {
         return;
       }
 

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

@@ -87,7 +87,7 @@ record CancelSignalReceived() implements AwaitResult {}
    * AbortedExecutionException}, which should be treated as a clean abort of the user code, not a
    * failure to re-report.
    */
-  @Nullable AwaitResult doAwait(UnresolvedFuture future);
+  AwaitResult doAwait(UnresolvedFuture future);
 
   sealed interface NotificationValue {