minibytes: add a way to zero-copy serde deserialize from Bytes to Bytes

Summary:
Add `Bytes::as_scoped_deserialize_hint` API so deseralizing from antoher
`Bytes` could use zero-copy for byte slices.

Reviewed By: muirdm

Differential Revision: D60189689

fbshipit-source-id: 5e50bc4d211196e0d11603b485a064094f4f6eee

eden/scm/lib/minibytes/Cargo.toml

@@ -13,6 +13,7 @@ serde = { version = "1.0.185", features = ["derive", "rc"] }
 
 [dev-dependencies]
 quickcheck = "1.0"
+serde_cbor = "0.11"
 
 [features]
 default = ["frombytes", "frommmap", "non-zerocopy-into"]

eden/scm/lib/minibytes/TARGETS

@@ -31,6 +31,7 @@ rust_library(
     ],
     test_deps = [
         "fbsource//third-party/rust:quickcheck",
+        "fbsource//third-party/rust:serde_cbor",
     ],
     # A test inside this target is using #[should_panic], setting the backtrace
     # to false here, otherwise the test binary will try to extract the backtrace

eden/scm/lib/minibytes/src/serde.rs

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use std::cell::RefCell;
 use std::fmt;
 
 use serde::de;
@@ -24,6 +25,32 @@ impl Serialize for Bytes {
 
 struct BytesVisitor;
 
+thread_local! {
+    static DESERIALIZE_HINT: RefCell<Option<Bytes>> = const { RefCell::new(None) };
+}
+
+fn set_deserialize_hint(bytes: Option<Bytes>) -> Option<Bytes> {
+    DESERIALIZE_HINT.with_borrow_mut(|f| {
+        let orig = f.take();
+        *f = bytes;
+        orig
+    })
+}
+
+impl Bytes {
+    /// Call `func` with a "deserialize hint" as an attempt to avoid `memcpy`s.
+    /// `func` is usually a serde deserialize function taking `self` as input.
+    ///
+    /// Only affects the current thread, with the assumption that serde
+    /// deserialize is usually single threaded.
+    pub fn as_deserialize_hint<R>(&self, func: impl Fn() -> R) -> R {
+        let orig = set_deserialize_hint(Some(self.clone()));
+        let result = func();
+        set_deserialize_hint(orig);
+        result
+    }
+}
+
 impl<'de> de::Visitor<'de> for BytesVisitor {
     type Value = Bytes;
 
@@ -32,7 +59,7 @@ impl<'de> de::Visitor<'de> for BytesVisitor {
     }
 
     fn visit_borrowed_bytes<E: de::Error>(self, v: &'de [u8]) -> Result<Self::Value, E> {
-        Ok(Bytes::copy_from_slice(v))
+        self.visit_bytes(v)
     }
 
     fn visit_borrowed_str<E: de::Error>(self, v: &'de str) -> Result<Self::Value, E> {
@@ -44,7 +71,11 @@ impl<'de> de::Visitor<'de> for BytesVisitor {
     }
 
     fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E> {
-        Ok(Bytes::copy_from_slice(v))
+        let bytes = DESERIALIZE_HINT.with_borrow(|parent_buffer| match parent_buffer {
+            Some(buf) => buf.slice_to_bytes(v),
+            None => Bytes::copy_from_slice(v),
+        });
+        Ok(bytes)
     }
 }
 
@@ -53,3 +84,33 @@ impl<'de> Deserialize<'de> for Bytes {
         deserializer.deserialize_bytes(BytesVisitor)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[derive(Deserialize, Serialize)]
+    struct S {
+        a: Bytes,
+        b: Bytes,
+    }
+
+    #[test]
+    fn test_deserialize_hint() {
+        let s1 = S {
+            a: Bytes::copy_from_slice(b"aaaa"),
+            b: Bytes::from_static(b"bbbb"),
+        };
+        let serialized = Bytes::from(serde_cbor::to_vec(&s1).unwrap());
+
+        // Deserialize directly - no zero copy.
+        let s2: S = serde_cbor::from_slice(&serialized).unwrap();
+        assert!(serialized.range_of_slice(s2.a.as_ref()).is_none());
+        assert!(serialized.range_of_slice(s2.b.as_ref()).is_none());
+
+        // Deserialize with hint - can be zero copy.
+        let s3: S = serialized.as_deserialize_hint(|| serde_cbor::from_slice(&serialized).unwrap());
+        assert!(serialized.range_of_slice(s3.a.as_ref()).is_some());
+        assert!(serialized.range_of_slice(s3.b.as_ref()).is_some());
+    }
+}