Always make a copy in MemoryStore to reduce long-lasting allocs

fix: #289

cas/store/BUILD

@@ -372,12 +372,16 @@ rust_test(
 rust_test(
     name = "memory_store_test",
     srcs = ["tests/memory_store_test.rs"],
+    env = {
+        "RUST_TEST_THREADS": "1", # This test can't run in parallel.
+    },
     deps = [
         ":memory_store",
         ":traits",
         "//config",
         "//util:common",
         "//util:error",
+        "@crate_index//:bytes",
         "@crate_index//:pretty_assertions",
         "@crate_index//:tokio",
     ],

cas/store/memory_store.rs

@@ -87,22 +87,20 @@ impl StoreTrait for MemoryStore {
             UploadSizeInfo::ExactSize(sz) => sz,
             UploadSizeInfo::MaxSize(sz) => sz,
         };
-        let buffer = reader
-            .collect_all_with_size_hint(max_size)
-            .await
-            .err_tip(|| "Failed to collect all bytes from reader in memory_store::update")?;
 
-        // Resize our buffer if our max_size was not accurate.
-        // The buffer might have reserved much more than the amount of data transferred.
-        // This will ensure we use less memory for the long term stored data.
-        let buffer = if buffer.len() != max_size {
+        // Internally Bytes might hold a reference to more data than just our data. To prevent
+        // this potential case, we make a full copy of our data for long-term storage.
+        let final_buffer = {
+            let buffer = reader
+                .collect_all_with_size_hint(max_size)
+                .await
+                .err_tip(|| "Failed to collect all bytes from reader in memory_store::update")?;
             let mut new_buffer = BytesMut::with_capacity(buffer.len());
             new_buffer.extend_from_slice(&buffer[..]);
             new_buffer.freeze()
-        } else {
-            buffer
         };
-        self.evicting_map.insert(digest, BytesWrapper(buffer)).await;
+
+        self.evicting_map.insert(digest, BytesWrapper(final_buffer)).await;
         Ok(())
     }
 

cas/store/tests/memory_store_test.rs

@@ -12,21 +12,30 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use error::Error;
+use std::io::Read;
 use std::pin::Pin;
 
+use bytes::{BufMut, BytesMut};
+
+use common::DigestInfo;
+use error::ResultExt;
+use memory_store::MemoryStore;
+use traits::StoreTrait;
+
+const VALID_HASH1: &str = "0123456789abcdef000000000000000000010000000000000123456789abcdef";
+const VALID_HASH2: &str = "0123456789abcdef000000000000000000020000000000000123456789abcdef";
+const VALID_HASH3: &str = "0123456789abcdef000000000000000000030000000000000123456789abcdef";
+const VALID_HASH4: &str = "0123456789abcdef000000000000000000040000000000000123456789abcdef";
+const TOO_LONG_HASH: &str = "0123456789abcdef000000000000000000010000000000000123456789abcdefff";
+const TOO_SHORT_HASH: &str = "100000000000000000000000000000000000000000000000000000000000001";
+const INVALID_HASH: &str = "g111111111111111111111111111111111111111111111111111111111111111";
+
 #[cfg(test)]
 mod memory_store_tests {
     use super::*;
     use pretty_assertions::assert_eq; // Must be declared in every module.
 
-    use error::Error;
-
-    use common::DigestInfo;
-    use memory_store::MemoryStore;
-    use traits::StoreTrait;
-
-    const VALID_HASH1: &str = "0123456789abcdef000000000000000000010000000000000123456789abcdef";
-
     #[tokio::test]
     async fn insert_one_item_then_update() -> Result<(), Error> {
         let store_owned = MemoryStore::new(&config::stores::MemoryStore::default());
@@ -68,9 +77,56 @@ mod memory_store_tests {
         Ok(())
     }
 
-    const TOO_LONG_HASH: &str = "0123456789abcdef000000000000000000010000000000000123456789abcdefff";
-    const TOO_SHORT_HASH: &str = "100000000000000000000000000000000000000000000000000000000000001";
-    const INVALID_HASH: &str = "g111111111111111111111111111111111111111111111111111111111111111";
+    // Regression test for: https://github.com/TraceMachina/turbo-cache/issues/289.
+    #[cfg(target_family = "unix")]
+    #[tokio::test]
+    async fn ensure_full_copy_of_bytes_is_made_test() -> Result<(), Error> {
+        let store_owned = MemoryStore::new(&config::stores::MemoryStore::default());
+        let store = Pin::new(&store_owned);
+
+        fn get_self_memory_usage() -> usize {
+            let mut content = String::new();
+            let _ = std::fs::File::open("/proc/self/statm")
+                .unwrap()
+                .read_to_string(&mut content)
+                .unwrap();
+            content
+                .split_whitespace()
+                .map(|s| s.parse::<usize>().unwrap())
+                .next()
+                .unwrap()
+        }
+
+        let initial_virtual_mem = get_self_memory_usage();
+        for (i, hash) in [VALID_HASH1, VALID_HASH2, VALID_HASH3, VALID_HASH4]
+            .into_iter()
+            .enumerate()
+        {
+            // User a variety of sizes increasing up to 10MB each iteration.
+            // We do this to reduce the chance of memory page size masking the potential bug.
+            let reserved_size = 10_usize.pow(i as u32 + 4);
+            let mut mut_data = BytesMut::with_capacity(reserved_size);
+            mut_data.put_bytes(i as u8, 1);
+            let data = mut_data.freeze();
+
+            let digest = DigestInfo::try_new(hash, data.len())?;
+            store
+                .update_oneshot(digest, data)
+                .await
+                .err_tip(|| "Could not update store")?;
+        }
+
+        let new_virtual_mem = get_self_memory_usage();
+        let memory_usage_increase_perc = new_virtual_mem as f32 / initial_virtual_mem as f32;
+        // Arbitrary value, this may be increased if we find out that this is too low
+        // for some kernels/operating systems.
+        const MAXIMUM_MEMORY_USAGE_INCREASE_PERC: f32 = 1.1; // 10% increase.
+        assert!(
+            memory_usage_increase_perc < MAXIMUM_MEMORY_USAGE_INCREASE_PERC,
+            "Memory usage increased by {memory_usage_increase_perc} perc, which is more than {MAXIMUM_MEMORY_USAGE_INCREASE_PERC} perc",
+        );
+        Ok(())
+    }
 
     #[tokio::test]
     async fn read_partial() -> Result<(), Error> {

gencargo/memory_store_test/Cargo.toml

@@ -19,6 +19,7 @@ path = "../../cas/store/tests/memory_store_test.rs"
 doctest = false
 
 [dependencies]
+bytes = { workspace = true }
 pretty_assertions = { workspace = true }
 tokio = { workspace = true }