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()
+        );
+    }
+}