chore(rate_limit): add rate limiting exemption for health check endpo…

…ints
TabbyML · Dec 19, 2024 · 684b985 · 684b985
1 parent c8735f5
commit 684b985
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git 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
@@ -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())