diff --git a/AGENTS.md b/AGENTS.md index 758747a..fdfa237 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -38,13 +38,13 @@ 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, 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`), 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/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, `with_conn` runs all rusqlite I/O in `spawn_blocking` | +| `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) | | `crates/xmedia-bot/src/state.rs` | `ChatStore`: parking_lot `Mutex` cache + SQLite write-through (`chat_state` table) | | `crates/xmedia-bot/src/link_cache.rs` | `LinkCache`: SQLite-backed cache (`link_cache` table) of successfully sent posts — raw caption fields + Telegram `file_id`s; repeat links re-send locally (no fetch/upload), TTL + prune, invalidated on permanent send failure | -| `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 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), `busy_timeout` on all connections | +| `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), `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 | | `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` | diff --git a/crates/x-media/src/lib.rs b/crates/x-media/src/lib.rs index 9fc0a31..97b0b4c 100644 --- a/crates/x-media/src/lib.rs +++ b/crates/x-media/src/lib.rs @@ -1,2 +1,9 @@ pub mod media; pub mod site; + +/// Prefix every temp file and temp dir this project creates, so a startup +/// sweep can recognise its own leftovers: a killed process leaves them behind +/// (`TempDir`/`NamedTempFile` clean up on drop, and a killed process runs no +/// destructors), and without a marker the only safe assumption about the OS +/// temp directory is "not mine". +pub const TEMP_FILE_PREFIX: &str = "tgxmb-"; diff --git a/crates/x-media/src/site/bsky/interface.rs b/crates/x-media/src/site/bsky/interface.rs index 8134efe..bd16a08 100644 --- a/crates/x-media/src/site/bsky/interface.rs +++ b/crates/x-media/src/site/bsky/interface.rs @@ -190,8 +190,14 @@ async fn resolve_bsky_video( return Err("bsky video has too many segments".to_string()); } - let frames_dir = tempfile::tempdir().map_err(|e| e.to_string())?; - let out_dir = tempfile::tempdir().map_err(|e| e.to_string())?; + let frames_dir = tempfile::Builder::new() + .prefix(crate::TEMP_FILE_PREFIX) + .tempdir() + .map_err(|e| e.to_string())?; + let out_dir = tempfile::Builder::new() + .prefix(crate::TEMP_FILE_PREFIX) + .tempdir() + .map_err(|e| e.to_string())?; let mut total: u64 = 0; let mut list = String::new(); for (i, seg) in segments.iter().enumerate() { diff --git a/crates/x-media/src/site/pixiv/api.rs b/crates/x-media/src/site/pixiv/api.rs index 0414714..db5bd69 100644 --- a/crates/x-media/src/site/pixiv/api.rs +++ b/crates/x-media/src/site/pixiv/api.rs @@ -221,6 +221,7 @@ impl PixivAPI { // memory: ugoira zips can be hundreds of MB, and the old // download_media_limited path spiked RAM up to the size cap. let mut zip_file = tempfile::Builder::new() + .prefix(crate::TEMP_FILE_PREFIX) .suffix(".zip") .tempfile() .map_err(|e| PixivError::Api(format!("temp zip failed: {e}")))?; @@ -233,8 +234,14 @@ impl PixivAPI { let frame_delays = metadata.frames.iter().map(|f| f.delay).collect::>(); let result = tokio::task::spawn_blocking(move || -> Result<(String, tempfile::TempDir), String> { - let frames_dir = tempfile::tempdir().map_err(|e| e.to_string())?; - let out_dir = tempfile::tempdir().map_err(|e| e.to_string())?; + let frames_dir = tempfile::Builder::new() + .prefix(crate::TEMP_FILE_PREFIX) + .tempdir() + .map_err(|e| e.to_string())?; + let out_dir = tempfile::Builder::new() + .prefix(crate::TEMP_FILE_PREFIX) + .tempdir() + .map_err(|e| e.to_string())?; // Extract frames to canonical zero-padded names; pixiv ugoira // frames are uniformly jpg or png per artwork. The zip is read diff --git a/crates/xmedia-bot/src/db.rs b/crates/xmedia-bot/src/db.rs index 0909430..a8aa316 100644 --- a/crates/xmedia-bot/src/db.rs +++ b/crates/xmedia-bot/src/db.rs @@ -124,9 +124,39 @@ pub fn open_store(path: &str) -> rusqlite::Result> { } let conn = open_db(path)?; schema_init(&conn)?; + migrate(&conn)?; Ok(Arc::new(DbPool::new(path))) } +/// Schema migrations, applied in order and tracked by `PRAGMA user_version` +/// (the index in this array + 1 is the version a statement brings the +/// database to). Append only — never edit or reorder an entry, or databases +/// already past it would skip or repeat work. +const MIGRATIONS: &[&str] = &[ + // 1: lease fencing. A worker's write-backs (`delete`/`reschedule`/the + // lease heartbeat) are guarded by the token it was leased with, so a + // lease that expired and was re-leased by another worker can no longer be + // written by its former holder — which used to duplicate a send or drop + // the new holder's retry state, silently. + "ALTER TABLE tasks ADD COLUMN lease_token TEXT", +]; + +/// Brings an existing database up to [`MIGRATIONS`]. Idempotent: a database +/// already at the latest version does no work. +fn migrate(conn: &Connection) -> rusqlite::Result<()> { + let version: i64 = conn.query_row("PRAGMA user_version", [], |row| row.get(0))?; + for (index, statement) in MIGRATIONS.iter().enumerate() { + let target = index as i64 + 1; + if version >= target { + continue; + } + conn.execute_batch(statement)?; + // `PRAGMA` does not take bind parameters; the value is our own index. + conn.execute_batch(&format!("PRAGMA user_version = {target}"))?; + } + Ok(()) +} + fn rusqlite_error(e: std::io::Error) -> rusqlite::Error { rusqlite::Error::ToSqlConversionFailure(Box::new(e)) } @@ -135,11 +165,11 @@ fn rusqlite_error(e: std::io::Error) -> rusqlite::Error { /// 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. +/// This is the **baseline** schema (version 0): a fresh database is created +/// exactly like this, and anything that must *change* an existing one is +/// appended to [`MIGRATIONS`] instead of being edited in here — otherwise a +/// database created before the change would never gain the new column and a +/// freshly created one would try to apply the migration a second time. pub fn schema_init(conn: &Connection) -> rusqlite::Result<()> { conn.execute_batch( "CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, payload TEXT NOT NULL, \ diff --git a/crates/xmedia-bot/src/main.rs b/crates/xmedia-bot/src/main.rs index 7da7beb..056f293 100644 --- a/crates/xmedia-bot/src/main.rs +++ b/crates/xmedia-bot/src/main.rs @@ -1,4 +1,5 @@ use dotenv::dotenv; +use std::time::Duration; use teloxide::dptree::endpoint; use teloxide::prelude::*; use teloxide::stop::StopToken; @@ -40,6 +41,54 @@ fn spawn_sigterm_handler(stop_token: StopToken) { #[cfg(not(unix))] fn spawn_sigterm_handler(_stop_token: StopToken) {} +/// A leftover temp file must be at least this old before the startup sweep +/// touches it. Orphans come from a *previous* run; anything younger could +/// belong to a second instance sharing the temp directory (a misconfiguration, +/// but one that must not cost it its in-flight download). +const ORPHAN_TEMP_AGE: Duration = Duration::from_secs(3600); + +/// Removes this project's own leftover temp entries (`x_media::TEMP_FILE_PREFIX`) +/// from `dir` once they are older than `older_than`. Returns how many were +/// removed. Entries that are not ours, or are too young, or cannot be dated, +/// are left alone: the OS temp directory is shared, and the marker prefix plus +/// the age gate are the only two things that make deleting here safe. +fn sweep_temp_dir(dir: &std::path::Path, older_than: Duration) -> usize { + let Ok(entries) = std::fs::read_dir(dir) else { + return 0; + }; + let cutoff = std::time::SystemTime::now() - older_than; + let mut removed = 0; + for entry in entries.flatten() { + let name = entry.file_name(); + if !name + .to_string_lossy() + .starts_with(x_media::TEMP_FILE_PREFIX) + { + continue; + } + let old_enough = entry + .metadata() + .and_then(|meta| meta.modified()) + .is_ok_and(|modified| modified < cutoff); + if !old_enough { + continue; + } + let path = entry.path(); + let result = if entry.file_type().is_ok_and(|kind| kind.is_dir()) { + std::fs::remove_dir_all(&path) + } else { + std::fs::remove_file(&path) + }; + match result { + Ok(()) => removed += 1, + // Not worth a warning per entry: a file another process removed + // first (or one we may not delete) is not a problem here. + Err(e) => log::debug!("could not remove orphaned temp entry {path:?}: {e}"), + } + } + removed +} + #[tokio::main] async fn main() { dotenv().ok(); @@ -57,6 +106,17 @@ async fn main() { .init(); log::info!("Starting bot"); + // Temp media (downloaded files, ugoira/remux dirs) is cleaned up by + // `TempDir`/`NamedTempFile` on drop — which a killed process never runs. + // Without this sweep every hard restart left its downloads behind (up to + // hundreds of MB each) and nothing could tell them apart from a live + // process's files or from anything else in the OS temp dir. See + // [`sweep_temp_dir`] for why the age gate makes that safe. + let orphans = sweep_temp_dir(&std::env::temp_dir(), ORPHAN_TEMP_AGE); + if orphans > 0 { + log::info!("swept {orphans} orphaned temp file(s) from a previous run"); + } + let bot = Bot::from_env(); // Force the queue workers' shared Bot to initialize now so a missing // token fails at startup, not on the first queued task. @@ -260,3 +320,53 @@ async fn main() { log::info!("Bot stopped"); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sweep_removes_only_our_old_temp_entries() { + let dir = tempfile::tempdir().unwrap(); + let old = std::time::SystemTime::now() - Duration::from_secs(7200); + let make = |name: &str, aged: bool| { + let path = dir.path().join(name); + std::fs::write(&path, b"x").unwrap(); + if aged { + let file = std::fs::File::options().write(true).open(&path).unwrap(); + file.set_modified(old).unwrap(); + } + path + }; + let ours_old = make(&format!("{}photo-old.jpg", x_media::TEMP_FILE_PREFIX), true); + let ours_fresh = make( + &format!("{}photo-new.jpg", x_media::TEMP_FILE_PREFIX), + false, + ); + let theirs = make("someone-elses-file", true); + + assert_eq!(sweep_temp_dir(dir.path(), Duration::from_secs(3600)), 1); + assert!(!ours_old.exists(), "an old leftover of ours is removed"); + assert!(ours_fresh.exists(), "a fresh file may belong to a live run"); + assert!( + theirs.exists(), + "files without our prefix are never touched" + ); + + // A caller with no age gate also reaches the directory branch (aging a + // *directory* is not portable, so the gate is what the first half + // above proves): the fresh dir and file go, the unrelated file stays. + let leftover_dir = dir + .path() + .join(format!("{}ugoira", x_media::TEMP_FILE_PREFIX)); + std::fs::create_dir(&leftover_dir).unwrap(); + std::fs::write(leftover_dir.join("frame.png"), b"x").unwrap(); + assert_eq!(sweep_temp_dir(dir.path(), Duration::ZERO), 2); + assert!( + !leftover_dir.exists(), + "leftover dirs go with their contents" + ); + assert!(!ours_fresh.exists(), "no age gate: ours, however fresh"); + assert!(theirs.exists()); + } +} diff --git a/crates/xmedia-bot/src/photo.rs b/crates/xmedia-bot/src/photo.rs index b73d72d..b228613 100644 --- a/crates/xmedia-bot/src/photo.rs +++ b/crates/xmedia-bot/src/photo.rs @@ -199,6 +199,7 @@ fn encode_jpeg(pix: &PixBuf, w: u32, h: u32) -> Result, String> { fn write_temp(bytes: &[u8], ext: &str) -> Result { let mut file = tempfile::Builder::new() + .prefix(x_media::TEMP_FILE_PREFIX) .suffix(&format!(".{ext}")) .tempfile() .map_err(|e| format!("temp file failed: {e}"))?; diff --git a/crates/xmedia-bot/src/queue.rs b/crates/xmedia-bot/src/queue.rs index aa2ba0f..9c8a41d 100644 --- a/crates/xmedia-bot/src/queue.rs +++ b/crates/xmedia-bot/src/queue.rs @@ -68,8 +68,19 @@ struct LeasedRow { id: String, payload: String, attempts: i32, + /// Random token for *this* lease. Every write-back the worker makes is + /// guarded by it, so a lease that expired (heartbeat starved, host + /// suspended) and was re-leased by another worker cannot be written by + /// its former holder. + lease_token: String, } +/// The row is no longer ours: its lease expired and another worker took it. +/// The former holder must not write anything back — a `delete` would erase the +/// new holder's row (or a `reschedule` would overwrite its retry state) — so +/// the attempt stops at the next heartbeat instead. +struct LeaseLost; + /// Owned worker state so the spawned loop does not borrow the queue handle. #[derive(Clone)] struct QueueWorker { @@ -252,15 +263,22 @@ impl PersistentTaskQueue { /// Last-resort terminal state for a row whose `DELETE` would not go through: /// `done` is invisible to `lease_next` (`status='pending'`), to the expiry /// sweep (`status='in_progress'`) and to the backlog line, so a task that -/// already ran cannot be leased and run again. -async fn mark_done(pool: &std::sync::Arc, id: &str) -> rusqlite::Result<()> { +/// already ran cannot be leased and run again. Token-guarded like every other +/// write-back: `Ok(false)` means the row was re-leased and is not ours to +/// tombstone. +async fn mark_done( + pool: &std::sync::Arc, + id: &str, + lease_token: &str, +) -> rusqlite::Result { let id = id.to_string(); + let lease_token = lease_token.to_string(); pool.with_conn(move |conn| { - conn.execute( - "UPDATE tasks SET status='done', locked_until=0 WHERE id = ?1", - params![id], + let affected = conn.execute( + "UPDATE tasks SET status='done', locked_until=0 WHERE id = ?1 AND lease_token = ?2", + params![id, lease_token], )?; - Ok(()) + Ok(affected == 1) }) .await } @@ -364,15 +382,17 @@ impl QueueWorker { } Err(e) => return Err(e), }; + let lease_token = format!("{:016x}", rand::random::()); tx.execute( - "UPDATE tasks SET status='in_progress', locked_until=?1 WHERE id=?2", - params![now + LOCK_TTL_SECONDS, id], + "UPDATE tasks SET status='in_progress', locked_until=?1, lease_token=?2 WHERE id=?3", + params![now + LOCK_TTL_SECONDS, lease_token, id], )?; tx.commit()?; Ok(Some(LeasedRow { id, payload, attempts, + lease_token, })) }) .await @@ -410,7 +430,7 @@ impl QueueWorker { Ok(value) => value, Err(e) => { log::error!("queue: unparseable payload for {}: {e}", row.id); - self.delete_row(&row.id).await; + self.delete_row(&row.id, &row.lease_token).await; (self.dead_letter)(Value::Null, format!("invalid stored payload: {e}")).await; return; } @@ -422,12 +442,28 @@ impl QueueWorker { row.attempts + 1 ); let attempt_started = std::time::Instant::now(); - let outcome = self.run_with_lease(&row.id, payload).await; + let outcome = match self + .run_with_lease(&row.id, &row.lease_token, payload) + .await + { + Ok(outcome) => outcome, + Err(LeaseLost) => { + // Another worker owns this row now and is delivering the same + // task: write nothing (no delete, no reschedule, no + // dead-letter) and leave it to them. + log::warn!( + "queue: lost the lease on {} {fields} (attempt {}); abandoning this attempt", + row.id, + row.attempts + 1 + ); + return; + } + }; let attempt_ms = attempt_started.elapsed().as_millis(); match outcome { Ok(()) => { log::debug!("task {} {fields} completed in {attempt_ms}ms", row.id); - self.delete_row(&row.id).await; + self.delete_row(&row.id, &row.lease_token).await; } Err(QueueError::Retryable { delay_seconds, @@ -444,7 +480,7 @@ impl QueueWorker { row.id, row.attempts + 1 ); - self.delete_row(&row.id).await; + self.delete_row(&row.id, &row.lease_token).await; (self.dead_letter)(payload, message).await; } else { let delay = scaled_retry_delay(delay_seconds, row.attempts); @@ -453,13 +489,13 @@ impl QueueWorker { row.id, row.attempts + 1 ); - self.reschedule(&row.id, payload, delay, row.attempts + 1) + self.reschedule(&row.id, &row.lease_token, payload, delay, row.attempts + 1) .await; } } Err(QueueError::Permanent { message, payload }) => { log::error!("dead-lettering {} {fields}: {message}", row.id); - self.delete_row(&row.id).await; + self.delete_row(&row.id, &row.lease_token).await; (self.dead_letter)(payload, message).await; } } @@ -470,7 +506,12 @@ impl QueueWorker { /// The heartbeat is part of this future, not a separate spawned task: if /// the worker task dies (panic) the heartbeat dies with it and the sweep /// recovers the row exactly as before. - async fn run_with_lease(&self, id: &str, payload: Value) -> Result<(), QueueError> { + async fn run_with_lease( + &self, + id: &str, + lease_token: &str, + payload: Value, + ) -> Result, LeaseLost> { let fut = (self.handler)(payload); tokio::pin!(fut); let mut interval = tokio::time::interval(Duration::from_secs(30)); @@ -478,23 +519,33 @@ impl QueueWorker { // just set by lease_next). interval.tick().await; let id_owned = id.to_string(); + let token_owned = lease_token.to_string(); loop { tokio::select! { - result = &mut fut => return result, + result = &mut fut => return Ok(result), _ = interval.tick() => { let now = now_f64(); let id = id_owned.clone(); + let token = token_owned.clone(); let result = self .pool .with_conn(move |conn| { conn.execute( - "UPDATE tasks SET locked_until=?1 WHERE id=?2 AND status='in_progress'", - params![now + LOCK_TTL_SECONDS, id], + "UPDATE tasks SET locked_until=?1 \ + WHERE id=?2 AND status='in_progress' AND lease_token=?3", + params![now + LOCK_TTL_SECONDS, id, token], ) }) .await; - if let Err(e) = result { - log::error!("queue lease heartbeat failed: {e}"); + match result { + // Still ours: the lease is extended. + Ok(1) => {} + // The row is no longer leased to us (another worker + // re-leased it, or it is gone): dropping the handler + // future here stops this attempt instead of racing the + // new holder through the same send. + Ok(_) => return Err(LeaseLost), + Err(e) => log::error!("queue lease heartbeat failed: {e}"), } } } @@ -511,18 +562,26 @@ impl QueueWorker { /// already ran can never be re-leased. Both writes failing is logged at /// error level with the row id, since that is the one case where a /// duplicate send stays possible. - async fn delete_row(&self, id: &str) { + async fn delete_row(&self, id: &str, lease_token: &str) { for attempt in 0..TERMINAL_WRITE_ATTEMPTS { - match self.try_delete_row(id).await { - Ok(()) => return, + match self.try_delete_row(id, lease_token).await { + Ok(true) => return, + // The row is not ours any more (re-leased while we worked): + // leaving it alone *is* the clean outcome — retrying or + // tombstoning here would erase the new holder's work. + Ok(false) => { + log::warn!("queue: row {id} was re-leased; not deleting it"); + return; + } Err(e) => { log::error!("queue delete failed (attempt {}): {e}", attempt + 1); tokio::time::sleep(terminal_write_backoff(attempt)).await; } } } - match mark_done(&self.pool, id).await { - Ok(()) => log::warn!("queue: row {id} marked done instead of deleted"), + match mark_done(&self.pool, id, lease_token).await { + Ok(true) => log::warn!("queue: row {id} marked done instead of deleted"), + Ok(false) => log::warn!("queue: row {id} was re-leased; nothing to tombstone"), Err(e) => log::error!( "queue: row {id} could not be deleted or marked done ({e}); \ the expiry sweep may run this finished task again" @@ -530,12 +589,17 @@ impl QueueWorker { } } - async fn try_delete_row(&self, id: &str) -> rusqlite::Result<()> { + /// `Ok(false)` when the `WHERE` matched no row — the lease is not ours. + async fn try_delete_row(&self, id: &str, lease_token: &str) -> rusqlite::Result { let id = id.to_string(); + let lease_token = lease_token.to_string(); self.pool .with_conn(move |conn| { - conn.execute("DELETE FROM tasks WHERE id = ?1", params![id])?; - Ok(()) + let affected = conn.execute( + "DELETE FROM tasks WHERE id = ?1 AND lease_token = ?2 AND status='in_progress'", + params![id, lease_token], + )?; + Ok(affected == 1) }) .await } @@ -547,30 +611,48 @@ impl QueueWorker { /// safe terminal fallback here (marking it done would drop the retry /// without telling anyone), so a persistent failure is logged loudly and /// the sweep's re-run — at-least-once, the documented trade — is named. - async fn reschedule(&self, id: &str, payload: Value, delay_seconds: f64, attempts: i32) { - let id = id.to_string(); + async fn reschedule( + &self, + id: &str, + lease_token: &str, + payload: Value, + delay_seconds: f64, + attempts: i32, + ) { + let row_id = id.to_string(); + let lease_token = lease_token.to_string(); let payload = payload.to_string(); let run_after = now_f64() + delay_seconds; let mut last_error = None; for attempt in 0..TERMINAL_WRITE_ATTEMPTS { - let id = id.clone(); + let id = row_id.clone(); + let lease_token = lease_token.clone(); let payload = payload.clone(); let result = self .pool .with_conn(move |conn| { - conn.execute( - "UPDATE tasks SET payload=?1, run_after=?2, attempts=?3, status='pending', locked_until=0 WHERE id=?4", - params![payload, run_after, attempts, id], + let affected = conn.execute( + "UPDATE tasks SET payload=?1, run_after=?2, attempts=?3, status='pending', locked_until=0 \ + WHERE id=?4 AND lease_token=?5 AND status='in_progress'", + params![payload, run_after, attempts, id, lease_token], )?; - Ok(()) + Ok(affected == 1) }) .await; match result { - Ok(()) => { + Ok(true) => { // Same permit semantics as enqueue: never lose the wakeup. self.notify.notify_one(); return; } + // Re-leased while we worked: the new holder owns the row and + // its retry, so writing our payload would overwrite progress. + Ok(false) => { + log::warn!( + "queue: row {row_id} was re-leased; not rescheduling it (the new holder decides)" + ); + return; + } Err(e) => { log::error!("queue reschedule failed (attempt {}): {e}", attempt + 1); last_error = Some(e.to_string()); @@ -579,7 +661,7 @@ impl QueueWorker { } } log::error!( - "queue: row {id} could not be rescheduled ({}); the expiry sweep will \ + "queue: row {row_id} could not be rescheduled ({}); the expiry sweep will \ re-run this attempt from its previous state", last_error.unwrap_or_default() ); @@ -605,6 +687,22 @@ mod tests { assert_eq!(scaled_retry_delay(1800.0, 1), 1800.0); } + /// Puts a row into the state a worker holds while running it. + async fn set_lease(queue: &PersistentTaskQueue, id: &str, token: &str) { + let (id, token) = (id.to_string(), token.to_string()); + queue + .pool + .with_conn(move |conn| { + conn.execute( + "UPDATE tasks SET status='in_progress', lease_token=?1 WHERE id=?2", + params![token, id], + )?; + Ok(()) + }) + .await + .unwrap(); + } + async fn new_queue() -> (PersistentTaskQueue, tempfile::TempDir) { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("queue.db"); @@ -672,7 +770,16 @@ mod tests { .with_conn(|conn| conn.query_row("SELECT id FROM tasks", [], |r| r.get(0))) .await .unwrap(); - mark_done(&queue.pool, &id).await.unwrap(); + // A token that is not the row's is refused: only the lease holder can + // write the row back. + assert!( + !mark_done(&queue.pool, &id, "someone-elses-token") + .await + .unwrap(), + "a foreign lease must not be able to tombstone the row" + ); + set_lease(&queue, &id, "ours").await; + assert!(mark_done(&queue.pool, &id, "ours").await.unwrap()); assert_eq!( queue.pending_backlog().await, diff --git a/crates/xmedia-bot/src/send/mod.rs b/crates/xmedia-bot/src/send/mod.rs index 0ecd352..de85bfd 100644 --- a/crates/xmedia-bot/src/send/mod.rs +++ b/crates/xmedia-bot/src/send/mod.rs @@ -899,7 +899,7 @@ mod tests { // A send failure names the post (the cache key) and the cause, so the // user knows which of their links died. let task = sequence_task("https://x.com/u/status/1"); - let text = super::post_send::failure_text(Some(&task), "retries exhausted"); + let text = super::post_send::failure_text(task.source_url(), "retries exhausted"); assert!(text.contains("twitter:1"), "{text}"); assert!(text.contains("retries exhausted"), "{text}"); @@ -912,7 +912,7 @@ mod tests { notify_chat_id: None, notify_message_id: None, }; - let text = super::post_send::failure_text(Some(&forward), "chat not found"); + let text = super::post_send::failure_text(forward.source_url(), "chat not found"); assert!(text.starts_with("Forward failed permanently"), "{text}"); assert!(text.contains("chat not found"), "{text}"); } diff --git a/crates/xmedia-bot/src/send/post_send.rs b/crates/xmedia-bot/src/send/post_send.rs index 416cff8..df020d4 100644 --- a/crates/xmedia-bot/src/send/post_send.rs +++ b/crates/xmedia-bot/src/send/post_send.rs @@ -321,7 +321,7 @@ pub(crate) async fn post_send_actions(ctx: &AppContext<'_>, task: &Task, message ctx.sender, notify_chat_id, notify_message_id, - &failure_text(Some(&task), "retry could not be queued"), + &failure_text(task.source_url(), "retry could not be queued"), ) .await; } @@ -435,15 +435,34 @@ async fn send_media_or_animation(ctx: &AppContext<'_>, task: &Task) -> Result, message: &str) -> String { - match task.and_then(|task| task.source_url()).map(log_key) { +pub(super) fn failure_text(source_url: Option<&str>, message: &str) -> String { + match source_url.map(log_key) { Some(key) => format!("Send failed permanently for {key}: {message}"), - // `ForwardMessages` carries no source URL: that failure is about the - // channel copy, not about a post. + // `ForwardMessages` carries no source URL (and neither does an + // unparsable payload): that failure is about the channel copy, not + // about a post. None => format!("Forward failed permanently: {message}"), } } +/// The post a stored payload is about, without parsing it into a [`Task`]: +/// used when the payload no longer deserializes (written by an older version, +/// or corrupted) but its identity fields are still readable. +fn payload_source_url(payload: &serde_json::Value) -> Option<&str> { + payload.get("source_url").and_then(|v| v.as_str()) +} + +/// Whether a stored payload was a *cached* send (see `Task::is_cached_send`), +/// read straight off the JSON — the unparsable case still has to know whether +/// a link-cache entry may be holding the media that failed. +fn payload_is_cached_send(payload: &serde_json::Value) -> bool { + payload + .get("cache_data") + .and_then(|data| data.get("media")) + .and_then(|media| media.as_array()) + .is_some_and(|media| !media.is_empty()) +} + /// Dead-letter callback wired to the queue in main: settles the task and /// notifies its chat. pub(crate) async fn dead_letter_notify( @@ -457,6 +476,18 @@ pub(crate) async fn dead_letter_notify( let task = serde_json::from_value::(payload.clone()).ok(); if let Some(task) = &task { settle_task(ctx, task, Settled::Failed).await; + } else { + // A payload that no longer parses (an older version's row shape, a + // corrupted one) still says which post it was about: drop the stale + // cache entry the same way, instead of leaving a bad file id to be + // re-sent forever — and name the post in the notification rather than + // reporting a *forward* failure for a send task. + if payload_is_cached_send(&payload) + && let Some(key) = payload_source_url(&payload).and_then(x_media::site::cache_key) + { + log::debug!("removing stale link cache entry for [key={key}]"); + ctx.link_cache.remove(&key).await; + } } let notify_chat_id = payload.get("notify_chat_id").and_then(|v| v.as_i64()); let notify_message_id = payload.get("notify_message_id").and_then(|v| v.as_i64()); @@ -464,7 +495,12 @@ pub(crate) async fn dead_letter_notify( ctx.sender, notify_chat_id, notify_message_id, - &failure_text(task.as_ref(), &message), + &failure_text( + task.as_ref() + .and_then(|task| task.source_url()) + .or_else(|| payload_source_url(&payload)), + &message, + ), ) .await; } diff --git a/crates/xmedia-bot/src/send/upload.rs b/crates/xmedia-bot/src/send/upload.rs index 3151d55..39c865f 100644 --- a/crates/xmedia-bot/src/send/upload.rs +++ b/crates/xmedia-bot/src/send/upload.rs @@ -75,6 +75,7 @@ async fn download_to_temp( }; let ext = sniff_ext(&bytes); let mut file = tempfile::Builder::new() + .prefix(x_media::TEMP_FILE_PREFIX) .suffix(&format!(".{ext}")) .tempfile() .map_err(|e| FallbackError::Permanent {