From 1e77bb04785d9439fe72d6caff3dceb062dfae5f Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Fri, 14 Aug 2026 19:36:41 +0800 Subject: [PATCH] refactor(db): share one DbPool across stores; merge schema init --- crates/xmedia-bot/src/db.rs | 33 ++++++++++++++++++++ crates/xmedia-bot/src/handlers.rs | 19 +++++++----- crates/xmedia-bot/src/link_cache.rs | 43 +++++++++++++------------- crates/xmedia-bot/src/queue.rs | 47 ++++++++--------------------- crates/xmedia-bot/src/state.rs | 33 ++++++-------------- 5 files changed, 90 insertions(+), 85 deletions(-) diff --git a/crates/xmedia-bot/src/db.rs b/crates/xmedia-bot/src/db.rs index 84dd5de..4ca2306 100644 --- a/crates/xmedia-bot/src/db.rs +++ b/crates/xmedia-bot/src/db.rs @@ -113,6 +113,39 @@ pub fn open_db(path: &str) -> rusqlite::Result { Ok(conn) } +/// Opens the shared DB file, runs the merged schema for all three tables and +/// returns a pool for it. One call per process in production (the stores +/// share the returned pool); tests call it per tempdir. +pub fn open_store(path: &str) -> rusqlite::Result> { + if let Some(parent) = std::path::Path::new(path).parent() + && !parent.as_os_str().is_empty() + { + std::fs::create_dir_all(parent).map_err(rusqlite_error)?; + } + let conn = open_db(path)?; + schema_init(&conn)?; + Ok(Arc::new(DbPool::new(path))) +} + +fn rusqlite_error(e: std::io::Error) -> rusqlite::Error { + rusqlite::Error::ToSqlConversionFailure(Box::new(e)) +} + +/// Creates the `tasks`, `chat_state` and `link_cache` tables (idempotent). +/// The three stores used to own their own schema; keeping it in one place +/// means one initialization for the whole database file. +pub fn schema_init(conn: &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, \ + locked_until REAL NOT NULL, created_at REAL NOT NULL); \ + CREATE INDEX IF NOT EXISTS idx_tasks_pending ON tasks(status, run_after); \ + CREATE TABLE IF NOT EXISTS chat_state (chat_id TEXT PRIMARY KEY, payload TEXT NOT NULL); \ + CREATE TABLE IF NOT EXISTS link_cache (url TEXT PRIMARY KEY, payload TEXT NOT NULL, \ + created_at REAL NOT NULL);", + ) +} + /// Unix timestamp in fractional seconds. Shared by the queue, chat store and /// link cache (previously four private copies). pub fn now_f64() -> f64 { diff --git a/crates/xmedia-bot/src/handlers.rs b/crates/xmedia-bot/src/handlers.rs index 63df14f..5a09b91 100644 --- a/crates/xmedia-bot/src/handlers.rs +++ b/crates/xmedia-bot/src/handlers.rs @@ -1,11 +1,11 @@ use crate::config::Config; -use crate::db::now_f64; +use crate::db::{self, now_f64}; use crate::link_cache::{CachedMediaKind, CachedPost, LinkCache}; use crate::queue::PersistentTaskQueue; use crate::send::{self, MediaItemPayload, Task}; use crate::state::{ChatData, ChatStore, unix_now}; use std::collections::HashSet; -use std::sync::LazyLock; +use std::sync::{Arc, LazyLock}; use teloxide::RequestError; use teloxide::prelude::*; use teloxide::types::{ @@ -79,12 +79,17 @@ pub async fn stop_url_workers() { } } -pub static CHAT_STORE: LazyLock = - LazyLock::new(|| ChatStore::open("data/task_queue.db").expect("failed to open chat store")); +/// One shared SQLite pool for the three stores (chat state, task queue, link +/// cache): a single pool bounds concurrent DB work on `data/task_queue.db` +/// instead of three independent pools competing for the same file. The schema +/// for all three tables is initialized once, here. +static DB: LazyLock> = + LazyLock::new(|| db::open_store("data/task_queue.db").expect("failed to open database")); + +pub static CHAT_STORE: LazyLock = LazyLock::new(|| ChatStore::new(Arc::clone(&DB))); pub static TASK_QUEUE: LazyLock = - LazyLock::new(|| PersistentTaskQueue::new("data/task_queue.db")); -pub static LINK_CACHE: LazyLock = - LazyLock::new(|| LinkCache::open("data/task_queue.db")); + LazyLock::new(|| PersistentTaskQueue::new(Arc::clone(&DB))); +pub static LINK_CACHE: LazyLock = LazyLock::new(|| LinkCache::new(Arc::clone(&DB))); pub static CONFIG: LazyLock = LazyLock::new(Config::load); #[derive(BotCommands, Clone)] diff --git a/crates/xmedia-bot/src/link_cache.rs b/crates/xmedia-bot/src/link_cache.rs index 839d24e..0e55479 100644 --- a/crates/xmedia-bot/src/link_cache.rs +++ b/crates/xmedia-bot/src/link_cache.rs @@ -9,8 +9,9 @@ //! by the periodic prune in `main`. use crate::db::now_f64; -use rusqlite::{Connection, params}; +use rusqlite::params; use serde::{Deserialize, Serialize}; +use std::sync::Arc; use std::time::Duration; #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] @@ -45,24 +46,16 @@ pub struct CachedPost { } /// SQLite-backed cache sharing `data/task_queue.db` with the queue and chat -/// state (same `open_db` pattern: busy timeout, `spawn_blocking` I/O). +/// state (same shared pool, see [`crate::db::open_store`]). pub struct LinkCache { - pool: crate::db::DbPool, + pool: Arc, } impl LinkCache { - pub fn open(db_path: &str) -> Self { - if let Ok(conn) = Connection::open(db_path) - && let Err(e) = conn.execute_batch( - "CREATE TABLE IF NOT EXISTS link_cache (url TEXT PRIMARY KEY, \ - payload TEXT NOT NULL, created_at REAL NOT NULL);", - ) - { - log::error!("failed to initialize link cache schema: {e}"); - } - Self { - pool: crate::db::DbPool::new(db_path), - } + /// Wraps the shared DB pool (the `link_cache` table lives in the merged + /// schema alongside `tasks` and `chat_state`). + pub fn new(pool: Arc) -> Self { + LinkCache { pool } } /// Returns the cached post if present and not expired; a stale entry is @@ -197,7 +190,9 @@ mod tests { #[tokio::test] async fn put_get_roundtrip() { let dir = tempfile::tempdir().unwrap(); - let cache = LinkCache::open(dir.path().join("c.db").to_str().unwrap()); + let cache = LinkCache::new( + crate::db::open_store(dir.path().join("c.db").to_str().unwrap()).unwrap(), + ); cache.put("twitter:1", &entry()).await; let got = cache.get("twitter:1", Duration::from_secs(3600)).await; assert!(got.is_some()); @@ -209,11 +204,13 @@ mod tests { #[tokio::test] async fn expired_entry_removed_on_read() { let dir = tempfile::tempdir().unwrap(); - let cache = LinkCache::open(dir.path().join("c.db").to_str().unwrap()); + let cache = LinkCache::new( + crate::db::open_store(dir.path().join("c.db").to_str().unwrap()).unwrap(), + ); cache.put("twitter:1", &entry()).await; // Force the row into the past so a 1s TTL expires it. { - let conn = Connection::open(dir.path().join("c.db")).unwrap(); + let conn = rusqlite::Connection::open(dir.path().join("c.db")).unwrap(); conn.execute("UPDATE link_cache SET created_at = created_at - 100", []) .unwrap(); } @@ -234,7 +231,9 @@ mod tests { #[tokio::test] async fn remove_and_prune() { let dir = tempfile::tempdir().unwrap(); - let cache = LinkCache::open(dir.path().join("c.db").to_str().unwrap()); + let cache = LinkCache::new( + crate::db::open_store(dir.path().join("c.db").to_str().unwrap()).unwrap(), + ); cache.put("twitter:1", &entry()).await; cache.put("pixiv:2", &entry()).await; cache.remove("twitter:1").await; @@ -251,7 +250,7 @@ mod tests { .is_some() ); { - let conn = Connection::open(dir.path().join("c.db")).unwrap(); + let conn = rusqlite::Connection::open(dir.path().join("c.db")).unwrap(); conn.execute("UPDATE link_cache SET created_at = created_at - 100", []) .unwrap(); } @@ -267,7 +266,9 @@ mod tests { #[tokio::test] async fn clear_one_entry_or_all() { let dir = tempfile::tempdir().unwrap(); - let cache = LinkCache::open(dir.path().join("c.db").to_str().unwrap()); + let cache = LinkCache::new( + crate::db::open_store(dir.path().join("c.db").to_str().unwrap()).unwrap(), + ); cache.put("twitter:1", &entry()).await; cache.put("pixiv:2", &entry()).await; // By key: only the matching row is removed. diff --git a/crates/xmedia-bot/src/queue.rs b/crates/xmedia-bot/src/queue.rs index b1b925e..6576486 100644 --- a/crates/xmedia-bot/src/queue.rs +++ b/crates/xmedia-bot/src/queue.rs @@ -7,7 +7,7 @@ use crate::db::now_f64; use parking_lot::Mutex; -use rusqlite::{Connection, TransactionBehavior, params}; +use rusqlite::{TransactionBehavior, params}; use serde_json::Value; use std::pin::Pin; use std::sync::Arc; @@ -81,34 +81,12 @@ fn scaled_retry_delay(base: f64, attempts: i32) -> f64 { (base * 2f64.powi(attempts)).min(300.0) } -fn ensure_schema(conn: &rusqlite::Connection) -> rusqlite::Result<()> { - conn.execute_batch( - "PRAGMA journal_mode=WAL; \ - 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 INDEX IF NOT EXISTS idx_tasks_pending ON tasks(status, run_after);", - ) -} - impl PersistentTaskQueue { - pub fn new(db_path: &str) -> Self { - // Ensure the parent dir and table exist even if only the queue (not - // ChatStore) is used — a fresh container without a mounted data dir - // must still be able to open the DB. - if let Some(parent) = std::path::Path::new(db_path).parent() - && !parent.as_os_str().is_empty() - && let Err(e) = std::fs::create_dir_all(parent) - { - log::error!("failed to create queue dir: {e}"); - } - if let Ok(conn) = Connection::open(db_path) - && let Err(e) = ensure_schema(&conn) - { - log::error!("failed to initialize queue schema: {e}"); - } + /// Wraps the shared DB pool; the schema is initialized once by + /// [`crate::db::open_store`] (all three stores share the pool). + pub fn new(pool: std::sync::Arc) -> Self { Self { - pool: std::sync::Arc::new(crate::db::DbPool::new(db_path)), + pool, notify: Arc::new(Notify::new()), stop: Arc::new(AtomicBool::new(false)), worker: Mutex::new(Vec::new()), @@ -424,7 +402,8 @@ mod tests { async fn new_queue() -> (PersistentTaskQueue, tempfile::TempDir) { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("queue.db"); - let queue = PersistentTaskQueue::new(path.to_str().unwrap()); + let pool = crate::db::open_store(path.to_str().unwrap()).unwrap(); + let queue = PersistentTaskQueue::new(pool); (queue, dir) } @@ -531,10 +510,11 @@ mod tests { async fn stale_in_progress_row_is_recovered_on_start() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("queue.db"); - // Insert a stale leased row directly (lease expired). + // Insert a stale leased row directly (lease expired). open_store runs + // the schema; the queue below shares the same pool. + let pool = crate::db::open_store(path.to_str().unwrap()).unwrap(); { - let conn = Connection::open(&path).unwrap(); - ensure_schema(&conn).unwrap(); + let conn = rusqlite::Connection::open(&path).unwrap(); conn.execute( "INSERT INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \ VALUES ('task_stale', '{\"s\":1}', 0, 0, 'in_progress', ?1, 0)", @@ -542,7 +522,7 @@ mod tests { ) .unwrap(); } - let queue = PersistentTaskQueue::new(path.to_str().unwrap()); + let queue = PersistentTaskQueue::new(pool); let calls = Arc::new(AtomicUsize::new(0)); let c = calls.clone(); queue @@ -578,8 +558,7 @@ mod tests { // Insert a stale leased row AFTER startup: without a runtime sweep it // would stay `in_progress` forever (only start() used to recover). { - let conn = Connection::open(queue.pool.path()).unwrap(); - ensure_schema(&conn).unwrap(); + let conn = rusqlite::Connection::open(queue.pool.path()).unwrap(); conn.execute( "INSERT INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \ VALUES ('task_stale_runtime', '{\"s\":1}', 0, 0, 'in_progress', ?1, 0)", diff --git a/crates/xmedia-bot/src/state.rs b/crates/xmedia-bot/src/state.rs index bbc8543..221eadf 100644 --- a/crates/xmedia-bot/src/state.rs +++ b/crates/xmedia-bot/src/state.rs @@ -5,7 +5,6 @@ use parking_lot::Mutex; 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}; @@ -38,7 +37,7 @@ pub struct ChatStore { /// Per-chat async locks serializing get→mutate→set so concurrent handler /// tasks (batch-forwards, callbacks) cannot clobber each other's writes. locks: Mutex>>>, - pool: crate::db::DbPool, + pool: Arc, } pub fn unix_now() -> i64 { @@ -49,26 +48,15 @@ pub fn unix_now() -> i64 { } impl ChatStore { - /// 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() - { - std::fs::create_dir_all(parent) - .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?; - } - let conn = crate::db::open_db(path)?; - conn.execute_batch( - "CREATE TABLE IF NOT EXISTS chat_state (chat_id TEXT PRIMARY KEY, payload TEXT NOT NULL);", - )?; - drop(conn); - Ok(ChatStore { + /// Wraps the shared DB pool (schema initialized once by + /// [`crate::db::open_store`]; the `chat_state` table lives in the merged + /// schema alongside `tasks` and `link_cache`). + pub fn new(pool: Arc) -> Self { + ChatStore { cache: Mutex::new(HashMap::new()), locks: Mutex::new(HashMap::new()), - pool: crate::db::DbPool::new(path), - }) + pool, + } } pub async fn get(&self, chat_id: i64) -> ChatData { @@ -209,9 +197,8 @@ mod tests { #[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 pool = crate::db::open_store(dir.path().join("s.db").to_str().unwrap()).unwrap(); + let store = std::sync::Arc::new(ChatStore::new(pool)); let mut handles = Vec::new(); for i in 0..4 { let store = Arc::clone(&store);