[ENH](storage): Decouple latency stopwatch from tracing

rust/storage/Cargo.toml

@@ -26,8 +26,8 @@ opentelemetry = { workspace = true }
 
 chroma-config = { workspace = true }
 chroma-error = { workspace = true }
-chroma-tracing = { workspace = true }
 chroma-types = { workspace = true }
 
 [dev-dependencies]
+opentelemetry_sdk = { workspace = true }
 rand_xorshift = { workspace = true }

rust/storage/src/admissioncontrolleds3.rs

@@ -1,3 +1,4 @@
+use crate::metrics::{StopWatchUnit, Stopwatch};
 use crate::object_storage::ObjectStorage;
 use crate::StorageError;
 use crate::{
@@ -11,7 +12,6 @@ use bytes::Bytes;
 use chroma_config::registry::Registry;
 use chroma_config::Configurable;
 use chroma_error::ChromaError;
-use chroma_tracing::util::Stopwatch;
 use futures::{stream, FutureExt, StreamExt, TryStreamExt};
 use opentelemetry::{global, metrics::Counter, KeyValue};
 use std::any::Any;
@@ -977,7 +977,7 @@ impl AdmissionControlledS3Storage {
                 let _lock_held_duration = Stopwatch::new(
                     &self.metrics.nac_lock_wait_duration_us,
                     &self.metrics.hostname_attribute,
-                    chroma_tracing::util::StopWatchUnit::Micros,
+                    StopWatchUnit::Micros,
                 );
                 let mut requests = match self.outstanding_read_requests.lock() {
                     Ok(requests) => requests,
@@ -1444,7 +1444,7 @@ impl CountBasedPolicy {
         let _stopwatch = Stopwatch::new(
             &self.metrics.nac_delay_secs,
             &priority_and_hostname_attr,
-            chroma_tracing::util::StopWatchUnit::Seconds,
+            StopWatchUnit::Seconds,
         );
         loop {
             let current_priority = priority.get_priority();

rust/storage/src/metrics.rs

@@ -1,6 +1,150 @@
 //! Metrics for storage operations.
 
-use opentelemetry::metrics::{Counter, Histogram};
+use std::time::{Duration, Instant};
+
+use opentelemetry::{
+    metrics::{Counter, Histogram},
+    KeyValue,
+};
+
+pub(crate) enum StopWatchUnit {
+    Micros,
+    Millis,
+    Seconds,
+}
+
+pub(crate) struct Stopwatch<'a> {
+    histogram: &'a Histogram<u64>,
+    attributes: &'a [KeyValue],
+    start: Instant,
+    unit: StopWatchUnit,
+    finished: bool,
+}
+
+impl<'a> Stopwatch<'a> {
+    pub(crate) fn new(
+        histogram: &'a Histogram<u64>,
+        attributes: &'a [KeyValue],
+        unit: StopWatchUnit,
+    ) -> Self {
+        Self {
+            histogram,
+            attributes,
+            start: Instant::now(),
+            unit,
+            finished: false,
+        }
+    }
+
+    pub(crate) fn finish(mut self) -> Duration {
+        let duration = self.start.elapsed();
+        self.record(duration);
+        self.finished = true;
+        duration
+    }
+
+    fn record(&self, duration: Duration) {
+        let elapsed = match self.unit {
+            StopWatchUnit::Micros => duration.as_micros() as u64,
+            StopWatchUnit::Millis => duration.as_millis() as u64,
+            StopWatchUnit::Seconds => duration.as_secs(),
+        };
+        self.histogram.record(elapsed, self.attributes);
+    }
+}
+
+impl Drop for Stopwatch<'_> {
+    fn drop(&mut self) {
+        if !self.finished {
+            self.record(self.start.elapsed());
+        }
+    }
+}
+
+#[cfg(test)]
+mod stopwatch_tests {
+    use std::sync::{Arc, Weak};
+
+    use opentelemetry::metrics::MeterProvider;
+    use opentelemetry_sdk::{
+        metrics::{
+            data::{Histogram, ResourceMetrics},
+            reader::MetricReader,
+            InstrumentKind, ManualReader, MetricResult, Pipeline, SdkMeterProvider, Temporality,
+        },
+        Resource,
+    };
+
+    use super::{StopWatchUnit, Stopwatch};
+
+    #[derive(Clone, Debug)]
+    struct SharedReader(Arc<dyn MetricReader>);
+
+    impl MetricReader for SharedReader {
+        fn register_pipeline(&self, pipeline: Weak<Pipeline>) {
+            self.0.register_pipeline(pipeline);
+        }
+
+        fn collect(&self, metrics: &mut ResourceMetrics) -> MetricResult<()> {
+            self.0.collect(metrics)
+        }
+
+        fn force_flush(&self) -> MetricResult<()> {
+            self.0.force_flush()
+        }
+
+        fn shutdown(&self) -> MetricResult<()> {
+            self.0.shutdown()
+        }
+
+        fn temporality(&self, kind: InstrumentKind) -> Temporality {
+            self.0.temporality(kind)
+        }
+    }
+
+    fn recorded_count(record: impl FnOnce(&opentelemetry::metrics::Histogram<u64>)) -> u64 {
+        let reader = SharedReader(Arc::new(ManualReader::default()));
+        let provider = SdkMeterProvider::builder()
+            .with_reader(reader.clone())
+            .build();
+        let histogram = provider
+            .meter("stopwatch-test")
+            .u64_histogram("operation-duration")
+            .build();
+
+        record(&histogram);
+        let mut exported = ResourceMetrics {
+            resource: Resource::empty(),
+            scope_metrics: Vec::new(),
+        };
+        reader
+            .collect(&mut exported)
+            .expect("metrics should collect");
+
+        let histogram = exported.scope_metrics[0].metrics[0]
+            .data
+            .as_any()
+            .downcast_ref::<Histogram<u64>>()
+            .expect("metric should be a u64 histogram");
+        histogram.data_points[0].count
+    }
+
+    #[test]
+    fn finish_records_the_duration_once() {
+        let count = recorded_count(|histogram| {
+            Stopwatch::new(histogram, &[], StopWatchUnit::Millis).finish();
+        });
+        assert_eq!(count, 1);
+    }
+
+    #[test]
+    fn drop_records_the_duration_once() {
+        let count = recorded_count(|histogram| {
+            let _stopwatch = Stopwatch::new(histogram, &[], StopWatchUnit::Millis);
+        });
+        assert_eq!(count, 1);
+    }
+}
 
 /// Metrics for tracking S3 and object storage operations.
 ///

rust/storage/src/object_storage.rs

@@ -5,11 +5,11 @@
 //! The `UpdateVersion` struct is serialized to JSON and stored as an ETag string.
 use std::{error::Error as StdError, sync::Arc, time::Duration};
 
+use crate::metrics::{StopWatchUnit, Stopwatch};
 use async_trait::async_trait;
 use bytes::{Bytes, BytesMut};
 use chroma_config::{registry::Registry, Configurable};
 use chroma_error::{source_chain_contains, ChromaError};
-use chroma_tracing::util::Stopwatch;
 use chroma_types::Cmek;
 use futures::stream::{self, StreamExt, TryStreamExt};
 use object_store::{
@@ -417,11 +417,7 @@ impl ObjectStorage {
             .map(|(start, end)| start..end);
 
         let mut buffer = BytesMut::zeroed(object_size as usize);
-        let stopwatch = Stopwatch::new(
-            &self.metrics.s3_get_latency_ms,
-            &[],
-            chroma_tracing::util::StopWatchUnit::Millis,
-        );
+        let stopwatch = Stopwatch::new(&self.metrics.s3_get_latency_ms, &[], StopWatchUnit::Millis);
         let get_part_futures = buffer
             .chunks_mut(self.download_part_size_bytes as usize)
             .zip(chunk_ranges)
@@ -465,11 +461,8 @@ impl ObjectStorage {
 
     async fn oneshot_get(&self, key: &str) -> Result<(Bytes, ETag), StorageError> {
         self.metrics.s3_get_count.add(1, &[]);
-        let _stopwatch = Stopwatch::new(
-            &self.metrics.s3_get_latency_ms,
-            &[],
-            chroma_tracing::util::StopWatchUnit::Millis,
-        );
+        let _stopwatch =
+            Stopwatch::new(&self.metrics.s3_get_latency_ms, &[], StopWatchUnit::Millis);
         let result = self.store.get_opts(&key.into(), Default::default()).await?;
         let update_version = UpdateVersion {
             e_tag: result.meta.e_tag.clone(),
@@ -503,11 +496,7 @@ impl ObjectStorage {
         let total_size_bytes = bytes.len() as u64;
         self.metrics.s3_put_count.add(1, &[]);
         self.metrics.s3_put_bytes.record(total_size_bytes, &[]);
-        let stopwatch = Stopwatch::new(
-            &self.metrics.s3_put_latency_ms,
-            &[],
-            chroma_tracing::util::StopWatchUnit::Millis,
-        );
+        let stopwatch = Stopwatch::new(&self.metrics.s3_put_latency_ms, &[], StopWatchUnit::Millis);
         let chunk_ranges = Self::partition(bytes.len() as u64, self.upload_part_size_bytes)
             .map(|(start, end)| start as usize..end as usize);
         let mut upload_handle = self
@@ -570,11 +559,7 @@ impl ObjectStorage {
         let total_size_bytes = bytes.len() as u64;
         self.metrics.s3_put_count.add(1, &[]);
         self.metrics.s3_put_bytes.record(total_size_bytes, &[]);
-        let stopwatch = Stopwatch::new(
-            &self.metrics.s3_put_latency_ms,
-            &[],
-            chroma_tracing::util::StopWatchUnit::Millis,
-        );
+        let stopwatch = Stopwatch::new(&self.metrics.s3_put_latency_ms, &[], StopWatchUnit::Millis);
 
         let result = self
             .store
@@ -678,7 +663,7 @@ impl ObjectStorage {
         let _stopwatch = Stopwatch::new(
             &self.metrics.s3_rename_latency_ms,
             &[],
-            chroma_tracing::util::StopWatchUnit::Millis,
+            StopWatchUnit::Millis,
         );
 
         self.store.rename(&src_key.into(), &dst_key.into()).await?;
@@ -687,11 +672,8 @@ impl ObjectStorage {
 
     pub async fn copy(&self, src_key: &str, dst_key: &str) -> Result<(), StorageError> {
         self.metrics.s3_copy_count.add(1, &[]);
-        let _stopwatch = Stopwatch::new(
-            &self.metrics.s3_copy_latency_ms,
-            &[],
-            chroma_tracing::util::StopWatchUnit::Millis,
-        );
+        let _stopwatch =
+            Stopwatch::new(&self.metrics.s3_copy_latency_ms, &[], StopWatchUnit::Millis);
         self.store.copy(&src_key.into(), &dst_key.into()).await?;
         Ok(())
     }
@@ -704,11 +686,8 @@ impl ObjectStorage {
         };
 
         self.metrics.s3_list_count.add(1, &[]);
-        let _stopwatch = Stopwatch::new(
-            &self.metrics.s3_list_latency_ms,
-            &[],
-            chroma_tracing::util::StopWatchUnit::Millis,
-        );
+        let _stopwatch =
+            Stopwatch::new(&self.metrics.s3_list_latency_ms, &[], StopWatchUnit::Millis);
 
         let list_stream = self.store.list(prefix_path.as_ref());