Merge branch 'abk/run-823-impose-max-chunk-store-size' into 'master'

RUN-823: Limit size of Wasm chunk store

Impose a limit on the total size of the Wasm chunk store for each
canister. This can be increased later if it's needed. 

See merge request dfinity-lab/public/ic!15752

rs/replicated_state/src/canister_state/system_state.rs

@@ -831,7 +831,7 @@ impl SystemState {
             initial_cycles,
             freeze_threshold,
             status,
-            WasmChunkStore::new_for_testing(),
+            WasmChunkStore::new_for_testing(wasm_chunk_store::DEFAULT_MAX_SIZE),
         )
     }
 

rs/replicated_state/src/canister_state/system_state/wasm_chunk_store.rs

@@ -8,6 +8,7 @@ use crate::{page_map::PageAllocatorFileDescriptor, PageMap};
 
 const PAGES_PER_CHUNK: u64 = 256;
 const CHUNK_SIZE: u64 = PAGES_PER_CHUNK * (PAGE_SIZE as u64);
+pub(crate) const DEFAULT_MAX_SIZE: NumBytes = NumBytes::new(100 * 1024 * 1024); // 100 MiB
 
 #[test]
 fn check_chunk_size() {
@@ -36,22 +37,26 @@ struct ChunkInfo {
 pub struct WasmChunkStore {
     data: PageMap,
     metadata: WasmChunkStoreMetadata,
+    /// Reject insertions that would grow the store beyond this size.
+    max_size: NumBytes,
 }
 
 impl WasmChunkStore {
     pub fn new(fd_factory: Arc<dyn PageAllocatorFileDescriptor>) -> Self {
         Self {
             data: PageMap::new(fd_factory),
             metadata: WasmChunkStoreMetadata::default(),
+            max_size: DEFAULT_MAX_SIZE,
         }
     }
 
     /// Creates a new `WasmChunkStore` that will use the temp file system for
     /// allocating new pages.
-    pub fn new_for_testing() -> Self {
+    pub fn new_for_testing(max_size: NumBytes) -> Self {
         Self {
             data: PageMap::new_for_testing(),
             metadata: WasmChunkStoreMetadata::default(),
+            max_size,
         }
     }
 
@@ -94,18 +99,31 @@ impl WasmChunkStore {
     /// succeed.
     pub fn can_insert_chunk(&self, chunk: &[u8]) -> Result<(), String> {
         if chunk.len() > CHUNK_SIZE as usize {
-            return Err(
-                format! {"Wasm chunk size {} exceeds the maximum chunk size of {}", chunk.len(), CHUNK_SIZE},
-            );
+            return Err(format!(
+                "Wasm chunk size {} exceeds the maximum chunk size of {}",
+                chunk.len(),
+                CHUNK_SIZE
+            ));
+        }
+        if self.metadata.chunks.len() as u64 * CHUNK_SIZE >= self.max_size.get() {
+            return Err(format!(
+                "Wasm chunk store has already reached maximum capacity of {} bytes",
+                self.max_size
+            ));
         }
         Ok(())
     }
 
     pub fn insert_chunk(&mut self, chunk: &[u8]) -> Result<[u8; 32], String> {
-        self.can_insert_chunk(chunk)?;
-
         let hash = ic_crypto_sha2::Sha256::hash(chunk);
 
+        // No changes needed if we already have the chunk
+        if self.metadata.chunks.contains_key(&hash) {
+            return Ok(hash);
+        }
+
+        self.can_insert_chunk(chunk)?;
+
         let index = self.metadata.chunks.len() as u64;
         let start_page = Self::page_index(index);
 
@@ -146,7 +164,11 @@ impl WasmChunkStore {
     }
 
     pub(crate) fn from_checkpoint(data: PageMap, metadata: WasmChunkStoreMetadata) -> Self {
-        Self { data, metadata }
+        Self {
+            data,
+            metadata,
+            max_size: DEFAULT_MAX_SIZE,
+        }
     }
 
     fn page_index(chunk_index: u64) -> PageIndex {
@@ -223,7 +245,7 @@ mod tests {
 
     #[test]
     fn store_and_retrieve_chunk() {
-        let mut store = WasmChunkStore::new_for_testing();
+        let mut store = WasmChunkStore::new_for_testing(DEFAULT_MAX_SIZE);
         let contents = [1, 2, 3].repeat(10_000);
         let hash = store.insert_chunk(&contents).unwrap();
         let round_trip_contents = get_chunk_as_vec(&store, hash);
@@ -232,7 +254,7 @@ mod tests {
 
     #[test]
     fn store_and_retrieve_empty_chunk() {
-        let mut store = WasmChunkStore::new_for_testing();
+        let mut store = WasmChunkStore::new_for_testing(DEFAULT_MAX_SIZE);
         let contents = vec![];
         let hash = store.insert_chunk(&contents).unwrap();
         let round_trip_contents = get_chunk_as_vec(&store, hash);
@@ -241,7 +263,7 @@ mod tests {
 
     #[test]
     fn error_when_chunk_exceeds_size_limit() {
-        let mut store = WasmChunkStore::new_for_testing();
+        let mut store = WasmChunkStore::new_for_testing(DEFAULT_MAX_SIZE);
         let contents = vec![0xab; chunk_size().get() as usize + 1];
         let result = store.insert_chunk(&contents);
         assert_eq!(
@@ -252,15 +274,15 @@ mod tests {
 
     #[test]
     fn can_insert_chunk_up_to_max_size() {
-        let mut store = WasmChunkStore::new_for_testing();
+        let mut store = WasmChunkStore::new_for_testing(DEFAULT_MAX_SIZE);
         let contents = vec![0xab; chunk_size().get() as usize];
         let result = store.insert_chunk(&contents);
         assert_matches!(result, Ok(_));
     }
 
     #[test]
     fn can_insert_and_retrieve_multiple_chunks() {
-        let mut store = WasmChunkStore::new_for_testing();
+        let mut store = WasmChunkStore::new_for_testing(DEFAULT_MAX_SIZE);
         let contents1 = vec![0xab; 1024];
         let hash1 = store.insert_chunk(&contents1).unwrap();
         let contents2 = vec![0x41; 1024];
@@ -273,6 +295,44 @@ mod tests {
         assert_eq!(contents2, round_trip_contents2);
     }
 
+    #[test]
+    fn cant_grow_beyond_max_size() {
+        let mut store = WasmChunkStore::new_for_testing(NumBytes::from(2 * 1024 * 1024));
+        let contents = vec![0xab; 1024];
+        let _hash = store.insert_chunk(&contents).unwrap();
+        let contents = vec![0xbc; 1024];
+        let _hash = store.insert_chunk(&contents).unwrap();
+        let contents = vec![0xcd; 1024];
+        store.insert_chunk(&contents).unwrap_err();
+    }
+
+    #[test]
+    fn inserting_same_chunk_doesnt_increase_size() {
+        // Store only has space for two chunks
+        let mut store = WasmChunkStore::new_for_testing(NumBytes::from(2 * 1024 * 1024));
+        let contents = vec![0xab; 1024];
+
+        // We can insert the same chunk many times because it doesn't take up
+        // new space in the store since it is already present.
+        let _hash = store.insert_chunk(&contents).unwrap();
+        let _hash = store.insert_chunk(&contents).unwrap();
+        let _hash = store.insert_chunk(&contents).unwrap();
+        let _hash = store.insert_chunk(&contents).unwrap();
+    }
+
+    #[test]
+    fn inserting_existing_chunk_succeeds_when_full() {
+        // Store only has space for two chunks
+        let mut store = WasmChunkStore::new_for_testing(NumBytes::from(2 * CHUNK_SIZE));
+
+        // We can insert the same chunk many times because it doesn't take up
+        // new space in the store since it is already present.
+        let _hash = store.insert_chunk(&[0xab; 10]).unwrap();
+        let _hash = store.insert_chunk(&[0xcd; 10]).unwrap();
+        // Store is now full, but inserting the same chunk again succeeds.
+        let _hash = store.insert_chunk(&[0xab; 10]).unwrap();
+    }
+
     mod proptest_tests {
         use super::*;
         use proptest::collection::vec as prop_vec;
@@ -282,8 +342,12 @@ mod tests {
 
         proptest! {
             #[test]
+            // Try chunks 2x as big as the size limit.
+            // If all inserts below the size limit succeeded, we'd expect 50 *
+            // .5 MiB = 25 MiB total. So set the max size below that to
+            // evenutally hit the size limit.
             fn insert_result_matches_can_insert(vecs in prop_vec((any::<u8>(), 0..2 * MB), 100)) {
-                let mut store = WasmChunkStore::new_for_testing();
+                let mut store = WasmChunkStore::new_for_testing(NumBytes::from(20 * MB as u64));
                 for (byte, length) in vecs {
                     let chunk = vec![byte; length];
                     let check = store.can_insert_chunk(&chunk);