Avoid transform after invoke, which appears to break CombineablePromise (#340)

jackkleeman · web-flow · commit 81888ce67568 · 2024-05-02T13:53:17.000+01:00
* Avoid transform after invoke, which appears to break CombineablePromise

This is a hotfix, a more meaningful fix would be to implement 'transform' for combineable promises so that they stay combineable

* Add tests for combining invokes

* InternalCombineablePromise should not be a wrapped promise
diff --git a/packages/restate-sdk/src/context_impl.ts b/packages/restate-sdk/src/context_impl.ts
@@ -74,10 +74,9 @@ export interface CallContext {
   delay?: number;
 }
 
-export type InternalCombineablePromise<T> = CombineablePromise<T> &
-  WrappedPromise<T> & {
-    journalIndex: number;
-  };
+export type InternalCombineablePromise<T> = CombineablePromise<T> & {
+  journalIndex: number;
+};
 
 export class ContextImpl implements ObjectContext {
   // here, we capture the context information for actions on the Restate context that
@@ -229,7 +228,8 @@ export class ContextImpl implements ObjectContext {
     method: string,
     data: Uint8Array,
     key?: string
-  ): InternalCombineablePromise<Uint8Array> {
+    /* eslint-disable-next-line @typescript-eslint/no-explicit-any */
+  ): InternalCombineablePromise<any> {
     this.checkState("invoke");
 
     const msg = new CallEntryMessage({
@@ -241,7 +241,7 @@ export class ContextImpl implements ObjectContext {
     return this.markCombineablePromise(
       this.stateMachine
         .handleUserCodeMessage(INVOKE_ENTRY_MESSAGE_TYPE, msg)
-        .transform((v) => v as Uint8Array)
+        .transform((v) => deserializeJson(v as Uint8Array))
     );
   }
 
@@ -280,9 +280,7 @@ export class ContextImpl implements ObjectContext {
           const route = prop as string;
           return (...args: unknown[]) => {
             const requestBytes = serializeJson(args.shift());
-            return this.invoke(name, route, requestBytes).transform(
-              (responseBytes) => deserializeJson(responseBytes)
-            );
+            return this.invoke(name, route, requestBytes);
           };
         },
       }
@@ -302,9 +300,7 @@ export class ContextImpl implements ObjectContext {
           const route = prop as string;
           return (...args: unknown[]) => {
             const requestBytes = serializeJson(args.shift());
-            return this.invoke(name, route, requestBytes, key).transform(
-              (responseBytes) => deserializeJson(responseBytes)
-            );
+            return this.invoke(name, route, requestBytes, key);
           };
         },
       }
@@ -670,20 +666,23 @@ export class ContextImpl implements ObjectContext {
       ]) as Promise<T>;
     };
 
-    return Object.defineProperties(p, {
-      [RESTATE_CTX_SYMBOL]: {
-        value: this,
-      },
-      journalIndex: {
-        value: journalIndex,
-      },
-      orTimeout: {
-        value: orTimeout.bind(this),
-      },
-    }) as InternalCombineablePromise<T>;
+    defineProperty(p, RESTATE_CTX_SYMBOL, this);
+    defineProperty(p, "journalIndex", journalIndex);
+    defineProperty(p, "orTimeout", orTimeout.bind(this));
+
+    return p;
   }
 }
 
