[storage] Draft for Envelope v2

# Summary
- Separate header, keys, and other envelope related
information from the actual payload of the envelope
- Ability to decode envelope header without decoding the full envelope
- Support different encoding for the payload (records) (flexbuffers, or bilrost)
- Delay decoding of payload until needed

# TODO:
- [ ] Decode v1 envelope into v2
- [ ] Use Envelope v2 in code

Author: muhamadazmy

Cargo.lock

@@ -12,6 +12,7 @@
 
 mod arc_encodings;
 mod nonzero;
+mod phantom_data;
 mod range;
 
 pub mod display_from_str;

crates/encoding/src/bilrost_encodings/mod.rs

@@ -0,0 +1,56 @@
+// Copyright (c) 2023 - 2025 Restate Software, Inc., Restate GmbH.
+// All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+use std::marker::PhantomData;
+
+use bilrost::{
+    DecodeErrorKind,
+    encoding::{EmptyState, ForOverwrite, Proxiable},
+};
+
+use crate::bilrost_encodings::RestateEncoding;
+
+struct PhantomDataTag;
+
+impl<T> Proxiable<PhantomDataTag> for PhantomData<T> {
+    type Proxy = ();
+
+    fn encode_proxy(&self) -> Self::Proxy {}
+
+    fn decode_proxy(&mut self, _: Self::Proxy) -> Result<(), DecodeErrorKind> {
+        Ok(())
+    }
+}
+
+impl<T> ForOverwrite<RestateEncoding, PhantomData<T>> for () {
+    fn for_overwrite() -> PhantomData<T> {
+        PhantomData
+    }
+}
+
+impl<T> EmptyState<RestateEncoding, PhantomData<T>> for () {
+    fn empty() -> PhantomData<T> {
+        PhantomData
+    }
+
+    fn is_empty(_: &PhantomData<T>) -> bool {
+        true
+    }
+
+    fn clear(_: &mut PhantomData<T>) {}
+}
+
+bilrost::delegate_proxied_encoding!(
+    use encoding (::bilrost::encoding::General)
+    to encode proxied type (PhantomData<T>)
+    using proxy tag (PhantomDataTag)
+    with encoding (RestateEncoding)
+    with generics (T)
+);

crates/encoding/src/bilrost_encodings/phantom_data.rs

@@ -16,11 +16,11 @@ use restate_types::invocation::InvocationEpoch;
 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::storage::{StoredRawEntry, StoredRawEntryHeader};
 use restate_types::time::MillisSinceEpoch;
+use restate_types::{flexbuffers_storage_encode_decode, journal_v2};
 use std::collections::HashSet;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -35,6 +35,8 @@ pub struct Effect {
     pub kind: EffectKind,
 }
 
+flexbuffers_storage_encode_decode!(Effect);
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 // todo: fix this and box the large variant (EffectKind is 320 bytes)

crates/invoker-api/src/effects.rs

@@ -349,15 +349,26 @@ impl From<InvocationUuid> for opentelemetry::trace::SpanId {
 /// Services are isolated by key. This means that there cannot be two concurrent
 /// invocations for the same service instance (service name, key).
 #[derive(
-    Eq, Hash, PartialEq, PartialOrd, Ord, Clone, Debug, serde::Serialize, serde::Deserialize,
+    Eq,
+    Hash,
+    PartialEq,
+    PartialOrd,
+    Ord,
+    Clone,
+    Debug,
+    serde::Serialize,
+    serde::Deserialize,
+    bilrost::Message,
 )]
 pub struct ServiceId {
     // TODO rename this to KeyedServiceId. This type can be used only by keyed service types (virtual objects and workflows)
     /// Identifies the grpc service
+    #[bilrost(1)]
     pub service_name: ByteString,
     /// Identifies the service instance for the given service name
+    #[bilrost(2)]
     pub key: ByteString,
-
+    #[bilrost(3)]
     partition_key: PartitionKey,
 }
 

crates/types/src/identifiers.rs

@@ -19,7 +19,10 @@ use crate::identifiers::{
 };
 use crate::journal_v2::{CompletionId, GetInvocationOutputResult, Signal};
 use crate::time::MillisSinceEpoch;
