re_perf_telemetry 0.36.0

// Some well-known headers that we want to re-surface all the way up to the root of the active trace.
//
// See `re_protos::headers`.

const RERUN_HTTP_HEADER_ENTRY_ID: &str = "x-rerun-entry-id";
const RERUN_HTTP_HEADER_CLIENT_VERSION: &str = "x-rerun-client-version";
const RERUN_HTTP_HEADER_SERVER_VERSION: &str = "x-rerun-server-version";

// Server-injected trace id, returned to the client in response headers.
// Mirrors `re_protos::trace_id_layer::RERUN_HTTP_HEADER_REQUEST_TRACE_ID`
// (kept as a string here to avoid the dependency).
const RERUN_HTTP_HEADER_REQUEST_TRACE_ID: &str = "x-request-trace-id";

// --- Telemetry middlewares ---

/// Implements [`tower_http::trace::MakeSpan`] where the trace name is the gRPC method name.
///
/// See details in `make_span` on how we ensure context propagation.
///
/// We keep track of a bunch of relevant in-house state associated with the span in `SpanMetadata`.
#[derive(Debug, Clone)]
pub struct GrpcMakeSpan {
    gauge: opentelemetry::metrics::Gauge<u64>,
    // unfortunately we can't have different implementation of `MakeSpan` as that creates a ripple effect
    // through the entire hierarchy of types of the redap client stack and its usage, hence to disable
    // the span creation, we create noop spans instead if telemetry is disabled at runtime
    create_noop_spans: bool,
}

impl GrpcMakeSpan {
    pub fn new() -> Self {
        // if telemetry is not explicitly enabled through an env var, we create noop spans
        let create_noop_spans = !std::env::var("TELEMETRY_ENABLED")
            .is_ok_and(|v| v == "1" || v.to_lowercase() == "true" || v.to_lowercase() == "yes");

        let meter = opentelemetry::global::meter("grpc");
        let gauge = meter
            .u64_gauge("grpc_make_span_state_size")
            .with_description("Size of the SpanMetadata state")
            .build();
        Self {
            gauge,
            create_noop_spans,
        }
    }
}

impl Default for GrpcMakeSpan {
    fn default() -> Self {
        Self::new()
    }
}

