Add on closed invocation stream hook to make sure we correctly shutdown pending coroutine context in case of a failure. (#466)

Author: slinkydeveloper

sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/HandlerRunner.kt

@@ -15,6 +15,7 @@ import dev.restate.serde.Serde
 import dev.restate.serde.SerdeFactory
 import io.opentelemetry.extension.kotlin.asContextElement
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.atomic.AtomicReference
 import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -98,6 +99,7 @@ internal constructor(
       handlerContext: HandlerContext,
       requestSerde: Serde<REQ>,
       responseSerde: Serde<RES>,
+      onClosedInvocationStreamHook: AtomicReference<Runnable>
   ): CompletableFuture<Slice> {
     val ctx: Context = ContextImpl(handlerContext, contextSerdeFactory)
 
@@ -109,42 +111,44 @@ internal constructor(
                 handlerContext.request().otelContext()!!.asContextElement())
 
     val completableFuture = CompletableFuture<Slice>()
-
-    scope.launch {
-      val serializedResult: Slice
-
-      try {
-        // Parse input
-        val req: REQ
-        try {
-          req = requestSerde.deserialize(handlerContext.request().body)
-        } catch (e: Throwable) {
-          LOG.warn("Error deserializing request", e)
-          completableFuture.completeExceptionally(
-              throw TerminalException(
-                  TerminalException.BAD_REQUEST_CODE, "Cannot deserialize request: " + e.message))
-          return@launch
-        }
-
-        // Execute user code
-        @Suppress("UNCHECKED_CAST") val res: RES = runner(ctx as CTX, req)
-
-        // Serialize output
-        try {
-          serializedResult = responseSerde.serialize(res)
-        } catch (e: Throwable) {
-          LOG.warn("Error when serializing response", e)
-          completableFuture.completeExceptionally(e)
-          return@launch
+    val job =
+        scope.launch {
+          val serializedResult: Slice
+
+          try {
+            // Parse input
+            val req: REQ
+            try {
+              req = requestSerde.deserialize(handlerContext.request().body)
+            } catch (e: Throwable) {
+              LOG.warn("Error deserializing request", e)
+              completableFuture.completeExceptionally(
+                  throw TerminalException(
+                      TerminalException.BAD_REQUEST_CODE,
+                      "Cannot deserialize request: " + e.message))
+              return@launch
+            }
+
+            // Execute user code
+            @Suppress("UNCHECKED_CAST") val res: RES = runner(ctx as CTX, req)
+
+            // Serialize output
+            try {
+              serializedResult = responseSerde.serialize(res)
+            } catch (e: Throwable) {
+              LOG.warn("Error when serializing response", e)
+              completableFuture.completeExceptionally(e)
+              return@launch
+            }
+          } catch (e: Throwable) {
+            completableFuture.completeExceptionally(e)
+            return@launch
+          }
+
+          // Complete callback
+          completableFuture.complete(serializedResult)
         }
-      } catch (e: Throwable) {
-        completableFuture.completeExceptionally(e)
-        return@launch
-      }
-
-      // Complete callback
-      completableFuture.complete(serializedResult)
-    }
+    onClosedInvocationStreamHook.set { job.cancel() }
 
     return completableFuture
   }

sdk-api/src/main/java/dev/restate/sdk/HandlerRunner.java

@@ -21,6 +21,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.jspecify.annotations.Nullable;
@@ -48,7 +49,10 @@ public class HandlerRunner<REQ, RES>
 
   @Override
   public CompletableFuture<Slice> run(
-      HandlerContext handlerContext, Serde<REQ> requestSerde, Serde<RES> responseSerde) {
+      HandlerContext handlerContext,
+      Serde<REQ> requestSerde,
+      Serde<RES> responseSerde,
+      AtomicReference<Runnable> onClosedInvocationStreamHook) {
     CompletableFuture<Slice> returnFuture = new CompletableFuture<>();
 
     // Wrap the executor for setting/unsetting the thread local

sdk-common/src/main/java/dev/restate/sdk/endpoint/definition/HandlerContext.java

@@ -48,8 +48,6 @@ public interface HandlerContext {
 
   CompletableFuture<Void> set(String name, Slice value);
 
-  // ----- Syscalls
-
   CompletableFuture<AsyncResult<Void>> timer(Duration duration, @Nullable String name);
 
   record CallResult(

sdk-common/src/main/java/dev/restate/sdk/endpoint/definition/HandlerRunner.java

@@ -11,6 +11,7 @@
 import dev.restate.common.Slice;
 import dev.restate.serde.Serde;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 
 public interface HandlerRunner<REQ, RES> {
   /**
@@ -27,5 +28,8 @@ public interface HandlerRunner<REQ, RES> {
   interface Options {}
 
   CompletableFuture<Slice> run(
-      HandlerContext handlerContext, Serde<REQ> requestSerde, Serde<RES> responseSerde);
+      HandlerContext handlerContext,
+      Serde<REQ> requestSerde,
+      Serde<RES> responseSerde,
+      AtomicReference<Runnable> onClosedInvocationStreamHook);
 }

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

@@ -16,6 +16,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Flow;
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.jspecify.annotations.Nullable;
@@ -30,6 +31,7 @@ final class RequestProcessorImpl implements RequestProcessor {
   private final Context otelContext;
   private final EndpointRequestHandler.LoggingContextSetter loggingContextSetter;
   private final Executor syscallsExecutor;
+  private final AtomicReference<Runnable> onHandlerTaskCancellation;
 
   @SuppressWarnings("unchecked")
   public RequestProcessorImpl(
@@ -45,14 +47,44 @@ public RequestProcessorImpl(
     this.loggingContextSetter = loggingContextSetter;
     this.handlerDefinition = (HandlerDefinition<Object, Object>) handlerDefinition;
     this.syscallsExecutor = syscallExecutor;
+    this.onHandlerTaskCancellation = new AtomicReference<>();
   }
 
   // Flow methods implementation
 
   @Override
   public void subscribe(Flow.Subscriber<? super Slice> subscriber) {
     LOG.trace("Start processing invocation");
-    this.stateMachine.subscribe(subscriber);
+    this.stateMachine.subscribe(
+        new Flow.Subscriber<>() {
+          @Override
+          public void onSubscribe(Flow.Subscription subscription) {
+            subscriber.onSubscribe(subscription);
+          }
+
+          @Override
+          public void onNext(Slice slice) {
+            subscriber.onNext(slice);
+          }
+
+          @Override
+          public void onError(Throwable throwable) {
+            Runnable cancelTask = onHandlerTaskCancellation.get();
+            if (cancelTask != null) {
+              cancelTask.run();
+            }
+            subscriber.onError(throwable);
+          }
+
+          @Override
+          public void onComplete() {
+            Runnable cancelTask = onHandlerTaskCancellation.get();
+            if (cancelTask != null) {
+              cancelTask.run();
+            }
+            subscriber.onComplete();
+          }
+        });
     stateMachine
         .waitForReady()
         .thenCompose(v -> this.onReady())
@@ -119,7 +151,8 @@ private CompletableFuture<Void> onReady() {
             .run(
                 contextInternal,
                 handlerDefinition.getRequestSerde(),
-                handlerDefinition.getResponseSerde());
+                handlerDefinition.getResponseSerde(),
+                onHandlerTaskCancellation);
 
     return userCodeFuture.handle(
         (slice, t) -> {