wal-protocol: introduce v2 Envelope with bilrost encoding and lazy payloads

crates/storage-api/src/deduplication_table/mod.rs

@@ -40,6 +40,7 @@ impl DedupInformation {
         }
     }
 
+    #[deprecated]
     pub fn ingress(producer_id: impl Into<ByteString>, sequence_number: MessageIndex) -> Self {
         DedupInformation {
             producer_id: ProducerId::Other(producer_id.into()),

crates/storage-api/src/vqueue_table/scheduler.rs

@@ -28,7 +28,7 @@ pub enum YieldReason {
     },
 }
 
-#[derive(Debug, bilrost::Oneof, bilrost::Message, derive_more::From)]
+#[derive(Debug, Clone, bilrost::Oneof, bilrost::Message, derive_more::From)]
 pub enum SchedulerAction {
     #[bilrost(empty)]
     Unknown,
@@ -40,23 +40,23 @@ pub enum SchedulerAction {
     Yield(YieldAction),
 }
 
-#[derive(Debug, bilrost::Message)]
+#[derive(Debug, Clone, bilrost::Message)]
 pub struct RunAction {
     #[bilrost(tag(1))]
     pub key: EntryKey,
     #[bilrost(tag(2))]
     pub wait_stats: WaitStats,
 }
 
-#[derive(Debug, bilrost::Message)]
+#[derive(Debug, Clone, bilrost::Message)]
 pub struct YieldAction {
     #[bilrost(tag(1))]
     pub key: EntryKey,
     #[bilrost(tag(2))]
     pub next_run_at: Option<RoughTimestamp>,
 }
 
-#[derive(Debug, bilrost::Message)]
+#[derive(Debug, Clone, bilrost::Message)]
 pub struct SchedulerDecisions {
     #[bilrost(tag(1))]
     pub qids: Vec<(VQueueId, Vec<SchedulerAction>)>,

crates/types/Cargo.toml

@@ -27,7 +27,7 @@ restate-clock = { workspace = true, features = ["prost-types", "jiff", "hlc"] }
 restate-encoding = { workspace = true }
 restate-errors = { workspace = true }
 restate-memory = { workspace = true }
-restate-platform = { workspace = true }
+restate-platform = { workspace = true, features = ["bilrost"] }
 restate-serde-util = { workspace = true }
 restate-sharding = { workspace = true, features = ["serde", "bilrost"] }
 restate-test-util = { workspace = true, optional = true }

crates/types/src/logs/mod.rs

@@ -327,6 +327,18 @@ impl<T: StorageEncode> BodyWithKeys<T> {
     pub fn into_inner(self) -> T {
         self.inner
     }
+
+    pub fn keys(&self) -> &Keys {
+        &self.keys
+    }
+
+    pub fn inner(&self) -> &T {
+        &self.inner
+    }
+
+    pub fn split(self) -> (Keys, T) {
+        (self.keys, self.inner)
+    }
 }
 
 impl<T> HasRecordKeys for BodyWithKeys<T>

crates/types/src/replication/nodeset.rs

@@ -406,6 +406,58 @@ fn write_nodes<'a>(
     write!(f, "]")
 }
 
+mod bilrost_encoding {
+    use bilrost::encoding::{EmptyState, ForOverwrite, Proxiable};
+
+    use crate::PlainNodeId;
+
+    use super::NodeSet;
+
+    impl Proxiable for NodeSet {
+        type Proxy = Vec<PlainNodeId>;
+
+        fn encode_proxy(&self) -> Self::Proxy {
+            Vec::from_iter(self.0.iter().cloned())
+        }
+
+        fn decode_proxy(&mut self, proxy: Self::Proxy) -> Result<(), bilrost::DecodeErrorKind> {
+            self.0.extend(proxy);
+            Ok(())
+        }
+    }
+
+    impl EmptyState<(), NodeSet> for () {
+        fn clear(val: &mut NodeSet) {
+            val.0.clear();
+        }
+        fn empty() -> NodeSet
+        where
+            NodeSet: Sized,
+        {
+            NodeSet::new()
+        }
+
+        fn is_empty(val: &NodeSet) -> bool {
+            val.is_empty()
+        }
+    }
+
+    impl ForOverwrite<(), NodeSet> for () {
+        fn for_overwrite() -> NodeSet
+        where
+            NodeSet: Sized,
+        {
+            NodeSet::new()
+        }
+    }
+
+    bilrost::delegate_proxied_encoding!(
+        use encoding (bilrost::encoding::GeneralPacked)
+        to encode proxied type (NodeSet)
+        with general encodings
+    );
+}
+
 #[cfg(test)]
 mod test {
     use ahash::HashMap;

crates/types/src/storage.rs

@@ -112,6 +112,39 @@ 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> {
+                $crate::storage::encode::encode_bilrost(self, buf)
+            }
+        }
+
+        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/wal-protocol/Cargo.toml

@@ -7,13 +7,20 @@ rust-version.workspace = true
 license.workspace = true
 publish = false
 
+[features]
+test-util = []
+
 [dependencies]
 restate-workspace-hack = { workspace = true }
 
+restate-encoding = { workspace = true }
+restate-limiter = { workspace = true }
 restate-storage-api = { workspace = true }
 restate-types = { workspace = true }
+restate-util-string = { workspace = true }
 restate-worker-api = { workspace = true, features = ["serde"] }
 
+anyhow = { workspace = true }
 bilrost = { workspace = true }
 bytes = { workspace = true }
 derive_more = { workspace = true, features = ["debug"] }

crates/wal-protocol/src/control.rs

@@ -9,28 +9,33 @@
 // by the Apache License, Version 2.0.
 
 use bytes::Bytes;
-
 use restate_storage_api::fsm_table::{CurrentReplicaSetState, NextReplicaSetState};
 use restate_types::identifiers::{LeaderEpoch, PartitionId};
-use restate_types::logs::{Keys, Lsn, SequenceNumber};
+use restate_types::logs::{HasRecordKeys, Keys, Lsn, SequenceNumber};
 use restate_types::partitions::PartitionConfiguration;
 use restate_types::partitions::state::{MemberState, ReplicaSetState};
 use restate_types::replication::{NodeSet, ReplicationProperty};
 use restate_types::schema::Schema;
 use restate_types::sharding::KeyRange;
 use restate_types::time::MillisSinceEpoch;
-use restate_types::{GenerationalNodeId, SemanticRestateVersion, Version, Versioned};
+use restate_types::{
+    GenerationalNodeId, SemanticRestateVersion, Version, Versioned, bilrost_storage_encode_decode,
+    flexbuffers_storage_encode_decode,
+};
 
 /// Announcing a new leader. This message can be written by any component to make the specified
 /// partition processor the leader.
-#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, bilrost::Message)]
 pub struct AnnounceLeader {
     /// Sender of the announce leader message.
     ///
     /// This became non-optional in v1.5. Noting that it has always been set in previous versions,
     /// it's safe to assume that it's always set.
+    #[bilrost(tag(1))]
     pub node_id: GenerationalNodeId,
+    #[bilrost(tag(2))]
     pub leader_epoch: LeaderEpoch,
+    #[bilrost(tag(3))]
     pub partition_key_range: KeyRange,
 
     /// Associated epoch metadata version
@@ -40,25 +45,40 @@ pub struct AnnounceLeader {
     ///
     /// *Since v1.6*
     #[serde(default, skip_serializing_if = "Option::is_none")]
+    #[bilrost(tag(4))]
     pub epoch_version: Option<Version>,
     /// Current replica set configuration at the time of the announcement.
     /// This field is optional for backward compatibility with older versions.
     /// *Since v1.6*
     #[serde(default, skip_serializing_if = "Option::is_none")]
+    #[bilrost(tag(5))]
     pub current_config: Option<CurrentReplicaSetConfiguration>,
     /// Next replica set configuration.
     /// *Since v1.6*
     #[serde(default, skip_serializing_if = "Option::is_none")]
+    #[bilrost(tag(6))]
     pub next_config: Option<NextReplicaSetConfiguration>,
 }
 
+bilrost_storage_encode_decode!(AnnounceLeader);
+
+impl HasRecordKeys for AnnounceLeader {
+    fn record_keys(&self) -> Keys {
+        Keys::RangeInclusive(self.partition_key_range.start()..=self.partition_key_range.end())
+    }
+}
+
 #[serde_with::serde_as]
-#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, bilrost::Message)]
 pub struct CurrentReplicaSetConfiguration {
+    #[bilrost(tag = 1)]
     pub version: Version,
+    #[bilrost(tag = 2)]
     pub replica_set: NodeSet,
+    #[bilrost(tag = 3)]
     pub modified_at: MillisSinceEpoch,
     #[serde_as(as = "serde_with::DisplayFromStr")]
+    #[bilrost(tag = 4)]
     pub replication: ReplicationProperty,
 }
 
@@ -87,9 +107,11 @@ impl CurrentReplicaSetConfiguration {
     }
 }
 
-#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, bilrost::Message)]
 pub struct NextReplicaSetConfiguration {
+    #[bilrost(tag(1))]
     pub version: Version,
+    #[bilrost(tag(2))]
     pub replica_set: NodeSet,
 }
 
@@ -138,7 +160,7 @@ fn new_replica_set_state(version: Version, node_set: &NodeSet) -> ReplicaSetStat
 /// Readers before v1.4.0 will crash when reading this command. For v1.4.0+, the barrier defines the
 /// minimum version of restate server that can progress after this command. It also updates the FSM
 /// in case command has been trimmed.
-#[derive(Debug, Clone, PartialEq, Eq, bilrost::Message, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, bilrost::Message, serde::Serialize, serde::Deserialize)]
 pub struct VersionBarrier {
     /// The minimum version required (inclusive) to progress after this barrier.
     pub version: SemanticRestateVersion,
@@ -147,21 +169,44 @@ pub struct VersionBarrier {
     pub partition_key_range: Keys,
 }
 
+bilrost_storage_encode_decode!(VersionBarrier);
+
+impl HasRecordKeys for VersionBarrier {
+    fn record_keys(&self) -> Keys {
+        self.partition_key_range.clone()
+    }
+}
+
 /// Updates the `PARTITION_DURABILITY` FSM variable to the given value. Note that durability
 /// only applies to partitions with the same `partition_id`. At replay time, the partition will
 /// ignore updates that are not targeted to its own ID.
 ///
 /// NOTE: The durability point is monotonically increasing.
 ///
 /// Since v1.4.2.
-#[derive(Debug, Clone, PartialEq, Eq, bilrost::Message, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, bilrost::Message, serde::Serialize, serde::Deserialize)]
 pub struct PartitionDurability {
+    #[bilrost(tag(1))]
     pub partition_id: PartitionId,
     /// The partition has applied this LSN durably to the replica-set and/or has been
     /// persisted in a snapshot in the snapshot repository.
+    #[bilrost(tag(2))]
     pub durable_point: Lsn,
     /// Timestamp which the durability point was updated
+    #[bilrost(tag(3))]
     pub modification_time: MillisSinceEpoch,
+    /// partition key range
+    #[bilrost(tag(4))]
+    #[serde(default)]
+    pub partition_key_range: Keys,
+}
+
+bilrost_storage_encode_decode!(PartitionDurability);
+
+impl HasRecordKeys for PartitionDurability {
+    fn record_keys(&self) -> Keys {
+        self.partition_key_range.clone()
+    }
 }
 
 /// Consistently store schema across partition replicas.
@@ -173,6 +218,14 @@ pub struct UpsertSchema {
     pub schema: Schema,
 }
 
+flexbuffers_storage_encode_decode!(UpsertSchema);
+
+impl HasRecordKeys for UpsertSchema {
+    fn record_keys(&self) -> Keys {
+        self.partition_key_range.clone()
+    }
+}
+
 /// Consistently distribute the cluster-global rule book across partition
 /// replicas. Each partition's leader observes a node-level cache of the
 /// rule book stored in the metadata store and proposes this command when

crates/wal-protocol/src/lib.rs

@@ -11,5 +11,6 @@
 pub mod control;
 pub mod timer;
 pub mod v1;
+pub mod v2;
 
 pub use v1::{Command, Destination, Envelope, Header, Source};