626 lines
22 KiB
Rust
626 lines
22 KiB
Rust
//! 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},
|
|
sync::Arc,
|
|
time::Duration,
|
|
};
|
|
|
|
use backoff::{backoff::Backoff, ExponentialBackoffBuilder};
|
|
use chrono::Utc;
|
|
use parking_lot::RwLock;
|
|
use tokio::{sync::mpsc, time::timeout};
|
|
|
|
use super::{
|
|
definition::{StepDefinition, WorkflowDefinition},
|
|
event::{EventBus, WorkflowEvent},
|
|
state::WorkflowStateStore,
|
|
types::*,
|
|
};
|
|
|
|
#[derive(Default)]
|
|
struct StepTracker {
|
|
completed: HashSet<StepId>,
|
|
failed: HashSet<StepId>,
|
|
skipped: HashSet<StepId>,
|
|
running: HashSet<StepId>,
|
|
}
|
|
|
|
impl StepTracker {
|
|
fn total_processed(&self) -> usize {
|
|
self.completed.len() + self.failed.len() + self.skipped.len()
|
|
}
|
|
|
|
fn is_step_processable(&self, step_id: &StepId) -> bool {
|
|
!self.completed.contains(step_id)
|
|
&& !self.failed.contains(step_id)
|
|
&& !self.skipped.contains(step_id)
|
|
&& !self.running.contains(step_id)
|
|
}
|
|
|
|
fn are_dependencies_satisfied(&self, depends_on: &[StepId]) -> bool {
|
|
depends_on
|
|
.iter()
|
|
.all(|dep| self.completed.contains(dep) || self.skipped.contains(dep))
|
|
}
|
|
|
|
fn has_failed_dependency(&self, depends_on: &[StepId]) -> bool {
|
|
depends_on.iter().any(|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<Duration> {
|
|
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<Duration> {
|
|
let next = self.current;
|
|
self.current = (self.current + self.increment).min(self.max);
|
|
Some(next)
|
|
}
|
|
}
|
|
|
|
/// Main workflow execution engine
|
|
pub struct WorkflowEngine {
|
|
definitions: Arc<RwLock<HashMap<WorkflowId, Arc<WorkflowDefinition>>>>,
|
|
state_store: WorkflowStateStore,
|
|
event_bus: Arc<EventBus>,
|
|
}
|
|
|
|
impl WorkflowEngine {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
definitions: Arc::new(RwLock::new(HashMap::new())),
|
|
state_store: WorkflowStateStore::new(),
|
|
event_bus: Arc::new(EventBus::new()),
|
|
}
|
|
}
|
|
|
|
/// 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.
|
|
///
|
|
/// # 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<Duration>,
|
|
interval: Option<Duration>,
|
|
) -> 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
|
|
|
|
tokio::spawn(async move {
|
|
let mut ticker = tokio::time::interval(interval);
|
|
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
|
|
loop {
|
|
ticker.tick().await;
|
|
state_store.cleanup_old_workflows(ttl);
|
|
}
|
|
})
|
|
}
|
|
|
|
/// Register a workflow definition
|
|
pub fn register_workflow(&self, mut definition: WorkflowDefinition) -> Result<(), String> {
|
|
// 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<EventBus> {
|
|
Arc::clone(&self.event_bus)
|
|
}
|
|
|
|
/// Get the state store
|
|
pub fn state_store(&self) -> &WorkflowStateStore {
|
|
&self.state_store
|
|
}
|
|
|
|
/// Start a new workflow instance
|
|
pub async fn start_workflow(
|
|
&self,
|
|
definition_id: WorkflowId,
|
|
context: WorkflowContext,
|
|
) -> WorkflowResult<WorkflowInstanceId> {
|
|
// Get workflow definition
|
|
let definition = {
|
|
let definitions = self.definitions.read();
|
|
definitions
|
|
.get(&definition_id)
|
|
.cloned()
|
|
.ok_or_else(|| WorkflowError::DefinitionNotFound(definition_id.clone()))?
|
|
};
|
|
|
|
// Create new workflow instance
|
|
let instance_id = context.instance_id;
|
|
let mut state = WorkflowState::new(instance_id, definition_id.clone());
|
|
state.status = WorkflowStatus::Running;
|
|
state.context = context;
|
|
|
|
// Initialize step states
|
|
for step in &definition.steps {
|
|
state
|
|
.step_states
|
|
.insert(step.id.clone(), StepState::default());
|
|
}
|
|
|
|
// Save initial state
|
|
self.state_store.save(state)?;
|
|
|
|
// Emit workflow started event
|
|
self.event_bus
|
|
.publish(WorkflowEvent::WorkflowStarted {
|
|
instance_id,
|
|
definition_id,
|
|
})
|
|
.await;
|
|
|
|
// Execute workflow in background
|
|
let engine = self.clone_for_execution();
|
|
let def = Arc::clone(&definition);
|
|
tokio::spawn(async move {
|
|
if let Err(e) = engine.execute_workflow(instance_id, def).await {
|
|
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<WorkflowDefinition>,
|
|
) -> WorkflowResult<()> {
|
|
let start_time = std::time::Instant::now();
|
|
let step_count = definition.steps.len();
|
|
|
|
let tracker: Arc<RwLock<StepTracker>> = 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<usize> = definition
|
|
.get_initial_step_indices()
|
|
.iter()
|
|
.copied()
|
|
.collect();
|
|
|
|
loop {
|
|
if self.state_store.is_cancelled(instance_id)? {
|
|
self.event_bus
|
|
.publish(WorkflowEvent::WorkflowCancelled { instance_id })
|
|
.await;
|
|
return Ok(());
|
|
}
|
|
|
|
// Find ready steps from pending_check (not all steps)
|
|
let (ready_step_indices, total_processed, running_count) = {
|
|
let t = tracker.read();
|
|
|
|
// Only check steps in pending_check, not all steps
|
|
let ready: Vec<usize> = pending_check
|
|
.drain(..)
|
|
.filter(|&idx| {
|
|
let step = &definition.steps[idx];
|
|
t.is_step_processable(&step.id)
|
|
&& t.are_dependencies_satisfied(&step.depends_on)
|
|
&& !t.has_failed_dependency(&step.depends_on)
|
|
})
|
|
.collect();
|
|
|
|
(ready, t.total_processed(), t.running.len())
|
|
};
|
|
|
|
// Check if we're done
|
|
if total_processed == step_count {
|
|
break;
|
|
}
|
|
|
|
// Handle blocked workflow (no ready steps, none running, but work remains)
|
|
if ready_step_indices.is_empty() && running_count == 0 && pending_check.is_empty() {
|
|
let failed_step = tracker.read().failed.iter().next().cloned();
|
|
let error_message = if failed_step.is_some() {
|
|
"Workflow failed due to step dependency failure".to_string()
|
|
} else {
|
|
"Workflow deadlocked: no steps ready and none running. This may indicate a scheduler bug.".to_string()
|
|
};
|
|
|
|
self.state_store.update(instance_id, |s| {
|
|
s.status = WorkflowStatus::Failed;
|
|
})?;
|
|
self.event_bus
|
|
.publish(WorkflowEvent::WorkflowFailed {
|
|
instance_id,
|
|
failed_step: failed_step
|
|
.unwrap_or_else(|| StepId::new("internal_scheduler")),
|
|
error: error_message,
|
|
})
|
|
.await;
|
|
return Ok(());
|
|
}
|
|
|
|
// Launch ready steps in parallel
|
|
for step_idx in ready_step_indices {
|
|
let step = &definition.steps[step_idx];
|
|
tracker.write().running.insert(step.id.clone());
|
|
|
|
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];
|
|
let result = engine
|
|
.execute_step_with_retry(instance_id, step, &def)
|
|
.await;
|
|
|
|
{
|
|
let mut t = tracker.write();
|
|
t.running.remove(&step_id);
|
|
|
|
match result {
|
|
Ok(StepResult::Success) => {
|
|
t.completed.insert(step_id.clone());
|
|
}
|
|
Ok(StepResult::Skip) => {
|
|
t.skipped.insert(step_id.clone());
|
|
}
|
|
Ok(StepResult::Failure) | Err(_) => match step.on_failure {
|
|
FailureAction::FailWorkflow | FailureAction::RetryIndefinitely => {
|
|
t.failed.insert(step_id.clone());
|
|
}
|
|
FailureAction::ContinueNextStep => {
|
|
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;
|
|
}
|
|
}) {
|
|
tracing::warn!(
|
|
step_id = %step_id,
|
|
error = ?e,
|
|
"Failed to update step state to Skipped"
|
|
);
|
|
}
|
|
t.skipped.insert(step_id.clone());
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
let signal = match result {
|
|
Ok(r) => r,
|
|
Err(_) => StepResult::Failure,
|
|
};
|
|
let _ = tx.send((step_id, signal)).await;
|
|
});
|
|
}
|
|
|
|
// Wait for at least one step to complete (if any running)
|
|
if !tracker.read().running.is_empty() {
|
|
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);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
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;
|
|
})?;
|
|
self.event_bus
|
|
.publish(WorkflowEvent::WorkflowFailed {
|
|
instance_id,
|
|
failed_step: step.clone(),
|
|
error: "One or more steps failed".to_string(),
|
|
})
|
|
.await;
|
|
} else {
|
|
self.state_store.update(instance_id, |s| {
|
|
s.status = WorkflowStatus::Completed;
|
|
})?;
|
|
|
|
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<StepResult> {
|
|
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)? {
|
|
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());
|
|
}
|
|
})?;
|
|
|
|
// 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)?;
|
|
|
|
// 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 = std::mem::replace(&mut context, WorkflowContext::new(instance_id));
|
|
})?;
|
|
|
|
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());
|
|
}
|
|
})?;
|
|
|
|
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());
|
|
}
|
|
}
|
|
})?;
|
|
|
|
// 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) -> Box<dyn Backoff + Send> {
|
|
match strategy {
|
|
BackoffStrategy::Fixed(duration) => Box::new(FixedBackoff(*duration)),
|
|
BackoffStrategy::Exponential { base, max } => {
|
|
let backoff = ExponentialBackoffBuilder::new()
|
|
.with_initial_interval(*base)
|
|
.with_max_interval(*max)
|
|
.with_max_elapsed_time(None)
|
|
.build();
|
|
Box::new(backoff)
|
|
}
|
|
BackoffStrategy::Linear { increment, max } => {
|
|
Box::new(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;
|
|
})?;
|
|
|
|
self.event_bus
|
|
.publish(WorkflowEvent::WorkflowCancelled { instance_id })
|
|
.await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get workflow status
|
|
pub fn get_status(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState> {
|
|
self.state_store.load(instance_id)
|
|
}
|
|
|
|
/// 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),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for WorkflowEngine {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl 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())
|
|
.field("state_count", &self.state_store.count())
|
|
.finish()
|
|
}
|
|
}
|