Use ImplicitCancellation feature from the shared core

Author: slinkydeveloper

packages/restate-sdk-cloudflare-workers/patches/vm/sdk_shared_core_wasm_bindings.d.ts

@@ -36,7 +36,7 @@ export interface WasmAwakeable {
 
 export type WasmAsyncResultValue = "NotReady" | "Empty" | { Success: Uint8Array } | { Failure: WasmFailure } | { StateKeys: string[] } | { InvocationId: string };
 
-export type WasmDoProgressResult = "AnyCompleted" | "ReadFromInput" | "WaitingPendingRun" | { ExecuteRun: number };
+export type WasmDoProgressResult = "AnyCompleted" | "ReadFromInput" | "WaitingPendingRun" | { ExecuteRun: number } | "CancelSignalReceived";
 
 export interface WasmCallHandle {
     invocation_id_completion_id: number;

packages/restate-sdk-cloudflare-workers/patches/vm/sdk_shared_core_wasm_bindings_bg.wasm

@@ -9,21 +9,12 @@
  * https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
  */
 
-import {
-  service,
-  endpoint,
-  type Context,
-  CombineablePromise,
-} from "@restatedev/restate-sdk";
-import { setTimeout } from "timers/promises";
+import { service, endpoint, type Context } from "@restatedev/restate-sdk";
 
 const greeter = service({
   name: "greeter",
   handlers: {
     greet: async (ctx: Context, name: string) => {
-      const p1 = ctx.sleep(110000);
-      const p2 = ctx.run("stuff-2", async () => setTimeout(10000));
-      await CombineablePromise.allSettled([p1, p2]);
       return `Hello ${name}`;
     },
   },

packages/restate-sdk-examples/src/greeter.ts

@@ -189,10 +189,9 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
       );
 
       // TODO eventually we return this promise back to the user
-      const invocationIdPromise = this.createInvocationIdPromise(
-        call_handles.invocation_id_completion_id
-      );
-      this.promisesExecutor.registerInvocationIdToCancel(invocationIdPromise);
+      // const invocationIdPromise = this.createInvocationIdPromise(
+      //   call_handles.invocation_id_completion_id
+      // );
 
       return call_handles.call_completion_id;
     }, completeUsing(SuccessWithSerde(responseSerde), Failure));
@@ -208,7 +207,7 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
         delay = BigInt(send.delay);
       }
 
-      const invocation_id_handle = vm.sys_send(
+      vm.sys_send(
         send.service,
         send.method,
         parameter,
@@ -222,9 +221,8 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
       ).invocation_id_completion_id;
 
       // TODO eventually we return this promise back to the user
-      const invocationIdPromise =
-        this.createInvocationIdPromise(invocation_id_handle);
-      this.promisesExecutor.registerInvocationIdToCancel(invocationIdPromise);
+      // const invocationIdPromise =
+      //   this.createInvocationIdPromise(invocation_id_handle);
     });
   }
 

packages/restate-sdk/src/context_impl.ts

@@ -17,34 +17,45 @@ export enum LogLevel {
   ERROR = 4,
 }
 export interface WasmFailure {
-    code: number;
-    message: string;
+  code: number;
+  message: string;
 }
 
 export interface WasmExponentialRetryConfig {
-    initial_interval: number | undefined;
-    factor: number;
-    max_interval: number | undefined;
-    max_attempts: number | undefined;
-    max_duration: number | undefined;
+  initial_interval: number | undefined;
+  factor: number;
+  max_interval: number | undefined;
+  max_attempts: number | undefined;
+  max_duration: number | undefined;
 }
 
 export interface WasmAwakeable {
-    id: string;
-    handle: number;
+  id: string;
+  handle: number;
 }
 
-export type WasmAsyncResultValue = "NotReady" | "Empty" | { Success: Uint8Array } | { Failure: WasmFailure } | { StateKeys: string[] } | { InvocationId: string };
+export type WasmAsyncResultValue =
+  | "NotReady"
+  | "Empty"
+  | { Success: Uint8Array }
+  | { Failure: WasmFailure }
+  | { StateKeys: string[] }
+  | { InvocationId: string };
 
