diff --git a/backend/.sqlx/query-2719f910142b32476a16025bb9836b0cab019ba0a436b330ea3a53fba4725f73.json b/backend/.sqlx/query-a3ccf362b4f6df400b3c7a084795dbf541eb14c5c374656ffb96da7283a2a6f1.json similarity index 57% rename from backend/.sqlx/query-2719f910142b32476a16025bb9836b0cab019ba0a436b330ea3a53fba4725f73.json rename to backend/.sqlx/query-a3ccf362b4f6df400b3c7a084795dbf541eb14c5c374656ffb96da7283a2a6f1.json index 89895ed0bd..bc5e2351f9 100644 --- a/backend/.sqlx/query-2719f910142b32476a16025bb9836b0cab019ba0a436b330ea3a53fba4725f73.json +++ b/backend/.sqlx/query-a3ccf362b4f6df400b3c7a084795dbf541eb14c5c374656ffb96da7283a2a6f1.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT concurrency_key FROM script WHERE hash = $1", + "query": "SELECT concurrency_key FROM script WHERE hash = $1 AND workspace_id = $2", "describe": { "columns": [ { @@ -11,12 +11,13 @@ ], "parameters": { "Left": [ - "Int8" + "Int8", + "Text" ] }, "nullable": [ true ] }, - "hash": "2719f910142b32476a16025bb9836b0cab019ba0a436b330ea3a53fba4725f73" + "hash": "a3ccf362b4f6df400b3c7a084795dbf541eb14c5c374656ffb96da7283a2a6f1" } diff --git a/backend/.sqlx/query-a875cb56485b812e9d4739afd0915067f7e5abe0ca0adf264b792fccf21e005b.json b/backend/.sqlx/query-a875cb56485b812e9d4739afd0915067f7e5abe0ca0adf264b792fccf21e005b.json new file mode 100644 index 0000000000..e264db5ef2 --- /dev/null +++ b/backend/.sqlx/query-a875cb56485b812e9d4739afd0915067f7e5abe0ca0adf264b792fccf21e005b.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT value->>'concurrency_key' FROM flow WHERE path = $1 AND workspace_id = $2", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "?column?", + "type_info": "Text" + } + ], + "parameters": { + "Left": [ + "Text", + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "a875cb56485b812e9d4739afd0915067f7e5abe0ca0adf264b792fccf21e005b" +} diff --git a/backend/migrations/20240409181630_add_concurrency_key_to_flow.down.sql b/backend/migrations/20240409181630_add_concurrency_key_to_flow.down.sql new file mode 100644 index 0000000000..050f94cee7 --- /dev/null +++ b/backend/migrations/20240409181630_add_concurrency_key_to_flow.down.sql @@ -0,0 +1,2 @@ +-- Add down migration script here +ALTER TABLE flow DROP COLUMN concurrency_key; diff --git a/backend/migrations/20240409181630_add_concurrency_key_to_flow.up.sql b/backend/migrations/20240409181630_add_concurrency_key_to_flow.up.sql new file mode 100644 index 0000000000..3ad350a983 --- /dev/null +++ b/backend/migrations/20240409181630_add_concurrency_key_to_flow.up.sql @@ -0,0 +1,2 @@ +-- Add up migration script here +ALTER TABLE flow ADD COLUMN concurrency_key VARCHAR(255); diff --git a/backend/windmill-api/src/flows.rs b/backend/windmill-api/src/flows.rs index 398a4faa7e..05f732358c 100644 --- a/backend/windmill-api/src/flows.rs +++ b/backend/windmill-api/src/flows.rs @@ -334,7 +334,7 @@ async fn create_flow( nf.draft_only, nf.tag, nf.dedicated_worker, - nf.visible_to_runner_only, + nf.visible_to_runner_only.unwrap_or(false), ) .execute(&mut tx) .await?; @@ -502,7 +502,7 @@ async fn update_flow( w_id, nf.tag, nf.dedicated_worker, - nf.visible_to_runner_only, + nf.visible_to_runner_only.unwrap_or(false), ) .execute(&mut tx) .await?; @@ -985,6 +985,7 @@ mod tests { cache_ttl: None, priority: None, early_return: None, + concurrency_key: None, }; let expect = serde_json::json!({ "modules": [ diff --git a/backend/windmill-common/src/flows.rs b/backend/windmill-common/src/flows.rs index 9d1d5caca6..ff4d023d7d 100644 --- a/backend/windmill-common/src/flows.rs +++ b/backend/windmill-common/src/flows.rs @@ -105,6 +105,8 @@ pub struct FlowValue { #[serde(skip_serializing_if = "Option::is_none")] // Priority at the flow level pub priority: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub concurrency_key: Option, } #[derive(Deserialize, Serialize, Debug, Clone)] diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index ee4ea1f694..21305fa52e 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -1657,12 +1657,12 @@ pub async fn pull( + Duration::try_seconds(i64::from(job_custom_concurrency_time_window_s)) .unwrap_or_default(); tracing::info!("Job '{}' from path '{}' with concurrency key '{}' has reached its concurrency limit of {} jobs run in the last {} seconds. This job will be re-queued for next execution at {}", - job_uuid, job_script_path, job_custom_concurrent_limit, job_concurrency_key, job_custom_concurrency_time_window_s, estimated_next_schedule_timestamp); + job_uuid, job_script_path, job_concurrency_key, job_custom_concurrent_limit, job_custom_concurrency_time_window_s, estimated_next_schedule_timestamp); let job_log_event = format!( "\nRe-scheduled job to {estimated_next_schedule_timestamp} due to concurrency limits with key {job_concurrency_key} and limit {job_custom_concurrent_limit} in the last {job_custom_concurrency_time_window_s} seconds", ); - let _ = append_logs(job_uuid, pulled_job.workspace_id, job_log_event, db); + let _ = append_logs(job_uuid, pulled_job.workspace_id, job_log_event, db).await; if rsmq.is_some() { // if let Some(ref mut rsmq) = tx.rsmq { // if using redis, only one message at a time can be poped from the queue. Process only this message and move to the next elligible job @@ -1884,50 +1884,63 @@ async fn pull_single_job_and_mark_as_running_no_concurrency_limit< } async fn concurrency_key(db: &Pool, queued_job: &QueuedJob) -> String { - if queued_job.is_flow() { - // custom concurrency keys are not yet supported for flows - queued_job.full_path_with_workspace() - } else { - let concurrency_key = sqlx::query_scalar!( - "SELECT concurrency_key FROM script WHERE hash = $1", - queued_job.script_hash.unwrap_or(ScriptHash(0)).0 + let r = if queued_job.is_flow() { + sqlx::query_scalar!( + "SELECT value->>'concurrency_key' FROM flow WHERE path = $1 AND workspace_id = $2", + queued_job.script_path, + queued_job.workspace_id ) .fetch_one(db) - .await; - match concurrency_key { - Ok(Some(custom_concurrency_key)) => { - let workspaced = - custom_concurrency_key.replace("$workspace", queued_job.workspace_id.as_str()); - if RE_ARG_TAG.is_match(&workspaced) { - let mut interpolated = workspaced.clone(); - for cap in RE_ARG_TAG.captures_iter(&workspaced) { - let arg_name = cap.get(1).unwrap().as_str(); - let arg_value = match queued_job.args.as_ref() { - Some(Json(args_map_json)) => match args_map_json.get(arg_name) { - Some(arg_value_raw) => { - serde_json::to_string(arg_value_raw).unwrap_or_default() - } - None => "".to_string(), - }, + .await + } else { + sqlx::query_scalar!( + "SELECT concurrency_key FROM script WHERE hash = $1 AND workspace_id = $2", + queued_job.script_hash.unwrap_or(ScriptHash(0)).0, + queued_job.workspace_id + ) + .fetch_one(db) + .await + }; + process_custom_concurrency_key(queued_job, r).await +} + +async fn process_custom_concurrency_key( + queued_job: &QueuedJob, + concurrency_key: Result, sqlx::Error>, +) -> String { + match concurrency_key { + Ok(Some(custom_concurrency_key)) => { + let workspaced = + custom_concurrency_key.replace("$workspace", queued_job.workspace_id.as_str()); + if RE_ARG_TAG.is_match(&workspaced) { + let mut interpolated = workspaced.clone(); + for cap in RE_ARG_TAG.captures_iter(&workspaced) { + let arg_name = cap.get(1).unwrap().as_str(); + let arg_value = match queued_job.args.as_ref() { + Some(Json(args_map_json)) => match args_map_json.get(arg_name) { + Some(arg_value_raw) => { + serde_json::to_string(arg_value_raw).unwrap_or_default() + } None => "".to_string(), - }; - interpolated = interpolated - .replace(format!("$args[{}]", arg_name).as_str(), arg_value.as_str()); - } - interpolated - } else { - workspaced + }, + None => "".to_string(), + }; + interpolated = interpolated + .replace(format!("$args[{}]", arg_name).as_str(), arg_value.as_str()); } + interpolated + } else { + workspaced } - Ok(None) => queued_job.full_path_with_workspace(), - _ => { - tracing::warn!( - "Unable to retrieve concurrency key for script {:?} | {:?}", - queued_job.script_path, - queued_job.script_hash - ); - queued_job.full_path_with_workspace() - } + } + Ok(None) => queued_job.full_path_with_workspace(), + _ => { + tracing::warn!( + "Unable to retrieve concurrency key for script {:?} | {:?}", + queued_job.script_path, + queued_job.script_hash + ); + queued_job.full_path_with_workspace() } } } @@ -3000,6 +3013,7 @@ pub async fn push<'c, T: Serialize + Send + Sync, R: rsmq_async::RsmqConnection skip_expr: None, cache_ttl: cache_ttl.map(|val| val as u32), early_return: None, + concurrency_key: None, priority: priority, }; ( diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 8d17ae0819..f6728e9742 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -2607,6 +2607,7 @@ async fn compute_next_flow_transform( cache_ttl: None, priority: None, early_return: None, + concurrency_key: None, }, path: Some(format!("{}/forloop", flow_job.script_path())), restarted_from: None, @@ -2705,6 +2706,7 @@ async fn compute_next_flow_transform( cache_ttl: None, priority: None, early_return: None, + concurrency_key: None, }, path: Some(format!( "{}/branchone-{}", @@ -2756,6 +2758,7 @@ async fn compute_next_flow_transform( cache_ttl: None, priority: None, early_return: None, + concurrency_key: None, }, path: Some(format!( "{}/branchall-{}", @@ -2823,6 +2826,7 @@ async fn compute_next_flow_transform( cache_ttl: None, priority: None, early_return: None, + concurrency_key: None, }, path: Some(format!( "{}/branchall-{}", @@ -2895,6 +2899,7 @@ async fn next_loop_iteration( cache_ttl: None, priority: None, early_return: None, + concurrency_key: None, }, path: inner_path, restarted_from: None, diff --git a/frontend/src/lib/components/ScriptBuilder.svelte b/frontend/src/lib/components/ScriptBuilder.svelte index 9fd3016e8a..3bef9fe01b 100644 --- a/frontend/src/lib/components/ScriptBuilder.svelte +++ b/frontend/src/lib/components/ScriptBuilder.svelte @@ -662,18 +662,21 @@ bind:seconds={script.concurrency_time_window_s} /> - -