Compress digest

LICENSE-3rdparty.csv

@@ -2,43 +2,30 @@ Component,Origin,License,Copyright
 allocator-api2,https://github.com/zakarumych/allocator-api2,MIT OR Apache-2.0,Zakarum <zaq.dev@icloud.com>
 anyhow,https://github.com/dtolnay/anyhow,MIT OR Apache-2.0,David Tolnay <dtolnay@gmail.com>
 async-trait,https://github.com/dtolnay/async-trait,MIT OR Apache-2.0,David Tolnay <dtolnay@gmail.com>
-bitflags,https://github.com/bitflags/bitflags,MIT OR Apache-2.0,The Rust Project Developers
 bytes,https://github.com/tokio-rs/bytes,MIT,"Carl Lerche <me@carllerche.com>, Sean McArthur <sean@seanmonstar.com>"
-cfg-if,https://github.com/rust-lang/cfg-if,MIT OR Apache-2.0,Alex Crichton <alex@alexcrichton.com>
+cfg-if,https://github.com/alexcrichton/cfg-if,MIT OR Apache-2.0,Alex Crichton <alex@alexcrichton.com>
 chacha20,https://github.com/RustCrypto/stream-ciphers,MIT OR Apache-2.0,RustCrypto Developers
 cpufeatures,https://github.com/RustCrypto/utils,MIT OR Apache-2.0,RustCrypto Developers
-either,https://github.com/rayon-rs/either,MIT OR Apache-2.0,The either Authors
-equivalent,https://github.com/indexmap-rs/equivalent,Apache-2.0 OR MIT,The equivalent Authors
+either,https://github.com/rayon-rs/either,MIT OR Apache-2.0,bluss
+equivalent,https://github.com/cuviper/equivalent,Apache-2.0 OR MIT,The equivalent Authors
 foldhash,https://github.com/orlp/foldhash,Zlib,Orson Peters <orsonpeters@gmail.com>
 futures-core,https://github.com/rust-lang/futures-rs,MIT OR Apache-2.0,The futures-core Authors
 futures-sink,https://github.com/rust-lang/futures-rs,MIT OR Apache-2.0,The futures-sink Authors
 getrandom,https://github.com/rust-random/getrandom,MIT OR Apache-2.0,The Rand Project Developers
-hashbrown,https://github.com/rust-lang/hashbrown,MIT OR Apache-2.0,Amanieu d'Antras <amanieu@gmail.com>
 hashbrown,https://github.com/rust-lang/hashbrown,MIT OR Apache-2.0,The hashbrown Authors
-heck,https://github.com/withoutboats/heck,MIT OR Apache-2.0,The heck Authors
-id-arena,https://github.com/fitzgen/id-arena,MIT OR Apache-2.0,"Nick Fitzgerald <fitzgen@gmail.com>, Aleksey Kladov <aleksey.kladov@gmail.com>"
-indexmap,https://github.com/indexmap-rs/indexmap,Apache-2.0 OR MIT,The indexmap Authors
 itertools,https://github.com/rust-itertools/itertools,MIT OR Apache-2.0,bluss
-itoa,https://github.com/dtolnay/itoa,MIT OR Apache-2.0,David Tolnay <dtolnay@gmail.com>
-leb128fmt,https://github.com/bluk/leb128fmt,MIT OR Apache-2.0,Bryant Luk <code@bryantluk.com>
 libc,https://github.com/rust-lang/libc,MIT OR Apache-2.0,The Rust Project Developers
-log,https://github.com/rust-lang/log,MIT OR Apache-2.0,The Rust Project Developers
 lru,https://github.com/jeromefroe/lru-rs,MIT,Jerome Froelich <jeromefroelic@hotmail.com>
