Skip to content

Commit

Permalink
Allow sync drops for 'sync_db_pools' connections.
Browse files Browse the repository at this point in the history
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`.
  • Loading branch information
SergioBenitez committed Aug 19, 2024
1 parent 7fdcf2d commit 327b1ad
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 12 deletions.
6 changes: 6 additions & 0 deletions contrib/sync_db_pools/lib/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
35 changes: 23 additions & 12 deletions contrib/sync_db_pools/lib/src/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}

Expand Down
25 changes: 25 additions & 0 deletions contrib/sync_db_pools/lib/tests/drop-with-connection.rs
Original file line number Diff line number Diff line change
@@ -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);
}

0 comments on commit 327b1ad

Please sign in to comment.