From bae85732ff7c70796c2defcd0430d64dedeb36f7 Mon Sep 17 00:00:00 2001 From: sqwishy Date: Fri, 19 Aug 2022 08:14:46 -0700 Subject: [PATCH] fix(backend): fixes forloop with 257 items only iterates once (#446) usize as u8 will mask out the higher bits, so 0x100 as u8 is 0x00, 0x101 is 0x01 and only iterates once, etc.... refactors push_next_flow_job a bit --- backend/src/jobs.rs | 9 + backend/src/worker.rs | 120 ++++++-- backend/src/worker_flow.rs | 541 +++++++++++++++++++------------------ 3 files changed, 386 insertions(+), 284 deletions(-) diff --git a/backend/src/jobs.rs b/backend/src/jobs.rs index da79d58f92..ce0553f3a1 100644 --- a/backend/src/jobs.rs +++ b/backend/src/jobs.rs @@ -93,6 +93,15 @@ pub struct QueuedJob { pub language: Option, } +impl QueuedJob { + pub fn script_path(&self) -> &str { + self.script_path + .as_ref() + .map(String::as_str) + .unwrap_or("NO_FLOW_PATH") + } +} + #[derive(Debug, sqlx::FromRow, Serialize)] struct CompletedJob { workspace_id: String, diff --git a/backend/src/worker.rs b/backend/src/worker.rs index b178bcfdbf..acab6cbd5a 100644 --- a/backend/src/worker.rs +++ b/backend/src/worker.rs @@ -1237,6 +1237,8 @@ pub async fn restart_zombie_jobs_periodically( mod tests { use sqlx::{postgres::PgListener, query_scalar}; + use serde_json::json; + use crate::{ db::DB, flows::{FlowModule, FlowModuleValue, FlowValue, InputTransform}, @@ -1542,34 +1544,104 @@ def main(): assert_eq!(result, serde_json::json!(9)); } - async fn run_job_in_new_worker_until_complete(db: &DB, job: JobPayload) -> serde_json::Value { - let (uuid, tx) = push( - db.begin().await.unwrap(), - "test-workspace", - job, - /* TODO None doesn't work for args for python because the setup transformer doesn't - * handle it */ - /* args */ - Some(Default::default()), - /* user */ "test-user", - /* permissioned_as */ Default::default(), - /* scheduled_for_o */ None, - /* schedule_path */ None, - /* parent_job */ None, - /* is_flow_step */ false, - ) - .await + #[sqlx::test(fixtures("base"))] + async fn test_iteration(db: DB) { + initialize_tracing().await; + + let flow: FlowValue = serde_json::from_value(serde_json::json!({ + "modules": [{ + "value": { + "type": "forloopflow", + "iterator": { "type": "javascript", "expr": "result.items" }, + "skip_failures": false, + "value": { + "modules": [{ + "input_transform": { + "n": { + "type": "javascript", + "expr": "previous_result.iter.value", + }, + }, + "value": { + "type": "rawscript", + "language": "python3", + "content": "def main(n):\n if 1 < n:\n raise StopIteration(n)", + }, + }], + } + }, + }], + })) .unwrap(); - tx.commit().await.unwrap(); + let result = RunJob::from(JobPayload::RawFlow { value: flow.clone(), path: None }) + .arg("items", json!([])) + .wait_until_complete(&db) + .await; + assert_eq!(result, serde_json::json!([])); - run_worker_until_complete(db, &uuid).await; - - query_scalar("SELECT result FROM completed_job WHERE id = $1") - .bind(&uuid) - .fetch_one(db) - .await + /* Don't actually test that this does 257 jobs or that will take forever. */ + let result = RunJob::from(JobPayload::RawFlow { value: flow.clone(), path: None }) + .arg("items", json!((0..257).collect::>())) + .wait_until_complete(&db) + .await; + assert!(matches!(result, Value::Object(_))); + assert!(result["error"] + .as_str() .unwrap() + .contains("StopIteration: 2")); + } + + struct RunJob { + payload: JobPayload, + args: Map, + } + + impl From for RunJob { + fn from(payload: JobPayload) -> Self { + Self { payload, args: Default::default() } + } + } + + impl RunJob { + fn arg>(mut self, k: S, v: serde_json::Value) -> Self { + self.args.insert(k.into(), v); + self + } + + /// push the job, spawn a worker, wait until the job is in completed_job + async fn wait_until_complete(self, db: &DB) -> serde_json::Value { + let RunJob { payload, args } = self; + + let (uuid, tx) = push( + db.begin().await.unwrap(), + "test-workspace", + payload, + Some(args), + /* user */ "test-user", + /* permissioned_as */ Default::default(), + /* scheduled_for_o */ None, + /* schedule_path */ None, + /* parent_job */ None, + /* is_flow_step */ false, + ) + .await + .unwrap(); + + tx.commit().await.unwrap(); + + run_worker_until_complete(db, &uuid).await; + + query_scalar("SELECT result FROM completed_job WHERE id = $1") + .bind(&uuid) + .fetch_one(db) + .await + .unwrap() + } + } + + async fn run_job_in_new_worker_until_complete(db: &DB, job: JobPayload) -> serde_json::Value { + RunJob::from(job).wait_until_complete(db).await } async fn run_worker_until_complete(db: &DB, wait_for: &uuid::Uuid) { diff --git a/backend/src/worker_flow.rs b/backend/src/worker_flow.rs index 74cf1bbe8b..d3fd596663 100644 --- a/backend/src/worker_flow.rs +++ b/backend/src/worker_flow.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use crate::{ db::DB, error::{self, Error}, - flows::{FlowModuleValue, FlowValue, InputTransform}, + flows::{FlowModule, FlowModuleValue, FlowValue, InputTransform}, jobs::{ add_completed_job, add_completed_job_error, get_queued_job, postprocess_queued_job, push, script_path_to_payload, JobPayload, QueuedJob, @@ -12,7 +12,9 @@ use crate::{ users::create_token_for_owner, worker, }; +use anyhow::Context; use async_recursion::async_recursion; +use futures::TryStreamExt; use serde::{Deserialize, Serialize}; use serde_json::{json, Map, Value}; use tracing::instrument; @@ -27,7 +29,7 @@ pub struct FlowStatus { #[derive(Serialize, Deserialize, Debug, Clone)] pub struct Iterator { - pub index: u8, + pub index: usize, pub itered: Vec, pub args: Map, } @@ -43,6 +45,16 @@ pub enum FlowStatusModule { Failure { job: Uuid, forloop_jobs: Option> }, } +impl FlowStatus { + pub fn new(f: &FlowValue) -> Self { + Self { + step: 0, + modules: vec![FlowStatusModule::WaitingForPriorSteps; f.modules.len()], + failure_module: FlowStatusModule::WaitingForPriorSteps, + } + } +} + #[async_recursion] #[instrument(level = "trace", skip_all)] pub async fn update_flow_status_after_job_completion( @@ -90,10 +102,10 @@ pub async fn update_flow_status_after_job_completion( module_status @ FlowStatusModule::InProgress { iterator: Some(Iterator { index, itered, .. }), .. - } if (index.to_owned() as usize) + 1 < itered.len() && (success || skip_loop_failures) => { + } if *index + 1 < itered.len() && (success || skip_loop_failures) => { (old_status.step, module_status.clone()) } - module_status @ _ => { + module_status => { let forloop_jobs = match module_status { FlowStatusModule::InProgress { forloop_jobs: Some(jobs), .. } => Some(jobs.clone()), _ => None, @@ -149,7 +161,6 @@ pub async fn update_flow_status_after_job_completion( let result = match new_status { FlowStatusModule::Success { forloop_jobs: Some(jobs), .. } => { - use futures::TryStreamExt; let results = sqlx::query_as( " SELECT result @@ -164,9 +175,9 @@ pub async fn update_flow_status_after_job_completion( .bind(w_id) .fetch(&mut tx) .map_ok(|(v,)| v) - .try_collect::>() + .try_collect::>() .await?; - serde_json::json!(results) + json!(results) } _ => result.clone(), }; @@ -261,11 +272,7 @@ async fn compute_stop_early(expr: String, result: serde_json::Value) -> error::R } pub fn init_flow_status(f: &FlowValue) -> FlowStatus { - FlowStatus { - step: 0, - modules: vec![FlowStatusModule::WaitingForPriorSteps; f.modules.len()], - failure_module: FlowStatusModule::WaitingForPriorSteps, - } + FlowStatus::new(f) } pub async fn update_flow_status_in_progress( @@ -314,13 +321,9 @@ async fn transform_input( let mut mapped = serde_json::Map::new(); for (key, val) in input_transform.into_iter() { - match val { - InputTransform::Static { value } => { - mapped.insert(key.to_string(), value.to_owned()); - () - } - _ => (), - }; + if let InputTransform::Static { value } = val { + mapped.insert(key.to_string(), value.to_owned()); + } } for (key, val) in input_transform.into_iter() { @@ -393,50 +396,46 @@ async fn push_next_flow_job( db: &sqlx::Pool, last_result: serde_json::Value, ) -> anyhow::Result<()> { - let flow_status_json = flow_job.flow_status.as_ref().ok_or_else(|| { - Error::InternalErr(format!("not found status for flow job {:?}", flow_job.id)) - })?; - let status = serde_json::from_value::(flow_status_json.to_owned())?; - let i = status.step as usize; + let status: FlowStatus = + serde_json::from_value::(flow_job.flow_status.clone().unwrap_or_default()) + .with_context(|| format!("parse flow status {}", flow_job.id))?; - if flow.modules.len() > i { - let module = &flow.modules[i]; - let mut tx = db.begin().await?; - let job_payload = match &module.value { - FlowModuleValue::Script { path: script_path } => { - script_path_to_payload(script_path, &mut tx, &flow_job.workspace_id).await? - } - FlowModuleValue::RawScript(raw_code) => { - let mut raw_code = raw_code.clone(); - if raw_code.path.is_none() { - raw_code.path = Some(format!( - "{}/{i}", - flow_job - .script_path - .as_ref() - .unwrap_or(&"NO_FLOW_PATH".to_owned()) - )); - } - JobPayload::Code(raw_code) - } - FlowModuleValue::ForloopFlow { value, .. } => JobPayload::RawFlow { - value: *(*value).to_owned(), - path: Some(format!( - "{}/{i}", - flow_job - .script_path - .as_ref() - .unwrap_or(&"NO_FLOW_PATH".to_owned()) - )), - }, - a @ _ => { - tracing::info!("Unrecognized module values {:?}", a); - Err(Error::BadRequest(format!( - "Unrecognized module values {:?}", - a - )))? - } - }; + let i = usize::try_from(status.step) + .with_context(|| format!("invalid module index {}", status.step))?; + + let module: &FlowModule = flow + .modules + .get(i) + .with_context(|| format!("no module at index {}", status.step))?; + + let status_module: FlowStatusModule = status + .modules + .get(i) + .cloned() + .with_context(|| format!("no status at index {}", status.step))?; + + tracing::debug!( + "PUSH: module: {:#?}, status: {:#?}", + module.value, + status_module + ); + + /* Don't evaluate `module.input_transform` after iteration has begun. Instead, args are + * carried through the Iterator by the InProgress variant. */ + + #[rustfmt::skip] + let compute_input_transform = !( matches!(&module.value, FlowModuleValue::ForloopFlow { .. }) + && matches!(&status_module, FlowStatusModule::InProgress { .. })); + + let mut args = if compute_input_transform { + let steps = status + .modules + .iter() + .map(|x| match x { + FlowStatusModule::Success { job, .. } => job.to_string(), + _ => "invalid step status".to_string(), + }) + .collect(); let token = create_token_for_owner( &db, @@ -449,211 +448,233 @@ async fn push_next_flow_job( &flow_job.created_by, ) .await?; - let input_transform = module.input_transform.clone(); - tracing::debug!( - "PUSH: module: {:#?}, status: {:#?}", - module.value, - status.modules[i] - ); - let (compute_input_transform, forloop_args, forloop_iterator) = match &module.value { - FlowModuleValue::ForloopFlow { iterator, .. } => match &status.modules[i] { - FlowStatusModule::WaitingForPriorSteps { .. } => { - let itered = match iterator { - InputTransform::Static { value } => value.clone(), - InputTransform::Javascript { expr } => { - let result = last_result.clone(); - eval_timeout( - expr.to_string(), - [("result".to_string(), result)].into(), - None, - vec![], - ) - .await? - } - }; - let mut args = Map::new(); - let mut iteration_map = Map::new(); - iteration_map.insert( - "index".to_string(), - serde_json::Value::Number(serde_json::Number::from(0)), - ); - iteration_map.insert("value".to_string(), itered[0].clone()); - args.insert("iter".to_string(), serde_json::Value::Object(iteration_map)); - match itered { - serde_json::Value::Array(arr) => { - (true, Some(args), Some((0 as u8, arr, vec![]))) - } - a @ _ => Err(Error::BadRequest(format!( - "Expected an array value, found: {:?}", - a - )))?, - } - } - FlowStatusModule::InProgress { - iterator: Some(Iterator { index, itered, args }), - forloop_jobs: Some(forloop_jobs), - .. - } if index.to_owned() + 1 < itered.len() as u8 => { - let mut args = args.clone(); - let nindex = index.to_owned() + 1; - let mut iteration_map = Map::new(); - iteration_map.insert( - "index".to_string(), - serde_json::Value::Number(serde_json::Number::from(nindex.to_owned())), - ); - iteration_map.insert( - "value".to_string(), - itered[nindex.to_owned() as usize].clone(), - ); - args.insert("iter".to_string(), serde_json::Value::Object(iteration_map)); - ( - false, - Some(args), - Some((nindex, itered.clone(), forloop_jobs.clone())), - ) - } - a @ _ => Err(Error::BadRequest(format!( - "Unrecognized module status for ForloopFlow {:?}", - a - )))?, - }, - _ => (true, None, None), - }; - - if let Some((_, vec, _)) = &forloop_iterator { - if vec.len() == 0 { - let new_job = sqlx::query_as::<_, QueuedJob>(&format!( - "UPDATE queue - SET - flow_status = jsonb_set(jsonb_set(flow_status, '{{modules, {}}}', $1), \ - '{{\"step\"}}', $2) - WHERE id = $3 - RETURNING *", - i - )) - .bind(serde_json::json!(FlowStatusModule::Success { - job: flow_job.id, - forloop_jobs: Some(vec![]) - })) - .bind(serde_json::json!(if flow.modules.len() > i + 1 { - i + 1 - } else { - i - })) - .bind(flow_job.id) - .fetch_one(&mut tx) - .await?; - tx.commit().await?; - if flow.modules.len() > i + 1 { - return Ok( - push_next_flow_job(&new_job, flow, schedule_path, db, json!([])).await?, - ); - } else { - add_completed_job( - db, - &new_job, - true, - false, - json!([]), - "Forloop completed without iteration".to_string(), - ) - .await?; - postprocess_queued_job( - false, - new_job.schedule_path, - new_job.script_path, - &new_job.workspace_id, - new_job.id, - db, - ) - .await?; - return Ok(()); - } - } - } - - let mut args = if compute_input_transform { - let steps = status - .modules - .into_iter() - .map(|x| match x { - FlowStatusModule::Success { job, forloop_jobs: _ } => job.to_string(), - _ => "invalid step status".to_string(), - }) - .collect(); - - transform_input( - &flow_job.args, - last_result.clone(), - &input_transform, - &flow_job.workspace_id, - &token, - steps, - ) - .await? - } else { - Map::new() - }; - let args = if forloop_args.is_some() { - args.extend(forloop_args.unwrap()); - if let Some(flow_args) = &flow_job.args { - match flow_args { - serde_json::Value::Object(obj) => { - for (k, v) in obj { - args.insert(k.to_string(), v.clone()); - } - } - _ => { - (Err(Error::BadRequest(format!( - "Expected an object value, found: {:?}", - flow_args - ))))? - } - } - } - args - } else { - args - }; - - let (uuid, mut tx) = push( - tx, + transform_input( + &flow_job.args, + last_result.clone(), + &module.input_transform, &flow_job.workspace_id, - job_payload, - Some(args.clone()), - &flow_job.created_by, - flow_job.permissioned_as.to_owned(), - None, - schedule_path, - Some(flow_job.id), - true, + &token, + steps, ) - .await?; + .await? + } else { + Map::new() + }; - let new_status = if let Some((index, itered, mut forloop_jobs)) = forloop_iterator { - forloop_jobs.push(uuid.to_owned()); - serde_json::json!(FlowStatusModule::InProgress { - job: uuid, - iterator: Some(Iterator { index: index, itered: itered, args: args }), + /* forloop modules are expected set `iter: { value: Value, index: usize }` as job arguments */ + + let next_loop_status: Option = if let FlowModuleValue::ForloopFlow { + iterator, + .. + } = &module.value + { + match status_module { + FlowStatusModule::WaitingForPriorSteps => { + /* Iterator is an InputTransform, evaluate it into an array. */ + let itered = iterator + .clone() + .evaluate_with(|| vec![("result".to_string(), last_result.clone())]) + .await? + .into_array() + .map_err(|not_array| { + Error::BadRequest(format!("Expected an array value, found: {not_array}")) + })?; + + let first = if let Some(first) = itered.first() { + first + } else { + /* Nothing to iterate, complete immediately and bail. */ + let next_step = i + .checked_add(1) + .filter(|i| (..flow.modules.len()).contains(i)); + + let new_job = sqlx::query_as::<_, QueuedJob>( + r#" + UPDATE queue + SET flow_status = JSONB_SET( + JSONB_SET(flow_status, ARRAY['modules', $1::TEXT], $2), + '{ step }', $3) + WHERE id = $4 + RETURNING * + "#, + ) + .bind(status.step) + .bind(json!(FlowStatusModule::Success { + job: flow_job.id, + forloop_jobs: Some(vec![]) + })) + .bind(json!(next_step.unwrap_or(i))) + .bind(flow_job.id) + .fetch_one(db) + .await?; + + return if next_step.is_some() { + push_next_flow_job(&new_job, flow, schedule_path, db, json!([])).await + } else { + let success = true; + let skipped = false; + let logs = "Forloop completed without iteration".to_string(); + let _uuid = + add_completed_job(db, &new_job, success, skipped, json!([]), logs) + .await?; + postprocess_queued_job( + false, + new_job.schedule_path, + new_job.script_path, + &new_job.workspace_id, + new_job.id, + db, + ) + .await?; + Ok(()) + }; + }; + + args.insert("iter".to_string(), json!({ "index": 0, "value": first })); + + Some(NextLoopStatus { index: 0, itered, forloop_jobs: vec![] }) + } + + FlowStatusModule::InProgress { + iterator: Some(Iterator { itered, index, args: iterator_args }), forloop_jobs: Some(forloop_jobs), - }) + .. + } => { + let (index, next) = index + .checked_add(1) + .and_then(|i| itered.get(i).map(|next| (i, next))) + /* we shouldn't get here because update_flow_status_after_job_completion + * should leave this state if there iteration is complete, but also it should + * be reasonable to just enter a completed state instead of failing, similar to + * iterating an empty list above */ + .with_context(|| format!("could not iterate index {index} of {itered:?}"))?; + + args.extend(iterator_args); + args.insert("iter".to_string(), json!({ "index": 0, "value": next })); + + Some(NextLoopStatus { index, itered, forloop_jobs }) + } + + _ => Err(Error::BadRequest(format!( + "Unrecognized module status for ForloopFlow {status_module:?}" + )))?, + } + } else { + None + }; + + if matches!(&module.value, FlowModuleValue::ForloopFlow { .. }) { + if let Some(value) = &flow_job.args { + value + .as_object() + .ok_or_else(|| { + Error::BadRequest(format!("Expected an object value, found: {value:?}")) + }) + .map(|map| args.extend(map.clone()))?; + } + } + + /* Finally, push the job into the queue */ + + let mut tx = db.begin().await?; + + let job_payload = match &module.value { + FlowModuleValue::Script { path: script_path } => { + script_path_to_payload(script_path, &mut tx, &flow_job.workspace_id).await? + } + FlowModuleValue::RawScript(raw_code) => { + let mut raw_code = raw_code.clone(); + if raw_code.path.is_none() { + raw_code.path = Some(format!("{}/{}", flow_job.script_path(), status.step)); + } + JobPayload::Code(raw_code) + } + FlowModuleValue::ForloopFlow { value, .. } => JobPayload::RawFlow { + value: (**value).clone(), + path: Some(format!("{}/{}", flow_job.script_path(), status.step)), + }, + a @ FlowModuleValue::Flow { .. } => { + tracing::info!("Unrecognized module values {:?}", a); + Err(Error::BadRequest(format!( + "Unrecognized module values {:?}", + a + )))? + } + }; + + let (uuid, mut tx) = push( + tx, + &flow_job.workspace_id, + job_payload, + Some(args.clone()), + &flow_job.created_by, + flow_job.permissioned_as.to_owned(), + None, + schedule_path, + Some(flow_job.id), + true, + ) + .await?; + + let new_status = + if let Some(NextLoopStatus { index, itered, mut forloop_jobs }) = next_loop_status { + forloop_jobs.push(uuid); + + FlowStatusModule::InProgress { + job: uuid, + iterator: Some(Iterator { index, itered, args }), + forloop_jobs: Some(forloop_jobs), + } } else { - serde_json::json!(FlowStatusModule::WaitingForExecutor { job: uuid }) + FlowStatusModule::WaitingForExecutor { job: uuid } }; - sqlx::query(&format!( - "UPDATE queue - SET - flow_status = jsonb_set(flow_status, '{{modules, {}}}', $1) - WHERE id = $2", - i - )) - .bind(new_status) - .bind(flow_job.id) - .execute(&mut tx) - .await?; - tx.commit().await?; + sqlx::query( + "UPDATE queue + SET flow_status = jsonb_set(flow_status, ARRAY['modules', $1::TEXT], $2) + WHERE id = $3", + ) + .bind(status.step) + .bind(json!(new_status)) + .bind(flow_job.id) + .execute(&mut tx) + .await?; + + tx.commit().await?; + + return Ok(()); + + /// Some state about the current/last forloop FlowStatusModule used to initialized the next + /// iteration's FlowStatusModule after pushing a job + struct NextLoopStatus { + index: usize, + itered: Vec, + forloop_jobs: Vec, + } +} + +impl InputTransform { + async fn evaluate_with(self, vars: F) -> anyhow::Result + where + F: FnOnce() -> Vec<(String, Value)>, + { + match self { + InputTransform::Static { value } => Ok(value), + InputTransform::Javascript { expr } => eval_timeout(expr, vars(), None, vec![]).await, + } + } +} + +trait IntoArray: Sized { + fn into_array(self) -> Result, Self>; +} + +impl IntoArray for Value { + fn into_array(self) -> Result, Self> { + match self { + Value::Array(array) => Ok(array), + not_array => Err(not_array), + } } - Ok(()) }