[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

@@ -11,16 +11,17 @@
 use crate::EffectKind::JournalEntryV2;
 use restate_types::deployment::PinnedDeployment;
 use restate_types::errors::InvocationError;
-use restate_types::identifiers::InvocationId;
+use restate_types::identifiers::{InvocationId, WithPartitionKey};
 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::logs::{HasRecordKeys, Keys};
 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 +36,14 @@ pub struct Effect {
     pub kind: EffectKind,
 }
 
+flexbuffers_storage_encode_decode!(Effect);
+
+impl HasRecordKeys for Effect {
+    fn record_keys(&self) -> restate_types::logs::Keys {
+        Keys::Single(self.invocation_id.partition_key())
+    }
+}
+
 #[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

@@ -12,25 +12,30 @@
 
 pub mod client;
 
+use std::borrow::Cow;
+use std::hash::Hash;
+use std::ops::Deref;
+use std::str::FromStr;
+use std::time::Duration;
+use std::{cmp, fmt};
+
+use bytes::Bytes;
+use bytestring::ByteString;
+use opentelemetry::trace::{SpanContext, SpanId, TraceFlags, TraceState};
+use serde_with::{DisplayFromStr, FromInto, serde_as};
+
 use crate::errors::InvocationError;
 use crate::identifiers::{
     DeploymentId, EntryIndex, IdempotencyId, InvocationId, PartitionKey,
     PartitionProcessorRpcRequestId, ServiceId, SubscriptionId, WithInvocationId, WithPartitionKey,
 };
 use crate::journal_v2::{CompletionId, GetInvocationOutputResult, Signal};
+use crate::logs::{HasRecordKeys, Keys};
 use crate::time::MillisSinceEpoch;
-use crate::{GenerationalNodeId, RestateVersion};
-
-use bytes::Bytes;
-use bytestring::ByteString;
-use opentelemetry::trace::{SpanContext, SpanId, TraceFlags, TraceState};
-use serde_with::{DisplayFromStr, FromInto, serde_as};
-use std::borrow::Cow;
-use std::hash::Hash;
-use std::ops::Deref;
-use std::str::FromStr;
-use std::time::Duration;
-use std::{cmp, fmt};
+use crate::{
+    GenerationalNodeId, RestateVersion, bilrost_storage_encode_decode,
+    flexbuffers_storage_encode_decode,
+};
 
 // Re-exporting opentelemetry [`TraceId`] to avoid having to import opentelemetry in all crates.
 pub use opentelemetry::trace::TraceId;
@@ -435,6 +440,14 @@ pub struct ServiceInvocation {
     pub restate_version: RestateVersion,
 }
 
+flexbuffers_storage_encode_decode!(ServiceInvocation);
+
+impl HasRecordKeys for ServiceInvocation {
+    fn record_keys(&self) -> Keys {
+        Keys::Single(self.invocation_id.partition_key())
+    }
+}
+
 #[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 #[serde(
     from = "serde_hacks::SubmitNotificationSink",
@@ -577,12 +590,20 @@ pub struct InvocationResponse {
     pub result: ResponseResult,
 }
 
+flexbuffers_storage_encode_decode!(InvocationResponse);
+
 impl WithInvocationId for InvocationResponse {
     fn invocation_id(&self) -> InvocationId {
         self.target.invocation_id()
     }
 }
 
+impl HasRecordKeys for InvocationResponse {
+    fn record_keys(&self) -> Keys {
+        Keys::Single(self.partition_key())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub enum ResponseResult {
     Success(Bytes),
@@ -623,12 +644,20 @@ pub struct GetInvocationOutputResponse {
     pub result: GetInvocationOutputResult,
 }
 
+bilrost_storage_encode_decode!(GetInvocationOutputResponse);
+
 impl WithInvocationId for GetInvocationOutputResponse {
     fn invocation_id(&self) -> InvocationId {
         self.target.invocation_id()
     }
 }
 
+impl HasRecordKeys for GetInvocationOutputResponse {
+    fn record_keys(&self) -> Keys {
+        Keys::Single(self.partition_key())
+    }
+}
+
 /// Definition of the sink where to send the result of a service invocation.
 #[derive(Debug, PartialEq, Eq, Clone, Hash, serde::Serialize, serde::Deserialize)]
 #[serde(
@@ -944,6 +973,14 @@ pub struct InvocationTermination {
     pub response_sink: Option<InvocationMutationResponseSink>,
 }
 
+flexbuffers_storage_encode_decode!(InvocationTermination);
+
+impl HasRecordKeys for InvocationTermination {
+    fn record_keys(&self) -> crate::logs::Keys {
+        Keys::Single(self.invocation_id.partition_key())
+    }
+}
+
 /// 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,12 +1000,20 @@ 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
     }
 }
 
+impl HasRecordKeys for PurgeInvocationRequest {
+    fn record_keys(&self) -> Keys {
+        Keys::Single(self.invocation_id.partition_key())
+    }
+}
+
 /// Message to resume an invocation.
 #[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
 pub struct ResumeInvocationRequest {
@@ -979,12 +1024,20 @@ 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
     }
 }
 
+impl HasRecordKeys for ResumeInvocationRequest {
+    fn record_keys(&self) -> Keys {
+        Keys::Single(self.invocation_id.partition_key())
+    }
+}
+
 /// Message to restart an invocation.
 #[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
 pub struct RestartAsNewInvocationRequest {
@@ -1001,12 +1054,20 @@ 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
     }
 }
 
+impl HasRecordKeys for RestartAsNewInvocationRequest {
+    fn record_keys(&self) -> Keys {
+        Keys::Single(self.invocation_id.partition_key())
+    }
+}
+
 // We use this struct instead of SpanContext as it is serialisable and allows us to use TraceStateDef
 #[serde_as]
 #[derive(Debug, PartialEq, Eq, Clone, serde::Serialize, serde::Deserialize)]
@@ -1320,25 +1381,41 @@ 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()
     }
 }
 
+impl HasRecordKeys for AttachInvocationRequest {
+    fn record_keys(&self) -> Keys {
+        Keys::Single(self.partition_key())
+    }
+}
+
 /// Represents a request to notify a signal to an invocation
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub struct NotifySignalRequest {
     pub invocation_id: InvocationId,
     pub signal: Signal,
 }
 
+flexbuffers_storage_encode_decode!(NotifySignalRequest);
+
 impl WithInvocationId for NotifySignalRequest {
     fn invocation_id(&self) -> InvocationId {
         self.invocation_id
     }
 }
 
+impl HasRecordKeys for NotifySignalRequest {
+    fn record_keys(&self) -> Keys {
+        Keys::Single(self.partition_key())
+    }
+}
+
 /// The invocation epoch represents the restarts count of the invocation, as seen from the Partition processor.
 pub type InvocationEpoch = u32;
 

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 TruncateOutboxRequest {
+    #[bilrost(1)]
+    pub index: MessageIndex,
+}
+
+bilrost_storage_encode_decode!(TruncateOutboxRequest);