Serialisation and compression module

ser changes

Author: sauravzg

grpc/Cargo.toml

@@ -14,9 +14,12 @@ allowed_external_types = [
 ]
 
 [features]
-default = ["dns", "_runtime-tokio", "protobuf"]
+default = ["dns", "_runtime-tokio", "protobuf", "gzip", "deflate", "zstd"]
 protobuf = ["dep:protobuf"]
 prost = ["dep:prost"]
+gzip = ["dep:flate2"]
+deflate = ["dep:flate2"]
+zstd = ["dep:zstd"]
 unsafe-downcast = []
 dns = ["dep:hickory-resolver", "_runtime-tokio"]
 # The following feature is used to ensure all modules use the runtime
@@ -60,6 +63,9 @@ tower = { version = "0.5.2", features = [
 ], optional = true }
 tower-service = "0.3.3"
 url = "2.5.0"
+arc-swap = "1.7"
+flate2 = { version = "1.0", optional = true }
+zstd = { version = "0.13", optional = true }
 trait-variant = "0.1"
 send-future = "0.1"
 

grpc/src/codec.rs

@@ -4,6 +4,9 @@ use tonic::{
     Status,
 };
 
+pub mod compression;
+pub mod serialization;
+
 /// An adapter for sending and receiving messages as bytes using tonic.
 /// Coding/decoding is handled within gRPC.
 /// TODO: Remove this when tonic allows access to bytes without requiring a

grpc/src/codec/compression.rs

@@ -0,0 +1,33 @@
+use bytes::{Buf, BufMut};
+use std::io;
+
+#[cfg(feature = "deflate")]
+pub mod deflate;
+#[cfg(feature = "gzip")]
+pub mod gzip;
+#[cfg(feature = "zstd")]
+pub mod zstd;
+
+pub mod registry;
+
+pub use self::registry::get_codec;
+
+/// A trait for identifying the encoding of a compression algorithm.
+pub trait Encoding {
+    /// The name of the compression algorithm, e.g., "gzip".
+    const NAME: &'static str;
+}
+
+/// A trait for compressing and decompressing data.
+pub trait Compressor: Send + Sync + 'static {
+    /// Compress data from `source` into `destination`.
+    fn compress(&self, source: &mut dyn Buf, destination: &mut dyn BufMut)
+        -> Result<(), io::Error>;
+
+    /// Decompress data from `source` into `destination`.
+    fn decompress(
+        &self,
+        source: &mut dyn Buf,
+        destination: &mut dyn BufMut,
+    ) -> Result<(), io::Error>;
+}

grpc/src/codec/compression/deflate.rs

