Use the new REQUIRES_ACK flag and the EntryAckMessage (#195)

Author: slinkydeveloper

proto/protocol.proto

@@ -1,13 +1,11 @@
-/*
- * Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
- *
- * This file is part of the Restate SDK for Node.js/TypeScript,
- * which is released under the MIT license.
- *
- * You can find a copy of the license in file LICENSE in the root
- * directory of this repository or package, or at
- * https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
- */
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate service protocol, which is
+// released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/service-protocol/blob/main/LICENSE
 
 syntax = "proto3";
 
@@ -44,7 +42,6 @@ message StartMessage {
 }
 
 // Type: 0x0000 + 1
-// Note: Acks to custom messages will have the `empty` field filled.
 message CompletionMessage {
   uint32 entry_index = 1;
 
@@ -55,6 +52,11 @@ 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 {

src/journal.ts

@@ -18,6 +18,7 @@ import {
   CLEAR_STATE_ENTRY_MESSAGE_TYPE,
   COMPLETE_AWAKEABLE_ENTRY_MESSAGE_TYPE,
   CompletionMessage,
+  EntryAckMessage,
   GET_STATE_ENTRY_MESSAGE_TYPE,
   GetStateEntryMessage,
   INVOKE_ENTRY_MESSAGE_TYPE,
@@ -191,16 +192,23 @@ export class Journal<I, O> {
       journalEntry.resolve(m.empty);
       this.pendingJournalEntries.delete(m.entryIndex);
     } else {
-      if (journalEntry.messageType === p.SIDE_EFFECT_ENTRY_MESSAGE_TYPE) {
-        // Just needs and ack without completion
-        journalEntry.resolve(undefined);
-        this.pendingJournalEntries.delete(m.entryIndex);
-      } else {
-        //TODO completion message without a value/failure/empty and message is not a side effect
-      }
+      //TODO completion message without a value/failure/empty
     }
   }
 
+  public handleEntryAckMessage(m: EntryAckMessage) {
+    // Get message at that entryIndex in pendingJournalEntries
+    const journalEntry = this.pendingJournalEntries.get(m.entryIndex);
+
+    if (journalEntry === undefined) {
+      return;
+    }
+
+    // Just needs an ack
+    journalEntry.resolve(undefined);
+    this.pendingJournalEntries.delete(m.entryIndex);
+  }
+
   private handleReplay(
     journalIndex: number,
     replayMessage: Message,

src/state_machine.ts

@@ -19,6 +19,7 @@ import { rlog } from "./utils/logger";
 import { clearTimeout } from "timers";
 import {
   COMPLETION_MESSAGE_TYPE,
+  ENTRY_ACK_MESSAGE_TYPE,
   ERROR_MESSAGE_TYPE,
   OUTPUT_STREAM_ENTRY_MESSAGE_TYPE,
   OutputStreamEntryMessage,
@@ -80,21 +81,30 @@ export class StateMachine<I, O> implements RestateStreamConsumer {
       return false;
     }
 
-    if (m.messageType !== COMPLETION_MESSAGE_TYPE) {
+    if (m.messageType === COMPLETION_MESSAGE_TYPE) {
+      rlog.debugJournalMessage(
+        this.invocation.logPrefix,
+        "Received completion message from Restate, adding to journal.",
+        m.messageType,
+        m.message
+      );
+      this.journal.handleRuntimeCompletionMessage(
+        m.message as p.CompletionMessage
+      );
+    } else if (m.messageType === ENTRY_ACK_MESSAGE_TYPE) {
+      rlog.debugJournalMessage(
+        this.invocation.logPrefix,
+        "Received entry ack message from Restate, adding to journal.",
+        m.messageType,
+        m.message
+      );
+      this.journal.handleEntryAckMessage(m.message as p.EntryAckMessage);
+    } else {
       throw RetryableError.protocolViolation(
-        `Received message of type ${m.messageType}. Can only accept completion messages after replay has finished.`
+        `Received message of type ${m.messageType}. Can only accept completion or acks messages after replay has finished.`
       );
     }
 
-    rlog.debugJournalMessage(
-      this.invocation.logPrefix,
-      "Received completion message from Restate, adding to journal.",
-      m.messageType,
-      m.message
-    );
-    this.journal.handleRuntimeCompletionMessage(
-      m.message as p.CompletionMessage
-    );
     // Remove lingering suspension timeouts, if we are not waiting for completions anymore
     if (
       this.suspensionTimeout !== undefined &&

src/types/protocol.ts

@@ -16,6 +16,7 @@ import {
   ClearStateEntryMessage,
   CompleteAwakeableEntryMessage,
   CompletionMessage,
+  EntryAckMessage,
   ErrorMessage,
   GetStateEntryMessage,
   InvokeEntryMessage,
@@ -43,13 +44,15 @@ export {
   SleepEntryMessage,
   StartMessage,
   SuspensionMessage,
+  EntryAckMessage,
 } from "../generated/proto/protocol";
 
 // Export the protocol message types as defined by the restate protocol.
 export const START_MESSAGE_TYPE = 0x0000n;
 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 POLL_INPUT_STREAM_ENTRY_MESSAGE_TYPE = 0x0400n;
 export const OUTPUT_STREAM_ENTRY_MESSAGE_TYPE = 0x0401n;
 export const GET_STATE_ENTRY_MESSAGE_TYPE = 0x0800n;
@@ -74,6 +77,7 @@ export const KNOWN_MESSAGE_TYPES = new Set([
   COMPLETION_MESSAGE_TYPE,
   SUSPENSION_MESSAGE_TYPE,
   ERROR_MESSAGE_TYPE,
+  ENTRY_ACK_MESSAGE_TYPE,
   POLL_INPUT_STREAM_ENTRY_MESSAGE_TYPE,
   OUTPUT_STREAM_ENTRY_MESSAGE_TYPE,
   GET_STATE_ENTRY_MESSAGE_TYPE,
@@ -92,6 +96,7 @@ export const PROTOBUF_MESSAGE_NAME_BY_TYPE = new Map<bigint, string>([
   [COMPLETION_MESSAGE_TYPE, "CompletionMessage"],
   [SUSPENSION_MESSAGE_TYPE, "SuspensionMessage"],
   [ERROR_MESSAGE_TYPE, "ErrorMessage"],
+  [ENTRY_ACK_MESSAGE_TYPE, "EntryAckMessage"],
   [POLL_INPUT_STREAM_ENTRY_MESSAGE_TYPE, "PollInputStreamEntryMessage"],
   [OUTPUT_STREAM_ENTRY_MESSAGE_TYPE, "OutputStreamEntryMessage"],
   [GET_STATE_ENTRY_MESSAGE_TYPE, "GetStateEntryMessage"],
@@ -111,6 +116,7 @@ const PROTOBUF_MESSAGES: Array<[bigint, any]> = [
   [COMPLETION_MESSAGE_TYPE, CompletionMessage],
   [SUSPENSION_MESSAGE_TYPE, SuspensionMessage],
   [ERROR_MESSAGE_TYPE, ErrorMessage],
+  [ENTRY_ACK_MESSAGE_TYPE, EntryAckMessage],
   [POLL_INPUT_STREAM_ENTRY_MESSAGE_TYPE, PollInputStreamEntryMessage],
   [OUTPUT_STREAM_ENTRY_MESSAGE_TYPE, OutputStreamEntryMessage],
   [GET_STATE_ENTRY_MESSAGE_TYPE, GetStateEntryMessage],
@@ -131,6 +137,7 @@ export type ProtocolMessage =
   | CompletionMessage
   | SuspensionMessage
   | ErrorMessage
+  | EntryAckMessage
   | PollInputStreamEntryMessage
   | OutputStreamEntryMessage
   | GetStateEntryMessage

src/types/types.ts

@@ -11,12 +11,23 @@
 
 import {
   AWAKEABLE_ENTRY_MESSAGE_TYPE,
+  BACKGROUND_INVOKE_ENTRY_MESSAGE_TYPE,
+  CLEAR_STATE_ENTRY_MESSAGE_TYPE,
+  COMPLETE_AWAKEABLE_ENTRY_MESSAGE_TYPE,
+  COMPLETION_MESSAGE_TYPE,
+  ENTRY_ACK_MESSAGE_TYPE,
+  ERROR_MESSAGE_TYPE,
   GET_STATE_ENTRY_MESSAGE_TYPE,
+  INVOKE_ENTRY_MESSAGE_TYPE,
   KNOWN_MESSAGE_TYPES,
+  OUTPUT_STREAM_ENTRY_MESSAGE_TYPE,
   POLL_INPUT_STREAM_ENTRY_MESSAGE_TYPE,
   ProtocolMessage,
+  SET_STATE_ENTRY_MESSAGE_TYPE,
+  SIDE_EFFECT_ENTRY_MESSAGE_TYPE,
   SLEEP_ENTRY_MESSAGE_TYPE,
   START_MESSAGE_TYPE,
+  SUSPENSION_MESSAGE_TYPE,
 } from "./protocol";
 
 export class Message {
@@ -53,19 +64,21 @@ class MessageType {
     return messageType == START_MESSAGE_TYPE;
   }
 
-  static isCustom(messageTypeId: bigint): boolean {
-    return !KNOWN_MESSAGE_TYPES.has(messageTypeId);
-  }
-
-  static hasRequiresAckFlag(messageTypeId: bigint): boolean {
-    return this.isCustom(messageTypeId);
+  static hasRequiresAckFlag(messageType: bigint): boolean {
+    return (
+      messageType !== START_MESSAGE_TYPE &&
+      messageType !== ERROR_MESSAGE_TYPE &&
+      messageType !== SUSPENSION_MESSAGE_TYPE &&
+      messageType !== ENTRY_ACK_MESSAGE_TYPE &&
+      messageType !== COMPLETION_MESSAGE_TYPE
+    );
   }
 }
 
 const CUSTOM_MESSAGE_MASK = BigInt(0xfc00);
 const COMPLETED_MASK = BigInt(0x0001_0000_0000);
 const VERSION_MASK = BigInt(0x03ff_0000_0000);
-const REQUIRES_ACK_MASK = BigInt(0x0001_0000_0000);
+const REQUIRES_ACK_MASK = BigInt(0x8000_0000_0000);
 
 // The header is exported but only for tests.
 export class Header {

test/protoutils.ts

@@ -41,6 +41,8 @@ import {
   SuspensionMessage,
   ERROR_MESSAGE_TYPE,
   ErrorMessage,
+  ENTRY_ACK_MESSAGE_TYPE,
+  EntryAckMessage,
 } from "../src/types/protocol";
 import { Message } from "../src/types/types";
 import { TestRequest, TestResponse } from "../src/generated/proto/test";
@@ -235,6 +237,15 @@ export function completionMessage(
   }
 }
 
+export function ackMessage(index: number): Message {
+  return new Message(
+    ENTRY_ACK_MESSAGE_TYPE,
+    EntryAckMessage.create({
+      entryIndex: index,
+    })
+  );
+}
+
 export function invokeMessage(
   serviceName: string,
   methodName: string,