[ENH] Add a delete_many call to the storage API. (#5020)

## Description of changes

This adds a delete-many call to storage that can batch a delete of many
keys to S3 storage.  The goal is to reduce GC's CPU overhead from many
serial deletes.

## Test plan

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

## Documentation Changes

N/A

Author: rescrv

rust/storage/src/admissioncontrolleds3.rs

@@ -489,8 +489,15 @@ impl AdmissionControlledS3Storage {
         self.storage.list_prefix(prefix).await
     }
 
-    pub async fn delete(&self, prefix: &str, options: DeleteOptions) -> Result<(), StorageError> {
-        self.storage.delete(prefix, options).await
+    pub async fn delete(&self, key: &str, options: DeleteOptions) -> Result<(), StorageError> {
+        self.storage.delete(key, options).await
+    }
+
+    pub async fn delete_many<S: AsRef<str> + std::fmt::Debug, I: IntoIterator<Item = S>>(
+        &self,
+        keys: I,
+    ) -> Result<crate::s3::DeletedObjects, StorageError> {
+        self.storage.delete_many(keys).await
     }
 }
 

rust/storage/src/lib.rs

@@ -316,6 +316,18 @@ impl Storage {
         }
     }
 
+    pub async fn delete_many<S: AsRef<str> + std::fmt::Debug, I: IntoIterator<Item = S>>(
+        &self,
+        keys: I,
+    ) -> Result<crate::s3::DeletedObjects, StorageError> {
+        match self {
+            Storage::ObjectStore(_) => Err(StorageError::NotImplemented),
+            Storage::S3(s3) => s3.delete_many(keys).await,
+            Storage::Local(_) => Err(StorageError::NotImplemented),
+            Storage::AdmissionControlledS3(ac) => ac.delete_many(keys).await,
+        }
+    }
+
     pub async fn rename(&self, src_key: &str, dst_key: &str) -> Result<(), StorageError> {
         match self {
             Storage::ObjectStore(object_store) => object_store.rename(src_key, dst_key).await,

rust/storage/src/s3.rs

@@ -23,8 +23,7 @@ use aws_sdk_s3::error::SdkError;
 use aws_sdk_s3::operation::create_bucket::CreateBucketError;
 use aws_sdk_s3::operation::get_object::GetObjectOutput;
 use aws_sdk_s3::primitives::ByteStream;
-use aws_sdk_s3::types::CompletedMultipartUpload;
-use aws_sdk_s3::types::CompletedPart;
+use aws_sdk_s3::types::{CompletedMultipartUpload, CompletedPart, Delete, ObjectIdentifier};
 use aws_smithy_types::byte_stream::Length;
 use bytes::Bytes;
 use chroma_config::registry::Registry;
@@ -43,6 +42,12 @@ use std::time::Duration;
 use tokio::io::AsyncReadExt;
 use tracing::Instrument;
 
+#[derive(Clone, Debug, Default)]
+pub struct DeletedObjects {
+    pub deleted: Vec<String>,
+    pub errors: Vec<StorageError>,
+}
+
 #[derive(Clone)]
 pub struct S3Storage {
     pub(super) bucket: String,
@@ -659,6 +664,67 @@ impl S3Storage {
         }
     }
 
+    #[tracing::instrument(skip(self, keys), level = "trace")]
+    pub async fn delete_many<S: AsRef<str> + std::fmt::Debug, I: IntoIterator<Item = S>>(
+        &self,
+        keys: I,
+    ) -> Result<DeletedObjects, StorageError> {
+        tracing::trace!("Deleting objects from S3");
+        let mut objects = vec![];
+        for key in keys {
+            objects.push(
+                ObjectIdentifier::builder()
+                    .key(key.as_ref())
+                    .build()
+                    .map_err(|err| StorageError::Generic {
+                        source: Arc::new(err),
+                    })?,
+            );
+        }
+        let delete = Delete::builder()
+            .set_objects(Some(objects))
+            .build()
+            .map_err(|err| StorageError::Generic {
+                source: Arc::new(err),
+            })?;
+
+        let req = self
+            .client
+            .delete_objects()
+            .bucket(&self.bucket)
+            .delete(delete);
+
+        match req.send().await {
+            Ok(resp) => {
+                let mut out = DeletedObjects::default();
+                for deleted in resp.deleted() {
+                    if let Some(key) = deleted.key.clone() {
+                        out.deleted.push(key);
+                    }
+                }
+                for error in resp.errors() {
+                    out.errors.push(if Some("NoSuchKey") == error.code() {
+                        StorageError::NotFound {
+                            path: error.key.clone().unwrap_or(String::new()),
+                            source: Arc::new(StorageError::Message {
+                                message: format!("{error:#?}"),
+                            }),
+                        }
+                    } else {
+                        StorageError::Message {
+                            message: format!("{error:#?}"),
+                        }
+                    });
+                }
+                tracing::trace!("Successfully deleted objects from S3");
+                Ok(out)
+            }
+            Err(e) => Err(StorageError::Generic {
+                source: Arc::new(e),
+            }),
+        }
+    }
+
     #[tracing::instrument(skip(self), level = "trace")]
     pub async fn rename(&self, src_key: &str, dst_key: &str) -> Result<(), StorageError> {
         // S3 doesn't have a native rename operation, so we need to copy and delete
@@ -1214,4 +1280,50 @@ mod tests {
             assert_eq!(format!("test/{:02x}", i), *result, "index = {}", i);
         }
     }
+
+    #[tokio::test]
+    async fn test_k8s_integration_delete_many() {
+        let storage = setup_with_bucket(1024 * 1024 * 8, 1024 * 1024 * 8).await;
+
+        // Create every other file (0, 2, 4, 6, 8, 10, 12, 14)
+        let mut created_files = Vec::new();
+        for i in (0..16).step_by(2) {
+            let key = format!("test/{:02x}", i);
+            storage
+                .oneshot_upload(
+                    &key,
+                    0,
+                    |_| Box::pin(ready(Ok(ByteStream::from(Bytes::new())))) as _,
+                    PutOptions {
+                        if_not_exists: true,
+                        if_match: None,
+                        priority: StorageRequestPriority::P0,
+                    },
+                )
+                .await
+                .unwrap();
+            created_files.push(key);
+        }
+
+        // Try to delete all files (0-15), including ones that don't exist
+        let mut all_keys = Vec::new();
+        for i in 0..16 {
+            all_keys.push(format!("test/{:02x}", i));
+        }
+        let all_key_refs: Vec<&str> = all_keys.iter().map(|s| s.as_str()).collect();
+
+        let delete_result = storage.delete_many(&all_key_refs).await.unwrap();
+
+        // Verify that every created file appears in the deleted files
+        for created_file in &created_files {
+            assert!(
+                delete_result.deleted.contains(created_file),
+                "Created file {} should be in deleted files",
+                created_file
+            );
+        }
+
+        eprintln!("Successfully deleted: {:#?}", delete_result.deleted);
+        eprintln!("Errors for non-existent files: {:#?}", delete_result.errors);
+    }
 }