From 791b1118c2cdec66bc38769a0726543b57e0d81e Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 24 Oct 2026 12:22:02 +0800 Subject: [PATCH 1/9] Handover skeleton: tw-handover crate, gateway hooks, control API 52 Contract for the seamless-upgrade work: the handover link and socket transfer (tw-handover), the gateway's snapshot and settlement hooks, and POST /handover/arm with Status.handover / Status.draining. Bodies are todo!() and are filled in by the parallel lanes. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 13 ++ Cargo.toml | 3 + crates/tw-api/src/ep.rs | 3 + crates/tw-api/src/lib.rs | 60 ++++++++- crates/tw-control/src/lib.rs | 3 + crates/tw-gateway/src/handover.rs | 52 ++++++++ crates/tw-gateway/src/lib.rs | 1 + crates/tw-handover/Cargo.toml | 27 ++++ crates/tw-handover/src/lib.rs | 197 ++++++++++++++++++++++++++++++ 9 files changed, 358 insertions(+), 1 deletion(-) create mode 100644 crates/tw-gateway/src/handover.rs create mode 100644 crates/tw-handover/Cargo.toml create mode 100644 crates/tw-handover/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index bb59dcc3..cf889771 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3871,6 +3871,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" 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/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..61003308 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![], }; diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index c28cb534..49613b82 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -232,6 +232,9 @@ async fn status(State(s): State) -> Json { providers: cfg.providers.len(), uptime_secs: s.started.elapsed().as_secs(), in_flight: s.gateway.live.count(), + // C2 接上交接状态 + handover: tw_api::HandoverState::Idle, + draining: 0, remote_control, gateway_reachable, }) diff --git a/crates/tw-gateway/src/handover.rs b/crates/tw-gateway/src/handover.rs new file mode 100644 index 00000000..e349a783 --- /dev/null +++ b/crates/tw-gateway/src/handover.rs @@ -0,0 +1,52 @@ +//! 交接时网关这一边要交出去、接进来的东西：内存状态的快照，排空期每个请求的结算。 +//! +//! 搬运这些的是 `tw-handover`，它只当成 JSON；格式由这里定义，和交接连接同一个 +//! 版本号（`tw_handover::VERSION`），**不做跨版本兼容**。 +//! +//! # 快照 +//! +//! 接管开始时旧 core 交出（[`AppState::handover_export`]），新 core 接进来 +//! （[`AppState::handover_import`]），新的就不从零开始：对话粘滞、加密历史与缓存 +//! 断点的拒收记录、熔断与健康、限流暂停、延迟与速度、负载均衡轮次、OAuth 访问 +//! 令牌缓存、密钥用量与在途预留，以及旧 core 此刻占着的并发名额（按密钥、按上游）。 +//! +//! # 排空期的名额 +//! +//! 两个进程同时转发时，名额和用量上限不能翻倍：快照里旧 core 占着的名额，新 core +//! 记成已占用；旧 core 每结束一个请求发一条结算（[`AppState::begin_drain`] 的 +//! 接收端），新 core 据此放掉（[`AppState::handover_settled`]）。交接连接断了， +//! 剩下的一并放掉（[`AppState::handover_release_all`]），之后由请求库重建用量校正。 + +use crate::state::AppState; + +impl AppState { + /// 旧 core：此刻的内存状态。 + pub fn handover_export(&self) -> serde_json::Value { + todo!("C3") + } + + /// 新 core：接进旧 core 的快照。**认不出的部分丢掉**，不让交接失败 —— 版本号 + /// 相同时不该发生，发生了也只是少接一样状态。 + pub fn handover_import(&self, _snapshot: &serde_json::Value) { + todo!("C3") + } + + /// 旧 core：提交之后调一次。从这一刻起： + /// + /// - 每结束一个请求，往返回的接收端发一条结算（密钥、上游、用量、费用、放掉的名额）； + /// - WebSocket 会话在当前这一轮结束后以 1012（服务重启）关闭，客户端重连到新 core。 + pub fn begin_drain(&self) -> tokio::sync::mpsc::UnboundedReceiver { + todo!("C3") + } + + /// 新 core：旧 core 结束了一个请求。 + pub fn handover_settled(&self, _settled: &serde_json::Value) { + todo!("C3") + } + + /// 新 core：交接连接断了（旧 core 退出或崩溃）。快照里记着的旧 core 的名额与 + /// 预留全部放掉。 + pub fn handover_release_all(&self) { + todo!("C3") + } +} 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-handover/Cargo.toml b/crates/tw-handover/Cargo.toml new file mode 100644 index 00000000..e8bce902 --- /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"] } + +[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..8551d9e3 --- /dev/null +++ b/crates/tw-handover/src/lib.rs @@ -0,0 +1,197 @@ +//! 新旧 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 的配置读进来了、运行时建好了） +/// 旧 → 新 Listeners + 网关、远程控制端口的套接字 +/// 新 → 旧 Commit （新 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, + Settled { + data: serde_json::Value, + }, + InFlight { + n: u64, + }, + Bye, + Fallback, +} + +/// 旧 core 那一头：arm 之后开着，等接替者连上来。 +pub struct Endpoint { + _todo: (), +} + +impl Endpoint { + /// 在数据目录里开交接端点。已有的同名文件（上次没退干净）先删掉。 + pub async fn bind(_dir: &Path) -> Result { + todo!("C1") + } + + /// 等一个接替者。**核对过身份才交出 [`Link`]**（unix：同一用户；Windows：口令）； + /// 不对的连接直接断开，接着等下一个。 + pub async fn accept(&self) -> Result { + todo!("C1") + } + + /// 关掉端点，删掉文件。被丢掉时也做同一件事。 + pub fn close(self) {} +} + +/// 新 core 那一头：去连锁的持有者。**`Ok(None)` = 它没有 arm**（端点不存在或连不上）， +/// 调用方照旧报「已有实例在运行」。 +pub async fn connect(_dir: &Path) -> Result> { + todo!("C1") +} + +/// 一条交接连接。 +pub struct Link { + _todo: (), +} + +impl Link { + /// 对端的 pid（unix 上从内核问到的；Windows 上取 `Hello` 里报的）。 + pub fn peer_pid(&self) -> Option { + todo!("C1") + } + + /// 发一条不带套接字的消息。 + pub async fn send(&mut self, _m: &Message) -> Result<()> { + todo!("C1") + } + + /// 发一条消息，附带这几个套接字（顺序即 `infos` 的顺序）。`peer_pid` 只在 + /// Windows 上用得到（`WSADuplicateSocketW` 要它）。**发送方自己的句柄不受影响**， + /// 照样能 accept，直到它自己关掉。 + pub async fn send_with( + &mut self, + _m: &Message, + _sockets: &[&Socket], + _peer_pid: u32, + ) -> Result<()> { + todo!("C1") + } + + /// 收一条消息和它附带的套接字。对端关闭连接时返回 `Io(UnexpectedEof)`。 + pub async fn recv(&mut self) -> Result<(Message, Vec)> { + todo!("C1") + } +} From f5106f8fdf46e9db4c7bc589aecae3e5501710a7 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 24 Oct 2026 12:32:52 +0800 Subject: [PATCH 2/9] tw-handover: endpoint, identity checks, framing and listener passing - Endpoint: a unix socket (0600) in the data directory on macOS/Linux; on Windows a random loopback port plus a 32-byte token from the system RNG. Closing removes the files only if they are still ours, so a draining core cannot delete its successor's endpoint. - accept checks the peer: same uid via getpeereid/SO_PEERCRED (pid from LOCAL_PEERPID/SO_PEERCRED); on Windows a token frame within 5 seconds, compared in constant time. connect returns None when nothing is armed. - Frames are a 4-byte big-endian length plus JSON, at most 16 MiB. - Listeners travel with SCM_RIGHTS on unix (sendmsg/recvmsg under async_io, CLOEXEC, MSG_CTRUNC checked, TCP vs unix told apart by address family) and with WSADuplicateSocketW/WSASocketW on Windows. - Tests: queued connections on a handed-over TCP or unix listener all reach the receiver after the sender closes its handle; socket order across frames; oversized and garbled frames; endpoint lifecycle. Co-Authored-By: Claude Opus 5.5 --- crates/tw-handover/Cargo.toml | 2 +- crates/tw-handover/src/lib.rs | 185 +++++++++++-- crates/tw-handover/src/tests.rs | 430 ++++++++++++++++++++++++++++++ crates/tw-handover/src/token.rs | 39 +++ crates/tw-handover/src/unix.rs | 379 ++++++++++++++++++++++++++ crates/tw-handover/src/windows.rs | 227 ++++++++++++++++ 6 files changed, 1246 insertions(+), 16 deletions(-) create mode 100644 crates/tw-handover/src/tests.rs create mode 100644 crates/tw-handover/src/token.rs create mode 100644 crates/tw-handover/src/unix.rs create mode 100644 crates/tw-handover/src/windows.rs diff --git a/crates/tw-handover/Cargo.toml b/crates/tw-handover/Cargo.toml index e8bce902..aa05e28f 100644 --- a/crates/tw-handover/Cargo.toml +++ b/crates/tw-handover/Cargo.toml @@ -20,7 +20,7 @@ thiserror = { workspace = true } libc = { workspace = true } [target.'cfg(windows)'.dependencies] -windows-sys = { workspace = true, features = ["Win32_Networking_WinSock", "Win32_Foundation"] } +windows-sys = { workspace = true, features = ["Win32_Networking_WinSock", "Win32_Foundation", "Win32_Security_Cryptography"] } [dev-dependencies] tokio = { workspace = true, features = ["rt", "macros", "test-util"] } diff --git a/crates/tw-handover/src/lib.rs b/crates/tw-handover/src/lib.rs index 8551d9e3..084eb633 100644 --- a/crates/tw-handover/src/lib.rs +++ b/crates/tw-handover/src/lib.rs @@ -135,21 +135,75 @@ pub enum Message { Fallback, } +/// 单帧上限。快照可能较大（几千条亲和、限额计数），但不该没有边：长度前缀 +/// 读错了时宁可报错，也别照着一个 4 GiB 的数去分配。 +const MAX_FRAME: usize = 16 << 20; + +/// 一帧最多附几个套接字。网关每张网卡一个、再加远程控制端口，远远够了； +/// 也在 Linux 一条辅助消息的上限（`SCM_MAX_FD` = 253）之内。 +const MAX_SOCKETS: usize = 64; + +/// 线上的一帧：消息外面包一层，带上套接字的附件信息。两端都在本 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 { - _todo: (), + listener: sys::Listener, } impl Endpoint { /// 在数据目录里开交接端点。已有的同名文件（上次没退干净）先删掉。 - pub async fn bind(_dir: &Path) -> Result { - todo!("C1") + pub async fn bind(dir: &Path) -> Result { + Ok(Endpoint { + listener: sys::Listener::bind(dir).await?, + }) } /// 等一个接替者。**核对过身份才交出 [`Link`]**（unix：同一用户；Windows：口令）； /// 不对的连接直接断开，接着等下一个。 pub async fn accept(&self) -> Result { - todo!("C1") + loop { + let Some((io, pid)) = self.listener.accept().await? else { + continue; + }; + #[allow(unused_mut)] // Windows 上要先在它上面读口令 + let mut link = Link::new(io, pid); + #[cfg(windows)] + { + // 口令是连上之后的第一帧。**限时**：一个连上来不说话的进程 + // 不能把门一直堵着；口令帧也不许长 —— 没过门的不配让我们分配 16 MiB + let frame = link.read_frame(sys::TOKEN_FRAME_MAX); + let token = tokio::time::timeout(sys::TOKEN_WAIT, frame).await; + match token { + Ok(Ok(t)) if token::eq(&t, self.listener.token()) => {} + _ => continue, + } + } + return Ok(link); + } } /// 关掉端点，删掉文件。被丢掉时也做同一件事。 @@ -158,24 +212,45 @@ impl Endpoint { /// 新 core 那一头：去连锁的持有者。**`Ok(None)` = 它没有 arm**（端点不存在或连不上）， /// 调用方照旧报「已有实例在运行」。 -pub async fn connect(_dir: &Path) -> Result> { - todo!("C1") +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 { - _todo: (), + 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 { - todo!("C1") + self.peer_pid } /// 发一条不带套接字的消息。 - pub async fn send(&mut self, _m: &Message) -> Result<()> { - todo!("C1") + pub async fn send(&mut self, m: &Message) -> Result<()> { + self.send_with(m, &[], 0).await } /// 发一条消息，附带这几个套接字（顺序即 `infos` 的顺序）。`peer_pid` 只在 @@ -183,15 +258,95 @@ impl Link { /// 照样能 accept，直到它自己关掉。 pub async fn send_with( &mut self, - _m: &Message, - _sockets: &[&Socket], - _peer_pid: u32, + m: &Message, + sockets: &[&Socket], + peer_pid: u32, ) -> Result<()> { - todo!("C1") + 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)> { - todo!("C1") + 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..63ec6731 --- /dev/null +++ b/crates/tw-handover/src/tests.rs @@ -0,0 +1,430 @@ +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)); +} 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..7d18be6c --- /dev/null +++ b/crates/tw-handover/src/unix.rs @@ -0,0 +1,379 @@ +//! 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, 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`（两边都补，不差这一次系统调用）。 +#[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 }) + } + + /// `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..3e471f3f --- /dev/null +++ b/crates/tw-handover/src/windows.rs @@ -0,0 +1,227 @@ +//! 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::time::Duration; + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +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, Result, Socket, TOKEN_FILE, token}; + +/// 一次最多读这么多。 +const CHUNK: usize = 64 * 1024; + +/// 连上来之后交口令的时限。 +pub(crate) const TOKEN_WAIT: Duration = Duration::from_secs(5); + +/// 口令帧的长度上限（口令本身是 64 个十六进制字符）。 +pub(crate) const TOKEN_FRAME_MAX: usize = 256; + +/// 口令的随机字节数（文件里是它的十六进制）。 +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?; + 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?; + // pid 要等对端的 Hello + Ok(Some((Stream { s }, None))) + } + + pub(crate) fn token(&self) -> &[u8] { + self.token.as_bytes() + } +} + +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); + }; + // 旧 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 是刚还原出来、只归我们所有的套接字 + Ok(Socket::Tcp(unsafe { + std::net::TcpListener::from_raw_socket(s as RawSocket) + })) +} + +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) +} From 9a80056751b07980c43860554ac8d1c369ad2cc2 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 24 Oct 2026 12:55:48 +0800 Subject: [PATCH 3/9] Handover C3: gateway snapshot, phantom limits, settlement, WebSocket drain The old core exports its in-memory state as JSON (handover_export) and the new core imports it (handover_import): conversation affinity and session grouping, sealed-history and cache-breakpoint refusals, breaker and recent success rates, rate-limit pauses, latency samples, pace ratios, load-balance rounds, OAuth access tokens, key and upstream usage books with in-flight reservations, and the concurrency places the old core holds per key and per upstream. Instants travel as remaining/elapsed milliseconds; unreadable sections or entries are dropped without failing the handover. The new core holds the old core's places as phantom passes and slots (debt when it has no free place), so the two processes together stay within each limit. Key places count queued requests too, so a request the old core had queued is not overtaken by a new one in the new core. begin_drain reports the full set of holds, then every place freed or taken and every reservation taken, settled, released or dropped; usage changes are buffered from the snapshot on so none are lost before commit. handover_settled applies them, handover_release_all drops whatever is left. While draining, Responses WebSocket connections close with 1012 once the turn in hand has finished; idle ones close at once. Co-Authored-By: Claude Opus 5.5 --- crates/tw-gateway/src/affinity.rs | 59 +++ crates/tw-gateway/src/balance.rs | 24 + crates/tw-gateway/src/cache_marks.rs | 19 + crates/tw-gateway/src/handover.rs | 408 +++++++++++++++- crates/tw-gateway/src/handover/feed.rs | 185 ++++++++ crates/tw-gateway/src/handover/tests.rs | 552 ++++++++++++++++++++++ crates/tw-gateway/src/handover/wire.rs | 274 +++++++++++ crates/tw-gateway/src/health.rs | 55 +++ crates/tw-gateway/src/key_limits/books.rs | 86 +++- crates/tw-gateway/src/key_limits/mod.rs | 240 ++++++++-- crates/tw-gateway/src/latency.rs | 25 + crates/tw-gateway/src/limits.rs | 83 +++- crates/tw-gateway/src/oauth.rs | 91 +++- crates/tw-gateway/src/pace.rs | 25 + crates/tw-gateway/src/seal.rs | 30 ++ crates/tw-gateway/src/session.rs | 24 + crates/tw-gateway/src/slots.rs | 79 +++- crates/tw-gateway/src/state.rs | 4 + crates/tw-gateway/src/throttle.rs | 74 +++ crates/tw-gateway/src/upstream_limits.rs | 309 +++++++++--- crates/tw-gateway/src/ws.rs | 36 +- crates/tw-gateway/src/ws/turn.rs | 5 + crates/tw-gateway/tests/handover_ws.rs | 197 ++++++++ 23 files changed, 2739 insertions(+), 145 deletions(-) create mode 100644 crates/tw-gateway/src/handover/feed.rs create mode 100644 crates/tw-gateway/src/handover/tests.rs create mode 100644 crates/tw-gateway/src/handover/wire.rs create mode 100644 crates/tw-gateway/tests/handover_ws.rs 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/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 index e349a783..2e849fd5 100644 --- a/crates/tw-gateway/src/handover.rs +++ b/crates/tw-gateway/src/handover.rs @@ -1,52 +1,426 @@ //! 交接时网关这一边要交出去、接进来的东西：内存状态的快照，排空期每个请求的结算。 //! -//! 搬运这些的是 `tw-handover`，它只当成 JSON；格式由这里定义，和交接连接同一个 +//! 搬运这些的是 `tw-handover`，它只当成 JSON；格式由这里定义（[`wire`]），和交接连接同一个 //! 版本号（`tw_handover::VERSION`），**不做跨版本兼容**。 //! //! # 快照 //! //! 接管开始时旧 core 交出（[`AppState::handover_export`]），新 core 接进来 -//! （[`AppState::handover_import`]），新的就不从零开始：对话粘滞、加密历史与缓存 -//! 断点的拒收记录、熔断与健康、限流暂停、延迟与速度、负载均衡轮次、OAuth 访问 -//! 令牌缓存、密钥用量与在途预留，以及旧 core 此刻占着的并发名额（按密钥、按上游）。 +//! （[`AppState::handover_import`]），新的就不从零开始：对话粘滞、会话归属、加密历史与缓存 +//! 断点的拒收记录、熔断与健康、限流暂停、延迟与速度、负载均衡轮次、OAuth 访问令牌缓存、 +//! 密钥和上游的用量与在途预留，以及旧 core 此刻占着的并发名额（按密钥、按上游）。 +//! +//! **不搬的**：模型目录、余额、订阅额度、GLM 额度 —— 新 core 起来后台重新取，期间显示为 +//! 获取中，搬过去的也只是几分钟前的样子；对话这一轮沿用的路由决定 —— 它认的是旧 core 那一份 +//! 引擎，换了进程一定对不上，同一轮留在原来那一家靠的是上次回答它的那一家；限流的排队和 +//! 探路名额、在途请求的中止柄和实时内容 —— 都是旧 core 手上的请求，跟着它们在旧 core 上结束； +//! 「已经报过」的那几张表（写回成功、凭据失效、额度用完）—— 新 core 的界面也是新开的，状况 +//! 还在的话再报一次。 //! //! # 排空期的名额 //! -//! 两个进程同时转发时，名额和用量上限不能翻倍：快照里旧 core 占着的名额，新 core -//! 记成已占用；旧 core 每结束一个请求发一条结算（[`AppState::begin_drain`] 的 -//! 接收端），新 core 据此放掉（[`AppState::handover_settled`]）。交接连接断了， -//! 剩下的一并放掉（[`AppState::handover_release_all`]），之后由请求库重建用量校正。 +//! 两个进程同时转发时，名额和用量上限不能翻倍：快照里旧 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 { - todo!("C3") + 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 的快照。**认不出的部分丢掉**，不让交接失败 —— 版本号 /// 相同时不该发生，发生了也只是少接一样状态。 - pub fn handover_import(&self, _snapshot: &serde_json::Value) { - todo!("C3") + /// + /// 只调一次，在第一个请求之前（快照在监听交过来之前就到了）。再调一次的话，上一次替旧 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 { - todo!("C3") + 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 结束了一个请求。 - pub fn handover_settled(&self, _settled: &serde_json::Value) { - todo!("C3") + /// 新 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) { - todo!("C3") + // 在锁外面丢：放掉一个通行证要拿闸自己的锁 + 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..f11b9b17 --- /dev/null +++ b/crates/tw-gateway/src/handover/tests.rs @@ -0,0 +1,552 @@ +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(), + "排空开始时没报占着的名额" + ); +} 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/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/oauth.rs b/crates/tw-gateway/src/oauth.rs index 20ad938e..f0e6e504 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())] @@ -135,6 +135,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, @@ -328,6 +337,86 @@ 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` 是快照导出到现在过了多久 + 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) + || g.contains_key(&e.upstream) + { + 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() 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)); +} From 2f5f81ad1a76306e10233d7f42030369248675f8 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 24 Oct 2026 13:11:20 +0800 Subject: [PATCH 4/9] Handover C2: arm, takeover, commit and drain across twcore, control plane and gateway listeners The state machine for upgrading without dropping connections, in twcore: - POST /handover/arm: refused over the remote port, in safe mode, while a ChatGPT or Z.ai sign-in is pending, or while a successor is taking over. Once armed the core stops watching --parent, opens the handover endpoint, moves stdout/stderr off the desktop app's pipe into twcore.handover.log, and waits (120 s by default); it relaunches the app once, then drains and exits the current way if nobody comes (H7). - Old side: Hello -> Control (the control plane stops accepting) -> Snapshot (after pending token write-backs and sync uploads) -> WantListeners (no time limit, H3) -> Listeners -> Commit within 30 s. Any failure before commit goes back to armed (F4); three failures after the listeners were handed over make the next successor fall back (F5). After commit it stops accepting (own descriptors only), stops background tasks, disowns the lock, sends Settled / InFlight while draining (30 min limit, then aborts the rest), and says Bye. - New side: when the lock holder is alive it connects to the handover endpoint; Refused or an unreadable reply makes it ask for Fallback and wait for the old core to exit (H2). It serves the control plane on the inherited socket, takes only the control plane while its configuration does not load (H3), accepts on the inherited listeners, commits within 20 s, claims the lock and follows the old core's drain. - The request history opens only after the previous core exits; records queue in memory meanwhile (10,000 records or 256 MB, oldest dropped). Request ids jump past the previous core's by a large gap so the two never overwrite each other's rows. - Background tasks run on their own runtime that is shut down at commit. Model probing and balance reads stay on the main runtime: they share the upstream connection pool with forwarding, and closing their runtime cut a stream in flight. tw-gateway listen.rs gains serve_with and Handoff (export, hold, follow, stop accepting); the remote control port and the local control plane get the same handles. Status reports handover and draining. End-to-end tests start two twcore processes; they are ignored until the handover link and gateway hooks are merged. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 3 + bin/twcore/Cargo.toml | 6 + bin/twcore/src/background.rs | 162 +++ bin/twcore/src/handover.rs | 1017 +++++++++++++++++++ bin/twcore/src/lockfile.rs | 57 +- bin/twcore/src/main.rs | 585 ++++++----- bin/twcore/src/record.rs | 440 ++++++++ bin/twcore/tests/handover.rs | 523 ++++++++++ crates/tw-api/msg-codes.txt | 4 + crates/tw-control/Cargo.toml | 2 + crates/tw-control/src/chatgpt.rs | 8 + crates/tw-control/src/handover.rs | 341 +++++++ crates/tw-control/src/lib.rs | 161 ++- crates/tw-control/src/remote.rs | 160 ++- crates/tw-control/src/rotation.rs | 64 +- crates/tw-control/src/zai.rs | 8 + crates/tw-control/tests/aliases.rs | 1 + crates/tw-control/tests/backups.rs | 1 + crates/tw-control/tests/bedrock.rs | 1 + crates/tw-control/tests/bundle.rs | 1 + crates/tw-control/tests/chatgpt.rs | 1 + crates/tw-control/tests/client_key.rs | 1 + crates/tw-control/tests/contract.rs | 1 + crates/tw-control/tests/dryrun.rs | 1 + crates/tw-control/tests/error_answers.rs | 1 + crates/tw-control/tests/gate.rs | 1 + crates/tw-control/tests/handover.rs | 150 +++ crates/tw-control/tests/history_search.rs | 2 + crates/tw-control/tests/in_flight.rs | 1 + crates/tw-control/tests/key_limits.rs | 1 + crates/tw-control/tests/keys.rs | 1 + crates/tw-control/tests/listen.rs | 1 + crates/tw-control/tests/live_content.rs | 1 + crates/tw-control/tests/live_state.rs | 1 + crates/tw-control/tests/manual_models.rs | 1 + crates/tw-control/tests/model_specs.rs | 1 + crates/tw-control/tests/models.rs | 1 + crates/tw-control/tests/plugin_defaults.rs | 1 + crates/tw-control/tests/plugins.rs | 1 + crates/tw-control/tests/pricing.rs | 1 + crates/tw-control/tests/remote.rs | 1 + crates/tw-control/tests/replay.rs | 1 + crates/tw-control/tests/resources.rs | 1 + crates/tw-control/tests/route_stats.rs | 3 + crates/tw-control/tests/routes.rs | 1 + crates/tw-control/tests/safe_mode.rs | 1 + crates/tw-control/tests/security.rs | 1 + crates/tw-control/tests/security_detail.rs | 1 + crates/tw-control/tests/sessions.rs | 1 + crates/tw-control/tests/speed.rs | 1 + crates/tw-control/tests/stored_bodies.rs | 1 + crates/tw-control/tests/summary.rs | 1 + crates/tw-control/tests/sync_devices.rs | 2 + crates/tw-control/tests/sync_regressions.rs | 1 + crates/tw-control/tests/transcript.rs | 1 + crates/tw-control/tests/turn_context.rs | 1 + crates/tw-control/tests/upstream_health.rs | 1 + crates/tw-control/tests/upstream_limits.rs | 1 + crates/tw-control/tests/zai.rs | 1 + crates/tw-gateway/src/listen.rs | 259 ++++- crates/tw-gateway/tests/listen_handoff.rs | 191 ++++ 61 files changed, 3865 insertions(+), 322 deletions(-) create mode 100644 bin/twcore/src/background.rs create mode 100644 bin/twcore/src/handover.rs create mode 100644 bin/twcore/src/record.rs create mode 100644 bin/twcore/tests/handover.rs create mode 100644 crates/tw-control/src/handover.rs create mode 100644 crates/tw-control/tests/handover.rs create mode 100644 crates/tw-gateway/tests/listen_handoff.rs diff --git a/Cargo.lock b/Cargo.lock index cf889771..8dcf7b4f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3767,6 +3767,7 @@ dependencies = [ "tw-engine", "tw-gateway", "tw-guard", + "tw-handover", "tw-link", "tw-observe", "tw-pack", @@ -4053,6 +4054,7 @@ dependencies = [ "axum", "chrono", "clap", + "futures", "http-body-util", "hyper", "hyper-util", @@ -4071,6 +4073,7 @@ dependencies = [ "tw-control", "tw-gateway", "tw-guard", + "tw-handover", "tw-link", "tw-observe", "tw-pack", 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..290728b8 --- /dev/null +++ b/bin/twcore/src/handover.rs @@ -0,0 +1,1017 @@ +//! 升级不断线：新旧 core 交接的状态机（方案第 4、6、7 节）。 +//! +//! 两个角色，同一个二进制： +//! +//! - **旧 core**（[`serve_arms`]）：桌面端 `POST /handover/arm` 之后脱离父进程、开交接端点、 +//! 等接替者；接替者连上来就把控制面、快照、监听交过去，提交之后停止 accept、停掉后台 +//! 事务、放弃锁，做完手上的请求再退出。 +//! - **新 core**（[`approach`] → [`Incoming`]）：取锁时发现持有者活着，就去连它的交接端点； +//! 连上了就在交过来的控制面上服务、读配置建运行时、要监听、提交，成为锁的持有者。 +//! +//! **提交是唯一的不可逆点。**提交之前任何一步出错，旧 core 回到待交接、照常服务（它始终 +//! 握着全部监听）；提交之后新 core 是唯一的「主」，旧 core 只做完手上的请求。任何情况下 +//! 都不出现两个进程都不监听的时刻。 +//! +//! # 时限 +//! +//! | 时限 | 多久 | 谁在等 | +//! | --- | --- | --- | +//! | 待交接 | arm 给的 `wait_secs`，缺省 120 秒 | 旧 core 等接替者连上来（H7） | +//! | 接管 | 30 秒（[`TAKEOVER_LIMIT`]） | 旧 core 交出监听之后等提交（F4） | +//! | 排空 | 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); +/// 连着几次交接在交出监听之后失败（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>>, + /// 请求库打开了没有。接替者要等它的前任退出才打开（见 `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 { + // 第一次 arm：开端点、脱离父进程 + let (mut endpoint, mut 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); + 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(())); + break ( + Some(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))); + } + } + }; + 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 Exit::HandedOver; + } + Given::Fallback => { + tracing::info!("the successor cannot take over; draining and exiting the current way"); + retire(core, FALLBACK_DRAIN).await; + return Exit::Retired; + } + Given::Abandoned { strike } => { + if strike { + strikes += 1; + } + // 回到待交接：控制面接着接，重新等一个周期（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) => 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 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(())); + } + } + } +} + +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)) + } + _ => return Given::NotASuccessor, + }; + core.handover().set_phase(HandoverState::TakingOver); + if version != tw_handover::VERSION || refuse { + // H2 / F5：退回现行流程。等它说一声 Fallback（它会马上说） + 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 + ) + }; + tracing::info!(%reason, "refusing the handover"); + let _ = link.send(&Message::Refused { reason }).await; + return match tokio::time::timeout(HELLO_WITHIN, link.recv()).await { + Ok(Ok((Message::Fallback, _))) => Given::Fallback, + _ => Given::Abandoned { strike: false }, + }; + } + + // ① 控制面。**交出去就不再从它上面接新连接**（见 `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(); + 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); + + // ② 快照。之前先把手上的收尾：正在写回的 refresh token、还没上传的本机改动 + 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: false }; + } + + // ③ 等它要监听。**不限时**（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) => { + tracing::info!("the successor left before taking the listeners: {e}"); + return Given::Abandoned { strike: false }; + } + } + } + + // ④ 监听：网关的每一个、远程控制端口 + 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） + match tokio::time::timeout(TAKEOVER_LIMIT, link.recv()).await { + Ok(Ok((Message::Commit, _))) => Given::Committed(link), + 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 } + } + } +} + +/// 提交之后：停止 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 之后的日志写到哪：数据目录里，每次 arm 从头写。 +pub const HANDOVER_LOG: &str = "twcore.handover.log"; + +/// 标准输出、错误输出改写进 [`HANDOVER_LOG`]（开不了就丢掉），标准输入接空设备。**再也 +/// 不碰桌面端的管道**：它退出之后写管道会出错。 +fn stdio_to_file(dir: &Path) { + let path = dir.join(HANDOVER_LOG); + let file = { + let mut o = std::fs::OpenOptions::new(); + o.write(true).create(true).truncate(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 交过来的网关与远程控制端口的监听。 +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, + }; + let snapshot = match tokio::time::timeout(HELLO_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, + }) +} + +/// 请旧 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 + } +} + +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; + }; + state.bus.resume_after(s.next_id.saturating_add(ID_GAP)); + state.handover_import(&s.gateway); + } + + /// 要监听（配置读进来了、运行时建好了）。 + pub async fn listeners(&mut self) -> Result { + self.link + .send(&Message::WantListeners) + .await + .context("the previous core left before handing over the listeners")?; + 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 这时候没了的话（F6）， + /// 套接字已经在手上，直接接着服务 + pub async fn commit(&mut self) { + if let Err(e) = self.link.send(&Message::Commit).await { + tracing::warn!("the previous core is gone at commit; serving on without it: {e}"); + } + } + + /// 提交之后跟着旧 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 的话等它的进程真没了再放：它说完 Bye 到退出之间手上已经没有请求了， + // 但名额要等它真的不可能再用才算放开。连接断了（崩了）也一样等一下 + let _ = wait_gone(self.predecessor, Duration::from_secs(60)).await; + 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_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..bfd7c6e6 100644 --- a/bin/twcore/src/lockfile.rs +++ b/bin/twcore/src/lockfile.rs @@ -10,6 +10,8 @@ use std::path::{Path, PathBuf}; #[derive(Debug)] pub struct LockFile { path: PathBuf, + /// 退出时删不删文件。**交接之后不删**：那时文件里写的已经是接替者的 pid + owned: bool, } #[derive(Debug)] @@ -35,14 +37,43 @@ impl LockFile { } // 到这儿说明：没有锁文件、内容不是数字、或者那个 pid 已经死了。 // 后两种都是「上次没退干净」，直接覆盖。 + 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("twcore.lock")) + } + + /// 放弃锁：退出时不删文件（交接提交之后，文件里马上会是接替者的 pid）。 + pub fn disown(mut self) { + self.owned = false; + } + + fn write(path: PathBuf) -> std::io::Result { let mut f = std::fs::File::create(&path)?; write!(f, "{}", std::process::id())?; - Ok(LockOutcome::Acquired(LockFile { path })) + Ok(LockFile { path, owned: true }) } } +/// 锁文件里此刻写的 pid（读不到、不是数字是 `None`）。 +#[cfg(test)] +pub fn holder(dir: &Path) -> Option { + std::fs::read_to_string(dir.join("twcore.lock")) + .ok() + .and_then(|t| t.trim().parse().ok()) +} + impl Drop for LockFile { fn drop(&mut self) { + if !self.owned { + return; + } // 尽力而为。删不掉也无所谓 —— 下次启动的僵尸检测会认出它。 let _ = std::fs::remove_file(&self.path); } @@ -133,6 +164,30 @@ mod tests { } } + #[test] + fn a_disowned_lock_leaves_the_file_for_the_successor() { + // 交接提交之后旧 core 放弃锁：它退出时文件里写的已经是接替者，不能删 + let d = tempfile::tempdir().unwrap(); + let p = d.path().join("twcore.lock"); + 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())); + } + #[test] fn dropping_the_lock_removes_the_file() { let d = tempfile::tempdir().unwrap(); diff --git a/bin/twcore/src/main.rs b/bin/twcore/src/main.rs index dc03ad21..cd6eeab3 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,26 @@ 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(); // **配置不存在就地生成**（第 1 步）。 // @@ -821,17 +845,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 +880,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 +889,19 @@ 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() { + inc.import(&state); + } // 观测这一层。**起不来不是致命的** —— 历史记录看不见，而网关 // 照常转发。所以这里所有的失败都只记一行日志。 @@ -880,66 +911,49 @@ 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 { + 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) = match &incoming { + None => { + 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) + } + // 接替者：旧 core 退出之后才打开请求库（R3，见 `record`）。用量从记录重建在打开的 + // 那一刻、这个进程排着的记录进库之前 + Some(inc) => { + let gw = state.clone(); + let (store, opened) = record::open_later( + &dir, + wiring, + inc.predecessor, + Box::new(move |db| tw_control::key_limits::rebuild(&gw, db)), + ); + 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); + } + }); + } + (store, opened) + } + }; + if store.is_some() { 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 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" - ); - } - tokio::time::sleep(std::time::Duration::from_secs(3600)).await; - } - }); } // 配置的唯一入口。UI、CLI、文件监听都从这里进。 @@ -954,54 +968,75 @@ 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); + 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,24 +1045,29 @@ 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); - }); + let (control_died, mut control_dead) = tokio::sync::oneshot::channel(); + { + let control = control.clone(); + 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 还在不在，哪个平台、 - // 谁来收养都一样。 + // 父进程守望：GUI 没了我们跟着退。轮询而不是用 kqueue，是因为 + // 这段代码要能在被重新 parent 之后仍然正确 —— macOS 上收养我们的 + // 是 launchd（getppid() 变成 1），Linux 桌面会话里多半是 + // `systemd --user`（它是 subreaper，getppid() 变成它的 pid，不是 1）。 + // 所以**不看 getppid()**，只看原来那个 pid 还在不在，哪个平台、 + // 谁来收养都一样。 + // + // **arm 之后关掉**（见 `handover`）：桌面端装好更新要重启，core 留下来等接替者 + let parent_watch = parent.map(|ppid| { tokio::spawn(async move { loop { tokio::time::sleep(std::time::Duration::from_secs(2)).await; @@ -1036,177 +1076,200 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> std::process::exit(0); } } - }); - } + }) + }); 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")) - } - msg = control_dead => { - anyhow::bail!("{}", msg.unwrap_or_else(|_| "the control plane stopped".into())) + 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(), + )) } - _ = 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() => 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")?; + inc.commit().await; + // 成为锁的持有者：写上自己的 pid（旧 core 已经放弃了锁） + match LockFile::claim(&dir) { + Ok(l) => lock = Some(l), + Err(e) => tracing::warn!("the lock file could not be taken over: {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 = bg + .lock() + .map_err(|_| anyhow::anyhow!("the background lock is poisoned"))? + .start(background::Start { + control: &control, + safe: false, + seeder: Some(s), + store: store.clone(), + store_opened: store_opened.clone(), + }) + .context("starting the background tasks")?; + sync = Some(started); + 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), + 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; + // 后台事务的运行时在异步上下文之外关（在里面丢掉一个运行时会 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..cfd52de6 --- /dev/null +++ b/bin/twcore/src/record.rs @@ -0,0 +1,440 @@ +//! 请求记录：观测层怎么接上网关。 +//! +//! **每一步失败都只是「没有历史记录」，不是「起不来」。**这条边界写死在代码形状里： +//! 这里交出去的是 `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（`predecessor`）的进程没了才打开请求库，之间的记录排队。 +/// +/// 交出去的是占位的那一份（内存库）；打开之后**原地换成真的**，控制面拿着的还是同一个 +/// `Arc`。`on_open` 在真的库打开、排队的记录写进去之前调一次（密钥用量从记录重建，和 +/// 平常启动时同一个时机：在这个进程自己的记录进库之前）。返回的 `watch` 在打开之后 +/// （或者打开失败、不再等之后）变成 `true`。 +pub fn open_later( + dir: &Path, + w: Wiring, + predecessor: u32, + on_open: Box, +) -> (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(), + ), + 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(predecessor, &mut events, &mut queue); + 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}" + ); + opened_tx.send_replace(true); + // 通道要有人收，否则总线那一头一直记「丢了」 + while events.blocking_recv().is_some() {} + return; + } + }; + resume_ids(&w.bus, &db); + on_open(&db); + let blobs_root = blobs.root().to_path_buf(); + *shared.blocking_lock() = tw_store::Recorder::new(db, blobs, w.pricing) + .reporting_to(w.bus) + .settling_to(w.settled); + spawn_body_writer(blobs_root, body_rx); + tw_store::task::record_plugin_runs(shared.clone(), run_rx); + 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 g = shared.blocking_lock(); + for ev in queue.drain() { + g.on_event(&ev); + } + } + 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) +} + +/// 旧 core 最多排空这么久（`handover::DRAIN_LIMIT` 加一点余量）。到了还没退，照样打开 +const PREDECESSOR_WAIT: Duration = Duration::from_secs(35 * 60); + +/// 等旧 core 的进程没了，之间把记录收进队列。 +fn wait_for_exit(pid: u32, events: &mut tw_observe::RecordFeed, queue: &mut Queue) { + let started = Instant::now(); + let mut last_check = Instant::now() - Duration::from_secs(1); + loop { + while let Some(ev) = events.try_recv() { + queue.push(ev); + } + queue.dropped += events.take_dropped(); + if last_check.elapsed() >= Duration::from_millis(500) { + last_check = Instant::now(); + if !crate::proc::alive(pid) { + return; + } + if started.elapsed() >= PREDECESSOR_WAIT { + tracing::warn!( + pid, + "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..c80bfa2b --- /dev/null +++ b/bin/twcore/tests/handover.rs @@ -0,0 +1,523 @@ +//! 升级不断线，端到端：起真的 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 { + 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\nproviders:\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 { + let log = std::fs::File::create(self.dir.path().join(format!("{name}.log"))).unwrap(); + self.command().stderr(log).spawn().unwrap() + } + + /// 和桌面端一样，stderr 接一根管道（桌面端退出时这根管道的读端就没了） + fn core_on_a_pipe(&self) -> Child { + self.command().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") + .stdout(Stdio::null()); + 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", "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)] +#[ignore = "等 C1/C3 集成"] +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(20), || 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].stderr.take()); + reap.0.push(bed.core("b")); + let b_pid = reap.0[1].id(); + assert!( + until(Duration::from_secs(20), || 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(20), || 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)] +#[ignore = "等 C1/C3 集成"] +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(20), || 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(); + 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)] +#[ignore = "等 C1/C3 集成"] +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(20), || 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)] +#[ignore = "等 C1/C3 集成"] +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(20), || 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); +} diff --git a/crates/tw-api/msg-codes.txt b/crates/tw-api/msg-codes.txt index c86f9a09..5e8493cd 100644 --- a/crates/tw-api/msg-codes.txt +++ b/crates/tw-api/msg-codes.txt @@ -178,6 +178,10 @@ 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.remote +control.handover.safe_mode control.header_no_value control.internal_error control.key_create_failed 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..cd108bfe 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() diff --git a/crates/tw-control/src/handover.rs b/crates/tw-control/src/handover.rs new file mode 100644 index 00000000..29974f04 --- /dev/null +++ b/crates/tw-control/src/handover.rs @@ -0,0 +1,341 @@ +//! 升级不断线：控制面这一边。 +//! +//! # 分工 +//! +//! 交接的状态机在 `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 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, + held: Mutex>, + inherited: Mutex>, +} + +impl Default for ControlListener { + fn default() -> Self { + Self { + accepting: watch::channel(true).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() + } +} + +/// 再复制一个句柄。 +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 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 49613b82..dd9adf02 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,9 +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(), - // C2 接上交接状态 - handover: tw_api::HandoverState::Idle, - draining: 0, + handover: s.handover.phase(), + draining: s.handover.draining(), remote_control, gateway_reachable, }) @@ -2028,13 +2031,103 @@ 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(), None), + Err(e) => tracing::warn!("the control plane could not accept a connection: {e}"), + }, + // 停下的那一刻放下手里这次 accept：不然它还会再接一个 + _ = accepting.wait_for(|on| !*on) => {} + } } } @@ -2043,7 +2136,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); @@ -2051,23 +2149,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 只有这一档。 @@ -2083,36 +2184,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。 diff --git a/crates/tw-control/src/remote.rs b/crates/tw-control/src/remote.rs index a45b7bde..688ebfba 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, + }) } /// 等到这个来源不再被放行：名单改了把它划出去，或者远程端口关了。 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..d1a7cd05 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() 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..e400bf14 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, diff --git a/crates/tw-control/tests/handover.rs b/crates/tw-control/tests/handover.rs new file mode 100644 index 00000000..eff4e626 --- /dev/null +++ b/crates/tw-control/tests/handover.rs @@ -0,0 +1,150 @@ +//! `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); +} 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/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/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, "退场的监听没放开端口"); +} From 83e76a8d708a0731af249422bb36ccd48ad74b7d Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 24 Oct 2026 13:32:39 +0800 Subject: [PATCH 5/9] Integrate handover lanes: enable end-to-end tests, fix Status JSON test and test wording Co-Authored-By: Claude Opus 5.5 --- bin/twcore/tests/handover.rs | 8 ++++---- crates/tw-api/src/lib.rs | 2 +- crates/tw-gateway/src/handover/tests.rs | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/bin/twcore/tests/handover.rs b/bin/twcore/tests/handover.rs index c80bfa2b..ae43edd6 100644 --- a/bin/twcore/tests/handover.rs +++ b/bin/twcore/tests/handover.rs @@ -265,7 +265,7 @@ async fn quick(c: &reqwest::Client, gw: SocketAddr) -> Result<(), String> { } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] -#[ignore = "等 C1/C3 集成"] + 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()]); @@ -415,7 +415,7 @@ async fn fake_successor(dir: &Path, version: u32) -> tw_handover::Link { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[ignore = "等 C1/C3 集成"] + 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")]); @@ -464,7 +464,7 @@ async fn a_successor_that_leaves_before_committing_leaves_the_old_core_serving() } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[ignore = "等 C1/C3 集成"] + 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")]); @@ -497,7 +497,7 @@ async fn a_different_handover_version_falls_back_to_drain_and_exit() { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[ignore = "等 C1/C3 集成"] + 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")]); diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index 61003308..a148f41f 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -7690,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-gateway/src/handover/tests.rs b/crates/tw-gateway/src/handover/tests.rs index f11b9b17..edf66652 100644 --- a/crates/tw-gateway/src/handover/tests.rs +++ b/crates/tw-gateway/src/handover/tests.rs @@ -270,7 +270,7 @@ async fn a_snapshot_round_trip_leaves_the_new_core_where_the_old_one_was() { 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("乙"), "熔断着的那一家又被放行了"); + 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); From d19f5c8a23c69689c505d3473d3ac49b2f9094ec Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 24 Oct 2026 13:43:02 +0800 Subject: [PATCH 6/9] Handover: hold OAuth renewals while two cores overlap; stop probing and balance reads once draining Renewing a refresh token invalidates the previous one, so two cores renewing independently leaves one of them signed out. The old core stops renewing when it hands its tokens over (resumes if the takeover fails); the successor starts with renewals held and takes them over at commit. Co-Authored-By: Claude Opus 5.5 --- bin/twcore/src/handover.rs | 6 +++++ bin/twcore/src/main.rs | 4 ++++ crates/tw-api/msg-codes.txt | 1 + crates/tw-gateway/src/balances.rs | 7 ++++++ crates/tw-gateway/src/models.rs | 7 ++++++ crates/tw-gateway/src/oauth.rs | 35 ++++++++++++++++++++++++++++- crates/tw-gateway/tests/oauth.rs | 37 +++++++++++++++++++++++++++++++ 7 files changed, 96 insertions(+), 1 deletion(-) diff --git a/bin/twcore/src/handover.rs b/bin/twcore/src/handover.rs index 290728b8..d787ee67 100644 --- a/bin/twcore/src/handover.rs +++ b/bin/twcore/src/handover.rs @@ -289,6 +289,8 @@ async fn run_arms( return 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 Exit::Retired; @@ -297,6 +299,8 @@ async fn run_arms( if strike { strikes += 1; } + // 接管没成：续期回到这一边 + core.state.oauth.hold_renewals(false); // 回到待交接：控制面接着接，重新等一个周期（F4） core.handover().control().resume(); core.handover().set_phase(HandoverState::Armed); @@ -407,6 +411,8 @@ async fn give(core: &Core, mut link: Link, refuse: bool) -> Given { if let Some(sync) = core.sync() { sync.flush().await; } + // 交出 access token 起暂停续期：两边各换一次 refresh token 必有一边作废（接管没成再恢复） + core.state.oauth.hold_renewals(true); 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}"); diff --git a/bin/twcore/src/main.rs b/bin/twcore/src/main.rs index cd6eeab3..996d42cd 100644 --- a/bin/twcore/src/main.rs +++ b/bin/twcore/src/main.rs @@ -900,6 +900,8 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> 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); } @@ -1156,6 +1158,8 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> .await .context("the listeners handed over could not be served in time")?; inc.commit().await; + // 从此续期由这一边做（旧 core 交出快照起已暂停） + state.oauth.hold_renewals(false); // 成为锁的持有者：写上自己的 pid（旧 core 已经放弃了锁） match LockFile::claim(&dir) { Ok(l) => lock = Some(l), diff --git a/crates/tw-api/msg-codes.txt b/crates/tw-api/msg-codes.txt index 5e8493cd..45737acd 100644 --- a/crates/tw-api/msg-codes.txt +++ b/crates/tw-api/msg-codes.txt @@ -365,6 +365,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-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/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 f0e6e504..f1e874f2 100644 --- a/crates/tw-gateway/src/oauth.rs +++ b/crates/tw-gateway/src/oauth.rs @@ -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() + } }, } } @@ -245,6 +255,8 @@ pub struct Cache { inner: Mutex>, /// 每个 provider 一把刷新锁（见模块文档第四条） flights: Mutex>>>, + /// 交接期间暂停续期（见 [`Self::hold_renewals`]） + held: std::sync::atomic::AtomicBool, } /// 换回来的东西。 @@ -469,6 +481,17 @@ 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); + } + /// 真的去换。**调用方持有这家的刷新锁**，所以这里不会有并发的第二次。 async fn refresh( &self, @@ -476,6 +499,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/tests/oauth.rs b/crates/tw-gateway/tests/oauth.rs index 59601571..19081904 100644 --- a/crates/tw-gateway/tests/oauth.rs +++ b/crates/tw-gateway/tests/oauth.rs @@ -516,3 +516,40 @@ 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); +} From 2a6d1eda6caa350e54dee52e003e6d236c4ade14 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 24 Oct 2026 14:08:03 +0800 Subject: [PATCH 7/9] tw-handover: check handover peers concurrently; keep Windows sockets out of child processes - Endpoint::accept now runs each connection's gate check in its own task (at most 64 at once, the oldest dropped when full) and hands over the first one that passes. On Windows the gate is a token within 5 seconds and anyone on the machine can reach the loopback port, so checking one connection at a time let silent connections keep the real successor out. Pending checks are aborted when accept returns or is cancelled. - Windows: clear HANDLE_FLAG_INHERIT on sockets restored from WSAPROTOCOL_INFOW, on the endpoint's listener and on both ends of the handover connection; std's CreateProcess always inherits handles. - Document the macOS window between recvmsg and FD_CLOEXEC. - Tests: silent and wrong-token connections (more than the limit) do not delay or impersonate the right one; the same through a real endpoint; on Windows, handover sockets are not inheritable. Co-Authored-By: Claude Opus 5.5 --- crates/tw-handover/src/lib.rs | 98 ++++++++++++++++++---- crates/tw-handover/src/tests.rs | 134 ++++++++++++++++++++++++++++++ crates/tw-handover/src/unix.rs | 13 ++- crates/tw-handover/src/windows.rs | 46 +++++++--- 4 files changed, 265 insertions(+), 26 deletions(-) diff --git a/crates/tw-handover/src/lib.rs b/crates/tw-handover/src/lib.rs index 084eb633..ca14078a 100644 --- a/crates/tw-handover/src/lib.rs +++ b/crates/tw-handover/src/lib.rs @@ -143,6 +143,46 @@ 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)] @@ -184,25 +224,53 @@ impl Endpoint { /// 等一个接替者。**核对过身份才交出 [`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 { - let Some((io, pid)) = self.listener.accept().await? else { - continue; - }; - #[allow(unused_mut)] // Windows 上要先在它上面读口令 - let mut link = Link::new(io, pid); - #[cfg(windows)] - { - // 口令是连上之后的第一帧。**限时**：一个连上来不说话的进程 - // 不能把门一直堵着；口令帧也不许长 —— 没过门的不配让我们分配 16 MiB - let frame = link.read_frame(sys::TOKEN_FRAME_MAX); - let token = tokio::time::timeout(sys::TOKEN_WAIT, frame).await; - match token { - Ok(Ok(t)) if token::eq(&t, self.listener.token()) => {} - _ => continue, + 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); + } } } - return Ok(link); } } diff --git a/crates/tw-handover/src/tests.rs b/crates/tw-handover/src/tests.rs index 63ec6731..94164008 100644 --- a/crates/tw-handover/src/tests.rs +++ b/crates/tw-handover/src/tests.rs @@ -428,3 +428,137 @@ 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/unix.rs b/crates/tw-handover/src/unix.rs index 7d18be6c..3b4bf7a4 100644 --- a/crates/tw-handover/src/unix.rs +++ b/crates/tw-handover/src/unix.rs @@ -13,7 +13,7 @@ use std::path::{Path, PathBuf}; use tokio::io::Interest; use tokio::net::{UnixListener, UnixStream}; -use crate::{ENDPOINT_FILE, Error, MAX_SOCKETS, Result, Socket}; +use crate::{ENDPOINT_FILE, Error, Gate, MAX_SOCKETS, Result, Socket}; /// 一次 `recvmsg` 最多读这么多正文。 const CHUNK: usize = 64 * 1024; @@ -25,6 +25,12 @@ 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")))] @@ -55,6 +61,11 @@ impl Listener { 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?; diff --git a/crates/tw-handover/src/windows.rs b/crates/tw-handover/src/windows.rs index 3e471f3f..e5d259b0 100644 --- a/crates/tw-handover/src/windows.rs +++ b/crates/tw-handover/src/windows.rs @@ -9,10 +9,12 @@ 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, @@ -21,16 +23,13 @@ use windows_sys::Win32::Security::Cryptography::{ BCRYPT_USE_SYSTEM_PREFERRED_RNG, BCryptGenRandom, }; -use crate::{ENDPOINT_FILE, Error, Result, Socket, TOKEN_FILE, token}; +use crate::{ENDPOINT_FILE, Error, Gate, Result, Socket, TOKEN_FILE, token}; /// 一次最多读这么多。 const CHUNK: usize = 64 * 1024; /// 连上来之后交口令的时限。 -pub(crate) const TOKEN_WAIT: Duration = Duration::from_secs(5); - -/// 口令帧的长度上限（口令本身是 64 个十六进制字符）。 -pub(crate) const TOKEN_FRAME_MAX: usize = 256; +const TOKEN_WAIT: Duration = Duration::from_secs(5); /// 口令的随机字节数（文件里是它的十六进制）。 const TOKEN_BYTES: usize = 32; @@ -46,6 +45,7 @@ pub(crate) struct Listener { 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 了没有，看见端口时口令必须已经在 @@ -62,12 +62,23 @@ impl Listener { /// 身份核对（口令）在 `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))) } - pub(crate) fn token(&self) -> &[u8] { - self.token.as_bytes() + #[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, + } } } @@ -98,6 +109,7 @@ pub(crate) async fn connect(dir: &Path) -> Result)>> 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))) } @@ -188,9 +200,23 @@ fn socket_from_info(hex: &str) -> Result { ))); } // SAFETY: s 是刚还原出来、只归我们所有的套接字 - Ok(Socket::Tcp(unsafe { - std::net::TcpListener::from_raw_socket(s as RawSocket) - })) + 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 { From 128f6de9a03b7e3845f7b2a4d5bc6b3a43aa2783 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 24 Oct 2026 14:36:50 +0800 Subject: [PATCH 8/9] Handover fixes from review: draining marker, commit ack, renewal pause, settlement while queued - A successor that exits while its predecessor still drains no longer deletes the lock. The successor records the draining core (pid and the request ids it may use) in twcore.draining; the lock is removed only if it still holds our pid and nobody is draining. A core that starts fresh and finds a live draining core opens the request history only after it exits, numbers requests past it and reports handover = draining. A successor that exits before the history opened logs how many queued events were not recorded. - Commits are acknowledged (Message::CommitAck). The old core replies before it stops accepting and never answers a commit after its 30 s limit; the successor counts as primary only on the ack or once the old process is gone, and otherwise exits. - OAuth renewals pause only after WantListeners, so a successor waiting for its configuration (H3) no longer blocks renewals on the core that is still forwarding. Pausing waits for renewals already in flight; the old core then sends a second, fresh snapshot (also re-imported against the real configuration after H3), and a later import replaces OAuth tokens from an earlier one. - While a successor queues records, a placeholder recorder settles them so key and upstream limits keep counting its own requests; on opening, the queue is replayed without settling and usage is rebuilt from the store. - Handing over the control plane closes connections already made to the old core, so clients reconnect to the successor. - An unreadable greeting is answered with Refused. - An arm is called off when the original parent is still alive at the deadline (the desktop app did not restart): endpoint closed, parent watch restored, back to idle. Arming starts a session of its own on unix. - Sign-ins that start after arming are refused (control.handover.login_refused). - The successor waits 30 s for each snapshot; a successor that leaves right after the first snapshot counts as a strike. - twcore.handover.log is appended to (rotated past 4 MB); the lock file is replaced atomically; nothing after the commit makes the successor exit; the successor releases the old core's slots only once its process is gone. Co-Authored-By: Claude Opus 5.5 --- bin/twcore/src/handover.rs | 302 ++++++++++++--- bin/twcore/src/lockfile.rs | 193 +++++++++- bin/twcore/src/main.rs | 160 +++++--- bin/twcore/src/record.rs | 93 +++-- bin/twcore/tests/handover.rs | 470 +++++++++++++++++++++++- crates/tw-api/msg-codes.txt | 1 + crates/tw-control/src/chatgpt.rs | 4 + crates/tw-control/src/handover.rs | 76 ++++ crates/tw-control/src/lib.rs | 39 +- crates/tw-control/src/remote.rs | 4 +- crates/tw-control/src/zai.rs | 4 + crates/tw-control/tests/gate.rs | 39 ++ crates/tw-control/tests/handover.rs | 31 ++ crates/tw-gateway/src/handover/tests.rs | 30 ++ crates/tw-gateway/src/oauth.rs | 36 +- crates/tw-gateway/tests/oauth.rs | 63 ++++ crates/tw-handover/src/lib.rs | 7 +- 17 files changed, 1389 insertions(+), 163 deletions(-) diff --git a/bin/twcore/src/handover.rs b/bin/twcore/src/handover.rs index d787ee67..e8ef844a 100644 --- a/bin/twcore/src/handover.rs +++ b/bin/twcore/src/handover.rs @@ -8,7 +8,7 @@ //! - **新 core**（[`approach`] → [`Incoming`]）：取锁时发现持有者活着，就去连它的交接端点； //! 连上了就在交过来的控制面上服务、读配置建运行时、要监听、提交，成为锁的持有者。 //! -//! **提交是唯一的不可逆点。**提交之前任何一步出错，旧 core 回到待交接、照常服务（它始终 +//! **提交是唯一的不可逆点**，以旧 core 回的 CommitAck 为准（它在时限之后收到的提交不回）。提交之前任何一步出错，旧 core 回到待交接、照常服务（它始终 //! 握着全部监听）；提交之后新 core 是唯一的「主」，旧 core 只做完手上的请求。任何情况下 //! 都不出现两个进程都不监听的时刻。 //! @@ -18,6 +18,7 @@ //! | --- | --- | --- | //! | 待交接 | 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） | //! @@ -52,6 +53,13 @@ pub const DRAIN_LIMIT: Duration = Duration::from_secs(30 * 60); 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`] 加上退出的余量）。 @@ -182,6 +190,9 @@ pub struct Core { 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, @@ -220,8 +231,22 @@ async fn run_arms( core: &Arc, arms: &mut tokio::sync::mpsc::Receiver, ) -> Exit { - // 第一次 arm：开端点、脱离父进程 - let (mut endpoint, mut params) = loop { + 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; }; @@ -231,14 +256,17 @@ async fn run_arms( // 往管道里写会出错，`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(())); - break ( - Some(ep), + return ( + ep, Params { wait: a.wait, relaunch: a.relaunch, @@ -250,7 +278,17 @@ async fn run_arms( 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 { @@ -286,14 +324,14 @@ async fn run_arms( match give(core, link, strikes >= MAX_STRIKES).await { Given::Committed(link) => { drain(core, link).await; - return Exit::HandedOver; + 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 Exit::Retired; + return Some(Exit::Retired); } Given::Abandoned { strike } => { if strike { @@ -310,23 +348,39 @@ async fn run_arms( Given::NotASuccessor => {} } } - _ = tokio::time::sleep_until(clock.deadline) => 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"); + _ = 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); - retire(core, FALLBACK_DRAIN).await; - return Exit::Retired; + 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 }; @@ -337,6 +391,45 @@ async fn run_arms( } } +/// 父进程守望：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), @@ -364,11 +457,21 @@ async fn give(core: &Core, mut link: Link, refuse: bool) -> Given { // 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：退回现行流程。等它说一声 Fallback（它会马上说） + // H2 / F5：退回现行流程 let reason = if refuse { format!("{MAX_STRIKES} handovers in a row failed after the listeners were handed over") } else { @@ -377,15 +480,10 @@ async fn give(core: &Core, mut link: Link, refuse: bool) -> Given { tw_handover::VERSION ) }; - tracing::info!(%reason, "refusing the handover"); - let _ = link.send(&Message::Refused { reason }).await; - return match tokio::time::timeout(HELLO_WITHIN, link.recv()).await { - Ok(Ok((Message::Fallback, _))) => Given::Fallback, - _ => Given::Abandoned { strike: false }, - }; + return refuse_with(&mut link, reason).await; } - // ① 控制面。**交出去就不再从它上面接新连接**（见 `tw_control::handover`） + // ① 控制面。**交出去就不再从它上面接新连接**，已经连着的也断开（见 `tw_control::handover`） let sock = match core.handover().control().export() { Ok(s) => s, Err(e) => { @@ -394,6 +492,7 @@ async fn give(core: &Core, mut link: Link, refuse: bool) -> Given { } }; core.handover().control().pause(); + core.handover().control().close_connections(); let control = Message::Control { info: ListenerInfo { role: Role::Control, @@ -406,20 +505,20 @@ async fn give(core: &Core, mut link: Link, refuse: bool) -> Given { } drop(sock); - // ② 快照。之前先把手上的收尾：正在写回的 refresh token、还没上传的本机改动 + // ② 第一份快照。之前把还没上传的本机改动传上去（同现行退出时的 `sync.flush`） background::settle_rotation(&core.background).await; if let Some(sync) = core.sync() { sync.flush().await; } - // 交出 access token 起暂停续期：两边各换一次 refresh token 必有一边作废（接管没成再恢复） - core.state.oauth.hold_renewals(true); 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: false }; + return Given::Abandoned { strike: true }; } + let sent = Instant::now(); - // ③ 等它要监听。**不限时**（H3：它的配置读不进来时在等人修），断了就是没成 + // ③ 等它要监听。**不限时**（H3：它的配置读不进来时在等人修），断了就是没成。**这一段照常 + // 续期**：这边还在转发，可能要等几个小时 loop { match link.recv().await { Ok((Message::WantListeners, _)) => break, @@ -428,13 +527,30 @@ async fn give(core: &Core, mut link: Link, refuse: bool) -> Given { tracing::warn!(?other, "an unexpected handover message; ignoring it") } Err(e) => { - tracing::info!("the successor left before taking the listeners: {e}"); - return Given::Abandoned { strike: false }; + // 快照刚发出去就断了：多半是它等快照等得不耐烦（或者起来就崩了），记一次； + // 等了很久才断的是在等配置修好的那一种（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() { @@ -477,9 +593,20 @@ async fn give(core: &Core, mut link: Link, refuse: bool) -> Given { drop(refs); drop(socks); - // ⑤ 等提交（F4：30 秒）。提交之前两边同时 accept 是允许的（H4） + // ⑥ 等提交（F4：30 秒）。提交之前两边同时 accept 是允许的（H4）。**过了时限就不再收**： + // 这条连接随之关掉，迟到的 Commit 没人回 CommitAck，接替者自己退出 match tokio::time::timeout(TAKEOVER_LIMIT, link.recv()).await { - Ok(Ok((Message::Commit, _))) => Given::Committed(link), + 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"); @@ -496,6 +623,16 @@ async fn give(core: &Core, mut link: Link, refuse: bool) -> Given { } } +/// 拒绝这一次交接（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"); @@ -581,16 +718,22 @@ fn abort_all(state: &tw_gateway::AppState) -> usize { .count() } -/// arm 之后的日志写到哪：数据目录里，每次 arm 从头写。 +/// 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.write(true).create(true).truncate(true); + o.append(true).create(true); #[cfg(unix)] { use std::os::unix::fs::OpenOptionsExt; @@ -719,6 +862,8 @@ pub struct Incoming { snapshot: Option, /// 旧 core 的 pid（锁文件里的那个） pub predecessor: u32, + /// 这边的请求号从这里起（旧 core 的号 + [`ID_GAP`]）。旧 core 用的都在它之下 + pub ids_from: u64, } /// 旧 core 交过来的网关与远程控制端口的监听。 @@ -770,7 +915,8 @@ pub async fn approach(dir: &Path, holder: u32) -> Approach { // 不说话：它正在交给另一个接替者 Err(_) => return Approach::NotArmed, }; - let snapshot = match tokio::time::timeout(HELLO_WITHIN, link.recv()).await { + // 旧 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"); @@ -784,6 +930,7 @@ pub async fn approach(dir: &Path, holder: u32) -> Approach { control, snapshot, predecessor: holder, + ids_from: 0, }) } @@ -807,6 +954,20 @@ async fn gone_or_busy(holder: u32) -> Approach { } } +/// 前任还在排空、却没有交接连接（它的接替者先退出了，这个进程是照常起来的）：状态报 +/// 「排空中」，它们都退了回到平常。在途数报不出来（没有连接），是 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 { @@ -829,16 +990,27 @@ impl Incoming { ); return; }; - state.bus.resume_after(s.next_id.saturating_add(ID_GAP)); + 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); } - /// 要监听（配置读进来了、运行时建好了）。 - pub async fn listeners(&mut self) -> Result { + /// 要监听（配置读进来了、运行时建好了）。**先收到第二份快照**：旧 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")? @@ -849,12 +1021,27 @@ impl Incoming { inherited(infos, socks) } - /// 提交：从此不再回退（H5）。**发不出去也算提交了**：旧 core 这时候没了的话（F6）， - /// 套接字已经在手上，直接接着服务 - pub async fn commit(&mut self) { - if let Err(e) = self.link.send(&Message::Commit).await { - tracing::warn!("the previous core is gone at commit; serving on without it: {e}"); + /// 提交（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 退出或崩溃）就放掉 @@ -874,9 +1061,14 @@ impl Incoming { Err(_) => break, } } - // 收到 Bye 的话等它的进程真没了再放：它说完 Bye 到退出之间手上已经没有请求了， - // 但名额要等它真的不可能再用才算放开。连接断了（崩了）也一样等一下 - let _ = wait_gone(self.predecessor, Duration::from_secs(60)).await; + // **等它的进程真没了再放**：说完 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）就别改回去 @@ -1015,6 +1207,12 @@ mod tests { 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 的时限要短一截 diff --git a/bin/twcore/src/lockfile.rs b/bin/twcore/src/lockfile.rs index bfd7c6e6..5bd983c8 100644 --- a/bin/twcore/src/lockfile.rs +++ b/bin/twcore/src/lockfile.rs @@ -4,9 +4,11 @@ //! 会表现为启动失败，而错误信息必须能说清是这个原因 —— 否则用户会以为 //! 自己电脑坏了。 -use std::io::Write; use std::path::{Path, PathBuf}; +/// 锁文件的名字。 +pub const LOCK_FILE: &str = "twcore.lock"; + #[derive(Debug)] pub struct LockFile { path: PathBuf, @@ -27,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() @@ -46,7 +48,7 @@ impl LockFile { /// [`LockFile::disown`]）。 pub fn claim(dir: &Path) -> std::io::Result { std::fs::create_dir_all(dir)?; - Self::write(dir.join("twcore.lock")) + Self::write(dir.join(LOCK_FILE)) } /// 放弃锁：退出时不删文件（交接提交之后，文件里马上会是接替者的 pid）。 @@ -54,17 +56,26 @@ impl LockFile { self.owned = false; } + /// **先写临时文件再改名**：改名是原子的，读锁的那一方看到的要么是旧的、要么是新的， + /// 不会读到写了一半的（空的）—— 空的会被当成「上次没退干净」直接覆盖 fn write(path: PathBuf) -> std::io::Result { - let mut f = std::fs::File::create(&path)?; - write!(f, "{}", std::process::id())?; + 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`）。 -#[cfg(test)] pub fn holder(dir: &Path) -> Option { - std::fs::read_to_string(dir.join("twcore.lock")) + std::fs::read_to_string(dir.join(LOCK_FILE)) .ok() .and_then(|t| t.trim().parse().ok()) } @@ -74,11 +85,108 @@ impl Drop for LockFile { 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::*; @@ -168,7 +276,7 @@ mod tests { fn a_disowned_lock_leaves_the_file_for_the_successor() { // 交接提交之后旧 core 放弃锁：它退出时文件里写的已经是接替者，不能删 let d = tempfile::tempdir().unwrap(); - let p = d.path().join("twcore.lock"); + let p = d.path().join(LOCK_FILE); let LockOutcome::Acquired(l) = LockFile::acquire(d.path()).unwrap() else { panic!("应该拿到锁"); }; @@ -186,12 +294,79 @@ mod tests { 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 996d42cd..83aaa576 100644 --- a/bin/twcore/src/main.rs +++ b/bin/twcore/src/main.rs @@ -812,6 +812,27 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> }, }; 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 步）。 // @@ -904,6 +925,9 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> state.oauth.hold_renewals(true); inc.import(&state); } + if let Some(below) = orphaned_ids { + state.bus.resume_after(below.saturating_add(handover::ID_GAP)); + } // 观测这一层。**起不来不是致命的** —— 历史记录看不见，而网关 // 照常转发。所以这里所有的失败都只记一行日志。 @@ -920,38 +944,40 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> runs: run_rx, settled: tw_control::key_limits::settle_hook(&state), }; - let (store, store_opened) = match &incoming { - None => { - 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) + 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()); } - // 接替者：旧 core 退出之后才打开请求库（R3，见 `record`）。用量从记录重建在打开的 - // 那一刻、这个进程排着的记录进库之前 - Some(inc) => { - let gw = state.clone(); - let (store, opened) = record::open_later( - &dir, - wiring, - inc.predecessor, - Box::new(move |db| tw_control::key_limits::rebuild(&gw, db)), - ); - 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); - } - }); - } - (store, opened) + (store, watch::channel(true).1) + } else { + // 接替者（或者前任还在排空）：前任都退了才打开请求库（R3，见 `record`）。用量从 + // 记录重建在打开的那一刻；排空标记在那时清掉 + let gw = state.clone(); + 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); + } + }); } + (store, opened) }; if store.is_some() { state.set_body_sink(body_tx); @@ -1012,6 +1038,10 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> }; // 安全模式没有数据面：没有可交接的，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 交过来的监听上服务；网关、远程控制端口提交之前不跟配置走 @@ -1069,17 +1099,7 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> // 谁来收养都一样。 // // **arm 之后关掉**（见 `handover`）：桌面端装好更新要重启，core 留下来等接替者 - let parent_watch = parent.map(|ppid| { - 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 parent_watch = parent.map(handover::watch_parent); if safe { // 安全模式：只起控制面。数据面不动，让用户还能改配置、回滚。 @@ -1131,7 +1151,7 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> } } let inherited = tokio::select! { - r = inc.listeners() => r?, + r = inc.listeners(&state) => r?, msg = &mut control_dead => { anyhow::bail!("{}", msg.unwrap_or_else(|_| "the control plane stopped".into())) } @@ -1157,32 +1177,49 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> tokio::time::timeout(handover::commit_within(), gateway.listening()) .await .context("the listeners handed over could not be served in time")?; - inc.commit().await; - // 从此续期由这一边做（旧 core 交出快照起已暂停） + // 旧 core 回了 CommitAck（或者它已经没了）才算提交；它还在却没回，这边退出， + // 守护重来（这时它仍握着全部监听） + inc.commit().await?; + // ── 从这里起不再回退（H5）。旧 core 已经停止 accept：**下面任何一步出错都只记 + // 一行、照常服务**，退出的话就没人在听了 + // + // 从此续期由这一边做（旧 core 交出第二份快照起已暂停） state.oauth.hold_renewals(false); - // 成为锁的持有者：写上自己的 pid（旧 core 已经放弃了锁） + // 成为锁的持有者：写上自己的 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 = bg - .lock() - .map_err(|_| anyhow::anyhow!("the background lock is poisoned"))? - .start(background::Start { - control: &control, - safe: false, - seeder: Some(s), - store: store.clone(), - store_opened: store_opened.clone(), - }) - .context("starting the background tasks")?; - sync = Some(started); + 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 } @@ -1199,6 +1236,7 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> watchers: Mutex::new(watchers), lock: Mutex::new(lock), parent: Mutex::new(parent_watch), + parent_pid: parent, store_opened, handed_over: Default::default(), }); @@ -1249,6 +1287,14 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> } } }); + // 前任还在写库、这边的库还没开就退了：排着的记录写不进去，**说一声丢了多少** + 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 { diff --git a/bin/twcore/src/record.rs b/bin/twcore/src/record.rs index cfd52de6..0947e5d5 100644 --- a/bin/twcore/src/record.rs +++ b/bin/twcore/src/record.rs @@ -70,17 +70,27 @@ pub fn open_now(dir: &Path, w: Wiring) -> Option { Some(recorder) } -/// 接替者：旧 core（`predecessor`）的进程没了才打开请求库，之间的记录排队。 +/// 接替者：旧 core（`predecessors`，交接链上可能不止一个）的进程都没了才打开请求库，之间的 +/// 记录排队。 /// /// 交出去的是占位的那一份（内存库）；打开之后**原地换成真的**，控制面拿着的还是同一个 -/// `Arc`。`on_open` 在真的库打开、排队的记录写进去之前调一次（密钥用量从记录重建，和 -/// 平常启动时同一个时机：在这个进程自己的记录进库之前）。返回的 `watch` 在打开之后 -/// （或者打开失败、不再等之后）变成 `true`。 +/// `Arc`。返回的 `watch` 在打开之后（或者打开失败、不再等之后）变成 `true`。 +/// +/// # 排队期间的结算 +/// +/// 用量上限要靠「记下一行就结算」（`settled`）把预留换成实数；不结算的预留过一会儿会被当成 +/// 走丢的放掉，这一期的用量里就少了这个进程自己的请求。所以**占位的那一份照样结算**：事件 +/// 先交给它（它记进内存库、按同一套价目表算钱、调同一个结算），再排进队列。 +/// +/// 真的库打开时**不重复算**：排着的事件交给一份**不结算、不报价钱**的新记录器写进库，写完 +/// 再从库里重建用量（`on_open`，库里这时有旧 core 的和这边排着的全部），最后才给它接上结算 +/// 和总线。重建是换掉而不是加上，排队期间结算进内存的那些被库里的同一批行换掉，一笔不多。 pub fn open_later( dir: &Path, w: Wiring, - predecessor: u32, + 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 { @@ -95,7 +105,9 @@ pub fn open_later( 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}" @@ -114,27 +126,22 @@ pub fn open_later( .name("tw-recorder".into()) .spawn(move || { let mut queue = Queue::default(); - wait_for_exit(predecessor, &mut events, &mut queue); + 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); - // 通道要有人收，否则总线那一头一直记「丢了」 - while events.blocking_recv().is_some() {} + // 通道要有人收，否则总线那一头一直记「丢了」。占位的那一份照样结算 + record_loop(&shared, &mut events); return; } }; resume_ids(&w.bus, &db); - on_open(&db); let blobs_root = blobs.root().to_path_buf(); - *shared.blocking_lock() = tw_store::Recorder::new(db, blobs, w.pricing) - .reporting_to(w.bus) - .settling_to(w.settled); - spawn_body_writer(blobs_root, body_rx); - tw_store::task::record_plugin_runs(shared.clone(), run_rx); let (n, dropped) = (queue.len(), queue.dropped); if dropped > 0 { tracing::warn!( @@ -142,12 +149,18 @@ pub fn open_later( "the request records queued while the previous core drained overflowed; the oldest {dropped} events were not recorded" ); } - { - let mut g = shared.blocking_lock(); - for ev in queue.drain() { - g.on_event(&ev); - } + // 排着的写进库：**不结算、不报价钱**（占位的那一份已经做过了） + 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); @@ -159,26 +172,58 @@ pub fn open_later( (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); -/// 等旧 core 的进程没了，之间把记录收进队列。 -fn wait_for_exit(pid: u32, events: &mut tw_observe::RecordFeed, queue: &mut Queue) { +/// 等前任的进程都没了，之间把记录交给占位的那一份（结算）、收进队列。 +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 !crate::proc::alive(pid) { + if pids.iter().all(|p| !crate::proc::alive(*p)) { return; } if started.elapsed() >= PREDECESSOR_WAIT { tracing::warn!( - pid, + ?pids, "the previous core is still running long past its drain limit; opening the request history anyway" ); return; diff --git a/bin/twcore/tests/handover.rs b/bin/twcore/tests/handover.rs index ae43edd6..0ad77378 100644 --- a/bin/twcore/tests/handover.rs +++ b/bin/twcore/tests/handover.rs @@ -131,11 +131,24 @@ struct Bed { 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\nproviders:\n - name: up\n base_url: http://{up}\n key: k\n protocol: anthropic\n" + "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 { @@ -149,13 +162,62 @@ impl Bed { } 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().stderr(log).spawn().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() } - /// 和桌面端一样，stderr 接一根管道（桌面端退出时这根管道的读端就没了） + /// 和桌面端一样，stdout、stderr 都接管道（桌面端退出时这两根管道的读端就没了） fn core_on_a_pipe(&self) -> Child { - self.command().stderr(Stdio::piped()).spawn().unwrap() + self.command() + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap() } fn command(&self) -> Command { @@ -167,8 +229,7 @@ impl Bed { .env_remove("HTTP_PROXY") .env_remove("HTTPS_PROXY") .env_remove("http_proxy") - .env_remove("https_proxy") - .stdout(Stdio::null()); + .env_remove("https_proxy"); c } @@ -199,7 +260,7 @@ impl Bed { fn logs(&self) -> String { let mut s = String::new(); - for n in ["a", "b", "twcore.handover"] { + 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")); } @@ -265,7 +326,6 @@ async fn quick(c: &reqwest::Client, gw: SocketAddr) -> Result<(), String> { } #[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()]); @@ -336,6 +396,7 @@ async fn a_handover_under_load_refuses_no_connection_and_finishes_the_long_strea 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(); @@ -415,7 +476,6 @@ async fn fake_successor(dir: &Path, version: u32) -> tw_handover::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")]); @@ -442,6 +502,9 @@ async fn a_successor_that_leaves_before_committing_leaves_the_old_core_serving() 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()); @@ -464,7 +527,6 @@ async fn a_successor_that_leaves_before_committing_leaves_the_old_core_serving() } #[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")]); @@ -497,7 +559,6 @@ async fn a_different_handover_version_falls_back_to_drain_and_exit() { } #[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")]); @@ -521,3 +582,390 @@ async fn without_a_successor_the_armed_core_drains_and_exits() { ); 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(20), || 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(20), || 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(20), || 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(20), || 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(20), || 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(20), || 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(20), || 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(20), || exited(&mut reap.0[0])).await, + "{}", + bed.logs() + ); + assert!( + until(Duration::from_secs(20), || 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 45737acd..0f7bd9f0 100644 --- a/crates/tw-api/msg-codes.txt +++ b/crates/tw-api/msg-codes.txt @@ -180,6 +180,7 @@ 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 diff --git a/crates/tw-control/src/chatgpt.rs b/crates/tw-control/src/chatgpt.rs index cd108bfe..00212605 100644 --- a/crates/tw-control/src/chatgpt.rs +++ b/crates/tw-control/src/chatgpt.rs @@ -240,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 index 29974f04..3a391aa4 100644 --- a/crates/tw-control/src/handover.rs +++ b/crates/tw-control/src/handover.rs @@ -140,6 +140,22 @@ pub(crate) fn refused_busy() -> Fail { ) } +/// 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!( @@ -210,6 +226,8 @@ pub(crate) async fn arm( /// 接着接。 pub struct ControlListener { accepting: watch::Sender, + /// 跳一下：已经连着的本机连接全部断开（见 [`ControlListener::close_connections`]） + closing: watch::Sender, held: Mutex>, inherited: Mutex>, } @@ -218,6 +236,7 @@ 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), } @@ -267,6 +286,32 @@ impl ControlListener { 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, + } } /// 再复制一个句柄。 @@ -319,6 +364,37 @@ mod tests { 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()); diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index dd9adf02..860d221e 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -2122,7 +2122,12 @@ async fn accept_loop( handover::until_accepting(&mut accepting).await; tokio::select! { r = listener.accept_one() => match r { - Ok(stream) => hand_off(stream, app.clone(), gate.clone(), None), + 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：不然它还会再接一个 @@ -2215,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, @@ -2252,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; }; @@ -2283,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 688ebfba..10fdf37b 100644 --- a/crates/tw-control/src/remote.rs +++ b/crates/tw-control/src/remote.rs @@ -552,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/zai.rs b/crates/tw-control/src/zai.rs index d1a7cd05..d3a75c3c 100644 --- a/crates/tw-control/src/zai.rs +++ b/crates/tw-control/src/zai.rs @@ -265,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/gate.rs b/crates/tw-control/tests/gate.rs index e400bf14..ec7258dc 100644 --- a/crates/tw-control/tests/gate.rs +++ b/crates/tw-control/tests/gate.rs @@ -255,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 index eff4e626..0763a0df 100644 --- a/crates/tw-control/tests/handover.rs +++ b/crates/tw-control/tests/handover.rs @@ -148,3 +148,34 @@ async fn status_reports_the_handover_phase_and_the_draining_count() { 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-gateway/src/handover/tests.rs b/crates/tw-gateway/src/handover/tests.rs index edf66652..1510fabf 100644 --- a/crates/tw-gateway/src/handover/tests.rs +++ b/crates/tw-gateway/src/handover/tests.rs @@ -550,3 +550,33 @@ async fn a_feed_that_overflowed_is_replaced_by_a_full_report_at_drain() { "排空开始时没报占着的名额" ); } + +/// 快照来两份（旧 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/oauth.rs b/crates/tw-gateway/src/oauth.rs index f1e874f2..623f4a0b 100644 --- a/crates/tw-gateway/src/oauth.rs +++ b/crates/tw-gateway/src/oauth.rs @@ -385,8 +385,12 @@ impl Cache { } /// 交接：接进旧 core 手里的 token。这里的配置认得它才接：生成它的那份配置和这里的一样， - /// 或者它手里的 refresh 就是这里配置里的（轮换之后写回了）。这里已经有的那一家不动。 - /// `lag_ms` 是快照导出到现在过了多久 + /// 或者它手里的 refresh 就是这里配置里的（轮换之后写回了）。`lag_ms` 是快照导出到现在 + /// 过了多久。 + /// + /// **这里已经有的那一家换掉**：提交之前续期归旧 core（这边暂停着），旧 core 手里的才是 + /// 准的。快照会来两份（交出监听之前再导出一次，那时旧 core 已经暂停续期），后一份要盖掉 + /// 前一份；这边从配置里读进来的那个也一样 pub(crate) fn handover_import( &self, entries: Vec, @@ -404,9 +408,7 @@ impl Cache { else { continue; }; - if (e.cfg != config_digest(cfg) && e.refresh != cfg.refresh) - || g.contains_key(&e.upstream) - { + if e.cfg != config_digest(cfg) && e.refresh != cfg.refresh { continue; } let after = |ms: u64| now + Duration::from_millis(ms).saturating_sub(lag); @@ -492,6 +494,30 @@ impl Cache { 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, diff --git a/crates/tw-gateway/tests/oauth.rs b/crates/tw-gateway/tests/oauth.rs index 19081904..54386415 100644 --- a/crates/tw-gateway/tests/oauth.rs +++ b/crates/tw-gateway/tests/oauth.rs @@ -553,3 +553,66 @@ async fn renewals_held_during_a_handover_never_reach_the_token_endpoint() { 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/src/lib.rs b/crates/tw-handover/src/lib.rs index 084eb633..b9648ace 100644 --- a/crates/tw-handover/src/lib.rs +++ b/crates/tw-handover/src/lib.rs @@ -97,8 +97,10 @@ pub enum Role { /// 旧 → 新 Refused（版本不同，到此为止）| Control + 控制面套接字 /// 旧 → 新 Snapshot /// 新 → 旧 WantListeners （新 core 的配置读进来了、运行时建好了） +/// 旧 → 新 Snapshot （再导出一次：暂停续期之后的 OAuth、此刻的状态） /// 旧 → 新 Listeners + 网关、远程控制端口的套接字 -/// 新 → 旧 Commit （新 core 已在全部监听上 accept；不可逆） +/// 新 → 旧 Commit （新 core 已在全部监听上 accept） +/// 旧 → 新 CommitAck （旧 core 收下了提交、不再 accept；从此不可逆） /// 旧 → 新 Settled* / InFlight* （排空期，每个请求结束一条） /// 旧 → 新 Bye （马上退出） /// 新 → 旧 Fallback （任何时候：请按现行方式排空后退出，新的自己绑） @@ -125,6 +127,9 @@ pub enum Message { infos: Vec, }, Commit, + /// 旧 core 收下了提交。**新 core 收到它（或者确认旧进程已经没了）才算提交**：旧 core 在 + /// 等提交的时限到了之后放弃、又收到一个迟到的 Commit 时不回它，两边就不会都以为自己是主 + CommitAck, Settled { data: serde_json::Value, }, From 73849f29660a2e9baaaf900cec0c4f5ed447b3ab Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:06:00 +0800 Subject: [PATCH 9/9] Update ctutils and tokio past RUSTSEC-2026-0385/0386; give the handover tests more time to start under load Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 28 ++++++++++++++-------------- bin/twcore/tests/handover.rs | 30 +++++++++++++++--------------- 2 files changed, 29 insertions(+), 29 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8dcf7b4f..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", @@ -4583,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/bin/twcore/tests/handover.rs b/bin/twcore/tests/handover.rs index 0ad77378..9b50ac57 100644 --- a/bin/twcore/tests/handover.rs +++ b/bin/twcore/tests/handover.rs @@ -330,7 +330,7 @@ async fn a_handover_under_load_refuses_no_connection_and_finishes_the_long_strea let bed = Bed::new(25).await; let mut reap = Reap(vec![bed.core_on_a_pipe()]); assert!( - until(Duration::from_secs(20), || bed.status().is_some()).await, + until(Duration::from_secs(60), || bed.status().is_some()).await, "a 没起来\n{}", bed.logs() ); @@ -401,7 +401,7 @@ async fn a_handover_under_load_refuses_no_connection_and_finishes_the_long_strea reap.0.push(bed.core("b")); let b_pid = reap.0[1].id(); assert!( - until(Duration::from_secs(20), || bed.lock_pid() == Some(b_pid)).await, + until(Duration::from_secs(60), || bed.lock_pid() == Some(b_pid)).await, "锁没交给 b\n{}", bed.logs() ); @@ -420,7 +420,7 @@ async fn a_handover_under_load_refuses_no_connection_and_finishes_the_long_strea // a 排空后退出 assert!( - until(Duration::from_secs(20), || exited(&mut reap.0[0])).await, + until(Duration::from_secs(60), || exited(&mut reap.0[0])).await, "a 没退\n{}", bed.logs() ); @@ -480,7 +480,7 @@ 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(20), || bed.status().is_some()).await, + until(Duration::from_secs(60), || bed.status().is_some()).await, "{}", bed.logs() ); @@ -531,7 +531,7 @@ 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(20), || bed.status().is_some()).await, + until(Duration::from_secs(60), || bed.status().is_some()).await, "{}", bed.logs() ); @@ -563,7 +563,7 @@ 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(20), || bed.status().is_some()).await, + until(Duration::from_secs(60), || bed.status().is_some()).await, "{}", bed.logs() ); @@ -606,7 +606,7 @@ 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(20), || bed.status().is_some()).await, + until(Duration::from_secs(60), || bed.status().is_some()).await, "{}", bed.logs() ); @@ -641,7 +641,7 @@ 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(20), || bed.status().is_some()).await, + until(Duration::from_secs(60), || bed.status().is_some()).await, "{}", bed.logs() ); @@ -750,7 +750,7 @@ async fn an_unreadable_greeting_is_answered_with_refused() { let bed = Bed::new(1).await; let _reap = Reap(vec![bed.core("a")]); assert!( - until(Duration::from_secs(20), || bed.status().is_some()).await, + until(Duration::from_secs(60), || bed.status().is_some()).await, "{}", bed.logs() ); @@ -801,7 +801,7 @@ async fn an_arm_is_called_off_when_the_desktop_app_did_not_restart() { let ppid = parent.id().to_string(); let mut reap = Reap(vec![bed.core_with("a", &["--parent", &ppid]), parent]); assert!( - until(Duration::from_secs(20), || bed.status().is_some()).await, + until(Duration::from_secs(60), || bed.status().is_some()).await, "{}", bed.logs() ); @@ -852,7 +852,7 @@ async fn a_successor_that_exits_while_its_predecessor_drains_leaves_the_next_cor let bed = Bed::with_daily_limit(450).await; // 长流约 90 秒 let mut reap = Reap(vec![bed.core("a")]); assert!( - until(Duration::from_secs(20), || bed.status().is_some()).await, + until(Duration::from_secs(60), || bed.status().is_some()).await, "{}", bed.logs() ); @@ -881,7 +881,7 @@ async fn a_successor_that_exits_while_its_predecessor_drains_leaves_the_next_cor reap.0.push(bed.core("b")); let b_pid = reap.0[1].id(); assert!( - until(Duration::from_secs(20), || bed.lock_pid() == Some(b_pid)).await, + until(Duration::from_secs(60), || bed.lock_pid() == Some(b_pid)).await, "{}", bed.logs() ); @@ -916,7 +916,7 @@ async fn a_successor_that_exits_while_its_predecessor_drains_leaves_the_next_cor reap.0.push(bed.core("c")); let c_pid = reap.0[2].id(); assert!( - until(Duration::from_secs(20), || bed + until(Duration::from_secs(60), || bed .status() .is_some_and(|s| s["pid"].as_u64() == Some(c_pid as u64))) .await, @@ -938,12 +938,12 @@ async fn a_successor_that_exits_while_its_predecessor_drains_leaves_the_next_cor let body = long.await.unwrap().expect("长流断了"); assert!(body.contains("message_stop"), "{body}"); assert!( - until(Duration::from_secs(20), || exited(&mut reap.0[0])).await, + until(Duration::from_secs(60), || exited(&mut reap.0[0])).await, "{}", bed.logs() ); assert!( - until(Duration::from_secs(20), || bed + until(Duration::from_secs(60), || bed .status() .is_some_and(|s| s["handover"] == "idle")) .await,