added metrics to block manager

Author: Tanuj Nayak

Cargo.lock

@@ -23,12 +23,14 @@ tokio = { workspace = true }
 num_cpus = { workspace = true }
 flatbuffers = { workspace = true }
 itertools = { workspace = true }
+opentelemetry = { workspace = true }
 
 chroma-error = { workspace = true }
 chroma-config = { workspace = true }
 chroma-storage = { workspace = true }
 chroma-cache = { workspace = true }
 chroma-types = { workspace = true }
+chroma-tracing = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true }

rust/blockstore/Cargo.toml

@@ -4,10 +4,11 @@ use super::{
     root::RootWriter,
     types::{ArrowWriteableKey, ArrowWriteableValue},
 };
-use chroma_error::ChromaError;
 use futures::{StreamExt, TryStreamExt};
 use uuid::Uuid;
 
+use chroma_error::ChromaError;
+
 pub struct ArrowBlockfileFlusher {
     block_manager: BlockManager,
     root_manager: RootManager,

rust/blockstore/src/arrow/flusher.rs

@@ -20,7 +20,9 @@ use chroma_error::{ChromaError, ErrorCodes};
 use chroma_storage::{
     admissioncontrolleds3::StorageRequestPriority, GetOptions, PutOptions, Storage,
 };
+use chroma_tracing::util::{get_current_trace_id, Stopwatch};
 use futures::{stream::FuturesUnordered, StreamExt};
+use opentelemetry::{global, KeyValue};
 use std::{
     collections::HashMap,
     sync::Arc,
@@ -263,6 +265,36 @@ impl ChromaError for ForkError {
     }
 }
 
+#[derive(Clone)]
+pub struct BlockMetrics {
+    pub commit_latency: opentelemetry::metrics::Histogram<u64>,
+    pub num_blocks_flushed: opentelemetry::metrics::Histogram<u64>,
+    pub flush_latency: opentelemetry::metrics::Histogram<u64>,
+}
+
+impl Default for BlockMetrics {
+    fn default() -> Self {
+        let meter = global::meter("chroma");
+        Self {
+            commit_latency: meter
+                .u64_histogram("block_commit_latency")
+                .with_description("Commit latency")
+                .with_unit("microseconds")
+                .build(),
+            num_blocks_flushed: meter
+                .u64_histogram("block_num_blocks_flushed")
+                .with_description("Number of blocks flushed")
+                .with_unit("blocks")
+                .build(),
+            flush_latency: meter
+                .u64_histogram("block_flush_latency")
+                .with_description("Flush latency")
+                .with_unit("microseconds")
+                .build(),
+        }
+    }
+}
+
 /// A simple local cache of Arrow-backed blocks, the blockfile provider passes this
 /// to the ArrowBlockfile when it creates a new blockfile. So that the blockfile can manage and access blocks
 /// # Note
@@ -274,6 +306,7 @@ pub(super) struct BlockManager {
     block_cache: Arc<dyn PersistentCache<Uuid, Block>>,
     storage: Storage,
     max_block_size_bytes: usize,
+    block_metrics: BlockMetrics,
 }
 
 impl BlockManager {
@@ -287,6 +320,7 @@ impl BlockManager {
             block_cache,
             storage,
             max_block_size_bytes,
+            block_metrics: BlockMetrics::default(),
         }
     }
 
@@ -317,6 +351,9 @@ impl BlockManager {
         &self,
         delta: impl Delta,
     ) -> Block {
+        let trace_id = get_current_trace_id().to_string();
+        let attribute = [KeyValue::new("trace_id", trace_id)];
+        let _stopwatch = Stopwatch::new(&self.block_metrics.commit_latency, &attribute);
         let delta_id = delta.id();
         let record_batch = delta.finish::<K, V>(None);
         let block = Block::from_record_batch(delta_id, record_batch);
@@ -387,6 +424,10 @@ impl BlockManager {
             }
         };
         let key = format!("block/{}", block.id);
+
+        let trace_id = get_current_trace_id().to_string();
+        let attribute = [KeyValue::new("trace_id", trace_id)];
+        let _stopwatch = Stopwatch::new(&self.block_metrics.flush_latency, &attribute);
         let block_bytes_len = bytes.len();
         let res = self
             .storage
@@ -403,6 +444,7 @@ impl BlockManager {
                     block.id,
                     block_bytes_len
                 );
+                self.block_metrics.num_blocks_flushed.record(1, &attribute);
             }
             Err(e) => {
                 tracing::info!("Error writing block to storage {}", e);

rust/blockstore/src/arrow/provider.rs

@@ -1,14 +1,12 @@
 use axum::extract::MatchedPath;
 use axum::http::{header, Request, Response};
 use axum::Router;
+use chroma_tracing::util::get_current_trace_id;
 use futures::future::BoxFuture;
-use opentelemetry::trace::TraceContextExt;
 use std::time::Duration;
 use tower::Service;
 use tower_http::trace::{MakeSpan, OnResponse, TraceLayer};
 
-use tracing_opentelemetry::OpenTelemetrySpanExt;
-
 #[derive(Clone)]
 struct RequestTracing;
 impl<B> MakeSpan<B> for RequestTracing {
@@ -81,12 +79,7 @@ where
         Box::pin(async move {
             let mut response: Response<Rs> = future.await?;
             if response.status().is_client_error() || response.status().is_server_error() {
-                let trace_id = tracing::Span::current()
-                    .context()
-                    .span()
-                    .span_context()
-                    .trace_id()
-                    .to_string();
+                let trace_id = get_current_trace_id().to_string();
                 let headers = response.headers_mut();
                 let header_val = trace_id.parse::<header::HeaderValue>();
                 if let Ok(val) = header_val {

rust/frontend/src/tower_tracing.rs

@@ -29,12 +29,14 @@ use bytes::Bytes;
 use chroma_config::registry::Registry;
 use chroma_config::Configurable;
 use chroma_error::ChromaError;
+use chroma_tracing::util::get_current_trace_id;
 use futures::future::BoxFuture;
 use futures::stream;
 use futures::FutureExt;
 use futures::Stream;
 use futures::StreamExt;
 use opentelemetry::metrics::Meter;
+use opentelemetry::KeyValue;
 use rand::Rng;
 use std::clone::Clone;
 use std::ops::Range;
@@ -45,15 +47,27 @@ use tracing::Instrument;
 
 #[derive(Clone, Debug)]
 pub struct S3StorageMetrics {
-    total_num_get_requests: opentelemetry::metrics::Counter<u64>,
+    num_get_requests: opentelemetry::metrics::Histogram<u64>,
+    num_put_requests: opentelemetry::metrics::Histogram<u64>,
+    num_delete_requests: opentelemetry::metrics::Histogram<u64>,
 }
 
 impl S3StorageMetrics {
     pub fn new(meter: Meter) -> Self {
         Self {
-            total_num_get_requests: meter
-                .u64_counter("s3_storage_total_num_get_requests")
-                .with_description("The number of GET requests to S3. This does not include metadata HEAD requests.")
+            num_get_requests: meter
+                .u64_histogram("s3_num_get_requests")
+                .with_description("The number of GET requests to S3.")
+                .with_unit("requests")
+                .build(),
+            num_put_requests: meter
+                .u64_histogram("s3_num_put_requests")
+                .with_description("The number of PUT requests to S3.")
+                .with_unit("requests")
+                .build(),
+            num_delete_requests: meter
+                .u64_histogram("s3_num_delete_requests")
+                .with_description("The number of DELETE requests to S3.")
                 .with_unit("requests")
                 .build(),
         }
@@ -144,7 +158,10 @@ impl S3Storage {
             .send()
             .instrument(tracing::trace_span!("cold S3 get"))
             .await;
-        self.metrics.total_num_get_requests.add(1, &[]);
+
+        let trace_id = get_current_trace_id().to_string();
+        let attribute = [KeyValue::new("trace_id", trace_id)];
+        self.metrics.num_get_requests.record(1, &attribute);
         match res {
             Ok(res) => {
                 let byte_stream = res.body;
@@ -200,6 +217,9 @@ impl S3Storage {
             .key(key)
             .send()
             .await;
+        let trace_id = get_current_trace_id().to_string();
+        let attribute = [KeyValue::new("trace_id", trace_id)];
+        self.metrics.num_get_requests.record(1, &attribute);
         let (content_length, e_tag) = match head_res {
             Ok(res) => match res.content_length {
                 Some(len) => (len, res.e_tag),
@@ -245,7 +265,9 @@ impl S3Storage {
             .send()
             .instrument(tracing::trace_span!("cold S3 get"))
             .await;
-        self.metrics.total_num_get_requests.add(1, &[]);
+        let trace_id = get_current_trace_id().to_string();
+        let attribute = [KeyValue::new("trace_id", trace_id)];
+        self.metrics.num_get_requests.record(1, &attribute);
         match res {
             Ok(output) => Ok(output),
             Err(e) => {
@@ -342,7 +364,6 @@ impl S3Storage {
         key: &str,
     ) -> Result<(Arc<Vec<u8>>, Option<ETag>), StorageError> {
         let (mut stream, e_tag) = self.get_stream_and_e_tag(key).await?;
-        let read_block_span = tracing::trace_span!("S3 read bytes to end");
         let buf = async {
             let mut buf: Vec<u8> = Vec::new();
             while let Some(res) = stream.next().await {
@@ -358,7 +379,6 @@ impl S3Storage {
             }
             Ok(Some(buf))
         }
-        .instrument(read_block_span)
         .await?;
         match buf {
             Some(buf) => Ok((Arc::new(buf), e_tag)),
@@ -444,13 +464,19 @@ impl S3Storage {
         ) -> BoxFuture<'static, Result<ByteStream, StorageError>>,
         options: PutOptions,
     ) -> Result<Option<ETag>, StorageError> {
-        if self.is_oneshot_upload(total_size_bytes) {
-            return self
-                .oneshot_upload(key, total_size_bytes, create_bytestream_fn, options)
-                .await;
-        }
-        self.multipart_upload(key, total_size_bytes, create_bytestream_fn, options)
-            .await
+        let res = {
+            if self.is_oneshot_upload(total_size_bytes) {
+                self.oneshot_upload(key, total_size_bytes, create_bytestream_fn, options)
+                    .await
+            } else {
+                self.multipart_upload(key, total_size_bytes, create_bytestream_fn, options)
+                    .await
+            }
+        };
+        let trace_id = get_current_trace_id().to_string();
+        let attribute = [KeyValue::new("trace_id", trace_id)];
+        self.metrics.num_put_requests.record(1, &attribute);
+        res
     }
 
     #[tracing::instrument(skip(self, create_bytestream_fn), level = "trace")]
@@ -492,6 +518,10 @@ impl S3Storage {
                 }
             }
         })?;
+
+        let trace_id = get_current_trace_id().to_string();
+        let attribute = [KeyValue::new("trace_id", trace_id)];
+        self.metrics.num_put_requests.record(1, &attribute);
         Ok(resp.e_tag.map(ETag))
     }
 
@@ -568,6 +598,9 @@ impl S3Storage {
                 source: Arc::new(err.into_service_error()),
             })?;
 
+        let trace_id = get_current_trace_id().to_string();
+        let attribute = [KeyValue::new("trace_id", trace_id)];
+        self.metrics.num_put_requests.record(1, &attribute);
         Ok(CompletedPart::builder()
             .e_tag(upload_part_res.e_tag.unwrap_or_default())
             .part_number(part_number)
@@ -656,7 +689,7 @@ impl S3Storage {
             None => req,
         };
 
-        match req.send().await {
+        let res = match req.send().await {
             Ok(_) => {
                 tracing::trace!(key = %key, "Successfully deleted object from S3");
                 Ok(())
@@ -667,7 +700,12 @@ impl S3Storage {
                     source: Arc::new(e.into_service_error()),
                 })
             }
-        }
+        };
+
+        let trace_id = get_current_trace_id().to_string();
+        let attribute = [KeyValue::new("trace_id", trace_id)];
+        self.metrics.num_delete_requests.record(1, &attribute);
+        res
     }
 
     #[tracing::instrument(skip(self), level = "trace")]

rust/storage/src/s3.rs

@@ -90,3 +90,8 @@ impl Drop for Stopwatch<'_> {
         self.0.record(elapsed, self.1);
     }
 }
+
+pub fn get_current_trace_id() -> TraceId {
+    let span = tracing::Span::current();
+    span.context().span().span_context().trace_id()
+}