//! 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). //! //! 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)?; conn.busy_timeout(Duration::from_secs(5))?; // WAL lets readers run alongside writer leases instead of blocking on // the rollback journal; the mode persists in the DB header, so the // idempotent pragma here and in ensure_schema only needs to win once. conn.pragma_update(None, "journal_mode", "WAL")?; 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)?; migrate(&conn)?; Ok(Arc::new(DbPool::new(path))) } /// Schema migrations, applied in order and tracked by `PRAGMA user_version` /// (the index in this array + 1 is the version a statement brings the /// database to). Append only — never edit or reorder an entry, or databases /// already past it would skip or repeat work. const MIGRATIONS: &[&str] = &[ // 1: lease fencing. A worker's write-backs (`delete`/`reschedule`/the // lease heartbeat) are guarded by the token it was leased with, so a // lease that expired and was re-leased by another worker can no longer be // written by its former holder — which used to duplicate a send or drop // the new holder's retry state, silently. "ALTER TABLE tasks ADD COLUMN lease_token TEXT", // 2: the 300 s sweep prunes the link cache by `created_at` // (`DELETE FROM link_cache WHERE created_at < ?`). Without an index that // is a full scan of every post sent inside the TTL window — up to a week // of them — on every sweep; the `url` primary key cannot serve it. "CREATE INDEX IF NOT EXISTS idx_link_cache_created_at ON link_cache(created_at)", ]; /// Brings an existing database up to [`MIGRATIONS`]. Idempotent: a database /// already at the latest version does no work. fn migrate(conn: &Connection) -> rusqlite::Result<()> { let version: i64 = conn.query_row("PRAGMA user_version", [], |row| row.get(0))?; for (index, statement) in MIGRATIONS.iter().enumerate() { let target = index as i64 + 1; if version >= target { continue; } conn.execute_batch(statement)?; // `PRAGMA` does not take bind parameters; the value is our own index. conn.execute_batch(&format!("PRAGMA user_version = {target}"))?; } Ok(()) } 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. /// /// This is the **baseline** schema (version 0): a fresh database is created /// exactly like this, and anything that must *change* an existing one is /// appended to [`MIGRATIONS`] instead of being edited in here — otherwise a /// database created before the change would never gain the new column and a /// freshly created one would try to apply the migration a second time. 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 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_secs_f64()) .unwrap_or(0.0) } /// Unix timestamp in whole seconds. Same clock as [`now_f64`], for fields /// that store integer seconds (chat-state expiry, edit prompts). pub fn unix_now() -> i64 { now_f64() as i64 } #[cfg(test)] mod tests { use super::*; use rusqlite::Connection; /// The schema as it shipped *before* the first migration: what an existing /// deployment has on disk when it starts on the new binary. Written out /// literally rather than derived from `schema_init`, so an edit to the /// baseline shows up here instead of being followed silently. const V0_SCHEMA: &str = "CREATE TABLE 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 idx_tasks_pending ON tasks(status, run_after); \ CREATE TABLE chat_state (chat_id TEXT PRIMARY KEY, payload TEXT NOT NULL); \ CREATE TABLE link_cache (url TEXT PRIMARY KEY, payload TEXT NOT NULL, \ created_at REAL NOT NULL);"; /// The migrations that have already shipped, verbatim. Appending is the only /// allowed change: editing one that a database has already applied leaves /// deployments on different schemas with nothing to notice it — the version /// counter says "done" and skips the new text. const SHIPPED_MIGRATIONS: &[&str] = &["ALTER TABLE tasks ADD COLUMN lease_token TEXT"]; fn columns(conn: &Connection, table: &str) -> Vec { let mut stmt = conn .prepare(&format!("PRAGMA table_info({table})")) .unwrap(); let mut names: Vec = stmt .query_map([], |row| row.get::<_, String>(1)) .unwrap() .map(Result::unwrap) .collect(); names.sort(); names } fn user_version(conn: &Connection) -> i64 { conn.query_row("PRAGMA user_version", [], |row| row.get(0)) .unwrap() } #[tokio::test] async fn a_pre_migration_database_upgrades_and_keeps_its_rows() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("old.db"); { let conn = Connection::open(&path).unwrap(); conn.execute_batch(V0_SCHEMA).unwrap(); conn.execute( "INSERT INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \ VALUES ('task_old', '{\"chat_id\":1}', 0, 0, 'pending', 0, 0)", [], ) .unwrap(); assert_eq!(user_version(&conn), 0, "the fixture starts un-migrated"); assert!( !columns(&conn, "tasks").contains(&"lease_token".to_string()), "the fixture is the pre-migration shape" ); } let pool = open_store(path.to_str().unwrap()).unwrap(); pool.with_conn(|conn| { assert_eq!(user_version(conn), MIGRATIONS.len() as i64); let mut expected = vec![ "id", "payload", "run_after", "attempts", "status", "locked_until", "created_at", "lease_token", ]; expected.sort(); assert_eq!( columns(conn, "tasks"), expected, "an upgrade must add the migration's column and nothing else" ); let payload: String = conn .query_row( "SELECT payload FROM tasks WHERE id = 'task_old'", [], |row| row.get(0), ) .unwrap(); assert_eq!(payload, "{\"chat_id\":1}", "rows survive the upgrade"); // The link-cache prune's index arrives with the migrations (the // baseline schema has none): without it every sweep scans the // whole table. let index: i64 = conn .query_row( "SELECT COUNT(*) FROM sqlite_master \ WHERE type = 'index' AND name = 'idx_link_cache_created_at'", [], |row| row.get(0), ) .unwrap(); assert_eq!(index, 1, "the migration's index must exist"); Ok(()) }) .await .unwrap(); } #[test] fn shipped_migrations_are_frozen() { assert!( MIGRATIONS.len() >= SHIPPED_MIGRATIONS.len(), "migrations were removed or reordered, not appended" ); for (index, (shipped, current)) in SHIPPED_MIGRATIONS.iter().zip(MIGRATIONS).enumerate() { assert_eq!( shipped, current, "migration {} already shipped: append a new one instead of editing it", index + 1 ); } } #[tokio::test] async fn a_fresh_database_lands_at_the_latest_version() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("fresh.db"); let pool = open_store(path.to_str().unwrap()).unwrap(); // Every migration is applied on creation, so a deployment that only ever // saw fresh databases is on the same schema as an upgraded one. pool.with_conn(|conn| { assert_eq!(user_version(conn), MIGRATIONS.len() as i64); Ok(()) }) .await .unwrap(); // Opening the same file again is a no-op (the version gate skips it). open_store(path.to_str().unwrap()).unwrap(); } }