Tracing support

tokio-postgres/Cargo.toml

@@ -58,6 +58,7 @@ postgres-protocol = { version = "0.6.5", path = "../postgres-protocol" }
 postgres-types = { version = "0.2.4", path = "../postgres-types" }
 tokio = { version = "1.27", features = ["io-util"] }
 tokio-util = { version = "0.7", features = ["codec"] }
+tracing = "0.1"
 rand = "0.8.5"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]

tokio-postgres/src/client.rs

@@ -9,12 +9,13 @@ use crate::simple_query::SimpleQueryStream;
 #[cfg(feature = "runtime")]
 use crate::tls::MakeTlsConnect;
 use crate::tls::TlsConnect;
+use crate::trace::{make_span_for_client, SpanOperation};
 use crate::types::{Oid, ToSql, Type};
 #[cfg(feature = "runtime")]
 use crate::Socket;
 use crate::{
-    copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken, CopyInSink, Error,
-    Row, SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
+    copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken, Config, CopyInSink,
+    Error, Row, SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
 };
 use bytes::{Buf, BytesMut};
 use fallible_iterator::FallibleIterator;
@@ -34,6 +35,7 @@ use std::task::{Context, Poll};
 #[cfg(feature = "runtime")]
 use std::time::Duration;
 use tokio::io::{AsyncRead, AsyncWrite};
+use tracing::Instrument;
 
 pub struct Responses {
     receiver: mpsc::Receiver<BackendMessages>,
@@ -85,6 +87,11 @@ pub struct InnerClient {
     sender: mpsc::UnboundedSender<Request>,
     cached_typeinfo: Mutex<CachedTypeInfo>,
 
+    db_user: String,
+    db_name: String,
+    #[cfg(feature = "runtime")]
+    socket_config: Option<SocketConfig>,
+
     /// A buffer to use when writing out postgres commands.
     buffer: Mutex<BytesMut>,
 }
@@ -103,6 +110,18 @@ impl InnerClient {
         })
     }
 
+    pub fn db_user(&self) -> &str {
+        &self.db_user
+    }
+
+    pub fn db_name(&self) -> &str {
+        &self.db_name
+    }
+
+    pub(crate) fn socket_config(&self) -> Option<&SocketConfig> {
+        self.socket_config.as_ref()
+    }
+
     pub fn typeinfo(&self) -> Option<Statement> {
         self.cached_typeinfo.lock().typeinfo.clone()
     }
