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>

1 files changed, 15 insertions, 1 deletions

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"