From 327b1ad064dfda2c75f3145a72e9241bbb2ff81d Mon Sep 17 00:00:00 2001
From: Sergio Benitez <sb@sergio.bz>
Date: Sun, 18 Aug 2024 15:44:30 -0700
Subject: [PATCH] Allow sync drops for 'sync_db_pools' connections.

Prior to this commit, connections from 'sync_db_pools' assumed that they
were being dropped in an async context. This is overwhelmingly the
common case as connections are typically dropped immediately after
request processing. Nothing requires this, however, so holding a
connection beyond the scope of the async context was possible (i.e. by
storing a connection in managed state). Given the connection's `Drop`
impl calls `spawn_blocking`, this resulted in a panic on drop.

This commit resolves the issue by modifying `Drop` so that it calls
`spawn_blocking` only when it is executing inside an async context. If
not, the connection is dropped normally, without `spawn_blocking`.
---
 contrib/sync_db_pools/lib/Cargo.toml          |  6 ++++
 contrib/sync_db_pools/lib/src/connection.rs   | 35 ++++++++++++-------
 .../lib/tests/drop-with-connection.rs         | 25 +++++++++++++
 3 files changed, 54 insertions(+), 12 deletions(-)
 create mode 100644 contrib/sync_db_pools/lib/tests/drop-with-connection.rs

diff --git a/contrib/sync_db_pools/lib/Cargo.toml b/contrib/sync_db_pools/lib/Cargo.toml
index 765ece2e..5e8dfa3f 100644
--- a/contrib/sync_db_pools/lib/Cargo.toml
+++ b/contrib/sync_db_pools/lib/Cargo.toml
@@ -48,5 +48,11 @@ default-features = false
 [build-dependencies]
 version_check = "0.9.1"
 
+[dev-dependencies.rocket]
+version = "0.6.0-dev"
+path = "../../../core/lib"
+default-features = false
+features = ["trace"]
+
 [package.metadata.docs.rs]
 all-features = true
diff --git a/contrib/sync_db_pools/lib/src/connection.rs b/contrib/sync_db_pools/lib/src/connection.rs
index bd5c2b4b..34f43e17 100644
--- a/contrib/sync_db_pools/lib/src/connection.rs
+++ b/contrib/sync_db_pools/lib/src/connection.rs
@@ -182,31 +182,42 @@ impl<K, C: Poolable> Drop for Connection<K, C> {
         let connection = self.connection.clone();
         let permit = self.permit.take();
 
-        // See same motivation above for this arrangement of spawn_blocking/block_on
-        tokio::task::spawn_blocking(move || {
-            let mut connection = tokio::runtime::Handle::current().block_on(async {
-                connection.lock_owned().await
-            });
+        // Only use spawn_blocking if the Tokio runtime is still available
+        if let Ok(handle) = tokio::runtime::Handle::try_current() {
+            // See above for motivation of this arrangement of spawn_blocking/block_on
+            handle.spawn_blocking(move || {
+                let mut connection = tokio::runtime::Handle::current()
+                    .block_on(async { connection.lock_owned().await });
 
-            if let Some(conn) = connection.take() {
+                if let Some(conn) = connection.take() {
+                    drop(conn);
+                }
+            });
+        } else {
+            warn!(type_name = std::any::type_name::<K>(),
+                "database connection is being dropped outside of an async context\n\
+                this means you have stored a connection beyond a request's lifetime\n\
+                this is not recommended: connections are not valid indefinitely\n\
+                instead, store a connection pool and get connections as needed");
+
+            if let Some(conn) = connection.blocking_lock().take() {
                 drop(conn);
             }
+        }
 
-            // Explicitly dropping the permit here so that it's only
-            // released after the connection is.
-            drop(permit);
-        });
+        // Explicitly drop permit here to release only after dropping connection.
+        drop(permit);
     }
 }
 
 impl<K, C: Poolable> Drop for ConnectionPool<K, C> {
     fn drop(&mut self) {
+        // Use spawn_blocking if the Tokio runtime is still available. Otherwise
+        // the pool will be dropped on the current thread.
         let pool = self.pool.take();
-        // Only use spawn_blocking if the Tokio runtime is still available
         if let Ok(handle) = tokio::runtime::Handle::try_current() {
             handle.spawn_blocking(move || drop(pool));
         }
-        // Otherwise the pool will be dropped on the current thread
     }
 }
 
diff --git a/contrib/sync_db_pools/lib/tests/drop-with-connection.rs b/contrib/sync_db_pools/lib/tests/drop-with-connection.rs
new file mode 100644
index 00000000..30dc2742
--- /dev/null
+++ b/contrib/sync_db_pools/lib/tests/drop-with-connection.rs
@@ -0,0 +1,25 @@
+#![cfg(feature = "diesel_sqlite_pool")]
+
+use rocket::figment::Figment;
+use rocket_sync_db_pools::database;
+
+#[database("example")]
+struct ExampleDb(diesel::SqliteConnection);
+
+#[test]
+fn can_drop_connection_in_sync_context() {
+    let conn = rocket::execute(async {
+        let figment = Figment::from(rocket::Config::debug_default())
+            .merge(("databases.example.url", ":memory:"));
+
+        let rocket = rocket::custom(figment)
+            .attach(ExampleDb::fairing())
+            .ignite().await
+            .expect("rocket");
+
+        ExampleDb::get_one(&rocket).await
+            .expect("attach => connection")
+    });
+
+    drop(conn);
+}