fix: prevent workers from being stuck on kill signal

This commit is contained in:
Ruben Fiszel
2025-07-13 11:05:46 +00:00
parent 28562a2bad
commit bc9ad754b0
2 changed files with 21 additions and 14 deletions

View File

@@ -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;
}
}

View File

@@ -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