wasmcloud-control-interface 3.0.1

// COPIED DIRECTLY FROM https://github.com/wasmCloud/weld/blob/wasmbus-rpc-v0.13.0/rpc-rs/src/otel.rs (minus unused functionality)

//! Contains helpers and code for enabling [OpenTelemetry](https://opentelemetry.io/) tracing for
//! wasmcloud-invocations. Please note that right now this is only supported for providers. This
//! module is only available with the `otel` feature enabled

use std::str::FromStr;

use async_nats::header::{HeaderMap, HeaderName, HeaderValue};
use opentelemetry::propagation::{Injector, TextMapPropagator};
use opentelemetry_sdk::propagation::TraceContextPropagator;
use tracing::span::Span;
use tracing_opentelemetry::OpenTelemetrySpanExt;

use crate::Result;

/// A convenience type that wraps a NATS [`HeaderMap`] and implements the [`Injector`] trait
#[derive(Debug, Default)]
#[non_exhaustive]
pub struct HeaderInjector {
    inner: HeaderMap,
}

impl HeaderInjector {
    /// Creates a new header injector from an iterator of string tuples
    pub fn from(headers: impl IntoIterator<Item = (String, String)>) -> Result<Self> {
        let inner = headers
            .into_iter()
            .map(
                |(h, v)| match (HeaderName::from_str(&h), HeaderValue::from_str(&v)) {
                    (Ok(h), Ok(v)) => Ok((h, v)),
                    (Err(e), _) => Err(format!("failed to build header name: {e}").into()),
                    (_, Err(e)) => Err(format!("failed to build header name: {e}").into()),
                },
            )
            .collect::<Result<HeaderMap>>()?;
        Ok(HeaderInjector { inner })
    }

    /// Creates a new injector using the given [`HeaderMap`]
    pub fn new(headers: HeaderMap) -> Self {
        HeaderInjector { inner: headers }
    }

    /// Convenience constructor that returns a new injector with the current span context already
    /// injected into the given header map
    pub fn new_with_span(headers: HeaderMap) -> Self {
        let mut header_map = Self::new(headers);
        header_map.inject_context();
        header_map
    }

    /// Convenience constructor that returns a new injector with the current span context already
    /// injected into a default [`HeaderMap`]
    pub fn default_with_span() -> Self {
        let mut header_map = Self::default();
        header_map.inject_context();
        header_map
    }

    /// Injects the current context from the span into the headers
    pub fn inject_context(&mut self) {
        let ctx_propagator = TraceContextPropagator::new();
        ctx_propagator.inject_context(&Span::current().context(), self);
    }
}

impl Injector for HeaderInjector {
    fn set(&mut self, key: &str, value: String) {
        self.inner.insert(key, value.as_ref());
    }
}

impl AsRef<HeaderMap> for HeaderInjector {
    fn as_ref(&self) -> &HeaderMap {
        &self.inner
    }
}

impl From<HeaderMap> for HeaderInjector {
    fn from(headers: HeaderMap) -> Self {
        HeaderInjector::new(headers)
    }
}

impl From<HeaderInjector> for HeaderMap {
    fn from(inj: HeaderInjector) -> Self {
        inj.inner
    }
}