impl<B> tower_http::trace::MakeSpan<B> for GrpcMakeSpan {
    fn make_span(&mut self, request: &http::Request<B>) -> tracing::Span {
        if self.create_noop_spans {
            return tracing::Span::none();
        }

        // Extract `OpenTelemetry` context from headers before creating the span
        let parent_ctx = opentelemetry::global::get_text_map_propagator(|prop| {
            prop.extract(&opentelemetry_http::HeaderExtractor(request.headers()))
        });

        // Pull the rerun session id out of the inbound `tracestate` (if any) so we can
        // record it directly as a span field at construction time, instead of relying on
        // a separate `tracing_subscriber::Layer` that records into a pre-declared field.
        //
        // Recording the value here means the field is populated before the span is exported,
        // and we sidestep the "silent no-op when the field wasn't pre-declared on the span"
        // pitfall of `Span::record`.
        let rerun_session_id = {
            use opentelemetry::trace::TraceContextExt as _;
            parent_ctx
                .span()
                .span_context()
                .trace_state()
                .get(crate::RERUN_SESSION_TRACESTATE_KEY)
                .and_then(crate::RerunTracingSessionId::parse)
                .map(String::from)
        };

        // This replaces the current tracing context with the extracted one, and it ensures that
        // any spans created within this scope will be children of the extracted context.
        // Note on the guard: guard must stay alive until after span creation so that tracing::span!()
        // can inherit the trace context. It is obviously dropped at the end of this function, but that's
        // ok as the future spans will inherit *this* span, not the old context. This makes sure that
        // we have propagation e2e (both client and server side), so change carefully.
        let _guard = parent_ctx.attach();

        let endpoint = request.uri().path().to_owned();
        // Split "/package.Service/Method" into rpc.service and rpc.method per OTel conventions.
        let (rpc_service, rpc_method) = endpoint
            .strip_prefix('/')
            .and_then(|s| s.split_once('/'))
            .unwrap_or(("", ""));
        let url = request
            .uri()
            .to_string()
            .strip_suffix(&endpoint)
            .map(ToOwned::to_owned);

        let client_version = request
            .headers()
            .get(RERUN_HTTP_HEADER_CLIENT_VERSION)
            .and_then(|v| v.to_str().ok().map(ToOwned::to_owned));

        let server_version = request
            .headers()
            .get(RERUN_HTTP_HEADER_SERVER_VERSION)
            .and_then(|v| v.to_str().ok().map(ToOwned::to_owned));

        let email = request
            .headers()
            .get("authorization")
            .and_then(|auth| auth.to_str().ok()?.strip_prefix("Bearer "))
            .and_then(|token| token.split('.').skip(1).take(1).next())
            .and_then(|data| {
                use base64::Engine as _;
                use base64::engine::general_purpose;
                general_purpose::STANDARD_NO_PAD.decode(data).ok()
            })
            .and_then(|data| {
                #[derive(serde::Deserialize)]
                struct TokenData {
                    sub: String,
                }
                serde_json::from_slice::<TokenData>(&data)
                    .ok()
                    .map(|data| data.sub)
            });

        let entry_id = request
            .headers()
            .get(RERUN_HTTP_HEADER_ENTRY_ID)
            .and_then(|v| v.to_str().ok().map(ToOwned::to_owned));

        // NOTE: Remember: the span we're creating here will propagate no matter what -- there is
        // no sampling at the `tracing` level, only at the `opentelemetry` level.
        // We use that fact to our advantage in order to carry a bunch of state around across all
        // the stages of the request (first response, first chunk, end-of-stream, etc).
        let span = tracing::span!(
            tracing::Level::INFO,
            "<request>",
            otel.name = %endpoint,
            url,
            method = %request.method(),

            // OTel semantic conventions for gRPC (https://opentelemetry.io/docs/specs/semconv/rpc/grpc/):
            rpc.system = "grpc",
            rpc.service = %rpc_service,
            rpc.method = %rpc_method,

            // The rerun session id, recorded as a top-level span field so it is queryable
            // in Tempo as `{ .rerun_session_id = "…" }`. Extracted from the inbound
            // `tracestate` header above. Empty when no `tracing_session()` is active on
            // the client.
            rerun_session_id = rerun_session_id.as_deref(),

            // The gRPC status code (e.g. "Ok", "AlreadyExists", "DeadlineExceeded").
            // Filled in later by `GrpcOnResponse` or `GrpcOnEos`, depending on the endpoint type (unary vs streaming).
            grpc_status = tracing::field::Empty,

            // The trace id reported back by the server in the `x-request-trace-id` response header.
            // Filled in client-side by `ClientOnResponse` so we can correlate client spans with the
            // server-side trace.
            server_trace_id = tracing::field::Empty,
        );

        let size = SpanMetadata::insert_opt(
            span.id(),
            SpanMetadata {
                endpoint,
                client_version,
                server_version,
                email,
                entry_id,
                first_chunk_returned: false,
                grpc_eos_classifier: None,
            },
            false,
        );
        self.gauge.record(size as _, &[]);

        span
    }
}

/// The global storage for [`SpanMetadata`]s.
///
/// Keeps track of relevant in-house context/metadata for all on-going gRPC spans.
///
/// We could also build a full-fledged `tracing::Subscriber` instead, but at this point I'd rather
/// _accomplish something_ instead of implementing yet another 50 layers of abstraction.
///
/// The state is written to and read from by our different gRPC middlewares. In particular,
/// [`GrpcOnEos`] is responsible for cleaning up dead entries.
static SPAN_METADATA: std::sync::OnceLock<
    parking_lot::RwLock<ahash::HashMap<tracing::span::Id, SpanMetadata>>,
> = std::sync::OnceLock::new();

/// Custom state/context/metadata that we associate with the spans we generate in our [`GrpcMakeSpan`] middleware.
///
/// All this state is stored in `SPAN_METADATA`.
#[derive(Debug, Clone)]
struct SpanMetadata {
    /// Which gRPC endpoint? Extracted from h2 headers.
    endpoint: String,

    /// The identity and semantic version advertised by the gRPC client.
    ///
    /// Extracted from h2 headers. See also `re_protos::headers::RERUN_HTTP_HEADER_CLIENT_VERSION`.
    client_version: Option<String>,

    /// The identity and semantic version advertised by the gRPC server.
    ///
    /// Extracted from h2 headers. See also `re_protos::headers::RERUN_HTTP_HEADER_SERVER_VERSION`.
    server_version: Option<String>,

    /// What email, if any? Extracted from h2 auth headers.
    email: Option<String>,

