Rate limit for NAC

chroma-core · Aug 4, 2024 · 197103a · 197103a
1 parent 648e0f8
commit 197103a
Show file tree

Hide file tree

Showing 14 changed files with 212 additions and 43 deletions.
diff --git a/rust/blockstore/src/arrow/block/delta.rs b/rust/blockstore/src/arrow/block/delta.rs
@@ -225,7 +225,8 @@ mod test {
         let path = tmp_dir.path().to_str().unwrap();
         let storage = Storage::Local(LocalStorage::new(path));
         let cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let block_manager = BlockManager::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -274,7 +275,8 @@ mod test {
         let path = tmp_dir.path().to_str().unwrap();
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let block_manager = BlockManager::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -337,7 +339,8 @@ mod test {
         let path = tmp_dir.path().to_str().unwrap();
         let storage = Storage::Local(LocalStorage::new(path));
         let cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let block_manager = BlockManager::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -380,7 +383,8 @@ mod test {
         let path = tmp_dir.path().to_str().unwrap();
         let storage = Storage::Local(LocalStorage::new(path));
         let cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let block_manager = BlockManager::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -421,7 +425,8 @@ mod test {
         let path = tmp_dir.path().to_str().unwrap();
         let storage = Storage::Local(LocalStorage::new(path));
         let cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let block_manager = BlockManager::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -490,7 +495,8 @@ mod test {
         let path = tmp_dir.path().to_str().unwrap();
         let storage = Storage::Local(LocalStorage::new(path));
         let cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let block_manager = BlockManager::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,

diff --git a/rust/blockstore/src/arrow/blockfile.rs b/rust/blockstore/src/arrow/blockfile.rs
@@ -591,7 +591,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -633,7 +634,8 @@ mod tests {
             let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
             let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
             let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-            let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+            let network_admission_control =
+                NetworkAdmissionControl::new_with_default_policy(storage.clone());
             let blockfile_provider = ArrowBlockfileProvider::new(
                 storage,
                 TEST_MAX_BLOCK_SIZE_BYTES,
@@ -700,7 +702,8 @@ mod tests {
             let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
             let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
             let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-            let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+            let network_admission_control =
+                NetworkAdmissionControl::new_with_default_policy(storage.clone());
             let blockfile_provider = ArrowBlockfileProvider::new(
                 storage,
                 TEST_MAX_BLOCK_SIZE_BYTES,
@@ -817,7 +820,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -859,7 +863,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -975,7 +980,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -1019,7 +1025,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -1058,7 +1065,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -1094,7 +1102,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -1142,7 +1151,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -1178,7 +1188,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -1232,7 +1243,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -1266,7 +1278,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -1331,7 +1344,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,
@@ -1372,7 +1386,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let block_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
         let sparse_index_cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let blockfile_provider = ArrowBlockfileProvider::new(
             storage,
             TEST_MAX_BLOCK_SIZE_BYTES,

diff --git a/rust/blockstore/src/arrow/concurrency_test.rs b/rust/blockstore/src/arrow/concurrency_test.rs
@@ -25,7 +25,8 @@ mod tests {
                 let sparse_index_cache = Cache::new(&CacheConfig::Lru(LruConfig {
                     capacity: SPARSE_INDEX_CACHE_CAPACITY,
                 }));
-                let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+                let network_admission_control =
+                    NetworkAdmissionControl::new_with_default_policy(storage.clone());
                 let blockfile_provider = ArrowBlockfileProvider::new(
                     storage,
                     TEST_MAX_BLOCK_SIZE_BYTES,

diff --git a/rust/index/src/hnsw_provider.rs b/rust/index/src/hnsw_provider.rs
@@ -487,7 +487,8 @@ mod tests {
 
         let storage = Storage::Local(LocalStorage::new(storage_dir.to_str().unwrap()));
         let cache = Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}));
-        let network_admission_control = NetworkAdmissionControl::new(storage.clone());
+        let network_admission_control =
+            NetworkAdmissionControl::new_with_default_policy(storage.clone());
         let provider =
             HnswIndexProvider::new(storage, hnsw_tmp_path, cache, network_admission_control);
         let segment = Segment {

diff --git a/rust/storage/src/config.rs b/rust/storage/src/config.rs
@@ -43,3 +43,13 @@ pub struct S3StorageConfig {
 pub struct LocalStorageConfig {
     pub root: String,
 }
+
+#[derive(Deserialize, Debug)]
+pub struct CountBasedPolicyConfig {
+    pub max_concurrent_requests: usize,
+}
+
+#[derive(Deserialize, Debug)]
+pub enum StorageAdmissionConfig {
+    CountBasedPolicy(CountBasedPolicyConfig),
+}
diff --git a/rust/storage/src/network_admission_control.rs b/rust/storage/src/network_admission_control.rs
@@ -1,9 +1,14 @@
+use crate::config::{CountBasedPolicyConfig, StorageAdmissionConfig};
+
 use super::{GetError, Storage};
+use async_trait::async_trait;
+use chroma_config::Configurable;
 use chroma_error::{ChromaError, ErrorCodes};
-use futures::{future::Shared, FutureExt, StreamExt};
+use futures::{future::Shared, FutureExt, StreamExt, TryFutureExt};
 use parking_lot::Mutex;
-use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};
+use std::{collections::HashMap, future::Future, marker::PhantomData, pin::Pin, sync::Arc};
 use thiserror::Error;
+use tokio::sync::{Semaphore, SemaphorePermit};
 use tracing::{Instrument, Span};
 
 #[derive(Clone)]
@@ -25,6 +30,7 @@ pub struct NetworkAdmissionControl {
             >,
         >,
     >,
+    rate_limiter: Arc<RateLimitPolicy>,
 }
 
 #[derive(Error, Debug, Clone)]
@@ -48,10 +54,18 @@ impl ChromaError for NetworkAdmissionControlError {
 }
 
 impl NetworkAdmissionControl {
-    pub fn new(storage: Storage) -> Self {
+    pub fn new_with_default_policy(storage: Storage) -> Self {
         Self {
             storage,
             outstanding_requests: Arc::new(Mutex::new(HashMap::new())),
+            rate_limiter: Arc::new(RateLimitPolicy::CountBasedPolicy(CountBasedPolicy::new(15))),
+        }
+    }
+    pub fn new(storage: Storage, policy: RateLimitPolicy) -> Self {
+        Self {
+            storage,
+            outstanding_requests: Arc::new(Mutex::new(HashMap::new())),
+            rate_limiter: Arc::new(policy),
         }
     }
 
@@ -108,6 +122,22 @@ impl NetworkAdmissionControl {
         }
     }
 
+    async fn enter(&self) -> SemaphorePermit<'_> {
+        match &*self.rate_limiter {
+            RateLimitPolicy::CountBasedPolicy(policy) => {
+                return policy.acquire().await;
+            }
+        }
+    }
+
+    async fn exit(&self, permit: SemaphorePermit<'_>) {
+        match &*self.rate_limiter {
+            RateLimitPolicy::CountBasedPolicy(policy) => {
+                policy.drop(permit).await;
+            }
+        }
+    }
+
     pub async fn get<F, R>(
         &self,
         key: String,
@@ -117,6 +147,8 @@ impl NetworkAdmissionControl {
         R: Future<Output = Result<(), Box<NetworkAdmissionControlError>>> + Send + 'static,
         F: (FnOnce(Vec<u8>) -> R) + Send + 'static,
     {
+        // Wait for permit.
+        let permit = self.enter().await;
         let future_to_await;
         {
             let mut requests = self.outstanding_requests.lock();
@@ -141,6 +173,63 @@ impl NetworkAdmissionControl {
             let mut requests = self.outstanding_requests.lock();
             requests.remove(&key);
         }
+        // Release permit.
+        self.exit(permit).await;
         res
     }
 }
+
+// Prefer enum dispatch over dyn since there could
+// only be a handful of these policies.
+#[derive(Debug)]
+enum RateLimitPolicy {
+    CountBasedPolicy(CountBasedPolicy),
+}
+
+#[derive(Debug)]
+struct CountBasedPolicy {
+    max_allowed_outstanding: usize,
+    remaining_tokens: Semaphore,
+}
+
+impl CountBasedPolicy {
+    fn new(max_allowed_outstanding: usize) -> Self {
+        Self {
+            max_allowed_outstanding,
+            remaining_tokens: Semaphore::new(max_allowed_outstanding),
+        }
+    }
+    async fn acquire(&self) -> SemaphorePermit<'_> {
+        let token_res = self.remaining_tokens.acquire().await;
+        match token_res {
+            Ok(token) => {
+                return token;
+            }
+            Err(e) => panic!("AcquireToken Failed {}", e),
+        }
+    }
+    async fn drop(&self, permit: SemaphorePermit<'_>) {
+        drop(permit);
+    }
+}
+
+pub async fn from_config(
+    config: &StorageAdmissionConfig,
+    storage: Storage,
+) -> Result<NetworkAdmissionControl, Box<dyn ChromaError>> {
+    match &config {
+        StorageAdmissionConfig::CountBasedPolicy(policy) => Ok(NetworkAdmissionControl::new(
+            storage,
+            RateLimitPolicy::CountBasedPolicy(CountBasedPolicy::try_from_config(policy).await?),
+        )),
+    }
+}
+
+#[async_trait]
+impl Configurable<CountBasedPolicyConfig> for CountBasedPolicy {
+    async fn try_from_config(
+        config: &CountBasedPolicyConfig,
+    ) -> Result<Self, Box<dyn ChromaError>> {
+        Ok(Self::new(config.max_concurrent_requests))
+    }
+}