From af5bb67f08e92c018111055dcef26d9e3bab665a Mon Sep 17 00:00:00 2001
From: Vlad Lazar <vlad@neon.tech>
Date: Mon, 2 Jun 2025 09:59:21 +0100
Subject: [PATCH] pageserver: more reactive wal receiver cancellation (#12076)

## Problem

If the wal receiver is cancelled, there's a 50% chance that it will
ingest yet more WAL.

## Summary of Changes

Always check cancellation first.
---
 pageserver/src/tenant/timeline/walreceiver.rs                   | 2 +-
 .../src/tenant/timeline/walreceiver/walreceiver_connection.rs   | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/pageserver/src/tenant/timeline/walreceiver.rs b/pageserver/src/tenant/timeline/walreceiver.rs
index 0f73eb839b..633c94a010 100644
--- a/pageserver/src/tenant/timeline/walreceiver.rs
+++ b/pageserver/src/tenant/timeline/walreceiver.rs
@@ -113,7 +113,7 @@ impl WalReceiver {
                 }
                 connection_manager_state.shutdown().await;
                 *loop_status.write().unwrap() = None;
-                debug!("task exits");
+                info!("task exits");
             }
             .instrument(info_span!(parent: None, "wal_connection_manager", tenant_id = %tenant_shard_id.tenant_id, shard_id = %tenant_shard_id.shard_slug(), timeline_id = %timeline_id))
         });
diff --git a/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs b/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs
index 52259f205b..249849ac4b 100644
--- a/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs
+++ b/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs
@@ -297,6 +297,7 @@ pub(super) async fn handle_walreceiver_connection(
     let mut expected_wal_start = startpoint;
     while let Some(replication_message) = {
         select! {
+            biased;
             _ = cancellation.cancelled() => {
                 debug!("walreceiver interrupted");
                 None