[wip] Introduce Envelope v2

crates/types/src/replication/nodeset.rs

@@ -406,6 +406,61 @@
     write!(f, "]")
 }
 
+mod bilrost_impl {
+    use bilrost::{
+        Canonicity,
+        encoding::{DistinguishedProxiable, 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.into_iter());
+            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

@@ -14,6 +14,7 @@ restate-storage-api = { workspace = true }
 restate-types = { 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

@@ -17,18 +17,24 @@ 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, PartialEq, Eq, 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
@@ -38,25 +44,34 @@ 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);
+
 #[serde_with::serde_as]
-#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, 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,
 }
 
@@ -85,9 +100,11 @@ impl CurrentReplicaSetConfiguration {
     }
 }
 
-#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, bilrost::Message)]
 pub struct NextReplicaSetConfiguration {
+    #[bilrost(tag(1))]
     pub version: Version,
+    #[bilrost(tag(2))]
     pub replica_set: NodeSet,
 }
 
@@ -145,6 +162,8 @@ pub struct VersionBarrier {
     pub partition_key_range: Keys,
 }
 
+bilrost_storage_encode_decode!(VersionBarrier);
+
 /// 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.
@@ -162,6 +181,8 @@ pub struct PartitionDurability {
     pub modification_time: MillisSinceEpoch,
 }
 
+bilrost_storage_encode_decode!(PartitionDurability);
+
 /// Consistently store schema across partition replicas.
 ///
 /// Since v1.6.0.
@@ -170,3 +191,5 @@ pub struct UpsertSchema {
     pub partition_key_range: Keys,
     pub schema: Schema,
 }
+
+flexbuffers_storage_encode_decode!(UpsertSchema);