Introduce EndMessage (#216)

* Update protocol.proto

* Introduce EndMessage

* Update tests

Author: slinkydeveloper

proto/protocol.proto

@@ -52,11 +52,6 @@ message CompletionMessage {
   };
 }
 
-// Type: 0x0000 + 4
-message EntryAckMessage {
-  uint32 entry_index = 1;
-}
-
 // Type: 0x0000 + 2
 // Implementations MUST send this message when suspending an invocation.
 message SuspensionMessage {
@@ -83,6 +78,16 @@ message ErrorMessage {
   string description = 3;
 }
 
+// Type: 0x0000 + 4
+message EntryAckMessage {
+  uint32 entry_index = 1;
+}
+
+// Type: 0x0000 + 5
+// Implementations MUST send this message when the invocation lifecycle ends.
+message EndMessage {
+}
+
 // --- Journal Entries ---
 
 // Every Completable JournalEntry has a result field, filled only and only if the entry is in DONE state.

src/state_machine.ts

@@ -19,6 +19,8 @@ import { rlog } from "./utils/logger";
 import { clearTimeout } from "timers";
 import {
   COMPLETION_MESSAGE_TYPE,
+  END_MESSAGE_TYPE,
+  EndMessage,
   ENTRY_ACK_MESSAGE_TYPE,
   ERROR_MESSAGE_TYPE,
   OUTPUT_STREAM_ENTRY_MESSAGE_TYPE,
@@ -282,6 +284,9 @@ export class StateMachine<I, O> implements RestateStreamConsumer {
             "Function completed successfully."
           );
 
+          // Mark the end of the invocation
+          this.send(new Message(END_MESSAGE_TYPE, EndMessage.create()));
+
           this.finish(value);
         } catch (e) {
           this.unhandledError(ensureError(e));
@@ -303,7 +308,7 @@ export class StateMachine<I, O> implements RestateStreamConsumer {
             "Function completed with an error: " + error.message
           );
 
-          this.finishWithError(error);
+          this.sendErrorAndFinish(error);
         } catch (ee) {
           this.unhandledError(ensureError(ee));
         }
@@ -314,7 +319,7 @@ export class StateMachine<I, O> implements RestateStreamConsumer {
     return this.invocationComplete.promise;
   }
 
-  private async finishWithError(e: Error) {
+  private async sendErrorAndFinish(e: Error) {
     if (e instanceof TerminalError) {
       this.sendTerminalError(e);
     } else {
@@ -354,6 +359,9 @@ export class StateMachine<I, O> implements RestateStreamConsumer {
     if (!this.journal.outputMsgWasReplayed()) {
       this.send(msg);
     }
+
+    // Mark the end of the invocation
+    this.send(new Message(END_MESSAGE_TYPE, EndMessage.create()));
   }
 
   private send(message: Message) {
@@ -464,7 +472,7 @@ export class StateMachine<I, O> implements RestateStreamConsumer {
   }
 
   public async notifyHandlerExecutionError(e: RetryableError | TerminalError) {
-    await this.finishWithError(e);
+    await this.sendErrorAndFinish(e);
   }
 
   /**

src/types/protocol.ts

@@ -18,6 +18,7 @@ import {
   CompletionMessage,
   EntryAckMessage,
   ErrorMessage,
+  EndMessage,
   GetStateEntryMessage,
   InvokeEntryMessage,
   OutputStreamEntryMessage,
@@ -36,6 +37,7 @@ export {
   CompleteAwakeableEntryMessage,
   CompletionMessage,
   ErrorMessage,
+  EndMessage,
   GetStateEntryMessage,
   InvokeEntryMessage,
   OutputStreamEntryMessage,
@@ -53,6 +55,7 @@ export const COMPLETION_MESSAGE_TYPE = 0x0001n;
 export const SUSPENSION_MESSAGE_TYPE = 0x0002n;
 export const ERROR_MESSAGE_TYPE = 0x0003n;
 export const ENTRY_ACK_MESSAGE_TYPE = 0x0004n;
+export const END_MESSAGE_TYPE = 0x0005n;
 export const POLL_INPUT_STREAM_ENTRY_MESSAGE_TYPE = 0x0400n;
 export const OUTPUT_STREAM_ENTRY_MESSAGE_TYPE = 0x0401n;
 export const GET_STATE_ENTRY_MESSAGE_TYPE = 0x0800n;
@@ -78,6 +81,7 @@ export const KNOWN_MESSAGE_TYPES = new Set([
   SUSPENSION_MESSAGE_TYPE,
   ERROR_MESSAGE_TYPE,
   ENTRY_ACK_MESSAGE_TYPE,
+  END_MESSAGE_TYPE,
   POLL_INPUT_STREAM_ENTRY_MESSAGE_TYPE,
   OUTPUT_STREAM_ENTRY_MESSAGE_TYPE,
   GET_STATE_ENTRY_MESSAGE_TYPE,
@@ -97,6 +101,7 @@ export const PROTOBUF_MESSAGE_NAME_BY_TYPE = new Map<bigint, string>([
   [SUSPENSION_MESSAGE_TYPE, "SuspensionMessage"],
   [ERROR_MESSAGE_TYPE, "ErrorMessage"],
   [ENTRY_ACK_MESSAGE_TYPE, "EntryAckMessage"],
+  [END_MESSAGE_TYPE, "EndMessage"],
   [POLL_INPUT_STREAM_ENTRY_MESSAGE_TYPE, "PollInputStreamEntryMessage"],
   [OUTPUT_STREAM_ENTRY_MESSAGE_TYPE, "OutputStreamEntryMessage"],
   [GET_STATE_ENTRY_MESSAGE_TYPE, "GetStateEntryMessage"],
@@ -117,6 +122,7 @@ const PROTOBUF_MESSAGES: Array<[bigint, any]> = [
   [SUSPENSION_MESSAGE_TYPE, SuspensionMessage],
   [ERROR_MESSAGE_TYPE, ErrorMessage],
   [ENTRY_ACK_MESSAGE_TYPE, EntryAckMessage],
+  [END_MESSAGE_TYPE, EndMessage],
   [POLL_INPUT_STREAM_ENTRY_MESSAGE_TYPE, PollInputStreamEntryMessage],
   [OUTPUT_STREAM_ENTRY_MESSAGE_TYPE, OutputStreamEntryMessage],
   [GET_STATE_ENTRY_MESSAGE_TYPE, GetStateEntryMessage],
@@ -138,6 +144,7 @@ export type ProtocolMessage =
   | SuspensionMessage
   | ErrorMessage
   | EntryAckMessage
+  | EndMessage
   | PollInputStreamEntryMessage
   | OutputStreamEntryMessage
   | GetStateEntryMessage

test/awakeable.test.ts

@@ -26,6 +26,7 @@ import {
   outputMessage,
   startMessage,
   suspensionMessage,
+  END_MESSAGE,
 } from "./protoutils";
 import { TestGreeter, TestResponse } from "../src/generated/proto/test";
 import { ProtocolMode } from "../src/generated/proto/discovery";
@@ -74,6 +75,7 @@ describe("AwakeableGreeter", () => {
     expect(result).toStrictEqual([
       awakeableMessage(),
       outputMessage(greetResponse(`Hello Francesco for ${getAwakeableId(1)}`)),
+      END_MESSAGE,
     ]);
   });
 
@@ -87,6 +89,7 @@ describe("AwakeableGreeter", () => {
     expect(result).toStrictEqual([
       awakeableMessage(),
       outputMessage(greetResponse(`Hello  for ${getAwakeableId(1)}`)),
+      END_MESSAGE,
     ]);
   });
 
@@ -102,6 +105,7 @@ describe("AwakeableGreeter", () => {
       outputMessage(
         greetResponse(`Hello [object Object] for ${getAwakeableId(1)}`)
       ),
+      END_MESSAGE,
     ]);
   });
 
@@ -117,9 +121,10 @@ describe("AwakeableGreeter", () => {
       ),
     ]).run();
 
-    expect(result.length).toStrictEqual(2);
+    expect(result.length).toStrictEqual(3);
     expect(result[0]).toStrictEqual(awakeableMessage());
     checkTerminalError(result[1], "Something went wrong");
+    expect(result[2]).toStrictEqual(END_MESSAGE);
   });
 
   it("handles replay with value", async () => {
@@ -131,6 +136,7 @@ describe("AwakeableGreeter", () => {
 
     expect(result).toStrictEqual([
       outputMessage(greetResponse(`Hello Francesco for ${getAwakeableId(1)}`)),
+      END_MESSAGE,
     ]);
   });
 
@@ -141,8 +147,9 @@ describe("AwakeableGreeter", () => {
       awakeableMessage(undefined, failure("Something went wrong")),
     ]).run();
 
-    expect(result.length).toStrictEqual(1);
+    expect(result.length).toStrictEqual(2);
     checkTerminalError(result[0], "Something went wrong");
+    expect(result[1]).toStrictEqual(END_MESSAGE);
   });
 
   it("fails on journal mismatch. Completed with CompleteAwakeable during replay.", async () => {

test/complete_awakeable.test.ts

@@ -22,6 +22,7 @@ import {
   outputMessage,
   startMessage,
   rejectAwakeableMessage,
+  END_MESSAGE,
 } from "./protoutils";
 import { describe, expect } from "@jest/globals";
 import { TestDriver } from "./testdriver";
@@ -49,6 +50,7 @@ describe("ResolveAwakeableGreeter", () => {
     expect(result).toStrictEqual([
       resolveAwakeableMessage(getAwakeableId(1), "hello"),
       outputMessage(greetResponse("Hello")),
+      END_MESSAGE,
     ]);
   });
 
@@ -61,6 +63,7 @@ describe("ResolveAwakeableGreeter", () => {
     expect(result).toStrictEqual([
       resolveAwakeableMessage(getAwakeableId(1), ""),
       outputMessage(greetResponse("Hello")),
+      END_MESSAGE,
     ]);
   });
 
@@ -71,7 +74,10 @@ describe("ResolveAwakeableGreeter", () => {
       resolveAwakeableMessage(getAwakeableId(1), "hello"),
     ]).run();
 
-    expect(result).toStrictEqual([outputMessage(greetResponse("Hello"))]);
+    expect(result).toStrictEqual([
+      outputMessage(greetResponse("Hello")),
+      END_MESSAGE,
+    ]);
   });
 
   it("handles replay with value empty string", async () => {
@@ -81,7 +87,10 @@ describe("ResolveAwakeableGreeter", () => {
       resolveAwakeableMessage(getAwakeableId(1), ""),
     ]).run();
 
-    expect(result).toStrictEqual([outputMessage(greetResponse("Hello"))]);
+    expect(result).toStrictEqual([
+      outputMessage(greetResponse("Hello")),
+      END_MESSAGE,
+    ]);
   });
 
   it("fails on journal mismatch. Completed with invoke during replay.", async () => {
@@ -138,6 +147,7 @@ describe("RejectAwakeableGreeter", () => {
     expect(result).toStrictEqual([
       rejectAwakeableMessage(getAwakeableId(1), "my bad error"),
       outputMessage(greetResponse("Hello")),
+      END_MESSAGE,
     ]);
   });
 });

test/eager_state.test.ts

@@ -19,6 +19,7 @@ import { TestDriver } from "./testdriver";
 import {
   clearStateMessage,
   completionMessage,
+  END_MESSAGE,
   getStateMessage,
   greetRequest,
   greetResponse,
@@ -55,6 +56,7 @@ describe("GetEmpty", () => {
     expect(result).toStrictEqual([
       getStateMessage("STATE", undefined, true),
       outputMessage(greetResponse("true")),
+      END_MESSAGE,
     ]);
   });
 
@@ -78,7 +80,10 @@ describe("GetEmpty", () => {
       ProtocolMode.BIDI_STREAM
     ).run();
 
-    expect(result).toStrictEqual([outputMessage(greetResponse("true"))]);
+    expect(result).toStrictEqual([
+      outputMessage(greetResponse("true")),
+      END_MESSAGE,
+    ]);
   });
 });
 
@@ -103,6 +108,7 @@ describe("Get", () => {
     expect(result).toStrictEqual([
       getStateMessage("STATE", "One"),
       outputMessage(greetResponse("One")),
+      END_MESSAGE,
     ]);
   });
 
@@ -116,6 +122,7 @@ describe("Get", () => {
     expect(result).toStrictEqual([
       getStateMessage("STATE", "One"),
       outputMessage(greetResponse("One")),
+      END_MESSAGE,
     ]);
   });
 
@@ -158,6 +165,7 @@ describe("GetAppendAndGet", () => {
       setStateMessage("STATE", "OneTwo"),
       getStateMessage("STATE", "OneTwo"),
       outputMessage(greetResponse("OneTwo")),
+      END_MESSAGE,
     ]);
   });
 
@@ -173,6 +181,7 @@ describe("GetAppendAndGet", () => {
       setStateMessage("STATE", "OneTwo"),
       getStateMessage("STATE", "OneTwo"),
       outputMessage(greetResponse("OneTwo")),
+      END_MESSAGE,
     ]);
   });
 });
@@ -202,6 +211,7 @@ describe("GetClearAndGet", () => {
       clearStateMessage("STATE"),
       getStateMessage("STATE", undefined, true),
       outputMessage(greetResponse("One-nothing")),
+      END_MESSAGE,
     ]);
   });
 
@@ -217,6 +227,7 @@ describe("GetClearAndGet", () => {
       clearStateMessage("STATE"),
       getStateMessage("STATE", undefined, true),
       outputMessage(greetResponse("One-nothing")),
+      END_MESSAGE,
     ]);
   });
 });
@@ -249,6 +260,7 @@ describe("MultipleGet", () => {
       getStateMessage("STATE", "One"),
       getStateMessage("STATE", "One"),
       outputMessage(greetResponse("One - One - One")),
+      END_MESSAGE,
     ]);
   });
 
@@ -264,6 +276,7 @@ describe("MultipleGet", () => {
       getStateMessage("STATE", "One"),
       getStateMessage("STATE", "One"),
       outputMessage(greetResponse("One - One - One")),
+      END_MESSAGE,
     ]);
   });
 });