From 351d4338b54f00e1a7e018090d76fa1b4aac1a12 Mon Sep 17 00:00:00 2001
From: Mario Savarese <mario.savarese@evervault.com>
Date: Mon, 20 Jul 2026 13:53:12 +0100
Subject: [PATCH] fix: only stop task when every clone of it has been dropped

This fixes an issue where the shutdown flag gets permanently flipped to true
(which signals that the polling task should be stopped) any time a cloned
pool is dropped. Because all the pool clones share the same `Arc<AtomicBool>`,
the first clone to be dropped incorrectly flips the flag, which is not
what the documentation implies.

This also adds a regression test to make sure the metric task surives
a cloned pool being dropped
---
 src/pool_metrics.rs   | 21 +++++++++++++--------
 tests/pool_metrics.rs | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/src/pool_metrics.rs b/src/pool_metrics.rs
index 82a9dc1..6b15a45 100644
--- a/src/pool_metrics.rs
+++ b/src/pool_metrics.rs
@@ -3,12 +3,13 @@
 //!
 //! Spawned by [`PoolBuilder::build()`](crate::PoolBuilder::build) when a pool name is
 //! configured and a runtime feature (e.g. `runtime-tokio`) is enabled. The task stops
-//! when the [`Pool`](crate::Pool) is dropped.
+//! only once the [`Pool`](crate::Pool) and every clone of it have been dropped.
 
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 
-/// Handle that signals the background polling task to stop when dropped.
+/// Handle that signals the background polling task to stop when the last clone is
+/// dropped.
 #[derive(Clone)]
 pub(crate) struct ShutdownHandle {
     flag: Arc<AtomicBool>,
@@ -16,7 +17,9 @@ pub(crate) struct ShutdownHandle {
 
 impl Drop for ShutdownHandle {
     fn drop(&mut self) {
-        self.flag.store(true, Ordering::Relaxed);
+        if Arc::strong_count(&self.flag) == 1 {
+            self.flag.store(true, Ordering::Relaxed);
+        }
     }
 }
 
@@ -61,8 +64,8 @@ pub(crate) fn spawn<R: crate::runtime::Runtime, DB: sqlx::Database>(
         count.record(used, &used_attrs);
     }
 
-    let shutdown = Arc::new(AtomicBool::new(false));
-    let flag = shutdown.clone();
+    let flag = Arc::new(AtomicBool::new(false));
+    let shutdown = Arc::downgrade(&flag);
 
     R::spawn(async move {
         let meter: Meter = opentelemetry::global::meter("sqlx-otel");
@@ -77,10 +80,12 @@ pub(crate) fn spawn<R: crate::runtime::Runtime, DB: sqlx::Database>(
 
         loop {
             R::sleep(interval).await;
-            if shutdown.load(Ordering::Relaxed) {
-                break;
+            match shutdown.upgrade() {
+                Some(flag) if !flag.load(Ordering::Relaxed) => {
+                    record(&count, &pool, &base_attrs);
+                }
+                _ => break,
             }
-            record(&count, &pool, &base_attrs);
         }
     });
 
diff --git a/tests/pool_metrics.rs b/tests/pool_metrics.rs
index 1f48866..e91003b 100644
--- a/tests/pool_metrics.rs
+++ b/tests/pool_metrics.rs
@@ -376,6 +376,42 @@ mod tokio_runtime {
         drop(conn);
     }
 
+    #[tokio::test]
+    #[serial]
+    async fn background_task_survives_clone_drop() {
+        let tel = common::TestTelemetry::install();
+        let raw = sqlx::SqlitePool::connect(":memory:").await.unwrap();
+        let pool = PoolBuilder::from(raw)
+            .with_pool_name(POOL_NAME)
+            .with_pool_metrics_interval(Duration::from_millis(50))
+            .build();
+
+        let clone = pool.clone();
+        // Dropping a clone must NOT stop the polling task
+        drop(clone);
+
+        let conn = pool.acquire().await.unwrap();
+
+        let used = poll_for(Duration::from_secs(2), || {
+            let snapshot = tel.metrics();
+            let used = gauge_value(
+                &snapshot,
+                "db.client.connection.count",
+                "db.client.connection.state",
+                "used",
+            )?;
+            if used >= 1 { Some(used) } else { None }
+        })
+        .await;
+
+        assert!(
+            used.is_some(),
+            "connection.count should keep updating after a clone is dropped"
+        );
+
+        drop(conn);
+    }
+
     #[tokio::test]
     #[serial]
     async fn no_pool_metrics_without_pool_name() {