From 6ab87f97dd30b1146d18443b8ea66844567399e5 Mon Sep 17 00:00:00 2001
From: Charles Hall <charles@computer.surgery>
Date: Sun, 22 Sep 2024 11:01:33 -0700
Subject: [PATCH] include traceresponse header if possible

This can help with debugging.
---
 book/changelog.md    |  3 +++
 src/cli/serve.rs     |  3 ++-
 src/observability.rs | 52 +++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/book/changelog.md b/book/changelog.md
index 26ed1d67..377e9ce8 100644
--- a/book/changelog.md
+++ b/book/changelog.md
@@ -248,3 +248,6 @@ This will be the first release of Grapevine since it was forked from Conduit
     [!102](https://gitlab.computer.surgery/matrix/grapevine/-/merge_requests/102))
 19. Allow configuring the served API components per listener.
     ([!109](https://gitlab.computer.surgery/matrix/grapevine/-/merge_requests/109))
+20. Include the [`traceresponse` header](https://w3c.github.io/trace-context/#traceresponse-header)
+    if OpenTelemetry Tracing is in use.
+    ([!112](https://gitlab.computer.surgery/matrix/grapevine/-/merge_requests/112))
diff --git a/src/cli/serve.rs b/src/cli/serve.rs
index 79a55482..0e646dc0 100644
--- a/src/cli/serve.rs
+++ b/src/cli/serve.rs
@@ -171,7 +171,8 @@ async fn run_server() -> Result<(), error::Serve> {
                 .try_into()
                 .expect("failed to convert max request size"),
         ))
-        .layer(axum::middleware::from_fn(observability::http_metrics_layer));
+        .layer(axum::middleware::from_fn(observability::http_metrics_layer))
+        .layer(axum::middleware::from_fn(observability::traceresponse_layer));
 
     let mut handles = Vec::new();
     let mut servers = JoinSet::new();
diff --git a/src/observability.rs b/src/observability.rs
index fb5df965..059ed436 100644
--- a/src/observability.rs
+++ b/src/observability.rs
@@ -18,9 +18,13 @@ use opentelemetry_sdk::{
 };
 use strum::{AsRefStr, IntoStaticStr};
 use tokio::time::Instant;
+use tracing::Span;
 use tracing_flame::{FlameLayer, FlushGuard};
+use tracing_opentelemetry::OtelData;
 use tracing_subscriber::{
-    layer::SubscriberExt, reload, EnvFilter, Layer, Registry,
+    layer::SubscriberExt,
+    registry::{LookupSpan, SpanData},
+    reload, EnvFilter, Layer, Registry,
 };
 
 use crate::{
@@ -138,8 +142,7 @@ fn make_backend<S, L, T>(
 >
 where
     L: Layer<S>,
-    S: tracing::Subscriber
-        + for<'span> tracing_subscriber::registry::LookupSpan<'span>,
+    S: tracing::Subscriber + for<'span> LookupSpan<'span>,
 {
     if !enable {
         return Ok((None, None, None));
@@ -423,3 +426,46 @@ pub(crate) async fn http_metrics_layer(req: Request, next: Next) -> Response {
         }
     }
 }
+
+/// Add `traceresponse` header if possible
+///
+/// See also <https://w3c.github.io/trace-context/#traceresponse-header>.
+pub(crate) async fn traceresponse_layer(req: Request, next: Next) -> Response {
+    let mut resp = next.run(req).await;
+
+    let ids = tracing::dispatcher::get_default(|dispatch| {
+        Span::current()
+            .id()
+            .and_then(|id| {
+                dispatch
+                    .downcast_ref::<Registry>()
+                    .and_then(|x| x.span_data(&id))
+            })
+            .and_then(|x| {
+                x.extensions()
+                    .get::<OtelData>()
+                    .and_then(|x| x.builder.trace_id.zip(x.builder.span_id))
+            })
+    });
+
+    if let Some((trace_id, span_id)) = ids {
+        let headers = resp.headers_mut();
+
+        headers.insert(
+            "traceresponse",
+            format!(
+                "{:02x}-{}-{}-{:02x}",
+                0,
+                trace_id,
+                span_id,
+                // Doesn't seem to be possible to get the SpanContext here, but
+                // this should be a fine default value
+                0,
+            )
+            .try_into()
+            .expect("traceresponse value should be a valid header value"),
+        );
+    }
+
+    resp
+}