@@ -0,0 +1,75 @@
+use crate::codec::compression::{Compressor, Encoding};
+use bytes::{Buf, BufMut};
+use flate2::{
+    bufread::{ZlibDecoder, ZlibEncoder},
+    Compression as FlateCompression,
+};
+use std::io;
+
+/// A deflate compression implementation.
+#[derive(Debug, Clone, Copy)]
+pub struct Deflate {
+    level: FlateCompression,
+}
+
+impl Deflate {
+    /// Creates a new deflate compression implementation.
+    pub fn new() -> Self {
+        Self {
+            level: FlateCompression::new(6),
+        }
+    }
+}
+
+impl Default for Deflate {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Compressor for Deflate {
+    fn compress(
+        &self,
+        source: &mut dyn Buf,
+        destination: &mut dyn BufMut,
+    ) -> Result<(), io::Error> {
+        let mut encoder = ZlibEncoder::new(source.reader(), self.level);
+        io::copy(&mut encoder, &mut destination.writer())?;
+        Ok(())
+    }
+
+    fn decompress(
+        &self,
+        source: &mut dyn Buf,
+        destination: &mut dyn BufMut,
+    ) -> Result<(), io::Error> {
+        let mut decoder = ZlibDecoder::new(source.reader());
+        io::copy(&mut decoder, &mut destination.writer())?;
+        Ok(())
+    }
+}
+
+impl Encoding for Deflate {
+    const NAME: &'static str = "deflate";
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use bytes::Bytes;
+
+    #[test]
+    fn deflate_compress_decompress() {
+        let compressor = Deflate::new();
+        let data = Bytes::from_static(b"hello world");
+        let mut compressed = Vec::new();
+        compressor
+            .compress(&mut data.clone(), &mut compressed)
+            .unwrap();
+        let mut decompressed = Vec::new();
+        compressor
+            .decompress(&mut compressed.as_slice(), &mut decompressed)
+            .unwrap();
+        assert_eq!(data, decompressed.as_slice());
+    }
+}

grpc/src/codec/compression/gzip.rs

@@ -0,0 +1,75 @@
+use crate::codec::compression::{Compressor, Encoding};
+use bytes::{Buf, BufMut};
+use flate2::{
+    bufread::{GzDecoder, GzEncoder},
+    Compression as FlateCompression,
+};
+use std::io;
+
+/// A gzip compression implementation.
+#[derive(Debug, Clone, Copy)]
+pub struct Gzip {
+    level: FlateCompression,
+}
+
+impl Gzip {
+    /// Creates a new gzip compression implementation.
+    pub fn new() -> Self {
+        Self {
+            level: FlateCompression::new(6),
+        }
+    }
+}
+
+impl Default for Gzip {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Compressor for Gzip {
+    fn compress(
+        &self,
+        source: &mut dyn Buf,
+        destination: &mut dyn BufMut,
+    ) -> Result<(), io::Error> {
+        let mut encoder = GzEncoder::new(source.reader(), self.level);
+        io::copy(&mut encoder, &mut destination.writer())?;
+        Ok(())
+    }
+
+    fn decompress(
+        &self,
+        source: &mut dyn Buf,
+        destination: &mut dyn BufMut,
+    ) -> Result<(), io::Error> {
+        let mut decoder = GzDecoder::new(source.reader());
+        io::copy(&mut decoder, &mut destination.writer())?;
+        Ok(())
+    }
+}
+
+impl Encoding for Gzip {
+    const NAME: &'static str = "gzip";
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use bytes::Bytes;
+
+    #[test]
+    fn gzip_compress_decompress() {
+        let compressor = Gzip::new();
+        let data = Bytes::from_static(b"hello world");
+        let mut compressed = Vec::new();
+        compressor
+            .compress(&mut data.clone(), &mut compressed)
+            .unwrap();
+        let mut decompressed = Vec::new();
+        compressor
+            .decompress(&mut compressed.as_slice(), &mut decompressed)
+            .unwrap();
+        assert_eq!(data, decompressed.as_slice());
+    }
+}

grpc/src/codec/compression/registry.rs

@@ -0,0 +1,123 @@
+use super::{Compressor, Encoding};
+use arc_swap::ArcSwap;
+use std::collections::HashMap;
+use std::sync::{Arc, LazyLock};
+
+/// A registry of compression implementations.
+#[derive(Default, Clone)]
+pub struct CompressionRegistry {
+    codecs: HashMap<&'static str, Arc<dyn Compressor>>,
+}
+
+// The global registry using ArcSwap for lock-free reads to implement something
+// closer to the RCU pattern.
+static GLOBAL_REGISTRY: LazyLock<ArcSwap<CompressionRegistry>> =
+    LazyLock::new(|| ArcSwap::from(Arc::new(CompressionRegistry::new())));
+
+/// Get a codec from the global registry.
+/// This operation is extremely fast and lock-free.
+pub fn get_codec(name: &str) -> Option<Arc<dyn Compressor>> {
+    GLOBAL_REGISTRY.load().get(name)
+}
+
+/// Add a new codec to the global registry using a copy-on-write strategy.
+pub fn add_codec(name: &'static str, codec: Arc<dyn Compressor>) {
+    let current_registry = GLOBAL_REGISTRY.load();
+    let new_registry = current_registry.with_codec(name, codec);
+    GLOBAL_REGISTRY.store(Arc::new(new_registry));
+}
+
+impl CompressionRegistry {
+    /// Creates a new compression registry with default codecs enabled by features.
+    pub fn new() -> Self {
+        let mut codecs: HashMap<&'static str, Arc<dyn Compressor>> = HashMap::new();
+
+        #[cfg(feature = "gzip")]
+        {
+            let gzip = Arc::new(super::gzip::Gzip::new());
+            codecs.insert(<super::gzip::Gzip as Encoding>::NAME, gzip);
+        }
+
+        #[cfg(feature = "deflate")]
+        {
+            let deflate = Arc::new(super::deflate::Deflate::new());
+            codecs.insert(<super::deflate::Deflate as Encoding>::NAME, deflate);
+        }
+
+        #[cfg(feature = "zstd")]
+        {
+            let zstd = Arc::new(super::zstd::Zstd::new());
+            codecs.insert(<super::zstd::Zstd as Encoding>::NAME, zstd);
+        }
+
+        Self { codecs }
+    }
+
+    /// Get a codec from this specific registry instance.
+    pub fn get(&self, name: &str) -> Option<Arc<dyn Compressor>> {
+        self.codecs.get(name).cloned()
+    }
+
+    /// Creates a new registry from an existing one, adding a new codec.
+    pub fn with_codec(&self, name: &'static str, codec: Arc<dyn Compressor>) -> Self {
+        // Clone the existing map of codecs
+        let mut new_codecs = self.codecs.clone();
+        // Add the new one
+        new_codecs.insert(name, codec);
+        // Return a new CompressionRegistry instance
+        Self { codecs: new_codecs }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::codec::compression::{Compressor, Encoding};
+    use bytes::{Buf, BufMut};
+    use std::io;
+
+    #[derive(Debug, Clone, Copy)]
+    struct MockCompression;
+
+    impl Compressor for MockCompression {
+        fn compress(
+            &self,
+            _source: &mut dyn Buf,
+            _destination: &mut dyn BufMut,
+        ) -> Result<(), io::Error> {
+            Ok(())
+        }
+
+        fn decompress(
+            &self,
+            _source: &mut dyn Buf,
+            _destination: &mut dyn BufMut,
+        ) -> Result<(), io::Error> {
+            Ok(())
+        }
+    }
+
+    impl Encoding for MockCompression {
+        const NAME: &'static str = "mock";
+    }
+
+    #[test]
+    fn registry_get_with() {
+        let registry = CompressionRegistry::new();
+        let registry = registry.with_codec("mock", Arc::new(MockCompression));
+        assert!(registry.get("mock").is_some());
+    }
+
+    #[test]
+    fn global_registry() {
+        #[cfg(feature = "gzip")]
+        assert!(get_codec("gzip").is_some());
+        #[cfg(feature = "deflate")]
+        assert!(get_codec("deflate").is_some());
+        #[cfg(feature = "zstd")]
+        assert!(get_codec("zstd").is_some());
+
+        add_codec("mock", Arc::new(MockCompression));
+        assert!(get_codec("mock").is_some());
+    }
+}