kalugin66/windmill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(backend): flow suspend resume (#522)

Browse Source 
Flow observes `suspend` setting and will wait for resume messages sent for the job before continuing to the next step in a flow.

Adds endpoints under workspaces at `/jobs/<cancel|resume>/<job-uuid>` to either cancel or resume the job with a payload. For POST requests to the endpoint, payload is a JSON document. For GET requests to the endpoints, the payload is a base64url encoded JSON document as the value of the payload query parameter.
This commit is contained in:
sqwishy
2022-09-14 11:46:57 -07:00
committed by GitHub
parent 7dd16de479
commit a4a583f4e3
14 changed files with 2658 additions and 1746 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
backend/Cargo.lock generated
View File

@@ -4481,6 +4481,7 @@ dependencies = [
"async-oauth2",
"async-recursion",
"axum",
"base64 0.11.0",
"chrono",
"console-subscriber",
"cron",

1
backend/Cargo.toml
View File

@@ -58,6 +58,7 @@ async-recursion = "^1"
swc_common = "^0"
swc_ecma_parser = "^0"
swc_ecma_ast = "^0"
base64 = "^0"
unicode-general-category = "^0"
sqlx = { version = "^0", features = ["macros", "offline", "migrate", "uuid", "json", "chrono", "postgres", "runtime-tokio-rustls"]}

5
backend/migrations/20220902211654_job-resume.down.sql Normal file
View File

@@ -0,0 +1,5 @@
DROP TABLE resume_job;
ALTER TABLE queue
DROP COLUMN suspend,
DROP COLUMN suspend_until;

16
backend/migrations/20220902211654_job-resume.up.sql Normal file
View File

@@ -0,0 +1,16 @@
CREATE TABLE resume_job (
id uuid NOT NULL,
job uuid NOT NULL,
flow uuid NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
value JSONB NOT NULL DEFAULT 'null'::jsonb
CHECK (length(value::text) < 10 * 1024),
is_cancel boolean NOT NULL default false,
PRIMARY KEY (id),
FOREIGN KEY (flow) REFERENCES queue(id) ON DELETE CASCADE
);
ALTER TABLE queue
ADD COLUMN suspend INTEGER NOT NULL DEFAULT 0,
ADD COLUMN suspend_until TIMESTAMPTZ;

62
backend/openapi.yaml
View File

@@ -2551,6 +2551,68 @@ paths:
schema:
type: string
/w/{workspace}/jobs/resume/{id}:
get:
summary: resume a job for a suspended flow
operationId: resumeSuspendedJob
tags:
- job
parameters:
- $ref: "#/components/parameters/WorkspaceId"
- $ref: "#/components/parameters/JobId"
- name: payload
in: query
schema:
type: object
responses:
"201"
post:
summary: resume a job for a suspended flow
operationId: resumeSuspendedJob
tags:
- job
parameters:
- $ref: "#/components/parameters/WorkspaceId"
- $ref: "#/components/parameters/JobId"
requestBody:
content:
application/json:
schema:
type: object
responses:
"201"
/w/{workspace}/jobs/cancel/{id}:
get:
summary: cancel a job for a suspended flow
operationId: cancelSuspendedJob
tags:
- job
parameters:
- $ref: "#/components/parameters/WorkspaceId"
- $ref: "#/components/parameters/JobId"
- name: payload
in: query
schema:
type: object
responses:
"201"
post:
summary: cancel a job for a suspended flow
operationId: cancelSuspendedJob
tags:
- job
parameters:
- $ref: "#/components/parameters/WorkspaceId"
- $ref: "#/components/parameters/JobId"
requestBody:
content:
application/json:
schema:
type: object
responses:
"201"
/schedules/preview:
post:
summary: preview schedule

3394
backend/sqlx-data.json
View File

File diff suppressed because it is too large Load Diff

10
backend/src/error.rs
View File

@@ -78,3 +78,13 @@ impl IntoResponse for Error {
.unwrap()
}
}
pub trait OrElseNotFound<T> {
fn or_else_not_found(self, s: impl ToString) -> Result<T>;
}
impl<T> OrElseNotFound<T> for Option<T> {
fn or_else_not_found(self, s: impl ToString) -> Result<T> {
self.ok_or_else(|| Error::NotFound(s.to_string()))
}
}

23
backend/src/fixtures/base.sql
View File

@@ -5,11 +5,12 @@ INSERT INTO workspace
(id, name, owner, domain)
VALUES ('test-workspace', 'test-workspace', 'test-user', null);
CREATE FUNCTION "notify_insert_on_completed_job" ()
RETURNS TRIGGER AS $$
BEGIN
PERFORM pg_notify('insert on completed_job', NEW.id::text);
RETURN new;
RETURN NEW;
END;
$$ LANGUAGE PLPGSQL;
@@ -17,3 +18,23 @@ $$ LANGUAGE PLPGSQL;
AFTER INSERT ON "completed_job"
FOR EACH ROW
EXECUTE FUNCTION "notify_insert_on_completed_job" ();
CREATE FUNCTION "notify_queue" ()
RETURNS TRIGGER AS $$
BEGIN
PERFORM pg_notify('queue', NEW.id::text);
RETURN NEW;
END;
$$ LANGUAGE PLPGSQL;
CREATE TRIGGER "notify_queue_after_insert"
AFTER INSERT ON "queue"
FOR EACH ROW
EXECUTE FUNCTION "notify_queue" ();
CREATE TRIGGER "notify_queue_after_flow_status_update"
AFTER UPDATE ON "queue"
FOR EACH ROW
WHEN (NEW.flow_status IS DISTINCT FROM OLD.flow_status)
EXECUTE FUNCTION "notify_queue" ();

7
backend/src/flows.rs
View File

@@ -159,6 +159,9 @@ pub struct FlowModule {
pub value: FlowModuleValue,
pub stop_after_if: Option<StopAfterIf>,
pub summary: Option<String>,
#[serde(default)]
#[serde(skip_serializing_if = "is_default")]
pub suspend: u16,
}
#[derive(Deserialize, Serialize, Debug, Clone, PartialEq)]
@@ -495,6 +498,7 @@ mod tests {
value: FlowModuleValue::Script { path: "test".to_string() },
stop_after_if: None,
summary: None,
suspend: Default::default(),
},
FlowModule {
input_transforms: HashMap::new(),
@@ -508,6 +512,7 @@ mod tests {
skip_if_stopped: false,
}),
summary: None,
suspend: Default::default(),
},
FlowModule {
input_transforms: [(
@@ -525,6 +530,7 @@ mod tests {
skip_if_stopped: false,
}),
summary: None,
suspend: Default::default(),
},
],
failure_module: Some(FlowModule {
@@ -535,6 +541,7 @@ mod tests {
skip_if_stopped: false,
}),
summary: None,
suspend: Default::default(),
}),
retry: Default::default(),
};

