Track all connection is dropped in deadpool-sqlite

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"] }

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<usize>,
 }
 
 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<Self::Type, Self::Error> {
         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(

crates/deadpool-sync/src/lib.rs

@@ -75,8 +75,14 @@ where
 {
     obj: Arc<Mutex<Option<T>>>,
     runtime: Runtime,
+    drop_callback: Option<Mutex<Box<dyn FnOnce() + Send + 'static>>>,
 }
 
+const _: () = {
+    const fn assert_send_sync<T: Send + Sync>() {}
+    assert_send_sync::<SyncWrapper<()>>();
+};
+
 // Implemented manually to avoid unnecessary trait bound on `E` type parameter.
 impl<T> fmt::Debug for SyncWrapper<T>
 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<SyncGuard<'_, T>, TryLockError<MutexGuard<'_, Option<T>>>> {
         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<F>(&mut self, drop_callback: F)
+    where
+        F: FnOnce() + Send + 'static,
+    {
+        let _ = self
+            .drop_callback
+            .replace(Mutex::new(Box::new(drop_callback)));
+    }
 }
 
 impl<T> Drop for SyncWrapper<T>
@@ -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();
     }