From de9abd129db13dcdf0e69e2c1e2d3aa558fb783a Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Fri, 28 Apr 2023 00:09:41 +0200 Subject: [PATCH] feat(backend): add EXIT_AFTER_NO_JOB_FOR_SECS for ephemeral workers --- README.md | 1 + backend/src/main.rs | 5 ++++- backend/windmill-common/src/lib.rs | 9 +++++++-- backend/windmill-worker/src/worker.rs | 19 ++++++++++++++++++- 4 files changed, 30 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 4df25439d3..189a7c6d22 100644 --- a/README.md +++ b/README.md @@ -342,6 +342,7 @@ it being synced automatically everyday. | WAIT_RESULT_FAST_POLL_INTERVAL_MS | 50 | The time in between polling for the run_wait_result endpoints in fast poll mode | Server | | WAIT_RESULT_SLOW_POLL_INTERVAL_MS | 200 | The time in between polling for the run_wait_result endpoints in fast poll mode | Server | | WAIT_RESULT_FAST_POLL_DURATION_SECS | 2 | The duration of fast poll mode before switching to slow poll | Server | +| EXIT_AFTER_NO_JOB_FOR_SECS | None | Exit worker if no job is received after duration in secs if defined | Worker | ## Run a local dev setup diff --git a/backend/src/main.rs b/backend/src/main.rs index 59dec19995..392c821941 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -102,7 +102,7 @@ async fn main() -> anyhow::Result<()> { } let (tx, rx) = tokio::sync::broadcast::channel::<()>(3); - let shutdown_signal = windmill_common::shutdown_signal(tx); + let shutdown_signal = windmill_common::shutdown_signal(tx.clone(), rx.resubscribe()); #[cfg(feature = "enterprise")] tracing::info!( @@ -169,6 +169,7 @@ Windmill Community Edition {GIT_VERSION} "WAIT_RESULT_FAST_POLL_DURATION_SECS", "WAIT_RESULT_SLOW_POLL_INTERVAL_MS", "WAIT_RESULT_FAST_POLL_INTERVAL_MS", + "EXIT_AFTER_NO_JOB_FOR_SECS", ]); if server_mode || num_workers > 0 { @@ -192,6 +193,8 @@ Windmill Community Edition {GIT_VERSION} rsmq.clone(), ) .await?; + tracing::info!("All workers exited."); + tx.send(())?; // signal server to shutdown } Ok(()) as anyhow::Result<()> }; diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index 7c670a4322..3d95b64086 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -47,7 +47,10 @@ lazy_static::lazy_static! { } #[cfg(feature = "tokio")] -pub async fn shutdown_signal(tx: tokio::sync::broadcast::Sender<()>) -> anyhow::Result<()> { +pub async fn shutdown_signal( + tx: tokio::sync::broadcast::Sender<()>, + mut rx: tokio::sync::broadcast::Receiver<()>, +) -> anyhow::Result<()> { use std::io; use tokio::signal::unix::SignalKind; @@ -61,6 +64,9 @@ pub async fn shutdown_signal(tx: tokio::sync::broadcast::Sender<()>) -> anyhow:: tokio::select! { _ = terminate() => {}, _ = tokio::signal::ctrl_c() => {}, + _ = rx.recv() => { + tracing::info!("shutdown monitor received killpill"); + }, } println!("signal received, starting graceful shutdown"); let _ = tx.send(()); @@ -169,7 +175,6 @@ pub async fn get_latest_deployed_hash_for_path<'c>( Ok((scripts::ScriptHash(script.hash), script.tag)) } - pub async fn get_latest_hash_for_path<'c>( db: &mut sqlx::Transaction<'c, sqlx::Postgres>, w_id: &str, diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 5759fe58c7..dfa3629c44 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -232,6 +232,11 @@ lazy_static::lazy_static! { .map(|e| Some(e)) .unwrap_or(None); + + pub static ref EXIT_AFTER_NO_JOB_FOR_SECS: Option = std::env::var("EXIT_AFTER_NO_JOB_FOR_SECS") + .ok() + .and_then(|x| x.parse::().ok()); + pub static ref CAN_PULL: Arc> = Arc::new(RwLock::new(())); } @@ -453,7 +458,7 @@ pub async fn run_worker = None; loop { if *METRICS_ENABLED { worker_busy.set(0); @@ -595,6 +600,8 @@ pub async fn run_worker { // println!("{:?}", SystemTime::now()); + last_executed_job = None; + let token = create_token_for_owner_in_bg(&db, &job).await; let language = job.language.clone(); let _timer = worker_execution_duration @@ -701,6 +708,16 @@ pub async fn run_worker { + if let Some(secs) = *EXIT_AFTER_NO_JOB_FOR_SECS { + if let Some(lj) = last_executed_job { + if lj.elapsed().as_secs() > secs { + tracing::info!(worker = %worker_name, "no job for {} seconds, exiting", secs); + return true; + } + } else { + last_executed_job = Some(Instant::now()); + } + } let _timer = if *METRICS_ENABLED { Some(Instant::now()) } else { None }; tokio::time::sleep(Duration::from_millis(*SLEEP_QUEUE)).await; _timer.map(|timer| {