232
backend/src/jobs.rs
View File

@@ -8,6 +8,7 @@
use axum::extract::Host;
use anyhow::Context;
use sql_builder::prelude::*;
use sqlx::{query_scalar, Postgres, Transaction};
use std::collections::HashMap;
@@ -16,23 +17,23 @@ use tracing::instrument;
use crate::{
audit::{audit_log, ActionKind},
db::{UserDB, DB},
error,
error::{to_anyhow, Error},
error::{self, to_anyhow, Error, OrElseNotFound},
flows::FlowValue,
schedule::get_schedule_opt,
scripts::{get_hub_script_by_path, ScriptHash, ScriptLang},
users::{owner_to_token_owner, Authed},
utils::{require_admin, Pagination, StripPath, now_from_db},
utils::{now_from_db, require_admin, Pagination, StripPath},
worker,
worker_flow::{init_flow_status, FlowStatus},
worker_flow::{init_flow_status, FlowStatus, FlowStatusModule},
};
use axum::{
extract::{Extension, Path, Query},
extract::{Extension, FromRequest, Path, Query},
response::{IntoResponse, Response},
routing::{get, post},
Json, Router,
};
use hyper::StatusCode;
use serde::{Deserialize, Serialize};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use serde_json::{Map, Value};
use sql_builder::SqlBuilder;
@@ -68,6 +69,14 @@ pub fn workspaced_service() -> Router {
.route("/getupdate/:id", get(get_job_update))
}
pub fn global_service() -> Router {
Router::new()
.route("/resume/:id", get(resume_suspended_job))
.route("/resume/:id", post(resume_suspended_job))
.route("/cancel/:id", get(cancel_suspended_job))
.route("/cancel/:id", post(cancel_suspended_job))
}
#[derive(Debug, sqlx::FromRow, Serialize, Clone)]
pub struct QueuedJob {
pub workspace_id: String,
@@ -103,18 +112,6 @@ impl QueuedJob {
.map(String::as_str)
.unwrap_or("NO_FLOW_PATH")
}
pub fn parse_raw_flow(&self) -> Option<FlowValue> {
self.raw_flow
.as_ref()
.and_then(|v| serde_json::from_value::<FlowValue>(v.clone()).ok())
}
pub fn parse_flow_status(&self) -> Option<FlowStatus> {
self.flow_status
.as_ref()
.and_then(|v| serde_json::from_value::<FlowStatus>(v.clone()).ok())
}
}
#[derive(Debug, sqlx::FromRow, Serialize)]
@@ -268,7 +265,7 @@ pub async fn run_wait_result_job_by_path(
let script_path = script_path.to_path();
let mut tx = user_db.clone().begin(&authed).await?;
let job_payload = script_path_to_payload(script_path, &mut tx, &w_id).await?;
let scheduled_for = run_query.get_scheduled_for(&mut tx).await?;
let scheduled_for = run_query.get_scheduled_for(&mut tx).await?;
let (uuid, tx) = push(
tx,
@@ -677,7 +674,7 @@ fn list_completed_jobs_query(
if let Some(sk) = &lq.is_skipped {
sqlb.and_where_eq("is_skipped", sk);
}
if let Some(fs) = &lq.is_flow_step {
if let Some(fs) = &lq.is_flow_step {
sqlb.and_where_eq("is_flow_step", fs);
}
if let Some(jk) = &lq.job_kinds {
@@ -959,6 +956,104 @@ pub async fn get_queued_job<'c>(
Ok(r)
}
pub async fn resume_suspended_job(
/* unauthed */
Extension(db): Extension<DB>,
Path((_, job)): Path<(String, Uuid)>,
QueryOrBody(value): QueryOrBody<serde_json::Value>,
) -> error::Result<StatusCode> {
NewResumeJob { job, value, is_cancel: false }
.insert(&db)
.await?;
Ok(StatusCode::CREATED)
}
pub async fn cancel_suspended_job(
/* unauthed */
Extension(db): Extension<DB>,
Path((_, job)): Path<(String, Uuid)>,
QueryOrBody(value): QueryOrBody<serde_json::Value>,
) -> error::Result<StatusCode> {
NewResumeJob { job, value, is_cancel: true }
.insert(&db)
.await?;
Ok(StatusCode::CREATED)
}
struct NewResumeJob {
pub job: Uuid,
pub value: Value,
pub is_cancel: bool,
}
impl NewResumeJob {
async fn insert(self, db: &DB) -> error::Result<()> {
insert_resume_job(db, self).await
}
}
async fn insert_resume_job(db: &DB, new: NewResumeJob) -> error::Result<()> {
let mut tx = db.begin().await?;
let flow = sqlx::query!(
r#"
SELECT id, flow_status, suspend
FROM queue
WHERE id = ( SELECT parent_job FROM queue WHERE id = $1
UNION
SELECT parent_job from completed_job WHERE id = $1 )
FOR UPDATE
"#,
new.job,
)
.fetch_optional(&mut tx)
.await?
.or_else_not_found(new.job)?;
sqlx::query!(
r#"
INSERT INTO resume_job
(id, job, flow, value, is_cancel)
VALUES ($1, $2, $3, $4, $5)
"#,
Uuid::from(Ulid::new()),
new.job,
flow.id,
new.value,
new.is_cancel,
)
.execute(&mut tx)
.await?;
/* If the flow is currently waiting to be resumed (`FlowStatusModule::WaitingForEvents`)
* the suspend column must be set to the number of resume messages waited on.
*
* The flow's queue row is locked in this transaction because to avoid race conditions around
* the suspend column.
* That is, a job needs one event but it hasn't arrived, a worker counts zero events before
* entering WaitingForEvents. Then this message arrives but the job isn't in WaitingForEvents
* yet so the suspend counter isn't updated. Then the job enters WaitingForEvents expecting
* one event to arrive based on the count that is no longer correct. */
if let Some(suspend) = (0 < flow.suspend).then(|| flow.suspend - 1) {
let status =
serde_json::from_value::<FlowStatus>(flow.flow_status.context("no flow status")?)
.context("deserialize flow status")?;
if matches!(status.current_step(), Some(FlowStatusModule::WaitingForEvents { job, .. }) if job == &new.job)
{
sqlx::query!(
"UPDATE queue SET suspend = $1 WHERE id = $2",
if new.is_cancel { 0 } else { suspend },
flow.id,
)
.execute(&mut tx)
.await?;
}
}
tx.commit().await?;
Ok(())
}
#[derive(Serialize)]
#[serde(tag = "type")]
enum Job {
@@ -1369,20 +1464,41 @@ pub async fn add_completed_job(
) -> Result<Uuid, Error> {
let job_id = queued_job.id.clone();
sqlx::query!(
"INSERT INTO completed_job as cj
(workspace_id, id, parent_job, created_by, created_at, started_at, duration_ms, success, \
script_hash, script_path, args, result, logs, \
raw_code, canceled, canceled_by, canceled_reason, job_kind, schedule_path, \
permissioned_as, flow_status, raw_flow, is_flow_step, is_skipped, language)
VALUES ($1, $2, $3, $4, $5, $6, EXTRACT(milliseconds FROM (now() - $6)), $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, \
$18, $19, $20, $21, $22, $23, $24)
"INSERT INTO completed_job AS cj
( workspace_id
, id
, parent_job
, created_by
, created_at
, started_at
, duration_ms
, success
, script_hash
, script_path
, args
, result
, logs
, raw_code
, canceled
, canceled_by
, canceled_reason
, job_kind
, schedule_path
, permissioned_as
, flow_status
, raw_flow
, is_flow_step
, is_skipped
, language )
VALUES ($1, $2, $3, $4, $5, $6, EXTRACT(milliseconds FROM (now() - $6)), $7, $8, $9,\
$10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24)
ON CONFLICT (id) DO UPDATE SET success = $7, result = $11, logs = concat(cj.logs, $12)",
queued_job.workspace_id,
queued_job.id,
queued_job.parent_job,
queued_job.created_by,
queued_job.created_at,
queued_job.started_at,
queued_job.started_at,
success,
queued_job.script_hash.map(|x| x.0),
queued_job.script_path,
@@ -1452,13 +1568,27 @@ pub async fn schedule_again_if_scheduled(
}
pub async fn pull(db: &DB) -> Result<Option<QueuedJob>, crate::Error> {
/* Jobs can be started if they:
* - haven't been started before,
* running = false
* - are flows with a step that needed resume,
* suspend_until is non-null
* and suspend = 0 when the resume messages are received
* or suspend_until <= now() if it has timed out */
let job: Option<QueuedJob> = sqlx::query_as::<_, QueuedJob>(
"UPDATE queue
SET running = true, started_at = now(), last_ping = now()
SET running = true
, started_at = coalesce(started_at, now())
, last_ping = now()
, suspend_until = null
WHERE id IN (
SELECT id
FROM queue
WHERE running = false AND scheduled_for <= now()
WHERE ( running = false
AND scheduled_for <= now())
OR (suspend_until IS NOT NULL
AND ( suspend <= 0
OR suspend_until <= now()))
ORDER BY scheduled_for
FOR UPDATE SKIP LOCKED
LIMIT 1
@@ -1486,3 +1616,45 @@ pub async fn delete_job(db: &DB, w_id: &str, job_id: Uuid) -> Result<(), crate::
tracing::debug!("Job {job_id} deletion was achieved with success: {job_removed}");
Ok(())
}
pub struct QueryOrBody<D>(pub D);
#[axum::async_trait]
impl<D, B> FromRequest<B> for QueryOrBody<D>
where
D: DeserializeOwned,
B: Send + axum::body::HttpBody,
<B as axum::body::HttpBody>::Data: Send,
<B as axum::body::HttpBody>::Error: Into<axum::BoxError>,
{
type Rejection = Response;
async fn from_request(
req: &mut axum::extract::RequestParts<B>,
) -> Result<Self, Self::Rejection> {
return if req.method() == axum::http::Method::GET {
let Query(InPayload { payload }) = Query::from_request(req)
.await
.map_err(IntoResponse::into_response)?;
decode_payload(payload)
.map(QueryOrBody)
.map_err(|err| (StatusCode::BAD_REQUEST, format!("{err:#?}")))
.map_err(IntoResponse::into_response)
} else {
Json::from_request(req)
.await
.map(|Json(v)| QueryOrBody(v))
.map_err(IntoResponse::into_response)
};
#[derive(Deserialize)]
struct InPayload {
payload: String,
}
fn decode_payload<D: DeserializeOwned, T: AsRef<[u8]>>(t: T) -> anyhow::Result<D> {
let vec = base64::decode_config(&t, base64::URL_SAFE).context("invalid base64")?;
serde_json::from_slice(vec.as_slice()).context("invalid json")
}
}
}

1
backend/src/lib.rs
View File

@@ -176,6 +176,7 @@ pub async fn run_server(
.nest("/schedules", schedule::global_service())
.route_layer(from_extractor::<users::Authed>())
.route_layer(from_extractor::<users::Tokened>())
.nest("/w/:workspace_id/jobs", jobs::global_service())
.nest(
"/auth",
users::make_unauthed_service().layer(Extension(argon2)),

469
backend/src/worker.rs
View File

@@ -1586,9 +1586,10 @@ pub async fn restart_zombie_jobs_periodically(
#[cfg(test)]
mod tests {
use sqlx::{postgres::PgListener, query_scalar};
use futures::Stream;
use serde_json::json;
use sqlx::{postgres::PgListener, query_scalar};
use uuid::Uuid;
use crate::{
db::DB,
@@ -1609,6 +1610,8 @@ mod tests {
/// it's important this is unique between tests as there is one prometheus registry and
/// run_worker shouldn't register the same metric with the same worker name more than once.
///
/// this must fit in varchar(50)
fn next_worker_name() -> String {
use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
@@ -1618,7 +1621,15 @@ mod tests {
// will be "main"... The id provides uniqueness & thread_name gives context.
let id = ID.fetch_add(1, SeqCst);
let thread = std::thread::current();
let thread_name = thread.name().unwrap_or("no thread name");
let thread_name = thread
.name()
.map(|s| {
s.len()
.checked_sub(39)
.and_then(|start| s.get(start..))
.unwrap_or(s)
})
.unwrap_or("no thread name");
format!("{id}/{thread_name}")
}
@@ -1640,6 +1651,7 @@ mod tests {
input_transforms: Default::default(),
stop_after_if: Default::default(),
summary: Default::default(),
suspend: Default::default(),
},
FlowModule {
value: FlowModuleValue::ForloopFlow {
@@ -1660,11 +1672,13 @@ mod tests {
.into(),
stop_after_if: Default::default(),
summary: Default::default(),
suspend: Default::default(),
}],
},
input_transforms: Default::default(),
stop_after_if: Default::default(),
summary: Default::default(),
suspend: Default::default(),
},
],
..Default::default()
@@ -1712,13 +1726,13 @@ mod tests {
let result = RunJob::from(job.clone())
.arg("n", json!(123))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert_eq!(json!("last step saw 123"), result);
let result = RunJob::from(job.clone())
.arg("n", json!(-123))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert_eq!(json!(-123), result);
}
@@ -1826,7 +1840,7 @@ func main(derp string) (string, error) {
language: ScriptLang::Go,
}))
.arg("derp", json!("world"))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert_eq!(result, serde_json::json!("hello world"));
@@ -2054,7 +2068,7 @@ def main():
let result = RunJob::from(JobPayload::RawFlow { value: flow.clone(), path: None })
.arg("n", json!(0))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert!(result["from failure module"]["error"]
.as_str()
@@ -2063,7 +2077,7 @@ def main():
let result = RunJob::from(JobPayload::RawFlow { value: flow.clone(), path: None })
.arg("n", json!(1))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert!(result["from failure module"]["error"]
.as_str()
@@ -2072,7 +2086,7 @@ def main():
let result = RunJob::from(JobPayload::RawFlow { value: flow.clone(), path: None })
.arg("n", json!(2))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert!(result["from failure module"]["error"]
.as_str()
@@ -2081,11 +2095,261 @@ def main():
let result = RunJob::from(JobPayload::RawFlow { value: flow.clone(), path: None })
.arg("n", json!(3))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert_eq!(json!({ "l": [0, 1, 2] }), result);
}
mod suspend_resume {
use super::*;
use futures::{Stream, StreamExt};
struct ApiServer {
addr: std::net::SocketAddr,
tx: tokio::sync::oneshot::Sender<()>,
task: tokio::task::JoinHandle<hyper::Result<()>>,
}
impl ApiServer {
async fn start(db: DB) -> Self {
let (tx, rx) = tokio::sync::oneshot::channel::<()>();
let sock = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = sock.local_addr().unwrap();
use crate::jobs::{cancel_suspended_job, resume_suspended_job};
use axum::routing::{get, post};
let router = axum::Router::new()
.route("/w/:workspace/jobs/resume/:id", get(resume_suspended_job))
.route("/w/:workspace/jobs/resume/:id", post(resume_suspended_job))
.route("/w/:workspace/jobs/cancel/:id", get(cancel_suspended_job))
.route("/w/:workspace/jobs/cancel/:id", post(cancel_suspended_job))
.layer(axum::extract::Extension(db));
let serve = axum::Server::from_tcp(sock.into_std().unwrap())
.unwrap()
.serve(router.into_make_service())
.with_graceful_shutdown(async { drop(rx.await) });
let task = tokio::task::spawn(serve);
return Self { addr, tx, task };
}
async fn close(self) -> hyper::Result<()> {
let Self { task, tx, .. } = self;
drop(tx);
task.await.unwrap()
}
}
async fn wait_until_flow_suspends(
flow: Uuid,
mut queue: impl Stream<Item = Uuid> + Unpin,
db: &DB,
) {
loop {
queue.by_ref().find(&flow).await.unwrap();
if query_scalar("SELECT suspend > 0 FROM queue WHERE id = $1")
.bind(flow)
.fetch_one(db)
.await
.unwrap()
{
break;
}
}
}
fn flow() -> FlowValue {
serde_json::from_value(serde_json::json!({
"modules": [{
"input_transform": {
"n": { "type": "javascript", "expr": "flow_input.n", },
"port": { "type": "javascript", "expr": "flow_input.port", },
"op": { "type": "javascript", "expr": "flow_input.op ?? 'resume'", },
},
"value": {
"type": "rawscript",
"language": "deno",
"content": "\
export async function main(n, port, op) {\
const job = Deno.env.get('WM_JOB_ID');
const r = await fetch(
`http://localhost:${port}/w/test-workspace/jobs/${op}/${job}`,\
{\
method: 'POST',\
body: JSON.stringify('from job'),\
headers: { 'content-type': 'application/json' }\
}\
);\
console.log(r);
return n + 1;\
}",
},
"suspend": 1,
}, {
"input_transform": {
"n": { "type": "javascript", "expr": "previous_result", },
"resume": { "type": "javascript", "expr": "resume", },
"resumes": { "type": "javascript", "expr": "resumes", },
},
"value": {
"type": "rawscript",
"language": "deno",
"content": "export function main(n, resume, resumes) { return { n: n + 1, resume, resumes } }"
},
"suspend": 1,
}, {
"input_transform": {
"last": { "type": "javascript", "expr": "previous_result", },
"resume": { "type": "javascript", "expr": "resume", },
"resumes": { "type": "javascript", "expr": "resumes", },
},
"value": {
"type": "rawscript",
"language": "deno",
"content": "export function main(last, resume, resumes) { return { last, resume, resumes } }"
},
}],
}))
.unwrap()
}
#[sqlx::test(fixtures("base"))]
async fn test(db: DB) {
initialize_tracing().await;
let server = ApiServer::start(db.clone()).await;
let port = server.addr.port();
let flow = RunJob::from(JobPayload::RawFlow { value: flow(), path: None })
.arg("n", json!(1))
.arg("port", json!(port))
.push(&db)
.await;
let mut completed = listen_for_completed_jobs(&db).await;
let queue = listen_for_queue(&db).await;
let db_ = db.clone();
in_test_worker(&db, async move {
let db = db_;
wait_until_flow_suspends(flow, queue, &db).await;
/* The first job resumes itself. */
let _first = completed.next().await.unwrap();
/* ... and send a request resume it. */
let second = completed.next().await.unwrap();
/* ImZyb20gdGVzdCIK = base64 "from test" */
reqwest::get(format!(
"http://localhost:{port}/w/test-workspace/jobs/resume/{second}?payload=ImZyb20gdGVzdCIK"
))
.await
.unwrap()
.error_for_status()
.unwrap();
completed.find(&flow).await.unwrap();
})
.await;
server.close().await.unwrap();
let result = completed_job_result(flow, &db).await;
assert_eq!(
json!({
"last": {
"resume": "from job",
"resumes": ["from job"],
"n": 3,
},
"resume": "from test",
"resumes": ["from test"],
}),
result
);
// ensure resumes are cleaned up through CASCADE when the flow is finished
assert_eq!(
0,
query_scalar::<_, i64>("SELECT count(*) FROM resume_job")
.fetch_one(&db)
.await
.unwrap()
);
}
#[sqlx::test(fixtures("base"))]
async fn cancel_from_job(db: DB) {
initialize_tracing().await;
let server = ApiServer::start(db.clone()).await;
let port = server.addr.port();
let result = RunJob::from(JobPayload::RawFlow { value: flow(), path: None })
.arg("n", json!(1))
.arg("op", json!("cancel"))
.arg("port", json!(port))
.run_until_complete(&db)
.await;
server.close().await.unwrap();
assert_eq!(json!("from job"), result);
}
#[sqlx::test(fixtures("base"))]
async fn cancel_after_suspend(db: DB) {
initialize_tracing().await;
let server = ApiServer::start(db.clone()).await;
let port = server.addr.port();
let flow = RunJob::from(JobPayload::RawFlow { value: flow(), path: None })
.arg("n", json!(1))
.arg("port", json!(port))
.push(&db)
.await;
let mut completed = listen_for_completed_jobs(&db).await;
let queue = listen_for_queue(&db).await;
let db_ = db.clone();
in_test_worker(&db, async move {
let db = db_;
wait_until_flow_suspends(flow, queue, &db).await;
/* The first job resumes itself. */
let _first = completed.next().await.unwrap();
/* ... and send a request resume it. */
let second = completed.next().await.unwrap();
reqwest::get(format!(
"http://localhost:{port}/w/test-workspace/jobs/cancel/{second}?payload=ImZyb20gdGVzdCIK"
))
.await
.unwrap()
.error_for_status()
.unwrap();
completed.find(&flow).await.unwrap();
})
.await;
server.close().await.unwrap();
let result = completed_job_result(flow, &db).await;
assert_eq!(json!("from test"), result);
}
}
mod retry {
use super::*;
@@ -2217,7 +2481,7 @@ def main(last, port):
let result = RunJob::from(JobPayload::RawFlow { value: flow(), path: None })
.arg("items", json!(["unused", "unused", "unused"]))
.arg("port", json!(server.addr.port()))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert_eq!(server.close().await, attempts);
@@ -2244,7 +2508,7 @@ def main(last, port):
let result = RunJob::from(JobPayload::RawFlow { value: flow(), path: None })
.arg("items", json!(["unused", "unused", "unused"]))
.arg("port", json!(server.addr.port()))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert_eq!(server.close().await, attempts);
@@ -2276,7 +2540,7 @@ def main(last, port):
let result = RunJob::from(JobPayload::RawFlow { value: flow(), path: None })
.arg("items", json!(["unused", "unused", "unused"]))
.arg("port", json!(server.addr.port()))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert_eq!(server.close().await, attempts);
@@ -2331,7 +2595,7 @@ def main(error, port):
let server = Server::start(responses).await;
let result = RunJob::from(JobPayload::RawFlow { value, path: None })
.arg("port", json!(server.addr.port()))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert_eq!(server.close().await, attempts);
@@ -2362,7 +2626,7 @@ def main(error, port):
.unwrap();
let result = RunJob::from(JobPayload::RawFlow { value, path: None })
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert_eq!(
result,
@@ -2401,14 +2665,14 @@ def main(error, port):
let result = RunJob::from(JobPayload::RawFlow { value: flow.clone(), path: None })
.arg("items", json!([]))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert_eq!(result, serde_json::json!([]));
/* 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::<Vec<_>>()))
.wait_until_complete(&db)
.run_until_complete(&db)
.await;
assert!(matches!(result, Value::Object(_)));
assert!(result["error"]
@@ -2434,8 +2698,7 @@ def main(error, port):
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 {
async fn push(self, db: &DB) -> Uuid {
let RunJob { payload, args } = self;
let (uuid, tx) = push(
@@ -2455,43 +2718,77 @@ def main(error, port):
tx.commit().await.unwrap();
run_worker_until_complete(db, &uuid).await;
uuid
}
query_scalar("SELECT result FROM completed_job WHERE id = $1")
.bind(&uuid)
.fetch_one(db)
.await
.unwrap()
/// push the job, spawn a worker, wait until the job is in completed_job
async fn run_until_complete(self, db: &DB) -> serde_json::Value {
let uuid = self.push(db).await;
let listener = listen_for_completed_jobs(db).await;
in_test_worker(db, listener.find(&uuid)).await;
completed_job_result(uuid, db).await
}
}
async fn run_job_in_new_worker_until_complete(db: &DB, job: JobPayload) -> serde_json::Value {
RunJob::from(job).wait_until_complete(db).await
RunJob::from(job).run_until_complete(db).await
}
async fn run_worker_until_complete(db: &DB, wait_for: &uuid::Uuid) {
let mut listener = PgListener::connect_with(db).await.unwrap();
listener.listen("insert on completed_job").await.unwrap();
/// Start a worker with a timeout and run a future, until the worker quits or we time out.
///
/// Cleans up the worker before resolving.
async fn in_test_worker<Fut: std::future::Future>(
db: &DB,
inner: Fut,
) -> <Fut as std::future::Future>::Output {
let (quit, worker) = spawn_test_worker(db);
let worker = tokio::time::timeout(std::time::Duration::from_secs(19), worker);
tokio::pin!(worker);
let res = tokio::select! {
biased;
res = inner => res,
res = &mut worker => match
res.expect("worker timed out")
.expect("worker panicked") {
_ => panic!("worker quit early"),
},
};
/* ensure the worker quits before we return */
drop(quit);
let _: () = worker
.await
.expect("worker timed out")
.expect("worker panicked");
res
}
fn spawn_test_worker(
db: &DB,
) -> (
tokio::sync::broadcast::Sender<()>,
tokio::task::JoinHandle<()>,
) {
let (tx, rx) = tokio::sync::broadcast::channel(1);
/* drop tx at the end of this block to close the channel and stop the worker */
let worker = {
let timeout = 4_000;
let worker_instance: &str = "test worker instance";
let worker_name: String = next_worker_name();
let i_worker: u64 = Default::default();
let num_workers: u64 = 2;
let ip: &str = Default::default();
let sleep_queue: u64 = DEFAULT_SLEEP_QUEUE / num_workers;
let worker_config = WorkerConfig {
base_internal_url: String::new(),
base_url: String::new(),
disable_nuser: false,
disable_nsjail: false,
keep_job_dir: false,
};
let db = db.to_owned();
let timeout = 4_000;
let worker_instance: &str = "test worker instance";
let worker_name: String = next_worker_name();
let i_worker: u64 = Default::default();
let num_workers: u64 = 2;
let ip: &str = Default::default();
let sleep_queue: u64 = DEFAULT_SLEEP_QUEUE / num_workers;
let worker_config = WorkerConfig {
base_internal_url: String::new(),
base_url: String::new(),
disable_nuser: false,
disable_nsjail: false,
keep_job_dir: false,
};
let future = async move {
run_worker(
&db,
timeout,
@@ -2504,30 +2801,62 @@ def main(error, port):
worker_config,
rx,
)
.await
};
let worker = tokio::time::timeout(std::time::Duration::from_secs(19), worker);
tokio::pin!(worker);
while wait_for
!= &tokio::select! {
biased;
notify = listener.recv() => notify,
res = &mut worker => if res.is_ok() {
panic!("worker quit early")
} else {
panic!("worker timed out")
},
}
.unwrap()
.payload()
.parse::<uuid::Uuid>()
.unwrap()
{}
/* ensure the worker quits before we return */
drop(tx);
worker.await.expect("worker timed out")
(tx, tokio::task::spawn(future))
}
async fn listen_for_completed_jobs(db: &DB) -> impl Stream<Item = Uuid> + Unpin {
listen_for_uuid_on(db, "insert on completed_job").await
}
async fn listen_for_queue(db: &DB) -> impl Stream<Item = Uuid> + Unpin {
listen_for_uuid_on(db, "queue").await
}
async fn listen_for_uuid_on(
db: &DB,
channel: &'static str,
) -> impl Stream<Item = Uuid> + Unpin {
let mut listener = PgListener::connect_with(db).await.unwrap();
listener.listen(channel).await.unwrap();
Box::pin(futures::stream::unfold(
listener,
|mut listener| async move {
let uuid = listener
.try_recv()
.await
.unwrap()
.expect("lost database connection")
.payload()
.parse::<Uuid>()
.expect("invalid uuid");
Some((uuid, listener))
},
))
}
async fn completed_job_result(uuid: Uuid, db: &DB) -> Value {
query_scalar("SELECT result FROM completed_job WHERE id = $1")
.bind(uuid)
.fetch_one(db)
.await
.unwrap()
}
#[axum::async_trait(?Send)]
trait StreamFind: futures::Stream + Unpin + Sized {
async fn find(self, item: &Self::Item) -> Option<Self::Item>
where
for<'l> &'l Self::Item: std::cmp::PartialEq,
{
use futures::{future::ready, StreamExt};
self.filter(|i| ready(i == item)).next().await
}
}
impl<T: futures::Stream + Unpin + Sized> StreamFind for T {}
}

177
backend/src/worker_flow.rs
View File

@@ -22,8 +22,11 @@ use serde_json::{json, Map, Value};
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 = /* six hours */ Duration::from_secs(6 * 60 * 60);
const MAX_RETRY_INTERVAL: Duration = HOURS.saturating_mul(6);
#[derive(Serialize, Deserialize, Debug)]
pub struct FlowStatus {
@@ -53,7 +56,7 @@ pub struct Iterator {
#[serde(tag = "type")]
pub enum FlowStatusModule {
WaitingForPriorSteps,
WaitingForEvent { event: String },
WaitingForEvents { count: u16, job: Uuid },
WaitingForExecutor { job: Uuid },
InProgress { job: Uuid, iterator: Option<Iterator>, forloop_jobs: Option<Vec<Uuid>> },
Success { job: Uuid, forloop_jobs: Option<Vec<Uuid>> },
@@ -69,6 +72,12 @@ impl FlowStatus {
retry: RetryStatus { fail_count: 0, previous_result: None },
}
}
/// current module status ... excluding failure_module
pub fn current_step(&self) -> Option<&FlowStatusModule> {
let i = usize::try_from(self.step).ok()?;
self.modules.get(i)
}
}
#[async_recursion]
@@ -234,14 +243,8 @@ pub async fn update_flow_status_after_job_completion(
false if skip_loop_failures => !is_last_step,
false
if next_retry(
&flow_job
.parse_raw_flow()
.map(|module| module.retry)
.unwrap_or_default(),
&flow_job
.parse_flow_status()
.map(|status| status.retry)
.unwrap_or_default(),
&flow_job.parse_raw_flow_retry().unwrap_or_default(),
&flow_job.parse_flow_status_retry().unwrap_or_default(),
)
.is_some() =>
{
@@ -404,6 +407,7 @@ pub async fn get_step_of_flow_status(db: &DB, id: Uuid) -> error::Result<i32> {
Ok(r)
}
/// resumes should be in order of timestamp ascending, so that more recent are at the end
#[instrument(level = "trace", skip_all)]
async fn transform_input(
flow_args: &Option<serde_json::Value>,
@@ -412,6 +416,7 @@ async fn transform_input(
workspace: &str,
token: &str,
steps: Vec<String>,
resumes: &[Value],
) -> anyhow::Result<Map<String, serde_json::Value>> {
let mut mapped = serde_json::Map::new();
@@ -433,6 +438,11 @@ async fn transform_input(
("params".to_string(), serde_json::json!(mapped)),
("previous_result".to_string(), previous_result),
("flow_input".to_string(), flow_input),
(
"resume".to_string(),
resumes.last().map(|v| json!(v)).unwrap_or_default(),
),
("resumes".to_string(), resumes.clone().into()),
],
Some(EvalCreds { workspace: workspace.to_string(), token: token.to_string() }),
steps.clone(),
@@ -533,6 +543,112 @@ async fn push_next_flow_job(
let mut scheduled_for_o = None;
let mut resume_messages: Vec<Value> = vec![];
/* (suspend / resume), when starting a module, if previous module has a
* non-zero `suspend` value, collect `resume_job`s for the previous module job.
*
* If there aren't enough, try again later. */
if matches!(
&status_module,
FlowStatusModule::WaitingForPriorSteps | FlowStatusModule::WaitingForEvents { .. }
) {
if let Some((count, last)) = needs_resume(&flow, &status) {
let mut tx = db.begin().await?;
/* Lock this row to prevent the suspend column getting out out of sync
* if a resume message arrives after we fetch and count them here.
*
* This only works because jobs::resume_job does the same thing. */
sqlx::query_scalar!(
"SELECT null FROM queue WHERE id = $1 FOR UPDATE",
flow_job.id
)
.fetch_one(&mut tx)
.await
.context("lock flow in queue")?;
let resumes = sqlx::query!(
"SELECT value, is_cancel FROM resume_job WHERE job = $1 ORDER BY created_at ASC",
last
)
.fetch_all(&mut tx)
.await?;
let is_cancelled = resumes
.iter()
.find(|r| r.is_cancel)
.map(|r| r.value.clone());
resume_messages.extend(resumes.into_iter().map(|r| r.value));
if is_cancelled.is_none() && resume_messages.len() >= count as usize {
/* If we are woken up after suspending, last_result will be the flow args, but we
* should use the result from the last job */
if let FlowStatusModule::WaitingForEvents { .. } = &status_module {
last_result =
sqlx::query_scalar!("SELECT result FROM completed_job WHERE id = $1", last)
.fetch_one(&mut tx)
.await?
.context("previous job result")?;
}
/* continue on and run this job! */
tx.commit().await?;
/* not enough messages to do this job, "park"/suspend until there are */
} else if is_cancelled.is_none()
&& matches!(&status_module, FlowStatusModule::WaitingForPriorSteps)
{
sqlx::query(
"
UPDATE queue
SET flow_status = JSONB_SET(flow_status, ARRAY['modules', flow_status->>'step'::text], $1)
, suspend = $2
, suspend_until = now() + $3
WHERE id = $4
",
)
.bind(json!(FlowStatusModule::WaitingForEvents { count, job: last }))
.bind(count as i32)
.bind(30 * MINUTES)
.bind(flow_job.id)
.execute(&mut tx)
.await?;
tx.commit().await?;
return Ok(());
/* cancelled or we're WaitingForEvents but we don't have enough messages (timed out) */
} else {
tx.commit().await?;
let success = false;
let skipped = false;
let logs = if is_cancelled.is_some() {
"Cancelled while waiting to be resumed"
} else {
"Timed out waiting to be resumed"
}
.to_string();
let result = is_cancelled.unwrap_or(json!({ "error": logs }));
let _uuid =
add_completed_job(db, &flow_job, success, skipped, result, logs).await?;
postprocess_queued_job(
false,
flow_job.schedule_path.clone(),
flow_job.script_path.clone(),
&flow_job.workspace_id,
flow_job.id,
db,
)
.await?;
return Ok(());
}
}
}
if matches!(&status_module, FlowStatusModule::Failure { .. }) {
if let Some((fail_count, retry_in)) = next_retry(&flow.retry, &status.retry) {
tracing::debug!(
@@ -642,6 +758,7 @@ async fn push_next_flow_job(
&flow_job.workspace_id,
&token,
steps.to_vec(),
resume_messages.as_slice(),
)
.await?
} else {
@@ -910,6 +1027,28 @@ impl InputTransform {
}
}
impl QueuedJob {
pub fn parse_raw_flow(&self) -> Option<FlowValue> {
self.raw_flow
.as_ref()
.and_then(|v| serde_json::from_value::<FlowValue>(v.clone()).ok())
}
pub fn parse_raw_flow_retry(&self) -> Option<Retry> {
self.parse_raw_flow().map(|module| module.retry)
}
pub fn parse_flow_status(&self) -> Option<FlowStatus> {
self.flow_status
.as_ref()
.and_then(|v| serde_json::from_value::<FlowStatus>(v.clone()).ok())
}
pub fn parse_flow_status_retry(&self) -> Option<RetryStatus> {
self.parse_flow_status().map(|status| status.retry)
}
}
trait IntoArray: Sized {
fn into_array(self) -> Result<Vec<Value>, Self>;
}
@@ -931,3 +1070,21 @@ fn from_now(duration: Duration) -> chrono::DateTime<chrono::Utc> {
.and_then(|d| chrono::Utc::now().checked_add_signed(d))
.unwrap_or(chrono::DateTime::<chrono::Utc>::MAX_UTC)
}
/// returns previous module non-zero suspend count and job
fn needs_resume(flow: &FlowValue, status: &FlowStatus) -> Option<(u16, Uuid)> {
let prev = usize::try_from(status.step)
.ok()
.and_then(|s| s.checked_sub(1))?;
let suspend = flow.modules.get(prev)?.suspend;
if suspend == 0 {
return None;
}
if let &FlowStatusModule::Success { job, .. } = status.modules.get(prev)? {
Some((suspend, job))
} else {
None
}
}

6
openflow.openapi.yaml
View File

@@ -84,6 +84,8 @@ components:
- expr
summary:
type: string
suspend:
type: integer
required:
- input_transforms
- value
@@ -239,8 +241,8 @@ components:
job:
type: string
format: uuid
event:
type: string
count:
type: integer
iterator:
type: object
properties: