chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

feat(workflow): add scheduled/delayed steps and conditional branching (#16980)

Browse Source
This commit is contained in:
Simo Lin
2026-01-12 15:06:12 -08:00
committed by GitHub
parent 6620548fd8
commit e0ac559ae1
3 changed files with 1025 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

114
sgl-model-gateway/src/workflow/definition.rs
View File

@@ -7,9 +7,11 @@ use std::{
time::Duration,
};
use chrono::{DateTime, Utc};
use super::{
executor::StepExecutor,
types::{FailureAction, RetryPolicy, StepId, WorkflowData, WorkflowId},
types::{FailureAction, RetryPolicy, StepId, WorkflowContext, WorkflowData, WorkflowId},
};
/// Errors that can occur during workflow validation
@@ -32,7 +34,17 @@ pub struct StepDefinition<D: WorkflowData> {
pub retry_policy: Option<RetryPolicy>,
pub timeout: Option<Duration>,
pub on_failure: FailureAction,
/// Dependencies that must ALL complete before this step runs
pub depends_on: Vec<StepId>,
/// Dependencies where ANY completing triggers this step (used with depends_on)
pub depends_on_any: Vec<StepId>,
/// Delay before starting the step (after dependencies satisfied)
pub delay: Option<Duration>,
/// Run step at or after this time (after dependencies satisfied)
pub scheduled_at: Option<DateTime<Utc>>,
/// Condition to evaluate; if false, step is skipped
#[allow(clippy::type_complexity)]
pub run_if: Option<Arc<dyn Fn(&WorkflowContext<D>) -> bool + Send + Sync>>,
}
impl<D: WorkflowData> fmt::Debug for StepDefinition<D> {
@@ -44,6 +56,10 @@ impl<D: WorkflowData> fmt::Debug for StepDefinition<D> {
.field("timeout", &self.timeout)
.field("on_failure", &self.on_failure)
.field("depends_on", &self.depends_on)
.field("depends_on_any", &self.depends_on_any)
.field("delay", &self.delay)
.field("scheduled_at", &self.scheduled_at)
.field("run_if", &self.run_if.as_ref().map(|_| "<condition>"))
.finish_non_exhaustive()
}
}
@@ -62,6 +78,10 @@ impl<D: WorkflowData> StepDefinition<D> {
timeout: None,
on_failure: FailureAction::FailWorkflow,
depends_on: Vec::new(),
depends_on_any: Vec::new(),
delay: None,
scheduled_at: None,
run_if: None,
}
}
@@ -81,12 +101,74 @@ impl<D: WorkflowData> StepDefinition<D> {
}
/// Set dependencies for this step.
/// The step will only run after all specified dependencies have completed successfully.
/// The step will only run after ALL specified dependencies have completed successfully.
/// Empty slice means no dependencies - step can run immediately in parallel with others.
pub fn depends_on(mut self, deps: &[&str]) -> Self {
self.depends_on = deps.iter().map(|s| StepId::new(*s)).collect();
self
}
/// Set "any of" dependencies for this step.
/// The step will run when ANY of these dependencies complete (in addition to
/// all `depends_on` dependencies).
///
/// **Combined semantics** (when both `depends_on` and `depends_on_any` are set):
/// - ALL `depends_on` must complete successfully, AND
/// - AT LEAST ONE `depends_on_any` must complete successfully
///
/// **Failure handling**:
/// - For `depends_on`: if ANY fails, this step is blocked
/// - For `depends_on_any`: only blocked if ALL fail (since we only need one)
///
/// **Skipped dependencies**: A skipped dependency (e.g., via `run_if`) counts as
/// "completed" for dependency satisfaction purposes.
pub fn depends_on_any(mut self, deps: &[&str]) -> Self {
self.depends_on_any = deps.iter().map(|s| StepId::new(*s)).collect();
self
}
/// Set a delay before starting the step (after dependencies are satisfied).
///
/// If both `delay` and `scheduled_at` are set, the step will wait for the
/// scheduled time AND THEN wait for the delay duration (they stack).
pub fn with_delay(mut self, delay: Duration) -> Self {
self.delay = Some(delay);
self
}
/// Schedule the step to run at or after the specified time.
/// The step will wait until this time even if dependencies are satisfied earlier.
///
/// If the scheduled time is in the past when the step becomes ready, it will
/// proceed immediately (a debug log is emitted).
///
/// If both `delay` and `scheduled_at` are set, the step will wait for the
/// scheduled time AND THEN wait for the delay duration (they stack).
pub fn scheduled_at(mut self, time: DateTime<Utc>) -> Self {
self.scheduled_at = Some(time);
self
}
/// Set a condition for running this step.
/// If the condition returns false, the step is skipped.
///
/// **Skipped step semantics**: When a step is skipped:
/// - It is marked as `StepStatus::Skipped`
/// - Downstream steps that depend on it (via `depends_on`) will consider it satisfied
/// - This allows conditional branches without blocking the workflow
///
/// **Error handling**: If the context cannot be retrieved to evaluate the condition,
/// the step will **fail** (not proceed blindly). This is a safety measure.
///
/// **Note**: The condition closure is not serializable, so workflows with `run_if`
/// cannot be persisted and resumed from external storage.
pub fn run_if<F>(mut self, condition: F) -> Self
where
F: Fn(&WorkflowContext<D>) -> bool + Send + Sync + 'static,
{
self.run_if = Some(Arc::new(condition));
self
}
}
/// Complete workflow definition
@@ -166,7 +248,7 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
let steps_map: HashMap<&StepId, &StepDefinition<D>> =
self.steps.iter().map(|s| (&s.id, s)).collect();
// Check all dependencies exist
// Check all dependencies exist (both depends_on and depends_on_any)
for step in &self.steps {
for dep in &step.depends_on {
if !steps_map.contains_key(dep) {
@@ -176,9 +258,17 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
});
}
}
for dep in &step.depends_on_any {
if !steps_map.contains_key(dep) {
return Err(ValidationError::MissingDependency {
step: step.id.clone(),
dependency: dep.clone(),
});
}
}
}
// Check for cycles using DFS
// Check for cycles using DFS (considers both dependency types)
let mut visited = HashSet::new();
let mut rec_stack = HashSet::new();
@@ -191,6 +281,7 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
}
// Build reverse dependency map: for each step, which steps depend on it?
// Include both depends_on and depends_on_any
self.reverse_deps.clear();
for (idx, step) in self.steps.iter().enumerate() {
for dep_id in &step.depends_on {
@@ -199,14 +290,21 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
.or_default()
.push(idx);
}
for dep_id in &step.depends_on_any {
self.reverse_deps
.entry(dep_id.clone())
.or_default()
.push(idx);
}
}
// Cache indices of steps with no dependencies (can start immediately)
// A step with only depends_on_any still needs at least one to complete
self.initial_step_indices = self
.steps
.iter()
.enumerate()
.filter(|(_, s)| s.depends_on.is_empty())
.filter(|(_, s)| s.depends_on.is_empty() && s.depends_on_any.is_empty())
.map(|(i, _)| i)
.collect();
@@ -231,12 +329,18 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
rec_stack.insert(step_id);
// O(1) lookup instead of linear search
// Check both depends_on and depends_on_any for cycles
if let Some(step) = steps_map.get(step_id) {
for dep in &step.depends_on {
if Self::has_cycle(dep, steps_map, visited, rec_stack) {
return true;
}
}
for dep in &step.depends_on_any {
if Self::has_cycle(dep, steps_map, visited, rec_stack) {
return true;
}
}
}
rec_stack.remove(step_id);

