From 2737e3356522c9210f646562473d024c318876ae Mon Sep 17 00:00:00 2001
From: Bojan Serafimov <bojan.serafimov7@gmail.com>
Date: Tue, 27 Dec 2022 15:10:56 -0500
Subject: [PATCH] Implement image_coverage (untested)

---
 pageserver/src/repository.rs           | 11 ++++++
 pageserver/src/tenant/bst_layer_map.rs | 18 ++++++++++
 pageserver/src/tenant/layer_map.rs     | 47 +++++++++++++++++++++++---
 3 files changed, 72 insertions(+), 4 deletions(-)

diff --git a/pageserver/src/repository.rs b/pageserver/src/repository.rs
index 586fd20886..092503b7c5 100644
--- a/pageserver/src/repository.rs
+++ b/pageserver/src/repository.rs
@@ -37,6 +37,17 @@ impl Key {
             | self.field6 as i128
     }
 
+    pub fn from_i128(x: i128) -> Self {
+        Key {
+            field1: ((x >> 120) & 0xf) as u8,
+            field2: ((x >> 104) & 0xFFFF) as u32,
+            field3: (x >> 72) as u32,
+            field4: (x >> 40) as u32,
+            field5: (x >> 32) as u8,
+            field6: x as u32,
+        }
+    }
+
     pub fn next(&self) -> Key {
         self.add(1)
     }
diff --git a/pageserver/src/tenant/bst_layer_map.rs b/pageserver/src/tenant/bst_layer_map.rs
index b9259e14c2..64cd71f3cc 100644
--- a/pageserver/src/tenant/bst_layer_map.rs
+++ b/pageserver/src/tenant/bst_layer_map.rs
@@ -125,6 +125,13 @@ impl<Value: Clone> PersistentLayerMap<Value> {
             .map(|(_, v)| v.clone())
     }
 
+    pub fn get_coverage(
+        self: &Self,
+        lsn: u64,
+    ) -> Option<&RedBlackTreeMapSync<i128, Option<(u64, Value)>>> {
+        Some(self.historic.range(..=lsn).rev().next()?.1)
+    }
+
     pub fn trim(self: &mut Self, begin: &u64) {
         self.historic.split_off(begin);
         self.head = self
@@ -371,6 +378,17 @@ impl<Value: Clone> RetroactiveLayerMap<Value> {
 
         self.map.query(key, lsn)
     }
+
+    pub fn get_coverage(
+        self: &Self,
+        lsn: u64,
+    ) -> Option<&RedBlackTreeMapSync<i128, Option<(u64, Value)>>> {
+        if !self.buffer.is_empty() {
+            panic!("rebuild pls")
+        }
+
+        self.map.get_coverage(lsn)
+    }
 }
 
 #[test]
diff --git a/pageserver/src/tenant/layer_map.rs b/pageserver/src/tenant/layer_map.rs
index d273f8cc80..fb1b67a27e 100644
--- a/pageserver/src/tenant/layer_map.rs
+++ b/pageserver/src/tenant/layer_map.rs
@@ -298,10 +298,7 @@ impl LayerMap {
                     lsn_floor = std::cmp::max(lsn_floor, image.get_lsn_range().start + 1)
                 }
             }
-            SearchResult {
-                layer,
-                lsn_floor,
-            }
+            SearchResult { layer, lsn_floor }
         }));
     }
 
@@ -483,6 +480,7 @@ impl LayerMap {
         key_range: &Range<Key>,
         lsn_range: &Range<Lsn>,
     ) -> Result<bool> {
+        // TODO implement using new index
         let mut range_remain = key_range.clone();
 
         loop {
@@ -524,12 +522,18 @@ impl LayerMap {
     }
 
     pub fn iter_historic_layers(&self) -> impl '_ + Iterator<Item = Arc<dyn Layer>> {
+        // TODO implement using new index
         self.historic_layers.iter().map(|e| e.layer.clone())
     }
 
     /// Find the last image layer that covers 'key', ignoring any image layers
     /// newer than 'lsn'.
     fn find_latest_image(&self, key: Key, lsn: Lsn) -> Option<Arc<dyn Layer>> {
+        let use_new_method = true;
+        if use_new_method {
+            return self.images.query(key.to_i128(), lsn.0);
+        }
+
         let mut candidate_lsn = Lsn(0);
         let mut candidate = None;
         let envelope = AABB::from_corners(
@@ -572,6 +576,40 @@ impl LayerMap {
         key_range: &Range<Key>,
         lsn: Lsn,
     ) -> Result<Vec<(Range<Key>, Option<Arc<dyn Layer>>)>> {
+        let use_new_method = true;
+        if use_new_method {
+            let bounds = match self.images.get_coverage(lsn.0) {
+                Some(x) => x,
+                None => return Ok(vec![]),
+            };
+
+            let start = key_range.start.to_i128();
+            let end = key_range.end.to_i128();
+
+            // Initialize loop variables
+            let mut coverage: Vec<(Range<Key>, Option<Arc<dyn Layer>>)> = vec![];
+            let mut current_key = start.clone();
+            let mut current_val = match bounds.range(..=start).rev().next() {
+                Some((_, Some((_, v)))) => Some(v.clone()),
+                Some((_, None)) => None,
+                None => None,
+            };
+
+            // Loop through the change events and push intervals
+            for (change_key, change_val) in bounds.range(start..end) {
+                let kr = Key::from_i128(current_key)..Key::from_i128(*change_key);
+                coverage.push((kr, current_val.take()));
+                current_key = change_key.clone();
+                current_val = change_val.as_ref().map(|(_, v)| v.clone());
+            }
+
+            // Add the final interval
+            let kr = Key::from_i128(current_key)..Key::from_i128(end);
+            coverage.push((kr, current_val.take()));
+
+            return Ok(coverage);
+        }
+
         let mut points = vec![key_range.start];
         let envelope = AABB::from_corners(
             [IntKey::from(key_range.start.to_i128()), IntKey::from(0)],
@@ -617,6 +655,7 @@ impl LayerMap {
     /// Count how many L1 delta layers there are that overlap with the
     /// given key and LSN range.
     pub fn count_deltas(&self, key_range: &Range<Key>, lsn_range: &Range<Lsn>) -> Result<usize> {
+        // TODO implement using new index
         let mut result = 0;
         if lsn_range.start >= lsn_range.end {
             return Ok(0);