add timing infos to response header

Author: MarinPostma

libsql-server/src/connection/libsql.rs

@@ -23,7 +23,7 @@ use crate::query_analysis::StmtKind;
 use crate::query_result_builder::{QueryBuilderConfig, QueryResultBuilder};
 use crate::replication::FrameNo;
 use crate::stats::{Stats, StatsUpdateMessage};
-use crate::{Result, BLOCKING_RT};
+use crate::{record_time, Result, BLOCKING_RT};
 
 use super::connection_manager::{
     ConnectionManager, InnerWalManager, ManagedConnectionWal, ManagedConnectionWalWrapper,
@@ -713,7 +713,10 @@ where
         builder: B,
         _replication_index: Option<FrameNo>,
     ) -> Result<B> {
-        self.execute(pgm, ctx, builder).await.map(|(b, _)| b)
+        record_time! {
+            "libsql_query_exec";
+            self.execute(pgm, ctx, builder).await.map(|(b, _)| b)
+        }
     }
 
     async fn describe(

libsql-server/src/http/user/mod.rs

@@ -6,6 +6,8 @@ mod listen;
 mod result_builder;
 mod trace;
 mod types;
+#[macro_use]
+pub mod timing;
 
 use std::path::Path;
 use std::sync::Arc;
@@ -16,7 +18,7 @@ use axum::http::request::Parts;
 use axum::http::HeaderValue;
 use axum::response::{Html, IntoResponse};
 use axum::routing::{get, post};
-use axum::Router;
+use axum::{middleware, Router};
 use axum_extra::middleware::option_layer;
 use base64::prelude::BASE64_STANDARD_NO_PAD;
 use base64::Engine;
@@ -37,6 +39,7 @@ use crate::connection::{Connection, RequestContext};
 use crate::error::Error;
 use crate::hrana;
 use crate::http::user::db_factory::MakeConnectionExtractorPath;
+use crate::http::user::timing::timings_middleware;
 use crate::http::user::types::HttpQuery;
 use crate::metrics::LEGACY_HTTP_CALL;
 use crate::namespace::NamespaceStore;
@@ -407,6 +410,7 @@ where
                 )
                 .route("/v1/jobs", get(handle_get_migrations))
                 .route("/v1/jobs/:job_id", get(handle_get_migration_details))
+                .layer(middleware::from_fn(timings_middleware))
                 .with_state(state);
 
             // Merge the grpc based axum router into our regular http router

libsql-server/src/http/user/timing.rs

@@ -0,0 +1,63 @@
+use std::fmt::Write as _;
+use std::sync::Arc;
+use std::time::Duration;
+
+use axum::http::Request;
+use axum::middleware::Next;
+use axum::response::Response;
+use hashbrown::HashMap;
+use parking_lot::Mutex;
+
+#[derive(Default, Clone, Debug)]
+pub struct Timings {
+    records: Arc<Mutex<HashMap<&'static str, Duration>>>,
+}
+
+impl Timings {
+    pub fn record(&self, k: &'static str, d: Duration) {
+        self.records.lock().insert(k, d);
+    }
+
+    fn format(&self) -> String {
+        let mut out = String::new();
+        let records = self.records.lock();
+        for (k, v) in records.iter() {
+            write!(&mut out, "{k};dur={v:?},").unwrap();
+        }
+        out
+    }
+}
+
+tokio::task_local! {
+    pub static TIMINGS: Timings;
+}
+
+#[macro_export]
+macro_rules! record_time {
+    ($k:literal; $($rest:tt)*) => {
+        {
+            let __before__ = std::time::Instant::now();
+            let __ret__ = {
+                $($rest)*
+            };
+            let _ = $crate::http::user::timing::TIMINGS.try_with(|t| t.record($k, __before__.elapsed()));
+            __ret__
+        }
+    };
+}
+
+pub(crate) async fn timings_middleware<B>(request: Request<B>, next: Next<B>) -> Response {
+    TIMINGS
+        .scope(Default::default(), async move {
+            let mut response = record_time! {
+                "query_total";
+                next.run(request).await
+            };
+            let timings = TIMINGS.get().format();
+            response
+                .headers_mut()
+                .insert("Server-Timing", timings.parse().unwrap());
+            response
+        })
+        .await
+}