From 0c243faf96c4ee9234bef5453c213b154b0781c9 Mon Sep 17 00:00:00 2001 From: khanova <32508607+khanova@users.noreply.github.com> Date: Thu, 16 Nov 2023 21:46:23 +0100 Subject: [PATCH] Proxy log pid hack (#5869) ## Problem Improve observability for the compute node. ## Summary of changes Log pid from the compute node. Doesn't work with pgbouncer. --- Cargo.lock | 10 +++---- Cargo.toml | 12 ++++---- proxy/src/compute.rs | 1 + proxy/src/proxy.rs | 2 +- proxy/src/serverless/conn_pool.rs | 40 +++++++++++++++++---------- proxy/src/serverless/sql_over_http.rs | 2 +- 6 files changed, 40 insertions(+), 27 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 841c60c7e4..f08ed2d744 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3221,7 +3221,7 @@ dependencies = [ [[package]] name = "postgres" version = "0.19.4" -source = "git+https://github.com/neondatabase/rust-postgres.git?rev=ce7260db5998fe27167da42503905a12e7ad9048#ce7260db5998fe27167da42503905a12e7ad9048" +source = "git+https://github.com/neondatabase/rust-postgres.git?rev=6ce32f791526e27533cab0232a6bb243b2c32584#6ce32f791526e27533cab0232a6bb243b2c32584" dependencies = [ "bytes", "fallible-iterator", @@ -3234,7 +3234,7 @@ dependencies = [ [[package]] name = "postgres-native-tls" version = "0.5.0" -source = "git+https://github.com/neondatabase/rust-postgres.git?rev=ce7260db5998fe27167da42503905a12e7ad9048#ce7260db5998fe27167da42503905a12e7ad9048" +source = "git+https://github.com/neondatabase/rust-postgres.git?rev=6ce32f791526e27533cab0232a6bb243b2c32584#6ce32f791526e27533cab0232a6bb243b2c32584" dependencies = [ "native-tls", "tokio", @@ -3245,7 +3245,7 @@ dependencies = [ [[package]] name = "postgres-protocol" version = "0.6.4" -source = "git+https://github.com/neondatabase/rust-postgres.git?rev=ce7260db5998fe27167da42503905a12e7ad9048#ce7260db5998fe27167da42503905a12e7ad9048" +source = "git+https://github.com/neondatabase/rust-postgres.git?rev=6ce32f791526e27533cab0232a6bb243b2c32584#6ce32f791526e27533cab0232a6bb243b2c32584" dependencies = [ "base64 0.20.0", "byteorder", @@ -3263,7 +3263,7 @@ dependencies = [ [[package]] name = "postgres-types" version = "0.2.4" -source = "git+https://github.com/neondatabase/rust-postgres.git?rev=ce7260db5998fe27167da42503905a12e7ad9048#ce7260db5998fe27167da42503905a12e7ad9048" +source = "git+https://github.com/neondatabase/rust-postgres.git?rev=6ce32f791526e27533cab0232a6bb243b2c32584#6ce32f791526e27533cab0232a6bb243b2c32584" dependencies = [ "bytes", "fallible-iterator", @@ -4933,7 +4933,7 @@ dependencies = [ [[package]] name = "tokio-postgres" version = "0.7.7" -source = "git+https://github.com/neondatabase/rust-postgres.git?rev=ce7260db5998fe27167da42503905a12e7ad9048#ce7260db5998fe27167da42503905a12e7ad9048" +source = "git+https://github.com/neondatabase/rust-postgres.git?rev=6ce32f791526e27533cab0232a6bb243b2c32584#6ce32f791526e27533cab0232a6bb243b2c32584" dependencies = [ "async-trait", "byteorder", diff --git a/Cargo.toml b/Cargo.toml index bfdb0442ab..7d7d74993e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -165,11 +165,11 @@ env_logger = "0.10" log = "0.4" ## Libraries from neondatabase/ git forks, ideally with changes to be upstreamed -postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="ce7260db5998fe27167da42503905a12e7ad9048" } -postgres-native-tls = { git = "https://github.com/neondatabase/rust-postgres.git", rev="ce7260db5998fe27167da42503905a12e7ad9048" } -postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="ce7260db5998fe27167da42503905a12e7ad9048" } -postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", rev="ce7260db5998fe27167da42503905a12e7ad9048" } -tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="ce7260db5998fe27167da42503905a12e7ad9048" } +postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="6ce32f791526e27533cab0232a6bb243b2c32584" } +postgres-native-tls = { git = "https://github.com/neondatabase/rust-postgres.git", rev="6ce32f791526e27533cab0232a6bb243b2c32584" } +postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="6ce32f791526e27533cab0232a6bb243b2c32584" } +postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", rev="6ce32f791526e27533cab0232a6bb243b2c32584" } +tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="6ce32f791526e27533cab0232a6bb243b2c32584" } ## Other git libraries heapless = { default-features=false, features=[], git = "https://github.com/japaric/heapless.git", rev = "644653bf3b831c6bb4963be2de24804acf5e5001" } # upstream release pending @@ -206,7 +206,7 @@ tonic-build = "0.9" # This is only needed for proxy's tests. # TODO: we should probably fork `tokio-postgres-rustls` instead. -tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="ce7260db5998fe27167da42503905a12e7ad9048" } +tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="6ce32f791526e27533cab0232a6bb243b2c32584" } ################# Binary contents sections diff --git a/proxy/src/compute.rs b/proxy/src/compute.rs index 53eb0e3a76..0741ad0623 100644 --- a/proxy/src/compute.rs +++ b/proxy/src/compute.rs @@ -248,6 +248,7 @@ impl ConnCfg { // connect_raw() will not use TLS if sslmode is "disable" let (client, connection) = self.0.connect_raw(stream, tls).await?; + tracing::Span::current().record("pid", &tracing::field::display(client.get_process_id())); let stream = connection.stream.into_inner(); info!( diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs index b03f95dc5f..b05b902303 100644 --- a/proxy/src/proxy.rs +++ b/proxy/src/proxy.rs @@ -514,7 +514,7 @@ pub fn invalidate_cache(node_info: console::CachedNodeInfo) -> compute::ConnCfg } /// Try to connect to the compute node once. -#[tracing::instrument(name = "connect_once", skip_all)] +#[tracing::instrument(name = "connect_once", fields(pid = tracing::field::Empty), skip_all)] async fn connect_to_compute_once( node_info: &console::CachedNodeInfo, timeout: time::Duration, diff --git a/proxy/src/serverless/conn_pool.rs b/proxy/src/serverless/conn_pool.rs index d09554a922..b753bc8918 100644 --- a/proxy/src/serverless/conn_pool.rs +++ b/proxy/src/serverless/conn_pool.rs @@ -208,14 +208,13 @@ impl GlobalConnPool { } else { info!("pool: reusing connection '{conn_info}'"); client.session.send(session_id)?; + tracing::Span::current().record( + "pid", + &tracing::field::display(client.inner.get_process_id()), + ); latency_timer.pool_hit(); latency_timer.success(); - return Ok(Client { - conn_id: client.conn_id, - inner: Some(client), - span: Span::current(), - pool, - }); + return Ok(Client::new(client, pool).await); } } else { let conn_id = uuid::Uuid::new_v4(); @@ -229,6 +228,12 @@ impl GlobalConnPool { ) .await }; + if let Ok(client) = &new_client { + tracing::Span::current().record( + "pid", + &tracing::field::display(client.inner.get_process_id()), + ); + } match &new_client { // clear the hash. it's no longer valid @@ -262,13 +267,8 @@ impl GlobalConnPool { } _ => {} } - - new_client.map(|inner| Client { - conn_id: inner.conn_id, - inner: Some(inner), - span: Span::current(), - pool, - }) + let new_client = new_client?; + Ok(Client::new(new_client, pool).await) } fn put(&self, conn_info: &ConnInfo, client: ClientInner) -> anyhow::Result<()> { @@ -394,7 +394,7 @@ impl ConnectMechanism for TokioMechanism<'_> { // Wake up the destination if needed. Code here is a bit involved because // we reuse the code from the usual proxy and we need to prepare few structures // that this code expects. -#[tracing::instrument(skip_all)] +#[tracing::instrument(fields(pid = tracing::field::Empty), skip_all)] async fn connect_to_compute( config: &config::ProxyConfig, conn_info: &ConnInfo, @@ -461,6 +461,7 @@ async fn connect_to_compute_once( .connect_timeout(timeout) .connect(tokio_postgres::NoTls) .await?; + tracing::Span::current().record("pid", &tracing::field::display(client.get_process_id())); let (tx, mut rx) = tokio::sync::watch::channel(session); @@ -547,6 +548,17 @@ pub struct Discard<'a> { } impl Client { + pub(self) async fn new( + inner: ClientInner, + pool: Option<(ConnInfo, Arc)>, + ) -> Self { + Self { + conn_id: inner.conn_id, + inner: Some(inner), + span: Span::current(), + pool, + } + } pub fn inner(&mut self) -> (&mut tokio_postgres::Client, Discard<'_>) { let Self { inner, diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 16736ac00d..4a9829e360 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -250,7 +250,7 @@ pub async fn handle( Ok(response) } -#[instrument(name = "sql-over-http", skip_all)] +#[instrument(name = "sql-over-http", fields(pid = tracing::field::Empty), skip_all)] async fn handle_inner( request: Request, sni_hostname: Option,