299
sgl-model-gateway/src/workflow/engine.rs
View File

@@ -35,6 +35,9 @@ struct StepTracker {
failed: HashSet<StepId>,
skipped: HashSet<StepId>,
running: HashSet<StepId>,
/// Steps waiting for delay/scheduled_at: maps step INDEX to ready time
/// Using index instead of StepId for O(w) iteration in the main loop
waiting_until: HashMap<usize, std::time::Instant>,
}
impl StepTracker {
@@ -42,22 +45,59 @@ impl StepTracker {
self.completed.len() + self.failed.len() + self.skipped.len()
}
fn is_step_processable(&self, step_id: &StepId) -> bool {
fn is_step_processable(&self, step_id: &StepId, step_idx: usize) -> bool {
!self.completed.contains(step_id)
&& !self.failed.contains(step_id)
&& !self.skipped.contains(step_id)
&& !self.running.contains(step_id)
&& !self.waiting_until.contains_key(&step_idx)
}
/// Get indices of waiting steps that are now ready to run - O(w) where w = waiting count
fn get_ready_waiting_indices(&self) -> Vec<usize> {
let now = std::time::Instant::now();
self.waiting_until
.iter()
.filter(|(_, &ready_at)| now >= ready_at)
.map(|(&idx, _)| idx)
.collect()
}
/// Mark a step as waiting until a specific time (by index)
fn set_waiting(&mut self, step_idx: usize, ready_at: std::time::Instant) {
self.waiting_until.insert(step_idx, ready_at);
}
/// Clear waiting status for a step (it's now ready to run)
fn clear_waiting(&mut self, step_idx: usize) {
self.waiting_until.remove(&step_idx);
}
/// Check if ALL dependencies are satisfied (completed or skipped)
fn are_dependencies_satisfied(&self, depends_on: &[StepId]) -> bool {
depends_on
.iter()
.all(|dep| self.completed.contains(dep) || self.skipped.contains(dep))
}
/// Check if ANY dependency is satisfied (completed or skipped)
/// Returns true if the list is empty (no "any" dependencies)
fn is_any_dependency_satisfied(&self, depends_on_any: &[StepId]) -> bool {
depends_on_any.is_empty()
|| depends_on_any
.iter()
.any(|dep| self.completed.contains(dep) || self.skipped.contains(dep))
}
/// Check if ANY dependency has failed (for depends_on - blocks if any fail)
fn has_failed_dependency(&self, depends_on: &[StepId]) -> bool {
depends_on.iter().any(|dep| self.failed.contains(dep))
}
/// Check if ALL dependencies have failed (for depends_on_any - blocks only if all fail)
fn have_all_any_deps_failed(&self, depends_on_any: &[StepId]) -> bool {
!depends_on_any.is_empty() && depends_on_any.iter().all(|dep| self.failed.contains(dep))
}
}
/// Fixed backoff that returns the same delay every time
@@ -264,6 +304,35 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
self.active_workflows.fetch_sub(1, Ordering::Release);
}
/// Calculate how long a step needs to wait based on delay and/or scheduled_at.
/// Returns None if no waiting is needed.
fn calculate_wait_duration(step: &StepDefinition<D>) -> Option<Duration> {
let now = Utc::now();
// Calculate wait time for scheduled_at
let schedule_wait = step.scheduled_at.and_then(|scheduled_time| {
if now < scheduled_time {
(scheduled_time - now).to_std().ok()
} else {
// Scheduled time is in the past
tracing::debug!(
step_id = %step.id,
scheduled_time = %scheduled_time,
"Step scheduled_at is in the past, proceeding immediately"
);
None
}
});
// Combine delay and schedule_wait (both apply if both are set)
match (step.delay, schedule_wait) {
(Some(delay), Some(schedule)) => Some(delay + schedule),
(Some(delay), None) => Some(delay),
(None, Some(schedule)) => Some(schedule),
(None, None) => None,
}
}
/// Create a guard that decrements active_workflows on drop.
/// This ensures the count is decremented even if a task panics.
fn active_workflow_guard(&self) -> ActiveWorkflowGuard {
@@ -428,22 +497,95 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
return Ok(());
}
// Find ready steps from pending_check (not all steps)
let (ready_step_indices, total_processed, running_count) = {
// Phase 0: Drain any pending completion signals to ensure dependents are added
// This prevents race conditions where tasks finish but their signals aren't processed
while let Ok((step_id, result)) = rx.try_recv() {
tracing::debug!(
step_id = %step_id,
result = ?result,
"Step completed (startup drain)"
);
if matches!(result, StepResult::Success | StepResult::Skip) {
for &dep_idx in definition.get_dependent_indices(&step_id) {
pending_check.push_back(dep_idx);
}
}
}
// Phase 1: Check waiting steps + deps-ready steps + blocked detection
// Single lock acquisition for all read operations
let (
newly_ready_from_wait,
deps_ready_indices,
total_processed,
current_running,
current_waiting,
) = {
let t = tracker.read();
// Only check steps in pending_check, not all steps
let ready: Vec<usize> = pending_check
// O(w) iteration over waiting_until keys instead of O(n) over all steps
let wait_ready: Vec<usize> = t.get_ready_waiting_indices();
// Check pending_check for dependency-satisfied steps
let deps_ready: Vec<usize> = pending_check
.drain(..)
.filter(|&idx| {
let step = &definition.steps[idx];
t.is_step_processable(&step.id)
t.is_step_processable(&step.id, idx)
&& t.are_dependencies_satisfied(&step.depends_on)
&& t.is_any_dependency_satisfied(&step.depends_on_any)
&& !t.has_failed_dependency(&step.depends_on)
&& !t.have_all_any_deps_failed(&step.depends_on_any)
})
.collect();
(ready, t.total_processed(), t.running.len())
(
wait_ready,
deps_ready,
t.total_processed(),
t.running.len(),
t.waiting_until.len(),
)
};
// Phase 2: Process waiting and deps-ready steps, update waiting_until
// Single write lock for all mutations
// Returns (ready_to_launch, steps_added_to_waiting)
let (ready_to_launch, steps_added_to_waiting) = {
let now = std::time::Instant::now();
let mut t = tracker.write();
let mut added_to_waiting = 0usize;
// Clear waiting status for ready steps
for &idx in &newly_ready_from_wait {
t.clear_waiting(idx);
}
// Collect steps ready to launch
let mut ready = newly_ready_from_wait;
for idx in deps_ready_indices {
let step = &definition.steps[idx];
let wait_duration = Self::calculate_wait_duration(step);
if let Some(duration) = wait_duration {
if duration > Duration::ZERO {
// Step needs to wait - add to waiting_until
let ready_at = now + duration;
tracing::debug!(
step_id = %step.id,
wait_ms = duration.as_millis(),
"Step waiting for delay/schedule"
);
t.set_waiting(idx, ready_at);
added_to_waiting += 1;
continue;
}
}
ready.push(idx);
}
(ready, added_to_waiting)
};
// Check if we're done
@@ -451,13 +593,23 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
break;
}
// Handle blocked workflow (no ready steps, none running, but work remains)
if ready_step_indices.is_empty() && running_count == 0 && pending_check.is_empty() {
// Handle blocked workflow (no ready steps, none running/waiting, but work remains)
// Use current_running/current_waiting from Phase 1, adjusted for Phase 2 changes.
// current_running may be stale-high (tasks completed since), but that's safe -
// we'd just do an extra loop iteration. current_waiting needs adjustment for
// steps we just added to waiting in Phase 2.
let effective_waiting = current_waiting + steps_added_to_waiting;
if ready_to_launch.is_empty()
&& current_running == 0
&& effective_waiting == 0
&& pending_check.is_empty()
{
let failed_step = tracker.read().failed.iter().next().cloned();
let error_message = if failed_step.is_some() {
"Workflow failed due to step dependency failure".to_string()
// Use &'static str to avoid allocation in common error paths
let error_message: &'static str = if failed_step.is_some() {
"Workflow failed due to step dependency failure"
} else {
"Workflow deadlocked: no steps ready and none running. This may indicate a scheduler bug.".to_string()
"Workflow deadlocked: no steps ready and none running"
};
self.state_store
@@ -470,19 +622,23 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
instance_id,
failed_step: failed_step
.unwrap_or_else(|| StepId::new("internal_scheduler")),
error: error_message,
error: error_message.to_string(),
})
.await;
return Ok(());
}
// Launch ready steps in parallel
let mut tasks_launched = 0;
for step_idx in ready_step_indices {
let step = &definition.steps[step_idx];
tracker.write().running.insert(step.id.clone());
tasks_launched += 1;
let tasks_launched = ready_to_launch.len();
if tasks_launched > 0 {
let mut t = tracker.write();
for &idx in &ready_to_launch {
t.running.insert(definition.steps[idx].id.clone());
}
}
for step_idx in ready_to_launch {
let step = &definition.steps[step_idx];
let engine = self.clone_for_execution();
let def = Arc::clone(&definition);
let step_id = step.id.clone();
@@ -491,6 +647,76 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
tokio::spawn(async move {
let step = &def.steps[step_idx];
// Evaluate run_if condition if present
if let Some(ref condition) = step.run_if {
match engine.state_store.get_context(instance_id).await {
Ok(ctx) => {
if !condition(&ctx) {
// Condition is false - skip this step
tracing::debug!(
step_id = %step_id,
"Step skipped due to run_if condition"
);
// Update tracker and send skip signal
{
let mut t = tracker.write();
t.running.remove(&step_id);
t.skipped.insert(step_id.clone());
}
// Update state store
let _ = engine
.state_store
.update(instance_id, |s| {
if let Some(step_state) =
s.step_states.get_mut(&step_id)
{
step_state.status = StepStatus::Skipped;
}
})
.await;
// Send skip signal
let _ = tx.try_send((step_id, StepResult::Skip));
return;
}
}
Err(e) => {
// Failed to get context - fail the step rather than proceeding blindly
tracing::error!(
step_id = %step_id,
error = ?e,
"Failed to get context for run_if evaluation, failing step"
);
// Update tracker
{
let mut t = tracker.write();
t.running.remove(&step_id);
t.failed.insert(step_id.clone());
}
// Update state store
let _ = engine
.state_store
.update(instance_id, |s| {
if let Some(step_state) = s.step_states.get_mut(&step_id) {
step_state.status = StepStatus::Failed;
step_state.last_error =
Some(format!("run_if context error: {e}"));
}
})
.await;
// Send failure signal
let _ = tx.try_send((step_id, StepResult::Failure));
return;
}
}
}
let result = engine
.execute_step_with_retry(instance_id, step, &def)
.await;
@@ -568,8 +794,17 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
});
}
let should_wait = tasks_launched > 0 || !tracker.read().running.is_empty();
if should_wait {
// Single lock read for both checks
let (has_running, has_waiting) = {
let t = tracker.read();
(
tasks_launched > 0 || !t.running.is_empty(),
!t.waiting_until.is_empty(),
)
};
if has_running {
// Wait for a step to complete; Phase 0 will drain any others
if let Some((completed_step_id, result)) = rx.recv().await {
tracing::debug!(
step_id = %completed_step_id,
@@ -585,6 +820,28 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
}
}
}
} else if has_waiting {
// No running tasks, but some steps are waiting for their delay/schedule
// Calculate how long to sleep based on the soonest waiting step
let sleep_duration = {
let t = tracker.read();
let now = std::time::Instant::now();
t.waiting_until
.values()
.filter_map(|&ready_at| {
if ready_at > now {
Some(ready_at - now)
} else {
None
}
})
.min()
.unwrap_or(Duration::from_millis(10))
};
// Sleep until the next step is ready (with a cap to allow cancellation checks)
let capped_sleep = sleep_duration.min(Duration::from_millis(100));
tokio::time::sleep(capped_sleep).await;
}
}
@@ -603,7 +860,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
.publish(WorkflowEvent::WorkflowFailed {
instance_id,
failed_step: step.clone(),
error: "One or more steps failed".to_string(),
error: "One or more steps failed".into(),
})
.await;
} else {

638
sgl-model-gateway/tests/workflow_test.rs
View File

@@ -68,6 +68,18 @@ impl StepExecutor<TestWorkflowData> for AlwaysSucceedStep {
}
}
struct AlwaysFailStep;
#[async_trait::async_trait]
impl StepExecutor<TestWorkflowData> for AlwaysFailStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
Ok(StepResult::Failure)
}
}
#[tokio::test]
async fn test_simple_workflow_execution() {
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
@@ -703,3 +715,629 @@ fn test_dag_validation_valid_workflow() {
let result = workflow.validate();
assert!(result.is_ok());
}
// ============================================================================
// Scheduled/Delayed Steps Tests (#24)
// ============================================================================
#[tokio::test]
async fn test_step_delay() {
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
// Create a workflow with a 100ms delay
let workflow = WorkflowDefinition::new("delay_workflow", "Delay Test").add_step(
StepDefinition::new("delayed_step", "Delayed Step", Arc::new(AlwaysSucceedStep))
.with_delay(Duration::from_millis(100)),
);
let workflow_id = workflow.id.clone();
engine.register_workflow(workflow).unwrap();
let start = std::time::Instant::now();
let instance_id = engine
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
// Wait for completion
engine
.wait_for_completion(instance_id, "test", Duration::from_secs(5))
.await
.unwrap();
let duration = start.elapsed();
// Verify delay was applied (should take at least 100ms)
// Note: wait_for_completion cleans up state, so we verify via timing
assert!(
duration >= Duration::from_millis(100),
"Step delay not applied, duration: {:?}",
duration
);
}
#[tokio::test]
async fn test_step_scheduled_at() {
use chrono::Utc;
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
// Schedule step to run 100ms in the future
let scheduled_time = Utc::now() + chrono::Duration::milliseconds(100);
let workflow = WorkflowDefinition::new("scheduled_workflow", "Scheduled Test").add_step(
StepDefinition::new(
"scheduled_step",
"Scheduled Step",
Arc::new(AlwaysSucceedStep),
)
.scheduled_at(scheduled_time),
);
let workflow_id = workflow.id.clone();
engine.register_workflow(workflow).unwrap();
let start = std::time::Instant::now();
let instance_id = engine
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
// Wait for completion
engine
.wait_for_completion(instance_id, "test", Duration::from_secs(5))
.await
.unwrap();
let duration = start.elapsed();
// Verify scheduled time was respected (should take at least 100ms)
// Note: wait_for_completion cleans up state, so we verify via timing
assert!(
duration >= Duration::from_millis(100),
"Scheduled time not respected, duration: {:?}",
duration
);
}
// ============================================================================
// Conditional Branching Tests (#25)
// ============================================================================
#[tokio::test]
async fn test_run_if_true() {
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
let executed = Arc::new(AtomicU32::new(0));
let executed_clone = Arc::clone(&executed);
struct TrackingStep {
counter: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl StepExecutor<TestWorkflowData> for TrackingStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
self.counter.fetch_add(1, Ordering::SeqCst);
Ok(StepResult::Success)
}
}
// Step with run_if that always returns true
let workflow = WorkflowDefinition::new("run_if_true_workflow", "Run If True Test").add_step(
StepDefinition::new(
"conditional_step",
"Conditional Step",
Arc::new(TrackingStep { counter: executed }),
)
.run_if(|_ctx| true),
);
let workflow_id = workflow.id.clone();
engine.register_workflow(workflow).unwrap();
let instance_id = engine
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
engine
.wait_for_completion(instance_id, "test", Duration::from_secs(5))
.await
.unwrap();
// Step should have executed (condition was true)
// Note: wait_for_completion cleans up state, so we verify via counter
assert_eq!(executed_clone.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_run_if_false() {
use tokio::time::sleep;
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
let executed = Arc::new(AtomicU32::new(0));
let executed_clone = Arc::clone(&executed);
struct TrackingStep {
counter: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl StepExecutor<TestWorkflowData> for TrackingStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
self.counter.fetch_add(1, Ordering::SeqCst);
Ok(StepResult::Success)
}
}
// Step with run_if that always returns false
let workflow = WorkflowDefinition::new("run_if_false_workflow", "Run If False Test").add_step(
StepDefinition::new(
"conditional_step",
"Conditional Step",
Arc::new(TrackingStep { counter: executed }),
)
.run_if(|_ctx| false),
);
let workflow_id = workflow.id.clone();
engine.register_workflow(workflow).unwrap();
let instance_id = engine
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
// Use polling to check status (don't use wait_for_completion which cleans up state)
let mut state = engine.get_status(instance_id).await.unwrap();
for _ in 0..50 {
if state.status != WorkflowStatus::Running && state.status != WorkflowStatus::Pending {
break;
}
sleep(Duration::from_millis(50)).await;
state = engine.get_status(instance_id).await.unwrap();
}
assert_eq!(state.status, WorkflowStatus::Completed);
// Step should NOT have executed (skipped due to run_if)
assert_eq!(executed_clone.load(Ordering::SeqCst), 0);
// Verify step was marked as skipped
let step_state = state
.step_states
.get(&StepId::new("conditional_step"))
.unwrap();
assert_eq!(step_state.status, StepStatus::Skipped);
}
#[tokio::test]
async fn test_run_if_context_based() {
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
// Step that sets test_key in context
struct SetKeyStep;
#[async_trait::async_trait]
impl StepExecutor<TestWorkflowData> for SetKeyStep {
async fn execute(
&self,
context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
context.data.test_key = Some("execute_next".to_string());
Ok(StepResult::Success)
}
}
let executed = Arc::new(AtomicU32::new(0));
let executed_clone = Arc::clone(&executed);
struct TrackingStep {
counter: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl StepExecutor<TestWorkflowData> for TrackingStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
self.counter.fetch_add(1, Ordering::SeqCst);
Ok(StepResult::Success)
}
}
// Workflow where second step only runs if first step sets the right key
let workflow = WorkflowDefinition::new("context_run_if_workflow", "Context Run If Test")
.add_step(StepDefinition::new(
"set_key_step",
"Set Key",
Arc::new(SetKeyStep),
))
.add_step(
StepDefinition::new(
"conditional_step",
"Conditional Step",
Arc::new(TrackingStep { counter: executed }),
)
.depends_on(&["set_key_step"])
.run_if(|ctx| ctx.data.test_key.as_deref() == Some("execute_next")),
);
let workflow_id = workflow.id.clone();
engine.register_workflow(workflow).unwrap();
let instance_id = engine
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
engine
.wait_for_completion(instance_id, "test", Duration::from_secs(5))
.await
.unwrap();
// Step should have executed because context had the right value
// Note: wait_for_completion cleans up state, so we verify via counter
assert_eq!(executed_clone.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_depends_on_any() {
// DAG: A and B run in parallel, C waits for ANY (not both)
// A ──┐
// ├──> C (any_of)
// B ──┘
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
let start_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
Arc::new(parking_lot::RwLock::new(Vec::new()));
let end_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
Arc::new(parking_lot::RwLock::new(Vec::new()));
// A takes 50ms, B takes 200ms
// C should start after A finishes (not wait for B)
let workflow = WorkflowDefinition::new("depends_on_any_workflow", "Depends On Any Test")
.add_step(StepDefinition::new(
"step_a",
"Step A",
Arc::new(TimingStep {
step_name: "step_a".to_string(),
duration_ms: 50,
start_times: Arc::clone(&start_times),
end_times: Arc::clone(&end_times),
}),
))
.add_step(StepDefinition::new(
"step_b",
"Step B",
Arc::new(TimingStep {
step_name: "step_b".to_string(),
duration_ms: 200,
start_times: Arc::clone(&start_times),
end_times: Arc::clone(&end_times),
}),
))
.add_step(
StepDefinition::new(
"step_c",
"Step C",
Arc::new(TimingStep {
step_name: "step_c".to_string(),
duration_ms: 50,
start_times: Arc::clone(&start_times),
end_times: Arc::clone(&end_times),
}),
)
.depends_on_any(&["step_a", "step_b"]),
);
let workflow_id = workflow.id.clone();
engine.register_workflow(workflow).unwrap();
let instance_id = engine
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
engine
.wait_for_completion(instance_id, "test", Duration::from_secs(5))
.await
.unwrap();
// Verify step C started after A finished but before B finished
// Note: wait_for_completion cleans up state, so we verify via timing
let starts = start_times.read();
let ends = end_times.read();
let c_start = starts.iter().find(|(n, _)| n == "step_c").unwrap().1;
let a_end = ends.iter().find(|(n, _)| n == "step_a").unwrap().1;
let b_end = ends.iter().find(|(n, _)| n == "step_b").unwrap().1;
assert!(
c_start >= a_end,
"Step C should start after Step A finishes"
);
assert!(
c_start < b_end,
"Step C should start before Step B finishes (any_of semantics)"
);
}
#[tokio::test]
async fn test_depends_on_any_combined_with_depends_on() {
// DAG: C requires ALL of [A] AND ANY of [B, D]
// A takes 50ms, B takes 100ms, D takes 200ms
// C should start after A AND (B or D) complete
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
let start_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
Arc::new(parking_lot::RwLock::new(Vec::new()));
let end_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
Arc::new(parking_lot::RwLock::new(Vec::new()));
let workflow = WorkflowDefinition::new("combined_deps_workflow", "Combined Dependencies Test")
.add_step(StepDefinition::new(
"step_a",
"Step A",
Arc::new(TimingStep {
step_name: "step_a".to_string(),
duration_ms: 50,
start_times: Arc::clone(&start_times),
end_times: Arc::clone(&end_times),
}),
))
.add_step(StepDefinition::new(
"step_b",
"Step B",
Arc::new(TimingStep {
step_name: "step_b".to_string(),
duration_ms: 100,
start_times: Arc::clone(&start_times),
end_times: Arc::clone(&end_times),
}),
))
.add_step(StepDefinition::new(
"step_d",
"Step D",
Arc::new(TimingStep {
step_name: "step_d".to_string(),
duration_ms: 200,
start_times: Arc::clone(&start_times),
end_times: Arc::clone(&end_times),
}),
))
.add_step(
StepDefinition::new(
"step_c",
"Step C",
Arc::new(TimingStep {
step_name: "step_c".to_string(),
duration_ms: 50,
start_times: Arc::clone(&start_times),
end_times: Arc::clone(&end_times),
}),
)
.depends_on(&["step_a"]) // Must wait for A
.depends_on_any(&["step_b", "step_d"]), // AND any of B or D
);
let workflow_id = workflow.id.clone();
engine.register_workflow(workflow).unwrap();
let instance_id = engine
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
engine
.wait_for_completion(instance_id, "test", Duration::from_secs(5))
.await
.unwrap();
// Verify step C started after both A AND B finished
// (B finishes at 100ms, which is after A at 50ms)
// Note: wait_for_completion cleans up state, so we verify via timing
let starts = start_times.read();
let ends = end_times.read();
let c_start = starts.iter().find(|(n, _)| n == "step_c").unwrap().1;
let a_end = ends.iter().find(|(n, _)| n == "step_a").unwrap().1;
let b_end = ends.iter().find(|(n, _)| n == "step_b").unwrap().1;
let d_end = ends.iter().find(|(n, _)| n == "step_d").unwrap().1;
assert!(
c_start >= a_end,
"Step C should start after Step A (depends_on)"
);
assert!(
c_start >= b_end || c_start >= d_end,
"Step C should start after at least one of B or D (depends_on_any)"
);
// Since B finishes first (100ms) and A finishes before B, C should start around 100ms
assert!(
c_start < d_end,
"Step C should start before D finishes (any_of semantics)"
);
}
#[test]
fn test_dag_validation_depends_on_any_missing() {
// Create a workflow with a missing depends_on_any dependency
let mut workflow = WorkflowDefinition::new("missing_any_dep_workflow", "Missing Any Dep Test")
.add_step(StepDefinition::new(
"step_a",
"Step A",
Arc::new(AlwaysSucceedStep),
))
.add_step(
StepDefinition::new("step_b", "Step B", Arc::new(AlwaysSucceedStep))
.depends_on_any(&["nonexistent_step"]),
);
let result = workflow.validate();
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
ValidationError::MissingDependency { .. }
));
}
#[tokio::test]
async fn test_depends_on_any_all_fail() {
// When ALL depends_on_any dependencies fail, the step should be blocked
// Workflow: A and B both fail, C depends_on_any([A, B])
// Expected: C should not run, workflow should fail
use tokio::time::sleep;
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
let c_executed = Arc::new(AtomicU32::new(0));
let c_executed_clone = Arc::clone(&c_executed);
struct TrackingStep {
counter: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl StepExecutor<TestWorkflowData> for TrackingStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
self.counter.fetch_add(1, Ordering::SeqCst);
Ok(StepResult::Success)
}
}
let workflow = WorkflowDefinition::new("all_any_fail_workflow", "All Any Fail Test")
.add_step(StepDefinition::new(
"step_a",
"Step A (fails)",
Arc::new(AlwaysFailStep),
))
.add_step(StepDefinition::new(
"step_b",
"Step B (fails)",
Arc::new(AlwaysFailStep),
))
.add_step(
StepDefinition::new(
"step_c",
"Step C (depends on any of A, B)",
Arc::new(TrackingStep {
counter: c_executed,
}),
)
.depends_on_any(&["step_a", "step_b"]),
);
let workflow_id = workflow.id.clone();
engine.register_workflow(workflow).unwrap();
let instance_id = engine
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
// Use polling to check status
let mut state = engine.get_status(instance_id).await.unwrap();
for _ in 0..50 {
if state.status != WorkflowStatus::Running && state.status != WorkflowStatus::Pending {
break;
}
sleep(Duration::from_millis(50)).await;
state = engine.get_status(instance_id).await.unwrap();
}
// Workflow should have failed (because all depends_on_any deps failed)
assert_eq!(state.status, WorkflowStatus::Failed);
// Step C should NOT have executed
assert_eq!(c_executed_clone.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn test_depends_on_any_one_fails_one_succeeds() {
// When only SOME depends_on_any dependencies fail (but at least one succeeds),
// the step should still run
// Workflow: A fails, B succeeds, C depends_on_any([A, B])
// Expected: C should run (B succeeded)
use tokio::time::sleep;
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
let c_executed = Arc::new(AtomicU32::new(0));
let c_executed_clone = Arc::clone(&c_executed);
struct TrackingStep {
counter: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl StepExecutor<TestWorkflowData> for TrackingStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
self.counter.fetch_add(1, Ordering::SeqCst);
Ok(StepResult::Success)
}
}
// Note: Step A uses ContinueNextStep so its failure doesn't fail the workflow.
// This is the correct way to model "any of" semantics where a failing path
// shouldn't fail the entire workflow if another path succeeds.
let workflow = WorkflowDefinition::new("one_any_fail_workflow", "One Any Fail Test")
.add_step(
StepDefinition::new("step_a", "Step A (fails)", Arc::new(AlwaysFailStep))
.with_failure_action(FailureAction::ContinueNextStep),
)
.add_step(StepDefinition::new(
"step_b",
"Step B (succeeds)",
Arc::new(AlwaysSucceedStep),
))
.add_step(
StepDefinition::new(
"step_c",
"Step C (depends on any of A, B)",
Arc::new(TrackingStep {
counter: c_executed,
}),
)
.depends_on_any(&["step_a", "step_b"]),
);
let workflow_id = workflow.id.clone();
engine.register_workflow(workflow).unwrap();
let instance_id = engine
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
// Use polling to check status
let mut state = engine.get_status(instance_id).await.unwrap();
for _ in 0..50 {
if state.status != WorkflowStatus::Running && state.status != WorkflowStatus::Pending {
break;
}
sleep(Duration::from_millis(50)).await;
state = engine.get_status(instance_id).await.unwrap();
}
// Workflow should have completed (B succeeded, so C could run)
assert_eq!(state.status, WorkflowStatus::Completed);
// Step C SHOULD have executed (because B succeeded)
assert_eq!(c_executed_clone.load(Ordering::SeqCst), 1);
}