Add assertion

Cargo.lock

@@ -601,6 +601,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"
@@ -1728,6 +1737,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.2"
@@ -2210,6 +2233,7 @@ dependencies = [
  "humantime",
  "humantime-serde",
  "hyper",
+ "im",
  "itertools",
  "metrics",
  "nix",
@@ -2728,6 +2752,15 @@ dependencies = [
  "getrandom",
 ]
 
+[[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.6.1"
@@ -3394,6 +3427,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

@@ -67,6 +67,7 @@ utils.workspace = true
 workspace_hack.workspace = true
 reqwest.workspace = true
 rpds.workspace = true
+im = "15.1.0"
 
 [dev-dependencies]
 criterion.workspace = true

pageserver/src/tenant/layer_map/layer_coverage.rs

@@ -1,8 +1,6 @@
 use std::ops::Range;
 
-// TODO the `im` crate has 20x more downloads and also has
-// persistent/immutable BTree. It also runs a bit faster but
-// results are not the same on some tests.
+use im::OrdMap;
 use rpds::RedBlackTreeMapSync;
 
 /// Data structure that can efficiently:
@@ -24,6 +22,7 @@ pub struct LayerCoverage<Value> {
     /// be Sync. Using nonsync might be faster, if we can work with
     /// that.
     nodes: RedBlackTreeMapSync<i128, Option<(u64, Value)>>,
+    im_nodes: OrdMap<i128, Option<(u64, Value)>>,
 }
 
 impl<T: Clone> Default for LayerCoverage<T> {
@@ -36,6 +35,7 @@ impl<Value: Clone> LayerCoverage<Value> {
     pub fn new() -> Self {
         Self {
             nodes: RedBlackTreeMapSync::default(),
+            im_nodes: OrdMap::default(),
         }
     }
 
@@ -49,6 +49,14 @@ impl<Value: Clone> LayerCoverage<Value> {
             None => None,
         };
         self.nodes.insert_mut(key, value);
+
+        let im_value = match self.im_nodes.range(..=key).last() {
+            Some((_, Some(v))) => Some(v.clone()),
+            Some((_, None)) => None,
+            None => None,
+        };
+        self.im_nodes.remove(&key);
+        self.im_nodes.insert(key, im_value);
     }
 
     /// Insert a layer.
@@ -96,19 +104,85 @@ impl<Value: Clone> LayerCoverage<Value> {
         for k in to_remove {
             self.nodes.remove_mut(&k);
         }
+
+
+        let mut to_update = Vec::new();
+        let mut to_remove = Vec::new();
+        let mut prev_covered = false;
+        for (k, node) in self.im_nodes.range(key.clone()) {
+            let needs_cover = match node {
+                None => true,
+                Some((h, _)) => h < &lsn.end,
+            };
+            if needs_cover {
+                match prev_covered {
+                    true => to_remove.push(*k),
+                    false => to_update.push(*k),
+                }
+            }
+            prev_covered = needs_cover;
+        }
+        if !prev_covered {
+            to_remove.push(key.end);
+        }
+        for k in to_update {
+            self.im_nodes.remove(&k);
+            self.im_nodes.insert(k, Some((lsn.end, value.clone())));
+        }
+        for k in to_remove {
+            self.im_nodes.remove(&k);
+        }
+
+    }
+
+    fn get_key_1(&self, key: i128) -> Option<u64> {
+        self.im_nodes
+            .get_prev(&key)?
+            .1
+            .as_ref()
+            .map(|(k, _)| k.clone())
+    }
+    fn get_key_2(&self, key: i128) -> Option<u64> {
+        self.im_nodes
+            .range(..=key)
+            .rev()
+            .next()?
+            .1
+            .as_ref()
+            .map(|(k, _)| k.clone())
     }
 
     /// Get the latest (by lsn.end) layer at a given key
     ///
     /// Complexity: O(log N)
     pub fn query(&self, key: i128) -> Option<Value> {
-        self.nodes
+
+        let k1 = self.get_key_1(key);
+        let k2 = self.get_key_2(key);
+        assert_eq!(k1, k2);
+
+
+        // self.im_nodes
+        //     .get_prev(&key)?
+        //     .1
+        //     .as_ref()
+        //     .map(|(_, v)| v.clone())
+
+        self.im_nodes
             .range(..=key)
             .rev()
             .next()?
             .1
             .as_ref()
             .map(|(_, v)| v.clone())
+
+        // self.nodes
+        //     .range(..=key)
+        //     .rev()
+        //     .next()?
+        //     .1
+        //     .as_ref()
+        //     .map(|(_, v)| v.clone())
     }
 
     /// Iterate the changes in layer coverage in a given range. You will likely
@@ -125,6 +199,7 @@ impl<Value: Clone> LayerCoverage<Value> {
     pub fn clone(&self) -> Self {
         Self {
             nodes: self.nodes.clone(),
+            im_nodes: self.im_nodes.clone(),
         }
     }
 }