feat: try_release_virtual_memory_buckets

Author: maksymar

src/memory_manager.rs

@@ -127,6 +127,20 @@ const HEADER_RESERVED_BYTES: usize = 32;
 /// Bucket MAX_NUM_BUCKETS                ↕ N pages
 /// ```
 ///
+/// # Bucket Release and Memory Reclamation
+///
+/// The memory manager now supports manual bucket release to address long-term memory cost concerns.
+/// When data structures are cleared, their underlying buckets can be explicitly released for reuse:
+///
+/// - **Safe Release**: Use `try_release_virtual_memory_buckets()` which checks if the virtual memory
+///   size is 0 before releasing buckets. This prevents accidental bucket release when data might
+///   still exist.
+/// - **Bucket Reuse**: Released buckets are automatically reused when new virtual memories need
+///   storage, reducing overall memory allocation.
+///
+/// Note: Calling `clear_new()` on data structures does not automatically shrink virtual memory size,
+/// so the safe release method provides an important safety check.
+///
 /// # Current Limitations and Future Improvements
 ///
 /// - Buckets are never deallocated once assigned to a memory ID, even when the memory becomes empty
@@ -162,6 +176,33 @@ impl<M: Memory> MemoryManager<M> {
         }
     }
 
+    /// Safely releases buckets only if the virtual memory appears to be empty.
+    ///
+    /// This method checks if the memory size is 0 before releasing buckets, providing a safety check
+    /// to avoid releasing buckets when data structures might still contain data.
+    /// Returns Ok(bucket_count) if buckets were released, or Err(current_size) if the memory
+    /// appears to still contain data.
+    ///
+    /// This method should be called when a data structure is completely cleared
+    /// to allow its memory buckets to be reused by other data structures.
+    ///
+    /// # Example
+    /// ```rust,ignore
+    /// let memory_manager = MemoryManager::init(DefaultMemoryImpl::default());
+    /// let memory_id = MemoryId::new(0);
+    ///
+    /// // After clearing a data structure
+    /// match memory_manager.try_release_virtual_memory_buckets(memory_id) {
+    ///     Ok(count) => println!("Released {} buckets", count),
+    ///     Err(size) => println!("Memory still has {} pages, cannot release", size),
+    /// }
+    /// ```
+    pub fn try_release_virtual_memory_buckets(&self, id: MemoryId) -> Result<usize, u64> {
+        self.inner
+            .borrow_mut()
+            .try_release_virtual_memory_buckets(id)
+    }
+
     /// Returns the underlying memory.
     ///
     /// # Returns
@@ -246,6 +287,9 @@ struct MemoryManagerInner<M: Memory> {
 
     /// A map mapping each managed memory to the bucket ids that are allocated to it.
     memory_buckets: Vec<Vec<BucketId>>,
+
+    /// A pool of free buckets that can be reused by any memory.
+    free_buckets: Vec<BucketId>,
 }
 
 impl<M: Memory> MemoryManagerInner<M> {
@@ -274,6 +318,7 @@ impl<M: Memory> MemoryManagerInner<M> {
             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: Vec::new(),
         };
 
         mem_mgr.save_header();
@@ -303,9 +348,13 @@ impl<M: Memory> MemoryManagerInner<M> {
         );
 
         let mut memory_buckets = vec![vec![]; MAX_NUM_MEMORIES as usize];
-        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));
+        let mut free_buckets = Vec::new();
+        for (bucket_idx, memory_id) in buckets.into_iter().enumerate() {
+            if memory_id != UNALLOCATED_BUCKET_MARKER {
+                memory_buckets[memory_id as usize].push(BucketId(bucket_idx as u16));
+            } else if (bucket_idx as u16) < header.num_allocated_buckets {
+                // This bucket was allocated but is now marked as unallocated, so it's free to reuse
+                free_buckets.push(BucketId(bucket_idx as u16));
             }
         }
 
