From 10cd64cf8dd8b4280882ac3ba0d89182ac1b3b14 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Wed, 14 Dec 2022 19:26:06 +0100 Subject: [PATCH] make TaskHandle::next_task_event cancellation-safe If we get cancelled before jh.await returns we've take()n the join handle but drop the result on the floor. Fix it by setting self.join_handle = None after the .await fixes https://github.com/neondatabase/neon/issues/3104 --- pageserver/src/walreceiver.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/pageserver/src/walreceiver.rs b/pageserver/src/walreceiver.rs index e627e9ecd0..74ede7c213 100644 --- a/pageserver/src/walreceiver.rs +++ b/pageserver/src/walreceiver.rs @@ -126,15 +126,21 @@ impl TaskHandle { match self.events_receiver.changed().await { Ok(()) => TaskEvent::Update((self.events_receiver.borrow()).clone()), Err(_task_channel_part_dropped) => { - TaskEvent::End(match self.join_handle.take() { + TaskEvent::End(match self.join_handle.as_mut() { Some(jh) => { if !jh.is_finished() { warn!("sender is dropped while join handle is still alive"); } - jh.await + let res = jh + .await .map_err(|e| anyhow::anyhow!("Failed to join task: {e}")) - .and_then(|x| x) + .and_then(|x| x); + + // For cancellation-safety, drop join_handle only after successful .await. + self.join_handle = None; + + res } None => { // Another option is to have an enum, join handle or result and give away the reference to it