diff --git a/backend/src/lib.rs b/backend/src/lib.rs index cdd89364e5..bcb3ae0975 100644 --- a/backend/src/lib.rs +++ b/backend/src/lib.rs @@ -199,14 +199,13 @@ pub async fn run_server( Ok(()) } -pub fn monitor_db(db: &DB, timeout: i32, tx: tokio::sync::broadcast::Sender<()>) { +pub fn monitor_db(db: &DB, timeout: i32, rx: tokio::sync::broadcast::Receiver<()>) { let db1 = db.clone(); let db2 = db.clone(); - let rx1 = tx.subscribe(); - let rx2 = tx.subscribe(); + let rx2 = rx.resubscribe(); - tokio::spawn(async move { worker::restart_zombie_jobs_periodically(&db1, timeout, rx1).await }); + tokio::spawn(async move { worker::restart_zombie_jobs_periodically(&db1, timeout, rx).await }); tokio::spawn(async move { users::delete_expired_items_perdiodically(&db2, rx2).await }); } @@ -219,7 +218,7 @@ pub async fn run_workers( base_url: String, disable_nuser: bool, disable_nsjail: bool, - tx: tokio::sync::broadcast::Sender<()>, + rx: tokio::sync::broadcast::Receiver<()>, ) -> anyhow::Result<()> { let instance_name = rd_string(5); @@ -234,8 +233,8 @@ pub async fn run_workers( let instance_name = instance_name.clone(); let worker_name = format!("dt-worker-{}-{}", &instance_name, rd_string(5)); let ip = ip.clone(); - let tx = tx.clone(); let base_url = base_url.clone(); + let rx = rx.resubscribe(); handles.push(tokio::spawn(async move { tracing::info!(addr = %addr.to_string(), worker = %worker_name, "starting worker"); worker::run_worker( @@ -250,7 +249,7 @@ pub async fn run_workers( &base_url, disable_nuser, disable_nsjail, - tx, + rx, ) .await })); diff --git a/backend/src/main.rs b/backend/src/main.rs index fc4cf0f8b0..bef233fd05 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -40,7 +40,7 @@ async fn main() -> anyhow::Result<()> { } let (tx, rx) = tokio::sync::broadcast::channel::<()>(3); - let shutdown_signal = windmill::shutdown_signal(tx.clone()); + let shutdown_signal = windmill::shutdown_signal(tx); if server_mode || monitor_mode || num_workers > 0 { let addr = SocketAddr::from(([0, 0, 0, 0], 8000)); @@ -61,7 +61,7 @@ async fn main() -> anyhow::Result<()> { server: "smtp.gmail.com".to_string(), password: std::env::var("SMTP_PASSWORD").unwrap_or("NOPASS".to_string()), }, - rx, + rx.resubscribe(), ) .await?; } @@ -100,7 +100,7 @@ async fn main() -> anyhow::Result<()> { base_url, disable_nuser, disable_nsjail, - tx.clone(), + rx.resubscribe(), ) .await?; } @@ -109,14 +109,14 @@ async fn main() -> anyhow::Result<()> { let monitor_f = async { if monitor_mode { - windmill::monitor_db(&db, timeout, tx.clone()); + windmill::monitor_db(&db, timeout, rx.resubscribe()); } Ok(()) as anyhow::Result<()> }; let metrics_f = async { match metrics_addr { - Some(addr) => windmill::serve_metrics(addr, tx.subscribe()) + Some(addr) => windmill::serve_metrics(addr, rx.resubscribe()) .await .map_err(anyhow::Error::from), None => Ok(()), diff --git a/backend/src/worker.rs b/backend/src/worker.rs index 26be75d0ab..785695f4cb 100644 --- a/backend/src/worker.rs +++ b/backend/src/worker.rs @@ -72,7 +72,7 @@ pub async fn run_worker( base_url: &str, disable_nuser: bool, disable_nsjail: bool, - tx: tokio::sync::broadcast::Sender<()>, + mut rx: tokio::sync::broadcast::Receiver<()>, ) { let worker_dir = format!("{TMP_DIR}/{worker_name}"); tracing::debug!(worker_dir = %worker_dir, worker_name = %worker_name, "Creating worker dir"); @@ -123,8 +123,6 @@ pub async fn run_worker( .expect("register prometheus metric"); let mut jobs_executed = 0; - let mut rx = tx.subscribe(); - drop(tx); loop { if last_ping.elapsed().as_secs() > NUM_SECS_ENV_CHECK { @@ -1402,7 +1400,7 @@ def main(): let mut listener = PgListener::connect_with(db).await.unwrap(); listener.listen("insert on completed_job").await.unwrap(); - let (tx, _rx) = tokio::sync::broadcast::channel(1); + let (tx, rx) = tokio::sync::broadcast::channel(1); /* drop tx at the end of this block to close the channel and stop the worker */ let worker = { @@ -1429,7 +1427,7 @@ def main(): base_url, disable_nuser, disable_nsjail, - tx.clone(), + rx, ) };