Getting current Handle in Drop (#1395)

Signed-off-by: Freyskeyd <[email protected]>

Author: Freyskeyd

sqlx-core/src/pool/connection.rs

@@ -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());
         }
     }

sqlx-core/src/postgres/listener.rs

@@ -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);
         }
     }
 }

sqlx-rt/src/lib.rs

@@ -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(