feat: flow concurrency limits support custom concurrency key

This commit is contained in:
Ruben Fiszel
2024-04-09 22:03:24 +02:00
parent 32cd206556
commit a55aad3003
12 changed files with 125 additions and 56 deletions
@@ -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"
}
@@ -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"
}
@@ -0,0 +1,2 @@
-- Add down migration script here
ALTER TABLE flow DROP COLUMN concurrency_key;
@@ -0,0 +1,2 @@
-- Add up migration script here
ALTER TABLE flow ADD COLUMN concurrency_key VARCHAR(255);
+3 -2
View File
@@ -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": [
+2
View File
@@ -105,6 +105,8 @@ pub struct FlowValue {
#[serde(skip_serializing_if = "Option::is_none")]
// Priority at the flow level
pub priority: Option<i16>,
#[serde(skip_serializing_if = "Option::is_none")]
pub concurrency_key: Option<String>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
+55 -41
View File
@@ -1657,12 +1657,12 @@ pub async fn pull<R: rsmq_async::RsmqConnection + Send + Clone>(
+ 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<Postgres>, 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<Option<String>, 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,
};
(
@@ -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,
@@ -662,18 +662,21 @@
bind:seconds={script.concurrency_time_window_s}
/>
</Label>
<Label label="Custom concurrency key">
<Label label="Custom concurrency key (optional)">
<svelte:fragment slot="header">
<Tooltip>
Concurrency keys are global, you can have them be workspace specific using
the variable `$workspace`. You can also use an argument's value using
`$args[name_of_arg]`</Tooltip
>
</svelte:fragment>
<input
disabled={!$enterpriseLicense}
type="text"
autofocus
bind:value={script.concurrency_key}
placeholder={`$workspace/script/${script.path}-$args[foo]`}
/>
<Tooltip
>Concurrency keys are global, you can have them be workspace specific using
the variable `$workspace`. You can also use an argument's value using
`$args[name_of_arg]`</Tooltip
>
</Label>
</div>
</Section>
@@ -510,10 +510,21 @@
bind:seconds={$flowStore.value.concurrency_time_window_s}
/>
</Label>
<Label label="Custom concurrency key">
<div class="text-tertiary text-xs"
>Custom concurrency keys can only be set as the setting of a workspace script</div
>
<Label label="Custom concurrency key (optional)">
<svelte:fragment slot="header">
<Tooltip>
Concurrency keys are global, you can have them be workspace specific using the
variable `$workspace`. You can also use an argument's value using
`$args[name_of_arg]`</Tooltip
>
</svelte:fragment>
<input
type="text"
autofocus
disabled={!$enterpriseLicense}
bind:value={$flowStore.value.concurrency_key}
placeholder={`$workspace/script/${$pathStore}-$args[foo]`}
/>
</Label>
</div>
</Section>
@@ -90,6 +90,9 @@ export function cleanInputs(flow: OpenFlow | any): OpenFlow & {
})
}
})
if (newFlow.value.concurrency_key == '') {
newFlow.value.concurrency_key = undefined
}
return newFlow
}
+2
View File
@@ -48,6 +48,8 @@ components:
type: boolean
concurrent_limit:
type: number
concurrency_key:
type: string
concurrency_time_window_s:
type: number
skip_expr: