mirror of
https://github.com/TheFunny/TelegramTwitterMediaBot.git
synced 2026-09-23 23:32:05 +00:00
perf: handle batch-forwarded URLs concurrently with queue workers
This commit is contained in:
@@ -30,9 +30,9 @@ The `x-media` library: `site::fetch(url)` dispatches (in order) twitter → bsky
|
||||
| `crates/x-media/src/site/<twitter\|pixiv\|bsky>/` | One directory per site: `mod.rs` (re-exports), `interface.rs` (PATTERN, `enabled()`, `fetch_from_url()`, site struct, `From<SiteStruct> 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<HashMap>` 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
|
||||
|
||||
@@ -90,7 +90,7 @@ pub async fn fetch(id: &str) -> Result<Tweet, FetchError> {
|
||||
{
|
||||
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
|
||||
|
||||
@@ -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<PersistentTaskQueue> =
|
||||
LazyLock::new(|| PersistentTaskQueue::new("data/task_queue.db"));
|
||||
pub static CONFIG: LazyLock<Config> = 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<Semaphore> = 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(())
|
||||
|
||||
@@ -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<Notify>,
|
||||
stop: Arc<AtomicBool>,
|
||||
worker: Mutex<Option<JoinHandle<()>>>,
|
||||
worker: Mutex<Vec<JoinHandle<()>>>,
|
||||
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<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(
|
||||
"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<H, F, D, G>(&self, handler: H, dead_letter: D)
|
||||
where
|
||||
@@ -114,21 +129,25 @@ impl PersistentTaskQueue {
|
||||
let dead_letter: Arc<DeadLetter> =
|
||||
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<LeasedRow> {
|
||||
let db_path = self.db_path.clone();
|
||||
tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<LeasedRow>> {
|
||||
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<f64> {
|
||||
let db_path = self.db_path.clone();
|
||||
tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<f64>> {
|
||||
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();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -74,6 +74,10 @@ impl ChatStore {
|
||||
let db_path = self.db_path.clone();
|
||||
let payload = tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<String>> {
|
||||
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],
|
||||
|
||||
Reference in New Issue
Block a user