proxy: add request context for observability and blocking (#6160)

## Summary of changes

### RequestMonitoring

We want to add an event stream with information on each request for
easier analysis than what we can do with diagnostic logs alone
(https://github.com/neondatabase/cloud/issues/8807). This
RequestMonitoring will keep a record of the final state of a request. On
drop it will be pushed into a queue to be uploaded.

Because this context is a bag of data, I don't want this information to
impact logic of request handling. I personally think that weakly typed
data (such as all these options) makes for spaghetti code. I will
however allow for this data to impact rate-limiting and blocking of
requests, as this does not _really_ change how a request is handled.

### Parquet

Each `RequestMonitoring` is flushed into a channel where it is converted
into `RequestData`, which is accumulated into parquet files. Each file
will have a certain number of rows per row group, and several row groups
will eventually fill up the file, which we then upload to S3.

We will also upload smaller files if they take too long to construct.

proxy/src/serverless/conn_pool.rs

@@ -13,7 +13,7 @@ use pq_proto::StartupMessageParams;
 use prometheus::{exponential_buckets, register_histogram, Histogram};
 use rand::Rng;
 use smol_str::SmolStr;
-use std::{collections::HashMap, net::IpAddr, pin::pin, sync::Arc, sync::Weak, time::Duration};
+use std::{collections::HashMap, pin::pin, sync::Arc, sync::Weak, time::Duration};
 use std::{
     fmt,
     task::{ready, Poll},
@@ -28,7 +28,8 @@ use tokio_postgres::{AsyncMessage, ReadyForQueryStatus};
 use crate::{
     auth::{self, backend::ComputeUserInfo, check_peer_addr_is_in_list},
     console,
-    metrics::{LatencyTimer, NUM_DB_CONNECTIONS_GAUGE},
+    context::RequestMonitoring,
+    metrics::NUM_DB_CONNECTIONS_GAUGE,
     proxy::{connect_compute::ConnectMechanism, neon_options},
     usage_metrics::{Ids, MetricCounter, USAGE_METRICS},
 };
@@ -309,13 +310,11 @@ impl GlobalConnPool {
 
     pub async fn get(
         self: &Arc<Self>,
+        ctx: &mut RequestMonitoring,
         conn_info: ConnInfo,
         force_new: bool,
-        session_id: uuid::Uuid,
-        peer_addr: IpAddr,
     ) -> anyhow::Result<Client> {
         let mut client: Option<ClientInner> = None;
-        let mut latency_timer = LatencyTimer::new("http");
 
         let mut hash_valid = false;
         let mut endpoint_pool = Weak::new();
@@ -360,23 +359,21 @@ impl GlobalConnPool {
                 info!(%conn_id, "pool: cached connection '{conn_info}' is closed, opening a new one");
                 connect_to_compute(
                     self.proxy_config,
+                    ctx,
                     &conn_info,
                     conn_id,
-                    session_id,
-                    latency_timer,
-                    peer_addr,
                     endpoint_pool.clone(),
                 )
                 .await
             } else {
                 info!("pool: reusing connection '{conn_info}'");
-                client.session.send(session_id)?;
+                client.session.send(ctx.session_id)?;
                 tracing::Span::current().record(
                     "pid",
                     &tracing::field::display(client.inner.get_process_id()),
                 );
-                latency_timer.pool_hit();
-                latency_timer.success();
+                ctx.latency_timer.pool_hit();
+                ctx.latency_timer.success();
                 return Ok(Client::new(client, conn_info, endpoint_pool).await);
             }
         } else {
@@ -384,11 +381,9 @@ impl GlobalConnPool {
             info!(%conn_id, "pool: opening a new connection '{conn_info}'");
             connect_to_compute(
                 self.proxy_config,
+                ctx,
                 &conn_info,
                 conn_id,
-                session_id,
-                latency_timer,
-                peer_addr,
                 endpoint_pool.clone(),
             )
             .await
@@ -483,7 +478,6 @@ impl GlobalConnPool {
 struct TokioMechanism<'a> {
     pool: Weak<RwLock<EndpointConnPool>>,
     conn_info: &'a ConnInfo,
-    session_id: uuid::Uuid,
     conn_id: uuid::Uuid,
     idle: Duration,
 }
@@ -496,15 +490,16 @@ impl ConnectMechanism for TokioMechanism<'_> {
 
     async fn connect_once(
         &self,
+        ctx: &mut RequestMonitoring,
         node_info: &console::CachedNodeInfo,
         timeout: time::Duration,
     ) -> Result<Self::Connection, Self::ConnectError> {
         connect_to_compute_once(
+            ctx,
             node_info,
             self.conn_info,
             timeout,
             self.conn_id,
-            self.session_id,
             self.pool.clone(),
             self.idle,
         )
@@ -520,11 +515,9 @@ impl ConnectMechanism for TokioMechanism<'_> {
 #[tracing::instrument(fields(pid = tracing::field::Empty), skip_all)]
 async fn connect_to_compute(
     config: &config::ProxyConfig,
+    ctx: &mut RequestMonitoring,
     conn_info: &ConnInfo,
     conn_id: uuid::Uuid,
-    session_id: uuid::Uuid,
-    latency_timer: LatencyTimer,
-    peer_addr: IpAddr,
     pool: Weak<RwLock<EndpointConnPool>>,
 ) -> anyhow::Result<ClientInner> {
     let tls = config.tls_config.as_ref();
@@ -536,12 +529,8 @@ async fn connect_to_compute(
         ("application_name", APP_NAME),
         ("options", conn_info.options.as_deref().unwrap_or("")),
     ]);
-    let creds = auth::ClientCredentials::parse(
-        &params,
-        Some(&conn_info.hostname),
-        common_names,
-        peer_addr,
-    )?;
+    let creds =
+        auth::ClientCredentials::parse(ctx, &params, Some(&conn_info.hostname), common_names)?;
 
     let creds =
         ComputeUserInfo::try_from(creds).map_err(|_| anyhow!("missing endpoint identifier"))?;
@@ -549,48 +538,48 @@ async fn connect_to_compute(
 
     let console_options = neon_options(&params);
 
-    let extra = console::ConsoleReqExtra {
-        session_id: uuid::Uuid::new_v4(),
-        application_name: APP_NAME.to_string(),
-        options: console_options,
-    };
     if !config.disable_ip_check_for_http {
-        let allowed_ips = backend.get_allowed_ips(&extra).await?;
-        if !check_peer_addr_is_in_list(&peer_addr, &allowed_ips) {
+        let allowed_ips = backend.get_allowed_ips(ctx).await?;
+        if !check_peer_addr_is_in_list(&ctx.peer_addr, &allowed_ips) {
             return Err(auth::AuthError::ip_address_not_allowed().into());
         }
     }
+    let extra = console::ConsoleReqExtra {
+        options: console_options,
+    };
     let node_info = backend
-        .wake_compute(&extra)
+        .wake_compute(ctx, &extra)
         .await?
         .context("missing cache entry from wake_compute")?;
 
+    ctx.set_project(node_info.aux.clone());
+
     crate::proxy::connect_compute::connect_to_compute(
+        ctx,
         &TokioMechanism {
             conn_id,
             conn_info,
-            session_id,
             pool,
             idle: config.http_config.pool_options.idle_timeout,
         },
         node_info,
         &extra,
         &backend,
-        latency_timer,
     )
     .await
 }
 
 async fn connect_to_compute_once(
+    ctx: &mut RequestMonitoring,
     node_info: &console::CachedNodeInfo,
     conn_info: &ConnInfo,
     timeout: time::Duration,
     conn_id: uuid::Uuid,
-    mut session: uuid::Uuid,
     pool: Weak<RwLock<EndpointConnPool>>,
     idle: Duration,
 ) -> Result<ClientInner, tokio_postgres::Error> {
     let mut config = (*node_info.config).clone();
+    let mut session = ctx.session_id;
 
     let (client, mut connection) = config
         .user(&conn_info.username)
@@ -601,7 +590,7 @@ async fn connect_to_compute_once(
         .await?;
 
     let conn_gauge = NUM_DB_CONNECTIONS_GAUGE
-        .with_label_values(&["http"])
+        .with_label_values(&[ctx.protocol])
         .guard();
 
     tracing::Span::current().record("pid", &tracing::field::display(client.get_process_id()));

proxy/src/serverless/sql_over_http.rs

@@ -1,4 +1,3 @@
-use std::net::IpAddr;
 use std::sync::Arc;
 
 use anyhow::bail;
@@ -14,6 +13,7 @@ use hyper::{Body, HeaderMap, Request};
 use serde_json::json;
 use serde_json::Map;
 use serde_json::Value;
+use smol_str::SmolStr;
 use tokio_postgres::error::DbError;
 use tokio_postgres::types::Kind;
 use tokio_postgres::types::Type;
@@ -29,6 +29,7 @@ use utils::http::error::ApiError;
 use utils::http::json::json_response;
 
 use crate::config::HttpConfig;
+use crate::context::RequestMonitoring;
 use crate::metrics::NUM_CONNECTION_REQUESTS_GAUGE;
 
 use super::conn_pool::ConnInfo;
@@ -121,6 +122,7 @@ fn json_array_to_pg_array(value: &Value) -> Option<String> {
 }
 
 fn get_conn_info(
+    ctx: &mut RequestMonitoring,
     headers: &HeaderMap,
     sni_hostname: Option<String>,
 ) -> Result<ConnInfo, anyhow::Error> {
@@ -146,10 +148,11 @@ fn get_conn_info(
         .next()
         .ok_or(anyhow::anyhow!("invalid database name"))?;
 
-    let username = connection_url.username();
+    let username = SmolStr::from(connection_url.username());
     if username.is_empty() {
         return Err(anyhow::anyhow!("missing username"));
     }
+    ctx.set_user(username.clone());
 
     let password = connection_url
         .password()
@@ -176,6 +179,9 @@ fn get_conn_info(
         }
     }
 
+    let hostname: SmolStr = hostname.into();
+    ctx.set_endpoint_id(Some(hostname.clone()));
+
     let pairs = connection_url.query_pairs();
 
     let mut options = Option::None;
@@ -188,9 +194,9 @@ fn get_conn_info(
     }
 
     Ok(ConnInfo {
-        username: username.into(),
+        username,
         dbname: dbname.into(),
-        hostname: hostname.into(),
+        hostname,
         password: password.into(),
         options,
     })
@@ -198,23 +204,15 @@ fn get_conn_info(
 
 // TODO: return different http error codes
 pub async fn handle(
+    config: &'static HttpConfig,
+    ctx: &mut RequestMonitoring,
     request: Request<Body>,
     sni_hostname: Option<String>,
     conn_pool: Arc<GlobalConnPool>,
-    session_id: uuid::Uuid,
-    peer_addr: IpAddr,
-    config: &'static HttpConfig,
 ) -> Result<Response<Body>, ApiError> {
     let result = tokio::time::timeout(
         config.request_timeout,
-        handle_inner(
-            config,
-            request,
-            sni_hostname,
-            conn_pool,
-            session_id,
-            peer_addr,
-        ),
+        handle_inner(config, ctx, request, sni_hostname, conn_pool),
     )
     .await;
     let mut response = match result {
@@ -297,11 +295,10 @@ pub async fn handle(
 #[instrument(name = "sql-over-http", fields(pid = tracing::field::Empty), skip_all)]
 async fn handle_inner(
     config: &'static HttpConfig,
+    ctx: &mut RequestMonitoring,
     request: Request<Body>,
     sni_hostname: Option<String>,
     conn_pool: Arc<GlobalConnPool>,
-    session_id: uuid::Uuid,
-    peer_addr: IpAddr,
 ) -> anyhow::Result<Response<Body>> {
     let _request_gauge = NUM_CONNECTION_REQUESTS_GAUGE
         .with_label_values(&["http"])
@@ -311,7 +308,7 @@ async fn handle_inner(
     // Determine the destination and connection params
     //
     let headers = request.headers();
-    let conn_info = get_conn_info(headers, sni_hostname)?;
+    let conn_info = get_conn_info(ctx, headers, sni_hostname)?;
 
     // Determine the output options. Default behaviour is 'false'. Anything that is not
     // strictly 'true' assumed to be false.
@@ -340,10 +337,12 @@ async fn handle_inner(
     let txn_read_only = headers.get(&TXN_READ_ONLY) == Some(&HEADER_VALUE_TRUE);
     let txn_deferrable = headers.get(&TXN_DEFERRABLE) == Some(&HEADER_VALUE_TRUE);
 
+    let paused = ctx.latency_timer.pause();
     let request_content_length = match request.body().size_hint().upper() {
         Some(v) => v,
         None => MAX_REQUEST_SIZE + 1,
     };
+    drop(paused);
 
     // we don't have a streaming request support yet so this is to prevent OOM
     // from a malicious user sending an extremely large request body
@@ -359,9 +358,7 @@ async fn handle_inner(
     let body = hyper::body::to_bytes(request.into_body()).await?;
     let payload: Payload = serde_json::from_slice(&body)?;
 
-    let mut client = conn_pool
-        .get(conn_info, !allow_pool, session_id, peer_addr)
-        .await?;
+    let mut client = conn_pool.get(ctx, conn_info, !allow_pool).await?;
 
     let mut response = Response::builder()
         .status(StatusCode::OK)
@@ -449,6 +446,7 @@ async fn handle_inner(
             }
         };
 
+    ctx.log();
     let metrics = client.metrics();
 
     // how could this possibly fail

proxy/src/serverless/websocket.rs

@@ -1,6 +1,7 @@
 use crate::{
     cancellation::CancelMap,
     config::ProxyConfig,
+    context::RequestMonitoring,
     error::io_error,
     proxy::{handle_client, ClientMode},
     rate_limiter::EndpointRateLimiter,
@@ -12,7 +13,6 @@ use hyper_tungstenite::{tungstenite::Message, HyperWebsocket, WebSocketStream};
 use pin_project_lite::pin_project;
 
 use std::{
-    net::IpAddr,
     pin::Pin,
     sync::Arc,
     task::{ready, Context, Poll},
@@ -130,22 +130,20 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AsyncBufRead for WebSocketRw<S> {
 }
 
 pub async fn serve_websocket(
-    websocket: HyperWebsocket,
     config: &'static ProxyConfig,
+    ctx: &mut RequestMonitoring,
+    websocket: HyperWebsocket,
     cancel_map: &CancelMap,
-    session_id: uuid::Uuid,
     hostname: Option<String>,
-    peer_addr: IpAddr,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
 ) -> anyhow::Result<()> {
     let websocket = websocket.await?;
     handle_client(
         config,
+        ctx,
         cancel_map,
-        session_id,
         WebSocketRw::new(websocket),
         ClientMode::Websockets { hostname },
-        peer_addr,
         endpoint_rate_limiter,
     )
     .await?;