mirror of
https://github.com/TheFunny/TelegramTwitterMediaBot.git
synced 2026-09-23 23:32:05 +00:00
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:
@@ -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 (`>` `<` `&` `'`) — 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` |
|
||||
|
||||
@@ -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-";
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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, \
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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}");
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user