YoursFunny/TelegramTwitterMediaBot
1
0
Fork 0
mirror of https://github.com/TheFunny/TelegramTwitterMediaBot.git synced 2026-09-23 23:32:05 +00:00
Code Releases Activity

refactor: share sqlite open/with_conn helpers in db.rs

Browse Source 
Converge the duplicated open_db (open + busy_timeout) and the
spawn_blocking + expect ceremony that every table access repeated
into one db.rs module. ChatStore no longer creates the tasks table
(schema ownership: queue.rs owns tasks, state.rs chat_state,
link_cache.rs link_cache). No schema or behavior change - all
CREATE TABLE statements are byte-identical, IF NOT EXISTS stays
idempotent, so existing data/task_queue.db files need no migration.
This commit is contained in:
YoursFunny
2026-08-07 16:00:12 +08:00
parent 4060a88031
commit b0ced34b4c
5 changed files with 134 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/xmedia-bot/src/db.rs
+37
View File
@@ -0,0 +1,37 @@
//! 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).
//!
//! 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.
use rusqlite::Connection;
use std::time::Duration;
/// Opens the shared DB with a busy timeout.
pub fn open_db(path: &str) -> rusqlite::Result<Connection> {
let conn = Connection::open(path)?;
conn.busy_timeout(Duration::from_secs(5))?;
Ok(conn)
}
/// 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")
}
crates/xmedia-bot/src/link_cache.rs
+27 -35
View File
@@ -49,12 +49,6 @@ pub struct LinkCache {
db_path: String, db_path: String,
} }
fn open_db(path: &str) -> rusqlite::Result<Connection> {
let conn = Connection::open(path)?;
conn.busy_timeout(Duration::from_secs(5))?;
Ok(conn)
}
impl LinkCache { impl LinkCache {
pub fn open(db_path: &str) -> Self { pub fn open(db_path: &str) -> Self {
if let Ok(conn) = Connection::open(db_path) if let Ok(conn) = Connection::open(db_path)
@@ -73,11 +67,9 @@ impl LinkCache {
/// Returns the cached post if present and not expired; a stale entry is /// Returns the cached post if present and not expired; a stale entry is
/// removed on the spot. /// removed on the spot.
pub async fn get(&self, key: &str, ttl: Duration) -> Option<CachedPost> { pub async fn get(&self, key: &str, ttl: Duration) -> Option<CachedPost> {
let db_path = self.db_path.clone();
let key = key.to_string(); let key = key.to_string();
let ttl = ttl.as_secs_f64(); let ttl = ttl.as_secs_f64();
tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<CachedPost>> { let result = crate::db::with_conn(&self.db_path, move |conn| {
let conn = open_db(&db_path)?;
let mut stmt = let mut stmt =
conn.prepare("SELECT payload, created_at FROM link_cache WHERE url = ?1")?; conn.prepare("SELECT payload, created_at FROM link_cache WHERE url = ?1")?;
let mut rows = stmt.query(params![key])?; let mut rows = stmt.query(params![key])?;
@@ -90,66 +82,66 @@ impl LinkCache {
conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?; conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?;
return Ok(None); return Ok(None);
} }
serde_json::from_str(&payload).map(Some).map_err(|e| { Ok(Some(serde_json::from_str::<CachedPost>(&payload).map_err(
rusqlite::Error::ToSqlConversionFailure(Box::new(e)) |e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)),
)?))
}) })
}) .await;
.await match result {
.expect("link cache read worker panicked") Ok(v) => v,
.unwrap_or_else(|e| { Err(e) => {
log::error!("link cache read failed: {e}"); log::error!("link cache read failed: {e}");
None None
}) }
}
} }
pub async fn put(&self, key: &str, post: &CachedPost) { pub async fn put(&self, key: &str, post: &CachedPost) {
let db_path = self.db_path.clone();
let key = key.to_string(); let key = key.to_string();
let payload = serde_json::to_string(post).expect("cached post serializes"); let payload = serde_json::to_string(post).expect("cached post serializes");
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { let result = crate::db::with_conn(&self.db_path, move |conn| {
let conn = open_db(&db_path)?;
conn.execute( conn.execute(
"INSERT OR REPLACE INTO link_cache (url, payload, created_at) VALUES (?1, ?2, ?3)", "INSERT OR REPLACE INTO link_cache (url, payload, created_at) VALUES (?1, ?2, ?3)",
params![key, payload, now_f64()], params![key, payload, now_f64()],
)?; )?;
Ok(()) Ok(())
}) })
.await .await;
.expect("link cache write worker panicked") if let Err(e) = result {
.unwrap_or_else(|e| log::error!("link cache write failed: {e}")); log::error!("link cache write failed: {e}");
}
} }
/// Drops an entry (e.g. a cached file id that turned out invalid). /// Drops an entry (e.g. a cached file id that turned out invalid).
pub async fn remove(&self, key: &str) { pub async fn remove(&self, key: &str) {
let db_path = self.db_path.clone();
let key = key.to_string(); let key = key.to_string();
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { let result = crate::db::with_conn(&self.db_path, move |conn| {
let conn = open_db(&db_path)?;
conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?; conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?;
Ok(()) Ok(())
}) })
.await .await;
.expect("link cache delete worker panicked") if let Err(e) = result {
.unwrap_or_else(|e| log::error!("link cache delete failed: {e}")); log::error!("link cache delete failed: {e}");
}
} }
/// Removes expired entries; returns how many were deleted. /// Removes expired entries; returns how many were deleted.
pub async fn prune(&self, ttl: Duration) -> usize { pub async fn prune(&self, ttl: Duration) -> usize {
let db_path = self.db_path.clone();
let cutoff = now_f64() - ttl.as_secs_f64(); let cutoff = now_f64() - ttl.as_secs_f64();
tokio::task::spawn_blocking(move || -> rusqlite::Result<usize> { let result = crate::db::with_conn(&self.db_path, move |conn| {
let conn = open_db(&db_path)?;
conn.execute( conn.execute(
"DELETE FROM link_cache WHERE created_at < ?1", "DELETE FROM link_cache WHERE created_at < ?1",
params![cutoff], params![cutoff],
) )
}) })
.await .await;
.expect("link cache prune worker panicked") match result {
.unwrap_or_else(|e| { Ok(n) => n,
Err(e) => {
log::error!("link cache prune failed: {e}"); log::error!("link cache prune failed: {e}");
0 0
}) }
}
} }
} }
crates/xmedia-bot/src/main.rs
+1
View File
@@ -8,6 +8,7 @@ use tokio::sync::watch;
use x_media::site; use x_media::site;
mod config; mod config;
mod db;
mod handlers; mod handlers;
mod link_cache; mod link_cache;
mod photo; mod photo;
crates/xmedia-bot/src/queue.rs
+38 -58
View File
@@ -6,11 +6,11 @@
//! replaced by dedicated columns. //! replaced by dedicated columns.
use parking_lot::Mutex; use parking_lot::Mutex;
use rusqlite::{params, Connection, TransactionBehavior}; use rusqlite::{Connection, TransactionBehavior, params};
use serde_json::Value; use serde_json::Value;
use std::pin::Pin; use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc; use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::{Duration, SystemTime, UNIX_EPOCH}; use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::Notify; use tokio::sync::Notify;
use tokio::task::JoinHandle; use tokio::task::JoinHandle;
@@ -29,15 +29,9 @@ const QUEUE_WORKERS: usize = 4;
pub enum QueueError { pub enum QueueError {
/// Reschedule with the given delay; after `MAX_RETRIES` attempts the task /// Reschedule with the given delay; after `MAX_RETRIES` attempts the task
/// is dead-lettered instead. /// is dead-lettered instead.
Retryable { Retryable { delay_seconds: f64, payload: Value },
delay_seconds: f64,
payload: Value,
},
/// Give up now. /// Give up now.
Permanent { Permanent { message: String, payload: Value },
message: String,
payload: Value,
},
} }
type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>; type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
@@ -74,16 +68,7 @@ fn now_f64() -> f64 {
.unwrap_or(0.0) .unwrap_or(0.0)
} }
/// Opens the queue DB with a busy timeout. Handler tasks enqueue while fn ensure_schema(conn: &rusqlite::Connection) -> rusqlite::Result<()> {
/// workers lease/update rows concurrently; without the timeout a concurrent
/// write fails immediately with SQLITE_BUSY and the operation is lost.
fn open_db(path: &str) -> rusqlite::Result<Connection> {
let conn = Connection::open(path)?;
conn.busy_timeout(Duration::from_secs(5))?;
Ok(conn)
}
fn ensure_schema(conn: &Connection) -> rusqlite::Result<()> {
conn.execute_batch( conn.execute_batch(
"CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, payload TEXT NOT NULL, \ "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, \ run_after REAL NOT NULL, attempts INTEGER NOT NULL, status TEXT NOT NULL, \
@@ -162,10 +147,8 @@ impl PersistentTaskQueue {
self.counter.fetch_add(1, Ordering::Relaxed) self.counter.fetch_add(1, Ordering::Relaxed)
); );
let payload = payload.to_string(); let payload = payload.to_string();
let db_path = self.db_path.clone();
log::info!("enqueued {id} (run_after {run_after:.1})"); log::info!("enqueued {id} (run_after {run_after:.1})");
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { crate::db::with_conn(&self.db_path, move |conn| {
let conn = open_db(&db_path)?;
conn.execute( conn.execute(
"INSERT OR REPLACE INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \ "INSERT OR REPLACE INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \
VALUES (?1, ?2, ?3, 0, 'pending', 0, ?4)", VALUES (?1, ?2, ?3, 0, 'pending', 0, ?4)",
@@ -173,8 +156,7 @@ impl PersistentTaskQueue {
)?; )?;
Ok(()) Ok(())
}) })
.await .await?;
.expect("queue insert worker panicked")?;
// Wake every sleeping worker: with several workers the one that finds // Wake every sleeping worker: with several workers the one that finds
// nothing due must not starve the newly inserted row. // nothing due must not starve the newly inserted row.
self.notify.notify_waiters(); self.notify.notify_waiters();
@@ -182,18 +164,17 @@ impl PersistentTaskQueue {
} }
async fn recover_stale(&self) { async fn recover_stale(&self) {
let db_path = self.db_path.clone(); let result = crate::db::with_conn(&self.db_path, move |conn| {
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> {
let conn = open_db(&db_path)?;
conn.execute( conn.execute(
"UPDATE tasks SET status='pending', locked_until=0 WHERE status='in_progress' AND locked_until < ?1", "UPDATE tasks SET status='pending', locked_until=0 WHERE status='in_progress' AND locked_until < ?1",
params![now_f64()], params![now_f64()],
)?; )?;
Ok(()) Ok(())
}) })
.await .await;
.expect("queue recovery worker panicked") if let Err(e) = result {
.unwrap_or_else(|e| log::error!("queue recovery failed: {e}")); log::error!("queue recovery failed: {e}");
}
} }
} }
@@ -225,9 +206,7 @@ impl QueueWorker {
/// Leases the oldest due row (sets it `in_progress` with a lock TTL). /// Leases the oldest due row (sets it `in_progress` with a lock TTL).
async fn lease_next(&self) -> Option<LeasedRow> { async fn lease_next(&self) -> Option<LeasedRow> {
let db_path = self.db_path.clone(); let result = crate::db::with_conn(&self.db_path, |conn| {
tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<LeasedRow>> {
let mut conn = open_db(&db_path)?;
// BEGIN IMMEDIATE: with several workers, a deferred transaction // BEGIN IMMEDIATE: with several workers, a deferred transaction
// that read before another worker's lease commit would fail with // that read before another worker's lease commit would fail with
// SQLITE_BUSY_SNAPSHOT. Taking the write lock up front serializes // SQLITE_BUSY_SNAPSHOT. Taking the write lock up front serializes
@@ -265,31 +244,34 @@ impl QueueWorker {
attempts, attempts,
})) }))
}) })
.await .await;
.expect("queue lease worker panicked") match result {
.unwrap_or_else(|e| { Ok(row) => row,
Err(e) => {
log::error!("queue lease failed: {e}"); log::error!("queue lease failed: {e}");
None None
}) }
}
} }
async fn earliest_run_after(&self) -> Option<f64> { async fn earliest_run_after(&self) -> Option<f64> {
let db_path = self.db_path.clone(); let result = crate::db::with_conn(&self.db_path, |conn| {
tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<f64>> { let mut stmt =
let conn = open_db(&db_path)?; conn.prepare("SELECT MIN(run_after) FROM tasks WHERE status='pending'")?;
let mut stmt = conn.prepare("SELECT MIN(run_after) FROM tasks WHERE status='pending'")?;
let mut rows = stmt.query([])?; let mut rows = stmt.query([])?;
match rows.next()? { match rows.next()? {
Some(row) => Ok(row.get::<_, Option<f64>>(0)?), Some(row) => Ok(row.get::<_, Option<f64>>(0)?),
None => Ok(None), None => Ok(None),
} }
}) })
.await .await;
.expect("queue timing worker panicked") match result {
.unwrap_or_else(|e| { Ok(v) => v,
Err(e) => {
log::error!("queue timing query failed: {e}"); log::error!("queue timing query failed: {e}");
None None
}) }
}
} }
async fn process(&self, row: LeasedRow) { async fn process(&self, row: LeasedRow) {
@@ -336,33 +318,31 @@ impl QueueWorker {
} }
async fn delete_row(&self, id: &str) { async fn delete_row(&self, id: &str) {
let db_path = self.db_path.clone();
let id = id.to_string(); let id = id.to_string();
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { let result = crate::db::with_conn(&self.db_path, move |conn| {
let conn = open_db(&db_path)?;
conn.execute("DELETE FROM tasks WHERE id = ?1", params![id])?; conn.execute("DELETE FROM tasks WHERE id = ?1", params![id])?;
Ok(()) Ok(())
}) })
.await .await;
.expect("queue delete worker panicked") if let Err(e) = result {
.unwrap_or_else(|e| log::error!("queue delete failed: {e}")); log::error!("queue delete failed: {e}");
}
} }
async fn reschedule(&self, id: &str, payload: Value, delay_seconds: f64, attempts: i32) { async fn reschedule(&self, id: &str, payload: Value, delay_seconds: f64, attempts: i32) {
let db_path = self.db_path.clone();
let id = id.to_string(); let id = id.to_string();
let payload = payload.to_string(); let payload = payload.to_string();
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { let result = crate::db::with_conn(&self.db_path, move |conn| {
let conn = open_db(&db_path)?;
conn.execute( conn.execute(
"UPDATE tasks SET payload=?1, run_after=?2, attempts=?3, status='pending', locked_until=0 WHERE id=?4", "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], params![payload, now_f64() + delay_seconds, attempts, id],
)?; )?;
Ok(()) Ok(())
}) })
.await .await;
.expect("queue reschedule worker panicked") if let Err(e) = result {
.unwrap_or_else(|e| log::error!("queue reschedule failed: {e}")); log::error!("queue reschedule failed: {e}");
}
self.notify.notify_waiters(); self.notify.notify_waiters();
} }
} }
crates/xmedia-bot/src/state.rs
+23 -25
View File
@@ -2,7 +2,7 @@
//! `data/task_queue.db`, shared with the task queue). //! `data/task_queue.db`, shared with the task queue).
use parking_lot::Mutex; use parking_lot::Mutex;
use rusqlite::{params, Connection}; use rusqlite::params;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::collections::HashMap; use std::collections::HashMap;
use std::path::Path; use std::path::Path;
@@ -45,7 +45,9 @@ pub fn unix_now() -> i64 {
} }
impl ChatStore { impl ChatStore {
/// Creates the parent directory and both tables (idempotent). /// 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<Self> { pub fn open(path: &str) -> rusqlite::Result<Self> {
if let Some(parent) = Path::new(path).parent() if let Some(parent) = Path::new(path).parent()
&& !parent.as_os_str().is_empty() && !parent.as_os_str().is_empty()
@@ -53,12 +55,9 @@ impl ChatStore {
std::fs::create_dir_all(parent) std::fs::create_dir_all(parent)
.map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?; .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;
} }
let conn = Connection::open(path)?; let conn = crate::db::open_db(path)?;
conn.execute_batch( conn.execute_batch(
"CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, payload TEXT NOT NULL, \ "CREATE TABLE IF NOT EXISTS chat_state (chat_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 TABLE IF NOT EXISTS chat_state (chat_id TEXT PRIMARY KEY, payload TEXT NOT NULL);",
)?; )?;
drop(conn); drop(conn);
Ok(ChatStore { Ok(ChatStore {
@@ -71,22 +70,19 @@ impl ChatStore {
if let Some(data) = self.cache.lock().get(&chat_id) { if let Some(data) = self.cache.lock().get(&chat_id) {
return data.clone(); return data.clone();
} }
let db_path = self.db_path.clone(); let chat_key = chat_id.to_string();
let payload = tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<String>> { let payload = crate::db::with_conn(&self.db_path, move |conn| {
let conn = Connection::open(&db_path)?;
// Concurrent handler tasks (batch-forwards) may write chat_state // Concurrent handler tasks (batch-forwards) may write chat_state
// while this read runs; without a busy timeout a write lock // while this read runs; the shared busy timeout handles the
// collision fails the query immediately. // write-lock collision instead of failing the query.
conn.busy_timeout(std::time::Duration::from_secs(5))?;
let mut stmt = conn.prepare("SELECT payload FROM chat_state WHERE chat_id = ?1")?; let mut stmt = conn.prepare("SELECT payload FROM chat_state WHERE chat_id = ?1")?;
let mut rows = stmt.query(params![chat_id.to_string()])?; let mut rows = stmt.query(params![chat_key])?;
match rows.next()? { match rows.next()? {
Some(row) => Ok(Some(row.get(0)?)), Some(row) => Ok(Some(row.get::<_, String>(0)?)),
None => Ok(None), None => Ok(None),
} }
}) })
.await .await
.expect("chat_state worker panicked")
.unwrap_or_else(|e| { .unwrap_or_else(|e| {
log::error!("chat_state read failed: {e}"); log::error!("chat_state read failed: {e}");
None None
@@ -101,19 +97,18 @@ impl ChatStore {
pub async fn set(&self, chat_id: i64, data: &ChatData) { pub async fn set(&self, chat_id: i64, data: &ChatData) {
self.cache.lock().insert(chat_id, data.clone()); self.cache.lock().insert(chat_id, data.clone());
let payload = serde_json::to_string(data).expect("chat state serializes"); let payload = serde_json::to_string(data).expect("chat state serializes");
let db_path = self.db_path.clone(); let chat_id = chat_id.to_string();
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { let result = crate::db::with_conn(&self.db_path, move |conn| {
let conn = Connection::open(&db_path)?;
conn.busy_timeout(std::time::Duration::from_secs(5))?;
conn.execute( conn.execute(
"INSERT OR REPLACE INTO chat_state (chat_id, payload) VALUES (?1, ?2)", "INSERT OR REPLACE INTO chat_state (chat_id, payload) VALUES (?1, ?2)",
params![chat_id.to_string(), payload], params![chat_id, payload],
)?; )?;
Ok(()) Ok(())
}) })
.await .await;
.expect("chat_state worker panicked") if let Err(e) = result {
.unwrap_or_else(|e| log::error!("chat_state write failed: {e}")); log::error!("chat_state write failed: {e}");
}
} }
/// Removes edit-before-forward records whose `created_at + ttl` is in the /// Removes edit-before-forward records whose `created_at + ttl` is in the
@@ -149,7 +144,10 @@ impl ChatStore {
self.set(chat_id, &data).await; self.set(chat_id, &data).await;
} }
if !removed.is_empty() { if !removed.is_empty() {
log::info!("pruned {} expired edit-before-forward record(s)", removed.len()); log::info!(
"pruned {} expired edit-before-forward record(s)",
removed.len()
);
} }
removed removed
} }