From 2fcce4526a239437221e37cfd4adfd4da616cf19 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Thu, 13 Aug 2026 06:14:55 +0200 Subject: [PATCH] feat: add EXIT_AFTER_N_JOBS worker mode for environment cleanup (#10671) * feat: add EXIT_AFTER_N_JOBS worker mode for environment cleanup Co-Authored-By: Claude Opus 5 (1M context) * fix: address review findings on the EXIT_AFTER_N_JOBS worker mode Co-Authored-By: Claude Opus 5 (1M context) * fix: address round-2 review findings on EXIT_AFTER_N_JOBS Co-Authored-By: Claude Opus 5 (1M context) * fix: address round-3 review findings on EXIT_AFTER_N_JOBS Co-Authored-By: Claude Opus 5 (1M context) * fix: bound WORKER_SUFFIX length and document the same-worker drain Co-Authored-By: Claude Opus 5 (1M context) * fix: validate the assembled worker name length Co-Authored-By: Claude Opus 5 (1M context) --------- Co-authored-by: Claude Opus 5 (1M context) --- README.md | 2 + ...b6ed918af2df012de33acb3befd3fcccbc257.json | 25 --- ...35baa410f93be3cac4b8dfd33ddbef446d3fe.json | 33 +++ backend/src/main.rs | 18 +- .../windmill-common/src/global_settings.rs | 2 + backend/windmill-common/src/utils.rs | 164 ++++++++++++++ backend/windmill-common/src/worker.rs | 58 ++++- backend/windmill-worker/src/worker.rs | 210 +++++++++++++++++- backend/windmill-worker/src/worker_utils.rs | 12 +- 9 files changed, 472 insertions(+), 52 deletions(-) delete mode 100644 backend/.sqlx/query-298fa4f8eb05b4c3f33b608b0cdb6ed918af2df012de33acb3befd3fcccbc257.json create mode 100644 backend/.sqlx/query-9d9fbcb598c582a29d65be87a1e35baa410f93be3cac4b8dfd33ddbef446d3fe.json diff --git a/README.md b/README.md index f51d184943..2d1d4d62ac 100644 --- a/README.md +++ b/README.md @@ -260,6 +260,8 @@ On self-hosted instances, you might want to import all the approved resource typ | NATIVE_MODE | false | Enable native mode: sets NUM_WORKERS=8, rejects non-native jobs (nativets, postgresql, mysql, etc.) | Worker | | SLEEP_QUEUE | 50 | The number of ms to sleep in between the last check for new jobs in the DB. It is multiplied by NUM_WORKERS such that in average, for one worker instance, there is one pull every SLEEP_QUEUE ms. | Worker | | KEEP_JOB_DIR | false | Keep the job directory after the job is done. Useful for debugging. | Worker | +| EXIT_AFTER_N_JOBS | None | Exit the worker process after it has executed that many jobs, so that a supervisor restarts it and no process runs more than that many, bar the steps of a same-worker flow it has started, which it always finishes (set it to 1 for a process per job; jobs handed to a dedicated worker, and the worker's own init and periodic scripts, do not count). For deployments that isolate executions by process lifetime rather than with nsjail; note that a container restart resets the process, not the container filesystem, so caches and `/tmp` survive it. The worker name is then derived from the hostname instead of being random, so the restarted worker keeps its row in the workers list (an agent worker keeps the row but restarts its job count). Use one worker per process: workers of one process share its environment, so the first to reach the limit shuts the others down too. | Worker | +| WORKER_SUFFIX | None | Pins the last part of the worker name, which is otherwise random, so that a restarted worker keeps its row in the workers list. Only needed when several worker processes of the same worker group run on one host, since the name is derived from the hostname: give each of them a distinct value, as two processes sharing one must never happen. At most 64 letters, digits and underscores; anything else is refused at startup. | Worker | | LICENSE_KEY (EE only) | None | License key checked at startup for the Enterprise Edition of Windmill | Worker | | SLACK_SIGNING_SECRET | None | The signing secret of your Slack app. See [Slack documentation](https://api.slack.com/authentication/verifying-requests-from-slack) | Server | | COOKIE_DOMAIN | None | The domain of the cookie. If not set, the cookie will be set by the browser based on the full origin | Server | diff --git a/backend/.sqlx/query-298fa4f8eb05b4c3f33b608b0cdb6ed918af2df012de33acb3befd3fcccbc257.json b/backend/.sqlx/query-298fa4f8eb05b4c3f33b608b0cdb6ed918af2df012de33acb3befd3fcccbc257.json deleted file mode 100644 index e0e7a16624..0000000000 --- a/backend/.sqlx/query-298fa4f8eb05b4c3f33b608b0cdb6ed918af2df012de33acb3befd3fcccbc257.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "INSERT INTO worker_ping (worker_instance, worker, ip, custom_tags, worker_group, dedicated_worker, dedicated_workers, wm_version, vcpus, memory, job_isolation, native_mode) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) ON CONFLICT (worker)\n DO UPDATE set ip = EXCLUDED.ip, custom_tags = EXCLUDED.custom_tags, worker_group = EXCLUDED.worker_group, dedicated_workers = EXCLUDED.dedicated_workers, native_mode = EXCLUDED.native_mode", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Varchar", - "Varchar", - "Varchar", - "TextArray", - "Varchar", - "Varchar", - "TextArray", - "Varchar", - "Int8", - "Int8", - "Text", - "Bool" - ] - }, - "nullable": [] - }, - "hash": "298fa4f8eb05b4c3f33b608b0cdb6ed918af2df012de33acb3befd3fcccbc257" -} diff --git a/backend/.sqlx/query-9d9fbcb598c582a29d65be87a1e35baa410f93be3cac4b8dfd33ddbef446d3fe.json b/backend/.sqlx/query-9d9fbcb598c582a29d65be87a1e35baa410f93be3cac4b8dfd33ddbef446d3fe.json new file mode 100644 index 0000000000..2c95e06391 --- /dev/null +++ b/backend/.sqlx/query-9d9fbcb598c582a29d65be87a1e35baa410f93be3cac4b8dfd33ddbef446d3fe.json @@ -0,0 +1,33 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO worker_ping (worker_instance, worker, ip, custom_tags, worker_group, dedicated_worker, dedicated_workers, wm_version, vcpus, memory, job_isolation, native_mode) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) ON CONFLICT (worker)\n DO UPDATE set ping_at = now(), worker_instance = EXCLUDED.worker_instance, ip = EXCLUDED.ip, custom_tags = EXCLUDED.custom_tags, worker_group = EXCLUDED.worker_group, dedicated_worker = EXCLUDED.dedicated_worker, dedicated_workers = EXCLUDED.dedicated_workers, wm_version = EXCLUDED.wm_version, vcpus = COALESCE(EXCLUDED.vcpus, worker_ping.vcpus), memory = COALESCE(EXCLUDED.memory, worker_ping.memory), job_isolation = EXCLUDED.job_isolation, native_mode = EXCLUDED.native_mode, current_job_id = NULL, current_job_workspace_id = NULL\n RETURNING jobs_executed", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "jobs_executed", + "type_info": "Int4" + } + ], + "parameters": { + "Left": [ + "Varchar", + "Varchar", + "Varchar", + "TextArray", + "Varchar", + "Varchar", + "TextArray", + "Varchar", + "Int8", + "Int8", + "Text", + "Bool" + ] + }, + "nullable": [ + false + ] + }, + "hash": "9d9fbcb598c582a29d65be87a1e35baa410f93be3cac4b8dfd33ddbef446d3fe" +} diff --git a/backend/src/main.rs b/backend/src/main.rs index 51febad6f0..43a92ce4bf 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -70,12 +70,12 @@ use windmill_common::{ stats_oss::schedule_stats, triggers::TriggerKind, utils::{ - create_default_worker_suffix, worker_name_with_suffix, Mode, GIT_VERSION, HOSTNAME, - MODE_AND_ADDONS, + checked_worker_name, resolve_worker_suffix, Mode, GIT_VERSION, HOSTNAME, MODE_AND_ADDONS, }, worker::{ - is_native_mode_from_env, reload_custom_tags_setting, Connection, HUB_CACHE_DIR, - HUB_RT_CACHE_DIR, NATIVE_MODE_RESOLVED, TMP_LOGS_DIR, WINDMILL_DIR, WORKER_GROUP, + is_native_mode_from_env, reload_custom_tags_setting, validate_worker_lifecycle_env, + Connection, HUB_CACHE_DIR, HUB_RT_CACHE_DIR, NATIVE_MODE_RESOLVED, TMP_LOGS_DIR, + WINDMILL_DIR, WORKER_GROUP, }, KillpillSender, DEFAULT_HUB_BASE_URL, INSTANCE_NAME, METRICS_ENABLED, }; @@ -782,6 +782,8 @@ async fn windmill_main() -> anyhow::Result<()> { } } + validate_worker_lifecycle_env()?; + let server_mode = !std::env::var("DISABLE_SERVER") .ok() .and_then(|x| x.parse::().ok()) @@ -813,7 +815,7 @@ async fn windmill_main() -> anyhow::Result<()> { "Creating http client for cluster using base internal url {}", agent_config.base_internal_url ); - let suffix = create_default_worker_suffix(&hostname); + let suffix = resolve_worker_suffix(&hostname, 1)?; ( Connection::Http(agent_config.build_http_client(&suffix)), Some(suffix), @@ -1332,7 +1334,7 @@ Windmill Community Edition {GIT_VERSION} let suffix = if i == 0 && first_suffix.is_some() { first_suffix.as_ref().unwrap().clone() } else { - create_default_worker_suffix(&hostname) + resolve_worker_suffix(&hostname, i as usize + 1)? }; let worker_conn = WorkerConn { @@ -1346,11 +1348,11 @@ Windmill Community Edition {GIT_VERSION} .build_http_client(&suffix), ) }, - worker_name: worker_name_with_suffix( + worker_name: checked_worker_name( mode == Mode::Agent, WORKER_GROUP.as_str(), &suffix, - ), + )?, }; workers.push(worker_conn); } diff --git a/backend/windmill-common/src/global_settings.rs b/backend/windmill-common/src/global_settings.rs index c9b5a46989..b7816d103a 100644 --- a/backend/windmill-common/src/global_settings.rs +++ b/backend/windmill-common/src/global_settings.rs @@ -294,6 +294,8 @@ pub const ENV_SETTINGS: &[&str] = &[ "WAIT_RESULT_SLOW_POLL_INTERVAL_MS", "WAIT_RESULT_FAST_POLL_INTERVAL_MS", "EXIT_AFTER_NO_JOB_FOR_SECS", + "EXIT_AFTER_N_JOBS", + "WORKER_SUFFIX", "REQUEST_SIZE_LIMIT", "CREATE_WORKSPACE_REQUIRE_SUPERADMIN", "GLOBAL_ERROR_HANDLER_PATH_IN_ADMINS_WORKSPACE", diff --git a/backend/windmill-common/src/utils.rs b/backend/windmill-common/src/utils.rs index 6d13515d1f..db0de177db 100644 --- a/backend/windmill-common/src/utils.rs +++ b/backend/windmill-common/src/utils.rs @@ -12,6 +12,7 @@ use crate::ee_oss::LICENSE_KEY_ID; use crate::ee_oss::{send_critical_alert, CriticalAlertKind}; use crate::error::{to_anyhow, Error, Result}; use crate::global_settings::UNIQUE_ID_SETTING; +use crate::worker::{EXIT_AFTER_N_JOBS, WORKER_SUFFIX}; use crate::DB; use anyhow::Context; use gethostname::gethostname; @@ -350,6 +351,9 @@ fn instance_name(hostname: &str) -> String { } const DEFAULT_WORKER_SUFFIX_LEN: usize = 5; +const MAX_WORKER_SUFFIX_LABEL_LEN: usize = 64; +/// `worker_ping.worker` is a `VARCHAR(255)`. +const MAX_WORKER_NAME_LEN: usize = 255; pub const SSH_AGENT_WORKER_SUFFIX: &'static str = "/ssh"; pub fn create_worker_suffix(hostname: &str, rd_string_len: usize) -> String { @@ -361,6 +365,68 @@ pub fn create_default_worker_suffix(hostname: &str) -> String { create_worker_suffix(hostname, DEFAULT_WORKER_SUFFIX_LEN) } +/// Same shape as [`create_default_worker_suffix`] but derived from the hostname and the +/// worker index instead of randomness, so the process gets the same worker name every time +/// it starts on that host. The index is folded into the digest rather than appended so that +/// the name still has exactly one suffix segment, which is what +/// [`retrieve_common_worker_prefix`] (the interactive shell tag) strips off. +fn create_stable_worker_suffix(hostname: &str, index: usize) -> String { + let digest = calculate_hash(&format!("{hostname}#{index}")); + format!( + "{}-{}", + instance_name(hostname), + &digest[..DEFAULT_WORKER_SUFFIX_LEN] + ) +} + +/// Suffix of the name of the `index`-th (1-based) worker of this process. +/// +/// A worker name is the primary key of its `worker_ping` row, so it decides whether a +/// restarted process appears as a new worker or resumes the previous one. Random by default +/// (two processes must never share a row); deterministic when the worker is expected to +/// restart in place, which is the case for `EXIT_AFTER_N_JOBS`. `WORKER_SUFFIX` overrides +/// both, for hosts running several worker processes of the same worker group: the hostname +/// alone cannot tell those apart. +pub fn resolve_worker_suffix(hostname: &str, index: usize) -> anyhow::Result { + Ok(match &*WORKER_SUFFIX { + Some(label) => create_labelled_worker_suffix(hostname, label, index)?, + None if EXIT_AFTER_N_JOBS.is_some() => create_stable_worker_suffix(hostname, index), + None => create_default_worker_suffix(hostname), + }) +} + +/// The operator's label only has to tell the worker processes of one host apart, so it is +/// appended to the stable suffix rather than replacing it: the digest is what keeps two hosts +/// whose names end on the same segment (`worker-east-1`, `worker-west-1`) from sharing an +/// identity, and it already folds in the worker index. A `-` in the label would add a segment +/// to the worker name, which [`retrieve_common_worker_prefix`] reads as the part to strip; +/// rejected rather than rewritten, since the point of the label is that two different ones +/// give two different names. +fn create_labelled_worker_suffix( + hostname: &str, + label: &str, + index: usize, +) -> anyhow::Result { + if !label.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') { + return Err(anyhow::anyhow!( + "WORKER_SUFFIX must only contain ASCII letters, digits and underscores, got '{label}'" + )); + } + // The worker name is a `VARCHAR(255)` primary key and a component of the worker + // directory's path: a label long enough to blow either only surfaces at the initial ping, + // which the worker `expect`s. + if label.len() > MAX_WORKER_SUFFIX_LABEL_LEN { + return Err(anyhow::anyhow!( + "WORKER_SUFFIX must be at most {MAX_WORKER_SUFFIX_LABEL_LEN} characters, got {}", + label.len() + )); + } + Ok(format!( + "{}_{label}", + create_stable_worker_suffix(hostname, index) + )) +} + pub fn worker_name_with_suffix(is_agent: bool, worker_group: &str, suffix: &str) -> String { if is_agent { format!("{}-{}-{}", AGENT_WORKER_NAME_PREFIX, worker_group, suffix) @@ -369,6 +435,26 @@ pub fn worker_name_with_suffix(is_agent: bool, worker_group: &str, suffix: &str) } } +/// The name is the `VARCHAR(255)` primary key of `worker_ping` and a component of the worker +/// directory's path, and every part of it comes from the environment (`WORKER_GROUP`, +/// hostname, `WORKER_SUFFIX`). A name that does not fit has to stop the process here rather +/// than at the directory it creates or the initial ping it `expect`s. +pub fn checked_worker_name( + is_agent: bool, + worker_group: &str, + suffix: &str, +) -> anyhow::Result { + let name = worker_name_with_suffix(is_agent, worker_group, suffix); + if name.len() > MAX_WORKER_NAME_LEN { + return Err(anyhow::anyhow!( + "worker name '{name}' is {} characters, more than the {MAX_WORKER_NAME_LEN} a worker \ + name may have: shorten WORKER_GROUP or WORKER_SUFFIX", + name.len() + )); + } + Ok(name) +} + pub fn retrieve_common_worker_prefix(worker_name: &str) -> String { let (prefix, _) = worker_name.rsplit_once('-').unzip(); @@ -1471,6 +1557,84 @@ pub fn truncate_with_ellipsis(s: &str, max_chars: usize) -> String { mod tests { use super::*; + /// A worker that restarts must land on the exact same name to reclaim its `worker_ping` + /// row, while still never colliding with the other workers of its own process. The + /// suffix must also stay a single `-` segment, which is what the interactive shell tag + /// strips off. + #[test] + fn stable_worker_suffix_is_per_host_and_per_index() { + let first = create_stable_worker_suffix("wm-worker-7d8f9c-abcde", 1); + assert_eq!( + first, + create_stable_worker_suffix("wm-worker-7d8f9c-abcde", 1) + ); + assert_ne!( + first, + create_stable_worker_suffix("wm-worker-7d8f9c-abcde", 2) + ); + assert_ne!( + first, + create_stable_worker_suffix("wm-worker-7d8f9c-fghij", 1) + ); + assert_eq!( + retrieve_common_worker_prefix(&worker_name_with_suffix(false, "default", &first)), + "wk-default-abcde" + ); + } + + /// The operator's label has to survive into the name for them to recognize the worker, + /// while staying one segment so it does not shift the interactive shell tag, and it must + /// still leave the workers of one process with distinct names. Two processes given + /// different labels must never end up with the same one, which is why a label that does + /// not fit a single segment is rejected instead of being rewritten into one that does. + #[test] + fn labelled_worker_suffix_stays_one_segment() { + let label = + |hostname, label, index| create_labelled_worker_suffix(hostname, label, index).unwrap(); + let first = label("wm-worker-abcde", "slot_a", 1); + assert!( + first.starts_with("abcde-") && first.ends_with("_slot_a"), + "{first}" + ); + assert_eq!(first, label("wm-worker-abcde", "slot_a", 1)); + // Neither another worker of this process, nor another label, nor a host whose name + // happens to end on the same segment, may land on this identity. + for other in [ + label("wm-worker-abcde", "slot_a", 2), + label("wm-worker-abcde", "slot_b", 1), + label("wm-worker-east-abcde", "slot_a", 1), + ] { + assert_ne!(first, other); + assert_eq!( + retrieve_common_worker_prefix(&worker_name_with_suffix(false, "default", &other)), + "wk-default-abcde" + ); + } + for rejected in [ + "slot-a", + "slot a", + "slot.a", + "slot/a", + &"s".repeat(MAX_WORKER_SUFFIX_LABEL_LEN + 1), + ] { + assert!( + create_labelled_worker_suffix("wm-worker-abcde", rejected, 1).is_err(), + "{rejected} should be rejected" + ); + } + } + + /// Every part of a worker name comes from the environment, so the name can only be kept + /// within what `worker_ping.worker` holds by checking the assembled thing. + #[test] + fn worker_name_longer_than_the_ping_key_is_refused() { + let suffix = "abcde-a1b2c_slot"; + let room = MAX_WORKER_NAME_LEN - worker_name_with_suffix(false, "", suffix).len(); + let fits = checked_worker_name(false, &"g".repeat(room), suffix).unwrap(); + assert_eq!(fits.len(), MAX_WORKER_NAME_LEN); + assert!(checked_worker_name(false, &"g".repeat(room + 1), suffix).is_err()); + } + /// The guards are only safe because they are never stricter than the DB /// constraints they front. Narrowing `\w` to ASCII reads equivalent and /// compiles, but would start rejecting paths that already deploy today. diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index 0fad3248cc..a94160be33 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -263,6 +263,31 @@ lazy_static::lazy_static! { pub static ref NO_LOGS: bool = std::env::var("NO_LOGS").ok().is_some_and(|x| x == "1" || x == "true"); + /// Shut the worker process down once it has executed this many jobs, so a supervisor + /// (docker restart policy, kubernetes, systemd, ...) restarts it on a pristine + /// environment. Meant for deployments that cannot sandbox jobs with nsjail and rely on + /// the process/container lifetime to isolate one execution from the next. `0` (or unset) + /// disables it. Only the jobs the worker's own main loop ran count: ones handed off to a + /// dedicated worker or a flow runner are executed by another task and never counted. + /// Workers of one process share that environment, so with `NUM_WORKERS > 1` the first of + /// them to reach the limit takes the whole process down, cancelling whatever the others + /// still run in a container or a dedicated worker. Run one worker per process. + /// A value that does not parse is rejected at startup by [`validate_worker_lifecycle_env`] + /// rather than read as "disabled" here: a deployment that isolates executions this way + /// would otherwise keep running with no isolation at all. + pub static ref EXIT_AFTER_N_JOBS: Option = std::env::var("EXIT_AFTER_N_JOBS") + .ok() + .and_then(|x| x.parse::().ok()) + .filter(|x| *x > 0); + + /// Replaces the random part of the worker name, which is what makes a restarted process + /// reclaim its `worker_ping` row rather than register as a new worker. Two worker + /// processes must never share it: it is only needed when several of them run on one host + /// under the same worker group, since the name is otherwise derived from the hostname. + pub static ref WORKER_SUFFIX: Option = std::env::var("WORKER_SUFFIX") + .ok() + .filter(|x| !x.is_empty()); + pub static ref NATIVE_MODE: bool = std::env::var("NATIVE_MODE").ok().is_some_and(|x| x == "1" || x == "true"); pub static ref LIMIT_WINDOWS_TO_1CU: bool = std::env::var("LIMIT_WINDOWS_TO_1CU").ok().is_some_and(|x| x == "1" || x == "true"); @@ -450,6 +475,18 @@ lazy_static::lazy_static! { pub static ref ROOT_CACHE_NOMOUNT_DIR: String = format!("{}/cache_nomount/", *WINDMILL_DIR); } +/// Refuses to start on an `EXIT_AFTER_N_JOBS` that does not parse. Silently ignoring it +/// would leave a worker meant to recycle its environment running forever without doing so, +/// which is exactly the guarantee the deployment set it for. +pub fn validate_worker_lifecycle_env() -> anyhow::Result<()> { + match std::env::var("EXIT_AFTER_N_JOBS") { + Ok(v) if !v.is_empty() && v.parse::().is_err() => Err(anyhow::anyhow!( + "EXIT_AFTER_N_JOBS must be a positive integer (or 0 to disable), got '{v}'" + )), + _ => Ok(()), + } +} + /// Whether native mode is forced by the environment (NATIVE_MODE=true env var or WORKER_GROUP=native). /// This does NOT account for native_mode set in the DB worker group config — for that, read /// `WORKER_CONFIG.native_mode` which combines all sources. @@ -1858,6 +1895,16 @@ pub async fn fetch_raw_script_from_app_query( .map(|r| RawScript { content: r.code, lock: r.lock, meta: None, modules: None }) } +/// Returns the number of jobs the row already accounted for: non-zero when a worker of the +/// same name pinged before, i.e. when this process is a restart of an earlier one (see +/// [`crate::utils::resolve_worker_suffix`]) and its counter is meant to keep climbing. +/// +/// Everything else describing the process is overwritten on such a restart — a row left +/// reporting the version or the isolation mode of the process that died would, for +/// `wm_version`, hold the instance-wide `MIN_VERSION` back forever, and one still naming the +/// job that process was killed mid-way through skews the zombie/OOM diagnostics that read it. +/// `started_at` and `jobs_executed` are the only two columns carried over, being the +/// continuity itself. pub async fn insert_ping_query( worker_instance: &str, worker_name: &str, @@ -1872,10 +1919,11 @@ pub async fn insert_ping_query( job_isolation: Option, native_mode: bool, db: &DB, -) -> anyhow::Result<()> { - sqlx::query!( +) -> anyhow::Result { + let previous_jobs_executed = sqlx::query_scalar!( "INSERT INTO worker_ping (worker_instance, worker, ip, custom_tags, worker_group, dedicated_worker, dedicated_workers, wm_version, vcpus, memory, job_isolation, native_mode) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) ON CONFLICT (worker) - DO UPDATE set ip = EXCLUDED.ip, custom_tags = EXCLUDED.custom_tags, worker_group = EXCLUDED.worker_group, dedicated_workers = EXCLUDED.dedicated_workers, native_mode = EXCLUDED.native_mode", + DO UPDATE set ping_at = now(), worker_instance = EXCLUDED.worker_instance, ip = EXCLUDED.ip, custom_tags = EXCLUDED.custom_tags, worker_group = EXCLUDED.worker_group, dedicated_worker = EXCLUDED.dedicated_worker, dedicated_workers = EXCLUDED.dedicated_workers, wm_version = EXCLUDED.wm_version, vcpus = COALESCE(EXCLUDED.vcpus, worker_ping.vcpus), memory = COALESCE(EXCLUDED.memory, worker_ping.memory), job_isolation = EXCLUDED.job_isolation, native_mode = EXCLUDED.native_mode, current_job_id = NULL, current_job_workspace_id = NULL + RETURNING jobs_executed", worker_instance, worker_name, ip, @@ -1889,9 +1937,9 @@ pub async fn insert_ping_query( job_isolation.as_deref(), native_mode, ) - .execute(db) + .fetch_one(db) .await?; - Ok(()) + Ok(previous_jobs_executed) } pub async fn update_worker_ping_from_job_query( diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index c42116a15f..eb35b38d4c 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -44,8 +44,9 @@ use windmill_common::{ schema::{should_validate_schema, SchemaValidator}, utils::{create_directory_async, WarnAfterExt}, worker::{ - is_allowed_file_location, make_pull_query, write_file, Connection, HttpClient, MAX_TIMEOUT, - MIN_PERIODIC_SCRIPT_INTERVAL_SECONDS, ROOT_CACHE_DIR, ROOT_CACHE_NOMOUNT_DIR, WINDMILL_DIR, + is_allowed_file_location, make_pull_query, write_file, Connection, HttpClient, + EXIT_AFTER_N_JOBS, MAX_TIMEOUT, MIN_PERIODIC_SCRIPT_INTERVAL_SECONDS, ROOT_CACHE_DIR, + ROOT_CACHE_NOMOUNT_DIR, WINDMILL_DIR, }, worker_group_job_stats::JobStatsMap, KillpillSender, @@ -99,8 +100,8 @@ use windmill_common::{ use windmill_queue::{ append_logs, canceled_job_to_result, empty_result, get_same_worker_job, pull, push_init_job, push_periodic_bash_job, CanceledBy, JobAndPerms, JobCompleted, MiniPulledJob, - PrecomputedAgentInfo, PulledJob, SameWorkerPayload, HTTP_CLIENT, INIT_SCRIPT_TAG, - PERIODIC_SCRIPT_TAG, + PrecomputedAgentInfo, PulledJob, SameWorkerPayload, HTTP_CLIENT, INIT_SCRIPT_PATH_PREFIX, + INIT_SCRIPT_TAG, PERIODIC_SCRIPT_PATH_PREFIX, PERIODIC_SCRIPT_TAG, }; #[cfg(feature = "prometheus")] @@ -2031,6 +2032,9 @@ pub enum JobOutcome { /// or was suspended waiting for child jobs (WAC v2). /// All of these leave the span `Status` `Unset`. Completed, + /// A valid cached result was found for the job's args and path, so it was + /// answered without running anything. + CompletedFromCache, /// Job was attempted but its execution returned an error; the failure has /// been dispatched to the result processor. `description` holds the /// truncated error string for the outer span's `Status.message`. @@ -2044,7 +2048,14 @@ impl JobOutcome { /// True when the job completed successfully on this worker. Used by /// callers that previously matched on `Ok(true)`. pub fn is_success(&self) -> bool { - matches!(self, Self::Completed) + matches!(self, Self::Completed | Self::CompletedFromCache) + } + + /// Whether anything ran here. Only a cached result is a true no-run: + /// `AlreadyCompleted` is raised when the queue row disappears *while* the + /// child process is running, so that job did execute, and was interrupted. + fn ran_on_this_worker(&self) -> bool { + !matches!(self, Self::CompletedFromCache) } } @@ -2058,7 +2069,7 @@ impl JobOutcome { /// description that reflects the actual cause. pub(crate) fn record_job_span_status(result: &windmill_common::error::Result) { let description = match result { - Ok(JobOutcome::Completed) => return, + Ok(JobOutcome::Completed) | Ok(JobOutcome::CompletedFromCache) => return, Ok(JobOutcome::Failed { description }) => description.clone(), Ok(JobOutcome::AlreadyCompleted) => "job already completed by another worker".to_string(), Err(err) => truncate_description(&err.to_string()), @@ -2354,12 +2365,128 @@ pub async fn create_job_dir(worker_directory: &str, job_id: impl Display) -> Str job_dir_path } +/// Whether running this job may leave anything behind in the worker's environment, which is +/// what `EXIT_AFTER_N_JOBS` counts. Flow orchestration, noop/identity steps and the warmup +/// job run no user code at all, and the init/periodic scripts are the worker's own setup: +/// counting any of them would burn a restart cycle without a single user job having run. +/// +/// The internal scripts are recognized by the path this very worker queues them under, and +/// not by their tag alone: a tag is only routing configuration, so a worker pulling +/// `init_script` could otherwise be fed user jobs that never age its environment. +fn dirties_worker_env( + kind: JobKind, + tag: &str, + runnable_path: Option<&str>, + job_id: Uuid, + rejected_before_run: bool, + worker_name: &str, +) -> bool { + let own_script = |own_tag: &str, path_prefix: &str| { + tag == own_tag + && runnable_path.is_some_and(|p| p.starts_with(&format!("{path_prefix}{worker_name}"))) + }; + !rejected_before_run + && !kind.is_flow() + && !matches!( + kind, + JobKind::Noop + | JobKind::Identity + | JobKind::UnassignedScript + | JobKind::UnassignedFlow + | JobKind::UnassignedSinglestepFlow + ) + && !own_script(INIT_SCRIPT_TAG, INIT_SCRIPT_PATH_PREFIX) + && !own_script(PERIODIC_SCRIPT_TAG, PERIODIC_SCRIPT_PATH_PREFIX) + && job_id != Uuid::nil() +} + +#[cfg(test)] +mod exit_after_n_jobs_tests { + use super::*; + + const WK: &str = "wk-default-host-a1b2c"; + + /// A worker with `EXIT_AFTER_N_JOBS=1` that counted its own init script would shut down + /// before ever running a user job, restart, and loop on that forever. Jobs rejected + /// before the executor ran are the same waste of a restart. + #[test] + fn worker_own_jobs_do_not_count() { + for (kind, tag, path) in [ + ( + JobKind::Script, + INIT_SCRIPT_TAG, + Some(format!("{INIT_SCRIPT_PATH_PREFIX}{WK}")), + ), + ( + JobKind::Script, + PERIODIC_SCRIPT_TAG, + Some(format!("{PERIODIC_SCRIPT_PATH_PREFIX}{WK}_1700000000")), + ), + (JobKind::Flow, "flow", Some("u/admin/f".to_string())), + (JobKind::Noop, "other", None), + (JobKind::UnassignedScript, "bash", Some("u/admin/s".into())), + ] { + assert!( + !dirties_worker_env(kind, tag, path.as_deref(), Uuid::from_u128(1), false, WK), + "{kind:?}/{tag} should not count" + ); + } + // The dedicated worker warmup job, which runs no user code. + assert!(!dirties_worker_env( + JobKind::Script, + "bash", + Some("u/admin/s"), + Uuid::nil(), + false, + WK + )); + // Cancelled, or errored before the executor ran. + assert!(!dirties_worker_env( + JobKind::Script, + "bash", + Some("u/admin/s"), + Uuid::from_u128(1), + true, + WK + )); + } + + /// A job whose queue row vanished mid-execution ran here all the same, and one that + /// failed left behind whatever it had written before it did. + #[test] + fn only_a_cached_result_means_nothing_ran() { + assert!(!JobOutcome::CompletedFromCache.ran_on_this_worker()); + assert!(JobOutcome::AlreadyCompleted.ran_on_this_worker()); + assert!(JobOutcome::Completed.ran_on_this_worker()); + assert!(JobOutcome::Failed { description: "boom".to_string() }.ran_on_this_worker()); + } + + /// The internal tags are ordinary routing tags: a worker can be configured to pull them, + /// and the user jobs it then runs must still age its environment. + #[test] + fn user_jobs_count_whatever_tag_they_are_routed_with() { + for (tag, path) in [ + ("bash", Some("u/admin/s")), + (INIT_SCRIPT_TAG, Some("u/admin/s")), + (PERIODIC_SCRIPT_TAG, Some("u/admin/s")), + // Another worker's init script would not be this one's setup either. + (INIT_SCRIPT_TAG, Some("init_script_wk-default-host-99999")), + (PERIODIC_SCRIPT_TAG, None), + ] { + assert!( + dirties_worker_env(JobKind::Script, tag, path, Uuid::from_u128(1), false, WK), + "{tag}/{path:?} should count" + ); + } + } +} + pub async fn run_worker( conn: &Connection, hostname: &str, worker_name: String, i_worker: u64, - _num_workers: u32, + num_workers: u32, ip: &str, mut killpill_rx: tokio::sync::broadcast::Receiver<()>, killpill_tx: KillpillSender, @@ -2456,7 +2583,7 @@ pub async fn run_worker( let mut last_ping = Instant::now() - Duration::from_secs(NUM_SECS_PING + 1); - insert_ping(hostname, &worker_name, ip, conn) + let previous_jobs_executed = insert_ping(hostname, &worker_name, ip, conn) .await .expect("initial ping could be sent"); @@ -2656,7 +2783,12 @@ pub async fn run_worker( // let counter = meter.u64_counter("jobs.execution").build(); let mut occupancy_metrics = OccupancyMetrics::new(start_time); - let mut jobs_executed = 0; + // Seeded from the ping row so a worker that reclaimed its name keeps counting from where + // the previous process left off instead of resetting the total shown for that worker. + let mut jobs_executed = previous_jobs_executed; + // Only jobs run by this process count towards EXIT_AFTER_N_JOBS: the point is the age of + // the environment, not the lifetime total. + let mut jobs_executed_in_env: u64 = 0; let is_dedicated_worker: bool = { let config = WORKER_CONFIG.load(); @@ -2667,6 +2799,25 @@ pub async fn run_worker( .is_some_and(|dws| !dws.is_empty()) }; + if EXIT_AFTER_N_JOBS.is_some() && i_worker == 1 { + if num_workers > 1 { + tracing::warn!( + worker = %worker_name, hostname = %hostname, + "EXIT_AFTER_N_JOBS is set but this process runs {num_workers} workers: they share \ + the environment it recycles, so the first one to reach the limit shuts the others \ + down as well, cancelling any job they still run in a container or a dedicated \ + worker. Run a single worker per process instead." + ); + } + if is_dedicated_worker { + tracing::warn!( + worker = %worker_name, hostname = %hostname, + "EXIT_AFTER_N_JOBS does not apply to the jobs this worker hands to its dedicated \ + workers: those run outside its main loop and are never counted." + ); + } + } + #[cfg(feature = "benchmark")] let benchmark_jobs: i32 = std::env::var("BENCHMARK_JOBS") .unwrap_or("5000".to_string()) @@ -3259,6 +3410,14 @@ pub async fn run_worker( last_executed_job = None; jobs_executed += 1; + let mut dirties_env = dirties_worker_env( + job.kind, + &job.tag, + job.runnable_path.as_deref(), + job.id, + job.canceled_by.is_some() || job.pre_run_error.is_some(), + &worker_name, + ); tracing::debug!(target: VERBOSE_TARGET, worker = %worker_name, hostname = %hostname, "started handling of job {}", job.id); @@ -3580,6 +3739,12 @@ pub async fn run_worker( ) .await; + // A result served from the cache went through the loop without running + // anything here. + dirties_env &= job_result + .as_ref() + .map_or(true, JobOutcome::ran_on_this_worker); + match job_result { Ok(ref outcome) if !outcome.is_success() && is_init_script => { tracing::error!("init script job failed, exiting"); @@ -3685,6 +3850,31 @@ pub async fn run_worker( } } + if let Some(max_jobs) = *EXIT_AFTER_N_JOBS { + if dirties_env { + jobs_executed_in_env += 1; + } + if jobs_executed_in_env >= max_jobs { + // Killpill rather than `break`: the main loop still has to drain the + // same-worker jobs it owns (a same-worker flow runs to its end here, past + // the limit) and let the background processor persist the results of what + // it just ran, before the process goes away. `send` reports whether this is + // the shutdown that got scheduled. + if killpill_tx.send() { + tracing::info!( + worker = %worker_name, hostname = %hostname, + "executed {jobs_executed_in_env} job(s), EXIT_AFTER_N_JOBS={max_jobs} \ + reached: shutting the worker process down so it restarts on a fresh environment" + ); + } + // `jobs_executed` only reaches the ping row every NUM_SECS_PING, which this + // process is about to exit before: force one after every job from here on, + // drained ones included, so the row the restarted worker reclaims counts + // the jobs this one ran. + last_ping = Instant::now() - Duration::from_secs(NUM_SECS_PING + 1); + } + } + #[cfg(feature = "benchmark")] { if started { @@ -4351,7 +4541,7 @@ pub async fn handle_queued_job( } } - return Ok(JobOutcome::Completed); + return Ok(JobOutcome::CompletedFromCache); } }; } diff --git a/backend/windmill-worker/src/worker_utils.rs b/backend/windmill-worker/src/worker_utils.rs index a36ad4c348..6f47b2d863 100644 --- a/backend/windmill-worker/src/worker_utils.rs +++ b/backend/windmill-worker/src/worker_utils.rs @@ -164,12 +164,14 @@ async fn update_worker_ping_full_inner( Ok(()) } +/// Registers the worker in `worker_ping` and returns the number of jobs already attributed +/// to that worker name, which is 0 unless this process reclaims the row of an earlier one. pub async fn insert_ping( worker_instance: &str, worker_name: &str, ip: &str, db: &Connection, -) -> anyhow::Result<()> { +) -> anyhow::Result { let (tags, dw, dws, native_mode) = { let wc = (**WORKER_CONFIG.load()).clone(); ( @@ -200,7 +202,7 @@ pub async fn insert_ping( match db { Connection::Sql(db) => { - insert_ping_query( + return insert_ping_query( worker_instance, worker_name, WORKER_GROUP.as_str(), @@ -215,7 +217,7 @@ pub async fn insert_ping( native_mode, db, ) - .await?; + .await; } Connection::Http(client) => { client @@ -248,7 +250,9 @@ pub async fn insert_ping( .await?; } } - Ok(()) + // The agent ping endpoint answers with nothing, so an agent worker always starts its + // counter from zero. + Ok(0) } pub async fn update_worker_ping_from_job(