From 1b6e2556b7170b4ef9e9c460e8900149259ae4af Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Wed, 8 Apr 2026 15:37:53 +0000 Subject: [PATCH] perf: add worker-side batch job pull from DB Co-Authored-By: Claude Opus 4.6 (1M context) --- backend/windmill-common/src/worker.rs | 78 ++++++++++++++++++++++++++- backend/windmill-queue/src/jobs.rs | 26 ++++++++- backend/windmill-worker/src/worker.rs | 51 ++++++++++++++++++ 3 files changed, 151 insertions(+), 4 deletions(-) diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index d4479e06d1..1c435f63b5 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -237,6 +237,12 @@ lazy_static::lazy_static! { pub static ref WORKER_PULL_QUERIES: Arc>> = Arc::new(RwLock::new(vec![])); pub static ref WORKER_SUSPENDED_PULL_QUERY: Arc> = Arc::new(RwLock::new("".to_string())); + pub static ref WORKER_BATCH_PULL_QUERIES: Arc>> = Arc::new(RwLock::new(vec![])); + + pub static ref BATCH_PULL_SIZE: i32 = std::env::var("BATCH_PULL_SIZE") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0); pub static ref SMTP_CONFIG: Arc>> = Arc::new(RwLock::new(None)); @@ -522,16 +528,84 @@ pub fn make_pull_query(tags: &[String]) -> String { pub async fn store_pull_query(wc: &WorkerConfig) { let mut queries = vec![]; + let mut batch_queries = vec![]; for tags in wc.priority_tags_sorted.iter() { if tags.tags.len() == 0 { tracing::error!("Empty tags in priority tags, skipping"); continue; } - let query = make_pull_query(&tags.tags); - queries.push(query); + queries.push(make_pull_query(&tags.tags)); + batch_queries.push(make_batch_pull_query(&tags.tags)); } let mut l = WORKER_PULL_QUERIES.write().await; *l = queries; + drop(l); + let mut l = WORKER_BATCH_PULL_QUERIES.write().await; + *l = batch_queries; +} + +/// Build a batch pull query that claims up to $2 jobs at once. +/// Uses $1 for worker_name and $2 for batch_size (i32). +pub fn make_batch_pull_query(tags: &[String]) -> String { + format_batch_pull_query(format!( + "SELECT id + FROM v2_job_queue + WHERE running = false + AND tag IN ({}) AND scheduled_for <= now() + AND id NOT IN (SELECT id FROM v2_job WHERE same_worker = true) + ORDER BY priority DESC NULLS LAST, scheduled_for + FOR UPDATE SKIP LOCKED + LIMIT $2", + tags.iter().map(|x| format!("'{x}'")).join(", ") + )) +} + +fn format_batch_pull_query(peek: String) -> String { + format!( + "WITH peek AS ( + {} + ), q AS NOT MATERIALIZED ( + UPDATE v2_job_queue SET + running = true, + started_at = coalesce(started_at, now()), + suspend_until = null, + worker = $1 + WHERE id IN (SELECT id FROM peek) + RETURNING + id, started_at, scheduled_for, + canceled_by, canceled_reason, worker, cache_ignore_s3_path, runnable_settings_handle + ), r AS NOT MATERIALIZED ( + UPDATE v2_job_runtime SET + ping = now() + WHERE id IN (SELECT id FROM peek) + ), j AS NOT MATERIALIZED ( + SELECT + id, workspace_id, parent_job, created_by, created_at, runnable_id, + runnable_path, args, kind, trigger, trigger_kind, + permissioned_as, permissioned_as_email, script_lang, + flow_innermost_root_job, root_job, flow_step_id, + same_worker, pre_run_error, visible_to_owner, tag, concurrent_limit, + concurrency_time_window_s, timeout, cache_ttl, priority, raw_code, raw_lock, + raw_flow, script_entrypoint_override, preprocessed + FROM v2_job + WHERE id IN (SELECT id FROM peek) + ) SELECT 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 + FROM q + INNER JOIN j ON q.id = j.id + LEFT JOIN v2_job_status f ON f.id = j.id + LEFT JOIN job_perms p ON p.job_id = j.id + LEFT JOIN v2_job pj ON j.parent_job = pj.id", + peek + ) } lazy_static::lazy_static! { diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index a7cb9e66af..31d594454c 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -443,6 +443,25 @@ pub async fn append_logs( } } +/// Pull up to `batch_size` jobs at once using the given batch pull query. +/// The query must use $1 for worker_name and $2 for batch_size (i32). +pub async fn batch_pull( + db: &Pool, + worker_name: &str, + batch_query: &str, + batch_size: i32, +) -> windmill_common::error::Result> { + let jobs = sqlx::query_as::<_, PulledJob>(batch_query) + .bind(worker_name) + .bind(batch_size) + .fetch_all(db) + .await + .map_err(|e| { + windmill_common::error::Error::InternalErr(format!("batch pull error: {e:#}")) + })?; + Ok(jobs) +} + pub const PERIODIC_SCRIPT_TAG: &str = "periodic_bash_script"; pub const INIT_SCRIPT_TAG: &str = "init_script"; pub const INIT_SCRIPT_PATH_PREFIX: &str = "init_script_"; @@ -3638,8 +3657,11 @@ pub fn resolve_debounce_key<'b>( .join(":"), )); - tracing::debug!("Original debounce key (len={}): {}", original_debounce_key.len(), original_debounce_key); - + tracing::debug!( + "Original debounce key (len={}): {}", + original_debounce_key.len(), + original_debounce_key + ); // If debounce_key is not too long (< 255 chars), keep it as is, otherwise hash it. // On cloud, we prepend "{workspace_id}:" so we must reserve space for that prefix diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 67b3580a40..d7b8566e4b 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -2075,6 +2075,10 @@ pub async fn run_worker( let mut last_suspend_first = Instant::now(); let mut killed_but_draining_same_worker_jobs = false; + let batch_pull_size = *windmill_common::worker::BATCH_PULL_SIZE; + let mut batch_pull_buffer: std::collections::VecDeque = + std::collections::VecDeque::new(); + let mut killpill_rx2 = killpill_rx.resubscribe(); loop { @@ -2318,6 +2322,53 @@ pub async fn run_worker( tokio::time::sleep(Duration::from_millis(200)).await; continue; } + } else if batch_pull_size > 0 && !batch_pull_buffer.is_empty() { + // Serve from batch pull buffer + Ok(batch_pull_buffer + .pop_front() + .map(|job| NextJob::Sql { flow_runners: None, job })) + } else if batch_pull_size > 0 && matches!(&conn, Connection::Sql(_)) { + // Batch pull: try to refill buffer from DB + let db = conn.as_sql().unwrap(); + let queries = windmill_common::worker::WORKER_BATCH_PULL_QUERIES + .read() + .await; + if queries.is_empty() { + drop(queries); + // Queries not populated yet, fall through to normal pull below + None + } else { + let mut pulled = Vec::new(); + for query in queries.iter() { + match windmill_queue::batch_pull(db, &worker_name, query, batch_pull_size) + .await + { + Ok(jobs) if !jobs.is_empty() => { + pulled = jobs; + break; + } + Ok(_) => {} + Err(e) => { + tracing::error!(worker = %worker_name, "batch pull error: {e:#}"); + } + } + } + drop(queries); + if pulled.is_empty() { + Some(Ok(None)) + } else { + let mut iter = pulled.into_iter(); + let first = iter.next(); + for job in iter { + batch_pull_buffer.push_back(job); + } + Some(Ok(first.map(|job| NextJob::Sql { flow_runners: None, job }))) + } + } + .unwrap_or_else(|| { + // Fall through: queries not ready yet + Ok(None) + }) } else { match &conn { Connection::Sql(db) => {