[ENH]: add metrics for span & metric exporters

Author: codetheweb

Cargo.lock

@@ -18,6 +18,9 @@ tracing-opentelemetry = { workspace = true }
 tracing-subscriber = { workspace = true }
 tracing = { workspace = true }
 uuid = { workspace = true }
+futures = { workspace = true }
+async-trait = { workspace = true }
+tokio = { workspace = true }
 
 chroma-system = { workspace = true }
 

rust/tracing/Cargo.toml

@@ -2,6 +2,8 @@
 // load/src/opentelemetry_config.rs file
 // Keep them in-sync manually.
 
+use crate::wrapped_metric_exporter::WrappedMetricExporter;
+use crate::wrapped_span_exporter::WrappedSpanExporter;
 use opentelemetry::trace::TracerProvider;
 use opentelemetry::{global, InstrumentationScope};
 use opentelemetry_otlp::WithExportConfig;
@@ -107,22 +109,25 @@ pub fn init_otel_layer(
     ]);
 
     // Prepare tracer.
-    let tracing_span_exporter = opentelemetry_otlp::SpanExporter::builder()
+    let tracing_span_exporter: WrappedSpanExporter<_> = opentelemetry_otlp::SpanExporter::builder()
         .with_tonic()
         .with_endpoint(otel_endpoint)
         .build()
-        .expect("could not build span exporter for tracing");
+        .expect("could not build span exporter for tracing")
+        .into();
     let trace_config = opentelemetry_sdk::trace::Config::default().with_resource(resource.clone());
     let tracer_provider = opentelemetry_sdk::trace::TracerProvider::builder()
         .with_batch_exporter(tracing_span_exporter, opentelemetry_sdk::runtime::Tokio)
         .with_config(trace_config)
         .build();
     let tracer = tracer_provider.tracer(service_name.clone());
