diff --git a/crates/deadpool-sqlite/Cargo.toml b/crates/deadpool-sqlite/Cargo.toml index 630e2628..a154f5fd 100644 --- a/crates/deadpool-sqlite/Cargo.toml +++ b/crates/deadpool-sqlite/Cargo.toml @@ -71,6 +71,7 @@ rusqlite = "0.38.0" serde = { package = "serde", version = "1.0", features = [ "derive", ], optional = true } +tokio = { version = "1.37", features = ["sync"] } [dev-dependencies] config = { version = "0.15", features = ["json"] } diff --git a/crates/deadpool-sqlite/src/lib.rs b/crates/deadpool-sqlite/src/lib.rs index 34329969..4472ea8c 100644 --- a/crates/deadpool-sqlite/src/lib.rs +++ b/crates/deadpool-sqlite/src/lib.rs @@ -53,6 +53,7 @@ pub struct Manager { config: Config, recycle_count: AtomicIsize, runtime: Runtime, + active_count: tokio::sync::watch::Sender, } impl Manager { @@ -64,8 +65,24 @@ impl Manager { config: config.clone(), recycle_count: AtomicIsize::new(0), runtime, + active_count: tokio::sync::watch::Sender::new(0), } } + + /// Retrieve the active `rusqlite::Connection` count. + #[must_use] + pub fn active_count(&self) -> usize { + *self.active_count.borrow() + } + + /// Wait for all `rusqlite::Connection` is closed in background. + pub async fn wait_all_inactive(&self) { + let _ = self + .active_count + .subscribe() + .wait_for(|count| *count == 0) + .await; + } } impl managed::Manager for Manager { @@ -74,7 +91,12 @@ impl managed::Manager for Manager { async fn create(&self) -> Result { let path = self.config.path.clone(); - SyncWrapper::new(self.runtime, move || rusqlite::Connection::open(path)).await + let mut conn = + SyncWrapper::new(self.runtime, move || rusqlite::Connection::open(path)).await?; + self.active_count.send_modify(|count| *count += 1); + let active_count = self.active_count.clone(); + conn.set_drop_callback(move || active_count.send_modify(|count| *count -= 1)); + Ok(conn) } async fn recycle( diff --git a/crates/deadpool-sync/src/lib.rs b/crates/deadpool-sync/src/lib.rs index 75c339bb..6959034a 100644 --- a/crates/deadpool-sync/src/lib.rs +++ b/crates/deadpool-sync/src/lib.rs @@ -75,8 +75,14 @@ where { obj: Arc>>, runtime: Runtime, + drop_callback: Option>>, } +const _: () = { + const fn assert_send_sync() {} + assert_send_sync::>(); +}; + // Implemented manually to avoid unnecessary trait bound on `E` type parameter. impl fmt::Debug for SyncWrapper where @@ -112,6 +118,7 @@ where result.map(|obj| Self { obj: Arc::new(Mutex::new(Some(obj))), runtime, + drop_callback: None, }) } @@ -157,6 +164,17 @@ where pub fn try_lock(&self) -> Result, TryLockError>>> { self.obj.try_lock().map(SyncGuard) } + + /// Set the callback when the actual obj is dropped. + // https://github.com/deadpool-rs/deadpool/issues/330 + pub fn set_drop_callback(&mut self, drop_callback: F) + where + F: FnOnce() + Send + 'static, + { + let _ = self + .drop_callback + .replace(Mutex::new(Box::new(drop_callback))); + } } impl Drop for SyncWrapper @@ -165,11 +183,18 @@ where { fn drop(&mut self) { let arc = self.obj.clone(); + let drop_callback = self.drop_callback.take(); + #[cfg(feature = "tracing")] + let span = tracing::Span::current(); // Drop the `rusqlite::Connection` inside a `spawn_blocking` // as the `drop` function of it can block. - spawn_blocking_background(self.runtime, move || match arc.lock() { - Ok(mut guard) => drop(guard.take()), - Err(e) => drop(e.into_inner().take()), + spawn_blocking_background(self.runtime, move || { + drop(arc.lock().unwrap_or_else(PoisonError::into_inner).take()); + if let Some(callback) = drop_callback { + #[cfg(feature = "tracing")] + let _span = span.enter(); + callback.into_inner().unwrap()(); + } }) .unwrap(); }