mirror of
https://github.com/TheFunny/TelegramTwitterMediaBot.git
synced 2026-09-23 23:32:05 +00:00
handlers: replace unbounded per-URL spawn with a bounded job channel
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.
This commit is contained in:
@@ -13,9 +13,52 @@ use teloxide::types::{
|
|||||||
MessageEntityKind, MessageId, ParseMode, Recipient, ReplyParameters,
|
MessageEntityKind, MessageId, ParseMode, Recipient, ReplyParameters,
|
||||||
};
|
};
|
||||||
use teloxide::utils::command::BotCommands;
|
use teloxide::utils::command::BotCommands;
|
||||||
use tokio::sync::Semaphore;
|
|
||||||
use x_media::media::Media;
|
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<parking_lot::Mutex<Option<tokio::sync::mpsc::Sender<UrlJob>>>> =
|
||||||
|
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::<UrlJob>(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<ChatStore> =
|
pub static CHAT_STORE: LazyLock<ChatStore> =
|
||||||
LazyLock::new(|| ChatStore::open("data/task_queue.db").expect("failed to open chat store"));
|
LazyLock::new(|| ChatStore::open("data/task_queue.db").expect("failed to open chat store"));
|
||||||
pub static TASK_QUEUE: LazyLock<PersistentTaskQueue> =
|
pub static TASK_QUEUE: LazyLock<PersistentTaskQueue> =
|
||||||
@@ -24,14 +67,6 @@ pub static LINK_CACHE: LazyLock<LinkCache> =
|
|||||||
LazyLock::new(|| LinkCache::open("data/task_queue.db"));
|
LazyLock::new(|| LinkCache::open("data/task_queue.db"));
|
||||||
pub static CONFIG: LazyLock<Config> = LazyLock::new(Config::load);
|
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)]
|
#[derive(BotCommands, Clone)]
|
||||||
#[command(
|
#[command(
|
||||||
rename_rule = "snake_case",
|
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());
|
log::info!("extracted {} URL(s): {urls:?}", urls.len());
|
||||||
}
|
}
|
||||||
for url in urls {
|
for url in urls {
|
||||||
let bot = bot.clone();
|
// Clone out of the lock: the parking_lot guard is !Send and must
|
||||||
let message = message.clone();
|
// not be held across the await below.
|
||||||
tokio::spawn(async move {
|
let Some(tx) = URL_JOBS.lock().clone() else {
|
||||||
// Held for the whole task; the semaphore is never closed.
|
log::warn!("url workers not started; dropping link");
|
||||||
let _permit = URL_TASKS.acquire().await.expect("URL semaphore closed");
|
break;
|
||||||
url_media(bot, &message, &url).await;
|
};
|
||||||
});
|
let _ = tx.send((bot.clone(), message.clone(), url)).await;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
respond(())
|
respond(())
|
||||||
|
|||||||
@@ -62,6 +62,10 @@ async fn main() {
|
|||||||
.await;
|
.await;
|
||||||
log::info!("task queue worker started");
|
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
|
// Pixiv login validation (user request): a failed login notifies the
|
||||||
// admin and disables pixiv for this process.
|
// admin and disables pixiv for this process.
|
||||||
if site::pixiv::enabled() {
|
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.
|
// Graceful stop (Ctrl+C / SIGTERM): stop the sweep, notify the admin, drain the queue.
|
||||||
log::info!("Stopping bot");
|
log::info!("Stopping bot");
|
||||||
let _ = stop_tx.send(true);
|
let _ = stop_tx.send(true);
|
||||||
|
handlers::stop_url_workers();
|
||||||
if let Some(admin) = CONFIG.admin_ids.first() {
|
if let Some(admin) = CONFIG.admin_ids.first() {
|
||||||
let _ = bot.send_message(ChatId(*admin), "Shutting down...").await;
|
let _ = bot.send_message(ChatId(*admin), "Shutting down...").await;
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user