fix(backend): concurrency limits preprocessor (#6727)

* fix(backend): concurrency limits preprocessor

* refactor

* fix concurrency on agent workers

* nit

* nit

* update ref

* fix ci

* chore: update ee ref

.github/workflows/backend-test.yml

@@ -54,7 +54,7 @@ jobs:
           toolchain: 1.90.0
       - name: Read EE repo commit hash
         run: |
-          echo "ee_repo_ref=$(cat ./backend/ee-repo-ref.txt)" >> "$GITHUB_ENV"
+          echo "ee_repo_ref=$(cat ./ee-repo-ref.txt)" >> "$GITHUB_ENV"
 
       - uses: actions/checkout@v4
         with:

backend/ee-repo-ref.txt

@@ -1 +1 @@
-fc153796d006d08777b5b778d4f370a60495c3af
+0835279921261f03720513c0f7ffabc4df44db6f

backend/windmill-api/src/agent_workers_oss.rs

@@ -17,7 +17,7 @@ use crate::db::DB;
 use axum::Router;
 
 #[cfg(not(feature = "private"))]
-pub fn global_service() -> Router {
+pub fn global_service(_job_completed_tx: windmill_worker::JobCompletedSender) -> Router {
     Router::new()
 }
 

backend/windmill-api/src/lib.rs

@@ -411,7 +411,7 @@ pub async fn run_server(
     };
 
     #[cfg(feature = "agent_worker_server")]
-    let (agent_workers_router, agent_workers_bg_processor, agent_workers_killpill_tx) =
+    let (agent_workers_router, agent_workers_bg_processor, agent_workers_job_completed_tx) =
         if server_mode {
             agent_workers_oss::workspaced_service(db.clone(), _base_internal_url.clone())
         } else {
@@ -535,7 +535,14 @@ pub async fn run_server(
                 .nest("/agent_workers", {
                     #[cfg(feature = "agent_worker_server")]
                     {
-                        agent_workers_oss::global_service().layer(Extension(agent_cache.clone()))
+                        if let Some(agent_workers_job_completed_tx) =
+                            agent_workers_job_completed_tx.clone()
+                        {
+                            agent_workers_oss::global_service(agent_workers_job_completed_tx)
+                                .layer(Extension(agent_cache.clone()))
+                        } else {
+                            Router::new()
+                        }
                     }
                     #[cfg(not(feature = "agent_worker_server"))]
                     {
@@ -705,8 +712,8 @@ pub async fn run_server(
     let server = server.with_graceful_shutdown(async move {
         killpill_rx.recv().await.ok();
         #[cfg(feature = "agent_worker_server")]
-        if let Some(agent_workers_killpill_tx) = agent_workers_killpill_tx {
-            if let Err(e) = agent_workers_killpill_tx.kill().await {
+        if let Some(agent_workers_job_completed_tx) = agent_workers_job_completed_tx {
+            if let Err(e) = agent_workers_job_completed_tx.kill().await {
                 tracing::error!("Error killing agent workers: {e:#}");
             }
         }

backend/windmill-queue/src/jobs.rs

@@ -2132,6 +2132,37 @@ pub async fn get_mini_pulled_job<'c>(
 pub struct PulledJobResult {
     pub job: Option<PulledJob>,
     pub suspended: bool,
+    pub missing_concurrency_key: bool,
+}
+
+pub enum PulledJobResultToJobErr {
+    MissingConcurrencyKey(JobCompleted),
+}
+
+impl PulledJobResult {
+    pub fn to_pulled_job(self) -> Result<Option<PulledJob>, PulledJobResultToJobErr> {
+        match self {
+            PulledJobResult { job: Some(job), missing_concurrency_key: true, .. } => Err(
+                PulledJobResultToJobErr::MissingConcurrencyKey(JobCompleted {
+                    preprocessed_args: None,
+                    job: Arc::new(job.job),
+                    success: false,
+                    result: Arc::new(windmill_common::worker::to_raw_value(&json!({
+                        "name": "InternalErr",
+                        "message": "The job has a concurrency limit but concurrency key couldn't be found. This is an unexpected behavior that should never happen. Please report this to support."}
+                    ))),
+                    result_columns: None,
+                    mem_peak: 0,
+                    cached_res_path: None,
+                    token: "".to_string(),
+                    canceled_by: None,
+                    duration: None,
+                    has_stream: Some(false),
+                }),
+            ),
+            PulledJobResult { job, .. } => Ok(job),
+        }
+    }
 }
 
 pub async fn pull(
@@ -2150,7 +2181,11 @@ pub async fn pull(
         }
         if pull_loop_count > 1000 {
             tracing::error!("Pull job loop count exceeded 1000, breaking");
-            return Ok(PulledJobResult { job: None, suspended: false });
+            return Ok(PulledJobResult {
+                job: None,
+                suspended: false,
+                missing_concurrency_key: false,
+            });
         }
         if let Some((query_suspended, query_no_suspend)) = query_o {
             let njob = {
@@ -2162,16 +2197,43 @@ pub async fn pull(
                         .fetch_optional(db)
                         .await?
                 };
-                if let Some(job) = job {
-                    PulledJobResult { job: Some(job), suspended: true }
+
+                let (job, suspended) = if let Some(job) = job {
+                    (Some(job), true)
                 } else {
                     let job = sqlx::query_as::<_, PulledJob>(query_no_suspend)
                         .bind(worker_name)
                         .fetch_optional(db)
                         .await?;
-                    PulledJobResult { job, suspended: false }
-                }
-            };
+                    (job, false)
+                };
+
+                #[cfg(all(feature = "enterprise", feature = "private"))]
+                let pulled_job_result = match job {
+                    Some(job) if job.concurrent_limit.is_some() => {
+                        let job = crate::jobs_ee::apply_concurrency_limit(
+                            db,
+                            pull_loop_count,
+                            suspended,
+                            job,
+                        )
+                        .await?;
+                        job.unwrap_or(PulledJobResult {
+                            job: None,
+                            suspended,
+                            missing_concurrency_key: false,
+                        })
+                    }
+                    _ => PulledJobResult { job, suspended, missing_concurrency_key: false },
+                };
+
+                #[cfg(not(all(feature = "enterprise", feature = "private")))]
+                let pulled_job_result =
+                    PulledJobResult { job, suspended, missing_concurrency_key: false };
+
+                Ok::<_, Error>(pulled_job_result)
+            }?;
+
             if let Some(job) = njob.job.as_ref() {
                 if job.is_flow() || job.is_dependency() {
                     let per_workspace = per_workspace_tag(&job.workspace_id).await;
@@ -2207,7 +2269,7 @@ pub async fn pull(
         .await?;
 
         let Some(job) = job else {
-            return Ok(PulledJobResult { job: None, suspended });
+            return Ok(PulledJobResult { job: None, suspended, missing_concurrency_key: false });
         };
 
         let has_concurent_limit = job.concurrent_limit.is_some();
@@ -2230,7 +2292,11 @@ pub async fn pull(
             if METRICS_ENABLED.load(std::sync::atomic::Ordering::Relaxed) {
                 QUEUE_PULL_COUNT.inc();
             }
-            return Ok(PulledJobResult { job: Some(pulled_job), suspended });
+            return Ok(PulledJobResult {
+                job: Some(pulled_job),
+                suspended,
+                missing_concurrency_key: false,
+            });
         }
 
         #[cfg(all(feature = "enterprise", feature = "private"))]

backend/windmill-worker/src/worker.rs

@@ -57,6 +57,7 @@ use std::{
     time::Duration,
 };
 use windmill_parser::MainArgSignature;
+use windmill_queue::PulledJobResultToJobErr;
 
 use uuid::Uuid;
 
@@ -821,7 +822,18 @@ pub fn start_interactive_worker_shell(
                         )
                         .await;
 
-                        job.map(|x| x.job.map(NextJob::Sql))
+                        match job {
+                            Ok(j) => match j.to_pulled_job() {
+                                Ok(j) => Ok(j.map(NextJob::Sql)),
+                                Err(PulledJobResultToJobErr::MissingConcurrencyKey(jc)) => {
+                                    if let Err(err) = job_completed_tx.send_job(jc, true).await {
+                                        tracing::error!("An error occurred while sending job completed (missing concurrency key): {:#?}", err)
+                                    }
+                                    Ok(None)
+                                }
+                            },
+                            Err(err) => Err(err),
+                        }
                     }
                     Connection::Http(client) => {
                         crate::agent_workers::pull_job(&client, None, Some(true))
@@ -1631,7 +1643,18 @@ pub async fn run_worker(
                                 }
                             }
                         }
-                        job.map(|x| x.job.map(NextJob::Sql))
+                        match job {
+                            Ok(pulled_job_result) => match pulled_job_result.to_pulled_job() {
+                                Ok(j) => Ok(j.map(NextJob::Sql)),
+                                Err(PulledJobResultToJobErr::MissingConcurrencyKey(jc)) => {
+                                    if let Err(err) = job_completed_tx.send_job(jc, true).await {
+                                        tracing::error!("An error occurred while sending job completed (missing concurrency key): {:#?}", err)
+                                    }
+                                    Ok(None)
+                                }
+                            },
+                            Err(err) => Err(err),
+                        }
                     }
                     Connection::Http(client) => crate::agent_workers::pull_job(&client, None, None)
                         .await
@@ -2370,14 +2393,16 @@ pub async fn handle_queued_job(
             | JobKind::Flow
             | JobKind::FlowDependencies,
             x,
-        ) => if x.map(|x| x.0).is_none_or(|x| is_special_codebase_hash(x)) {
-            Some(
-                cache::job::fetch_preview(conn, &job.id, raw_lock, raw_code, raw_flow.clone())
-                    .await?,
-            ) 
-        } else {
+        ) => {
+            if x.map(|x| x.0).is_none_or(|x| is_special_codebase_hash(x)) {
+                Some(
+                    cache::job::fetch_preview(conn, &job.id, raw_lock, raw_code, raw_flow.clone())
+                        .await?,
+                )
+            } else {
                 None
-        },
+            }
+        }
         _ => None,
     };
 
@@ -2869,7 +2894,9 @@ async fn handle_code_execution_job(
         ScriptMetadata { language, envs, codebase, schema_validator, schema },
     ) = match job.kind {
         JobKind::Preview => {
-            let codebase = job.runnable_id.and_then(|x| hash_to_codebase_id(&job.id.to_string(), x.0));
+            let codebase = job
+                .runnable_id
+                .and_then(|x| hash_to_codebase_id(&job.id.to_string(), x.0));
             if codebase.is_none() && job.runnable_id.is_some() {
                 (arc_data, arc_metadata) =
                     cache::script::fetch(conn, job.runnable_id.unwrap()).await?;

backend/windmill-worker/src/worker_flow.rs

@@ -1242,14 +1242,14 @@ pub async fn update_flow_status_after_job_completion_internal(
                 if let Some(t) = tag {
                     tag = Some(interpolate_args(t, &args, &flow_job.workspace_id));
                 }
-            } else if let Some(ck) = concurrency_key {
+            } else if concurrent_limit.is_some() {
                 let mut tx = db.begin().await?;
                 insert_concurrency_key(
                     &flow_job.workspace_id,
                     &PushArgs::from(&HashMap::new()),
                     &flow_job.runnable_path,
                     JobKind::Flow,
-                    Some(ck),
+                    concurrency_key,
                     &mut tx,
                     flow,
                 )