From ab2306002a7b1ace6b56cc099a9938d17b91f784 Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Wed, 5 Aug 2026 01:26:24 +0800 Subject: [PATCH] perf: handle batch-forwarded URLs concurrently with queue workers --- AGENTS.md | 4 +- crates/x-media/src/site/twitter/interface.rs | 2 +- crates/xmedia-bot/src/handlers.rs | 17 ++++- crates/xmedia-bot/src/queue.rs | 79 +++++++++++++------- crates/xmedia-bot/src/state.rs | 5 ++ 5 files changed, 76 insertions(+), 31 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 6fd7f21..318b616 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -30,9 +30,9 @@ The `x-media` library: `site::fetch(url)` dispatches (in order) twitter → bsky | `crates/x-media/src/site//` | One directory per site: `mod.rs` (re-exports), `interface.rs` (PATTERN, `enabled()`, `fetch_from_url()`, site struct, `From for Fetched`), `model.rs` (serde DTOs). Pixiv adds `api.rs` (auth + transport); twitter adds `auth.rs` (logged-in GraphQL `TweetDetail` fallback for NSFW tweets, gated on `TWITTER_AUTH_TOKEN`) | | `crates/xmedia-bot/src/main.rs` | Entry point: env/log init, queue worker start, pixiv validation, 300 s edit-expiry sweep, dptree handler tree, webhook vs polling dispatch | | `crates/xmedia-bot/src/config.rs` | Manual env parsing into `Config` | -| `crates/xmedia-bot/src/handlers.rs` | `Command` enum (teloxide `BotCommands`), message/inline/callback handlers, URL extraction, global statics | +| `crates/xmedia-bot/src/handlers.rs` | `Command` enum (teloxide `BotCommands`), message/inline/callback handlers, URL extraction, global statics; per-URL work spawned with a `Semaphore(8)` cap (teloxide's per-chat workers are sequential — batch-forwards need concurrency) | | `crates/xmedia-bot/src/state.rs` | `ChatStore`: parking_lot `Mutex` cache + SQLite write-through (`chat_state` table) | -| `crates/xmedia-bot/src/queue.rs` | `PersistentTaskQueue`: SQLite-backed single-worker queue (`tasks` table) | +| `crates/xmedia-bot/src/queue.rs` | `PersistentTaskQueue`: SQLite-backed queue (`tasks` table), `QUEUE_WORKERS = 4` concurrent workers (lease via `BEGIN IMMEDIATE` + `locked_until` TTL), retry→dead-letter, `Notify::notify_waiters` wakeup, `busy_timeout` on all connections | | `crates/xmedia-bot/src/send.rs` | Media senders, upload fallback, error classification, queue task handlers | ## Development Commands diff --git a/crates/x-media/src/site/twitter/interface.rs b/crates/x-media/src/site/twitter/interface.rs index e3dbe2e..43ce94c 100644 --- a/crates/x-media/src/site/twitter/interface.rs +++ b/crates/x-media/src/site/twitter/interface.rs @@ -90,7 +90,7 @@ pub async fn fetch(id: &str) -> Result { { return Err(FetchError::Sensitive); } - Ok(Tweet::from_syndication_json(&text).map_err(FetchError::Json)?) + Tweet::from_syndication_json(&text).map_err(FetchError::Json) } /// The syndication token: JS `((id / 1e15) * PI).toString(36)` (the diff --git a/crates/xmedia-bot/src/handlers.rs b/crates/xmedia-bot/src/handlers.rs index 8b133b2..1f1f174 100644 --- a/crates/xmedia-bot/src/handlers.rs +++ b/crates/xmedia-bot/src/handlers.rs @@ -5,6 +5,7 @@ use crate::state::{ChatStore, unix_now}; use std::collections::HashSet; use std::sync::LazyLock; use teloxide::prelude::*; +use tokio::sync::Semaphore; use teloxide::types::{ CallbackQuery, ChatAction, ChatId, ChatKind, InlineQuery, InlineQueryResult, InlineQueryResultMpeg4Gif, InlineQueryResultPhoto, InlineQueryResultVideo, Message, @@ -21,6 +22,14 @@ pub static TASK_QUEUE: LazyLock = LazyLock::new(|| PersistentTaskQueue::new("data/task_queue.db")); pub static CONFIG: LazyLock = LazyLock::new(Config::load); +/// Cap on concurrent per-URL processing. teloxide dispatches updates to a +/// per-chat worker that handles them sequentially, so a batch-forward of many +/// messages would otherwise be processed one at a time (fetch + send each, +/// roughly a second per message). Moving the work into spawned tasks trades +/// per-chat reply ordering for throughput; the semaphore bounds how many run +/// at once so a big burst cannot hammer Telegram's rate limits. +static URL_TASKS: LazyLock = LazyLock::new(|| Semaphore::new(8)); + #[derive(BotCommands, Clone)] #[command(rename_rule = "snake_case", description = "")] enum Command { @@ -481,7 +490,13 @@ pub async fn message_handler(bot: Bot, message: Message) -> Result<(), RequestEr log::info!("extracted {} URL(s): {urls:?}", urls.len()); } for url in urls { - url_media(bot.clone(), &message, &url).await; + let bot = bot.clone(); + let message = message.clone(); + tokio::spawn(async move { + // Held for the whole task; the semaphore is never closed. + let _permit = URL_TASKS.acquire().await.expect("URL semaphore closed"); + url_media(bot, &message, &url).await; + }); } } respond(()) diff --git a/crates/xmedia-bot/src/queue.rs b/crates/xmedia-bot/src/queue.rs index a9cdfe3..cefe025 100644 --- a/crates/xmedia-bot/src/queue.rs +++ b/crates/xmedia-bot/src/queue.rs @@ -6,7 +6,7 @@ //! replaced by dedicated columns. use parking_lot::Mutex; -use rusqlite::{params, Connection}; +use rusqlite::{params, Connection, TransactionBehavior}; use serde_json::Value; use std::pin::Pin; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; @@ -18,6 +18,12 @@ use tokio::task::JoinHandle; pub const MAX_RETRIES: u32 = 2; pub const LOCK_TTL_SECONDS: f64 = 120.0; +/// Number of concurrent worker loops. Tasks are independent (retries and +/// forward resumes); leases serialize row claims via SQLite transactions, so +/// extra workers drain backlogs faster. Each worker can be mid-send to +/// Telegram at the same time as handler tasks, so keep this modest. +const QUEUE_WORKERS: usize = 4; + /// What a handler returns instead of throwing. The payload it carries is the /// (possibly updated) task state to persist for the next attempt. pub enum QueueError { @@ -42,7 +48,7 @@ pub struct PersistentTaskQueue { db_path: String, notify: Arc, stop: Arc, - worker: Mutex>>, + worker: Mutex>>, counter: AtomicU64, } @@ -68,6 +74,15 @@ fn now_f64() -> f64 { .unwrap_or(0.0) } +/// Opens the queue DB with a busy timeout. Handler tasks enqueue while +/// 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 { + let conn = Connection::open(path)?; + conn.busy_timeout(Duration::from_secs(5))?; + Ok(conn) +} + fn ensure_schema(conn: &Connection) -> rusqlite::Result<()> { conn.execute_batch( "CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, payload TEXT NOT NULL, \ @@ -96,12 +111,12 @@ impl PersistentTaskQueue { db_path: db_path.to_string(), notify: Arc::new(Notify::new()), stop: Arc::new(AtomicBool::new(false)), - worker: Mutex::new(None), + worker: Mutex::new(Vec::new()), counter: AtomicU64::new(0), } } - /// Starts the worker loop. Also recovers rows left `in_progress` by a + /// Starts the worker loops. Also recovers rows left `in_progress` by a /// previous process (lease expired). pub async fn start(&self, handler: H, dead_letter: D) where @@ -114,21 +129,25 @@ impl PersistentTaskQueue { let dead_letter: Arc = Arc::new(move |payload, message| Box::pin(dead_letter(payload, message))); self.recover_stale().await; - let worker = QueueWorker { - db_path: self.db_path.clone(), - notify: Arc::clone(&self.notify), - stop: Arc::clone(&self.stop), - handler, - dead_letter, - }; - let worker = tokio::spawn(worker.run_loop()); - *self.worker.lock() = Some(worker); + let mut handles = Vec::with_capacity(QUEUE_WORKERS); + for _ in 0..QUEUE_WORKERS { + let worker = QueueWorker { + db_path: self.db_path.clone(), + notify: Arc::clone(&self.notify), + stop: Arc::clone(&self.stop), + handler: Arc::clone(&handler), + dead_letter: Arc::clone(&dead_letter), + }; + handles.push(tokio::spawn(worker.run_loop())); + } + *self.worker.lock() = handles; } pub async fn stop(&self) { self.stop.store(true, Ordering::Relaxed); - self.notify.notify_one(); - if let Some(handle) = self.worker.lock().take() { + self.notify.notify_waiters(); + let handles = std::mem::take(&mut *self.worker.lock()); + for handle in handles { let _ = handle.await; } } @@ -145,8 +164,8 @@ impl PersistentTaskQueue { let payload = payload.to_string(); let db_path = self.db_path.clone(); log::info!("enqueued {id} (run_after {run_after:.1})"); - let result = tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = Connection::open(&db_path)?; + tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { + let conn = open_db(&db_path)?; 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)", @@ -156,14 +175,16 @@ impl PersistentTaskQueue { }) .await .expect("queue insert worker panicked")?; - self.notify.notify_one(); - Ok(result) + // Wake every sleeping worker: with several workers the one that finds + // nothing due must not starve the newly inserted row. + self.notify.notify_waiters(); + Ok(()) } async fn recover_stale(&self) { let db_path = self.db_path.clone(); tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = Connection::open(&db_path)?; + let conn = open_db(&db_path)?; conn.execute( "UPDATE tasks SET status='pending', locked_until=0 WHERE status='in_progress' AND locked_until < ?1", params![now_f64()], @@ -206,11 +227,15 @@ impl QueueWorker { async fn lease_next(&self) -> Option { let db_path = self.db_path.clone(); tokio::task::spawn_blocking(move || -> rusqlite::Result> { - let mut conn = Connection::open(&db_path)?; - let tx = conn.transaction()?; + let mut conn = open_db(&db_path)?; + // 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 + // leases and re-reads the freshest committed state. + let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; let now = now_f64(); let row = tx.query_row( - "SELECT id, payload, attempts FROM tasks WHERE status='pending' AND run_after <= ?1 \ + "SELECT id, payload, attempts FROM tasks WHERE status='pending' AND run_after <= ?1 AND locked_until <= ?1 \ ORDER BY run_after LIMIT 1", params![now], |r| { @@ -251,7 +276,7 @@ impl QueueWorker { async fn earliest_run_after(&self) -> Option { let db_path = self.db_path.clone(); tokio::task::spawn_blocking(move || -> rusqlite::Result> { - let conn = Connection::open(&db_path)?; + let conn = open_db(&db_path)?; let mut stmt = conn.prepare("SELECT MIN(run_after) FROM tasks WHERE status='pending'")?; let mut rows = stmt.query([])?; match rows.next()? { @@ -314,7 +339,7 @@ impl QueueWorker { let db_path = self.db_path.clone(); let id = id.to_string(); tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = Connection::open(&db_path)?; + let conn = open_db(&db_path)?; conn.execute("DELETE FROM tasks WHERE id = ?1", params![id])?; Ok(()) }) @@ -328,7 +353,7 @@ impl QueueWorker { let id = id.to_string(); let payload = payload.to_string(); tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { - let conn = Connection::open(&db_path)?; + let conn = open_db(&db_path)?; 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], @@ -338,7 +363,7 @@ impl QueueWorker { .await .expect("queue reschedule worker panicked") .unwrap_or_else(|e| log::error!("queue reschedule failed: {e}")); - self.notify.notify_one(); + self.notify.notify_waiters(); } } diff --git a/crates/xmedia-bot/src/state.rs b/crates/xmedia-bot/src/state.rs index 932703d..2318da8 100644 --- a/crates/xmedia-bot/src/state.rs +++ b/crates/xmedia-bot/src/state.rs @@ -74,6 +74,10 @@ impl ChatStore { let db_path = self.db_path.clone(); let payload = tokio::task::spawn_blocking(move || -> rusqlite::Result> { let conn = Connection::open(&db_path)?; + // Concurrent handler tasks (batch-forwards) may write chat_state + // while this read runs; without a busy timeout a write lock + // collision fails the query immediately. + 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 rows = stmt.query(params![chat_id.to_string()])?; match rows.next()? { @@ -100,6 +104,7 @@ impl ChatStore { let db_path = self.db_path.clone(); tokio::task::spawn_blocking(move || -> rusqlite::Result<()> { let conn = Connection::open(&db_path)?; + conn.busy_timeout(std::time::Duration::from_secs(5))?; conn.execute( "INSERT OR REPLACE INTO chat_state (chat_id, payload) VALUES (?1, ?2)", params![chat_id.to_string(), payload],