    /// What entry ID, if any? Extracted from h2 Rerun extension headers.
    entry_id: Option<String>,

    /// Has the gRPC stream associated with this span streamed back its first chunk of data yet?
    ///
    /// This is set by our [`GrpcOnFirstBodyChunk`] middleware.
    first_chunk_returned: bool,

    /// If the gRPC stream's failure outcome is to be determined by its response stream, this will
    /// tell us how.
    ///
    /// This is set by our [`GrpcOnResponse`] middleware.
    grpc_eos_classifier: Option<tower_http::classify::GrpcEosErrorsAsFailures>,
}

impl Default for SpanMetadata {
    fn default() -> Self {
        Self {
            endpoint: "undefined".to_owned(),
            client_version: None,
            server_version: None,
            email: None,
            entry_id: None,
            first_chunk_returned: false,
            grpc_eos_classifier: None,
        }
    }
}

/// Number of in-flight gRPC requests/streams, labeled by `endpoint`.
///
/// An entry lives in `SPAN_METADATA` for exactly the lifetime of a request (inserted in
/// [`GrpcMakeSpan`]'s `make_span`, removed at end-of-stream, on immediate error, or on span close),
/// so counting `+1` on a genuine insert and `-1` on a real removal makes this gauge exactly the
/// number of live entries per endpoint — balanced by construction, no matter which termination path
/// a request takes. It is a non-monotonic sum, so it exports as a gauge.
fn requests_in_flight() -> &'static opentelemetry::metrics::UpDownCounter<i64> {
    static INSTANCE: std::sync::OnceLock<opentelemetry::metrics::UpDownCounter<i64>> =
        std::sync::OnceLock::new();
    INSTANCE.get_or_init(|| {
        opentelemetry::global::meter("grpc")
            .i64_up_down_counter("grpc_requests_in_flight")
            .with_description("Number of in-flight gRPC requests/streams, by endpoint")
            .build()
    })
}

impl SpanMetadata {
    /// Returns the new size of the map.
    #[expect(clippy::needless_pass_by_value)]
    fn insert(span_id: tracing::span::Id, metadata: Self, expect_conflict: bool) -> usize {
        let endpoint = metadata.endpoint.clone();
        let (is_overwrite, new_len) = {
            let mut state = SPAN_METADATA.get_or_init(Default::default).write();
            let is_overwrite = state.insert(span_id.clone(), metadata).is_some();
            let new_len = state.len();
            (is_overwrite, new_len)
        };

        if is_overwrite && !expect_conflict {
            tracing::warn!(id=?span_id, "overwritten span metadata -- this should never happen");
        }

        // Only a genuine new entry (not an in-place update of an existing request's metadata) adds
        // an in-flight request; the matching -1 happens in `remove`/`remove_silent`.
        if !is_overwrite {
            requests_in_flight().add(1, &[opentelemetry::KeyValue::new("endpoint", endpoint)]);
        }

        new_len
    }

    /// Returns the new size of the map.
    fn insert_opt(
        span_id: Option<tracing::span::Id>,
        metadata: Self,
        expect_conflict: bool,
    ) -> usize {
        if let Some(span_id) = span_id {
            Self::insert(span_id, metadata, expect_conflict)
        } else {
            SPAN_METADATA.get_or_init(Default::default).read().len()
        }
    }

    fn get(span_id: &tracing::span::Id) -> Option<Self> {
        let md = SPAN_METADATA
            .get()
            .and_then(|spans| spans.read().get(span_id).cloned());

        if md.is_none() {
            tracing::warn!(id=?span_id, "missing span metadata -- this should never happen");
        }

        md
    }

    fn get_opt(span_id: Option<&tracing::span::Id>) -> Option<Self> {
        span_id.and_then(Self::get)
    }

    fn remove(span_id: &tracing::span::Id) -> Option<Self> {
        let md = SPAN_METADATA
            .get()
            .and_then(|spans| spans.write().remove(span_id));

        if let Some(md) = &md {
            requests_in_flight().add(
                -1,
                &[opentelemetry::KeyValue::new(
                    "endpoint",
                    md.endpoint.clone(),
                )],
            );
        } else {
            tracing::warn!(id=?span_id, "missing span metadata -- this should never happen");
        }

        md
    }

    fn remove_opt(span_id: Option<&tracing::span::Id>) -> Option<Self> {
        span_id.and_then(Self::remove)
    }

