add tracing to connect

Author: titanous

tokio-postgres/src/client.rs

@@ -9,7 +9,7 @@ use crate::simple_query::SimpleQueryStream;
 #[cfg(feature = "runtime")]
 use crate::tls::MakeTlsConnect;
 use crate::tls::TlsConnect;
-use crate::trace::{make_span, SpanOperation};
+use crate::trace::{make_span_for_client, SpanOperation};
 use crate::types::{Oid, ToSql, Type};
 #[cfg(feature = "runtime")]
 use crate::Socket;
@@ -258,7 +258,7 @@ impl Client {
         query: &str,
         parameter_types: &[Type],
     ) -> Result<Statement, Error> {
-        let span = make_span(&self.inner, SpanOperation::Prepare);
+        let span = make_span_for_client(&self.inner, SpanOperation::Prepare);
 
         prepare::prepare(&self.inner, query, parameter_types)
             .instrument(span)

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,9 +98,14 @@ 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, config);

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/prepare.rs

@@ -2,7 +2,7 @@ use crate::client::InnerClient;
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::error::SqlState;
-use crate::trace::{make_span, SpanOperation};
+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};
@@ -116,7 +116,7 @@ fn prepare_rec<'a>(
     query: &'a str,
     types: &'a [Type],
 ) -> Pin<Box<dyn Future<Output = Result<Statement, Error>> + 'a + Send>> {
-    let span = make_span(client, SpanOperation::Prepare);
+    let span = make_span_for_client(client, SpanOperation::Prepare);
     Box::pin(prepare(client, query, types).instrument(span))
 }
 

tokio-postgres/src/query.rs

@@ -1,7 +1,7 @@
 use crate::client::{InnerClient, Responses};
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
-use crate::trace::{make_span, SpanOperation};
+use crate::trace::{make_span_for_client, record_error, SpanOperation, record_ok};
 use crate::types::{BorrowToSql, IsNull};
 use crate::{Error, Portal, Row, Statement};
 use bytes::{Bytes, BytesMut};
