From 0a8f616565521677067a7dc35cb997a1c8ab33ed Mon Sep 17 00:00:00 2001 From: Jarvis Date: Tue, 4 Aug 2026 17:03:49 +0800 Subject: [PATCH 1/2] feat(ratelimit): scheduled suspension windows for rate limit policies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A RateLimitPolicy can now carry `schedules` — recurring wall-clock windows (weekly days or explicit dates, HH:MM range, IANA timezone) during which the policy is suspended: the quota gate skips it and enforcement resumes automatically when the window closes. Windows are a union; cross-midnight windows belong to their start day; the bucket key never changes, so suspension does not reset the surrounding window's counts. Both policy-iteration sites honor suspension (reserve_layers and the routing/ensemble reserve_model_only path). Evaluation converts UTC to the schedule's local wall clock (total conversion — DST-safe), and any malformed field fails toward enforcing. Rows without the field keep enforcing unchanged; cp-api omits the field when empty so schedule-less rows stay parseable by pre-`schedules` strict data planes. A policy that does carry schedules requires data planes on this version (older strict loaders drop the whole row). LiteLLM has no equivalent to compare against: its rate limits are inline key/team/user fields with no standalone policy object and no time-based scheduling. --- Cargo.lock | 35 ++ Cargo.toml | 1 + crates/aisix-core/Cargo.toml | 1 + .../src/models/rate_limit_policy.rs | 371 ++++++++++++++++++ crates/aisix-core/src/models/schema.rs | 18 +- crates/aisix-proxy/src/lib.rs | 156 ++++++++ crates/aisix-proxy/src/quota.rs | 13 +- .../resources/rate_limit_policy.schema.json | 76 ++++ tests/e2e/src/cases/ratelimit-e2e.test.ts | 148 +++++++ 9 files changed, 814 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b89e4d40..347f0b6e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -99,6 +99,7 @@ dependencies = [ "anyhow", "arc-swap", "chrono", + "chrono-tz", "config", "dashmap", "hex", @@ -1380,6 +1381,16 @@ dependencies = [ "windows-link", ] +[[package]] +name = "chrono-tz" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" +dependencies = [ + "chrono", + "phf", +] + [[package]] name = "clap" version = "4.6.1" @@ -3378,6 +3389,24 @@ dependencies = [ "indexmap 2.14.0", ] +[[package]] +name = "phf" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_shared" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06005508882fb681fd97892ecff4b7fd0fee13ef1aa569f8695dae7ab9099981" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project" version = "1.1.11" @@ -4573,6 +4602,12 @@ dependencies = [ "time", ] +[[package]] +name = "siphasher" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" + [[package]] name = "sketches-ddsketch" version = "0.3.1" diff --git a/Cargo.toml b/Cargo.toml index 4704ead9..83ec4804 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -112,6 +112,7 @@ yaml-rust2 = "0.8" # Time / IDs chrono = { version = "0.4", features = ["serde"] } +chrono-tz = "0.10" uuid = { version = "1.11", features = ["v4", "v5", "serde"] } hostname = "0.4" diff --git a/crates/aisix-core/Cargo.toml b/crates/aisix-core/Cargo.toml index 717095b9..4b0b51bc 100644 --- a/crates/aisix-core/Cargo.toml +++ b/crates/aisix-core/Cargo.toml @@ -21,6 +21,7 @@ schemars = { workspace = true, features = ["chrono"] } # ApiKey::expires_at parses RFC 3339 timestamps for key-expiry # enforcement in the proxy auth path (#933). chrono.workspace = true +chrono-tz.workspace = true once_cell.workspace = true # ApiKey::hash_bearer (prd-09a §9A.7B.4) is the canonical SHA-256 the # proxy uses to compare an incoming Authorization: Bearer diff --git a/crates/aisix-core/src/models/rate_limit_policy.rs b/crates/aisix-core/src/models/rate_limit_policy.rs index 4681622a..a5ce14e2 100644 --- a/crates/aisix-core/src/models/rate_limit_policy.rs +++ b/crates/aisix-core/src/models/rate_limit_policy.rs @@ -15,6 +15,7 @@ //! reserves under `policy:<scope>:<scope_ref>:<policy_id>` — with the //! member's `user_id` appended for the `team_member` scope. +use chrono::{DateTime, Datelike, NaiveDate, Timelike, Utc}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; @@ -74,6 +75,136 @@ impl std::fmt::Display for PolicyWindow { } } +/// Day-of-week selector for a [`PolicySchedule`], in the schedule's +/// timezone. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(rename_all = "lowercase")] +pub enum ScheduleWeekday { + Mon, + Tue, + Wed, + Thu, + Fri, + Sat, + Sun, +} + +impl ScheduleWeekday { + fn as_chrono(self) -> chrono::Weekday { + match self { + Self::Mon => chrono::Weekday::Mon, + Self::Tue => chrono::Weekday::Tue, + Self::Wed => chrono::Weekday::Wed, + Self::Thu => chrono::Weekday::Thu, + Self::Fri => chrono::Weekday::Fri, + Self::Sat => chrono::Weekday::Sat, + Self::Sun => chrono::Weekday::Sun, + } + } +} + +/// One recurring wall-clock window during which the owning policy is +/// suspended (not enforced). Days are selected by `days_of_week` OR by +/// an explicit `dates` list (exactly one selector; the JSON Schema's +/// injected `oneOf` enforces this — see +/// [`crate::models::schema::rate_limit_policy_root_schema`]), evaluated +/// in `timezone`. A window whose `end_time` ≤ `start_time` crosses +/// midnight and belongs to its **start** day: `days_of_week: [fri], +/// 22:00 → 09:00` covers Friday 22:00 through Saturday 09:00. +#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)] +pub struct PolicySchedule { + /// IANA timezone the window's wall-clock fields are interpreted in + /// (e.g. `Asia/Shanghai`). + #[schemars(length(min = 1))] + pub timezone: String, + /// Recurring weekly selector: the window opens on each listed day. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + #[schemars(length(min = 1))] + pub days_of_week: Vec<ScheduleWeekday>, + /// Explicit calendar dates (`YYYY-MM-DD`, in `timezone`) the window + /// opens on — for holidays and other irregular days. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + #[schemars(length(min = 1), inner(regex(pattern = r"^\d{4}-\d{2}-\d{2}$")))] + pub dates: Vec<String>, + /// Window start, `HH:MM` wall clock (inclusive). + #[schemars(regex(pattern = r"^([01]\d|2[0-3]):[0-5]\d$"))] + pub start_time: String, + /// Window end, `HH:MM` wall clock (exclusive); `24:00` = end of + /// day. End ≤ start crosses into the following day. + #[schemars(regex(pattern = r"^([01]\d|2[0-3]):[0-5]\d$|^24:00$"))] + pub end_time: String, +} + +/// Parse `HH:MM` to minutes since midnight. `24:00` (end-of-day, only +/// meaningful as an exclusive end bound) is admitted when `allow_2400`. +fn parse_hhmm(s: &str, allow_2400: bool) -> Option<u32> { + if allow_2400 && s == "24:00" { + return Some(24 * 60); + } + let (h, m) = s.split_once(':')?; + if h.len() != 2 || m.len() != 2 { + return None; + } + let h: u32 = h.parse().ok()?; + let m: u32 = m.parse().ok()?; + if h > 23 || m > 59 { + return None; + } + Some(h * 60 + m) +} + +impl PolicySchedule { + /// Whether the instant `now` falls inside this window. Evaluation + /// converts UTC → the schedule's local wall clock, which is total + /// (no DST ambiguity — that only exists in the local→UTC + /// direction). Any unparseable field makes the window non-matching, + /// so a malformed schedule fails toward *enforcing* the policy. + fn matches(&self, now: DateTime<Utc>) -> bool { + let Ok(tz) = self.timezone.parse::<chrono_tz::Tz>() else { + return false; + }; + let Some(start) = parse_hhmm(&self.start_time, false) else { + return false; + }; + let Some(end) = parse_hhmm(&self.end_time, true) else { + return false; + }; + let local = now.with_timezone(&tz); + let minute = local.time().hour() * 60 + local.time().minute(); + let today = local.date_naive(); + if start < end { + // Same-day window [start, end). + self.selects_day(today) && minute >= start && minute < end + } else if start > end { + // Cross-midnight window owned by the start day: + // [start, 24:00) on a selected day D, plus [00:00, end) on D+1. + (self.selects_day(today) && minute >= start) + || (today + .pred_opt() + .is_some_and(|yesterday| self.selects_day(yesterday)) + && minute < end) + } else { + // start == end: rejected at the write path; empty here. + false + } + } + + /// Whether `date` (local to the schedule's timezone) is selected. + /// `days_of_week` takes precedence if a row ever carries both + /// selectors (the schema forbids that shape). + fn selects_day(&self, date: NaiveDate) -> bool { + if !self.days_of_week.is_empty() { + return self + .days_of_week + .iter() + .any(|d| d.as_chrono() == date.weekday()); + } + self.dates + .iter() + .any(|s| NaiveDate::parse_from_str(s, "%Y-%m-%d") == Ok(date)) + } +} + #[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)] pub struct RateLimitPolicy { #[schemars(length(min = 1))] @@ -88,11 +219,29 @@ pub struct RateLimitPolicy { #[serde(default, skip_serializing_if = "Option::is_none")] #[schemars(range(min = 1))] pub max_tokens: Option<u64>, + /// Recurring windows during which this policy is suspended — the + /// quota gate skips it while `now` falls in any listed window and + /// enforcement resumes automatically afterwards (AISIX-Cloud#1104). + /// The counters' bucket key never changes, so suspension does not + /// reset the surrounding window's counts. Empty/absent = always + /// enforced. cp-api omits the field when empty, keeping + /// schedule-less rows parseable by pre-`schedules` strict data + /// planes. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub schedules: Vec<PolicySchedule>, #[serde(skip)] pub(crate) runtime_id: String, } +impl RateLimitPolicy { + /// Whether the policy is suspended at `now` — true when any + /// schedule window contains the instant (windows are a union). + pub fn suspended_at(&self, now: DateTime<Utc>) -> bool { + self.schedules.iter().any(|s| s.matches(now)) + } +} + /// The one cross-field invariant `schemars` can't derive: a policy must cap at /// least one of `max_requests` / `max_tokens`. Injected as a top-level `anyOf` /// by [`crate::models::schema::rate_limit_policy_root_schema`]. @@ -103,6 +252,17 @@ pub fn rate_limit_policy_any_of() -> Value { ]) } +/// The [`PolicySchedule`] day-selector XOR `schemars` can't derive: exactly +/// one of `days_of_week` / `dates` must be present. Injected into the +/// `PolicySchedule` definition by +/// [`crate::models::schema::rate_limit_policy_root_schema`]. +pub fn policy_schedule_one_of() -> Value { + json!([ + { "required": ["days_of_week"] }, + { "required": ["dates"] } + ]) +} + impl Resource for RateLimitPolicy { fn id(&self) -> &str { &self.runtime_id @@ -182,6 +342,217 @@ mod tests { assert_eq!(RateLimitPolicy::kind(), "rate_limit_policies"); } + // --- schedules (AISIX-Cloud#1104) -------------------------------- + + /// Parse an RFC3339 instant for schedule sweeps. + fn at(s: &str) -> DateTime<Utc> { + DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc) + } + + fn policy_with_schedules(schedules: Value) -> RateLimitPolicy { + serde_json::from_value(json!({ + "name": "sched", + "scope": "team", + "scope_ref": "t1", + "window": "minute", + "max_requests": 1, + "schedules": schedules + })) + .unwrap() + } + + #[test] + fn no_schedules_never_suspends_and_old_rows_deserialize() { + // Pre-`schedules` etcd rows carry no field — they must keep + // enforcing unchanged after upgrade. + let p: RateLimitPolicy = serde_json::from_str( + r#"{"name":"x","scope":"team","scope_ref":"t1","window":"minute","max_requests":1}"#, + ) + .unwrap(); + assert!(p.schedules.is_empty()); + assert!(!p.suspended_at(at("2026-08-08T02:00:00Z"))); + } + + #[test] + fn weekend_full_day_window_matches_in_its_timezone() { + // 2026-08-08 is a Saturday. Beijing Sat 10:00 = 02:00Z. + let p = policy_with_schedules(json!([{ + "timezone": "Asia/Shanghai", + "days_of_week": ["sat", "sun"], + "start_time": "00:00", + "end_time": "24:00" + }])); + assert!(p.suspended_at(at("2026-08-08T02:00:00Z")), "Sat 10:00 CST"); + // 23:59 local still inside a 24:00-bounded day. + assert!(p.suspended_at(at("2026-08-08T15:59:00Z")), "Sat 23:59 CST"); + assert!(!p.suspended_at(at("2026-08-10T02:00:00Z")), "Mon 10:00 CST"); + } + + #[test] + fn same_instant_differs_by_timezone() { + // 2026-08-08T02:00Z is Saturday 10:00 in Beijing but still + // Friday 22:00 in New York — the selector must follow the + // schedule's timezone, not UTC. + let entry = |tz: &str| { + policy_with_schedules(json!([{ + "timezone": tz, + "days_of_week": ["sat"], + "start_time": "00:00", + "end_time": "24:00" + }])) + }; + assert!(entry("Asia/Shanghai").suspended_at(at("2026-08-08T02:00:00Z"))); + assert!(!entry("America/New_York").suspended_at(at("2026-08-08T02:00:00Z"))); + } + + #[test] + fn cross_midnight_window_belongs_to_its_start_day() { + // Workday off-peak: Mon–Fri 22:00 → next morning 09:00 (CST). + let p = policy_with_schedules(json!([{ + "timezone": "Asia/Shanghai", + "days_of_week": ["mon", "tue", "wed", "thu", "fri"], + "start_time": "22:00", + "end_time": "09:00" + }])); + // 2026-08-07 is a Friday (CST = UTC+8). + assert!(!p.suspended_at(at("2026-08-07T13:59:00Z")), "Fri 21:59"); + assert!( + p.suspended_at(at("2026-08-07T14:00:00Z")), + "Fri 22:00 opens" + ); + assert!(p.suspended_at(at("2026-08-07T15:00:00Z")), "Fri 23:00"); + // Saturday morning is the spill-over of Friday's window... + assert!(p.suspended_at(at("2026-08-08T00:59:00Z")), "Sat 08:59"); + assert!( + !p.suspended_at(at("2026-08-08T01:00:00Z")), + "Sat 09:00 closes" + ); + // ...but Saturday evening opens nothing (sat not selected). + assert!(!p.suspended_at(at("2026-08-08T15:00:00Z")), "Sat 23:00"); + // And Monday 08:00 stays enforced: its night belongs to Sunday, + // which is not selected. (Cover Sunday nights by adding "sun".) + assert!(!p.suspended_at(at("2026-08-10T00:00:00Z")), "Mon 08:00"); + assert!(p.suspended_at(at("2026-08-10T14:30:00Z")), "Mon 22:30"); + assert!(p.suspended_at(at("2026-08-11T00:30:00Z")), "Tue 08:30"); + assert!( + !p.suspended_at(at("2026-08-11T01:00:00Z")), + "Tue 09:00 closes" + ); + } + + #[test] + fn dates_selector_matches_explicit_holidays() { + let p = policy_with_schedules(json!([{ + "timezone": "Asia/Shanghai", + "dates": ["2026-10-01", "2026-10-02"], + "start_time": "00:00", + "end_time": "24:00" + }])); + assert!(p.suspended_at(at("2026-10-01T04:00:00Z")), "Oct 1 noon CST"); + assert!(p.suspended_at(at("2026-10-02T04:00:00Z")), "Oct 2 noon CST"); + assert!( + !p.suspended_at(at("2026-10-03T04:00:00Z")), + "Oct 3 noon CST" + ); + // CST midnight boundary: Sep 30 16:00Z is Oct 1 00:00 CST. + assert!( + p.suspended_at(at("2026-09-30T16:00:00Z")), + "Oct 1 00:00 CST" + ); + assert!( + !p.suspended_at(at("2026-09-30T15:59:00Z")), + "Sep 30 23:59 CST" + ); + } + + #[test] + fn dated_cross_midnight_spills_into_the_next_day() { + let p = policy_with_schedules(json!([{ + "timezone": "Asia/Shanghai", + "dates": ["2026-10-01"], + "start_time": "22:00", + "end_time": "02:00" + }])); + assert!( + p.suspended_at(at("2026-10-01T15:30:00Z")), + "Oct 1 23:30 CST" + ); + assert!( + p.suspended_at(at("2026-10-01T17:00:00Z")), + "Oct 2 01:00 CST" + ); + assert!( + !p.suspended_at(at("2026-10-01T18:00:00Z")), + "Oct 2 02:00 CST" + ); + assert!( + !p.suspended_at(at("2026-10-02T15:30:00Z")), + "Oct 2 23:30 CST" + ); + } + + #[test] + fn schedule_union_any_match_suspends() { + let p = policy_with_schedules(json!([ + { + "timezone": "Asia/Shanghai", + "days_of_week": ["sat", "sun"], + "start_time": "00:00", + "end_time": "24:00" + }, + { + "timezone": "Asia/Shanghai", + "days_of_week": ["mon", "tue", "wed", "thu", "fri", "sun"], + "start_time": "22:00", + "end_time": "09:00" + } + ])); + // Monday 08:00 CST — covered by the second entry via Sunday. + assert!(p.suspended_at(at("2026-08-10T00:00:00Z"))); + // Saturday noon — covered by the first entry only. + assert!(p.suspended_at(at("2026-08-08T04:00:00Z"))); + // Monday noon — covered by neither. + assert!(!p.suspended_at(at("2026-08-10T04:00:00Z"))); + } + + #[test] + fn malformed_schedule_fields_fail_toward_enforcing() { + // Unknown timezone / bad times / start==end: the window simply + // never matches, so the policy keeps enforcing (conservative). + for sched in [ + json!([{"timezone": "Not/AZone", "days_of_week": ["sat"], + "start_time": "00:00", "end_time": "24:00"}]), + json!([{"timezone": "Asia/Shanghai", "days_of_week": ["sat"], + "start_time": "9:00", "end_time": "10:00"}]), + json!([{"timezone": "Asia/Shanghai", "days_of_week": ["sat"], + "start_time": "09:00", "end_time": "09:00"}]), + ] { + let p = policy_with_schedules(sched); + assert!(!p.suspended_at(at("2026-08-08T02:00:00Z"))); + } + } + + #[test] + fn empty_schedules_are_omitted_from_serialization() { + // cp-api relies on this shape: schedule-less rows stay + // byte-compatible with pre-`schedules` data planes. + let p: RateLimitPolicy = serde_json::from_str( + r#"{"name":"x","scope":"team","scope_ref":"t1","window":"minute","max_requests":1}"#, + ) + .unwrap(); + let out = serde_json::to_value(&p).unwrap(); + assert!(out.get("schedules").is_none()); + + let with = policy_with_schedules(json!([{ + "timezone": "Asia/Shanghai", + "days_of_week": ["sat"], + "start_time": "00:00", + "end_time": "24:00" + }])); + let out = serde_json::to_value(&with).unwrap(); + assert_eq!(out["schedules"][0]["timezone"], "Asia/Shanghai"); + } + #[test] fn resource_name_returns_scope_ref() { let mut p: RateLimitPolicy = serde_json::from_str( diff --git a/crates/aisix-core/src/models/schema.rs b/crates/aisix-core/src/models/schema.rs index 9779cfd5..936e6dad 100644 --- a/crates/aisix-core/src/models/schema.rs +++ b/crates/aisix-core/src/models/schema.rs @@ -1037,12 +1037,22 @@ fn branch_kind(branch: &serde_json::Map<String, Value>) -> Option<&str> { /// ([`super::rate_limit_policy::rate_limit_policy_any_of`]). pub fn rate_limit_policy_root_schema() -> Value { let mut schema = struct_root_schema::<crate::models::RateLimitPolicy>(false); - schema + let obj = schema .as_object_mut() - .expect("rate_limit_policy root schema is a JSON object") + .expect("rate_limit_policy root schema is a JSON object"); + obj.insert( + "anyOf".to_string(), + super::rate_limit_policy::rate_limit_policy_any_of(), + ); + // The schedule day-selector XOR is the same kind of cross-field + // invariant, one level down in the definitions. + obj.get_mut("definitions") + .and_then(|d| d.get_mut("PolicySchedule")) + .and_then(Value::as_object_mut) + .expect("rate_limit_policy schema defines PolicySchedule") .insert( - "anyOf".to_string(), - super::rate_limit_policy::rate_limit_policy_any_of(), + "oneOf".to_string(), + super::rate_limit_policy::policy_schedule_one_of(), ); schema } diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index d9b32898..c507e2b8 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -3314,6 +3314,162 @@ data: [DONE]\n\n" ); } + /// A schedule window that always matches "now" (every weekday, all + /// day) so suspension is deterministic under the system clock. + fn always_on_schedule() -> serde_json::Value { + serde_json::json!([{ + "timezone": "UTC", + "days_of_week": ["mon", "tue", "wed", "thu", "fri", "sat", "sun"], + "start_time": "00:00", + "end_time": "24:00" + }]) + } + + /// A schedule window that can never match (a fixed past date). + fn never_on_schedule() -> serde_json::Value { + serde_json::json!([{ + "timezone": "UTC", + "dates": ["2000-01-01"], + "start_time": "00:00", + "end_time": "24:00" + }]) + } + + /// While inside a scheduled suspension window the policy reserves + /// nothing; once the schedule no longer matches, enforcement resumes + /// on the unchanged bucket (AISIX-Cloud#1104). + #[tokio::test] + async fn scheduled_suspension_pauses_policy_until_window_closes() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/anything")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true}))) + .mount(&upstream) + .await; + + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(openai_test_bridge())); + let snap = new_snap(&upstream.uri()); + snap.models.insert(model_entry("my-video-model")); + let policy_json = |schedules: serde_json::Value| { + serde_json::json!({ + "name": "video-cap", + "scope": "model", + "scope_ref": "model-id-1", + "window": "minute", + "max_requests": 1, + "schedules": schedules + }) + }; + snap.rate_limit_policies.insert(ResourceEntry::new( + "pol-1", + serde_json::from_value(policy_json(always_on_schedule())).unwrap(), + 1, + )); + snap.apikeys + .insert(apikey_entry("sk-caller", &["my-video-model"])); + let state = build_state(snap, hub); + + let make_req = || { + Request::builder() + .method("POST") + .uri("/passthrough/openai/anything") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(r#"{"model":"my-video-model","input":"x"}"#)) + .unwrap() + }; + + // Suspended: max_requests=1 would reject the second call — both pass. + for _ in 0..2 { + let resp = run(build_router(state.clone()), make_req()).await; + assert_eq!( + resp.status(), + StatusCode::OK, + "suspended policy must not gate requests", + ); + } + + // Swap in a schedule that no longer matches (as the loader does + // when the window closes relative to a fresh evaluation). + state + .snapshot + .load() + .rate_limit_policies + .insert(ResourceEntry::new( + "pol-1", + serde_json::from_value(policy_json(never_on_schedule())).unwrap(), + 2, + )); + + let resp = run(build_router(state.clone()), make_req()).await; + assert_eq!(resp.status(), StatusCode::OK); + let resp = run(build_router(state.clone()), make_req()).await; + assert_eq!( + resp.status(), + StatusCode::TOO_MANY_REQUESTS, + "policy must enforce again outside its suspension windows", + ); + } + + /// The routing/ensemble per-target path (`reserve_model_only`) + /// iterates the policy table independently of `reserve_layers`, so + /// it must honor scheduled suspensions too (AISIX-Cloud#1104). + #[tokio::test] + async fn reserve_model_only_honors_scheduled_suspension() { + let hub = Arc::new(Hub::new()); + let snap = new_snap("http://127.0.0.1:1"); + let model = model_entry("mg-member"); + let target = model.value.clone(); + snap.models.insert(model); + let policy_json = |schedules: serde_json::Value| { + serde_json::json!({ + "name": "member-cap", + "scope": "model", + "scope_ref": "model-id-1", + "window": "minute", + "max_requests": 1, + "schedules": schedules + }) + }; + snap.rate_limit_policies.insert(ResourceEntry::new( + "pol-1", + serde_json::from_value(policy_json(always_on_schedule())).unwrap(), + 1, + )); + let state = build_state(snap, hub); + + // Suspended: max_requests=1 would deny the second reservation + // (pre_commit counts stick even when the reservation drops + // uncommitted) — both succeed because nothing is reserved. + for _ in 0..2 { + let r = quota::reserve_model_only(&state, "mg-member", "model-id-1", &target).await; + assert!(r.is_ok(), "suspended policy must reserve nothing"); + } + + state + .snapshot + .load() + .rate_limit_policies + .insert(ResourceEntry::new( + "pol-1", + serde_json::from_value(policy_json(never_on_schedule())).unwrap(), + 2, + )); + + assert!( + quota::reserve_model_only(&state, "mg-member", "model-id-1", &target) + .await + .is_ok() + ); + assert!( + quota::reserve_model_only(&state, "mg-member", "model-id-1", &target) + .await + .is_err(), + "policy outside its windows must throttle the second reservation", + ); + } + /// The tunnel forwards bodies verbatim, so callers typically name the /// provider-native id (`model_name`), not the gateway alias /// (`display_name`). The limit must bind either way, and the bucket is diff --git a/crates/aisix-proxy/src/quota.rs b/crates/aisix-proxy/src/quota.rs index a0668eb5..fea3c34f 100644 --- a/crates/aisix-proxy/src/quota.rs +++ b/crates/aisix-proxy/src/quota.rs @@ -177,8 +177,15 @@ async fn reserve_layers( // Layer 4+: Rate limit policies from snapshot. let snap = state.snapshot.load(); + let now = chrono::Utc::now(); for entry in snap.rate_limit_policies.entries() { let policy = &entry.value; + // Inside a scheduled suspension window the policy reserves + // nothing; enforcement resumes automatically when the window + // closes, on the unchanged bucket (AISIX-Cloud#1104). + if policy.suspended_at(now) { + continue; + } let applies = match policy.scope { PolicyScope::ApiKey => policy.scope_ref == auth.entry.id, PolicyScope::Model => model_rl.is_some_and(|m| policy.scope_ref == m.entry_id), @@ -316,9 +323,13 @@ pub(crate) async fn reserve_model_only( // `model`-scope rate-limit policies for this model. (model scope never // buckets per-user, so the base bucket key suffices — no auth needed.) let snap = state.snapshot.load(); + let now = chrono::Utc::now(); for entry in snap.rate_limit_policies.entries() { let policy = &entry.value; - if policy.scope != PolicyScope::Model || policy.scope_ref != model_entry_id { + if policy.scope != PolicyScope::Model + || policy.scope_ref != model_entry_id + || policy.suspended_at(now) + { continue; } let rl = policy_to_rate_limit(policy); diff --git a/schemas/resources/rate_limit_policy.schema.json b/schemas/resources/rate_limit_policy.schema.json index 57c9e425..485bce04 100644 --- a/schemas/resources/rate_limit_policy.schema.json +++ b/schemas/resources/rate_limit_policy.schema.json @@ -14,6 +14,62 @@ } ], "definitions": { + "PolicySchedule": { + "additionalProperties": false, + "description": "One recurring wall-clock window during which the owning policy is suspended (not enforced). Days are selected by `days_of_week` OR by an explicit `dates` list (exactly one selector; the JSON Schema's injected `oneOf` enforces this — see [`crate::models::schema::rate_limit_policy_root_schema`]), evaluated in `timezone`. A window whose `end_time` ≤ `start_time` crosses midnight and belongs to its **start** day: `days_of_week: [fri], 22:00 → 09:00` covers Friday 22:00 through Saturday 09:00.", + "oneOf": [ + { + "required": [ + "days_of_week" + ] + }, + { + "required": [ + "dates" + ] + } + ], + "properties": { + "dates": { + "description": "Explicit calendar dates (`YYYY-MM-DD`, in `timezone`) the window opens on — for holidays and other irregular days.", + "items": { + "pattern": "^\\d{4}-\\d{2}-\\d{2}$", + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + "days_of_week": { + "description": "Recurring weekly selector: the window opens on each listed day.", + "items": { + "$ref": "#/definitions/ScheduleWeekday" + }, + "minItems": 1, + "type": "array" + }, + "end_time": { + "description": "Window end, `HH:MM` wall clock (exclusive); `24:00` = end of day. End ≤ start crosses into the following day.", + "pattern": "^([01]\\d|2[0-3]):[0-5]\\d$|^24:00$", + "type": "string" + }, + "start_time": { + "description": "Window start, `HH:MM` wall clock (inclusive).", + "pattern": "^([01]\\d|2[0-3]):[0-5]\\d$", + "type": "string" + }, + "timezone": { + "description": "IANA timezone the window's wall-clock fields are interpreted in (e.g. `Asia/Shanghai`).", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "end_time", + "start_time", + "timezone" + ], + "type": "object" + }, "PolicyScope": { "description": "Subject a [`RateLimitPolicy`] targets, paired with `scope_ref`.", "enum": [ @@ -33,6 +89,19 @@ "hour" ], "type": "string" + }, + "ScheduleWeekday": { + "description": "Day-of-week selector for a [`PolicySchedule`], in the schedule's timezone.", + "enum": [ + "mon", + "tue", + "wed", + "thu", + "fri", + "sat", + "sun" + ], + "type": "string" } }, "properties": { @@ -50,6 +119,13 @@ "minLength": 1, "type": "string" }, + "schedules": { + "description": "Recurring windows during which this policy is suspended — the quota gate skips it while `now` falls in any listed window and enforcement resumes automatically afterwards (AISIX-Cloud#1104). The counters' bucket key never changes, so suspension does not reset the surrounding window's counts. Empty/absent = always enforced. cp-api omits the field when empty, keeping schedule-less rows parseable by pre-`schedules` strict data planes.", + "items": { + "$ref": "#/definitions/PolicySchedule" + }, + "type": "array" + }, "scope": { "$ref": "#/definitions/PolicyScope" }, diff --git a/tests/e2e/src/cases/ratelimit-e2e.test.ts b/tests/e2e/src/cases/ratelimit-e2e.test.ts index 12e1ae0f..44239a5c 100644 --- a/tests/e2e/src/cases/ratelimit-e2e.test.ts +++ b/tests/e2e/src/cases/ratelimit-e2e.test.ts @@ -139,3 +139,151 @@ describe("rate limit e2e: RPM=1 second call gets 429", () => { expect(retryAfterSeconds).toBeLessThanOrEqual(60); }); }); + +// E2E: RateLimitPolicy.schedules — recurring suspension windows +// (AISIX-Cloud#1104), driven through etcd exactly as cp-api writes +// them. Windows are picked relative to "now" so the test is +// deterministic without waiting for wall-clock boundaries: +// - an all-week 00:00–24:00 window is always active → suspended +// - a fixed past date (2000-01-01) never matches → enforced +// Also pins the upgrade contract: a pre-`schedules` row (field absent) +// enforces unchanged, and toggling schedules never touches the bucket, +// so the window's burned count survives a suspend/resume cycle. +const SCHED_CALLER = "sk-rlp-sched-e2e-caller"; +const SCHED_KEY_ID = "c0000000-0000-0000-0000-000000000011"; +const SCHED_POLICY_ID = "d0000000-0000-0000-0000-000000000011"; + +describe("rate limit policy schedules e2e (AISIX-Cloud#1104)", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let seed: SeedClient | undefined; + let etcdReachable = false; + + const policyDoc = ( + schedules?: Array<Record<string, unknown>>, + ): Record<string, unknown> => ({ + name: "sched-cap", + scope: "api_key", + scope_ref: SCHED_KEY_ID, + window: "minute", + max_requests: 1, + // cp-api omits `schedules` when empty so schedule-less rows stay + // parseable by pre-`schedules` strict data planes. + ...(schedules ? { schedules } : {}), + }); + + const alwaysOn = [ + { + timezone: "UTC", + days_of_week: ["mon", "tue", "wed", "thu", "fri", "sat", "sun"], + start_time: "00:00", + end_time: "24:00", + }, + ]; + const neverOn = [ + { + timezone: "UTC", + dates: ["2000-01-01"], + start_time: "00:00", + end_time: "24:00", + }, + ]; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + upstream = await startOpenAiUpstream(); + app = await spawnApp(); + seed = new SeedClient(etcd, app.etcdPrefix); + + const pk = await seed.createProviderKey({ + display_name: "rlp-sched-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: "rlp-sched", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + // api_key scope matches on the key's etcd entry id, so the key + // needs a fixed id — seed it straight to etcd. + await etcd.put( + `${app.etcdPrefix}/api_keys/${SCHED_KEY_ID}`, + JSON.stringify({ + key_hash: createHash("sha256").update(SCHED_CALLER).digest("hex"), + allowed_models: ["rlp-sched"], + }), + ); + // Start from the pre-`schedules` shape (field absent). + await seed.update("rate_limit_policies", SCHED_POLICY_ID, policyDoc()); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + test("suspension window pauses the policy; leaving it resumes the same bucket", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + + const probe = new ProxyClient(app.proxyUrl, SCHED_CALLER); + await waitConfigPropagation(async () => { + const res = await probe.listModels(); + if (res.status !== 200) return false; + const data = (res.body as { data?: Array<{ id?: string }> }).data ?? []; + return data.some((m) => m.id === "rlp-sched"); + }); + + const client = new OpenAI({ + apiKey: SCHED_CALLER, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + const callStatus = async (): Promise<number> => { + try { + await client.chat.completions.create({ + model: "rlp-sched", + messages: [{ role: "user", content: "hi" }], + }); + return 200; + } catch (e) { + if (e instanceof APIError) return e.status ?? -1; + throw e; + } + }; + + // Keep the burn → suspend → resume sequence inside one fixed + // minute window so the final 429 provably reuses the burned count. + await awaitWindowHeadroom(30); + + // Pre-`schedules` row enforces: first call burns the single slot. + expect(await callStatus()).toBe(200); + expect(await callStatus()).toBe(429); + + // Enter a suspension window → propagation lands when a call passes + // again. Suspended probes reserve nothing, so counts stay intact. + await seed.update( + "rate_limit_policies", + SCHED_POLICY_ID, + policyDoc(alwaysOn), + ); + await waitConfigPropagation(async () => (await callStatus()) === 200); + + // Leave the window (schedule no longer matches). The bucket still + // holds the burned slot from this minute, so enforcement resumes + // as 429 — suspension must not reset quotas. + await seed.update( + "rate_limit_policies", + SCHED_POLICY_ID, + policyDoc(neverOn), + ); + await waitConfigPropagation(async () => (await callStatus()) === 429); + }); +}); From bb052d9f6b4da0994dafcca8736f067feb4b1a91 Mon Sep 17 00:00:00 2001 From: Jarvis <jarvis@api7.ai> Date: Tue, 4 Aug 2026 17:26:42 +0800 Subject: [PATCH 2/2] docs(ratelimit): precise three-case schedule time-bound wording CodeRabbit review: start==end is an empty window, not a write-path rejection on the DP side (only cp-api validates the shape); spell out the three comparisons in the struct docs the schema description inherits. --- .../aisix-core/src/models/rate_limit_policy.rs | 16 +++++++++++----- schemas/resources/rate_limit_policy.schema.json | 4 ++-- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/crates/aisix-core/src/models/rate_limit_policy.rs b/crates/aisix-core/src/models/rate_limit_policy.rs index a5ce14e2..c2abe991 100644 --- a/crates/aisix-core/src/models/rate_limit_policy.rs +++ b/crates/aisix-core/src/models/rate_limit_policy.rs @@ -108,9 +108,12 @@ impl ScheduleWeekday { /// an explicit `dates` list (exactly one selector; the JSON Schema's /// injected `oneOf` enforces this — see /// [`crate::models::schema::rate_limit_policy_root_schema`]), evaluated -/// in `timezone`. A window whose `end_time` ≤ `start_time` crosses -/// midnight and belongs to its **start** day: `days_of_week: [fri], -/// 22:00 → 09:00` covers Friday 22:00 through Saturday 09:00. +/// in `timezone`. Time bounds compare as wall-clock minutes: +/// `start_time < end_time` is a same-day window; `start_time > +/// end_time` crosses midnight and belongs to its **start** day +/// (`days_of_week: [fri], 22:00 → 09:00` covers Friday 22:00 through +/// Saturday 09:00); equal times are an empty window that never +/// matches. #[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)] pub struct PolicySchedule { /// IANA timezone the window's wall-clock fields are interpreted in @@ -130,7 +133,8 @@ pub struct PolicySchedule { #[schemars(regex(pattern = r"^([01]\d|2[0-3]):[0-5]\d$"))] pub start_time: String, /// Window end, `HH:MM` wall clock (exclusive); `24:00` = end of - /// day. End ≤ start crosses into the following day. + /// day. An end before the start crosses into the following day; an + /// end equal to the start is an empty window that never matches. #[schemars(regex(pattern = r"^([01]\d|2[0-3]):[0-5]\d$|^24:00$"))] pub end_time: String, } @@ -184,7 +188,9 @@ impl PolicySchedule { .is_some_and(|yesterday| self.selects_day(yesterday)) && minute < end) } else { - // start == end: rejected at the write path; empty here. + // start == end: an empty window (cp-api rejects the shape; + // the DP's own write surface admits it and it matches + // nothing). false } } diff --git a/schemas/resources/rate_limit_policy.schema.json b/schemas/resources/rate_limit_policy.schema.json index 485bce04..a1d8a722 100644 --- a/schemas/resources/rate_limit_policy.schema.json +++ b/schemas/resources/rate_limit_policy.schema.json @@ -16,7 +16,7 @@ "definitions": { "PolicySchedule": { "additionalProperties": false, - "description": "One recurring wall-clock window during which the owning policy is suspended (not enforced). Days are selected by `days_of_week` OR by an explicit `dates` list (exactly one selector; the JSON Schema's injected `oneOf` enforces this — see [`crate::models::schema::rate_limit_policy_root_schema`]), evaluated in `timezone`. A window whose `end_time` ≤ `start_time` crosses midnight and belongs to its **start** day: `days_of_week: [fri], 22:00 → 09:00` covers Friday 22:00 through Saturday 09:00.", + "description": "One recurring wall-clock window during which the owning policy is suspended (not enforced). Days are selected by `days_of_week` OR by an explicit `dates` list (exactly one selector; the JSON Schema's injected `oneOf` enforces this — see [`crate::models::schema::rate_limit_policy_root_schema`]), evaluated in `timezone`. Time bounds compare as wall-clock minutes: `start_time < end_time` is a same-day window; `start_time > end_time` crosses midnight and belongs to its **start** day (`days_of_week: [fri], 22:00 → 09:00` covers Friday 22:00 through Saturday 09:00); equal times are an empty window that never matches.", "oneOf": [ { "required": [ @@ -48,7 +48,7 @@ "type": "array" }, "end_time": { - "description": "Window end, `HH:MM` wall clock (exclusive); `24:00` = end of day. End ≤ start crosses into the following day.", + "description": "Window end, `HH:MM` wall clock (exclusive); `24:00` = end of day. An end before the start crosses into the following day; an end equal to the start is an empty window that never matches.", "pattern": "^([01]\\d|2[0-3]):[0-5]\\d$|^24:00$", "type": "string" },