[service-protocol] Add support for suspension message V7

Summary:
This adds support to handle which suspension message to decode
based on the protocol version. It also introduce the Suspension message V7
based on the new protocol.proto schema

Update the invocation Suspended state to store and handle
the awaiting on field

Author: muhamadazmy

.github/workflows/ci.yml

@@ -228,6 +228,23 @@ jobs:
       restateCommit: ${{ github.event.pull_request.head.sha || github.sha }}
       serviceImage: "ghcr.io/restatedev/test-services-node:main"
 
+  sdk-typescript-protocol-v7:
+    name: Run SDK-Typescript integration tests with Protocol V7
+    permissions:
+      contents: read
+      issues: read
+      checks: write
+      pull-requests: write
+      actions: read
+    secrets: inherit
+    needs: docker
+    uses: restatedev/sdk-typescript/.github/workflows/integration.yaml@main
+    with:
+      restateCommit: ${{ github.event.pull_request.head.sha || github.sha }}
+      envVars: |
+        RESTATE_experimental_allow_protocol_v7=true
+      testArtifactOutput: sdk-typescript-protocol-v7
+  
   sdk-rust:
     name: Run SDK-Rust integration tests
     permissions:

crates/invoker-api/src/effects.rs

@@ -8,23 +8,24 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::collections::HashSet;
+
 use restate_memory::NonZeroByteCount;
 use restate_types::deployment::PinnedDeployment;
 use restate_types::errors::InvocationError;
 use restate_types::identifiers::InvocationId;
 use restate_types::journal::EntryIndex;
 use restate_types::journal::enriched::EnrichedRawEntry;
 use restate_types::journal_events::raw::RawEvent;
-use restate_types::journal_v2;
 use restate_types::journal_v2::CommandIndex;
 use restate_types::journal_v2::raw::RawEntry;
+use restate_types::journal_v2::{self, UnresolvedFuture};
 use restate_types::storage::{StoredRawEntry, StoredRawEntryHeader};
 use restate_types::time::MillisSinceEpoch;
