diff --git a/crates/rustmail/src/commands/add_reminder/common.rs b/crates/rustmail/src/commands/add_reminder/common.rs index b98c87c..3792223 100644 --- a/crates/rustmail/src/commands/add_reminder/common.rs +++ b/crates/rustmail/src/commands/add_reminder/common.rs @@ -5,7 +5,7 @@ use chrono::Local; use serenity::all::{ChannelId, CommandInteraction, Context, GuildId, Message, RoleId, UserId}; use sqlx::SqlitePool; use std::collections::{HashMap, HashSet}; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use std::time::Duration; use tokio::select; use tokio::sync::watch::Receiver; @@ -109,6 +109,9 @@ pub async fn send_register_confirmation_from_command( .await; } +/// One timer per reminder id. +static REMINDER_TIMERS: LazyLock> = LazyLock::new(KeyedTasks::new); + pub fn spawn_reminder( reminder: &Reminder, reminder_id: Option, @@ -123,7 +126,7 @@ pub fn spawn_reminder( let reminder = reminder.clone(); let mut shutdown_rx = shutdown.as_ref().clone(); - tokio::spawn(async move { + REMINDER_TIMERS.spawn(reminder.id, async move { let now = Local::now().timestamp(); let delay_duration = if reminder.data.trigger_time > now { reminder.data.trigger_time - now diff --git a/crates/rustmail/src/commands/delete/common.rs b/crates/rustmail/src/commands/delete/common.rs index ae38811..cb91573 100644 --- a/crates/rustmail/src/commands/delete/common.rs +++ b/crates/rustmail/src/commands/delete/common.rs @@ -82,7 +82,7 @@ pub async fn delete_dm_message(ctx: &Context, user_id: i64, message_ids: &Messag if let Some(dm_msg_id) = &message_ids.dm_message_id && let Ok(msg_id) = dm_msg_id.parse::() { - if let Ok(user) = ctx.http.get_user(UserId::new(user_id as u64)).await { + if let Ok(user) = UserId::new(user_id as u64).to_user(ctx).await { if let Ok(dm_channel) = user.create_dm_channel(&ctx.http).await { match dm_channel.message(&ctx.http, msg_id).await { Ok(dm_message) => { diff --git a/crates/rustmail/src/commands/new_thread/slash_command/new_thread.rs b/crates/rustmail/src/commands/new_thread/slash_command/new_thread.rs index 9e2e931..7d30375 100644 --- a/crates/rustmail/src/commands/new_thread/slash_command/new_thread.rs +++ b/crates/rustmail/src/commands/new_thread/slash_command/new_thread.rs @@ -112,7 +112,7 @@ impl RegistrableCommand for NewThreadCommand { } }; - let user = match ctx.http.get_user(user_id).await { + let user = match user_id.to_user(&ctx).await { Ok(user) => user, Err(_) => return Err(ModmailError::Discord(DiscordError::UserNotFound)), }; diff --git a/crates/rustmail/src/commands/new_thread/text_command/new_thread.rs b/crates/rustmail/src/commands/new_thread/text_command/new_thread.rs index ab4f5fe..e3d6fef 100644 --- a/crates/rustmail/src/commands/new_thread/text_command/new_thread.rs +++ b/crates/rustmail/src/commands/new_thread/text_command/new_thread.rs @@ -28,7 +28,7 @@ pub async fn new_thread( let user_id = user_id.unwrap(); - let user = match ctx.http.get_user(user_id).await { + let user = match user_id.to_user(&ctx).await { Ok(user) => user, Err(_) => return Err(ModmailError::Discord(DiscordError::UserNotFound)), }; diff --git a/crates/rustmail/src/commands/release/slash_command/release.rs b/crates/rustmail/src/commands/release/slash_command/release.rs index 926c441..dbacf60 100644 --- a/crates/rustmail/src/commands/release/slash_command/release.rs +++ b/crates/rustmail/src/commands/release/slash_command/release.rs @@ -97,12 +97,11 @@ impl RegistrableCommand for ReleaseCommand { } }; ticket_status.taken_by = None; - let _ = update_thread_status_db( - &thread.id.to_string(), - &ticket_status, - &db_pool, - ) - .await; + if let Err(e) = + update_thread_status_db(&thread.id, &ticket_status, &db_pool).await + { + eprintln!("Failed to update thread status for {}: {e:?}", thread.id); + } let applied = update_thread_status_ui(&ctx, &ticket_status) .await diff --git a/crates/rustmail/src/commands/release/text_command/release.rs b/crates/rustmail/src/commands/release/text_command/release.rs index c344f64..d471552 100644 --- a/crates/rustmail/src/commands/release/text_command/release.rs +++ b/crates/rustmail/src/commands/release/text_command/release.rs @@ -50,7 +50,10 @@ pub async fn release( } }; ticket_status.taken_by = None; - let _ = update_thread_status_db(&thread.id, &ticket_status, &db_pool).await; + if let Err(e) = update_thread_status_db(&thread.id, &ticket_status, &db_pool).await + { + eprintln!("Failed to update thread status for {}: {e:?}", thread.id); + } let applied = update_thread_status_ui(&ctx, &ticket_status) .await diff --git a/crates/rustmail/src/commands/rename/slash_command/rename.rs b/crates/rustmail/src/commands/rename/slash_command/rename.rs index 03426a0..8086b0b 100644 --- a/crates/rustmail/src/commands/rename/slash_command/rename.rs +++ b/crates/rustmail/src/commands/rename/slash_command/rename.rs @@ -110,7 +110,11 @@ impl RegistrableCommand for RenameCommand { }; ticket_status.label = new_label.clone(); - let _ = update_thread_status_db(&thread.id, &ticket_status, &db_pool).await; + if let Err(e) = + update_thread_status_db(&thread.id, &ticket_status, &db_pool).await + { + eprintln!("Failed to update thread status for {}: {e:?}", thread.id); + } let has_label = new_label.is_some(); let applied = update_thread_status_ui(&ctx, &ticket_status) diff --git a/crates/rustmail/src/commands/rename/text_command/rename.rs b/crates/rustmail/src/commands/rename/text_command/rename.rs index 8027f84..337b721 100644 --- a/crates/rustmail/src/commands/rename/text_command/rename.rs +++ b/crates/rustmail/src/commands/rename/text_command/rename.rs @@ -46,7 +46,9 @@ pub async fn rename_ticket( ticket_status.label = Some(new_label.clone()); } - let _ = update_thread_status_db(&thread.id, &ticket_status, &db_pool).await; + if let Err(e) = update_thread_status_db(&thread.id, &ticket_status, &db_pool).await { + eprintln!("Failed to update thread status for {}: {e:?}", thread.id); + } let has_label = ticket_status.label.is_some(); let applied = update_thread_status_ui(&ctx, &ticket_status) diff --git a/crates/rustmail/src/commands/take/slash_command/take.rs b/crates/rustmail/src/commands/take/slash_command/take.rs index 9256613..b9ee367 100644 --- a/crates/rustmail/src/commands/take/slash_command/take.rs +++ b/crates/rustmail/src/commands/take/slash_command/take.rs @@ -98,7 +98,11 @@ impl RegistrableCommand for TakeCommand { } }; ticket_status.taken_by = Some(command.user.id.to_string()); - let _ = update_thread_status_db(&thread.id, &ticket_status, &db_pool).await; + if let Err(e) = + update_thread_status_db(&thread.id, &ticket_status, &db_pool).await + { + eprintln!("Failed to update thread status for {}: {e:?}", thread.id); + } let applied = update_thread_status_ui(&ctx, &ticket_status) .await diff --git a/crates/rustmail/src/commands/take/text_command/take.rs b/crates/rustmail/src/commands/take/text_command/take.rs index a8c10d6..d756f8b 100644 --- a/crates/rustmail/src/commands/take/text_command/take.rs +++ b/crates/rustmail/src/commands/take/text_command/take.rs @@ -49,7 +49,10 @@ pub async fn take( } }; ticket_status.taken_by = Some(msg.author.id.to_string()); - let _ = update_thread_status_db(&thread.id, &ticket_status, &db_pool).await; + if let Err(e) = update_thread_status_db(&thread.id, &ticket_status, &db_pool).await + { + eprintln!("Failed to update thread status for {}: {e:?}", thread.id); + } let applied = update_thread_status_ui(&ctx, &ticket_status) .await diff --git a/crates/rustmail/src/db/operations/init.rs b/crates/rustmail/src/db/operations/init.rs index 727e29f..f997574 100644 --- a/crates/rustmail/src/db/operations/init.rs +++ b/crates/rustmail/src/db/operations/init.rs @@ -28,7 +28,10 @@ pub async fn init_database() -> Result { .busy_timeout(Duration::from_secs(5)); let pool = SqlitePoolOptions::new() - .max_connections(30) + // SQLite has a single writer; extra connections only add lock contention. + .max_connections(10) + // Fail fast instead of piling up callers for 30s when the pool is exhausted. + .acquire_timeout(Duration::from_secs(10)) .connect_with(connect_options) .await?; diff --git a/crates/rustmail/src/db/operations/messages.rs b/crates/rustmail/src/db/operations/messages.rs index f14b40e..14040e0 100644 --- a/crates/rustmail/src/db/operations/messages.rs +++ b/crates/rustmail/src/db/operations/messages.rs @@ -26,7 +26,7 @@ pub async fn insert_staff_message( let user_id = staff_user_id.get() as i64; let user_name = staff_user_id - .to_user(&ctx.http) + .to_user(ctx) .await .map(|user| user.name) .unwrap_or_else(|_| "Unknown".to_string()); @@ -273,7 +273,7 @@ pub async fn insert_internal_message( let user_name = msg .author .id - .to_user(&ctx.http) + .to_user(ctx) .await .map(|user| user.name) .unwrap_or_else(|_| "Unknown".to_string()); diff --git a/crates/rustmail/src/db/operations/threads.rs b/crates/rustmail/src/db/operations/threads.rs index 45e4653..d846b7c 100644 --- a/crates/rustmail/src/db/operations/threads.rs +++ b/crates/rustmail/src/db/operations/threads.rs @@ -119,7 +119,7 @@ pub async fn create_thread_for_user( let thread_id = Uuid::new_v4().to_string(); let silent_int: i64 = if silent { 1 } else { 0 }; - let res = match sqlx::query( + match sqlx::query( "INSERT INTO threads (id, user_id, user_name, channel_id, silent) VALUES (?, ?, ?, ?, ?)", ) .bind(&thread_id) @@ -127,32 +127,28 @@ pub async fn create_thread_for_user( .bind(user_name) .bind(&channel_id) .bind(silent_int) - .execute(&pool.clone()) + .execute(pool) .await { - Ok(_) => Ok(thread_id.clone()), + Ok(_) => {} Err(Error::Database(db_err)) if db_err.code() == Some(std::borrow::Cow::Borrowed("2067")) => { - if let Some(existing_thread_id) = - sqlx::query_scalar("SELECT id FROM threads WHERE user_id = ? AND status = 1") - .bind(user_id) - .fetch_optional(pool) - .await? - { - Ok(existing_thread_id) - } else { - Err(Error::Database(db_err)) - } + // The user already has an open thread: reuse it. Its status row already exists. + return sqlx::query_scalar("SELECT id FROM threads WHERE user_id = ? AND status = 1") + .bind(user_id) + .fetch_optional(pool) + .await? + .ok_or(Error::Database(db_err)); } - Err(e) => Err(e), - }; + Err(e) => return Err(e), + } let channel_id = channel.id.get() as i64; let user_id_str = user_id.to_string(); let timestamp = Utc::now().timestamp(); - let _ = match sqlx::query!( + if let Err(e) = sqlx::query!( "INSERT INTO thread_status (thread_id, channel_id, owner_id, taken_by, last_message_by, last_message_at) VALUES (?, ?, ?, ?, ?, ?)", thread_id, channel_id, @@ -161,28 +157,14 @@ pub async fn create_thread_for_user( "user", timestamp ) - .execute(&pool.clone()) + .execute(pool) .await { - Ok(_) => Ok(thread_id), - Err(Error::Database(db_err)) - if db_err.code() == Some(std::borrow::Cow::Borrowed("2067")) => - { - if let Some(existing_thread_id) = - sqlx::query_scalar("SELECT id FROM threads WHERE user_id = ? AND status = 1") - .bind(user_id) - .fetch_optional(pool) - .await? - { - Ok(existing_thread_id) - } else { - Err(Error::Database(db_err)) - } - } - Err(e) => Err(e), - }; + // The thread itself is usable; only the channel-name status is missing. + eprintln!("Failed to create status for thread {thread_id}: {e:?}"); + } - res + Ok(thread_id) } pub async fn is_thread_silent(thread_id: &str, pool: &SqlitePool) -> bool { @@ -210,6 +192,9 @@ pub async fn mark_thread_engaged(thread_id: &str, pool: &SqlitePool) { } } +/// Marks an open thread as closed. Returns `false` if it was already closed, so +/// concurrent close paths (command, scheduled closure, channel deletion) don't +/// overwrite each other's closing metadata. pub async fn close_thread( thread_id: &str, closed_by: &str, @@ -217,11 +202,11 @@ pub async fn close_thread( category_name: &str, required_permissions: u64, pool: &SqlitePool, -) -> ModmailResult<()> { +) -> ModmailResult { let closed_at = Utc::now().timestamp(); - let required_permissions = required_permissions.clone().to_string(); + let required_permissions = required_permissions.to_string(); - sqlx::query!( + let res = sqlx::query( r#" UPDATE threads SET @@ -231,19 +216,19 @@ pub async fn close_thread( category_id = ?, category_name = ?, required_permissions = ? - WHERE id = ? + WHERE id = ? AND status = 1 "#, - closed_at, - closed_by, - category_id, - category_name, - required_permissions, - thread_id ) + .bind(closed_at) + .bind(closed_by) + .bind(category_id) + .bind(category_name) + .bind(required_permissions) + .bind(thread_id) .execute(pool) .await?; - Ok(()) + Ok(res.rows_affected() > 0) } pub async fn thread_exists_by_user(user_id: UserId, pool: &SqlitePool) -> bool { @@ -432,26 +417,17 @@ pub async fn allocate_next_message_number( thread_id: &str, pool: &SqlitePool, ) -> Result { - let mut tx = pool.begin().await?; - - let current: Option = - sqlx::query_scalar("SELECT next_message_number FROM threads WHERE id = ?") - .bind(thread_id) - .fetch_optional(&mut *tx) - .await?; - - let num = current.unwrap_or(1); - - sqlx::query!( - "UPDATE threads SET next_message_number = next_message_number + 1 WHERE id = ?", - thread_id + // A single statement is atomic in SQLite. A deferred read-then-write transaction + // fails immediately with SQLITE_BUSY (ignoring busy_timeout) when another + // connection writes between the read and the write. + let current: Option = sqlx::query_scalar( + "UPDATE threads SET next_message_number = next_message_number + 1 WHERE id = ? RETURNING next_message_number - 1", ) - .execute(&mut *tx) + .bind(thread_id) + .fetch_optional(pool) .await?; - tx.commit().await?; - - Ok(num as u64) + Ok(current.unwrap_or(1) as u64) } pub async fn is_orphaned_thread_channel( @@ -580,3 +556,104 @@ pub async fn update_thread_status_db( Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; + use std::collections::HashSet; + use std::time::Duration; + + struct TestDb { + pool: SqlitePool, + path: std::path::PathBuf, + } + + impl TestDb { + async fn cleanup(self) { + self.pool.close().await; + for suffix in ["", "-wal", "-shm"] { + let _ = std::fs::remove_file(format!("{}{suffix}", self.path.display())); + } + } + } + + async fn test_db() -> TestDb { + let path = std::env::temp_dir().join(format!("rustmail-test-{}.sqlite", Uuid::new_v4())); + let options = SqliteConnectOptions::new() + .filename(&path) + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(Duration::from_secs(5)); + let pool = SqlitePoolOptions::new() + .max_connections(10) + .connect_with(options) + .await + .expect("open test database"); + sqlx::migrate!("../../migrations") + .run(&pool) + .await + .expect("run migrations"); + TestDb { pool, path } + } + + async fn insert_open_thread(pool: &SqlitePool, thread_id: &str, user_id: i64) { + sqlx::query( + "INSERT INTO threads (id, user_id, user_name, channel_id, silent) VALUES (?, ?, 'user', ?, 0)", + ) + .bind(thread_id) + .bind(user_id) + .bind(format!("{}", 1000 + user_id)) + .execute(pool) + .await + .expect("insert thread"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn concurrent_message_numbers_are_unique() { + let db = test_db().await; + let pool = db.pool.clone(); + insert_open_thread(&pool, "thread-a", 1).await; + + let tasks: Vec<_> = (0..40) + .map(|_| { + let pool = pool.clone(); + tokio::spawn(async move { allocate_next_message_number("thread-a", &pool).await }) + }) + .collect(); + + let mut numbers = HashSet::new(); + for task in tasks { + let number = task.await.unwrap().expect("allocation must not fail"); + assert!(numbers.insert(number), "duplicate message number {number}"); + } + assert_eq!(numbers, (1..=40).collect::>()); + db.cleanup().await; + } + + #[tokio::test] + async fn closing_twice_only_closes_once() { + let db = test_db().await; + let pool = db.pool.clone(); + insert_open_thread(&pool, "thread-b", 2).await; + + assert!( + close_thread("thread-b", "staff-1", "", "", 0, &pool) + .await + .unwrap() + ); + assert!( + !close_thread("thread-b", "staff-2", "", "", 0, &pool) + .await + .unwrap() + ); + + let closed_by: String = sqlx::query_scalar("SELECT closed_by FROM threads WHERE id = ?") + .bind("thread-b") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(closed_by, "staff-1"); + db.cleanup().await; + } +} diff --git a/crates/rustmail/src/handlers/guild_members_handler.rs b/crates/rustmail/src/handlers/guild_members_handler.rs index e4538dd..51e6343 100644 --- a/crates/rustmail/src/handlers/guild_members_handler.rs +++ b/crates/rustmail/src/handlers/guild_members_handler.rs @@ -141,7 +141,7 @@ impl EventHandler for GuildMembersHandler { _ => 0u64, }; - let _ = close_thread( + if let Err(e) = close_thread( &thread.id, &closed_by, &category_id, @@ -149,6 +149,12 @@ impl EventHandler for GuildMembersHandler { required_permissions, pool, ) - .await; + .await + { + eprintln!( + "Failed to close thread {} after member left: {e:?}", + thread.id + ); + } } } diff --git a/crates/rustmail/src/handlers/ready_handler.rs b/crates/rustmail/src/handlers/ready_handler.rs index 01e2ce7..3f661f9 100644 --- a/crates/rustmail/src/handlers/ready_handler.rs +++ b/crates/rustmail/src/handlers/ready_handler.rs @@ -25,6 +25,17 @@ pub struct ReadyHandler { pub bot_state: Arc>, backfill_started: Arc, update_checker_started: Arc, + startup_tasks_started: Arc, + recovery_running: Arc, +} + +/// Clears an `AtomicBool` when dropped, so the flag is released even if the task panics. +struct FlagGuard(Arc); + +impl Drop for FlagGuard { + fn drop(&mut self) { + self.0.store(false, Ordering::SeqCst); + } } impl ReadyHandler { @@ -41,6 +52,8 @@ impl ReadyHandler { bot_state, backfill_started: Arc::new(AtomicBool::new(false)), update_checker_started: Arc::new(AtomicBool::new(false)), + startup_tasks_started: Arc::new(AtomicBool::new(false)), + recovery_running: Arc::new(AtomicBool::new(false)), } } } @@ -68,16 +81,35 @@ impl EventHandler for ReadyHandler { ctx.set_activity(Option::from(ActivityData::playing(&self.config.bot.status))); + // Discord sends READY again on every new gateway session (invalidated session, + // reconnect...). Long-lived tasks must only be started once per bot run, otherwise + // each reconnect adds another copy of every loop and timer. + let first_ready = self + .startup_tasks_started + .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) + .is_ok(); + + // Message recovery is useful after every reconnect, but must not overlap. + let recovery_guard = self + .recovery_running + .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + .then(|| FlagGuard(self.recovery_running.clone())); + tokio::spawn({ let ctx = ctx.clone(); let config = config.clone(); async move { - let recovery_results = recover_missing_messages(&ctx, &config).await; - send_recovery_summary(&ctx, &config, &recovery_results).await; - sync_features(&ctx, &config).await; - hydrate_scheduled_closures(&ctx, &config).await; - hydrate_pending_category_selections(&ctx, &config).await; + if let Some(_guard) = recovery_guard { + let recovery_results = recover_missing_messages(&ctx, &config).await; + send_recovery_summary(&ctx, &config, &recovery_results).await; + } + if first_ready { + sync_features(&ctx, &config).await; + hydrate_scheduled_closures(&ctx, &config).await; + hydrate_pending_category_selections(&ctx, &config).await; + } } }); @@ -114,9 +146,13 @@ impl EventHandler for ReadyHandler { }); } - load_reminders(&ctx, &self.config, &pool.clone(), self.shutdown.clone()).await; - - update_threads_status(&ctx, &pool.clone()); + if first_ready { + println!("[startup] Starting background tasks"); + load_reminders(&ctx, &self.config, &pool.clone(), self.shutdown.clone()).await; + update_threads_status(&ctx, &pool.clone(), (*self.shutdown).clone()); + } else { + println!("[startup] New gateway session, background tasks already running"); + } let staff_guild_id = GuildId::new(self.config.bot.get_staff_guild_id()); let community_guild_id = GuildId::new(self.config.bot.get_community_guild_id()); @@ -157,36 +193,46 @@ impl EventHandler for ReadyHandler { } } -fn update_threads_status(ctx: &Context, pool: &SqlitePool) { +fn update_threads_status(ctx: &Context, pool: &SqlitePool, mut shutdown: Receiver) { tokio::spawn({ let ctx = ctx.clone(); let pool = pool.clone(); async move { - let mut interval = interval(Duration::from_secs(60 * 10)); + let mut interval = interval(STATUS_REFRESH_INTERVAL); interval.tick().await; loop { + if *shutdown.borrow() { + return; + } + let tickets_status = get_all_thread_status(&pool).await; for ticket in tickets_status.iter() { let ctx = ctx.clone(); let ticket = ticket.clone(); - if let Err(e) = update_thread_status_ui(&ctx, &ticket).await { + if let Err(e) = refresh_thread_status_ui(&ctx, &ticket).await { eprintln!( "Failed to update thread status for channel {}: {:?}", ticket.channel_id, e ); } - tokio::time::sleep(Duration::from_millis(500)).await; + tokio::select! { + _ = tokio::time::sleep(Duration::from_millis(500)) => {} + _ = shutdown.changed() => return, + } } println!("Updated {} ticket statuses", tickets_status.len()); - interval.tick().await; + tokio::select! { + _ = interval.tick() => {} + _ = shutdown.changed() => return, + } } } }); diff --git a/crates/rustmail/src/modules/categories.rs b/crates/rustmail/src/modules/categories.rs index 7b3e591..49f67f0 100644 --- a/crates/rustmail/src/modules/categories.rs +++ b/crates/rustmail/src/modules/categories.rs @@ -11,6 +11,7 @@ use serenity::all::{ }; use serenity::builder::{CreateActionRow, CreateButton}; use std::collections::HashMap; +use std::sync::LazyLock; use tokio::time::{Duration, sleep}; pub const CATEGORY_BUTTON_MAX_PER_ROW: usize = 5; @@ -155,13 +156,16 @@ pub async fn maybe_start_category_selection(ctx: &Context, config: &Config, msg: true } +/// One selection timeout per user id. +static CATEGORY_TIMEOUTS: LazyLock> = LazyLock::new(KeyedTasks::new); + pub fn schedule_category_timeout(ctx: &Context, config: &Config, user_id: i64, expires_at: i64) { let ctx_clone = ctx.clone(); let config_clone = config.clone(); let now = Utc::now().timestamp(); let delay = (expires_at - now).max(0) as u64; - tokio::spawn(async move { + CATEGORY_TIMEOUTS.spawn(user_id, async move { if delay > 0 { sleep(Duration::from_secs(delay)).await; } diff --git a/crates/rustmail/src/modules/reminders.rs b/crates/rustmail/src/modules/reminders.rs index e20f525..b721338 100644 --- a/crates/rustmail/src/modules/reminders.rs +++ b/crates/rustmail/src/modules/reminders.rs @@ -17,7 +17,14 @@ pub async fn load_reminders( }); for reminder in reminders { - spawn_reminder(&reminder, None, ctx, config, pool, shutdown.clone()); + spawn_reminder( + &reminder, + Some(reminder.id), + ctx, + config, + pool, + shutdown.clone(), + ); } println!("All pending reminders have been scheduled."); } diff --git a/crates/rustmail/src/modules/scheduled_closures.rs b/crates/rustmail/src/modules/scheduled_closures.rs index a8cda10..cbd2053 100644 --- a/crates/rustmail/src/modules/scheduled_closures.rs +++ b/crates/rustmail/src/modules/scheduled_closures.rs @@ -3,86 +3,139 @@ use crate::prelude::db::*; use crate::prelude::utils::*; use chrono::Utc; use serenity::all::{ChannelId, Context, UserId}; +use sqlx::SqlitePool; +use std::sync::LazyLock; use tokio::time::{Duration, sleep}; +/// One closure timer per thread id. +static CLOSURE_TIMERS: LazyLock> = LazyLock::new(KeyedTasks::new); + pub fn schedule_one(ctx: &Context, config: &Config, thread_id: String, close_at: i64) { let now = Utc::now().timestamp(); let delay_secs = (close_at - now).max(0) as u64; let ctx_clone = ctx.clone(); let config_clone = config.clone(); - tokio::spawn(async move { + CLOSURE_TIMERS.spawn(thread_id.clone(), async move { if delay_secs > 0 { sleep(Duration::from_secs(delay_secs)).await; } - if let Some(pool) = config_clone.db_pool.as_ref() - && let Ok(Some(current)) = get_scheduled_closure(&thread_id, pool).await - { - if current.close_at <= Utc::now().timestamp() { - if let Some(thread) = get_thread_by_id(&thread_id, pool).await { - let channel_id = ChannelId::new(thread.channel_id.parse::().unwrap_or(0)); - let user_id = UserId::new(thread.user_id as u64); - - let _ = close_thread( - &thread_id, - ¤t.closed_by, - ¤t.category_id, - ¤t.category_name, - current.required_permissions.parse::().unwrap_or(0), - pool, - ) - .await; - let _ = delete_scheduled_closure(&thread_id, pool).await; - - if config_clone.bot.enable_rustmail_logs - && let Some(logs_channel_id) = config_clone.bot.logs_channel_id - { - let base_url = config_clone - .bot - .redirect_url - .trim_end_matches("/api/auth/callback") - .trim_end_matches('/'); - - let panel_url = format!("{}/panel/tickets/{}", base_url, thread.id); - - let mut params = std::collections::HashMap::new(); - params.insert("username".to_string(), thread.user_name.clone()); - params.insert("user_id".to_string(), thread.user_id.to_string()); - params.insert("panel_url".to_string(), panel_url); - - let _ = MessageBuilder::system_message(&ctx_clone, &config_clone) - .translated_content("logs.ticket_closed", Some(¶ms), None, None) - .await - .to_channel(ChannelId::new(logs_channel_id)) - .send(true) - .await; - } - - let effective_silent = - current.silent || is_thread_silent(&thread_id, pool).await; - if !effective_silent { - let _ = MessageBuilder::system_message(&ctx_clone, &config_clone) - .content(&config_clone.bot.close_message) - .to_user(user_id) - .send(true) - .await; - } - let _ = channel_id.delete(&ctx_clone.http).await; - } else { - let _ = delete_scheduled_closure(&thread_id, pool).await; + let Some(pool) = config_clone.db_pool.as_ref() else { + return; + }; + let current = match get_scheduled_closure(&thread_id, pool).await { + Ok(Some(current)) => current, + Ok(None) => return, + Err(e) => { + eprintln!("Failed to load scheduled closure for thread {thread_id}: {e:?}"); + return; + } + }; + + if current.close_at > Utc::now().timestamp() { + schedule_one( + &ctx_clone, + &config_clone, + thread_id.clone(), + current.close_at, + ); + return; + } + + match get_thread_by_id(&thread_id, pool).await { + Some(thread) => { + execute_scheduled_closure(&ctx_clone, &config_clone, &thread, ¤t, pool).await + } + None => { + if let Err(e) = delete_scheduled_closure(&thread_id, pool).await { + eprintln!("Failed to delete scheduled closure for thread {thread_id}: {e:?}"); } - } else { - schedule_one( - &ctx_clone, - &config_clone, - thread_id.clone(), - current.close_at, - ); } } }); } +async fn execute_scheduled_closure( + ctx: &Context, + config: &Config, + thread: &Thread, + closure: &ScheduledClosure, + pool: &SqlitePool, +) { + let closed = match close_thread( + &thread.id, + &closure.closed_by, + &closure.category_id, + &closure.category_name, + closure.required_permissions.parse::().unwrap_or(0), + pool, + ) + .await + { + Ok(closed) => closed, + Err(e) => { + // The closure stays in the database and is retried on the next startup. + eprintln!( + "Failed to close thread {} for scheduled closure: {e:?}", + thread.id + ); + return; + } + }; + + if let Err(e) = delete_scheduled_closure(&thread.id, pool).await { + eprintln!( + "Failed to delete scheduled closure for thread {}: {e:?}", + thread.id + ); + } + + if !closed { + // Already closed by another path, which handled the logs, message and channel. + return; + } + + if config.bot.enable_rustmail_logs + && let Some(logs_channel_id) = config.bot.logs_channel_id + { + let base_url = config + .bot + .redirect_url + .trim_end_matches("/api/auth/callback") + .trim_end_matches('/'); + + let panel_url = format!("{}/panel/tickets/{}", base_url, thread.id); + + let mut params = std::collections::HashMap::new(); + params.insert("username".to_string(), thread.user_name.clone()); + params.insert("user_id".to_string(), thread.user_id.to_string()); + params.insert("panel_url".to_string(), panel_url); + + let _ = MessageBuilder::system_message(ctx, config) + .translated_content("logs.ticket_closed", Some(¶ms), None, None) + .await + .to_channel(ChannelId::new(logs_channel_id)) + .send(true) + .await; + } + + let effective_silent = closure.silent || is_thread_silent(&thread.id, pool).await; + if !effective_silent && thread.user_id > 0 { + let _ = MessageBuilder::system_message(ctx, config) + .content(&config.bot.close_message) + .to_user(UserId::new(thread.user_id as u64)) + .send(true) + .await; + } + + if let Ok(channel_id) = thread.channel_id.parse::() + && channel_id != 0 + { + let _ = ChannelId::new(channel_id).delete(&ctx.http).await; + } +} + +/// Re-arms every scheduled closure after startup; overdue ones fire immediately. pub async fn hydrate_scheduled_closures(ctx: &Context, config: &Config) { let Some(pool) = config.db_pool.as_ref() else { return; @@ -97,59 +150,6 @@ pub async fn hydrate_scheduled_closures(ctx: &Context, config: &Config) { }; for sc in list { - if let Some(thread) = get_thread_by_id(&sc.thread_id, pool).await { - if sc.close_at <= Utc::now().timestamp() { - let channel_id = ChannelId::new(thread.channel_id.parse::().unwrap_or(0)); - let user_id = UserId::new(thread.user_id as u64); - let _ = close_thread( - &thread.id, - &sc.closed_by, - &sc.category_id, - &sc.category_name, - sc.required_permissions.parse::().unwrap_or(0), - pool, - ) - .await; - let _ = delete_scheduled_closure(&thread.id, pool).await; - - if config.bot.enable_rustmail_logs - && let Some(logs_channel_id) = config.bot.logs_channel_id - { - let base_url = config - .bot - .redirect_url - .trim_end_matches("/api/auth/callback") - .trim_end_matches('/'); - - let panel_url = format!("{}/panel/tickets/{}", base_url, thread.id); - - let mut params = std::collections::HashMap::new(); - params.insert("username".to_string(), thread.user_name.clone()); - params.insert("user_id".to_string(), thread.user_id.to_string()); - params.insert("panel_url".to_string(), panel_url); - - let _ = MessageBuilder::system_message(ctx, config) - .translated_content("logs.ticket_closed", Some(¶ms), None, None) - .await - .to_channel(ChannelId::new(logs_channel_id)) - .send(true) - .await; - } - - let effective_silent = sc.silent || is_thread_silent(&thread.id, pool).await; - if !effective_silent { - let _ = MessageBuilder::system_message(ctx, config) - .content(&config.bot.close_message) - .to_user(user_id) - .send(true) - .await; - } - let _ = channel_id.delete(&ctx.http).await; - } else { - schedule_one(ctx, config, sc.thread_id.clone(), sc.close_at); - } - } else { - let _ = delete_scheduled_closure(&sc.thread_id, pool).await; - } + schedule_one(ctx, config, sc.thread_id, sc.close_at); } } diff --git a/crates/rustmail/src/modules/threads.rs b/crates/rustmail/src/modules/threads.rs index fb216ee..6cfb5b1 100644 --- a/crates/rustmail/src/modules/threads.rs +++ b/crates/rustmail/src/modules/threads.rs @@ -28,7 +28,7 @@ pub async fn create_or_get_thread_for_user( } }; - let username = match user_id.to_user(&ctx.http).await { + let username = match user_id.to_user(ctx).await { Ok(u) => u.name.clone(), Err(_) => user_id.get().to_string(), }; @@ -43,12 +43,23 @@ pub async fn create_or_get_thread_for_user( .create_channel(&ctx.http, channel_builder) .await?; - let thread_id = create_thread_for_user(&channel, user_id.get() as i64, &username, false, pool) - .await - .map_err(|e| { + let thread_id = match create_thread_for_user( + &channel, + user_id.get() as i64, + &username, + false, + pool, + ) + .await + { + Ok(id) => id, + Err(e) => { eprintln!("Error creating thread: {}", e); - e - })?; + // Don't leave an orphan channel behind in the inbox category. + let _ = channel.delete(&ctx.http).await; + return Err(e.into()); + } + }; if let Some(cat_id) = ticket_category_id && let Err(e) = set_thread_category(&thread_id, Some(cat_id), pool).await @@ -276,7 +287,7 @@ pub async fn handle_thread_modal_interaction( } }; - let user = match user_id.to_user(&ctx.http).await { + let user = match user_id.to_user(ctx).await { Ok(user) => user, Err(_) => { eprintln!("Failed to fetch user by ID: {}", user_id); diff --git a/crates/rustmail/src/modules/threads_status.rs b/crates/rustmail/src/modules/threads_status.rs index ab42364..2af055c 100644 --- a/crates/rustmail/src/modules/threads_status.rs +++ b/crates/rustmail/src/modules/threads_status.rs @@ -4,16 +4,84 @@ use chrono::Utc; use serenity::all::{ChannelId, UserId}; use serenity::builder::EditChannel; use serenity::client::Context; -use std::time::Duration; +use std::collections::{HashMap, VecDeque}; +use std::sync::{LazyLock, Mutex}; +use std::time::{Duration, Instant}; use tokio::time::timeout; -/// Updates the Discord channel name to reflect the ticket state. -/// -/// Returns `true` if the edit was applied immediately, `false` if it timed out -/// (likely rate-limited) but is still running in the background. -pub async fn update_thread_status_ui(ctx: &Context, ticket: &TicketState) -> ModmailResult { - let channel = ChannelId::new(ticket.channel_id as u64); +/// Discord allows 2 channel renames per 10 minutes per channel. Anything beyond that is +/// held in serenity's ratelimiter, which also delays every other request on the same +/// channel (fetching it, editing it...). We never send a rename we know will be limited. +const RENAME_WINDOW: Duration = Duration::from_secs(10 * 60 + 15); +const RENAMES_PER_WINDOW: usize = 2; + +/// Interval of the periodic status refresh: one rename window, so each refresh finds the +/// slot it used last time free again. +pub const STATUS_REFRESH_INTERVAL: Duration = RENAME_WINDOW; + +#[derive(Default)] +struct RenameState { + recent: VecDeque, + in_flight: bool, + last_name: Option, +} + +static RENAME_STATES: LazyLock>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +/// Reserves a rename slot for `channel_id` if fewer than `max_recent` renames happened in the +/// current window and no rename is already pending. +fn try_reserve_rename(channel_id: u64, max_recent: usize) -> bool { + let mut states = RENAME_STATES.lock().expect("rename state lock poisoned"); + let now = Instant::now(); + if states.len() > 1024 { + states.retain(|_, s| { + s.in_flight + || s.recent + .back() + .is_some_and(|t| now.duration_since(*t) < RENAME_WINDOW) + }); + } + + let state = states.entry(channel_id).or_default(); + while state + .recent + .front() + .is_some_and(|t| now.duration_since(*t) >= RENAME_WINDOW) + { + state.recent.pop_front(); + } + + if state.in_flight || state.recent.len() >= max_recent { + return false; + } + + state.in_flight = true; + state.recent.push_back(now); + true +} + +/// Ends a pending rename; `applied_name` is the name Discord accepted, if any. +fn release_rename(channel_id: u64, applied_name: Option) { + let mut states = RENAME_STATES.lock().expect("rename state lock poisoned"); + if let Some(state) = states.get_mut(&channel_id) { + state.in_flight = false; + if applied_name.is_some() { + state.last_name = applied_name; + } + } +} + +fn is_current_name(channel_id: u64, name: &str) -> bool { + let states = RENAME_STATES.lock().expect("rename state lock poisoned"); + states + .get(&channel_id) + .and_then(|s| s.last_name.as_deref()) + .is_some_and(|n| n == name) +} + +fn build_channel_name(ctx: &Context, ticket: &TicketState, owner_name: &str) -> String { let color = match ticket.last_message_by { TicketAuthor::Staff => "🔵", TicketAuthor::User => "🔴", @@ -27,9 +95,6 @@ pub async fn update_thread_status_ui(ctx: &Context, ticket: &TicketState) -> Mod format!("{}h", minutes / 60) }; - let owner_id = ticket.owner_id.parse().unwrap_or(0); - let owner_name = UserId::new(owner_id).to_user(&ctx.http).await?.name; - let mut staff_name_part = String::new(); if let Some(staff_id) = &ticket.taken_by { let staff_id_u64 = staff_id.parse().unwrap_or(0); @@ -43,7 +108,7 @@ pub async fn update_thread_status_ui(ctx: &Context, ticket: &TicketState) -> Mod } } - let name = if let Some(label) = &ticket.label { + if let Some(label) = &ticket.label { // Fixed parts: color(1 emoji) + "・" + "・" + owner + staff + "・" + time // We truncate the label so the total stays <= 100 chars. let without_label = format!("{color}・・{owner_name}{staff_name_part}・{time_str}"); @@ -56,18 +121,64 @@ pub async fn update_thread_status_ui(ctx: &Context, ticket: &TicketState) -> Mod } } else { format!("{color}・{owner_name}{staff_name_part}・{time_str}") + } +} + +/// Updates the Discord channel name to reflect the ticket state, after a staff action. +/// +/// Returns `true` if the channel shows the new name, `false` if the rename was deferred +/// (rate limited); the periodic status refresh will apply it later. +pub async fn update_thread_status_ui(ctx: &Context, ticket: &TicketState) -> ModmailResult { + apply_thread_status_ui(ctx, ticket, RENAMES_PER_WINDOW).await +} + +/// Periodic variant: only uses one rename slot per window so staff actions keep one. +pub async fn refresh_thread_status_ui(ctx: &Context, ticket: &TicketState) -> ModmailResult { + apply_thread_status_ui(ctx, ticket, RENAMES_PER_WINDOW - 1).await +} + +async fn apply_thread_status_ui( + ctx: &Context, + ticket: &TicketState, + max_recent: usize, +) -> ModmailResult { + let channel_id_u64 = ticket.channel_id as u64; + if channel_id_u64 == 0 { + return Ok(true); + } + let channel = ChannelId::new(channel_id_u64); + + let owner_id: u64 = ticket.owner_id.parse().unwrap_or(0); + let owner_name = if owner_id == 0 { + "unknown".to_string() + } else { + // Passing `ctx` (not `&ctx.http`) lets serenity answer from its cache first. + // GET /users/{id} shares a single ratelimit bucket for the whole bot. + UserId::new(owner_id).to_user(ctx).await?.name }; + let name = build_channel_name(ctx, ticket, &owner_name); + + if is_current_name(channel_id_u64, &name) { + return Ok(true); + } + + if !try_reserve_rename(channel_id_u64, max_recent) { + return Ok(false); + } + let http = ctx.http.clone(); - let channel_id_log = ticket.channel_id; // Spawn the edit so it is never cancelled — even if the caller drops us, - // the request keeps running until Discord replies. + // the request keeps running until Discord replies. The in-flight flag keeps + // a second edit from queuing behind it. let handle = tokio::spawn(async move { - match channel.edit(&http, EditChannel::new().name(&name)).await { + let result = channel.edit(&http, EditChannel::new().name(&name)).await; + release_rename(channel_id_u64, result.is_ok().then_some(name)); + match result { Ok(_) => Ok(()), Err(e) => { - eprintln!("Failed to edit channel {}: {:?}", channel_id_log, e); + eprintln!("Failed to edit channel {}: {:?}", channel_id_u64, e); Err(e) } } @@ -79,18 +190,62 @@ pub async fn update_thread_status_ui(ctx: &Context, ticket: &TicketState) -> Mod Ok(Ok(Ok(_))) => Ok(true), Ok(Ok(Err(e))) => Err(e.into()), Ok(Err(e)) => { - eprintln!( - "Edit task panicked for channel {}: {:?}", - ticket.channel_id, e - ); + release_rename(channel_id_u64, None); + eprintln!("Edit task panicked for channel {}: {:?}", channel_id_u64, e); Ok(true) } Err(_) => { eprintln!( "Timeout editing channel {} — continuing in background", - ticket.channel_id + channel_id_u64 ); Ok(false) } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rename_budget_is_enforced_per_channel() { + let channel = 9_000_000_000_000_001; + + assert!(try_reserve_rename(channel, 2)); + // A rename is still pending: nothing else may queue behind it. + assert!(!try_reserve_rename(channel, 2)); + + release_rename(channel, None); + assert!(try_reserve_rename(channel, 2)); + release_rename(channel, None); + + // Two renames used in the window. + assert!(!try_reserve_rename(channel, 2)); + // Other channels are unaffected. + assert!(try_reserve_rename(channel + 1, 2)); + release_rename(channel + 1, None); + } + + #[test] + fn applied_name_is_remembered() { + let channel = 9_000_000_000_000_201; + + assert!(!is_current_name(channel, "🔴・user・0m")); + assert!(try_reserve_rename(channel, 2)); + release_rename(channel, Some("🔴・user・0m".to_string())); + assert!(is_current_name(channel, "🔴・user・0m")); + assert!(!is_current_name(channel, "🔵・user・0m")); + } + + #[test] + fn periodic_refresh_leaves_a_slot_for_staff_actions() { + let channel = 9_000_000_000_000_101; + + assert!(try_reserve_rename(channel, RENAMES_PER_WINDOW - 1)); + release_rename(channel, None); + assert!(!try_reserve_rename(channel, RENAMES_PER_WINDOW - 1)); + assert!(try_reserve_rename(channel, RENAMES_PER_WINDOW)); + release_rename(channel, None); + } +} diff --git a/crates/rustmail/src/utils/message/message_builder.rs b/crates/rustmail/src/utils/message/message_builder.rs index 32d5de2..6749878 100644 --- a/crates/rustmail/src/utils/message/message_builder.rs +++ b/crates/rustmail/src/utils/message/message_builder.rs @@ -223,7 +223,7 @@ impl<'a> MessageBuilder<'a> { } async fn get_user_avatar_url(&self, user_id: UserId) -> String { - match user_id.to_user(&self.ctx.http).await { + match user_id.to_user(&self.ctx).await { Ok(user) => user .avatar_url() .unwrap_or_else(|| user.default_avatar_url()), diff --git a/crates/rustmail/src/utils/mod.rs b/crates/rustmail/src/utils/mod.rs index 2c41d90..df39c08 100644 --- a/crates/rustmail/src/utils/mod.rs +++ b/crates/rustmail/src/utils/mod.rs @@ -1,11 +1,13 @@ pub mod command; pub mod conversion; pub mod message; +pub mod task; pub mod thread; pub mod time; pub use command::*; pub use conversion::*; pub use message::*; +pub use task::*; pub use thread::*; pub use time::*; diff --git a/crates/rustmail/src/utils/task/keyed_tasks.rs b/crates/rustmail/src/utils/task/keyed_tasks.rs new file mode 100644 index 0000000..2e5ca65 --- /dev/null +++ b/crates/rustmail/src/utils/task/keyed_tasks.rs @@ -0,0 +1,107 @@ +use std::collections::HashMap; +use std::hash::Hash; +use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; +use tokio::task::AbortHandle; + +/// Background tasks indexed by key. Spawning a task for a key that already has one aborts +/// the previous task, so a timer (scheduled closure, reminder...) can never run twice, +/// however many times it gets scheduled. +pub struct KeyedTasks { + tasks: Mutex>, + next_generation: AtomicU64, +} + +impl KeyedTasks { + pub fn new() -> Self { + Self { + tasks: Mutex::new(HashMap::new()), + next_generation: AtomicU64::new(0), + } + } + + pub fn spawn(&'static self, key: K, future: F) + where + F: Future + Send + 'static, + { + let generation = self.next_generation.fetch_add(1, Ordering::Relaxed); + let cleanup_key = key.clone(); + + // The lock is held across the spawn so the task can't finish (and clean up) + // before its entry is inserted. + let mut tasks = self.tasks.lock().expect("keyed tasks lock poisoned"); + let handle = tokio::spawn(async move { + future.await; + self.finish(&cleanup_key, generation); + }); + + if let Some((_, previous)) = tasks.insert(key, (generation, handle.abort_handle())) { + // When a task reschedules itself, this aborts the caller; it has no await + // left at that point, so it simply completes. + previous.abort(); + } + } + + fn finish(&self, key: &K, generation: u64) { + let mut tasks = self.tasks.lock().expect("keyed tasks lock poisoned"); + if tasks.get(key).is_some_and(|(g, _)| *g == generation) { + tasks.remove(key); + } + } + + #[cfg(test)] + fn len(&self) -> usize { + self.tasks.lock().unwrap().len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::LazyLock; + use std::sync::atomic::AtomicUsize; + use std::time::Duration; + + #[tokio::test] + async fn rescheduling_a_key_replaces_the_previous_task() { + static TASKS: LazyLock> = LazyLock::new(KeyedTasks::new); + static RUNS: AtomicUsize = AtomicUsize::new(0); + + for _ in 0..5 { + TASKS.spawn(1, async { + tokio::time::sleep(Duration::from_millis(50)).await; + RUNS.fetch_add(1, Ordering::SeqCst); + }); + } + TASKS.spawn(2, async { + RUNS.fetch_add(1, Ordering::SeqCst); + }); + + tokio::time::sleep(Duration::from_millis(200)).await; + + assert_eq!(RUNS.load(Ordering::SeqCst), 2); + assert_eq!(TASKS.len(), 0); + } + + #[tokio::test] + async fn a_task_can_reschedule_itself() { + static TASKS: LazyLock> = LazyLock::new(KeyedTasks::new); + static RUNS: AtomicUsize = AtomicUsize::new(0); + + fn schedule(remaining: usize) { + TASKS.spawn(7, async move { + tokio::time::sleep(Duration::from_millis(10)).await; + RUNS.fetch_add(1, Ordering::SeqCst); + if remaining > 0 { + schedule(remaining - 1); + } + }); + } + + schedule(2); + tokio::time::sleep(Duration::from_millis(200)).await; + + assert_eq!(RUNS.load(Ordering::SeqCst), 3); + assert_eq!(TASKS.len(), 0); + } +} diff --git a/crates/rustmail/src/utils/task/mod.rs b/crates/rustmail/src/utils/task/mod.rs new file mode 100644 index 0000000..2ca8d00 --- /dev/null +++ b/crates/rustmail/src/utils/task/mod.rs @@ -0,0 +1,3 @@ +pub mod keyed_tasks; + +pub use keyed_tasks::*;