From 024dfd50b39d22861913966f6f58bbb7663e6348 Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Mon, 21 Sep 2026 01:03:15 +0800 Subject: [PATCH] fix: bound the inline state map, test the 300s sweep, add a bot-wide send budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three gaps the last audit list named, all in the "resource growth, background timers and limits nobody watches" class. **Idle inline-query entries are pruned.** `DebounceStates` had no eviction at all: one entry per user who ever used inline mode, forever, while the rate limiter's buckets and the chat store both prune in the 300s sweep. Entries now carry a `last_seen` stamp and `prune_idle_states()` drops the ones idle past 300s — the window Telegram caches an inline answer for (`cache_time(300)`), after which a repeat reaches the bot again and has to be answered fresh, so the entry would only suppress a fetch the user is waiting for. The boundary is tested through `prune_idle_at(now, idle_for)` so it does not depend on ageing a monotonic clock. **The 300s sweep is a function, and tested.** It was an inline `tokio::spawn` block: the expiry edit (the only part that talks to Telegram) had no test at all. It is now `periodic_sweep(sender, chat_store, link_cache, task_queue, config, stop)`, which also prunes the inline entries, driven in a test with `start_paused` — the loop's own timer fires the tick, exactly one expired prompt is rewritten in place, a live one keeps its record and buttons. The interval is pinned as a constant because no assertion on the edits can see it (a shorter one produces the same single edit; the paused clock can jump past the boundary while a tick's DB work is in flight). To make the edit reachable at all, `edit_message_text` joined the `MediaSender` trait (Bot impl + mock recording), which is also what keeps `main.rs`'s remaining `Bot` calls unambiguous. `main.rs` leaves the "untested modules" list except for startup/shutdown and the dispatcher tree. **The bot-wide send budget exists.** Telegram throttles a bot in total (~30 msg/s) as well as per chat; only the per-chat bucket existed, so a batch forward fanned out over many chats was unguarded and earned 429s the queue then retried. `acquire_global` charges the same spend against a single shared bucket at the three paced sites (`send_media_group`, `send_animation`, `copy_messages`). The unpaced ones (`send_message`, the edits, the toasts) stay unpaced on purpose: they are one call per action, far below the ceiling, and pacing a user-visible reply would delay it. Not covered: that the send paths call it (they need a real `Bot`), which is the same structural gap as the dispatcher tree. Also: the startup token-exchange decision is now `startup_validation(result)` instead of living inside the `Site::validate` future, so "a 5xx while the container comes up must not disable pixiv" is asserted as a decision — the message the admin gets plus `enabled()` unchanged. The rejected-credential half is deliberately not exercised: it calls `disable()`, a process-wide flag with no reset, and a test touching it would order-couple every other pixiv test. Verified: `cargo fmt`, `cargo clippy --workspace --all-targets --locked -- -D warnings`, `cargo test --workspace --locked` (184 passed, 14 ignored) — plus mutations, each confirmed to fail the relevant test: the sweep not being driven on its timer, the interval shortened to 60s, and (earlier) the queue sweep's missing wake-up. Dropped an empty leftover `crates/x-media/tests/` directory while there (never tracked by git). --- AGENTS.md | 8 +- crates/x-media/src/site/pixiv/interface.rs | 59 +++-- crates/xmedia-bot/src/ctx.rs | 4 + crates/xmedia-bot/src/handlers/inline.rs | 55 ++++- crates/xmedia-bot/src/handlers/mod.rs | 1 + crates/xmedia-bot/src/main.rs | 239 ++++++++++++++++----- crates/xmedia-bot/src/media_sender.rs | 55 +++++ crates/xmedia-bot/src/rate_limit.rs | 47 +++- 8 files changed, 384 insertions(+), 84 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 0d00236..fd249bd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -38,7 +38,7 @@ The `x-media` library: `site::fetch(url)` dispatches through the `SITES` registr |---|---| | `crates/x-media/src/` | Fetch library. `site/mod.rs` = dispatcher + `Fetched`/`FetchError`/`download_media`/`media_size`; `media.rs` = `Media` enum; `examples/fetch.rs` = end-to-end usage sample | | `crates/x-media/src/site//` | One directory per site: `mod.rs` (re-exports), `interface.rs` (PATTERN, `enabled()`, `fetch_from_url()`, `cache_key`/`is_retryable`/`media_headers`, unit struct `Site` implementing `site::Site`, `From for Fetched`), `model.rs` (serde DTOs). Pixiv adds `api.rs` (auth + transport); twitter adds `auth.rs` (logged-in GraphQL `TweetDetail` fallback for NSFW tweets, gated on `TWITTER_AUTH_TOKEN`; without the token a withheld tweet stays `FetchError::Sensitive` and the bot reports it as age-restricted instead of "no media"). Misskey targets misskey.io only (`POST /api/notes/show`, 400+`NO_SUCH_NOTE` → NotFound). Bilibili fetches dynamics (images/animated images only — an attached video degrades to its cover, and its title stands in for the post text, which AV dynamics do not have) from `/x/polymer/web-dynamic/v1/detail` sent with `features=itemOpusStyle` (without that flag the legacy serialization drops an image/text post's body and headline entirely — `desc` comes back `null`; the adapter still parses the legacy `major.draw`/`desc`/`archive` shapes as a fallback). No WBI signature is involved; device cookies `buvid3`/`buvid4` are fetched automatically from `/x/frontend/finger/spi` because bilibili's `-352` risk control starts rejecting plain requests, `BILIBILI_COOKIE` is the escalation when an IP stays blocked; `b23.tv` short links are deliberately unmatched. Twitter's `from_syndication_json` HTML-decodes the API text — syndication and GraphQL `full_text` both arrive pre-escaped (`>` `<` `&` `'`) — so the stored text is raw and the caption escap… -| `crates/xmedia-bot/src/main.rs` | Entry point: env/log init, command registration (`register_commands` — `setMyCommands` plus the profile description texts), shared `send::BOT` force-init, startup sweep of this project's leftover temp files (`x_media::TEMP_FILE_PREFIX` + an age gate, since a killed process runs no destructors), queue worker start, site login validation (`site::validate_all`), 300 s edit-expiry sweep (expired prompts are rewritten in place to `EDIT_PROMPT_EXPIRED_TEXT` with an empty keyboard — an edit, never a new message, so a background timer cannot wake a chat), dptree handler tree, webhook vs polling dispatch | +| `crates/xmedia-bot/src/main.rs` | Entry point: env/log init, command registration (`register_commands` — `setMyCommands` plus the profile description texts), shared `send::BOT` force-init, startup sweep of this project's leftover temp files (`x_media::TEMP_FILE_PREFIX` + an age gate, since a killed process runs no destructors), queue worker start, site login validation (`site::validate_all`), `periodic_sweep` (`SWEEP_INTERVAL` 300 s): expired prompts are rewritten in place to `EDIT_PROMPT_EXPIRED_TEXT` with an empty keyboard — an edit, never a new message, so a background timer cannot wake a chat — plus the link-cache prune, the idle rate-limit buckets and the idle inline-query entries, and the queue backlog line (only when non-empty). Takes its collaborators rather than the statics so its loop is testable with a paused clock, dptree handler tree, webhook vs polling dispatch | | `crates/xmedia-bot/src/config.rs` | Manual env parsing into `Config` | | `crates/xmedia-bot/src/db.rs` | `DbPool`: one shared SQLite connection pool (`POOL_SIZE = 4`, WAL, busy_timeout) for all three tables over `$DATA_DIR/task_queue.db` (default `data/`) — the three stores share it; `open_store` creates file + schema and then applies the `PRAGMA user_version` migration chain (`MIGRATIONS` + `migrate` — append-only; `schema_init` is the version-0 baseline and must not gain columns an existing database would never receive), `with_conn` runs all rusqlite I/O in `spawn_blocking` | | `crates/xmedia-bot/src/handlers/` | Handler modules: `mod.rs` (message entry point, `reply`, `log_key`, the group-only `GROUP_LINK_HINT` for a supported link posted outside a private chat), `commands.rs` (teloxide `BotCommands` enum + command executor, incl. `/test ` (send-only) / `/debug ` (parse-only) and the admin-only `/bot_dict` state dump; `/set_format` rejects unknown `{…}` placeholders and resets with `-`), `urls.rs` (URL extraction + bounded job channel (256) drained by `URL_WORKERS = 8` workers (`start_url_workers`) — backpressure instead of unbounded spawns; teloxide's per-chat workers are sequential — batch-forwards need concurrency), `inline.rs`/`callback.rs` (inline queries / edit-before-forward buttons, incl. `skip`; a forward that fails retryably is both queued *and* settles the prompt — the queued row carries the message ids itself, and a prompt left live let a second Confirm copy the same messages twice and let Skip answer "nothing was forwarded" while the row still delivered), `statics.rs` (global statics) | @@ -47,8 +47,8 @@ The `x-media` library: `site::fetch(url)` dispatches through the `SITES` registr | `crates/xmedia-bot/src/queue.rs` | `PersistentTaskQueue`: SQLite-backed queue (`tasks` table), `QUEUE_WORKERS = 4` concurrent workers (lease via `BEGIN IMMEDIATE` + `locked_until` TTL), retry→dead-letter, a `lease_token` fence: `lease_next` stamps a random token and every write-back (heartbeat, `delete`, `reschedule`, `mark_done`) is guarded by it, so a lease that expired and was re-leased cannot be written by its former holder — a lost lease stops the attempt instead; a finished row's `DELETE`/reschedule retried and a failed delete falling back to a `done` tombstone (the lease query and the sweep only look at `pending`/`in_progress`, so a task that already ran cannot be resurrected and re-run), `notify_one` worker wakeup plus a separate `Notify` for the 30 s lease-expiry sweep (a shared one let the sweep steal the workers' wakeup permit; the sweep does notify the workers after it actually recovered a row, since a recovered task is due immediately while every worker may be parked on `notify` with no pending row to sleep on), `busy_timeout` on all connections | | `crates/xmedia-bot/src/ctx.rs` | `AppContext`: the injected collaborators (`sender` + `ChatStore`/`PersistentTaskQueue`/`LinkCache`/`Config`), `from_statics` for production and the `CONTEXT` static the worker closures hold. `test_support::TestStores` backs handler tests with a tempdir store set, and the module also carries the fixtures those tests share — the canonical cached post (`cached_photo`), the edit-before-forward prompt (`seed_prompt` with its `PROMPT_ID`/`FORWARDED_ID`) and a scripted API error (`api_error`) — so no two test modules keep their own copies | | `crates/xmedia-bot/src/send/` | `send/mod.rs`: `Task`/`MediaItemPayload` payloads, `SendError`/`Classification`, `send_media_sequence`/`send_animation`/`forward_messages`; `send/input_media.rs`: payload → `InputFile`/`InputMedia` + `build_media_group` (caption on the first item only); `send/upload.rs`: the download-and-reupload fallback (`prepare_upload_item`/`send_batch_via_upload`, photo downscale handoff); `send/post_send.rs`: link-cache write, `KEEP_ALIVE` registry, `settle_task`, `post_send_actions`, `handle_task`/`dead_letter_notify` | -| `crates/xmedia-bot/src/media_sender.rs` | `MediaSender` trait: the user-flow surface (`send_media_group`/`send_animation`/`copy_messages`/`send_message`/`answer_callback_query`/`edit_message_caption`/`delete_message`/`send_chat_action`) implemented by teloxide `Bot` (per-chat rate-limited) and by a recording `MockSender` in tests. Admin/setup APIs (`get_chat`, `set_my_commands`, …) stay on the concrete `Bot` | -| `crates/xmedia-bot/src/rate_limit.rs` | Per-chat token bucket (`CAPACITY = 20`, ~20 msg/min refill) paced before sends reach the API so batch forwards don't trip flood control | +| `crates/xmedia-bot/src/media_sender.rs` | `MediaSender` trait: the user-flow surface (`send_media_group`/`send_animation`/`copy_messages`/`send_message`/`answer_callback_query`/`edit_message_text`/`edit_message_caption`/`delete_message`/`send_chat_action`) implemented by teloxide `Bot` (per-chat rate-limited) and by a recording `MockSender` in tests. Admin/setup APIs (`get_chat`, `set_my_commands`, …) stay on the concrete `Bot` | +| `crates/xmedia-bot/src/rate_limit.rs` | Two token buckets paced before sends reach the API so batch forwards don't trip flood control: one per chat (`CAPACITY = 20`, ~20 msg/min refill) and one bot-wide (`acquire_global`, 30/s — Telegram's per-bot ceiling, invisible to any per-chat bucket and only binding when a batch fans out over many chats). `prune_idle` drops the per-chat buckets that refilled while unheld | ## Development Commands @@ -110,5 +110,5 @@ Docker: `docker build -t tgxmb .` then `docker run --rm -d --name tgxmb --env-fi - Live-network tests exist in `site/twitter/interface.rs` (5), `site/bsky/interface.rs` (1), `site/misskey/interface.rs` (1), `site/bilibili/interface.rs` (5), `site/pixiv/api.rs` (1); `photo.rs` adds one `#[ignore = "heavy: …"]` test. `site/mod.rs` also has a **token-gated but not `#[ignore]`d** pixiv download test (`download_media_pixiv_original_with_referer`): it hits `i.pximg.net` whenever `PIXIV_REFRESH_TOKEN` is set, so a local `cargo test --workspace` is not fully offline and can flake on a pixiv CDN body timeout. `disabled_site_is_reported_not_ignored` (same file) is gated the other way round: it asserts `fetch` answers `FetchError::Disabled { site: "pixiv" }` for a pixiv link and early-returns when `PIXIV_REFRESH_TOKEN` **is** set (the site is then enabled). Test gating convention (enforced by `.github/workflows/ci.yml`): pure unit tests always run; live-network tests carry `#[ignore = "live network: ..."]` (run via `cargo test --workspace -- --ignored live`); token-gated pixiv tests early-return when `PIXIV_REFRESH_TOKEN` is absent **or empty** (an unset GitHub secret arrives as `""` — `is_err()` alone would run them tokenless and fail), and the bilibili live tests early-return when the API answers risk control (`-352`, which bilibili applies per IP by request volume). Run the full offline suite with `cargo test --workspace`. - Fixtures are inline `serde_json::json!` builder fns (`fixture()`, `thread_json()`, `illust_json()`), not files. The shared `CLIENT` sets `pool_max_idle_per_host(0)` under `#[cfg(test)]` to avoid cross-runtime `DispatchGone`. - **CI** — `.github/workflows/ci.yml` (actions pinned to commit SHAs, `--locked` on every cargo invocation, `concurrency` cancels superseded runs, `RUST_BACKTRACE=1`) runs `cargo fmt --check` + `cargo clippy --workspace --all-targets --locked -- -D warnings` + `cargo test --workspace --locked` + a release-profile `cargo build --release --locked` + an `actions-rust-lang/audit` dependency-vulnerability gate (offline, no secrets, on every push/PR) and a `live` job (schedule/manual/tag only, `-p x-media` since every network/secret-gated test lives there, `continue-on-error`) for the `#[ignore]`d live + token tests. `.github/workflows/docker.yml` builds and pushes the image on master/tag and runs a **build-only check on pull requests touching the build inputs** (`Dockerfile`, entrypoint, manifests, `.dockerignore`); a release tag must match both crate versions or the build stops, and `FFMPEG_URL`/`FFMPEG_SHA256` are taken from repository variables when set (a release can pin an exact ffmpeg build). `.github/dependabot.yml` keeps crates, the pinned actions and the Docker base images current. -- Untested and hard to test without a mock seam: `main.rs`, `config.rs`, `db.rs`, `handlers/statics.rs`, `media_sender.rs` (holds the `MockSender` itself); in `x-media`: `media.rs`, `lib.rs`, all `model.rs`. The `commands.rs` *executor* needs a real `Bot` (only its pure report builder is tested). Everything else — `handlers/{mod,callback,inline,urls}.rs`, `send/*`, `ctx.rs`, `state.rs`, `queue.rs`, `link_cache.rs`, `rate_limit.rs` — is driven through `TestStores`/`ctx::test_support` and the scripted `MockSender`. +- Untested and hard to test without a mock seam: `config.rs`, `db.rs`, `handlers/statics.rs`, `media_sender.rs` (holds the `MockSender` itself, so only the trait's mock side is exercised); `main.rs` is covered where it was split out (`periodic_sweep`, `sweep_temp_dir`) but not for startup/shutdown or the dispatcher tree; in `x-media`: `media.rs`, `lib.rs`, all `model.rs`. The `commands.rs` *executor* needs a real `Bot` (only its pure report builder is tested). Everything else — `handlers/{mod,callback,inline,urls}.rs`, `send/*`, `ctx.rs`, `state.rs`, `queue.rs`, `link_cache.rs`, `rate_limit.rs` — is driven through `TestStores`/`ctx::test_support` and the scripted `MockSender`. - No coverage tracking. diff --git a/crates/x-media/src/site/pixiv/interface.rs b/crates/x-media/src/site/pixiv/interface.rs index 9896ddc..98aba91 100644 --- a/crates/x-media/src/site/pixiv/interface.rs +++ b/crates/x-media/src/site/pixiv/interface.rs @@ -46,23 +46,25 @@ impl Site for PixivSite { } fn validate(&self) -> SiteFuture<'static, (), String> { - Box::pin(async { - match super::api::validate().await { - Ok(()) => Ok(()), - Err(e) => { - // A rejected credential disables pixiv for the rest of - // this process (it will not fix itself). A bad *moment* — - // a 5xx or a network error while the container comes up — - // must not: disabling on any error turned every later - // pixiv link into "pixiv support is disabled". - if pixiv_error_is_retryable(&e) { - return Err(format!("{e} (transient — pixiv stays enabled)")); - } - super::api::disable(); - Err(format!("{e}")) - } - } - }) + Box::pin(async { startup_validation(super::api::validate().await) }) + } +} + +/// Turns the startup token exchange's outcome into what the bot reports, and +/// disables pixiv only for a rejected credential. A bad *moment* — a 5xx or a +/// network error while the container comes up — must not disable it: disabling +/// on any error turned every later pixiv link into "support is disabled". +/// Separate from the network call so the decision is testable. +fn startup_validation(result: Result<(), PixivError>) -> Result<(), String> { + match result { + Ok(()) => Ok(()), + Err(e) if pixiv_error_is_retryable(&e) => { + Err(format!("{e} (transient — pixiv stays enabled)")) + } + Err(e) => { + super::api::disable(); + Err(format!("{e}")) + } } } @@ -422,6 +424,29 @@ mod tests { } } + #[test] + fn startup_validation_keeps_the_site_enabled_on_a_bad_moment() { + use super::super::api; + + // The startup decision, not the retry policy: a 5xx/429 while the + // container comes up must leave pixiv enabled and say so in the message + // the admin gets. The rejected-credential half is not exercised here — + // it calls `disable()`, a process-wide flag with no reset, so a test + // touching it would order-couple every other pixiv test (the predicate + // it keys on is covered by the table below). + for err in [PixivError::Status(429), PixivError::Status(503)] { + let enabled_before = api::enabled(); + let message = startup_validation(Err(err)).unwrap_err(); + assert!(message.contains("stays enabled"), "{message}"); + assert_eq!( + api::enabled(), + enabled_before, + "a bad moment must not disable the site" + ); + } + assert!(startup_validation(Ok(())).is_ok()); + } + #[test] fn is_retryable_classifies_transient_and_permanent() { // Transient: network errors, explicit transient, pixiv 429/5xx. diff --git a/crates/xmedia-bot/src/ctx.rs b/crates/xmedia-bot/src/ctx.rs index d06bd9d..e1e57d7 100644 --- a/crates/xmedia-bot/src/ctx.rs +++ b/crates/xmedia-bot/src/ctx.rs @@ -157,6 +157,10 @@ pub(crate) mod test_support { &self.link_cache } + pub(crate) fn task_queue(&self) -> &PersistentTaskQueue { + &self.task_queue + } + /// Rows persisted in the task queue: what "queued for retry" looks like /// from the outside. pub(crate) async fn queued_tasks(&self) -> i64 { diff --git a/crates/xmedia-bot/src/handlers/inline.rs b/crates/xmedia-bot/src/handlers/inline.rs index 268d37f..a705f80 100644 --- a/crates/xmedia-bot/src/handlers/inline.rs +++ b/crates/xmedia-bot/src/handlers/inline.rs @@ -20,6 +20,12 @@ use x_media::media::Media; /// post id. Only answer once the query has been stable for this long. const INLINE_DEBOUNCE: std::time::Duration = std::time::Duration::from_millis(800); +/// How long a debounce entry is worth keeping: the window Telegram caches an +/// inline answer for (`answer_inline_query` asks for `cache_time(300)`). Past +/// it a repeat is sent to the bot again and has to be answered fresh, so the +/// entry would only suppress a fetch the user is waiting for. +const INLINE_STATE_TTL: std::time::Duration = std::time::Duration::from_secs(300); + /// Last seen inline query per user and whether it was already answered. /// Guards the debounce timer: a repeat of an answered query is served by /// Telegram's inline cache (see `cache_time`), not by another fetch. Keyed by @@ -28,6 +34,10 @@ const INLINE_DEBOUNCE: std::time::Duration = std::time::Duration::from_millis(80 struct InlineDebounceState { query: String, answered: bool, + /// When a query last touched this entry, so the periodic sweep can drop + /// one per user who ever used inline mode (the map had no eviction at all, + /// unlike the rate limiter's buckets and the chat store). + last_seen: std::time::Instant, } #[derive(Default)] @@ -49,11 +59,21 @@ impl DebounceStates { InlineDebounceState { query: query.to_string(), answered: false, + last_seen: std::time::Instant::now(), }, ); true } + /// Drops entries no query has touched for `idle_for`. Split from the clock + /// so the boundary is testable without ageing a monotonic instant. + fn prune_idle_at(&mut self, now: std::time::Instant, idle_for: std::time::Duration) -> usize { + let before = self.0.len(); + self.0 + .retain(|_, state| now.saturating_duration_since(state.last_seen) < idle_for); + before - self.0.len() + } + /// Claims the answer for the user's newest query; false when a newer query /// superseded it or the answer was already claimed. fn claim(&mut self, user_id: u64, query: &str) -> bool { @@ -64,6 +84,7 @@ impl DebounceStates { return false; } state.answered = true; + state.last_seen = std::time::Instant::now(); true } @@ -73,10 +94,19 @@ impl DebounceStates { && state.query == query { state.answered = false; + state.last_seen = std::time::Instant::now(); } } } +/// Drops debounce entries idle for [`INLINE_STATE_TTL`]; the 300 s sweep calls +/// this next to the rate limiter's prune. Returns how many were dropped. +pub(crate) fn prune_idle_states() -> usize { + INLINE_DEBOUNCE_STATE + .lock() + .prune_idle_at(std::time::Instant::now(), INLINE_STATE_TTL) +} + static INLINE_DEBOUNCE_STATE: LazyLock> = LazyLock::new(|| parking_lot::Mutex::new(DebounceStates::default())); @@ -212,7 +242,7 @@ async fn answer_inline_query(bot: Bot, query: InlineQuery) -> Result break, - _ = tokio::time::sleep(std::time::Duration::from_secs(300)) => {} - } - let ttl = CONFIG.edit_message_ttl; - let removed = CHAT_STORE.prune_expired(ttl).await; - let pruned = LINK_CACHE.prune(CONFIG.link_cache_ttl).await; - if pruned > 0 { - log::info!("link cache: pruned {pruned} expired entr(ies)"); - } - let idle_limiters = crate::rate_limit::prune_idle(); - if idle_limiters > 0 { - log::debug!("rate limiter: dropped {idle_limiters} idle bucket(s)"); - } - // Only speaks up when the queue is not empty: a healthy bot - // has nothing to report, and a periodic "0 pending" line is - // noise that hides the lines that matter. - if let Some((pending, oldest_run_after)) = TASK_QUEUE.pending_backlog().await { - let overdue = crate::db::now_f64() - oldest_run_after; - if overdue >= 0.0 { - log::info!( - "queue: {pending} pending task(s), oldest {overdue:.0}s overdue" - ); - } else { - log::info!( - "queue: {pending} pending task(s), oldest retry in {:.0}s", - -overdue - ); - } - } - for (chat_id, prompt_message_id) in removed { - // Rewritten in place, not announced: the sweep is a - // background timer, and a fresh message would wake the chat - // up to a full TTL later about a prompt the user already - // walked away from. The edit drops the buttons too. If the - // prompt was already deleted this fails with a 400 - // "message to edit not found" — log and ignore. - if let Err(e) = bot - .edit_message_text( - ChatId(chat_id), - MessageId(prompt_message_id as i32), - send::EDIT_PROMPT_EXPIRED_TEXT, - ) - .reply_markup(InlineKeyboardMarkup::default()) - .await - { - log::info!("edit-expiry sweep: prompt message gone: {e}"); - } - } - } + periodic_sweep( + &bot, + &CHAT_STORE, + &LINK_CACHE, + &TASK_QUEUE, + &CONFIG, + stop_rx, + ) + .await; }); } @@ -324,6 +282,79 @@ async fn main() { } } +/// How often [`periodic_sweep`] runs. +const SWEEP_INTERVAL: Duration = Duration::from_secs(300); + +/// The background sweep: rewrites the expired edit prompts in place, prunes the +/// link cache, the idle rate-limit buckets and the idle inline-query entries, +/// and reports the queue only when it is not empty. +/// +/// Takes its collaborators instead of reaching for the statics so a test can +/// drive a tick with a paused clock: a sleeping task nothing drives is how the +/// queue's own sweep kept a missing worker wake-up. +async fn periodic_sweep( + sender: &dyn crate::media_sender::MediaSender, + chat_store: &crate::state::ChatStore, + link_cache: &crate::link_cache::LinkCache, + task_queue: &crate::queue::PersistentTaskQueue, + config: &crate::config::Config, + mut stop: watch::Receiver, +) { + loop { + tokio::select! { + _ = stop.changed() => break, + _ = tokio::time::sleep(SWEEP_INTERVAL) => {} + } + let removed = chat_store.prune_expired(config.edit_message_ttl).await; + let pruned = link_cache.prune(config.link_cache_ttl).await; + if pruned > 0 { + log::info!("link cache: pruned {pruned} expired entr(ies)"); + } + let idle_limiters = crate::rate_limit::prune_idle(); + if idle_limiters > 0 { + log::debug!("rate limiter: dropped {idle_limiters} idle bucket(s)"); + } + // Entries past Telegram's own inline cache window: a repeat is sent to + // the bot again anyway, so keeping them would suppress a fetch the user + // is waiting for (and the map grew one entry per user, forever). + let idle_inline = handlers::prune_idle_states(); + if idle_inline > 0 { + log::debug!("inline queries: dropped {idle_inline} idle entry(ies)"); + } + // Only speaks up when the queue is not empty: a healthy bot has nothing + // to report, and a periodic "0 pending" line is noise that hides the + // lines that matter. + if let Some((pending, oldest_run_after)) = task_queue.pending_backlog().await { + let overdue = crate::db::now_f64() - oldest_run_after; + if overdue >= 0.0 { + log::info!("queue: {pending} pending task(s), oldest {overdue:.0}s overdue"); + } else { + log::info!( + "queue: {pending} pending task(s), oldest retry in {:.0}s", + -overdue + ); + } + } + for (chat_id, prompt_message_id) in removed { + // Rewritten in place, not announced: the sweep is a background + // timer, and a fresh message would wake the chat up to a full TTL + // later about a prompt the user already walked away from. The edit + // drops the buttons too. If the prompt was already deleted this + // fails with a 400 "message to edit not found" — log and ignore. + if let Err(e) = sender + .edit_message_text( + ChatId(chat_id), + MessageId(prompt_message_id as i32), + send::EDIT_PROMPT_EXPIRED_TEXT.to_string(), + ) + .await + { + log::info!("edit-expiry sweep: prompt message gone: {e}"); + } + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -372,4 +403,102 @@ mod tests { assert!(!ours_fresh.exists(), "no age gate: ours, however fresh"); assert!(theirs.exists()); } + + /// The sweep's tick: an expired prompt is rewritten in place (buttons + /// dropped) while a live one is left alone. Driven through the loop's own + /// timer on a paused clock — the loop is what a hand-called helper would + /// leave untested, which is how the queue's sweep kept a missing wake-up. + #[tokio::test(start_paused = true)] + async fn the_sweep_expires_only_the_prompts_past_their_ttl() { + use crate::ctx::test_support::{ + FORWARDED_ID, PROMPT_ID, TestStores, api_error, seed_prompt, + }; + use crate::media_sender::test_support::MockSender; + use crate::state::EditMessage; + + // The interval is pinned here because no assertion on the edits can see + // it: a shorter interval produces the same single edit (the record is + // gone after the first tick), and the paused clock can jump past the + // boundary while a tick's DB work is in flight. + assert_eq!(SWEEP_INTERVAL, Duration::from_secs(300)); + + let config = crate::config::Config::load(); + let stores = TestStores::new(); + let sender = MockSender::scripted(vec![], || { + api_error("Bad Request: message to edit not found") + }); + let ctx = stores.ctx(&sender); + // Chat 1 holds a prompt past its ttl; chat 2 a live one. + let stale = crate::db::unix_now() - config.edit_message_ttl.as_secs() as i64 - 1; + seed_prompt(&ctx, "", stale).await; + stores + .chat_store() + .update(2, |data| { + data.edit_message.insert( + PROMPT_ID, + EditMessage { + url: "https://x.com/u/status/1".into(), + chat_id: 2, + forward_message_ids: vec![FORWARDED_ID], + template: String::new(), + created_at: crate::db::unix_now(), + }, + ); + }) + .await; + + let (stop_tx, stop_rx) = watch::channel(false); + let sweep = periodic_sweep( + &sender, + stores.chat_store(), + stores.link_cache(), + stores.task_queue(), + &config, + stop_rx, + ); + tokio::pin!(sweep); + + // One second short of the interval: nothing has been touched. The + // select is what polls the loop (a pinned future nobody awaits never + // runs), and the paused clock makes this the loop's own timer. + tokio::select! { + _ = &mut sweep => unreachable!("the sweep only returns on stop"), + _ = tokio::time::sleep(SWEEP_INTERVAL - Duration::from_secs(1)) => {} + } + assert!( + sender.edited_texts().is_empty(), + "the sweep ran before its interval" + ); + + // The second that crosses the interval: the tick fires. + tokio::select! { + _ = &mut sweep => unreachable!("the sweep only returns on stop"), + _ = tokio::time::sleep(Duration::from_secs(2)) => {} + } + + assert_eq!( + sender.edited_texts(), + vec![(1, PROMPT_ID, send::EDIT_PROMPT_EXPIRED_TEXT.to_string())], + "exactly the expired prompt, rewritten in place" + ); + assert!( + !ctx.chat_store + .get(1) + .await + .edit_message + .contains_key(&PROMPT_ID), + "the expired record is gone" + ); + assert!( + ctx.chat_store + .get(2) + .await + .edit_message + .contains_key(&PROMPT_ID), + "a live prompt keeps its record and its buttons" + ); + + stop_tx.send(true).unwrap(); + sweep.await; + } } diff --git a/crates/xmedia-bot/src/media_sender.rs b/crates/xmedia-bot/src/media_sender.rs index 2d45bac..25bd65a 100644 --- a/crates/xmedia-bot/src/media_sender.rs +++ b/crates/xmedia-bot/src/media_sender.rs @@ -68,6 +68,16 @@ pub trait MediaSender: Send + Sync { text: Option, ) -> BoxFuture<'_, Result<(), RequestError>>; + /// Rewrites a message's text and drops its inline keyboard: the + /// edit-expiry sweep rewriting a prompt whose record expired (a button left + /// behind could only answer "Expired"). + fn edit_message_text( + &self, + chat_id: ChatId, + message_id: MessageId, + text: String, + ) -> BoxFuture<'_, Result<(), RequestError>>; + /// Rewrites a message's caption, always with HTML parse mode (every caller /// in this bot renders escaped HTML: templates and edit-before-forward /// links). @@ -106,6 +116,9 @@ impl MediaSender for Bot { crate::rate_limit::limiter_for(chat_id.0) .acquire(items.len() as f64) .await; + // Same spend against the bot-wide budget: a fan-out over chats is + // invisible to the per-chat buckets. + crate::rate_limit::acquire_global(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) @@ -124,6 +137,7 @@ impl MediaSender for Bot { ) -> BoxFuture<'a, Result> { Box::pin(async move { crate::rate_limit::limiter_for(chat_id.0).acquire(1.0).await; + crate::rate_limit::acquire_global(1.0).await; let mut request = ::send_animation(self, chat_id, file) .caption(caption) .parse_mode(ParseMode::Html) @@ -147,6 +161,7 @@ impl MediaSender for Bot { crate::rate_limit::limiter_for(to.0) .acquire(ids.len() as f64) .await; + crate::rate_limit::acquire_global(ids.len() as f64).await; ::copy_messages(self, to, from, ids).await }) } @@ -185,6 +200,20 @@ impl MediaSender for Bot { }) } + fn edit_message_text( + &self, + chat_id: ChatId, + message_id: MessageId, + text: String, + ) -> BoxFuture<'_, Result<(), RequestError>> { + Box::pin(async move { + ::edit_message_text(self, chat_id, message_id, text) + .reply_markup(InlineKeyboardMarkup::default()) + .await + .map(|_| ()) + }) + } + fn edit_message_caption( &self, chat_id: ChatId, @@ -260,6 +289,8 @@ pub(crate) mod test_support { messages: Mutex>, captions: Mutex>, answers: Mutex>>, + /// `(chat, message, text)` of every text rewrite, in order. + edited_texts: Mutex>, /// Builds the error every `*Err` outcome returns (RequestError is not /// cloneable, so the factory recreates it per call). error: Box RequestError + Send + Sync>, @@ -280,6 +311,7 @@ pub(crate) mod test_support { messages: Mutex::new(Vec::new()), captions: Mutex::new(Vec::new()), answers: Mutex::new(Vec::new()), + edited_texts: Mutex::new(Vec::new()), error: Box::new(error), } } @@ -305,6 +337,11 @@ pub(crate) mod test_support { self.answers.lock().clone() } + /// `(chat, message, text)` of every `edit_message_text`, in order. + pub(crate) fn edited_texts(&self) -> Vec<(i64, i64, String)> { + self.edited_texts.lock().clone() + } + fn next(&self, kind: &'static str) -> Outcome { self.calls.lock().push(kind); let script = self.script.lock(); @@ -411,6 +448,24 @@ pub(crate) mod test_support { }) } + fn edit_message_text( + &self, + chat_id: ChatId, + message_id: MessageId, + text: String, + ) -> BoxFuture<'_, Result<(), RequestError>> { + // Always succeeds: the only caller is the expiry sweep, which + // tolerates a failure (a prompt the user already deleted), so the + // script stays free for the call the test is about. + Box::pin(async move { + self.calls.lock().push("edit_message_text"); + self.edited_texts + .lock() + .push((chat_id.0, message_id.0 as i64, text)); + Ok(()) + }) + } + fn edit_message_caption( &self, _chat_id: ChatId, diff --git a/crates/xmedia-bot/src/rate_limit.rs b/crates/xmedia-bot/src/rate_limit.rs index 30adea0..e8188bc 100644 --- a/crates/xmedia-bot/src/rate_limit.rs +++ b/crates/xmedia-bot/src/rate_limit.rs @@ -1,12 +1,13 @@ //! 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.). +//! Telegram throttles bots on two budgets: one per chat (roughly 20 +//! messages/min for channels/groups) and a bot-wide one (~30 messages per +//! second). Both are smoothed here *before* the burst reaches the API — the +//! per-chat bucket charges one token per message, and [`acquire_global`] +//! charges the same spend against the bot-wide budget, which no per-chat +//! bucket can see (a forward fanned out over many chats spends one token in +//! each and nothing anywhere). The queue retry stays as the safety net for +//! whatever neither bucket models. use parking_lot::Mutex; use std::collections::HashMap; @@ -19,6 +20,12 @@ const CAPACITY: f64 = 20.0; /// Sustained refill: ~20 messages per minute. const REFILL_PER_SEC: f64 = 20.0 / 60.0; +/// The bot-wide budget: Telegram allows roughly 30 messages per second for a +/// bot in total, independently of the per-chat limits. Set to the documented +/// ceiling, so it only ever binds on a cross-chat burst. +const GLOBAL_CAPACITY: f64 = 30.0; +const GLOBAL_REFILL_PER_SEC: f64 = 30.0; + struct State { /// Current token balance; may go negative (debt from an acquire larger /// than the capacity, repaid by subsequent refills). @@ -107,6 +114,17 @@ pub fn limiter_for(chat_id: i64) -> Arc { .clone() } +/// The one bucket every chat shares: Telegram's bot-wide budget. +static GLOBAL_LIMITER: LazyLock = + LazyLock::new(|| TokenBucket::new(GLOBAL_CAPACITY, GLOBAL_REFILL_PER_SEC)); + +/// Waits for `n` messages' worth of the bot-wide budget. Called by the send +/// paths next to their per-chat [`limiter_for`]: at ~30/s it does not bind on +/// a single chat, but a batch fanned out over many chats has no other guard. +pub async fn acquire_global(n: f64) { + GLOBAL_LIMITER.acquire(n).await; +} + /// Drops limiters that are idle (refilled to capacity, so the chat has not /// sent recently) and are not still held by an in-flight sender. The map /// would otherwise keep one bucket per chat that ever sent media, forever. @@ -160,6 +178,21 @@ mod tests { ); } + #[tokio::test(start_paused = true)] + async fn the_global_budget_is_paced_and_shared() { + // Drain the process-wide budget (no other test touches it: the send + // paths that use it are mocked), then prove the next message waits for + // the refill instead of going out instantly. + acquire_global(GLOBAL_CAPACITY).await; + let start = tokio::time::Instant::now(); + acquire_global(1.0).await; + assert!( + start.elapsed() >= Duration::from_secs_f64(1.0 / GLOBAL_REFILL_PER_SEC), + "a fanned-out burst must be paced: elapsed {:?}", + start.elapsed() + ); + } + #[tokio::test(start_paused = true)] async fn prune_idle_drops_full_unheld_buckets_only() { // Held by this task: kept even at full capacity, a sender has it.