feat(tvix/tracing): http propagation for axum

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>
This commit is contained in:
Simon Hauser 2024-07-02 12:50:43 +02:00
parent fdc0cf0c94
commit 1515a970be
8 changed files with 207 additions and 11 deletions

View file

@ -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"

View file

@ -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"