From 7507d137dee7bbc4335d9710e03e69bb25fae426 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Mon, 24 Jun 2024 18:25:57 +0000
Subject: [PATCH] WIP

---
 pageserver/src/bin/pageserver.rs  |  2 +-
 pageserver/src/lib.rs             | 21 +++++++++++----------
 pageserver/src/walredo.rs         | 13 +++++++------
 pageserver/src/walredo/process.rs |  4 ----
 4 files changed, 19 insertions(+), 21 deletions(-)

diff --git a/pageserver/src/bin/pageserver.rs b/pageserver/src/bin/pageserver.rs
index db5e0ae788..b29031c628 100644
--- a/pageserver/src/bin/pageserver.rs
+++ b/pageserver/src/bin/pageserver.rs
@@ -394,7 +394,7 @@ fn start_pageserver(
         deletion_workers.spawn_with(BACKGROUND_RUNTIME.handle());
     }
 
-    // Set up global walredo manager state
+    // Set up global tracking of walredo processes
     let walredo_global_state = BACKGROUND_RUNTIME.block_on(
         pageserver::walredo::GlobalState::spawn(conf, shutdown_pageserver.clone()),
     );
diff --git a/pageserver/src/lib.rs b/pageserver/src/lib.rs
index 2c47d6cfe0..cdfd6d1907 100644
--- a/pageserver/src/lib.rs
+++ b/pageserver/src/lib.rs
@@ -82,6 +82,17 @@ pub async fn shutdown_pageserver(
     )
     .await;
 
+    // In theory, walredo processes are tenant-scoped and should have been shut down after
+    // tenant manager shutdown above.
+    // In practive, we have lingering walredo processes even after pageserver shutdowns that
+    // don't hit the systemd TimeoutSec timeout of 10 seconds (i.e., that log `Shut down successfully completed` below).
+    timed(
+        walredo_global_state.wait_shutdown_complete(),
+        "wait for walredo processes to exit",
+        Duration::from_secs(1),
+    )
+    .await;
+
     // Shut down any page service tasks: any in-progress work for particular timelines or tenants
     // should already have been canclled via mgr::shutdown_all_tenants
     timed(
@@ -91,16 +102,6 @@ pub async fn shutdown_pageserver(
     )
     .await;
 
-    // The caller of this function already cancelled the `shutdown_pageserver` cancellation token,
-    // to which all the per-tenant walredo _manager_ methods are sensitive.
-    // This here is just to make sure the underlying walredo _processes_ are gone.
-    timed(
-        walredo_global_state.wait_shutdown_complete().await,
-        "wait for walredo processes to exit",
-        Duration::from_secs(1),
-    )
-    .await;
-
     // Best effort to persist any outstanding deletions, to avoid leaking objects
     deletion_queue.shutdown(Duration::from_secs(5)).await;
 
diff --git a/pageserver/src/walredo.rs b/pageserver/src/walredo.rs
index d7b712bb8e..005d551cb7 100644
--- a/pageserver/src/walredo.rs
+++ b/pageserver/src/walredo.rs
@@ -84,10 +84,8 @@ impl GlobalState {
         self.spawn_gate.close().await;
         info!("all walredo processes have been killed and no new ones will be spawned");
     }
-    pub(self) fn is_shutdown_requested(self: &Arc<Self>) -> bool {
-        self.shutdown_bool.load(Ordering::Relaxed)
-    }
-    pub(crate) fn wait_shutdown_complete(self: &Arc<Self>) {
+
+    pub(crate) async fn wait_shutdown_complete(self: &Arc<Self>) {
         assert!(
             self.shutdown.is_cancelled(),
             "must cancel the `shutdown` token before waiting, otherwise we will wait forever"
@@ -319,7 +317,10 @@ impl PostgresRedoManager {
             let result = proc
                 .apply_wal_records(rel, blknum, &base_img, records, wal_redo_timeout)
                 .await
-                .context("apply_wal_records");
+                .map_err(|e| match e {
+                    Error::Cancelled => Error::Cancelled,
+                    Error::Other(e) => Error::Other(e.context("apply_wal_records")),
+                });
             if matches!(result, Err(Error::Cancelled)) {
                 // bail asap and also avoid log noise due to the error reporting below
                 return Err(Error::Cancelled);
@@ -396,7 +397,7 @@ impl PostgresRedoManager {
             }
             n_attempts += 1;
             if n_attempts > MAX_RETRY_ATTEMPTS || result.is_ok() {
-                return result.map_err(Error::Other);
+                return result;
             }
         }
     }
diff --git a/pageserver/src/walredo/process.rs b/pageserver/src/walredo/process.rs
index c07f0a24d2..1a178708eb 100644
--- a/pageserver/src/walredo/process.rs
+++ b/pageserver/src/walredo/process.rs
@@ -218,10 +218,6 @@ impl WalRedoProcess {
     ) -> Result<Bytes, super::Error> {
         debug_assert_current_span_has_tenant_id();
 
-        if self.global_state.is_shutdown_requested() {
-            return Err(super::Error::Cancelled);
-        }
-
         let tag = protocol::BufferTag { rel, blknum };
 
         // Serialize all the messages to send the WAL redo process first.