diff --git a/.github/workflows/backend-test.yml b/.github/workflows/backend-test.yml index 466224b0e4..0f82ec30b9 100644 --- a/.github/workflows/backend-test.yml +++ b/.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: diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 90ab852aec..537f161c12 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -fc153796d006d08777b5b778d4f370a60495c3af \ No newline at end of file +0835279921261f03720513c0f7ffabc4df44db6f \ No newline at end of file diff --git a/backend/windmill-api/src/agent_workers_oss.rs b/backend/windmill-api/src/agent_workers_oss.rs index dabb2cb079..1459efef1d 100644 --- a/backend/windmill-api/src/agent_workers_oss.rs +++ b/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() } diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index 3ec3c1b176..630ba354ba 100644 --- a/backend/windmill-api/src/lib.rs +++ b/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:#}"); } } diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 77e2f212e0..189581446d 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -2132,6 +2132,37 @@ pub async fn get_mini_pulled_job<'c>( pub struct PulledJobResult { pub job: Option, pub suspended: bool, + pub missing_concurrency_key: bool, +} + +pub enum PulledJobResultToJobErr { + MissingConcurrencyKey(JobCompleted), +} + +impl PulledJobResult { + pub fn to_pulled_job(self) -> Result, 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"))] diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 5e407f3857..ad630d3577 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/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?; diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 75fc1df076..1587b95844 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/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, )