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:
		
							parent
							
								
									fdc0cf0c94
								
							
						
					
					
						commit
						1515a970be
					
				
					 8 changed files with 207 additions and 11 deletions
				
			
		
							
								
								
									
										22
									
								
								tvix/Cargo.lock
									
										
									
										generated
									
									
									
								
							
							
						
						
									
										22
									
								
								tvix/Cargo.lock
									
										
									
										generated
									
									
									
								
							|  | @ -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", | ||||
|  |  | |||
							
								
								
									
										116
									
								
								tvix/Cargo.nix
									
										
									
									
									
								
							
							
						
						
									
										116
									
								
								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 <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" ]; | ||||
|           } | ||||
|           { | ||||
|  | @ -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"; | ||||
|  |  | |||
|  | @ -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" | ||||
|  |  | |||
|  | @ -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" | ||||
|  |  | |||
|  | @ -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 | ||||
|  |  | |||
|  | @ -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" ]; | ||||
|   }; | ||||
| }) | ||||
|  |  | |||
							
								
								
									
										48
									
								
								tvix/tracing/src/propagate/axum.rs
									
										
									
									
									
										Normal file
									
								
							
							
						
						
									
										48
									
								
								tvix/tracing/src/propagate/axum.rs
									
										
									
									
									
										Normal file
									
								
							|  | @ -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() | ||||
|     } | ||||
| } | ||||
|  | @ -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; | ||||
|  |  | |||
		Loading…
	
	Add table
		Add a link
		
	
		Reference in a new issue