From 6911e9146e1e207880208622933f37294ca18162 Mon Sep 17 00:00:00 2001
From: YoursFunny <admin@yoursfunny.top>
Date: Thu, 13 Aug 2026 23:15:14 +0800
Subject: [PATCH] fix(handlers): stop URL workers by closing the job channel

The old stop only set an atomic flag checked between jobs: a worker
blocked in recv() never woke (the channel was never closed), and queued
jobs were neither drained nor abandoned in a defined way despite the
"drains up to 256 jobs" comment. Now stop_url_workers sets the flag,
drops the sender so blocked recv() calls wake with None, and awaits the
worker JoinHandles (each finishes its in-flight job first). main awaits
it inside the existing 30s shutdown timeout.
---
 crates/xmedia-bot/src/handlers.rs | 28 ++++++++++++++++++++++++----
 crates/xmedia-bot/src/main.rs     |  2 +-
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/crates/xmedia-bot/src/handlers.rs b/crates/xmedia-bot/src/handlers.rs
index 7b6a389..071b82b 100644
--- a/crates/xmedia-bot/src/handlers.rs
+++ b/crates/xmedia-bot/src/handlers.rs
@@ -26,6 +26,10 @@ static URL_JOBS: LazyLock<parking_lot::Mutex<Option<tokio::sync::mpsc::Sender<Ur
 /// Set by main's shutdown sequence; workers stop pulling new jobs.
 static URL_STOP: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
 
+/// JoinHandles of the URL workers, awaited by [`stop_url_workers`].
+static URL_WORKER_HANDLES: LazyLock<parking_lot::Mutex<Option<Vec<tokio::task::JoinHandle<()>>>>> =
+    LazyLock::new(|| parking_lot::Mutex::new(None));
+
 /// 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;
@@ -40,9 +44,10 @@ 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));
+    let mut handles = Vec::with_capacity(URL_WORKERS);
     for _ in 0..URL_WORKERS {
         let rx = std::sync::Arc::clone(&rx);
-        tokio::spawn(async move {
+        handles.push(tokio::spawn(async move {
             while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) {
                 let job = rx.lock().await.recv().await;
                 match job {
@@ -50,13 +55,28 @@ pub async fn start_url_workers() {
                     None => break,
                 }
             }
-        });
+        }));
     }
+    *URL_WORKER_HANDLES.lock() = Some(handles);
 }
 
-/// Stops URL workers (drains up to the 256 queued jobs, then exits).
-pub fn stop_url_workers() {
+/// Stops the URL workers: sets the stop flag, drops the job channel (so
+/// workers blocked in \`recv()\` wake with \`None\` and exit) and awaits the
+/// worker tasks. Each worker finishes its in-flight job first; jobs still
+/// queued in the channel are abandoned (the old implementation neither
+/// drained them nor woke blocked workers — it only set a flag checked
+/// between jobs).
+pub async fn stop_url_workers() {
     URL_STOP.store(true, std::sync::atomic::Ordering::Relaxed);
+    // Dropping the sender makes every worker's recv() return None.
+    *URL_JOBS.lock() = None;
+    // Take the handles first so the lock guard drops before the awaits.
+    let handles = URL_WORKER_HANDLES.lock().take();
+    if let Some(handles) = handles {
+        for handle in handles {
+            let _ = handle.await;
+        }
+    }
 }
 
 pub static CHAT_STORE: LazyLock<ChatStore> =
diff --git a/crates/xmedia-bot/src/main.rs b/crates/xmedia-bot/src/main.rs
index f4be25f..6f32e66 100644
--- a/crates/xmedia-bot/src/main.rs
+++ b/crates/xmedia-bot/src/main.rs
@@ -185,7 +185,7 @@ async fn main() {
     const SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
     let shutdown = async {
         let _ = stop_tx.send(true);
-        handlers::stop_url_workers();
+        handlers::stop_url_workers().await;
         if let Some(admin) = CONFIG.admin_ids.first() {
             let _ = bot.send_message(ChatId(*admin), "Shutting down...").await;
         }