From 16ed53fead03e51f8b2b01d216c7cbe3a2faef66 Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Sat, 8 Aug 2026 20:15:40 +0800 Subject: [PATCH] handlers: replace unbounded per-URL spawn with a bounded job channel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 8-permit semaphore was acquired inside the spawned task, so a burst queued unlimited tasks (each cloning Bot+Message) and nothing tracked them at shutdown — in-flight sends fired after the stop notice. URL work now flows through a 256-slot mpsc drained by 8 workers started from main; a full channel backpressures the per-chat handler, and shutdown sets URL_STOP so workers stop pulling. --- crates/xmedia-bot/src/handlers.rs | 67 +++++++++++++++++++++++-------- crates/xmedia-bot/src/main.rs | 5 +++ 2 files changed, 56 insertions(+), 16 deletions(-) diff --git a/crates/xmedia-bot/src/handlers.rs b/crates/xmedia-bot/src/handlers.rs index 10707c9..c799766 100644 --- a/crates/xmedia-bot/src/handlers.rs +++ b/crates/xmedia-bot/src/handlers.rs @@ -13,9 +13,52 @@ use teloxide::types::{ MessageEntityKind, MessageId, ParseMode, Recipient, ReplyParameters, }; use teloxide::utils::command::BotCommands; -use tokio::sync::Semaphore; use x_media::media::Media; +/// One URL job: bot handle + the message + the extracted URL. +type UrlJob = (Bot, Message, String); +/// Bounded channel of URL jobs drained by [`start_url_workers`]. The bound +/// caps both queued memory and shutdown backlog; a full channel applies +/// backpressure to the per-chat handler instead of spawning unbounded tasks. +static URL_JOBS: LazyLock>>> = + LazyLock::new(|| parking_lot::Mutex::new(None)); +/// Set by main's shutdown sequence; workers stop pulling new jobs. +static URL_STOP: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + +/// Worker count draining URL jobs; keeps the old 8-permit concurrency cap +/// while bounding how many jobs can be queued at all. +const URL_WORKERS: usize = 8; + +/// Starts the URL job workers (called once from main after the queue starts). +/// 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); the workers add throughput, and FIFO order preserves per-message +/// URL order. +pub async fn start_url_workers() { + let (tx, rx) = tokio::sync::mpsc::channel::(256); + *URL_JOBS.lock() = Some(tx); + let rx = std::sync::Arc::new(tokio::sync::Mutex::new(rx)); + for _ in 0..URL_WORKERS { + let rx = std::sync::Arc::clone(&rx); + tokio::spawn(async move { + while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) { + let job = rx.lock().await.recv().await; + match job { + Some((bot, message, url)) => url_media(bot, &message, &url).await, + None => break, + } + } + }); + } +} + +/// Stops URL workers (drains up to the 256 queued jobs, then exits). +pub fn stop_url_workers() { + URL_STOP.store(true, std::sync::atomic::Ordering::Relaxed); +} + pub static CHAT_STORE: LazyLock = LazyLock::new(|| ChatStore::open("data/task_queue.db").expect("failed to open chat store")); pub static TASK_QUEUE: LazyLock = @@ -24,14 +67,6 @@ pub static LINK_CACHE: LazyLock = LazyLock::new(|| LinkCache::open("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", @@ -721,13 +756,13 @@ pub async fn message_handler(bot: Bot, message: Message) -> Result<(), RequestEr log::info!("extracted {} URL(s): {urls:?}", urls.len()); } for url in urls { - 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; - }); + // Clone out of the lock: the parking_lot guard is !Send and must + // not be held across the await below. + let Some(tx) = URL_JOBS.lock().clone() else { + log::warn!("url workers not started; dropping link"); + break; + }; + let _ = tx.send((bot.clone(), message.clone(), url)).await; } } respond(()) diff --git a/crates/xmedia-bot/src/main.rs b/crates/xmedia-bot/src/main.rs index 39c0de5..5e67779 100644 --- a/crates/xmedia-bot/src/main.rs +++ b/crates/xmedia-bot/src/main.rs @@ -62,6 +62,10 @@ async fn main() { .await; log::info!("task queue worker started"); + // URL job workers: bounded channel + fixed pool for per-URL work. + handlers::start_url_workers().await; + log::info!("url workers started"); + // Pixiv login validation (user request): a failed login notifies the // admin and disables pixiv for this process. if site::pixiv::enabled() { @@ -174,6 +178,7 @@ async fn main() { // Graceful stop (Ctrl+C / SIGTERM): stop the sweep, notify the admin, drain the queue. log::info!("Stopping bot"); let _ = stop_tx.send(true); + handlers::stop_url_workers(); if let Some(admin) = CONFIG.admin_ids.first() { let _ = bot.send_message(ChatId(*admin), "Shutting down...").await; }