@@ -63,12 +63,11 @@ where
 async fn start_traced(client: &InnerClient, span: &Span, buf: Bytes) -> Result<Responses, Error> {
     match start(client, buf).instrument(span.clone()).await {
         Ok(response) => {
-            span.record("otel.status_code", "OK");
+            record_ok(span);
             Ok(response)
         }
         Err(e) => {
-            span.record("otel.status_code", "ERROR");
-            span.record("exception.message", e.to_string());
+            record_error(span, &e);
             Err(e)
         }
     }
@@ -84,8 +83,7 @@ where
     I: IntoIterator<Item = P>,
     I::IntoIter: ExactSizeIterator,
 {
-    let span = make_span(client, SpanOperation::Query);
-    span.record("db.operation", "query");
+    let span = make_span_for_client(client, SpanOperation::Query);
 
     let buf = encode_with_logs(client, &span, &statement, params)?;
 
@@ -108,7 +106,7 @@ pub async fn query_portal(
     portal: &Portal,
     max_rows: i32,
 ) -> Result<RowStream, Error> {
-    let span = make_span(client, SpanOperation::Portal);
+    let span = make_span_for_client(client, SpanOperation::Portal);
     span.record("db.statement", portal.statement().query());
 
     let buf = client.with_buf(|buf| {
@@ -151,7 +149,7 @@ where
     I: IntoIterator<Item = P>,
     I::IntoIter: ExactSizeIterator,
 {
-    let span = make_span(client, SpanOperation::Execute);
+    let span = make_span_for_client(client, SpanOperation::Execute);
 
     let buf = encode_with_logs(client, &span, &statement, params)?;
 

tokio-postgres/src/trace.rs

@@ -1,8 +1,12 @@
 use tracing::{field::Empty, Level, Span};
 
-use crate::client::{Addr, InnerClient};
+use crate::{
+    client::{Addr, InnerClient, SocketConfig},
+    Error,
+};
 
 pub(crate) enum SpanOperation {
+    Connect,
     Query,
     Portal,
     Execute,
@@ -12,6 +16,7 @@ pub(crate) enum SpanOperation {
 impl SpanOperation {
     const fn name(&self) -> &'static str {
         match self {
+            SpanOperation::Connect => "connect",
             SpanOperation::Query => "query",
             SpanOperation::Portal => "portal",
             SpanOperation::Execute => "execute",
@@ -21,6 +26,7 @@ impl SpanOperation {
 
     const fn level(&self) -> Level {
         match self {
+            SpanOperation::Connect => Level::DEBUG,
             SpanOperation::Query => Level::DEBUG,
             SpanOperation::Portal => Level::DEBUG,
             SpanOperation::Execute => Level::DEBUG,
@@ -76,50 +82,77 @@ macro_rules! span_dynamic_lvl {
     };
 }
 
-pub(crate) fn make_span(client: &InnerClient, operation: SpanOperation) -> Span {
-    let span = span_dynamic_lvl!(operation.level(), "query",
+pub(crate) fn make_span_for_client(client: &InnerClient, operation: SpanOperation) -> Span {
+    let span = make_span(operation, client.db_user(), client.db_name());
+    record_socket_config(&span, client.socket_config());
+    span
+}
+
+pub(crate) fn make_span(operation: SpanOperation, db_user: &str, db_name: &str) -> Span {
+    span_dynamic_lvl!(
+        operation.level(),
+        "query",
         db.system = "postgresql",
         server.address = Empty,
         server.socket.address = Empty,
         server.port = Empty,
         "network.type" = Empty,
         network.transport = Empty,
-        db.name = %client.db_name(),
-        db.user = %client.db_user(),
-        otel.name = %format!("PSQL {} {}", operation.name(), client.db_name()),
+        db.name = db_name,
+        db.user = db_user,
+        otel.name = format!("PSQL {} {}", operation.name(), db_name),
         otel.kind = "Client",
-
         // to set when output
         otel.status_code = Empty,
         exception.message = Empty,
         db.operation = operation.name(),
-
         // for queries
         db.statement = Empty,
         db.statement.params = Empty,
         db.sql.rows_affected = Empty,
-    );
+        // for connections
+        db.connect.attempt = Empty,
+        db.connect.timing.dns_lookup_ns = Empty,
+        db.connect.timing.tcp_handshake_ns = Empty,
+        db.connect.timing.tls_handshake_ns = Empty,
+        db.connect.timing.auth_ns = Empty,
+    )
+}
+
+pub(crate) fn record_socket_config(span: &Span, socket_config: Option<&SocketConfig>) {
+    if let Some(s) = socket_config {
+        record_connect_info(span, &s.addr, s.hostname.as_deref(), s.port);
+    }
+}
+
+pub(crate) fn record_connect_info(span: &Span, addr: &Addr, hostname: Option<&str>, port: u16) {
     // only executes if span passed the filter
     span.in_scope(|| {
-        if let Some(socket_config) = client.socket_config() {
-            if let Some(hostname) = socket_config.hostname.as_deref() {
-                span.record("server.address", hostname);
+        if let Some(hostname) = hostname {
+            span.record("server.address", hostname);
+        }
+        match addr {
+            Addr::Tcp(addr) => {
+                span.record("server.socket.address", addr.to_string());
+                let network_type = if addr.is_ipv4() { "ipv4" } else { "ipv6" };
+                span.record("network.type", network_type);
+                span.record("network.transport", "tcp");
             }
-            match &socket_config.addr {
-                Addr::Tcp(addr) => {
-                    span.record("server.socket.address", addr.to_string());
-                    let network_type = if addr.is_ipv4() { "ipv4" } else { "ipv6" };
-                    span.record("network.type", network_type);
-                    span.record("network.transport", "tcp");
-                }
-                #[cfg(unix)]
-                Addr::Unix(path) => {
-                    span.record("server.socket.address", path.to_string_lossy().into_owned());
-                    span.record("network.type", "unix");
-                }
+            #[cfg(unix)]
+            Addr::Unix(path) => {
+                span.record("server.socket.address", path.to_string_lossy().into_owned());
+                span.record("network.type", "unix");
             }
-            span.record("server.port", socket_config.port);
         }
+        span.record("server.port", port);
     });
-    span
+}
+
+pub(crate) fn record_error(span: &Span, e: &Error) {
+    span.record("otel.status_code", "ERROR")
+        .record("exception.message", tracing::field::display(e));
+}
+
+pub(crate) fn record_ok(span: &Span) {
+    span.record("otel.status_code", "OK");
 }