    /// Silently removes metadata for a span, returning `Some` if it existed.
    ///
    /// Unlike [`Self::remove`], this does NOT warn when the metadata is missing.
    /// Used by [`SpanMetadataCleanupLayer::on_close`] where the metadata has typically
    /// already been removed by [`GrpcOnEos`] or [`GrpcOnResponse`].
    fn remove_silent(span_id: &tracing::span::Id) -> Option<Self> {
        let spans = SPAN_METADATA.get()?;
        let md = spans.write().remove(span_id);
        if let Some(md) = &md {
            requests_in_flight().add(
                -1,
                &[opentelemetry::KeyValue::new(
                    "endpoint",
                    md.endpoint.clone(),
                )],
            );
        }
        md
    }
}

// ---

/// Implements a [`tower_http::trace::OnRequest`] middleware.
#[derive(Clone)]
pub struct GrpcOnRequest {}

impl GrpcOnRequest {
    #[expect(clippy::new_without_default)] // future-proofing
    pub fn new() -> Self {
        Self {}
    }
}

impl<B> tower_http::trace::OnRequest<B> for GrpcOnRequest {
    fn on_request(&mut self, request: &http::Request<B>, span: &tracing::Span) {
        let Some(span_metadata) = SpanMetadata::get_opt(span.id().as_ref()) else {
            tracing::info!(
                uri = %request.uri(),
                "grpc_on_request with unknown span.id"
            );
            return;
        };

        let SpanMetadata {
            endpoint,
            client_version,
            server_version,
            email,
            entry_id: dataset_id,
            first_chunk_returned: _,
            grpc_eos_classifier: _,
        } = span_metadata.clone();

        let client_version = client_version.as_deref().unwrap_or("undefined");
        let server_version = server_version.as_deref().unwrap_or("undefined");
        let email = email.as_deref().unwrap_or("undefined");
        let dataset_id = dataset_id.as_deref().unwrap_or("undefined");

        // NOTE: repeat all these attributes so services such as CloudWatch, which don't really
        // support OTLP, can actually see them.
        tracing::info!(%endpoint, %client_version, %server_version, %email, %dataset_id, "grpc_on_request");
    }
}

// ---

/// Implements a [`tower_http::trace::OnResponse`] middleware.
#[derive(Clone)]
pub struct GrpcOnResponse {
    histogram: opentelemetry::metrics::Histogram<f64>,
    eos_counter: opentelemetry::metrics::Counter<u64>,
}

#[derive(Default)]
pub struct GrpcOnResponseOptions {
    pub fully_qualified_route_names: Option<Vec<String>>,
}

impl GrpcOnResponse {
    pub fn new(options: GrpcOnResponseOptions) -> Self {
        let meter = opentelemetry::global::meter("grpc");
        let histogram = meter
            .f64_histogram("grpc_on_response_ms")
            .with_description("Latency percentiles for all gRPC endpoints (\"time to response\")")
            .build();
        let eos_counter = meter
            .u64_counter("grpc_on_eos")
            .with_description("End-of-stream counter for all gRPC endpoints")
            .build();

        if let Some(fully_qualified_rout_names) = options.fully_qualified_route_names {
            for endpoint in &fully_qualified_rout_names {
                eos_counter.add(
                    0,
                    &GrpcEosMetadata {
                        endpoint: endpoint.clone(),
                        grpc_status: String::new(),
                        client_version: None,
                        server_version: None,
                        email: None,
                    }
                    .otel_kvs(),
                );
            }
        }
        Self {
            histogram,
            eos_counter,
        }
    }
}

struct GrpcEosMetadata {
    endpoint: String,
    grpc_status: String,
    client_version: Option<String>,
    server_version: Option<String>,
    email: Option<String>,
}

impl GrpcEosMetadata {
    fn otel_kvs(self) -> Vec<opentelemetry::KeyValue> {
        vec![
            opentelemetry::KeyValue::new("endpoint", self.endpoint),
            opentelemetry::KeyValue::new("grpc_status", self.grpc_status),
            opentelemetry::KeyValue::new(
                "client_version",
                self.client_version
                    .as_deref()
                    .unwrap_or("undefined")
                    .to_owned(),
            ),
            opentelemetry::KeyValue::new(
                "server_version",
                self.server_version
                    .as_deref()
                    .unwrap_or("undefined")
                    .to_owned(),
            ),
            opentelemetry::KeyValue::new(
                "email",
                self.email.as_deref().unwrap_or("undefined").to_owned(),
            ),
        ]
    }
}