-use std::collections::HashSet;
 
 use crate::EffectKind::JournalEntryV2;
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub struct Effect {
     pub invocation_id: InvocationId,
@@ -37,7 +38,7 @@ pub struct Effect {
     pub kind: EffectKind,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 // todo: fix this and box the large variant (EffectKind is 320 bytes)
 #[allow(clippy::large_enum_variant)]
@@ -62,8 +63,16 @@ pub enum EffectKind {
         event: RawEvent,
     },
     SuspendedV2 {
+        /// Flattened set of notification ids
+        /// can be thought of as `UnresolvedFuture::unknown(waiting_for_notifications)`
         waiting_for_notifications: HashSet<journal_v2::NotificationId>,
     },
+    // Introduced in Restate v1.7. With the new service-protocol v7
+    SuspendedV3 {
+        /// Future tree describing the notifications this invocation is waiting on.
+        /// Introduced in Restate v1.7 (protocol version V7). `None` for older invocations.
+        awaiting_on: UnresolvedFuture,
+    },
     Paused {
         paused_event: RawEvent,
     },

crates/invoker-impl/src/invocation_task/mod.rs

@@ -46,9 +46,8 @@ use restate_types::identifiers::{InvocationId, PartitionLeaderEpoch};
 use restate_types::invocation::InvocationTarget;
 use restate_types::journal::EntryIndex;
 use restate_types::journal::enriched::EnrichedRawEntry;
-use restate_types::journal_v2;
 use restate_types::journal_v2::raw::RawNotification;
-use restate_types::journal_v2::{CommandIndex, NotificationId};
+use restate_types::journal_v2::{self, CommandIndex, NotificationId, UnresolvedFuture};
 use restate_types::live::Live;
 use restate_types::schema::deployment::DeploymentResolver;
 use restate_types::schema::invocation_target::InvocationTargetResolver;
@@ -85,6 +84,10 @@ const SERVICE_PROTOCOL_VERSION_V5: HeaderValue =
 const SERVICE_PROTOCOL_VERSION_V6: HeaderValue =
     HeaderValue::from_static("application/vnd.restate.invocation.v6");
 
+#[allow(clippy::declare_interior_mutable_const)]
+const SERVICE_PROTOCOL_VERSION_V7: HeaderValue =
+    HeaderValue::from_static("application/vnd.restate.invocation.v7");
+
 #[allow(clippy::declare_interior_mutable_const)]
 const X_RESTATE_SERVER: HeaderName = HeaderName::from_static("x-restate-server");
 
@@ -180,6 +183,7 @@ pub(super) enum InvocationTaskOutputInner {
     Closed,
     Suspended(HashSet<EntryIndex>),
     SuspendedV2(HashSet<NotificationId>),
+    SuspendedV3(UnresolvedFuture),
     Failed(InvokerError, LocalMemoryPool),
     /// The invocation task yielded due to memory pressure.
     /// The budget was dropped, returning memory to the global pool.
@@ -258,6 +262,8 @@ pub(super) struct InvocationTask<EE, DMR> {
 
     // throttling
     action_token_bucket: Option<TokenBucket>,
+
+    allow_protocol_v7: bool,
 }
 
 /// This is needed to split the run_internal in multiple loop functions and have shortcircuiting.
@@ -266,6 +272,7 @@ enum TerminalLoopState<T> {
     Closed,
     Suspended(HashSet<EntryIndex>),
     SuspendedV2(HashSet<NotificationId>),
+    SuspendedV3(UnresolvedFuture),
     Failed(InvokerError),
     /// Memory budget exhausted — the invocation should yield.
     ShouldYield(InvocationMemoryExhausted),
@@ -305,6 +312,7 @@ macro_rules! shortcircuit {
             TerminalLoopState::Closed => return TerminalLoopState::Closed,
             TerminalLoopState::Suspended(v) => return TerminalLoopState::Suspended(v),
             TerminalLoopState::SuspendedV2(v) => return TerminalLoopState::SuspendedV2(v),
+            TerminalLoopState::SuspendedV3(v) => return TerminalLoopState::SuspendedV3(v),
             TerminalLoopState::ShouldYield(oom) => return TerminalLoopState::ShouldYield(oom),
             TerminalLoopState::Failed(e) => return TerminalLoopState::Failed(e),
         }
@@ -333,6 +341,7 @@ where
         invoker_tx: mpsc::UnboundedSender<InvocationTaskOutput>,
         invoker_rx: mpsc::UnboundedReceiver<Notification>,
         action_token_bucket: Option<TokenBucket>,
+        allow_protocol_v7: bool,
     ) -> Self {
         Self {
             client,
@@ -350,6 +359,7 @@ where
             message_size_warning,
             retry_count_since_last_stored_entry,
             action_token_bucket,
+            allow_protocol_v7,
         }
     }
 
@@ -388,6 +398,7 @@ where
             TerminalLoopState::Closed => InvocationTaskOutputInner::Closed,
             TerminalLoopState::Suspended(v) => InvocationTaskOutputInner::Suspended(v),
             TerminalLoopState::SuspendedV2(v) => InvocationTaskOutputInner::SuspendedV2(v),
+            TerminalLoopState::SuspendedV3(v) => InvocationTaskOutputInner::SuspendedV3(v),
             TerminalLoopState::Failed(e) => {
                 // Best effort to release excessive memory. Note there can still be effects in flight
                 // that are being replicated and thereby occupy memory. Best if we periodically check
@@ -478,13 +489,16 @@ where
                 );
 
                 let chosen_service_protocol_version = shortcircuit!(
-                    ServiceProtocolVersion::pick(&deployment.supported_protocol_versions,)
-                        .ok_or_else(|| {
-                            InvokerError::IncompatibleServiceEndpoint(
-                                deployment.id,
-                                deployment.supported_protocol_versions.clone(),
-                            )
-                        })
+                    ServiceProtocolVersion::pick(
+                        &deployment.supported_protocol_versions,
+                        self.allow_protocol_v7
+                    )
+                    .ok_or_else(|| {
+                        InvokerError::IncompatibleServiceEndpoint(
+                            deployment.id,
+                            deployment.supported_protocol_versions.clone(),
+                        )
+                    })
                 );
 
                 (
@@ -600,6 +614,7 @@ fn service_protocol_version_to_header_value(
         ServiceProtocolVersion::V4 => SERVICE_PROTOCOL_VERSION_V4,
         ServiceProtocolVersion::V5 => SERVICE_PROTOCOL_VERSION_V5,
         ServiceProtocolVersion::V6 => SERVICE_PROTOCOL_VERSION_V6,
+        ServiceProtocolVersion::V7 => SERVICE_PROTOCOL_VERSION_V7,
     }
 }
 

crates/invoker-impl/src/invocation_task/service_protocol_runner_v4.rs

@@ -8,7 +8,6 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use std::collections::HashSet;
 use std::num::NonZeroUsize;
 use std::ops::Deref;
 use std::pin::Pin;
@@ -38,7 +37,7 @@ use restate_service_protocol_v4::entry_codec::ServiceProtocolV4Codec;
 use restate_service_protocol_v4::message_codec::{
     Decoder, Encoder, Message, MessageHeader, MessageType, StateEntry, proto,
 };
-use restate_types::errors::InvocationError;
+use restate_types::errors::{GenericError, InvocationError};
 use restate_types::identifiers::InvocationId;
 use restate_types::identifiers::ServiceId;
 use restate_types::invocation::{
@@ -51,7 +50,7 @@ use restate_types::journal_v2::command::{
 };
 use restate_types::journal_v2::raw::{RawCommand, RawEntry, RawNotification};
 use restate_types::journal_v2::{
-    CommandIndex, CommandType, Entry, EntryType, NotificationId, RunCompletion, RunResult, SignalId,
+    CommandIndex, CommandType, Entry, EntryType, RunCompletion, RunResult, UnresolvedFuture,
 };
 use restate_types::schema::deployment::{Deployment, DeploymentType, ProtocolType};
 use restate_types::schema::invocation_target::{DeploymentStatus, InvocationTargetResolver};
@@ -849,6 +848,7 @@ where
                 MessageType::CommandAck,
             )),
             Message::Suspension(suspension) => self.handle_suspension_message(suspension),
+            Message::AwaitingOn(awaiting_on) => self.handle_awaiting_on_message(awaiting_on),
             Message::Error(e) => self.handle_error_message(e),
             Message::End(_) => TerminalLoopState::Closed,
 
@@ -1171,34 +1171,45 @@ where
         }
     }
 
