From 0a577600fdb64446559bc637f98c0d8bab1fe598 Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Wed, 16 Sep 2026 21:11:39 +0800 Subject: [PATCH] fix: repair four correctness defects in the send/state/handler paths - send: the PHOTO_INVALID_DIMENSIONS marker never matched (the description is lower-cased, the marker was not), so oversized photos sent by URL were classified Permanent instead of taking the download-and-downscale fallback. - inline: the debounce state was one global slot, so a second user's query cancelled the first user's pending answer entirely; it is now per user. - state: prune_expired wrote back a stale snapshot without the per-chat lock, clobbering a concurrent update() (lost edit-message record -> "Expired"); it now re-reads and prunes under the same lock update() uses. - send/queue: a task dead-lettered on retry exhaustion kept its keep-alive temp media (ugoira MP4) alive until process exit; dead_letter_notify now releases it, and enqueue_retry releases when the enqueue itself fails. Also folds the duplicated retry enqueue in post_send_actions into send::enqueue_retry (single clock source, single place that releases). Tests: +6 (marker, per-user debounce x3, keep-alive release, prune contract x2); the marker and keep-alive cases were verified to fail before the fix. cargo fmt/clippy clean, 52 + 69 tests pass. --- crates/xmedia-bot/src/handlers/callback.rs | 3 +- crates/xmedia-bot/src/handlers/inline.rs | 147 +++++++++++++++----- crates/xmedia-bot/src/handlers/urls.rs | 11 +- crates/xmedia-bot/src/send.rs | 69 ++++++++-- crates/xmedia-bot/src/state.rs | 153 +++++++++++++++------ 5 files changed, 278 insertions(+), 105 deletions(-) diff --git a/crates/xmedia-bot/src/handlers/callback.rs b/crates/xmedia-bot/src/handlers/callback.rs index 66fc34d..a4244f7 100644 --- a/crates/xmedia-bot/src/handlers/callback.rs +++ b/crates/xmedia-bot/src/handlers/callback.rs @@ -1,7 +1,6 @@ //! Callback query handling: the edit-before-forward prompt's "forward" and //! "template|" buttons. -use super::urls::enqueue_retry; use super::{CHAT_STORE, CONFIG, TASK_QUEUE}; use crate::db::unix_now; use crate::send::{self, Task}; @@ -83,7 +82,7 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<() task, }) => { log::info!("forward queued for retry in {delay_seconds:.1}s"); - enqueue_retry(&TASK_QUEUE, *task, delay_seconds).await; + send::enqueue_retry(&TASK_QUEUE, *task, delay_seconds).await; bot.answer_callback_query(callback_query_id) .text("Forward queued for retry.") .await?; diff --git a/crates/xmedia-bot/src/handlers/inline.rs b/crates/xmedia-bot/src/handlers/inline.rs index 7dca6eb..aa0e041 100644 --- a/crates/xmedia-bot/src/handlers/inline.rs +++ b/crates/xmedia-bot/src/handlers/inline.rs @@ -3,6 +3,7 @@ //! inline cache instead of re-fetching. use super::log_key; +use std::collections::HashMap; use std::sync::LazyLock; use teloxide::RequestError; use teloxide::prelude::*; @@ -19,16 +20,65 @@ use x_media::media::Media; /// post id. Only answer once the query has been stable for this long. const INLINE_DEBOUNCE: std::time::Duration = std::time::Duration::from_millis(800); -/// Last seen inline query and whether it was already answered. Guards the -/// debounce timer: a repeat of an answered query is served by Telegram's -/// inline cache (see `cache_time`), not by another fetch. +/// Last seen inline query per user and whether it was already answered. +/// Guards the debounce timer: a repeat of an answered query is served by +/// Telegram's inline cache (see `cache_time`), not by another fetch. Keyed by +/// user id — a single shared slot would let one user's typing burst (or a +/// different user's query) cancel another user's pending answer. struct InlineDebounceState { query: String, answered: bool, } -static INLINE_DEBOUNCE_STATE: LazyLock>> = - LazyLock::new(|| parking_lot::Mutex::new(None)); +#[derive(Default)] +struct DebounceStates(HashMap); + +impl DebounceStates { + /// Records `query` as the user's newest query. Returns false when it is a + /// repeat whose answer already went out (Telegram's inline cache serves + /// it; re-fetching would only hit the source site again). + fn note(&mut self, user_id: u64, query: &str) -> bool { + if let Some(prev) = self.0.get(&user_id) + && prev.query == query + && prev.answered + { + return false; + } + self.0.insert( + user_id, + InlineDebounceState { + query: query.to_string(), + answered: false, + }, + ); + true + } + + /// Claims the answer for the user's newest query; false when a newer query + /// superseded it or the answer was already claimed. + fn claim(&mut self, user_id: u64, query: &str) -> bool { + let Some(state) = self.0.get_mut(&user_id) else { + return false; + }; + if state.query != query || state.answered { + return false; + } + state.answered = true; + true + } + + /// Releases a claimed-but-unsent answer so a repeat can retry the fetch. + fn release(&mut self, user_id: u64, query: &str) { + if let Some(state) = self.0.get_mut(&user_id) + && state.query == query + { + state.answered = false; + } + } +} + +static INLINE_DEBOUNCE_STATE: LazyLock> = + LazyLock::new(|| parking_lot::Mutex::new(DebounceStates::default())); pub async fn inline_query_handler(bot: Bot, query: InlineQuery) -> Result<(), RequestError> { if query.query.is_empty() { @@ -41,48 +91,23 @@ pub async fn inline_query_handler(bot: Bot, query: InlineQuery) -> Result<(), Re // Debounce: record the query and answer only after it has been stable for // INLINE_DEBOUNCE (the timer below). An already-answered repeat of the // same query is left to Telegram's inline cache instead of re-fetching. - { - let mut state = INLINE_DEBOUNCE_STATE.lock(); - if let Some(prev) = state.as_ref() - && prev.query == query.query - && prev.answered - { - return respond(()); - } - *state = Some(InlineDebounceState { - query: query.query.clone(), - answered: false, - }); + let user_id = query.from.id.0; + if !INLINE_DEBOUNCE_STATE.lock().note(user_id, &query.query) { + return respond(()); } let query_text = query.query.clone(); tokio::spawn(async move { tokio::time::sleep(INLINE_DEBOUNCE).await; - // Only the last query of a typing burst survives: earlier timers see - // the query changed and give up without answering. - { - let mut state = INLINE_DEBOUNCE_STATE.lock(); - let Some(state) = state.as_mut() else { - return; - }; - if state.query != query_text || state.answered { - return; - } - // Claim the answer so a repeat of the same query cannot start a - // second fetch; reset below when no answer was produced. - state.answered = true; + // Only the user's last query of a typing burst survives: earlier + // timers see the query changed and give up without answering. + if !INLINE_DEBOUNCE_STATE.lock().claim(user_id, &query_text) { + return; } match answer_inline_query(bot, query).await { Ok(true) => {} // No results produced (or nothing to answer): let a repeat of the // same query retry the fetch. - Ok(false) | Err(_) => { - let mut state = INLINE_DEBOUNCE_STATE.lock(); - if let Some(state) = state.as_mut() - && state.query == query_text - { - state.answered = false; - } - } + Ok(false) | Err(_) => INLINE_DEBOUNCE_STATE.lock().release(user_id, &query_text), } }); respond(()) @@ -159,3 +184,49 @@ async fn answer_inline_query(bot: Bot, query: InlineQuery) -> Result MediaItemPayload { } } -pub(crate) async fn enqueue_retry(queue: &PersistentTaskQueue, task: Task, delay_seconds: f64) { - let payload = serde_json::to_value(task).expect("task serializes"); - let run_after = now_f64() + delay_seconds; - if let Err(e) = queue.enqueue(payload, run_after).await { - log::error!("failed to enqueue retry: {e}"); - } -} - /// Sends a task and handles the outcome: post-send actions on success, retry /// enqueue on retryable failure, reply + link-cache invalidation on /// permanent failure (a stale cached file id must not repeat forever). @@ -216,7 +207,7 @@ async fn dispatch_send( "send for [key={}] failed, queued for retry in {delay_seconds:.1}s", log_key(url) ); - enqueue_retry(ctx.task_queue, *task, delay_seconds).await; + send::enqueue_retry(ctx.task_queue, *task, delay_seconds).await; let _ = reply( ctx.sender, chat_id, diff --git a/crates/xmedia-bot/src/send.rs b/crates/xmedia-bot/src/send.rs index de16a8c..8358982 100644 --- a/crates/xmedia-bot/src/send.rs +++ b/crates/xmedia-bot/src/send.rs @@ -3,12 +3,12 @@ //! URL is blocked by hotlink protection; the bot downloads the file itself //! and uploads it via multipart). -use crate::db::unix_now; +use crate::db::{now_f64, unix_now}; use crate::handlers::{CHAT_STORE, LINK_CACHE, TASK_QUEUE, log_key}; use crate::link_cache::{CachedMedia, CachedMediaKind, CachedPost, LinkCache}; use crate::media_sender::MediaSender; use crate::photo::{self, MAX_UPLOAD_BYTES, PhotoPrep}; -use crate::queue::QueueError; +use crate::queue::{PersistentTaskQueue, QueueError}; use crate::state::EditMessage; use rand::Rng; use serde::{Deserialize, Serialize}; @@ -334,7 +334,7 @@ pub fn is_media_fetch_failure(e: &ApiError) -> bool { "timeout", // Oversized photos (width + height > 10000 px) are rejected on URL // sends too; route them to the download-and-resize fallback. - "PHOTO_INVALID_DIMENSIONS", + "photo_invalid_dimensions", ]; let description = e.to_string().to_lowercase(); MARKERS.iter().any(|marker| description.contains(marker)) @@ -1279,15 +1279,7 @@ pub async fn post_send_actions(sender: &dyn MediaSender, task: &Task, message_id delay_seconds, task, }) => { - let payload = serde_json::to_value(task).expect("task serializes"); - let run_after = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs_f64()) - .unwrap_or(0.0) - + delay_seconds; - if let Err(e) = TASK_QUEUE.enqueue(payload, run_after).await { - log::error!("failed to enqueue forward retry: {e}"); - } + enqueue_retry(&TASK_QUEUE, *task, delay_seconds).await; } Err(SendError::Permanent { message, .. }) => { notify_failure( @@ -1302,6 +1294,18 @@ pub async fn post_send_actions(sender: &dyn MediaSender, task: &Task, message_id } } +/// Enqueues a task for a later attempt (retry / forward resume). When the +/// enqueue itself fails the task can never be sent again, so its keep-alive +/// temp media is released instead of leaking until process exit. +pub async fn enqueue_retry(queue: &PersistentTaskQueue, task: Task, delay_seconds: f64) { + let payload = serde_json::to_value(&task).expect("task serializes"); + let run_after = now_f64() + delay_seconds; + if let Err(e) = queue.enqueue(payload, run_after).await { + log::error!("failed to enqueue retry: {e}"); + release_keep_alive(&task); + } +} + /// Queue entry point: parses the stored task and dispatches. pub async fn handle_task(payload: serde_json::Value) -> Result<(), QueueError> { let task: Task = match serde_json::from_value(payload.clone()) { @@ -1381,6 +1385,13 @@ async fn send_media_or_animation( /// Dead-letter callback wired to the queue in main: notifies the task's chat. pub async fn dead_letter_notify(payload: serde_json::Value, message: String) { + // A dead-lettered task never runs again. The queue dead-letters retry + // exhaustion itself (the handler is not called again), so this is the + // only place that sees the final payload — release the keep-alive temp + // media the fetch pipeline handed over, or it lives until process exit. + if let Ok(task) = serde_json::from_value::(payload.clone()) { + release_keep_alive(&task); + } 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()); if notify_chat_id.is_some() { @@ -1485,6 +1496,9 @@ mod tests { "Bad Request: EMPTY_WEB_MEDIA", "Bad Request: webpage_curl_failed", "Bad Request: request timeout", + // Telegram sends the code in upper case; the comparison is against + // the lower-cased description. + "Bad Request: PHOTO_INVALID_DIMENSIONS: width and height must be <= 10000", ] { let api = ApiError::Unknown(description.to_string()); assert!(is_media_fetch_failure(&api), "{description}"); @@ -1811,4 +1825,35 @@ mod tests { Err(SendError::Permanent { .. }) )); } + + #[tokio::test] + async fn dead_letter_releases_keep_alive_temp_media() { + // A task that exhausts its retries is dead-lettered by the queue + // without the handler running again: the keep-alive temp dir the + // fetch pipeline handed over must not outlive the task. + let dir = tempfile::tempdir().unwrap(); + let file = dir.path().join("ugoira.mp4"); + std::fs::write(&file, b"not-a-real-mp4").unwrap(); + let mut task = sequence_task(file.to_str().unwrap()); + if let Task::SendMediaSequence { + notify_chat_id, + notify_message_id, + .. + } = &mut task + { + // No chat to notify → no Bot is built by the notify path. + *notify_chat_id = None; + *notify_message_id = None; + } + let dir_path = dir.path().to_path_buf(); + KEEP_ALIVE.lock().push(dir); + + let payload = serde_json::to_value(&task).unwrap(); + dead_letter_notify(payload, "task failed after 2 retries".into()).await; + + assert!( + !KEEP_ALIVE.lock().iter().any(|dir| dir.path() == dir_path), + "dead-lettered task kept its temp media alive" + ); + } } diff --git a/crates/xmedia-bot/src/state.rs b/crates/xmedia-bot/src/state.rs index 177e314..ab77b48 100644 --- a/crates/xmedia-bot/src/state.rs +++ b/crates/xmedia-bot/src/state.rs @@ -102,19 +102,22 @@ impl ChatStore { } } + /// The per-chat async lock serializing get→mutate→set cycles. + fn lock_for(&self, chat_id: i64) -> Arc> { + self.locks + .lock() + .entry(chat_id) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() + } + /// Serializes a get→mutate→set cycle per chat: concurrent handler tasks /// (the batch-forward design spawns several per chat) each snapshot the /// same `ChatData` and last-writer-wins would silently drop mutations, /// e.g. a second `edit_message` record. The per-chat lock makes the /// cycle atomic. Returns the closure's result. pub async fn update(&self, chat_id: i64, f: impl FnOnce(&mut ChatData) -> R) -> R { - let lock = { - let mut locks = self.locks.lock(); - locks - .entry(chat_id) - .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) - .clone() - }; + let lock = self.lock_for(chat_id); let _guard = lock.lock().await; let mut data = self.get(chat_id).await; let r = f(&mut data); @@ -128,43 +131,45 @@ impl ChatStore { pub async fn prune_expired(&self, ttl: Duration) -> Vec<(i64, i64)> { let now = unix_now(); let ttl_secs = ttl.as_secs() as i64; - let mut removed = Vec::new(); - // Chats with no live edit records: evicted from the cache (and their - // per-chat lock) so the cache stays bounded to active prompts. The DB - // keeps the row; the next get() reloads it. - let mut evicted_chats = Vec::new(); - let changed: Vec<(i64, ChatData)> = { - let mut cache = self.cache.lock(); - let mut out = Vec::new(); - for (chat_id, data) in cache.iter_mut() { - let keys: Vec = data.edit_message.keys().copied().collect(); - let mut kept = HashMap::new(); - for key in keys { - if let Some(entry) = data.edit_message.get(&key) { - if entry.created_at + ttl_secs > now { - kept.insert(key, entry.clone()); - } else { - removed.push((*chat_id, key)); - } - } - } - if kept.len() != data.edit_message.len() { - // Persist the pruned row (removes expired records from - // the DB too, not just the cache). - data.edit_message = kept; - out.push((*chat_id, data.clone())); - } - if data.edit_message.is_empty() { - evicted_chats.push(*chat_id); - } - } - // Lock order: update() takes the per-chat lock before the cache - // lock, so prune must not hold the cache lock while taking locks. - drop(cache); - out + // Chats that may have an expired record, from a cache snapshot; the + // pruning itself re-reads and writes under the per-chat lock below + // (see the eviction note). Takes no lock of its own, so a chat + // appearing later is simply picked up by the next sweep. + let candidates: Vec = { + let cache = self.cache.lock(); + cache + .iter() + .filter(|(_, data)| { + data.edit_message + .values() + .any(|entry| entry.created_at + ttl_secs <= now) + }) + .map(|(chat_id, _)| *chat_id) + .collect() }; - for (chat_id, data) in changed { - self.set(chat_id, &data).await; + let mut removed = Vec::new(); + let mut evicted_chats = Vec::new(); + for chat_id in candidates { + let lock = self.lock_for(chat_id); + let _guard = lock.lock().await; + let mut data = self.get(chat_id).await; + let before = data.edit_message.len(); + data.edit_message.retain(|key, entry| { + if entry.created_at + ttl_secs > now { + return true; + } + removed.push((chat_id, *key)); + false + }); + if data.edit_message.len() != before { + self.set(chat_id, &data).await; + } + // Chats with no live edit records: evicted from the cache (and + // their per-chat lock) so the cache stays bounded to active + // prompts. The DB keeps the row; the next get() reloads it. + if data.edit_message.is_empty() { + evicted_chats.push(chat_id); + } } if !evicted_chats.is_empty() { let mut cache = self.cache.lock(); @@ -223,4 +228,66 @@ mod tests { "concurrent get→mutate→set must not drop records" ); } + + fn edit_entry(chat_id: i64, created_at: i64) -> EditMessage { + EditMessage { + url: "https://x.com/u/status/1".into(), + chat_id, + forward_message_ids: vec![9], + template: String::new(), + created_at, + } + } + + #[tokio::test] + async fn prune_removes_only_expired_records() { + let dir = tempfile::tempdir().unwrap(); + let pool = crate::db::open_store(dir.path().join("p.db").to_str().unwrap()).unwrap(); + let store = ChatStore::new(pool); + let now = unix_now(); + store + .update(7, |data| { + data.template.insert("t".into(), "[]".into()); + data.edit_message.insert(1, edit_entry(7, now - 3600)); + data.edit_message.insert(2, edit_entry(7, now)); + }) + .await; + + let removed = store.prune_expired(Duration::from_secs(60)).await; + + assert_eq!(removed, vec![(7, 1)]); + let data = store.get(7).await; + assert!(data.edit_message.contains_key(&2), "live record pruned"); + assert_eq!( + data.template.get("t").map(String::as_str), + Some("[]"), + "unrelated state lost by the prune" + ); + } + + #[tokio::test] + async fn prune_eviction_keeps_the_persisted_state() { + // Every record expires → the chat is evicted from the cache; the + // pruned state must already be in the DB when that happens. + let dir = tempfile::tempdir().unwrap(); + let pool = crate::db::open_store(dir.path().join("p.db").to_str().unwrap()).unwrap(); + let store = ChatStore::new(pool); + store + .update(8, |data| { + data.template.insert("keep".into(), "[]".into()); + data.edit_message.insert(1, edit_entry(8, 0)); + }) + .await; + + let removed = store.prune_expired(Duration::from_secs(60)).await; + + assert_eq!(removed, vec![(8, 1)]); + let data = store.get(8).await; + assert!(data.edit_message.is_empty()); + assert_eq!( + data.template.get("keep").map(String::as_str), + Some("[]"), + "eviction dropped state the DB never received" + ); + } }