-memchr,https://github.com/BurntSushi/memchr,Unlicense OR MIT,"Andrew Gallant <jamslam@gmail.com>, bluss"
 mio,https://github.com/tokio-rs/mio,MIT,"Carl Lerche <me@carllerche.com>, Thomas de Zeeuw <thomasdezeeuw@gmail.com>, Tokio Contributors <team@tokio.rs>"
 once_cell,https://github.com/matklad/once_cell,MIT OR Apache-2.0,Aleksey Kladov <aleksey.kladov@gmail.com>
 pin-project-lite,https://github.com/taiki-e/pin-project-lite,Apache-2.0 OR MIT,The pin-project-lite Authors
-prettyplease,https://github.com/dtolnay/prettyplease,MIT OR Apache-2.0,David Tolnay <dtolnay@gmail.com>
 proc-macro2,https://github.com/dtolnay/proc-macro2,MIT OR Apache-2.0,"David Tolnay <dtolnay@gmail.com>, Alex Crichton <alex@alexcrichton.com>"
 quote,https://github.com/dtolnay/quote,MIT OR Apache-2.0,David Tolnay <dtolnay@gmail.com>
 r-efi,https://github.com/r-efi/r-efi,MIT OR Apache-2.0 OR LGPL-2.1-or-later,The r-efi Authors
 rand,https://github.com/rust-random/rand,MIT OR Apache-2.0,"The Rand Project Developers, The Rust Project Developers"
 rand_core,https://github.com/rust-random/rand_core,MIT OR Apache-2.0,The Rand Project Developers
-semver,https://github.com/dtolnay/semver,MIT OR Apache-2.0,David Tolnay <dtolnay@gmail.com>
 serde,https://github.com/serde-rs/serde,MIT OR Apache-2.0,"Erick Tryzelaar <erick.tryzelaar@gmail.com>, David Tolnay <dtolnay@gmail.com>"
-serde_core,https://github.com/serde-rs/serde,MIT OR Apache-2.0,"Erick Tryzelaar <erick.tryzelaar@gmail.com>, David Tolnay <dtolnay@gmail.com>"
 serde_derive,https://github.com/serde-rs/serde,MIT OR Apache-2.0,"Erick Tryzelaar <erick.tryzelaar@gmail.com>, David Tolnay <dtolnay@gmail.com>"
-serde_json,https://github.com/serde-rs/json,MIT OR Apache-2.0,"Erick Tryzelaar <erick.tryzelaar@gmail.com>, David Tolnay <dtolnay@gmail.com>"
 socket2,https://github.com/rust-lang/socket2,MIT OR Apache-2.0,"Alex Crichton <alex@alexcrichton.com>, Thomas de Zeeuw <thomasdezeeuw@gmail.com>"
 syn,https://github.com/dtolnay/syn,MIT OR Apache-2.0,David Tolnay <dtolnay@gmail.com>
 tokio,https://github.com/tokio-rs/tokio,MIT,Tokio Contributors <team@tokio.rs>
@@ -49,23 +36,10 @@ tracing,https://github.com/tokio-rs/tracing,MIT,"Eliza Weisman <eliza@buoyant.io
 tracing-attributes,https://github.com/tokio-rs/tracing,MIT,"Tokio Contributors <team@tokio.rs>, Eliza Weisman <eliza@buoyant.io>, David Barsky <dbarsky@amazon.com>"
 tracing-core,https://github.com/tokio-rs/tracing,MIT,Tokio Contributors <team@tokio.rs>
 unicode-ident,https://github.com/dtolnay/unicode-ident,(MIT OR Apache-2.0) AND Unicode-3.0,David Tolnay <dtolnay@gmail.com>
-unicode-xid,https://github.com/unicode-rs/unicode-xid,MIT OR Apache-2.0,"erick.tryzelaar <erick.tryzelaar@gmail.com>, kwantam <kwantam@gmail.com>, Manish Goregaokar <manishsmail@gmail.com>"
 valuable,https://github.com/tokio-rs/valuable,MIT,The valuable Authors
 wasi,https://github.com/bytecodealliance/wasi,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,The Cranelift Project Developers
