From fb6569c8808b947873a6349642cf7785a12cf0d0 Mon Sep 17 00:00:00 2001
From: Bojan Serafimov <bojan.serafimov7@gmail.com>
Date: Thu, 5 Jan 2023 02:50:48 -0500
Subject: [PATCH] Implement get_difficulty_map

---
 pageserver/src/tenant/coverage.rs         |  6 ++++
 pageserver/src/tenant/latest_layer_map.rs |  4 +++
 pageserver/src/tenant/layer_map.rs        | 34 +++++++++++++++++++++--
 3 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/pageserver/src/tenant/coverage.rs b/pageserver/src/tenant/coverage.rs
index cc56b31bfe..20a45a9880 100644
--- a/pageserver/src/tenant/coverage.rs
+++ b/pageserver/src/tenant/coverage.rs
@@ -96,6 +96,12 @@ impl<Value: Clone> Coverage<Value> {
             .map(|(k, v)| (k.clone(), v.as_ref().map(|x| x.1.clone())))
     }
 
+    pub fn iter(self: &Self) -> impl '_ + Iterator<Item = (i128, Option<Value>)> {
+        self.head
+            .iter()
+            .map(|(k, v)| (k.clone(), v.as_ref().map(|x| x.1.clone())))
+    }
+
     pub fn clone(self: &Self) -> Self {
         Self {
             head: self.head.clone(),
diff --git a/pageserver/src/tenant/latest_layer_map.rs b/pageserver/src/tenant/latest_layer_map.rs
index b67f7ea3bf..2e3c2850f9 100644
--- a/pageserver/src/tenant/latest_layer_map.rs
+++ b/pageserver/src/tenant/latest_layer_map.rs
@@ -44,6 +44,10 @@ impl<Value: Clone> LatestLayerMap<Value> {
         self.image_coverage.range(key)
     }
 
+    pub fn image_iter(self: &Self) -> impl '_ + Iterator<Item = (i128, Option<Value>)> {
+        self.image_coverage.iter()
+    }
+
     pub fn delta_coverage(
         self: &Self,
         key: Range<i128>,
diff --git a/pageserver/src/tenant/layer_map.rs b/pageserver/src/tenant/layer_map.rs
index 63a6019c60..50653820fd 100644
--- a/pageserver/src/tenant/layer_map.rs
+++ b/pageserver/src/tenant/layer_map.rs
@@ -329,8 +329,38 @@ where
     /// result for the entire partitioning at once allows this function to be more
     /// efficient, and further optimization is possible by using iterators instead,
     /// to allow early return.
-    pub fn get_difficulty_map(&self, _lsn: Lsn, _partitioning: &KeyPartitioning) -> Vec<usize> {
-        todo!()
+    pub fn get_difficulty_map(&self, lsn: Lsn, partitioning: &KeyPartitioning) -> Vec<usize> {
+        // TODO This is a naive implementation. Perf improvements to do:
+        // 1. Instead of calling self.image_coverage and self.count_deltas,
+        //    iterate the image and delta coverage only once.
+        // 2. Implement early return when the difficulty exceeds a threshold.
+        partitioning
+            .parts
+            .iter()
+            .map(|part| {
+                let mut difficulty = 0;
+                for range in &part.ranges {
+                    for (img_range, last_img) in self
+                        .image_coverage(range, lsn)
+                        .expect("why would this err?")
+                    {
+                        let img_lsn = if let Some(last_img) = last_img {
+                            last_img.get_lsn_range().end
+                        } else {
+                            Lsn(0)
+                        };
+
+                        if img_lsn < lsn {
+                            let num_deltas = self
+                                .count_deltas(&img_range, &(img_lsn..lsn))
+                                .expect("why would this err lol?");
+                            difficulty = std::cmp::max(difficulty, num_deltas);
+                        }
+                    }
+                }
+                difficulty
+            })
+            .collect()
     }
 
     /// Return all L0 delta layers