Store non-full pages in a BTreeSet, not a Vec

And sort them by number of available var-len granules.
This prevents an accidentally quadratic behavior where,
for a table where the average row contains many var-len granules,
after inserting a large number of rows, there would be a large number of pages in `non_full_pages`
each of which had enough space for at least one fixed-len row part,
but insufficient space for an actual row in practice due to insufficient var-len granules.
Each insertion would then do a linear scan over `non_full_pages`
before either inserting into the last page or allocating a new page which went to the end.

Now, non-full pages are stored in a `BTreeSet` sorted by the number of free var-len granules,
and the search for a useable page is done with a `BTreeSet::range` iterator for only the pages with enough granules.
I think there may still be an off-by-one-ish bug here,
where a page may have enough bytes in the gap that it could either store the fixed-len part
or the var-len granules, but not both,
but this fix hopefully will suffice for now.

Author: gefjon

crates/memory-usage/src/lib.rs

@@ -126,6 +126,13 @@ impl<K: MemoryUsage, V: MemoryUsage> MemoryUsage for std::collections::BTreeMap<
     }
 }
 
+impl<T: MemoryUsage> MemoryUsage for std::collections::BTreeSet<T> {
+    fn heap_usage(&self) -> usize {
+        // NB: this is best-effort, since we don't have a `capacity()` method on `BTreeMap`.
+        self.len() * mem::size_of::<T>() + self.iter().map(|t| t.heap_usage()).sum::<usize>()
+    }
+}
+
 #[cfg(feature = "smallvec")]
 impl<A: smallvec::Array> MemoryUsage for smallvec::SmallVec<A>
 where

crates/table/src/page.rs

@@ -382,6 +382,13 @@ impl PageHeader {
     pub(super) fn present_rows_storage_ptr_for_test(&self) -> *const () {
         self.fixed.present_rows.storage().as_ptr().cast()
     }
+
+    /// Returns the number of var-len granules available for allocation,
+    /// including those in the "gap" between the fixed-len and var-len part of the page.
+    fn available_var_len_granules(&self) -> usize {
+        self.var.freelist_len as usize
+            + VarLenGranule::space_to_granules(gap_remaining_size(self.var.first, self.fixed.last))
+    }
 }
 
 /// Fixed-length row portions must be at least large enough to store a `FreeCellRef`.
@@ -1195,6 +1202,12 @@ impl Page {
         self.header.var.num_granules as usize
     }
 
+    /// Returns the number of var-len granules free to store data,
+    /// including those in the "gap" between the fixed-len and var-len part of the page.
+    pub fn available_var_len_granules(&self) -> usize {
+        self.header.available_var_len_granules()
+    }
+
     #[cfg(test)]
     /// # Safety
     ///

crates/table/src/pages.rs

@@ -9,6 +9,7 @@ use super::var_len::VarLenMembers;
 use core::ops::{ControlFlow, Deref, Index, IndexMut};
 use spacetimedb_sats::layout::Size;
 use spacetimedb_sats::memory_usage::MemoryUsage;
+use std::collections::BTreeSet;
 use std::ops::DerefMut;
 use thiserror::Error;
 
@@ -39,8 +40,10 @@ impl IndexMut<PageIndex> for Pages {
 pub struct Pages {
     /// The collection of pages under management.
     pages: Vec<Box<Page>>,
-    /// The set of pages that aren't yet full.
-    non_full_pages: Vec<PageIndex>,
+
+    /// The set of pages that aren't yet full,
+    /// sorted by the number of var-len granules available in each page.
+    non_full_pages: BTreeSet<(usize, PageIndex)>,
 }
 
 impl MemoryUsage for Pages {
@@ -78,7 +81,9 @@ impl Pages {
             page.clear();
         }
         // Mark every page non-full.
-        self.non_full_pages = (0..self.pages.len()).map(|idx| PageIndex(idx as u64)).collect();
+        self.non_full_pages = (0..self.pages.len())
+            .map(|idx| (self.pages[idx].available_var_len_granules(), PageIndex(idx as u64)))
+            .collect();
     }
 
     /// Get a reference to fixed-len row data.
@@ -113,14 +118,15 @@ impl Pages {
 
     /// Mark the page at `idx` as non-full.
     pub fn mark_page_non_full(&mut self, idx: PageIndex) {
-        self.non_full_pages.push(idx);
+        let available_granules = self.pages[idx.idx()].available_var_len_granules();
+        self.non_full_pages.insert((available_granules, idx));
     }
 
     /// If the page at `page_index` is not full,
     /// add it to the non-full set so that later insertions can access it.
     pub fn maybe_mark_page_non_full(&mut self, page_index: PageIndex, fixed_row_size: Size) {
         if !self[page_index].is_full(fixed_row_size) {
-            self.non_full_pages.push(page_index);
+            self.mark_page_non_full(page_index);
         }
     }
 
@@ -151,14 +157,13 @@ impl Pages {
         fixed_row_size: Size,
         num_var_len_granules: usize,
     ) -> Result<PageIndex, Error> {
-        if let Some((page_idx_idx, page_idx)) = self
+        if let Some((page_num_free_granules, page_idx)) = self
             .non_full_pages
-            .iter()
+            .range((num_var_len_granules, PageIndex(0))..)
             .copied()
-            .enumerate()
             .find(|(_, page_idx)| self[*page_idx].has_space_for_row(fixed_row_size, num_var_len_granules))
         {
-            self.non_full_pages.swap_remove(page_idx_idx);
+            self.non_full_pages.remove(&(page_num_free_granules, page_idx));
             return Ok(page_idx);
         }
 
@@ -358,7 +363,9 @@ impl Pages {
         self.non_full_pages = pages
             .iter()
             .enumerate()
-            .filter_map(|(idx, page)| (!page.is_full(fixed_row_size)).then_some(PageIndex(idx as _)))
+            .filter_map(|(idx, page)| {
+                (!page.is_full(fixed_row_size)).then_some((page.available_var_len_granules(), PageIndex(idx as _)))
+            })
             .collect();
         self.pages = pages;
     }