impl<B> tower_http::trace::OnResponse<B> for GrpcOnResponse {
    fn on_response(
        self,
        response: &http::Response<B>,
        latency: std::time::Duration,
        span: &tracing::Span,
    ) {
        let Some(span_metadata) = SpanMetadata::get_opt(span.id().as_ref()) else {
            return;
        };

        let SpanMetadata {
            endpoint,
            client_version,
            server_version,
            email,
            entry_id: dataset_id,
            first_chunk_returned: _,
            grpc_eos_classifier: _,
        } = span_metadata.clone();

        // Record response metrics and optionally set the final `grpc_status` span attribute.
        // Pass `Some(code)` when the final gRPC status is known, `None` when it will be
        // determined later by `GrpcOnEos` (streaming responses).
        let record = |span: &tracing::Span, grpc_code: Option<tonic::Code>| {
            let grpc_status = if let Some(grpc_code) = grpc_code {
                format!("{grpc_code:?}")
            } else {
                "<pending>".to_owned()
            };
            span.record("grpc_status", grpc_status.as_str());
            let http_status = response.status().as_str().to_owned();

            let client_version = client_version.as_deref().unwrap_or("undefined");
            let server_version = server_version.as_deref().unwrap_or("undefined");
            let email = email.as_deref().unwrap_or("undefined");
            let dataset_id = dataset_id.as_deref().unwrap_or("undefined");

            // NOTE: repeat all these attributes so services such as CloudWatch, which don't really
            // support OTLP, can actually see them.
            match grpc_status.as_str() {
                "Ok" => {
                    tracing::info!(%endpoint, %grpc_status, %http_status, %client_version, %server_version, %email, %dataset_id, ?latency, "grpc_on_response");
                }
                "<pending>" => {
                    // Streaming response — final gRPC status will be logged by `GrpcOnEos`.
                    tracing::debug!(%endpoint, %grpc_status, %http_status, %client_version, %server_version, %email, %dataset_id, ?latency, "grpc_on_response");
                }
                _ => {
                    tracing::error!(%endpoint, %grpc_status, %http_status, %client_version, %server_version, %email, %dataset_id, ?latency, "grpc_on_response");
                }
            }

            self.histogram.record(
                latency.as_secs_f64() * 1000.0,
                &[
                    opentelemetry::KeyValue::new("endpoint", endpoint.clone()),
                    opentelemetry::KeyValue::new("grpc_status", grpc_status),
                    opentelemetry::KeyValue::new("http_status", http_status),
                    opentelemetry::KeyValue::new("server_version", server_version.to_owned()),
                ],
            );
        };

        use tower_http::classify::ClassifyResponse as _;
        let classified =
            tower_http::classify::GrpcErrorsAsFailures::new().classify_response(response);
        match classified {
            tower_http::classify::ClassifiedResponse::Ready(Err(err)) => {
                let grpc_code = match err {
                    tower_http::classify::GrpcFailureClass::Code(code) => {
                        tonic::Code::from_i32(code.into())
                    }
                    tower_http::classify::GrpcFailureClass::Error(err) => {
                        tonic::Status::from_error(err.into()).code()
                    }
                };
                record(span, Some(grpc_code));

                // For immediate errors, emit grpc_on_eos counter here since on_eos won't be called
                let grpc_status = format!("{grpc_code:?}"); // NOTE: The debug repr is the enum variant name (e.g. DeadlineExceeded).

                self.eos_counter.add(
                    1,
                    &GrpcEosMetadata {
                        endpoint: endpoint.clone(),
                        grpc_status,
                        client_version,
                        server_version,
                        email,
                    }
                    .otel_kvs(),
                );

                // Remove metadata since on_eos won't be called for immediate errors
                SpanMetadata::remove_opt(span.id().as_ref());
            }

            tower_http::classify::ClassifiedResponse::Ready(Ok(())) => {
                record(span, Some(tonic::Code::Ok));
            }

            tower_http::classify::ClassifiedResponse::RequiresEos(eos) => {
                // Final gRPC status is unknown until end-of-stream; `GrpcOnEos` will set it.
                record(span, None);
                SpanMetadata::insert_opt(
                    span.id(),
                    SpanMetadata {
                        grpc_eos_classifier: Some(eos),
                        ..span_metadata
                    },
                    true,
                );
            }
        }
    }
}

