use orbit_ir::WorkflowIr;
use orbit_runtime::engine::{ContractAnnouncement, RegisterWorker};
use orbit_runtime::model::{
CompensationOutcome, ExecutionState, TaskError, UnitStatus, WorkflowStatus,
};
use orbit_runtime::{Engine, EngineConfig, Store};
use serde_json::json;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
const SRC: &str = r#"
task "a" { version "1.0.0" input { n: Int } output { n: Int } }
task "undo" { version "1.0.0" input { n: Int } output { ok: Bool } }
workflow Child { version "1.0.0" input { n: Int } output { n: Int }
step c { worker "a" with { n: input.n } } return { n: c.n } }
workflow Main {
version "1.0.0"
input { ns: List<Int> }
step each {
for n in input.ns
worker "a"
with { n: n }
retry { attempts: 2, initial_delay: 1s }
compensate { worker "undo" with { n: each.n } }
}
step child { workflow "Child" after each with { n: len(each) } }
step last { worker "a" after child with { n: child.n } }
}
"#;
struct Db {
admin: String,
name: String,
url: String,
}
impl Db {
fn create(test: &str) -> Option<Db> {
let full = std::env::var("ORBIT_TEST_POSTGRES_URL").ok()?;
let (admin, query) = match full.split_once('?') {
Some((a, q)) => (a.to_string(), format!("?{q}")),
None => (full.clone(), String::new()),
};
let name = format!("orbit_test_{}_{test}", std::process::id());
let mut c = postgres::Client::connect(&admin, postgres::NoTls)
.expect("connect to ORBIT_TEST_POSTGRES_URL");
c.batch_execute(&format!("DROP DATABASE IF EXISTS {name}"))
.unwrap();
c.batch_execute(&format!("CREATE DATABASE {name}")).unwrap();
let base = admin.rsplit_once('/').map(|(b, _)| b).unwrap_or(&admin);
let url = format!("{base}/{name}{query}");
Some(Db { admin, name, url })
}
}
impl Drop for Db {
fn drop(&mut self) {
if let Ok(mut c) = postgres::Client::connect(&self.admin, postgres::NoTls) {
let _ = c.batch_execute(&format!(
"DROP DATABASE IF EXISTS {} WITH (FORCE)",
self.name
));
}
}
}
fn engine(url: &str, clock: &Arc<AtomicU64>) -> Engine {
let c = clock.clone();
Engine::open(
Store::connect(url).unwrap(),
EngineConfig::default(),
Box::new(move || c.load(Ordering::SeqCst)),
)
.unwrap()
}
fn workflows() -> Vec<WorkflowIr> {
let c = orbit_compiler::compile_str("pg.orbit", SRC);
assert!(!c.has_errors(), "{}", c.render_diagnostics(false));
c.workflows
}
fn worker(e: &mut Engine) -> String {
e.register_worker(RegisterWorker {
worker_id: "pg-worker".into(),
runtime: "rust".into(),
sdk: None,
contracts: ["a", "undo"]
.iter()
.map(|t| ContractAnnouncement {
task_name: t.to_string(),
contract_version: "1.0.0".into(),
})
.collect(),
})
.unwrap()
.worker_session_id
}
macro_rules! db_or_skip {
($name:literal) => {
match Db::create($name) {
Some(db) => db,
None => {
eprintln!("skipping: set ORBIT_TEST_POSTGRES_URL to run PostgreSQL tests");
return;
}
}
};
}
#[test]
fn full_lifecycle_survives_restarts_on_postgres() {
let db = db_or_skip!("lifecycle");
let clock = Arc::new(AtomicU64::new(1_000_000));
let mut e = engine(&db.url, &clock);
for wf in workflows() {
assert!(e.register_definition(wf.clone()).unwrap());
assert!(
!e.register_definition(wf).unwrap(),
"idempotent re-registration"
);
}
let id = e
.start(
"Main",
None,
json!({ "ns": [1, 2] }),
Some("order-1".into()),
)
.unwrap()
.execution_id;
assert!(
!e.start(
"Main",
None,
json!({ "ns": [1, 2] }),
Some("order-1".into())
)
.unwrap()
.created
);
let w = worker(&mut e);
let t0 = e.poll(&w).unwrap().unwrap();
let t1 = e.poll(&w).unwrap().unwrap();
e.fail(&t0.token, TaskError::new("FLAKY", "try again", true))
.unwrap();
e.complete(&t1.token, json!({ "n": 20 })).unwrap();
let before = serde_json::to_value(e.execution(&id).unwrap()).unwrap();
drop(e);
let mut e = engine(&db.url, &clock);
assert_eq!(
serde_json::to_value(e.execution(&id).unwrap()).unwrap(),
before
);
clock.fetch_add(1_000, Ordering::SeqCst);
e.tick().unwrap();
let w = worker(&mut e);
let retry = e.poll(&w).unwrap().unwrap();
assert_eq!((retry.step_id.as_str(), retry.attempt), ("each[0]", 2));
e.complete(&retry.token, json!({ "n": 10 })).unwrap();
e.tick().unwrap();
let child_task = e.poll(&w).unwrap().unwrap();
assert_eq!(child_task.execution_id, "order-1.child");
e.fail(&child_task.token, TaskError::new("BROKEN", "no", false))
.unwrap();
e.tick().unwrap();
let mut compensated = vec![];
while let Some(t) = e.poll(&w).unwrap() {
compensated.push(t.step_id.clone());
e.complete(&t.token, json!({ "ok": true })).unwrap();
}
assert_eq!(
compensated,
vec!["compensate:each[0]", "compensate:each[1]"]
);
let s = e.execution(&id).unwrap();
assert_eq!(s.status, WorkflowStatus::Failed);
assert_eq!(s.compensation, Some(CompensationOutcome::Completed));
assert_eq!(s.units["each"].status(), UnitStatus::Succeeded);
let history = e.history(&id).unwrap();
assert!(history
.iter()
.enumerate()
.all(|(i, h)| h.seq == i as u64 + 1));
let replayed = ExecutionState::replay(&id, &history);
assert_eq!(
serde_json::to_value(replayed).unwrap(),
serde_json::to_value(&s).unwrap()
);
let listed = e.list_executions(Some("FAILED"), Some("Main"), 10).unwrap();
assert_eq!(listed.len(), 1);
assert!(e
.store_stats()
.unwrap()
.iter()
.any(|(st, n)| st == "FAILED" && *n >= 1));
assert_eq!(e.workers().unwrap()[0].record.worker_id, "pg-worker");
}
#[test]
fn definitions_are_immutable_on_postgres() {
let db = db_or_skip!("immutable");
let clock = Arc::new(AtomicU64::new(1));
let mut e = engine(&db.url, &clock);
let wfs = workflows();
e.register_definition(wfs[0].clone()).unwrap();
let mut changed = wfs[0].clone();
changed.description = Some("different".into());
let changed = changed.seal();
assert_eq!(
e.register_definition(changed).unwrap_err().code,
"DEFINITION_CONFLICT"
);
}