From db0deb8457c904ca9dd98ffba9bc5a49d56a1640 Mon Sep 17 00:00:00 2001
From: John Spray <john@neon.tech>
Date: Thu, 3 Aug 2023 10:01:12 +0100
Subject: [PATCH] pageserver: respect task_mgr cancellation in metrics task

This previously relied on seeing a channel close, when
Tenant is destroyed: this task ran beyond Tenant::shutdown,
whereas the idea of that shutdown function is that all the
per-tenant background tasks are joined when it completes.

Instead, stop this task as soon as background tasks for the
Tenant are cancelled, making the behavior of shutdown() much
more obvious and bringing the task into line with how we
do shutdown in other background tasks.
---
 pageserver/src/tenant.rs | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs
index adb3e86ffd..5d05eccd67 100644
--- a/pageserver/src/tenant.rs
+++ b/pageserver/src/tenant.rs
@@ -2084,6 +2084,8 @@ impl Tenant {
             &format!("state metrics collector for tenant {tenant_id}"),
             false,
             async move {
+            let cancel = task_mgr::shutdown_token();
+
             let tid = tenant_id.to_string();
 
             fn inspect_state(state: &TenantState) -> ([&'static str; 1], bool) {
@@ -2113,7 +2115,12 @@ impl Tenant {
                 let current = TENANT_STATE_METRIC.with_label_values(labels);
                 current.inc();
 
-                if rx.changed().await.is_err() {
+                let changed = tokio::select! {
+                    changed = rx.changed() => {changed},
+                    _ = cancel.cancelled() => {return Ok(())}
+                };
+
+                if changed.is_err() {
                     // tenant has been dropped; decrement the counter because a tenant with that
                     // state is no longer in tenant map, but allow any broken set item to exist
                     // still.