Queue requests on upstream 429; upstream usage limits

crates/tw-api/msg-codes.txt

@@ -80,6 +80,8 @@ config.key_limit_cache_reads
 config.key_limit_cost_too_small
 config.key_limit_duplicate
 config.key_limit_empty
+config.key_limit_every
+config.key_limit_every_calendar
 config.key_limit_not_positive
 config.key_limit_retention
 config.key_limit_two_measures
@@ -128,6 +130,15 @@ config.unknown_price_sheet
 config.unknown_rule
 config.unknown_variant
 config.unparsable passthrough
+config.upstream_limit_cache_reads
+config.upstream_limit_cost_too_small
+config.upstream_limit_duplicate
+config.upstream_limit_empty
+config.upstream_limit_every
+config.upstream_limit_every_calendar
+config.upstream_limit_not_positive
+config.upstream_limit_retention
+config.upstream_limit_two_measures
 config.zero_concurrency
 control.account_service_not_json
 control.account_service_refused
@@ -316,10 +327,13 @@ gw.count_tokens.bedrock_upstream
 gw.files.unsupported
 gw.internal
 gw.internal_panic
+gw.key_limit.cost_every
 gw.key_limit.cost_per_period
 gw.key_limit.cost_rolling
+gw.key_limit.requests_every
 gw.key_limit.requests_per_period
 gw.key_limit.requests_rolling
+gw.key_limit.tokens_every
 gw.key_limit.tokens_per_period
 gw.key_limit.tokens_rolling
 gw.listen.addr_unavailable
@@ -422,6 +436,8 @@ gw.upstream.eventstream_broken
 gw.upstream.forward_failed
 gw.upstream.idle_timeout
 gw.upstream.idle_timeout_mid_stream
+gw.upstream.paused
+gw.upstream.period_used_up
 gw.upstream.rate_limited
 gw.upstream.sign_failed
 gw.upstream.status
@@ -431,6 +447,11 @@ gw.upstream.stream_exception
 gw.upstream.stream_opening_error
 gw.upstream.timeout
 gw.upstream.unreachable
+gw.upstream.window_full
+gw.wait.mixed
+gw.wait.rate_limited
+gw.wait.used_up
+gw.wait.window_full
 gw.ws.alias_unserved
 gw.ws.alias_unserved_rewritten
 gw.ws.bad_header

crates/tw-config/src/failover.rs

@@ -1,4 +1,4 @@
-//! 故障转移：一家上游失败之后停用多久、多久没有内容就不再等它、上游满着时最多等多久。
+//! 故障转移：一家上游失败之后停用多久、多久没有内容就不再等它、一个请求最多等多久。
 
 use serde::{Deserialize, Serialize};
 
