tonic with tracing, opentelemetry, and context propagation

service.rs

use opentelemetry::api::propagation::text_propagator::HttpTextFormat;
use tracing::{debug, info};
use tracing_futures::Instrument as _;
use tracing_opentelemetry::OpenTelemetrySpanExt as _;

#[tonic::async_trait]
impl crate::protobuf::kubernetes_server::Kubernetes for Service {
   async fn kubernetes_deployment_component_get(
        &self,
        mut request: tonic::Request<crate::protobuf::KubernetesDeploymentComponentGetRequest>,
    ) -> std::result::Result<
        tonic::Response<crate::protobuf::KubernetesDeploymentComponentGetReply>,
        tonic::Status,
    > {
        let trace_propagator =
            opentelemetry::api::trace::trace_context_propagator::TraceContextPropagator::new();
        let span_context = {
            let metadata_wrapper = TonicMetaWrapper(request.metadata_mut());
            trace_propagator.extract(&metadata_wrapper)
        };
        let span = tracing::span!(
            tracing::Level::INFO,
            "kubernetes.kubernetes_deployment_component_get",
            metadata.content_type = tracing::field::Empty,
            authenticated = tracing::field::Empty,
            userId = tracing::field::Empty,
            billingAccountId = tracing::field::Empty,
            http.user_agent = tracing::field::Empty,
        );
        span.set_parent(span_context);

        {
            let metadata = request.metadata();
            if let Some(raw_value) = metadata.get("authenticated") {
                let value = raw_value.to_str().unwrap_or("unserializable");
                span.record("authenticated", &tracing::field::display(value));
            }
            if let Some(raw_value) = metadata.get("userid") {
                let value = raw_value.to_str().unwrap_or("unserializable");
                span.record("userId", &tracing::field::display(value));
            }
            if let Some(raw_value) = metadata.get("billingAccountId") {
                let value = raw_value.to_str().unwrap_or("unserializable");
                span.record("billingAccountId", &tracing::field::display(value));
            }
            if let Some(raw_value) = metadata.get("user-agent") {
                let value = raw_value.to_str().unwrap_or("unserializable");
                span.record("http.user_agent", &tracing::field::display(value));
            }
        }

        async {
            si_account::authorize::authnz(
                &self.db,
                &request,
                "kubernetes_deployment_component_get",
            )
            .await?;

            let inner = request.into_inner();
            let id = inner
                .id
                .ok_or_else(|| si_data::DataError::RequiredField("id".to_string()))?;

            let output = crate::protobuf::KubernetesDeploymentComponent::get(&self.db, &id).await?;

            Ok(tonic::Response::new(
                crate::protobuf::KubernetesDeploymentComponentGetReply { item: Some(output) },
            ))
        }
        .instrument(span)
        .await
    }
}

struct TonicMetaWrapper<'a>(&'a mut tonic::metadata::MetadataMap);

impl<'a> opentelemetry::api::propagation::Carrier for TonicMetaWrapper<'a> {
    fn get(&self, key: &'static str) -> Option<&str> {
        let raw_value = self.0.get(key)?;
        match raw_value.to_str() {
            Ok(value) => Some(value),
            Err(_e) => {
                debug!("Cannot extract header for trace parent, not a string");
                None
            }
        }
    }

    fn set(&mut self, key: &'static str, raw_value: String) {
        let value = match tonic::metadata::MetadataValue::from_str(&raw_value) {
            Ok(value) => value,
            Err(_e) => {
                debug!("Cannot insert header for trace parent, not a string");
                debug!("Inserting the empty string");
                tonic::metadata::MetadataValue::from_str("").unwrap()
            }
        };
        self.0.insert(key, value);
    }
}