fix(retry): fence the queue lease, clean up after a kill, name dead-lettered posts

P2 (hardening) of the retry audit, closing the report's remaining findings.

- Lease fencing. `lease_next` now stamps a random `lease_token`, and every
  write-back a worker makes (the 30s heartbeat, `delete_row`, `reschedule`,
  `mark_done`) is guarded by it. A lease that expired while its holder was
  stalled and was then re-leased used to let *both* holders write the same row:
  one duplicated the send, the other silently discarded the new holder's retry
  (a 0-row update was not even logged). Now a worker that no longer holds the
  lease drops its attempt at the next heartbeat and writes nothing. Reaching
  existing databases needed a migration chain, which `db.rs` had been
  pre-committed to: `MIGRATIONS` + `migrate` track `PRAGMA user_version`, with
  `schema_init` as the version-0 baseline. Verified on a database created
  before this change: user_version 0 -> 1, column added, rows intact.
- Dead-letter notifications no longer mislabel an unparsable payload. A row
  whose payload no longer deserializes as a `Task` (an older version's shape,
  corruption) used to skip the cache invalidation *and* report "Forward failed
  permanently" for a send task, because both were derived from the parsed
  value. The identity now comes off the raw JSON, so the stale link-cache entry
  is dropped and the message names the post.
- Temp files are marked and swept. Every temp file/dir the project creates now
  carries `x_media::TEMP_FILE_PREFIX`, and startup removes entries with that
  prefix older than an hour — a killed process leaves its downloads (up to
  hundreds of MB) behind because no destructor runs, and the age gate keeps the
  sweep away from a second instance's in-flight files. Verified live: the log
  reports the sweep, an aged leftover goes, a fresh prefixed file and an
  unrelated file stay.
This commit is contained in:
2026-09-20 21:34:45 +08:00
parent 0a82ca5a42
commit 4cdf618c25
11 changed files with 363 additions and 58 deletions
+3 -3
View File
@@ -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/<twitter\|pixiv\|bsky\|misskey\|bilibili>/` | One directory per site: `mod.rs` (re-exports), `interface.rs` (PATTERN, `enabled()`, `fetch_from_url()`, `cache_key`/`is_retryable`/`media_headers`, unit struct `<Name>Site` implementing `site::Site`, `From<SiteStruct> 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 (`&gt;` `&lt;` `&amp;` `&#39;`) — 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 <url>` (send-only) / `/debug <url>` (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<HashMap>` 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` |
+7
View File
@@ -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-";
+8 -2
View File
@@ -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() {
+9 -2
View File
@@ -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::<Vec<_>>();
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
+35 -5
View File
@@ -124,9 +124,39 @@ pub fn open_store(path: &str) -> rusqlite::Result<Arc<DbPool>> {
}
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, \
+110
View File
@@ -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());
}
}
+1
View File
@@ -199,6 +199,7 @@ fn encode_jpeg(pix: &PixBuf, w: u32, h: u32) -> Result<Vec<u8>, String> {
fn write_temp(bytes: &[u8], ext: &str) -> Result<NamedTempFile, String> {
let mut file = tempfile::Builder::new()
.prefix(x_media::TEMP_FILE_PREFIX)
.suffix(&format!(".{ext}"))
.tempfile()
.map_err(|e| format!("temp file failed: {e}"))?;
+145 -38
View File
@@ -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<crate::db::DbPool>, 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<crate::db::DbPool>,
id: &str,
lease_token: &str,
) -> rusqlite::Result<bool> {
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::<u64>());
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<Result<(), QueueError>, 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<bool> {
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,
+2 -2
View File
@@ -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}");
}
+42 -6
View File
@@ -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<Ve
/// User-facing text for a task that will never run again: which link died and
/// why. The raw error alone left the user guessing which post it was about.
pub(super) fn failure_text(task: Option<&Task>, 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::<Task>(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;
}
+1
View File
@@ -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 {