From 684b98504d6d87d14159f17dd79f2fc38a1e5084 Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Thu, 19 Dec 2024 09:31:09 +0800 Subject: [PATCH 1/3] chore(rate_limit): add rate limiting exemption for health check endpoints --- ee/tabby-webserver/src/rate_limit.rs | 7 ++++++- ee/tabby-webserver/src/service/mod.rs | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/ee/tabby-webserver/src/rate_limit.rs b/ee/tabby-webserver/src/rate_limit.rs index c85bd5f7c8c3..ddbc0f5b0352 100644 --- a/ee/tabby-webserver/src/rate_limit.rs +++ b/ee/tabby-webserver/src/rate_limit.rs @@ -20,7 +20,12 @@ impl Default for UserRateLimiter { } impl UserRateLimiter { - pub async fn is_allowed(&self, user_id: &str) -> bool { + pub async fn is_allowed(&self, uri: &axum::http::Uri, user_id: &str) -> bool { + // Do not limit health check requests. + if uri.path().ends_with("/v1/health") || uri.path().ends_with("/v1beta/health") { + return true + } + let mut rate_limiters = self.rate_limiters.lock().await; let rate_limiter = rate_limiters.cache_get_or_set_with(user_id.to_string(), || { // Create a new rate limiter for this user. diff --git a/ee/tabby-webserver/src/service/mod.rs b/ee/tabby-webserver/src/service/mod.rs index eb3300e6310a..ae0f60ca7261 100644 --- a/ee/tabby-webserver/src/service/mod.rs +++ b/ee/tabby-webserver/src/service/mod.rs @@ -230,7 +230,7 @@ impl WorkerService for ServerContext { if let Some(user) = user { // Apply rate limiting when `user` is not none. - if !self.user_rate_limiter.is_allowed(&user).await { + if !self.user_rate_limiter.is_allowed(request.uri(), &user).await { return axum::response::Response::builder() .status(StatusCode::TOO_MANY_REQUESTS) .body(Body::empty()) From 2c32784ee7413ce02113767c6550d01bdbb17500 Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Thu, 19 Dec 2024 09:35:06 +0800 Subject: [PATCH 2/3] update --- ee/tabby-webserver/src/rate_limit.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/ee/tabby-webserver/src/rate_limit.rs b/ee/tabby-webserver/src/rate_limit.rs index ddbc0f5b0352..f313a98bb5e7 100644 --- a/ee/tabby-webserver/src/rate_limit.rs +++ b/ee/tabby-webserver/src/rate_limit.rs @@ -55,12 +55,18 @@ mod tests { let user_id = "test_user"; let rate_limiter = UserRateLimiter::default(); + let uri: axum::http::Uri = "/v1/completions".parse().unwrap(); + let healthcheck_uri: axum::http::Uri = "/v1/health".parse().unwrap(); + // Test that the first `USER_REQUEST_LIMIT_PER_MINUTE` requests are allowed for _ in 0..USER_REQUEST_LIMIT_PER_MINUTE { - assert!(rate_limiter.is_allowed(user_id).await); + assert!(rate_limiter.is_allowed(&uri, user_id).await); } // Test that the 201st request is not allowed - assert!(!rate_limiter.is_allowed(user_id).await); + assert!(!rate_limiter.is_allowed(&uri, user_id).await); + + // Test that health check requests are not limited + assert!(rate_limiter.is_allowed(&healthcheck_uri, user_id).await); } } From 7870c308fabb4e45a7cbec18156ee60764cfeab8 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 19 Dec 2024 01:44:38 +0000 Subject: [PATCH 3/3] [autofix.ci] apply automated fixes --- ee/tabby-webserver/src/rate_limit.rs | 2 +- ee/tabby-webserver/src/service/mod.rs | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/ee/tabby-webserver/src/rate_limit.rs b/ee/tabby-webserver/src/rate_limit.rs index f313a98bb5e7..16abdda8a172 100644 --- a/ee/tabby-webserver/src/rate_limit.rs +++ b/ee/tabby-webserver/src/rate_limit.rs @@ -23,7 +23,7 @@ impl UserRateLimiter { pub async fn is_allowed(&self, uri: &axum::http::Uri, user_id: &str) -> bool { // Do not limit health check requests. if uri.path().ends_with("/v1/health") || uri.path().ends_with("/v1beta/health") { - return true + return true; } let mut rate_limiters = self.rate_limiters.lock().await; diff --git a/ee/tabby-webserver/src/service/mod.rs b/ee/tabby-webserver/src/service/mod.rs index ae0f60ca7261..f67734392f57 100644 --- a/ee/tabby-webserver/src/service/mod.rs +++ b/ee/tabby-webserver/src/service/mod.rs @@ -230,7 +230,11 @@ impl WorkerService for ServerContext { if let Some(user) = user { // Apply rate limiting when `user` is not none. - if !self.user_rate_limiter.is_allowed(request.uri(), &user).await { + if !self + .user_rate_limiter + .is_allowed(request.uri(), &user) + .await + { return axum::response::Response::builder() .status(StatusCode::TOO_MANY_REQUESTS) .body(Body::empty())