From fe4ef121b6c4ffeaf9175001c463cfdc598164d1 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Thu, 25 May 2023 10:44:26 +0200
Subject: [PATCH] use tokio::sync::watch::Receiver::wait_for

---
 pageserver/src/tenant.rs | 23 ++++++++++-------------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs
index f58fc138ee..5fba3026cd 100644
--- a/pageserver/src/tenant.rs
+++ b/pageserver/src/tenant.rs
@@ -1704,15 +1704,12 @@ impl Tenant {
     ///
     /// This function is not cancel-safe!
     pub async fn set_stopping(&self) {
-        // Get the rx before checking state inside send_if_modified.
-        // This way, when we later rx.changed().await, we won't have missed
-        // any state changes.
+        // `Activating` is a transient state during which no external state transitions are supported.
         let mut rx = self.state.subscribe();
-        while *rx.borrow() == TenantState::Activating {
-            rx.changed()
-                .await
-                .expect("we're a method on Tenant, so, we're keeping self.state alive here");
-        }
+        rx.wait_for(|state| state != TenantState::Activating)
+            .await
+            .expect("cannot drop self.state while on a &self method");
+
         let mut stopping = false;
         self.state.send_modify(|current_state| {
             match current_state {
@@ -1750,12 +1747,12 @@ impl Tenant {
     }
 
     pub async fn set_broken(&self, reason: String) {
+        // `Activating` is a transient state during which no external state transitions are supported.
         let mut rx = self.state.subscribe();
-        while *rx.borrow() == TenantState::Activating {
-            rx.changed()
-                .await
-                .expect("we're a method on Tenant, so, we're keeping self.state alive here");
-        }
+        rx.wait_for(|state| state != TenantState::Activating)
+            .await
+            .expect("cannot drop self.state while on a &self method");
+
         self.state.send_modify(|current_state| {
             match *current_state {
                 TenantState::Activating => {