/// Implements a [`tower_http::trace::OnBodyChunk`] middleware, but only accounts for the first one.
#[derive(Clone)]
pub struct GrpcOnFirstBodyChunk {
    histogram: opentelemetry::metrics::Histogram<f64>,
}

impl GrpcOnFirstBodyChunk {
    #[expect(clippy::new_without_default)] // future-proofing
    pub fn new() -> Self {
        let meter = opentelemetry::global::meter("grpc");
        let histogram = meter
            .f64_histogram("grpc_on_first_body_chunk_ms")
            .with_description(
                "Latency percentiles for all gRPC endpoints (\"time to first chunk\")",
            )
            .build();
        Self { histogram }
    }
}

impl<B> tower_http::trace::OnBodyChunk<B> for GrpcOnFirstBodyChunk {
    fn on_body_chunk(&mut self, _: &B, latency: std::time::Duration, span: &tracing::Span) {
        let Some(span_metadata) = SpanMetadata::get_opt(span.id().as_ref()) else {
            return;
        };

        let SpanMetadata {
            endpoint,
            client_version,
            server_version,
            email,
            entry_id: dataset_id,
            first_chunk_returned,
            grpc_eos_classifier: _,
        } = span_metadata.clone();

        if !first_chunk_returned {
            let client_version = client_version.as_deref().unwrap_or("undefined");
            let server_version = server_version.as_deref().unwrap_or("undefined");
            let email = email.as_deref().unwrap_or("undefined");
            let dataset_id = dataset_id.as_deref().unwrap_or("undefined");

            // NOTE: repeat all these attributes so services such as CloudWatch, which don't really
            // support OTLP, can actually see them.
            tracing::debug!(%endpoint, %client_version, %server_version, %email, %dataset_id, ?latency, "grpc_on_first_body_chunk");

            self.histogram.record(
                latency.as_secs_f64() * 1000.0,
                &[
                    opentelemetry::KeyValue::new("endpoint", endpoint),
                    opentelemetry::KeyValue::new("server_version", server_version.to_owned()),
                ],
            );

            SpanMetadata::insert_opt(
                span.id(),
                SpanMetadata {
                    first_chunk_returned: true,
                    ..span_metadata
                },
                true,
            );
        }
    }
}

/// Implements a [`tower_http::trace::OnEos`] middleware.
///
/// Note that even unary endpoints are implemented as streams internally, and will therefore be
/// impacted by this middleware. This is especially important at this middleware is responsible for
/// GC'ing the contents of `SPAN_METADATA`.
#[derive(Clone)]
pub struct GrpcOnEos {
    counter: opentelemetry::metrics::Counter<u64>,
}

impl GrpcOnEos {
    #[expect(clippy::new_without_default)] // future-proofing
    pub fn new() -> Self {
        let meter = opentelemetry::global::meter("grpc");
        let counter = meter
            .u64_counter("grpc_on_eos")
            .with_description("End-of-stream counter for all gRPC endpoints")
            .build();
        Self { counter }
    }
}

impl tower_http::trace::OnEos for GrpcOnEos {
    fn on_eos(
        self,
        trailers: Option<&http::HeaderMap>,
        duration: std::time::Duration,
        span: &tracing::Span,
    ) {
        let Some(span_metadata) = SpanMetadata::remove_opt(span.id().as_ref()) else {
            return;
        };

        let SpanMetadata {
            endpoint,
            client_version,
            server_version,
            email,
            entry_id: dataset_id,
            first_chunk_returned: _,
            grpc_eos_classifier,
        } = span_metadata;

        let grpc_code = if let Some(classifier) = grpc_eos_classifier {
            use tower_http::classify::ClassifyEos as _;
            match classifier.classify_eos(trailers) {
                Ok(()) => tonic::Code::Ok,
                Err(err) => match err {
                    tower_http::classify::GrpcFailureClass::Code(code) => {
                        tonic::Code::from_i32(code.into())
                    }
                    tower_http::classify::GrpcFailureClass::Error(err) => {
                        tonic::Status::from_error(err.into()).code()
                    }
                },
            }
        } else {
            tracing::warn!(
                endpoint,
                client_version,
                server_version,
                email,
                dataset_id,
                "couldn't determine gRPC EOS status code"
            );
            tonic::Code::Unknown
        };
        let grpc_status = format!("{grpc_code:?}"); // NOTE: The debug repr is the enum variant name (e.g. DeadlineExceeded).
        span.record("grpc_status", &grpc_status);

        let client_version = client_version.as_deref().unwrap_or("undefined");
        let server_version = server_version.as_deref().unwrap_or("undefined");
        let email = email.as_deref().unwrap_or("undefined");
        let dataset_id = dataset_id.as_deref().unwrap_or("undefined");

        // NOTE: repeat all these attributes so services such as CloudWatch, which don't really
        // support OTLP, can actually see them.
        if grpc_status == "Ok" {
            tracing::info!(%endpoint, %grpc_status, %client_version, %server_version, %email, %dataset_id, ?duration, "grpc_on_eos");
        } else {
            tracing::error!(%endpoint, %grpc_status, %client_version, %server_version, %email, %dataset_id, ?duration, "grpc_on_eos");
        }

        self.counter.add(
            1,
            &[
                opentelemetry::KeyValue::new("endpoint", endpoint),
                opentelemetry::KeyValue::new("grpc_status", grpc_status),
                opentelemetry::KeyValue::new("client_version", client_version.to_owned()),
                opentelemetry::KeyValue::new("server_version", server_version.to_owned()),
                opentelemetry::KeyValue::new("email", email.to_owned()),
            ],
        );
    }
}

