fix: only stop task when every clone of it has been dropped

src/pool_metrics.rs

@@ -3,20 +3,23 @@
 //!
 //! 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>,
 }
 
 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);
         }
     });
 

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() {