From edb32c23b49fac3d95f360d318d0a9fd465752ef Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Thu, 13 Aug 2026 22:12:34 +0800 Subject: [PATCH] perf(db): reuse SQLite connections via a small per-store pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every DB operation (queue lease/enqueue, chat_state get/set, link_cache read/write) used to open a fresh connection — including the busy timeout and WAL pragma — then close it, on every message, URL job and callback. Replace with DbPool: a tiny pool (4 connections max, semaphore-bounded concurrency for backpressure) whose with_conn() method runs the closure on a pooled connection inside spawn_blocking. Steady-state cost of an operation is a list pop + semaphore acquire instead of a connection open. --- crates/xmedia-bot/src/db.rs | 119 ++++++++++++++++++++++------ crates/xmedia-bot/src/link_cache.rs | 14 ++-- crates/xmedia-bot/src/queue.rs | 26 +++--- crates/xmedia-bot/src/state.rs | 8 +- 4 files changed, 120 insertions(+), 47 deletions(-) diff --git a/crates/xmedia-bot/src/db.rs b/crates/xmedia-bot/src/db.rs index 6d783a4..84dd5de 100644 --- a/crates/xmedia-bot/src/db.rs +++ b/crates/xmedia-bot/src/db.rs @@ -2,16 +2,106 @@ //! (`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. +//! All I/O runs inside `spawn_blocking` via [`DbPool::with_conn`] — rusqlite +//! connections are not Send-friendly to hold across an await point, and +//! blocking the async executor stalls every handler. Connections are reused +//! through a small per-store pool instead of opening a fresh connection per +//! operation: WAL lets readers run alongside writer leases, and the pool's +//! semaphore bounds how many DB operations run concurrently, giving natural +//! backpressure on hot paths (every message / URL / callback touches +//! chat_state or the link cache). +use parking_lot::Mutex; use rusqlite::Connection; +use std::sync::Arc; use std::time::Duration; +/// Upper bound on pooled (reused) connections and on concurrent DB +/// operations per store. Small on purpose: the queue's `BEGIN IMMEDIATE` +/// leases serialize writes anyway, and WAL readers rarely need more. +const POOL_SIZE: usize = 4; + +/// A tiny connection pool for one SQLite file. Connections are checked out +/// on a blocking thread and returned afterwards; `acquire` opens a new +/// connection only when the idle list is empty, so the steady-state cost of +/// an operation is a list pop instead of a fresh open (+ busy timeout + WAL +/// pragma). The semaphore caps the number of concurrent operations, so a +/// burst of handlers queues up instead of opening unbounded connections. +pub struct DbPool { + // Arc so [`DbPool::with_conn`] can hand an owned handle to + // `spawn_blocking` without borrowing across the await point. + inner: Arc, +} + +struct PoolInner { + path: String, + permits: tokio::sync::Semaphore, + idle: Mutex>, +} + +impl DbPool { + pub fn new(path: &str) -> Self { + DbPool { + inner: Arc::new(PoolInner { + path: path.to_string(), + permits: tokio::sync::Semaphore::new(POOL_SIZE), + idle: Mutex::new(Vec::new()), + }), + } + } + + /// Runs `f` against a pooled connection on a blocking thread, returning + /// the closure's result. Owns the semaphore + `spawn_blocking` + + /// `expect` ceremony shared by every table access; the caller maps + /// errors to its own log line. + pub async fn with_conn(&self, f: F) -> rusqlite::Result + where + T: Send + 'static, + F: FnOnce(&mut Connection) -> rusqlite::Result + Send + 'static, + { + let _permit = self + .inner + .permits + .acquire() + .await + .expect("db pool semaphore closed"); + let inner = Arc::clone(&self.inner); + tokio::task::spawn_blocking(move || { + let mut conn = inner.acquire()?; + let result = f(&mut conn); + inner.release(conn); + result + }) + .await + .expect("db worker panicked") + } + + /// The database file this pool serves (used by tests that need a raw + /// connection, e.g. to seed rows directly). + #[cfg(test)] + pub fn path(&self) -> &str { + &self.inner.path + } +} + +impl PoolInner { + /// Reuses an idle connection or opens a fresh one. + fn acquire(&self) -> rusqlite::Result { + if let Some(conn) = self.idle.lock().pop() { + return Ok(conn); + } + open_db(&self.path) + } + + /// Returns a connection to the pool (dropped when the pool is full). + fn release(&self, conn: Connection) { + let mut idle = self.idle.lock(); + if idle.len() < POOL_SIZE { + idle.push(conn); + } + } +} + /// Opens the shared DB with a busy timeout. pub fn open_db(path: &str) -> rusqlite::Result { let conn = Connection::open(path)?; @@ -31,20 +121,3 @@ pub fn now_f64() -> f64 { .map(|d| d.as_secs_f64()) .unwrap_or(0.0) } - -/// 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 cf25850..a5eb43c 100644 --- a/crates/xmedia-bot/src/link_cache.rs +++ b/crates/xmedia-bot/src/link_cache.rs @@ -47,7 +47,7 @@ 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). pub struct LinkCache { - db_path: String, + pool: crate::db::DbPool, } impl LinkCache { @@ -61,7 +61,7 @@ impl LinkCache { log::error!("failed to initialize link cache schema: {e}"); } Self { - db_path: db_path.to_string(), + pool: crate::db::DbPool::new(db_path), } } @@ -70,7 +70,7 @@ impl LinkCache { pub async fn get(&self, key: &str, ttl: Duration) -> Option { let key = key.to_string(); let ttl = ttl.as_secs_f64(); - let result = crate::db::with_conn(&self.db_path, move |conn| { + let result = self.pool.with_conn(move |conn| { let mut stmt = conn.prepare("SELECT payload, created_at FROM link_cache WHERE url = ?1")?; let mut rows = stmt.query(params![key])?; @@ -100,7 +100,7 @@ impl LinkCache { pub async fn put(&self, key: &str, post: &CachedPost) { let key = key.to_string(); let payload = serde_json::to_string(post).expect("cached post serializes"); - let result = crate::db::with_conn(&self.db_path, move |conn| { + let result = self.pool.with_conn(move |conn| { conn.execute( "INSERT OR REPLACE INTO link_cache (url, payload, created_at) VALUES (?1, ?2, ?3)", params![key, payload, now_f64()], @@ -116,7 +116,7 @@ impl LinkCache { /// Drops an entry (e.g. a cached file id that turned out invalid). pub async fn remove(&self, key: &str) { let key = key.to_string(); - let result = crate::db::with_conn(&self.db_path, move |conn| { + let result = self.pool.with_conn(move |conn| { conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?; Ok(()) }) @@ -129,7 +129,7 @@ impl LinkCache { /// Removes expired entries; returns how many were deleted. pub async fn prune(&self, ttl: Duration) -> usize { let cutoff = now_f64() - ttl.as_secs_f64(); - let result = crate::db::with_conn(&self.db_path, move |conn| { + let result = self.pool.with_conn(move |conn| { conn.execute( "DELETE FROM link_cache WHERE created_at < ?1", params![cutoff], @@ -149,7 +149,7 @@ impl LinkCache { /// `key` is `None`. Returns how many rows were removed. pub async fn clear(&self, key: Option<&str>) -> usize { let key = key.map(str::to_string); - let result = crate::db::with_conn(&self.db_path, move |conn| match &key { + let result = self.pool.with_conn(move |conn| match &key { Some(key) => conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key]), None => conn.execute("DELETE FROM link_cache", []), }) diff --git a/crates/xmedia-bot/src/queue.rs b/crates/xmedia-bot/src/queue.rs index 6250ff7..8606f72 100644 --- a/crates/xmedia-bot/src/queue.rs +++ b/crates/xmedia-bot/src/queue.rs @@ -40,7 +40,7 @@ type Handler = dyn Fn(Value) -> BoxFuture<'static, Result<(), QueueError>> + Sen type DeadLetter = dyn Fn(Value, String) -> BoxFuture<'static, ()> + Send + Sync; pub struct PersistentTaskQueue { - db_path: String, + pool: std::sync::Arc, notify: Arc, stop: Arc, worker: Mutex>>, @@ -56,7 +56,7 @@ struct LeasedRow { /// Owned worker state so the spawned loop does not borrow the queue handle. #[derive(Clone)] struct QueueWorker { - db_path: String, + pool: std::sync::Arc, notify: Arc, stop: Arc, handler: Arc, @@ -108,7 +108,7 @@ impl PersistentTaskQueue { log::error!("failed to initialize queue schema: {e}"); } Self { - db_path: db_path.to_string(), + pool: std::sync::Arc::new(crate::db::DbPool::new(db_path)), notify: Arc::new(Notify::new()), stop: Arc::new(AtomicBool::new(false)), worker: Mutex::new(Vec::new()), @@ -132,7 +132,7 @@ impl PersistentTaskQueue { let mut handles = Vec::with_capacity(QUEUE_WORKERS + 1); for _ in 0..QUEUE_WORKERS { let worker = QueueWorker { - db_path: self.db_path.clone(), + pool: std::sync::Arc::clone(&self.pool), notify: Arc::clone(&self.notify), stop: Arc::clone(&self.stop), handler: Arc::clone(&handler), @@ -145,7 +145,7 @@ impl PersistentTaskQueue { // the same notify as the workers, so enqueue and stop interrupt the // sleep; the first interval tick fires immediately (harmless extra // recovery at startup). - let sweep_db_path = self.db_path.clone(); + let sweep_pool = std::sync::Arc::clone(&self.pool); let sweep_notify = Arc::clone(&self.notify); let sweep_stop = Arc::clone(&self.stop); handles.push(tokio::spawn(async move { @@ -161,7 +161,7 @@ impl PersistentTaskQueue { break; } let result = - crate::db::with_conn(&sweep_db_path, move |conn| recover_update(conn)).await; + sweep_pool.with_conn(move |conn| recover_update(conn)).await; if let Err(e) = result { log::error!("queue sweep failed: {e}"); } @@ -190,7 +190,7 @@ impl PersistentTaskQueue { ); let payload = payload.to_string(); log::info!("enqueued {id} (run_after {run_after:.1})"); - crate::db::with_conn(&self.db_path, move |conn| { + self.pool.with_conn(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)", @@ -212,7 +212,7 @@ impl PersistentTaskQueue { } async fn recover_sweep(&self) { - let result = crate::db::with_conn(&self.db_path, move |conn| recover_update(conn)).await; + let result = self.pool.with_conn(move |conn| recover_update(conn)).await; if let Err(e) = result { log::error!("queue recovery failed: {e}"); } @@ -267,7 +267,7 @@ impl QueueWorker { /// Leases the oldest due row (sets it `in_progress` with a lock TTL). /// Errors are surfaced so the caller can back off instead of spinning. async fn lease_next(&self) -> Result, rusqlite::Error> { - crate::db::with_conn(&self.db_path, |conn| { + self.pool.with_conn(|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 @@ -309,7 +309,7 @@ impl QueueWorker { } async fn earliest_run_after(&self) -> Option { - let result = crate::db::with_conn(&self.db_path, |conn| { + let result = self.pool.with_conn(|conn| { let mut stmt = conn.prepare("SELECT MIN(run_after) FROM tasks WHERE status='pending'")?; let mut rows = stmt.query([])?; @@ -374,7 +374,7 @@ impl QueueWorker { async fn delete_row(&self, id: &str) { let id = id.to_string(); - let result = crate::db::with_conn(&self.db_path, move |conn| { + let result = self.pool.with_conn(move |conn| { conn.execute("DELETE FROM tasks WHERE id = ?1", params![id])?; Ok(()) }) @@ -387,7 +387,7 @@ impl QueueWorker { async fn reschedule(&self, id: &str, payload: Value, delay_seconds: f64, attempts: i32) { let id = id.to_string(); let payload = payload.to_string(); - let result = crate::db::with_conn(&self.db_path, move |conn| { + let result = self.pool.with_conn(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], @@ -575,7 +575,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.db_path).unwrap(); + let conn = 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) \ diff --git a/crates/xmedia-bot/src/state.rs b/crates/xmedia-bot/src/state.rs index c7fd55d..d26dcde 100644 --- a/crates/xmedia-bot/src/state.rs +++ b/crates/xmedia-bot/src/state.rs @@ -38,7 +38,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>>>, - db_path: String, + pool: crate::db::DbPool, } pub fn unix_now() -> i64 { @@ -67,7 +67,7 @@ impl ChatStore { Ok(ChatStore { cache: Mutex::new(HashMap::new()), locks: Mutex::new(HashMap::new()), - db_path: path.to_string(), + pool: crate::db::DbPool::new(path), }) } @@ -76,7 +76,7 @@ impl ChatStore { return data.clone(); } let chat_key = chat_id.to_string(); - let payload = crate::db::with_conn(&self.db_path, move |conn| { + let payload = self.pool.with_conn(move |conn| { // Concurrent handler tasks (batch-forwards) may write chat_state // while this read runs; the shared busy timeout handles the // write-lock collision instead of failing the query. @@ -103,7 +103,7 @@ impl ChatStore { self.cache.lock().insert(chat_id, data.clone()); let payload = serde_json::to_string(data).expect("chat state serializes"); let chat_id = chat_id.to_string(); - let result = crate::db::with_conn(&self.db_path, move |conn| { + let result = self.pool.with_conn(move |conn| { conn.execute( "INSERT OR REPLACE INTO chat_state (chat_id, payload) VALUES (?1, ?2)", params![chat_id, payload],