From dc69d5f37efcb93fe7b10d31c565b3e66a21759e Mon Sep 17 00:00:00 2001
From: Vlad Lazar <vlad@neon.tech>
Date: Fri, 2 Feb 2024 13:21:18 +0000
Subject: [PATCH] pageserver_api: allow for merging of keyspaces

---
 Cargo.lock                          |  1 +
 libs/pageserver_api/Cargo.toml      |  1 +
 libs/pageserver_api/src/keyspace.rs | 29 +++++++++++++++++++++++++++++
 3 files changed, 31 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index b2b2777408..26067d4e38 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3520,6 +3520,7 @@ dependencies = [
  "enum-map",
  "hex",
  "humantime-serde",
+ "itertools",
  "postgres_ffi",
  "rand 0.8.5",
  "serde",
diff --git a/libs/pageserver_api/Cargo.toml b/libs/pageserver_api/Cargo.toml
index 902af21965..938910caea 100644
--- a/libs/pageserver_api/Cargo.toml
+++ b/libs/pageserver_api/Cargo.toml
@@ -21,6 +21,7 @@ hex.workspace = true
 thiserror.workspace = true
 humantime-serde.workspace = true
 chrono.workspace = true
+itertools.workspace = true
 
 workspace_hack.workspace = true
 
diff --git a/libs/pageserver_api/src/keyspace.rs b/libs/pageserver_api/src/keyspace.rs
index e04a5a5617..afb872e63f 100644
--- a/libs/pageserver_api/src/keyspace.rs
+++ b/libs/pageserver_api/src/keyspace.rs
@@ -2,6 +2,7 @@ use postgres_ffi::BLCKSZ;
 use std::ops::Range;
 
 use crate::key::Key;
+use itertools::Itertools;
 
 ///
 /// Represents a set of Keys, in a compact form.
@@ -63,6 +64,34 @@ impl KeySpace {
         KeyPartitioning { parts }
     }
 
+    /// Merge another keyspace into the current one.
+    /// Note: the keyspaces must not ovelap (enforced via assertions)
+    pub fn merge(&mut self, other: &KeySpace) {
+        let all_ranges = self
+            .ranges
+            .iter()
+            .merge_by(other.ranges.iter(), |lhs, rhs| lhs.start < rhs.start);
+
+        let mut accum = KeySpaceAccum::new();
+        let mut prev: Option<&Range<Key>> = None;
+        for range in all_ranges {
+            if let Some(prev) = prev {
+                let overlap =
+                    std::cmp::max(range.start, prev.start) < std::cmp::min(range.end, prev.end);
+                assert!(
+                    !overlap,
+                    "Attempt to merge ovelapping keyspaces: {:?} overlaps {:?}",
+                    prev, range
+                );
+            }
+
+            accum.add_range(range.clone());
+            prev = Some(range);
+        }
+
+        self.ranges = accum.to_keyspace().ranges;
+    }
+
     /// Update the keyspace such that it doesn't contain any range
     /// that is overlapping with `other`. This can involve splitting or
     /// removing of existing ranges.