diff --git a/crates/tw-api/msg-codes.txt b/crates/tw-api/msg-codes.txt index bbd8d2d5..ce2d8f5d 100644 --- a/crates/tw-api/msg-codes.txt +++ b/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 diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index f3b459ab..4e8a19d3 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -192,6 +192,37 @@ slug_enum! { } } +slug_enum! { + /// 一家上游为什么暂时发不出去([`Throttle::kind`])。 + pub enum ThrottleKind { + /// 上游回了 429(不是额度用完):这家暂停着,到 `until_ms` 为止。请求排队等它 + RateLimited = "rate_limited", + /// 这家一条分钟、小时的用量上限([`ProviderView::limits`])满了:最近的用量滑出窗口 + /// 之后又能发,`until_ms` 是算出来的那一刻。请求排队等它 + WindowFull = "window_full", + /// 这家一条天、周、月的用量上限用完了:到下一期开始(`until_ms`)之前跳过它,不等 + PeriodUsedUp = "period_used_up", + } +} + +/// 一家上游暂时发不出去的状态([`ProviderView::throttle`]、[`Event::ThrottleChanged`])。 +/// +/// **不是失败**:上游没坏,不进熔断的账,也不算进成功率。几种同时成立时给 `until_ms` 最晚的 +/// 那一种 —— 到那一刻之前它都发不出去。只在内存里:core 重启之后限流的暂停没有了,滚动窗口 +/// 从空的开始,天、周、月照请求记录算回来。 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct Throttle { + pub kind: ThrottleKind, + /// 什么时候又能发(Unix 毫秒)。限流是暂停结束的时刻,`window_full` 是算出来的那一刻, + /// `period_used_up` 是下一期开始的时刻 + pub until_ms: u64, + /// 是哪一条用量上限:[`ProviderView::limits`] 里的下标。`window_full` / `period_used_up` + /// 才有,限流没有 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub limit: Option, +} + slug_enum! { /// 熔断器的状态变成了什么。 pub enum BreakerState { @@ -492,6 +523,15 @@ slug_enum! { /// 它的并发数满了(`max_concurrent`):这一跳没有发出去,换了下一家。只在尝试链里 /// 出现([`AttemptView::skipped`]) Busy = "busy", + /// 它回过 429,正暂停着([`ThrottleKind::RateLimited`]):这一跳没有发出去,换了下一家。 + /// 只在尝试链里出现 + Throttled = "throttled", + /// 它一条分钟、小时的用量上限满了([`ThrottleKind::WindowFull`]):这一跳没有发出去,换了 + /// 下一家。只在尝试链里出现 + WindowFull = "window_full", + /// 它一条天、周、月的用量上限用完了([`ThrottleKind::PeriodUsedUp`]):跳过,不等。只在 + /// 尝试链里出现 + PeriodUsedUp = "period_used_up", } } @@ -1004,7 +1044,24 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// 版本时配置的来源是 `sync`,改了网关密钥时发 `GatewayKeysChanged`。[`ProviderView`] 多了 /// `sync_excluded`(不同步的上游)和 `sync_url_changed`(地址经同步变更,改之前的地址)。 /// 同步的码是 `sync.*`。安全模式下不自动同步(`paused`) -pub const CONTROL_API_VERSION: u32 = 50; +/// +/// **51 起上游被限流时请求排队,上游有了用量上限**:[`FailoverView`] 的 `slot_wait_secs` 改名 +/// `wait_secs`(一个请求合计最多等多久,默认 120、最多 600;等空位、等限流暂停、等用量窗口都 +/// 算在里面)。[`ProviderView`] 多了 `throttle`([`Throttle`]:`kind` 是 [`ThrottleKind`] 的 +/// `rate_limited` / `window_full` / `period_used_up`,`until_ms` 是又能发的时刻,`limit` 是 +/// 满了的那一条在 `limits` 里的下标;几样都有时报恢复得最晚的那一样;能发时没有)和 `limits` +/// ([`KeyLimitView`] 的列表,和密钥的一样带着此刻用了多少);[`ProviderInput`] 多了 `limits` +/// ([`KeyLimitInput`],整份替换)。[`KeyLimitView`]、[`KeyLimitInput`] 多了 `every`(滚动窗口 +/// 是几个 `per`:分钟 1 到 60、小时 1 到 24,天、周、月是 1)。新事件 +/// [`Event::ThrottleChanged`](一家上游开始、结束、换了暂时发不出去的样子,`throttle` 为空是 +/// 又能发了;不进通知)、[`Event::RequestQueued`](一个请求开始排队:在等哪几家、最多等到 +/// 什么时候)和 [`Event::RequestDequeued`](它不排队了:等到了、被中止了、等不到了),后两条也在 +/// `InFlightRequest::events` 里。[`ServeSkip`] 多了 `throttled`、`window_full`、`period_used_up`,[`AttemptView`] 多了 `pause_ms`(这一跳被限流,这家因此暂停了多久); +/// 带着 `pause_ms` 的、回了 429 的一跳不算进上游的用量上限([`AttemptView::counts_for_upstream`])。 +/// 新码:`config.key_limit_every*`、`config.upstream_limit_*`、`gw.key_limit.*_every`、 +/// `gw.upstream.paused`、`gw.upstream.window_full`、`gw.upstream.period_used_up`、`gw.wait.*`。 +/// 照 50 写的界面读不到 `wait_secs`,保存密钥的上限时没有 `every` 会被拒 +pub const CONTROL_API_VERSION: u32 = 51; /// `/live` 的速度看最近多久(见 [`LiveView::tokens_per_sec`]),毫秒 pub const LIVE_RATE_WINDOW_MS: u64 = 3_000; @@ -1294,6 +1351,26 @@ pub enum Event { #[serde(default, skip_serializing_if = "Option::is_none")] answered_model: Option, }, + /// 这个请求开始排队了:它能去的上游此刻都发不出去(满着、限流暂停着、用量上限满了), + /// 它在等其中先能发的那一家。留在一家的对话等那一家时也发。 + /// + /// **一次排队发一条**,不排了发一条 [`Event::RequestDequeued`]:一个请求可以排好几次(发出去 + /// 又被限流了,回来接着等)。在 [`InFlightRequest::events`] 里,半路才来听的一方照样看得出它 + /// 在排队。 + RequestQueued { + id: u64, + /// 在等哪几家,按候选的顺序 + upstreams: Vec, + /// 最多等到什么时候(Unix 毫秒):这个请求的等待期限(`failover.wait_secs`) + until_ms: u64, + at_ms: u64, + }, + /// 这个请求不排队了:等到了一家、要发出去了,被手动中止了,或者等不到了。和它之前的那条 + /// [`Event::RequestQueued`] 成对;路由事件要等整条尝试链走完才发,界面靠这一条知道它不在 + /// 排队了。也在 [`InFlightRequest::events`] 里。 + /// + /// 客户端在排队时走了的不发:紧跟着的结局(取消)就是它的结束 + RequestDequeued { id: u64, at_ms: u64 }, /// 路由决定完了,尝试链也走完了。 /// /// **单独一个事件,因为成功和失败两条路都要发它。**挂在 @@ -1534,6 +1611,18 @@ pub enum Event { state: BreakerState, at_ms: u64, }, + /// 某家上游暂时发不出去了(限流暂停、用量上限满了或用完了)、换了一种、又能发了。 + /// + /// 和 [`ProviderView::throttle`] 说的是同一件事:一个是现状,一个是变化。**不挂在请求上**, + /// 号是新取的。它自己会好:界面显示它,不发通知。 + ThrottleChanged { + id: u64, + provider: String, + /// 现在的样子。又能发了是空 + #[serde(default, skip_serializing_if = "Option::is_none")] + throttle: Option, + at_ms: u64, + }, /// 某家上游的模型清单开始获取了,或者获取完了(列出来了、没列出来、 /// 没问到)。 /// @@ -1814,14 +1903,21 @@ pub struct AttemptView { /// 出来的(请求解不开)没有。别的结果都没有:接下请求的那一跳的用量在结局里 #[serde(default, skip_serializing_if = "Option::is_none")] pub usage: Option, - /// 这一跳等了多少毫秒才轮到一个空位:这家设了 `max_concurrent` 而它满着。不算在 `ms` - /// 里。没等的没有 + /// 这一跳等了多少毫秒才轮到:这家满着(`max_concurrent`)、限流暂停着、分钟或小时的用量 + /// 上限满了。不算在 `ms` 里。没等的没有 #[serde(default, skip_serializing_if = "Option::is_none")] pub queued_ms: Option, - /// 这一跳为什么没发出去:`busy`(这家满着,换了下一家;等过它的话 `queued_ms` 是等了 - /// 多久)。这时 `outcome` 是 `error`,`error` 是同一件事的那句话。发出去了的没有 + /// 这一跳为什么没发出去:`busy`(这家满着)、`throttled`(限流暂停着)、`window_full`(分钟、 + /// 小时的用量上限满了),换了下一家或者等着它,等过它的话 `queued_ms` 是等了多久; + /// `period_used_up`(天、周、月的用量上限用完了),跳过。这时 `outcome` 是 `error`,`error` + /// 是同一件事的那句话。发出去了的没有 #[serde(default, skip_serializing_if = "Option::is_none")] pub skipped: Option, + /// 这一跳收到了 429(或者上游在回答的第一段内容之前报了限流),这家因此暂停了多少毫秒(见 + /// [`ThrottleKind::RateLimited`])。请求接着换下一家或者等这家,客户端察觉不到。有它的一跳 + /// 不算进这家的用量上限([`Self::counts_for_upstream`])。别的结果没有 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub pause_ms: Option, /// 这一跳走的代理:这家上游配置的代理名(跟随系统的是 `system`),直连是 None。没发出去 /// 的一跳(满着、被拒、转换不了……)是它要走的那个 —— 连不上的那一跳,问题可能就出在 /// 这个代理上。WebSocket 只连直连的上游,一律是 None @@ -1867,7 +1963,37 @@ impl AttemptView { } } +impl AttemptView { + /// 这一跳算进这家上游的用量上限(见 `tw_gateway::upstream_limits`):发到了上游 + /// ([`Self::sent`]),上游没回 429,也没在回答开头报限流([`Self::pause_ms`])。被限流的 + /// 那一跳上游没接,请求另找一家或者等它,不该把这家的上限用掉。**网关在内存里也是这么数的**: + /// 这几跳发出时占的当场还回去,和重启后从记录里加回来的是同一个数 + pub fn counts_for_upstream(&self) -> bool { + self.sent() && self.status != Some(429) && self.pause_ms.is_none() + } +} + impl RoutingView { + /// 算进上游用量上限的那几跳各是哪一家,按先后([`AttemptView::counts_for_upstream`])。 + /// 同一家可以出现几次:每一跳算一个请求 + pub fn upstream_hops(&self) -> Vec { + self.attempts + .iter() + .filter(|a| a.counts_for_upstream()) + .map(|a| a.provider.clone()) + .collect() + } + + /// 这个请求的 token、费用算在哪一家的用量上限上:链的最后一跳,它算进上限的话 + /// ([`AttemptView::counts_for_upstream`])。这一行的用量就是它的:转移在第一次成功时就 + /// 停,前面那几跳没有用量 + pub fn served_by(&self) -> Option<&str> { + self.attempts + .last() + .filter(|a| a.counts_for_upstream()) + .map(|a| a.provider.as_str()) + } + /// 这个请求可能发到了上游:尝试链上有一跳发出去了([`AttemptView::sent`])。**尝试链是空的 /// 时看它怎么收场**(`failed`):失败的是网关在发往哪一家之前就拒了(规则、内容过滤、 /// 用量上限);没失败的(客户端走了)是还没等到路由事件 —— 那时请求可能正在上游那里, @@ -2026,6 +2152,9 @@ impl Event { | Event::ContentMatched { id, .. } | Event::RequestPriced { id, .. } | Event::HealthChanged { id, .. } + | Event::ThrottleChanged { id, .. } + | Event::RequestQueued { id, .. } + | Event::RequestDequeued { id, .. } | Event::ModelsChanged { id, .. } | Event::ProxyChanged { id, .. } | Event::AuthChanged { id, .. } @@ -2295,9 +2424,9 @@ pub struct FailoverView { /// 上游多少秒没有内容就不再等它(无响应超时):从请求发出去起算,每来一段内容重新 /// 计时。客户端还什么都没收到时换下一家,已经收到一部分的报错收尾 pub idle_timeout_secs: u64, - /// 一个请求合计最多等多少秒:等密钥的分钟、小时上限空出名额,和等满着(`max_concurrent`) - /// 的上游空出位置,共用这一段。0 是不等 - pub slot_wait_secs: u64, + /// 一个请求合计最多等多少秒:等密钥的分钟、小时上限空出名额,和等暂时发不出去的上游(满着、 + /// 限流暂停着、用量上限满了),共用这一段。0 是不等 + pub wait_secs: u64, } /// 每项防护各在哪一档:`off` / `observe` / `enforce`。 @@ -2374,6 +2503,10 @@ pub struct ProviderView { pub disabled: bool, /// `ok` / `open`(熔断中) pub health: Health, + /// 这家暂时发不出去:限流暂停着、用量上限满了或用完了(见 [`Throttle`])。能发是空。变化报 + /// [`Event::ThrottleChanged`] + #[serde(default, skip_serializing_if = "Option::is_none")] + pub throttle: Option, /// 上游拒绝了凭据:最近一次得到答复的请求回的是这个状态码(401 / 403)。 /// 没被拒是空的,被拒之后有请求成功了也是空的。 /// @@ -2394,6 +2527,8 @@ pub struct ProviderView { /// 同时最多发给这家几个请求。不限是空 #[serde(default, skip_serializing_if = "Option::is_none")] pub max_concurrent: Option, + /// 用量上限,按配置里的顺序,各带此刻用了多少(发出去还没结束的按估算算)。没设的是空的 + pub limits: Vec, /// 配置里的 `balance:`,原样:编辑对话框保存时交回去([`ProviderInput::balance`]) pub balance_setting: BalanceSetting, /// 最近一次读到的余额。**是空的:这家没有余额可读**(关掉了、账号上游和 Bedrock、 @@ -2697,11 +2832,14 @@ slug_enum! { } } -/// 一条用量上限,和它此刻用了多少。 +/// 一条用量上限,和它此刻用了多少。密钥的([`ClientView::limits`])和上游的 +/// ([`ProviderView::limits`])都是它。 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] pub struct KeyLimitView { pub per: LimitPer, + /// 滚动窗口是几个 `per`:`hour` 的 5 是最近 5 小时。没写的是 1;天、周、月一律是 1 + pub every: u32, pub measure: LimitMeasure, /// 上限:请求数、token 数,费用是微分 pub max: u64, @@ -2718,11 +2856,13 @@ pub struct KeyLimitView { pub reached: bool, } -/// 新建、保存密钥时的一条用量上限。 +/// 新建、保存密钥(和上游)时的一条用量上限。 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] pub struct KeyLimitInput { pub per: LimitPer, + /// 滚动窗口是几个 `per`:`minute` 1 到 60、`hour` 1 到 24。天、周、月交 1 + pub every: u32, pub measure: LimitMeasure, /// 上限:请求数、token 数,费用是微分。要大于 0 pub max: u64, @@ -3354,6 +3494,9 @@ pub struct ProviderInput { /// 同时最多发给这家几个请求,1 到 1000。不给就是不限 #[serde(default, skip_serializing_if = "Option::is_none")] pub max_concurrent: Option, + /// 用量上限,整份替换。不给 = 一条都没有 + #[serde(default)] + pub limits: Vec, /// 停用 #[serde(default)] pub disabled: bool, @@ -6688,10 +6831,41 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, } } + /// 算进上游用量上限的那几跳:发到了、没被限流。回了 429 的、在回答开头报了限流的 + /// (带着 `pause_ms`)上游没接,不算 —— 网关在内存里也是这么数的 + #[test] + fn a_rate_limited_hop_does_not_count_toward_the_upstreams_limit() { + use AttemptOutcome::*; + let served = attempt(Served, Some(200), None); + assert!(served.counts_for_upstream()); + let mut paused = attempt(Status, Some(429), None); + paused.pause_ms = Some(5_000); + assert!(!paused.counts_for_upstream()); + // 额度用完的 429 不暂停,也不算 + assert!(!attempt(Status, Some(429), None).counts_for_upstream()); + let mut opening = attempt(Error, None, Some("gw.upstream.stream_opening_error")); + assert!(opening.counts_for_upstream()); + opening.pause_ms = Some(5_000); + assert!(!opening.counts_for_upstream()); + let mut turn = attempt(Served, Some(200), None); + turn.pause_ms = Some(5_000); + assert!(!turn.counts_for_upstream()); + + let mut b = served.clone(); + b.provider = "b".into(); + let routing = RoutingView { + attempts: vec![paused, opening, b], + ..Default::default() + }; + assert_eq!(routing.upstream_hops(), vec!["b".to_string()]); + assert_eq!(routing.served_by(), Some("b")); + } + /// 发到了上游的一跳:上游回了话,或者发出去之后才失败。没发出去的(满着、被拒、连不上) /// 不算。认的码都是 core 真发得出的 #[test] diff --git a/crates/tw-config/src/failover.rs b/crates/tw-config/src/failover.rs index 23dbfcd5..9175bfd9 100644 --- a/crates/tw-config/src/failover.rs +++ b/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() diff --git a/crates/tw-config/src/lib.rs b/crates/tw-config/src/lib.rs index ec17e088..c56c477f 100644 --- a/crates/tw-config/src/lib.rs +++ b/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, + /// 用量上限:每 N 分钟、每 N 小时、每天、每周、每月最多发给这家多少个请求、多少 token、 + /// 花多少钱。**每一条都要过**,不写就不限。写法和密钥的一样([`limits`])。 + /// + /// **照着这家自己的限额写**(中转站、订阅套餐的「每 5 小时多少次」):用满了就不往这家 + /// 发,而不是发过去挨一个 429。滚动窗口满了的,请求换下一家或者等它空出来;天、周、月 + /// 用满了的,到下一期之前只换不等(见 `tw_gateway::upstream_limits`) + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub limits: Vec, /// 停用。**配置原样留着**:不参与路由,它的模型也不出现在 /// `/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}; diff --git a/crates/tw-config/src/limits.rs b/crates/tw-config/src/limits.rs index 9ae5492e..916e287b 100644 --- a/crates/tw-config/src/limits.rs +++ b/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,7 +36,8 @@ impl LimitPer { } } - /// 滚动的那两种:窗口多长,毫秒。自然周期是 None + /// 滚动的那两种:一个单位多长,毫秒。自然周期是 None。一条上限的窗口还要乘上它的 + /// `every`,见 [`KeyLimit::window_ms`] pub fn rolling_ms(&self) -> Option { match self { LimitPer::Minute => Some(60_000), @@ -40,6 +45,16 @@ impl LimitPer { _ => None, } } + + /// `every` 最多写多少:一小时以内的分钟、一天以内的小时。再长就是写天。自然周期不能写 + /// `every`,是 None + pub fn every_max(&self) -> Option { + 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, /// 请求数。**写成有符号的**:写了负数时报的是「要大于 0」,不是一句解析器的话 #[serde(default, skip_serializing_if = "Option::is_none")] pub requests: Option, @@ -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 { + self.per + .rolling_ms() + .map(|w| w.saturating_mul(i64::from(self.every_or_one().max(1)))) + } + /// 写了的那几种量。校验要求恰好一种 pub fn measures(&self) -> Vec { 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::("{per: day, request: 3}").is_err()); diff --git a/crates/tw-config/src/repair.rs b/crates/tw-config/src/repair.rs index d7d9dd1b..d5667819 100644 --- a/crates/tw-config/src/repair.rs +++ b/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"; diff --git a/crates/tw-config/src/validate.rs b/crates/tw-config/src/validate.rs index c8117b48..33721d6a 100644 --- a/crates/tw-config/src/validate.rs +++ b/crates/tw-config/src/validate.rs @@ -164,6 +164,63 @@ pub enum ValidationError { days: u64, min: u64, }, + /// `every` 超出了这种时间单位能写的范围(分钟 1 到 60、小时 1 到 24) + #[error("{}", self.msg())] + KeyLimitEvery { + key: String, + per: &'static str, + every: u32, + max: u32, + }, + /// 天、周、月写了 `every`:自然周期没有「每 3 天」 + #[error("{}", self.msg())] + KeyLimitEveryCalendar { key: String, per: &'static str }, + // 上游的用量上限(`providers[].limits`):和密钥的一条一条对应,码另起一套 —— 句子里 + // 说的是上游,不是密钥 + #[error("{}", self.msg())] + UpstreamLimitEmpty { upstream: String }, + #[error("{}", self.msg())] + UpstreamLimitTwoMeasures { upstream: String, measures: String }, + #[error("{}", self.msg())] + UpstreamLimitNotPositive { + upstream: String, + per: &'static str, + measure: &'static str, + value: String, + }, + #[error("{}", self.msg())] + UpstreamLimitDuplicate { + upstream: String, + per: &'static str, + measure: &'static str, + }, + #[error("{}", self.msg())] + UpstreamLimitCacheReads { + upstream: String, + measure: &'static str, + }, + #[error("{}", self.msg())] + UpstreamLimitCostTooSmall { + upstream: String, + per: &'static str, + value: String, + }, + #[error("{}", self.msg())] + UpstreamLimitRetention { + upstream: String, + per: &'static str, + days: u64, + min: u64, + }, + #[error("{}", self.msg())] + UpstreamLimitEvery { + upstream: String, + per: &'static str, + every: u32, + max: u32, + }, + #[error("{}", self.msg())] + UpstreamLimitEveryCalendar { upstream: String, per: &'static str }, /// `sync` 一节写得不对([`crate::sync::SyncError`]) #[error("{}", self.msg())] Sync(crate::sync::SyncError), @@ -464,6 +521,94 @@ impl ValidationError { After a restart the total for the {per} is added up again from the request \ records, so row_days has to be at least {min}" ), + KeyLimitEvery { + key, + per, + every, + max, + } => msg!( + "config.key_limit_every", key = key, per = per, every = every, max = max => + "a limit per {per} of gateway key `{key}` has every: {every}. With per: {per}, \ + every takes 1 to {max}" + ), + KeyLimitEveryCalendar { key, per } => msg!( + "config.key_limit_every_calendar", key = key, per = per => + "a limit per {per} of gateway key `{key}` sets every, which only per: minute and \ + per: hour take. A day, a week and a month start again at local midnight" + ), + UpstreamLimitEmpty { upstream } => msg!( + "config.upstream_limit_empty", upstream = upstream => + "a limit of upstream `{upstream}` names nothing to count. Each entry under limits \ + takes one of requests, tokens or cost" + ), + UpstreamLimitTwoMeasures { upstream, measures } => msg!( + "config.upstream_limit_two_measures", upstream = upstream, measures = measures => + "a limit of upstream `{upstream}` names {measures} together. Each entry takes one \ + of requests, tokens or cost; write one entry for each" + ), + UpstreamLimitNotPositive { + upstream, + per, + measure, + value, + } => msg!( + "config.upstream_limit_not_positive", upstream = upstream, per = per, + measure = measure, value = value => + "the {measure} limit per {per} of upstream `{upstream}` is {value}; it has to be \ + more than 0" + ), + UpstreamLimitDuplicate { + upstream, + per, + measure, + } => msg!( + "config.upstream_limit_duplicate", upstream = upstream, per = per, + measure = measure => + "upstream `{upstream}` has two {measure} limits per {per}. Keep one of them" + ), + UpstreamLimitCacheReads { upstream, measure } => msg!( + "config.upstream_limit_cache_reads", upstream = upstream, measure = measure => + "the {measure} limit of upstream `{upstream}` sets cache_reads, which only a \ + tokens limit takes" + ), + UpstreamLimitCostTooSmall { + upstream, + per, + value, + } => msg!( + "config.upstream_limit_cost_too_small", upstream = upstream, per = per, + value = value => + "the cost limit per {per} of upstream `{upstream}` is {value}; it has to be at \ + least 0.01" + ), + UpstreamLimitRetention { + upstream, + per, + days, + min, + } => msg!( + "config.upstream_limit_retention", upstream = upstream, per = per, days = days, + min = min => + "upstream `{upstream}` has a limit per {per}, and retention.row_days is {days}. \ + After a restart the total for the {per} is added up again from the request \ + records, so row_days has to be at least {min}" + ), + UpstreamLimitEvery { + upstream, + per, + every, + max, + } => msg!( + "config.upstream_limit_every", upstream = upstream, per = per, every = every, + max = max => + "a limit per {per} of upstream `{upstream}` has every: {every}. With per: {per}, \ + every takes 1 to {max}" + ), + UpstreamLimitEveryCalendar { upstream, per } => msg!( + "config.upstream_limit_every_calendar", upstream = upstream, per = per => + "a limit per {per} of upstream `{upstream}` sets every, which only per: minute \ + and per: hour take. A day, a week and a month start again at local midnight" + ), Sync(e) => e.msg(), } } @@ -546,6 +691,7 @@ pub fn validate(cfg: &Config) -> Result<(), ValidationError> { } } check_manual_models(&p.name, &p.models)?; + check_upstream_limits(p, cfg.retention.row_days)?; for (model, spec) in &p.model_specs { check_model_spec(&p.name, model, Some(spec))?; } @@ -736,27 +882,56 @@ fn row_days_needed(per: crate::LimitPer) -> Option { /// 费用上限最小是多少美元:再小的话换成微分是 0(不到一微分),或者小到一个请求就用完 const COST_MIN: f64 = 0.01; -/// 一把密钥的用量上限写得对不对。 +/// 一条用量上限哪儿写得不对。密钥的和上游的写法、规则一样,报的码各是各的(见 +/// [`check_limits`]、[`check_upstream_limits`]) +enum LimitFault { + Empty, + TwoMeasures(String), + NotPositive { + per: &'static str, + measure: &'static str, + value: String, + }, + CostTooSmall { + per: &'static str, + value: String, + }, + CacheReads(&'static str), + Duplicate { + per: &'static str, + measure: &'static str, + }, + Retention { + per: &'static str, + days: u64, + min: u64, + }, + Every { + per: &'static str, + every: u32, + max: u32, + }, + EveryCalendar(&'static str), +} + +/// 一组用量上限写得对不对。 /// /// **每一条恰好数一种量**:一条里写两种,读的人分不清是「都要过」还是「过一个就行」。 /// 同一段时间、同一种量写两条(缓存读算不算进去也一样)的,两条里总有一条不起作用, /// 而用户以为它在起作用。 -fn check_limits(c: &crate::Client, row_days: u64) -> Result<(), ValidationError> { - let key = || c.name.clone(); - for (i, l) in c.limits.iter().enumerate() { +fn check_limit_entries(limits: &[crate::KeyLimit], row_days: u64) -> Result<(), LimitFault> { + for (i, l) in limits.iter().enumerate() { let measures = l.measures(); let measure = match measures.as_slice() { - [] => return Err(ValidationError::KeyLimitEmpty { key: key() }), + [] => return Err(LimitFault::Empty), [one] => *one, many => { - return Err(ValidationError::KeyLimitTwoMeasures { - key: key(), - measures: many - .iter() + return Err(LimitFault::TwoMeasures( + many.iter() .map(|m| m.word()) .collect::>() .join(" and "), - }); + )); } }; // 写了 `.nan`、`.inf` 的费用也在这里挡:前者比什么都不大,后者换不成微分 @@ -771,36 +946,43 @@ fn check_limits(c: &crate::Client, row_days: u64) -> Result<(), ValidationError> crate::LimitMeasure::Tokens => l.tokens.unwrap_or_default().to_string(), crate::LimitMeasure::Cost => l.cost.unwrap_or_default().to_string(), }; - return Err(ValidationError::KeyLimitNotPositive { - key: key(), + return Err(LimitFault::NotPositive { per: l.per.word(), measure: measure.word(), value, }); } if let Some(x) = l.cost.filter(|x| *x < COST_MIN) { - return Err(ValidationError::KeyLimitCostTooSmall { - key: key(), + return Err(LimitFault::CostTooSmall { per: l.per.word(), value: x.to_string(), }); } if l.cache_reads && measure != crate::LimitMeasure::Tokens { - return Err(ValidationError::KeyLimitCacheReads { - key: key(), - measure: measure.word(), - }); + return Err(LimitFault::CacheReads(measure.word())); } - if c.limits[..i].iter().any(|o| o.same_as(l)) { - return Err(ValidationError::KeyLimitDuplicate { - key: key(), + // 「每 0 分钟」没有意思;「每 90 分钟」写成每 2 小时之外,就是写错了单位 + if let Some(every) = l.every { + match l.per.every_max() { + None => return Err(LimitFault::EveryCalendar(l.per.word())), + Some(max) if every == 0 || every > max => { + return Err(LimitFault::Every { + per: l.per.word(), + every, + max, + }); + } + Some(_) => {} + } + } + if limits[..i].iter().any(|o| o.same_as(l)) { + return Err(LimitFault::Duplicate { per: l.per.word(), measure: measure.word(), }); } if let Some(min) = row_days_needed(l.per).filter(|min| row_days < *min) { - return Err(ValidationError::KeyLimitRetention { - key: key(), + return Err(LimitFault::Retention { per: l.per.word(), days: row_days, min, @@ -810,6 +992,95 @@ fn check_limits(c: &crate::Client, row_days: u64) -> Result<(), ValidationError> Ok(()) } +/// 一把密钥的用量上限写得对不对([`check_limit_entries`])。 +fn check_limits(c: &crate::Client, row_days: u64) -> Result<(), ValidationError> { + let key = c.name.clone(); + check_limit_entries(&c.limits, row_days).map_err(|f| match f { + LimitFault::Empty => ValidationError::KeyLimitEmpty { key }, + LimitFault::TwoMeasures(measures) => ValidationError::KeyLimitTwoMeasures { key, measures }, + LimitFault::NotPositive { + per, + measure, + value, + } => ValidationError::KeyLimitNotPositive { + key, + per, + measure, + value, + }, + LimitFault::CostTooSmall { per, value } => { + ValidationError::KeyLimitCostTooSmall { key, per, value } + } + LimitFault::CacheReads(measure) => ValidationError::KeyLimitCacheReads { key, measure }, + LimitFault::Duplicate { per, measure } => { + ValidationError::KeyLimitDuplicate { key, per, measure } + } + LimitFault::Retention { per, days, min } => ValidationError::KeyLimitRetention { + key, + per, + days, + min, + }, + LimitFault::Every { per, every, max } => ValidationError::KeyLimitEvery { + key, + per, + every, + max, + }, + LimitFault::EveryCalendar(per) => ValidationError::KeyLimitEveryCalendar { key, per }, + }) +} + +/// 一家上游的用量上限写得对不对:和密钥的同一套规则([`check_limit_entries`]),天、周、月 +/// 一样要请求记录留够那一期 —— 重启之后上游这一期用了多少也是从记录里加回来的。 +fn check_upstream_limits(p: &crate::Provider, row_days: u64) -> Result<(), ValidationError> { + let upstream = p.name.clone(); + check_limit_entries(&p.limits, row_days).map_err(|f| match f { + LimitFault::Empty => ValidationError::UpstreamLimitEmpty { upstream }, + LimitFault::TwoMeasures(measures) => { + ValidationError::UpstreamLimitTwoMeasures { upstream, measures } + } + LimitFault::NotPositive { + per, + measure, + value, + } => ValidationError::UpstreamLimitNotPositive { + upstream, + per, + measure, + value, + }, + LimitFault::CostTooSmall { per, value } => ValidationError::UpstreamLimitCostTooSmall { + upstream, + per, + value, + }, + LimitFault::CacheReads(measure) => { + ValidationError::UpstreamLimitCacheReads { upstream, measure } + } + LimitFault::Duplicate { per, measure } => ValidationError::UpstreamLimitDuplicate { + upstream, + per, + measure, + }, + LimitFault::Retention { per, days, min } => ValidationError::UpstreamLimitRetention { + upstream, + per, + days, + min, + }, + LimitFault::Every { per, every, max } => ValidationError::UpstreamLimitEvery { + upstream, + per, + every, + max, + }, + LimitFault::EveryCalendar(per) => { + ValidationError::UpstreamLimitEveryCalendar { upstream, per } + } + }) +} + /// 别名表写得对不对,和整份配置的校验是同一套(界面预览一个别名时也用它)。 /// **模型名在不在上游清单里不在这里查**:清单是运行时问来的, /// 上游一时没列出来不该让整份配置不收。 @@ -1585,9 +1856,7 @@ groups: ("idle_timeout_secs", |f| { f.idle_timeout_secs = crate::MAX_IDLE_TIMEOUT_SECS + 1 }), - ("slot_wait_secs", |f| { - f.slot_wait_secs = crate::MAX_SLOT_WAIT_SECS + 1 - }), + ("wait_secs", |f| f.wait_secs = crate::MAX_WAIT_SECS + 1), ]; for (want, bend) in cases { let mut x = base.clone(); @@ -1597,10 +1866,12 @@ groups: other => panic!("{want} 该被拒,实际 {other:?}"), } } - // 等空位写 0 是不等,不是写错 - let mut x = base.clone(); - x.failover.slot_wait_secs = 0; - assert!(validate(&x).is_ok()); + // 等待写 0 是不等,不是写错;上限本身也收 + for secs in [0, crate::MAX_WAIT_SECS] { + let mut x = base.clone(); + x.failover.wait_secs = secs; + assert!(validate(&x).is_ok(), "{secs}"); + } // 无响应超时的两头都收 for secs in [crate::MIN_IDLE_TIMEOUT_SECS, crate::MAX_IDLE_TIMEOUT_SECS] { let mut x = base.clone(); @@ -1793,6 +2064,119 @@ groups: ("k", "day", "0.005") ); assert!(with("[{per: day, cost: 0.01}]", 90).is_ok()); + // every:分钟 1 到 60、小时 1 到 24;天、周、月不能写 + assert!( + with( + "[{per: minute, every: 60, requests: 3}, {per: hour, every: 24, requests: 3}, \ + {per: hour, every: 5, cost: 1}, {per: hour, cost: 1}]", + 90 + ) + .is_ok(), + "每 5 小时和每小时是两条" + ); + for bad in [ + "[{per: minute, every: 0, requests: 3}]", + "[{per: minute, every: 61, requests: 3}]", + "[{per: hour, every: 25, requests: 3}]", + ] { + assert_eq!(code(bad), "config.key_limit_every", "{bad}"); + } + let m = with("[{per: minute, every: 90, requests: 3}]", 90).unwrap_err(); + assert_eq!( + (m.arg("key"), m.arg("per"), m.arg("every"), m.arg("max")), + ("k", "minute", "90", "60") + ); + for per in ["day", "week", "month"] { + let m = with(&format!("[{{per: {per}, every: 1, requests: 3}}]"), 90).unwrap_err(); + assert_eq!(m.code, "config.key_limit_every_calendar", "{per}"); + assert_eq!(m.arg("per"), per); + } + assert_eq!( + code("[{per: hour, every: 5, cost: 3}, {per: hour, every: 5, cost: 5}]"), + "config.key_limit_duplicate" + ); + assert_eq!( + code("[{per: hour, every: 1, cost: 3}, {per: hour, cost: 5}]"), + "config.key_limit_duplicate", + "写明的 1 和不写是同一条" + ); + } + + /// 上游的用量上限:和密钥的同一套规则,报的是上游自己的码 + #[test] + fn upstream_limits_follow_the_key_rules_with_their_own_codes() { + let with = |limits: &str, row_days: u64| { + let mut up = p("u", "https://api.example.com"); + up.limits = serde_yaml_ng::from_str(limits).unwrap(); + let mut cfg = cfg(vec![c("k", "tw-1")], vec![up]); + cfg.retention.row_days = row_days; + validate(&cfg).map_err(|e| e.msg()) + }; + let code = |limits: &str| with(limits, 90).unwrap_err().code; + assert!( + with( + "[{per: hour, every: 5, requests: 300}, {per: minute, tokens: 100000}, \ + {per: day, tokens: 900000, cache_reads: true}, {per: month, cost: 100}]", + 90 + ) + .is_ok() + ); + for (bad, want) in [ + ("[{per: day}]", "config.upstream_limit_empty"), + ( + "[{per: day, requests: 3, cost: 1}]", + "config.upstream_limit_two_measures", + ), + ( + "[{per: day, requests: 0}]", + "config.upstream_limit_not_positive", + ), + ( + "[{per: day, cost: .nan}]", + "config.upstream_limit_not_positive", + ), + ( + "[{per: day, cost: 0.005}]", + "config.upstream_limit_cost_too_small", + ), + ( + "[{per: day, requests: 3, cache_reads: true}]", + "config.upstream_limit_cache_reads", + ), + ( + "[{per: day, cost: 3}, {per: day, cost: 5}]", + "config.upstream_limit_duplicate", + ), + ( + "[{per: minute, every: 61, requests: 3}]", + "config.upstream_limit_every", + ), + ( + "[{per: hour, every: 0, requests: 3}]", + "config.upstream_limit_every", + ), + ( + "[{per: month, every: 2, requests: 3}]", + "config.upstream_limit_every_calendar", + ), + ] { + assert_eq!(code(bad), want, "{bad}"); + } + let m = with("[{per: hour, every: 30, tokens: 5}]", 90).unwrap_err(); + assert_eq!( + ( + m.arg("upstream"), + m.arg("per"), + m.arg("every"), + m.arg("max") + ), + ("u", "hour", "30", "24") + ); + // 天、周、月要请求记录留够那一期:重启之后从记录里加回来 + let m = with("[{per: week, requests: 3}]", 6).unwrap_err(); + assert_eq!(m.code, "config.upstream_limit_retention"); + assert_eq!((m.arg("days"), m.arg("min")), ("6", "7")); + assert!(with("[{per: hour, requests: 3}]", 0).is_ok()); } #[test] @@ -2043,6 +2427,59 @@ mod msg_codes { days: 30, min: 31, }, + KeyLimitEvery { + key: "k".into(), + per: "minute", + every: 90, + max: 60, + }, + KeyLimitEveryCalendar { + key: "k".into(), + per: "day", + }, + UpstreamLimitEmpty { + upstream: "u".into(), + }, + UpstreamLimitTwoMeasures { + upstream: "u".into(), + measures: "requests and cost".into(), + }, + UpstreamLimitNotPositive { + upstream: "u".into(), + per: "day", + measure: "cost", + value: "0".into(), + }, + UpstreamLimitDuplicate { + upstream: "u".into(), + per: "day", + measure: "cost", + }, + UpstreamLimitCacheReads { + upstream: "u".into(), + measure: "cost", + }, + UpstreamLimitCostTooSmall { + upstream: "u".into(), + per: "day", + value: "0.005".into(), + }, + UpstreamLimitRetention { + upstream: "u".into(), + per: "month", + days: 30, + min: 31, + }, + UpstreamLimitEvery { + upstream: "u".into(), + per: "hour", + every: 0, + max: 24, + }, + UpstreamLimitEveryCalendar { + upstream: "u".into(), + per: "week", + }, ]; check( "config.", diff --git a/crates/tw-config/tests/manual/schema.rs b/crates/tw-config/tests/manual/schema.rs index bd25378c..b0ecadfe 100644 --- a/crates/tw-config/tests/manual/schema.rs +++ b/crates/tw-config/tests/manual/schema.rs @@ -526,8 +526,17 @@ pub fn sections() -> Vec
{ Kind::Enum(limit_pers), Def::Required, t( - "The period. `minute` and `hour` are rolling (the last 60 seconds, the last 60 minutes); `day`, `week` and `month` start again at local midnight, on Monday and on the 1st.", - "按多长一段时间算。`minute`、`hour` 是滚动的(最近 60 秒、最近 60 分钟);`day`、`week`、`month` 在本地时间的零点、周一零点、每月一号零点重新算。", + "The period. `minute` and `hour` are rolling (the last 60 seconds, the last 60 minutes, or `every` of them); `day`, `week` and `month` start again at local midnight, on Monday and on the 1st.", + "按多长一段时间算。`minute`、`hour` 是滚动的(最近 60 秒、最近 60 分钟,或者 `every` 个这样的单位);`day`、`week`、`month` 在本地时间的零点、周一零点、每月一号零点重新算。", + ), + ), + row( + "every", + Kind::Int, + Def::Unset, + t( + "How many minutes or hours the rolling window spans: `{per: hour, every: 5}` is the last 5 hours. From 1 to 60 for `minute`, from 1 to 24 for `hour`; `day`, `week` and `month` do not take it. Unset: 1.", + "滚动窗口有几分钟、几小时:`{per: hour, every: 5}` 是最近 5 小时。`minute` 取 1 到 60,`hour` 取 1 到 24;`day`、`week`、`month` 不能写。不写:1。", ), ), row( @@ -713,8 +722,17 @@ pub fn sections() -> Vec
{ Kind::Int, Def::Unset, t( - "Most requests sent to this upstream at the same time, from 1 to 1000. When it is full, a conversation that stays on it waits for a free slot and other requests go to the next upstream; see `failover.slot_wait_secs`. Unset: no limit.", - "同时发给这家的请求最多几个,取值 1 到 1000。满了的时候,留在这家的对话等空位,别的请求换下一家;等多久见 `failover.slot_wait_secs`。不写:不限。", + "Most requests sent to this upstream at the same time, from 1 to 1000. When it is full, a conversation that stays on it waits for a free slot and other requests go to the next upstream; see `failover.wait_secs`. Unset: no limit.", + "同时发给这家的请求最多几个,取值 1 到 1000。满了的时候,留在这家的对话等空位,别的请求换下一家;等多久见 `failover.wait_secs`。不写:不限。", + ), + ), + row( + "limits", + Kind::Objs("providers[].limits[]"), + Def::Unset, + t( + "Usage limits of this upstream, written like its own quota: requests, tokens or cost every few minutes or hours, per day, week or month. Every one has to have room before a request is sent to it. When a rolling window is full, requests go to the next upstream or wait for room; when a day, week or month is used up, requests go to the next upstream until the next one starts. Unset: no limit.", + "这家上游的用量上限,照着它自己的限额写:每几分钟、每几小时、每天、每周或每月的请求数、token 数或费用。每一条都有余量才往这家发。滚动窗口满了,请求换下一家或者等它空出来;天、周、月用满了,到下一期之前请求换下一家。不写:不限。", ), ), row( @@ -789,6 +807,66 @@ pub fn sections() -> Vec
{ ), ], }, + Section { + path: "providers[].limits[]", + ty: checked!(UsageLimit, "{per: day}"), + rows: vec![ + row( + "per", + Kind::Enum(limit_pers), + Def::Required, + t( + "The period. `minute` and `hour` are rolling (the last 60 seconds, the last 60 minutes, or `every` of them); `day`, `week` and `month` start again at local midnight, on Monday and on the 1st.", + "按多长一段时间算。`minute`、`hour` 是滚动的(最近 60 秒、最近 60 分钟,或者 `every` 个这样的单位);`day`、`week`、`month` 在本地时间的零点、周一零点、每月一号零点重新算。", + ), + ), + row( + "every", + Kind::Int, + Def::Unset, + t( + "How many minutes or hours the rolling window spans: `{per: hour, every: 5}` is the last 5 hours. From 1 to 60 for `minute`, from 1 to 24 for `hour`; `day`, `week` and `month` do not take it. Unset: 1.", + "滚动窗口有几分钟、几小时:`{per: hour, every: 5}` 是最近 5 小时。`minute` 取 1 到 60,`hour` 取 1 到 24;`day`、`week`、`month` 不能写。不写:1。", + ), + ), + row( + "requests", + Kind::Int, + Def::Unset, + t( + "At most this many requests sent to this upstream. Each attempt that reaches it counts, unless it answers 429; token counts do not count.", + "最多往这家发这么多个请求。每一次发到了这家的尝试都算,它回 429 的不算;数 token 的请求不算。", + ), + ), + row( + "tokens", + Kind::Int, + Def::Unset, + t( + "At most this many tokens in the answers of this upstream: uncached input, cache writes and output.", + "这家回答的请求最多用这么多 token:未命中缓存的输入、写入缓存的和输出。", + ), + ), + row( + "cost", + Kind::Num, + Def::Unset, + t( + "At most this much, in US dollars, as recorded for the requests this upstream answers; at least 0.01. Models without a price count as 0, and so does everything with `billing: free`.", + "这家回答的请求最多花这么多美元,按每个请求记下的费用算;至少 0.01。没有价格的模型算 0,`billing: free` 时全都算 0。", + ), + ), + row( + "cache_reads", + Kind::Bool, + Def::Is("false"), + t( + "Count cache reads too. Only for a `tokens` limit.", + "把读取缓存的 token 也算进去。只有 `tokens` 上限能写。", + ), + ), + ], + }, Section { path: "providers[].aws", // 字段都可选:写访问密钥还是写 profile,由写法检查管二选一 @@ -1483,8 +1561,8 @@ pub fn sections() -> Vec
{ Kind::Int, Def::Is("3600"), t( - "A rate-limited upstream is set aside for the time its `Retry-After` gives, at most this many seconds. Without `Retry-After` it counts as a failure without a stated reason.", - "被限流的上游按它给的 `Retry-After` 停用,最多这么多秒。没有 `Retry-After` 的按没有说明原因的失败计。", + "Longest pause, in seconds, after an upstream answers 429 (rate limited). The pause lasts as long as the upstream says (`retry-after-ms`, `Retry-After`, or a delay in the error body); when it says nothing, it backs off by the upstream's consecutive 429s: 5, 10, 20 and 40 seconds, then 60 seconds each time. A pause is not a failure: requests wait for the upstream within `wait_secs`, and the circuit breaker does not count it.", + "上游回 429(限流)后暂停这家的最长秒数。上游给出了等待时间的(`retry-after-ms`、`Retry-After` 或错误正文里的等待时间)按它给的暂停;没有给出的,按这家连续被限流的次数退避:5、10、20、40 秒,之后每次 60 秒。暂停不算失败:请求在 `wait_secs` 之内等这家恢复,熔断也不计入。", ), ), row( @@ -1497,12 +1575,12 @@ pub fn sections() -> Vec
{ ), ), row( - "slot_wait_secs", + "wait_secs", Kind::Int, - Def::Is("30"), + Def::Is("120"), t( - "Seconds a request waits in all, counted once the key's own `max_concurrent` lets it in: for a key's `minute` or `hour` limit to free up, and for a free slot on upstreams at their `max_concurrent`. A key limit that does not free up in time refuses the request; without an upstream slot in time it goes to the next upstream, or, when every candidate is full, is answered with 429. `0`: never wait. From 0 to 300.", - "一个请求合计最多等的秒数,从过了密钥自己的 `max_concurrent` 时算起:等密钥的 `minute`、`hour` 上限空出名额,和等并发数满了(`max_concurrent`)的上游空出位置,都算在里面。密钥的上限到时空不出来就拒绝;等不到上游的空位就换下一家,候选全满时回 429。`0`:不等。取值 0 到 300。", + "Seconds a request waits in all, counted once the key's own `max_concurrent` lets it in: for a key's `minute` or `hour` limit to free up, and for an upstream that cannot take it yet — at its `max_concurrent`, paused after answering 429, or at one of its `minute` or `hour` usage limits. A key limit that does not free up in time refuses the request. An upstream that cannot take it is skipped for the next one, and when none is left the request waits for the first to become available; a conversation that stays on an upstream waits for it. When nothing became available in time, the client gets 429 with `Retry-After`. `0`: never wait. From 0 to 600.", + "一个请求合计最多等的秒数,从过了密钥自己的 `max_concurrent` 时算起:等密钥的 `minute`、`hour` 上限空出名额,和等暂时不能发的上游 —— 并发数满了(`max_concurrent`)、回了 429 正在暂停、到了这家 `minute` 或 `hour` 的用量上限 —— 都算在里面。密钥的上限到时空不出来就拒绝。上游暂时不能发的先换下一家,没有下一家了就等最先恢复的那一家;留在某一家的对话等这一家。到时都没有恢复,客户端收到带 `Retry-After` 的 429。`0`:不等。取值 0 到 600。", ), ), ], diff --git a/crates/tw-control/src/key_limits.rs b/crates/tw-control/src/key_limits.rs index 08df8fce..7aaba866 100644 --- a/crates/tw-control/src/key_limits.rs +++ b/crates/tw-control/src/key_limits.rs @@ -1,6 +1,7 @@ -//! 密钥的用量上限和存储层之间的三根线:每记下一行请求就结算,启动时从请求记录里把这一天、 +//! 用量上限和存储层之间的三根线:每记下一行请求就结算,启动时从请求记录里把这一天、 //! 这一周、这个月用了多少加回来,一期的开头变了(到了下一期、机器换了时区)时把新的那一期 -//! 重新加起来(见 `tw_gateway::key_limits`)。 +//! 重新加起来。密钥的(`tw_gateway::key_limits`)和上游的(`tw_gateway::upstream_limits`) +//! 各接一套,接法一样。 //! //! **网关和存储层互不依赖**(两件平级的事),线在这里接:除了 twcore,控制面是唯一同时 //! 看得见两边的地方。twcore 起来时接上,测试也从这里接。 @@ -9,32 +10,31 @@ use std::sync::Arc; use tw_gateway::key_limits::Recorded; -/// 存储层每记下一行,交给网关结算。 +/// 存储层每记下一行,交给网关结算:密钥的和上游的各一份。 pub fn settle_hook(gw: &tw_gateway::AppState) -> tw_store::SettleHook { let limits = gw.key_limits.clone(); + let upstreams = gw.upstream_limits.clone(); Arc::new(move |s: &tw_store::Settled| { - limits.settle( - s.id, - s.at_ms, - &Recorded { - client: s.client.clone(), - path: s.path.clone(), - local: s.local, - reached: s.reached, - requests: 1, - input: s.input, - output: s.output, - cache_read: s.cache_read, - cache_write: s.cache_write, - // 算不出钱的(没有价格、没有用量)算 0:费用上限只管得住有价格的 - cost_micros: s.cost_micros.unwrap_or(0), - }, - ) + let rec = Recorded { + client: s.client.clone(), + path: s.path.clone(), + local: s.local, + reached: s.reached, + requests: 1, + input: s.input, + output: s.output, + cache_read: s.cache_read, + cache_write: s.cache_write, + // 算不出钱的(没有价格、没有用量)算 0:费用上限只管得住有价格的 + cost_micros: s.cost_micros.unwrap_or(0), + }; + limits.settle(s.id, s.at_ms, &rec); + upstreams.settle(s.id, s.at_ms, &s.hops, s.served.as_deref(), &rec); }) } -/// 从请求记录里把每把密钥这一期用了多少加回来。**在第一个请求之前调**。读不了库就 -/// 从 0 起,只记一行 —— 观测挂了,转发照常。 +/// 从请求记录里把每把密钥、每家上游这一期用了多少加回来。**在第一个请求之前调**。读不了 +/// 库就从 0 起,只记一行 —— 观测挂了,转发照常。 pub fn rebuild(gw: &tw_gateway::AppState, db: &tw_store::Db) { gw.key_limits.rebuild(|since| match db.key_usage_since(since) { Ok(rows) => rows.into_iter().map(recorded).collect(), @@ -45,6 +45,16 @@ pub fn rebuild(gw: &tw_gateway::AppState, db: &tw_store::Db) { Vec::new() } }); + gw.upstream_limits + .rebuild(|since| match db.upstream_usage_since(since) { + Ok(rows) => rows.into_iter().map(upstream_recorded).collect(), + Err(e) => { + tracing::warn!( + "the usage of the upstreams could not be read back, so their limits count from zero: {e}" + ); + Vec::new() + } + }); } /// 一期的开头变了:从请求记录里把新的这一期重新加起来(见 @@ -55,6 +65,7 @@ pub fn rebuild(gw: &tw_gateway::AppState, db: &tw_store::Db) { /// 在另起的任务上:要它的那一刻可能正拿着这把锁(结算一行时发现到了下一期)。读不了库就 /// 从 0 起,只记一行 pub fn follow(gw: &tw_gateway::AppState, store: Arc>) { + follow_upstreams(gw, store.clone()); // 弱引用:网关的账拿着这个办法,办法再拿着账就是一个圈,谁都放不掉 let limits = Arc::downgrade(&gw.key_limits); // **接上时的运行时,记在这里**:结算在存储层自己的记录线程上,那条线程不在运行时里,在 @@ -89,6 +100,60 @@ pub fn follow(gw: &tw_gateway::AppState, store: Arc>) { + let limits = Arc::downgrade(&gw.upstream_limits); + let here = tokio::runtime::Handle::try_current().ok(); + gw.upstream_limits + .reread_with(Arc::new(move |m: tw_gateway::key_limits::Moved| { + let (limits, store) = (limits.clone(), store.clone()); + let Some(rt) = here + .clone() + .or_else(|| tokio::runtime::Handle::try_current().ok()) + else { + return; + }; + rt.spawn(async move { + let rec = store.lock().await; + let Some(limits) = limits.upgrade() else { + return; + }; + match rec.db().upstream_usage_since(m.start) { + Ok(rows) => { + let rows: Vec<(String, Recorded)> = + rows.into_iter().map(upstream_recorded).collect(); + limits.reread(&m, &rows); + } + Err(e) => tracing::warn!( + upstream = %m.key, + "the usage of an upstream for its new period could not be read back, \ + so it counts from zero: {e}" + ), + } + }); + })); +} + +/// 一家上游、一个路径的那一组:请求数是算数的那几跳,用量是它接下的那些 +fn upstream_recorded(u: tw_store::UpstreamUsage) -> (String, Recorded) { + let n = |v: i64| v.max(0) as u64; + ( + u.upstream, + Recorded { + client: String::new(), + path: u.path, + local: false, + reached: true, + requests: n(u.requests), + input: n(u.input_tokens), + output: n(u.output_tokens), + cache_read: n(u.cache_read_tokens), + cache_write: n(u.cache_write_tokens), + cost_micros: u.cost_micros, + }, + ) +} + fn recorded(u: tw_store::KeyUsage) -> Recorded { let n = |v: i64| v.max(0) as u64; Recorded { diff --git a/crates/tw-control/src/keys.rs b/crates/tw-control/src/keys.rs index c2819ae1..d2007922 100644 --- a/crates/tw-control/src/keys.rs +++ b/crates/tw-control/src/keys.rs @@ -348,12 +348,14 @@ fn to_client( }) } -/// 界面上的一条上限 → 配置里的写法。费用在界面上是微分,配置里写美元 -fn limit(l: &tw_api::KeyLimitInput) -> tw_config::KeyLimit { +/// 界面上的一条上限 → 配置里的写法。费用在界面上是微分,配置里写美元。`every` 是 1 的不写 +/// (那是默认);别的照写,范围由校验说。密钥的、上游的都用它 +pub(crate) fn limit(l: &tw_api::KeyLimitInput) -> tw_config::KeyLimit { let max = l.max.min(i64::MAX as u64) as i64; let is = |m: tw_api::LimitMeasure| l.measure == m; tw_config::KeyLimit { per: l.per.into(), + every: (l.every != 1).then_some(l.every), requests: is(tw_api::LimitMeasure::Requests).then_some(max), tokens: is(tw_api::LimitMeasure::Tokens).then_some(max), cost: is(tw_api::LimitMeasure::Cost).then(|| max as f64 / 1_000_000.0), diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index 3e2b89a0..c28cb534 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -423,7 +423,7 @@ async fn overview(State(s): State) -> Json { quota_pause_secs: f.quota_pause_secs, rate_limit_max_pause_secs: f.rate_limit_max_pause_secs, idle_timeout_secs: f.idle_timeout_secs, - slot_wait_secs: f.slot_wait_secs, + wait_secs: f.wait_secs, } }, listen: tw_api::ListenView { @@ -513,6 +513,7 @@ fn provider_view( } tw_gateway::health::State::Open => tw_api::Health::Open, }, + throttle: s.gateway.throttles.view(&p.name), auth_rejected: s.gateway.auth_rejected(&p.name), writeback_failed: s.gateway.writeback_failed(&p.name), billing: p.billing.into(), @@ -522,6 +523,7 @@ fn provider_view( .collect(), pricing: p.pricing.clone(), max_concurrent: p.max_concurrent, + limits: s.gateway.upstream_limits.view(&p.name, &p.limits), balance_setting: p.balance.into(), balance: s.gateway.balance_of(p), signed_in: p.signed_in.map(Into::into), diff --git a/crates/tw-control/src/resources.rs b/crates/tw-control/src/resources.rs index d8f27477..7430a5a8 100644 --- a/crates/tw-control/src/resources.rs +++ b/crates/tw-control/src/resources.rs @@ -87,6 +87,13 @@ async fn update_provider( }) .await .map_err(apply_fail)?; + // 改了名:用量上限的账跟着走(今天、这周、这个月用了多少,在路上的预留)。存储层那些行 + // 记的是旧名字,重启之后加回来的只认新名字下的。限流暂停、退避和排着的队也跟着走:改个名 + // 不该让一家正在限流的上游马上又挨发 + if req.provider.name != name { + s.gateway.upstream_limits.rename(&name, &req.provider.name); + s.gateway.throttles.rename(&name, &req.provider.name); + } Ok(Json(tw_api::ConfigWritten { version })) } @@ -554,6 +561,8 @@ fn to_provider( // 沿用原来的,改名时跟着这一项走 model_specs: existing.map(|e| e.model_specs.clone()).unwrap_or_default(), max_concurrent: input.max_concurrent, + // 写得对不对由配置校验说,和手写的同一套 + limits: input.limits.iter().map(crate::keys::limit).collect(), disabled: input.disabled, balance: input.balance.map(Into::into).unwrap_or_default(), signed_in, diff --git a/crates/tw-control/tests/in_flight.rs b/crates/tw-control/tests/in_flight.rs index 336c172c..498d3388 100644 --- a/crates/tw-control/tests/in_flight.rs +++ b/crates/tw-control/tests/in_flight.rs @@ -100,6 +100,7 @@ async fn it_gives_every_running_request_with_what_has_happened_to_it_so_far() { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing: tw_api::Billing::PerToken, diff --git a/crates/tw-control/tests/key_limits.rs b/crates/tw-control/tests/key_limits.rs index 082abe0e..febbdf1b 100644 --- a/crates/tw-control/tests/key_limits.rs +++ b/crates/tw-control/tests/key_limits.rs @@ -143,6 +143,7 @@ fn request(id: u64, key: &str, at: i64, input: u64) -> Vec { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing: tw_api::Billing::PerToken, @@ -298,8 +299,8 @@ async fn saving_a_key_writes_its_limits_and_a_rename_keeps_what_it_used() { "PUT", "/keys/k", serde_json::json!({ "key": { "name": "k2", "limits": [ - { "per": "day", "measure": "cost", "max": 250_000 }, - { "per": "hour", "measure": "tokens", "max": 1000, "cache_reads": true }, + { "per": "day", "every": 1, "measure": "cost", "max": 250_000 }, + { "per": "hour", "every": 5, "measure": "tokens", "max": 1000, "cache_reads": true }, ] } }), ) .await; @@ -309,18 +310,24 @@ async fn saving_a_key_writes_its_limits_and_a_rename_keeps_what_it_used() { assert_eq!(k2.limits[0].cost, Some(0.25)); assert_eq!(k2.limits[1].tokens, Some(1000)); assert!(k2.limits[1].cache_reads); + // every 是 1 的不写进配置,别的照写,读回来也是它 + assert_eq!((k2.limits[0].every, k2.limits[1].every), (None, Some(5))); // 改了名:今天花掉的还在,而且已经超过了新的上限 let list = keys(&b).await; let k2 = key(&list, "k2"); assert_eq!(k2["limits"][0]["used"], 300_000); assert_eq!(k2["limits"][0]["reached"], true); + assert_eq!( + (&k2["limits"][0]["every"], &k2["limits"][1]["every"]), + (&serde_json::json!(1), &serde_json::json!(5)) + ); let (st, v) = call( &b.app, "PUT", "/keys/k2", serde_json::json!({ "key": { "name": "k2", "limits": [ - { "per": "day", "measure": "requests", "max": 0 }, + { "per": "day", "every": 1, "measure": "requests", "max": 0 }, ] } }), ) .await; @@ -334,7 +341,7 @@ async fn saving_a_key_writes_its_limits_and_a_rename_keeps_what_it_used() { "PUT", "/keys/k2", serde_json::json!({ "key": { "name": "k2", "limits": [ - { "per": "day", "measure": "cost", "max": 1_000_000, "cache_reads": true }, + { "per": "day", "every": 1, "measure": "cost", "max": 1_000_000, "cache_reads": true }, ] } }), ) .await; @@ -447,6 +454,7 @@ fn hop(outcome: tw_api::AttemptOutcome, code: &str, busy: bool) -> tw_api::Attem usage: None, queued_ms: None, skipped: busy.then_some(tw_api::ServeSkip::Busy), + pause_ms: None, proxy: None, } } diff --git a/crates/tw-control/tests/live_state.rs b/crates/tw-control/tests/live_state.rs index 1977842d..1096f958 100644 --- a/crates/tw-control/tests/live_state.rs +++ b/crates/tw-control/tests/live_state.rs @@ -157,6 +157,7 @@ async fn a_request_still_running_can_be_opened_and_becomes_whole_when_it_ends() usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing: tw_api::Billing::PerToken, diff --git a/crates/tw-control/tests/replay.rs b/crates/tw-control/tests/replay.rs index 9aa19138..1d55c2f1 100644 --- a/crates/tw-control/tests/replay.rs +++ b/crates/tw-control/tests/replay.rs @@ -390,6 +390,7 @@ fn routed(hops: &[(&str, Option<&str>)]) -> Option { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }) .collect(); diff --git a/crates/tw-control/tests/resources.rs b/crates/tw-control/tests/resources.rs index 481eb3ce..78d624a6 100644 --- a/crates/tw-control/tests/resources.rs +++ b/crates/tw-control/tests/resources.rs @@ -185,7 +185,8 @@ async fn failover_settings_show_their_defaults_and_take_an_edit() { assert_eq!(f["failures_to_pause"], 3, "{body}"); assert_eq!(f["pause_secs"], 60); assert_eq!(f["idle_timeout_secs"], 300); - assert_eq!(f["slot_wait_secs"], 30); + assert_eq!(f["wait_secs"], 120); + assert!(f.get("slot_wait_secs").is_none(), "{body}"); assert!(f.get("stream_start_wait_secs").is_none(), "{body}"); assert!(f.get("next_on_slow_start").is_none(), "{body}"); @@ -224,12 +225,12 @@ async fn failover_settings_show_their_defaults_and_take_an_edit() { "PATCH", "/config", serde_json::json!({ - "ops": [{ "op": "replace", "path": "/failover/slot_wait_secs", "value": 0 }], + "ops": [{ "op": "replace", "path": "/failover/wait_secs", "value": 0 }], }), ) .await; assert_eq!(st, StatusCode::OK, "{body}"); - assert_eq!(b.parsed().failover.slot_wait_secs, 0); + assert_eq!(b.parsed().failover.wait_secs, 0); } /// 无响应超时:概览给出真在用的数,改了就生效;删掉的两项写不进去 diff --git a/crates/tw-control/tests/upstream_limits.rs b/crates/tw-control/tests/upstream_limits.rs new file mode 100644 index 00000000..d4c264be --- /dev/null +++ b/crates/tw-control/tests/upstream_limits.rs @@ -0,0 +1,424 @@ +//! 上游的用量上限在控制面上:新建、保存上游时写进配置,概览里每一条用了多少;改名带着 +//! 用量走;存储层每记下一行就结算(算数的每一跳一个请求,用量记在接下它的那一家);重启 +//! 之后从请求记录里加回来。 + +use std::sync::Arc; + +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use tower::ServiceExt; +use tw_api::Event; +use tw_control::{ConfigManager, ControlState}; + +const CONFIG: &str = "version: 1 +listen: + control: + key: c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00 +clients: + - name: default + key: tw-aaaa +providers: + - name: 中转 + base_url: http://127.0.0.1:9 + key: sk-x + protocol: anthropic + models: [claude-sonnet-4-5] + limits: + - { per: day, requests: 10 } + - { per: hour, every: 5, cost: 1 } + - name: 备用 + base_url: http://127.0.0.1:8 + key: sk-y + protocol: anthropic + models: [claude-sonnet-4-5] +"; + +/// 中午,东八区。天的上限不会碰巧在测试跑到一半时过零点 +const NOON: &str = "2026-10-05T12:00:00+08:00"; + +fn ms(s: &str) -> i64 { + chrono::DateTime::parse_from_rfc3339(s) + .unwrap() + .timestamp_millis() +} + +struct Bed { + _dir: tempfile::TempDir, + gw: tw_gateway::AppState, + app: axum::Router, + rec: Arc>, + path: std::path::PathBuf, +} + +/// `db`:上一次运行留下的请求记录(重启);没有就是一个空库 +fn bed(db: Option) -> Bed { + let d = tempfile::tempdir().unwrap(); + let p = d.path().join("config.yaml"); + std::fs::write(&p, CONFIG).unwrap(); + let cfg = tw_config::try_parse(CONFIG).unwrap(); + let mut gw = tw_gateway::AppState::new(cfg).unwrap(); + gw.set_key_limits_clock(Arc::new(tw_gateway::key_limits::TestClock::new( + ms(NOON), + 8 * 3600, + ))); + let db = db.unwrap_or_else(|| tw_store::Db::in_memory().unwrap()); + // 重启:第一个请求之前把这一期加回来,和 twcore 起来时一样 + tw_control::key_limits::rebuild(&gw, &db); + let rec = tw_store::Recorder::new( + db, + tw_store::Blobs::new(d.path().join("blobs")), + gw.pricing.clone(), + ) + .settling_to(tw_control::key_limits::settle_hook(&gw)); + let rec = Arc::new(tokio::sync::Mutex::new(rec)); + let state = ControlState { + shutdown: Default::default(), + remote: Default::default(), + cfg: Arc::new(ConfigManager::new(p.clone(), gw.clone(), gw.bus.clone())), + gateway: gw.clone(), + store: Some(rec.clone()), + started: std::time::Instant::now(), + price_updater: Default::default(), + chatgpt: Default::default(), + zai: Default::default(), + }; + Bed { + gw, + app: tw_control::router(state), + _dir: d, + rec, + path: p, + } +} + +fn attempt(provider: &str, status: u16) -> tw_api::AttemptView { + tw_api::AttemptView { + provider: provider.into(), + model: None, + outcome: if status == 200 { + tw_api::AttemptOutcome::Served + } else { + tw_api::AttemptOutcome::Status + }, + status: Some(status), + error: None, + ms: 5, + usage: None, + queued_ms: None, + skipped: None, + pause_ms: None, + proxy: None, + } +} + +/// 一个跑完的请求的几条事件:开始于 `at`,尝试链是 `attempts`,按 Sonnet 4.5 的价($3/M +/// 输入) +fn request(id: u64, at: i64, attempts: Vec, input: u64) -> Vec { + vec![ + Event::RequestStarted { + id, + client: "default".into(), + client_hint: None, + session: None, + peer: None, + key_masked: None, + route: "default".into(), + rule: "r".into(), + group: None, + rewritten_by: vec![], + provider: attempts[0].provider.clone(), + billing: tw_api::Billing::PerToken, + model: "claude-sonnet-4-5".into(), + method: "POST".into(), + path: "/v1/messages".into(), + input_estimate: None, + session_log_bytes: None, + at_ms: at as u64, + }, + Event::RequestRouted { + id, + route: "default".into(), + rule: "r".into(), + group: None, + rewritten_by: vec![], + denied_by: None, + affinity: None, + attempts, + billing: tw_api::Billing::PerToken, + egress: None, + }, + Event::RequestFinished { + id, + model: String::new(), + status: 200, + sent_bytes: 0, + received_bytes: 1, + duration_ms: 1, + usage: Some(tw_api::UsageView { + input, + output: 0, + cache_read: 0, + cache_write: 0, + cache_1h: false, + }), + tokens_per_sec: None, + answered_model: None, + }, + ] +} + +async fn call( + app: &axum::Router, + method: &str, + path: &str, + body: serde_json::Value, +) -> (StatusCode, serde_json::Value) { + let r = app + .clone() + .oneshot( + Request::builder() + .method(method) + .uri(path) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let st = r.status(); + let b = axum::body::to_bytes(r.into_body(), 1 << 20).await.unwrap(); + let text = String::from_utf8_lossy(&b).to_string(); + ( + st, + serde_json::from_str(&text).unwrap_or(serde_json::Value::String(text)), + ) +} + +/// 概览里的那一家 +async fn upstream(b: &Bed, name: &str) -> serde_json::Value { + let (st, v) = call(&b.app, "GET", "/overview", serde_json::Value::Null).await; + assert_eq!(st, StatusCode::OK, "{v}"); + v["providers"] + .as_array() + .unwrap() + .iter() + .find(|p| p["name"] == name) + .cloned() + .unwrap_or_else(|| panic!("没有上游「{name}」:{v}")) +} + +fn config(b: &Bed) -> tw_config::Config { + tw_config::try_parse(&std::fs::read_to_string(&b.path).unwrap()).unwrap() +} + +/// 存储层每记下一行就结算:发到「中转」的每一跳算一个请求(回 429 的不算),费用算在接下 +/// 请求的那一家。概览里每一条带着设置和用量。 +#[tokio::test] +async fn the_overview_shows_what_each_upstream_limit_has_used() { + let b = bed(None); + // 「中转」回了 429,换「备用」接下:中转不算 + let events = [ + request( + 1, + ms(NOON), + vec![attempt("中转", 429), attempt("备用", 200)], + 10, + ), + // 「中转」接下:一个请求,十万输入 $0.30 + request(2, ms(NOON), vec![attempt("中转", 200)], 100_000), + // 「中转」回 500,换「备用」接下:中转算一个请求,费用不是它的 + request( + 3, + ms(NOON), + vec![attempt("中转", 500), attempt("备用", 200)], + 100_000, + ), + ]; + for e in events.iter().flatten() { + b.rec.lock().await.on_event(e); + } + let u = upstream(&b, "中转").await; + let limits = &u["limits"]; + assert_eq!(limits[0]["per"], "day"); + assert_eq!(limits[0]["every"], 1); + assert_eq!(limits[0]["measure"], "requests"); + assert_eq!(limits[0]["max"], 10); + assert_eq!(limits[0]["used"], 2); + assert_eq!(limits[0]["resets_at_ms"], ms("2026-10-06T00:00:00+08:00")); + assert_eq!(limits[1]["per"], "hour"); + assert_eq!(limits[1]["every"], 5); + assert_eq!(limits[1]["measure"], "cost"); + assert_eq!(limits[1]["used"], 300_000, "十万输入 × $3/M"); + assert!( + limits[1].get("resets_at_ms").is_none(), + "滚动的没有重置时刻" + ); + assert_eq!(limits[1]["reached"], false); + // 没设上限的,列表是空的 + assert_eq!(upstream(&b, "备用").await["limits"], serde_json::json!([])); +} + +/// 新建、保存上游时写进配置:every 是 1 的不写,费用在界面上是微分、配置里写美元;写错了 +/// 是配置校验的那句话。改名之后用过的跟着走。 +#[tokio::test] +async fn saving_an_upstream_writes_its_limits_and_a_rename_keeps_what_it_used() { + let b = bed(None); + let (st, v) = call( + &b.app, + "POST", + "/providers", + serde_json::json!({ "provider": { + "name": "新的", + "base_url": "https://relay.example", + "key": "sk-new", + "protocol": "anthropic", + "limits": [ + { "per": "minute", "every": 30, "measure": "requests", "max": 60 }, + { "per": "week", "every": 1, "measure": "cost", "max": 2_500_000 }, + ], + } }), + ) + .await; + assert_eq!(st, StatusCode::OK, "{v}"); + let cfg = config(&b); + let p = cfg.providers.iter().find(|p| p.name == "新的").unwrap(); + assert_eq!( + (p.limits[0].per, p.limits[0].every, p.limits[0].requests), + (tw_config::LimitPer::Minute, Some(30), Some(60)) + ); + assert_eq!((p.limits[1].every, p.limits[1].cost), (None, Some(2.5))); + let text = std::fs::read_to_string(&b.path).unwrap(); + assert!(text.contains("every: 30"), "{text}"); + let u = upstream(&b, "新的").await; + assert_eq!( + (&u["limits"][0]["every"], &u["limits"][1]["max"]), + (&serde_json::json!(30), &serde_json::json!(2_500_000)) + ); + + // 「中转」用掉一个请求、刚回了 429 暂停着,然后改名、换上限:用量和暂停都跟着走 + for e in request(1, ms(NOON), vec![attempt("中转", 200)], 10) { + b.rec.lock().await.on_event(&e); + } + let paused = b.gw.throttles.limited( + "中转", + std::time::Instant::now(), + Some(std::time::Duration::from_secs(60)), + ); + assert_eq!(paused.as_secs(), 60); + let (st, v) = call( + &b.app, + "PUT", + &tw_api::ep::UpdateProvider::path("中转"), + serde_json::json!({ "provider": { + "name": "中转二", + "base_url": "http://127.0.0.1:9", + "key": "sk-x", + "protocol": "anthropic", + "models": ["claude-sonnet-4-5"], + "limits": [ { "per": "day", "every": 1, "measure": "requests", "max": 1 } ], + } }), + ) + .await; + assert_eq!(st, StatusCode::OK, "{v}"); + let u = upstream(&b, "中转二").await; + assert_eq!(u["limits"][0]["used"], 1); + assert_eq!(u["limits"][0]["reached"], true); + assert!(matches!( + b.gw.throttles.peek("中转二", &[]), + tw_gateway::throttle::Peek::Wait { + kind: tw_api::ThrottleKind::RateLimited, + .. + } + )); + assert_eq!( + b.gw.throttles.peek("中转", &[]), + tw_gateway::throttle::Peek::Open + ); + assert_eq!(config(&b).providers[0].limits.len(), 1); + + // 写错了:校验的那句话,码是上游自己的 + for (limit, code) in [ + ( + serde_json::json!({ "per": "day", "every": 1, "measure": "requests", "max": 0 }), + "config.upstream_limit_not_positive", + ), + ( + serde_json::json!({ "per": "day", "every": 3, "measure": "requests", "max": 5 }), + "config.upstream_limit_every_calendar", + ), + ( + serde_json::json!({ "per": "hour", "every": 25, "measure": "requests", "max": 5 }), + "config.upstream_limit_every", + ), + ( + serde_json::json!({ "per": "day", "every": 1, "measure": "cost", "max": 1_000_000, + "cache_reads": true }), + "config.upstream_limit_cache_reads", + ), + ] { + let (st, v) = call( + &b.app, + "PUT", + &tw_api::ep::UpdateProvider::path("中转二"), + serde_json::json!({ "provider": { + "name": "中转二", + "base_url": "http://127.0.0.1:9", + "key": "sk-x", + "protocol": "anthropic", + "limits": [limit], + } }), + ) + .await; + assert_eq!(st, StatusCode::BAD_REQUEST, "{v}"); + assert!(v.to_string().contains(code), "{code}: {v}"); + } +} + +/// 重启:这一天、这一周、这个月的数从请求记录里加回来,按尝试链数请求、按接下它的那一家记 +/// 用量。前一天的不算进今天。 +#[tokio::test] +async fn a_restart_adds_this_periods_upstream_usage_back_from_the_records() { + let d = tempfile::tempdir().unwrap(); + let file = d.path().join("data.db"); + let mut earlier = tw_store::Recorder::new( + tw_store::Db::open(&file).unwrap(), + tw_store::Blobs::new(d.path().join("blobs")), + tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), + ); + let events = [ + request( + 1, + ms("2026-10-05T09:00:00+08:00"), + vec![attempt("中转", 200)], + 100_000, + ), + request( + 2, + ms("2026-10-05T10:00:00+08:00"), + vec![ + attempt("中转", 500), + attempt("中转", 429), + attempt("备用", 200), + ], + 100_000, + ), + // 前一天的:今天不算 + request( + 3, + ms("2026-10-04T23:00:00+08:00"), + vec![attempt("中转", 200)], + 100_000, + ), + ]; + for e in events.iter().flatten() { + earlier.on_event(e); + } + drop(earlier); + let b = bed(Some(tw_store::Db::open(&file).unwrap())); + let u = upstream(&b, "中转").await; + assert_eq!( + u["limits"][0]["used"], 2, + "今天两跳算数:接下的一跳、回 500 的一跳" + ); + assert_eq!(u["limits"][1]["used"], 0, "滚动窗口从空的开始"); +} diff --git a/crates/tw-gateway/src/affinity.rs b/crates/tw-gateway/src/affinity.rs index b1698d0d..57de1abd 100644 --- a/crates/tw-gateway/src/affinity.rs +++ b/crates/tw-gateway/src/affinity.rs @@ -14,7 +14,8 @@ //! - **留在「实际回答的那一家」**,不是路由排的头一个:故障转移之后接下回答的备选, //! 就是之后留下的那一家 —— 缓存在它那儿。 //! - **那一家不可用(熔断着、冷却中)就放开**,照常排序:留在一家发不出去的上游 -//! 上,省下的缓存一分钱都兑现不了。 +//! 上,省下的缓存一分钱都兑现不了。用量上限这一期用完了的、限流暂停到这个请求等不到的 +//! 一样放开(见 [`crate::throttle`]);满着的、暂停一会儿就好的照旧留,请求等它。 //! //! 对话按「路由 + 对话标识」记(见 [`Conversation`]),回答的那一家还要对得上策略组: //! 同一段对话这一轮被规则送去另一个组,上次那一家的缓存和这一组无关。 @@ -199,7 +200,8 @@ impl Affinity { /// 交回留下的理由。 /// /// `group` 是这次规则交给的策略组;`available` 说一家现在能不能进候选链(熔断、 - /// 冷却)。只剩一个候选时没有「留不留」可言。 + /// 冷却),和这段对话还该不该留在它那儿(用量上限用完了、限流暂停得太久)。只剩一个候选时 + /// 没有「留不留」可言。 pub fn stay( &self, c: &Conversation, diff --git a/crates/tw-gateway/src/ending.rs b/crates/tw-gateway/src/ending.rs index c6c4d725..6ebb1840 100644 --- a/crates/tw-gateway/src/ending.rs +++ b/crates/tw-gateway/src/ending.rs @@ -76,6 +76,8 @@ pub struct Ending { /// 上游在流里报的错:哪一家、原话。**有它就不是成功** —— 响应头是 200,回答却断在了 /// 半路 upstream_error: Option<(String, String)>, + /// 上游报错的那一帧的 `data`:回程照它认是不是限流(见 [`Ending::upstream_fault`]) + upstream_fault: Option, /// 上游回的不是 2xx、原样交给了客户端(见 [`Ending::refused`])。**有它就不是成功**: /// 客户端拿到的是上游的错误,不是回答 refusal: Option, @@ -167,12 +169,15 @@ struct Watch { } impl Watch { - /// 这些帧里上游报的第一个错 - fn error_in(&self, frames: Vec) -> Option<(String, String)> { - frames - .iter() - .find_map(|f| tw_dialect::convert::stream_error(self.dialect, f)) - .map(|message| (self.provider.clone(), message)) + /// 这些帧里上游报的第一个错:哪一家、原话,和那一帧的 `data` + fn error_in( + &self, + frames: Vec, + ) -> Option<((String, String), String)> { + frames.iter().find_map(|f| { + tw_dialect::convert::stream_error(self.dialect, f) + .map(|message| ((self.provider.clone(), message), f.data.clone())) + }) } } @@ -315,6 +320,7 @@ impl Ending { lap: None, watch: None, upstream_error: None, + upstream_fault: None, refusal: None, answer: None, abort: None, @@ -379,6 +385,13 @@ impl Ending { self.opened.is_some() } + /// 上游在流里报了错:那一帧的 `data`(上游的格式,[`Ending::streaming`] 时给的)。回程照它 + /// 认是不是限流 —— 答到一半才报的限流,回答照样以错误收尾,这一家照样暂停(见 + /// [`crate::throttle`]) + pub fn upstream_fault(&self) -> Option<&str> { + self.upstream_fault.as_deref() + } + /// 上游 `provider` 回的不是 2xx,原样交给了客户端(4xx 是请求本身的问题,或者没有 /// 下一家可换了;3xx 交还客户端,由它决定跟不跟)。 /// @@ -485,8 +498,9 @@ impl Ending { return; }; let frames = w.frames.feed(chunk); - if let Some(found) = w.error_in(frames) { + if let Some((found, data)) = w.error_in(frames) { self.upstream_error = Some(found); + self.upstream_fault = Some(data); self.watch = None; } } @@ -598,7 +612,12 @@ impl Ending { // 最后一帧后面不带空行的上游:收尾时再看一眼 if let Some(mut w) = self.watch.take() { let frames = w.frames.flush(); - self.upstream_error = self.upstream_error.take().or(w.error_in(frames)); + if self.upstream_error.is_none() + && let Some((found, data)) = w.error_in(frames) + { + self.upstream_error = Some(found); + self.upstream_fault = Some(data); + } } if let Some((upstream, message)) = self.upstream_error.take() { let message: String = message.chars().take(500).collect(); diff --git a/crates/tw-gateway/src/error.rs b/crates/tw-gateway/src/error.rs index c23e3a40..997b13cb 100644 --- a/crates/tw-gateway/src/error.rs +++ b/crates/tw-gateway/src/error.rs @@ -59,7 +59,7 @@ pub enum Source { /// 上游都满着时告诉客户端过几秒再来(`Retry-After`)。 /// -/// **不按 `slot_wait_secs` 算**:那么久已经在网关里等过了,再让客户端干等同样久没有意义。 +/// **不按 `wait_secs` 算**:那么久已经在网关里等过了,再让客户端干等同样久没有意义。 /// 重试进来照样排队等空位,所以这个数只管客户端别立刻打回来 pub const BUSY_RETRY_AFTER_SECS: u64 = 5; diff --git a/crates/tw-gateway/src/failure.rs b/crates/tw-gateway/src/failure.rs index 1006769d..6d33d0d1 100644 --- a/crates/tw-gateway/src/failure.rs +++ b/crates/tw-gateway/src/failure.rs @@ -141,31 +141,39 @@ fn used_up_reset(headers: &HeaderMap, now_ms: u64) -> Option> { Some(full.iter().filter_map(|w| w.resets_at_ms).max()) } -/// `Retry-After`:秒数,或者一个 HTTP 日期。 +/// 按秒说的「多久之后再来」短于这个的当没说(`Retry-After: 0`、`"retryDelay": "0s"`):照它 +/// 暂停等于不暂停,被限流的请求当场又发、又被限流,一直空转到等待期限(见 [`crate::throttle`]) +pub const MIN_RETRY_AFTER: Duration = Duration::from_secs(1); + +/// `Retry-After`:秒数,或者一个 HTTP 日期。不到一秒([`MIN_RETRY_AFTER`])的当没说 pub fn retry_after(headers: &HeaderMap, now_ms: u64) -> Option { let v = headers .get(http::header::RETRY_AFTER)? .to_str() .ok()? .trim(); - if let Ok(secs) = v.parse::() { - return Some(Duration::from_secs(secs)); - } - let at = chrono::DateTime::parse_from_rfc2822(v) - .ok()? - .timestamp_millis(); - let left = u64::try_from(at).ok()?.checked_sub(now_ms)?; - Some(Duration::from_millis(left)) + let left = match v.parse::() { + Ok(secs) => Duration::from_secs(secs), + Err(_) => { + let at = chrono::DateTime::parse_from_rfc2822(v) + .ok()? + .timestamp_millis(); + Duration::from_millis(u64::try_from(at).ok()?.checked_sub(now_ms)?) + } + }; + (left >= MIN_RETRY_AFTER).then_some(left) } -/// Gemini 把「多久之后再来」写在正文里:`"retryDelay": "33s"`。 -fn retry_delay_in_body(body: &[u8]) -> Option { +/// Gemini 把「多久之后再来」写在正文里:`"retryDelay": "33s"`。不到一秒的当没说;大得离谱的 +/// 照最长的算(交给调用方封顶),**不能让换算 panic** +pub fn retry_delay_in_body(body: &[u8]) -> Option { let v: serde_json::Value = serde_json::from_slice(body).ok()?; let details = v.pointer("/error/details")?.as_array()?; details.iter().find_map(|d| { let s = d.get("retryDelay")?.as_str()?; let secs: f64 = s.strip_suffix('s')?.parse().ok()?; - (secs.is_finite() && secs >= 0.0).then(|| Duration::from_secs_f64(secs)) + (secs.is_finite() && secs >= MIN_RETRY_AFTER.as_secs_f64()) + .then(|| Duration::try_from_secs_f64(secs).unwrap_or(Duration::MAX)) }) } @@ -348,6 +356,29 @@ mod tests { ); } + #[test] + fn a_retry_after_under_a_second_counts_as_unsaid_and_an_absurd_one_does_not_panic() { + // 照 0 秒暂停等于不暂停:当没说,按退避算 + let h = headers(&[("retry-after", "0")]); + assert_eq!(retry_after(&h, NOW), None); + let past = chrono::DateTime::from_timestamp_millis((NOW - 5_000) as i64) + .unwrap() + .to_rfc2822(); + assert_eq!(retry_after(&headers(&[("retry-after", &past)]), NOW), None); + let body = |d: &str| format!(r#"{{"error":{{"details":[{{"retryDelay":"{d}"}}]}}}}"#); + assert_eq!(retry_delay_in_body(body("0s").as_bytes()), None); + assert_eq!(retry_delay_in_body(body("0.5s").as_bytes()), None); + assert_eq!( + retry_delay_in_body(body("1.5s").as_bytes()), + Some(Duration::from_millis(1500)) + ); + // 大得离谱的照最长的算,由调用方封顶 + assert_eq!( + retry_delay_in_body(body("1e300s").as_bytes()), + Some(Duration::MAX) + ); + } + #[test] fn retry_after_reads_an_http_date() { let at = chrono::DateTime::from_timestamp_millis((NOW + 90_000) as i64) diff --git a/crates/tw-gateway/src/health.rs b/crates/tw-gateway/src/health.rs index 098e839a..f3560513 100644 --- a/crates/tw-gateway/src/health.rs +++ b/crates/tw-gateway/src/health.rs @@ -4,9 +4,13 @@ //! 这里是桌面版的**规矩**: //! //! **停用多久看失败的原因**(见 [`crate::failure`])。上游说了原因的 —— 余额 -//! 不足、额度用完、限流并给了 `Retry-After` —— 一次就停用,停到那个原因该过去 -//! 的时候;说不出原因的(5xx、连不上)连续几次才停,停的时长每次翻倍,成功一次 -//! 就回到起点。几个数都在配置的 `failover` 里。 +//! 不足、额度用完 —— 一次就停用,停到那个原因该过去的时候;说不出原因的(5xx、 +//! 连不上)连续几次才停,停的时长每次翻倍,成功一次就回到起点。几个数都在配置的 +//! `failover` 里。 +//! +//! **限流(429)不在这里**:它不是失败,是让等一会儿 —— 那一家暂停,请求排队等它(见 +//! [`crate::throttle`])。记进熔断的话,一家只是被限流的上游连着几次就被当成坏了,唯一的 +//! 候选还会被自己停用。 //! //! 两条硬边界,它们比参数重要得多: //! @@ -215,22 +219,22 @@ impl Health { } /// 记一次失败,按原因停用。返回状态变化。 + /// + /// **限流不记**([`Cause::RateLimited`]):不进熔断、不算进成功率,那一家由 + /// [`crate::throttle`] 暂停。调用方该把它交给那边([`crate::throttle::Throttles::limited`]) pub fn record_cause(&self, name: &str, cause: Cause) -> Option { let s = self.settings(); let now = (self.clock)(); let secs = |n: u64| now.saturating_add(i64::try_from(n.saturating_mul(1000)).unwrap_or(i64::MAX)); let until = match cause { - Cause::ModelUnavailable => return None, + Cause::ModelUnavailable | Cause::RateLimited { .. } => return None, Cause::NoBalance => secs(s.no_balance_pause_secs), Cause::QuotaUsedUp { resets_at_ms } => resets_at_ms .and_then(|t| i64::try_from(t).ok()) .filter(|t| *t > now) .unwrap_or_else(|| secs(s.quota_pause_secs)), - Cause::RateLimited { - retry_after: Some(d), - } if !d.is_zero() => secs(d.as_secs().max(1).min(s.rate_limit_max_pause_secs)), - Cause::RateLimited { .. } | Cause::AuthRejected | Cause::Unexplained => { + Cause::AuthRejected | Cause::Unexplained => { return self.record_failure(name); } }; @@ -253,9 +257,10 @@ impl Health { /// (`tw_engine::balance_factors`)。 /// /// 成败的判据就是熔断的判据,不另起一套:上面几个 `record_*` 记一次,这里就记一次。 - /// 于是 5xx、连不上、超时、限流、额度用完、没钱了、凭据被拒或取不到、流在第一段内容 - /// 之前断了都算失败;请求本身的问题(别的 4xx)算这家答上了;这家没有这个模型不记 —— - /// 它对别的模型照样好好的。客户端中途走了的不经过这里,也不记。 + /// 于是 5xx、连不上、超时、额度用完、没钱了、凭据被拒或取不到、流在第一段内容 + /// 之前断了都算失败;请求本身的问题(别的 4xx)算这家答上了;这家没有这个模型、限流 + /// 不记 —— 前者它对别的模型照样好好的,后者是让等一会儿(见 [`crate::throttle`])。 + /// 客户端中途走了的不经过这里,也不记。 /// /// **换下一家、却不是这家的错的**(它只是慢,或者正忙)不该调 `record_*`:调了,它在 /// 熔断和这里都会被算成一次失败。 @@ -542,15 +547,6 @@ mod tests { h.record_cause("c", Cause::QuotaUsedUp { resets_at_ms: None }); assert_eq!(h.cooldown_left("c"), Some(Duration::from_secs(3600))); - // Retry-After 太长的按上限 - h.record_cause( - "d", - Cause::RateLimited { - retry_after: Some(Duration::from_secs(99_999)), - }, - ); - assert_eq!(h.cooldown_left("d"), Some(Duration::from_secs(3600))); - // 到点了:能用,定时器看到 0;成功之后什么都不记得 advance(&now, 1_800_000); assert!(h.is_available("a")); @@ -560,17 +556,24 @@ mod tests { } #[test] - fn a_missing_model_or_an_unstated_rate_limit_does_not_set_it_aside_at_once() { + fn a_missing_model_or_a_rate_limit_never_sets_it_aside() { let (h, _) = clocked(); assert_eq!(h.record_cause("a", Cause::ModelUnavailable), None); assert!(h.is_available("a")); - // 没说多久的限流按说不出原因的失败算:连续三次才停 - for _ in 0..2 { - h.record_cause("b", Cause::RateLimited { retry_after: None }); + // 限流不是失败,说没说多久都一样:多少次都不停用,也不算进连续失败的次数 + for retry_after in [None, Some(Duration::from_secs(30))] { + for _ in 0..10 { + assert_eq!( + h.record_cause("b", Cause::RateLimited { retry_after }), + None + ); + } } assert!(h.is_available("b")); + assert_eq!(h.cooldown_left("b"), None); + h.record_cause("b", Cause::Unexplained); h.record_cause("b", Cause::Unexplained); - assert!(!h.is_available("b")); + assert!(h.is_available("b"), "两次说不出原因的失败还不够"); } #[test] @@ -639,26 +642,27 @@ mod tests { for cause in [ Cause::NoBalance, Cause::QuotaUsedUp { resets_at_ms: None }, - Cause::RateLimited { - retry_after: Some(Duration::from_secs(30)), - }, - Cause::RateLimited { retry_after: None }, Cause::AuthRejected, Cause::Unexplained, ] { h.record_cause("a", cause); } h.record_failure("a"); - assert_eq!(rate(&h, "a"), Some(0.0), "七次失败"); + assert_eq!(rate(&h, "a"), Some(0.0), "五次失败"); + // 限流不是失败,也不算答上了:不记 + for retry_after in [None, Some(Duration::from_secs(30))] { + h.record_cause("a", Cause::RateLimited { retry_after }); + } + assert_eq!(rate(&h, "a"), Some(0.0), "还是五次"); // 这家没有这个模型:不算它坏了,也不算它答上了 for _ in 0..10 { h.record_cause("b", Cause::ModelUnavailable); } assert_eq!(rate(&h, "b"), None); - for _ in 0..3 { + for _ in 0..5 { h.record_success("a"); } - assert_eq!(rate(&h, "a"), Some(0.3)); + assert_eq!(rate(&h, "a"), Some(0.5)); } #[test] diff --git a/crates/tw-gateway/src/key_limits/books.rs b/crates/tw-gateway/src/key_limits/books.rs new file mode 100644 index 00000000..9d71324e --- /dev/null +++ b/crates/tw-gateway/src/key_limits/books.rs @@ -0,0 +1,349 @@ +//! 用量的账:密钥的用量上限([`super::KeyLimits`])和上游的([`crate::upstream_limits`])共用的 +//! 那一半 —— 一期一期的合计、滚动窗口里的每一笔、在跑的预留,和从它们算出来的「此刻用了 +//! 多少」「什么时候空出来」。 +//! +//! **两边数法一样,记在谁的名下不一样**:密钥的记在密钥名下,上游的记在上游名下。什么时候 +//! 记上、记多少、用满了怎么办,各自决定。 + +use std::collections::{HashMap, HashSet, VecDeque}; + +use tw_config::{KeyLimit, LimitMeasure, LimitPer}; + +use super::{Clock, Moved}; + +/// 自然周期。下标就是 [`Books::periods`] 的下标 +pub(crate) const CALENDAR: [LimitPer; 3] = [LimitPer::Day, LimitPer::Week, LimitPer::Month]; + +/// 请求已经结束、存储层却还没交来它那一行时,它的预留再留多久。 +/// +/// 存储层落后、丢了事件时那一行永远不会来:不放掉的话,那份预留一直占着额度。存储层 +/// 正常时几毫秒就到 +pub(crate) const GRACE_MS: i64 = 60_000; + +pub(crate) fn calendar_index(per: LimitPer) -> usize { + CALENDAR.iter().position(|p| *p == per).unwrap_or(0) +} + +/// 一份用量。四样分开记,各条上限各取各的。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(crate) struct Amount { + pub requests: i64, + /// 没走缓存的输入 + 写进缓存的 + 输出 + pub tokens: i64, + pub cache_read: i64, + /// 微分 + pub cost: i64, +} + +impl Amount { + pub fn of(&self, m: LimitMeasure, cache_reads: bool) -> i64 { + match m { + LimitMeasure::Requests => self.requests, + LimitMeasure::Tokens if cache_reads => self.tokens.saturating_add(self.cache_read), + LimitMeasure::Tokens => self.tokens, + LimitMeasure::Cost => self.cost, + } + } + pub fn add(&mut self, o: &Amount) { + self.requests = self.requests.saturating_add(o.requests); + self.tokens = self.tokens.saturating_add(o.tokens); + self.cache_read = self.cache_read.saturating_add(o.cache_read); + self.cost = self.cost.saturating_add(o.cost); + } +} + +/// 一期(天、周、月)的合计。 +#[derive(Debug, Clone, Copy)] +pub(crate) struct Period { + pub start: i64, + pub end: i64, + pub sum: Amount, +} + +/// 一个名字(一把密钥、一家上游)的账。**跨重载存活**,按名字记 —— 存储层那一行记的也是 +/// 名字。 +#[derive(Debug, Default)] +pub(crate) struct Books { + /// 天、周、月,和 [`CALENDAR`] 同一个顺序。到了下一期就换一本新的 + pub periods: [Option; 3], + /// 滚动窗口用:最近的每一笔(请求数在记上它的那一刻,token 和费用在结算时),按记下的 + /// 先后。**只有设了分钟、小时上限的才记**,留到最长的那个窗口为止 + pub recent: VecDeque<(i64, Amount)>, + /// 开头变了、换了一本新的那几期(哪一种、新的开头),等着从请求记录里重新加(见 + /// [`super::KeyLimits::reread_with`]) + pub moved: Vec<(LimitPer, i64)>, +} + +impl Books { + /// `now` 所在的那一期。开头和记着的不一样(到了下一期、换了时区)就换一本新的,记下它 + /// 要从请求记录里重新加;加回来之前先从 0 起。从没记过这一种的,从 0 起就是对的 + /// —— 启动时从库里加回来过(`rebuild`),之后的每一行都结算过 + pub fn period(&mut self, clock: &dyn Clock, per: LimitPer, now: i64) -> &mut Period { + let (start, end) = clock.period(per, now); + let slot = &mut self.periods[calendar_index(per)]; + match slot { + Some(p) if p.start == start => {} + _ => { + if slot.is_some() { + self.moved.push((per, start)); + } + *slot = Some(Period { + start, + end, + sum: Amount::default(), + }) + } + } + slot.as_mut().expect("set just above") + } +} + +/// 一个在跑的请求(密钥的),或者一跳发出去的请求(上游的)占着的那一份。 +#[derive(Debug)] +pub(crate) struct Reservation { + /// 记在谁的名下:密钥名,或者上游名 + pub key: String, + /// 记上它的那一刻。算在哪一期、在不在滚动窗口里看它 + pub at_ms: i64, + /// 输入 token 的估算 + pub tokens: i64, + /// 输入费用的估算,微分 + pub cost: i64, + /// 它是哪个请求。还没交给请求号之前没有 + pub request: Option, + /// 发现这个请求已经结束的那一刻(见 [`GRACE_MS`]) + pub closed_since: Option, +} + +impl Reservation { + pub fn new(key: &str, at_ms: i64, tokens: u64, cost_micros: i64) -> Self { + Self { + key: key.to_string(), + at_ms, + tokens: tokens.min(i64::MAX as u64) as i64, + cost: cost_micros.max(0), + request: None, + closed_since: None, + } + } +} + +/// 一组上限里最长的那个滚动窗口。没有就是 None +pub(crate) fn longest_window(limits: &[KeyLimit]) -> Option { + limits.iter().filter_map(KeyLimit::window_ms).max() +} + +/// 一条上限此刻用了多少:记下的,加上在跑的请求的预留 +pub(crate) fn used( + books: &mut HashMap, + held: &HashMap, + clock: &dyn Clock, + key: &str, + l: &KeyLimit, + now: i64, +) -> i64 { + let (m, cr) = (l.measure(), l.cache_reads); + let reserved = |r: &Reservation| match m { + LimitMeasure::Requests => 1, + LimitMeasure::Tokens => r.tokens, + LimitMeasure::Cost => r.cost, + }; + match l.window_ms() { + None => { + let p = *books + .entry(key.to_string()) + .or_default() + .period(clock, l.per, now); + let held: i64 = held + .values() + .filter(|r| r.key == key && r.at_ms >= p.start && r.at_ms < p.end) + .map(reserved) + .sum(); + p.sum.of(m, cr).saturating_add(held) + } + Some(w) => { + let cutoff = now - w; + let done: i64 = books + .get(key) + .map(|b| { + b.recent + .iter() + .filter(|(t, _)| *t > cutoff) + .map(|(_, a)| a.of(m, cr)) + .sum() + }) + .unwrap_or(0); + // 请求数在记上的那一刻已经进了窗口,预留里只算 token 和费用 + let held: i64 = held + .values() + .filter(|r| r.key == key && r.at_ms > cutoff) + .map(|r| match m { + LimitMeasure::Requests => 0, + _ => reserved(r), + }) + .sum(); + done.saturating_add(held) + } + } +} + +/// 一条滚动窗口上限用了 `used`(不小于上限)时,什么时候空出一个位置:窗口里的每一笔按 +/// 时间先后滑出去,滑到用量低于上限的那一刻。Unix 毫秒 +pub(crate) fn free_at( + books: &HashMap, + held: &HashMap, + key: &str, + l: &KeyLimit, + used: i64, + now: i64, +) -> i64 { + let Some(w) = l.window_ms() else { + return now; + }; + let (m, cr, max) = (l.measure(), l.cache_reads, l.max()); + let cutoff = now - w; + let mut items: Vec<(i64, i64)> = books + .get(key) + .map(|b| { + b.recent + .iter() + .filter(|(t, _)| *t > cutoff) + .map(|(t, a)| (*t, a.of(m, cr))) + .collect() + }) + .unwrap_or_default(); + if m != LimitMeasure::Requests { + items.extend( + held.values() + .filter(|r| r.key == key && r.at_ms > cutoff) + .map(|r| { + let v = if m == LimitMeasure::Tokens { + r.tokens + } else { + r.cost + }; + (r.at_ms, v) + }), + ); + } + items.retain(|(_, v)| *v > 0); + items.sort_by_key(|(t, _)| *t); + let mut left = used; + let mut free_at = now + w; + for (t, v) in items { + left -= v; + if left < max { + free_at = t + w; + break; + } + } + free_at +} + +/// 在滚动窗口里记上一个请求(这组上限有分钟、小时的才记),顺手把滑出去的扔掉 +pub(crate) fn record_request( + books: &mut HashMap, + key: &str, + limits: &[KeyLimit], + now: i64, +) { + if let Some(w) = longest_window(limits) { + let books = books.entry(key.to_string()).or_default(); + books.recent.push_back(( + now, + Amount { + requests: 1, + ..Default::default() + }, + )); + while books.recent.front().is_some_and(|(t, _)| *t <= now - w) { + books.recent.pop_front(); + } + } +} + +/// 撤掉记上这份预留时在滚动窗口里记的那一个请求([`record_request`] 记在那一刻,和预留的 +/// `at_ms` 是同一个数)。请求数一个一个都一样,撤哪一个都行。已经滑出窗口的不用撤 +pub(crate) fn unrecord(books: &mut HashMap, r: &Reservation) { + let Some(b) = books.get_mut(&r.key) else { + return; + }; + if let Some(i) = b + .recent + .iter() + .position(|(t, a)| *t == r.at_ms && a.requests > 0) + { + b.recent[i].1.requests -= 1; + if b.recent[i].1 == Amount::default() { + b.recent.remove(i); + } + } +} + +/// 放掉结束了太久、存储层却一直没交来那一行的预留(见 [`GRACE_MS`])。交回放掉的那几份: +/// 预留的号、请求号 +pub(crate) fn sweep_held( + held: &mut HashMap, + bus: &tw_observe::EventBus, + now: i64, +) -> Vec<(u64, u64)> { + let mut gone = Vec::new(); + for (seq, r) in held.iter_mut() { + let Some(id) = r.request else { continue }; + if bus.is_open(id) { + r.closed_since = None; + continue; + } + let since = *r.closed_since.get_or_insert(now); + if now - since >= GRACE_MS { + gone.push((*seq, id)); + } + } + for (seq, _) in &gone { + held.remove(seq); + } + gone +} + +/// 把滚动窗口里滑出去的扔掉;此刻没有分钟、小时上限的,整个清空 +pub(crate) fn trim_windows( + books: &mut HashMap, + limits: &HashMap>, + now: i64, +) { + let windows: HashMap<&str, i64> = limits + .iter() + .filter_map(|(k, ls)| longest_window(ls).map(|w| (k.as_str(), w))) + .collect(); + for (key, b) in books.iter_mut() { + match windows.get(key.as_str()) { + Some(w) => { + while b.recent.front().is_some_and(|(t, _)| *t <= now - w) { + b.recent.pop_front(); + } + } + None => b.recent.clear(), + } + } +} + +/// 这个名字刚换了新的那几期([`Books::moved`])里还没要过的:记下要过了,交回去,**在锁 +/// 外面**去请求记录里重新加 +pub(crate) fn take_moved( + books: &mut HashMap, + asked: &mut HashSet, + key: &str, +) -> Vec { + let Some(b) = books.get_mut(key) else { + return Vec::new(); + }; + std::mem::take(&mut b.moved) + .into_iter() + .map(|(per, start)| Moved { + key: key.to_string(), + per, + start, + }) + .filter(|m| asked.insert(m.clone())) + .collect() +} diff --git a/crates/tw-gateway/src/key_limits/mod.rs b/crates/tw-gateway/src/key_limits/mod.rs index 477b4719..16016ed0 100644 --- a/crates/tw-gateway/src/key_limits/mod.rs +++ b/crates/tw-gateway/src/key_limits/mod.rs @@ -12,7 +12,7 @@ //! //! **怎么拒、怎么等。**天、周、月是自然周期:用满了就拒,到下一期之前重试也没用。分钟、 //! 小时是滚动窗口:用满了先看下一个空位多久之后空出来,等得到(在这个请求的等待期限之前, -//! 期限见 [`slot_wait`])就等,等不到就拒,并说清楚多久之后再来。顺序见 [`crate::server`] +//! 期限见 [`request_wait`])就等,等不到就拒,并说清楚多久之后再来。顺序见 [`crate::server`] //! 的管线第 3 步:自然周期 → 并发上限 → 滚动窗口。 //! //! **在跑的怎么算。**准入时按输入估一个数占着(预留:输入 token 的估算,和按头一个候选算 @@ -26,35 +26,29 @@ //! 一个跑了三分钟的请求,它的输出要等它跑完才知道,按开始的时刻记的话,「每分钟多少 //! token」永远数不到它。 -use std::collections::{HashMap, HashSet, VecDeque}; +use std::collections::{HashMap, HashSet}; use std::sync::{Arc, Mutex, PoisonError}; use std::time::Duration; use tw_config::{KeyLimit, LimitMeasure, LimitPer}; use tw_types::{Msg, msg}; +pub(crate) mod books; mod clock; #[cfg(test)] mod tests; +use books::{Amount, Books, CALENDAR, Period, Reservation, calendar_index, longest_window}; pub use clock::{Clock, SystemClock, TestClock}; -/// 自然周期。下标就是 [`Books::periods`] 的下标 -const CALENDAR: [LimitPer; 3] = [LimitPer::Day, LimitPer::Week, LimitPer::Month]; - -/// 请求已经结束、存储层却还没交来它那一行时,它的预留再留多久。 -/// -/// 存储层落后、丢了事件时那一行永远不会来:不放掉的话,那份预留一直占着这把密钥的 -/// 额度。存储层正常时几毫秒就到 -const GRACE_MS: i64 = 60_000; - -/// 一个请求最多等多久:`failover.slot_wait_secs`,0 是不等。 +/// 一个请求最多等多久:`failover.wait_secs`,0 是不等。 /// -/// **一个请求合起来算**:等滚动窗口的空位在准入时、发给哪一家之前,等上游空位 -/// (`crate::slots`)在之后的那几跳里,两段共用准入时定下的同一个期限(见 -/// `server::pipeline::admission`)。各给一份的话,两样都碰上的请求能等两倍那么久 -pub fn slot_wait(cfg: &tw_config::Config) -> Duration { - Duration::from_secs(cfg.failover.slot_wait_secs) +/// **一个请求合起来算**:等滚动窗口的空位在准入时、发给哪一家之前,等上游(满着的空位 +/// `crate::slots`、限流暂停和上游的用量上限 `crate::throttle`)在之后的那几跳里,几段共用准入时 +/// 定下的同一个期限(见 `server::pipeline::admission`)。各给一份的话,几样都碰上的请求能等 +/// 好几倍那么久 +pub fn request_wait(cfg: &tw_config::Config) -> Duration { + Duration::from_secs(cfg.failover.wait_secs) } /// 这个路径的请求算不算:数 token 的不算(Anthropic 的 `count_tokens`、Gemini 的 @@ -67,79 +61,6 @@ pub fn uncounted(path: &str) -> bool { p.strip_prefix("/v1").unwrap_or(p) == "/responses/input_tokens" } -/// 一份用量。四样分开记,各条上限各取各的。 -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -struct Amount { - requests: i64, - /// 没走缓存的输入 + 写进缓存的 + 输出 - tokens: i64, - cache_read: i64, - /// 微分 - cost: i64, -} - -impl Amount { - fn of(&self, m: LimitMeasure, cache_reads: bool) -> i64 { - match m { - LimitMeasure::Requests => self.requests, - LimitMeasure::Tokens if cache_reads => self.tokens.saturating_add(self.cache_read), - LimitMeasure::Tokens => self.tokens, - LimitMeasure::Cost => self.cost, - } - } - fn add(&mut self, o: &Amount) { - self.requests = self.requests.saturating_add(o.requests); - self.tokens = self.tokens.saturating_add(o.tokens); - self.cache_read = self.cache_read.saturating_add(o.cache_read); - self.cost = self.cost.saturating_add(o.cost); - } -} - -/// 一期(天、周、月)的合计。 -#[derive(Debug, Clone, Copy)] -struct Period { - start: i64, - end: i64, - sum: Amount, -} - -/// 一把密钥的账。**跨重载存活**,按密钥的名字记 —— 存储层那一行记的也是名字。 -#[derive(Debug, Default)] -struct Books { - /// 天、周、月,和 [`CALENDAR`] 同一个顺序。到了下一期就换一本新的 - periods: [Option; 3], - /// 滚动窗口用:最近的每一笔(请求数在准入时,token 和费用在结算时),按记下的先后。 - /// **只有设了分钟、小时上限的密钥才记**,留到最长的那个窗口为止 - recent: VecDeque<(i64, Amount)>, - /// 开头变了、换了一本新的那几期(哪一种、新的开头),等着从请求记录里重新加(见 - /// [`KeyLimits::reread_with`]) - moved: Vec<(LimitPer, i64)>, -} - -impl Books { - /// `now` 所在的那一期。开头和记着的不一样(到了下一期、换了时区)就换一本新的,记下它 - /// 要从请求记录里重新加;加回来之前先从 0 起。这把密钥从没记过这一种的,从 0 起就是对的 - /// —— 启动时从库里加回来过([`KeyLimits::rebuild`]),之后的每一行都结算过 - fn period(&mut self, clock: &dyn Clock, per: LimitPer, now: i64) -> &mut Period { - let (start, end) = clock.period(per, now); - let slot = &mut self.periods[calendar_index(per)]; - match slot { - Some(p) if p.start == start => {} - _ => { - if slot.is_some() { - self.moved.push((per, start)); - } - *slot = Some(Period { - start, - end, - sum: Amount::default(), - }) - } - } - slot.as_mut().expect("set just above") - } -} - /// 开头变了的一期:哪把密钥、哪一种、新的开头。从请求记录里重新加它([`KeyLimits::reread_with`]) #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct Moved { @@ -153,26 +74,6 @@ pub struct Moved { /// [`KeyLimits::reread`] pub type Reread = Arc; -fn calendar_index(per: LimitPer) -> usize { - CALENDAR.iter().position(|p| *p == per).unwrap_or(0) -} - -/// 一个在跑的请求占着的那一份。 -#[derive(Debug)] -struct Reservation { - key: String, - /// 准入的那一刻。算在哪一期、在不在滚动窗口里看它 - at_ms: i64, - /// 输入 token 的估算 - tokens: i64, - /// 头一个候选的输入费用估算,微分 - cost: i64, - /// 它是哪个请求。准入之后、开始事件之前还没有号 - request: Option, - /// 发现这个请求已经结束的那一刻(见 [`GRACE_MS`]) - closed_since: Option, -} - /// 报过的那一档:哪把密钥、哪一条、哪一期、八成还是到顶。**上限改了就是另一条**, /// 重新报 #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -241,7 +142,7 @@ impl Recorded { self.reached && !self.local && !uncounted(&self.path) } - fn amount(&self) -> Amount { + pub(crate) fn amount(&self) -> Amount { let n = |v: u64| v.min(i64::MAX as u64) as i64; Amount { requests: n(self.requests), @@ -284,6 +185,7 @@ impl Refusal { pub fn msg(&self) -> Msg { let key = self.key.clone(); let per = self.limit.per.word(); + let every = self.limit.every_or_one(); let measure = self.limit.measure(); let (max, used) = match measure { LimitMeasure::Cost => (dollars(self.limit.max()), dollars(self.used)), @@ -308,6 +210,30 @@ impl Refusal { "Gateway key `{key}` has reached its limit of {max} per {per}: {used} spent so \ far. It resets at {resets}." ), + // 每 5 小时这种:句子里说「每 5 个小时」,参数多一个 every,另起一套码 + (None, measure) if every > 1 => { + let retry = self.retry_after_ms.div_ceil(1000).max(1); + match measure { + LimitMeasure::Requests => msg!( + "gw.key_limit.requests_every", key = key, max = max, every = every, + per = per, used = used, retry = retry => + "Gateway key `{key}` has reached its limit of {max} requests every {every} \ + {per}s: {used} in the last {every} {per}s. Try again in {retry} s." + ), + LimitMeasure::Tokens => msg!( + "gw.key_limit.tokens_every", key = key, max = max, every = every, + per = per, used = used, retry = retry => + "Gateway key `{key}` has reached its limit of {max} tokens every {every} \ + {per}s: {used} in the last {every} {per}s. Try again in {retry} s." + ), + LimitMeasure::Cost => msg!( + "gw.key_limit.cost_every", key = key, max = max, every = every, + per = per, used = used, retry = retry => + "Gateway key `{key}` has reached its limit of {max} every {every} {per}s: \ + {used} spent in the last {every} {per}s. Try again in {retry} s." + ), + } + } (None, measure) => { let retry = self.retry_after_ms.div_ceil(1000).max(1); match measure { @@ -395,7 +321,7 @@ pub struct KeyLimits { /// 一期的开头变了时从请求记录里重新加([`Self::reread_with`])。没接上(没有存储层)是 None: /// 那时从 0 起 reread: Mutex>, - /// 报「到了八成、到了上限」,和认请求有没有结束(见 [`GRACE_MS`]) + /// 报「到了八成、到了上限」,和认请求有没有结束(见 [`books::GRACE_MS`]) bus: tw_observe::EventBus, } @@ -452,21 +378,7 @@ impl KeyLimits { /// 这把密钥刚换了新的那几期([`Books::moved`])里还没要过的:记下要过了,交回去,**在锁外面** /// 交给 [`Self::ask_reread`] fn moved(&self, g: &mut Inner, key: &str) -> Vec { - let Some(b) = g.books.get_mut(key) else { - return Vec::new(); - }; - let moved: Vec = std::mem::take(&mut b.moved) - .into_iter() - .map(|(per, start)| Moved { - key: key.to_string(), - per, - start, - }) - .collect(); - moved - .into_iter() - .filter(|m| g.asked.insert(m.clone())) - .collect() + books::take_moved(&mut g.books, &mut g.asked, key) } /// 去请求记录里重新加这几期。没接上记录的,从 0 起 @@ -534,7 +446,7 @@ impl KeyLimits { /// 准入第一步:天、周、月的上限,**用满了就拒**。在并发上限之前:一个这一期已经 /// 用满的请求不用先排一轮队 pub fn calendar(&self, key: &str, limits: &[KeyLimit]) -> Result<(), Box> { - if !limits.iter().any(|l| l.per.rolling_ms().is_none()) { + if !limits.iter().any(|l| l.window_ms().is_none()) { return Ok(()); } let now = self.clock.now_ms(); @@ -568,7 +480,7 @@ impl KeyLimits { } /// 准入第三步:分钟、小时的上限。用满了先等下一个空位,最多等到 `until`(这个请求的 - /// 等待期限,见 [`slot_wait`]);等不到就拒,带着多久之后能再来。**过了就记上**:这个 + /// 等待期限,见 [`request_wait`]);等不到就拒,带着多久之后能再来。**过了就记上**:这个 /// 请求算进滚动窗口,输入的估算占上预留。 /// /// 天、周、月在这里再看一遍:等并发名额、等空位的那一阵,别的请求可能把它用满了。 @@ -632,7 +544,7 @@ impl KeyLimits { let held = g.by_request.remove(&id).and_then(|s| g.held.remove(&s)); if !rec.counts() { if let Some(r) = held { - unrecord(&mut g, &r); + books::unrecord(&mut g.books, &r); } return; } @@ -719,13 +631,13 @@ impl KeyLimits { let max = l.max(); tw_api::KeyLimitView { per: l.per.into(), + every: l.every_or_one(), measure: l.measure().into(), max: max.max(0) as u64, cache_reads: l.cache_reads, used: used.max(0) as u64, resets_at_ms: l - .per - .rolling_ms() + .window_ms() .is_none() .then(|| self.clock.period(l.per, now).1.max(0) as u64), reached: used >= max, @@ -742,52 +654,7 @@ impl KeyLimits { /// 一条上限此刻用了多少:记下的,加上在跑的请求的预留 fn used(&self, g: &mut Inner, key: &str, l: &KeyLimit, now: i64) -> i64 { - let (m, cr) = (l.measure(), l.cache_reads); - let reserved = |r: &Reservation| match m { - LimitMeasure::Requests => 1, - LimitMeasure::Tokens => r.tokens, - LimitMeasure::Cost => r.cost, - }; - match l.per.rolling_ms() { - None => { - let p = *g.books.entry(key.to_string()).or_default().period( - self.clock.as_ref(), - l.per, - now, - ); - let held: i64 = g - .held - .values() - .filter(|r| r.key == key && r.at_ms >= p.start && r.at_ms < p.end) - .map(reserved) - .sum(); - p.sum.of(m, cr).saturating_add(held) - } - Some(w) => { - let cutoff = now - w; - let books = g.books.get(key); - let done: i64 = books - .map(|b| { - b.recent - .iter() - .filter(|(t, _)| *t > cutoff) - .map(|(_, a)| a.of(m, cr)) - .sum() - }) - .unwrap_or(0); - // 请求数在准入时已经进了窗口,预留里只算 token 和费用 - let held: i64 = g - .held - .values() - .filter(|r| r.key == key && r.at_ms > cutoff) - .map(|r| match m { - LimitMeasure::Requests => 0, - _ => reserved(r), - }) - .sum(); - done.saturating_add(held) - } - } + books::used(&mut g.books, &g.held, self.clock.as_ref(), key, l, now) } fn calendar_locked( @@ -797,7 +664,7 @@ impl KeyLimits { limits: &[KeyLimit], now: i64, ) -> Result<(), Box> { - for l in limits.iter().filter(|l| l.per.rolling_ms().is_none()) { + for l in limits.iter().filter(|l| l.window_ms().is_none()) { let used = self.used(g, key, l, now); if used >= l.max() { let end = self.clock.period(l.per, now).1; @@ -822,55 +689,12 @@ impl KeyLimits { now: i64, ) -> Option { let mut worst: Option = None; - for l in limits { - let Some(w) = l.per.rolling_ms() else { - continue; - }; + for l in limits.iter().filter(|l| l.window_ms().is_some()) { let used = self.used(g, key, l, now); - let max = l.max(); - if used < max { + if used < l.max() { continue; } - // 窗口里的每一笔按时间先后滑出去,滑到用量低于上限的那一刻就是空位 - let (m, cr) = (l.measure(), l.cache_reads); - let cutoff = now - w; - let mut items: Vec<(i64, i64)> = g - .books - .get(key) - .map(|b| { - b.recent - .iter() - .filter(|(t, _)| *t > cutoff) - .map(|(t, a)| (*t, a.of(m, cr))) - .collect() - }) - .unwrap_or_default(); - if m != LimitMeasure::Requests { - items.extend( - g.held - .values() - .filter(|r| r.key == key && r.at_ms > cutoff) - .map(|r| { - let v = if m == LimitMeasure::Tokens { - r.tokens - } else { - r.cost - }; - (r.at_ms, v) - }), - ); - } - items.retain(|(_, v)| *v > 0); - items.sort_by_key(|(t, _)| *t); - let mut left = used; - let mut free_at = now + w; - for (t, v) in items { - left -= v; - if left < max { - free_at = t + w; - break; - } - } + let free_at = books::free_at(&g.books, &g.held, key, l, used, now); let wait = (free_at - now).max(1) as u64; if worst.as_ref().is_none_or(|r| wait > r.retry_after_ms) { worst = Some(Refusal { @@ -887,32 +711,11 @@ impl KeyLimits { /// 记上这个请求:滚动窗口里的一个请求,和它的预留。交回预留的号 fn reserve(&self, g: &mut Inner, key: &str, limits: &[KeyLimit], ask: Ask, now: i64) -> u64 { - if let Some(w) = longest_window(limits) { - let books = g.books.entry(key.to_string()).or_default(); - books.recent.push_back(( - now, - Amount { - requests: 1, - ..Default::default() - }, - )); - while books.recent.front().is_some_and(|(t, _)| *t <= now - w) { - books.recent.pop_front(); - } - } + books::record_request(&mut g.books, key, limits, now); g.seq += 1; let seq = g.seq; - g.held.insert( - seq, - Reservation { - key: key.to_string(), - at_ms: now, - tokens: ask.tokens.min(i64::MAX as u64) as i64, - cost: ask.cost_micros.max(0), - request: None, - closed_since: None, - }, - ); + g.held + .insert(seq, Reservation::new(key, now, ask.tokens, ask.cost_micros)); seq } @@ -932,44 +735,17 @@ impl KeyLimits { if let Some(id) = r.request { g.by_request.remove(&id); } - unrecord(&mut g, &r); + books::unrecord(&mut g.books, &r); } } - /// 放掉结束了太久、存储层却一直没交来那一行的预留(见 [`GRACE_MS`])。顺手把滚动 + /// 放掉结束了太久、存储层却一直没交来那一行的预留(见 [`books::GRACE_MS`])。顺手把滚动 /// 窗口里滑出去的扔掉 fn sweep(&self, g: &mut Inner, now: i64) { - let mut gone = Vec::new(); - for (seq, r) in g.held.iter_mut() { - let Some(id) = r.request else { continue }; - if self.bus.is_open(id) { - r.closed_since = None; - continue; - } - let since = *r.closed_since.get_or_insert(now); - if now - since >= GRACE_MS { - gone.push((*seq, id)); - } - } - for (seq, id) in gone { - g.held.remove(&seq); + for (_, id) in books::sweep_held(&mut g.held, &self.bus, now) { g.by_request.remove(&id); } - let windows: HashMap = g - .limits - .iter() - .filter_map(|(k, ls)| longest_window(ls).map(|w| (k.clone(), w))) - .collect(); - for (key, b) in g.books.iter_mut() { - match windows.get(key) { - Some(w) => { - while b.recent.front().is_some_and(|(t, _)| *t <= now - w) { - b.recent.pop_front(); - } - } - None => b.recent.clear(), - } - } + books::trim_windows(&mut g.books, &g.limits, now); } /// 天、周、月的上限到了八成、到了顶:**每一期、每一档报一次**。只看记下的数 —— @@ -979,7 +755,7 @@ impl KeyLimits { return Vec::new(); }; let mut out = Vec::new(); - for l in limits.iter().filter(|l| l.per.rolling_ms().is_none()) { + for l in limits.iter().filter(|l| l.window_ms().is_none()) { let p = *g.books.entry(key.to_string()).or_default().period( self.clock.as_ref(), l.per, @@ -1065,29 +841,6 @@ impl KeyLimits { } } -/// 撤掉准入时给这个请求在滚动窗口里记的那一个请求([`KeyLimits::reserve`] 记在准入的那一刻, -/// 和预留的 `at_ms` 是同一个数)。请求数一个一个都一样,撤哪一个都行。已经滑出窗口的不用撤 -fn unrecord(g: &mut Inner, r: &Reservation) { - let Some(b) = g.books.get_mut(&r.key) else { - return; - }; - if let Some(i) = b - .recent - .iter() - .position(|(t, a)| *t == r.at_ms && a.requests > 0) - { - b.recent[i].1.requests -= 1; - if b.recent[i].1 == Amount::default() { - b.recent.remove(i); - } - } -} - -/// 一把密钥的分钟、小时上限里最长的那个窗口。没有就是 None -fn longest_window(limits: &[KeyLimit]) -> Option { - limits.iter().filter_map(|l| l.per.rolling_ms()).max() -} - /// 这把密钥用得到、却没有价格的模型(`ClientView.unpriced_models`)。 /// /// **和 `GET /v1/models` 同一份清单**:目录里这把密钥的 `allow` 放行的名称(不挑格式 —— diff --git a/crates/tw-gateway/src/key_limits/tests.rs b/crates/tw-gateway/src/key_limits/tests.rs index 38c87f7a..6668a103 100644 --- a/crates/tw-gateway/src/key_limits/tests.rs +++ b/crates/tw-gateway/src/key_limits/tests.rs @@ -255,7 +255,7 @@ async fn a_rolling_minute_waits_for_its_next_slot_or_refuses() { assert_eq!(b.limits.clock.now_ms(), t0 + 60_000, "等到第一个滑出窗口"); } -/// 等待期限是整个请求的那一个(见 [`slot_wait`]),调用方给的是那一刻,不是从此刻起再给 +/// 等待期限是整个请求的那一个(见 [`request_wait`]),调用方给的是那一刻,不是从此刻起再给 /// 一整段:期限之前空不出来的就拒 #[tokio::test(start_paused = true)] async fn a_rolling_wait_ends_at_the_requests_own_deadline() { @@ -299,6 +299,72 @@ async fn tokens_in_a_rolling_hour_count_when_they_are_settled() { assert_eq!(r.msg().code, "gw.key_limit.tokens_rolling"); } +/// `every`:窗口是几个单位那么长。每 5 分钟 2 个请求,第一个要到 5 分钟时才滑出去 +#[tokio::test(start_paused = true)] +async fn every_makes_the_rolling_window_that_many_units_long() { + let b = bed( + "2026-10-05T10:00:00+08:00", + "[{per: minute, every: 5, requests: 2}]", + ); + let t0 = b.now(); + b.run(1, Ask::default()).await.unwrap(); + tokio::time::advance(Duration::from_secs(60)).await; + b.run(2, Ask::default()).await.unwrap(); + // 一分钟之后第一个还在窗口里:每分钟的话它已经滑出去了 + tokio::time::advance(Duration::from_secs(60)).await; + let r = b.try_admit(Ask::default()).await.unwrap_err(); + assert_eq!(r.retry_after_ms, 180_000, "第一个在 5 分钟时滑出去"); + let m = r.msg(); + assert_eq!(m.code, "gw.key_limit.requests_every"); + assert_eq!(m.arg("every"), "5"); + assert_eq!( + m.text, + "Gateway key `k` has reached its limit of 2 requests every 5 minutes: 2 in the last 5 \ + minutes. Try again in 180 s." + ); + let v = b.limits.view("k", &b.set); + assert_eq!((v[0].every, v[0].used, v[0].reached), (5, 2, true)); + // 等得到就等:到第一个滑出去为止 + let hold = b + .limits + .admit("k", &b.set, Ask::default(), Duration::from_secs(200)) + .await + .unwrap(); + drop(hold); + assert_eq!(b.now(), t0 + 300_000); +} + +#[tokio::test(start_paused = true)] +async fn every_on_tokens_and_cost_names_its_window_too() { + let b = bed( + "2026-10-05T10:00:00+08:00", + "[{per: hour, every: 3, tokens: 100}, {per: hour, every: 2, cost: 1}]", + ); + b.run(1, ask(0, 0)).await.unwrap(); + b.done(1, row(60, 60, 0, 2_000_000)); + let r = b.try_admit(Ask::default()).await.unwrap_err(); + // 两条都满了:报等得久的那条 + assert_eq!(r.retry_after_ms, 3 * 3_600_000); + assert_eq!(r.msg().code, "gw.key_limit.tokens_every"); + // 两小时之后费用那条滑出去了,token 那条还满着 + tokio::time::advance(Duration::from_secs(2 * 3600)).await; + let r = b.try_admit(Ask::default()).await.unwrap_err(); + assert_eq!(r.retry_after_ms, 3_600_000); + let b = bed( + "2026-10-05T10:00:00+08:00", + "[{per: hour, every: 2, cost: 1}]", + ); + b.run(1, ask(0, 0)).await.unwrap(); + b.done(1, row(1, 1, 0, 1_500_000)); + let m = b.try_admit(Ask::default()).await.unwrap_err().msg(); + assert_eq!(m.code, "gw.key_limit.cost_every"); + assert_eq!( + m.text, + "Gateway key `k` has reached its limit of $1.00 every 2 hours: $1.50 spent in the last 2 \ + hours. Try again in 7200 s." + ); +} + // ───────────────────────────────────────────────── 预留与结算 #[tokio::test(start_paused = true)] @@ -676,4 +742,5 @@ async fn the_view_says_how_much_is_left_and_when_it_resets() { (v[1].used, v[1].reached, v[1].resets_at_ms), (1, true, None) ); + assert_eq!((v[0].every, v[1].every), (1, 1)); } diff --git a/crates/tw-gateway/src/lib.rs b/crates/tw-gateway/src/lib.rs index 819d115f..ec4370e2 100644 --- a/crates/tw-gateway/src/lib.rs +++ b/crates/tw-gateway/src/lib.rs @@ -56,8 +56,10 @@ pub mod session; pub mod slots; mod splice; pub mod state; +pub mod throttle; pub mod traffic; pub mod translate; +pub mod upstream_limits; pub mod wire; pub mod ws; diff --git a/crates/tw-gateway/src/server.rs b/crates/tw-gateway/src/server.rs index b96db449..2ad1b029 100644 --- a/crates/tw-gateway/src/server.rs +++ b/crates/tw-gateway/src/server.rs @@ -385,6 +385,7 @@ pub(crate) fn hop( usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, } } @@ -412,6 +413,86 @@ pub(crate) fn hop_busy( usage: None, queued_ms, skipped: Some(tw_api::ServeSkip::Busy), + pause_ms: None, + proxy: None, + } +} + +/// 这家暂时发不出去、没发出去的一跳(见 [`crate::throttle`]):限流暂停着,或者分钟、小时的用量 +/// 上限满了。`until_ms` 是预计又能发的那一刻(说不出的是 None),`queued_ms`:等过它的话等了多久。 +/// +/// **不是失败**,和满着的一样:不停用、不进熔断的账,这一行只说明请求为什么去了下一家、或者在等它。 +pub(crate) fn hop_waiting( + provider: &str, + model: Option, + kind: tw_api::ThrottleKind, + until_ms: Option, + queued_ms: Option, + started: std::time::Instant, +) -> tw_api::AttemptView { + let secs = until_ms + .map(|t| t.saturating_sub(now_ms())) + .unwrap_or(crate::throttle::PROBE_GRACE.as_millis() as u64) + .div_ceil(1000) + .max(1); + let (error, skipped) = match kind { + tw_api::ThrottleKind::WindowFull => ( + msg!( + "gw.upstream.window_full", upstream = provider, secs = secs => + "Upstream `{upstream}` is at one of its per-minute or per-hour usage limits; it has \ + room again in about {secs} s." + ), + tw_api::ServeSkip::WindowFull, + ), + _ => ( + msg!( + "gw.upstream.paused", upstream = provider, secs = secs => + "Upstream `{upstream}` answered 429 (rate limited) and is paused; it takes requests \ + again in about {secs} s." + ), + tw_api::ServeSkip::Throttled, + ), + }; + tw_api::AttemptView { + provider: provider.to_string(), + model, + outcome: tw_api::AttemptOutcome::Error, + status: None, + error: Some(error), + ms: started.elapsed().as_millis() as u64, + usage: None, + queued_ms, + skipped: Some(skipped), + pause_ms: None, + proxy: None, + } +} + +/// 这家一条天、周、月的用量上限用完了、跳过了的一跳。`until_ms` 是它重置的那一刻 +pub(crate) fn hop_used_up( + provider: &str, + model: Option, + until_ms: u64, + started: std::time::Instant, +) -> tw_api::AttemptView { + use crate::key_limits::Clock; + let resets = crate::key_limits::SystemClock.show(i64::try_from(until_ms).unwrap_or(i64::MAX)); + tw_api::AttemptView { + provider: provider.to_string(), + model, + outcome: tw_api::AttemptOutcome::Error, + status: None, + error: Some(msg!( + "gw.upstream.period_used_up", + upstream = provider, resets = resets, resets_at_ms = until_ms => + "Upstream `{upstream}` has used up one of its usage limits for the current day, week \ + or month; it resets at {resets}." + )), + ms: started.elapsed().as_millis() as u64, + usage: None, + queued_ms: None, + skipped: Some(tw_api::ServeSkip::PeriodUsedUp), + pause_ms: None, proxy: None, } } @@ -433,6 +514,7 @@ pub(crate) fn hop_failed( usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, } } diff --git a/crates/tw-gateway/src/server/pipeline.rs b/crates/tw-gateway/src/server/pipeline.rs index 61f20f39..bd15a3fc 100644 --- a/crates/tw-gateway/src/server/pipeline.rs +++ b/crates/tw-gateway/src/server/pipeline.rs @@ -80,7 +80,7 @@ struct Started { /// 那些字节时(同格式直通、一个字节都没改),换的时候照它换,不再找一遍(见 [`hop`]) hits: Option>, /// 这个请求最多等到什么时候:准入时定下(见 [`admission`]),等密钥的分钟、小时上限和 - /// 等上游的空位共用这一段(`failover.slot_wait_secs`) + /// 等上游的空位共用这一段(`failover.wait_secs`) wait_until: tokio::time::Instant, } @@ -941,13 +941,22 @@ pub(super) fn arrange<'a>( facts_rt = Some(f); } // 留在上次回答这段对话的那一家:同一轮里一律留,跨轮看缓存值不值得留。**排在 - // 策略组排序之后** —— 该留的时候盖过策略,放开的时候策略照常说了算 + // 策略组排序之后** —— 该留的时候盖过策略,放开的时候策略照常说了算。那一家熔断着、这一期 + // 用完了、限流暂停到这个请求等不到的时候(见 `crate::throttle`)放开 + let deadline_ms = now.saturating_add(rt.config.failover.wait_secs.saturating_mul(1000)); let stayed = conv.and_then(|c| { state.affinity.stay( c, decision.via_group.as_deref(), &mut decision.candidates, - |p| state.health.is_available(p), + |p| { + state.health.is_available(p) + && state.throttles.keeps_conversation( + p, + hop::usage_limits_of(rt, p), + deadline_ms, + ) + }, now, ) }); diff --git a/crates/tw-gateway/src/server/pipeline/admission.rs b/crates/tw-gateway/src/server/pipeline/admission.rs index b6daa367..1d862ddd 100644 --- a/crates/tw-gateway/src/server/pipeline/admission.rs +++ b/crates/tw-gateway/src/server/pipeline/admission.rs @@ -9,7 +9,7 @@ //! 说清多久之后再来。过了就把这个请求记上,按输入估一个数占着,等存储层记下它那一行时 //! 换成实数。 //! -//! **等待期限一个请求只有一个**:过了并发上限那一刻起算 `failover.slot_wait_secs`,这里等 +//! **等待期限一个请求只有一个**:过了并发上限那一刻起算 `failover.wait_secs`,这里等 //! 滚动窗口的空位、之后等上游的空位(见 `hop`)都算在里面 —— 两段各给一份的话,一个请求能 //! 等两倍那么久。并发上限那一段不算:它等前面的请求结束,不拒绝,等多久由客户端决定(见 //! [`crate::limits`])。 @@ -59,7 +59,7 @@ pub(super) async fn admit( } let pass = state.gate.acquire(&req.client_name, max_concurrent).await; // 等待期限从这里起算:之后的两段等待共用它 - let wait_until = tokio::time::Instant::now() + crate::key_limits::slot_wait(&rt.config); + let wait_until = tokio::time::Instant::now() + crate::key_limits::request_wait(&rt.config); let ask = if limits.is_empty() { Ask::default() } else { diff --git a/crates/tw-gateway/src/server/pipeline/hop.rs b/crates/tw-gateway/src/server/pipeline/hop.rs index 62ac0e0a..36d6d9b0 100644 --- a/crates/tw-gateway/src/server/pipeline/hop.rs +++ b/crates/tw-gateway/src/server/pipeline/hop.rs @@ -49,7 +49,8 @@ pub(super) struct Served<'a> { pub(super) slot: crate::slots::Slot, /// 这一跳发出去的那一刻。这一家的快慢样本从这里算起(见 [`crate::latency`]) pub(super) sent_at: std::time::Instant, - /// 这一家答没答上还不知道:最后一家的开头不压,成败的账等回程看到第一段内容才记(见 + /// 这一家答没答上还不知道:最后一家的开头没压(整包的、等不起的),或者压着时到点了还没有 + /// 内容、在开头报了限流以外的错,原样交了出去。成败的账等回程看到第一段内容才记(见 /// `relay`)。没有内容就超时了记失败,和压着开头时一样 pub(super) unconfirmed: bool, } @@ -161,49 +162,90 @@ pub(super) async fn try_upstreams<'a>( let quiet = super::idle::Quiet::of(state, rt); // 还没看的候选,按顺序 let mut queue: std::collections::VecDeque<&String> = started.alive.iter().collect(); - // 满着没发的那几家(见 `crate::slots`),按候选的顺序。候选都看过一遍还没有着落时, - // 在它们里面等先空出来的那一家 - let mut busy: Vec<&String> = Vec::new(); - // 等空位等到什么时候:准入时定下的那一刻(见 `super::admission`),**整个请求共用** - // —— 等密钥的分钟、小时上限已经用掉的,这里就少等那么久。等的时候客户端一个字节都 - // 收不到 + // 暂时发不出去、等得到的那几家(见 `crate::slots`、`crate::throttle`):满着的、限流暂停着 + // 的、分钟或小时的用量上限满了的,和跳过时它是哪一样,按候选的顺序。候选都看过一遍还没有 + // 着落时,在它们里面等先能发的那一家 + let mut waiting: Vec<(&String, Wait)> = Vec::new(); + // 这一期用完了、跳过了的那几家,和它们重置的时刻。一家都发不出去时,429 里说到它们 + let mut spent: Vec<(&String, u64)> = Vec::new(); + // 等到什么时候:准入时定下的那一刻(见 `super::admission`),**整个请求共用** —— 等密钥的 + // 分钟、小时上限已经用掉的,这里就少等那么久。流式的回答等过一阵先交出 `200` 和保活(见 + // `super::commit`),整包的等的时候客户端一个字节都收不到 let until = started.wait_until; - // 等过空位的那一跳在尝试链上的位置和等了多久。那一跳怎么收场都只进一行,进了之后补上 + let deadline_ms = now_ms() + + until + .saturating_duration_since(tokio::time::Instant::now()) + .as_millis() as u64; + // 这个请求在哪几家的队里排着(见 `crate::throttle`)。**丢掉就离开**:客户端走了、被中止了, + // 后面的不用替它等 + let mut waiter = state.throttles.waiter(); + // 等过的那一跳在尝试链上的位置和等了多久。那一跳怎么收场都只进一行,进了之后补上 let mut queued: Option<(usize, u64)> = None; - // 等到最后,剩下的候选还都满着 + // 等到最后,剩下的候选还都发不出去 let mut stalled = false; // 有一跳真的发到了上游、失败了(回了错误、流在内容之前断了、连不上、没有内容超时了)。 - // 等满着的那几家等到最后的话,交出去的是这一家的错,不是「都满着」(见下面的 `stalled`) + // 等发不出去的那几家等到最后的话,交出去的是这一家的错,不是「都发不出去」(见下面的 + // `stalled`)。**回 429 的不算**:那是让等一会儿,不是失败 let mut reached = false; loop { stamp_queued(&mut chain, queued.take()); - // 下一个候选。都看过了,就在满着的那几家里等先空出来的:等到的那一家带着占到的位置 + // 下一个候选。都看过了,就在等得到的那几家里等先能发的:等到的那一家带着占到的位置 let (name, mut granted) = match queue.pop_front() { Some(name) => (name, None), - None if busy.is_empty() => break, + None if waiting.is_empty() => break, None => { + // 每一家都要到等待期限之后才能发:不用干等,当场说清多久之后再来。这个请求能等的 + // 已经等完了(或者配置的是不等)也一样 —— **它没排过队**,不报排队(界面会显示成 + // 「排队中」) + if tokio::time::Instant::now() >= until + || hopeless(state, rt, &waiting, deadline_ms) + { + stalled = true; + break; + } + let names: Vec<(&str, &[tw_config::KeyLimit])> = waiting + .iter() + .map(|(n, _)| (n.as_str(), usage_limits_of(rt, n))) + .collect(); + queued_up(state, id, &names, deadline_ms); let t = std::time::Instant::now(); - let freed = tokio::select! { + let ready = tokio::select! { biased; - freed = state.slots.first_free(&busy, until) => freed, - // 等空位的时候被手动中止:一个字节都还没发出去,尝试链上不多一跳 + ready = crate::throttle::first_ready(state, &mut waiter, &names, until) => ready, + // 等的时候被手动中止:一个字节都还没发出去,尝试链上不多一跳 _ = req.abort.wait() => { + dequeued(state, id); halt = Some(GatewayError::aborted()); break; } }; - match freed { - Some((k, slot)) => (busy.remove(k), Some((slot, t.elapsed()))), - None => { + dequeued(state, id); + match ready { + crate::throttle::Outcome::Ready(k, ready) => { + (waiting.remove(k).0, Some((ready, t.elapsed()))) + } + crate::throttle::Outcome::Late => { stalled = true; break; } + // 等着的都不用等了:这一期用完了的记下来(429 里说到它们),停用了的交给熔断 + // (尝试链上已经有跳过它们的那一行) + crate::throttle::Outcome::Gone => { + for (n, _) in std::mem::take(&mut waiting) { + if let crate::throttle::Peek::UsedUp { until_ms, .. } = + state.throttles.peek(n, usage_limits_of(rt, n)) + { + spent.push((n, until_ms)); + } + } + continue; + } } } }; - // 后面没有别的候选了。满着、还等得到的那几家也算在后面 - let last = queue.is_empty() && busy.is_empty(); + // 后面没有别的候选了。暂时发不出去、还等得到的那几家也算在后面 + let last = queue.is_empty() && waiting.is_empty(); let Some(provider) = rt.config.providers.iter().find(|p| &p.name == name) else { // 校验时挡过一次,能到这儿说明配置在运行中被换过。 last_err = Some(GatewayError::config(msg!( @@ -336,58 +378,132 @@ pub(super) async fn try_upstreams<'a>( } } - // 并发上限:这一跳要发出去了,先占这家一个位置(见 `crate::slots`)。**在插件和 - // 转换之前**:满着没发的一跳不跑插件、不报转换。等不是失败:不停用、不进熔断的账。 + // 这一家此刻发得出去吗(见 `crate::throttle`):限流暂停着、用量上限满了或用完了,和 + // 并发上限(见 `crate::slots`)—— 要发出去了,先占这家一个位置。**在插件和转换之前**: + // 没发的一跳不跑插件、不报转换。等不是失败:不停用、不进熔断的账。 // - // 数 token 不占位置:它不跑模型,一眨眼就回来;为它排在几个长回答后面等上半分钟、 - // 最后回一个 429,客户端连上下文还剩多少都看不到 - let slot = match granted.take() { - Some((slot, waited)) => { + // 数 token 不排队、不占位置:它不跑模型,一眨眼就回来;为它排在几个长回答后面等上半分钟、 + // 最后回一个 429,客户端连上下文还剩多少都看不到。这一家暂时发不出去的,由网关估 + let limits = crate::throttle::usage_limits(provider); + let (slot, mut probe) = match granted.take() { + Some((ready, waited)) => { queued = Some((chain.len(), waited.as_millis() as u64)); hop_started = std::time::Instant::now(); - slot + (ready.slot, ready.probe) } - None if counting => crate::slots::Slot::free(), - None => { - let mut waited = None; - let mut slot = state.slots.try_take(&provider.name); - // 这段对话留在这家是为了它的缓存:等它空出来,等不到再换下一家(缓存就丢在 - // 这家了)。别的候选满着当场跳过 - if slot.is_none() && started.choice.stayed_on.as_ref() == Some(name) { - // 这个请求能等的已经等完了(或者配置的是不等):不再等 - if tokio::time::Instant::now() < until { - let t = std::time::Instant::now(); - slot = tokio::select! { - biased; - slot = state.slots.take_by(&provider.name, until) => slot, - _ = req.abort.wait() => { - halt = Some(GatewayError::aborted()); - break; - } - }; - waited = Some(t.elapsed().as_millis() as u64); - hop_started = std::time::Instant::now(); - } + None if counting => { + if state.throttles.peek(&provider.name, limits) != crate::throttle::Peek::Open { + chain.push(estimated_hop(provider, None, None, hop_started)); + estimated = Some(provider); + break; } - match slot { - Some(slot) => { - queued = waited.map(|ms| (chain.len(), ms)); - slot - } - None => { - chain.push(crate::server::hop_busy( + (crate::slots::Slot::free(), None) + } + None => { + let model = Some(sent.clone()).filter(asked_other); + // 这段对话留在这家是为了它的缓存:等它,等不到再换下一家(缓存就丢在这家了)。 + // 别的候选发不出去当场跳过,记下来,候选都看过之后再等 + let sticky = started.choice.stayed_on.as_ref() == Some(name); + let wait = match state + .throttles + .gate(&provider.name, limits, Some(waiter.id)) + { + // 这一期用完了:到下一期之前等多久都一样,跳过,不等 + crate::throttle::Gate::UsedUp { until_ms, .. } => { + chain.push(crate::server::hop_used_up( &provider.name, - Some(sent.clone()).filter(asked_other), - state.slots.limit(&provider.name).unwrap_or_default(), - waited, + model, + until_ms, hop_started, )); - busy.push(name); + spent.push((name, until_ms)); continue; } + crate::throttle::Gate::Go(probe) => { + match state.slots.try_take(&provider.name) { + Some(slot) => Ok((slot, probe)), + None => Err((Wait::Busy, None)), + } + } + crate::throttle::Gate::Wait { kind, until_ms } => { + Err((Wait::of(kind), until_ms)) + } + }; + match wait { + Ok(go) => go, + Err((why, until_ms)) => { + // 留下的那一家等得到才等:说得出什么时候恢复、又晚于这个请求的等待期限的, + // 等也白等。这个请求能等的已经等完了(或者配置的是不等)也不再等 + let mut waited = None; + if sticky + && until_ms.is_none_or(|t| t <= deadline_ms) + && tokio::time::Instant::now() < until + { + let one = [(provider.name.as_str(), limits)]; + queued_up(state, id, &one, deadline_ms); + let t = std::time::Instant::now(); + let ready = tokio::select! { + biased; + ready = crate::throttle::first_ready(state, &mut waiter, &one, until) => ready, + _ = req.abort.wait() => { + dequeued(state, id); + halt = Some(GatewayError::aborted()); + break; + } + }; + dequeued(state, id); + waited = Some(t.elapsed().as_millis() as u64); + hop_started = std::time::Instant::now(); + if let crate::throttle::Outcome::Ready(_, ready) = ready { + queued = waited.map(|ms| (chain.len(), ms)); + (ready.slot, ready.probe) + } else if let crate::throttle::Peek::UsedUp { until_ms, .. } = + state.throttles.peek(&provider.name, limits) + { + // 等的时候这一期用完了:跳过,不再等它 + chain.push(crate::server::hop_used_up( + &provider.name, + model, + until_ms, + hop_started, + )); + spent.push((name, until_ms)); + continue; + } else { + chain.push(skipped_hop( + state, + provider, + model, + (why, until_ms), + waited, + hop_started, + )); + waiting.push((name, why)); + continue; + } + } else { + chain.push(skipped_hop( + state, + provider, + model, + (why, until_ms), + waited, + hop_started, + )); + waiting.push((name, why)); + continue; + } + } } } }; + // 这一家的用量上限:要发出去了,请求数此刻记上,token 和费用按请求估一个数占着(见 + // `crate::upstream_limits`)。没发出去的还回去(丢掉就还),发出去了的见下面的 [`Spent`]。 + // 数 token 的不算 + let usage_hold = (!counting && !limits.is_empty()).then(|| { + let ask = usage_ask(state, reading, provider, &sent, limits); + state.upstream_limits.reserve(&provider.name, limits, &ask) + }); // 插件的请求钩子:管这一跳的从客户端的原话起改。**拒绝的是整个请求**,不换下一家 let plugged = match super::plug::attempt( @@ -568,6 +684,17 @@ pub(super) async fn try_upstreams<'a>( // `super::idle`)。之前等空位、跑插件的时间都不算 let sent_at = std::time::Instant::now(); let quiet_by = quiet.after(sent_at); + // 这一家的答复也看这一刻(见 `crate::throttle`):最近一次 429 之前发出去的不算数。 + // 暂停结束后探路的就是这一跳的话,从这一刻起计它的宽限 + if let Some(p) = probe.as_mut() { + p.sent(); + } + // 发出去了:这一跳在这一家的用量上限上**算一个请求**,怎么收场都算(失败的、超时的、被中止 + // 的上游都已经收到了它),只有回了 429 的、其实没发出去的(连不上、签不了名)还回去 + let mut usage_hold = Spent { + hold: usage_hold, + id, + }; // 发出去、等响应头。**等着的这个 future 只活在这一块里**:放弃这一家、被中止时它跟着 // 丢掉,连接随之断开 @@ -620,14 +747,18 @@ pub(super) async fn try_upstreams<'a>( }; let sent = match waited { Waited::Done(sent) => { - // 上游的响应头(重发过的是最后那一遍的) + // 上游的响应头(重发过的是最后那一遍的)。这一家在不在限流(见 `crate::throttle`) + // 等看清了这个答复再说:压着开头的回答,流的开头还可能报限流 if let Ok(r) = &sent { tap.answered(r); } sent } - // 被手动中止:这一跳到此为止,不再换下一家。不停用、不算失败 + // 被手动中止:这一跳到此为止,不再换下一家。不停用、不算失败;探路的名额还回去 Waited::Aborted => { + if let Some(p) = probe.take() { + p.withdraw(); + } chain.push(super::idle::aborted_hop( &provider.name, model.clone(), @@ -640,6 +771,7 @@ pub(super) async fn try_upstreams<'a>( // 响应头都没等到:放弃这一家,换下一家;没有下一家了,交出去的就是这个超时 Waited::Quiet => { reached = true; + state.throttles.released(&provider.name, sent_at); last_err = Some(gave_up( state, started, @@ -663,6 +795,7 @@ pub(super) async fn try_upstreams<'a>( // 同格式的上游没实现数 token(不少中转只做了生成回答):网关自己估。 // **在判断换不换之前**:404 平时是「这家没有这个模型」,要换下一家 if counting && crate::count::unsupported(status) { + state.throttles.answered(&provider.name, sent_at); note_health( &state.bus, &state.health, @@ -685,17 +818,54 @@ pub(super) async fn try_upstreams<'a>( state.note_glm_429_body(id, provider, &head); } let verdict = crate::failure::classify(status, &headers, &head, now_ms()); + if !matches!(verdict, Verdict::Failed(Cause::RateLimited { .. })) { + to_throttle(state, &provider.name, sent_at, verdict); + } // **最后一家的 4xx 交给客户端**:没有下一家可换了,上游自己的原话 // (「invalid x-api-key」)比我们转述的一句「回了 401」有用得多。 // 5xx 和 429 照旧报我们的话,带着尝试链 let hand_on = last && status < 500 && status != 429; match verdict { + // 限流(额度用完不算):**不是失败**。这一家暂停(见 `crate::throttle`),这个请求 + // 记下它、接着看下一家,候选都看过了就等它恢复 —— 客户端什么都还没收到,再发 + // 一遍它察觉不到。不进熔断的账,也不当成交出去的那条错误 + Verdict::Failed(Cause::RateLimited { .. }) => { + tap.said(&head); + state.note_quota(id, &provider.name, &headers); + state.note_proxy_ok(&provider.proxy); + state.glm_traffic(provider); + let hint = crate::throttle::pause_hint(&headers, &head, now_ms()); + let pause = state.throttles.limited(&provider.name, sent_at, hint); + usage_hold.release(); + let mut a = hop( + &provider.name, + model.clone(), + tw_api::AttemptOutcome::Status, + status, + hop_started, + ); + a.pause_ms = Some(pause.as_millis() as u64); + chain.push(a); + // 数 token 不等:网关估 + if counting { + chain.push(estimated_hop(provider, model.clone(), None, hop_started)); + estimated = Some(provider); + break; + } + // 探路的又被限流了:它下一次等的时候回到这一家的队头,后面的照旧排在它 + // 后面。**先去试别的候选时不占着**(见 `throttle::Waiter::keep_place`) + if probe.is_some() { + waiter.keep_place(&provider.name); + } + waiting.push((name, Wait::Throttled)); + continue; + } Verdict::Failed(cause) if !hand_on => { // 换下一家:这一跳的回答就是读来认原因的那一截 tap.said(&head); - // 5xx、限流、没钱了、额度用完、凭据被拒、没有这个模型:换一家有 - // 意义,那边是另一把密钥、另一个账户。停用多久看原因。后面只剩满着 - // 的那几家也一样等它们:空出来的那一家可能答得上 + // 5xx、没钱了、额度用完、凭据被拒、没有这个模型:换一家有意义,那边是 + // 另一把密钥、另一个账户。停用多久看原因。后面只剩暂时发不出去的那几家 + // 也一样等它们:先能发的那一家可能答得上 reached = true; // // 交给客户端的那一跳,下面这几样由回程(`relay`)去记 @@ -707,6 +877,11 @@ pub(super) async fn try_upstreams<'a>( state.glm_traffic(provider); // 凭据被拒是要让用户知道的状态,**换了下一家也一样** state.note_auth(&provider.name, status); + // 回了 429 的(额度用完)上游没接:不算进这家的用量上限,和记录里一样 + // (`tw_api::AttemptView::counts_for_upstream`) + if status == 429 { + usage_hold.release(); + } let cause = known_reset(state, &provider.name, cause); note_health( &state.bus, @@ -722,7 +897,8 @@ pub(super) async fn try_upstreams<'a>( hop_started, )); // **429 要保住 429。**塌成 502 的话,客户端会当成「服务器 - // 坏了」而不是「该退避了」,而它们该做的事完全不同。 + // 坏了」而不是「该退避了」,而它们该做的事完全不同。到这儿的 429 是额度 + // 用完、没钱了(限流在上面那一支) last_err = Some(if status == 429 { GatewayError::rate_limited(msg!( "gw.upstream.rate_limited", upstream = provider.name.clone() => @@ -778,14 +954,30 @@ pub(super) async fn try_upstreams<'a>( } } Ok(r) => { - // 回答的开头:**第一段内容到之前**上游在流里报的错照样换下一家,一直没有内容 - // 也换(见 `opening`)。最后一家不压 —— 没有下一家可换,压着只会让客户端晚 - // 一点看到同样的东西;它不出声由回程按无响应超时收尾(见 `relay`) + // 回答的开头:**第一段内容到之前**上游在流里报的错照样换下一家(报的是限流就排队 + // 等它,和回 429 一样),一直没有内容也换(见 `opening`)。 + // + // 最后一家只在这一家可能在限流时压(探路的、最近被限流过的,见 + // `throttle::Throttles::wary`),而且只压流式的、这个请求又还等得起的:流的开头报的 + // 限流要暂停这一家、排队等它恢复,客户端察觉不到(不压的话那个错误直接交给了客户端)。 + // 别的时候不压:没有下一家可换,客户端该当场拿到响应头和开头的事件;整包的回答不在 + // 正文里报限流,等不起的也排不了队。不压的由回程按无响应超时收尾(见 `relay`) let hold = hold_of(req, reading, provider, &out, &r); - // 最后一家不压开头:它答没答上(成败的账)要等回程看到第一段内容才知道 - let unconfirmed = last && hold.is_some(); - let r = match hold.filter(|_| !last) { - None => r, + let held = hold.as_ref().is_some_and(|h| { + !last + || (matches!(h, Hold::Stream(..)) + && tokio::time::Instant::now() < until + && state.throttles.wary(&provider.name)) + }); + // 不压开头的最后一家:它答没答上(成败的账)要等回程看到第一段内容才知道 + let mut unconfirmed = last && hold.is_some() && !held; + let r = match hold.filter(|_| held) { + None => { + // 不压开头:响应头就是它的答复,这一家没在限流(探路的过了,连续限流的次数 + // 归零) + state.throttles.answered(&provider.name, sent_at); + r + } Some(hold) => { let status = r.status().as_u16(); let watching = async move { @@ -804,6 +996,9 @@ pub(super) async fn try_upstreams<'a>( }; match opened { None => { + if let Some(p) = probe.take() { + p.withdraw(); + } chain.push(super::idle::aborted_hop( &provider.name, model.clone(), @@ -813,30 +1008,43 @@ pub(super) async fn try_upstreams<'a>( halt = Some(GatewayError::aborted()); break; } - Some(super::opening::Opening::Go(r)) => r, - // 到点了还没有内容:放弃这一家。**响应跟着这一轮循环丢掉**,和上游的连接 - // 随之断开,它不再接着生成 + // 内容来了:**这时**才是这一家的答复。响应头是 200、流里接着报限流的 + // 上游,响应头不说明它好了 + Some(super::opening::Opening::Go(r)) => { + state.throttles.answered(&provider.name, sent_at); + r + } + // 到点了还没有内容:这一跳没有结论 Some(super::opening::Opening::Slow { response, usage }) => { - let status = response.status().as_u16(); - drop(response); - reached = true; - // 上游回了话,说明代理是通的 - state.note_proxy_ok(&provider.proxy); - last_err = Some(gave_up( - state, - started, - reading, - Abandoned { - provider, - model: model.clone(), - status: Some(status), - usage, - quiet, - hop_started, - }, - &mut chain, - )); - continue; + state.throttles.released(&provider.name, sent_at); + if last { + // 最后一家:和不压开头时一样交给回程,由它按无响应超时收尾 + unconfirmed = true; + response + } else { + // 放弃这一家。**响应跟着这一轮循环丢掉**,和上游的连接随之 + // 断开,它不再接着生成 + let status = response.status().as_u16(); + drop(response); + reached = true; + // 上游回了话,说明代理是通的 + state.note_proxy_ok(&provider.proxy); + last_err = Some(gave_up( + state, + started, + reading, + Abandoned { + provider, + model: model.clone(), + status: Some(status), + usage, + quiet, + hop_started, + }, + &mut chain, + )); + continue; + } } Some(super::opening::Opening::Failed { status, @@ -845,45 +1053,83 @@ pub(super) async fn try_upstreams<'a>( kind, message, response, - }) => match crate::failure::classify(status, &headers, &body, now_ms()) - { - Verdict::ClientError => response, - Verdict::Failed(cause) => { - reached = true; + }) => { + let verdict = + crate::failure::classify(status, &headers, &body, now_ms()); + let limited = + matches!(verdict, Verdict::Failed(Cause::RateLimited { .. })); + if !limited { + to_throttle(state, &provider.name, sent_at, verdict); + } + if verdict == Verdict::ClientError || (last && !limited) { + // 请求本身的问题,或者最后一家报了限流以外的错:原样交出去,和 + // 不压开头时一样(后者的成败由回程记) + unconfirmed |= verdict != Verdict::ClientError; + response + } else { // 换下一家:这一跳的回答就是上游在开头报的那个错 tap.said(&body); state.note_quota(id, &provider.name, &headers); state.note_proxy_ok(&provider.proxy); - let cause = known_reset(state, &provider.name, cause); - note_health( - &state.bus, - &state.health, - &provider.name, - state.health.record_cause(&provider.name, cause), - ); let said = msg!( "gw.upstream.stream_opening_error", upstream = provider.name.clone(), kind = kind, message = message => "Upstream `{upstream}` started the answer and reported an error \ before any content ({kind}): {message}" ); - let err = if status == 429 { - GatewayError::rate_limited(said) - } else { - GatewayError::upstream(said) - }; - chain.push(hop_failed( + let mut a = hop_failed( &provider.name, model.clone(), - err.detail.clone(), + said.clone(), hop_started, - )); - last_err = Some(err); + ); + match verdict { + // 流开头报的限流:和回 429 一样,暂停这一家,这个请求接着看 + // 下一家或者等它(见上面回 429 的那一支) + Verdict::Failed(Cause::RateLimited { .. }) => { + let hint = crate::throttle::pause_hint( + &headers, + &body, + now_ms(), + ); + let pause = state.throttles.limited( + &provider.name, + sent_at, + hint, + ); + usage_hold.release(); + a.pause_ms = Some(pause.as_millis() as u64); + chain.push(a); + if probe.is_some() { + waiter.keep_place(&provider.name); + } + waiting.push((name, Wait::Throttled)); + } + Verdict::Failed(cause) => { + reached = true; + let cause = known_reset(state, &provider.name, cause); + note_health( + &state.bus, + &state.health, + &provider.name, + state.health.record_cause(&provider.name, cause), + ); + chain.push(a); + last_err = Some(if status == 429 { + GatewayError::rate_limited(said) + } else { + GatewayError::upstream(said) + }); + } + // 请求本身的问题:上面原样交出去了 + Verdict::ClientError => {} + } continue; } - }, + } Some(super::opening::Opening::Broken(err)) => { reached = true; + state.throttles.released(&provider.name, sent_at); note_health( &state.bus, &state.health, @@ -940,18 +1186,35 @@ pub(super) async fn try_upstreams<'a>( state.health.record_failure(&provider.name), ); let err = match e { + // 连不上:上游一个字节都没收到,这一跳不算 + SendError::Http(e) if e.is_connect() => { + usage_hold.release(); + reached = true; + // 没有结论:探路的不再探,交给熔断(见 `crate::throttle`) + state.throttles.released(&provider.name, sent_at); + state.check_proxy(&provider.proxy); + forward::map_reqwest_error(e) + } SendError::Http(e) => { reached = true; + state.throttles.released(&provider.name, sent_at); // 连不上的可能是代理而不是上游 —— 检一次那个代理,说清是哪一件事 state.check_proxy(&provider.proxy); forward::map_reqwest_error(e) } - SendError::Sign(e) => GatewayError::config(msg!( - "gw.upstream.sign_failed", - upstream = provider.name.clone(), detail = e.0 => - "The request to upstream `{upstream}` could not be signed with its AWS \ - access keys: {detail}" - )), + // 签不了名:没发出去,探路的名额还回去 + SendError::Sign(e) => { + usage_hold.release(); + if let Some(p) = probe.take() { + p.withdraw(); + } + GatewayError::config(msg!( + "gw.upstream.sign_failed", + upstream = provider.name.clone(), detail = e.0 => + "The request to upstream `{upstream}` could not be signed with its AWS \ + access keys: {detail}" + )) + } }; chain.push(hop_failed( &provider.name, @@ -1036,23 +1299,33 @@ pub(super) async fn try_upstreams<'a>( )))); } - // 剩下的候选等过了还都满着,**一家都没发到过**:429,带 `Retry-After`。请求本身没问题, - // 过一会儿再来就发得出去;报成哪一家的失败都不对,它们一个字节都没收到。 + // 剩下的候选等到最后还都发不出去(满着、限流暂停着、用量上限满了),**一家都没发到过**: + // 429,带 `Retry-After`。请求本身没问题,过一会儿再来就发得出去;报成哪一家的失败都不对, + // 它们一个字节都没收到 —— 回了 429 的也算没收到:它让等一会儿。候选全都是这一期用完了的也 + // 一样,那时到重置之前别再试(见 `crate::throttle::refusal`)。 // - // 发到过的话,交出去的是最后那一家失败的原因,和没有满着的上游时一样:乙回了 401、甲一直 + // 发到过的话,交出去的是最后那一家失败的原因,和没有发不出去的上游时一样:乙回了 401、甲一直 // 满着,客户端该看到的是乙的错 —— 一个「都满着」的 429 会让它退避了再试,而该修的是乙的 // 密钥 - if stalled && !reached { - let upstreams = busy + if !reached && (stalled || (last_err.is_none() && served.is_none() && !spent.is_empty())) { + let now = now_ms(); + let mut stuck: Vec<(usize, String, crate::throttle::Stuck)> = waiting .iter() - .map(|n| format!("`{n}`")) - .collect::>() - .join(", "); - return Err(GatewayError::busy(msg!( - "gw.busy_all", upstreams = upstreams => - "Every upstream that can serve this request is at its concurrency limit \ - (max_concurrent): {upstreams}. None had a free slot in time; try again shortly." - ))); + .map(|(n, why)| (n, stuck_now(state, rt, n, *why, now))) + .chain( + spent + .iter() + .map(|(n, t)| (n, crate::throttle::Stuck::UsedUp(*t))), + ) + .map(|(n, s)| { + let at = started.alive.iter().position(|a| a == *n); + (at.unwrap_or(usize::MAX), n.to_string(), s) + }) + .collect(); + stuck.sort_by_key(|(at, _, _)| *at); + let stuck: Vec<(String, crate::throttle::Stuck)> = + stuck.into_iter().map(|(_, n, s)| (n, s)).collect(); + return Err(crate::throttle::refusal(&stuck, now)); } let Some(served) = served else { let mut err = last_err.unwrap_or_else(|| { @@ -1102,6 +1375,213 @@ fn stamp_queued(chain: &mut [tw_api::AttemptView], queued: Option<(usize, u64)>) } } +/// 发出去了的一跳在这一家用量上限上占着的(见 [`crate::upstream_limits`])。**丢掉就交给请求号**, +/// 这一跳算一个请求、等结算:不管它怎么收场,上游都收到了它。回了 429 的、其实没发出去的调 +/// [`Spent::release`] 还回去 +struct Spent { + hold: Option, + id: u64, +} + +impl Spent { + fn release(&mut self) { + if let Some(h) = self.hold.take() { + h.release(); + } + } +} + +impl Drop for Spent { + fn drop(&mut self) { + if let Some(h) = self.hold.take() { + h.bind(self.id); + } + } +} + +/// 一家暂时发不出去、等得到:跳过它时它是哪一样。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Wait { + /// 满着(`max_concurrent`) + Busy, + /// 限流暂停着 + Throttled, + /// 分钟、小时的用量上限满了 + WindowFull, +} + +impl Wait { + fn of(kind: tw_api::ThrottleKind) -> Self { + match kind { + tw_api::ThrottleKind::WindowFull => Wait::WindowFull, + _ => Wait::Throttled, + } + } +} + +/// 这一家的用量上限(见 [`crate::throttle::usage_limits`])。配置里没有它的是空的 +pub(super) fn usage_limits_of<'r>(rt: &'r Runtime, name: &str) -> &'r [tw_config::KeyLimit] { + rt.config + .providers + .iter() + .find(|p| p.name == name) + .map(crate::throttle::usage_limits) + .unwrap_or_default() +} + +/// 跳过了的、等过没等到的一跳:满着的、限流暂停着的、用量上限满了的。`until` 是那一刻看到的 +/// 预计恢复的时刻 +fn skipped_hop( + state: &AppState, + provider: &tw_config::Provider, + model: Option, + (why, until_ms): (Wait, Option), + queued_ms: Option, + started: std::time::Instant, +) -> tw_api::AttemptView { + let name = &provider.name; + match why { + Wait::Busy => crate::server::hop_busy( + name, + model, + state.slots.limit(name).unwrap_or_default(), + queued_ms, + started, + ), + Wait::Throttled => crate::server::hop_waiting( + name, + model, + tw_api::ThrottleKind::RateLimited, + until_ms, + queued_ms, + started, + ), + Wait::WindowFull => crate::server::hop_waiting( + name, + model, + tw_api::ThrottleKind::WindowFull, + until_ms, + queued_ms, + started, + ), + } +} + +/// 要等的这几家,**每一家都要到等待期限之后才恢复**:说得出什么时候、又晚于期限。满着的随时 +/// 可能空出来,说不出什么时候的(探路的还没回话)也等得到,都不算 +fn hopeless(state: &AppState, rt: &Runtime, waiting: &[(&String, Wait)], deadline_ms: u64) -> bool { + waiting.iter().all(|(n, why)| { + *why != Wait::Busy + && match state.throttles.peek(n, usage_limits_of(rt, n)) { + crate::throttle::Peek::Wait { + until_ms: Some(t), .. + } => t > deadline_ms, + crate::throttle::Peek::UsedUp { .. } => true, + _ => false, + } + }) +} + +/// 一家等到最后还发不出去:此刻它是哪一样、预计什么时候恢复(见 [`crate::throttle::refusal`]) +fn stuck_now( + state: &AppState, + rt: &Runtime, + name: &str, + why: Wait, + now: u64, +) -> crate::throttle::Stuck { + use crate::throttle::{Peek, Stuck}; + if why == Wait::Busy { + return Stuck::Busy; + } + match state.throttles.peek(name, usage_limits_of(rt, name)) { + Peek::UsedUp { until_ms, .. } => Stuck::UsedUp(until_ms), + Peek::Wait { + kind: tw_api::ThrottleKind::WindowFull, + until_ms, + } => Stuck::WindowFull(until_ms), + Peek::Wait { until_ms, .. } => Stuck::RateLimited(until_ms), + // 到点那一刻刚好恢复了:马上再来就行 + Peek::Open if why == Wait::WindowFull => Stuck::WindowFull(Some(now)), + Peek::Open => Stuck::RateLimited(Some(now)), + } +} + +/// 这个请求开始排队了(见 [`tw_api::Event::RequestQueued`]):在等哪几家,最多等到什么时候 +fn queued_up( + state: &AppState, + id: u64, + names: &[(&str, &[tw_config::KeyLimit])], + deadline_ms: u64, +) { + state.bus.emit(tw_api::Event::RequestQueued { + id, + upstreams: names.iter().map(|(n, _)| n.to_string()).collect(), + until_ms: deadline_ms, + at_ms: now_ms(), + }); +} + +/// 这个请求不排队了(见 [`tw_api::Event::RequestDequeued`]):等到了要发出去、被中止了、等不到了 +fn dequeued(state: &AppState, id: u64) { + state.bus.emit(tw_api::Event::RequestDequeued { + id, + at_ms: now_ms(), + }); +} + +/// 一跳有了不是限流的结论,告诉 [`crate::throttle`]:额度用完、没钱了是等不到 —— 不再探路,交给 +/// 熔断([`crate::throttle::Throttles::released`]);别的是这一家的答复,它没在限流 +fn to_throttle(state: &AppState, name: &str, sent_at: std::time::Instant, verdict: Verdict) { + match verdict { + Verdict::Failed(Cause::QuotaUsedUp { .. } | Cause::NoBalance) => { + state.throttles.released(name, sent_at) + } + _ => state.throttles.answered(name, sent_at), + } +} + +/// 这一跳在这一家的用量上限上占多少:输入 token 的估算(解不开的请求没有估算,占 0),和设了 +/// 费用上限时按发给它的名字算的输入费用。没有价格、不计费的是 0 —— 和结算时一样(见 +/// `super::admission`) +fn usage_ask( + state: &AppState, + reading: &crate::client_api::Reading, + provider: &tw_config::Provider, + model: &str, + limits: &[tw_config::KeyLimit], +) -> crate::key_limits::Ask { + let tokens = if matches!(reading.decoded, Some(Ok(_))) { + reading.facts.input_tokens + } else { + 0 + }; + let priced = limits + .iter() + .any(|l| l.measure() == tw_config::LimitMeasure::Cost); + let cost_micros = if priced && tokens > 0 { + let usage = tw_pricing::Usage { + input: tokens, + ..Default::default() + }; + crate::quote::quote( + &state.pricing.load(), + &provider.name, + model, + &usage, + provider.billing, + ) + .cost_micros + .unwrap_or(0) + } else { + 0 + }; + crate::key_limits::Ask { + tokens, + cost_micros, + } +} + /// 这一家和客户端是同一种格式(配置里没写格式的也算:照原样发过去) fn same_format(req: &Inbound, provider: &tw_config::Provider) -> bool { match (req.api, provider.effective_protocol()) { @@ -1137,6 +1617,7 @@ fn estimated_hop( usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, } } diff --git a/crates/tw-gateway/src/server/pipeline/idle.rs b/crates/tw-gateway/src/server/pipeline/idle.rs index 3f06425d..75a35386 100644 --- a/crates/tw-gateway/src/server/pipeline/idle.rs +++ b/crates/tw-gateway/src/server/pipeline/idle.rs @@ -98,6 +98,7 @@ pub(super) fn abandoned( // 等过空位的话,等了多久由尝试链补上(`stamp_queued`) queued_ms: None, skipped: None, + pause_ms: None, proxy: None, } } @@ -120,6 +121,7 @@ pub(super) fn aborted_hop( usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, } } diff --git a/crates/tw-gateway/src/server/pipeline/opening.rs b/crates/tw-gateway/src/server/pipeline/opening.rs index 02412173..f9ea1ee0 100644 --- a/crates/tw-gateway/src/server/pipeline/opening.rs +++ b/crates/tw-gateway/src/server/pipeline/opening.rs @@ -11,9 +11,11 @@ //! 和无响应超时是同一个判据(见 [`crate::pulse`])。 //! //! 等到无响应超时那一刻(`failover.idle_timeout_secs`,从请求发出去算起)还没有内容是 -//! [`Opening::Slow`]:调用方放弃这一家、换下一家。**最后一家不在这里等**(见 `hop`):没有 -//! 下一家可换,压着它只会让客户端晚一点看到同样的东西。开头压得太多([`HOLD_LIMIT`])也不 -//! 再压,交出去由回程照常计时。压着的这段时间里客户端等的是响应头:流式的回答过了 +//! [`Opening::Slow`]:调用方放弃这一家、换下一家。**最后一家只在它可能在限流时在这里等**(流式 +//! 的、请求还等得起,见 `hop`):开头报的限流要暂停这一家、排队等它恢复;别的错、到点了还没有 +//! 内容,没有下一家可换,照原样交给回程。别的时候最后一家不等:压着它只会让客户端晚一点看到 +//! 同样的东西。 +//! 开头压得太多([`HOLD_LIMIT`])也不再压,交出去由回程照常计时。压着的这段时间里客户端等的是响应头:流式的回答过了 //! [`crate::OPENING_HOLD`] 先拿到 `200` 和保活注释,上游的事件照旧压在这里(见 //! [`super::commit`])。 //! diff --git a/crates/tw-gateway/src/server/pipeline/relay.rs b/crates/tw-gateway/src/server/pipeline/relay.rs index 6b6eec36..926b2e42 100644 --- a/crates/tw-gateway/src/server/pipeline/relay.rs +++ b/crates/tw-gateway/src/server/pipeline/relay.rs @@ -183,6 +183,9 @@ pub(super) fn respond( let mut pulse = crate::pulse::Pulse::new(plan.is_sse.then_some(upstream_dialect)); let abort = req.abort.clone(); let (bus, health) = (state.bus.clone(), state.health.clone()); + // 上游在流里报了限流(答到一半的也算):这一家暂停(见 `crate::throttle`)。最近一次 429 之后 + // 发出去的才算连续的又一次,按这一跳发出去的那一刻比 + let throttles = state.throttles.clone(); // 这个请求结束时记一笔,到了时候重读这一家的余额(见 `crate::balances`) let passing = state.balance_passing(&upstream_name); let stream = async_stream::stream! { @@ -327,6 +330,21 @@ pub(super) fn respond( if let Some(err) = denied { broke = Some(err.in_dialect(dialect)); } + // 上游在流里报的错是限流:回答照样以错误收尾(客户端已经收到了东西,换不了),这一家暂停 + if let Some((code, body)) = ending + .upstream_fault() + .and_then(|data| super::opening::stream_error(upstream_dialect, data)) + && let crate::failure::Verdict::Failed(crate::failure::Cause::RateLimited { + retry_after, + }) = crate::failure::classify( + code, + &http::HeaderMap::new(), + &body, + crate::server::now_ms(), + ) + { + throttles.limited(&upstream_name, sent_at, retry_after); + } // 响应体留档和结束事件都在 `ending` 里:三种结局要交出去的是同一份 // 东西,分开写就会有一种漏掉 // 一段真内容都没有就好好收了尾:也是答上了 diff --git a/crates/tw-gateway/src/server/upgrade.rs b/crates/tw-gateway/src/server/upgrade.rs index 50f57e77..ffca81a7 100644 --- a/crates/tw-gateway/src/server/upgrade.rs +++ b/crates/tw-gateway/src/server/upgrade.rs @@ -286,7 +286,7 @@ pub(super) async fn ws_upgrade( &client_name, limits, Default::default(), - crate::key_limits::slot_wait(&rt.config), + crate::key_limits::request_wait(&rt.config), ) .await { diff --git a/crates/tw-gateway/src/slots.rs b/crates/tw-gateway/src/slots.rs index 11a164d4..bba179c1 100644 --- a/crates/tw-gateway/src/slots.rs +++ b/crates/tw-gateway/src/slots.rs @@ -9,8 +9,8 @@ //! - 别的候选满着:当场跳过,试下一家; //! - 候选都满着:在它们里面等先空出来的那一家,等不到就回 429。 //! -//! 等多久是 `failover.slot_wait_secs`,**一个请求合起来算**,准入时等密钥的分钟、小时上限 -//! 用掉的也算在里面(见 `crate::key_limits::slot_wait`):等的时候客户端一个字节都收不到。 +//! 等多久是 `failover.wait_secs`,**一个请求合起来算**,准入时等密钥的分钟、小时上限 +//! 用掉的也算在里面(见 `crate::key_limits::request_wait`):等的时候客户端一个字节都收不到。 //! **等不是失败**:满着的上游不停用、不进熔断的账。 //! //! 一个位置从发出请求占到回答交完、或者客户端走掉,由 [`Slot`] 的 Drop 还回去。和每把 diff --git a/crates/tw-gateway/src/state.rs b/crates/tw-gateway/src/state.rs index 5a56b067..1a4704f5 100644 --- a/crates/tw-gateway/src/state.rs +++ b/crates/tw-gateway/src/state.rs @@ -158,6 +158,13 @@ pub struct AppState { /// **跨重载存活**,理由和并发闸门一样:改一条路由规则不该让今天花的钱归零,在跑的 /// 请求的预留也不该变成孤儿。上限本身跟着配置换([`crate::key_limits::KeyLimits::configure`]) pub key_limits: Arc, + /// 每家上游的用量上限(见 [`crate::upstream_limits`]):这家今天、这周、这个月用了多少, + /// 最近几分钟、几小时用了多少,在路上的那几跳占着多少。**跨重载存活**,理由和密钥的 + /// 一样;上限本身跟着配置换([`crate::upstream_limits::UpstreamLimits::configure`]) + pub upstream_limits: Arc, + /// 每家上游暂时发不出去的样子:限流暂停、用量上限,和等它们的请求排的队(见 + /// [`crate::throttle`])。**跨重载存活**:改一条规则不该让一家正被限流的上游马上又挨一遍 + pub throttles: Arc, /// 一家上游不在当前运行时里时顶上的 Client(直连,不读系统代理) pub http: reqwest::Client, /// 观测事件往这里丢。没有订阅者时是零成本的 —— 数据面不该知道有 @@ -303,12 +310,21 @@ impl AppState { let bus = tw_observe::EventBus::new(); let key_limits = Arc::new(crate::key_limits::KeyLimits::new(bus.clone())); key_limits.configure(&rt.config); + let upstream_limits = Arc::new(crate::upstream_limits::UpstreamLimits::new(bus.clone())); + upstream_limits.configure(&rt.config); + let throttles = Arc::new(crate::throttle::Throttles::new( + bus.clone(), + upstream_limits.clone(), + )); + throttles.configure(&rt.config); let state = Self { local_time: Arc::new(tw_engine::LocalTime::now), rt: Arc::new(arc_swap::ArcSwap::from_pointee(rt)), gate: Default::default(), slots, key_limits, + upstream_limits, + throttles, http, bus, health, @@ -395,12 +411,22 @@ impl AppState { self.local_time = clock; } - /// 换一个看用量上限的时钟(测试把时间拨到零点前后)。**账从空的开始**:只在测试里、 - /// 第一个请求之前调 + /// 换一个看用量上限的时钟(测试把时间拨到零点前后):密钥的和上游的用同一个。**账从 + /// 空的开始**,上游的暂停和排队也是:只在测试里、第一个请求之前调 pub fn set_key_limits_clock(&mut self, clock: Arc) { - let limits = crate::key_limits::KeyLimits::with_clock(self.bus.clone(), clock); + let limits = crate::key_limits::KeyLimits::with_clock(self.bus.clone(), clock.clone()); limits.configure(&self.config()); self.key_limits = Arc::new(limits); + let upstream = crate::upstream_limits::UpstreamLimits::with_clock(self.bus.clone(), clock); + upstream.configure(&self.config()); + self.upstream_limits = Arc::new(upstream); + // 暂停和排队看的是同一份用量账:跟着换 + let throttles = Arc::new(crate::throttle::Throttles::new( + self.bus.clone(), + self.upstream_limits.clone(), + )); + throttles.configure(&self.config()); + self.throttles = throttles; } /// 直接换一份插件进去,配置照旧:正在跑的请求用完它们手上那一份,新请求看到的是 @@ -463,11 +489,21 @@ impl AppState { } else { Default::default() }, - // 并发数满着的也不参加这一轮:它们会被当场跳过(见 `crate::slots`) + // 并发数满着的也不参加这一轮:它们会被当场跳过(见 `crate::slots`)。限流暂停着的、 + // 用量上限满了或用完了的一样(见 `crate::throttle`)—— 跳过的不记账,真发出去的那一家记 busy: if balanced { candidates .iter() - .filter(|p| self.slots.is_full(p)) + .filter(|p| { + self.slots.is_full(p) || { + let limits = providers + .iter() + .find(|x| &x.name == *p) + .map(crate::throttle::usage_limits) + .unwrap_or_default(); + self.throttles.peek(p, limits) != crate::throttle::Peek::Open + } + }) .cloned() .collect() } else { @@ -563,6 +599,8 @@ impl AppState { self.health.configure(&next.config.failover); self.slots.configure(&next.config.providers); self.key_limits.configure(&next.config); + self.upstream_limits.configure(&next.config); + self.throttles.configure(&next.config); self.rt.store(Arc::new(next)); self.announce_broken(broken); // 模型汇总马上按新配置重算:删掉、停用的上游的模型必须立刻消失(列表 diff --git a/crates/tw-gateway/src/throttle.rs b/crates/tw-gateway/src/throttle.rs new file mode 100644 index 00000000..ad820e9f --- /dev/null +++ b/crates/tw-gateway/src/throttle.rs @@ -0,0 +1,1967 @@ +//! 上游暂时发不出去:回了 429 之后的暂停,用量上限(见 [`crate::upstream_limits`]),和等它们的 +//! 请求排的队。 +//! +//! **这些都不是失败。**上游没坏,只是要等一会儿:不进熔断的账,不算进成功率,`load-balance` +//! 也不因此少分给它(见 [`crate::health`])。请求先换下一家;没有下一家了,在它们里面等先能发的 +//! 那一家,等到这个请求的等待期限(`failover.wait_secs`)为止(见 `server::pipeline::hop`)。 +//! +//! 四种样子,前三种等得到,最后一种不等: +//! +//! - **满着**(`max_concurrent`):哪个请求结束了就空出来,事先不知道是什么时候(见 +//! [`crate::slots`]); +//! - **限流暂停**:上游回了 429(不是额度用完)。暂停多久按上游说的 —— `retry-after-ms`、 +//! `Retry-After`、正文里的等待时间([`pause_hint`])—— 没说就按这一家连续被限流的次数退避 +//! ([`backoff`]),最多 `failover.rate_limit_max_pause_secs`。说了等于没说的(`Retry-After: 0`) +//! 当没说; +//! - **滚动窗口满了**:分钟、小时的用量上限,最近的用量滑出窗口就又能发,什么时候是算得出的; +//! - **这一期用完了**:天、周、月的用量上限,到下一期之前等多久都一样 —— 跳过,不等。 +//! +//! **每家一个队,先来先发**:留在这一家的对话和别的请求排在同一个队里,不插队。一个请求在几家 +//! 的队里同时排着,哪一家先放它就去哪一家,别的队它当场离开([`Waiter`])。暂停结束时**只放 +//! 队头一个过去探路**([`Probe`]):它拿到不是 429 的答复(压着开头的回答要等到第一段内容,见 +//! `hop`),或者发出去 [`PROBE_GRACE`] 还没有结论,这一家就算好了,后面的照并发和用量上限放;它 +//! 又被限流,这一家接着暂停(退避接着往上走),它下一次等的时候回到队头。它没有结论 —— 连不上、 +//! 没有内容超时了、回的是额度用完 —— 不再探路,排着的照常放,这一家怎么样交给熔断 +//! ([`Throttles::released`])。滚动窗口空出来不用探路:用量是本地数的,空出多少放多少。 +//! +//! 客户端走了、请求被手动中止了:当场离开队,什么都不占。 +//! +//! **只在内存里**:跨配置重载存活,重启就没了 —— 重启之后第一个请求照常发,上游还在限流的话 +//! 再回一个 429,这里重新暂停。暂停多久按单调时钟和墙上的时钟里走得多的那个算:系统时间被往回拨, +//! 暂停不跟着变长;合上盖子睡过去的时间算进暂停(单调时钟那时不走)。 +//! +//! 状态变了(开始暂停、暂停结束、用量上限满了或空出来)报 [`tw_api::Event::ThrottleChanged`], +//! 界面显示它;它自己会好,不发通知。 + +use std::collections::{HashMap, VecDeque}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::time::{Duration, Instant}; + +use http::HeaderMap; +use tokio::sync::watch; +use tw_api::ThrottleKind; +use tw_config::KeyLimit; + +use crate::upstream_limits::LimitState; + +/// 暂停结束后放过去探路的那一个,发出去这么久还没有结论就当这一家好了。上游要是还在限流, +/// 多半当场就回 429;一个长回答的开头到得慢,不该让后面的一直等。 +/// +/// **这不是知道的恢复时刻**:探路的多半早就有了结论。等不等得到(`hop` 的 `hopeless`、会话 +/// 粘性)把它当成说不出;只有交给客户端的 `Retry-After` 拿它估一个数([`refusal`]) +pub const PROBE_GRACE: Duration = Duration::from_secs(10); + +/// 上游说的暂停短于这个的当没说([`Throttles::limited`]):照它暂停等于不暂停,探路的当场又发、 +/// 又被限流,一直空转到这个请求的等待期限。按秒说的(`Retry-After`、正文里的)不到一秒就当没说 +/// (见 [`crate::failure::MIN_RETRY_AFTER`]) +const MIN_HINT: Duration = Duration::from_millis(1); + +/// 上游没说暂停多久时,按这一家连续被限流的次数退避:第一次 5 秒,第二次 10 秒……第五次起 +/// 每次 60 秒 +const BACKOFF_SECS: [u64; 5] = [5, 10, 20, 40, 60]; + +/// 一家最近一次被限流之后这么久之内,它还「可能在限流」([`Throttles::wary`]) +const WARY_FOR: Duration = Duration::from_secs(300); + +/// 两个时钟走的差,到了这么多才当是系统时间被拨过(见 [`Span::end_ms`]);两个时钟换算出来的 +/// 时刻差不到这么多的当是同一刻(见 [`Throttles::limit_state`])。先后读两个时钟,差会抖一两毫秒 +const SLACK_MS: u64 = 1_000; + +/// 等一个时刻到来时,最多睡这么久就起来再看一眼。**tokio 的定时器按单调时钟走**,合上盖子睡着的 +/// 时候它不走:醒来时暂停按墙上的时钟早就结束了,定时器却还要睡完剩下的 +const NAP: Duration = Duration::from_secs(1); + +/// 一家上游的用量上限写在哪儿(`providers[].limits`,见 [`crate::upstream_limits`]) +pub fn usage_limits(p: &tw_config::Provider) -> &[KeyLimit] { + &p.limits +} + +/// 这一家连续第 `strikes` 次被限流、上游又没说暂停多久时,暂停多久 +pub fn backoff(strikes: u32) -> Duration { + let i = (strikes.max(1) - 1).min(BACKOFF_SECS.len() as u32 - 1) as usize; + Duration::from_secs(BACKOFF_SECS[i]) +} + +/// 上游在 429 里说了暂停多久没有。**按这个顺序取第一个**:`retry-after-ms`(毫秒,Anthropic、 +/// OpenAI 的 SDK 认它)、`Retry-After`(秒数或 HTTP 日期)、正文里的等待时间(Gemini 的 +/// `retryDelay`)。都没有是 None,按退避算。说了等于没说的(不到一毫秒的 `retry-after-ms`, +/// 不到一秒的另外两样)当没说,往下看 +pub fn pause_hint(headers: &HeaderMap, body: &[u8], now_ms: u64) -> Option { + retry_after_ms(headers) + .or_else(|| crate::failure::retry_after(headers, now_ms)) + .or_else(|| crate::failure::retry_delay_in_body(body)) +} + +fn retry_after_ms(headers: &HeaderMap) -> Option { + let v: f64 = headers + .get("retry-after-ms")? + .to_str() + .ok()? + .trim() + .parse() + .ok()?; + // 大得离谱的(`1e300`)照最长的算、交给最长暂停封顶:浮点换整数是饱和的,**不能 panic** + (v.is_finite() && v >= MIN_HINT.as_millis() as f64).then(|| Duration::from_millis(v as u64)) +} + +/// 一家上游此刻能不能发、对这一个请求来说([`Throttles::gate`])。 +pub enum Gate { + /// 能发。暂停刚结束时,这一个是探路的(带着 [`Probe`]) + Go(Option), + /// 等得到:限流暂停着(`rate_limited`)或者滚动窗口满了(`window_full`)。`until_ms` 是 + /// 预计又能发的那一刻;前面还有人排着、探路的还没有结论时说不出,是 None + Wait { + kind: ThrottleKind, + until_ms: Option, + }, + /// 这一期用完了:跳过,不等。`limit` 是哪一条上限 + UsedUp { until_ms: u64, limit: u32 }, +} + +/// 只看不占的 [`Gate`]:不拿探路的名额 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Peek { + Open, + Wait { + kind: ThrottleKind, + until_ms: Option, + }, + UsedUp { + until_ms: u64, + limit: u32, + }, +} + +/// 墙上的时钟,Unix 毫秒 +type Wall = Arc u64 + Send + Sync>; + +/// 单调时钟 +type Mono = Arc Instant + Send + Sync>; + +/// 同一刻,在两个时钟上 +#[derive(Debug, Clone, Copy)] +struct Moment { + mono: Instant, + wall: u64, +} + +impl Moment { + /// 从这一刻到 `now` 过了多久:**两个时钟里走得多的那个**。合上盖子睡着的时候单调时钟不走 + /// (macOS、Linux 都是),墙上的照走 —— 睡着的时间要算进去;系统时间被往回拨,墙上的倒退、 + /// 单调的照走 —— 拨回去也拉不长 + fn since(&self, now: Moment) -> Duration { + now.mono + .saturating_duration_since(self.mono) + .max(Duration::from_millis(now.wall.saturating_sub(self.wall))) + } +} + +/// 从某一刻起的一段时间:限流暂停、探路的宽限。**两个时钟哪个先走完就结束**(见 +/// [`Moment::since`]) +#[derive(Debug, Clone, Copy)] +struct Span { + from: Moment, + len: Duration, +} + +impl Span { + /// 还剩多久。结束了是 0 + fn left(&self, now: Moment) -> Duration { + self.len.saturating_sub(self.from.since(now)) + } + + fn over(&self, now: Moment) -> bool { + self.left(now).is_zero() + } + + /// 结束的那一刻,墙上的(报出去的 `until_ms`)。系统时间被往回拨过 —— 墙上的比单调的少走了 + /// [`SLACK_MS`] 以上 —— 起点在墙上跟着拨回去:报出去的时刻按拨过的算,还要等多久不变,之后 + /// 再报还是这一刻。**不按两个时钟的差现算**:先后读两个时钟,差会抖,报出去的时刻跟着抖的话, + /// 同一个暂停会报好几次([`Throttles::refresh`] 比的是报过的样子) + fn end_ms(&mut self, now: Moment) -> u64 { + let mono = now + .mono + .saturating_duration_since(self.from.mono) + .as_millis() as i64; + let wall = now.wall as i64 - self.from.wall as i64; + if wall.saturating_add(SLACK_MS as i64) < mono { + self.from.wall = now.wall.saturating_sub(mono as u64); + } + self.from.wall.saturating_add(self.len.as_millis() as u64) + } +} + +/// 每家上游的暂停、排队和探路。**跨重载存活**,按上游的名字记。 +/// +/// 两个时钟(见 [`Moment`]):暂停、宽限、「最近被限流过」**按两个里走得多的那个算**,合上盖子 +/// 睡过去的时间算进去,系统时间往回拨也拉不长。报出去的时刻(`until_ms`、事件的 `at_ms`)、请求 +/// 的等待期限、上游说的 HTTP 日期是墙上的(Unix 毫秒);一跳是什么时候发出去的看单调的 +pub struct Throttles { + ups: Mutex>, + limits: Arc, + bus: tw_observe::EventBus, + settings: Mutex, + wall: Wall, + mono: Mono, + /// 排队的、探路的各自的号 + next: AtomicU64, + /// 看着用量上限的那些后台任务([`Self::configure`]),一家一个。**配置换了就全换掉**: + /// 删掉、改了名的那一家,它的任务不会再被叫醒 + watchers: Mutex>, +} + +#[derive(Default)] +struct Settings { + max_pause: Duration, + /// 每家的用量上限:报状态、等窗口空出来时要看 + limits: HashMap>, +} + +/// 一家上游记着的。 +struct Up { + /// 暂停着的话,从哪一刻起、多久 + pause: Option, + /// 最近一次收到 429 的时刻。**在它之前发出去的请求的答复不算数**:一阵同时被限流的请求 + /// 只算一次,暂停之前发出去的一个长回答到了响应头,也不说明这家好了 + limited_at: Option, + /// 连续被限流的次数(退避按它算)。拿到别的答复就归零 + strikes: u32, + /// 暂停结束了、还没有结论:只放一个过去探路 + probing: bool, + /// 探路的那一个:它的号,和发出去的时刻(还没发出去是 None) + probe: Option<(u64, Option)>, + /// 排着队的,先来的在前 + queue: VecDeque, + /// 队是为什么排起来的:后来的、不在队里的请求看到的就是这个 + queue_kind: ThrottleKind, + /// 有变化就加一:排队的请求盯着它 + tx: watch::Sender, + /// 最近一次报出去的样子([`tw_api::Event::ThrottleChanged`]) + shown: Option, +} + +impl Up { + fn new() -> Self { + Self { + pause: None, + limited_at: None, + strikes: 0, + probing: false, + probe: None, + queue: VecDeque::new(), + queue_kind: ThrottleKind::RateLimited, + tx: watch::channel(0).0, + shown: None, + } + } + + fn bump(&self) { + self.tx.send_modify(|v| *v = v.wrapping_add(1)); + } + + /// 此刻还在暂停着的话,那一段 + fn paused(&mut self, now: Moment) -> Option<&mut Span> { + self.pause.as_mut().filter(|p| !p.over(now)) + } + + /// 探路的发出去 [`PROBE_GRACE`] 还没有结论:当这一家好了 + fn settle_probe(&mut self, now: Moment) { + if self.probing + && let Some((_, Some(sent))) = self.probe + && grace(sent).over(now) + { + self.probing = false; + self.probe = None; + self.bump(); + } + } + + /// 这个号排着的话离开队 + fn leave(&mut self, id: u64) { + if let Some(at) = self.queue.iter().position(|w| *w == id) { + self.queue.remove(at); + // 队头走了:下一个该看看轮没轮到它 + if at == 0 { + self.bump(); + } + } + } +} + +/// `sent` 发出去的探路的那一个的宽限 +fn grace(sent: Moment) -> Span { + Span { + from: sent, + len: PROBE_GRACE, + } +} + +/// 拿着这张探路名额的那一家。**认号不认名字**:拿到名额之后这一家可能改了名([`Throttles::rename`]), +/// 名额跟着走,[`Probe`] 手上的还是旧名字 +fn holding(map: &mut HashMap, ticket: u64) -> Option<(&String, &mut Up)> { + map.iter_mut() + .find(|(_, up)| matches!(up.probe, Some((t, _)) if t == ticket)) +} + +impl Throttles { + pub fn new( + bus: tw_observe::EventBus, + limits: Arc, + ) -> Self { + Self::with_clocks( + bus, + limits, + Arc::new(crate::server::now_ms), + Arc::new(Instant::now), + ) + } + + fn with_clocks( + bus: tw_observe::EventBus, + limits: Arc, + wall: Wall, + mono: Mono, + ) -> Self { + Self { + ups: Mutex::new(HashMap::new()), + limits, + bus, + settings: Mutex::new(Settings { + max_pause: Duration::from_secs( + tw_config::Failover::default().rate_limit_max_pause_secs, + ), + ..Default::default() + }), + wall, + mono, + next: AtomicU64::new(1), + watchers: Mutex::new(Vec::new()), + } + } + + fn lock(&self) -> MutexGuard<'_, HashMap> { + // 锁中毒了照样拿里面的表:少暂停一次,好过一个不转发的网关 + self.ups.lock().unwrap_or_else(|p| p.into_inner()) + } + + fn settings(&self) -> MutexGuard<'_, Settings> { + self.settings.lock().unwrap_or_else(|p| p.into_inner()) + } + + /// 此刻,墙上的时钟(Unix 毫秒) + fn now(&self) -> u64 { + (self.wall)() + } + + /// 此刻,在两个时钟上 + fn moment(&self) -> Moment { + Moment { + mono: (self.mono)(), + wall: self.now(), + } + } + + /// 照这份配置换设置:最长暂停多久、每家的用量上限。**暂停着的照旧暂停着**。有用量上限的 + /// 那几家各起一个后台任务看着,状态变了就报(见 [`Self::refresh`]);上一份配置的那些当场 + /// 收工 —— 删掉、改了名的那一家的用量不会再变,它的任务不收工就一直挂着,重载一次多一批 + pub fn configure(self: &Arc, cfg: &tw_config::Config) { + let names: Vec = { + let mut s = self.settings(); + s.max_pause = Duration::from_secs(cfg.failover.rate_limit_max_pause_secs); + s.limits = cfg + .providers + .iter() + .filter(|p| !usage_limits(p).is_empty()) + .map(|p| (p.name.clone(), usage_limits(p).to_vec())) + .collect(); + s.limits.keys().cloned().collect() + }; + // 配置换了,用量上限可能删了:报一遍现在的样子 + let known: Vec = self.lock().keys().cloned().collect(); + for n in known.iter().chain(&names) { + self.refresh(n); + } + let mut watchers = self.watchers.lock().unwrap_or_else(|p| p.into_inner()); + for w in watchers.drain(..) { + w.abort(); + } + let Ok(rt) = tokio::runtime::Handle::try_current() else { + return; + }; + for name in names { + let me = Arc::downgrade(self); + let task = rt.spawn(async move { + loop { + // **等的时候不攥着它**:它没了(网关关了)任务跟着收工,见 `Drop` + let (limits, nap) = { + let Some(t) = me.upgrade() else { return }; + let nap = t.view(&name).map(|v| t.nap(v.until_ms)); + (t.limits.clone(), nap) + }; + let wake = async { + match nap { + Some(d) => tokio::time::sleep(d).await, + None => std::future::pending().await, + } + }; + tokio::select! { + _ = limits.changed(&name) => {} + _ = wake => {} + } + let Some(t) = me.upgrade() else { return }; + t.refresh(&name); + t.bump(&name); + } + }); + watchers.push(task.abort_handle()); + } + } + + /// 等墙上的 `at_ms` 到来,这一次睡多久:离它还有多久,最多 [`NAP`]。已经过了的是 0 + fn nap(&self, at_ms: u64) -> Duration { + Duration::from_millis(at_ms.saturating_sub(self.now())).min(NAP) + } + + /// 这一家的用量上限此刻的样子([`crate::upstream_limits::UpstreamLimits::state`]),时刻换成 + /// 这里的时钟:用量上限按它自己的时钟算(测试会拨它),暂停、等待期限、`Retry-After` 都按 + /// 这里的。两个时钟各读一次;**差不到 [`SLACK_MS`] 的当是同一个**,时刻原样交回 —— 不然先后 + /// 读的两个时钟差一毫秒,报出去的时刻就差一毫秒,同一个样子报好几次 + fn limit_state(&self, name: &str, limits: &[KeyLimit]) -> LimitState { + let theirs = self.limits.now_ms(); + let mine = self.now() as i64; + let same = mine.abs_diff(theirs) <= SLACK_MS; + let convert = |until_ms: u64| { + if same { + until_ms + } else { + let left = (until_ms as i64).saturating_sub(theirs).max(0); + mine.saturating_add(left).max(0) as u64 + } + }; + match self.limits.state(name, limits) { + LimitState::Free => LimitState::Free, + LimitState::WindowFull { until_ms, limit } => LimitState::WindowFull { + until_ms: convert(until_ms), + limit, + }, + LimitState::PeriodUsedUp { until_ms, limit } => LimitState::PeriodUsedUp { + until_ms: convert(until_ms), + limit, + }, + } + } + + /// 这一家的用量上限([`Self::configure`] 记下的) + fn limits_of(&self, name: &str) -> Vec { + self.settings() + .limits + .get(name) + .cloned() + .unwrap_or_default() + } + + fn bump(&self, name: &str) { + if let Some(up) = self.lock().get(name) { + up.bump(); + } + } + + /// 这一家有变化时会动的那个。排队的请求盯着它(见 [`first_ready`]) + pub fn watch(&self, name: &str) -> watch::Receiver { + self.lock() + .entry(name.to_string()) + .or_insert_with(Up::new) + .tx + .subscribe() + } + + /// 这一家此刻对这一个请求(`waiter`,在排队的是它的号)能不能发。**能发、又该探路时拿走 + /// 探路的名额**([`Gate::Go`] 带着 [`Probe`])。`limits` 是这一家的用量上限 + pub fn gate(self: &Arc, name: &str, limits: &[KeyLimit], waiter: Option) -> Gate { + match self.check(name, limits, waiter, true) { + (Peek::Open, ticket) => Gate::Go(ticket.map(|ticket| Probe { + throttles: self.clone(), + ticket, + sent: false, + })), + (Peek::Wait { kind, until_ms }, _) => Gate::Wait { kind, until_ms }, + (Peek::UsedUp { until_ms, limit }, _) => Gate::UsedUp { until_ms, limit }, + } + } + + /// 只看不占:一个不在任何队里的请求此刻看到的样子。数 token 的请求、排序、会话粘性看它 + pub fn peek(&self, name: &str, limits: &[KeyLimit]) -> Peek { + self.check(name, limits, None, false).0 + } + + fn check( + &self, + name: &str, + limits: &[KeyLimit], + waiter: Option, + take: bool, + ) -> (Peek, Option) { + let window = match self.limit_state(name, limits) { + LimitState::PeriodUsedUp { until_ms, limit } => { + return (Peek::UsedUp { until_ms, limit }, None); + } + LimitState::WindowFull { until_ms, limit: _ } => Some(until_ms), + LimitState::Free => None, + }; + let now = self.moment(); + let mut map = self.lock(); + let Some(up) = map.get_mut(name) else { + return match window { + Some(t) => ( + Peek::Wait { + kind: ThrottleKind::WindowFull, + until_ms: Some(t), + }, + None, + ), + None => (Peek::Open, None), + }; + }; + up.settle_probe(now); + if let Some(pause) = up.paused(now) { + let until = pause.end_ms(now).max(window.unwrap_or(0)); + return ( + Peek::Wait { + kind: ThrottleKind::RateLimited, + until_ms: Some(until), + }, + None, + ); + } + if let Some(t) = window { + return ( + Peek::Wait { + kind: ThrottleKind::WindowFull, + until_ms: Some(t), + }, + None, + ); + } + // 先来先发:前面有人排着,就轮不到它 + if let Some(head) = up.queue.front() + && Some(*head) != waiter + { + return ( + Peek::Wait { + kind: up.queue_kind, + until_ms: None, + }, + None, + ); + } + if up.probing { + return match up.probe { + // 探路的已经在路上了(或者拿了名额、正要发):等它的结论。**什么时候有结论说不出** + // —— 宽限期的结束只是最晚的那一刻(见 [`PROBE_GRACE`]) + Some(_) => ( + Peek::Wait { + kind: ThrottleKind::RateLimited, + until_ms: None, + }, + None, + ), + None if take => { + let ticket = self.next.fetch_add(1, Ordering::Relaxed); + up.probe = Some((ticket, None)); + (Peek::Open, Some(ticket)) + } + None => ( + Peek::Wait { + kind: ThrottleKind::RateLimited, + until_ms: Some(now.wall), + }, + None, + ), + }; + } + (Peek::Open, None) + } + + /// 这一家此刻**可能在限流**:在探路、连续被限流还没有别的答复,或者最近 + /// ([`WARY_FOR`])被限流过。最后一家的流式回答这时也压着开头(见 `hop`):流的开头报的 + /// 限流要排队等它恢复,不交给客户端。一直好好的上游不压 —— 客户端照常当场拿到响应头和 + /// 开头的那几个事件 + pub fn wary(&self, name: &str) -> bool { + let now = self.moment(); + self.lock().get(name).is_some_and(|up| { + up.probing || up.strikes > 0 || up.limited_at.is_some_and(|t| t.since(now) < WARY_FOR) + }) + } + + /// 留在这一家的对话还留不留(见 [`crate::affinity`]):这一期用完了的、限流暂停要到 + /// `deadline_ms`(这个请求的等待期限)之后才结束的,放开 —— 和熔断着的一样,留在一家发不出去 + /// 的上游上,省下的缓存一分钱都兑现不了。等得到的照旧留,到时候等它 + pub fn keeps_conversation(&self, name: &str, limits: &[KeyLimit], deadline_ms: u64) -> bool { + match self.peek(name, limits) { + Peek::UsedUp { .. } => false, + Peek::Wait { + kind: ThrottleKind::RateLimited, + until_ms: Some(t), + } => t <= deadline_ms, + _ => true, + } + } + + /// 发给这一家的、`sent` 发出去的一跳被限流了(回了 429,或者在流的开头报了限流;额度用完 + /// 不算):暂停这一家。`hint` 是上游说的暂停多久([`pause_hint`]),没说、说了等于没说 + /// (不到 [`MIN_HINT`])按退避([`backoff`])。**已经暂停着的延到两个里晚的那个**。交回从此刻起 + /// 还要暂停多久 + pub fn limited( + self: &Arc, + name: &str, + sent: Instant, + hint: Option, + ) -> Duration { + let now = self.moment(); + let max = self.settings().max_pause; + let hint = hint.filter(|d| *d >= MIN_HINT); + let pause = { + let mut map = self.lock(); + let up = map.entry(name.to_string()).or_insert_with(Up::new); + // 上一次 429 之后才发出去的又被限流了,才是连续的又一次。同一阵里被限流的那几个只算 + // 一次 + if up.limited_at.is_none_or(|t| sent > t.mono) { + up.strikes = up.strikes.saturating_add(1); + } + let len = hint.unwrap_or_else(|| backoff(up.strikes)).min(max); + let fresh = Span { from: now, len }; + let pause = match up.pause { + Some(p) if p.left(now) >= len => p, + _ => fresh, + }; + up.pause = Some(pause); + up.limited_at = Some(now); + // 探路的被拒了,或者别人在探路时被拒了:暂停结束后重新探 + up.probing = true; + up.probe = None; + up.bump(); + pause + }; + self.refresh(name); + // 暂停结束那一刻也是一次变化,没有哪个请求会触发它 + self.at(name, pause); + pause.left(now) + } + + /// 发给这一家的、`sent` 发出去的一跳有了不是限流的答复:连续限流的次数归零;在探路的话, + /// 这一家好了,排着的照并发和用量上限放。**最近一次 429 之前发出去的不算**。 + /// + /// 答复是什么时候算的看调用方:压着开头的回答(见 `hop`、`ws::turn`)要等到第一段内容 —— + /// 响应头是 200、流里接着报限流的上游,响应头不说明它好了 + pub fn answered(&self, name: &str, sent: Instant) { + let changed = { + let mut map = self.lock(); + let Some(up) = map.get_mut(name) else { return }; + if up.limited_at.is_some_and(|t| sent <= t.mono) { + return; + } + up.strikes = 0; + let changed = up.probing; + up.probing = false; + up.probe = None; + if changed { + up.bump(); + } + changed + }; + if changed { + self.refresh(name); + } + } + + /// 发给这一家的、`sent` 发出去的一跳**没有结论**:连不上、发出去之后断了、到无响应超时还 + /// 没有内容,或者回的 429 是额度用完、没钱了(等不到,见 [`crate::health`])。在探路的话不再 + /// 探:排着的照常放,这一家怎么样交给熔断 —— 不然后面的要干等完宽限期,再一拥而上,打到一家 + /// 熔断刚停用的上游上(排着的看得见它停用了,见 [`first_ready`])。连续限流的次数不动:这不是 + /// 一个不是限流的答复。最近一次 429 之前发出去的不算 + pub fn released(&self, name: &str, sent: Instant) { + { + let mut map = self.lock(); + let Some(up) = map.get_mut(name) else { return }; + if !up.probing || up.limited_at.is_some_and(|t| sent <= t.mono) { + return; + } + up.probing = false; + up.probe = None; + up.bump(); + } + self.refresh(name); + } + + /// 上游改了名(控制面改的):暂停、退避、排着的队、探路的跟着走,和用量上限的账一样(见 + /// [`crate::upstream_limits::UpstreamLimits::rename`])。探路的名额跟着走也不会丢:[`Probe`] + /// 认号不认名字。等旧名字的叫醒:它们手上是改名之前的配置,再看时旧名字下什么都没有,照常发 + pub fn rename(self: &Arc, from: &str, to: &str) { + if from == to { + return; + } + let now = self.moment(); + let (pause, sent) = { + let mut map = self.lock(); + let Some(old) = map.remove(from) else { return }; + old.bump(); + // 新名字下多半什么都没有;改名和重载之间已经有请求用新名字发过的,两份并起来 + let up = map.entry(to.to_string()).or_insert_with(Up::new); + if up.queue.is_empty() { + up.queue_kind = old.queue_kind; + } + up.pause = match (up.pause, old.pause) { + (Some(a), Some(b)) => Some(if b.left(now) > a.left(now) { b } else { a }), + (a, b) => a.or(b), + }; + up.limited_at = match (up.limited_at, old.limited_at) { + (Some(a), Some(b)) => Some(if b.mono > a.mono { b } else { a }), + (a, b) => a.or(b), + }; + up.strikes = up.strikes.max(old.strikes); + up.probing |= old.probing; + if up.probe.is_none() { + up.probe = old.probe; + } + for w in old.queue { + if !up.queue.contains(&w) { + up.queue.push_back(w); + } + } + up.bump(); + (up.pause, up.probe.and_then(|(_, sent)| sent)) + }; + self.refresh(to); + // 到点时的那一次变化原来挂在旧名字上 + if let Some(pause) = pause.filter(|p| !p.over(now)) { + self.at(to, pause); + } + if let Some(sent) = sent { + self.at(to, grace(sent)); + } + } + + /// `span` 结束时重报一次这一家的样子、叫醒排队的。**睡一会儿起来看一眼**([`NAP`]):睡着的 + /// 时候单调时钟不走,只按它睡,醒来后还要再等完剩下的 + fn at(self: &Arc, name: &str, span: Span) { + let Ok(rt) = tokio::runtime::Handle::try_current() else { + return; + }; + let (me, name) = (Arc::downgrade(self), name.to_string()); + rt.spawn(async move { + loop { + let left = match me.upgrade() { + Some(t) => span.left(t.moment()), + None => return, + }; + if left.is_zero() { + break; + } + tokio::time::sleep(left.min(NAP)).await; + } + if let Some(t) = me.upgrade() { + t.refresh(&name); + t.bump(&name); + } + }); + } + + /// 这一家此刻暂时发不出去的样子:限流暂停、滚动窗口满了、这一期用完了。**几样都有时给 + /// 又能发得最晚的那一样**。能发是 None。探路那几秒不算:暂停已经结束了 + pub fn view(&self, name: &str) -> Option { + let now = self.moment(); + let paused = self + .lock() + .get_mut(name) + .and_then(|up| up.paused(now).map(|p| p.end_ms(now))) + .map(|until_ms| tw_api::Throttle { + kind: ThrottleKind::RateLimited, + until_ms, + limit: None, + }); + let limited = match self.limit_state(name, &self.limits_of(name)) { + LimitState::Free => None, + LimitState::WindowFull { until_ms, limit } => Some(tw_api::Throttle { + kind: ThrottleKind::WindowFull, + until_ms, + limit: Some(limit), + }), + LimitState::PeriodUsedUp { until_ms, limit } => Some(tw_api::Throttle { + kind: ThrottleKind::PeriodUsedUp, + until_ms, + limit: Some(limit), + }), + }; + match (paused, limited) { + (Some(a), Some(b)) => Some(if b.until_ms > a.until_ms { b } else { a }), + (a, b) => a.or(b), + } + } + + /// 这一家的样子和上次报的不一样了就报一条 [`tw_api::Event::ThrottleChanged`] + pub fn refresh(&self, name: &str) { + let now = self.view(name); + { + let mut map = self.lock(); + let up = map.entry(name.to_string()).or_insert_with(Up::new); + if up.shown == now { + return; + } + up.shown = now.clone(); + } + let id = self.bus.next_id(); + self.bus.emit(tw_api::Event::ThrottleChanged { + id, + provider: name.to_string(), + throttle: now, + at_ms: self.now(), + }); + } + + /// 一个要排队的请求。**丢掉就离开它排着的每一个队** + pub fn waiter(self: &Arc) -> Waiter { + Waiter { + throttles: self.clone(), + id: self.next.fetch_add(1, Ordering::Relaxed), + joined: Vec::new(), + owed: Vec::new(), + } + } + + fn join(&self, name: &str, id: u64, kind: ThrottleKind, front: bool) { + let mut map = self.lock(); + let up = map.entry(name.to_string()).or_insert_with(Up::new); + if up.queue.is_empty() { + up.queue_kind = kind; + } + match up.queue.iter().position(|w| *w == id) { + Some(0) => {} + Some(at) if front => { + up.queue.remove(at); + up.queue.push_front(id); + } + Some(_) => {} + None if front => up.queue.push_front(id), + None => up.queue.push_back(id), + } + } + + /// 离开这一家的队 + fn leave(&self, name: &str, id: u64) { + if let Some(up) = self.lock().get_mut(name) { + up.leave(id); + } + } + + /// 离开每一家的队。**认号不认名字**:排着的时候这一家可能改了名([`Self::rename`]) + fn leave_everywhere(&self, id: u64) { + for up in self.lock().values_mut() { + up.leave(id); + } + } + + #[cfg(test)] + fn queue(&self, name: &str) -> Vec { + self.lock() + .get(name) + .map(|u| u.queue.iter().copied().collect()) + .unwrap_or_default() + } +} + +impl Drop for Throttles { + fn drop(&mut self) { + for w in self + .watchers + .get_mut() + .unwrap_or_else(|p| p.into_inner()) + .drain(..) + { + w.abort(); + } + } +} + +/// 一个在排队的请求。在哪几家的队里排着,**丢掉就全部离开**:客户端走了、被手动中止了, +/// 后面的不用替它等。 +pub struct Waiter { + throttles: Arc, + pub id: u64, + joined: Vec, + /// 探过路、又被限流了的那几家:下一次等的时候回到它们的队头([`Self::keep_place`]) + owed: Vec, +} + +impl Waiter { + /// 排到这一家的队尾。已经在队里的不动 + pub fn join(&mut self, name: &str, kind: ThrottleKind) { + self.throttles.join(name, self.id, kind, false); + if !self.joined.iter().any(|n| n == name) { + self.joined.push(name.to_string()); + } + } + + /// 排到这一家的队头:探路的又被限流了,回到它原来的位置 + fn join_front(&mut self, name: &str) { + self.throttles + .join(name, self.id, ThrottleKind::RateLimited, true); + if !self.joined.iter().any(|n| n == name) { + self.joined.push(name.to_string()); + } + } + + /// 探路的又被限流了:它在这一家的队头的位置留着,**下一次等的时候才回去**([`first_ready`])。 + /// 它先去试别的候选、在别处发出去的话,不该在这里占着队头 —— 这一家的暂停结束了,后面的却要 + /// 等它在别处的那一跳 + pub fn keep_place(&mut self, name: &str) { + if !self.owed.iter().any(|n| n == name) { + self.owed.push(name.to_string()); + } + } + + /// 离开这一家的队:不等它了 + fn leave(&mut self, name: &str) { + if let Some(at) = self.joined.iter().position(|n| n == name) { + self.joined.remove(at); + self.throttles.leave(name, self.id); + } + } + + /// 离开排着的每一个队:去了一家,或者不等了 + pub fn leave_all(&mut self) { + if !std::mem::take(&mut self.joined).is_empty() { + self.throttles.leave_everywhere(self.id); + } + } +} + +impl Drop for Waiter { + fn drop(&mut self) { + self.leave_all(); + } +} + +/// 暂停结束后放过去探路的那一个。**发出去时说一声**([`Probe::sent`]):从那一刻起 +/// [`PROBE_GRACE`] 还没有结论,这一家就算好了。没发出去就丢掉的(没占到位置、被中止了),名额 +/// 还回去,下一个接着探。 +/// +/// **名额认号不认名字**([`holding`]):拿到名额、还没发出去的时候这一家改了名,名额跟着新名字 +/// 走,这里说的发出去、还回去照样落在它身上 —— 不然新名字下一直挂着一张没人管的名额,宽限期 +/// 不从它算起,这一家再也放不出一个请求 +pub struct Probe { + throttles: Arc, + ticket: u64, + sent: bool, +} + +impl Probe { + /// 这一跳发出去了 + pub fn sent(&mut self) { + self.sent = true; + let now = self.throttles.moment(); + let name = { + let mut map = self.throttles.lock(); + let Some((name, up)) = holding(&mut map, self.ticket) else { + return; + }; + if let Some((_, sent)) = &mut up.probe { + *sent = Some(now); + } + up.bump(); + name.clone() + }; + self.throttles.at(&name, grace(now)); + } + + /// 发出去了、却不作数:被手动中止了、签不了名。名额还回去,下一个接着探 + pub fn withdraw(mut self) { + self.sent = true; + if let Some((_, up)) = holding(&mut self.throttles.lock(), self.ticket) { + up.probe = None; + up.bump(); + } + } + + /// 名额还有效:拿到之后这一家没有又被限流 + pub fn valid(&self) -> bool { + let now = self.throttles.moment(); + holding(&mut self.throttles.lock(), self.ticket) + .is_some_and(|(_, up)| up.paused(now).is_none()) + } +} + +impl Drop for Probe { + fn drop(&mut self) { + if self.sent { + return; + } + if let Some((_, up)) = holding(&mut self.throttles.lock(), self.ticket) + && matches!(up.probe, Some((_, None))) + { + up.probe = None; + up.bump(); + } + } +} + +/// 等到了的那一家:占到的位置,和探路的名额(该探路时)。 +pub struct Ready { + pub slot: crate::slots::Slot, + pub probe: Option, +} + +/// 等的结局([`first_ready`])。 +pub(crate) enum Outcome { + /// 等到了:它在 `names` 里的位置,和占到的 + Ready(usize, Ready), + /// 到等待期限都没有 + Late, + /// 不用等了:等着的每一家要么这一期用完了,要么熔断着、额度用完了(还有别的能用时,见 + /// [`first_ready`])。到期限之前也发不出去 + Gone, +} + +/// 在 `names`(每一家和它的用量上限)里等先能发的那一家:不暂停、用量上限没满、排到了它、 +/// 有空位。交回它在 `names` 里的位置和 [`Ready`];等到 `until` 都没有是 [`Outcome::Late`]。 +/// **过了 `until` 就不再放**:等到了之后再发出去,就过了这个请求的等待期限。 +/// +/// 等的时候 `waiter` 排在等得到的那几家的队里(探过路又被限流的那几家排在队头,见 +/// [`Waiter::keep_place`]),**等到了、等不到都离开所有的队**:等不到的请求接着去试别的候选, +/// 不该还占着队头。满着的那几家照 [`crate::slots`] 的规矩等空位(不排这里的队)。**同时能发的 +/// 按 `names` 的顺序**,也就是候选的顺序。 +/// +/// 等的时候这一期用完了的不再等:到期限之前也重置不了。熔断着、额度用完了的([`crate::health`]) +/// 也不再等 —— **还有别的能用时**;全都停用着的照样等,和挑候选时一样(`Health::filter`),一家 +/// 都不剩的时候宁可试一试。一家都不剩是 [`Outcome::Gone`],当场交回 +pub(crate) async fn first_ready( + state: &crate::AppState, + waiter: &mut Waiter, + names: &[(&str, &[KeyLimit])], + until: tokio::time::Instant, +) -> Outcome { + let got = wait_ready(state, waiter, names, until).await; + waiter.leave_all(); + got +} + +async fn wait_ready( + state: &crate::AppState, + waiter: &mut Waiter, + names: &[(&str, &[KeyLimit])], + until: tokio::time::Instant, +) -> Outcome { + let throttles = &state.throttles; + // 探过路、又被限流了的:回到那几家的队头 + for (name, _) in names { + if let Some(at) = waiter.owed.iter().position(|n| n == name) { + waiter.owed.remove(at); + waiter.join_front(name); + } + } + // 拿到了探路的名额、在等空位的那几家:名额攥着,空出位置就用它发 + let mut held: HashMap = HashMap::new(); + loop { + if tokio::time::Instant::now() >= until { + return Outcome::Late; + } + // **先盯上再看**:看的时候发生的变化不会漏掉 + let mut changes: Vec> = + names.iter().map(|(n, _)| throttles.watch(n)).collect(); + for c in &mut changes { + c.borrow_and_update(); + } + let any_alive = names.iter().any(|(n, _)| state.health.is_available(n)); + let mut slot_waits: Vec = Vec::new(); + let mut wake: Option = None; + let mut gone = 0; + for (i, (name, limits)) in names.iter().enumerate() { + if held.get(&i).is_some_and(|p| !p.valid()) { + held.remove(&i); + } + // 停用了:交给熔断照常管,不再等它 + if any_alive && !state.health.is_available(name) { + held.remove(&i); + waiter.leave(name); + gone += 1; + continue; + } + let probe = match held.remove(&i) { + Some(p) => Some(p), + None => match throttles.gate(name, limits, Some(waiter.id)) { + Gate::Go(probe) => probe, + Gate::Wait { kind, until_ms } => { + waiter.join(name, kind); + if let Some(t) = until_ms { + wake = Some(wake.map_or(t, |w| w.min(t))); + } + continue; + } + // 等的时候这一期用完了:不等它了,到期限之前也重置不了 + Gate::UsedUp { .. } => { + waiter.leave(name); + gone += 1; + continue; + } + }, + }; + match state.slots.try_take(name) { + Some(slot) => return Outcome::Ready(i, Ready { slot, probe }), + None => { + if let Some(p) = probe { + held.insert(i, p); + } + slot_waits.push(i); + } + } + } + if gone == names.len() { + return Outcome::Gone; + } + let full: Vec<&str> = slot_waits.iter().map(|i| names[*i].0).collect(); + let changed = + futures::future::select_all(changes.iter_mut().map(|c| Box::pin(c.changed()))); + let windows = futures::future::select_all( + names + .iter() + .map(|(n, _)| Box::pin(state.upstream_limits.changed(n))), + ); + let woken = async { + match wake { + Some(at) => tokio::time::sleep(throttles.nap(at)).await, + None => std::future::pending().await, + } + }; + tokio::select! { + biased; + got = state.slots.first_free(&full, until), if !full.is_empty() => { + let Some((k, slot)) = got else { return Outcome::Late }; + let i = slot_waits[k]; + let (name, limits) = names[i]; + // 空出位置的这一刻它还能发吗:等的时候它可能又被限流了 + let probe = match held.remove(&i).filter(Probe::valid) { + Some(p) => Some(p), + None => match throttles.gate(name, limits, Some(waiter.id)) { + Gate::Go(probe) => probe, + _ => continue, + }, + }; + return Outcome::Ready(i, Ready { slot, probe }); + } + _ = changed => {} + _ = windows => {} + _ = woken => {} + _ = tokio::time::sleep_until(until) => return Outcome::Late, + } + } +} + +/// 等到最后还发不出去的一家,为什么([`refusal`])。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Stuck { + /// 满着(`max_concurrent`) + Busy, + /// 限流暂停着:预计又能发的那一刻,不知道是 None + RateLimited(Option), + /// 滚动窗口满了:同上 + WindowFull(Option), + /// 这一期用完了:重置的那一刻 + UsedUp(u64), +} + +/// 没有哪一家在等待期限之前能发:客户端收到的 429。**说出每一家、各是为什么**;`Retry-After` +/// 是最先能发的那一家还要多久(满着的按 [`crate::error::BUSY_RETRY_AFTER_SECS`] 算),可以重试。 +/// 全都是这一期用完了的,`Retry-After` 是最早的重置,**到那之前别重试**(`x-should-retry: false`, +/// 见 [`crate::error::Retry`])。都只是满着的,和以前一样是 [`crate::error::Source::Busy`] +pub fn refusal(stuck: &[(String, Stuck)], now_ms: u64) -> crate::GatewayError { + use crate::GatewayError; + use tw_types::msg; + let names = |f: &dyn Fn(&Stuck) -> bool| { + stuck + .iter() + .filter(|(_, s)| f(s)) + .map(|(n, _)| format!("`{n}`")) + .collect::>() + }; + let busy = names(&|s| matches!(s, Stuck::Busy)); + let rate_limited = names(&|s| matches!(s, Stuck::RateLimited(_))); + let window_full = names(&|s| matches!(s, Stuck::WindowFull(_))); + let used_up = names(&|s| matches!(s, Stuck::UsedUp(_))); + let all = |v: &Vec| v.len() == stuck.len(); + let list = |v: &Vec| { + if v.is_empty() { + "—".to_string() + } else { + v.join(", ") + } + }; + if all(&busy) { + return GatewayError::busy(msg!( + "gw.busy_all", upstreams = list(&busy) => + "Every upstream that can serve this request is at its concurrency limit \ + (max_concurrent): {upstreams}. None had a free slot in time; try again shortly." + )); + } + if all(&used_up) { + let reset = stuck + .iter() + .filter_map(|(_, s)| match s { + Stuck::UsedUp(t) => Some(*t), + _ => None, + }) + .min() + .unwrap_or(now_ms); + return GatewayError::rate_limited(msg!( + "gw.wait.used_up", upstreams = list(&used_up) => + "Every upstream that can serve this request has used up one of its usage limits for \ + the current day, week or month: {upstreams}. Retrying before the earliest reset does \ + not help." + )) + .with_retry(crate::error::Retry { + after_ms: reset.saturating_sub(now_ms).max(1), + until_reset: true, + }); + } + // 最先能发的那一家还要多久。说不出的(探路的还没回话、前面有人排着)按探路的宽限算 + let after_ms = stuck + .iter() + .filter_map(|(_, s)| match s { + Stuck::Busy => Some(crate::error::BUSY_RETRY_AFTER_SECS * 1000), + Stuck::RateLimited(t) | Stuck::WindowFull(t) => { + Some(t.map_or(PROBE_GRACE.as_millis() as u64, |t| t.saturating_sub(now_ms))) + } + Stuck::UsedUp(_) => None, + }) + .min() + .unwrap_or(PROBE_GRACE.as_millis() as u64) + .max(1); + let why = if all(&rate_limited) { + msg!( + "gw.wait.rate_limited", upstreams = list(&rate_limited) => + "Every upstream that can serve this request is paused after answering 429 (rate \ + limited): {upstreams}. None took requests again in time; try again shortly." + ) + } else if all(&window_full) { + msg!( + "gw.wait.window_full", upstreams = list(&window_full) => + "Every upstream that can serve this request is at one of its per-minute or per-hour \ + usage limits: {upstreams}. None had room again in time; try again shortly." + ) + } else { + msg!( + "gw.wait.mixed", + rate_limited = list(&rate_limited), window_full = list(&window_full), + used_up = list(&used_up), busy = list(&busy) => + "No upstream that can serve this request became available in time. Paused after \ + answering 429: {rate_limited}. At a per-minute or per-hour usage limit: \ + {window_full}. Used up for the current day, week or month: {used_up}. At the \ + concurrency limit: {busy}. Try again shortly." + ) + }; + GatewayError::rate_limited(why).with_retry(crate::error::Retry { + after_ms, + until_reset: false, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::AtomicU64; + + const T0: u64 = 1_800_000_000_000; + + /// 单调时钟上的零点在墙上 `T0` 之前这么久:测试里说「`T0` 之前一毫秒发出去的」也说得出 + const HEAD: u64 = 3_600_000; + + /// 拨得动的两个时钟 + struct Clocks { + wall: Arc, + /// 单调时钟,从 `base` 起的毫秒 + mono: Arc, + base: Instant, + } + + impl Clocks { + fn new() -> Self { + Self { + wall: Arc::new(AtomicU64::new(T0)), + mono: Arc::new(AtomicU64::new(HEAD)), + base: Instant::now(), + } + } + + fn throttles(&self) -> Arc { + let bus = tw_observe::EventBus::new(); + let limits = Arc::new(crate::upstream_limits::UpstreamLimits::new(bus.clone())); + self.throttles_with(bus, limits) + } + + fn throttles_with( + &self, + bus: tw_observe::EventBus, + limits: Arc, + ) -> Arc { + let (wall, mono, base) = (self.wall.clone(), self.mono.clone(), self.base); + Arc::new(Throttles::with_clocks( + bus, + limits, + Arc::new(move || wall.load(Ordering::SeqCst)), + Arc::new(move || base + Duration::from_millis(mono.load(Ordering::SeqCst))), + )) + } + + /// 两个时钟一起走 + fn advance(&self, ms: u64) { + self.wall.fetch_add(ms, Ordering::SeqCst); + self.mono.fetch_add(ms, Ordering::SeqCst); + } + + /// 墙上的此刻 + fn now(&self) -> u64 { + self.wall.load(Ordering::SeqCst) + } + + /// 墙上 `ms` 那一刻在单调时钟上是哪一刻(两个时钟一起走、没拨过时) + fn at(&self, ms: u64) -> Instant { + self.base + Duration::from_millis(ms + HEAD - T0) + } + + /// 单调时钟上的此刻 + fn sent(&self) -> Instant { + self.base + Duration::from_millis(self.mono.load(Ordering::SeqCst)) + } + } + + fn clocked() -> (Arc, Clocks) { + let c = Clocks::new(); + (c.throttles(), c) + } + + fn headers(pairs: &[(&'static str, &str)]) -> HeaderMap { + let mut h = HeaderMap::new(); + for (k, v) in pairs { + h.insert(*k, v.parse().unwrap()); + } + h + } + + fn is_go(g: &Gate) -> bool { + matches!(g, Gate::Go(_)) + } + + #[test] + fn the_pause_takes_what_the_upstream_says_in_order() { + let gemini = br#"{"error":{"details":[{"retryDelay":"33s"}]}}"#; + // 毫秒的那个最准,先看它 + let h = headers(&[("retry-after-ms", "1500"), ("retry-after", "30")]); + assert_eq!( + pause_hint(&h, gemini, T0), + Some(Duration::from_millis(1500)) + ); + // 然后是 Retry-After:秒数,或者一个 HTTP 日期 + let h = headers(&[("retry-after", "30")]); + assert_eq!(pause_hint(&h, gemini, T0), Some(Duration::from_secs(30))); + let at = chrono::DateTime::from_timestamp_millis((T0 + 90_000) as i64) + .unwrap() + .to_rfc2822(); + let h = headers(&[("retry-after", &at)]); + assert_eq!(pause_hint(&h, b"{}", T0), Some(Duration::from_secs(90))); + // 头里都没说,看正文 + assert_eq!( + pause_hint(&HeaderMap::new(), gemini, T0), + Some(Duration::from_secs(33)) + ); + // 写坏了的不算,接着往下看 + let h = headers(&[("retry-after-ms", "soon"), ("retry-after", "7")]); + assert_eq!(pause_hint(&h, b"", T0), Some(Duration::from_secs(7))); + assert_eq!(pause_hint(&HeaderMap::new(), b"slow down", T0), None); + } + + #[test] + fn without_a_hint_the_pause_backs_off_and_any_other_answer_starts_over() { + let (t, c) = clocked(); + let mut got = Vec::new(); + for _ in 0..7 { + // 每一次都是上一次暂停结束之后才发出去的那一个(探路的) + c.advance(1); + let d = t.limited("甲", c.sent(), None); + got.push(d.as_secs()); + c.advance(d.as_millis() as u64); + } + assert_eq!(got, [5, 10, 20, 40, 60, 60, 60]); + // 拿到一个不是 429 的答复:从头退避 + c.advance(1); + t.answered("甲", c.sent()); + c.advance(1); + assert_eq!(t.limited("甲", c.sent(), None), Duration::from_secs(5)); + } + + #[test] + fn a_burst_of_429s_counts_once_and_a_later_one_extends_the_pause() { + let (t, c) = clocked(); + // 同一阵发出去的五个都被限流:只算一次,暂停 5 秒 + for _ in 0..5 { + t.limited("甲", c.at(T0), None); + } + assert!(matches!( + t.peek("甲", &[]), + Peek::Wait { kind: ThrottleKind::RateLimited, until_ms: Some(u) } if u == T0 + 5_000 + )); + // 暂停中又一个说了 30 秒:延到晚的那个 + c.advance(1_000); + let d = t.limited("甲", c.at(T0), Some(Duration::from_secs(30))); + assert_eq!(d, Duration::from_secs(30)); + // 说得更短的不往前提 + let d = t.limited("甲", c.at(T0), Some(Duration::from_secs(1))); + assert_eq!(d, Duration::from_secs(30)); + // 最长暂停按配置封顶 + let (t, c) = clocked(); + let d = t.limited("乙", c.at(T0), Some(Duration::from_secs(99_999))); + assert_eq!(d, Duration::from_secs(3600)); + } + + #[test] + fn after_a_pause_only_one_request_goes_ahead_until_it_is_answered() { + let (t, c) = clocked(); + t.limited("甲", c.at(T0), None); + assert!(!is_go(&t.gate("甲", &[], None)), "暂停中"); + c.advance(5_000); + // 暂停结束:头一个去探路 + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], None) else { + panic!("暂停结束了该放一个去探路"); + }; + // 它在路上的时候别的都等 + assert!(matches!( + t.gate("甲", &[], None), + Gate::Wait { + kind: ThrottleKind::RateLimited, + until_ms: None + } + )); + probe.sent(); + let sent = c.sent(); + // 什么时候有结论说不出:宽限期的结束不是知道的恢复时刻 + assert!(matches!( + t.peek("甲", &[]), + Peek::Wait { + kind: ThrottleKind::RateLimited, + until_ms: None + } + )); + assert!(t.keeps_conversation("甲", &[], c.now() + 1_000)); + // 它拿到了响应头,不是 429:放开 + t.answered("甲", sent); + assert!(matches!(t.gate("甲", &[], None), Gate::Go(None))); + } + + #[test] + fn a_probe_that_is_limited_again_pauses_longer_and_one_without_a_verdict_passes_after_the_grace() + { + let (t, c) = clocked(); + t.limited("甲", c.at(T0), None); + c.advance(5_000); + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], None) else { + panic!() + }; + probe.sent(); + // 探路的又被限流:退避接着往上走 + assert_eq!(t.limited("甲", c.sent(), None), Duration::from_secs(10)); + assert!(!probe.valid()); + drop(probe); + c.advance(10_000); + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], None) else { + panic!() + }; + probe.sent(); + drop(probe); + // 一直没有结论:10 秒之后当它好了 + c.advance(9_999); + assert!(!is_go(&t.gate("甲", &[], None))); + c.advance(1); + assert!(matches!(t.gate("甲", &[], None), Gate::Go(None))); + } + + #[test] + fn a_probe_dropped_before_it_was_sent_lets_the_next_one_go() { + let (t, c) = clocked(); + t.limited("甲", c.at(T0), None); + c.advance(5_000); + let Gate::Go(Some(probe)) = t.gate("甲", &[], None) else { + panic!() + }; + // 没占到位置、被中止了:名额还回去 + drop(probe); + assert!(matches!(t.gate("甲", &[], None), Gate::Go(Some(_)))); + } + + #[test] + fn an_answer_to_a_request_sent_before_the_429_does_not_end_the_pause() { + let (t, c) = clocked(); + c.advance(1_000); + t.limited("甲", c.at(T0 + 1_000), None); + c.advance(5_000); + // 暂停之前发出去的长回答这时才到响应头:不说明这家好了 + t.answered("甲", c.at(T0)); + assert!(matches!(t.gate("甲", &[], None), Gate::Go(Some(_)))); + // **按单调时钟比**,不差一毫秒:429 那一刻之后才发出去的,哪怕只晚一点点,也是新的一跳 + let (t, c) = clocked(); + t.limited("甲", c.sent(), None); + let after = c.sent() + Duration::from_micros(1); + c.advance(5_000); + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], None) else { + panic!() + }; + probe.sent(); + t.answered("甲", after); + assert!(matches!(t.gate("甲", &[], None), Gate::Go(None))); + } + + #[test] + fn the_queue_lets_them_go_in_order_and_a_leaver_makes_room() { + let (t, c) = clocked(); + t.limited("甲", c.at(T0), None); + let (mut a, mut b) = (t.waiter(), t.waiter()); + a.join("甲", ThrottleKind::RateLimited); + b.join("甲", ThrottleKind::RateLimited); + assert_eq!(t.queue("甲"), [a.id, b.id]); + c.advance(5_000); + // 后来的、不在队里的不插队 + assert!(matches!( + t.gate("甲", &[], None), + Gate::Wait { + kind: ThrottleKind::RateLimited, + until_ms: None + } + )); + assert!(!is_go(&t.gate("甲", &[], Some(b.id))), "b 排在 a 后面"); + // 队头探路 + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], Some(a.id)) else { + panic!() + }; + a.leave_all(); + probe.sent(); + // 探路的又被限流:它回到队头 + t.limited("甲", c.sent(), None); + a.join_front("甲"); + assert_eq!(t.queue("甲"), [a.id, b.id]); + // a 不等了(客户端走了):b 排到了头 + drop(a); + assert_eq!(t.queue("甲"), [b.id]); + c.advance(10_000); + assert!(matches!(t.gate("甲", &[], Some(b.id)), Gate::Go(Some(_)))); + } + + #[tokio::test] + async fn every_change_is_reported_once() { + let (t, c) = clocked(); + let mut rx = t.bus.subscribe(); + t.limited("甲", c.at(T0), None); + t.limited("甲", c.at(T0), None); + c.advance(5_000); + t.refresh("甲"); + t.refresh("甲"); + let mut got = Vec::new(); + while let Ok(ev) = rx.try_recv() { + if let tw_api::Event::ThrottleChanged { + provider, throttle, .. + } = ev + { + got.push((provider, throttle)); + } + } + assert_eq!( + got, + vec![ + ( + "甲".to_string(), + Some(tw_api::Throttle { + kind: ThrottleKind::RateLimited, + until_ms: T0 + 5_000, + limit: None, + }) + ), + ("甲".to_string(), None), + ] + ); + } + + /// 用量上限的时钟:比墙上的慢 `behind` 毫秒 + struct Behind { + wall: Arc, + behind: u64, + } + + impl crate::key_limits::Clock for Behind { + fn now_ms(&self) -> i64 { + (self.wall.load(Ordering::SeqCst) - self.behind) as i64 + } + fn period(&self, per: tw_config::LimitPer, at_ms: i64) -> (i64, i64) { + crate::key_limits::SystemClock.period(per, at_ms) + } + fn show(&self, at_ms: i64) -> String { + crate::key_limits::SystemClock.show(at_ms) + } + } + + /// 用量上限的时钟和这里的先后读,差一毫秒:报出去的时刻**不跟着抖**,同一个样子只报一次 + #[tokio::test] + async fn a_full_window_is_reported_once_although_the_two_clocks_are_read_apart() { + let c = Clocks::new(); + let bus = tw_observe::EventBus::new(); + let limits = Arc::new(crate::upstream_limits::UpstreamLimits::with_clock( + bus.clone(), + Arc::new(Behind { + wall: c.wall.clone(), + behind: 1, + }), + )); + let t = c.throttles_with(bus.clone(), limits.clone()); + let set: Vec = serde_yaml_ng::from_str("[{per: minute, requests: 1}]").unwrap(); + let cfg = tw_config::Config { + providers: vec![tw_config::Provider { + name: "甲".into(), + base_url: "https://a.example".into(), + limits: set.clone(), + ..Default::default() + }], + ..Default::default() + }; + limits.configure(&cfg); + t.configure(&cfg); + let mut rx = bus.subscribe(); + let _hold = limits.reserve("甲", &set, &crate::key_limits::Ask::default()); + let until = T0 - 1 + 60_000; + for _ in 0..3 { + assert_eq!(t.view("甲").map(|v| v.until_ms), Some(until)); + t.refresh("甲"); + c.advance(1); + } + let shown: Vec<_> = std::iter::from_fn(|| rx.try_recv().ok()) + .filter_map(|ev| match ev { + tw_api::Event::ThrottleChanged { throttle, .. } => Some(throttle), + _ => None, + }) + .collect(); + assert_eq!(shown.len(), 1, "{shown:?}"); + } + + /// 重载一次不多一批后台任务:上一份配置的那些收工,删掉了用量上限的那一家不留任务 + #[tokio::test] + async fn reloading_does_not_pile_up_watchers() { + let (t, _c) = clocked(); + let with = |limits: &str| tw_config::Config { + providers: vec![tw_config::Provider { + name: "甲".into(), + base_url: "https://a.example".into(), + limits: serde_yaml_ng::from_str(limits).unwrap(), + ..Default::default() + }], + ..Default::default() + }; + let alive = || async { + for _ in 0..10 { + tokio::task::yield_now().await; + } + tokio::runtime::Handle::current() + .metrics() + .num_alive_tasks() + }; + let before = alive().await; + for _ in 0..5 { + t.configure(&with("[{per: minute, requests: 1}]")); + } + assert_eq!(alive().await, before + 1); + // 用量上限删了(改名、删掉上游也一样):它的任务不再被叫醒,要当场收工 + t.configure(&with("[]")); + assert_eq!(alive().await, before); + // 网关没了,任务跟着收工 + t.configure(&with("[{per: minute, requests: 1}]")); + assert_eq!(alive().await, before + 1); + drop(t); + assert_eq!(alive().await, before); + } + + #[test] + fn the_final_429_names_each_upstream_and_when_to_come_back() { + let n = |v: &[(&str, Stuck)]| { + v.iter() + .map(|(a, b)| (a.to_string(), *b)) + .collect::>() + }; + // 都在暂停:最先恢复的那一家还要多久,可以重试 + let e = refusal( + &n(&[ + ("甲", Stuck::RateLimited(Some(T0 + 30_000))), + ("乙", Stuck::RateLimited(Some(T0 + 1_500))), + ]), + T0, + ); + assert_eq!(e.detail.code, "gw.wait.rate_limited"); + assert_eq!(e.detail.arg("upstreams"), "`甲`, `乙`"); + assert_eq!(e.source, crate::error::Source::RateLimited); + assert_eq!( + e.retry, + Some(crate::error::Retry { + after_ms: 1_500, + until_reset: false + }) + ); + // 全都是这一期用完了:最早的重置,别重试 + let e = refusal( + &n(&[ + ("甲", Stuck::UsedUp(T0 + 7_200_000)), + ("乙", Stuck::UsedUp(T0 + 3_600_000)), + ]), + T0, + ); + assert_eq!(e.detail.code, "gw.wait.used_up"); + assert_eq!( + e.retry, + Some(crate::error::Retry { + after_ms: 3_600_000, + until_reset: true + }) + ); + // 只是满着:和以前一样 + let e = refusal(&n(&[("甲", Stuck::Busy)]), T0); + assert_eq!(e.detail.code, "gw.busy_all"); + assert_eq!(e.source, crate::error::Source::Busy); + assert_eq!(e.retry.unwrap().after_ms, 5_000); + // 几样都有:每一家各在它那一类里;用完了的不算进多久之后再来 + let e = refusal( + &n(&[ + ("甲", Stuck::WindowFull(Some(T0 + 20_000))), + ("乙", Stuck::UsedUp(T0 + 1_000)), + ("丙", Stuck::Busy), + ]), + T0, + ); + assert_eq!(e.detail.code, "gw.wait.mixed"); + assert_eq!(e.detail.arg("window_full"), "`甲`"); + assert_eq!(e.detail.arg("used_up"), "`乙`"); + assert_eq!(e.detail.arg("busy"), "`丙`"); + assert_eq!(e.detail.arg("rate_limited"), "—"); + assert_eq!( + e.retry, + Some(crate::error::Retry { + after_ms: 5_000, + until_reset: false + }) + ); + // 已经过了的时刻不给 0 + let e = refusal(&n(&[("甲", Stuck::WindowFull(Some(T0 - 5)))]), T0); + assert_eq!(e.retry.unwrap().after_ms, 1); + } + + #[test] + fn a_hint_that_says_nothing_falls_back_to_the_backoff_and_an_absurd_one_is_capped() { + let (t, c) = clocked(); + // `retry-after-ms: 0`、不到一毫秒:照它暂停等于不暂停,按退避算 + assert_eq!( + t.limited("甲", c.at(T0), Some(Duration::ZERO)), + Duration::from_secs(5) + ); + c.advance(5_000); + assert_eq!( + t.limited("甲", c.at(T0 + 5_000), Some(Duration::from_micros(400))), + Duration::from_secs(10) + ); + let h = headers(&[("retry-after-ms", "0"), ("retry-after", "0")]); + assert_eq!(pause_hint(&h, b"", T0), None); + let h = headers(&[("retry-after-ms", "0.4"), ("retry-after", "3")]); + assert_eq!(pause_hint(&h, b"", T0), Some(Duration::from_secs(3))); + // 大得离谱的不 panic,照最长暂停封顶 + let h = headers(&[("retry-after-ms", "1e300")]); + let hint = pause_hint(&h, b"", T0); + assert!(hint.is_some_and(|d| d > Duration::from_secs(3600))); + let (t, c) = clocked(); + assert_eq!(t.limited("乙", c.at(T0), hint), Duration::from_secs(3600)); + } + + #[test] + fn the_pause_runs_on_the_monotonic_clock_when_the_system_time_moves_back() { + let (t, c) = clocked(); + assert_eq!(t.limited("甲", c.sent(), None), Duration::from_secs(5)); + // 系统时间往回拨了一小时:暂停照样只剩 5 秒,报出去的时刻按拨过的算 + c.wall.fetch_sub(3_600_000, Ordering::SeqCst); + assert!(matches!( + t.peek("甲", &[]), + Peek::Wait { until_ms: Some(u), .. } if u == T0 - 3_600_000 + 5_000 + )); + assert_eq!(t.view("甲").unwrap().until_ms, T0 - 3_600_000 + 5_000); + c.mono.fetch_add(4_999, Ordering::SeqCst); + assert!(!is_go(&t.gate("甲", &[], None))); + c.mono.fetch_add(1, Ordering::SeqCst); + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], None) else { + panic!("暂停该结束了"); + }; + // 探路的宽限也拉不长:系统时间再往回拨,宽限照样 10 秒 + probe.sent(); + c.wall.fetch_sub(3_600_000, Ordering::SeqCst); + c.mono.fetch_add(10_000, Ordering::SeqCst); + assert!(matches!(t.gate("甲", &[], None), Gate::Go(None))); + } + + /// 合上盖子睡着的时候单调时钟不走、墙上的走:**睡过去的时间算进暂停**,醒来时该结束的已经 + /// 结束了,不再等完剩下的(最长要等一小时) + #[test] + fn a_pause_slept_through_is_over_when_the_machine_wakes() { + let (t, c) = clocked(); + assert_eq!( + t.limited("甲", c.sent(), Some(Duration::from_secs(600))), + Duration::from_secs(600) + ); + // 睡了五分钟:还剩五分钟,报出去的结束时刻不变 + c.wall.fetch_add(300_000, Ordering::SeqCst); + assert!(matches!( + t.peek("甲", &[]), + Peek::Wait { until_ms: Some(u), .. } if u == T0 + 600_000 + )); + // 醒着又过了五分钟(两个时钟一起走):结束 + c.advance(299_999); + assert!(!is_go(&t.gate("甲", &[], None))); + c.advance(1); + assert!(t.view("甲").is_none()); + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], None) else { + panic!("暂停该结束了"); + }; + // 一下子睡过整个暂停:醒来就结束 + probe.sent(); + let sent = c.sent(); + c.advance(1); + assert_eq!( + t.limited( + "甲", + sent + Duration::from_millis(1), + Some(Duration::from_secs(3_000)) + ), + Duration::from_secs(3_000) + ); + c.wall.fetch_add(3_000_000, Ordering::SeqCst); + assert!(t.view("甲").is_none()); + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], None) else { + panic!("睡过了整个暂停,醒来就该结束"); + }; + // 探路的宽限也一样:发出去之后睡着了,醒来就过了宽限 + probe.sent(); + assert!(!is_go(&t.gate("甲", &[], None))); + c.wall.fetch_add(10_000, Ordering::SeqCst); + assert!(matches!(t.gate("甲", &[], None), Gate::Go(None))); + } + + #[test] + fn a_probe_without_a_verdict_hands_the_queue_back_without_waiting_for_the_grace() { + let (t, c) = clocked(); + t.limited("甲", c.at(T0), None); + c.advance(5_000); + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], None) else { + panic!() + }; + probe.sent(); + let sent = c.sent(); + assert!(!is_go(&t.gate("甲", &[], None))); + // 暂停之前发出去的那一跳连不上:不是探路的,不算 + t.released("甲", c.at(T0 - 1)); + assert!(!is_go(&t.gate("甲", &[], None))); + // 探路的连不上(或者回的是额度用完):不再探,排着的照常放,交给熔断 + t.released("甲", sent); + assert!(matches!(t.gate("甲", &[], None), Gate::Go(None))); + // 连续限流的次数没有归零:下一次接着往上退避 + c.advance(1); + assert_eq!(t.limited("甲", c.sent(), None), Duration::from_secs(10)); + // 被中止的探路:名额还回去,下一个接着探 + c.advance(10_000); + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], None) else { + panic!() + }; + probe.sent(); + assert!(!is_go(&t.gate("甲", &[], None))); + probe.withdraw(); + assert!(matches!(t.gate("甲", &[], None), Gate::Go(Some(_)))); + } + + #[test] + fn an_upstream_is_wary_while_probing_and_for_a_while_after_a_429() { + let (t, c) = clocked(); + assert!(!t.wary("甲"), "没被限流过"); + t.limited("甲", c.at(T0), None); + assert!(t.wary("甲")); + c.advance(5_000); + t.answered("甲", c.at(T0 + 5_000)); + // 探路的过了、连续的次数归零:最近被限流过,还要当心 + assert!(t.wary("甲")); + c.advance(295_000); + assert!(!t.wary("甲")); + } + + #[tokio::test] + async fn a_renamed_upstream_keeps_its_pause_and_its_queue() { + let (t, c) = clocked(); + let mut rx = t.bus.subscribe(); + t.limited("甲", c.at(T0), Some(Duration::from_secs(30))); + let mut a = t.waiter(); + a.join("甲", ThrottleKind::RateLimited); + let old = t.watch("甲"); + t.rename("甲", "乙"); + // 暂停和队跟着新名字走;旧名字下什么都没有了 + assert!(matches!( + t.peek("乙", &[]), + Peek::Wait { kind: ThrottleKind::RateLimited, until_ms: Some(u) } if u == T0 + 30_000 + )); + assert_eq!(t.queue("乙"), [a.id]); + assert_eq!(t.peek("甲", &[]), Peek::Open); + // 等旧名字的被叫醒了 + assert!(old.has_changed().unwrap_or(true)); + // 新名字报一次现在的样子 + let mut shown = Vec::new(); + while let Ok(ev) = rx.try_recv() { + if let tw_api::Event::ThrottleChanged { + provider, throttle, .. + } = ev + { + shown.push((provider, throttle.map(|t| t.until_ms))); + } + } + assert_eq!( + shown.last(), + Some(&("乙".to_string(), Some(T0 + 30_000))), + "{shown:?}" + ); + // 它排着的时候改了名:离开时照样从新名字的队里走掉 + drop(a); + assert!(t.queue("乙").is_empty()); + // 退避跟着走:暂停结束后再被限流是第二次 + c.advance(30_000); + assert_eq!( + t.limited("乙", c.at(T0 + 30_000), None), + Duration::from_secs(10) + ); + } + + /// 探路的名额拿到了、还没发出去时改了名:名额跟着新名字走,[`Probe`] 手上的是旧名字也照样 + /// 落在它身上 —— 发出去从那一刻起算宽限,没发出去就还回去 + #[test] + fn a_probe_ticket_follows_a_rename() { + for send in [false, true] { + let (t, c) = clocked(); + t.limited("甲", c.at(T0), None); + c.advance(5_000); + let Gate::Go(Some(mut probe)) = t.gate("甲", &[], None) else { + panic!() + }; + t.rename("甲", "乙"); + assert!(probe.valid()); + assert!(!is_go(&t.gate("乙", &[], None)), "探路的还没有结论"); + if send { + probe.sent(); + drop(probe); + c.advance(PROBE_GRACE.as_millis() as u64); + assert!(matches!(t.gate("乙", &[], None), Gate::Go(None))); + } else { + drop(probe); + assert!(matches!(t.gate("乙", &[], None), Gate::Go(Some(_)))); + } + } + } +} diff --git a/crates/tw-gateway/src/upstream_limits.rs b/crates/tw-gateway/src/upstream_limits.rs new file mode 100644 index 00000000..d37e1ebf --- /dev/null +++ b/crates/tw-gateway/src/upstream_limits.rs @@ -0,0 +1,633 @@ +//! 每家上游的用量上限:数、看、预留、结算。 +//! +//! 写法和密钥的一样(`tw_config::limits`,写在 `providers[].limits`),数法也一样(账本和 +//! 密钥的共用,见 [`crate::key_limits::books`])。**不一样的是数哪几跳、到了怎么办。** +//! +//! **数什么。**请求数:发到这家的每一跳算一个,成功失败都算,**这家回了 429 的不算** —— 它 +//! 没接,请求换一家或者等它,不该把这家的上限用掉;没发出去的也不算,数 token 的请求不算。 +//! token、费用:接下请求的那一跳(尝试链的最后一跳)的用量,记在那一家名下,算法和密钥的 +//! 一样(没走缓存的输入 + 写进缓存的 + 输出,开了 `cache_reads` 再加从缓存读的;费用是记下的, +//! 没有价格、不计费的算 0)。**都按存储层记下的那一行算**([`UpstreamLimits::settle`],哪几跳 +//! 算看 `tw_api::RoutingView::upstream_hops`、用量算谁的看 `served_by`):重启之后从库里加 +//! 回来的([`UpstreamLimits::rebuild`])和此刻内存里的是同一份数。 +//! +//! **到了怎么办。**这里不拒绝谁,只说这家此刻能不能发([`UpstreamLimits::state`]):分钟、 +//! 小时的窗口满了,什么时候空出来算得出来,是等得到的;天、周、月用满了,到下一期之前只能 +//! 跳过。换一家、排队、等多久由管线决定(`server::pipeline::hop`),它在 +//! [`UpstreamLimits::changed`] 上等这家空出来。 +//! +//! **在跑的怎么算。**一跳发出去的那一刻按输入估一个数占着([`UpstreamLimits::reserve`]:一个 +//! 请求、输入 token 的估算、按这家算的输入费用)。这一跳回了 429、没发出去,整份还回去 +//! ([`UpstreamHold::release`],丢掉也一样);别的结局交给请求号([`UpstreamHold::bind`]), +//! 存储层记下那一行时换成实数。几个请求同时发往一家时,超出上限的最多是在路上的那些。 +//! +//! **什么时候算到哪一期。**和密钥的一样:天、周、月按请求开始的时刻归期,一期的开头变了就从 +//! 请求记录里把新的这一期重新加起来([`UpstreamLimits::reread_with`]);滚动窗口里请求数记在 +//! 发出的那一刻,token 和费用记在结算的那一刻。 + +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; +use std::time::Duration; + +use tokio::sync::Notify; +use tw_config::KeyLimit; + +use crate::key_limits::books::{self, Amount, Books, CALENDAR, Period, Reservation}; +use crate::key_limits::{Ask, Clock, Moved, Recorded, Reread, SystemClock, uncounted}; + +/// 一家上游此刻能不能发。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LimitState { + Free, + /// 滚动窗口满了:到 `until_ms`(Unix 毫秒)空出位置。`limit` 是满了的那一条在这家 + /// `limits` 里的下标;几条都满了的,是空出来得最晚的那一条 + WindowFull { + until_ms: u64, + limit: u32, + }, + /// 天、周、月用满了:`until_ms` 是下一期开始的时刻。`limit` 同上;几条都满了的,是重新 + /// 开始得最晚的那一条。**它比窗口满了要紧**:两样都有时报它 —— 等不到 + PeriodUsedUp { + until_ms: u64, + limit: u32, + }, +} + +#[derive(Default)] +struct Inner { + books: HashMap, + /// 在路上的那几跳的预留,按一个自己的号 + held: HashMap, + /// 请求号 → 它发出去的那几跳的预留的号。**一个请求可以有好几跳**,同一家也可以有两跳 + by_request: HashMap>, + seq: u64, + /// 此刻配置里每家上游的上限。结算时要不要记滚动窗口、[`UpstreamLimits::changed`] 算 + /// 什么时候再看一眼,看它 + limits: HashMap>, + /// 已经要过、还没读回来的几期,一期只要一次(见 `KeyLimits` 的同名字段) + asked: HashSet, + /// 等这家变化的([`UpstreamLimits::changed`]) + wake: HashMap>, +} + +struct Shared { + inner: Mutex, + clock: Arc, + /// 一期的开头变了时从请求记录里重新加([`UpstreamLimits::reread_with`])。没接上是 None: + /// 那时从 0 起 + reread: Mutex>, + /// 认请求有没有结束(见 [`books::GRACE_MS`]) + bus: tw_observe::EventBus, +} + +impl Shared { + fn lock(&self) -> MutexGuard<'_, Inner> { + self.inner.lock().unwrap_or_else(PoisonError::into_inner) + } + + /// 放掉结束了太久、存储层一直没交来那一行的预留;顺手把滚动窗口里滑出去的扔掉。交回 + /// 放掉了预留的那几家 + fn sweep(&self, g: &mut Inner, now: i64) -> Vec { + let mut touched = Vec::new(); + let before: HashMap = + g.held.iter().map(|(s, r)| (*s, r.key.clone())).collect(); + for (seq, id) in books::sweep_held(&mut g.held, &self.bus, now) { + if let Some(seqs) = g.by_request.get_mut(&id) { + seqs.retain(|s| *s != seq); + if seqs.is_empty() { + g.by_request.remove(&id); + } + } + if let Some(k) = before.get(&seq) { + touched.push(k.clone()); + } + } + books::trim_windows(&mut g.books, &g.limits, now); + touched + } + + /// 叫醒等这几家的(在锁外面调) + fn wake<'a>(&self, names: impl IntoIterator) { + let notifies: Vec> = { + let g = self.lock(); + names + .into_iter() + .filter_map(|n| g.wake.get(n).cloned()) + .collect() + }; + for n in notifies { + n.notify_waiters(); + } + } + + /// 去请求记录里重新加这几期。没接上记录的,从 0 起 + fn ask_reread(&self, moved: Vec) { + if moved.is_empty() { + return; + } + let f = self + .reread + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone(); + match f { + Some(f) => moved.into_iter().for_each(|m| f(m)), + None => self.lock().asked.clear(), + } + } + + fn bind(&self, seq: u64, request: u64) { + let mut g = self.lock(); + if let Some(r) = g.held.get_mut(&seq) { + r.request = Some(request); + g.by_request.entry(request).or_default().push(seq); + } + } + + /// 一跳回了 429、没发出去([`UpstreamHold`] 的 Drop):预留和滚动窗口里的那一个请求都 + /// 还回去 + fn release(&self, seq: u64) { + let key = { + let mut g = self.lock(); + let g = &mut *g; + let Some(r) = g.held.remove(&seq) else { + return; + }; + if let Some(id) = r.request + && let Some(seqs) = g.by_request.get_mut(&id) + { + seqs.retain(|s| *s != seq); + if seqs.is_empty() { + g.by_request.remove(&id); + } + } + books::unrecord(&mut g.books, &r); + r.key + }; + self.wake([key.as_str()]); + } +} + +/// 一跳发出去时占着的那一份([`UpstreamLimits::reserve`])。**丢掉就还回去**:回了 429、 +/// 没发出去的那一跳不会一直占着这家的额度。上游接下了(回的不是 429,或者发出去之后才 +/// 失败)就交给请求号([`Self::bind`]),等存储层结算。 +#[must_use = "dropping it gives the reservation back"] +pub struct UpstreamHold { + owner: Option>, + seq: u64, +} + +impl UpstreamHold { + /// 没有预留:这家没设上限 + pub fn none() -> Self { + Self { + owner: None, + seq: 0, + } + } + + /// 这一跳算数:上游回了 429 以外的答复,或者发出去之后才失败(超时、断了)。预留跟着 + /// 请求号,等存储层记下那一行时换成实数;请求数留在窗口里 + pub fn bind(mut self, request: u64) { + if let Some(owner) = self.owner.take() { + owner.bind(self.seq, request); + } + } + + /// 这一跳不算:上游回了 429,或者没发出去。整份还回去 —— 和丢掉它一样,写出来是为了 + /// 读的人看得见 + pub fn release(self) {} +} + +impl std::fmt::Debug for UpstreamHold { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UpstreamHold") + .field("held", &self.owner.is_some()) + .finish() + } +} + +impl Drop for UpstreamHold { + fn drop(&mut self) { + if let Some(owner) = self.owner.take() { + owner.release(self.seq); + } + } +} + +/// 全部上游的用量账。**跨重载存活**(在 `AppState` 里):改一条路由规则不该让这家今天用的 +/// 归零。按上游的名字记。 +pub struct UpstreamLimits { + shared: Arc, +} + +/// 这一行算不算进上游的用量:网关自己答的、数 token 的不算。发到了哪几家看尝试链 +fn counts(rec: &Recorded) -> bool { + !rec.local && !uncounted(&rec.path) +} + +impl UpstreamLimits { + pub fn new(bus: tw_observe::EventBus) -> Self { + Self::with_clock(bus, Arc::new(SystemClock)) + } + + pub fn with_clock(bus: tw_observe::EventBus, clock: Arc) -> Self { + Self { + shared: Arc::new(Shared { + inner: Mutex::default(), + clock, + reread: Mutex::default(), + bus, + }), + } + } + + fn now(&self) -> i64 { + self.shared.clock.now_ms() + } + + /// 此刻,按这里的时钟(Unix 毫秒)。[`LimitState`] 里的时刻都是这个时钟上的 + pub fn now_ms(&self) -> i64 { + self.now() + } + + /// 接上请求记录:一期的开头变了(到了下一期、机器换了时区)时,用 `f` 把新的这一期从 + /// 记录里重新加起来,加好了交回 [`Self::reread`]。和 `KeyLimits::reread_with` 一样:`f` 在 + /// 锁外面调、不该等,读回来之前这一期先从 0 起,同一期只要一次 + pub fn reread_with(&self, f: Reread) { + *self + .shared + .reread + .lock() + .unwrap_or_else(PoisonError::into_inner) = Some(f); + } + + /// 读回来了:`rows` 是从 `m.start` 起记下的那些,按上游分好(存储层的 + /// `upstream_usage_since`),`m.key` 是上游名。这一期的数换成它们加起来的;这一期的开头 + /// 又变了的话不换。调用方要保证读的时候没有结算在进行(见 `KeyLimits::reread`) + pub fn reread(&self, m: &Moved, rows: &[(String, Recorded)]) { + let mut sum = Amount::default(); + for (_, r) in rows.iter().filter(|(up, r)| *up == m.key && counts(r)) { + sum.add(&r.amount()); + } + { + let mut g = self.shared.lock(); + g.asked.remove(m); + let Some(p) = g + .books + .get_mut(&m.key) + .and_then(|b| b.periods[books::calendar_index(m.per)].as_mut()) + .filter(|p| p.start == m.start) + else { + return; + }; + p.sum = sum; + } + self.shared.wake([m.key.as_str()]); + } + + /// 配置换了:记下每家此刻的上限。**账不动**:上限改了,已经用掉的照样算。等着的都叫醒 + /// 看一眼:上限可能松了 + pub fn configure(&self, cfg: &tw_config::Config) { + let limits = cfg + .providers + .iter() + .filter(|p| !p.limits.is_empty()) + .map(|p| (p.name.clone(), p.limits.clone())) + .collect(); + let all: Vec> = { + let mut g = self.shared.lock(); + g.limits = limits; + g.wake.values().cloned().collect() + }; + for n in all { + n.notify_waiters(); + } + } + + /// 上游改了名(控制面改的):账跟着走。手改配置文件改的名认不出来 —— 那和删掉一家、 + /// 新建一家看起来一样,新的那家从 0 起 + pub fn rename(&self, from: &str, to: &str) { + if from == to { + return; + } + let woken = { + let mut g = self.shared.lock(); + if let Some(b) = g.books.remove(from) { + g.books.insert(to.to_string(), b); + } + for r in g.held.values_mut() { + if r.key == from { + r.key = to.to_string(); + } + } + if let Some(l) = g.limits.remove(from) { + g.limits.insert(to.to_string(), l); + } + g.wake.remove(from) + }; + // 等旧名字的叫醒:它们再看时用的是新名字 + if let Some(n) = woken { + n.notify_waiters(); + } + } + + /// 这家此刻能不能发。**在路上的也算**:发出去还没结束的按预留算,和 [`Self::view`] 是 + /// 同一个数。天、周、月用满了的先报(等不到);都没满再看分钟、小时的窗口 + pub fn state(&self, upstream: &str, limits: &[KeyLimit]) -> LimitState { + if limits.is_empty() { + return LimitState::Free; + } + let now = self.now(); + let clock = self.shared.clock.as_ref(); + let (state, moved, touched) = { + let mut g = self.shared.lock(); + let g = &mut *g; + let touched = self.shared.sweep(g, now); + let mut period: Option<(i64, u32)> = None; + let mut window: Option<(i64, u32)> = None; + for (i, l) in limits.iter().enumerate() { + let used = books::used(&mut g.books, &g.held, clock, upstream, l, now); + if used < l.max() { + continue; + } + let i = i as u32; + match l.window_ms() { + None => { + let end = clock.period(l.per, now).1; + if period.is_none_or(|(t, _)| end > t) { + period = Some((end, i)); + } + } + Some(_) => { + let at = books::free_at(&g.books, &g.held, upstream, l, used, now); + if window.is_none_or(|(t, _)| at > t) { + window = Some((at, i)); + } + } + } + } + let ms = |t: i64| t.max(0) as u64; + let state = match (period, window) { + (Some((t, limit)), _) => LimitState::PeriodUsedUp { + until_ms: ms(t), + limit, + }, + (None, Some((t, limit))) => LimitState::WindowFull { + until_ms: ms(t), + limit, + }, + (None, None) => LimitState::Free, + }; + ( + state, + books::take_moved(&mut g.books, &mut g.asked, upstream), + touched, + ) + }; + self.shared.wake(touched.iter().map(String::as_str)); + self.shared.ask_reread(moved); + state + } + + /// 一跳发出去的那一刻:记上一个请求(有分钟、小时上限的进滚动窗口),按输入估的 token 和 + /// 费用占着。**不看满没满** —— 那是 [`Self::state`] 的事,发不发由调用方定。这家没设上限的 + /// 什么都不记 + pub fn reserve(&self, upstream: &str, limits: &[KeyLimit], ask: &Ask) -> UpstreamHold { + if limits.is_empty() { + return UpstreamHold::none(); + } + let now = self.now(); + let seq = { + let mut g = self.shared.lock(); + let g = &mut *g; + books::record_request(&mut g.books, upstream, limits, now); + g.seq += 1; + g.held.insert( + g.seq, + Reservation::new(upstream, now, ask.tokens, ask.cost_micros), + ); + g.seq + }; + // 记上的这一个可能把窗口、这一期用满了:看着这家的(报状态的)再看一眼 + self.shared.wake([upstream]); + UpstreamHold { + owner: Some(self.shared.clone()), + seq, + } + } + + /// 存储层记下了一行:这个请求的几跳的预留换成实数。`at_ms` 是请求开始的时刻,算在哪一期 + /// 看它。`hops` 是算数的那几跳各发给了哪一家(`tw_api::RoutingView::upstream_hops`), + /// `served` 是用量记在哪一家(`served_by`)。 + /// + /// 记录里算数的那几跳,每一跳给那家记一个请求;用量记在 `served` 名下。**记录里不算的那 + /// 几跳把发出时记上的也还回去**:网关自己答的、数 token 的请求,和记录里没有的那几跳 + pub fn settle( + &self, + id: u64, + at_ms: i64, + hops: &[String], + served: Option<&str>, + rec: &Recorded, + ) { + let now = self.now(); + let clock = self.shared.clock.as_ref(); + let (touched, moved) = { + let mut g = self.shared.lock(); + let g = &mut *g; + let held: Vec = g + .by_request + .remove(&id) + .unwrap_or_default() + .into_iter() + .filter_map(|s| g.held.remove(&s)) + .collect(); + let mut touched: HashSet = held.iter().map(|r| r.key.clone()).collect(); + let counted = counts(rec); + // 滚动窗口里的请求数:发出时就记上了。记录里算数的那几跳留着,别的撤掉 + let mut left: Vec<&str> = if counted { + hops.iter().map(String::as_str).collect() + } else { + Vec::new() + }; + for r in &held { + match left.iter().position(|h| *h == r.key) { + Some(i) => { + left.swap_remove(i); + } + None => books::unrecord(&mut g.books, r), + } + } + if counted { + let one = Amount { + requests: 1, + ..Default::default() + }; + for up in hops { + add_to_periods( + g.books.entry(up.clone()).or_default(), + clock, + &one, + at_ms, + now, + ); + touched.insert(up.clone()); + } + if let Some(up) = served { + let amount = Amount { + requests: 0, + ..rec.amount() + }; + let rolling = g + .limits + .get(up) + .and_then(|ls| books::longest_window(ls)) + .is_some(); + let b = g.books.entry(up.to_string()).or_default(); + add_to_periods(b, clock, &amount, at_ms, now); + if rolling { + b.recent.push_back((now, amount)); + } + touched.insert(up.to_string()); + } + } + let moved: Vec = touched + .iter() + .flat_map(|k| books::take_moved(&mut g.books, &mut g.asked, k)) + .collect(); + (touched, moved) + }; + self.shared.wake(touched.iter().map(String::as_str)); + self.shared.ask_reread(moved); + } + + /// 重启之后把天、周、月的数从请求记录里加回来。`since(t)` 给出从 `t` 起每家上游记下的 + /// 那些(存储层的 `upstream_usage_since`:请求数是算数的那几跳,用量是它接下的那些)。 + /// 分钟、小时的窗口从空的开始 + pub fn rebuild(&self, mut since: impl FnMut(i64) -> Vec<(String, Recorded)>) { + let now = self.now(); + let mut sums: Vec<(usize, i64, i64, HashMap)> = Vec::new(); + for (i, per) in CALENDAR.into_iter().enumerate() { + let (start, end) = self.shared.clock.period(per, now); + let mut by_up: HashMap = HashMap::new(); + for (up, r) in since(start).iter().filter(|(_, r)| counts(r)) { + by_up.entry(up.clone()).or_default().add(&r.amount()); + } + sums.push((i, start, end, by_up)); + } + let names: Vec = { + let mut g = self.shared.lock(); + for (i, start, end, by_up) in sums { + for b in g.books.values_mut() { + b.periods[i] = Some(Period { + start, + end, + sum: Amount::default(), + }); + } + for (up, sum) in by_up { + g.books.entry(up).or_default().periods[i] = Some(Period { start, end, sum }); + } + } + g.wake.keys().cloned().collect() + }; + self.shared.wake(names.iter().map(String::as_str)); + } + + /// 一家的每条上限此刻用了多少(`ProviderView::limits`)。在路上的按预留算,和 + /// [`Self::state`] 看的是同一个数 + pub fn view(&self, upstream: &str, limits: &[KeyLimit]) -> Vec { + if limits.is_empty() { + return Vec::new(); + } + let now = self.now(); + let clock = self.shared.clock.as_ref(); + let (view, moved, touched) = { + let mut g = self.shared.lock(); + let g = &mut *g; + let touched = self.shared.sweep(g, now); + let view = limits + .iter() + .map(|l| { + let used = books::used(&mut g.books, &g.held, clock, upstream, l, now); + let max = l.max(); + tw_api::KeyLimitView { + per: l.per.into(), + every: l.every_or_one(), + measure: l.measure().into(), + max: max.max(0) as u64, + cache_reads: l.cache_reads, + used: used.max(0) as u64, + resets_at_ms: l + .window_ms() + .is_none() + .then(|| clock.period(l.per, now).1.max(0) as u64), + reached: used >= max, + } + }) + .collect(); + ( + view, + books::take_moved(&mut g.books, &mut g.asked, upstream), + touched, + ) + }; + self.shared.wake(touched.iter().map(String::as_str)); + self.shared.ask_reread(moved); + view + } + + /// 等到这家的 [`Self::state`] 可能变了:有预留还回来、有一行结算了、配置换了、改了名, + /// 或者到了它满着的那条上限空出来(下一期开始)的时刻。**可能没变** —— 醒了再看一眼 + /// `state`,还满着就再等。按此刻配置里这家的上限算时刻 + pub async fn changed(&self, upstream: &str) { + let notify = self + .shared + .lock() + .wake + .entry(upstream.to_string()) + .or_default() + .clone(); + let notified = notify.notified(); + tokio::pin!(notified); + // 先登记再看:看完之后、等之前发生的变化也叫得醒它 + notified.as_mut().enable(); + let limits = self + .shared + .lock() + .limits + .get(upstream) + .cloned() + .unwrap_or_default(); + let until = match self.state(upstream, &limits) { + LimitState::Free => None, + LimitState::WindowFull { until_ms, .. } | LimitState::PeriodUsedUp { until_ms, .. } => { + Some(until_ms) + } + }; + match until { + Some(t) => { + let ms = (t as i64 - self.now()).max(0) as u64; + tokio::select! { + _ = &mut notified => {} + _ = tokio::time::sleep(Duration::from_millis(ms)) => {} + } + } + None => notified.await, + } + } +} + +/// 一笔记进天、周、月:`at_ms`(请求开始的时刻)落在 `now` 所在的那一期里的才记 +fn add_to_periods(b: &mut Books, clock: &dyn Clock, amount: &Amount, at_ms: i64, now: i64) { + for per in CALENDAR { + let p = b.period(clock, per, now); + if at_ms >= p.start && at_ms < p.end { + p.sum.add(amount); + } + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/tw-gateway/src/upstream_limits/tests.rs b/crates/tw-gateway/src/upstream_limits/tests.rs new file mode 100644 index 00000000..9b3df17f --- /dev/null +++ b/crates/tw-gateway/src/upstream_limits/tests.rs @@ -0,0 +1,478 @@ +use super::*; + +use crate::key_limits::TestClock; +use tw_config::LimitPer; + +/// 东八区。2026-10-05 是周一 +const CST: i32 = 8 * 3600; + +fn at(s: &str) -> i64 { + chrono::DateTime::parse_from_rfc3339(s) + .unwrap() + .timestamp_millis() +} + +fn parse(yaml: &str) -> Vec { + serde_yaml_ng::from_str(yaml).unwrap() +} + +struct Bed { + limits: UpstreamLimits, + bus: tw_observe::EventBus, + set: Vec, +} + +/// 上游 `u` 带着这几条上限(`v` 没有上限),时钟从 `now` 起 +fn bed(now: &str, yaml: &str) -> Bed { + let bus = tw_observe::EventBus::new(); + let limits = UpstreamLimits::with_clock(bus.clone(), Arc::new(TestClock::new(at(now), CST))); + let set = parse(yaml); + limits.configure(&config(&set)); + Bed { limits, bus, set } +} + +fn config(set: &[KeyLimit]) -> tw_config::Config { + tw_config::Config { + providers: vec![ + tw_config::Provider { + name: "u".into(), + base_url: "https://u.example".into(), + limits: set.to_vec(), + ..Default::default() + }, + tw_config::Provider { + name: "v".into(), + base_url: "https://v.example".into(), + ..Default::default() + }, + ], + ..Default::default() + } +} + +impl Bed { + fn now(&self) -> i64 { + self.limits.now() + } + + fn state(&self) -> LimitState { + self.limits.state("u", &self.set) + } + + /// 一跳发往 `u`,上游接下了:交给请求号 + fn send(&self, id: u64, ask: Ask) { + self.bus.emit(started(id)); + self.limits.reserve("u", &self.set, &ask).bind(id); + } + + /// 存储层记下了这一行:请求在总线上结束,然后结算。`hops` 是算数的那几跳 + fn done(&self, id: u64, hops: &[&str], served: Option<&str>, rec: Recorded) { + self.bus.emit(finished(id)); + let hops: Vec = hops.iter().map(|h| h.to_string()).collect(); + self.limits.settle(id, self.now(), &hops, served, &rec); + } + + fn used(&self) -> Vec { + self.limits + .view("u", &self.set) + .iter() + .map(|v| v.used) + .collect() + } +} + +fn started(id: u64) -> tw_api::Event { + tw_api::Event::RequestStarted { + id, + client: "k".into(), + client_hint: None, + session: None, + peer: None, + key_masked: None, + route: "default".into(), + rule: "r".into(), + group: None, + rewritten_by: vec![], + provider: "u".into(), + billing: tw_api::Billing::PerToken, + model: "m".into(), + method: "POST".into(), + path: "/v1/messages".into(), + input_estimate: None, + session_log_bytes: None, + at_ms: 0, + } +} + +fn finished(id: u64) -> tw_api::Event { + tw_api::Event::RequestFinished { + id, + model: "m".into(), + status: 200, + sent_bytes: 0, + received_bytes: 0, + duration_ms: 0, + usage: None, + tokens_per_sec: None, + answered_model: None, + } +} + +/// 记下的一行:发往了上游、成功了 +fn row(input: u64, output: u64, cache_read: u64, cost: i64) -> Recorded { + Recorded { + client: "k".into(), + path: "/v1/messages".into(), + local: false, + reached: true, + requests: 1, + input, + output, + cache_read, + cache_write: 0, + cost_micros: cost, + } +} + +fn ask(tokens: u64, cost: i64) -> Ask { + Ask { + tokens, + cost_micros: cost, + } +} + +// ───────────────────────────────────────────────── 滚动窗口 + +#[tokio::test(start_paused = true)] +async fn a_full_window_says_when_it_frees_and_frees_then() { + let b = bed( + "2026-10-05T10:00:00+08:00", + "[{per: minute, every: 5, requests: 2}]", + ); + let t0 = b.now(); + assert_eq!(b.state(), LimitState::Free); + b.send(1, Ask::default()); + tokio::time::advance(Duration::from_secs(60)).await; + b.send(2, Ask::default()); + // 两个都发出去了:窗口满了,第一个在 5 分钟时滑出去 + assert_eq!( + b.state(), + LimitState::WindowFull { + until_ms: (t0 + 300_000) as u64, + limit: 0 + } + ); + // 结算不改请求数:发出去就算 + b.done(1, &["u"], Some("u"), row(10, 10, 0, 0)); + b.done(2, &["u"], Some("u"), row(10, 10, 0, 0)); + assert_eq!(b.used(), [2]); + assert!(matches!(b.state(), LimitState::WindowFull { .. })); + tokio::time::advance(Duration::from_secs(240)).await; + assert_eq!(b.state(), LimitState::Free, "第一个滑出去了"); + assert_eq!(b.used(), [1]); +} + +#[tokio::test(start_paused = true)] +async fn tokens_count_for_the_upstream_that_answered_when_settled() { + let b = bed("2026-10-05T10:00:00+08:00", "[{per: hour, tokens: 1000}]"); + // 估了 600 占着 + b.send(1, ask(600, 0)); + assert_eq!(b.used(), [600]); + assert_eq!(b.state(), LimitState::Free); + // 跑了 50 分钟才结束,用了 900:记在结束的那一刻 + tokio::time::advance(Duration::from_secs(50 * 60)).await; + let settled = b.now(); + b.done(1, &["u"], Some("u"), row(500, 400, 0, 0)); + assert_eq!(b.used(), [900]); + b.send(2, ask(200, 0)); + assert_eq!( + b.state(), + LimitState::WindowFull { + until_ms: (settled + 3_600_000) as u64, + limit: 0 + }, + "900 那一笔滑出去才空" + ); + // 用量记在接下请求的那一家:别的跳只算请求数 + b.done(2, &["v", "u"], Some("v"), row(5_000, 5_000, 0, 0)); + assert_eq!(b.used(), [900]); +} + +// ───────────────────────────────────────────────── 自然周期 + +#[tokio::test(start_paused = true)] +async fn a_used_up_day_is_skipped_until_local_midnight() { + let b = bed( + "2026-10-05T23:59:00+08:00", + "[{per: minute, requests: 1}, {per: day, cost: 1}]", + ); + b.send(1, ask(0, 400_000)); + b.done(1, &["u"], Some("u"), row(1, 1, 0, 1_200_000)); + // 天和分钟都满了:天的先报,等不到 + assert_eq!( + b.state(), + LimitState::PeriodUsedUp { + until_ms: at("2026-10-06T00:00:00+08:00") as u64, + limit: 1 + } + ); + let v = b.limits.view("u", &b.set); + assert_eq!( + (v[1].used, v[1].reached, v[1].resets_at_ms), + ( + 1_200_000, + true, + Some(at("2026-10-06T00:00:00+08:00") as u64) + ) + ); + assert_eq!((v[0].every, v[0].resets_at_ms), (1, None)); + // 过了零点,新的一天从 0 起;分钟那条也滑出去了 + tokio::time::advance(Duration::from_secs(61)).await; + assert_eq!(b.state(), LimitState::Free); + assert_eq!(b.used(), [0, 0]); +} + +#[tokio::test(start_paused = true)] +async fn requests_in_flight_fill_a_period_before_they_are_settled() { + let b = bed("2026-10-05T10:00:00+08:00", "[{per: week, requests: 2}]"); + b.send(1, Ask::default()); + b.send(2, Ask::default()); + assert_eq!( + b.state(), + LimitState::PeriodUsedUp { + until_ms: at("2026-10-12T00:00:00+08:00") as u64, + limit: 0 + }, + "下周一零点" + ); + b.done(1, &["u"], Some("u"), row(1, 1, 0, 0)); + b.done(2, &["u"], Some("u"), row(1, 1, 0, 0)); + assert_eq!(b.used(), [2]); +} + +// ───────────────────────────────────────────────── 预留、还回、结算 + +#[tokio::test(start_paused = true)] +async fn a_hop_answered_429_or_never_sent_gives_everything_back() { + let b = bed( + "2026-10-05T10:00:00+08:00", + "[{per: minute, requests: 1}, {per: day, tokens: 100}]", + ); + let hold = b.limits.reserve("u", &b.set, &ask(80, 0)); + assert_eq!(b.used(), [1, 80]); + assert!(matches!(b.state(), LimitState::WindowFull { .. })); + hold.release(); + assert_eq!(b.used(), [0, 0]); + assert_eq!(b.state(), LimitState::Free); + // 丢掉也一样 + drop(b.limits.reserve("u", &b.set, &ask(80, 0))); + assert_eq!(b.used(), [0, 0]); + // 没设上限的那家什么都不记 + let none = b.limits.reserve("v", &[], &ask(80, 0)); + assert!(format!("{none:?}").contains("false")); +} + +#[tokio::test(start_paused = true)] +async fn the_record_decides_which_hops_count() { + let b = bed( + "2026-10-05T10:00:00+08:00", + "[{per: minute, requests: 5}, {per: day, requests: 5}]", + ); + // 一个请求发了三跳到 u:头一跳回了 429 由管线还回去了,第二跳超时(算),第三跳接下了 + b.bus.emit(started(1)); + b.limits.reserve("u", &b.set, &Ask::default()).release(); + b.limits.reserve("u", &b.set, &Ask::default()).bind(1); + b.limits.reserve("u", &b.set, &Ask::default()).bind(1); + assert_eq!(b.used(), [2, 2]); + b.done(1, &["u", "u"], Some("u"), row(1, 1, 0, 0)); + assert_eq!(b.used(), [2, 2], "两跳各算一个"); + // 记录里没有的那一跳把窗口里记上的也还回去 + b.send(2, Ask::default()); + b.send(2, Ask::default()); + b.done(2, &["u"], Some("u"), row(1, 1, 0, 0)); + assert_eq!(b.used(), [3, 3]); + // 数 token 的请求、网关自己答的不算 + b.send(3, Ask::default()); + let mut count = row(0, 0, 0, 0); + count.path = "/v1/messages/count_tokens".into(); + b.done(3, &["u"], Some("u"), count); + assert_eq!(b.used(), [3, 3]); +} + +#[tokio::test(start_paused = true)] +async fn a_reservation_whose_row_never_comes_is_let_go_after_its_request_ended() { + let b = bed("2026-10-05T10:00:00+08:00", "[{per: day, tokens: 100}]"); + b.send(1, ask(80, 0)); + b.bus.emit(finished(1)); + assert_eq!(b.used(), [80], "刚结束:再等存储层一会儿"); + tokio::time::advance(Duration::from_millis(books::GRACE_MS as u64)).await; + assert_eq!(b.used(), [0]); +} + +#[tokio::test(start_paused = true)] +async fn cache_reads_count_only_when_asked_and_cost_is_what_was_recorded() { + let b = bed( + "2026-10-05T10:00:00+08:00", + "[{per: day, tokens: 1000000}, {per: day, tokens: 1000000, cache_reads: true}, \ + {per: hour, every: 2, cost: 5}]", + ); + b.send(1, ask(0, 0)); + b.done(1, &["u"], Some("u"), row(100, 50, 10_000, 30_000)); + assert_eq!(b.used(), [150, 10_150, 30_000]); +} + +// ───────────────────────────────────────────────── 重启、换期、改名 + +#[tokio::test(start_paused = true)] +async fn a_restart_adds_the_periods_back_from_the_store() { + let b = bed( + "2026-10-07T12:00:00+08:00", + "[{per: day, requests: 10}, {per: week, tokens: 10000}, {per: month, cost: 50}]", + ); + let day = at("2026-10-07T00:00:00+08:00"); + let week = at("2026-10-05T00:00:00+08:00"); + let month = at("2026-10-01T00:00:00+08:00"); + let mut asked = Vec::new(); + b.limits.rebuild(|since| { + asked.push(since); + // 今天:u 3 跳、用量 100;这周早些时候再加 2 跳、用量 200;这个月再早些再加 1 跳、 + // $2。v 的不算在 u 头上;数 token 的不算 + let mut rows = vec![( + "u".to_string(), + Recorded { + requests: 3, + ..row(50, 50, 0, 1_000_000) + }, + )]; + if since <= week { + rows.push(( + "u".to_string(), + Recorded { + requests: 2, + ..row(100, 100, 0, 0) + }, + )); + } + if since <= month { + rows.push(("u".to_string(), row(0, 0, 0, 2_000_000))); + } + rows.push(("v".to_string(), row(9_000, 9_000, 0, 9_000_000))); + let mut count = row(5_000, 0, 0, 0); + count.path = "/v1/messages/count_tokens".into(); + rows.push(("u".to_string(), count)); + rows + }); + assert_eq!(asked, [day, week, month]); + assert_eq!(b.used(), [3, 300, 3_000_000]); + assert_eq!(b.state(), LimitState::Free); + // 加回来的之上照常记 + b.send(1, Ask::default()); + b.done(1, &["u"], Some("u"), row(10, 10, 0, 0)); + assert_eq!(b.used(), [4, 320, 3_000_000]); +} + +#[tokio::test(start_paused = true)] +async fn a_period_whose_start_moved_is_read_back_once() { + let b = bed("2026-10-05T23:59:30+08:00", "[{per: day, requests: 5}]"); + let asked: Arc>> = Arc::default(); + let log = asked.clone(); + b.limits + .reread_with(Arc::new(move |m| log.lock().unwrap().push(m))); + b.send(1, Ask::default()); + b.done(1, &["u"], Some("u"), row(1, 1, 0, 0)); + assert_eq!(b.used(), [1]); + // 过了零点:换新的一天,要一次读回 + tokio::time::advance(Duration::from_secs(60)).await; + assert_eq!(b.used(), [0]); + assert_eq!(b.used(), [0]); + let m = { + let got = asked.lock().unwrap(); + assert_eq!(got.len(), 1, "一期只要一次"); + got[0].clone() + }; + assert_eq!( + (m.key.as_str(), m.per, m.start), + ("u", LimitPer::Day, at("2026-10-06T00:00:00+08:00")) + ); + b.limits.reread( + &m, + &[ + ( + "u".to_string(), + Recorded { + requests: 4, + ..row(1, 1, 0, 0) + }, + ), + ("v".to_string(), row(1, 1, 0, 0)), + ], + ); + assert_eq!(b.used(), [4]); +} + +#[tokio::test(start_paused = true)] +async fn a_renamed_upstream_keeps_what_it_used() { + let b = bed("2026-10-05T10:00:00+08:00", "[{per: day, requests: 3}]"); + b.send(1, Ask::default()); + b.done(1, &["u"], Some("u"), row(1, 1, 0, 0)); + b.send(2, Ask::default()); + b.limits.rename("u", "u2"); + assert_eq!(b.limits.view("u2", &b.set)[0].used, 2); + assert_eq!(b.limits.view("u", &b.set)[0].used, 0); +} + +// ───────────────────────────────────────────────── 等它变化 + +/// 叫没叫醒:给 `changed` 一点时间,看它返回了没有 +async fn woke(f: impl std::future::Future, within: Duration) -> bool { + tokio::time::timeout(within, f).await.is_ok() +} + +#[tokio::test(start_paused = true)] +async fn changed_wakes_when_a_reservation_comes_back() { + let b = bed("2026-10-05T10:00:00+08:00", "[{per: day, requests: 1}]"); + let hold = b.limits.reserve("u", &b.set, &Ask::default()); + assert!(matches!(b.state(), LimitState::PeriodUsedUp { .. })); + let limits = &b.limits; + let wait = limits.changed("u"); + tokio::pin!(wait); + // 还没变:一秒之内不醒 + assert!(!woke(wait.as_mut(), Duration::from_secs(1)).await); + hold.release(); + assert!(woke(wait, Duration::from_millis(1)).await); + assert_eq!(b.state(), LimitState::Free); +} + +#[tokio::test(start_paused = true)] +async fn changed_wakes_when_the_window_slides() { + let b = bed("2026-10-05T10:00:00+08:00", "[{per: minute, requests: 1}]"); + b.send(1, Ask::default()); + let t0 = b.now(); + tokio::time::advance(Duration::from_secs(20)).await; + // 40 秒之后滑出去:到点就醒,不用谁叫 + b.limits.changed("u").await; + assert_eq!(b.now(), t0 + 60_000); + assert_eq!(b.state(), LimitState::Free); +} + +#[tokio::test(start_paused = true)] +async fn changed_wakes_when_a_settle_or_new_limits_arrive() { + let b = bed("2026-10-05T10:00:00+08:00", "[{per: hour, tokens: 100}]"); + b.send(1, ask(500, 0)); + assert!(matches!(b.state(), LimitState::WindowFull { .. })); + let wait = b.limits.changed("u"); + tokio::pin!(wait); + assert!(!woke(wait.as_mut(), Duration::from_secs(1)).await); + // 结算:估多了,实际只用了 10 + b.done(1, &["u"], Some("u"), row(5, 5, 0, 0)); + assert!(woke(wait, Duration::from_millis(1)).await); + assert_eq!(b.state(), LimitState::Free); + // 配置换了也叫醒:上限可能松了 + b.send(2, ask(500, 0)); + let wait = b.limits.changed("u"); + tokio::pin!(wait); + assert!(!woke(wait.as_mut(), Duration::from_secs(1)).await); + b.limits + .configure(&config(&parse("[{per: hour, tokens: 100000}]"))); + assert!(woke(wait, Duration::from_millis(1)).await); +} diff --git a/crates/tw-gateway/src/ws.rs b/crates/tw-gateway/src/ws.rs index 85924f71..cec8401c 100644 --- a/crates/tw-gateway/src/ws.rs +++ b/crates/tw-gateway/src/ws.rs @@ -847,8 +847,9 @@ pub(crate) async fn proxy( } }; // 这一家接没接下这条连接,和 HTTP 那条路一跳的成败记在同一笔账上(见 `crate::health`): - // 连不上、回了 5xx 是失败,凭据被拒、限流按原因停用,请求本身的问题算这一家答上了。每一轮 - // 一行的连接接下了不记 —— 它的每一轮各记各的(见 `turn`) + // 连不上、回了 5xx 是失败,凭据被拒按原因停用,请求本身的问题算这一家答上了;回 429 的不是 + // 失败,这一家暂停(见 `crate::throttle`)。每一轮一行的连接接下了不记 —— 它的每一轮各记 + // 各的(见 `turn`) let change = match &connected { Ok(_) if turns.is_some() => None, Ok(_) => state.health.record_success(name), @@ -860,6 +861,10 @@ pub(crate) async fn proxy( &[], crate::server::now_ms(), ) { + crate::failure::Verdict::Failed(crate::failure::Cause::RateLimited { retry_after }) => { + state.throttles.limited(name, hop_started, retry_after); + None + } crate::failure::Verdict::Failed(cause) => state.health.record_cause(name, cause), crate::failure::Verdict::ClientError => state.health.record_success(name), }, diff --git a/crates/tw-gateway/src/ws/turn.rs b/crates/tw-gateway/src/ws/turn.rs index e493abc5..3345e6be 100644 --- a/crates/tw-gateway/src/ws/turn.rs +++ b/crates/tw-gateway/src/ws/turn.rs @@ -14,10 +14,11 @@ //! (见 `server::upgrade`)。 //! //! **上限和并发按轮算,和 HTTP 那条路的一个请求同一套**([`admit`]):天、周、月用满了就拒, -//! 密钥的并发上限等前面的结束,分钟、小时等得到就等;然后占这一家的一个位置 -//! (`max_concurrent`)—— 这条连接只连着这一家,没有下一家可换:满着就等,等不到回 -//! `response.failed`。等分钟、小时的空位和等这一家的位置**共用这一轮的一个期限**,和 HTTP -//! 那条路一样。被拒的这一轮替它回一个 `response.failed`,连接照常。占着的(密钥的 +//! 密钥的并发上限等前面的结束,分钟、小时等得到就等;然后等这一家能发 —— 这条连接只连着这一家, +//! 没有下一家可换:满着(`max_concurrent`)、限流暂停着、分钟或小时的用量上限满了就等它(见 +//! [`crate::throttle`]),排在它的队里,等不到、这一期用完了回 `response.failed`。等分钟、小时的 +//! 空位和等这一家**共用这一轮的一个期限**,和 HTTP 那条路一样。上游在这一轮里回了限流的错误, +//! 这一家暂停,错误照常交给客户端:这一帧已经发出去了,换不了。被拒的这一轮替它回一个 `response.failed`,连接照常。占着的(密钥的 //! 并发通行证、这一家的位置)从发出去占到这一次回答完、或者连接断了;**闲着的连接什么都 //! 不占**。 //! @@ -38,7 +39,8 @@ use crate::ending::Ending; use crate::error::GatewayError; use crate::server::Choice; use crate::state::AppState; -use tw_types::{Msg, msg}; +use crate::throttle::Stuck; +use tw_types::Msg; /// 一条连接上开始一个请求时都一样的那几项:谁发的、从哪儿来、哪条路径。整条连接一行的 /// (Realtime 和别的路径)和每一轮一行的(Responses)都由它开始。 @@ -266,6 +268,15 @@ pub(crate) struct Turn { /// 密钥的并发通行证和这一家的位置:**跟着这一轮走**,回答完了、连接断了就还 _pass: crate::limits::Pass, _slot: crate::slots::Slot, + /// 暂停结束后探路的就是这一轮的话,探路的名额(见 [`crate::throttle::Probe`])。发出去时说一声 + probe: Option, + /// 这一帧发出去的时刻:这一家的答复是不是最近一次 429 之后的,看它 + sent_at: Instant, + /// 这一轮在这一家用量上限上占着的(见 [`crate::upstream_limits`]):发出去了的这一轮**怎么收场 + /// 都交给这一轮的号**等结算(上游收到了它),被限流了、没发出去的还回去 + usage: Option, + /// 上游在这一轮里报了限流,这一家已经暂停了。**一轮只暂停一次** + paused: bool, /// 这一轮结束时记一笔,到了时候重读这一家的余额(见 [`crate::balances`]) _passing: crate::balances::Passing, } @@ -278,6 +289,10 @@ impl Turn { r.attempt.model = model; r.since = Instant::now(); } + self.sent_at = Instant::now(); + if let Some(p) = self.probe.as_mut() { + p.sent(); + } } /// 这一轮发给上游的一帧,载荷 `bytes` 字节:`response.create` 那一帧,和它在答的时候客户端 @@ -309,6 +324,14 @@ impl Turn { self.response = response.map(str::to_string); } self.errored |= kind == Some("error"); + // 收尾的这一帧报的错。和 HTTP 那条路开头报错的一跳同一个判据 + let dialect = tw_dialect::ir::Dialect::Responses; + let fault = super::ends_turn(kind) + .then(|| crate::server::stream_fault(&self.state, &self.provider, dialect, text)) + .flatten(); + let limited = matches!(fault, Some(crate::failure::Cause::RateLimited { .. })); + // **在报路由事件之前**:这一轮在内容之前就被限流了的,尝试链上那一跳要带着暂停多久 + self.pause_on(fault); if !self.responded { self.responded = true; self.routed(); @@ -327,16 +350,59 @@ impl Turn { if self.judged { return; } + // 这一家在不在限流(见 `crate::throttle`)也在这里说,**不在头一帧**:`response.created` + // 之后接着报限流的上游,头一帧不说明它好了 —— 和 HTTP 那条路压着开头的一跳一样 if self.ending.as_ref().is_some_and(Ending::has_content) { - // 内容到了:这一家答上了。之后流里再报错也不改 —— HTTP 那条路的一跳也是开头 - // 一过就记成功 + // 内容到了:这一家答上了,没在限流(探路的过了,连续限流的次数归零)。之后流里再报错 + // 也不改 —— HTTP 那条路的一跳也是开头一过就记成功 + self.state.throttles.answered(&self.provider, self.sent_at); self.judge(|h, p| h.record_success(p)); } else if super::ends_turn(kind) { - // 一段内容都没有就收了尾:和 HTTP 那条路开头报错的一跳同一个判据 - let dialect = tw_dialect::ir::Dialect::Responses; - match crate::server::stream_fault(&self.state, &self.provider, dialect, text) { - Some(cause) => self.judge(|h, p| h.record_cause(p, cause)), - None => self.judge(|h, p| h.record_success(p)), + // 一段内容都没有就收了尾。限流不是失败(上面暂停过了),不记;额度用完、没钱了等不到, + // 探路的不再探,交给熔断 + match fault { + Some(_) if limited => self.judged = true, + Some(cause) => { + if matches!( + cause, + crate::failure::Cause::QuotaUsedUp { .. } + | crate::failure::Cause::NoBalance + ) { + self.state.throttles.released(&self.provider, self.sent_at); + } else { + self.state.throttles.answered(&self.provider, self.sent_at); + } + self.judge(|h, p| h.record_cause(p, cause)) + } + None => { + self.state.throttles.answered(&self.provider, self.sent_at); + self.judge(|h, p| h.record_success(p)) + } + } + } + } + + /// 上游在这一轮里报了限流(`fault`):这一家暂停。**错误照常交给客户端**:这一帧已经发出去 + /// 了。答到一半才报的也暂停。 + /// + /// 这一轮在用量上限上占着的,**还没有内容、路由事件也还没报的才还回去**,尝试链上那一跳带着 + /// `pause_ms`:记录里不算它(`tw_api::AttemptView::counts_for_upstream`),内存里也不算,两边 + /// 是同一个数。已经有了内容的(上游答了)、或者路由事件已经报出去了(`response.created` 之后 + /// 才报的)照常算一个请求 —— 记录里那一跳就是这样,这里跟着它 + fn pause_on(&mut self, fault: Option) { + if let Some(crate::failure::Cause::RateLimited { retry_after }) = fault + && !std::mem::replace(&mut self.paused, true) + { + let pause = self + .state + .throttles + .limited(&self.provider, self.sent_at, retry_after); + let fresh = !self.ending.as_ref().is_some_and(Ending::has_content); + if fresh && let Some(r) = self.route.as_mut() { + r.attempt.pause_ms = Some(pause.as_millis() as u64); + if let Some(h) = self.usage.take() { + h.release(); + } } } } @@ -355,6 +421,15 @@ impl Turn { crate::server::note_health(&self.state.bus, health, &self.provider, change); } + /// 发出去了的这一轮在这一家的用量上限上算一个请求:占着的交给这一轮的号,等存储层结算。 + /// **在报结局之前**:存储层先收到结局、先结算的话,这一份没交给号、没人认领,要在内存里多占 + /// 一阵(见 [`crate::upstream_limits`]) + fn bind_usage(&mut self) { + if let Some(h) = self.usage.take() { + h.bind(self.id); + } + } + fn routed(&mut self) { let Some(r) = self.route.take() else { return }; let mut attempt = r.attempt; @@ -368,6 +443,7 @@ impl Turn { /// 那句(见 [`Ending::streaming`]) fn finish(mut self) { self.routed(); + self.bind_usage(); let Some(e) = self.ending.take() else { return }; match self.cut.take() { Some(why) => e.failed(tw_api::FailureSource::Denied, why), @@ -380,8 +456,13 @@ impl Turn { pub(crate) fn fail(mut self, source: tw_api::FailureSource, why: Msg) { if source == tw_api::FailureSource::Upstream { self.judge(|h, p| h.record_failure(p)); + // 没有结论:探路的不再探,交给熔断(见 `crate::throttle::Throttles::released`) + if !self.ending.as_ref().is_some_and(Ending::has_content) { + self.state.throttles.released(&self.provider, self.sent_at); + } } self.routed(); + self.bind_usage(); if let Some(e) = self.ending.take() { e.failed(source, why); } @@ -399,6 +480,17 @@ impl Turn { if source == tw_api::FailureSource::Upstream { self.judge(|h, p| h.record_failure(p)); } + // 没发出去:用量上限上占着的还回去,探路的名额也是 + if let Some(h) = self.usage.take() { + h.release(); + } + if let Some(p) = self.probe.take() { + p.withdraw(); + } + // 写不过去是上游断了:没有结论,探路的不再探,交给熔断 + if source == tw_api::FailureSource::Upstream { + self.state.throttles.released(&self.provider, self.sent_at); + } if let Some(r) = self.route.take() { self.state.bus.emit(routed( self.id, @@ -418,6 +510,9 @@ impl Drop for Turn { /// 这一行没有尝试链。结局随后由它自己的 Drop 报成取消 fn drop(&mut self) { self.routed(); + // 发出去了的这一轮(连接断了):在这一家的用量上限上算一个请求。结局(取消)由它自己的 + // Drop 随后报,交给号在那之前 + self.bind_usage(); } } @@ -481,10 +576,11 @@ pub(crate) enum NotAdmitted { /// 3. 分钟、小时的上限 —— 下一个空位在这一轮的等待期限之前空出来就等,等不到就拒;过了 /// 就按输入的估算占着,存储层记下这一行时换成实数; /// 4. 开始:发开始事件,内容过滤的结论挂在这一轮上报(拒绝的切断连接); -/// 5. 这一家的位置(`max_concurrent`)—— 满着就等,等到同一个期限;这条连接只连着 -/// 这一家,等不到就是这一家忙,回 429 那句话。 +/// 5. 这一家能发 —— 满着(`max_concurrent`)、限流暂停着、分钟或小时的用量上限满了就等, +/// 排在它的队里,等到同一个期限(见 [`crate::throttle`]);这条连接只连着这一家,等不到、 +/// 这一期用完了,回 429 那句话(和 HTTP 那条路一家都发不出去时同一句)。 /// -/// **等待期限一轮只有一个**:过了密钥的并发上限那一刻起算 `failover.slot_wait_secs`,第 3 步、 +/// **等待期限一轮只有一个**:过了密钥的并发上限那一刻起算 `failover.wait_secs`,第 3 步、 /// 第 5 步的等待都算在里面,和 HTTP 那条路一样(见 `server::pipeline::admission`)。 /// /// 被上限拒的、等不到位置的**照样留一行**,流量里看得见它为什么没发出去。拿到的通行证和 @@ -506,8 +602,8 @@ pub(crate) async fn admit(a: Admit) -> Result { } let pass = state.gate.acquire(key_name, max_concurrent).await; // 等待期限从这里起算:等滚动窗口的空位、等这一家的位置共用它(见 - // `crate::key_limits::slot_wait`) - let wait_until = tokio::time::Instant::now() + crate::key_limits::slot_wait(&rt.config); + // `crate::key_limits::request_wait`) + let wait_until = tokio::time::Instant::now() + crate::key_limits::request_wait(&rt.config); let ask = ask(state, &a, limits); let hold = match state .key_limits @@ -542,37 +638,128 @@ pub(crate) async fn admit(a: Admit) -> Result { } let model = Some(a.sent.clone()).filter(|m| !m.is_empty() && *m != a.requested); let hop_started = Instant::now(); - let (slot, queued_ms) = match state.slots.try_take(&line.provider) { - Some(slot) => (slot, None), - None => { + let provider = rt.config.providers.iter().find(|p| p.name == line.provider); + let limits = provider + .map(crate::throttle::usage_limits) + .unwrap_or_default(); + let deadline_ms = crate::server::now_ms() + + wait_until + .saturating_duration_since(tokio::time::Instant::now()) + .as_millis() as u64; + let mut waiter = state.throttles.waiter(); + let at_once = match state + .throttles + .gate(&line.provider, limits, Some(waiter.id)) + { + crate::throttle::Gate::Go(probe) => match state.slots.try_take(&line.provider) { + Some(slot) => Ok((slot, probe)), + None => Err(Stuck::Busy), + }, + crate::throttle::Gate::Wait { kind, until_ms } => Err(match kind { + tw_api::ThrottleKind::WindowFull => Stuck::WindowFull(until_ms), + _ => Stuck::RateLimited(until_ms), + }), + crate::throttle::Gate::UsedUp { until_ms, .. } => Err(Stuck::UsedUp(until_ms)), + }; + let (slot, probe, queued_ms) = match at_once { + Ok((slot, probe)) => (slot, probe, None), + Err(stuck) => { let mut waited = None; - let mut slot = None; - // 这一轮能等的已经等完了(或者配置的是不等):不再等 - if tokio::time::Instant::now() < wait_until { - slot = state.slots.take_by(&line.provider, wait_until).await; + let mut ready = None; + // 等得到才等:这一期用完了的、说得出什么时候恢复又晚于期限的不等。这一轮能等的已经 + // 等完了(或者配置的是不等)也不再等 + let hopeful = match stuck { + Stuck::Busy => true, + Stuck::RateLimited(t) | Stuck::WindowFull(t) => t.is_none_or(|t| t <= deadline_ms), + Stuck::UsedUp(_) => false, + }; + if hopeful && tokio::time::Instant::now() < wait_until { + let one = [(line.provider.as_str(), limits)]; + state.bus.emit(tw_api::Event::RequestQueued { + id, + upstreams: vec![line.provider.clone()], + until_ms: deadline_ms, + at_ms: crate::server::now_ms(), + }); + ready = + Some(crate::throttle::first_ready(state, &mut waiter, &one, wait_until).await); + state.bus.emit(tw_api::Event::RequestDequeued { + id, + at_ms: crate::server::now_ms(), + }); waited = Some(hop_started.elapsed().as_millis() as u64); } - match slot { - Some(slot) => (slot, waited), - None => { - let limit = state.slots.limit(&line.provider).unwrap_or_default(); - let busy = - crate::server::hop_busy(&line.provider, model, limit, waited, hop_started); + match ready { + Some(crate::throttle::Outcome::Ready(_, ready)) => { + (ready.slot, ready.probe, waited) + } + _ => { + let now = crate::server::now_ms(); + // 此刻它是哪一样:等的时候可能换了 + let stuck = match (stuck, state.throttles.peek(&line.provider, limits)) { + (Stuck::Busy, _) => Stuck::Busy, + (_, crate::throttle::Peek::UsedUp { until_ms, .. }) => { + Stuck::UsedUp(until_ms) + } + ( + _, + crate::throttle::Peek::Wait { + kind: tw_api::ThrottleKind::WindowFull, + until_ms, + }, + ) => Stuck::WindowFull(until_ms), + (_, crate::throttle::Peek::Wait { until_ms, .. }) => { + Stuck::RateLimited(until_ms) + } + (stuck, crate::throttle::Peek::Open) => stuck, + }; + let hop = match stuck { + Stuck::Busy => crate::server::hop_busy( + &line.provider, + model, + state.slots.limit(&line.provider).unwrap_or_default(), + waited, + hop_started, + ), + Stuck::RateLimited(t) => crate::server::hop_waiting( + &line.provider, + model, + tw_api::ThrottleKind::RateLimited, + t, + waited, + hop_started, + ), + Stuck::WindowFull(t) => crate::server::hop_waiting( + &line.provider, + model, + tw_api::ThrottleKind::WindowFull, + t, + waited, + hop_started, + ), + Stuck::UsedUp(t) => { + crate::server::hop_used_up(&line.provider, model, t, hop_started) + } + }; state .bus - .emit(routed(id, choice, vec![busy], tw_api::Billing::PerToken)); - // 和 HTTP 那条路候选全满时同一句:能服务它的就这一家 - let err = GatewayError::busy(msg!( - "gw.busy_all", upstreams = format!("`{}`", line.provider) => - "Every upstream that can serve this request is at its concurrency limit \ - (max_concurrent): {upstreams}. None had a free slot in time; try again shortly." - )); + .emit(routed(id, choice, vec![hop], tw_api::Billing::PerToken)); + // 和 HTTP 那条路一家都发不出去时同一句:能服务它的就这一家 + let err = crate::throttle::refusal(&[(line.provider.clone(), stuck)], now); ending.failed(err.source.into(), err.detail.clone()); return Err(NotAdmitted::Failed(err)); } } } }; + // 这一家的用量上限:请求数此刻记上,token 和费用按这一帧估一个数占着(见 + // `crate::upstream_limits`) + let usage = (!limits.is_empty()).then(|| { + let asked = self::ask(state, &a, limits); + state + .upstream_limits + .reserve(&line.provider, limits, &asked) + }); let mut attempt = crate::server::hop( &line.provider, model, @@ -600,6 +787,10 @@ pub(crate) async fn admit(a: Admit) -> Result { cut: None, _pass: pass, _slot: slot, + probe, + sent_at: Instant::now(), + usage, + paused: false, _passing: state.balance_passing(&line.provider), }) } diff --git a/crates/tw-gateway/tests/endings.rs b/crates/tw-gateway/tests/endings.rs index b3b63208..b49a4a59 100644 --- a/crates/tw-gateway/tests/endings.rs +++ b/crates/tw-gateway/tests/endings.rs @@ -597,9 +597,14 @@ async fn a_phase_two_deny_after_the_request_started_is_reported_as_denied() { /// 每一家上游都没接住。**只报一次失败,`source` 是它真正的原因** —— 被限流 /// 就是 `rate_limited`,和客户端收到的 `x-thinkwatch-error` 是同一个词。 +/// +/// 限流的上游请求会等它恢复(见 `tw_gateway::throttle`):这里不让等(`wait_secs: 0`), +/// 被限流的那一次就是结局 #[tokio::test] async fn a_request_every_upstream_refused_is_failed_once_for_the_reason_it_was_refused() { - let (gw, mut events) = serve(cfg(provider(refusing_upstream(429).await))).await; + let mut c = cfg(provider(refusing_upstream(429).await)); + c.failover.wait_secs = 0; + let (gw, mut events) = serve(c).await; let r = post(gw).send().await.unwrap(); assert_eq!(r.status(), 429); let header = r.headers().get("x-thinkwatch-error").cloned(); diff --git a/crates/tw-gateway/tests/idle_timeout.rs b/crates/tw-gateway/tests/idle_timeout.rs index df1143d4..a8c0b3d4 100644 --- a/crates/tw-gateway/tests/idle_timeout.rs +++ b/crates/tw-gateway/tests/idle_timeout.rs @@ -166,22 +166,22 @@ fn provider(name: &str, up: &Upstream, protocol: Protocol) -> Provider { const HOLD: Duration = Duration::from_millis(300); const KEEPALIVE_EVERY: Duration = Duration::from_millis(150); -/// 起网关,交回数据面的状态。`slot_wait_secs` 是等空位的期限 +/// 起网关,交回数据面的状态。`wait_secs` 是等空位的期限 async fn gateway_with( providers: Vec, - slot_wait_secs: u64, + wait_secs: u64, ) -> ( SocketAddr, tokio::sync::broadcast::Receiver, tw_gateway::AppState, ) { - gateway_held(providers, slot_wait_secs, Duration::from_secs(10)).await + gateway_held(providers, wait_secs, Duration::from_secs(10)).await } /// 起网关,流式回答最多压 `hold` 不给响应头 async fn gateway_held( providers: Vec, - slot_wait_secs: u64, + wait_secs: u64, hold: Duration, ) -> ( SocketAddr, @@ -193,14 +193,14 @@ async fn gateway_held( key: "tw-k".into(), ..Default::default() }; - gateway_as(client, providers, slot_wait_secs, hold).await + gateway_as(client, providers, wait_secs, hold).await } /// 起网关,密钥是 `client`,流式回答最多压 `hold` 不给响应头 async fn gateway_as( client: Client, providers: Vec, - slot_wait_secs: u64, + wait_secs: u64, hold: Duration, ) -> ( SocketAddr, @@ -213,7 +213,7 @@ async fn gateway_as( providers, failover: Failover { idle_timeout_secs: 30, - slot_wait_secs, + wait_secs, ..Default::default() }, ..Default::default() diff --git a/crates/tw-gateway/tests/key_limits.rs b/crates/tw-gateway/tests/key_limits.rs index d4608330..d7c78e00 100644 --- a/crates/tw-gateway/tests/key_limits.rs +++ b/crates/tw-gateway/tests/key_limits.rs @@ -192,18 +192,18 @@ async fn a_refused_request_is_recorded_like_other_refusals() { ))); } -/// 一分钟一个:第二个等不到空位(要等快一分钟,超过 30 秒),拒,并说准多久之后再来, -/// 可以重试。 +/// 一小时一个:第二个等不到空位(要等快一小时,超过等待期限 `wait_secs` 的 120 秒),拒,并说 +/// 准多久之后再来,可以重试。 #[tokio::test] async fn a_rolling_limit_says_exactly_when_to_retry() { let (up, _hits) = upstream().await; - let (gw, _rx) = gateway(up, "[{per: minute, requests: 1}]").await; + let (gw, _rx) = gateway(up, "[{per: hour, requests: 1}]").await; assert_eq!(anthropic(gw, "/v1/messages").await.status(), 200); let r = anthropic(gw, "/v1/messages").await; assert_eq!(r.status(), 429); assert_eq!(head(&r, "x-should-retry").as_deref(), Some("true")); let ms: u64 = head(&r, "retry-after-ms").unwrap().parse().unwrap(); - assert!((50_000..=60_000).contains(&ms), "{ms}"); + assert!((3_500_000..=3_600_000).contains(&ms), "{ms}"); let secs: u64 = head(&r, "retry-after").unwrap().parse().unwrap(); assert_eq!(secs, ms.div_ceil(1000)); let v: serde_json::Value = r.json().await.unwrap(); @@ -211,7 +211,7 @@ async fn a_rolling_limit_says_exactly_when_to_retry() { v["error"]["message"] .as_str() .unwrap() - .contains("1 requests per minute"), + .contains("1 requests per hour"), "{v}" ); } diff --git a/crates/tw-gateway/tests/success_rate.rs b/crates/tw-gateway/tests/success_rate.rs index 9b9f60fe..0ef531c8 100644 --- a/crates/tw-gateway/tests/success_rate.rs +++ b/crates/tw-gateway/tests/success_rate.rs @@ -62,8 +62,10 @@ async fn gateway(providers: Vec) -> (SocketAddr, Arc) { ..Default::default() }], providers, + // 不等:一家回 429 之后暂停着,请求当场以 429 收场,不在这里干等它 failover: Failover { failures_to_pause: 100, + wait_secs: 0, ..Default::default() }, ..Default::default() @@ -116,9 +118,11 @@ async fn an_upstream_that_answers_counts_as_a_success() { } #[tokio::test] -async fn server_errors_and_rate_limits_count_against_the_upstream() { +async fn server_errors_count_against_the_upstream_and_rate_limits_do_not() { assert_eq!(after_five(503, "{}").await, Some(0.0)); - assert_eq!(after_five(429, "{}").await, Some(0.0)); + // 限流不是失败:这一家暂停,请求等它(见 `tw_gateway::throttle`),不算进成功率。等不到的 + // 那几个请求放弃了,也不算它答上了 + assert_eq!(after_five(429, "{}").await, None); // 凭据被拒:换一家有意义,那边是另一把密钥 assert_eq!(after_five(401, "{}").await, Some(0.0)); } diff --git a/crates/tw-gateway/tests/throttle_queue.rs b/crates/tw-gateway/tests/throttle_queue.rs new file mode 100644 index 00000000..33caeb8f --- /dev/null +++ b/crates/tw-gateway/tests/throttle_queue.rs @@ -0,0 +1,912 @@ +//! 上游回 429 之后:这一家暂停,请求排队等它(`tw_gateway::throttle`),端到端。 +//! +//! 假上游按脚本一个个地答:先回几个 429(带着暂停多久),之后照常答。它记下每个请求到的时刻 +//! —— 暂停结束时只放一个过去探路、探路的有了结论别的才跟上,看的就是这些时刻。暂停写成几百 +//! 毫秒(`retry-after-ms`),测试不用干等几秒的退避。 + +use std::collections::VecDeque; +use std::net::SocketAddr; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use axum::Router; +use futures::StreamExt; +use tw_api::{AttemptOutcome, AttemptView, Event, ServeSkip, Stay, ThrottleKind}; + +/// 假上游的一次答复 +#[derive(Clone)] +struct Reply { + status: u16, + headers: Vec<(&'static str, String)>, + body: String, + /// 响应头之前等多久 + delay_ms: u64, + /// 流式的回答:响应头之后一段一段地给,每一段之前等多久。有的话 `body` 不用 + chunks: Vec<(u64, String)>, +} + +/// 照常答的那一种:读了 5000 token 的缓存(这段对话跨轮值得留下) +fn ok() -> Reply { + Reply { + status: 200, + headers: Vec::new(), + body: r#"{"type":"message","role":"assistant","model":"claude-sonnet-5","content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn","usage":{"input_tokens":3,"cache_read_input_tokens":5000,"output_tokens":1}}"#.into(), + delay_ms: 0, + chunks: Vec::new(), + } +} + +/// 限流:Anthropic 的样子,`headers` 说暂停多久 +fn limited(headers: &[(&'static str, &str)]) -> Reply { + Reply { + status: 429, + headers: headers.iter().map(|(k, v)| (*k, v.to_string())).collect(), + body: r#"{"type":"error","error":{"type":"rate_limit_error","message":"Number of request tokens has exceeded your per-minute rate limit"}}"#.into(), + delay_ms: 0, + chunks: Vec::new(), + } +} + +/// Anthropic 流式回答里的一个事件 +fn event(name: &str, data: &str) -> String { + format!("event: {name}\ndata: {data}\n\n") +} + +fn message_start() -> String { + event( + "message_start", + r#"{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-5","content":[],"stop_reason":null,"usage":{"input_tokens":3,"output_tokens":0}}}"#, + ) +} + +/// 流式地照常答:开头、一段内容「ok」、收尾 +fn streamed() -> Reply { + let chunks = [ + message_start(), + event( + "content_block_start", + r#"{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#, + ), + event( + "content_block_delta", + r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}"#, + ), + event( + "content_block_stop", + r#"{"type":"content_block_stop","index":0}"#, + ), + event( + "message_delta", + r#"{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}}"#, + ), + event("message_stop", r#"{"type":"message_stop"}"#), + ]; + Reply { + status: 200, + headers: Vec::new(), + body: String::new(), + delay_ms: 0, + chunks: chunks.into_iter().map(|c| (0, c)).collect(), + } +} + +/// 响应头是 200、流也开了,开头之后 `after_ms` 在流里报限流(一段内容都还没有),说 1 秒之后 +/// 再来 +fn limited_in_the_stream(after_ms: u64) -> Reply { + Reply { + chunks: vec![ + (0, message_start()), + ( + after_ms, + event( + "error", + r#"{"type":"error","error":{"type":"rate_limit_error","message":"slow down","details":[{"retryDelay":"1s"}]}}"#, + ), + ), + ], + ..streamed() + } +} + +/// 一家假上游:先按 `script` 的顺序答,答完了一律答 `then` +struct Up { + addr: SocketAddr, + /// 每个请求到的时刻 + hits: Arc>>, +} + +impl Up { + fn hits(&self) -> usize { + self.hits.lock().unwrap().len() + } + fn times(&self) -> Vec { + self.hits.lock().unwrap().clone() + } +} + +async fn upstream(script: Vec, then: Reply) -> Up { + let script = Arc::new(Mutex::new(VecDeque::from(script))); + let hits = Arc::new(Mutex::new(Vec::new())); + let h = hits.clone(); + let app = Router::new().fallback(axum::routing::any(move || { + let (script, h, then) = (script.clone(), h.clone(), then.clone()); + async move { + h.lock().unwrap().push(Instant::now()); + let r = script.lock().unwrap().pop_front().unwrap_or(then); + tokio::time::sleep(Duration::from_millis(r.delay_ms)).await; + let streams = !r.chunks.is_empty(); + let mut b = axum::response::Response::builder().status(r.status).header( + "content-type", + if streams { + "text/event-stream" + } else { + "application/json" + }, + ); + for (k, v) in &r.headers { + b = b.header(*k, v.as_str()); + } + if !streams { + return b.body(axum::body::Body::from(r.body)).unwrap(); + } + let chunks = futures::stream::iter(r.chunks).then(|(ms, c)| async move { + tokio::time::sleep(Duration::from_millis(ms)).await; + Ok::<_, std::io::Error>(bytes::Bytes::from(c)) + }); + b.body(axum::body::Body::from_stream(chunks)).unwrap() + } + })); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + Up { addr, hits } +} + +/// 一家或者两家(按顺序的组「池」)上游,默认路由把请求都交给它们 +fn config(ups: &[&Up], wait_secs: u64) -> tw_config::Config { + let names = ["甲", "乙"]; + let mut yaml = + "version: 1\ndefault_route: 默认\nclients:\n - name: me\n key: tw-k\nproviders:\n" + .to_string(); + for (n, u) in names.iter().zip(ups) { + yaml += &format!( + " - name: {n}\n base_url: http://{}\n key: sk-x\n protocol: anthropic\n", + u.addr + ); + } + let members = names[..ups.len()].join(", "); + yaml += &format!( + "groups:\n - name: 池\n type: fallback\n providers: [{members}]\nroutes:\n - name: 默认\n rules:\n - name: 全部\n to: 池\nfailover:\n wait_secs: {wait_secs}\n" + ); + serde_yaml_ng::from_str(&yaml).unwrap() +} + +/// 总线上发生过的事件,后台一直收着 +#[derive(Clone)] +struct Log(Arc>>); + +impl Log { + async fn until(&self, what: &str, f: impl Fn(&[Event]) -> Option) -> T { + for _ in 0..500 { + if let Some(t) = f(&self.0.lock().unwrap()) { + return t; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("10 秒内没等到{what}:{:#?}", self.0.lock().unwrap()); + } + + /// 客户端写的模型是 `model` 的那个请求的号 + async fn id(&self, model: &str) -> u64 { + self.until(&format!(" {model} 开始"), |evs| { + evs.iter().find_map(|e| match e { + Event::RequestStarted { id, model: m, .. } if m == model => Some(*id), + _ => None, + }) + }) + .await + } + + /// 那个请求的尝试链,和这段对话留在上次那一家的理由 + async fn routed(&self, model: &str) -> (Vec, Option) { + let id = self.id(model).await; + self.until(&format!(" {model} 的路由"), |evs| { + evs.iter().find_map(|e| match e { + Event::RequestRouted { + id: i, + attempts, + affinity, + .. + } if *i == id => Some((attempts.clone(), affinity.as_ref().and_then(|a| a.stayed))), + _ => None, + }) + }) + .await + } + + fn all(&self) -> Vec { + self.0.lock().unwrap().clone() + } +} + +async fn serve(cfg: tw_config::Config) -> (SocketAddr, tw_gateway::AppState, Log) { + let state = tw_gateway::AppState::new(cfg).unwrap(); + let mut rx = state.bus.subscribe(); + let log = Log(Arc::default()); + let into = log.0.clone(); + tokio::spawn(async move { + while let Ok(ev) = rx.recv().await { + into.lock().unwrap().push(ev); + } + }); + let addr = tw_gateway::serve(state.clone(), ([127, 0, 0, 1], 0).into()) + .await + .unwrap(); + (addr, state, log) +} + +/// 一个整包的请求。`conversation` 给了的话带着会话头,`turn` 是这段对话的第几轮(用户说了 +/// 几次话)。**开头那句按对话写**:开头一字不差的请求会被认成同一段对话(见 +/// `tw_gateway::affinity::identity`),留在上次回答它的那一家 +fn ask( + gw: SocketAddr, + model: &str, + conversation: Option<&str>, + turn: usize, +) -> reqwest::RequestBuilder { + let mut messages = Vec::new(); + for i in 0..turn { + if i > 0 { + messages.push(r#"{"role":"assistant","content":"ok"}"#.to_string()); + } + let who = conversation.unwrap_or(model); + messages.push(format!(r#"{{"role":"user","content":"{who}:第 {i} 句"}}"#)); + } + let body = format!( + r#"{{"model":"{model}","max_tokens":16,"system":"你是一个助手","messages":[{}]}}"#, + messages.join(",") + ); + let mut r = reqwest::Client::builder() + .no_proxy() + .build() + .unwrap() + .post(format!("http://{gw}/v1/messages")) + .header("x-api-key", "tw-k") + .header("anthropic-version", "2023-06-01") + .body(body); + if let Some(c) = conversation { + r = r.header("x-claude-code-session-id", c); + } + r +} + +/// 一个流式的请求 +fn ask_stream(gw: SocketAddr, model: &str) -> reqwest::RequestBuilder { + reqwest::Client::builder() + .no_proxy() + .build() + .unwrap() + .post(format!("http://{gw}/v1/messages")) + .header("x-api-key", "tw-k") + .header("anthropic-version", "2023-06-01") + .body(format!( + r#"{{"model":"{model}","max_tokens":16,"stream":true,"messages":[{{"role":"user","content":"{model}"}}]}}"# + )) +} + +/// 直接让一家暂停这么久,像它刚回了一个 429 +fn pause(state: &tw_gateway::AppState, name: &str, ms: u64) { + state.throttles.limited( + name, + std::time::Instant::now(), + Some(Duration::from_millis(ms)), + ); +} + +#[tokio::test] +async fn a_lone_upstream_that_rate_limits_is_waited_for_and_then_answers() { + let up = upstream(vec![limited(&[("retry-after-ms", "400")])], ok()).await; + let (gw, state, log) = serve(config(&[&up], 10)).await; + let t = Instant::now(); + let r = tokio::spawn(ask(gw, "一个", None, 1).send()); + // 在排队:半路才来看的一方也看得出它在等哪一家 + let id = log.id("一个").await; + tokio::time::sleep(Duration::from_millis(150)).await; + let flying = state.bus.in_flight(); + let me = flying.requests.iter().find(|r| r.id == id).unwrap(); + assert!( + me.events.iter().any(|e| matches!( + e, + Event::RequestQueued { upstreams, .. } if upstreams == &["甲".to_string()] + )), + "{:#?}", + me.events + ); + let r = r.await.unwrap().unwrap(); + assert_eq!(r.status(), 200, "客户端察觉不到那一次 429"); + assert!(r.text().await.unwrap().contains("\"ok\"")); + assert!( + t.elapsed() >= Duration::from_millis(380), + "{:?}", + t.elapsed() + ); + let times = up.times(); + assert_eq!(times.len(), 2); + assert!( + times[1] - times[0] >= Duration::from_millis(380), + "没等暂停结束就又发了:{:?}", + times[1] - times[0] + ); + + let (chain, _) = log.routed("一个").await; + assert_eq!(chain.len(), 2, "{chain:#?}"); + assert_eq!( + (chain[0].outcome, chain[0].status), + (AttemptOutcome::Status, Some(429)) + ); + assert!( + chain[0] + .pause_ms + .is_some_and(|ms| (300..=400).contains(&ms)), + "{chain:#?}" + ); + assert_eq!(chain[1].outcome, AttemptOutcome::Served); + assert_eq!(chain[1].pause_ms, None); + assert!( + chain[1].queued_ms.is_some_and(|ms| ms >= 250), + "等暂停的那一段该记在发出去的那一跳上:{chain:#?}" + ); + // 暂停不是失败:不熔断;开始暂停、暂停结束各报一次 + let evs = log.all(); + assert!(!evs.iter().any(|e| matches!(e, Event::HealthChanged { .. }))); + let changes: Vec> = evs + .iter() + .filter_map(|e| match e { + Event::ThrottleChanged { + provider, throttle, .. + } if provider == "甲" => Some(throttle.as_ref().map(|t| t.kind)), + _ => None, + }) + .collect(); + assert_eq!(changes, [Some(ThrottleKind::RateLimited), None], "{evs:#?}"); + let queued = evs + .iter() + .filter(|e| matches!(e, Event::RequestQueued { id: i, .. } if *i == id)) + .count(); + assert_eq!(queued, 1, "一次排队报一条"); + // 不排了也报一条:在路由事件之前(路由事件要等整条尝试链走完) + let at = |f: &dyn Fn(&Event) -> bool| evs.iter().position(f); + let queued = at(&|e| matches!(e, Event::RequestQueued { id: i, .. } if *i == id)).unwrap(); + let dequeued = + at(&|e| matches!(e, Event::RequestDequeued { id: i, .. } if *i == id)).expect("没报不排了"); + let routed = at(&|e| matches!(e, Event::RequestRouted { id: i, .. } if *i == id)).unwrap(); + assert!(queued < dequeued && dequeued < routed, "{evs:#?}"); + assert_eq!( + evs.iter() + .filter(|e| matches!(e, Event::RequestDequeued { id: i, .. } if *i == id)) + .count(), + 1 + ); +} + +#[tokio::test] +async fn after_a_pause_one_request_goes_first_and_the_rest_follow_its_answer() { + // 头一个被限流 300 毫秒;之后的每个都要 400 毫秒才回响应头 + let slow = Reply { + delay_ms: 400, + ..ok() + }; + let up = upstream(vec![limited(&[("retry-after-ms", "300")])], slow).await; + let (gw, _state, _log) = serve(config(&[&up], 10)).await; + let first = tokio::spawn(ask(gw, "一", None, 1).send()); + tokio::time::sleep(Duration::from_millis(100)).await; + // 暂停着的时候又来了两个:排队 + let more: Vec<_> = ["二", "三"] + .into_iter() + .map(|m| tokio::spawn(ask(gw, m, None, 1).send())) + .collect(); + assert_eq!(first.await.unwrap().unwrap().status(), 200); + for r in more { + assert_eq!(r.await.unwrap().unwrap().status(), 200); + } + let times = up.times(); + assert_eq!(times.len(), 4, "429 一次,之后三个各一次"); + // 暂停结束:只放一个过去探路。它的响应头 400 毫秒后才到,别的等它 + assert!(times[1] - times[0] >= Duration::from_millis(280)); + for later in ×[2..] { + assert!( + *later - times[1] >= Duration::from_millis(350), + "探路的还没有结论别的就发了:{:?}", + *later - times[1] + ); + } +} + +#[tokio::test] +async fn when_nothing_recovers_in_time_the_client_gets_a_429_saying_when() { + let up = upstream(Vec::new(), limited(&[("retry-after", "30")])).await; + let (gw, _state, log) = serve(config(&[&up], 2)).await; + let t = Instant::now(); + let r = ask(gw, "一", None, 1).send().await.unwrap(); + // 暂停 30 秒,这个请求最多等 2 秒:不用干等,当场说 + assert!( + t.elapsed() < Duration::from_millis(1500), + "{:?}", + t.elapsed() + ); + assert_eq!(r.status(), 429); + let h = |k: &str| r.headers().get(k).map(|v| v.to_str().unwrap().to_string()); + assert_eq!(h("x-thinkwatch-error").as_deref(), Some("rate_limited")); + assert_eq!(h("x-should-retry").as_deref(), Some("true")); + let secs: u64 = h("retry-after").unwrap().parse().unwrap(); + assert!((29..=30).contains(&secs), "{secs}"); + let ms: u64 = h("retry-after-ms").unwrap().parse().unwrap(); + assert!((28_000..=30_000).contains(&ms), "{ms}"); + let body: serde_json::Value = r.json().await.unwrap(); + let said = body["error"]["message"].as_str().unwrap(); + assert!( + said.contains("paused after answering 429 (rate limited): `甲`"), + "{said}" + ); + assert_eq!(up.hits(), 1); + let failed = log + .until("失败", |evs| { + evs.iter().find_map(|e| match e { + Event::RequestFailed { + source, message, .. + } => Some((*source, message.code.clone())), + _ => None, + }) + }) + .await; + assert_eq!(failed.0.slug(), "rate_limited"); + assert_eq!(failed.1, "gw.wait.rate_limited"); + + // 还暂停着:下一个请求不发给它,尝试链上说它在暂停 + let r = ask(gw, "二", None, 1).send().await.unwrap(); + assert_eq!(r.status(), 429); + assert_eq!(up.hits(), 1, "暂停着还发了"); + let (chain, _) = log.routed("二").await; + assert_eq!(chain.len(), 1, "{chain:#?}"); + assert_eq!(chain[0].skipped, Some(ServeSkip::Throttled)); + assert_eq!( + chain[0].error.as_ref().map(|m| m.code.as_str()), + Some("gw.upstream.paused") + ); +} + +#[tokio::test] +async fn counting_tokens_is_estimated_while_the_upstream_is_paused() { + let up = upstream(vec![limited(&[("retry-after-ms", "60000")])], ok()).await; + let (gw, state, log) = serve(config(&[&up], 10)).await; + let count = |model: &str| { + reqwest::Client::builder() + .no_proxy() + .build() + .unwrap() + .post(format!("http://{gw}/v1/messages/count_tokens")) + .header("x-api-key", "tw-k") + .header("anthropic-version", "2023-06-01") + .body(format!( + r#"{{"model":"{model}","messages":[{{"role":"user","content":"数一数这句话"}}]}}"# + )) + }; + // 数 token 的请求被限流了:这一家暂停,这一次由网关估,不排队 + let t = Instant::now(); + let r = count("一").send().await.unwrap(); + assert_eq!(r.status(), 200); + assert_eq!(r.headers()["x-thinkwatch-local"], "1"); + assert!(t.elapsed() < Duration::from_secs(2), "{:?}", t.elapsed()); + let (chain, _) = log.routed("一").await; + assert_eq!(chain.len(), 2, "{chain:#?}"); + assert_eq!(chain[0].status, Some(429)); + assert!(chain[0].pause_ms.is_some()); + assert_eq!(chain[1].outcome, AttemptOutcome::Estimated); + assert_eq!(up.hits(), 1); + assert!(state.throttles.view("甲").is_some(), "该暂停着"); + + // 还暂停着:不发,直接估 + let r = count("二").send().await.unwrap(); + assert_eq!(r.status(), 200); + assert_eq!(r.headers()["x-thinkwatch-local"], "1"); + let body: serde_json::Value = r.json().await.unwrap(); + assert!( + body["input_tokens"].as_u64().is_some_and(|n| n > 0), + "{body}" + ); + assert_eq!(up.hits(), 1); + let (chain, _) = log.routed("二").await; + assert_eq!(chain.len(), 1, "{chain:#?}"); + assert_eq!(chain[0].outcome, AttemptOutcome::Estimated); +} + +#[tokio::test] +async fn aborting_a_queued_request_takes_it_out_of_the_queue() { + let up = upstream(Vec::new(), ok()).await; + let (gw, state, log) = serve(config(&[&up], 10)).await; + pause(&state, "甲", 800); + let r = tokio::spawn(ask(gw, "一", None, 1).send()); + let id = log.id("一").await; + tokio::time::sleep(Duration::from_millis(100)).await; + state.aborts.request(id).expect("还在跑"); + let r = r.await.unwrap().unwrap(); + assert_eq!(r.status(), 499); + let source = log + .until("中止", |evs| { + evs.iter().find_map(|e| match e { + Event::RequestFailed { id: i, source, .. } if *i == id => Some(*source), + _ => None, + }) + }) + .await; + assert_eq!(source.slug(), "aborted"); + assert_eq!(up.hits(), 0); + // 被中止也是不排了 + assert!( + log.all() + .iter() + .any(|e| matches!(e, Event::RequestDequeued { id: i, .. } if *i == id)) + ); + // 它不在队里了:暂停结束后下一个请求当场就发(它要是还排在队头,后面的要一直等它) + tokio::time::sleep(Duration::from_millis(800)).await; + let t = Instant::now(); + let r = ask(gw, "二", None, 1).send().await.unwrap(); + assert_eq!(r.status(), 200); + assert!( + t.elapsed() < Duration::from_millis(500), + "{:?}", + t.elapsed() + ); +} + +#[tokio::test] +async fn a_conversation_waits_for_its_paused_upstream_and_a_new_one_moves_on() { + let (a, b) = ( + upstream(Vec::new(), ok()).await, + upstream(Vec::new(), ok()).await, + ); + let (gw, state, log) = serve(config(&[&a, &b], 5)).await; + // 这段对话由甲回答,读了 5000 token 的缓存:下一轮留在甲 + let r = ask(gw, "第一轮", Some("对话"), 1).send().await.unwrap(); + assert_eq!(r.status(), 200); + r.text().await.unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + + // 甲暂停半秒:这段对话等它;新的对话当场换乙 + pause(&state, "甲", 500); + let t = Instant::now(); + let stays = tokio::spawn(ask(gw, "第二轮", Some("对话"), 2).send()); + let fresh = ask(gw, "新的", Some("别的对话"), 1).send().await.unwrap(); + assert_eq!(fresh.status(), 200); + assert!(t.elapsed() < Duration::from_millis(400), "新的对话等了甲"); + let (chain, _) = log.routed("新的").await; + assert_eq!(chain.len(), 2, "{chain:#?}"); + assert_eq!(chain[0].provider, "甲"); + assert_eq!(chain[0].skipped, Some(ServeSkip::Throttled)); + assert_eq!(chain[0].queued_ms, None, "新的对话不等"); + assert_eq!( + (chain[1].provider.as_str(), chain[1].outcome), + ("乙", AttemptOutcome::Served) + ); + + let r = stays.await.unwrap().unwrap(); + assert_eq!(r.status(), 200); + r.text().await.unwrap(); + let (chain, stayed) = log.routed("第二轮").await; + assert_eq!(stayed, Some(Stay::Cache)); + assert_eq!(chain.len(), 1, "{chain:#?}"); + assert_eq!( + (chain[0].provider.as_str(), chain[0].outcome), + ("甲", AttemptOutcome::Served) + ); + assert!(chain[0].queued_ms.is_some_and(|ms| ms >= 350), "{chain:#?}"); + tokio::time::sleep(Duration::from_millis(100)).await; + + // 甲要暂停一分钟,这个请求最多等 5 秒:等也白等,不再留在甲,照常排序 + pause(&state, "甲", 60_000); + let t = Instant::now(); + let r = ask(gw, "第三轮", Some("对话"), 3).send().await.unwrap(); + assert_eq!(r.status(), 200); + assert!( + t.elapsed() < Duration::from_millis(500), + "{:?}", + t.elapsed() + ); + let (chain, stayed) = log.routed("第三轮").await; + assert_eq!(stayed, None, "这段对话该放开甲"); + assert_eq!(chain.len(), 2, "{chain:#?}"); + assert_eq!(chain[0].skipped, Some(ServeSkip::Throttled)); + assert_eq!(chain[1].provider, "乙"); +} + +#[tokio::test] +async fn a_rate_limited_upstream_moves_the_request_on_and_stays_out_of_the_breaker() { + // 甲一直限流,乙照常:每个请求都去了乙,甲暂停之后不再挨发 + let a = upstream(Vec::new(), limited(&[("retry-after", "30")])).await; + let b = upstream(Vec::new(), ok()).await; + let (gw, state, log) = serve(config(&[&a, &b], 10)).await; + for m in ["一", "二", "三", "四"] { + let r = ask(gw, m, None, 1).send().await.unwrap(); + assert_eq!(r.status(), 200); + } + assert_eq!(a.hits(), 1, "暂停着还发给了甲"); + assert_eq!(b.hits(), 4); + let (chain, _) = log.routed("一").await; + assert_eq!(chain[0].status, Some(429)); + assert!( + chain[0].pause_ms.is_some_and(|ms| ms > 29_000), + "{chain:#?}" + ); + let (chain, _) = log.routed("二").await; + assert_eq!(chain[0].skipped, Some(ServeSkip::Throttled), "{chain:#?}"); + // 限流不进熔断,也不算进成功率 + assert!(state.health.is_available("甲")); + assert!(state.health.success_rates(&["甲".to_string()]).is_empty()); + assert!( + !log.all() + .iter() + .any(|e| matches!(e, Event::HealthChanged { .. })) + ); +} + +/// 探路的又被限流了,回到甲的队头,接着去试乙(乙要 5 秒)。**它在乙那儿的时候不该挡着甲的队**: +/// 甲的暂停 300 毫秒后就结束了,排在后面的请求该轮到探路,而不是等到期限回 429 +#[tokio::test] +async fn a_probe_limited_again_does_not_hold_the_queue_while_it_tries_another_upstream() { + let a = upstream( + vec![ + limited(&[("retry-after-ms", "300")]), + limited(&[("retry-after-ms", "300")]), + ], + ok(), + ) + .await; + let slow = Reply { + delay_ms: 5000, + ..ok() + }; + let b = upstream(vec![ok()], slow).await; + let mut cfg = config(&[&a, &b], 2); + cfg.providers[1].max_concurrent = Some(1); + let (gw, _state, _log) = serve(cfg).await; + // 甲 429(暂停 300 毫秒),乙当场答 + let r = ask(gw, "零", None, 1).send().await.unwrap(); + assert_eq!(r.status(), 200); + r.text().await.unwrap(); + tokio::time::sleep(Duration::from_millis(400)).await; + // 探路的又被限流,去了乙(占着乙唯一的位置 5 秒) + let probe = tokio::spawn(ask(gw, "一", None, 1).send()); + tokio::time::sleep(Duration::from_millis(150)).await; + // 甲暂停着、乙满着:在两家里等 + let t = Instant::now(); + let r = ask(gw, "二", None, 1).send().await.unwrap(); + assert_eq!( + r.status(), + 200, + "甲早就能发了,却等了 {:?},甲只收到 {} 个请求", + t.elapsed(), + a.hits() + ); + let _ = probe.await; +} + +/// 上游写了一个大得离谱的 `retry-after-ms`:照上限暂停,不能把请求弄崩 +#[tokio::test] +async fn an_absurd_retry_after_ms_is_capped_rather_than_crashing_the_request() { + let up = upstream(vec![limited(&[("retry-after-ms", "1e300")])], ok()).await; + let (gw, _state, _log) = serve(config(&[&up], 1)).await; + let r = ask(gw, "一", None, 1) + .send() + .await + .expect("网关该回话,不是断开连接"); + assert_eq!(r.status(), 429); +} + +/// 上游一直回 429 和 `Retry-After: 0`:不能不停地重发,也不能过了等待期限还在发 +#[tokio::test] +async fn a_zero_retry_after_neither_loops_nor_outlives_the_deadline() { + let up = upstream(Vec::new(), limited(&[("retry-after", "0")])).await; + let (gw, _state, _log) = serve(config(&[&up], 2)).await; + // 最多等 2 秒的请求,5 秒还没结束就是停不下来了 + let r = tokio::time::timeout(Duration::from_secs(5), ask(gw, "一", None, 1).send()).await; + let hits = up.hits(); + assert!(r.is_ok(), "过了期限 3 秒还在重发,已经发了 {hits} 次"); + assert!(hits < 20, "两秒里发了 {hits} 次"); +} + +/// 探路的发出去了还没有结论:宽限期的结束不是「知道的恢复时刻」。期限比宽限期近的请求 +/// 该等探路的结论,不该当场回 429 +#[tokio::test] +async fn a_request_waits_for_a_probe_in_flight_even_when_its_deadline_is_within_the_grace() { + let slow = Reply { + delay_ms: 1000, + ..ok() + }; + let up = upstream(vec![limited(&[("retry-after-ms", "300")])], slow).await; + let (gw, _state, _log) = serve(config(&[&up], 4)).await; + // 头一个:429,300 毫秒后探路,1300 毫秒时探路的有了结论 + let first = tokio::spawn(ask(gw, "一", None, 1).send()); + tokio::time::sleep(Duration::from_millis(500)).await; + let t = Instant::now(); + let r = ask(gw, "二", None, 1).send().await.unwrap(); + assert_eq!( + r.status(), + 200, + "{:?} 就回了 429,探路的不到一秒后就过了", + t.elapsed() + ); + assert_eq!(first.await.unwrap().unwrap().status(), 200); +} + +/// 回答压着开头的时候,响应头是 200 不说明这一家好了:探路的拿到 200、流里接着报限流,**它的 +/// 结论是又被限流**。排在后面的不能在响应头到的那一刻就一拥而上,退避也不归零 +#[tokio::test] +async fn a_probe_whose_stream_reports_a_rate_limit_after_a_200_head_does_not_open_the_queue() { + let up = upstream( + vec![ + limited(&[("retry-after-ms", "300")]), + // 探路的:响应头当场到,限流 300 毫秒之后才在流里报 + limited_in_the_stream(300), + ], + streamed(), + ) + .await; + let (gw, _state, log) = serve(config(&[&up], 10)).await; + let first = tokio::spawn(ask_stream(gw, "一").send()); + tokio::time::sleep(Duration::from_millis(100)).await; + let more: Vec<_> = ["二", "三"] + .into_iter() + .map(|m| tokio::spawn(ask_stream(gw, m).send())) + .collect(); + for r in std::iter::once(first).chain(more) { + let r = r.await.unwrap().unwrap(); + assert_eq!(r.status(), 200); + let text = r.text().await.unwrap(); + assert!(text.contains("text_delta"), "{text}"); + assert!( + !text.contains("rate_limit_error"), + "客户端看到了限流:{text}" + ); + } + let times = up.times(); + assert_eq!(times.len(), 5, "429、探路的、再探一次、后面两个"); + for later in ×[2..] { + assert!( + *later - times[1] >= Duration::from_millis(1_200), + "探路的在流里又被限流了,后面的却跟着发了:{:?}", + *later - times[1] + ); + } + let (chain, _) = log.routed("一").await; + assert_eq!(chain.len(), 3, "{chain:#?}"); + assert!( + chain[1] + .pause_ms + .is_some_and(|ms| (900..=1_000).contains(&ms)), + "{chain:#?}" + ); + assert_eq!(chain[2].outcome, AttemptOutcome::Served); +} + +/// 只有一家、它刚限流过:流的开头(内容之前)又报了限流,和回 429 一样暂停、排队、再发, +/// **客户端察觉不到**。一直好好的最后一家不压开头(客户端当场拿到响应头),见 `hop` +#[tokio::test] +async fn a_rate_limit_inside_the_stream_before_any_content_is_invisible_on_the_last_upstream() { + let up = upstream( + vec![ + limited(&[("retry-after-ms", "300")]), + ok(), + limited_in_the_stream(0), + ], + streamed(), + ) + .await; + let (gw, _state, log) = serve(config(&[&up], 10)).await; + // 先被限流一次(等过之后照常答了):这一家最近限流过 + assert_eq!(ask(gw, "零", None, 1).send().await.unwrap().status(), 200); + let r = ask_stream(gw, "一").send().await.unwrap(); + assert_eq!(r.status(), 200); + let text = r.text().await.unwrap(); + assert!(text.contains("text_delta"), "{text}"); + assert!( + !text.contains("rate_limit_error"), + "客户端看到了限流:{text}" + ); + let times = up.times(); + assert_eq!(times.len(), 4); + assert!(times[3] - times[2] >= Duration::from_millis(950)); + let (chain, _) = log.routed("一").await; + assert_eq!(chain.len(), 2, "{chain:#?}"); + assert_eq!(chain[0].outcome, AttemptOutcome::Error); + assert!( + chain[0] + .pause_ms + .is_some_and(|ms| (900..=1_000).contains(&ms)), + "{chain:#?}" + ); + assert_eq!(chain[1].outcome, AttemptOutcome::Served); + assert!(chain[1].queued_ms.is_some_and(|ms| ms >= 900), "{chain:#?}"); +} + +/// 探路的拿到的 429 是额度用完了(等不到):它没有「好了没有」的结论,**排在后面的不用干等完 +/// 宽限期**,照常发,这一家怎么样由熔断说了算 +#[tokio::test] +async fn a_probe_answered_with_a_used_up_quota_lets_the_queue_go_at_once() { + let used_up = Reply { + body: r#"{"type":"error","error":{"type":"rate_limit_error","message":"You have reached your usage limit for this plan"}}"#.into(), + delay_ms: 400, + ..limited(&[]) + }; + let up = upstream(vec![limited(&[("retry-after-ms", "300")]), used_up], ok()).await; + let (gw, _state, _log) = serve(config(&[&up], 20)).await; + let first = tokio::spawn(ask(gw, "一", None, 1).send()); + // 探路的发出去了、还没有结论:这时来的排在后面等它 + tokio::time::sleep(Duration::from_millis(450)).await; + let t = Instant::now(); + let r = ask(gw, "二", None, 1).send().await.unwrap(); + assert_eq!(r.status(), 200); + assert!( + t.elapsed() < Duration::from_secs(2), + "额度用完的探路之后还干等了宽限期:{:?}", + t.elapsed() + ); + // 探路的那个请求自己收到的是额度用完 + assert_eq!(first.await.unwrap().unwrap().status(), 429); + assert_eq!(up.hits(), 3); +} + +/// 暂停结束、探路的名额拿到了还没发出去(插件、取密钥、等空位的时候)上游改了名:名额的 +/// 主人认的是旧名字,新名字下**不能一直挂着这张发不出去的名额** —— 它不会有结论,宽限期也 +/// 不从它算起,这一家就再也放不出一个请求 +#[tokio::test] +async fn renaming_while_an_unsent_probe_ticket_is_out_does_not_wedge_the_upstream() { + let up = upstream(Vec::new(), ok()).await; + let (_gw, state, _log) = serve(config(&[&up], 10)).await; + pause(&state, "甲", 50); + tokio::time::sleep(Duration::from_millis(100)).await; + let ticket = match state.throttles.gate("甲", &[], None) { + tw_gateway::throttle::Gate::Go(Some(p)) => p, + _ => panic!("暂停结束后头一个该拿到探路的名额"), + }; + state.throttles.rename("甲", "丙"); + // 名额的主人放弃了(没发出去),或者照旧名字说它发出去了:新名字下都该接着有人能探路 + drop(ticket); + tokio::time::sleep(tw_gateway::throttle::PROBE_GRACE + Duration::from_millis(500)).await; + assert!( + matches!( + state.throttles.gate("丙", &[], None), + tw_gateway::throttle::Gate::Go(_) + ), + "改名之后 {:?} 还发不出去:{:?}", + tw_gateway::throttle::PROBE_GRACE, + state.throttles.peek("丙", &[]) + ); +} + +/// `failover.wait_secs: 0` 是不等:一家满着时当场回 429。**它没排过队**,不该报 +/// `RequestQueued`(界面会把它显示成「排队中」) +#[tokio::test] +async fn a_request_that_is_not_allowed_to_wait_is_never_reported_as_queued() { + let slow = Reply { + delay_ms: 1000, + ..ok() + }; + let up = upstream(Vec::new(), slow).await; + let mut cfg = config(&[&up], 0); + cfg.providers[0].max_concurrent = Some(1); + let (gw, _state, log) = serve(cfg).await; + let first = tokio::spawn(ask(gw, "一", None, 1).send()); + tokio::time::sleep(Duration::from_millis(200)).await; + let r = ask(gw, "二", None, 1).send().await.unwrap(); + assert_eq!(r.status(), 429); + let id = log.id("二").await; + let queued: Vec = log + .all() + .into_iter() + .filter(|e| { + matches!(e, Event::RequestQueued { id: i, .. } | Event::RequestDequeued { id: i, .. } if *i == id) + }) + .collect(); + assert!(queued.is_empty(), "不等的请求报了排队:{queued:#?}"); + assert_eq!(first.await.unwrap().unwrap().status(), 200); +} diff --git a/crates/tw-gateway/tests/upstream_limits_queue.rs b/crates/tw-gateway/tests/upstream_limits_queue.rs new file mode 100644 index 00000000..2325fa02 --- /dev/null +++ b/crates/tw-gateway/tests/upstream_limits_queue.rs @@ -0,0 +1,501 @@ +//! 上游的用量上限接在管线上(`tw_gateway::upstream_limits` + `tw_gateway::throttle`),端到端。 +//! +//! 分钟、小时的窗口满了:请求等它空出来再发;天、周、月用完了:跳过,都用完了回 429、到重置 +//! 之前别重试;回了 429 的、没发出去的那一跳不算进这家的上限。时钟可以往前拨([`Ahead`]): +//! 拨到第一个请求之后 59 秒,这一分钟的窗口不到一秒就空出来,测试不用干等一分钟。 + +use std::collections::VecDeque; +use std::net::SocketAddr; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use axum::Router; +use tw_api::{AttemptOutcome, AttemptView, Event, ServeSkip, ThrottleKind}; +use tw_gateway::key_limits::Clock; + +/// 假上游的一次答复 +#[derive(Clone)] +struct Reply { + status: u16, + headers: Vec<(&'static str, String)>, + body: String, +} + +fn ok() -> Reply { + Reply { + status: 200, + headers: Vec::new(), + body: r#"{"type":"message","role":"assistant","model":"claude-sonnet-5","content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn","usage":{"input_tokens":3,"output_tokens":1}}"#.into(), + } +} + +/// 限流:暂停 `ms` 毫秒 +fn limited(ms: u64) -> Reply { + Reply { + status: 429, + headers: vec![("retry-after-ms", ms.to_string())], + body: r#"{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}"# + .into(), + } +} + +/// 一家假上游:先按 `script` 的顺序答,答完了一律答 `then`。数着收到了几个请求 +struct Up { + addr: SocketAddr, + hits: Arc>>, +} + +impl Up { + fn hits(&self) -> usize { + self.hits.lock().unwrap().len() + } +} + +async fn upstream(script: Vec, then: Reply) -> Up { + let script = Arc::new(Mutex::new(VecDeque::from(script))); + let hits = Arc::new(Mutex::new(Vec::new())); + let h = hits.clone(); + let app = Router::new().fallback(axum::routing::any(move || { + let (script, h, then) = (script.clone(), h.clone(), then.clone()); + async move { + h.lock().unwrap().push(Instant::now()); + let r = script.lock().unwrap().pop_front().unwrap_or(then); + let mut b = axum::response::Response::builder() + .status(r.status) + .header("content-type", "application/json"); + for (k, v) in &r.headers { + b = b.header(*k, v.as_str()); + } + b.body(axum::body::Body::from(r.body)).unwrap() + } + })); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + Up { addr, hits } +} + +/// 一个没人听的地址:连不上 +async fn nobody() -> SocketAddr { + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + drop(l); + addr +} + +/// 系统时钟,可以往前拨 +struct Ahead(AtomicI64); + +impl Ahead { + fn new() -> Arc { + Arc::new(Self(AtomicI64::new(0))) + } + fn jump(&self, ms: i64) { + self.0.fetch_add(ms, Ordering::SeqCst); + } +} + +impl Clock for Ahead { + fn now_ms(&self) -> i64 { + tw_gateway::key_limits::SystemClock.now_ms() + self.0.load(Ordering::SeqCst) + } + fn period(&self, per: tw_config::LimitPer, at_ms: i64) -> (i64, i64) { + tw_gateway::key_limits::SystemClock.period(per, at_ms) + } + fn show(&self, at_ms: i64) -> String { + tw_gateway::key_limits::SystemClock.show(at_ms) + } +} + +/// 甲(和乙)按顺序组成「池」,默认路由把请求都交给它们。每一家带着自己的 `limits`(YAML) +fn config(ups: &[(SocketAddr, &str)], wait_secs: u64) -> tw_config::Config { + let names = ["甲", "乙"]; + let mut yaml = + "version: 1\ndefault_route: 默认\nclients:\n - name: me\n key: tw-k\nproviders:\n" + .to_string(); + for (n, (addr, limits)) in names.iter().zip(ups) { + yaml += &format!( + " - name: {n}\n base_url: http://{addr}\n key: sk-x\n protocol: anthropic\n limits: {limits}\n" + ); + } + let members = names[..ups.len()].join(", "); + yaml += &format!( + "groups:\n - name: 池\n type: fallback\n providers: [{members}]\nroutes:\n - name: 默认\n rules:\n - name: 全部\n to: 池\nfailover:\n wait_secs: {wait_secs}\n" + ); + serde_yaml_ng::from_str(&yaml).unwrap() +} + +/// 总线上发生过的事件,后台一直收着 +#[derive(Clone)] +struct Log(Arc>>); + +impl Log { + async fn until(&self, what: &str, f: impl Fn(&[Event]) -> Option) -> T { + for _ in 0..500 { + if let Some(t) = f(&self.0.lock().unwrap()) { + return t; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("10 秒内没等到{what}:{:#?}", self.0.lock().unwrap()); + } + + async fn id(&self, model: &str) -> u64 { + self.until(&format!(" {model} 开始"), |evs| { + evs.iter().find_map(|e| match e { + Event::RequestStarted { id, model: m, .. } if m == model => Some(*id), + _ => None, + }) + }) + .await + } + + async fn routed(&self, model: &str) -> Vec { + let id = self.id(model).await; + self.until(&format!(" {model} 的路由"), |evs| { + evs.iter().find_map(|e| match e { + Event::RequestRouted { + id: i, attempts, .. + } if *i == id => Some(attempts.clone()), + _ => None, + }) + }) + .await + } + + async fn failed_code(&self, model: &str) -> String { + let id = self.id(model).await; + self.until(&format!(" {model} 失败"), |evs| { + evs.iter().find_map(|e| match e { + Event::RequestFailed { id: i, message, .. } if *i == id => { + Some(message.code.clone()) + } + _ => None, + }) + }) + .await + } + + fn all(&self) -> Vec { + self.0.lock().unwrap().clone() + } +} + +async fn serve( + cfg: tw_config::Config, + clock: Arc, +) -> (SocketAddr, tw_gateway::AppState, Log) { + let mut state = tw_gateway::AppState::new(cfg).unwrap(); + state.set_key_limits_clock(clock); + let mut rx = state.bus.subscribe(); + let log = Log(Arc::default()); + let into = log.0.clone(); + tokio::spawn(async move { + while let Ok(ev) = rx.recv().await { + into.lock().unwrap().push(ev); + } + }); + let addr = tw_gateway::serve(state.clone(), ([127, 0, 0, 1], 0).into()) + .await + .unwrap(); + (addr, state, log) +} + +async fn ask(gw: SocketAddr, model: &str) -> reqwest::Response { + let body = format!( + r#"{{"model":"{model}","max_tokens":16,"messages":[{{"role":"user","content":"{model}"}}]}}"# + ); + reqwest::Client::builder() + .no_proxy() + .build() + .unwrap() + .post(format!("http://{gw}/v1/messages")) + .header("x-api-key", "tw-k") + .header("anthropic-version", "2023-06-01") + .body(body) + .send() + .await + .unwrap() +} + +fn head(r: &reqwest::Response, h: &str) -> Option { + r.headers().get(h).map(|v| v.to_str().unwrap().to_string()) +} + +/// 这一家第 `i` 条上限此刻用了多少(`ProviderView::limits` 里的那个数) +fn used(state: &tw_gateway::AppState, name: &str, i: usize) -> u64 { + let cfg = state.config(); + let p = cfg.providers.iter().find(|p| p.name == name).unwrap(); + state.upstream_limits.view(name, &p.limits)[i].used +} + +/// 一分钟一个请求:第二个请求**等窗口空出来再发**,客户端照常拿到回答。尝试链上先是跳过的 +/// 那一跳(`window_full`),再是等过之后发出去的那一跳;这一家一满就报「用量已满」 +#[tokio::test] +async fn a_full_minute_window_makes_the_request_wait_and_then_send() { + let up = upstream(Vec::new(), ok()).await; + let clock = Ahead::new(); + let (gw, state, log) = serve( + config(&[(up.addr, "[{per: minute, requests: 1}]")], 10), + clock.clone(), + ) + .await; + assert_eq!(ask(gw, "第一个").await.status(), 200); + assert_eq!(used(&state, "甲", 0), 1); + // 用满了就报:不用等哪个请求结束 + log.until("用量已满", |evs| { + evs.iter().find_map(|e| match e { + Event::ThrottleChanged { + provider, + throttle: Some(t), + .. + } if provider == "甲" && t.kind == ThrottleKind::WindowFull => Some(()), + _ => None, + }) + }) + .await; + // 拨到第一个请求之后 59.2 秒:再过不到一秒滑出去 + clock.jump(59_200); + let throttle = state.throttles.view("甲").expect("still full"); + assert_eq!(throttle.kind, ThrottleKind::WindowFull); + assert_eq!(throttle.limit, Some(0)); + + let t = Instant::now(); + let r = ask(gw, "第二个").await; + assert_eq!(r.status(), 200); + let took = t.elapsed(); + assert!( + took >= Duration::from_millis(500) && took < Duration::from_secs(3), + "该等窗口空出来:{took:?}" + ); + assert_eq!(up.hits(), 2); + + let chain = log.routed("第二个").await; + assert_eq!(chain.len(), 2, "{chain:#?}"); + assert_eq!(chain[0].skipped, Some(ServeSkip::WindowFull), "{chain:#?}"); + assert_eq!(chain[1].outcome, AttemptOutcome::Served); + assert!(chain[1].queued_ms.is_some_and(|ms| ms >= 400), "{chain:#?}"); + let id = log.id("第二个").await; + assert!( + log.all().iter().any(|e| matches!( + e, + Event::RequestQueued { id: i, upstreams, .. } if *i == id && upstreams == &["甲".to_string()] + )), + "没报排队" + ); + // 空出来的那个位置当场给了它:这一家又满了,报的是新的恢复时刻 + let fulls: Vec = log + .all() + .iter() + .filter_map(|e| match e { + Event::ThrottleChanged { + provider, + throttle: Some(t), + .. + } if provider == "甲" && t.kind == ThrottleKind::WindowFull => Some(t.until_ms), + _ => None, + }) + .collect(); + assert!( + fulls.len() >= 2 && fulls.last() > fulls.first(), + "{fulls:?}" + ); + // 用量上限满了不是失败:不熔断 + assert!( + !log.all() + .iter() + .any(|e| matches!(e, Event::HealthChanged { .. })) + ); +} + +/// `every: 5`:窗口是最近 5 分钟。一分钟之后还满着;离空出来比这个请求能等的还久,当场 429, +/// `Retry-After` 是还要多久;拨到快满 5 分钟时,等一下就发得出去 +#[tokio::test] +async fn an_every_five_minutes_window_spans_five_minutes() { + let up = upstream(Vec::new(), ok()).await; + let clock = Ahead::new(); + let (gw, _state, log) = serve( + config(&[(up.addr, "[{per: minute, every: 5, requests: 1}]")], 2), + clock.clone(), + ) + .await; + assert_eq!(ask(gw, "第一个").await.status(), 200); + clock.jump(61_000); + + let t = Instant::now(); + let r = ask(gw, "第二个").await; + assert!(t.elapsed() < Duration::from_secs(1), "等不到就不该干等"); + assert_eq!(r.status(), 429); + assert_eq!(head(&r, "x-should-retry").as_deref(), Some("true")); + let after: u64 = head(&r, "retry-after").unwrap().parse().unwrap(); + assert!((237..=240).contains(&after), "Retry-After {after}"); + let after_ms: u64 = head(&r, "retry-after-ms").unwrap().parse().unwrap(); + assert!( + (236_000..=239_000).contains(&after_ms), + "retry-after-ms {after_ms}" + ); + assert_eq!(log.failed_code("第二个").await, "gw.wait.window_full"); + assert_eq!(up.hits(), 1); + + // 离第一个请求 5 分钟差 800 毫秒 + clock.jump(238_200); + let r = ask(gw, "第三个").await; + assert_eq!(r.status(), 200); + assert_eq!(up.hits(), 2); +} + +/// 一天一个请求:甲今天用完了,第二个请求**跳过它**、发给乙,尝试链上记着 +/// `period_used_up`;两家都用完了,429 说清是哪几家,到重置之前别重试 +/// (`x-should-retry: false`,没有 `retry-after-ms`) +#[tokio::test] +async fn a_used_up_day_is_skipped_and_all_used_up_is_final() { + let (a, b) = ( + upstream(Vec::new(), ok()).await, + upstream(Vec::new(), ok()).await, + ); + let noon = chrono::DateTime::parse_from_rfc3339("2026-10-05T12:00:00+08:00") + .unwrap() + .timestamp_millis(); + let clock = Arc::new(tw_gateway::key_limits::TestClock::new(noon, 8 * 3600)); + let day = "[{per: day, requests: 1}]"; + let (gw, state, log) = serve(config(&[(a.addr, day), (b.addr, day)], 10), clock).await; + + assert_eq!(ask(gw, "第一个").await.status(), 200); + assert_eq!((a.hits(), b.hits()), (1, 0)); + let throttle = state.throttles.view("甲").expect("used up"); + assert_eq!(throttle.kind, ThrottleKind::PeriodUsedUp); + assert_eq!(throttle.limit, Some(0)); + + let t = Instant::now(); + assert_eq!(ask(gw, "第二个").await.status(), 200); + assert!(t.elapsed() < Duration::from_secs(1), "用完了的不该等"); + assert_eq!((a.hits(), b.hits()), (1, 1)); + let chain = log.routed("第二个").await; + assert_eq!(chain.len(), 2, "{chain:#?}"); + assert_eq!( + (chain[0].provider.as_str(), chain[0].skipped), + ("甲", Some(ServeSkip::PeriodUsedUp)) + ); + assert_eq!( + (chain[1].provider.as_str(), chain[1].outcome), + ("乙", AttemptOutcome::Served) + ); + + let t = Instant::now(); + let r = ask(gw, "第三个").await; + assert!(t.elapsed() < Duration::from_secs(1), "用完了的不该等"); + assert_eq!(r.status(), 429); + assert_eq!(head(&r, "x-should-retry").as_deref(), Some("false")); + assert_eq!(head(&r, "retry-after-ms"), None); + // 中午到零点 + let after: u64 = head(&r, "retry-after").unwrap().parse().unwrap(); + assert!((43_190..=43_200).contains(&after), "Retry-After {after}"); + let v: serde_json::Value = r.json().await.unwrap(); + let said = v["error"]["message"].as_str().unwrap(); + assert!(said.contains("`甲`, `乙`"), "{said}"); + assert_eq!(log.failed_code("第三个").await, "gw.wait.used_up"); + assert_eq!((a.hits(), b.hits()), (1, 1)); +} + +/// 回了 429 的那一跳上游没接:**不算进这家的请求数上限**,预留当场还回去。一分钟一个请求的 +/// 上游先回 429、暂停结束后接下:这个请求照常拿到回答,这一分钟只用了一个 +#[tokio::test] +async fn a_rate_limited_hop_does_not_count_toward_the_upstreams_limit() { + let up = upstream(vec![limited(300)], ok()).await; + let (gw, state, log) = serve( + config(&[(up.addr, "[{per: minute, requests: 1}]")], 10), + Ahead::new(), + ) + .await; + let r = ask(gw, "一个").await; + assert_eq!(r.status(), 200); + assert_eq!(up.hits(), 2); + assert_eq!(used(&state, "甲", 0), 1, "回了 429 的那一跳不算"); + let chain = log.routed("一个").await; + assert_eq!(chain.len(), 2, "{chain:#?}"); + assert_eq!(chain[0].status, Some(429)); + assert!(chain[0].pause_ms.is_some()); + assert!(!chain[0].counts_for_upstream()); + assert!(chain[1].counts_for_upstream()); +} + +/// 没发出去的、回了 429 换了下一家的:甲占着的都还回去,用量上限上什么都没留下 +#[tokio::test] +async fn a_reservation_is_given_back_when_the_hop_is_refused_or_never_sent() { + // 甲回 429(暂停很久),请求换乙 + let (a, b) = ( + upstream(vec![limited(60_000)], ok()).await, + upstream(Vec::new(), ok()).await, + ); + let limit = "[{per: minute, requests: 1}, {per: day, tokens: 1000000}]"; + let (gw, state, _log) = serve( + config(&[(a.addr, limit), (b.addr, limit)], 10), + Ahead::new(), + ) + .await; + assert_eq!(ask(gw, "一个").await.status(), 200); + assert_eq!((a.hits(), b.hits()), (1, 1)); + assert_eq!(used(&state, "甲", 0), 0); + assert_eq!(used(&state, "甲", 1), 0, "token 的预留也还回去"); + assert_eq!(used(&state, "乙", 0), 1); + + // 甲连不上:一个字节都没发出去 + let b = upstream(Vec::new(), ok()).await; + let (gw, state, _log) = serve( + config(&[(nobody().await, limit), (b.addr, limit)], 10), + Ahead::new(), + ) + .await; + assert_eq!(ask(gw, "一个").await.status(), 200); + assert_eq!(b.hits(), 1); + assert_eq!(used(&state, "甲", 0), 0); + assert_eq!(used(&state, "甲", 1), 0); + assert_eq!(used(&state, "乙", 0), 1); +} + +/// 一个请求在等分钟的窗口空出来,等的时候这一家今天的也用完了:**不再干等到期限**,当场回 +/// 「这一期用完了」的 429,到重置之前别重试 +#[tokio::test] +async fn a_request_waiting_for_a_window_stops_at_once_when_the_day_is_used_up_meanwhile() { + let up = upstream(Vec::new(), ok()).await; + let clock = Ahead::new(); + let limits = "[{per: minute, requests: 1}, {per: day, requests: 2}]"; + let (gw, state, log) = serve(config(&[(up.addr, limits)], 120), clock.clone()).await; + assert_eq!(ask(gw, "第一个").await.status(), 200); + // 这一分钟的窗口再过 30 秒空出来:这个请求等得到,排队 + clock.jump(30_000); + let t = Instant::now(); + let second = tokio::spawn(async move { ask(gw, "第二个").await }); + let id = log.id("第二个").await; + log.until("排队", |evs| { + evs.iter() + .any(|e| matches!(e, Event::RequestQueued { id: i, .. } if *i == id)) + .then_some(()) + }) + .await; + // 别处又发了一个:今天的两个用完了 + let cfg = state.config(); + let p = cfg.providers.iter().find(|p| p.name == "甲").unwrap(); + let _hold = + state + .upstream_limits + .reserve("甲", &p.limits, &tw_gateway::key_limits::Ask::default()); + let r = second.await.unwrap(); + assert!( + t.elapsed() < Duration::from_secs(3), + "该当场说用完了,却等了 {:?}", + t.elapsed() + ); + assert_eq!(r.status(), 429); + assert_eq!(head(&r, "x-should-retry").as_deref(), Some("false")); + assert_eq!(head(&r, "retry-after-ms"), None); + assert_eq!(log.failed_code("第二个").await, "gw.wait.used_up"); + assert_eq!(up.hits(), 1); + assert!( + log.all() + .iter() + .any(|e| matches!(e, Event::RequestDequeued { id: i, .. } if *i == id)), + "没报不排了" + ); +} diff --git a/crates/tw-gateway/tests/upstream_slots.rs b/crates/tw-gateway/tests/upstream_slots.rs index ecf43fd9..b042cae9 100644 --- a/crates/tw-gateway/tests/upstream_slots.rs +++ b/crates/tw-gateway/tests/upstream_slots.rs @@ -139,7 +139,7 @@ routes: - name: 全部 to: 池 failover: - slot_wait_secs: {wait_secs} + wait_secs: {wait_secs} ", a.addr, limit(limits.0), diff --git a/crates/tw-gateway/tests/ws.rs b/crates/tw-gateway/tests/ws.rs index d70d174a..06aff559 100644 --- a/crates/tw-gateway/tests/ws.rs +++ b/crates/tw-gateway/tests/ws.rs @@ -1128,6 +1128,7 @@ fn id_of(e: &Event) -> Option { | Event::RequestFinished { id, .. } | Event::RequestFailed { id, .. } | Event::RequestCancelled { id, .. } + | Event::RequestQueued { id, .. } | Event::ContentMatched { id, .. } => Some(*id), _ => None, } @@ -1332,7 +1333,7 @@ async fn a_used_up_key_limit_fails_the_turn_and_the_connection_stays_usable() { } /// 每分钟的上限:下一个空位在等得到的时候空出来,这一轮**等它**,然后照常发出;等不到的 -/// (`slot_wait_secs: 0`)当场替它回 `response.failed`,说清多久之后再来 +/// (`wait_secs: 0`)当场替它回 `response.failed`,说清多久之后再来 #[tokio::test] async fn a_rolling_key_limit_waits_within_the_turn_or_refuses_it() { let (up, seen) = responder(None).await; @@ -1365,7 +1366,7 @@ async fn a_rolling_key_limit_waits_within_the_turn_or_refuses_it() { up, |c| { c.clients[0].limits = serde_yaml_ng::from_str("[{per: minute, requests: 1}]").unwrap(); - c.failover.slot_wait_secs = 0; + c.failover.wait_secs = 0; }, Some(JumpClock::new()), ) @@ -1451,7 +1452,7 @@ async fn a_keys_max_concurrent_is_held_per_turn() { assert_eq!(frames.last().unwrap()["type"], "response.completed"); } -/// 这一家满着:这一轮**等它空出位置**(最多 `slot_wait_secs`),空出来了照常发,尝试链上记着 +/// 这一家满着:这一轮**等它空出位置**(最多 `wait_secs`),空出来了照常发,尝试链上记着 /// 等了多久;等不到替它回 `response.failed`(忙,可以重试),流量里留一行,连接照常 #[tokio::test] async fn a_turn_waits_for_a_full_upstream_and_fails_as_busy_when_none_frees() { @@ -1460,7 +1461,7 @@ async fn a_turn_waits_for_a_full_upstream_and_fails_as_busy_when_none_frees() { up, |c| { c.providers[0].max_concurrent = Some(1); - c.failover.slot_wait_secs = 1; + c.failover.wait_secs = 1; }, None, ) @@ -1533,6 +1534,70 @@ async fn a_turn_waits_for_a_full_upstream_and_fails_as_busy_when_none_frees() { assert_eq!(frames.last().unwrap()["type"], "response.completed"); } +/// 这一家回过 429、正暂停着:这一轮**等暂停结束再发**(`response.create` 推迟了,客户端照常拿到 +/// 回答),尝试链上记着等了多久,排队报一条;暂停到这一轮等不到的时候,当场替它回 +/// `response.failed`(限流,可以重试),连接照常 +#[tokio::test] +async fn a_turn_waits_for_a_paused_upstream_and_fails_when_the_pause_outlasts_it() { + let (up, seen) = responder(None).await; + let (gw, mut rx, state) = turns_gateway(up, |c| c.failover.wait_secs = 2, None).await; + let mut c = connect(gw).await; + let now = std::time::Instant::now; + state + .throttles + .limited("up", now(), Some(Duration::from_millis(400))); + let t = std::time::Instant::now(); + c.send(create("one")).await.unwrap(); + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(seen.lock().unwrap().is_empty(), "暂停着就发了"); + let frames = one_answer(&mut c).await; + assert_eq!(frames.last().unwrap()["type"], "response.completed"); + assert!( + t.elapsed() >= Duration::from_millis(380), + "{:?}", + t.elapsed() + ); + let evs = requests(&mut rx, 1).await; + assert!( + evs.iter().any(|e| matches!( + e, + Event::RequestQueued { upstreams, .. } if upstreams == &["up".to_string()] + )), + "{evs:?}" + ); + let queued = evs.iter().find_map(|e| match e { + Event::RequestRouted { attempts, .. } => attempts[0].queued_ms, + _ => None, + }); + assert!(queued.is_some_and(|ms| ms >= 150), "{evs:?}"); + + // 暂停一分钟,这一轮最多等两秒:不等,当场说 + state + .throttles + .limited("up", now(), Some(Duration::from_secs(60))); + let t = std::time::Instant::now(); + c.send(create("two")).await.unwrap(); + let frames = one_answer(&mut c).await; + assert!(t.elapsed() < Duration::from_secs(1), "{:?}", t.elapsed()); + let failed = frames.last().unwrap(); + assert_eq!(failed["type"], "response.failed", "{frames:?}"); + assert_eq!(failed["response"]["error"]["code"], "rate_limit_exceeded"); + assert_eq!(seen.lock().unwrap().len(), 1); + let evs = requests(&mut rx, 1).await; + let hop = evs + .iter() + .find_map(|e| match e { + Event::RequestRouted { attempts, .. } => Some(attempts[0].clone()), + _ => None, + }) + .unwrap_or_else(|| panic!("{evs:?}")); + assert_eq!(hop.skipped, Some(tw_api::ServeSkip::Throttled)); + match evs.last().unwrap() { + Event::RequestFailed { message, .. } => assert_eq!(message.code, "gw.wait.rate_limited"), + other => panic!("{other:?}"), + } +} + // ---------------------------------------------------------------- 快慢和成败 /// 一轮的剧本:收到第 n 个 `response.create`(从 1 数)时回的每一帧,和发它之前先等的毫秒数 @@ -1699,6 +1764,152 @@ async fn each_turn_feeds_the_upstreams_speed_and_success_like_an_http_hop() { assert_eq!(rate(), Some(7.0 / 8.0), "客户端走了被算进了成败"); } +/// 上游在一轮的第一帧就报了限流:这一家暂停,尝试链上那一跳带着暂停多久(`pause_ms`), +/// **不算进这家的用量上限** —— 内存里的数和记录里重新加的(`AttemptView::counts_for_upstream`) +/// 是同一个。内容到了之后才报的照常算一个请求 +#[tokio::test] +async fn a_turn_rate_limited_before_any_content_does_not_count_toward_the_upstreams_limit() { + fn script(n: usize) -> Vec<(u64, serde_json::Value)> { + let id = format!("resp_{n}"); + match n { + 1 => vec![(0, error("rate_limit_exceeded"))], + _ => vec![ + (0, created(&id)), + (0, delta()), + (0, failed(&id, "rate_limit_exceeded")), + ], + } + } + let up = scripted(script).await; + let (gw, mut rx, state) = turns_gateway( + up, + |c| { + c.providers[0].limits = + serde_yaml_ng::from_str("[{per: minute, requests: 5}]").unwrap(); + c.failover.wait_secs = 10; + // 退避的 5 秒截成 1 秒:测试不用干等 + c.failover.rate_limit_max_pause_secs = 1; + }, + None, + ) + .await; + let used = || { + let cfg = state.config(); + state.upstream_limits.view("up", &cfg.providers[0].limits)[0].used + }; + let mut c = connect(gw).await; + c.send(create("one")).await.unwrap(); + assert_eq!(until_end(&mut c).await.last().unwrap()["type"], "error"); + let evs = requests(&mut rx, 1).await; + let hop = evs + .iter() + .find_map(|e| match e { + Event::RequestRouted { attempts, .. } => Some(attempts[0].clone()), + _ => None, + }) + .unwrap_or_else(|| panic!("{evs:?}")); + assert!(hop.pause_ms.is_some_and(|ms| ms > 0), "{hop:?}"); + assert!(!hop.counts_for_upstream()); + assert_eq!(used(), 0, "限流了的这一轮不算"); + assert_eq!( + state.throttles.view("up").map(|t| t.kind), + Some(tw_api::ThrottleKind::RateLimited) + ); + + // 暂停结束后的下一轮:答了内容才报限流,算一个请求 + c.send(create("two")).await.unwrap(); + let frames = until_end(&mut c).await; + assert_eq!(frames.last().unwrap()["type"], "response.failed"); + let evs = requests(&mut rx, 1).await; + let hop = evs + .iter() + .find_map(|e| match e { + Event::RequestRouted { attempts, .. } => Some(attempts[0].clone()), + _ => None, + }) + .unwrap_or_else(|| panic!("{evs:?}")); + assert_eq!(hop.pause_ms, None, "{hop:?}"); + assert!(hop.counts_for_upstream()); + assert_eq!(used(), 1); +} + +/// 暂停结束后探路的那一轮:上游先回了 `response.created`,接着在内容之前报限流。**头一帧不说明 +/// 这一家好了**:排在后面的那一轮不能在 `response.created` 到的那一刻就跟着发,要等这一家再次 +/// 暂停完 +#[tokio::test] +async fn a_probing_turn_limited_after_response_created_keeps_the_next_turn_waiting() { + let times: Arc>> = Arc::default(); + let t = times.clone(); + let app = Router::new().route( + "/backend-api/codex/responses", + axum::routing::any(move |ws: WebSocketUpgrade| { + let times = t.clone(); + async move { + ws.on_upgrade(move |mut sock| async move { + while let Some(Ok(m)) = sock.recv().await { + if !matches!(m, Message::Text(_)) { + continue; + } + let n = { + let mut t = times.lock().unwrap(); + t.push(std::time::Instant::now()); + t.len() + }; + let id = format!("resp_{n}"); + // 头一个(探路的):开头当场到,300 毫秒后报限流 + let frames = if n == 1 { + vec![(0, created(&id)), (300, error("rate_limit_exceeded"))] + } else { + vec![(0, created(&id)), (0, delta()), (0, completed(&id))] + }; + for (wait, f) in frames { + tokio::time::sleep(Duration::from_millis(wait)).await; + if sock + .send(Message::Text(f.to_string().into())) + .await + .is_err() + { + return; + } + } + } + }) + } + }), + ); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let up = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + let (gw, _rx, state) = turns_gateway( + up, + // 退避截成 1 秒:测试不用干等 + |c| c.failover.rate_limit_max_pause_secs = 1, + None, + ) + .await; + state.throttles.limited( + "up", + std::time::Instant::now(), + Some(Duration::from_millis(300)), + ); + let (mut a, mut b) = (connect(gw).await, connect(gw).await); + a.send(create("one")).await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + b.send(create("two")).await.unwrap(); + assert_eq!(until_end(&mut a).await.last().unwrap()["type"], "error"); + assert_eq!( + until_end(&mut b).await.last().unwrap()["type"], + "response.completed" + ); + let times = times.lock().unwrap().clone(); + assert_eq!(times.len(), 2); + assert!( + times[1] - times[0] >= Duration::from_millis(1_000), + "探路的那一轮在 response.created 之后被限流了,后面的却跟着发了:{:?}", + times[1] - times[0] + ); +} + /// 一轮回到一半上游就收了连接:帧里写着 `CLOSE` 的,回了开头之后发关闭帧;写着 `DROP` 的, /// 回了开头之后直接断开(没有关闭帧);写着 `HOLD` 的,回了开头之后一直不说完。别的照常答完 async fn closing() -> SocketAddr { @@ -1891,7 +2102,7 @@ async fn a_turns_key_limit_wait_and_slot_wait_share_one_budget() { |c| { c.clients[0].limits = serde_yaml_ng::from_str("[{per: minute, requests: 1}]").unwrap(); c.providers[0].max_concurrent = Some(1); - c.failover.slot_wait_secs = 2; + c.failover.wait_secs = 2; }, Some(clock.clone()), ) diff --git a/crates/tw-observe/src/bus.rs b/crates/tw-observe/src/bus.rs index 8925909e..690667ff 100644 --- a/crates/tw-observe/src/bus.rs +++ b/crates/tw-observe/src/bus.rs @@ -99,7 +99,10 @@ fn about_the_request(ev: &tw_api::Event) -> bool { | E::Translated { .. } | E::SecretsFound { .. } | E::ContentMatched { .. } - | E::ToolCallFlagged { .. } => true, + | E::ToolCallFlagged { .. } + // 在排队:半路才来的一方要看得出它在等哪几家,和它是不是已经不排了 + | E::RequestQueued { .. } + | E::RequestDequeued { .. } => true, // 开始和三种结局由 `track_at` 自己管 E::RequestStarted { .. } | E::RequestFinished { .. } @@ -119,6 +122,8 @@ fn about_the_request(ev: &tw_api::Event) -> bool { // 插件出错是一条通知:它的号是新取的,出错的请求在 `request_id` 里 | E::PluginFailed { .. } | E::HealthChanged { .. } + // 一家上游暂时发不出去:说的是那一家现在的样子(`/overview` 答得了) + | E::ThrottleChanged { .. } | E::ModelsChanged { .. } | E::ProxyChanged { .. } | E::AuthChanged { .. } @@ -502,6 +507,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }, tw_api::AttemptView { @@ -514,6 +520,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }, ], diff --git a/crates/tw-pack/src/model/tests.rs b/crates/tw-pack/src/model/tests.rs index 79c0af6c..2df5eb61 100644 --- a/crates/tw-pack/src/model/tests.rs +++ b/crates/tw-pack/src/model/tests.rs @@ -49,6 +49,10 @@ providers: pricing: 中转 models: [m1] max_concurrent: 4 + limits: + - per: hour + every: 5 + requests: 100 - name: chatgpt base_url: https://chatgpt.com/backend-api/codex oauth: @@ -206,7 +210,11 @@ fn an_upstream_splits_into_four_units_by_what_its_fields_are_about() { ); assert_eq!(relay.units["models"], json!({"models": ["m1"]})); assert_eq!(relay.units["billing"], json!({"pricing": "中转"})); - assert_eq!(relay.units["other"], json!({"max_concurrent": 4})); + // 用量上限跟着上游走(分享、导出、同步都带着),和并发上限在同一个单位里 + assert_eq!( + relay.units["other"], + json!({"max_concurrent": 4, "limits": [{"per": "hour", "every": 5, "requests": 100}]}) + ); // 写明的默认值(proxy: direct)和没写一样;请求头按写的顺序 let official = &s.items[&id(ItemKind::Provider, "官方")]; assert_eq!( diff --git a/crates/tw-store/src/db.rs b/crates/tw-store/src/db.rs index bde21467..b059be96 100644 --- a/crates/tw-store/src/db.rs +++ b/crates/tw-store/src/db.rs @@ -192,6 +192,25 @@ pub struct KeyUsage { pub cost_micros: i64, } +/// 一家上游从某一刻起用了多少(见 [`Db::upstream_usage_since`])。 +/// +/// **请求数按尝试链数,用量按接下它的那一家**:和上游的用量上限结算一行时 +/// (`tw_gateway::upstream_limits`)是同一个判断。按路径分开:数 token 的请求在网关那边剔掉 +#[derive(Debug, Clone, Default, PartialEq)] +pub struct UpstreamUsage { + pub upstream: String, + pub path: String, + /// 算数的那几跳([`tw_api::RoutingView::upstream_hops`])里发给这家的有几跳 + pub requests: i64, + /// 这家接下的那些请求([`tw_api::RoutingView::served_by`])的用量合计 + pub input_tokens: i64, + pub output_tokens: i64, + pub cache_read_tokens: i64, + pub cache_write_tokens: i64, + /// 记下的费用合计,微分。算不出来的那几行算 0 + pub cost_micros: i64, +} + /// 装上 `tw_reached(路由那一列, 失败没有)`:这一行的请求可能发到了上游 /// ([`tw_api::RoutingView::reached_upstream`])。**判断写在 tw-api 那一处**,记下一行时交给 /// 网关的([`crate::Settled::reached`])和从库里加回来的是同一个。路由那一列读不出来的当作 @@ -1528,6 +1547,58 @@ impl Db { rows.collect::, _>>().map_err(Into::into) } + /// 每家上游从 `since_ms` 起用了多少:重启之后(和一期的开头变了的时候),上游的用量上限 + /// 按它把这一天、这一周、这个月的数加回来。 + /// + /// **尝试链在路由那一列的 JSON 里**,一行一行读出来数(和 `tw_reached` 一样按 tw-api 的 + /// 判断):算数的每一跳给那一家记一个请求,这一行的用量记在接下它的那一家。网关自己答的、 + /// 没有路由那一列的(没走到上游)不在里面;路由那一列读不出来的当作尝试链是空的 + pub fn upstream_usage_since(&self, since_ms: i64) -> Result, DbError> { + let mut st = self.conn.prepare( + "SELECT path, routing, + COALESCE(input_tokens, 0), COALESCE(output_tokens, 0), + COALESCE(cache_read_tokens, 0), COALESCE(cache_write_tokens, 0), + COALESCE(cost_micros, 0) + FROM requests + WHERE at_ms >= ?1 AND local = 0 AND routing IS NOT NULL", + )?; + let mut rows = st.query(params![since_ms])?; + let mut by: std::collections::BTreeMap<(String, String), UpstreamUsage> = + std::collections::BTreeMap::new(); + /// 一家、一个路径的那一组 + fn entry<'a>( + by: &'a mut std::collections::BTreeMap<(String, String), UpstreamUsage>, + upstream: &str, + path: &str, + ) -> &'a mut UpstreamUsage { + by.entry((upstream.to_string(), path.to_string())) + .or_insert_with(|| UpstreamUsage { + upstream: upstream.to_string(), + path: path.to_string(), + ..Default::default() + }) + } + while let Some(r) = rows.next()? { + let path: String = r.get(0)?; + let Ok(routing) = serde_json::from_str::(&r.get::<_, String>(1)?) + else { + continue; + }; + for up in routing.upstream_hops() { + entry(&mut by, &up, &path).requests += 1; + } + if let Some(up) = routing.served_by() { + let u = entry(&mut by, up, &path); + u.input_tokens += r.get::<_, i64>(2)?; + u.output_tokens += r.get::<_, i64>(3)?; + u.cache_read_tokens += r.get::<_, i64>(4)?; + u.cache_write_tokens += r.get::<_, i64>(5)?; + u.cost_micros += r.get::<_, i64>(6)?; + } + } + Ok(by.into_values().collect()) + } + /// 各条路由走了多少请求、各条规则命中了多少,以及记录从哪一刻起是全的(见 /// [`tw_api::RouteStats`])。 pub fn route_stats(&self, since_ms: i64, until_ms: i64) -> Result { @@ -2467,6 +2538,71 @@ pub(crate) mod tests { ); } + #[test] + fn upstream_usage_counts_hops_and_puts_usage_on_the_upstream_that_answered() { + let db = Db::in_memory().unwrap(); + let t0 = 1_000_000_000i64; + // 头一跳 a 超时(发出去了,算),第二跳 b 回 429(不算),第三跳 c 接下 + let failover = r#"{"route":"default","rule":"r","rewritten_by":[],"attempts":[ + {"provider":"a","outcome":"error","error":{"code":"gw.upstream.timeout","args":{},"text":"x"},"ms":5}, + {"provider":"b","outcome":"status","status":429,"ms":5}, + {"provider":"c","outcome":"served","status":200,"ms":5}]}"#; + let served = r#"{"route":"default","rule":"r","rewritten_by":[], + "attempts":[{"provider":"c","outcome":"served","status":200,"ms":5}]}"#; + // 都满着:跳过的那一跳没发出去 + let busy = r#"{"route":"default","rule":"r","rewritten_by":[],"attempts":[{"provider":"c", + "outcome":"error","error":{"code":"gw.busy_upstream","args":{},"text":"x"}, + "ms":0,"skipped":"busy"}]}"#; + let mut one = row(1, t0); + one.routing = Some(failover.into()); + let mut two = row(2, t0 + 1); + two.routing = Some(served.into()); + two.cost_micros = None; + let mut early = row(3, t0 - 1); + early.routing = Some(served.into()); + let mut skipped = row(4, t0 + 2); + skipped.routing = Some(busy.into()); + let mut local = row(5, t0 + 3); + local.routing = Some(served.into()); + local.local = true; + let mut count = row(6, t0 + 4); + count.routing = Some(served.into()); + count.path = "/v1/messages/count_tokens".into(); + let no_routing = row(7, t0 + 5); + for r in [&one, &two, &early, &skipped, &local, &count, &no_routing] { + db.insert(r).unwrap(); + } + let got = db.upstream_usage_since(t0).unwrap(); + let find = |up: &str, path: &str| { + got.iter() + .find(|u| u.upstream == up && u.path == path) + .cloned() + }; + assert_eq!(got.len(), 3, "{got:?}"); + let a = find("a", "/v1/messages").unwrap(); + assert_eq!( + (a.requests, a.input_tokens), + (1, 0), + "超时的那一跳算一个请求,用量不是它的" + ); + assert!(find("b", "/v1/messages").is_none(), "回 429 的不算"); + let c = find("c", "/v1/messages").unwrap(); + assert_eq!( + ( + c.requests, + c.input_tokens, + c.output_tokens, + c.cache_read_tokens, + c.cache_write_tokens, + c.cost_micros + ), + (2, 2000, 1000, 400, 0, 12_000), + "早于那一刻的、本地答的、跳过的不算;算不出钱的算 0" + ); + // 数 token 的按路径分开交出去,网关那边剔掉 + assert_eq!(find("c", "/v1/messages/count_tokens").unwrap().requests, 1); + } + #[test] fn a_window_narrows_the_history_and_no_window_means_the_most_recent() { let db = Db::in_memory().unwrap(); @@ -3496,6 +3632,7 @@ mod cost_state_tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], ..Default::default() @@ -3996,6 +4133,7 @@ mod cost { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, } }) diff --git a/crates/tw-store/src/lib.rs b/crates/tw-store/src/lib.rs index 9f5cb492..820bc3d9 100644 --- a/crates/tw-store/src/lib.rs +++ b/crates/tw-store/src/lib.rs @@ -22,6 +22,7 @@ pub mod transcript; pub use blobs::{Blobs, Which}; pub use db::{ Db, DbError, KeyUsage, Latency, PluginRunRow, RequestRow, SecurityEvent, Summary, TokenRate, + UpstreamUsage, }; pub use recorder::{Recorder, SettleHook, Settled, price_source}; pub use task::StoredBody; diff --git a/crates/tw-store/src/recorder.rs b/crates/tw-store/src/recorder.rs index e94b1dcc..dac196c8 100644 --- a/crates/tw-store/src/recorder.rs +++ b/crates/tw-store/src/recorder.rs @@ -105,7 +105,8 @@ pub struct Recorder { settled: Option, } -/// 记下的一行里,密钥的用量上限要的那几样(见 `tw_gateway::key_limits`)。 +/// 记下的一行里,密钥和上游的用量上限要的那几样(见 `tw_gateway::key_limits`、 +/// `tw_gateway::upstream_limits`)。 /// /// **和库里那一行是同一份数**:用量、费用就是写进去的那些,重启之后从库里加回来的和 /// 此刻结算的一样多。 @@ -128,6 +129,10 @@ pub struct Settled { pub cache_write: u64, /// 记下的费用,微分。算不出来的(没有价格、没有用量)是 None pub cost_micros: Option, + /// 算进上游用量上限的那几跳各发给了哪一家([`tw_api::RoutingView::upstream_hops`]) + pub hops: Vec, + /// 这一行的用量记在哪一家上游的用量上限上([`tw_api::RoutingView::served_by`]) + pub served: Option, } /// 结算往哪儿交。由 twcore 接到网关的密钥用量上限上(`tw_control::key_limits`) @@ -687,6 +692,11 @@ impl Recorder { // 某家上游熔断了 —— 是「现在什么情况」,不是「刚才发生过什么」。 // 这张表只装后者。 | Event::HealthChanged { .. } + // 一家上游暂时发不出去,同上:是现状 + | Event::ThrottleChanged { .. } + // 一个请求开始排队:它等了多久在尝试链(`queued_ms`)里,这一条不用落库 + | Event::RequestQueued { .. } + | Event::RequestDequeued { .. } | Event::ModelsChanged { .. } | Event::ProxyChanged { .. } | Event::AuthChanged { .. } @@ -817,6 +827,8 @@ impl Recorder { cache_read: u.map_or(0, |u| u.cache_read), cache_write: u.map_or(0, |u| u.cache_write), cost_micros, + hops: p.routing.upstream_hops(), + served: p.routing.served_by().map(str::to_string), }); } self.write(RequestRow { @@ -978,6 +990,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }, tw_api::AttemptView { @@ -990,6 +1003,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: Some("机场".into()), }, ], @@ -1055,6 +1069,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing: tw_api::Billing::PerToken, @@ -1089,6 +1104,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }, tw_api::AttemptView { @@ -1101,6 +1117,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }, ], @@ -1154,6 +1171,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing: tw_api::Billing::PerToken, @@ -1282,6 +1300,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing: tw_api::Billing::PerToken, @@ -1608,6 +1627,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing: tw_api::Billing::Free, @@ -1650,6 +1670,7 @@ mod tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing: tw_api::Billing::Free, @@ -1731,6 +1752,7 @@ mod billing_tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing, @@ -1823,6 +1845,7 @@ mod billing_tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing, @@ -2269,6 +2292,7 @@ mod security_detail_tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing: tw_api::Billing::PerToken, @@ -2564,6 +2588,7 @@ mod translation_tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }) .collect(), @@ -3256,6 +3281,7 @@ mod settle_hook_tests { usage: None, queued_ms: None, skipped: None, + pause_ms: None, proxy: None, }], billing: tw_api::Billing::PerToken, diff --git a/docs/config.md b/docs/config.md index f0788e96..31f55de8 100644 --- a/docs/config.md +++ b/docs/config.md @@ -350,9 +350,11 @@ minute, hour, day, week or month. Each entry counts one of the three; a key can have several, and a request has to pass every one. `minute` and `hour` are rolling: the last 60 seconds, the last 60 minutes. +With `every`, the window spans several of them: `{ per: hour, every: 5 }` is +the last 5 hours. When one is used up, a request waits for the next free slot if it frees -within `failover.slot_wait_secs` (30 seconds by default), and is refused -otherwise. Any later wait for a busy upstream comes out of the same time. +within `failover.wait_secs` (120 seconds by default), and is refused +otherwise. Any later wait for an upstream comes out of the same time. `day`, `week` and `month` follow the calendar in the time zone of the machine twcore runs on and start again at midnight, on Monday and on the 1st. When one is used up, requests are refused until it starts again. If the machine's time zone changes, the current day, week and month are added up @@ -383,7 +385,8 @@ all its answers count when it closes. | Field | Type | Default | Description | |---|---|---|---| -| `per` | `minute` \| `hour` \| `day` \| `week` \| `month` | **required** | The period. `minute` and `hour` are rolling (the last 60 seconds, the last 60 minutes); `day`, `week` and `month` start again at local midnight, on Monday and on the 1st. | +| `per` | `minute` \| `hour` \| `day` \| `week` \| `month` | **required** | The period. `minute` and `hour` are rolling (the last 60 seconds, the last 60 minutes, or `every` of them); `day`, `week` and `month` start again at local midnight, on Monday and on the 1st. | +| `every` | integer | — | How many minutes or hours the rolling window spans: `{per: hour, every: 5}` is the last 5 hours. From 1 to 60 for `minute`, from 1 to 24 for `hour`; `day`, `week` and `month` do not take it. Unset: 1. | | `requests` | integer | — | At most this many requests. Token counts, answers the gateway gives itself and requests that never reach an upstream do not count. | | `tokens` | integer | — | At most this many tokens: uncached input, cache writes and output. | | `cost` | number | — | At most this much, in US dollars, as recorded for each request; at least 0.01. Models without a price and upstreams with `billing: free` count as 0. | @@ -424,7 +427,8 @@ Upstreams: the APIs requests are forwarded to. | `billing` | `per-token` \| `free` | `per-token` | `per-token`: cost is usage times the price in the upstream's price sheet, subscription accounts included. `free`: cost is recorded as 0. | | `pricing` | string | — | Name of a price sheet under `pricing.sheets`. Unset: the default price table. | | `model_specs` | map of model id → [`providers[].model_specs.*`](#cfg-providers-model_specs) | `{}` | Context window, output limit, reasoning and image input of single models of this upstream, written by hand, by exact model id. They take precedence over the price table: for models it does not know, or gets wrong. | -| `max_concurrent` | integer | — | Most requests sent to this upstream at the same time, from 1 to 1000. When it is full, a conversation that stays on it waits for a free slot and other requests go to the next upstream; see `failover.slot_wait_secs`. Unset: no limit. | +| `max_concurrent` | integer | — | Most requests sent to this upstream at the same time, from 1 to 1000. When it is full, a conversation that stays on it waits for a free slot and other requests go to the next upstream; see `failover.wait_secs`. Unset: no limit. | +| `limits` | list of [`providers[].limits[]`](#cfg-providers-limits) | — | Usage limits of this upstream, written like its own quota: requests, tokens or cost every few minutes or hours, per day, week or month. Every one has to have room before a request is sent to it. When a rolling window is full, requests go to the next upstream or wait for room; when a day, week or month is used up, requests go to the next upstream until the next one starts. Unset: no limit. | | `disabled` | bool | `false` | Take the upstream out of routing and out of the model list, and keep its configuration. | | `balance` | `auto` \| `off` \| `openrouter` \| `deepseek` \| `moonshot` \| `sub2api` \| `newapi` \| `thinkwatch` | `auto` | Where the upstream's balance is read from, with its own key and proxy: `openrouter`, `deepseek`, `moonshot`, `sub2api`, `newapi` (New API and One API) or `thinkwatch` (a ThinkWatch Enterprise gateway). `auto`: recognized from `base_url` for OpenRouter, DeepSeek and Moonshot; any other address, except Anthropic's, OpenAI's and Gemini's own, is asked once whether it is a Sub2API, ThinkWatch Enterprise or New API relay, and the answer is kept in memory only. `off`: no balance. ChatGPT and Z.ai accounts and Bedrock have no balance whatever is written here. | | `signed_in` | `zai` \| `bigmodel` | — | Which account sign-in the `key` came from: `zai` (Z.ai) or `bigmodel` (BigModel). Signing in to the account in the desktop app writes it; replacing the key there removes it. It does not change how requests are sent. Unset: the key was written by hand. | @@ -507,12 +511,59 @@ slot on the upstream when it is sent, and gives it back when the answer has been passed on in full or the client has gone. When the upstream is full, a conversation that stays on it to reuse its prompt cache waits for a slot; any other request goes straight to the next upstream. How long a request -waits is `failover.slot_wait_secs`. Waiting is not a failure: the upstream is +waits is `failover.wait_secs`. Waiting is not a failure: the upstream is not set aside. Requests that only count tokens do not take a slot. On a Responses WebSocket connection, each `response.create` takes a slot from when it is sent until its answer ends, and so does a key's `max_concurrent`; an idle connection takes none. +#### `providers[].limits` + +Usage limits for an upstream, written like the quota the upstream itself +enforces: at most so many requests, tokens or dollars every few minutes or +hours, per day, week or month. The entries are written like a key's `limits`; +an upstream can have several, and every one has to have room before a request +is sent to it. + +Each attempt sent to the upstream counts as a request, whether it succeeds or +fails, unless the upstream answers it with 429. Tokens and cost count for the +attempt that is answered. Requests that only count tokens do not count. An +attempt still running counts with an estimate of its input until it is +recorded. + +When a `minute` or `hour` window is full, the upstream is passed over like a +busy one: a conversation that stays on it waits for room, other requests go to +the next upstream, and when no candidate has room a request waits for the +first one that does, within the time a request may wait (see `failover`). When +a `day`, `week` or `month` is used up, requests go to the next upstream until +the period starts again; none waits for it. Being full is not a failure: the +upstream is not set aside. After a restart the day, week and month are added up +again from the request records, so `retention.row_days` has to cover the +period, as for a key; minute and hour windows start empty. + + + + +| Field | Type | Default | Description | +|---|---|---|---| +| `per` | `minute` \| `hour` \| `day` \| `week` \| `month` | **required** | The period. `minute` and `hour` are rolling (the last 60 seconds, the last 60 minutes, or `every` of them); `day`, `week` and `month` start again at local midnight, on Monday and on the 1st. | +| `every` | integer | — | How many minutes or hours the rolling window spans: `{per: hour, every: 5}` is the last 5 hours. From 1 to 60 for `minute`, from 1 to 24 for `hour`; `day`, `week` and `month` do not take it. Unset: 1. | +| `requests` | integer | — | At most this many requests sent to this upstream. Each attempt that reaches it counts, unless it answers 429; token counts do not count. | +| `tokens` | integer | — | At most this many tokens in the answers of this upstream: uncached input, cache writes and output. | +| `cost` | number | — | At most this much, in US dollars, as recorded for the requests this upstream answers; at least 0.01. Models without a price count as 0, and so does everything with `billing: free`. | +| `cache_reads` | bool | `false` | Count cache reads too. Only for a `tokens` limit. | + + +```yaml +providers: + - name: relay + base_url: https://relay.example.com + key: sk-relay + limits: + - { per: hour, every: 5, requests: 200 } + - { per: week, cost: 50 } +``` + #### `providers[].oauth` @@ -1071,9 +1122,9 @@ mode, `off` included. An upstream that fails is set aside for a while, so that the next requests go straight to the next candidate. How long depends on the reason the -upstream gives: an insufficient balance waits for a top-up, a used-up quota -waits until the moment the upstream says it resets, and a rate limit usually -passes within seconds. A request with a single candidate is never affected. +upstream gives: an insufficient balance waits for a top-up, and a used-up +quota waits until the moment the upstream says it resets. A request with a +single candidate is never affected. On a Responses WebSocket connection, each `response.create` counts here like a request: one that fails because of the upstream before any content arrives counts as a failure, and so does a connection the upstream refuses or that @@ -1122,16 +1173,39 @@ failover: idle_timeout_secs: 600 ``` -When upstreams are at their `max_concurrent`, a request waits for a free slot -for at most `slot_wait_secs` in all. The same time also covers waiting for a -key's `minute` or `hour` limit, so a request never waits longer than -`slot_wait_secs` for the two together. An ongoing conversation waits for the -upstream it stays on and, if no slot frees in time, moves on to the next one, -where its cache starts over. A new conversation skips a full upstream at once. -When every candidate is full, the request waits for whichever frees first; if -none does, the client gets a 429 with `Retry-After` saying the upstreams are -busy. If an upstream did receive the request and failed, the client gets that -failure instead. +An upstream that answers 429 (rate limited) is not set aside but paused: +for as long as it says (`retry-after-ms`, `Retry-After`, or a delay in the +error body; under one second, or under one millisecond for `retry-after-ms`, +counts as not said), and otherwise for 5, 10, 20 and 40 seconds, then 60 +seconds each time, by its consecutive 429s; any other answer starts the count +again. The +pause is at most `rate_limit_max_pause_secs`, and the circuit breaker and +`load-balance` do not count it as a failure. The request that got the 429 +moves to the next candidate, or waits for the pause to end; the client sees +nothing of it. A streamed answer that reports a rate limit before any content +counts as a 429 too; only on a request's last candidate that has not been rate +limited in the last five minutes does the client get that error (the upstream +is paused all the same). When the pause ends, one request goes first, and the +others follow once the upstream has accepted it (for a streamed answer, once +its content starts), or at once if that request gets no answer, times out or +finds the upstream's quota used up. A pause lasts until twcore restarts at +most. + +A request waits for upstreams that cannot take it yet for at most `wait_secs` +in all: an upstream at its `max_concurrent`, one paused after a 429, and one +at its `minute` or `hour` usage limit. The same time also covers waiting for +a key's `minute` or `hour` limit, so a request never waits longer than +`wait_secs` for all of them together. An ongoing conversation waits for the +upstream it stays on and, if it does not become available in time, moves on +to the next one, where its cache starts over. A new conversation skips such an +upstream at once. An upstream that has used up a day, week or month usage +limit is skipped. When no candidate is left, the request waits for whichever +becomes available first; streamed answers keep the client's connection open +with the `200` and keep-alives described above. If none becomes available, the +client gets a 429 with `Retry-After` and `retry-after-ms` set to when the +first one will; when every candidate has used up its period, `Retry-After` +is the earliest reset and `x-should-retry` is `false`. If an upstream did +receive the request and failed, the client gets that failure instead. @@ -1143,9 +1217,9 @@ failure instead. | `max_pause_secs` | integer | `600` | Upper bound on the doubled pause, in seconds; not less than `pause_secs`. | | `no_balance_pause_secs` | integer | `1800` | Seconds to set aside an upstream that reports an insufficient balance. | | `quota_pause_secs` | integer | `3600` | Seconds to set aside an upstream whose quota is used up when it does not say when the quota resets. When it does, the upstream is set aside until then. | -| `rate_limit_max_pause_secs` | integer | `3600` | A rate-limited upstream is set aside for the time its `Retry-After` gives, at most this many seconds. Without `Retry-After` it counts as a failure without a stated reason. | +| `rate_limit_max_pause_secs` | integer | `3600` | Longest pause, in seconds, after an upstream answers 429 (rate limited). The pause lasts as long as the upstream says (`retry-after-ms`, `Retry-After`, or a delay in the error body); when it says nothing, it backs off by the upstream's consecutive 429s: 5, 10, 20 and 40 seconds, then 60 seconds each time. A pause is not a failure: requests wait for the upstream within `wait_secs`, and the circuit breaker does not count it. | | `idle_timeout_secs` | integer | `300` | Seconds an upstream may go without sending content before the gateway stops waiting for it. Counted from the moment the request is sent and started again by every piece of content: text, reasoning and tool calls count, keep-alives do not. A whole (non-streamed) answer counts from sending to the complete answer. Before any content has reached the client, the upstream counts as failed and the request moves to the next one; with none left, the client gets a timeout error. After content has reached the client, the answer ends with an error. From 30 to 3600. | -| `slot_wait_secs` | integer | `30` | Seconds a request waits in all, counted once the key's own `max_concurrent` lets it in: for a key's `minute` or `hour` limit to free up, and for a free slot on upstreams at their `max_concurrent`. A key limit that does not free up in time refuses the request; without an upstream slot in time it goes to the next upstream, or, when every candidate is full, is answered with 429. `0`: never wait. From 0 to 300. | +| `wait_secs` | integer | `120` | Seconds a request waits in all, counted once the key's own `max_concurrent` lets it in: for a key's `minute` or `hour` limit to free up, and for an upstream that cannot take it yet — at its `max_concurrent`, paused after answering 429, or at one of its `minute` or `hour` usage limits. A key limit that does not free up in time refuses the request. An upstream that cannot take it is skipped for the next one, and when none is left the request waits for the first to become available; a conversation that stays on an upstream waits for it. When nothing became available in time, the client gets 429 with `Retry-After`. `0`: never wait. From 0 to 600. | ### `aliases` @@ -1257,12 +1331,13 @@ group shares out requests by the result in the same way as above. group has its weight multiplied by four, by at most ten and at least a tenth. - `health`: upstreams that fail less get a larger share. It looks at the - last 50 requests within the past 30 minutes. Server errors, rate limits, - used-up quota or balance, rejected credentials, timeouts (including + last 50 requests within the past 30 minutes. Server errors, used-up quota + or balance, rejected credentials, timeouts (including `failover.idle_timeout_secs` before any content) and connection errors count as failures; errors caused by the request itself do not, and neither - does a client that cancels, a request aborted by hand, or an upstream - skipped because it is at its `max_concurrent`. An upstream that keeps failing keeps a twentieth of its + does a rate limit (429), a client that cancels, a request aborted by hand, + or an upstream skipped because it is at its `max_concurrent` or at a usage + limit. An upstream that keeps failing keeps a twentieth of its weight, so it still gets the occasional request and its recovery is noticed; one that fails outright is set aside by [`failover`](#cfg-failover) as before. diff --git a/docs/config.zh-CN.md b/docs/config.zh-CN.md index 751d2ed5..529f4b55 100644 --- a/docs/config.zh-CN.md +++ b/docs/config.zh-CN.md @@ -243,9 +243,10 @@ clients: 一把密钥的用量上限:每分钟、每小时、每天、每周或每月最多多少个请求、多少 token、多少美元。 每一条只数其中一种;一把密钥可以有好几条,请求要通过每一条。 -`minute`、`hour` 是滚动的:最近 60 秒、最近 60 分钟。用满之后,下一个空位在 -`failover.slot_wait_secs`(默认 30 秒)之内空出来,请求就等它;等不到就拒绝。之后若还要等 -上游空出位置,用的是同一段时间里剩下的部分。`day`、`week`、`month` 按 twcore 所在机器的 +`minute`、`hour` 是滚动的:最近 60 秒、最近 60 分钟;写了 `every` 就是几个这样的单位, +`{ per: hour, every: 5 }` 是最近 5 小时。用满之后,下一个空位在 +`failover.wait_secs`(默认 120 秒)之内空出来,请求就等它;等不到就拒绝。之后若还要等 +上游,用的是同一段时间里剩下的部分。`day`、`week`、`month` 按 twcore 所在机器的 本地时区算,在零点、周一零点、每月一号零点重新开始;用满之后,到重新开始之前的请求一律拒绝。 机器换了时区时,当前这一天、这一周、这个月的用量按新的时区从请求记录里重新加起来。 @@ -267,7 +268,8 @@ Responses 的 WebSocket 连接上,每个 `response.create` 各算一个请求 | 字段 | 类型 | 默认值 | 说明 | |---|---|---|---| -| `per` | `minute` \| `hour` \| `day` \| `week` \| `month` | **必填** | 按多长一段时间算。`minute`、`hour` 是滚动的(最近 60 秒、最近 60 分钟);`day`、`week`、`month` 在本地时间的零点、周一零点、每月一号零点重新算。 | +| `per` | `minute` \| `hour` \| `day` \| `week` \| `month` | **必填** | 按多长一段时间算。`minute`、`hour` 是滚动的(最近 60 秒、最近 60 分钟,或者 `every` 个这样的单位);`day`、`week`、`month` 在本地时间的零点、周一零点、每月一号零点重新算。 | +| `every` | 整数 | — | 滚动窗口有几分钟、几小时:`{per: hour, every: 5}` 是最近 5 小时。`minute` 取 1 到 60,`hour` 取 1 到 24;`day`、`week`、`month` 不能写。不写:1。 | | `requests` | 整数 | — | 最多这么多个请求。数 token 的请求、网关自己答的、没有发到上游的不算。 | | `tokens` | 整数 | — | 最多这么多 token:未命中缓存的输入、写入缓存的和输出。 | | `cost` | 数字 | — | 最多花这么多美元,按每个请求记下的费用算;至少 0.01。没有价格的模型、`billing: free` 的上游算 0。 | @@ -308,7 +310,8 @@ clients: | `billing` | `per-token` \| `free` | `per-token` | `per-token`:费用为用量乘以所选价目表中的单价,订阅账号同样如此。`free`:费用记为 0。 | | `pricing` | 字符串 | — | `pricing.sheets` 中某张价目表的名字。不写:默认价目表。 | | `model_specs` | 映射: 模型 ID → [`providers[].model_specs.*`](#cfg-providers-model_specs) | `{}` | 手写这家上游某些模型的上下文窗口、输出上限、会不会推理、收不收图,按模型 ID 完全匹配。写了就优先于价目表,用于价目表里没有或写错的模型。 | -| `max_concurrent` | 整数 | — | 同时发给这家的请求最多几个,取值 1 到 1000。满了的时候,留在这家的对话等空位,别的请求换下一家;等多久见 `failover.slot_wait_secs`。不写:不限。 | +| `max_concurrent` | 整数 | — | 同时发给这家的请求最多几个,取值 1 到 1000。满了的时候,留在这家的对话等空位,别的请求换下一家;等多久见 `failover.wait_secs`。不写:不限。 | +| `limits` | 对象列表,见 [`providers[].limits[]`](#cfg-providers-limits) | — | 这家上游的用量上限,照着它自己的限额写:每几分钟、每几小时、每天、每周或每月的请求数、token 数或费用。每一条都有余量才往这家发。滚动窗口满了,请求换下一家或者等它空出来;天、周、月用满了,到下一期之前请求换下一家。不写:不限。 | | `disabled` | 布尔 | `false` | 不参与路由,模型也不出现在模型列表里;配置原样保留。 | | `balance` | `auto` \| `off` \| `openrouter` \| `deepseek` \| `moonshot` \| `sub2api` \| `newapi` \| `thinkwatch` | `auto` | 从哪里读这家上游的余额,用它自己的密钥、走它自己的代理:`openrouter`、`deepseek`、`moonshot`、`sub2api`、`newapi`(New API 和 One API)或 `thinkwatch`(ThinkWatch 企业网关)。`auto`:OpenRouter、DeepSeek、Moonshot 按 `base_url` 识别;其他地址(Anthropic、OpenAI、Gemini 的官方地址除外)询问一次它是不是 Sub2API、ThinkWatch 企业网关或 New API 中转站,结果只保存在内存中。`off`:不读余额。ChatGPT、Z.ai 账号和 Bedrock 不读余额,与这里写什么无关。 | | `signed_in` | `zai` \| `bigmodel` | — | `key` 来自哪一家的账号登录:`zai`(Z.ai)或 `bigmodel`(BigModel)。在桌面端登录账号时写入,在桌面端更换密钥时删除。不影响请求的发送方式。不写:密钥是手动填写的。 | @@ -348,7 +351,46 @@ Bedrock 上只为 AWS 列出支持提示缓存的 Claude 模型标注:Claude 3 ChatGPT 账号上游(`protocol: chatgpt`)只接受桌面应用登录得到的凭据,不能手写。不支持 Claude 和 Google 的订阅登录,请使用 API 密钥。 -有的中转站和账号同时只接受几个请求,多出来的直接拒绝。`max_concurrent` 让网关守住这个数:请求发出时占用这家的一个位置,回答完整交给客户端、或者客户端断开时归还。这家满了的时候,为复用提示缓存而留在这家的对话等空位,别的请求直接换下一家。最多等多久由 `failover.slot_wait_secs` 决定。等待不算失败,这家不会因此停用。只计算 token 数的请求不占位置。Responses 的 WebSocket 连接上,每个 `response.create` 从发出起占一个位置,直到它的回答结束,密钥的 `max_concurrent` 也一样;空闲的连接不占位置。 +有的中转站和账号同时只接受几个请求,多出来的直接拒绝。`max_concurrent` 让网关守住这个数:请求发出时占用这家的一个位置,回答完整交给客户端、或者客户端断开时归还。这家满了的时候,为复用提示缓存而留在这家的对话等空位,别的请求直接换下一家。最多等多久由 `failover.wait_secs` 决定。等待不算失败,这家不会因此停用。只计算 token 数的请求不占位置。Responses 的 WebSocket 连接上,每个 `response.create` 从发出起占一个位置,直到它的回答结束,密钥的 `max_concurrent` 也一样;空闲的连接不占位置。 + +#### `providers[].limits` + +一家上游的用量上限,照着这家自己的限额写:每几分钟、每几小时、每天、每周或每月最多多少个 +请求、多少 token、多少美元。写法和密钥的 `limits` 一样;一家上游可以有好几条,每一条都有 +余量才往这家发。 + +每一次发到这家的尝试算一个请求,成功失败都算,这家回 429 的不算。token 和费用算在得到回答 +的那一次尝试上。只计算 token 数的请求不算。还在进行的尝试按输入的估算计入,记下之后换成 +实数。 + +`minute`、`hour` 的窗口满了,这家和并发数满了一样被跳过:留在这家的对话等它空出余量,别的 +请求换下一家;候选都没有余量时,请求等最先空出来的那一家,最多等一个请求能等的时间(见 +`failover`)。`day`、`week`、`month` 用满之后,到下一期开始之前请求换下一家,不等它。满了 +不算失败,这家不会因此停用。重启之后,这一天、这一周、这个月的用量从请求记录里重新加起来, +所以 `retention.row_days` 要留够那一期,和密钥的一样;分钟、小时的窗口从空的开始。 + + + + +| 字段 | 类型 | 默认值 | 说明 | +|---|---|---|---| +| `per` | `minute` \| `hour` \| `day` \| `week` \| `month` | **必填** | 按多长一段时间算。`minute`、`hour` 是滚动的(最近 60 秒、最近 60 分钟,或者 `every` 个这样的单位);`day`、`week`、`month` 在本地时间的零点、周一零点、每月一号零点重新算。 | +| `every` | 整数 | — | 滚动窗口有几分钟、几小时:`{per: hour, every: 5}` 是最近 5 小时。`minute` 取 1 到 60,`hour` 取 1 到 24;`day`、`week`、`month` 不能写。不写:1。 | +| `requests` | 整数 | — | 最多往这家发这么多个请求。每一次发到了这家的尝试都算,它回 429 的不算;数 token 的请求不算。 | +| `tokens` | 整数 | — | 这家回答的请求最多用这么多 token:未命中缓存的输入、写入缓存的和输出。 | +| `cost` | 数字 | — | 这家回答的请求最多花这么多美元,按每个请求记下的费用算;至少 0.01。没有价格的模型算 0,`billing: free` 时全都算 0。 | +| `cache_reads` | 布尔 | `false` | 把读取缓存的 token 也算进去。只有 `tokens` 上限能写。 | + + +```yaml +providers: + - name: relay + base_url: https://relay.example.com + key: sk-relay + limits: + - { per: hour, every: 5, requests: 200 } + - { per: week, cost: 50 } +``` #### `providers[].oauth` @@ -802,8 +844,8 @@ security: ### `failover` 上游失败后会停用一段时间,接下来的请求直接交给下一个候选。停用多久取决于上游 -给出的原因:余额不足要等充值,额度用完要等到上游说的重置时刻,限流通常几秒钟就 -过去。只有一个候选的请求不受影响。Responses 的 WebSocket 连接上,每个 `response.create` +给出的原因:余额不足要等充值,额度用完要等到上游说的重置时刻。只有一个候选的请求 +不受影响。Responses 的 WebSocket 连接上,每个 `response.create` 在这里和一个请求一样算:第一段内容到达之前因为上游出错而失败的,算作一次失败;上游 拒绝或连不上的连接也一样。 @@ -834,12 +876,25 @@ failover: idle_timeout_secs: 600 ``` -上游的并发数满了(`max_concurrent`)时,一个请求等空位合计最多 `slot_wait_secs` -秒。等密钥的 `minute`、`hour` 上限也算在这段时间里,两样加起来不超过 `slot_wait_secs`。 -进行中的对话等它留在的那一家,到时还没有空位就换下一家,缓存在那边从头建; -新的对话遇到满着的上游直接跳过。候选全满时,请求等先空出来的那一家;都没有空出来, -客户端收到 429 和 `Retry-After`,说明上游都忙。如果有上游收到过这个请求并且失败了, -客户端收到的是那次失败。 +上游回 429(限流)时不停用,而是暂停:上游给出了等待时间的(`retry-after-ms`、 +`Retry-After` 或错误正文里的等待时间;不到一秒的、`retry-after-ms` 不到一毫秒的当作没有 +给出)暂停那么久,没有给出的按这家连续被限流的次数 +暂停 5、10、20、40 秒,之后每次 60 秒;上游回了别的答复就重新计数。暂停最长 +`rate_limit_max_pause_secs` 秒,熔断和 `load-balance` 都不把它算作失败。收到 429 的 +请求换下一个候选,或者等暂停结束,客户端察觉不到。流式回答在内容之前报的限流也按 429 +处理;只有请求的最后一个候选、而且这家五分钟内没被限流过时,客户端才会收到这个错误(这家照样 +暂停)。暂停结束时先放一个请求过去,这家接受了它(流式回答要等到内容开始),其余的请求再 +跟上;这个请求没有得到答复、超时,或者遇到这家额度用完,其余的请求也照常发出。暂停只保留 +到 twcore 重启为止。 + +一个请求等暂时不能发的上游合计最多 `wait_secs` 秒:并发数满了(`max_concurrent`)的、 +回了 429 正在暂停的、到了 `minute` 或 `hour` 用量上限的都算。等密钥的 `minute`、`hour` +上限也算在这段时间里,几样加起来不超过 `wait_secs`。进行中的对话等它留在的那一家,到时 +还没有恢复就换下一家,缓存在那边从头建;新的对话遇到这样的上游直接跳过。用完了天、周、 +月用量上限的上游也跳过。没有候选了,请求等最先恢复的那一家;流式回答照上文先收到 `200` +和保活,连接不会因此断开。都没有恢复,客户端收到 429,`Retry-After` 和 `retry-after-ms` +是最先恢复的那一家还要多久;候选全都用完了这一期的,`Retry-After` 是最早的重置时刻, +`x-should-retry` 是 `false`。如果有上游收到过这个请求并且失败了,客户端收到的是那次失败。 @@ -851,9 +906,9 @@ failover: | `max_pause_secs` | 整数 | `600` | 翻倍后的停用上限,单位秒,不小于 `pause_secs`。 | | `no_balance_pause_secs` | 整数 | `1800` | 上游报告余额不足时停用的秒数。 | | `quota_pause_secs` | 整数 | `3600` | 上游报告额度用完、但没有给出重置时间时停用的秒数。给出了重置时间的,停用到那一刻。 | -| `rate_limit_max_pause_secs` | 整数 | `3600` | 被限流的上游按它给的 `Retry-After` 停用,最多这么多秒。没有 `Retry-After` 的按没有说明原因的失败计。 | +| `rate_limit_max_pause_secs` | 整数 | `3600` | 上游回 429(限流)后暂停这家的最长秒数。上游给出了等待时间的(`retry-after-ms`、`Retry-After` 或错误正文里的等待时间)按它给的暂停;没有给出的,按这家连续被限流的次数退避:5、10、20、40 秒,之后每次 60 秒。暂停不算失败:请求在 `wait_secs` 之内等这家恢复,熔断也不计入。 | | `idle_timeout_secs` | 整数 | `300` | 上游多少秒没有发出内容,网关就不再等它。从请求发出的那一刻算起,每来一段内容重新计时:正文、推理、工具调用都算,心跳不算。整包(非流式)的回答从发出算到整份回来。还没有内容交给客户端时,这家上游记一次失败,请求换到下一家;没有下一家了,客户端收到超时错误。已经有内容交给客户端的,回答以一条错误收尾。取值 30 到 3600。 | -| `slot_wait_secs` | 整数 | `30` | 一个请求合计最多等的秒数,从过了密钥自己的 `max_concurrent` 时算起:等密钥的 `minute`、`hour` 上限空出名额,和等并发数满了(`max_concurrent`)的上游空出位置,都算在里面。密钥的上限到时空不出来就拒绝;等不到上游的空位就换下一家,候选全满时回 429。`0`:不等。取值 0 到 300。 | +| `wait_secs` | 整数 | `120` | 一个请求合计最多等的秒数,从过了密钥自己的 `max_concurrent` 时算起:等密钥的 `minute`、`hour` 上限空出名额,和等暂时不能发的上游 —— 并发数满了(`max_concurrent`)、回了 429 正在暂停、到了这家 `minute` 或 `hour` 的用量上限 —— 都算在里面。密钥的上限到时空不出来就拒绝。上游暂时不能发的先换下一家,没有下一家了就等最先恢复的那一家;留在某一家的对话等这一家。到时都没有恢复,客户端收到带 `Retry-After` 的 429。`0`:不等。取值 0 到 600。 | ### `aliases` @@ -921,7 +976,7 @@ groups: - `weights`(默认):只按权重。 - `latency`:越快的上游分得越多。快慢看典型的从发出请求到回答第一段内容的时间,与 `url-test` 使用同一份测量。比组内居中者快一倍的上游,权重乘以四;最多乘以十,最少乘以十分之一。 -- `health`:越少失败的上游分得越多。依据是最近 30 分钟内的最近 50 次请求:服务器错误、限流、额度或余额用尽、凭据被拒、超时和连接失败算作失败;请求本身导致的错误不算,客户端取消、手动中止、因并发数满了(`max_concurrent`)而跳过也不算;超过 `failover.idle_timeout_secs` 仍没有内容算作超时。经常失败的上游至少保留权重的二十分之一,仍会偶尔分到请求,以便发现它已经恢复;完全失败的上游照旧由 [`failover`](#cfg-failover) 暂停。 +- `health`:越少失败的上游分得越多。依据是最近 30 分钟内的最近 50 次请求:服务器错误、额度或余额用尽、凭据被拒、超时和连接失败算作失败;请求本身导致的错误不算,限流(429)、客户端取消、手动中止、因并发数满了(`max_concurrent`)或到了用量上限而跳过也不算;超过 `failover.idle_timeout_secs` 仍没有内容算作超时。经常失败的上游至少保留权重的二十分之一,仍会偶尔分到请求,以便发现它已经恢复;完全失败的上游照旧由 [`failover`](#cfg-failover) 暂停。 - `latency-health`:两个系数相乘。 快慢只在流式回答上测,从请求发给这家上游的那一刻算起:之前的等待、之前失败的上游都不算在内。因超过 `failover.idle_timeout_secs` 没有内容而被放弃的上游,按等满的那段时间计。Responses 的 WebSocket 连接上,每个 `response.create` 在快慢和成败上都算一个请求,快慢从上游开始回答它的那一刻算起。