test(server): Add tests for throughput rate limits (#254)

lcian · web-flow · commit a9f3b741eb32 · 2026-01-09T09:55:07.000Z
diff --git a/objectstore-server/src/config.rs b/objectstore-server/src/config.rs
@@ -933,6 +933,7 @@ mod tests {
     use secrecy::ExposeSecret;
 
     use crate::killswitches::Killswitch;
+    use crate::rate_limits::{BandwidthLimits, RateLimits, ThroughputLimits, ThroughputRule};
 
     use super::*;
 
@@ -1187,4 +1188,64 @@ mod tests {
             Ok(())
         });
     }
+
+    #[test]
+    fn configure_rate_limits_with_yaml() {
+        let mut tempfile = tempfile::NamedTempFile::new().unwrap();
+        tempfile
+            .write_all(
+                br#"
+                rate_limits:
+                  throughput:
+                    global_rps: 1000
+                    burst: 100
+                    usecase_pct: 50
+                    scope_pct: 25
+                    rules:
+                      - usecase: "high_priority"
+                        scopes:
+                          - ["org", "123"]
+                        rps: 500
+                      - scopes:
+                          - ["org", "456"]
+                          - ["project", "789"]
+                        pct: 10
+                "#,
+            )
+            .unwrap();
+
+        figment::Jail::expect_with(|_jail| {
+            let expected = RateLimits {
+                throughput: ThroughputLimits {
+                    global_rps: Some(1000),
+                    burst: 100,
+                    usecase_pct: Some(50),
+                    scope_pct: Some(25),
+                    rules: vec![
+                        ThroughputRule {
+                            usecase: Some("high_priority".to_string()),
+                            scopes: vec![("org".to_string(), "123".to_string())],
+                            rps: Some(500),
+                            pct: None,
+                        },
+                        ThroughputRule {
+                            usecase: None,
+                            scopes: vec![
+                                ("org".to_string(), "456".to_string()),
+                                ("project".to_string(), "789".to_string()),
+                            ],
+                            rps: None,
+                            pct: Some(10),
+                        },
+                    ],
+                },
+                bandwidth: BandwidthLimits::default(),
+            };
+
+            let config = Config::load(Some(tempfile.path())).unwrap();
+            assert_eq!(config.rate_limits, expected);
+
+            Ok(())
+        });
+    }
 }
