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

3 files changed, 75 insertions, 0 deletions

diff --git a/tvix/tracing/src/lib.rs b/tvix/tracing/src/lib.rs
index 68a417998a24..b965ca4a3dbc 100644
--- a/tvix/tracing/src/lib.rs
+++ b/tvix/tracing/src/lib.rs
@@ -9,6 +9,7 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilte
 use opentelemetry::{trace::Tracer, KeyValue};
 #[cfg(feature = "otlp")]
 use opentelemetry_sdk::{
+    propagation::TraceContextPropagator,
     resource::{ResourceDetector, SdkProvidedResourceDetector},
     trace::BatchConfigBuilder,
     Resource,
@@ -16,6 +17,9 @@ use opentelemetry_sdk::{
 #[cfg(feature = "tracy")]
 use tracing_tracy::TracyLayer;
 
+#[cfg(feature = "tonic")] // TODO or http
+pub mod propagate;
+
 lazy_static! {
     pub static ref PB_PROGRESS_STYLE: ProgressStyle = ProgressStyle::with_template(
         "{span_child_prefix} {wide_msg} {bar:10} ({elapsed}) {pos:>7}/{len:7}"
@@ -186,6 +190,9 @@ impl TracingBuilder {
         #[cfg(feature = "otlp")]
         {
             if let Some(service_name) = self.service_name {
+                // register a text map propagator for trace propagation
+                opentelemetry::global::set_text_map_propagator(TraceContextPropagator::new());
+
                 let (tracer, tx) = gen_otlp_tracer(service_name.to_string());
                 // Create a tracing layer with the configured tracer
                 let layer = tracing_opentelemetry::layer().with_tracer(tracer);
diff --git a/tvix/tracing/src/propagate/mod.rs b/tvix/tracing/src/propagate/mod.rs
new file mode 100644
index 000000000000..42c532e9d8cf
--- /dev/null
+++ b/tvix/tracing/src/propagate/mod.rs
@@ -0,0 +1,9 @@
+#[cfg(feature = "tonic")]
+pub mod tonic;
+
+// TODO: Helper library for reqwest. We could use
+// https://github.com/TrueLayer/reqwest-middleware/tree/main/reqwest-tracing to realise this
+
+// 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
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)
+}