-wasip2,https://github.com/bytecodealliance/wasi-rs,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,The wasip2 Authors
-wasip3,https://github.com/bytecodealliance/wasi-rs,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,The wasip3 Authors
-wasm-encoder,https://github.com/bytecodealliance/wasm-tools/tree/main/crates/wasm-encoder,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,Nick Fitzgerald <fitzgen@gmail.com>
-wasm-metadata,https://github.com/bytecodealliance/wasm-tools/tree/main/crates/wasm-metadata,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,The wasm-metadata Authors
-wasmparser,https://github.com/bytecodealliance/wasm-tools/tree/main/crates/wasmparser,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,Yury Delendik <ydelendik@mozilla.com>
 windows-link,https://github.com/microsoft/windows-rs,MIT OR Apache-2.0,The windows-link Authors
 windows-sys,https://github.com/microsoft/windows-rs,MIT OR Apache-2.0,The windows-sys Authors
-wit-bindgen,https://github.com/bytecodealliance/wit-bindgen,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,Alex Crichton <alex@alexcrichton.com>
-wit-bindgen-core,https://github.com/bytecodealliance/wit-bindgen,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,Alex Crichton <alex@alexcrichton.com>
-wit-bindgen-rust,https://github.com/bytecodealliance/wit-bindgen,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,Alex Crichton <alex@alexcrichton.com>
-wit-bindgen-rust-macro,https://github.com/bytecodealliance/wit-bindgen,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,Alex Crichton <alex@alexcrichton.com>
-wit-component,https://github.com/bytecodealliance/wasm-tools/tree/main/crates/wit-component,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,Peter Huene <peter@huene.dev>
-wit-parser,https://github.com/bytecodealliance/wasm-tools/tree/main/crates/wit-parser,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,Alex Crichton <alex@alexcrichton.com>
-zmij,https://github.com/dtolnay/zmij,MIT,David Tolnay <dtolnay@gmail.com>
 zstd,https://github.com/gyscos/zstd-rs,MIT,Alexandre Bury <alexandre.bury@gmail.com>
 zstd-safe,https://github.com/gyscos/zstd-rs,MIT OR Apache-2.0,Alexandre Bury <alexandre.bury@gmail.com>
 zstd-sys,https://github.com/gyscos/zstd-rs,MIT OR Apache-2.0,Alexandre Bury <alexandre.bury@gmail.com>

chitchat-test/src/main.rs

@@ -3,7 +3,9 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
 use chitchat::transport::UdpTransport;
-use chitchat::{Chitchat, ChitchatConfig, ChitchatId, FailureDetectorConfig, spawn_chitchat};
+use chitchat::{
+    Chitchat, ChitchatConfig, ChitchatId, FailureDetectorConfig, ProtocolVersion, spawn_chitchat,
+};
 use chitchat_test::{ApiResponse, SetKeyValueResponse};
 use cool_id_generator::Size;
 use poem::listener::TcpListener;
@@ -112,6 +114,7 @@ async fn main() -> anyhow::Result<()> {
         marked_for_deletion_grace_period: Duration::from_secs(60),
         catchup_callback: None,
         extra_liveness_predicate: None,
+        protocol_version: ProtocolVersion::V1,
     };
     let chitchat_handler = spawn_chitchat(config, Vec::new(), &UdpTransport).await?;
     let chitchat = chitchat_handler.chitchat();

chitchat/src/configuration.rs

@@ -3,7 +3,7 @@
 use std::net::SocketAddr;
 use std::time::Duration;
 
-use crate::{ChitchatId, FailureDetectorConfig, NodeState};
+use crate::{ChitchatId, FailureDetectorConfig, NodeState, ProtocolVersion};
 
 /// An optional user-defined callback executed when the self node is lagging behind.
 pub type CatchupCallback = Box<dyn Fn() + Send>;
