laburnum 1.17.0

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

use std::collections::HashMap;

pub mod exporter;
pub mod prefixer;
use opentelemetry::{
  Context,
  propagation::{
    Extractor,
    Injector,
  },
};
// Re-export from otel

pub const OTEL_NAME: &str = "otel.name";
pub const OTEL_ORIGINAL_NAME: &str = "otel.original_name";
pub const OTEL_KIND: &str = "otel.kind";
pub const OTEL_STATUS_CODE: &str = "otel.status_code";
pub const OTEL_STATUS_MESSAGE: &str = "otel.status_message";

/// W3C Trace Context fields for distributed tracing propagation.
///
/// This struct holds the `traceparent` and `tracestate` headers as defined
/// by the W3C Trace Context specification.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct TraceContext {
  pub traceparent: Option<String>,
  pub tracestate:  Option<String>,
}

impl TraceContext {
  /// Create a new TraceContext with the given values.
  #[must_use]
  pub fn new(traceparent: Option<String>, tracestate: Option<String>) -> Self {
    Self {
      traceparent,
      tracestate,
    }
  }

  // #[must_use]
  // pub fn from_span(span: opentelemetry::Span) -> Self {
  //   use opentelemetry::global;

  //   let mut map = HashMap::new();
  //   let cx = span.context();
  //   global::get_text_map_propagator(|propagator| {
  //     propagator.inject_context(&cx, &mut JsonRpcInjector::new(&mut map));
  //   });

  //   Self {
  //     traceparent: map.remove("traceparent"),
  //     tracestate: map.remove("tracestate"),
  //   }
  // }

  /// Inject trace context from the current tracing span.
  ///
  /// Uses the globally registered OpenTelemetry propagator to extract
  /// trace context from the current span and return it as a TraceContext.
  #[must_use]
  pub fn from_current_span() -> Self {
    use opentelemetry::global;

    let mut map = HashMap::new();
    let cx = Context::current();
    global::get_text_map_propagator(|propagator| {
      propagator.inject_context(&cx, &mut JsonRpcInjector::new(&mut map));
    });

    Self {
      traceparent: map.remove("traceparent"),
      tracestate:  map.remove("tracestate"),
    }
  }

  // /// Extract trace context and set it as parent on the current span.
  // ///
  // /// Uses the globally registered OpenTelemetry propagator to parse
  // /// the trace context and set it as the parent of the current span.
  // pub fn apply_to_current_span(&self) {
  //   let _ = self.apply(tracing::Span::current());
  // }

  /// Returns true if both traceparent and tracestate are None.
  #[must_use]
  pub fn is_empty(&self) -> bool {
    self.traceparent.is_none() && self.tracestate.is_none()
  }

  pub fn attach(&self) -> opentelemetry::ContextGuard {
    use opentelemetry::global;

    let mut map = HashMap::new();
    if let Some(tp) = &self.traceparent {
      map.insert("traceparent".to_string(), tp.clone());
    }
    if let Some(ts) = &self.tracestate {
      map.insert("tracestate".to_string(), ts.clone());
    }

    let parent_cx = global::get_text_map_propagator(|propagator| {
      propagator.extract(&JsonRpcExtractor::new(&map))
    });

    parent_cx.attach()
  }
}

/// Wrapper for HashMap that implements OpenTelemetry Injector trait.
///
/// Used to inject trace context into JSON-RPC messages.
pub struct JsonRpcInjector<'a> {
  map: &'a mut HashMap<String, String>,
}

impl<'a> JsonRpcInjector<'a> {
  pub fn new(map: &'a mut HashMap<String, String>) -> Self {
    Self { map }
  }
}

impl Injector for JsonRpcInjector<'_> {
  fn set(&mut self, key: &str, value: String) {
    self.map.insert(key.to_string(), value);
  }
}

/// Wrapper for HashMap that implements OpenTelemetry Extractor trait.
///
/// Used to extract trace context from JSON-RPC messages.
pub struct JsonRpcExtractor<'a> {
  map: &'a HashMap<String, String>,
}

impl<'a> JsonRpcExtractor<'a> {
  pub fn new(map: &'a HashMap<String, String>) -> Self {
    Self { map }
  }
}

impl Extractor for JsonRpcExtractor<'_> {
  fn get(&self, key: &str) -> Option<&str> {
    self.map.get(key).map(|s| s.as_str())
  }

  fn keys(&self) -> Vec<&str> {
    self.map.keys().map(|s| s.as_str()).collect()
  }
}

#[cfg(test)]
mod tests {
  use super::*;

  #[test]
  fn trace_context_default_is_empty() {
    let ctx = TraceContext::default();
    assert!(ctx.is_empty());
    assert_eq!(ctx.traceparent, None);
    assert_eq!(ctx.tracestate, None);
  }

  #[test]
  fn trace_context_with_traceparent_is_not_empty() {
    let ctx = TraceContext::new(
      Some(
        "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01".to_string(),
      ),
      None,
    );
    assert!(!ctx.is_empty());
  }

  #[test]
  fn trace_context_with_tracestate_is_not_empty() {
    let ctx =
      TraceContext::new(None, Some("rojo=00f067aa0ba902b7".to_string()));
    assert!(!ctx.is_empty());
  }

  #[test]
  fn trace_context_with_both_is_not_empty() {
    let ctx = TraceContext::new(
      Some(
        "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01".to_string(),
      ),
      Some("rojo=00f067aa0ba902b7,congo=t61rcWkgMzE".to_string()),
    );
    assert!(!ctx.is_empty());
    assert_eq!(
      ctx.traceparent,
      Some(
        "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01".to_string()
      )
    );
    assert_eq!(
      ctx.tracestate,
      Some("rojo=00f067aa0ba902b7,congo=t61rcWkgMzE".to_string())
    );
  }
}