Rebase

Author: sanketkedia

rust/frontend/src/executor/distributed.rs

@@ -62,6 +62,7 @@ struct Metrics {
     delete_retries_counter: Counter<u64>,
     count_retries_counter: Counter<u64>,
     query_retries_counter: Counter<u64>,
+    search_retries_counter: Counter<u64>,
     get_retries_counter: Counter<u64>,
     add_retries_counter: Counter<u64>,
     update_retries_counter: Counter<u64>,
@@ -122,6 +123,7 @@ impl Configurable<(config::DistributedExecutorConfig, System)> for DistributedEx
         let delete_retries_counter = meter.u64_counter("delete_retries").build();
         let count_retries_counter = meter.u64_counter("count_retries").build();
         let query_retries_counter = meter.u64_counter("query_retries").build();
+        let search_retries_counter = meter.u64_counter("search_retries").build();
         let get_retries_counter = meter.u64_counter("get_retries").build();
         let add_retries_counter = meter.u64_counter("add_retries").build();
         let update_retries_counter = meter.u64_counter("update_retries").build();
@@ -135,6 +137,7 @@ impl Configurable<(config::DistributedExecutorConfig, System)> for DistributedEx
             add_retries_counter,
             update_retries_counter,
             upsert_retries_counter,
+            search_retries_counter,
         });
         Ok(Self {
             client_assigner,
@@ -395,7 +398,16 @@ impl DistributedExecutor {
         Ok(res.into_inner().try_into()?)
     }
 
-    pub async fn search(&mut self, plan: Search) -> Result<SearchResult, ExecutorError> {
+    pub async fn search<F, Fut>(
+        &mut self,
+        plan: Search,
+        replan_closure: F,
+    ) -> Result<SearchResult, ExecutorError>
+    where
+        F: Fn(tonic::Code) -> Fut,
+        Fut: Future<Output = Result<Search, Box<dyn ChromaError>>>,
+    {
+        let retry_count = Arc::new(AtomicUsize::new(0));
         // Get the collection ID from the plan
         let collection_id = &plan
             .scan
@@ -413,21 +425,57 @@ impl DistributedExecutor {
         let request: chroma_types::chroma_proto::SearchPlan = plan.try_into()?;
 
         let attempt_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
+        let last_error = std::sync::Arc::new(parking_lot::Mutex::new(tonic::Code::Ok));
         let config = self.client_selection_config.clone();
         let res = {
             let attempt_count = attempt_count.clone();
             (|| async {
                 let current_attempt =
                     attempt_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                 let is_retry = current_attempt > 0;
+                if is_retry {
+                    let last_error_code = *last_error.lock();
+                    let replan_search =
+                        replan_closure(last_error_code)
+                            .await
+                            .map_err(|e| -> tonic::Status {
+                                Status::new(
+                                    e.code().into(),
+                                    format!("Failed to replan search {:?}", e),
+                                )
+                            })?;
+                    return choose_query_client_weighted(&clients, &config, is_retry)?
+                        .search(Request::new(replan_search.try_into()?))
+                        .await;
+                }
                 choose_query_client_weighted(&clients, &config, is_retry)?
                     .search(Request::new(request.clone()))
                     .await
             })
             .retry(self.backoff)
             .when(is_retryable_error)
+            .notify(|e, _| {
+                let mut last_error = last_error.lock();
+                *last_error = e.code();
+                tracing::info!(
+                    "Retrying search for collection {}, error {:?}",
+                    collection_id,
+                    e
+                );
+                retry_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+            })
+            .adjust(|e, d| {
+                if e.code() == Code::NotFound {
+                    return Some(Duration::from_micros(0));
+                }
+                d
+            })
             .await?
         };
+        self.metrics.search_retries_counter.add(
+            retry_count.load(std::sync::atomic::Ordering::Relaxed) as u64,
+            &[],
+        );
         Ok(res.into_inner().try_into()?)
     }
 

rust/frontend/src/executor/mod.rs

@@ -71,9 +71,19 @@ impl Executor {
             Executor::Local(local_executor) => local_executor.knn(plan, replan_closure).await,
         }
     }
-    pub async fn search(&mut self, plan: Search) -> Result<SearchResult, ExecutorError> {
+    pub async fn search<F, Fut>(
+        &mut self,
+        plan: Search,
+        replan_closure: F,
+    ) -> Result<SearchResult, ExecutorError>
+    where
+        F: Fn(tonic::Code) -> Fut,
+        Fut: Future<Output = Result<Search, Box<dyn ChromaError>>>,
+    {
         match self {
-            Executor::Distributed(distributed_executor) => distributed_executor.search(plan).await,
+            Executor::Distributed(distributed_executor) => {
+                distributed_executor.search(plan, replan_closure).await
+            }
             Executor::Local(local_executor) => local_executor.search(plan).await,
         }
     }

rust/frontend/src/impls/service_based_frontend.rs

@@ -75,7 +75,6 @@ struct Metrics {
     create_tenant_retries_counter: Counter<u64>,
     update_tenant_retries_counter: Counter<u64>,
     get_collection_with_segments_counter: Counter<u64>,
-    search_retries_counter: Counter<u64>,
     metering_fork_counter: Counter<u64>,
     metering_read_counter: Counter<u64>,
     metering_write_counter: Counter<u64>,
@@ -113,7 +112,6 @@ impl ServiceBasedFrontend {
         let add_retries_counter = meter.u64_counter("add_retries").build();
         let update_retries_counter = meter.u64_counter("update_retries").build();
         let upsert_retries_counter = meter.u64_counter("upsert_retries").build();
-        let search_retries_counter = meter.u64_counter("search_retries").build();
         let metering_fork_counter = meter.u64_counter("metering_events_sent.fork").with_description("The number of fork metering events sent by the frontend to the metering event receiver.").build();
         let metering_read_counter = meter.u64_counter("metering_events_sent.read").with_description("The number of read metering events sent by the frontend to the metering event receiver.").build();
         let metering_write_counter = meter.u64_counter("metering_events_sent.write").with_description("The number of write metering events sent by the frontend to the metering event receiver.").build();
@@ -162,7 +160,6 @@ impl ServiceBasedFrontend {
             create_db_retries_counter,
             delete_db_retries_counter,
             delete_collection_retries_counter,
-            search_retries_counter,
             metering_fork_counter,
             metering_read_counter,
             metering_write_counter,
@@ -1577,7 +1574,9 @@ impl ServiceBasedFrontend {
         };
 
         if let Some(event) = read_event {
-            event.submit().await;
+            if let Ok(()) = event.submit().await {
+                self.metrics.metering_read_counter.add(1, &[]);
+            }
         }
 
         Ok(records)
@@ -1996,10 +1995,8 @@ impl ServiceBasedFrontend {
         // TODO: The dispatch logic is mostly the same for count/get/query/search, we should consider unifying them
         // Get collection and segments once for all queries
         let collection_and_segments = self
-            .collections_with_segments_provider
-            .get_collection_with_segments(request.collection_id)
-            .await
-            .map_err(|err| QueryError::Other(Box::new(err) as Box<dyn ChromaError>))?;
+            .retryable_get_collection_with_segments(request.collection_id)
+            .await?;
 
         let latest_collection_logical_size_bytes = collection_and_segments
             .collection
@@ -2032,8 +2029,36 @@ impl ServiceBasedFrontend {
             payloads: request.searches,
         };
 
+        let collection_id = search_plan
+            .scan
+            .collection_and_segments
+            .collection
+            .collection_id;
+
         // Execute the single search plan using the executor
-        let result = self.executor.search(search_plan).await?;
+        let result = self
+            .executor
+            .search(search_plan.clone(), |code: tonic::Code| {
+                let mut provider = self.collections_with_segments_provider.clone();
+                let mut search_replanned = search_plan.clone();
+                async move {
+                    if code == tonic::Code::NotFound {
+                        provider
+                            .collections_with_segments_cache
+                            .remove(&collection_id)
+                            .await;
+                        let collection_and_segments = provider
+                            .get_collection_with_segments(collection_id)
+                            .await
+                            .map_err(|err| Box::new(err) as Box<dyn ChromaError>)?;
+                        search_replanned.scan = Scan {
+                            collection_and_segments,
+                        };
+                    }
+                    Ok(search_replanned)
+                }
+            })
+            .await?;
 
         // Calculate return bytes (approximate size of the response)
         let return_bytes = result.size_bytes();
@@ -2079,50 +2104,7 @@ impl ServiceBasedFrontend {
     }
 
     pub async fn search(&mut self, request: SearchRequest) -> Result<SearchResponse, QueryError> {
-        // TODO: The retry logic is mostly the same for count/get/query/search, we should consider unifying them
-        let retries = Arc::new(AtomicUsize::new(0));
-        let search_to_retry = || {
-            let mut self_clone = self.clone();
-            let request_clone = request.clone();
-            let cache_clone = self
-                .collections_with_segments_provider
-                .collections_with_segments_cache
-                .clone();
-            async move {
-                let res = self_clone.retryable_search(request_clone).await;
-                match res {
-                    Ok(res) => Ok(res),
-                    Err(e) => {
-                        if e.code() == ErrorCodes::NotFound {
-                            tracing::info!(
-                                "Invalidating cache for collection {}",
-                                request.collection_id
-                            );
-                            cache_clone.remove(&request.collection_id).await;
-                        }
-                        Err(e)
-                    }
-                }
-            }
-        };
-        let res = search_to_retry
-            .retry(self.collections_with_segments_provider.get_retry_backoff())
-            // NOTE: Transport level errors will manifest as unknown errors, and they should also be retried
-            .when(|e| matches!(e.code(), ErrorCodes::NotFound | ErrorCodes::Unknown))
-            .notify(|_, _| {
-                let retried = retries.fetch_add(1, Ordering::Relaxed);
-                if retried > 0 {
-                    tracing::info!(
-                        "Retrying search() request for collection {}",
-                        request.collection_id
-                    );
-                }
-            })
-            .await;
-        self.metrics
-            .search_retries_counter
-            .add(retries.load(Ordering::Relaxed) as u64, &[]);
-        res
+        self.retryable_search(request).await
     }
 
     pub async fn healthcheck(&self) -> HealthCheckResponse {