RSCBC-32: Reduce allocations on KV hot paths

Changes:
- Get result values use Bytes instead of Vec<u8>, eliminating a full
  document-body memcpy on every get/get_and_lock/get_and_touch/get_meta
  by propagating zero-copy Bytes from the codec through to the SDK layer
- RetryRequest.retry_reasons is now Option<HashSet>, only allocating on
  the first retry attempt instead of every operation
- Compressor.compress returns Cow<[u8]> instead of borrowing from self,
  removing heap state from StdCompressor entirely
- MutationToken bucket_name uses Arc<str> instead of String, replacing a
  per-mutation String::clone with an atomic ref-count bump
- Add unit tests for compression manager

Author: chvck

sdk/couchbase-core/src/compressionmanager.rs

@@ -16,25 +16,24 @@
  *
  */
 
+use snap::raw::Encoder;
+use std::borrow::Cow;
 use std::fmt::Debug;
 use std::marker::PhantomData;
 
-use snap::raw::Encoder;
-
 use crate::error;
 use crate::error::ErrorKind;
 use crate::memdx::datatype::DataTypeFlag;
 use crate::options::agent::{CompressionConfig, CompressionMode};
 
 pub(crate) trait Compressor: Send + Sync + Debug {
     fn new(compression_config: &CompressionConfig) -> Self;
-    // This is a bit of a weird signature,  but it allows us to avoid allocations when no compression occurs.
     fn compress<'a>(
-        &'a mut self,
+        &mut self,
         connection_supports_snappy: bool,
         datatype: DataTypeFlag,
         input: &'a [u8],
-    ) -> error::Result<(&'a [u8], u8)>;
+    ) -> error::Result<(Cow<'a, [u8]>, u8)>;
 }
 
 #[derive(Debug)]
@@ -64,12 +63,8 @@ pub(crate) struct StdCompressor {
     compression_enabled: bool,
     compression_min_size: usize,
     compression_min_ratio: f64,
-
-    compressed_value: Vec<u8>,
 }
 
