From 924d91c47de8ed6588bcbf6124bba235fa252fcd Mon Sep 17 00:00:00 2001
From: Bojan Serafimov <bojan.serafimov7@gmail.com>
Date: Thu, 1 Dec 2022 16:42:24 -0500
Subject: [PATCH] Bench layer map using persistent range queries

---
 pageserver/benches/bench_layer_map.rs         | 37 ++++++++++++-----
 .../src/tenant/segment_tree_layer_map.rs      | 41 ++++++++++++++++---
 2 files changed, 62 insertions(+), 16 deletions(-)

diff --git a/pageserver/benches/bench_layer_map.rs b/pageserver/benches/bench_layer_map.rs
index 760626c817..33a0fe1fe6 100644
--- a/pageserver/benches/bench_layer_map.rs
+++ b/pageserver/benches/bench_layer_map.rs
@@ -280,32 +280,49 @@ fn bench_sequential(c: &mut Criterion) {
             lsn: Lsn(i),
         };
         layer_map.insert_historic(Arc::new(layer));
-        stlm.insert(10 * i32, 10 * i32 + 1);
     }
+    println!("Finished layer map init in {:?}", now.elapsed());
 
-    // Manually measure runtime without criterion because criterion
-    // has a minimum sample size of 10 and I don't want to run it 10 times.
-    println!("Finished init in {:?}", now.elapsed());
+    let now = Instant::now();
+    for i in 0..100_000 {
+        // TODO try inserting a super-wide layer in between every 10 to reflect
+        //      what often happens with L1 layers that include non-rel changes.
+        //      Maybe do that as a separate test.
+        let i32 = (i as u32) % 100;
+        let zero = Key::from_hex("000000000000000000000000000000000000").unwrap();
+        let layer = DummyImage {
+            key_range: zero.add(10 * i32)..zero.add(10 * i32 + 1),
+            lsn: Lsn(i),
+        };
+        stlm.insert(10 * i32, 10 * i32 + 1, i as u32, format!("Layer {}", i));
+    }
+    println!("Finished persistent segment tree init in {:?}", now.elapsed());
 
     // Choose 100 uniformly random queries
     let rng = &mut StdRng::seed_from_u64(1);
     let queries: Vec<(Key, Lsn)> = uniform_query_pattern(&layer_map)
-        .choose_multiple(rng, 1)
+        .choose_multiple(rng, 100)
         .copied()
         .collect();
 
     // Define and name the benchmark function
-    c.bench_function("sequential_uniform_queries", |b| {
-        // Run the search queries
+    let mut group = c.benchmark_group("sequential_uniform_queries");
+    group.bench_function("current_code", |b| {
         b.iter(|| {
             for q in queries.clone().into_iter() {
                 layer_map.search(q.0, q.1).unwrap();
             }
         });
     });
+    group.bench_function("persistent_segment_tree", |b| {
+        b.iter(|| {
+            for q in queries.clone().into_iter() {
+                stlm.query(q.0.field6, q.1.0 as u32);
+            }
+        });
+    });
+    group.finish();
 }
 
-criterion_group!(group_1, bench_from_captest_env);
-criterion_group!(group_2, bench_from_real_project);
 criterion_group!(group_3, bench_sequential);
-criterion_main!(group_1, group_2, group_3);
+criterion_main!(group_3);
diff --git a/pageserver/src/tenant/segment_tree_layer_map.rs b/pageserver/src/tenant/segment_tree_layer_map.rs
index 558962a9f2..f24abc577f 100644
--- a/pageserver/src/tenant/segment_tree_layer_map.rs
+++ b/pageserver/src/tenant/segment_tree_layer_map.rs
@@ -1,11 +1,12 @@
 use persistent_range_query::naive::{IndexableKey, NaiveVecStorage};
 use persistent_range_query::ops::SameElementsInitializer;
-use persistent_range_query::segment_tree::{MidpointableKey, PersistentSegmentTree};
+use persistent_range_query::segment_tree::{MidpointableKey, PersistentSegmentTree, PersistentSegmentTreeVersion};
 use persistent_range_query::{
     LazyRangeInitializer, PersistentVecStorage, RangeModification, RangeQueryResult,
     VecReadableVersion,
 };
 use std::cmp::Ordering;
+use std::collections::BTreeMap;
 use std::ops::Range;
 
 #[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
@@ -194,26 +195,54 @@ impl LazyRangeInitializer<LayerMapInformation, PageIndex> for SameElementsInitia
     }
 }
 
+type Head = PersistentSegmentTree<LayerMapModification, SameElementsInitializer<()>, PageIndex>;
+type Version = PersistentSegmentTreeVersion<LayerMapModification, SameElementsInitializer<()>, PageIndex>;
+
 pub struct STLM {
-    s: PersistentVecStorage<LayerMapModification, SameElementsInitializer<()>, PageIndex>,
+    head: Head,
+    historic: BTreeMap<u32, Version>,
 }
 
+/// Layer map (good enough for benchmarks) implemented using persistent segment tree
 impl STLM {
     pub fn new() -> Self {
         STLM {
-            s: PersistentVecStorage::new(
+            head: PersistentSegmentTree::new(
                 PageIndex(0)..PageIndex(100),
                 SameElementsInitializer::new(()),
             ),
+            historic: BTreeMap::default(),
         }
     }
 
-    pub fn insert(key_begin: i32, key_end: i32) {
-        s.modify(
+    pub fn insert(self: &mut Self, key_begin: u32, key_end: u32, lsn: u32, value: String) {
+        self.head.modify(
             &(PageIndex(key_begin)..PageIndex(key_end)),
-            &LayerMapModification::add_image_layer("Img0..70"),
+            &LayerMapModification::add_image_layer(value),
         );
+        self.historic.insert(lsn, self.head.freeze());
     }
+
+    pub fn query(self: &Self, key: u32, lsn: u32) -> Option<String> {
+        let version = self.historic.range(0..=lsn).rev().next()?.1;
+        let info = version.get(&(PageIndex(key)..PageIndex(key + 1)));
+        info.last_image_layer.map(|s| s.clone())
+    }
+}
+
+fn test_stlm() {
+    let mut stlm = STLM::new();
+    stlm.insert(0, 5, 100, "layer 1".to_string());
+    stlm.insert(3, 9, 110, "layer 2".to_string());
+
+    dbg!(stlm.query(1, 105));
+    dbg!(stlm.query(4, 105));
+    dbg!(stlm.query(4, 115));
+}
+
+#[test]
+fn test_stlm_() {
+    test_stlm()
 }
 
 fn test_layer_map<