@@ -41,6 +41,10 @@ pub struct ChitchatConfig {
     // It can be used for instance, to only surface the nodes that are both alive according
     // to the failure detector, but also have a given set of required keys.
     pub extra_liveness_predicate: Option<ExtraLivenessPredicate>,
+    /// Protocol version used to exchange messages between peers. The latest version is V1, which
+    /// compresses the digest. This MUST stay `V0` until every node in the cluster runs a version
+    /// able to decode the newer format.
+    pub protocol_version: ProtocolVersion,
 }
 
 impl ChitchatConfig {
@@ -58,6 +62,7 @@ impl ChitchatConfig {
             marked_for_deletion_grace_period: Duration::from_secs(10_000),
             catchup_callback: None,
             extra_liveness_predicate: None,
+            protocol_version: ProtocolVersion::V0,
         }
     }
 }
@@ -77,6 +82,7 @@ impl Default for ChitchatConfig {
             marked_for_deletion_grace_period: Duration::from_secs(3_600 * 2), // 2h
             catchup_callback: None,
             extra_liveness_predicate: None,
+            protocol_version: ProtocolVersion::V0,
         }
     }
 }

chitchat/src/digest.rs

@@ -1,8 +1,12 @@
 use std::collections::BTreeMap;
 
+use crate::message::ProtocolVersion;
 use crate::serialize::*;
 use crate::{ChitchatId, Heartbeat, Version};
 
+/// Block size used by the compressed digest stream.
+const DIGEST_BLOCK_THRESHOLD: u16 = 16_384u16;
+
 #[derive(Debug, Clone, Copy, Default, Eq, PartialEq)]
 pub(crate) struct NodeDigest {
     pub(crate) heartbeat: Heartbeat,
@@ -37,6 +41,15 @@ impl Deserializable for NodeDigest {
     }
 }
 
+/// Decodes one `(ChitchatId, NodeDigest)` entry of the compressed digest stream.
+impl Deserializable for (ChitchatId, NodeDigest) {
+    fn deserialize(buf: &mut &[u8]) -> anyhow::Result<Self> {
+        let chitchat_id = ChitchatId::deserialize(buf)?;
+        let node_digest = NodeDigest::deserialize(buf)?;
+        Ok((chitchat_id, node_digest))
+    }
+}
+
 /// A digest represents is a piece of information summarizing
 /// the staleness of one peer's data.
 ///
@@ -63,28 +76,81 @@ impl Digest {
         };
         self.node_digests.insert(node, node_digest);
     }
+
+    pub fn sample_for_test(num_nodes: u16) -> Digest {
+        let mut digest = Digest::default();
+        for i in 0..num_nodes {
+            let node = ChitchatId::for_local_test(10_000 + i);
+            digest.add_node(
+                node,
+                Heartbeat(100 + i as u64),
+                i as u64 % 7,
+                1_000 + i as u64,
+            );
+        }
+        digest
+    }
 }
 
