Use im

Cargo.lock

@@ -635,6 +635,15 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
+[[package]]
+name = "bitmaps"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "031043d04099746d8db04daf1fa424b2bc8bd69d92b25962dcde24da39ab64a2"
+dependencies = [
+ "typenum",
+]
+
 [[package]]
 name = "block-buffer"
 version = "0.10.3"
@@ -1826,6 +1835,20 @@ dependencies = [
  "unicode-normalization",
 ]
 
+[[package]]
+name = "im"
+version = "15.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d0acd33ff0285af998aaf9b57342af478078f53492322fafc47450e09397e0e9"
+dependencies = [
+ "bitmaps",
+ "rand_core",
+ "rand_xoshiro",
+ "sized-chunks",
+ "typenum",
+ "version_check",
+]
+
 [[package]]
 name = "indexmap"
 version = "1.9.1"
@@ -2337,6 +2360,7 @@ dependencies = [
  "humantime",
  "humantime-serde",
  "hyper",
+ "im",
  "itertools",
  "metrics",
  "nix 0.25.0",
@@ -2963,6 +2987,15 @@ dependencies = [
  "rand_core",
 ]
 
+[[package]]
+name = "rand_xoshiro"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa"
+dependencies = [
+ "rand_core",
+]
+
 [[package]]
 name = "rayon"
 version = "1.5.3"
@@ -3597,6 +3630,16 @@ version = "0.3.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7bd3e3206899af3f8b12af284fafc038cc1dc2b41d1b89dd17297221c5d225de"
 
+[[package]]
+name = "sized-chunks"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "16d69225bde7a69b235da73377861095455d298f2b970996eec25ddbb42b3d1e"
+dependencies = [
+ "bitmaps",
+ "typenum",
+]
+
 [[package]]
 name = "slab"
 version = "0.4.7"

pageserver/Cargo.toml

@@ -70,6 +70,7 @@ tenant_size_model = { path = "../libs/tenant_size_model" }
 utils = { path = "../libs/utils" }
 workspace_hack = { version = "0.1", path = "../workspace_hack" }
 rpds = "0.12.0"
+im = "15.1.0"
 
 [dev-dependencies]
 criterion = "0.4"

pageserver/src/tenant/bst_layer_map.rs

@@ -1,10 +1,9 @@
 use std::collections::BTreeMap;
 use std::ops::Range;
-use std::sync::Arc;
 
-// TODO the `im` crate has 20x more downloads and also has
-// persistent/immutable BTree. See if it's better.
+// TODO drop rpds. So far `im` looks 30% faster.
 use rpds::RedBlackTreeMapSync;
+use im::OrdMap;
 
 /// Layer map implemented using persistent/immutable binary search tree.
 /// It supports historical queries, but no retroactive inserts. For that
@@ -20,12 +19,12 @@ pub struct PersistentLayerMap<Value> {
     /// TODO Merge historic with retroactive, into HistoricLayerMap
     /// TODO Maintain a pair of heads, one for images, one for deltas.
     ///      This way we can query both of them with one BTreeMap query.
-    head: RedBlackTreeMapSync<i128, Option<(u64, Value)>>,
+    head: OrdMap<i128, Option<(u64, Value)>>,
 
     /// All previous states of `self.head`
     ///
     /// TODO: Sorted Vec + binary search could be slightly faster.
-    historic: BTreeMap<u64, RedBlackTreeMapSync<i128, Option<(u64, Value)>>>,
+    historic: BTreeMap<u64, OrdMap<i128, Option<(u64, Value)>>>,
 }
 
 impl<Value: std::fmt::Debug> std::fmt::Debug for PersistentLayerMap<Value> {
@@ -44,7 +43,7 @@ impl<T: Clone> Default for PersistentLayerMap<T> {
 impl<Value: Clone> PersistentLayerMap<Value> {
     pub fn new() -> Self {
         Self {
-            head: RedBlackTreeMapSync::default(),
+            head: OrdMap::default(),
             historic: BTreeMap::default(),
         }
     }
@@ -56,7 +55,7 @@ impl<Value: Clone> PersistentLayerMap<Value> {
             Some((_, None)) => None,
             None => None,
         };
-        self.head.insert_mut(key, value);
+        self.head.insert(key, value);
     }
 
     pub fn insert(self: &mut Self, key: Range<i128>, lsn: Range<u64>, value: Value) {
@@ -104,10 +103,10 @@ impl<Value: Clone> PersistentLayerMap<Value> {
         }
         for k in to_update {
             self.head
-                .insert_mut(k.clone(), Some((lsn.end.clone(), value.clone())));
+                .insert(k.clone(), Some((lsn.end.clone(), value.clone())));
         }
         for k in to_remove {
-            self.head.remove_mut(&k);
+            self.head.remove(&k);
         }
 
         // Remember history. Clone is O(1)
@@ -117,9 +116,11 @@ impl<Value: Clone> PersistentLayerMap<Value> {
     pub fn query(self: &Self, key: i128, lsn: u64) -> Option<Value> {
         let version = self.historic.range(0..=lsn).rev().next()?.1;
         version
-            .range(0..=key)
-            .rev()
-            .next()?
+            .get_prev(&key)?
+            // .range(0..=key).rev().next()?
+            // NOTE The canonical way to do this in other crates is
+            //      `.range(0..=key).rev.next()` and `im` supports this
+            //      API but it's 2x slower than `.get_prev(&key)`.
             .1
             .as_ref()
             .map(|(_, v)| v.clone())

pageserver/src/tenant/layer_map.rs

@@ -248,19 +248,19 @@ impl LayerMap {
     /// layer.
     ///
     pub fn search(&self, key: Key, end_lsn: Lsn) -> Result<Option<SearchResult>> {
-        // let old = self.search_old(key, end_lsn)?;
+        let old = self.search_old(key, end_lsn)?;
         let new = self.search_new(key, end_lsn)?;
-        // match (&old, &new) {
-        //     (None, None) => {}
-        //     (None, Some(_)) => panic!("returned Some, expected None"),
-        //     (Some(_), None) => panic!("returned None, expected Some"),
-        //     (Some(old), Some(new)) => {
-        //         // TODO be more verbose and flexible
-        //         let context = format!("query: key {}, end_lsn: {}", key, end_lsn);
-        //         assert_eq!(old.layer.filename(), new.layer.filename(), "{}", context);
-        //         assert_eq!(old.lsn_floor, new.lsn_floor, "{}", context);
-        //     }
-        // }
+        match (&old, &new) {
+            (None, None) => {}
+            (None, Some(_)) => panic!("returned Some, expected None"),
+            (Some(_), None) => panic!("returned None, expected Some"),
+            (Some(old), Some(new)) => {
+                // TODO be more verbose and flexible
+                let context = format!("query: key {}, end_lsn: {}", key, end_lsn);
+                assert_eq!(old.layer.filename(), new.layer.filename(), "{}", context);
+                assert_eq!(old.lsn_floor, new.lsn_floor, "{}", context);
+            }
+        }
         return Ok(new);
     }