@@ -315,6 +364,7 @@ impl<M: Memory> MemoryManagerInner<M> {
             bucket_size_in_pages: header.bucket_size_in_pages,
             memory_sizes_in_pages: header.memory_sizes_in_pages,
             memory_buckets,
+            free_buckets,
         }
     }
 
@@ -345,7 +395,11 @@ impl<M: Memory> MemoryManagerInner<M> {
         let required_buckets = self.num_buckets_needed(new_size);
         let new_buckets_needed = required_buckets - current_buckets;
 
-        if new_buckets_needed + self.allocated_buckets as u64 > MAX_NUM_BUCKETS {
+        // Check if we have enough buckets available (either already allocated or can allocate new ones)
+        let available_free_buckets = self.free_buckets.len() as u64;
+        let new_buckets_to_allocate = new_buckets_needed.saturating_sub(available_free_buckets);
+
+        if self.allocated_buckets as u64 + new_buckets_to_allocate > MAX_NUM_BUCKETS {
             // Exceeded the memory that can be managed.
             return -1;
         }
@@ -354,7 +408,16 @@ impl<M: Memory> MemoryManagerInner<M> {
         // Allocate new buckets as needed.
         memory_bucket.reserve(new_buckets_needed as usize);
         for _ in 0..new_buckets_needed {
-            let new_bucket_id = BucketId(self.allocated_buckets);
+            let new_bucket_id = if let Some(free_bucket) = self.free_buckets.pop() {
+                // Reuse a bucket from the free pool
+                free_bucket
+            } else {
+                // Allocate a new bucket
+                let bucket = BucketId(self.allocated_buckets);
+                self.allocated_buckets += 1;
+                bucket
+            };
+
             memory_bucket.push(new_bucket_id);
 
             // Write in stable store that this bucket belongs to the memory with the provided `id`.
@@ -363,8 +426,6 @@ impl<M: Memory> MemoryManagerInner<M> {
                 bucket_allocations_address(new_bucket_id).get(),
                 &[id.0],
             );
-
-            self.allocated_buckets += 1;
         }
 
         // Grow the underlying memory if necessary.
@@ -386,6 +447,41 @@ impl<M: Memory> MemoryManagerInner<M> {
         old_size as i64
     }
 
+    /// Safely releases buckets only if the virtual memory appears to be empty.
+    ///
+    /// Checks if the memory size is 0 before releasing buckets. This provides a safety check
+    /// to avoid releasing buckets when data structures might still contain data.
+    ///
+    /// Returns Ok(bucket_count) if buckets were released, or Err(current_size) if the memory
+    /// appears to still contain data.
+    fn try_release_virtual_memory_buckets(&mut self, id: MemoryId) -> Result<usize, u64> {
+        let current_size = self.memory_sizes_in_pages[id.0 as usize];
+
+        if current_size == 0 {
+            // Memory appears empty, safe to release buckets
+            let memory_buckets = &mut self.memory_buckets[id.0 as usize];
+            let bucket_count = memory_buckets.len();
+
+            // Mark all buckets as unallocated in stable storage and move to free pool
+            for &bucket_id in memory_buckets.iter() {
+                write(
+                    &self.memory,
+                    bucket_allocations_address(bucket_id).get(),
+                    &[UNALLOCATED_BUCKET_MARKER],
+                );
+                self.free_buckets.push(bucket_id);
+            }
+
+            // Clear the memory's bucket list (size is already 0)
+            memory_buckets.clear();
+            self.save_header();
+
+            Ok(bucket_count)
+        } else {
+            Err(current_size)
+        }
+    }
+
     fn write(&self, id: MemoryId, offset: u64, src: &[u8], bucket_cache: &BucketCache) {
         if let Some(real_address) = bucket_cache.get(VirtualSegment {
             address: offset.into(),
@@ -1110,3 +1206,6 @@ mod test {
         );
     }
 }
+
+#[cfg(test)]
+mod bucket_release_tests;