[ENH] Support list_prefix operations for S3 and AC/S3.

rust/garbage_collector/src/garbage_collector_orchestrator.rs

@@ -527,6 +527,7 @@ mod tests {
     use chroma_storage::config::{
         ObjectStoreBucketConfig, ObjectStoreConfig, ObjectStoreType, StorageConfig,
     };
+    use chroma_storage::GetOptions;
     use chroma_sysdb::{GrpcSysDbConfig, SysDbConfig};
     use chroma_system::System;
     use std::str::FromStr;
@@ -751,7 +752,7 @@ mod tests {
 
     async fn get_hnsw_index_ids(storage: &Storage) -> Vec<Uuid> {
         storage
-            .list_prefix("hnsw")
+            .list_prefix("hnsw", GetOptions::default())
             .await
             .unwrap()
             .into_iter()
@@ -935,7 +936,7 @@ mod tests {
             .unwrap();
 
         let deleted_hnsw_files_before_test: Vec<_> = storage
-            .list_prefix("gc")
+            .list_prefix("gc", GetOptions::default())
             .await
             .unwrap()
             .into_iter()
@@ -1060,7 +1061,7 @@ mod tests {
 
         // Verify that "deleted" files are renamed with the "gc" prefix
         let deleted_hnsw_files: Vec<_> = storage
-            .list_prefix("gc")
+            .list_prefix("gc", GetOptions::default())
             .await
             .unwrap()
             .into_iter()

rust/garbage_collector/tests/proptest_helpers/garbage_collector_under_test.rs

@@ -48,7 +48,11 @@ impl Drop for GarbageCollectorUnderTest {
         self.runtime.block_on(async {
             self.sysdb.reset().await.unwrap();
 
-            let files = self.storage.list_prefix("").await.unwrap();
+            let files = self
+                .storage
+                .list_prefix("", GetOptions::default())
+                .await
+                .unwrap();
             if files.is_empty() {
                 return;
             }
@@ -472,7 +476,7 @@ impl StateMachineTest for GarbageCollectorUnderTest {
         let file_ref_counts = ref_state.get_file_ref_counts();
         let files_on_disk = ref_state
             .runtime
-            .block_on(state.storage.list_prefix(""))
+            .block_on(state.storage.list_prefix("", GetOptions::default()))
             .unwrap()
             .into_iter()
             .collect::<HashSet<_>>();

rust/storage/src/admissioncontrolleds3.rs

@@ -482,6 +482,16 @@ impl AdmissionControlledS3Storage {
         // Akin to a HEAD request; no AC.
         self.storage.copy(src_key, dst_key).await
     }
+
+    pub async fn list_prefix(
+        &self,
+        prefix: &str,
+        options: GetOptions,
+    ) -> Result<Vec<String>, StorageError> {
+        let atomic_priority = Arc::new(AtomicUsize::new(options.priority.as_usize()));
+        let _permit = self.rate_limiter.enter(atomic_priority, None).await;
+        self.storage.list_prefix(prefix).await
+    }
 }
 
 #[async_trait]

rust/storage/src/lib.rs

@@ -320,16 +320,16 @@ impl Storage {
         }
     }
 
-    pub async fn list_prefix(&self, prefix: &str) -> Result<Vec<String>, StorageError> {
+    pub async fn list_prefix(
+        &self,
+        prefix: &str,
+        options: GetOptions,
+    ) -> Result<Vec<String>, StorageError> {
         match self {
             Storage::Local(local) => local.list_prefix(prefix).await,
-            Storage::S3(_) => {
-                unimplemented!("list_prefix not implemented for S3")
-            }
+            Storage::S3(s3) => s3.list_prefix(prefix).await,
             Storage::ObjectStore(object_store) => object_store.list_prefix(prefix).await,
-            Storage::AdmissionControlledS3(_) => {
-                unimplemented!("list_prefix not implemented for AdmissionControlledS3")
-            }
+            Storage::AdmissionControlledS3(acs3) => acs3.list_prefix(prefix, options).await,
         }
     }
 }

rust/storage/src/s3.rs

@@ -686,6 +686,33 @@ impl S3Storage {
             }
         }
     }
+
+    pub async fn list_prefix(&self, prefix: &str) -> Result<Vec<String>, StorageError> {
+        let mut outs = self
+            .client
+            .list_objects_v2()
+            .bucket(&self.bucket)
+            .set_max_keys(Some(1000))
+            .prefix(prefix)
+            .into_paginator()
+            .send();
+        let mut paths = vec![];
+        while let Some(result) = outs.next().await {
+            let output = result.map_err(|err| StorageError::Generic {
+                source: Arc::new(err),
+            })?;
+            for object in output.contents() {
+                if let Some(key) = object.key() {
+                    paths.push(key.to_string());
+                } else {
+                    return Err(StorageError::Message {
+                        message: format!("list on prefix {:?} led to empty key", prefix),
+                    });
+                }
+            }
+        }
+        Ok(paths)
+    }
 }
 
 #[async_trait]
@@ -1131,4 +1158,32 @@ mod tests {
         let bytes = storage.get("test2/00").await.unwrap();
         assert_eq!("ABC123XYZ".as_bytes(), bytes.as_slice());
     }
+
+    #[tokio::test]
+    async fn test_k8s_integration_list_prefix() {
+        let storage = setup_with_bucket(1024 * 1024 * 8, 1024 * 1024 * 8).await;
+        for i in 0..16 {
+            storage
+                .oneshot_upload(
+                    &format!("test/{:02x}", i),
+                    0,
+                    |_| Box::pin(ready(Ok(ByteStream::from(Bytes::new())))) as _,
+                    PutOptions {
+                        if_not_exists: true,
+                        if_match: None,
+                        priority: StorageRequestPriority::P0,
+                    },
+                )
+                .await
+                .unwrap();
+        }
+
+        let mut results = storage.list_prefix("test/").await.unwrap();
+        results.sort();
+        eprintln!("Results of listing (should be 0x00..0xff inclusive): {results:#?}");
+        assert_eq!(16, results.len());
+        for (i, result) in results.iter().enumerate() {
+            assert_eq!(format!("test/{:02x}", i), *result, "index = {}", i);
+        }
+    }
 }