Add a thread-local cache for hot btree pages

CHANGELOG.md

@@ -57,6 +57,9 @@
   optional, and lets container types holding that key store shorter separators.
 * Shorten separators for array keys, when the element type is variable width. Tables with such
   keys use slightly less space, and lookups are faster.
+* Improve multithreaded read scaling. Each thread now caches the level below the root of the
+  btrees it reads, removing lock and reference-count contention on those hot pages. Random reads
+  are about 1.2x faster with 16 threads and 1.6x faster with 32.
 * Fix a crash shortly after a commit being able to silently roll that commit back during
   recovery, if `check_integrity()` had previously repaired the database.
 * Fix iterators silently omitting data when iteration continues after an error. An iterator

src/db.rs

@@ -5,6 +5,7 @@ use crate::tree_store::ReadOnlyBackend;
 use crate::tree_store::{
     AllocationPolicy, BtreeHeader, InternalTableDefinition, PAGE_SIZE, PageHint, PageNumber,
     PageResolver, ShrinkPolicy, TableTree, TableType, TransactionalMemory,
+    release_thread_local_page_cache,
 };
 use crate::types::{Key, Value};
 use crate::{
@@ -341,6 +342,10 @@ impl TransactionGuard {
             }
         }
     }
+
+    pub(crate) fn is_read(&self) -> bool {
+        matches!(self, Self::Read { .. })
+    }
 }
 
 impl Drop for TransactionGuard {
@@ -349,7 +354,12 @@ impl Drop for TransactionGuard {
             Self::Read {
                 tracker,
                 transaction_id,
-            } => tracker.deallocate_read_transaction(*transaction_id),
+            } => {
+                // Clear the thread local cache, so that it doesn't leak cache space if this
+                // thread doesn't run another transaction
+                release_thread_local_page_cache();
+                tracker.deallocate_read_transaction(*transaction_id);
+            }
             Self::Write {
                 tracker,
                 transaction_id,

src/tree_store/btree.rs

@@ -11,13 +11,16 @@ use crate::tree_store::btree_cursor::BtreeCursorMut;
 use crate::tree_store::btree_cursor::{CursorMut, Position};
 use crate::tree_store::btree_iters::{bounds_are_empty, encode_bounds};
 use crate::tree_store::btree_mutator::MutateHelper;
-use crate::tree_store::page_store::{Page, PageImpl, PageMut};
+use crate::tree_store::page_store::{
+    CacheTag, Page, PageImpl, PageMut, ThreadLocalPageCache, with_thread_local_page_cache,
+};
 use crate::tree_store::{
     AccessGuardMutInPlace, AllPageNumbersBtreeIter, BtreeCursorRange, BtreeExtractIf,
     PageAllocator, PageHint, PageNumber, PageNumberHashMap, PageResolver, PageTracker,
 };
 use crate::types::{Key, MutInPlaceValue, Value};
 use crate::{AccessGuard, Result, StorageError};
+use alloc::rc::Rc;
 use alloc::string::ToString;
 use alloc::sync::Arc;
 use alloc::vec;
@@ -1005,11 +1008,40 @@ impl RawBtree {
     }
 }
 
+// A page visited during a lookup: the borrowed root, one from the shared page cache, or this
+// thread's private copy of a branch page
+enum NodePage<'a> {
+    Root(&'a PageImpl),
+    Shared(PageImpl),
+    Local(PageNumber, Rc<[u8]>),
+}
+
+impl Page for NodePage<'_> {
+    fn memory(&self) -> &[u8] {
+        match self {
+            Self::Root(page) => page.memory(),
+            Self::Shared(page) => page.memory(),
+            Self::Local(_, bytes) => bytes,
+        }
+    }
+
+    fn get_page_number(&self) -> PageNumber {
+        match self {
+            Self::Root(page) => page.get_page_number(),
+            Self::Shared(page) => page.get_page_number(),
+            Self::Local(page_number, _) => *page_number,
+        }
+    }
+}
+
 pub(crate) struct Btree<K: Key + 'static, V: Value + 'static> {
     mem: PageResolver,
     transaction_guard: Arc<TransactionGuard>,
     // Cache of the root page to avoid repeated lookups
     cached_root: Option<PageImpl>,
+    // None for write transactions, which can free and reallocate pages within one generation
+    cache_tag: Option<CacheTag>,
+    cache_max_pages: usize,
     root: Option<BtreeHeader>,
     hint: PageHint,
     _key_type: PhantomData<K>,
@@ -1028,10 +1060,20 @@ impl<K: Key, V: Value> Btree<K, V> {
         } else {
             None
         };
+        let cache_tag = if guard.is_read() {
+            Some(CacheTag {
+                instance: mem.cache_instance_id(),
+                generation: guard.id().raw_id(),
+            })
+        } else {
+            None
+        };
         Ok(Self {
+            cache_max_pages: mem.thread_local_cache_pages(),
             mem,
             transaction_guard: guard,
             cached_root,
+            cache_tag,
             root,
             hint,
             _key_type: PhantomData,
@@ -1084,40 +1126,90 @@ impl<K: Key, V: Value> Btree<K, V> {
     }
 
     pub(crate) fn get(&self, key: &K::SelfType<'_>) -> Result<Option<AccessGuard<'static, V>>> {
-        if let Some(ref root_page) = self.cached_root {
-            self.get_helper(root_page, K::as_bytes(key).as_ref())
+        let Some(ref root_page) = self.cached_root else {
+            return Ok(None);
+        };
+        let key_bytes = K::as_bytes(key);
+        let query = key_bytes.as_ref();
+        if let Some(tag) = self.cache_tag {
+            with_thread_local_page_cache(|cache| {
+                self.get_helper(&mut cache.map(|cache| (cache, tag)), root_page, query)
+            })
         } else {
-            Ok(None)
+            self.get_helper(&mut None, root_page, query)
         }
     }
 
     // Returns the value for the queried key, if present
-    fn get_helper(&self, page: &PageImpl, query: &[u8]) -> Result<Option<AccessGuard<'static, V>>> {
-        let mut descended: Option<PageImpl> = None;
+    fn get_helper(
+        &self,
+        cache: &mut Option<(&mut ThreadLocalPageCache, CacheTag)>,
+        root: &PageImpl,
+        query: &[u8],
+    ) -> Result<Option<AccessGuard<'static, V>>> {
+        let mut node = NodePage::Root(root);
+        // Counting the root's children as 1
+        let mut child_depth = 1;
         for _ in 0..MAX_BTREE_DEPTH {
-            let page = descended.as_ref().unwrap_or(page);
-            let child_page = match page.memory()[0] {
+            let child_number = match node.memory()[0] {
                 LEAF => {
                     let accessor =
-                        LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width());
-                    return Ok(accessor.find_key::<K>(query).map(|entry_index| {
-                        let (start, end) = accessor.value_range(entry_index).unwrap();
-                        AccessGuard::with_page(page.clone(), start..end)
-                    }));
+                        LeafAccessor::new(node.memory(), K::fixed_width(), V::fixed_width());
+                    let Some(entry_index) = accessor.find_key::<K>(query) else {
+                        return Ok(None);
+                    };
+                    let (start, end) = accessor.value_range(entry_index).unwrap();
+                    let page = match node {
+                        NodePage::Root(page) => page.clone(),
+                        NodePage::Shared(page) => page,
+                        NodePage::Local(..) => unreachable!("only branch pages are cached"),
+                    };
+                    return Ok(Some(AccessGuard::with_page(page, start..end)));
                 }
                 BRANCH => {
-                    let accessor = BranchAccessor::new(page, K::fixed_width());
+                    let accessor = BranchAccessor::new(&node, K::fixed_width());
                     accessor.child_for_key::<K>(query).1
                 }
                 _ => unreachable!(),
             };
-            descended = Some(self.mem.get_page(child_page, self.hint)?);
+            node = self.get_page_for_read(cache, child_number, child_depth)?;
+            child_depth += 1;
         }
         Err(StorageError::Corrupted(
             "Btree exceeded maximum depth".to_string(),
         ))
     }
 
+    // Levels below the root served from the thread local cache. Only the first, because a level
+    // holds roughly a page's worth of children per page of the level above it: the one below the
+    // root is small enough to cache whole, and the next is already too large. Caching a level
+    // that does not fit is worse than not caching it, since each miss copies a page
+    const MAX_CACHED_DEPTH: u32 = 1;
+
+    // Fetches a page during a lookup, from this thread's cache when it holds it. Only branch
+    // pages are cached, because an AccessGuard returned for a leaf holds the PageImpl
+    fn get_page_for_read<'a>(
+        &'a self,
+        cache: &mut Option<(&mut ThreadLocalPageCache, CacheTag)>,
+        page_number: PageNumber,
+        depth: u32,
+    ) -> Result<NodePage<'a>> {
+        if depth <= Self::MAX_CACHED_DEPTH
+            && page_number.page_order == 0
+            && let Some((cache, tag)) = cache
+        {
+            if let Some(bytes) = cache.get(*tag, page_number) {
+                return Ok(NodePage::Local(page_number, bytes));
+            }
+            let page = self.mem.get_page(page_number, self.hint)?;
+            if page.memory()[0] == BRANCH {
+                cache.insert(*tag, self.cache_max_pages, page_number, page.memory());
+            }
+            return Ok(NodePage::Shared(page));
+        }
+        Ok(NodePage::Shared(self.mem.get_page(page_number, self.hint)?))
+    }
+
     pub(crate) fn first(
         &self,
     ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {

src/tree_store/mod.rs

@@ -30,6 +30,7 @@ pub(crate) use page_store::{
     AllocationPolicy, FILE_FORMAT_VERSION3, MAX_PAIR_LENGTH, MAX_VALUE_LENGTH, PAGE_SIZE, Page,
     PageAllocator, PageHint, PageNumber, PageNumberHashMap, PageNumberHashSet, PageResolver,
     PageTracker, SerializedSavepoint, ShrinkPolicy, TransactionalMemory,
+    release_thread_local_page_cache,
 };
 pub use page_store::{InMemoryBackend, Savepoint};
 pub(crate) use table_tree::{PageListMut, TableTree, TableTreeMut};

src/tree_store/page_store/cached_file.rs

@@ -328,6 +328,10 @@ impl PagedCachedFile {
         &self.write_buffer[stripe]
     }
 
+    pub(super) fn max_cache_size(&self) -> usize {
+        self.max_cache_size
+    }
+
     #[allow(clippy::unused_self)]
     pub(crate) fn cache_stats(&self) -> CacheStats {
         #[cfg(not(feature = "cache_metrics"))]

src/tree_store/page_store/mod.rs

@@ -12,6 +12,7 @@ mod lru_cache;
 mod page_manager;
 mod region;
 mod savepoint;
+mod thread_local_page_cache;
 #[allow(clippy::pedantic, dead_code)]
 mod xxh3;
 
@@ -27,6 +28,10 @@ pub(crate) use page_manager::{
 };
 pub use savepoint::Savepoint;
 pub(crate) use savepoint::SerializedSavepoint;
+pub(crate) use thread_local_page_cache::release_thread_local_page_cache;
+pub(super) use thread_local_page_cache::{
+    CacheTag, ThreadLocalPageCache, with_thread_local_page_cache,
+};
 
 pub(super) use base::{PageImpl, PageMut};
 pub(super) use xxh3::hash128_with_seed;

src/tree_store/page_store/page_manager.rs

@@ -28,6 +28,28 @@ use core::marker::PhantomData;
 use core::mem;
 use core::sync::atomic::{AtomicBool, Ordering};
 
+// Assigns every TransactionalMemory, and every in-place reload of one, an id that no other
+// instance in this process has held, which is what lets a thread local page cache tell whose
+// pages it is holding. The id must be unique process wide, not just within one database:
+// thread local caches outlive the databases that filled them, and every database a thread
+// reads shares one cache, so an id unique only within a database would let one database's
+// pages match another's tag. Ids are never reused -- incrementing a u64 once per instance and
+// per reload cannot wrap in practice -- so a tag from a dropped database matches nothing.
+//
+// A mutex rather than an atomic, because 64 bit atomics are unavailable on some of the targets
+// redb supports, and this is taken only when a database is opened or reloaded.
+static NEXT_CACHE_INSTANCE_ID: Mutex<u64> = Mutex::new(1);
+
+fn next_cache_instance_id() -> u64 {
+    let mut next = NEXT_CACHE_INSTANCE_ID.lock().unwrap();
+    let id = *next;
+    *next += 1;
+    id
+}
+
+// Enough for the top levels of a large btree. 512KiB with 4KiB pages
+const MAX_THREAD_LOCAL_CACHE_PAGES: usize = 128;
+
 // The region header is optional in the v3 file format
 // It's an artifact of the v2 file format, so we initialize new databases without headers to save space
 const NO_HEADER: u32 = 0;
@@ -104,6 +126,14 @@ impl PageResolver {
         self.mem.get_page(page_number, hint)
     }
 
+    pub(crate) fn cache_instance_id(&self) -> u64 {
+        self.mem.cache_instance_id
+    }
+
+    pub(crate) fn thread_local_cache_pages(&self) -> usize {
+        self.mem.thread_local_cache_pages()
+    }
+
     pub(crate) fn count_allocated_pages(&self) -> Result<u64> {
         self.mem.count_allocated_pages()
     }
@@ -500,6 +530,10 @@ impl UnpersistedState {
 }
 
 pub(crate) struct TransactionalMemory {
+    // Identifies this instance's pages in the thread local cache. A transaction id alone is
+    // not enough, because a reload can rewind it and later commits then reuse ids that named
+    // different data
+    cache_instance_id: u64,
     unpersisted: Mutex<UnpersistedState>,
     storage: PagedCachedFile,
     state: Mutex<InMemoryState>,
@@ -650,6 +684,7 @@ impl TransactionalMemory {
         assert!(page_size >= DB_HEADER_SIZE);
 
         Ok(Self {
+            cache_instance_id: next_cache_instance_id(),
             unpersisted: Mutex::new(UnpersistedState::default()),
             storage,
             state: Mutex::new(state),
@@ -741,6 +776,9 @@ impl TransactionalMemory {
         // leave cached pages that disagree with the file.
         self.storage.discard_write_buffer();
         self.storage.invalidate_cache_all();
+        // Invalidate the thread local caches too, since the state being loaded can give a
+        // transaction id already read under this instance different contents
+        self.cache_instance_id = next_cache_instance_id();
         self.storage.sync_file()?;
 
         let header_bytes = self.storage.read_direct(0, DB_HEADER_SIZE)?;
@@ -1668,6 +1706,14 @@ impl TransactionalMemory {
         self.page_size.try_into().unwrap()
     }
 
+    // Pages each thread may cache. Bounded by the configured cache size, so that a small one
+    // shrinks it and set_cache_size(0) disables it, but this is a per thread bound: threads
+    // reading concurrently can hold this many pages each, beyond the configured size
+    fn thread_local_cache_pages(&self) -> usize {
+        let budget = self.storage.max_cache_size() / (self.page_size as usize);
+        budget.min(MAX_THREAD_LOCAL_CACHE_PAGES)
+    }
+
     pub(crate) fn close(&self) -> Result {
         let shutdown_result = self.flush_shutdown_header();
         // The backend's close() contract guarantees it is called exactly once, so it must be