diff options
author | Simon Hauser <simon.hauser@helsinki-systems.de> | 2024-06-20T09·39+0200 |
---|---|---|
committer | Simon Hauser <simon.hauser@helsinki-systems.de> | 2024-06-20T19·21+0000 |
commit | 639a00e2ab1f2a6530477d2d7f91cccbc9f70746 (patch) | |
tree | 4f953d503a6ee0c3d04c623c0eed13f7381bd699 /tvix/tracing/src/propagate/tonic.rs | |
parent | 2b20d8d82dd424f2cb457c0cdef3ab3e98512117 (diff) |
feat(tvix/tracing): gRPC trace context propagation r/8299
This introduces optional helper function in tvix/tracing for trace propagation and uses these helper in the `tvix-store`. The GRPCBlobService, GRPCDirectoryService and GRPCPathInfoService now accept a generic client, meaning the client can be generated with either `::new` or `::with_interceptor`. This was tested and validated by starting a `tvix-store daemon` and `tvix-store import`. Change-Id: I4b194483bf09266820104b4b56e4a135dca2b77a Reviewed-on: https://cl.tvl.fyi/c/depot/+/11863 Reviewed-by: flokli <flokli@flokli.de> Tested-by: BuildkiteCI
Diffstat (limited to '')
-rw-r--r-- | tvix/tracing/src/propagate/tonic.rs | 59 |
1 files changed, 59 insertions, 0 deletions
diff --git a/tvix/tracing/src/propagate/tonic.rs b/tvix/tracing/src/propagate/tonic.rs new file mode 100644 index 000000000000..af643eb70e36 --- /dev/null +++ b/tvix/tracing/src/propagate/tonic.rs @@ -0,0 +1,59 @@ +use tonic::{ + metadata::{MetadataKey, MetadataMap, MetadataValue}, + Status, +}; +use tracing::{warn, Span}; + +#[cfg(feature = "otlp")] +use opentelemetry::{global, propagation::Injector}; +#[cfg(feature = "otlp")] +use opentelemetry_http::HeaderExtractor; +#[cfg(feature = "otlp")] +use tracing_opentelemetry::OpenTelemetrySpanExt; + +/// 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: http::Request<B>) -> 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())) + }); + Span::current().set_parent(parent_context); + } + request +} + +#[cfg(feature = "otlp")] +struct MetadataInjector<'a>(&'a mut MetadataMap); + +#[cfg(feature = "otlp")] +impl Injector for MetadataInjector<'_> { + fn set(&mut self, key: &str, value: String) { + match MetadataKey::from_bytes(key.as_bytes()) { + Ok(key) => match MetadataValue::try_from(&value) { + Ok(value) => { + self.0.insert(key, value); + } + Err(error) => warn!(value, error = format!("{error:#}"), "parse metadata value"), + }, + Err(error) => warn!(key, error = format!("{error:#}"), "parse metadata key"), + } + } +} + +/// Trace context propagation: send the trace context by injecting it into the metadata of the given +/// request. This only injects the current span if the otlp feature is also enabled. +pub fn send_trace<T>(mut request: tonic::Request<T>) -> Result<tonic::Request<T>, Status> { + #[cfg(feature = "otlp")] + { + global::get_text_map_propagator(|propagator| { + let context = Span::current().context(); + propagator.inject_context(&context, &mut MetadataInjector(request.metadata_mut())) + }); + } + Ok(request) +} |