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>

8 files changed, 207 insertions, 11 deletions

diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index 08b06663eecd..4b731460eb4e 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -2262,6 +2262,8 @@ dependencies = [
  "tokio-util",
  "tonic",
  "tonic-build",
+ "tower",
+ "tower-http 0.5.2",
  "tracing",
  "tracing-subscriber",
  "tvix-castore",
@@ -4351,6 +4353,23 @@ dependencies = [
 ]
 
 [[package]]
+name = "tower-http"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e9cd434a998747dd2c4276bc96ee2e0c7a2eadf3cae88e52be55a05fa9053f5"
+dependencies = [
+ "bitflags 2.4.2",
+ "bytes",
+ "http 1.1.0",
+ "http-body 1.0.0",
+ "http-body-util",
+ "pin-project-lite",
+ "tower-layer",
+ "tower-service",
+ "tracing",
+]
+
+[[package]]
 name = "tower-layer"
 version = "0.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -4755,7 +4774,7 @@ dependencies = [
  "tonic-build",
  "tonic-reflection",
  "tower",
- "tower-http",
+ "tower-http 0.4.4",
  "tracing",
  "tracing-indicatif",
  "tvix-castore",
@@ -4768,6 +4787,7 @@ dependencies = [
 name = "tvix-tracing"
 version = "0.1.0"
 dependencies = [
+ "axum 0.7.5",
  "http 0.2.11",
  "indicatif",
  "lazy_static",
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index 104c3dc825e1..58c60b0cd1b0 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -6942,6 +6942,15 @@ rec {
             features = [ "tls" "tls-roots" ];
           }
           {
+            name = "tower";
+            packageId = "tower";
+          }
+          {
+            name = "tower-http";
+            packageId = "tower-http 0.5.2";
+            features = [ "trace" ];
+          }
+          {
             name = "tracing";
             packageId = "tracing";
           }
@@ -6960,7 +6969,7 @@ rec {
           {
             name = "tvix-tracing";
             packageId = "tvix-tracing";
-            features = [ "tonic" ];
+            features = [ "tonic" "axum" ];
           }
           {
             name = "url";
@@ -13669,7 +13678,7 @@ rec {
         };
         resolvedDefaultFeatures = [ "__common" "balance" "buffer" "default" "discover" "futures-core" "futures-util" "indexmap" "limit" "load" "log" "make" "pin-project" "pin-project-lite" "rand" "ready-cache" "slab" "timeout" "tokio" "tokio-util" "tracing" "util" ];
       };
-      "tower-http" = rec {
+      "tower-http 0.4.4" = rec {
         crateName = "tower-http";
         version = "0.4.4";
         edition = "2018";
@@ -13769,6 +13778,99 @@ rec {
         };
         resolvedDefaultFeatures = [ "default" "trace" "tracing" ];
       };
+      "tower-http 0.5.2" = rec {
+        crateName = "tower-http";
+        version = "0.5.2";
+        edition = "2018";
+        sha256 = "1xakj3x0anp55gjqibiwvzma5iz0w9pcjsr7qk97sx4qm4sd970y";
+        authors = [
+          "Tower Maintainers <team@tower-rs.com>"
+        ];
+        dependencies = [
+          {
+            name = "bitflags";
+            packageId = "bitflags 2.4.2";
+          }
+          {
+            name = "bytes";
+            packageId = "bytes";
+          }
+          {
+            name = "http";
+            packageId = "http 1.1.0";
+          }
+          {
+            name = "http-body";
+            packageId = "http-body 1.0.0";
+          }
+          {
+            name = "http-body-util";
+            packageId = "http-body-util";
+          }
+          {
+            name = "pin-project-lite";
+            packageId = "pin-project-lite";
+          }
+          {
+            name = "tower-layer";
+            packageId = "tower-layer";
+          }
+          {
+            name = "tower-service";
+            packageId = "tower-service";
+          }
+          {
+            name = "tracing";
+            packageId = "tracing";
+            optional = true;
+            usesDefaultFeatures = false;
+          }
+        ];
+        devDependencies = [
+          {
+            name = "bytes";
+            packageId = "bytes";
+          }
+        ];
+        features = {
+          "async-compression" = [ "dep:async-compression" ];
+          "auth" = [ "base64" "validate-request" ];
+          "base64" = [ "dep:base64" ];
+          "catch-panic" = [ "tracing" "futures-util/std" ];
+          "compression-br" = [ "async-compression/brotli" "futures-core" "tokio-util" "tokio" ];
+          "compression-deflate" = [ "async-compression/zlib" "futures-core" "tokio-util" "tokio" ];
+          "compression-full" = [ "compression-br" "compression-deflate" "compression-gzip" "compression-zstd" ];
+          "compression-gzip" = [ "async-compression/gzip" "futures-core" "tokio-util" "tokio" ];
+          "compression-zstd" = [ "async-compression/zstd" "futures-core" "tokio-util" "tokio" ];
+          "decompression-br" = [ "async-compression/brotli" "futures-core" "tokio-util" "tokio" ];
+          "decompression-deflate" = [ "async-compression/zlib" "futures-core" "tokio-util" "tokio" ];
+          "decompression-full" = [ "decompression-br" "decompression-deflate" "decompression-gzip" "decompression-zstd" ];
+          "decompression-gzip" = [ "async-compression/gzip" "futures-core" "tokio-util" "tokio" ];
+          "decompression-zstd" = [ "async-compression/zstd" "futures-core" "tokio-util" "tokio" ];
+          "follow-redirect" = [ "futures-util" "iri-string" "tower/util" ];
+          "fs" = [ "futures-util" "tokio/fs" "tokio-util/io" "tokio/io-util" "dep:http-range-header" "mime_guess" "mime" "percent-encoding" "httpdate" "set-status" "futures-util/alloc" "tracing" ];
+          "full" = [ "add-extension" "auth" "catch-panic" "compression-full" "cors" "decompression-full" "follow-redirect" "fs" "limit" "map-request-body" "map-response-body" "metrics" "normalize-path" "propagate-header" "redirect" "request-id" "sensitive-headers" "set-header" "set-status" "timeout" "trace" "util" "validate-request" ];
+          "futures-core" = [ "dep:futures-core" ];
+          "futures-util" = [ "dep:futures-util" ];
+          "httpdate" = [ "dep:httpdate" ];
+          "iri-string" = [ "dep:iri-string" ];
+          "metrics" = [ "tokio/time" ];
+          "mime" = [ "dep:mime" ];
+          "mime_guess" = [ "dep:mime_guess" ];
+          "percent-encoding" = [ "dep:percent-encoding" ];
+          "request-id" = [ "uuid" ];
+          "timeout" = [ "tokio/time" ];
+          "tokio" = [ "dep:tokio" ];
+          "tokio-util" = [ "dep:tokio-util" ];
+          "tower" = [ "dep:tower" ];
+          "trace" = [ "tracing" ];
+          "tracing" = [ "dep:tracing" ];
+          "util" = [ "tower" ];
+          "uuid" = [ "dep:uuid" ];
+          "validate-request" = [ "mime" ];
+        };
+        resolvedDefaultFeatures = [ "default" "trace" "tracing" ];
+      };
       "tower-layer" = rec {
         crateName = "tower-layer";
         version = "0.3.2";
@@ -15335,7 +15437,7 @@ rec {
           }
           {
             name = "tower-http";
-            packageId = "tower-http";
+            packageId = "tower-http 0.4.4";
             features = [ "trace" ];
           }
           {
@@ -15414,6 +15516,11 @@ rec {
         src = lib.cleanSourceWith { filter = sourceFilter; src = ./tracing; };
         dependencies = [
           {
+            name = "axum";
+            packageId = "axum 0.7.5";
+            optional = true;
+          }
+          {
             name = "http";
             packageId = "http 0.2.11";
             optional = true;
@@ -15494,12 +15601,13 @@ rec {
           }
         ];
         features = {
+          "axum" = [ "dep:axum" ];
           "otlp" = [ "dep:tracing-opentelemetry" "dep:opentelemetry" "dep:opentelemetry-otlp" "dep:opentelemetry_sdk" "dep:opentelemetry-http" "reqwest-tracing?/opentelemetry_0_22" ];
           "reqwest" = [ "dep:reqwest-tracing" ];
           "tonic" = [ "dep:tonic" "dep:http" ];
           "tracy" = [ "dep:tracing-tracy" ];
         };
-        resolvedDefaultFeatures = [ "default" "otlp" "reqwest" "tonic" "tracy" ];
+        resolvedDefaultFeatures = [ "axum" "default" "otlp" "reqwest" "tonic" "tracy" ];
       };
       "typeid" = rec {
         crateName = "typeid";
diff --git a/tvix/nar-bridge/Cargo.toml b/tvix/nar-bridge/Cargo.toml
index 7dc3a82848b6..dd38d1b30e9f 100644
--- a/tvix/nar-bridge/Cargo.toml
+++ b/tvix/nar-bridge/Cargo.toml
@@ -5,6 +5,8 @@ edition = "2021"
 
 [dependencies]
 axum = { version = "0.7.5", features = ["http2"] }
+tower = "0.4.13"
+tower-http = { version = "0.5", features = ["trace"] }
 bytes = "1.4.0"
 clap = { version = "4.0", features = ["derive", "env"] }
 data-encoding = "2.3.3"
@@ -19,7 +21,7 @@ tokio-util = { version = "0.7.9", features = ["io", "io-util", "compat"] }
 tonic = { version = "0.11.0", features = ["tls", "tls-roots"] }
 tvix-castore = { path = "../castore" }
 tvix-store = { path = "../store" }
-tvix-tracing = { path = "../tracing", features = ["tonic"] }
+tvix-tracing = { path = "../tracing", features = ["tonic", "axum"] }
 tracing = "0.1.37"
 tracing-subscriber = "0.3.16"
 url = "2.4.0"
diff --git a/tvix/nar-bridge/src/bin/nar-bridge.rs b/tvix/nar-bridge/src/bin/nar-bridge.rs
index 1e42237754e5..5aa0113dcd0c 100644
--- a/tvix/nar-bridge/src/bin/nar-bridge.rs
+++ b/tvix/nar-bridge/src/bin/nar-bridge.rs
@@ -1,5 +1,7 @@
 use clap::Parser;
 use nar_bridge::AppState;
+use tower::ServiceBuilder;
+use tower_http::trace::{DefaultMakeSpan, TraceLayer};
 use tracing::info;
 
 /// Expose the Nix HTTP Binary Cache protocol for a tvix-store.
@@ -57,7 +59,19 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
 
     let state = AppState::new(blob_service, directory_service, path_info_service);
 
-    let app = nar_bridge::gen_router(cli.priority).with_state(state);
+    let app = nar_bridge::gen_router(cli.priority)
+        .layer(
+            ServiceBuilder::new()
+                .layer(
+                    TraceLayer::new_for_http().make_span_with(
+                        DefaultMakeSpan::new()
+                            .level(tracing::Level::INFO)
+                            .include_headers(true),
+                    ),
+                )
+                .map_request(tvix_tracing::propagate::axum::accept_trace),
+        )
+        .with_state(state);
 
     let listen_address = &cli.listen_args.listen_address.unwrap_or_else(|| {
         "[::]:8000"
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;