@@ -29,8 +29,9 @@ pub struct Failover {
     /// 上游说额度用完、又没说什么时候重置时，停用多少秒。说了的就停到那一刻
     #[serde(default = "d_quota_pause_secs")]
     pub quota_pause_secs: u64,
-    /// 限流时按上游给的 `Retry-After` 停用，**最多停多少秒**。没给就当作说不出
-    /// 原因的失败
+    /// 上游回 429（限流）之后这一家暂停多久，**最多多少秒**。上游说了多久的（`retry-after-ms`、
+    /// `Retry-After`、正文里的等待时间）按它说的，没说的按这一家连续被限流的次数退避：5、10、
+    /// 20、40 秒，之后每次 60 秒。**暂停不是停用**：请求排队等它，不进熔断的账
     #[serde(default = "d_rate_limit_max_pause_secs")]
     pub rate_limit_max_pause_secs: u64,
     /// 上游多少秒没有内容就不再等它（无响应超时）。从请求发给它的那一刻算起，**每来一段
@@ -43,11 +44,11 @@ pub struct Failover {
     #[serde(default = "d_idle_timeout_secs")]
     pub idle_timeout_secs: u64,
     /// 一个请求最多等多少秒，**整个请求合起来算**：准入时等密钥的分钟、小时上限空出名额，
-    /// 之后上游的并发数满了（`providers[].max_concurrent`）时等空位，共用这一段。留在那一家
-    /// 的对话等它空出来，候选都满了时等先空出来的那一家；等不到的换下一家，或者回 429。
-    /// 0 是不等
-    #[serde(default = "d_slot_wait_secs")]
-    pub slot_wait_secs: u64,
+    /// 之后等上游 —— 并发数满了（`providers[].max_concurrent`）、被限流暂停着、分钟或小时的
+    /// 用量上限满了 —— 共用这一段。留在那一家的对话等它，候选都等不上时等先能发的那一家；
+    /// 等不到的换下一家，或者回 429。0 是不等
+    #[serde(default = "d_wait_secs")]
+    pub wait_secs: u64,
 }
 
 fn d_failures_to_pause() -> u32 {
@@ -71,8 +72,8 @@ fn d_rate_limit_max_pause_secs() -> u64 {
 fn d_idle_timeout_secs() -> u64 {
     300
 }
-fn d_slot_wait_secs() -> u64 {
-    30
+fn d_wait_secs() -> u64 {
+    120
 }
 
 impl Default for Failover {
@@ -85,7 +86,7 @@ impl Default for Failover {
             quota_pause_secs: d_quota_pause_secs(),
             rate_limit_max_pause_secs: d_rate_limit_max_pause_secs(),
             idle_timeout_secs: d_idle_timeout_secs(),
-            slot_wait_secs: d_slot_wait_secs(),
+            wait_secs: d_wait_secs(),
         }
     }
 }
@@ -100,8 +101,9 @@ pub const MIN_IDLE_TIMEOUT_SECS: u64 = 30;
 /// 无响应超时最多写多少秒：一小时。再长就等于不设
 pub const MAX_IDLE_TIMEOUT_SECS: u64 = 3600;
 
-/// 等空位最多写多少秒。等的时候客户端一个字节都收不到，再长的话它先超时了
-pub const MAX_SLOT_WAIT_SECS: u64 = 300;
+/// 一个请求最多写等多少秒。流式的回答等过一阵先交出响应头和保活，整包的回答等的时候
+/// 客户端一个字节都收不到 —— 再长的话它先超时了
+pub const MAX_WAIT_SECS: u64 = 600;
 
 impl Failover {
     /// 不在允许范围里的第一项：字段名、写的值、下限、上限。
@@ -139,7 +141,7 @@ impl Failover {
                 MIN_IDLE_TIMEOUT_SECS,
                 MAX_IDLE_TIMEOUT_SECS,
             ),
-            ("slot_wait_secs", self.slot_wait_secs, 0, MAX_SLOT_WAIT_SECS),
+            ("wait_secs", self.wait_secs, 0, MAX_WAIT_SECS),
         ];
         fields
             .into_iter()

crates/tw-config/src/lib.rs

@@ -37,7 +37,7 @@ mod wire;
 pub use aliases::{Alias, Aliases};
 pub use credential::{CredentialError, Header, Headers, Secret, SecretResolveError, auth_header};
 pub use init::{generate_control_key, generate_initial, generate_key};
-pub use limits::{KeyLimit, LimitMeasure, LimitPer};
+pub use limits::{KeyLimit, LimitMeasure, LimitPer, UsageLimit};
 pub use model_specs::{ModelSpec, ResolvedSpec, Sourced, SpecSource};
 pub use plugins::Plugin;
 pub use proxy::{DIRECT, OnProxyFail, Proxy, ProxyKind, SYSTEM};
@@ -199,6 +199,7 @@ impl Default for Provider {
             pricing: None,
             model_specs: std::collections::BTreeMap::new(),
             max_concurrent: None,
+            limits: Vec::new(),
             disabled: false,
             balance: BalanceSetting::Auto,
             signed_in: None,
@@ -848,6 +849,14 @@ pub struct Provider {
     /// 这家的对话等它空出来，别的请求换下一家（见 `tw_gateway::slots`）
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub max_concurrent: Option<u32>,
+    /// 用量上限：每 N 分钟、每 N 小时、每天、每周、每月最多发给这家多少个请求、多少 token、
+    /// 花多少钱。**每一条都要过**，不写就不限。写法和密钥的一样（[`limits`]）。
+    ///
+    /// **照着这家自己的限额写**（中转站、订阅套餐的「每 5 小时多少次」）：用满了就不往这家
+    /// 发，而不是发过去挨一个 429。滚动窗口满了的，请求换下一家或者等它空出来；天、周、月
+    /// 用满了的，到下一期之前只换不等（见 `tw_gateway::upstream_limits`）
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub limits: Vec<UsageLimit>,
     /// 停用。**配置原样留着**：不参与路由，它的模型也不出现在
     /// `/v1/models` 里。要暂时不用一家上游时，比删掉再重新填一遍凭据好。
     #[serde(default, skip_serializing_if = "is_default")]
@@ -1300,7 +1309,7 @@ pub fn write(path: &Path, cfg: &Config) -> Result<(), WriteError> {
 }
 
 pub use failover::{
-    Failover, MAX_IDLE_TIMEOUT_SECS, MAX_PAUSE_SECS, MAX_SLOT_WAIT_SECS, MIN_IDLE_TIMEOUT_SECS,
+    Failover, MAX_IDLE_TIMEOUT_SECS, MAX_PAUSE_SECS, MAX_WAIT_SECS, MIN_IDLE_TIMEOUT_SECS,
 };
 pub use probes::{ClientProbes, ProbeAction};
 pub use reload::{Rejected, Stage, stand_in, try_parse};

crates/tw-config/src/limits.rs

@@ -1,10 +1,14 @@
-//! 一把网关密钥的用量上限：每分钟、每小时、每天、每周、每月最多多少个请求、多少 token、
-//! 花多少钱。
+//! 用量上限：每 N 分钟、每 N 小时、每天、每周、每月最多多少个请求、多少 token、花多少钱。
 //!
-//! **上限挂在密钥上**，不挂在上游上：要管住的是「这台机器上的这个脚本」，不是「这家
-//! 上游」。一把密钥可以有好几条，**每一条都要过**。
+//! **同一种写法挂在两处**，管的是两件事：
 //!
-//! 怎么数、怎么等、到了怎么拒，在 `tw_gateway::key_limits`；这里只有写法。
+//! - 网关密钥上（`clients[].limits`）：管住「这台机器上的这个脚本」。到了就拒绝这把密钥的
+//!   请求，见 `tw_gateway::key_limits`。
+//! - 上游上（`providers[].limits`）：照着那家上游自己的限额写（中转站、订阅套餐常有「每
+//!   5 小时多少次」）。到了就不往那家发，换下一家或者等它空出来，见
+//!   `tw_gateway::upstream_limits`。
+//!
+//! 一处可以有好几条，**每一条都要过**。这里只有写法。
 
 use serde::{Deserialize, Serialize};
 
@@ -32,14 +36,25 @@ impl LimitPer {
         }
     }
 
-    /// 滚动的那两种：窗口多长，毫秒。自然周期是 None
+    /// 滚动的那两种：一个单位多长，毫秒。自然周期是 None。一条上限的窗口还要乘上它的
+    /// `every`，见 [`KeyLimit::window_ms`]
     pub fn rolling_ms(&self) -> Option<i64> {
         match self {
             LimitPer::Minute => Some(60_000),
             LimitPer::Hour => Some(3_600_000),
             _ => None,
         }
     }
+
+    /// `every` 最多写多少：一小时以内的分钟、一天以内的小时。再长就是写天。自然周期不能写
+    /// `every`，是 None
+    pub fn every_max(&self) -> Option<u32> {
+        match self {
+            LimitPer::Minute => Some(60),
+            LimitPer::Hour => Some(24),
+            _ => None,
+        }
+    }
 }
 
 /// 一条上限数的是什么。
@@ -60,15 +75,23 @@ impl LimitMeasure {
     }
 }
 
-/// 一条上限：`{ per: day, cost: 5 }`。**一条只写一种量**（请求数、token、费用三选一），
-/// 校验时查。
+/// 一条上限：`{ per: day, cost: 5 }`、`{ per: hour, every: 5, requests: 300 }`。**一条只写
+/// 一种量**（请求数、token、费用三选一），校验时查。
 ///
 /// 三种量写成三个可选字段而不是 `{ measure: cost, max: 5 }`：手写的人读
 /// `requests: 30` 不用再对一遍哪个数是什么单位。
+///
+/// 名字里的 Key 是它先有的那一处；挂在上游上的是同一个类型（[`UsageLimit`]），写法、校验
+/// 一样
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(deny_unknown_fields)]
 pub struct KeyLimit {
     pub per: LimitPer,
+    /// 滚动窗口是几个单位：`per: hour, every: 5` 是最近 5 小时。**只有 `minute`（1 到 60）、
+    /// `hour`（1 到 24）能写**，不写是 1。天、周、月是自然的，没有「每 3 天」：三天从哪天
+    /// 算起说不清
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub every: Option<u32>,
     /// 请求数。**写成有符号的**：写了负数时报的是「要大于 0」，不是一句解析器的话
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub requests: Option<i64>,
@@ -84,7 +107,27 @@ pub struct KeyLimit {
     pub cache_reads: bool,
 }
 
+/// 挂在上游上的用量上限：和密钥的是同一种写法，换个名字读起来清楚
+pub type UsageLimit = KeyLimit;
+
 impl KeyLimit {
+    /// 滚动窗口是几个单位。不写是 1；自然周期一律是 1
+    pub fn every_or_one(&self) -> u32 {
+        match self.per.rolling_ms() {
+            Some(_) => self.every.unwrap_or(1),
+            None => 1,
+        }
+    }
+
+    /// 滚动窗口多长，毫秒：一个单位乘上 `every`。自然周期是 None。
+    ///
+    /// **数滚动窗口的地方一律用它**，不用 [`LimitPer::rolling_ms`]：那只是一个单位
+    pub fn window_ms(&self) -> Option<i64> {
+        self.per
+            .rolling_ms()
+            .map(|w| w.saturating_mul(i64::from(self.every_or_one().max(1))))
+    }
+
     /// 写了的那几种量。校验要求恰好一种
     pub fn measures(&self) -> Vec<LimitMeasure> {
         let mut out = Vec::new();
@@ -117,9 +160,11 @@ impl KeyLimit {
         }
     }
 
-    /// 两条算不算同一条：同一段时间、同一种量、缓存读算不算进去也一样
+    /// 两条算不算同一条：同一段时间（`every` 也一样）、同一种量、缓存读算不算进去也一样。
+    /// 每 1 小时和每 5 小时是两条：管的是两件事
     pub fn same_as(&self, other: &KeyLimit) -> bool {
         self.per == other.per
+            && self.every_or_one() == other.every_or_one()
             && self.measure() == other.measure()
             && self.cache_reads == other.cache_reads
     }
@@ -176,6 +221,7 @@ mod tests {
         // 默认值不写回去
         let out = serde_yaml_ng::to_string(&KeyLimit {
             per: LimitPer::Hour,
+            every: None,
             requests: Some(3),
             tokens: None,
             cost: None,
@@ -185,6 +231,32 @@ mod tests {
         assert_eq!(out.trim(), "per: hour\nrequests: 3");
     }
 
+    #[test]
+    fn every_multiplies_the_rolling_window() {
+        let l: KeyLimit = serde_yaml_ng::from_str("{per: hour, every: 5, requests: 300}").unwrap();
+        assert_eq!(l.every, Some(5));
+        assert_eq!(l.window_ms(), Some(5 * 3_600_000));
+        let l: KeyLimit = serde_yaml_ng::from_str("{per: minute, requests: 30}").unwrap();
+        assert_eq!((l.every_or_one(), l.window_ms()), (1, Some(60_000)));
+        // 自然周期没有窗口；写了 every 的由校验拒，这里不当真
+        let l: KeyLimit = serde_yaml_ng::from_str("{per: day, every: 3, requests: 1}").unwrap();
+        assert_eq!((l.every_or_one(), l.window_ms()), (1, None));
+        // 每 1 小时和每 5 小时是两条；写明的 1 和不写是同一条
+        let one: KeyLimit = serde_yaml_ng::from_str("{per: hour, requests: 3}").unwrap();
+        let five: KeyLimit = serde_yaml_ng::from_str("{per: hour, every: 5, requests: 9}").unwrap();
+        let also_one: KeyLimit =
+            serde_yaml_ng::from_str("{per: hour, every: 1, requests: 9}").unwrap();
+        assert!(!one.same_as(&five));
+        assert!(one.same_as(&also_one));
+        // 不写不写回去，写了照写
+        assert!(!serde_yaml_ng::to_string(&one).unwrap().contains("every"));
+        assert!(
+            serde_yaml_ng::to_string(&five)
+                .unwrap()
+                .contains("every: 5")
+        );
+    }
+
     #[test]
     fn a_typo_is_refused_rather_than_ignored() {
         assert!(serde_yaml_ng::from_str::<KeyLimit>("{per: day, request: 3}").is_err());

crates/tw-config/src/repair.rs

@@ -279,6 +279,26 @@ mod tests {
         assert_eq!(repair(&only).unwrap().text, GOOD);
     }
 
+    /// `slot_wait_secs` 改名成了 `wait_secs`（不只等空位了，也等限流暂停、用量上限）：写着旧名的
+    /// 老配置读不进来，一键修复删掉它。**旧值不搬过去**：等待照新的默认值走
+    #[test]
+    fn the_renamed_slot_wait_is_offered_for_deletion_without_carrying_its_value() {
+        let bad = format!("{GOOD}failover:\n  pause_secs: 90\n  slot_wait_secs: 45\n");
+        let r = repair(&bad).expect("该修得了");
+        let got: Vec<_> = r
+            .fixes
+            .iter()
+            .map(|f| (f.kind, f.field.as_str(), f.value.as_deref()))
+            .collect();
+        assert_eq!(
+            got,
+            vec![(FixKind::UnknownField, "failover.slot_wait_secs", Some("45"))]
+        );
+        assert_eq!(r.text, format!("{GOOD}failover:\n  pause_secs: 90\n"));
+        let cfg = try_parse(&r.text).unwrap();
+        assert_eq!(cfg.failover.wait_secs, 120);
+    }
+
     #[test]
     fn a_misspelled_field_in_a_list_item_is_removed() {
         let bad = "version: 1\nlisten:\n  control:\n    key: c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00\nclients:\n  - name: c\n    key: tw-k\n    colour: red\n";