diff --git a/proxy/src/auth/backend.rs b/proxy/src/auth/backend.rs index c08a3cf87a44..4b8ebae86fec 100644 --- a/proxy/src/auth/backend.rs +++ b/proxy/src/auth/backend.rs @@ -190,7 +190,10 @@ async fn auth_quirks( Err(info) => { let res = hacks::password_hack_no_authentication(info, client, &mut ctx.latency_timer) .await?; - ctx.set_endpoint_id(Some(res.info.endpoint.clone())); + + ctx.set_endpoint_id(res.info.endpoint.clone()); + tracing::Span::current().record("ep", &tracing::field::display(&res.info.endpoint)); + (res.info, Some(res.keys)) } Ok(info) => (info, None), diff --git a/proxy/src/auth/backend/link.rs b/proxy/src/auth/backend/link.rs index fd054334e8ae..d8ae362c03c1 100644 --- a/proxy/src/auth/backend/link.rs +++ b/proxy/src/auth/backend/link.rs @@ -98,6 +98,7 @@ pub(super) async fn authenticate( ctx.set_user(db_info.user.into()); ctx.set_project(db_info.aux.clone()); + tracing::Span::current().record("ep", &tracing::field::display(&db_info.aux.endpoint_id)); // Backwards compatibility. pg_sni_proxy uses "--" in domain names // while direct connections do not. Once we migrate to pg_sni_proxy diff --git a/proxy/src/auth/credentials.rs b/proxy/src/auth/credentials.rs index 5bf7667a1f80..875baaec4780 100644 --- a/proxy/src/auth/credentials.rs +++ b/proxy/src/auth/credentials.rs @@ -126,7 +126,11 @@ impl ComputeUserInfoMaybeEndpoint { }), } .transpose()?; - ctx.set_endpoint_id(endpoint.clone()); + + if let Some(ep) = &endpoint { + ctx.set_endpoint_id(ep.clone()); + tracing::Span::current().record("ep", &tracing::field::display(ep)); + } info!(%user, project = endpoint.as_deref(), "credentials"); if sni.is_some() { @@ -150,7 +154,7 @@ impl ComputeUserInfoMaybeEndpoint { Ok(Self { user, - endpoint_id: endpoint.map(EndpointId::from), + endpoint_id: endpoint, options, }) } diff --git a/proxy/src/context.rs b/proxy/src/context.rs index ed2ed5e36761..e2b0294cd3f0 100644 --- a/proxy/src/context.rs +++ b/proxy/src/context.rs @@ -89,13 +89,11 @@ impl RequestMonitoring { self.project = Some(x.project_id); } - pub fn set_endpoint_id(&mut self, endpoint_id: Option) { - self.endpoint_id = endpoint_id.or_else(|| self.endpoint_id.clone()); - if let Some(ep) = &self.endpoint_id { - crate::metrics::CONNECTING_ENDPOINTS - .with_label_values(&[self.protocol]) - .measure(&ep); - } + pub fn set_endpoint_id(&mut self, endpoint_id: EndpointId) { + crate::metrics::CONNECTING_ENDPOINTS + .with_label_values(&[self.protocol]) + .measure(&endpoint_id); + self.endpoint_id = Some(endpoint_id); } pub fn set_application(&mut self, app: Option) { diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs index 6e803733d0d8..4aa1f3590dff 100644 --- a/proxy/src/proxy.rs +++ b/proxy/src/proxy.rs @@ -77,6 +77,13 @@ pub async fn task_main( let cancel_map = Arc::clone(&cancel_map); let endpoint_rate_limiter = endpoint_rate_limiter.clone(); + let session_span = info_span!( + "handle_client", + ?session_id, + peer_addr = tracing::field::Empty, + ep = tracing::field::Empty, + ); + connections.spawn( async move { info!("accepted postgres client connection"); @@ -107,15 +114,11 @@ pub async fn task_main( ) .await } - .instrument(info_span!( - "handle_client", - ?session_id, - peer_addr = tracing::field::Empty - )) .unwrap_or_else(move |e| { // Acknowledge that the task has finished with an error. - error!(?session_id, "per-client task finished with an error: {e:#}"); - }), + error!("per-client task finished with an error: {e:#}"); + }) + .instrument(session_span), ); } @@ -212,8 +215,6 @@ pub async fn handle_client( Err(e) => stream.throw_error(e).await?, }; - ctx.set_endpoint_id(user_info.get_endpoint()); - // check rate limit if let Some(ep) = user_info.get_endpoint() { if !endpoint_rate_limiter.check(ep) { diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 1e2ddaa2ff51..27c21342213f 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -189,7 +189,7 @@ fn get_conn_info( } let endpoint = endpoint_sni(hostname, &tls.common_names)?.context("malformed endpoint")?; - ctx.set_endpoint_id(Some(endpoint.clone())); + ctx.set_endpoint_id(endpoint.clone()); let pairs = connection_url.query_pairs();