-export type WasmDoProgressResult = "AnyCompleted" | "ReadFromInput" | "WaitingPendingRun" | { ExecuteRun: number };
+export type WasmDoProgressResult =
+  | "AnyCompleted"
+  | "ReadFromInput"
+  | "WaitingPendingRun"
+  | { ExecuteRun: number }
+  | "CancelSignalReceived";
 
 export interface WasmCallHandle {
-    invocation_id_completion_id: number;
-    call_completion_id: number;
+  invocation_id_completion_id: number;
+  call_completion_id: number;
 }
 
 export interface WasmSendHandle {
-    invocation_id_completion_id: number;
+  invocation_id_completion_id: number;
 }
 
 export class WasmHeader {
@@ -91,8 +102,21 @@ export class WasmVM {
   sys_clear_state(key: string): void;
   sys_clear_all_state(): void;
   sys_sleep(millis: bigint): number;
-  sys_call(service: string, handler: string, buffer: Uint8Array, key: string | null | undefined, headers: WasmHeader[]): WasmCallHandle;
-  sys_send(service: string, handler: string, buffer: Uint8Array, key: string | null | undefined, headers: WasmHeader[], delay?: bigint | null): WasmSendHandle;
+  sys_call(
+    service: string,
+    handler: string,
+    buffer: Uint8Array,
+    key: string | null | undefined,
+    headers: WasmHeader[]
+  ): WasmCallHandle;
+  sys_send(
+    service: string,
+    handler: string,
+    buffer: Uint8Array,
+    key: string | null | undefined,
+    headers: WasmHeader[],
+    delay?: bigint | null
+  ): WasmSendHandle;
   sys_awakeable(): WasmAwakeable;
   sys_complete_awakeable_success(id: string, buffer: Uint8Array): void;
   sys_complete_awakeable_failure(id: string, value: WasmFailure): void;
@@ -103,11 +127,16 @@ export class WasmVM {
   sys_run(name: string): number;
   propose_run_completion_success(handle: number, buffer: Uint8Array): void;
   propose_run_completion_failure(handle: number, value: WasmFailure): void;
-  propose_run_completion_failure_transient(handle: number, error_message: string, error_stacktrace: string | null | undefined, attempt_duration: bigint, config: WasmExponentialRetryConfig): void;
+  propose_run_completion_failure_transient(
+    handle: number,
+    error_message: string,
+    error_stacktrace: string | null | undefined,
+    attempt_duration: bigint,
+    config: WasmExponentialRetryConfig
+  ): void;
   sys_cancel_invocation(target_invocation_id: string): void;
   sys_write_output_success(buffer: Uint8Array): void;
   sys_write_output_failure(value: WasmFailure): void;
   sys_end(): void;
   is_processing(): boolean;
 }
-

packages/restate-sdk/src/endpoint/handlers/vm/sdk_shared_core_wasm_bindings.d.ts

@@ -12,7 +12,6 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
 import type { CombineablePromise } from "./context.js";
-import * as vm from "./endpoint/handlers/vm/sdk_shared_core_wasm_bindings.js";
 import { CancelledError, TimeoutError } from "./types/errors.js";
 import { CompletablePromise } from "./utils/completable_promise.js";
 import type { ContextImpl, RunClosuresTracker } from "./context_impl.js";
@@ -80,7 +79,7 @@ abstract class AbstractRestatePromise<T> implements RestatePromise<T> {
     this.pollingPromise =
       this.pollingPromise ||
       this[RESTATE_CTX_SYMBOL].promisesExecutor
-        .doProgress(true, this)
+        .doProgress(this)
         .catch(() => {});
     return this.publicPromiseOrCancelPromise().then(onfulfilled, onrejected);
   }
@@ -94,7 +93,7 @@ abstract class AbstractRestatePromise<T> implements RestatePromise<T> {
     this.pollingPromise =
       this.pollingPromise ||
       this[RESTATE_CTX_SYMBOL].promisesExecutor
-        .doProgress(true, this)
+        .doProgress(this)
         .catch(() => {});
     return this.publicPromiseOrCancelPromise().catch(onrejected);
   }
@@ -103,7 +102,7 @@ abstract class AbstractRestatePromise<T> implements RestatePromise<T> {
     this.pollingPromise =
       this.pollingPromise ||
       this[RESTATE_CTX_SYMBOL].promisesExecutor
-        .doProgress(true, this)
+        .doProgress(this)
         .catch(() => {});
     return this.publicPromiseOrCancelPromise().finally(onfinally);
   }
@@ -282,9 +281,6 @@ export class RestatePendingPromise<T> implements RestatePromise<T> {
  * Promises executor, gluing VM with I/O and Promises given to user space.
  */
 export class PromisesExecutor {
-  private readonly invocationIdsToCancel: Array<RestateSinglePromise<string>> =
-    [];
-
   constructor(
     private readonly coreVm: vm.WasmVM,
     private readonly inputPump: InputPump,
@@ -293,52 +289,13 @@ export class PromisesExecutor {
     private readonly errorCallback: (e: any) => void
   ) {}
 
-  async doProgress(
-    enableImplicitCancellation: boolean,
-    restatePromise: RestatePromise<unknown>
-  ) {
+  async doProgress(restatePromise: RestatePromise<unknown>) {
     // Only the first time try process output
     await this.outputPump.awaitNextProgress();
-    await this.doProgressInner(enableImplicitCancellation, restatePromise);
-  }
-
-  registerInvocationIdToCancel(p: RestateSinglePromise<string>) {
-    this.invocationIdsToCancel.push(p);
+    await this.doProgressInner(restatePromise);
   }
 
-  private async doProgressInner(
-    enableImplicitCancellation: boolean,
-    restatePromise: RestatePromise<unknown>
-  ) {
-    if (enableImplicitCancellation) {
-      const cancellation_notification = this.coreVm.take_notification(
-        vm.cancel_handle()
-      );
-      if (cancellation_notification !== "NotReady") {
-        // Cancel child invocations!
-        try {
-          const invocationIds: string[] = [];
-          // Await all the invocation id promises
-          for (const invocationId of this.invocationIdsToCancel) {
-            // This will not cancel, because I flipped the cancellation by taking it with take_notification
-            invocationIds.push(await invocationId);
-          }
-          // Now cancel all of them boom!
-          for (const invocationId of invocationIds) {
-            this.coreVm.sys_cancel_invocation(invocationId);
-          }
-        } catch (e) {
-          // Not good
-          this.errorCallback(e);
-          return;
-        }
-
-        // Now flag this promise as cancelled and boom!
-        restatePromise.tryCancel();
-        return;
-      }
-    }
-
+  private async doProgressInner(restatePromise: RestatePromise<unknown>) {
     // Try complete the promise
     restatePromise.tryComplete();
 
@@ -356,10 +313,6 @@ export class PromisesExecutor {
           // Completed, we're good!
           return;
         }
-        if (enableImplicitCancellation) {
-          // We want to be woken up on the cancellation marker too!
-          handles.push(vm.cancel_handle());
-        }
         const doProgressResult = this.coreVm.do_progress(
           new Uint32Array(handles)
         );
@@ -372,6 +325,9 @@ export class PromisesExecutor {
         } else if (doProgressResult === "WaitingPendingRun") {
           // Wait for any of the pending run to complete
           await this.runClosuresTracker.awaitNextCompletedRun();
+        } else if (doProgressResult === "CancelSignalReceived") {
+          restatePromise.tryCancel();
+          return;
         } else {
           // We need to execute a run closure
           this.runClosuresTracker.executeRun(doProgressResult.ExecuteRun);
@@ -380,7 +336,7 @@ export class PromisesExecutor {
         }
 
         // Recursion
-        await this.doProgress(enableImplicitCancellation, restatePromise);
+        await this.doProgress(restatePromise);
       } catch (e) {
         // Not good, this is a retryable error.
         this.errorCallback(e);

packages/restate-sdk/src/endpoint/handlers/vm/sdk_shared_core_wasm_bindings.js

@@ -320,6 +320,8 @@ pub enum WasmDoProgressResult {
     WaitingPendingRun,
     /// The SDK should execute a pending run
     ExecuteRun(#[tsify(type = "number")] WasmNotificationHandle),
+    /// Got cancel signal
+    CancelSignalReceived,
 }
 
 impl From<DoProgressResponse> for WasmDoProgressResult {
@@ -329,6 +331,7 @@ impl From<DoProgressResponse> for WasmDoProgressResult {
             DoProgressResponse::ReadFromInput => WasmDoProgressResult::ReadFromInput,
             DoProgressResponse::WaitingPendingRun => WasmDoProgressResult::WaitingPendingRun,
             DoProgressResponse::ExecuteRun(n) => WasmDoProgressResult::ExecuteRun(n.into()),
+            DoProgressResponse::CancelSignalReceived => WasmDoProgressResult::CancelSignalReceived
         }
     }
 }
@@ -392,7 +395,7 @@ impl WasmVM {
         let log_dispatcher = Dispatch::new(log_subscriber(log_level, Some(logger_id)));
 
         let vm = tracing::dispatcher::with_default(&log_dispatcher, || {
-            CoreVM::new(WasmHeaderList::from(headers), VMOptions {})
+            CoreVM::new(WasmHeaderList::from(headers), VMOptions::default())
         })?;
 
         Ok(Self { vm, log_dispatcher })