fix: BTreeMap memory leak when deallocating nodes with overflows (#212)

# Problem
When storing unbounded types in a `BTreeMap`, a node is represented as a
linked list of "memory chunks". It was discovered recently that when we
deallocate a node, in some cases only the first memory chunk is
deallocated, and the rest of the memory chunks remain (incorrectly)
allocated, causing a memory leak.

# Solution
Change the logic for deallocating nodes to ensure that all of a node's
memory chunks are deallocated. Tests have been added to prevent
regressions of this nature moving forward.

proptest-regressions/storable/tests.txt

@@ -0,0 +1,9 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc bba96dcb90414b34f8abdf74f80e9b15da65b9276309c25baf124b6268f22b80 # shrinks to v = Some(([], [0, 0], []))
+cc 5fff4549ed43e16909e38ae3ddb942e40822a3859b7f07c74b2f5da825c1f572 # shrinks to v = Some((0, [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], []))
+cc b70e6693f80d4af46e8d66ffcabf319448163991202c2d33c604c8b0dd0bfd1b # shrinks to v = Some((0, [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], []))

src/btreemap.rs

@@ -663,7 +663,7 @@ where
                             );
 
                             // Deallocate the empty node.
-                            self.allocator.deallocate(node.address());
+                            node.deallocate(&mut self.allocator);
                             self.root_addr = NULL;
                         } else {
                             node.save(self.allocator_mut());
@@ -953,9 +953,10 @@ where
                             node.remove_child(idx);
 
                             if node.entries_len() == 0 {
-                                self.allocator.deallocate(node.address());
+                                let node_address = node.address();
+                                node.deallocate(&mut self.allocator);
 
-                                if node.address() == self.root_addr {
+                                if node_address == self.root_addr {
                                     // Update the root.
                                     self.root_addr = left_sibling.address();
                                     self.save();
@@ -981,9 +982,10 @@ where
                             node.remove_child(idx);
 
                             if node.entries_len() == 0 {
-                                self.allocator.deallocate(node.address());
+                                let node_address = node.address();
+                                node.deallocate(&mut self.allocator);
 
-                                if node.address() == self.root_addr {
+                                if node_address == self.root_addr {
                                     // Update the root.
                                     self.root_addr = right_sibling.address();
                                     self.save();
@@ -1199,10 +1201,7 @@ where
     //   [1, 2, 3, 4, 5, 6, 7] (stored in the `into` node)
     //   `source` is deallocated.
     fn merge(&mut self, source: Node<K>, mut into: Node<K>, median: Entry<K>) -> Node<K> {
-        let source_address = source.address();
-        into.merge(source, median, self.memory());
-        into.save(self.allocator_mut());
-        self.allocator.deallocate(source_address);
+        into.merge(source, median, &mut self.allocator);
         into
     }
 
@@ -3085,4 +3084,46 @@ mod test {
 
         assert_eq!(header_actual, header_expected);
     }
+
+    #[test]
+    fn deallocating_node_with_overflows() {
+        let mem = make_memory();
+        let mut btree: BTreeMap<Vec<u8>, Vec<u8>, _> = BTreeMap::new(mem.clone());
+
+        // No allocated chunks yet.
+        assert_eq!(btree.allocator.num_allocated_chunks(), 0);
+
+        // Insert and remove an entry that's large and requires overflow pages.
+        btree.insert(vec![0; 10_000], vec![]);
+
+        // At least two chunks should be allocated.
+        // One for the node itself and at least one overflow page.
+        assert!(btree.allocator.num_allocated_chunks() >= 2);
+        btree.remove(&vec![0; 10_000]);
+
+        // All chunks have been deallocated.
+        assert_eq!(btree.allocator.num_allocated_chunks(), 0);
+    }
+
+    #[test]
+    fn repeatedly_deallocating_nodes_with_overflows() {
+        let mem = make_memory();
+        let mut btree: BTreeMap<Vec<u8>, Vec<u8>, _> = BTreeMap::new(mem.clone());
+
+        // No allocated chunks yet.
+        assert_eq!(btree.allocator.num_allocated_chunks(), 0);
+
+        for _ in 0..100 {
+            for i in 0..100 {
+                btree.insert(vec![i; 10_000], vec![]);
+            }
+
+            for i in 0..100 {
+                btree.remove(&vec![i; 10_000]);
+            }
+        }
+
+        // All chunks have been deallocated.
+        assert_eq!(btree.allocator.num_allocated_chunks(), 0);
+    }
 }

src/btreemap/node.rs

@@ -289,13 +289,18 @@ impl<K: Storable + Ord + Clone> Node<K> {
     ///   * `self` and `source` are of the same node type.
     ///
     /// POSTCONDITION:
-    ///   * `source` is empty (no entries and no children).
+    ///   * `source` is deallocated.
     ///   * all the entries of `source`, as well as the median, are merged into `self`, in sorted
     ///      order.
-    pub fn merge<M: Memory>(&mut self, mut source: Node<K>, median: Entry<K>, memory: &M) {
+    pub fn merge<M: Memory>(
+        &mut self,
+        mut source: Node<K>,
+        median: Entry<K>,
+        allocator: &mut Allocator<M>,
+    ) {
         // Load all the values from the source node first, as they will be moved out.
         for i in 0..source.entries_len() {
-            source.value(i, memory);
+            source.value(i, allocator.memory());
         }
 
         if source.key(0) > self.key(0) {
@@ -308,10 +313,13 @@ impl<K: Storable + Ord + Clone> Node<K> {
             Self::append(&mut source, self, median);
 
             // Move the entries and children into self.
-            self.keys = source.keys;
-            self.encoded_values = source.encoded_values;
-            self.children = source.children;
+            self.keys = core::mem::take(&mut source.keys);
+            self.encoded_values = core::mem::take(&mut source.encoded_values);
+            self.children = core::mem::take(&mut source.children);
         }
+
+        self.save(allocator);
+        source.deallocate(allocator);
     }
 
     // Appends the entries and children of node `b` into `a`, along with the median entry.
@@ -416,6 +424,15 @@ impl<K: Storable + Ord + Clone> Node<K> {
         self.pop_entry(memory)
             .expect("An initially full node cannot be empty")
     }
+
+    /// Deallocates a node.
+    pub fn deallocate<M: Memory>(self, allocator: &mut Allocator<M>) {
+        for overflow in self.overflows.into_iter() {
+            allocator.deallocate(overflow);
+        }
+
+        allocator.deallocate(self.address);
+    }
 }
 
 // A transient data structure for reading/writing metadata into/from stable memory.

src/btreemap/proptests.rs

@@ -159,6 +159,25 @@ fn iter_count_test(#[strategy(0..250u8)] start: u8, #[strategy(#start..255u8)] e
     });
 }
 
+#[proptest]
+fn no_memory_leaks(#[strategy(pvec(pvec(0..u8::MAX, 100..10_000), 100))] keys: Vec<Vec<u8>>) {
+    let mem = make_memory();
+    let mut btree = BTreeMap::new(mem);
+
+    // Insert entries.
+    for k in keys.iter() {
+        btree.insert(k.clone(), ());
+    }
+
+    // Remove entries.
+    for k in keys.iter() {
+        btree.remove(k);
+    }
+
+    // After inserting and deleting all the entries, there should be no allocated chunks.
+    assert_eq!(btree.allocator.num_allocated_chunks(), 0);
+}
+
 // Given an operation, executes it on the given stable btreemap and standard btreemap, verifying
 // that the result of the operation is equal in both btrees.
 fn execute_operation<M: Memory>(