fix: improve flow cancellation (#4013)

* fix: improve flow cancellation

* fix: sqlx

* fix: cancel shortcut when no parent and not running
This commit is contained in:
HugoCasa
2024-07-03 14:50:52 +02:00
committed by GitHub
parent 957a0a65b9
commit 5e23cc26d0
7 changed files with 140 additions and 189 deletions
@@ -0,0 +1,22 @@
{
"db_name": "PostgreSQL",
"query": "SELECT id FROM queue WHERE id = ANY($1) AND schedule_path IS NULL",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "id",
"type_info": "Uuid"
}
],
"parameters": {
"Left": [
"UuidArray"
]
},
"nullable": [
false
]
},
"hash": "273d275be89516b135d7846d179c84ba0d684a620e9e5c0c87f82bdc9cd57dbe"
}
@@ -1,34 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "SELECT id, is_flow_step, running FROM queue WHERE id = ANY($1) AND schedule_path IS NULL",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "id",
"type_info": "Uuid"
},
{
"ordinal": 1,
"name": "is_flow_step",
"type_info": "Bool"
},
{
"ordinal": 2,
"name": "running",
"type_info": "Bool"
}
],
"parameters": {
"Left": [
"UuidArray"
]
},
"nullable": [
false,
true,
false
]
},
"hash": "4c529e29a0eb084a92f4e952506d701592bc707fd8a8cb3311bf146be6078f26"
}
@@ -1,15 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "UPDATE queue SET canceled = true, canceled_reason = 'http connection broke', canceled_by = queue.created_by WHERE id = $1 AND workspace_id = $2",
"describe": {
"columns": [],
"parameters": {
"Left": [
"Uuid",
"Text"
]
},
"nullable": []
},
"hash": "a96dd57b127a1adbdca13867a76d7fd9f9a9c76b02bcebd2830d928821f3abc0"
}
@@ -1,23 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "UPDATE queue SET canceled = true, canceled_by = $1, scheduled_for = now(), suspend = 0 WHERE id = $2 RETURNING 1 as one",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "one",
"type_info": "Int4"
}
],
"parameters": {
"Left": [
"Varchar",
"Uuid"
]
},
"nullable": [
null
]
},
"hash": "d4fb94ee8198592c24e85d29078feb5220ab367e339510ee0e83bb7b5abfd184"
}
+10 -8
View File
@@ -1123,13 +1123,15 @@ async fn handle_zombie_flows(
let id = flow.id.clone();
let last_ping = flow.last_ping.clone();
let now = now_from_db(db).await?;
cancel_zombie_flow_job(
db,
flow,
&rsmq,
format!("Flow {} cancelled as it was hanging in between 2 steps. Last ping: {last_ping:?} (now: {now})", id),
)
.await?;
let reason = format!(
"{} was hanging in between 2 steps. Last ping: {last_ping:?} (now: {now})",
if flow.is_flow_step && flow.parent_job.is_some() {
format!("Flow was cancelled because subflow {id}")
} else {
format!("Flow {id} was cancelled because it")
}
);
cancel_zombie_flow_job(db, flow, &rsmq, reason).await?;
}
}
@@ -1189,7 +1191,7 @@ async fn cancel_zombie_flow_job(
tx,
db,
rsmq.clone(),
false,
true,
false,
)
.await?;
+77 -70
View File
@@ -76,8 +76,8 @@ use windmill_common::{METRICS_DEBUG_ENABLED, METRICS_ENABLED};
use windmill_common::{get_latest_deployed_hash_for_path, BASE_URL};
use windmill_queue::{
add_completed_job_error, append_logs, get_queued_job, get_result_by_id_from_running_flow,
job_is_complete, push, CanceledBy, DecodeQueries, PushArgs, PushIsolationLevel,
cancel_job, get_queued_job, get_result_by_id_from_running_flow, job_is_complete, push,
DecodeQueries, PushArgs, PushIsolationLevel,
};
#[cfg(feature = "prometheus")]
@@ -1290,66 +1290,49 @@ async fn list_queue_jobs(
Ok(Json(jobs))
}
#[derive(Deserialize, FromRow)]
struct JobToCancel {
id: Uuid,
is_flow_step: Option<bool>,
running: bool,
}
async fn cancel_jobs(
jobs: Vec<JobToCancel>,
jobs: Vec<Uuid>,
db: &DB,
username: &str,
w_id: &str,
rsmq: Option<rsmq_async::MultiplexedRsmq>,
) -> error::JsonResult<Vec<Uuid>> {
let mut uuids = vec![];
for j in jobs.iter() {
let r = sqlx::query!(
"UPDATE queue SET canceled = true, canceled_by = $1, scheduled_for = now(), suspend = 0 WHERE id = $2 RETURNING 1 as one",
username,
j.id,
)
.fetch_optional(db)
.await;
if r.as_ref().is_ok_and(|x| x.is_some()) {
uuids.push(j.id);
if !j.running && !j.is_flow_step.unwrap_or(false) {
let e = serde_json::json!({"message": format!("Job canceled: cancel_all by {username}"), "name": "Canceled", "reason": "cancel_all", "canceler": username});
let job_running = get_queued_job(&j.id, &w_id, &db).await?;
if let Some(job_running) = job_running {
append_logs(
&j.id,
w_id,
format!("canceled by {username}: cancel_all"),
db.clone(),
)
.await;
let add_job = add_completed_job_error(
&db,
&job_running,
job_running.mem_peak.unwrap_or(0),
Some(CanceledBy {
username: Some(username.to_string()),
reason: Some("cancel_all".to_string()),
}),
e,
rsmq.clone(),
"server",
true,
)
.await;
if let Err(e) = add_job {
tracing::error!("Failed to add canceled job: {}", e);
}
for job_id in jobs.into_iter() {
let rsmq = rsmq.clone();
match tokio::time::timeout(tokio::time::Duration::from_secs(5), async move {
let tx = db.begin().await?;
let (tx, _) = windmill_queue::cancel_job(
username,
None,
job_id.clone(),
w_id,
tx,
db,
rsmq,
false,
false,
)
.await?;
tx.commit().await?;
Ok::<_, anyhow::Error>(())
})
.await
{
Ok(result) => match result {
Ok(_) => {
uuids.push(job_id);
}
Err(e) => {
tracing::error!("Failed to cancel job {:?}: {:?}", job_id, e);
}
},
Err(_) => {
tracing::error!(
"Timeout while trying to cancel job {:?} after 5 seconds",
job_id
);
}
} else {
tracing::error!("Failed to cancel job: {:?} {:?}", j.id, r.err());
}
}
@@ -1368,9 +1351,8 @@ async fn cancel_selection(
require_admin(authed.is_admin, &authed.username)?;
let mut tx = user_db.begin(&authed).await?;
let jobs_to_cancel = sqlx::query_as!(
JobToCancel,
"SELECT id, is_flow_step, running FROM queue WHERE id = ANY($1) AND schedule_path IS NULL",
let jobs_to_cancel = sqlx::query_scalar!(
"SELECT id FROM queue WHERE id = ANY($1) AND schedule_path IS NULL",
&jobs
)
.fetch_all(&mut *tx)
@@ -2935,6 +2917,7 @@ struct Guard {
id: Uuid,
w_id: String,
db: DB,
username: String,
}
impl Drop for Guard {
@@ -2943,16 +2926,33 @@ impl Drop for Guard {
let id = self.id;
let w_id = self.w_id.clone();
let db = self.db.clone();
let username = self.username.clone();
tracing::info!("http connection broke, marking job {id} as canceled");
tokio::spawn(async move {
let _ = sqlx::query!(
"UPDATE queue SET canceled = true, canceled_reason = 'http connection broke', canceled_by = queue.created_by WHERE id = $1 AND workspace_id = $2",
id,
w_id
)
.execute(&db)
.await;
let cancel_f = async {
let tx = db.begin().await?;
let (tx, _) = cancel_job(
&username,
Some("http connection broke".to_string()),
id,
&w_id,
tx,
&db,
None,
false,
false,
)
.await?;
tx.commit().await?;
Ok::<_, anyhow::Error>(())
};
if let Err(e) = cancel_f.await {
tracing::error!(
"Error marking job as canceled after http connection broke: {e}"
);
}
});
}
}
@@ -2969,6 +2969,7 @@ async fn run_wait_result(
uuid: Uuid,
w_id: String,
node_id_for_empty_return: Option<String>,
username: &str,
) -> error::Result<Response> {
let mut result = None;
let timeout = SERVER_CONFIG.read().await.timeout_wait_result.clone();
@@ -2978,7 +2979,13 @@ async fn run_wait_result(
(timeout * 1000) as u64
};
let mut g = Guard { done: false, id: uuid, w_id: w_id.clone(), db: db.clone() };
let mut g = Guard {
done: false,
id: uuid,
w_id: w_id.clone(),
db: db.clone(),
username: username.to_string(),
};
let fast_poll_duration = *WAIT_RESULT_FAST_POLL_DURATION_SECS as u64 * 1000;
let mut accumulated_delay = 0 as u64;
@@ -3219,7 +3226,7 @@ pub async fn run_wait_result_job_by_path_get(
.await?;
tx.commit().await?;
let wait_result = run_wait_result(&db, uuid, w_id, None).await;
let wait_result = run_wait_result(&db, uuid, w_id, None, &authed.username).await;
if delete_after_use.unwrap_or(false) {
delete_job_metadata_after_use(&db, uuid).await?;
}
@@ -3339,7 +3346,7 @@ async fn run_wait_result_script_by_path_internal(
.await?;
tx.commit().await?;
let wait_result = run_wait_result(&db, uuid, w_id, None).await;
let wait_result = run_wait_result(&db, uuid, w_id, None, &authed.username).await;
if delete_after_use.unwrap_or(false) {
delete_job_metadata_after_use(&db, uuid).await?;
}
@@ -3420,7 +3427,7 @@ pub async fn run_wait_result_script_by_hash(
.await?;
tx.commit().await?;
let wait_result = run_wait_result(&db, uuid, w_id, None).await;
let wait_result = run_wait_result(&db, uuid, w_id, None, &authed.username).await;
if delete_after_use.unwrap_or(false) {
delete_job_metadata_after_use(&db, uuid).await?;
}
@@ -3503,7 +3510,7 @@ async fn run_wait_result_flow_by_path_internal(
.await?;
tx.commit().await?;
run_wait_result(&db, uuid, w_id, early_return).await
run_wait_result(&db, uuid, w_id, early_return, &authed.username).await
}
async fn run_preview_script(
@@ -3792,7 +3799,7 @@ async fn run_dependencies_job(
.await?;
tx.commit().await?;
let wait_result = run_wait_result(&db, uuid, w_id, None).await;
let wait_result = run_wait_result(&db, uuid, w_id, None, &authed.username).await;
wait_result
}
@@ -3847,7 +3854,7 @@ async fn run_flow_dependencies_job(
.await?;
tx.commit().await?;
let wait_result = run_wait_result(&db, uuid, w_id, None).await;
let wait_result = run_wait_result(&db, uuid, w_id, None, &authed.username).await;
wait_result
}
+31 -39
View File
@@ -144,22 +144,7 @@ pub async fn cancel_single_job<'c>(
rsmq: Option<rsmq_async::MultiplexedRsmq>,
force_cancel: bool,
) -> error::Result<(Transaction<'c, Postgres>, Option<Uuid>)> {
if ((job_running.running || job_running.root_job.is_some()) || (job_running.is_flow()))
&& !force_cancel
{
let id = sqlx::query_scalar!(
"UPDATE queue SET canceled = true, canceled_by = $1, canceled_reason = $2, scheduled_for = now(), suspend = 0 WHERE id = $3 AND workspace_id = $4 AND canceled = false RETURNING id",
username,
reason,
job_running.id,
w_id
)
.fetch_optional(&mut *tx)
.await?;
if let Some(id) = id {
tracing::info!("Soft cancelling job {}", id);
}
} else {
if force_cancel || (job_running.parent_job.is_none() && !job_running.running) {
let username = username.to_string();
let job_running = job_running.clone();
let w_id = w_id.to_string();
@@ -193,6 +178,19 @@ pub async fn cancel_single_job<'c>(
tracing::error!("Failed to add canceled job: {}", e);
}
});
} else {
let id: Option<Uuid> = sqlx::query_scalar!(
"UPDATE queue SET canceled = true, canceled_by = $1, canceled_reason = $2, scheduled_for = now(), suspend = 0 WHERE id = $3 AND workspace_id = $4 AND canceled = false RETURNING id",
username,
reason,
job_running.id,
w_id
)
.fetch_optional(&mut *tx)
.await?;
if let Some(id) = id {
tracing::info!("Soft cancelling job {}", id);
}
}
if let Some(mut rsmq) = rsmq.clone() {
rsmq.change_message_visibility(&job_running.tag, &job_running.id.to_string(), 0)
@@ -225,10 +223,25 @@ pub async fn cancel_job<'c>(
"You are not logged in and this job was not created by an anonymous user like you so you cannot cancel it".to_string(),
));
}
let job = job.unwrap();
let mut job = job.unwrap();
if force_cancel {
// if force canceling a flow step, make sure we force cancel from the highest parent
loop {
if job.parent_job.is_none() {
break;
}
match get_queued_job_tx(job.parent_job.unwrap(), &w_id, &mut tx).await? {
Some(j) => {
job = j;
}
None => break,
}
}
}
// get all children
let mut jobs = vec![id];
let mut jobs = vec![job.id];
let mut jobs_to_cancel = vec![];
while !jobs.is_empty() {
let p_job = jobs.pop();
@@ -257,27 +270,6 @@ pub async fn cancel_job<'c>(
.await?;
tx = ntx;
// soft cancel parent if force cancel and job is a flow step
if force_cancel && job.is_flow_step {
if let Some(parent_job_id) = job.parent_job {
let job = get_queued_job_tx(parent_job_id, &w_id, &mut tx).await?;
if let Some(job) = job {
let (ntx, _) = cancel_single_job(
username,
reason.clone(),
&job,
w_id,
tx,
db,
rsmq.clone(),
false,
)
.await?;
tx = ntx;
}
}
}
// cancel children
for job_id in jobs_to_cancel {
let job = get_queued_job_tx(job_id, &w_id, &mut tx).await?;