-impl StdCompressor {}
-
 impl Compressor for StdCompressor {
     fn new(compression_config: &CompressionConfig) -> Self {
         let (compression_enabled, compression_min_size, compression_min_ratio) =
@@ -85,33 +80,31 @@ impl Compressor for StdCompressor {
             compression_enabled,
             compression_min_size,
             compression_min_ratio,
-
-            compressed_value: Vec::new(),
         }
     }
 
     fn compress<'a>(
-        &'a mut self,
+        &mut self,
         connection_supports_snappy: bool,
         datatype: DataTypeFlag,
         input: &'a [u8],
-    ) -> error::Result<(&'a [u8], u8)> {
+    ) -> error::Result<(Cow<'a, [u8]>, u8)> {
         if !connection_supports_snappy || !self.compression_enabled {
-            return Ok((input, u8::from(datatype)));
+            return Ok((Cow::Borrowed(input), u8::from(datatype)));
         }
 
         let datatype = u8::from(datatype);
 
         // If the packet is already compressed then we don't want to compress it again.
         if datatype & u8::from(DataTypeFlag::Compressed) != 0 {
-            return Ok((input, datatype));
+            return Ok((Cow::Borrowed(input), datatype));
         }
 
         let packet_size = input.len();
 
         // Only compress values that are large enough to be worthwhile.
         if packet_size <= self.compression_min_size {
-            return Ok((input, datatype));
+            return Ok((Cow::Borrowed(input), datatype));
         }
 
         let mut encoder = Encoder::new();
@@ -121,14 +114,135 @@ impl Compressor for StdCompressor {
 
         // Only return the compressed value if the ratio of compressed:original is small enough.
         if compressed_value.len() as f64 / packet_size as f64 > self.compression_min_ratio {
-            return Ok((input, datatype));
+            return Ok((Cow::Borrowed(input), datatype));
         }
 
-        self.compressed_value = compressed_value;
-
         Ok((
-            &self.compressed_value,
+            Cow::Owned(compressed_value),
             datatype | u8::from(DataTypeFlag::Compressed),
         ))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::borrow::Cow;
+
+    fn enabled_config(min_size: usize, min_ratio: f64) -> CompressionConfig {
+        CompressionConfig::new(CompressionMode::Enabled {
+            min_size,
+            min_ratio,
+        })
+    }
+
+    fn disabled_config() -> CompressionConfig {
+        CompressionConfig::new(CompressionMode::Disabled)
+    }
+
+    #[test]
+    fn disabled_compression_returns_input_unchanged() {
+        let mut compressor = StdCompressor::new(&disabled_config());
+        let input = b"hello world";
+
+        let (output, dt) = compressor
+            .compress(true, DataTypeFlag::Json, input)
+            .unwrap();
+
+        assert!(matches!(output, Cow::Borrowed(_)));
+        assert_eq!(&*output, input.as_slice());
+        assert_eq!(dt, u8::from(DataTypeFlag::Json));
+    }
+
+    #[test]
+    fn connection_without_snappy_returns_input_unchanged() {
+        let mut compressor = StdCompressor::new(&enabled_config(0, 1.0));
+        let input = b"hello world";
+
+        let (output, dt) = compressor
+            .compress(false, DataTypeFlag::Json, input)
+            .unwrap();
+
+        assert!(matches!(output, Cow::Borrowed(_)));
+        assert_eq!(&*output, input.as_slice());
+    }
+
+    #[test]
+    fn already_compressed_returns_input_unchanged() {
+        let mut compressor = StdCompressor::new(&enabled_config(0, 1.0));
+        let input = b"already compressed data";
+
+        let (output, dt) = compressor
+            .compress(true, DataTypeFlag::Compressed, input)
+            .unwrap();
+
+        assert!(matches!(output, Cow::Borrowed(_)));
+        assert_eq!(&*output, input.as_slice());
+        assert_eq!(dt, u8::from(DataTypeFlag::Compressed));
+    }
+
+    #[test]
+    fn input_below_min_size_returns_input_unchanged() {
+        let mut compressor = StdCompressor::new(&enabled_config(1024, 1.0));
+        let input = b"small";
+
+        let (output, dt) = compressor
+            .compress(true, DataTypeFlag::Json, input)
+            .unwrap();
+
+        assert!(matches!(output, Cow::Borrowed(_)));
+        assert_eq!(&*output, input.as_slice());
+        assert_eq!(dt, u8::from(DataTypeFlag::Json));
+    }
+
+    #[test]
+    fn compressible_input_returns_owned_with_compressed_flag() {
+        let mut compressor = StdCompressor::new(&enabled_config(0, 1.0));
+        // Highly compressible: repeated bytes.
+        let input = vec![b'a'; 4096];
+
+        let (output, dt) = compressor
+            .compress(true, DataTypeFlag::Json, &input)
+            .unwrap();
+
+        assert!(matches!(output, Cow::Owned(_)));
+        assert!(output.len() < input.len());
+        assert_eq!(
+            dt,
+            u8::from(DataTypeFlag::Json) | u8::from(DataTypeFlag::Compressed)
+        );
+
+        // Verify it round-trips through snappy.
+        let decompressed = snap::raw::Decoder::new().decompress_vec(&output).unwrap();
+        assert_eq!(decompressed, input);
+    }
+
+    #[test]
+    fn poor_ratio_returns_input_unchanged() {
+        // Set a very aggressive ratio that compressed output can't beat.
+        let mut compressor = StdCompressor::new(&enabled_config(0, 0.01));
+        let input = vec![b'a'; 256];
+
+        let (output, dt) = compressor
+            .compress(true, DataTypeFlag::Json, &input)
+            .unwrap();
+
+        assert!(matches!(output, Cow::Borrowed(_)));
+        assert_eq!(&*output, input.as_slice());
+        assert_eq!(dt, u8::from(DataTypeFlag::Json));
+    }
+
+    #[test]
+    fn compression_manager_creates_compressor() {
+        let manager = CompressionManager::<StdCompressor>::new(enabled_config(0, 1.0));
+        let mut compressor = manager.compressor();
+        let input = vec![b'x'; 4096];
+
+        let (output, _dt) = compressor
+            .compress(true, DataTypeFlag::None, &input)
+            .unwrap();
+
+        assert!(matches!(output, Cow::Owned(_)));
+        assert!(output.len() < input.len());
+    }
+}

sdk/couchbase-core/src/crudcomponent.rs

@@ -130,7 +130,7 @@ impl<
                         key: opts.key,
                         vbucket_id,
                         flags: opts.flags,
-                        value,
+                        value: &value,
                         datatype,
                         expiry: opts.expiry,
                         preserve_expiry: opts.preserve_expiry,
@@ -195,7 +195,7 @@ impl<
                         Error::new_contextual_memdx_error(e)
                     })
                     .map_ok(|resp| GetResult {
-                        value: resp.value.to_vec(),
+                        value: resp.value,
                         datatype: resp.datatype,
                         cas: resp.cas,
                         flags: resp.flags,
@@ -233,7 +233,7 @@ impl<
                         Error::new_contextual_memdx_error(e)
                     })
                     .map_ok(|resp| GetMetaResult {
-                        value: resp.value.to_vec(),
+                        value: resp.value,
                         datatype: resp.datatype,
                         server_duration: resp.server_duration,
                         expiry: resp.expiry,
@@ -323,7 +323,7 @@ impl<
                         Error::new_contextual_memdx_error(e)
                     })
                     .map_ok(|resp| GetAndLockResult {
-                        value: resp.value.to_vec(),
+                        value: resp.value,
                         datatype: resp.datatype,
                         cas: resp.cas,
                         flags: resp.flags,
@@ -362,7 +362,7 @@ impl<
                         Error::new_contextual_memdx_error(e)
                     })
                     .map_ok(|resp| GetAndTouchResult {
-                        value: resp.value.to_vec(),
+                        value: resp.value,
                         datatype: resp.datatype,
                         cas: resp.cas,
                         flags: resp.flags,
@@ -469,7 +469,7 @@ impl<
                         key: opts.key,
                         vbucket_id,
                         flags: opts.flags,
-                        value,
+                        value: &value,
                         datatype,
                         expiry: opts.expiry,
                         on_behalf_of: None,
@@ -531,7 +531,7 @@ impl<
                         key: opts.key,
                         vbucket_id,
                         flags: opts.flags,
-                        value,
+                        value: &value,
                         datatype,
                         expiry: opts.expiry,
                         preserve_expiry: opts.preserve_expiry,
@@ -594,7 +594,7 @@ impl<
                         collection_id,
                         key: opts.key,
                         vbucket_id,
-                        value,
+                        value: &value,
                         datatype,
                         cas: opts.cas,
                         on_behalf_of: None,
@@ -655,7 +655,7 @@ impl<
                         collection_id,
                         key: opts.key,
                         vbucket_id,
-                        value,
+                        value: &value,
                         datatype,
                         cas: opts.cas,
                         on_behalf_of: None,

sdk/couchbase-core/src/results/kv.rs

@@ -16,13 +16,15 @@
  *
  */
 
+use bytes::Bytes;
+
 use crate::error;
 use crate::mutationtoken::MutationToken;
 use std::time::Duration;
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct GetResult {
-    pub value: Vec<u8>,
+    pub value: Bytes,
     pub flags: u32,
     pub datatype: u8,
     pub cas: u64,
@@ -32,7 +34,7 @@ pub struct GetResult {
 pub struct GetMetaResult {
     pub cas: u64,
     pub flags: u32,
-    pub value: Vec<u8>,
+    pub value: Bytes,
     pub datatype: u8,
     pub server_duration: Option<Duration>,
     pub expiry: u32,
@@ -54,15 +56,15 @@ pub struct DeleteResult {
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct GetAndLockResult {
-    pub value: Vec<u8>,
+    pub value: Bytes,
     pub flags: u32,
     pub datatype: u8,
     pub cas: u64,
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct GetAndTouchResult {
-    pub value: Vec<u8>,
+    pub value: Bytes,
     pub flags: u32,
     pub datatype: u8,
     pub cas: u64,