From 135d16a34f2a05b075a7939d129d9076f4356c06 Mon Sep 17 00:00:00 2001
From: Simon Paitrault <freyskeyd@gmail.com>
Date: Mon, 30 Aug 2021 23:11:49 +0200
Subject: [PATCH] Getting current Handle in Drop (#1395)

Signed-off-by: Freyskeyd <simon.paitrault@gmail.com>
---
 sqlx-core/src/pool/connection.rs   |  6 ++++++
 sqlx-core/src/postgres/listener.rs | 14 +++++++++++---
 sqlx-rt/src/lib.rs                 |  2 +-
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/sqlx-core/src/pool/connection.rs b/sqlx-core/src/pool/connection.rs
index 643f8df1c1..88864566c1 100644
--- a/sqlx-core/src/pool/connection.rs
+++ b/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());
         }
     }
diff --git a/sqlx-core/src/postgres/listener.rs b/sqlx-core/src/postgres/listener.rs
index b9968f0295..36a1cb5f5c 100644
--- a/sqlx-core/src/postgres/listener.rs
+++ b/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);
         }
     }
 }
diff --git a/sqlx-rt/src/lib.rs b/sqlx-rt/src/lib.rs
index 4b5cd77823..39c50855ec 100644
--- a/sqlx-rt/src/lib.rs
+++ b/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(