mirror of
https://github.com/TheFunny/TelegramTwitterMediaBot.git
synced 2026-09-23 23:32:05 +00:00
fix(queue): heartbeat the lease so long tasks are not re-processed
The lease was set once to now + LOCK_TTL_SECONDS (120 s) with no renewal. Tasks that legitimately take longer — slow CDN downloads, ugoira encodes, rate-limited batch forwards (a 100-message channel copy waits ~4 min on the per-chat token bucket) — had their lease expire mid-run; the 30 s expiry sweep flipped the row back to pending and another worker processed it again, double-sending. run_with_lease now drives the handler through tokio::select! and refreshes locked_until every 30 s while it runs. The heartbeat lives in the same future as the handler, so a panicking worker still lets the sweep recover the row (no leaked task keeping the lease fresh forever).
This commit is contained in:
@@ -307,6 +307,11 @@ impl QueueWorker {
|
||||
}
|
||||
}
|
||||
|
||||
/// Processes one leased row, keeping the lease alive while the handler
|
||||
/// runs. Without the heartbeat a task longer than [`LOCK_TTL_SECONDS`]
|
||||
/// (slow download, ugoira encode, rate-limited batch forward) would have
|
||||
/// its lease expire mid-run; the expiry sweep would flip the row back to
|
||||
/// `pending` and another worker would process it again — duplicate sends.
|
||||
async fn process(&self, row: LeasedRow) {
|
||||
let payload: Value = match serde_json::from_str(&row.payload) {
|
||||
Ok(value) => value,
|
||||
@@ -318,7 +323,8 @@ impl QueueWorker {
|
||||
}
|
||||
};
|
||||
log::debug!("processing {} (attempt {})", row.id, row.attempts + 1);
|
||||
match (self.handler)(payload).await {
|
||||
let outcome = self.run_with_lease(&row.id, payload).await;
|
||||
match outcome {
|
||||
Ok(()) => {
|
||||
log::debug!("task {} completed", row.id);
|
||||
self.delete_row(&row.id).await;
|
||||
@@ -351,6 +357,42 @@ impl QueueWorker {
|
||||
}
|
||||
}
|
||||
|
||||
/// Drives the handler to completion, refreshing the row's `locked_until`
|
||||
/// every 30 s so the expiry sweep never re-leases a still-running task.
|
||||
/// 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> {
|
||||
let fut = (self.handler)(payload);
|
||||
tokio::pin!(fut);
|
||||
let mut interval = tokio::time::interval(Duration::from_secs(30));
|
||||
// The first interval tick fires immediately; skip it (the lease was
|
||||
// just set by lease_next).
|
||||
interval.tick().await;
|
||||
let id_owned = id.to_string();
|
||||
loop {
|
||||
tokio::select! {
|
||||
result = &mut fut => return result,
|
||||
_ = interval.tick() => {
|
||||
let now = now_f64();
|
||||
let id = id_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],
|
||||
)
|
||||
})
|
||||
.await;
|
||||
if let Err(e) = result {
|
||||
log::error!("queue lease heartbeat failed: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_row(&self, id: &str) {
|
||||
let id = id.to_string();
|
||||
let result = self
|
||||
|
||||
Reference in New Issue
Block a user