Rate limit for NAC

chroma-core · Aug 16, 2024 · dc5f9cf · dc5f9cf
1 parent b2012fe
commit dc5f9cf
Show file tree

Hide file tree

Showing 4 changed files with 120 additions and 3 deletions.
diff --git a/rust/storage/src/admissioncontrolleds3.rs b/rust/storage/src/admissioncontrolleds3.rs
@@ -1,6 +1,6 @@
 use super::GetError;
 use crate::{
-    config::StorageConfig,
+    config::{CountBasedPolicyConfig, RateLimitingConfig, StorageConfig},
     s3::{S3GetError, S3PutError, S3Storage, StorageConfigError},
     stream::ByteStreamItem,
 };
@@ -11,6 +11,7 @@ use futures::{future::Shared, FutureExt, Stream, StreamExt};
 use parking_lot::Mutex;
 use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};
 use thiserror::Error;
+use tokio::sync::{Semaphore, SemaphorePermit};
 use tracing::{Instrument, Span};
 
 /// Wrapper over s3 storage that provides proxy features such as
@@ -38,6 +39,7 @@ pub struct AdmissionControlledS3Storage {
             >,
         >,
     >,
+    rate_limiter: Arc<RateLimitPolicy>,
 }
 
 #[derive(Error, Debug, Clone)]
@@ -55,10 +57,19 @@ impl ChromaError for AdmissionControlledS3StorageError {
 }
 
 impl AdmissionControlledS3Storage {
-    pub fn new(storage: S3Storage) -> Self {
+    pub fn new_with_default_policy(storage: S3Storage) -> Self {
         Self {
             storage,
             outstanding_requests: Arc::new(Mutex::new(HashMap::new())),
+            rate_limiter: Arc::new(RateLimitPolicy::CountBasedPolicy(CountBasedPolicy::new(15))),
+        }
+    }
+
+    pub fn new(storage: S3Storage, policy: RateLimitPolicy) -> Self {
+        Self {
+            storage,
+            outstanding_requests: Arc::new(Mutex::new(HashMap::new())),
+            rate_limiter: Arc::new(policy),
         }
     }
 
@@ -143,10 +154,27 @@ impl AdmissionControlledS3Storage {
         }
     }
 
+    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(
         &self,
         key: String,
     ) -> Result<Arc<Vec<u8>>, AdmissionControlledS3StorageError> {
+        let permit = self.enter().await;
         let future_to_await;
         {
             let mut requests = self.outstanding_requests.lock();
@@ -170,6 +198,8 @@ impl AdmissionControlledS3Storage {
             let mut requests = self.outstanding_requests.lock();
             requests.remove(&key);
         }
+        // Release permit.
+        self.exit(permit).await;
         res
     }
 
@@ -190,11 +220,60 @@ impl Configurable<StorageConfig> for AdmissionControlledS3Storage {
                 let s3_storage =
                     S3Storage::try_from_config(&StorageConfig::S3(nacconfig.s3_config.clone()))
                         .await?;
-                return Ok(Self::new(s3_storage));
+                let policy =
+                    RateLimitPolicy::try_from_config(&nacconfig.rate_limiting_policy).await?;
+                return Ok(Self::new(s3_storage, policy));
             }
             _ => {
                 return Err(Box::new(StorageConfigError::InvalidStorageConfig));
             }
         }
     }
 }