pub type ServerTelemetryLayer = tower_http::trace::TraceLayer<
    tower_http::trace::GrpcMakeClassifier,
    GrpcMakeSpan,
    GrpcOnRequest,
    GrpcOnResponse,
    GrpcOnFirstBodyChunk,
    GrpcOnEos,
>;

#[derive(Default)]
pub struct TelemetryLayerOptions {
    pub fully_qualified_route_names: Option<Vec<String>>,
}

/// Creates a new [`tower::Layer`] middleware that automatically:
/// * Traces gRPC requests and responses.
/// * Logs all gRPC responses (status, latency, etc).
/// * Measures all gRPC responses (status, latency, etc).
///
/// Note on callback behavior based on testing:
/// * unary endpoint success - `on_response` and `on_eos` called
/// * unary endpoint failure - both `on_response` (error handling path) and `on_failure` called. We currently don't have `on_failure` implemented.
/// * streaming endpoint success - `on_response` and `on_eos` called
/// * streaming endpoint immediate error (no stream started)  - `on_response` (error handling path) called and `on_failure` (again, not implemented)
/// * streaming endpoint mid stream error - `on_response` called (but only initially with OK code, no error detected here), `on_eos` called (and existing error handling logic executed).
///   `on_failure` is not called here and from code inspection it seems that is only called for immediate failures and polling frame errors (transport errors, although not verified with testing)
pub fn new_server_telemetry_layer(options: TelemetryLayerOptions) -> ServerTelemetryLayer {
    tower_http::trace::TraceLayer::new_for_grpc()
        .make_span_with(GrpcMakeSpan::new())
        .on_request(GrpcOnRequest::new())
        .on_response(GrpcOnResponse::new(GrpcOnResponseOptions {
            fully_qualified_route_names: options.fully_qualified_route_names,
        }))
        .on_body_chunk(GrpcOnFirstBodyChunk::new())
        .on_eos(GrpcOnEos::new())
}

/// Implements a [`tower_http::trace::OnResponse`] middleware for the gRPC client.
///
/// Records the server-reported trace id (from the `x-request-trace-id` response header)
/// onto the client span, so client-side traces can be correlated with the server-side trace.
#[derive(Debug, Clone, Default)]
pub struct ClientOnResponse {}

impl ClientOnResponse {
    pub fn new() -> Self {
        Self {}
    }
}

impl<B> tower_http::trace::OnResponse<B> for ClientOnResponse {
    fn on_response(
        self,
        response: &http::Response<B>,
        _latency: std::time::Duration,
        span: &tracing::Span,
    ) {
        if let Some(trace_id) = response
            .headers()
            .get(RERUN_HTTP_HEADER_REQUEST_TRACE_ID)
            .and_then(|v| v.to_str().ok())
        {
            span.record("server_trace_id", trace_id);
        }
    }
}

pub type ClientTelemetryLayer = tower::layer::util::Stack<
    tonic::service::interceptor::InterceptorLayer<TracingInjectorInterceptor>,
    tower::layer::util::Stack<
        tower_http::trace::TraceLayer<
            tower_http::trace::GrpcMakeClassifier,
            GrpcMakeSpan,
            tower_http::trace::DefaultOnRequest,
            ClientOnResponse,
        >,
        tower::layer::util::Identity,
    >,
