diff --git a/backend/migrations/20221005022757_same_worker.down.sql b/backend/migrations/20221005022757_same_worker.down.sql new file mode 100644 index 0000000000..85e0ac0391 --- /dev/null +++ b/backend/migrations/20221005022757_same_worker.down.sql @@ -0,0 +1,2 @@ +-- Add down migration script here +ALTER TABLE queue DROP COLUMN same_worker; diff --git a/backend/migrations/20221005022757_same_worker.up.sql b/backend/migrations/20221005022757_same_worker.up.sql new file mode 100644 index 0000000000..0f71520b44 --- /dev/null +++ b/backend/migrations/20221005022757_same_worker.up.sql @@ -0,0 +1,2 @@ +-- Add up migration script here +ALTER TABLE queue ADD COLUMN same_worker BOOLEAN DEFAULT FALSE; diff --git a/backend/sqlx-data.json b/backend/sqlx-data.json index ec834e8888..b353e02b3a 100644 --- a/backend/sqlx-data.json +++ b/backend/sqlx-data.json @@ -384,53 +384,6 @@ }, "query": "SELECT * FROM workspace LIMIT $1 OFFSET $2" }, - "12d3b30b0f5f5bd9601565de5bdeadc3d99a946d2d91cf2bcba84bd161726d51": { - "describe": { - "columns": [ - { - "name": "id", - "ordinal": 0, - "type_info": "Uuid" - }, - { - "name": "workspace_id", - "ordinal": 1, - "type_info": "Varchar" - }, - { - "name": "last_ping", - "ordinal": 2, - "type_info": "Timestamptz" - } - ], - "nullable": [ - false, - false, - false - ], - "parameters": { - "Left": [ - "Text", - { - "Custom": { - "kind": { - "Enum": [ - "script", - "preview", - "flow", - "dependencies", - "flowpreview", - "script_hub" - ] - }, - "name": "job_kind" - } - } - ] - } - }, - "query": "UPDATE queue SET running = false WHERE last_ping < now() - ($1 || ' seconds')::interval AND running = true AND job_kind != $2 RETURNING id, workspace_id, last_ping" - }, "13a10a404e892d6975d4913424880a75bfb2c6dfb5134c270c7c7225acb05db4": { "describe": { "columns": [ @@ -632,66 +585,6 @@ }, "query": "SELECT canceled FROM queue WHERE id = $1" }, - "2420cb110a116dfbc6b8658a6d2d35db60c6aadd157488cb72eb9116ae7e9f54": { - "describe": { - "columns": [ - { - "name": "id", - "ordinal": 0, - "type_info": "Uuid" - } - ], - "nullable": [ - false - ], - "parameters": { - "Left": [ - "Varchar", - "Uuid", - "Uuid", - "Varchar", - "Varchar", - "Timestamptz", - "Int8", - "Varchar", - "Text", - "Jsonb", - { - "Custom": { - "kind": { - "Enum": [ - "script", - "preview", - "flow", - "dependencies", - "flowpreview", - "script_hub" - ] - }, - "name": "job_kind" - } - }, - "Varchar", - "Jsonb", - "Jsonb", - "Bool", - { - "Custom": { - "kind": { - "Enum": [ - "python3", - "deno", - "go" - ] - }, - "name": "script_lang" - } - } - ] - } - }, - "query": "INSERT INTO queue\n (workspace_id, id, parent_job, created_by, permissioned_as, scheduled_for, \n script_hash, script_path, raw_code, args, job_kind, schedule_path, raw_flow, flow_status, is_flow_step, language)\n VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16) RETURNING id" - }, "255aafff962738317f3227ae4eb871830d89b4c12c73d8dbabe6836da124e54d": { "describe": { "columns": [ @@ -1650,6 +1543,68 @@ }, "query": "SELECT EXISTS(SELECT 1 FROM usr WHERE workspace_id = $1 AND username = $2)" }, + "752a9a95a4fbf741b0013923bfb5efff4db5ccfa8d9e6d5f821fc230fe790cc4": { + "describe": { + "columns": [ + { + "name": "id", + "ordinal": 0, + "type_info": "Uuid" + } + ], + "nullable": [ + false + ], + "parameters": { + "Left": [ + "Varchar", + "Uuid", + "Bool", + "Uuid", + "Varchar", + "Varchar", + "Timestamptz", + "Int8", + "Varchar", + "Text", + "Jsonb", + { + "Custom": { + "kind": { + "Enum": [ + "script", + "preview", + "flow", + "dependencies", + "flowpreview", + "script_hub" + ] + }, + "name": "job_kind" + } + }, + "Varchar", + "Jsonb", + "Jsonb", + "Bool", + { + "Custom": { + "kind": { + "Enum": [ + "python3", + "deno", + "go" + ] + }, + "name": "script_lang" + } + }, + "Bool" + ] + } + }, + "query": "INSERT INTO queue\n (workspace_id, id, running, parent_job, created_by, permissioned_as, scheduled_for, \n script_hash, script_path, raw_code, args, job_kind, schedule_path, raw_flow, flow_status, is_flow_step, language, started_at, same_worker)\n VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, CASE WHEN $3 THEN now() END, $18) RETURNING id" + }, "765c18d77412cbb4474f4074d583b9b44681f3b9f58754662ac07a3a3470a3c5": { "describe": { "columns": [ @@ -2887,6 +2842,53 @@ }, "query": "INSERT INTO group_\n VALUES ($1, 'all', 'The group that always contains all users of this workspace')" }, + "cc34f0298cfdaffe892a286e6b4390c287d33298e3342ccd1e10d48a9c3cc5c0": { + "describe": { + "columns": [ + { + "name": "id", + "ordinal": 0, + "type_info": "Uuid" + }, + { + "name": "workspace_id", + "ordinal": 1, + "type_info": "Varchar" + }, + { + "name": "last_ping", + "ordinal": 2, + "type_info": "Timestamptz" + } + ], + "nullable": [ + false, + false, + false + ], + "parameters": { + "Left": [ + "Text", + { + "Custom": { + "kind": { + "Enum": [ + "script", + "preview", + "flow", + "dependencies", + "flowpreview", + "script_hub" + ] + }, + "name": "job_kind" + } + } + ] + } + }, + "query": "UPDATE queue SET running = false WHERE last_ping < now() - ($1 || ' seconds')::interval AND running = true AND job_kind != $2 AND same_worker = false RETURNING id, workspace_id, last_ping" + }, "d2dcf69b20488d610599c309862722f805049e479035be6a416d05d73528a8e1": { "describe": { "columns": [], diff --git a/backend/src/flows.rs b/backend/src/flows.rs index 98bd229d14..9f7c30695c 100644 --- a/backend/src/flows.rs +++ b/backend/src/flows.rs @@ -76,6 +76,9 @@ pub struct FlowValue { pub modules: Vec, #[serde(default)] pub failure_module: Option, + #[serde(default)] + #[serde(skip_serializing_if = "is_default")] + pub same_worker: bool, } #[derive(Deserialize, Serialize, Debug, Clone)] @@ -546,6 +549,7 @@ mod tests { suspend: Default::default(), retry: None, }), + same_worker: false, }; let expect = serde_json::json!({ "modules": [ diff --git a/backend/src/jobs.rs b/backend/src/jobs.rs index 25b00e92f7..26f9bd851b 100644 --- a/backend/src/jobs.rs +++ b/backend/src/jobs.rs @@ -24,7 +24,9 @@ use crate::{ users::{owner_to_token_owner, Authed}, utils::{now_from_db, require_admin, Pagination, StripPath}, worker, - worker_flow::{init_flow_status, FlowStatus, FlowStatusModule}, + worker_flow::{ + init_flow_status, FlowStatus, FlowStatusModule, MAX_RETRY_ATTEMPTS, MAX_RETRY_INTERVAL, + }, }; use axum::{ extract::{Extension, FromRequest, Path, Query}, @@ -103,6 +105,7 @@ pub struct QueuedJob { pub raw_flow: Option, pub is_flow_step: bool, pub language: Option, + pub same_worker: bool, } impl QueuedJob { @@ -188,6 +191,7 @@ pub async fn run_flow_by_path( None, run_query.parent_job, false, + false, ) .await?; tx.commit().await?; @@ -217,6 +221,7 @@ pub async fn run_job_by_path( None, run_query.parent_job, false, + false, ) .await?; tx.commit().await?; @@ -278,6 +283,7 @@ pub async fn run_wait_result_job_by_path( None, run_query.parent_job, false, + false, ) .await?; tx.commit().await?; @@ -308,6 +314,7 @@ pub async fn run_wait_result_job_by_hash( None, run_query.parent_job, false, + false, ) .await?; tx.commit().await?; @@ -374,6 +381,7 @@ pub async fn run_job_by_hash( None, run_query.parent_job, false, + false, ) .await?; tx.commit().await?; @@ -426,6 +434,7 @@ async fn run_preview_job( None, None, false, + false, ) .await?; tx.commit().await?; @@ -452,6 +461,7 @@ async fn run_preview_flow_job( None, None, false, + false, ) .await?; tx.commit().await?; @@ -1154,6 +1164,7 @@ impl From for Job { raw_flow: None, is_flow_step: uj.is_flow_step, language: uj.language, + same_worker: false, }), t => panic!("job type {} not valid", t), } @@ -1208,6 +1219,7 @@ pub async fn push<'c>( schedule_path: Option, parent_job: Option, is_flow_step: bool, + mut same_worker: bool, ) -> Result<(Uuid, Transaction<'c, Postgres>), Error> { let scheduled_for = scheduled_for_o.unwrap_or_else(chrono::Utc::now); let args_json = args.map(serde_json::Value::Object); @@ -1361,16 +1373,45 @@ pub async fn push<'c>( } }; + let mut is_running = same_worker; + if let Some(flow) = raw_flow.as_ref() { + is_running = false; + same_worker = same_worker || flow.same_worker; + if flow.modules.len() == 0 { + Err(Error::BadRequest(format!( + "A flow needs at least one module to run" + )))?; + } + + for module in flow.modules.iter() { + if let Some(retry) = &module.retry { + if retry.max_attempts() > MAX_RETRY_ATTEMPTS { + Err(Error::BadRequest(format!( + "retry attempts exceeds the maximum of {MAX_RETRY_ATTEMPTS}" + )))? + } + + if matches!(retry.max_interval(), Some(interval) if interval > MAX_RETRY_INTERVAL) { + let max = MAX_RETRY_INTERVAL.as_secs(); + Err(Error::BadRequest(format!( + "retry interval exceeds the maximum of {max} seconds" + )))? + } + } + } + } + let flow_status = raw_flow.as_ref().map(init_flow_status); let uuid = sqlx::query_scalar!( "INSERT INTO queue - (workspace_id, id, parent_job, created_by, permissioned_as, scheduled_for, + (workspace_id, id, running, parent_job, created_by, permissioned_as, scheduled_for, script_hash, script_path, raw_code, args, job_kind, schedule_path, raw_flow, \ - flow_status, is_flow_step, language) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16) \ + flow_status, is_flow_step, language, started_at, same_worker) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, CASE WHEN $3 THEN now() END, $18) \ RETURNING id", workspace_id, job_id, + is_running, parent_job, user, permissioned_as, @@ -1384,7 +1425,8 @@ pub async fn push<'c>( raw_flow.map(|f| serde_json::json!(f)), flow_status.map(|f| serde_json::json!(f)), is_flow_step, - language: ScriptLang + language: ScriptLang, + same_worker ) .fetch_one(&mut tx) .await @@ -1429,9 +1471,9 @@ pub async fn add_completed_job_error( queued_job: &QueuedJob, logs: String, e: E, - metrics: &worker::Metrics, + metrics: Option, ) -> Result<(Uuid, Map), Error> { - metrics.jobs_failed.inc(); + metrics.map(|m| m.jobs_failed.inc()); let mut output_map = serde_json::Map::new(); output_map.insert( "error".to_string(), diff --git a/backend/src/lib.rs b/backend/src/lib.rs index 217c7fc3bd..e51107325e 100644 --- a/backend/src/lib.rs +++ b/backend/src/lib.rs @@ -212,7 +212,7 @@ pub fn monitor_db(db: &DB, timeout: i32, rx: tokio::sync::broadcast::Receiver<() let rx2 = rx.resubscribe(); - tokio::spawn(async move { worker::restart_zombie_jobs_periodically(&db1, timeout, rx).await }); + tokio::spawn(async move { worker::handle_zombie_jobs_periodically(&db1, timeout, rx).await }); tokio::spawn(async move { users::delete_expired_items_perdiodically(&db2, rx2).await }); } diff --git a/backend/src/oauth2.rs b/backend/src/oauth2.rs index efa6c0c46a..af07dff5c9 100644 --- a/backend/src/oauth2.rs +++ b/backend/src/oauth2.rs @@ -678,6 +678,7 @@ async fn slack_command( None, None, false, + false, ) .await?; tx.commit().await?; diff --git a/backend/src/schedule.rs b/backend/src/schedule.rs index c6a7e0ea67..22944d9d76 100644 --- a/backend/src/schedule.rs +++ b/backend/src/schedule.rs @@ -119,6 +119,7 @@ pub async fn push_scheduled_job<'c>( Some(schedule.path), None, false, + false, ) .await?; Ok(tx) diff --git a/backend/src/scripts.rs b/backend/src/scripts.rs index 73844e8a3e..eb1a3c52d3 100644 --- a/backend/src/scripts.rs +++ b/backend/src/scripts.rs @@ -455,6 +455,7 @@ async fn create_script( None, None, false, + false, ) .await?; tx diff --git a/backend/src/worker.rs b/backend/src/worker.rs index b1ee50bf81..8b2cabac6d 100644 --- a/backend/src/worker.rs +++ b/backend/src/worker.rs @@ -39,7 +39,10 @@ use tokio::{ fs::{DirBuilder, File}, io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}, process::{Child, Command}, - sync::watch, + sync::{ + mpsc::{self, Sender}, + watch, + }, time::{interval, sleep, Instant, MissedTickBehavior}, }; @@ -67,6 +70,8 @@ const NSJAIL_CONFIG_RUN_GO_CONTENT: &str = include_str!("../../nsjail/run.go.con const NSJAIL_CONFIG_RUN_DENO_CONTENT: &str = include_str!("../../nsjail/run.deno.config.proto"); const MAX_LOG_SIZE: u32 = 200000; const GO_REQ_SPLITTER: &str = "//go.sum"; + +#[derive(Clone)] pub struct Metrics { pub jobs_failed: prometheus::IntCounter, } @@ -171,6 +176,8 @@ pub async fn run_worker( pip_trusted_host, }; + let (same_worker_tx, mut same_worker_rx) = mpsc::channel::(5); + loop { if last_ping.elapsed().as_secs() > NUM_SECS_ENV_CHECK { sqlx::query!( @@ -185,7 +192,24 @@ pub async fn run_worker( last_ping = Instant::now(); } - match pull(db).await { + let next_job = tokio::select! { + biased; + _ = rx.recv() => { + println!("received killpill for worker {}", i_worker); + break; + }, + Some(job_id) = same_worker_rx.recv() => { + sqlx::query_as::<_, QueuedJob>("SELECT * FROM queue WHERE id = $1") + .bind(job_id) + .fetch_optional(db) + .await + .map_err(|_| Error::InternalErr("Impossible to fetch same_worker job".to_string())) + }, + job = pull(&db) => + job, + }; + + match next_job { Ok(Some(job)) => { let label_values = [ &job.workspace_id, @@ -203,94 +227,133 @@ pub async fn run_worker( tracing::info!(worker = %worker_name, id = %job.id, "fetched job {}", job.id); + let job_dir = format!("{worker_dir}/{}", job.id); + + DirBuilder::new() + .create(&job_dir) + .await + .expect("could not create job dir"); + + let same_worker = job.same_worker; + let is_flow = job.job_kind == JobKind::Flow || job.job_kind == JobKind::FlowPreview; + + if is_flow && same_worker { + DirBuilder::new() + .create(&format!("{job_dir}/shared")) + .await + .expect("could not create shared dir"); + } + if let Some(err) = handle_queued_job( job.clone(), db, timeout, &worker_name, &worker_dir, + &job_dir, &worker_config, - &metrics, + metrics.clone(), &envs, + same_worker_tx.clone(), ) .await .err() { - let m = add_completed_job_error( - db, - &job, - "Unexpected error during job execution:\n".to_string(), - &err, - &metrics, - ) - .await - .map(|(_, m)| m) - .unwrap_or_else(|_| Map::new()); - - let _ = postprocess_queued_job( - job.is_flow_step, - job.schedule_path.clone(), - job.script_path.clone(), - &job.workspace_id, - job.id, + handle_job_error( db, + job, + err, + Some(metrics), + false, + same_worker_tx.clone(), + &worker_dir, + !worker_config.keep_job_dir, ) .await; - - if let Some(parent_job_id) = job.parent_job { - let updated_flow = update_flow_status_after_job_completion( - db, - &job, - false, - serde_json::Value::Object(m), - &metrics, - ) - .await; - if let Err(err) = updated_flow { - if let Ok(mut tx) = db.begin().await { - if let Ok(Some(parent_job)) = - get_queued_job(parent_job_id, &job.workspace_id, &mut tx).await - { - let _ = add_completed_job_error( - db, - &parent_job, - format!("Unexpected error during flow job error handling:\n{err}") - , - err, - &metrics, - ) - .await; - - let _ = postprocess_queued_job( - parent_job.is_flow_step, - parent_job.schedule_path.clone(), - parent_job.script_path.clone(), - &job.workspace_id, - parent_job.id, - db, - ) - .await; - } - } - } - } - tracing::error!(job_id = %job.id, err = err.alt(), "error handling job: {} {} {}", job.id, job.workspace_id, job.created_by); }; + + if !worker_config.keep_job_dir && !(is_flow && same_worker) { + let _ = tokio::fs::remove_dir_all(job_dir).await; + } } - Ok(None) => (), + Ok(None) => tokio::time::sleep(Duration::from_millis(sleep_queue * num_workers)).await, Err(err) => { tracing::error!(worker = %worker_name, "run_worker: pulling jobs: {}", err); } }; + } +} - tokio::select! { - _ = tokio::time::sleep(Duration::from_millis(sleep_queue * num_workers)) => (), - _ = rx.recv() => { - println!("received killpill for worker {}", i_worker); - break; +async fn handle_job_error( + db: &DB, + job: QueuedJob, + err: Error, + metrics: Option, + unrecoverable: bool, + same_worker_tx: Sender, + worker_dir: &str, + keep_job_dir: bool, +) { + let m = add_completed_job_error( + db, + &job, + "Unexpected error during job execution:\n".to_string(), + &err, + metrics.clone(), + ) + .await + .map(|(_, m)| m) + .unwrap_or_else(|_| Map::new()); + let _ = postprocess_queued_job( + job.is_flow_step, + job.schedule_path.clone(), + job.script_path.clone(), + &job.workspace_id, + job.id, + db, + ) + .await; + if let Some(parent_job_id) = job.parent_job { + let updated_flow = update_flow_status_after_job_completion( + db, + &job, + false, + serde_json::Value::Object(m), + metrics.clone(), + unrecoverable, + same_worker_tx, + worker_dir, + keep_job_dir, + ) + .await; + if let Err(err) = updated_flow { + if let Ok(mut tx) = db.begin().await { + if let Ok(Some(parent_job)) = + get_queued_job(parent_job_id, &job.workspace_id, &mut tx).await + { + let _ = add_completed_job_error( + db, + &parent_job, + format!("Unexpected error during flow job error handling:\n{err}"), + err, + metrics, + ) + .await; + + let _ = postprocess_queued_job( + parent_job.is_flow_step, + parent_job.schedule_path.clone(), + parent_job.script_path.clone(), + &job.workspace_id, + parent_job.id, + db, + ) + .await; + } } } } + tracing::error!(job_id = %job.id, err = err.alt(), "error handling job: {} {} {}", job.id, job.workspace_id, job.created_by); } async fn insert_initial_ping(worker_instance: &str, worker_name: &str, ip: &str, db: &DB) { @@ -317,15 +380,18 @@ struct Envs { pip_extra_index_url: Option, pip_trusted_host: Option, } + async fn handle_queued_job( job: QueuedJob, db: &sqlx::Pool, timeout: i32, worker_name: &str, worker_dir: &str, + job_dir: &str, worker_config: &WorkerConfig, - metrics: &Metrics, + metrics: Metrics, envs: &Envs, + same_worker_tx: Sender, ) -> crate::error::Result<()> { let job_id = job.id; let w_id = &job.workspace_id.clone(); @@ -333,7 +399,7 @@ async fn handle_queued_job( match job.job_kind { JobKind::FlowPreview | JobKind::Flow => { let args = job.args.clone().unwrap_or(Value::Null); - handle_flow(&job, db, args).await?; + handle_flow(&job, db, args, same_worker_tx).await?; } _ => { let mut logs = "".to_string(); @@ -352,10 +418,11 @@ async fn handle_queued_job( let execution = handle_job( &job, + &job_dir, db, timeout, worker_name, - worker_dir, + worker_dir.clone(), &mut logs, &mut last_line, worker_config, @@ -367,19 +434,34 @@ async fn handle_queued_job( Ok(r) => { add_completed_job(db, &job, true, false, r.clone(), logs).await?; if job.is_flow_step { - update_flow_status_after_job_completion(db, &job, true, r, metrics).await?; + update_flow_status_after_job_completion( + db, + &job, + true, + r, + Some(metrics.clone()), + false, + same_worker_tx.clone(), + worker_dir, + worker_config.keep_job_dir, + ) + .await?; } } Err(e) => { let (_, output_map) = - add_completed_job_error(db, &job, logs, e, &metrics).await?; + add_completed_job_error(db, &job, logs, e, Some(metrics.clone())).await?; if job.is_flow_step { update_flow_status_after_job_completion( db, &job, false, serde_json::Value::Object(output_map), - metrics, + Some(metrics), + false, + same_worker_tx, + worker_dir, + worker_config.keep_job_dir, ) .await?; } @@ -450,6 +532,7 @@ async fn transform_json_value( #[allow(clippy::too_many_arguments)] async fn handle_job( job: &QueuedJob, + job_dir: &str, db: &DB, timeout: i32, worker_name: &str, @@ -468,12 +551,6 @@ async fn handle_job( ); logs.push_str(&format!("job {} on worker {}\n", &job.id, &worker_name)); - let job_dir = format!("{worker_dir}/{}", job.id); - - DirBuilder::new() - .create(&job_dir) - .await - .expect("could not create job dir"); let mut status: Result = Err(Error::InternalErr("job not started".to_string())); @@ -482,7 +559,7 @@ async fn handle_job( handle_dependency_job( job, logs, - &job_dir, + job_dir, &mut status, db, last_line, @@ -494,7 +571,7 @@ async fn handle_job( handle_code_execution_job( job, db, - &job_dir, + job_dir, worker_dir, logs, &mut status, @@ -506,10 +583,6 @@ async fn handle_job( .await?; } - if !worker_config.keep_job_dir { - tokio::fs::remove_dir_all(job_dir).await?; - } - if status.is_ok() && status.as_ref().unwrap().success() { let result = serde_json::from_str::(last_line).map_err(|e| { Error::ExecutionErr(format!( @@ -523,9 +596,9 @@ async fn handle_job( let err = match status { Ok(_) => { let s = format!( - "Error during execution of the script\nlast 5 logs lines:\n{}", + "Error during execution of the script\nlast 10 logs lines:\n{}", logs.lines() - .skip(logs.lines().count().max(5) - 5) + .skip(logs.lines().count().max(10) - 10) .join("\n") ); logs.push_str("\n\n--- ERROR ---\n"); @@ -540,7 +613,7 @@ async fn handle_job( async fn handle_code_execution_job( job: &QueuedJob, db: &sqlx::Pool, - job_dir: &String, + job_dir: &str, worker_dir: &str, logs: &mut String, status: &mut Result, @@ -585,6 +658,24 @@ async fn handle_code_execution_job( &lang_str, job.id ); + + let shared_mount = if job.same_worker { + format!( + r#" +mount {{ + src: "{worker_dir}/{}/shared" + dst: "/shared" + is_bind: true + rw: true +}} + "#, + job.parent_job.ok_or(Error::ExecutionErr( + "no parent job, required for same worker job".to_string() + ))?, + ) + } else { + "".to_string() + }; match language { None => { return Err(Error::ExecutionErr( @@ -606,6 +697,7 @@ async fn handle_code_execution_job( last_line, timeout, &inner_content, + &shared_mount, ) .await? } @@ -621,6 +713,7 @@ async fn handle_code_execution_job( timeout, status, last_line, + &shared_mount, ) .await?; } @@ -637,6 +730,7 @@ async fn handle_code_execution_job( requirements_o, status, last_line, + &shared_mount, ) .await? } @@ -661,10 +755,11 @@ async fn handle_go_job( db: &sqlx::Pool, inner_content: &str, timeout: i32, - job_dir: &String, + job_dir: &str, requirements_o: Option, status: &mut Result, last_line: &mut String, + shared_mount: &str, ) -> Result<(), Error> { //go does not like executing modules at temp root let job_dir = &format!("{job_dir}/go"); @@ -784,7 +879,8 @@ func main() {{ &NSJAIL_CONFIG_RUN_GO_CONTENT .replace("{JOB_DIR}", job_dir) .replace("{CACHE_DIR}", GO_CACHE_DIR) - .replace("{CLONE_NEWUSER}", &(!disable_nuser).to_string()), + .replace("{CLONE_NEWUSER}", &(!disable_nuser).to_string()) + .replace("{SHARED_MOUNT}", shared_mount), ) .await?; @@ -829,11 +925,12 @@ async fn handle_deno_job( logs: &mut String, job: &QueuedJob, db: &sqlx::Pool, - job_dir: &String, + job_dir: &str, inner_content: &String, timeout: i32, status: &mut Result, last_line: &mut String, + shared_mount: &str, ) -> Result<(), Error> { logs.push_str("\n\n--- DENO CODE EXECUTION ---\n"); set_logs(logs, job.id, db).await; @@ -892,7 +989,8 @@ run(); &NSJAIL_CONFIG_RUN_DENO_CONTENT .replace("{JOB_DIR}", job_dir) .replace("{CACHE_DIR}", DENO_CACHE_DIR) - .replace("{CLONE_NEWUSER}", &(!disable_nuser).to_string()), + .replace("{CLONE_NEWUSER}", &(!disable_nuser).to_string()) + .replace("{SHARED_MOUNT}", shared_mount), ) .await?; Command::new(nsjail_path) @@ -952,7 +1050,7 @@ async fn handle_python_job( .. }: &Envs, requirements_o: Option, - job_dir: &String, + job_dir: &str, worker_dir: &str, worker_name: &str, job: &QueuedJob, @@ -962,6 +1060,7 @@ async fn handle_python_job( last_line: &mut String, timeout: i32, inner_content: &String, + shared_mount: &str, ) -> Result<(), Error> { let requirements = requirements_o.ok_or_else(|| Error::InternalErr(format!("lockfile missing")))?; @@ -1136,7 +1235,8 @@ print(res_json) "run.config.proto", &NSJAIL_CONFIG_RUN_PYTHON3_CONTENT .replace("{JOB_DIR}", job_dir) - .replace("{CLONE_NEWUSER}", &(!disable_nuser).to_string()), + .replace("{CLONE_NEWUSER}", &(!disable_nuser).to_string()) + .replace("{SHARED_MOUNT}", shared_mount), ) .await?; } else { @@ -1197,7 +1297,7 @@ async fn create_dependencies_dir(job_dir: &str) { async fn handle_dependency_job( job: &QueuedJob, logs: &mut String, - job_dir: &String, + job_dir: &str, status: &mut error::Result, db: &sqlx::Pool, last_line: &mut String, @@ -1305,7 +1405,7 @@ async fn install_go_dependencies( job_id: &Uuid, code: &str, logs: &mut String, - job_dir: &String, + job_dir: &str, status: &mut Result, db: &sqlx::Pool, last_line: &mut String, @@ -1356,7 +1456,7 @@ async fn install_go_dependencies( } } -async fn gen_go_mymod(code: &str, job_dir: &String) -> error::Result<()> { +async fn gen_go_mymod(code: &str, job_dir: &str) -> error::Result<()> { let code = &format!("package inner; {code}").replace("func main(", "func Inner_main("); let mymod_dir = format!("{job_dir}/inner"); @@ -1705,14 +1805,14 @@ async fn append_logs(job_id: uuid::Uuid, logs: impl AsRef, db: impl Borrow< } } -pub async fn restart_zombie_jobs_periodically( +pub async fn handle_zombie_jobs_periodically( db: &DB, timeout: i32, mut rx: tokio::sync::broadcast::Receiver<()>, ) { loop { let restarted = sqlx::query!( - "UPDATE queue SET running = false WHERE last_ping < now() - ($1 || ' seconds')::interval AND running = true AND job_kind != $2 RETURNING id, workspace_id, last_ping", + "UPDATE queue SET running = false WHERE last_ping < now() - ($1 || ' seconds')::interval AND running = true AND job_kind != $2 AND same_worker = false RETURNING id, workspace_id, last_ping", (timeout * 5).to_string(), JobKind::Flow: JobKind, ) @@ -1730,6 +1830,39 @@ pub async fn restart_zombie_jobs_periodically( ); } + let timeouts = sqlx::query_as::<_, QueuedJob>( + "SELECT * FROM queue WHERE last_ping < now() - ($1 || ' seconds')::interval AND running = true AND job_kind != $2 AND same_worker = true", + ) + .bind((timeout * 5).to_string()) + .bind(JobKind::Flow) + .fetch_all(db) + .await + .ok() + .unwrap_or_else(|| vec![]); + + for job in timeouts { + tracing::info!( + "timedouts zombie same_worker job {} {}", + job.id, + job.workspace_id, + ); + + // since the job is unrecoverable, the same worker queue should never be sent anything + let (same_worker_tx_never_used, _same_worker_rx_never_used) = mpsc::channel::(1); + + let _ = handle_job_error( + db, + job, + error::Error::ExecutionErr("Same worker job timed out".to_string()), + None, + true, + same_worker_tx_never_used, + "", + true, + ) + .await; + } + tokio::select! { _ = tokio::time::sleep(Duration::from_secs(60)) => (), _ = rx.recv() => { @@ -1796,51 +1929,54 @@ mod tests { let numbers = "export function main() { return [1, 2, 3]; }"; let doubles = "export function main(n) { return n * 2; }"; - let flow = FlowValue { - modules: vec![ - FlowModule { - value: FlowModuleValue::RawScript(RawCode { - language: ScriptLang::Deno, - content: numbers.to_string(), - path: None, - }), - input_transforms: Default::default(), - stop_after_if: Default::default(), - summary: Default::default(), - suspend: Default::default(), - retry: None, - }, - FlowModule { - value: FlowModuleValue::ForloopFlow { - iterator: InputTransform::Javascript { expr: "result".to_string() }, - skip_failures: false, - modules: vec![FlowModule { - value: FlowModuleValue::RawScript(RawCode { - language: ScriptLang::Deno, - content: doubles.to_string(), - path: None, - }), - input_transforms: [( - "n".to_string(), - InputTransform::Javascript { - expr: "previous_result.iter.value".to_string(), - }, - )] - .into(), - stop_after_if: Default::default(), - summary: Default::default(), - suspend: Default::default(), - retry: None, - }], + let flow = { + FlowValue { + modules: vec![ + FlowModule { + value: FlowModuleValue::RawScript(RawCode { + language: ScriptLang::Deno, + content: numbers.to_string(), + path: None, + }), + input_transforms: Default::default(), + stop_after_if: Default::default(), + summary: Default::default(), + suspend: Default::default(), + retry: None, }, - input_transforms: Default::default(), - stop_after_if: Default::default(), - summary: Default::default(), - suspend: Default::default(), - retry: None, - }, - ], - ..Default::default() + FlowModule { + value: FlowModuleValue::ForloopFlow { + iterator: InputTransform::Javascript { expr: "result".to_string() }, + skip_failures: false, + modules: vec![FlowModule { + value: FlowModuleValue::RawScript(RawCode { + language: ScriptLang::Deno, + content: doubles.to_string(), + path: None, + }), + input_transforms: [( + "n".to_string(), + InputTransform::Javascript { + expr: "previous_result.iter.value".to_string(), + }, + )] + .into(), + stop_after_if: Default::default(), + summary: Default::default(), + suspend: Default::default(), + retry: None, + }], + }, + input_transforms: Default::default(), + stop_after_if: Default::default(), + summary: Default::default(), + suspend: Default::default(), + retry: None, + }, + ], + same_worker: false, + ..Default::default() + } }; let job = JobPayload::RawFlow { value: flow, path: None }; @@ -1852,6 +1988,120 @@ mod tests { } } + #[sqlx::test(fixtures("base"))] + async fn test_deno_flow_same_worker(db: DB) { + initialize_tracing().await; + + let write_file = r#"export async function main(loop: boolean, i: number) { + await Deno.writeTextFile("/shared/file.txt", `${loop} ${i}`); + }"# + .to_string(); + + let read_file = r#"export async function main() { + return await Deno.readTextFile("/shared/file.txt"); + }"# + .to_string(); + + let flow = FlowValue { + modules: vec![ + FlowModule { + value: FlowModuleValue::RawScript(RawCode { + language: ScriptLang::Deno, + content: write_file.clone(), + path: None, + }), + input_transforms: [ + ( + "loop".to_string(), + InputTransform::Static { value: json!(false) }, + ), + ("i".to_string(), InputTransform::Static { value: json!(1) }), + ] + .into(), + stop_after_if: Default::default(), + summary: Default::default(), + suspend: Default::default(), + retry: None, + }, + FlowModule { + value: FlowModuleValue::ForloopFlow { + iterator: InputTransform::Static { value: json!([1, 2, 3]) }, + skip_failures: false, + modules: vec![ + FlowModule { + value: FlowModuleValue::RawScript(RawCode { + language: ScriptLang::Deno, + content: write_file, + path: None, + }), + input_transforms: [ + ( + "i".to_string(), + InputTransform::Javascript { + expr: "previous_result.iter.value".to_string(), + }, + ), + ( + "loop".to_string(), + InputTransform::Static { value: json!(true) }, + ), + ] + .into(), + stop_after_if: Default::default(), + summary: Default::default(), + suspend: Default::default(), + retry: None, + }, + FlowModule { + value: FlowModuleValue::RawScript(RawCode { + language: ScriptLang::Deno, + content: read_file.clone(), + path: None, + }), + input_transforms: [].into(), + stop_after_if: Default::default(), + summary: Default::default(), + suspend: Default::default(), + retry: None, + }, + ], + }, + input_transforms: Default::default(), + stop_after_if: Default::default(), + summary: Default::default(), + suspend: Default::default(), + retry: None, + }, + FlowModule { + value: FlowModuleValue::RawScript(RawCode { + language: ScriptLang::Deno, + content: r#"export async function main(loops: string[]) { + return await Deno.readTextFile("/shared/file.txt") + ","+ loops; + }"# + .to_string(), + path: None, + }), + input_transforms: [( + "loops".to_string(), + InputTransform::Javascript { expr: "previous_result".to_string() }, + )] + .into(), + stop_after_if: Default::default(), + summary: Default::default(), + suspend: Default::default(), + retry: None, + }, + ], + same_worker: true, + ..Default::default() + }; + + let job = JobPayload::RawFlow { value: flow, path: None }; + + let result = run_job_in_new_worker_until_complete(&db, job.clone()).await; + assert_eq!(result, serde_json::json!("false 1,true 1,true 2,true 3")); + } + #[sqlx::test(fixtures("base"))] async fn test_stop_after_if(db: DB) { initialize_tracing().await; @@ -2771,35 +3021,11 @@ def main(error, port): json!({ "recv": 42, "from failure module": { - "error": "\ - Error during execution of the script\nlast 5 logs lines:\n \ - File \"/tmp/main.py\", line 14, in \n \ - res = inner_script.main(**kwargs)\n \ - File \"/tmp/inner.py\", line 5, in main\n \ - return sock.recv(1)[0]\nIndexError: index out of range" + "error": "Error during execution of the script\nlast 10 logs lines:\n\n\n--- PYTHON CODE EXECUTION ---\n\nTraceback (most recent call last):\n File \"/tmp/main.py\", line 14, in \n res = inner_script.main(**kwargs)\n File \"/tmp/inner.py\", line 5, in main\n return sock.recv(1)[0]\nIndexError: index out of range", } }) ); } - - #[sqlx::test(fixtures("base"))] - async fn bad_values_max(db: DB) { - let value = serde_json::from_value(json!({ - "modules": [{ - "value": { "type": "rawscript", "language": "python3", "content": "asdf" }, - "retry": { "exponential": { "attempts": 50, "seconds": 60 } }, - }], - })) - .unwrap(); - - let result = RunJob::from(JobPayload::RawFlow { value, path: None }) - .run_until_complete(&db) - .await; - assert_eq!( - result, - json!({"error": "Bad request: retry interval exceeds the maximum of 21600 seconds"}) - ) - } } #[sqlx::test(fixtures("base"))] @@ -2879,6 +3105,7 @@ def main(error, port): /* schedule_path */ None, /* parent_job */ None, /* is_flow_step */ false, + /* running */ false, ) .await .unwrap(); @@ -2953,7 +3180,10 @@ def main(error, port): base_url: String::new(), disable_nuser: false, disable_nsjail: false, - keep_job_dir: false, + keep_job_dir: std::env::var("KEEP_JOB_DIR") + .ok() + .and_then(|x| x.parse::().ok()) + .unwrap_or(false), }; let future = async move { run_worker( diff --git a/backend/src/worker_flow.rs b/backend/src/worker_flow.rs index 9d8a5e258e..55680ad220 100644 --- a/backend/src/worker_flow.rs +++ b/backend/src/worker_flow.rs @@ -19,14 +19,15 @@ use async_recursion::async_recursion; use futures::TryStreamExt; use serde::{Deserialize, Serialize}; use serde_json::{json, Map, Value}; +use tokio::sync::mpsc::Sender; use tracing::instrument; use uuid::Uuid; const MINUTES: Duration = Duration::from_secs(60); const HOURS: Duration = MINUTES.saturating_mul(60); -const MAX_RETRY_ATTEMPTS: u16 = 1000; -const MAX_RETRY_INTERVAL: Duration = HOURS.saturating_mul(6); +pub const MAX_RETRY_ATTEMPTS: u16 = 1000; +pub const MAX_RETRY_INTERVAL: Duration = HOURS.saturating_mul(6); #[derive(Serialize, Deserialize, Debug)] pub struct FlowStatus { @@ -87,7 +88,11 @@ pub async fn update_flow_status_after_job_completion( job: &QueuedJob, success: bool, result: serde_json::Value, - metrics: &worker::Metrics, + metrics: Option, + unrecoverable: bool, + same_worker_tx: Sender, + worker_dir: &str, + keep_job_dir: bool, ) -> error::Result<()> { tracing::debug!("HANDLE FLOW: {job:?} {success} {result:?}"); @@ -240,6 +245,7 @@ pub async fn update_flow_status_after_job_completion( _ if stop_early => false, _ if flow_job.canceled => false, true => !is_last_step, + false if unrecoverable => false, false if skip_loop_failures => !is_last_step, false if next_retry( @@ -278,14 +284,14 @@ pub async fn update_flow_status_after_job_completion( .await?; true } else { - match handle_flow(&flow_job, db, result.clone()).await { + match handle_flow(&flow_job, db, result.clone(), same_worker_tx.clone()).await { Err(err) => { let _ = add_completed_job_error( db, &flow_job, "Unexpected error during flow chaining:\n".to_string(), err, - metrics, + metrics.clone(), ) .await; true @@ -305,9 +311,21 @@ pub async fn update_flow_status_after_job_completion( ) .await?; + if flow_job.same_worker && !keep_job_dir { + let _ = tokio::fs::remove_dir_all(format!("{worker_dir}/{}", flow_job.id)).await; + } + if flow_job.parent_job.is_some() { return Ok(update_flow_status_after_job_completion( - db, &flow_job, success, result, metrics, + db, + &flow_job, + success, + result, + metrics, + false, + same_worker_tx.clone(), + worker_dir, + keep_job_dir, ) .await?); } @@ -470,6 +488,7 @@ pub async fn handle_flow( flow_job: &QueuedJob, db: &sqlx::Pool, last_result: serde_json::Value, + same_worker_tx: Sender, ) -> anyhow::Result<()> { let value = flow_job .raw_flow @@ -478,37 +497,7 @@ pub async fn handle_flow( .to_owned(); let flow = serde_json::from_value::(value.to_owned())?; - if flow.modules.len() == 0 { - Err(Error::BadRequest(format!( - "A flow needs at least one module to run" - )))?; - } - - for module in flow.modules.iter() { - if let Some(retry) = &module.retry { - if retry.max_attempts() > MAX_RETRY_ATTEMPTS { - Err(Error::BadRequest(format!( - "retry attempts exceeds the maximum of {MAX_RETRY_ATTEMPTS}" - )))? - } - - if matches!(retry.max_interval(), Some(interval) if interval > MAX_RETRY_INTERVAL) { - let max = MAX_RETRY_INTERVAL.as_secs(); - Err(Error::BadRequest(format!( - "retry interval exceeds the maximum of {max} seconds" - )))? - } - } - } - - push_next_flow_job( - flow_job, - flow, - flow_job.schedule_path.clone(), - db, - last_result, - ) - .await?; + push_next_flow_job(flow_job, flow, db, last_result, same_worker_tx).await?; Ok(()) } @@ -517,9 +506,9 @@ pub async fn handle_flow( async fn push_next_flow_job( flow_job: &QueuedJob, flow: FlowValue, - schedule_path: Option, db: &sqlx::Pool, mut last_result: serde_json::Value, + same_worker_tx: Sender, ) -> anyhow::Result<()> { let status: FlowStatus = serde_json::from_value::(flow_job.flow_status.clone().unwrap_or_default()) @@ -841,7 +830,7 @@ async fn push_next_flow_job( .await?; return if next_step.is_some() { - push_next_flow_job(&new_job, flow, schedule_path, db, json!([])).await + push_next_flow_job(&new_job, flow, db, json!([]), same_worker_tx).await } else { let success = true; let skipped = false; @@ -925,6 +914,7 @@ async fn push_next_flow_job( value: FlowValue { modules: (*modules).clone(), failure_module: flow.failure_module.clone(), + same_worker: flow.same_worker, }, path: Some(format!("{}/{}", flow_job.script_path(), status.step)), }, @@ -937,6 +927,8 @@ async fn push_next_flow_job( } }; + let continue_on_same_worker = + flow.same_worker && !matches!(job_payload, JobPayload::RawFlow { .. }); let (uuid, mut tx) = push( tx, &flow_job.workspace_id, @@ -945,9 +937,10 @@ async fn push_next_flow_job( &flow_job.created_by, flow_job.permissioned_as.to_owned(), scheduled_for_o, - schedule_path, + flow_job.schedule_path.clone(), Some(flow_job.id), true, + continue_on_same_worker, ) .await?; @@ -982,6 +975,9 @@ async fn push_next_flow_job( tx.commit().await?; + if continue_on_same_worker { + same_worker_tx.send(uuid).await?; + } return Ok(()); /// Some state about the current/last forloop FlowStatusModule used to initialized the next diff --git a/frontend/src/lib/components/flows/content/FlowLoop.svelte b/frontend/src/lib/components/flows/content/FlowLoop.svelte index f8972c4c7c..072a8e3b0e 100644 --- a/frontend/src/lib/components/flows/content/FlowLoop.svelte +++ b/frontend/src/lib/components/flows/content/FlowLoop.svelte @@ -39,15 +39,15 @@ ).pickableProperties - -
- -
-
+
+ +
+ +
- - -
+ + +
{#if mod.value.type === 'forloopflow'} Iterator expression @@ -179,6 +179,6 @@ -
-
- +
+
diff --git a/frontend/src/lib/components/flows/content/FlowSettings.svelte b/frontend/src/lib/components/flows/content/FlowSettings.svelte index 6830e128a9..7d6b0a7261 100644 --- a/frontend/src/lib/components/flows/content/FlowSettings.svelte +++ b/frontend/src/lib/components/flows/content/FlowSettings.svelte @@ -9,6 +9,7 @@ import FlowCard from '../common/FlowCard.svelte' import FlowSchedules from './FlowSchedules.svelte' import SvelteMarkdown from 'svelte-markdown' + import { Toggle } from 'flowbite-svelte' export let initialPath: string @@ -19,6 +20,7 @@ Metadata Schedule + Same Worker @@ -77,6 +79,17 @@ + + + Same Worker + + + diff --git a/nsjail/run.deno.config.proto b/nsjail/run.deno.config.proto index 03e9275b89..c3ea501706 100644 --- a/nsjail/run.deno.config.proto +++ b/nsjail/run.deno.config.proto @@ -112,6 +112,7 @@ mount { mandatory: false } +{SHARED_MOUNT} iface_no_lo: true diff --git a/nsjail/run.go.config.proto b/nsjail/run.go.config.proto index 5df40fa99b..f41b98d2c0 100644 --- a/nsjail/run.go.config.proto +++ b/nsjail/run.go.config.proto @@ -126,6 +126,8 @@ mount { mandatory: false } +{SHARED_MOUNT} + envar: "GOPATH=/tmp/.cache/go" envar: "HOME=/tmp/go" diff --git a/nsjail/run.python3.config.proto b/nsjail/run.python3.config.proto index 524e3a9bf0..336383a9be 100644 --- a/nsjail/run.python3.config.proto +++ b/nsjail/run.python3.config.proto @@ -115,6 +115,8 @@ mount { is_bind: true } +{SHARED_MOUNT} + iface_no_lo: true envar: "PYTHONPATH=/tmp/dependencies" diff --git a/openflow.openapi.yaml b/openflow.openapi.yaml index 003d4b0add..6a21cd9656 100644 --- a/openflow.openapi.yaml +++ b/openflow.openapi.yaml @@ -42,6 +42,8 @@ components: $ref: "#/components/schemas/FlowModule" failure_module: $ref: "#/components/schemas/FlowModule" + same_worker: + type: boolean required: - modules