Add CombineablePromise.map (#506)

Author: slinkydeveloper

packages/restate-e2e-services/exclusions.yaml

@@ -1,11 +1 @@
-exclusions:
-  "alwaysSuspending":
-    - "dev.restate.sdktesting.tests.Combinators.awakeableOrTimeoutUsingAwaitAny"
-  "default":
-    - "dev.restate.sdktesting.tests.Combinators.awakeableOrTimeoutUsingAwaitAny"
-  "singleThreadSinglePartition":
-    - "dev.restate.sdktesting.tests.Combinators.awakeableOrTimeoutUsingAwaitAny"
-  "threeNodes":
-    - "dev.restate.sdktesting.tests.Combinators.awakeableOrTimeoutUsingAwaitAny"
-  "threeNodesAlwaysSuspending":
-    - "dev.restate.sdktesting.tests.Combinators.awakeableOrTimeoutUsingAwaitAny"
+exclusions: {}

packages/restate-e2e-services/src/virtual_object_command_interpreter.ts

@@ -96,13 +96,7 @@ function parseAwaitableCommand(
     case "createAwakeable":
       return createAwakeable(ctx, command.awakeableKey);
     case "sleep":
-      // TODO yes this is an incorrect cast, but for now this is fine, type coercion in TS FTW.
-      //  We need a mapper function in our promise type to make this working.
-      //  The kotlin code looks like this:
-      //  ctx.timer(this.timeoutMillis.milliseconds).map { "sleep" }
-      return ctx.sleep(
-        command.timeoutMillis
-      ) as unknown as CombineablePromise<string>;
+      return ctx.sleep(command.timeoutMillis).map(() => "sleep");
     case "runThrowTerminalException":
       return ctx.run<string>(() => {
         throw new TerminalError(command.reason);

packages/restate-sdk/src/context.ts

@@ -25,6 +25,7 @@ import type {
   Serde,
 } from "@restatedev/restate-sdk-core";
 import { ContextImpl } from "./context_impl.js";
+import type { TerminalError } from "./types/errors.js";
 
 /**
  * Represents the original request as sent to this handler.
@@ -651,6 +652,24 @@ export type CombineablePromise<T> = Promise<T> & {
    * This is a lower-bound.
    */
   orTimeout(millis: number): CombineablePromise<T>;
+
+  /**
+   * Creates a new {@link CombineablePromise} that maps the result of this promise with
+   * the provided `mapper`, once this promise is fulfilled.
+   *
+   * **NOTE**: You **MUST** use this API when you need to map the result of a
+   * {@link CombineablePromise} without `await`ing it, rather than using {@link Promise.then}.
+   * {@link Promise.then} is used by Restate to distinguish when awaiting an asynchronous operation,
+   * thus calling `.then` on several Restate promises can lead to concurrency issues.
+   *
+   * @param mapper the function to execute when this promise is fulfilled.
+   * If the promise completed successfully, `value` is provided as input, otherwise `failure` is provided as input.
+   * If this mapper returns a value, this value will be used to resolve the returned {@link CombineablePromise}.
+   * If the mapper throws a {@link TerminalError}, this error will be used to reject the returned {@link CombineablePromise}.
+   */
+  map<U>(
+    mapper: (value?: T, failure?: TerminalError) => U
+  ): CombineablePromise<U>;
 };
 
 /**

packages/restate-sdk/src/promises.ts

@@ -17,7 +17,12 @@ import type {
   InvocationPromise,
 } from "./context.js";
 import type * as vm from "./endpoint/handlers/vm/sdk_shared_core_wasm_bindings.js";
-import { CancelledError, TimeoutError } from "./types/errors.js";
+import {
+  CancelledError,
+  RestateError,
+  TerminalError,
+  TimeoutError,
+} from "./types/errors.js";
 import { CompletablePromise } from "./utils/completable_promise.js";
 import type { ContextImpl, RunClosuresTracker } from "./context_impl.js";
 import { setTimeout } from "timers/promises";
@@ -136,6 +141,12 @@ abstract class AbstractRestatePromise<T> implements RestatePromise<T> {
     ) as CombineablePromise<T>;
   }
 
+  map<U>(
+    mapper: (value?: T, failure?: TerminalError) => U
+  ): CombineablePromise<U> {
+    return new RestateMappedPromise(this[RESTATE_CTX_SYMBOL], this, mapper);
+  }
+
   tryCancel() {
     this.cancelPromise.reject(new CancelledError());
   }
@@ -288,6 +299,10 @@ export class RestatePendingPromise<T> implements RestatePromise<T> {
     return this;
   }
 
+  map<U>(): CombineablePromise<U> {
+    return this as unknown as CombineablePromise<U>;
+  }
+
   tryCancel(): void {}
   tryComplete(): void {}
   uncompletedLeaves(): number[] {
@@ -313,6 +328,58 @@ export class InvocationPendingPromise<T>
   }
 }
 
+export class RestateMappedPromise<T, U> extends AbstractRestatePromise<U> {
+  private publicPromiseMapper: (
+    value?: T,
+    failure?: TerminalError
+  ) => Promise<U>;
+
+  constructor(
+    ctx: ContextImpl,
+    readonly inner: RestatePromise<T>,
+    mapper: (value?: T, failure?: TerminalError) => U
+  ) {
+    super(ctx);
+    this.publicPromiseMapper = (value?: T, failure?: TerminalError) => {
+      try {
+        return Promise.resolve(mapper(value, failure));
+      } catch (e) {
+        if (e instanceof TerminalError) {
+          return Promise.reject(e);
+        } else {
+          ctx.handleInvocationEndError(e);
+          return pendingPromise();
+        }
+      }
+    };
+  }
+
+  tryComplete(): void {
+    this.inner.tryComplete();
+  }
+
+  uncompletedLeaves(): number[] {
+    return this.inner.uncompletedLeaves();
+  }
+
+  publicPromise(): Promise<U> {
+    const promiseMapper = this.publicPromiseMapper;
+    return this.inner.publicPromise().then(
+      (t) => promiseMapper(t, undefined),
+      (error) => {
+        if (error instanceof RestateError) {
+          return promiseMapper(undefined, error);
+        } else {
+          // Something else, just re-throw it
+          throw error;
+        }
+      }
+    );
+  }
+
+  readonly [Symbol.toStringTag] = "RestateMappedPromise";
+}
+
 /**
  * Promises executor, gluing VM with I/O and Promises given to user space.
  */