-use crate::{GenerationalNodeId, RestateVersion};
+use crate::{
+    GenerationalNodeId, RestateVersion, bilrost_storage_encode_decode,
+    flexbuffers_storage_encode_decode,
+};
 
 use bytes::Bytes;
 use bytestring::ByteString;
@@ -435,6 +438,8 @@ pub struct ServiceInvocation {
     pub restate_version: RestateVersion,
 }
 
+flexbuffers_storage_encode_decode!(ServiceInvocation);
+
 #[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 #[serde(
     from = "serde_hacks::SubmitNotificationSink",
@@ -577,6 +582,8 @@ pub struct InvocationResponse {
     pub result: ResponseResult,
 }
 
+flexbuffers_storage_encode_decode!(InvocationResponse);
+
 impl WithInvocationId for InvocationResponse {
     fn invocation_id(&self) -> InvocationId {
         self.target.invocation_id()
@@ -623,6 +630,8 @@ pub struct GetInvocationOutputResponse {
     pub result: GetInvocationOutputResult,
 }
 
+bilrost_storage_encode_decode!(GetInvocationOutputResponse);
+
 impl WithInvocationId for GetInvocationOutputResponse {
     fn invocation_id(&self) -> InvocationId {
         self.target.invocation_id()
@@ -944,6 +953,8 @@ pub struct InvocationTermination {
     pub response_sink: Option<InvocationMutationResponseSink>,
 }
 
+flexbuffers_storage_encode_decode!(InvocationTermination);
+
 /// Flavor of the termination. Can be kill (hard stop) or graceful cancel.
 #[derive(
     Debug, Clone, Copy, Eq, PartialEq, serde::Serialize, serde::Deserialize, bilrost::Enumeration,
@@ -963,6 +974,8 @@ pub struct PurgeInvocationRequest {
     pub response_sink: Option<InvocationMutationResponseSink>,
 }
 
+flexbuffers_storage_encode_decode!(PurgeInvocationRequest);
+
 impl WithInvocationId for PurgeInvocationRequest {
     fn invocation_id(&self) -> InvocationId {
         self.invocation_id
@@ -979,6 +992,8 @@ pub struct ResumeInvocationRequest {
     pub response_sink: Option<InvocationMutationResponseSink>,
 }
 
+flexbuffers_storage_encode_decode!(ResumeInvocationRequest);
+
 impl WithInvocationId for ResumeInvocationRequest {
     fn invocation_id(&self) -> InvocationId {
         self.invocation_id
@@ -1001,6 +1016,8 @@ pub struct RestartAsNewInvocationRequest {
     pub response_sink: Option<InvocationMutationResponseSink>,
 }
 
+flexbuffers_storage_encode_decode!(RestartAsNewInvocationRequest);
+
 impl WithInvocationId for RestartAsNewInvocationRequest {
     fn invocation_id(&self) -> InvocationId {
         self.invocation_id
@@ -1320,6 +1337,8 @@ pub struct AttachInvocationRequest {
     pub response_sink: ServiceInvocationResponseSink,
 }
 
+flexbuffers_storage_encode_decode!(AttachInvocationRequest);
+
 impl WithPartitionKey for AttachInvocationRequest {
     fn partition_key(&self) -> PartitionKey {
         self.invocation_query.partition_key()
@@ -1333,6 +1352,8 @@ pub struct NotifySignalRequest {
     pub signal: Signal,
 }
 
+flexbuffers_storage_encode_decode!(NotifySignalRequest);
+
 impl WithInvocationId for NotifySignalRequest {
     fn invocation_id(&self) -> InvocationId {
         self.invocation_id

crates/types/src/invocation/mod.rs

@@ -10,7 +10,7 @@
 
 //! This module defines types used for the internal messaging between Restate components.
 
-use crate::identifiers::PartitionId;
+use crate::{bilrost_storage_encode_decode, identifiers::PartitionId};
 
 /// Wrapper that extends a message with its target peer to which the message should be sent.
 pub type PartitionTarget<Msg> = (PartitionId, Msg);
@@ -29,3 +29,11 @@ pub enum AckKind {
         last_known_seq_number: MessageIndex,
     },
 }
+
+#[derive(Debug, Clone, Copy, bilrost::Message)]
+pub struct MessageIndexRecord {
+    #[bilrost(1)]
+    pub index: MessageIndex,
+}
+
+bilrost_storage_encode_decode!(MessageIndexRecord);

crates/types/src/message.rs

@@ -16,21 +16,27 @@ use bytes::Bytes;
 use serde_with::serde_as;
 use sha2::{Digest, Sha256};
 
+use crate::bilrost_storage_encode_decode;
 use crate::identifiers::ServiceId;
 
 #[serde_as]
 /// ExternalStateMutation
 ///
 /// represents an external request to mutate a user's state.
-#[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize, bilrost::Message)]
 pub struct ExternalStateMutation {
+    #[bilrost(1)]
     pub service_id: ServiceId,
+    #[bilrost(2)]
     pub version: Option<String>,
     // flexbuffers only supports string-keyed maps :-( --> so we store it as vector of kv pairs
+    #[bilrost(3)]
     #[serde_as(as = "serde_with::Seq<(_, _)>")]
     pub state: HashMap<Bytes, Bytes>,
 }
 
+bilrost_storage_encode_decode!(ExternalStateMutation);
+
 /// # StateMutationVersion
 ///
 /// This type represents a user state version. This implementation hashes canonically the raw key-value

crates/types/src/state_mut.rs

@@ -27,6 +27,8 @@ use crate::journal_v2::raw::{RawEntry, RawEntryError, TryFromEntry};
 use crate::journal_v2::{Decoder, EntryMetadata, EntryType};
 use crate::time::MillisSinceEpoch;
 
+pub use tracing;
+
 #[derive(Debug, thiserror::Error)]
 pub enum StorageEncodeError {
     #[error("encoding failed: {0}")]
@@ -187,7 +189,7 @@ macro_rules! flexbuffers_storage_encode_decode {
                 Self: Sized,
             {
                 $crate::storage::decode::decode_serde(buf, kind).map_err(|err| {
-                    ::tracing::error!(%err, "{} decode failure (decoding {})", kind, stringify!($name));
+                    $crate::storage::tracing::error!(%err, "{} decode failure (decoding {})", kind, stringify!($name));
                     err
                 })
 
@@ -196,6 +198,40 @@ macro_rules! flexbuffers_storage_encode_decode {
     };
 }
 
+/// Implements the [`StorageEncode`] and [`StorageDecode`] by encoding/decoding the implementing
+/// type using [`bilrost`].
+#[macro_export]
+macro_rules! bilrost_storage_encode_decode {
+    ($name:tt) => {
+        impl $crate::storage::StorageEncode for $name {
+            fn default_codec(&self) -> $crate::storage::StorageCodecKind {
+                $crate::storage::StorageCodecKind::Bilrost
+            }
+
+            fn encode(
+                &self,
+                buf: &mut ::bytes::BytesMut,
+            ) -> Result<(), $crate::storage::StorageEncodeError> {
+                bytes::BufMut::put(buf, $crate::storage::encode::encode_bilrost(self));
+                Ok(())
+            }
+        }
+
+        impl $crate::storage::StorageDecode for $name {
+            fn decode<B: ::bytes::Buf>(
+                buf: &mut B,
+                kind: $crate::storage::StorageCodecKind,
+            ) -> Result<Self, $crate::storage::StorageDecodeError>
+            where
+                Self: Sized,
+            {
+                debug_assert_eq!(kind, $crate::storage::StorageCodecKind::Bilrost);
+                $crate::storage::decode::decode_bilrost(buf)
+            }
+        }
+    };
+}
+
 /// A polymorphic container of a buffer or a cached storage-encodeable object
 #[derive(Clone, derive_more::Debug, BilrostAs)]
 #[bilrost_as(dto::PolyBytes)]

crates/types/src/storage.rs

@@ -13,7 +13,7 @@ serde = ["dep:serde", "enum-map/serde", "bytestring/serde", "restate-invoker-api
 
 [dependencies]
 restate-workspace-hack = { workspace = true }
-
+restate-encoding = { workspace = true }
 restate-invoker-api = { workspace = true }
 restate-storage-api = { workspace = true }
 restate-types = { workspace = true }