feat: Add Server-Timing header support

Co-authored-by: contact <[email protected]>

Author: cursoragent

src/render.rs

@@ -80,7 +80,7 @@ pub enum PageContext {
 /// Handles the first SQL statements, before the headers have been sent to
 pub struct HeaderContext {
     app_state: Arc<AppState>,
-    request_context: RequestContext,
+    pub request_context: RequestContext,
     pub writer: ResponseWriter,
     response: HttpResponseBuilder,
     has_status: bool,
@@ -368,7 +368,13 @@ impl HeaderContext {
         Ok(PageContext::Header(self))
     }
 
-    async fn start_body(self, data: JsonValue) -> anyhow::Result<PageContext> {
+    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));
+            }
+        }
         let html_renderer =
             HtmlRenderContext::new(self.app_state, self.request_context, self.writer, data)
                 .await
@@ -382,6 +388,12 @@ 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.response.finish()
     }
 }

src/webserver/http.rs

@@ -7,6 +7,7 @@ use crate::webserver::content_security_policy::ContentSecurityPolicy;
 use crate::webserver::database::execute_queries::stop_at_first_error;
 use crate::webserver::database::{execute_queries::stream_query_results_with_conn, DbItem};
 use crate::webserver::http_request_info::extract_request_info;
+use crate::webserver::server_timing::ServerTiming;
 use crate::webserver::ErrorWithStatus;
 use crate::{AppConfig, AppState, ParsedSqlFile, DEFAULT_404_FILE};
 use actix_web::dev::{fn_service, ServiceFactory, ServiceRequest};
@@ -46,6 +47,7 @@ pub struct RequestContext {
     pub is_embedded: bool,
     pub source_path: PathBuf,
     pub content_security_policy: ContentSecurityPolicy,
+    pub server_timing: Option<ServerTiming>,
 }
 
 async fn stream_response(stream: impl Stream<Item = DbItem>, mut renderer: AnyRenderBodyContext) {
@@ -104,7 +106,14 @@ 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::FinishedQuery => {
@@ -167,21 +176,32 @@ async fn render_sql(
     let app_state = srv_req
         .app_data::<web::Data<AppState>>()
         .ok_or_else(|| ErrorInternalServerError("no state"))?
-        .clone() // Cheap reference count increase
+        .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))?;
     log::debug!("Received a request with the following parameters: {req_param:?}");
 
+    if let Some(ref mut timing) = server_timing {
+        timing.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 request_context = RequestContext {
             is_embedded: req_param.get_variables.contains_key("_sqlpage_embed"),
             source_path,
             content_security_policy: ContentSecurityPolicy::with_random_nonce(),
+            server_timing,
         };
         let mut conn = None;
         let database_entries_stream =

src/webserver/mod.rs

@@ -38,6 +38,7 @@ pub mod http_client;
 pub mod http_request_info;
 mod https;
 pub mod request_variables;
+pub mod server_timing;
 
 pub use database::Database;
 pub use error_with_status::ErrorWithStatus;

src/webserver/server_timing.rs

@@ -0,0 +1,43 @@
+use std::time::Instant;
+
+#[derive(Debug, Clone)]
+pub struct ServerTiming {
+    start: Instant,
+    events: Vec<TimingEvent>,
+}
+
+#[derive(Debug, Clone)]
+struct TimingEvent {
+    name: &'static str,
+    duration_ms: f64,
+}
+
+impl ServerTiming {
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            start: Instant::now(),
+            events: Vec::new(),
+        }
+    }
+
+    pub fn record(&mut self, name: &'static str) {
+        let duration_ms = self.start.elapsed().as_secs_f64() * 1000.0;
+        self.events.push(TimingEvent { name, duration_ms });
+    }
+
+    #[must_use]
+    pub fn as_header_value(&self) -> String {
+        self.events
+            .iter()
+            .map(|event| format!("{};dur={:.2}", event.name, event.duration_ms))
+            .collect::<Vec<_>>()
+            .join(", ")
+    }
+}
+
+impl Default for ServerTiming {
+    fn default() -> Self {
+        Self::new()
+    }
+}