Introduce orTimeout(millis) API, to easily combine a CombineablePromise with a sleep to implement timeouts.

Author: slinkydeveloper

src/restate_context.ts

@@ -17,7 +17,18 @@ import { Client, SendClient } from "./types/router";
 /**
  * A promise that can be combined using Promise combinators in RestateContext.
  */
-export type CombineablePromise<T> = Promise<T> & { __combineable: void };
+export type CombineablePromise<T> = Promise<T> & {
+  __combineable: void;
+
+  /**
+   * Creates a promise that awaits for the current promise up to the specified timeout duration.
+   * If the timeout is fired, this Promise will be rejected with a {@link TimeoutError}.
+   *
+   * @param millis duration of the sleep in millis.
+   * This is a lower-bound.
+   */
+  orTimeout(millis: number): Promise<T>;
+};
 
 /**
  * Base Restate context, which contains all operations that are the same in the gRPC-based API

src/restate_context_impl.ts

@@ -48,6 +48,7 @@ import {
   TerminalError,
   ensureError,
   errorToFailureWithTerminal,
+  TimeoutError,
 } from "./types/errors";
 import { jsonSerialize, jsonDeserialize } from "./utils/utils";
 import { Empty } from "./generated/google/protobuf/empty";
@@ -66,6 +67,7 @@ import {
   newJournalEntryPromiseId,
   PromiseId,
 } from "./promise_combinator_tracker";
+import { WrappedPromise } from "./utils/promises";
 
 export enum CallContexType {
   None,
@@ -345,7 +347,7 @@ export class RestateGrpcContextImpl implements RestateGrpcContext {
     return this.markCombineablePromise(this.sleepInternal(millis));
   }
 
-  private sleepInternal(millis: number): Promise<void> {
+  private sleepInternal(millis: number): WrappedPromise<void> {
     return this.stateMachine.handleUserCodeMessage<void>(
       SLEEP_ENTRY_MESSAGE_TYPE,
       SleepEntryMessage.create({ wakeUpTime: Date.now() + millis })
@@ -519,12 +521,34 @@ export class RestateGrpcContextImpl implements RestateGrpcContext {
   private markCombineablePromise<T>(
     p: Promise<T>
   ): InternalCombineablePromise<T> {
+    const journalIndex = this.stateMachine.getUserCodeJournalIndex();
+    const orTimeout = (millis: number): Promise<T> => {
+      const sleepPromise = this.sleepInternal(millis).transform(() => {
+        throw new TimeoutError();
+      });
+      const sleepPromiseIndex = this.stateMachine.getUserCodeJournalIndex();
+
+      return this.stateMachine.createCombinator(Promise.race.bind(Promise), [
+        {
+          id: newJournalEntryPromiseId(journalIndex),
+          promise: p,
+        },
+        {
+          id: newJournalEntryPromiseId(sleepPromiseIndex),
+          promise: sleepPromise,
+        },
+      ]);
+    };
+
     return Object.defineProperties(p, {
       __combineable: {
         value: undefined,
       },
       journalIndex: {
-        value: this.stateMachine.getUserCodeJournalIndex(),
+        value: journalIndex,
+      },
+      orTimeout: {
+        value: orTimeout.bind(this),
       },
     }) as InternalCombineablePromise<T>;
   }

src/types/errors.ts

@@ -202,6 +202,12 @@ export class TerminalError extends RestateError {
   }
 }
 
+export class TimeoutError extends TerminalError {
+  constructor() {
+    super("Timeout occurred", { errorCode: ErrorCodes.DEADLINE_EXCEEDED });
+  }
+}
+
 // Leads to Restate retries
 export class RetryableError extends RestateError {
   constructor(message: string, options?: { errorCode?: number; cause?: any }) {

test/promise_combinators.test.ts

@@ -30,6 +30,7 @@ import {
 } from "./protoutils";
 import { TestGreeter, TestResponse } from "../src/generated/proto/test";
 import { SLEEP_ENTRY_MESSAGE_TYPE } from "../src/types/protocol";
+import { TimeoutError } from "../src/types/errors";
 
 class AwakeableSleepRaceGreeter implements TestGreeter {
   async greet(): Promise<TestResponse> {
@@ -236,3 +237,61 @@ describe("AwakeableSleepRaceInterleavedWithSideEffectGreeter", () => {
     ]);
   });
 });
+
+class AwakeableOrTimeoutGreeter implements TestGreeter {
+  async greet(): Promise<TestResponse> {
+    const ctx = restate.useContext(this);
+
+    const { promise } = ctx.awakeable<string>();
+    try {
+      const result = await promise.orTimeout(100);
+      return TestResponse.create({
+        greeting: `Hello ${result}`,
+      });
+    } catch (e) {
+      if (e instanceof TimeoutError) {
+        return TestResponse.create({
+          greeting: `Hello timed-out`,
+        });
+      }
+    }
+
+    throw new Error("Unexpected result");
+  }
+}
+
+describe("AwakeableOrTimeoutGreeter", () => {
+  it("handles completion of awakeable", async () => {
+    const result = await new TestDriver(new AwakeableOrTimeoutGreeter(), [
+      startMessage(),
+      inputMessage(greetRequest("Till")),
+      completionMessage(1, JSON.stringify("Francesco")),
+    ]).run();
+
+    expect(result.length).toStrictEqual(5);
+    expect(result[0]).toStrictEqual(awakeableMessage());
+    expect(result[1].messageType).toStrictEqual(SLEEP_ENTRY_MESSAGE_TYPE);
+    expect(result.slice(2)).toStrictEqual([
+      combinatorEntryMessage(0, [1]),
+      outputMessage(greetResponse(`Hello Francesco`)),
+      END_MESSAGE,
+    ]);
+  });
+
+  it("handles completion of sleep", async () => {
+    const result = await new TestDriver(new AwakeableOrTimeoutGreeter(), [
+      startMessage(),
+      inputMessage(greetRequest("Till")),
+      completionMessage(2, undefined, true),
+    ]).run();
+
+    expect(result.length).toStrictEqual(5);
+    expect(result[0]).toStrictEqual(awakeableMessage());
+    expect(result[1].messageType).toStrictEqual(SLEEP_ENTRY_MESSAGE_TYPE);
+    expect(result.slice(2)).toStrictEqual([
+      combinatorEntryMessage(0, [2]),
+      outputMessage(greetResponse(`Hello timed-out`)),
+      END_MESSAGE,
+    ]);
+  });
+});