From e3307851d78f2fd149ed64bca940af5966810736 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Wed, 4 Sep 2024 18:48:48 +0000
Subject: [PATCH] [DNM] benchmark: add ability to force a walredo roundtrip
 even if in-proces is possible

---
 pageserver/src/tenant/timeline.rs | 10 ++++++++++
 pageserver/src/walredo.rs         | 18 +++++++++++++++---
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs
index 6eadf9a564..a4f7906762 100644
--- a/pageserver/src/tenant/timeline.rs
+++ b/pageserver/src/tenant/timeline.rs
@@ -5085,6 +5085,16 @@ impl Timeline {
                     img_lsn,
                     request_lsn,
                 );
+                if *crate::walredo::ADD_PROCESS_PING_TO_IN_PROCESS {
+                    self.walredo_mgr
+                        .as_ref()
+                        .context("timeline has no walredo manager")
+                        .map_err(PageReconstructError::WalRedo)?
+                        .ping(self.pg_version)
+                        .await
+                        .unwrap();
+                }
+
                 Ok(img.clone())
             } else {
                 Err(PageReconstructError::from(anyhow!(
diff --git a/pageserver/src/walredo.rs b/pageserver/src/walredo.rs
index 632dbf37ee..aaac7f0bd2 100644
--- a/pageserver/src/walredo.rs
+++ b/pageserver/src/walredo.rs
@@ -171,7 +171,8 @@ impl PostgresRedoManager {
 
             if rec_neon != batch_neon {
                 let result = if batch_neon {
-                    self.apply_batch_neon(key, lsn, img, &records[batch_start..i])
+                    self.apply_batch_neon(key, lsn, img, &records[batch_start..i], pg_version)
+                        .await
                 } else {
                     self.apply_batch_postgres(
                         key,
@@ -192,7 +193,8 @@ impl PostgresRedoManager {
         }
         // last batch
         if batch_neon {
-            self.apply_batch_neon(key, lsn, img, &records[batch_start..])
+            self.apply_batch_neon(key, lsn, img, &records[batch_start..], pg_version)
+                .await
         } else {
             self.apply_batch_postgres(
                 key,
@@ -494,12 +496,13 @@ impl PostgresRedoManager {
     ///
     /// Process a batch of WAL records using bespoken Neon code.
     ///
-    fn apply_batch_neon(
+    async fn apply_batch_neon(
         &self,
         key: Key,
         lsn: Lsn,
         base_img: Option<Bytes>,
         records: &[(Lsn, NeonWalRecord)],
+        pg_version: u32,
     ) -> Result<Bytes, Error> {
         let start_time = Instant::now();
 
@@ -529,6 +532,10 @@ impl PostgresRedoManager {
             lsn
         );
 
+        if *ADD_PROCESS_PING_TO_IN_PROCESS {
+            self.ping(pg_version).await.unwrap();
+        }
+
         Ok(page.freeze())
     }
 
@@ -545,6 +552,11 @@ impl PostgresRedoManager {
     }
 }
 
+// add a walredo roundtrip to simulate latency overhead
+pub(crate) static ADD_PROCESS_PING_TO_IN_PROCESS: Lazy<bool> = Lazy::new(|| {
+    utils::env::var("NEON_PAGESERVER_WALREDO_ADD_PROCESS_PING_TO_IN_PROCESS").unwrap_or(false)
+});
+
 #[cfg(test)]
 mod tests {
     use super::PostgresRedoManager;