[ENH]: add request timing to metering (#4877)

This PR uses the `chroma-metering` library introduced in
#4868 to add execution timing
to requests received on the frontend.

⚠️ #4868 should be merged first.

- [x] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes
N/A

Author: c-gamble

rust/frontend/src/impls/service_based_frontend.rs

@@ -7,7 +7,7 @@ use chroma_config::{registry, Configurable};
 use chroma_error::{ChromaError, ErrorCodes};
 use chroma_log::{LocalCompactionManager, LocalCompactionManagerConfig, Log};
 use chroma_metering::{
-    CollectionForkContext, CollectionReadContext, CollectionWriteContext, Enterable,
+    CollectionForkContext, CollectionReadContext, CollectionWriteContext, Enterable, FinishRequest,
     FtsQueryLength, LatestCollectionLogicalSizeBytes, LogSizeBytes, MetadataPredicateCount,
     MeterEvent, PulledLogSizeBytes, QueryEmbeddingCount, ReturnBytes, WriteAction,
 };
@@ -40,10 +40,13 @@ use chroma_types::{
 };
 use opentelemetry::global;
 use opentelemetry::metrics::Counter;
-use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
 use std::{collections::HashSet, time::Duration};
+use std::{
+    sync::atomic::{AtomicUsize, Ordering},
+    time::Instant,
+};
 
 use super::utils::to_records;
 
@@ -637,9 +640,9 @@ impl ServiceBasedFrontend {
             .set_collection_with_segments(collection_and_segments)
             .await;
 
-        // Attach the collection's latest logical size to the metering context
+        // Attach metadata to the metering context
         chroma_metering::with_current(|context| {
-            context.latest_collection_logical_size_bytes(latest_collection_logical_size_bytes)
+            context.latest_collection_logical_size_bytes(latest_collection_logical_size_bytes);
         });
 
         // TODO: Submit event after the response is sent
@@ -752,8 +755,11 @@ impl ServiceBasedFrontend {
             .add_retries_counter
             .add(retries.load(Ordering::Relaxed) as u64, &[]);
 
-        // Attach the log size in bytes to the metering context
-        chroma_metering::with_current(|context| context.log_size_bytes(log_size_bytes));
+        // Attach metadata to the metering context
+        chroma_metering::with_current(|context| {
+            context.log_size_bytes(log_size_bytes);
+            context.finish_request(Instant::now());
+        });
 
         // TODO: Submit event after the response is sent
         match res {
@@ -834,8 +840,11 @@ impl ServiceBasedFrontend {
             .update_retries_counter
             .add(retries.load(Ordering::Relaxed) as u64, &[]);
 
-        // Attach the log size in bytes to the metering context
-        chroma_metering::with_current(|context| context.log_size_bytes(log_size_bytes));
+        // Attach metadata to the metering context
+        chroma_metering::with_current(|context| {
+            context.log_size_bytes(log_size_bytes);
+            context.finish_request(Instant::now());
+        });
 
         // TODO: Submit event after the response is sent
         match res {
@@ -918,8 +927,11 @@ impl ServiceBasedFrontend {
             .upsert_retries_counter
             .add(retries.load(Ordering::Relaxed) as u64, &[]);
 
-        // Attach the log size in bytes to the metering context
-        chroma_metering::with_current(|context| context.log_size_bytes(log_size_bytes));
+        // Attach metadata to the metering context
+        chroma_metering::with_current(|context| {
+            context.log_size_bytes(log_size_bytes);
+            context.finish_request(Instant::now());
+        });
 
         // TODO: Submit event after the response is sent
         match res {
@@ -1041,9 +1053,10 @@ impl ServiceBasedFrontend {
             None
         };
 
-        // NOTE(c-gamble): See note above for additional context, but this is a non-standard pattern
-        // and we are only implementing metering for delete in this manner because delete currently
-        // produces two metering events.
+        if let Some(event) = read_event {
+            event.submit().await;
+        }
+
         let collection_write_context_container =
             chroma_metering::create::<CollectionWriteContext>(CollectionWriteContext::new(
                 tenant_id.clone(),
@@ -1071,12 +1084,10 @@ impl ServiceBasedFrontend {
                 }
             })?;
 
-        if let Some(event) = read_event {
-            event.submit().await;
-        }
-
-        // Attach the log size bytes to the write context
-        chroma_metering::with_current(|context| context.log_size_bytes(log_size_bytes));
+        // Attach metadata to the write context
+        chroma_metering::with_current(|context| {
+            context.log_size_bytes(log_size_bytes);
+        });
 
         // TODO: Submit event after the response is sent
         match chroma_metering::close::<CollectionWriteContext>() {
@@ -1296,6 +1307,7 @@ impl ServiceBasedFrontend {
             context.pulled_log_size_bytes(get_result.pulled_log_bytes);
             context.latest_collection_logical_size_bytes(latest_collection_logical_size_bytes);
             context.return_bytes(return_bytes);
+            context.finish_request(Instant::now());
         });
 
         // TODO: Submit event after the response is sent
@@ -1424,6 +1436,7 @@ impl ServiceBasedFrontend {
             context.pulled_log_size_bytes(query_result.pulled_log_bytes);
             context.latest_collection_logical_size_bytes(latest_collection_logical_size_bytes);
             context.return_bytes(return_bytes);
+            context.finish_request(Instant::now());
         });
 
         // TODO: Submit event after the response is sent

rust/frontend/src/server.rs

@@ -6,8 +6,8 @@ use axum::{
     Json, Router, ServiceExt,
 };
 use chroma_metering::{
-    CollectionForkContext, CollectionReadContext, CollectionWriteContext, MeteredFutureExt,
-    ReadAction, WriteAction,
+    CollectionForkContext, CollectionReadContext, CollectionWriteContext, Enterable,
+    MeteredFutureExt, ReadAction, StartRequest, WriteAction,
 };
 use chroma_system::System;
 use chroma_types::{
@@ -29,11 +29,11 @@ use opentelemetry::global;
 use opentelemetry::metrics::{Counter, Meter};
 use opentelemetry::KeyValue;
 use serde::{Deserialize, Serialize};
-use std::str::FromStr;
 use std::sync::{
     atomic::{AtomicBool, Ordering},
     Arc,
 };
+use std::{str::FromStr, time::Instant};
 use tokio::{select, signal};
 use tower_http::cors::CorsLayer;
 use utoipa::openapi::security::{ApiKey, ApiKeyValue, SecurityScheme};
@@ -1340,6 +1340,12 @@ async fn collection_add(
             WriteAction::Add,
         ));
 
+    metering_context_container.enter();
+
+    chroma_metering::with_current(|context| {
+        context.start_request(Instant::now());
+    });
+
     tracing::info!(name: "collection_add", tenant_name = %tenant, database_name = %database, collection_id = %collection_id, num_ids = %payload.ids.len());
     let request = chroma_types::AddCollectionRecordsRequest::try_new(
         tenant,
@@ -1443,6 +1449,12 @@ async fn collection_update(
             WriteAction::Update,
         ));
 
+    metering_context_container.enter();
+
+    chroma_metering::with_current(|context| {
+        context.start_request(Instant::now());
+    });
+
     tracing::info!(name: "collection_update", tenant_name = %tenant, database_name = %database, collection_id = %collection_id, num_ids = %payload.ids.len());
     let request = chroma_types::UpdateCollectionRecordsRequest::try_new(
         tenant,
@@ -1553,6 +1565,12 @@ async fn collection_upsert(
             WriteAction::Upsert,
         ));
 
+    metering_context_container.enter();
+
+    chroma_metering::with_current(|context| {
+        context.start_request(Instant::now());
+    });
+
     tracing::info!(name: "collection_upsert", tenant_name = %tenant, database_name = %database, collection_id = %collection_id, num_ids = %payload.ids.len());
     let request = chroma_types::UpsertCollectionRecordsRequest::try_new(
         tenant,
@@ -1834,6 +1852,12 @@ async fn collection_get(
             ReadAction::Get,
         ));
 
+    metering_context_container.enter();
+
+    chroma_metering::with_current(|context| {
+        context.start_request(Instant::now());
+    });
+
     tracing::info!(
         name: "collection_get",
         num_ids = payload.ids.as_ref().map_or(0, |ids| ids.len()),
@@ -1888,6 +1912,7 @@ pub struct QueryRequestPayload {
         ("offset" = Option<u32>, Query, description = "Offset for pagination")
     )
 )]
+
 async fn collection_query(
     headers: HeaderMap,
     Path((tenant, database, collection_id)): Path<(String, String, String)>,
@@ -1950,6 +1975,12 @@ async fn collection_query(
             ReadAction::Query,
         ));
 
+    metering_context_container.enter();
+
+    chroma_metering::with_current(|context| {
+        context.start_request(Instant::now());
+    });
+
     tracing::info!(
         name: "collection_query",
         num_ids = payload.ids.as_ref().map_or(0, |ids| ids.len()),

rust/metering-macros/src/lib.rs

@@ -90,6 +90,8 @@ pub fn initialize_metering(raw_token_stream: proc_macro::TokenStream) -> proc_ma
             ReceiverAlreadyInitializedError,
             #[error("Failed to downcast context to provided type")]
             DowncastError,
+            #[error("RwLock poisoned when attempting to read or write")]
+            RwLockPoisonedError,
         }
 
         /// A type alias for a shared, boxed, metering context