From 52d0ca944b77812e1a0e60033878dc11785e9e9d Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sun, 25 Jan 2026 18:20:58 -0500 Subject: [PATCH] [smg] import official wfaas crate to replace self managed one (#17724) --- sgl-model-gateway/Cargo.toml | 5 +- .../benches/special_token_search.rs | 53 - sgl-model-gateway/src/lib.rs | 2 +- sgl-model-gateway/src/workflow/definition.rs | 364 ----- sgl-model-gateway/src/workflow/engine.rs | 1243 --------------- sgl-model-gateway/src/workflow/event.rs | 282 ---- sgl-model-gateway/src/workflow/executor.rs | 152 -- sgl-model-gateway/src/workflow/mod.rs | 15 - sgl-model-gateway/src/workflow/state.rs | 204 --- sgl-model-gateway/src/workflow/types.rs | 277 ---- sgl-model-gateway/tests/workflow_test.rs | 1343 ----------------- 11 files changed, 2 insertions(+), 3938 deletions(-) delete mode 100644 sgl-model-gateway/benches/special_token_search.rs delete mode 100644 sgl-model-gateway/src/workflow/definition.rs delete mode 100644 sgl-model-gateway/src/workflow/engine.rs delete mode 100644 sgl-model-gateway/src/workflow/event.rs delete mode 100644 sgl-model-gateway/src/workflow/executor.rs delete mode 100644 sgl-model-gateway/src/workflow/mod.rs delete mode 100644 sgl-model-gateway/src/workflow/state.rs delete mode 100644 sgl-model-gateway/src/workflow/types.rs delete mode 100644 sgl-model-gateway/tests/workflow_test.rs diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index 1bbabda10..94375b188 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -85,6 +85,7 @@ openai-protocol = { version = "1.0.0", features = ["axum"] } tool-parser = "1.0.0" llm-tokenizer = "1.0.0" smg-auth = "1.0.0" +wfaas = "1.0.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std"] } rustls-pemfile = "2.2" openssl = "0.10.73" @@ -153,10 +154,6 @@ tonic-v12 = { version = "0.12.3", package = "tonic" } serial_test = "3.0" rsa = { version = "0.9", features = ["sha2"] } -[[bench]] -name = "special_token_search" -harness = false -path = "benches/special_token_search.rs" [[bench]] name = "wasm_middleware_latency" harness = false diff --git a/sgl-model-gateway/benches/special_token_search.rs b/sgl-model-gateway/benches/special_token_search.rs deleted file mode 100644 index 2c57a44b0..000000000 --- a/sgl-model-gateway/benches/special_token_search.rs +++ /dev/null @@ -1,53 +0,0 @@ -use aho_corasick::AhoCorasick; -use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput}; - -fn find_special_token_boundaries_naive( - text: &str, - special_tokens: &[String], -) -> Vec<(usize, usize)> { - let mut boundaries = Vec::new(); - for token in special_tokens { - let mut start = 0; - while let Some(pos) = text[start..].find(token) { - let actual_pos = start + pos; - boundaries.push((actual_pos, actual_pos + token.len())); - start = actual_pos + token.len(); - } - } - boundaries.sort_by_key(|b| b.0); - boundaries -} - -fn find_special_token_boundaries_aho(text: &str, ac: &AhoCorasick) -> Vec<(usize, usize)> { - ac.find_iter(text) - .map(|mat| (mat.start(), mat.end())) - .collect() -} - -fn bench_token_search_comparison(c: &mut Criterion) { - let mut group = c.benchmark_group("token_boundary_search"); - let text = "User: Hello! Assistant: How can I help you today? ".repeat(1000); - - for token_count in [5, 50] { - let special_tokens: Vec = (0..token_count) - .map(|i| format!("<|stop_sequence_{}|>", i)) - .collect(); - - let ac = AhoCorasick::new(&special_tokens).unwrap(); - group.throughput(Throughput::Bytes(text.len() as u64)); - - group.bench_function(format!("naive_tokens_{}", token_count), |b| { - b.iter(|| { - find_special_token_boundaries_naive(black_box(&text), black_box(&special_tokens)) - }) - }); - - group.bench_function(format!("aho_tokens_{}", token_count), |b| { - b.iter(|| find_special_token_boundaries_aho(black_box(&text), black_box(&ac))) - }); - } - group.finish(); -} - -criterion_group!(benches, bench_token_search_comparison); -criterion_main!(benches); diff --git a/sgl-model-gateway/src/lib.rs b/sgl-model-gateway/src/lib.rs index 2693b364e..a75f0e5a4 100644 --- a/sgl-model-gateway/src/lib.rs +++ b/sgl-model-gateway/src/lib.rs @@ -19,4 +19,4 @@ pub use llm_tokenizer as tokenizer; pub use tool_parser; pub mod version; pub mod wasm; -pub mod workflow; +pub use wfaas as workflow; diff --git a/sgl-model-gateway/src/workflow/definition.rs b/sgl-model-gateway/src/workflow/definition.rs deleted file mode 100644 index c1f2ba0eb..000000000 --- a/sgl-model-gateway/src/workflow/definition.rs +++ /dev/null @@ -1,364 +0,0 @@ -//! Workflow definition types - -use std::{ - collections::{HashMap, HashSet}, - fmt, - sync::Arc, - time::Duration, -}; - -use chrono::{DateTime, Utc}; - -use super::{ - executor::StepExecutor, - types::{FailureAction, RetryPolicy, StepId, WorkflowContext, WorkflowData, WorkflowId}, -}; - -/// A condition function that determines whether a step should run. -pub type StepCondition = Arc) -> bool + Send + Sync>; - -/// Errors that can occur during workflow validation -#[derive(Debug, Clone, thiserror::Error)] -pub enum ValidationError { - /// A step depends on another step that doesn't exist - #[error("Step '{step}' depends on non-existent step '{dependency}'")] - MissingDependency { step: StepId, dependency: StepId }, - - /// A cycle was detected in the workflow DAG - #[error("Cycle detected involving step '{0}'")] - CycleDetected(StepId), -} - -/// Definition of a single step within a workflow -pub struct StepDefinition { - pub id: StepId, - pub name: String, - pub executor: Arc>, - pub retry_policy: Option, - pub timeout: Option, - pub on_failure: FailureAction, - /// Dependencies that must ALL complete before this step runs - pub depends_on: Vec, - /// Dependencies where ANY completing triggers this step (used with depends_on) - pub depends_on_any: Vec, - /// Delay before starting the step (after dependencies satisfied) - pub delay: Option, - /// Run step at or after this time (after dependencies satisfied) - pub scheduled_at: Option>, - /// Condition to evaluate; if false, step is skipped - pub run_if: Option>, -} - -impl fmt::Debug for StepDefinition { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("StepDefinition") - .field("id", &self.id) - .field("name", &self.name) - .field("retry_policy", &self.retry_policy) - .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(|_| "")) - .finish_non_exhaustive() - } -} - -impl StepDefinition { - pub fn new( - id: impl Into, - name: impl Into, - executor: Arc>, - ) -> Self { - Self { - id: StepId::new(id.into()), - name: name.into(), - executor, - retry_policy: None, - timeout: None, - on_failure: FailureAction::FailWorkflow, - depends_on: Vec::new(), - depends_on_any: Vec::new(), - delay: None, - scheduled_at: None, - run_if: None, - } - } - - pub fn with_retry(mut self, policy: RetryPolicy) -> Self { - self.retry_policy = Some(policy); - self - } - - pub fn with_timeout(mut self, timeout: Duration) -> Self { - self.timeout = Some(timeout); - self - } - - pub fn with_failure_action(mut self, action: FailureAction) -> Self { - self.on_failure = action; - self - } - - /// Set dependencies for this step. - /// 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) -> 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(mut self, condition: F) -> Self - where - F: Fn(&WorkflowContext) -> bool + Send + Sync + 'static, - { - self.run_if = Some(Arc::new(condition)); - self - } -} - -/// Complete workflow definition -pub struct WorkflowDefinition { - pub id: WorkflowId, - pub name: String, - pub steps: Vec>, - pub default_retry_policy: RetryPolicy, - pub default_timeout: Duration, - /// Pre-computed reverse dependencies: step_id -> indices of steps that depend on it - reverse_deps: HashMap>, - /// Pre-computed indices of steps with no dependencies (can start immediately) - initial_step_indices: Vec, -} - -impl fmt::Debug for WorkflowDefinition { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("WorkflowDefinition") - .field("id", &self.id) - .field("name", &self.name) - .field("steps", &self.steps) - .field("default_retry_policy", &self.default_retry_policy) - .field("default_timeout", &self.default_timeout) - .finish_non_exhaustive() - } -} - -impl WorkflowDefinition { - pub fn new(id: impl Into, name: impl Into) -> Self { - Self { - id: WorkflowId::new(id.into()), - name: name.into(), - steps: Vec::new(), - default_retry_policy: RetryPolicy::default(), - default_timeout: Duration::from_secs(300), // 5 minutes - reverse_deps: HashMap::new(), - initial_step_indices: Vec::new(), - } - } - - pub fn add_step(mut self, step: StepDefinition) -> Self { - self.steps.push(step); - self - } - - pub fn with_default_retry(mut self, policy: RetryPolicy) -> Self { - self.default_retry_policy = policy; - self - } - - pub fn with_default_timeout(mut self, timeout: Duration) -> Self { - self.default_timeout = timeout; - self - } - - /// Get the retry policy for a step (step-specific or default) - pub fn get_retry_policy<'a>(&'a self, step: &'a StepDefinition) -> &'a RetryPolicy { - step.retry_policy - .as_ref() - .unwrap_or(&self.default_retry_policy) - } - - /// Get the timeout for a step (step-specific or default) - pub fn get_timeout(&self, step: &StepDefinition) -> Duration { - step.timeout.unwrap_or(self.default_timeout) - } - - /// Validate the workflow DAG structure and build dependency graph. - /// Returns an error if: - /// - A step depends on a non-existent step - /// - There's a cycle in the dependencies - /// - /// On success, pre-computes reverse dependencies for O(1) dependent lookup. - #[must_use = "validation result should be checked"] - pub fn validate(&mut self) -> Result<(), ValidationError> { - // Build HashMap for O(1) lookup instead of O(n) linear search - let steps_map: HashMap<&StepId, &StepDefinition> = - self.steps.iter().map(|s| (&s.id, s)).collect(); - - // 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) { - return Err(ValidationError::MissingDependency { - step: step.id.clone(), - dependency: dep.clone(), - }); - } - } - 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 (considers both dependency types) - let mut visited = HashSet::new(); - let mut rec_stack = HashSet::new(); - - for step in &self.steps { - if !visited.contains(&step.id) - && Self::has_cycle(&step.id, &steps_map, &mut visited, &mut rec_stack) - { - return Err(ValidationError::CycleDetected(step.id.clone())); - } - } - - // 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 { - self.reverse_deps - .entry(dep_id.clone()) - .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() && s.depends_on_any.is_empty()) - .map(|(i, _)| i) - .collect(); - - Ok(()) - } - - /// DFS helper for cycle detection with O(1) HashMap lookup - fn has_cycle<'a>( - step_id: &'a StepId, - steps_map: &HashMap<&'a StepId, &'a StepDefinition>, - visited: &mut HashSet<&'a StepId>, - rec_stack: &mut HashSet<&'a StepId>, - ) -> bool { - if rec_stack.contains(step_id) { - return true; // Back edge found - cycle! - } - if visited.contains(step_id) { - return false; // Already fully processed - } - - visited.insert(step_id); - 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); - false - } - - /// Get indices of steps that depend on the given step - pub fn get_dependent_indices(&self, step_id: &StepId) -> &[usize] { - self.reverse_deps - .get(step_id) - .map(|v| v.as_slice()) - .unwrap_or(&[]) - } - - /// Get indices of steps with no dependencies - pub fn get_initial_step_indices(&self) -> &[usize] { - &self.initial_step_indices - } -} diff --git a/sgl-model-gateway/src/workflow/engine.rs b/sgl-model-gateway/src/workflow/engine.rs deleted file mode 100644 index dec72d93f..000000000 --- a/sgl-model-gateway/src/workflow/engine.rs +++ /dev/null @@ -1,1243 +0,0 @@ -//! Workflow execution engine -//! -//! Supports DAG-based parallel execution of workflow steps. -//! Steps with no dependencies run in parallel, steps with dependencies -//! wait for all dependencies to complete successfully. - -use std::{ - collections::{HashMap, HashSet, VecDeque}, - marker::PhantomData, - sync::{ - atomic::{AtomicUsize, Ordering}, - Arc, - }, - time::Duration, -}; - -use backoff::{backoff::Backoff, ExponentialBackoffBuilder}; -use chrono::Utc; -use parking_lot::RwLock; -use tokio::{ - sync::{mpsc, watch}, - time::timeout, -}; - -use super::{ - definition::{StepDefinition, WorkflowDefinition}, - event::{EventBus, WorkflowEvent}, - state::{InMemoryStore, StateStore}, - types::*, -}; - -#[derive(Default)] -struct StepTracker { - completed: HashSet, - failed: HashSet, - skipped: HashSet, - running: HashSet, - /// 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, -} - -impl StepTracker { - fn total_processed(&self) -> usize { - self.completed.len() + self.failed.len() + self.skipped.len() - } - - 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 { - 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 -struct FixedBackoff(Duration); - -impl Backoff for FixedBackoff { - fn reset(&mut self) {} - - fn next_backoff(&mut self) -> Option { - Some(self.0) - } -} - -/// Linear backoff that increases delay by a fixed amount each retry -struct LinearBackoff { - current: Duration, - increment: Duration, - max: Duration, -} - -impl LinearBackoff { - fn new(increment: Duration, max: Duration) -> Self { - Self { - current: increment, - increment, - max, - } - } -} - -impl Backoff for LinearBackoff { - fn reset(&mut self) { - self.current = self.increment; - } - - fn next_backoff(&mut self) -> Option { - let next = self.current; - self.current = (self.current + self.increment).min(self.max); - Some(next) - } -} - -/// Enum-based backoff implementation to avoid heap allocation -enum BackoffImpl { - Fixed(FixedBackoff), - Exponential(backoff::ExponentialBackoff), - Linear(LinearBackoff), -} - -impl BackoffImpl { - fn next_backoff(&mut self) -> Option { - match self { - BackoffImpl::Fixed(b) => b.next_backoff(), - BackoffImpl::Exponential(b) => b.next_backoff(), - BackoffImpl::Linear(b) => b.next_backoff(), - } - } -} - -/// Main workflow execution engine -/// -/// # Type Parameters -/// -/// * `D` - The workflow data type that implements `WorkflowData` -/// * `S` - The state store implementation (defaults to `InMemoryStore`) -/// -/// # Graceful Shutdown -/// -/// The engine supports graceful shutdown via [`shutdown()`](Self::shutdown): -/// -/// ```ignore -/// // Trigger shutdown - stops accepting new workflows -/// engine.shutdown(); -/// -/// // Wait for all running workflows to complete (with timeout) -/// if !engine.wait_for_shutdown(Duration::from_secs(30)).await { -/// // Force cancel remaining workflows -/// engine.force_cancel_all().await; -/// } -/// ``` -pub struct WorkflowEngine = InMemoryStore> { - definitions: Arc>>>>, - state_store: S, - event_bus: Arc, - /// Shutdown signal sender - when true, engine is shutting down - shutdown_tx: Arc>, - /// Shutdown signal receiver for cloning to tasks - shutdown_rx: watch::Receiver, - /// Count of active workflow executions - active_workflows: Arc, - _phantom: PhantomData, -} - -impl WorkflowEngine> { - pub fn new() -> Self { - Self::with_store(InMemoryStore::new()) - } -} - -impl + 'static> WorkflowEngine { - /// Create a new workflow engine with a custom state store - pub fn with_store(state_store: S) -> Self { - let (shutdown_tx, shutdown_rx) = watch::channel(false); - Self { - definitions: Arc::new(RwLock::new(HashMap::new())), - state_store, - event_bus: Arc::new(EventBus::new()), - shutdown_tx: Arc::new(shutdown_tx), - shutdown_rx, - active_workflows: Arc::new(AtomicUsize::new(0)), - _phantom: PhantomData, - } - } - - /// Check if the engine is shutting down - pub fn is_shutting_down(&self) -> bool { - *self.shutdown_rx.borrow() - } - - /// Initiate graceful shutdown - /// - /// This will: - /// - Stop accepting new workflows (start_workflow will return an error) - /// - Stop the cleanup task - /// - Allow running workflows to complete - /// - /// Use [`wait_for_shutdown`](Self::wait_for_shutdown) to wait for completion. - pub fn shutdown(&self) { - let _ = self.shutdown_tx.send(true); - tracing::info!("Workflow engine shutdown initiated"); - } - - /// Wait for all active workflows to complete - /// - /// Returns `true` if all workflows completed within the timeout, - /// `false` if the timeout was reached with workflows still running. - /// - /// Uses simple polling - appropriate for shutdown which happens once per process. - pub async fn wait_for_shutdown(&self, timeout_duration: Duration) -> bool { - let start = tokio::time::Instant::now(); - - loop { - let active = self.active_workflows.load(Ordering::Acquire); - if active == 0 { - tracing::info!("All workflows completed, shutdown complete"); - return true; - } - - if start.elapsed() >= timeout_duration { - tracing::warn!( - remaining_workflows = active, - "Shutdown timeout reached with workflows still running" - ); - return false; - } - - tracing::debug!( - active_workflows = active, - "Waiting for workflows to complete" - ); - - tokio::time::sleep(Duration::from_millis(50)).await; - } - } - - /// Force cancel all running workflows - /// - /// This should be called after `wait_for_shutdown` times out if you need - /// to ensure all workflows are stopped. Note that this cancels workflows - /// at the state level; running steps may still complete. - pub async fn force_cancel_all(&self) -> usize { - let active_states = match self.state_store.list_active().await { - Ok(states) => states, - Err(e) => { - tracing::error!(error = ?e, "Failed to list active workflows for force cancel"); - return 0; - } - }; - - let mut cancelled = 0; - for state in active_states { - if let Err(e) = self.cancel_workflow(state.instance_id).await { - tracing::warn!( - instance_id = %state.instance_id, - error = ?e, - "Failed to cancel workflow during force shutdown" - ); - } else { - cancelled += 1; - } - } - - tracing::info!(cancelled_count = cancelled, "Force cancelled workflows"); - cancelled - } - - /// Get the number of currently active workflow executions - pub fn active_workflow_count(&self) -> usize { - self.active_workflows.load(Ordering::Acquire) - } - - /// Decrement active workflow count - fn workflow_finished(&self) { - 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) -> Option { - 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 { - ActiveWorkflowGuard { - active_workflows: Arc::clone(&self.active_workflows), - } - } - - /// Start a background task to periodically clean up old workflow states - /// - /// This prevents unbounded memory growth by removing completed/failed workflows - /// that are older than the specified TTL. - /// - /// The task will automatically stop when [`shutdown()`](Self::shutdown) is called. - /// - /// # Arguments - /// - /// * `ttl` - Time-to-live for terminal workflows (default: 1 hour) - /// * `interval` - How often to run cleanup (default: 5 minutes) - /// - /// # Returns - /// - /// A join handle for the cleanup task that can be used to stop it. - pub fn start_cleanup_task( - &self, - ttl: Option, - interval: Option, - ) -> tokio::task::JoinHandle<()> { - let state_store = self.state_store.clone(); - let ttl = ttl.unwrap_or(Duration::from_secs(3600)); // 1 hour default - let interval = interval.unwrap_or(Duration::from_secs(300)); // 5 minutes default - let mut shutdown_rx = self.shutdown_rx.clone(); - - tokio::spawn(async move { - let mut ticker = tokio::time::interval(interval); - ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - - loop { - tokio::select! { - _ = ticker.tick() => { - state_store.cleanup_old_workflows(ttl).await; - } - _ = shutdown_rx.changed() => { - tracing::info!("Cleanup task stopping due to shutdown"); - break; - } - } - } - }) - } - - /// Register a workflow definition - #[must_use = "registration result should be checked"] - pub fn register_workflow( - &self, - mut definition: WorkflowDefinition, - ) -> Result<(), super::definition::ValidationError> { - // Validate DAG and build dependency graph once at registration - definition.validate()?; - - let id = definition.id.clone(); - self.definitions.write().insert(id, Arc::new(definition)); - Ok(()) - } - - /// Get the event bus for subscribing to workflow events - pub fn event_bus(&self) -> Arc { - Arc::clone(&self.event_bus) - } - - /// Get the state store - pub fn state_store(&self) -> &S { - &self.state_store - } - - /// Start a new workflow instance - /// - /// Returns `Err(WorkflowError::ShuttingDown)` if the engine is shutting down. - #[must_use = "workflow instance ID should be stored or awaited"] - pub async fn start_workflow( - &self, - definition_id: WorkflowId, - data: D, - ) -> WorkflowResult { - // Guard increments counter and decrements on drop unless committed. - // This handles all error paths automatically. - let guard = StartGuard::new(self); - - if self.is_shutting_down() { - return Err(WorkflowError::ShuttingDown); - } - - let definition = self - .definitions - .read() - .get(&definition_id) - .cloned() - .ok_or_else(|| WorkflowError::DefinitionNotFound(definition_id.clone()))?; - - let instance_id = WorkflowInstanceId::new(); - let mut state = WorkflowState::new(instance_id, definition_id.clone(), data); - state.status = WorkflowStatus::Running; - - for step in &definition.steps { - state - .step_states - .insert(step.id.clone(), StepState::default()); - } - - self.state_store.save(state).await?; - - self.event_bus - .publish(WorkflowEvent::WorkflowStarted { - instance_id, - definition_id, - }) - .await; - - // Commit the guard - from here the spawned task takes ownership of the count - guard.commit(); - - let engine = self.clone_for_execution(); - let def = Arc::clone(&definition); - tokio::spawn(async move { - let _guard = engine.active_workflow_guard(); - let result = engine.execute_workflow(instance_id, def).await; - if let Err(e) = result { - tracing::error!(instance_id = %instance_id, error = ?e, "Workflow execution failed"); - } - }); - - Ok(instance_id) - } - - /// Execute a workflow with DAG-based parallel execution - /// - /// Uses event-driven readiness: instead of scanning all steps each iteration, - /// we only check steps whose dependencies just completed. - async fn execute_workflow( - &self, - instance_id: WorkflowInstanceId, - definition: Arc>, - ) -> WorkflowResult<()> { - let start_time = std::time::Instant::now(); - let step_count = definition.steps.len(); - - let tracker: Arc> = Arc::new(RwLock::new(StepTracker::default())); - let (tx, mut rx) = mpsc::channel::<(StepId, StepResult)>(step_count.max(1)); - - // Initialize with steps that have no dependencies (O(1) lookup) - let mut pending_check: VecDeque = definition - .get_initial_step_indices() - .iter() - .copied() - .collect(); - - loop { - if self.state_store.is_cancelled(instance_id).await? { - self.event_bus - .publish(WorkflowEvent::WorkflowCancelled { instance_id }) - .await; - return Ok(()); - } - - // 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(); - - // O(w) iteration over waiting_until keys instead of O(n) over all steps - let wait_ready: Vec = t.get_ready_waiting_indices(); - - // Check pending_check for dependency-satisfied steps - let deps_ready: Vec = pending_check - .drain(..) - .filter(|&idx| { - let step = &definition.steps[idx]; - 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(); - - ( - 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 - if total_processed == step_count { - break; - } - - // 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(); - // 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" - }; - - self.state_store - .update(instance_id, |s| { - s.status = WorkflowStatus::Failed; - }) - .await?; - self.event_bus - .publish(WorkflowEvent::WorkflowFailed { - instance_id, - failed_step: failed_step - .unwrap_or_else(|| StepId::new("internal_scheduler")), - error: error_message.to_string(), - }) - .await; - return Ok(()); - } - - // Launch ready steps in parallel - 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(); - let tx = tx.clone(); - let tracker = Arc::clone(&tracker); - - 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; - - let signal = match &result { - Ok(r) => *r, - Err(_) => StepResult::Failure, - }; - - // Track whether we need to update state to Skipped after releasing lock - let needs_skip_update = { - let mut t = tracker.write(); - t.running.remove(&step_id); - - let needs_update = match result { - Ok(StepResult::Success) => { - t.completed.insert(step_id.clone()); - false - } - Ok(StepResult::Skip) => { - t.skipped.insert(step_id.clone()); - false - } - Ok(StepResult::Failure) | Err(_) => match step.on_failure { - FailureAction::FailWorkflow | FailureAction::RetryIndefinitely => { - t.failed.insert(step_id.clone()); - false - } - FailureAction::ContinueNextStep => { - t.skipped.insert(step_id.clone()); - true // Need to update state store after releasing lock - } - }, - }; - - if let Err(e) = tx.try_send((step_id.clone(), signal)) { - use mpsc::error::TrySendError; - match e { - TrySendError::Full(_) => { - tracing::error!( - step_id = %step_id, - "Channel full when sending step completion - this is a bug" - ); - } - TrySendError::Closed(_) => { - tracing::debug!( - step_id = %step_id, - "Channel closed, workflow likely cancelled" - ); - } - } - } - - needs_update - }; - - // Perform async state update after releasing the tracker lock - if needs_skip_update { - if let Err(e) = 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 - { - tracing::warn!( - step_id = %step_id, - error = ?e, - "Failed to update step state to Skipped" - ); - } - } - }); - } - - // 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, - result = ?result, - "Step completed" - ); - - // Add dependents of completed step to pending_check (O(1) lookup) - // Only if the step succeeded or was skipped (not failed) - if matches!(result, StepResult::Success | StepResult::Skip) { - for &dep_idx in definition.get_dependent_indices(&completed_step_id) { - pending_check.push_back(dep_idx); - } - } - } - } 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; - } - } - - let failed_step = { - let t = tracker.read(); - t.failed.iter().next().cloned() - }; - - if let Some(ref step) = failed_step { - self.state_store - .update(instance_id, |s| { - s.status = WorkflowStatus::Failed; - }) - .await?; - self.event_bus - .publish(WorkflowEvent::WorkflowFailed { - instance_id, - failed_step: step.clone(), - error: "One or more steps failed".into(), - }) - .await; - } else { - self.state_store - .update(instance_id, |s| { - s.status = WorkflowStatus::Completed; - }) - .await?; - - let duration = start_time.elapsed(); - self.event_bus - .publish(WorkflowEvent::WorkflowCompleted { - instance_id, - duration, - }) - .await; - } - - Ok(()) - } - - /// Execute a step with retry logic - async fn execute_step_with_retry( - &self, - instance_id: WorkflowInstanceId, - step: &StepDefinition, - definition: &WorkflowDefinition, - ) -> WorkflowResult { - let retry_policy = definition.get_retry_policy(step); - let step_timeout = definition.get_timeout(step); - - let mut attempt = 1; - let max_attempts = if matches!(step.on_failure, FailureAction::RetryIndefinitely) { - u32::MAX - } else { - retry_policy.max_attempts - }; - - let mut backoff = Self::create_backoff(&retry_policy.backoff); - - loop { - if self.state_store.is_cancelled(instance_id).await? { - return Err(WorkflowError::Cancelled(instance_id)); - } - - // Update step state - self.state_store - .update(instance_id, |s| { - s.current_step = Some(step.id.clone()); - if let Some(step_state) = s.step_states.get_mut(&step.id) { - step_state.status = if attempt == 1 { - StepStatus::Running - } else { - StepStatus::Retrying - }; - step_state.attempt = attempt; - step_state.started_at = Some(Utc::now()); - } - }) - .await?; - - // Emit step started event - self.event_bus - .publish(WorkflowEvent::StepStarted { - instance_id, - step_id: step.id.clone(), - attempt, - }) - .await; - - let mut context = self.state_store.get_context(instance_id).await?; - - // Execute step with timeout - let step_start = std::time::Instant::now(); - let result = timeout(step_timeout, step.executor.execute(&mut context)).await; - - let step_duration = step_start.elapsed(); - - self.state_store - .update(instance_id, |s| { - s.context = context.clone(); - }) - .await?; - - match result { - Ok(Ok(StepResult::Success)) => { - // Step succeeded - self.state_store - .update(instance_id, |s| { - if let Some(step_state) = s.step_states.get_mut(&step.id) { - step_state.status = StepStatus::Succeeded; - step_state.completed_at = Some(Utc::now()); - } - }) - .await?; - - self.event_bus - .publish(WorkflowEvent::StepSucceeded { - instance_id, - step_id: step.id.clone(), - duration: step_duration, - }) - .await; - - // Call on_success hook - if let Err(e) = step.executor.on_success(&context).await { - tracing::warn!(step_id = %step.id, error = ?e, "on_success hook failed"); - } - - return Ok(StepResult::Success); - } - Ok(Ok(StepResult::Skip)) => { - return Ok(StepResult::Skip); - } - Ok(Ok(StepResult::Failure)) | Ok(Err(_)) | Err(_) => { - let (error_msg, should_retry) = match result { - Ok(Err(e)) => { - let msg = format!("{}", e); - let retryable = step.executor.is_retryable(&e); - (msg, retryable) - } - Err(_) => ( - format!("Step timeout after {:?}", step_timeout), - true, // Timeouts are retryable - ), - _ => ("Step failed".to_string(), false), - }; - - let will_retry = should_retry && attempt < max_attempts; - - // Update step state - self.state_store - .update(instance_id, |s| { - if let Some(step_state) = s.step_states.get_mut(&step.id) { - step_state.status = if will_retry { - StepStatus::Retrying - } else { - StepStatus::Failed - }; - step_state.last_error = Some(error_msg.clone()); - if !will_retry { - step_state.completed_at = Some(Utc::now()); - } - } - }) - .await?; - - // Emit step failed event - self.event_bus - .publish(WorkflowEvent::StepFailed { - instance_id, - step_id: step.id.clone(), - error: error_msg.clone(), - will_retry, - }) - .await; - - if will_retry { - // Calculate backoff delay - let delay = backoff - .next_backoff() - .unwrap_or_else(|| Duration::from_secs(1)); - - self.event_bus - .publish(WorkflowEvent::StepRetrying { - instance_id, - step_id: step.id.clone(), - attempt: attempt + 1, - delay, - }) - .await; - - tokio::time::sleep(delay).await; - attempt += 1; - } else { - // No more retries, call on_failure hook - // Create a generic error for the hook - let hook_error = WorkflowError::StepFailed { - step_id: step.id.clone(), - message: error_msg, - }; - if let Err(hook_err) = step.executor.on_failure(&context, &hook_error).await - { - tracing::warn!(step_id = %step.id, error = ?hook_err, "on_failure hook failed"); - } - - return Ok(StepResult::Failure); - } - } - } - } - } - - fn create_backoff(strategy: &BackoffStrategy) -> BackoffImpl { - match strategy { - BackoffStrategy::Fixed(duration) => BackoffImpl::Fixed(FixedBackoff(*duration)), - BackoffStrategy::Exponential { base, max } => { - let backoff = ExponentialBackoffBuilder::new() - .with_initial_interval(*base) - .with_max_interval(*max) - .with_max_elapsed_time(None) - .build(); - BackoffImpl::Exponential(backoff) - } - BackoffStrategy::Linear { increment, max } => { - BackoffImpl::Linear(LinearBackoff::new(*increment, *max)) - } - } - } - - /// Cancel a running workflow - pub async fn cancel_workflow(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> { - self.state_store - .update(instance_id, |s| { - s.status = WorkflowStatus::Cancelled; - }) - .await?; - - self.event_bus - .publish(WorkflowEvent::WorkflowCancelled { instance_id }) - .await; - - Ok(()) - } - - /// Get workflow status - pub async fn get_status( - &self, - instance_id: WorkflowInstanceId, - ) -> WorkflowResult> { - self.state_store.load(instance_id).await - } - - /// Wait for a workflow to complete with adaptive polling - /// - /// Returns Ok with success message on completion, Err on failure/timeout/cancellation. - /// Automatically cleans up terminal workflow states. - pub async fn wait_for_completion( - &self, - instance_id: WorkflowInstanceId, - label: &str, - timeout_duration: Duration, - ) -> Result { - let start = std::time::Instant::now(); - let mut poll_interval = Duration::from_millis(100); - let max_poll_interval = Duration::from_millis(2000); - let poll_backoff = Duration::from_millis(200); - - loop { - if start.elapsed() > timeout_duration { - return Err(format!( - "Workflow timeout after {}s for {}", - timeout_duration.as_secs(), - label - )); - } - - let state = self - .get_status(instance_id) - .await - .map_err(|e| format!("Failed to get workflow status: {:?}", e))?; - - let result = match state.status { - WorkflowStatus::Completed => { - Ok(format!("{} completed successfully via workflow", label)) - } - WorkflowStatus::Failed => { - let current_step = state.current_step.as_ref(); - let step_name = current_step - .map(|s| s.to_string()) - .unwrap_or_else(|| "unknown".to_string()); - let error_msg = current_step - .and_then(|step_id| state.step_states.get(step_id)) - .and_then(|s| s.last_error.as_deref()) - .unwrap_or("Unknown error"); - Err(format!( - "Workflow failed at step {}: {}", - step_name, error_msg - )) - } - WorkflowStatus::Cancelled => Err(format!("Workflow cancelled for {}", label)), - WorkflowStatus::Pending | WorkflowStatus::Paused | WorkflowStatus::Running => { - tokio::time::sleep(poll_interval).await; - poll_interval = (poll_interval + poll_backoff).min(max_poll_interval); - continue; - } - }; - - self.state_store.cleanup_if_terminal(instance_id).await; - return result; - } - } - - /// Clone engine for async execution - fn clone_for_execution(&self) -> Self { - Self { - definitions: Arc::clone(&self.definitions), - state_store: self.state_store.clone(), - event_bus: Arc::clone(&self.event_bus), - shutdown_tx: Arc::clone(&self.shutdown_tx), - shutdown_rx: self.shutdown_rx.clone(), - active_workflows: Arc::clone(&self.active_workflows), - _phantom: PhantomData, - } - } -} - -/// RAII guard that decrements active_workflows count on drop. -/// Ensures proper cleanup even if a workflow task panics. -struct ActiveWorkflowGuard { - active_workflows: Arc, -} - -impl Drop for ActiveWorkflowGuard { - fn drop(&mut self) { - self.active_workflows.fetch_sub(1, Ordering::Release); - } -} - -/// RAII guard for start_workflow that increments on creation and decrements on drop -/// unless commit() is called. Handles all error paths automatically. -struct StartGuard<'a, D: WorkflowData, S: StateStore + 'static> { - engine: &'a WorkflowEngine, - committed: bool, -} - -impl<'a, D: WorkflowData, S: StateStore + 'static> StartGuard<'a, D, S> { - fn new(engine: &'a WorkflowEngine) -> Self { - engine.active_workflows.fetch_add(1, Ordering::AcqRel); - Self { - engine, - committed: false, - } - } - - fn commit(mut self) { - self.committed = true; - } -} - -impl + 'static> Drop for StartGuard<'_, D, S> { - fn drop(&mut self) { - if !self.committed { - self.engine.workflow_finished(); - } - } -} - -/// Clone implementation for internal use. -/// -/// **Note**: This creates a shallow clone that shares state with the original engine. -/// Both engines will share the same: -/// - Workflow definitions -/// - State store -/// - Event bus -/// - Shutdown signal -/// - Active workflow counter -/// -/// This is intentional for spawning async tasks that need access to the engine. -/// For most use cases, prefer sharing the engine via `Arc` rather -/// than cloning. -impl + 'static> Clone for WorkflowEngine { - fn clone(&self) -> Self { - self.clone_for_execution() - } -} - -impl Default for WorkflowEngine> { - fn default() -> Self { - Self::new() - } -} - -impl + 'static> std::fmt::Debug for WorkflowEngine { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("WorkflowEngine") - .field("definitions_count", &self.definitions.read().len()) - .finish() - } -} diff --git a/sgl-model-gateway/src/workflow/event.rs b/sgl-model-gateway/src/workflow/event.rs deleted file mode 100644 index e24d2307e..000000000 --- a/sgl-model-gateway/src/workflow/event.rs +++ /dev/null @@ -1,282 +0,0 @@ -//! Workflow event system for observability and monitoring - -use std::{sync::Arc, time::Duration}; - -use async_trait::async_trait; -use tokio::sync::RwLock; -use tracing::{error, info, warn}; - -use super::types::{StepId, WorkflowId, WorkflowInstanceId}; - -/// Default timeout for subscriber event handlers -const DEFAULT_SUBSCRIBER_TIMEOUT: Duration = Duration::from_secs(30); - -/// Events emitted by the workflow engine -#[derive(Debug, Clone)] -pub enum WorkflowEvent { - WorkflowStarted { - instance_id: WorkflowInstanceId, - definition_id: WorkflowId, - }, - StepStarted { - instance_id: WorkflowInstanceId, - step_id: StepId, - attempt: u32, - }, - StepSucceeded { - instance_id: WorkflowInstanceId, - step_id: StepId, - duration: Duration, - }, - StepFailed { - instance_id: WorkflowInstanceId, - step_id: StepId, - error: String, - will_retry: bool, - }, - StepRetrying { - instance_id: WorkflowInstanceId, - step_id: StepId, - attempt: u32, - delay: Duration, - }, - WorkflowCompleted { - instance_id: WorkflowInstanceId, - duration: Duration, - }, - WorkflowFailed { - instance_id: WorkflowInstanceId, - failed_step: StepId, - error: String, - }, - WorkflowCancelled { - instance_id: WorkflowInstanceId, - }, -} - -/// Trait for subscribing to workflow events -#[async_trait] -pub trait EventSubscriber: Send + Sync { - async fn on_event(&self, event: &WorkflowEvent); -} - -/// Event bus for publishing and subscribing to workflow events -/// -/// # Subscriber Isolation -/// -/// Each subscriber is notified in a separate spawned task with a timeout. -/// This ensures that: -/// - A slow subscriber doesn't block other subscribers -/// - A panicking subscriber doesn't affect other subscribers -/// - Event publishing returns quickly regardless of subscriber behavior -pub struct EventBus { - subscribers: Arc>>>, - /// Timeout for each subscriber's event handler - subscriber_timeout: Duration, -} - -impl EventBus { - pub fn new() -> Self { - Self { - subscribers: Arc::new(RwLock::new(Vec::new())), - subscriber_timeout: DEFAULT_SUBSCRIBER_TIMEOUT, - } - } - - /// Create an EventBus with a custom subscriber timeout - pub fn with_timeout(timeout: Duration) -> Self { - Self { - subscribers: Arc::new(RwLock::new(Vec::new())), - subscriber_timeout: timeout, - } - } - - /// Subscribe to workflow events - pub async fn subscribe(&self, subscriber: Arc) { - self.subscribers.write().await.push(subscriber); - } - - /// Unsubscribe from workflow events - /// - /// Removes the subscriber by Arc pointer equality. - /// Returns true if the subscriber was found and removed. - pub async fn unsubscribe(&self, subscriber: &Arc) -> bool { - let mut subs = self.subscribers.write().await; - let len_before = subs.len(); - subs.retain(|s| !Arc::ptr_eq(s, subscriber)); - subs.len() < len_before - } - - /// Publish an event to all subscribers concurrently - /// - /// Each subscriber is notified in a separate spawned task with a timeout. - /// This method returns after spawning all notification tasks, without - /// waiting for subscribers to complete (fire-and-forget). - /// - /// Subscriber failures (timeout or panic) are logged but don't affect - /// other subscribers or the caller. - pub async fn publish(&self, event: WorkflowEvent) { - let subscribers: Vec<_> = self.subscribers.read().await.iter().cloned().collect(); - let timeout = self.subscriber_timeout; - - for (idx, subscriber) in subscribers.into_iter().enumerate() { - let event = event.clone(); - tokio::spawn(async move { - let result = tokio::time::timeout(timeout, subscriber.on_event(&event)).await; - match result { - Ok(()) => {} - Err(_) => { - warn!( - subscriber_index = idx, - timeout_secs = timeout.as_secs(), - "Event subscriber timed out" - ); - } - } - }); - } - } - - /// Publish an event and wait for all subscribers to complete - /// - /// Unlike `publish`, this method waits for all subscribers to finish - /// (or timeout). Use this when you need to ensure all subscribers - /// have processed the event before continuing. - pub async fn publish_and_wait(&self, event: WorkflowEvent) { - let subscribers: Vec<_> = self.subscribers.read().await.iter().cloned().collect(); - let timeout = self.subscriber_timeout; - - let handles: Vec<_> = subscribers - .into_iter() - .enumerate() - .map(|(idx, subscriber)| { - let event = event.clone(); - tokio::spawn(async move { - let result = tokio::time::timeout(timeout, subscriber.on_event(&event)).await; - if result.is_err() { - warn!( - subscriber_index = idx, - timeout_secs = timeout.as_secs(), - "Event subscriber timed out" - ); - } - }) - }) - .collect(); - - // Wait for all spawned tasks, ignoring individual failures (panics) - for handle in handles { - let _ = handle.await; - } - } -} - -impl Default for EventBus { - fn default() -> Self { - Self::new() - } -} - -impl std::fmt::Debug for EventBus { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("EventBus").finish_non_exhaustive() - } -} - -/// Logging subscriber that logs events using tracing -pub struct LoggingSubscriber; - -#[async_trait] -impl EventSubscriber for LoggingSubscriber { - async fn on_event(&self, event: &WorkflowEvent) { - match event { - WorkflowEvent::WorkflowStarted { - instance_id, - definition_id, - } => { - info!( - instance_id = %instance_id, - definition_id = %definition_id, - "Workflow started" - ); - } - WorkflowEvent::StepStarted { - instance_id, - step_id, - attempt, - } => { - info!( - instance_id = %instance_id, - step_id = %step_id, - attempt = attempt, - "Step started" - ); - } - WorkflowEvent::StepSucceeded { - instance_id, - step_id, - duration, - } => { - info!( - instance_id = %instance_id, - step_id = %step_id, - duration_ms = duration.as_millis(), - "Step succeeded" - ); - } - WorkflowEvent::StepFailed { - instance_id, - step_id, - error, - will_retry, - } => { - warn!( - instance_id = %instance_id, - step_id = %step_id, - error = error, - will_retry = will_retry, - "Step failed" - ); - } - WorkflowEvent::StepRetrying { - instance_id, - step_id, - attempt, - delay, - } => { - info!( - instance_id = %instance_id, - step_id = %step_id, - attempt = attempt, - delay_ms = delay.as_millis(), - "Step retrying" - ); - } - WorkflowEvent::WorkflowCompleted { - instance_id, - duration, - } => { - info!( - instance_id = %instance_id, - duration_ms = duration.as_millis(), - "Workflow completed" - ); - } - WorkflowEvent::WorkflowFailed { - instance_id, - failed_step, - error, - } => { - error!( - instance_id = %instance_id, - failed_step = %failed_step, - error = error, - "Workflow failed" - ); - } - WorkflowEvent::WorkflowCancelled { instance_id } => { - info!(instance_id = %instance_id, "Workflow cancelled"); - } - } - } -} diff --git a/sgl-model-gateway/src/workflow/executor.rs b/sgl-model-gateway/src/workflow/executor.rs deleted file mode 100644 index 9dc50468e..000000000 --- a/sgl-model-gateway/src/workflow/executor.rs +++ /dev/null @@ -1,152 +0,0 @@ -//! Step executor trait and implementations - -use async_trait::async_trait; - -use super::types::{StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult}; - -/// Trait for executing individual workflow steps -#[async_trait] -pub trait StepExecutor: Send + Sync { - /// Execute the step with the given context - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult; - - /// Check if an error is retry-able - /// - /// Override this method to customize which errors should trigger retries. - /// By default, all errors are considered retry-able. - fn is_retryable(&self, _error: &WorkflowError) -> bool { - true - } - - /// Called when the step succeeds - /// - /// This hook allows steps to perform cleanup or additional actions - /// after successful execution. - async fn on_success(&self, _context: &WorkflowContext) -> WorkflowResult<()> { - Ok(()) - } - - /// Called when the step fails after all retries - /// - /// This hook allows steps to perform cleanup or compensation logic - /// when the step cannot complete successfully. - async fn on_failure( - &self, - _context: &WorkflowContext, - _error: &WorkflowError, - ) -> WorkflowResult<()> { - Ok(()) - } -} - -/// Simple function-based step executor -pub struct FunctionStep -where - D: WorkflowData, - F: Fn( - &mut WorkflowContext, - ) -> std::pin::Pin< - Box> + Send + '_>, - > + Send - + Sync, -{ - func: F, - _phantom: std::marker::PhantomData, -} - -impl FunctionStep -where - D: WorkflowData, - F: Fn( - &mut WorkflowContext, - ) -> std::pin::Pin< - Box> + Send + '_>, - > + Send - + Sync, -{ - pub fn new(func: F) -> Self { - Self { - func, - _phantom: std::marker::PhantomData, - } - } -} - -#[async_trait] -impl StepExecutor for FunctionStep -where - D: WorkflowData, - F: Fn( - &mut WorkflowContext, - ) -> std::pin::Pin< - Box> + Send + '_>, - > + Send - + Sync, -{ - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - (self.func)(context).await - } -} - -#[cfg(test)] -mod tests { - use serde::{Deserialize, Serialize}; - - use super::*; - use crate::workflow::types::WorkflowInstanceId; - - #[derive(Debug, Clone, Serialize, Deserialize)] - struct TestData { - value: i32, - } - - impl WorkflowData for TestData { - fn workflow_type() -> &'static str { - "test" - } - } - - struct TestStep { - should_succeed: bool, - } - - #[async_trait] - impl StepExecutor for TestStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - if self.should_succeed { - Ok(StepResult::Success) - } else { - Err(WorkflowError::StepFailed { - step_id: crate::workflow::types::StepId::new("test"), - message: "test error".to_string(), - }) - } - } - } - - #[tokio::test] - async fn test_step_executor_success() { - let step = TestStep { - should_succeed: true, - }; - let mut context = WorkflowContext::new(WorkflowInstanceId::new(), TestData { value: 42 }); - - let result = step.execute(&mut context).await; - assert!(result.is_ok()); - assert_eq!(result.unwrap(), StepResult::Success); - } - - #[tokio::test] - async fn test_step_executor_failure() { - let step = TestStep { - should_succeed: false, - }; - let mut context = WorkflowContext::new(WorkflowInstanceId::new(), TestData { value: 42 }); - - let result = step.execute(&mut context).await; - assert!(result.is_err()); - } -} diff --git a/sgl-model-gateway/src/workflow/mod.rs b/sgl-model-gateway/src/workflow/mod.rs deleted file mode 100644 index 9290b3c05..000000000 --- a/sgl-model-gateway/src/workflow/mod.rs +++ /dev/null @@ -1,15 +0,0 @@ -//! Workflow engine for managing multi-step operations - -mod definition; -mod engine; -mod event; -mod executor; -mod state; -pub mod types; - -pub use definition::{StepDefinition, ValidationError, WorkflowDefinition}; -pub use engine::WorkflowEngine; -pub use event::{EventBus, EventSubscriber, LoggingSubscriber, WorkflowEvent}; -pub use executor::{FunctionStep, StepExecutor}; -pub use state::{InMemoryStore, StateStore}; -pub use types::*; diff --git a/sgl-model-gateway/src/workflow/state.rs b/sgl-model-gateway/src/workflow/state.rs deleted file mode 100644 index 127159949..000000000 --- a/sgl-model-gateway/src/workflow/state.rs +++ /dev/null @@ -1,204 +0,0 @@ -//! Workflow state management - -use std::{collections::HashMap, marker::PhantomData, sync::Arc, time::Duration}; - -use async_trait::async_trait; -use parking_lot::RwLock; - -use super::types::{ - WorkflowContext, WorkflowData, WorkflowError, WorkflowInstanceId, WorkflowResult, - WorkflowState, WorkflowStatus, -}; - -/// Trait for workflow state persistence. -/// -/// Implement this trait to provide custom storage backends (e.g., PostgreSQL, Redis). -/// The default implementation is `InMemoryStore` which keeps state in memory. -#[async_trait] -pub trait StateStore: Send + Sync + Clone { - /// Save workflow state - async fn save(&self, state: WorkflowState) -> WorkflowResult<()>; - - /// Load workflow state by instance ID - async fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult>; - - /// Update workflow state using a closure - async fn update(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()> - where - F: FnOnce(&mut WorkflowState) + Send; - - /// Delete workflow state - async fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()>; - - /// List all active workflows (Running or Pending) - async fn list_active(&self) -> WorkflowResult>>; - - /// List all workflows - async fn list_all(&self) -> WorkflowResult>>; - - /// Check if workflow is cancelled without loading full state - async fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult; - - /// Clean up old completed/failed/cancelled workflows beyond a time threshold - async fn cleanup_old_workflows(&self, ttl: Duration) -> usize; - - /// Get just the workflow context without cloning the entire state - async fn get_context( - &self, - instance_id: WorkflowInstanceId, - ) -> WorkflowResult>; - - /// Clean up a specific workflow immediately if it's in a terminal state - /// Returns true if the workflow was removed, false otherwise - async fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool; -} - -/// In-memory state storage for workflow instances -#[derive(Clone)] -pub struct InMemoryStore { - states: Arc>>>, - _phantom: PhantomData, -} - -impl InMemoryStore { - pub fn new() -> Self { - Self { - states: Arc::new(RwLock::new(HashMap::new())), - _phantom: PhantomData, - } - } - - /// Get count of workflows by status - pub fn count_by_status(&self, status: WorkflowStatus) -> usize { - self.states - .read() - .values() - .filter(|s| s.status == status) - .count() - } - - /// Get total count of all workflows - pub fn count(&self) -> usize { - self.states.read().len() - } -} - -impl Default for InMemoryStore { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl StateStore for InMemoryStore { - async fn save(&self, state: WorkflowState) -> WorkflowResult<()> { - self.states.write().insert(state.instance_id, state); - Ok(()) - } - - async fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult> { - self.states - .read() - .get(&instance_id) - .cloned() - .ok_or(WorkflowError::NotFound(instance_id)) - } - - async fn update(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()> - where - F: FnOnce(&mut WorkflowState) + Send, - { - let mut states = self.states.write(); - let state = states - .get_mut(&instance_id) - .ok_or(WorkflowError::NotFound(instance_id))?; - f(state); - state.updated_at = chrono::Utc::now(); - Ok(()) - } - - async fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> { - self.states.write().remove(&instance_id); - Ok(()) - } - - async fn list_active(&self) -> WorkflowResult>> { - let states = self.states.read(); - Ok(states - .values() - .filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending)) - .cloned() - .collect()) - } - - async fn list_all(&self) -> WorkflowResult>> { - let states = self.states.read(); - Ok(states.values().cloned().collect()) - } - - async fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult { - self.states - .read() - .get(&instance_id) - .map(|s| s.status == WorkflowStatus::Cancelled) - .ok_or(WorkflowError::NotFound(instance_id)) - } - - async fn cleanup_old_workflows(&self, ttl: Duration) -> usize { - let now = chrono::Utc::now(); - let mut states = self.states.write(); - let initial_count = states.len(); - - states.retain(|_, state| { - // Keep active workflows - if matches!( - state.status, - WorkflowStatus::Running | WorkflowStatus::Pending | WorkflowStatus::Paused - ) { - return true; - } - - // For terminal workflows, check age - let age = now - .signed_duration_since(state.updated_at) - .to_std() - .unwrap_or_default(); - age < ttl - }); - - let removed_count = initial_count - states.len(); - if removed_count > 0 { - tracing::info!( - removed = removed_count, - remaining = states.len(), - "Cleaned up old workflow states" - ); - } - removed_count - } - - async fn get_context( - &self, - instance_id: WorkflowInstanceId, - ) -> WorkflowResult> { - self.states - .read() - .get(&instance_id) - .map(|s| s.context.clone()) - .ok_or(WorkflowError::NotFound(instance_id)) - } - - async fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool { - let mut states = self.states.write(); - if let Some(state) = states.get(&instance_id) { - if matches!( - state.status, - WorkflowStatus::Completed | WorkflowStatus::Failed | WorkflowStatus::Cancelled - ) { - states.remove(&instance_id); - return true; - } - } - false - } -} diff --git a/sgl-model-gateway/src/workflow/types.rs b/sgl-model-gateway/src/workflow/types.rs deleted file mode 100644 index 0ba58a8bb..000000000 --- a/sgl-model-gateway/src/workflow/types.rs +++ /dev/null @@ -1,277 +0,0 @@ -//! Core workflow types and definitions - -use std::{collections::HashMap, fmt, time::Duration}; - -use chrono::{DateTime, Utc}; -use serde::{de::DeserializeOwned, Deserialize, Serialize}; -use uuid::Uuid; - -/// Trait for workflow data that can be passed through workflow steps. -/// -/// Implementing this trait allows your data type to be used as the typed -/// context for a workflow. The data must be serializable for state persistence. -/// -/// # Example -/// -/// ```ignore -/// #[derive(Debug, Clone, Serialize, Deserialize)] -/// pub struct MyWorkflowData { -/// pub config: MyConfig, -/// pub result: Option, -/// #[serde(skip, default)] -/// pub app_context: Option>, -/// } -/// -/// impl WorkflowData for MyWorkflowData { -/// fn workflow_type() -> &'static str { "my_workflow" } -/// } -/// ``` -pub trait WorkflowData: Serialize + DeserializeOwned + Send + Sync + Clone + 'static { - /// Human-readable name for logging and identification - fn workflow_type() -> &'static str; -} - -/// Unique identifier for a workflow definition -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub struct WorkflowId(String); - -impl WorkflowId { - pub fn new(id: impl Into) -> Self { - Self(id.into()) - } -} - -impl fmt::Display for WorkflowId { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}", self.0) - } -} - -/// Unique identifier for a workflow instance -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub struct WorkflowInstanceId(Uuid); - -impl WorkflowInstanceId { - pub fn new() -> Self { - Self(Uuid::new_v4()) - } -} - -impl Default for WorkflowInstanceId { - fn default() -> Self { - Self::new() - } -} - -impl fmt::Display for WorkflowInstanceId { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}", self.0) - } -} - -/// Unique identifier for a workflow step -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub struct StepId(String); - -impl StepId { - pub fn new(id: impl Into) -> Self { - Self(id.into()) - } -} - -impl fmt::Display for StepId { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}", self.0) - } -} - -/// Retry policy configuration -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct RetryPolicy { - pub max_attempts: u32, - pub backoff: BackoffStrategy, -} - -impl Default for RetryPolicy { - fn default() -> Self { - Self { - max_attempts: 3, - backoff: BackoffStrategy::Exponential { - base: Duration::from_secs(1), - max: Duration::from_secs(30), - }, - } - } -} - -/// Backoff strategy for retries -#[derive(Debug, Clone, Serialize, Deserialize)] -pub enum BackoffStrategy { - /// Fixed delay between retries - Fixed(Duration), - /// Exponential backoff with base and max duration - Exponential { base: Duration, max: Duration }, - /// Linear backoff with increment and max duration - Linear { increment: Duration, max: Duration }, -} - -/// Action to take when a step fails -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum FailureAction { - /// Stop the entire workflow - FailWorkflow, - /// Skip this step and continue to the next - ContinueNextStep, - /// Keep retrying indefinitely until manual intervention - RetryIndefinitely, -} - -/// Workflow execution status -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum WorkflowStatus { - Pending, - Running, - Paused, - Completed, - Failed, - Cancelled, -} - -/// Step execution status -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum StepStatus { - Pending, - Running, - Succeeded, - Failed, - Retrying, - Skipped, -} - -/// State of a workflow step -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct StepState { - pub status: StepStatus, - pub attempt: u32, - pub last_error: Option, - pub started_at: Option>, - pub completed_at: Option>, -} - -impl Default for StepState { - fn default() -> Self { - Self { - status: StepStatus::Pending, - attempt: 0, - last_error: None, - started_at: None, - completed_at: None, - } - } -} - -/// Workflow instance state -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(bound( - serialize = "D: Serialize", - deserialize = "D: serde::de::DeserializeOwned" -))] -pub struct WorkflowState { - pub instance_id: WorkflowInstanceId, - pub definition_id: WorkflowId, - pub status: WorkflowStatus, - pub current_step: Option, - pub step_states: HashMap, - pub context: WorkflowContext, - pub created_at: DateTime, - pub updated_at: DateTime, -} - -impl WorkflowState { - pub fn new(instance_id: WorkflowInstanceId, definition_id: WorkflowId, data: D) -> Self { - let now = Utc::now(); - Self { - instance_id, - definition_id, - status: WorkflowStatus::Pending, - current_step: None, - step_states: HashMap::new(), - context: WorkflowContext::new(instance_id, data), - created_at: now, - updated_at: now, - } - } -} - -/// Shared context passed between workflow steps. -/// -/// The context contains typed workflow data that is fully serializable, -/// enabling state persistence and workflow recovery. -/// -/// # Type Parameter -/// -/// `D` - The workflow-specific data type implementing `WorkflowData`. -/// This type holds all the state needed by workflow steps and must be -/// serializable (except for fields marked with `#[serde(skip)]`). -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(bound( - serialize = "D: Serialize", - deserialize = "D: serde::de::DeserializeOwned" -))] -pub struct WorkflowContext { - pub instance_id: WorkflowInstanceId, - pub data: D, -} - -impl WorkflowContext { - pub fn new(instance_id: WorkflowInstanceId, data: D) -> Self { - Self { instance_id, data } - } -} - -/// Result returned by a step execution -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum StepResult { - Success, - Failure, - Skip, -} - -/// Error kinds for workflow operations -#[derive(Debug, Clone, thiserror::Error)] -pub enum WorkflowError { - #[error("Workflow not found: {0}")] - NotFound(WorkflowInstanceId), - - #[error("Workflow definition not found: {0}")] - DefinitionNotFound(WorkflowId), - - #[error("Step failed: {step_id} - {message}")] - StepFailed { step_id: StepId, message: String }, - - #[error("Step timeout: {step_id}")] - StepTimeout { step_id: StepId }, - - #[error("Workflow cancelled: {0}")] - Cancelled(WorkflowInstanceId), - - #[error("Invalid state transition: {from:?} -> {to:?}")] - InvalidStateTransition { - from: WorkflowStatus, - to: WorkflowStatus, - }, - - #[error("Context value not found: {0}")] - ContextValueNotFound(String), - - #[error("Type mismatch: expected {expected}, got {actual}")] - TypeMismatch { - expected: &'static str, - actual: &'static str, - }, - - #[error("Engine is shutting down, not accepting new workflows")] - ShuttingDown, -} - -pub type WorkflowResult = Result; diff --git a/sgl-model-gateway/tests/workflow_test.rs b/sgl-model-gateway/tests/workflow_test.rs deleted file mode 100644 index a70509f8c..000000000 --- a/sgl-model-gateway/tests/workflow_test.rs +++ /dev/null @@ -1,1343 +0,0 @@ -//! Integration tests for workflow engine - -use std::{ - sync::{ - atomic::{AtomicU32, Ordering}, - Arc, - }, - time::Duration, -}; - -use serde::{Deserialize, Serialize}; -use smg::workflow::*; -use tokio::time::sleep; - -/// Test workflow data type for integration tests. -#[derive(Debug, Clone, Serialize, Deserialize, Default)] -struct TestWorkflowData { - /// Execution count for tracking step invocations - pub execution_count: u32, - /// Test key for context sharing tests - pub test_key: Option, -} - -impl WorkflowData for TestWorkflowData { - fn workflow_type() -> &'static str { - "test_workflow" - } -} - -// Test step that counts invocations -struct CountingStep { - counter: Arc, - should_succeed_after: u32, -} - -#[async_trait::async_trait] -impl StepExecutor for CountingStep { - async fn execute( - &self, - context: &mut WorkflowContext, - ) -> WorkflowResult { - let count = self.counter.fetch_add(1, Ordering::SeqCst) + 1; - - // Store count in context - context.data.execution_count = count; - - if count >= self.should_succeed_after { - Ok(StepResult::Success) - } else { - Err(WorkflowError::StepFailed { - step_id: StepId::new("counting_step"), - message: format!("Not ready yet, attempt {}", count), - }) - } - } -} - -// Test step that always succeeds -struct AlwaysSucceedStep; - -#[async_trait::async_trait] -impl StepExecutor for AlwaysSucceedStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - Ok(StepResult::Success) - } -} - -struct AlwaysFailStep; - -#[async_trait::async_trait] -impl StepExecutor for AlwaysFailStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - Ok(StepResult::Failure) - } -} - -#[tokio::test] -async fn test_simple_workflow_execution() { - let engine: WorkflowEngine = WorkflowEngine::new(); - - // Subscribe to events for logging - engine - .event_bus() - .subscribe(Arc::new(LoggingSubscriber)) - .await; - - // Create a simple workflow - let workflow = WorkflowDefinition::new("test_workflow", "Simple Test Workflow") - .add_step(StepDefinition::new( - "step1", - "First Step", - Arc::new(AlwaysSucceedStep), - )) - .add_step(StepDefinition::new( - "step2", - "Second Step", - Arc::new(AlwaysSucceedStep), - )); - - let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow).unwrap(); - - // Start workflow - let instance_id = engine - .start_workflow(workflow_id, TestWorkflowData::default()) - .await - .unwrap(); - - // Wait for completion - sleep(Duration::from_millis(100)).await; - - // Check status - let state = engine.get_status(instance_id).await.unwrap(); - assert_eq!(state.status, WorkflowStatus::Completed); - assert_eq!(state.step_states.len(), 2); -} - -#[tokio::test] -async fn test_workflow_with_retry() { - let engine: WorkflowEngine = WorkflowEngine::new(); - engine - .event_bus() - .subscribe(Arc::new(LoggingSubscriber)) - .await; - - let counter = Arc::new(AtomicU32::new(0)); - - // Create workflow with retry logic - let workflow = WorkflowDefinition::new("retry_workflow", "Workflow with Retry").add_step( - StepDefinition::new( - "retry_step", - "Step that retries", - Arc::new(CountingStep { - counter: Arc::clone(&counter), - should_succeed_after: 3, - }), - ) - .with_retry(RetryPolicy { - max_attempts: 5, - backoff: BackoffStrategy::Fixed(Duration::from_millis(10)), - }) - .with_timeout(Duration::from_secs(5)), - ); - - let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow).unwrap(); - - // Start workflow - let instance_id = engine - .start_workflow(workflow_id, TestWorkflowData::default()) - .await - .unwrap(); - - // Wait for completion - sleep(Duration::from_millis(500)).await; - - // Check that step was retried and eventually succeeded - let state = engine.get_status(instance_id).await.unwrap(); - assert_eq!(state.status, WorkflowStatus::Completed); - - let step_state = state.step_states.get(&StepId::new("retry_step")).unwrap(); - assert_eq!(step_state.status, StepStatus::Succeeded); - assert_eq!(step_state.attempt, 3); // Should have taken 3 attempts - - // Verify counter - assert_eq!(counter.load(Ordering::SeqCst), 3); -} - -#[tokio::test] -async fn test_workflow_failure_after_max_retries() { - let engine: WorkflowEngine = WorkflowEngine::new(); - engine - .event_bus() - .subscribe(Arc::new(LoggingSubscriber)) - .await; - - let counter = Arc::new(AtomicU32::new(0)); - - // Create workflow that will fail - let workflow = WorkflowDefinition::new("failing_workflow", "Workflow that Fails").add_step( - StepDefinition::new( - "failing_step", - "Step that always fails", - Arc::new(CountingStep { - counter: Arc::clone(&counter), - should_succeed_after: 10, // Will never succeed within max_attempts - }), - ) - .with_retry(RetryPolicy { - max_attempts: 3, - backoff: BackoffStrategy::Fixed(Duration::from_millis(10)), - }) - .with_failure_action(FailureAction::FailWorkflow), - ); - - let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow).unwrap(); - - // Start workflow - let instance_id = engine - .start_workflow(workflow_id, TestWorkflowData::default()) - .await - .unwrap(); - - // Wait for completion - sleep(Duration::from_millis(500)).await; - - // Check that workflow failed - let state = engine.get_status(instance_id).await.unwrap(); - assert_eq!(state.status, WorkflowStatus::Failed); - - let step_state = state.step_states.get(&StepId::new("failing_step")).unwrap(); - assert_eq!(step_state.status, StepStatus::Failed); - assert_eq!(step_state.attempt, 3); // Should have tried 3 times - - // Verify counter - assert_eq!(counter.load(Ordering::SeqCst), 3); -} - -#[tokio::test] -async fn test_workflow_continue_on_failure() { - let engine: WorkflowEngine = WorkflowEngine::new(); - engine - .event_bus() - .subscribe(Arc::new(LoggingSubscriber)) - .await; - - let counter = Arc::new(AtomicU32::new(0)); - - // Create workflow where first step fails but workflow continues - let workflow = WorkflowDefinition::new("continue_workflow", "Continue on Failure") - .add_step( - StepDefinition::new( - "failing_step", - "Step that fails", - Arc::new(CountingStep { - counter: Arc::clone(&counter), - should_succeed_after: 10, - }), - ) - .with_retry(RetryPolicy { - max_attempts: 2, - backoff: BackoffStrategy::Fixed(Duration::from_millis(10)), - }) - .with_failure_action(FailureAction::ContinueNextStep), - ) - .add_step(StepDefinition::new( - "success_step", - "Step that succeeds", - Arc::new(AlwaysSucceedStep), - )); - - let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow).unwrap(); - - // Start workflow - let instance_id = engine - .start_workflow(workflow_id, TestWorkflowData::default()) - .await - .unwrap(); - - // Wait for completion - sleep(Duration::from_millis(500)).await; - - // Workflow should complete despite first step failing - let state = engine.get_status(instance_id).await.unwrap(); - assert_eq!(state.status, WorkflowStatus::Completed); - - // First step should be skipped - let step1_state = state.step_states.get(&StepId::new("failing_step")).unwrap(); - assert_eq!(step1_state.status, StepStatus::Skipped); - - // Second step should succeed - let step2_state = state.step_states.get(&StepId::new("success_step")).unwrap(); - assert_eq!(step2_state.status, StepStatus::Succeeded); -} - -#[tokio::test] -async fn test_workflow_context_sharing() { - let engine: WorkflowEngine = WorkflowEngine::new(); - - struct ContextWriterStep { - value: String, - } - - #[async_trait::async_trait] - impl StepExecutor for ContextWriterStep { - async fn execute( - &self, - context: &mut WorkflowContext, - ) -> WorkflowResult { - context.data.test_key = Some(self.value.clone()); - Ok(StepResult::Success) - } - } - - struct ContextReaderStep { - expected_value: String, - } - - #[async_trait::async_trait] - impl StepExecutor for ContextReaderStep { - async fn execute( - &self, - context: &mut WorkflowContext, - ) -> WorkflowResult { - let value = context - .data - .test_key - .as_ref() - .ok_or_else(|| WorkflowError::ContextValueNotFound("test_key".to_string()))?; - - if value == &self.expected_value { - Ok(StepResult::Success) - } else { - Err(WorkflowError::StepFailed { - step_id: StepId::new("reader"), - message: format!("Expected {}, got {}", self.expected_value, value), - }) - } - } - } - - let workflow = WorkflowDefinition::new("context_workflow", "Context Sharing Test") - .add_step(StepDefinition::new( - "writer", - "Write to context", - Arc::new(ContextWriterStep { - value: "test_value".to_string(), - }), - )) - .add_step(StepDefinition::new( - "reader", - "Read from context", - Arc::new(ContextReaderStep { - expected_value: "test_value".to_string(), - }), - )); - - let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow).unwrap(); - - let instance_id = engine - .start_workflow(workflow_id, TestWorkflowData::default()) - .await - .unwrap(); - - sleep(Duration::from_millis(100)).await; - - let state = engine.get_status(instance_id).await.unwrap(); - assert_eq!(state.status, WorkflowStatus::Completed); -} - -// ============================================================================ -// DAG / Parallel Execution Tests -// ============================================================================ - -// Step that records when it starts and ends (for testing parallel execution) -struct TimingStep { - step_name: String, - duration_ms: u64, - start_times: Arc>>, - end_times: Arc>>, -} - -#[async_trait::async_trait] -impl StepExecutor for TimingStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - let start = std::time::Instant::now(); - self.start_times - .write() - .push((self.step_name.clone(), start)); - - sleep(Duration::from_millis(self.duration_ms)).await; - - let end = std::time::Instant::now(); - self.end_times.write().push((self.step_name.clone(), end)); - - Ok(StepResult::Success) - } -} - -#[tokio::test] -async fn test_parallel_execution_no_dependencies() { - // Steps without dependencies should run in parallel - let engine: WorkflowEngine = WorkflowEngine::new(); - - let start_times: Arc>> = - Arc::new(parking_lot::RwLock::new(Vec::new())); - let end_times: Arc>> = - Arc::new(parking_lot::RwLock::new(Vec::new())); - - // Three steps, each taking 100ms, no dependencies - // If parallel: ~100ms total - // If sequential: ~300ms total - let workflow = WorkflowDefinition::new("parallel_workflow", "Parallel Test") - .add_step(StepDefinition::new( - "step_a", - "Step A", - Arc::new(TimingStep { - step_name: "step_a".to_string(), - duration_ms: 100, - 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_c", - "Step C", - Arc::new(TimingStep { - step_name: "step_c".to_string(), - duration_ms: 100, - start_times: Arc::clone(&start_times), - end_times: Arc::clone(&end_times), - }), - )); - - let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow).unwrap(); - - let overall_start = std::time::Instant::now(); - let instance_id = engine - .start_workflow(workflow_id, TestWorkflowData::default()) - .await - .unwrap(); - - // Wait for completion - give enough time for async scheduling - for _ in 0..50 { - sleep(Duration::from_millis(50)).await; - let state = engine.get_status(instance_id).await.unwrap(); - if state.status != WorkflowStatus::Running { - break; - } - } - - let overall_duration = overall_start.elapsed(); - - let state = engine.get_status(instance_id).await.unwrap(); - assert_eq!(state.status, WorkflowStatus::Completed); - - // Check that all steps completed - assert_eq!(end_times.read().len(), 3); - - // Verify parallel execution: all steps should start around the same time - let starts = start_times.read(); - let first_start = starts.iter().map(|(_, t)| t).min().unwrap(); - let last_start = starts.iter().map(|(_, t)| t).max().unwrap(); - - // All starts should be within 100ms of each other (allowing for scheduling variance) - let start_spread = last_start.duration_since(*first_start); - assert!( - start_spread < Duration::from_millis(100), - "Steps did not start in parallel, spread: {:?}", - start_spread - ); - - // Total duration should be less than sequential (300ms) - use generous threshold - assert!( - overall_duration < Duration::from_millis(500), - "Parallel execution took too long: {:?}", - overall_duration - ); -} - -#[tokio::test] -async fn test_dag_with_dependencies() { - // DAG: A and B run in parallel, C waits for both - // A ──┐ - // ├──> C - // B ──┘ - let engine: WorkflowEngine = WorkflowEngine::new(); - - let start_times: Arc>> = - Arc::new(parking_lot::RwLock::new(Vec::new())); - let end_times: Arc>> = - Arc::new(parking_lot::RwLock::new(Vec::new())); - - let workflow = WorkflowDefinition::new("dag_workflow", "DAG 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_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", "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(); - - // Poll until workflow completes (or timeout) - for _ in 0..50 { - sleep(Duration::from_millis(50)).await; - let state = engine.get_status(instance_id).await.unwrap(); - if state.status != WorkflowStatus::Running { - break; - } - } - - let state = engine.get_status(instance_id).await.unwrap(); - assert_eq!(state.status, WorkflowStatus::Completed); - - // Verify step C started after both A and B finished - 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 started before Step A finished"); - assert!(c_start >= b_end, "Step C started before Step B finished"); -} - -#[tokio::test] -async fn test_dag_dependency_failure_blocks_dependents() { - // If step A fails with FailWorkflow, step B (depends on A) should not run - let engine: WorkflowEngine = WorkflowEngine::new(); - - let b_executed = Arc::new(AtomicU32::new(0)); - - struct FailingStep; - - #[async_trait::async_trait] - impl StepExecutor for FailingStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - Err(WorkflowError::StepFailed { - step_id: StepId::new("failing"), - message: "Intentional failure".to_string(), - }) - } - - fn is_retryable(&self, _error: &WorkflowError) -> bool { - false // Disable retries for this test - } - } - - struct TrackingStep { - counter: Arc, - } - - #[async_trait::async_trait] - impl StepExecutor for TrackingStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - self.counter.fetch_add(1, Ordering::SeqCst); - Ok(StepResult::Success) - } - } - - let workflow = WorkflowDefinition::new("blocked_workflow", "Blocked Test") - .add_step( - StepDefinition::new("step_a", "Step A", Arc::new(FailingStep)) - .with_failure_action(FailureAction::FailWorkflow), - ) - .add_step( - StepDefinition::new( - "step_b", - "Step B", - Arc::new(TrackingStep { - counter: Arc::clone(&b_executed), - }), - ) - .depends_on(&["step_a"]), - ); - - let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow).unwrap(); - - let instance_id = engine - .start_workflow(workflow_id, TestWorkflowData::default()) - .await - .unwrap(); - - // Poll until workflow completes (or timeout) - for _ in 0..50 { - sleep(Duration::from_millis(50)).await; - let state = engine.get_status(instance_id).await.unwrap(); - if state.status != WorkflowStatus::Running { - break; - } - } - - let state = engine.get_status(instance_id).await.unwrap(); - assert_eq!(state.status, WorkflowStatus::Failed); - - // Step B should not have executed - assert_eq!(b_executed.load(Ordering::SeqCst), 0); -} - -#[test] -fn test_dag_validation_cycle_detection() { - // Create a workflow with a cycle: A -> B -> C -> A - let mut workflow = WorkflowDefinition::new("cyclic_workflow", "Cyclic Test") - .add_step( - StepDefinition::new("step_a", "Step A", Arc::new(AlwaysSucceedStep)) - .depends_on(&["step_c"]), - ) - .add_step( - StepDefinition::new("step_b", "Step B", Arc::new(AlwaysSucceedStep)) - .depends_on(&["step_a"]), - ) - .add_step( - StepDefinition::new("step_c", "Step C", Arc::new(AlwaysSucceedStep)) - .depends_on(&["step_b"]), - ); - - let result = workflow.validate(); - assert!(result.is_err()); - assert!(matches!( - result.unwrap_err(), - ValidationError::CycleDetected(_) - )); -} - -#[test] -fn test_dag_validation_missing_dependency() { - // Create a workflow with a missing dependency - let mut workflow = WorkflowDefinition::new("missing_dep_workflow", "Missing 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(&["nonexistent_step"]), - ); - - let result = workflow.validate(); - assert!(result.is_err()); - assert!(matches!( - result.unwrap_err(), - ValidationError::MissingDependency { .. } - )); -} - -#[test] -fn test_dag_validation_valid_workflow() { - // Create a valid DAG workflow - let mut workflow = WorkflowDefinition::new("valid_workflow", "Valid Test") - .add_step(StepDefinition::new( - "step_a", - "Step A", - Arc::new(AlwaysSucceedStep), - )) - .add_step(StepDefinition::new( - "step_b", - "Step B", - Arc::new(AlwaysSucceedStep), - )) - .add_step( - StepDefinition::new("step_c", "Step C", Arc::new(AlwaysSucceedStep)) - .depends_on(&["step_a", "step_b"]), - ) - .add_step( - StepDefinition::new("step_d", "Step D", Arc::new(AlwaysSucceedStep)) - .depends_on(&["step_c"]), - ); - - let result = workflow.validate(); - assert!(result.is_ok()); -} - -// ============================================================================ -// Scheduled/Delayed Steps Tests (#24) -// ============================================================================ - -#[tokio::test] -async fn test_step_delay() { - let engine: WorkflowEngine = 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 = 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 = WorkflowEngine::new(); - - let executed = Arc::new(AtomicU32::new(0)); - let executed_clone = Arc::clone(&executed); - - struct TrackingStep { - counter: Arc, - } - - #[async_trait::async_trait] - impl StepExecutor for TrackingStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - 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 = WorkflowEngine::new(); - - let executed = Arc::new(AtomicU32::new(0)); - let executed_clone = Arc::clone(&executed); - - struct TrackingStep { - counter: Arc, - } - - #[async_trait::async_trait] - impl StepExecutor for TrackingStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - 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 = WorkflowEngine::new(); - - // Step that sets test_key in context - struct SetKeyStep; - - #[async_trait::async_trait] - impl StepExecutor for SetKeyStep { - async fn execute( - &self, - context: &mut WorkflowContext, - ) -> WorkflowResult { - 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, - } - - #[async_trait::async_trait] - impl StepExecutor for TrackingStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - 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 = WorkflowEngine::new(); - - let start_times: Arc>> = - Arc::new(parking_lot::RwLock::new(Vec::new())); - let end_times: Arc>> = - 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 = WorkflowEngine::new(); - - let start_times: Arc>> = - Arc::new(parking_lot::RwLock::new(Vec::new())); - let end_times: Arc>> = - 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 = WorkflowEngine::new(); - - let c_executed = Arc::new(AtomicU32::new(0)); - let c_executed_clone = Arc::clone(&c_executed); - - struct TrackingStep { - counter: Arc, - } - - #[async_trait::async_trait] - impl StepExecutor for TrackingStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - 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 = WorkflowEngine::new(); - - let c_executed = Arc::new(AtomicU32::new(0)); - let c_executed_clone = Arc::clone(&c_executed); - - struct TrackingStep { - counter: Arc, - } - - #[async_trait::async_trait] - impl StepExecutor for TrackingStep { - async fn execute( - &self, - _context: &mut WorkflowContext, - ) -> WorkflowResult { - 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); -}