diff --git a/Cargo.lock b/Cargo.lock index bb59dcc3..d558265b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -130,7 +130,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -141,7 +141,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -974,9 +974,9 @@ dependencies = [ [[package]] name = "ctutils" -version = "0.4.2" +version = "0.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +checksum = "38d28fb6aa7dd680b3ffd1a3e64e115c177bdba63c0ae7f9cd48129ff05c4fa0" dependencies = [ "cmov", ] @@ -1090,7 +1090,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2153,7 +2153,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -2463,7 +2463,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2798,7 +2798,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2856,7 +2856,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3195,7 +3195,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -3286,7 +3286,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3405,9 +3405,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.53.1" +version = "1.54.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +checksum = "b52fd3255258d26fac7264ab5a76b53c13367136dc011d3d6ad4f8ed070396a3" dependencies = [ "bytes", "libc", @@ -3767,6 +3767,7 @@ dependencies = [ "tw-engine", "tw-gateway", "tw-guard", + "tw-handover", "tw-link", "tw-observe", "tw-pack", @@ -3871,6 +3872,19 @@ dependencies = [ "tw-dialect", ] +[[package]] +name = "tw-handover" +version = "0.81.0" +dependencies = [ + "libc", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.20", + "tokio", + "windows-sys 0.61.2", +] + [[package]] name = "tw-link" version = "0.81.0" @@ -4040,6 +4054,7 @@ dependencies = [ "axum", "chrono", "clap", + "futures", "http-body-util", "hyper", "hyper-util", @@ -4058,6 +4073,7 @@ dependencies = [ "tw-control", "tw-gateway", "tw-guard", + "tw-handover", "tw-link", "tw-observe", "tw-pack", @@ -4567,7 +4583,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 563c06af..9ebedc58 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,8 @@ members = [ "crates/tw-sync", # 控制通道的握手与加密。桌面端按 tag 依赖它，所以只依赖契约层和 tw-yaml "crates/tw-link", + # 新旧 core 之间的交接：传监听套接字、搬运快照与结算（升级不断线） + "crates/tw-handover", # ── bin ────────────────────────────────────────────── "bin/twcore", ] @@ -83,6 +85,7 @@ tw-link = { path = "crates/tw-link" } tw-plugin = { path = "crates/tw-plugin" } tw-pack = { path = "crates/tw-pack" } tw-sync = { path = "crates/tw-sync" } +tw-handover = { path = "crates/tw-handover" } # ── 版本与企业版 workspace 对齐，便于反向依赖时不打架 ──── tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "time", "sync", "io-util", "signal"] } diff --git a/bin/twcore/Cargo.toml b/bin/twcore/Cargo.toml index fc898f4b..66c556dd 100644 --- a/bin/twcore/Cargo.toml +++ b/bin/twcore/Cargo.toml @@ -16,6 +16,8 @@ categories.workspace = true [dependencies] tw-config = { workspace = true } tw-control = { workspace = true } +# 升级不断线：新旧 core 之间的交接连接 +tw-handover = { workspace = true } # `twcore export` / `import` / `backup`：配置包和本机备份 tw-pack = { workspace = true } tw-gateway = { workspace = true } @@ -53,6 +55,8 @@ libc = { workspace = true } windows-sys = { workspace = true, features = [ "Win32_Foundation", "Win32_System_Threading", + # 交接时把标准输出、错误输出从桌面端的管道上挪开（SetStdHandle） + "Win32_System_Console", ] } [dev-dependencies] @@ -60,3 +64,5 @@ tempfile = "3" # `twcore upgrade` 的测试对着一个假的 GitHub 跑 axum = { workspace = true } serde_json = { workspace = true } +# 交接的端到端测试：假上游流式回答，测试自己扮接替者 +futures = { workspace = true } diff --git a/bin/twcore/src/background.rs b/bin/twcore/src/background.rs new file mode 100644 index 00000000..38d3279b --- /dev/null +++ b/bin/twcore/src/background.rs @@ -0,0 +1,162 @@ +//! 后台事务：OAuth 写回、同步、价目表刷新、默认插件、模型目录、余额、日志回收。 +//! +//! # 为什么放在一个单独的运行时上 +//! +//! 交接提交之后，旧 core 只做完手上的请求，**后台事务一律交给新 core**（方案 6.1）：两边 +//! 各自续期会让对方的 refresh token 作废，两边各自同步、写配置会互相覆盖。可这些任务大多 +//! 是起了就不管的循环，没有一个能叫停的把手。把它们都起在这个运行时上，提交时 +//! [`Background::stop`] 一次关掉整个运行时，任务连同它们派生的任务一起停下，转发那一边 +//! （主运行时）不受影响。 +//! +//! **模型目录、余额例外**，见 [`Background::start`] 里的说明：它们和转发共用连接池， +//! 不能跟着这个运行时一起关。 +//! +//! **OAuth 写回例外地先收尾**：它手上那份新 refresh token 要是丢了，新 core 读到的就是一张 +//! 已经作废的（见 [`tw_control::rotation::Rotation::settle`]）。 + +use std::sync::Arc; +use std::time::Duration; + +use tokio::sync::watch; + +use crate::record::Store; + +/// 收尾 OAuth 写回最多等多久 +const ROTATION_SETTLE: Duration = Duration::from_secs(5); + +#[derive(Default)] +pub struct Background { + rt: Option, + rotation: Option>, +} + +/// 起后台事务要的东西。 +pub struct Start<'a> { + pub control: &'a tw_control::ControlState, + /// 安全模式：只起写回、同步（暂停着）、价目表；不起默认插件、模型目录、余额 + pub safe: bool, + pub seeder: Option, + pub store: Option, + /// 请求库打开了没有（接替者要等旧 core 退出才打开，见 `record::open_later`）。 + /// 日志回收等它 + pub store_opened: watch::Receiver, +} + +impl Background { + /// 起来。返回同步服务（退出前要 `flush`）。 + pub fn start(&mut self, s: Start<'_>) -> std::io::Result> { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_name("tw-background") + .enable_all() + .build()?; + let sync = { + // 下面这些函数各自 `tokio::spawn`：在这个运行时里调，任务就落在这里 + let _here = rt.enter(); + if let Some(store) = s.store.clone() { + spawn_gc(s.control.gateway.clone(), store, s.store_opened.clone()); + } + // 凭据轮换要写回 config.yaml。**不是人发起的配置写入只有两种**， + // 这是一种（另一种是默认插件）—— 理由是服务器换发新 refresh + // token 的那一刻旧的就作废了，不写回等于让配置文件从那一秒起就是坏的。 + self.rotation = Some(Arc::new(tw_control::rotation::spawn(s.control.clone()))); + // 多设备同步：启动时一轮，之后按本机改动、定时、文件夹变化、唤醒。安全模式下不自动同步 + let sync = tw_control::sync::start(s.control, s.safe); + // 定期刷新默认价目表（`pricing.auto_update`，默认开） + tw_control::pricing::spawn(s.control.clone()); + if !s.safe + && let Some(seeder) = s.seeder + { + tw_control::plugins::defaults::spawn(seeder, s.control.cfg.clone()); + } + sync + }; + if !s.safe { + // **这两个不在这个运行时上**，起在调用方的（主）运行时：它们和转发共用上游的 + // 连接池，连接的驱动任务落在哪个运行时，就随那个运行时一起没 —— 关掉这个运行时 + // 会掐断在途请求正在复用的连接（交接的端到端测试里，长流就断在提交那一刻）。 + // 所以交接之后它们在旧 core 上不停，只是多几次探测；要停得网关给一个开关 + // + // 模型目录：后台去问每个上游有哪些模型。**不挡启动** —— + // 探测要打网络，而网关不该因为一次探测慢而起不来。 + tw_gateway::models::spawn(s.control.gateway.clone()); + // 余额：启动时读一遍，之后按节奏读（见 `tw_gateway::balances`）。**不挡启动，不挡转发** + tw_gateway::balances::spawn(s.control.gateway.clone()); + } + self.rt = Some(rt); + Ok(sync) + } + + /// OAuth 写回的把手（交接开始前收尾用） + pub fn rotation(&self) -> Option> { + self.rotation.clone() + } + + /// 全部停下（交接提交之后）。**先在锁外收尾 OAuth 写回**（见模块文档），再调这个。 + pub fn stop(&mut self) { + self.rotation = None; + if let Some(rt) = self.rt.take() { + rt.shutdown_background(); + tracing::info!("background tasks stopped; the successor runs them from now on"); + } + } + + /// 进程退出时交出运行时，在异步上下文之外关掉（在里面丢掉一个运行时会 panic） + pub fn take_runtime(&mut self) -> Option { + self.rt.take() + } +} + +/// 收尾 OAuth 写回。 +pub async fn settle_rotation(bg: &std::sync::Mutex) { + let r = bg.lock().ok().and_then(|g| g.rotation()); + if let Some(r) = r { + r.settle(ROTATION_SETTLE).await; + } +} + +/* + 日志回收。**启动先跑一次，之后每小时一次。** + + 只在启动时跑的话，一台一直开着的机器永远轮不到它；只定时跑的 + 话，关机过夜的那台永远等不到那一小时。两个都要。 + + 每一轮都重新读配置，所以改完期限不用重启 —— 下一轮就按新的算。 + 扫的是目录和一条 DELETE，空跑一次的代价可以忽略。 + + **在阻塞线程上跑，删正文时不拿记录库的锁**（见 `tw_store::task::gc`）：一次删几个 + GB 可以要好几秒，以前那期间这把锁一直攥着，记录和查询都在等。 + + **接替者等请求库真的打开了才跑**：之前那是个占位的内存库，正文目录却是真的，旧 core + 还在用它。 +*/ +fn spawn_gc(gw: tw_gateway::AppState, rec: Store, mut opened: watch::Receiver) { + tokio::spawn(async move { + let _ = opened.wait_for(|o| *o).await; + loop { + let r = gw.config().retention.clone(); + let freed = { + 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!( + freed_bytes = freed, + body_days = r.body_days, + row_days = r.row_days, + "reclaimed expired request logs" + ); + } + tokio::time::sleep(Duration::from_secs(3600)).await; + } + }); +} diff --git a/bin/twcore/src/handover.rs b/bin/twcore/src/handover.rs new file mode 100644 index 00000000..e8ef844a --- /dev/null +++ b/bin/twcore/src/handover.rs @@ -0,0 +1,1221 @@ +//! 升级不断线：新旧 core 交接的状态机（方案第 4、6、7 节）。 +//! +//! 两个角色，同一个二进制： +//! +//! - **旧 core**（[`serve_arms`]）：桌面端 `POST /handover/arm` 之后脱离父进程、开交接端点、 +//! 等接替者；接替者连上来就把控制面、快照、监听交过去，提交之后停止 accept、停掉后台 +//! 事务、放弃锁，做完手上的请求再退出。 +//! - **新 core**（[`approach`] → [`Incoming`]）：取锁时发现持有者活着，就去连它的交接端点； +//! 连上了就在交过来的控制面上服务、读配置建运行时、要监听、提交，成为锁的持有者。 +//! +//! **提交是唯一的不可逆点**，以旧 core 回的 CommitAck 为准（它在时限之后收到的提交不回）。提交之前任何一步出错，旧 core 回到待交接、照常服务（它始终 +//! 握着全部监听）；提交之后新 core 是唯一的「主」，旧 core 只做完手上的请求。任何情况下 +//! 都不出现两个进程都不监听的时刻。 +//! +//! # 时限 +//! +//! | 时限 | 多久 | 谁在等 | +//! | --- | --- | --- | +//! | 待交接 | arm 给的 `wait_secs`，缺省 120 秒 | 旧 core 等接替者连上来（H7） | +//! | 接管 | 30 秒（[`TAKEOVER_LIMIT`]） | 旧 core 交出监听之后等提交（F4） | +//! | 确认 | 10 秒（[`ACK_WITHIN`]） | 新 core 发出提交之后等 CommitAck | +//! | 排空 | 30 分钟（[`DRAIN_LIMIT`]） | 旧 core 提交之后等手上的请求（H6） | +//! | 退回 | 180 秒（[`FALLBACK_DRAIN`]） | 版本不同时旧 core 等请求结束再退（H2） | +//! +//! **接管时限从交出监听算起，不从 hello 算起**（和方案 9.4 的写法不同）：新 core 配置读 +//! 不进来时（H3），它只接管控制面、等用户把配置修好，这可以是几分钟、几小时 —— 这段 +//! 时间旧 core 照常转发，没有理由让交接失败。这一段的失败（新 core 崩了）靠交接连接断开 +//! 发现。 + +use std::net::SocketAddr; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use anyhow::{Context, Result}; +use tokio::sync::watch; +use tokio::time::Instant; +use tw_api::HandoverState; +use tw_handover::{Link, ListenerInfo, Message, Role, Socket}; + +use crate::background::{self, Background}; +use crate::lockfile::LockFile; + +/// 旧 core 交出监听之后等提交多久（F4）。 +pub const TAKEOVER_LIMIT: Duration = Duration::from_secs(30); +/// 新 core 收到监听之后，自己最晚什么时候提交。**比旧 core 等的短一截**：旧 core 等到点 +/// 放弃、新 core 同一刻发出提交的话，两边都会以为自己是主 +const COMMIT_WITHIN: Duration = Duration::from_secs(20); +/// 旧 core 提交之后排空多久（H6）。 +pub const DRAIN_LIMIT: Duration = Duration::from_secs(30 * 60); +/// 退回现行流程时旧 core 等请求结束多久（H2；和桌面端现行的等待一样长）。 +pub const FALLBACK_DRAIN: Duration = Duration::from_secs(180); +/// 连上之后对方多久不说话就算了：新 core 等旧 core 的第一句、旧 core 等 `Hello`。 +const HELLO_WITHIN: Duration = Duration::from_secs(10); +/// 交出 access token 之前等正在进行的续期换完最多多久 +const RENEWAL_SETTLE: Duration = Duration::from_secs(10); +/// 新 core 等快照多久。**要盖得住旧 core 发快照之前的收尾**：OAuth 写回最多 5 秒、上传本机 +/// 改动最多 8 秒（`sync::FLUSH_WITHIN`）、暂停续期最多 [`RENEWAL_SETTLE`]，再留余量 +const SNAPSHOT_WITHIN: Duration = Duration::from_secs(30); +/// 新 core 发出提交之后等 CommitAck 多久 +const ACK_WITHIN: Duration = Duration::from_secs(10); +/// 连着几次交接在交出监听之后失败（F4、F5），下一个接替者就请它退回现行流程（F5）。 +pub const MAX_STRIKES: u32 = 3; +/// 新 core 请旧 core 退回之后，等它退出多久（[`FALLBACK_DRAIN`] 加上退出的余量）。 +const FALLBACK_EXIT_WITHIN: Duration = Duration::from_secs(200); + +/// 接替者的请求号从旧 core 的号往后跳这么多。 +/// +/// **两个进程同时在发请求号**：旧 core 排空期间还在记自己的请求，新 core 的记录要等它退出 +/// 才进库（见 `record`），库写的是 `INSERT OR REPLACE`，号撞了就是覆盖。快照里带着旧 core +/// 此刻的号，新 core 从它往后跳一大截：旧 core 从快照到退出用掉的号（交接前的重叠、排空 +/// 期间的事件）远到不了这么多。 +pub const ID_GAP: u64 = 10_000_000; + +// ─────────────────────────────────────────────────────────── 快照外面那层 + +/// 交接连接上的快照：网关的那份（格式归 `tw_gateway::handover`）外面包一层 core 自己的。 +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +struct Snapshot { + gateway: serde_json::Value, + /// 旧 core 此刻发到了哪个请求号（见 [`ID_GAP`]） + next_id: u64, +} + +fn wrap(gateway: serde_json::Value, next_id: u64) -> serde_json::Value { + serde_json::to_value(Snapshot { gateway, next_id }).unwrap_or(serde_json::Value::Null) +} + +/// 拆开。**认不出就当没有**：版本号相同时不该发生，发生了也只是少接一样状态 +fn unwrap(v: serde_json::Value) -> Option { + serde_json::from_value(v).ok() +} + +// ─────────────────────────────────────────────────────────── 待交接的时钟（纯逻辑） + +/// 待交接时等不到接替者怎么办（H7）。 +#[derive(Debug, Clone, PartialEq)] +pub struct ArmClock { + pub deadline: Instant, + wait: Duration, + can_relaunch: bool, + relaunched: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Expiry { + /// 重新打开桌面端一次，再等一个周期 + Relaunch, + /// 不等了：按现行方式排空后退出 + GiveUp, +} + +impl ArmClock { + pub fn new(now: Instant, wait: Duration, can_relaunch: bool) -> Self { + Self { + deadline: now + wait, + wait, + can_relaunch, + relaunched: false, + } + } + + /// 到点了。**只重开一次** + pub fn expired(&mut self, now: Instant) -> Expiry { + if self.can_relaunch && !self.relaunched { + self.relaunched = true; + self.deadline = now + self.wait; + Expiry::Relaunch + } else { + Expiry::GiveUp + } + } + + /// 一次交接没成、回到待交接：重新等一个周期（重开过的不再重开） + pub fn restart(&mut self, now: Instant) { + self.deadline = now + self.wait; + } +} + +// ─────────────────────────────────────────────────────────── 排空判定（纯逻辑） + +/// 排空期什么时候算做完：网关的在途计数归零、走 HTTP 的连接都结束，**连着几拍都是这样** +/// （WebSocket 以 1012 关掉的那一帧、最后一条结算要有工夫发出去）。 +#[derive(Debug, Default)] +pub struct Drained { + calm: u32, + last: Option, +} + +/// 连着几拍（每拍 [`DRAIN_TICK`]） +const CALM_TICKS: u32 = 3; +const DRAIN_TICK: Duration = Duration::from_millis(200); + +impl Drained { + /// 这一拍看到的在途数和还在服务的连接数。返回（要不要报一条新的在途数，做完没有） + pub fn tick(&mut self, in_flight: usize, serving: usize) -> (Option, bool) { + let report = (self.last != Some(in_flight)).then_some(in_flight); + self.last = Some(in_flight); + if in_flight == 0 && serving == 0 { + self.calm += 1; + } else { + self.calm = 0; + } + (report, self.calm >= CALM_TICKS) + } +} + +// ─────────────────────────────────────────────────────────── 旧 core + +/// 进程为什么该退了。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Exit { + /// 交给了接替者、手上的做完了。后台事务早已停下，**不再同步、不碰配置** + HandedOver, + /// 没交出去（等不到接替者、版本不同）：按现行方式排空后退出，这时还是主 + Retired, +} + +/// 两个角色都要碰的那些东西。`main` 建一份，交接的任务拿一份。 +pub struct Core { + pub dir: PathBuf, + pub state: tw_gateway::AppState, + pub control: tw_control::ControlState, + pub gateway: tw_gateway::listen::Handoff, + pub background: Arc>, + pub sync: Mutex>>, + /// 配置文件、插件目录的监听。**丢掉就不再盯**：交接提交之后旧 core 不再读配置 + pub watchers: Mutex>>, + pub lock: Mutex>, + /// `--parent` 守望。arm 之后关掉：桌面端要重启，core 不能跟着退 + pub parent: Mutex>>, + /// 原来的父进程（`--parent`）。arm 之后等到点时它还在，说明桌面端没有重启（更新失败、 + /// 取消了），撤销 arm（见 [`run_arms`]） + pub parent_pid: Option, + /// 请求库打开了没有。接替者要等它的前任退出才打开（见 `record::open_later`）； + /// **它自己排空完也要等这个**才说 `Bye`（H8：库只在最老的那个退出之后才由下一个打开） + pub store_opened: watch::Receiver, + /// 交了出去 + pub handed_over: AtomicBool, +} + +impl Core { + fn handover(&self) -> &tw_control::handover::Handover { + &self.control.handover + } + + /// 同步服务（这个进程是接替者时，提交之后才有） + fn sync(&self) -> Option> { + self.sync.lock().ok().and_then(|g| g.clone()) + } +} + +/// 跑旧 core 那一半：接 arm、等接替者、交接、排空。返回的接收端在该退出时收到 [`Exit`]。 +pub fn serve_arms(core: Arc) -> tokio::sync::oneshot::Receiver { + let (tx, rx) = tokio::sync::oneshot::channel(); + let mut arms = core.handover().arm_requests(); + tokio::spawn(async move { + let exit = run_arms(&core, &mut arms).await; + let _ = tx.send(exit); + }); + rx +} + +struct Params { + wait: Duration, + relaunch: Option, +} + +async fn run_arms( + core: &Arc, + arms: &mut tokio::sync::mpsc::Receiver, +) -> Exit { + loop { + let (endpoint, params) = first_arm(core, arms).await; + match armed(core, arms, endpoint, params).await { + Some(exit) => return exit, + // 撤销了：回到平常，等下一次 arm + None => continue, + } + } +} + +/// 等第一次 arm：开端点、脱离父进程。 +async fn first_arm( + core: &Core, + arms: &mut tokio::sync::mpsc::Receiver, +) -> (tw_handover::Endpoint, Params) { + loop { + let Some(a) = arms.recv().await else { + return std::future::pending().await; + }; + match tw_handover::Endpoint::bind(&core.dir).await { + Ok(ep) => { + // **先把标准输出、错误输出从桌面端的管道上挪开**：桌面端马上要退出，那之后再 + // 往管道里写会出错，`eprintln!` 和日志的出错处理会因此 panic，排空中的旧 core + // 就挂了 + stdio_to_file(&core.dir); + // **自己另起一个会话**：桌面端退出时给整个进程组发的信号（SIGTERM、SIGHUP） + // 不能把等着交接、排空中的 core 一起带走 + leave_process_group(); + if let Some(h) = core.parent.lock().ok().and_then(|mut g| g.take()) { + h.abort(); + tracing::info!("armed for a handover; no longer exiting with the parent"); + } + core.handover().set_phase(HandoverState::Armed); + let _ = a.reply.send(Ok(())); + return ( + ep, + Params { + wait: a.wait, + relaunch: a.relaunch, + }, + ); + } + Err(e) => { + tracing::warn!("the handover endpoint could not be opened: {e}"); + let _ = a.reply.send(Err(tw_control::handover::endpoint_failed(&e))); + } + } + } +} + +/// 待交接：等接替者、交接。`None` = 撤销了 arm（桌面端没有重启），回到平常 +async fn armed( + core: &Arc, + arms: &mut tokio::sync::mpsc::Receiver, + endpoint: tw_handover::Endpoint, + mut params: Params, +) -> Option { + let mut endpoint = Some(endpoint); + let mut clock = ArmClock::new(Instant::now(), params.wait, params.relaunch.is_some()); + let mut strikes = 0u32; + loop { + if endpoint.is_none() { + match tw_handover::Endpoint::bind(&core.dir).await { + Ok(ep) => endpoint = Some(ep), + // 开不回来：等不到接替者，到点照 H7 收场 + Err(e) => tracing::warn!("the handover endpoint could not be reopened: {e}"), + } + } + let accepted = async { + match endpoint.as_ref() { + Some(ep) => ep.accept().await, + None => std::future::pending().await, + } + }; + tokio::select! { + r = accepted => { + let link = match r { + Ok(l) => l, + Err(e) => { + tracing::warn!("the handover endpoint failed: {e}"); + endpoint = None; + tokio::time::sleep(Duration::from_secs(1)).await; + continue; + } + }; + // **接到一个就关掉端点**（H8）：同一时刻至多一个接替者；再来的连不上，按「已有 + // 实例在运行」退避重试 + if let Some(ep) = endpoint.take() { + ep.close(); + } + match give(core, link, strikes >= MAX_STRIKES).await { + Given::Committed(link) => { + drain(core, link).await; + return Some(Exit::HandedOver); + } + Given::Fallback => { + // 接替者自己绑端口、起来之后才续期；这边排空期间仍是唯一在转发的那个 + core.state.oauth.hold_renewals(false); + tracing::info!("the successor cannot take over; draining and exiting the current way"); + retire(core, FALLBACK_DRAIN).await; + return Some(Exit::Retired); + } + Given::Abandoned { strike } => { + if strike { + strikes += 1; + } + // 接管没成：续期回到这一边 + core.state.oauth.hold_renewals(false); + // 回到待交接：控制面接着接，重新等一个周期（F4） + core.handover().control().resume(); + core.handover().set_phase(HandoverState::Armed); + clock.restart(Instant::now()); + tracing::info!(strikes, "the handover did not complete; serving on and waiting for the next successor"); + } + Given::NotASuccessor => {} + } + } + _ = tokio::time::sleep_until(clock.deadline) => { + // 桌面端根本没重启（更新失败、取消了）：**撤销 arm**，不重开它、不退出 + if let Some(ppid) = core.parent_pid + && crate::proc::alive(ppid) + { + tracing::info!(ppid, "the desktop app did not restart; the handover is called off"); + if let Some(ep) = endpoint.take() { + ep.close(); + } + if let Ok(mut g) = core.parent.lock() { + *g = Some(watch_parent(ppid)); + } + core.handover().set_phase(HandoverState::Idle); + return None; + } + match clock.expired(Instant::now()) { + Expiry::Relaunch => { + if let Some(r) = ¶ms.relaunch { + relaunch(r); + } + } + Expiry::GiveUp => { + // H7：不留孤儿进程 + tracing::info!("no successor arrived; draining and exiting the current way"); + if let Some(ep) = endpoint.take() { + ep.close(); + } + core.handover().set_phase(HandoverState::Idle); + retire(core, FALLBACK_DRAIN).await; + return Some(Exit::Retired); + } + } + } + Some(a) = arms.recv() => { + // 又 arm 了一次（桌面端重试）：按新的参数重新等 + params = Params { wait: a.wait, relaunch: a.relaunch }; + clock = ArmClock::new(Instant::now(), params.wait, params.relaunch.is_some()); + let _ = a.reply.send(Ok(())); + } + } + } +} + +/// 父进程守望：GUI 没了我们跟着退。轮询而不是用 kqueue，是因为 +/// 这段代码要能在被重新 parent 之后仍然正确 —— macOS 上收养我们的 +/// 是 launchd（getppid() 变成 1），Linux 桌面会话里多半是 +/// `systemd --user`（它是 subreaper，getppid() 变成它的 pid，不是 1）。 +/// 所以**不看 getppid()**，只看原来那个 pid 还在不在，哪个平台、 +/// 谁来收养都一样。 +pub fn watch_parent(ppid: u32) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + loop { + tokio::time::sleep(Duration::from_secs(2)).await; + if !crate::proc::alive(ppid) { + tracing::info!(ppid, "the parent exited; exiting with it"); + std::process::exit(0); + } + } + }) +} + +/// 另起一个会话（unix 的 `setsid`）。已经是进程组组长时会失败，只记一行 +#[cfg(unix)] +fn leave_process_group() { + // SAFETY: getsid/getpid/setsid 都不碰内存，只问、只改这个进程的会话和进程组 + let me = unsafe { libc::getpid() }; + if unsafe { libc::getsid(0) } == me { + // 已经是自己会话的头了（上一次 arm 时另起的） + return; + } + if unsafe { libc::setsid() } == -1 { + tracing::info!( + "could not start a session of its own: {}", + std::io::Error::last_os_error() + ); + } +} + +/// Windows 上没有进程组信号这回事：桌面端退出不会顺带结束 core +#[cfg(windows)] +fn leave_process_group() {} + +enum Given { + /// 提交了：从此只做完手上的请求 + Committed(Link), + /// 接替者请我们按现行方式排空后退出（H2） + Fallback, + /// 没成，回到待交接。`strike`：在交出监听之后失败的（F4、F5），记一次 + Abandoned { strike: bool }, + /// 连上来的不说 `Hello`（不是接替者，或者它已经放弃了） + NotASuccessor, +} + +/// 把控制面、快照、监听交给接替者，等它提交。 +async fn give(core: &Core, mut link: Link, refuse: bool) -> Given { + let hello = tokio::time::timeout(HELLO_WITHIN, link.recv()).await; + let (version, peer) = match hello { + Ok(Ok(( + Message::Hello { + version, + pid, + core_version, + }, + _, + ))) => { + tracing::info!(pid, %core_version, "a successor connected"); + // Windows 上 peer_pid 要等收到 Hello 才有（就是它报的那个） + (version, link.peer_pid().unwrap_or(pid)) + } + // **读不懂第一帧也回一句**：将来 Hello 改了样子，新的那个据此退回现行流程，而不是 + // 只看见连接断了 + Ok(Err(tw_handover::Error::Malformed(e))) => { + core.handover().set_phase(HandoverState::TakingOver); + return refuse_with( + &mut link, + format!("the greeting could not be read here ({e})"), + ) + .await; + } + _ => return Given::NotASuccessor, + }; + core.handover().set_phase(HandoverState::TakingOver); + if version != tw_handover::VERSION || refuse { + // H2 / F5：退回现行流程 + let reason = if refuse { + format!("{MAX_STRIKES} handovers in a row failed after the listeners were handed over") + } else { + format!( + "handover version {} here, {version} on the successor", + tw_handover::VERSION + ) + }; + return refuse_with(&mut link, reason).await; + } + + // ① 控制面。**交出去就不再从它上面接新连接**，已经连着的也断开（见 `tw_control::handover`） + let sock = match core.handover().control().export() { + Ok(s) => s, + Err(e) => { + tracing::warn!("the control plane listener could not be handed over: {e}"); + return Given::Abandoned { strike: true }; + } + }; + core.handover().control().pause(); + core.handover().control().close_connections(); + let control = Message::Control { + info: ListenerInfo { + role: Role::Control, + want: None, + }, + }; + if let Err(e) = link.send_with(&control, &[&sock], peer).await { + tracing::warn!("the control plane listener could not be handed over: {e}"); + return Given::Abandoned { strike: true }; + } + drop(sock); + + // ② 第一份快照。之前把还没上传的本机改动传上去（同现行退出时的 `sync.flush`） + background::settle_rotation(&core.background).await; + if let Some(sync) = core.sync() { + sync.flush().await; + } + let snapshot = wrap(core.state.handover_export(), core.state.bus.next_id()); + if let Err(e) = link.send(&Message::Snapshot { data: snapshot }).await { + tracing::warn!("the snapshot could not be sent to the successor: {e}"); + return Given::Abandoned { strike: true }; + } + let sent = Instant::now(); + + // ③ 等它要监听。**不限时**（H3：它的配置读不进来时在等人修），断了就是没成。**这一段照常 + // 续期**：这边还在转发，可能要等几个小时 + loop { + match link.recv().await { + Ok((Message::WantListeners, _)) => break, + Ok((Message::Fallback, _)) => return Given::Fallback, + Ok((other, _)) => { + tracing::warn!(?other, "an unexpected handover message; ignoring it") + } + Err(e) => { + // 快照刚发出去就断了：多半是它等快照等得不耐烦（或者起来就崩了），记一次； + // 等了很久才断的是在等配置修好的那一种（H3），不记 + let strike = sent.elapsed() < SNAPSHOT_WITHIN + Duration::from_secs(5); + tracing::info!( + strike, + "the successor left before taking the listeners: {e}" + ); + return Given::Abandoned { strike }; + } + } + } + + // ④ 第二份快照：**先暂停续期、等正在换的换完**，再收尾写回，再导出。交出去的 access + // token 和 refresh token 从此只有接替者在用（接管没成再恢复）。顺带也是此刻的全部状态： + // 接替者在 H3 里等了很久的话，第一份早就过时了，而且那时它是对着临时顶上的配置接的 + core.state.oauth.hold_renewals_settled(RENEWAL_SETTLE).await; + background::settle_rotation(&core.background).await; + let snapshot = wrap(core.state.handover_export(), core.state.bus.next_id()); + if let Err(e) = link.send(&Message::Snapshot { data: snapshot }).await { + tracing::warn!("the snapshot could not be sent to the successor: {e}"); + return Given::Abandoned { strike: true }; + } + + // ⑤ 监听：网关的每一个、远程控制端口 + let mut infos = Vec::new(); + let mut socks = Vec::new(); + match core.gateway.export() { + Ok(gw) => { + for (want, l) in gw { + infos.push(ListenerInfo { + role: Role::Gateway, + want: Some(want), + }); + socks.push(Socket::Tcp(l)); + } + } + Err(e) => { + tracing::warn!("the gateway listeners could not be handed over: {e}"); + return Given::Abandoned { strike: true }; + } + } + match core.control.remote.export() { + Ok(Some((want, l))) => { + infos.push(ListenerInfo { + role: Role::Remote, + want: Some(want), + }); + socks.push(Socket::Tcp(l)); + } + Ok(None) => {} + Err(e) => { + tracing::warn!("the remote control port could not be handed over: {e}"); + return Given::Abandoned { strike: true }; + } + } + let refs: Vec<&Socket> = socks.iter().collect(); + if let Err(e) = link + .send_with(&Message::Listeners { infos }, &refs, peer) + .await + { + tracing::warn!("the listeners could not be handed over: {e}"); + return Given::Abandoned { strike: true }; + } + drop(refs); + drop(socks); + + // ⑥ 等提交（F4：30 秒）。提交之前两边同时 accept 是允许的（H4）。**过了时限就不再收**： + // 这条连接随之关掉，迟到的 Commit 没人回 CommitAck，接替者自己退出 + match tokio::time::timeout(TAKEOVER_LIMIT, link.recv()).await { + Ok(Ok((Message::Commit, _))) => { + // 先回 CommitAck 再停止 accept：回不出去（接替者已经没了）就当没提交，这边的 + // 监听还都在，接着服务。回出去了才停 —— 接替者收到它才把自己当主 + match link.send(&Message::CommitAck).await { + Ok(()) => Given::Committed(link), + Err(e) => { + tracing::warn!("the commit could not be acknowledged: {e}"); + Given::Abandoned { strike: true } + } + } + } + Ok(Ok((Message::Fallback, _))) => Given::Fallback, + Ok(Ok((other, _))) => { + tracing::warn!(?other, "expected a commit from the successor"); + Given::Abandoned { strike: true } + } + Ok(Err(e)) => { + tracing::info!("the successor left before committing: {e}"); + Given::Abandoned { strike: true } + } + Err(_) => { + tracing::warn!("the successor did not commit within {TAKEOVER_LIMIT:?}"); + Given::Abandoned { strike: true } + } + } +} + +/// 拒绝这一次交接（H2、F5、读不懂的招呼），等它说一声 Fallback（它会马上说）。 +async fn refuse_with(link: &mut Link, reason: String) -> Given { + tracing::info!(%reason, "refusing the handover"); + let _ = link.send(&Message::Refused { reason }).await; + match tokio::time::timeout(HELLO_WITHIN, link.recv()).await { + Ok(Ok((Message::Fallback, _))) => Given::Fallback, + _ => Given::Abandoned { strike: false }, + } +} + +/// 提交之后：停止 accept、停掉后台事务、放弃锁，做完手上的请求，说 `Bye`。 +async fn drain(core: &Core, mut link: Link) { + tracing::info!("committed; the successor takes new requests from now on"); + core.handed_over.store(true, Ordering::SeqCst); + // 停止 accept：**只丢自己的句柄**，内核队列里没被接受的连接由接替者取走 + core.gateway.stop_accepting(); + // 远程控制端口：从它进来的连接现在断开，远端按断线重连连到接替者 + core.control.remote.stop_accepting(); + let mut settled = core.state.begin_drain(); + background::settle_rotation(&core.background).await; + if let Ok(mut bg) = core.background.lock() { + bg.stop(); + } + // 不再读配置：新版本的配置格式可能变了 + if let Ok(mut w) = core.watchers.lock() { + w.clear(); + } + // 锁归接替者：退出时不删文件 + if let Some(l) = core.lock.lock().ok().and_then(|mut g| g.take()) { + l.disown(); + } + + let deadline = Instant::now() + DRAIN_LIMIT; + let mut drained = Drained::default(); + let mut tick = tokio::time::interval(DRAIN_TICK); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + Some(s) = settled.recv() => { + // 接替者没了（F7）也照样排空，发不出去就算了 + let _ = link.send(&Message::Settled { data: s }).await; + } + _ = tick.tick() => { + let (report, done) = drained.tick(core.state.live.count(), core.gateway.serving()); + if let Some(n) = report { + let _ = link.send(&Message::InFlight { n: n as u64 }).await; + } + if done { + break; + } + } + _ = tokio::time::sleep_until(deadline) => { + // H6：到点了还没结束的中断掉，记成中断 + let n = abort_all(&core.state); + tracing::warn!(n, "the drain limit passed; aborting the requests still running"); + // 让结局进记录 + tokio::time::sleep(Duration::from_secs(2)).await; + break; + } + } + } + while let Ok(s) = settled.try_recv() { + let _ = link.send(&Message::Settled { data: s }).await; + } + let _ = link.send(&Message::InFlight { n: 0 }).await; + // H8：我们自己也是接替者、前任还没退的话，等它（请求库要由最老的那个退出之后才打开） + let mut opened = core.store_opened.clone(); + if !*opened.borrow() { + tracing::info!("waiting for the core before this one to exit before leaving"); + let _ = tokio::time::timeout(DRAIN_LIMIT, opened.wait_for(|o| *o)).await; + } + let _ = link.send(&Message::Bye).await; + tracing::info!("drained; exiting"); +} + +/// 没交出去：照常 accept，等手上的请求结束（最多 `limit`），然后退出。和桌面端现行的 +/// 「等请求结束再重启」一样。 +async fn retire(core: &Core, limit: Duration) { + let deadline = Instant::now() + limit; + while core.state.live.count() > 0 && Instant::now() < deadline { + tokio::time::sleep(DRAIN_TICK).await; + } +} + +/// 中止所有在跑的请求。返回几个 +fn abort_all(state: &tw_gateway::AppState) -> usize { + state + .bus + .in_flight() + .requests + .iter() + .filter(|r| state.aborts.request(r.id).is_ok()) + .count() +} + +/// arm 之后的日志写到哪：数据目录里，**接着往后写**（上一次交接的日志留着，排查时要对照）。 +pub const HANDOVER_LOG: &str = "twcore.handover.log"; +/// 它长过这么多就挪成上一份（`.log.1`） +const HANDOVER_LOG_MAX: u64 = 4 * 1024 * 1024; + +/// 标准输出、错误输出改写进 [`HANDOVER_LOG`]（开不了就丢掉），标准输入接空设备。**再也 +/// 不碰桌面端的管道**：它退出之后写管道会出错。 +fn stdio_to_file(dir: &Path) { + let path = dir.join(HANDOVER_LOG); + // 接着写，但不无限长：过了几 MB 就把它挪成上一份，从头开始 + if std::fs::metadata(&path).is_ok_and(|m| m.len() > HANDOVER_LOG_MAX) { + let _ = std::fs::rename(&path, path.with_extension("log.1")); + } + let file = { + let mut o = std::fs::OpenOptions::new(); + o.append(true).create(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + o.mode(0o600); + } + o.open(&path) + }; + let file = match file { + Ok(f) => f, + Err(e) => { + tracing::warn!(path = %path.display(), "the handover log could not be opened; output is discarded: {e}"); + match std::fs::OpenOptions::new().write(true).open(null_device()) { + Ok(f) => f, + Err(_) => return, + } + } + }; + redirect(file); +} + +#[cfg(unix)] +fn null_device() -> &'static str { + "/dev/null" +} + +#[cfg(windows)] +fn null_device() -> &'static str { + "NUL" +} + +#[cfg(unix)] +fn redirect(file: std::fs::File) { + use std::os::fd::AsRawFd; + let fd = file.as_raw_fd(); + // SAFETY: dup2 把 1、2 号描述符原子地换成这个文件；之后 `file` 关掉自己那个也不影响 + unsafe { + libc::dup2(fd, libc::STDOUT_FILENO); + libc::dup2(fd, libc::STDERR_FILENO); + if let Ok(null) = std::fs::File::open("/dev/null") { + libc::dup2(null.as_raw_fd(), libc::STDIN_FILENO); + } + } +} + +#[cfg(windows)] +fn redirect(file: std::fs::File) { + use std::os::windows::io::IntoRawHandle; + use windows_sys::Win32::System::Console::{STD_ERROR_HANDLE, STD_OUTPUT_HANDLE, SetStdHandle}; + // 句柄交给进程的标准句柄表，**不再关**（std 每次写都现取标准句柄） + let h = file.into_raw_handle(); + // SAFETY: 换的是本进程的标准句柄；h 是一个开着的文件句柄，之后一直开着 + unsafe { + SetStdHandle(STD_OUTPUT_HANDLE, h as _); + SetStdHandle(STD_ERROR_HANDLE, h as _); + } +} + +/// 重新打开桌面端一次（H7）。**和 core 脱开**：core 随后退出时不能把它带走。 +fn relaunch(r: &tw_api::Relaunch) { + use std::process::{Command, Stdio}; + let mut cmd = match r { + tw_api::Relaunch::BundleId { id } => { + let mut c = Command::new("open"); + c.arg("-b").arg(id); + c + } + tw_api::Relaunch::Exe { path, args } => { + let mut c = Command::new(path); + c.args(args); + detach(&mut c); + c + } + }; + cmd.stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()); + match cmd.spawn() { + Ok(mut child) => { + tracing::info!(?r, "no successor yet; reopening the desktop app once"); + // 收尸，不留僵尸进程 + std::thread::spawn(move || { + let _ = child.wait(); + }); + } + Err(e) => tracing::warn!(?r, "the desktop app could not be reopened: {e}"), + } +} + +#[cfg(unix)] +fn detach(c: &mut std::process::Command) { + use std::os::unix::process::CommandExt; + // SAFETY: setsid 是 async-signal-safe 的，fork 之后、exec 之前调它是标准做法 + unsafe { + c.pre_exec(|| { + libc::setsid(); + Ok(()) + }); + } +} + +#[cfg(windows)] +fn detach(c: &mut std::process::Command) { + use std::os::windows::process::CommandExt; + const DETACHED_PROCESS: u32 = 0x0000_0008; + const CREATE_NEW_PROCESS_GROUP: u32 = 0x0000_0200; + c.creation_flags(DETACHED_PROCESS | CREATE_NEW_PROCESS_GROUP); +} + +// ─────────────────────────────────────────────────────────── 新 core + +/// 去连锁的持有者，结果是三种之一。 +pub enum Approach { + /// 它没有 arm（或者正忙着交给别人）：照旧报「已有实例在运行」 + NotArmed, + /// 连上了，控制面和快照到手 + Proceed(Incoming), + /// 它已经退了（退回现行流程，或者交出监听之前就崩了，F6）：照常启动 + StartFresh, +} + +/// 接替者手上的交接连接。 +pub struct Incoming { + link: Link, + /// 控制面的监听（旧 core 交过来的） + pub control: Option, + snapshot: Option, + /// 旧 core 的 pid（锁文件里的那个） + pub predecessor: u32, + /// 这边的请求号从这里起（旧 core 的号 + [`ID_GAP`]）。旧 core 用的都在它之下 + pub ids_from: u64, +} + +/// 旧 core 交过来的网关与远程控制端口的监听。 +pub struct Inherited { + pub gateway: Vec<(SocketAddr, std::net::TcpListener)>, + pub remote: Option<(SocketAddr, std::net::TcpListener)>, +} + +/// 新 core 取锁时发现持有者活着：连它的交接端点（H1）。 +pub async fn approach(dir: &Path, holder: u32) -> Approach { + let mut link = match tw_handover::connect(dir).await { + Ok(Some(l)) => l, + Ok(None) => return Approach::NotArmed, + Err(e) => { + tracing::warn!("the handover endpoint is there but could not be used: {e}"); + return Approach::NotArmed; + } + }; + let hello = Message::Hello { + version: tw_handover::VERSION, + core_version: env!("CARGO_PKG_VERSION").to_string(), + pid: std::process::id(), + }; + if let Err(e) = link.send(&hello).await { + tracing::warn!("the running core did not take the handover greeting: {e}"); + return gone_or_busy(holder).await; + } + let control = match tokio::time::timeout(HELLO_WITHIN, link.recv()).await { + Ok(Ok((Message::Control { .. }, mut socks))) if socks.len() == 1 => socks.pop(), + Ok(Ok((Message::Refused { reason }, _))) => { + tracing::info!(%reason, "the running core cannot hand over; asking it to drain and exit the current way"); + return fall_back(link, holder).await; + } + // 认不出它的话：版本不同的另一种样子（H2） + Ok(Err(tw_handover::Error::Malformed(e))) => { + tracing::info!( + "the running core speaks another handover version ({e}); asking it to drain and exit the current way" + ); + return fall_back(link, holder).await; + } + Ok(Ok((other, _))) => { + tracing::warn!(?other, "an unexpected first handover message"); + return fall_back(link, holder).await; + } + Ok(Err(e)) => { + tracing::info!("the running core closed the handover: {e}"); + return gone_or_busy(holder).await; + } + // 不说话：它正在交给另一个接替者 + Err(_) => return Approach::NotArmed, + }; + // 旧 core 发快照之前要收尾（写回、上传本机改动），等得够久（见 SNAPSHOT_WITHIN） + let snapshot = match tokio::time::timeout(SNAPSHOT_WITHIN, link.recv()).await { + Ok(Ok((Message::Snapshot { data }, _))) => unwrap(data), + other => { + tracing::warn!(ok = other.is_ok(), "the snapshot did not arrive"); + // 控制面的套接字丢掉：旧的那个回到待交接、接着用它 + drop(control); + return gone_or_busy(holder).await; + } + }; + Approach::Proceed(Incoming { + link, + control, + snapshot, + predecessor: holder, + ids_from: 0, + }) +} + +/// 请旧 core 按现行方式排空后退出，等它退（H2）。 +async fn fall_back(mut link: Link, holder: u32) -> Approach { + let _ = link.send(&Message::Fallback).await; + if wait_gone(holder, FALLBACK_EXIT_WITHIN).await { + Approach::StartFresh + } else { + tracing::warn!(pid = holder, "the running core did not exit in time"); + Approach::NotArmed + } +} + +/// 交接连接没走下去：旧 core 没了就照常启动（F6），还在就算「已有实例」（守护退避重试）。 +async fn gone_or_busy(holder: u32) -> Approach { + if wait_gone(holder, Duration::from_secs(2)).await { + Approach::StartFresh + } else { + Approach::NotArmed + } +} + +/// 前任还在排空、却没有交接连接（它的接替者先退出了，这个进程是照常起来的）：状态报 +/// 「排空中」，它们都退了回到平常。在途数报不出来（没有连接），是 0 +pub fn follow_vanished(pids: Vec, handover: Arc) { + handover.set_phase(HandoverState::Draining); + tokio::spawn(async move { + for pid in pids { + let _ = wait_gone(pid, DRAIN_LIMIT + Duration::from_secs(5 * 60)).await; + } + if handover.phase() == HandoverState::Draining { + handover.set_phase(HandoverState::Idle); + } + }); +} + +async fn wait_gone(pid: u32, within: Duration) -> bool { + let deadline = Instant::now() + within; + loop { + if !crate::proc::alive(pid) { + return true; + } + if Instant::now() >= deadline { + return false; + } + tokio::time::sleep(Duration::from_millis(200)).await; + } +} + +impl Incoming { + /// 接进旧 core 的快照。请求号从它的号往后跳（见 [`ID_GAP`]） + pub fn import(&mut self, state: &tw_gateway::AppState) { + let Some(s) = self.snapshot.take() else { + tracing::warn!( + "the snapshot from the previous core was not readable; starting from scratch" + ); + return; + }; + self.ids_from = self.ids_from.max(s.next_id.saturating_add(ID_GAP)); + state.bus.resume_after(self.ids_from); + state.handover_import(&s.gateway); + } + + /// 要监听（配置读进来了、运行时建好了）。**先收到第二份快照**：旧 core 暂停了续期之后 + /// 导出的那一份，这时接进来（盖掉第一份；H3 时第一份还是对着临时顶上的配置接的） + pub async fn listeners(&mut self, state: &tw_gateway::AppState) -> Result { + self.link + .send(&Message::WantListeners) + .await + .context("the previous core left before handing over the listeners")?; + let (m, _) = tokio::time::timeout(SNAPSHOT_WITHIN, self.link.recv()) + .await + .context("the previous core did not send the snapshot in time")? + .context("the previous core left before handing over the listeners")?; + let Message::Snapshot { data } = m else { + anyhow::bail!("expected a snapshot from the previous core, got {m:?}"); + }; + self.snapshot = unwrap(data); + self.import(state); + let (m, socks) = tokio::time::timeout(HELLO_WITHIN, self.link.recv()) + .await + .context("the previous core did not hand over the listeners in time")? + .context("the previous core left before handing over the listeners")?; + let Message::Listeners { infos } = m else { + anyhow::bail!("expected the listeners from the previous core, got {m:?}"); + }; + inherited(infos, socks) + } + + /// 提交（H5）。**旧 core 回了 CommitAck、或者它的进程已经没了（F6，套接字已经在手上） + /// 才算数**；它还活着却没回（它等提交的时限已过、回到了待交接），这边就不能当主 —— + /// `Err`，调用方退出，守护重来一次 + pub async fn commit(&mut self) -> Result<()> { + let sent = self.link.send(&Message::Commit).await; + if sent.is_ok() { + loop { + match tokio::time::timeout(ACK_WITHIN, self.link.recv()).await { + Ok(Ok((Message::CommitAck, _))) => return Ok(()), + Ok(Ok((other, _))) => { + tracing::debug!(?other, "ignoring a handover message before the commit ack") + } + _ => break, + } + } + } + if wait_gone(self.predecessor, Duration::from_secs(2)).await { + tracing::warn!("the previous core is gone at commit; serving on without it"); + return Ok(()); + } + anyhow::bail!("the previous core did not acknowledge the commit and is still running") + } + + /// 提交之后跟着旧 core 排空：结算、在途数、`Bye`。连接断了（旧 core 退出或崩溃）就放掉 + /// 快照里记着的旧 core 的名额，等它的进程没了再回到平常。 + pub fn follow_drain( + mut self, + state: tw_gateway::AppState, + handover: Arc, + ) { + tokio::spawn(async move { + loop { + match self.link.recv().await { + Ok((Message::Settled { data }, _)) => state.handover_settled(&data), + Ok((Message::InFlight { n }, _)) => handover.set_draining(n as usize), + Ok((Message::Bye, _)) => break, + Ok((other, _)) => tracing::debug!(?other, "ignoring a handover message"), + Err(_) => break, + } + } + // **等它的进程真没了再放**：说完 Bye 到退出之间、或者连接异常断开（它的进程可能还 + // 在跑、手上还有请求），名额要等它真的不可能再用才算放开。最多等到它的排空上限 + if !wait_gone(self.predecessor, DRAIN_LIMIT + Duration::from_secs(5 * 60)).await { + tracing::warn!( + pid = self.predecessor, + "the previous core is still running past its drain limit" + ); + } + state.handover_release_all(); + handover.set_draining(0); + // 期间又被 arm 了（H8）就别改回去 + if handover.phase() == HandoverState::Draining { + handover.set_phase(HandoverState::Idle); + } + tracing::info!("the previous core has finished draining"); + }); + } + + /// 等交接连接断开（安全模式的接替者用：它不要监听，只是一直握着连接） + pub async fn hold(mut self) { + while self.link.recv().await.is_ok() {} + } +} + +/// 按角色分开收到的监听（`infos` 和套接字按顺序一一对应）。 +fn inherited(infos: Vec, socks: Vec) -> Result { + anyhow::ensure!( + infos.len() == socks.len(), + "{} listeners described, {} handed over", + infos.len(), + socks.len() + ); + let mut out = Inherited { + gateway: Vec::new(), + remote: None, + }; + for (info, s) in infos.into_iter().zip(socks) { + // Windows 上只有 TCP 一种，这个 match 在那儿只有一支 + #[allow(clippy::infallible_destructuring_match)] + let l = match s { + Socket::Tcp(l) => l, + #[cfg(unix)] + Socket::Unix(_) => { + anyhow::bail!( + "a {:?} listener was handed over as a unix socket", + info.role + ) + } + }; + let want = info + .want + .with_context(|| format!("a {:?} listener without its address", info.role))?; + match info.role { + Role::Gateway => out.gateway.push((want, l)), + Role::Remote => out.remote = Some((want, l)), + Role::Control => anyhow::bail!("a control listener among the listeners"), + } + } + Ok(out) +} + +/// 新 core 收到监听之后最晚什么时候提交（见 [`COMMIT_WITHIN`]） +pub fn commit_within() -> Duration { + COMMIT_WITHIN +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_arm_clock_relaunches_once_then_gives_up() { + let t0 = Instant::now(); + let w = Duration::from_secs(120); + let mut c = ArmClock::new(t0, w, true); + assert_eq!(c.deadline, t0 + w); + assert_eq!(c.expired(t0 + w), Expiry::Relaunch); + assert_eq!(c.deadline, t0 + w * 2, "重开之后再等一个周期"); + assert_eq!(c.expired(t0 + w * 2), Expiry::GiveUp, "只重开一次"); + } + + #[test] + fn without_a_relaunch_the_arm_clock_gives_up_at_once() { + let t0 = Instant::now(); + let mut c = ArmClock::new(t0, Duration::from_secs(5), false); + assert_eq!(c.expired(t0), Expiry::GiveUp); + } + + #[test] + fn a_failed_takeover_waits_a_fresh_period_but_does_not_relaunch_again() { + let t0 = Instant::now(); + let w = Duration::from_secs(10); + let mut c = ArmClock::new(t0, w, true); + assert_eq!(c.expired(t0 + w), Expiry::Relaunch); + c.restart(t0 + w * 5); + assert_eq!(c.deadline, t0 + w * 6); + assert_eq!(c.expired(c.deadline), Expiry::GiveUp); + } + + #[test] + fn draining_reports_changes_and_needs_a_few_calm_ticks() { + let mut d = Drained::default(); + assert_eq!(d.tick(2, 1), (Some(2), false)); + assert_eq!(d.tick(2, 1), (None, false), "没变不报"); + assert_eq!(d.tick(0, 1), (Some(0), false), "还有连接没结束"); + assert_eq!(d.tick(0, 0), (None, false)); + assert_eq!(d.tick(0, 0), (None, false)); + assert_eq!(d.tick(0, 0), (None, true), "连着几拍都空了才算做完"); + let mut d = Drained::default(); + d.tick(0, 0); + d.tick(0, 0); + assert_eq!(d.tick(1, 0), (Some(1), false), "中间又来了一个就重新数"); + assert_eq!(d.tick(0, 0), (Some(0), false)); + } + + #[test] + fn the_snapshot_wrapper_round_trips_and_tolerates_garbage() { + let g = serde_json::json!({"affinity": [1, 2, 3]}); + let s = unwrap(wrap(g.clone(), 42)).expect("拆得开"); + assert_eq!(s.gateway, g); + assert_eq!(s.next_id, 42); + assert!(unwrap(serde_json::json!("不是快照")).is_none()); + } + + #[test] + fn listeners_are_sorted_by_role_and_a_count_mismatch_is_an_error() { + let l = || std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let a: SocketAddr = "127.0.0.1:8788".parse().unwrap(); + let r: SocketAddr = "0.0.0.0:9788".parse().unwrap(); + let infos = vec![ + ListenerInfo { + role: Role::Gateway, + want: Some(a), + }, + ListenerInfo { + role: Role::Remote, + want: Some(r), + }, + ]; + let got = inherited(infos.clone(), vec![Socket::Tcp(l()), Socket::Tcp(l())]).unwrap(); + assert_eq!(got.gateway.len(), 1); + assert_eq!(got.gateway[0].0, a); + assert_eq!(got.remote.as_ref().map(|x| x.0), Some(r)); + assert!(inherited(infos, vec![Socket::Tcp(l())]).is_err()); + } + + #[test] + fn the_successor_waits_for_the_snapshot_longer_than_the_old_core_may_take_to_send_it() { + // 旧 core 发快照之前：OAuth 写回最多 5 秒、上传本机改动最多 8 秒、暂停续期最多 10 秒 + const { assert!(SNAPSHOT_WITHIN.as_secs() >= 5 + 8 + RENEWAL_SETTLE.as_secs()) }; + } + + #[test] + fn the_successor_commits_well_before_the_previous_core_gives_up() { + // 两边同一刻到点的话会各自以为是主：新 core 的时限要短一截 + const { assert!(COMMIT_WITHIN.as_secs() + 5 <= TAKEOVER_LIMIT.as_secs()) }; + } +} diff --git a/bin/twcore/src/lockfile.rs b/bin/twcore/src/lockfile.rs index ce67c03c..5bd983c8 100644 --- a/bin/twcore/src/lockfile.rs +++ b/bin/twcore/src/lockfile.rs @@ -4,12 +4,16 @@ //! 会表现为启动失败，而错误信息必须能说清是这个原因 —— 否则用户会以为 //! 自己电脑坏了。 -use std::io::Write; use std::path::{Path, PathBuf}; +/// 锁文件的名字。 +pub const LOCK_FILE: &str = "twcore.lock"; + #[derive(Debug)] pub struct LockFile { path: PathBuf, + /// 退出时删不删文件。**交接之后不删**：那时文件里写的已经是接替者的 pid + owned: bool, } #[derive(Debug)] @@ -25,7 +29,7 @@ impl LockFile { /// 尝试拿锁。发现陈旧的锁文件（进程已经没了）就接管它。 pub fn acquire(dir: &Path) -> std::io::Result { std::fs::create_dir_all(dir)?; - let path = dir.join("twcore.lock"); + let path = dir.join(LOCK_FILE); if let Ok(text) = std::fs::read_to_string(&path) && let Ok(pid) = text.trim().parse::() && pid != std::process::id() @@ -35,19 +39,154 @@ impl LockFile { } // 到这儿说明：没有锁文件、内容不是数字、或者那个 pid 已经死了。 // 后两种都是「上次没退干净」，直接覆盖。 - let mut f = std::fs::File::create(&path)?; - write!(f, "{}", std::process::id())?; - Ok(LockOutcome::Acquired(LockFile { path })) + Ok(LockOutcome::Acquired(Self::write(path)?)) + } + + /// 接替者在交接提交时成为锁的持有者：**不看现在写的是谁**，直接写上自己。 + /// + /// 此刻文件里写的是旧 core 的 pid、它也还活着（在排空），但它已经放弃了锁（见 + /// [`LockFile::disown`]）。 + pub fn claim(dir: &Path) -> std::io::Result { + std::fs::create_dir_all(dir)?; + Self::write(dir.join(LOCK_FILE)) } + + /// 放弃锁：退出时不删文件（交接提交之后，文件里马上会是接替者的 pid）。 + pub fn disown(mut self) { + self.owned = false; + } + + /// **先写临时文件再改名**：改名是原子的，读锁的那一方看到的要么是旧的、要么是新的， + /// 不会读到写了一半的（空的）—— 空的会被当成「上次没退干净」直接覆盖 + fn write(path: PathBuf) -> std::io::Result { + write_atomic(&path, std::process::id().to_string().as_bytes())?; + Ok(LockFile { path, owned: true }) + } +} + +/// 先写同目录下的临时文件，再改名盖过去。 +fn write_atomic(path: &Path, bytes: &[u8]) -> std::io::Result<()> { + let tmp = path.with_extension(format!("{}.tmp", std::process::id())); + std::fs::write(&tmp, bytes)?; + std::fs::rename(&tmp, path).inspect_err(|_| { + let _ = std::fs::remove_file(&tmp); + }) +} + +/// 锁文件里此刻写的 pid（读不到、不是数字是 `None`）。 +pub fn holder(dir: &Path) -> Option { + std::fs::read_to_string(dir.join(LOCK_FILE)) + .ok() + .and_then(|t| t.trim().parse().ok()) } impl Drop for LockFile { fn drop(&mut self) { + if !self.owned { + return; + } + let Some(dir) = self.path.parent() else { + return; + }; + // **只删自己的**：文件里已经是别人的 pid（交接提交了、接替者写上了自己），删了它 + // 下一个起来的就看不见锁 + if holder(dir) != Some(std::process::id()) { + return; + } + // **还有前任在排空就不删**（它还在写请求库）：留下这个锁（进程没了就是陈旧的）， + // 下一个起来的照常拿到锁，再从排空标记看见前任，等它退了才开库 + if !draining(dir).is_empty() { + return; + } // 尽力而为。删不掉也无所谓 —— 下次启动的僵尸检测会认出它。 let _ = std::fs::remove_file(&self.path); } } +// ─────────────────────────────────────────────────────────── 排空中的前任 + +/// 排空标记的文件名：交接之后还在排空的旧 core。 +pub const DRAINING_FILE: &str = "twcore.draining"; + +/// 标记多老就不认了（旧 core 最多排空 30 分钟，加点余量）：pid 被系统回收给了别的进程时， +/// 不至于一直等下去 +const DRAINING_STALE_MS: u64 = 35 * 60 * 1000; + +/// 一个交接之后还在排空的旧 core。 +/// +/// **为什么要记**：接替者在它排空完之前退出（关掉、崩溃、跟着父进程退）时，下一个起来的 +/// 看不见交接连接，只看得见锁。不记的话，它以为没有别人、马上开请求库 —— 和还在写库的旧 core +/// 一起写，请求号互相覆盖。 +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct Draining { + pub pid: u32, + /// 它（和它之前的）用的请求号都在这之下 + pub ids_below: u64, + /// 记下的时刻，Unix 毫秒 + pub since_ms: u64, +} + +impl Draining { + /// 此刻记下的 + pub fn new(pid: u32, ids_below: u64) -> Self { + Self { + pid, + ids_below, + since_ms: now_ms(), + } + } +} + +fn now_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} + +fn read_draining(dir: &Path) -> Vec { + std::fs::read(dir.join(DRAINING_FILE)) + .ok() + .and_then(|b| serde_json::from_slice(&b).ok()) + .unwrap_or_default() +} + +/// 此刻还在排空的前任（进程活着、标记不太老）。 +pub fn draining(dir: &Path) -> Vec { + let now = now_ms(); + read_draining(dir) + .into_iter() + .filter(|d| d.pid != std::process::id()) + .filter(|d| now.saturating_sub(d.since_ms) < DRAINING_STALE_MS && crate::proc::alive(d.pid)) + .collect() +} + +/// 记下一个刚交出去、开始排空的前任（同一个 pid 的旧记录换掉）。不在了的顺手清掉 +pub fn mark_draining(dir: &Path, d: Draining) -> std::io::Result<()> { + let mut all = draining(dir); + all.retain(|x| x.pid != d.pid); + all.push(d); + write_draining(dir, &all) +} + +/// 清掉已经不在了的。一个都不剩就删文件 +pub fn prune_draining(dir: &Path) { + let all = draining(dir); + let _ = write_draining(dir, &all); +} + +fn write_draining(dir: &Path, all: &[Draining]) -> std::io::Result<()> { + let path = dir.join(DRAINING_FILE); + if all.is_empty() { + return match std::fs::remove_file(&path) { + Err(e) if e.kind() != std::io::ErrorKind::NotFound => Err(e), + _ => Ok(()), + }; + } + let bytes = serde_json::to_vec(all).map_err(std::io::Error::other)?; + write_atomic(&path, &bytes) +} + #[cfg(test)] mod tests { use super::*; @@ -133,10 +272,101 @@ mod tests { } } + #[test] + fn a_disowned_lock_leaves_the_file_for_the_successor() { + // 交接提交之后旧 core 放弃锁：它退出时文件里写的已经是接替者，不能删 + let d = tempfile::tempdir().unwrap(); + let p = d.path().join(LOCK_FILE); + let LockOutcome::Acquired(l) = LockFile::acquire(d.path()).unwrap() else { + panic!("应该拿到锁"); + }; + l.disown(); + assert!(p.exists(), "放弃了锁，文件还该在"); + } + + #[test] + fn claiming_writes_our_pid_over_a_live_holder() { + let d = tempfile::tempdir().unwrap(); + let mut child = sleeping_child(); + std::fs::write(d.path().join("twcore.lock"), child.id().to_string()).unwrap(); + let claimed = LockFile::claim(d.path()); + let _ = child.kill(); + let _ = child.wait(); + let _l = claimed.unwrap(); + assert_eq!(holder(d.path()), Some(std::process::id())); + // 先写临时文件再改名：临时文件不留下 + let left: Vec<_> = std::fs::read_dir(d.path()) + .unwrap() + .flatten() + .map(|e| e.file_name().to_string_lossy().into_owned()) + .filter(|n| n.ends_with(".tmp")) + .collect(); + assert!(left.is_empty(), "{left:?}"); + } + + #[test] + fn a_lock_rewritten_by_someone_else_is_not_removed() { + let d = tempfile::tempdir().unwrap(); + let p = d.path().join(LOCK_FILE); + let LockOutcome::Acquired(l) = LockFile::acquire(d.path()).unwrap() else { + panic!("应该拿到锁"); + }; + // 交接提交了、接替者写上了它自己 + std::fs::write(&p, "4242").unwrap(); + drop(l); + assert_eq!(std::fs::read_to_string(&p).unwrap(), "4242", "删了别人的锁"); + } + + #[test] + fn a_lock_is_kept_while_a_predecessor_is_still_draining() { + let d = tempfile::tempdir().unwrap(); + let p = d.path().join(LOCK_FILE); + let mut child = sleeping_child(); + mark_draining(d.path(), Draining::new(child.id(), 123)).unwrap(); + let LockOutcome::Acquired(l) = LockFile::acquire(d.path()).unwrap() else { + panic!("应该拿到锁"); + }; + drop(l); + let kept = p.exists(); + let seen = draining(d.path()); + let _ = child.kill(); + let _ = child.wait(); + assert!(kept, "前任还在排空，锁被删了"); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].ids_below, 123); + // 前任没了：标记清掉，锁照常删 + prune_draining(d.path()); + assert!(!d.path().join(DRAINING_FILE).exists()); + let LockOutcome::Acquired(l) = LockFile::acquire(d.path()).unwrap() else { + panic!("应该拿到锁"); + }; + drop(l); + assert!(!p.exists()); + } + + #[test] + fn dead_or_stale_draining_marks_are_ignored() { + let d = tempfile::tempdir().unwrap(); + let dead = { + let mut c = sleeping_child(); + let pid = c.id(); + c.kill().unwrap(); + c.wait().unwrap(); + pid + }; + mark_draining(d.path(), Draining::new(dead, 1)).unwrap(); + assert!(draining(d.path()).is_empty(), "进程没了还算在排空"); + let mut old = Draining::new(std::process::id() + 1, 1); + old.pid = 1; // 一个一定活着、又不是我们的进程 + old.since_ms = 0; + mark_draining(d.path(), old).unwrap(); + assert!(draining(d.path()).is_empty(), "很久以前的标记还在算"); + } + #[test] fn dropping_the_lock_removes_the_file() { let d = tempfile::tempdir().unwrap(); - let p = d.path().join("twcore.lock"); + let p = d.path().join(LOCK_FILE); { let _l = LockFile::acquire(d.path()).unwrap(); assert!(p.exists()); diff --git a/bin/twcore/src/main.rs b/bin/twcore/src/main.rs index dc03ad21..83aaa576 100644 --- a/bin/twcore/src/main.rs +++ b/bin/twcore/src/main.rs @@ -5,17 +5,23 @@ //! 干净退出。 use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; use anyhow::{Context, Result}; use clap::{Parser, Subcommand}; +use tokio::sync::watch; +mod background; mod call; +mod handover; mod lockfile; mod pack_cli; mod proc; +mod record; mod sync_cmd; mod upgrade; +use background::Background; use lockfile::{LockFile, LockOutcome}; #[derive(Parser)] @@ -768,6 +774,15 @@ fn cmd_check(path: &Path) -> Result<()> { } } +fn already_running(dir: &Path, pid: u32) -> anyhow::Error { + anyhow::anyhow!( + "已有 twcore 实例正在运行（pid {pid}）。\n\ + 两个实例无法共用同一个网关端口，请先停止该实例。\n\ + 如果该进程实际已退出，请删除 {} 后重试", + dir.join("twcore.lock").display() + ) +} + fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> Result<()> { let dir = path .parent() @@ -777,17 +792,47 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> // 端口文件）都会顺手 `create_dir_all`，谁先到谁建 —— 而 Windows 上「谁 // 建的」决定了它的 ACL，也就决定了那份明文密钥同机的其他用户读不读得到。 tw_config::private_dir::create(&dir).with_context(|| format!("creating {}", dir.display()))?; - let _lock = match LockFile::acquire(&dir)? { - LockOutcome::Acquired(l) => l, - LockOutcome::AlreadyRunning { pid } => { - anyhow::bail!( - "已有 twcore 实例正在运行（pid {pid}）。\n\ - 两个实例无法共用同一个网关端口，请先停止该实例。\n\ - 如果该进程实际已退出，请删除 {} 后重试", - dir.join("twcore.lock").display() - ); - } + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build()?; + // **锁被一个活着的 core 占着时，先问它是不是在等接替者**（升级不断线，见 `handover`）。 + // 是的话接管它；不是的话照旧报「已有实例在运行」 + let (mut lock, mut incoming) = match LockFile::acquire(&dir)? { + LockOutcome::Acquired(l) => (Some(l), None), + LockOutcome::AlreadyRunning { pid } => match rt.block_on(handover::approach(&dir, pid)) { + handover::Approach::Proceed(inc) => { + tracing::info!(pid, "taking over from the running core"); + (None, Some(inc)) + } + handover::Approach::StartFresh => match LockFile::acquire(&dir)? { + LockOutcome::Acquired(l) => (Some(l), None), + LockOutcome::AlreadyRunning { pid } => return Err(already_running(&dir, pid)), + }, + handover::Approach::NotArmed => return Err(already_running(&dir, pid)), + }, }; + let taking_over = incoming.is_some(); + // **还在排空的前任**（见 `lockfile::Draining`）：接替者在它排空完之前退了，这个进程是照常 + // 拿到的锁，可前任还在写请求库。这时候和接替者一样：等它们都退了才开库，请求号跳过它们的 + let draining = lockfile::draining(&dir); + let mut predecessors: Vec = draining.iter().map(|d| d.pid).collect(); + if let Some(inc) = &incoming { + predecessors.push(inc.predecessor); + } + predecessors.sort_unstable(); + predecessors.dedup(); + let orphaned_ids = (!taking_over) + .then(|| draining.iter().map(|d| d.ids_below).max()) + .flatten(); + if orphaned_ids.is_some() { + tracing::warn!( + ?predecessors, + "a core that handed over is still draining; the request history opens once it exits" + ); + } + // 还没进库的记录有几条（库开之前退出时说一声） + let pending = record::Pending::default(); + let pending_outer = pending.clone(); // **配置不存在就地生成**（第 1 步）。 // @@ -821,17 +866,20 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> // // 顶上的是一份临时的空配置，只有控制面的钥匙取自原文（桌面端从同一个文件读它）。原文连 // 钥匙都找不到（YAML 坏到解析不了）就照旧退出：那时起了控制面也没人进得来 + // + // **接管时也一样**（H3）：新版本读不进这份配置，就只接管控制面、等配置修好，网关的监听 + // 留在旧 core 上按旧配置转发 let (cfg, stood_in) = match tw_config::load(path) { Ok(cfg) => (cfg, None), Err(e) => { - let standin = safe + let standin = (safe || taking_over) .then(|| std::fs::read_to_string(path).ok()) .flatten() .and_then(|text| tw_config::stand_in(&text)); match standin { Some((cfg, r)) => { tracing::error!( - "safe mode: the configuration does not load ({r}); serving the control plane with a stand-in" + "the configuration does not load ({r}); serving the control plane with a stand-in" ); (cfg, Some(r)) } @@ -853,9 +901,6 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> // 听回环，网卡出现了再补上，换了地址也跟着换，见 `tw_gateway::serve_at` let listen_at = format!("{}:{}", listen.bind, listen.port); - let rt = tokio::runtime::Builder::new_multi_thread() - .enable_all() - .build()?; // 控制面听在哪由平台决定，见 `tw_api::control::Address`。 let endpoint = tw_api::control::Address::in_dir(&dir); // **在起任何东西之前问**。等到 bind 失败时，网关已经在监听、客户端 @@ -865,12 +910,24 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> // 一份，而控制面可能在主循环开始等之前就收到那条请求。 let shutdown = tw_control::Shutdown::default(); let config_path = path.to_path_buf(); - rt.block_on(async move { + // 后台事务的运行时（见 `background`）。**在这里留一份**：进程退出时要在异步上下文之外关掉它 + let bg = Arc::new(Mutex::new(Background::default())); + let bg_outer = bg.clone(); + let r = rt.block_on(async move { let state = tw_gateway::AppState::new(cfg.clone()) .map_err(|e| anyhow::anyhow!("{}", e.message()))?; // 默认价目表：上次联网刷新存下的那份，或者内置的。自定义价目表已经 // 随配置进了价格簿 state.set_price_table(tw_control::pricing::load_table(&config_path)); + // 接替者：旧 core 的内存状态接进来，**在第一个请求之前** + if let Some(inc) = incoming.as_mut() { + // 提交之前续期归旧 core（见 `Cache::hold_renewals`） + state.oauth.hold_renewals(true); + inc.import(&state); + } + if let Some(below) = orphaned_ids { + state.bus.resume_after(below.saturating_add(handover::ID_GAP)); + } // 观测这一层。**起不来不是致命的** —— 历史记录看不见，而网关 // 照常转发。所以这里所有的失败都只记一行日志。 @@ -880,66 +937,51 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> let (body_tx, body_rx) = tw_gateway::bodies::channel(); // 插件在每个请求上的运行记录，和正文同一个道理：网关交出去，存储层落库 let (run_tx, run_rx) = tokio::sync::mpsc::channel(tw_gateway::plugin::RUN_CHANNEL_CAP); - let store = build_store( - &dir, - state.bus.clone(), - state.pricing.clone(), - body_rx, - run_rx, - tw_control::key_limits::settle_hook(&state), - ); - if let Some(rec) = &store { - state.set_body_sink(body_tx); - state.set_plugin_sink(run_tx); - // 密钥的用量上限：这一天、这一周、这个月已经用了多少，从记录里加回来。**在第一个 - // 请求之前**（网关还没开始听） - tw_control::key_limits::rebuild(&state, rec.lock().await.db()); - // 之后一期的开头变了（到了下一期、换了时区），从记录里把新的那一期加起来 - tw_control::key_limits::follow(&state, rec.clone()); - } - - /* - 日志回收。**启动先跑一次，之后每小时一次。** - - 只在启动时跑的话，一台一直开着的机器永远轮不到它；只定时跑的 - 话，关机过夜的那台永远等不到那一小时。两个都要。 - - 每一轮都重新读配置，所以改完期限不用重启 —— 下一轮就按新的算。 - 扫的是目录和一条 DELETE，空跑一次的代价可以忽略。 - - **在阻塞线程上跑，删正文时不拿记录库的锁**（见 `tw_store::task::gc`）：一次删几个 - GB 可以要好几秒，以前那期间这把锁一直攥着，记录和查询都在等。 - */ - if let Some(rec) = store.clone() { + let wiring = record::Wiring { + bus: state.bus.clone(), + pricing: state.pricing.clone(), + bodies: body_rx, + runs: run_rx, + settled: tw_control::key_limits::settle_hook(&state), + }; + let (store, store_opened) = if predecessors.is_empty() { + let store = record::open_now(&dir, wiring); + if let Some(rec) = &store { + // 密钥的用量上限：这一天、这一周、这个月已经用了多少，从记录里加回来。**在第一个 + // 请求之前**（网关还没开始听） + tw_control::key_limits::rebuild(&state, rec.lock().await.db()); + // 之后一期的开头变了（到了下一期、换了时区），从记录里把新的那一期加起来 + tw_control::key_limits::follow(&state, rec.clone()); + } + (store, watch::channel(true).1) + } else { + // 接替者（或者前任还在排空）：前任都退了才打开请求库（R3，见 `record`）。用量从 + // 记录重建在打开的那一刻；排空标记在那时清掉 let gw = state.clone(); - tokio::spawn(async move { - loop { - let r = gw.config().retention.clone(); - let freed = { - 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!( - freed_bytes = freed, - body_days = r.body_days, - row_days = r.row_days, - "reclaimed expired request logs" - ); + let marks = dir.clone(); + let (store, opened) = record::open_later( + &dir, + wiring, + predecessors.clone(), + Box::new(move |db| { + tw_control::key_limits::rebuild(&gw, db); + lockfile::prune_draining(&marks); + }), + pending, + ); + if let Some(rec) = store.clone() { + let (gw, mut opened) = (state.clone(), opened.clone()); + tokio::spawn(async move { + if opened.wait_for(|o| *o).await.is_ok() { + tw_control::key_limits::follow(&gw, rec); } - tokio::time::sleep(std::time::Duration::from_secs(3600)).await; - } - }); + }); + } + (store, opened) + }; + if store.is_some() { + state.set_body_sink(body_tx); + state.set_plugin_sink(run_tx); } // 配置的唯一入口。UI、CLI、文件监听都从这里进。 @@ -954,54 +996,79 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> // 默认插件（随 core 发的那几个）：没给过的装上（停用着），没动过的换成新版。 // 启动时在控制面起来之前走一遍，界面第一次取插件就看得到它们；之后每换入一份 // 配置再走一遍。**不挡启动**：哪个没办成只记一行、说一声。安全模式不走 —— - // 那时只有控制面，不替人往配置里写东西 - if !safe { - let seeder = tw_control::plugins::defaults::Seeder::shipped(); - seeder.seed(&manager).await; - tw_control::plugins::defaults::spawn(seeder, manager.clone()); + // 那时只有控制面，不替人往配置里写东西。**接替者提交之后才走**：之前写配置的是旧 core + let mut seeder = None; + if !safe && !taking_over { + let s = tw_control::plugins::defaults::Seeder::shipped(); + s.seed(&manager).await; + seeder = Some(s); } // **监听要留着** —— 扔掉它就停止监听，而那个失效是静默的。 // 起不来不是致命的：手改文件不会自动生效，但界面和 CLI 照常能用， // 所以说一句就继续。 - let _watch = match tw_control::spawn_watcher(manager.clone()) { - Ok(w) => Some(w), + let mut watchers: Vec> = Vec::new(); + match tw_control::spawn_watcher(manager.clone()) { + Ok(w) => watchers.push(Box::new(w)), Err(e) => { tracing::warn!("the configuration file cannot be watched, so a hand edit will not take effect on its own: {e}"); - None } - }; + } // 插件目录也盯着：**插件文件被改了，那个插件马上停用**，不等下一次改配置。 // 盯不住时退回到每次换配置时重算哈希，所以同样只说一句 - let _plugin_watch = match tw_control::plugins::spawn_watcher(state.clone(), manager.path()) - { - Ok(w) => Some(w), + match tw_control::plugins::spawn_watcher(state.clone(), manager.path()) { + Ok(w) => watchers.push(Box::new(w)), Err(e) => { tracing::warn!("the plugin directory cannot be watched, so a changed plugin file is noticed only at the next configuration change: {e}"); - None } - }; + } // 控制面无论如何都要起来 —— **网关挂了的时候，用户最需要的恰恰 // 是能改配置**。安全模式就是「只有这一半」。 let control = tw_control::ControlState { shutdown: shutdown.clone(), remote: Default::default(), + handover: Default::default(), started: std::time::Instant::now(), gateway: state.clone(), - cfg: manager, - store, + cfg: manager.clone(), + store: store.clone(), price_updater: Default::default(), chatgpt: Default::default(), zai: Default::default(), }; - // 凭据轮换要写回 config.yaml。**不是人发起的配置写入只有两种**， - // 这是一种（另一种是上面的默认插件）—— 理由是服务器换发新 refresh - // token 的那一刻旧的就作废了，不写回等于让配置文件从那一秒起就是坏的。 - tw_control::rotation::spawn(control.clone()); - // 多设备同步：启动时一轮，之后按本机改动、定时、文件夹变化、唤醒。安全模式下不自动同步 - let sync = tw_control::sync::start(&control, safe); - // 定期刷新默认价目表（`pricing.auto_update`，默认开） - tw_control::pricing::spawn(control.clone()); + // 安全模式没有数据面：没有可交接的，arm 一律拒绝 + control.handover.set_no_data_plane(safe); + // 前任还在排空、却没有交接连接（它的接替者先退了）：照「前任在排空」起来，等它退了回到平常 + if orphaned_ids.is_some() { + handover::follow_vanished(predecessors.clone(), control.handover.clone()); + } + let gateway = tw_gateway::listen::Handoff::default(); + if let Some(inc) = incoming.as_mut() { + // 接替者：控制面在旧 core 交过来的监听上服务；网关、远程控制端口提交之前不跟配置走 + control.handover.set_phase(tw_api::HandoverState::TakingOver); + if let Some(s) = inc.control.take() { + control.handover.control().inherit(s); + } + control.remote.hold(); + gateway.hold(); + } + // 后台事务（写回、同步、价目表、默认插件、模型目录、余额、日志回收）。**接替者提交 + // 之后才起**：之前做这些事的是旧 core + let mut sync = None; + if !taking_over { + let started = bg + .lock() + .map_err(|_| anyhow::anyhow!("the background lock is poisoned"))? + .start(background::Start { + control: &control, + safe, + seeder: seeder.take(), + store: store.clone(), + store_opened: store_opened.clone(), + }) + .context("starting the background tasks")?; + sync = Some(started); + } let at = endpoint.clone(); // **控制面没了就得退，不能只记一行日志。** @@ -1010,203 +1077,249 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> // 也关不掉。守护那边的心跳会因此失败，然后按重启阶梯反复拉起来， // 而用户看到的是「core 连续失败」——真正的原因（比如 socket 路径 // 太长）只在日志里躺着。退出让那句话有机会走到人眼前。 - let (control_died, control_dead) = tokio::sync::oneshot::channel(); - tokio::spawn(async move { - let r = tw_control::serve(control, &at).await; - let msg = match r { - Err(e) => format!("{e}"), - // serve 正常返回意味着 accept 循环结束了，同样是没了 - Ok(()) => "the control plane ended unexpectedly".to_string(), - }; - let _ = control_died.send(msg); - }); - - if let Some(ppid) = parent { - // 父进程守望：GUI 没了我们跟着退。轮询而不是用 kqueue，是因为 - // 这段代码要能在被重新 parent 之后仍然正确 —— macOS 上收养我们的 - // 是 launchd（getppid() 变成 1），Linux 桌面会话里多半是 - // `systemd --user`（它是 subreaper，getppid() 变成它的 pid，不是 1）。 - // 所以**不看 getppid()**，只看原来那个 pid 还在不在，哪个平台、 - // 谁来收养都一样。 + let (control_died, mut control_dead) = tokio::sync::oneshot::channel(); + { + let control = control.clone(); tokio::spawn(async move { - loop { - tokio::time::sleep(std::time::Duration::from_secs(2)).await; - if !proc::alive(ppid) { - tracing::info!(ppid, "the parent exited; exiting with it"); - std::process::exit(0); - } - } + let r = tw_control::serve(control, &at).await; + let msg = match r { + Err(e) => format!("{e}"), + // serve 正常返回意味着 accept 循环结束了，同样是没了 + Ok(()) => "the control plane ended unexpectedly".to_string(), + }; + let _ = control_died.send(msg); }); } + // 父进程守望：GUI 没了我们跟着退。轮询而不是用 kqueue，是因为 + // 这段代码要能在被重新 parent 之后仍然正确 —— macOS 上收养我们的 + // 是 launchd（getppid() 变成 1），Linux 桌面会话里多半是 + // `systemd --user`（它是 subreaper，getppid() 变成它的 pid，不是 1）。 + // 所以**不看 getppid()**，只看原来那个 pid 还在不在，哪个平台、 + // 谁来收养都一样。 + // + // **arm 之后关掉**（见 `handover`）：桌面端装好更新要重启，core 留下来等接替者 + let parent_watch = parent.map(handover::watch_parent); + if safe { // 安全模式：只起控制面。数据面不动，让用户还能改配置、回滚。 // **这时候控制面就是全部** —— 它没了，这个进程一件事都干不了。 + // + // 安全模式的接替者只接管控制面，**一直握着交接连接**：松手的话旧 core 回到待交接， + // 又从控制面上接新连接了 tracing::warn!("safe mode: starting the control plane only, not the data plane"); + let held = async { + match incoming { + Some(inc) => inc.hold().await, + None => std::future::pending().await, + } + }; tokio::select! { msg = control_dead => { anyhow::bail!("{}", msg.unwrap_or_else(|_| "the control plane stopped".into())) } _ = shutdown_requested(&shutdown) => {} + _ = held => {} } return Ok(()); } - // 模型目录：后台去问每个上游有哪些模型。**不挡启动** —— - // 探测要打网络，而网关不该因为一次探测慢而起不来。 - tw_gateway::models::spawn(state.clone()); - // 余额：启动时读一遍，之后按节奏读（见 `tw_gateway::balances`）。**不挡启动，不挡转发** - tw_gateway::balances::spawn(state.clone()); - - tracing::info!(listen = %listen_at, "starting"); - tokio::select! { - r = tw_gateway::serve_at(state, listen, !overridden) => { - r.with_context(|| format!("{listen_at} could not be listened on. If the port is taken, check whether an earlier instance has fully exited")) + let follow = !overridden; + let mut serve = match incoming { + None => { + tracing::info!(listen = %listen_at, "starting"); + tokio::spawn(tw_gateway::listen::serve_with( + state.clone(), + listen, + follow, + gateway.clone(), + Vec::new(), + )) } - msg = control_dead => { - anyhow::bail!("{}", msg.unwrap_or_else(|_| "the control plane stopped".into())) - } - _ = shutdown_requested(&shutdown) => { - tracing::info!("got a shutdown signal"); - // 还没上传的本机改动，退出前上传（等不了太久：没传成的下次启动时传） - sync.flush().await; - Ok(()) + Some(mut inc) => { + // H3：配置读不进来就只接管控制面，等配置修好（换入成功）再接着走 + if stood_in.is_some() { + tracing::warn!("taking over the control plane only until the configuration loads; the previous core keeps forwarding"); + tokio::select! { + _ = config_fixed(&manager) => { + tracing::info!("the configuration loads now; taking over the gateway"); + } + msg = &mut control_dead => { + anyhow::bail!("{}", msg.unwrap_or_else(|_| "the control plane stopped".into())) + } + _ = shutdown_requested(&shutdown) => return Ok(()), + } + } + let inherited = tokio::select! { + r = inc.listeners(&state) => r?, + msg = &mut control_dead => { + anyhow::bail!("{}", msg.unwrap_or_else(|_| "the control plane stopped".into())) + } + _ = shutdown_requested(&shutdown) => return Ok(()), + }; + // 配置换入过（H3 修好了），按此刻的配置跟 + let listen = if follow { + state.config().listen.gateway.clone() + } else { + listen + }; + let serve = tokio::spawn(tw_gateway::listen::serve_with( + state.clone(), + listen, + follow, + gateway.clone(), + inherited.gateway, + )); + if let Some((want, l)) = inherited.remote { + control.remote.adopt(want, l); + } + // 已经在全部监听上 accept 了才提交（不可逆，H5） + tokio::time::timeout(handover::commit_within(), gateway.listening()) + .await + .context("the listeners handed over could not be served in time")?; + // 旧 core 回了 CommitAck（或者它已经没了）才算提交；它还在却没回，这边退出， + // 守护重来（这时它仍握着全部监听） + inc.commit().await?; + // ── 从这里起不再回退（H5）。旧 core 已经停止 accept：**下面任何一步出错都只记 + // 一行、照常服务**，退出的话就没人在听了 + // + // 从此续期由这一边做（旧 core 交出第二份快照起已暂停） + state.oauth.hold_renewals(false); + // 成为锁的持有者：写上自己的 pid（旧 core 已经放弃了锁），记下它在排空 + match LockFile::claim(&dir) { + Ok(l) => lock = Some(l), + Err(e) => tracing::warn!("the lock file could not be taken over: {e}"), + } + if let Err(e) = lockfile::mark_draining( + &dir, + lockfile::Draining::new(inc.predecessor, inc.ids_from), + ) { + tracing::warn!("the draining core could not be noted down: {e}"); + } + gateway.follow(); + control.remote.follow(); + control.handover.set_phase(tw_api::HandoverState::Draining); + tracing::info!("took over; the previous core is draining"); + let s = tw_control::plugins::defaults::Seeder::shipped(); + s.seed(&manager).await; + let started = match bg.lock() { + Ok(mut g) => g + .start(background::Start { + control: &control, + safe: false, + seeder: Some(s), + store: store.clone(), + store_opened: store_opened.clone(), + }) + .map_err(|e| e.to_string()), + Err(_) => Err("the background lock is poisoned".to_string()), + }; + match started { + Ok(started) => sync = Some(started), + Err(e) => tracing::error!( + "the background tasks could not start, so token write-back, sync and price updates are off this run; forwarding is unaffected: {e}" + ), + } + inc.follow_drain(state.clone(), control.handover.clone()); + serve } - } - }) -} - -/// 建观测层。**每一步失败都只是「没有历史记录」，不是「起不来」。** -/// -/// 这条边界值得写死在代码形状里：这个函数返回 `Option`，而不是 -/// `Result` —— 调用方连处理错误的机会都不该有，因为没有任何一种 -/// 处理方式是「不转发了」。 -fn build_store( - dir: &Path, - // **收整条总线，不只是一个订阅端。**存储层算完价钱要往回报一条 - // （见 `Event::RequestPriced`）—— 它是这条链上唯一知道单价的地方。 - bus: tw_observe::EventBus, - // **和网关同一份价格簿**，不是一份副本：改了价目表，下一个结束的请求就按新价算 - pricing: tw_pricing::Shared, - bodies: tokio::sync::mpsc::Receiver, - runs: tokio::sync::mpsc::Receiver, - // 每记下一行请求，交给网关的密钥用量上限结算（见 `tw_control::key_limits`） - settled: tw_store::SettleHook, -) -> Option>> { - // **存储层自己的那条通道，不订阅界面的广播**（见 `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) => { - tracing::warn!( - "request recording could not start, so nothing is recorded this run; forwarding is unaffected: {e}" - ); - return None; - } - }; - /* - **接着上次的号往下发。**这里是唯一同时看得见总线和库的地方， - 而且它跑在 `serve` 之前（见 `cmd_serve`）—— 第一个请求拿到号 - 的时候，计数器已经归位了。 + }; - 读不到就照常跑：那顶多是这一次重启又压掉几条历史，比起因为 - 一次查询失败而不记录任何东西要好得多。 - */ - match db.last_request_id() { - Ok(0) => {} - Ok(last) => { - tracing::debug!(last, "request ids resume from where they left off"); - bus.resume_after(last); - } - Err(e) => tracing::warn!( - "could not read the last request id, so this run may overwrite the oldest records: {e}" - ), - } - /* - 两边的 body 结构在这里对接，**落盘的不是原文**：脱敏规则认得出的值在这里换掉、 - 打码（`BodyRecord::for_disk`，见 `tw_gateway::bodies`）。放在阻塞线程上 —— - 一份 4 MB 的正文要扫好几遍，占着异步线程的话，同一个线程上的转发都得等它。 + // 旧 core 那一半：接 arm、等接替者、交接、排空（见 `handover`） + let core = Arc::new(handover::Core { + dir: dir.clone(), + state: state.clone(), + control: control.clone(), + gateway: gateway.clone(), + background: bg.clone(), + sync: Mutex::new(sync.clone()), + watchers: Mutex::new(watchers), + lock: Mutex::new(lock), + parent: Mutex::new(parent_watch), + parent_pid: parent, + store_opened, + handed_over: Default::default(), + }); + let mut exit = handover::serve_arms(core.clone()); - 交给存储层的这一头**只留一个空位**。等着落盘的正文由网关那一头按字节记账 - （`bodies::QUEUED_MAX`），一份正文的额度要等存储层收下它才还回去；这里再开一个 - 大口子的话，攒在这里的那些就没人管了。 - */ - let (tx, rx) = tokio::sync::mpsc::channel(1); - let mut bodies = bodies; - tokio::spawn(async move { - while let Some(b) = bodies.recv().await { - let Ok(disk) = tokio::task::spawn_blocking(move || b.for_disk()).await else { - continue; - }; - let tw_gateway::bodies::ForDisk { - id, - at_ms, - kind, - body, - original_len, - held, - } = disk; - let mapped = tw_store::StoredBody { - id, - at_ms, - which: match kind { - tw_gateway::bodies::BodyKind::Request => tw_store::Which::Request, - tw_gateway::bodies::BodyKind::Response => tw_store::Which::Response, - tw_gateway::bodies::BodyKind::AfterPlugins => tw_store::Which::AfterPlugins, - tw_gateway::bodies::BodyKind::UpstreamRequest => { - tw_store::Which::UpstreamRequest + let mut serving = true; + let mut armed_task = true; + loop { + tokio::select! { + r = &mut serve, if serving => match r { + // 交接提交之后停止了 accept：在跑的请求还在，等交接那一边说退 + Ok(Ok(())) => serving = false, + Ok(Err(e)) => { + return Err(e).with_context(|| format!("{listen_at} could not be listened on. If the port is taken, check whether an earlier instance has fully exited")); + } + Err(e) => anyhow::bail!("the gateway stopped: {e}"), + }, + msg = &mut control_dead => { + anyhow::bail!("{}", msg.unwrap_or_else(|_| "the control plane stopped".into())) + } + // 系统的退出信号、`/shutdown`：照旧立即退出（F9；`/shutdown` 的语义不变） + _ = shutdown_requested(&shutdown) => { + tracing::info!("got a shutdown signal"); + // 还没上传的本机改动，退出前上传（等不了太久：没传成的下次启动时传）。 + // **交出去了就不传**：同步已经归接替者 + if !core.handed_over.load(std::sync::atomic::Ordering::SeqCst) + && let Some(sync) = &sync + { + sync.flush().await; + } + return Ok(()); + } + e = &mut exit, if armed_task => match e { + Ok(e) => { + if e == handover::Exit::Retired + && let Some(sync) = &sync + { + sync.flush().await; + } + return Ok(()); + } + // 交接那一路自己没了（panic）：不是退出的理由，照常服务 + Err(_) => { + tracing::error!("the handover task ended unexpectedly; serving on without handovers"); + armed_task = false; } - tw_gateway::bodies::BodyKind::ClientResponse => tw_store::Which::ClientResponse, - tw_gateway::bodies::BodyKind::Heads => tw_store::Which::Heads, }, - body, - original_len, - }; - if tx.send(mapped).await.is_err() { - return; } - // 存储层收下了：额度还回去 - drop(held); } }); - let recorder = tw_store::task::spawn( - // 算完价钱往回报一条 —— 见 `Event::RequestPriced`。这里是唯一 - // 同时看得见总线和存储层的地方，所以接线在这儿完成。 - tw_store::Recorder::new(db, blobs, pricing) - .reporting_to(bus) - .settling_to(settled), - events, - rx, - ); - // 插件的运行记录同样在这里对接：网关那边的一次运行，换成存储层的一行 - let (run_tx, run_rx) = tokio::sync::mpsc::channel(tw_gateway::plugin::RUN_CHANNEL_CAP); - let mut runs = runs; - tokio::spawn(async move { - while let Some(r) = runs.recv().await { - let row = tw_store::PluginRunRow { - request_id: r.request_id as i64, - at_ms: r.at_ms as i64, - plugin_id: r.run.plugin_id, - plugin_name: r.run.plugin_name, - hook: r.run.hook, - outcome: r.run.outcome, - error: r.run.error, - cpu_us: r.run.cpu_us.min(i64::MAX as u64) as i64, - detail: r.run.detail.map(|d| d.to_string()), - }; - if run_tx.send(row).await.is_err() { - return; + // 前任还在写库、这边的库还没开就退了：排着的记录写不进去，**说一声丢了多少** + let left = pending_outer.get(); + if left > 0 { + tracing::warn!( + events = left, + "exiting before the previous core finished draining; {left} queued request events are not recorded" + ); + } + // 后台事务的运行时在异步上下文之外关（在里面丢掉一个运行时会 panic） + let leftover = bg_outer.lock().ok().and_then(|mut g| g.take_runtime()); + if let Some(bg_rt) = leftover { + bg_rt.shutdown_timeout(std::time::Duration::from_secs(1)); + } + r +} + +/// 等配置读得进来了（H3）：换入了一份、不再是被拒的那一份。 +async fn config_fixed(manager: &tw_control::ConfigManager) { + let mut events = manager.bus().subscribe(); + loop { + match events.recv().await { + Ok(tw_api::Event::ConfigReloaded { .. }) + | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { + if manager.rejected().is_none() { + return; + } + } + Ok(_) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + return std::future::pending().await; } } - }); - tw_store::task::record_plugin_runs(recorder.clone(), run_rx); - Some(recorder) + } } /// 等一个「该退了」。 diff --git a/bin/twcore/src/record.rs b/bin/twcore/src/record.rs new file mode 100644 index 00000000..0947e5d5 --- /dev/null +++ b/bin/twcore/src/record.rs @@ -0,0 +1,485 @@ +//! 请求记录：观测层怎么接上网关。 +//! +//! **每一步失败都只是「没有历史记录」，不是「起不来」。**这条边界写死在代码形状里： +//! 这里交出去的是 `Option`，而不是 `Result` —— 调用方连处理错误的机会都不该有，因为 +//! 没有任何一种处理方式是「不转发了」。 +//! +//! # 接替者晚一点打开（升级不断线） +//! +//! 交接之后旧 core 还在排空，**它还在写自己的记录**。新 core 这时候打开请求库，两个进程 +//! 一起写；新版本提高了 SCHEMA 的话，新的一打开就把旧库重建了，旧 core 写的全丢。所以 +//! 接替者（[`open_later`]）先给控制面一个占位的（内存库，查历史是空的），等旧 core 的 +//! 进程没了才打开真的库；之间的记录在内存里排队（[`Queue`]），满了丢最早的、记个数。 +//! +//! 正文不排队：存储层不收，网关那一头按字节记账的队列（`bodies::QUEUED_MAX`）满了自己丢， +//! 一样不挡转发。 + +use std::collections::VecDeque; +use std::path::Path; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use tokio::sync::watch; + +pub type Store = Arc>; + +/// 观测层接上网关的那几头。 +pub struct Wiring { + /// **收整条总线，不只是一个订阅端。**存储层算完价钱要往回报一条 + /// （见 `Event::RequestPriced`）—— 它是这条链上唯一知道单价的地方。 + pub bus: tw_observe::EventBus, + /// **和网关同一份价格簿**，不是一份副本：改了价目表，下一个结束的请求就按新价算 + pub pricing: tw_pricing::Shared, + pub bodies: tokio::sync::mpsc::Receiver, + pub runs: tokio::sync::mpsc::Receiver, + /// 每记下一行请求，交给网关的密钥用量上限结算（见 `tw_control::key_limits`） + pub settled: tw_store::SettleHook, +} + +/// 现在就打开（平常的启动）。 +pub fn open_now(dir: &Path, w: Wiring) -> Option { + // **存储层自己的那条通道，不订阅界面的广播**（见 `EventBus::record_feed`）：广播上 + // 掉队丢的是界面上的几个点，记录丢的是那几个请求 + let Some(events) = w.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) => { + tracing::warn!( + "request recording could not start, so nothing is recorded this run; forwarding is unaffected: {e}" + ); + return None; + } + }; + resume_ids(&w.bus, &db); + let tx = bodies_to_disk(w.bodies); + let recorder = tw_store::task::spawn( + // 算完价钱往回报一条 —— 见 `Event::RequestPriced`。这里是唯一 + // 同时看得见总线和存储层的地方，所以接线在这儿完成。 + tw_store::Recorder::new(db, blobs, w.pricing) + .reporting_to(w.bus) + .settling_to(w.settled), + events, + tx, + ); + tw_store::task::record_plugin_runs(recorder.clone(), plugin_runs(w.runs)); + Some(recorder) +} + +/// 接替者：旧 core（`predecessors`，交接链上可能不止一个）的进程都没了才打开请求库，之间的 +/// 记录排队。 +/// +/// 交出去的是占位的那一份（内存库）；打开之后**原地换成真的**，控制面拿着的还是同一个 +/// `Arc`。返回的 `watch` 在打开之后（或者打开失败、不再等之后）变成 `true`。 +/// +/// # 排队期间的结算 +/// +/// 用量上限要靠「记下一行就结算」（`settled`）把预留换成实数；不结算的预留过一会儿会被当成 +/// 走丢的放掉，这一期的用量里就少了这个进程自己的请求。所以**占位的那一份照样结算**：事件 +/// 先交给它（它记进内存库、按同一套价目表算钱、调同一个结算），再排进队列。 +/// +/// 真的库打开时**不重复算**：排着的事件交给一份**不结算、不报价钱**的新记录器写进库，写完 +/// 再从库里重建用量（`on_open`，库里这时有旧 core 的和这边排着的全部），最后才给它接上结算 +/// 和总线。重建是换掉而不是加上，排队期间结算进内存的那些被库里的同一批行换掉，一笔不多。 +pub fn open_later( + dir: &Path, + w: Wiring, + predecessors: Vec, + on_open: Box, + pending: Pending, +) -> (Option, watch::Receiver) { + let (opened_tx, opened) = watch::channel(false); + let Some(mut events) = w.bus.record_feed() else { + tracing::warn!( + "request recording is already attached to the event bus; not starting a second one" + ); + opened_tx.send_replace(true); + return (None, opened); + }; + let placeholder = match tw_store::Db::in_memory() { + Ok(db) => tw_store::Recorder::new( + db, + tw_store::Blobs::new(dir.join("blobs")), + w.pricing.clone(), + ) + .reporting_to(w.bus.clone()) + .settling_to(w.settled.clone()), + Err(e) => { + tracing::warn!( + "request recording could not start, so nothing is recorded this run; forwarding is unaffected: {e}" + ); + opened_tx.send_replace(true); + return (None, opened); + } + }; + let store: Store = Arc::new(tokio::sync::Mutex::new(placeholder)); + // 正文这一头先不收（见模块文档）；插件的运行记录也等库开了再接 + let body_rx = bodies_to_disk(w.bodies); + let run_rx = plugin_runs(w.runs); + let shared = store.clone(); + let dir = dir.to_path_buf(); + let spawned = std::thread::Builder::new() + .name("tw-recorder".into()) + .spawn(move || { + let mut queue = Queue::default(); + wait_for_exit(&predecessors, &mut events, &mut queue, &shared, &pending); + let (db, blobs) = match tw_store::open(&dir) { + Ok(x) => x, + Err(e) => { + tracing::warn!( + "request recording could not start, so nothing is recorded this run; forwarding is unaffected: {e}" + ); + pending.set(0); + opened_tx.send_replace(true); + // 通道要有人收，否则总线那一头一直记「丢了」。占位的那一份照样结算 + record_loop(&shared, &mut events); + return; + } + }; + resume_ids(&w.bus, &db); + let blobs_root = blobs.root().to_path_buf(); + let (n, dropped) = (queue.len(), queue.dropped); + if dropped > 0 { + tracing::warn!( + dropped, + "the request records queued while the previous core drained overflowed; the oldest {dropped} events were not recorded" + ); + } + // 排着的写进库：**不结算、不报价钱**（占位的那一份已经做过了） + let mut real = tw_store::Recorder::new(db, blobs, w.pricing); + for ev in queue.drain() { + real.on_event(&ev); + } + // 用量以库为准：旧 core 的、这边排着的都在里面了 + on_open(real.db()); + let real = real.reporting_to(w.bus).settling_to(w.settled); + *shared.blocking_lock() = real; + pending.set(0); + spawn_body_writer(blobs_root, body_rx); + tw_store::task::record_plugin_runs(shared.clone(), run_rx); + tracing::info!(queued = n, "the previous core has exited; request recording is open"); + opened_tx.send_replace(true); + record_loop(&shared, &mut events); + }); + if let Err(e) = spawned { + tracing::warn!("the recorder thread could not start, so it records nothing this run: {e}"); + return (None, opened); + } + (Some(store), opened) +} + +/// 还排着、没进库的记录有几条（事件数）。**进程在库打开之前退出时**照它说一声丢了多少： +/// 前任还在写库，这时候不能开库 +#[derive(Clone, Default)] +pub struct Pending(Arc); + +impl Pending { + pub fn get(&self) -> usize { + self.0.load(std::sync::atomic::Ordering::Relaxed) + } + + fn set(&self, n: usize) { + self.0.store(n, std::sync::atomic::Ordering::Relaxed); + } +} + +/// 旧 core 最多排空这么久（`handover::DRAIN_LIMIT` 加一点余量）。到了还没退，照样打开 +const PREDECESSOR_WAIT: Duration = Duration::from_secs(35 * 60); + +/// 等前任的进程都没了，之间把记录交给占位的那一份（结算）、收进队列。 +fn wait_for_exit( + pids: &[u32], + events: &mut tw_observe::RecordFeed, + queue: &mut Queue, + placeholder: &Store, + pending: &Pending, +) { + let started = Instant::now(); + let mut last_check = Instant::now() - Duration::from_secs(1); + loop { + let mut got = Vec::new(); + while let Some(ev) = events.try_recv() { + got.push(ev); + } + if !got.is_empty() { + let mut g = placeholder.blocking_lock(); + for ev in &got { + g.on_event(ev); + } + } + for ev in got { + queue.push(ev); + } + queue.dropped += events.take_dropped(); + pending.set(queue.len()); + if last_check.elapsed() >= Duration::from_millis(500) { + last_check = Instant::now(); + if pids.iter().all(|p| !crate::proc::alive(*p)) { + return; + } + if started.elapsed() >= PREDECESSOR_WAIT { + tracing::warn!( + ?pids, + "the previous core is still running long past its drain limit; opening the request history anyway" + ); + return; + } + } + std::thread::sleep(Duration::from_millis(50)); + } +} + +/// 记录线程拿着锁一次最多吃几条事件（同 `tw_store::task`）。 +const BATCH: usize = 64; + +/// 平常的记录循环，和 `tw_store::task::spawn` 里那条一样。 +fn record_loop(shared: &Store, events: &mut tw_observe::RecordFeed) { + while let Some(ev) = events.blocking_recv() { + { + let mut g = shared.blocking_lock(); + g.on_event(&ev); + for _ in 1..BATCH { + let Some(ev) = events.try_recv() else { break }; + g.on_event(&ev); + } + } + // **落后了就丢，但要说出来。**通道满了丢的是观测记录，而等着它会拖慢转发 + let dropped = events.take_dropped(); + if dropped > 0 { + tracing::warn!( + dropped, + "recording fell behind; {dropped} events were not recorded" + ); + } + } +} + +/// 一个请求大约几条事件（开始、响应头、首字、路由、结局、价钱）。排队的上限按请求数说， +/// 数的是事件 +const EVENTS_PER_RECORD: usize = 6; +/// 排队最多多少条记录（方案 9.4） +pub const QUEUE_MAX_RECORDS: usize = 10_000; +/// 排队最多多少字节（方案 9.4） +pub const QUEUE_MAX_BYTES: usize = 256 * 1024 * 1024; + +/// 接替者等旧 core 退出时排着的记录。**满了丢最早的**，记个数。 +pub struct Queue { + items: VecDeque<(tw_api::Event, usize)>, + bytes: usize, + max_events: usize, + max_bytes: usize, + pub dropped: u64, +} + +impl Default for Queue { + fn default() -> Self { + Self::with_limits(QUEUE_MAX_RECORDS * EVENTS_PER_RECORD, QUEUE_MAX_BYTES) + } +} + +impl Queue { + pub fn with_limits(max_events: usize, max_bytes: usize) -> Self { + Self { + items: VecDeque::new(), + bytes: 0, + max_events, + max_bytes, + dropped: 0, + } + } + + pub fn push(&mut self, ev: tw_api::Event) { + // 按序列化之后的长度记：事件里没有正文，最长的是带错误原因、内容命中的那几条 + let size = serde_json::to_vec(&ev).map(|v| v.len()).unwrap_or(256); + while !self.items.is_empty() + && (self.items.len() >= self.max_events || self.bytes + size > self.max_bytes) + { + if let Some((_, n)) = self.items.pop_front() { + self.bytes -= n; + self.dropped += 1; + } + } + if size > self.max_bytes { + self.dropped += 1; + return; + } + self.bytes += size; + self.items.push_back((ev, size)); + } + + pub fn len(&self) -> usize { + self.items.len() + } + + pub fn drain(&mut self) -> impl Iterator + '_ { + self.bytes = 0; + self.items.drain(..).map(|(e, _)| e) + } +} + +/// **接着上次的号往下发。**这里是唯一同时看得见总线和库的地方。读不到就照常跑：那顶多 +/// 是这一次重启又压掉几条历史，比起因为一次查询失败而不记录任何东西要好得多。 +fn resume_ids(bus: &tw_observe::EventBus, db: &tw_store::Db) { + match db.last_request_id() { + Ok(0) => {} + Ok(last) => { + tracing::debug!(last, "request ids resume from where they left off"); + bus.resume_after(last); + } + Err(e) => tracing::warn!( + "could not read the last request id, so this run may overwrite the oldest records: {e}" + ), + } +} + +/* + 两边的 body 结构在这里对接，**落盘的不是原文**：脱敏规则认得出的值在这里换掉、 + 打码（`BodyRecord::for_disk`，见 `tw_gateway::bodies`）。放在阻塞线程上 —— + 一份 4 MB 的正文要扫好几遍，占着异步线程的话，同一个线程上的转发都得等它。 + + 交给存储层的这一头**只留一个空位**。等着落盘的正文由网关那一头按字节记账 + （`bodies::QUEUED_MAX`），一份正文的额度要等存储层收下它才还回去；这里再开一个 + 大口子的话，攒在这里的那些就没人管了。 +*/ +fn bodies_to_disk( + mut bodies: tokio::sync::mpsc::Receiver, +) -> tokio::sync::mpsc::Receiver { + let (tx, rx) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + while let Some(b) = bodies.recv().await { + let Ok(disk) = tokio::task::spawn_blocking(move || b.for_disk()).await else { + continue; + }; + let tw_gateway::bodies::ForDisk { + id, + at_ms, + kind, + body, + original_len, + held, + } = disk; + let mapped = tw_store::StoredBody { + id, + at_ms, + which: match kind { + tw_gateway::bodies::BodyKind::Request => tw_store::Which::Request, + tw_gateway::bodies::BodyKind::Response => tw_store::Which::Response, + tw_gateway::bodies::BodyKind::AfterPlugins => tw_store::Which::AfterPlugins, + tw_gateway::bodies::BodyKind::UpstreamRequest => { + tw_store::Which::UpstreamRequest + } + tw_gateway::bodies::BodyKind::ClientResponse => tw_store::Which::ClientResponse, + tw_gateway::bodies::BodyKind::Heads => tw_store::Which::Heads, + }, + body, + original_len, + }; + if tx.send(mapped).await.is_err() { + return; + } + // 存储层收下了：额度还回去 + drop(held); + } + }); + rx +} + +/// 接替者的正文写盘线程（平常由 `tw_store::task::spawn` 起）。 +fn spawn_body_writer( + root: std::path::PathBuf, + mut rx: tokio::sync::mpsc::Receiver, +) { + let blobs = tw_store::Blobs::new(root); + let spawned = std::thread::Builder::new() + .name("tw-bodies".into()) + .spawn(move || { + while let Some(b) = rx.blocking_recv() { + blobs.put_with_len(b.at_ms, b.id as i64, b.which, &b.body, b.original_len); + } + }); + if let Err(e) = spawned { + tracing::warn!("the tw-bodies thread could not start, so it records nothing this run: {e}"); + } +} + +/// 插件的运行记录同样在这里对接：网关那边的一次运行，换成存储层的一行 +fn plugin_runs( + mut runs: tokio::sync::mpsc::Receiver, +) -> tokio::sync::mpsc::Receiver { + let (run_tx, run_rx) = tokio::sync::mpsc::channel(tw_gateway::plugin::RUN_CHANNEL_CAP); + tokio::spawn(async move { + while let Some(r) = runs.recv().await { + let row = tw_store::PluginRunRow { + request_id: r.request_id as i64, + at_ms: r.at_ms as i64, + plugin_id: r.run.plugin_id, + plugin_name: r.run.plugin_name, + hook: r.run.hook, + outcome: r.run.outcome, + error: r.run.error, + cpu_us: r.run.cpu_us.min(i64::MAX as u64) as i64, + detail: r.run.detail.map(|d| d.to_string()), + }; + if run_tx.send(row).await.is_err() { + return; + } + } + }); + run_rx +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ev(id: u64) -> tw_api::Event { + tw_api::Event::ConfigReloaded { + id, + version: "v".into(), + origin: tw_api::ConfigOrigin::External, + at_ms: 0, + } + } + + fn ids(q: &mut Queue) -> Vec { + q.drain() + .map(|e| match e { + tw_api::Event::ConfigReloaded { id, .. } => id, + _ => unreachable!(), + }) + .collect() + } + + #[test] + fn a_full_queue_drops_the_oldest_and_counts_them() { + let mut q = Queue::with_limits(3, usize::MAX); + for i in 1..=5 { + q.push(ev(i)); + } + assert_eq!(q.dropped, 2); + assert_eq!(ids(&mut q), [3, 4, 5], "留下的是最新的"); + } + + #[test] + fn the_byte_limit_also_drops_the_oldest() { + let one = serde_json::to_vec(&ev(1)).unwrap().len(); + let mut q = Queue::with_limits(usize::MAX, one * 2 + one / 2); + for i in 1..=4 { + q.push(ev(i)); + } + assert_eq!(q.len(), 2); + assert_eq!(q.dropped, 2); + assert_eq!(ids(&mut q), [3, 4]); + } + + #[test] + fn the_default_limits_are_the_documented_ones() { + let q = Queue::default(); + assert_eq!(q.max_events, QUEUE_MAX_RECORDS * EVENTS_PER_RECORD); + assert_eq!(q.max_bytes, 256 * 1024 * 1024); + } +} diff --git a/bin/twcore/tests/handover.rs b/bin/twcore/tests/handover.rs new file mode 100644 index 00000000..9b50ac57 --- /dev/null +++ b/bin/twcore/tests/handover.rs @@ -0,0 +1,971 @@ +//! 升级不断线，端到端：起真的 twcore 进程，在压测中交接。 +//! +//! - 交接前开始的长流在旧 core 上完整结束；交接期间持续发的请求**没有一个被拒绝或重置** +//! （方案 V1）；旧 core 排空后退出，锁归新 core。 +//! - 接替者在提交之前走掉：旧 core 回到待交接、控制面接着接、网关照常（F4）。 +//! - 版本号不同：旧 core 拒绝，接替者请它退回，它排空后退出（H2）。 +//! - 等不到接替者：到点排空退出，不留孤儿进程（H7）。 +//! +//! 这几条要跑真的交接连接和网关钩子，整合之前标着 ignore。 + +use std::net::SocketAddr; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command, Stdio}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use axum::Router; +use axum::routing::post; + +const BIN: &str = env!("CARGO_BIN_EXE_twcore"); + +/// 数据目录。**路径要短**：控制面的 unix socket 路径有 104 字节的上限，macOS 的临时目录 +/// 本身就占了一大半 +fn short_dir() -> tempfile::TempDir { + #[cfg(unix)] + { + tempfile::Builder::new() + .prefix("twho") + .tempdir_in("/tmp") + .unwrap() + } + #[cfg(not(unix))] + { + tempfile::tempdir().unwrap() + } +} + +fn spare_port() -> u16 { + std::net::TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() + .port() +} + +/// 假的 Anthropic 上游：请求体里有 `slow` 就慢慢流（每 200 毫秒一段，`chunks` 段）， +/// 否则立刻答一条。 +async fn upstream(chunks: usize) -> SocketAddr { + let app = Router::new().route( + "/v1/messages", + post(move |body: String| async move { + if !body.contains("slow") { + return axum::response::Response::builder() + .header("content-type", "application/json") + .body(axum::body::Body::from( + r#"{"id":"msg_1","type":"message","role":"assistant","model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#, + )) + .unwrap(); + } + let stream = async_stream(chunks); + axum::response::Response::builder() + .header("content-type", "text/event-stream") + .body(axum::body::Body::from_stream(stream)) + .unwrap() + }), + ); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let a = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + a +} + +fn async_stream( + chunks: usize, +) -> impl futures::Stream> { + let (tx, rx) = tokio::sync::mpsc::channel::>(4); + tokio::spawn(async move { + let ev = |name: &str, data: &str| format!("event: {name}\ndata: {data}\n\n"); + let _ = tx + .send(Ok(ev( + "message_start", + r#"{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"m","content":[],"stop_reason":null,"usage":{"input_tokens":1,"output_tokens":0}}}"#, + ))) + .await; + let _ = tx + .send(Ok(ev( + "content_block_start", + r#"{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#, + ))) + .await; + for i in 0..chunks { + tokio::time::sleep(Duration::from_millis(200)).await; + let d = format!( + r#"{{"type":"content_block_delta","index":0,"delta":{{"type":"text_delta","text":"part{i} "}}}}"# + ); + if tx.send(Ok(ev("content_block_delta", &d))).await.is_err() { + return; + } + } + let _ = tx + .send(Ok(ev( + "content_block_stop", + r#"{"type":"content_block_stop","index":0}"#, + ))) + .await; + let _ = tx + .send(Ok(ev( + "message_delta", + r#"{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":5}}"#, + ))) + .await; + let _ = tx + .send(Ok(ev("message_stop", r#"{"type":"message_stop"}"#))) + .await; + }); + tokio_stream_from(rx) +} + +/// mpsc 接收端当成 Stream（不为这一处引 tokio-stream） +fn tokio_stream_from( + rx: tokio::sync::mpsc::Receiver, +) -> impl futures::Stream { + futures::stream::unfold(rx, |mut rx| async move { rx.recv().await.map(|x| (x, rx)) }) +} + +struct Bed { + dir: tempfile::TempDir, + gateway: SocketAddr, +} + +impl Bed { + async fn new(chunks: usize) -> Bed { + Self::with_client(chunks, "").await + } + + /// 密钥 `t` 带一条按天数请求的上限（看用量从哪儿来、算没算重） + async fn with_daily_limit(chunks: usize) -> Bed { + Self::with_client( + chunks, + " limits:\n - { per: day, requests: 100000 }\n", + ) + .await + } + + async fn with_client(chunks: usize, client_extra: &str) -> Bed { + let up = upstream(chunks).await; + let dir = short_dir(); + let port = spare_port(); + let cfg = format!( + "version: 1\nlisten:\n gateway:\n port: {port}\nclients:\n - name: t\n key: tw-test\n{client_extra}providers:\n - name: up\n base_url: http://{up}\n key: k\n protocol: anthropic\n" + ); + std::fs::write(dir.path().join("config.yaml"), cfg).unwrap(); + Bed { + dir, + gateway: SocketAddr::from(([127, 0, 0, 1], port)), + } + } + + fn config(&self) -> PathBuf { + self.dir.path().join("config.yaml") + } + + fn core(&self, name: &str) -> Child { + self.core_with(name, &[]) + } + + /// 多带几个 `serve` 的参数（比如 `--parent`） + fn core_with(&self, name: &str, args: &[&str]) -> Child { + // 日志写在 stdout（tracing 的默认），stderr 留给 panic：两个都收 + let log = std::fs::File::create(self.dir.path().join(format!("{name}.log"))).unwrap(); + self.command() + .args(args) + .stdout(log.try_clone().unwrap()) + .stderr(log) + .spawn() + .unwrap() + } + + fn arm(&self, wait_secs: u64) -> Option { + self.call( + "POST", + "/handover/arm", + Some(&format!(r#"{{"wait_secs":{wait_secs},"relaunch":null}}"#)), + ) + } + + /// 此刻答控制面的那个 core 记着的请求号（最近的在前） + fn history_ids(&self) -> Vec { + self.call("GET", "/history?limit=500", None) + .and_then(|v| { + v.as_array().map(|rows| { + rows.iter() + .filter_map(|r| r["id"].as_u64()) + .collect::>() + }) + }) + .unwrap_or_default() + } + + /// 密钥 `t` 那条按天上限此刻用了多少 + fn used_today(&self) -> Option { + self.call("GET", "/keys", None)? + .as_array()? + .iter() + .find(|k| k["name"] == "t")?["limits"][0]["used"] + .as_u64() + } + + fn draining_marked(&self) -> bool { + self.dir.path().join("twcore.draining").exists() + } + + /// 和桌面端一样，stdout、stderr 都接管道（桌面端退出时这两根管道的读端就没了） + fn core_on_a_pipe(&self) -> Child { + self.command() + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap() + } + + fn command(&self) -> Command { + let mut c = Command::new(BIN); + c.arg("--config") + .arg(self.config()) + .arg("serve") + .env("TWCORE_LOG", "info") + .env_remove("HTTP_PROXY") + .env_remove("HTTPS_PROXY") + .env_remove("http_proxy") + .env_remove("https_proxy"); + c + } + + /// `twcore call`：控制面上发一个请求，拿回 JSON + fn call(&self, method: &str, path: &str, data: Option<&str>) -> Option { + let mut c = Command::new(BIN); + c.arg("--config") + .arg(self.config()) + .args(["call", path, "-X", method]); + if let Some(d) = data { + c.args(["-d", d]); + } + let out = c.output().ok()?; + serde_json::from_slice(&out.stdout).ok() + } + + fn status(&self) -> Option { + self.call("GET", "/status", None) + } + + fn lock_pid(&self) -> Option { + std::fs::read_to_string(self.dir.path().join("twcore.lock")) + .ok()? + .trim() + .parse() + .ok() + } + + fn logs(&self) -> String { + let mut s = String::new(); + for n in ["a", "b", "c", "twcore.handover"] { + if let Ok(t) = std::fs::read_to_string(self.dir.path().join(format!("{n}.log"))) { + s.push_str(&format!("── {n} ──\n{t}\n")); + } + } + s + } +} + +/// 等到 `f` 成立（最多 `within`） +async fn until(within: Duration, mut f: impl FnMut() -> bool) -> bool { + let t = Instant::now(); + while t.elapsed() < within { + if f() { + return true; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + false +} + +fn exited(c: &mut Child) -> bool { + matches!(c.try_wait(), Ok(Some(_))) +} + +/// 收尸（测试失败时也别把进程留下） +struct Reap(Vec); + +impl Drop for Reap { + fn drop(&mut self) { + for c in &mut self.0 { + let _ = c.kill(); + let _ = c.wait(); + } + } +} + +/// 每个请求一条新连接：**连接被拒绝、被重置都会算出来**（复用连接的话看不出监听那一边） +fn fresh_client() -> reqwest::Client { + reqwest::Client::builder() + .pool_max_idle_per_host(0) + .no_proxy() + .build() + .unwrap() +} + +async fn quick(c: &reqwest::Client, gw: SocketAddr) -> Result<(), String> { + let r = c + .post(format!("http://{gw}/v1/messages")) + .header("x-api-key", "tw-test") + .header("content-type", "application/json") + .body(r#"{"model":"m","max_tokens":5,"messages":[{"role":"user","content":"hi"}]}"#) + .timeout(Duration::from_secs(10)) + .send() + .await + .map_err(|e| format!("{e:?}"))?; + let status = r.status(); + r.bytes().await.map_err(|e| format!("{e:?}"))?; + if status.is_success() { + Ok(()) + } else { + Err(format!("{status}")) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn a_handover_under_load_refuses_no_connection_and_finishes_the_long_stream() { + let bed = Bed::new(25).await; + let mut reap = Reap(vec![bed.core_on_a_pipe()]); + assert!( + until(Duration::from_secs(60), || bed.status().is_some()).await, + "a 没起来\n{}", + bed.logs() + ); + let a_pid = bed.status().unwrap()["pid"].as_u64().unwrap() as u32; + assert!( + until(Duration::from_secs(5), || std::net::TcpStream::connect( + bed.gateway + ) + .is_ok()) + .await + ); + + // 一条长流，交接前开始（约 5 秒） + let gw = bed.gateway; + let long = tokio::spawn(async move { + let r = fresh_client() + .post(format!("http://{gw}/v1/messages")) + .header("x-api-key", "tw-test") + .header("content-type", "application/json") + .body(r#"{"model":"m","max_tokens":5,"stream":true,"messages":[{"role":"user","content":"slow"}]}"#) + .timeout(Duration::from_secs(60)) + .send() + .await + .map_err(|e| format!("{e:?}"))?; + let body = r.text().await.map_err(|e| format!("{e:?}"))?; + Ok::(body) + }); + + // 压测：几路并发，每个请求一条新连接 + let stop = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let (ok, bad) = ( + Arc::new(AtomicUsize::new(0)), + Arc::new(std::sync::Mutex::new(Vec::new())), + ); + let mut load = Vec::new(); + for _ in 0..4 { + let (stop, ok, bad) = (stop.clone(), ok.clone(), bad.clone()); + load.push(tokio::spawn(async move { + let c = fresh_client(); + while !stop.load(Ordering::Relaxed) { + match quick(&c, gw).await { + Ok(()) => { + ok.fetch_add(1, Ordering::Relaxed); + } + Err(e) => bad.lock().unwrap().push(e), + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + })); + } + tokio::time::sleep(Duration::from_millis(500)).await; + + // arm，然后起接替者（桌面端的顺序：arm → 自己重启 → 新的桌面端起新 core） + let armed = bed.call( + "POST", + "/handover/arm", + Some(r#"{"wait_secs":30,"relaunch":null}"#), + ); + assert_eq!( + armed.as_ref().and_then(|v| v["version"].as_u64()), + Some(1), + "{armed:?}\n{}", + bed.logs() + ); + // 桌面端退出：a 的 stderr 那根管道没人读了。a 之后再写日志不能挂（日志已经改写进数据目录） + drop(reap.0[0].stdout.take()); + drop(reap.0[0].stderr.take()); + reap.0.push(bed.core("b")); + let b_pid = reap.0[1].id(); + assert!( + until(Duration::from_secs(60), || bed.lock_pid() == Some(b_pid)).await, + "锁没交给 b\n{}", + bed.logs() + ); + // 控制面现在是 b 在答 + let st = bed.status().unwrap(); + assert_eq!(st["pid"].as_u64(), Some(b_pid as u64)); + assert!( + matches!(st["handover"].as_str(), Some("draining" | "idle")), + "{st}" + ); + + // 长流在 a 上完整结束 + let body = long.await.unwrap().expect("长流断了"); + assert!(body.contains("part24"), "长流没收全：{body}"); + assert!(body.contains("message_stop"), "长流没收尾：{body}"); + + // a 排空后退出 + assert!( + until(Duration::from_secs(60), || exited(&mut reap.0[0])).await, + "a 没退\n{}", + bed.logs() + ); + tokio::time::sleep(Duration::from_millis(500)).await; + stop.store(true, Ordering::Relaxed); + for l in load { + l.await.unwrap(); + } + let bad = bad.lock().unwrap().clone(); + assert!( + bad.is_empty(), + "交接期间有请求失败：{bad:?}\n{}", + bed.logs() + ); + assert!(ok.load(Ordering::Relaxed) > 20, "压测没压起来"); + assert_eq!(bed.lock_pid(), Some(b_pid), "a 退出时把锁删了"); + // a 是自己排空完退的，不是写坏了管道挂掉的；它 arm 之后的日志在数据目录里 + let status = reap.0[0].try_wait().unwrap().unwrap(); + assert!( + status.success(), + "a 不是正常退出的：{status:?}\n{}", + bed.logs() + ); + let after = std::fs::read_to_string(bed.dir.path().join("twcore.handover.log")).unwrap(); + assert!(after.contains("drained"), "{after}"); + assert_ne!(a_pid, b_pid); + // 平常了 + assert!( + until(Duration::from_secs(10), || { + bed.status() + .is_some_and(|s| s["handover"] == "idle" && s["draining"] == 0) + }) + .await, + "{:?}", + bed.status() + ); +} + +/// 自己扮接替者：连交接端点，打招呼，收下控制面和快照。 +async fn fake_successor(dir: &Path, version: u32) -> tw_handover::Link { + let mut link = tw_handover::connect(dir) + .await + .unwrap() + .expect("arm 之后交接端点该开着"); + link.send(&tw_handover::Message::Hello { + version, + core_version: "test".into(), + pid: std::process::id(), + }) + .await + .unwrap(); + link +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_successor_that_leaves_before_committing_leaves_the_old_core_serving() { + let bed = Bed::new(1).await; + let mut reap = Reap(vec![bed.core("a")]); + assert!( + until(Duration::from_secs(60), || bed.status().is_some()).await, + "{}", + bed.logs() + ); + let armed = bed.call( + "POST", + "/handover/arm", + Some(r#"{"wait_secs":30,"relaunch":null}"#), + ); + assert!(armed.is_some_and(|v| v["version"] == 1)); + + // 要监听、收下，然后不提交就走（F4）：块结束时连接和套接字都丢掉 + { + let mut link = fake_successor(bed.dir.path(), tw_handover::VERSION).await; + let (m, socks) = link.recv().await.unwrap(); + assert!(matches!(m, tw_handover::Message::Control { .. }), "{m:?}"); + assert_eq!(socks.len(), 1); + let (m, _) = link.recv().await.unwrap(); + assert!(matches!(m, tw_handover::Message::Snapshot { .. }), "{m:?}"); + link.send(&tw_handover::Message::WantListeners) + .await + .unwrap(); + // 第二份快照：旧 core 暂停续期之后导出的那一份 + let (m, _) = link.recv().await.unwrap(); + assert!(matches!(m, tw_handover::Message::Snapshot { .. }), "{m:?}"); + let (m, socks) = link.recv().await.unwrap(); + assert!(matches!(m, tw_handover::Message::Listeners { .. }), "{m:?}"); + assert!(!socks.is_empty()); + } + + // 旧的回到待交接：控制面又是它在答，网关照常 + assert!( + until(Duration::from_secs(10), || bed + .status() + .is_some_and(|s| s["handover"] == "armed")) + .await, + "{:?}\n{}", + bed.status(), + bed.logs() + ); + assert!(quick(&fresh_client(), bed.gateway).await.is_ok()); + assert!(!exited(&mut reap.0[0])); + // 锁还是它的 + assert_eq!(bed.lock_pid(), Some(reap.0[0].id())); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_different_handover_version_falls_back_to_drain_and_exit() { + let bed = Bed::new(1).await; + let mut reap = Reap(vec![bed.core("a")]); + assert!( + until(Duration::from_secs(60), || bed.status().is_some()).await, + "{}", + bed.logs() + ); + assert!( + bed.call( + "POST", + "/handover/arm", + Some(r#"{"wait_secs":30,"relaunch":null}"#) + ) + .is_some() + ); + let mut link = fake_successor(bed.dir.path(), tw_handover::VERSION + 1).await; + let (m, socks) = link.recv().await.unwrap(); + assert!(matches!(m, tw_handover::Message::Refused { .. }), "{m:?}"); + assert!(socks.is_empty()); + link.send(&tw_handover::Message::Fallback).await.unwrap(); + // H2：没有在跑的请求，马上退 + assert!( + until(Duration::from_secs(10), || exited(&mut reap.0[0])).await, + "{}", + bed.logs() + ); + // 按现行方式退出：锁文件删了，下一个照常启动 + assert_eq!(bed.lock_pid(), None); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn without_a_successor_the_armed_core_drains_and_exits() { + let bed = Bed::new(1).await; + let mut reap = Reap(vec![bed.core("a")]); + assert!( + until(Duration::from_secs(60), || bed.status().is_some()).await, + "{}", + bed.logs() + ); + let armed = bed.call( + "POST", + "/handover/arm", + Some(r#"{"wait_secs":1,"relaunch":null}"#), + ); + assert!(armed.is_some_and(|v| v["version"] == 1)); + assert_eq!(bed.status().unwrap()["handover"], "armed"); + // H7：等一个周期、没有可重开的，排空退出 + assert!( + until(Duration::from_secs(10), || exited(&mut reap.0[0])).await, + "{}", + bed.logs() + ); + assert_eq!(bed.lock_pid(), None); +} + +/// 扮接替者走到拿到监听为止：招呼、控制面、第一份快照、要监听、第二份快照、监听。 +async fn take_listeners(dir: &Path) -> (tw_handover::Link, Vec) { + let mut link = fake_successor(dir, tw_handover::VERSION).await; + let (m, _control) = link.recv().await.unwrap(); + assert!(matches!(m, tw_handover::Message::Control { .. }), "{m:?}"); + let (m, _) = link.recv().await.unwrap(); + assert!(matches!(m, tw_handover::Message::Snapshot { .. }), "{m:?}"); + link.send(&tw_handover::Message::WantListeners) + .await + .unwrap(); + let (m, _) = link.recv().await.unwrap(); + assert!(matches!(m, tw_handover::Message::Snapshot { .. }), "{m:?}"); + let (m, socks) = link.recv().await.unwrap(); + assert!(matches!(m, tw_handover::Message::Listeners { .. }), "{m:?}"); + (link, socks) +} + +/// 提交要确认（CommitAck）：旧 core 收到 Commit 先回确认，再停止 accept、排空、说 Bye 退出。 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn the_old_core_acknowledges_the_commit_then_drains_and_says_bye() { + let bed = Bed::new(1).await; + let mut reap = Reap(vec![bed.core("a")]); + assert!( + until(Duration::from_secs(60), || bed.status().is_some()).await, + "{}", + bed.logs() + ); + assert!(bed.arm(30).is_some_and(|v| v["version"] == 1)); + let (mut link, _socks) = take_listeners(bed.dir.path()).await; + link.send(&tw_handover::Message::Commit).await.unwrap(); + let (m, _) = link.recv().await.unwrap(); + assert_eq!(m, tw_handover::Message::CommitAck, "提交之后头一条该是确认"); + // 排空：没有在跑的请求，马上说 Bye + let bye = tokio::time::timeout(Duration::from_secs(15), async { + loop { + match link.recv().await { + Ok((tw_handover::Message::Bye, _)) => return true, + Ok(_) => {} + Err(_) => return false, + } + } + }) + .await; + assert_eq!(bye.ok(), Some(true), "{}", bed.logs()); + assert!( + until(Duration::from_secs(10), || exited(&mut reap.0[0])).await, + "{}", + bed.logs() + ); + assert!(reap.0[0].try_wait().unwrap().unwrap().success()); +} + +/// 旧 core 等提交的时限过了就回到待交接：**迟到的 Commit 不回确认**、连接关掉，它照常服务。 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_commit_that_comes_after_the_old_core_gave_up_is_not_acknowledged() { + let bed = Bed::new(1).await; + let mut reap = Reap(vec![bed.core("a")]); + assert!( + until(Duration::from_secs(60), || bed.status().is_some()).await, + "{}", + bed.logs() + ); + assert!(bed.arm(120).is_some_and(|v| v["version"] == 1)); + let (mut link, socks) = take_listeners(bed.dir.path()).await; + drop(socks); + // 旧 core 等 30 秒 + tokio::time::sleep(Duration::from_secs(32)).await; + let _ = link.send(&tw_handover::Message::Commit).await; + let late = tokio::time::timeout(Duration::from_secs(5), link.recv()).await; + assert!( + !matches!(late, Ok(Ok((tw_handover::Message::CommitAck, _)))), + "放弃之后又确认了一个迟到的提交" + ); + assert!( + until(Duration::from_secs(10), || bed + .status() + .is_some_and(|s| s["handover"] == "armed")) + .await, + "{:?}\n{}", + bed.status(), + bed.logs() + ); + assert!(quick(&fresh_client(), bed.gateway).await.is_ok()); + assert!(!exited(&mut reap.0[0])); +} + +/// 接替者发出提交、旧 core 却不确认（它还活着）：**接替者不能当主**，退出、不改锁。 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_successor_whose_commit_is_not_acknowledged_exits_without_taking_the_lock() { + let bed = Bed::new(1).await; + // 这个测试进程扮旧 core：锁写自己的 pid，开交接端点 + std::fs::write( + bed.dir.path().join("twcore.lock"), + std::process::id().to_string(), + ) + .unwrap(); + let endpoint = tw_handover::Endpoint::bind(bed.dir.path()).await.unwrap(); + let mut reap = Reap(vec![bed.core("b")]); + + let mut link = tokio::time::timeout(Duration::from_secs(20), endpoint.accept()) + .await + .expect("接替者没连上来") + .unwrap(); + let (m, _) = link.recv().await.unwrap(); + let tw_handover::Message::Hello { pid, .. } = m else { + panic!("头一条该是 Hello：{m:?}"); + }; + let peer = link.peer_pid().unwrap_or(pid); + #[cfg(unix)] + let control = tw_handover::Socket::Unix( + std::os::unix::net::UnixListener::bind(bed.dir.path().join("ctl.sock")).unwrap(), + ); + #[cfg(not(unix))] + let control = tw_handover::Socket::Tcp(std::net::TcpListener::bind("127.0.0.1:0").unwrap()); + let info = tw_handover::ListenerInfo { + role: tw_handover::Role::Control, + want: None, + }; + link.send_with(&tw_handover::Message::Control { info }, &[&control], peer) + .await + .unwrap(); + let snapshot = serde_json::json!({ "gateway": {}, "next_id": 1 }); + link.send(&tw_handover::Message::Snapshot { + data: snapshot.clone(), + }) + .await + .unwrap(); + let (m, _) = tokio::time::timeout(Duration::from_secs(20), link.recv()) + .await + .expect("没要监听") + .unwrap(); + assert_eq!(m, tw_handover::Message::WantListeners, "{}", bed.logs()); + link.send(&tw_handover::Message::Snapshot { data: snapshot }) + .await + .unwrap(); + let gw = tw_handover::Socket::Tcp(std::net::TcpListener::bind(bed.gateway).unwrap()); + let infos = vec![tw_handover::ListenerInfo { + role: tw_handover::Role::Gateway, + want: Some(bed.gateway), + }]; + link.send_with(&tw_handover::Message::Listeners { infos }, &[&gw], peer) + .await + .unwrap(); + let (m, _) = tokio::time::timeout(Duration::from_secs(25), link.recv()) + .await + .expect("没提交") + .unwrap(); + assert_eq!(m, tw_handover::Message::Commit); + // 不确认，连接也不关：接替者自己等不到就退出 + assert!( + until(Duration::from_secs(25), || exited(&mut reap.0[0])).await, + "没确认的接替者还在跑\n{}", + bed.logs() + ); + assert!(!reap.0[0].try_wait().unwrap().unwrap().success()); + assert_eq!(bed.lock_pid(), Some(std::process::id()), "没确认就改了锁"); + drop(link); +} + +/// 读不懂的第一帧（将来 Hello 改了样子）：**也回 Refused**，不是一声不响地断开。 +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn an_unreadable_greeting_is_answered_with_refused() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let bed = Bed::new(1).await; + let _reap = Reap(vec![bed.core("a")]); + assert!( + until(Duration::from_secs(60), || bed.status().is_some()).await, + "{}", + bed.logs() + ); + assert!(bed.arm(30).is_some_and(|v| v["version"] == 1)); + let mut s = tokio::net::UnixStream::connect(bed.dir.path().join(tw_handover::ENDPOINT_FILE)) + .await + .unwrap(); + let junk = br#"{"hello from the future":true}"#; + s.write_all(&(junk.len() as u32).to_be_bytes()) + .await + .unwrap(); + s.write_all(junk).await.unwrap(); + let mut len = [0u8; 4]; + tokio::time::timeout(Duration::from_secs(10), s.read_exact(&mut len)) + .await + .expect("读不懂的招呼没有回话") + .unwrap(); + let mut body = vec![0u8; u32::from_be_bytes(len) as usize]; + s.read_exact(&mut body).await.unwrap(); + let body = String::from_utf8_lossy(&body); + assert!(body.contains("refused"), "{body}"); + drop(s); + // 之后照旧待交接 + assert!( + until(Duration::from_secs(10), || bed + .status() + .is_some_and(|s| s["handover"] == "armed")) + .await, + "{}", + bed.logs() + ); +} + +/// arm 之后桌面端没有重启（更新失败、取消）：到点时原来的父进程还在，**撤销 arm**、恢复 +/// 父进程守望，不重开、不退出。arm 时另起会话（unix），日志接着往后写。 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn an_arm_is_called_off_when_the_desktop_app_did_not_restart() { + let bed = Bed::new(1).await; + // 扮桌面端：一个活着的进程 + #[cfg(windows)] + let parent = Command::new("ping") + .args(["-n", "120", "127.0.0.1"]) + .stdout(Stdio::null()) + .spawn() + .unwrap(); + #[cfg(not(windows))] + let parent = Command::new("sleep").arg("120").spawn().unwrap(); + let ppid = parent.id().to_string(); + let mut reap = Reap(vec![bed.core_with("a", &["--parent", &ppid]), parent]); + assert!( + until(Duration::from_secs(60), || bed.status().is_some()).await, + "{}", + bed.logs() + ); + for round in 0..2 { + assert!( + bed.arm(1).is_some_and(|v| v["version"] == 1), + "第 {round} 次" + ); + assert_eq!(bed.status().unwrap()["handover"], "armed"); + assert!( + until(Duration::from_secs(8), || bed + .status() + .is_some_and(|s| s["handover"] == "idle")) + .await, + "父进程还在，arm 没撤销\n{}", + bed.logs() + ); + assert!(!exited(&mut reap.0[0]), "撤销 arm 时退出了"); + assert!(quick(&fresh_client(), bed.gateway).await.is_ok()); + } + #[cfg(unix)] + { + let a = reap.0[0].id() as i32; + // SAFETY: getsid 只是问 + assert_eq!(unsafe { libc::getsid(a) }, a, "arm 之后没有另起会话"); + } + let log = std::fs::read_to_string(bed.dir.path().join("twcore.handover.log")).unwrap(); + assert_eq!( + log.matches("no longer exiting with the parent").count(), + 2, + "两次 arm 的日志该都在（接着往后写）：{log}" + ); + // 父进程守望恢复了：父进程没了，它跟着退 + let _ = reap.0[1].kill(); + let _ = reap.0[1].wait(); + assert!( + until(Duration::from_secs(10), || exited(&mut reap.0[0])).await, + "父进程没了还在跑\n{}", + bed.logs() + ); +} + +/// 一条长流跑着的时候走一遍：A 交给 B；B 在 A 排空期间自己的请求照样结算（过了预留的宽限期 +/// 也还算在用量里）；B 在 A 排空完之前退出，**锁和排空标记都留着**；C 照常起来，看见 A 还在 +/// 排空，就以「前任在排空」起来、请求号跳过去、等 A 退了才开库；开库之后用量以库为准、不重算。 +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn a_successor_that_exits_while_its_predecessor_drains_leaves_the_next_core_waiting() { + let bed = Bed::with_daily_limit(450).await; // 长流约 90 秒 + let mut reap = Reap(vec![bed.core("a")]); + assert!( + until(Duration::from_secs(60), || bed.status().is_some()).await, + "{}", + bed.logs() + ); + let a_pid = reap.0[0].id(); + let client = fresh_client(); + for _ in 0..3 { + quick(&client, bed.gateway).await.unwrap(); + } + let gw = bed.gateway; + let long = tokio::spawn(async move { + let r = fresh_client() + .post(format!("http://{gw}/v1/messages")) + .header("x-api-key", "tw-test") + .header("content-type", "application/json") + .body(r#"{"model":"m","max_tokens":5,"stream":true,"messages":[{"role":"user","content":"slow"}]}"#) + .timeout(Duration::from_secs(200)) + .send() + .await + .map_err(|e| format!("{e:?}"))?; + r.text().await.map_err(|e| format!("{e:?}")) + }); + tokio::time::sleep(Duration::from_millis(500)).await; + + // A → B + assert!(bed.arm(60).is_some_and(|v| v["version"] == 1)); + reap.0.push(bed.core("b")); + let b_pid = reap.0[1].id(); + assert!( + until(Duration::from_secs(60), || bed.lock_pid() == Some(b_pid)).await, + "{}", + bed.logs() + ); + assert_eq!(bed.status().unwrap()["handover"], "draining"); + assert!(bed.draining_marked(), "没记下 A 在排空"); + + // B 自己的请求：A 排空期间库没开，照样结算（过了 60 秒的宽限期也不丢） + for _ in 0..5 { + quick(&client, bed.gateway).await.unwrap(); + } + let _ = bed.used_today(); + tokio::time::sleep(Duration::from_secs(65)).await; + let used = bed.used_today().unwrap(); + assert!( + used >= 3 + 5, + "B 自己的请求没算进用量：{used}\n{}", + bed.logs() + ); + + // B 在 A 排空完之前退出 + assert!(bed.call("POST", "/shutdown", None).is_some()); + assert!( + until(Duration::from_secs(10), || exited(&mut reap.0[1])).await, + "{}", + bed.logs() + ); + assert!(!exited(&mut reap.0[0]), "A 不该这时就退了（长流还在跑）"); + assert_eq!(bed.lock_pid(), Some(b_pid), "A 还在排空，B 把锁删了"); + assert!(bed.draining_marked(), "排空标记没了"); + + // C 照常起来：前任在排空 + reap.0.push(bed.core("c")); + let c_pid = reap.0[2].id(); + assert!( + until(Duration::from_secs(60), || bed + .status() + .is_some_and(|s| s["pid"].as_u64() == Some(c_pid as u64))) + .await, + "{}", + bed.logs() + ); + assert_eq!(bed.lock_pid(), Some(c_pid)); + assert_eq!(bed.status().unwrap()["handover"], "draining"); + for _ in 0..2 { + quick(&client, bed.gateway).await.unwrap(); + } + let early = bed.history_ids(); + assert!( + early.iter().all(|id| *id > 15_000_000), + "C 的请求号没跳过前任的：{early:?}" + ); + + // 长流在 A 上收完，A 退出；C 回到平常、开库 + let body = long.await.unwrap().expect("长流断了"); + assert!(body.contains("message_stop"), "{body}"); + assert!( + until(Duration::from_secs(60), || exited(&mut reap.0[0])).await, + "{}", + bed.logs() + ); + assert!( + until(Duration::from_secs(60), || bed + .status() + .is_some_and(|s| s["handover"] == "idle")) + .await, + "{}", + bed.logs() + ); + assert!(until(Duration::from_secs(10), || !bed.draining_marked()).await); + // 库里：A 的 3 个短请求和那条长流、C 的 2 个（B 的那 5 个没进库：它退出时库还没开） + let ids = bed.history_ids(); + let mut unique = ids.clone(); + unique.sort_unstable(); + unique.dedup(); + assert_eq!(unique.len(), ids.len(), "请求号撞了：{ids:?}"); + assert_eq!(ids.len(), 3 + 1 + 2, "{ids:?}\n{}", bed.logs()); + // 用量以库为准：不把 C 排队时结算过的再算一遍 + assert_eq!(bed.used_today(), Some(3 + 1 + 2), "{}", bed.logs()); + assert_ne!(a_pid, c_pid); + assert!( + std::fs::read_to_string(bed.dir.path().join("b.log")) + .unwrap() + .contains("queued request events are not recorded"), + "B 退出时没说丢了几条\n{}", + bed.logs() + ); +} diff --git a/crates/tw-api/msg-codes.txt b/crates/tw-api/msg-codes.txt index c86f9a09..0f7bd9f0 100644 --- a/crates/tw-api/msg-codes.txt +++ b/crates/tw-api/msg-codes.txt @@ -178,6 +178,11 @@ control.group.weight_not_load_balance control.group.weight_not_member control.group.weight_out_of_range control.group_in_use +control.handover.busy +control.handover.endpoint_failed +control.handover.login_refused +control.handover.remote +control.handover.safe_mode control.header_no_value control.internal_error control.key_create_failed @@ -361,6 +366,7 @@ gw.oauth.backoff.no_token gw.oauth.backoff.status gw.oauth.backoff.unreachable gw.oauth.expired +gw.oauth.held gw.oauth.no_token gw.oauth.not_configured gw.oauth.rotation_no_manager diff --git a/crates/tw-api/src/ep.rs b/crates/tw-api/src/ep.rs index fa991be1..5bb4319a 100644 --- a/crates/tw-api/src/ep.rs +++ b/crates/tw-api/src/ep.rs @@ -12,6 +12,9 @@ endpoints! { Status: GET "/status", () => api::Status; /// 请网关退出（202） Shutdown: POST "/shutdown", () => api::Msg; + /// 请 core 等接替者：脱离父进程、开交接端点（升级不断线）。被拒时是结构化错误， + /// 码 `control.handover.*`，桌面端退回「等请求结束再重启」 + HandoverArm: POST "/handover/arm", api::HandoverArm => api::HandoverArmed; /// 本机的网卡，一张一行 Interfaces: GET "/interfaces", () => Vec; /// 实时事件流 diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index 8b100d50..a148f41f 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -1061,7 +1061,12 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// 新码：`config.key_limit_every*`、`config.upstream_limit_*`、`gw.key_limit.*_every`、 /// `gw.upstream.paused`、`gw.upstream.window_full`、`gw.upstream.period_used_up`、`gw.wait.*`。 /// 照 50 写的界面读不到 `wait_secs`，保存密钥的上限时没有 `every` 会被拒 -pub const CONTROL_API_VERSION: u32 = 51; +/// +/// **52 起升级不断线**：`POST /handover/arm`（[`HandoverArm`] → [`HandoverArmed`]）请 core 脱离 +/// 父进程、开交接端点、等接替者；[`Status`] 多了 `handover`（[`HandoverState`]）和 `draining` +/// （排空中的旧 core 手上还有几个请求，只有接管了它的新 core 报，平时是 0）。 +/// 新码：`control.handover.*` +pub const CONTROL_API_VERSION: u32 = 52; /// `/live` 的速度看最近多久（见 [`LiveView::tokens_per_sec`]），毫秒 pub const LIVE_RATE_WINDOW_MS: u64 = 3_000; @@ -1097,6 +1102,11 @@ pub struct Status { /// /// 重启网关之前要看它 —— 重启会掐断所有还没结束的流。 pub in_flight: usize, + /// 交接走到了哪一步。见 [`HandoverState`] + pub handover: HandoverState, + /// 排空中的旧 core 手上还有几个请求（它每结束一个就报一次）。**只有接管了它的新 core + /// 报**，平时是 0。界面把它加进在途个数，不单独提示 + pub draining: usize, /// 远程控制端口此刻的样子。 pub remote_control: RemoteControlView, /// 别的机器连网关该用的地址（`地址:端口`），**不含回环**。 @@ -1111,6 +1121,52 @@ pub struct Status { pub gateway_reachable: Vec, } +/// 交接走到了哪一步（[`Status::handover`]）。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +#[serde(rename_all = "snake_case")] +pub enum HandoverState { + /// 平常 + #[default] + Idle, + /// 已 arm：脱离了父进程，交接端点开着，等接替者 + Armed, + /// 接替者连上来了，还没提交 + TakingOver, + /// 已提交：这个 core 是接替者，旧的那个还在排空（[`Status::draining`]） + Draining, +} + +/// `POST /handover/arm` 的请求体。 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct HandoverArm { + /// 等接替者多久，秒。缺省 120 + pub wait_secs: Option, + /// 等不到时怎么重新打开桌面端（只试一次，之后再等一个 `wait_secs`，仍等不到就排空退出）。 + /// 空 = 不重新打开 + pub relaunch: Option, +} + +/// 重新打开桌面端的方式。 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum Relaunch { + /// macOS：`open -b ` + BundleId { id: String }, + /// Windows、Linux：直接起这个可执行文件，和 core 脱开 + Exe { path: String, args: Vec }, +} + +/// `POST /handover/arm` 的应答。 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct HandoverArmed { + /// 交接连接的版本号（`tw_handover::VERSION`）。接替者的不同就不交接，退回现行流程 + pub version: u32, +} + /// 远程控制端口（`listen.control.remote`）。 #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] @@ -7618,6 +7674,8 @@ mod tests { providers: 1, uptime_secs: 0, in_flight: 3, + handover: HandoverState::Draining, + draining: 2, remote_control: RemoteControlView::default(), gateway_reachable: vec![], }; @@ -7632,7 +7690,7 @@ mod tests { // 而不是「gateway_addr 是空字符串」这种约定。 let json = r#"{"api_version":1,"version":"x","pid":1,"gateway_addr":null, "config_path":"/x","clients":0,"providers":0,"uptime_secs":0, - "in_flight":0,"gateway_reachable":[], + "in_flight":0,"handover":"idle","draining":0,"gateway_reachable":[], "remote_control":{"enabled":false,"addr":null,"allow_from":[],"reachable":[]}}"#; let s: Status = serde_json::from_str(json).unwrap(); assert!(s.gateway_addr.is_none()); diff --git a/crates/tw-control/Cargo.toml b/crates/tw-control/Cargo.toml index 69eb079c..8da9bb56 100644 --- a/crates/tw-control/Cargo.toml +++ b/crates/tw-control/Cargo.toml @@ -17,6 +17,8 @@ tw-api = { workspace = true } tw-types = { workspace = true } # 每条控制面连接先握手 tw-link = { workspace = true } +# 升级不断线：控制面的监听交给接替的 core（套接字的类型在那里） +tw-handover = { workspace = true } # 列网卡用 getifaddrs —— 比引一个 crate 便宜 libc = { workspace = true } # 配置包（导出、导入、本机备份、恢复）。文件内容在线上是 base64 diff --git a/crates/tw-control/src/chatgpt.rs b/crates/tw-control/src/chatgpt.rs index 3627647a..00212605 100644 --- a/crates/tw-control/src/chatgpt.rs +++ b/crates/tw-control/src/chatgpt.rs @@ -101,6 +101,14 @@ impl Accounts { } } + /// 有一次登录正在等用户授权。**这时不交接**（升级不断线）：登录要在这个 core 上走完 + pub fn in_progress(&self) -> bool { + self.current.lock().is_ok_and(|g| { + g.as_ref() + .is_some_and(|c| c.status.status == tw_api::LoginStatus::Pending) + }) + } + fn set_status(&self, status: tw_api::ChatgptLoginStatus) { if let Ok(mut g) = self.current.lock() && let Some(c) = g.as_mut() @@ -232,6 +240,10 @@ async fn start( State(s): State, Json(req): Json, ) -> Result, Fail> { + // arm 之后不开始新的登录：回调端口不交接，登录要走完才能交（见 `handover`） + if crate::handover::refuses_logins(&s.handover) { + return Err(crate::handover::login_refused()); + } let want = Want::read(&s, &req)?; // 还没完成的上一次登录作废：回调端口要让出来 let prev = s.chatgpt.current.lock().ok().and_then(|mut g| g.take()); diff --git a/crates/tw-control/src/handover.rs b/crates/tw-control/src/handover.rs new file mode 100644 index 00000000..3a391aa4 --- /dev/null +++ b/crates/tw-control/src/handover.rs @@ -0,0 +1,417 @@ +//! 升级不断线：控制面这一边。 +//! +//! # 分工 +//! +//! 交接的状态机在 `twcore` 里（只有它同时看得见锁、监听、后台任务和交接连接）。这里 +//! 只放控制面要用的三样： +//! +//! - `POST /handover/arm` 的处理函数：先挡掉不该交接的（远程来的、安全模式、登录进行中）， +//! 过了门交给 `twcore`（[`Handover::arm_requests`]），等它把交接端点开好再答应； +//! - [`Status::handover`](tw_api::Status::handover) / `draining` 报的那两个数； +//! - 本机控制面的监听（[`ControlListener`]）：交出去、停止 accept、从交过来的上面服务。 +//! +//! # 为什么交出控制面时旧的要停止 accept +//! +//! 网关的监听两边短暂地同时 accept 没关系（每个连接只会被一方接受，哪边接都一样转发）。 +//! 控制面不行：新的桌面端要和新 core 说话，落到旧 core 上看到的是旧版本、旧协议。所以 +//! 控制面一交出去，旧的就不再从它上面接新连接；交接没成（新的提交前崩了），再接着接。 + +use std::sync::Mutex; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::Duration; + +use axum::Json; +use axum::extract::State; +use axum::http::StatusCode; +use tokio::sync::{mpsc, oneshot, watch}; +use tw_api::HandoverState; +use tw_types::{Msg, msg}; + +use crate::{ControlState, Fail, fail}; + +/// 等接替者多久，缺省。 +pub const DEFAULT_WAIT: Duration = Duration::from_secs(120); + +/// `twcore` 开交接端点最多等多久。开一个 socket 是毫秒级的事；等不到说明出了别的问题， +/// 桌面端退回现行流程 +const ARM_REPLY_WITHIN: Duration = Duration::from_secs(10); + +/// 一次 arm：交给 `twcore`，它开好交接端点、脱离父进程之后回话。 +pub struct Arm { + pub wait: Duration, + pub relaunch: Option, + /// `Err` 时那句话原样回给桌面端 + pub reply: oneshot::Sender>, +} + +/// 交接在控制面上的样子。`ControlState` 里一份。 +pub struct Handover { + phase: watch::Sender, + draining: AtomicUsize, + /// 这个 core 没有数据面（安全模式），没有可交接的 + no_data_plane: AtomicBool, + arms: Mutex>>, + control: ControlListener, +} + +impl Default for Handover { + fn default() -> Self { + Self { + phase: watch::channel(HandoverState::Idle).0, + draining: AtomicUsize::new(0), + no_data_plane: AtomicBool::new(false), + arms: Mutex::new(None), + control: ControlListener::default(), + } + } +} + +impl Handover { + pub fn phase(&self) -> HandoverState { + *self.phase.borrow() + } + + pub fn set_phase(&self, p: HandoverState) { + self.phase.send_replace(p); + } + + /// 跟着交接走到哪一步（测试和 `twcore` 等某一步用） + pub fn watch_phase(&self) -> watch::Receiver { + self.phase.subscribe() + } + + /// 排空中的旧 core 手上还有几个请求（[`tw_api::Status::draining`]） + pub fn draining(&self) -> usize { + self.draining.load(Ordering::Relaxed) + } + + pub fn set_draining(&self, n: usize) { + self.draining.store(n, Ordering::Relaxed); + } + + /// 这个 core 没有数据面：arm 一律拒绝（`control.handover.safe_mode`） + pub fn set_no_data_plane(&self, yes: bool) { + self.no_data_plane.store(yes, Ordering::Relaxed); + } + + /// `twcore` 接 arm 请求的那一头。**只有一个**：再调一次换掉上一个。没人接时 arm 被拒 + pub fn arm_requests(&self) -> mpsc::Receiver { + let (tx, rx) = mpsc::channel(4); + if let Ok(mut g) = self.arms.lock() { + *g = Some(tx); + } + rx + } + + /// 本机控制面的监听 + pub fn control(&self) -> &ControlListener { + &self.control + } +} + +/// 拒绝 arm 的几种情形，说给人听（码是契约，桌面端据此退回现行流程）。 +pub(crate) fn refused_remote() -> Fail { + fail( + StatusCode::FORBIDDEN, + msg!( + "control.handover.remote" => + "A handover cannot be prepared over a remote connection." + ), + ) +} + +pub(crate) fn refused_safe_mode() -> Fail { + fail( + StatusCode::CONFLICT, + msg!( + "control.handover.safe_mode" => + "The gateway is in safe mode and is not forwarding, so there is nothing to hand over." + ), + ) +} + +pub(crate) fn refused_busy() -> Fail { + fail( + StatusCode::CONFLICT, + msg!( + "control.handover.busy" => + "An account sign-in or another handover is in progress; try again once it has finished." + ), + ) +} + +/// arm 之后要开始一次账号登录：不行（回调端口不交接，登录要在接替者上走）。 +pub(crate) fn login_refused() -> Fail { + fail( + StatusCode::CONFLICT, + msg!( + "control.handover.login_refused" => + "The gateway is about to be updated; sign in again once the update has finished." + ), + ) +} + +/// 正在交接（已 arm、或者接替者正在接管）：这时候不开始新的账号登录 +pub(crate) fn refuses_logins(h: &Handover) -> bool { + matches!(h.phase(), HandoverState::Armed | HandoverState::TakingOver) +} + +/// 交接端点开不了（arm 答应不了）。 +pub fn endpoint_failed(detail: impl std::fmt::Display) -> Msg { + msg!( + "control.handover.endpoint_failed", detail = detail => + "The handover endpoint could not be opened: {detail}" + ) +} + +/// `POST /handover/arm` +pub(crate) async fn arm( + State(s): State, + Json(req): Json, +) -> Result, Fail> { + // 远程来的不行：服务器上的 core 归服务管理器管，也没有桌面端会来接替它 + if crate::remote::is_remote() { + return Err(refused_remote()); + } + let h = &s.handover; + if h.no_data_plane.load(Ordering::Relaxed) { + return Err(refused_safe_mode()); + } + // **登录进行中不交接**：回调端口（1455/1457）不交接，登录要在这个 core 上走完 + if s.chatgpt.in_progress() || s.zai.in_progress() { + return Err(refused_busy()); + } + // 同一时刻至多一个交接（H8）：已经有接替者在接管了 + if h.phase() == HandoverState::TakingOver { + return Err(refused_busy()); + } + let wait = req + .wait_secs + .filter(|s| *s > 0) + .map(Duration::from_secs) + .unwrap_or(DEFAULT_WAIT); + let Some(tx) = h.arms.lock().ok().and_then(|g| g.clone()) else { + // 没人接（这个进程不跑交接）：和安全模式一样，没有可交接的 + return Err(refused_safe_mode()); + }; + let (reply, answer) = oneshot::channel(); + let sent = tx + .send(Arm { + wait, + relaunch: req.relaunch, + reply, + }) + .await; + if sent.is_err() { + return Err(refused_safe_mode()); + } + match tokio::time::timeout(ARM_REPLY_WITHIN, answer).await { + Ok(Ok(Ok(()))) => { + tracing::info!(wait_secs = wait.as_secs(), "armed for a handover"); + Ok(Json(tw_api::HandoverArmed { + version: tw_handover::VERSION, + })) + } + Ok(Ok(Err(m))) => Err(fail(StatusCode::INTERNAL_SERVER_ERROR, m)), + _ => Err(crate::internal( + "the handover endpoint did not open in time", + )), + } +} + +/// 本机控制面的监听：交出去、停止 accept、从交过来的上面服务。 +/// +/// `serve` 起来时把一个复制出来的句柄放在这里（同一个内核套接字），交接时交出去的是 +/// 它的又一个复制；停止 accept 只是 accept 循环不再去接，句柄一直留着 —— 交接没成时 +/// 接着接。 +pub struct ControlListener { + accepting: watch::Sender, + /// 跳一下：已经连着的本机连接全部断开（见 [`ControlListener::close_connections`]） + closing: watch::Sender, + held: Mutex>, + inherited: Mutex>, +} + +impl Default for ControlListener { + fn default() -> Self { + Self { + accepting: watch::channel(true).0, + closing: watch::channel(0).0, + held: Mutex::new(None), + inherited: Mutex::new(None), + } + } +} + +impl ControlListener { + /// 新 core：控制面在旧 core 交过来的这个上面服务，不自己绑。**在 `serve` 之前调**。 + pub fn inherit(&self, s: tw_handover::Socket) { + if let Ok(mut g) = self.inherited.lock() { + *g = Some(s); + } + } + + pub(crate) fn take_inherited(&self) -> Option { + self.inherited.lock().ok().and_then(|mut g| g.take()) + } + + pub(crate) fn hold(&self, s: tw_handover::Socket) { + if let Ok(mut g) = self.held.lock() { + *g = Some(s); + } + } + + /// 旧 core：交出去的那个（复制出来的句柄）。控制面还没起来时是错误 + pub fn export(&self) -> std::io::Result { + let g = self + .held + .lock() + .map_err(|_| std::io::Error::other("the control listener lock is poisoned"))?; + match g.as_ref() { + Some(s) => clone_socket(s), + None => Err(std::io::Error::other("the control plane is not listening")), + } + } + + /// 不再从控制面接新连接（已经连着的不动） + pub fn pause(&self) { + self.accepting.send_replace(false); + } + + /// 接着接 + pub fn resume(&self) { + self.accepting.send_replace(true); + } + + pub(crate) fn accepting(&self) -> watch::Receiver { + self.accepting.subscribe() + } + + /// 断开已经连着的本机控制面连接（交出控制面时调）。 + /// + /// **为什么断开而不是只拒绝改配置**：连着的多半是新桌面端在交接之前连上来的（它的事件流、 + /// 请求都落在这个旧 core 上），只拒绝改动的话它照样从这里读状态、收事件，看到的是一个已经 + /// 不再转发新请求的进程；断开之后它按断线重连，连到的就是接替者。命令行的一次 `twcore call` + /// 碰巧在这时候被断，重来一次就好。交接没成时控制面恢复 accept，断开的那些重连回到这里 + pub fn close_connections(&self) { + self.closing.send_modify(|n| *n += 1); + } + + pub(crate) fn closing(&self) -> watch::Receiver { + self.closing.subscribe() + } +} + +/// 等「断开」那一下。没有信号（测试）就一直等 +pub(crate) async fn closed(rx: &mut Option>) { + match rx { + Some(rx) => { + if rx.changed().await.is_err() { + std::future::pending::<()>().await; + } + } + None => std::future::pending().await, + } +} + +/// 再复制一个句柄。 +pub fn clone_socket(s: &tw_handover::Socket) -> std::io::Result { + Ok(match s { + tw_handover::Socket::Tcp(l) => tw_handover::Socket::Tcp(l.try_clone()?), + #[cfg(unix)] + tw_handover::Socket::Unix(l) => tw_handover::Socket::Unix(l.try_clone()?), + }) +} + +/// 等到可以接（`accepting` 为真）。通道没了就当一直能接。 +pub(crate) async fn until_accepting(rx: &mut watch::Receiver) { + let _ = rx.wait_for(|on| *on).await; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn pausing_and_resuming_the_control_plane_is_seen_by_the_accept_loop() { + let c = ControlListener::default(); + let mut rx = c.accepting(); + // 默认能接 + tokio::time::timeout(Duration::from_secs(1), until_accepting(&mut rx)) + .await + .expect("默认就该接"); + c.pause(); + assert!( + tokio::time::timeout(Duration::from_millis(50), until_accepting(&mut rx)) + .await + .is_err(), + "停了还在接" + ); + c.resume(); + tokio::time::timeout(Duration::from_secs(1), until_accepting(&mut rx)) + .await + .expect("恢复之后该接"); + } + + #[tokio::test] + async fn arming_over_the_remote_port_is_refused() { + // 远程连接进来的请求带着「远程」的标记：处理函数第一步就挡掉，不看别的 + let refused = + crate::remote::as_remote(async { crate::remote::is_remote().then(refused_remote) }) + .await + .expect("远程连接上该被挡掉"); + assert_eq!(refused.0, StatusCode::FORBIDDEN); + assert_eq!(refused.1.0.code, "control.handover.remote"); + } + + #[test] + fn only_a_core_being_handed_over_refuses_sign_ins() { + let h = Handover::default(); + for (p, refused) in [ + (HandoverState::Idle, false), + (HandoverState::Armed, true), + (HandoverState::TakingOver, true), + // 交接完成、前任在排空的新 core 是主：照常能登 + (HandoverState::Draining, false), + ] { + h.set_phase(p); + assert_eq!(refuses_logins(&h), refused, "{p:?}"); + } + } + + #[tokio::test] + async fn closing_wakes_the_connections_and_no_signal_waits_forever() { + let c = ControlListener::default(); + let mut rx = Some(c.closing()); + c.close_connections(); + tokio::time::timeout(Duration::from_secs(1), closed(&mut rx)) + .await + .expect("断开的信号没到"); + let mut none = None; + assert!( + tokio::time::timeout(Duration::from_millis(50), closed(&mut none)) + .await + .is_err() + ); + } + + #[test] + fn exporting_before_listening_is_an_error_not_a_panic() { + assert!(ControlListener::default().export().is_err()); + } + + #[test] + fn an_exported_socket_is_the_same_listener() { + let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = l.local_addr().unwrap(); + let c = ControlListener::default(); + c.hold(tw_handover::Socket::Tcp(l)); + #[allow(clippy::infallible_destructuring_match)] + let dup = match c.export().unwrap() { + tw_handover::Socket::Tcp(l) => l, + #[cfg(unix)] + tw_handover::Socket::Unix(_) => panic!("交出去的不是 TCP"), + }; + assert_eq!(dup.local_addr().unwrap(), addr, "复制出来的是同一个套接字"); + } +} diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index c28cb534..860d221e 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -31,6 +31,7 @@ mod contract; pub mod diagnostics; pub mod dryrun; mod gate; +pub mod handover; pub mod key_limits; pub mod keys; pub mod listen; @@ -73,6 +74,8 @@ pub struct ControlState { pub shutdown: Shutdown, /// 远程控制端口此刻在听哪儿、谁被晾着。 pub remote: Arc, + /// 升级不断线：交接走到哪一步、本机控制面的监听（见 [`handover`]） + pub handover: Arc, } impl ControlState { @@ -95,6 +98,7 @@ pub fn router(state: ControlState) -> Router { Router::new() .at(ep::Status, status) .at(ep::Shutdown, ask_shutdown) + .at(ep::HandoverArm, handover::arm) .at(ep::Interfaces, interfaces) .merge(keys::router()) .merge(listen::router()) @@ -232,6 +236,8 @@ async fn status(State(s): State) -> Json { providers: cfg.providers.len(), uptime_secs: s.started.elapsed().as_secs(), in_flight: s.gateway.live.count(), + handover: s.handover.phase(), + draining: s.handover.draining(), remote_control, gateway_reachable, }) @@ -2025,13 +2031,108 @@ pub async fn serve(state: ControlState, at: &tw_api::control::Address) -> Result use tw_api::control::Address; let gate = gate::Gate::new(&state); let app = router(state.clone()); + let handover = state.handover.clone(); tokio::spawn(remote::follow(state, app.clone(), gate.clone())); + // **交接过来的控制面不自己绑**：新 core 在旧 core 交过来的那个上面服务（socket 文件、 + // Windows 的端口文件都不动 —— 桌面端认的就是它们） + if let Some(s) = handover.control().take_inherited() { + return serve_inherited(app, gate, &handover, s, at).await; + } match at { #[cfg(unix)] - Address::Socket(path) => serve_socket(app, gate, path).await, + Address::Socket(path) => serve_socket(app, gate, &handover, path).await, #[cfg(not(unix))] Address::Socket(path) => Err(ControlError::NoUnixSockets { path: path.clone() }), - Address::Loopback { port_file } => serve_loopback(app, gate, port_file).await, + Address::Loopback { port_file } => serve_loopback(app, gate, &handover, port_file).await, + } +} + +/// 在旧 core 交过来的监听上服务。 +// Windows 上只有 TCP 一种，下面的 match 在那儿只有一支 +#[allow(clippy::match_single_binding, clippy::infallible_destructuring_match)] +async fn serve_inherited( + app: Router, + gate: gate::Gate, + handover: &handover::Handover, + s: tw_handover::Socket, + at: &tw_api::control::Address, +) -> Result<(), ControlError> { + let path = match at { + tw_api::control::Address::Socket(p) => p.clone(), + tw_api::control::Address::Loopback { port_file } => port_file.clone(), + }; + let bind_err = |source| ControlError::Bind { + path: path.clone(), + source, + }; + match s { + tw_handover::Socket::Tcp(l) => { + handover + .control() + .hold(tw_handover::Socket::Tcp(l.try_clone().map_err(bind_err)?)); + let listener = tw_gateway::listen::adopt_listener(l).map_err(bind_err)?; + tracing::info!("the control plane is serving on the listener handed over"); + accept_loop(&listener, app, gate, handover).await + } + #[cfg(unix)] + tw_handover::Socket::Unix(l) => { + handover + .control() + .hold(tw_handover::Socket::Unix(l.try_clone().map_err(bind_err)?)); + // 收到的那一头只能设成非阻塞，**绝不改回阻塞**：O_NONBLOCK 是两个进程共享的 + l.set_nonblocking(true).map_err(bind_err)?; + let listener = tokio::net::UnixListener::from_std(l).map_err(bind_err)?; + tracing::info!("the control plane is serving on the socket handed over"); + accept_loop(&listener, app, gate, handover).await + } + } +} + +/// 两种监听器的 accept。 +trait Accept { + type Io: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin + 'static; + fn accept_one(&self) -> impl std::future::Future> + Send; +} + +impl Accept for tokio::net::TcpListener { + type Io = tokio::net::TcpStream; + async fn accept_one(&self) -> std::io::Result { + self.accept().await.map(|(s, _)| s) + } +} + +#[cfg(unix)] +impl Accept for tokio::net::UnixListener { + type Io = tokio::net::UnixStream; + async fn accept_one(&self) -> std::io::Result { + self.accept().await.map(|(s, _)| s) + } +} + +/// accept 循环。**交接时停下**（[`handover::ControlListener::pause`]）：停的时候不去接， +/// 监听器本身留着，恢复了接着接 —— 见 [`handover`] 的模块文档。 +async fn accept_loop( + listener: &L, + app: Router, + gate: gate::Gate, + handover: &handover::Handover, +) -> Result<(), ControlError> { + let mut accepting = handover.control().accepting(); + loop { + handover::until_accepting(&mut accepting).await; + tokio::select! { + r = listener.accept_one() => match r { + Ok(stream) => hand_off( + stream, + app.clone(), + gate.clone(), + Local(Some(handover.control().closing())), + ), + Err(e) => tracing::warn!("the control plane could not accept a connection: {e}"), + }, + // 停下的那一刻放下手里这次 accept：不然它还会再接一个 + _ = accepting.wait_for(|on| !*on) => {} + } } } @@ -2040,7 +2141,12 @@ pub async fn serve(state: ControlState, at: &tw_api::control::Address) -> Result /// 陈旧的 socket 文件直接删掉重建 —— 它和 lock 文件不一样，没有「另一个 /// 实例可能还在用」的歧义：单实例锁已经在上一步挡住了。 #[cfg(unix)] -async fn serve_socket(app: Router, gate: gate::Gate, path: &Path) -> Result<(), ControlError> { +async fn serve_socket( + app: Router, + gate: gate::Gate, + handover: &handover::Handover, + path: &Path, +) -> Result<(), ControlError> { socket_path_fits(path)?; if path.exists() { let _ = std::fs::remove_file(path); @@ -2048,23 +2154,26 @@ async fn serve_socket(app: Router, gate: gate::Gate, path: &Path) -> Result<(), if let Some(dir) = path.parent() { let _ = std::fs::create_dir_all(dir); } - let listener = tokio::net::UnixListener::bind(path).map_err(|source| ControlError::Bind { + let bind_err = |source| ControlError::Bind { path: path.to_path_buf(), source, - })?; + }; + // **先建 std 的监听，复制一份给 tokio，原件留着交接用**（同一个内核套接字） + let std_listener = std::os::unix::net::UnixListener::bind(path).map_err(bind_err)?; { use std::os::unix::fs::PermissionsExt; // 0700：只有当前用户能连。握手那道门在它之外，不是替代它 —— // 两道都在，而只有这一道是平台给的。 let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700)); } + let mine = std_listener.try_clone().map_err(bind_err)?; + mine.set_nonblocking(true).map_err(bind_err)?; + let listener = tokio::net::UnixListener::from_std(mine).map_err(bind_err)?; + handover + .control() + .hold(tw_handover::Socket::Unix(std_listener)); tracing::info!(path = %path.display(), "the control plane is listening"); - loop { - match listener.accept().await { - Ok((stream, _)) => hand_off(stream, app.clone(), gate.clone(), None), - Err(e) => tracing::warn!("the control plane could not accept a connection: {e}"), - } - } + accept_loop(&listener, app, gate, handover).await } /// 回环 TCP 上的控制面。Windows 只有这一档。 @@ -2080,36 +2189,30 @@ async fn serve_socket(app: Router, gate: gate::Gate, path: &Path) -> Result<(), async fn serve_loopback( app: Router, gate: gate::Gate, + handover: &handover::Handover, port_file: &Path, ) -> Result<(), ControlError> { use std::net::Ipv4Addr; if let Some(dir) = port_file.parent() { let _ = std::fs::create_dir_all(dir); } - let listener = tokio::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) - .await - .map_err(|source| ControlError::Bind { - path: port_file.to_path_buf(), - source, - })?; - let port = listener - .local_addr() - .map_err(|source| ControlError::Bind { - path: port_file.to_path_buf(), - source, - })? - .port(); + let bind_err = |source| ControlError::Bind { + path: port_file.to_path_buf(), + source, + }; + let std_listener = std::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).map_err(bind_err)?; + let port = std_listener.local_addr().map_err(bind_err)?.port(); + let mine = std_listener.try_clone().map_err(bind_err)?; + let listener = tw_gateway::listen::adopt_listener(mine).map_err(bind_err)?; std::fs::write(port_file, port.to_string()).map_err(|source| ControlError::PortFile { path: port_file.to_path_buf(), source, })?; + handover + .control() + .hold(tw_handover::Socket::Tcp(std_listener)); tracing::info!(port, "the control plane is listening on loopback"); - loop { - match listener.accept().await { - Ok((stream, _)) => hand_off(stream, app.clone(), gate.clone(), None), - Err(e) => tracing::warn!("the control plane could not accept a connection: {e}"), - } - } + accept_loop(&listener, app, gate, handover).await } /// 一条连接：先握手，握上了交给 hyper。 @@ -2117,10 +2220,14 @@ async fn serve_loopback( /// **每种传输共用**：它们的差别只在怎么拿到这个流，拿到之后的每一件事 /// （握手、协议协商、错误怎么记）都该一模一样 —— 写两遍就是两遍会漂。 /// 远程端口接进来的连接也走这里，多带一个 [`RemoteConn`]。 -fn hand_off(stream: S, app: Router, gate: gate::Gate, remote: Option) +fn hand_off(stream: S, app: Router, gate: gate::Gate, remote: impl Into) where S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin + 'static, { + let (remote, mut closing) = match remote.into() { + Conn::Remote(r) => (Some(r), None), + Conn::Local(c) => (None, c), + }; tokio::spawn(async move { let link = match gate.admit(stream).await { Ok(l) => l, @@ -2154,6 +2261,10 @@ where _ = gate.key_changed_from(&link.key) => { tracing::info!("the control key changed; closing a connection made with the previous one"); } + // 控制面交给了接替者：已经连着的这一条也断开（见 `handover::ControlListener::close_connections`） + _ = handover::closed(&mut closing) => { + tracing::info!("the control plane was handed over; closing a connection to this core"); + } } return; }; @@ -2185,6 +2296,28 @@ where }); } +/// 一条连接是从哪儿进来的。 +pub(crate) enum Conn { + Remote(RemoteConn), + /// 本机的通道。带着的是「控制面交出去了」的信号（测试里没有） + Local(Option>), +} + +/// 本机的连接 +pub(crate) struct Local(pub(crate) Option>); + +impl From for Conn { + fn from(l: Local) -> Self { + Conn::Local(l.0) + } +} + +impl From for Conn { + fn from(r: RemoteConn) -> Self { + Conn::Remote(r) + } +} + /// 从远程端口进来的一条连接多带的东西。 pub(crate) struct RemoteConn { /// 握手没成：记一笔（节流） diff --git a/crates/tw-control/src/remote.rs b/crates/tw-control/src/remote.rs index a45b7bde..10fdf37b 100644 --- a/crates/tw-control/src/remote.rs +++ b/crates/tw-control/src/remote.rs @@ -101,6 +101,12 @@ pub struct Remote { slots: Arc, /// 每换入一次配置跳一下。已经连着的远程连接听它，名单把自己划出去了就断开 reloads: watch::Sender, + /// 交接时跟不跟配置走（和网关的监听同一个规矩，见 `tw_gateway::listen::Handoff`） + mode: watch::Sender, + /// 此刻在听的那个的复制句柄，交接时交出去。**弱引用**：攥着它，关掉的端口就放不开 + held: Mutex)>>, + /// 旧 core 交过来、还没起服务的那个 + adopted: Mutex>, } impl Default for Remote { @@ -110,6 +116,60 @@ impl Default for Remote { strikes: Mutex::default(), slots: Arc::new(Semaphore::new(MAX_CONNECTIONS)), reloads: watch::channel(0).0, + mode: watch::channel(Mode::Follow).0, + held: Mutex::default(), + adopted: Mutex::default(), + } + } +} + +/// 交接时远程端口听谁的。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Mode { + Follow, + /// 新 core 提交之前：只在交过来的那个上面 accept，不跟配置换 + Hold, + /// 旧 core 提交之后：停下，从它进来的连接一起断开（远端按断线重连连到新 core） + Stop, +} + +impl Remote { + /// 先别跟配置开关远程端口（新 core 提交之前）。**在控制面起来之前调** + pub fn hold(&self) { + self.mode.send_replace(Mode::Hold); + } + + /// 旧 core 交过来的远程端口：马上在它上面 accept + pub fn adopt(&self, want: SocketAddr, l: std::net::TcpListener) { + if let Ok(mut g) = self.adopted.lock() { + *g = Some((want, l)); + } + // 叫醒跟着的那一路：值不变也要叫 + self.mode.send_modify(|_| {}); + } + + /// 从现在起跟着配置走（新 core 提交之后）。交过来的地址和配置不一样就换过去 + pub fn follow(&self) { + self.mode.send_replace(Mode::Follow); + } + + /// 停下（旧 core 提交之后）：不再 accept，从它进来的连接一起断开。**只丢自己的句柄** + pub fn stop_accepting(&self) { + self.mode.send_replace(Mode::Stop); + } + + /// 此刻在听的那个：配置算出来的地址，和一个复制出来的句柄 + pub fn export(&self) -> std::io::Result> { + let held = self.held.lock().ok().and_then(|g| g.clone()); + match held.and_then(|(want, l)| l.upgrade().map(|l| (want, l))) { + Some((want, l)) => Ok(Some((want, l.try_clone()?))), + None => Ok(None), + } + } + + fn set_held(&self, run: Option<&Running>) { + if let Ok(mut g) = self.held.lock() { + *g = run.map(|r| (r.want, Arc::downgrade(&r.dup))); } } } @@ -196,6 +256,8 @@ pub fn reachable(addr: SocketAddr) -> Vec { /// 进来的连接也一起断开**：关掉远程端口的意思就是现在谁都不能从远程进来。 struct Running { want: SocketAddr, + /// 复制出来的句柄，交接时交出去 + dup: Arc, stop: watch::Sender<()>, /// accept 循环。**它结束了，监听的 socket 才真的放开** —— 同一个端口换地址 /// 时要等到这一刻才能绑新的 @@ -218,27 +280,85 @@ impl Running { pub(crate) async fn follow(state: ControlState, app: Router, gate: Gate) { use tokio::sync::broadcast::error::RecvError; let mut events = state.bus().subscribe(); + let mut mode = state.remote.mode.subscribe(); let mut running: Option = None; loop { - apply(&state, &app, &gate, &mut running).await; - // 已经连着的远程连接各自对一遍名单（见 `revoked`） - state.remote.reloads.send_modify(|n| *n += 1); + let now = *mode.borrow_and_update(); + if now == Mode::Stop { + state.remote.set_held(None); + if let Some(run) = running.take() { + tracing::info!("handed over; the remote control port no longer accepts here"); + run.close().await; + } + return; + } + adopt(&state, &app, &gate, &mut running).await; + if now == Mode::Follow { + apply(&state, &app, &gate, &mut running).await; + state.remote.set_held(running.as_ref()); + // 已经连着的远程连接各自对一遍名单（见 `revoked`） + state.remote.reloads.send_modify(|n| *n += 1); + } loop { - match events.recv().await { - Ok(tw_api::Event::ConfigReloaded { .. }) | Err(RecvError::Lagged(_)) => break, - Ok(_) => {} - Err(RecvError::Closed) => return, + tokio::select! { + ev = events.recv() => match ev { + // 提交之前不跟配置走（见 `Mode::Hold`） + Ok(tw_api::Event::ConfigReloaded { .. }) | Err(RecvError::Lagged(_)) + if now == Mode::Follow => break, + Ok(_) | Err(RecvError::Lagged(_)) => {} + Err(RecvError::Closed) => return, + }, + r = mode.changed() => { + if r.is_err() { + return; + } + break; + } } } } } +/// 旧 core 交过来的那个：在它上面起服务（替掉手上的，提交之前手上本来就没有）。 +async fn adopt(state: &ControlState, app: &Router, gate: &Gate, running: &mut Option) { + let Some((want, l)) = state.remote.adopted.lock().ok().and_then(|mut g| g.take()) else { + return; + }; + let listener = match tw_gateway::listen::adopt_listener(l) { + Ok(l) => l, + Err(e) => { + tracing::warn!(%want, %e, "the remote control port handed over could not be used"); + return; + } + }; + let actual = listener.local_addr().ok(); + state.remote.set_held(None); + if let Some(old) = running.take() { + old.close().await; + } + match start(state, app, gate, want, listener) { + Ok(run) => { + *running = Some(run); + state.remote.set_held(running.as_ref()); + state.remote.set_listening(Listening { + addr: actual, + error: None, + }); + tracing::info!(addr = ?actual, "the remote control port is serving on the listener handed over"); + } + Err(e) => { + tracing::warn!(%want, %e, "the remote control port handed over could not be used") + } + } +} + async fn apply(state: &ControlState, app: &Router, gate: &Gate, running: &mut Option) { let r = &state.remote; let cfg = state.config(); let want = match &cfg.listen.control.remote { Some(x) if x.enabled => x.clone(), _ => { + r.set_held(None); if running.take().is_some() { tracing::info!("the remote control port is closed"); } @@ -279,6 +399,7 @@ async fn apply(state: &ControlState, app: &Router, gate: &Gate, running: &mut Op { let old = running.take().expect("checked just above"); let old_want = old.want; + r.set_held(None); old.close().await; match TcpListener::bind(addr).await { Ok(l) => l, @@ -288,7 +409,7 @@ async fn apply(state: &ControlState, app: &Router, gate: &Gate, running: &mut Op match TcpListener::bind(old_want).await { Ok(l) => { let actual = l.local_addr().ok(); - *running = Some(start(state, app, gate, old_want, l)); + *running = start(state, app, gate, old_want, l).ok(); r.set_listening(Listening { addr: actual, error, @@ -314,10 +435,21 @@ async fn apply(state: &ControlState, app: &Router, gate: &Gate, running: &mut Op }; let actual = listener.local_addr().ok(); // 旧的那个关掉：不再接新连接，从它进来的也断开 + r.set_held(None); if let Some(old) = running.take() { old.close().await; } - *running = Some(start(state, app, gate, addr, listener)); + *running = match start(state, app, gate, addr, listener) { + Ok(run) => Some(run), + Err(e) => { + tracing::warn!(%addr, %e, "the remote control port cannot be listened on"); + r.set_listening(Listening { + addr: None, + error: Some(tw_gateway::listen::bind_failure(addr, &e)), + }); + return; + } + }; r.set_listening(Listening { addr: actual, error: None, @@ -331,7 +463,8 @@ fn start( gate: &Gate, want: SocketAddr, listener: TcpListener, -) -> Running { +) -> std::io::Result { + let dup = Arc::new(tw_gateway::listen::dup_listener(&listener)?); let (stop, stopped) = watch::channel(()); let task = tokio::spawn(accept_loop( listener, @@ -340,7 +473,12 @@ fn start( app.clone(), gate.clone(), )); - Running { want, stop, task } + Ok(Running { + want, + dup, + stop, + task, + }) } /// 等到这个来源不再被放行：名单改了把它划出去，或者远程端口关了。 @@ -414,12 +552,12 @@ async fn accept_loop( stream, app.clone(), gate.clone(), - Some(crate::RemoteConn { + crate::RemoteConn { on_failure: Box::new(move || remote.strike(ip, Instant::now())), closed: stopped.clone(), revoked: Box::pin(revoked(state.clone(), reloads, ip)), _slot: slot, - }), + }, ); } } diff --git a/crates/tw-control/src/rotation.rs b/crates/tw-control/src/rotation.rs index 3015165e..30463b7d 100644 --- a/crates/tw-control/src/rotation.rs +++ b/crates/tw-control/src/rotation.rs @@ -44,12 +44,69 @@ pub const CHANNEL_CAP: usize = 8; const TRIES: usize = 3; const GAP: Duration = Duration::from_millis(400); +/// 写回那一路的把手：交接之前等它把手上的写完（升级不断线）。 +pub struct Rotation { + /// 看通道里还排着几份（不拿它发东西） + tx: tokio::sync::mpsc::Sender, + busy: std::sync::Arc, +} + +impl Rotation { + /// 等到手上没有要写回的：排着的写完、正在写的写完。最多等 `within`。 + /// + /// **交接之前调**：新 core 读的是配置文件里的 refresh token，旧 core 刚换到、还没写回 + /// 的那一份要是丢了，新 core 手里的就是一张已经作废的。 + pub async fn settle(&self, within: Duration) { + use std::sync::atomic::Ordering; + // **连着两次看到闲着才算**：取出一份和记下「在忙」之间有一瞬间两边都看不见它 + let quiet = async { + let mut calm = 0; + loop { + let queued = self.tx.max_capacity() - self.tx.capacity(); + if queued == 0 && !self.busy.load(Ordering::SeqCst) { + calm += 1; + if calm >= 2 { + return; + } + } else { + calm = 0; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + }; + if tokio::time::timeout(within, quiet).await.is_err() { + tracing::warn!("a token write-back was still in progress when the handover went ahead"); + } + } +} + /// 起一个任务收刷新结果并写回文件。**控制面起来之后调一次。** -pub fn spawn(state: crate::ControlState) { +pub fn spawn(state: crate::ControlState) -> Rotation { let (tx, mut rx) = tokio::sync::mpsc::channel::(CHANNEL_CAP); - state.gateway.set_renewal_sink(tx); + state.gateway.set_renewal_sink(tx.clone()); + let busy = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let handle = Rotation { + tx, + busy: busy.clone(), + }; tokio::spawn(async move { - while let Some(r) = rx.recv().await { + use std::sync::atomic::Ordering; + loop { + // **先记「在忙」再取**：取出来到写完之间，`settle` 要看得见这一份 + busy.store(true, Ordering::SeqCst); + let r = match rx.try_recv() { + Ok(r) => r, + Err(_) => { + busy.store(false, Ordering::SeqCst); + match rx.recv().await { + Some(r) => { + busy.store(true, Ordering::SeqCst); + r + } + None => return, + } + } + }; let (ok, detail) = persist(&state, &r).await; if r.refresh.is_some() { // 换发了 refresh token：成没成都要报（「成功」只报一次，由网关那边管） @@ -59,6 +116,7 @@ pub fn spawn(state: crate::ControlState) { } } }); + handle } /// 真正写那一次。返回（成没成，那句话）。 diff --git a/crates/tw-control/src/zai.rs b/crates/tw-control/src/zai.rs index e3c1ef94..d3a75c3c 100644 --- a/crates/tw-control/src/zai.rs +++ b/crates/tw-control/src/zai.rs @@ -155,6 +155,14 @@ impl Accounts { } } + /// 有一次登录正在等用户授权。**这时不交接**（升级不断线）：登录要在这个 core 上走完 + pub fn in_progress(&self) -> bool { + self.current.lock().is_ok_and(|g| { + g.as_ref() + .is_some_and(|c| c.status.status == tw_api::LoginStatus::Pending) + }) + } + fn set_status(&self, status: tw_api::ZaiLoginStatus) { if let Ok(mut g) = self.current.lock() && let Some(c) = g.as_mut() @@ -257,6 +265,10 @@ async fn start( State(s): State, Json(req): Json, ) -> Result, Fail> { + // arm 之后不开始新的登录：回调端口不交接，登录要走完才能交（见 `handover`） + if crate::handover::refuses_logins(&s.handover) { + return Err(crate::handover::login_refused()); + } let want = Want::read(&s, &req)?; // 还没完成的上一次登录作废 let prev = s.zai.current.lock().ok().and_then(|mut g| g.take()); diff --git a/crates/tw-control/tests/aliases.rs b/crates/tw-control/tests/aliases.rs index f7b44831..889f520a 100644 --- a/crates/tw-control/tests/aliases.rs +++ b/crates/tw-control/tests/aliases.rs @@ -156,6 +156,7 @@ fn bed(yaml: &str, rows: Option<&[tw_store::db::RequestRow]>) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store, diff --git a/crates/tw-control/tests/backups.rs b/crates/tw-control/tests/backups.rs index 39f480f4..20a26811 100644 --- a/crates/tw-control/tests/backups.rs +++ b/crates/tw-control/tests/backups.rs @@ -84,6 +84,7 @@ fn state_for( ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: mgr, gateway: gw, store: None, diff --git a/crates/tw-control/tests/bedrock.rs b/crates/tw-control/tests/bedrock.rs index 4fded445..fb6e8de7 100644 --- a/crates/tw-control/tests/bedrock.rs +++ b/crates/tw-control/tests/bedrock.rs @@ -43,6 +43,7 @@ fn bed(yaml: &str) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/bundle.rs b/crates/tw-control/tests/bundle.rs index cace6b0d..baa2db50 100644 --- a/crates/tw-control/tests/bundle.rs +++ b/crates/tw-control/tests/bundle.rs @@ -29,6 +29,7 @@ fn bed_with_store(store: Option, endpoints: Endpoints) -> Bed let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/client_key.rs b/crates/tw-control/tests/client_key.rs index a21ece38..a8492497 100644 --- a/crates/tw-control/tests/client_key.rs +++ b/crates/tw-control/tests/client_key.rs @@ -27,6 +27,7 @@ fn bed() -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store, diff --git a/crates/tw-control/tests/contract.rs b/crates/tw-control/tests/contract.rs index 45a5228c..567fc363 100644 --- a/crates/tw-control/tests/contract.rs +++ b/crates/tw-control/tests/contract.rs @@ -24,6 +24,7 @@ fn app() -> (tempfile::TempDir, axum::Router) { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/dryrun.rs b/crates/tw-control/tests/dryrun.rs index ee8ba232..6bb56c2d 100644 --- a/crates/tw-control/tests/dryrun.rs +++ b/crates/tw-control/tests/dryrun.rs @@ -66,6 +66,7 @@ fn app_and_gateway(text: &str) -> (tempfile::TempDir, axum::Router, tw_gateway:: let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/error_answers.rs b/crates/tw-control/tests/error_answers.rs index 04cd465d..b965c04b 100644 --- a/crates/tw-control/tests/error_answers.rs +++ b/crates/tw-control/tests/error_answers.rs @@ -140,6 +140,7 @@ async fn world() -> World { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), gw.bus.clone())), gateway: gw.clone(), store: Some(store.clone()), diff --git a/crates/tw-control/tests/gate.rs b/crates/tw-control/tests/gate.rs index 71c45f50..ec7258dc 100644 --- a/crates/tw-control/tests/gate.rs +++ b/crates/tw-control/tests/gate.rs @@ -50,6 +50,7 @@ fn bed() -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, @@ -254,6 +255,45 @@ async fn a_new_key_closes_the_connections_made_with_the_old_one() { assert_eq!(get(&mut s, "/status").await.0, StatusCode::OK); } +/// 交出控制面（升级不断线）：已经连着这个 core 的本机连接断开、不再接新的；交接没成、恢复 +/// accept 之后又连得上。连着的多半是新桌面端交接之前连上来的，断开它才会重连到接替者 +#[tokio::test] +async fn handing_over_the_control_plane_closes_the_connections_to_this_core() { + let b = bed(); + let at = Address::Loopback { + port_file: b.dir.path().join("control.port"), + }; + listen(&b, at.clone()).await; + let (mut s, conn) = open(&at, &key(KEY)).await.unwrap(); + assert_eq!(get(&mut s, "/status").await.0, StatusCode::OK); + + b.state.handover.control().pause(); + b.state.handover.control().close_connections(); + tokio::time::timeout(Duration::from_secs(5), conn) + .await + .expect("交出控制面之后，连着的那一条该被断开") + .unwrap(); + // 停着 accept：新连接握不上手（没人接） + let port: u16 = std::fs::read_to_string(b.dir.path().join("control.port")) + .unwrap() + .trim() + .parse() + .unwrap(); + let stalled = tokio::time::timeout(Duration::from_millis(300), async { + let tcp = tokio::net::TcpStream::connect(("127.0.0.1", port)) + .await + .unwrap(); + tw_link::connect(tcp, &key(KEY), "gate test").await + }) + .await; + assert!(stalled.is_err(), "停了 accept 还握得上手"); + + // 交接没成：接着接 + b.state.handover.control().resume(); + let (mut s, _) = open(&at, &key(KEY)).await.unwrap(); + assert_eq!(get(&mut s, "/status").await.0, StatusCode::OK); +} + /// 配置里的钥匙写坏了（外部改动）：配置不收，旧钥匙继续管用。 #[tokio::test] async fn a_broken_key_in_the_file_is_rejected_and_the_old_one_keeps_working() { diff --git a/crates/tw-control/tests/handover.rs b/crates/tw-control/tests/handover.rs new file mode 100644 index 00000000..0763a0df --- /dev/null +++ b/crates/tw-control/tests/handover.rs @@ -0,0 +1,181 @@ +//! `POST /handover/arm` 与 `Status.handover` / `draining`（升级不断线）。 +//! +//! 交接本身在 `twcore` 里；这里盯控制面这一边说的话：哪些情形拒绝、拒绝时的码（桌面端 +//! 据此退回现行流程），答应时交给 `twcore` 的参数，以及 `/status` 报的两个数。 + +use std::sync::Arc; +use std::time::Duration; + +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use tower::ServiceExt; +use tw_control::{ConfigManager, ControlState}; + +const BASE: &str = "version: 1\nlisten:\n control:\n key: c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00\nclients:\n - name: 我\n key: tw-一把钥匙就够\nproviders: []\n"; + +fn state() -> (tempfile::TempDir, ControlState) { + let d = tempfile::tempdir().unwrap(); + let p = d.path().join("config.yaml"); + std::fs::write(&p, BASE).unwrap(); + let cfg: tw_config::Config = serde_yaml_ng::from_str(BASE).unwrap(); + let gw = tw_gateway::AppState::new(cfg).unwrap(); + let bus = gw.bus.clone(); + let state = ControlState { + shutdown: Default::default(), + remote: Default::default(), + handover: Default::default(), + cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), + gateway: gw, + store: None, + started: std::time::Instant::now(), + price_updater: Default::default(), + chatgpt: Default::default(), + zai: Default::default(), + }; + (d, state) +} + +async fn arm(s: &ControlState, body: &str) -> (StatusCode, serde_json::Value) { + let resp = tw_control::router(s.clone()) + .oneshot( + Request::post("/handover/arm") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = resp.status(); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + (status, serde_json::from_slice(&body).unwrap()) +} + +async fn status(s: &ControlState) -> tw_api::Status { + let resp = tw_control::router(s.clone()) + .oneshot(Request::get("/status").body(Body::empty()).unwrap()) + .await + .unwrap(); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + serde_json::from_slice(&body).unwrap() +} + +#[tokio::test] +async fn arming_hands_the_request_to_the_core_and_answers_with_the_handover_version() { + let (_d, s) = state(); + let mut arms = s.handover.arm_requests(); + let core = tokio::spawn(async move { + let a = arms.recv().await.unwrap(); + let got = (a.wait, a.relaunch.clone()); + a.reply.send(Ok(())).unwrap(); + got + }); + let (code, body) = arm( + &s, + r#"{"wait_secs":null,"relaunch":{"kind":"bundle_id","id":"app.thinkwatch.lite"}}"#, + ) + .await; + assert_eq!(code, StatusCode::OK, "{body}"); + let armed: tw_api::HandoverArmed = serde_json::from_value(body).unwrap(); + assert_eq!(armed.version, tw_handover::VERSION); + let (wait, relaunch) = core.await.unwrap(); + assert_eq!(wait, Duration::from_secs(120), "缺省等 120 秒"); + assert_eq!( + relaunch, + Some(tw_api::Relaunch::BundleId { + id: "app.thinkwatch.lite".into() + }) + ); +} + +#[tokio::test] +async fn a_core_without_a_data_plane_refuses_with_safe_mode() { + let (_d, s) = state(); + let _arms = s.handover.arm_requests(); + s.handover.set_no_data_plane(true); + let (code, body) = arm(&s, r#"{"wait_secs":5,"relaunch":null}"#).await; + assert_eq!(code, StatusCode::CONFLICT); + assert_eq!(body["code"], "control.handover.safe_mode", "{body}"); +} + +#[tokio::test] +async fn a_core_that_does_not_run_handovers_refuses_rather_than_pretending() { + // 没有人接 arm（比如测试里、或者这个进程不跑交接）：说清楚，不假装答应了 + let (_d, s) = state(); + let (code, body) = arm(&s, r#"{"wait_secs":5,"relaunch":null}"#).await; + assert_eq!(code, StatusCode::CONFLICT); + assert_eq!(body["code"], "control.handover.safe_mode", "{body}"); +} + +#[tokio::test] +async fn a_second_handover_while_one_is_taking_over_is_busy() { + let (_d, s) = state(); + let _arms = s.handover.arm_requests(); + s.handover.set_phase(tw_api::HandoverState::TakingOver); + let (code, body) = arm(&s, r#"{"wait_secs":5,"relaunch":null}"#).await; + assert_eq!(code, StatusCode::CONFLICT); + assert_eq!(body["code"], "control.handover.busy", "{body}"); +} + +#[tokio::test] +async fn an_endpoint_that_cannot_open_is_reported_with_its_reason() { + let (_d, s) = state(); + let mut arms = s.handover.arm_requests(); + tokio::spawn(async move { + let a = arms.recv().await.unwrap(); + let _ = a + .reply + .send(Err(tw_control::handover::endpoint_failed("disk full"))); + }); + let (code, body) = arm(&s, r#"{"wait_secs":5,"relaunch":null}"#).await; + assert_eq!(code, StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(body["code"], "control.handover.endpoint_failed", "{body}"); +} + +#[tokio::test] +async fn status_reports_the_handover_phase_and_the_draining_count() { + let (_d, s) = state(); + let st = status(&s).await; + assert_eq!(st.handover, tw_api::HandoverState::Idle); + assert_eq!(st.draining, 0); + s.handover.set_phase(tw_api::HandoverState::Draining); + s.handover.set_draining(3); + let st = status(&s).await; + assert_eq!(st.handover, tw_api::HandoverState::Draining); + assert_eq!(st.draining, 3); + assert_eq!(st.api_version, 52); +} + +/// arm 之后（以及接替者正在接管时）不开始新的账号登录：回调端口不交接。（平常、排空中的 +/// 新 core 照常能登：见 `handover::tests` 里那条，这里不真的去登） +#[tokio::test] +async fn a_sign_in_that_starts_after_arming_is_refused() { + let (_d, s) = state(); + for phase in [ + tw_api::HandoverState::Armed, + tw_api::HandoverState::TakingOver, + ] { + s.handover.set_phase(phase); + for path in ["/chatgpt/login", "/zai/login"] { + let resp = tw_control::router(s.clone()) + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .body(Body::from("{}")) + .unwrap(), + ) + .await + .unwrap(); + let status = resp.status(); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap_or_default(); + assert_eq!(status, StatusCode::CONFLICT, "{path} {phase:?}: {body}"); + assert_eq!(body["code"], "control.handover.login_refused", "{body}"); + } + } +} diff --git a/crates/tw-control/tests/history_search.rs b/crates/tw-control/tests/history_search.rs index 026617e2..bd3e4bee 100644 --- a/crates/tw-control/tests/history_search.rs +++ b/crates/tw-control/tests/history_search.rs @@ -105,6 +105,7 @@ fn app( let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: Some(Arc::new(tokio::sync::Mutex::new(rec))), @@ -226,6 +227,7 @@ async fn without_a_store_the_search_says_recording_is_unavailable() { let app = tw_control::router(ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/in_flight.rs b/crates/tw-control/tests/in_flight.rs index 498d3388..97c60be6 100644 --- a/crates/tw-control/tests/in_flight.rs +++ b/crates/tw-control/tests/in_flight.rs @@ -23,6 +23,7 @@ fn app() -> (tempfile::TempDir, tw_observe::EventBus, axum::Router) { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus.clone())), gateway: gw, store: None, diff --git a/crates/tw-control/tests/key_limits.rs b/crates/tw-control/tests/key_limits.rs index febbdf1b..8adb609f 100644 --- a/crates/tw-control/tests/key_limits.rs +++ b/crates/tw-control/tests/key_limits.rs @@ -84,6 +84,7 @@ fn bed_with(db: Option, clock: Arc Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/listen.rs b/crates/tw-control/tests/listen.rs index 68524f97..86078d81 100644 --- a/crates/tw-control/tests/listen.rs +++ b/crates/tw-control/tests/listen.rs @@ -39,6 +39,7 @@ fn bed(yaml: &str) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw.clone(), store: None, diff --git a/crates/tw-control/tests/live_content.rs b/crates/tw-control/tests/live_content.rs index 77f12aba..a48b32d8 100644 --- a/crates/tw-control/tests/live_content.rs +++ b/crates/tw-control/tests/live_content.rs @@ -202,6 +202,7 @@ async fn world(providers: &str, mode: &str) -> World { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), gw.bus.clone())), gateway: gw.clone(), store: Some(store.clone()), diff --git a/crates/tw-control/tests/live_state.rs b/crates/tw-control/tests/live_state.rs index 1096f958..0cb6f4b0 100644 --- a/crates/tw-control/tests/live_state.rs +++ b/crates/tw-control/tests/live_state.rs @@ -29,6 +29,7 @@ fn control( let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw.clone(), store, diff --git a/crates/tw-control/tests/manual_models.rs b/crates/tw-control/tests/manual_models.rs index d5d3dbe4..27dda919 100644 --- a/crates/tw-control/tests/manual_models.rs +++ b/crates/tw-control/tests/manual_models.rs @@ -31,6 +31,7 @@ fn bed(yaml: &str) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/model_specs.rs b/crates/tw-control/tests/model_specs.rs index 520f8905..6f87d0bb 100644 --- a/crates/tw-control/tests/model_specs.rs +++ b/crates/tw-control/tests/model_specs.rs @@ -31,6 +31,7 @@ fn bed(yaml: &str) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/models.rs b/crates/tw-control/tests/models.rs index 327c6ab4..b0bea6e7 100644 --- a/crates/tw-control/tests/models.rs +++ b/crates/tw-control/tests/models.rs @@ -32,6 +32,7 @@ fn bed(yaml: &str) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/plugin_defaults.rs b/crates/tw-control/tests/plugin_defaults.rs index 83e2bf37..a0505cef 100644 --- a/crates/tw-control/tests/plugin_defaults.rs +++ b/crates/tw-control/tests/plugin_defaults.rs @@ -170,6 +170,7 @@ fn open(dir: PathBuf, engine: Arc) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: mgr.clone(), gateway: gw.clone(), store: None, diff --git a/crates/tw-control/tests/plugins.rs b/crates/tw-control/tests/plugins.rs index e9ea6a69..af8c44d5 100644 --- a/crates/tw-control/tests/plugins.rs +++ b/crates/tw-control/tests/plugins.rs @@ -160,6 +160,7 @@ fn bed_with(sub: &str, fake: bool) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw.clone(), store: Some(store.clone()), diff --git a/crates/tw-control/tests/pricing.rs b/crates/tw-control/tests/pricing.rs index b8312987..2f37b1a7 100644 --- a/crates/tw-control/tests/pricing.rs +++ b/crates/tw-control/tests/pricing.rs @@ -68,6 +68,7 @@ fn bed_with(yaml: &str, updater: Updater) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/remote.rs b/crates/tw-control/tests/remote.rs index 6d63b383..6a91fba5 100644 --- a/crates/tw-control/tests/remote.rs +++ b/crates/tw-control/tests/remote.rs @@ -73,6 +73,7 @@ async fn bed(text: String) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/replay.rs b/crates/tw-control/tests/replay.rs index 1d55c2f1..34030445 100644 --- a/crates/tw-control/tests/replay.rs +++ b/crates/tw-control/tests/replay.rs @@ -127,6 +127,7 @@ fn app_of( let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: Some(Arc::new(tokio::sync::Mutex::new(rec))), diff --git a/crates/tw-control/tests/resources.rs b/crates/tw-control/tests/resources.rs index 78d624a6..8b078171 100644 --- a/crates/tw-control/tests/resources.rs +++ b/crates/tw-control/tests/resources.rs @@ -62,6 +62,7 @@ fn bed(yaml: &str) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw.clone(), store: None, diff --git a/crates/tw-control/tests/route_stats.rs b/crates/tw-control/tests/route_stats.rs index be859251..2671bf20 100644 --- a/crates/tw-control/tests/route_stats.rs +++ b/crates/tw-control/tests/route_stats.rs @@ -107,6 +107,7 @@ async fn world(text: &str) -> World { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), gw.bus.clone())), gateway: gw.clone(), store: Some(store.clone()), @@ -289,6 +290,7 @@ async fn a_store_with_no_history_yet_says_nothing_is_covered() { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), gw.bus.clone())), gateway: gw, store: Some(store), @@ -335,6 +337,7 @@ async fn without_a_store_it_says_so() { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), gw.bus.clone())), gateway: gw, store: None, diff --git a/crates/tw-control/tests/routes.rs b/crates/tw-control/tests/routes.rs index f2b22480..21e40fb5 100644 --- a/crates/tw-control/tests/routes.rs +++ b/crates/tw-control/tests/routes.rs @@ -87,6 +87,7 @@ fn bed(yaml: &str) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/safe_mode.rs b/crates/tw-control/tests/safe_mode.rs index fb2fd5af..8b6215ec 100644 --- a/crates/tw-control/tests/safe_mode.rs +++ b/crates/tw-control/tests/safe_mode.rs @@ -42,6 +42,7 @@ fn safe_mode(yaml: &str) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: mgr.clone(), gateway: gw.clone(), store: None, diff --git a/crates/tw-control/tests/security.rs b/crates/tw-control/tests/security.rs index d1957557..61e30b60 100644 --- a/crates/tw-control/tests/security.rs +++ b/crates/tw-control/tests/security.rs @@ -128,6 +128,7 @@ fn bed_with(yaml: &str, seed: impl FnOnce(&tw_store::Db)) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: Some(Arc::new(tokio::sync::Mutex::new(rec))), diff --git a/crates/tw-control/tests/security_detail.rs b/crates/tw-control/tests/security_detail.rs index 956bb796..3070d75f 100644 --- a/crates/tw-control/tests/security_detail.rs +++ b/crates/tw-control/tests/security_detail.rs @@ -121,6 +121,7 @@ security: let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), gw.bus.clone())), gateway: gw.clone(), store: Some(store.clone()), diff --git a/crates/tw-control/tests/sessions.rs b/crates/tw-control/tests/sessions.rs index f17e747a..3b5b3b1f 100644 --- a/crates/tw-control/tests/sessions.rs +++ b/crates/tw-control/tests/sessions.rs @@ -74,6 +74,7 @@ fn app(rows: &[tw_store::db::RequestRow]) -> (tempfile::TempDir, axum::Router) { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: Some(Arc::new(tokio::sync::Mutex::new(rec))), diff --git a/crates/tw-control/tests/speed.rs b/crates/tw-control/tests/speed.rs index 437df2fb..dbd38787 100644 --- a/crates/tw-control/tests/speed.rs +++ b/crates/tw-control/tests/speed.rs @@ -26,6 +26,7 @@ fn bed(yaml: &str) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/stored_bodies.rs b/crates/tw-control/tests/stored_bodies.rs index 3f3f1823..fdcbe25d 100644 --- a/crates/tw-control/tests/stored_bodies.rs +++ b/crates/tw-control/tests/stored_bodies.rs @@ -134,6 +134,7 @@ async fn world(mode: &str) -> World { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), gw.bus.clone())), gateway: gw.clone(), store: Some(store.clone()), diff --git a/crates/tw-control/tests/summary.rs b/crates/tw-control/tests/summary.rs index 5ff62560..1b08027a 100644 --- a/crates/tw-control/tests/summary.rs +++ b/crates/tw-control/tests/summary.rs @@ -73,6 +73,7 @@ fn app(rows: &[tw_store::db::RequestRow]) -> (tempfile::TempDir, axum::Router) { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: Some(Arc::new(tokio::sync::Mutex::new(rec))), diff --git a/crates/tw-control/tests/sync_devices.rs b/crates/tw-control/tests/sync_devices.rs index 1214173b..54daf193 100644 --- a/crates/tw-control/tests/sync_devices.rs +++ b/crates/tw-control/tests/sync_devices.rs @@ -43,6 +43,7 @@ impl Dev { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, @@ -616,6 +617,7 @@ async fn in_safe_mode_the_sync_location_still_lists_its_versions() { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(path, gw.clone(), bus).standing_in(&rejected)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/sync_regressions.rs b/crates/tw-control/tests/sync_regressions.rs index 74490dc4..c7b2e653 100644 --- a/crates/tw-control/tests/sync_regressions.rs +++ b/crates/tw-control/tests/sync_regressions.rs @@ -39,6 +39,7 @@ impl Dev { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-control/tests/transcript.rs b/crates/tw-control/tests/transcript.rs index dfd45098..a739ea48 100644 --- a/crates/tw-control/tests/transcript.rs +++ b/crates/tw-control/tests/transcript.rs @@ -88,6 +88,7 @@ fn app( let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: Some(Arc::new(tokio::sync::Mutex::new(rec))), diff --git a/crates/tw-control/tests/turn_context.rs b/crates/tw-control/tests/turn_context.rs index feb388d7..33ab817e 100644 --- a/crates/tw-control/tests/turn_context.rs +++ b/crates/tw-control/tests/turn_context.rs @@ -85,6 +85,7 @@ fn app( let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: Some(Arc::new(tokio::sync::Mutex::new(rec))), diff --git a/crates/tw-control/tests/upstream_health.rs b/crates/tw-control/tests/upstream_health.rs index 43e47a44..396ddedc 100644 --- a/crates/tw-control/tests/upstream_health.rs +++ b/crates/tw-control/tests/upstream_health.rs @@ -80,6 +80,7 @@ fn app(rows: &[tw_store::db::RequestRow]) -> (tempfile::TempDir, axum::Router) { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: Some(Arc::new(tokio::sync::Mutex::new(rec))), diff --git a/crates/tw-control/tests/upstream_limits.rs b/crates/tw-control/tests/upstream_limits.rs index d4c264be..8e7526ed 100644 --- a/crates/tw-control/tests/upstream_limits.rs +++ b/crates/tw-control/tests/upstream_limits.rs @@ -74,6 +74,7 @@ fn bed(db: Option) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p.clone(), gw.clone(), gw.bus.clone())), gateway: gw.clone(), store: Some(rec.clone()), diff --git a/crates/tw-control/tests/zai.rs b/crates/tw-control/tests/zai.rs index 7838e550..936ee4e3 100644 --- a/crates/tw-control/tests/zai.rs +++ b/crates/tw-control/tests/zai.rs @@ -283,6 +283,7 @@ async fn bed(extra: impl FnOnce(&Endpoints) -> String) -> Bed { let state = ControlState { shutdown: Default::default(), remote: Default::default(), + handover: Default::default(), cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), gateway: gw, store: None, diff --git a/crates/tw-gateway/src/affinity.rs b/crates/tw-gateway/src/affinity.rs index b027b2dd..1952620f 100644 --- a/crates/tw-gateway/src/affinity.rs +++ b/crates/tw-gateway/src/affinity.rs @@ -345,6 +345,65 @@ impl Affinity { f(e); } + /// 交接：每段对话上次由谁回答、封存的内容是谁的（见 [`crate::handover`]）。**这一轮沿用的 + /// 路由决定不交**：它只在旧 core 那一份配置下有效，换了进程一定对不上 —— 同一轮里留在原来 + /// 那一家靠的是回答它的那一家（[`Self::stay`]），规则重新求一次值 + pub(crate) fn handover_export(&self) -> Vec { + use crate::handover::wire::{AffinityAnswer, AffinityEntry}; + let mut out: Vec = self + .lock() + .iter() + .map(|(key, e)| AffinityEntry { + key: key.clone(), + touched_ms: e.touched_ms, + answer: e.answer.as_ref().map(|a| AffinityAnswer { + group: a.group.clone(), + provider: a.provider.clone(), + turn: a.turn, + at_ms: a.at_ms, + cache: a.cache, + }), + holder: e.holder.clone(), + sealed: e.sealed, + }) + .collect(); + // 最近动过的在前：接的那一边满了，丢的是最久没动静的 + out.sort_by_key(|e| std::cmp::Reverse(e.touched_ms)); + out + } + + /// 交接：接进旧 core 记着的。这里已经记着、而且更新的那段对话不动 + pub(crate) fn handover_import(&self, entries: Vec) { + let mut map = self.lock(); + for e in entries { + if map + .get(&e.key) + .is_some_and(|m| m.touched_ms >= e.touched_ms) + { + continue; + } + if !map.contains_key(&e.key) && map.len() >= CAP { + continue; + } + map.insert( + e.key, + Entry { + touched_ms: e.touched_ms, + held: None, + answer: e.answer.map(|a| Answer { + group: a.group, + provider: a.provider, + turn: a.turn, + at_ms: a.at_ms, + cache: a.cache, + }), + holder: e.holder, + sealed: e.sealed, + }, + ); + } + } + #[cfg(test)] fn tracked(&self) -> usize { self.lock().len() diff --git a/crates/tw-gateway/src/balance.rs b/crates/tw-gateway/src/balance.rs index 53b1727b..2e35b5ef 100644 --- a/crates/tw-gateway/src/balance.rs +++ b/crates/tw-gateway/src/balance.rs @@ -54,6 +54,30 @@ impl Balance { Turn { rounds, group: g } } + /// 交接：每个组轮到哪儿了（见 [`crate::handover`]） + pub(crate) fn handover_export(&self) -> Vec { + self.lock() + .iter() + .map(|(group, r)| crate::handover::wire::BalanceRound { + group: group.clone(), + members: r.members.clone(), + current: r.current.clone(), + }) + .collect() + } + + /// 交接：接进旧 core 轮到的地方。成员和权重和这里的配置对不上的，下一次排时自己从头轮 + /// （见 [`Self::turn`]）；这里已经轮过的组不动 + pub(crate) fn handover_import(&self, rounds: Vec) { + let mut map = self.lock(); + for r in rounds { + map.entry(r.group).or_insert(Round { + members: r.members, + current: r.current, + }); + } + } + /// 组 `g` 此刻的当前权重，**只看不记**。试算用：它和数据面读同一份，说完了不改变答案 pub fn peek(&self, g: &Group) -> HashMap { self.lock() diff --git a/crates/tw-gateway/src/balances.rs b/crates/tw-gateway/src/balances.rs index 95bda5a2..0dd362cf 100644 --- a/crates/tw-gateway/src/balances.rs +++ b/crates/tw-gateway/src/balances.rs @@ -375,13 +375,20 @@ pub async fn run(http: &reqwest::Client, step: Step, base_url: &str, auth: &str) } /// 后台：启动时读一遍，之后按节奏读（见 [`tracker`]）。有请求结束时被提前叫醒。 +/// +/// **排空开始就停**（交接之后的旧 core）：余额由接替者读 pub fn spawn(state: crate::AppState) { tokio::spawn(async move { + let mut draining = state.handover.draining(); loop { + if *draining.borrow() { + return; + } let _ = state.start_due_balances(); tokio::select! { _ = tokio::time::sleep(TICK) => {} _ = state.balances.wake.notified() => {} + _ = draining.changed() => {} } } }); diff --git a/crates/tw-gateway/src/cache_marks.rs b/crates/tw-gateway/src/cache_marks.rs index daccc7b4..f2803232 100644 --- a/crates/tw-gateway/src/cache_marks.rs +++ b/crates/tw-gateway/src/cache_marks.rs @@ -46,6 +46,25 @@ impl Refused { by.contains_key(&(upstream.to_string(), model.to_string())) } + /// 交接：全部拒过的（见 [`crate::handover`]） + pub(crate) fn handover_export(&self) -> Vec { + let by = self.by.lock().unwrap_or_else(|p| p.into_inner()); + by.iter() + .map(|((upstream, model), at)| crate::handover::wire::CacheMark { + upstream: upstream.clone(), + model: model.clone(), + at_ms: *at, + }) + .collect() + } + + /// 交接：接进旧 core 记着的 + pub(crate) fn handover_import(&self, marks: Vec) { + for m in marks { + self.note(&m.upstream, &m.model, m.at_ms); + } + } + #[cfg(test)] fn len(&self) -> usize { self.by.lock().unwrap().len() diff --git a/crates/tw-gateway/src/handover.rs b/crates/tw-gateway/src/handover.rs new file mode 100644 index 00000000..2e849fd5 --- /dev/null +++ b/crates/tw-gateway/src/handover.rs @@ -0,0 +1,426 @@ +//! 交接时网关这一边要交出去、接进来的东西：内存状态的快照，排空期每个请求的结算。 +//! +//! 搬运这些的是 `tw-handover`，它只当成 JSON；格式由这里定义（[`wire`]），和交接连接同一个 +//! 版本号（`tw_handover::VERSION`），**不做跨版本兼容**。 +//! +//! # 快照 +//! +//! 接管开始时旧 core 交出（[`AppState::handover_export`]），新 core 接进来 +//! （[`AppState::handover_import`]），新的就不从零开始：对话粘滞、会话归属、加密历史与缓存 +//! 断点的拒收记录、熔断与健康、限流暂停、延迟与速度、负载均衡轮次、OAuth 访问令牌缓存、 +//! 密钥和上游的用量与在途预留，以及旧 core 此刻占着的并发名额（按密钥、按上游）。 +//! +//! **不搬的**：模型目录、余额、订阅额度、GLM 额度 —— 新 core 起来后台重新取，期间显示为 +//! 获取中，搬过去的也只是几分钟前的样子；对话这一轮沿用的路由决定 —— 它认的是旧 core 那一份 +//! 引擎，换了进程一定对不上，同一轮留在原来那一家靠的是上次回答它的那一家；限流的排队和 +//! 探路名额、在途请求的中止柄和实时内容 —— 都是旧 core 手上的请求，跟着它们在旧 core 上结束； +//! 「已经报过」的那几张表（写回成功、凭据失效、额度用完）—— 新 core 的界面也是新开的，状况 +//! 还在的话再报一次。 +//! +//! # 排空期的名额 +//! +//! 两个进程同时转发时，名额和用量上限不能翻倍：快照里旧 core 占着的名额，新 core 记成已占用 +//! （「替它占着」，见 [`crate::limits::Gate::phantom`]、[`crate::slots::Slots::phantom`]）；旧 core +//! 每放掉一个名额、每结算一份用量发一条结算（[`AppState::begin_drain`] 的接收端，格式见 +//! [`wire::Event`]），新 core 据此放掉（[`AppState::handover_settled`]）。交接连接断了，剩下的 +//! 一并放掉（[`AppState::handover_release_all`]），之后由请求库重建用量校正。 +//! +//! **密钥的名额连排着队等的也算**：旧 core 在提交之前收下、还在排队的请求，排在新 core 的新 +//! 请求前面 —— 不算的话，它前面那个结束、新 core 放掉一个名额的那一刻，新 core 的新请求和它 +//! 会同时进去。 +//! +//! **从快照到排空开始之间**，旧 core 照常转发（接管还没提交），名额的变化不报：排空开始时 +//! 重报一份此刻占着的全部（[`wire::Event::GateHolds`] 等），新 core 照它换掉手上记着的。用量的 +//! 变化从快照那一刻起就攒着（[`feed::Feed`]），排空开始时一起交出去：快照之后结算的那些请求 +//! 用了多少，新 core 一笔不落。 +//! +//! # WebSocket +//! +//! 排空开始后，Responses 的连接在手上这一轮答完时以 1012（服务重启）关闭，闲着的当场关闭： +//! 客户端重连，连到新 core 上（见 `ws`）。整条连接一行的那几种（Realtime 等）分不出一轮 +//! 一轮，照常开到客户端或上游收场。 + +use std::collections::HashMap; +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; + +use serde::de::DeserializeOwned; +use tokio::sync::{mpsc, watch}; + +use crate::state::AppState; + +pub(crate) mod feed; +#[cfg(test)] +mod tests; +pub(crate) mod wire; + +/// 交接时网关这一边记着的。**跨重载存活**（在 `AppState` 里） +pub(crate) struct Handover { + /// 导出快照时开的那一路和它的接收端：用量的变化从那一刻起攒在里面，排空开始时交出去 + pending: Mutex, mpsc::UnboundedReceiver)>>, + /// 排空开始了（旧 core）：WebSocket 看着它收场 + drain: watch::Sender, + /// 替旧 core 占着的名额（新 core） + phantoms: Mutex, +} + +impl Default for Handover { + fn default() -> Self { + Self { + pending: Mutex::default(), + drain: watch::channel(false).0, + phantoms: Mutex::default(), + } + } +} + +/// 替旧 core 占着的名额：按密钥的通行证、按上游的位置。**丢掉就放掉**，和别的一样 +#[derive(Default)] +struct Phantoms { + gate: HashMap>, + slots: HashMap>, +} + +impl Handover { + fn phantoms(&self) -> MutexGuard<'_, Phantoms> { + self.phantoms.lock().unwrap_or_else(PoisonError::into_inner) + } + + /// 看着排空开始的那一下（`ws` 用） + pub(crate) fn draining(&self) -> watch::Receiver { + self.drain.subscribe() + } +} + +/// 快照里的一节：读不懂就丢掉这一节，记一行 +fn section( + m: &serde_json::Map, + name: &str, +) -> Option { + let v = m.get(name)?; + match serde_json::from_value(v.clone()) { + Ok(t) => Some(t), + Err(e) => { + tracing::warn!( + section = name, + "a part of the handover snapshot could not be read and is left out: {e}" + ); + None + } + } +} + +/// 快照里一节的每一条：读不懂的那一条丢掉，别的照接 +fn entries( + m: &serde_json::Map, + name: &str, +) -> Vec { + let Some(items) = m.get(name).and_then(serde_json::Value::as_array) else { + return Vec::new(); + }; + let out: Vec = items + .iter() + .filter_map(|v| serde_json::from_value(v.clone()).ok()) + .collect(); + if out.len() < items.len() { + tracing::warn!( + section = name, + left_out = items.len() - out.len(), + "some entries of the handover snapshot could not be read and are left out" + ); + } + out +} + +impl AppState { + /// 旧 core：此刻的内存状态。 + /// + /// 从这一刻起用量的变化（预留、结算）攒着，[`Self::begin_drain`] 时一起交出去。**再调一次 + /// 从头攒**：上一次接管没成、换了一个接替者，上一次攒的那些没人要了 + pub fn handover_export(&self) -> serde_json::Value { + let cfg = self.config(); + let (feed, rx) = feed::Feed::new(); + let key_limits = self.key_limits.handover_export(Some(feed.clone())); + let upstream_limits = self.upstream_limits.handover_export(Some(feed.clone())); + *self + .handover + .pending + .lock() + .unwrap_or_else(PoisonError::into_inner) = Some((feed, rx)); + let mut m = serde_json::Map::new(); + let mut put = |name: &str, v: serde_json::Result| match v { + Ok(v) => { + m.insert(name.to_string(), v); + } + Err(e) => tracing::warn!( + section = name, + "a part of the handover snapshot could not be written: {e}" + ), + }; + put(wire::AT_MS, Ok(crate::server::now_ms().into())); + put( + wire::AFFINITY, + serde_json::to_value(self.affinity.handover_export()), + ); + put( + wire::SEALS, + serde_json::to_value(self.seals.handover_export()), + ); + put( + wire::CACHE_MARKS, + serde_json::to_value(self.cache_marks.handover_export()), + ); + put( + wire::SESSIONS, + serde_json::to_value(self.sessions.handover_export()), + ); + put( + wire::HEALTH, + serde_json::to_value(self.health.handover_export()), + ); + put( + wire::THROTTLES, + serde_json::to_value(self.throttles.handover_export()), + ); + put( + wire::LATENCY, + serde_json::to_value(self.latency.handover_export()), + ); + put( + wire::PACE, + serde_json::to_value(self.pace.handover_export()), + ); + put( + wire::BALANCE, + serde_json::to_value(self.balance.handover_export()), + ); + put( + wire::OAUTH, + serde_json::to_value(self.oauth.handover_export(&cfg.providers)), + ); + put(wire::KEY_LIMITS, serde_json::to_value(key_limits)); + put(wire::UPSTREAM_LIMITS, serde_json::to_value(upstream_limits)); + put(wire::GATE, serde_json::to_value(self.gate.claims.counts())); + put( + wire::SLOTS, + serde_json::to_value(self.slots.claims.counts()), + ); + serde_json::Value::Object(m) + } + + /// 新 core：接进旧 core 的快照。**认不出的部分丢掉**，不让交接失败 —— 版本号 + /// 相同时不该发生，发生了也只是少接一样状态。 + /// + /// 只调一次，在第一个请求之前（快照在监听交过来之前就到了）。再调一次的话，上一次替旧 core + /// 占着的名额先全部放掉 + pub fn handover_import(&self, snapshot: &serde_json::Value) { + let Some(m) = snapshot.as_object() else { + tracing::warn!("the handover snapshot is not an object and is ignored"); + return; + }; + self.handover_release_all(); + let lag_ms = m + .get(wire::AT_MS) + .and_then(serde_json::Value::as_u64) + .map_or(0, |at| crate::server::now_ms().saturating_sub(at)); + let cfg = self.config(); + self.affinity.handover_import(entries(m, wire::AFFINITY)); + self.seals.handover_import(entries(m, wire::SEALS)); + self.cache_marks + .handover_import(entries(m, wire::CACHE_MARKS)); + self.sessions.handover_import(entries(m, wire::SESSIONS)); + if let Some(h) = section(m, wire::HEALTH) { + self.health.handover_import(h); + } + self.throttles + .handover_import(entries(m, wire::THROTTLES), lag_ms); + self.latency.handover_import(entries(m, wire::LATENCY)); + self.pace.handover_import(entries(m, wire::PACE)); + self.balance.handover_import(entries(m, wire::BALANCE)); + self.oauth + .handover_import(entries(m, wire::OAUTH), &cfg.providers, lag_ms); + if let Some(b) = section(m, wire::KEY_LIMITS) { + self.key_limits.handover_import(b); + } + if let Some(b) = section(m, wire::UPSTREAM_LIMITS) { + self.upstream_limits.handover_import(b); + } + if let Some(n) = section::>(m, wire::GATE) { + self.hold_gate(&n); + } + if let Some(n) = section::>(m, wire::SLOTS) { + self.hold_slots(&n); + } + } + + /// 旧 core：提交之后调一次。从这一刻起： + /// + /// - 每放掉一个名额、每结算一份用量，往返回的接收端发一条结算（格式见 [`wire::Event`]）； + /// 第一批是此刻占着的全部名额和预留，之前攒着的用量变化排在它们前面； + /// - WebSocket 会话在当前这一轮结束后以 1012（服务重启）关闭，客户端重连到新 core。 + /// + /// **只调一次**。再调的话交回一个新的接收端，之前那个从此收不到东西 + pub fn begin_drain(&self) -> tokio::sync::mpsc::UnboundedReceiver { + let pending = self + .handover + .pending + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take(); + let (feed, rx) = match pending { + Some((f, rx)) if !f.lost() => (f, rx), + // 没导出过快照，或者攒满了不再攒：从现在起报，开头那份全部照报 + _ => feed::Feed::new(), + }; + feed.start_draining(); + self.gate.claims.drain_into(feed.clone()); + self.slots.claims.drain_into(feed.clone()); + self.key_limits.handover_drain(feed.clone()); + self.upstream_limits.handover_drain(feed); + self.handover.drain.send_replace(true); + rx + } + + /// 新 core：旧 core 发来一条结算（[`Self::begin_drain`] 的接收端收到的那一条）。读不懂的 + /// 丢掉 + pub fn handover_settled(&self, settled: &serde_json::Value) { + let ev: wire::Event = match serde_json::from_value(settled.clone()) { + Ok(ev) => ev, + Err(e) => { + tracing::debug!("a handover settlement could not be read and is ignored: {e}"); + return; + } + }; + match &ev { + wire::Event::GateHolds { n } => self.hold_gate(n), + wire::Event::Gate { key, d } => self.step_gate(key, *d), + wire::Event::SlotHolds { n } => self.hold_slots(n), + wire::Event::Slot { upstream, d } => self.step_slot(upstream, *d), + ev => { + if !self.key_limits.handover_apply(ev) { + self.upstream_limits.handover_apply(ev); + } + } + } + } + + /// 新 core：交接连接断了（旧 core 退出或崩溃）。快照里记着的旧 core 的名额与 + /// 预留全部放掉。 + pub fn handover_release_all(&self) { + // 在锁外面丢：放掉一个通行证要拿闸自己的锁 + let dropped = std::mem::take(&mut *self.handover.phantoms()); + drop(dropped); + self.key_limits.handover_release_all(); + self.upstream_limits.handover_release_all(); + } + + // ------------------------------------------------------------ 替旧 core 占着的名额 + + /// 这把密钥此刻配置里的并发上限。没设的不用替谁占 + fn key_limit(&self, key: &str) -> Option { + self.config() + .clients + .iter() + .find(|c| c.name == key) + .and_then(|c| c.max_concurrent) + } + + /// 替旧 core 占着的密钥名额换成 `want` 那么多 + fn hold_gate(&self, want: &HashMap) { + let mut dropped = Vec::new(); + { + let mut ph = self.handover.phantoms(); + let stale: Vec = ph + .gate + .keys() + .filter(|k| !want.contains_key(*k)) + .cloned() + .collect(); + for k in stale { + dropped.extend(ph.gate.remove(&k).unwrap_or_default()); + } + for (key, n) in want { + let Some(limit) = self.key_limit(key) else { + dropped.extend(ph.gate.remove(key).unwrap_or_default()); + continue; + }; + let held = ph.gate.entry(key.clone()).or_default(); + while held.len() > *n { + dropped.extend(held.pop()); + } + while held.len() < *n { + held.push(self.gate.phantom(key, limit)); + } + } + ph.gate.retain(|_, v| !v.is_empty()); + } + drop(dropped); + } + + fn step_gate(&self, key: &str, d: i64) { + let dropped = { + let mut ph = self.handover.phantoms(); + if d > 0 { + if let Some(limit) = self.key_limit(key) { + let pass = self.gate.phantom(key, limit); + ph.gate.entry(key.to_string()).or_default().push(pass); + } + None + } else { + let held = ph.gate.get_mut(key); + let one = held.and_then(Vec::pop); + ph.gate.retain(|_, v| !v.is_empty()); + one + } + }; + drop(dropped); + } + + /// 替旧 core 占着的上游位置换成 `want` 那么多 + fn hold_slots(&self, want: &HashMap) { + let mut dropped = Vec::new(); + { + let mut ph = self.handover.phantoms(); + let stale: Vec = ph + .slots + .keys() + .filter(|k| !want.contains_key(*k)) + .cloned() + .collect(); + for k in stale { + dropped.extend(ph.slots.remove(&k).unwrap_or_default()); + } + for (name, n) in want { + let held = ph.slots.entry(name.clone()).or_default(); + while held.len() > *n { + dropped.extend(held.pop()); + } + while held.len() < *n { + // 这家此刻不限并发：不用替谁占 + let Some(slot) = self.slots.phantom(name) else { + break; + }; + held.push(slot); + } + } + ph.slots.retain(|_, v| !v.is_empty()); + } + drop(dropped); + } + + fn step_slot(&self, name: &str, d: i64) { + let dropped = { + let mut ph = self.handover.phantoms(); + if d > 0 { + if let Some(slot) = self.slots.phantom(name) { + ph.slots.entry(name.to_string()).or_default().push(slot); + } + None + } else { + let one = ph.slots.get_mut(name).and_then(Vec::pop); + ph.slots.retain(|_, v| !v.is_empty()); + one + } + }; + drop(dropped); + } +} diff --git a/crates/tw-gateway/src/handover/feed.rs b/crates/tw-gateway/src/handover/feed.rs new file mode 100644 index 00000000..c289cb18 --- /dev/null +++ b/crates/tw-gateway/src/handover/feed.rs @@ -0,0 +1,185 @@ +//! 旧 core 往交接连接上报变化的那一路（[`Feed`]），和名额的计数（[`Claims`]）。 + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; + +use tokio::sync::mpsc; + +use super::wire::Event; + +/// 导出快照之后、排空开始之前最多攒多少条。 +/// +/// **攒着是为了排空开始时一条不漏**：快照是接管一开始导出的，之后到提交那一刻，旧 core 照常 +/// 转发，用量上限的预留和结算还在变。交接正常时这一段不过一两秒；交接失败退回待交接、又迟迟 +/// 没有下一个接替者时，攒下的东西没人要 —— 到了这个数就不攒了，排空开始时改发一份此刻占着的 +/// 全部（`*_holds`），这一段的用量等新 core 打开请求库后从记录里重建 +const PENDING_MAX: usize = 20_000; + +/// 一路结算。**发不阻塞**：持着各自的锁发（见各个调用处），同一样东西的变化按发生的先后进来。 +pub(crate) struct Feed { + tx: mpsc::UnboundedSender, + /// 排空开始之前攒了几条 + pending: AtomicUsize, + /// 排空开始了：接收端有人在读，不再数 + draining: AtomicBool, + /// 攒满了，后面的不再攒 + lost: AtomicBool, +} + +impl std::fmt::Debug for Feed { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Feed") + .field("draining", &self.draining.load(Ordering::Relaxed)) + .field("lost", &self.lost.load(Ordering::Relaxed)) + .finish_non_exhaustive() + } +} + +impl Feed { + pub(crate) fn new() -> (Arc, mpsc::UnboundedReceiver) { + let (tx, rx) = mpsc::unbounded_channel(); + let feed = Arc::new(Self { + tx, + pending: AtomicUsize::new(0), + draining: AtomicBool::new(false), + lost: AtomicBool::new(false), + }); + (feed, rx) + } + + pub(crate) fn send(&self, ev: &Event) { + if self.lost.load(Ordering::Relaxed) { + return; + } + if !self.draining.load(Ordering::Relaxed) + && self.pending.fetch_add(1, Ordering::Relaxed) >= PENDING_MAX + { + self.lost.store(true, Ordering::Relaxed); + return; + } + // 接收端没了（交接连接断了）：发不出去就算了，没有人再要它 + let _ = self.tx.send(ev.to_value()); + } + + pub(crate) fn start_draining(&self) { + self.draining.store(true, Ordering::Relaxed); + } + + pub(crate) fn lost(&self) -> bool { + self.lost.load(Ordering::Relaxed) + } + + #[cfg(test)] + pub(crate) fn mark_lost(&self) { + self.lost.store(true, Ordering::Relaxed); + } +} + +/// 发到这一路上（没接上就什么都不做） +pub(crate) fn send(feed: &Option>, ev: impl FnOnce() -> Event) { + if let Some(f) = feed { + f.send(&ev()); + } +} + +/// 名额是密钥的还是上游的 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub(crate) enum Kind { + #[default] + Gate, + Slot, +} + +/// 每个名字此刻占着几个名额（密钥的连排着队等的也算）。导出快照、排空开始时报的就是它； +/// 排空开始之后每变一次报一次（[`Event::Gate`]、[`Event::Slot`]）。 +/// +/// **和信号量分开数**：信号量说不出「此刻有几个在等」，也分不清空着的位置是上限调小了欠着的 +/// 还是真空着。计数和报出去的变化在同一把锁里，排空开始时报的那份全部和之后的变化前后接得上。 +#[derive(Default)] +pub(crate) struct Claims { + kind: Kind, + inner: Mutex, +} + +#[derive(Default)] +struct ClaimsInner { + n: HashMap, + feed: Option>, +} + +impl Claims { + pub(crate) fn new(kind: Kind) -> Self { + Self { + kind, + inner: Mutex::default(), + } + } + + fn lock(&self) -> MutexGuard<'_, ClaimsInner> { + self.inner.lock().unwrap_or_else(PoisonError::into_inner) + } + + fn event(&self, name: &str, d: i64) -> Event { + match self.kind { + Kind::Gate => Event::Gate { + key: name.to_string(), + d, + }, + Kind::Slot => Event::Slot { + upstream: name.to_string(), + d, + }, + } + } + + /// 占上一个。**丢掉交回的那一个就放掉** + pub(crate) fn begin(self: &Arc, name: &str) -> Claim { + let mut g = self.lock(); + *g.n.entry(name.to_string()).or_default() += 1; + send(&g.feed, || self.event(name, 1)); + Claim { + claims: self.clone(), + name: name.to_string(), + } + } + + fn end(&self, name: &str) { + let mut g = self.lock(); + if let Some(n) = g.n.get_mut(name) { + *n = n.saturating_sub(1); + if *n == 0 { + g.n.remove(name); + } + } + send(&g.feed, || self.event(name, -1)); + } + + /// 此刻每个名字占着几个 + pub(crate) fn counts(&self) -> HashMap { + self.lock().n.clone() + } + + /// 排空开始：报一份此刻的全部，之后的每一次变化接着报 + pub(crate) fn drain_into(&self, feed: Arc) { + let mut g = self.lock(); + let n = g.n.clone(); + feed.send(&match self.kind { + Kind::Gate => Event::GateHolds { n }, + Kind::Slot => Event::SlotHolds { n }, + }); + g.feed = Some(feed); + } +} + +/// 占着的一个名额。**丢掉就放掉**，跟着通行证（[`crate::limits::Pass`]、[`crate::slots::Slot`]）走 +pub(crate) struct Claim { + claims: Arc, + name: String, +} + +impl Drop for Claim { + fn drop(&mut self) { + self.claims.end(&self.name); + } +} diff --git a/crates/tw-gateway/src/handover/tests.rs b/crates/tw-gateway/src/handover/tests.rs new file mode 100644 index 00000000..1510fabf --- /dev/null +++ b/crates/tw-gateway/src/handover/tests.rs @@ -0,0 +1,582 @@ +use std::time::Duration; + +use super::*; +use crate::affinity::{Conversation, Turn}; +use crate::key_limits::{Ask, Recorded}; + +fn parse(yaml: &str) -> Vec { + serde_yaml_ng::from_str(yaml).unwrap() +} + +fn key_limits() -> Vec { + parse("- {per: day, tokens: 100000}") +} + +fn upstream_limits() -> Vec { + parse("- {per: day, requests: 1000}") +} + +/// 一小时之后过期，RFC 3339 +fn in_an_hour() -> String { + (chrono::Utc::now() + chrono::Duration::hours(1)) + .to_rfc3339_opts(chrono::SecondsFormat::Secs, true) +} + +fn oauth(access: &str) -> tw_config::OAuth { + tw_config::OAuth { + access: Some(access.into()), + expires_at: Some(in_an_hour()), + refresh: "rt-1".into(), + endpoint: "https://auth.invalid/token".into(), + client_id: None, + client_secret: None, + refresh_before: None, + } +} + +/// 密钥 `k`（并发上限 `max`、每天 10 万 token），上游 `甲`（OAuth、并发上限 1、每天 1000 个 +/// 请求）和 `乙` +fn config(max: usize, access: &str) -> tw_config::Config { + tw_config::Config { + clients: vec![tw_config::Client { + name: "k".into(), + key: "tw-k".into(), + max_concurrent: Some(max), + limits: key_limits(), + ..Default::default() + }], + providers: vec![ + tw_config::Provider { + name: "甲".into(), + base_url: "https://a.invalid".into(), + oauth: Some(oauth(access)), + max_concurrent: Some(1), + limits: upstream_limits(), + ..Default::default() + }, + tw_config::Provider { + name: "乙".into(), + base_url: "https://b.invalid".into(), + key: Some("sk-b".into()), + ..Default::default() + }, + ], + ..Default::default() + } +} + +/// 旧 core 和新 core，同一份配置 +fn pair() -> (AppState, AppState) { + ( + AppState::new(config(1, "at-1")).unwrap(), + AppState::new(config(1, "at-1")).unwrap(), + ) +} + +/// 过一遍交接连接：写成字符串再读回来 +fn over_the_wire(v: serde_json::Value) -> serde_json::Value { + serde_json::from_str(&v.to_string()).unwrap() +} + +/// 让等着的有机会走一步 +async fn settle() { + tokio::time::sleep(Duration::from_millis(30)).await; +} + +/// 排空开始：旧 core 报的每一条都交给新 core，像交接连接那样按顺序 +fn forward(old: &AppState, new: &AppState) -> tokio::task::JoinHandle<()> { + let mut rx = old.begin_drain(); + let new = new.clone(); + tokio::spawn(async move { + while let Some(v) = rx.recv().await { + new.handover_settled(&over_the_wire(v)); + } + }) +} + +fn group() -> tw_engine::Group { + tw_engine::Group { + name: "池子".into(), + kind: tw_engine::GroupType::LoadBalance, + providers: [("甲", 2), ("乙", 1)] + .iter() + .map(|(n, w)| tw_engine::Member { + name: n.to_string(), + weight: *w, + }) + .collect(), + selected: None, + balance_by: Default::default(), + } +} + +fn recorded(input: u64, output: u64) -> Recorded { + Recorded { + client: "k".into(), + path: "/v1/messages".into(), + reached: true, + requests: 1, + input, + output, + ..Default::default() + } +} + +fn used(s: &AppState) -> u64 { + s.key_limits.view("k", &key_limits())[0].used +} + +fn upstream_used(s: &AppState) -> u64 { + s.upstream_limits.view("甲", &upstream_limits())[0].used +} + +/// 一节，数组按各条的样子排好：两边的 HashMap 顺序不同 +fn sorted(snap: &serde_json::Value, name: &str) -> serde_json::Value { + let mut v = snap[name].clone(); + if let Some(a) = v.as_array_mut() { + a.sort_by_key(|x| x.to_string()); + } + v +} + +/// **快照往返**：旧 core 导出、新 core 接进来之后，新 core 看到的和旧 core 一样 —— 再从新 core +/// 导出一遍，搬过去的每一节都一字不差；路由、熔断、限流、用量、名额照旧起作用。 +#[tokio::test] +async fn a_snapshot_round_trip_leaves_the_new_core_where_the_old_one_was() { + let old = AppState::new(config(1, "at-1")).unwrap(); + // 新 core 的配置里 access token 是旧的那个（写回还没落地）：接到的缓存让它用旧 core 手上的 + let new = AppState::new(config(1, "at-0")).unwrap(); + let now = crate::server::now_ms(); + + let c = Conversation::new( + "default", + "对话", + Turn { + number: 1, + within: false, + }, + ); + old.affinity + .ticket(&c, Some("池子".into()), "甲") + .answered(4096); + assert_eq!(old.affinity.holder(&c, now).as_deref(), Some("甲")); + old.seals.note("对话", "乙", &[[7u8; 32]], now); + old.cache_marks.note("甲", "claude-opus", now); + let session = old.sessions.assign("fp", now); + for _ in 0..3 { + old.health.record_failure("乙"); + } + assert!(!old.health.is_available("乙")); + old.throttles.limited( + "甲", + std::time::Instant::now(), + Some(Duration::from_secs(120)), + ); + for ms in [100, 120, 140] { + old.latency.record("甲", ms); + } + old.pace.learn("甲", "m", 100, 150); + let g = group(); + { + let members = g.names(); + let t = old.balance.turn(&g); + let f = tw_engine::Facts { + current_weight: t.current(), + ..Default::default() + }; + t.charge(&members, &f, "甲"); + } + let (token, _) = old + .oauth + .token("甲", &oauth("at-1"), &old.http) + .await + .unwrap(); + assert_eq!(token, "at-1"); + // 一个结算过的（150 token），一个在跑的（预留 600） + let done = old + .key_limits + .admit("k", &key_limits(), Ask::default(), Duration::ZERO) + .await + .unwrap(); + done.bind(6); + old.key_limits + .settle(6, crate::server::now_ms() as i64, &recorded(100, 50)); + let running = old + .key_limits + .admit( + "k", + &key_limits(), + Ask { + tokens: 600, + cost_micros: 0, + }, + Duration::ZERO, + ) + .await + .unwrap(); + running.bind(7); + let hop = old + .upstream_limits + .reserve("甲", &upstream_limits(), &Ask::default()); + hop.bind(7); + let _pass = old.gate.acquire("k", Some(1)).await; + let _slot = old.slots.try_take("甲").unwrap(); + + let snap = over_the_wire(old.handover_export()); + new.handover_import(&snap); + + // 搬过去的一字不差 + let again = over_the_wire(new.handover_export()); + for name in [ + wire::AFFINITY, + wire::SEALS, + wire::CACHE_MARKS, + wire::SESSIONS, + wire::LATENCY, + wire::PACE, + wire::BALANCE, + wire::GATE, + wire::SLOTS, + ] { + assert_eq!(sorted(&again, name), sorted(&snap, name), "{name}"); + } + assert_eq!(again[wire::HEALTH], snap[wire::HEALTH]); + let books = |s: &serde_json::Value, n: &str| { + let mut b = s[n]["books"].clone(); + b.as_array_mut().unwrap().sort_by_key(|x| x.to_string()); + b + }; + assert_eq!( + books(&again, wire::KEY_LIMITS), + books(&snap, wire::KEY_LIMITS) + ); + assert_eq!( + books(&again, wire::UPSTREAM_LIMITS), + books(&snap, wire::UPSTREAM_LIMITS) + ); + assert_eq!(again[wire::KEY_LIMITS]["held"].as_array().unwrap().len(), 1); + assert_eq!(again[wire::OAUTH][0]["access"], "at-1"); + let left = |s: &serde_json::Value| s[wire::THROTTLES][0]["pause_left_ms"].as_u64().unwrap(); + assert!(left(&again) <= left(&snap) && left(&snap) - left(&again) < 2_000); + + // 照旧起作用 + let mut candidates = vec!["乙".to_string(), "甲".to_string()]; + let stay = new + .affinity + .stay(&c, Some("池子"), &mut candidates, |_| true, now); + assert_eq!(stay, Some(tw_api::Stay::Cache)); + assert_eq!(candidates[0], "甲"); + assert_eq!(new.affinity.holder(&c, now).as_deref(), Some("甲")); + assert!(new.seals.of("对话", "乙").unwrap().contains(&[7u8; 32])); + assert!(new.cache_marks.refused("甲", "claude-opus")); + assert_eq!(new.sessions.assign("fp", now + 1), session); + assert!(!new.health.is_available("乙"), "熔断中的上游又被放行了"); + let paused = new.throttles.view("甲").expect("限流暂停丢了"); + assert_eq!(paused.kind, tw_api::ThrottleKind::RateLimited); + assert!(paused.until_ms > now + 100_000); + assert_eq!(new.latency.typical("甲"), Some(120)); + assert_eq!(new.pace.ratio("甲", "m"), 1.5); + assert_eq!(new.balance.peek(&g), old.balance.peek(&g)); + let (token, renewed) = new + .oauth + .token("甲", &oauth("at-0"), &new.http) + .await + .unwrap(); + assert_eq!(token, "at-1", "新 core 没用上旧 core 手上的 access token"); + assert!(renewed.is_none()); + assert_eq!(used(&new), 750, "结算过的 150 加在跑的预留 600"); + assert_eq!(upstream_used(&new), 1); + assert!( + tokio::time::timeout(Duration::from_millis(50), new.gate.acquire("k", Some(1))) + .await + .is_err(), + "旧 core 占着的密钥名额没算上" + ); + assert!( + new.slots.try_take("甲").is_none(), + "旧 core 占着的上游位置没算上" + ); +} + +/// **V5：密钥并发上限 1 时，排空期两个进程合计不超过 1。**旧 core 一个在跑、一个排着队；新 core +/// 新来的请求等到旧 core 的两个都结束 —— 旧 core 那个排着队的先进去，新的不能趁前一个结束时 +/// 一起进去。 +#[tokio::test] +async fn v5_with_a_limit_of_one_the_two_cores_together_never_run_two() { + let (old, new) = pair(); + let running = old.gate.acquire("k", Some(1)).await; + let o = old.clone(); + let queued = tokio::spawn(async move { o.gate.acquire("k", Some(1)).await }); + settle().await; + + new.handover_import(&over_the_wire(old.handover_export())); + let _fwd = forward(&old, &new); + let n = new.clone(); + let fresh = tokio::spawn(async move { n.gate.acquire("k", Some(1)).await }); + settle().await; + assert!( + !fresh.is_finished(), + "旧 core 有一个在跑，新 core 又放进去一个" + ); + + drop(running); + settle().await; + assert!(queued.is_finished(), "旧 core 排着的那一个该进去了"); + assert!( + !fresh.is_finished(), + "旧 core 排着的那一个刚进去，新 core 的也进去了：合计 2" + ); + + drop(queued.await.unwrap()); + settle().await; + assert!(fresh.is_finished(), "旧 core 都结束了，新 core 的还在等"); + let _p = fresh.await.unwrap(); + // 名额没有因为交接多出来或者少掉 + assert!( + tokio::time::timeout(Duration::from_millis(50), new.gate.acquire("k", Some(1))) + .await + .is_err() + ); +} + +/// 旧 core 排着队、客户端走了：那个名额跟着放掉，新 core 不用等一个永远不会结束的请求 +#[tokio::test] +async fn a_request_that_gave_up_while_queued_in_the_old_core_frees_its_place() { + let (old, new) = pair(); + let running = old.gate.acquire("k", Some(1)).await; + let o = old.clone(); + let queued = tokio::spawn(async move { o.gate.acquire("k", Some(1)).await }); + settle().await; + new.handover_import(&over_the_wire(old.handover_export())); + let _fwd = forward(&old, &new); + + queued.abort(); + settle().await; + let n = new.clone(); + let fresh = tokio::spawn(async move { n.gate.acquire("k", Some(1)).await }); + settle().await; + assert!(!fresh.is_finished(), "旧 core 还有一个在跑"); + drop(running); + settle().await; + assert!(fresh.is_finished()); +} + +/// 交接连接断了：替旧 core 占着的名额和预留一并放掉 +#[tokio::test] +async fn release_all_gives_back_everything_held_for_the_old_core() { + let (old, new) = pair(); + let _pass = old.gate.acquire("k", Some(1)).await; + let _slot = old.slots.try_take("甲").unwrap(); + let hold = old + .key_limits + .admit( + "k", + &key_limits(), + Ask { + tokens: 600, + cost_micros: 0, + }, + Duration::ZERO, + ) + .await + .unwrap(); + hold.bind(7); + new.handover_import(&over_the_wire(old.handover_export())); + assert_eq!(used(&new), 600); + + new.handover_release_all(); + tokio::time::timeout(Duration::from_millis(200), new.gate.acquire("k", Some(1))) + .await + .expect("放掉之后还在等"); + assert!(new.slots.try_take("甲").is_some()); + assert_eq!(used(&new), 0); +} + +/// 旧 core 占着的比新 core 的上限还多（旧 core 那边上限是 2）：多出来的记成欠着，新 core 等它们 +/// 全部结束；结束之后新 core 的上限照旧是 1，不多不少 +#[tokio::test] +async fn more_held_than_the_new_limit_is_carried_as_debt() { + let old = AppState::new(config(2, "at-1")).unwrap(); + let new = AppState::new(config(1, "at-1")).unwrap(); + let a = old.gate.acquire("k", Some(2)).await; + let b = old.gate.acquire("k", Some(2)).await; + new.handover_import(&over_the_wire(old.handover_export())); + let _fwd = forward(&old, &new); + let n = new.clone(); + let fresh = tokio::spawn(async move { n.gate.acquire("k", Some(1)).await }); + settle().await; + drop(a); + settle().await; + assert!(!fresh.is_finished(), "旧 core 还有一个在跑"); + drop(b); + settle().await; + assert!(fresh.is_finished()); + let _held = fresh.await.unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(50), new.gate.acquire("k", Some(1))) + .await + .is_err(), + "欠账还完之后上限变大了" + ); +} + +/// 上游的位置：旧 core 那一跳结束才空出来 +#[tokio::test] +async fn an_upstream_place_held_by_the_old_core_frees_when_its_hop_ends() { + let (old, new) = pair(); + let slot = old.slots.try_take("甲").unwrap(); + new.handover_import(&over_the_wire(old.handover_export())); + assert!(new.slots.try_take("甲").is_none()); + let _fwd = forward(&old, &new); + settle().await; + assert!(new.slots.try_take("甲").is_none()); + drop(slot); + settle().await; + assert!(new.slots.try_take("甲").is_some()); +} + +/// 用量：预留在旧 core 结算时换成实数；快照之后才占上的预留、之后退回的，新 core 都跟着变 +#[tokio::test] +async fn reservations_turn_into_real_usage_when_the_old_core_settles() { + let (old, new) = pair(); + let ask = |tokens| Ask { + tokens, + cost_micros: 0, + }; + let running = old + .key_limits + .admit("k", &key_limits(), ask(600), Duration::ZERO) + .await + .unwrap(); + running.bind(7); + let hop = old + .upstream_limits + .reserve("甲", &upstream_limits(), &Ask::default()); + hop.bind(7); + new.handover_import(&over_the_wire(old.handover_export())); + assert_eq!(used(&new), 600); + assert_eq!(upstream_used(&new), 1); + + // 快照之后、排空之前：又一个准入了（还没开始），之后被丢掉 + let late = old + .key_limits + .admit("k", &key_limits(), ask(300), Duration::ZERO) + .await + .unwrap(); + let fwd = forward(&old, &new); + settle().await; + assert_eq!(used(&new), 900, "快照之后占上的预留没跟过来"); + drop(late); + settle().await; + assert_eq!(used(&new), 600, "退回的预留还算着"); + + let at = crate::server::now_ms() as i64; + old.key_limits.settle(7, at, &recorded(100, 50)); + old.upstream_limits + .settle(7, at, &["甲".to_string()], Some("甲"), &recorded(100, 50)); + settle().await; + assert_eq!(used(&new), 150, "预留没换成实数"); + assert_eq!(upstream_used(&new), 1); + assert_eq!(used(&old), 150); + fwd.abort(); +} + +/// 一段接一段地交（H8）：替更早的 core 占着的名额，交给下一个时照样算上 +#[tokio::test] +async fn names_held_for_an_earlier_core_are_handed_on() { + let (a, b) = pair(); + let _pass = a.gate.acquire("k", Some(1)).await; + b.handover_import(&over_the_wire(a.handover_export())); + let snap = b.handover_export(); + assert_eq!(snap[wire::GATE]["k"], 1); +} + +/// 认不出的部分丢掉，不 panic：读不懂的一节、一节里读不懂的一条、不认识的一节 +#[tokio::test] +async fn unreadable_parts_of_a_snapshot_are_left_out() { + let (_, new) = pair(); + let snap = serde_json::json!({ + "at_ms": "昨天", + "health": "坏了", + "pace": [{"upstream": "甲", "model": "m", "ratio": 1.25}, {"什么": 1}], + "gate": {"k": "一"}, + "以后才有的": [1, 2, 3], + }); + new.handover_import(&snap); + new.handover_import(&serde_json::json!("not even an object")); + assert_eq!(new.pace.ratio("甲", "m"), 1.25); + new.handover_settled(&serde_json::json!({"k": "以后才有的"})); + new.handover_settled(&serde_json::json!(42)); + new.handover_settled(&serde_json::json!({"k": "key_settle", "seq": 99, "at_ms": 0, "rec": {}})); +} + +/// 快照之后迟迟没有排空、攒的变化太多：不再攒，排空开始时照样报一份此刻占着的全部 +#[tokio::test] +async fn a_feed_that_overflowed_is_replaced_by_a_full_report_at_drain() { + let (feed, _rx) = feed::Feed::new(); + for _ in 0..=20_000 { + feed.send(&wire::Event::KeyDrop { seq: 0 }); + } + assert!(feed.lost()); + + let (old, new) = pair(); + let hold = old + .key_limits + .admit( + "k", + &key_limits(), + Ask { + tokens: 600, + cost_micros: 0, + }, + Duration::ZERO, + ) + .await + .unwrap(); + hold.bind(7); + let _ = old.handover_export(); + // 当成攒满了 + if let Some((f, _)) = old.handover.pending.lock().unwrap().as_ref() { + f.mark_lost(); + } + let _pass = old.gate.acquire("k", Some(1)).await; + let _fwd = forward(&old, &new); + settle().await; + assert_eq!(used(&new), 600, "排空开始时没报在跑的预留"); + assert!( + tokio::time::timeout(Duration::from_millis(50), new.gate.acquire("k", Some(1))) + .await + .is_err(), + "排空开始时没报占着的名额" + ); +} + +/// 快照来两份（旧 core 交出监听之前再导出一次，那时它已暂停续期）：**后一份的 OAuth 盖掉 +/// 前一份的**，新 core 用的是旧 core 最后手上的那个 access token +#[tokio::test] +async fn a_later_snapshot_replaces_the_oauth_tokens_of_an_earlier_one() { + let new = AppState::new(config(1, "at-1")).unwrap(); + // 第一份：旧 core 手上是 at-1 + let first = AppState::new(config(1, "at-1")).unwrap(); + first + .oauth + .token("甲", &oauth("at-1"), &first.http) + .await + .unwrap(); + new.handover_import(&over_the_wire(first.handover_export())); + // 第二份：之间旧 core 换到了 at-2（同一份配置、同一个 refresh token） + let second = AppState::new(config(1, "at-2")).unwrap(); + second + .oauth + .token("甲", &oauth("at-2"), &second.http) + .await + .unwrap(); + new.handover_import(&over_the_wire(second.handover_export())); + new.oauth.hold_renewals(true); + let (token, _) = new + .oauth + .token("甲", &oauth("at-1"), &new.http) + .await + .unwrap(); + assert_eq!(token, "at-2", "后一份快照没盖掉前一份"); +} diff --git a/crates/tw-gateway/src/handover/wire.rs b/crates/tw-gateway/src/handover/wire.rs new file mode 100644 index 00000000..0c098857 --- /dev/null +++ b/crates/tw-gateway/src/handover/wire.rs @@ -0,0 +1,274 @@ +//! 快照与结算的格式。**只在交接连接上走**，和 `tw_handover::VERSION` 同一个版本号，不做 +//! 跨版本兼容：改了这里任何一样都升那个号。 +//! +//! 时刻一律是墙上的时钟（Unix 毫秒）或者「还剩多久 / 过去多久」的毫秒数 —— `Instant` 出不了 +//! 这个进程。快照里的「还剩多久」从导出的那一刻（[`AT_MS`] 那一节）算起，接进来时扣掉路上 +//! 用掉的那一点。 + +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; + +use crate::key_limits::Recorded; +use crate::key_limits::books::{Amount, Period}; + +/// 快照里各节的名字。**一节读不懂只丢那一节**（见 [`crate::state::AppState::handover_import`]） +pub(crate) const AT_MS: &str = "at_ms"; +pub(crate) const AFFINITY: &str = "affinity"; +pub(crate) const SEALS: &str = "seals"; +pub(crate) const CACHE_MARKS: &str = "cache_marks"; +pub(crate) const SESSIONS: &str = "sessions"; +pub(crate) const HEALTH: &str = "health"; +pub(crate) const THROTTLES: &str = "throttles"; +pub(crate) const LATENCY: &str = "latency"; +pub(crate) const PACE: &str = "pace"; +pub(crate) const BALANCE: &str = "balance"; +pub(crate) const OAUTH: &str = "oauth"; +pub(crate) const KEY_LIMITS: &str = "key_limits"; +pub(crate) const UPSTREAM_LIMITS: &str = "upstream_limits"; +pub(crate) const GATE: &str = "gate"; +pub(crate) const SLOTS: &str = "slots"; + +// ------------------------------------------------------------ 快照的各节 + +/// 一段对话记着的（[`crate::affinity`]）。**这一轮沿用的路由决定不搬**：它认的是旧 core 那一份 +/// 引擎（配置换了就不沿用），换了进程一定对不上；同一轮留在原来那一家靠的是 `answer` +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct AffinityEntry { + pub key: String, + pub touched_ms: u64, + pub answer: Option, + pub holder: Option, + pub sealed: bool, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct AffinityAnswer { + pub group: Option, + pub provider: String, + pub turn: u32, + pub at_ms: u64, + pub cache: u64, +} + +/// 一段对话里一家拒过的别家封存的推理（[`crate::seal`]）。指纹写成十六进制 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct SealEntry { + pub conversation: String, + pub upstream: String, + pub last_ms: u64, + pub seals: Vec, +} + +/// 拒过自动标的缓存断点的（上游，模型）（[`crate::cache_marks`]） +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct CacheMark { + pub upstream: String, + pub model: String, + pub at_ms: u64, +} + +/// 一段对话此刻归在哪一次会话（[`crate::session`]）。不搬的话，交接时正在进行的任务在流量里 +/// 被切成两段 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct SessionEntry { + pub fp: String, + pub id: String, + pub last_ms: u64, +} + +/// 熔断与最近的成败（[`crate::health`]） +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub(crate) struct Health { + pub entries: Vec, + pub recent: Vec, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct HealthEntry { + pub name: String, + pub breaker: tw_breaker::Breaker, + pub trips: u32, + pub held_until_ms: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct HealthRecent { + pub name: String, + /// （时刻，成没成），先的在前 + pub samples: Vec<(i64, bool)>, +} + +/// 一家的限流暂停（[`crate::throttle`]）。**排着的队和探路的名额不搬**：那是旧 core 手上的请求 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct ThrottleEntry { + pub upstream: String, + /// 暂停还剩多久。没在暂停是 None + pub pause_left_ms: Option, + /// 最近一次被限流是多久之前 + pub limited_ago_ms: Option, + pub strikes: u32, + pub probing: bool, +} + +/// 一家的首段内容用时样本（[`crate::latency`]） +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct LatencyEntry { + pub upstream: String, + pub real: Vec, + pub seed: Option, +} + +/// 流里估数和实数之比（[`crate::pace`]） +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct PaceEntry { + pub upstream: String, + pub model: String, + pub ratio: f64, +} + +/// 一个 `load-balance` 组轮到哪儿了（[`crate::balance`]） +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct BalanceRound { + pub group: String, + pub members: Vec, + pub current: HashMap, +} + +/// 一家的 OAuth access token（[`crate::oauth`]）。**原文**：只在交接连接上走（同一个用户、 +/// 本机），和 config.yaml 里写着的是同一份东西 +#[derive(Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct OauthEntry { + pub upstream: String, + /// 生成它的那份配置（refresh、端点、client）的指纹，blake3。新 core 的配置对得上才接 + pub cfg: String, + pub access: String, + pub refresh: String, + /// 还有多久该去刷。服务器没说有效期是 None + pub renew_in_ms: Option, + /// 是多久之前从 token 端点换来的。从配置里读的是 None + pub age_ms: Option, + pub failed: Option, +} + +impl std::fmt::Debug for OauthEntry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OauthEntry") + .field("upstream", &self.upstream) + .field("renew_in_ms", &self.renew_in_ms) + .field("age_ms", &self.age_ms) + .field("failed", &self.failed) + .finish_non_exhaustive() + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct OauthFailed { + pub cause: crate::oauth::OauthError, + pub retry_in_ms: u64, +} + +/// 用量的账（密钥的、上游的同一个样子，见 [`crate::key_limits::books`]）：每个名字的天、周、 +/// 月和滚动窗口，和此刻在跑的预留 +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub(crate) struct LimitBooks { + pub books: Vec, + pub held: Vec, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct NamedBooks { + pub name: String, + pub periods: [Option; 3], + pub recent: Vec<(i64, Amount)>, +} + +/// 一份在跑的预留。`seq` 是**发出它的那个 core** 给的号，结算时认它 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct Held { + pub seq: u64, + pub key: String, + pub at_ms: i64, + pub tokens: i64, + pub cost: i64, +} + +// ------------------------------------------------------------ 排空期的结算 + +/// 排空期旧 core 发来的一条结算（`tw_handover::Message::Settled` 的 `data`）。 +/// +/// **一条是一样东西的一次变化**，不是一个请求一条：名额在请求结束（或者换下一家、排队时客户端 +/// 走了）的那一刻放掉，用量要等存储层记下那一行才知道 —— 两样不在同一刻，按请求凑成一条的话， +/// 要么名额压到用量出来才放，要么用量丢了。 +/// +/// - `*_holds`：旧 core 此刻占着的全部（排空开始时发一次），新 core 照它换掉手上记着的； +/// - `gate` / `slot`：一把密钥、一家上游的名额多占（`d` 是 1）或放掉（`d` 是 -1）了一个。密钥的 +/// 名额**连排着队等的也算**：它们是排在前面的，新 core 不该趁它们还在等时把位置给了别人； +/// - `key_*` / `up_*`：用量上限的预留占上、结算成实数、退回、过期作废。结算带着密钥、上游、 +/// token 用量和费用（[`Recorded`]）。 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "k", rename_all = "snake_case")] +pub(crate) enum Event { + GateHolds { + n: HashMap, + }, + Gate { + key: String, + d: i64, + }, + SlotHolds { + n: HashMap, + }, + Slot { + upstream: String, + d: i64, + }, + KeyHolds { + held: Vec, + }, + /// `windowed`：占上时滚动窗口里记了这一个请求 + KeyReserve { + held: Held, + windowed: bool, + }, + /// 存储层记下了一行。`seq` 是它占着的预留（没有预留的是 None：那把密钥没设上限） + KeySettle { + seq: Option, + at_ms: i64, + rec: Recorded, + }, + /// 预留整份退回，滚动窗口里那一个请求也撤掉（准入之后、开始之前就丢掉了） + KeyRelease { + seq: u64, + }, + /// 预留作废（请求结束了太久、存储层一直没交来那一行） + KeyDrop { + seq: u64, + }, + UpHolds { + held: Vec, + }, + UpReserve { + held: Held, + windowed: bool, + }, + UpSettle { + seqs: Vec, + at_ms: i64, + hops: Vec, + served: Option, + rec: Recorded, + }, + UpRelease { + seq: u64, + }, + UpDrop { + seq: u64, + }, +} + +impl Event { + pub(crate) fn to_value(&self) -> serde_json::Value { + serde_json::to_value(self).unwrap_or(serde_json::Value::Null) + } +} diff --git a/crates/tw-gateway/src/health.rs b/crates/tw-gateway/src/health.rs index f3560513..0bbe7caf 100644 --- a/crates/tw-gateway/src/health.rs +++ b/crates/tw-gateway/src/health.rs @@ -295,6 +295,61 @@ impl Health { } } + /// 交接：每家的熔断和最近的成败（见 [`crate::handover`]）。时刻都是墙上的，原样交 + pub(crate) fn handover_export(&self) -> crate::handover::wire::Health { + use crate::handover::wire::{HealthEntry, HealthRecent}; + let entries = self + .map + .lock() + .map(|m| { + m.iter() + .map(|(name, e)| HealthEntry { + name: name.clone(), + breaker: e.breaker, + trips: e.trips, + held_until_ms: e.held_until_ms, + }) + .collect() + }) + .unwrap_or_default(); + let recent = self + .recent + .lock() + .map(|r| { + r.0.iter() + .map(|(name, w)| HealthRecent { + name: name.clone(), + samples: w.iter().copied().collect(), + }) + .collect() + }) + .unwrap_or_default(); + crate::handover::wire::Health { entries, recent } + } + + /// 交接：接进旧 core 的熔断和成败。这里已经有记录的那一家不动 —— 接进来之前它已经在这里 + /// 成过败过，那是更新的事 + pub(crate) fn handover_import(&self, h: crate::handover::wire::Health) { + if let Ok(mut m) = self.map.lock() { + for e in h.entries { + m.entry(e.name).or_insert(Entry { + breaker: e.breaker, + trips: e.trips, + held_until_ms: e.held_until_ms, + }); + } + } + if let Ok(mut r) = self.recent.lock() { + for w in h.recent { + let mut samples: VecDeque<(i64, bool)> = w.samples.into_iter().collect(); + while samples.len() > RECENT_MAX { + samples.pop_front(); + } + r.0.entry(w.name).or_insert(samples); + } + } + } + /// 从候选里挑出还能用的，**并说明是不是 fail-open**。 /// /// 返回的第二个值为真时表示「全都熔断了，我们还是把原列表还给你」—— diff --git a/crates/tw-gateway/src/key_limits/books.rs b/crates/tw-gateway/src/key_limits/books.rs index 9d71324e..9a1c9aa4 100644 --- a/crates/tw-gateway/src/key_limits/books.rs +++ b/crates/tw-gateway/src/key_limits/books.rs @@ -25,7 +25,7 @@ pub(crate) fn calendar_index(per: LimitPer) -> usize { } /// 一份用量。四样分开记，各条上限各取各的。 -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub(crate) struct Amount { pub requests: i64, /// 没走缓存的输入 + 写进缓存的 + 输出 @@ -53,7 +53,7 @@ impl Amount { } /// 一期（天、周、月）的合计。 -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize)] pub(crate) struct Period { pub start: i64, pub end: i64, @@ -347,3 +347,85 @@ pub(crate) fn take_moved( .filter(|m| asked.insert(m.clone())) .collect() } + +// ------------------------------------------------------------ 交接（见 `crate::handover`） + +/// 一份预留写成交接连接上的样子。`seq` 是这里给它的号 +pub(crate) fn held_wire(seq: u64, r: &Reservation) -> crate::handover::wire::Held { + crate::handover::wire::Held { + seq, + key: r.key.clone(), + at_ms: r.at_ms, + tokens: r.tokens, + cost: r.cost, + } +} + +/// 旧 core 的一份预留，在这里替它记着。**没有请求号**：它不会被当成「结束了太久」放掉（这里 +/// 的事件总线上没有那个请求），只等旧 core 的结算 +pub(crate) fn held_reservation(h: &crate::handover::wire::Held) -> Reservation { + Reservation { + key: h.key.clone(), + at_ms: h.at_ms, + tokens: h.tokens, + cost: h.cost.max(0), + request: None, + closed_since: None, + } +} + +/// 在滚动窗口里记上 `at_ms` 那一刻的一个请求（旧 core 记的那一个，见 [`record_request`]） +pub(crate) fn record_at(books: &mut HashMap, key: &str, at_ms: i64) { + books.entry(key.to_string()).or_default().recent.push_back(( + at_ms, + Amount { + requests: 1, + ..Default::default() + }, + )); +} + +/// 账和预留写成交接连接上的样子 +pub(crate) fn export( + books: &HashMap, + held: &HashMap, +) -> crate::handover::wire::LimitBooks { + let mut held: Vec<_> = held.iter().map(|(s, r)| held_wire(*s, r)).collect(); + held.sort_by_key(|h| h.seq); + crate::handover::wire::LimitBooks { + books: books + .iter() + .map(|(name, b)| crate::handover::wire::NamedBooks { + name: name.clone(), + periods: b.periods, + recent: b.recent.iter().copied().collect(), + }) + .collect(), + held, + } +} + +/// 把旧 core 的账并进来。**并，不是换**：这里已经记着的（多半什么都没有 —— 新 core 接进快照 +/// 时还没转发过）照样算。同一期的合计相加，滚动窗口里的每一笔按时刻排进去；开头不一样的 +/// 一期（交接正好跨过零点）留这里的 +pub(crate) fn import( + books: &mut HashMap, + theirs: Vec, +) { + for nb in theirs { + let b = books.entry(nb.name).or_default(); + for (mine, theirs) in b.periods.iter_mut().zip(nb.periods) { + let Some(theirs) = theirs else { continue }; + match mine { + Some(p) if p.start == theirs.start => p.sum.add(&theirs.sum), + Some(_) => {} + None => *mine = Some(theirs), + } + } + if !nb.recent.is_empty() { + let mut all: Vec<(i64, Amount)> = b.recent.drain(..).chain(nb.recent).collect(); + all.sort_by_key(|(t, _)| *t); + b.recent = all.into(); + } + } +} diff --git a/crates/tw-gateway/src/key_limits/mod.rs b/crates/tw-gateway/src/key_limits/mod.rs index 16016ed0..48d5c93a 100644 --- a/crates/tw-gateway/src/key_limits/mod.rs +++ b/crates/tw-gateway/src/key_limits/mod.rs @@ -101,6 +101,12 @@ struct Inner { /// 已经要过、还没读回来的几期。**一期只要一次**：读回来之前每个请求都会看到同一个开头， /// 不该每个请求去库里查一遍 asked: HashSet, + /// 交接：预留和结算往哪儿报（旧 core，见 [`crate::handover`]）。**在这把锁里报**：快照和之后 + /// 的变化前后接得上 + feed: Option>, + /// 交接：替旧 core 记着的预留，旧 core 那边的号 → 这里的号（新 core）。它们没有请求号，不会 + /// 被当成「结束了太久」放掉，只等旧 core 的结算 + imported: HashMap, } /// 准入时一个请求要占的：输入 token 的估算，和按头一个候选算的输入费用。 @@ -111,7 +117,7 @@ pub struct Ask { } /// 存储层记下的一行（结算），或者库里加起来的一组（重建）：算不算、算多少，看的是这几样。 -#[derive(Debug, Clone, Default, PartialEq)] +#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)] pub struct Recorded { /// 密钥的名字 pub client: String, @@ -541,44 +547,68 @@ impl KeyLimits { let now = self.clock.now_ms(); let (events, moved) = { let mut g = self.lock(); - let held = g.by_request.remove(&id).and_then(|s| g.held.remove(&s)); - if !rec.counts() { - if let Some(r) = held { - books::unrecord(&mut g.books, &r); - } - return; - } - let key = held.map_or_else(|| rec.client.clone(), |r| r.key); - let amount = rec.amount(); - let rolling = g - .limits - .get(&key) - .and_then(|ls| longest_window(ls)) - .is_some(); - let clock = self.clock.as_ref(); - let books = g.books.entry(key.clone()).or_default(); - for per in CALENDAR { - let p = books.period(clock, per, now); - if at_ms >= p.start && at_ms < p.end { - p.sum.add(&amount); - } - } - // 请求数在准入时已经进了滚动窗口 - if rolling { - books.recent.push_back(( - now, - Amount { - requests: 0, - ..amount - }, - )); - } - (self.alerts(&mut g, &key, now), self.moved(&mut g, &key)) + let held = g + .by_request + .remove(&id) + .and_then(|s| g.held.remove(&s).map(|r| (s, r))); + self.settle_locked(&mut g, held, at_ms, rec, now) }; self.emit(events); self.ask_reread(moved); } + /// 结算一行：`held` 是它占着的预留和预留的号。交回要报的提醒和换了新一期的那几期 + fn settle_locked( + &self, + g: &mut Inner, + held: Option<(u64, Reservation)>, + at_ms: i64, + rec: &Recorded, + now: i64, + ) -> (Vec, Vec) { + if held.is_some() || rec.counts() { + let seq = held.as_ref().map(|(s, _)| *s); + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::KeySettle { + seq, + at_ms, + rec: rec.clone(), + }); + } + let held = held.map(|(_, r)| r); + if !rec.counts() { + if let Some(r) = held { + books::unrecord(&mut g.books, &r); + } + return (Vec::new(), Vec::new()); + } + let key = held.map_or_else(|| rec.client.clone(), |r| r.key); + let amount = rec.amount(); + let rolling = g + .limits + .get(&key) + .and_then(|ls| longest_window(ls)) + .is_some(); + let clock = self.clock.as_ref(); + let books = g.books.entry(key.clone()).or_default(); + for per in CALENDAR { + let p = books.period(clock, per, now); + if at_ms >= p.start && at_ms < p.end { + p.sum.add(&amount); + } + } + // 请求数在准入时已经进了滚动窗口 + if rolling { + books.recent.push_back(( + now, + Amount { + requests: 0, + ..amount + }, + )); + } + (self.alerts(g, &key, now), self.moved(g, &key)) + } + /// 重启之后把天、周、月的数从请求记录里加回来。`since(t)` 给出从 `t` 起每把密钥 /// 记下的那些（存储层的 `key_usage_since`）。分钟、小时的窗口从空的开始。 /// @@ -650,6 +680,130 @@ impl KeyLimits { view } + // ------------------------------------------------------------ 交接 + + /// 旧 core：此刻的账和在跑的预留（替更早的 core 记着的也在里面），之后的变化报给 `feed`。 + /// 两件事在同一把锁里做：快照之后的每一次变化都在 `feed` 上，一次不漏、一次不重 + pub(crate) fn handover_export( + &self, + feed: Option>, + ) -> crate::handover::wire::LimitBooks { + let mut g = self.lock(); + let out = books::export(&g.books, &g.held); + g.feed = feed; + out + } + + /// 旧 core：排空开始。报一份此刻在跑的全部预留，之后的变化接着报 + pub(crate) fn handover_drain(&self, feed: Arc) { + let mut g = self.lock(); + feed.send(&crate::handover::wire::Event::KeyHolds { + held: books::export(&HashMap::new(), &g.held).held, + }); + g.feed = Some(feed); + } + + /// 新 core：接进旧 core 的账。天、周、月和滚动窗口并进来，在跑的预留记成替它占着的 + pub(crate) fn handover_import(&self, b: crate::handover::wire::LimitBooks) { + let mut g = self.lock(); + let g = &mut *g; + books::import(&mut g.books, b.books); + for h in &b.held { + Self::phantom(g, h, false); + } + } + + /// 替旧 core 记上一份预留。`windowed`：旧 core 占上它时滚动窗口里记了一个请求，这里也记上 + fn phantom(g: &mut Inner, h: &crate::handover::wire::Held, windowed: bool) { + if g.imported.contains_key(&h.seq) { + return; + } + g.seq += 1; + let seq = g.seq; + let r = books::held_reservation(h); + if windowed { + books::record_at(&mut g.books, &r.key, r.at_ms); + } + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::KeyReserve { + held: books::held_wire(seq, &r), + windowed, + }); + g.held.insert(seq, r); + g.imported.insert(h.seq, seq); + } + + /// 替旧 core 记着的这一份不用再记了：作废，不撤滚动窗口里那一个请求（它确实发生过） + fn drop_phantom(g: &mut Inner, theirs: u64) { + if let Some(seq) = g.imported.remove(&theirs) + && g.held.remove(&seq).is_some() + { + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::KeyDrop { seq }); + } + } + + /// 新 core：旧 core 报来一条变化（[`crate::handover::wire::Event`] 里 `key_*` 那几种）。别的 + /// 不认，交回 false + pub(crate) fn handover_apply(&self, ev: &crate::handover::wire::Event) -> bool { + use crate::handover::wire::Event; + let now = self.clock.now_ms(); + let (events, moved) = { + let mut g = self.lock(); + let g = &mut *g; + match ev { + Event::KeyHolds { held } => { + let want: HashSet = held.iter().map(|h| h.seq).collect(); + let stale: Vec = g + .imported + .keys() + .filter(|s| !want.contains(s)) + .copied() + .collect(); + for s in stale { + Self::drop_phantom(g, s); + } + for h in held { + Self::phantom(g, h, false); + } + (Vec::new(), Vec::new()) + } + Event::KeyReserve { held, windowed } => { + Self::phantom(g, held, *windowed); + (Vec::new(), Vec::new()) + } + Event::KeySettle { seq, at_ms, rec } => { + let held = seq + .and_then(|s| g.imported.remove(&s)) + .and_then(|s| g.held.remove(&s).map(|r| (s, r))); + self.settle_locked(g, held, *at_ms, rec, now) + } + Event::KeyRelease { seq } => { + if let Some(s) = g.imported.remove(seq) { + Self::release_locked(g, s); + } + (Vec::new(), Vec::new()) + } + Event::KeyDrop { seq } => { + Self::drop_phantom(g, *seq); + (Vec::new(), Vec::new()) + } + _ => return false, + } + }; + self.emit(events); + self.ask_reread(moved); + true + } + + /// 新 core：交接连接断了。替旧 core 记着的预留全部作废；它的请求用了多少，等请求库打开后 + /// 从记录里重建 + pub(crate) fn handover_release_all(&self) { + let mut g = self.lock(); + let theirs: Vec = g.imported.keys().copied().collect(); + for s in theirs { + Self::drop_phantom(&mut g, s); + } + } + // ------------------------------------------------------------ 锁里面的 /// 一条上限此刻用了多少：记下的，加上在跑的请求的预留 @@ -714,8 +868,12 @@ impl KeyLimits { books::record_request(&mut g.books, key, limits, now); g.seq += 1; let seq = g.seq; - g.held - .insert(seq, Reservation::new(key, now, ask.tokens, ask.cost_micros)); + let r = Reservation::new(key, now, ask.tokens, ask.cost_micros); + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::KeyReserve { + held: books::held_wire(seq, &r), + windowed: longest_window(limits).is_some(), + }); + g.held.insert(seq, r); seq } @@ -731,19 +889,27 @@ impl KeyLimits { /// 滚动窗口里的那一个请求都还回去 fn release(&self, seq: u64) { let mut g = self.lock(); + Self::release_locked(&mut g, seq); + } + + fn release_locked(g: &mut Inner, seq: u64) { if let Some(r) = g.held.remove(&seq) { if let Some(id) = r.request { g.by_request.remove(&id); } books::unrecord(&mut g.books, &r); + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::KeyRelease { + seq, + }); } } /// 放掉结束了太久、存储层却一直没交来那一行的预留（见 [`books::GRACE_MS`]）。顺手把滚动 /// 窗口里滑出去的扔掉 fn sweep(&self, g: &mut Inner, now: i64) { - for (_, id) in books::sweep_held(&mut g.held, &self.bus, now) { + for (seq, id) in books::sweep_held(&mut g.held, &self.bus, now) { g.by_request.remove(&id); + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::KeyDrop { seq }); } books::trim_windows(&mut g.books, &g.limits, now); } diff --git a/crates/tw-gateway/src/latency.rs b/crates/tw-gateway/src/latency.rs index e35d6bc8..5c8c1185 100644 --- a/crates/tw-gateway/src/latency.rs +++ b/crates/tw-gateway/src/latency.rs @@ -136,6 +136,31 @@ impl Latency { self.batch(names, Window::measured) } + /// 交接：每家的样本和垫的底（见 [`crate::handover`]） + pub(crate) fn handover_export(&self) -> Vec { + let g = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + g.iter() + .map(|(name, w)| crate::handover::wire::LatencyEntry { + upstream: name.clone(), + real: w.real.iter().copied().collect(), + seed: w.seed, + }) + .collect() + } + + /// 交接：接进旧 core 的样本。这里已经有样本的那一家不动 + pub(crate) fn handover_import(&self, entries: Vec) { + let mut g = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + for e in entries { + let mut real: VecDeque = e.real.into_iter().collect(); + while real.len() > WINDOW { + real.pop_front(); + } + let seed = e.seed.filter(|_| real.len() < MIN_SAMPLES); + g.entry(e.upstream).or_insert(Window { real, seed }); + } + } + fn batch(&self, names: &[String], of: fn(&Window) -> Option) -> HashMap { let g = self.inner.lock().expect("lock not poisoned"); names diff --git a/crates/tw-gateway/src/lib.rs b/crates/tw-gateway/src/lib.rs index ec4370e2..9f3a85b6 100644 --- a/crates/tw-gateway/src/lib.rs +++ b/crates/tw-gateway/src/lib.rs @@ -28,6 +28,7 @@ pub mod fixture; pub mod forward; pub mod glm; pub mod guard; +pub mod handover; pub mod health; pub mod hint; pub mod inflate; diff --git a/crates/tw-gateway/src/limits.rs b/crates/tw-gateway/src/limits.rs index 4d3bb522..581d4753 100644 --- a/crates/tw-gateway/src/limits.rs +++ b/crates/tw-gateway/src/limits.rs @@ -20,6 +20,8 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore}; #[derive(Default)] pub struct Gate { per_key: Mutex>>, + /// 每把密钥此刻占着、等着的通行证有几张（交接时报给接替的 core，见 [`crate::handover`]） + pub(crate) claims: Arc, } /// 一把密钥的闸。 @@ -38,18 +40,42 @@ struct Books { /// 拿到的通行证。**Drop 时自动归还** —— 手工释放迟早会在某条错误路径上 /// 漏掉，而漏掉的表现是并发数只减不增，最后这把密钥的请求一起卡死。 pub struct Pass { - held: Option<(Arc, OwnedSemaphorePermit)>, + held: Option<(Arc, Hand)>, + /// 交接时数的那一张（见 [`crate::handover::feed::Claims`]）。排着队时就占上了 + _claim: Option, +} + +/// 通行证手里拿着的：信号量的一个位置，或者一笔欠账。 +/// +/// **欠账只有替旧 core 占的名额才有**（[`Gate::phantom`]）：交接时旧 core 占着的比这里此刻空着 +/// 的还多（它那边调小过上限、排着队的也算），多出来的那几个记成欠着 —— 和调小上限时还没收回的 +/// 那几个（`owed`）是同一本账：之后谁还回一个位置，先抵掉欠账。 +enum Hand { + Permit(OwnedSemaphorePermit), + Debt, } impl Drop for Pass { fn drop(&mut self) { - let Some((key, permit)) = self.held.take() else { + let Some((key, hand)) = self.held.take() else { return; }; let mut books = key.books.lock().unwrap(); - if books.owed > 0 { - books.owed -= 1; - permit.forget(); + match hand { + Hand::Permit(permit) => { + if books.owed > 0 { + books.owed -= 1; + permit.forget(); + } + } + // 欠着的那一个不用还了：还没抵掉就勾销，抵掉过（有人替它少还了一个位置）就补回去 + Hand::Debt => { + if books.owed > 0 { + books.owed -= 1; + } else { + key.sem.add_permits(1); + } + } } } } @@ -59,8 +85,28 @@ impl Gate { /// 没设就不等。 pub async fn acquire(&self, key: &str, limit: Option) -> Pass { let Some(limit) = limit else { - return Pass { held: None }; + return Pass { + held: None, + _claim: None, + }; }; + let k = self.key(key, limit); + // **排着队就算占上了**：交接时它排在接替的 core 的新请求前面（见 [`crate::handover`]）。 + // 客户端走了、这一份等待被丢掉时跟着放掉 + let claim = self.claims.begin(key); + let permit = k + .sem + .clone() + .acquire_owned() + .await + .expect("the semaphore is never closed"); + Pass { + held: Some((k, Hand::Permit(permit))), + _claim: Some(claim), + } + } + + fn key(&self, key: &str, limit: usize) -> Arc { let k = self .per_key .lock() @@ -74,14 +120,25 @@ impl Gate { }) .clone(); k.resize(limit); - let permit = k - .sem - .clone() - .acquire_owned() - .await - .expect("the semaphore is never closed"); + k + } + + /// 替交接前的旧 core 占一张（见 [`crate::handover`]）：**不等**。有空位就占一个，没有就记一笔 + /// 欠账 —— 旧 core 那个请求此刻确实在跑，这里新来的得等它结束，哪怕这样一来这一阵记着的 + /// 比上限还多。丢掉就放掉，和别的通行证一样 + pub(crate) fn phantom(&self, key: &str, limit: usize) -> Pass { + let k = self.key(key, limit); + let claim = self.claims.begin(key); + let hand = match k.sem.clone().try_acquire_owned() { + Ok(permit) => Hand::Permit(permit), + Err(_) => { + k.books.lock().unwrap().owed += 1; + Hand::Debt + } + }; Pass { - held: Some((k, permit)), + held: Some((k, hand)), + _claim: Some(claim), } } } diff --git a/crates/tw-gateway/src/listen.rs b/crates/tw-gateway/src/listen.rs index f5c1c923..d558b51f 100644 --- a/crates/tw-gateway/src/listen.rs +++ b/crates/tw-gateway/src/listen.rs @@ -23,10 +23,25 @@ //! //! 界面上的地址以前是启动时记的一次。换了端口，网关已经在新端口上服务， //! 界面还写着旧的 —— 用户分不清是没生效还是没刷新。 +//! +//! # 交接（升级不断线） +//! +//! 新旧 core 交接时，监听的套接字原样交给新的那个（见 `tw-handover`），**不是重新绑**。 +//! 这一边要做的三件事都在 [`Handoff`] 上：交出此刻在听的每一个（[`Handoff::export`]， +//! 复制出来的句柄）、从交过来的那些起服务（[`serve_with`]）、以及旧的那个停止 accept +//! （[`Handoff::stop_accepting`]）。 +//! +//! **停止 accept 只丢自己的句柄。**交出去的是同一个内核套接字：axum 丢下监听器关的 +//! 只是这个进程里的那个描述符，新的那个手里的照样在听，内核队列里还没被接受的连接 +//! 由它取走。unix 上一个套接字要等最后一个描述符关掉才关；Windows 上 +//! `WSADuplicateSocketW` 复制出来的也是一样（最后一个 `closesocket` 才真正关）。 use std::net::SocketAddr; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::net::TcpListener; +use tokio::sync::watch; use tw_types::{Msg, msg}; use crate::server::{now_ms, router}; @@ -118,12 +133,148 @@ fn overlaps_ours(a: SocketAddr, held: &[SocketAddr]) -> bool { }) } +/// 复制一个监听器的句柄：**同一个内核套接字**，多一个描述符。交接时交出去的是它， +/// 自己手里的那个照样 accept。 +pub fn dup_listener(l: &TcpListener) -> std::io::Result { + #[cfg(unix)] + { + use std::os::fd::AsFd; + Ok(std::net::TcpListener::from(l.as_fd().try_clone_to_owned()?)) + } + #[cfg(windows)] + { + use std::os::windows::io::AsSocket; + Ok(std::net::TcpListener::from( + l.as_socket().try_clone_to_owned()?, + )) + } +} + +/// 从交过来的套接字起一个 tokio 的监听器。**先设成非阻塞**：交过来的那一头不保证。 +pub fn adopt_listener(l: std::net::TcpListener) -> std::io::Result { + l.set_nonblocking(true)?; + TcpListener::from_std(l) +} + +/// 交接时监听那一边听谁的：照常跟着配置走、先别动、还是停下。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Mode { + Follow, + /// 新 core 还没提交：从交过来的监听上 accept，但**不跟配置换监听** —— 提交之前 + /// 旧的那个也还在这些套接字上 accept，这时候换掉、让出端口都不对 + Hold, + /// 旧 core 提交之后：全部停止 accept，**只丢自己的句柄** + Stop, +} + +/// 交接用的把手：交出监听、停止 accept、等起来。[`serve_with`] 拿一份，交接那一边拿 +/// 另一份（克隆出来的是同一个）。 +#[derive(Clone)] +pub struct Handoff(Arc); + +struct HandoffInner { + mode: watch::Sender, + /// 此刻在听的：配置算出来的地址，和复制出来的句柄。换监听时跟着换。 + /// + /// **弱引用**：句柄归 [`Bound`]。这里要是攥着一份，退场的监听器就关不掉 —— 内核 + /// 套接字还有一个描述符开着，端口一直占着，「旧的先让出来」那一步就永远绑不上 + held: std::sync::Mutex)>>, + /// 起始那一套起来了 + up: watch::Sender, + /// 还没结束的 axum 服务（停止 accept 之后，它要等手上的连接都结束才结束） + serving: Arc, +} + +impl Default for Handoff { + fn default() -> Self { + Self(Arc::new(HandoffInner { + mode: watch::channel(Mode::Follow).0, + held: Default::default(), + up: watch::channel(false).0, + serving: Default::default(), + })) + } +} + +impl Handoff { + /// 先别跟配置换监听（新 core 提交之前）。**在 [`serve_with`] 之前调。** + pub fn hold(&self) { + self.0.mode.send_replace(Mode::Hold); + } + + /// 从现在起跟着配置走。交过来的地址和配置里的不一样时，**马上**按现有的换监听 + /// 规矩换过去（新 core 提交之后调）。 + pub fn follow(&self) { + self.0.mode.send_if_modified(|m| { + let was = *m == Mode::Hold; + if was { + *m = Mode::Follow; + } + was + }); + } + + /// 全部停止 accept（旧 core 提交之后调）。在跑的连接照常跑完；**共享的套接字不关** + /// （见模块文档）。之后 [`serve_with`] 返回 `Ok(())`，不再换监听。 + pub fn stop_accepting(&self) { + self.0.mode.send_replace(Mode::Stop); + } + + /// 此刻在听的每一个：配置算出来的地址，和一个复制出来的句柄（交给新 core）。 + pub fn export(&self) -> std::io::Result> { + let held = self.0.held.lock().map(|g| g.clone()).unwrap_or_default(); + held.iter() + .filter_map(|(want, l)| l.upgrade().map(|l| (*want, l))) + .map(|(want, l)| Ok((want, l.try_clone()?))) + .collect() + } + + /// 等到起始那一套起来了。 + pub async fn listening(&self) { + let mut rx = self.0.up.subscribe(); + let _ = rx.wait_for(|up| *up).await; + } + + /// 还没结束的服务有几个。停止 accept 之后它们各自等手上的连接结束；**归零 = 走 HTTP + /// 的连接都结束了**（WebSocket 升级之后不归它们管，看网关的在途计数）。 + pub fn serving(&self) -> usize { + self.0.serving.load(Ordering::SeqCst) + } + + fn publish(&self, bound: &[Bound]) { + if let Ok(mut g) = self.0.held.lock() { + *g = bound + .iter() + .map(|b| (b.want, Arc::downgrade(&b.dup))) + .collect(); + } + } +} + +/// 一个服务在跑，结束时减一。 +struct Serving(Arc); + +impl Serving { + fn enter(n: &Arc) -> Self { + n.fetch_add(1, Ordering::SeqCst); + Self(n.clone()) + } +} + +impl Drop for Serving { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } +} + /// 一个地址上的服务。 struct Bound { /// 配置算出来的地址。比对「变没变」用它 —— 端口写 0 时实际拿到的每次都不一样 want: SocketAddr, /// 真的绑上的 actual: SocketAddr, + /// 复制出来的句柄，交接时交出去（见 [`Handoff::export`]） + dup: Arc, /// 发一下（或者丢掉）就停止接新连接，已经在跑的请求自己跑完 stop: tokio::sync::oneshot::Sender<()>, /// 端口真正放掉的那一刻：axum 丢下了监听器 @@ -175,13 +326,21 @@ async fn bind_all( Ok(out) } -fn start(state: &AppState, want: SocketAddr, listener: TcpListener) -> std::io::Result { +fn start( + state: &AppState, + handoff: &Handoff, + want: SocketAddr, + listener: TcpListener, +) -> std::io::Result { let actual = listener.local_addr()?; + let dup = Arc::new(dup_listener(&listener)?); tracing::info!(%actual, "the gateway is listening"); let (stop, stopped) = tokio::sync::oneshot::channel::<()>(); let (release, released) = tokio::sync::oneshot::channel::<()>(); let st = state.clone(); + let serving = Serving::enter(&handoff.0.serving); tokio::spawn(async move { + let _serving = serving; use axum::serve::ListenerExt; let listener = Tracked { inner: listener, @@ -208,11 +367,29 @@ fn start(state: &AppState, want: SocketAddr, listener: TcpListener) -> std::io:: Ok(Bound { want, actual, + dup, stop, released, }) } +/// 记下此刻在听的（界面看的那份，和交接要交出去的那份）。 +fn publish(state: &AppState, handoff: &Handoff, bound: &[Bound], error: Option) { + handoff.publish(bound); + state.set_listening(snapshot(bound, error)); +} + +/// 全部停止 accept。**只丢自己的句柄**（见模块文档）；在跑的连接各自跑完。 +fn stop_all(handoff: &Handoff, bound: Vec) { + for b in bound { + tracing::info!(addr = %b.actual, "handed over; no longer accepting here"); + let _ = b.stop.send(()); + } + if let Ok(mut g) = handoff.0.held.lock() { + g.clear(); + } +} + fn snapshot(bound: &[Bound], error: Option) -> Listening { Listening { addrs: bound.iter().map(|b| b.actual).collect(), @@ -228,7 +405,7 @@ pub(crate) async fn serve_detached( let listener = TcpListener::bind(want) .await .map_err(|e| std::io::Error::new(e.kind(), format!("{}", bind_failure(want, &e))))?; - let bound = start(&state, want, listener)?; + let bound = start(&state, &Handoff::default(), want, listener)?; let actual = bound.actual; state.set_listening(snapshot(std::slice::from_ref(&bound), None)); // 丢掉 `Bound` 就停止接新连接：把它留在一个不会结束的任务里 @@ -286,11 +463,37 @@ pub async fn serve_at( listen: tw_config::GatewayListen, follow: bool, ) -> std::io::Result<()> { - let (want, mut note) = plan(&listen); - let mut bound = Vec::with_capacity(want.len()); + serve_with(state, listen, follow, Handoff::default(), Vec::new()).await +} + +/// 同 [`serve_at`]，多两样交接用的东西： +/// +/// - `handoff`：交出监听、停止 accept（旧 core），先别跟配置换监听（新 core 提交之前， +/// 见 [`Handoff::hold`]）； +/// - `inherited`：旧 core 交过来的监听（配置算出来的地址，和套接字）。**不空就不绑**， +/// 直接在它们上面 accept；地址和这边的配置不一样的，跟配置走时（[`Handoff::follow`]） +/// 按现有的换监听规矩换过去。 +/// +/// [`Handoff::stop_accepting`] 之后返回 `Ok(())`。 +pub async fn serve_with( + state: AppState, + listen: tw_config::GatewayListen, + follow: bool, + handoff: Handoff, + inherited: Vec<(SocketAddr, std::net::TcpListener)>, +) -> std::io::Result<()> { + let (want, mut note) = if inherited.is_empty() { + plan(&listen) + } else { + (Vec::new(), None) + }; + let mut bound = Vec::with_capacity(want.len().max(inherited.len())); + for (a, l) in inherited { + bound.push(start(&state, &handoff, a, adopt_listener(l)?)?); + } for a in want { match TcpListener::bind(a).await { - Ok(l) => bound.push(start(&state, a, l)?), + Ok(l) => bound.push(start(&state, &handoff, a, l)?), Err(e) if absent(a, &e) => { tracing::warn!(%a, %e, "not an address of this machine yet; listening without it until it is"); note = Some(bind_failure(a, &e)); @@ -306,7 +509,30 @@ pub async fn serve_at( if let Some(e) = ¬e { tracing::warn!(%e, "listening on loopback only for now; the rest is added once it appears"); } - state.set_listening(snapshot(&bound, note)); + publish(&state, &handoff, &bound, note); + handoff.0.up.send_replace(true); + // 提交之前先别动（见 `Mode::Hold`）；停下就停下。跟配置走的那一刻马上对一遍： + // 交过来的地址可能和这边的配置不一样 + let mut mode = handoff.0.mode.subscribe(); + let mut held = false; + loop { + let now = *mode.borrow_and_update(); + match now { + Mode::Follow => break, + Mode::Stop => { + stop_all(&handoff, bound); + return Ok(()); + } + Mode::Hold => held = true, + } + if mode.changed().await.is_err() { + stop_all(&handoff, bound); + return Ok(()); + } + } + if held { + state.relisten(); + } // 眼下这些监听器是照着哪份设置绑的。**设置没变、网卡却解析不出来了**，是 // 我们跟着的那张网卡没了（WSL 关了）：退回只听回环，等它回来。设置变了 // 而新的解析不出来，那是另一回事，见下面 @@ -321,6 +547,13 @@ pub async fn serve_at( let asked = tokio::select! { _ = state.relisten_signal().notified() => true, _ = recheck.tick() => false, + r = mode.changed() => { + if r.is_err() || *mode.borrow_and_update() == Mode::Stop { + stop_all(&handoff, bound); + return Ok(()); + } + continue; + } }; let listen = if follow { state.runtime().config.listen.gateway.clone() @@ -349,7 +582,7 @@ pub async fn serve_at( if asked { tracing::error!(%e, "the new listen address cannot be resolved; keeping the current one"); } - state.set_listening(snapshot(&bound, Some(unresolved(&e)))); + publish(&state, &handoff, &bound, Some(unresolved(&e))); continue; } }; @@ -357,7 +590,7 @@ pub async fn serve_at( if want == have { // 通知来了但地址没变（比如又改回去了）。上一次没换成的那句话到此作废 following = spec; - state.set_listening(snapshot(&bound, note)); + publish(&state, &handoff, &bound, note); continue; } // **还要的地址原样留着，只绑新增的** —— 从「仅本机」换到「局域网」时， @@ -395,14 +628,14 @@ pub async fn serve_at( // 能听的照听，`addrs` 说的就是此刻真在听的 for w in back { match TcpListener::bind(w).await { - Ok(l) => bound.push(start(&state, w, l)?), + Ok(l) => bound.push(start(&state, &handoff, w, l)?), Err(e) => { tracing::error!(%w, %e, "the previous listen address could not be taken back") } } } bound.sort_by_key(|b| have.iter().position(|w| *w == b.want)); - state.set_listening(snapshot(&bound, Some(bind_failure(a, &e)))); + publish(&state, &handoff, &bound, Some(bind_failure(a, &e))); continue; } } @@ -415,18 +648,18 @@ pub async fn serve_at( } bound = keep.into_iter().chain(gone).collect(); bound.sort_by_key(|b| have.iter().position(|w| *w == b.want)); - state.set_listening(snapshot(&bound, Some(bind_failure(a, &e)))); + publish(&state, &handoff, &bound, Some(bind_failure(a, &e))); continue; } }; bound = keep; for (a, l) in listeners { - bound.push(start(&state, a, l)?); + bound.push(start(&state, &handoff, a, l)?); } following = spec; // 配置里写的那个排在最后，和 `addrs()` 同一个顺序 bound.sort_by_key(|b| want.iter().position(|w| *w == b.want)); - state.set_listening(snapshot(&bound, note)); + publish(&state, &handoff, &bound, note); } } diff --git a/crates/tw-gateway/src/models.rs b/crates/tw-gateway/src/models.rs index 2c289297..3ea843be 100644 --- a/crates/tw-gateway/src/models.rs +++ b/crates/tw-gateway/src/models.rs @@ -558,10 +558,16 @@ pub async fn refresh_one(state: &AppState, name: &str) -> Option { /// /// **`url-test` 的垫底跟着它一起跑**：两者都是「启动时打一次网络、之后靠 /// 真实流量」，而且都不该挡住启动。 +/// +/// **排空开始就停**（交接之后的旧 core）：探测和测速由接替者做，两边各打一遍只是重复的网络 pub fn spawn(state: AppState) { tokio::spawn(async move { let mut seeded_at: Option = None; + let mut draining = state.handover.draining(); loop { + if *draining.borrow() { + return; + } let now = now_ms(); if seeded_at.is_none_or(|t| now.saturating_sub(t) >= REFRESH_EVERY.as_millis() as u64) { crate::latency::seed_url_test(&state).await; @@ -574,6 +580,7 @@ pub fn spawn(state: AppState) { tokio::select! { _ = tokio::time::sleep(TICK) => {} _ = state.models.wake.notified() => {} + _ = draining.changed() => {} } } }); diff --git a/crates/tw-gateway/src/oauth.rs b/crates/tw-gateway/src/oauth.rs index 20ad938e..623f4a0b 100644 --- a/crates/tw-gateway/src/oauth.rs +++ b/crates/tw-gateway/src/oauth.rs @@ -53,7 +53,7 @@ const EXPIRED_BACKOFF: Duration = Duration::from_secs(24 * 3600); /// 换 access token 失败了。**英文只写一遍**：`Display` 就是 [`OauthError::msg`] 的原句， /// 界面拿码去翻；端点回的正文（已打码）和连接库的原话留在参数里。 -#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error, serde::Serialize, serde::Deserialize)] pub enum OauthError { /// 连不上 token 端点。`detail` 是连接库的原话 #[error("{}", self.msg())] @@ -69,6 +69,9 @@ pub enum OauthError { /// refresh token 已经作废。**重试没有用**，要重新登录或换一个 refresh token #[error("{}", self.msg())] Expired { status: u16, body: String }, + /// 交接期间暂停续期（见 [`Cache::hold_renewals`]），手上又没有能用的 access token + #[error("{}", self.msg())] + Held, } impl OauthError { @@ -91,6 +94,11 @@ impl OauthError { OauthError::NoToken => msg!( "gw.oauth.no_token" => "The token endpoint's response has no access_token." ), + OauthError::Held => msg!( + "gw.oauth.held" => + "The sign-in token cannot be renewed while the gateway is being updated. \ + Retry the request in a moment." + ), OauthError::Expired { status, body } => msg!( "gw.oauth.expired", status = status, body = body => "The OAuth credential has expired; sign in again or replace the refresh token. \ @@ -113,7 +121,9 @@ impl OauthError { endpoint's response has no access_token." ), // 不会退避在这两种上：作废的另记，退避不套退避 - OauthError::Backoff { .. } | OauthError::Expired { .. } => cause.msg(), + OauthError::Backoff { .. } | OauthError::Expired { .. } | OauthError::Held => { + cause.msg() + } }, } } @@ -135,6 +145,15 @@ fn fingerprint(cfg: &OAuth) -> u64 { h.finish() } +/// 交接时比的配置指纹（[`Cache::handover_export`]）。**和 [`fingerprint`] 管的是同几样**，但 +/// 要跨进程、跨版本比：std 的 hasher 不保证换个编译器版本还算出同一个数 +fn config_digest(cfg: &OAuth) -> String { + let parts = serde_json::json!([cfg.refresh, cfg.endpoint, cfg.client_id, cfg.client_secret]); + blake3::hash(parts.to_string().as_bytes()) + .to_hex() + .to_string() +} + struct Live { /// **生成它的那份配置**的指纹。对不上未必作废 —— 见 `usable_for` fp: u64, @@ -236,6 +255,8 @@ pub struct Cache { inner: Mutex>, /// 每个 provider 一把刷新锁（见模块文档第四条） flights: Mutex>>>, + /// 交接期间暂停续期（见 [`Self::hold_renewals`]） + held: std::sync::atomic::AtomicBool, } /// 换回来的东西。 @@ -328,6 +349,88 @@ impl Cache { Some((f.error().msg(), f.expired())) } + /// 交接：每家手里能用的 access token（见 [`crate::handover`]），免得新 core 一起来就各换一次 + /// —— 会轮换的服务器每换一次作废一个 refresh token，配置文件跟着写一次。**只交还认得这份 + /// 配置的**（[`Live::usable_for`]）；时刻换成「还有多久」「多久之前」 + pub(crate) fn handover_export( + &self, + providers: &[tw_config::Provider], + ) -> Vec { + let now = Instant::now(); + let ms = |d: Duration| d.as_millis().min(u128::from(u64::MAX)) as u64; + let g = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + g.iter() + .filter_map(|(name, live)| { + let cfg = providers.iter().find(|p| &p.name == name)?.oauth.as_ref()?; + if !live.usable_for(cfg) { + return None; + } + Some(crate::handover::wire::OauthEntry { + upstream: name.clone(), + cfg: config_digest(cfg), + access: live.access.clone(), + refresh: live.refresh.clone(), + renew_in_ms: live.renew_at.map(|t| ms(t.saturating_duration_since(now))), + age_ms: live.obtained.map(|t| ms(now.saturating_duration_since(t))), + failed: live + .failed + .as_ref() + .map(|f| crate::handover::wire::OauthFailed { + cause: f.cause.clone(), + retry_in_ms: ms(f.until.saturating_duration_since(now)), + }), + }) + }) + .collect() + } + + /// 交接：接进旧 core 手里的 token。这里的配置认得它才接：生成它的那份配置和这里的一样， + /// 或者它手里的 refresh 就是这里配置里的（轮换之后写回了）。`lag_ms` 是快照导出到现在 + /// 过了多久。 + /// + /// **这里已经有的那一家换掉**：提交之前续期归旧 core（这边暂停着），旧 core 手里的才是 + /// 准的。快照会来两份（交出监听之前再导出一次，那时旧 core 已经暂停续期），后一份要盖掉 + /// 前一份；这边从配置里读进来的那个也一样 + pub(crate) fn handover_import( + &self, + entries: Vec, + providers: &[tw_config::Provider], + lag_ms: u64, + ) { + let now = Instant::now(); + let lag = Duration::from_millis(lag_ms); + let mut g = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + for e in entries { + let Some(cfg) = providers + .iter() + .find(|p| p.name == e.upstream) + .and_then(|p| p.oauth.as_ref()) + else { + continue; + }; + if e.cfg != config_digest(cfg) && e.refresh != cfg.refresh { + continue; + } + let after = |ms: u64| now + Duration::from_millis(ms).saturating_sub(lag); + g.insert( + e.upstream, + Live { + fp: fingerprint(cfg), + access: e.access, + renew_at: e.renew_in_ms.map(after), + refresh: e.refresh, + obtained: e + .age_ms + .and_then(|ms| now.checked_sub(Duration::from_millis(ms) + lag)), + failed: e.failed.map(|f| Failure { + cause: f.cause, + until: after(f.retry_in_ms), + }), + }, + ); + } + } + fn flight(&self, provider: &str) -> Arc> { self.flights .lock() @@ -380,6 +483,41 @@ impl Cache { Ok(None) } + /// 暂停或恢复续期。 + /// + /// **交接期间新旧两个 core 不能各自续期**：服务器换发新的 refresh token 的那一刻旧的就 + /// 作废，两边各换一次，必有一边拿着作废的令牌（ChatGPT 会因此吊销整个登录）。所以旧 core + /// 交出快照（里面有它手上的 access token）起暂停，接管没成再恢复；新 core 在提交之前暂停， + /// 提交之后由它一家续期。暂停时照旧用手上那个 access token，过期了就让请求失败（[`OauthError::Held`]） + /// —— 比把整个登录弄丢好得多，而这只发生在令牌恰好在交接那几十秒里到期的时候 + pub fn hold_renewals(&self, hold: bool) { + self.held.store(hold, std::sync::atomic::Ordering::SeqCst); + } + + /// 暂停续期，**并等正在换的那几家换完**：每家的刷新锁都拿一下再放。暂停之后进来的 + /// 换不了（拿到锁时看见暂停着），暂停之前已经在换的拿着锁 —— 等它回来，它换到的 + /// access token 和新 refresh token 才在手里、才交得出去。最多等 `within` + pub async fn hold_renewals_settled(&self, within: Duration) { + self.hold_renewals(true); + let flights: Vec>> = self + .flights + .lock() + .expect("lock not poisoned") + .values() + .cloned() + .collect(); + let all = async { + for f in flights { + drop(f.lock().await); + } + }; + if tokio::time::timeout(within, all).await.is_err() { + tracing::warn!( + "a sign-in token renewal was still running when renewals were paused for the handover" + ); + } + } + /// 真的去换。**调用方持有这家的刷新锁**，所以这里不会有并发的第二次。 async fn refresh( &self, @@ -387,6 +525,16 @@ impl Cache { cfg: &OAuth, http: &reqwest::Client, ) -> Result<(String, Option), OauthError> { + if self.held.load(std::sync::atomic::Ordering::SeqCst) { + let g = self.inner.lock().expect("lock not poisoned"); + return g + .get(provider) + .filter(|l| l.usable_for(cfg) && !l.access.is_empty()) + .map(|l| l.access.clone()) + .or_else(|| cfg.access.clone().filter(|a| !a.is_empty())) + .map(|a| (a, None)) + .ok_or(OauthError::Held); + } // 用手里那个 refresh（可能是服务器换过的），没有就用配置里的 let refresh = { let g = self.inner.lock().expect("lock not poisoned"); diff --git a/crates/tw-gateway/src/pace.rs b/crates/tw-gateway/src/pace.rs index 6fc81a92..71ddbd84 100644 --- a/crates/tw-gateway/src/pace.rs +++ b/crates/tw-gateway/src/pace.rs @@ -36,6 +36,31 @@ impl Pace { .unwrap_or(1.0) } + /// 交接：每家每个模型校准到的比例（见 [`crate::handover`]） + pub(crate) fn handover_export(&self) -> Vec { + let r = self.ratios.lock().unwrap_or_else(|p| p.into_inner()); + r.iter() + .map( + |((upstream, model), ratio)| crate::handover::wire::PaceEntry { + upstream: upstream.clone(), + model: model.clone(), + ratio: *ratio, + }, + ) + .collect() + } + + /// 交接：接进旧 core 校准到的。这里已经校准过的不动；不在范围里的（读坏了）丢掉 + pub(crate) fn handover_import(&self, entries: Vec) { + let mut r = self.ratios.lock().unwrap_or_else(|p| p.into_inner()); + for e in entries + .into_iter() + .filter(|e| (RANGE.0..=RANGE.1).contains(&e.ratio)) + { + r.entry((e.upstream, e.model)).or_insert(e.ratio); + } + } + /// 一个流式请求结束了：估出来 `estimated` 个（没乘比例的），上游说是 `actual` 个 pub fn learn(&self, provider: &str, model: &str, estimated: u64, actual: u64) { if estimated < MIN_SAMPLE || actual == 0 { diff --git a/crates/tw-gateway/src/seal.rs b/crates/tw-gateway/src/seal.rs index f907f176..41537029 100644 --- a/crates/tw-gateway/src/seal.rs +++ b/crates/tw-gateway/src/seal.rs @@ -73,6 +73,36 @@ impl Refused { by.get(&(conversation.to_string(), upstream.to_string())) .map(|(_, s)| s.clone()) } + + /// 交接：全部拒过的（见 [`crate::handover`]）。指纹写成十六进制 + pub(crate) fn handover_export(&self) -> Vec { + let by = self.by.lock().unwrap_or_else(|p| p.into_inner()); + by.iter() + .map( + |((conversation, upstream), (last, seals))| crate::handover::wire::SealEntry { + conversation: conversation.clone(), + upstream: upstream.clone(), + last_ms: *last, + seals: seals + .iter() + .map(|s| blake3::Hash::from(*s).to_hex().to_string()) + .collect(), + }, + ) + .collect() + } + + /// 交接：接进旧 core 记着的，和这里的并起来。读不懂的指纹丢掉 + pub(crate) fn handover_import(&self, entries: Vec) { + for e in entries { + let seals: Vec = e + .seals + .iter() + .filter_map(|h| blake3::Hash::from_hex(h).ok().map(|h| *h.as_bytes())) + .collect(); + self.note(&e.conversation, &e.upstream, &seals, e.last_ms); + } + } } /// 上游这个 400 是在拒绝封存的推理吗。`wire` 是发给它的那一份的格式。认法见 diff --git a/crates/tw-gateway/src/session.rs b/crates/tw-gateway/src/session.rs index b4377c95..71894233 100644 --- a/crates/tw-gateway/src/session.rs +++ b/crates/tw-gateway/src/session.rs @@ -66,6 +66,30 @@ impl Sessions { e.0.clone() } + /// 交接：每段对话此刻归在哪一次会话（见 [`crate::handover`]）。不交的话，交接时正在进行的 + /// 任务在流量里被切成两段 + pub(crate) fn handover_export(&self) -> Vec { + let open = self.open.lock().unwrap_or_else(|p| p.into_inner()); + open.iter() + .map(|(fp, (id, last))| crate::handover::wire::SessionEntry { + fp: fp.clone(), + id: id.clone(), + last_ms: *last, + }) + .collect() + } + + /// 交接：接进旧 core 记着的。这里已经记着、而且更新的不动 + pub(crate) fn handover_import(&self, entries: Vec) { + let mut open = self.open.lock().unwrap_or_else(|p| p.into_inner()); + for e in entries { + if open.get(&e.fp).is_some_and(|(_, last)| *last >= e.last_ms) { + continue; + } + open.insert(e.fp, (e.id, e.last_ms)); + } + } + #[cfg(test)] fn tracked(&self) -> usize { self.open.lock().unwrap().len() diff --git a/crates/tw-gateway/src/slots.rs b/crates/tw-gateway/src/slots.rs index bba179c1..e8f06ac9 100644 --- a/crates/tw-gateway/src/slots.rs +++ b/crates/tw-gateway/src/slots.rs @@ -28,15 +28,30 @@ use std::sync::{Arc, Mutex}; use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError}; use tokio::time::Instant; -#[derive(Default)] pub struct Slots { per_upstream: Mutex>>, + /// 每家此刻占着几个位置（交接时报给接替的 core，见 [`crate::handover`]） + pub(crate) claims: Arc, +} + +impl Default for Slots { + fn default() -> Self { + Self { + per_upstream: Mutex::default(), + claims: Arc::new(crate::handover::feed::Claims::new( + crate::handover::feed::Kind::Slot, + )), + } + } } /// 一家上游的位置。 struct Pool { + /// 这家的名字：交接时按它数（[`Slot`] 的 `claim`） + name: String, sem: Arc, books: Mutex, + claims: Arc, } struct Books { @@ -49,31 +64,56 @@ struct Books { /// 占着的一个位置。**Drop 时自动归还** —— 手工释放迟早会在某条错误路径上漏掉，漏掉的 /// 表现是这家的位置只减不增，最后发给它的请求全都在等。没设上限的上游给的是空的一个。 pub struct Slot { - held: Option<(Arc, OwnedSemaphorePermit)>, + held: Option<(Arc, Hand)>, + /// 交接时数的那一个（见 [`crate::handover::feed::Claims`]） + _claim: Option, +} + +/// 占着的：信号量的一个位置，或者一笔欠账（替旧 core 占的、这里此刻没空着的那几个，见 +/// [`Slots::phantom`]；和密钥的闸里的是同一回事，见 [`crate::limits`]） +enum Hand { + Permit(OwnedSemaphorePermit), + Debt, } impl Slot { /// 这家不限并发：什么都不占 pub fn free() -> Self { - Self { held: None } + Self { + held: None, + _claim: None, + } } fn of(pool: Arc, permit: OwnedSemaphorePermit) -> Self { + let claim = pool.claims.begin(&pool.name); Self { - held: Some((pool, permit)), + held: Some((pool, Hand::Permit(permit))), + _claim: Some(claim), } } } impl Drop for Slot { fn drop(&mut self) { - let Some((pool, permit)) = self.held.take() else { + let Some((pool, hand)) = self.held.take() else { return; }; let mut books = pool.books.lock().unwrap_or_else(|p| p.into_inner()); - if books.owed > 0 { - books.owed -= 1; - permit.forget(); + match hand { + Hand::Permit(permit) => { + if books.owed > 0 { + books.owed -= 1; + permit.forget(); + } + } + Hand::Debt => { + if books.owed > 0 { + books.owed -= 1; + } else if !pool.sem.is_closed() { + pool.sem.add_permits(1); + } + } } } } @@ -117,8 +157,10 @@ impl Slots { per.insert( p.name.clone(), Arc::new(Pool { + name: p.name.clone(), sem: Arc::new(Semaphore::new(limit)), books: Mutex::new(Books { limit, owed: 0 }), + claims: self.claims.clone(), }), ); } @@ -154,6 +196,27 @@ impl Slots { } } + /// 替交接前的旧 core 占一个（见 [`crate::handover`]）：**不等**。有空位就占，没有就记一笔欠账 + /// —— 旧 core 发给这家的那一跳此刻确实在路上。这家不限并发是 None：什么都不用占 + pub(crate) fn phantom(&self, name: &str) -> Option { + let pool = self.pool(name)?; + if pool.sem.is_closed() { + return None; + } + let hand = match pool.sem.clone().try_acquire_owned() { + Ok(permit) => Hand::Permit(permit), + Err(_) => { + pool.books.lock().unwrap_or_else(|p| p.into_inner()).owed += 1; + Hand::Debt + } + }; + let claim = pool.claims.begin(&pool.name); + Some(Slot { + held: Some((pool, hand)), + _claim: Some(claim), + }) + } + /// 等这一家空出一个位置，最多等到 `until`。到点还满着是 `None`。 pub async fn take_by(&self, name: &str, until: Instant) -> Option { self.first_free(&[name], until).await.map(|(_, slot)| slot) diff --git a/crates/tw-gateway/src/state.rs b/crates/tw-gateway/src/state.rs index 1a4704f5..3586bbf1 100644 --- a/crates/tw-gateway/src/state.rs +++ b/crates/tw-gateway/src/state.rs @@ -287,6 +287,9 @@ pub struct AppState { pub opening_hold: std::time::Duration, /// 响应头先交出去之后隔多久发一行保活（[`crate::KEEPALIVE_EVERY`]）。**测试会把它调短** pub keepalive_every: std::time::Duration, + /// 新旧 core 交接时网关这一边的东西：排空的标记、替旧 core 占着的名额（见 + /// [`crate::handover`]）。**跨重载存活** + pub(crate) handover: Arc, } impl AppState { @@ -372,6 +375,7 @@ impl AppState { idle_tick: std::time::Duration::from_secs(1), opening_hold: crate::OPENING_HOLD, keepalive_every: crate::KEEPALIVE_EVERY, + handover: Default::default(), }; // 手写的清单马上可用；向上游问是后台的事，不挡启动 state.publish_catalog(); diff --git a/crates/tw-gateway/src/throttle.rs b/crates/tw-gateway/src/throttle.rs index ad820e9f..9ee7d67a 100644 --- a/crates/tw-gateway/src/throttle.rs +++ b/crates/tw-gateway/src/throttle.rs @@ -742,6 +742,80 @@ impl Throttles { } } + /// 交接：每家的限流暂停（见 [`crate::handover`]）。时刻换成「还剩多久」「过去多久」： + /// 单调时钟出不了这个进程。**排着的队和探路的名额不交**：那是旧 core 手上的请求 + pub(crate) fn handover_export(&self) -> Vec { + let now = self.moment(); + let ms = |d: Duration| d.as_millis().min(u128::from(u64::MAX)) as u64; + self.lock() + .iter() + .filter_map(|(name, up)| { + let pause_left_ms = up.pause.filter(|p| !p.over(now)).map(|p| ms(p.left(now))); + let limited_ago_ms = up.limited_at.map(|t| ms(t.since(now))); + if pause_left_ms.is_none() + && limited_ago_ms.is_none() + && up.strikes == 0 + && !up.probing + { + return None; + } + Some(crate::handover::wire::ThrottleEntry { + upstream: name.clone(), + pause_left_ms, + limited_ago_ms, + strikes: up.strikes, + probing: up.probing, + }) + }) + .collect() + } + + /// 交接：接进旧 core 的限流暂停。`lag_ms` 是快照导出到现在过了多久，从剩下的里扣掉。这里 + /// 已经被限流过的那一家不动 —— 那是更新的事。暂停结束那一刻照样报一次（[`Self::at`]） + pub(crate) fn handover_import( + self: &Arc, + entries: Vec, + lag_ms: u64, + ) { + let now = self.moment(); + let back = |ago_ms: u64| { + let d = Duration::from_millis(ago_ms); + Moment { + mono: now.mono.checked_sub(d).unwrap_or(now.mono), + wall: now.wall.saturating_sub(ago_ms), + } + }; + let mut touched = Vec::new(); + { + let mut map = self.lock(); + for e in entries { + let up = map.entry(e.upstream.clone()).or_insert_with(Up::new); + if up.pause.is_some() || up.limited_at.is_some() { + continue; + } + up.pause = e + .pause_left_ms + .map(|ms| ms.saturating_sub(lag_ms)) + .filter(|ms| *ms > 0) + .map(|ms| Span { + from: now, + len: Duration::from_millis(ms), + }); + up.limited_at = e.limited_ago_ms.map(|ago| back(ago.saturating_add(lag_ms))); + up.strikes = e.strikes; + up.probing = e.probing; + up.bump(); + touched.push((e.upstream, up.pause)); + } + } + for (name, pause) in touched { + self.refresh(&name); + if let Some(p) = pause { + self.at(&name, p); + } + } + } + /// `span` 结束时重报一次这一家的样子、叫醒排队的。**睡一会儿起来看一眼**（[`NAP`]）：睡着的 /// 时候单调时钟不走，只按它睡，醒来后还要再等完剩下的 fn at(self: &Arc, name: &str, span: Span) { diff --git a/crates/tw-gateway/src/upstream_limits.rs b/crates/tw-gateway/src/upstream_limits.rs index d37e1ebf..f636824b 100644 --- a/crates/tw-gateway/src/upstream_limits.rs +++ b/crates/tw-gateway/src/upstream_limits.rs @@ -68,6 +68,10 @@ struct Inner { asked: HashSet, /// 等这家变化的（[`UpstreamLimits::changed`]） wake: HashMap>, + /// 交接：预留和结算往哪儿报（旧 core，见 [`crate::handover`]），在这把锁里报 + feed: Option>, + /// 交接：替旧 core 记着的预留，旧 core 那边的号 → 这里的号（新 core） + imported: HashMap, } struct Shared { @@ -92,6 +96,7 @@ impl Shared { let before: HashMap = g.held.iter().map(|(s, r)| (*s, r.key.clone())).collect(); for (seq, id) in books::sweep_held(&mut g.held, &self.bus, now) { + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::UpDrop { seq }); if let Some(seqs) = g.by_request.get_mut(&id) { seqs.retain(|s| *s != seq); if seqs.is_empty() { @@ -149,23 +154,29 @@ impl Shared { fn release(&self, seq: u64) { let key = { let mut g = self.lock(); - let g = &mut *g; - let Some(r) = g.held.remove(&seq) else { + let Some(key) = Self::release_locked(&mut g, seq) else { return; }; - if let Some(id) = r.request - && let Some(seqs) = g.by_request.get_mut(&id) - { - seqs.retain(|s| *s != seq); - if seqs.is_empty() { - g.by_request.remove(&id); - } - } - books::unrecord(&mut g.books, &r); - r.key + key }; self.wake([key.as_str()]); } + + /// 交回还回去的那一家 + fn release_locked(g: &mut Inner, seq: u64) -> Option { + let r = g.held.remove(&seq)?; + if let Some(id) = r.request + && let Some(seqs) = g.by_request.get_mut(&id) + { + seqs.retain(|s| *s != seq); + if seqs.is_empty() { + g.by_request.remove(&id); + } + } + books::unrecord(&mut g.books, &r); + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::UpRelease { seq }); + Some(r.key) + } } /// 一跳发出去时占着的那一份（[`UpstreamLimits::reserve`]）。**丢掉就还回去**：回了 429、 @@ -403,10 +414,12 @@ impl UpstreamLimits { let g = &mut *g; books::record_request(&mut g.books, upstream, limits, now); g.seq += 1; - g.held.insert( - g.seq, - Reservation::new(upstream, now, ask.tokens, ask.cost_micros), - ); + let r = Reservation::new(upstream, now, ask.tokens, ask.cost_micros); + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::UpReserve { + held: books::held_wire(g.seq, &r), + windowed: books::longest_window(limits).is_some(), + }); + g.held.insert(g.seq, r); g.seq }; // 记上的这一个可能把窗口、这一期用满了：看着这家的（报状态的）再看一眼 @@ -432,74 +445,238 @@ impl UpstreamLimits { rec: &Recorded, ) { let now = self.now(); - let clock = self.shared.clock.as_ref(); let (touched, moved) = { let mut g = self.shared.lock(); let g = &mut *g; - let held: Vec = g + let held: Vec<(u64, Reservation)> = g .by_request .remove(&id) .unwrap_or_default() .into_iter() - .filter_map(|s| g.held.remove(&s)) + .filter_map(|s| g.held.remove(&s).map(|r| (s, r))) .collect(); - let mut touched: HashSet = held.iter().map(|r| r.key.clone()).collect(); - let counted = counts(rec); - // 滚动窗口里的请求数：发出时就记上了。记录里算数的那几跳留着，别的撤掉 - let mut left: Vec<&str> = if counted { - hops.iter().map(String::as_str).collect() - } else { - Vec::new() - }; - for r in &held { - match left.iter().position(|h| *h == r.key) { - Some(i) => { - left.swap_remove(i); - } - None => books::unrecord(&mut g.books, r), + self.settle_locked(g, held, at_ms, hops, served, rec, now) + }; + self.shared.wake(touched.iter().map(String::as_str)); + self.shared.ask_reread(moved); + } + + /// 结算一行：`held` 是它那几跳占着的预留和预留的号。交回动过的那几家和换了新一期的那几期 + #[allow(clippy::too_many_arguments)] + fn settle_locked( + &self, + g: &mut Inner, + held: Vec<(u64, Reservation)>, + at_ms: i64, + hops: &[String], + served: Option<&str>, + rec: &Recorded, + now: i64, + ) -> (HashSet, Vec) { + let clock = self.shared.clock.as_ref(); + if !held.is_empty() || counts(rec) { + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::UpSettle { + seqs: held.iter().map(|(s, _)| *s).collect(), + at_ms, + hops: hops.to_vec(), + served: served.map(str::to_string), + rec: rec.clone(), + }); + } + let held: Vec = held.into_iter().map(|(_, r)| r).collect(); + let mut touched: HashSet = held.iter().map(|r| r.key.clone()).collect(); + let counted = counts(rec); + // 滚动窗口里的请求数：发出时就记上了。记录里算数的那几跳留着，别的撤掉 + let mut left: Vec<&str> = if counted { + hops.iter().map(String::as_str).collect() + } else { + Vec::new() + }; + for r in &held { + match left.iter().position(|h| *h == r.key) { + Some(i) => { + left.swap_remove(i); } + None => books::unrecord(&mut g.books, r), + } + } + if counted { + let one = Amount { + requests: 1, + ..Default::default() + }; + for up in hops { + add_to_periods( + g.books.entry(up.clone()).or_default(), + clock, + &one, + at_ms, + now, + ); + touched.insert(up.clone()); } - if counted { - let one = Amount { - requests: 1, - ..Default::default() + if let Some(up) = served { + let amount = Amount { + requests: 0, + ..rec.amount() }; - for up in hops { - add_to_periods( - g.books.entry(up.clone()).or_default(), - clock, - &one, - at_ms, - now, - ); - touched.insert(up.clone()); + let rolling = g + .limits + .get(up) + .and_then(|ls| books::longest_window(ls)) + .is_some(); + let b = g.books.entry(up.to_string()).or_default(); + add_to_periods(b, clock, &amount, at_ms, now); + if rolling { + b.recent.push_back((now, amount)); } - if let Some(up) = served { - let amount = Amount { - requests: 0, - ..rec.amount() - }; - let rolling = g - .limits - .get(up) - .and_then(|ls| books::longest_window(ls)) - .is_some(); - let b = g.books.entry(up.to_string()).or_default(); - add_to_periods(b, clock, &amount, at_ms, now); - if rolling { - b.recent.push_back((now, amount)); + touched.insert(up.to_string()); + } + } + let moved: Vec = touched + .iter() + .flat_map(|k| books::take_moved(&mut g.books, &mut g.asked, k)) + .collect(); + (touched, moved) + } + + // ------------------------------------------------------------ 交接（见 `crate::handover`） + + /// 旧 core：此刻的账和在路上的预留，之后的变化报给 `feed`（同一把锁里，见 + /// `KeyLimits::handover_export`） + pub(crate) fn handover_export( + &self, + feed: Option>, + ) -> crate::handover::wire::LimitBooks { + let mut g = self.shared.lock(); + let out = books::export(&g.books, &g.held); + g.feed = feed; + out + } + + /// 旧 core：排空开始。报一份此刻在路上的全部预留，之后的变化接着报 + pub(crate) fn handover_drain(&self, feed: Arc) { + let mut g = self.shared.lock(); + feed.send(&crate::handover::wire::Event::UpHolds { + held: books::export(&HashMap::new(), &g.held).held, + }); + g.feed = Some(feed); + } + + /// 新 core：接进旧 core 的账，在路上的预留记成替它占着的 + pub(crate) fn handover_import(&self, b: crate::handover::wire::LimitBooks) { + let names: Vec = { + let mut g = self.shared.lock(); + let g = &mut *g; + books::import(&mut g.books, b.books); + for h in &b.held { + Self::phantom(g, h, false); + } + g.books.keys().cloned().collect() + }; + self.shared.wake(names.iter().map(String::as_str)); + } + + fn phantom(g: &mut Inner, h: &crate::handover::wire::Held, windowed: bool) { + if g.imported.contains_key(&h.seq) { + return; + } + g.seq += 1; + let seq = g.seq; + let r = books::held_reservation(h); + if windowed { + books::record_at(&mut g.books, &r.key, r.at_ms); + } + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::UpReserve { + held: books::held_wire(seq, &r), + windowed, + }); + g.held.insert(seq, r); + g.imported.insert(h.seq, seq); + } + + /// 替旧 core 记着的这一份作废（不撤滚动窗口里那一个请求）。交回它是哪一家的 + fn drop_phantom(g: &mut Inner, theirs: u64) -> Option { + let seq = g.imported.remove(&theirs)?; + let r = g.held.remove(&seq)?; + crate::handover::feed::send(&g.feed, || crate::handover::wire::Event::UpDrop { seq }); + Some(r.key) + } + + /// 新 core：旧 core 报来一条变化（`up_*` 那几种）。别的不认，交回 false + pub(crate) fn handover_apply(&self, ev: &crate::handover::wire::Event) -> bool { + use crate::handover::wire::Event; + let now = self.now(); + let (touched, moved) = { + let mut g = self.shared.lock(); + let g = &mut *g; + match ev { + Event::UpHolds { held } => { + let want: HashSet = held.iter().map(|h| h.seq).collect(); + let stale: Vec = g + .imported + .keys() + .filter(|s| !want.contains(s)) + .copied() + .collect(); + let mut touched: HashSet = stale + .into_iter() + .filter_map(|s| Self::drop_phantom(g, s)) + .collect(); + for h in held { + Self::phantom(g, h, false); + touched.insert(h.key.clone()); } - touched.insert(up.to_string()); + (touched, Vec::new()) + } + Event::UpReserve { held, windowed } => { + Self::phantom(g, held, *windowed); + (HashSet::from([held.key.clone()]), Vec::new()) + } + Event::UpSettle { + seqs, + at_ms, + hops, + served, + rec, + } => { + let held: Vec<(u64, Reservation)> = seqs + .iter() + .filter_map(|s| g.imported.remove(s)) + .filter_map(|s| g.held.remove(&s).map(|r| (s, r))) + .collect(); + self.settle_locked(g, held, *at_ms, hops, served.as_deref(), rec, now) + } + Event::UpRelease { seq } => { + let touched = g + .imported + .remove(seq) + .and_then(|s| Shared::release_locked(g, s)); + (touched.into_iter().collect(), Vec::new()) } + Event::UpDrop { seq } => ( + Self::drop_phantom(g, *seq).into_iter().collect(), + Vec::new(), + ), + _ => return false, } - let moved: Vec = touched - .iter() - .flat_map(|k| books::take_moved(&mut g.books, &mut g.asked, k)) - .collect(); - (touched, moved) }; self.shared.wake(touched.iter().map(String::as_str)); self.shared.ask_reread(moved); + true + } + + /// 新 core：交接连接断了。替旧 core 记着的预留全部作废 + pub(crate) fn handover_release_all(&self) { + let touched: Vec = { + let mut g = self.shared.lock(); + let theirs: Vec = g.imported.keys().copied().collect(); + theirs + .into_iter() + .filter_map(|s| Self::drop_phantom(&mut g, s)) + .collect() + }; + self.shared.wake(touched.iter().map(String::as_str)); } /// 重启之后把天、周、月的数从请求记录里加回来。`since(t)` 给出从 `t` 起每家上游记下的 diff --git a/crates/tw-gateway/src/ws.rs b/crates/tw-gateway/src/ws.rs index 091b7c25..d86f00e3 100644 --- a/crates/tw-gateway/src/ws.rs +++ b/crates/tw-gateway/src/ws.rs @@ -1100,8 +1100,14 @@ enum End { Broke(Msg), /// 被防护切断了：回答里的工具调用命中了切断规则，或者客户端发来的一帧被内容过滤拒了 Cut(Msg), + /// 这个 core 在排空（交给了新版本的 core，见 [`crate::handover`]），手上这一轮答完了：以 + /// 1012（服务重启）关闭，客户端重连到新的那一个 + Restart, } +/// 排空时关连接用的关闭码：服务重启（RFC 6455 7.4.1 之后 IANA 登记的 1012）。客户端照它重连 +pub(crate) const CLOSE_RESTART: u16 = 1012; + /// 在等准入的那一轮（见 [`turn::admit`]）：等到了交回这一轮和这一帧接下来要用的。 type Waiting = std::pin::Pin< Box, Next)> + Send>, @@ -1135,6 +1141,11 @@ async fn pump( let mut waiting: Option = None; // 等的时候客户端接着发来的帧：排在那一轮后面，轮到了按顺序处理 let mut held: std::collections::VecDeque = Default::default(); + // 排空（见 `crate::handover`）：Responses 的连接手上没有在跑的轮次了就以 1012 关闭 —— 闲着的 + // 当场关，在答的答完这一轮再关。整条连接一行的那几种分不出一轮一轮，照常开着 + let mut drain = state.handover.draining(); + let mut draining = *drain.borrow_and_update(); + let mut drain_gone = false; let end = 'pump: loop { while waiting.is_none() && let Some(m) = held.pop_front() @@ -1153,7 +1164,21 @@ async fn pump( break 'pump end; } } + if draining + && waiting.is_none() + && held.is_empty() + && p.turns.as_ref().is_some_and(turn::Turns::is_empty) + { + break End::Restart; + } tokio::select! { + changed = drain.changed(), if !draining && !drain_gone => { + match changed { + Ok(()) => draining = *drain.borrow_and_update(), + // 发的那一头没了（网关整个收了）：不再看 + Err(_) => drain_gone = true, + } + } msg = c_rx.next() => { let Some(Ok(m)) = msg else { break End::Closed }; if waiting.is_some() { @@ -1254,13 +1279,22 @@ async fn pump( _ => t.clear(), } } + let restart = matches!(end, End::Restart); if let Some(ending) = ending { match end { - End::Closed | End::UpstreamClosed => ending.finished(101), + End::Closed | End::UpstreamClosed | End::Restart => ending.finished(101), End::Broke(why) => ending.failed(tw_api::FailureSource::Upstream, why), End::Cut(why) => ending.failed(tw_api::FailureSource::Denied, why), } } + if restart { + let _ = c_tx + .send(Message::Close(Some(axum::extract::ws::CloseFrame { + code: CLOSE_RESTART, + reason: "the gateway is restarting; reconnect".into(), + }))) + .await; + } let _ = c_tx.close().await; let _ = u_tx.close().await; } diff --git a/crates/tw-gateway/src/ws/turn.rs b/crates/tw-gateway/src/ws/turn.rs index 65542d8b..cf00af22 100644 --- a/crates/tw-gateway/src/ws/turn.rs +++ b/crates/tw-gateway/src/ws/turn.rs @@ -187,6 +187,11 @@ impl Turns { self.queue.front_mut() } + /// 一轮都没有在跑 + pub(crate) fn is_empty(&self) -> bool { + self.queue.is_empty() + } + pub(crate) fn front_id(&self) -> Option { self.queue.front().map(|t| t.id) } diff --git a/crates/tw-gateway/tests/handover_ws.rs b/crates/tw-gateway/tests/handover_ws.rs new file mode 100644 index 00000000..15bab018 --- /dev/null +++ b/crates/tw-gateway/tests/handover_ws.rs @@ -0,0 +1,197 @@ +//! 交接时的 WebSocket：旧 core 排空开始后，Responses 的连接**答完手上这一轮**以 1012（服务重启） +//! 关闭，闲着的当场关闭 —— 客户端重连，连到新 core 上。 + +use std::net::SocketAddr; +use std::time::Duration; + +use axum::Router; +use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; +use futures::{SinkExt, StreamExt}; +use tokio_tungstenite::tungstenite::Message as WsMsg; +use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; +use tw_config::{Client, Config, Provider}; + +/// 像 Responses 的 WebSocket 那样回答的上游：每个 `response.create` 回 created、一段文字，等 +/// `hold` 变成 true 再回 completed +async fn responder(hold: tokio::sync::watch::Receiver) -> SocketAddr { + let app = Router::new().route( + "/backend-api/codex/responses", + axum::routing::any(move |ws: WebSocketUpgrade| { + let hold = hold.clone(); + async move { ws.on_upgrade(move |sock| answer(sock, hold)) } + }), + ); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + addr +} + +async fn answer(mut sock: WebSocket, mut hold: tokio::sync::watch::Receiver) { + let mut n = 0; + while let Some(Ok(m)) = sock.recv().await { + if !matches!(m, Message::Text(_)) { + continue; + } + n += 1; + let id = format!("resp_{n}"); + let head = [ + serde_json::json!({"type":"response.created","response":{"id":id,"status":"in_progress","model":"gpt-5","output":[]}}), + serde_json::json!({"type":"response.output_text.delta","output_index":0,"content_index":0,"item_id":"msg","delta":"hello"}), + ]; + for f in head { + if sock + .send(Message::Text(f.to_string().into())) + .await + .is_err() + { + return; + } + } + while !*hold.borrow_and_update() { + if hold.changed().await.is_err() { + return; + } + } + let done = serde_json::json!({"type":"response.completed","response":{"id":id,"status":"completed","model":"gpt-5","output":[],"usage":{"input_tokens":10,"output_tokens":2,"total_tokens":12}}}); + if sock + .send(Message::Text(done.to_string().into())) + .await + .is_err() + { + return; + } + } +} + +async fn gateway(up: SocketAddr) -> (SocketAddr, tw_gateway::AppState) { + let cfg = Config { + version: 1, + clients: vec![Client { + name: "codex".into(), + key: "tw-wskey".into(), + ..Default::default() + }], + providers: vec![Provider { + name: "up".into(), + base_url: format!("http://{up}"), + key: Some("sk-upstream".into()), + protocol: Some(tw_config::Protocol::OpenaiResponses), + ..Default::default() + }], + ..Default::default() + }; + let state = tw_gateway::AppState::new(cfg).unwrap(); + let addr = tw_gateway::serve(state.clone(), ([127, 0, 0, 1], 0).into()) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(60)).await; + (addr, state) +} + +type Socket = + tokio_tungstenite::WebSocketStream>; + +async fn connect(gw: SocketAddr) -> Socket { + use tokio_tungstenite::tungstenite::client::IntoClientRequest; + let mut req = format!("ws://{gw}/backend-api/codex/responses") + .into_client_request() + .unwrap(); + req.headers_mut() + .insert("x-api-key", "tw-wskey".parse().unwrap()); + let (s, _) = tokio_tungstenite::connect_async(req).await.unwrap(); + s +} + +fn create() -> WsMsg { + WsMsg::Text( + serde_json::json!({ + "type": "response.create", + "model": "gpt-5", + "input": [{"role": "user", "content": [{"type": "input_text", "text": "hi"}]}] + }) + .to_string() + .into(), + ) +} + +async fn next(c: &mut Socket) -> WsMsg { + tokio::time::timeout(Duration::from_secs(5), c.next()) + .await + .expect("nothing came") + .expect("the connection ended without a close frame") + .unwrap() +} + +/// 读到这一轮的第一段文字为止：上游已经在回答这一轮了 +async fn until_text(c: &mut Socket) { + loop { + if let WsMsg::Text(t) = next(c).await + && t.contains("output_text.delta") + { + return; + } + } +} + +/// 读文字帧，直到关闭帧：交回读到的每一帧的 `type`，和关闭码 +async fn until_close(c: &mut Socket) -> (Vec, Option) { + let mut kinds = Vec::new(); + loop { + match next(c).await { + WsMsg::Text(t) => { + let v: serde_json::Value = serde_json::from_str(&t).unwrap(); + kinds.push(v["type"].as_str().unwrap_or_default().to_string()); + } + WsMsg::Close(f) => return (kinds, f.map(|f| f.code)), + _ => {} + } + } +} + +/// 闲着的连接：排空一开始就以 1012 关闭 +#[tokio::test] +async fn an_idle_connection_is_closed_with_service_restart_when_draining_begins() { + let (_tx, hold) = tokio::sync::watch::channel(true); + let up = responder(hold).await; + let (gw, state) = gateway(up).await; + let mut c = connect(gw).await; + // 先答完一轮：连接是活的，只是此刻闲着 + c.send(create()).await.unwrap(); + loop { + if let WsMsg::Text(t) = next(&mut c).await + && t.contains("response.completed") + { + break; + } + } + + let _rx = state.begin_drain(); + let (kinds, code) = until_close(&mut c).await; + assert!(kinds.is_empty(), "{kinds:?}"); + assert_eq!(code, Some(CloseCode::Restart)); +} + +/// 在答的连接：这一轮照常答完（客户端收到 completed），然后才以 1012 关闭 +#[tokio::test] +async fn a_connection_in_the_middle_of_a_turn_finishes_it_before_closing() { + let (tx, hold) = tokio::sync::watch::channel(false); + let up = responder(hold).await; + let (gw, state) = gateway(up).await; + let mut c = connect(gw).await; + c.send(create()).await.unwrap(); + until_text(&mut c).await; + + let _rx = state.begin_drain(); + // 排空开始了，这一轮还没答完：连接不能关 + assert!( + tokio::time::timeout(Duration::from_millis(200), c.next()) + .await + .is_err(), + "这一轮还没答完，连接就有动静了" + ); + tx.send(true).unwrap(); + let (kinds, code) = until_close(&mut c).await; + assert_eq!(kinds, ["response.completed"]); + assert_eq!(code, Some(CloseCode::Restart)); +} diff --git a/crates/tw-gateway/tests/listen_handoff.rs b/crates/tw-gateway/tests/listen_handoff.rs new file mode 100644 index 00000000..55ce17ee --- /dev/null +++ b/crates/tw-gateway/tests/listen_handoff.rs @@ -0,0 +1,191 @@ +//! 交接时的监听（升级不断线）：交出去的是同一个内核套接字，旧的停止 accept 只丢自己的 +//! 句柄，新的照样在听；新的提交之前不跟配置换监听。 +//! +//! 这里不跨进程：两个「core」是同一个进程里的两份 `serve_with`，套接字用复制出来的句柄 +//! 交 —— 跨进程那一段（SCM_RIGHTS / WSADuplicateSocketW）是 `tw-handover` 的事，内核 +//! 看到的东西是一样的：同一个套接字的两个描述符。 + +use std::net::SocketAddr; +use std::time::Duration; + +use tw_config::{Client, Config}; +use tw_gateway::listen::Handoff; + +fn cfg(port: u16) -> Config { + let mut c = Config { + clients: vec![Client { + name: "c".into(), + key: "tw-k".into(), + ..Default::default() + }], + ..Default::default() + }; + c.listen.gateway.port = port; + c +} + +/// 网关答不答（没有上游时也会答一句错误：答了就是在听） +async fn answers(addr: SocketAddr) -> bool { + reqwest::Client::new() + .get(format!("http://{addr}/v1/models")) + .header("x-api-key", "tw-k") + .timeout(Duration::from_secs(2)) + .send() + .await + .is_ok() +} + +#[tokio::test] +async fn the_old_side_stops_accepting_and_the_shared_socket_keeps_listening() { + let c = cfg(0); + let old = tw_gateway::AppState::new(c.clone()).unwrap(); + let new = tw_gateway::AppState::new(c.clone()).unwrap(); + let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = l.local_addr().unwrap(); + + // 旧的在一个现成的监听上服务（起始那一套也可以是交过来的） + let old_hand = Handoff::default(); + let serving_old = tokio::spawn(tw_gateway::listen::serve_with( + old.clone(), + c.listen.gateway.clone(), + false, + old_hand.clone(), + vec![(addr, l)], + )); + old_hand.listening().await; + assert!(answers(addr).await, "旧的没在听"); + assert_eq!(old.listening().primary(), Some(addr)); + + // 交出去：复制出来的句柄，同一个地址 + let exported = old_hand.export().unwrap(); + assert_eq!(exported.len(), 1); + assert_eq!(exported[0].0, addr); + assert_eq!(exported[0].1.local_addr().unwrap(), addr); + + // 新的在交过来的上面服务，提交之前先不跟配置走 + let new_hand = Handoff::default(); + new_hand.hold(); + tokio::spawn(tw_gateway::listen::serve_with( + new.clone(), + c.listen.gateway.clone(), + true, + new_hand.clone(), + exported, + )); + new_hand.listening().await; + + // 提交：旧的停止 accept。它的服务结束、句柄丢了，**套接字还在听** + old_hand.stop_accepting(); + tokio::time::timeout(Duration::from_secs(5), serving_old) + .await + .expect("停止 accept 之后 serve_with 该返回") + .unwrap() + .unwrap(); + for _ in 0..20 { + if old_hand.serving() == 0 { + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert_eq!(old_hand.serving(), 0, "旧的服务没结束"); + assert!(old_hand.export().unwrap().is_empty(), "停了还交得出监听"); + for _ in 0..10 { + assert!(answers(addr).await, "旧的丢了句柄，端口跟着关了"); + } + new_hand.follow(); +} + +#[tokio::test] +async fn a_held_listener_does_not_follow_the_config_until_told() { + let spare = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let elsewhere = spare.local_addr().unwrap().port(); + drop(spare); + + let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = l.local_addr().unwrap(); + // 配置说的是另一个端口：交过来的地址和配置不一样 + let c = cfg(elsewhere); + let state = tw_gateway::AppState::new(c.clone()).unwrap(); + let hand = Handoff::default(); + hand.hold(); + tokio::spawn(tw_gateway::listen::serve_with( + state.clone(), + c.listen.gateway.clone(), + true, + hand.clone(), + vec![(addr, l)], + )); + hand.listening().await; + // 叫它对一遍配置：提交之前不该动 + state.relisten(); + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!( + state.listening().primary(), + Some(addr), + "提交之前就换了监听" + ); + + // 提交之后跟配置走：马上换过去 + hand.follow(); + let moved = SocketAddr::from(([127, 0, 0, 1], elsewhere)); + for _ in 0..40 { + if state.listening().primary() == Some(moved) { + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert_eq!( + state.listening().primary(), + Some(moved), + "跟配置走之后没换过去" + ); + assert!(answers(moved).await); +} + +#[tokio::test] +async fn a_released_listener_is_not_kept_alive_by_the_handoff() { + // 换监听时退场的那个要真的放开端口：交接把手里攥着一份的话，「旧的先让出来」那一步 + // 就永远绑不上 + let p = { + let s = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + s.local_addr().unwrap().port() + }; + let mut c = cfg(p); + let state = tw_gateway::AppState::new(c.clone()).unwrap(); + let hand = Handoff::default(); + tokio::spawn(tw_gateway::listen::serve_with( + state.clone(), + c.listen.gateway.clone(), + true, + hand.clone(), + Vec::new(), + )); + hand.listening().await; + let q = { + let s = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + s.local_addr().unwrap().port() + }; + c.listen.gateway.port = q; + state.reload(c).unwrap(); + for _ in 0..40 { + if state.listening().primary().map(|a| a.port()) == Some(q) { + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert_eq!(state.listening().primary().map(|a| a.port()), Some(q)); + let exported = hand.export().unwrap(); + assert_eq!(exported.len(), 1); + assert_eq!(exported[0].0.port(), q, "交出去的该是新的那个"); + drop(exported); + // 旧端口放开了：绑得上 + let mut freed = false; + for _ in 0..40 { + if std::net::TcpListener::bind(("127.0.0.1", p)).is_ok() { + freed = true; + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert!(freed, "退场的监听没放开端口"); +} diff --git a/crates/tw-gateway/tests/oauth.rs b/crates/tw-gateway/tests/oauth.rs index 59601571..54386415 100644 --- a/crates/tw-gateway/tests/oauth.rs +++ b/crates/tw-gateway/tests/oauth.rs @@ -516,3 +516,103 @@ async fn an_error_body_from_the_token_endpoint_is_masked_before_being_reported() ); assert!(e.contains(""), "抹掉了但没说抹掉了什么：{e}"); } + +/// 交接期间暂停续期：**一次都不去换**，手上没有 access token 的请求失败，恢复之后照常换。 +/// +/// 两个 core 各换一次 refresh token，必有一边拿着作废的那个 —— ChatGPT 会因此吊销整个登录。 +#[tokio::test] +async fn renewals_held_during_a_handover_never_reach_the_token_endpoint() { + let t = Arc::new(Token { + expires_in: Some(3600), + ..Default::default() + }); + let token_addr = start_token_endpoint(t.clone()).await; + let (up, seen) = start_upstream().await; + let (gw, s) = start_gateway(vec![oauth_provider( + "p", + up, + &format!("http://{token_addr}/token"), + )]) + .await; + + s.oauth.hold_renewals(true); + assert_ne!(ask(gw).await, 200, "暂停续期时没有 access token 也放行了"); + assert_eq!( + t.calls.load(Ordering::SeqCst), + 0, + "暂停续期时仍去换了 token" + ); + assert!(seen.lock().unwrap().is_empty(), "没有令牌的请求到了上游"); + + s.oauth.hold_renewals(false); + assert_eq!(ask(gw).await, 200); + assert_eq!(t.calls.load(Ordering::SeqCst), 1); + + // 已经有一个能用的：暂停期间照用，不换 + s.oauth.hold_renewals(true); + assert_eq!(ask(gw).await, 200); + assert_eq!(t.calls.load(Ordering::SeqCst), 1); +} + +/// 交出 access token 之前（升级不断线）：暂停续期**并等正在换的那一次换完**。暂停之前已经 +/// 发出去的那次换回来的才是旧 core 手上最新的，要等它；暂停之后进来的一次都不去换。 +#[tokio::test] +async fn pausing_renewals_waits_for_the_renewal_already_in_flight() { + // 慢的 token 端点：一次要 600 毫秒 + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let c = calls.clone(); + let app = axum::Router::new().route( + "/token", + axum::routing::post(move || { + let c = c.clone(); + async move { + c.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(std::time::Duration::from_millis(600)).await; + axum::Json(serde_json::json!({ + "access_token": "at-slow", + "expires_in": 3600, + })) + } + }), + ); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + + let cache = Arc::new(tw_gateway::oauth::Cache::new()); + let cfg = OAuth { + access: None, + expires_at: None, + refresh: "rt".into(), + endpoint: format!("http://{addr}/token"), + client_id: None, + client_secret: None, + refresh_before: None, + }; + let renewing = { + let (cache, cfg) = (cache.clone(), cfg.clone()); + tokio::spawn(async move { cache.token("p", &cfg, &reqwest::Client::new()).await }) + }; + // 等它真的发出去 + while calls.load(Ordering::SeqCst) == 0 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + let t0 = std::time::Instant::now(); + cache + .hold_renewals_settled(std::time::Duration::from_secs(5)) + .await; + assert!( + t0.elapsed() >= std::time::Duration::from_millis(300), + "没等正在换的那一次就交出去了" + ); + let (token, _) = renewing.await.unwrap().unwrap(); + assert_eq!(token, "at-slow"); + // 暂停之后：手上那个照用，不再去换 + let (token, renewed) = cache + .token("p", &cfg, &reqwest::Client::new()) + .await + .unwrap(); + assert_eq!(token, "at-slow"); + assert!(renewed.is_none()); + assert_eq!(calls.load(Ordering::SeqCst), 1, "暂停之后又去换了"); +} diff --git a/crates/tw-handover/Cargo.toml b/crates/tw-handover/Cargo.toml new file mode 100644 index 00000000..aa05e28f --- /dev/null +++ b/crates/tw-handover/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "tw-handover" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +documentation.workspace = true +readme.workspace = true +description = "Hands the listening sockets and runtime state of a running core to its successor" + +[dependencies] +tokio = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } + +[target.'cfg(unix)'.dependencies] +libc = { workspace = true } + +[target.'cfg(windows)'.dependencies] +windows-sys = { workspace = true, features = ["Win32_Networking_WinSock", "Win32_Foundation", "Win32_Security_Cryptography"] } + +[dev-dependencies] +tokio = { workspace = true, features = ["rt", "macros", "test-util"] } +tempfile = "3" diff --git a/crates/tw-handover/src/lib.rs b/crates/tw-handover/src/lib.rs new file mode 100644 index 00000000..f7bb05e2 --- /dev/null +++ b/crates/tw-handover/src/lib.rs @@ -0,0 +1,425 @@ +//! 新旧 core 之间的交接：把正在监听的套接字和运行时状态交给接替它的那一个。 +//! +//! # 为什么传套接字，而不是各自绑一次 +//! +//! 传过去的是**同一个内核套接字**，不是重新绑定。旧的停止 accept 之后，队列里 +//! 还没被接受的连接由新的取走，一个都不丢；两边短暂地同时 accept 也没关系， +//! 每个连接只会被一方接受。各自绑（`SO_REUSEPORT`）做不到这一点：Linux 上旧的 +//! 那个关掉时，它队列里的连接会被重置；Windows 上没有等价的东西。 +//! +//! # 交接连接 +//! +//! 旧 core 被 arm 之后在数据目录里开一个**只为交接用**的端点（[`Endpoint`]）， +//! 新 core 取锁时发现锁的持有者还活着，就去连它（[`connect`]）。连不上 = 没有 +//! arm，照旧报「已有实例在运行」。 +//! +//! - macOS / Linux：数据目录下的 unix socket（[`ENDPOINT_FILE`]），套接字随 +//! `SCM_RIGHTS` 辅助消息走；对端用 `getpeereid` / `SO_PEERCRED` 核对是同一个用户。 +//! - Windows：回环上的随机端口，号码写进 [`ENDPOINT_FILE`]；连上之后先交 +//! [`TOKEN_FILE`] 里的口令。套接字由发送方对接收方的 pid 调 +//! `WSADuplicateSocketW`，把得到的 `WSAPROTOCOL_INFOW` 随消息发过去，接收方 +//! `WSASocketW` 还原。 +//! +//! 消息是带长度前缀的 JSON（[`Message`]）。 +//! +//! # 不做兼容 +//! +//! [`VERSION`] 是这条连接自己的版本号，和 `tw-api` 的协议号无关。**格式有任何 +//! 改动都升它**；双方不同就不交接，退回「等请求结束再重启」（[`Message::Refused`]）。 +//! 快照、结算的内容由网关定义，这里只当成 JSON 搬运，也跟着这个号走。 + +use std::net::SocketAddr; +use std::path::Path; + +/// 交接连接的版本号。**格式有任何改动都升它**，见模块文档。 +pub const VERSION: u32 = 1; + +/// 交接端点在数据目录里的文件名：unix 上是 socket 本身，Windows 上是写着端口号的文件。 +#[cfg(unix)] +pub const ENDPOINT_FILE: &str = "twcore.handover.sock"; +#[cfg(not(unix))] +pub const ENDPOINT_FILE: &str = "twcore.handover.port"; + +/// Windows 上交接连接的口令文件。回环端口挡不住同机的任何进程，门靠它。 +#[cfg(not(unix))] +pub const TOKEN_FILE: &str = "twcore.handover.token"; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error(transparent)] + Io(#[from] std::io::Error), + /// 对端发来的不是一条看得懂的消息 + #[error("the handover peer sent something unreadable: {0}")] + Malformed(String), + /// 对端不是同一个用户，或者口令不对 + #[error("the handover peer is not allowed: {0}")] + Denied(String), + /// 套接字传不过去（句柄无效、权限不足、对端 pid 不对） + #[error("a listening socket could not be handed over: {0}")] + Transfer(String), +} + +pub type Result = std::result::Result; + +/// 交接的一个监听。 +#[derive(Debug)] +pub enum Socket { + Tcp(std::net::TcpListener), + #[cfg(unix)] + Unix(std::os::unix::net::UnixListener), +} + +/// 一个监听是做什么的，以及它对应配置里的哪个地址。随 [`Message::Listeners`]、 +/// [`Message::Control`] 走，和消息里附带的套接字**按顺序一一对应**。 +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct ListenerInfo { + pub role: Role, + /// 配置算出来的地址（网关监听比对「变没变」用它）。控制面的 unix socket 没有 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub want: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Role { + /// 数据面 + Gateway, + /// 远程控制端口（tw-link） + Remote, + /// 本机控制面 + Control, +} + +/// 交接连接上的一条消息。顺序见方案第 4 节： +/// +/// ```text +/// 新 → 旧 Hello +/// 旧 → 新 Refused（版本不同，到此为止）| Control + 控制面套接字 +/// 旧 → 新 Snapshot +/// 新 → 旧 WantListeners （新 core 的配置读进来了、运行时建好了） +/// 旧 → 新 Snapshot （再导出一次：暂停续期之后的 OAuth、此刻的状态） +/// 旧 → 新 Listeners + 网关、远程控制端口的套接字 +/// 新 → 旧 Commit （新 core 已在全部监听上 accept） +/// 旧 → 新 CommitAck （旧 core 收下了提交、不再 accept；从此不可逆） +/// 旧 → 新 Settled* / InFlight* （排空期，每个请求结束一条） +/// 旧 → 新 Bye （马上退出） +/// 新 → 旧 Fallback （任何时候：请按现行方式排空后退出，新的自己绑） +/// ``` +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +#[serde(tag = "t", rename_all = "snake_case")] +pub enum Message { + Hello { + version: u32, + core_version: String, + pid: u32, + }, + Refused { + reason: String, + }, + Control { + info: ListenerInfo, + }, + Snapshot { + data: serde_json::Value, + }, + WantListeners, + Listeners { + infos: Vec, + }, + Commit, + /// 旧 core 收下了提交。**新 core 收到它（或者确认旧进程已经没了）才算提交**：旧 core 在 + /// 等提交的时限到了之后放弃、又收到一个迟到的 Commit 时不回它，两边就不会都以为自己是主 + CommitAck, + Settled { + data: serde_json::Value, + }, + InFlight { + n: u64, + }, + Bye, + Fallback, +} + +/// 单帧上限。快照可能较大（几千条亲和、限额计数），但不该没有边：长度前缀 +/// 读错了时宁可报错，也别照着一个 4 GiB 的数去分配。 +const MAX_FRAME: usize = 16 << 20; + +/// 一帧最多附几个套接字。网关每张网卡一个、再加远程控制端口，远远够了； +/// 也在 Linux 一条辅助消息的上限（`SCM_MAX_FD` = 253）之内。 +const MAX_SOCKETS: usize = 64; + +/// 同时在过门的连接最多几个，见 [`Endpoint::admit`]。 +const MAX_CHECKS: usize = 64; + +/// 端点的门：连上来的连接要过了它才交出去。 +#[derive(Clone)] +enum Gate { + /// unix：同一用户已经由内核核对过（`Listener::accept` 里），直接放行 + #[cfg(unix)] + Open, + /// 第一帧必须是这个口令，`wait` 之内交上来 + #[cfg(any(windows, test))] + Token { + token: std::sync::Arc<[u8]>, + wait: std::time::Duration, + }, +} + +/// 口令帧的长度上限（口令本身是 64 个十六进制字符）。没过门的连接不配让我们 +/// 分配 16 MiB。 +#[cfg(any(windows, test))] +const TOKEN_FRAME_MAX: usize = 256; + +impl Gate { + async fn check(self, link: Link) -> Option { + match self { + #[cfg(unix)] + Gate::Open => Some(link), + #[cfg(any(windows, test))] + Gate::Token { token, wait } => { + let mut link = link; + let frame = link.read_frame(TOKEN_FRAME_MAX); + match tokio::time::timeout(wait, frame).await { + Ok(Ok(t)) if token::eq(&t, &token) => Some(link), + _ => None, + } + } + } + } +} + +/// 线上的一帧：消息外面包一层，带上套接字的附件信息。两端都在本 crate， +/// 不对外。 +#[derive(serde::Serialize)] +struct EnvelopeOut<'a> { + m: &'a Message, + /// 随这一帧走的套接字个数（unix：`SCM_RIGHTS` 里的句柄数） + #[serde(skip_serializing_if = "is_zero")] + n: usize, + /// Windows：每个套接字的 `WSAPROTOCOL_INFOW`，十六进制 + #[serde(skip_serializing_if = "Vec::is_empty")] + w: Vec, +} + +#[derive(serde::Deserialize)] +struct EnvelopeIn { + m: Message, + #[serde(default)] + n: usize, + #[serde(default)] + w: Vec, +} + +fn is_zero(n: &usize) -> bool { + *n == 0 +} + +/// 旧 core 那一头：arm 之后开着，等接替者连上来。 +pub struct Endpoint { + listener: sys::Listener, +} + +impl Endpoint { + /// 在数据目录里开交接端点。已有的同名文件（上次没退干净）先删掉。 + pub async fn bind(dir: &Path) -> Result { + Ok(Endpoint { + listener: sys::Listener::bind(dir).await?, + }) + } + + /// 等一个接替者。**核对过身份才交出 [`Link`]**（unix：同一用户；Windows：口令）； + /// 不对的连接直接断开，接着等下一个。 + /// + /// 可以被取消（放在 `select!` 里）：正在核对的连接随之丢掉。 + pub async fn accept(&self) -> Result { + self.admit(self.listener.gate()).await + } + + /// 接连接、过门，第一个过了门的交出去。 + /// + /// **各个连接并发地过门**：Windows 的门是「5 秒内交口令」，而回环端口同机谁都 + /// 连得上 —— 一个一个核对的话，反复连上来不说话就能把真正的接替者一直挡在外面。 + /// 同时核对的个数有上限，满了就**丢掉最早的那个**再接新的：真正的接替者连上 + /// 就交口令（回环上是微秒级），要把它挤掉，得在这点时间里连进来上限那么多个。 + /// 满了停下不接是不行的：那样占满名额就能一直堵着。 + /// + /// unix 的门是一次 `getsockopt`，当场就有结论，不存在被拖住的问题；走同一条 + /// 路只是为了两边一套代码。 + async fn admit(&self, gate: Gate) -> Result { + use std::collections::VecDeque; + use tokio::task::{AbortHandle, JoinSet}; + + // JoinSet 被丢掉时（返回、取消）会中止里面所有任务，没过门的连接跟着关掉 + let mut checks: JoinSet> = JoinSet::new(); + let mut order: VecDeque = VecDeque::new(); + loop { + tokio::select! { + r = self.listener.accept() => { + let Some((io, pid)) = r? else { + continue; + }; + if order.len() >= MAX_CHECKS { + order.pop_front().expect("full").abort(); + } + let gate = gate.clone(); + order.push_back(checks.spawn(gate.check(Link::new(io, pid)))); + } + Some(done) = checks.join_next_with_id(), if !checks.is_empty() => { + let (id, link) = match done { + Ok((id, link)) => (id, link), + // 被挤掉的那个 + Err(e) => (e.id(), None), + }; + order.retain(|h| h.id() != id); + if let Some(link) = link { + return Ok(link); + } + } + } + } + } + + /// 关掉端点，删掉文件。被丢掉时也做同一件事。 + pub fn close(self) {} +} + +/// 新 core 那一头：去连锁的持有者。**`Ok(None)` = 它没有 arm**（端点不存在或连不上）， +/// 调用方照旧报「已有实例在运行」。 +pub async fn connect(dir: &Path) -> Result> { + let Some((io, pid)) = sys::connect(dir).await? else { + return Ok(None); + }; + #[allow(unused_mut)] // Windows 上要先在它上面交口令 + let mut link = Link::new(io, pid); + #[cfg(windows)] + { + let token = sys::read_token(dir)?; + link.write_frame(token.as_bytes(), &[]).await?; + } + Ok(Some(link)) +} + +/// 一条交接连接。 +pub struct Link { + io: sys::Stream, + /// 读进来、还没拼成整帧的字节 + rbuf: Vec, + peer_pid: Option, +} + +impl Link { + fn new(io: sys::Stream, peer_pid: Option) -> Link { + Link { + io, + rbuf: Vec::new(), + peer_pid, + } + } + + /// 对端的 pid（unix 上从内核问到的；Windows 上取 `Hello` 里报的）。 + pub fn peer_pid(&self) -> Option { + self.peer_pid + } + + /// 发一条不带套接字的消息。 + pub async fn send(&mut self, m: &Message) -> Result<()> { + self.send_with(m, &[], 0).await + } + + /// 发一条消息，附带这几个套接字（顺序即 `infos` 的顺序）。`peer_pid` 只在 + /// Windows 上用得到（`WSADuplicateSocketW` 要它）。**发送方自己的句柄不受影响**， + /// 照样能 accept，直到它自己关掉。 + pub async fn send_with( + &mut self, + m: &Message, + sockets: &[&Socket], + peer_pid: u32, + ) -> Result<()> { + if sockets.len() > MAX_SOCKETS { + return Err(Error::Transfer(format!( + "{} sockets in one message, at most {MAX_SOCKETS}", + sockets.len() + ))); + } + let env = EnvelopeOut { + m, + n: sockets.len(), + w: sys::export(sockets, peer_pid)?, + }; + let body = serde_json::to_vec(&env).map_err(|e| Error::Malformed(e.to_string()))?; + self.write_frame(&body, sockets).await + } + + /// 收一条消息和它附带的套接字。对端关闭连接时返回 `Io(UnexpectedEof)`。 + pub async fn recv(&mut self) -> Result<(Message, Vec)> { + let body = self.read_frame(MAX_FRAME).await?; + let env: EnvelopeIn = + serde_json::from_slice(&body).map_err(|e| Error::Malformed(e.to_string()))?; + if env.n > MAX_SOCKETS { + return Err(Error::Malformed(format!( + "{} sockets in one message", + env.n + ))); + } + let sockets = self.io.import(env.n, &env.w)?; + #[cfg(windows)] + if let Message::Hello { pid, .. } = &env.m { + // Windows 问不到对端 pid，只能信它自己报的 —— 对端已经交过口令， + // 是同一个用户；报错了 pid 也只是复制不出套接字 + self.peer_pid = Some(*pid); + } + Ok((env.m, sockets)) + } + + /// 写一帧：4 字节大端长度 + 正文。unix 上套接字附在这一帧的第一个字节上。 + async fn write_frame(&mut self, body: &[u8], sockets: &[&Socket]) -> Result<()> { + if body.len() > MAX_FRAME { + return Err(Error::Malformed(format!( + "a message of {} bytes is over the {MAX_FRAME}-byte limit", + body.len() + ))); + } + let mut frame = Vec::with_capacity(4 + body.len()); + frame.extend_from_slice(&(body.len() as u32).to_be_bytes()); + frame.extend_from_slice(body); + self.io.write_frame(&frame, sockets).await + } + + /// 读一整帧的正文，长于 `max` 的报 `Malformed`。 + async fn read_frame(&mut self, max: usize) -> Result> { + loop { + if self.rbuf.len() >= 4 { + let len = u32::from_be_bytes(self.rbuf[..4].try_into().expect("4 bytes")) as usize; + if len > max { + return Err(Error::Malformed(format!( + "a message of {len} bytes is over the {max}-byte limit" + ))); + } + if self.rbuf.len() >= 4 + len { + let body = self.rbuf[4..4 + len].to_vec(); + self.rbuf.drain(..4 + len); + return Ok(body); + } + // 知道要多长了，一次留够，别边读边翻倍 + self.rbuf.reserve(4 + len - self.rbuf.len()); + } + if self.io.read_some(&mut self.rbuf).await? == 0 { + return Err(Error::Io(std::io::ErrorKind::UnexpectedEof.into())); + } + } + } +} + +#[cfg(any(windows, test))] +mod token; + +#[cfg(unix)] +#[path = "unix.rs"] +mod sys; + +#[cfg(windows)] +#[path = "windows.rs"] +mod sys; + +#[cfg(test)] +mod tests; diff --git a/crates/tw-handover/src/tests.rs b/crates/tw-handover/src/tests.rs new file mode 100644 index 00000000..94164008 --- /dev/null +++ b/crates/tw-handover/src/tests.rs @@ -0,0 +1,564 @@ +use std::collections::BTreeSet; +use std::io::{Read, Write}; +use std::net::{Ipv4Addr, TcpListener, TcpStream}; +use std::time::{Duration, Instant}; + +use super::*; + +/// unix socket 的路径有长度上限（macOS 104 字节），系统临时目录可能很深， +/// 所以 unix 上固定放 /tmp 底下。 +fn tempdir() -> tempfile::TempDir { + #[cfg(unix)] + return tempfile::Builder::new() + .prefix("twho") + .tempdir_in("/tmp") + .unwrap(); + #[cfg(not(unix))] + return tempfile::tempdir().unwrap(); +} + +/// 开端点、连上去，返回（旧 core 那一头，新 core 那一头）。 +async fn pair(dir: &Path) -> (Endpoint, Link, Link) { + let ep = Endpoint::bind(dir).await.unwrap(); + let (old, new) = tokio::join!(ep.accept(), connect(dir)); + (ep, old.unwrap(), new.unwrap().expect("armed")) +} + +/// Windows 上 `Socket` 只有 TCP 一种，模式匹配是不可反驳的。 +#[allow(irrefutable_let_patterns)] +fn into_tcp(s: Socket) -> TcpListener { + let Socket::Tcp(l) = s else { + panic!("expected a TCP listener") + }; + l +} + +fn me() -> u32 { + std::process::id() +} + +fn hello() -> Message { + Message::Hello { + version: VERSION, + core_version: "0.0.0-test".into(), + pid: me(), + } +} + +/// 在收到的监听上把排队的连接都接出来，读出每个连接发来的那一个字节。 +fn drain_ids(l: &TcpListener, want: usize) -> BTreeSet { + l.set_nonblocking(true).unwrap(); + let deadline = Instant::now() + Duration::from_secs(5); + let mut ids = BTreeSet::new(); + while ids.len() < want { + match l.accept() { + Ok((mut s, _)) => { + s.set_nonblocking(false).unwrap(); + let mut b = [0u8]; + s.read_exact(&mut b).unwrap(); + ids.insert(b[0]); + } + Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => { + assert!( + Instant::now() < deadline, + "only {} of {want} connections arrived", + ids.len() + ); + std::thread::sleep(Duration::from_millis(10)); + } + Err(e) => panic!("{e}"), + } + } + ids +} + +#[tokio::test] +async fn not_armed_is_none() { + let dir = tempdir(); + assert!(connect(dir.path()).await.unwrap().is_none()); +} + +#[cfg(unix)] +#[tokio::test] +async fn stale_socket_file_is_none() { + let dir = tempdir(); + let path = dir.path().join(ENDPOINT_FILE); + // 绑过、没人在听了，文件还在 —— 上次没退干净的样子 + drop(std::os::unix::net::UnixListener::bind(&path).unwrap()); + assert!(path.exists()); + assert!(connect(dir.path()).await.unwrap().is_none()); + // 重新 bind 时把残留删掉 + let ep = Endpoint::bind(dir.path()).await.unwrap(); + let (old, new) = tokio::join!(ep.accept(), connect(dir.path())); + assert!(old.is_ok() && new.unwrap().is_some()); +} + +#[tokio::test] +async fn bind_connect_close() { + let dir = tempdir(); + let (ep, old, new) = pair(dir.path()).await; + let file = dir.path().join(ENDPOINT_FILE); + assert!(file.exists()); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mode = std::fs::metadata(&file).unwrap().permissions().mode(); + assert_eq!(mode & 0o777, 0o600); + // 同一进程：两头问到的都是自己 + assert_eq!(old.peer_pid(), Some(me())); + assert_eq!(new.peer_pid(), Some(me())); + } + #[cfg(windows)] + { + assert!(dir.path().join(TOKEN_FILE).exists()); + assert_eq!(old.peer_pid(), None); + } + drop((old, new)); + ep.close(); + assert!(!file.exists()); + #[cfg(windows)] + assert!(!dir.path().join(TOKEN_FILE).exists()); + assert!(connect(dir.path()).await.unwrap().is_none()); +} + +#[tokio::test] +async fn closing_does_not_remove_a_newer_endpoint() { + // 排空中的旧 core 关端点时，新 core 可能已经在同一路径上开了自己的（H8） + let dir = tempdir(); + let first = Endpoint::bind(dir.path()).await.unwrap(); + let second = Endpoint::bind(dir.path()).await.unwrap(); + first.close(); + assert!(dir.path().join(ENDPOINT_FILE).exists()); + let (old, new) = tokio::join!(second.accept(), connect(dir.path())); + assert!(old.is_ok() && new.unwrap().is_some()); + second.close(); + assert!(!dir.path().join(ENDPOINT_FILE).exists()); +} + +#[tokio::test] +async fn messages_round_trip() { + let dir = tempdir(); + let (_ep, mut old, mut new) = pair(dir.path()).await; + + new.send(&hello()).await.unwrap(); + let (m, s) = old.recv().await.unwrap(); + assert_eq!(m, hello()); + assert!(s.is_empty()); + // Windows 上 pid 从 Hello 来 + assert_eq!(old.peer_pid(), Some(me())); + + // 大快照：一帧要分好几次读 + let big = serde_json::json!({ "affinity": "x".repeat(3 << 20) }); + let from_old = [ + Message::Refused { + reason: "version".into(), + }, + Message::Control { + info: ListenerInfo { + role: Role::Control, + want: None, + }, + }, + Message::Snapshot { data: big }, + Message::Listeners { + infos: vec![ + ListenerInfo { + role: Role::Gateway, + want: Some("127.0.0.1:8080".parse().unwrap()), + }, + ListenerInfo { + role: Role::Remote, + want: Some("[::]:9000".parse().unwrap()), + }, + ], + }, + Message::Settled { + data: serde_json::json!({ "key": "k", "cost": 1.5 }), + }, + Message::InFlight { n: 3 }, + Message::Bye, + ]; + let from_new = [Message::WantListeners, Message::Commit, Message::Fallback]; + + // 一口气发完再收：多帧挤在一次读里也要拆得对 + let send = async { + for m in &from_old { + old.send(m).await.unwrap(); + } + old + }; + let recv = async { + let mut got = Vec::new(); + for _ in 0..from_old.len() { + got.push(new.recv().await.unwrap().0); + } + (new, got) + }; + let (mut old, (mut new, got)) = tokio::join!(send, recv); + assert_eq!(got, from_old); + + for m in &from_new { + new.send(m).await.unwrap(); + } + for m in &from_new { + assert_eq!(&old.recv().await.unwrap().0, m); + } + + drop(new); + match old.recv().await { + Err(Error::Io(e)) => assert_eq!(e.kind(), std::io::ErrorKind::UnexpectedEof), + other => panic!("expected EOF, got {other:?}"), + } +} + +/// **关键性质**：排在队列里、还没被 accept 的连接，随套接字一起交过去，一个不丢； +/// 发送方交出之后自己照样能 accept，关掉自己的句柄也不影响接收方。 +#[tokio::test] +async fn queued_tcp_connections_survive_the_handover() { + let dir = tempdir(); + let (_ep, mut old, mut new) = pair(dir.path()).await; + + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap(); + let addr = listener.local_addr().unwrap(); + const N: u8 = 6; + let clients: Vec = (0..N) + .map(|i| { + let mut c = TcpStream::connect(addr).unwrap(); + c.write_all(&[i]).unwrap(); + c + }) + .collect(); + + let socket = Socket::Tcp(listener); + let infos = vec![ListenerInfo { + role: Role::Gateway, + want: Some(addr), + }]; + old.send_with( + &Message::Listeners { + infos: infos.clone(), + }, + &[&socket], + me(), + ) + .await + .unwrap(); + let (m, mut sockets) = new.recv().await.unwrap(); + assert_eq!(m, Message::Listeners { infos }); + assert_eq!(sockets.len(), 1); + let received = into_tcp(sockets.pop().unwrap()); + assert_eq!(received.local_addr().unwrap(), addr); + + // 发送方的句柄不受影响：它还能接一个 + let original = into_tcp(socket); + let (mut s, _) = original.accept().unwrap(); + let mut b = [0u8]; + s.read_exact(&mut b).unwrap(); + drop(original); + + let mut ids = drain_ids(&received, (N - 1) as usize); + ids.insert(b[0]); + assert_eq!(ids, (0..N).collect::>()); + + // 交接之后来的新连接也落在接收方 + let mut late = TcpStream::connect(addr).unwrap(); + late.write_all(&[N]).unwrap(); + assert_eq!(drain_ids(&received, 1), BTreeSet::from([N])); + drop(clients); +} + +#[cfg(unix)] +#[tokio::test] +async fn queued_unix_connections_survive_the_handover() { + use std::os::unix::net::{UnixListener, UnixStream}; + + let dir = tempdir(); + let (_ep, mut old, mut new) = pair(dir.path()).await; + + let path = dir.path().join("control.sock"); + let listener = UnixListener::bind(&path).unwrap(); + let clients: Vec = (0..4u8) + .map(|i| { + let mut c = UnixStream::connect(&path).unwrap(); + c.write_all(&[i]).unwrap(); + c + }) + .collect(); + + let socket = Socket::Unix(listener); + let info = ListenerInfo { + role: Role::Control, + want: None, + }; + old.send_with(&Message::Control { info }, &[&socket], 0) + .await + .unwrap(); + drop(socket); + let (_, mut sockets) = new.recv().await.unwrap(); + let Some(Socket::Unix(received)) = sockets.pop() else { + panic!("expected a unix listener"); + }; + let mut ids = BTreeSet::new(); + for _ in 0..4 { + let (mut s, _) = received.accept().unwrap(); + let mut b = [0u8]; + s.read_exact(&mut b).unwrap(); + ids.insert(b[0]); + } + assert_eq!(ids, (0..4).collect::>()); + drop(clients); +} + +/// 几个套接字一帧发、前后夹着不带套接字的消息：顺序和归属都要对得上。 +#[tokio::test] +async fn several_sockets_keep_their_order() { + let dir = tempdir(); + let (_ep, mut old, mut new) = pair(dir.path()).await; + + let a = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap(); + let b = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap(); + let c = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap(); + let addrs: Vec<_> = [&a, &b, &c] + .iter() + .map(|l| l.local_addr().unwrap()) + .collect(); + let (a, b, c) = (Socket::Tcp(a), Socket::Tcp(b), Socket::Tcp(c)); + + old.send(&Message::InFlight { n: 1 }).await.unwrap(); + old.send_with(&Message::Commit, &[&a, &b], me()) + .await + .unwrap(); + old.send(&Message::InFlight { n: 2 }).await.unwrap(); + old.send_with(&Message::Bye, &[&c], me()).await.unwrap(); + + let port = |s: &Socket| match s { + Socket::Tcp(l) => l.local_addr().unwrap(), + #[cfg(unix)] + Socket::Unix(_) => panic!("expected TCP"), + }; + let (m, s) = new.recv().await.unwrap(); + assert_eq!((m, s.len()), (Message::InFlight { n: 1 }, 0)); + let (m, s) = new.recv().await.unwrap(); + assert_eq!(m, Message::Commit); + assert_eq!(s.iter().map(port).collect::>(), addrs[..2]); + let (m, s) = new.recv().await.unwrap(); + assert_eq!((m, s.len()), (Message::InFlight { n: 2 }, 0)); + let (m, s) = new.recv().await.unwrap(); + assert_eq!(m, Message::Bye); + assert_eq!(s.iter().map(port).collect::>(), addrs[2..]); +} + +#[tokio::test] +async fn oversized_message_is_not_sent() { + let dir = tempdir(); + let (_ep, mut old, _new) = pair(dir.path()).await; + let data = serde_json::Value::String("x".repeat(MAX_FRAME)); + assert!(matches!( + old.send(&Message::Snapshot { data }).await, + Err(Error::Malformed(_)) + )); +} + +#[cfg(unix)] +#[tokio::test] +async fn oversized_frame_is_rejected() { + use tokio::io::AsyncWriteExt; + + let dir = tempdir(); + let ep = Endpoint::bind(dir.path()).await.unwrap(); + let raw = tokio::net::UnixStream::connect(dir.path().join(ENDPOINT_FILE)); + let (old, raw) = tokio::join!(ep.accept(), raw); + let (mut old, mut raw) = (old.unwrap(), raw.unwrap()); + raw.write_all(&((MAX_FRAME as u32) + 1).to_be_bytes()) + .await + .unwrap(); + assert!(matches!(old.recv().await, Err(Error::Malformed(_)))); +} + +#[cfg(unix)] +#[tokio::test] +async fn garbage_is_malformed() { + use tokio::io::AsyncWriteExt; + + let dir = tempdir(); + let ep = Endpoint::bind(dir.path()).await.unwrap(); + let raw = tokio::net::UnixStream::connect(dir.path().join(ENDPOINT_FILE)); + let (old, raw) = tokio::join!(ep.accept(), raw); + let (mut old, mut raw) = (old.unwrap(), raw.unwrap()); + let body = br#"{"m":{"t":"no_such_message"}}"#; + raw.write_all(&(body.len() as u32).to_be_bytes()) + .await + .unwrap(); + raw.write_all(body).await.unwrap(); + assert!(matches!(old.recv().await, Err(Error::Malformed(_)))); +} + +#[test] +fn token_comparison() { + assert!(token::eq(b"abcd", b"abcd")); + assert!(!token::eq(b"abcd", b"abce")); + assert!(!token::eq(b"abcd", b"abc")); + assert!(!token::eq(b"", b"a")); + assert!(token::eq(b"", b"")); +} + +#[test] +fn hex_round_trip() { + let bytes: Vec = (0..=255).collect(); + let s = token::hex(&bytes); + assert_eq!(s.len(), 512); + assert_eq!(token::unhex(&s).unwrap(), bytes); + assert_eq!(token::unhex("0aFf").unwrap(), vec![0x0a, 0xff]); + assert!(token::unhex("abc").is_none()); + assert!(token::unhex("zz").is_none()); +} + +/// 网关会把连接放进 spawn 出去的任务里用：这几个 future 必须是 Send。 +#[allow(dead_code)] +fn futures_are_send(ep: &Endpoint, link: &mut Link, dir: &Path) { + fn send(_: T) {} + send(Endpoint::bind(dir)); + send(ep.accept()); + send(connect(dir)); + send(link.recv()); +} + +#[allow(dead_code)] +fn send_with_is_send(link: &mut Link, s: &Socket) { + fn send(_: T) {} + send(link.send_with(&Message::Bye, &[s], 0)); +} + +/// 连上端点但什么都不说：只是一条原始连接。 +#[cfg(unix)] +async fn raw_connect(dir: &Path) -> tokio::net::UnixStream { + tokio::net::UnixStream::connect(dir.join(ENDPOINT_FILE)) + .await + .unwrap() +} + +#[cfg(windows)] +async fn raw_connect(dir: &Path) -> tokio::net::TcpStream { + let port: u16 = std::fs::read_to_string(dir.join(ENDPOINT_FILE)) + .unwrap() + .trim() + .parse() + .unwrap(); + tokio::net::TcpStream::connect((Ipv4Addr::LOCALHOST, port)) + .await + .unwrap() +} + +fn frame(body: &[u8]) -> Vec { + let mut f = (body.len() as u32).to_be_bytes().to_vec(); + f.extend_from_slice(body); + f +} + +/// 口令门**并发**核对：一堆连上来不说话的（比同时核对的上限还多）、一个交错口令的， +/// 都挡不住随后交对口令的那一个，也不会被错放进来。Windows 的门就是这个，这里在 +/// 各个平台上都用它跑一遍。 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_gate_is_not_held_up() { + use tokio::io::AsyncWriteExt; + + let dir = tempdir(); + let ep = std::sync::Arc::new(Endpoint::bind(dir.path()).await.unwrap()); + let gate = Gate::Token { + token: std::sync::Arc::from(&b"secret"[..]), + wait: Duration::from_secs(5), + }; + let start = Instant::now(); + let admitted = tokio::spawn({ + let ep = ep.clone(); + async move { ep.admit(gate).await } + }); + + let mut silent = Vec::new(); + for _ in 0..(MAX_CHECKS + 36) { + silent.push(raw_connect(dir.path()).await); + } + let mut wrong = raw_connect(dir.path()).await; + wrong.write_all(&frame(b"guess!")).await.unwrap(); + wrong + .write_all(&frame(br#"{"m":{"t":"fallback"}}"#)) + .await + .unwrap(); + + let mut right = raw_connect(dir.path()).await; + right.write_all(&frame(b"secret")).await.unwrap(); + right + .write_all(&frame(br#"{"m":{"t":"bye"}}"#)) + .await + .unwrap(); + + let mut link = tokio::time::timeout(Duration::from_secs(3), admitted) + .await + .expect("held up by silent connections") + .unwrap() + .unwrap(); + assert!(start.elapsed() < Duration::from_secs(3)); + // 放进来的是交对口令的那一个 + assert_eq!(link.recv().await.unwrap().0, Message::Bye); + drop((silent, wrong)); +} + +/// 真端点上的同一件事：Windows 上门是口令，unix 上门由内核当场判定。 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn accept_is_not_held_up_by_silent_connections() { + let dir = tempdir(); + let ep = std::sync::Arc::new(Endpoint::bind(dir.path()).await.unwrap()); + let accepted = tokio::spawn({ + let ep = ep.clone(); + async move { ep.accept().await } + }); + #[cfg(windows)] + let silent = { + let mut v = Vec::new(); + for _ in 0..(MAX_CHECKS + 36) { + v.push(raw_connect(dir.path()).await); + } + v + }; + let mut new = connect(dir.path()).await.unwrap().expect("armed"); + new.send(&hello()).await.unwrap(); + let mut old = tokio::time::timeout(Duration::from_secs(3), accepted) + .await + .expect("held up") + .unwrap() + .unwrap(); + assert_eq!(old.recv().await.unwrap().0, hello()); + #[cfg(windows)] + drop(silent); +} + +/// Windows：端点的监听、还原出来的监听都不能被子进程继承。 +#[cfg(windows)] +#[tokio::test] +async fn sockets_are_not_inheritable() { + use std::os::windows::io::AsRawSocket; + use windows_sys::Win32::Foundation::{GetHandleInformation, HANDLE, HANDLE_FLAG_INHERIT}; + + fn inheritable(raw: std::os::windows::io::RawSocket) -> bool { + let mut flags = 0u32; + // SAFETY: raw 是有效的套接字句柄，出参是本地变量 + assert_ne!( + unsafe { GetHandleInformation(raw as HANDLE, &mut flags) }, + 0 + ); + flags & HANDLE_FLAG_INHERIT != 0 + } + + let dir = tempdir(); + let (ep, mut old, mut new) = pair(dir.path()).await; + assert!(!inheritable(ep.listener.raw())); + + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap(); + let socket = Socket::Tcp(listener); + old.send_with(&Message::Commit, &[&socket], me()) + .await + .unwrap(); + let (_, mut sockets) = new.recv().await.unwrap(); + let received = into_tcp(sockets.pop().unwrap()); + assert!(!inheritable(received.as_raw_socket())); +} diff --git a/crates/tw-handover/src/token.rs b/crates/tw-handover/src/token.rs new file mode 100644 index 00000000..f2970c31 --- /dev/null +++ b/crates/tw-handover/src/token.rs @@ -0,0 +1,39 @@ +//! Windows 交接口令的纯函数部分：编码与比较。单独放，是为了在 Mac 上也能测。 + +/// 常数时间比较：耗时只取决于长度，不取决于从第几个字节开始不同。 +/// 长度不同直接判否 —— 口令的长度是固定的，不是秘密。 +pub(crate) fn eq(a: &[u8], b: &[u8]) -> bool { + if a.len() != b.len() { + return false; + } + let diff = a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)); + std::hint::black_box(diff) == 0 +} + +pub(crate) fn hex(bytes: &[u8]) -> String { + const DIGITS: &[u8; 16] = b"0123456789abcdef"; + let mut s = String::with_capacity(bytes.len() * 2); + for b in bytes { + s.push(DIGITS[(b >> 4) as usize] as char); + s.push(DIGITS[(b & 0xf) as usize] as char); + } + s +} + +pub(crate) fn unhex(s: &str) -> Option> { + fn digit(c: u8) -> Option { + match c { + b'0'..=b'9' => Some(c - b'0'), + b'a'..=b'f' => Some(c - b'a' + 10), + b'A'..=b'F' => Some(c - b'A' + 10), + _ => None, + } + } + let s = s.as_bytes(); + if s.len() % 2 != 0 { + return None; + } + s.chunks_exact(2) + .map(|p| Some(digit(p[0])? << 4 | digit(p[1])?)) + .collect() +} diff --git a/crates/tw-handover/src/unix.rs b/crates/tw-handover/src/unix.rs new file mode 100644 index 00000000..3b4bf7a4 --- /dev/null +++ b/crates/tw-handover/src/unix.rs @@ -0,0 +1,390 @@ +//! macOS / Linux：数据目录下的 unix socket，套接字随 `SCM_RIGHTS` 走。 +//! +//! tokio 的 `UnixStream` 不收发辅助消息，所以读写都自己调 `recvmsg` / `sendmsg`， +//! 外面套 `async_io` —— 它在 `WouldBlock` 时让出、等就绪再试，不会堵住运行时。 + +use std::collections::VecDeque; +use std::io; +use std::mem; +use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd}; +use std::os::unix::fs::{MetadataExt, PermissionsExt}; +use std::path::{Path, PathBuf}; + +use tokio::io::Interest; +use tokio::net::{UnixListener, UnixStream}; + +use crate::{ENDPOINT_FILE, Error, Gate, MAX_SOCKETS, Result, Socket}; + +/// 一次 `recvmsg` 最多读这么多正文。 +const CHUNK: usize = 64 * 1024; + +#[cfg(any(target_os = "linux", target_os = "android"))] +const SEND_FLAGS: libc::c_int = libc::MSG_NOSIGNAL; +#[cfg(not(any(target_os = "linux", target_os = "android")))] +const SEND_FLAGS: libc::c_int = 0; + +/// Linux 能让内核直接给收到的句柄打上 close-on-exec；macOS 没有这个标志， +/// 收到后补一次 `fcntl`（两边都补，不差这一次系统调用）。 +/// +/// **macOS 上因此有一个窗口**：`recvmsg` 返回到 `fcntl` 之间，别的线程如果恰好 +/// fork + exec，子进程会带走一份监听，我们退出后它还占着端口。没有干净的修法 +/// （macOS 既没有 `MSG_CMSG_CLOEXEC`，也不能让 fork 等这里），而这个窗口可以接受： +/// 它只在交接时打开一次、只有几条指令长；core 平时几乎不起子进程，接管阶段 +/// （读配置、建运行时）更不会。 +#[cfg(any(target_os = "linux", target_os = "android"))] +const RECV_FLAGS: libc::c_int = libc::MSG_CMSG_CLOEXEC; +#[cfg(not(any(target_os = "linux", target_os = "android")))] +const RECV_FLAGS: libc::c_int = 0; + +pub(crate) struct Listener { + listener: UnixListener, + path: PathBuf, + /// 绑定时那个 socket 文件的 (dev, ino)。删文件前比一比：**排空中的旧 core + /// 不能删掉新 core 后来在同一路径上开的端点**（方案 H8：排空中再更新） + id: Option<(u64, u64)>, +} + +impl Listener { + pub(crate) async fn bind(dir: &Path) -> Result { + let path = dir.join(ENDPOINT_FILE); + match std::fs::remove_file(&path) { + Ok(()) => {} + Err(e) if e.kind() == io::ErrorKind::NotFound => {} + Err(e) => return Err(e.into()), + } + let listener = UnixListener::bind(&path)?; + // 0600：只有当前用户能连。accept 时还会问内核对端是谁 —— 两道都在 + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?; + let id = std::fs::symlink_metadata(&path) + .ok() + .map(|m| (m.dev(), m.ino())); + Ok(Listener { listener, path, id }) + } + + /// 内核已经在 `accept` 里核对过用户了，门是开的。 + pub(crate) fn gate(&self) -> Gate { + Gate::Open + } + + /// `Ok(None)` = 连上来的不是同一个用户，已经断开。 + pub(crate) async fn accept(&self) -> Result)>> { + let (s, _) = self.listener.accept().await?; + match peer(&s) { + Ok((uid, pid)) if uid == own_uid() => Ok(Some((Stream::new(s), pid))), + _ => Ok(None), + } + } +} + +impl Drop for Listener { + fn drop(&mut self) { + let same = std::fs::symlink_metadata(&self.path) + .ok() + .map(|m| (m.dev(), m.ino())); + if same.is_some() && same == self.id { + let _ = std::fs::remove_file(&self.path); + } + } +} + +pub(crate) async fn connect(dir: &Path) -> Result)>> { + // 文件不在、没人在听（上次没删干净的残留）、连不上 —— 都是「没有 arm」 + let Ok(s) = UnixStream::connect(dir.join(ENDPOINT_FILE)).await else { + return Ok(None); + }; + let (uid, pid) = peer(&s)?; + if uid != own_uid() { + return Err(Error::Denied(format!( + "the handover endpoint belongs to uid {uid}" + ))); + } + Ok(Some((Stream::new(s), pid))) +} + +fn own_uid() -> u32 { + // SAFETY: geteuid 总是成功，没有参数 + unsafe { libc::geteuid() } +} + +/// 问内核对端的 (uid, pid)。 +#[cfg(any(target_os = "macos", target_os = "ios"))] +fn peer(s: &UnixStream) -> io::Result<(u32, Option)> { + let fd = s.as_raw_fd(); + let mut uid: libc::uid_t = 0; + let mut gid: libc::gid_t = 0; + // SAFETY: fd 是有效的 unix socket，两个出参都是本地变量 + if unsafe { libc::getpeereid(fd, &mut uid, &mut gid) } != 0 { + return Err(io::Error::last_os_error()); + } + let mut pid: libc::pid_t = 0; + let mut len = mem::size_of::() as libc::socklen_t; + // SAFETY: 出参是一个 pid_t，长度如实给出 + let ok = unsafe { + libc::getsockopt( + fd, + libc::SOL_LOCAL, + libc::LOCAL_PEERPID, + (&mut pid as *mut libc::pid_t).cast(), + &mut len, + ) + } == 0; + Ok((uid, (ok && pid > 0).then_some(pid as u32))) +} + +#[cfg(any(target_os = "linux", target_os = "android"))] +fn peer(s: &UnixStream) -> io::Result<(u32, Option)> { + let mut cred = libc::ucred { + pid: 0, + uid: 0, + gid: 0, + }; + let mut len = mem::size_of::() as libc::socklen_t; + // SAFETY: 出参是一个 ucred，长度如实给出 + if unsafe { + libc::getsockopt( + s.as_raw_fd(), + libc::SOL_SOCKET, + libc::SO_PEERCRED, + (&mut cred as *mut libc::ucred).cast(), + &mut len, + ) + } != 0 + { + return Err(io::Error::last_os_error()); + } + Ok((cred.uid, (cred.pid > 0).then_some(cred.pid as u32))) +} + +#[cfg(not(any( + target_os = "macos", + target_os = "ios", + target_os = "linux", + target_os = "android" +)))] +fn peer(s: &UnixStream) -> io::Result<(u32, Option)> { + let c = s.peer_cred()?; + Ok((c.uid(), c.pid().and_then(|p| u32::try_from(p).ok()))) +} + +pub(crate) struct Stream { + s: UnixStream, + /// 收到、还没认领的句柄。每帧的句柄附在它的第一个字节上，帧又是按序的， + /// 所以按先来后到认领就对得上 —— 哪怕一次读进了好几帧 + fds: VecDeque, +} + +impl Stream { + fn new(s: UnixStream) -> Stream { + Stream { + s, + fds: VecDeque::new(), + } + } + + pub(crate) async fn write_frame(&mut self, frame: &[u8], sockets: &[&Socket]) -> Result<()> { + let fds: Vec = sockets + .iter() + .map(|s| match s { + Socket::Tcp(l) => l.as_raw_fd(), + Socket::Unix(l) => l.as_raw_fd(), + }) + .collect(); + let fd = self.s.as_raw_fd(); + let mut off = 0; + let mut attach: &[RawFd] = &fds; + while off < frame.len() { + let rest = &frame[off..]; + let n = self + .s + .async_io(Interest::WRITABLE, || sendmsg(fd, rest, attach)) + .await?; + if n == 0 { + return Err(Error::Io(io::ErrorKind::WriteZero.into())); + } + // 只要发出去了一个字节，句柄就跟着它走了；后面的字节不再带 + attach = &[]; + off += n; + } + Ok(()) + } + + pub(crate) async fn read_some(&mut self, buf: &mut Vec) -> Result { + buf.reserve(CHUNK); + let fd = self.s.as_raw_fd(); + let spare = buf.spare_capacity_mut(); + // 地址按整数存：裸指针不是 Send，跨 await 拿着它会让整个 Link 的 + // future 不能 spawn + let (dst, cap) = (spare.as_mut_ptr() as usize, spare.len()); + let (n, fds, truncated) = self + .s + .async_io(Interest::READABLE, || recvmsg(fd, dst as *mut u8, cap)) + .await?; + // SAFETY: recvmsg 刚往 spare 的前 n 个字节里写了数据 + unsafe { buf.set_len(buf.len() + n) }; + if truncated { + // 收到的那几个已经在 fds 里，随它一起关掉；这条连接上的句柄对不上了 + drop(fds); + return Err(Error::Transfer( + "the sockets did not fit the ancillary buffer".into(), + )); + } + self.fds.extend(fds); + Ok(n) + } + + /// 认领这一帧的 n 个句柄，还原成监听。 + pub(crate) fn import(&mut self, n: usize, _w: &[String]) -> Result> { + if self.fds.len() < n { + return Err(Error::Transfer(format!( + "the message carries {n} sockets but {} arrived", + self.fds.len() + ))); + } + let fds: Vec = self.fds.drain(..n).collect(); + fds.into_iter().map(socket_from_fd).collect() + } +} + +/// unix 上套接字走辅助消息，帧里不用带别的。 +pub(crate) fn export(_sockets: &[&Socket], _peer_pid: u32) -> Result> { + Ok(Vec::new()) +} + +/// 按地址族把句柄还原成 TCP 或 unix 监听。 +fn socket_from_fd(fd: OwnedFd) -> Result { + let raw = fd.as_raw_fd(); + let mut ty: libc::c_int = 0; + let mut len = mem::size_of::() as libc::socklen_t; + // SAFETY: 出参是一个 c_int，长度如实给出 + if unsafe { + libc::getsockopt( + raw, + libc::SOL_SOCKET, + libc::SO_TYPE, + (&mut ty as *mut libc::c_int).cast(), + &mut len, + ) + } != 0 + { + return Err(Error::Transfer(format!( + "a received handle is not a socket: {}", + io::Error::last_os_error() + ))); + } + if ty != libc::SOCK_STREAM { + return Err(Error::Transfer(format!( + "a received socket is not a stream socket (type {ty})" + ))); + } + // SAFETY: sockaddr_storage 全零是合法值；出参长度如实给出 + let mut addr: libc::sockaddr_storage = unsafe { mem::zeroed() }; + let mut len = mem::size_of::() as libc::socklen_t; + if unsafe { + libc::getsockname( + raw, + (&mut addr as *mut libc::sockaddr_storage).cast(), + &mut len, + ) + } != 0 + { + return Err(Error::Transfer(io::Error::last_os_error().to_string())); + } + match addr.ss_family as libc::c_int { + libc::AF_INET | libc::AF_INET6 => Ok(Socket::Tcp(std::net::TcpListener::from(fd))), + libc::AF_UNIX => Ok(Socket::Unix(std::os::unix::net::UnixListener::from(fd))), + family => Err(Error::Transfer(format!( + "a received socket has address family {family}" + ))), + } +} + +fn sendmsg(fd: RawFd, data: &[u8], fds: &[RawFd]) -> io::Result { + let mut iov = libc::iovec { + iov_base: data.as_ptr() as *mut libc::c_void, + iov_len: data.len(), + }; + // SAFETY: msghdr 全零是合法值（没有 name、没有 control） + let mut msg: libc::msghdr = unsafe { mem::zeroed() }; + msg.msg_iov = &mut iov; + msg.msg_iovlen = 1; + let mut control: Vec; + if !fds.is_empty() { + let payload = mem::size_of_val(fds); + // SAFETY: 两个宏只做算术 + let (space, cmsg_len) = unsafe { + ( + libc::CMSG_SPACE(payload as libc::c_uint) as usize, + libc::CMSG_LEN(payload as libc::c_uint) as usize, + ) + }; + // 用 u64 数组，是为了让 cmsghdr 按它要的方式对齐 + control = vec![0u64; space.div_ceil(8)]; + msg.msg_control = control.as_mut_ptr().cast(); + msg.msg_controllen = space as _; + // SAFETY: control 至少 CMSG_SPACE(payload) 字节且已对齐，第一个头必然在里面； + // CMSG_DATA 之后正好放得下 payload 字节 + unsafe { + let c = libc::CMSG_FIRSTHDR(&msg); + (*c).cmsg_level = libc::SOL_SOCKET; + (*c).cmsg_type = libc::SCM_RIGHTS; + (*c).cmsg_len = cmsg_len as _; + std::ptr::copy_nonoverlapping(fds.as_ptr().cast::(), libc::CMSG_DATA(c), payload); + } + } + // SAFETY: msg 指向的 iov、control 都活到这次调用结束 + let n = unsafe { libc::sendmsg(fd, &msg, SEND_FLAGS) }; + if n < 0 { + return Err(io::Error::last_os_error()); + } + Ok(n as usize) +} + +/// 返回（读到的字节数，收到的句柄，辅助数据是否被截断）。 +fn recvmsg(fd: RawFd, dst: *mut u8, cap: usize) -> io::Result<(usize, Vec, bool)> { + let mut iov = libc::iovec { + iov_base: dst.cast(), + iov_len: cap, + }; + // SAFETY: 只做算术 + let space = + unsafe { libc::CMSG_SPACE((MAX_SOCKETS * mem::size_of::()) as libc::c_uint) } + as usize; + let mut control = vec![0u64; space.div_ceil(8)]; + // SAFETY: msghdr 全零是合法值 + let mut msg: libc::msghdr = unsafe { mem::zeroed() }; + msg.msg_iov = &mut iov; + msg.msg_iovlen = 1; + msg.msg_control = control.as_mut_ptr().cast(); + msg.msg_controllen = space as _; + // SAFETY: dst 往后 cap 个字节可写（调用方保证），control 活到调用结束 + let n = unsafe { libc::recvmsg(fd, &mut msg, RECV_FLAGS) }; + if n < 0 { + return Err(io::Error::last_os_error()); + } + // **先把收到的句柄都接住**，哪怕接下来要报错 —— 它们已经在本进程里了， + // 不接住就漏掉，接住了出错时随 OwnedFd 一起关 + let mut fds = Vec::new(); + // SAFETY: 遍历内核刚填好的 control；每个头的长度由内核给出 + unsafe { + let header = libc::CMSG_LEN(0) as usize; + let mut c = libc::CMSG_FIRSTHDR(&msg); + while !c.is_null() { + if (*c).cmsg_level == libc::SOL_SOCKET && (*c).cmsg_type == libc::SCM_RIGHTS { + let count = + ((*c).cmsg_len as usize).saturating_sub(header) / mem::size_of::(); + let data = libc::CMSG_DATA(c).cast::(); + for i in 0..count { + let raw = std::ptr::read_unaligned(data.add(i)); + fds.push(OwnedFd::from_raw_fd(raw)); + } + } + c = libc::CMSG_NXTHDR(&msg, c); + } + } + for f in &fds { + // SAFETY: f 是刚收到的有效句柄 + unsafe { libc::fcntl(f.as_raw_fd(), libc::F_SETFD, libc::FD_CLOEXEC) }; + } + let truncated = msg.msg_flags & libc::MSG_CTRUNC != 0; + Ok((n as usize, fds, truncated)) +} diff --git a/crates/tw-handover/src/windows.rs b/crates/tw-handover/src/windows.rs new file mode 100644 index 00000000..e5d259b0 --- /dev/null +++ b/crates/tw-handover/src/windows.rs @@ -0,0 +1,253 @@ +//! Windows：回环上的随机端口 + 口令；套接字用 `WSADuplicateSocketW` / +//! `WSASocketW` 复制到对端进程。 +//! +//! **Winsock 不用自己初始化**：两端在用到这些函数之前，都已经经由 std / tokio +//! 建过 TCP 套接字（端点本身、交接连接本身），`WSAStartup` 早就调过了。 + +use std::io; +use std::mem; +use std::net::{Ipv4Addr, SocketAddr}; +use std::os::windows::io::{AsRawSocket, FromRawSocket, RawSocket}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use windows_sys::Win32::Foundation::{HANDLE, HANDLE_FLAG_INHERIT, SetHandleInformation}; +use windows_sys::Win32::Networking::WinSock::{ + FROM_PROTOCOL_INFO, INVALID_SOCKET, SOCKET, WSA_FLAG_NO_HANDLE_INHERIT, WSA_FLAG_OVERLAPPED, + WSADuplicateSocketW, WSAGetLastError, WSAPROTOCOL_INFOW, WSASocketW, +}; +use windows_sys::Win32::Security::Cryptography::{ + BCRYPT_USE_SYSTEM_PREFERRED_RNG, BCryptGenRandom, +}; + +use crate::{ENDPOINT_FILE, Error, Gate, Result, Socket, TOKEN_FILE, token}; + +/// 一次最多读这么多。 +const CHUNK: usize = 64 * 1024; + +/// 连上来之后交口令的时限。 +const TOKEN_WAIT: Duration = Duration::from_secs(5); + +/// 口令的随机字节数（文件里是它的十六进制）。 +const TOKEN_BYTES: usize = 32; + +pub(crate) struct Listener { + listener: TcpListener, + dir: PathBuf, + port: u16, + /// 十六进制口令，即对端要交上来的那一帧 + token: String, +} + +impl Listener { + pub(crate) async fn bind(dir: &Path) -> Result { + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await?; + no_inherit(listener.as_raw_socket())?; + let port = listener.local_addr()?.port(); + let token = token::hex(&random(TOKEN_BYTES)?); + // **先口令后端口**：对端以端口文件为准判断 arm 了没有，看见端口时口令必须已经在 + write_atomic(&dir.join(TOKEN_FILE), &token)?; + write_atomic(&dir.join(ENDPOINT_FILE), &port.to_string())?; + Ok(Listener { + listener, + dir: dir.to_path_buf(), + port, + token, + }) + } + + /// 身份核对（口令）在 `Endpoint::accept` 里做，这里只管接。 + pub(crate) async fn accept(&self) -> Result)>> { + let (s, _) = self.listener.accept().await?; + if no_inherit(s.as_raw_socket()).is_err() { + return Ok(None); + } + // pid 要等对端的 Hello + Ok(Some((Stream { s }, None))) + } + + #[cfg(test)] + pub(crate) fn raw(&self) -> RawSocket { + self.listener.as_raw_socket() + } + + pub(crate) fn gate(&self) -> Gate { + Gate::Token { + token: Arc::from(self.token.as_bytes()), + wait: TOKEN_WAIT, + } + } +} + +impl Drop for Listener { + fn drop(&mut self) { + // 内容还是自己写的才删：排空中的旧 core 不能删掉新 core 后来开的端点（H8） + let port = self.dir.join(ENDPOINT_FILE); + if std::fs::read_to_string(&port).is_ok_and(|s| s.trim() == self.port.to_string()) { + let _ = std::fs::remove_file(&port); + } + let token = self.dir.join(TOKEN_FILE); + if std::fs::read_to_string(&token).is_ok_and(|s| s.trim() == self.token) { + let _ = std::fs::remove_file(&token); + } + } +} + +pub(crate) async fn connect(dir: &Path) -> Result)>> { + let Ok(text) = std::fs::read_to_string(dir.join(ENDPOINT_FILE)) else { + return Ok(None); + }; + let Ok(port) = text.trim().parse::() else { + return Ok(None); + }; + if !dir.join(TOKEN_FILE).exists() { + return Ok(None); + } + let Ok(s) = TcpStream::connect(SocketAddr::from((Ipv4Addr::LOCALHOST, port))).await else { + return Ok(None); + }; + no_inherit(s.as_raw_socket())?; + // 旧 core 的 pid 这一头问不到（它不发 Hello）；调用方要的话看锁文件 + Ok(Some((Stream { s }, None))) +} + +pub(crate) fn read_token(dir: &Path) -> Result { + Ok(std::fs::read_to_string(dir.join(TOKEN_FILE))? + .trim() + .to_string()) +} + +pub(crate) struct Stream { + s: TcpStream, +} + +impl Stream { + pub(crate) async fn write_frame(&mut self, frame: &[u8], _sockets: &[&Socket]) -> Result<()> { + // 套接字已经在 export 里变成帧里的 WSAPROTOCOL_INFOW 了 + self.s.write_all(frame).await?; + Ok(()) + } + + pub(crate) async fn read_some(&mut self, buf: &mut Vec) -> Result { + buf.reserve(CHUNK); + Ok(self.s.read_buf(buf).await?) + } + + pub(crate) fn import(&mut self, n: usize, w: &[String]) -> Result> { + if w.len() != n { + return Err(Error::Malformed(format!( + "the message carries {n} sockets but {} protocol infos", + w.len() + ))); + } + w.iter().map(|hex| socket_from_info(hex)).collect() + } +} + +/// 对 `peer_pid` 复制每个套接字，得到它的 `WSAPROTOCOL_INFOW`（十六进制）。 +/// 发送方自己的句柄不受影响。 +pub(crate) fn export(sockets: &[&Socket], peer_pid: u32) -> Result> { + sockets + .iter() + .map(|s| { + let Socket::Tcp(l) = s; + let raw = l.as_raw_socket() as SOCKET; + // SAFETY: WSAPROTOCOL_INFOW 是纯数据结构，全零合法；出参由系统填满 + let mut info: WSAPROTOCOL_INFOW = unsafe { mem::zeroed() }; + if unsafe { WSADuplicateSocketW(raw, peer_pid, &mut info) } != 0 { + return Err(Error::Transfer(format!( + "WSADuplicateSocketW for pid {peer_pid} failed: {}", + wsa_error() + ))); + } + // SAFETY: 把一个已初始化的纯数据结构按字节看 + let bytes = unsafe { + std::slice::from_raw_parts( + (&info as *const WSAPROTOCOL_INFOW).cast::(), + mem::size_of::(), + ) + }; + Ok(token::hex(bytes)) + }) + .collect() +} + +fn socket_from_info(hex: &str) -> Result { + let bytes = token::unhex(hex) + .filter(|b| b.len() == mem::size_of::()) + .ok_or_else(|| Error::Malformed("a socket's protocol info is not readable".into()))?; + // SAFETY: 长度刚核对过；WSAPROTOCOL_INFOW 是纯数据结构，任何字节都是合法值 + let info: WSAPROTOCOL_INFOW = + unsafe { std::ptr::read_unaligned(bytes.as_ptr().cast::()) }; + // SAFETY: info 活到调用结束 + let s = unsafe { + WSASocketW( + FROM_PROTOCOL_INFO, + FROM_PROTOCOL_INFO, + FROM_PROTOCOL_INFO, + &info, + 0, + WSA_FLAG_OVERLAPPED | WSA_FLAG_NO_HANDLE_INHERIT, + ) + }; + if s == INVALID_SOCKET { + return Err(Error::Transfer(format!( + "WSASocketW failed: {}", + wsa_error() + ))); + } + // SAFETY: s 是刚还原出来、只归我们所有的套接字 + let listener = unsafe { std::net::TcpListener::from_raw_socket(s as RawSocket) }; + // WSA_FLAG_NO_HANDLE_INHERIT 对 FROM_PROTOCOL_INFO 还原出来的套接字是否生效， + // 文档没有说；std 的 CreateProcess 又总是 bInheritHandles=TRUE —— 漏进子进程的 + // 监听会在我们退出后一直占着端口。所以再明确清一次继承标志 + no_inherit(s as RawSocket) + .map_err(|e| Error::Transfer(format!("SetHandleInformation failed: {e}")))?; + Ok(Socket::Tcp(listener)) +} + +/// 不让子进程继承这个套接字（std 起子进程时 `bInheritHandles` 总是 TRUE， +/// 能不能继承全看句柄自己的标志）。 +pub(crate) fn no_inherit(raw: RawSocket) -> io::Result<()> { + // SAFETY: raw 是本进程里有效的套接字句柄；只改它的继承标志 + if unsafe { SetHandleInformation(raw as HANDLE, HANDLE_FLAG_INHERIT, 0) } == 0 { + return Err(io::Error::last_os_error()); + } + Ok(()) +} + +fn wsa_error() -> io::Error { + // SAFETY: 没有参数，只读本线程的错误码 + io::Error::from_raw_os_error(unsafe { WSAGetLastError() }) +} + +/// 系统随机源（`BCryptGenRandom`，系统首选的 RNG）。 +fn random(n: usize) -> Result> { + let mut buf = vec![0u8; n]; + // SAFETY: buf 有 n 个字节可写；算法句柄为空配 SYSTEM_PREFERRED 是文档给的用法 + let status = unsafe { + BCryptGenRandom( + std::ptr::null_mut(), + buf.as_mut_ptr(), + n as u32, + BCRYPT_USE_SYSTEM_PREFERRED_RNG, + ) + }; + if status != 0 { + return Err(Error::Io(io::Error::other(format!( + "BCryptGenRandom failed: {status:#x}" + )))); + } + Ok(buf) +} + +/// 先写临时文件再改名，对端不会读到写了一半的内容。 +fn write_atomic(path: &Path, text: &str) -> io::Result<()> { + let mut tmp = path.as_os_str().to_owned(); + tmp.push(".tmp"); + std::fs::write(&tmp, text)?; + std::fs::rename(&tmp, path) +}