From 0bb037240d7a76f9ad74c38e806dbea09ddffbd2 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Mon, 25 Nov 2024 10:30:07 +0100
Subject: [PATCH] logging to debug test_pageserver_restarts_under_worload

---
 pageserver/src/lib.rs          | 19 +++++++++++++++++++
 pageserver/src/page_service.rs | 23 ++++++++++++++++++++---
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/pageserver/src/lib.rs b/pageserver/src/lib.rs
index ef6711397a..ff6af3566c 100644
--- a/pageserver/src/lib.rs
+++ b/pageserver/src/lib.rs
@@ -356,6 +356,25 @@ async fn timed<Fut: std::future::Future>(
     }
 }
 
+/// Like [`timed`], but the warning timeout only starts after `cancel` has been cancelled.
+async fn timed_after_cancellation<Fut: std::future::Future>(
+    fut: Fut,
+    name: &str,
+    warn_at: std::time::Duration,
+    cancel: &CancellationToken,
+) -> <Fut as std::future::Future>::Output {
+    let mut fut = std::pin::pin!(fut);
+
+    tokio::select! {
+        _ = cancel.cancelled() => {
+            timed(fut, name, warn_at).await
+        }
+        ret = &mut fut => {
+            ret
+        }
+    }
+}
+
 #[cfg(test)]
 mod timed_tests {
     use super::timed;
diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs
index 88f636c7da..74dd7eddaf 100644
--- a/pageserver/src/page_service.rs
+++ b/pageserver/src/page_service.rs
@@ -44,7 +44,6 @@ use utils::{
 };
 
 use crate::auth::check_permission;
-use crate::basebackup;
 use crate::basebackup::BasebackupError;
 use crate::config::PageServerConf;
 use crate::context::{DownloadBehavior, RequestContext};
@@ -62,6 +61,7 @@ use crate::tenant::timeline::{self, WaitLsnError};
 use crate::tenant::GetTimelineError;
 use crate::tenant::PageReconstructError;
 use crate::tenant::Timeline;
+use crate::{basebackup, timed_after_cancellation};
 use pageserver_api::key::rel_block_to_key;
 use pageserver_api::reltag::{BlockNumber, RelTag, SlruKind};
 use postgres_ffi::pg_constants::DEFAULTTABLESPACE_OID;
@@ -1101,6 +1101,8 @@ impl PageServerHandler {
             protocol_pipelining_mode,
         } = pipelining_config;
 
+        let cancel = self.cancel.clone();
+
         let (requests_tx, mut requests_rx) = tokio::sync::mpsc::channel(1);
         let read_messages = {
             let cancel = self.cancel.child_token();
@@ -1244,8 +1246,23 @@ impl PageServerHandler {
         let executor_res;
         match protocol_pipelining_mode {
             PageServiceProtocolPipeliningMode::ConcurrentFutures => {
-                (read_messages_res, _, executor_res) =
-                    tokio::join!(read_messages, batcher, executor);
+                (read_messages_res, _, executor_res) = {
+                    macro_rules! timed {
+                        ($fut:expr, $what:literal) => {
+                            timed_after_cancellation(
+                                $fut,
+                                $what,
+                                Duration::from_millis(100),
+                                &cancel,
+                            )
+                        };
+                    }
+                    tokio::join!(
+                        timed!(read_messages, "read-messages"),
+                        timed!(batcher, "batcher"),
+                        timed!(executor, "executor"),
+                    )
+                }
             }
             PageServiceProtocolPipeliningMode::Tasks => {
                 // cancelled via sensitivity to self.cancel