[ENH]: Add retry for add, update and upsert (#4457)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
  - Adds exponential retry to coll.add(), coll.update() and coll.upsert() for duplicate key errors returned by postgres
- New functionality
  - ...

## Test plan
_How are these changes tested?_
- [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Documentation Changes
None

Author: sanketkedia

go/pkg/log/repository/log.go

@@ -3,11 +3,13 @@ package repository
 import (
 	"context"
 	"errors"
+	"fmt"
 	"time"
 
 	log "github.com/chroma-core/chroma/go/pkg/log/store/db"
 	"github.com/chroma-core/chroma/go/pkg/log/sysdb"
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgxpool"
 	trace_log "github.com/pingcap/log"
 	"go.uber.org/zap"
@@ -68,6 +70,14 @@ func (r *LogRepository) InsertRecords(ctx context.Context, collectionId string,
 	}
 	insertCount, err = queriesWithTx.InsertRecord(ctx, params)
 	if err != nil {
+		var pgErr *pgconn.PgError
+		// This is a retryable error and should be retried by upstream. It happens
+		// when two concurrent adds to the same collection happen.
+		if errors.As(err, &pgErr) && pgErr.Code == "23505" {
+			trace_log.Error("Duplicate key error while inserting into record_log", zap.String("collectionId", collectionId), zap.String("detail", pgErr.Detail))
+			err = status.Error(codes.AlreadyExists, fmt.Sprintf("Duplicate key error while inserting into record_log: %s", pgErr.Detail))
+			return
+		}
 		trace_log.Error("Error in inserting records to record_log table", zap.Error(err), zap.String("collectionId", collectionId))
 		return
 	}

go/pkg/log/repository/log_test.go

@@ -2,13 +2,16 @@ package repository
 
 import (
 	"context"
+	"errors"
 	"testing"
 	"time"
 
 	"github.com/chroma-core/chroma/go/pkg/log/configuration"
+	log "github.com/chroma-core/chroma/go/pkg/log/store/db"
 	"github.com/chroma-core/chroma/go/pkg/log/sysdb"
 	"github.com/chroma-core/chroma/go/pkg/types"
 	libs2 "github.com/chroma-core/chroma/go/shared/libs"
+	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/suite"
@@ -79,6 +82,34 @@ func (suite *LogTestSuite) TestGarbageCollection() {
 	assert.Equal(suite.t, int64(1), records[0].Offset, "Failed to run garbage collection")
 }
 
+func (suite *LogTestSuite) TestUniqueConstraintPushLogs() {
+	ctx := context.Background()
+	collectionId := types.NewUniqueID()
+
+	records := [][]byte{
+		{1, 2, 3},
+		{4, 5, 6},
+	}
+	params := make([]log.InsertRecordParams, 2)
+	for i, record := range records {
+		offset := 1
+		params[i] = log.InsertRecordParams{
+			CollectionID: collectionId.String(),
+			Record:       record,
+			Offset:       int64(offset),
+			Timestamp:    time.Now().UnixNano(),
+		}
+	}
+	_, err := suite.lr.queries.InsertRecord(ctx, params)
+	assert.Error(suite.t, err, "Failed to insert records")
+	var pgErr *pgconn.PgError
+	if errors.As(err, &pgErr) {
+		assert.Equal(suite.t, "23505", pgErr.Code, "Expected SQLSTATE 23505 for duplicate key")
+	} else {
+		assert.Fail(suite.t, "Expected pgconn.PgError but got different error", err)
+	}
+}
+
 func TestLogTestSuite(t *testing.T) {
 	testSuite := new(LogTestSuite)
 	testSuite.t = t

rust/frontend/src/impls/service_based_frontend.rs

@@ -2,7 +2,7 @@ use crate::{
     config::FrontendConfig, executor::Executor, types::errors::ValidationError,
     CollectionsWithSegmentsProvider,
 };
-use backon::Retryable;
+use backon::{ExponentialBuilder, Retryable};
 use chroma_config::{registry, Configurable};
 use chroma_error::{ChromaError, ErrorCodes};
 use chroma_log::{LocalCompactionManager, LocalCompactionManagerConfig, Log};
@@ -35,10 +35,10 @@ use chroma_types::{
 };
 use opentelemetry::global;
 use opentelemetry::metrics::Counter;
-use std::collections::HashSet;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
+use std::{collections::HashSet, time::Duration};
 
 use super::utils::to_records;
 
@@ -49,6 +49,9 @@ struct Metrics {
     count_retries_counter: Counter<u64>,
     query_retries_counter: Counter<u64>,
     get_retries_counter: Counter<u64>,
+    add_retries_counter: Counter<u64>,
+    update_retries_counter: Counter<u64>,
+    upsert_retries_counter: Counter<u64>,
 }
 
 #[derive(Clone, Debug)]
@@ -61,6 +64,7 @@ pub struct ServiceBasedFrontend {
     max_batch_size: u32,
     metrics: Arc<Metrics>,
     default_knn_index: KnnIndex,
+    retries_builder: ExponentialBuilder,
 }
 
 impl ServiceBasedFrontend {
@@ -78,14 +82,32 @@ impl ServiceBasedFrontend {
         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 get_retries_counter = meter.u64_counter("query_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();
+        let upsert_retries_counter = meter.u64_counter("upsert_retries").build();
         let metrics = Arc::new(Metrics {
             fork_retries_counter,
             delete_retries_counter,
             count_retries_counter,
             query_retries_counter,
             get_retries_counter,
+            add_retries_counter,
+            update_retries_counter,
+            upsert_retries_counter,
         });
+        // factor: 2.0,
+        // min_delay_ms: 100,
+        // max_delay_ms: 5000,
+        // max_attempts: 5,
+        // jitter: true,
+        // TODO(Sanket): Ideally config for this.
+        let retries_builder = ExponentialBuilder::default()
+            .with_max_times(5)
+            .with_factor(2.0)
+            .with_max_delay(Duration::from_millis(5000))
+            .with_min_delay(Duration::from_millis(100))
+            .with_jitter();
         ServiceBasedFrontend {
             allow_reset,
             executor,
@@ -95,6 +117,7 @@ impl ServiceBasedFrontend {
             max_batch_size,
             metrics,
             default_knn_index,
+            retries_builder,
         }
     }
 
@@ -630,6 +653,14 @@ impl ServiceBasedFrontend {
         res
     }
 
+    pub async fn retryable_push_logs(
+        &mut self,
+        collection_id: CollectionUuid,
+        records: Vec<OperationRecord>,
+    ) -> Result<(), Box<dyn ChromaError>> {
+        self.log_client.push_logs(collection_id, records).await
+    }
+
     pub async fn add(
         &mut self,
         AddCollectionRecordsRequest {
@@ -656,16 +687,29 @@ impl ServiceBasedFrontend {
             to_records(ids, embeddings, documents, uris, metadatas, Operation::Add)
                 .map_err(|err| Box::new(err) as Box<dyn ChromaError>)?;
 
-        self.log_client
-            .push_logs(collection_id, records)
-            .await
-            .map_err(|err| {
-                if err.code() == ErrorCodes::Unavailable {
-                    AddCollectionRecordsError::Backoff
-                } else {
-                    AddCollectionRecordsError::Other(Box::new(err) as _)
+        let retries = Arc::new(AtomicUsize::new(0));
+        let add_to_retry = || {
+            let mut self_clone = self.clone();
+            let records_clone = records.clone();
+            async move {
+                self_clone
+                    .retryable_push_logs(collection_id, records_clone)
+                    .await
+            }
+        };
+        let res = add_to_retry
+            .retry(self.retries_builder)
+            .when(|e| matches!(e.code(), ErrorCodes::AlreadyExists))
+            .notify(|_, _| {
+                let retried = retries.fetch_add(1, Ordering::Relaxed);
+                if retried > 0 {
+                    tracing::info!("Retrying add() request for collection {}", collection_id);
                 }
-            })?;
+            })
+            .await;
+        self.metrics
+            .add_retries_counter
+            .add(retries.load(Ordering::Relaxed) as u64, &[]);
 
         // TODO: Submit event after the response is sent
         MeterEvent::CollectionWrite {
@@ -678,7 +722,16 @@ impl ServiceBasedFrontend {
         .submit()
         .await;
 
-        Ok(AddCollectionRecordsResponse {})
+        match res {
+            Ok(()) => Ok(AddCollectionRecordsResponse {}),
+            Err(e) => {
+                if e.code() == ErrorCodes::AlreadyExists {
+                    Err(AddCollectionRecordsError::Backoff)
+                } else {
+                    Err(AddCollectionRecordsError::Other(Box::new(e) as _))
+                }
+            }
+        }
     }
 
     pub async fn update(
@@ -711,16 +764,29 @@ impl ServiceBasedFrontend {
         )
         .map_err(|err| Box::new(err) as Box<dyn ChromaError>)?;
 
-        self.log_client
-            .push_logs(collection_id, records)
-            .await
-            .map_err(|err| {
-                if err.code() == ErrorCodes::Unavailable {
-                    UpdateCollectionRecordsError::Backoff
-                } else {
-                    UpdateCollectionRecordsError::Other(Box::new(err) as _)
+        let retries = Arc::new(AtomicUsize::new(0));
+        let add_to_retry = || {
+            let mut self_clone = self.clone();
+            let records_clone = records.clone();
+            async move {
+                self_clone
+                    .retryable_push_logs(collection_id, records_clone)
+                    .await
+            }
+        };
+        let res = add_to_retry
+            .retry(self.retries_builder)
+            .when(|e| matches!(e.code(), ErrorCodes::AlreadyExists))
+            .notify(|_, _| {
+                let retried = retries.fetch_add(1, Ordering::Relaxed);
+                if retried > 0 {
+                    tracing::info!("Retrying update() request for collection {}", collection_id);
                 }
-            })?;
+            })
+            .await;
+        self.metrics
+            .update_retries_counter
+            .add(retries.load(Ordering::Relaxed) as u64, &[]);
 
         // TODO: Submit event after the response is sent
         MeterEvent::CollectionWrite {
@@ -733,7 +799,16 @@ impl ServiceBasedFrontend {
         .submit()
         .await;
 
-        Ok(UpdateCollectionRecordsResponse {})
+        match res {
+            Ok(()) => Ok(UpdateCollectionRecordsResponse {}),
+            Err(e) => {
+                if e.code() == ErrorCodes::AlreadyExists {
+                    Err(UpdateCollectionRecordsError::Backoff)
+                } else {
+                    Err(UpdateCollectionRecordsError::Other(Box::new(e) as _))
+                }
+            }
+        }
     }
 
     pub async fn upsert(
@@ -768,16 +843,29 @@ impl ServiceBasedFrontend {
         )
         .map_err(|err| Box::new(err) as Box<dyn ChromaError>)?;
 
-        self.log_client
-            .push_logs(collection_id, records)
-            .await
-            .map_err(|err| {
-                if err.code() == ErrorCodes::Unavailable {
-                    UpsertCollectionRecordsError::Backoff
-                } else {
-                    UpsertCollectionRecordsError::Other(Box::new(err) as _)
+        let retries = Arc::new(AtomicUsize::new(0));
+        let add_to_retry = || {
+            let mut self_clone = self.clone();
+            let records_clone = records.clone();
+            async move {
+                self_clone
+                    .retryable_push_logs(collection_id, records_clone)
+                    .await
+            }
+        };
+        let res = add_to_retry
+            .retry(self.retries_builder)
+            .when(|e| matches!(e.code(), ErrorCodes::AlreadyExists))
+            .notify(|_, _| {
+                let retried = retries.fetch_add(1, Ordering::Relaxed);
+                if retried > 0 {
+                    tracing::info!("Retrying upsert() request for collection {}", collection_id);
                 }
-            })?;
+            })
+            .await;
+        self.metrics
+            .upsert_retries_counter
+            .add(retries.load(Ordering::Relaxed) as u64, &[]);
 
         // TODO: Submit event after the response is sent
         MeterEvent::CollectionWrite {
@@ -790,7 +878,16 @@ impl ServiceBasedFrontend {
         .submit()
         .await;
 
-        Ok(UpsertCollectionRecordsResponse {})
+        match res {
+            Ok(()) => Ok(UpsertCollectionRecordsResponse {}),
+            Err(e) => {
+                if e.code() == ErrorCodes::AlreadyExists {
+                    Err(UpsertCollectionRecordsError::Backoff)
+                } else {
+                    Err(UpsertCollectionRecordsError::Other(Box::new(e) as _))
+                }
+            }
+        }
     }
 
     pub async fn retryable_delete(

rust/log/src/grpc_log.rs

@@ -52,7 +52,7 @@ pub enum GrpcPushLogsError {
 impl ChromaError for GrpcPushLogsError {
     fn code(&self) -> ErrorCodes {
         match self {
-            GrpcPushLogsError::Backoff => ErrorCodes::Unavailable,
+            GrpcPushLogsError::Backoff => ErrorCodes::AlreadyExists,
             GrpcPushLogsError::FailedToPushLogs(_) => ErrorCodes::Internal,
             GrpcPushLogsError::ConversionError(_) => ErrorCodes::Internal,
         }
@@ -315,7 +315,9 @@ impl GrpcLog {
             .push_logs(request)
             .await
             .map_err(|err| {
-                if err.code() == ErrorCodes::Unavailable.into() {
+                if err.code() == ErrorCodes::Unavailable.into()
+                    || err.code() == ErrorCodes::AlreadyExists.into()
+                {
                     GrpcPushLogsError::Backoff
                 } else {
                     err.into()