chore: add Open Telemetry attributes to grpc spans

Cargo.toml

@@ -26,6 +26,7 @@ version      = "0.8.0"
 
 [workspace.dependencies]
 assert_matches            = { version = "1.5" }
+http                      = { version = "1.2" }
 itertools                 = { version = "0.14" }
 miden-air                 = { version = "0.12" }
 miden-lib                 = { git = "https://github.com/0xPolygonMiden/miden-base", branch = "next" }
@@ -44,6 +45,8 @@ thiserror                 = { version = "2.0", default-features = false }
 tokio                     = { version = "1.40", features = ["rt-multi-thread"] }
 tokio-stream              = { version = "0.1" }
 tonic                     = { version = "0.12" }
+tower                     = { version = "0.5" }
+tower-http                = { version = "0.6", features = ["trace"] }
 tracing                   = { version = "0.1" }
 tracing-subscriber        = { version = "0.3", features = ["env-filter", "fmt", "json"] }
 url                       = { version = "2.5", features = ["serde"] }

bin/faucet/Cargo.toml

@@ -34,8 +34,8 @@ thiserror        = { workspace = true }
 tokio            = { workspace = true, features = ["fs"] }
 toml             = { version = "0.8" }
 tonic            = { workspace = true }
-tower            = "0.5"
-tower-http       = { version = "0.6", features = ["cors", "set-header", "trace"] }
+tower            = { workspace = true }
+tower-http       = { workspace = true, features = ["cors", "set-header", "trace"] }
 tracing          = { workspace = true }
 url              = { workspace = true }
 

crates/block-producer/Cargo.toml

@@ -24,7 +24,7 @@ itertools             = { workspace = true }
 miden-block-prover    = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "next" }
 miden-lib             = { workspace = true }
 miden-node-proto      = { workspace = true }
-miden-node-utils      = { workspace = true }
+miden-node-utils      = { workspace = true, features = ["testing"] }
 miden-objects         = { workspace = true }
 miden-processor       = { workspace = true }
 miden-tx              = { workspace = true }
@@ -34,7 +34,8 @@ serde                 = { version = "1.0", features = ["derive"] }
 thiserror             = { workspace = true }
 tokio                 = { workspace = true, features = ["macros", "net", "rt-multi-thread", "sync", "time"] }
 tokio-stream          = { workspace = true, features = ["net"] }
-tonic                 = { workspace = true }
+tonic                 = { workspace = true, features = ["transport"] }
+tower-http            = { workspace = true, features = ["util"] }
 tracing               = { workspace = true }
 url                   = { workspace = true }
 

crates/block-producer/src/server.rs