@@ -190,11 +209,16 @@ impl Client {
         ssl_mode: SslMode,
         process_id: i32,
         secret_key: i32,
+        config: &Config,
     ) -> Client {
         Client {
             inner: Arc::new(InnerClient {
                 sender,
                 cached_typeinfo: Default::default(),
+                db_user: config.user.clone().unwrap_or_default(),
+                db_name: config.dbname.clone().unwrap_or_default(),
+                #[cfg(feature = "runtime")]
+                socket_config: None,
                 buffer: Default::default(),
             }),
             #[cfg(feature = "runtime")]
@@ -211,6 +235,9 @@ impl Client {
 
     #[cfg(feature = "runtime")]
     pub(crate) fn set_socket_config(&mut self, socket_config: SocketConfig) {
+        if let Some(inner) = Arc::get_mut(&mut self.inner) {
+            inner.socket_config = Some(socket_config.clone());
+        }
         self.socket_config = Some(socket_config);
     }
 
@@ -231,7 +258,11 @@ impl Client {
         query: &str,
         parameter_types: &[Type],
     ) -> Result<Statement, Error> {
-        prepare::prepare(&self.inner, query, parameter_types).await
+        let span = make_span_for_client(&self.inner, SpanOperation::Prepare);
+
+        prepare::prepare(&self.inner, query, parameter_types)
+            .instrument(span)
+            .await
     }
 
     /// Executes a statement, returning a vector of the resulting rows.

tokio-postgres/src/connect.rs

@@ -3,12 +3,15 @@ use crate::config::{Host, LoadBalanceHosts, TargetSessionAttrs};
 use crate::connect_raw::connect_raw;
 use crate::connect_socket::connect_socket;
 use crate::tls::MakeTlsConnect;
+use crate::trace::{make_span, record_connect_info, record_error, record_ok, SpanOperation};
 use crate::{Client, Config, Connection, Error, SimpleQueryMessage, Socket};
 use futures_util::{future, pin_mut, Future, FutureExt, Stream};
 use rand::seq::SliceRandom;
 use std::task::Poll;
+use std::time::Instant;
 use std::{cmp, io};
 use tokio::net;
+use tracing::Instrument as _;
 
 pub async fn connect<T>(
     mut tls: T,
@@ -93,24 +96,51 @@ async fn connect_host<T>(
 where
     T: MakeTlsConnect<Socket>,
 {
+    let mut span = make_span(
+        SpanOperation::Connect,
+        config.user.as_deref().unwrap_or_default(),
+        config.dbname.as_deref().unwrap_or_default(),
+    );
     match host {
         Host::Tcp(host) => {
+            let dns_start = Instant::now();
             let mut addrs = net::lookup_host((&*host, port))
+                .instrument(span.clone())
                 .await
                 .map_err(Error::connect)?
                 .collect::<Vec<_>>();
+            span.record(
+                "db.connect.timing.dns_lookup_ns",
+                dns_start.elapsed().as_nanos() as u64,
+            );
 
             if config.load_balance_hosts == LoadBalanceHosts::Random {
                 addrs.shuffle(&mut rand::thread_rng());
             }
 
             let mut last_err = None;
-            for addr in addrs {
-                match connect_once(Addr::Tcp(addr.ip()), hostname.as_deref(), port, tls, config)
+            for (i, addr) in addrs.iter().enumerate() {
+                if i > 0 {
+                    span = make_span(
+                        SpanOperation::Connect,
+                        config.user.as_deref().unwrap_or_default(),
+                        config.dbname.as_deref().unwrap_or_default(),
+                    );
+                }
+
+                let addr = Addr::Tcp(addr.ip());
+                span.record("db.connect.attempt", i);
+                record_connect_info(&span, &addr, hostname.as_deref(), port);
+                match connect_once(addr, hostname.as_deref(), port, tls, config)
+                    .instrument(span.clone())
                     .await
                 {
-                    Ok(stream) => return Ok(stream),
+                    Ok((client, conn)) => {
+                        record_ok(&span);
+                        return Ok((client, conn));
+                    }
                     Err(e) => {
+                        record_error(&span, &e);
                         last_err = Some(e);
                         continue;
                     }
@@ -126,7 +156,21 @@ where
         }
         #[cfg(unix)]
         Host::Unix(path) => {
-            connect_once(Addr::Unix(path), hostname.as_deref(), port, tls, config).await
+            let addr = Addr::Unix(path);
+            record_connect_info(&span, &addr, hostname.as_deref(), port);
+            match connect_once(addr, hostname.as_deref(), port, tls, config)
+                .instrument(span.clone())
+                .await
+            {
+                Ok((client, conn)) => {
+                    record_ok(&span);
+                    Ok((client, conn))
+                }
+                Err(e) => {
+                    record_error(&span, &e);
+                    Err(e)
+                }
+            }
         }
     }
 }

tokio-postgres/src/connect_raw.rs

@@ -17,8 +17,10 @@ use std::collections::{HashMap, VecDeque};
 use std::io;
 use std::pin::Pin;
 use std::task::{Context, Poll};
+use std::time::Instant;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_util::codec::Framed;
+use tracing::Span;
 
 pub struct StartupStream<S, T> {
     inner: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
@@ -96,12 +98,17 @@ where
         delayed: VecDeque::new(),
     };
 
+    let start_time = Instant::now();
     startup(&mut stream, config).await?;
     authenticate(&mut stream, config).await?;
     let (process_id, secret_key, parameters) = read_info(&mut stream).await?;
+    Span::current().record(
+        "db.connect.timing.auth_ns",
+        start_time.elapsed().as_nanos() as u64,
+    );
 
     let (sender, receiver) = mpsc::unbounded();
-    let client = Client::new(sender, config.ssl_mode, process_id, secret_key);
+    let client = Client::new(sender, config.ssl_mode, process_id, secret_key, config);
     let connection = Connection::new(stream.inner, stream.delayed, parameters, receiver);
 
     Ok((client, connection))

tokio-postgres/src/connect_socket.rs

@@ -4,11 +4,12 @@ use crate::{Error, Socket};
 use socket2::{SockRef, TcpKeepalive};
 use std::future::Future;
 use std::io;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 use tokio::net::TcpStream;
 #[cfg(unix)]
 use tokio::net::UnixStream;
 use tokio::time;
+use tracing::Span;
 
 pub(crate) async fn connect_socket(
     addr: &Addr,
@@ -21,8 +22,13 @@ pub(crate) async fn connect_socket(
 ) -> Result<Socket, Error> {
     match addr {
         Addr::Tcp(ip) => {
+            let connect_start = Instant::now();
             let stream =
                 connect_with_timeout(TcpStream::connect((*ip, port)), connect_timeout).await?;
+            Span::current().record(
+                "db.connect.timing.tcp_handshake_ns",
+                connect_start.elapsed().as_nanos() as u64,
+            );
 
             stream.set_nodelay(true).map_err(Error::connect)?;
 

tokio-postgres/src/connect_tls.rs

@@ -5,7 +5,9 @@ use crate::tls::TlsConnect;
 use crate::Error;
 use bytes::BytesMut;
 use postgres_protocol::message::frontend;
+use std::time::Instant;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+use tracing::Span;
 
 pub async fn connect_tls<S, T>(
     mut stream: S,
@@ -25,6 +27,7 @@ where
         SslMode::Prefer | SslMode::Require => {}
     }
 
+    let start_time = Instant::now();
     let mut buf = BytesMut::new();
     frontend::ssl_request(&mut buf);
     stream.write_all(&buf).await.map_err(Error::io)?;
@@ -48,6 +51,10 @@ where
         .connect(stream)
         .await
         .map_err(|e| Error::tls(e.into()))?;
+    Span::current().record(
+        "db.connect.timing.tls_handshake_ns",
+        start_time.elapsed().as_nanos() as u64,
+    );
 
     Ok(MaybeTlsStream::Tls(stream))
 }

tokio-postgres/src/lib.rs

@@ -176,6 +176,7 @@ mod socket;
 mod statement;
 pub mod tls;
 mod to_statement;
+mod trace;
 mod transaction;
 mod transaction_builder;
 pub mod types;

tokio-postgres/src/prepare.rs

@@ -2,6 +2,7 @@ use crate::client::InnerClient;
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::error::SqlState;
+use crate::trace::{make_span_for_client, SpanOperation};
 use crate::types::{Field, Kind, Oid, Type};
 use crate::{query, slice_iter};
 use crate::{Column, Error, Statement};
@@ -15,6 +16,7 @@ use std::future::Future;
 use std::pin::Pin;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
+use tracing::Instrument;
 
 const TYPEINFO_QUERY: &str = "\
 SELECT t.typname, t.typtype, t.typelem, r.rngsubtype, t.typbasetype, n.nspname, t.typrelid
@@ -100,15 +102,22 @@ pub async fn prepare(
         }
     }
 
-    Ok(Statement::new(client, name, parameters, columns))
+    Ok(Statement::new(
+        client,
+        name,
+        query.to_string(),
+        parameters,
+        columns,
+    ))
 }
 
 fn prepare_rec<'a>(
     client: &'a Arc<InnerClient>,
     query: &'a str,
     types: &'a [Type],
 ) -> Pin<Box<dyn Future<Output = Result<Statement, Error>> + 'a + Send>> {
-    Box::pin(prepare(client, query, types))
+    let span = make_span_for_client(client, SpanOperation::Prepare);
+    Box::pin(prepare(client, query, types).instrument(span))
 }
 
 fn encode(client: &InnerClient, name: &str, query: &str, types: &[Type]) -> Result<Bytes, Error> {