laburnum 1.17.0

An LSP framework for building language servers and compilers, powered by an incremental query tree with content-addressed storage, task-based dataflow, and parallel queries.
Documentation
// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

use {
  dashmap::DashSet,
  opentelemetry::{
    KeyValue,
    trace::TraceId,
  },
  opentelemetry_sdk::{
    error::OTelSdkResult,
    trace::{
      SpanData,
      SpanExporter,
    },
  },
  std::fmt::Debug,
};

pub struct PrefixingExporter<E> {
  inner:          E,
  prefix:         String,
  seen_trace_ids: DashSet<TraceId>,
  debug:          bool,
}

impl<E> PrefixingExporter<E> {
  pub fn new(prefix: impl Into<String>, inner: E) -> Self {
    Self {
      inner,
      prefix: prefix.into(),
      seen_trace_ids: DashSet::new(),
      debug: std::env::var_os("LABURNUM_OTEL_DEBUG").is_some(),
    }
  }

  fn should_skip_prefix(key: &str) -> bool {
    const SKIP_PREFIXES: &[&str] = &[
      "otel.",
      "telemetry.",
      "service.",
      "host.",
      "process.",
      "container.",
      "k8s.",
      "cloud.",
      "deployment.",
      "device.",
      "os.",
      "browser.",
      "webengine.",
      "http.",
      "url.",
      "server.",
      "client.",
      "network.",
      "db.",
      "rpc.",
      "messaging.",
      "faas.",
      "exception.",
      "thread.",
      "code.",
      "enduser.",
      "peer.",
      "net.",
      "gen_ai.",
      "aws.",
      "gcp.",
      "azure.",
    ];

    const SKIP_EXACT: &[&str] = &[
      "level",
      "target",
      "file",
      "line",
      "module_path",
      "busy_ns",
      "idle_ns",
    ];

    SKIP_PREFIXES.iter().any(|p| key.starts_with(p))
      || SKIP_EXACT.contains(&key)
  }

  fn prefix_attributes(&self, mut span: SpanData) -> SpanData {
    let prefixed: Vec<KeyValue> = span
      .attributes
      .drain(..)
      .map(|kv| {
        let key_str = kv.key.as_str();
        if Self::should_skip_prefix(key_str) {
          kv
        } else {
          KeyValue::new(format!("{}.{}", self.prefix, key_str), kv.value)
        }
      })
      .collect();
    span.attributes = prefixed;
    span
  }
}

impl<E: Debug> Debug for PrefixingExporter<E> {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    f.debug_struct("PrefixingExporter")
      .field("prefix", &self.prefix)
      .field("inner", &self.inner)
      .finish()
  }
}

impl<E: SpanExporter + Send + 'static> SpanExporter for PrefixingExporter<E> {
  async fn export(&self, batch: Vec<SpanData>) -> OTelSdkResult {
    if self.debug {
      eprintln!("Exporting batch of {} spans", batch.len());
      for span in &batch {
        let trace_id = span.span_context.trace_id();
        eprintln!(
          "span_name={}, trace_id={}, span_id={}",
          span.name,
          trace_id,
          span.span_context.span_id()
        );
        if self.seen_trace_ids.insert(trace_id) {
          eprintln!("trace_id: {}", trace_id);
        }
      }
    }

    let prefixed: Vec<SpanData> = batch
      .into_iter()
      .map(|span| self.prefix_attributes(span))
      .collect();

    self.inner.export(prefixed).await
  }

  fn shutdown(&mut self) -> OTelSdkResult {
    self.inner.shutdown()
  }

  fn force_flush(&mut self) -> OTelSdkResult {
    self.inner.force_flush()
  }
}