refactor(db): share one DbPool across stores; merge schema init

crates/xmedia-bot/src/db.rs

@@ -113,6 +113,39 @@ pub fn open_db(path: &str) -> rusqlite::Result<Connection> {
     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<Arc<DbPool>> {
+    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 {

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<ChatStore> =
+/// One shared SQLite pool for the three stores (chat state, task queue, link
-    LazyLock::new(|| ChatStore::open("data/task_queue.db").expect("failed to open chat store"));
+/// 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<Arc<db::DbPool>> =
+    LazyLock::new(|| db::open_store("data/task_queue.db").expect("failed to open database"));
+
+pub static CHAT_STORE: LazyLock<ChatStore> = LazyLock::new(|| ChatStore::new(Arc::clone(&DB)));
 pub static TASK_QUEUE: LazyLock<PersistentTaskQueue> =
-    LazyLock::new(|| PersistentTaskQueue::new("data/task_queue.db"));
+    LazyLock::new(|| PersistentTaskQueue::new(Arc::clone(&DB)));
-pub static LINK_CACHE: LazyLock<LinkCache> =
+pub static LINK_CACHE: LazyLock<LinkCache> = LazyLock::new(|| LinkCache::new(Arc::clone(&DB)));
-    LazyLock::new(|| LinkCache::open("data/task_queue.db"));
 pub static CONFIG: LazyLock<Config> = LazyLock::new(Config::load);
 
 #[derive(BotCommands, Clone)]

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<crate::db::DbPool>,
 }
 
 impl LinkCache {
-    pub fn open(db_path: &str) -> Self {
+    /// Wraps the shared DB pool (the `link_cache` table lives in the merged
-        if let Ok(conn) = Connection::open(db_path)
+    /// schema alongside `tasks` and `chat_state`).
-            && let Err(e) = conn.execute_batch(
+    pub fn new(pool: Arc<crate::db::DbPool>) -> Self {
-                "CREATE TABLE IF NOT EXISTS link_cache (url TEXT PRIMARY KEY, \
+        LinkCache { pool }
-                 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),
-        }
     }
 
     /// 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.

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 {
+    /// Wraps the shared DB pool; the schema is initialized once by
-        // Ensure the parent dir and table exist even if only the queue (not
+    /// [`crate::db::open_store`] (all three stores share the pool).
-        // ChatStore) is used — a fresh container without a mounted data dir
+    pub fn new(pool: std::sync::Arc<crate::db::DbPool>) -> Self {
-        // 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}");
-        }
         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();
+            let conn = rusqlite::Connection::open(&path).unwrap();
-            ensure_schema(&conn).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();
+            let conn = rusqlite::Connection::open(queue.pool.path()).unwrap();
-            ensure_schema(&conn).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)",

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<HashMap<i64, Arc<tokio::sync::Mutex<()>>>>,
-    pool: crate::db::DbPool,
+    pool: Arc<crate::db::DbPool>,
 }
 
 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).
+    /// Wraps the shared DB pool (schema initialized once by
-    /// The shared `tasks` / `link_cache` tables are owned by `queue.rs` and
+    /// [`crate::db::open_store`]; the `chat_state` table lives in the merged
-    /// `link_cache.rs` respectively.
+    /// schema alongside `tasks` and `link_cache`).
-    pub fn open(path: &str) -> rusqlite::Result<Self> {
+    pub fn new(pool: Arc<crate::db::DbPool>) -> Self {
-        if let Some(parent) = Path::new(path).parent()
+        ChatStore {
-            && !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 {
             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(
+        let pool = crate::db::open_store(dir.path().join("s.db").to_str().unwrap()).unwrap();
-            ChatStore::open(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);