diff --git a/crates/xmedia-bot/src/db.rs b/crates/xmedia-bot/src/db.rs new file mode 100644 index 0000000..89dcdcb --- /dev/null +++ b/crates/xmedia-bot/src/db.rs @@ -0,0 +1,37 @@ +//! Shared SQLite plumbing for the three tables in `data/task_queue.db` +//! (`tasks` in queue.rs, `chat_state` in state.rs, `link_cache` in +//! link_cache.rs). +//! +//! Every operation opens its own short-lived connection with a busy timeout: +//! handler tasks enqueue while workers lease/update rows concurrently, and +//! without the timeout a concurrent write fails immediately with SQLITE_BUSY +//! and the operation is lost. All I/O runs inside `spawn_blocking` via +//! [`with_conn`] — rusqlite connections are not Send-friendly to hold across +//! an await point, and blocking the async executor stalls every handler. + +use rusqlite::Connection; +use std::time::Duration; + +/// Opens the shared DB with a busy timeout. +pub fn open_db(path: &str) -> rusqlite::Result { + let conn = Connection::open(path)?; + conn.busy_timeout(Duration::from_secs(5))?; + Ok(conn) +} + +/// Runs `f` against a fresh connection on a blocking thread, returning the +/// closure's result. Owns the `spawn_blocking` + `expect` ceremony shared by +/// every table access; the caller maps errors to its own log line. +pub async fn with_conn(path: &str, f: F) -> rusqlite::Result +where + T: Send + 'static, + F: FnOnce(&mut Connection) -> rusqlite::Result + Send + 'static, +{ + let path = path.to_string(); + tokio::task::spawn_blocking(move || { + let mut conn = open_db(&path)?; + f(&mut conn) + }) + .await + .expect("db worker panicked") +} diff --git a/crates/xmedia-bot/src/link_cache.rs b/crates/xmedia-bot/src/link_cache.rs index be79cb4..8a927bd 100644 --- a/crates/xmedia-bot/src/link_cache.rs +++ b/crates/xmedia-bot/src/link_cache.rs @@ -49,12 +49,6 @@ pub struct LinkCache { db_path: String, } -fn open_db(path: &str) -> rusqlite::Result { - let conn = Connection::open(path)?; - conn.busy_timeout(Duration::from_secs(5))?; - Ok(conn) -} - impl LinkCache { pub fn open(db_path: &str) -> Self { if let Ok(conn) = Connection::open(db_path) @@ -73,11 +67,9 @@ impl LinkCache { /// Returns the cached post if present and not expired; a stale entry is /// removed on the spot. pub async fn get(&self, key: &str, ttl: Duration) -> Option { - let db_path = self.db_path.clone(); let key = key.to_string(); let ttl = ttl.as_secs_f64(); - tokio::task::spawn_blocking(move || -> rusqlite::Result> { - let conn = open_db(&db_path)?; + let result = crate::db::with_conn(&self.db_path, move |conn| { let mut stmt = conn.prepare("SELECT payload, created_at FROM link_cache WHERE url = ?1")?; let mut rows = stmt.query(params![key])?; @@ -90,66 +82,66 @@ impl LinkCache { conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?; return Ok(None); } - serde_json::from_str(&payload).map(Some).map_err(|e| { - rusqlite::Error::ToSqlConversionFailure(Box::new(e)) - }) - }) - .await - .expect("link cache read worker panicked") - .unwrap_or_else(|e| { - log::error!("link cache read failed: {e}"); - None + Ok(Some(serde_json::from_str::(&payload).map_err( + |e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)), + )?)) }) + .await; + match result { + Ok(v) => v, + Err(e) => { + log::error!("link cache read failed: {e}"); + None + } + } } pub async fn put(&self, key: &str, post: &CachedPost) { - let db_path = self.db_path.clone(); let key = key.to_string(); let payload = serde_json::to_string(post).expect("cached post serializes"); - tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = open_db(&db_path)?; + let result = crate::db::with_conn(&self.db_path, move |conn| { conn.execute( "INSERT OR REPLACE INTO link_cache (url, payload, created_at) VALUES (?1, ?2, ?3)", params![key, payload, now_f64()], )?; Ok(()) }) - .await - .expect("link cache write worker panicked") - .unwrap_or_else(|e| log::error!("link cache write failed: {e}")); + .await; + if let Err(e) = result { + log::error!("link cache write failed: {e}"); + } } /// Drops an entry (e.g. a cached file id that turned out invalid). pub async fn remove(&self, key: &str) { - let db_path = self.db_path.clone(); let key = key.to_string(); - tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = open_db(&db_path)?; + let result = crate::db::with_conn(&self.db_path, move |conn| { conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?; Ok(()) }) - .await - .expect("link cache delete worker panicked") - .unwrap_or_else(|e| log::error!("link cache delete failed: {e}")); + .await; + if let Err(e) = result { + log::error!("link cache delete failed: {e}"); + } } /// Removes expired entries; returns how many were deleted. pub async fn prune(&self, ttl: Duration) -> usize { - let db_path = self.db_path.clone(); let cutoff = now_f64() - ttl.as_secs_f64(); - tokio::task::spawn_blocking(move || -> rusqlite::Result { - let conn = open_db(&db_path)?; + let result = crate::db::with_conn(&self.db_path, move |conn| { conn.execute( "DELETE FROM link_cache WHERE created_at < ?1", params![cutoff], ) }) - .await - .expect("link cache prune worker panicked") - .unwrap_or_else(|e| { - log::error!("link cache prune failed: {e}"); - 0 - }) + .await; + match result { + Ok(n) => n, + Err(e) => { + log::error!("link cache prune failed: {e}"); + 0 + } + } } } diff --git a/crates/xmedia-bot/src/main.rs b/crates/xmedia-bot/src/main.rs index 3e52d4f..aa4b271 100644 --- a/crates/xmedia-bot/src/main.rs +++ b/crates/xmedia-bot/src/main.rs @@ -8,6 +8,7 @@ use tokio::sync::watch; use x_media::site; mod config; +mod db; mod handlers; mod link_cache; mod photo; diff --git a/crates/xmedia-bot/src/queue.rs b/crates/xmedia-bot/src/queue.rs index cefe025..d03d6fc 100644 --- a/crates/xmedia-bot/src/queue.rs +++ b/crates/xmedia-bot/src/queue.rs @@ -6,11 +6,11 @@ //! replaced by dedicated columns. use parking_lot::Mutex; -use rusqlite::{params, Connection, TransactionBehavior}; +use rusqlite::{Connection, TransactionBehavior, params}; use serde_json::Value; use std::pin::Pin; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::Notify; use tokio::task::JoinHandle; @@ -29,15 +29,9 @@ const QUEUE_WORKERS: usize = 4; pub enum QueueError { /// Reschedule with the given delay; after `MAX_RETRIES` attempts the task /// is dead-lettered instead. - Retryable { - delay_seconds: f64, - payload: Value, - }, + Retryable { delay_seconds: f64, payload: Value }, /// Give up now. - Permanent { - message: String, - payload: Value, - }, + Permanent { message: String, payload: Value }, } type BoxFuture<'a, T> = Pin + Send + 'a>>; @@ -74,16 +68,7 @@ fn now_f64() -> f64 { .unwrap_or(0.0) } -/// Opens the queue DB with a busy timeout. Handler tasks enqueue while -/// workers lease/update rows concurrently; without the timeout a concurrent -/// write fails immediately with SQLITE_BUSY and the operation is lost. -fn open_db(path: &str) -> rusqlite::Result { - let conn = Connection::open(path)?; - conn.busy_timeout(Duration::from_secs(5))?; - Ok(conn) -} - -fn ensure_schema(conn: &Connection) -> rusqlite::Result<()> { +fn ensure_schema(conn: &rusqlite::Connection) -> rusqlite::Result<()> { conn.execute_batch( "CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, payload TEXT NOT NULL, \ run_after REAL NOT NULL, attempts INTEGER NOT NULL, status TEXT NOT NULL, \ @@ -162,10 +147,8 @@ impl PersistentTaskQueue { self.counter.fetch_add(1, Ordering::Relaxed) ); let payload = payload.to_string(); - let db_path = self.db_path.clone(); log::info!("enqueued {id} (run_after {run_after:.1})"); - tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = open_db(&db_path)?; + crate::db::with_conn(&self.db_path, move |conn| { conn.execute( "INSERT OR REPLACE INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \ VALUES (?1, ?2, ?3, 0, 'pending', 0, ?4)", @@ -173,8 +156,7 @@ impl PersistentTaskQueue { )?; Ok(()) }) - .await - .expect("queue insert worker panicked")?; + .await?; // Wake every sleeping worker: with several workers the one that finds // nothing due must not starve the newly inserted row. self.notify.notify_waiters(); @@ -182,18 +164,17 @@ impl PersistentTaskQueue { } async fn recover_stale(&self) { - let db_path = self.db_path.clone(); - tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = open_db(&db_path)?; + let result = crate::db::with_conn(&self.db_path, move |conn| { conn.execute( "UPDATE tasks SET status='pending', locked_until=0 WHERE status='in_progress' AND locked_until < ?1", params![now_f64()], )?; Ok(()) }) - .await - .expect("queue recovery worker panicked") - .unwrap_or_else(|e| log::error!("queue recovery failed: {e}")); + .await; + if let Err(e) = result { + log::error!("queue recovery failed: {e}"); + } } } @@ -225,9 +206,7 @@ impl QueueWorker { /// Leases the oldest due row (sets it `in_progress` with a lock TTL). async fn lease_next(&self) -> Option { - let db_path = self.db_path.clone(); - tokio::task::spawn_blocking(move || -> rusqlite::Result> { - let mut conn = open_db(&db_path)?; + let result = crate::db::with_conn(&self.db_path, |conn| { // BEGIN IMMEDIATE: with several workers, a deferred transaction // that read before another worker's lease commit would fail with // SQLITE_BUSY_SNAPSHOT. Taking the write lock up front serializes @@ -265,31 +244,34 @@ impl QueueWorker { attempts, })) }) - .await - .expect("queue lease worker panicked") - .unwrap_or_else(|e| { - log::error!("queue lease failed: {e}"); - None - }) + .await; + match result { + Ok(row) => row, + Err(e) => { + log::error!("queue lease failed: {e}"); + None + } + } } async fn earliest_run_after(&self) -> Option { - let db_path = self.db_path.clone(); - tokio::task::spawn_blocking(move || -> rusqlite::Result> { - let conn = open_db(&db_path)?; - let mut stmt = conn.prepare("SELECT MIN(run_after) FROM tasks WHERE status='pending'")?; + let result = crate::db::with_conn(&self.db_path, |conn| { + let mut stmt = + conn.prepare("SELECT MIN(run_after) FROM tasks WHERE status='pending'")?; let mut rows = stmt.query([])?; match rows.next()? { Some(row) => Ok(row.get::<_, Option>(0)?), None => Ok(None), } }) - .await - .expect("queue timing worker panicked") - .unwrap_or_else(|e| { - log::error!("queue timing query failed: {e}"); - None - }) + .await; + match result { + Ok(v) => v, + Err(e) => { + log::error!("queue timing query failed: {e}"); + None + } + } } async fn process(&self, row: LeasedRow) { @@ -336,33 +318,31 @@ impl QueueWorker { } async fn delete_row(&self, id: &str) { - let db_path = self.db_path.clone(); let id = id.to_string(); - tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = open_db(&db_path)?; + let result = crate::db::with_conn(&self.db_path, move |conn| { conn.execute("DELETE FROM tasks WHERE id = ?1", params![id])?; Ok(()) }) - .await - .expect("queue delete worker panicked") - .unwrap_or_else(|e| log::error!("queue delete failed: {e}")); + .await; + if let Err(e) = result { + log::error!("queue delete failed: {e}"); + } } async fn reschedule(&self, id: &str, payload: Value, delay_seconds: f64, attempts: i32) { - let db_path = self.db_path.clone(); let id = id.to_string(); let payload = payload.to_string(); - tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = open_db(&db_path)?; + let result = crate::db::with_conn(&self.db_path, move |conn| { conn.execute( "UPDATE tasks SET payload=?1, run_after=?2, attempts=?3, status='pending', locked_until=0 WHERE id=?4", params![payload, now_f64() + delay_seconds, attempts, id], )?; Ok(()) }) - .await - .expect("queue reschedule worker panicked") - .unwrap_or_else(|e| log::error!("queue reschedule failed: {e}")); + .await; + if let Err(e) = result { + log::error!("queue reschedule failed: {e}"); + } self.notify.notify_waiters(); } } diff --git a/crates/xmedia-bot/src/state.rs b/crates/xmedia-bot/src/state.rs index 2318da8..335eb53 100644 --- a/crates/xmedia-bot/src/state.rs +++ b/crates/xmedia-bot/src/state.rs @@ -2,7 +2,7 @@ //! `data/task_queue.db`, shared with the task queue). use parking_lot::Mutex; -use rusqlite::{params, Connection}; +use rusqlite::params; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::path::Path; @@ -45,7 +45,9 @@ pub fn unix_now() -> i64 { } impl ChatStore { - /// Creates the parent directory and both tables (idempotent). + /// Creates the parent directory and the `chat_state` table (idempotent). + /// The shared `tasks` / `link_cache` tables are owned by `queue.rs` and + /// `link_cache.rs` respectively. pub fn open(path: &str) -> rusqlite::Result { if let Some(parent) = Path::new(path).parent() && !parent.as_os_str().is_empty() @@ -53,12 +55,9 @@ impl ChatStore { std::fs::create_dir_all(parent) .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?; } - let conn = Connection::open(path)?; + let conn = crate::db::open_db(path)?; conn.execute_batch( - "CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, payload TEXT NOT NULL, \ - run_after REAL NOT NULL, attempts INTEGER NOT NULL, status TEXT NOT NULL, \ - locked_until REAL NOT NULL, created_at REAL NOT NULL); \ - CREATE TABLE IF NOT EXISTS chat_state (chat_id TEXT PRIMARY KEY, payload TEXT NOT NULL);", + "CREATE TABLE IF NOT EXISTS chat_state (chat_id TEXT PRIMARY KEY, payload TEXT NOT NULL);", )?; drop(conn); Ok(ChatStore { @@ -71,22 +70,19 @@ impl ChatStore { if let Some(data) = self.cache.lock().get(&chat_id) { return data.clone(); } - let db_path = self.db_path.clone(); - let payload = tokio::task::spawn_blocking(move || -> rusqlite::Result> { - let conn = Connection::open(&db_path)?; + let chat_key = chat_id.to_string(); + let payload = crate::db::with_conn(&self.db_path, move |conn| { // Concurrent handler tasks (batch-forwards) may write chat_state - // while this read runs; without a busy timeout a write lock - // collision fails the query immediately. - conn.busy_timeout(std::time::Duration::from_secs(5))?; + // while this read runs; the shared busy timeout handles the + // write-lock collision instead of failing the query. let mut stmt = conn.prepare("SELECT payload FROM chat_state WHERE chat_id = ?1")?; - let mut rows = stmt.query(params![chat_id.to_string()])?; + let mut rows = stmt.query(params![chat_key])?; match rows.next()? { - Some(row) => Ok(Some(row.get(0)?)), + Some(row) => Ok(Some(row.get::<_, String>(0)?)), None => Ok(None), } }) .await - .expect("chat_state worker panicked") .unwrap_or_else(|e| { log::error!("chat_state read failed: {e}"); None @@ -101,19 +97,18 @@ impl ChatStore { pub async fn set(&self, chat_id: i64, data: &ChatData) { self.cache.lock().insert(chat_id, data.clone()); let payload = serde_json::to_string(data).expect("chat state serializes"); - let db_path = self.db_path.clone(); - tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = Connection::open(&db_path)?; - conn.busy_timeout(std::time::Duration::from_secs(5))?; + let chat_id = chat_id.to_string(); + let result = crate::db::with_conn(&self.db_path, move |conn| { conn.execute( "INSERT OR REPLACE INTO chat_state (chat_id, payload) VALUES (?1, ?2)", - params![chat_id.to_string(), payload], + params![chat_id, payload], )?; Ok(()) }) - .await - .expect("chat_state worker panicked") - .unwrap_or_else(|e| log::error!("chat_state write failed: {e}")); + .await; + if let Err(e) = result { + log::error!("chat_state write failed: {e}"); + } } /// Removes edit-before-forward records whose `created_at + ttl` is in the @@ -149,7 +144,10 @@ impl ChatStore { self.set(chat_id, &data).await; } if !removed.is_empty() { - log::info!("pruned {} expired edit-before-forward record(s)", removed.len()); + log::info!( + "pruned {} expired edit-before-forward record(s)", + removed.len() + ); } removed }