@@ -7,14 +7,15 @@ use miden_node_proto::generated::{
 use miden_node_utils::{
     errors::ApiError,
     formatting::{format_input_notes, format_output_notes},
-    tracing::grpc::OtelInterceptor,
+    tracing::grpc::{block_producer_trace_fn, OtelInterceptor},
 };
 use miden_objects::{
     block::BlockNumber, transaction::ProvenTransaction, utils::serde::Deserializable,
 };
 use tokio::{net::TcpListener, sync::Mutex};
 use tokio_stream::wrappers::TcpListenerStream;
 use tonic::Status;
+use tower_http::trace::TraceLayer;
 use tracing::{debug, info, instrument};
 
 use crate::{
@@ -211,8 +212,9 @@ impl BlockProducerRpcServer {
     }
 
     async fn serve(self, listener: TcpListener) -> Result<(), tonic::transport::Error> {
+        // Build the gRPC server with the API service and trace layer.
         tonic::transport::Server::builder()
-            .trace_fn(miden_node_utils::tracing::grpc::block_producer_trace_fn)
+            .layer(TraceLayer::new_for_grpc().make_span_with(block_producer_trace_fn))
             .add_service(api_server::ApiServer::new(self))
             .serve_with_incoming(TcpListenerStream::new(listener))
             .await

crates/store/Cargo.toml

@@ -28,6 +28,7 @@ thiserror          = { workspace = true }
 tokio              = { workspace = true, features = ["fs", "macros", "net", "rt-multi-thread"] }
 tokio-stream       = { workspace = true, features = ["net"] }
 tonic              = { workspace = true }
+tower-http         = { workspace = true, features = ["util"] }
 tracing            = { workspace = true }
 url                = { workspace = true }
 

crates/store/src/server/mod.rs

@@ -1,9 +1,10 @@
 use std::sync::Arc;
 
 use miden_node_proto::generated::store::api_server;
-use miden_node_utils::errors::ApiError;
+use miden_node_utils::{errors::ApiError, tracing::grpc::store_trace_fn};
 use tokio::net::TcpListener;
 use tokio_stream::wrappers::TcpListenerStream;
+use tower_http::trace::TraceLayer;
 use tracing::info;
 
 use crate::{blocks::BlockStore, config::StoreConfig, db::Db, state::State, COMPONENT};
@@ -61,8 +62,9 @@ impl Store {
     ///
     /// Note: this blocks until the server dies.
     pub async fn serve(self) -> Result<(), ApiError> {
+        // Build the gRPC server with the API service and trace layer.
         tonic::transport::Server::builder()
-            .trace_fn(miden_node_utils::tracing::grpc::store_trace_fn)
+            .layer(TraceLayer::new_for_grpc().make_span_with(store_trace_fn))
             .add_service(self.api_service)
             .serve_with_incoming(TcpListenerStream::new(self.listener))
             .await

crates/utils/Cargo.toml

@@ -17,16 +17,18 @@ workspace = true
 [features]
 # Enables depedencies intended for build script generation of version metadata.
 vergen = ["dep:vergen", "dep:vergen-gitcl"]
+# Enables utility functions for testing traces created by some other crate's stack.
+testing = ["dep:tokio"]
 
 [dependencies]
 anyhow                = { version = "1.0" }
 figment               = { version = "0.10", features = ["env", "toml"] }
-http                  = { version = "1.2" }
+http                  = { workspace = true }
 itertools             = { workspace = true }
 miden-objects         = { workspace = true }
 opentelemetry         = { version = "0.28" }
 opentelemetry-otlp    = { version = "0.28", default-features = false, features = ["grpc-tonic", "tls-roots", "trace"] }
-opentelemetry_sdk     = { version = "0.28", features = ["rt-tokio"] }
+opentelemetry_sdk     = { version = "0.28", features = ["rt-tokio", "testing"] }
 rand                  = { workspace = true }
 serde                 = { version = "1.0", features = ["derive"] }
 thiserror             = { workspace = true }
@@ -35,6 +37,9 @@ tracing               = { workspace = true }
 tracing-forest        = { version = "0.1", optional = true, features = ["chrono"] }
 tracing-opentelemetry = { version = "0.29" }
 tracing-subscriber    = { workspace = true }
+
+# Optional dependencies enabled by `testing` feature.
+tokio = { workspace = true, optional = true }
 # Optional dependencies enabled by `vergen` feature.
 # This must match the version expected by `vergen-gitcl`.
 vergen       = { "version" = "9.0", optional = true }

crates/utils/src/logging.rs

@@ -3,7 +3,7 @@ use std::str::FromStr;
 use anyhow::Result;
 use opentelemetry::trace::TracerProvider as _;
 use opentelemetry_otlp::WithTonicConfig;
-use opentelemetry_sdk::propagation::TraceContextPropagator;
+use opentelemetry_sdk::{propagation::TraceContextPropagator, trace::SpanExporter};
 use tracing::subscriber::Subscriber;
 use tracing_opentelemetry::OpenTelemetryLayer;
 use tracing_subscriber::{
@@ -39,25 +39,49 @@ pub fn setup_tracing(otel: OpenTelemetry) -> Result<()> {
     // Note: open-telemetry requires a tokio-runtime, so this _must_ be lazily evaluated (aka not
     // `then_some`) to avoid crashing sync callers (with OpenTelemetry::Disabled set). Examples of
     // such callers are tests with logging enabled.
-    let otel_layer = otel.is_enabled().then(open_telemetry_layer);
+    let otel_layer = {
+        if otel.is_enabled() {
+            let exporter = opentelemetry_otlp::SpanExporter::builder()
+                .with_tonic()
+                .with_tls_config(tonic::transport::ClientTlsConfig::new().with_native_roots())
+                .build()?;
+            Some(open_telemetry_layer(exporter))
+        } else {
+            None
+        }
+    };
 
     let subscriber = Registry::default()
         .with(stdout_layer().with_filter(env_or_default_filter()))
         .with(otel_layer.with_filter(env_or_default_filter()));
     tracing::subscriber::set_global_default(subscriber).map_err(Into::into)
 }
 
-fn open_telemetry_layer<S>() -> Box<dyn tracing_subscriber::Layer<S> + Send + Sync + 'static>
+#[cfg(feature = "testing")]
+pub fn setup_test_tracing() -> Result<(
+    tokio::sync::mpsc::UnboundedReceiver<opentelemetry_sdk::trace::SpanData>,
+    tokio::sync::mpsc::UnboundedReceiver<()>,
+)> {
+    opentelemetry::global::set_text_map_propagator(TraceContextPropagator::new());
+
+    let (exporter, rx_export, rx_shutdown) =
+        opentelemetry_sdk::testing::trace::new_tokio_test_exporter();
+
+    let otel_layer = open_telemetry_layer(exporter);
+    let subscriber = Registry::default()
+        .with(stdout_layer().with_filter(env_or_default_filter()))
+        .with(otel_layer.with_filter(env_or_default_filter()));
+    tracing::subscriber::set_global_default(subscriber)?;
+    Ok((rx_export, rx_shutdown))
+}
+
+fn open_telemetry_layer<S>(
+    exporter: impl SpanExporter + 'static,
+) -> Box<dyn tracing_subscriber::Layer<S> + Send + Sync + 'static>
 where
     S: Subscriber + Sync + Send,
     for<'a> S: tracing_subscriber::registry::LookupSpan<'a>,
 {
-    let exporter = opentelemetry_otlp::SpanExporter::builder()
-        .with_tonic()
-        .with_tls_config(tonic::transport::ClientTlsConfig::new().with_native_roots())
-        .build()
-        .unwrap();
-
     let tracer = opentelemetry_sdk::trace::SdkTracerProvider::builder()
         .with_batch_exporter(exporter)
         .build();

crates/utils/src/tracing/grpc.rs

@@ -1,19 +1,21 @@
-use tracing_opentelemetry::OpenTelemetrySpanExt;
+use super::OpenTelemetrySpanExt;
 
 /// A [`trace_fn`](tonic::transport::server::Server) implementation for the block producer which
 /// adds open-telemetry information to the span.
 ///
 /// Creates an `info` span following the open-telemetry standard: `block-producer.rpc/{method}`.
 /// Additionally also pulls in remote tracing context which allows the server trace to be connected
 /// to the client's origin trace.
-pub fn block_producer_trace_fn(request: &http::Request<()>) -> tracing::Span {
+pub fn block_producer_trace_fn<T>(request: &http::Request<T>) -> tracing::Span {
     let span = if let Some("SubmitProvenTransaction") = request.uri().path().rsplit('/').next() {
         tracing::info_span!("block-producer.rpc/SubmitProvenTransaction")
     } else {
         tracing::info_span!("block-producer.rpc/Unknown")
     };
 
-    add_otel_span_attributes(span, request)
+    span.set_parent(request);
+    span.set_http_attributes(request);
+    span
 }
 
 /// A [`trace_fn`](tonic::transport::server::Server) implementation for the store which adds
@@ -22,7 +24,7 @@ pub fn block_producer_trace_fn(request: &http::Request<()>) -> tracing::Span {
 /// Creates an `info` span following the open-telemetry standard: `store.rpc/{method}`. Additionally
 /// also pulls in remote tracing context which allows the server trace to be connected to the
 /// client's origin trace.
-pub fn store_trace_fn(request: &http::Request<()>) -> tracing::Span {
+pub fn store_trace_fn<T>(request: &http::Request<T>) -> tracing::Span {
     let span = match request.uri().path().rsplit('/').next() {
         Some("ApplyBlock") => tracing::info_span!("store.rpc/ApplyBlock"),
         Some("CheckNullifiers") => tracing::info_span!("store.rpc/CheckNullifiers"),
@@ -41,26 +43,8 @@ pub fn store_trace_fn(request: &http::Request<()>) -> tracing::Span {
         _ => tracing::info_span!("store.rpc/Unknown"),
     };
 
-    add_otel_span_attributes(span, request)
-}
-
-/// Adds remote tracing context to the span.
-///
-/// Could be expanded in the future by adding in more open-telemetry properties.
-fn add_otel_span_attributes(span: tracing::Span, request: &http::Request<()>) -> tracing::Span {
-    // Pull the open-telemetry parent context using the HTTP extractor. We could make a more
-    // generic gRPC extractor by utilising the gRPC metadata. However that
-    //     (a) requires cloning headers,
-    //     (b) we would have to write this ourselves, and
-    //     (c) gRPC metadata is transferred using HTTP headers in any case.
-    use tracing_opentelemetry::OpenTelemetrySpanExt;
-    let otel_ctx = opentelemetry::global::get_text_map_propagator(|propagator| {
-        propagator.extract(&MetadataExtractor(&tonic::metadata::MetadataMap::from_headers(
-            request.headers().clone(),
-        )))
-    });
-    span.set_parent(otel_ctx);
-
+    span.set_parent(request);
+    span.set_http_attributes(request);
     span
 }
 
@@ -82,26 +66,6 @@ impl tonic::service::Interceptor for OtelInterceptor {
     }
 }
 
-struct MetadataExtractor<'a>(&'a tonic::metadata::MetadataMap);
-impl opentelemetry::propagation::Extractor for MetadataExtractor<'_> {
-    /// Get a value for a key from the `MetadataMap`.  If the value can't be converted to &str,
-    /// returns None
-    fn get(&self, key: &str) -> Option<&str> {
-        self.0.get(key).and_then(|metadata| metadata.to_str().ok())
-    }
-
-    /// Collect all the keys from the `MetadataMap`.
-    fn keys(&self) -> Vec<&str> {
-        self.0
-            .keys()
-            .map(|key| match key {
-                tonic::metadata::KeyRef::Ascii(v) => v.as_str(),
-                tonic::metadata::KeyRef::Binary(v) => v.as_str(),
-            })
-            .collect::<Vec<_>>()
-    }
-}
-
 struct MetadataInjector<'a>(&'a mut tonic::metadata::MetadataMap);
 impl opentelemetry::propagation::Injector for MetadataInjector<'_> {
     /// Set a key and value in the `MetadataMap`.  Does nothing if the key or value are not valid

crates/utils/src/tracing/span_ext.rs

@@ -1,4 +1,5 @@
 use core::time::Duration;
+use std::net::SocketAddr;
 
 use miden_objects::{block::BlockNumber, Digest};
 use opentelemetry::{trace::Status, Key, Value};
@@ -8,6 +9,16 @@ pub trait ToValue {
     fn to_value(&self) -> Value;
 }
 
+impl ToValue for Option<SocketAddr> {
+    fn to_value(&self) -> Value {
+        if let Some(socket_addr) = self {
+            socket_addr.to_string().into()
+        } else {
+            "no_remote_addr".into()
+        }
+    }
+}
+
 impl ToValue for Duration {
     fn to_value(&self) -> Value {
         self.as_secs_f64().into()
@@ -48,21 +59,53 @@ impl ToValue for i64 {
 ///
 /// This is a sealed trait. It and cannot be implemented outside of this module.
 pub trait OpenTelemetrySpanExt: private::Sealed {
+    fn set_parent<T>(&self, request: &http::Request<T>);
     fn set_attribute(&self, key: impl Into<Key>, value: impl ToValue);
+    fn set_http_attributes<T>(&self, request: &http::Request<T>);
     fn set_error(&self, err: &dyn std::error::Error);
+    fn context(&self) -> opentelemetry::Context;
 }
 
 impl<S> OpenTelemetrySpanExt for S
 where
     S: tracing_opentelemetry::OpenTelemetrySpanExt,
 {
+    /// Sets the parent context by extracting HTTP metadata from the request.
+    fn set_parent<T>(&self, request: &http::Request<T>) {
+        // Pull the open-telemetry parent context using the HTTP extractor. We could make a more
+        // generic gRPC extractor by utilising the gRPC metadata. However that
+        //     (a) requires cloning headers,
+        //     (b) we would have to write this ourselves, and
+        //     (c) gRPC metadata is transferred using HTTP headers in any case.
+        let otel_ctx = opentelemetry::global::get_text_map_propagator(|propagator| {
+            propagator.extract(&MetadataExtractor(&tonic::metadata::MetadataMap::from_headers(
+                request.headers().clone(),
+            )))
+        });
+        tracing_opentelemetry::OpenTelemetrySpanExt::set_parent(self, otel_ctx);
+    }
+
+    /// Returns the context of `Span`.
+    fn context(&self) -> opentelemetry::Context {
+        tracing_opentelemetry::OpenTelemetrySpanExt::context(self)
+    }
+
     /// Sets an attribute on `Span`.
     ///
     /// Implementations for `ToValue` should be added to this crate (miden-node-utils).
     fn set_attribute(&self, key: impl Into<Key>, value: impl ToValue) {
         tracing_opentelemetry::OpenTelemetrySpanExt::set_attribute(self, key, value.to_value());
     }
 
+    /// Sets standard attributes to the `Span` based on an associated HTTP request.
+    fn set_http_attributes<T>(&self, request: &http::Request<T>) {
+        let remote_addr = request
+            .extensions()
+            .get::<tonic::transport::server::TcpConnectInfo>()
+            .and_then(tonic::transport::server::TcpConnectInfo::remote_addr);
+        OpenTelemetrySpanExt::set_attribute(self, "remote_addr", remote_addr);
+    }
+
     /// Sets a status on `Span` based on an error.
     fn set_error(&self, err: &dyn std::error::Error) {
         // Coalesce all sources into one string.
@@ -82,3 +125,24 @@ mod private {
     pub trait Sealed {}
     impl<S> Sealed for S where S: tracing_opentelemetry::OpenTelemetrySpanExt {}
 }
+
+/// Facilitates Open Telemetry metadata extraction for Tonic `MetadataMap`.
+struct MetadataExtractor<'a>(&'a tonic::metadata::MetadataMap);
+impl opentelemetry::propagation::Extractor for MetadataExtractor<'_> {
+    /// Get a value for a key from the `MetadataMap`.  If the value can't be converted to &str,
+    /// returns None
+    fn get(&self, key: &str) -> Option<&str> {
+        self.0.get(key).and_then(|metadata| metadata.to_str().ok())
+    }
+
+    /// Collect all the keys from the `MetadataMap`.
+    fn keys(&self) -> Vec<&str> {
+        self.0
+            .keys()
+            .map(|key| match key {
+                tonic::metadata::KeyRef::Ascii(v) => v.as_str(),
+                tonic::metadata::KeyRef::Binary(v) => v.as_str(),
+            })
+            .collect::<Vec<_>>()
+    }
+}