From 63fb9dddc217c3a379c30483ced9ce34bf808ac4 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:02:50 +0800 Subject: [PATCH 01/16] feat(dryrun): each load-balanced candidate carries its share of new conversations The desktop app's dry-run dialog worked out each member's share itself (weight x balance factor, skipping tripped upstreams). That is the data plane's algorithm copied into the UI, and the copies had already drifted: it ignored members whose concurrency limit is full, which the engine leaves out of the round, and it multiplied the raw factor where the engine rounds the effective weight to thousandths with a floor of one. tw-engine gains `weighted::shares`, built on the same `round` that picks the next leader: each member's effective weight over the round's total, None for members sitting the round out (paused or busy, unless all are). Smooth weighted round-robin leads each member exactly its effective weight times per cycle, so this is the long-run split the data plane produces from the current numbers. `tw_api::DryRunCandidate` gains `share: Option` (0..=1), present only for load-balance groups; the shares of one group sum to 1. Co-Authored-By: Claude Opus 5.5 --- crates/tw-api/src/lib.rs | 6 +++ crates/tw-control/src/dryrun.rs | 24 ++++++--- crates/tw-control/tests/dryrun.rs | 82 +++++++++++++++++++++++++++++++ crates/tw-engine/src/weighted.rs | 62 +++++++++++++++++++++++ 4 files changed, 167 insertions(+), 7 deletions(-) diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index 122ae583..6feafc87 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -5066,6 +5066,12 @@ pub struct DryRunCandidate { /// 分得少，没有样本的那一项算 1。只在 `balance_by` 不是 `weights` 的 `load-balance` 里有 #[serde(default, skip_serializing_if = "Option::is_none")] pub balance_factor: Option, + /// 照此刻的数字轮下去，它分到的新对话占这个组的多少，0 到 1：有效权重（权重 × 系数） + /// 除以这一轮参加的各家之和，和数据面挑排头是同一份数字。只在 `load-balance` 里有， + /// 同一个组里有份额的各家加起来是 1。熔断着、冷却着、并发数满着的这一轮不参加，没有 + /// 份额 —— 全都这样时都参加 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub share: Option, } /// 一个要转换格式的候选上游。 diff --git a/crates/tw-control/src/dryrun.rs b/crates/tw-control/src/dryrun.rs index 98f2563d..6ea33aca 100644 --- a/crates/tw-control/src/dryrun.rs +++ b/crates/tw-control/src/dryrun.rs @@ -51,17 +51,18 @@ fn order_like_the_data_plane( } /// 一家候选的排序依据：典型的快慢（从发出去到回答的第一段内容）、最近的成功率、按它们算出 -/// 的系数。**只给顺序真用到的那几样**（[`runtime_facts`] 只取用得上的）—— `url-test` 看快慢； -/// `load-balance` 按 `balance_by` 看快慢、成败，系数乘在权重上。用不着的、没有样本的 -/// 是 None。 +/// 的系数，和 `load-balance` 照这些数字分给它的份额。**只给顺序真用到的那几样** +/// （[`runtime_facts`] 只取用得上的）—— `url-test` 看快慢；`load-balance` 按 `balance_by` +/// 看快慢、成败，系数乘在权重上。用不着的、没有样本的是 None。 struct Basis { ttfb_ms: Option, success_rate: Option, balance_factor: Option, + share: Option, } -/// 每一家候选的排序依据，按 `d.candidates` 的名字查。系数和数据面一样按这份候选、 -/// 这份数字算（`tw_engine::balance_factors`） +/// 每一家候选的排序依据，按 `d.candidates` 的名字查。系数、份额和数据面一样按这份候选、 +/// 这份数字算（`tw_engine::balance_factors`、`tw_engine::weighted::shares`） fn bases( d: &tw_engine::Decision, runtime: Option<&(&tw_engine::Group, tw_engine::Facts)>, @@ -69,11 +70,18 @@ fn bases( let Some((g, f)) = runtime else { return Default::default(); }; - let factors = if g.kind == tw_engine::GroupType::LoadBalance && !g.balance_by.is_weights() { + let balanced = g.kind == tw_engine::GroupType::LoadBalance; + let factors = if balanced && !g.balance_by.is_weights() { tw_engine::balance_factors(g.balance_by, &d.candidates, f) } else { Vec::new() }; + // 份额只按权重分时也有：权重本身就是比例 + let shares = if balanced { + tw_engine::weighted::shares(g, &d.candidates, f) + } else { + Vec::new() + }; d.candidates .iter() .enumerate() @@ -82,6 +90,7 @@ fn bases( ttfb_ms: f.ttfb_ms.get(name).copied(), success_rate: f.success.get(name).copied(), balance_factor: factors.get(i).copied(), + share: shares.get(i).copied().flatten(), }; (name.clone(), basis) }) @@ -325,7 +334,7 @@ pub async fn dry_run( // `load-balance` 的候选带上各自的权重：排头的为什么是它，一半在这个数里 let balanced = group.filter(|g| g.kind == tw_engine::GroupType::LoadBalance); // 每一家收到的模型名，和为什么不是请求里写的那个；顺序看运行时数字的，再加上 - // 每一家的那几个数字（权重、快慢、成功率、系数）—— 「为什么轮到它」 + // 每一家的那几个数字（权重、快慢、成功率、系数、份额）—— 「为什么轮到它」 // 要能从这里看出来 let mut basis_of = bases(&d, runtime.as_ref()); out.candidate_models = out @@ -344,6 +353,7 @@ pub async fn dry_run( ttfb_ms: basis.as_ref().and_then(|b| b.ttfb_ms), success_rate: basis.as_ref().and_then(|b| b.success_rate), balance_factor: basis.as_ref().and_then(|b| b.balance_factor), + share: basis.as_ref().and_then(|b| b.share), } }) .collect(); diff --git a/crates/tw-control/tests/dryrun.rs b/crates/tw-control/tests/dryrun.rs index 7cdebd2d..cefc9aac 100644 --- a/crates/tw-control/tests/dryrun.rs +++ b/crates/tw-control/tests/dryrun.rs @@ -804,6 +804,88 @@ async fn only_the_numbers_the_order_uses_are_shown() { ); } +/// 每一家的份额：`load-balance` 组里各家照此刻的数字分到的新对话的比例，加起来是 1。 +/// 界面不再自己拿权重和系数去乘 —— 那是数据面的算法，各算各的迟早对不上（停着的、满着的 +/// 哪些不参加，系数取整之后是多少） +#[tokio::test] +async fn each_balanced_candidate_carries_the_share_the_data_plane_gives_it() { + let shares = |r: &tw_api::DryRunResult| { + r.candidate_models + .iter() + .map(|c| (c.provider.clone(), c.share)) + .collect::>() + }; + let weighted = CFG.replace( + " providers: [官方, 中转]\n", + " providers: [{ name: 官方, weight: 3 }, 中转]\n", + ); + + // 只按权重分：份额就是权重的比例 + let (_d, app, gw) = app_and_gateway(&weighted); + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + let s = shares(&r); + assert_eq!((s["官方"], s["中转"]), (Some(0.75), Some(0.25)), "{r:?}"); + + // 熔断着的这一轮不参加：没有份额，剩下的那一家拿全部 + while gw.health.is_available("中转") { + gw.health.record_failure("中转"); + } + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + assert_eq!(r.circuit_open, ["中转"]); + let s = shares(&r); + assert_eq!((s["官方"], s["中转"]), (Some(1.0), None), "{r:?}"); + // 全都熔断着：都参加，照权重分（数据面照样一家家试） + while gw.health.is_available("官方") { + gw.health.record_failure("官方"); + } + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + let s = shares(&r); + assert_eq!((s["官方"], s["中转"]), (Some(0.75), Some(0.25)), "{r:?}"); + + // 并发数满着的同样不参加 + let busy = weighted.replace( + " base_url: https://relay.example.com\n", + " base_url: https://relay.example.com\n max_concurrent: 1\n", + ); + let (_d, app, gw) = app_and_gateway(&busy); + let held = gw.slots.try_take("中转").expect("中转有空位"); + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + let s = shares(&r); + assert_eq!((s["官方"], s["中转"]), (Some(1.0), None), "{r:?}"); + drop(held); + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + assert_eq!(shares(&r)["中转"], Some(0.25), "空出来之后回到轮里"); + + // 按快慢、成败分：有效权重是权重 × 系数放大一千倍取整，份额按它算 + let (_d, app, gw) = app_and_gateway(&weighted.replace( + " type: load-balance\n", + " type: load-balance\n balance_by: latency\n", + )); + for _ in 0..3 { + gw.latency.record("官方", 400); + gw.latency.record("中转", 100); + } + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + // 中位数 250：官方 3 × (250/400)² → 1172，中转 1 × (250/100)² → 6250 + let s = shares(&r); + let (a, b) = (s["官方"].unwrap(), s["中转"].unwrap()); + assert!((a - 1172.0 / 7422.0).abs() < 1e-12, "{r:?}"); + assert!((b - 6250.0 / 7422.0).abs() < 1e-12, "{r:?}"); + assert!((a + b - 1.0).abs() < 1e-12); + + // 不经过 `load-balance` 的没有份额：指定的上游、按顺序的组 + let r = run(&app, r#"{"model":"claude-sonnet-4-5","cache":true}"#).await; + assert_eq!(shares(&r)["官方"], None); + for kind in ["fallback", "url-test", "cheapest"] { + let (_d, app) = app_with(&CFG.replace("type: load-balance", &format!("type: {kind}"))); + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + assert!( + r.candidate_models.iter().all(|c| c.share.is_none()), + "{kind}: {r:?}" + ); + } +} + /// 按快慢、成败分时，试算说的排头就是数据面下一个新对话真的去的那一家：同一份延迟表、 /// 成功率、轮询状态，同一个有效权重。起真网关、真上游，每发一个新对话之前先试算一次， /// 一次都不能对不上 —— 真请求会添新的首字节样本和成败，系数一直在变，对得上才说明两边 diff --git a/crates/tw-engine/src/weighted.rs b/crates/tw-engine/src/weighted.rs index 951c5922..f50c6fb0 100644 --- a/crates/tw-engine/src/weighted.rs +++ b/crates/tw-engine/src/weighted.rs @@ -238,6 +238,28 @@ pub fn advance(g: &Group, members: &[String], f: &Facts, leader: &str) -> HashMa out } +/// 长期看每一家排头的份额，0 到 1，和 `members` 一一对应：**这一轮参加的**（[`round`]， +/// 停着的、满着的不算，全不能参加时都算）按各自的有效权重（[`effective`]）分，加起来是 1； +/// 不参加的是 `None`。 +/// +/// 平滑加权轮询每排有效权重之和那么多次，各家正好各排头自己的有效权重那么多次，所以份额 +/// 就是有效权重 ÷ 这一轮的总和。和 [`lead`] 读的是同一份 [`round`]：试算说的份额，就是 +/// 数据面照此刻的数字轮下去各家分到的新对话的比例（进行中的对话留在原来那一家，那一份 +/// 由之后的新对话补回来，见 [`advance`]）。 +pub fn shares(g: &Group, members: &[String], f: &Facts) -> Vec> { + let round = round(g, members, f); + let total: i64 = round.iter().map(|(_, w)| w).sum(); + members + .iter() + .map(|m| { + round + .iter() + .find(|(r, _)| *r == m.as_str()) + .map(|(_, w)| *w as f64 / total as f64) + }) + .collect() +} + #[cfg(test)] mod tests { use super::*; @@ -563,6 +585,46 @@ mod tests { assert_eq!((got["慢"], got["快"]), (1172, 1563)); } + /// 份额就是一整圈里各家排头的比例：轮一整圈数出来的次数 ÷ 圈长，和 [`shares`] 给的一样。 + /// 停着的、满着的不参加，份额是空的，剩下的加起来还是 1；全都不能参加时都算 + #[test] + fn the_share_is_how_often_each_member_leads_over_a_full_round() { + let g = balanced(&[("慢", 3), ("快", 1), ("停", 2)], BalanceBy::Latency); + let members = names(&g); + let mut f = Facts { + ttfb_ms: ttfb(&[("慢", 400), ("快", 100)]), + paused: ["停".to_string()].into_iter().collect(), + ..Default::default() + }; + let got = shares(&g, &members, &f); + assert_eq!(got[2], None, "停着的不参加"); + // 中位数 250：慢 3 × 0.390625 → 1172，快 1 × 6.25 → 6250 + let n = 1172 + 6250; + let led = count(&g, &mut f, n); + for (i, m) in members.iter().take(2).enumerate() { + let want = led[m] as f64 / n as f64; + assert!((got[i].unwrap() - want).abs() < 1e-12, "{m}: {got:?}"); + } + assert!((got[0].unwrap() + got[1].unwrap() - 1.0).abs() < 1e-12); + + // 满着的同样不参加 + let f = Facts { + busy: ["快".to_string()].into_iter().collect(), + ..Default::default() + }; + assert_eq!(shares(&g, &members, &f), [Some(0.6), None, Some(0.4)]); + // 全都停着、满着：都参加，按权重分 + let f = Facts { + paused: ["慢".to_string(), "停".to_string()].into_iter().collect(), + busy: ["快".to_string()].into_iter().collect(), + ..Default::default() + }; + assert_eq!( + shares(&g, &members, &f), + [Some(0.5), Some(1.0 / 6.0), Some(1.0 / 3.0)] + ); + } + /// 各家系数一样时，次序和只按权重分一模一样：7:3 照样穿插着来，不是先七后三 #[test] fn equal_factors_keep_the_exact_interleaving() { From 876e1343865214403bb3f938a902618bc3c77301 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:16:11 +0800 Subject: [PATCH 02/16] fix(store): body GC trims the newest day instead of emptying it Over the size cap, GC deleted whole days oldest-first and did not stop at the newest one. A single day can exceed body_max_bytes (each request keeps up to 4 MiB of request and 4 MiB of response), and then every hourly pass wiped today's directory, including the session being viewed and the request that had just finished. Older days still go whole. Once only the newest day is left over the cap, its oldest requests are removed, all files of one request together (request id order is start order), until the total fits. Bodies are now written to a temporary file and renamed into place, with the `.len` marker written first. Readers (detail, transcript, search) do not take the recorder's lock, so they could read a body halfway through being written and take it for an unreadable or complete one. Co-Authored-By: Claude Opus 5.5 --- crates/tw-store/src/blobs.rs | 160 +++++++++++++++++++++++++++++------ 1 file changed, 136 insertions(+), 24 deletions(-) diff --git a/crates/tw-store/src/blobs.rs b/crates/tw-store/src/blobs.rs index 4644a9ee..f1bae15b 100644 --- a/crates/tw-store/src/blobs.rs +++ b/crates/tw-store/src/blobs.rs @@ -3,6 +3,10 @@ //! **不进 SQLite。**body 可能几百 KB（长上下文），塞进库里会让它膨胀到 //! GB 级 —— VACUUM 慢、备份慢、回收难。存文件系统按天分目录，过期直接 //! 删掉整个目录，**回收成本 O(1)**。 +//! +//! **一份正文要么整个在、要么不在**：先写到旁边一个临时文件，写完再换上去（见 +//! [`Blobs::put`]）。读的一方（详情、对话记录、按正文找）和写的一方不排队，以前读到一半 +//! 写着的文件会当成一份读不懂的正文。 use std::path::{Path, PathBuf}; @@ -99,15 +103,8 @@ impl Blobs { // 截断而不是跳过：**开头那几 KB 是最有用的部分**（模型名、system // prompt、工具定义都在前面），而完整存下来会挤掉别人的。 let slice = &body[..body.len().min(MAX_ONE)]; - match write_private(&p, slice) { - Ok(()) => { - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let _ = std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o600)); - } - true - } + match replace_private(&p, slice) { + Ok(()) => true, Err(e) => { tracing::debug!(path = %p.display(), "the request body could not be written: {e}"); false @@ -131,6 +128,9 @@ impl Blobs { /// `original_len` 比**真正存下的**长时才另记一个 `.len` —— 截断可能发生在交来之前 /// （网关只攒了开头），也可能发生在这里（`body` 比 [`MAX_ONE`] 长）。以前只看前一种： /// 一个 5 MB 的请求体存下 4 MB，却没有一处说它被截过，重放照样把半截 JSON 发了出去。 + /// + /// **`.len` 先写**：正文一出现，读的一方就该知道它是不是截过的。反过来的话，中间有一刻 + /// 读到的是一份看起来完整的半截正文。 pub fn put_with_len( &self, at_ms: i64, @@ -139,22 +139,17 @@ impl Blobs { body: &[u8], original_len: usize, ) -> bool { - if !self.put(at_ms, id, which, body) { - return false; - } if original_len > body.len().min(MAX_ONE) { let p = self .path_for(at_ms, id, which) .with_extension(format!("{}.len", which.suffix())); - if write_private(&p, original_len.to_string().as_bytes()).is_ok() { - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let _ = std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o600)); - } + if let Some(dir) = p.parent() + && create_private_dirs(dir).is_ok() + { + let _ = replace_private(&p, original_len.to_string().as_bytes()); } } - true + self.put(at_ms, id, which, body) } /// 原始长度。**没有这个文件说明没截断** —— 那时 body 自己的长度 @@ -186,13 +181,18 @@ impl Blobs { /// 盘上最老的那一天从哪一刻起（那天零点，UTC）。一天都没有是 `None`。 /// /// **按正文找时用它划界**：比它早的请求一份正文都不会有，不必一条一条去盘上问 —— - /// 三个月的记录里，正文只占最近几天。回收是整天整天地删（先删最老的），所以划在 - /// 最老的那一天上是准的。 + /// 三个月的记录里，正文只占最近几天。回收先删最老的，整天整天地删；只剩最新的一天时 + /// 才从那一天里最早的请求删起。所以比它早的一份都没有，它之后的照样要去盘上问。 pub fn oldest_ms(&self) -> Option { self.days().first().and_then(|(name, _)| start_of_day(name)) } /// 回收：先按天数删，再按总量删。返回删掉的字节数。 + /// + /// **按总量删时，最新的那一天不整个删。**从最老的那天整天整天地删；删到只剩最新的一天 + /// 还超，就在这一天里从最早的请求删起（[`trim_oldest`]）。以前它也整个删：一天的正文就 + /// 超过上限时（一个请求最多存 8 MB 多），每小时的回收都把当天的目录清空一次，连同正在 + /// 看的那次会话。 pub fn gc(&self, now_ms: i64, keep_days: u64, max_bytes: u64) -> u64 { let cutoff = day_of(now_ms - (keep_days as i64) * 86_400_000); let mut freed = 0; @@ -207,10 +207,13 @@ impl Blobs { remaining.push((name, path, size)); } } - // 还超总量的话，继续从最旧的开始删。**整目录删，不删单个文件** —— - // 半天的记录比没有记录更难解释（「为什么上午的请求点开是空的」）。 let mut total: u64 = remaining.iter().map(|(_, _, s)| *s).sum(); - for (_, path, size) in &remaining { + let Some(((_, newest, _), older)) = remaining.split_last() else { + return freed; + }; + // 还超总量的话，继续从最旧的开始删。**旧的整目录删，不删单个文件** —— + // 半天的记录比没有记录更难解释（「为什么上午的请求点开是空的」）。 + for (_, path, size) in older { if total <= max_bytes { break; } @@ -219,6 +222,9 @@ impl Blobs { freed += size; } } + if total > max_bytes { + freed += trim_oldest(newest, total - max_bytes); + } freed } @@ -254,6 +260,60 @@ fn write_private(p: &Path, bytes: &[u8]) -> std::io::Result<()> { opts.open(p)?.write_all(bytes) } +/// 换上一份只给自己看的文件：**先写到旁边的临时文件，写完再改名换上去**。改名是一步完成 +/// 的，读的一方要么读到旧的（或者没有），要么读到写完的那一份，不会读到一半。 +fn replace_private(p: &Path, bytes: &[u8]) -> std::io::Result<()> { + let mut tmp = p.as_os_str().to_owned(); + tmp.push(".tmp"); + let tmp = PathBuf::from(tmp); + let done = write_private(&tmp, bytes).and_then(|()| std::fs::rename(&tmp, p)); + if done.is_err() { + let _ = std::fs::remove_file(&tmp); + } + #[cfg(unix)] + if done.is_ok() { + // 临时文件早就在（上次写到一半断了）的话，新建时给的 0600 不生效，这里补上 + use std::os::unix::fs::PermissionsExt; + let _ = std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o600)); + } + done +} + +/// 在一天的目录里从最早的请求删起，删够 `excess` 字节为止。返回删掉的字节数。 +/// +/// **一个请求的几份一起删**（请求体、回答、插件改过的、`.len`）：按文件名开头的请求号归在 +/// 一起，号小的先删 —— 号是按请求开始的先后发的。认不出请求号的文件不动。 +fn trim_oldest(day: &Path, excess: u64) -> u64 { + let Ok(rd) = std::fs::read_dir(day) else { + return 0; + }; + let mut by_id: std::collections::BTreeMap> = Default::default(); + for e in rd.flatten() { + let name = e.file_name(); + let Some(id) = name + .to_str() + .and_then(|n| n.split_once('.')) + .and_then(|(id, _)| id.parse::().ok()) + else { + continue; + }; + let size = e.metadata().map(|m| m.len()).unwrap_or(0); + by_id.entry(id).or_default().push((e.path(), size)); + } + let mut freed = 0; + for files in by_id.into_values() { + if freed >= excess { + break; + } + for (path, size) in files { + if std::fs::remove_file(&path).is_ok() { + freed += size; + } + } + } + freed +} + fn looks_like_a_day(n: &str) -> bool { let b = n.as_bytes(); b.len() == 10 @@ -463,6 +523,58 @@ mod tests { assert_eq!(left[0].0, day_of(now)); } + /// **一天的正文就超过上限时，那一天不整个删。**以前每小时的回收都把当天的目录清空： + /// 正在看的那次会话、刚结束的那个请求，点开都是空的。现在删到只剩最新的一天，就在这一天 + /// 里从最早的请求删起，一个请求的几份一起删 + #[test] + fn gc_trims_the_newest_day_from_its_oldest_requests_instead_of_emptying_it() { + let (_d, b) = setup(); + let now = 100 * DAY; + // 前一天的整个删 + b.put(now - DAY, 1, Which::Request, &vec![b'x'; 1000]); + // 当天五个请求，各有请求体和回答，2000 字节一个；第 10 个还记着原本多长 + for id in 10..15 { + b.put(now + id, id, Which::Request, &vec![b'x'; 1000]); + b.put(now + id, id, Which::Response, &vec![b'x'; 1000]); + } + b.put_with_len(now, 10, Which::Request, &vec![b'x'; 1000], 9999); + let today = b.root().join(day_of(now)); + let before = b.total_bytes(); + + let freed = b.gc(now, 30, 3500); + + assert_eq!(b.days().len(), 1, "{:?}", b.days()); + assert!(today.is_dir(), "当天的目录被整个删了"); + assert!(b.total_bytes() <= 3500, "{}", b.total_bytes()); + assert_eq!(freed, before - b.total_bytes()); + // 最早的四个连同 `.len` 一起没了，最新的那个还在 + for id in 10..14 { + assert!(b.get(now, id, Which::Request).is_none(), "{id}"); + assert!(b.get(now, id, Which::Response).is_none(), "{id}"); + } + assert_eq!(b.original_len(now, 10, Which::Request), None); + assert!(b.get(now, 14, Which::Request).is_some()); + assert!(b.get(now, 14, Which::Response).is_some()); + + // 下一小时没有新写入：已经在上限之内，什么都不再删 + assert_eq!(b.gc(now + 3_600_000, 30, 3500), 0); + assert!(b.get(now, 14, Which::Response).is_some()); + } + + /// 写到一半的正文读不到：先写临时文件再换上去，换完不留临时文件 + #[test] + fn a_body_is_swapped_in_whole_and_leaves_no_temporary_file_behind() { + let (_d, b) = setup(); + assert!(b.put(0, 1, Which::Response, b"first")); + assert!(b.put(0, 1, Which::Response, b"second, longer")); + assert_eq!(b.get(0, 1, Which::Response).unwrap(), b"second, longer"); + let names: Vec = std::fs::read_dir(b.root().join(day_of(0))) + .unwrap() + .map(|e| e.unwrap().file_name().to_string_lossy().to_string()) + .collect(); + assert_eq!(names, ["1.res"]); + } + #[test] fn a_stray_directory_the_user_put_there_is_left_alone() { // 这是用户自己的磁盘。**不认识的东西不删。** From d6b6786d076adbff3087bd8593650c1f5251666d Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:16:11 +0800 Subject: [PATCH 03/16] perf(store): index-bounded session list, key last-seen and route hits Measured on a synthetic 270k-row, 222 MB request database (90 days, release build; `history_cost` reproduces it): - sessions(None, 200): 160-200 ms -> 12 ms. It grouped every row of the retention period and sorted the groups. It now walks a partial index (at_ms, session) newest-first until it has `limit` sessions, checks a session's first turn only when its last one is after the window, and aggregates just those sessions through the session index. A test compares it with the old grouping query across windows and limits. - Db::session(id): direct lookup for the session detail (was: take the newest 500 sessions and search them). - last_seen_by_client: 78 ms -> 0.02 ms, a loose index scan over a partial (client, at_ms) index. The index is partial (`client > ''`) so the windowed GROUP BY client queries keep walking the time index instead of scanning it in key order. - route_hits: 22 ms -> 6 ms, and it no longer reads 7.6 MB of routing JSON: route, rule, rewrites and denial live in an expression index (json_extract, written at insert), so the query reads only the index. Rows whose routing does not parse stay out of it and are still written. The indexes are created with CREATE INDEX IF NOT EXISTS on every open, not through SCHEMA: they do not change what a row looks like, other versions read and write the database as before, and a SCHEMA bump would wipe everyone's request history. First open of an existing 270k-row database builds them in about a second; the file grows ~15%. Co-Authored-By: Claude Opus 5.5 --- crates/tw-store/src/db.rs | 718 +++++++++++++++++++++++++++++++++----- 1 file changed, 638 insertions(+), 80 deletions(-) diff --git a/crates/tw-store/src/db.rs b/crates/tw-store/src/db.rs index b809c39c..b55c36db 100644 --- a/crates/tw-store/src/db.rs +++ b/crates/tw-store/src/db.rs @@ -50,6 +50,20 @@ const NO_PRICE: &str = "(cost_micros IS NULL AND input_tokens IS NOT NULL \ const NO_USAGE: &str = "(cost_micros IS NULL AND input_tokens IS NULL AND error IS NULL \ AND billing = 'per-token' AND (cancelled = 1 OR status < 300))"; +/// 数路由命中要的那几项，从路由那一列（`tw_api::RoutingView` 的 JSON）里取（见 +/// [`Db::route_hits`]）。**建索引的和查询的是同一串**：表达式一字不差，SQLite 才用索引里 +/// 存好的值，不回表去解 JSON。 +/// +/// 用 `json_extract` 不用 `->>`：索引写在库文件里，读这个库的 SQLite 都得认得它。 +const ROUTE: &str = "json_extract(routing, '$.route')"; +const RULE: &str = "json_extract(routing, '$.rule')"; +/// 一个 JSON 数组的原文（`[]`、`["关思考"]`） +const REWRITTEN_BY: &str = "json_extract(routing, '$.rewritten_by')"; +const DENIED_BY: &str = "json_extract(routing, '$.denied_by')"; +/// 经过路由、路由那一列解得开的行。**解不开的不进索引**：一行坏掉的记录不该让整张表出不来， +/// 也不该让它自己写不进去 +const ROUTED: &str = "local = 0 AND json_valid(routing)"; + #[derive(Debug, thiserror::Error)] pub enum DbError { #[error("{path} could not be opened: {source}")] @@ -272,9 +286,42 @@ impl Db { if found == 0 { db.create()?; } + db.indexes(); Ok(db) } + /// 表之外另加的索引，**每次打开都走一遍**（`IF NOT EXISTS`，建过了就是空操作）。 + /// + /// **它们不算 schema，加它们不加 [`SCHEMA`]**：索引不改变一行长什么样，少了它查询照样 + /// 对、只是慢，多了它别的版本照样读写（SQLite 自己维护）。加 [`SCHEMA`] 却要清掉每个人 + /// 的请求历史 —— 为了快一点不值得。已有的库第一次打开时要建一遍（二十七万行一秒上下）。 + /// + /// **建不成只记一行日志**：那时查询照样对，只是慢。 + fn indexes(&self) { + // - `requests_client`：每把密钥最后一次用在什么时候，一把一把地跳着取（见 + // `last_seen_by_client`）。**只给那一条查询用**：条件里要写上 `client > ''` 才用得上 + // 它。按时间窗、按密钥分组的那些（密钥的用量、按密钥的费用）用了它，会把整个索引按 + // 密钥扫一遍，比沿着时刻走慢得多 + // - `requests_session_at`：最近活动过的会话，按时间倒着走，只读索引（见 `sessions`） + // - `requests_routed`：各条路由、规则命中了多少要的那几项，落库时从路由那一列里取出来 + // 存在索引里。数的时候只读索引，不再把一周几 MB 的 JSON 读出来解一遍（见 + // `route_hits`） + let sql = format!( + "CREATE INDEX IF NOT EXISTS requests_client ON requests (client, at_ms) + WHERE client > ''; + CREATE INDEX IF NOT EXISTS requests_session_at ON requests (at_ms, session) + WHERE session IS NOT NULL AND local = 0; + CREATE INDEX IF NOT EXISTS requests_routed + ON requests (at_ms, {ROUTE}, {RULE}, {REWRITTEN_BY}, {DENIED_BY}, error IS NOT NULL) + WHERE {ROUTED};" + ); + if let Err(e) = self.conn.execute_batch(&sql) { + tracing::warn!( + "the request history could not be indexed, so some views load slowly: {e}" + ); + } + } + /// 建表。**只有这一份，没有迁移**：表的样子一变，改这里、[`SCHEMA`] 加一， /// 版本对不上的旧库整个重建。 /// @@ -535,6 +582,52 @@ pub struct SessionRow { pub errors: i64, } +/// 一次会话汇总的那几列，和 [`session_row`] 一列对一列。`GROUP BY session` 之后用 +fn session_columns() -> String { + format!( + "session, + client, + MIN(at_ms), MAX(at_ms), COUNT(*), + COALESCE(SUM(cost_micros), 0), + COALESCE(SUM({NO_PRICE}), 0), + COALESCE(SUM(input_tokens), 0), + COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(cache_read_tokens), 0), + COALESCE(SUM(cache_write_tokens), 0), + COALESCE(SUM(cache_saved_micros), 0), + COALESCE(MAX(input_tokens), 0), + GROUP_CONCAT(DISTINCT model), + SUM(CASE WHEN error IS NOT NULL THEN 1 ELSE 0 END), + COALESCE(SUM(CASE WHEN cost_estimated = 1 THEN cost_micros ELSE 0 END), 0), + COUNT(cost_micros), + COALESCE(SUM({NO_USAGE}), 0)" + ) +} + +/// 读出 [`session_columns`] 的一行 +fn session_row(r: &rusqlite::Row) -> rusqlite::Result { + Ok(SessionRow { + id: r.get(0)?, + client: r.get(1)?, + started_ms: r.get(2)?, + ended_ms: r.get(3)?, + turns: r.get(4)?, + cost_micros: r.get(5)?, + unpriced_turns: r.get(6)?, + input_tokens: r.get(7)?, + output_tokens: r.get(8)?, + cache_read_tokens: r.get(9)?, + cache_write_tokens: r.get(10)?, + cache_saved_micros: r.get(11)?, + peak_input_tokens: r.get(12)?, + models: r.get::<_, Option>(13)?.unwrap_or_default(), + errors: r.get(14)?, + cost_micros_estimated: r.get(15)?, + priced_turns: r.get(16)?, + no_usage_turns: r.get(17)?, + }) +} + /// 会话里的一轮。上下文增长曲线和成本瀑布画的就是它。 #[derive(Debug, Clone, PartialEq)] pub struct TurnRow { @@ -568,63 +661,93 @@ impl Db { /// /// **筛的是会话，不是轮次。**把轮次按时间筛掉再聚合的话，一次跨过 /// 窗口边界的任务会少算几轮、少算一截钱 —— 而「那次重构花了多少」 - /// 问的是整次任务，不是它落在某个窗口里的那一段。所以先整体聚合， - /// 再按「有没有任何一轮落在窗口里」留下整条。 + /// 问的是整次任务，不是它落在某个窗口里的那一段。所以留下的会话整条聚合， + /// 按「首尾之间和窗口有没有重叠」留。 + /// + /// **先挑会话，再聚合挑中的那几次**（[`Db::recent_sessions`]）。以前把三个月里每一行按 + /// 会话分组、排序之后再截前几百条：二十七万行要两百毫秒，而界面每结束一个请求就来要一次。 pub fn sessions( &self, within: Option<(i64, i64)>, limit: usize, ) -> Result, DbError> { + let picked = self.recent_sessions(within, limit)?; + if picked.is_empty() { + return Ok(Vec::new()); + } + // 挑中的会话号作为一个 JSON 数组交进去：几百个占位符拼不出一条好读的 SQL + let picked = serde_json::Value::from(picked).to_string(); let mut st = self.conn.prepare(&format!( - "SELECT session, - client, - MIN(at_ms), MAX(at_ms), COUNT(*), - COALESCE(SUM(cost_micros), 0), - COALESCE(SUM({NO_PRICE}), 0), - COALESCE(SUM(input_tokens), 0), - COALESCE(SUM(output_tokens), 0), - COALESCE(SUM(cache_read_tokens), 0), - COALESCE(SUM(cache_write_tokens), 0), - COALESCE(SUM(cache_saved_micros), 0), - COALESCE(MAX(input_tokens), 0), - GROUP_CONCAT(DISTINCT model), - SUM(CASE WHEN error IS NOT NULL THEN 1 ELSE 0 END), - COALESCE(SUM(CASE WHEN cost_estimated = 1 THEN cost_micros ELSE 0 END), 0), - COUNT(cost_micros), - COALESCE(SUM({NO_USAGE}), 0) - FROM requests - WHERE session IS NOT NULL AND local = 0 + "SELECT {} FROM requests + WHERE session IN (SELECT value FROM json_each(?1)) AND local = 0 GROUP BY session - HAVING MAX(at_ms) >= ?1 AND MIN(at_ms) <= ?2 - ORDER BY MAX(at_ms) DESC - LIMIT ?3" + ORDER BY MAX(at_ms) DESC, session DESC", + session_columns() ))?; - let (from, to) = within.unwrap_or((i64::MIN, i64::MAX)); - let rows = st.query_map([from, to, limit as i64], |r| { - Ok(SessionRow { - id: r.get(0)?, - client: r.get(1)?, - started_ms: r.get(2)?, - ended_ms: r.get(3)?, - turns: r.get(4)?, - cost_micros: r.get(5)?, - unpriced_turns: r.get(6)?, - input_tokens: r.get(7)?, - output_tokens: r.get(8)?, - cache_read_tokens: r.get(9)?, - cache_write_tokens: r.get(10)?, - cache_saved_micros: r.get(11)?, - peak_input_tokens: r.get(12)?, - models: r.get::<_, Option>(13)?.unwrap_or_default(), - errors: r.get(14)?, - cost_micros_estimated: r.get(15)?, - priced_turns: r.get(16)?, - no_usage_turns: r.get(17)?, - }) - })?; + let rows = st.query_map([picked], session_row)?; Ok(rows.collect::, _>>()?) } + /// 一次会话的汇总。没有这次会话（或者它的每一轮都是本地应答的）是 None。 + /// + /// **按会话号直接取**，不从列表里找：以前取最近五百次会话再从里面挑，更早的会话点开是 404。 + pub fn session(&self, id: &str) -> Result, DbError> { + Ok(self + .conn + .prepare(&format!( + "SELECT {} FROM requests WHERE session = ?1 AND local = 0 GROUP BY session", + session_columns() + ))? + .query_row([id], session_row) + .optional()?) + } + + /// 最近活动过的 `limit` 次会话，最后一次活动最近的在前；`within` 给了就只要首尾之间和它 + /// 重叠的。 + /// + /// **只走索引，不按会话分组整张表**：按时间倒着走（`requests_session_at`，只有带会话、 + /// 不是本地应答的行），一次会话第一次碰到的那一行就是它最后一次活动，碰够了就停。窗口的 + /// 起点之前不必再走：最后一次活动比它早的会话不和窗口重叠。最后一次活动在窗口终点之后的， + /// 再按会话号取它的第一轮（`requests_session`）看是不是不晚于终点。 + fn recent_sessions( + &self, + within: Option<(i64, i64)>, + limit: usize, + ) -> Result, DbError> { + let (from, to) = within.unwrap_or((i64::MIN, i64::MAX)); + let mut out = Vec::new(); + if limit == 0 { + return Ok(out); + } + let mut walk = self.conn.prepare( + "SELECT session, at_ms FROM requests + WHERE session IS NOT NULL AND local = 0 AND at_ms >= ?1 + ORDER BY at_ms DESC, session DESC", + )?; + let mut first = self.conn.prepare( + "SELECT at_ms FROM requests WHERE session = ?1 AND local = 0 + ORDER BY at_ms LIMIT 1", + )?; + let mut seen = std::collections::HashSet::new(); + let mut rows = walk.query([from])?; + while let Some(r) = rows.next()? { + let session = r.get_ref(0)?.as_str().map_err(rusqlite::Error::from)?; + if seen.contains(session) { + continue; + } + let last: i64 = r.get(1)?; + let overlaps = last <= to || first.query_row([session], |r| r.get::<_, i64>(0))? <= to; + seen.insert(session.to_string()); + if overlaps { + out.push(session.to_string()); + if out.len() == limit { + break; + } + } + } + Ok(out) + } + /// 一次会话里的每一轮，**按时间正序** —— 曲线是从左往右画的。 pub fn turns(&self, session: &str) -> Result, DbError> { let mut st = self.conn.prepare( @@ -672,10 +795,22 @@ impl Db { /// 只有带着为那个客户端生成的密钥的请求能证明。按请求头里自报的客户端 /// 标识分组的话，那个标识谁都能写 —— 「这把钥匙还有没有人在用」「接好了 /// 没有」都只能按密钥问。 + /// + /// **一把一把地跳着取**（`requests_client`）：取下一把比上一把大的密钥名，再取它最大的 + /// 时刻，各是一次索引查找。按密钥分组的话要把三个月里每一行过一遍，而密钥只有几把。 + /// 每个子查询都写着 `client > ''`：那个索引只收这些行，条件里没有它就用不上。 pub fn last_seen_by_client(&self) -> Result, DbError> { let mut st = self.conn.prepare( - "SELECT client, MAX(at_ms) FROM requests - WHERE client <> '' GROUP BY client", + "WITH RECURSIVE keys(client) AS ( + SELECT MIN(client) FROM requests WHERE client > '' + UNION ALL + SELECT (SELECT MIN(client) FROM requests WHERE client > keys.client AND client > '') + FROM keys WHERE keys.client IS NOT NULL + ) + SELECT client, + (SELECT MAX(at_ms) FROM requests + WHERE requests.client = keys.client AND client > '') + FROM keys WHERE client IS NOT NULL", )?; let rows = st.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?; Ok(rows.collect::, _>>()?) @@ -1311,15 +1446,10 @@ impl Db { /// **按每一行记下的路由算**，不按现在的配置推：请求走的是它那一刻的路由和 /// 规则。一个请求算在这几条规则上，每条只算一次：决定去向的那一条、附加了 /// 改写的每一条、选定上游之后拒绝了它的那一条。本地应答的没有路由，不算。 + /// + /// **只读索引**（`requests_routed`）：要的那几项落库时就从路由那一列里取出来存在索引里， + /// 路由图开着时每十秒来一次，不再把一周几 MB 的 JSON 读出来一行一行地解。 fn route_hits(&self, since_ms: i64, until_ms: i64) -> Result, DbError> { - /// 数命中要的那几项。尝试链不用解 - #[derive(serde::Deserialize)] - struct Routed { - route: String, - rule: String, - rewritten_by: Vec, - denied_by: Option, - } #[derive(Default)] struct Tally { requests: i64, @@ -1334,44 +1464,52 @@ impl Db { self.last_ms = self.last_ms.max(at_ms); } } - let mut st = self.conn.prepare( - "SELECT at_ms, error IS NOT NULL, routing FROM requests - WHERE at_ms >= ?1 AND at_ms < ?2 AND local = 0 AND routing IS NOT NULL", - )?; - let rows = st.query_map(params![since_ms, until_ms], |r| { - Ok(( - r.get::<_, i64>(0)?, - r.get::<_, bool>(1)?, - r.get::<_, String>(2)?, - )) - })?; - let mut routes: std::collections::BTreeMap< - String, - (Tally, std::collections::BTreeMap), - > = Default::default(); - for row in rows { - let (at_ms, failed, json) = row?; - // 解不开的那一行不算。**一条坏掉的记录不该让整张表出不来** - let Ok(r) = serde_json::from_str::(&json) else { + let mut st = self.conn.prepare(&format!( + "SELECT at_ms, error IS NOT NULL, {ROUTE}, {RULE}, {REWRITTEN_BY}, {DENIED_BY} + FROM requests + WHERE at_ms >= ?1 AND at_ms < ?2 AND {ROUTED}" + ))?; + let mut routes: Named<(Tally, Named)> = Named::default(); + let mut rows = st.query(params![since_ms, until_ms])?; + while let Some(r) = rows.next()? { + let at_ms: i64 = r.get(0)?; + let failed: bool = r.get(1)?; + // 缺了哪一项的那一行不算。**一条坏掉的记录不该让整张表出不来** + let (Ok(route), Ok(rule), Ok(rewritten_by), Ok(denied_by)) = ( + r.get_ref(2)?.as_str(), + r.get_ref(3)?.as_str(), + r.get_ref(4)?.as_str(), + r.get_ref(5)?.as_str_or_null(), + ) else { continue; }; - let (total, rules) = routes.entry(r.route).or_default(); + // 绝大多数请求没有附加改写，不必解 + let rewritten_by: Vec = match rewritten_by { + "[]" => Vec::new(), + j => match serde_json::from_str(j) { + Ok(names) => names, + Err(_) => continue, + }, + }; + let (total, rules) = routes.get(route); total.add(at_ms, failed); - let decider = rules.entry(r.rule.clone()).or_default(); + let decider = rules.get(rule); decider.add(at_ms, failed); decider.decided += 1; - let mut counted = vec![r.rule]; - for name in r.rewritten_by.into_iter().chain(r.denied_by) { + let mut counted = vec![rule]; + for name in rewritten_by.iter().map(String::as_str).chain(denied_by) { if !counted.contains(&name) { - rules.entry(name.clone()).or_default().add(at_ms, failed); + rules.get(name).add(at_ms, failed); counted.push(name); } } } let mut out: Vec = routes + .items .into_iter() .map(|(route, (total, rules))| { let mut rules: Vec = rules + .items .into_iter() .map(|(rule, t)| tw_api::RuleHits { rule, @@ -1412,6 +1550,36 @@ impl Db { } } +/// 按名字数的一组东西。**名字只在第一次见到时拷一份**：数路由命中时一周两万行，每行都按 +/// 名字找一遍（见 [`Db::route_hits`]） +struct Named { + at: std::collections::HashMap, + items: Vec<(String, T)>, +} + +impl Default for Named { + fn default() -> Self { + Self { + at: Default::default(), + items: Vec::new(), + } + } +} + +impl Named { + fn get(&mut self, name: &str) -> &mut T { + let i = match self.at.get(name) { + Some(&i) => i, + None => { + self.at.insert(name.to_string(), self.items.len()); + self.items.push((name.to_string(), T::default())); + self.items.len() - 1 + } + }; + &mut self.items[i].1 + } +} + /// 一个插件在一个请求上的一次运行，落库的样子（`plugin_runs` 一行，少了 `seq`： /// 它在写入时按这个请求已有的行数定）。 #[derive(Debug, Clone, PartialEq)] @@ -2014,6 +2182,87 @@ pub(crate) mod tests { assert!(db.sessions(Some((20_000, 30_000)), 10).unwrap().is_empty()); } + /// 列表是最后一次活动最近的那几次会话，每一次都整条聚合 —— 挑会话时只看了最近的那几行， + /// 聚合的却是它的每一轮。窗口按首尾和它重叠不重叠留：最后一次活动在窗口之后、开始在窗口 + /// 之内的留，整个在窗口之后、之前的不留 + #[test] + fn the_list_is_the_most_recently_active_sessions_each_counted_whole() { + let db = Db::in_memory().unwrap(); + // （会话，时刻）：a 最早开始、最后结束；b 整个在中间；c 在 a 的两头之间断断续续 + for (id, (session, at)) in [ + ("a", 1_000), + ("c", 2_000), + ("b", 3_000), + ("b", 4_000), + ("c", 5_000), + ("c", 6_000), + ("a", 9_000), + ] + .into_iter() + .enumerate() + { + let mut r = row(id as i64 + 1, at); + r.session = Some(session.into()); + db.insert(&r).unwrap(); + } + // 本地应答的不算，哪怕它最新 + let mut probe = row(99, 99_000); + probe.session = Some("b".into()); + probe.local = true; + db.insert(&probe).unwrap(); + + let ids = |got: Vec| got.into_iter().map(|s| s.id).collect::>(); + let two = db.sessions(None, 2).unwrap(); + assert_eq!( + two.iter() + .map(|s| (s.id.as_str(), s.turns)) + .collect::>(), + [("a", 2), ("c", 3)] + ); + assert_eq!(two[0].started_ms, 1_000); + assert_eq!(ids(db.sessions(None, 10).unwrap()), ["a", "c", "b"]); + assert!(db.sessions(None, 0).unwrap().is_empty()); + + // 窗口 [3500, 4500]：三次都和它重叠（a、c 跨过它，b 在里面） + assert_eq!( + ids(db.sessions(Some((3_500, 4_500)), 10).unwrap()), + ["a", "c", "b"] + ); + // [0, 1500]：只有 a 开始得那么早 + assert_eq!(ids(db.sessions(Some((0, 1_500)), 10).unwrap()), ["a"]); + // [6500, 8000]：b、c 那时都结束了，a 还没结束 + assert_eq!(ids(db.sessions(Some((6_500, 8_000)), 10).unwrap()), ["a"]); + // 窗口在所有会话之后 + assert!(db.sessions(Some((10_000, 20_000)), 10).unwrap().is_empty()); + } + + /// 会话详情按会话号直接取：多老、前面压着多少次会话都取得到。以前从最近五百次里找， + /// 更早的点开是 404 + #[test] + fn one_session_is_found_by_its_id_however_old() { + let db = Db::in_memory().unwrap(); + for i in 0..600i64 { + let mut r = row(i + 1, 1_000 + i); + r.session = Some(format!("s{i}")); + db.insert(&r).unwrap(); + } + let mut second = row(1_000, 5_000); + second.session = Some("s0".into()); + db.insert(&second).unwrap(); + + let s0 = db.session("s0").unwrap().unwrap(); + assert_eq!((s0.turns, s0.started_ms, s0.ended_ms), (2, 1_000, 5_000)); + assert_eq!(db.session("s599").unwrap().unwrap().turns, 1); + assert_eq!(db.session("nope").unwrap(), None); + + // 只有本地应答的不算一次会话，和列表一样 + let mut probe = row(2_000, 6_000); + probe.session = Some("probes".into()); + probe.local = true; + db.insert(&probe).unwrap(); + assert_eq!(db.session("probes").unwrap(), None); + } + #[test] fn the_last_request_id_is_where_the_next_run_has_to_start() { let db = Db::in_memory().unwrap(); @@ -2489,6 +2738,10 @@ pub(crate) mod tests { r.client_hint = hint.map(|s| s.to_string()); db.insert(&r).unwrap(); } + // 没有密钥的（本地来的、旧记录）不算一把 + let mut keyless = row(9, 900); + keyless.client = String::new(); + db.insert(&keyless).unwrap(); let mut got = db.last_seen_by_client().unwrap(); got.sort(); assert_eq!( @@ -2497,6 +2750,37 @@ pub(crate) mod tests { ); } + /// 另加的索引不算 schema：同一版本、还没有这些索引的库照样打开，行都在，缺的索引补上 + /// —— 加 SCHEMA 的话，每个人的请求历史都要清掉 + #[test] + fn a_database_without_the_extra_indexes_keeps_its_rows_and_gains_them() { + const EXTRA: [&str; 3] = ["requests_client", "requests_session_at", "requests_routed"]; + let d = tempfile::tempdir().unwrap(); + let path = d.path().join("data.db"); + { + let db = Db::open(&path).unwrap(); + db.insert(&row(1, 100)).unwrap(); + for name in EXTRA { + db.conn.execute(&format!("DROP INDEX {name}"), []).unwrap(); + } + } + let db = Db::open(&path).unwrap(); + assert_eq!(db.count().unwrap(), 1); + let names: Vec = db + .conn + .prepare( + "SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'requests'", + ) + .unwrap() + .query_map([], |r| r.get(0)) + .unwrap() + .collect::>() + .unwrap(); + for name in EXTRA { + assert!(names.iter().any(|n| n == name), "{name}: {names:?}"); + } + } + #[test] fn a_database_from_another_version_is_refused_rather_than_read() { // 不迁移，也不硬读：版本对不上就说出来，怎么处理由 `crate::open` 定 @@ -3193,6 +3477,27 @@ mod route_stats_tests { } /// 窗口外的、本地应答的、没有路由的不算；**一条解不开的记录不挡住别的** + /// 路由那一列解不开的那一行照样写得进去（索引不收它），数命中时不算。**一条坏掉的记录 + /// 不该让整张表出不来，也不该让它自己写不进去** + #[test] + fn a_row_whose_routing_does_not_parse_is_kept_and_left_out_of_the_count() { + let db = Db::in_memory().unwrap(); + let mut broken = row(1, 100); + broken.routing = Some("{\"route\": ".into()); + db.insert(&broken).unwrap(); + // 解得开、却缺了必有的一项的也不算 + let mut partial = row(2, 150); + partial.routing = Some(r#"{"route":"默认"}"#.into()); + db.insert(&partial).unwrap(); + db.insert(&routed(3, 200, "默认", "兜底", &[], None)) + .unwrap(); + assert_eq!(db.count().unwrap(), 3); + let got = db.route_stats(0, 1_000).unwrap(); + assert_eq!(got.routes.len(), 1, "{got:?}"); + assert_eq!(got.routes[0].requests, 1); + assert_eq!(hits(&got.routes[0], "兜底").decided, 1); + } + #[test] fn only_routed_requests_inside_the_window_count() { let db = Db::in_memory().unwrap(); @@ -3286,3 +3591,256 @@ mod route_stats_tests { assert_eq!(db.route_stats(700, 700).unwrap().covered_since_ms, None); } } + +/// 一份和真实库一样大的请求库上，界面常问的那几条查询各花多久。 +/// +/// 默认不跑：要先写二十七万行。发布构建下跑（`TW_BENCH_DB` 给一个路径的话库留在那里， +/// 下次接着用）： +/// `cargo test --release -p tw-store history_cost -- --ignored --nocapture` +#[cfg(test)] +mod cost { + use super::tests::row; + use super::*; + + /// 九十天、每天三千条，和留得最久的记录一样长 + const ROWS: i64 = 270_000; + const NOW: i64 = 1_790_000_000_000; + const DAY: i64 = 86_400_000; + + /// 不引随机数的库：一个线性同余就够造数据 + struct Lcg(u64); + impl Lcg { + fn next(&mut self) -> u64 { + self.0 = self + .0 + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + self.0 >> 33 + } + fn below(&mut self, n: u64) -> u64 { + self.next() % n + } + } + + const CLIENTS: &[&str] = &["claude-code", "codex", "cursor", "opencode", "gemini", "pi"]; + const ROUTES: &[&str] = &["默认", "工作", "便宜"]; + const RULES: &[&str] = &[ + "兜底", + "opus 走中转", + "haiku 本地", + "长上下文", + "关思考", + "拒绝 gpt-4", + ]; + const PROVIDERS: &[&str] = &["官方", "中转 A", "中转 B", "OpenRouter", "Bedrock"]; + const MODELS: &[&str] = &[ + "claude-sonnet-4-5", + "claude-opus-4-1", + "gpt-5-codex", + "gemini-2.5-pro", + ]; + + fn routing(g: &mut Lcg) -> String { + let hops = if g.below(10) < 4 { 2 + g.below(2) } else { 1 }; + let attempts = (0..hops) + .map(|i| { + let last = i + 1 == hops; + tw_api::AttemptView { + provider: PROVIDERS[g.below(PROVIDERS.len() as u64) as usize].into(), + model: (g.below(4) == 0).then(|| "claude-sonnet-4-5-20250929".into()), + outcome: if last { + tw_api::AttemptOutcome::Served + } else { + tw_api::AttemptOutcome::Error + }, + status: last.then_some(200), + error: (!last).then(|| Msg { + code: "gw.upstream.timeout".into(), + args: Default::default(), + text: "The upstream answered 529 Overloaded: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"},\"request_id\":\"req_011CTxxxxxxxxxxxxxxxxxxx\"}; trying the next one." + .into(), + }), + ms: 200 + g.below(30_000), + usage: None, + queued_ms: None, + skipped: None, + } + }) + .collect(); + serde_json::to_string(&tw_api::RoutingView { + route: ROUTES[g.below(ROUTES.len() as u64) as usize].into(), + rule: RULES[g.below(RULES.len() as u64) as usize].into(), + group: Some("__all__".into()), + rewritten_by: (0..g.below(3)) + .map(|_| RULES[g.below(RULES.len() as u64) as usize].to_string()) + .collect(), + denied_by: (g.below(50) == 0).then(|| RULES[5].into()), + affinity: Some(tw_api::AffinityView { + held_route: g.below(2) == 0, + stayed: None, + }), + attempts, + }) + .unwrap() + } + + /// 造一份：会话两三个交错着进行，一次几轮到几百轮；少数请求认不出会话（WebSocket） + /// 或者是网关自己答的 + fn fill(db: &Db, rows: i64) { + let mut g = Lcg(7); + // （会话号，密钥，还剩几轮） + let mut open: Vec<(String, &str, u64)> = Vec::new(); + db.conn.execute_batch("BEGIN").unwrap(); + for id in 1..=rows { + let at = NOW - 90 * DAY + id * (90 * DAY / rows); + let mut r = row(id, at); + r.price_source = Some(r#"{"kind":"default","date":"2026-09-30"}"#.into()); + r.model = MODELS[g.below(MODELS.len() as u64) as usize].into(); + r.sent_model = r.model.clone(); + r.provider = PROVIDERS[g.below(PROVIDERS.len() as u64) as usize].into(); + match g.below(100) { + 0..=2 => { + r.local = true; + r.path = "count_tokens".into(); + r.client = CLIENTS[g.below(CLIENTS.len() as u64) as usize].into(); + } + 3..=9 => { + r.client = CLIENTS[g.below(CLIENTS.len() as u64) as usize].into(); + r.routing = Some(routing(&mut g)); + } + _ => { + while open.len() < 3 { + let len = if g.below(20) == 0 { + 300 + g.below(400) + } else { + 1 + g.below(120) + }; + open.push(( + format!("{:012x}-{at}", g.next()), + CLIENTS[g.below(CLIENTS.len() as u64) as usize], + len, + )); + } + let i = g.below(open.len() as u64) as usize; + r.session = Some(open[i].0.clone()); + r.client = open[i].1.into(); + r.routing = Some(routing(&mut g)); + open[i].2 -= 1; + if open[i].2 == 0 { + open.swap_remove(i); + } + } + } + if g.below(30) == 0 { + r.error = Some(super::tests::upstream_failed("upstream returned 529")); + } + db.insert(&r).unwrap(); + } + db.conn.execute_batch("COMMIT").unwrap(); + } + + /// 以前的会话列表：整张表按会话分组，排序之后截前几条。新的那条要和它一模一样 + fn grouping_everything(db: &Db, within: Option<(i64, i64)>, limit: usize) -> Vec { + let (from, to) = within.unwrap_or((i64::MIN, i64::MAX)); + db.conn + .prepare(&format!( + "SELECT {} FROM requests + WHERE session IS NOT NULL AND local = 0 + GROUP BY session + HAVING MAX(at_ms) >= ?1 AND MIN(at_ms) <= ?2 + ORDER BY MAX(at_ms) DESC, session DESC + LIMIT ?3", + session_columns() + )) + .unwrap() + .query_map([from, to, limit as i64], session_row) + .unwrap() + .collect::>() + .unwrap() + } + + /// 挑会话再聚合，和整张表分组再截前几条，在各种窗口、各种条数下给的是同一张表 + #[test] + fn the_list_is_what_grouping_every_row_gives() { + let db = Db::in_memory().unwrap(); + fill(&db, 6_000); + for limit in [1, 7, 50, 2_000] { + for within in [ + None, + Some((NOW - 7 * DAY, NOW)), + Some((NOW - 31 * DAY, NOW - 30 * DAY)), + Some((NOW - 60 * DAY, NOW - 59 * DAY + 1)), + Some((NOW + DAY, NOW + 2 * DAY)), + Some((i64::MIN, NOW - 89 * DAY)), + ] { + assert_eq!( + db.sessions(within, limit).unwrap(), + grouping_everything(&db, within, limit), + "{within:?} {limit}" + ); + } + } + } + + fn time(what: &str, mut f: impl FnMut() -> T) -> T { + let mut out = f(); + let mut best = std::time::Duration::MAX; + for _ in 0..5 { + let t0 = std::time::Instant::now(); + out = f(); + best = best.min(t0.elapsed()); + } + eprintln!("{what:<48} {best:>12.2?}"); + out + } + + #[test] + #[ignore] + fn history_cost() { + let dir = tempfile::tempdir().unwrap(); + let path = std::env::var_os("TW_BENCH_DB") + .map(std::path::PathBuf::from) + .unwrap_or_else(|| dir.path().join("data.db")); + let t0 = std::time::Instant::now(); + let db = Db::open(&path).unwrap(); + eprintln!( + "{:<48} {:>12.2?}", + "open (builds missing indexes)", + t0.elapsed() + ); + if db.count().unwrap() < ROWS { + fill(&db, ROWS); + } + let size = std::fs::metadata(&path).map(|m| m.len()).unwrap_or(0); + eprintln!( + "{} rows, {} MB on disk", + db.count().unwrap(), + size / 1_000_000 + ); + for (what, within) in [ + ("sessions (no window, 200)", None), + ("sessions (last 7 days, 200)", Some((NOW - 7 * DAY, NOW))), + ( + "sessions (a day a month ago, 200)", + Some((NOW - 31 * DAY, NOW - 30 * DAY)), + ), + ] { + let got = time(what, || db.sessions(within, 200).unwrap()); + let grouped = time(" the same by grouping every row", || { + grouping_everything(&db, within, 200) + }); + assert_eq!(got, grouped, "{what}"); + } + // 会话详情要的那一次：最老的一次会话 + let oldest = db.sessions(None, 2000).unwrap().pop().unwrap(); + let one = time("one session (oldest of 2000)", || { + db.session(&oldest.id).unwrap() + }); + assert_eq!(one, Some(oldest.clone())); + time("one session's turns", || db.turns(&oldest.id).unwrap()); + time("last_seen_by_client", || db.last_seen_by_client().unwrap()); + time("route_stats (last 7 days)", || { + db.route_stats(NOW - 7 * DAY, NOW).unwrap() + }); + } +} From 7ae1fe5be5aabaa39173ed9dbf176b48373d2414 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:19:47 +0800 Subject: [PATCH 04/16] fix(control): session detail finds any session by id; store reads off the async workers GET /sessions/{id} took the newest 500 sessions and searched them, so an older session listed further down answered 404. It now asks for that session directly (Db::session). The session list, session detail, route stats, storage status and the overview's body total now run on the blocking pool (`on_store`): the queries and the walk over every stored body file used to run on a runtime worker, and the body walk did so while holding the recorder's lock. The body total no longer takes the lock at all; the body directory is only a path. Co-Authored-By: Claude Opus 5.5 --- crates/tw-control/src/lib.rs | 89 ++++++++++++++++++----------- crates/tw-control/tests/sessions.rs | 31 ++++++++++ 2 files changed, 86 insertions(+), 34 deletions(-) diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index 9ee2f0fe..c552f737 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -305,10 +305,15 @@ async fn overview(State(s): State) -> Json { let cfg = s.config(); let cfg = &*cfg; let engine = cfg.engine(); - // 正文现在占了多少。**扫的是目录，所以放在这儿算一次**，不塞进 - // 下面那个视图表达式里 + // 正文现在占了多少。**扫的是目录**：每份正文问一次大小，几万份要几十毫秒，所以在阻塞 + // 线程上扫，也不拿着记录库的锁扫 —— 正文目录只是一个路径 let body_bytes_now = match &s.store { - Some(st) => st.lock().await.blobs().total_bytes(), + Some(st) => { + let root = st.lock().await.blobs().root().to_path_buf(); + tokio::task::spawn_blocking(move || tw_store::Blobs::new(root).total_bytes()) + .await + .unwrap_or(0) + } None => 0, }; Json(tw_api::Overview { @@ -724,8 +729,8 @@ async fn route_stats( ) -> Result, Fail> { let (from, to) = range(q.from_ms, q.to_ms); let store = need_store(&s)?; - let g = store.lock().await; - Ok(Json(g.db().route_stats(from, to).map_err(records)?)) + let stats = on_store(store, move |g| g.db().route_stats(from, to)).await?; + Ok(Json(stats.map_err(records)?)) } /// 一段时间的汇总。不给参数就是「今天」。 @@ -1226,11 +1231,17 @@ async fn storage(State(s): State) -> Json { forwarding_affected: false, }); }; - let g = store.lock().await; + let root = store.lock().await.blobs().root().to_path_buf(); + let rows = on_store(store, |g| g.db().count().unwrap_or(0)) + .await + .unwrap_or(0); + let blob_bytes = tokio::task::spawn_blocking(move || tw_store::Blobs::new(root).total_bytes()) + .await + .unwrap_or(0); Json(tw_api::StorageStatus { recording: true, - rows: g.db().count().unwrap_or(0), - blob_bytes: g.blobs().total_bytes(), + rows, + blob_bytes, // **永远是 false。**观测挂了，代理照跑。哪天有人想改成 // true，先回去读那一节。 forwarding_affected: false, @@ -1252,6 +1263,21 @@ pub(crate) fn need_store( }) } +/// 在阻塞线程上拿着记录库做 `f`。 +/// +/// **查库、扫正文目录都不占异步线程**：一条聚合几毫秒到几十毫秒，同一个线程上的别的请求 +/// （控制面的、转发的）都得等它。锁也在那个线程上等：记录器正拿着它写一批行时，等的是那个 +/// 阻塞线程，不是异步线程。 +pub(crate) async fn on_store( + store: &Arc>, + f: impl FnOnce(&tw_store::Recorder) -> T + Send + 'static, +) -> Result { + let store = store.clone(); + tokio::task::spawn_blocking(move || f(&store.blocking_lock())) + .await + .map_err(internal) +} + fn history_row( r: tw_store::RequestRow, security: Vec, @@ -1652,15 +1678,13 @@ async fn sessions( let Some(store) = &s.store else { return Json(Vec::new()); }; - let g = store.lock().await; - Json( - g.db() - .sessions(within(&q), list_limit(&q)) - .unwrap_or_default() - .iter() - .map(session_view) - .collect(), - ) + let (within, limit) = (within(&q), list_limit(&q)); + let rows = on_store(store, move |g| g.db().sessions(within, limit)) + .await + .ok() + .and_then(Result::ok) + .unwrap_or_default(); + Json(rows.iter().map(session_view).collect()) } async fn session_detail( @@ -1673,23 +1697,20 @@ async fn session_detail( msg!("control.store_off" => "Request recording is not running."), )); }; - let g = store.lock().await; - let session = g - .db() - .sessions(None, 500) - .unwrap_or_default() - .iter() - .find(|x| x.id == id) - .map(session_view) - .ok_or_else(|| no_such_session(&id))?; - let turns = g - .db() - .turns(&id) - .unwrap_or_default() - .iter() - .map(turn_view) - .collect(); - Ok(Json(tw_api::SessionDetail { session, turns })) + // **按会话号直接取**：以前取最近五百次会话再从里面找，更早的会话点开是 404 + let (session, turns) = { + let id = id.clone(); + on_store(store, move |g| { + Ok::<_, tw_store::DbError>((g.db().session(&id)?, g.db().turns(&id)?)) + }) + .await? + .map_err(records)? + }; + let session = session.ok_or_else(|| no_such_session(&id))?; + Ok(Json(tw_api::SessionDetail { + session: session_view(&session), + turns: turns.iter().map(turn_view).collect(), + })) } fn no_such_session(id: &str) -> Fail { diff --git a/crates/tw-control/tests/sessions.rs b/crates/tw-control/tests/sessions.rs index 5c47be1c..084dc6c9 100644 --- a/crates/tw-control/tests/sessions.rs +++ b/crates/tw-control/tests/sessions.rs @@ -194,3 +194,34 @@ async fn the_endpoints_take_a_time_window_off_the_query_string() { let none = get(&app, "/sessions?from_ms=9000&to_ms=9999").await; assert!(none.as_array().unwrap().is_empty(), "{none}"); } + +/// **会话详情按会话号取，多老都取得到。**以前它取最近五百次会话再从里面找：更早的会话 +/// 在列表里点得到（翻到那么远的话），点开却是 404 +#[tokio::test] +async fn a_session_behind_the_newest_five_hundred_still_opens() { + let mut rows = vec![turn(1, "per-token"), turn(2, "per-token")]; + for i in 0..600 { + let mut r = turn(100 + i, "per-token"); + r.session = Some(format!("later-{i}")); + rows.push(r); + } + let (_d, app) = app(&rows); + + let detail = get(&app, "/sessions/s1").await; + assert_eq!(detail["session"]["id"], "s1"); + assert_eq!(detail["session"]["turns"], 2, "{detail}"); + assert_eq!(detail["turns"].as_array().unwrap().len(), 2); + + // 没有的照样是 404 + let r = app + .clone() + .oneshot( + Request::builder() + .uri("/sessions/nope") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(r.status(), StatusCode::NOT_FOUND); +} From 7a82e5410aa22f4222d5d890097d33806668a410 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:31:31 +0800 Subject: [PATCH 05/16] perf(store): recorder on its own threads and channel; GC deletes bodies without the lock The recorder ran as tokio tasks doing synchronous work on runtime workers under one tokio mutex: SQLite inserts per event, body writes of up to 4 MiB plus chmods, plugin-run inserts. The hourly GC held that same lock while deleting day directories, which can take seconds. - Events, bodies and plugin runs are each handled on a named OS thread (tw-recorder, tw-bodies, tw-plugin-runs). The event thread takes the lock once per batch of up to 64 events. Body writes no longer take the lock at all; the body directory is only a path. - The recorder no longer subscribes to the UI's 1024-slot broadcast, where lagging meant those requests were missing from history for good. EventBus::record_feed gives it its own bounded channel (16384 slots, room for ~2000 requests of backlog). emit() try_sends into it, so forwarding still never waits; when it is full the event is dropped and counted, and the recorder logs how many. - task::gc deletes expired and over-cap bodies without the lock, then takes it only to prune rows; twcore runs it on the blocking pool. Recorder::gc became Recorder::prune (rows only). - key_limits::follow keeps the runtime handle it was attached on: settling now happens on the recorder thread, where Handle::try_current() fails, so a period change found while settling would never have been read back from the records. Co-Authored-By: Claude Opus 5.5 --- bin/twcore/src/main.rs | 34 ++++--- crates/tw-control/src/key_limits.rs | 8 +- crates/tw-control/tests/error_answers.rs | 2 +- crates/tw-control/tests/key_limits.rs | 46 ++++++++++ crates/tw-control/tests/route_stats.rs | 6 +- crates/tw-control/tests/stored_bodies.rs | 2 +- crates/tw-control/tests/ws_turns.rs | 4 +- crates/tw-observe/src/bus.rs | 107 ++++++++++++++++++++++- crates/tw-observe/src/lib.rs | 2 +- crates/tw-store/src/recorder.rs | 33 ++++--- crates/tw-store/src/task.rs | 98 +++++++++++++++------ 11 files changed, 282 insertions(+), 60 deletions(-) diff --git a/bin/twcore/src/main.rs b/bin/twcore/src/main.rs index 9cf231dc..0db6e860 100644 --- a/bin/twcore/src/main.rs +++ b/bin/twcore/src/main.rs @@ -840,6 +840,9 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> 每一轮都重新读配置，所以改完期限不用重启 —— 下一轮就按新的算。 扫的是目录和一条 DELETE，空跑一次的代价可以忽略。 + + **在阻塞线程上跑，删正文时不拿记录库的锁**（见 `tw_store::task::gc`）：一次删几个 + GB 可以要好几秒，以前那期间这把锁一直攥着，记录和查询都在等。 */ if let Some(rec) = store.clone() { let gw = state.clone(); @@ -847,16 +850,18 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> loop { let r = gw.config().retention.clone(); let freed = { - let g = rec.lock().await; - g.gc( - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as i64) - .unwrap_or(0), - r.body_days, - r.row_days, - r.body_max_bytes, - ) + let rec = rec.clone(); + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0); + let (body_days, row_days, max) = + (r.body_days, r.row_days, r.body_max_bytes); + tokio::task::spawn_blocking(move || { + tw_store::task::gc(&rec, now, body_days, row_days, max) + }) + .await + .unwrap_or(0) }; if freed > 0 { tracing::info!( @@ -1013,7 +1018,14 @@ fn build_store( // 每记下一行请求，交给网关的密钥用量上限结算（见 `tw_control::key_limits`） settled: tw_store::SettleHook, ) -> Option>> { - let events = bus.subscribe(); + // **存储层自己的那条通道，不订阅界面的广播**（见 `EventBus::record_feed`）：广播上 + // 掉队丢的是界面上的几个点，记录丢的是那几个请求 + let Some(events) = bus.record_feed() else { + tracing::warn!( + "request recording is already attached to the event bus; not starting a second one" + ); + return None; + }; let (db, blobs) = match tw_store::open(dir) { Ok(opened) => opened, Err(e) => { diff --git a/crates/tw-control/src/key_limits.rs b/crates/tw-control/src/key_limits.rs index 891d447e..08df8fce 100644 --- a/crates/tw-control/src/key_limits.rs +++ b/crates/tw-control/src/key_limits.rs @@ -57,10 +57,16 @@ pub fn rebuild(gw: &tw_gateway::AppState, db: &tw_store::Db) { pub fn follow(gw: &tw_gateway::AppState, store: Arc>) { // 弱引用：网关的账拿着这个办法，办法再拿着账就是一个圈，谁都放不掉 let limits = Arc::downgrade(&gw.key_limits); + // **接上时的运行时，记在这里**：结算在存储层自己的记录线程上，那条线程不在运行时里，在 + // 那儿问「当前的运行时」问不到，新的一期就永远不会从记录里加回来 + let here = tokio::runtime::Handle::try_current().ok(); gw.key_limits .reread_with(Arc::new(move |m: tw_gateway::key_limits::Moved| { let (limits, store) = (limits.clone(), store.clone()); - let Ok(rt) = tokio::runtime::Handle::try_current() else { + let Some(rt) = here + .clone() + .or_else(|| tokio::runtime::Handle::try_current().ok()) + else { return; }; rt.spawn(async move { diff --git a/crates/tw-control/tests/error_answers.rs b/crates/tw-control/tests/error_answers.rs index 35697a55..73248ccb 100644 --- a/crates/tw-control/tests/error_answers.rs +++ b/crates/tw-control/tests/error_answers.rs @@ -131,7 +131,7 @@ async fn world() -> World { tw_store::Blobs::new(blobs.clone()), tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), ), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let state = ControlState { diff --git a/crates/tw-control/tests/key_limits.rs b/crates/tw-control/tests/key_limits.rs index 4d2032b4..f518e93c 100644 --- a/crates/tw-control/tests/key_limits.rs +++ b/crates/tw-control/tests/key_limits.rs @@ -587,3 +587,49 @@ async fn a_time_zone_change_adds_the_new_day_up_again_from_the_records() { } assert_eq!(used(&b), 4); } + +/// 同上，只是**发现这一期变了的是一次结算**，而结算在存储层自己的记录线程上 —— 那条线程 +/// 不在异步运行时里。以前读回新一期的任务要在「当前的运行时」上起，在那儿起不来，这一天 +/// 就从 0 起了 +#[tokio::test] +async fn a_period_change_found_while_settling_on_the_recorder_thread_is_read_back() { + let clock = Arc::new(Moving { + east: tw_gateway::key_limits::TestClock::new(ms(NOON), 8 * 3600), + west: tw_gateway::key_limits::TestClock::new(ms(NOON), -5 * 3600), + moved: Default::default(), + }); + let b = bed_with(None, clock.clone()); + tw_control::key_limits::follow(&b.gw, b.rec.clone()); + for e in [ + request(1, "plain", ms("2026-10-04T23:00:00+08:00"), 10), + request(2, "plain", ms("2026-10-05T09:00:00+08:00"), 10), + ] + .iter() + .flatten() + { + b.rec.lock().await.on_event(e); + } + let limits: Vec = + serde_yaml_ng::from_str("[{per: day, requests: 100}]").unwrap(); + let used = |b: &Bed| b.gw.key_limits.view("plain", &limits)[0].used; + assert_eq!(used(&b), 1, "东八区的今天"); + clock.moved.store(true, std::sync::atomic::Ordering::SeqCst); + // 换了时区之后第一件事是结算一个请求，在一条不属于运行时的线程上 + let rec = b.rec.clone(); + std::thread::spawn(move || { + for e in request(3, "plain", ms(NOON), 10) { + rec.blocking_lock().on_event(&e); + } + }) + .join() + .unwrap(); + let mut seen = 0; + for _ in 0..100 { + seen = used(&b); + if seen == 3 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + assert_eq!(seen, 3, "西五区的今天：三个请求都在里面"); +} diff --git a/crates/tw-control/tests/route_stats.rs b/crates/tw-control/tests/route_stats.rs index 1f84e97f..e373bb0f 100644 --- a/crates/tw-control/tests/route_stats.rs +++ b/crates/tw-control/tests/route_stats.rs @@ -87,7 +87,7 @@ async fn every_route_and_rule_counts_what_went_through_it_including_denials() { std::fs::write(&p, &text).unwrap(); let cfg: tw_config::Config = serde_yaml_ng::from_str(&text).unwrap(); let gw = tw_gateway::AppState::new(cfg).unwrap(); - // 和 twcore 一样：存储层订阅总线，事件落库 + // 和 twcore 一样：存储层从总线给它的那条通道收事件，落库 let (_bodies, rx) = tokio::sync::mpsc::channel(8); let store = tw_store::task::spawn( tw_store::Recorder::new( @@ -95,7 +95,7 @@ async fn every_route_and_rule_counts_what_went_through_it_including_denials() { tw_store::Blobs::new(d.path().join("blobs")), tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), ), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let state = ControlState { @@ -187,7 +187,7 @@ async fn a_store_with_no_history_yet_says_nothing_is_covered() { tw_store::Blobs::new(d.path().join("blobs")), tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), ), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let state = ControlState { diff --git a/crates/tw-control/tests/stored_bodies.rs b/crates/tw-control/tests/stored_bodies.rs index c1861802..6dd0e484 100644 --- a/crates/tw-control/tests/stored_bodies.rs +++ b/crates/tw-control/tests/stored_bodies.rs @@ -125,7 +125,7 @@ async fn world(mode: &str) -> World { tw_store::Blobs::new(blobs.clone()), tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), ), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let state = ControlState { diff --git a/crates/tw-control/tests/ws_turns.rs b/crates/tw-control/tests/ws_turns.rs index 7553b31a..7f63bfe1 100644 --- a/crates/tw-control/tests/ws_turns.rs +++ b/crates/tw-control/tests/ws_turns.rs @@ -79,7 +79,7 @@ providers: gw.pricing.clone(), ) .settling_to(tw_control::key_limits::settle_hook(&gw)), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let addr = tw_gateway::serve(gw.clone(), ([127, 0, 0, 1], 0).into()) @@ -243,7 +243,7 @@ providers: gw.pricing.clone(), ) .settling_to(tw_control::key_limits::settle_hook(&gw)), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let addr = tw_gateway::serve(gw.clone(), ([127, 0, 0, 1], 0).into()) diff --git a/crates/tw-observe/src/bus.rs b/crates/tw-observe/src/bus.rs index d0dd2f61..d51119b4 100644 --- a/crates/tw-observe/src/bus.rs +++ b/crates/tw-observe/src/bus.rs @@ -7,16 +7,23 @@ //! //! 注意这条豁免**不适用于成本记账**：那类数据丢了账单永久对不 //! 上，必须走别的路径。这里只走「丢了只是图上少个点」的东西。 +//! +//! **存储层不订阅广播**，它有自己的一条有界通道（[`EventBus::record_feed`]）：广播上 +//! 掉队丢掉的是界面上的几个点，记录丢掉的是那几个请求，永远不在库里了。 use std::collections::{BTreeMap, HashMap, VecDeque}; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, OnceLock}; use std::time::{Duration, Instant}; -use tokio::sync::broadcast; +use tokio::sync::{broadcast, mpsc}; const CAPACITY: usize = 1024; +/// 存储层那条通道有多少格（见 [`EventBus::record_feed`]）。一个请求五六条事件，这么多格 +/// 够存储层落后两千多个请求 +pub const RECORD_CAPACITY: usize = 16 * 1024; + /// 生成速率看最近这么久里跑完的请求 pub const RATE_WINDOW: Duration = Duration::from_secs(60); @@ -26,6 +33,39 @@ pub struct EventBus { next_id: Arc, /// 在跑的和最近跑完的。见 [`EventBus::in_flight`]、[`EventBus::live`] tally: Arc>, + /// 存储层自己的那条通道，接上了才有。见 [`EventBus::record_feed`] + record: Arc>, +} + +/// 往存储层那条通道里送的一头。 +struct RecordSink { + tx: mpsc::Sender, + /// 通道满着、没送进去的条数，存储层拿去记一行日志 + dropped: Arc, +} + +/// 存储层收事件的那一头（[`EventBus::record_feed`]）。 +pub struct RecordFeed { + rx: mpsc::Receiver, + dropped: Arc, +} + +impl RecordFeed { + /// 等下一条。总线没了是 None。**会挡住当前线程**：给存储层自己的记录线程用，不在 + /// 异步线程上调 + pub fn blocking_recv(&mut self) -> Option { + self.rx.blocking_recv() + } + + /// 已经到了的下一条，不等 + pub fn try_recv(&mut self) -> Option { + self.rx.try_recv().ok() + } + + /// 上次问到现在，通道满着丢掉了几条 + pub fn take_dropped(&self) -> u64 { + self.dropped.swap(0, Ordering::Relaxed) + } } /// 按事件数出来的现状。 @@ -124,6 +164,7 @@ impl EventBus { tx, next_id: Arc::new(AtomicU64::new(1)), tally: Arc::default(), + record: Arc::default(), } } @@ -149,11 +190,38 @@ impl EventBus { /// /// **先记账再发**（见 [`EventBus::in_flight`]）：订阅者收到一个结局的 /// 时候，快照里已经没有它了。 + /// + /// 存储层那条通道满着时**丢，不等**：转发不等观测。 pub fn emit(&self, ev: tw_api::Event) { self.track(&ev); + if let Some(sink) = self.record.get() + && let Err(mpsc::error::TrySendError::Full(_)) = sink.tx.try_send(ev.clone()) + { + sink.dropped.fetch_add(1, Ordering::Relaxed); + } let _ = self.tx.send(ev); } + /// 存储层收事件的那条通道：**有界，和界面的广播分开**，之后发的每一条都进来。 + /// + /// 广播满了丢最老的，掉队的订阅者跳过一截。对界面那是少画几个点，它还会收到 + /// `EventsDropped` 自己对账；对记录那是几个请求永远不在库里。以前存储层就订阅广播， + /// 和界面挤同一个 [`CAPACITY`] 格的环。走自己这条的话，[`RECORD_CAPACITY`] 格平常 + /// 一条都不丢；真满了照样丢（转发不等观测），存储层说一声丢了几条。 + /// + /// **只给一次**，第二次是 None：两个记录器各分到一半事件，两边都缺。 + pub fn record_feed(&self) -> Option { + let (tx, rx) = mpsc::channel(RECORD_CAPACITY); + let dropped = Arc::new(AtomicU64::new(0)); + self.record + .set(RecordSink { + tx, + dropped: dropped.clone(), + }) + .ok()?; + Some(RecordFeed { rx, dropped }) + } + /// 开始的记下，有了结局的划掉。 /// /// **这里什么都不能 panic。**`Ending` 在 Drop 里也会走到这儿，那时可能 @@ -753,6 +821,41 @@ mod tests { assert!(c > a); } + /// **存储层不和界面挤一个环。**界面那个订阅者落后到被跳过一截时，存储层那条通道里 + /// 一条不少、顺序不乱 —— 以前它也订阅广播，界面落后的时候它多半也落后，丢掉的那几个 + /// 请求永远不在库里 + #[test] + fn the_store_gets_every_event_while_a_lagging_subscriber_skips_some() { + let b = EventBus::new(); + let mut ui = b.subscribe(); + let mut feed = b.record_feed().unwrap(); + let n = CAPACITY as u64 * 4; + for id in 0..n { + b.emit(finished(id, 0, None)); + } + assert!(matches!( + ui.try_recv(), + Err(broadcast::error::TryRecvError::Lagged(_)) + )); + let got: Vec = std::iter::from_fn(|| feed.try_recv().map(|e| e.id())).collect(); + assert_eq!(got, (0..n).collect::>()); + assert_eq!(feed.take_dropped(), 0); + } + + /// 存储层那条通道满了：**丢，不挡**，丢了几条说得出来。只给一次 + #[test] + fn a_full_store_feed_drops_and_counts_rather_than_blocking() { + let b = EventBus::new(); + let mut feed = b.record_feed().unwrap(); + assert!(b.record_feed().is_none(), "给了第二个记录器"); + for id in 0..(RECORD_CAPACITY as u64 + 10) { + b.emit(finished(id, 0, None)); + } + assert_eq!(feed.take_dropped(), 10); + assert_eq!(feed.take_dropped(), 0, "问过一次就清零"); + assert_eq!(feed.try_recv().map(|e| e.id()), Some(0)); + } + #[tokio::test] async fn a_slow_subscriber_lags_but_does_not_block_the_producer() { // 这是关键性质：观测端跟不上，数据面**不能**因此变慢。 diff --git a/crates/tw-observe/src/lib.rs b/crates/tw-observe/src/lib.rs index 0c4bd591..dd312d29 100644 --- a/crates/tw-observe/src/lib.rs +++ b/crates/tw-observe/src/lib.rs @@ -6,4 +6,4 @@ pub mod bus; -pub use bus::EventBus; +pub use bus::{EventBus, RecordFeed}; diff --git a/crates/tw-store/src/recorder.rs b/crates/tw-store/src/recorder.rs index 42f97333..fa180f53 100644 --- a/crates/tw-store/src/recorder.rs +++ b/crates/tw-store/src/recorder.rs @@ -808,16 +808,15 @@ impl Recorder { } } - /// 定期回收。**返回删了多少字节**，调用方记一行日志就够了。 - pub fn gc(&self, now_ms: i64, keep_days: u64, metadata_keep_days: u64, max_bytes: u64) -> u64 { - let freed = self.blobs.gc(now_ms, keep_days, max_bytes); - let cutoff = now_ms - (metadata_keep_days as i64) * 86_400_000; + /// 删掉过期的记录行（`keep_days` 天以前的）。正文的回收不在这里：它不需要这把锁，见 + /// [`crate::task::gc`]。 + pub fn prune(&self, now_ms: i64, keep_days: u64) { + let cutoff = now_ms - (keep_days as i64) * 86_400_000; match self.db.prune_before(cutoff) { Ok(n) if n > 0 => tracing::info!(rows = n, "cleaned up expired request rows"), Err(e) => tracing::debug!("the request rows could not be cleaned up: {e}"), _ => {} } - freed } } @@ -2327,8 +2326,13 @@ mod cancellation_tests { r.db().insert(&crate::db::tests::row(1, old)).unwrap(); r.record_body(old, 1, Which::Request, b"an old body", 11); + let shared = tokio::sync::Mutex::new(r); + let gc = |keep_days, metadata_keep_days| { + crate::task::gc(&shared, now, keep_days, metadata_keep_days, u64::MAX) + }; // 正文留 7 天、记录留 90 天：正文该没了，记录还在 - r.gc(now, 7, 90, u64::MAX); + gc(7, 90); + let r = shared.blocking_lock(); assert_eq!( r.db().recent(None, 10).unwrap().len(), 1, @@ -2339,9 +2343,17 @@ mod cancellation_tests { "正文过了它自己的期限却还在" ); + drop(r); // 记录的期限也到了才删行 - r.gc(now, 7, 7, u64::MAX); - assert!(r.db().recent(None, 10).unwrap().is_empty()); + gc(7, 7); + assert!( + shared + .blocking_lock() + .db() + .recent(None, 10) + .unwrap() + .is_empty() + ); } /// 总量上限是给突发准备的：按天算出来的占用取决于用量，而用量会有 @@ -2358,9 +2370,10 @@ mod cancellation_tests { let before = r.blobs().total_bytes(); assert!(before > 0); // 一天都没过期，但总量超了 —— 从最旧的整天开始删 - let freed = r.gc(now, 30, 30, 4096); + let shared = tokio::sync::Mutex::new(r); + let freed = crate::task::gc(&shared, now, 30, 30, 4096); assert!(freed > 0, "总量超了却什么都没删"); - assert!(r.blobs().total_bytes() < before); + assert!(shared.blocking_lock().blobs().total_bytes() < before); } /// 会话里的那一轮也要看得出是取消的。否则在每轮花费里，它就是一轮 diff --git a/crates/tw-store/src/task.rs b/crates/tw-store/src/task.rs index bfd6d2fd..6ffb92e6 100644 --- a/crates/tw-store/src/task.rs +++ b/crates/tw-store/src/task.rs @@ -1,15 +1,22 @@ -//! 起一个后台任务，把事件流写进库里。 +//! 起几个后台线程，把事件流写进库里、把正文写到盘上。 //! -//! **它订阅事件总线，而不是被数据面直接调用。**这条边界是刻意的：写库 +//! **它收事件总线上的事件，而不是被数据面直接调用。**这条边界是刻意的：写库 //! 慢了、库锁住了、磁盘满了，转发这条路上一行代码都不会等它。 +//! +//! **跑在自己的线程上，不在异步运行时里。**写库、写几 MB 的正文、`chmod` 都是会挡住线程的 +//! 调用：放在运行时的工作线程上，同一个线程上排着的转发和控制面请求都得等它。 use std::sync::Arc; use tokio::sync::Mutex; -use crate::blobs::Which; +use crate::blobs::{Blobs, Which}; use crate::recorder::Recorder; +/// 记录线程拿着锁一次最多吃几条事件。**拿着锁的时候查询在等**：一条几十微秒，一批也就 +/// 几毫秒；一条一拿又白白多出许多次抢锁 +const BATCH: usize = 64; + /// 一份等着落盘的 body。 /// /// **这个类型属于 store，不属于网关。**让 store 依赖网关会把「观测」 @@ -29,14 +36,14 @@ pub struct StoredBody { } /// 把插件的运行记录写进库里，**和正文同一个道理**：数据面只管交出去，写库在这条 -/// 任务上。记录和请求那一行各走各的路，落库不分先后（见 `Db::insert_plugin_run`）。 +/// 线程上。记录和请求那一行各走各的路，落库不分先后（见 `Db::insert_plugin_run`）。 pub fn record_plugin_runs( recorder: Arc>, mut runs: tokio::sync::mpsc::Receiver, ) { - tokio::spawn(async move { - while let Some(r) = runs.recv().await { - recorder.lock().await.record_plugin_run(&r); + thread("tw-plugin-runs", move || { + while let Some(r) = runs.blocking_recv() { + recorder.blocking_lock().record_plugin_run(&r); } }); } @@ -44,38 +51,73 @@ pub fn record_plugin_runs( /// 起来。返回的 handle 给别的地方查历史用 —— **同一个 Recorder**， /// 不是第二个连接：两个连接会让「刚写进去的还查不到」变成可能。 /// -/// **回收不在这里。**它按配置里的 `retention` 走（`twcore` 每小时读一次配置，调 -/// [`Recorder::gc`]）。这里以前另跑着一个按写死的 7 天、90 天、2 GB 回收的循环，和按 -/// 配置的那个同时跑：配得比出厂更宽的期限和上限，每小时被它削回出厂值一次。 +/// 事件从总线给存储层的那条通道来（[`tw_observe::EventBus::record_feed`]），**不订阅界面的 +/// 广播**：那是个一千来格的环，界面落后时存储层多半也落后，丢掉的那几个请求永远不在库里。 +/// +/// **正文不拿锁**：正文目录只是一个路径，写一份 4 MB 的正文时，记录和查询不必等它。 +/// +/// **回收不在这里。**它按配置里的 `retention` 走（`twcore` 每小时读一次配置，调 [`gc`]）。 +/// 这里以前另跑着一个按写死的 7 天、90 天、2 GB 回收的循环，和按配置的那个同时跑：配得比出厂 +/// 更宽的期限和上限，每小时被它削回出厂值一次。 pub fn spawn( recorder: Recorder, - mut rx: tokio::sync::broadcast::Receiver, + mut events: tw_observe::RecordFeed, mut bodies: tokio::sync::mpsc::Receiver, ) -> Arc> { + let blobs = Blobs::new(recorder.blobs().root().to_path_buf()); let shared = Arc::new(Mutex::new(recorder)); - let r = shared.clone(); - tokio::spawn(async move { - while let Some(b) = bodies.recv().await { - // **写盘在这条任务上，不在转发那条路上。**慢磁盘只会让 - // 通道积压然后丢，不会让请求变慢。 - r.lock() - .await - .record_body(b.at_ms, b.id, b.which, &b.body, b.original_len); + thread("tw-bodies", move || { + // 慢磁盘只会让通道积压然后丢，不会让请求变慢 + while let Some(b) = bodies.blocking_recv() { + blobs.put_with_len(b.at_ms, b.id as i64, b.which, &b.body, b.original_len); } }); let r = shared.clone(); - tokio::spawn(async move { - loop { - match rx.recv().await { - Ok(ev) => r.lock().await.on_event(&ev), - // **落后了就跳过，不要退出。**广播通道满的时候丢的是 - // 观测记录，而退出会让之后的记录全部丢失。 - Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { - tracing::warn!("recording fell behind; {n} requests were not recorded"); + thread("tw-recorder", move || { + while let Some(ev) = events.blocking_recv() { + { + let mut g = r.blocking_lock(); + g.on_event(&ev); + for _ in 1..BATCH { + let Some(ev) = events.try_recv() else { break }; + g.on_event(&ev); } - Err(tokio::sync::broadcast::error::RecvError::Closed) => return, + } + // **落后了就丢，但要说出来。**通道满了丢的是观测记录，而等着它会拖慢转发 + let dropped = events.take_dropped(); + if dropped > 0 { + tracing::warn!( + dropped, + "recording fell behind; {dropped} events were not recorded" + ); } } }); shared } + +/// 回收一轮：过期的正文、超出总量的正文、过期的记录行。返回删掉了多少字节的正文。 +/// +/// **删正文的时候不拿锁**：正文目录只是一个路径，一次删几个 GB 可以要好几秒，而那期间记录 +/// 和查询都要这把锁。删行要拿：行在库里，一小时一轮删的只是这一小时过期的那些。 +/// +/// 会挡住当前线程：在阻塞线程上调。 +pub fn gc( + shared: &Mutex, + now_ms: i64, + keep_days: u64, + metadata_keep_days: u64, + max_bytes: u64, +) -> u64 { + let blobs = Blobs::new(shared.blocking_lock().blobs().root().to_path_buf()); + let freed = blobs.gc(now_ms, keep_days, max_bytes); + shared.blocking_lock().prune(now_ms, metadata_keep_days); + freed +} + +/// 起一条有名字的线程。**起不来只记一行**：那时这一路什么都不记，转发照常 +fn thread(name: &str, f: impl FnOnce() + Send + 'static) { + if let Err(e) = std::thread::Builder::new().name(name.into()).spawn(f) { + tracing::warn!("the {name} thread could not start, so it records nothing this run: {e}"); + } +} From 440f3af17665c340df320d034900e8af3c4fbb4f Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:32:30 +0800 Subject: [PATCH 06/16] perf(watch): watch only a directory's own entries on macOS, with kqueue notify's default macOS backend is FSEvents, which always watches the whole tree below a path. `RecursiveMode::NonRecursive` only drops the deeper events inside our process, after the kernel has already woken us, with zero latency and one event per file. The desktop app watches `~` itself (that is where `~/.claude.json` lives): an audit counted 7-14 raw events a second under `~` while idle, none of them relevant, and writes under `~/Library/Caches` cost the app measurable CPU. tw-watch also backs core's config and plugin hot-reload. On macOS tw-watch now does it with kqueue: - each watched directory gets an O_EVTONLY vnode watch, so only its own entries being added, removed or renamed wake us; on a wake it lists the directory and fires when a relevant entry appeared, disappeared or was replaced (inode changed); - each relevant regular file (or symlink to one) also gets a watch, for in-place writes, appends and truncation, which do not touch the directory; after an atomic replace the directory listing re-arms it on the new inode; - attribute-only changes count only when size or mtime moved (truncate reports nothing else; chmod and reads do not count); - a watched directory that is removed or renamed away counts, and its parent is watched until the name reappears, so a replaced or recreated directory keeps being watched, as with FSEvents. notify's own kqueue backend (`macos_kqueue`) is not usable here: on the first change in a non-recursively watched directory it treats the first entry it has not seen as new and watches it recursively, which for `~` means opening all of `~/Library`. Each watch costs a file descriptor and apps started from Finder get a soft limit of 256, so all watches in a process together hold at most half the soft limit (capped at 1024). A watch that would exceed it falls back to FSEvents as a whole, at start or later, rather than silently missing in-place writes. Measured with a probe binary (release build, 200 ms debounce): - 1,500 writes 3 levels below a watched directory over ~9 s: before 51-54 ms CPU and ~890 context switches, after 0.4-0.5 ms and 9-13; both reported the one relevant change afterwards. - watching `~` for 30 s idle, both probes side by side: before 8.8-11.1 ms CPU and ~233 context switches, after 0.1-0.2 ms and 3. Linux (inotify) and Windows (ReadDirectoryChangesW) keep using notify unchanged. The public API is unchanged. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 1 + crates/tw-watch/Cargo.toml | 4 + crates/tw-watch/src/kqueue.rs | 912 +++++++++++++++++++++++++++++ crates/tw-watch/src/lib.rs | 221 ++++++- crates/tw-watch/tests/fd_budget.rs | 122 ++++ 5 files changed, 1244 insertions(+), 16 deletions(-) create mode 100644 crates/tw-watch/src/kqueue.rs create mode 100644 crates/tw-watch/tests/fd_budget.rs diff --git a/Cargo.lock b/Cargo.lock index 4827b683..a8893b02 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3424,6 +3424,7 @@ dependencies = [ name = "tw-watch" version = "0.64.1" dependencies = [ + "libc", "notify", "tempfile", "thiserror", diff --git a/crates/tw-watch/Cargo.toml b/crates/tw-watch/Cargo.toml index 0390cfe0..c5ee0818 100644 --- a/crates/tw-watch/Cargo.toml +++ b/crates/tw-watch/Cargo.toml @@ -15,6 +15,10 @@ notify = { workspace = true } tokio = { workspace = true } thiserror = { workspace = true } +# macOS 上不递归的监听自己用 kqueue 写（见 src/kqueue.rs） +[target.'cfg(target_os = "macos")'.dependencies] +libc = { workspace = true } + [dev-dependencies] tempfile = "3" tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } diff --git a/crates/tw-watch/src/kqueue.rs b/crates/tw-watch/src/kqueue.rs new file mode 100644 index 00000000..fad86033 --- /dev/null +++ b/crates/tw-watch/src/kqueue.rs @@ -0,0 +1,912 @@ +//! macOS 上的不递归监听：kqueue。 +//! +//! **为什么不用 notify 在 macOS 上默认的 FSEvents。**FSEvents 只会盯一整棵树： +//! `RecursiveMode::NonRecursive` 是 notify 在我们进程里把深处的事件丢掉，而内核照样把 +//! 底下每一次写都送过来（零延迟、按文件报），每一次都叫醒我们。桌面端要盯 `~` 本身 +//! （`~/.claude.json` 在那儿），实测空闲时 `~` 底下每秒七到十四个事件，没有一个相关 —— +//! 「空闲时接近零」就这么没了。 +//! +//! **也不用 notify 的 kqueue 后端**（`macos_kqueue`）：它不递归地盯一个目录时，目录里 +//! 一有变化，就把目录里第一个它没盯过的条目当成「新建的」，并且**递归地**盯上它 —— +//! 盯着 `~` 的话，第一次有人在 `~` 里建个文件，它就去把 `~/Library` 整棵树挨个打开。 +//! +//! 这里的做法： +//! +//! - **目录**：打开（`O_EVTONLY`：只为了收事件，不妨碍卸载）挂一个 vnode 事件。只有目录 +//! 自己的条目增、删、改名（包括 rename 进来的原子保存）才叫醒我们，子目录里再怎么变 +//! 都不会。醒了之后列一遍目录，只看算数的条目：多了、少了、换了一个（inode 变了）才算 +//! 一次。 +//! - **算数的普通文件**：也各挂一个。原地写（不经过 rename 的保存、追加、截断）不动目录， +//! 只有盯着文件本身才看得见。文件被换掉之后，手上那个 fd 指的还是旧的那个 inode —— +//! 它报一个删除或改名，目录那边重新列一遍，就地挂到新的上。 +//! - **只看内容和存在性**：属性变了（`NOTE_ATTRIB`）只有大小或修改时间跟着变才算。截断 +//! 只报这一个；改权限不算，读文件本来就不报。 +//! +//! **fd 有名额。**每盯一个目录或文件占一个 fd，而从访达打开的应用默认只有 256 个。整个 +//! 进程里所有监听合起来最多占软上限的一半（[`budget`]）；超了，这个监听整个换回 FSEvents —— +//! 费电，但一个事件都不漏。只盯一部分文件、悄悄漏掉它们的原地写，比费电更糟。 +//! +//! **盯的是路径，不是那个 inode**，和 FSEvents 一样：盯着的目录自己被删掉、挪走时算一次 +//! （`relevant` 问的是目录自己的路径），然后改盯它的上一层，等这个名字再出现 —— 整个目录 +//! 被换掉（rename 一个新的过来）、删了又建，都接着盯新的那个，出现时也算一次。上一层也 +//! 没了，才不再盯。 + +use std::collections::{BTreeSet, HashMap}; +use std::ffi::{CString, OsString}; +use std::io; +use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd}; +use std::os::unix::ffi::OsStrExt; +use std::os::unix::fs::{DirEntryExt, MetadataExt}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::mpsc::Sender; +use std::sync::{Arc, Mutex, PoisonError}; + +use crate::{Relevant, WatchError}; + +/// 目录上要的事件：条目变了（`NOTE_WRITE`），和它自己没了、被挪走 +const DIR_EVENTS: u32 = libc::NOTE_WRITE | GONE; + +/// 文件上要的事件：内容变了、属性变了（[`Loop::stamp_moved`] 再筛），和它自己没了、被挪走 +const FILE_EVENTS: u32 = libc::NOTE_WRITE | libc::NOTE_EXTEND | libc::NOTE_ATTRIB | GONE; + +/// 手上这个 vnode 不再是那个路径上的东西了：删了、改名挪走了、文件系统卸载了 +const GONE: u32 = libc::NOTE_DELETE | libc::NOTE_RENAME | libc::NOTE_REVOKE; + +/// 叫事件循环停下的那个用户事件（`EVFILT_USER`；和 vnode 事件不在一个名字空间里） +const STOP: usize = 0; + +/// fd 名额的硬顶：软上限再高，所有监听合起来也不占更多 +const MAX_HELD: usize = 1024; + +/// 整个进程里所有监听此刻占着的 fd +static HELD: AtomicUsize = AtomicUsize::new(0); + +/// 所有监听合起来最多占多少个 fd：软上限的一半，最多 [`MAX_HELD`]。 +/// +/// **按进程算，不按监听算**：一个进程里可能同时有好几个（桌面端换一份监听时，新的起来了 +/// 旧的才放掉）。每次都现问软上限：进程中途调高了，名额跟着多。 +fn budget() -> usize { + let mut r = libc::rlimit { + rlim_cur: 0, + rlim_max: 0, + }; + // SAFETY: getrlimit 只往传进去的结构体里写 + let soft = if unsafe { libc::getrlimit(libc::RLIMIT_NOFILE, &mut r) } == 0 { + r.rlim_cur + } else { + 256 + }; + usize::try_from(soft / 2).unwrap_or(MAX_HELD).min(MAX_HELD) +} + +/// 一个占着名额的、只为了收事件打开的 fd。放掉它，fd 关上（挂在上面的事件跟着摘掉）， +/// 名额还回去。 +struct Held(OwnedFd); + +impl Held { + /// 打开 `path`。名额用完了是 `Ok(None)`。 + fn open(path: &Path) -> io::Result> { + let reserved = HELD.fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| { + (n < budget()).then_some(n + 1) + }); + if reserved.is_err() { + return Ok(None); + } + let release = || HELD.fetch_sub(1, Ordering::AcqRel); + let Ok(c) = CString::new(path.as_os_str().as_bytes()) else { + release(); + return Err(io::ErrorKind::InvalidInput.into()); + }; + // O_NONBLOCK：万一打开的是个管道（条目刚被换成了别的东西），也不会卡在这里 + // SAFETY: c 是以 NUL 结尾的路径；返回的 fd 由下面的 OwnedFd 接管 + let fd = unsafe { + libc::open( + c.as_ptr(), + libc::O_EVTONLY | libc::O_CLOEXEC | libc::O_NONBLOCK, + ) + }; + if fd < 0 { + let e = io::Error::last_os_error(); + release(); + return Err(e); + } + // SAFETY: fd 是刚打开的、只有这里拿着 + Ok(Some(Held(unsafe { OwnedFd::from_raw_fd(fd) }))) + } + + fn raw(&self) -> RawFd { + self.0.as_raw_fd() + } + + /// fstat：用的是它的类型、(dev, ino)、大小和修改时间 + fn stat(&self) -> io::Result { + // SAFETY: stat 是纯输出参数，全零是合法的初值 + let mut st: libc::stat = unsafe { std::mem::zeroed() }; + // SAFETY: fd 活着；fstat 只往 st 里写 + if unsafe { libc::fstat(self.raw(), &mut st) } != 0 { + return Err(io::Error::last_os_error()); + } + Ok(st) + } +} + +impl Drop for Held { + fn drop(&mut self) { + HELD.fetch_sub(1, Ordering::AcqRel); + } +} + +/// 一个活着的 kqueue 监听；名额不够时是退回去的 FSEvents。 +pub struct Watch { + _backend: Backend, +} + +/// 拿着它就是在盯；扔掉就停 +enum Backend { + Kqueue { + _running: Running, + }, + Notify { + _fsevents: notify::RecommendedWatcher, + }, +} + +/// 跑着的事件循环：kqueue 本身，和跑循环的线程 +struct Running { + kq: Arc, + thread: Option>, + /// 事件循环醒过几次（测试数它：深处的改动不该叫醒它） + #[cfg(test)] + wakeups: Arc, +} + +impl Drop for Running { + /// **等线程收完再返回**：它手上的 fd 这时都已经关了、名额都还了，调用方紧接着再起一个 + /// 监听不会撞上名额 + fn drop(&mut self) { + let _ = control( + self.kq.as_raw_fd(), + &[event(STOP, libc::EVFILT_USER, 0, libc::NOTE_TRIGGER)], + ); + if let Some(t) = self.thread.take() { + let _ = t.join(); + } + } +} + +impl Watch { + /// 事件循环到现在醒过几次。退回 FSEvents 的是 `None` + #[cfg(test)] + pub(crate) fn wakeups(&self) -> Option { + match &self._backend { + Backend::Kqueue { _running: r } => Some(r.wakeups.load(Ordering::Relaxed)), + Backend::Notify { .. } => None, + } + } +} + +/// 盯住 `dirs`：见模块文档。每有一处算数的改动，往 `raw_tx` 里发一个。 +/// +/// **返回时就已经盯上了**：目录和文件都挂好、列过一遍，之后的改动一个都不会漏。 +pub(crate) fn watch( + dirs: &[PathBuf], + relevant: Relevant, + raw_tx: Sender<()>, +) -> Result { + let start_error = |path: &Path, e: io::Error| WatchError::Start { + path: path.to_path_buf(), + source: notify::Error::io(e), + }; + let first = dirs.first().map(PathBuf::as_path).unwrap_or(Path::new("")); + // SAFETY: kqueue() 不碰内存；返回的 fd 由 OwnedFd 接管 + let kq = unsafe { libc::kqueue() }; + if kq < 0 { + return Err(start_error(first, io::Error::last_os_error())); + } + // SAFETY: kq 是刚建的、只有这里拿着 + let kq = Arc::new(unsafe { OwnedFd::from_raw_fd(kq) }); + // fork 出去的子进程不该拿着它 + // SAFETY: kq 活着；FD_CLOEXEC 只改这个 fd 自己的标志 + unsafe { libc::fcntl(kq.as_raw_fd(), libc::F_SETFD, libc::FD_CLOEXEC) }; + control( + kq.as_raw_fd(), + &[event( + STOP, + libc::EVFILT_USER, + libc::EV_ADD | libc::EV_CLEAR, + 0, + )], + ) + .map_err(|e| start_error(first, e))?; + + let mut l = Loop { + kq: kq.clone(), + dirs: dirs + .iter() + .map(|p| Dir { + path: p.clone(), + fd: None, + parent: None, + seen: HashMap::new(), + }) + .collect(), + files: HashMap::new(), + by_fd: HashMap::new(), + relevant: Arc::new(Mutex::new(relevant)), + capped: false, + #[cfg(test)] + wakeups: Arc::new(AtomicUsize::new(0)), + }; + let mut over = false; + for (i, path) in dirs.iter().enumerate() { + match l.open_dir(i) { + Ok(true) => {} + Ok(false) => { + over = true; + break; + } + Err(e) => return Err(start_error(path, e)), + } + } + // 先挂好再列：列的时候已经有改动的话，事件也已经排在那儿了 + if !over { + over = (0..l.dirs.len()).any(|i| l.rescan(i).is_err()); + } + if over { + // 一开始就不够：手上的放掉，整个用 FSEvents + let relevant = l.handover(); + drop(l); + return Ok(Watch { + _backend: Backend::Notify { + _fsevents: crate::by_notify(dirs, relevant, raw_tx)?, + }, + }); + } + + #[cfg(test)] + let wakeups = l.wakeups.clone(); + let thread = std::thread::Builder::new() + .name("tw-watch kqueue".into()) + .spawn(move || l.run(&raw_tx)) + .map_err(|e| start_error(first, e))?; + Ok(Watch { + _backend: Backend::Kqueue { + _running: Running { + kq, + thread: Some(thread), + #[cfg(test)] + wakeups, + }, + }, + }) +} + +/// 一个盯着的目录 +struct Dir { + path: PathBuf, + /// 没了、挪走了，那里还没有新的时是 None：这时盯着 `parent` + fd: Option, + /// 它没了的时候盯着的上一层：等这个名字再出现 + parent: Option, + /// 上次列出来的、算数的条目：名字 → (inode, 是什么) + seen: HashMap, +} + +/// 一个条目是什么（readdir 给的，不另外 stat） +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Kind { + File, + Symlink, + Other, +} + +/// 一个盯着的文件（目录里算数的普通文件，或者指向普通文件的符号链接） +struct File { + fd: Held, + /// 它在哪个目录里（[`Loop::dirs`] 的下标） + dir: usize, + /// 挂上时它是哪个文件（跟完符号链接） + dev: u64, + ino: u64, + /// (大小, 修改时间的秒, 纳秒)：属性变了时拿来比，见 [`Loop::stamp_moved`] + stamp: (i64, i64, i64), +} + +/// 一个 fd 挂的是谁 +enum Target { + Dir(usize), + /// 第几个目录没了之后盯着的上一层 + Parent(usize), + File(PathBuf), +} + +/// 名额用完了：整个换回 FSEvents +struct OverBudget; + +/// 事件循环为什么停下 +enum Exit { + /// 监听被扔掉了，或者去抖那一头没了 + Stop, + /// 名额用完了 + OverBudget, +} + +struct Loop { + kq: Arc, + dirs: Vec, + files: HashMap, + by_fd: HashMap, + /// 退回 FSEvents 时要交给它，交不成还得自己接着用：所以是共用的 + relevant: Arc>, + /// 退不回 FSEvents（它起不来）：名额不够的就不盯，能盯多少盯多少 + capped: bool, + #[cfg(test)] + wakeups: Arc, +} + +impl Loop { + fn run(mut self, raw_tx: &Sender<()>) { + loop { + if let Exit::Stop = self.serve(raw_tx) { + return; + } + // **先起 FSEvents，再放手上的**：中间那一下不会有改动没人看。此刻没了的目录不盯 + // （FSEvents 盯不存在的路径会报错，整个起不来） + let alive: Vec = self + .dirs + .iter() + .map(|d| d.path.clone()) + .filter(|p| p.is_dir()) + .collect(); + // 名额用完的那一批改动没算完：当作有过一次 + let _ = raw_tx.send(()); + match crate::by_notify(&alive, self.handover(), raw_tx.clone()) { + Ok(fsevents) => { + self.files.clear(); + self.by_fd.clear(); + self.dirs.clear(); + self.wait_for_stop(); + drop(fsevents); + return; + } + // 起不来就接着用 kqueue，盯得住多少算多少 + Err(_) => self.capped = true, + } + } + } + + /// 交给 FSEvents 的那一份 `relevant`：和这里用的是同一个 + fn handover(&self) -> Relevant { + let r = self.relevant.clone(); + Box::new(move |p| (r.lock().unwrap_or_else(PoisonError::into_inner))(p)) + } + + /// `p` 算不算数 + fn relevant(&self, p: &Path) -> bool { + (self.relevant.lock().unwrap_or_else(PoisonError::into_inner))(p) + } + + /// 收事件，直到被叫停或者名额用完 + fn serve(&mut self, raw_tx: &Sender<()>) -> Exit { + let mut buf = [event(0, 0, 0, 0); 64]; + loop { + let n = match wait(self.kq.as_raw_fd(), &mut buf) { + Ok(n) => n, + Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, + Err(_) => return Exit::Stop, + }; + #[cfg(test)] + self.wakeups.fetch_add(1, Ordering::Relaxed); + match self.handle(&buf[..n]) { + Ok(false) => {} + Ok(true) => { + if raw_tx.send(()).is_err() { + return Exit::Stop; + } + } + Err(exit) => return exit, + } + } + } + + /// 处理一批事件，有算数的改动是 `Ok(true)`。 + /// + /// **分两遍**：第一遍只记下要做什么，第二遍才关 fd、重新打开。关掉的 fd 号马上会被 + /// 新打开的复用，一边处理一边开关的话，同一批里排在后面的、关于旧文件的事件会被 + /// 算到新文件头上。 + fn handle(&mut self, events: &[libc::kevent]) -> Result { + let mut changed = false; + let mut rescan = BTreeSet::new(); + let mut gone_dirs = Vec::new(); + let mut gone_parents = Vec::new(); + let mut gone_files = Vec::new(); + for ev in events { + if ev.filter == libc::EVFILT_USER { + return Err(Exit::Stop); + } + let Ok(fd) = RawFd::try_from(ev.ident) else { + continue; + }; + match self.by_fd.get(&fd) { + Some(Target::Dir(i)) => { + let i = *i; + if ev.fflags & GONE != 0 { + changed |= self.relevant(&self.dirs[i].path); + gone_dirs.push(i); + } + rescan.insert(i); + } + // 上一层的条目变了：等的那个名字可能出现了 + Some(Target::Parent(i)) => { + let i = *i; + if ev.fflags & GONE != 0 { + gone_parents.push(i); + } + rescan.insert(i); + } + Some(Target::File(path)) => { + let path = path.clone(); + let Some(f) = self.files.get_mut(&path) else { + continue; + }; + // 写过就算；大小和修改时间每次都记下新的，之后单报属性时才比得准 + let wrote = ev.fflags & (libc::NOTE_WRITE | libc::NOTE_EXTEND) != 0; + let moved = + ev.fflags & (libc::NOTE_WRITE | libc::NOTE_EXTEND | libc::NOTE_ATTRIB) != 0 + && Self::stamp_moved(f); + changed |= wrote || moved; + if ev.fflags & GONE != 0 { + changed = true; + rescan.insert(f.dir); + gone_files.push(path); + } + } + None => {} + } + } + for path in gone_files { + self.drop_file(&path); + } + for i in gone_dirs { + self.drop_dir(i); + } + for i in gone_parents { + self.drop_parent(i); + } + for i in rescan { + changed |= self.rescan(i).map_err(|OverBudget| Exit::OverBudget)?; + } + Ok(changed) + } + + /// 属性变了的文件，大小或修改时间是不是也变了（变了就记下新的）。截断只报属性； + /// 改权限、改属主也报属性，但它们不算 + fn stamp_moved(f: &mut File) -> bool { + let Ok(st) = f.fd.stat() else { + return false; + }; + let now = stamp(&st); + let moved = now != f.stamp; + f.stamp = now; + moved + } + + /// 打开第 `i` 个目录、挂上事件。名额用完了是 `Ok(false)`。 + fn open_dir(&mut self, i: usize) -> io::Result { + let Some(fd) = Held::open(&self.dirs[i].path)? else { + return Ok(false); + }; + if fd.stat()?.st_mode & libc::S_IFMT != libc::S_IFDIR { + return Err(io::Error::from_raw_os_error(libc::ENOTDIR)); + } + arm(self.kq.as_raw_fd(), fd.raw(), DIR_EVENTS)?; + self.by_fd.insert(fd.raw(), Target::Dir(i)); + self.dirs[i].fd = Some(fd); + Ok(true) + } + + /// 第 `i` 个目录自己没了、被挪走：手上那个 fd 不再是那个路径上的东西，放掉 + fn drop_dir(&mut self, i: usize) { + if let Some(fd) = self.dirs[i].fd.take() { + self.by_fd.remove(&fd.raw()); + } + } + + /// 第 `i` 个目录的上一层也没了、被挪走：不再等 + fn drop_parent(&mut self, i: usize) { + if let Some(fd) = self.dirs[i].parent.take() { + self.by_fd.remove(&fd.raw()); + } + } + + /// 没了的第 `i` 个目录：按原来的路径再打开。**先盯上一层，再试着打开** —— 反过来的话， + /// 两步之间换上来的那个就没人看见了。打开了就不用再盯上一层；打不开就等上一层报信 + fn revive(&mut self, i: usize) -> Result<(), OverBudget> { + let over = |capped: bool| if capped { Ok(()) } else { Err(OverBudget) }; + if self.dirs[i].parent.is_none() + && let Some(up) = self.dirs[i].path.parent().map(Path::to_path_buf) + && !up.as_os_str().is_empty() + { + match Held::open(&up) { + Ok(Some(fd)) => { + if arm(self.kq.as_raw_fd(), fd.raw(), DIR_EVENTS).is_ok() { + self.by_fd.insert(fd.raw(), Target::Parent(i)); + self.dirs[i].parent = Some(fd); + } + } + Ok(None) => over(self.capped)?, + // 上一层也没了：不再盯 + Err(_) => {} + } + } + match self.open_dir(i) { + Ok(true) => self.drop_parent(i), + Ok(false) => over(self.capped)?, + Err(_) => {} + } + Ok(()) + } + + /// 重新列第 `i` 个目录：算数的条目多了、少了、换了，是 `Ok(true)`。顺带把要盯的文件 + /// 对上：新的挂上，换掉的重新挂，不在了的放掉。 + /// + /// 目录自己没了的，先按原来的路径再打开（[`Self::revive`]）：又有了就接着盯它，它的 + /// 出现也算一次；没有的话，算数的条目一个都不剩了。 + fn rescan(&mut self, i: usize) -> Result { + let mut changed = false; + if self.dirs[i].fd.is_none() { + self.revive(i)?; + if self.dirs[i].fd.is_some() { + changed |= self.relevant(&self.dirs[i].path); + } + } + let now = match &self.dirs[i].fd { + Some(_) => self.list(&self.dirs[i].path), + None => HashMap::new(), + }; + changed |= now != self.dirs[i].seen; + self.dirs[i].seen = now; + self.sync_files(i)?; + Ok(changed) + } + + /// 目录里算数的条目 + fn list(&self, dir: &Path) -> HashMap { + let Ok(rd) = std::fs::read_dir(dir) else { + return HashMap::new(); + }; + rd.flatten() + .filter(|e| self.relevant(&dir.join(e.file_name()))) + .map(|e| { + // readdir 给的类型；给不出来时 std 自己补一次 lstat + let kind = match e.file_type() { + Ok(t) if t.is_file() => Kind::File, + Ok(t) if t.is_symlink() => Kind::Symlink, + _ => Kind::Other, + }; + (e.file_name(), (e.ino(), kind)) + }) + .collect() + } + + /// 把第 `i` 个目录里要盯的文件对上它此刻的样子（[`Dir::seen`]） + fn sync_files(&mut self, i: usize) -> Result<(), OverBudget> { + let dir = &self.dirs[i]; + let want: Vec<(PathBuf, u64, Kind)> = dir + .seen + .iter() + .filter(|(_, (_, kind))| *kind != Kind::Other) + .map(|(name, (ino, kind))| (dir.path.join(name), *ino, *kind)) + .collect(); + let stale: Vec = self + .files + .iter() + .filter(|(p, f)| f.dir == i && !want.iter().any(|(w, ..)| w == *p)) + .map(|(p, _)| p.clone()) + .collect(); + for p in stale { + self.drop_file(&p); + } + for (path, ino, kind) in want { + if let Some(f) = self.files.get(&path) { + // 还是挂着的那一个就不动。符号链接按它指向的那个文件比 + let same = match kind { + Kind::Symlink => { + std::fs::metadata(&path).is_ok_and(|m| (m.dev(), m.ino()) == (f.dev, f.ino)) + } + _ => ino == f.ino, + }; + if same { + continue; + } + self.drop_file(&path); + } + // 符号链接指向的不是普通文件（目录、设备、不存在）就不盯 + if kind == Kind::Symlink && !std::fs::metadata(&path).is_ok_and(|m| m.is_file()) { + continue; + } + self.add_file(i, path)?; + } + Ok(()) + } + + /// 挂上一个文件。打不开（刚好又没了、没有权限）就算了：它的增删改名目录那边照样看得见 + fn add_file(&mut self, dir: usize, path: PathBuf) -> Result<(), OverBudget> { + let fd = match Held::open(&path) { + Ok(Some(fd)) => fd, + Ok(None) if self.capped => return Ok(()), + Ok(None) => return Err(OverBudget), + Err(_) => return Ok(()), + }; + let Ok(st) = fd.stat() else { return Ok(()) }; + if st.st_mode & libc::S_IFMT != libc::S_IFREG { + return Ok(()); + } + if arm(self.kq.as_raw_fd(), fd.raw(), FILE_EVENTS).is_err() { + return Ok(()); + } + self.by_fd.insert(fd.raw(), Target::File(path.clone())); + self.files.insert( + path, + File { + dir, + dev: st.st_dev as u64, + ino: st.st_ino, + stamp: stamp(&st), + fd, + }, + ); + Ok(()) + } + + fn drop_file(&mut self, path: &Path) { + if let Some(f) = self.files.remove(path) { + self.by_fd.remove(&f.fd.raw()); + } + } + + /// 只等叫停的那个事件 + fn wait_for_stop(&self) { + let mut buf = [event(0, 0, 0, 0); 4]; + loop { + match wait(self.kq.as_raw_fd(), &mut buf) { + Ok(n) if buf[..n].iter().any(|e| e.filter == libc::EVFILT_USER) => return, + Ok(_) => {} + Err(e) if e.kind() == io::ErrorKind::Interrupted => {} + Err(_) => return, + } + } + } +} + +/// (大小, 修改时间的秒, 纳秒) +fn stamp(st: &libc::stat) -> (i64, i64, i64) { + (st.st_size, st.st_mtime, st.st_mtime_nsec) +} + +fn event(ident: usize, filter: i16, flags: u16, fflags: u32) -> libc::kevent { + libc::kevent { + ident, + filter, + flags, + fflags, + data: 0, + udata: std::ptr::null_mut(), + } +} + +/// 在 `fd` 上挂 vnode 事件。`EV_CLEAR`：报过一次就清掉，不会同一件事一直报 +fn arm(kq: RawFd, fd: RawFd, fflags: u32) -> io::Result<()> { + let ident = usize::try_from(fd).map_err(|_| io::Error::from_raw_os_error(libc::EBADF))?; + control( + kq, + &[event( + ident, + libc::EVFILT_VNODE, + libc::EV_ADD | libc::EV_CLEAR, + fflags, + )], + ) +} + +/// 改 kqueue 里挂着的事件，不收 +fn control(kq: RawFd, changes: &[libc::kevent]) -> io::Result<()> { + let n = libc::c_int::try_from(changes.len()).unwrap_or(libc::c_int::MAX); + // SAFETY: changes 活着、长度如实；不收事件，输出缓冲区是空的 + let r = unsafe { + libc::kevent( + kq, + changes.as_ptr(), + n, + std::ptr::null_mut(), + 0, + std::ptr::null(), + ) + }; + if r < 0 { + return Err(io::Error::last_os_error()); + } + Ok(()) +} + +/// 等到有事件为止（不设超时），收进 `buf`，交回收到几个 +fn wait(kq: RawFd, buf: &mut [libc::kevent]) -> io::Result { + let cap = libc::c_int::try_from(buf.len()).unwrap_or(libc::c_int::MAX); + // SAFETY: buf 活着、可写、长度如实；不改挂着的事件 + let n = unsafe { + libc::kevent( + kq, + std::ptr::null(), + 0, + buf.as_mut_ptr(), + cap, + std::ptr::null(), + ) + }; + usize::try_from(n).map_err(|_| io::Error::last_os_error()) +} + +#[cfg(test)] +mod tests { + use std::path::Path; + use std::time::Duration; + + use crate::tests::{Signal, WINDOW, append, only_md, recv, save_by_rename, scratch, settle}; + + /// 事件循环醒过几次 + fn wakeups(w: &crate::Watch) -> usize { + w._inner.wakeups().expect("退回了 FSEvents") + } + + /// **这一条就是用 kqueue 的原因。**子目录里再怎么写，事件循环一次都不醒 —— FSEvents + /// 在这里每一次写都会叫醒我们，再在进程里丢掉。而盯着的目录自己的条目变了，照样醒 + #[tokio::test] + async fn changes_deeper_down_do_not_even_wake_it() { + let d = scratch(); + let deep = d.path().join("history/2026/10"); + std::fs::create_dir_all(&deep).unwrap(); + std::fs::write(d.path().join("a.md"), "0\n").unwrap(); + let (w, mut rx) = crate::watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + for i in 0..200 { + std::fs::write(deep.join(format!("{i}.md")), "x\n").unwrap(); + } + std::fs::create_dir_all(deep.join("more/and/more")).unwrap(); + std::fs::remove_dir_all(d.path().join("history/2026")).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_millis(600)).await, + Signal::Quiet + ); + assert_eq!(wakeups(&w), 0, "深处的改动叫醒了事件循环"); + + // 这一层的条目变了才醒；不算数的条目醒了也不发信号 + std::fs::write(d.path().join("x.log"), "x\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_millis(600)).await, + Signal::Quiet + ); + assert!(wakeups(&w) > 0, "这一层的改动没叫醒它"); + } + + /// 属性变了只有大小、修改时间跟着变才算：改权限不算，`touch` 算 + #[tokio::test] + async fn only_size_and_modification_time_count_among_attributes() { + use std::os::unix::fs::PermissionsExt as _; + let d = scratch(); + let p = d.path().join("a.md"); + std::fs::write(&p, "0\n").unwrap(); + let (_w, mut rx) = crate::watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o600)).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_millis(600)).await, + Signal::Quiet, + "改权限" + ); + let later = std::time::SystemTime::now() + Duration::from_secs(5); + std::fs::File::options() + .write(true) + .open(&p) + .unwrap() + .set_modified(later) + .unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "touch" + ); + } + + /// 符号链接指过去的文件也盯：原地写算数；它被换掉之后，重新挂到换上来的那一个 + /// （`~/.claude.json` 指进 dotfiles 仓库的那种布局） + #[tokio::test] + async fn a_symlinked_file_is_followed_to_where_it_points() { + let d = scratch(); + let (watched, elsewhere) = (d.path().join("home"), d.path().join("dotfiles")); + std::fs::create_dir_all(&watched).unwrap(); + std::fs::create_dir_all(&elsewhere).unwrap(); + let real = elsewhere.join("a.md"); + std::fs::write(&real, "0\n").unwrap(); + std::os::unix::fs::symlink(&real, watched.join("a.md")).unwrap(); + let (_w, mut rx) = crate::watch(std::slice::from_ref(&watched), WINDOW, only_md).unwrap(); + + append(&real, "1\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "追加" + ); + settle(&mut rx).await; + save_by_rename(&real, "2\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "换文件" + ); + settle(&mut rx).await; + append(&real, "3\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "换过之后的追加" + ); + } + + /// 盯着的目录被整个换掉（rename 一个新目录过来）、删了过一阵又建出来：目录自己算数的话 + /// 各算一次，之后接着盯新的那个。盯的是路径，和 FSEvents 一样 + #[tokio::test] + async fn a_watched_directory_that_is_replaced_or_recreated_is_followed() { + let d = scratch(); + let dir = d.path().join("skills"); + std::fs::create_dir_all(&dir).unwrap(); + let me = dir.clone(); + let (_w, mut rx) = crate::watch(std::slice::from_ref(&dir), WINDOW, move |p: &Path| { + p == me || only_md(p) + }) + .unwrap(); + + let next = d.path().join("skills.new"); + std::fs::create_dir_all(&next).unwrap(); + std::fs::rename(&dir, d.path().join("skills.old")).unwrap(); + std::fs::rename(&next, &dir).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "目录被挪走" + ); + settle(&mut rx).await; + std::fs::write(dir.join("a.md"), "0\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "换上来的目录没盯上" + ); + settle(&mut rx).await; + // 挪走的那个不再盯 + std::fs::write(d.path().join("skills.old/b.md"), "0\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_millis(600)).await, + Signal::Quiet + ); + + std::fs::remove_dir_all(&dir).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "删除" + ); + settle(&mut rx).await; + std::fs::create_dir(&dir).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "再建" + ); + settle(&mut rx).await; + std::fs::write(dir.join("c.md"), "0\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "再建出来的目录没盯上" + ); + } +} diff --git a/crates/tw-watch/src/lib.rs b/crates/tw-watch/src/lib.rs index 3ca8b18f..1a99809e 100644 --- a/crates/tw-watch/src/lib.rs +++ b/crates/tw-watch/src/lib.rs @@ -9,10 +9,14 @@ //! //! **一、盯目录，不盯文件。**编辑器保存不是「往原文件写」，而是「写一个临时文件 //! 再 rename 过去」。rename 之后原来那个 inode 就没人指了 —— 盯着文件的监听会在 -//! 第一次保存之后**永久失效**，而且不报任何错。 +//! 第一次保存之后**永久失效**，而且不报任何错。（macOS 的 kqueue 看不见目录里文件的 +//! 原地写，所以那边算数的文件**另外**也盯着，换掉之后由目录那边重新挂上，见 `kqueue` 模块。） //! //! **二、不递归。**配置旁边就是历史目录；`~/.claude/` 底下有每分钟都在变的东西。 -//! 递归盯它们等于给自己找一个永不停歇的事件源。 +//! 递归盯它们等于给自己找一个永不停歇的事件源。**不递归要做到内核那一层**：macOS 上 +//! notify 默认的 FSEvents 只会盯整棵树，「不递归」是在我们进程里把深处的事件丢掉 —— +//! 丢之前已经被叫醒了。所以 macOS 上用 kqueue（`kqueue` 模块），Linux 的 inotify、 +//! Windows 的 ReadDirectoryChangesW 本来就只报这一层，照旧用 notify。 //! //! **三、去抖。**一次保存常常触发好几个事件（建临时文件、rename、改属性、改 //! mtime）。不聚合的话，一次保存会触发好几轮重读。 @@ -20,6 +24,9 @@ use std::path::{Path, PathBuf}; use std::time::Duration; +#[cfg(target_os = "macos")] +mod kqueue; + #[derive(Debug, thiserror::Error)] pub enum WatchError { #[error("changes to {path} could not be watched: {source}")] @@ -31,9 +38,17 @@ pub enum WatchError { /// 一个活着的监听。**扔掉它就停止监听** —— 所以调用方必须留着它。 pub struct Watch { - _inner: notify::RecommendedWatcher, + _inner: Inner, } +#[cfg(target_os = "macos")] +type Inner = kqueue::Watch; +#[cfg(not(target_os = "macos"))] +type Inner = notify::RecommendedWatcher; + +/// 哪些路径算数（[`watch`] 的 `relevant`） +type Relevant = Box bool + Send + 'static>; + /// 盯住 `dirs`（不递归），改动涉及的路径里有一个让 `relevant` 说是的，就算一次； /// 一次之后 `debounce` 之内的都并进去，然后往通道里发**一个**信号。 /// @@ -43,6 +58,11 @@ pub struct Watch { /// 只看内容和存在性的变化（建、改、删）。**访问时间之类的不算** —— 备份和索引 /// 工具会大量产生它们。 /// +/// `relevant` 拿到的是改动的那个条目的路径（盯着的目录自己没了、被挪走时是目录自己）。 +/// 写法各平台不完全一样：inotify、kqueue 是 `dirs` 里那个目录接上条目名，FSEvents（macOS +/// 上名额不够时退回去用它）是跟完符号链接的真实路径 —— 要按整条路径比的，盯的目录先 +/// 换成真实路径，两种写法就一样了。 +/// /// 通道里最多攒一个信号：满了就丢。信号的意思是「去看一眼」，攒两个和攒一个是 /// 一回事。 pub fn watch( @@ -50,10 +70,29 @@ pub fn watch( debounce: Duration, relevant: impl Fn(&Path) -> bool + Send + 'static, ) -> Result<(Watch, tokio::sync::mpsc::Receiver<()>), WatchError> { - use notify::{EventKind, RecursiveMode, Watcher as _}; - let (tx, rx) = tokio::sync::mpsc::channel(1); let (raw_tx, raw_rx) = std::sync::mpsc::channel::<()>(); + #[cfg(target_os = "macos")] + let inner = kqueue::watch(dirs, Box::new(relevant), raw_tx)?; + #[cfg(not(target_os = "macos"))] + let inner = by_notify(dirs, Box::new(relevant), raw_tx)?; + + // 监听的回调跑在它自己的线程上，这里把它接到 tokio 上并去抖。 + std::thread::spawn(move || debounce_loop(&raw_rx, debounce, &tx)); + + Ok((Watch { _inner: inner }, rx)) +} + +/// 用 notify 给这个平台推荐的那一种盯：Linux 的 inotify、Windows 的 +/// ReadDirectoryChangesW。macOS 上是 FSEvents，只在 kqueue 的 fd 名额用完时退回来用它 +/// （见 `kqueue` 模块）。每有一处算数的改动，往 `raw_tx` 里发一个。 +fn by_notify( + dirs: &[PathBuf], + relevant: Relevant, + raw_tx: std::sync::mpsc::Sender<()>, +) -> Result { + use notify::{EventKind, RecursiveMode, Watcher as _}; + let mut w = notify::recommended_watcher(move |ev: notify::Result| { let Ok(ev) = ev else { return }; if !matches!( @@ -77,11 +116,7 @@ pub fn watch( source, })?; } - - // notify 的回调跑在它自己的线程上，这里把它接到 tokio 上并去抖。 - std::thread::spawn(move || debounce_loop(&raw_rx, debounce, &tx)); - - Ok((Watch { _inner: w }, rx)) + Ok(w) } /// 去抖：收到一个事件之后，把间隔不到 `window` 的后续事件全部并进来，静下来 @@ -110,21 +145,21 @@ fn debounce_loop( } #[cfg(test)] -mod tests { +pub(crate) mod tests { use super::*; - const WINDOW: Duration = Duration::from_millis(200); + pub(crate) const WINDOW: Duration = Duration::from_millis(200); /// 等一个信号，**区分三种结果**：收到了、窗口内没动静、通道关了（监听线程 /// 没了之后永远不会再有信号 —— 不分开的话它和「收到了」分不清）。 #[derive(Debug, PartialEq)] - enum Signal { + pub(crate) enum Signal { Got, Quiet, Closed, } - async fn recv(rx: &mut tokio::sync::mpsc::Receiver<()>, within: Duration) -> Signal { + pub(crate) async fn recv(rx: &mut tokio::sync::mpsc::Receiver<()>, within: Duration) -> Signal { match tokio::time::timeout(within, rx.recv()).await { Ok(Some(())) => Signal::Got, Ok(None) => Signal::Closed, @@ -132,7 +167,7 @@ mod tests { } } - fn only_md(p: &Path) -> bool { + pub(crate) fn only_md(p: &Path) -> bool { p.extension().is_some_and(|e| e == "md") } @@ -145,7 +180,7 @@ mod tests { /// 是对的，错的是测试。内存盘上的写不碰磁盘，前提才真的成立。 /// /// macOS 没有 `/dev/shm`，用系统的临时目录 —— 那边从来没有因此失败过。 - fn scratch() -> tempfile::TempDir { + pub(crate) fn scratch() -> tempfile::TempDir { let shm = Path::new("/dev/shm"); if cfg!(target_os = "linux") && shm.is_dir() { tempfile::tempdir_in(shm).unwrap() @@ -244,4 +279,158 @@ mod tests { }; assert_eq!(path, missing); } + + /// 调用方把它放在跨线程共用的地方（桌面端放在应用状态的 Mutex 里） + #[test] + fn a_watch_can_move_between_threads() { + fn send() {} + send::(); + } + + /// 等前面那些改动的信号都收完、静下来：之后每一步要的是它自己的那个信号 + pub(crate) async fn settle(rx: &mut tokio::sync::mpsc::Receiver<()>) { + while recv(rx, WINDOW * 3).await == Signal::Got {} + } + + /// 往文件末尾追加：原地写，不经过 rename + pub(crate) fn append(p: &Path, text: &str) { + use std::io::Write as _; + let mut f = std::fs::OpenOptions::new().append(true).open(p).unwrap(); + f.write_all(text.as_bytes()).unwrap(); + } + + /// 编辑器那种保存：写一个临时文件，rename 过去 + pub(crate) fn save_by_rename(p: &Path, text: &str) { + let tmp = p.with_extension("swp"); + std::fs::write(&tmp, text).unwrap(); + std::fs::rename(&tmp, p).unwrap(); + } + + #[tokio::test] + async fn every_save_by_rename_is_seen() { + let d = scratch(); + let p = d.path().join("a.md"); + std::fs::write(&p, "0\n").unwrap(); + let (_w, mut rx) = watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + for i in 1..=3 { + save_by_rename(&p, &format!("{i}\n")); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "第 {i} 次 rename 保存没被发现" + ); + settle(&mut rx).await; + } + } + + /// 原地写也算：追加、截断都不经过 rename。**换过一次文件之后还看得见** —— 盯着文件本身 + /// 的实现（macOS）手上的是换掉之前的那个 inode，要重新挂到新的上 + #[tokio::test] + async fn writing_in_place_is_seen_also_after_the_file_was_replaced() { + let d = scratch(); + let p = d.path().join("a.md"); + std::fs::write(&p, "0\n").unwrap(); + let (_w, mut rx) = watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + + append(&p, "1\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "追加" + ); + settle(&mut rx).await; + + save_by_rename(&p, "2\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "换文件" + ); + settle(&mut rx).await; + + append(&p, "3\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "换过之后的追加" + ); + settle(&mut rx).await; + + // 只截断、不写：大小变了，别的什么都没动 + std::fs::File::create(&p).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "截断" + ); + } + + /// 后来才有的文件：建出来算一次，之后原地写也算 + #[tokio::test] + async fn a_file_created_later_is_seen_and_so_are_its_writes() { + let d = scratch(); + let p = d.path().join("a.md"); + let (_w, mut rx) = watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + std::fs::write(&p, "0\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "新建" + ); + settle(&mut rx).await; + append(&p, "1\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "追加" + ); + } + + /// 删掉算一次，再建出来又算一次，再建出来的那个原地写也算 + #[tokio::test] + async fn a_deleted_file_that_comes_back_is_seen_again() { + let d = scratch(); + let p = d.path().join("a.md"); + std::fs::write(&p, "0\n").unwrap(); + let (_w, mut rx) = watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + std::fs::remove_file(&p).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "删除" + ); + settle(&mut rx).await; + std::fs::write(&p, "1\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "再建" + ); + settle(&mut rx).await; + append(&p, "2\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "追加" + ); + } + + /// **读文件不算。**调用方收到信号就去读：读也算的话，读一次就又叫醒自己一次 + #[tokio::test] + async fn reading_a_watched_file_stays_quiet() { + let d = scratch(); + let p = d.path().join("a.md"); + std::fs::write(&p, "0\n").unwrap(); + let (_w, mut rx) = watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + for _ in 0..20 { + std::fs::read_to_string(&p).unwrap(); + std::fs::metadata(&p).unwrap(); + } + std::fs::read_dir(d.path()).unwrap().for_each(drop); + assert_eq!( + recv(&mut rx, Duration::from_millis(800)).await, + Signal::Quiet, + "读文件把它叫醒了" + ); + } } diff --git a/crates/tw-watch/tests/fd_budget.rs b/crates/tw-watch/tests/fd_budget.rs new file mode 100644 index 00000000..9b0244f6 --- /dev/null +++ b/crates/tw-watch/tests/fd_budget.rs @@ -0,0 +1,122 @@ +//! macOS 上 fd 名额紧的时候：kqueue 每盯一个目录、一份文件占一个 fd，而从访达打开的应用 +//! 只有 256 个。名额不够就整个换回 FSEvents —— 费电，但一个改动都不漏。 +//! +//! **单独一个测试进程、只有一个测试**：要调低整个进程的 fd 软上限、数整个进程开着的 fd， +//! 和别的测试放在一起会互相连累。 +#![cfg(target_os = "macos")] + +use std::path::{Path, PathBuf}; +use std::time::Duration; + +const WINDOW: Duration = Duration::from_millis(200); + +fn only_md(p: &Path) -> bool { + p.extension().is_some_and(|e| e == "md") +} + +/// 此刻整个进程开着几个 fd（列 `/dev/fd` 自己也占一个，每次都一样） +fn open_fds() -> usize { + std::fs::read_dir("/dev/fd").unwrap().count() +} + +/// 把 fd 软上限调到 `n`：监听最多能占一半 +fn set_soft_limit(n: libc::rlim_t) { + let mut r = libc::rlimit { + rlim_cur: 0, + rlim_max: 0, + }; + // SAFETY: 两个调用都只读写传进去的结构体 + unsafe { + assert_eq!(libc::getrlimit(libc::RLIMIT_NOFILE, &mut r), 0); + r.rlim_cur = n; + assert_eq!(libc::setrlimit(libc::RLIMIT_NOFILE, &r), 0); + } +} + +fn append(p: &Path, text: &str) { + use std::io::Write as _; + let mut f = std::fs::OpenOptions::new().append(true).open(p).unwrap(); + f.write_all(text.as_bytes()).unwrap(); +} + +/// `n` 个目录，各有 `files` 份算数的文件 +fn dirs(root: &Path, prefix: &str, n: usize, files: usize) -> Vec { + (0..n) + .map(|i| { + let d = root.join(format!("{prefix}{i}")); + std::fs::create_dir_all(&d).unwrap(); + for j in 0..files { + std::fs::write(d.join(format!("{j}.md")), "0\n").unwrap(); + } + d + }) + .collect() +} + +async fn got(rx: &mut tokio::sync::mpsc::Receiver<()>) -> bool { + matches!( + tokio::time::timeout(Duration::from_secs(5), rx.recv()).await, + Ok(Some(())) + ) +} + +/// 等前面那些改动的信号都收完、静下来 +async fn settle(rx: &mut tokio::sync::mpsc::Receiver<()>) { + while let Ok(Some(())) = tokio::time::timeout(WINDOW * 3, rx.recv()).await {} +} + +#[test] +fn with_too_few_fds_it_falls_back_to_fsevents_and_misses_nothing() { + // 名额 32 + set_soft_limit(64); + let rt = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + rt.block_on(async { + let root = tempfile::tempdir().unwrap(); + let before = open_fds(); + + // 8 个目录、16 份文件，名额够：都用 kqueue 盯，用完都还回去 + let some = dirs(root.path(), "some", 8, 2); + let (w, mut rx) = tw_watch::watch(&some, WINDOW, only_md).unwrap(); + assert_eq!( + open_fds(), + before + 1 + 8 + 16, + "kqueue、8 个目录、16 份文件" + ); + append(&some[3].join("1.md"), "1\n"); + assert!(got(&mut rx).await, "原地写"); + settle(&mut rx).await; + + // 又多出 10 份：超了名额，换回 FSEvents，手上的都放掉（只剩 kqueue 自己） + for j in 2..12 { + std::fs::write(some[0].join(format!("{j}.md")), "0\n").unwrap(); + } + assert!(got(&mut rx).await, "超名额的那一批新文件"); + settle(&mut rx).await; + assert_eq!(open_fds(), before + 1, "换回 FSEvents 之后还拿着 fd"); + // 换过去之后照样一个不漏：新文件的原地写、rename 保存 + append(&some[0].join("7.md"), "1\n"); + assert!(got(&mut rx).await, "换回 FSEvents 之后的原地写"); + settle(&mut rx).await; + std::fs::write(some[5].join("x.tmp"), "2\n").unwrap(); + std::fs::rename(some[5].join("x.tmp"), some[5].join("0.md")).unwrap(); + assert!(got(&mut rx).await, "换回 FSEvents 之后的 rename 保存"); + drop(w); + assert_eq!(open_fds(), before, "扔掉之后还拿着 fd"); + + // 一开始就不够：40 个目录，直接用 FSEvents + let many = dirs(root.path(), "many", 40, 1); + let (w, mut rx) = tw_watch::watch(&many, WINDOW, only_md).unwrap(); + assert_eq!(open_fds(), before, "名额不够还占着 fd"); + settle(&mut rx).await; + append(&many[39].join("0.md"), "1\n"); + assert!(got(&mut rx).await, "FSEvents 盯着的原地写"); + drop(w); + + // 名额还回来了：又能用 kqueue 盯 + let (_w, _rx) = tw_watch::watch(&some[..2], WINDOW, only_md).unwrap(); + assert!(open_fds() > before + 2, "名额没还回来"); + }); +} From 99b74a4e27deaf753627c4c3fbdd78d0c1872421 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:42:32 +0800 Subject: [PATCH 07/16] perf(transcript): read only new turns, and let the app ask for them Every GET /sessions/{id}/transcript re-read and re-parsed every stored body of the session, and the app asks again on every new turn: for a 200-turn session of ~1.5 MB requests that is ~300 MB per turn. Measured with `transcript_cost` (300 turns, 367 MB of bodies, release build): the whole transcript took ~830 ms; with one new turn it now takes 6.4 ms, and asking from `settled_turns` takes 0.07 ms and sends 27 KB instead of 1.5 MB. tw-store: transcript::Cache (one per Recorder, Recorder::transcripts) keeps, for the last four sessions read, the request ids read, the chain state after them and the unmasked turns, and continues from there when those ids are still the session's first rows. A middle insertion or a pruned head fails that check and reads from the start. Only the settled prefix is kept: a turn that ended under two minutes ago with a gap may still get its body, so it and everything after it are read again. GC invalidates the cache whenever it deleted a body (a generation counter keeps a read that raced the GC from being stored). tw-api, protocol addition (no version bump here): - SessionTranscript takes TranscriptQuery { from_turn?: u32 }. Present: `turns` holds only the turns with index >= from_turn (empty past the end); absent or 0: the whole transcript, as before. - Transcript gains `total_turns` (number of turns in the session) and `settled_turns` (leading turns that will not change any more: bodies final, not rewritten by a later turn's tool-call ids, and no running request of the session will land before them). The app keeps its first `settled_turns` turns and asks again with that as from_turn. Co-Authored-By: Claude Opus 5.5 --- crates/tw-api/src/ep.rs | 5 +- crates/tw-api/src/lib.rs | 26 +- crates/tw-api/src/ts.rs | 9 +- crates/tw-control/src/lib.rs | 57 +++- crates/tw-control/tests/transcript.rs | 43 +++ crates/tw-store/src/recorder.rs | 10 + crates/tw-store/src/task.rs | 72 +++- crates/tw-store/src/transcript/mod.rs | 473 +++++++++++++++++++++++++- 8 files changed, 660 insertions(+), 35 deletions(-) diff --git a/crates/tw-api/src/ep.rs b/crates/tw-api/src/ep.rs index c4d6c3d6..aaf98bfa 100644 --- a/crates/tw-api/src/ep.rs +++ b/crates/tw-api/src/ep.rs @@ -64,8 +64,9 @@ endpoints! { Fixture: GET "/request/{id}/fixture" [id], () => String, text; Sessions: GET "/sessions", api::ListQuery => Vec; SessionDetail: GET "/sessions/{id}" [id], () => api::SessionDetail; - /// 一次会话读成一段对话：每一轮新说的话、回答、工具调用和结果（已脱敏） - SessionTranscript: GET "/sessions/{id}/transcript" [id], () => api::Transcript; + /// 一次会话读成一段对话：每一轮新说的话、回答、工具调用和结果（已脱敏）。可以只要 + /// 从某一轮起的那些 + SessionTranscript: GET "/sessions/{id}/transcript" [id], api::TranscriptQuery => api::Transcript; // ─────────────────────────────────────────────── 测速、回放、试路由 SpeedQuote: POST "/speed/quote", api::SpeedRunRequest => api::SpeedQuote; diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index 122ae583..33fe5624 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -4697,6 +4697,9 @@ pub struct SessionDetail { /// 太大的只留开头，没存下来的也有。读不到的地方，那一轮的 `gaps` 说出来。 /// /// **已脱敏**，和请求详情里的正文同一套打码。图片只说类型和大小，从不带数据。 +/// +/// **可以只要后面的几轮**（[`TranscriptQuery::from_turn`]）：会话开着时每来一轮重读一次， +/// 前面不会再变的那些（`settled_turns`）不必再传一遍。 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] pub struct Transcript { @@ -4705,10 +4708,31 @@ pub struct Transcript { /// Gemini 的 `systemInstruction`，Chat 和 Responses 还有开头连着的 system、developer /// 消息，几段之间空一行。没有是 null pub system: Option, - /// 和 [`SessionDetail::turns`] 同样的请求，同样的顺序 + /// 这次会话一共几轮，就是 [`SessionDetail::turns`] 有几条。给了 + /// [`TranscriptQuery::from_turn`] 时，`turns` 是其中从那一轮起的那些 + pub total_turns: u32, + /// 开头这么多轮不会再变了：它们的正文都落了盘（或者不会再来），后面的请求不会再改写 + /// 它们，这次会话里还在跑的请求落库时也不会排到它们前面。**下一次从这里要起** + /// （`from_turn = settled_turns`），留着手上的前这么多轮，换掉后面的。 + /// + /// 后面那几轮还会变：最后一轮的回答可能还没落盘、工具调用的号会被下一轮换成客户端 + /// 记下的那个、一个开始得早结束得晚的请求会插到它们中间 + pub settled_turns: u32, + /// 和 [`SessionDetail::turns`] 同样的请求，同样的顺序。给了 + /// [`TranscriptQuery::from_turn`] 的，只有从那一轮起的那些 pub turns: Vec, } +/// `GET /sessions/{id}/transcript` 的查询串。 +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct TranscriptQuery { + /// 只要从这一轮起的那些（从 0 数，[`Transcript::turns`] 里的第几条）。不给是整段； + /// 比总轮数还大的，`turns` 是空的 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub from_turn: Option, +} + /// 对话里的一轮，就是会话里的一个请求。 /// /// 客户端每一轮都把整段历史发上来：请求 i 的消息 = 请求 i-1 的消息 + 上一轮的回答 + 新的 diff --git a/crates/tw-api/src/ts.rs b/crates/tw-api/src/ts.rs index 6ba442ac..4843bf7f 100644 --- a/crates/tw-api/src/ts.rs +++ b/crates/tw-api/src/ts.rs @@ -301,7 +301,12 @@ mod tests { #[test] fn a_transcript_is_turns_of_messages_and_parts() { let ts = typescript(); - assert!(ts.contains(" SessionTranscript: { req: null; res: Transcript };")); + assert!(ts.contains(" SessionTranscript: { req: TranscriptQuery; res: Transcript };")); + assert!( + decl_of(&ts, "TranscriptQuery").contains("from_turn?: number"), + "{}", + decl_of(&ts, "TranscriptQuery") + ); assert!(ts.contains( " SessionTranscript: { method: \"GET\", path: \"/sessions/{id}/transcript\", params: [\"id\"], format: \"json\" }," )); @@ -309,6 +314,8 @@ mod tests { for field in [ "session: string", "system: string | null", + "total_turns: number", + "settled_turns: number", "turns: Array", ] { assert!(transcript.contains(field), "{field}: {transcript}"); diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index c552f737..59af3e47 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -1720,28 +1720,53 @@ fn no_such_session(id: &str) -> Fail { ) } -/// 一次会话读成一段对话（见 `tw_store::transcript`）。 +/// 一次会话读成一段对话（见 `tw_store::transcript`）。给了 `from_turn` 的只交出从那一轮起的 +/// 那些（见 [`tw_api::Transcript::settled_turns`]）。 /// /// **放到阻塞线程上跑**，库只在取行的时候锁一下：几百轮的会话要读几百份正文、解析几百 MB -/// 的 JSON，和按正文找是同一个道理（见 [`history_search`]）—— 记录和正文落盘走的是同一把锁。 +/// 的 JSON，和按正文找是同一个道理（见 [`history_search`]）—— 记录落库走的是同一把锁。读过 +/// 的记着（`Recorder::transcripts`），下一次只读新来的那几轮。 async fn session_transcript( State(s): State, axum::extract::Path(id): axum::extract::Path, + axum::extract::Query(q): axum::extract::Query, ) -> Result, Fail> { - let store = need_store(&s)?; - let (rows, blobs) = { - let g = store.lock().await; - let rows = g.db().session_requests(&id).map_err(records)?; - // 正文目录只是一个路径，读它不需要锁 - (rows, tw_store::Blobs::new(g.blobs().root().to_path_buf())) - }; - if rows.is_empty() { - return Err(no_such_session(&id)); - } - let transcript = - tokio::task::spawn_blocking(move || tw_store::transcript::build(&id, &rows, &blobs)) - .await - .map_err(internal)?; + let store = need_store(&s)?.clone(); + // 这次会话里还在跑的请求：落库时排在它开始的那一刻，不一定排在最后 + let running = s + .gateway + .bus + .live() + .running + .into_iter() + .filter(|r| r.session.as_deref() == Some(id.as_str())) + .map(|r| (r.at_ms as i64, r.id as i64)) + .min(); + let from_turn = q.from_turn.unwrap_or(0) as usize; + let transcript = tokio::task::spawn_blocking(move || { + let (rows, blobs, transcripts) = { + let g = store.blocking_lock(); + let rows = g.db().session_requests(&id).map_err(records)?; + // 正文目录只是一个路径，读它不需要锁 + let blobs = tw_store::Blobs::new(g.blobs().root().to_path_buf()); + (rows, blobs, g.transcripts()) + }; + if rows.is_empty() { + return Err(no_such_session(&id)); + } + Ok(transcripts.read( + &blobs, + &tw_store::transcript::Ask { + session: &id, + rows: &rows, + from_turn, + running, + now_ms: now_ms(), + }, + )) + }) + .await + .map_err(internal)??; Ok(Json(transcript)) } diff --git a/crates/tw-control/tests/transcript.rs b/crates/tw-control/tests/transcript.rs index 8a11dde4..baeaec18 100644 --- a/crates/tw-control/tests/transcript.rs +++ b/crates/tw-control/tests/transcript.rs @@ -148,6 +148,9 @@ async fn the_transcript_is_the_shape_the_desktop_app_reads() { json!({ "session": "s1", "system": "sys", + "total_turns": 2, + // 第二轮没有回答（不会再被改写），两轮都早就结束了 + "settled_turns": 2, "turns": [ { "id": "1", @@ -220,3 +223,43 @@ async fn a_session_whose_bodies_are_gone_still_lists_its_turns() { assert_eq!(x["input"], json!([])); } } + +/// `from_turn`：只要从那一轮起的那些，形状不变，`total_turns` 说一共几轮。比总轮数还大的 +/// 是空的。不给和给 0 一样，是整段 +#[tokio::test] +async fn from_turn_gives_the_turns_from_there_on() { + let first = + json!({"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "一"}]}); + let answer = json!({"type": "message", "content": [{"type": "text", "text": "回一"}]}); + let second = json!({"model": "claude-sonnet-4-5", "messages": [ + {"role": "user", "content": "一"}, + {"role": "assistant", "content": "回一"}, + {"role": "user", "content": "二"} + ]}); + let (_d, app) = app( + &[turn(1), turn(2)], + &[(1, Some(first), Some(answer)), (2, Some(second), None)], + ); + + let (st, all) = get(&app, "/sessions/s1/transcript").await; + assert_eq!(st, StatusCode::OK, "{all}"); + let (_, zero) = get(&app, "/sessions/s1/transcript?from_turn=0").await; + assert_eq!(zero, all); + + let (st, tail) = get(&app, "/sessions/s1/transcript?from_turn=1").await; + assert_eq!(st, StatusCode::OK, "{tail}"); + assert_eq!(tail["total_turns"], 2); + assert_eq!(tail["settled_turns"], all["settled_turns"]); + assert_eq!(tail["system"], all["system"]); + assert_eq!(tail["turns"], json!([all["turns"][1]])); + assert_eq!(tail["turns"][0]["id"], "2"); + + let (st, none) = get(&app, "/sessions/s1/transcript?from_turn=9").await; + assert_eq!(st, StatusCode::OK); + assert_eq!(none["total_turns"], 2); + assert_eq!(none["turns"], json!([])); + + // 不是数的 `from_turn` 是请求写错了 + let (st, body) = get(&app, "/sessions/s1/transcript?from_turn=x").await; + assert_eq!(st, StatusCode::BAD_REQUEST, "{body}"); +} diff --git a/crates/tw-store/src/recorder.rs b/crates/tw-store/src/recorder.rs index fa180f53..bf077d62 100644 --- a/crates/tw-store/src/recorder.rs +++ b/crates/tw-store/src/recorder.rs @@ -93,6 +93,9 @@ pub struct Recorder { /// /// `None` 表示没人要听（测试、以及不带总线的调用方）。 bus: Option, + /// 最近几次会话读成对话读到了哪儿（见 [`crate::transcript::Cache`]）。正文回收删了东西 + /// 就作废（[`crate::task::gc`]） + transcripts: std::sync::Arc, /// 每记下一行请求就交一份 [`Settled`] 出去：密钥的用量上限拿它把预留换成实数。 /// /// **直接调，不走总线。**总线上丢了事件的话，这一行也就不在库里，重启之后从库里 @@ -153,6 +156,7 @@ impl Recorder { pricing, inflight: HashMap::new(), bus: None, + transcripts: Default::default(), settled: None, } } @@ -230,6 +234,12 @@ impl Recorder { &self.blobs } + /// 读对话记录用的那一份记忆（见 [`crate::transcript::Cache`]）。读的时候不拿这把锁： + /// 几百轮的会话要读几百份正文 + pub fn transcripts(&self) -> std::sync::Arc { + self.transcripts.clone() + } + /// 记一次请求体或响应体。 /// /// **在事件之外单独走** —— body 不进事件流：那是个广播通道，每个 diff --git a/crates/tw-store/src/task.rs b/crates/tw-store/src/task.rs index 6ffb92e6..e59fbd96 100644 --- a/crates/tw-store/src/task.rs +++ b/crates/tw-store/src/task.rs @@ -109,8 +109,15 @@ pub fn gc( metadata_keep_days: u64, max_bytes: u64, ) -> u64 { - let blobs = Blobs::new(shared.blocking_lock().blobs().root().to_path_buf()); + let (blobs, transcripts) = { + let g = shared.blocking_lock(); + (Blobs::new(g.blobs().root().to_path_buf()), g.transcripts()) + }; let freed = blobs.gc(now_ms, keep_days, max_bytes); + // 删掉的正文可能在记着的对话里：记着的那些作废，下次重读 + if freed > 0 { + transcripts.invalidate(); + } shared.blocking_lock().prune(now_ms, metadata_keep_days); freed } @@ -121,3 +128,66 @@ fn thread(name: &str, f: impl FnOnce() + Send + 'static) { tracing::warn!("the {name} thread could not start, so it records nothing this run: {e}"); } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::transcript::Ask; + + const NOW: i64 = 1_790_000_000_000; + const DAY: i64 = 86_400_000; + + /// 回收删了正文，记着的对话跟着作废：之后读出来的那一轮说出缺了什么，而不是删掉之前的 + /// 样子 + #[test] + fn reclaiming_bodies_forgets_the_transcripts_read_from_them() { + let d = tempfile::tempdir().unwrap(); + let rec = Recorder::new( + crate::Db::in_memory().unwrap(), + Blobs::new(d.path().join("blobs")), + tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), + ); + let mut row = crate::db::tests::row(1, NOW); + row.session = Some("s".into()); + rec.db().insert(&row).unwrap(); + let request = + br#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hi"}]}"#; + let answer = br#"{"type":"message","content":[{"type":"text","text":"hello"}]}"#; + rec.record_body(NOW, 1, Which::Request, request, request.len()); + rec.record_body(NOW, 1, Which::Response, answer, answer.len()); + let shared = Mutex::new(rec); + let read = || { + let (rows, blobs, transcripts) = { + let g = shared.blocking_lock(); + ( + g.db().session_requests("s").unwrap(), + Blobs::new(g.blobs().root().to_path_buf()), + g.transcripts(), + ) + }; + transcripts.read( + &blobs, + &Ask { + session: "s", + rows: &rows, + from_turn: 0, + running: None, + now_ms: NOW + 3 * DAY, + }, + ) + }; + let first = read(); + assert!(first.turns[0].gaps.is_empty(), "{first:?}"); + assert!(!first.turns[0].output.is_empty()); + + assert!(gc(&shared, NOW + 3 * DAY, 1, 90, u64::MAX) > 0); + let again = read(); + assert_eq!( + again.turns[0].gaps, + [ + tw_api::TranscriptGap::RequestMissing, + tw_api::TranscriptGap::ResponseMissing + ] + ); + } +} diff --git a/crates/tw-store/src/transcript/mod.rs b/crates/tw-store/src/transcript/mod.rs index d5a8e676..0ad55314 100644 --- a/crates/tw-store/src/transcript/mod.rs +++ b/crates/tw-store/src/transcript/mod.rs @@ -18,11 +18,16 @@ //! //! 几百轮的会话，每个请求几 MB。每份正文读一次、解析一次，读完这一轮就丢；从一个请求留到 //! 下一个的只有每条消息 16 字节的指纹和系统提示。 +//! +//! **读过的不再读**（[`Cache`]）：界面开着一次会话时，每来一轮就要一次，从头读的话，两百轮 +//! 的会话每来一轮就要读、解析几百 MB。记着最近几次会话读到了哪儿（读完时的 [`Chain`] 和读出 +//! 来的几轮），下一次只读新来的请求。只记定下来的那一截：正文还可能变的那几轮每次都重读。 mod answer; mod read; use std::collections::HashSet; +use std::sync::Arc; use serde_json::Value; use tw_api::{TranscriptGap, TranscriptMessage, TranscriptPart, TranscriptRole, TranscriptTurn}; @@ -33,29 +38,207 @@ use crate::db::RequestRow; use crate::search::text::{client_dialect, dialect_of}; use read::{Input, Piece, Role}; -/// 读成一段对话。`rows` 是这次会话的请求，和会话详情同样的顺序 +/// 一轮结束这么久之后就定下来了：它的正文要么已经在盘上，要么不会再来。正文和那一行各走 +/// 各的路落盘（见 `crate::task`），先后差不了几毫秒，排队排得再久也到不了这么久 +const SETTLE_MS: i64 = 120_000; + +/// 记着最近几次会话读到了哪儿。一次几百轮的会话记下的是几 MB 的文字 +const CACHED: usize = 4; + +/// 整段读一遍，不记。`rows` 是这次会话的请求，和会话详情同样的顺序 /// （[`crate::Db::session_requests`]）。 /// /// **已脱敏**，和请求详情里的正文同一套打码（`tw_secret::mask_body`）。 pub fn build(session: &str, rows: &[RequestRow], blobs: &Blobs) -> tw_api::Transcript { - let mut t = tw_api::Transcript { - session: session.to_string(), - system: None, - turns: Vec::with_capacity(rows.len()), - }; - let mut chain = Chain::default(); - for row in rows { - chain.turn(row, blobs, &mut t); + Cache::default().read( + blobs, + &Ask { + session, + rows, + from_turn: 0, + running: None, + now_ms: i64::MAX, + }, + ) +} + +/// 要读哪一次会话、交出去哪几轮（见 [`Cache::read`]）。 +pub struct Ask<'a> { + pub session: &'a str, + /// 这次会话的请求，和会话详情同样的顺序（[`crate::Db::session_requests`]） + pub rows: &'a [RequestRow], + /// 只交出从这一轮起的那些（从 0 数）。0 是整段 + pub from_turn: usize, + /// 这次会话里还在跑的请求里最早开始的那个：（开始的时刻，请求号）。它落库时排在开始的 + /// 那一刻，不一定排在最后：排在它后面的几轮还没定下来（见 + /// [`tw_api::Transcript::settled_turns`]） + pub running: Option<(i64, i64)>, + /// 此刻，Unix 毫秒 + pub now_ms: i64, +} + +/// 最近几次会话读到了哪儿（见 [`Cache::read`]）。存储层一份（[`crate::Recorder::transcripts`]）。 +#[derive(Default)] +pub struct Cache { + kept: std::sync::Mutex, +} + +#[derive(Default)] +struct Kept { + /// 正文回收每删掉一次东西就加一。读到一半它变了的，读出来的就不记 + generation: u64, + /// 最近用过的在最后 + sessions: Vec<(String, Arc)>, +} + +/// 一次会话读到了哪儿：读过的那几个请求，读完它们时的 [`Chain`]，读出来的几轮（没打码）。 +#[derive(Clone, Default)] +struct Progress { + ids: Vec, + chain: Chain, + doc: Doc, +} + +/// 读出来的东西：系统提示和每一轮。**没打码**：前后比对、换工具调用的号都要原文 +#[derive(Clone, Default)] +struct Doc { + system: Option, + turns: Vec, +} + +impl Cache { + /// 正文回收删掉了东西：记着的都不作数了 + pub fn invalidate(&self) { + let mut k = self.kept(); + k.generation += 1; + k.sessions.clear(); + } + + /// 锁中毒了照样用里面的：记着的东西坏不了谁，最多重读一遍 + fn kept(&self) -> std::sync::MutexGuard<'_, Kept> { + self.kept.lock().unwrap_or_else(|p| p.into_inner()) + } + + /// 读成一段对话。**已脱敏**，和请求详情里的正文同一套打码（`tw_secret::mask_body`）。 + /// + /// **接着上次读到的地方往下读**：记着的那几个请求还是这次会话开头的那几个，就只读后面 + /// 新来的；中间插进来一个（开始得早、结束得晚的请求）、开头的过期删掉了，都从头读。 + /// 记下的是定下来的那一截：最后几轮的正文可能还在路上，它们下一次还要重读。 + pub fn read(&self, blobs: &Blobs, ask: &Ask) -> tw_api::Transcript { + let (generation, kept) = { + let mut k = self.kept(); + let kept = k + .sessions + .iter() + .position(|(s, _)| s == ask.session) + .map(|i| { + let hit = k.sessions.remove(i); + let p = hit.1.clone(); + k.sessions.push(hit); + p + }); + (k.generation, kept) + }; + let base = kept + .filter(|p| { + p.ids.len() <= ask.rows.len() + && p.ids.iter().zip(ask.rows).all(|(id, row)| *id == row.id) + }) + .unwrap_or_default(); + let fresh = &ask.rows[base.ids.len()..]; + // 定下来的那一截有多长。记着的那一截本来就是定下来的 + let (whole, settled) = if fresh.is_empty() { + let n = base.ids.len(); + (base, n) + } else { + let (whole, keep) = continue_reading(&base, fresh, blobs, ask.now_ms); + let n = keep.ids.len(); + self.keep(ask.session, generation, keep); + (whole, n) + }; + + let total = whole.doc.turns.len(); + // 后面的请求还会改写的那一轮：最后一个读得懂的请求的回答（见 `same_calls`） + let rewritable = match whole.chain.last { + Some((at, true)) => at, + _ => total, + }; + let before_running = ask.running.map_or(total, |r| { + ask.rows.partition_point(|row| (row.at_ms, row.id) < r) + }); + let from = ask.from_turn.min(total); + let mut t = tw_api::Transcript { + session: ask.session.to_string(), + system: whole.doc.system.clone(), + total_turns: total as u32, + settled_turns: settled.min(rewritable).min(before_running) as u32, + turns: whole.doc.turns[from..].to_vec(), + }; + mask(&mut t); + t + } + + /// 记下一次会话读到的地方。正文回收在这期间删过东西的不记：读的时候它们可能还在 + fn keep(&self, session: &str, generation: u64, p: Arc) { + let mut k = self.kept(); + if k.generation != generation { + return; + } + k.sessions.retain(|(s, _)| s != session); + k.sessions.push((session.to_string(), p)); + if k.sessions.len() > CACHED { + k.sessions.remove(0); + } } - mask(&mut t); - t +} + +/// 从 `base` 读到的地方接着读 `fresh`。返回读完的样子，和其中定下来的那一截（要记下的）。 +/// +/// 一轮定下来了：它结束得够久（[`SETTLE_MS`]），或者它什么都不缺。还缺着正文、又刚结束的 +/// 那一轮，正文可能还在路上；从它起都不算定下来。 +fn continue_reading( + base: &Progress, + fresh: &[RequestRow], + blobs: &Blobs, + now_ms: i64, +) -> (Arc, Arc) { + let mut p = base.clone(); + let mut settled: Option = None; + for row in fresh { + let ended = row.at_ms.saturating_add(row.duration_ms.unwrap_or(0)); + let recent = now_ms.saturating_sub(ended) < SETTLE_MS; + // 这一个可能定不下来：先留着读它之前的样子。读它只会改它自己和上一个读得懂的那一轮 + let before = (settled.is_none() && recent).then(|| { + let last = p.chain.last.map(|(at, _)| (at, p.doc.turns[at].clone())); + (p.chain.clone(), p.doc.system.clone(), last) + }); + p.chain.turn(row, blobs, &mut p.doc); + p.ids.push(row.id); + if let Some((chain, system, last)) = before + && p.doc.turns.last().is_some_and(|t| !t.gaps.is_empty()) + { + let n = p.ids.len() - 1; + let mut turns = p.doc.turns[..n].to_vec(); + if let Some((at, turn)) = last { + turns[at] = turn; + } + settled = Some(Progress { + ids: p.ids[..n].to_vec(), + chain, + doc: Doc { system, turns }, + }); + } + } + let p = Arc::new(p); + let settled = settled.map_or_else(|| p.clone(), Arc::new); + (p, settled) } /// 一条消息比对用的指纹 type Fp = [u8; 16]; /// 读到哪儿了：上一个读得懂的请求留下来的东西。 -#[derive(Default)] +#[derive(Clone, Default)] struct Chain { /// 它的每条消息的指纹（只有推理的消息不比，不在里面）。还没有读得懂的请求是 None prev: Option>, @@ -68,7 +251,7 @@ struct Chain { } impl Chain { - fn turn(&mut self, row: &RequestRow, blobs: &Blobs, t: &mut tw_api::Transcript) { + fn turn(&mut self, row: &RequestRow, blobs: &Blobs, t: &mut Doc) { let mut turn = TranscriptTurn { id: row.id.to_string(), restart: false, @@ -119,7 +302,7 @@ impl Chain { client: Dialect, v: &Value, turn: &mut TranscriptTurn, - t: &mut tw_api::Transcript, + t: &mut Doc, freeform: &mut HashSet, ) { let read::Body { @@ -1844,6 +2027,239 @@ mod tests { assert!(t.turns[0].output.is_empty() && t.turns[0].gaps.is_empty()); } + // ───────────────────────────────────────────────── 读过的不再读 + + /// 很久以后的此刻：每一轮都早就结束了 + const LATER: i64 = NOW + 86_400_000; + + impl Disk { + /// 用 `cache` 读：从 `from_turn` 起，此刻是 `now_ms` + fn read(&self, cache: &Cache, from_turn: usize, now_ms: i64) -> Transcript { + self.read_running(cache, from_turn, None, now_ms) + } + + fn read_running( + &self, + cache: &Cache, + from_turn: usize, + running: Option<(i64, i64)>, + now_ms: i64, + ) -> Transcript { + let rows = self.db.session_requests("s").unwrap(); + cache.read( + &self.blobs, + &Ask { + session: "s", + rows: &rows, + from_turn, + running, + now_ms, + }, + ) + } + + /// 第 `id` 轮结束之后过了 `ms` 毫秒（`row` 给的每一轮都跑了 4 秒） + fn after(&self, id: i64, ms: i64) -> i64 { + NOW + id * 1000 + 4000 + ms + } + } + + /// 一段带工具调用的对话：第 `k` 轮的请求和回答。回答里工具调用的号和客户端记下的不一样 + /// （转换过格式的上游就是这样）：读下一轮时，上一轮回答里的号换成客户端的那个 + fn claude_turn(k: usize) -> (Value, Vec) { + let mut messages = vec![user("把 bug 修了")]; + for i in 0..k { + messages.push(json!({"role": "assistant", "content": [ + {"type": "text", "text": format!("读 {i}")}, + {"type": "tool_use", "id": format!("t{i}"), "name": "Read", "input": {"n": i}}]})); + messages.push(json!({"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": format!("t{i}"), "content": format!("文件 {i}")}]})); + } + let answer = anthropic_stream(&[ + json!({"type": "text", "text": format!("读 {k}")}), + json!({"type": "tool_use", "id": format!("up{k}"), "name": "Read", "input": {"n": k}}), + ]); + (anthropic(&messages), answer) + } + + /// **接着读和从头读是同一段对话**：每来一轮读一次，每一次都和整段从头读的一样；只要 + /// 后面几轮的，就是整段的那一截。中间夹着一个数 token 的调用 + #[test] + fn reading_on_from_where_it_stopped_gives_what_reading_it_all_gives() { + let mut d = Disk::new(); + let cache = Cache::default(); + for k in 0..8 { + let (request, answer) = claude_turn(k); + d.turn("/v1/messages", &request, &answer); + if k == 3 { + d.turn("/v1/messages/count_tokens", &request, b"{}"); + } + let full = d.transcript(); + let got = d.read(&cache, 0, LATER); + assert_eq!(got.turns, full.turns, "第 {k} 轮"); + assert_eq!(got.system, full.system); + assert_eq!(got.total_turns as usize, full.turns.len()); + assert!(got.turns.iter().skip(1).all(|t| !t.restart), "{got:?}"); + // 上一轮回答里的号换成了客户端记下的，最后一轮的还是上游给的 + if k > 0 { + let ids = |t: &TranscriptTurn| -> Vec { + t.output + .iter() + .filter_map(|p| match p { + P::ToolCall { id, .. } => Some(id.clone()), + _ => None, + }) + .collect() + }; + let answered: Vec> = got + .turns + .iter() + .map(ids) + .filter(|x| !x.is_empty()) + .collect(); + assert_eq!(answered[k - 1], [format!("t{}", k - 1)]); + assert_eq!(answered[k], [format!("up{k}")]); + } + for from in [0, 1, full.turns.len() - 1, full.turns.len(), 99] { + let part = d.read(&cache, from, LATER); + assert_eq!( + part.turns, + full.turns[from.min(full.turns.len())..], + "{from}" + ); + assert_eq!(part.total_turns, got.total_turns); + } + } + } + + /// **读过的不再读，正文回收之后才重读。**读过一遍之后正文没了（这里直接删掉），记着的 + /// 照样交出去 —— 说明它真的没再去读；作废之后重读，那几轮说出缺了什么 + #[test] + fn what_was_read_is_not_read_again_until_bodies_are_reclaimed() { + let mut d = Disk::new(); + let cache = Cache::default(); + for k in 0..2 { + let (request, answer) = claude_turn(k); + d.turn("/v1/messages", &request, &answer); + } + let first = d.read(&cache, 0, LATER); + assert!(first.turns.iter().all(|t| t.gaps.is_empty()), "{first:?}"); + // 整天的正文都删掉：两天之后、只留一天 + d.blobs.gc(NOW + 2 * 86_400_000, 1, u64::MAX); + assert!(d.blobs.get(NOW + 1000, 1, Which::Request).is_none()); + assert_eq!(d.read(&cache, 0, LATER), first, "记着的没用上"); + + cache.invalidate(); + let again = d.read(&cache, 0, LATER); + for t in &again.turns { + assert_eq!( + t.gaps, + [Gap::RequestMissing, Gap::ResponseMissing], + "{again:?}" + ); + } + } + + /// 刚结束、回答还没落盘的那一轮**不记下**：回答一落盘，下一次就读得到，后面那一轮跟着 + /// 对得上（上一轮的回答交出去了，它的输入里不再重复那条助手消息）。结束得够久还缺着的， + /// 就是真的缺了 + #[test] + fn a_turn_whose_answer_is_still_on_its_way_is_read_again() { + let mut d = Disk::new(); + let cache = Cache::default(); + let (r0, a0) = claude_turn(0); + d.turn("/v1/messages", &r0, &a0); + let (r1, a1) = claude_turn(1); + let id = d.put( + "/v1/messages", + Some(r1.to_string().as_bytes()), + None, + |_| {}, + ); + let now = d.after(id, 500); + + let t = d.read(&cache, 0, now); + assert_eq!(t.turns[1].gaps, [Gap::ResponseMissing]); + assert_eq!(t.settled_turns, 1, "回答还在路上的那一轮算定下来了"); + // 下一轮来了，回答还没到 + let (r2, a2) = claude_turn(2); + d.turn("/v1/messages", &r2, &a2); + let t = d.read(&cache, 0, d.after(id + 1, 500)); + assert_eq!(t.turns[1].gaps, [Gap::ResponseMissing]); + assert_eq!(t.settled_turns, 1); + + // 回答落盘了：重读那一轮，和从头读的一样 + assert!(d.blobs.put(NOW + id * 1000, id, Which::Response, &a1)); + let t = d.read(&cache, 0, d.after(id + 1, 600)); + assert_eq!(t.turns, d.transcript().turns); + assert!(t.turns.iter().all(|x| x.gaps.is_empty()), "{t:?}"); + assert_eq!(t.total_turns, 3); + // 最后一轮的回答下一轮还会改写（工具调用的号），它还不算定下来 + assert_eq!(t.settled_turns, 2); + + // 回答一直没来：结束得够久之后它就是缺了，算定下来 + let (r3, _) = claude_turn(3); + let late = d.put( + "/v1/messages", + Some(r3.to_string().as_bytes()), + None, + |_| {}, + ); + let t = d.read(&cache, 0, d.after(late, 1000)); + assert_eq!(t.settled_turns, 3); + let t = d.read(&cache, 0, d.after(late, SETTLE_MS)); + assert_eq!(t.turns[3].gaps, [Gap::ResponseMissing]); + assert_eq!(t.settled_turns, 4, "缺着回答的最后一轮不会再被改写"); + } + + /// 定下来的几轮也要排在还在跑的请求前面：一个开始得早、结束得晚的请求落库时插在它开始的 + /// 那一刻，它后面的几轮都要往后挪一格 + #[test] + fn turns_after_a_request_still_running_are_not_settled() { + let mut d = Disk::new(); + let cache = Cache::default(); + for k in 0..3 { + let (request, answer) = claude_turn(k); + d.turn("/v1/messages", &request, &answer); + } + let t = d.read(&cache, 0, LATER); + assert_eq!((t.total_turns, t.settled_turns), (3, 2)); + // 第二轮之前开始的一个请求还在跑 + let t = d.read_running(&cache, 0, Some((NOW + 2 * 1000 - 1, 99)), LATER); + assert_eq!(t.settled_turns, 1); + // 在最后一轮之后开始的不碍事 + let t = d.read_running(&cache, 0, Some((NOW + 10_000, 99)), LATER); + assert_eq!(t.settled_turns, 2); + } + + /// 一个请求插进了读过的几轮中间（开始得早、结束得晚）：记着的那一截对不上了，从头读 + #[test] + fn a_request_landing_between_turns_already_read_is_read_from_the_start() { + let mut d = Disk::new(); + let cache = Cache::default(); + let (r0, a0) = claude_turn(0); + let (r1, a1) = claude_turn(1); + let (r2, a2) = claude_turn(2); + d.turn("/v1/messages", &r0, &a0); + d.turn("/v1/messages", &r2, &a2); + d.read(&cache, 0, LATER); + // 第二轮开始在前两轮之间，结束在它们之后 + d.put( + "/v1/messages", + Some(r1.to_string().as_bytes()), + Some(&a1), + |r| r.at_ms = NOW + 1500, + ); + let got = d.read(&cache, 0, LATER); + let full = d.transcript(); + assert_eq!(got.turns, full.turns); + assert_eq!( + got.turns.iter().map(|t| t.id.as_str()).collect::>(), + ["1", "3", "2"] + ); + assert!(got.turns.iter().all(|t| !t.restart), "{got:?}"); + } + // ───────────────────────────────────────────────── 量一量 /// 读一次几百轮的会话要多久：照 Claude Code 的样子造一段 300 轮的对话，每一轮的请求带着 @@ -1928,5 +2344,34 @@ mod tests { .all(|x| x.gaps.is_empty() && x.output.len() == 3) ); } + + // 界面开着这次会话：读过前 299 轮，第 300 轮来了 + let cache = Cache::default(); + let read = |rows: &[RequestRow], from_turn| { + let t0 = std::time::Instant::now(); + let t = cache.read( + &d.blobs, + &Ask { + session: "s", + rows, + from_turn, + running: None, + now_ms: i64::MAX, + }, + ); + (t0.elapsed(), t) + }; + let (took, before) = read(&rows[..rows.len() - 1], 0); + eprintln!("first read of {} turns: {took:?}", before.total_turns); + let (took, whole) = read(&rows, 0); + eprintln!("one more turn, whole transcript again: {took:?}"); + assert_eq!(whole.turns, build("s", &rows, &d.blobs).turns); + let (took, tail) = read(&rows, before.settled_turns as usize); + eprintln!( + "one more turn, from settled_turns ({}): {took:?}, {} turns, {} KB of JSON", + before.settled_turns, + tail.turns.len(), + serde_json::to_string(&tail).unwrap().len() / 1000 + ); } } From fbf9a73463e0983d07e1e35bde43447febb10145 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:55:47 +0800 Subject: [PATCH 08/16] perf: find markers in request bodies with memchr The local-answer check looks for five markers in every request body before deciding whether to parse it, and it did so with `windows().any()`: about 16 ms of synchronous CPU for a 2 MiB Claude Code request, on the async worker, before anything else happened. `strip_carried` did the same on every same-format hop, and the usage sniffer on every streamed chunk. memchr is already in the tree through serde_json and regex, so adding it as a workspace dependency compiles nothing new. memmem gives the same first match: 0.26 ms for the five markers over 2 MiB, 0.04 ms for `strip_carried`'s. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 2 ++ Cargo.toml | 2 ++ crates/tw-dialect/Cargo.toml | 1 + crates/tw-dialect/src/convert.rs | 5 ++--- crates/tw-dialect/src/usage.rs | 5 +---- crates/tw-gateway/Cargo.toml | 1 + crates/tw-gateway/src/clientprobe.rs | 10 +--------- 7 files changed, 10 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4827b683..dcf2d4ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3252,6 +3252,7 @@ dependencies = [ name = "tw-dialect" version = "0.64.1" dependencies = [ + "memchr", "serde", "serde_json", ] @@ -3287,6 +3288,7 @@ dependencies = [ "http-body-util", "hyper", "hyper-util", + "memchr", "rand 0.10.2", "reqwest", "rustls", diff --git a/Cargo.toml b/Cargo.toml index 8391e8b6..b1ad8111 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -103,6 +103,8 @@ flate2 = "1" # regex-automata / regex-syntax 已经在依赖树里（tracing 的 EnvFilter）， # 所以加它的增量很小 regex = "1" +# 在整个请求体里找一小段字节（SIMD）。serde_json、regex 本来就依赖它，不多编任何东西 +memchr = "2" # 保格式地改 TOML。接管 Codex 要写它的 config.toml，而重新序列化 # 会把用户的注释和排版一起抹掉 toml_edit = "0.25" diff --git a/crates/tw-dialect/Cargo.toml b/crates/tw-dialect/Cargo.toml index 7a3e72a2..06a6fd99 100644 --- a/crates/tw-dialect/Cargo.toml +++ b/crates/tw-dialect/Cargo.toml @@ -13,3 +13,4 @@ description = "Request, response and stream conversion between Anthropic Message [dependencies] serde = { workspace = true } serde_json = { workspace = true } +memchr = { workspace = true } diff --git a/crates/tw-dialect/src/convert.rs b/crates/tw-dialect/src/convert.rs index d0775178..793b36e1 100644 --- a/crates/tw-dialect/src/convert.rs +++ b/crates/tw-dialect/src/convert.rs @@ -852,9 +852,8 @@ impl Collector { /// `tw1.` 前缀的签名不是这个上游签发的：Anthropic 会以签名无效拒绝整个请求， /// 对话从此卡死。所以直通前把它们去掉 —— 那段推理本来就不是这个上游产生的。 pub fn strip_carried(client: Dialect, body: &[u8]) -> Option> { - if !body.windows(CARRIED.len()).any(|w| w == CARRIED.as_bytes()) { - return None; - } + // 每一跳同格式直通都要看一眼整个请求体：按字节找（SIMD），绝大多数请求在这里就回去了 + memchr::memmem::find(body, CARRIED.as_bytes())?; let mut v: Value = serde_json::from_slice(body).ok()?; let carried = |x: Option<&Value>| { x.and_then(Value::as_str) diff --git a/crates/tw-dialect/src/usage.rs b/crates/tw-dialect/src/usage.rs index 248fd7a5..dab624a0 100644 --- a/crates/tw-dialect/src/usage.rs +++ b/crates/tw-dialect/src/usage.rs @@ -492,10 +492,7 @@ fn find_at(hay: &[u8], needle: &[u8], from: usize) -> Option { if from >= hay.len() || needle.len() > hay.len() - from { return None; } - hay[from..] - .windows(needle.len()) - .position(|w| w == needle) - .map(|i| i + from) + memchr::memmem::find(&hay[from..], needle).map(|i| i + from) } /// 从 `open`（一个 `{`）开始配对，返回对应 `}` 的下标。 diff --git a/crates/tw-gateway/Cargo.toml b/crates/tw-gateway/Cargo.toml index b2859387..3ddbe4c1 100644 --- a/crates/tw-gateway/Cargo.toml +++ b/crates/tw-gateway/Cargo.toml @@ -39,6 +39,7 @@ tokio-rustls = { workspace = true } base64 = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +memchr = { workspace = true } blake3 = { workspace = true } sha2 = { workspace = true } rand = { workspace = true } diff --git a/crates/tw-gateway/src/clientprobe.rs b/crates/tw-gateway/src/clientprobe.rs index ea981f20..449f5699 100644 --- a/crates/tw-gateway/src/clientprobe.rs +++ b/crates/tw-gateway/src/clientprobe.rs @@ -62,7 +62,7 @@ const WARMUP_BODY: &str = "Warmup"; /// `is_claude_code` 由调用方判断。`max_tokens: 1` 那条**必须**同时要求 /// 它，否则会误伤别人真实的 `max_tokens: 1` 请求。 pub fn classify(body: &Bytes, is_claude_code: bool) -> Option { - let has = |m: &str| memfind(body, m.as_bytes()); + let has = |m: &str| memchr::memmem::find(body, m.as_bytes()).is_some(); // 快速排除。四个 B 类标记各有一句独特的原文；A 类里 warmup 也有， // 而 health check 只能靠 `max_tokens` 这个字段名先粗筛。 @@ -107,14 +107,6 @@ pub fn classify(body: &Bytes, is_claude_code: bool) -> Option { None } -/// 朴素子串查找。请求体是几十 KB 的量级，不值得为它引一个 SIMD 库。 -fn memfind(hay: &[u8], needle: &[u8]) -> bool { - if needle.is_empty() || hay.len() < needle.len() { - return false; - } - hay.windows(needle.len()).any(|w| w == needle) -} - /// 把 system 和所有消息里的文本摊平了看。 fn all_text(obj: &serde_json::Map) -> String { let mut out = String::new(); From 86b84e38a8e52ec615db05c07231911a9d820d98 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:56:54 +0800 Subject: [PATCH 09/16] perf(tw-guard): cut the secret scan to a fifth The redaction scan runs over every request body even in observe mode (22 ms per 2 MiB, on the request path and again before storage). About a quarter of it was SipHash: three rules were looked up in the enabled set for every token. Another sixth was memcmp: every token was compared with twenty-odd key prefixes. The rest was the tokenizer decoding the body character by character. - Ask whether the per-token rules are on once per scan, as `want_jwt` already did. - Reject a token before the prefix loop when it is shorter than the shortest key any rule can match or starts with a byte no prefix starts with. Both bounds are computed from BUILTINS at compile time, and a test checks the shortcut against the full comparison for every rule set. - Walk the body by byte. Tokens, backslashes and quotes are all ASCII and every byte of a multi-byte character is >= 0x80, so cutting by byte gives the same tokens; a test compares it with the old char-based cutter on random text with escapes, CJK, emoji and combining marks. - Find PEM headers and `://` with memmem. - `look` in enforce mode built a fully replaced copy of the body only to keep its ledger. Number the values without writing the copy. `look_hits` also hands back the hits it found, so a caller that scans the same bytes again (the desktop gateway, before storing the body) can reuse them. Results are unchanged: 2 MiB Claude Code-shaped body, 29-54 ms -> 3.7 ms; 8 MiB, 71-92 ms -> 15 ms. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 1 + crates/tw-guard/Cargo.toml | 1 + crates/tw-guard/src/redact/flow.rs | 85 ++++++- crates/tw-guard/src/redact/replace.rs | 14 ++ crates/tw-guard/src/redact/rules.rs | 322 +++++++++++++++++++++++--- 5 files changed, 381 insertions(+), 42 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index dcf2d4ed..419fa025 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3326,6 +3326,7 @@ version = "0.64.1" dependencies = [ "base64", "bytes", + "memchr", "regex", "serde", "serde_json", diff --git a/crates/tw-guard/Cargo.toml b/crates/tw-guard/Cargo.toml index e37cbffd..9a35191b 100644 --- a/crates/tw-guard/Cargo.toml +++ b/crates/tw-guard/Cargo.toml @@ -16,6 +16,7 @@ base64 = { workspace = true } bytes = { workspace = true } serde_json = { workspace = true } regex = { workspace = true } +memchr = { workspace = true } thiserror = { workspace = true } serde_yaml_ng = { workspace = true } tw-dialect = { workspace = true } diff --git a/crates/tw-guard/src/redact/flow.rs b/crates/tw-guard/src/redact/flow.rs index 42e8a2eb..ef8aeea7 100644 --- a/crates/tw-guard/src/redact/flow.rs +++ b/crates/tw-guard/src/redact/flow.rs @@ -128,25 +128,51 @@ pub fn look(mode: Mode, rules: &RuleSet, body: &[u8]) -> (Vec, Ledger) /// 请求），改过的这一份接着原文那本账编，同一个值还是同一个号，新出现的值接着往后编。 /// 不在替换档时账本是空的，`seed` 用不上。 pub fn look_from(mode: Mode, rules: &RuleSet, body: &[u8], seed: Ledger) -> (Vec, Ledger) { - let empty = || Ledger::new(Scheme::SECRET); + let l = look_hits(mode, rules, body, seed); + (l.found, l.ledger) +} + +/// [`look_from`] 的结果，连同找到的那些命中。 +#[derive(Debug, Clone)] +pub struct Look { + /// 报出去的记录（见 [`look`]） + pub found: Vec, + /// 这个请求的账本。不在替换档时是空的 + pub ledger: Ledger, + /// 这一遍在 `body` 上找到的命中（同 [`hits`]，区间是 `body` 里的）。没找的是 `None`： + /// 关着、一条规则都没开、不是 UTF-8。 + /// + /// 同一份正文之后还要再找一遍的调用方（桌面版落盘前打码）拿它省掉那一遍：规则一样、 + /// 正文一样，找出来的就一样 + pub hits: Option>, +} + +/// [`look_from`]，把找到的命中一起交回去（见 [`Look::hits`]）。 +pub fn look_hits(mode: Mode, rules: &RuleSet, body: &[u8], seed: Ledger) -> Look { + let empty = || Look { + found: Vec::new(), + ledger: Ledger::new(Scheme::SECRET), + hits: None, + }; if !mode.detects() || rules.is_empty() { - return (Vec::new(), empty()); + return empty(); } let Ok(text) = std::str::from_utf8(body) else { - return (Vec::new(), empty()); + return empty(); }; let hits = hits(text, rules); let found = crate::redact::rules::findings(text, &hits); - if !mode.acts() { - return (found, empty()); - } - let seed = seed.avoiding(text); - let ledger = if hits.is_empty() { - seed + let ledger = if !mode.acts() { + Ledger::new(Scheme::SECRET) } else { - crate::redact::replace::apply(text, &hits, seed).ledger + // 只要账本：换过的那一份在每一跳发出去之前才写（见 `replace`），这里不抄一份 + crate::redact::replace::number(text, &hits, seed.avoiding(text)) }; - (found, ledger) + Look { + found, + ledger, + hits: Some(hits), + } } /// 替换档下换掉要发出去的这一份，**接着 `ledger` 的账**（见 [`look`]）。返回换过的体和 @@ -398,6 +424,43 @@ mod tests { assert!(l.is_empty()); } + /// 看一遍只编号、不抄一份：账本和真换一遍得到的一样（让开原文里写着的占位符、同一个值 + /// 一个号），交回的命中就是 [`hits`] 找到的 + #[test] + fn looking_numbers_like_replacing_and_hands_back_its_hits() { + let other = "ghp_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; + let body = format!( + r#"{{"system":"<> {KEY}","messages":[{{"content":"{other} {KEY}"}}]}}"# + ); + let rules = RuleSet::defaults(); + let l = look_hits(Mode::Enforce, &rules, body.as_bytes(), fresh()); + let found = hits(&body, &rules); + assert_eq!(found.len(), 3); + assert_eq!(l.hits.as_deref(), Some(&found[..])); + let replaced = crate::redact::replace::apply(&body, &found, fresh().avoiding(&body)); + assert_eq!(l.ledger.table(), replaced.ledger.table()); + assert_eq!(l.ledger.len(), 2); + assert_eq!((l.found.clone(), l.ledger.table().clone()), { + let (f, ledger) = look(Mode::Enforce, &rules, body.as_bytes()); + (f, ledger.table().clone()) + }); + // 观察档照样找、照样交回命中，只是不编号 + let o = look_hits(Mode::Observe, &rules, body.as_bytes(), fresh()); + assert_eq!((o.found, o.hits), (l.found, l.hits)); + assert!(o.ledger.is_empty()); + // 关着的、不是 UTF-8 的不找 + assert!( + look_hits(Mode::Off, &rules, body.as_bytes(), fresh()) + .hits + .is_none() + ); + assert!( + look_hits(Mode::Enforce, &rules, &[0xff, 0xfe], fresh()) + .hits + .is_none() + ); + } + /// 连接串里的占位符长得像口令，可它不是凭据：不再换一次、不报、原样留着 #[test] fn a_placeholder_where_a_password_would_be_is_not_a_password() { diff --git a/crates/tw-guard/src/redact/replace.rs b/crates/tw-guard/src/redact/replace.rs index af29e3d9..3c7880a2 100644 --- a/crates/tw-guard/src/redact/replace.rs +++ b/crates/tw-guard/src/redact/replace.rs @@ -207,6 +207,20 @@ pub fn apply(text: &str, hits: &[Hit], mut ledger: Ledger) -> Redacted { Redacted { text: out, ledger } } +/// 只编号、不抄：[`apply`] 会给 `hits` 发的那些号，不写出换过的那一份。只要账本时用它 +/// （拦截档下看一遍客户端原文，见 [`crate::redact::flow::look`]）—— 换过的那一份在每一跳 +/// 发出去之前才写。 +pub(crate) fn number(text: &str, hits: &[Hit], mut ledger: Ledger) -> Ledger { + let default = ledger.scheme.label; + for h in hits { + ledger.issue( + &text[h.bytes.clone()], + h.label.as_deref().unwrap_or(default), + ); + } + ledger +} + /// 扫 + 换，一步到位。`text` 是 JSON 请求体（见 [`crate::redact::rules::scan`]）。 pub fn redact(text: &str, rules: &RuleSet, ledger: Ledger) -> Redacted { let hits = crate::redact::rules::scan(text, rules); diff --git a/crates/tw-guard/src/redact/rules.rs b/crates/tw-guard/src/redact/rules.rs index a4ef6ec5..648714ea 100644 --- a/crates/tw-guard/src/redact/rules.rs +++ b/crates/tw-guard/src/redact/rules.rs @@ -523,12 +523,68 @@ impl RuleSet { } } -fn is_tok(c: char) -> bool { - c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' +/// 一个字节是不是 token 的一部分。**token 只由 ASCII 组成**，所以按字节切和按字符切切出来的 +/// 一样：多字节字符的每个字节都不小于 0x80，既落不进 token，也不是反斜杠和引号。 +const fn is_tok(b: u8) -> bool { + b.is_ascii_alphanumeric() || b == b'-' || b == b'_' || b == b'.' } -/// 一段 token 是哪种 API 密钥。**只看开着的那几条。** -fn classify_token(tok: &str, set: &RuleSet) -> Option<&'static str> { +/// [`is_tok`] 查表：每个请求体的每个字节都要问一次 +const TOK: [bool; 256] = { + let mut t = [false; 256]; + let mut b = 0; + while b < 256 { + t[b] = is_tok(b as u8); + b += 1; + } + t +}; + +/// 一把密钥最短有多长：前缀规则里最短的那一种（前缀加上至少要有的尾巴），和 OpenAI 老式 +/// 密钥的下限，取小的。比它短的 token 哪一种都不是 —— 一段正文里绝大多数 token 在这里就 +/// 回去了 +const KEY_MIN: usize = { + let mut min = OPENAI_MIN; + let mut i = 0; + while i < BUILTINS.len() { + if let Matcher::Prefix { prefix, min_tail } = BUILTINS[i].matcher + && prefix.len() + min_tail < min + { + min = prefix.len() + min_tail; + } + i += 1; + } + min +}; + +/// 一把密钥能以哪些字节开头：每条前缀规则的头一个字节，和 OpenAI 老式密钥的 `s` +const KEY_FIRST: [bool; 256] = { + let mut t = [false; 256]; + t[b's' as usize] = true; + let mut i = 0; + while i < BUILTINS.len() { + if let Matcher::Prefix { prefix, .. } = BUILTINS[i].matcher { + t[prefix.as_bytes()[0] as usize] = true; + } + i += 1; + } + t +}; + +/// 一段 token 是哪种 API 密钥。**只看开着的那几条。**`openai`：OpenAI 老式密钥那条开着 +/// （扫之前问一次，不在每个 token 上查一遍表）。 +/// +/// 太短的、开头哪条前缀都对不上的直接回去：[`classify_prefixed`] 认得出的 token 都过得了 +/// 这两关（[`KEY_MIN`]、[`KEY_FIRST`]）。 +fn classify_token(tok: &str, set: &RuleSet, openai: bool) -> Option<&'static str> { + if tok.len() < KEY_MIN || !KEY_FIRST[usize::from(tok.as_bytes()[0])] { + return None; + } + classify_prefixed(tok, set, openai) +} + +/// [`classify_token`] 的判据本身。 +fn classify_prefixed(tok: &str, set: &RuleSet, openai: bool) -> Option<&'static str> { for b in BUILTINS { if let Matcher::Prefix { prefix, min_tail } = b.matcher && let Some(tail) = tok.strip_prefix(prefix) @@ -539,7 +595,7 @@ fn classify_token(tok: &str, set: &RuleSet) -> Option<&'static str> { return set.is_on(b.id).then_some(b.id); } } - if set.is_on("openai-api-key") + if openai && let Some(tail) = tok.strip_prefix("sk-") && tok.len() >= OPENAI_MIN && !tail.starts_with("ant-") @@ -569,34 +625,43 @@ fn classify_token(tok: &str, set: &RuleSet) -> Option<&'static str> { /// /// 回调的第三个参数说这个 token 在不在一个 JSON 字符串里面（数没转义的引号； /// 文本是 JSON 时才有意义）。 +/// +/// **按字节走，不按字符解码**：token、反斜杠、引号都是 ASCII（见 [`is_tok`]），多字节字符 +/// 的字节在这里只是分隔。反斜杠后面跟着的是一个多字节字符时，跳过的只是它的头一个字节， +/// 剩下的几个照样只是分隔 —— 切出来的和按字符切的一样（测试拿两种切法对过）。token 的两头 +/// 都落在 ASCII 字节上，切片一定在字符边界上。 fn for_each_token(text: &str, mut f: impl FnMut(&str, Range, bool)) { + let b = text.as_bytes(); let mut start: Option = None; let mut quoted = false; - let mut chars = text.char_indices().peekable(); - while let Some((i, c)) = chars.next() { - if c == '\\' { - if let Some(s) = start.take() { - f(&text[s..i], s..i, quoted); - } - // 反斜杠后面那个字符是转义的一部分；`\u` 再带四位十六进制 - if let Some((_, 'u')) = chars.next() { - for _ in 0..4 { - if chars.next_if(|(_, h)| h.is_ascii_hexdigit()).is_none() { - break; - } - } - } - continue; - } - if is_tok(c) { + let mut i = 0; + while i < b.len() { + let c = b[i]; + if TOK[usize::from(c)] { start.get_or_insert(i); + i += 1; continue; } if let Some(s) = start.take() { f(&text[s..i], s..i, quoted); } - if c == '"' { - quoted = !quoted; + i += 1; + match c { + // 反斜杠后面那个字符是转义的一部分；`\u` 再带四位十六进制 + b'\\' => { + if b.get(i) == Some(&b'u') { + i += 1; + let mut n = 0; + while n < 4 && b.get(i).is_some_and(u8::is_ascii_hexdigit) { + i += 1; + n += 1; + } + } else if i < b.len() { + i += 1; + } + } + b'"' => quoted = !quoted, + _ => {} } } if let Some(s) = start { @@ -615,7 +680,8 @@ const PEM_KEY: &str = "PRIVATE KEY-----"; /// 空壳 PEM，它会以为文件坏了然后建议你重新生成一把 —— 那是在帮倒忙。 fn private_keys(text: &str, out: &mut Vec) { let mut from = 0; - while let Some(b) = text[from..].find(PEM_BEGIN) { + // 按字节找（SIMD）：整个请求体都要过一遍，而它几乎从来不在里面 + while let Some(b) = memchr::memmem::find(&text.as_bytes()[from..], PEM_BEGIN.as_bytes()) { let begin = from + b; // **不能按换行找那一行的结尾。**请求体是 JSON，里面的换行是 // `\n` 两个字符，不是一个换行符 —— 按换行找的话整个 PEM 会落在 @@ -663,6 +729,10 @@ fn is_b64url(s: &str) -> bool { /// 那个解码是必需的：光看「三段用点分开的 base64」会把版本号、文件路径、 /// 甚至 `a.b.c` 这种普通标识符全算上。 fn looks_like_jwt(tok: &str) -> bool { + // 首段至少 8 个字符，后两段各至少一个，加两个点：短于 12 的不可能是 + if tok.len() < 12 { + return false; + } // 先数点，不把三段收集起来：每个 token 都要过这一关，绝大多数在这里就回去了 if tok.bytes().filter(|&b| b == b'.').count() != 2 || !tok.split('.').all(is_b64url) { return false; @@ -695,7 +765,7 @@ fn looks_like_jwt(tok: &str) -> bool { fn conn_strings(text: &str, out: &mut Vec) { let bytes = text.as_bytes(); let mut from = 0; - while let Some(i) = text[from..].find("://") { + while let Some(i) = memchr::memmem::find(&bytes[from..], b"://") { let sep = from + i; let after = sep + 3; // userinfo 到 `@` 为止；`@` 必须在同一段里（不能跨空白或斜杠） @@ -752,14 +822,16 @@ fn is_rfc1918(tok: &str) -> bool { /// **不含回环。**`127.0.0.1` 和 `localhost` 是这台机器自己 —— 而用户问 /// 的很可能正是「我本地这个服务为什么连不上」，把它换成占位符等于把问题 /// 本身藏起来了。何况我们的网关自己就住在那儿。 -fn internal_token(tok: &str, set: &RuleSet) -> Option<&'static str> { - if set.is_on("internal-ip") && is_rfc1918(tok) { +/// +/// `ip`、`domain`：那两条开没开（扫之前问一次，不在每个 token 上查一遍表）。 +fn internal_token(tok: &str, ip: bool, domain: bool) -> Option<&'static str> { + if ip && is_rfc1918(tok) { return Some("internal-ip"); } // 后缀不分大小写地比，**不先转成小写**：每个 token 都要过这一关，转一次就是 // 一次分配 let t = tok.as_bytes(); - if set.is_on("internal-domain") + if domain && t.len() > 7 && t.contains(&b'.') && INTERNAL_SUFFIXES.iter().any(|s| { @@ -1264,6 +1336,11 @@ pub fn scan(text: &str, set: &RuleSet) -> Vec { if set.is_on("email") { emails(text, &mut out); } + // 每个 token 都要问的几条开没开，扫之前问一次：在每个 token 上查一遍表，光哈希就占了 + // 整遍扫描的两成 + let want_openai = set.is_on("openai-api-key"); + let want_ip = set.is_on("internal-ip"); + let want_domain = set.is_on("internal-domain"); let want_jwt = set.is_on("jwt"); let want_id = set.is_on("cn-resident-id"); let want_card = set.is_on("bank-card"); @@ -1291,7 +1368,7 @@ pub fn scan(text: &str, set: &RuleSet) -> Vec { // 字符串外面的东西把在接的那一串截断 spaced.finish(text, &mut out); } - if let Some(id) = classify_token(tok, set) { + if let Some(id) = classify_token(tok, set, want_openai) { out.push(builtin_hit(id, span)); return; } @@ -1299,7 +1376,7 @@ pub fn scan(text: &str, set: &RuleSet) -> Vec { out.push(builtin_hit("jwt", span)); return; } - if let Some(id) = internal_token(tok, set) { + if let Some(id) = internal_token(tok, want_ip, want_domain) { out.push(builtin_hit(id, span)); } }); @@ -2672,4 +2749,187 @@ mod tests { ); assert_eq!(crate::redact::replace::restore(&r.text, &r.ledger), body); } + + /// 以前的切法：按字符解码着走。按字节走的 [`for_each_token`] 切出来的要和它一模一样 + fn tokens_by_char(text: &str) -> Vec<(String, Range, bool)> { + let mut out = Vec::new(); + let mut f = |t: &str, r: Range, q: bool| out.push((t.to_string(), r, q)); + let mut start: Option = None; + let mut quoted = false; + let mut chars = text.char_indices().peekable(); + while let Some((i, c)) = chars.next() { + if c == '\\' { + if let Some(s) = start.take() { + f(&text[s..i], s..i, quoted); + } + if let Some((_, 'u')) = chars.next() { + for _ in 0..4 { + if chars.next_if(|(_, h)| h.is_ascii_hexdigit()).is_none() { + break; + } + } + } + continue; + } + if c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' { + start.get_or_insert(i); + continue; + } + if let Some(s) = start.take() { + f(&text[s..i], s..i, quoted); + } + if c == '"' { + quoted = !quoted; + } + } + if let Some(s) = start { + f(&text[s..], s..text.len(), quoted); + } + out + } + + fn tokens_by_byte(text: &str) -> Vec<(String, Range, bool)> { + let mut out = Vec::new(); + for_each_token(text, |t, r, q| out.push((t.to_string(), r, q))); + out + } + + /// 随机拼出来的一段：ASCII 的字母数字和标点、引号、反斜杠和各种转义（含不完整的 + /// `\u`、反斜杠后面跟着中文和 emoji、结尾一个落单的反斜杠）、中文、emoji、组合字符 + fn random_text(next: &mut impl FnMut(usize) -> usize) -> String { + const PIECES: &[&str] = &[ + "a", + "Z", + "9", + "-", + "_", + ".", + " ", + ",", + ":", + "{", + "}", + "[", + "]", + "\"", + "\\", + "\\\\", + "\\n", + "\\t", + "\\\"", + "\\u", + "\\u00e9", + "\\u5bc6", + "\\uD83D\\uDE00", + "\\u12", + "\\uzz", + "\\é", + "\\密", + "\\😀", + "u", + "00e9", + "sk-ant-", + "AKIA", + "eyJ", + "密钥", + "：", + "😀", + "👩‍💻", + "e\u{301}", + "\u{200B}", + "\u{7f}", + "\t", + "\n", + "é", + "→", + "ghp_", + "@", + "/", + "+", + "=", + ]; + let n = next(40); + (0..n).map(|_| PIECES[next(PIECES.len())]).collect() + } + + #[test] + fn cutting_by_byte_cuts_exactly_where_cutting_by_char_did() { + let mut seed = 0x2545_f491_4f6c_dd1d_u64; + let mut next = |n: usize| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + (seed % n as u64) as usize + }; + for _ in 0..20_000 { + let t = random_text(&mut next); + assert_eq!(tokens_by_byte(&t), tokens_by_char(&t), "{t:?}"); + } + for t in [ + "", + "\\", + "a\\", + "\\u", + "\\u0", + "\\😀x", + "\"a\\\"b\"c", + "密\\钥sk", + "x\u{1F600}", + ] { + assert_eq!(tokens_by_byte(t), tokens_by_char(t), "{t:?}"); + } + } + + /// 太短的、开头哪条前缀都对不上的 token 直接回去：哪一条规则开着、关着，结论都和逐条 + /// 比前缀的一样 + #[test] + fn the_quick_rejections_never_drop_a_key() { + let mut seed = 0x9e37_79b9_7f4a_7c15_u64; + let mut next = |n: usize| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + (seed % n as u64) as usize + }; + const TAIL: &[u8] = b"abcXYZ0189-_."; + let mut heads: Vec<&str> = BUILTINS + .iter() + .filter_map(|b| match b.matcher { + Matcher::Prefix { prefix, .. } => Some(prefix), + _ => None, + }) + .collect(); + heads.extend(["sk-", "s", "x", "", "Sk-", "ak"]); + let sets = [ + all(), + RuleSet::defaults(), + RuleSet::none(), + RuleSet::only(&["openai-api-key"]), + RuleSet::only(&["anthropic-api-key", "aws-access-key-id"]), + ]; + for _ in 0..20_000 { + let mut tok = heads[next(heads.len())].to_string(); + for _ in 0..next(50) { + tok.push(TAIL[next(TAIL.len())] as char); + } + if tok.is_empty() { + continue; + } + for set in &sets { + for openai in [true, false] { + assert_eq!( + classify_token(&tok, set, openai), + classify_prefixed(&tok, set, openai), + "{tok}" + ); + } + } + } + // 最短的那种刚好够长时认得出 + assert_eq!(KEY_MIN, "AKIA".len() + 12); + assert_eq!( + classify_token("AKIAABCDEFGHIJKL", &all(), true), + Some("aws-access-key-id") + ); + } } From 270c154c3189e8692d18a232ab21478a288446e7 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:57:54 +0800 Subject: [PATCH 10/16] perf(tw-guard): content filter decodes only the characters it can match MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With the default rules (two code-point rules, three phrases, observe mode) the content filter cost 13-16 ms of a 2 MiB Claude Code request besides parsing it. Claude Code's file reads put `→` after every line number, so nearly every tool result is non-ASCII and missed the existing all-ASCII shortcut: - A code-point rule whose code points are all above U+007F cannot match an ASCII byte. Skip ASCII eight bytes at a time and decode only the other characters. - The lowercase copy for phrase rules came from `str::to_lowercase`, which goes character by character after the first non-ASCII one. Lowercase ASCII runs in bulk and the other characters one by one; `Σ`, the only context-dependent mapping, hands the whole text back to the standard library. - Find phrases with memmem, non-overlapping like `match_indices`. Tests compare each against the old way on random text. Content filter on a 2 MiB body: 15-18 ms -> 4.8 ms including the parse (1.9 ms). Co-Authored-By: Claude Opus 5.5 --- crates/tw-guard/src/content.rs | 189 +++++++++++++++++++++++++++++++-- 1 file changed, 183 insertions(+), 6 deletions(-) diff --git a/crates/tw-guard/src/content.rs b/crates/tw-guard/src/content.rs index 79e4e7b5..54e96a3a 100644 --- a/crates/tw-guard/src/content.rs +++ b/crates/tw-guard/src/content.rs @@ -259,10 +259,54 @@ impl<'a> Lower<'a> { Lower { text, lower: None } } fn get(&mut self) -> &str { - self.lower.get_or_insert_with(|| self.text.to_lowercase()) + self.lower.get_or_insert_with(|| lowercase(self.text)) } } +/// 和 `text.to_lowercase()` 一样的结果，快一些：ASCII 的一段整段转，别的字符逐个转（标准库 +/// 碰到头一个 ASCII 以外的字符之后就逐个字符地转了，夹着中文、`→` 的正文几乎整段都是）。 +/// +/// 按上下文转的只有 `Σ`（在词尾是 `ς`）：正文里有它就整段交给标准库。 +fn lowercase(text: &str) -> String { + if memchr::memmem::find(text.as_bytes(), "Σ".as_bytes()).is_some() { + return text.to_lowercase(); + } + let b = text.as_bytes(); + let mut out = String::with_capacity(text.len()); + let mut i = 0; + while i < b.len() { + let start = i; + i = ascii_end(b, i); + if i > start { + let at = out.len(); + out.push_str(&text[start..i]); + out[at..].make_ascii_lowercase(); + } + // 停在一个多字节字符的头一个字节上：字符边界 + if let Some(c) = text.get(i..).and_then(|rest| rest.chars().next()) { + out.extend(c.to_lowercase()); + i += c.len_utf8(); + } + } + out +} + +/// 从 `i` 起连着的 ASCII 字节到哪儿为止。八个一组地看 +fn ascii_end(b: &[u8], mut i: usize) -> usize { + const HIGH: u64 = 0x8080_8080_8080_8080; + while let Some(word) = b.get(i..i + 8) { + let word = u64::from_ne_bytes(word.try_into().unwrap_or_default()); + if word & HIGH != 0 { + break; + } + i += 8; + } + while i < b.len() && b[i].is_ascii() { + i += 1; + } + i +} + /// 一条规则在一段正文里的全部命中 #[derive(Debug, Default)] struct Found { @@ -289,7 +333,9 @@ fn find_contains(text: &str, lower: &str, needle: &str) -> Vec> { return out; } if lower.len() == text.len() { - for (at, _) in lower.match_indices(needle) { + // 按字节找（SIMD）：每段正文、每条规则都要从头找到尾。和 `match_indices` 一样是 + // 一个接一个、互不重叠的 + for at in memchr::memmem::find_iter(lower.as_bytes(), needle.as_bytes()) { // 区间落回原文要对齐到字符边界 let start = floor(text, at); let end = ceil(text, at + needle.len()); @@ -373,19 +419,46 @@ fn find_points(p: &Codepoints, text: &str, escapes: bool) -> Found { if p.min() > 0x7F && text.is_ascii() && (!escapes || !text.contains("\\u")) { return out; } - for (r, c) in chars(text, escapes) { - if !p.contains(c) { - continue; - } + let mut hit = |r: Range| { out.count += 1; match out.ranges.last_mut() { Some(last) if last.end == r.start => last.end = r.end, _ => out.ranges.push(r), } + }; + // 码位都在 ASCII 以外、又不认转义：ASCII 的字节一个都命中不了，跳过去，只解码别的 + // 字符。夹着中文、`→` 的正文（工具读回来的代码）就不再逐个字符地解一遍 + if p.min() > 0x7F && !escapes { + for (r, c) in non_ascii(text) { + if p.contains(c) { + hit(r); + } + } + } else { + for (r, c) in chars(text, escapes) { + if p.contains(c) { + hit(r); + } + } } out } +/// 一段正文里 ASCII 以外的每一个字符和它的字节区间，按出现的先后。ASCII 的字节八个一组 +/// 地跳过（[`ascii_end`]）。 +fn non_ascii(text: &str) -> impl Iterator, char)> + '_ { + let b = text.as_bytes(); + let mut i = 0; + std::iter::from_fn(move || { + i = ascii_end(b, i); + // 停在一个多字节字符的头一个字节上：字符边界 + let c = text.get(i..)?.chars().next()?; + let at = i; + i += c.len_utf8(); + Some((at..i, c)) + }) +} + /// 一处命中。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct Hit { @@ -978,4 +1051,108 @@ mod tests { assert_eq!(find_points(&p, &text, false).count, 0); assert!(text.is_ascii(), "例子里不该有真的不可见字符"); } + + /// 码位都在 ASCII 以外时只解码 ASCII 以外的字符：找到的和逐个字符看的一样 —— 几段、 + /// 每段从哪到哪、一共几个 + #[test] + fn skipping_ascii_finds_what_looking_at_every_character_found() { + let mut seed = 0x2545_f491_4f6c_dd1d_u64; + let mut next = |n: usize| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + (seed % n as u64) as usize + }; + // 码位用 `char::from_u32` 拼：测试里直接写出来的转义，经过某些编辑工具会变成真字符 + let ch = |n: u32| char::from_u32(n).unwrap().to_string(); + let pieces: Vec = [ + "a", + "Z", + " ", + "0", + "\"", + "{", + "→", + "密钥", + "é", + "😀", + "abcdefghijklmnop", + ] + .iter() + .map(|s| s.to_string()) + .chain( + [ + 0x200B, 0x202E, 0x2066, 0xE0041, 0xE007F, 0xE0000, 0x7F, 0x80, 0xFEFF, + ] + .into_iter() + .map(ch), + ) + .collect(); + let sets = [ + "U+E0000–U+E007F", + "U+202A–U+202E, U+2066–U+2069", + "U+200B, U+FEFF", + "U+0080–U+10FFFF", + ] + .map(|p| Codepoints::parse(p).unwrap()); + for _ in 0..5_000 { + let text: String = (0..next(60)) + .map(|_| pieces[next(pieces.len())].as_str()) + .collect(); + for p in &sets { + assert!(p.min() > 0x7F); + let fast = find_points(p, &text, false); + let mut slow = Found::default(); + for (r, c) in chars(&text, false) { + if p.contains(c) { + slow.count += 1; + match slow.ranges.last_mut() { + Some(last) if last.end == r.start => last.end = r.end, + _ => slow.ranges.push(r), + } + } + } + assert_eq!( + (fast.count, fast.ranges), + (slow.count, slow.ranges), + "{text:?}" + ); + } + } + } + + /// 快一些的小写和标准库的一字不差：`Σ` 在词尾、词中，会变长的（`İ`）、四个字节的、 + /// 组合字符、开尔文符号这类转成 ASCII 的 + #[test] + fn lowercasing_by_runs_gives_what_the_standard_library_gives() { + let mut seed = 0x9e37_79b9_7f4a_7c15_u64; + let mut next = |n: usize| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + (seed % n as u64) as usize + }; + // 码位用 `char::from_u32` 拼：测试里直接写出来的转义，经过某些编辑工具会变成真字符 + let ch = |n: u32| char::from_u32(n).unwrap().to_string(); + let pieces: Vec = ["A", "z", " ", "IGNORE", "Previous", ".", "→", "密钥", "😀"] + .iter() + .map(|s| s.to_string()) + .chain( + [ + 0x03A3, 0x0391, 0x0130, 0x212A, 0x1E9E, 0x00DF, 0x01C5, 0xFB00, 0x0301, 0x2167, + 0x10400, 0x00C9, 0x0049, 0x0131, 0x03C2, + ] + .into_iter() + .map(ch), + ) + .collect(); + for _ in 0..20_000 { + let text: String = (0..next(30)) + .map(|_| pieces[next(pieces.len())].as_str()) + .collect(); + assert_eq!(lowercase(&text), text.to_lowercase(), "{text:?}"); + } + let long = format!("{}{}", "Ignore Previous ".repeat(10), ch(0x212A)); + assert_eq!(lowercase(&long), long.to_lowercase()); + } } From bd97e901254733c999197e290b91b20bf31fbf5e Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 21:03:42 +0800 Subject: [PATCH 11/16] perf(gateway): parse a request once; cut the client identity by bytes A Claude Code request body was parsed three times before its first hop: for the routing facts, for the content filter, and again on every hop to drop `metadata.user_id` -- which Claude Code always sends and `forward_client_identity` is off for by default. The local-answer check parsed it a fourth time whenever the body mentions `haiku`, which Claude Code's Task tool schema lists among its models. The pipeline now parses the body once and hands the value to the local-answer check, the routing facts and the content filter (`tw_guard::content::screen_value`, which only copies the value when it has to strip something). The identity is cut out of the body bytes instead of parsing and writing it back. Besides the cost (4-6.6 ms and a full copy per hop at 2 MiB), writing it back was a fidelity bug: serde_json without `preserve_order` sorts every object's keys, so every Claude Code request reached the upstream with its tool schemas, tool inputs and top level reordered, which also changes the bytes an upstream's prompt cache keys on. Now the forwarded body is the client's body less that one member (and its comma, or the whole `metadata` once it is empty). Keys are compared decoded, as serde_json does. A key written twice falls back to the old parse and write, since serde_json keeps the last one. The cut relies on the body being valid JSON, which the pipeline already knows from its parse (`Reading::json`); a body that does not parse is still forwarded untouched. A random test checks that the cut means what the old rewrite meant, and a gateway test that a 1.3 MB Claude Code body reaches the upstream byte for byte less `metadata`. Co-Authored-By: Claude Opus 5.5 --- crates/tw-gateway/src/client_api.rs | 4 + crates/tw-gateway/src/clientprobe.rs | 22 +- crates/tw-gateway/src/egress.rs | 456 ++++++++++++++++++- crates/tw-gateway/src/guard.rs | 10 + crates/tw-gateway/src/server/pipeline.rs | 51 ++- crates/tw-gateway/src/server/pipeline/hop.rs | 11 +- crates/tw-gateway/tests/headers_and_paths.rs | 28 ++ crates/tw-guard/src/content.rs | 2 +- crates/tw-guard/src/content/screen.rs | 17 +- 9 files changed, 566 insertions(+), 35 deletions(-) diff --git a/crates/tw-gateway/src/client_api.rs b/crates/tw-gateway/src/client_api.rs index 3c1b5229..8b72581a 100644 --- a/crates/tw-gateway/src/client_api.rs +++ b/crates/tw-gateway/src/client_api.rs @@ -224,6 +224,9 @@ pub struct Reading { pub facts: tw_engine::RequestFacts, /// DeepSeek Harness 发的请求（见 [`tw_dialect::harness`]）。要看请求头，由管线填 pub harness: Option, + /// 请求体解得开（是 JSON）。之后每一跳按字节改它的地方（见 [`crate::egress`]）凭这个 + /// 知道不用再解一遍 + pub json: bool, } pub fn read(path: &str, query: Option<&str>, body: Option<&serde_json::Value>) -> Reading { @@ -259,6 +262,7 @@ pub fn read(path: &str, query: Option<&str>, body: Option<&serde_json::Value>) - decoded, facts, harness: None, + json: body.is_some(), } } diff --git a/crates/tw-gateway/src/clientprobe.rs b/crates/tw-gateway/src/clientprobe.rs index 449f5699..db8ba3a0 100644 --- a/crates/tw-gateway/src/clientprobe.rs +++ b/crates/tw-gateway/src/clientprobe.rs @@ -54,14 +54,15 @@ const TOPIC_MARK: &str = "Analyze if this message indicates a new conversation t const SUGGEST_MARK: &str = "[SUGGESTION MODE:"; const WARMUP_BODY: &str = "Warmup"; -/// 认一认这是不是客户端自己的辅助请求。 +/// 认一认这是不是客户端自己的辅助请求。`parsed` 是 `body` 解析出来的 JSON（解不开是 +/// `None`）：管线为路由本来就要解析一遍，这里不再解一遍。 /// -/// **先做整包字节扫描快速排除**，不命中就直接返回，连 JSON 都不解析。 +/// **先做整包字节扫描快速排除**，不命中就直接返回，连解析出来的 JSON 都不看。 /// 绝大多数请求会在这一步出去 —— 而这段代码跑在每一个请求的关键路径上。 /// /// `is_claude_code` 由调用方判断。`max_tokens: 1` 那条**必须**同时要求 /// 它，否则会误伤别人真实的 `max_tokens: 1` 请求。 -pub fn classify(body: &Bytes, is_claude_code: bool) -> Option { +pub fn classify(body: &[u8], parsed: Option<&Value>, is_claude_code: bool) -> Option { let has = |m: &str| memchr::memmem::find(body, m.as_bytes()).is_some(); // 快速排除。四个 B 类标记各有一句独特的原文；A 类里 warmup 也有， @@ -75,8 +76,7 @@ pub fn classify(body: &Bytes, is_claude_code: bool) -> Option { return None; } - let v: Value = serde_json::from_slice(body).ok()?; - let obj = v.as_object()?; + let obj = parsed?.as_object()?; // B 类的三条：标记出现在 system 或者第一个 user 消息里就算。 // **先判 B 类** —— 一个带标记的请求即使 max_tokens 是 1，它也是 @@ -331,6 +331,12 @@ mod tests { Bytes::from(s.to_string()) } + /// 和管线一样：先解析一遍，再认 + fn classify(body: &Bytes, is_claude_code: bool) -> Option { + let parsed = serde_json::from_slice::(body).ok(); + super::classify(body, parsed.as_ref(), is_claude_code) + } + // ── 该认出来的 ────────────────────────────────────────────────── #[test] fn a_health_check_is_recognised() { @@ -441,10 +447,14 @@ mod tests { #[test] fn an_ordinary_request_does_not_even_get_parsed() { - // 快速排除跑在每个请求的关键路径上。这条测的是它确实先排除了。 + // 快速排除跑在每个请求的关键路径上。这条测的是它确实先排除了：原文里一个标记都 + // 没有，解析出来的 JSON 就不看（这里故意交一份「是预热」的进去） let body = b(r#"{"model":"claude-sonnet-4-5","max_tokens":8000, "messages":[{"role":"user","content":"帮我看看这段代码"}]}"#); assert_eq!(classify(&body, true), None); + let warmup: Value = + serde_json::json!({"messages": [{"role": "user", "content": "Warmup"}]}); + assert_eq!(super::classify(&body, Some(&warmup), true), None); } #[test] diff --git a/crates/tw-gateway/src/egress.rs b/crates/tw-gateway/src/egress.rs index be407ae9..40c63e5f 100644 --- a/crates/tw-gateway/src/egress.rs +++ b/crates/tw-gateway/src/egress.rs @@ -150,19 +150,41 @@ pub fn gateway_headers(hop: &Hop, client: &HeaderMap) -> Vec<(String, String)> { /// （[`crate::chatgpt::IDENTITY_METADATA`]）；会话、线程这些照发 /// /// 返回去掉之后的请求体；没有可去的是 `None`，请求体一个字节都不动。这不是为了兼容去掉 -/// 请求的内容，所以不算丢弃的字段 +/// 请求的内容，所以不算丢弃的字段。 +/// +/// **只剪掉那几个成员，别的字节一个不动**（见 [`cut_members`]）：解析再写回去的话，没开 +/// `preserve_order` 的 serde_json 会把每个对象的键按字母重排 —— 工具定义、工具参数、整个 +/// 请求都变了样（上游的提示缓存按字节认），而 Claude Code 的请求每一个都带着 `user_id`。 +/// 剪不了的怪样子（同一个键写了两遍）才照旧解析、去掉、写回。 +/// +/// `body` 要是解得开的 JSON：调用方解过它（见管线的 `read`），这里不再解一遍 pub fn strip_body_identity(dialect: Dialect, body: &Bytes) -> Option { - // 大多数请求没有这些字段：先便宜地看一眼，免得每个请求都解析、重写一遍请求体 - let text = std::str::from_utf8(body).ok()?; + // 大多数请求没有这些字段：先便宜地看一眼，免得每个请求都把整个请求体走一遍 + let (parent, keys, drop_empty): (&str, &[&str], bool) = match dialect { + Dialect::Anthropic => ("metadata", &["user_id"], true), + Dialect::Responses => ("client_metadata", &crate::chatgpt::IDENTITY_METADATA, false), + _ => return None, + }; let marks: &[&str] = match dialect { Dialect::Anthropic => &["\"user_id\""], - Dialect::Responses => &crate::chatgpt::IDENTITY_METADATA, - _ => return None, + _ => keys, }; - if !marks.iter().any(|m| text.contains(m)) { + if !marks + .iter() + .any(|m| memchr::memmem::find(body, m.as_bytes()).is_some()) + { return None; } - let mut v: Value = serde_json::from_str(text).ok()?; + match cut_members(body, parent, keys, drop_empty) { + Cut::Done(out) => Some(Bytes::from(out)), + Cut::Nothing => None, + Cut::Unusual => reserialize(dialect, body), + } +} + +/// 解析、去掉、写回：[`strip_body_identity`] 剪不了的时候走这条（键会按字母重排）。 +fn reserialize(dialect: Dialect, body: &Bytes) -> Option { + let mut v: Value = serde_json::from_slice(body).ok()?; let obj = v.as_object_mut()?; let removed = match dialect { Dialect::Anthropic => { @@ -186,6 +208,210 @@ pub fn strip_body_identity(dialect: Dialect, body: &Bytes) -> Option { serde_json::to_vec(&v).ok().map(Bytes::from) } +/// [`cut_members`] 的结果。 +#[derive(Debug, PartialEq)] +enum Cut { + /// 剪过的请求体 + Done(Vec), + /// 没有要剪的 + Nothing, + /// 同一个键写了两遍，或者不像 JSON 对象：剪的结果说不准和解析出来的一样，交给 + /// [`reserialize`] + Unusual, +} + +/// 顶层对象里 `parent` 那个对象，剪掉它的 `keys` 成员（连同分隔它的逗号）；剪空了而 +/// `drop_empty` 的话，`parent` 整个剪掉。别的字节原样留着。 +/// +/// 只走顶层和 `parent` 这两层：别的值整个跳过，字符串里按字节找下一个引号或反斜杠。键里 +/// 写了转义的（`"user\u005fid"`）照解析出来的样子比，和 serde_json 认的一样。 +fn cut_members(body: &[u8], parent: &str, keys: &[&str], drop_empty: bool) -> Cut { + let Some(top) = json::object(body, json::ws(body, 0)) else { + return Cut::Unusual; + }; + let mut parents = top.iter().filter(|m| m.is(body, parent)); + let (Some(p), None) = (parents.next(), parents.next()) else { + // 一个都没有就没得剪；有两个的话 serde_json 留后一个，交给它 + return if top.iter().any(|m| m.is(body, parent)) { + Cut::Unusual + } else { + Cut::Nothing + }; + }; + if body[p.value.start] != b'{' { + return Cut::Nothing; + } + let Some(inner) = json::object(body, p.value.start) else { + return Cut::Unusual; + }; + let gone: Vec = (0..inner.len()) + .filter(|&i| keys.iter().any(|k| inner[i].is(body, k))) + .collect(); + if gone.is_empty() { + return Cut::Nothing; + } + // 同一个键写了两遍 + if keys + .iter() + .any(|k| gone.iter().filter(|&&i| inner[i].is(body, k)).count() > 1) + { + return Cut::Unusual; + } + let cuts = if drop_empty && gone.len() == inner.len() { + let at = top + .iter() + .position(|m| m.key.start == p.key.start) + .unwrap_or_default(); + json::cuts(&top, &[at]) + } else { + json::cuts(&inner, &gone) + }; + let mut out = Vec::with_capacity(body.len()); + let mut at = 0; + for c in cuts { + out.extend_from_slice(&body[at..c.start]); + at = c.end; + } + out.extend_from_slice(&body[at..]); + Cut::Done(out) +} + +/// 在 JSON 原文上认出一个对象的成员，好按字节剪。**只给解得开的 JSON 用**：写坏了的地方 +/// 不报错，只保证不越界。 +mod json { + use std::ops::Range; + + /// 对象里的一个成员：键（连同引号）和值在原文里的区间 + #[derive(Debug, Clone)] + pub(super) struct Member { + pub(super) key: Range, + pub(super) value: Range, + } + + impl Member { + /// 键解出来是不是 `name` + pub(super) fn is(&self, b: &[u8], name: &str) -> bool { + let raw = &b[self.key.start + 1..self.key.end - 1]; + if !raw.contains(&b'\\') { + return raw == name.as_bytes(); + } + serde_json::from_slice::(&b[self.key.clone()]).is_ok_and(|k| k == name) + } + } + + /// 从 `i` 起跳过空白 + pub(super) fn ws(b: &[u8], mut i: usize) -> usize { + while i < b.len() && matches!(b[i], b' ' | b'\t' | b'\n' | b'\r') { + i += 1; + } + i + } + + /// `b[i]` 是一个字符串的开头引号：返回结尾引号之后的位置 + fn string_end(b: &[u8], i: usize) -> Option { + if b.get(i) != Some(&b'"') { + return None; + } + let mut j = i + 1; + loop { + j += memchr::memchr2(b'"', b'\\', b.get(j..)?)?; + if b[j] == b'"' { + return Some(j + 1); + } + j += 2; + } + } + + /// `b[i]` 起的一个值：返回它之后的位置 + fn value_end(b: &[u8], i: usize) -> Option { + match *b.get(i)? { + b'"' => string_end(b, i), + b'{' | b'[' => { + let mut depth = 0usize; + let mut j = i; + loop { + match *b.get(j)? { + b'"' => { + j = string_end(b, j)?; + continue; + } + b'{' | b'[' => depth += 1, + b'}' | b']' => { + depth -= 1; + if depth == 0 { + return Some(j + 1); + } + } + _ => {} + } + j += 1; + } + } + // 数、true、false、null：到下一个分隔为止 + _ => { + let end = b[i..] + .iter() + .position(|c| matches!(c, b',' | b'}' | b']' | b' ' | b'\t' | b'\n' | b'\r')) + .map_or(b.len(), |n| i + n); + (end > i).then_some(end) + } + } + } + + /// `b[open]` 是一个对象的 `{`：它的成员，按先后 + pub(super) fn object(b: &[u8], open: usize) -> Option> { + if b.get(open) != Some(&b'{') { + return None; + } + let mut out = Vec::new(); + let mut i = ws(b, open + 1); + if b.get(i) == Some(&b'}') { + return Some(out); + } + loop { + let key = i..string_end(b, i)?; + let colon = ws(b, key.end); + if b.get(colon) != Some(&b':') { + return None; + } + let v = ws(b, colon + 1); + let value = v..value_end(b, v)?; + let next = ws(b, value.end); + out.push(Member { key, value }); + match *b.get(next)? { + b',' => i = ws(b, next + 1), + b'}' => return Some(out), + _ => return None, + } + } + } + + /// 剪掉 `members` 里第 `gone` 几个（按先后）要剪的区间，按先后、互不重叠。剪完还是 + /// 合法的 JSON：剪掉一个成员连同它后面的逗号；排在最后的，连同它前面的逗号 + pub(super) fn cuts(members: &[Member], gone: &[usize]) -> Vec> { + let mut out = Vec::new(); + let mut k = 0; + while k < gone.len() { + // 连着要剪的一串 + let first = gone[k]; + let mut last = first; + while k + 1 < gone.len() && gone[k + 1] == last + 1 { + k += 1; + last += 1; + } + k += 1; + out.push(if last + 1 < members.len() { + members[first].key.start..members[last + 1].key.start + } else if first > 0 { + members[first - 1].value.end..members[last].value.end + } else { + members[first].key.start..members[last].value.end + }); + } + out + } +} + #[cfg(test)] mod tests { use super::*; @@ -437,4 +663,220 @@ mod tests { ); assert_eq!(strip_body_identity(Dialect::Chat, &body), None); } + + /// 去掉 `user_id` 之后，请求体别的字节一个不动：键的先后（工具定义、工具参数、整个请求）、 + /// 空白、转义都照客户端写的。以前解析再写回，每个对象的键都按字母重排了 + #[test] + fn the_forwarded_body_is_the_client_body_less_the_identity() { + let tools = r#""tools":[{"name":"Read","description":"读文件 \"quoted\" \\ path","input_schema":{"type":"object","properties":{"file_path":{"type":"string"},"offset":{"type":"number"},"limit":{"type":"number"}},"required":["file_path"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}}]"#; + let messages = r#""messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_01","content":"{\"user_id\":\"not this one\"} }]"}]},{"role":"assistant","content":[{"type":"tool_use","id":"toolu_02","name":"Read","input":{"zeta":1,"alpha":[2,{"b":3,"a":4}]}}]}]"#; + let identity = r#""metadata":{"user_id":"user_abc_account_123_session_456"}"#; + for (before, after) in [ + // Claude Code 的顺序：metadata 在工具定义后面 + ( + format!( + r#"{{"model":"claude-sonnet-5",{messages},"system":"s",{tools},{identity},"max_tokens":32000,"stream":true}}"# + ), + format!( + r#"{{"model":"claude-sonnet-5",{messages},"system":"s",{tools},"max_tokens":32000,"stream":true}}"# + ), + ), + // 排在最前、最后 + ( + format!(r#"{{{identity},{messages}}}"#), + format!(r#"{{{messages}}}"#), + ), + ( + format!(r#"{{{messages},{identity}}}"#), + format!(r#"{{{messages}}}"#), + ), + // 带空白、换行的 + ( + format!("{{\n {messages},\n {identity} ,\n \"stream\": true\n}}"), + format!("{{\n {messages},\n \"stream\": true\n}}"), + ), + // metadata 里还有别的：只剪 user_id + ( + r#"{"metadata":{"user_id":"u","note":"x"},"messages":[]}"#.to_string(), + r#"{"metadata":{"note":"x"},"messages":[]}"#.to_string(), + ), + ( + r#"{"metadata": {"note": "x", "user_id": "u"}, "messages": []}"#.to_string(), + r#"{"metadata": {"note": "x"}, "messages": []}"#.to_string(), + ), + ( + r#"{"metadata":{"a":1,"user_id":{"nested":["}"]},"b":2}}"#.to_string(), + r#"{"metadata":{"a":1,"b":2}}"#.to_string(), + ), + // 键写成转义的：serde_json 认它是 user_id，这里也认（原文里别处得写着 `"user_id"`， + // 先看的那一眼才放它过去，和以前一样）。反斜杠由 `char::from(92)` 拼：测试里直接 + // 写出来的转义，经过某些编辑工具会变成真字符 + ( + format!( + r#"{{"metadata":{{"user{b}u005fid":"u"}},"user_id":1}}"#, + b = char::from(92) + ), + r#"{"user_id":1}"#.to_string(), + ), + ] { + let out = strip_body_identity(Dialect::Anthropic, &Bytes::from(before.clone())) + .unwrap_or_else(|| panic!("nothing stripped from {before}")); + assert_eq!(std::str::from_utf8(&out).unwrap(), after); + // 和解析、去掉、写回的意思一样 + assert_eq!( + serde_json::from_slice::(&out).unwrap(), + serde_json::from_slice::( + &reserialize(Dialect::Anthropic, &Bytes::from(before)).unwrap() + ) + .unwrap() + ); + } + + let codex = r#"{"model":"gpt-5.5","input":[],"client_metadata":{"x-codex-turn-metadata":"{}","session_id":"conv-1","x-codex-installation-id":"inst-1"},"stream":true}"#; + let out = strip_body_identity(Dialect::Responses, &Bytes::from(codex)).unwrap(); + assert_eq!( + std::str::from_utf8(&out).unwrap(), + r#"{"model":"gpt-5.5","input":[],"client_metadata":{"session_id":"conv-1"},"stream":true}"# + ); + // 剪空了也留着这个对象，和以前一样 + let codex = r#"{"client_metadata":{"x-codex-installation-id":"inst-1","x-codex-turn-metadata":"{}"}}"#; + let out = strip_body_identity(Dialect::Responses, &Bytes::from(codex)).unwrap(); + assert_eq!( + std::str::from_utf8(&out).unwrap(), + r#"{"client_metadata":{}}"# + ); + } + + /// 同一个键写了两遍：剪的结果说不准和 serde_json 认的一样（它留后一个），照旧解析、 + /// 去掉、写回 + #[test] + fn a_key_written_twice_falls_back_to_reserializing() { + for body in [ + r#"{"metadata":{"user_id":"a"},"x":1,"metadata":{"user_id":"b","n":1}}"#, + r#"{"metadata":{"user_id":"a","user_id":"b"},"x":1}"#, + ] { + let body = Bytes::from(body); + assert_eq!( + strip_body_identity(Dialect::Anthropic, &body), + reserialize(Dialect::Anthropic, &body) + ); + } + assert_eq!( + cut_members( + br#"{"metadata":1,"metadata":{"user_id":"b"}}"#, + "metadata", + &["user_id"], + true + ), + Cut::Unusual + ); + } + + /// 随机拼出来的请求：剪出来的和解析、去掉、写回的意思一样，剪没剪也一样；剪过的那一份 + /// 只少了几段，别的字节按原样、原来的先后都在 + #[test] + fn cutting_means_what_reserializing_meant() { + struct Rng(u64); + impl Rng { + fn below(&mut self, n: usize) -> usize { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 % n as u64) as usize + } + fn pick<'a>(&mut self, xs: &[&'a str]) -> &'a str { + xs[self.below(xs.len())] + } + } + let mut r = Rng(0x51_7cc1_b727_220a); + const KEYS: &[&str] = &[ + "\"user_id\"", + "\"user\\u005fid\"", + "\"note\"", + "\"a\"", + "\"x-codex-installation-id\"", + "\"x-codex-turn-metadata\"", + "\"session_id\"", + ]; + const VALUES: &[&str] = &[ + "1", + "-2.5e3", + "true", + "null", + "\"s\"", + "\"q\\\"}]{[,\\\\\"", + "\"\\\"user_id\\\"\"", + "[]", + "{}", + "[1,{\"user_id\":\"deep\"},\"]\"]", + "{\"metadata\":{\"user_id\":\"deep\"}}", + ]; + const WS: &[&str] = &["", " ", "\n ", "\t"]; + let inner = |r: &mut Rng| { + let members: Vec = (0..r.below(4)) + .map(|_| { + let (k, a, b, v) = (r.pick(KEYS), r.pick(WS), r.pick(WS), r.pick(VALUES)); + format!("{k}{a}:{b}{v}") + }) + .collect(); + format!("{{{}}}", members.join(",")) + }; + for _ in 0..5_000 { + let mut top = Vec::new(); + for _ in 0..r.below(5) { + let (k, v) = match r.below(4) { + 0 => ("\"metadata\"", inner(&mut r)), + 1 => ("\"client_metadata\"", inner(&mut r)), + 2 => ("\"messages\"", r.pick(VALUES).to_string()), + _ => (r.pick(KEYS), r.pick(VALUES).to_string()), + }; + let (a, b, c, d) = (r.pick(WS), r.pick(WS), r.pick(WS), r.pick(WS)); + top.push(format!("{a}{k}{b}:{c}{v}{d}")); + } + let (a, b) = (r.pick(WS), r.pick(WS)); + let body = Bytes::from(format!("{a}{{{}}}{b}", top.join(","))); + if serde_json::from_slice::(&body).is_err() { + continue; + } + for d in [Dialect::Anthropic, Dialect::Responses] { + let cut = strip_body_identity(d, &body); + let old = reserialize(d, &body); + // 旧的写法先看一眼原文里有没有那几个字：没有就不解析 + let marked = match d { + Dialect::Anthropic => memchr::memmem::find(&body, b"\"user_id\"").is_some(), + _ => crate::chatgpt::IDENTITY_METADATA + .iter() + .any(|m| memchr::memmem::find(&body, m.as_bytes()).is_some()), + }; + let old = old.filter(|_| marked); + let value = |b: &Option| { + b.as_ref() + .map(|b| serde_json::from_slice::(b).unwrap()) + }; + assert_eq!( + value(&cut), + value(&old), + "{d:?} {}", + String::from_utf8_lossy(&body) + ); + // 剪过的那一份是原文的一个子序列：别的字节原样、按原来的先后 + if let Some(cut) = &cut + && cut_members(&body, "metadata", &["user_id"], true) != Cut::Unusual + && cut_members( + &body, + "client_metadata", + &crate::chatgpt::IDENTITY_METADATA, + false, + ) != Cut::Unusual + { + let mut rest = body.iter(); + assert!( + cut.iter().all(|c| rest.any(|b| b == c)), + "{}", + String::from_utf8_lossy(&body) + ); + } + } + } + } } diff --git a/crates/tw-gateway/src/guard.rs b/crates/tw-gateway/src/guard.rs index d4eb6c2c..72bac1ed 100644 --- a/crates/tw-gateway/src/guard.rs +++ b/crates/tw-gateway/src/guard.rs @@ -140,6 +140,16 @@ pub fn screen(s: &Screen, dialect: tw_dialect::ir::Dialect, body: &[u8]) -> Scre tw_guard::content::screen(s.mode, &s.rules, dialect, body) } +/// [`screen`]，请求体已经解析好了：`v` 是它解出来的样子（见 +/// [`tw_guard::content::screen_value`]）。 +pub fn screen_value( + s: &Screen, + dialect: tw_dialect::ir::Dialect, + v: &serde_json::Value, +) -> Screening { + tw_guard::content::screen_value(s.mode, &s.rules, dialect, v) +} + /// 把一次查下来的结论报出去：每条命中的规则一条 [`tw_api::Event::ContentMatched`]，挂在 /// 请求 `id` 上。要拒绝时返回告诉客户端的那句话。 pub fn report( diff --git a/crates/tw-gateway/src/server/pipeline.rs b/crates/tw-gateway/src/server/pipeline.rs index 07f56865..972433f2 100644 --- a/crates/tw-gateway/src/server/pipeline.rs +++ b/crates/tw-gateway/src/server/pipeline.rs @@ -71,7 +71,10 @@ pub(super) async fn pipeline( live: crate::live::Pass, ending: &mut Option, ) -> Result { - let intent = match probe(&state, &rt, &req) { + // 请求体**只解析一次**：本地应答的判定、路由事实、对话的指纹、内容过滤都用这一份（见 + // `read`）。解不开是 None —— 照样往下走，同格式直通照样发 + let parsed = serde_json::from_slice::(&req.body).ok(); + let intent = match probe(&state, &rt, &req, parsed.as_ref()) { Probe::Answered(resp) => return Ok(resp), Probe::Intent(intent) => intent, }; @@ -89,7 +92,7 @@ pub(super) async fn pipeline( // 管线第 2 步：读出路由事实，路由。**看的是客户端的原话**：插件的请求钩子排在路由 // 之后（每发往一个上游跑一次，见 `plug`），左右不了请求去哪一家 - let (mut reading, fp) = read(&req, intent); + let (mut reading, fp) = read(&req, intent, parsed.as_ref()); let conv = conversation(&rt, &req, &reading, fp.as_deref()); let (choice, decision) = match route(&state, &rt, &req, &reading, conv.as_ref())? { Routed::Go(choice, decision) => (choice, decision), @@ -137,7 +140,7 @@ pub(super) async fn pipeline( .await?; // 管线第 4 步：内容过滤先下结论，不发事件。删过的话，后面一律用删过的那一份 - let screening = screen(&rt, &mut req, &mut reading); + let screening = screen(&rt, &mut req, &mut reading, parsed.as_ref()); let started = start( &state, &rt, @@ -291,10 +294,17 @@ enum Probe { /// 于是断网时健康检查照样失败，白白丢掉这个功能最有价值的场景。 /// 同样的理由让它排在「一个 provider 都没有」那一条之前：那一条也是 /// 一种「没有可用上游」，而本地应答本来就不需要上游。 -fn probe(state: &AppState, rt: &Runtime, req: &Inbound) -> Probe { - let Some(kind) = - crate::clientprobe::classify(&req.body, is_claude_code(&req.client_name, &req.headers)) - else { +fn probe( + state: &AppState, + rt: &Runtime, + req: &Inbound, + parsed: Option<&serde_json::Value>, +) -> Probe { + let Some(kind) = crate::clientprobe::classify( + &req.body, + parsed, + is_claude_code(&req.client_name, &req.headers), + ) else { return Probe::Intent(String::new()); }; use tw_config::ProbeAction::*; @@ -321,27 +331,30 @@ fn probe(state: &AppState, rt: &Runtime, req: &Inbound) -> Probe { /// 管线第 2 步的前半：读出路由事实，和这段对话的指纹。 /// -/// **只解析一次，指纹也只算一次。**路由要它，认对话（见 `crate::affinity`）要指纹， -/// 开始事件归会话也要指纹，而 body 可能有几百 KB —— 解两遍、哈希两遍是白付一份钱。 +/// **只解析一次，指纹也只算一次。**`parsed` 是管线开头解出来的那一份（本地应答的判定、 +/// 内容过滤用的也是它）。路由要它，认对话（见 `crate::affinity`）要指纹，开始事件归会话 +/// 也要指纹，而 body 可能有几 MB —— 解两遍、哈希两遍是白付一份钱。 /// /// 生成回答的请求解码成中间表示，**四种格式的客户端读出同一份路由事实**。 /// body 解不开时用空的性质走兜底规则。**不要因此拒绝请求** —— 我们的解析器 /// 不认识的东西，上游可能完全认识（只有需要转换时才用得上解码结果） -fn read(req: &Inbound, intent: String) -> (crate::client_api::Reading, Option) { - let parsed = serde_json::from_slice::(&req.body).ok(); - let mut reading = - crate::client_api::read(req.uri.path(), req.query.as_deref(), parsed.as_ref()); +fn read( + req: &Inbound, + intent: String, + parsed: Option<&serde_json::Value>, +) -> (crate::client_api::Reading, Option) { + let mut reading = crate::client_api::read(req.uri.path(), req.query.as_deref(), parsed); reading.facts.client = req.client_name.clone(); reading.facts.intent = intent; reading.harness = tw_dialect::harness::detect( req.headers .get(axum::http::header::USER_AGENT) .and_then(|v| v.to_str().ok()), - parsed.as_ref(), + parsed, ); // 认出「这几十个请求是同一次任务」。**认不出来就是 None** —— 硬凑一个会把 // 互不相干的请求并成一个「会话」 - let fp = parsed.as_ref().and_then(crate::session::fingerprint); + let fp = parsed.and_then(crate::session::fingerprint); (reading, fp) } @@ -965,17 +978,21 @@ fn open( /// 模型，查了只会在真正的请求之前把同一处命中多记一遍、还可能把计数请求拒掉（理由见 /// [`crate::client_api::ClientApi::screened`]）。在原文上查，中间表示解不开的请求照样查 /// （同格式直通照样发，上游可能认得它）。 +/// +/// `parsed` 是请求体解析出来的样子（管线开头解的那一份），解不开的是 None：解不开的请求体 +/// 没有调用方的正文可言，什么都不报。 fn screen( rt: &Runtime, req: &mut Inbound, reading: &mut crate::client_api::Reading, + parsed: Option<&serde_json::Value>, ) -> tw_guard::content::Screening { let screened = crate::client_api::ClientApi::screened(req.uri.path()); - let Some(api) = req.api.filter(|_| screened) else { + let (Some(api), Some(parsed)) = (req.api.filter(|_| screened), parsed) else { return Default::default(); }; let dialect = api.dialect(); - let sc = crate::guard::screen(&crate::guard::Screen::of(rt), dialect, &req.body); + let sc = crate::guard::screen_value(&crate::guard::Screen::of(rt), dialect, parsed); if let Some(body) = &sc.body { req.body = body.clone(); if reading.decoded.is_some() { diff --git a/crates/tw-gateway/src/server/pipeline/hop.rs b/crates/tw-gateway/src/server/pipeline/hop.rs index ece5c1de..bd6442d1 100644 --- a/crates/tw-gateway/src/server/pipeline/hop.rs +++ b/crates/tw-gateway/src/server/pipeline/hop.rs @@ -64,6 +64,9 @@ struct Asked<'r> { body: &'r Bytes, path: &'r str, decoded: Option<&'r Result>, + /// `body` 解得开（是 JSON）：客户端的原话照管线开头解的那一遍，插件改过的那一份是 + /// 网关自己写出来的，一定解得开 + json: bool, } impl<'r> Asked<'r> { @@ -77,11 +80,13 @@ impl<'r> Asked<'r> { body: &r.body, path: &r.path, decoded: r.decoded.as_ref(), + json: true, }, None => Asked { body: &req.body, path: req.uri.path(), decoded: reading.decoded.as_ref(), + json: reading.json, }, } } @@ -1133,6 +1138,7 @@ fn successor<'r>( body: &req.body, path: req.uri.path(), decoded: reading.decoded.as_ref(), + json: reading.json, }; rest.any(|name| { let Some(provider) = rt.config.providers.iter().find(|p| &p.name == name) else { @@ -1232,8 +1238,11 @@ fn prepare( } else if provider.forward_client_identity { out } else { - // 客户端自动填的身份字段不发（见 `egress` 模块）。**不算丢弃的字段** + // 客户端自动填的身份字段不发（见 `egress` 模块）。**不算丢弃的字段**。按字节 + // 剪，要请求体解得开：上面几步要么原样交回、要么照解出来的写回，解不开的照旧 + // 解不开、原样发 client_dialect + .filter(|_| asked.json) .and_then(|d| crate::egress::strip_body_identity(d, &out)) .unwrap_or(out) }; diff --git a/crates/tw-gateway/tests/headers_and_paths.rs b/crates/tw-gateway/tests/headers_and_paths.rs index 764876bd..843e08a0 100644 --- a/crates/tw-gateway/tests/headers_and_paths.rs +++ b/crates/tw-gateway/tests/headers_and_paths.rs @@ -440,6 +440,34 @@ async fn an_anthropic_upstream_gets_the_request_and_nothing_about_the_client() { assert_eq!(v["messages"][0]["content"], "hi"); } +/// 去掉 `user_id` 之后，上游收到的就是客户端发的那些字节：键的先后（工具定义、工具参数、 +/// 整个请求）、空白、转义一个都没动。以前解析再写回，每个对象的键都按字母重排了 —— 而 +/// Claude Code 的请求每一个都带着 `user_id`。 +/// +/// 体过了 1 MiB、跑在多线程的运行时上：这么大的请求体，管线上那几步是挪出异步线程算的 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn the_upstream_gets_the_client_body_byte_for_byte_less_the_identity() { + let (up, seen) = start_upstream().await; + let gw = start_gateway(anthropic(up)).await; + let file = "fn main() {\n println!(\"读了 \\\"一个\\\" 文件\");\n}\n".repeat(20_000); + let file = serde_json::to_string(&file).unwrap(); + let head = format!( + r#"{{"model":"claude-sonnet-4-5","messages":[{{"role":"user","content":[{{"type":"tool_result","tool_use_id":"toolu_01","content":{file}}}]}},{{"role":"assistant","content":[{{"type":"tool_use","id":"toolu_02","name":"Read","input":{{"zeta":1,"alpha":[2,{{"b":3,"a":4}}]}}}}]}}], "tools":[{{"name":"Read","input_schema":{{"type":"object","properties":{{"file_path":{{"type":"string"}},"limit":{{"type":"number"}}}},"required":["file_path"],"$schema":"http://json-schema.org/draft-07/schema#"}}}}],"# + ); + let identity = r#""metadata":{"user_id":"user_abc_account_123_session_456"},"#; + let tail = r#""max_tokens":16,"stream":false}"#; + let body = format!("{head}{identity}{tail}"); + assert!(body.len() > 1024 * 1024, "{}", body.len()); + post(gw, "/v1/messages", CLAUDE_CODE, &body).await; + + let g = seen.lock().unwrap(); + assert!( + g.body == format!("{head}{tail}").as_bytes(), + "the upstream got a different body: {}", + String::from_utf8_lossy(&g.body[g.body.len().saturating_sub(300)..]) + ); +} + #[tokio::test] async fn an_upstream_that_asks_for_it_gets_the_clients_own_identity() { let (up, seen) = start_upstream().await; diff --git a/crates/tw-guard/src/content.rs b/crates/tw-guard/src/content.rs index 54e96a3a..12ceba1d 100644 --- a/crates/tw-guard/src/content.rs +++ b/crates/tw-guard/src/content.rs @@ -34,7 +34,7 @@ mod screen; pub(crate) use codepoints::visible as codepoints_visible; pub use codepoints::{CodepointError, Codepoints, MAX_ITEMS as MAX_CODEPOINT_ITEMS}; pub(crate) use screen::evaluate; -pub use screen::{Outcome, ScreenHit, Screening, screen, screen_text}; +pub use screen::{Outcome, ScreenHit, Screening, screen, screen_text, screen_value}; /// 一条规则怎么认。 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize)] diff --git a/crates/tw-guard/src/content/screen.rs b/crates/tw-guard/src/content/screen.rs index 91337f92..bd7bc38f 100644 --- a/crates/tw-guard/src/content/screen.rs +++ b/crates/tw-guard/src/content/screen.rs @@ -88,17 +88,26 @@ pub fn screen(mode: Mode, rules: &Rules, dialect: Dialect, body: &[u8]) -> Scree if !mode.detects() || rules.is_empty() { return Screening::default(); } - let Ok(mut v) = serde_json::from_slice::(body) else { + let Ok(v) = serde_json::from_slice::(body) else { return Screening::default(); }; - let spots = tw_dialect::caller::spots(dialect, &v); + screen_value(mode, rules, dialect, &v) +} + +/// [`screen`]，请求体已经解析好了（调用方为别的事解过一遍）：`v` 是它解出来的样子，不再 +/// 解一遍。删过的请求体照它写出来，和 [`screen`] 写的一样。 +pub fn screen_value(mode: Mode, rules: &Rules, dialect: Dialect, v: &Value) -> Screening { + if !mode.detects() || rules.is_empty() { + return Screening::default(); + } + let spots = tw_dialect::caller::spots(dialect, v); if spots.is_empty() { return Screening::default(); } let e = { let segments: Vec<(&str, bool)> = spots .iter() - .map(|s| (s.get(&v).unwrap_or_default(), s.in_tool_result)) + .map(|s| (s.get(v).unwrap_or_default(), s.in_tool_result)) .collect(); evaluate(rules, &segments, Scope::default()) }; @@ -107,6 +116,8 @@ pub fn screen(mode: Mode, rules: &Rules, dialect: Dialect, body: &[u8]) -> Scree && out.refused.is_none() && let Some(texts) = e.texts { + // 要删的时候才抄一份来改 + let mut v = v.clone(); for (s, t) in spots.iter().zip(texts) { if let Some(x) = s.get_mut(&mut v) { *x = t; From cfd48f8015c38e690854e9141629a9c8da8faa23 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 21:08:47 +0800 Subject: [PATCH 12/16] perf(gateway): reuse the request-path secret scan The request body was scanned for secrets on the request path and then, for storage, scanned again in full on the blocking pool (the scan plus `mask_body` was 29 ms per 2 MiB). In enforce mode each hop scanned again too, even when the hop sends exactly the client's bytes. - The hits from the request-path look travel with the stored request body (`BodyRecord::found`). They are used only when the record holds the whole body and it is UTF-8, so the text masked before storage is the text that was scanned; a body cut to the storage window is scanned again as before. - In enforce mode a hop whose body is the very buffer the client sent (same-format passthrough, nothing rewritten) is replaced from those hits (`guard::replace_found`). Any other hop body -- translated, rewritten by a rule or plugin, identity cut -- is scanned as before. Same rules, same bytes, same hits: tests check both paths against a fresh scan. Co-Authored-By: Claude Opus 5.5 --- crates/tw-gateway/src/bodies.rs | 80 ++++++++++++++++++- crates/tw-gateway/src/guard.rs | 66 +++++++++++++++ crates/tw-gateway/src/server/pipeline.rs | 40 ++++++---- .../src/server/pipeline/admission.rs | 4 +- crates/tw-gateway/src/server/pipeline/hop.rs | 11 ++- 5 files changed, 177 insertions(+), 24 deletions(-) diff --git a/crates/tw-gateway/src/bodies.rs b/crates/tw-gateway/src/bodies.rs index 1c83f322..df5c11aa 100644 --- a/crates/tw-gateway/src/bodies.rs +++ b/crates/tw-gateway/src/bodies.rs @@ -25,7 +25,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use bytes::Bytes; use tw_guard::redact::replace::{Ledger, Scheme}; -use tw_guard::redact::rules::{Kind, Rule, RuleSet}; +use tw_guard::redact::rules::{Hit, Kind, Rule, RuleSet}; /// 打码要多看的那一截。 /// @@ -101,12 +101,25 @@ impl Redaction { /// /// **读的时候还会再打一遍**，打第二遍不再改动什么（`mask_body` 认得自己打过的码）。 pub fn apply(&self, text: &str) -> String { + self.apply_found(text, None) + } + + /// [`apply`](Self::apply)，`found` 是按这套规则在 `text` 上已经找过的命中（[`crate::guard::hits`] + /// 的结果）：不再找一遍。没找过的是 `None` + fn apply_found(&self, text: &str, found: Option<&[Hit]>) -> String { // 我们自己的占位符不算（见 `crate::guard::hits`）：回答里回显的、重放过来的，原样存 - let hits = crate::guard::hits(text, &self.rules); + let fresh; + let hits = match found { + Some(hits) => hits, + None => { + fresh = crate::guard::hits(text, &self.rules); + &fresh[..] + } + }; let sent: HashMap<&str, &str> = self.ledger.replacements().collect(); let mut out = String::with_capacity(text.len()); let mut at = 0; - for h in &hits { + for h in hits { out.push_str(&text[at..h.bytes.start]); let value = &text[h.bytes.clone()]; match sent.get(value) { @@ -148,6 +161,8 @@ pub struct BodyRecord { pub original_len: usize, /// 落盘之前怎么换、怎么打码 pub redaction: Redaction, + /// 请求路上按 `redaction` 的规则在这份正文上已经找过的命中（见 [`Self::found`]） + found: Option>, /// 占着的那份额度（见 [`QUEUED_MAX`]） held: Option, } @@ -168,10 +183,21 @@ impl BodyRecord { body, original_len, redaction, + found: None, held: None, } } + /// 请求路上出站脱敏已经按 `redaction` 的规则在 `body` 上找过一遍（见 + /// [`crate::guard::look_hits`]）：落盘前打码用它找到的，不再把整份正文扫一遍。 + /// + /// 只在交来的是整份正文、它又是 UTF-8 时用得上 —— 那时落盘前看的文字就是请求路上看的 + /// 那一份；截过的、不是 UTF-8 的照旧再找（[`Self::for_disk`]）。 + pub fn found(mut self, hits: Option>) -> Self { + self.found = hits; + self + } + /// 落盘的那一份：换过、打过码（[`Redaction::apply`]），和要记下的原本长度。 /// /// **在阻塞线程上调**：一份 4 MB 的正文要扫好几遍。 @@ -180,7 +206,13 @@ impl BodyRecord { /// 那样，而二进制的正文里没有能看的东西。 pub fn for_disk(self) -> ForDisk { let window = &self.body[..self.body.len().min(WINDOW)]; - let body = self.redaction.apply(&String::from_utf8_lossy(window)); + let text = String::from_utf8_lossy(window); + // 请求路上找过的那一遍看的就是这些字：整份都在窗口里，原文就是 UTF-8 + let entire = window.len() == self.body.len() && self.original_len <= self.body.len(); + let found = self + .found + .filter(|_| entire && matches!(text, std::borrow::Cow::Borrowed(_))); + let body = self.redaction.apply_found(&text, found.as_deref()); // 截过的（交来的只是开头）报原本的长度。没截过的就是换过、打过码的这一份的长度： // 比存储层的上限还长的，由存储层截、由它记下（`tw_store::Blobs::put_with_len`） let whole = self.original_len.max(self.body.len()); @@ -266,6 +298,8 @@ pub fn offer(sink: &Option, mut rec: BodyRecord) { let Some(s) = sink else { return }; if rec.body.len() > WINDOW { rec.body = Bytes::copy_from_slice(&rec.body[..WINDOW]); + // 找过的是整份；落盘前看的只是开头这一截，要重新找（见 `BodyRecord::found`） + rec.found = None; } let n = rec.body.len(); // 先占额度，占不下就丢 @@ -461,6 +495,44 @@ mod tests { serde_json::from_str::(&stored).expect("存下来的还是 JSON"); } + /// 请求路上找过的命中拿来打码，落盘的和自己再找一遍的一字不差；交来的只是开头一截时 + /// 不用它（它说的是整份），照旧自己找 + #[test] + fn the_hits_found_on_the_way_in_are_used_only_for_the_whole_body() { + let body = format!( + r#"{{"messages":[{{"role":"user","content":"key {KEY} 库 postgres://app:hunter2hunter2@db/x"}}]}}"# + ); + for mode in [ + tw_config::SecurityMode::Observe, + tw_config::SecurityMode::Enforce, + ] { + let rules = RuleSet::defaults(); + let seen = crate::guard::look_hits(mode, &rules, body.as_bytes()); + assert_eq!(seen.hits.as_ref().map(Vec::len), Some(2)); + let r = Redaction { + rules: Arc::new(rules), + ledger: seen.ledger, + }; + let again = written(record(BodyKind::Request, &body, r.clone())); + let reused = + written(record(BodyKind::Request, &body, r).found(seen.hits.map(Arc::from))); + assert_eq!(reused, again); + assert!(!reused.contains(KEY), "{reused}"); + } + // 说「什么都没找到」的命中：整份交来时照它（证明真的用上了），截过的不照它 + let lie = || Some(Arc::from(Vec::new())); + let whole = written(record(BodyKind::Request, &body, Redaction::default()).found(lie())); + assert!(whole.contains("sk-an…AAAA"), "{whole}"); + assert!( + whole.contains("hunter2hunter2"), + "规则那一道真的没再找：{whole}" + ); + let mut cut = record(BodyKind::Request, &body, Redaction::default()).found(lie()); + cut.original_len = body.len() + 1; + let cut = written(cut); + assert!(!cut.contains("hunter2hunter2"), "{cut}"); + } + #[test] fn what_the_rules_miss_is_masked_by_its_shape() { // 网关自己的钥匙（`tw-`）没有一条脱敏规则认：读的时候那一道兜住它，写的时候也一样 diff --git a/crates/tw-gateway/src/guard.rs b/crates/tw-gateway/src/guard.rs index 72bac1ed..b0379ae8 100644 --- a/crates/tw-gateway/src/guard.rs +++ b/crates/tw-gateway/src/guard.rs @@ -51,6 +51,17 @@ pub fn look(mode: Mode, rules: &RuleSet, body: &[u8]) -> (Vec, Ledger) flow::look(mode, rules, body) } +/// [`look`]，连同找到的命中（见 [`flow::Look::hits`]）：同一份正文落盘前打码时不再找一遍 +/// （[`crate::bodies::BodyRecord::found`]）。 +pub fn look_hits(mode: Mode, rules: &RuleSet, body: &[u8]) -> flow::Look { + flow::look_hits( + mode, + rules, + body, + Ledger::new(tw_guard::redact::replace::Scheme::SECRET), + ) +} + /// [`look`]，接着 `seed` 的账编号（见 [`flow::look_from`]）。 pub fn look_from(mode: Mode, rules: &RuleSet, body: &[u8], seed: Ledger) -> (Vec, Ledger) { flow::look_from(mode, rules, body, seed) @@ -66,6 +77,30 @@ pub fn replace( flow::replace(mode, rules, body, ledger) } +/// [`replace`]，`found` 是按这套规则在 `body` 上已经找过的命中（[`look_hits`] 在同一份字节上 +/// 找的）：照它换，不再找一遍。没找过的是 `None`，和 [`replace`] 一样。 +pub fn replace_found( + mode: Mode, + rules: &RuleSet, + body: bytes::Bytes, + ledger: &Ledger, + found: Option<&[Hit]>, +) -> (bytes::Bytes, Ledger) { + let Some(hits) = found.filter(|_| mode.acts() && !rules.is_empty()) else { + return replace(mode, rules, body, ledger); + }; + // 没命中就原样返回，连一次拷贝都不做 + if hits.is_empty() { + return (body, ledger.clone()); + } + // 找得到命中的一定是 UTF-8（见 `flow::look_hits`） + let Ok(text) = std::str::from_utf8(&body) else { + return replace(mode, rules, body, ledger); + }; + let r = tw_guard::redact::replace::apply(text, hits, ledger.clone()); + (bytes::Bytes::from(r.text), r.ledger) +} + /// `after` 里 `before` 没有的那些值：插件写进请求里的（见 [`crate::plugin::request`]）。 /// /// 按规则和打过码的样子比：同一个值在两份里打出来的码一样。客户端原话里就有的值，开头 @@ -445,6 +480,37 @@ mod tests { assert_eq!(ledger.len(), 150); } + /// 照开头那一遍找到的命中换，和自己再找一遍换的一字不差、账本一样 + #[test] + fn replacing_with_the_hits_already_found_is_replacing() { + let other = "ghp_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; + let body = bytes::Bytes::from(format!( + r#"{{"system":"<> {KEY}","messages":[{{"content":"{other} {KEY}"}}]}}"# + )); + let rules = RuleSet::defaults(); + let seen = look_hits(Mode::Enforce, &rules, &body); + let hits = seen.hits.unwrap(); + let again = replace(Mode::Enforce, &rules, body.clone(), &seen.ledger); + let known = replace_found( + Mode::Enforce, + &rules, + body.clone(), + &seen.ledger, + Some(&hits), + ); + assert_eq!(known.0, again.0); + assert_eq!(known.1.table(), again.1.table()); + assert!(!String::from_utf8_lossy(&known.0).contains(KEY)); + // 观察档不换，找过也不换 + let (out, _) = replace_found(Mode::Observe, &rules, body.clone(), &fresh(), Some(&hits)); + assert_eq!(out, body); + // 什么都没找到：原样，不拷贝 + let plain = bytes::Bytes::from_static(b"{\"messages\":[]}"); + let (out, l) = replace_found(Mode::Enforce, &rules, plain.clone(), &fresh(), Some(&[])); + assert_eq!(out.as_ptr(), plain.as_ptr()); + assert!(l.is_empty()); + } + /// 查表比和两两比，留下的一模一样：同一条规则下打码一样的算见过，换一条规则就不算 #[test] fn more_found_keeps_exactly_what_comparing_every_pair_kept() { diff --git a/crates/tw-gateway/src/server/pipeline.rs b/crates/tw-gateway/src/server/pipeline.rs index 972433f2..5b1bcb07 100644 --- a/crates/tw-gateway/src/server/pipeline.rs +++ b/crates/tw-gateway/src/server/pipeline.rs @@ -59,6 +59,9 @@ struct Started { ledger: tw_guard::redact::replace::Ledger, /// 出站脱敏在客户端原文里找到的。插件改过的那一跳只再报插件写进来的（见 [`plug`]） found: Vec, + /// 拦截档下，出站脱敏在客户端原文（`req.body`）里找到的命中。一跳发出去的就是原文的 + /// 那些字节时（同格式直通、一个字节都没改），换的时候照它换，不再找一遍（见 [`hop`]） + hits: Option>, /// 这个请求最多等到什么时候：准入时定下（见 [`admission`]），等密钥的分钟、小时上限和 /// 等上游的空位共用这一段（`failover.slot_wait_secs`） wait_until: tokio::time::Instant, @@ -101,8 +104,8 @@ pub(super) async fn pipeline( Routed::Refused(choice, why) => { let to = ("", tw_api::Billing::PerToken); // 一个字节都没发出去，也没什么可报的；存下来的请求照样按这一档换、打码 - let (_, ledger) = look(&rt, &req); - let redaction = redaction(&rt, ledger); + let seen = look(&rt, &req); + let redaction = redaction(&rt, seen.ledger); let (id, _) = open( &state, &req, @@ -111,7 +114,7 @@ pub(super) async fn pipeline( to, fp.as_deref(), ending, - redaction, + (redaction, seen.hits.map(Into::into)), ); state.bus.emit(super::routed_nowhere(id, choice)); return Err(why); @@ -830,7 +833,9 @@ fn start( // 是同一条记录**，差别只在换没换 —— 真正的替换在每一跳发出去之前做， // 那一跳的请求体可能是转换过格式的。拦截档下账本在这里就编好号：每一跳、 // 存下来的那份请求都按它换，同一个值处处是同一个占位符 - let (found, ledger) = look(rt, req); + let seen = look(rt, req); + let (found, ledger) = (seen.found, seen.ledger); + let hits: Option> = seen.hits.map(Into::into); let (id, at_ms) = open( state, req, @@ -839,7 +844,7 @@ fn start( (first, billing.into()), fp, ending, - redaction(rt, ledger.clone()), + (redaction(rt, ledger.clone()), hits.clone()), ); let redact_mode = rt.config.security.redact.mode; if !found.is_empty() { @@ -859,6 +864,7 @@ fn start( conversation: crate::affinity::identity(&req.headers, fp), ledger, found, + hits: hits.filter(|_| redact_mode.acts()), wait_until, } } @@ -867,14 +873,11 @@ fn start( /// /// 插件的密钥映射按同一份原文、同一个找法编号（见 [`crate::plugin::bridge`]）：插件看到的 /// 占位符和这本账里的是同一个号。插件往某一跳写进新的值，那一跳接着编（见 [`plug`]）。 -fn look( - rt: &Runtime, - req: &Inbound, -) -> ( - Vec, - tw_guard::redact::replace::Ledger, -) { - crate::guard::look(rt.config.security.redact.mode, &rt.redact, &req.body) +/// +/// 找到的命中跟着请求体交去留档：落盘前打码不再把同一份正文找一遍（见 +/// [`crate::bodies::BodyRecord::found`]）。 +fn look(rt: &Runtime, req: &Inbound) -> tw_guard::redact::flow::Look { + crate::guard::look_hits(rt.config.security.redact.mode, &rt.redact, &req.body) } /// 这个请求的正文落盘之前怎么换、怎么打码：此刻生效的规则，和这个请求的账本。 @@ -887,7 +890,8 @@ fn redaction(rt: &Runtime, ledger: tw_guard::redact::replace::Ledger) -> crate:: /// 发 `RequestStarted`、把这个请求欠着的结局放进 `ending`、把请求体交去留档， /// 交回这个请求的号和开始的时刻。`to` 是要发往的那一家和它怎么收钱；一家都不会去的 -/// （被规则拒绝了）是空的名字。`redaction` 是请求体、响应体落盘之前怎么换、打码。 +/// （被规则拒绝了）是空的名字。`redaction` 是请求体、响应体落盘之前怎么换、打码，和出站 +/// 脱敏在请求体上已经找到的命中（见 [`look`]）。 /// /// **会话在这里定**（见 [`crate::session::Sessions`]）：开始事件带着它，落库的 /// 那一行记的也是它。 @@ -900,7 +904,10 @@ fn open( to: (&str, tw_api::Billing), fp: Option<&str>, ending: &mut Option, - redaction: crate::bodies::Redaction, + (redaction, found): ( + crate::bodies::Redaction, + Option>, + ), ) -> (u64, u64) { let facts = &reading.facts; let id = state.bus.next_id(); @@ -961,7 +968,8 @@ fn open( req.body.clone(), req.body.len(), redaction, - ), + ) + .found(found), ); (id, at_ms) } diff --git a/crates/tw-gateway/src/server/pipeline/admission.rs b/crates/tw-gateway/src/server/pipeline/admission.rs index 41f28cfd..b899de3f 100644 --- a/crates/tw-gateway/src/server/pipeline/admission.rs +++ b/crates/tw-gateway/src/server/pipeline/admission.rs @@ -140,7 +140,7 @@ fn refused( tracing::info!(key = %req.client_name, per = r.limit.per.word(), "a usage limit of the key refused the request"); let why = r.error(); // 一个字节都没发出去；存下来的请求照样按这一档换、打码 - let (_, ledger) = look(rt, req); + let seen = look(rt, req); let (id, _) = open( state, req, @@ -149,7 +149,7 @@ fn refused( ("", tw_api::Billing::PerToken), fp, ending, - redaction(rt, ledger), + (redaction(rt, seen.ledger), seen.hits.map(Into::into)), ); state .bus diff --git a/crates/tw-gateway/src/server/pipeline/hop.rs b/crates/tw-gateway/src/server/pipeline/hop.rs index bd6442d1..b2a7623f 100644 --- a/crates/tw-gateway/src/server/pipeline/hop.rs +++ b/crates/tw-gateway/src/server/pipeline/hop.rs @@ -439,8 +439,15 @@ pub(super) async fn try_upstreams<'a>( .rewritten .as_ref() .map_or(&started.ledger, |r| &r.ledger); - let (body, ledger) = - crate::guard::replace(rt.config.security.redact.mode, &rt.redact, unsealed, seed); + // 发出去的就是客户端原文的那些字节（同格式直通、一个字节都没改）：开头那一遍在 + // 它上面找到的就是这一跳要换的，不再找一遍 + let known = started.hits.as_deref().filter(|_| { + plugged.rewritten.is_none() + && unsealed.as_ptr() == req.body.as_ptr() + && unsealed.len() == req.body.len() + }); + let mode = rt.config.security.redact.mode; + let (body, ledger) = crate::guard::replace_found(mode, &rt.redact, unsealed, seed, known); // 用这个 provider 自己的 Client —— 它带着该走的代理。**在取密钥 // 之前拿到**：OAuth 换 token 也要走这条代理。 From 84ea86606bb01a7bd6b52404cc4d6bd71e3904f8 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 21:16:51 +0800 Subject: [PATCH 13/16] perf(gateway): work on large request bodies off the async workers Parsing, the local-answer check, the routing facts, the content filter, the secret scan and each hop's rewrite all ran inline on the tokio worker that owns the connection. Every other connection on that worker -- a streaming answer being relayed -- stopped for as long: with one worker, SSE gaps of 470-610 ms while 8 MiB requests came in. For bodies of 1 MiB or more these steps now run in `block_in_place`, which hands the worker's other tasks to another thread for the duration. Smaller bodies stay inline, where the hand-off would cost more than the work. On a current-thread runtime (the tests) `block_in_place` panics, so it is used only on a multi-thread one. With one gateway worker and 8 MiB requests arriving alongside an SSE relay, the longest gap between relayed chunks fell from 470-610 ms to 17-39 ms, and the gateway took 45-55 such requests in the time it used to take 8-12. Co-Authored-By: Claude Opus 5.5 --- crates/tw-gateway/src/server/pipeline.rs | 68 ++++++++++++++----- .../src/server/pipeline/admission.rs | 2 +- crates/tw-gateway/src/server/pipeline/hop.rs | 42 ++++++------ 3 files changed, 76 insertions(+), 36 deletions(-) diff --git a/crates/tw-gateway/src/server/pipeline.rs b/crates/tw-gateway/src/server/pipeline.rs index 5b1bcb07..f2e8addf 100644 --- a/crates/tw-gateway/src/server/pipeline.rs +++ b/crates/tw-gateway/src/server/pipeline.rs @@ -76,8 +76,12 @@ pub(super) async fn pipeline( ) -> Result { // 请求体**只解析一次**：本地应答的判定、路由事实、对话的指纹、内容过滤都用这一份（见 // `read`）。解不开是 None —— 照样往下走，同格式直通照样发 - let parsed = serde_json::from_slice::(&req.body).ok(); - let intent = match probe(&state, &rt, &req, parsed.as_ref()) { + let (parsed, probed) = heavy(&req.body, || { + let parsed = serde_json::from_slice::(&req.body).ok(); + let probed = probe(&state, &rt, &req, parsed.as_ref()); + (parsed, probed) + }); + let intent = match probed { Probe::Answered(resp) => return Ok(resp), Probe::Intent(intent) => intent, }; @@ -95,7 +99,7 @@ pub(super) async fn pipeline( // 管线第 2 步：读出路由事实，路由。**看的是客户端的原话**：插件的请求钩子排在路由 // 之后（每发往一个上游跑一次，见 `plug`），左右不了请求去哪一家 - let (mut reading, fp) = read(&req, intent, parsed.as_ref()); + let (mut reading, fp) = heavy(&req.body, || read(&req, intent, parsed.as_ref())); let conv = conversation(&rt, &req, &reading, fp.as_deref()); let (choice, decision) = match route(&state, &rt, &req, &reading, conv.as_ref())? { Routed::Go(choice, decision) => (choice, decision), @@ -104,7 +108,7 @@ pub(super) async fn pipeline( Routed::Refused(choice, why) => { let to = ("", tw_api::Billing::PerToken); // 一个字节都没发出去，也没什么可报的；存下来的请求照样按这一档换、打码 - let seen = look(&rt, &req); + let seen = heavy(&req.body, || look(&rt, &req)); let redaction = redaction(&rt, seen.ledger); let (id, _) = open( &state, @@ -143,18 +147,24 @@ pub(super) async fn pipeline( .await?; // 管线第 4 步：内容过滤先下结论，不发事件。删过的话，后面一律用删过的那一份 - let screening = screen(&rt, &mut req, &mut reading, parsed.as_ref()); - let started = start( - &state, - &rt, - &req, - &reading, - choice, - wait_until, - &decision, - fp.as_deref(), - ending, - ); + let size = req.body.len(); + let (screening, started) = heavy_for(size, || { + let screening = screen(&rt, &mut req, &mut reading, parsed.as_ref()); + // 解析出来的那一份到这里就用完了：删过字的话它也不再是请求体的样子 + drop(parsed); + let started = start( + &state, + &rt, + &req, + &reading, + choice, + wait_until, + &decision, + fp.as_deref(), + ending, + ); + (screening, started) + }); // 用量上限的预留跟着请求号走，等存储层记下这一行时换成实数 hold.bind(started.id); // 结论挂在请求号上报。**拒绝的也在开始之后**：被拒是一次来源为 `denied` 的失败， @@ -1023,6 +1033,32 @@ fn screen( sc } +/// 请求体到这么大，管线上整份解它、扫它的那几步挪出异步线程（见 [`heavy`]）。 +const HEAVY_BODY: usize = 1024 * 1024; + +/// 在整个请求体上做的 CPU 活：解析、本地应答的判定、内容过滤、出站脱敏、每一跳改写请求体。 +/// +/// **大的请求体挪出异步线程做**（`block_in_place`）：一个 8 MB 的请求要算几十毫秒，就地算的话， +/// 同一个线程上别的连接 —— 正在流的回答 —— 跟着停住这么久。挪出去时这个线程上排着的活交给 +/// 别的线程接着干。小的就地做：挪一次的开销比它本身还大。 +/// +/// 只在多线程的运行时上挪：单线程的运行时（测试用的就是它）里 `block_in_place` 会 panic。 +pub(super) fn heavy(body: &[u8], f: impl FnOnce() -> T) -> T { + heavy_for(body.len(), f) +} + +/// [`heavy`]，按请求体的长度：要在 `f` 里改请求体的调用方先量好长度 +pub(super) fn heavy_for(size: usize, f: impl FnOnce() -> T) -> T { + use tokio::runtime::{Handle, RuntimeFlavor}; + let threads = + Handle::try_current().is_ok_and(|h| h.runtime_flavor() == RuntimeFlavor::MultiThread); + if size >= HEAVY_BODY && threads { + tokio::task::block_in_place(f) + } else { + f() + } +} + /// 这个请求是 Claude Code 发的吗。 /// /// **`max_tokens: 1` 那条判定必须同时要求它**，否则会误伤别人 diff --git a/crates/tw-gateway/src/server/pipeline/admission.rs b/crates/tw-gateway/src/server/pipeline/admission.rs index b899de3f..e6cc42a1 100644 --- a/crates/tw-gateway/src/server/pipeline/admission.rs +++ b/crates/tw-gateway/src/server/pipeline/admission.rs @@ -140,7 +140,7 @@ fn refused( tracing::info!(key = %req.client_name, per = r.limit.per.word(), "a usage limit of the key refused the request"); let why = r.error(); // 一个字节都没发出去；存下来的请求照样按这一档换、打码 - let seen = look(rt, req); + let seen = super::heavy(&req.body, || look(rt, req)); let (id, _) = open( state, req, diff --git a/crates/tw-gateway/src/server/pipeline/hop.rs b/crates/tw-gateway/src/server/pipeline/hop.rs index b2a7623f..2c29542d 100644 --- a/crates/tw-gateway/src/server/pipeline/hop.rs +++ b/crates/tw-gateway/src/server/pipeline/hop.rs @@ -407,15 +407,17 @@ pub(super) async fn try_upstreams<'a>( let model = Some(plugged.model.clone().unwrap_or(sent)).filter(asked_other); let asked = Asked::of(req, reading, &plugged); - let out = match prepare( - state, - req, - reading, - &asked, - provider, - &effective_set, - Some(id), - ) { + let out = match super::heavy(asked.body, || { + prepare( + state, + req, + reading, + &asked, + provider, + &effective_set, + Some(id), + ) + }) { Ok(out) => out, Err(err) => { chain.push(hop_failed( @@ -429,8 +431,6 @@ pub(super) async fn try_upstreams<'a>( } }; - // 这一家在这段对话里拒过的别家封存的推理：发之前先去掉（见 `crate::seal`） - let unsealed = unseal_upfront(state, req, started, provider, &out); // 出站脱敏的拦截档：换掉**这一跳真正发出去的那一份**（可能转换过 // 格式）。规则是全局的，每一跳换掉的是同一批东西；**接着原文那本账换**， // 同一个值在每一跳、在存下来的那份请求里都是同一个占位符。插件改过的一跳接着 @@ -439,15 +439,19 @@ pub(super) async fn try_upstreams<'a>( .rewritten .as_ref() .map_or(&started.ledger, |r| &r.ledger); - // 发出去的就是客户端原文的那些字节（同格式直通、一个字节都没改）：开头那一遍在 - // 它上面找到的就是这一跳要换的，不再找一遍 - let known = started.hits.as_deref().filter(|_| { - plugged.rewritten.is_none() - && unsealed.as_ptr() == req.body.as_ptr() - && unsealed.len() == req.body.len() + let (body, ledger) = super::heavy(&out.body, || { + // 这一家在这段对话里拒过的别家封存的推理：发之前先去掉（见 `crate::seal`） + let unsealed = unseal_upfront(state, req, started, provider, &out); + // 发出去的就是客户端原文的那些字节（同格式直通、一个字节都没改）：开头那一遍在 + // 它上面找到的就是这一跳要换的，不再找一遍 + let known = started.hits.as_deref().filter(|_| { + plugged.rewritten.is_none() + && unsealed.as_ptr() == req.body.as_ptr() + && unsealed.len() == req.body.len() + }); + let mode = rt.config.security.redact.mode; + crate::guard::replace_found(mode, &rt.redact, unsealed, seed, known) }); - let mode = rt.config.security.redact.mode; - let (body, ledger) = crate::guard::replace_found(mode, &rt.redact, unsealed, seed, known); // 用这个 provider 自己的 Client —— 它带着该走的代理。**在取密钥 // 之前拿到**：OAuth 换 token 也要走这条代理。 From 0a3463650c1c6f0d5b7b3e21b20e9bf70514fb3e Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 21:31:39 +0800 Subject: [PATCH 14/16] fix(gateway): routing rewrites change only the fields they set When a routing rule (or an alias rename) changes the model, the output limit or turns thinking off, `forward::apply_set` parsed the body and wrote it back whole. serde_json without `preserve_order` sorts every object's keys, so the upstream got the client's tool schemas, tool inputs and top level reordered -- a different byte string for its prompt cache on every such request, the very cache loss the module warns about. The rewrite now diffs the parsed body before and after the change and edits the original text: a changed value is replaced in place (objects on both sides are compared one level further in, so Gemini's `generationConfig` keeps its other keys and their order), a removed member is cut with its comma, an added one is appended to its object. Everything else -- order, whitespace, escapes -- is the client's. A rewrite that changes nothing sends the original buffer. Only when a key is written twice in an object the change reaches does it fall back to writing the body back, since serde_json keeps the last one. The member scanner moves from `egress` to a shared `splice` module. Why not turn on `preserve_order`: it is a feature of the shared serde_json, so it would also reorder every JSON the gateway, the control-plane API and ThinkWatch Enterprise produce (tw-dialect and tw-guard inherit it), and `Map::remove` would become `swap_remove`, moving the last key into the removed one's place -- removing `thinking` would still reorder the body. The splice keeps the change on the one path that rewrites, which is off by default. Tests: the exact bytes for an Anthropic body with a nested tool schema, an appended limit and Gemini's nested config; a random test over five formats and random rewrites checks that the result means what writing it back meant and that untouched members keep their bytes and order. Co-Authored-By: Claude Opus 5.5 --- crates/tw-gateway/src/egress.rs | 144 +------------ crates/tw-gateway/src/forward.rs | 248 +++++++++++++++++++++-- crates/tw-gateway/src/lib.rs | 1 + crates/tw-gateway/src/splice.rs | 334 +++++++++++++++++++++++++++++++ 4 files changed, 571 insertions(+), 156 deletions(-) create mode 100644 crates/tw-gateway/src/splice.rs diff --git a/crates/tw-gateway/src/egress.rs b/crates/tw-gateway/src/egress.rs index 40c63e5f..c294ebc0 100644 --- a/crates/tw-gateway/src/egress.rs +++ b/crates/tw-gateway/src/egress.rs @@ -226,7 +226,7 @@ enum Cut { /// 只走顶层和 `parent` 这两层：别的值整个跳过，字符串里按字节找下一个引号或反斜杠。键里 /// 写了转义的（`"user\u005fid"`）照解析出来的样子比，和 serde_json 认的一样。 fn cut_members(body: &[u8], parent: &str, keys: &[&str], drop_empty: bool) -> Cut { - let Some(top) = json::object(body, json::ws(body, 0)) else { + let Some(top) = crate::splice::object(body, crate::splice::ws(body, 0)) else { return Cut::Unusual; }; let mut parents = top.iter().filter(|m| m.is(body, parent)); @@ -241,7 +241,7 @@ fn cut_members(body: &[u8], parent: &str, keys: &[&str], drop_empty: bool) -> Cu if body[p.value.start] != b'{' { return Cut::Nothing; } - let Some(inner) = json::object(body, p.value.start) else { + let Some(inner) = crate::splice::object(body, p.value.start) else { return Cut::Unusual; }; let gone: Vec = (0..inner.len()) @@ -262,9 +262,9 @@ fn cut_members(body: &[u8], parent: &str, keys: &[&str], drop_empty: bool) -> Cu .iter() .position(|m| m.key.start == p.key.start) .unwrap_or_default(); - json::cuts(&top, &[at]) + crate::splice::cuts(&top, &[at]) } else { - json::cuts(&inner, &gone) + crate::splice::cuts(&inner, &gone) }; let mut out = Vec::with_capacity(body.len()); let mut at = 0; @@ -276,142 +276,6 @@ fn cut_members(body: &[u8], parent: &str, keys: &[&str], drop_empty: bool) -> Cu Cut::Done(out) } -/// 在 JSON 原文上认出一个对象的成员，好按字节剪。**只给解得开的 JSON 用**：写坏了的地方 -/// 不报错，只保证不越界。 -mod json { - use std::ops::Range; - - /// 对象里的一个成员：键（连同引号）和值在原文里的区间 - #[derive(Debug, Clone)] - pub(super) struct Member { - pub(super) key: Range, - pub(super) value: Range, - } - - impl Member { - /// 键解出来是不是 `name` - pub(super) fn is(&self, b: &[u8], name: &str) -> bool { - let raw = &b[self.key.start + 1..self.key.end - 1]; - if !raw.contains(&b'\\') { - return raw == name.as_bytes(); - } - serde_json::from_slice::(&b[self.key.clone()]).is_ok_and(|k| k == name) - } - } - - /// 从 `i` 起跳过空白 - pub(super) fn ws(b: &[u8], mut i: usize) -> usize { - while i < b.len() && matches!(b[i], b' ' | b'\t' | b'\n' | b'\r') { - i += 1; - } - i - } - - /// `b[i]` 是一个字符串的开头引号：返回结尾引号之后的位置 - fn string_end(b: &[u8], i: usize) -> Option { - if b.get(i) != Some(&b'"') { - return None; - } - let mut j = i + 1; - loop { - j += memchr::memchr2(b'"', b'\\', b.get(j..)?)?; - if b[j] == b'"' { - return Some(j + 1); - } - j += 2; - } - } - - /// `b[i]` 起的一个值：返回它之后的位置 - fn value_end(b: &[u8], i: usize) -> Option { - match *b.get(i)? { - b'"' => string_end(b, i), - b'{' | b'[' => { - let mut depth = 0usize; - let mut j = i; - loop { - match *b.get(j)? { - b'"' => { - j = string_end(b, j)?; - continue; - } - b'{' | b'[' => depth += 1, - b'}' | b']' => { - depth -= 1; - if depth == 0 { - return Some(j + 1); - } - } - _ => {} - } - j += 1; - } - } - // 数、true、false、null：到下一个分隔为止 - _ => { - let end = b[i..] - .iter() - .position(|c| matches!(c, b',' | b'}' | b']' | b' ' | b'\t' | b'\n' | b'\r')) - .map_or(b.len(), |n| i + n); - (end > i).then_some(end) - } - } - } - - /// `b[open]` 是一个对象的 `{`：它的成员，按先后 - pub(super) fn object(b: &[u8], open: usize) -> Option> { - if b.get(open) != Some(&b'{') { - return None; - } - let mut out = Vec::new(); - let mut i = ws(b, open + 1); - if b.get(i) == Some(&b'}') { - return Some(out); - } - loop { - let key = i..string_end(b, i)?; - let colon = ws(b, key.end); - if b.get(colon) != Some(&b':') { - return None; - } - let v = ws(b, colon + 1); - let value = v..value_end(b, v)?; - let next = ws(b, value.end); - out.push(Member { key, value }); - match *b.get(next)? { - b',' => i = ws(b, next + 1), - b'}' => return Some(out), - _ => return None, - } - } - } - - /// 剪掉 `members` 里第 `gone` 几个（按先后）要剪的区间，按先后、互不重叠。剪完还是 - /// 合法的 JSON：剪掉一个成员连同它后面的逗号；排在最后的，连同它前面的逗号 - pub(super) fn cuts(members: &[Member], gone: &[usize]) -> Vec> { - let mut out = Vec::new(); - let mut k = 0; - while k < gone.len() { - // 连着要剪的一串 - let first = gone[k]; - let mut last = first; - while k + 1 < gone.len() && gone[k + 1] == last + 1 { - k += 1; - last += 1; - } - k += 1; - out.push(if last + 1 < members.len() { - members[first].key.start..members[last + 1].key.start - } else if first > 0 { - members[first - 1].value.end..members[last].value.end - } else { - members[first].key.start..members[last].value.end - }); - } - out - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/tw-gateway/src/forward.rs b/crates/tw-gateway/src/forward.rs index 1af97948..99cfe6e1 100644 --- a/crates/tw-gateway/src/forward.rs +++ b/crates/tw-gateway/src/forward.rs @@ -75,29 +75,58 @@ pub fn gemini_path_with_model(path: &str, model: &str) -> String { /// 字段名按客户端的格式写（[`tw_dialect::params`]）：同一个「最大输出」在四种格式里叫 /// 四个名字。认不出格式时按 Anthropic 写。Gemini 的模型在路径里，见 /// [`gemini_path_with_model`]。 +/// +/// **只动改了的那几个字段，别的字节一个不动**（见 [`crate::splice`]）：解析再写回的话， +/// 每个对象的键都按字母重排了，工具定义、工具参数一起变了样。改出来和原来一样的（规则 +/// 写的值客户端本来就是这么写的）原样发。改到的对象里同一个键写了两遍的，才照旧写回。 pub fn apply_set( body: &Bytes, set: &tw_engine::SetAction, client: Option, ) -> Bytes { - use tw_dialect::params; if set.is_empty() { return body.clone(); } - let Ok(mut v) = serde_json::from_slice::(body) else { + let Ok(before) = serde_json::from_slice::(body) else { // 解不开就别动。我们的解析器不认识的东西，上游可能完全认识。 tracing::warn!("the request body is not JSON; skipping the parameter rewrites"); return body.clone(); }; - if !v.is_object() { + if !before.is_object() { return body.clone(); } + let mut v = before.clone(); + set_params(&mut v, set, client); + match crate::splice::rewrite(body, &before, &v) { + crate::splice::Rewrite::Changed(b) => return Bytes::from(b), + crate::splice::Rewrite::Same => return body.clone(), + crate::splice::Rewrite::Unusual => {} + } + match serde_json::to_vec(&v) { + Ok(b) => Bytes::from(b), + // 序列化不该失败，但真失败了宁可发原文也不要发半个 body + Err(e) => { + tracing::error!( + "the rewritten request body could not be serialized; sending the original instead: {e}" + ); + body.clone() + } + } +} + +/// 把 `set` 点名的参数写进解析出来的请求体 `v`（见 [`apply_set`]）。 +fn set_params( + v: &mut serde_json::Value, + set: &tw_engine::SetAction, + client: Option, +) { + use tw_dialect::params; let client = client.unwrap_or(tw_dialect::ir::Dialect::Anthropic); if let Some(m) = &set.model { - params::set_model(client, &mut v, m); + params::set_model(client, v, m); } if let Some(t) = set.max_tokens { - params::set_max_output_tokens(client, &mut v, t); + params::set_max_output_tokens(client, v, t); } match set.thinking { // 开启思考需要一个 budget，而我们没有一个合理的值可以编。 @@ -105,19 +134,9 @@ pub fn apply_set( Some(true) => tracing::warn!( "set.thinking: true is not supported yet (it needs budget_tokens); ignoring it" ), - Some(false) => params::disable_thinking(client, &mut v), + Some(false) => params::disable_thinking(client, v), None => {} } - match serde_json::to_vec(&v) { - Ok(b) => Bytes::from(b), - // 序列化不该失败，但真失败了宁可发原文也不要发半个 body - Err(e) => { - tracing::error!( - "the rewritten request body could not be serialized; sending the original instead: {e}" - ); - body.clone() - } - } } pub fn map_reqwest_error(e: reqwest::Error) -> GatewayError { @@ -272,4 +291,201 @@ mod tests { "/v1beta/models/small:streamGenerateContent" ); } + + fn set( + model: Option<&str>, + max_tokens: Option, + thinking: Option, + ) -> tw_engine::SetAction { + tw_engine::SetAction { + model: model.map(str::to_string), + max_tokens, + thinking, + ..Default::default() + } + } + + /// 以前的写法：解析、改、整份写回（键按字母重排） + fn whole(body: &[u8], s: &tw_engine::SetAction, d: Option) -> Bytes { + let mut v: serde_json::Value = serde_json::from_slice(body).unwrap(); + set_params(&mut v, s, d); + Bytes::from(serde_json::to_vec(&v).unwrap()) + } + + /// 规则改了模型、输出上限、关了思考：上游收到的是客户端的那些字节，只有这几个字段变了 + /// —— 工具定义里的键、工具参数里的键、整个请求的先后都照客户端写的 + #[test] + fn a_rewritten_body_is_the_client_body_but_for_the_fields_set() { + use tw_dialect::ir::Dialect; + let tools = r#""tools":[{"name":"Read","description":"读 \"文件\"","input_schema":{"type":"object","properties":{"file_path":{"type":"string"},"offset":{"type":"number"},"limit":{"type":"number"}},"required":["file_path"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}}]"#; + let messages = r#""messages":[{"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"Read","input":{"zeta":1,"alpha":{"b":2,"a":3}}}]}]"#; + let body = format!( + r#"{{"model":"claude-opus-5",{messages},"system":"s",{tools},"metadata":{{"user_id":"u"}},"max_tokens":32000,"thinking":{{"type":"enabled","budget_tokens":31999}},"stream":true}}"# + ); + let out = apply_set( + &Bytes::from(body.clone()), + &set(Some("claude-haiku-5"), Some(4096), Some(false)), + Some(Dialect::Anthropic), + ); + assert_eq!( + std::str::from_utf8(&out).unwrap(), + format!( + r#"{{"model":"claude-haiku-5",{messages},"system":"s",{tools},"metadata":{{"user_id":"u"}},"max_tokens":4096,"stream":true}}"# + ) + ); + // 客户端没写的输出上限补在最后 + let out = apply_set( + &Bytes::from(format!(r#"{{"model":"m",{tools}}}"#)), + &set(None, Some(10), None), + Some(Dialect::Anthropic), + ); + assert_eq!( + std::str::from_utf8(&out).unwrap(), + format!(r#"{{"model":"m",{tools},"max_tokens":10}}"#) + ); + // Gemini 的输出上限和思考在 generationConfig 里：只动里面那两个，别的照旧 + let gemini = format!( + r#"{{"contents":[],{tools},"generationConfig": {{"temperature": 0.2, "thinkingConfig": {{"thinkingBudget": 9}}, "maxOutputTokens": 9, "topK": 3}}}}"# + ); + let out = apply_set( + &Bytes::from(gemini), + &set(Some("ignored"), Some(512), Some(false)), + Some(Dialect::Gemini), + ); + assert_eq!( + std::str::from_utf8(&out).unwrap(), + format!( + r#"{{"contents":[],{tools},"generationConfig": {{"temperature": 0.2, "maxOutputTokens": 512, "topK": 3}}}}"# + ) + ); + } + + /// 改出来和原来一样（规则写的值客户端本来就这么写）：原样发，连拷贝都没有 + #[test] + fn a_set_that_changes_nothing_sends_the_same_bytes() { + let b = Bytes::from_static(br#"{ "max_tokens" : 100, "b":{"z":1,"a":2} }"#); + let out = apply_set(&b, &set(None, Some(100), None), None); + assert_eq!(out.as_ptr(), b.as_ptr()); + } + + /// 改到的对象里同一个键写了两遍：剪的结果说不准和 serde_json 认的一样（它留后一个）， + /// 照旧整份写回 + #[test] + fn a_key_written_twice_falls_back_to_writing_the_body_back() { + let b = br#"{"model":"a","x":1,"model":"b"}"#; + let s = set(Some("c"), None, None); + assert_eq!( + apply_set(&Bytes::from_static(b), &s, None), + whole(b, &s, None) + ); + } + + /// 随机拼出来的请求、随机的改写、五种格式：改出来的和整份写回的意思一样；没改到的 + /// 顶层成员原样、按原来的先后都在 + #[test] + fn rewriting_in_place_means_what_writing_back_meant() { + use tw_dialect::ir::Dialect; + struct Rng(u64); + impl Rng { + fn below(&mut self, n: usize) -> usize { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 % n as u64) as usize + } + fn pick<'a>(&mut self, xs: &[&'a str]) -> &'a str { + xs[self.below(xs.len())] + } + } + let mut r = Rng(0x2545_f491_4f6c_dd1d); + // 转义的键（`"model"`）由 `char::from(92)` 拼：测试里直接写出来的转义，经过 + // 某些编辑工具会变成真字符 + let escaped_model = format!("\"mod{}u0065l\"", char::from(92)); + let config = [ + r#"{"temperature":0.2,"maxOutputTokens":9,"thinkingConfig":{"thinkingBudget":1}}"#, + r#"{"max_output_tokens":9,"thinking_config":{"a":1},"top_k":2}"#, + r#"{ }"#, + r#"{"thinkingConfig":1,"thinkingConfig":2}"#, + r#""not an object""#, + ]; + let members = [ + "\"model\":\"big\"", + "\"max_tokens\":9", + "\"max_completion_tokens\":7", + "\"max_output_tokens\":5", + "\"thinking\":{\"type\":\"enabled\",\"budget_tokens\":9}", + "\"reasoning_effort\":\"high\"", + "\"reasoning\":{\"effort\":\"low\"}", + "\"inferenceConfig\":{\"temperature\":1,\"maxTokens\":3}", + "\"additionalModelRequestFields\":{\"thinking\":{\"type\":\"enabled\"},\"k\":1}", + "\"tools\":[{\"input_schema\":{\"type\":\"object\",\"properties\":{\"b\":{},\"a\":{}}}}]", + "\"messages\":[{\"role\":\"user\",\"content\":\"max_tokens \\\"model\\\" }{\"}]", + ]; + const WS: &[&str] = &["", " ", "\n "]; + let dialects = [ + Dialect::Anthropic, + Dialect::Chat, + Dialect::Responses, + Dialect::Gemini, + Dialect::Bedrock, + ]; + let (mut changed, mut unusual) = (0, 0); + for _ in 0..5_000 { + let mut top: Vec = Vec::new(); + for _ in 0..r.below(7) { + let m = match r.below(14) { + 0 => format!("\"generationConfig\":{}", r.pick(&config)), + 1 => format!("\"generation_config\":{}", r.pick(&config)), + 2 => format!("{escaped_model}:\"esc\""), + _ => r.pick(&members).to_string(), + }; + let (a, b) = (r.pick(WS), r.pick(WS)); + top.push(format!("{a}{m}{b}")); + } + let body = format!("{{{}}}", top.join(",")); + let s = set( + [None, Some("small"), Some("big")][r.below(3)], + [None, Some(9), Some(512)][r.below(3)], + [None, Some(false), Some(true)][r.below(3)], + ); + if s.is_empty() { + continue; + } + let d = Some(dialects[r.below(dialects.len())]); + let out = apply_set(&Bytes::from(body.clone()), &s, d); + let value = |b: &[u8]| serde_json::from_slice::(b).unwrap(); + assert_eq!( + value(&out), + value(&whole(body.as_bytes(), &s, d)), + "{body} {s:?} {d:?}" + ); + let before = value(body.as_bytes()); + let mut after = before.clone(); + set_params(&mut after, &s, d); + if crate::splice::rewrite(body.as_bytes(), &before, &after) + == crate::splice::Rewrite::Unusual + { + unusual += 1; + continue; + } + changed += 1; + // 没改到的顶层成员：原文那一段照样在，先后不变 + let b = body.as_bytes(); + let mut from = 0; + for m in crate::splice::object(b, 0).unwrap() { + let name: String = serde_json::from_slice(&b[m.key.clone()]).unwrap(); + if before.get(&name) != after.get(&name) { + continue; + } + let piece = &body[m.key.start..m.value.end]; + let at = std::str::from_utf8(&out).unwrap()[from..] + .find(piece) + .unwrap_or_else(|| { + panic!("{piece} moved in {}", String::from_utf8_lossy(&out)) + }); + from += at + piece.len(); + } + } + assert!(changed > 1_000 && unusual > 50, "{changed} {unusual}"); + } } diff --git a/crates/tw-gateway/src/lib.rs b/crates/tw-gateway/src/lib.rs index 3a111a64..7cb53c2f 100644 --- a/crates/tw-gateway/src/lib.rs +++ b/crates/tw-gateway/src/lib.rs @@ -45,6 +45,7 @@ pub mod sent; pub mod server; pub mod session; pub mod slots; +mod splice; pub mod state; pub mod translate; pub mod wire; diff --git a/crates/tw-gateway/src/splice.rs b/crates/tw-gateway/src/splice.rs new file mode 100644 index 00000000..268b230e --- /dev/null +++ b/crates/tw-gateway/src/splice.rs @@ -0,0 +1,334 @@ +//! 在 JSON 原文上按字节改请求体：只动改了的那几个成员，别的字节一个不动。 +//! +//! 解析再写回不行：没开 `preserve_order` 的 serde_json 会把每个对象的键按字母重排 —— +//! 工具定义、工具参数、整个请求都变了样，而上游的提示缓存按字节认。所以去掉客户端的身份 +//! 字段（[`crate::egress`]）、按路由规则改参数（[`crate::forward::apply_set`]）都在原文上 +//! 剪、换、补。 +//! +//! **只给解得开的 JSON 用**：调用方解析过它。写坏了的地方这里不报错，只保证不越界。 + +use std::borrow::Cow; +use std::ops::Range; + +use serde_json::{Map, Value}; + +/// 对象里的一个成员：键（连同引号）和值在原文里的区间 +#[derive(Debug, Clone)] +pub(crate) struct Member { + pub(crate) key: Range, + pub(crate) value: Range, +} + +impl Member { + /// 键解出来的样子：写了转义的照 serde_json 解 + fn name<'b>(&self, b: &'b [u8]) -> Option> { + let raw = &b[self.key.start + 1..self.key.end - 1]; + if !raw.contains(&b'\\') { + return std::str::from_utf8(raw).ok().map(Cow::Borrowed); + } + serde_json::from_slice::(&b[self.key.clone()]) + .ok() + .map(Cow::Owned) + } + + /// 键解出来是不是 `name` + pub(crate) fn is(&self, b: &[u8], name: &str) -> bool { + self.name(b).is_some_and(|k| k == name) + } +} + +/// 从 `i` 起跳过空白 +pub(crate) fn ws(b: &[u8], mut i: usize) -> usize { + while i < b.len() && matches!(b[i], b' ' | b'\t' | b'\n' | b'\r') { + i += 1; + } + i +} + +/// `b[i]` 是一个字符串的开头引号：返回结尾引号之后的位置 +fn string_end(b: &[u8], i: usize) -> Option { + if b.get(i) != Some(&b'"') { + return None; + } + let mut j = i + 1; + loop { + j += memchr::memchr2(b'"', b'\\', b.get(j..)?)?; + if b[j] == b'"' { + return Some(j + 1); + } + j += 2; + } +} + +/// `b[i]` 起的一个值：返回它之后的位置 +fn value_end(b: &[u8], i: usize) -> Option { + match *b.get(i)? { + b'"' => string_end(b, i), + b'{' | b'[' => { + let mut depth = 0usize; + let mut j = i; + loop { + match *b.get(j)? { + b'"' => { + j = string_end(b, j)?; + continue; + } + b'{' | b'[' => depth += 1, + b'}' | b']' => { + depth -= 1; + if depth == 0 { + return Some(j + 1); + } + } + _ => {} + } + j += 1; + } + } + // 数、true、false、null：到下一个分隔为止 + _ => { + let end = b[i..] + .iter() + .position(|c| matches!(c, b',' | b'}' | b']' | b' ' | b'\t' | b'\n' | b'\r')) + .map_or(b.len(), |n| i + n); + (end > i).then_some(end) + } + } +} + +/// `b[open]` 是一个对象的 `{`：它的成员，按先后 +pub(crate) fn object(b: &[u8], open: usize) -> Option> { + if b.get(open) != Some(&b'{') { + return None; + } + let mut out = Vec::new(); + let mut i = ws(b, open + 1); + if b.get(i) == Some(&b'}') { + return Some(out); + } + loop { + let key = i..string_end(b, i)?; + let colon = ws(b, key.end); + if b.get(colon) != Some(&b':') { + return None; + } + let v = ws(b, colon + 1); + let value = v..value_end(b, v)?; + let next = ws(b, value.end); + out.push(Member { key, value }); + match *b.get(next)? { + b',' => i = ws(b, next + 1), + b'}' => return Some(out), + _ => return None, + } + } +} + +/// 剪掉 `members` 里第 `gone` 几个（按先后）要剪的区间，按先后、互不重叠。剪完还是 +/// 合法的 JSON：剪掉一个成员连同它后面的逗号；排在最后的，连同它前面的逗号 +pub(crate) fn cuts(members: &[Member], gone: &[usize]) -> Vec> { + let mut out = Vec::new(); + let mut k = 0; + while k < gone.len() { + // 连着要剪的一串 + let first = gone[k]; + let mut last = first; + while k + 1 < gone.len() && gone[k + 1] == last + 1 { + k += 1; + last += 1; + } + k += 1; + out.push(if last + 1 < members.len() { + members[first].key.start..members[last + 1].key.start + } else if first > 0 { + members[first - 1].value.end..members[last].value.end + } else { + members[first].key.start..members[last].value.end + }); + } + out +} + +/// [`rewrite`] 的结果。 +#[derive(Debug, PartialEq)] +pub(crate) enum Rewrite { + /// 改了的那一份 + Changed(Vec), + /// 改完和原来一样：原文一个字节都不用动 + Same, + /// 改到的对象里同一个键写了两遍（serde_json 留后一个，剪的结果说不准和它一样），或者 + /// 原文不像一个 JSON 对象：交给调用方解析、写回 + Unusual, +} + +/// 原文 `body`（解出来是 `before`）改成 `after` 的样子，**只动改了的地方**： +/// +/// - 值变了的成员只换它的值，两边都是对象的往里一层一层地比，只换里面变了的； +/// - 没了的成员连同分隔它的逗号剪掉； +/// - 新加的成员补在那个对象的末尾（按键名的先后）。 +/// +/// 没变的成员、它们的先后、空白和转义都照原文。改出来的那一份解析出来就是 `after`（测试 +/// 拿它和解析、写回的那一份比过）。 +pub(crate) fn rewrite(body: &[u8], before: &Value, after: &Value) -> Rewrite { + let (Some(before), Some(after)) = (before.as_object(), after.as_object()) else { + return Rewrite::Unusual; + }; + let mut edits = Vec::new(); + if diff(body, ws(body, 0), before, after, &mut edits).is_none() { + return Rewrite::Unusual; + } + if edits.is_empty() { + return Rewrite::Same; + } + // 同一处起的：先补（区间是空的），再剪、再换 + edits.sort_by_key(|(r, _): &(Range, Vec)| (r.start, r.end)); + let mut out = Vec::with_capacity(body.len() + 64); + let mut at = 0; + for (r, text) in edits { + debug_assert!(r.start >= at, "edits overlap"); + out.extend_from_slice(&body[at..r.start]); + out.extend_from_slice(&text); + at = r.end; + } + out.extend_from_slice(&body[at..]); + Rewrite::Changed(out) +} + +/// `b[open]` 处的对象从 `before` 改成 `after` 要做的改动，记进 `edits`（原文的区间换成 +/// 什么）。同一个键写了两遍、或者原文对不上的是 `None` +fn diff( + b: &[u8], + open: usize, + before: &Map, + after: &Map, + edits: &mut Vec<(Range, Vec)>, +) -> Option<()> { + let members = object(b, open)?; + let names: Vec> = members.iter().map(|m| m.name(b)).collect::>()?; + let mut sorted: Vec<&str> = names.iter().map(|n| n.as_ref()).collect(); + sorted.sort_unstable(); + if sorted.windows(2).any(|w| w[0] == w[1]) { + return None; + } + let mut gone = Vec::new(); + for (i, (m, name)) in members.iter().zip(&names).enumerate() { + let Some(now) = after.get(name.as_ref()) else { + gone.push(i); + continue; + }; + match (before.get(name.as_ref()), now) { + (Some(was), now) if was == now => {} + (Some(Value::Object(was)), Value::Object(now)) if b[m.value.start] == b'{' => { + diff(b, m.value.start, was, now, edits)?; + } + _ => edits.push((m.value.clone(), serde_json::to_vec(now).ok()?)), + } + } + let mut added = Vec::new(); + for (k, v) in after { + if !names.iter().any(|n| n == k) { + added.push(format!( + "{}:{}", + serde_json::to_string(k).ok()?, + serde_json::to_string(v).ok()? + )); + } + } + let added = added.join(","); + let kept = members.len() - gone.len(); + if !added.is_empty() && kept == 0 && !members.is_empty() { + // 原来的成员都没了：整段换成新加的 + let last = members.len() - 1; + edits.push(( + members[0].key.start..members[last].value.end, + added.into_bytes(), + )); + return Some(()); + } + for c in cuts(&members, &gone) { + edits.push((c, Vec::new())); + } + if !added.is_empty() { + match members + .iter() + .enumerate() + .rev() + .find(|(i, _)| !gone.contains(i)) + { + Some((_, m)) => { + let at = m.value.end; + edits.push((at..at, format!(",{added}").into_bytes())); + } + // 空对象：写在 `{` 后面 + None => edits.push((open + 1..open + 1, added.into_bytes())), + } + } + Some(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rewritten(body: &str, after: impl FnOnce(&mut Value)) -> Rewrite { + let before: Value = serde_json::from_str(body).unwrap(); + let mut now = before.clone(); + after(&mut now); + rewrite(body.as_bytes(), &before, &now) + } + + fn text(r: Rewrite) -> String { + match r { + Rewrite::Changed(b) => String::from_utf8(b).unwrap(), + other => panic!("{other:?}"), + } + } + + #[test] + fn only_what_changed_is_touched() { + let body = r#"{ "zeta": 1, "model": "opus", "tools": [{"name":"t","input_schema":{"type":"object","properties":{"b":{},"a":{}},"required":["b"]}}], "conf": {"y": 1, "x": 2, "drop": 3}, "gone": true }"#; + let out = text(rewritten(body, |v| { + v["model"] = "haiku".into(); + v["conf"]["x"] = 20.into(); + v["conf"].as_object_mut().unwrap().remove("drop"); + v["conf"]["new"] = "n".into(); + v.as_object_mut().unwrap().remove("gone"); + v["added"] = 7.into(); + })); + assert_eq!( + out, + r#"{ "zeta": 1, "model": "haiku", "tools": [{"name":"t","input_schema":{"type":"object","properties":{"b":{},"a":{}},"required":["b"]}}], "conf": {"y": 1, "x": 20,"new":"n"},"added":7 }"# + ); + } + + #[test] + fn nothing_changed_is_the_same_and_an_emptied_object_takes_the_new_members() { + assert_eq!(rewritten(r#"{"a":1}"#, |_| {}), Rewrite::Same); + assert_eq!( + text(rewritten(r#"{"c":{"old":1}}"#, |v| v["c"] = + serde_json::json!({"new": 2}))), + r#"{"c":{"new":2}}"# + ); + assert_eq!( + text(rewritten(r#"{"c":{ }}"#, |v| v["c"]["k"] = 1.into())), + r#"{"c":{"k":1 }}"# + ); + assert_eq!( + text(rewritten(r#"{}"#, |v| v["k"] = 1.into())), + r#"{"k":1}"# + ); + } + + #[test] + fn a_key_written_twice_where_something_changes_is_unusual() { + assert_eq!( + rewritten(r#"{"a":1,"a":2}"#, |v| v["b"] = 1.into()), + Rewrite::Unusual + ); + // 写了两遍的在没改到的那一层里：照原文留着，和客户端发的一样 + assert_eq!( + text(rewritten(r#"{"m":{"a":1,"a":2},"x":1}"#, |v| v["x"] = 2.into())), + r#"{"m":{"a":1,"a":2},"x":2}"# + ); + } +} From b79c4e7fe9a632544aeee672407bbd3244dfe4fe Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 21:37:37 +0800 Subject: [PATCH 15/16] Control-plane protocol 41: incremental transcripts and dry-run shares Co-Authored-By: Claude Opus 5.5 --- crates/tw-api/src/lib.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index 25599f7c..db1a1903 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -798,7 +798,14 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// **40 起模型规格多了会不会推理、收不收图**：[`ModelRow`] 多了 `reasoning`、`image_input` /// 和它们的来源，[`ModelSpecSave`] 多了 `reasoning`、`image_input`（四项都空才是删掉）。 /// 照 39 写的界面保存规格时会把这两项丢掉。 -pub const CONTROL_API_VERSION: u32 = 40; +/// +/// **41 起会话的对话视图可以只取新的轮次，试算给出每家分到的份额**： +/// `GET /sessions/{id}/transcript` 的请求从 `()` 变成 [`TranscriptQuery`]（`from_turn`， +/// 不带就是整份），[`Transcript`] 多了 `total_turns` 和 `settled_turns`（前面这么多轮 +/// 不会再变，下次从这里取）；[`DryRunCandidate`] 多了 `share`（负载均衡组里这一家按 +/// 当前权重和系数分到的份额，别的组和这一轮不参与的是 None）。照 40 写的界面取对话视图 +/// 时传 `null` 会被拒。 +pub const CONTROL_API_VERSION: u32 = 41; #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] From b251606b99bd8274d6a343d5559b008711435f60 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:13:18 +0800 Subject: [PATCH 16/16] Reserve watch slots with a compare-exchange loop; the stored-bodies test skips bodies still being written fetch_update is deprecated on the current stable toolchain. Bodies are written to a .tmp file and renamed into place, so a directory listing can see a file that is gone by the time it is read (seen on Windows). Co-Authored-By: Claude Opus 5.5 --- crates/tw-control/tests/stored_bodies.rs | 11 +++++++++-- crates/tw-watch/src/kqueue.rs | 15 ++++++++++----- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/crates/tw-control/tests/stored_bodies.rs b/crates/tw-control/tests/stored_bodies.rs index 6dd0e484..d561ffa4 100644 --- a/crates/tw-control/tests/stored_bodies.rs +++ b/crates/tw-control/tests/stored_bodies.rs @@ -212,8 +212,15 @@ fn files(root: &std::path::Path) -> Vec<(String, String)> { .flatten() .flatten() { - let text = String::from_utf8_lossy(&std::fs::read(f.path()).unwrap()).to_string(); - out.push((f.file_name().to_string_lossy().to_string(), text)); + let name = f.file_name().to_string_lossy().to_string(); + // 正文先写进 `.tmp` 再改名：还在写的那份不算，列目录之后才改名走的也跳过 + if name.ends_with(".tmp") { + continue; + } + let Ok(bytes) = std::fs::read(f.path()) else { + continue; + }; + out.push((name, String::from_utf8_lossy(&bytes).to_string())); } } out diff --git a/crates/tw-watch/src/kqueue.rs b/crates/tw-watch/src/kqueue.rs index fad86033..0e165465 100644 --- a/crates/tw-watch/src/kqueue.rs +++ b/crates/tw-watch/src/kqueue.rs @@ -87,11 +87,16 @@ struct Held(OwnedFd); impl Held { /// 打开 `path`。名额用完了是 `Ok(None)`。 fn open(path: &Path) -> io::Result> { - let reserved = HELD.fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| { - (n < budget()).then_some(n + 1) - }); - if reserved.is_err() { - return Ok(None); + // 先占名额再开 fd。手写比较交换：`fetch_update` 在新版标准库里改了名 + let mut n = HELD.load(Ordering::Acquire); + loop { + if n >= budget() { + return Ok(None); + } + match HELD.compare_exchange_weak(n, n + 1, Ordering::AcqRel, Ordering::Acquire) { + Ok(_) => break, + Err(now) => n = now, + } } let release = || HELD.fetch_sub(1, Ordering::AcqRel); let Ok(c) = CString::new(path.as_os_str().as_bytes()) else {