Skip to content

Commit

Permalink
Getting current Handle in Drop (#1395)
Browse files Browse the repository at this point in the history
Signed-off-by: Freyskeyd <simon.paitrault@gmail.com>
  • Loading branch information
Freyskeyd committed Aug 30, 2021
1 parent ad81e35 commit 135d16a
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 4 deletions.
6 changes: 6 additions & 0 deletions sqlx-core/src/pool/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,12 @@ impl<DB: Database> PoolConnection<DB> {
impl<DB: Database> Drop for PoolConnection<DB> {
fn drop(&mut self) {
if self.live.is_some() {
#[cfg(not(feature = "_rt-async-std"))]
if let Ok(handle) = sqlx_rt::Handle::try_current() {
handle.spawn(self.return_to_pool());
}

#[cfg(feature = "_rt-async-std")]
sqlx_rt::spawn(self.return_to_pool());
}
}
Expand Down
14 changes: 11 additions & 3 deletions sqlx-core/src/postgres/listener.rs
Original file line number Diff line number Diff line change
Expand Up @@ -260,15 +260,23 @@ impl PgListener {
impl Drop for PgListener {
fn drop(&mut self) {
if let Some(mut conn) = self.connection.take() {
// Unregister any listeners before returning the connection to the pool.
sqlx_rt::spawn(async move {
let fut = async move {
let _ = conn.execute("UNLISTEN *").await;

// inline the drop handler from `PoolConnection` so it doesn't try to spawn another task
// otherwise, it may trigger a panic if this task is dropped because the runtime is going away:
// https://github.com/launchbadge/sqlx/issues/1389
conn.return_to_pool().await;
});
};

// Unregister any listeners before returning the connection to the pool.
#[cfg(not(feature = "_rt-async-std"))]
if let Ok(handle) = sqlx_rt::Handle::try_current() {
handle.spawn(fut);
}

#[cfg(feature = "_rt-async-std")]
sqlx_rt::spawn(fut);
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion sqlx-rt/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ pub use native_tls;
))]
pub use tokio::{
self, fs, io::AsyncRead, io::AsyncReadExt, io::AsyncWrite, io::AsyncWriteExt, io::ReadBuf,
net::TcpStream, task::spawn, task::yield_now, time::sleep, time::timeout,
net::TcpStream, runtime::Handle, task::spawn, task::yield_now, time::sleep, time::timeout,
};

#[cfg(all(
Expand Down

0 comments on commit 135d16a

Please sign in to comment.