diff --git a/README.md b/README.md index e14e461a..14686d22 100644 --- a/README.md +++ b/README.md @@ -40,16 +40,16 @@ Stable structures are able to work directly in stable memory because each data s its own memory. When initializing a stable structure, a memory is provided that the data structure can use to store its data. -### Basic Usage +Here are some basic examples: -Here's a basic example: +### Example: BTreeMap ```rust use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; -let mut map: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); +let mut map: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); -map.insert(1, "hello".to_string()); -assert_eq!(map.get(&1), Some("hello".to_string())); +map.insert(1, 2); +assert_eq!(map.get(&1), Some(2)); ``` Memories are abstracted with the [Memory] trait, and stable structures can work with any storage @@ -58,30 +58,39 @@ This includes stable memory, a vector ([VectorMemory]), or even a flat file ([Fi The example above initializes a [BTreeMap] with a [DefaultMemoryImpl], which maps to stable memory when used in a canister and to a [VectorMemory] otherwise. -### Memory Isolation Requirement +### Example: BTreeSet -> **⚠️ CRITICAL:** Stable structures **MUST NOT** share memories! -> Each memory must belong to only one stable structure. +The `BTreeSet` is a stable set implementation based on a B-Tree. It allows efficient insertion, deletion, and lookup of unique elements. +```rust +use ic_stable_structures::{BTreeSet, DefaultMemoryImpl}; +let mut set: BTreeSet = BTreeSet::new(DefaultMemoryImpl::default()); + +set.insert(42); +assert!(set.contains(&42)); +assert_eq!(set.pop_first(), Some(42)); +assert!(set.is_empty()); +``` + + +Note that **stable structures cannot share memories.** +Each memory must belong to only one stable structure. For example, this fails when run in a canister: -```rust,ignore +```no_run use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; -let mut map_a: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); -let mut map_b: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); +let mut map_1: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); +let mut map_2: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); -map_a.insert(1, b'A'); -map_b.insert(1, b'B'); -assert_eq!(map_a.get(&1), Some(b'A')); // ❌ FAILS: Returns b'B' due to shared memory! -assert_eq!(map_b.get(&1), Some(b'B')); // ✅ Succeeds, but corrupted map_a +map_1.insert(1, 2); +map_2.insert(1, 3); +assert_eq!(map_1.get(&1), Some(2)); // This assertion fails. ``` -It fails because both `map_a` and `map_b` are using the same stable memory under the hood, and so changes in `map_b` end up changing or corrupting `map_a`. +It fails because both `map_1` and `map_2` are using the same stable memory under the hood, and so changes in `map_1` end up changing or corrupting `map_2`. -### Using MemoryManager - -To address this issue, we use the [MemoryManager](memory_manager::MemoryManager), which takes a single memory and creates up to 255 virtual memories for our use. -Here's the above failing example, but fixed: +To address this issue, we make use of the [MemoryManager](memory_manager::MemoryManager), which takes a single memory and creates up to 255 virtual memories for our disposal. +Here's the above failing example, but fixed by using the [MemoryManager](memory_manager::MemoryManager): ```rust use ic_stable_structures::{ @@ -89,77 +98,12 @@ use ic_stable_structures::{ BTreeMap, DefaultMemoryImpl, }; let mem_mgr = MemoryManager::init(DefaultMemoryImpl::default()); -let (mem_id_a, mem_id_b) = (MemoryId::new(0), MemoryId::new(1)); -let mut map_a: BTreeMap = BTreeMap::init(mem_mgr.get(mem_id_a)); -let mut map_b: BTreeMap = BTreeMap::init(mem_mgr.get(mem_id_b)); - -map_a.insert(1, b'A'); -map_b.insert(1, b'B'); -assert_eq!(map_a.get(&1), Some(b'A')); // ✅ Succeeds: Each map has its own memory -assert_eq!(map_b.get(&1), Some(b'B')); // ✅ Succeeds: No data corruption -``` - -### Memory Reclamation - -During migrations you often create a new structure (B) and copy data from an existing one (A). Without reclamation, this can double memory usage even after A is no longer needed. - -Bucket IDs are an internal implementation detail — hidden and not user-controllable — and each virtual memory must receive bucket IDs in strictly ascending order. Because of this, reuse of freed buckets is guaranteed when allocating into a newly created (empty) structure. For existing structures, reuse may or may not work: it succeeds only if there is a free bucket with an ID greater than the structure’s current maximum; otherwise a new bucket is allocated. - -Example: A = `[0, 4, 5]`, B = `[1, 2, 3]`. After releasing A, `free = [0, 4, 5]`. When B grows, it can’t take `0` (must be `> 3`) but can take `4` → `B = [1, 2, 3, 4]`, `free = [0, 5]`. - -**Recommendation:** for predictable reuse migrate into a newly created structure rather than relying on reuse with a populated one. - -> **⚠️ CRITICAL SAFETY REQUIREMENT:** -> - **MUST** drop the original structure object before calling `reclaim_memory`. -> - **NEVER** use the original structure after reclamation — doing so corrupts data. - -The `MemoryManager` provides a `reclaim_memory` method to efficiently handle these scenarios: - -```rust -use ic_stable_structures::{ - memory_manager::{MemoryId, MemoryManager}, - BTreeMap, DefaultMemoryImpl, Memory, -}; +let mut map_1: BTreeMap = BTreeMap::init(mem_mgr.get(MemoryId::new(0))); +let mut map_2: BTreeMap = BTreeMap::init(mem_mgr.get(MemoryId::new(1))); -let mem = DefaultMemoryImpl::default(); -let mem_mgr = MemoryManager::init(mem.clone()); -let (mem_id_a, mem_id_b) = (MemoryId::new(0), MemoryId::new(1)); - -// ======================================== -// Scenario 1: WITHOUT reclamation -// ======================================== -let mut map_a: BTreeMap = BTreeMap::init(mem_mgr.get(mem_id_a)); -map_a.insert(1, b'A'); // Populate map A with data -let data = map_a.get(&1); // Extract data for migration -map_a.clear_new(); // A is now empty -drop(map_a); // Memory stays allocated to mem_id_a -let actual_size_before_migration = mem.size(); - -let mut map_b: BTreeMap = BTreeMap::new(mem_mgr.get(mem_id_b)); -map_b.insert(1, data.unwrap()); // B allocates NEW memory -let actual_size_after_migration = mem.size(); - // Result: ~2x memory usage - // Memory allocation grew (waste) -assert!(actual_size_before_migration < actual_size_after_migration); - -// ======================================== -// Scenario 2: WITH reclamation -// ======================================== -let mut map_a: BTreeMap = BTreeMap::init(mem_mgr.get(mem_id_a)); -map_a.insert(1, b'A'); // Populate map A with data -let data = map_a.get(&1); // Extract data for migration -map_a.clear_new(); // A is now empty -drop(map_a); // Drop A completely -let actual_size_before_migration = mem.size(); -mem_mgr.reclaim_memory(mem_id_a); // Free A's memory buckets for reuse - -// Reusing free memory buckets works best on newly created structures -let mut map_b: BTreeMap = BTreeMap::new(mem_mgr.get(mem_id_b)); -map_b.insert(1, data.unwrap()); // B reuses A's reclaimed memory buckets -let actual_size_after_migration = mem.size(); - // Result: 1x memory usage - // Memory allocation stayed the same (no waste) -assert!(actual_size_before_migration == actual_size_after_migration); +map_1.insert(1, 2); +map_2.insert(1, 3); +assert_eq!(map_1.get(&1), Some(2)); // Succeeds, as expected. ``` ## Example Canister @@ -191,7 +135,7 @@ thread_local! { RefCell::new(MemoryManager::init(DefaultMemoryImpl::default())); // Initialize a `StableBTreeMap` with `MemoryId(0)`. - static MAP: RefCell> = RefCell::new( + static MAP: RefCell> = RefCell::new( StableBTreeMap::init( MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(0))), ) @@ -200,32 +144,32 @@ thread_local! { // Retrieves the value associated with the given key if it exists. #[ic_cdk_macros::query] -fn get(key: u64) -> Option { +fn get(key: u128) -> Option { MAP.with(|p| p.borrow().get(&key)) } // Inserts an entry into the map and returns the previous value of the key if it exists. #[ic_cdk_macros::update] -fn insert(key: u64, value: String) -> Option { +fn insert(key: u128, value: u128) -> Option { MAP.with(|p| p.borrow_mut().insert(key, value)) } ``` ### More Examples -- [Basic Example](https://github.com/dfinity/stable-structures/tree/main/examples/src/basic_example): Simple usage patterns +- [Basic Example](https://github.com/dfinity/stable-structures/tree/main/examples/src/basic_example) (the one above) - [Quickstart Example](https://github.com/dfinity/stable-structures/tree/main/examples/src/quick_start): Ideal as a template when developing a new canister - [Custom Types Example](https://github.com/dfinity/stable-structures/tree/main/examples/src/custom_types_example): Showcases storing your own custom types ## Combined Persistence -If your project uses only stable structures, memory can expand in size without requiring `pre_upgrade`/`post_upgrade` hooks. +If your project exclusively relies on stable structures, the memory can expand in size without the requirement of `pre_upgrade`/`post_upgrade` hooks. -However, if you also need to serialize/deserialize heap data, you must use the memory manager to avoid conflicts. To combine both approaches effectively, refer to the [Quickstart Example](https://github.com/dfinity/stable-structures/tree/main/examples/src/quick_start) for guidance. +However, it's important to note that if you also intend to perform serialization/deserialization of the heap data, utilizing the memory manager becomes necessary. To effectively combine both approaches, refer to the [Quickstart Example](https://github.com/dfinity/stable-structures/tree/main/examples/src/quick_start) for guidance. ## Fuzzing -Stable structures require strong guarantees to work reliably and scale over millions of operations. To that extent, we use fuzzing to emulate such operations on the available data structures. +Stable structures requires strong guarantees to work reliably and scale over millions of operations. To that extent, we use fuzzing to emulate such operations on the available data structures. To run a fuzzer locally, ```sh diff --git a/benchmarks/btreemap/canbench_results.yml b/benchmarks/btreemap/canbench_results.yml index 95cfd69b..2dabed3e 100644 --- a/benchmarks/btreemap/canbench_results.yml +++ b/benchmarks/btreemap/canbench_results.yml @@ -975,35 +975,35 @@ benches: btreemap_v2_mem_manager_insert_blob512_u64: total: calls: 1 - instructions: 3127680632 + instructions: 3127680452 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_v2_mem_manager_insert_u64_blob512: total: calls: 1 - instructions: 607288554 + instructions: 607288370 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_v2_mem_manager_insert_u64_u64: total: calls: 1 - instructions: 520545876 + instructions: 520545864 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_v2_mem_manager_insert_u64_vec512: total: calls: 1 - instructions: 834100785 + instructions: 834100595 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_v2_mem_manager_insert_vec512_u64: total: calls: 1 - instructions: 1964405636 + instructions: 1964405448 heap_increase: 0 stable_memory_increase: 0 scopes: {} diff --git a/benchmarks/io_chunks/canbench_results.yml b/benchmarks/io_chunks/canbench_results.yml index f9166ffd..6c2e9160 100644 --- a/benchmarks/io_chunks/canbench_results.yml +++ b/benchmarks/io_chunks/canbench_results.yml @@ -2,126 +2,126 @@ benches: read_chunks_btreemap_1: total: calls: 1 - instructions: 148723635 + instructions: 148723585 heap_increase: 1601 stable_memory_increase: 0 scopes: {} read_chunks_btreemap_1k: total: calls: 1 - instructions: 498267990 + instructions: 498228206 heap_increase: 0 stable_memory_increase: 0 scopes: {} read_chunks_btreemap_1m: total: calls: 1 - instructions: 40947358379 + instructions: 40940569325 heap_increase: 0 stable_memory_increase: 0 scopes: {} read_chunks_stable_1: total: calls: 1 - instructions: 104859145 + instructions: 104859143 heap_increase: 0 stable_memory_increase: 0 scopes: {} read_chunks_stable_1k: total: calls: 1 - instructions: 104985765 + instructions: 104985739 heap_increase: 0 stable_memory_increase: 0 scopes: {} read_chunks_stable_1m: total: calls: 1 - instructions: 230002765 + instructions: 230002739 heap_increase: 0 stable_memory_increase: 0 scopes: {} read_chunks_vec_1: total: calls: 1 - instructions: 104859762 + instructions: 104859760 heap_increase: 0 stable_memory_increase: 0 scopes: {} read_chunks_vec_1k: total: calls: 1 - instructions: 105830202 + instructions: 105826498 heap_increase: 0 stable_memory_increase: 0 scopes: {} read_chunks_vec_1m: total: calls: 1 - instructions: 1014586404 + instructions: 1010905944 heap_increase: 0 stable_memory_increase: 0 scopes: {} write_chunks_btreemap_1: total: calls: 1 - instructions: 357208825 + instructions: 357205397 heap_increase: 13 stable_memory_increase: 1536 scopes: {} write_chunks_btreemap_1k: total: calls: 1 - instructions: 4187350965 + instructions: 4187119879 heap_increase: 2 stable_memory_increase: 1536 scopes: {} write_chunks_btreemap_1m: total: calls: 1 - instructions: 83670464159 + instructions: 83659829857 heap_increase: 0 stable_memory_increase: 3072 scopes: {} write_chunks_stable_1: total: calls: 1 - instructions: 130472086 + instructions: 130471968 heap_increase: 0 stable_memory_increase: 1664 scopes: {} write_chunks_stable_1k: total: calls: 1 - instructions: 130598863 + instructions: 130598745 heap_increase: 0 stable_memory_increase: 1664 scopes: {} write_chunks_stable_1m: total: calls: 1 - instructions: 255406776 + instructions: 255406658 heap_increase: 0 stable_memory_increase: 1664 scopes: {} write_chunks_vec_1: total: calls: 1 - instructions: 549903682 + instructions: 549903573 heap_increase: 0 stable_memory_increase: 1536 scopes: {} write_chunks_vec_1k: total: calls: 1 - instructions: 562259611 + instructions: 562257515 heap_increase: 0 stable_memory_increase: 1536 scopes: {} write_chunks_vec_1m: total: calls: 1 - instructions: 1896596401 + instructions: 1896593101 heap_increase: 0 stable_memory_increase: 1536 scopes: {} diff --git a/benchmarks/memory_manager/canbench_results.yml b/benchmarks/memory_manager/canbench_results.yml index bd8ad406..bf482bf0 100644 --- a/benchmarks/memory_manager/canbench_results.yml +++ b/benchmarks/memory_manager/canbench_results.yml @@ -9,14 +9,14 @@ benches: memory_manager_grow: total: calls: 1 - instructions: 348618009 + instructions: 347433966 heap_increase: 2 stable_memory_increase: 32000 scopes: {} memory_manager_overhead: total: calls: 1 - instructions: 1181984472 + instructions: 1181977502 heap_increase: 0 stable_memory_increase: 8320 scopes: {} diff --git a/docs/src/concepts/memory-manager.md b/docs/src/concepts/memory-manager.md index 4fa27edd..888fd317 100644 --- a/docs/src/concepts/memory-manager.md +++ b/docs/src/concepts/memory-manager.md @@ -22,83 +22,16 @@ use ic_stable_structures::{ let mem_mgr = MemoryManager::init(DefaultMemoryImpl::default()); // Create two separate BTreeMaps, each with its own virtual memory -let (mem_id_a, mem_id_b) = (MemoryId::new(0), MemoryId::new(1)); -let mut map_a: BTreeMap = BTreeMap::init(mem_mgr.get(mem_id_a)); -let mut map_b: BTreeMap = BTreeMap::init(mem_mgr.get(mem_id_b)); +let mut map_1: BTreeMap = BTreeMap::init(mem_mgr.get(MemoryId::new(0))); +let mut map_2: BTreeMap = BTreeMap::init(mem_mgr.get(MemoryId::new(1))); // Demonstrate independent operation of the two maps -map_a.insert(1, b'A'); -map_b.insert(1, b'B'); -assert_eq!(map_a.get(&1), Some(b'A')); // ✅ Succeeds: Each map has its own memory -assert_eq!(map_b.get(&1), Some(b'B')); // ✅ Succeeds: No data corruption +map_1.insert(1, 2); +map_2.insert(1, 3); +assert_eq!(map_1.get(&1), Some(2)); // Succeeds as expected ``` ```admonish warning "" -**⚠️ CRITICAL:** Stable structures **MUST NOT** share memories! -Each memory instance must be assigned to exactly one stable structure. -``` - -## Memory Reclamation - -During migrations you often create a new structure (B) and copy data from an existing one (A). Without reclamation, this can double memory usage even after A is no longer needed. - -Bucket IDs are an internal implementation detail — hidden and not user-controllable — and each virtual memory must receive bucket IDs in strictly ascending order. Because of this, reuse of freed buckets is guaranteed when allocating into a newly created (empty) structure. For existing structures, reuse may or may not work: it succeeds only if there is a free bucket with an ID greater than the structure’s current maximum; otherwise a new bucket is allocated. - -Example: A = `[0, 4, 5]`, B = `[1, 2, 3]`. After releasing A, `free = [0, 4, 5]`. When B grows, it can’t take `0` (must be `> 3`) but can take `4` → `B = [1, 2, 3, 4]`, `free = [0, 5]`. - -**Recommendation:** for predictable reuse migrate into a newly created structure rather than relying on reuse with a populated one. - -```admonish info "" -**⚠️ CRITICAL SAFETY REQUIREMENT:** -- **MUST** drop the original structure object before calling `reclaim_memory`. -- **NEVER** use the original structure after reclamation — doing so corrupts data. -``` - -The `MemoryManager` provides a `reclaim_memory` method to efficiently handle these scenarios: - -```rust -use ic_stable_structures::{ - memory_manager::{MemoryId, MemoryManager}, - BTreeMap, DefaultMemoryImpl, Memory, -}; - -let mem = DefaultMemoryImpl::default(); -let mem_mgr = MemoryManager::init(mem.clone()); -let (mem_id_a, mem_id_b) = (MemoryId::new(0), MemoryId::new(1)); - -// ======================================== -// Scenario 1: WITHOUT reclamation -// ======================================== -let mut map_a: BTreeMap = BTreeMap::init(mem_mgr.get(mem_id_a)); -map_a.insert(1, b'A'); // Populate map A with data -let data = map_a.get(&1); // Extract data for migration -map_a.clear_new(); // A is now empty -drop(map_a); // Memory stays allocated to mem_id_a -let actual_size_before_migration = mem.size(); - -let mut map_b: BTreeMap = BTreeMap::new(mem_mgr.get(mem_id_b)); -map_b.insert(1, data.unwrap()); // B allocates NEW memory -let actual_size_after_migration = mem.size(); - // Result: ~2x memory usage - // Memory allocation grew (waste) -assert!(actual_size_before_migration < actual_size_after_migration); - -// ======================================== -// Scenario 2: WITH reclamation -// ======================================== -let mut map_a: BTreeMap = BTreeMap::init(mem_mgr.get(mem_id_a)); -map_a.insert(1, b'A'); // Populate map A with data -let data = map_a.get(&1); // Extract data for migration -map_a.clear_new(); // A is now empty -drop(map_a); // Drop A completely -let actual_size_before_migration = mem.size(); -mem_mgr.reclaim_memory(mem_id_a); // Free A's memory buckets for reuse - -// Reusing free memory buckets works best on newly created structures -let mut map_b: BTreeMap = BTreeMap::new(mem_mgr.get(mem_id_b)); -map_b.insert(1, data.unwrap()); // B reuses A's reclaimed memory buckets -let actual_size_after_migration = mem.size(); - // Result: 1x memory usage - // Memory allocation stayed the same (no waste) -assert!(actual_size_before_migration == actual_size_after_migration); +Virtual memories from the `MemoryManager` cannot be shared between stable structures. +Each memory instance should be assigned to exactly one stable structure. ``` diff --git a/docs/src/concepts/memory-trait.md b/docs/src/concepts/memory-trait.md index 1eedef6a..588e30b1 100644 --- a/docs/src/concepts/memory-trait.md +++ b/docs/src/concepts/memory-trait.md @@ -55,15 +55,12 @@ Here's how to initialize a stable `BTreeMap` using `DefaultMemoryImpl`: ```rust use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; -let mut map: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); - -map.insert(1, "hello".to_string()); -assert_eq!(map.get(&1), Some("hello".to_string())); +let mut map: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); ``` ```admonish warning "" -**⚠️ CRITICAL:** Stable structures **MUST NOT** share memories! -Each memory must belong to only one stable structure. +**Important**: Stable structures cannot share memories. +Each memory must be dedicated to a single stable structure. ``` While the above example works correctly, it demonstrates a potential issue: the `BTreeMap` will use the entire stable memory. @@ -72,16 +69,15 @@ For example, the following code will fail in a canister: ```rust use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; -let mut map_a: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); -let mut map_b: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); +let mut map_1: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); +let mut map_2: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); -map_a.insert(1, b'A'); -map_b.insert(1, b'B'); -assert_eq!(map_a.get(&1), Some(b'A')); // ❌ FAILS: Returns b'B' due to shared memory! -assert_eq!(map_b.get(&1), Some(b'B')); // ✅ Succeeds, but corrupted map_a +map_1.insert(1, 2); +map_2.insert(1, 3); +assert_eq!(map_1.get(&1), Some(2)); // This assertion fails. ``` -The code fails because both `map_a` and `map_b` are using the same stable memory. +The code fails because both `map_1` and `map_2` are using the same stable memory. This causes changes in one map to affect or corrupt the other. To solve this problem, the library provides the [MemoryManager](./memory-manager.md), which creates up to 255 virtual memories from a single memory instance. diff --git a/src/btreemap.rs b/src/btreemap.rs index 66a56035..c6f4ae92 100644 --- a/src/btreemap.rs +++ b/src/btreemap.rs @@ -112,8 +112,8 @@ const PAGE_SIZE_VALUE_MARKER: u32 = u32::MAX; /// /// ## Multiple BTreeMaps and Memory Management /// -/// > **⚠️ CRITICAL:** Stable structures **MUST NOT** share memories! -/// > Each memory must belong to only one stable structure. +/// **Important**: Each stable structure requires its own designated memory region. Attempting to +/// initialize multiple structures with the same memory will lead to data corruption. /// /// ### What NOT to do: /// @@ -121,14 +121,13 @@ const PAGE_SIZE_VALUE_MARKER: u32 = u32::MAX; /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; /// /// // ERROR: Using the same memory for multiple BTreeMaps will corrupt data -/// let mut map_a: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); -/// let mut map_b: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); +/// let mut map_1: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); +/// let mut map_2: BTreeMap = BTreeMap::init(DefaultMemoryImpl::default()); /// -/// map_a.insert(1, b'A'); -/// map_b.insert(1, b'B'); -/// // This assertion would fail: changes to map_b corrupt map_a's data -/// assert_eq!(map_a.get(&1), Some(b'A')); // ❌ FAILS: Returns b'B' due to shared memory! -/// assert_eq!(map_b.get(&1), Some(b'B')); // ✅ Succeeds, but corrupted map_a +/// map_1.insert(1, "two".to_string()); +/// map_2.insert(1, "three".to_string()); +/// // This assertion would fail: changes to map_2 corrupt map_1's data +/// assert_eq!(map_1.get(&1), Some("two".to_string())); /// ``` /// /// ### Correct approach using MemoryManager: @@ -138,15 +137,18 @@ const PAGE_SIZE_VALUE_MARKER: u32 = u32::MAX; /// memory_manager::{MemoryId, MemoryManager}, /// BTreeMap, DefaultMemoryImpl, /// }; -/// let mem_mgr = MemoryManager::init(DefaultMemoryImpl::default()); -/// let (mem_id_a, mem_id_b) = (MemoryId::new(0), MemoryId::new(1)); -/// let mut map_a: BTreeMap = BTreeMap::init(mem_mgr.get(mem_id_a)); -/// let mut map_b: BTreeMap = BTreeMap::init(mem_mgr.get(mem_id_b)); /// -/// map_a.insert(1, b'A'); -/// map_b.insert(1, b'B'); -/// assert_eq!(map_a.get(&1), Some(b'A')); // ✅ Succeeds: Each map has its own memory -/// assert_eq!(map_b.get(&1), Some(b'B')); // ✅ Succeeds: No data corruption +/// // Initialize the memory manager with a single memory +/// let memory_manager = MemoryManager::init(DefaultMemoryImpl::default()); +/// +/// // Get separate virtual memories for each BTreeMap +/// let mut map_1: BTreeMap = BTreeMap::init(memory_manager.get(MemoryId::new(0))); +/// let mut map_2: BTreeMap = BTreeMap::init(memory_manager.get(MemoryId::new(1))); +/// +/// map_1.insert(1, "two".to_string()); +/// map_2.insert(1, "three".to_string()); +/// // Now this works as expected +/// assert_eq!(map_1.get(&1), Some("two".to_string())); /// ``` /// /// The [`MemoryManager`](crate::memory_manager::MemoryManager) creates up to 255 virtual memories @@ -707,15 +709,6 @@ where /// Removes all elements from the map. // TODO: In next major release (v1.0), rename this method to `clear` to follow // standard Rust collection naming conventions. - /// - /// # Safety Note for Memory Reclamation - /// If using manual memory reclamation via `MemoryManager::reclaim_memory()`: - /// 1. **MANDATORY**: Drop this BTreeMap object first (let it go out of scope) - /// 2. Call `reclaim_memory()` on the memory manager - /// 3. Create new structures as needed - /// - /// Using this BTreeMap after memory reclamation causes data corruption. - /// Note: You can still call `clear_new()` if you need to clear data without memory reclamation. pub fn clear_new(&mut self) { self.root_addr = NULL; self.length = 0; diff --git a/src/btreeset.rs b/src/btreeset.rs index 781c5531..d4853b15 100644 --- a/src/btreeset.rs +++ b/src/btreeset.rs @@ -352,15 +352,6 @@ where /// set.clear(); /// assert!(set.is_empty()); /// ``` - /// - /// # Safety Note for Memory Reclamation - /// If using manual memory reclamation via `MemoryManager::reclaim_memory()`: - /// 1. **MANDATORY**: Drop this BTreeSet object first (let it go out of scope) - /// 2. Call `reclaim_memory()` on the memory manager - /// 3. Create new structures as needed - /// - /// Using this BTreeSet after memory reclamation causes data corruption. - /// Note: You can still call `clear()` if you need to clear data without memory reclamation. pub fn clear(&mut self) { self.map.clear_new(); } diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 8d82d52e..463e4e0a 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -22,22 +22,22 @@ //! let mem_mgr = MemoryManager::init(DefaultMemoryImpl::default()); //! //! // Create different memories, each with a unique ID. -//! let memory_a = mem_mgr.get(MemoryId::new(0)); -//! let memory_b = mem_mgr.get(MemoryId::new(1)); +//! let memory_0 = mem_mgr.get(MemoryId::new(0)); +//! let memory_1 = mem_mgr.get(MemoryId::new(1)); //! //! // Each memory can be used independently. -//! memory_a.grow(1); -//! memory_a.write(0, &[1, 2, 3]); +//! memory_0.grow(1); +//! memory_0.write(0, &[1, 2, 3]); //! -//! memory_b.grow(1); -//! memory_b.write(0, &[4, 5, 6]); +//! memory_1.grow(1); +//! memory_1.write(0, &[4, 5, 6]); //! //! let mut bytes = vec![0; 3]; -//! memory_a.read(0, &mut bytes); +//! memory_0.read(0, &mut bytes); //! assert_eq!(bytes, vec![1, 2, 3]); //! //! let mut bytes = vec![0; 3]; -//! memory_b.read(0, &mut bytes); +//! memory_1.read(0, &mut bytes); //! assert_eq!(bytes, vec![4, 5, 6]); //! ``` use crate::{ @@ -46,8 +46,6 @@ use crate::{ write, write_struct, Memory, WASM_PAGE_SIZE, }; use std::cell::{Cell, RefCell}; -use std::collections::BTreeSet; -use std::ops::Bound::{Excluded, Unbounded}; use std::rc::Rc; const MAGIC: &[u8; 3] = b"MGR"; @@ -129,13 +127,12 @@ const HEADER_RESERVED_BYTES: usize = 32; /// Bucket MAX_NUM_BUCKETS ↕ N pages /// ``` /// -/// # Current Limitations & Safety Requirements +/// # Current Limitations and Future Improvements /// -/// - **Manual memory reclamation** - call `reclaim_memory()` after dropping structures -/// - **Structure invalidation** - original structures become invalid after memory reclamation -/// - **Mandatory drop** - you MUST drop original structures BEFORE reclaiming memory -/// - **No safety verification** - user discipline required to prevent data corruption -/// - **Incorrect usage** - using structures after memory reclamation causes data corruption +/// - Buckets are never deallocated once assigned to a memory ID, even when the memory becomes empty +/// - Clearing data structures (BTreeMap, Vec) does not reclaim underlying bucket allocations +/// - Long-running canisters may accumulate unused buckets, increasing storage costs +/// - Future versions may consider automatic bucket reclamation and memory compaction features pub struct MemoryManager { inner: Rc>>, } @@ -165,24 +162,6 @@ impl MemoryManager { } } - /// Reclaims memory allocated to the specified virtual memory for reuse. - /// - /// **CRITICAL SAFETY REQUIREMENT**: Drop the original structure object first! - /// - /// **Usage Pattern:** - /// ```rust,ignore - /// drop(map); // 1. Drop the structure object first - /// let pages = mem_mgr.reclaim_memory(memory_id); // 2. Reclaim memory - /// let new_map = BTreeMap::new(mem_mgr.get(memory_id)); // 3. Create new structure - /// ``` - /// - /// **DANGER**: Using the original structure after reclamation causes data corruption. - /// - /// Returns the number of pages reclaimed (0 if none allocated). - pub fn reclaim_memory(&self, id: MemoryId) -> u64 { - self.inner.borrow_mut().reclaim_memory(id) - } - /// Returns the underlying memory. /// /// # Returns @@ -267,11 +246,6 @@ struct MemoryManagerInner { /// A map mapping each managed memory to the bucket ids that are allocated to it. memory_buckets: Vec>, - - /// A pool of free buckets that can be reused by any memory. - /// Maintained in increasing order (lowest bucket IDs at the front) - /// to ensure optimal memory locality when reusing buckets. - free_buckets: BTreeSet, } impl MemoryManagerInner { @@ -300,7 +274,6 @@ impl MemoryManagerInner { memory_sizes_in_pages: [0; MAX_NUM_MEMORIES as usize], memory_buckets: vec![vec![]; MAX_NUM_MEMORIES as usize], bucket_size_in_pages, - free_buckets: BTreeSet::new(), }; mem_mgr.save_header(); @@ -330,15 +303,9 @@ impl MemoryManagerInner { ); let mut memory_buckets = vec![vec![]; MAX_NUM_MEMORIES as usize]; - let mut free_buckets = BTreeSet::new(); - for (bucket_idx, memory_id) in buckets.into_iter().enumerate() { - let bucket_id = BucketId(bucket_idx as u16); - if memory_id != UNALLOCATED_BUCKET_MARKER { - memory_buckets[memory_id as usize].push(bucket_id); - } else if bucket_id.0 < header.num_allocated_buckets { - // Only buckets with indices less than `num_allocated_buckets` were ever allocated. - // If such a bucket is now marked as unallocated, it is free to reuse. - free_buckets.insert(bucket_id); + for (bucket_idx, memory) in buckets.into_iter().enumerate() { + if memory != UNALLOCATED_BUCKET_MARKER { + memory_buckets[memory as usize].push(BucketId(bucket_idx as u16)); } } @@ -348,7 +315,6 @@ impl MemoryManagerInner { bucket_size_in_pages: header.bucket_size_in_pages, memory_sizes_in_pages: header.memory_sizes_in_pages, memory_buckets, - free_buckets, } } @@ -379,11 +345,7 @@ impl MemoryManagerInner { let required_buckets = self.num_buckets_needed(new_size); let new_buckets_needed = required_buckets - current_buckets; - // Check if we have enough buckets available (either already allocated or can allocate new ones) - let free_bucket_count = self.free_buckets.len() as u64; - let new_buckets_to_allocate = new_buckets_needed.saturating_sub(free_bucket_count); - - if self.allocated_buckets as u64 + new_buckets_to_allocate > MAX_NUM_BUCKETS { + if new_buckets_needed + self.allocated_buckets as u64 > MAX_NUM_BUCKETS { // Exceeded the memory that can be managed. return -1; } @@ -392,44 +354,17 @@ impl MemoryManagerInner { // Allocate new buckets as needed. memory_bucket.reserve(new_buckets_needed as usize); for _ in 0..new_buckets_needed { - // Try to reuse a free bucket. - let maybe_free_bucket = match memory_bucket.last() { - // Fresh memory — smallest free bucket. - None => self.free_buckets.pop_first(), - - // Existing memory — smallest free bucket > current max. - Some(max_b) => { - if let Some(candidate) = self - .free_buckets - .range((Excluded(max_b), Unbounded)) - .next() - .copied() - { - self.free_buckets.take(&candidate) - } else { - None - } - } - }; - - // If there is no free bucket available, allocate a new one. - let new_bucket_id = maybe_free_bucket.unwrap_or_else(|| { - let bucket = BucketId(self.allocated_buckets); - self.allocated_buckets += 1; - bucket - }); - - // Assert ascending order (strictly increasing). - if let Some(prev_max) = memory_bucket.last().copied() { - debug_assert!(prev_max < new_bucket_id); - } - + let new_bucket_id = BucketId(self.allocated_buckets); memory_bucket.push(new_bucket_id); + + // Write in stable store that this bucket belongs to the memory with the provided `id`. write( &self.memory, bucket_allocations_address(new_bucket_id).get(), &[id.0], ); + + self.allocated_buckets += 1; } // Grow the underlying memory if necessary. @@ -451,43 +386,6 @@ impl MemoryManagerInner { old_size as i64 } - /// Reclaims memory for the specified virtual memory, marking buckets unallocated - /// and adding them to the free pool. Resets memory size to 0. - /// - /// **CRITICAL**: This invalidates any data structures using this memory ID. - /// Drop the original structure object BEFORE calling this method. - /// - /// Returns the number of pages reclaimed (0 if none allocated). - fn reclaim_memory(&mut self, id: MemoryId) -> u64 { - let memory_buckets = &mut self.memory_buckets[id.0 as usize]; - let bucket_count = memory_buckets.len(); - - if bucket_count == 0 { - return 0; - } - - // Mark all buckets as unallocated in stable storage and collect them - for &bucket_id in memory_buckets.iter() { - write( - &self.memory, - bucket_allocations_address(bucket_id).get(), - &[UNALLOCATED_BUCKET_MARKER], - ); - self.free_buckets.insert(bucket_id); - } - - // Clear the memory's bucket list - memory_buckets.clear(); - - // Reset the memory size to 0 - self.memory_sizes_in_pages[id.0 as usize] = 0; - - self.save_header(); - - // Return pages reclaimed (bucket count * pages per bucket) - bucket_count as u64 * self.bucket_size_in_pages as u64 - } - fn write(&self, id: MemoryId, offset: u64, src: &[u8], bucket_cache: &BucketCache) { if let Some(real_address) = bucket_cache.get(VirtualSegment { address: offset.into(), @@ -700,7 +598,7 @@ impl MemoryId { } // Referring to a bucket. -#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)] +#[derive(Clone, Copy, Debug, PartialEq)] struct BucketId(u16); fn bucket_allocations_address(id: BucketId) -> Address { @@ -1211,228 +1109,4 @@ mod test { None ); } - - #[test] - fn reclaim_empty_memory_returns_zero() { - let mem_mgr = MemoryManager::init(make_memory()); - let memory = mem_mgr.get(MemoryId::new(0)); - - assert_eq!(memory.size(), 0); - assert_eq!(mem_mgr.reclaim_memory(MemoryId::new(0)), 0); - assert_eq!(memory.size(), 0); - } - - #[test] - fn reclaim_returns_correct_page_count() { - let mem_mgr = MemoryManager::init(make_memory()); - let memory = mem_mgr.get(MemoryId::new(0)); - - let pages_allocated = BUCKET_SIZE_IN_PAGES * 2; - memory.grow(pages_allocated); - - let pages_reclaimed = mem_mgr.reclaim_memory(MemoryId::new(0)); - assert_eq!(pages_reclaimed, pages_allocated); - } - - #[test] - fn reclaim_resets_memory_size_to_zero() { - let mem = make_memory(); - let mem_mgr = MemoryManager::init(mem.clone()); - let memory = mem_mgr.get(MemoryId::new(0)); - - memory.grow(BUCKET_SIZE_IN_PAGES * 2); - assert_eq!(memory.size(), BUCKET_SIZE_IN_PAGES * 2); - - mem_mgr.reclaim_memory(MemoryId::new(0)); - assert_eq!(memory.size(), 0); - - // Verify state persists after reload - let mem_mgr_reloaded = MemoryManager::init(mem); - let memory_reloaded = mem_mgr_reloaded.get(MemoryId::new(0)); - assert_eq!(memory_reloaded.size(), 0); - } - - #[test] - fn memory_growth_without_reclaim_uses_new_buckets() { - let mem = make_memory(); - let mem_mgr = MemoryManager::init(mem.clone()); - - // Baseline test: allocate two memories without reclaiming - should use separate buckets - // This contrasts with reclaim tests where buckets can be reused - let memory_0 = mem_mgr.get(MemoryId::new(0)); - let memory_1 = mem_mgr.get(MemoryId::new(1)); - - memory_0.grow(BUCKET_SIZE_IN_PAGES); - memory_0.write(0, b"data_0"); - let size_after_first = mem.size(); - - memory_1.grow(BUCKET_SIZE_IN_PAGES); - memory_1.write(0, b"data_1"); - - // Second allocation should grow underlying memory - assert!(mem.size() > size_after_first); - - // Verify distinct data in separate buckets - let mut buf = vec![0u8; 6]; - memory_0.read(0, &mut buf); - assert_eq!(&buf, b"data_0"); - memory_1.read(0, &mut buf); - assert_eq!(&buf, b"data_1"); - } - - #[test] - fn memory_reuses_buckets_after_reclaim() { - let mem = make_memory(); - let mem_mgr = MemoryManager::init(mem.clone()); - - // Allocate and reclaim first memory - let memory_0 = mem_mgr.get(MemoryId::new(0)); - memory_0.grow(BUCKET_SIZE_IN_PAGES); - memory_0.write(0, b"original"); - let size_after_allocation = mem.size(); - - mem_mgr.reclaim_memory(MemoryId::new(0)); - drop(memory_0); // Drop invalidated memory - - // Allocate second memory - should reuse bucket without growing - let memory_1 = mem_mgr.get(MemoryId::new(1)); - memory_1.grow(BUCKET_SIZE_IN_PAGES); - assert_eq!(mem.size(), size_after_allocation); - - // Verify bucket reuse by reading original data - let mut buf = vec![0u8; 8]; - memory_1.read(0, &mut buf); - assert_eq!(&buf, b"original"); - - // Verify state persists after reload - let mem_mgr_reloaded = MemoryManager::init(mem); - let memory_1_reloaded = mem_mgr_reloaded.get(MemoryId::new(1)); - assert_eq!(memory_1_reloaded.size(), BUCKET_SIZE_IN_PAGES); - memory_1_reloaded.read(0, &mut buf); - assert_eq!(&buf, b"original"); - } - - #[test] - fn conservative_bucket_reuse_prevents_corruption() { - let mem = make_memory(); - let mem_mgr = MemoryManager::init(mem.clone()); - - // Create memories with buckets 0 and 1 - let memory_0 = mem_mgr.get(MemoryId::new(0)); - let memory_1 = mem_mgr.get(MemoryId::new(1)); - memory_0.grow(BUCKET_SIZE_IN_PAGES); - memory_1.grow(BUCKET_SIZE_IN_PAGES); - - // Reclaim memory 0 (frees bucket 0) - mem_mgr.reclaim_memory(MemoryId::new(0)); - - // Memory 1 cannot reuse lower-ID bucket 0 - must allocate bucket 2 - let initial_size = mem_mgr.inner.borrow().memory.size(); - memory_1.grow(BUCKET_SIZE_IN_PAGES); - assert!(mem_mgr.inner.borrow().memory.size() > initial_size); - - // Verify conservative reuse persists after reload - let mem_mgr_reloaded = MemoryManager::init(mem); - let memory_1_reloaded = mem_mgr_reloaded.get(MemoryId::new(1)); - assert_eq!(memory_1_reloaded.size(), BUCKET_SIZE_IN_PAGES * 2); - } - - #[test] - fn optimal_match_conservative_bucket_reuse_prevents_corruption() { - let mem = make_memory(); - let mem_mgr = MemoryManager::init(mem.clone()); - - // Create memories with buckets [0, 2] and [1] - let memory_0 = mem_mgr.get(MemoryId::new(0)); - let memory_1 = mem_mgr.get(MemoryId::new(1)); - memory_0.grow(BUCKET_SIZE_IN_PAGES); // 0 - memory_1.grow(BUCKET_SIZE_IN_PAGES); // 1 - memory_0.grow(BUCKET_SIZE_IN_PAGES); // 2 - - // Reclaim memory 0 (frees buckets 0 and 2) - mem_mgr.reclaim_memory(MemoryId::new(0)); - - // Memory 1 cannot reuse lower-ID bucket 0, but can reuse bucket 2 -- no waste - let initial_size = mem_mgr.inner.borrow().memory.size(); - memory_1.grow(BUCKET_SIZE_IN_PAGES); - assert!(mem_mgr.inner.borrow().memory.size() == initial_size); - - // Verify conservative reuse persists after reload - let mem_mgr_reloaded = MemoryManager::init(mem); - let memory_1_reloaded = mem_mgr_reloaded.get(MemoryId::new(1)); - assert_eq!(memory_1_reloaded.size(), BUCKET_SIZE_IN_PAGES * 2); - } - - #[test] - fn reclaim_frees_buckets_for_reuse_in_order() { - let mem_mgr = MemoryManager::init(make_memory()); - - // Allocate buckets 0, 1, 2 - let memory_0 = mem_mgr.get(MemoryId::new(0)); - let memory_1 = mem_mgr.get(MemoryId::new(1)); - let memory_2 = mem_mgr.get(MemoryId::new(2)); - - memory_0.grow(BUCKET_SIZE_IN_PAGES); - memory_0.write(0, b"bucket_0"); - memory_1.grow(BUCKET_SIZE_IN_PAGES); - memory_1.write(0, b"bucket_1"); - memory_2.grow(BUCKET_SIZE_IN_PAGES); - memory_2.write(0, b"bucket_2"); - - // Reclaim buckets 2 and 0 (out of order) to test free bucket sorting - mem_mgr.reclaim_memory(MemoryId::new(2)); - mem_mgr.reclaim_memory(MemoryId::new(0)); - - // New allocations should reuse reclaimed buckets 0, 2 in increasing order - let memory_3 = mem_mgr.get(MemoryId::new(3)); - let memory_4 = mem_mgr.get(MemoryId::new(4)); - memory_3.grow(BUCKET_SIZE_IN_PAGES); - memory_4.grow(BUCKET_SIZE_IN_PAGES); - - let mut buf = vec![0u8; 8]; - memory_3.read(0, &mut buf); - assert_eq!(&buf, b"bucket_0"); // Reused bucket 0 first - memory_4.read(0, &mut buf); - assert_eq!(&buf, b"bucket_2"); // Reused bucket 2 second - } - - #[test] - fn bucket_ordering_preserved_across_reload() { - let mem = make_memory(); - let mem_mgr = MemoryManager::init_with_bucket_size(mem.clone(), 1); - - // Create memories with buckets 0 and 1 - let memory_a = mem_mgr.get(MemoryId::new(0)); - let memory_b = mem_mgr.get(MemoryId::new(1)); - memory_a.grow(1); - memory_a.write(0, b"a_data"); - memory_b.grow(1); - memory_b.write(0, b"b_data"); - - // Reclaim memory A and grow memory B - // Note: Conservative reuse prevents Memory B from getting bucket 0 (since 0 < 1) - // So Memory B gets a new bucket 2, resulting in buckets [1, 2] - mem_mgr.reclaim_memory(MemoryId::new(0)); - memory_b.grow(1); - - // Verify memory B still reads from its first bucket - let mut buf = vec![0u8; 6]; - memory_b.read(0, &mut buf); - assert_eq!(&buf, b"b_data"); - - // Critical: After reload, memory B should still read same data - let mem_mgr_reloaded = MemoryManager::init_with_bucket_size(mem, 1); - let memory_b_reloaded = mem_mgr_reloaded.get(MemoryId::new(1)); - memory_b_reloaded.read(0, &mut buf); - assert_eq!(&buf, b"b_data"); - } } - -#[cfg(test)] -mod memory_reclaim_core_tests; - -#[cfg(test)] -mod memory_reclaim_btreemap_tests; - -#[cfg(test)] -mod memory_reclaim_vec_tests; diff --git a/src/memory_manager/memory_reclaim_btreemap_tests.rs b/src/memory_manager/memory_reclaim_btreemap_tests.rs deleted file mode 100644 index ff28d4d2..00000000 --- a/src/memory_manager/memory_reclaim_btreemap_tests.rs +++ /dev/null @@ -1,190 +0,0 @@ -//! Migration scenario tests for BTreeMap with memory reclamation. -//! -//! These tests demonstrate real-world migration patterns where users move data -//! from one structure to another. They show how memory reclamation prevents memory -//! waste during common migration scenarios, and most importantly, demonstrate -//! the data corruption bug and its safe usage solution. -//! -//! **CRITICAL SAFETY REQUIREMENTS**: -//! All memory reclamation operations require mandatory Rust object drop BEFORE reclamation. -//! Using original data structures after memory reclamation causes data corruption. -//! See MemoryManager documentation for proper usage patterns. - -use super::{MemoryId, MemoryManager}; -use crate::{btreemap::BTreeMap, vec_mem::VectorMemory, Memory}; - -/// Helper function to create a blob that triggers bucket allocation -fn large_value(id: u64) -> Vec { - let mut data = vec![0u8; 1024]; // 1KB value - // Fill with pattern based on id to make values unique - for (i, byte) in data.iter_mut().enumerate() { - *byte = ((id + i as u64) % 256) as u8; - } - data -} - -#[test] -fn migration_without_reclaim_wastes_buckets() { - // Scenario: Populate A → Drop A without memory reclamation → Populate B - // Result: B cannot reuse A's buckets, causing memory waste (growth) - let (a, b) = (MemoryId::new(0), MemoryId::new(1)); - let mock_stable_memory = VectorMemory::default(); - let mm = MemoryManager::init(mock_stable_memory.clone()); - - // Allocate in A - let mut a_map = BTreeMap::init(mm.get(a)); - for i in 0u64..50 { - a_map.insert(i, large_value(i)); - } - assert_eq!(a_map.len(), 50); - - // Drop A without releasing buckets - drop(a_map); - let stable_before = mock_stable_memory.size(); - - // Allocate in B → should need new buckets since A's aren't reclaimed - let mut b_map = BTreeMap::init(mm.get(b)); - for i in 0u64..50 { - b_map.insert(i, large_value(i + 100)); - } - let stable_after = mock_stable_memory.size(); - - // Verify: B allocated new buckets, stable memory grew (waste) - assert_eq!(b_map.len(), 50); - assert!(stable_after > stable_before, "Stable memory grew (waste)"); -} - -#[test] -fn migration_with_reclaim_reuses_buckets() { - // Scenario: Populate A → Drop A with memory reclamation → Populate B - // Result: B reuses A's reclaimed buckets, preventing memory waste (no growth) - let (a, b) = (MemoryId::new(0), MemoryId::new(1)); - let mock_stable_memory = VectorMemory::default(); - let mm = MemoryManager::init(mock_stable_memory.clone()); - - // Allocate in A - let mut a_map = BTreeMap::init(mm.get(a)); - for i in 0u64..50 { - a_map.insert(i, large_value(i)); - } - assert_eq!(a_map.len(), 50); - - // MANDATORY: Drop the Rust object before reclaiming memory - drop(a_map); - - // Reclaim the memory after dropping the object - let pages_reclaimed = mm.reclaim_memory(a); - assert!(pages_reclaimed > 0); - let stable_before = mock_stable_memory.size(); - - // Allocate in B → should reuse A's reclaimed buckets - let mut b_map = BTreeMap::init(mm.get(b)); - for i in 0u64..50 { - b_map.insert(i, large_value(i + 100)); - } - let stable_after = mock_stable_memory.size(); - - // Verify: B reused A's buckets, stable memory unchanged (reuse) - assert_eq!(b_map.len(), 50); - assert!( - stable_after <= stable_before, - "Stable memory unchanged (reuse)" - ); -} - -/// **DANGER**: This test demonstrates data corruption from unsafe memory reclamation usage. -/// This shows what happens when you DON'T drop the original object before memory reclamation. -#[test] -fn data_corruption_without_mandatory_drop() { - let (a, b) = (MemoryId::new(0), MemoryId::new(1)); - let mm = MemoryManager::init(VectorMemory::default()); - - // Create BTreeMap A with test data - let mut map_a = BTreeMap::init(mm.get(a)); - map_a.insert(1u64, 100u64); - assert_eq!(map_a.get(&1u64).unwrap(), 100u64); - - // DANGEROUS: Reclaim memory but keep map_a alive - mm.reclaim_memory(a); - - // Create BTreeMap B - reuses A's reclaimed buckets - let mut map_b = BTreeMap::init(mm.get(b)); - map_b.insert(2u64, 200u64); - assert_eq!(map_b.get(&2u64).unwrap(), 200u64); - - // CORRUPTION: map_a and map_b now share the same underlying memory - // This can manifest in different ways - either seeing shared data or allocation corruption - - // First check if map_a can see map_b's data (shared memory corruption) - if let Some(corrupted_data) = map_a.get(&2u64) { - assert_eq!( - corrupted_data, 200u64, - "CORRUPTION: map_a sees map_b's data!" - ); - return; // Corruption demonstrated through shared data - } - - // If shared data isn't visible, try to trigger allocation corruption - let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - map_a.insert(3u64, 300u64); - map_a.get(&3u64) - })); - - // Should either panic or show corruption through shared allocations - match result { - Ok(_) => { - // If it succeeds, check if map_b can see the new data (shared allocation) - if map_b.get(&3u64).is_some() { - println!("CORRUPTION: Both maps operating on the same reclaimed memory space"); - } - } - Err(_) => { - // Expected: panic due to allocator corruption - println!("CORRUPTION: Panic due to memory corruption - this proves the bug"); - } - } - - // This test proves why objects MUST be dropped before memory reclamation -} - -/// **SAFE**: This test demonstrates the correct way to use memory reclamation. -/// This shows how mandatory drop prevents data corruption. -#[test] -fn safe_usage_with_mandatory_drop() { - let (a, b) = (MemoryId::new(0), MemoryId::new(1)); - let mm = MemoryManager::init(VectorMemory::default()); - - // Create and populate BTreeMap A - let mut map_a = BTreeMap::init(mm.get(a)); - map_a.insert(1u64, 100u64); - assert_eq!(map_a.get(&1u64).unwrap(), 100u64); - - // MANDATORY: Drop the Rust object before reclaiming memory - drop(map_a); - - // Reclaim the memory after dropping the object - let pages_reclaimed = mm.reclaim_memory(a); - assert!(pages_reclaimed > 0); - - // Create BTreeMap B - safely reuses A's reclaimed buckets - let mut map_b = BTreeMap::init(mm.get(b)); - map_b.insert(2u64, 200u64); - assert_eq!(map_b.get(&2u64).unwrap(), 200u64); - - // Create new BTreeMap on same memory ID A - safe after proper drop - let mut map_a_new = BTreeMap::init(mm.get(a)); - map_a_new.insert(3u64, 300u64); - assert_eq!(map_a_new.get(&3u64).unwrap(), 300u64); - - // Verify maps are completely independent - no corruption - assert_eq!(map_b.len(), 1, "Map B should have 1 element"); - assert_eq!(map_a_new.len(), 1, "Map A new should have 1 element"); - assert_eq!(map_b.get(&2u64).unwrap(), 200u64); - assert_eq!(map_a_new.get(&3u64).unwrap(), 300u64); - - // Both maps can grow independently without corruption - map_a_new.insert(4u64, 400u64); - map_b.insert(5u64, 500u64); - assert_eq!(map_a_new.len(), 2); - assert_eq!(map_b.len(), 2); -} diff --git a/src/memory_manager/memory_reclaim_core_tests.rs b/src/memory_manager/memory_reclaim_core_tests.rs deleted file mode 100644 index c5088e28..00000000 --- a/src/memory_manager/memory_reclaim_core_tests.rs +++ /dev/null @@ -1,132 +0,0 @@ -//! Core memory reclamation functionality tests for MemoryManager. -//! -//! These tests verify the basic memory management operations without dependency -//! on specific data structures. They test the fundamental bucket allocation, -//! reclamation, and reuse mechanisms. -//! -//! **CRITICAL SAFETY REQUIREMENTS**: -//! All memory reclamation operations require mandatory Rust object drop BEFORE reclamation. -//! Using original data structures after memory reclamation causes data corruption. -//! See MemoryManager documentation for proper usage patterns. - -use super::{MemoryId, MemoryManager}; -use crate::{btreemap::BTreeMap, vec_mem::VectorMemory, Memory}; - -/// Helper to create a minimal data structure for bucket allocation -fn allocate_buckets_via_btreemap( - mm: &MemoryManager, - memory_id: MemoryId, - count: u64, -) { - let mut map = BTreeMap::init(mm.get(memory_id)); - for i in 0..count { - map.insert(i, vec![42u8; 2000]); // 2KB blob to trigger bucket allocation - } - drop(map); // Drop the structure before memory reclamation -} - -#[test] -fn reclaim_empty_memory_returns_zero() { - // Test: Reclaiming memory from unused memory should return 0 - let memory_id = MemoryId::new(0); - let mock_stable_memory = VectorMemory::default(); - let mm = MemoryManager::init(mock_stable_memory); - - // Try to reclaim memory from empty memory - let pages_reclaimed = mm.reclaim_memory(memory_id); - - // Should return 0 since no pages were allocated - assert_eq!(pages_reclaimed, 0, "Empty memory should reclaim 0 pages"); -} - -#[test] -fn double_reclaim_returns_zero_second_time() { - // Test: Reclaiming the same memory twice should return 0 on second attempt - let memory_id = MemoryId::new(0); - let mock_stable_memory = VectorMemory::default(); - let mm = MemoryManager::init(mock_stable_memory); - - // Allocate buckets and clear structure - allocate_buckets_via_btreemap(&mm, memory_id, 50); - - // First reclamation should return non-zero - let first_reclaim = mm.reclaim_memory(memory_id); - assert!(first_reclaim > 0, "First reclaim should return pages > 0"); - - // Second reclamation should return 0 (nothing left to reclaim) - let second_reclaim = mm.reclaim_memory(memory_id); - assert_eq!(second_reclaim, 0, "Second reclaim should return 0 pages"); -} - -#[test] -fn reclaim_only_affects_target_memory() { - // Test: Reclaiming memory A should not affect memory B's buckets - let (a, b) = (MemoryId::new(0), MemoryId::new(1)); - let mock_stable_memory = VectorMemory::default(); - let mm = MemoryManager::init(mock_stable_memory); - - // Allocate buckets in both memories - allocate_buckets_via_btreemap(&mm, a, 30); - allocate_buckets_via_btreemap(&mm, b, 30); - - // Reclaim A's memory - let pages_reclaimed = mm.reclaim_memory(a); - assert!(pages_reclaimed > 0, "Should reclaim A's pages"); - - // Verify B still has its memory (can't be reclaimed again) - let b_reclaim_attempt = mm.reclaim_memory(b); - assert!( - b_reclaim_attempt > 0, - "B should still have reclaimable pages" - ); -} - -#[test] -fn multiple_reclaim_cycles() { - // Test: Multiple allocate→reclaim cycles should work consistently - let memory_id = MemoryId::new(0); - let mock_stable_memory = VectorMemory::default(); - let mm = MemoryManager::init(mock_stable_memory); - - // Perform several cycles of allocate→reclaim - for cycle in 0..5 { - // Allocate buckets - allocate_buckets_via_btreemap(&mm, memory_id, 40); - - // Reclaim memory - let pages_reclaimed = mm.reclaim_memory(memory_id); - - assert!( - pages_reclaimed > 0, - "Cycle {} should reclaim >0 pages, got {}", - cycle, - pages_reclaimed - ); - } -} - -#[test] -fn bucket_reuse_prevents_memory_growth() { - // Test: Reclaimed buckets are actually reused by subsequent allocations - let (a, b) = (MemoryId::new(0), MemoryId::new(1)); - let mock_stable_memory = VectorMemory::default(); - let mm = MemoryManager::init(mock_stable_memory.clone()); - - // Allocate buckets in A - allocate_buckets_via_btreemap(&mm, a, 50); - - // Reclaim A's memory - let pages_reclaimed = mm.reclaim_memory(a); - assert!(pages_reclaimed > 0); - let stable_before = mock_stable_memory.size(); - - // Allocate same amount in B → should reuse A's reclaimed buckets - allocate_buckets_via_btreemap(&mm, b, 50); - let stable_after = mock_stable_memory.size(); - - // Verify: stable memory should not grow significantly (reuse occurred) - assert!( - stable_after <= stable_before + 1, // Allow minimal growth for structure overhead - "Stable memory should not grow when reusing reclaimed buckets" - ); -} diff --git a/src/memory_manager/memory_reclaim_vec_tests.rs b/src/memory_manager/memory_reclaim_vec_tests.rs deleted file mode 100644 index d5bd00b4..00000000 --- a/src/memory_manager/memory_reclaim_vec_tests.rs +++ /dev/null @@ -1,198 +0,0 @@ -//! Migration scenario tests for Vec with memory reclamation. -//! -//! These tests demonstrate real-world migration patterns where users move data -//! from one structure to another. They show how memory reclamation prevents memory -//! waste during common migration scenarios, and most importantly, demonstrate -//! the data corruption bug and its safe usage solution for Vec structures. -//! -//! **CRITICAL SAFETY REQUIREMENTS**: -//! All memory reclamation operations require mandatory Rust object drop BEFORE reclamation. -//! Using original data structures after memory reclamation causes data corruption. -//! See MemoryManager documentation for proper usage patterns. -//! -//! **Vec-Specific Notes**: -//! Vec doesn't have a clear() method. Since we drop the object before memory reclamation, -//! there's no need to clear the data first. - -use super::{MemoryId, MemoryManager}; -use crate::{vec::Vec as StableVec, vec_mem::VectorMemory, Memory}; - -/// Helper function to create data that triggers bucket allocation -fn large_data(id: u64) -> [u8; 1024] { - let mut data = [0u8; 1024]; - // Fill with pattern based on id to make data unique - for (i, byte) in data.iter_mut().enumerate() { - *byte = ((id + i as u64) % 256) as u8; - } - data -} - -#[test] -fn migration_without_reclaim_wastes_buckets() { - // Scenario: Populate A → Clear A without memory reclamation → Populate B - // Result: B cannot reuse A's buckets, causing memory waste (growth) - let (a, b) = (MemoryId::new(0), MemoryId::new(1)); - let mock_stable_memory = VectorMemory::default(); - let mm = MemoryManager::init(mock_stable_memory.clone()); - - // Allocate in A - let vec_a = StableVec::init(mm.get(a)); - for i in 0u64..50 { - vec_a.push(&large_data(i)); - } - assert_eq!(vec_a.len(), 50); - - // "Clear" A by creating new instance (overwrites data) without releasing buckets - let vec_a: StableVec<[u8; 1024], _> = StableVec::new(mm.get(a)); - assert_eq!(vec_a.len(), 0); - let stable_before = mock_stable_memory.size(); - - // Allocate in B → should need new buckets since A's aren't reclaimed - let vec_b = StableVec::init(mm.get(b)); - for i in 0u64..50 { - vec_b.push(&large_data(i + 100)); - } - let stable_after = mock_stable_memory.size(); - - // Verify: vecs correct, stable memory grew (waste) - assert_eq!(vec_a.len(), 0); - assert_eq!(vec_b.len(), 50); - assert!(stable_after > stable_before, "Stable memory grew (waste)"); -} - -#[test] -fn migration_with_reclaim_reuses_buckets() { - // Scenario: Populate A → Clear A with memory reclamation → Populate B - // Result: B reuses A's reclaimed buckets, preventing memory waste (no growth) - let (a, b) = (MemoryId::new(0), MemoryId::new(1)); - let mock_stable_memory = VectorMemory::default(); - let mm = MemoryManager::init(mock_stable_memory.clone()); - - // Allocate in A - let vec_a_original = StableVec::init(mm.get(a)); - for i in 0u64..50 { - vec_a_original.push(&large_data(i)); - } - assert_eq!(vec_a_original.len(), 50); - - // MANDATORY: Drop the Rust object before reclaiming memory - drop(vec_a_original); - - // Reclaim the memory after dropping the object - let pages_reclaimed = mm.reclaim_memory(a); - assert!(pages_reclaimed > 0); - let stable_before = mock_stable_memory.size(); - - // Allocate in B → should reuse A's reclaimed buckets - let vec_b = StableVec::init(mm.get(b)); - for i in 0u64..50 { - vec_b.push(&large_data(i + 100)); - } - let stable_after = mock_stable_memory.size(); - - // Verify: B reused A's buckets, stable memory unchanged (reuse) - assert_eq!(vec_b.len(), 50); - assert!( - stable_after <= stable_before, - "Stable memory unchanged (reuse)" - ); -} - -/// **DANGER**: This test demonstrates data corruption from unsafe memory reclamation usage. -/// This shows what happens when you DON'T drop the original object after memory reclamation. -#[test] -fn data_corruption_without_mandatory_drop() { - let (a, b) = (MemoryId::new(0), MemoryId::new(1)); - let mm = MemoryManager::init(VectorMemory::default()); - - // Create Vec A with test data - let vec_a = StableVec::init(mm.get(a)); - vec_a.push(&1u64); - assert_eq!(vec_a.get(0).unwrap(), 1u64); - - // "Clear" by creating new instance, but keep vec_a alive (DANGEROUS!) - let vec_a: StableVec = StableVec::new(mm.get(a)); - assert_eq!(vec_a.len(), 0); - mm.reclaim_memory(a); - - // Create Vec B - reuses A's reclaimed buckets - let vec_b = StableVec::init(mm.get(b)); - vec_b.push(&2u64); - assert_eq!(vec_b.get(0).unwrap(), 2u64); - - // CORRUPTION: vec_a and vec_b now share the same underlying memory - // This can manifest in different ways - either seeing shared data or allocation corruption - - // First check if vec_a can see vec_b's data (shared memory corruption) - if !vec_a.is_empty() { - if let Some(corrupted_data) = vec_a.get(0) { - assert_eq!(corrupted_data, 2u64, "CORRUPTION: vec_a sees vec_b's data!"); - return; // Corruption demonstrated through shared data - } - } - - // If shared data isn't visible, try to trigger allocation corruption - let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - vec_a.push(&3u64); - vec_a.get(0) - })); - - // Should either panic or show corruption through shared allocations - match result { - Ok(_) => { - // If it succeeds, check if vec_b can see the new data (shared allocation) - if !vec_a.is_empty() && vec_b.len() > vec_a.len() { - // Check if data appears in unexpected places - println!("CORRUPTION: Both vecs operating on the same reclaimed memory space"); - } - } - Err(_) => { - // Expected: panic due to allocator corruption - println!("CORRUPTION: Panic due to memory corruption - this proves the bug"); - } - } - - // This test proves why objects MUST be dropped after memory reclamation -} - -/// **SAFE**: This test demonstrates the correct way to use memory reclamation. -/// This shows how mandatory drop prevents data corruption. -#[test] -fn safe_usage_with_mandatory_drop() { - let (a, b) = (MemoryId::new(0), MemoryId::new(1)); - let mm = MemoryManager::init(VectorMemory::default()); - - // Create and populate Vec A - let vec_a = StableVec::init(mm.get(a)); - vec_a.push(&1u64); - assert_eq!(vec_a.get(0).unwrap(), 1u64); - - // MANDATORY: Drop the Rust object before reclaiming memory - drop(vec_a); - - // Reclaim the memory after dropping the object - let pages_reclaimed = mm.reclaim_memory(a); - assert!(pages_reclaimed > 0); - - // Create Vec B - safely reuses A's reclaimed buckets - let vec_b = StableVec::init(mm.get(b)); - vec_b.push(&2u64); - assert_eq!(vec_b.get(0).unwrap(), 2u64); - - // Create new Vec on same memory ID A - safe after proper drop - let vec_a_new = StableVec::init(mm.get(a)); - vec_a_new.push(&3u64); - assert_eq!(vec_a_new.get(0).unwrap(), 3u64); - - // Verify vecs are completely independent - no corruption - assert_eq!(vec_b.len(), 1, "Vec B should have 1 element"); - assert_eq!(vec_a_new.len(), 1, "Vec A new should have 1 element"); - assert_eq!(vec_b.get(0).unwrap(), 2u64); - assert_eq!(vec_a_new.get(0).unwrap(), 3u64); - - // Both vecs can grow independently without corruption - vec_a_new.push(&4u64); - vec_b.push(&5u64); - assert_eq!(vec_a_new.len(), 2); - assert_eq!(vec_b.len(), 2); -}