From 06a969295679ffb157667fb4af6e9a502ede69f1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arpad=20M=C3=BCller?= <arpad@neon.tech>
Date: Fri, 27 Oct 2023 05:47:16 +0200
Subject: [PATCH] Adopt ExpCounter

---
 libs/utils/src/exp_counter.rs       | 10 +++++-----
 pageserver/src/pgdatadir_mapping.rs | 26 +++++++++++++++++---------
 2 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/libs/utils/src/exp_counter.rs b/libs/utils/src/exp_counter.rs
index 3c2069f46b..566ce2c3d3 100644
--- a/libs/utils/src/exp_counter.rs
+++ b/libs/utils/src/exp_counter.rs
@@ -23,7 +23,7 @@ pub struct ExpCounter {
 }
 
 impl ExpCounter {
-    pub fn new(max: u64) -> Self {
+    pub fn with_max(max: u64) -> Self {
         Self {
             max,
             base: 0,
@@ -103,28 +103,28 @@ mod test {
     #[test]
     fn to_64() {
         let max = 64;
-        let mut list = ExpCounter::new(max).collect::<Vec<_>>();
+        let mut list = ExpCounter::with_max(max).collect::<Vec<_>>();
         assert_eq!(dupes_and_missing(&mut list, max), (0, 0));
     }
 
     #[test]
     fn to_100() {
         let max = 100;
-        let mut list = ExpCounter::new(max).collect::<Vec<_>>();
+        let mut list = ExpCounter::with_max(max).collect::<Vec<_>>();
         assert_eq!(dupes_and_missing(&mut list, max), (0, 0));
     }
 
     #[test]
     fn to_127() {
         let max = 127;
-        let mut list = ExpCounter::new(max).collect::<Vec<_>>();
+        let mut list = ExpCounter::with_max(max).collect::<Vec<_>>();
         assert_eq!(dupes_and_missing(&mut list, max), (0, 0));
     }
 
     #[test]
     fn to_12345() {
         let max = 12345;
-        let mut list = ExpCounter::new(max).collect::<Vec<_>>();
+        let mut list = ExpCounter::with_max(max).collect::<Vec<_>>();
         assert_eq!(dupes_and_missing(&mut list, max), (0, 0));
     }
 }
diff --git a/pageserver/src/pgdatadir_mapping.rs b/pageserver/src/pgdatadir_mapping.rs
index 1b010abb46..89739b4e49 100644
--- a/pageserver/src/pgdatadir_mapping.rs
+++ b/pageserver/src/pgdatadir_mapping.rs
@@ -23,6 +23,7 @@ use std::ops::ControlFlow;
 use std::ops::Range;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, trace, warn};
+use utils::exp_counter::ExpCounter;
 use utils::{bin_ser::BeSer, lsn::Lsn};
 
 /// Block number within a relation or SLRU. This matches PostgreSQL's BlockNumber type.
@@ -342,7 +343,7 @@ impl Timeline {
         let mut found_larger = false;
 
         // This is a search budget to not make the search too expensive
-        let mut budget = 1000u16;
+        let mut budget = 10_000u16;
         while low < high {
             // cannot overflow, high and low are both smaller than u64::MAX / 2
             // this always holds: low <= mid_start < high
@@ -351,23 +352,30 @@ impl Timeline {
 
             let mut max = None;
 
-            // Linear search for an lsn that has a commit timestamp.
+            // Search for an lsn that has a commit timestamp.
+            // The search with `ExpCounter` will eventually try all LSNs
+            // in the range between mid_start and high, which is important
+            // when we set high to mid if we have exhausted all possibilities.
             // We don't do an explosive search as we want to prove that
             // every single lsn up to high is giving inconclusive results.
             // This can only be done by trying all lsns.
-            while max.is_none() && mid < high {
+            for offs in ExpCounter::with_max(high - mid_start) {
+                mid = offs + mid_start;
+
+                // Do the query for mid + 1 instead of mid so that we can make definite statements
+                // about low (low's invariant: always points to commit before or at search_timestamp).
+                max = self
+                    .get_max_timestamp_for_lsn(Lsn((mid + 1) * 8), ctx)
+                    .await?;
+                if max.is_some() {
+                    break;
+                }
                 // Do some limiting to make sure the query does not become too expensive.
                 if let Some(new_budget) = budget.checked_sub(1) {
                     budget = new_budget;
                 } else {
                     return Ok(LsnForTimestamp::NoData(min_lsn));
                 }
-                // Do the query for mid + 1 instead of mid so that we can make definite statements
-                // about low (low's invariant: always points to commit before or at search_timestamp).
-                max = self
-                    .get_max_timestamp_for_lsn(Lsn((mid + 1) * 8), ctx)
-                    .await?;
-                mid += 1;
             }
 
             if let Some(max) = max {