feat(send): add per-chat token bucket rate limiting

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"] }

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, \

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;
 

crates/xmedia-bot/src/media_sender.rs

@@ -73,6 +73,11 @@ impl MediaSender for Bot {
         items: Vec<InputMedia>,
     ) -> BoxFuture<'_, Result<Vec<Message>, 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;
             // `<Bot as Requester>::` disambiguates from this trait's same-named
             // method (teloxide's API lives in the `Requester` trait).
             <Bot as Requester>::send_media_group(self, chat_id, items)
@@ -90,6 +95,7 @@ impl MediaSender for Bot {
         file: InputFile,
     ) -> BoxFuture<'a, Result<Message, RequestError>> {
         Box::pin(async move {
+            crate::rate_limit::limiter_for(chat_id.0).acquire(1.0).await;
             let mut request = <Bot as Requester>::send_animation(self, chat_id, file)
                 .caption(caption)
                 .parse_mode(ParseMode::Html)
@@ -107,7 +113,14 @@ impl MediaSender for Bot {
         from: ChatId,
         ids: Vec<MessageId>,
     ) -> BoxFuture<'_, Result<Vec<MessageId>, RequestError>> {
-        Box::pin(async move { <Bot as Requester>::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;
+            <Bot as Requester>::copy_messages(self, to, from, ids).await
+        })
     }
 
     fn send_message(

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<State>,
+}
+
+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<Mutex<HashMap<i64, Arc<TokenBucket>>>> =
+    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<TokenBucket> {
+    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()
+        );
+    }
+}

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<State> }
 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<TokenBucket>; // 每频道一个桶
 ```
 
-- 按频道粒度（`HashMap<ChatId, Arc<TokenBucket>>`），在 `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`（`<Bot as Requester>::` 消歧）；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 变更时落地。