-impl Serializable for Digest {
-    fn serialize(&self, buf: &mut Vec<u8>) {
+impl Digest {
+    /// Serializes the digest using the wire format selected by `protocol_version`.
+    pub(crate) fn serialize(&self, protocol_version: ProtocolVersion, buf: &mut Vec<u8>) {
+        match protocol_version {
+            ProtocolVersion::V0 => self.serialize_v0(buf),
+            ProtocolVersion::V1 => self.serialize_v1(buf),
+        }
+    }
+
+    pub(crate) fn serialized_len(&self, protocol_version: ProtocolVersion) -> usize {
+        match protocol_version {
+            ProtocolVersion::V0 => {
+                let mut len = (self.node_digests.len() as u16).serialized_len();
+                for (chitchat_id, node_digest) in &self.node_digests {
+                    len += chitchat_id.serialized_len();
+                    len += node_digest.serialized_len();
+                }
+                len
+            }
+            ProtocolVersion::V1 => {
+                let mut buf = Vec::new();
+                self.serialize_v1(&mut buf);
+                buf.len()
+            }
+        }
+    }
+
+    pub(crate) fn deserialize(
+        protocol_version: ProtocolVersion,
+        buf: &mut &[u8],
+    ) -> anyhow::Result<Self> {
+        match protocol_version {
+            ProtocolVersion::V0 => Self::deserialize_v0(buf),
+            ProtocolVersion::V1 => Self::deserialize_v1(buf),
+        }
+    }
+
+    /// Serializes the digest using the V0 (uncompressed) wire format.
+    fn serialize_v0(&self, buf: &mut Vec<u8>) {
         (self.node_digests.len() as u16).serialize(buf);
         for (chitchat_id, node_digest) in &self.node_digests {
             chitchat_id.serialize(buf);
             node_digest.serialize(buf);
         }
     }
-    fn serialized_len(&self) -> usize {
-        let mut len = (self.node_digests.len() as u16).serialized_len();
+
+    /// Serializes the digest using the V1 (compressed) wire format: a zstd block
+    /// stream of `(ChitchatId, NodeDigest)` entries.
+    fn serialize_v1(&self, buf: &mut Vec<u8>) {
+        let mut stream_writer =
+            CompressedStreamWriter::with_block_threshold(DIGEST_BLOCK_THRESHOLD);
         for (chitchat_id, node_digest) in &self.node_digests {
-            len += chitchat_id.serialized_len();
-            len += node_digest.serialized_len();
+            stream_writer.append(chitchat_id);
+            stream_writer.append(node_digest);
         }
-        len
+        buf.extend_from_slice(&stream_writer.finish());
     }
-}
 
-impl Deserializable for Digest {
-    fn deserialize(buf: &mut &[u8]) -> anyhow::Result<Self> {
+    fn deserialize_v0(buf: &mut &[u8]) -> anyhow::Result<Self> {
         let num_nodes = u16::deserialize(buf)?;
         let mut node_digests: BTreeMap<ChitchatId, NodeDigest> = Default::default();
 
@@ -95,14 +161,29 @@ impl Deserializable for Digest {
         }
         Ok(Digest { node_digests })
     }
+
+    fn deserialize_v1(buf: &mut &[u8]) -> anyhow::Result<Self> {
+        let node_digests = deserialize_stream::<(ChitchatId, NodeDigest)>(buf)?
+            .into_iter()
+            .collect();
+        Ok(Digest { node_digests })
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use crate::digest::{Digest, NodeDigest};
+    use crate::message::ProtocolVersion;
     use crate::serialize::test_serdeser_aux;
     use crate::{ChitchatId, Heartbeat};
 
+    fn digest_serdeser_aux(digest: &Digest, protocol_version: ProtocolVersion) -> Digest {
+        let mut buf = Vec::new();
+        digest.serialize(protocol_version, &mut buf);
+        assert_eq!(buf.len(), digest.serialized_len(protocol_version));
+        Digest::deserialize(protocol_version, &mut &buf[..]).unwrap()
+    }
+
     #[test]
     fn test_node_digest_serialization() {
         let node_digest = NodeDigest {
@@ -122,6 +203,27 @@ mod tests {
         digest.add_node(node1, Heartbeat(101), 1, 11);
         digest.add_node(node2, Heartbeat(102), 20, 12);
         digest.add_node(node3, Heartbeat(103), 0, 13);
-        test_serdeser_aux(&digest, 104);
+
+        let mut buf = Vec::new();
+        digest.serialize(ProtocolVersion::V0, &mut buf);
+        assert_eq!(buf.len(), 104);
+        assert_eq!(digest_serdeser_aux(&digest, ProtocolVersion::V0), digest);
+    }
+
+    #[test]
+    fn test_digest_compressed_roundtrip() {
+        for num_nodes in [0u16, 1, 3, 500] {
+            let digest = Digest::sample_for_test(num_nodes);
+            assert_eq!(digest_serdeser_aux(&digest, ProtocolVersion::V1), digest);
+        }
+    }
+
+    #[test]
+    fn test_digest_v0_and_v1_decode_identically() {
+        let digest = Digest::sample_for_test(200);
+        assert_eq!(
+            digest_serdeser_aux(&digest, ProtocolVersion::V0),
+            digest_serdeser_aux(&digest, ProtocolVersion::V1),
+        );
     }
 }