perf(db): reuse SQLite connections via a small per-store pool

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.
This commit is contained in:
2026-08-13 22:12:34 +08:00
parent 4a467641aa
commit edb32c23b4
4 changed files with 120 additions and 47 deletions
+96 -23
View File
@@ -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<PoolInner>,
}
struct PoolInner {
path: String,
permits: tokio::sync::Semaphore,
idle: Mutex<Vec<Connection>>,
}
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<T, F>(&self, f: F) -> rusqlite::Result<T>
where
T: Send + 'static,
F: FnOnce(&mut Connection) -> rusqlite::Result<T> + 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<Connection> {
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<Connection> {
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<T, F>(path: &str, f: F) -> rusqlite::Result<T>
where
T: Send + 'static,
F: FnOnce(&mut Connection) -> rusqlite::Result<T> + 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")
}
+7 -7
View File
@@ -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<CachedPost> {
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", []),
})
+13 -13
View File
@@ -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<crate::db::DbPool>,
notify: Arc<Notify>,
stop: Arc<AtomicBool>,
worker: Mutex<Vec<JoinHandle<()>>>,
@@ -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<crate::db::DbPool>,
notify: Arc<Notify>,
stop: Arc<AtomicBool>,
handler: Arc<Handler>,
@@ -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<Option<LeasedRow>, 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<f64> {
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) \
+4 -4
View File
@@ -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<HashMap<i64, Arc<tokio::sync::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],