From bc9ad754b0c39848bbe8da61370619b8edb7d4f7 Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@windmill.dev>
Date: Sun, 13 Jul 2025 11:05:46 +0000
Subject: [PATCH] fix: prevent workers from being stuck on kill signal

---
 .../windmill-worker/src/result_processor.rs   |  6 +++-
 backend/windmill-worker/src/worker.rs         | 29 ++++++++++---------
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/backend/windmill-worker/src/result_processor.rs b/backend/windmill-worker/src/result_processor.rs
index ce314a00e6..c3e23411dd 100644
--- a/backend/windmill-worker/src/result_processor.rs
+++ b/backend/windmill-worker/src/result_processor.rs
@@ -142,6 +142,9 @@ pub fn start_background_processor(
 
         //if we have been killed, we want to drain the queue of jobs
         while let Some(sr) = {
+            if has_been_killed {
+                tracing::info!("bg processor is killed, draining. same_worker_queue_size: {}, unbounded_rx: {}, bounded_rx: {}", same_worker_queue_size.load(Ordering::SeqCst), unbounded_rx.len(), bounded_rx.len())
+            }
             if has_been_killed && same_worker_queue_size.load(Ordering::SeqCst) == 0 {
                 unbounded_rx
                     .try_recv()
@@ -157,8 +160,8 @@ pub fn start_background_processor(
                     result = bounded_rx.recv_async() => {
                         result.ok().map(JobCompletedRx::JobCompleted)
                     }
-
                     _ = killpill_rx.recv() => {
+                        tracing::info!("bg processor received killpill signal, queuing killpill job");
                         Some(JobCompletedRx::Killpill)
                     }
                 }
@@ -251,6 +254,7 @@ pub fn start_background_processor(
                     }
                 }
                 JobCompletedRx::Killpill => {
+                    tracing::info!("killpill job received, processing only same worker jobs");
                     has_been_killed = true;
                 }
             }
diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs
index 8f50053b60..b3364d962c 100644
--- a/backend/windmill-worker/src/worker.rs
+++ b/backend/windmill-worker/src/worker.rs
@@ -1324,19 +1324,10 @@ pub async fn run_worker(
     let mut killed_but_draining_same_worker_jobs = false;
 
     let mut killpill_rx2 = killpill_rx.resubscribe();
+    
     loop {
         #[cfg(feature = "enterprise")]
         {
-            if let Ok(_) = killpill_rx.try_recv() {
-                tracing::info!(worker = %worker_name, hostname = %hostname, "killpill received on worker waiting for valid key");
-                if send_result.is_some() {
-                    job_completed_tx
-                        .kill()
-                        .await
-                        .expect("send kill to job completed tx");
-                }
-                break;
-            }
             let valid_key = *LICENSE_KEY_VALID.read().await;
 
             if !valid_key {
@@ -1344,8 +1335,16 @@ pub async fn run_worker(
                     worker = %worker_name, hostname = %hostname,
                     "Invalid license key, workers require a valid license key, sleeping for 10s waiting for valid key to be set"
                 );
-                tokio::time::sleep(Duration::from_secs(10)).await;
-                continue;
+                tokio::select! {
+                    _ = tokio::time::sleep(Duration::from_secs(10)) => {
+                        tracing::info!(worker = %worker_name, hostname = %hostname, "sleeping for 10s waiting for valid key to be set");
+                        continue;
+                    }
+                    _ = killpill_rx.recv() => {
+                        tracing::info!(worker = %worker_name, hostname = %hostname, "killpill received while waiting for valid key, exiting");
+                        break;
+                    }
+                }
             }
         }
 
@@ -1459,9 +1458,13 @@ pub async fn run_worker(
                         .map_err(|e| error::Error::InternalErr(e.to_string()))
                         .map(|x: Option<JobAndPerms>| x.map(|y| NextJob::Http(y))),
                 }
-            } else if let Ok(_) = killpill_rx.try_recv() {
+            } else if match killpill_rx.try_recv() {
+                Ok(_) | Err(broadcast::error::TryRecvError::Closed) => true,
+                _ => false,
+            } {
                 if !killed_but_draining_same_worker_jobs {
                     killed_but_draining_same_worker_jobs = true;
+                    tracing::info!(worker = %worker_name, hostname = %hostname, "killpill received in worker main loop, sending killpill job");
                     job_completed_tx
                         .kill()
                         .await