feat(tvix/tracing): http propagation for axum r/8384

It introduces a new accept_trace function for axum0.7 which can be used
to accept a header trace from a received request. This function can be
used for tonic 0.12 once that version is released, and the specific
`accept_trace` function within `tvix_tracing::propagate::tonic` can then
be removed.

This also integrates http propagation into the nar_bridge crate.

Change-Id: I46dcc797d494bb3977c2633753e7060d88d29129
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11925
Reviewed-by: Brian Olsen <me@griff.name>
Tested-by: BuildkiteCI
Reviewed-by: Simon Hauser <simon.hauser@helsinki-systems.de>
Reviewed-by: flokli <flokli@flokli.de>

4 files changed, 56 insertions, 4 deletions

diff --git a/tvix/tracing/Cargo.toml b/tvix/tracing/Cargo.toml
index 80892bf7a09f..db1626d26d52 100644
--- a/tvix/tracing/Cargo.toml
+++ b/tvix/tracing/Cargo.toml
@@ -24,6 +24,8 @@ http  = { version = "0.2.11", optional = true }
 
 reqwest-tracing = { version = "0.4.8", default-features = false, optional = true }
 
+axum = { version = "0.7.5", optional = true }
+
 [features]
 default = []
 otlp = [
@@ -44,6 +46,9 @@ tonic = [
 reqwest = [
   "dep:reqwest-tracing",
 ]
+axum = [
+  "dep:axum",
+]
 
 [lints]
 workspace = true
diff --git a/tvix/tracing/default.nix b/tvix/tracing/default.nix
index a4ac9de1c8bc..ef1985cb47fb 100644
--- a/tvix/tracing/default.nix
+++ b/tvix/tracing/default.nix
@@ -6,6 +6,6 @@
   meta.ci.targets = lib.filter (x: lib.hasPrefix "with-features" x || x == "no-features") (lib.attrNames passthru);
   passthru = depot.tvix.utils.mkFeaturePowerset {
     inherit (old) crateName;
-    features = [ "otlp" "tracy" "tonic" "reqwest" ];
+    features = [ "otlp" "tracy" "tonic" "reqwest" "axum" ];
   };
 })
diff --git a/tvix/tracing/src/propagate/axum.rs b/tvix/tracing/src/propagate/axum.rs
new file mode 100644
index 000000000000..6d012f449762
--- /dev/null
+++ b/tvix/tracing/src/propagate/axum.rs
@@ -0,0 +1,48 @@
+#[cfg(feature = "otlp")]
+use opentelemetry::{global, propagation::Extractor};
+#[cfg(feature = "otlp")]
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+
+// TODO: accept_trace can be shared with tonic, as soon as tonic upstream has a release with
+// support for axum07. Latest master already has support for axum07 but there is not release yet:
+// https://github.com/hyperium/tonic/pull/1740
+
+/// Trace context propagation: associate the current span with the otlp trace of the given request,
+/// if any and valid. This only sets the parent trace if the otlp feature is also enabled.
+pub fn accept_trace<B>(request: axum::http::Request<B>) -> axum::http::Request<B> {
+    // we only extract and set a parent trace if otlp feature is enabled, otherwise this feature is
+    // an noop and we return the request as is
+    #[cfg(feature = "otlp")]
+    {
+        // Current context, if no or invalid data is received.
+        let parent_context = global::get_text_map_propagator(|propagator| {
+            propagator.extract(&HeaderExtractor(request.headers()))
+        });
+        tracing::Span::current().set_parent(parent_context);
+    }
+    request
+}
+
+/// Helper for extracting headers from HTTP Requests. This is used for OpenTelemetry context
+/// propagation over HTTP.
+#[cfg(feature = "otlp")]
+struct HeaderExtractor<'a>(&'a axum::http::HeaderMap);
+
+#[cfg(feature = "otlp")]
+impl<'a> Extractor for HeaderExtractor<'a> {
+    /// Get a value for a key from the HeaderMap.  If the value is not valid ASCII, returns None.
+    fn get(&self, key: &str) -> Option<&str> {
+        self.0.get(key).and_then(|v| {
+            let s = v.to_str();
+            if let Err(ref error) = s {
+                tracing::warn!(%error, ?v, "cannot convert header value to ASCII")
+            };
+            s.ok()
+        })
+    }
+
+    /// Collect all the keys from the HeaderMap.
+    fn keys(&self) -> Vec<&str> {
+        self.0.keys().map(|k| k.as_str()).collect()
+    }
+}
diff --git a/tvix/tracing/src/propagate/mod.rs b/tvix/tracing/src/propagate/mod.rs
index 9a7e4332b637..2e56a832e5b7 100644
--- a/tvix/tracing/src/propagate/mod.rs
+++ b/tvix/tracing/src/propagate/mod.rs
@@ -4,6 +4,5 @@ pub mod tonic;
 #[cfg(feature = "reqwest")]
 pub mod reqwest;
 
-// TODO: Helper library for axum or another http server, see
-// https://github.com/hseeberger/hello-tracing-rs/blob/main/hello-tracing-common/src/otel/http.rs
-// as an example and we can reuse tonic::accept_trace fun, at least for a tower::ServiceBuilder
+#[cfg(feature = "axum")]
+pub mod axum;