feat: Add Server-Timing header to track request performance

Co-authored-by: contact <contact@ophir.dev>

Author: cursoragent

src/render.rs

@@ -368,13 +368,15 @@ impl HeaderContext {
         Ok(PageContext::Header(self))
     }
 
-    async fn start_body(mut self, data: JsonValue) -> anyhow::Result<PageContext> {
-        if let Some(ref timing) = self.request_context.server_timing {
-            let header_value = timing.as_header_value();
-            if !header_value.is_empty() {
-                self.response.insert_header(("Server-Timing", header_value));
-            }
+    fn add_server_timing_header(&mut self) {
+        let header_value = self.request_context.server_timing.as_header_value();
+        if !header_value.is_empty() {
+            self.response.insert_header(("Server-Timing", header_value));
         }
+    }
+
+    async fn start_body(mut self, data: JsonValue) -> anyhow::Result<PageContext> {
+        self.add_server_timing_header();
         let html_renderer =
             HtmlRenderContext::new(self.app_state, self.request_context, self.writer, data)
                 .await
@@ -388,12 +390,7 @@ impl HeaderContext {
     }
 
     pub fn close(mut self) -> HttpResponse {
-        if let Some(ref timing) = self.request_context.server_timing {
-            let header_value = timing.as_header_value();
-            if !header_value.is_empty() {
-                self.response.insert_header(("Server-Timing", header_value));
-            }
-        }
+        self.add_server_timing_header();
         self.response.finish()
     }
 }

src/webserver/database/execute_queries.rs

@@ -52,13 +52,13 @@ pub fn stream_query_results_with_conn<'a>(
         for res in &sql_file.statements {
             match res {
                 ParsedStatement::CsvImport(csv_import) => {
-                    let connection = take_connection(&request.app_state.db, db_connection).await?;
+                    let connection = take_connection(&request.app_state.db, db_connection, request).await?;
                     log::debug!("Executing CSV import: {csv_import:?}");
                     run_csv_import(connection, csv_import, request).await.with_context(|| format!("Failed to import the CSV file {:?} into the table {:?}", csv_import.uploaded_file, csv_import.table_name))?;
                 },
                 ParsedStatement::StmtWithParams(stmt) => {
                     let query = bind_parameters(stmt, request, db_connection).await?;
-                    let connection = take_connection(&request.app_state.db, db_connection).await?;
+                    let connection = take_connection(&request.app_state.db, db_connection, request).await?;
                     log::trace!("Executing query {:?}", query.sql);
                     let mut stream = connection.fetch_many(query);
                     let mut error = None;
@@ -192,7 +192,7 @@ async fn execute_set_variable_query<'a>(
     source_file: &Path,
 ) -> anyhow::Result<()> {
     let query = bind_parameters(statement, request, db_connection).await?;
-    let connection = take_connection(&request.app_state.db, db_connection).await?;
+    let connection = take_connection(&request.app_state.db, db_connection, request).await?;
     log::debug!(
         "Executing query to set the {variable:?} variable: {:?}",
         query.sql
@@ -276,13 +276,15 @@ fn vars_and_name<'a, 'b>(
 async fn take_connection<'a>(
     db: &'a Database,
     conn: &'a mut DbConn,
+    request: &RequestInfo,
 ) -> anyhow::Result<&'a mut PoolConnection<sqlx::Any>> {
     if let Some(c) = conn {
         return Ok(c);
     }
     match db.connection.acquire().await {
         Ok(c) => {
             log::debug!("Acquired a database connection");
+            request.server_timing.borrow_mut().record("db_conn");
             *conn = Some(c);
             Ok(conn.as_mut().unwrap())
         }

src/webserver/http.rs

@@ -47,7 +47,7 @@ pub struct RequestContext {
     pub is_embedded: bool,
     pub source_path: PathBuf,
     pub content_security_policy: ContentSecurityPolicy,
-    pub server_timing: Option<ServerTiming>,
+    pub server_timing: ServerTiming,
 }
 
 async fn stream_response(stream: impl Stream<Item = DbItem>, mut renderer: AnyRenderBodyContext) {
@@ -106,16 +106,12 @@ async fn build_response_header_and_stream<S: Stream<Item = DbItem>>(
     let writer = ResponseWriter::new(sender);
     let mut head_context = HeaderContext::new(app_state, request_context, writer);
     let mut stream = Box::pin(database_entries);
-    let mut first_row = true;
     while let Some(item) = stream.next().await {
-        if first_row {
-            if let Some(ref mut timing) = head_context.request_context.server_timing {
-                timing.record("query");
-            }
-            first_row = false;
-        }
         let page_context = match item {
-            DbItem::Row(data) => head_context.handle_row(data).await?,
+            DbItem::Row(data) => {
+                head_context.request_context.server_timing.record("row");
+                head_context.handle_row(data).await?
+            }
             DbItem::FinishedQuery => {
                 log::debug!("finished query");
                 continue;
@@ -172,36 +168,31 @@ enum ResponseWithWriter<S> {
 async fn render_sql(
     srv_req: &mut ServiceRequest,
     sql_file: Arc<ParsedSqlFile>,
+    server_timing: ServerTiming,
 ) -> actix_web::Result<HttpResponse> {
     let app_state = srv_req
         .app_data::<web::Data<AppState>>()
         .ok_or_else(|| ErrorInternalServerError("no state"))?
         .clone()
         .into_inner();
 
-    let mut server_timing = if !app_state.config.environment.is_prod() {
-        Some(ServerTiming::new())
-    } else {
-        None
-    };
-
-    let mut req_param = extract_request_info(srv_req, Arc::clone(&app_state))
-        .await
-        .map_err(|e| anyhow_err_to_actix(e, &app_state))?;
+    let mut req_param =
+        extract_request_info(srv_req, Arc::clone(&app_state), server_timing.clone())
+            .await
+            .map_err(|e| anyhow_err_to_actix(e, &app_state))?;
     log::debug!("Received a request with the following parameters: {req_param:?}");
 
-    if let Some(ref mut timing) = server_timing {
-        timing.record("parse");
-    }
+    req_param.server_timing.borrow_mut().record("parse");
 
     let (resp_send, resp_recv) = tokio::sync::oneshot::channel::<HttpResponse>();
     let source_path: PathBuf = sql_file.source_path.clone();
     actix_web::rt::spawn(async move {
+        let server_timing_for_context = req_param.server_timing.borrow().clone();
         let request_context = RequestContext {
             is_embedded: req_param.get_variables.contains_key("_sqlpage_embed"),
             source_path,
             content_security_policy: ContentSecurityPolicy::with_random_nonce(),
-            server_timing,
+            server_timing: server_timing_for_context,
         };
         let mut conn = None;
         let database_entries_stream =
@@ -295,13 +286,18 @@ async fn process_sql_request(
     sql_path: PathBuf,
 ) -> actix_web::Result<HttpResponse> {
     let app_state: &web::Data<AppState> = req.app_data().expect("app_state");
+    let mut server_timing = ServerTiming::new(!app_state.config.environment.is_prod());
+    server_timing.record("request");
+
     let sql_file = app_state
         .sql_file_cache
         .get_with_privilege(app_state, &sql_path, false)
         .await
         .with_context(|| format!("Unable to read SQL file \"{}\"", sql_path.display()))
         .map_err(|e| anyhow_err_to_actix(e, app_state))?;
-    render_sql(req, sql_file).await
+    server_timing.record("sql_file");
+
+    render_sql(req, sql_file, server_timing).await
 }
 
 async fn serve_file(

src/webserver/http_request_info.rs

@@ -42,6 +42,7 @@ pub struct RequestInfo {
     pub clone_depth: u8,
     pub raw_body: Option<Vec<u8>>,
     pub oidc_claims: Option<OidcClaims>,
+    pub server_timing: std::cell::RefCell<super::server_timing::ServerTiming>,
 }
 
 impl RequestInfo {
@@ -62,6 +63,7 @@ impl RequestInfo {
             clone_depth: self.clone_depth + 1,
             raw_body: self.raw_body.clone(),
             oidc_claims: self.oidc_claims.clone(),
+            server_timing: std::cell::RefCell::new(self.server_timing.borrow().clone()),
         }
     }
 }
@@ -78,6 +80,7 @@ impl Clone for RequestInfo {
 pub(crate) async fn extract_request_info(
     req: &mut ServiceRequest,
     app_state: Arc<AppState>,
+    server_timing: super::server_timing::ServerTiming,
 ) -> anyhow::Result<RequestInfo> {
     let (http_req, payload) = req.parts_mut();
     let method = http_req.method().clone();
@@ -123,6 +126,7 @@ pub(crate) async fn extract_request_info(
         clone_depth: 0,
         raw_body,
         oidc_claims,
+        server_timing: std::cell::RefCell::new(server_timing),
     })
 }
 

src/webserver/server_timing.rs

@@ -2,6 +2,7 @@ use std::time::Instant;
 
 #[derive(Debug, Clone)]
 pub struct ServerTiming {
+    enabled: bool,
     start: Instant,
     events: Vec<TimingEvent>,
 }
@@ -14,20 +15,32 @@ struct TimingEvent {
 
 impl ServerTiming {
     #[must_use]
-    pub fn new() -> Self {
+    pub fn new(enabled: bool) -> Self {
         Self {
+            enabled,
             start: Instant::now(),
             events: Vec::new(),
         }
     }
 
     pub fn record(&mut self, name: &'static str) {
+        if !self.enabled {
+            return;
+        }
         let duration_ms = self.start.elapsed().as_secs_f64() * 1000.0;
         self.events.push(TimingEvent { name, duration_ms });
     }
 
+    #[must_use]
+    pub fn is_enabled(&self) -> bool {
+        self.enabled
+    }
+
     #[must_use]
     pub fn as_header_value(&self) -> String {
+        if !self.enabled {
+            return String::new();
+        }
         self.events
             .iter()
             .map(|event| format!("{};dur={:.2}", event.name, event.duration_ms))
@@ -38,6 +51,77 @@ impl ServerTiming {
 
 impl Default for ServerTiming {
     fn default() -> Self {
-        Self::new()
+        Self::new(false)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::thread;
+    use std::time::Duration;
+
+    #[test]
+    fn test_disabled_timing() {
+        let mut timing = ServerTiming::new(false);
+        assert!(!timing.is_enabled());
+        timing.record("event1");
+        timing.record("event2");
+        assert_eq!(timing.as_header_value(), "");
+    }
+
+    #[test]
+    fn test_enabled_timing() {
+        let mut timing = ServerTiming::new(true);
+        assert!(timing.is_enabled());
+        timing.record("event1");
+        thread::sleep(Duration::from_millis(10));
+        timing.record("event2");
+        let header = timing.as_header_value();
+        assert!(header.contains("event1;dur="));
+        assert!(header.contains("event2;dur="));
+        assert!(header.contains(", "));
+    }
+
+    #[test]
+    fn test_timing_values_increase() {
+        let mut timing = ServerTiming::new(true);
+        timing.record("first");
+        thread::sleep(Duration::from_millis(5));
+        timing.record("second");
+        assert_eq!(timing.events.len(), 2);
+        assert!(timing.events[1].duration_ms > timing.events[0].duration_ms);
+    }
+
+    #[test]
+    fn test_default_is_disabled() {
+        let timing = ServerTiming::default();
+        assert!(!timing.is_enabled());
+    }
+
+    #[test]
+    fn test_header_format() {
+        let mut timing = ServerTiming::new(true);
+        timing.events.push(TimingEvent {
+            name: "test",
+            duration_ms: 123.456,
+        });
+        let header = timing.as_header_value();
+        assert_eq!(header, "test;dur=123.46");
+    }
+
+    #[test]
+    fn test_multiple_events_format() {
+        let mut timing = ServerTiming::new(true);
+        timing.events.push(TimingEvent {
+            name: "first",
+            duration_ms: 10.5,
+        });
+        timing.events.push(TimingEvent {
+            name: "second",
+            duration_ms: 25.75,
+        });
+        let header = timing.as_header_value();
+        assert_eq!(header, "first;dur=10.50, second;dur=25.75");
     }
 }