>;

/// Creates a new [`tower::Layer`] middleware that automatically:
/// * Traces gRPC requests and responses.
/// * Logs all gRPC responses (status, latency, etc).
/// * Measures all gRPC responses (status, latency, etc).
//
// TODO(cmc): at the moment there's little value to have anything beyond traces on the client, but
// we ultimately can add all the same things that we have on the server as we need them.
pub fn new_client_telemetry_layer() -> ClientTelemetryLayer {
    let trace_layer = tower_http::trace::TraceLayer::new_for_grpc()
        // Note: we're actually disabling all DEBUG level logs for `tower` in re_log, so if you want to enable it
        // you'll need to adjust that as well. See crates/utils/re_log/src/lib.rs
        .on_failure(DefaultOnFailure::new().level(tracing::Level::DEBUG))
        .make_span_with(GrpcMakeSpan::new())
        .on_response(ClientOnResponse::new());

    tower::ServiceBuilder::new()
        .layer(trace_layer)
        .layer(TracingInjectorInterceptor::new_layer())
        .into_inner()
}

// --- Propagation middlewares ---

/// This implements a [`tonic::service::Interceptor`] that injects trace/span metadata into the
/// request headers, according to W3C standards.
///
/// This trace/span information is extracted from the currently opened [`tracing::Span`], then
/// converting to the `OpenTelemetry` format, and finally injected into the request headers, thereby
/// propagating the trace across network boundaries.
///
#[derive(Default, Clone)]
pub struct TracingInjectorInterceptor;

impl TracingInjectorInterceptor {
    /// Creates a new [`tower::Layer`] middleware that automatically applies the injector.
    ///
    /// See also [`new_client_telemetry_layer`].
    pub fn new_layer() -> tonic::service::interceptor::InterceptorLayer<Self> {
        tonic::service::interceptor::InterceptorLayer::new(Self)
    }
}

impl tonic::service::Interceptor for TracingInjectorInterceptor {
    fn call(&mut self, mut req: tonic::Request<()>) -> tonic::Result<tonic::Request<()>> {
        struct MetadataMap<'a>(&'a mut tonic::metadata::MetadataMap);

        impl opentelemetry::propagation::Injector for MetadataMap<'_> {
            fn set(&mut self, key: &str, value: String) {
                if let Ok(key) = tonic::metadata::MetadataKey::from_bytes(key.as_bytes())
                    && let Ok(val) = tonic::metadata::MetadataValue::try_from(&value)
                {
                    self.0.insert(key, val);
                }
            }
        }

        // Grab the trace information from `tracing`, and convert that into `opentelemetry`.
        use tracing_opentelemetry::OpenTelemetrySpanExt as _;
        let cx = tracing::Span::current().context();

        // Inject the opentelemetry-formatted trace information into the headers.
        opentelemetry::global::get_text_map_propagator(|propagator| {
            propagator.inject_context(&cx, &mut MetadataMap(req.metadata_mut()));
        });

        Ok(req)
    }
}

// ---

use tower_http::trace::DefaultOnFailure;
use tracing::Subscriber;
use tracing::span::Id;
use tracing_subscriber::Layer;
use tracing_subscriber::layer::Context;

/// A [`tracing_subscriber::Layer`] that cleans up `SpanMetadata` entries when spans close.
///
/// In the normal flow, metadata is removed by [`GrpcOnEos`] `on_eos` (streaming responses)
/// or [`GrpcOnResponse`] `on_response` (immediate errors). However, if the client disconnects
/// or a transport error occurs, `on_eos` may never be called, leaving stale entries behind.
///
/// Because the `tracing` crate recycles span IDs after a span closes, stale entries cause
/// "overwritten span metadata" warnings when the recycled ID is reused by a new request.
///
/// `on_close` is called when all clones of a span are dropped, and critically, **before** the
/// span ID is recycled. This guarantees any leaked metadata is cleaned up before the ID can
/// be reused.
#[derive(Default)]
pub struct SpanMetadataCleanupLayer {
    _private: (),
}

impl<S> Layer<S> for SpanMetadataCleanupLayer
where
    S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
{
    fn on_close(&self, id: Id, _ctx: Context<'_, S>) {
        SpanMetadata::remove_silent(&id);
    }
}