use crate::value::Value;
#[derive(Debug, Clone)]
pub struct ExecutionPlanSpec {
pub groups: Vec<Vec<usize>>,
pub concurrency: i64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PolicyKind {
Fail,
Retry,
Continue,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RelationKind {
Single,
Connection,
}
#[derive(Debug, Clone)]
pub struct OpSpec {
pub id: String,
pub parent: Option<usize>,
pub bind_field: Option<String>,
pub relation_kind: Option<RelationKind>,
pub policy: Option<String>,
}
#[derive(Debug, Clone)]
pub enum ExecOutcome {
Ok(Value),
Error(String),
}
#[derive(Debug, Clone)]
enum OpState {
Ok(Value),
Skipped,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PlanFailureCode {
OpFailed,
UnknownPolicy,
InvalidPlan,
}
impl PlanFailureCode {
pub fn as_str(self) -> &'static str {
match self {
PlanFailureCode::OpFailed => "OP_FAILED",
PlanFailureCode::UnknownPolicy => "UNKNOWN_POLICY",
PlanFailureCode::InvalidPlan => "INVALID_PLAN",
}
}
}
#[derive(Debug, Clone)]
pub struct PlanFailure {
pub code: PlanFailureCode,
pub message: String,
}
impl std::fmt::Display for PlanFailure {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}: {}", self.code.as_str(), self.message)
}
}
impl std::error::Error for PlanFailure {}
fn fail<T>(code: PlanFailureCode, message: impl Into<String>) -> Result<T, PlanFailure> {
Err(PlanFailure {
code,
message: message.into(),
})
}
#[derive(Debug)]
pub struct RunResult {
states: Vec<Option<OpState>>,
pub executed: Vec<String>,
pub skipped: Vec<String>,
ops_ref: Vec<(String, Option<RelationKind>)>,
}
impl RunResult {
pub fn final_tree(&self) -> Vec<(String, Value)> {
let mut tree = Vec::new();
for (i, st) in self.states.iter().enumerate() {
match st {
Some(OpState::Ok(v)) => tree.push((self.ops_ref[i].0.clone(), v.clone())),
Some(OpState::Skipped) => tree.push((
self.ops_ref[i].0.clone(),
unproduced_value(self.ops_ref[i].1),
)),
None => {}
}
}
tree
}
}
fn unproduced_value(kind: Option<RelationKind>) -> Value {
if kind == Some(RelationKind::Connection) {
Value::Obj(vec![
("items".into(), Value::Arr(vec![])),
("cursor".into(), Value::Null),
])
} else {
Value::Null
}
}
fn parse_policy(raw: &Option<String>) -> Result<PolicyKind, PlanFailure> {
match raw.as_deref() {
None | Some("fail") => Ok(PolicyKind::Fail),
Some("retry") => Ok(PolicyKind::Retry),
Some("continue") => Ok(PolicyKind::Continue),
Some(other) => fail(
PlanFailureCode::UnknownPolicy,
format!("unknown policy kind: {other}"),
),
}
}
pub fn run_plan<F>(
plan: Option<&ExecutionPlanSpec>,
ops: &[OpSpec],
mut exec: F,
) -> Result<RunResult, PlanFailure>
where
F: FnMut(&OpSpec, Option<&Value>) -> ExecOutcome,
{
let n = ops.len();
let mut states: Vec<Option<OpState>> = vec![None; n];
let mut executed_idx: Vec<usize> = Vec::new();
let mut skipped_idx: Vec<usize> = Vec::new();
let stages: Vec<Vec<usize>> = match plan {
Some(p) => p.groups.clone(),
None => (0..n).map(|i| vec![i]).collect(),
};
validate_coverage(&stages, n)?;
for stage in &stages {
let mut ordered = stage.clone();
ordered.sort_unstable(); for &i in &ordered {
let st = run_op(
i,
&ops[i],
&states,
&mut exec,
&mut executed_idx,
&mut skipped_idx,
)?;
states[i] = Some(st);
}
}
executed_idx.sort_unstable();
skipped_idx.sort_unstable();
Ok(RunResult {
executed: executed_idx.iter().map(|&i| ops[i].id.clone()).collect(),
skipped: skipped_idx.iter().map(|&i| ops[i].id.clone()).collect(),
ops_ref: ops
.iter()
.map(|o| (o.id.clone(), o.relation_kind))
.collect(),
states,
})
}
pub fn run_plan_parallel<F>(
plan: Option<&ExecutionPlanSpec>,
ops: &[OpSpec],
exec: F,
) -> Result<RunResult, PlanFailure>
where
F: Fn(&OpSpec, Option<&Value>) -> ExecOutcome + Sync,
{
let n = ops.len();
let mut states: Vec<Option<OpState>> = vec![None; n];
let mut executed_idx: Vec<usize> = Vec::new();
let mut skipped_idx: Vec<usize> = Vec::new();
let stages: Vec<Vec<usize>> = match plan {
Some(p) => p.groups.clone(),
None => (0..n).map(|i| vec![i]).collect(),
};
let concurrency = plan.map(|p| p.concurrency).unwrap_or(1);
validate_coverage(&stages, n)?;
for stage in &stages {
let mut ordered = stage.clone();
ordered.sort_unstable();
let has_intra_deps = ordered
.iter()
.any(|&i| matches!(ops[i].parent, Some(p) if ordered.contains(&p)));
if concurrency <= 1 || ordered.len() <= 1 || has_intra_deps {
for &i in &ordered {
let st = run_op(
i,
&ops[i],
&states,
&mut |op: &OpSpec, b: Option<&Value>| exec(op, b),
&mut executed_idx,
&mut skipped_idx,
)?;
states[i] = Some(st);
}
continue;
}
let pres: Vec<Preflight> = ordered
.iter()
.map(|&i| preflight_op(&ops[i], &states))
.collect();
let jobs: Vec<usize> = (0..ordered.len())
.filter(|&k| matches!(pres[k], Preflight::Exec { .. }))
.collect();
let workers = (concurrency.max(1) as usize).min(jobs.len());
let cursor = std::sync::atomic::AtomicUsize::new(0);
let slots: Vec<std::sync::Mutex<Option<ExecOutcome>>> = (0..ordered.len())
.map(|_| std::sync::Mutex::new(None))
.collect();
std::thread::scope(|scope| {
for _ in 0..workers {
scope.spawn(|| loop {
let j = cursor.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if j >= jobs.len() {
break;
}
let k = jobs[j];
let i = ordered[k];
let bound = match &pres[k] {
Preflight::Exec { bound, .. } => bound.as_ref(),
_ => unreachable!(),
};
let outcome = exec(&ops[i], bound);
*slots[k].lock().unwrap() = Some(outcome);
});
}
});
for (k, &i) in ordered.iter().enumerate() {
match &pres[k] {
Preflight::InvalidPolicy(msg) => {
return fail(PlanFailureCode::UnknownPolicy, msg.clone())
}
Preflight::Skip => {
skipped_idx.push(i);
states[i] = Some(OpState::Skipped);
}
Preflight::Exec { policy, .. } => {
let outcome = slots[k]
.lock()
.unwrap()
.take()
.expect("dispatched op must have an outcome");
let st = interpret_outcome(
i,
&ops[i],
*policy,
outcome,
&mut executed_idx,
&mut skipped_idx,
)?;
states[i] = Some(st);
}
}
}
}
executed_idx.sort_unstable();
skipped_idx.sort_unstable();
Ok(RunResult {
executed: executed_idx.iter().map(|&i| ops[i].id.clone()).collect(),
skipped: skipped_idx.iter().map(|&i| ops[i].id.clone()).collect(),
ops_ref: ops
.iter()
.map(|o| (o.id.clone(), o.relation_kind))
.collect(),
states,
})
}
enum Preflight {
InvalidPolicy(String),
Skip,
Exec {
policy: PolicyKind,
bound: Option<Value>,
},
}
fn preflight_op(op: &OpSpec, states: &[Option<OpState>]) -> Preflight {
let policy = match parse_policy(&op.policy) {
Ok(p) => p,
Err(e) => return Preflight::InvalidPolicy(e.message), };
let mut bound_value: Option<Value> = None;
if let Some(pidx) = op.parent {
match states.get(pidx).and_then(|s| s.as_ref()) {
Some(OpState::Ok(pv)) => {
if let Some(field) = &op.bind_field {
let bound = pv.obj_get(field);
match bound {
None | Some(Value::Null) => return Preflight::Skip,
Some(b) => bound_value = Some(b.clone()),
}
}
}
_ => {
return Preflight::Skip;
}
}
}
Preflight::Exec {
policy,
bound: bound_value,
}
}
fn interpret_outcome(
i: usize,
op: &OpSpec,
policy: PolicyKind,
outcome: ExecOutcome,
executed_idx: &mut Vec<usize>,
skipped_idx: &mut Vec<usize>,
) -> Result<OpState, PlanFailure> {
executed_idx.push(i);
match outcome {
ExecOutcome::Ok(v) => Ok(OpState::Ok(v)),
ExecOutcome::Error(e) => match policy {
PolicyKind::Fail => fail(
PlanFailureCode::OpFailed,
format!("operation '{}' failed under 'fail' policy: {e}", op.id),
),
PolicyKind::Retry => fail(
PlanFailureCode::OpFailed,
format!(
"operation '{}' failed under 'retry' policy (exhausted): {e}",
op.id
),
),
PolicyKind::Continue => {
skipped_idx.push(i);
executed_idx.pop(); Ok(OpState::Skipped)
}
},
}
}
fn run_op<F>(
i: usize,
op: &OpSpec,
states: &[Option<OpState>],
exec: &mut F,
executed_idx: &mut Vec<usize>,
skipped_idx: &mut Vec<usize>,
) -> Result<OpState, PlanFailure>
where
F: FnMut(&OpSpec, Option<&Value>) -> ExecOutcome,
{
let (policy, bound_value) = match preflight_op(op, states) {
Preflight::InvalidPolicy(msg) => return fail(PlanFailureCode::UnknownPolicy, msg),
Preflight::Skip => {
skipped_idx.push(i);
return Ok(OpState::Skipped);
}
Preflight::Exec { policy, bound } => (policy, bound),
};
let outcome = exec(op, bound_value.as_ref());
interpret_outcome(i, op, policy, outcome, executed_idx, skipped_idx)
}
fn validate_coverage(stages: &[Vec<usize>], n: usize) -> Result<(), PlanFailure> {
let mut seen = std::collections::HashSet::new();
for stage in stages {
for &i in stage {
if i >= n {
return fail(
PlanFailureCode::InvalidPlan,
format!("stage index {i} out of range [0,{n})"),
);
}
if !seen.insert(i) {
return fail(
PlanFailureCode::InvalidPlan,
format!("operation {i} appears in more than one stage"),
);
}
}
}
if seen.len() != n {
return fail(
PlanFailureCode::InvalidPlan,
format!(
"plan does not cover all {n} operations (covered {})",
seen.len()
),
);
}
Ok(())
}