+    fn handle_awaiting_on_message(
+        &mut self,
+        _awaiting_on: proto::AwaitingOnMessage,
+    ) -> TerminalLoopState<()> {
+        // this message should mark this invocation as suspendable.
+        // if it's not running any side effects.
+
+        // todo(azmy): Handle awaiting on message"
+        TerminalLoopState::Continue(())
+    }
+
     fn handle_suspension_message(
         &mut self,
         suspension: proto::SuspensionMessage,
     ) -> TerminalLoopState<()> {
-        let suspension_indexes: HashSet<_> = suspension
-            .waiting_completions
-            .into_iter()
-            .map(NotificationId::for_completion)
-            .chain(
-                suspension
-                    .waiting_signals
-                    .into_iter()
-                    .map(SignalId::for_index)
-                    .map(NotificationId::for_signal),
-            )
-            .chain(
-                suspension
-                    .waiting_named_signals
-                    .into_iter()
-                    .map(|s| SignalId::for_name(s.into()))
-                    .map(NotificationId::for_signal),
-            )
-            .collect();
-        // We currently don't support empty suspension_indexes set
-        if suspension_indexes.is_empty() {
+        let Some(awaiting_on) = suspension.awaiting_on else {
             return TerminalLoopState::Failed(InvokerError::EmptySuspensionMessage);
+        };
+
+        let future: UnresolvedFuture = match awaiting_on.try_into().map_err(GenericError::from) {
+            Ok(future) => future,
+            Err(err) => return TerminalLoopState::Failed(InvokerError::EncodingV2(err.into())),
+        };
+
+        // We currently don't support empty future set
+        if future.is_empty() {
+            return TerminalLoopState::Failed(InvokerError::EmptySuspensionMessage);
+        }
+
+        match self.service_protocol_version {
+            ServiceProtocolVersion::Unspecified => unreachable!(),
+            ServiceProtocolVersion::V1
+            | ServiceProtocolVersion::V2
+            | ServiceProtocolVersion::V3
+            | ServiceProtocolVersion::V4
+            | ServiceProtocolVersion::V5
+            | ServiceProtocolVersion::V6 => TerminalLoopState::SuspendedV2(future.flatten()),
+            ServiceProtocolVersion::V7 => TerminalLoopState::SuspendedV3(future),
         }
-        TerminalLoopState::SuspendedV2(suspension_indexes)
     }
 
     fn handle_error_message(&mut self, error: proto::ErrorMessage) -> TerminalLoopState<()> {