fix(queue): gate post-send side effects on the row's lease

A handler whose row was re-leased while it sat between heartbeats (a
>=120 s stall) still ran the channel forward, opened the edit prompt and
released the keep-alive temp media the new holder was using; only the
terminal write was token-guarded, which left the side effects to run
twice. The lease is now exposed to the handler on a task-local
(Lease/lease_still_held) and the completion path re-reads it, since the
loss may have gone unnoticed by the heartbeat. The forward retry is gated
the same way.
This commit is contained in:
2026-10-14 22:10:36 +08:00
parent 632cbb7c89
commit 3b955cf321
4 changed files with 243 additions and 20 deletions
+1 -1
View File
@@ -44,7 +44,7 @@ The `x-media` library: `site::fetch(url)` dispatches through the `SITES` registr
| `crates/xmedia-bot/src/handlers/` | Handler modules: `mod.rs` (message entry point, `reply`, `log_key`, the group-only `GROUP_LINK_HINT` for a supported link posted outside a private chat), `url_workers.rs`/`repair.rs` (worker pool; startup repair), `commands.rs` (teloxide `BotCommands` enum + command executor, incl. `/test <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 + the per-URL pipeline; one *shared* in-flight fetch per cache key (`shared_fetch`+`IN_FLIGHT_FETCHES`: a second chat, a batch forward or a retry asking for the same post meanwhile waits for the first caller's result, the entry is dropped the moment the fetch settles so nothing is ever answered from an old fetch, and a waiter whose sharer was cancelled fetches for itself); plus the startup repair `repair_lost_local_media`, whose decision (`needs_refetch`) and rewrite (`apply_refresh`) are pure and tested while the fetch itself is a live test), `url_workers.rs` (the bounded job channel (256) and its `URL_WORKERS = 8` supervised workers, `start_url_workers`/`stop_url_workers` — backpressure instead of unbounded spawns; teloxide's per-chat workers are sequential, so batch-forwards need this concurrency), `repair.rs` (startup `repair_lost_local_media` with its `needs_refetch`/`apply_refresh`/`refetch`), `inline.rs`/`callback.rs` (inline queries / edit-before-forward buttons, incl. `skip`; a forward that fails retryably is both queued *and* settles the prompt — the queued row carries the message ids itself, and a prompt left live let a second Confirm copy the same messages twice and let Skip answer "nothing was forwarded" while the row still delivered; a Permanent failure that already copied a batch settles the prompt too, since a second Confirm would rebuild the task from offset 0 and re-copy those messages, while one that copied nothing keeps it live for a retry), `statics.rs` (global statics) |
| `crates/xmedia-bot/src/state.rs` | `ChatStore`: parking_lot `Mutex<HashMap>` cache + SQLite write-through (`chat_state` table); the 300 s sweep's `prune_expired` evicts any chat with no live edit-before-forward prompt, so the cache (and the per-chat lock map) stays bounded to active prompts — durable settings reload from the DB on next use; `get` degrades a failed read to defaults for display/control flow, while `try_get` is the strict read the send pipeline uses — an unreadable row drops the send (and tells the user) instead of silently ignoring the chat's forward settings and persisting those defaults into the retry payload |
| `crates/xmedia-bot/src/link_cache.rs` | `LinkCache`: SQLite-backed cache (`link_cache` table) of successfully sent posts — raw caption fields + the source media URLs + Telegram `file_id`s; repeat links re-send locally (no fetch/upload), TTL + prune; a permanent send failure *degrades* the entry instead of dropping it (the file ids go, the URLs stay, so the next request re-sends from those without a fetch), and a degraded entry that fails again is removed |
| `crates/xmedia-bot/src/queue.rs` | `PersistentTaskQueue`: SQLite-backed queue (`tasks` table), `QUEUE_WORKERS = 4` concurrent workers (lease via `BEGIN IMMEDIATE` + `locked_until` TTL), retry→dead-letter, a `lease_token` fence: `lease_next` stamps a random token and every write-back (heartbeat, `delete`, `reschedule`, `mark_done`) is guarded by it, so a lease that expired and was re-leased cannot be written by its former holder — a lost lease stops the attempt instead; a finished row's `DELETE`/reschedule retried and a failed delete falling back to a `done` tombstone (the lease query and the sweep only look at `pending`/`in_progress`, so a task that already ran cannot be resurrected and re-run), `runnable_rows`/`replace_payload` (the startup repair's read/rewrite path: it runs before the workers exist, which is why it needs no lease token), `notify_one` worker wakeup plus a separate `Notify` for the 30 s lease-expiry sweep (a shared one let the sweep steal the workers' wakeup permit; the sweep does notify the workers after it actually recovered a row, since a recovered task is due immediately while every worker may be parked on `notify` with no pending row to sleep on), `busy_timeout` on all connections |
| `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 (and the handler verifies ownership with a fresh read — `Lease`/`lease_still_held`, carried on a task-local — before running side effects a re-lease would duplicate: a channel forward, an edit prompt, the keep-alive release); a finished row's `DELETE`/reschedule retried and a failed delete falling back to a `done` tombstone (the lease query and the sweep only look at `pending`/`in_progress`, so a task that already ran cannot be resurrected and re-run), `runnable_rows`/`replace_payload` (the startup repair's read/rewrite path: it runs before the workers exist, which is why it needs no lease token), `notify_one` worker wakeup plus a separate `Notify` for the 30 s lease-expiry sweep (a shared one let the sweep steal the workers' wakeup permit; the sweep does notify the workers after it actually recovered a row, since a recovered task is due immediately while every worker may be parked on `notify` with no pending row to sleep on), `busy_timeout` on all connections |
| `crates/xmedia-bot/src/ctx.rs` | `AppContext`: the injected collaborators (`sender` + `ChatStore`/`PersistentTaskQueue`/`LinkCache`/`Config`), `from_statics` for production and the `CONTEXT` static the worker closures hold. `test_support::TestStores` backs handler tests with a tempdir store set, and the module also carries the fixtures those tests share — the canonical cached post (`cached_photo`), the edit-before-forward prompt (`seed_prompt` with its `PROMPT_ID`/`FORWARDED_ID`) and a scripted API error (`api_error`) — so no two test modules keep their own copies |
| `crates/xmedia-bot/src/send/` | `send/mod.rs`: `Task`/`MediaItemPayload` payloads (`media: MediaRef`, i.e. `Source` URL-or-path vs `FileId` — one field used to carry both with a flag), `send_media_sequence`/`send_animation`/`forward_messages`; `send/error.rs`: the Bot API error policy (`SendError`/`Classification`, `classify_request_error`, the media-fetch/size markers); `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/{mod.rs,test_support.rs}` | `MediaSender` trait: the user-flow surface (`send_media_group`/`send_animation`/`copy_messages`/`send_message`/`answer_callback_query`/`edit_message_text`/`edit_message_caption`/`delete_message`/`send_chat_action`/`answer_inline_query`) 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`. `test_support.rs` (cfg(test)-only) holds the scripted `MockSender` and `fake_api` (the stand-in API the real-`Bot` tests drive) |
+125 -2
View File
@@ -41,6 +41,73 @@ fn heartbeat_failure_is_fatal(secs_since_confirmed: f64) -> bool {
secs_since_confirmed + HEARTBEAT_SECONDS >= LOCK_TTL_SECONDS
}
tokio::task_local! {
/// The lease of the attempt running on this worker task. A handler reads it
/// (via [`lease_still_held`]) to verify with a fresh DB read that it still
/// owns the row before running side effects a re-lease would duplicate.
static CURRENT_LEASE: Lease;
}
/// Handle to the row lease one queue attempt runs under.
#[derive(Clone)]
pub struct Lease {
pool: Arc<crate::db::DbPool>,
id: String,
token: String,
}
impl Lease {
/// Whether this attempt still owns the row. A fresh read on purpose: while
/// a stalled attempt sits between heartbeats its lease can expire and be
/// re-leased by the sweep, and only the DB knows.
pub async fn still_held(&self) -> bool {
let id = self.id.clone();
let token = self.token.clone();
self.pool
.with_conn_or(
log::Level::Warn,
"queue lease check failed",
false,
move |conn| {
conn.query_row(
"SELECT EXISTS(SELECT 1 FROM tasks \
WHERE id=?1 AND status='in_progress' AND lease_token=?2)",
params![id, token],
|row| row.get::<_, bool>(0),
)
},
)
.await
}
}
/// Whether the attempt running on this worker task still owns its row. Answers
/// `true` outside a queue worker (a direct handler call, as in tests), so a
/// caller never skips its effects for want of a lease.
pub async fn lease_still_held() -> bool {
match CURRENT_LEASE.try_with(|lease| lease.clone()) {
Ok(lease) => lease.still_held().await,
Err(_) => true,
}
}
/// Runs `f` with a lease installed as the current worker lease. Test-only seam
/// for driving a handler's lease check without a live worker loop.
#[cfg(test)]
pub(crate) async fn scope_lease_for_test<F: Future>(
pool: Arc<crate::db::DbPool>,
id: &str,
token: &str,
f: F,
) -> F::Output {
let lease = Lease {
pool,
id: id.to_string(),
token: token.to_string(),
};
CURRENT_LEASE.scope(lease, f).await
}
/// Number of concurrent worker loops. Tasks are independent (retries and
/// forward resumes); leases serialize row claims via SQLite transactions, so
/// extra workers drain backlogs faster. Each worker can be mid-send to
@@ -535,8 +602,19 @@ impl QueueWorker {
row.attempts + 1
);
let attempt_started = std::time::Instant::now();
let outcome = match self
.run_with_lease(&row.id, &row.lease_token, payload)
// The attempt runs under its own lease handle, exposed to the handler
// through a task-local: a handler about to run duplicate-able side
// effects verifies it still owns the row first (see `lease_still_held`).
let lease = Lease {
pool: Arc::clone(&self.pool),
id: row.id.clone(),
token: row.lease_token.clone(),
};
let outcome = match CURRENT_LEASE
.scope(
lease,
self.run_with_lease(&row.id, &row.lease_token, payload),
)
.await
{
Ok(outcome) => outcome,
@@ -832,6 +910,51 @@ mod tests {
assert!(heartbeat_failure_is_fatal(90.0));
}
/// The lease handle a handler sees is a live view of the row: once the row
/// is re-leased — what a stalled attempt hits — `lease_still_held` says so,
/// which is what gates the duplicate-able side effects (forward, prompt).
#[tokio::test]
async fn lease_still_held_is_false_once_the_row_is_re_leased() {
let (queue, _dir) = new_queue().await;
let pool = Arc::clone(&queue.pool);
let seen = Arc::new(Mutex::new(Vec::new()));
let seen_handler = Arc::clone(&seen);
let pool_handler = Arc::clone(&pool);
queue
.start(
move |_payload| {
let seen = Arc::clone(&seen_handler);
let pool = Arc::clone(&pool_handler);
async move {
let held_before = lease_still_held().await;
seen.lock().push(held_before);
// Steal the lease, as the sweep does once a stalled
// attempt let it expire.
pool.with_conn(|conn| {
conn.execute(
"UPDATE tasks SET lease_token='other' WHERE status='in_progress'",
[],
)
})
.await
.unwrap();
let held_after = lease_still_held().await;
seen.lock().push(held_after);
Ok(())
}
},
|_payload, _message| async {},
)
.await;
queue
.enqueue(serde_json::json!({"n": 1}), now_f64())
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(300)).await;
assert_eq!(&*seen.lock(), &[true, false]);
queue.stop().await;
}
async fn new_queue() -> (PersistentTaskQueue, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("queue.db");
+79
View File
@@ -1681,6 +1681,85 @@ mod tests {
assert_eq!(payload["message_ids"], serde_json::json!([10, 11]));
}
/// A re-leased row's handler must not run the forward: the new holder owns
/// that copy, so running it here too would deliver the messages twice.
#[tokio::test]
async fn a_forward_skips_the_copy_once_the_lease_is_gone() {
// The empty script panics on any call, so a copy that sneaks through
// fails loudly rather than hiding behind an assertion.
let sender = MockSender::scripted(vec![], || api_error("unused"));
let stores = TestStores::new();
let ctx = stores.ctx(&sender);
let task = Task::ForwardMessages {
from_chat_id: 1,
to_chat_id: 2,
message_ids: vec![10, 11],
forward_offset: 0,
notify_chat_id: None,
notify_message_id: None,
};
let payload = serde_json::to_value(&task).unwrap();
// A lease the tasks table does not hold: `still_held` reads false.
let pool = crate::db::open_store(stores.db_path()).unwrap();
let outcome =
crate::queue::scope_lease_for_test(pool, "gone", "tok", handle_task(&ctx, payload))
.await;
assert!(outcome.is_ok());
assert!(sender.calls().is_empty(), "no copy may run");
}
/// The completion gate: a media send whose row was re-leased must not open
/// the edit-before-forward prompt (nor forward) — the new holder does, so
/// doing it here too would prompt and forward the post a second time.
#[tokio::test]
async fn a_media_send_skips_post_send_actions_once_the_lease_is_gone() {
let sender = MockSender::scripted(vec![Outcome::GroupOk, Outcome::MessageOk], || {
api_error("unused")
});
let stores = TestStores::new();
let ctx = stores.ctx(&sender);
let mut task = sent_task(Some(2), false);
if let Task::SendMediaSequence {
media_batches,
edit_before_forward,
..
} = &mut task
{
// Two items so the send is one `send_media_group`; the prompt on so
// the completion path has an observable call — a group send alone
// reports no message ids, so a forward would copy nothing.
*media_batches = vec![vec![
photo_item("https://p/1.jpg", false, false),
photo_item("https://p/2.jpg", false, false),
]];
*edit_before_forward = true;
}
let payload = serde_json::to_value(&task).unwrap();
let pool = crate::db::open_store(stores.db_path()).unwrap();
let outcome =
crate::queue::scope_lease_for_test(pool, "gone", "tok", handle_task(&ctx, payload))
.await;
assert!(outcome.is_ok());
assert_eq!(
sender.calls(),
vec!["send_media_group"],
"the edit prompt must not open"
);
// Control: the same task under a live lease does open the prompt — so
// the assertion above cannot pass because the prompt was never
// observable in the first place.
let live = MockSender::scripted(vec![Outcome::GroupOk, Outcome::MessageOk], || {
api_error("unused")
});
let live_ctx = stores.ctx(&live);
let outcome = handle_task(&live_ctx, serde_json::to_value(&task).unwrap()).await;
assert!(outcome.is_ok());
assert_eq!(live.calls(), vec!["send_media_group", "send_message"]);
}
#[tokio::test]
async fn post_send_notifies_a_permanent_forward_failure() {
let sender = MockSender::scripted(vec![Outcome::CopyErr, Outcome::MessageErr], || {
+38 -17
View File
@@ -480,28 +480,49 @@ pub(crate) async fn handle_task(
// whole sequence (every batch) completed, so the channel forward
// and the edit-before-forward prompt must not be lost just
// because the send needed a retry.
//
// Unless the row was re-leased while this attempt sat between
// heartbeats (a ≥120 s stall): the new holder runs the same
// effects, so running them here too would copy to the channel a
// second time and open a second prompt. The keep-alive release is
// skipped with them — the payload (and so its temp media) is the
// same one the new holder is using.
if !crate::queue::lease_still_held().await {
log::warn!(
"queue: row re-leased before post-send actions; leaving the forward and the edit prompt to the new holder"
);
return Ok(());
}
post_send_actions(ctx, &task, message_ids).await;
settle_task(ctx, &task, Settled::Sent).await;
Ok(())
}
Task::ForwardMessages { .. } => match forward_messages(ctx, &task).await {
Ok(()) => Ok(()),
Err(SendError::Retryable {
delay_seconds,
task,
}) => Err(QueueError::Retryable {
delay_seconds,
payload: serde_json::to_value(task).expect("task serializes"),
}),
Err(SendError::Permanent { message, task }) => {
// Settled by `dead_letter_notify`, which the queue invokes for
// this payload.
Err(QueueError::Permanent {
message,
payload: serde_json::to_value(task).expect("task serializes"),
})
Task::ForwardMessages { .. } => {
// The copy *is* the side effect here, so it is the thing to gate:
// a re-leased row is already being copied by its new holder.
if !crate::queue::lease_still_held().await {
log::warn!("queue: row re-leased before the forward; leaving it to the new holder");
return Ok(());
}
},
match forward_messages(ctx, &task).await {
Ok(()) => Ok(()),
Err(SendError::Retryable {
delay_seconds,
task,
}) => Err(QueueError::Retryable {
delay_seconds,
payload: serde_json::to_value(task).expect("task serializes"),
}),
Err(SendError::Permanent { message, task }) => {
// Settled by `dead_letter_notify`, which the queue invokes
// for this payload.
Err(QueueError::Permanent {
message,
payload: serde_json::to_value(task).expect("task serializes"),
})
}
}
}
}
}