+// wraps defineProperty such that it informs tsc of the correct type of its output
+function defineProperty<Obj extends object, Key extends PropertyKey, T>(
+  obj: Obj,
+  prop: Key,
+  value: T
+): asserts obj is Obj & Readonly<Record<Key, T>> {
+  Object.defineProperty(obj, prop, { value });
+}
+
 function unpack<T>(
   a: string | RunAction<T>,
   b?: RunAction<T>
diff --git a/packages/restate-sdk/test/promise_combinators.test.ts b/packages/restate-sdk/test/promise_combinators.test.ts
@@ -11,7 +11,12 @@
 
 import { describe, expect } from "@jest/globals";
 import * as restate from "../src/public_api";
-import { TestDriver, TestGreeter, TestResponse } from "./testdriver";
+import {
+  GreeterApi,
+  TestDriver,
+  TestGreeter,
+  TestResponse,
+} from "./testdriver";
 import {
   awakeableMessage,
   completionMessage,
@@ -27,6 +32,7 @@ import {
   sleepMessage,
   sideEffectMessage,
   ackMessage,
+  invokeMessage,
 } from "./protoutils";
 import {
   COMBINATOR_ENTRY_MESSAGE,
@@ -35,6 +41,7 @@ import {
 import { TimeoutError } from "../src/types/errors";
 import { CombineablePromise } from "../src/context";
 import { Empty } from "../src/generated/proto/protocol_pb";
+import { Buffer } from "node:buffer";
 
 class AwakeableSleepRaceGreeter implements TestGreeter {
   async greet(ctx: restate.ObjectContext): Promise<TestResponse> {
@@ -398,3 +405,182 @@ describe("AwakeableOrTimeoutGreeter", () => {
     ]);
   });
 });
+
+class InvokeRaceGreeter implements TestGreeter {
+  async greet(ctx: restate.ObjectContext): Promise<TestResponse> {
+    const jack = ctx.objectClient(GreeterApi, "Jack").greet({ name: "Jack" });
+    const jill = ctx.objectClient(GreeterApi, "Jill").greet({ name: "Jill" });
+
+    const result = await CombineablePromise.race([jack, jill]);
+
+    return TestResponse.create({
+      greeting: result.greeting,
+    });
+  }
+}
+
+describe("InvokeRaceGreeter", () => {
+  const jackInvoke = invokeMessage(
+    "greeter",
+    "greet",
+    Buffer.from(JSON.stringify({ name: "Jack" })),
+    undefined,
+    undefined,
+    "Jack"
+  );
+  const jillInvoke = invokeMessage(
+    "greeter",
+    "greet",
+    Buffer.from(JSON.stringify({ name: "Jill" })),
+    undefined,
+    undefined,
+    "Jill"
+  );
+
+  it("should suspend without completions", async () => {
+    const result = await new TestDriver(new InvokeRaceGreeter(), [
+      startMessage(),
+      inputMessage(greetRequest("Till")),
+    ]).run();
+
+    expect(result.length).toStrictEqual(3);
+    expect(result[0]).toStrictEqual(jackInvoke);
+    expect(result[1]).toStrictEqual(jillInvoke);
+    expect(result[2]).toStrictEqual(suspensionMessage([1, 2]));
+  });
+
+  it("handles completion of first invoke", async () => {
+    const result = await new TestDriver(new InvokeRaceGreeter(), [
+      startMessage(),
+      inputMessage(greetRequest("Till")),
+      completionMessage(1, greetResponse("Hi Jack")),
+      ackMessage(3),
+    ]).run();
+
+    expect(result.length).toStrictEqual(5);
+    expect(result[0]).toStrictEqual(jackInvoke);
+    expect(result[1]).toStrictEqual(jillInvoke);
+    expect(result.slice(2)).toStrictEqual([
+      combinatorEntryMessage(0, [1]),
+      outputMessage(greetResponse("Hi Jack")),
+      END_MESSAGE,
+    ]);
+  });
+
+  it("handles completion of second invoke", async () => {
+    const result = await new TestDriver(new InvokeRaceGreeter(), [
+      startMessage(),
+      inputMessage(greetRequest("Till")),
+      completionMessage(2, greetResponse("Hi Jill")),
+      ackMessage(3),
+    ]).run();
+
+    expect(result.length).toStrictEqual(5);
+    expect(result[0]).toStrictEqual(jackInvoke);
+    expect(result[1]).toStrictEqual(jillInvoke);
+    expect(result.slice(2)).toStrictEqual([
+      combinatorEntryMessage(0, [2]),
+      outputMessage(greetResponse(`Hi Jill`)),
+      END_MESSAGE,
+    ]);
+  });
+
+  it("handles replay of the first invoke", async () => {
+    const result = await new TestDriver(new InvokeRaceGreeter(), [
+      startMessage(),
+      inputMessage(greetRequest("Till")),
+      invokeMessage(
+        "greeter",
+        "greet",
+        Buffer.from(JSON.stringify({ name: "Jack" })),
+        greetResponse("Hi Jack"),
+        undefined,
+        "Jack"
+      ),
+      ackMessage(3),
+    ]).run();
+
+    expect(result.length).toStrictEqual(4);
+    expect(result[0]).toStrictEqual(jillInvoke);
+    expect(result.slice(1)).toStrictEqual([
+      combinatorEntryMessage(0, [1]),
+      outputMessage(greetResponse("Hi Jack")),
+      END_MESSAGE,
+    ]);
+  });
+
+  it("handles replay of both invokes", async () => {
+    const result = await new TestDriver(new InvokeRaceGreeter(), [
+      startMessage(),
+      inputMessage(greetRequest("Till")),
+      invokeMessage(
+        "greeter",
+        "greet",
+        Buffer.from(JSON.stringify({ name: "Jack" })),
+        greetResponse("Hi Jack"),
+        undefined,
+        "Jack"
+      ),
+      invokeMessage(
+        "greeter",
+        "greet",
+        Buffer.from(JSON.stringify({ name: "Jill" })),
+        greetResponse("Hi Jill"),
+        undefined,
+        "Jill"
+      ),
+      ackMessage(3),
+    ]).run();
+
+    expect(result).toStrictEqual([
+      // The first invoke will be chosen because Promise.race will pick the first promise, in case both are resolved
+      combinatorEntryMessage(0, [1, 2]),
+      outputMessage(greetResponse("Hi Jack")),
+      END_MESSAGE,
+    ]);
+  });
+
+  it("handles replay of the combinator with first invoke completed", async () => {
+    const result = await new TestDriver(new InvokeRaceGreeter(), [
+      startMessage(),
+      inputMessage(greetRequest("Till")),
+      invokeMessage(
+        "greeter",
+        "greet",
+        Buffer.from(JSON.stringify({ name: "Jack" })),
+        greetResponse("Hi Jack"),
+        undefined,
+        "Jack"
+      ),
+      jillInvoke,
+      combinatorEntryMessage(0, [1]),
+    ]).run();
+
+    expect(result).toStrictEqual([
+      outputMessage(greetResponse("Hi Jack")),
+      END_MESSAGE,
+    ]);
+  });
+
+  it("handles replay of the combinator with second invoke completed", async () => {
+    const result = await new TestDriver(new InvokeRaceGreeter(), [
+      startMessage(),
+      inputMessage(greetRequest("Till")),
+      jackInvoke,
+      invokeMessage(
+        "greeter",
+        "greet",
+        Buffer.from(JSON.stringify({ name: "Jill" })),
+        greetResponse("Hi Jill"),
+        undefined,
+        "Jill"
+      ),
+      combinatorEntryMessage(0, [2]),
+    ]).run();
+
+    expect(result).toStrictEqual([
+      outputMessage(greetResponse("Hi Jill")),
+      END_MESSAGE,
+    ]);
+  });
+});
diff --git a/packages/restate-sdk/test/testdriver.ts b/packages/restate-sdk/test/testdriver.ts
@@ -23,7 +23,11 @@ import { StateMachine } from "../src/state_machine";
 import { InvocationBuilder } from "../src/invocation";
 import { EndpointImpl } from "../src/endpoint/endpoint_impl";
 import { ObjectContext } from "../src/context";
-import { ServiceDefinition, object } from "../src/public_api";
+import {
+  object,
+  VirtualObjectDefinition,
+  VirtualObject,
+} from "../src/public_api";
 import { ProtocolMode } from "../src/types/discovery";
 
 export type TestRequest = {
@@ -38,11 +42,10 @@ export const TestResponse = {
   create: (test: TestResponse): TestResponse => test,
 };
 
-export type GreetType = {
-  greet: (key: string, arg: TestRequest) => Promise<TestResponse>;
-};
-
-export const GreeterApi: ServiceDefinition<"greeter", GreetType> = {
+export const GreeterApi: VirtualObjectDefinition<
+  "greeter",
+  VirtualObject<TestGreeter>
+> = {
   name: "greeter",
 };
 
@@ -188,7 +191,7 @@ export class TestDriver implements Connection {
     rlog.debug(
       `Adding result to the result array. Message type: ${
         msg.messageType
-      }, message: 
+      }, message:
         ${
           msg.message instanceof Uint8Array
             ? (msg.message as Uint8Array).toString()