+
+// 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);
+    }
+}
+
+#[async_trait]
+impl Configurable<RateLimitingConfig> for RateLimitPolicy {
+    async fn try_from_config(config: &RateLimitingConfig) -> Result<Self, Box<dyn ChromaError>> {
+        match &config {
+            RateLimitingConfig::CountBasedPolicy(count_policy) => {
+                return Ok(RateLimitPolicy::CountBasedPolicy(CountBasedPolicy::new(
+                    count_policy.max_concurrent_requests,
+                )));
+            }
+        }
+    }
+}
diff --git a/rust/storage/src/config.rs b/rust/storage/src/config.rs
@@ -49,4 +49,15 @@ pub struct LocalStorageConfig {
 #[derive(Deserialize, Debug, Clone)]
 pub struct AdmissionControlledS3StorageConfig {
     pub s3_config: S3StorageConfig,
+    pub rate_limiting_policy: RateLimitingConfig,
+}
+
+#[derive(Deserialize, Debug, Clone)]
+pub struct CountBasedPolicyConfig {
+    pub max_concurrent_requests: usize,
+}
+
+#[derive(Deserialize, Debug, Clone)]
+pub enum RateLimitingConfig {
+    CountBasedPolicy(CountBasedPolicyConfig),
 }
diff --git a/rust/worker/chroma_config.yaml b/rust/worker/chroma_config.yaml
@@ -30,6 +30,9 @@ query_service:
                 connect_timeout_ms: 5000
                 request_timeout_ms: 30000 # 1 minute
                 upload_part_size_bytes: 536870912 # 512MiB
+            rate_limiting_policy:
+                CountBasedPolicy:
+                    max_concurrent_requests: 15
     log:
         Grpc:
             host: "logservice.chroma"
@@ -84,6 +87,9 @@ compaction_service:
                 connect_timeout_ms: 5000
                 request_timeout_ms: 60000 # 1 minute
                 upload_part_size_bytes: 536870912 # 512MiB
+            rate_limiting_policy:
+                CountBasedPolicy:
+                    max_concurrent_requests: 15
     log:
         Grpc:
             host: "logservice.chroma"

diff --git a/rust/worker/src/config.rs b/rust/worker/src/config.rs
@@ -171,6 +171,9 @@ mod tests {
                                 connect_timeout_ms: 5000
                                 request_timeout_ms: 1000
                                 upload_part_size_bytes: 8388608
+                            rate_limiting_policy:
+                                CountBasedPolicy:
+                                    max_concurrent_requests: 15
                     log:
                         Grpc:
                             host: "localhost"
@@ -225,6 +228,9 @@ mod tests {
                                 connect_timeout_ms: 5000
                                 request_timeout_ms: 1000
                                 upload_part_size_bytes: 8388608
+                            rate_limiting_policy:
+                                CountBasedPolicy:
+                                    max_concurrent_requests: 15
                     log:
                         Grpc:
                             host: "localhost"
@@ -305,6 +311,9 @@ mod tests {
                                 connect_timeout_ms: 5000
                                 request_timeout_ms: 1000
                                 upload_part_size_bytes: 8388608
+                            rate_limiting_policy:
+                                CountBasedPolicy:
+                                    max_concurrent_requests: 15
                     log:
                         Grpc:
                             host: "localhost"
@@ -359,6 +368,9 @@ mod tests {
                                 connect_timeout_ms: 5000
                                 request_timeout_ms: 1000
                                 upload_part_size_bytes: 8388608
+                            rate_limiting_policy:
+                                CountBasedPolicy:
+                                    max_concurrent_requests: 15
                     log:
                         Grpc:
                             host: "localhost"
@@ -457,6 +469,9 @@ mod tests {
                                 connect_timeout_ms: 5000
                                 request_timeout_ms: 1000
                                 upload_part_size_bytes: 8388608
+                            rate_limiting_policy:
+                                CountBasedPolicy:
+                                    max_concurrent_requests: 15
                     log:
                         Grpc:
                             host: "localhost"
@@ -511,6 +526,9 @@ mod tests {
                                 connect_timeout_ms: 5000
                                 request_timeout_ms: 1000
                                 upload_part_size_bytes: 8388608
+                            rate_limiting_policy:
+                                CountBasedPolicy:
+                                    max_concurrent_requests: 15
                     log:
                         Grpc:
                             host: "localhost"
@@ -607,6 +625,9 @@ mod tests {
                                 connect_timeout_ms: 5000
                                 request_timeout_ms: 1000
                                 upload_part_size_bytes: 8388608
+                            rate_limiting_policy:
+                                CountBasedPolicy:
+                                    max_concurrent_requests: 15
                     log:
                         Grpc:
                             host: "localhost"