From 493760bb7abe908392fd8a374021f98b6c8dc3d8 Mon Sep 17 00:00:00 2001
From: Vlad Lazar <vlad@neon.tech>
Date: Thu, 25 Jan 2024 10:47:49 +0000
Subject: [PATCH] pageserver: add ranged layer coverage for all intersected
 layers

---
 .../src/tenant/layer_map/layer_coverage.rs    | 39 +++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/pageserver/src/tenant/layer_map/layer_coverage.rs b/pageserver/src/tenant/layer_map/layer_coverage.rs
index 1d9101d3d1..7d830ad006 100644
--- a/pageserver/src/tenant/layer_map/layer_coverage.rs
+++ b/pageserver/src/tenant/layer_map/layer_coverage.rs
@@ -129,6 +129,45 @@ impl<Value: Clone> LayerCoverage<Value> {
             .map(|(k, v)| (*k, v.as_ref().map(|x| x.1.clone())))
     }
 
+    /// Returns an iterator which includes all coverage changes for layers that intersect
+    /// with the provided range.
+    pub fn range_overlaps(
+        &self,
+        key_range: &Range<i128>,
+    ) -> itertools::Either<
+        impl Iterator<Item = (i128, Option<Value>)> + '_,
+        impl Iterator<Item = (i128, Option<Value>)> + '_,
+    >
+    where
+        Value: Eq,
+    {
+        let first_change = self.query(key_range.start);
+        match first_change {
+            Some(change) => {
+                // If the start of the range is covered, we have to deal with two cases:
+                // 1. Start of the range is aligned with the start of a layer.
+                // In this case the return of `self.range` will contain the layer which aligns with the start of the key range.
+                // We advance said iterator to avoid duplicating the first change.
+                // 2. Start of the range is not aligned with the start of a layer.
+                let range = key_range.start..key_range.end;
+                let mut range_coverage = self.range(range).peekable();
+                if range_coverage
+                    .peek()
+                    .is_some_and(|c| c.1 == Some(change.clone()))
+                {
+                    range_coverage.next();
+                }
+                itertools::Either::Left(
+                    std::iter::once((key_range.start, Some(change))).chain(range_coverage),
+                )
+            }
+            None => {
+                let range = key_range.start..key_range.end;
+                let coverage = self.range(range);
+                itertools::Either::Right(coverage)
+            }
+        }
+    }
     /// O(1) clone
     pub fn clone(&self) -> Self {
         Self {