-    let fastrace_span_exporter = opentelemetry_otlp::SpanExporter::builder()
-        .with_tonic()
-        .with_endpoint(otel_endpoint)
-        .build()
-        .expect("could not build span exporter for fastrace");
+    let fastrace_span_exporter: WrappedSpanExporter<_> =
+        opentelemetry_otlp::SpanExporter::builder()
+            .with_tonic()
+            .with_endpoint(otel_endpoint)
+            .build()
+            .expect("could not build span exporter for fastrace")
+            .into();
     fastrace::set_reporter(
         fastrace_opentelemetry::OpenTelemetryReporter::new(
             fastrace_span_exporter,
@@ -134,13 +139,14 @@ pub fn init_otel_layer(
     );
 
     // Prepare meter.
-    let metric_exporter = opentelemetry_otlp::MetricExporter::builder()
+    let metric_exporter: WrappedMetricExporter<_> = opentelemetry_otlp::MetricExporter::builder()
         .with_tonic()
         .with_endpoint(
             std::env::var("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT").unwrap_or(otel_endpoint.clone()),
         )
         .build()
-        .expect("could not build metric exporter");
+        .expect("could not build metric exporter")
+        .into();
 
     let reader = opentelemetry_sdk::metrics::PeriodicReader::builder(
         metric_exporter,

rust/tracing/src/init_tracer.rs

@@ -2,6 +2,8 @@
 pub mod grpc_tower;
 pub mod init_tracer;
 pub mod util;
+mod wrapped_metric_exporter;
+mod wrapped_span_exporter;
 
 #[cfg(feature = "grpc")]
 pub use grpc_tower::*;

rust/tracing/src/lib.rs

@@ -0,0 +1,81 @@
+use opentelemetry::{global, metrics::Counter, KeyValue};
+use opentelemetry_sdk::metrics::{
+    data::ResourceMetrics, exporter::PushMetricExporter, MetricResult,
+};
+use std::sync::LazyLock;
+
+/// A wrapper around an OpenTelemetry MetricExporter that counts the number of export calls & tracks errors.
+#[derive(Debug)]
+pub struct WrappedMetricExporter<E> {
+    inner: E,
+    counter: LazyLock<Counter<u64>>,
+}
+
+impl<E> WrappedMetricExporter<E> {
+    pub fn new(inner: E) -> Self {
+        Self {
+            inner,
+            // We use a LazyLock so that the counter is only initialized after the global exporter is set up.
+            counter: LazyLock::new(|| {
+                global::meter("chroma.tracing")
+                    .u64_counter("metric_exporter_calls")
+                    .with_description("Counts the number of metric exporter calls")
+                    .build()
+            }),
+        }
+    }
+}
+
+impl From<opentelemetry_otlp::MetricExporter>
+    for WrappedMetricExporter<opentelemetry_otlp::MetricExporter>
+{
+    fn from(exporter: opentelemetry_otlp::MetricExporter) -> Self {
+        WrappedMetricExporter::new(exporter)
+    }
+}
+
+#[async_trait::async_trait]
+impl<E> PushMetricExporter for WrappedMetricExporter<E>
+where
+    E: PushMetricExporter + Send + Sync + 'static,
+{
+    async fn export(&self, metrics: &mut ResourceMetrics) -> MetricResult<()> {
+        let counter = self.counter.clone();
+        println!("Exporting metrics: {:?}", metrics);
+        let result = self.inner.export(metrics).await;
+        match &result {
+            Ok(_) => counter.add(1, &[KeyValue::new("status", "success")]),
+            Err(err) => {
+                let error_name = match err {
+                    opentelemetry_sdk::metrics::MetricError::Other(_) => "other",
+                    opentelemetry_sdk::metrics::MetricError::Config(_) => "config",
+                    opentelemetry_sdk::metrics::MetricError::ExportErr(_) => "export_failed",
+                    opentelemetry_sdk::metrics::MetricError::InvalidInstrumentConfiguration(_) => {
+                        "invalid_instrument_configuration"
+                    }
+                    _ => "unknown",
+                };
+                counter.add(
+                    1,
+                    &[
+                        KeyValue::new("status", "error"),
+                        KeyValue::new("error", error_name.to_string()),
+                    ],
+                );
+            }
+        }
+        result
+    }
+
+    async fn force_flush(&self) -> MetricResult<()> {
+        self.inner.force_flush().await
+    }
+
+    fn shutdown(&self) -> MetricResult<()> {
+        self.inner.shutdown()
+    }
+
+    fn temporality(&self) -> opentelemetry_sdk::metrics::Temporality {
+        self.inner.temporality()
+    }
+}

rust/tracing/src/wrapped_metric_exporter.rs

@@ -0,0 +1,81 @@
+use futures::future::BoxFuture;
+use opentelemetry::{global, metrics::Counter, KeyValue};
+use opentelemetry_sdk::export::trace::SpanExporter;
+use std::sync::{Arc, LazyLock};
+use tokio::sync::Mutex;
+
+/// A wrapper around an OpenTelemetry SpanExporter that counts the number of export calls & tracks errors.
+#[derive(Debug)]
+pub struct WrappedSpanExporter<E> {
+    // The documentation for the SpanExporter trait specifies that "This function will never be called concurrently for the same exporter instance" so using a Mutex should not limit performance.
+    inner: Arc<Mutex<E>>,
+    // We use a LazyLock so that the counter is only initialized after the global exporter is set up.
+    counter: LazyLock<Counter<u64>>,
+}
+
+impl From<opentelemetry_otlp::SpanExporter>
+    for WrappedSpanExporter<opentelemetry_otlp::SpanExporter>
+{
+    fn from(exporter: opentelemetry_otlp::SpanExporter) -> Self {
+        WrappedSpanExporter::new(exporter)
+    }
+}
+
+impl<E> WrappedSpanExporter<E> {
+    pub fn new(inner: E) -> Self {
+        Self {
+            inner: Arc::new(Mutex::new(inner)),
+            counter: LazyLock::new(|| {
+                global::meter("chroma.tracing")
+                    .u64_counter("span_exporter_calls")
+                    .with_description("Counts the number of span exporter calls")
+                    .build()
+            }),
+        }
+    }
+}
+
+impl<E> SpanExporter for WrappedSpanExporter<E>
+where
+    E: SpanExporter + Send + Sync + 'static,
+{
+    fn export(
+        &mut self,
+        batch: Vec<opentelemetry_sdk::export::trace::SpanData>,
+    ) -> BoxFuture<'static, opentelemetry_sdk::export::trace::ExportResult> {
+        let exporter = self.inner.clone();
+        let counter = self.counter.clone();
+        Box::pin(async move {
+            let mut exporter = exporter.lock().await;
+            let result = exporter.export(batch).await;
+            match &result {
+                Ok(_) => counter.add(1, &[KeyValue::new("status", "success")]),
+                Err(err) => {
+                    println!("caught error in span exporter: {:?}", err);
+                    let error_name = match err {
+                        opentelemetry::trace::TraceError::ExportFailed(_) => "export_failed",
+                        opentelemetry::trace::TraceError::ExportTimedOut(_) => "timeout",
+                        opentelemetry::trace::TraceError::TracerProviderAlreadyShutdown => {
+                            "shutdown"
+                        }
+                        opentelemetry::trace::TraceError::Other(_) => "other",
+                        _ => "unknown",
+                    };
+                    counter.add(
+                        1,
+                        &[
+                            KeyValue::new("status", "error"),
+                            KeyValue::new("error", error_name.to_string()),
+                        ],
+                    );
+                }
+            }
+
+            result
+        })
+    }
+
+    fn shutdown(&mut self) {
+        self.inner.blocking_lock().shutdown()
+    }
+}