postgres: Use current tracing span when dropping PgListener

launchbadge · Mar 5, 2025 · 5524ccc · 5524ccc
1 parent 3a20a92
commit 5524ccc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/sqlx-core/src/ext/async_stream.rs b/sqlx-core/src/ext/async_stream.rs
@@ -121,7 +121,7 @@ impl<'a, T> Stream for TryAsyncStream<'a, T> {
 #[macro_export]
 macro_rules! try_stream {
     ($($block:tt)*) => {
-        $crate::ext::async_stream::TryAsyncStream::new(move |yielder| async move {
+        $crate::ext::async_stream::TryAsyncStream::new(move |yielder| ::tracing::Instrument::in_current_span(async move {
             // Anti-footgun: effectively pins `yielder` to this future to prevent any accidental
             // move to another task, which could deadlock.
             let yielder = &yielder;
@@ -133,6 +133,6 @@ macro_rules! try_stream {
             }
 
             $($block)*
-        })
+        }))
     }
 }
diff --git a/sqlx-postgres/src/listener.rs b/sqlx-postgres/src/listener.rs
@@ -9,6 +9,7 @@ use futures_util::{FutureExt, StreamExt, TryFutureExt, TryStreamExt};
 use sqlx_core::acquire::Acquire;
 use sqlx_core::transaction::Transaction;
 use sqlx_core::Either;
+use tracing::Instrument;
 
 use crate::describe::Describe;
 use crate::error::Error;
@@ -366,7 +367,7 @@ impl Drop for PgListener {
             };
 
             // Unregister any listeners before returning the connection to the pool.
-            crate::rt::spawn(fut);
+            crate::rt::spawn(fut.in_current_span());
         }
     }
 }