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); + } +}