From 1515a970bedbb6d7b5e8f966dddd0d8fff9bb03c Mon Sep 17 00:00:00 2001 From: Simon Hauser Date: Tue, 2 Jul 2024 12:50:43 +0200 Subject: [PATCH] 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 Tested-by: BuildkiteCI Reviewed-by: Simon Hauser Reviewed-by: flokli --- tvix/Cargo.lock | 22 ++++- tvix/Cargo.nix | 116 +++++++++++++++++++++++++- tvix/nar-bridge/Cargo.toml | 4 +- tvix/nar-bridge/src/bin/nar-bridge.rs | 16 +++- tvix/tracing/Cargo.toml | 5 ++ tvix/tracing/default.nix | 2 +- tvix/tracing/src/propagate/axum.rs | 48 +++++++++++ tvix/tracing/src/propagate/mod.rs | 5 +- 8 files changed, 207 insertions(+), 11 deletions(-) create mode 100644 tvix/tracing/src/propagate/axum.rs diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock index 08b06663e..4b731460e 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", @@ -4350,6 +4352,23 @@ dependencies = [ "tracing", ] +[[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" @@ -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 104c3dc82..58c60b0cd 100644 --- a/tvix/Cargo.nix +++ b/tvix/Cargo.nix @@ -6941,6 +6941,15 @@ rec { packageId = "tonic"; 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 " + ]; + 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" ]; } { @@ -15413,6 +15515,11 @@ rec { edition = "2021"; src = lib.cleanSourceWith { filter = sourceFilter; src = ./tracing; }; dependencies = [ + { + name = "axum"; + packageId = "axum 0.7.5"; + optional = true; + } { name = "http"; packageId = "http 0.2.11"; @@ -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 7dc3a8284..dd38d1b30 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 1e4223775..5aa0113dc 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> { 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 80892bf7a..db1626d26 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 a4ac9de1c..ef1985cb4 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 000000000..6d012f449 --- /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(request: axum::http::Request) -> axum::http::Request { + // 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 9a7e4332b..2e56a832e 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;