From 01e097a8b692fa9d432900bb8019cf4939a7238a Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Mon, 21 Sep 2026 18:28:41 +0800 Subject: [PATCH] refactor(handlers): split urls.rs into workers, pipeline and startup repair urls.rs carried five reasons to change: the job channel and its worker pool, the single-flight fetch, URL parsing, the per-URL pipeline, and the startup repair of queued retries. The two with their own lifecycle move out: - url_workers.rs: the bounded channel, its supervised pool and start/stop_url_workers (the pipeline stays in urls.rs, which the workers call). - repair.rs: needs_refetch/apply_refresh/refetch/repair_lost_local_media with their tests, moved whole (the live one keeps its #[ignore]). Also folds the two byte-identical render_fields -> CachedPost mappings (urls.rs and repair.rs) into urls::cached_snapshot, and moves the shared permanent_error test fixture into ctx::test_support. --- AGENTS.md | 2 +- crates/xmedia-bot/src/ctx.rs | 7 + crates/xmedia-bot/src/handlers/mod.rs | 9 +- crates/xmedia-bot/src/handlers/repair.rs | 355 +++++++++++++ crates/xmedia-bot/src/handlers/url_workers.rs | 102 ++++ crates/xmedia-bot/src/handlers/urls.rs | 489 +----------------- 6 files changed, 495 insertions(+), 469 deletions(-) create mode 100644 crates/xmedia-bot/src/handlers/repair.rs create mode 100644 crates/xmedia-bot/src/handlers/url_workers.rs diff --git a/AGENTS.md b/AGENTS.md index 74145f8..673637d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -41,7 +41,7 @@ The `x-media` library: `site::fetch(url)` dispatches through the `SITES` registr | `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), startup repair of queued retries whose local media did not survive a restart (`handlers::repair_lost_local_media`, before any worker can lease: those rows are re-fetched from their `source_url`), 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 — `db.rs`'s tests pin a pre-migration database upgrading intact, the shipped migration text frozen (appending is the only allowed change) and a fresh database landing at the latest version), `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; one *shared* in-flight fetch per cache key (`shared_fetch`+`IN_FLIGHT_FETCHES`: a second chat, a batch forward or a retry asking for the same post meanwhile waits for the first caller's result, the entry is dropped the moment the fetch settles so nothing is ever answered from an old fetch, and a waiter whose sharer was cancelled fetches for itself); plus the startup repair `repair_lost_local_media`, whose decision (`needs_refetch`) and rewrite (`apply_refresh`) are pure and tested while the fetch itself is a live test), `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) | +| `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), `url_workers.rs`/`repair.rs` (worker pool; startup repair), `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 + the per-URL pipeline; one *shared* in-flight fetch per cache key (`shared_fetch`+`IN_FLIGHT_FETCHES`: a second chat, a batch forward or a retry asking for the same post meanwhile waits for the first caller's result, the entry is dropped the moment the fetch settles so nothing is ever answered from an old fetch, and a waiter whose sharer was cancelled fetches for itself); plus the startup repair `repair_lost_local_media`, whose decision (`needs_refetch`) and rewrite (`apply_refresh`) are pure and tested while the fetch itself is a live test), `url_workers.rs` (the bounded job channel (256) and its `URL_WORKERS = 8` supervised workers, `start_url_workers`/`stop_url_workers` — backpressure instead of unbounded spawns; teloxide's per-chat workers are sequential, so batch-forwards need this concurrency), `repair.rs` (startup `repair_lost_local_media` with its `needs_refetch`/`apply_refresh`/`refetch`), `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) | | `crates/xmedia-bot/src/state.rs` | `ChatStore`: parking_lot `Mutex` cache + SQLite write-through (`chat_state` table); the 300 s sweep's `prune_expired` evicts any chat with no live edit-before-forward prompt, so the cache (and the per-chat lock map) stays bounded to active prompts — durable settings reload from the DB on next use | | `crates/xmedia-bot/src/link_cache.rs` | `LinkCache`: SQLite-backed cache (`link_cache` table) of successfully sent posts — raw caption fields + the source media URLs + Telegram `file_id`s; repeat links re-send locally (no fetch/upload), TTL + prune; a permanent send failure *degrades* the entry instead of dropping it (the file ids go, the URLs stay, so the next request re-sends from those without a fetch), and a degraded entry that fails again is removed | | `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), `runnable_rows`/`replace_payload` (the startup repair's read/rewrite path: it runs before the workers exist, which is why it needs no lease token), `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 | diff --git a/crates/xmedia-bot/src/ctx.rs b/crates/xmedia-bot/src/ctx.rs index 9449963..603e3db 100644 --- a/crates/xmedia-bot/src/ctx.rs +++ b/crates/xmedia-bot/src/ctx.rs @@ -65,6 +65,13 @@ pub(crate) mod test_support { RequestError::Api(ApiError::Unknown(message.to_string())) } + /// The API error a caption edit that changes nothing answers with — what + /// the mocks script for a permanent send failure. A `fn` pointer, so it can + /// be handed to `MockSender::scripted` as-is. + pub(crate) fn permanent_error() -> RequestError { + api_error("Bad Request: message is not modified") + } + /// One photo payload item: `media` in the two flags the tests vary (no /// smaller variant, since that is the field most tests leave alone — /// `send`'s own tests build that case directly). diff --git a/crates/xmedia-bot/src/handlers/mod.rs b/crates/xmedia-bot/src/handlers/mod.rs index f1b7391..dde1044 100644 --- a/crates/xmedia-bot/src/handlers/mod.rs +++ b/crates/xmedia-bot/src/handlers/mod.rs @@ -9,18 +9,20 @@ mod callback; mod commands; mod inline; +mod repair; mod statics; +mod url_workers; mod urls; pub use callback::callback_query_handler; pub use commands::register_commands; pub use inline::inline_query_handler; pub(crate) use inline::prune_idle_states; +pub(crate) use repair::repair_lost_local_media; /// The resolved `$DATA_DIR/task_queue.db` path, for the startup config line. pub(crate) use statics::db_path; pub use statics::{CHAT_STORE, CONFIG, LINK_CACHE, TASK_QUEUE}; -pub(crate) use urls::repair_lost_local_media; -pub use urls::{start_url_workers, stop_url_workers}; +pub use url_workers::{start_url_workers, stop_url_workers}; use crate::ctx::AppContext; use crate::media_sender::MediaSender; @@ -29,7 +31,8 @@ use teloxide::RequestError; use teloxide::prelude::*; use teloxide::types::{ChatId, Message, MessageId}; use teloxide::utils::command::BotCommands; -use urls::{URL_JOBS, extract_urls}; +use url_workers::URL_JOBS; +use urls::extract_urls; /// Reply to a message by id, keeping the reply decoration even if the /// original was already deleted. diff --git a/crates/xmedia-bot/src/handlers/repair.rs b/crates/xmedia-bot/src/handlers/repair.rs new file mode 100644 index 0000000..5f653df --- /dev/null +++ b/crates/xmedia-bot/src/handlers/repair.rs @@ -0,0 +1,355 @@ +//! Startup repair, run before any queue worker exists: a queued retry whose +//! local media (a ugoira MP4, a bsky remux, a downloaded temp file) did not +//! survive the restart can never succeed, because the registry that kept those +//! files alive (`send::KEEP_ALIVE`) is in memory. Those rows are re-fetched +//! from their post instead of dead-lettering the user's link. + +use super::log_key; +use super::urls::{cached_snapshot, media_to_payload}; +use crate::ctx::AppContext; +use crate::link_cache::CachedPost; +use crate::send::{self, Delivery, MediaItemPayload, Task}; + +// ── Startup repair: queued retries whose local media did not survive ─────── + +/// A post's fresh media plus the caption and cache snapshot that go with them: +/// what [`refetch`] hands [`apply_refresh`]. Plain data, so the rewrite below +/// can be tested without a network fetch (which cannot be faked here: +/// [`x_media::site::Fetched`] keeps a private field and is not constructible +/// outside its crate). +struct Refetched { + caption: String, + items: Vec, + cache_data: Option, +} + +/// Whether a queued task should have its post re-fetched, because it still +/// wants a local file (ugoira MP4, a bsky remux, a downloaded temp file) that is +/// gone. Those files live in the system temp dir and the registry that keeps +/// them alive for the retry (`send::KEEP_ALIVE`) is in memory, so a restart +/// takes all of them — a retry that needs one can only dead-letter. +/// +/// A partially delivered album is left alone: its remaining batches cannot be +/// reconciled with a fresh media list without risking a second copy of what the +/// user already received. +fn needs_refetch(task: &Task) -> bool { + if let Task::SendMediaSequence { + batch_index, + sent_message_ids, + .. + } = task + && (*batch_index > 0 || !sent_message_ids.is_empty()) + { + return false; + } + task.local_media_paths().iter().any(|path| !path.exists()) +} + +/// Rebuilds the task from the fresh media, keeping its delivery envelope (chat, +/// reply, forward/edit settings, notify targets): the retry that was queued must +/// still deliver the same way, whoever asked for it. +fn apply_refresh(task: &Task, fresh: &Refetched) -> Option { + let chat_id = task.chat_id()?; + let (edit_before_forward, forward_channel_id) = match task { + Task::SendMediaSequence { + edit_before_forward, + forward_channel_id, + .. + } + | Task::SendAnimation { + edit_before_forward, + forward_channel_id, + .. + } => (*edit_before_forward, *forward_channel_id), + Task::ForwardMessages { .. } => return None, + }; + let reply_to_message_id = match task { + Task::SendMediaSequence { + reply_to_message_id, + .. + } + | Task::SendAnimation { + reply_to_message_id, + .. + } => *reply_to_message_id, + Task::ForwardMessages { .. } => return None, + }; + let (notify_chat_id, notify_message_id) = task.notify_target(); + Some(Task::from_items( + Delivery { + chat_id, + reply_to_message_id, + edit_before_forward, + forward_channel_id, + notify_chat_id, + notify_message_id, + }, + task.source_url()?.to_string(), + fresh.caption.clone(), + fresh.items.clone(), + fresh.cache_data.clone(), + )) +} + +/// Fetches the post again and maps it into [`Refetched`]: the same mapping the +/// fresh-fetch path uses (per-site caption format from the chat, render fields +/// for the link-cache snapshot), so a repaired task looks like a first send. +async fn refetch( + ctx: &AppContext<'_>, + chat_id: i64, + url: &str, +) -> Result, x_media::site::FetchError> { + let Some(fetched) = x_media::site::fetch(url).await? else { + return Ok(None); + }; + if fetched.media.is_empty() { + return Ok(None); + } + let chat_data = ctx.chat_store.get(chat_id).await; + let format = chat_data.format_for(fetched.site_id); + let caption = fetched.caption_with(&format); + let cache_data = cached_snapshot(&fetched); + let items: Vec = fetched + .media + .iter() + .map(|media| media_to_payload(media, fetched.sensitive)) + .collect(); + // The re-fetch may produce a fresh local file (ugoira / bsky remux): hand it + // to the same keep-alive registry the first fetch uses. + if let Some(dir) = fetched.keep_alive() { + send::KEEP_ALIVE.lock().push(dir); + } + Ok(Some(Refetched { + caption, + items, + cache_data, + })) +} + +/// Re-fetches every queued task whose local media did not survive the restart, +/// so the user's link is still delivered instead of dead-lettering on a file +/// that cannot come back. Returns how many rows were rewritten. +/// +/// Startup only, before the queue workers start: no worker can lease a row while +/// this writes, which is what lets it replace payloads without the lease-token +/// guard every worker write-back carries. +pub(crate) async fn repair_lost_local_media(ctx: &AppContext<'_>) -> usize { + let mut repaired = 0; + for (id, payload) in ctx.task_queue.runnable_rows().await { + let Ok(task) = serde_json::from_str::(&payload) else { + continue; + }; + if !needs_refetch(&task) { + continue; + } + let (Some(url), Some(chat_id)) = (task.source_url().map(str::to_string), task.chat_id()) + else { + continue; + }; + match refetch(ctx, chat_id, &url).await { + Ok(Some(fresh)) => { + let Some(updated) = apply_refresh(&task, &fresh) else { + continue; + }; + let updated = serde_json::to_value(&updated).expect("task serializes"); + if ctx.task_queue.replace_payload(&id, &updated).await { + repaired += 1; + log::info!( + "startup repair: re-fetched [key={}] for chat={chat_id} (its local media did not survive the restart)", + log_key(&url) + ); + } + } + // The post is gone or withheld now: the retry could not have + // delivered anything either, so say why instead of letting it + // dead-letter on a missing file. + Ok(None) | Err(_) => { + let (notify_chat_id, notify_message_id) = task.notify_target(); + log::warn!( + "startup repair: [key={}] for chat={chat_id} needed a re-fetch and none was possible", + log_key(&url) + ); + send::notify_failure( + ctx.sender, + notify_chat_id, + notify_message_id, + &format!( + "{} — the media held for retry was lost when the bot restarted and the post could not be fetched again. Please send the link again.", + log_key(&url) + ), + ) + .await; + } + } + } + repaired +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ctx::test_support::{TestStores, permanent_error, photo_item}; + use crate::media_sender::test_support::MockSender; + + fn queued_task(media: &str, batch_index: usize, sent: Vec) -> Task { + Task::SendMediaSequence { + chat_id: 1, + reply_to_message_id: 2, + caption: "cap".into(), + media_batches: vec![vec![photo_item(media, false, false)]], + batch_index, + sent_message_ids: sent, + source_url: "https://x.com/u/status/1".into(), + edit_before_forward: true, + forward_channel_id: Some(2), + notify_chat_id: Some(1), + notify_message_id: Some(2), + cache_data: None, + } + } + + #[test] + fn only_tasks_missing_a_local_file_need_a_refetch() { + // A URL send needs nothing. + assert!(!needs_refetch(&queued_task("https://cdn/1.jpg", 0, vec![]))); + // A local path that is still there (a survived temp file) needs nothing. + let dir = tempfile::tempdir().unwrap(); + let alive = dir.path().join("ugoira.mp4"); + std::fs::write(&alive, b"x").unwrap(); + assert!(!needs_refetch(&queued_task( + alive.to_str().unwrap(), + 0, + vec![] + ))); + // A local path the restart took away does. + assert!(needs_refetch(&queued_task( + "/nonexistent-ugoira.mp4", + 0, + vec![] + ))); + // A partially delivered album is left to its own retry path. + assert!(!needs_refetch(&queued_task( + "/nonexistent-ugoira.mp4", + 1, + vec![7] + ))); + assert!(!needs_refetch(&queued_task( + "/nonexistent-ugoira.mp4", + 0, + vec![7] + ))); + // A channel copy holds no media. + assert!(!needs_refetch(&Task::ForwardMessages { + from_chat_id: 1, + to_chat_id: 2, + message_ids: vec![3], + notify_chat_id: None, + notify_message_id: None, + })); + } + + #[test] + fn apply_refresh_keeps_the_delivery_envelope() { + let task = queued_task("/nonexistent-ugoira.mp4", 0, vec![]); + let fresh = Refetched { + caption: "fresh caption".into(), + items: vec![photo_item("https://cdn/fresh.jpg", true, false)], + cache_data: None, + }; + match apply_refresh(&task, &fresh).expect("a repairable task") { + Task::SendMediaSequence { + chat_id, + reply_to_message_id, + caption, + media_batches, + batch_index, + sent_message_ids, + source_url, + edit_before_forward, + forward_channel_id, + notify_chat_id, + notify_message_id, + .. + } => { + // Same delivery: chat, reply, forward/edit settings, notify. + assert_eq!((chat_id, reply_to_message_id), (1, 2)); + assert!(edit_before_forward); + assert_eq!(forward_channel_id, Some(2)); + assert_eq!((notify_chat_id, notify_message_id), (Some(1), Some(2))); + assert_eq!(source_url, "https://x.com/u/status/1"); + // Fresh media, and nothing of it counted as sent yet. + assert_eq!(caption, "fresh caption"); + assert!( + matches!( + &media_batches[0][0], + MediaItemPayload::Photo { media, .. } if media == "https://cdn/fresh.jpg" + ), + "fresh media must replace the lost local file" + ); + assert!(matches!( + media_batches[0][0], + MediaItemPayload::Photo { + has_spoiler: true, + .. + } + )); + assert_eq!((batch_index, sent_message_ids.len()), (0, 0)); + } + other => panic!("expected a media sequence, got {other:?}"), + } + } + + /// The whole repair against a real post: a queued row whose media is a local + /// file the restart took away is re-fetched from its `source_url` and + /// rewritten in place, so the retry can still deliver it. + #[tokio::test] + #[ignore = "live network: requires outbound HTTPS to public.api.bsky.app"] + async fn live_repair_refetches_a_lost_local_media_row() { + let stores = TestStores::new(); + // An empty script: the repair must not need to tell the user anything. + let sender = MockSender::scripted(vec![], permanent_error); + let ctx = stores.ctx(&sender); + let mut task = queued_task("/nonexistent-ugoira.mp4", 0, vec![]); + if let Task::SendMediaSequence { source_url, .. } = &mut task { + *source_url = "https://bsky.app/profile/fu-futa.bsky.social/post/3laoveufjv224".into(); + } + stores + .task_queue() + .enqueue(serde_json::to_value(&task).unwrap(), crate::db::now_f64()) + .await + .unwrap(); + + assert_eq!(repair_lost_local_media(&ctx).await, 1); + + let updated: Task = serde_json::from_value(stores.queued_payload().await).unwrap(); + match updated { + Task::SendMediaSequence { + media_batches, + batch_index, + sent_message_ids, + caption, + .. + } => { + let media: Vec = media_batches + .iter() + .flatten() + .map(|item| match item { + MediaItemPayload::Photo { media, .. } + | MediaItemPayload::Video { media, .. } + | MediaItemPayload::Animation { media, .. } => media.clone(), + }) + .collect(); + assert!(!media.is_empty(), "the fresh fetch yielded no media"); + assert!( + media.iter().all(|m| m.starts_with("http")), + "the retry must be uploadable from URLs again: {media:?}" + ); + assert_eq!((batch_index, sent_message_ids.len()), (0, 0)); + assert!(!caption.is_empty()); + } + other => panic!("expected a repaired media sequence, got {other:?}"), + } + // The post was re-read, not re-delivered: nothing was sent. + assert!(sender.calls().is_empty(), "{:?}", sender.calls()); + } +} diff --git a/crates/xmedia-bot/src/handlers/url_workers.rs b/crates/xmedia-bot/src/handlers/url_workers.rs new file mode 100644 index 0000000..4dc34c1 --- /dev/null +++ b/crates/xmedia-bot/src/handlers/url_workers.rs @@ -0,0 +1,102 @@ +//! The URL job channel and its worker pool: a bounded queue (backpressure +//! instead of unbounded spawns) drained by [`URL_WORKERS`] supervised workers. +//! +//! teloxide's per-chat workers are sequential, so a batch forward needs its own +//! concurrency: this is where a link handed over by `handlers::mod` actually +//! reaches the pipeline. + +use super::urls::{PostSend, url_media}; +use crate::ctx::CONTEXT; +use std::sync::LazyLock; +use teloxide::types::Message; + +/// One URL job: the message + the extracted URL (the sender and stores come +/// from the shared [`AppContext`], assembled from statics inside the worker). +type UrlJob = (Message, String); +/// Bounded channel of URL jobs drained by [`start_url_workers`]. The bound +/// caps both queued memory and shutdown backlog; a full channel applies +/// backpressure to the per-chat handler instead of spawning unbounded tasks. +pub(crate) static URL_JOBS: LazyLock< + parking_lot::Mutex>>, +> = LazyLock::new(|| parking_lot::Mutex::new(None)); +/// Set by main's shutdown sequence; workers stop pulling new jobs. +pub(crate) static URL_STOP: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + +/// JoinHandles of the URL workers, awaited by [`stop_url_workers`]. +static URL_WORKER_HANDLES: LazyLock>>>> = + LazyLock::new(|| parking_lot::Mutex::new(None)); + +/// Worker count draining URL jobs; keeps the old 8-permit concurrency cap +/// while bounding how many jobs can be queued at all. +const URL_WORKERS: usize = 8; + +/// Starts the URL job workers (called once from main after the queue starts). +/// teloxide dispatches updates to a per-chat worker that handles them +/// sequentially, so a batch-forward of many messages would otherwise be +/// processed one at a time (fetch + send each, roughly a second per +/// message); the workers add throughput, and FIFO order preserves per-message +/// URL order. +pub async fn start_url_workers() { + let (tx, rx) = tokio::sync::mpsc::channel::(256); + *URL_JOBS.lock() = Some(tx); + let rx = std::sync::Arc::new(tokio::sync::Mutex::new(rx)); + let mut handles = Vec::with_capacity(URL_WORKERS); + for _ in 0..URL_WORKERS { + let rx = std::sync::Arc::clone(&rx); + handles.push(tokio::spawn(async move { + // Supervised like the queue workers: a panic inside a worker + // (a handler, a poisoned lock) used to kill it for good and + // silently shrink the pool — the remaining workers keep the + // channel drained, so nothing else surfaces the loss. The job the + // panicking worker held is lost; the panic is not. + while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) { + let rx = std::sync::Arc::clone(&rx); + if let Err(e) = tokio::spawn(async move { + while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) { + let job = rx.lock().await.recv().await; + match job { + Some((message, url)) => { + url_media( + &CONTEXT, + message.chat.id.0, + message.id.0 as i64, + &url, + PostSend::FromChat, + ) + .await; + } + None => break, + } + } + }) + .await + { + log::error!("url worker panicked, restarting: {e}"); + } + } + })); + } + *URL_WORKER_HANDLES.lock() = Some(handles); +} + +/// Stops the URL workers: sets the stop flag, drops the job channel (so +/// workers blocked in \`recv()\` wake with \`None\` and exit) and awaits the +/// worker tasks. Each worker finishes its in-flight job first; jobs still +/// queued in the channel are abandoned (the old implementation neither +/// drained them nor woke blocked workers — it only set a flag checked +/// between jobs). +pub async fn stop_url_workers() { + URL_STOP.store(true, std::sync::atomic::Ordering::Relaxed); + // Dropping the sender makes every worker's recv() return None. + *URL_JOBS.lock() = None; + // Take the handles first so the lock guard drops before the awaits. + let handles = URL_WORKER_HANDLES.lock().take(); + if let Some(handles) = handles { + for handle in handles { + if let Err(e) = handle.await { + log::error!("url worker panicked at shutdown: {e}"); + } + } + } +} diff --git a/crates/xmedia-bot/src/handlers/urls.rs b/crates/xmedia-bot/src/handlers/urls.rs index d79d947..e9d92ef 100644 --- a/crates/xmedia-bot/src/handlers/urls.rs +++ b/crates/xmedia-bot/src/handlers/urls.rs @@ -2,7 +2,7 @@ //! worker pool, link-cache fast path, fetch, task build and send dispatch. use super::{log_key, reply}; -use crate::ctx::{AppContext, CONTEXT}; +use crate::ctx::AppContext; use crate::link_cache::{CachedMedia, CachedMediaKind, CachedPost}; use crate::media_sender::MediaSender; use crate::send::{self, Delivery, MediaItemPayload, Task}; @@ -14,27 +14,6 @@ use teloxide::RequestError; use teloxide::types::{ChatAction, ChatId, Message, MessageEntityKind, MessageId}; use x_media::media::Media; -/// One URL job: the message + the extracted URL (the sender and stores come -/// from the shared [`AppContext`], assembled from statics inside the worker). -type UrlJob = (Message, String); -/// Bounded channel of URL jobs drained by [`start_url_workers`]. The bound -/// caps both queued memory and shutdown backlog; a full channel applies -/// backpressure to the per-chat handler instead of spawning unbounded tasks. -pub(crate) static URL_JOBS: LazyLock< - parking_lot::Mutex>>, -> = LazyLock::new(|| parking_lot::Mutex::new(None)); -/// Set by main's shutdown sequence; workers stop pulling new jobs. -pub(crate) static URL_STOP: std::sync::atomic::AtomicBool = - std::sync::atomic::AtomicBool::new(false); - -/// JoinHandles of the URL workers, awaited by [`stop_url_workers`]. -static URL_WORKER_HANDLES: LazyLock>>>> = - LazyLock::new(|| parking_lot::Mutex::new(None)); - -/// Worker count draining URL jobs; keeps the old 8-permit concurrency cap -/// while bounding how many jobs can be queued at all. -const URL_WORKERS: usize = 8; - /// One fetch per post at a time, keyed by the normalized cache key. Two chats /// posting the same link at the same moment (or a batch forward and a queued /// retry) used to run two full fetches: two sets of source requests, and for an @@ -120,76 +99,6 @@ impl Drop for InFlightFetch<'_, T> { } } -/// Starts the URL job workers (called once from main after the queue starts). -/// teloxide dispatches updates to a per-chat worker that handles them -/// sequentially, so a batch-forward of many messages would otherwise be -/// processed one at a time (fetch + send each, roughly a second per -/// message); the workers add throughput, and FIFO order preserves per-message -/// URL order. -pub async fn start_url_workers() { - let (tx, rx) = tokio::sync::mpsc::channel::(256); - *URL_JOBS.lock() = Some(tx); - let rx = std::sync::Arc::new(tokio::sync::Mutex::new(rx)); - let mut handles = Vec::with_capacity(URL_WORKERS); - for _ in 0..URL_WORKERS { - let rx = std::sync::Arc::clone(&rx); - handles.push(tokio::spawn(async move { - // Supervised like the queue workers: a panic inside a worker - // (a handler, a poisoned lock) used to kill it for good and - // silently shrink the pool — the remaining workers keep the - // channel drained, so nothing else surfaces the loss. The job the - // panicking worker held is lost; the panic is not. - while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) { - let rx = std::sync::Arc::clone(&rx); - if let Err(e) = tokio::spawn(async move { - while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) { - let job = rx.lock().await.recv().await; - match job { - Some((message, url)) => { - url_media( - &CONTEXT, - message.chat.id.0, - message.id.0 as i64, - &url, - PostSend::FromChat, - ) - .await; - } - None => break, - } - } - }) - .await - { - log::error!("url worker panicked, restarting: {e}"); - } - } - })); - } - *URL_WORKER_HANDLES.lock() = Some(handles); -} - -/// Stops the URL workers: sets the stop flag, drops the job channel (so -/// workers blocked in \`recv()\` wake with \`None\` and exit) and awaits the -/// worker tasks. Each worker finishes its in-flight job first; jobs still -/// queued in the channel are abandoned (the old implementation neither -/// drained them nor woke blocked workers — it only set a flag checked -/// between jobs). -pub async fn stop_url_workers() { - URL_STOP.store(true, std::sync::atomic::Ordering::Relaxed); - // Dropping the sender makes every worker's recv() return None. - *URL_JOBS.lock() = None; - // Take the handles first so the lock guard drops before the awaits. - let handles = URL_WORKER_HANDLES.lock().take(); - if let Some(handles) = handles { - for handle in handles { - if let Err(e) = handle.await { - log::error!("url worker panicked at shutdown: {e}"); - } - } - } -} - /// Extracts URL and text-link entities (text + caption), deduped in order. /// /// The offset work (`parse_entities` turning entities into slices of the @@ -253,7 +162,27 @@ fn thumbnail_for(media: &Media) -> Option { } } -fn media_to_payload(media: &Media, sensitive: bool) -> MediaItemPayload { +/// The link-cache snapshot of a freshly fetched post: its caption and raw +/// render fields, with no media yet — the send fills in the Telegram file ids +/// and persists the entry. `None` for a post that carries no render data +/// (nothing to rebuild a caption format from later). +pub(super) fn cached_snapshot(fetched: &x_media::site::Fetched) -> Option { + fetched + .render_fields() + .map(|(author, author_url, title, content, tags)| CachedPost { + url: fetched.source_url.clone(), + caption: fetched.caption.clone(), + title: title.to_string(), + content: content.to_string(), + author: author.to_string(), + author_url: author_url.to_string(), + tags: tags.to_string(), + sensitive: fetched.sensitive, + media: vec![], + }) +} + +pub(super) fn media_to_payload(media: &Media, sensitive: bool) -> MediaItemPayload { let fallback_url = media.smaller_url().map(str::to_string); match media { // A gif inside a group becomes a video item; a lone gif takes the @@ -707,20 +636,7 @@ async fn url_media_inner( let caption = fetched.caption_with(&format); // Raw render data for the link cache; the send fills in the // Telegram file ids and persists the entry. - let cache_data = - fetched - .render_fields() - .map(|(author, author_url, title, content, tags)| CachedPost { - url: fetched.source_url.clone(), - caption: fetched.caption.clone(), - title: title.to_string(), - content: content.to_string(), - author: author.to_string(), - author_url: author_url.to_string(), - tags: tags.to_string(), - sensitive: fetched.sensitive, - media: vec![], - }); + let cache_data = cached_snapshot(fetched); let items: Vec = fetched .media .iter() @@ -752,207 +668,12 @@ async fn url_media_inner( } } -// ── Startup repair: queued retries whose local media did not survive ─────── - -/// A post's fresh media plus the caption and cache snapshot that go with them: -/// what [`refetch`] hands [`apply_refresh`]. Plain data, so the rewrite below -/// can be tested without a network fetch (which cannot be faked here: -/// [`x_media::site::Fetched`] keeps a private field and is not constructible -/// outside its crate). -struct Refetched { - caption: String, - items: Vec, - cache_data: Option, -} - -/// Whether a queued task should have its post re-fetched, because it still -/// wants a local file (ugoira MP4, a bsky remux, a downloaded temp file) that is -/// gone. Those files live in the system temp dir and the registry that keeps -/// them alive for the retry (`send::KEEP_ALIVE`) is in memory, so a restart -/// takes all of them — a retry that needs one can only dead-letter. -/// -/// A partially delivered album is left alone: its remaining batches cannot be -/// reconciled with a fresh media list without risking a second copy of what the -/// user already received. -fn needs_refetch(task: &Task) -> bool { - if let Task::SendMediaSequence { - batch_index, - sent_message_ids, - .. - } = task - && (*batch_index > 0 || !sent_message_ids.is_empty()) - { - return false; - } - task.local_media_paths().iter().any(|path| !path.exists()) -} - -/// Rebuilds the task from the fresh media, keeping its delivery envelope (chat, -/// reply, forward/edit settings, notify targets): the retry that was queued must -/// still deliver the same way, whoever asked for it. -fn apply_refresh(task: &Task, fresh: &Refetched) -> Option { - let chat_id = task.chat_id()?; - let (edit_before_forward, forward_channel_id) = match task { - Task::SendMediaSequence { - edit_before_forward, - forward_channel_id, - .. - } - | Task::SendAnimation { - edit_before_forward, - forward_channel_id, - .. - } => (*edit_before_forward, *forward_channel_id), - Task::ForwardMessages { .. } => return None, - }; - let reply_to_message_id = match task { - Task::SendMediaSequence { - reply_to_message_id, - .. - } - | Task::SendAnimation { - reply_to_message_id, - .. - } => *reply_to_message_id, - Task::ForwardMessages { .. } => return None, - }; - let (notify_chat_id, notify_message_id) = task.notify_target(); - Some(Task::from_items( - Delivery { - chat_id, - reply_to_message_id, - edit_before_forward, - forward_channel_id, - notify_chat_id, - notify_message_id, - }, - task.source_url()?.to_string(), - fresh.caption.clone(), - fresh.items.clone(), - fresh.cache_data.clone(), - )) -} - -/// Fetches the post again and maps it into [`Refetched`]: the same mapping the -/// fresh-fetch path uses (per-site caption format from the chat, render fields -/// for the link-cache snapshot), so a repaired task looks like a first send. -async fn refetch( - ctx: &AppContext<'_>, - chat_id: i64, - url: &str, -) -> Result, x_media::site::FetchError> { - let Some(fetched) = x_media::site::fetch(url).await? else { - return Ok(None); - }; - if fetched.media.is_empty() { - return Ok(None); - } - let chat_data = ctx.chat_store.get(chat_id).await; - let format = chat_data.format_for(fetched.site_id); - let caption = fetched.caption_with(&format); - let cache_data = fetched - .render_fields() - .map(|(author, author_url, title, content, tags)| CachedPost { - url: fetched.source_url.clone(), - caption: fetched.caption.clone(), - title: title.to_string(), - content: content.to_string(), - author: author.to_string(), - author_url: author_url.to_string(), - tags: tags.to_string(), - sensitive: fetched.sensitive, - media: vec![], - }); - let items: Vec = fetched - .media - .iter() - .map(|media| media_to_payload(media, fetched.sensitive)) - .collect(); - // The re-fetch may produce a fresh local file (ugoira / bsky remux): hand it - // to the same keep-alive registry the first fetch uses. - if let Some(dir) = fetched.keep_alive() { - send::KEEP_ALIVE.lock().push(dir); - } - Ok(Some(Refetched { - caption, - items, - cache_data, - })) -} - -/// Re-fetches every queued task whose local media did not survive the restart, -/// so the user's link is still delivered instead of dead-lettering on a file -/// that cannot come back. Returns how many rows were rewritten. -/// -/// Startup only, before the queue workers start: no worker can lease a row while -/// this writes, which is what lets it replace payloads without the lease-token -/// guard every worker write-back carries. -pub(crate) async fn repair_lost_local_media(ctx: &AppContext<'_>) -> usize { - let mut repaired = 0; - for (id, payload) in ctx.task_queue.runnable_rows().await { - let Ok(task) = serde_json::from_str::(&payload) else { - continue; - }; - if !needs_refetch(&task) { - continue; - } - let (Some(url), Some(chat_id)) = (task.source_url().map(str::to_string), task.chat_id()) - else { - continue; - }; - match refetch(ctx, chat_id, &url).await { - Ok(Some(fresh)) => { - let Some(updated) = apply_refresh(&task, &fresh) else { - continue; - }; - let updated = serde_json::to_value(&updated).expect("task serializes"); - if ctx.task_queue.replace_payload(&id, &updated).await { - repaired += 1; - log::info!( - "startup repair: re-fetched [key={}] for chat={chat_id} (its local media did not survive the restart)", - log_key(&url) - ); - } - } - // The post is gone or withheld now: the retry could not have - // delivered anything either, so say why instead of letting it - // dead-letter on a missing file. - Ok(None) | Err(_) => { - let (notify_chat_id, notify_message_id) = task.notify_target(); - log::warn!( - "startup repair: [key={}] for chat={chat_id} needed a re-fetch and none was possible", - log_key(&url) - ); - send::notify_failure( - ctx.sender, - notify_chat_id, - notify_message_id, - &format!( - "{} — the media held for retry was lost when the bot restarted and the post could not be fetched again. Please send the link again.", - log_key(&url) - ), - ) - .await; - } - } - } - repaired -} - #[cfg(test)] mod tests { use super::*; - use crate::ctx::test_support::{TestStores, api_error, cached_photo, photo_item}; + use crate::ctx::test_support::{TestStores, cached_photo, permanent_error}; use crate::media_sender::test_support::{MockSender, Outcome}; use std::time::Duration; - use teloxide::RequestError; - - /// The API error a caption edit that changes nothing answers with — what - /// the mocks script for a permanent send failure. A `fn` pointer, so it - /// can be handed to `MockSender::scripted` as-is. - fn permanent_error() -> RequestError { - api_error("Bad Request: message is not modified") - } #[tokio::test] async fn cache_hit_sends_file_ids_and_degrades_on_permanent_failure() { @@ -1359,168 +1080,6 @@ mod tests { assert!(dedupe_urls(vec![]).is_empty()); } - fn queued_task(media: &str, batch_index: usize, sent: Vec) -> Task { - Task::SendMediaSequence { - chat_id: 1, - reply_to_message_id: 2, - caption: "cap".into(), - media_batches: vec![vec![photo_item(media, false, false)]], - batch_index, - sent_message_ids: sent, - source_url: "https://x.com/u/status/1".into(), - edit_before_forward: true, - forward_channel_id: Some(2), - notify_chat_id: Some(1), - notify_message_id: Some(2), - cache_data: None, - } - } - - #[test] - fn only_tasks_missing_a_local_file_need_a_refetch() { - // A URL send needs nothing. - assert!(!needs_refetch(&queued_task("https://cdn/1.jpg", 0, vec![]))); - // A local path that is still there (a survived temp file) needs nothing. - let dir = tempfile::tempdir().unwrap(); - let alive = dir.path().join("ugoira.mp4"); - std::fs::write(&alive, b"x").unwrap(); - assert!(!needs_refetch(&queued_task( - alive.to_str().unwrap(), - 0, - vec![] - ))); - // A local path the restart took away does. - assert!(needs_refetch(&queued_task( - "/nonexistent-ugoira.mp4", - 0, - vec![] - ))); - // A partially delivered album is left to its own retry path. - assert!(!needs_refetch(&queued_task( - "/nonexistent-ugoira.mp4", - 1, - vec![7] - ))); - assert!(!needs_refetch(&queued_task( - "/nonexistent-ugoira.mp4", - 0, - vec![7] - ))); - // A channel copy holds no media. - assert!(!needs_refetch(&Task::ForwardMessages { - from_chat_id: 1, - to_chat_id: 2, - message_ids: vec![3], - notify_chat_id: None, - notify_message_id: None, - })); - } - - #[test] - fn apply_refresh_keeps_the_delivery_envelope() { - let task = queued_task("/nonexistent-ugoira.mp4", 0, vec![]); - let fresh = Refetched { - caption: "fresh caption".into(), - items: vec![photo_item("https://cdn/fresh.jpg", true, false)], - cache_data: None, - }; - match apply_refresh(&task, &fresh).expect("a repairable task") { - Task::SendMediaSequence { - chat_id, - reply_to_message_id, - caption, - media_batches, - batch_index, - sent_message_ids, - source_url, - edit_before_forward, - forward_channel_id, - notify_chat_id, - notify_message_id, - .. - } => { - // Same delivery: chat, reply, forward/edit settings, notify. - assert_eq!((chat_id, reply_to_message_id), (1, 2)); - assert!(edit_before_forward); - assert_eq!(forward_channel_id, Some(2)); - assert_eq!((notify_chat_id, notify_message_id), (Some(1), Some(2))); - assert_eq!(source_url, "https://x.com/u/status/1"); - // Fresh media, and nothing of it counted as sent yet. - assert_eq!(caption, "fresh caption"); - assert!( - matches!( - &media_batches[0][0], - MediaItemPayload::Photo { media, .. } if media == "https://cdn/fresh.jpg" - ), - "fresh media must replace the lost local file" - ); - assert!(matches!( - media_batches[0][0], - MediaItemPayload::Photo { - has_spoiler: true, - .. - } - )); - assert_eq!((batch_index, sent_message_ids.len()), (0, 0)); - } - other => panic!("expected a media sequence, got {other:?}"), - } - } - - /// The whole repair against a real post: a queued row whose media is a local - /// file the restart took away is re-fetched from its `source_url` and - /// rewritten in place, so the retry can still deliver it. - #[tokio::test] - #[ignore = "live network: requires outbound HTTPS to public.api.bsky.app"] - async fn live_repair_refetches_a_lost_local_media_row() { - let stores = TestStores::new(); - // An empty script: the repair must not need to tell the user anything. - let sender = MockSender::scripted(vec![], permanent_error); - let ctx = stores.ctx(&sender); - let mut task = queued_task("/nonexistent-ugoira.mp4", 0, vec![]); - if let Task::SendMediaSequence { source_url, .. } = &mut task { - *source_url = "https://bsky.app/profile/fu-futa.bsky.social/post/3laoveufjv224".into(); - } - stores - .task_queue() - .enqueue(serde_json::to_value(&task).unwrap(), crate::db::now_f64()) - .await - .unwrap(); - - assert_eq!(repair_lost_local_media(&ctx).await, 1); - - let updated: Task = serde_json::from_value(stores.queued_payload().await).unwrap(); - match updated { - Task::SendMediaSequence { - media_batches, - batch_index, - sent_message_ids, - caption, - .. - } => { - let media: Vec = media_batches - .iter() - .flatten() - .map(|item| match item { - MediaItemPayload::Photo { media, .. } - | MediaItemPayload::Video { media, .. } - | MediaItemPayload::Animation { media, .. } => media.clone(), - }) - .collect(); - assert!(!media.is_empty(), "the fresh fetch yielded no media"); - assert!( - media.iter().all(|m| m.starts_with("http")), - "the retry must be uploadable from URLs again: {media:?}" - ); - assert_eq!((batch_index, sent_message_ids.len()), (0, 0)); - assert!(!caption.is_empty()); - } - other => panic!("expected a repaired media sequence, got {other:?}"), - } - // The post was re-read, not re-delivered: nothing was sent. - assert!(sender.calls().is_empty(), "{:?}", sender.calls()); - } - #[test] fn fetch_errors_map_to_distinct_user_messages() { use x_media::site::FetchError;