diff --git a/crates/xmedia-bot/src/handlers.rs b/crates/xmedia-bot/src/handlers.rs index 8da3ada..10707c9 100644 --- a/crates/xmedia-bot/src/handlers.rs +++ b/crates/xmedia-bot/src/handlers.rs @@ -238,9 +238,11 @@ async fn execute_command( Command::SetForwardChannel(channel) => { let result = match set_forward_channel_handler(bot, message, channel).await { Ok(channel_id) => { - let mut chat_data = CHAT_STORE.get(message.chat.id.0).await; - chat_data.forward_channel_id = Some(channel_id); - CHAT_STORE.set(message.chat.id.0, &chat_data).await; + CHAT_STORE + .update(message.chat.id.0, |data| { + data.forward_channel_id = Some(channel_id); + }) + .await; "Add successfully.".to_string() } Err(SetForwardChannelError::EmptyParameter) => { @@ -264,31 +266,34 @@ async fn execute_command( } Command::RemoveForwardChannel => { let chat_id = message.chat.id.0; - let mut chat_data = CHAT_STORE.get(chat_id).await; - let text = if chat_data.forward_channel_id.is_some() { - chat_data.forward_channel_id = None; - CHAT_STORE.set(chat_id, &chat_data).await; - "Remove successfully.".to_string() - } else { - "No channel to remove.".to_string() - }; + let text = CHAT_STORE + .update(chat_id, |data| { + if data.forward_channel_id.is_some() { + data.forward_channel_id = None; + "Remove successfully.".to_string() + } else { + "No channel to remove.".to_string() + } + }) + .await; reply(bot.clone(), message.clone(), text).await?; } Command::EditBeforeForward => { let chat_id = message.chat.id.0; - let mut chat_data = CHAT_STORE.get(chat_id).await; - let text = if chat_data.forward_channel_id.is_none() { - "Please enable forward channel first.".to_string() - } else if chat_data.edit_before_forward { - chat_data.edit_before_forward = false; - chat_data.edit_message.clear(); - CHAT_STORE.set(chat_id, &chat_data).await; - "Disable edit before forward.".to_string() - } else { - chat_data.edit_before_forward = true; - CHAT_STORE.set(chat_id, &chat_data).await; - "Enable edit before forward.".to_string() - }; + let text = CHAT_STORE + .update(chat_id, |data| { + if data.forward_channel_id.is_none() { + "Please enable forward channel first.".to_string() + } else if data.edit_before_forward { + data.edit_before_forward = false; + data.edit_message.clear(); + "Disable edit before forward.".to_string() + } else { + data.edit_before_forward = true; + "Enable edit before forward.".to_string() + } + }) + .await; reply(bot.clone(), message.clone(), text).await?; } Command::SetTemplate(name) => { @@ -302,11 +307,14 @@ async fn execute_command( } else if name.is_empty() { "Please provide a name for the template.".to_string() } else { - let mut chat_data = CHAT_STORE.get(chat_id).await; - chat_data - .template - .insert(name, html_escape::encode_text(reply_text).into_owned()); - CHAT_STORE.set(chat_id, &chat_data).await; + CHAT_STORE + .update(chat_id, |data| { + data.template.insert( + name, + html_escape::encode_text(reply_text).into_owned(), + ); + }) + .await; "Template set.".to_string() } } @@ -344,9 +352,11 @@ async fn execute_command( .await?; return Ok(()); } - let mut chat_data = CHAT_STORE.get(chat_id).await; - chat_data.message_format.insert(site.to_string(), format); - CHAT_STORE.set(chat_id, &chat_data).await; + CHAT_STORE + .update(chat_id, |data| { + data.message_format.insert(site.to_string(), format); + }) + .await; reply(bot.clone(), message.clone(), "Format set.").await?; } Command::ClearCache(arg) => { @@ -793,7 +803,7 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<() let chat_id = message.chat().id.0; let prompt_message_id = message.id().0 as i64; let ttl_secs = CONFIG.edit_message_ttl.as_secs() as i64; - let mut chat_data = CHAT_STORE.get(chat_id).await; + let chat_data = CHAT_STORE.get(chat_id).await; let edit = chat_data.edit_message.get(&prompt_message_id).cloned(); let Some(edit) = edit else { log::info!( @@ -807,8 +817,11 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<() }; // Lazy expiry: a stale record (past the TTL, not yet swept) is dropped. if edit.created_at + ttl_secs <= unix_now() { - chat_data.edit_message.remove(&prompt_message_id); - CHAT_STORE.set(chat_id, &chat_data).await; + CHAT_STORE + .update(chat_id, |data| { + data.edit_message.remove(&prompt_message_id); + }) + .await; bot.answer_callback_query(callback_query_id) .text("Expired") .await?; @@ -844,8 +857,11 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<() let _ = bot .delete_message(ChatId(chat_id), MessageId(prompt_message_id as i32)) .await; - chat_data.edit_message.remove(&prompt_message_id); - CHAT_STORE.set(chat_id, &chat_data).await; + CHAT_STORE + .update(chat_id, |data| { + data.edit_message.remove(&prompt_message_id); + }) + .await; } Err(send::SendError::Retryable { delay_seconds, @@ -884,10 +900,13 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<() .caption(template_html) .parse_mode(ParseMode::Html) .await; - if let Some(entry) = chat_data.edit_message.get_mut(&prompt_message_id) { - entry.template = name.to_string(); - } - CHAT_STORE.set(chat_id, &chat_data).await; + CHAT_STORE + .update(chat_id, |data| { + if let Some(entry) = data.edit_message.get_mut(&prompt_message_id) { + entry.template = name.to_string(); + } + }) + .await; log::info!("template '{name}' applied to prompt {prompt_message_id}"); } bot.answer_callback_query(callback_query_id).await?; diff --git a/crates/xmedia-bot/src/send.rs b/crates/xmedia-bot/src/send.rs index da8ba64..655fbce 100644 --- a/crates/xmedia-bot/src/send.rs +++ b/crates/xmedia-bot/src/send.rs @@ -1048,8 +1048,7 @@ pub async fn post_send_actions(bot: &Bot, task: &Task, message_ids: Vec) { }; if edit_before_forward { - let mut chat_data = CHAT_STORE.get(chat_id).await; - let keyboard = build_edit_markup(&chat_data.template); + let keyboard = build_edit_markup(&CHAT_STORE.get(chat_id).await.template); match bot .send_message(ChatId(chat_id), "Reply to edit message.") .reply_markup(keyboard) @@ -1064,17 +1063,22 @@ pub async fn post_send_actions(bot: &Bot, task: &Task, message_ids: Vec) { prompt.id.0, message_ids.len() ); - chat_data.edit_message.insert( - prompt.id.0 as i64, - EditMessage { - url: source_url, - chat_id, - forward_message_ids: message_ids, - template: String::new(), - created_at: unix_now(), - }, - ); - CHAT_STORE.set(chat_id, &chat_data).await; + let prompt_id = prompt.id.0 as i64; + let source_url = source_url.clone(); + CHAT_STORE + .update(chat_id, move |data| { + data.edit_message.insert( + prompt_id, + EditMessage { + url: source_url, + chat_id, + forward_message_ids: message_ids, + template: String::new(), + created_at: unix_now(), + }, + ); + }) + .await; } Err(e) => log::error!("failed to send edit prompt: {e}"), } diff --git a/crates/xmedia-bot/src/state.rs b/crates/xmedia-bot/src/state.rs index 335eb53..372134c 100644 --- a/crates/xmedia-bot/src/state.rs +++ b/crates/xmedia-bot/src/state.rs @@ -6,6 +6,7 @@ use rusqlite::params; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::path::Path; +use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; #[derive(Serialize, Deserialize, Default, Clone, Debug)] @@ -34,6 +35,9 @@ pub struct EditMessage { pub struct ChatStore { /// In-memory cache; the DB is the source of truth on first access. cache: Mutex>, + /// Per-chat async locks serializing get→mutate→set so concurrent handler + /// tasks (batch-forwards, callbacks) cannot clobber each other's writes. + locks: Mutex>>>, db_path: String, } @@ -62,6 +66,7 @@ impl ChatStore { drop(conn); Ok(ChatStore { cache: Mutex::new(HashMap::new()), + locks: Mutex::new(HashMap::new()), db_path: path.to_string(), }) } @@ -111,6 +116,26 @@ impl ChatStore { } } + /// 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 _guard = lock.lock().await; + let mut data = self.get(chat_id).await; + let r = f(&mut data); + self.set(chat_id, &data).await; + r + } + /// Removes edit-before-forward records whose `created_at + ttl` is in the /// past. Returns the removed `(chat_id, prompt_message_id)` pairs so the /// caller can clear the prompt's buttons. @@ -152,3 +177,45 @@ impl ChatStore { removed } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn concurrent_updates_do_not_lose_edit_records() { + let dir = tempfile::tempdir().unwrap(); + let store = std::sync::Arc::new( + ChatStore::open(dir.path().join("s.db").to_str().unwrap()).unwrap(), + ); + let mut handles = Vec::new(); + for i in 0..4 { + let store = Arc::clone(&store); + handles.push(tokio::spawn(async move { + store + .update(1001, |data| { + data.edit_message.insert( + i, + EditMessage { + url: format!("https://x.com/u/status/{i}"), + chat_id: 1001, + forward_message_ids: vec![i], + template: String::new(), + created_at: 0, + }, + ); + }) + .await; + })); + } + for h in handles { + h.await.unwrap(); + } + let data = store.get(1001).await; + assert_eq!( + data.edit_message.len(), + 4, + "concurrent get→mutate→set must not drop records" + ); + } +}