From 8c8d60df09d258768fe2adb7fd1e7a9bab10b60f Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 17 Oct 2026 14:18:27 +0800 Subject: [PATCH] Live speed: /live reports how fast running requests generate right now /live used to give the speed of requests that finished in the last minute, so during a long request the menu bar had nothing to show. The gateway now keeps reading a successful stream after its first token and estimates the output as it arrives (text, reasoning text and tool arguments, with the same rule as the input estimate), multiplied by a per-upstream, per-model ratio. Each streaming request that finishes with reported usage calibrates that ratio (moving average, bounded 0.5..2, skipped when reasoning was hidden). Gemini's cumulative usage in the stream is used when it is higher. Each running request has a meter on the event bus (250 ms slots, 3 s window). /live gives every running request's tokens_per_sec and output_tokens (null before the first token and for non-streaming), and the total of all running requests: null when nothing runs, 0 when requests run but none is generating. The per-request speed recorded at the end is unchanged. Control protocol 49. Co-Authored-By: Claude Opus 5.5 --- crates/tw-api/src/lib.rs | 30 ++- crates/tw-control/tests/in_flight.rs | 18 +- crates/tw-engine/src/facts.rs | 14 +- crates/tw-engine/src/lib.rs | 4 +- crates/tw-gateway/src/ending.rs | 169 ++++++++++++--- crates/tw-gateway/src/lib.rs | 1 + crates/tw-gateway/src/pace.rs | 80 ++++++++ .../tw-gateway/src/server/pipeline/relay.rs | 2 +- crates/tw-gateway/src/state.rs | 3 + crates/tw-gateway/src/ws/turn.rs | 6 +- crates/tw-gateway/tests/live_speed.rs | 141 +++++++++++++ crates/tw-observe/src/bus.rs | 192 +++++++----------- crates/tw-observe/src/lib.rs | 2 + crates/tw-observe/src/meter.rs | 169 +++++++++++++++ 14 files changed, 663 insertions(+), 168 deletions(-) create mode 100644 crates/tw-gateway/src/pace.rs create mode 100644 crates/tw-gateway/tests/live_speed.rs create mode 100644 crates/tw-observe/src/meter.rs diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index 9396f4dc..e80b0174 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -940,7 +940,16 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// `replaced`，[`OutcomeDetail`] 的 `replaced` 多了可选的 `segments`（换了几段）和 /// `replacement`（换成的文字），回答里的命中 `direction` 是 `response`；[`SecurityCounts`] 多了 /// `content_replaced`。照 47 写的界面不认 `placeholder`、`replace` 这两个动作。 -pub const CONTROL_API_VERSION: u32 = 48; +/// +/// **49 起 `/live` 的速度是此刻的**：[`LiveView::tokens_per_sec`] 从「最近一分钟跑完的请求」 +/// 改成在跑的请求最近 [`LIVE_RATE_WINDOW_MS`] 毫秒的合计生成速度，按流里收到的内容现估； +/// 没有在跑的请求是 null，有请求却都还没开始生成是 0。[`RunningView`] 多了 `tokens_per_sec` +/// （这一个请求此刻的速度）和 `output_tokens`（到目前为止估的输出 token），第一个 token 之前、 +/// 非流式的是 null。 +pub const CONTROL_API_VERSION: u32 = 49; + +/// `/live` 的速度看最近多久（见 [`LiveView::tokens_per_sec`]），毫秒 +pub const LIVE_RATE_WINDOW_MS: u64 = 3_000; #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] @@ -1956,7 +1965,7 @@ impl Event { } } -/// 此刻的实时读数（`GET /live`）：在跑的请求，和最近跑完的请求生成得多快。 +/// 此刻的实时读数（`GET /live`）：在跑的请求，和它们此刻生成得多快。 /// /// **按事件数出来的，core 里只数一份。**以前菜单栏自己听事件流、自己记在跑的和 /// 跑完的：半路才开始听的那一段要拿 `/in-flight` 对账，丢过事件要重对，而这些 @@ -1966,10 +1975,13 @@ impl Event { pub struct LiveView { /// 开始了、还没有结局的请求，开始得早的在前。和 `/in-flight` 是同一批 pub running: Vec, - /// 最近一分钟跑完的请求平均每秒生成多少 token：每个请求的速度（见 - /// `RequestFinished::tokens_per_sec`）按它生成用的时间加权，也就是这些请求 - /// 吐出的 token 合计除以生成用的时间合计。**只数有速度的**：非流式的不算。这一分钟里 - /// 没有的是空，不是 0 + /// 此刻每秒生成多少 token：在跑的请求各自的速度（[`RunningView::tokens_per_sec`]）加起来。 + /// **没有在跑的请求是空**；有请求在跑、却都还在等第一个 token（或者是非流式的）是 0。 + /// + /// 是估的：流里每到一段正文、推理文字、工具参数就按字数估 token（同输入的估算），再乘上 + /// 这家上游这个模型此前实测的比例 —— 每个流式请求结束时，拿上游报的输出 token 数校准 + /// 一次。上游在流里报累计数的（Gemini）用它报的。**请求结束后记下的速度 + /// （`RequestFinished::tokens_per_sec`）照旧按上游报的实数算**，和这里不是一个数 pub tokens_per_sec: Option, } @@ -2000,6 +2012,12 @@ pub struct RunningView { /// 实际接下它的上游：尝试链里接下了的那一跳。**路由报出结论之前没有** /// （还在等上游的响应头，或者正在故障转移） pub upstream: Option, + /// 这个请求此刻每秒生成多少 token：最近 [`LIVE_RATE_WINDOW_MS`] 毫秒（刚开口不到这么久 + /// 的，从第一个 token 算起）里估出的输出除以这段时间，估法见 [`LiveView::tokens_per_sec`]。 + /// 停下来不吐字的时候往下掉到 0。**第一个 token 之前、非流式的没有** + pub tokens_per_sec: Option, + /// 到目前为止估的输出 token（同上的估法）。**第一个 token 之前、非流式的没有** + pub output_tokens: Option, } /// 此刻还在跑的请求（`GET /in-flight`）：每一个到目前为止的事件，和 core 此刻的 diff --git a/crates/tw-control/tests/in_flight.rs b/crates/tw-control/tests/in_flight.rs index 3e7adb62..336c172c 100644 --- a/crates/tw-control/tests/in_flight.rs +++ b/crates/tw-control/tests/in_flight.rs @@ -165,13 +165,13 @@ async fn live(router: axum::Router) -> serde_json::Value { serde_json::from_slice(&body).unwrap() } -/// `/live`：菜单栏要的在跑的请求和生成速率，由 core 数好。 +/// `/live`：菜单栏要的在跑的请求和此刻的生成速度，由 core 数好。 #[tokio::test] async fn live_gives_the_running_requests_and_the_generation_rate() { let (_d, bus, router) = app(); let idle = live(router.clone()).await; assert_eq!(idle["running"], serde_json::json!([])); - // 这一分钟里没有跑完的：是空，不是 0 + // 没有在跑的：是空，不是 0 assert!(idle["tokens_per_sec"].is_null(), "{idle}"); bus.emit(started(1, "claude-sonnet-5")); @@ -199,12 +199,22 @@ async fn live_gives_the_running_requests_and_the_generation_rate() { answered_model: None, }); bus.emit(started(2, "gpt-5.5")); + bus.emit(started(3, "gpt-5.5")); + // 3 开口了、记了 40 个；2 还在等第一个 token + bus.meter(3).unwrap().add(40); let v = live(router).await; - assert_eq!(v["tokens_per_sec"], 50, "{v}"); - assert_eq!(v["running"].as_array().unwrap().len(), 1); + // **跑完的不算**，只看在跑的：2 没开口不算速度，合计就是 3 的。3 刚开口，40 个除以 + // 开口以来的时间（至少一格 250 毫秒，所以最多 160）；跑得慢的机器上读得晚一点，数小一点 + let rate = v["running"][1]["tokens_per_sec"].as_u64().unwrap(); + assert!((1..=160).contains(&rate), "{v}"); + assert_eq!(v["tokens_per_sec"], rate, "{v}"); + assert_eq!(v["running"].as_array().unwrap().len(), 2); + assert_eq!(v["running"][1]["output_tokens"], 40, "{v}"); let r = &v["running"][0]; assert_eq!(r["id"], 2); + assert!(r["tokens_per_sec"].is_null(), "{v}"); + assert!(r["output_tokens"].is_null(), "{v}"); assert_eq!(r["client"], "我"); assert_eq!(r["client_hint"], "claude-code"); assert_eq!(r["model"], "gpt-5.5"); diff --git a/crates/tw-engine/src/facts.rs b/crates/tw-engine/src/facts.rs index 6d2468d9..3c97e990 100644 --- a/crates/tw-engine/src/facts.rs +++ b/crates/tw-engine/src/facts.rs @@ -235,15 +235,19 @@ fn part(text: &mut Text, p: &Part) -> u64 { 0 } -/// 攒着的文字：ASCII 字节数和其余字符数分开数。 -#[derive(Default)] -struct Text { +/// 攒着的文字：ASCII 字节数和其余字符数分开数。估 token 和输入同一个算法（见 +/// [`estimate_tokens`]）。 +/// +/// **一段一段喂、最后一起算**：流里的一段常常只有两三个字符，每段各自取整的话，英文会 +/// 被估高好几倍。网关读着回答的流估已经生成了多少，用的也是它 +#[derive(Debug, Default, Clone, Copy)] +pub struct Text { ascii: u64, other: u64, } impl Text { - fn add(&mut self, s: &str) { + pub fn add(&mut self, s: &str) { self.add_bytes(s.as_bytes()); } @@ -275,7 +279,7 @@ impl Text { let _ = serde_json::to_writer(Count(self), v); } - fn tokens(&self) -> u64 { + pub fn tokens(&self) -> u64 { self.ascii.div_ceil(4) + self.other } } diff --git a/crates/tw-engine/src/lib.rs b/crates/tw-engine/src/lib.rs index 3dd814f4..f9dab4ca 100644 --- a/crates/tw-engine/src/lib.rs +++ b/crates/tw-engine/src/lib.rs @@ -13,6 +13,8 @@ pub use engine::{ RuleNotes, SetAction, Target, balance_factors, has_catch_all, is_builtin_group, notes, order_by, scalar_name, }; -pub use facts::{RequestFacts, TokenParts, estimate_parts, estimate_strings, estimate_tokens}; +pub use facts::{ + RequestFacts, Text as TextTokens, TokenParts, estimate_parts, estimate_strings, estimate_tokens, +}; pub use time::{LocalTime, Window}; pub use weighted::Member; diff --git a/crates/tw-gateway/src/ending.rs b/crates/tw-gateway/src/ending.rs index 9ec6856c..c6c4d725 100644 --- a/crates/tw-gateway/src/ending.rs +++ b/crates/tw-gateway/src/ending.rs @@ -63,9 +63,9 @@ pub struct Ending { /// 一条连接上每一次回答报的用量加起来（[`Ending::add_usage`]）。有它就不看嗅探器 total: Option, tap: ResponseTap, - /// 认第一个 token 的。**只在上游回的是成功的流时才有**（见 [`Ending::streaming`]）， - /// 认出来就扔掉 —— 之后的字节不必再解析 - first: Option, + /// 读着回答的流：认第一个 token，之后估生成了多少（`/live` 的速度）。**只在上游回的是 + /// 成功的流时才有**（见 [`Ending::streaming`]） + stream: Option, /// 第一个 token 是什么时候、以什么开的头 opened: Option, /// 第一个 token 到的时候给回答的这一家记一个快慢样本（见 [`Ending::timed`]） @@ -176,11 +176,95 @@ impl Watch { } } -/// 还在等第一个 token。 -struct FirstToken { +/// 读着回答的流。先认第一个 token；认出来之后接着读，每段输出按字数估 token，记到这个 +/// 请求的速度表上（`/live`，见 [`tw_observe::Meter`]）。 +struct Stream { reader: Reader, - /// 在那之前有过推理块，却一个推理的字都没有：推理被隐藏了 + /// 上游的格式。Gemini 在流里一段段报累计的输出数，有它就用它 + upstream: ir::Dialect, + /// 在第一个 token 之前有过推理块，却一个推理的字都没有：推理被隐藏了 hidden_thought: bool, + /// 认出第一个 token 了 + opened: bool, + /// 第一个 token 起的正文、推理文字、工具参数 + text: tw_engine::TextTokens, + /// 一个推理块开了头还没见到字（见 [`Stream::count`]） + thought_pending: bool, + /// 有过一个字都没有的推理块：上游报的输出里有看不见的推理，不拿来校准 + thought_hidden_later: bool, + /// 见过推理的字 + thought_text: bool, + /// 上游在流里报的累计输出数（Gemini） + reported: u64, + /// 已经记到速度表上的 + added: u64, + /// 这个请求的速度表。总线上已经没有它（结局报过了）的话是空的 + meter: Option>, + /// 估数乘多少（见 [`crate::pace`]），和结束时校准记到哪儿 + pace: std::sync::Arc, + provider: String, + ratio: f64, +} + +impl Stream { + /// 第一个 token 之后的一个事件：输出的字数进账，Gemini 报的累计数记下 + fn count(&mut self, e: &ir::Event) { + match e { + ir::Event::Delta { + delta: ir::Delta::Text(t) | ir::Delta::ToolInput(t), + .. + } => self.text.add(t), + ir::Event::Delta { + delta: ir::Delta::Thinking(t), + .. + } => { + self.text.add(t); + if !t.is_empty() { + self.thought_text = true; + self.thought_pending = false; + } + } + ir::Event::BlockStart { + kind: ir::BlockKind::Thinking, + .. + } => self.thought_pending = true, + ir::Event::BlockStop { .. } if self.thought_pending => { + self.thought_pending = false; + self.thought_hidden_later = true; + } + ir::Event::Usage(u) if self.upstream == ir::Dialect::Gemini => { + self.reported = self.reported.max(u.output); + } + _ => {} + } + } + + /// 这一块读完了：估到目前为止一共多少，多出来的记到速度表上。开口那一刻哪怕估出 0 个 + /// 也记一笔 —— 速度从那一刻算起 + fn tally(&mut self) { + if !self.opened { + return; + } + let estimated = (self.text.tokens() as f64 * self.ratio).round() as u64; + let total = estimated.max(self.reported); + let Some(m) = &self.meter else { return }; + if total > self.added || self.added == 0 { + m.add(total.saturating_sub(self.added)); + self.added = total.max(self.added); + } + } + + /// 结束了，上游报了用量：拿它校准这家这个模型的估法。**有看不见的推理就不校**：上游报 + /// 的输出里有推理 token，流里却没有它们的字，比出来的比例会偏高 + fn learn(&self, model: &str, usage: &Usage) { + let hidden = self.hidden_thought + || self.thought_hidden_later + || (usage.reasoning > 0 && !self.thought_text); + if self.opened && !hidden { + self.pace + .learn(&self.provider, model, self.text.tokens(), usage.output); + } + } } /// 回答的这一家的快慢样本记到哪儿、从什么时候算起（见 [`crate::latency`]）。 @@ -226,7 +310,7 @@ impl Ending { sniffer: Sniffer::new(), total: None, tap: ResponseTap::new(), - first: None, + stream: None, opened: None, lap: None, watch: None, @@ -249,10 +333,30 @@ impl Ending { /// /// **流里报了错的不是成功。**上游先回 200、写到一半才报 `overloaded_error` 的，客户端 /// 拿到的是半截回答加一个错误；记成成功的话，这家上游看起来从不出错。 - pub fn streaming(&mut self, upstream: ir::Dialect, provider: &str) { - self.first = Some(FirstToken { + /// + /// 认出第一个 token 之后接着读，估生成了多少记到这个请求的速度表上（`/live`），估数乘上 + /// `pace` 里这家这个模型的比例；成功结束时按上游报的实数校准它（见 [`crate::pace`]）。 + pub fn streaming( + &mut self, + upstream: ir::Dialect, + provider: &str, + pace: &std::sync::Arc, + ) { + self.stream = Some(Stream { reader: Reader::new(upstream), + upstream, hidden_thought: false, + opened: false, + text: Default::default(), + thought_pending: false, + thought_hidden_later: false, + thought_text: false, + reported: 0, + added: 0, + meter: self.bus.meter(self.id), + pace: pace.clone(), + provider: provider.to_string(), + ratio: pace.ratio(provider, &self.model), }); self.watch = Some(Watch { provider: provider.to_string(), @@ -387,26 +491,32 @@ impl Ending { } } - /// 在这一块里找第一个 token。找到就报，之后不再解析。 + /// 读这一块：认第一个 token，认出来就报；之后的输出估数记到速度表上（见 [`Stream`]）。 fn spot(&mut self, chunk: &[u8]) { - let Some(f) = self.first.as_mut() else { + let Some(f) = self.stream.as_mut() else { return; }; let mut opened = None; for e in f.reader.feed(chunk) { - match opening(&e) { - Some(with_thought) => { - opened = Some((with_thought, f.hidden_thought)); - break; + if !f.opened { + match opening(&e) { + Some(with_thought) => { + f.opened = true; + opened = Some((with_thought, f.hidden_thought)); + } + None if hides_thought(&e) => { + f.hidden_thought = true; + continue; + } + None => continue, } - None if hides_thought(&e) => f.hidden_thought = true, - None => {} } + f.count(&e); } + f.tally(); let Some((with_thought, hidden_thought)) = opened else { return; }; - self.first = None; if let Some(lap) = self.lap.take() { lap.latency .record(&lap.provider, crate::latency::ms(lap.sent.elapsed())); @@ -502,6 +612,11 @@ impl Ending { return; } let (usage, answered_model) = self.settle(); + if let (Some(s), Some(u)) = (self.stream.take(), usage.as_ref()) + && (200..300).contains(&status) + { + s.learn(&self.model, u); + } if let Some(t) = self.answer.take() && (200..300).contains(&status) { @@ -857,7 +972,7 @@ mod tests { let bus = tw_observe::EventBus::new(); let mut rx = bus.subscribe(); let mut e = responding(&bus); - e.streaming(ir::Dialect::Anthropic, "官方"); + e.streaming(ir::Dialect::Anthropic, "官方", &Default::default()); e.feed(MESSAGE_START); // 错误帧被切成一个字节一个字节地到 for b in OVERLOADED { @@ -891,7 +1006,7 @@ mod tests { let bus = tw_observe::EventBus::new(); let mut rx = bus.subscribe(); let mut e = responding(&bus); - e.streaming(ir::Dialect::Chat, "中转"); + e.streaming(ir::Dialect::Chat, "中转", &Default::default()); e.feed(b"data: {\"error\":{\"message\":\"Provider disconnected\"}}"); e.finished(200); assert!( @@ -1233,7 +1348,7 @@ mod tests { Some(tx), ); e.responded(200); - e.streaming(ir::Dialect::Responses, "chatgpt"); + e.streaming(ir::Dialect::Responses, "chatgpt", &Default::default()); wire(&mut e, COMPLETED); e.delivered(); drop(e); @@ -1264,7 +1379,7 @@ mod tests { let bus = tw_observe::EventBus::new(); let mut rx = bus.subscribe(); let mut e = responding(&bus); - e.streaming(ir::Dialect::Responses, "chatgpt"); + e.streaming(ir::Dialect::Responses, "chatgpt", &Default::default()); e.feed(b"event: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"code\":\"server_error\",\"message\":\"boom\"}}}\n\n"); e.delivered(); drop(e); @@ -1345,7 +1460,7 @@ mod tests { let bus = tw_observe::EventBus::new(); let mut rx = bus.subscribe(); let mut e = responding(&bus); - e.streaming(ir::Dialect::Anthropic, "up"); + e.streaming(ir::Dialect::Anthropic, "up", &Default::default()); e.feed(MESSAGE_START); e.feed(TEXT_BLOCK); assert!(drain(&mut rx).is_empty(), "开场帧被当成了第一个 token"); @@ -1376,7 +1491,7 @@ mod tests { for _ in 0..3 { let mut e = Ending::new(bus.clone(), 7, MODEL.into(), arrived, 1_000, None); e.responded(200); - e.streaming(ir::Dialect::Anthropic, "up"); + e.streaming(ir::Dialect::Anthropic, "up", &Default::default()); let sent = Instant::now() - std::time::Duration::from_millis(300); e.timed(lap(&latency, sent)); e.feed(MESSAGE_START); @@ -1397,7 +1512,7 @@ mod tests { let latency = std::sync::Arc::new(crate::latency::Latency::new()); for _ in 0..3 { let mut e = responding(&bus); - e.streaming(ir::Dialect::Anthropic, "up"); + e.streaming(ir::Dialect::Anthropic, "up", &Default::default()); e.timed(lap(&latency, Instant::now())); e.feed(MESSAGE_START); e.feed(b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n"); @@ -1418,7 +1533,7 @@ mod tests { let bus = tw_observe::EventBus::new(); let mut rx = bus.subscribe(); let mut e = responding(&bus); - e.streaming(ir::Dialect::Anthropic, "up"); + e.streaming(ir::Dialect::Anthropic, "up", &Default::default()); e.feed(MESSAGE_START); e.feed(b"event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"t1\",\"name\":\"Bash\",\"input\":{}}}\n\n"); assert_eq!(first_tokens(&drain(&mut rx)).len(), 1); @@ -1548,7 +1663,7 @@ mod tests { let bus = tw_observe::EventBus::new(); let mut rx = bus.subscribe(); let mut e = responding(&bus); - e.streaming(ir::Dialect::Responses, "up"); + e.streaming(ir::Dialect::Responses, "up", &Default::default()); let frames = [ r#"{"type":"response.created","response":{"id":"r","status":"in_progress","model":"gpt-5","output":[]}}"#, r#"{"type":"response.output_item.added","output_index":0,"item":{"type":"message","id":"m","role":"assistant","content":[]}}"#, diff --git a/crates/tw-gateway/src/lib.rs b/crates/tw-gateway/src/lib.rs index 808cbecf..819d115f 100644 --- a/crates/tw-gateway/src/lib.rs +++ b/crates/tw-gateway/src/lib.rs @@ -42,6 +42,7 @@ pub mod live; pub mod models; pub mod oauth; pub mod outbound; +pub mod pace; pub mod panic_note; pub mod plugin; pub mod probe; diff --git a/crates/tw-gateway/src/pace.rs b/crates/tw-gateway/src/pace.rs new file mode 100644 index 00000000..6fc81a92 --- /dev/null +++ b/crates/tw-gateway/src/pace.rs @@ -0,0 +1,80 @@ +//! 流里估出来的 token 和上游报的实数差多少，按上游 + 模型记。 +//! +//! `/live` 的速度是读着回答的流现估的（见 [`crate::ending::Ending::streaming`]）：每段正文按 +//! 字数估 token，和估输入同一个算法（ASCII 四个字节一个，其余一个字一个）。各家分词不同， +//! 中文差得尤其多。**每个流式请求结束时拿上游报的输出数校准一次**：实数除以估数，按 +//! 滑动平均记下，之后同一家同一个模型的估数乘上它。重启之后从头校准，第一个请求用的是 +//! 没校准过的估法。 +//! +//! 只管 `/live` 的那个数。请求结束后记下的速度按上游报的实数算，用不着它。 + +use std::collections::HashMap; +use std::sync::Mutex; + +/// 新的一个样本占多少分量 +const WEIGHT: f64 = 0.3; +/// 比例限在这个范围里：估法再差也差不出两倍，超出的多半是这一次估漏了什么（隐藏的推理） +const RANGE: (f64, f64) = (0.5, 2.0); +/// 估出来不到这么多个的不拿来校准：几个 token 的回答，差一两个就是几成 +const MIN_SAMPLE: u64 = 50; + +#[derive(Debug, Default)] +pub struct Pace { + ratios: Mutex>, +} + +impl Pace { + pub fn new() -> Self { + Self::default() + } + + /// 这家上游这个模型，估数该乘多少。没校准过是 1 + pub fn ratio(&self, provider: &str, model: &str) -> f64 { + let r = self.ratios.lock().unwrap_or_else(|p| p.into_inner()); + r.get(&(provider.to_string(), model.to_string())) + .copied() + .unwrap_or(1.0) + } + + /// 一个流式请求结束了：估出来 `estimated` 个（没乘比例的），上游说是 `actual` 个 + pub fn learn(&self, provider: &str, model: &str, estimated: u64, actual: u64) { + if estimated < MIN_SAMPLE || actual == 0 { + return; + } + let sample = (actual as f64 / estimated as f64).clamp(RANGE.0, RANGE.1); + let mut r = self.ratios.lock().unwrap_or_else(|p| p.into_inner()); + let entry = r.entry((provider.to_string(), model.to_string())); + entry + .and_modify(|old| *old = *old * (1.0 - WEIGHT) + sample * WEIGHT) + .or_insert(sample); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_first_sample_is_taken_as_it_is_and_later_ones_move_it_gradually() { + let p = Pace::new(); + assert_eq!(p.ratio("a", "m"), 1.0); + p.learn("a", "m", 100, 150); + assert_eq!(p.ratio("a", "m"), 1.5); + p.learn("a", "m", 100, 100); + assert!((p.ratio("a", "m") - 1.35).abs() < 1e-9); + // 别家、别的模型不受影响 + assert_eq!(p.ratio("a", "n"), 1.0); + assert_eq!(p.ratio("b", "m"), 1.0); + } + + #[test] + fn small_and_wild_samples_are_kept_in_bounds() { + let p = Pace::new(); + p.learn("a", "m", 10, 100); + assert_eq!(p.ratio("a", "m"), 1.0, "too small to judge"); + p.learn("a", "m", 100, 1_000); + assert_eq!(p.ratio("a", "m"), 2.0); + p.learn("b", "m", 100, 0); + assert_eq!(p.ratio("b", "m"), 1.0); + } +} diff --git a/crates/tw-gateway/src/server/pipeline/relay.rs b/crates/tw-gateway/src/server/pipeline/relay.rs index a5a548c6..f1efb403 100644 --- a/crates/tw-gateway/src/server/pipeline/relay.rs +++ b/crates/tw-gateway/src/server/pipeline/relay.rs @@ -140,7 +140,7 @@ pub(super) fn respond( let upstream_dialect = upstream_dialect(req, session.as_ref()); // 成功的流才有「第一个 token」：整包的一起到，错误不是回答 if generates && plan.is_sse && status.is_success() { - ending.streaming(upstream_dialect, &provider.name); + ending.streaming(upstream_dialect, &provider.name, &state.pace); // `url-test` 和按快慢分的 `load-balance` 的判据：从这一跳发出去到第一个 token（见 // `crate::latency`）。**只有成功的流有**：一个 500 在十毫秒内返回，会让最坏的上游 // 看起来最快；整包的回答分不出排队和说话 diff --git a/crates/tw-gateway/src/state.rs b/crates/tw-gateway/src/state.rs index b4109137..5a56b067 100644 --- a/crates/tw-gateway/src/state.rs +++ b/crates/tw-gateway/src/state.rs @@ -210,6 +210,8 @@ pub struct AppState { /// /// **跨重载存活**：改一条规则不该让所有上游回到「没测过」。 pub latency: Arc, + /// 流里估的 token 和上游实数之比，按上游 + 模型（见 [`crate::pace`]）。`/live` 的速度乘它 + pub pace: Arc, /// 价格簿：默认价目表 + 自定义价目表 + 哪个上游用哪张。 /// /// **全进程只有这一份。**`cheapest` 排序、记账、测速报价、回放都从这 @@ -320,6 +322,7 @@ impl AppState { listening: Arc::new(std::sync::Mutex::new(Default::default())), oauth: Arc::new(crate::oauth::Cache::new()), latency: Arc::new(crate::latency::Latency::new()), + pace: Arc::new(crate::pace::Pace::new()), pricing: tw_pricing::shared(tw_pricing::PriceBook::new( Arc::new(tw_pricing::Table::builtin().unwrap_or_else(|e| { // **加载不了不能挡住启动**：那时成本显示「未知」， diff --git a/crates/tw-gateway/src/ws/turn.rs b/crates/tw-gateway/src/ws/turn.rs index cf20ec47..e493abc5 100644 --- a/crates/tw-gateway/src/ws/turn.rs +++ b/crates/tw-gateway/src/ws/turn.rs @@ -527,7 +527,11 @@ pub(crate) async fn admit(a: Admit) -> Result { at_ms: a.at_ms, }); hold.bind(id); - ending.streaming(tw_dialect::ir::Dialect::Responses, &line.provider); + ending.streaming( + tw_dialect::ir::Dialect::Responses, + &line.provider, + &state.pace, + ); // 拒绝的也在开始之后：被拒是一次来源为 `denied` 的失败，一个字节都不发 let details = a.screened.clone(); if let Some(why) = diff --git a/crates/tw-gateway/tests/live_speed.rs b/crates/tw-gateway/tests/live_speed.rs new file mode 100644 index 00000000..e98fb76f --- /dev/null +++ b/crates/tw-gateway/tests/live_speed.rs @@ -0,0 +1,141 @@ +//! `/live` 的速度是此刻的：流还在吐的时候就有，按收到的输出现估；结束时按上游报的实数 +//! 校准这家这个模型的估法。 +//! +//! 起真网关、真上游：要验的是「流到一半」那一刻总线上的读数，单元测试里只能假装。 + +use std::net::SocketAddr; +use std::time::Duration; + +use axum::Router; +use axum::routing::any; +use bytes::Bytes; +use tw_config::{Client, Config, Protocol, Provider}; + +const MESSAGE_START: &[u8] = b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":50,\"output_tokens\":1}}}\n\n"; +const TEXT_BLOCK: &[u8] = b"event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n"; +/// 八个 ASCII 字节：估出来两个 token +const TEXT_DELTA: &[u8] = b"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"words ab\"}}\n\n"; +/// 上游说一共输出了 300 个 +const MESSAGE_END: &[u8] = b"event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":300}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + +/// 上游读完输入要多久：第一个字在这之后 +const PREFILL: Duration = Duration::from_millis(400); + +async fn listen(app: Router) -> SocketAddr { + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let a = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + a +} + +/// 开场帧马上发，读完输入之后每 `every` 吐一段，一共 `deltas` 段，最后报累计输出 +fn streaming(deltas: usize, every: Duration) -> Router { + Router::new().fallback(any(move || async move { + let body = async_stream::stream! { + yield Ok::<_, std::io::Error>(Bytes::from_static(MESSAGE_START)); + tokio::time::sleep(PREFILL).await; + yield Ok(Bytes::from_static(TEXT_BLOCK)); + for _ in 0..deltas { + yield Ok(Bytes::from_static(TEXT_DELTA)); + tokio::time::sleep(every).await; + } + yield Ok(Bytes::from_static(MESSAGE_END)); + }; + axum::response::Response::builder() + .header("content-type", "text/event-stream") + .body(axum::body::Body::from_stream(body)) + .unwrap() + })) +} + +fn cfg(upstream: SocketAddr) -> Config { + Config { + clients: vec![Client { + name: "me".into(), + key: "tw-k".into(), + ..Default::default() + }], + providers: vec![Provider { + name: "up".into(), + base_url: format!("http://{upstream}"), + key: Some("sk-x".into()), + protocol: Some(Protocol::Anthropic), + ..Default::default() + }], + ..Default::default() + } +} + +struct Gateway { + addr: SocketAddr, + bus: tw_observe::EventBus, + pace: std::sync::Arc, +} + +async fn serve(cfg: Config) -> Gateway { + let state = tw_gateway::AppState::new(cfg).unwrap(); + let (bus, pace) = (state.bus.clone(), state.pace.clone()); + let addr = tw_gateway::serve(state, ([127, 0, 0, 1], 0).into()) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + Gateway { addr, bus, pace } +} + +async fn ask(gw: SocketAddr) -> String { + reqwest::Client::new() + .post(format!("http://{gw}/v1/messages")) + .header("x-api-key", "tw-k") + .body(r#"{"model":"claude-sonnet-5","stream":true,"messages":[]}"#) + .send() + .await + .unwrap() + .text() + .await + .unwrap() +} + +/// **流到一半就有速度。**还在读输入的时候有请求在跑、却没有字：合计是 0，这个请求没有 +/// 速度；开口之后每秒约 20 个（每 100 毫秒一段、一段估两个） +#[tokio::test] +async fn the_speed_is_there_while_the_stream_is_still_running() { + let up = listen(streaming(30, Duration::from_millis(100))).await; + let gw = serve(cfg(up)).await; + let request = tokio::spawn(ask(gw.addr)); + + tokio::time::sleep(Duration::from_millis(200)).await; + let live = gw.bus.live(); + assert_eq!(live.running.len(), 1, "{live:?}"); + assert_eq!(live.tokens_per_sec, Some(0)); + assert_eq!(live.running[0].tokens_per_sec, None); + assert_eq!(live.running[0].output_tokens, None); + + tokio::time::sleep(Duration::from_millis(1_600)).await; + let live = gw.bus.live(); + let r = &live.running[0]; + let rate = r.tokens_per_sec.expect("opened by now"); + assert!((12..=30).contains(&rate), "{rate} tokens/s"); + assert_eq!(live.tokens_per_sec, Some(rate)); + let out = r.output_tokens.unwrap(); + assert!((20..=40).contains(&out), "{out} tokens so far"); + + let text = request.await.unwrap(); + assert!(text.contains("message_stop"), "{text}"); + tokio::time::sleep(Duration::from_millis(100)).await; + let live = gw.bus.live(); + assert!(live.running.is_empty()); + assert_eq!(live.tokens_per_sec, None); +} + +/// **结束时按上游的实数校准。**估出 60 个（30 段、每段两个），上游说 300 个：之后这家这个 +/// 模型的估数乘 2（比例最多到 2） +#[tokio::test] +async fn a_finished_stream_teaches_the_estimate() { + let up = listen(streaming(30, Duration::from_millis(5))).await; + let gw = serve(cfg(up)).await; + assert_eq!(gw.pace.ratio("up", "claude-sonnet-5"), 1.0); + let text = ask(gw.addr).await; + assert!(text.contains("message_stop"), "{text}"); + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(gw.pace.ratio("up", "claude-sonnet-5"), 2.0); +} diff --git a/crates/tw-observe/src/bus.rs b/crates/tw-observe/src/bus.rs index 1c9c845d..feff35b9 100644 --- a/crates/tw-observe/src/bus.rs +++ b/crates/tw-observe/src/bus.rs @@ -11,22 +11,21 @@ //! **存储层不订阅广播**，它有自己的一条有界通道（[`EventBus::record_feed`]）：广播上 //! 掉队丢掉的是界面上的几个点，记录丢掉的是那几个请求，永远不在库里了。 -use std::collections::{BTreeMap, HashMap, VecDeque}; +use std::collections::BTreeMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; -use std::time::{Duration, Instant}; +use std::time::Instant; use tokio::sync::{broadcast, mpsc}; +use crate::meter::Meter; + const CAPACITY: usize = 1024; /// 存储层那条通道有多少格（见 [`EventBus::record_feed`]）。一个请求五六条事件，这么多格 /// 够存储层落后两千多个请求 pub const RECORD_CAPACITY: usize = 16 * 1024; -/// 生成速率看最近这么久里跑完的请求 -pub const RATE_WINDOW: Duration = Duration::from_secs(60); - #[derive(Clone)] pub struct EventBus { tx: broadcast::Sender, @@ -73,10 +72,6 @@ impl RecordFeed { struct Tally { /// 开始了、还没有结局的请求，按 id open: BTreeMap, - /// 在跑的请求第一个 token 用了多久，毫秒。生成用时要从总耗时里减掉它 - ttft: HashMap, - /// 最近跑完的：（什么时候结束的，第一个 token 之后吐了多少 token，生成用了多少毫秒） - done: VecDeque<(Instant, u64, u64)>, } /// 一个还没有结局的请求。 @@ -85,6 +80,8 @@ struct Open { since: Instant, /// 到目前为止关于它的事件，按发生的先后，第一条是开始事件（见 [`about_the_request`]） events: Vec, + /// 它此刻生成得多快。数据面读着流往里记（见 [`EventBus::meter`]） + meter: Arc, } /// 这条事件说的是它挂着的那个请求本身吗 —— 是的话，半路才来的一方要重建那个请求， @@ -140,18 +137,6 @@ fn now_ms() -> u64 { .unwrap_or(0) } -impl Tally { - fn forget(&mut self, now: Instant) { - while self - .done - .front() - .is_some_and(|(at, _, _)| now.saturating_duration_since(*at) > RATE_WINDOW) - { - self.done.pop_front(); - } - } -} - impl Default for EventBus { fn default() -> Self { Self::new() @@ -242,38 +227,14 @@ impl EventBus { Open { since: now, events: vec![ev.clone()], + meter: Arc::new(Meter::new(now)), }, ); } - E::RequestFirstToken { id, ttft_ms } => { - if let Some(o) = t.open.get_mut(id) { - o.events.push(ev.clone()); - t.ttft.insert(*id, *ttft_ms); - } - } - E::RequestFinished { - id, - duration_ms, - tokens_per_sec, - .. - } => { - t.open.remove(id); - let ttft = t.ttft.remove(id); - // **只数有速度的。**速度怎么算、什么时候没有，由网关在结局里定（见 - // `RequestFinished::tokens_per_sec`）：非流式的没有第一个 token，取消的、 - // 断掉的输出不全。这里只按生成用的时间把它们合起来 - if let (Some(rate), Some(ttft)) = (tokens_per_sec, ttft) { - let gen_ms = duration_ms.saturating_sub(ttft); - if *rate > 0 && gen_ms > 0 { - t.done - .push_back((now, u64::from(*rate) * gen_ms / 1000, gen_ms)); - } - } - t.forget(now); - } - E::RequestFailed { id, .. } | E::RequestCancelled { id, .. } => { + E::RequestFinished { id, .. } + | E::RequestFailed { id, .. } + | E::RequestCancelled { id, .. } => { t.open.remove(id); - t.ttft.remove(id); } // 别的事件挂着的 id 要么是一个请求的，要么是它自己取的号 —— 号从同一个 // 计数器里取，不会和一个在跑的请求撞上 @@ -326,33 +287,41 @@ impl EventBus { t.open.contains_key(&id) } - /// 此刻的实时读数：在跑的请求（和 [`EventBus::in_flight`] 同一批），和最近 - /// 一分钟跑完的请求平均每秒生成多少 token。 + /// 这个在跑的请求的速度表：数据面读着流，每读到一段输出就往里记估出来的 token（见 + /// [`Meter::add`]），`/live` 从这里读。**开始事件之后才有**，结局之后就没了 —— 记在一份 + /// 已经划掉的表上不碍事，只是没人再读 + pub fn meter(&self, id: u64) -> Option> { + let t = self.tally.lock().unwrap_or_else(|p| p.into_inner()); + t.open.get(&id).map(|o| o.meter.clone()) + } + + /// 此刻的实时读数：在跑的请求（和 [`EventBus::in_flight`] 同一批），各自此刻生成得 + /// 多快，和加起来的速度。 /// - /// **量的是生成的快慢**：每个请求的速度按它生成用的时间（总耗时减去第一个 token 的 - /// 时刻）加权，也就是吐出的 token 合计除以生成用的时间合计。拿「一分钟里输出了多少」去除以六十的话，一个跑了一分钟才 - /// 结束的长请求，要等它结束之后才把速率一次性摊进来，数字跟着请求的长短忽高 - /// 忽低。 + /// **量的是此刻**：每个请求看最近 [`crate::meter::WINDOW`] 里生成了多少（见 [`Meter`]）。 + /// 以前看的是最近一分钟跑完的请求，一个跑几分钟的长请求要等它结束才有数，跑着的时候 + /// 菜单栏上是空的。 + /// + /// 没有在跑的请求，速度是空；有请求在跑、却都还没开口（等第一个 token，或者非流式的）， + /// 是 0 —— 有活在干，只是此刻没有字出来 pub fn live(&self) -> tw_api::LiveView { self.live_at(Instant::now()) } fn live_at(&self, now: Instant) -> tw_api::LiveView { - let mut t = self.tally.lock().unwrap_or_else(|p| p.into_inner()); - t.forget(now); - let (tokens, ms) = t - .done - .iter() - .fold((0u64, 0u64), |(tok, ms), (_, out, gen_ms)| { - (tok + out, ms + gen_ms) - }); + let t = self.tally.lock().unwrap_or_else(|p| p.into_inner()); let mut running: Vec = t.open.values().filter_map(|o| running(o, now)).collect(); running.sort_by_key(|r| (r.at_ms, r.id)); + let tokens_per_sec = (!running.is_empty()).then(|| { + running + .iter() + .filter_map(|r| r.tokens_per_sec) + .fold(0u32, u32::saturating_add) + }); tw_api::LiveView { running, - tokens_per_sec: (ms > 0) - .then(|| (tokens.saturating_mul(1000) / ms).min(u32::MAX as u64) as u32), + tokens_per_sec, } } @@ -400,6 +369,7 @@ fn running(o: &Open, now: Instant) -> Option { .map(|a| a.provider.clone()), _ => None, }); + let reading = o.meter.read_at(now); Some(tw_api::RunningView { id: *id, client: client.clone(), @@ -413,11 +383,15 @@ fn running(o: &Open, now: Instant) -> Option { rule: rule.clone(), group: group.clone(), upstream, + tokens_per_sec: reading.map(|r| r.tokens_per_sec), + output_tokens: reading.map(|r| r.total), }) } #[cfg(test)] mod tests { + use std::time::Duration; + use super::*; fn started(id: u64) -> tw_api::Event { @@ -650,74 +624,43 @@ mod tests { } } - /// **几个请求按生成用的时间加权**：吐出的 token 合计除以生成用的时间合计。 + /// **没有在跑的就没有速度**，不是 0；有请求在跑却都还没开口，是 0 #[test] - fn the_rate_weighs_each_request_by_its_generation_time() { + fn the_speed_is_empty_when_idle_and_zero_when_nothing_is_generating() { let b = EventBus::new(); let now = Instant::now(); - // 第一个 token 在 1 秒，之后 2 秒、每秒 50 个 + assert_eq!(b.live_at(now).tokens_per_sec, None); b.track_at(&started(1), now); - b.track_at(&first_token(1, 1_000), now); - b.track_at(&finished(1, 3_000, Some(50)), now); - assert_eq!(b.live_at(now).tokens_per_sec, Some(50)); - // 又一个：1 秒、每秒 200 个。合起来 300 个 / 3 秒，不是两个速度的平均 125 - b.track_at(&started(2), now); - b.track_at(&first_token(2, 500), now); - b.track_at(&finished(2, 1_500, Some(200)), now); - assert_eq!(b.live_at(now).tokens_per_sec, Some(100)); - } - - /// 一分钟之前跑完的不算；**没有就是没有**，不是 0 - #[test] - fn an_idle_minute_has_no_rate() { - let b = EventBus::new(); - let then = Instant::now(); - assert_eq!(b.live_at(then).tokens_per_sec, None); - b.track_at(&started(1), then); - b.track_at(&first_token(1, 100), then); - b.track_at(&finished(1, 1_100, Some(40)), then); - assert_eq!(b.live_at(then).tokens_per_sec, Some(40)); - let later = then + RATE_WINDOW + Duration::from_secs(1); - assert_eq!(b.live_at(later).tokens_per_sec, None); + let live = b.live_at(now + Duration::from_secs(5)); + assert_eq!(live.tokens_per_sec, Some(0)); + assert_eq!(live.running[0].tokens_per_sec, None); + assert_eq!(live.running[0].output_tokens, None); } - /// 没有速度的（非流式：响应头要等整段生成完才到，没有第一个 token）、取消的，都不 - /// 拿来算。**非流式的混进来会把数顶上天**：总耗时减去响应头只剩读 body 的几毫秒 + /// 在跑的请求各自的速度加起来；结局之后它就不在里面了 #[test] - fn only_requests_with_a_rate_make_the_rate() { + fn the_speed_is_the_sum_of_the_running_requests() { let b = EventBus::new(); - let now = Instant::now(); - b.track_at(&started(1), now); + let t = Instant::now(); + b.track_at(&started(1), t); + b.track_at(&started(2), t); + let (m1, m2) = (b.meter(1).unwrap(), b.meter(2).unwrap()); + m1.add_at(50, t + Duration::from_millis(1_000)); + m1.add_at(50, t + Duration::from_millis(2_000)); + m2.add_at(30, t + Duration::from_millis(1_000)); + m2.add_at(30, t + Duration::from_millis(2_000)); + let live = b.live_at(t + Duration::from_millis(2_000)); + assert_eq!(live.running[0].tokens_per_sec, Some(100)); + assert_eq!(live.running[0].output_tokens, Some(100)); + assert_eq!(live.running[1].tokens_per_sec, Some(60)); + assert_eq!(live.tokens_per_sec, Some(160)); b.track_at( - &tw_api::Event::RequestHeaders { - id: 1, - status: 200, - ttfb_ms: 1_995, - }, - now, + &finished(1, 2_000, Some(50)), + t + Duration::from_millis(2_000), ); - b.track_at(&finished(1, 2_000, None), now); - b.track_at(&started(2), now); - b.track_at(&first_token(2, 100), now); - b.track_at( - &tw_api::Event::RequestCancelled { - id: 2, - model: "m".into(), - status: Some(200), - sent_bytes: Some(0), - received_bytes: Some(1), - duration_ms: 2_000, - usage: Some(tw_api::UsageView { - output: 100, - ..Default::default() - }), - answered_model: None, - }, - now, - ); - let live = b.live_at(now); - assert_eq!(live.tokens_per_sec, None); - assert!(live.running.is_empty()); + assert!(b.meter(1).is_none()); + let live = b.live_at(t + Duration::from_millis(2_000)); + assert_eq!(live.tokens_per_sec, Some(60)); } /// 菜单的「进行中」要知道是谁、用什么模型、跑了多久；开始得早的在前 @@ -761,6 +704,9 @@ mod tests { group: Some("__all__".into()), // 路由还没报出结论：不知道最后是哪家接下的 upstream: None, + // 还没开口 + tokens_per_sec: None, + output_tokens: None, } ); } diff --git a/crates/tw-observe/src/lib.rs b/crates/tw-observe/src/lib.rs index dd312d29..72373de3 100644 --- a/crates/tw-observe/src/lib.rs +++ b/crates/tw-observe/src/lib.rs @@ -5,5 +5,7 @@ //! 一个没人观测的网关必须能照常转发。 pub mod bus; +pub mod meter; pub use bus::{EventBus, RecordFeed}; +pub use meter::Meter; diff --git a/crates/tw-observe/src/meter.rs b/crates/tw-observe/src/meter.rs new file mode 100644 index 00000000..0bdac9f1 --- /dev/null +++ b/crates/tw-observe/src/meter.rs @@ -0,0 +1,169 @@ +//! 一个在跑的请求此刻生成得多快（`/live` 的速度，见 `tw_api::LiveView::tokens_per_sec`）。 +//! +//! 数据面每读到一段输出就记一笔（[`Meter::add`]），记的是估出来的 token 数；`/live` 来问时 +//! 看最近 [`WINDOW`] 里记了多少。**按格记，不按笔记**：一条流一秒几十笔，留着每一笔 +//! 既费内存又要逐笔加；按 [`SLOT`] 一格攒着，窗口里最多 [`SLOTS`] 格。 +//! +//! 一个请求一份，开始事件时由总线建好（见 `EventBus::meter`），结局时跟着请求一起划掉。 + +use std::collections::VecDeque; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +/// 速度看最近多久 +pub const WINDOW: Duration = Duration::from_millis(tw_api::LIVE_RATE_WINDOW_MS); +/// 一格多长 +const SLOT: Duration = Duration::from_millis(250); +/// 窗口里有几格 +const SLOTS: u64 = tw_api::LIVE_RATE_WINDOW_MS / 250; + +#[derive(Debug)] +pub struct Meter { + inner: Mutex, +} + +#[derive(Debug)] +struct Inner { + /// 格号从这一刻数起 + origin: Instant, + /// 第一笔的时刻。**之前没有速度**：还在等第一个 token + first: Option, + /// 到目前为止一共记了多少 + total: u64, + /// （格号，这一格里记了多少），格号从小到大，只留窗口里的 + slots: VecDeque<(u64, u64)>, +} + +/// 此刻的读数 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Reading { + pub tokens_per_sec: u32, + pub total: u64, +} + +impl Meter { + pub fn new(now: Instant) -> Self { + Self { + inner: Mutex::new(Inner { + origin: now, + first: None, + total: 0, + slots: VecDeque::new(), + }), + } + } + + /// 又生成了 `tokens` 个。**0 也算开了口**：第一个 token 的那一段可能估出 0 个（一个 + /// 空格），之后的速度要从这一刻算起 + pub fn add(&self, tokens: u64) { + self.add_at(tokens, Instant::now()); + } + + pub fn add_at(&self, tokens: u64, now: Instant) { + let mut m = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + m.first.get_or_insert(now); + m.total = m.total.saturating_add(tokens); + let slot = m.slot(now); + match m.slots.back_mut() { + Some((s, n)) if *s == slot => *n = n.saturating_add(tokens), + _ => m.slots.push_back((slot, tokens)), + } + m.forget(slot); + } + + /// 到目前为止一共记了多少 + pub fn total(&self) -> u64 { + self.inner.lock().unwrap_or_else(|p| p.into_inner()).total + } + + /// 此刻的速度：窗口里记的除以窗口的长度 —— 开口还不到一个窗口的，除以开口以来的时间 + /// （至少一格，刚开口那一下不至于除出个天文数字）。还没开口的没有读数 + pub fn read_at(&self, now: Instant) -> Option { + let mut m = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + let first = m.first?; + let slot = m.slot(now); + m.forget(slot); + let tokens: u64 = m.slots.iter().map(|(_, n)| n).sum(); + let span = now + .saturating_duration_since(first) + .clamp(SLOT, WINDOW) + .as_millis() as u64; + Some(Reading { + tokens_per_sec: (tokens.saturating_mul(1000) / span).min(u64::from(u32::MAX)) as u32, + total: m.total, + }) + } +} + +impl Inner { + fn slot(&self, now: Instant) -> u64 { + (now.saturating_duration_since(self.origin).as_millis() / SLOT.as_millis()) as u64 + } + + /// 掉出窗口的格扔掉：此刻在第 `slot` 格，留它和它之前的 `SLOTS - 1` 格 + fn forget(&mut self, slot: u64) { + while self.slots.front().is_some_and(|(s, _)| *s + SLOTS <= slot) { + self.slots.pop_front(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ms(n: u64) -> Duration { + Duration::from_millis(n) + } + + #[test] + fn nothing_before_the_first_token() { + let t = Instant::now(); + let m = Meter::new(t); + assert_eq!(m.read_at(t + ms(5_000)), None); + // 开了口、估出 0 个：从这一刻起有速度，是 0 + m.add_at(0, t + ms(5_000)); + assert_eq!( + m.read_at(t + ms(6_000)), + Some(Reading { + tokens_per_sec: 0, + total: 0 + }) + ); + } + + /// 开口不到一个窗口：除以开口以来的时间，不是整个窗口 —— 不然刚开口的一秒只显示三分之一 + #[test] + fn a_fresh_stream_is_divided_by_the_time_since_it_opened() { + let t = Instant::now(); + let m = Meter::new(t); + m.add_at(10, t + ms(1_000)); + m.add_at(40, t + ms(1_500)); + m.add_at(50, t + ms(2_000)); + assert_eq!(m.read_at(t + ms(2_000)).unwrap().tokens_per_sec, 100); + } + + /// 窗口之外的不算；停下来不吐字的时候一路掉到 0 + #[test] + fn the_window_forgets_and_a_stalled_stream_falls_to_zero() { + let t = Instant::now(); + let m = Meter::new(t); + for i in 0..20 { + m.add_at(30, t + ms(i * 250)); + } + // 第 5 秒：窗口是 2 秒到 5 秒那 12 格，每格 30 个 + let r = m.read_at(t + ms(4_999)).unwrap(); + assert_eq!(r.tokens_per_sec, 120); + assert_eq!(r.total, 600); + assert_eq!(m.read_at(t + ms(20_000)).unwrap().tokens_per_sec, 0); + assert_eq!(m.read_at(t + ms(20_000)).unwrap().total, 600); + } + + #[test] + fn the_very_first_instant_is_not_divided_by_zero() { + let t = Instant::now(); + let m = Meter::new(t); + m.add_at(5, t); + assert_eq!(m.read_at(t).unwrap().tokens_per_sec, 20); + } +}