From ec6fb422bd83e64c53a8ebd55c76100698fbd126 Mon Sep 17 00:00:00 2001
From: Vlad Lazar <vlad@neon.tech>
Date: Fri, 2 Feb 2024 14:10:56 +0000
Subject: [PATCH] pageserver: take key space as param for get get vectored

Move the copying up front since the vectored read path will
need the key space in any case.
---
 pageserver/src/basebackup.rs      |  5 +----
 pageserver/src/tenant/timeline.rs | 13 +++++--------
 2 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/pageserver/src/basebackup.rs b/pageserver/src/basebackup.rs
index 7edfab75d4..c862816b80 100644
--- a/pageserver/src/basebackup.rs
+++ b/pageserver/src/basebackup.rs
@@ -261,10 +261,7 @@ where
             let mut slru_builder = SlruSegmentsBuilder::new(&mut self.ar);
 
             for part in slru_partitions.parts {
-                let blocks = self
-                    .timeline
-                    .get_vectored(&part.ranges, self.lsn, self.ctx)
-                    .await?;
+                let blocks = self.timeline.get_vectored(part, self.lsn, self.ctx).await?;
 
                 for (key, block) in blocks {
                     slru_builder.add_block(&key, block?).await?;
diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs
index c8aaeb115d..a1736f0dd8 100644
--- a/pageserver/src/tenant/timeline.rs
+++ b/pageserver/src/tenant/timeline.rs
@@ -14,7 +14,7 @@ use enumset::EnumSet;
 use fail::fail_point;
 use itertools::Itertools;
 use pageserver_api::{
-    keyspace::{key_range_size, KeySpaceAccum},
+    keyspace::KeySpaceAccum,
     models::{
         DownloadRemoteLayersTaskInfo, DownloadRemoteLayersTaskSpawnRequest, EvictionPolicy,
         LayerMapInfo, TimelineState,
@@ -701,7 +701,7 @@ impl Timeline {
     /// which actually vectorizes the read path.
     pub(crate) async fn get_vectored(
         &self,
-        key_ranges: &[Range<Key>],
+        keyspace: KeySpace,
         lsn: Lsn,
         ctx: &RequestContext,
     ) -> Result<BTreeMap<Key, Result<Bytes, PageReconstructError>>, GetVectoredError> {
@@ -709,10 +709,7 @@ impl Timeline {
             return Err(GetVectoredError::InvalidLsn(lsn));
         }
 
-        let key_count = key_ranges
-            .iter()
-            .map(|range| key_range_size(range) as u64)
-            .sum();
+        let key_count = keyspace.total_size().try_into().unwrap();
         if key_count > Timeline::MAX_GET_VECTORED_KEYS {
             return Err(GetVectoredError::Oversized(key_count));
         }
@@ -722,7 +719,7 @@ impl Timeline {
             .map(|t| t.start_timer());
 
         let mut values = BTreeMap::new();
-        for range in key_ranges {
+        for range in keyspace.ranges {
             let mut key = range.start;
             while key != range.end {
                 assert!(!self.shard_identity.is_key_disposable(&key));
@@ -3480,7 +3477,7 @@ impl Timeline {
                         {
                             let results = self
                                 .get_vectored(
-                                    &key_request_accum.consume_keyspace().ranges,
+                                    key_request_accum.consume_keyspace(),
                                     lsn,
                                     ctx,
                                 )