From e6ba178983a048f2e6065c633af2770cce3e6771 Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Sat, 15 Aug 2026 00:34:59 +0800 Subject: [PATCH] feat(send): add per-chat token bucket rate limiting --- crates/xmedia-bot/Cargo.toml | 5 +- crates/xmedia-bot/src/db.rs | 6 ++ crates/xmedia-bot/src/main.rs | 1 + crates/xmedia-bot/src/media_sender.rs | 15 ++- crates/xmedia-bot/src/rate_limit.rs | 136 ++++++++++++++++++++++++++ docs/architecture-refactor.md | 38 ++++--- 6 files changed, 185 insertions(+), 16 deletions(-) create mode 100644 crates/xmedia-bot/src/rate_limit.rs diff --git a/crates/xmedia-bot/Cargo.toml b/crates/xmedia-bot/Cargo.toml index 3b41890..236853d 100644 --- a/crates/xmedia-bot/Cargo.toml +++ b/crates/xmedia-bot/Cargo.toml @@ -5,7 +5,7 @@ edition = "2024" [dependencies] teloxide = { version = "0.17", default-features = false, features = ["webhooks-axum", "macros", "rustls", "ctrlc_handler"] } -tokio = { version = "1.40", features = ["rt-multi-thread", "macros"] } +tokio = { version = "1.40", features = ["rt-multi-thread", "macros", "time"] } serde = { version = "1", features = ["derive"] } serde_json = "1" log = "0.4" @@ -23,3 +23,6 @@ zune-jpeg = "0.5" fast_image_resize = "6" jpeg-encoder = "0.7" x-media = { path = "../x-media" } + +[dev-dependencies] +tokio = { version = "1.40", features = ["test-util"] } diff --git a/crates/xmedia-bot/src/db.rs b/crates/xmedia-bot/src/db.rs index 4ca2306..d2f647a 100644 --- a/crates/xmedia-bot/src/db.rs +++ b/crates/xmedia-bot/src/db.rs @@ -134,6 +134,12 @@ fn rusqlite_error(e: std::io::Error) -> rusqlite::Error { /// Creates the `tasks`, `chat_state` and `link_cache` tables (idempotent). /// The three stores used to own their own schema; keeping it in one place /// means one initialization for the whole database file. +/// +/// ⚠️ Schema-change reminder (deferred, see `docs/architecture-refactor.md` +/// §5): this is a plain `CREATE TABLE IF NOT EXISTS` with no versioning. +/// Before any column/table change that must migrate existing databases, land +/// the `PRAGMA user_version` migration chain first (`MIGRATIONS: &[&str]` + +/// `migrate(conn)`), then restructure this function. pub fn schema_init(conn: &Connection) -> rusqlite::Result<()> { conn.execute_batch( "CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, payload TEXT NOT NULL, \ diff --git a/crates/xmedia-bot/src/main.rs b/crates/xmedia-bot/src/main.rs index 4d451cc..65994e6 100644 --- a/crates/xmedia-bot/src/main.rs +++ b/crates/xmedia-bot/src/main.rs @@ -14,6 +14,7 @@ mod link_cache; mod media_sender; mod photo; mod queue; +mod rate_limit; mod send; mod state; diff --git a/crates/xmedia-bot/src/media_sender.rs b/crates/xmedia-bot/src/media_sender.rs index d0b5725..c1b103e 100644 --- a/crates/xmedia-bot/src/media_sender.rs +++ b/crates/xmedia-bot/src/media_sender.rs @@ -73,6 +73,11 @@ impl MediaSender for Bot { items: Vec, ) -> BoxFuture<'_, Result, RequestError>> { Box::pin(async move { + // Pace media sends per chat (one token per item) so bursts do not + // trip Telegram's flood control. + crate::rate_limit::limiter_for(chat_id.0) + .acquire(items.len() as f64) + .await; // `::` disambiguates from this trait's same-named // method (teloxide's API lives in the `Requester` trait). ::send_media_group(self, chat_id, items) @@ -90,6 +95,7 @@ impl MediaSender for Bot { file: InputFile, ) -> BoxFuture<'a, Result> { Box::pin(async move { + crate::rate_limit::limiter_for(chat_id.0).acquire(1.0).await; let mut request = ::send_animation(self, chat_id, file) .caption(caption) .parse_mode(ParseMode::Html) @@ -107,7 +113,14 @@ impl MediaSender for Bot { from: ChatId, ids: Vec, ) -> BoxFuture<'_, Result, RequestError>> { - Box::pin(async move { ::copy_messages(self, to, from, ids).await }) + Box::pin(async move { + // Channel forwards are the burstiest path (batch copies); pace + // them per message against the channel's budget. + crate::rate_limit::limiter_for(to.0) + .acquire(ids.len() as f64) + .await; + ::copy_messages(self, to, from, ids).await + }) } fn send_message( diff --git a/crates/xmedia-bot/src/rate_limit.rs b/crates/xmedia-bot/src/rate_limit.rs new file mode 100644 index 0000000..034b012 --- /dev/null +++ b/crates/xmedia-bot/src/rate_limit.rs @@ -0,0 +1,136 @@ +//! Per-chat token-bucket rate limiting. +//! +//! Telegram throttles bots that burst past a chat's message budget +//! (roughly 20 messages/min for channels/groups); today the bot absorbs +//! those 429s with queue retries. This limiter smooths the burst *before* +//! it reaches the API: media sends to a chat consume one token per +//! message, refilled at [`REFILL_PER_SEC`], so a batch forward paces itself +//! instead of tripping flood control. The queue retry stays as the safety +//! net for limits this bucket does not model (global per-bot limits etc.). + +use parking_lot::Mutex; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::LazyLock; +use std::time::Duration; + +/// Burst capacity: how many messages may be sent at once without waiting. +const CAPACITY: f64 = 20.0; +/// Sustained refill: ~20 messages per minute. +const REFILL_PER_SEC: f64 = 20.0 / 60.0; + +struct State { + /// Current token balance; may go negative (debt from an acquire larger + /// than the capacity, repaid by subsequent refills). + tokens: f64, + last_refill: tokio::time::Instant, +} + +/// A token bucket: at most `CAPACITY` tokens accumulate, refilled at +/// `REFILL_PER_SEC`. [`TokenBucket::acquire`] consumes `n` tokens, waiting +/// for the deficit (a single acquire may exceed the capacity and goes into +/// debt, which the refill repays). +pub struct TokenBucket { + capacity: f64, + refill_per_sec: f64, + state: Mutex, +} + +impl TokenBucket { + fn new(capacity: f64, refill_per_sec: f64) -> Self { + TokenBucket { + capacity, + refill_per_sec, + state: Mutex::new(State { + tokens: capacity, + last_refill: tokio::time::Instant::now(), + }), + } + } + + /// Waits until `n` tokens are available, consuming them. The wait is + /// bounded: the deficit is committed as debt and repaid over time, so a + /// large acquire returns once its share of the refill budget has passed. + pub async fn acquire(&self, n: f64) { + // The parking_lot guard is confined to this block: only the plain + // `wait` duration crosses the await (a guard across an await point + // would make the future !Send). + let wait = { + let mut state = self.state.lock(); + let now = tokio::time::Instant::now(); + let elapsed = now + .saturating_duration_since(state.last_refill) + .as_secs_f64(); + // Refill up to the capacity; a debt (negative balance) is repaid + // before any surplus accumulates. + state.tokens = (state.tokens + elapsed * self.refill_per_sec).min(self.capacity); + state.last_refill = now; + if state.tokens >= n { + state.tokens -= n; + return; + } + // Commit the whole consumption now; the caller proceeds once the + // deficit's worth of refill time has passed. + let debt = n - state.tokens; + state.tokens = -debt; + debt / self.refill_per_sec + }; + tokio::time::sleep(Duration::from_secs_f64(wait)).await; + } +} + +/// One limiter per chat, created on first use. Per-chat so one chat's burst +/// never throttles another. +static LIMITERS: LazyLock>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +/// Returns the shared limiter for a chat, creating it on first use. +pub fn limiter_for(chat_id: i64) -> Arc { + LIMITERS + .lock() + .entry(chat_id) + .or_insert_with(|| Arc::new(TokenBucket::new(CAPACITY, REFILL_PER_SEC))) + .clone() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn limiter_for_reuses_the_per_chat_bucket() { + let a = limiter_for(1); + let b = limiter_for(1); + let c = limiter_for(2); + assert!(Arc::ptr_eq(&a, &b), "same chat → same bucket"); + assert!(!Arc::ptr_eq(&a, &c), "different chat → different bucket"); + } + + #[tokio::test(start_paused = true)] + async fn burst_is_consumed_instantly_then_refill_waits() { + let bucket = TokenBucket::new(3.0, 1.0); + // A burst within capacity passes without waiting. + bucket.acquire(3.0).await; + // The bucket is empty now; one token needs 1s of refill. + let start = tokio::time::Instant::now(); + bucket.acquire(1.0).await; + assert!( + start.elapsed() >= Duration::from_secs(1), + "elapsed {:?}", + start.elapsed() + ); + } + + #[tokio::test(start_paused = true)] + async fn acquire_larger_than_capacity_waits_for_the_deficit() { + let bucket = TokenBucket::new(2.0, 1.0); + // 5 tokens with a capacity of 2: the 3-token deficit takes 3s. + let start = tokio::time::Instant::now(); + bucket.acquire(5.0).await; + assert!( + start.elapsed() >= Duration::from_secs(3), + "elapsed {:?}", + start.elapsed() + ); + } +} diff --git a/docs/architecture-refactor.md b/docs/architecture-refactor.md index d7d9283..32dba43 100644 --- a/docs/architecture-refactor.md +++ b/docs/architecture-refactor.md @@ -1,6 +1,7 @@ # 架构优化设计：可测试性接缝 + handlers 拆分 -> 状态：**阶段 A、B 已实施**（A: `c9e72fd`，B: `50206a9`）；C、D 为可选后续。 +> 状态：**阶段 A、B、C 已实施**（A: `c9e72fd`，B: `50206a9` + `ae69d72`，C: +> rate_limit 提交）；**D 已延迟**——待下次数据库 schema 变化时实施（见 §5）。 > 目标：把仓库最大的测试空白（`handlers.rs`/`send.rs` 的发送与分派逻辑）补上 > 可测试接缝，并把 ~1100 行的 handlers 单体拆成模块。 @@ -74,26 +75,34 @@ impl MediaSender for Bot { /* 委托现有 teloxide 调用 */ } **风险**：中。动 `send.rs`/`handlers.rs` 签名（约 15 处调用点），行为不变。 **不做**：`main.rs` 的 teloxide 装配不抽象（那是真正的胶水，无测试价值）。 -## 4. 阶段 C（可选）：主动限流 +## 4. 阶段 C：主动限流（已实施） 批量转发时的突发会触发 Telegram 频道限速，现在靠 `RetryAfter → 队列重试` 被动 -应对。新增轻量令牌桶（`rate_limit.rs`，~50 行）： +应对。新增轻量令牌桶（`rate_limit.rs`）： ```rust -pub struct TokenBucket { /* capacity, refill_rate, state */ } +pub struct TokenBucket { capacity, refill_per_sec, state: Mutex } impl TokenBucket { - pub async fn acquire(&self, n: u64) -> Duration; // 等待时长（或 Notify 唤醒） + pub async fn acquire(&self, n: f64); // 按 n 个 token 等待并消费 } +pub fn limiter_for(chat_id: i64) -> Arc; // 每频道一个桶 ``` -- 按频道粒度（`HashMap>`），在 `send_media_group`/ - `copy_messages` 前置 `acquire`。 -- 收益：减少 429 → 重试 → 死信；风险低，独立模块。 -- 不做的理由（若选不做）：当前重试链路已能自愈，容量可按需再加。 +- 默认 `CAPACITY = 20`、`REFILL_PER_SEC = 20/60`（约 20 msg/min）； + 单次 acquire 可超出容量（记为债务，由后续 refill 偿还）。 +- 挂点：`MediaSender for Bot` 的 `send_media_group`（按 items 数）、 + `copy_messages`（按 ids 数）、`send_animation`（1 token）前置 `acquire`； + MockSender 不受影响（测试不经过限流）。 +- 收益：减少 429 → 重试 → 死信；队列重试仍是全局限速的安全网。 +- 风险：低，独立模块；`tokio::time`（paused-clock 可测）。 -## 5. 阶段 D（可选）：DB 版本化迁移 +## 5. 阶段 D：DB 版本化迁移（**已延迟**） -`schema_init` 是 `CREATE TABLE IF NOT EXISTS`，无版本概念。改为： +> ⚠️ **待办提醒**：本阶段**推迟到下次数据库 schema 变化时实施**（给 +> `link_cache`/`chat_state`/`tasks` 加列、改结构等）。当前 `schema_init` 是 +> `CREATE TABLE IF NOT EXISTS`，无版本概念；一旦需要迁移已有线上库，必须先落地 +> 本方案（`PRAGMA user_version` 迁移链）再改 schema。`db.rs` 的 `schema_init` +> 处已留注释指向这里。 ```rust // db.rs @@ -127,7 +136,8 @@ pub fn migrate(conn: &Connection) -> rusqlite::Result<()> { |---|---|---| | A | `c9e72fd` | handlers 拆为 `{mod, statics, commands, urls, inline, callback}` | | B | `50206a9` | `media_sender.rs`：`trait MediaSender` + `impl for Bot`（`::` 消歧）；send.rs 8 处签名改 `&dyn MediaSender`；`MockSender` 测试覆盖兜底触发与错误分类（+5 测试） | -| B 待办 | — | url_media 的 `AppContext` 注入（sender/store/queue/cache），解锁 url_media 全链路测试 | -| C / D | — | 可选后续 | +| B | `ae69d72` | `AppContext` 注入 `url_media`（sender/store/queue/cache），url_media 全链路测试（缓存命中/失效/成功/不支持 URL，+3 测试） | +| C | rate_limit 提交 | `rate_limit.rs` 令牌桶 + 每频道注册表；`MediaSender for Bot` 的 group/copy/animation 前置 `acquire`（+3 测试） | +| D | — | **已延迟**：待下次数据库 schema 变化时实施（见 §5） | -每阶段独立合入；A、B 为核心，C、D 可选。 +A、B、C 为核心并已实施；D 在 schema 变更时落地。