chore: Improve tracing, especially around running queries (#71)

sbernauer · web-flow · commit 37c1a29525a0 · 2025-03-12T11:53:12.000Z
* chore: Improve tracing for running queries

* Add TraceLayer

* changelog
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,8 +7,11 @@ All notable changes to this project will be documented in this file.
 ### Added
 
 - Support configuring compression for OTEL ([#70]).
+- Improve tracing details by adding a `tower_http::trace::TraceLayer` that creates spans for every HTTP request ([#71]).
+- Improve tracing for running queries on Trino, adding spans for the request to Trino and parsing ([#71]).
 
 [#70]: https://github.com/stackabletech/trino-lb/pull/70
+[#71]: https://github.com/stackabletech/trino-lb/pull/71
 
 ## [0.4.1] - 2025-03-03
 
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -82,6 +82,7 @@ sqlx = { version = "0.8.2", features = [
 strum = { version = "0.27", features = ["derive"] }
 tokio = "1.39"
 tower = "0.5"
+tower-http = { version = "0.6", features = ["tracing"] }
 tracing = "0.1"
 tracing-opentelemetry = "0.25"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
diff --git a/trino-lb-core/src/trino_api.rs b/trino-lb-core/src/trino_api.rs
@@ -82,6 +82,7 @@ pub struct Stat {
 
 impl TrinoQueryApiResponse {
     #[instrument(
+        skip(query),
         fields(trino_lb_addr = %trino_lb_addr),
     )]
     pub fn new_from_queued_query(
@@ -151,6 +152,7 @@ impl TrinoQueryApiResponse {
     }
 
     #[instrument(
+        skip(self),
         fields(trino_lb_addr = %trino_lb_addr),
     )]
     pub fn change_next_uri_to_trino_lb(&mut self, trino_lb_addr: &Url) -> Result<(), Error> {
diff --git a/trino-lb-persistence/src/in_memory.rs b/trino-lb-persistence/src/in_memory.rs
@@ -54,7 +54,7 @@ impl Default for InMemoryPersistence {
 }
 
 impl Persistence for InMemoryPersistence {
-    #[instrument(skip(self))]
+    #[instrument(skip(self, queued_query))]
     async fn store_queued_query(&self, queued_query: QueuedQuery) -> Result<(), super::Error> {
         let mut queued_queries = self.queued_queries.write().await;
         queued_queries.insert(queued_query.id.clone(), queued_query);
@@ -74,15 +74,15 @@ impl Persistence for InMemoryPersistence {
             .clone())
     }
 
-    #[instrument(skip(self))]
+    #[instrument(skip(self, queued_query))]
     async fn remove_queued_query(&self, queued_query: &QueuedQuery) -> Result<(), super::Error> {
         let mut queued_queries = self.queued_queries.write().await;
         queued_queries.remove(&queued_query.id);
 
         Ok(())
     }
 
-    #[instrument(skip(self))]
+    #[instrument(skip(self, query))]
     async fn store_query(&self, query: TrinoQuery) -> Result<(), super::Error> {
         let mut queries = self.queries.write().await;
         queries.insert(query.id.clone(), query);
diff --git a/trino-lb-persistence/src/postgres/mod.rs b/trino-lb-persistence/src/postgres/mod.rs
@@ -135,7 +135,7 @@ struct HeaderMapWrapper {
 }
 
 impl Persistence for PostgresPersistence {
-    #[instrument(skip(self))]
+    #[instrument(skip(self, queued_query))]
     async fn store_queued_query(&self, queued_query: QueuedQuery) -> Result<(), super::Error> {
         query!(
             r#"INSERT INTO queued_queries (id, query, headers, creation_time, last_accessed, cluster_group)
@@ -185,7 +185,7 @@ impl Persistence for PostgresPersistence {
         Ok(queued_query)
     }
 
-    #[instrument(skip(self))]
+    #[instrument(skip(self, queued_query))]
     async fn remove_queued_query(&self, queued_query: &QueuedQuery) -> Result<(), super::Error> {
         query!(
             r#"DELETE FROM queued_queries
@@ -199,7 +199,7 @@ impl Persistence for PostgresPersistence {
         Ok(())
     }
 
-    #[instrument(skip(self))]
+    #[instrument(skip(self, query))]
     async fn store_query(&self, query: TrinoQuery) -> Result<(), super::Error> {
         query!(
             r#"INSERT INTO queries (id, trino_cluster, trino_endpoint, creation_time, delivered_time)
diff --git a/trino-lb-persistence/src/redis/mod.rs b/trino-lb-persistence/src/redis/mod.rs
@@ -174,7 +174,7 @@ impl<R> Persistence for RedisPersistence<R>
 where
     R: AsyncCommands + Clone,
 {
-    #[instrument(skip(self))]
+    #[instrument(skip(self, queued_query))]
     async fn store_queued_query(&self, queued_query: QueuedQuery) -> Result<(), super::Error> {
         let key = queued_query_key(&queued_query.id);
         let value = bincode::serialize(&queued_query).context(SerializeToBinarySnafu)?;
@@ -210,7 +210,7 @@ where
         Ok(bincode::deserialize(&value).context(DeserializeFromBinarySnafu)?)
     }
 
-    #[instrument(skip(self))]
+    #[instrument(skip(self, queued_query))]
     async fn remove_queued_query(&self, queued_query: &QueuedQuery) -> Result<(), super::Error> {
         let key = queued_query_key(&queued_query.id);
         let mut connection = self.connection();
@@ -225,7 +225,7 @@ where
         Ok(())
     }
 
-    #[instrument(skip(self))]
+    #[instrument(skip(self, query))]
     async fn store_query(&self, query: TrinoQuery) -> Result<(), super::Error> {
         let key = query_key(&query.id);
         let value = bincode::serialize(&query).context(SerializeToBinarySnafu)?;
diff --git a/trino-lb/Cargo.toml b/trino-lb/Cargo.toml
@@ -48,6 +48,7 @@ snafu.workspace = true
 strum.workspace = true
 tokio.workspace = true
 tower.workspace = true
+tower-http.workspace = true
 tracing-opentelemetry.workspace = true
 tracing-subscriber.workspace = true
 tracing.workspace = true
diff --git a/trino-lb/src/cluster_group_manager.rs b/trino-lb/src/cluster_group_manager.rs
@@ -9,7 +9,7 @@ use futures::future::try_join_all;
 use http::{HeaderMap, StatusCode};
 use reqwest::Client;
 use snafu::{OptionExt, ResultExt, Snafu};
-use tracing::{debug, instrument};
+use tracing::{debug, info_span, instrument, Instrument};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 use trino_lb_core::{
     config::Config, sanitization::Sanitize, trino_api::TrinoQueryApiResponse,
@@ -37,6 +37,12 @@ pub enum Error {
     #[snafu(display("Failed to decode Trino API response"))]
     DecodeTrinoResponse { source: reqwest::Error },
 
+    #[snafu(display("Failed to get the bytes of the Trino API response"))]
+    GetTrinoResponseBytes { source: reqwest::Error },
+
+    #[snafu(display("Failed to parse Trino API response as JSON"))]
+    ParseTrinoResponse { source: serde_json::Error },
+
     #[snafu(display("Configuration error: A specific Trino cluster can only be part of a single clusterGroup. Please make sure the Trino cluster {cluster_name:?} only is part of a single clusterGroup."))]
     ConfigErrorTrinoClusterInMultipleClusterGroups { cluster_name: String },
 
@@ -111,6 +117,7 @@ impl IntoResponse for SendToTrinoResponse {
 }
 
 impl ClusterGroupManager {
+    // Intentionally including the config here, this is only logged on startup
     #[instrument(skip(persistence))]
     pub fn new(
         persistence: Arc<PersistenceImplementation>,
@@ -152,7 +159,10 @@ impl ClusterGroupManager {
         })
     }
 
-    #[instrument(skip(self))]
+    #[instrument(
+        skip(self, cluster),
+        fields(cluster.name, headers = ?headers.sanitize())
+    )]
     pub async fn send_query_to_cluster(
         &self,
         query: String,
@@ -188,7 +198,7 @@ impl ClusterGroupManager {
             let body = response
                 .bytes()
                 .await
-                .context(DecodeTrinoResponseSnafu)?
+                .context(GetTrinoResponseBytesSnafu)?
                 .into();
             return Ok(SendToTrinoResponse::Unauthorized { headers, body });
         }
@@ -217,12 +227,19 @@ impl ClusterGroupManager {
             .get(next_uri)
             .headers(headers)
             .send()
+            .instrument(info_span!("Send HTTP get to Trino"))
             .await
             .context(ContactTrinoPostQuerySnafu)?;
         let headers = response.headers();
-
         let headers = filter_to_trino_headers(headers);
-        let trino_query_api_response = response.json().await.context(DecodeTrinoResponseSnafu)?;
+
+        let bytes = response
+            .bytes()
+            .instrument(info_span!("Get response bytes"))
+            .await
+            .context(GetTrinoResponseBytesSnafu)?;
+        let trino_query_api_response = info_span!("Parse JSON response", bytes = bytes.len())
+            .in_scope(|| serde_json::from_slice(&bytes).context(ParseTrinoResponseSnafu))?;
 
         Ok((trino_query_api_response, headers))
     }
diff --git a/trino-lb/src/http_server/mod.rs b/trino-lb/src/http_server/mod.rs
@@ -15,6 +15,7 @@ use axum_server::{tls_rustls::RustlsConfig, Handle};
 use futures::FutureExt;
 use snafu::{OptionExt, ResultExt, Snafu};
 use tokio::time::sleep;
+use tower_http::trace::TraceLayer;
 use tracing::info;
 use trino_lb_persistence::PersistenceImplementation;
 
@@ -122,6 +123,7 @@ pub async fn start_http_server(
         )
         .route("/ui/index.html", get(ui::index::get_ui_index))
         .route("/ui/query.html", get(ui::query::get_ui_query))
+        .layer(TraceLayer::new_for_http())
         .with_state(app_state);
 
     if tls_config.enabled {
diff --git a/trino-lb/src/http_server/v1/statement.rs b/trino-lb/src/http_server/v1/statement.rs
@@ -235,7 +235,7 @@ pub async fn get_trino_executing_statement(
     handle_query_running_on_trino(&state, headers, query_id, uri.path()).await
 }
 
-#[instrument(skip(state))]
+#[instrument(skip(state, queued_query))]
 async fn queue_or_hand_over_query(
     state: &Arc<AppState>,
     queued_query: QueuedQuery,
diff --git a/trino-lb/src/maintenance/query_count_fetcher.rs b/trino-lb/src/maintenance/query_count_fetcher.rs
@@ -28,6 +28,7 @@ pub struct QueryCountFetcher {
 }
 
 impl QueryCountFetcher {
+    // Intentionally including the config here, this is only logged on startup
     #[instrument(skip(persistence, metrics))]
     pub fn new(
         persistence: Arc<PersistenceImplementation>,
@@ -154,7 +155,7 @@ impl QueryCountFetcher {
     /// - In case we know the cluster is not reachable (e.g. the cluster is turned off or currently
     ///   starting), store a query count of zero (0) to avoid dangling clusters with non-zero query
     ///   counts.
-    #[instrument(skip(self, cluster), fields(cluster_name = cluster.name))]
+    #[instrument(skip(self, cluster), fields(cluster.name))]
     async fn process_cluster(&self, cluster: &TrinoClusterConfig, state: ClusterState) {
         match state {
             ClusterState::Ready | ClusterState::Unhealthy | ClusterState::Draining { .. } => {
@@ -180,7 +181,7 @@ impl QueryCountFetcher {
         }
     }
 
-    #[instrument(skip(self, cluster), fields(cluster_name = cluster.name))]
+    #[instrument(skip(self, cluster), fields(cluster.name))]
     async fn fetch_and_store_query_count(&self, cluster: &TrinoClusterConfig) {
         let cluster_info =
             get_cluster_info(&cluster.endpoint, self.ignore_certs, &cluster.credentials).await;
diff --git a/trino-lb/src/routing/client_tags.rs b/trino-lb/src/routing/client_tags.rs
@@ -25,6 +25,7 @@ pub struct ClientTagsRouter {
 }
 
 impl ClientTagsRouter {
+    // Intentionally including the config here, this is only logged on startup
     #[instrument(name = "ClientTagsRouter::new")]
     pub fn new(
         config: &ClientTagsRouterConfig,
diff --git a/trino-lb/src/routing/explain_costs.rs b/trino-lb/src/routing/explain_costs.rs
@@ -27,6 +27,7 @@ pub struct ExplainCostsRouter {
 }
 
 impl ExplainCostsRouter {
+    // Intentionally including the config here, this is only logged on startup
     #[instrument(name = "ExplainCostsRouter::new")]
     pub fn new(
         config: &ExplainCostsRouterConfig,
diff --git a/trino-lb/src/routing/mod.rs b/trino-lb/src/routing/mod.rs
@@ -42,6 +42,7 @@ pub struct Router {
 }
 
 impl Router {
+    // Intentionally including the config here, this is only logged on startup
     #[instrument]
     pub fn new(config: &Config) -> Result<Self, Error> {
         let mut routers = Vec::with_capacity(config.routers.len());
diff --git a/trino-lb/src/routing/python_script.rs b/trino-lb/src/routing/python_script.rs
@@ -41,6 +41,7 @@ pub struct PythonScriptRouter {
 }
 
 impl PythonScriptRouter {
+    // Intentionally including the config here, this is only logged on startup
     #[instrument(name = "PythonScriptRouter::new")]
     pub fn new(
         config: &PythonScriptRouterConfig,
diff --git a/trino-lb/src/routing/trino_routing_group_header.rs b/trino-lb/src/routing/trino_routing_group_header.rs
@@ -11,6 +11,7 @@ pub struct TrinoRoutingGroupHeaderRouter {
 }
 
 impl TrinoRoutingGroupHeaderRouter {
+    // Intentionally including the config here, this is only logged on startup
     #[instrument(name = "TrinoRoutingGroupHeaderRouter::new")]
     pub fn new(
         config: &TrinoRoutingGroupHeaderRouterConfig,
diff --git a/trino-lb/src/scaling/mod.rs b/trino-lb/src/scaling/mod.rs
@@ -134,6 +134,7 @@ pub struct Scaler {
 }
 
 impl Scaler {
+    // Intentionally including the config here, this is only logged on startup
     #[instrument(skip(persistence))]
     pub async fn new(
         config: &Config,
@@ -553,7 +554,7 @@ impl Scaler {
     #[instrument(
         name = "Scaler::apply_cluster_target_state",
         skip(self, cluster),
-        fields(%cluster.name)
+        fields(cluster.name)
     )]
     async fn apply_cluster_target_state(
         self: Arc<Self>,
@@ -606,7 +607,7 @@ impl Scaler {
         Ok(())
     }
 
-    #[instrument(skip(self))]
+    #[instrument(skip(self, scaling_config))]
     fn get_current_min_cluster_count(
         &self,
         scaling_config: &TrinoClusterGroupAutoscaling,
diff --git a/trino-lb/src/scaling/stackable.rs b/trino-lb/src/scaling/stackable.rs
@@ -138,6 +138,7 @@ struct StackableTrinoCluster {
 }
 
 impl StackableScaler {
+    // Intentionally including the config here, this is only logged on startup
     #[instrument(name = "StackableScaler::new")]
     pub async fn new(
         config: &StackableScalerConfig,
diff --git a/trino-lb/src/tracing.rs b/trino-lb/src/tracing.rs
@@ -99,7 +99,6 @@ fn otel_tracer(tracing_config: &TrinoLbTracingConfig) -> Result<trace::Tracer, E
             trace::Config::default()
                 .with_sampler(Sampler::AlwaysOn)
                 .with_id_generator(RandomIdGenerator::default())
-                .with_max_events_per_span(64)
                 .with_max_attributes_per_span(16)
                 .with_max_events_per_span(16)
                 .with_resource(Resource::new(vec![KeyValue::new(
diff --git a/trino-lb/src/trino_client/cluster_info.rs b/trino-lb/src/trino_client/cluster_info.rs
@@ -52,7 +52,7 @@ pub struct ClusterInfo {
     pub total_cpu_time_secs: u64,
 }
 
-#[instrument]
+#[instrument(skip(credentials))]
 pub async fn get_cluster_info(
     endpoint: &Url,
     ignore_certs: bool,
diff --git a/trino-lb/src/trino_client/mod.rs b/trino-lb/src/trino_client/mod.rs
@@ -139,7 +139,7 @@ impl TrinoClient {
     }
 }
 
-#[instrument]
+#[instrument(skip(config))]
 fn trino_client_builder_from_config(config: &TrinoClientConfig) -> Result<ClientBuilder, Error> {
     Ok(ClientBuilder::new(
         &config.username,
