From 9f40cdca6258a9218770ebe1e7b89179e7d6667f Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Tue, 29 Sep 2026 15:24:39 +0200 Subject: [PATCH 1/3] feat: add pull_batch to claim jobs for many waiting workers at once (#11350) * feat: add pull_batch to claim jobs for many waiting workers at once Co-Authored-By: Claude Opus 5.5 (1M context) * fix: keep jobs admitted by earlier batch passes when a re-pull fails Co-Authored-By: Claude Opus 5.5 (1M context) * fix: keep suspended flows first on every batch re-pull pass Co-Authored-By: Claude Opus 5.5 (1M context) --------- Co-authored-by: Claude Opus 5.5 (1M context) --- backend/tests/batch_pull.rs | 278 +++++++++++++++++++++ backend/windmill-common/src/worker.rs | 159 +++++++++---- backend/windmill-queue/src/jobs.rs | 331 ++++++++++++++++++++------ 3 files changed, 652 insertions(+), 116 deletions(-) create mode 100644 backend/tests/batch_pull.rs diff --git a/backend/tests/batch_pull.rs b/backend/tests/batch_pull.rs new file mode 100644 index 0000000000..184b85dd06 --- /dev/null +++ b/backend/tests/batch_pull.rs @@ -0,0 +1,278 @@ +//! Pins what `pull_batch` promises a caller serving many waiting workers from one claim: every +//! claimed job is marked running under exactly one of those workers, concurrent batches never +//! claim the same job, the queues are walked in the same order a single pull walks them, and +//! the batch peek keeps the ordered index scan the single pull relies on. + +use std::collections::HashSet; + +use serde_json::Value; +use sqlx::{Pool, Postgres}; +use uuid::Uuid; +use windmill_common::worker::{make_batch_pull_query, PullQueue}; +use windmill_queue::{pull_batch, PulledJobResult}; + +const W_ID: &str = "test-workspace"; + +async fn queue_job(db: &Pool, tag: &str, suspend: Option) -> anyhow::Result { + let id = Uuid::new_v4(); + sqlx::query( + "INSERT INTO v2_job (id, workspace_id, created_by, created_at, permissioned_as, \ + permissioned_as_email, kind, tag, args, visible_to_owner) \ + VALUES ($1, $2, 'test-user', now(), 'u/test-user', 'test@windmill.dev', 'noop', $3, \ + '{}', true)", + ) + .bind(id) + .bind(W_ID) + .bind(tag) + .execute(db) + .await?; + // A suspended flow sits running with `suspend_until` set; `suspend = 0` makes it resumable. + sqlx::query( + "INSERT INTO v2_job_queue (id, workspace_id, scheduled_for, running, tag, suspend, \ + suspend_until) \ + VALUES ($1, $2, now() - interval '1 second', $3, $4, COALESCE($5, 0), \ + CASE WHEN $5::int IS NULL THEN NULL ELSE now() + interval '1 day' END)", + ) + .bind(id) + .bind(W_ID) + .bind(suspend.is_some()) + .bind(tag) + .bind(suspend) + .execute(db) + .await?; + sqlx::query("INSERT INTO v2_job_runtime (id) VALUES ($1)") + .bind(id) + .execute(db) + .await?; + Ok(id) +} + +fn names(prefix: &str, n: usize) -> Vec { + (0..n).map(|i| format!("{prefix}-{i}")).collect() +} + +fn job_id(res: &PulledJobResult) -> Uuid { + res.job + .as_ref() + .expect("an admitted result carries its job") + .id +} + +#[sqlx::test(fixtures("base"))] +async fn batch_pull_assigns_each_job_to_one_waiting_worker( + db: Pool, +) -> anyhow::Result<()> { + for _ in 0..30 { + queue_job(&db, "batch", None).await?; + } + let groups = vec![vec!["batch".to_string()]]; + + // Three batches race over one queue; a name listed twice still gets one job. + let mut a = names("a", 8); + a.push("a-0".to_string()); + let (b, c) = (names("b", 8), names("c", 8)); + let (ra, rb, rc) = tokio::join!( + pull_batch(&db, &groups, &a, false), + pull_batch(&db, &groups, &b, false), + pull_batch(&db, &groups, &c, false), + ); + let claimed: Vec<(String, PulledJobResult)> = [ra?, rb?, rc?].into_iter().flatten().collect(); + assert_eq!( + claimed.len(), + 24, + "30 queued jobs cover all 24 distinct workers" + ); + let workers: HashSet<&String> = claimed.iter().map(|(w, _)| w).collect(); + let jobs: HashSet = claimed.iter().map(|(_, r)| job_id(r)).collect(); + assert_eq!(workers.len(), 24, "no worker is handed two jobs"); + assert_eq!(jobs.len(), 24, "no job is handed to two workers"); + + for (worker, res) in &claimed { + let (running, row_worker): (bool, Option) = + sqlx::query_as("SELECT running, worker FROM v2_job_queue WHERE id = $1") + .bind(job_id(res)) + .fetch_one(&db) + .await?; + assert!(running); + assert_eq!( + row_worker.as_ref(), + Some(worker), + "job is running under its worker" + ); + } + + // The result is the wire type a remote caller hands to its worker. + let (_, first) = &claimed[0]; + let round_trip: PulledJobResult = serde_json::from_str(&serde_json::to_string(first)?)?; + assert_eq!(job_id(&round_trip), job_id(first)); + + // Six jobs remain for ten workers: four go without. + let rest = pull_batch(&db, &groups, &names("d", 10), false).await?; + assert_eq!(rest.len(), 6); + Ok(()) +} + +#[sqlx::test(fixtures("base"))] +async fn batch_pull_walks_suspended_then_priority_groups(db: Pool) -> anyhow::Result<()> { + let resumable = queue_job(&db, "low", Some(0)).await?; + let mut high = HashSet::new(); + for _ in 0..3 { + high.insert(queue_job(&db, "high", None).await?); + queue_job(&db, "low", None).await?; + } + let groups = vec![vec!["high".to_string()], vec!["low".to_string()]]; + + let claimed = pull_batch(&db, &groups, &names("w", 5), true).await?; + assert_eq!(claimed.len(), 5); + let suspended: Vec = claimed + .iter() + .filter(|(_, r)| r.suspended) + .map(|(_, r)| job_id(r)) + .collect(); + assert_eq!( + suspended, + vec![resumable], + "the resumable flow is claimed first" + ); + let ready: HashSet = claimed + .iter() + .filter(|(_, r)| !r.suspended) + .map(|(_, r)| job_id(r)) + .collect(); + assert!( + high.is_subset(&ready), + "the high group drains before the low one" + ); + assert_eq!(ready.len(), 4); + Ok(()) +} + +// A job over its concurrency limit must go back to the queue rather than sit claimed, and the +// worker it was claimed for is pulled for again in the same order: a resumable flow still wins +// over a ready job on the re-pull, as it does on every loop of a single pull. +#[cfg(feature = "enterprise")] +#[sqlx::test(fixtures("base"))] +async fn batch_pull_requeues_over_limit_job_and_serves_its_worker( + db: Pool, +) -> anyhow::Result<()> { + let limited = queue_job(&db, "batch", Some(0)).await?; + sqlx::query( + "UPDATE v2_job SET kind = 'script', runnable_path = 'f/test/limited', \ + concurrent_limit = 1, concurrency_time_window_s = 0 WHERE id = $1", + ) + .bind(limited) + .execute(&db) + .await?; + // Claimed first, so the bounce happens before the other jobs are reached. + sqlx::query("UPDATE v2_job_queue SET priority = 10 WHERE id = $1") + .bind(limited) + .execute(&db) + .await?; + sqlx::query("INSERT INTO concurrency_key (key, job_id) VALUES ('limited-key', $1)") + .bind(limited) + .execute(&db) + .await?; + sqlx::query( + "INSERT INTO concurrency_counter (concurrency_id, job_uuids) \ + VALUES ('limited-key', jsonb_build_object(gen_random_uuid()::text, '{}'::jsonb))", + ) + .execute(&db) + .await?; + let resumable = queue_job(&db, "batch", Some(0)).await?; + queue_job(&db, "batch", None).await?; + + let groups = vec![vec!["batch".to_string()]]; + let claimed = pull_batch(&db, &groups, &names("w", 1), true).await?; + assert_eq!(claimed.len(), 1); + assert_eq!(claimed[0].0, "w-0"); + assert_eq!( + job_id(&claimed[0].1), + resumable, + "the worker is served the other resumable flow" + ); + assert!(claimed[0].1.suspended); + + let (running, rescheduled): (bool, bool) = + sqlx::query_as("SELECT running, scheduled_for > now() FROM v2_job_queue WHERE id = $1") + .bind(limited) + .fetch_one(&db) + .await?; + assert!(!running, "the over-limit job is back in the queue"); + assert!(rescheduled, "the over-limit job waits for a free slot"); + Ok(()) +} + +// Admission failing for one claimed job must hand it back to the queue: it is already marked +// running under its worker, which will never be told about it. +#[sqlx::test(fixtures("base"))] +async fn batch_pull_requeues_job_whose_admission_fails(db: Pool) -> anyhow::Result<()> { + let broken = queue_job(&db, "batch", None).await?; + // A settings handle with no `runnable_settings` row makes the settings lookup fail. + sqlx::query("UPDATE v2_job_queue SET runnable_settings_handle = -42 WHERE id = $1") + .bind(broken) + .execute(&db) + .await?; + + let groups = vec![vec!["batch".to_string()]]; + let claimed = pull_batch(&db, &groups, &names("w", 1), false).await?; + assert!(claimed.is_empty()); + + let (running, started): (bool, bool) = + sqlx::query_as("SELECT running, started_at IS NOT NULL FROM v2_job_queue WHERE id = $1") + .bind(broken) + .fetch_one(&db) + .await?; + assert!(!running && !started, "the job is back in the queue, not stranded"); + Ok(()) +} + +/// Depth-first walk of an `EXPLAIN (FORMAT JSON)` plan tree. +fn nodes(plan: &Value, out: &mut Vec) { + out.push(plan.clone()); + for child in plan["Plans"].as_array().unwrap_or(&vec![]) { + nodes(child, out); + } +} + +// The batch LIMIT is a bind parameter, so a cached generic plan cannot see it and the planner +// falls back to guessing a fraction of the matching rows. On a deep backlog, a guess that tips +// it into a bitmap scan plus sort reads every queued job of the tag on each pull. +#[sqlx::test(fixtures("base"))] +async fn batch_pull_generic_plan_walks_the_queue_index(db: Pool) -> anyhow::Result<()> { + sqlx::query( + "INSERT INTO v2_job_queue (id, workspace_id, scheduled_for, running, tag) + SELECT gen_random_uuid(), 'test-workspace', now() - make_interval(secs => i % 1000), + false, CASE WHEN i % 10 = 0 THEN 'python3' ELSE 'deno' END + FROM generate_series(1, 200000) i", + ) + .execute(&db) + .await?; + sqlx::query("ANALYZE v2_job_queue").execute(&db).await?; + + let mut conn = db.acquire().await?; + sqlx::query("SET plan_cache_mode = force_generic_plan") + .execute(&mut *conn) + .await?; + let query = make_batch_pull_query(&["python3".to_string()], PullQueue::Ready); + sqlx::query(&format!("PREPARE batch_pull(text[]) AS {query}")) + .execute(&mut *conn) + .await?; + let explained: Value = + sqlx::query_scalar("EXPLAIN (FORMAT JSON) EXECUTE batch_pull(ARRAY['w-0', 'w-1'])") + .fetch_one(&mut *conn) + .await?; + + let mut all = vec![]; + nodes(&explained[0]["Plan"], &mut all); + let pretty = serde_json::to_string_pretty(&explained)?; + assert!( + all.iter() + .any(|n| n["Node Type"] == "Index Scan" && n["Index Name"] == "queue_sort_v2"), + "batch peek does not walk queue_sort_v2 in order:\n{pretty}" + ); + assert!( + !all.iter().any(|n| n["Node Type"] == "Sort"), + "batch peek sorts the backlog:\n{pretty}" + ); + Ok(()) +} diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index d480a89851..449fbd7606 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -774,6 +774,60 @@ impl From<&Pool> for Connection { } } +/// The columns every pull returns, read by name into `windmill_queue::PulledJob`. Expects `q` +/// (the claimed `v2_job_queue` rows), `j` (their `v2_job` rows), `f`, `p` and `pj` in scope. +const PULLED_JOB_COLUMNS: &str = "j.id, j.workspace_id, j.parent_job, j.created_by, q.started_at, q.scheduled_for, + j.runnable_id, j.runnable_path, j.args, q.canceled_by, + q.canceled_reason, j.kind, j.trigger, j.trigger_kind, j.permissioned_as, + f.flow_status, j.script_lang, + j.same_worker, j.pre_run_error, j.visible_to_owner, + j.tag, j.concurrent_limit, j.concurrency_time_window_s, j.flow_innermost_root_job, j.root_job, + j.timeout, j.flow_step_id, j.cache_ttl, q.cache_ignore_s3_path, q.runnable_settings_handle, j.priority, j.raw_code, j.raw_lock, j.raw_flow, + j.script_entrypoint_override, j.preprocessed, COALESCE(pj.runnable_path, j.args->>'_FLOW_PATH') as parent_runnable_path, + COALESCE(p.email, j.permissioned_as_email) as permissioned_as_email, p.username as permissioned_as_username, p.is_admin as permissioned_as_is_admin, + p.is_operator as permissioned_as_is_operator, p.groups as permissioned_as_groups, p.folders as permissioned_as_folders, p.end_user_email as permissioned_as_end_user_email"; + +fn sql_tag_list(tags: &[String]) -> String { + tags.iter() + .map(|x| format!("'{}'", x.replace('\'', "''"))) + .join(", ") +} + +/// Selects the next runnable jobs for `tags`, highest priority first. `exclude_overloaded` names +/// the text[] bind parameter holding workspaces that workspace fairness currently caps. +fn ready_peek(tags: &[String], exclude_overloaded: Option<&str>, limit: &str) -> String { + let fairness = exclude_overloaded + .map(|param| format!("\n AND workspace_id <> ALL({param}::text[])")) + .unwrap_or_default(); + format!( + "SELECT id + FROM v2_job_queue + WHERE running = false AND tag IN ({}) AND scheduled_for <= now(){fairness} + ORDER BY priority DESC NULLS LAST, scheduled_for + FOR UPDATE SKIP LOCKED + LIMIT {limit}", + sql_tag_list(tags) + ) +} + +// The `CASE` is `suspend <= 0 OR suspend_until <= now()` written as one indexable +// expression, equivalent only under the `suspend_until IS NOT NULL` guard. It must stay in +// sync with `queue_suspended_v2` (migration 20260826202939): if it no longer matches, the +// test silently reverts to a heap filter over every suspended row on every worker poll. +fn suspended_peek(tags: &[String], limit: &str) -> String { + format!( + "SELECT id + FROM v2_job_queue + WHERE suspend_until IS NOT NULL + AND (CASE WHEN suspend <= 0 THEN '-infinity'::timestamptz ELSE suspend_until END) <= now() + AND tag IN ({}) + ORDER BY priority DESC NULLS LAST, created_at + FOR UPDATE SKIP LOCKED + LIMIT {limit}", + sql_tag_list(tags) + ) +} + fn format_pull_query(peek: String) -> String { let r = format!( "WITH peek AS ( @@ -803,16 +857,7 @@ fn format_pull_query(peek: String) -> String { raw_flow, script_entrypoint_override, preprocessed FROM v2_job WHERE id = (SELECT id FROM peek) - ) SELECT j.id, j.workspace_id, j.parent_job, j.created_by, started_at, scheduled_for, - j.runnable_id, j.runnable_path, j.args, canceled_by, - canceled_reason, j.kind, j.trigger, j.trigger_kind, j.permissioned_as, - flow_status, j.script_lang, - j.same_worker, j.pre_run_error, j.visible_to_owner, - j.tag, j.concurrent_limit, j.concurrency_time_window_s, j.flow_innermost_root_job, j.root_job, - j.timeout, j.flow_step_id, j.cache_ttl, q.cache_ignore_s3_path, q.runnable_settings_handle, j.priority, j.raw_code, j.raw_lock, j.raw_flow, - j.script_entrypoint_override, j.preprocessed, COALESCE(pj.runnable_path, j.args->>'_FLOW_PATH') as parent_runnable_path, - COALESCE(p.email, j.permissioned_as_email) as permissioned_as_email, p.username as permissioned_as_username, p.is_admin as permissioned_as_is_admin, - p.is_operator as permissioned_as_is_operator, p.groups as permissioned_as_groups, p.folders as permissioned_as_folders, p.end_user_email as permissioned_as_end_user_email + ) SELECT {PULLED_JOB_COLUMNS} FROM q, j LEFT JOIN v2_job_status f USING (id) LEFT JOIN job_perms p ON p.job_id = j.id @@ -824,22 +869,63 @@ fn format_pull_query(peek: String) -> String { r } -// The `CASE` is `suspend <= 0 OR suspend_until <= now()` written as one indexable -// expression, equivalent only under the `suspend_until IS NOT NULL` guard. It must stay in -// sync with `queue_suspended_v2` (migration 20260826202939): if it no longer matches, the -// test silently reverts to a heap filter over every suspended row on every worker poll. +/// Claims up to `cardinality($1)` jobs in one statement and assigns the i-th claimed job to the +/// worker named `$1[i]`, so each row is marked running under the worker that will run it. +/// Returns the `PULLED_JOB_COLUMNS` plus `assigned_worker`. +fn format_batch_pull_query(peek: String) -> String { + format!( + "WITH peek AS MATERIALIZED ( + {peek} + ), assign AS ( + SELECT id, ($1::text[])[row_number() OVER ()] AS assigned_worker FROM peek + ), q AS ( + UPDATE v2_job_queue SET + running = true, + started_at = coalesce(started_at, now()), + suspend_until = null, + worker = assign.assigned_worker + FROM assign + WHERE v2_job_queue.id = assign.id + RETURNING + v2_job_queue.id, started_at, scheduled_for, + canceled_by, canceled_reason, worker, cache_ignore_s3_path, runnable_settings_handle + ), r AS ( + UPDATE v2_job_runtime SET + ping = now() + WHERE id = ANY(ARRAY(SELECT id FROM q)) + ) SELECT {PULLED_JOB_COLUMNS}, q.worker AS assigned_worker + FROM q + JOIN v2_job j ON j.id = q.id + LEFT JOIN v2_job_status f ON f.id = q.id + LEFT JOIN job_perms p ON p.job_id = q.id + LEFT JOIN v2_job pj ON j.parent_job = pj.id" + ) +} + +/// Which part of the queue a pull claims from. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PullQueue { + /// Jobs ready to start. + Ready, + /// Jobs ready to start, excluding the workspaces bound as `$2::text[]`. + ReadyFairness, + /// Suspended flows whose resume condition is met. + Suspended, +} + +/// Batch variant of the pull queries: binds the claiming workers' names as `$1::text[]`, which +/// must be distinct, and claims at most one job per name. +pub fn make_batch_pull_query(tags: &[String], queue: PullQueue) -> String { + const LIMIT: &str = "cardinality($1::text[])"; + format_batch_pull_query(match queue { + PullQueue::Ready => ready_peek(tags, None, LIMIT), + PullQueue::ReadyFairness => ready_peek(tags, Some("$2"), LIMIT), + PullQueue::Suspended => suspended_peek(tags, LIMIT), + }) +} + pub fn make_suspended_pull_query(tags: &[String]) -> String { - format_pull_query(format!( - "SELECT id - FROM v2_job_queue - WHERE suspend_until IS NOT NULL - AND (CASE WHEN suspend <= 0 THEN '-infinity'::timestamptz ELSE suspend_until END) <= now() - AND tag IN ({}) - ORDER BY priority DESC NULLS LAST, created_at - FOR UPDATE SKIP LOCKED - LIMIT 1", - tags.iter().map(|x| format!("'{x}'")).join(", ") - )) + format_pull_query(suspended_peek(tags, "1")) } // pub async fn make_suspended pub async fn store_suspended_pull_query(wc: &WorkerConfig) { @@ -852,16 +938,7 @@ pub async fn store_suspended_pull_query(wc: &WorkerConfig) { } pub fn make_pull_query(tags: &[String]) -> String { - let query = format_pull_query(format!( - "SELECT id - FROM v2_job_queue - WHERE running = false AND tag IN ({}) AND scheduled_for <= now() - ORDER BY priority DESC NULLS LAST, scheduled_for - FOR UPDATE SKIP LOCKED - LIMIT 1", - tags.iter().map(|x| format!("'{x}'")).join(", ") - )); - query + format_pull_query(ready_peek(tags, None, "1")) } // Variant of `make_pull_query` that additionally excludes jobs whose workspace_id is in the @@ -872,17 +949,7 @@ pub fn make_pull_query(tags: &[String]) -> String { // `pub(crate)` because only `store_pull_query` consumes it; the resulting query string is what // crosses crate boundaries via `WORKER_PULL_QUERIES_FAIRNESS`. pub(crate) fn make_pull_query_fairness(tags: &[String]) -> String { - let query = format_pull_query(format!( - "SELECT id - FROM v2_job_queue - WHERE running = false AND tag IN ({}) AND scheduled_for <= now() - AND workspace_id <> ALL($2::text[]) - ORDER BY priority DESC NULLS LAST, scheduled_for - FOR UPDATE SKIP LOCKED - LIMIT 1", - tags.iter().map(|x| format!("'{x}'")).join(", ") - )); - query + format_pull_query(ready_peek(tags, Some("$2"), "1")) } pub async fn store_pull_query(wc: &WorkerConfig) { diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 7b552fb02c..8f65638790 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -78,9 +78,9 @@ use windmill_common::{ users::{SUPERADMIN_NOTIFICATION_EMAIL, SUPERADMIN_SECRET_EMAIL}, utils::{not_found_if_none, report_critical_error, StripPath, WarnAfterExt}, worker::{ - to_raw_value, CLOUD_HOSTED, DISABLE_FLOW_SCRIPT, NO_LOGS, PREVIEW_TAGS_OVERRIDE, - WORKER_PULL_QUERIES, WORKER_PULL_QUERIES_FAIRNESS, WORKER_SUSPENDED_PULL_QUERY, - WORKSPACE_FAIRNESS_OVERLOADED, + make_batch_pull_query, to_raw_value, PullQueue, CLOUD_HOSTED, DISABLE_FLOW_SCRIPT, NO_LOGS, + PREVIEW_TAGS_OVERRIDE, WORKER_PULL_QUERIES, WORKER_PULL_QUERIES_FAIRNESS, + WORKER_SUSPENDED_PULL_QUERY, WORKSPACE_FAIRNESS_OVERLOADED, }, DB, METRICS_ENABLED, }; @@ -3570,7 +3570,7 @@ impl MiniPulledJob { } } -#[derive(sqlx::FromRow, Debug, Clone)] +#[derive(sqlx::FromRow, Debug, Clone, Serialize, Deserialize)] pub struct PulledJob { #[sqlx(flatten)] pub job: MiniPulledJob, @@ -3837,7 +3837,7 @@ pub async fn get_queued_job_v2<'c>( Ok(job) } -#[derive(Debug)] +#[derive(Debug, Serialize, Deserialize)] pub struct PulledJobResult { pub job: Option, pub suspended: bool, @@ -4228,6 +4228,13 @@ async fn clone_runnable(j: &mut PulledJob, db: &DB) -> error::Result<()> { Ok(()) } +// Cap to bound DB work per pull cycle. Each iteration on an over-limit job runs +// the full apply_concurrency_limit query stack (~7 queries), so without a tight +// cap a single pull() can issue thousands of queries when the queue is full of +// jobs sharing one over-limit concurrency key. Capping at 10 lets workers skip +// a few stale jobs in healthy conditions while preventing storm amplification. +const PULL_LOOP_LIMIT: i32 = 10; + // TODO: Factorize /// Pull the job from queue pub async fn pull( @@ -4242,14 +4249,9 @@ pub async fn pull( let mut pull_loop_count = 0; loop { pull_loop_count += 1; - // Cap to bound DB work per pull cycle. Each iteration on an over-limit job runs - // the full apply_concurrency_limit query stack (~7 queries), so without a tight - // cap a single pull() can issue thousands of queries when the queue is full of - // jobs sharing one over-limit concurrency key. Capping at 10 lets workers skip - // a few stale jobs in healthy conditions while preventing storm amplification. - if pull_loop_count > 10 { + if pull_loop_count > PULL_LOOP_LIMIT { tracing::warn!( - "Pull job loop count exceeded 10, backing off (likely concurrency re-queue storm)" + "Pull job loop count exceeded {PULL_LOOP_LIMIT}, backing off (likely concurrency re-queue storm)" ); return Ok(PulledJobResult { job: None, @@ -4389,68 +4391,257 @@ pub async fn pull( }); }; - let concurrency_settings = windmill_common::runnable_settings::prefetch_cached_from_handle( - job.runnable_settings_handle, - db, - ) - .await? - .1 - .maybe_fallback(None, job.concurrent_limit, job.concurrency_time_window_s); - - let has_concurent_limit = - has_active_concurrency_limit(concurrency_settings.concurrent_limit); - - #[cfg(not(feature = "enterprise"))] - if has_concurent_limit && !job.is_dependency() { - tracing::error!("Concurrent limits are an EE feature only, ignoring constraints") - } - - #[cfg(not(feature = "enterprise"))] - let has_concurent_limit = job.is_dependency() - && has_active_concurrency_limit(job.concurrent_limit) - && cfg!(feature = "private") - && !*WMDEBUG_NO_DEBOUNCING; - // if we don't have private flag, we don't have concurrency limit - - // concurrency check. If more than X jobs for this path are already running, we re-queue and pull another job from the queue - let pulled_job = job; - if pulled_job.runnable_path.is_none() - || !has_concurent_limit - || pulled_job.canceled_by.is_some() - { - #[cfg(feature = "prometheus")] - if METRICS_ENABLED.load(std::sync::atomic::Ordering::Relaxed) { - QUEUE_PULL_COUNT.inc(); - } - otel_incr_queue_pull_count(); - return Ok(PulledJobResult { - job: Some(pulled_job), - suspended, - missing_concurrency_key: false, - error_while_preprocessing: None, - }); - } - - #[cfg(feature = "private")] - if cfg!(feature = "enterprise") || (pulled_job.is_dependency() && !*WMDEBUG_NO_DEBOUNCING) { - if let Some(pulled_job_res) = timeout( - Duration::from_secs(15), - crate::jobs_ee::apply_concurrency_limit( - db, - pull_loop_count, - suspended, - pulled_job, - &concurrency_settings, - ), - ) - .await?? - { - return Ok(pulled_job_res); - } + if let Some(pulled_job_res) = admit_pulled_job(db, job, suspended, pull_loop_count).await? { + return Ok(pulled_job_res); } } } +/// Applies a just-claimed job's concurrency limit, if it has one. `None` means the job was over +/// its limit and has been re-queued for later, so the claiming worker should pull again. +async fn admit_pulled_job( + db: &Pool, + job: PulledJob, + suspended: bool, + pull_loop_count: i32, +) -> windmill_common::error::Result> { + let concurrency_settings = windmill_common::runnable_settings::prefetch_cached_from_handle( + job.runnable_settings_handle, + db, + ) + .await? + .1 + .maybe_fallback(None, job.concurrent_limit, job.concurrency_time_window_s); + + let has_concurent_limit = has_active_concurrency_limit(concurrency_settings.concurrent_limit); + + #[cfg(not(feature = "enterprise"))] + if has_concurent_limit && !job.is_dependency() { + tracing::error!("Concurrent limits are an EE feature only, ignoring constraints") + } + + #[cfg(not(feature = "enterprise"))] + let has_concurent_limit = job.is_dependency() + && has_active_concurrency_limit(job.concurrent_limit) + && cfg!(feature = "private") + && !*WMDEBUG_NO_DEBOUNCING; + // if we don't have private flag, we don't have concurrency limit + + // concurrency check. If more than X jobs for this path are already running, we re-queue and pull another job from the queue + let pulled_job = job; + if pulled_job.runnable_path.is_none() + || !has_concurent_limit + || pulled_job.canceled_by.is_some() + { + #[cfg(feature = "prometheus")] + if METRICS_ENABLED.load(std::sync::atomic::Ordering::Relaxed) { + QUEUE_PULL_COUNT.inc(); + } + otel_incr_queue_pull_count(); + return Ok(Some(PulledJobResult { + job: Some(pulled_job), + suspended, + missing_concurrency_key: false, + error_while_preprocessing: None, + })); + } + + #[cfg(feature = "private")] + if cfg!(feature = "enterprise") || (pulled_job.is_dependency() && !*WMDEBUG_NO_DEBOUNCING) { + return Ok(timeout( + Duration::from_secs(15), + crate::jobs_ee::apply_concurrency_limit( + db, + pull_loop_count, + suspended, + pulled_job, + &concurrency_settings, + ), + ) + .await??); + } + #[cfg(not(feature = "private"))] + let _ = pull_loop_count; + Ok(None) +} + +/// Claims jobs for several waiting workers in one statement per queue walked, instead of one +/// pull per worker. Each distinct name in `worker_names` gets at most one job, marked running +/// under that name; a worker missing from the result found nothing to run. +/// +/// Walks the queues the way [`pull`] does for one worker: when `suspend_first`, suspended flows +/// ready to resume on any of the tags; then `tag_groups` from highest priority down, with the +/// workspace-fairness share of the workers first restricted to uncapped workspaces. A job over +/// its concurrency limit is re-queued for later and its worker is pulled for again. +/// +/// No authorization happens here: jobs from every workspace matching the tags are claimed and +/// returned with their args and permissions. Callers must authenticate each worker named and +/// pass only tags that worker is allowed to serve. +pub async fn pull_batch( + db: &Pool, + tag_groups: &[Vec], + worker_names: &[String], + suspend_first: bool, +) -> windmill_common::error::Result> { + let mut waiting: Vec = worker_names.iter().unique().cloned().collect(); + let mut admitted = vec![]; + for pull_loop_count in 1..=PULL_LOOP_LIMIT { + if waiting.is_empty() { + break; + } + let claimed = match claim_batch(db, tag_groups, &waiting, suspend_first).await + { + Ok(claimed) => claimed, + Err(e) if admitted.is_empty() => return Err(e), + // Jobs admitted by an earlier pass are running under their workers: return them. + Err(e) => { + tracing::error!( + "batch pull re-pull failed, returning the jobs admitted so far: {e:#}" + ); + break; + } + }; + let mut requeued = false; + for (worker_name, job, suspended) in claimed { + waiting.retain(|w| w != &worker_name); + let job_id = job.id; + match admit_pulled_job(db, job, suspended, pull_loop_count).await { + Ok(Some(res)) => admitted.push((worker_name, res)), + Ok(None) => { + waiting.push(worker_name); + requeued = true; + } + // The other claimed jobs are already running under their workers, so an error + // here must not abort the batch: hand this one back to the queue instead. + Err(e) => { + tracing::error!( + "error admitting batch-pulled job {job_id}, re-queuing it: {e:#}" + ); + if let Err(e) = sqlx::query( + "WITH ping AS ( + UPDATE v2_job_runtime SET ping = null WHERE id = $1 + ) + UPDATE v2_job_queue SET running = false, started_at = null + WHERE id = $1 AND worker = $2 AND running = true", + ) + .bind(job_id) + .bind(&worker_name) + .execute(db) + .await + { + tracing::error!("could not re-queue batch-pulled job {job_id}: {e:#}"); + } + } + } + } + if !requeued { + break; + } + if pull_loop_count == PULL_LOOP_LIMIT { + tracing::warn!( + "Batch pull loop count reached {PULL_LOOP_LIMIT}, backing off (likely concurrency re-queue storm)" + ); + } + } + Ok(admitted) +} + +/// One pass over the queues for `worker_names`, returning `(worker, job, suspended)` for each +/// job claimed. A failing statement ends the pass but keeps what was already claimed, since +/// those jobs are marked running and would otherwise be stranded until the zombie monitor. +async fn claim_batch( + db: &Pool, + tag_groups: &[Vec], + worker_names: &[String], + suspend_first: bool, +) -> windmill_common::error::Result> { + let mut claimed = vec![]; + let mut waiting = worker_names.to_vec(); + let pass: windmill_common::error::Result<()> = async { + let tags: Vec = tag_groups.iter().flatten().unique().cloned().collect(); + if suspend_first && !tags.is_empty() { + for (worker, job) in + claim_from_queue(db, &tags, PullQueue::Suspended, &[], &waiting).await? + { + waiting.retain(|w| w != &worker); + claimed.push((worker, job, true)); + } + } + + crate::workspace_fairness::maybe_refresh_overloaded(db); + let overloaded = WORKSPACE_FAIRNESS_OVERLOADED.load_full(); + if !overloaded.is_empty() { + // Same admission draw `pull` makes once per worker, made here once per waiting worker. + let mut fair: Vec = waiting + .iter() + .filter(|_| !crate::workspace_fairness::should_admit_capped()) + .cloned() + .collect(); + for tags in tag_groups.iter().filter(|t| !t.is_empty()) { + if fair.is_empty() { + break; + } + for (worker, job) in claim_from_queue( + db, + tags, + PullQueue::ReadyFairness, + overloaded.as_slice(), + &fair, + ) + .await? + { + fair.retain(|w| w != &worker); + waiting.retain(|w| w != &worker); + claimed.push((worker, job, false)); + } + } + } + + for tags in tag_groups.iter().filter(|t| !t.is_empty()) { + if waiting.is_empty() { + break; + } + for (worker, job) in claim_from_queue(db, tags, PullQueue::Ready, &[], &waiting).await? + { + waiting.retain(|w| w != &worker); + claimed.push((worker, job, false)); + } + } + Ok(()) + } + .await; + match pass { + Err(e) if claimed.is_empty() => Err(e), + Err(e) => { + tracing::error!( + "batch pull pass failed after claiming {} jobs: {e:#}", + claimed.len() + ); + Ok(claimed) + } + Ok(()) => Ok(claimed), + } +} + +async fn claim_from_queue( + db: &Pool, + tags: &[String], + queue: PullQueue, + overloaded: &[String], + worker_names: &[String], +) -> windmill_common::error::Result> { + use sqlx::{FromRow, Row}; + let query = make_batch_pull_query(tags, queue); + let mut query = sqlx::query(&query).bind(worker_names); + if queue == PullQueue::ReadyFairness { + query = query.bind(overloaded); + } + let rows = timeout(Duration::from_secs(15), query.fetch_all(db)).await??; + rows.iter() + .map(|row| Ok((row.try_get("assigned_worker")?, PulledJob::from_row(row)?))) + .collect() +} + async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>( db: &Pool, suspend_first: bool, From 3936191f46367653179a0820bb8281e815a9dedb Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Tue, 29 Sep 2026 15:25:35 +0200 Subject: [PATCH 2/3] =?UTF-8?q?Revert=20"feat:=20add=20pull=5Fbatch=20to?= =?UTF-8?q?=20claim=20jobs=20for=20many=20waiting=20workers=20at=20once=20?= =?UTF-8?q?(=E2=80=A6"=20(#11412)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 9f40cdca6258a9218770ebe1e7b89179e7d6667f. --- backend/tests/batch_pull.rs | 278 ----------------------- backend/windmill-common/src/worker.rs | 159 ++++--------- backend/windmill-queue/src/jobs.rs | 311 +++++--------------------- 3 files changed, 106 insertions(+), 642 deletions(-) delete mode 100644 backend/tests/batch_pull.rs diff --git a/backend/tests/batch_pull.rs b/backend/tests/batch_pull.rs deleted file mode 100644 index 184b85dd06..0000000000 --- a/backend/tests/batch_pull.rs +++ /dev/null @@ -1,278 +0,0 @@ -//! Pins what `pull_batch` promises a caller serving many waiting workers from one claim: every -//! claimed job is marked running under exactly one of those workers, concurrent batches never -//! claim the same job, the queues are walked in the same order a single pull walks them, and -//! the batch peek keeps the ordered index scan the single pull relies on. - -use std::collections::HashSet; - -use serde_json::Value; -use sqlx::{Pool, Postgres}; -use uuid::Uuid; -use windmill_common::worker::{make_batch_pull_query, PullQueue}; -use windmill_queue::{pull_batch, PulledJobResult}; - -const W_ID: &str = "test-workspace"; - -async fn queue_job(db: &Pool, tag: &str, suspend: Option) -> anyhow::Result { - let id = Uuid::new_v4(); - sqlx::query( - "INSERT INTO v2_job (id, workspace_id, created_by, created_at, permissioned_as, \ - permissioned_as_email, kind, tag, args, visible_to_owner) \ - VALUES ($1, $2, 'test-user', now(), 'u/test-user', 'test@windmill.dev', 'noop', $3, \ - '{}', true)", - ) - .bind(id) - .bind(W_ID) - .bind(tag) - .execute(db) - .await?; - // A suspended flow sits running with `suspend_until` set; `suspend = 0` makes it resumable. - sqlx::query( - "INSERT INTO v2_job_queue (id, workspace_id, scheduled_for, running, tag, suspend, \ - suspend_until) \ - VALUES ($1, $2, now() - interval '1 second', $3, $4, COALESCE($5, 0), \ - CASE WHEN $5::int IS NULL THEN NULL ELSE now() + interval '1 day' END)", - ) - .bind(id) - .bind(W_ID) - .bind(suspend.is_some()) - .bind(tag) - .bind(suspend) - .execute(db) - .await?; - sqlx::query("INSERT INTO v2_job_runtime (id) VALUES ($1)") - .bind(id) - .execute(db) - .await?; - Ok(id) -} - -fn names(prefix: &str, n: usize) -> Vec { - (0..n).map(|i| format!("{prefix}-{i}")).collect() -} - -fn job_id(res: &PulledJobResult) -> Uuid { - res.job - .as_ref() - .expect("an admitted result carries its job") - .id -} - -#[sqlx::test(fixtures("base"))] -async fn batch_pull_assigns_each_job_to_one_waiting_worker( - db: Pool, -) -> anyhow::Result<()> { - for _ in 0..30 { - queue_job(&db, "batch", None).await?; - } - let groups = vec![vec!["batch".to_string()]]; - - // Three batches race over one queue; a name listed twice still gets one job. - let mut a = names("a", 8); - a.push("a-0".to_string()); - let (b, c) = (names("b", 8), names("c", 8)); - let (ra, rb, rc) = tokio::join!( - pull_batch(&db, &groups, &a, false), - pull_batch(&db, &groups, &b, false), - pull_batch(&db, &groups, &c, false), - ); - let claimed: Vec<(String, PulledJobResult)> = [ra?, rb?, rc?].into_iter().flatten().collect(); - assert_eq!( - claimed.len(), - 24, - "30 queued jobs cover all 24 distinct workers" - ); - let workers: HashSet<&String> = claimed.iter().map(|(w, _)| w).collect(); - let jobs: HashSet = claimed.iter().map(|(_, r)| job_id(r)).collect(); - assert_eq!(workers.len(), 24, "no worker is handed two jobs"); - assert_eq!(jobs.len(), 24, "no job is handed to two workers"); - - for (worker, res) in &claimed { - let (running, row_worker): (bool, Option) = - sqlx::query_as("SELECT running, worker FROM v2_job_queue WHERE id = $1") - .bind(job_id(res)) - .fetch_one(&db) - .await?; - assert!(running); - assert_eq!( - row_worker.as_ref(), - Some(worker), - "job is running under its worker" - ); - } - - // The result is the wire type a remote caller hands to its worker. - let (_, first) = &claimed[0]; - let round_trip: PulledJobResult = serde_json::from_str(&serde_json::to_string(first)?)?; - assert_eq!(job_id(&round_trip), job_id(first)); - - // Six jobs remain for ten workers: four go without. - let rest = pull_batch(&db, &groups, &names("d", 10), false).await?; - assert_eq!(rest.len(), 6); - Ok(()) -} - -#[sqlx::test(fixtures("base"))] -async fn batch_pull_walks_suspended_then_priority_groups(db: Pool) -> anyhow::Result<()> { - let resumable = queue_job(&db, "low", Some(0)).await?; - let mut high = HashSet::new(); - for _ in 0..3 { - high.insert(queue_job(&db, "high", None).await?); - queue_job(&db, "low", None).await?; - } - let groups = vec![vec!["high".to_string()], vec!["low".to_string()]]; - - let claimed = pull_batch(&db, &groups, &names("w", 5), true).await?; - assert_eq!(claimed.len(), 5); - let suspended: Vec = claimed - .iter() - .filter(|(_, r)| r.suspended) - .map(|(_, r)| job_id(r)) - .collect(); - assert_eq!( - suspended, - vec![resumable], - "the resumable flow is claimed first" - ); - let ready: HashSet = claimed - .iter() - .filter(|(_, r)| !r.suspended) - .map(|(_, r)| job_id(r)) - .collect(); - assert!( - high.is_subset(&ready), - "the high group drains before the low one" - ); - assert_eq!(ready.len(), 4); - Ok(()) -} - -// A job over its concurrency limit must go back to the queue rather than sit claimed, and the -// worker it was claimed for is pulled for again in the same order: a resumable flow still wins -// over a ready job on the re-pull, as it does on every loop of a single pull. -#[cfg(feature = "enterprise")] -#[sqlx::test(fixtures("base"))] -async fn batch_pull_requeues_over_limit_job_and_serves_its_worker( - db: Pool, -) -> anyhow::Result<()> { - let limited = queue_job(&db, "batch", Some(0)).await?; - sqlx::query( - "UPDATE v2_job SET kind = 'script', runnable_path = 'f/test/limited', \ - concurrent_limit = 1, concurrency_time_window_s = 0 WHERE id = $1", - ) - .bind(limited) - .execute(&db) - .await?; - // Claimed first, so the bounce happens before the other jobs are reached. - sqlx::query("UPDATE v2_job_queue SET priority = 10 WHERE id = $1") - .bind(limited) - .execute(&db) - .await?; - sqlx::query("INSERT INTO concurrency_key (key, job_id) VALUES ('limited-key', $1)") - .bind(limited) - .execute(&db) - .await?; - sqlx::query( - "INSERT INTO concurrency_counter (concurrency_id, job_uuids) \ - VALUES ('limited-key', jsonb_build_object(gen_random_uuid()::text, '{}'::jsonb))", - ) - .execute(&db) - .await?; - let resumable = queue_job(&db, "batch", Some(0)).await?; - queue_job(&db, "batch", None).await?; - - let groups = vec![vec!["batch".to_string()]]; - let claimed = pull_batch(&db, &groups, &names("w", 1), true).await?; - assert_eq!(claimed.len(), 1); - assert_eq!(claimed[0].0, "w-0"); - assert_eq!( - job_id(&claimed[0].1), - resumable, - "the worker is served the other resumable flow" - ); - assert!(claimed[0].1.suspended); - - let (running, rescheduled): (bool, bool) = - sqlx::query_as("SELECT running, scheduled_for > now() FROM v2_job_queue WHERE id = $1") - .bind(limited) - .fetch_one(&db) - .await?; - assert!(!running, "the over-limit job is back in the queue"); - assert!(rescheduled, "the over-limit job waits for a free slot"); - Ok(()) -} - -// Admission failing for one claimed job must hand it back to the queue: it is already marked -// running under its worker, which will never be told about it. -#[sqlx::test(fixtures("base"))] -async fn batch_pull_requeues_job_whose_admission_fails(db: Pool) -> anyhow::Result<()> { - let broken = queue_job(&db, "batch", None).await?; - // A settings handle with no `runnable_settings` row makes the settings lookup fail. - sqlx::query("UPDATE v2_job_queue SET runnable_settings_handle = -42 WHERE id = $1") - .bind(broken) - .execute(&db) - .await?; - - let groups = vec![vec!["batch".to_string()]]; - let claimed = pull_batch(&db, &groups, &names("w", 1), false).await?; - assert!(claimed.is_empty()); - - let (running, started): (bool, bool) = - sqlx::query_as("SELECT running, started_at IS NOT NULL FROM v2_job_queue WHERE id = $1") - .bind(broken) - .fetch_one(&db) - .await?; - assert!(!running && !started, "the job is back in the queue, not stranded"); - Ok(()) -} - -/// Depth-first walk of an `EXPLAIN (FORMAT JSON)` plan tree. -fn nodes(plan: &Value, out: &mut Vec) { - out.push(plan.clone()); - for child in plan["Plans"].as_array().unwrap_or(&vec![]) { - nodes(child, out); - } -} - -// The batch LIMIT is a bind parameter, so a cached generic plan cannot see it and the planner -// falls back to guessing a fraction of the matching rows. On a deep backlog, a guess that tips -// it into a bitmap scan plus sort reads every queued job of the tag on each pull. -#[sqlx::test(fixtures("base"))] -async fn batch_pull_generic_plan_walks_the_queue_index(db: Pool) -> anyhow::Result<()> { - sqlx::query( - "INSERT INTO v2_job_queue (id, workspace_id, scheduled_for, running, tag) - SELECT gen_random_uuid(), 'test-workspace', now() - make_interval(secs => i % 1000), - false, CASE WHEN i % 10 = 0 THEN 'python3' ELSE 'deno' END - FROM generate_series(1, 200000) i", - ) - .execute(&db) - .await?; - sqlx::query("ANALYZE v2_job_queue").execute(&db).await?; - - let mut conn = db.acquire().await?; - sqlx::query("SET plan_cache_mode = force_generic_plan") - .execute(&mut *conn) - .await?; - let query = make_batch_pull_query(&["python3".to_string()], PullQueue::Ready); - sqlx::query(&format!("PREPARE batch_pull(text[]) AS {query}")) - .execute(&mut *conn) - .await?; - let explained: Value = - sqlx::query_scalar("EXPLAIN (FORMAT JSON) EXECUTE batch_pull(ARRAY['w-0', 'w-1'])") - .fetch_one(&mut *conn) - .await?; - - let mut all = vec![]; - nodes(&explained[0]["Plan"], &mut all); - let pretty = serde_json::to_string_pretty(&explained)?; - assert!( - all.iter() - .any(|n| n["Node Type"] == "Index Scan" && n["Index Name"] == "queue_sort_v2"), - "batch peek does not walk queue_sort_v2 in order:\n{pretty}" - ); - assert!( - !all.iter().any(|n| n["Node Type"] == "Sort"), - "batch peek sorts the backlog:\n{pretty}" - ); - Ok(()) -} diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index 449fbd7606..d480a89851 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -774,60 +774,6 @@ impl From<&Pool> for Connection { } } -/// The columns every pull returns, read by name into `windmill_queue::PulledJob`. Expects `q` -/// (the claimed `v2_job_queue` rows), `j` (their `v2_job` rows), `f`, `p` and `pj` in scope. -const PULLED_JOB_COLUMNS: &str = "j.id, j.workspace_id, j.parent_job, j.created_by, q.started_at, q.scheduled_for, - j.runnable_id, j.runnable_path, j.args, q.canceled_by, - q.canceled_reason, j.kind, j.trigger, j.trigger_kind, j.permissioned_as, - f.flow_status, j.script_lang, - j.same_worker, j.pre_run_error, j.visible_to_owner, - j.tag, j.concurrent_limit, j.concurrency_time_window_s, j.flow_innermost_root_job, j.root_job, - j.timeout, j.flow_step_id, j.cache_ttl, q.cache_ignore_s3_path, q.runnable_settings_handle, j.priority, j.raw_code, j.raw_lock, j.raw_flow, - j.script_entrypoint_override, j.preprocessed, COALESCE(pj.runnable_path, j.args->>'_FLOW_PATH') as parent_runnable_path, - COALESCE(p.email, j.permissioned_as_email) as permissioned_as_email, p.username as permissioned_as_username, p.is_admin as permissioned_as_is_admin, - p.is_operator as permissioned_as_is_operator, p.groups as permissioned_as_groups, p.folders as permissioned_as_folders, p.end_user_email as permissioned_as_end_user_email"; - -fn sql_tag_list(tags: &[String]) -> String { - tags.iter() - .map(|x| format!("'{}'", x.replace('\'', "''"))) - .join(", ") -} - -/// Selects the next runnable jobs for `tags`, highest priority first. `exclude_overloaded` names -/// the text[] bind parameter holding workspaces that workspace fairness currently caps. -fn ready_peek(tags: &[String], exclude_overloaded: Option<&str>, limit: &str) -> String { - let fairness = exclude_overloaded - .map(|param| format!("\n AND workspace_id <> ALL({param}::text[])")) - .unwrap_or_default(); - format!( - "SELECT id - FROM v2_job_queue - WHERE running = false AND tag IN ({}) AND scheduled_for <= now(){fairness} - ORDER BY priority DESC NULLS LAST, scheduled_for - FOR UPDATE SKIP LOCKED - LIMIT {limit}", - sql_tag_list(tags) - ) -} - -// The `CASE` is `suspend <= 0 OR suspend_until <= now()` written as one indexable -// expression, equivalent only under the `suspend_until IS NOT NULL` guard. It must stay in -// sync with `queue_suspended_v2` (migration 20260826202939): if it no longer matches, the -// test silently reverts to a heap filter over every suspended row on every worker poll. -fn suspended_peek(tags: &[String], limit: &str) -> String { - format!( - "SELECT id - FROM v2_job_queue - WHERE suspend_until IS NOT NULL - AND (CASE WHEN suspend <= 0 THEN '-infinity'::timestamptz ELSE suspend_until END) <= now() - AND tag IN ({}) - ORDER BY priority DESC NULLS LAST, created_at - FOR UPDATE SKIP LOCKED - LIMIT {limit}", - sql_tag_list(tags) - ) -} - fn format_pull_query(peek: String) -> String { let r = format!( "WITH peek AS ( @@ -857,7 +803,16 @@ fn format_pull_query(peek: String) -> String { raw_flow, script_entrypoint_override, preprocessed FROM v2_job WHERE id = (SELECT id FROM peek) - ) SELECT {PULLED_JOB_COLUMNS} + ) SELECT j.id, j.workspace_id, j.parent_job, j.created_by, started_at, scheduled_for, + j.runnable_id, j.runnable_path, j.args, canceled_by, + canceled_reason, j.kind, j.trigger, j.trigger_kind, j.permissioned_as, + flow_status, j.script_lang, + j.same_worker, j.pre_run_error, j.visible_to_owner, + j.tag, j.concurrent_limit, j.concurrency_time_window_s, j.flow_innermost_root_job, j.root_job, + j.timeout, j.flow_step_id, j.cache_ttl, q.cache_ignore_s3_path, q.runnable_settings_handle, j.priority, j.raw_code, j.raw_lock, j.raw_flow, + j.script_entrypoint_override, j.preprocessed, COALESCE(pj.runnable_path, j.args->>'_FLOW_PATH') as parent_runnable_path, + COALESCE(p.email, j.permissioned_as_email) as permissioned_as_email, p.username as permissioned_as_username, p.is_admin as permissioned_as_is_admin, + p.is_operator as permissioned_as_is_operator, p.groups as permissioned_as_groups, p.folders as permissioned_as_folders, p.end_user_email as permissioned_as_end_user_email FROM q, j LEFT JOIN v2_job_status f USING (id) LEFT JOIN job_perms p ON p.job_id = j.id @@ -869,63 +824,22 @@ fn format_pull_query(peek: String) -> String { r } -/// Claims up to `cardinality($1)` jobs in one statement and assigns the i-th claimed job to the -/// worker named `$1[i]`, so each row is marked running under the worker that will run it. -/// Returns the `PULLED_JOB_COLUMNS` plus `assigned_worker`. -fn format_batch_pull_query(peek: String) -> String { - format!( - "WITH peek AS MATERIALIZED ( - {peek} - ), assign AS ( - SELECT id, ($1::text[])[row_number() OVER ()] AS assigned_worker FROM peek - ), q AS ( - UPDATE v2_job_queue SET - running = true, - started_at = coalesce(started_at, now()), - suspend_until = null, - worker = assign.assigned_worker - FROM assign - WHERE v2_job_queue.id = assign.id - RETURNING - v2_job_queue.id, started_at, scheduled_for, - canceled_by, canceled_reason, worker, cache_ignore_s3_path, runnable_settings_handle - ), r AS ( - UPDATE v2_job_runtime SET - ping = now() - WHERE id = ANY(ARRAY(SELECT id FROM q)) - ) SELECT {PULLED_JOB_COLUMNS}, q.worker AS assigned_worker - FROM q - JOIN v2_job j ON j.id = q.id - LEFT JOIN v2_job_status f ON f.id = q.id - LEFT JOIN job_perms p ON p.job_id = q.id - LEFT JOIN v2_job pj ON j.parent_job = pj.id" - ) -} - -/// Which part of the queue a pull claims from. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum PullQueue { - /// Jobs ready to start. - Ready, - /// Jobs ready to start, excluding the workspaces bound as `$2::text[]`. - ReadyFairness, - /// Suspended flows whose resume condition is met. - Suspended, -} - -/// Batch variant of the pull queries: binds the claiming workers' names as `$1::text[]`, which -/// must be distinct, and claims at most one job per name. -pub fn make_batch_pull_query(tags: &[String], queue: PullQueue) -> String { - const LIMIT: &str = "cardinality($1::text[])"; - format_batch_pull_query(match queue { - PullQueue::Ready => ready_peek(tags, None, LIMIT), - PullQueue::ReadyFairness => ready_peek(tags, Some("$2"), LIMIT), - PullQueue::Suspended => suspended_peek(tags, LIMIT), - }) -} - +// The `CASE` is `suspend <= 0 OR suspend_until <= now()` written as one indexable +// expression, equivalent only under the `suspend_until IS NOT NULL` guard. It must stay in +// sync with `queue_suspended_v2` (migration 20260826202939): if it no longer matches, the +// test silently reverts to a heap filter over every suspended row on every worker poll. pub fn make_suspended_pull_query(tags: &[String]) -> String { - format_pull_query(suspended_peek(tags, "1")) + format_pull_query(format!( + "SELECT id + FROM v2_job_queue + WHERE suspend_until IS NOT NULL + AND (CASE WHEN suspend <= 0 THEN '-infinity'::timestamptz ELSE suspend_until END) <= now() + AND tag IN ({}) + ORDER BY priority DESC NULLS LAST, created_at + FOR UPDATE SKIP LOCKED + LIMIT 1", + tags.iter().map(|x| format!("'{x}'")).join(", ") + )) } // pub async fn make_suspended pub async fn store_suspended_pull_query(wc: &WorkerConfig) { @@ -938,7 +852,16 @@ pub async fn store_suspended_pull_query(wc: &WorkerConfig) { } pub fn make_pull_query(tags: &[String]) -> String { - format_pull_query(ready_peek(tags, None, "1")) + let query = format_pull_query(format!( + "SELECT id + FROM v2_job_queue + WHERE running = false AND tag IN ({}) AND scheduled_for <= now() + ORDER BY priority DESC NULLS LAST, scheduled_for + FOR UPDATE SKIP LOCKED + LIMIT 1", + tags.iter().map(|x| format!("'{x}'")).join(", ") + )); + query } // Variant of `make_pull_query` that additionally excludes jobs whose workspace_id is in the @@ -949,7 +872,17 @@ pub fn make_pull_query(tags: &[String]) -> String { // `pub(crate)` because only `store_pull_query` consumes it; the resulting query string is what // crosses crate boundaries via `WORKER_PULL_QUERIES_FAIRNESS`. pub(crate) fn make_pull_query_fairness(tags: &[String]) -> String { - format_pull_query(ready_peek(tags, Some("$2"), "1")) + let query = format_pull_query(format!( + "SELECT id + FROM v2_job_queue + WHERE running = false AND tag IN ({}) AND scheduled_for <= now() + AND workspace_id <> ALL($2::text[]) + ORDER BY priority DESC NULLS LAST, scheduled_for + FOR UPDATE SKIP LOCKED + LIMIT 1", + tags.iter().map(|x| format!("'{x}'")).join(", ") + )); + query } pub async fn store_pull_query(wc: &WorkerConfig) { diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 8f65638790..7b552fb02c 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -78,9 +78,9 @@ use windmill_common::{ users::{SUPERADMIN_NOTIFICATION_EMAIL, SUPERADMIN_SECRET_EMAIL}, utils::{not_found_if_none, report_critical_error, StripPath, WarnAfterExt}, worker::{ - make_batch_pull_query, to_raw_value, PullQueue, CLOUD_HOSTED, DISABLE_FLOW_SCRIPT, NO_LOGS, - PREVIEW_TAGS_OVERRIDE, WORKER_PULL_QUERIES, WORKER_PULL_QUERIES_FAIRNESS, - WORKER_SUSPENDED_PULL_QUERY, WORKSPACE_FAIRNESS_OVERLOADED, + to_raw_value, CLOUD_HOSTED, DISABLE_FLOW_SCRIPT, NO_LOGS, PREVIEW_TAGS_OVERRIDE, + WORKER_PULL_QUERIES, WORKER_PULL_QUERIES_FAIRNESS, WORKER_SUSPENDED_PULL_QUERY, + WORKSPACE_FAIRNESS_OVERLOADED, }, DB, METRICS_ENABLED, }; @@ -3570,7 +3570,7 @@ impl MiniPulledJob { } } -#[derive(sqlx::FromRow, Debug, Clone, Serialize, Deserialize)] +#[derive(sqlx::FromRow, Debug, Clone)] pub struct PulledJob { #[sqlx(flatten)] pub job: MiniPulledJob, @@ -3837,7 +3837,7 @@ pub async fn get_queued_job_v2<'c>( Ok(job) } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug)] pub struct PulledJobResult { pub job: Option, pub suspended: bool, @@ -4228,13 +4228,6 @@ async fn clone_runnable(j: &mut PulledJob, db: &DB) -> error::Result<()> { Ok(()) } -// Cap to bound DB work per pull cycle. Each iteration on an over-limit job runs -// the full apply_concurrency_limit query stack (~7 queries), so without a tight -// cap a single pull() can issue thousands of queries when the queue is full of -// jobs sharing one over-limit concurrency key. Capping at 10 lets workers skip -// a few stale jobs in healthy conditions while preventing storm amplification. -const PULL_LOOP_LIMIT: i32 = 10; - // TODO: Factorize /// Pull the job from queue pub async fn pull( @@ -4249,9 +4242,14 @@ pub async fn pull( let mut pull_loop_count = 0; loop { pull_loop_count += 1; - if pull_loop_count > PULL_LOOP_LIMIT { + // Cap to bound DB work per pull cycle. Each iteration on an over-limit job runs + // the full apply_concurrency_limit query stack (~7 queries), so without a tight + // cap a single pull() can issue thousands of queries when the queue is full of + // jobs sharing one over-limit concurrency key. Capping at 10 lets workers skip + // a few stale jobs in healthy conditions while preventing storm amplification. + if pull_loop_count > 10 { tracing::warn!( - "Pull job loop count exceeded {PULL_LOOP_LIMIT}, backing off (likely concurrency re-queue storm)" + "Pull job loop count exceeded 10, backing off (likely concurrency re-queue storm)" ); return Ok(PulledJobResult { job: None, @@ -4391,255 +4389,66 @@ pub async fn pull( }); }; - if let Some(pulled_job_res) = admit_pulled_job(db, job, suspended, pull_loop_count).await? { - return Ok(pulled_job_res); - } - } -} - -/// Applies a just-claimed job's concurrency limit, if it has one. `None` means the job was over -/// its limit and has been re-queued for later, so the claiming worker should pull again. -async fn admit_pulled_job( - db: &Pool, - job: PulledJob, - suspended: bool, - pull_loop_count: i32, -) -> windmill_common::error::Result> { - let concurrency_settings = windmill_common::runnable_settings::prefetch_cached_from_handle( - job.runnable_settings_handle, - db, - ) - .await? - .1 - .maybe_fallback(None, job.concurrent_limit, job.concurrency_time_window_s); - - let has_concurent_limit = has_active_concurrency_limit(concurrency_settings.concurrent_limit); - - #[cfg(not(feature = "enterprise"))] - if has_concurent_limit && !job.is_dependency() { - tracing::error!("Concurrent limits are an EE feature only, ignoring constraints") - } - - #[cfg(not(feature = "enterprise"))] - let has_concurent_limit = job.is_dependency() - && has_active_concurrency_limit(job.concurrent_limit) - && cfg!(feature = "private") - && !*WMDEBUG_NO_DEBOUNCING; - // if we don't have private flag, we don't have concurrency limit - - // concurrency check. If more than X jobs for this path are already running, we re-queue and pull another job from the queue - let pulled_job = job; - if pulled_job.runnable_path.is_none() - || !has_concurent_limit - || pulled_job.canceled_by.is_some() - { - #[cfg(feature = "prometheus")] - if METRICS_ENABLED.load(std::sync::atomic::Ordering::Relaxed) { - QUEUE_PULL_COUNT.inc(); - } - otel_incr_queue_pull_count(); - return Ok(Some(PulledJobResult { - job: Some(pulled_job), - suspended, - missing_concurrency_key: false, - error_while_preprocessing: None, - })); - } - - #[cfg(feature = "private")] - if cfg!(feature = "enterprise") || (pulled_job.is_dependency() && !*WMDEBUG_NO_DEBOUNCING) { - return Ok(timeout( - Duration::from_secs(15), - crate::jobs_ee::apply_concurrency_limit( - db, - pull_loop_count, - suspended, - pulled_job, - &concurrency_settings, - ), + let concurrency_settings = windmill_common::runnable_settings::prefetch_cached_from_handle( + job.runnable_settings_handle, + db, ) - .await??); - } - #[cfg(not(feature = "private"))] - let _ = pull_loop_count; - Ok(None) -} + .await? + .1 + .maybe_fallback(None, job.concurrent_limit, job.concurrency_time_window_s); -/// Claims jobs for several waiting workers in one statement per queue walked, instead of one -/// pull per worker. Each distinct name in `worker_names` gets at most one job, marked running -/// under that name; a worker missing from the result found nothing to run. -/// -/// Walks the queues the way [`pull`] does for one worker: when `suspend_first`, suspended flows -/// ready to resume on any of the tags; then `tag_groups` from highest priority down, with the -/// workspace-fairness share of the workers first restricted to uncapped workspaces. A job over -/// its concurrency limit is re-queued for later and its worker is pulled for again. -/// -/// No authorization happens here: jobs from every workspace matching the tags are claimed and -/// returned with their args and permissions. Callers must authenticate each worker named and -/// pass only tags that worker is allowed to serve. -pub async fn pull_batch( - db: &Pool, - tag_groups: &[Vec], - worker_names: &[String], - suspend_first: bool, -) -> windmill_common::error::Result> { - let mut waiting: Vec = worker_names.iter().unique().cloned().collect(); - let mut admitted = vec![]; - for pull_loop_count in 1..=PULL_LOOP_LIMIT { - if waiting.is_empty() { - break; + let has_concurent_limit = + has_active_concurrency_limit(concurrency_settings.concurrent_limit); + + #[cfg(not(feature = "enterprise"))] + if has_concurent_limit && !job.is_dependency() { + tracing::error!("Concurrent limits are an EE feature only, ignoring constraints") } - let claimed = match claim_batch(db, tag_groups, &waiting, suspend_first).await + + #[cfg(not(feature = "enterprise"))] + let has_concurent_limit = job.is_dependency() + && has_active_concurrency_limit(job.concurrent_limit) + && cfg!(feature = "private") + && !*WMDEBUG_NO_DEBOUNCING; + // if we don't have private flag, we don't have concurrency limit + + // concurrency check. If more than X jobs for this path are already running, we re-queue and pull another job from the queue + let pulled_job = job; + if pulled_job.runnable_path.is_none() + || !has_concurent_limit + || pulled_job.canceled_by.is_some() { - Ok(claimed) => claimed, - Err(e) if admitted.is_empty() => return Err(e), - // Jobs admitted by an earlier pass are running under their workers: return them. - Err(e) => { - tracing::error!( - "batch pull re-pull failed, returning the jobs admitted so far: {e:#}" - ); - break; - } - }; - let mut requeued = false; - for (worker_name, job, suspended) in claimed { - waiting.retain(|w| w != &worker_name); - let job_id = job.id; - match admit_pulled_job(db, job, suspended, pull_loop_count).await { - Ok(Some(res)) => admitted.push((worker_name, res)), - Ok(None) => { - waiting.push(worker_name); - requeued = true; - } - // The other claimed jobs are already running under their workers, so an error - // here must not abort the batch: hand this one back to the queue instead. - Err(e) => { - tracing::error!( - "error admitting batch-pulled job {job_id}, re-queuing it: {e:#}" - ); - if let Err(e) = sqlx::query( - "WITH ping AS ( - UPDATE v2_job_runtime SET ping = null WHERE id = $1 - ) - UPDATE v2_job_queue SET running = false, started_at = null - WHERE id = $1 AND worker = $2 AND running = true", - ) - .bind(job_id) - .bind(&worker_name) - .execute(db) - .await - { - tracing::error!("could not re-queue batch-pulled job {job_id}: {e:#}"); - } - } - } - } - if !requeued { - break; - } - if pull_loop_count == PULL_LOOP_LIMIT { - tracing::warn!( - "Batch pull loop count reached {PULL_LOOP_LIMIT}, backing off (likely concurrency re-queue storm)" - ); - } - } - Ok(admitted) -} - -/// One pass over the queues for `worker_names`, returning `(worker, job, suspended)` for each -/// job claimed. A failing statement ends the pass but keeps what was already claimed, since -/// those jobs are marked running and would otherwise be stranded until the zombie monitor. -async fn claim_batch( - db: &Pool, - tag_groups: &[Vec], - worker_names: &[String], - suspend_first: bool, -) -> windmill_common::error::Result> { - let mut claimed = vec![]; - let mut waiting = worker_names.to_vec(); - let pass: windmill_common::error::Result<()> = async { - let tags: Vec = tag_groups.iter().flatten().unique().cloned().collect(); - if suspend_first && !tags.is_empty() { - for (worker, job) in - claim_from_queue(db, &tags, PullQueue::Suspended, &[], &waiting).await? - { - waiting.retain(|w| w != &worker); - claimed.push((worker, job, true)); + #[cfg(feature = "prometheus")] + if METRICS_ENABLED.load(std::sync::atomic::Ordering::Relaxed) { + QUEUE_PULL_COUNT.inc(); } + otel_incr_queue_pull_count(); + return Ok(PulledJobResult { + job: Some(pulled_job), + suspended, + missing_concurrency_key: false, + error_while_preprocessing: None, + }); } - crate::workspace_fairness::maybe_refresh_overloaded(db); - let overloaded = WORKSPACE_FAIRNESS_OVERLOADED.load_full(); - if !overloaded.is_empty() { - // Same admission draw `pull` makes once per worker, made here once per waiting worker. - let mut fair: Vec = waiting - .iter() - .filter(|_| !crate::workspace_fairness::should_admit_capped()) - .cloned() - .collect(); - for tags in tag_groups.iter().filter(|t| !t.is_empty()) { - if fair.is_empty() { - break; - } - for (worker, job) in claim_from_queue( + #[cfg(feature = "private")] + if cfg!(feature = "enterprise") || (pulled_job.is_dependency() && !*WMDEBUG_NO_DEBOUNCING) { + if let Some(pulled_job_res) = timeout( + Duration::from_secs(15), + crate::jobs_ee::apply_concurrency_limit( db, - tags, - PullQueue::ReadyFairness, - overloaded.as_slice(), - &fair, - ) - .await? - { - fair.retain(|w| w != &worker); - waiting.retain(|w| w != &worker); - claimed.push((worker, job, false)); - } - } - } - - for tags in tag_groups.iter().filter(|t| !t.is_empty()) { - if waiting.is_empty() { - break; - } - for (worker, job) in claim_from_queue(db, tags, PullQueue::Ready, &[], &waiting).await? + pull_loop_count, + suspended, + pulled_job, + &concurrency_settings, + ), + ) + .await?? { - waiting.retain(|w| w != &worker); - claimed.push((worker, job, false)); + return Ok(pulled_job_res); } } - Ok(()) } - .await; - match pass { - Err(e) if claimed.is_empty() => Err(e), - Err(e) => { - tracing::error!( - "batch pull pass failed after claiming {} jobs: {e:#}", - claimed.len() - ); - Ok(claimed) - } - Ok(()) => Ok(claimed), - } -} - -async fn claim_from_queue( - db: &Pool, - tags: &[String], - queue: PullQueue, - overloaded: &[String], - worker_names: &[String], -) -> windmill_common::error::Result> { - use sqlx::{FromRow, Row}; - let query = make_batch_pull_query(tags, queue); - let mut query = sqlx::query(&query).bind(worker_names); - if queue == PullQueue::ReadyFairness { - query = query.bind(overloaded); - } - let rows = timeout(Duration::from_secs(15), query.fetch_all(db)).await??; - rows.iter() - .map(|row| Ok((row.try_get("assigned_worker")?, PulledJob::from_row(row)?))) - .collect() } async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>( From 054109d8559cc0a41ef361786c4d1d99eab9c502 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Tue, 29 Sep 2026 15:26:41 +0200 Subject: [PATCH 3/3] fix: create workspace forks in the background, with progress (#11406) * fix: create a fork in the background so a proxy timeout cannot cut it create_fork copies the whole workspace inside the request, which can run past the route timeout of an ingress in front of Windmill (Envoy's 15s default), and the UI then reports a failure for a fork still being made. create_fork?background=true now returns once the request is validated and records the copy in workspace_fork_creation, which the new fork_creation_status endpoint reads. The wizard and AI-session forks use it. Co-Authored-By: Claude Opus 5.5 (1M context) * fix: drain background forks on shutdown and fork in the background from the CLI Co-Authored-By: Claude Opus 5.5 (1M context) * fix: settle fork polls by the fork's existence, adopt in-flight creations Co-Authored-By: Claude Opus 5.5 (1M context) * feat: show which part of the copy a fork being created is in The background copy reports its phase (data tables, settings, resources, scripts, flows, apps, drafts, triggers) through a watch channel. The heartbeat task records it on the fork's creation row as soon as it changes, and fork_creation_status returns it. The fork wizard shows it on its button, and the CLI logs each step. Co-Authored-By: Claude Opus 5.5 (1M context) * fix: join a retried fork creation server-side, settle status by the fork's existence A retry from the same user and parent joins the creation in flight instead of being refused, so clients no longer match the refusal's wording, and another requester can never adopt it. The status route reports a fork that exists under its parent as completed, whatever its run's record says. Clients give up on failing polls after a time window rather than a count, which a rolling deploy can exhaust. Co-Authored-By: Claude Opus 5.5 (1M context) * fix: drop fork-creation joins and existence probes A second request for a fork being created is refused again; only the AI-session fork, whose request never varies, waits for its own earlier one. Clients recognise a server without background forks by its synchronous answer instead of probing for a workspace by id, which could name another parent's fork. Co-Authored-By: Claude Opus 5.5 (1M context) * fix: poll a background fork by the id of its own attempt create_fork?background=true answers with a creation id, and the status route reads that attempt only, for the user who started it. A retry that reuses the fork id is a new attempt, so a poller never reads another attempt's outcome. The AI-session fork no longer adopts a creation in flight, which it could not tell apart from someone else's. Co-Authored-By: Claude Opus 5.5 (1M context) * fix: record a fork's completion in its own commit, resume a session's fork after reload The attempt is marked complete in the transaction that creates the fork, so the status never infers completion from a workspace that may belong to another request. An AI session keeps the creation id on its pending fork and, after a reload mid-copy, waits for that attempt instead of requesting the fork again. Co-Authored-By: Claude Opus 5.5 (1M context) --------- Co-authored-by: Claude Opus 5.5 (1M context) --- backend/Cargo.lock | 1 + backend/Cargo.toml | 2 +- ...929092707_workspace_fork_creation.down.sql | 1 + ...60929092707_workspace_fork_creation.up.sql | 20 + .../tests/fork_background.rs | 120 ++++++ backend/windmill-api-workspaces/Cargo.toml | 1 + .../windmill-api-workspaces/src/workspaces.rs | 346 +++++++++++++++--- backend/windmill-api/openapi.yaml | 43 ++- backend/windmill-api/src/lib.rs | 9 + cli/src/commands/workspace/fork.ts | 46 ++- .../sessions/sessionState.svelte.ts | 33 +- .../components/sessions/sessionState.test.ts | 33 ++ .../CreateWorkspaceInner.svelte | 23 +- frontend/src/lib/utils/forkCreation.test.ts | 63 ++++ frontend/src/lib/utils/forkCreation.ts | 57 +++ 15 files changed, 730 insertions(+), 68 deletions(-) create mode 100644 backend/migrations/20260929092707_workspace_fork_creation.down.sql create mode 100644 backend/migrations/20260929092707_workspace_fork_creation.up.sql create mode 100644 backend/windmill-api-integration-tests/tests/fork_background.rs create mode 100644 frontend/src/lib/utils/forkCreation.test.ts create mode 100644 frontend/src/lib/utils/forkCreation.ts diff --git a/backend/Cargo.lock b/backend/Cargo.lock index e80d0edf08..4d73a953b8 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -15551,6 +15551,7 @@ dependencies = [ "strum", "tokio", "tokio-postgres", + "tokio-util", "tracing", "url", "uuid", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 8571b0b99c..27d8aeac8e 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -519,7 +519,7 @@ serde_urlencoded = "^0" astral-tokio-tar = "^0.6.2" tempfile = "^3" x509-parser = "^0.16" -tokio-util = { version = "=0.7.17", features = ["io"] } +tokio-util = { version = "=0.7.17", features = ["io", "rt"] } json-pointer = "^0" itertools = "^0.14.0" regex = "^1" diff --git a/backend/migrations/20260929092707_workspace_fork_creation.down.sql b/backend/migrations/20260929092707_workspace_fork_creation.down.sql new file mode 100644 index 0000000000..924855e129 --- /dev/null +++ b/backend/migrations/20260929092707_workspace_fork_creation.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS workspace_fork_creation; diff --git a/backend/migrations/20260929092707_workspace_fork_creation.up.sql b/backend/migrations/20260929092707_workspace_fork_creation.up.sql new file mode 100644 index 0000000000..8e3a678419 --- /dev/null +++ b/backend/migrations/20260929092707_workspace_fork_creation.up.sql @@ -0,0 +1,20 @@ +-- A fork being created in the background, and how that ended. The fork's own workspace row only +-- exists once the copy commits, so this is where a failure is reported to whoever polls for it. +-- `heartbeat_at` is refreshed while the copy runs: an unfinished row that stops being refreshed +-- belongs to a server that went away mid-copy. `step` names the part of the copy it is in. +-- `creation_id` names one attempt: the row is reused when a failed or abandoned one is retried, and +-- whoever polls reads only the attempt it started. +CREATE TABLE workspace_fork_creation ( + fork_workspace_id VARCHAR(50) PRIMARY KEY, + parent_workspace_id VARCHAR(50) NOT NULL, + creation_id UUID NOT NULL, + created_by VARCHAR(255) NOT NULL, + started_at TIMESTAMPTZ NOT NULL DEFAULT now(), + heartbeat_at TIMESTAMPTZ NOT NULL DEFAULT now(), + step TEXT, + finished_at TIMESTAMPTZ, + error TEXT +); + +GRANT ALL ON workspace_fork_creation TO windmill_user; +GRANT ALL ON workspace_fork_creation TO windmill_admin; diff --git a/backend/windmill-api-integration-tests/tests/fork_background.rs b/backend/windmill-api-integration-tests/tests/fork_background.rs new file mode 100644 index 0000000000..0e82fb0d70 --- /dev/null +++ b/backend/windmill-api-integration-tests/tests/fork_background.rs @@ -0,0 +1,120 @@ +use serde_json::{json, Value}; +use sqlx::{Pool, Postgres}; + +use windmill_test_utils::*; + +async fn wait_for_fork(client: &reqwest::Client, base_url: &str, creation_id: &str) -> Value { + for _ in 0..100 { + let status: Value = client + .get(format!("{base_url}/fork_creation_status/{creation_id}")) + .header("Authorization", "Bearer SECRET_TOKEN") + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + if status["status"] != "running" { + return status; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + panic!("fork creation {creation_id} still running"); +} + +/// A fork created in the background is reported by `fork_creation_status`, including a failure +/// that happens after the request has already returned. +#[sqlx::test(migrations = "../migrations", fixtures("base"))] +async fn test_fork_created_in_background(db: Pool) -> anyhow::Result<()> { + initialize_tracing().await; + let server = ApiServer::start(db.clone()).await?; + let base_url = format!( + "http://localhost:{}/api/w/test-workspace/workspaces", + server.addr.port() + ); + let client = reqwest::Client::new(); + let create = |id: &str| { + client + .post(format!("{base_url}/create_fork?background=true")) + .header("Authorization", "Bearer SECRET_TOKEN") + .json(&json!({ "id": id, "name": id })) + .send() + }; + + let resp = create("wm-fork-bg-ok").await?; + assert_eq!(resp.status(), 200); + let ok_creation = resp.text().await?; + let status = wait_for_fork(&client, &base_url, &ok_creation).await; + assert_eq!(status, json!({ "status": "completed" })); + // Nobody but the user who started a creation reads it. + let resp = client + .get(format!("{base_url}/fork_creation_status/{ok_creation}")) + .header("Authorization", "Bearer SECRET_TOKEN_2") + .send() + .await?; + assert_eq!(resp.status(), 404); + let exists: bool = + sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM workspace WHERE id = 'wm-fork-bg-ok')") + .fetch_one(&db) + .await?; + assert!(exists); + // Under CE's workspace cap. + sqlx::query("UPDATE workspace SET deleted = true WHERE id = 'wm-fork-bg-ok'") + .execute(&db) + .await?; + + // Hold the fork's write lock so the request is validated and returns, then commit a workspace + // with the same id: the background write fails on it. + let mut blocker = db.begin().await?; + sqlx::query("SELECT pg_advisory_xact_lock(hashtext('fork:wm-fork-bg-err'))") + .execute(&mut *blocker) + .await?; + let resp = create("wm-fork-bg-err").await?; + assert_eq!(resp.status(), 200); + let err_creation = resp.text().await?; + let resp = create("wm-fork-bg-err").await?; + assert_eq!(resp.status(), 400, "a second creation while the first is in flight"); + sqlx::query("INSERT INTO workspace (id, name, owner) VALUES ('wm-fork-bg-err', 'x', 'x')") + .execute(&mut *blocker) + .await?; + blocker.commit().await?; + let status = wait_for_fork(&client, &base_url, &err_creation).await; + assert_eq!(status["status"], "failed"); + assert!( + status["error"].as_str().unwrap().contains("workspace_pkey"), + "{status}" + ); + + // A retry of the same id is a new attempt: the failed one does not read its outcome. + sqlx::query("DELETE FROM workspace WHERE id = 'wm-fork-bg-err'") + .execute(&db) + .await?; + let resp = create("wm-fork-bg-err").await?; + assert_eq!(resp.status(), 200); + let retry_creation = resp.text().await?; + assert_eq!( + wait_for_fork(&client, &base_url, &retry_creation).await, + json!({ "status": "completed" }) + ); + let resp = client + .get(format!("{base_url}/fork_creation_status/{err_creation}")) + .header("Authorization", "Bearer SECRET_TOKEN") + .send() + .await?; + assert_eq!(resp.status(), 404); + + // Without the flag the fork is created before the answer, which clients asking for a + // background fork recognise a server without it by. + sqlx::query("UPDATE workspace SET deleted = true WHERE id = 'wm-fork-bg-err'") + .execute(&db) + .await?; + let resp = client + .post(format!("{base_url}/create_fork")) + .header("Authorization", "Bearer SECRET_TOKEN") + .json(&json!({ "id": "wm-fork-sync", "name": "wm-fork-sync" })) + .send() + .await?; + assert_eq!(resp.text().await?, "Created forked workspace wm-fork-sync"); + + Ok(()) +} diff --git a/backend/windmill-api-workspaces/Cargo.toml b/backend/windmill-api-workspaces/Cargo.toml index 7ca1f692ad..9f05a26f28 100644 --- a/backend/windmill-api-workspaces/Cargo.toml +++ b/backend/windmill-api-workspaces/Cargo.toml @@ -50,6 +50,7 @@ sha2.workspace = true sqlx.workspace = true tokio.workspace = true tokio-postgres.workspace = true +tokio-util.workspace = true tracing.workspace = true url.workspace = true uuid.workspace = true diff --git a/backend/windmill-api-workspaces/src/workspaces.rs b/backend/windmill-api-workspaces/src/workspaces.rs index 458e63a973..31d049d3c0 100644 --- a/backend/windmill-api-workspaces/src/workspaces.rs +++ b/backend/windmill-api-workspaces/src/workspaces.rs @@ -176,6 +176,10 @@ pub fn workspaced_service() -> Router { .route("/leave", post(leave_workspace)) .route("/get_workspace_name", get(get_workspace_name)) .route("/create_fork", post(create_workspace_fork)) + .route( + "/fork_creation_status/{creation_id}", + get(get_fork_creation_status), + ) .route("/attach_dev_workspace", post(attach_dev_workspace)) .route("/detach_dev_workspace", post(detach_dev_workspace)) .route("/get_dev_workspace", get(get_dev_workspace)) @@ -6924,7 +6928,9 @@ async fn clone_workspace_data( source_workspace_id: &str, target_workspace_id: &str, authed: &ApiAuthed, + progress: &ForkProgress, ) -> Result<()> { + progress.step("Copying settings, folders and groups"); // Clone workspace settings (merge with existing basic settings) update_workspace_settings(tx, source_workspace_id, target_workspace_id).await?; @@ -6946,6 +6952,7 @@ async fn clone_workspace_data( // Clone groups clone_groups(tx, source_workspace_id, target_workspace_id).await?; + progress.step("Copying resources and variables"); // Clone resource types clone_resource_types(tx, source_workspace_id, target_workspace_id).await?; @@ -6955,6 +6962,7 @@ async fn clone_workspace_data( // Clone variables (including external secret backend replication) clone_variables(tx, db, source_workspace_id, target_workspace_id).await?; + progress.step("Copying scripts"); // Clone scripts with new hashes clone_scripts(tx, source_workspace_id, target_workspace_id).await?; @@ -6972,18 +6980,21 @@ async fn clone_workspace_data( clone_metric_catalog(tx, source_workspace_id, target_workspace_id).await?; clone_asset_usages_and_triggers(tx, source_workspace_id, target_workspace_id).await?; + progress.step("Copying flows"); // Clone flows with new versions clone_flows(tx, source_workspace_id, target_workspace_id).await?; // Clone flow nodes clone_flow_nodes(tx, source_workspace_id, target_workspace_id).await?; + progress.step("Copying apps"); // Clone apps with new IDs and app scripts let _app_id_mapping = clone_apps(tx, source_workspace_id, target_workspace_id, authed).await?; // Clone raw apps clone_raw_apps(tx, source_workspace_id, target_workspace_id).await?; + progress.step("Copying drafts and dependencies"); // Clone the forker's own per-user drafts (plus the legacy NULL-email // workspace draft, if any) so they keep their pending edits in the // fork. Other users' drafts are intentionally NOT cloned — they don't @@ -9209,6 +9220,7 @@ async fn create_workspace_fork( authed: ApiAuthed, Extension(db): Extension, Path(parent_workspace_id): Path, + Query(q): Query, Json(nw): Json, ) -> Result { enforce_fork_depth(&db, &parent_workspace_id, 0).await?; @@ -9314,63 +9326,279 @@ async fn create_workspace_fork( // Refused here, before any database exists; `make_copies` checks again under the locks. validate_forked_datatables(&db, &authed, &parent_workspace_id, &nw).await?; - // Detached from the request: a client that goes away while the copies are made must still end - // with the fork created or no copy left behind, and dropping the handler's future would skip - // that cleanup. - tokio::spawn(async move { - let fork_id = nw.id.clone(); - let copies = crate::datatable_clone::make_copies( - &db, - &authed, - &parent_workspace_id, - ©_requests(&nw.forked_datatables), - ) - .await?; - let replayed: Vec = copies - .iter() - .filter(|c| c.replayed) - .map(|c| c.behavior) - .collect(); - match write_workspace_fork( + if !q.background.unwrap_or(false) { + // Detached from the request: a client that goes away while the copies are made must still + // end with the fork created or no copy left behind, and dropping the handler's future + // would skip that cleanup. + return tokio::spawn(make_workspace_fork( + db, + authed, + parent_workspace_id, + nw, + dev_workspace_label, + ForkProgress::default(), + )) + .await + .map_err(|e| { + Error::internal_err(format!("Creating the fork stopped unexpectedly: {e}")) + })?; + } + + let fork_id = nw.id.clone(); + sqlx::query( + "DELETE FROM workspace_fork_creation + WHERE COALESCE(finished_at, heartbeat_at) < now() - interval '7 days'", + ) + .execute(&db) + .await?; + let claimed = sqlx::query( + "INSERT INTO workspace_fork_creation + (fork_workspace_id, parent_workspace_id, created_by, creation_id) + VALUES ($1, $2, $3, $5) + ON CONFLICT (fork_workspace_id) DO UPDATE SET + parent_workspace_id = EXCLUDED.parent_workspace_id, + created_by = EXCLUDED.created_by, + creation_id = EXCLUDED.creation_id, + started_at = now(), heartbeat_at = now(), step = NULL, finished_at = NULL, + error = NULL + WHERE workspace_fork_creation.finished_at IS NOT NULL + OR workspace_fork_creation.heartbeat_at < now() - $4 * interval '1 second' + RETURNING creation_id", + ) + .bind(&fork_id) + .bind(&parent_workspace_id) + .bind(&authed.email) + .bind(FORK_HEARTBEAT_STALE_SECS as f64) + .bind(Uuid::new_v4()) + .fetch_optional(&db) + .await?; + // Every write of this attempt is scoped to its `creation_id`: one taken over as abandoned must + // not report over the attempt that took its place. + let Some(creation_id) = claimed.map(|r| r.get::("creation_id")) else { + return Err(Error::BadRequest(format!( + "workspace '{fork_id}' is already being created" + ))); + }; + + BACKGROUND_FORKS.spawn(async move { + let (progress_tx, mut progress_rx) = tokio::sync::watch::channel(""); + // The one writer of the run's row while it copies: it records a new step as soon as the + // copy reports it, and refreshes the heartbeat in between. + let heartbeat = { + let (db, fork_id) = (db.clone(), fork_id.clone()); + tokio::spawn(async move { + let mut interval = + tokio::time::interval(std::time::Duration::from_secs(FORK_HEARTBEAT_SECS)); + let mut reporting = true; + loop { + tokio::select! { + _ = interval.tick() => {} + changed = progress_rx.changed(), if reporting => reporting = changed.is_ok(), + } + let step = *progress_rx.borrow_and_update(); + if let Err(e) = sqlx::query( + "UPDATE workspace_fork_creation SET heartbeat_at = now(), step = $3 + WHERE fork_workspace_id = $1 AND creation_id = $2", + ) + .bind(&fork_id) + .bind(creation_id) + .bind((!step.is_empty()).then_some(step)) + .execute(&db) + .await + { + tracing::warn!("fork '{fork_id}': could not record progress: {e}"); + } + } + }) + }; + let made = tokio::spawn(make_workspace_fork( db.clone(), authed, parent_workspace_id, nw, dev_workspace_label, - &copies, + ForkProgress { creation_id: Some(creation_id), steps: Some(progress_tx) }, + )) + .await + .map_err(|e| Error::internal_err(format!("Creating the fork stopped unexpectedly: {e}"))) + .and_then(|r| r); + heartbeat.abort(); + let error = made.err().map(|e| { + tracing::error!("Creating fork '{fork_id}' failed: {e}"); + e.to_string() + }); + // A fork that committed is already recorded complete, even when its commit's + // acknowledgement was lost and the attempt reads as an error here. + if let Err(e) = sqlx::query( + "UPDATE workspace_fork_creation SET finished_at = now(), error = $3 + WHERE fork_workspace_id = $1 AND creation_id = $2 AND finished_at IS NULL", ) + .bind(&fork_id) + .bind(creation_id) + .bind(error) + .execute(&db) .await { - Ok(message) => { - for behavior in replayed { - windmill_common::feature_usage::log_feature_usage( - "datatable", - "clone_replayed", - match behavior { - DataTableForkBehavior::SchemaOnly => "schema_only", - _ => "schema_and_data", - }, - ); - } - Ok(message) - } - // A commit whose acknowledgement was lost can still have committed, and a concurrent - // request for the same id can have: once the write has settled, the cleanup keeps - // whichever copies a committed fork names, and drops the rest. - Err(e) => match wait_for_fork_write(&db, &fork_id).await { - Ok(()) => Err(crate::datatable_clone::drop_copies_after(&db, copies, e).await), - Err(settle) => { - tracing::error!( - "Could not tell whether fork '{fork_id}' was created, so its copies were \ - kept: {settle}" - ); - Err(e) - } - }, + tracing::error!("fork '{fork_id}': could not record how its creation ended: {e}"); } - }) + }); + Ok(creation_id.to_string()) +} + +/// How often a fork created in the background records that its copy is still running, and how long +/// without that record its creation counts as abandoned. +const FORK_HEARTBEAT_SECS: u64 = 10; +const FORK_HEARTBEAT_STALE_SECS: u64 = 60; + +/// Forks being created in the background. The server waits for them on shutdown: a copy cut short +/// never drops the data table databases it already made, and a retry under the same id is refused +/// while they exist. +pub static BACKGROUND_FORKS: std::sync::LazyLock = + std::sync::LazyLock::new(tokio_util::task::TaskTracker::new); + +#[derive(Deserialize)] +struct CreateWorkspaceForkQuery { + /// Return once the request is validated and create the fork in the background, answering with + /// the id `fork_creation_status` reads this attempt by. Without it the answer is + /// `Created forked workspace `, which clients take as a server without background forks. + background: Option, +} + +/// The attempt a fork created in the background belongs to, and where its copy is, for its status +/// to report. +#[derive(Default)] +struct ForkProgress { + creation_id: Option, + steps: Option>, +} + +impl ForkProgress { + fn step(&self, step: &'static str) { + if let Some(tx) = &self.steps { + tx.send_replace(step); + } + } +} + +#[derive(Serialize)] +struct ForkCreationStatus { + /// `running`, `completed` or `failed`. + status: &'static str, + /// The part of the copy a running fork is in. + #[serde(skip_serializing_if = "Option::is_none")] + step: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +async fn get_fork_creation_status( + authed: ApiAuthed, + Extension(db): Extension, + Path((parent_workspace_id, creation_id)): Path<(String, Uuid)>, +) -> JsonResult { + let row = sqlx::query( + "SELECT c.created_by, c.finished_at IS NOT NULL AS finished, c.error, c.step, + c.heartbeat_at < now() - $3 * interval '1 second' AS stale + FROM workspace_fork_creation c + WHERE c.creation_id = $1 AND c.parent_workspace_id = $2", + ) + .bind(creation_id) + .bind(&parent_workspace_id) + .bind(FORK_HEARTBEAT_STALE_SECS as f64) + .fetch_optional(&db) + .await?; + // Only the user who started it reads an attempt: its id is not handed to anyone else. + let Some(row) = row.filter(|r| r.get::("created_by") == authed.email) else { + return Err(Error::NotFound(format!( + "no fork creation {creation_id} from '{parent_workspace_id}'" + ))); + }; + let error: Option = row.get("error"); + let status = if row.get::("finished") { + match error { + Some(error) => ForkCreationStatus { status: "failed", step: None, error: Some(error) }, + None => ForkCreationStatus { status: "completed", step: None, error: None }, + } + } else if row.get::("stale") { + ForkCreationStatus { + status: "failed", + step: None, + error: Some( + "the server creating the fork stopped before it finished; check the server logs \ + and whether the fork exists before retrying" + .to_string(), + ), + } + } else { + ForkCreationStatus { status: "running", step: row.get("step"), error: None } + }; + Ok(Json(status)) +} + +/// Copy the data tables the fork gets, then write the fork itself, dropping the copies again if +/// that write fails. +async fn make_workspace_fork( + db: DB, + authed: ApiAuthed, + parent_workspace_id: String, + nw: CreateWorkspaceFork, + dev_workspace_label: Option, + progress: ForkProgress, +) -> Result { + let fork_id = nw.id.clone(); + if !nw.forked_datatables.is_empty() { + progress.step("Copying data tables"); + } + let copies = crate::datatable_clone::make_copies( + &db, + &authed, + &parent_workspace_id, + ©_requests(&nw.forked_datatables), + ) + .await?; + let replayed: Vec = copies + .iter() + .filter(|c| c.replayed) + .map(|c| c.behavior) + .collect(); + match write_workspace_fork( + db.clone(), + authed, + parent_workspace_id, + nw, + dev_workspace_label, + &copies, + &progress, + ) .await - .map_err(|e| Error::internal_err(format!("Creating the fork stopped unexpectedly: {e}")))? + { + Ok(message) => { + for behavior in replayed { + windmill_common::feature_usage::log_feature_usage( + "datatable", + "clone_replayed", + match behavior { + DataTableForkBehavior::SchemaOnly => "schema_only", + _ => "schema_and_data", + }, + ); + } + Ok(message) + } + // A commit whose acknowledgement was lost can still have committed, and a concurrent + // request for the same id can have: once the write has settled, the cleanup keeps + // whichever copies a committed fork names, and drops the rest. + Err(e) => match wait_for_fork_write(&db, &fork_id).await { + Ok(()) => Err(crate::datatable_clone::drop_copies_after(&db, copies, e).await), + Err(settle) => { + tracing::error!( + "Could not tell whether fork '{fork_id}' was created, so its copies were \ + kept: {settle}" + ); + Err(e) + } + }, + } } /// The database a data table of the fork or dev workspace `fork_id` is copied into, as the wizard @@ -9466,6 +9694,7 @@ async fn write_workspace_fork( nw: CreateWorkspaceFork, dev_workspace_label: Option, copies: &[crate::datatable_clone::MadeCopy], + progress: &ForkProgress, ) -> Result { let mut tx: Transaction<'_, Postgres> = db.begin().await?; // Held until this transaction ends: after an error, the copies' cleanup waits on it, so it reads @@ -9568,8 +9797,15 @@ async fn write_workspace_fork( .await?; // Clone all data from the parent workspace using Rust implementation - if let Err(e) = - clone_workspace_data(&mut tx, &db, &parent_workspace_id, &forked_id, &authed).await + if let Err(e) = clone_workspace_data( + &mut tx, + &db, + &parent_workspace_id, + &forked_id, + &authed, + progress, + ) + .await { // A genuine `\u0000` in a source `json` value (`app_version.value` / // `flow_version.schema`) aborts the clone when it is re-encoded to jsonb: @@ -9600,6 +9836,7 @@ async fn write_workspace_fork( // enabled=false. Disabled rows have no side effects (no listener // attaches, no cron fires) so this is safe by construction. The user // re-enables in the fork, with parent-conflict warnings on enable. + progress.step("Copying triggers and schedules"); clone_triggers_and_schedules(&mut tx, &parent_workspace_id, &forked_id).await?; repoint_unresolvable_cloned_identities(&mut tx, &forked_id, &authed).await?; @@ -9747,6 +9984,16 @@ async fn write_workspace_fork( (!copied.is_empty()).then(|| [("copied_datatables", copied.as_str())].into()), ) .await?; + // Committed with the fork, so its attempt reads complete exactly when the fork exists. + if let Some(creation_id) = progress.creation_id { + sqlx::query( + "UPDATE workspace_fork_creation SET finished_at = now(), error = NULL + WHERE creation_id = $1", + ) + .bind(creation_id) + .execute(&mut *tx) + .await?; + } tx.commit().await?; // A pre-creation lookup could have cached an EMPTY ancestor chain for this id, which @@ -9760,6 +10007,7 @@ async fn write_workspace_fork( windmill_common::workspaces::invalidate_protection_rules_cache(&parent_workspace_id); } + // Clients that ask for a background fork recognise a server without it by this answer. Ok(format!("Created forked workspace {}", &forked_id)) } diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml index cf373241d4..e7afd78ceb 100644 --- a/backend/windmill-api/openapi.yaml +++ b/backend/windmill-api/openapi.yaml @@ -1398,6 +1398,14 @@ paths: - workspace parameters: - $ref: "#/components/parameters/WorkspaceId" + - name: background + description: | + return once the request is validated and create the fork in the background, answering + with the creation id getForkCreationStatus reads + in: query + required: false + schema: + type: boolean requestBody: description: new forked workspace required: true @@ -1407,12 +1415,45 @@ paths: $ref: "#/components/schemas/CreateWorkspaceFork" responses: "201": - description: forked workspace created + description: forked workspace created, or with background set, the id of its creation content: text/plain: schema: type: string + /w/{workspace}/workspaces/fork_creation_status/{creation_id}: + get: + summary: get the status of a fork created in the background + operationId: getForkCreationStatus + tags: + - workspace + parameters: + - $ref: "#/components/parameters/WorkspaceId" + - name: creation_id + in: path + required: true + schema: + type: string + format: uuid + responses: + "200": + description: fork creation status + content: + application/json: + schema: + type: object + properties: + status: + type: string + enum: [running, completed, failed] + step: + type: string + description: the part of the copy a running fork is in + error: + type: string + required: + - status + /w/{workspace}/workspaces/attach_dev_workspace: post: summary: attach an existing workspace as the dev workspace of this (prod) workspace diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index 156c6ca898..556163ce1c 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -1279,6 +1279,15 @@ pub async fn run_server( }); server.await?; + let forks = &windmill_api_workspaces::workspaces::BACKGROUND_FORKS; + forks.close(); + if !forks.is_empty() { + tracing::info!( + "server off. waiting for {} fork(s) being created", + forks.len() + ); + forks.wait().await; + } #[cfg(feature = "agent_worker_server")] for (i, bg_processor) in agent_workers_bg_processor.into_iter().enumerate() { tracing::info!("server off. shutting down agent worker bg processor {i}"); diff --git a/cli/src/commands/workspace/fork.ts b/cli/src/commands/workspace/fork.ts index 0aca0984ff..eedd9c0842 100644 --- a/cli/src/commands/workspace/fork.ts +++ b/cli/src/commands/workspace/fork.ts @@ -13,6 +13,7 @@ import { renameCurrentGitBranch, } from "../../utils/git.ts"; import process from "node:process"; +import { apiErrorMessage } from "../../utils/utils.ts"; import { WM_FORK_PREFIX } from "../../core/constants.ts"; import { tryResolveBranchWorkspace } from "../../core/context.ts"; import { @@ -22,6 +23,38 @@ import { readConfigFile, } from "../../core/conf.ts"; +/** Wait for a fork creation started in the background, by the id its request answered with. */ +async function waitForForkCreation(parentWorkspace: string, creationId: string) { + let failingSince: number | undefined; + let lastStep: string | undefined; + while (true) { + let result: + | Awaited> + | undefined; + try { + result = await wmill.getForkCreationStatus({ + workspace: parentWorkspace, + creationId, + }); + failingSince = undefined; + } catch (e) { + // A failed poll says nothing about the fork: it can be lost to a proxy, or reach a replica + // that has no status route yet during a rolling deploy. + failingSince ??= Date.now(); + if (Date.now() - failingSince >= 3 * 60 * 1000) throw e; + } + if (result?.status === "running" && result.step && result.step !== lastStep) { + lastStep = result.step; + log.info(colors.gray(`${lastStep}...`)); + } + if (result?.status === "completed") return; + if (result?.status === "failed") { + throw new Error(result.error ?? "Unknown error"); + } + await new Promise((resolve) => setTimeout(resolve, 1500)); + } +} + async function createWorkspaceFork( opts: GlobalOptions & { createWorkspaceName: string | undefined; @@ -376,8 +409,9 @@ async function createWorkspaceFork( // --- Create the fork workspace --- try { - const result = await wmill.createWorkspaceFork({ + const response = await wmill.createWorkspaceFork({ workspace: workspace.workspaceId, + background: true, requestBody: { id: trueWorkspaceId, name: opts.createWorkspaceName ?? workspaceName ?? trueWorkspaceId, @@ -385,11 +419,17 @@ async function createWorkspaceFork( forked_datatables: forkedDatatables, }, }); + // A server without background forks ignores the flag and answers once the fork is created. + if (!response.startsWith("Created forked workspace")) { + await waitForForkCreation(workspace.workspaceId, response); + } - log.info(colors.green(`✅ ${result}`)); + log.info(colors.green(`✅ Created forked workspace ${trueWorkspaceId}`)); } catch (error) { log.error( - colors.red(`Failed to create forked workspace: ${(error as Error).message}`), + colors.red( + `Failed to create forked workspace: ${apiErrorMessage(error) ?? (error as Error).message}`, + ), ); throw error; } diff --git a/frontend/src/lib/components/sessions/sessionState.svelte.ts b/frontend/src/lib/components/sessions/sessionState.svelte.ts index c5a2f28455..e81df02120 100644 --- a/frontend/src/lib/components/sessions/sessionState.svelte.ts +++ b/frontend/src/lib/components/sessions/sessionState.svelte.ts @@ -45,6 +45,7 @@ export function syncWorkspaceTo(workspaceId: string | undefined): void { switchWorkspace(workspaceId) } import { WorkspaceService } from '$lib/gen' +import { createWorkspaceForkAndWait, waitForForkCreation } from '$lib/utils/forkCreation' import { sendUserToast } from '$lib/toast' import { onUserChange } from '$lib/userScopedStorage' @@ -76,6 +77,9 @@ export type PendingFork = { id: string // Display name shown in the workspace bar. name: string + // The background creation this fork started, so a reload mid-copy waits for it rather than + // requesting the fork again, which the server refuses while that creation runs. + creation_id?: string } export type SessionSummarySource = 'placeholder' | 'generated' | 'manual' @@ -1180,7 +1184,10 @@ export async function commitSessionWorkspace( `Community edition is limited to ${CE_MAX_NON_ADMIN_WORKSPACES + 1} workspaces — archive a workspace or pick one to run in` ) } - const newId = await materializeFork(fork) + const newId = await materializeFork(fork, (creation_id) => { + s.pending_fork = { ...fork, creation_id } + void putSession(s) + }) if (!newId) { // Real failure (not a recovered duplicate). Drop the pending // fork so the session falls through to the workspace-pick @@ -1297,13 +1304,27 @@ export function setGeneratedSessionSummary( // lost). Adopt it silently instead of re-POSTing — the API would // otherwise reject with workspace_pkey. Likewise, if the API returns a // duplicate-key error we refresh the store and adopt the existing row. -export async function materializeFork(fork: PendingFork): Promise { +export async function materializeFork( + fork: PendingFork, + onCreationStarted?: (creationId: string) => void +): Promise { if (get(userWorkspaces).some((w) => w.id === fork.id)) return fork.id try { - await WorkspaceService.createWorkspaceFork({ - workspace: fork.parent_workspace_id, - requestBody: { id: fork.id, name: fork.name } - }) + let resumed = false + if (fork.creation_id) { + // An attempt that failed or is gone is requested again below. + resumed = await waitForForkCreation(fork.parent_workspace_id, fork.creation_id).then( + () => true, + () => false + ) + } + if (!resumed) { + await createWorkspaceForkAndWait( + fork.parent_workspace_id, + { id: fork.id, name: fork.name }, + { onCreationStarted } + ) + } usersWorkspaceStore.set(await WorkspaceService.listUserWorkspaces()) sendUserToast(`Created fork ${fork.name}`) return fork.id diff --git a/frontend/src/lib/components/sessions/sessionState.test.ts b/frontend/src/lib/components/sessions/sessionState.test.ts index 5cf3e3269e..62d2b8d39f 100644 --- a/frontend/src/lib/components/sessions/sessionState.test.ts +++ b/frontend/src/lib/components/sessions/sessionState.test.ts @@ -37,6 +37,7 @@ vi.mock('$lib/gen', async (orig) => { WorkspaceService: { ...actual.WorkspaceService, createWorkspaceFork: vi.fn().mockRejectedValue(new Error('fork creation failed')), + getForkCreationStatus: vi.fn().mockResolvedValue({ status: 'completed' }), listUserWorkspaces: vi.fn().mockResolvedValue([]) } } @@ -108,6 +109,38 @@ describe('commitSessionWorkspace — fork-creation failure', () => { }) }) +describe('commitSessionWorkspace — fork creation resumed after a reload', () => { + it('waits for the creation the pending fork started instead of requesting it again', async () => { + const id = 'test-commit-fork-resume' + const prevLicense = get(enterpriseLicense) + enterpriseLicense.set('test-license') + vi.mocked(WorkspaceService.createWorkspaceFork).mockClear() + sessionState.sessions.push({ + id, + name: 'fork-resume', + createdAt: 0, + pending_fork: { + parent_workspace_id: 'parent_ws', + id: 'wm-fork-resume', + name: 'resume', + creation_id: '5c3e1b1e-0000-4000-8000-000000000000' + } + } as Session) + try { + expect(await commitSessionWorkspace(id, 'parent_ws')).toBe('wm-fork-resume') + expect(WorkspaceService.getForkCreationStatus).toHaveBeenCalledWith({ + workspace: 'parent_ws', + creationId: '5c3e1b1e-0000-4000-8000-000000000000' + }) + expect(WorkspaceService.createWorkspaceFork).not.toHaveBeenCalled() + } finally { + const i = sessionState.sessions.findIndex((x) => x.id === id) + if (i >= 0) sessionState.sessions.splice(i, 1) + enterpriseLicense.set(prevLicense) + } + }) +}) + // Set the workspace list to a given number of non-'admins' workspaces so the // commit-path's CE workspace-cap check (mirror of backend _check_nb_of_workspaces) // can be exercised. Returns a restore fn. diff --git a/frontend/src/lib/components/workspaceSettings/CreateWorkspaceInner.svelte b/frontend/src/lib/components/workspaceSettings/CreateWorkspaceInner.svelte index d0f86b36ad..e6a0dce4fd 100644 --- a/frontend/src/lib/components/workspaceSettings/CreateWorkspaceInner.svelte +++ b/frontend/src/lib/components/workspaceSettings/CreateWorkspaceInner.svelte @@ -70,6 +70,7 @@ import WorkspaceScopeTrigger from '../WorkspaceScopeTrigger.svelte' import DarkModeToggle from '../sidebar/DarkModeToggle.svelte' import ForkDucklakeSection from './ForkDucklakeSection.svelte' + import { createWorkspaceForkAndWait } from '$lib/utils/forkCreation' interface Props { isFork?: boolean @@ -329,6 +330,7 @@ } let forkCreationLoading = $state(false) + let forkCreationStep = $state('Creating branch') let forkCreationError = $state('') let errorMsgs: string[] = $state([]) let failedSyncJobs: string[] = $state([]) @@ -376,6 +378,7 @@ async function forkWorkspace(prefixed_id: string): Promise { if (baseWorkspaceId) { forkCreationLoading = true + forkCreationStep = 'Creating branch' errorMsgs = [] failedSyncJobs = [] forkCreationError = '' @@ -470,9 +473,10 @@ } try { - await WorkspaceService.createWorkspaceFork({ - workspace: baseWorkspaceId!, - requestBody: { + forkCreationStep = 'Creating fork' + await createWorkspaceForkAndWait( + baseWorkspaceId!, + { id: prefixed_id, name, color: colorEnabled && workspaceColor ? workspaceColor : undefined, @@ -483,13 +487,15 @@ lock_prod_deploy: createAsDevWorkspace && effectiveLockProdDeploy, lock_prod_forking: createAsDevWorkspace && effectiveLockProdForking, copy_members: copyMembers - } - }) + }, + { onStep: (step) => (forkCreationStep = step) } + ) } catch (e) { + const msg = e?.body ?? e?.message ?? e ?? 'Unknown error' forkCreationError = `Failed to create fork '${prefixed_id}'` - errorMsgs.push(e?.body ?? e ?? 'Unknown error') + errorMsgs.push(msg) forkCreationLoading = false - sendUserToast(`Could not create fork '${prefixed_id}' ${e}`, true) + sendUserToast(`Could not create fork '${prefixed_id}' ${msg}`, true) return } @@ -1115,7 +1121,8 @@ {:else} {/if} diff --git a/frontend/src/lib/utils/forkCreation.test.ts b/frontend/src/lib/utils/forkCreation.test.ts new file mode 100644 index 0000000000..bf180c6c6f --- /dev/null +++ b/frontend/src/lib/utils/forkCreation.test.ts @@ -0,0 +1,63 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +vi.mock('$lib/gen', () => ({ + WorkspaceService: { + createWorkspaceFork: vi.fn(), + getForkCreationStatus: vi.fn() + } +})) + +import { WorkspaceService } from '$lib/gen' +import { createWorkspaceForkAndWait } from './forkCreation' + +const svc = vi.mocked(WorkspaceService) +const fork = { id: 'wm-fork-x', name: 'x' } + +async function run(onStep?: (step: string) => void) { + const p = createWorkspaceForkAndWait('ws', fork, { onStep }) + const settled = p.then( + () => 'completed', + (e) => `failed: ${e.message}` + ) + await vi.runAllTimersAsync() + return settled +} + +describe('createWorkspaceForkAndWait', () => { + beforeEach(() => { + vi.useFakeTimers() + vi.resetAllMocks() + svc.createWorkspaceFork.mockResolvedValue('5c3e1b1e-0000-4000-8000-000000000000') + }) + afterEach(() => vi.useRealTimers()) + + it('reports each step while the fork runs, then how it ended', async () => { + svc.getForkCreationStatus + .mockResolvedValueOnce({ status: 'running' }) + .mockResolvedValueOnce({ status: 'running', step: 'Copying flows' }) + .mockResolvedValueOnce({ status: 'failed', error: 'boom' }) + const steps: string[] = [] + expect(await run((s) => steps.push(s))).toBe('failed: boom') + expect(steps).toEqual(['Copying flows']) + }) + + it('returns at once when the server created the fork synchronously', async () => { + // A server without background forks ignores the flag. + svc.createWorkspaceFork.mockResolvedValue('Created forked workspace wm-fork-x') + expect(await run()).toBe('completed') + expect(svc.getForkCreationStatus).not.toHaveBeenCalled() + }) + + it('keeps polling through failed polls', async () => { + svc.getForkCreationStatus + .mockRejectedValueOnce(Object.assign(new Error('Not Found'), { status: 404 })) + .mockRejectedValueOnce(new Error('upstream request timeout')) + .mockResolvedValueOnce({ status: 'completed' }) + expect(await run()).toBe('completed') + }) + + it('gives up once polls keep failing for a fork that never appears', async () => { + svc.getForkCreationStatus.mockRejectedValue(new Error('upstream request timeout')) + expect(await run()).toBe('failed: upstream request timeout') + }) +}) diff --git a/frontend/src/lib/utils/forkCreation.ts b/frontend/src/lib/utils/forkCreation.ts new file mode 100644 index 0000000000..d0fb4e5edf --- /dev/null +++ b/frontend/src/lib/utils/forkCreation.ts @@ -0,0 +1,57 @@ +import { WorkspaceService, type CreateWorkspaceFork } from '$lib/gen' + +const POLL_INTERVAL_MS = 1500 +// How long polls may keep failing, which a rolling deploy's replicas without the status route cause. +const MAX_FAILING_POLLS_MS = 3 * 60 * 1000 + +/** + * Create a fork and resolve once it exists, rejecting with the reason it could not be created. + * + * The copy runs on the server after the request returns: a large workspace takes longer to copy + * than the timeout of many proxies in front of Windmill, which would otherwise cut the request. + * `onCreationStarted` receives the id `waitForForkCreation` resumes the wait by, and `onStep` the + * part of the copy the server reports it is in. + */ +export async function createWorkspaceForkAndWait( + parentWorkspace: string, + fork: CreateWorkspaceFork, + opts: { onCreationStarted?: (creationId: string) => void; onStep?: (step: string) => void } = {} +): Promise { + const response = await WorkspaceService.createWorkspaceFork({ + workspace: parentWorkspace, + background: true, + requestBody: fork + }) + // A server without background forks ignores the flag and answers once the fork is created. + if (response.startsWith('Created forked workspace')) return + opts.onCreationStarted?.(response) + await waitForForkCreation(parentWorkspace, response, opts.onStep) +} + +/** Wait for a fork creation started in the background, by the id its request answered with. */ +export async function waitForForkCreation( + parentWorkspace: string, + creationId: string, + onStep?: (step: string) => void +): Promise { + let failingSince: number | undefined + while (true) { + let result: Awaited> | undefined + try { + result = await WorkspaceService.getForkCreationStatus({ + workspace: parentWorkspace, + creationId + }) + failingSince = undefined + } catch (e) { + // A failed poll says nothing about the fork: it can be lost to the same proxy the fork + // runs in the background to avoid, or reach a replica that has no status route yet. + failingSince ??= Date.now() + if (Date.now() - failingSince >= MAX_FAILING_POLLS_MS) throw e + } + if (result?.status === 'running' && result.step) onStep?.(result.step) + if (result?.status === 'completed') return + if (result?.status === 'failed') throw new Error(result.error ?? 'Unknown error') + await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS)) + } +}