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