diff --git a/objectstore-server/src/extractors/id.rs b/objectstore-server/src/extractors/id.rs
@@ -56,10 +56,12 @@ impl FromRequestParts<ServiceState> for Xt<ObjectId> {
         sentry::configure_scope(|s| s.set_extra("key", id.key().into()));
 
         if state.config.killswitches.matches(id.context()) {
+            tracing::debug!("Request rejected due to killswitches");
             return Err(ObjectRejection::Killswitched);
         }
 
         if !state.rate_limiter.check(id.context()) {
+            tracing::debug!("Request rejected due to rate limits");
             return Err(ObjectRejection::RateLimited);
         }
 
@@ -121,10 +123,12 @@ impl FromRequestParts<ServiceState> for Xt<ObjectContext> {
         populate_sentry_context(&context);
 
         if state.config.killswitches.matches(&context) {
+            tracing::debug!("Request rejected due to killswitches");
             return Err(ObjectRejection::Killswitched);
         }
 
         if !state.rate_limiter.check(&context) {
+            tracing::debug!("Request rejected due to rate limits");
             return Err(ObjectRejection::RateLimited);
         }
 
diff --git a/objectstore-server/src/rate_limits.rs b/objectstore-server/src/rate_limits.rs
@@ -6,15 +6,15 @@ use objectstore_types::scope::Scopes;
 use serde::{Deserialize, Serialize};
 
 /// Rate limits for objectstore.
-#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq)]
 pub struct RateLimits {
     /// Limits the number of requests per second per service instance.
     pub throughput: ThroughputLimits,
     /// Limits the concurrent bandwidth per service instance.
     pub bandwidth: BandwidthLimits,
 }
 
-#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq)]
 pub struct ThroughputLimits {
     /// The overall maximum number of requests per second per service instance.
     ///
@@ -45,7 +45,7 @@ pub struct ThroughputLimits {
     pub rules: Vec<ThroughputRule>,
 }
 
-#[derive(Clone, Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
 pub struct ThroughputRule {
     /// Optional usecase to match.
     ///
@@ -91,7 +91,7 @@ impl ThroughputRule {
     }
 }
 
-#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq)]
 pub struct BandwidthLimits {
     /// The overall maximum bandwidth (in bytes per second) per service instance.
     ///
@@ -237,15 +237,15 @@ struct TokenBucket {
 }
 
 impl TokenBucket {
-    /// Creates a new token bucket with the specified rate limit and burst capacity.
+    /// Creates a new, full token bucket with the specified rate limit and burst capacity.
     ///
     /// - `rps`: tokens refilled per second (sustained rate limit)
     /// - `burst`: initial tokens and burst allowance above sustained rate
     pub fn new(rps: u32, burst: u32) -> Self {
         Self {
             refill_rate: rps as f64,
             capacity: (rps + burst) as f64,
-            tokens: burst as f64,
+            tokens: (rps + burst) as f64,
             last_update: Instant::now(),
         }
     }
diff --git a/objectstore-server/tests/limits.rs b/objectstore-server/tests/limits.rs
@@ -8,6 +8,7 @@ use std::collections::BTreeMap;
 use anyhow::Result;
 use objectstore_server::config::{AuthZ, Config};
 use objectstore_server::killswitches::{Killswitch, Killswitches};
+use objectstore_server::rate_limits::{RateLimits, ThroughputLimits, ThroughputRule};
 use objectstore_test::server::TestServer;
 
 #[tokio::test]
@@ -51,3 +52,221 @@ async fn test_killswitches() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_throughput_global_rps_limit() -> Result<()> {
+    let server = TestServer::with_config(Config {
+        rate_limits: RateLimits {
+            throughput: ThroughputLimits {
+                global_rps: Some(2),
+                burst: 1,
+                ..Default::default()
+            },
+            ..Default::default()
+        },
+        ..Default::default()
+    })
+    .await;
+
+    let client = reqwest::Client::new();
+
+    // First three requests without waiting should succeed, using up both the regular and burst budget
+    for _ in 0..3 {
+        let response = client
+            .get(server.url("/v1/objects/test/org=1/nonexistent"))
+            .send()
+            .await?;
+        assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+    }
+
+    // Fourth request without waiting should be rate limited
+    let response = client
+        .get(server.url("/v1/objects/test/org=1/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
+
+    // Refill bucket
+    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+
+    // After waiting, request succeeds
+    let response = client
+        .get(server.url("/v1/objects/test/org=1/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_throughput_usecase_pct_limit() -> Result<()> {
+    let server = TestServer::with_config(Config {
+        rate_limits: RateLimits {
+            throughput: ThroughputLimits {
+                global_rps: Some(100),
+                burst: 0,
+                usecase_pct: Some(2),
+                ..Default::default()
+            },
+            ..Default::default()
+        },
+        ..Default::default()
+    })
+    .await;
+
+    let client = reqwest::Client::new();
+
+    // First two requests to the same usecase without waiting should succeed
+    for _ in 0..2 {
+        let response = client
+            .get(server.url("/v1/objects/test/org=1/nonexistent"))
+            .send()
+            .await?;
+        assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+    }
+
+    // Third request to the same usecase without waiting should fail
+    let response = client
+        .get(server.url("/v1/objects/test/org=1/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
+
+    // Request to a different usecase should succeed
+    let response = client
+        .get(server.url("/v1/objects/other/org=1/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+
+    // Refill bucket
+    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+
+    // After waiting, request to first usecase should succeed again
+    let response = client
+        .get(server.url("/v1/objects/test/org=1/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_throughput_scope_pct_limit() -> Result<()> {
+    let server = TestServer::with_config(Config {
+        rate_limits: RateLimits {
+            throughput: ThroughputLimits {
+                global_rps: Some(100),
+                burst: 0,
+                scope_pct: Some(2),
+                ..Default::default()
+            },
+            ..Default::default()
+        },
+        ..Default::default()
+    })
+    .await;
+
+    let client = reqwest::Client::new();
+
+    // First two requests to the same scope without waiting should succeed
+    for _ in 0..2 {
+        let response = client
+            .get(server.url("/v1/objects/test/org=1/nonexistent"))
+            .send()
+            .await?;
+        assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+    }
+
+    // Third request to the same scope without waiting should fail
+    let response = client
+        .get(server.url("/v1/objects/test/org=1/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
+
+    // Request to a different scope should succeed
+    let response = client
+        .get(server.url("/v1/objects/test/org=2/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+
+    // Refill bucket
+    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+
+    // After waiting, request to first scope should succeed again
+    let response = client
+        .get(server.url("/v1/objects/test/org=1/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_throughput_rule() -> Result<()> {
+    let server = TestServer::with_config(Config {
+        rate_limits: RateLimits {
+            throughput: ThroughputLimits {
+                global_rps: None,
+                burst: 0,
+                rules: vec![ThroughputRule {
+                    usecase: Some("restricted".to_string()),
+                    scopes: vec![("org".to_string(), "42".to_string())],
+                    rps: Some(1),
+                    pct: None,
+                }],
+                ..Default::default()
+            },
+            ..Default::default()
+        },
+        ..Default::default()
+    })
+    .await;
+
+    let client = reqwest::Client::new();
+
+    // First request matching rule should succeed
+    let response = client
+        .get(server.url("/v1/objects/restricted/org=42/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+
+    // Second request matching rule should fail
+    let response = client
+        .get(server.url("/v1/objects/restricted/org=42/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
+
+    // Different usecase should not be affected by rule
+    let response = client
+        .get(server.url("/v1/objects/other/org=42/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+
+    // Same usecase but different scope should not be affected by rule
+    let response = client
+        .get(server.url("/v1/objects/restricted/org=43/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+
+    // Refill bucket
+    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+
+    // After waiting, request matching rule should succeed again
+    let response = client
+        .get(server.url("/v1/objects/restricted/org=42/nonexistent"))
+        .send()
+        .await?;
+    assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
+
+    Ok(())
+}

Original file line number	Diff line number	Diff line change
`@@ -56,10 +56,12 @@ impl FromRequestParts<ServiceState> for Xt<ObjectId> {`
`56`	`56`	`sentry::configure_scope(\|s\| s.set_extra("key", id.key().into()));`
`57`	`57`
`58`	`58`	`if state.config.killswitches.matches(id.context()) {`
	`59`	`+ tracing::debug!("Request rejected due to killswitches");`
`59`	`60`	`return Err(ObjectRejection::Killswitched);`
`60`	`61`	`}`
`61`	`62`
`62`	`63`	`if !state.rate_limiter.check(id.context()) {`
	`64`	`+ tracing::debug!("Request rejected due to rate limits");`
`63`	`65`	`return Err(ObjectRejection::RateLimited);`
`64`	`66`	`}`
`65`	`67`
`@@ -121,10 +123,12 @@ impl FromRequestParts<ServiceState> for Xt<ObjectContext> {`
`121`	`123`	`populate_sentry_context(&context);`
`122`	`124`
`123`	`125`	`if state.config.killswitches.matches(&context) {`
	`126`	`+ tracing::debug!("Request rejected due to killswitches");`
`124`	`127`	`return Err(ObjectRejection::Killswitched);`
`125`	`128`	`}`
`126`	`129`
`127`	`130`	`if !state.rate_limiter.check(&context) {`
	`131`	`+ tracing::debug!("Request rejected due to rate limits");`
`128`	`132`	`return Err(ObjectRejection::RateLimited);`
`129`	`133`	`}`
`130`	`134`