all

backend/Cargo.lock

@@ -14435,6 +14435,7 @@ dependencies = [
  "itertools 0.14.0",
  "lazy_static",
  "prometheus",
+ "rand 0.9.0",
  "regex",
  "reqwest 0.12.15",
  "serde",

backend/windmill-common/src/bench.rs

@@ -235,7 +235,7 @@ pub async fn benchmark_init(benchmark_jobs: i32, db: &DB) {
                 )
                 .fetch_all(&mut *tx)
                 .await.unwrap_or_else(|_e| panic!("failed to insert noop jobs (1)"));
-                sqlx::query!("INSERT INTO v2_job_queue (id, workspace_id, scheduled_for, tag) SELECT unnest($1::uuid[]), $2, now(), $3", &uuids, "admins", "deno")
+                sqlx::query!("INSERT INTO v2_job_queue_partitioned (id, workspace_id, scheduled_for, tag, shard_id) SELECT unnest($1::uuid[]), $2, now(), $3, (floor(random() * 4 + 1))::int", &uuids, "admins", "deno")
                 .execute(&mut *tx)
                 .await.unwrap_or_else(|_e| panic!("failed to insert noop jobs (2)"));
                 sqlx::query!(
@@ -244,7 +244,7 @@ pub async fn benchmark_init(benchmark_jobs: i32, db: &DB) {
                 )
                 .execute(&mut *tx)
                 .await
-                .unwrap_or_else(|_e| panic!("failed to insert noop jobs (3)"));
+                .unwrap_or_else(|_e| panic!("failed to insert noop jobs (3): {_e:#}"));
             }
         }
         tx.commit().await.unwrap();

backend/windmill-common/src/worker.rs

@@ -240,12 +240,12 @@ fn format_pull_query(peek: String) -> String {
         "WITH peek AS (
             {}
         ), q AS NOT MATERIALIZED (
-            UPDATE v2_job_queue SET
+            UPDATE v2_job_queue_partitioned SET
                 running = true,
                 started_at = coalesce(started_at, now()),
                 suspend_until = null,
                 worker = $1
-            WHERE id = (SELECT id FROM peek)
+            WHERE id = (SELECT id FROM peek) AND shard_id = $2
             RETURNING
                 started_at, scheduled_for,
                 canceled_by, canceled_reason, worker
@@ -309,12 +309,12 @@ pub async fn store_suspended_pull_query(wc: &WorkerConfig) {
 pub fn make_pull_query(tags: &[String]) -> String {
     format_pull_query(format!(
         "SELECT id
-        FROM v2_job_queue
-        WHERE running = false AND tag IN ({}) AND scheduled_for <= now()
+        FROM v2_job_queue_partitioned
+        WHERE running = false AND tag IN ({}) AND scheduled_for <= now() AND shard_id = $2
         ORDER BY priority DESC NULLS LAST, scheduled_for
         FOR UPDATE SKIP LOCKED
         LIMIT 1",
-        tags.iter().map(|x| format!("'{x}'")).join(", ")
+        tags.iter().map(|x| format!("'{x}'")).join(", "),
     ))
 }
 

backend/windmill-queue/Cargo.toml

@@ -16,6 +16,7 @@ benchmark = ["windmill-common/benchmark"]
 prometheus = ["dep:prometheus"]
 
 [dependencies]
+rand.workspace = true
 windmill-audit.workspace = true
 windmill-common = { workspace = true, default-features = false }
 anyhow.workspace = true

backend/windmill-queue/src/jobs.rs

@@ -727,7 +727,7 @@ pub async fn add_completed_job<T: Serialize + Send + Sync + ValidableJson>(
                         WHEN $2::BOOL THEN 'success'::job_status
                         ELSE 'failure'::job_status END AS status,
                         q.worker
-                FROM v2_job_queue q LEFT JOIN v2_job_status USING (id) WHERE q.id = $1
+                FROM v2_job_queue_partitioned q LEFT JOIN v2_job_status USING (id) WHERE q.id = $1
             ON CONFLICT (id) DO UPDATE SET status = EXCLUDED.status, result = $3 RETURNING duration_ms AS \"duration_ms!\"",
             /* $1 */ queued_job.id,
             /* $2 */ success,
@@ -2572,7 +2572,7 @@ async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>(
             return Ok((None, false));
         }
 
-        let r = if suspend_first {
+        let r = if suspend_first && false {
             // tracing::info!("Pulling job with query: {}", query);
             sqlx::query_as::<_, PulledJob>(&query)
                 .bind(worker_name)
@@ -2602,6 +2602,7 @@ async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>(
 
                 let r = sqlx::query_as::<_, PulledJob>(query)
                     .bind(worker_name)
+                    .bind(((rand::random::<i8>() % 4) + 1) as i32)
                     .fetch_optional(db)
                     .await?;
 
@@ -3066,7 +3067,7 @@ pub async fn delete_job<'c>(
     }
 
     let job_removed =
-        sqlx::query_scalar!("DELETE FROM v2_job_queue WHERE id = $1 RETURNING 1", job_id,)
+        sqlx::query_scalar!("DELETE FROM v2_job_queue_partitioned WHERE id = $1 RETURNING 1", job_id)
             .fetch_optional(&mut *tx)
             .await;
 
@@ -4239,8 +4240,8 @@ pub async fn push<'c, 'd>(
     tracing::debug!("Pushing job {job_id} with tag {tag}, schedule_path {schedule_path:?}, script_path: {script_path:?}, email {email}, workspace_id {workspace_id}");
     let uuid = sqlx::query_scalar!(
         "INSERT INTO v2_job_queue
-            (workspace_id, id, running, scheduled_for, started_at, tag, priority)
-            VALUES ($1, $2, $3, COALESCE($4, now()), CASE WHEN $3 THEN now() END, $5, $6) \
+            (workspace_id, id, running, scheduled_for, started_at, tag, priority, shard_id)
+            VALUES ($1, $2, $3, COALESCE($4, now()), CASE WHEN $3 THEN now() END, $5, $6, $7) \
          RETURNING id AS \"id!\"",
         workspace_id,
         job_id,
@@ -4248,6 +4249,16 @@ pub async fn push<'c, 'd>(
         scheduled_for_o,
         tag,
         final_priority,
+        if std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_nanos()
+            % 2 == 0
+        {
+            1
+        } else {
+            2
+        }
     )
     .fetch_one(&mut *tx)
     .warn_after_seconds(1)

backend/windmill-worker/src/worker.rs

@@ -1002,10 +1002,8 @@ pub async fn run_worker(
 
     #[cfg(feature = "benchmark")]
     {
-        if i_worker == 1 {
-            if let Some(db) = conn.as_sql() {
-                benchmark_init(benchmark_jobs, db).await;
-            }
+        if let Some(db) = conn.as_sql() {
+            benchmark_init(benchmark_jobs, db).await;
         }
     }