987 lines
35 KiB
Rust
987 lines
35 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},
|
|
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<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)
|
|
}
|
|
}
|
|
|
|
/// 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<Duration> {
|
|
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<D>`)
|
|
///
|
|
/// # 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<D: WorkflowData, S: StateStore<D> = InMemoryStore<D>> {
|
|
definitions: Arc<RwLock<HashMap<WorkflowId, Arc<WorkflowDefinition<D>>>>>,
|
|
state_store: S,
|
|
event_bus: Arc<EventBus>,
|
|
/// Shutdown signal sender - when true, engine is shutting down
|
|
shutdown_tx: Arc<watch::Sender<bool>>,
|
|
/// Shutdown signal receiver for cloning to tasks
|
|
shutdown_rx: watch::Receiver<bool>,
|
|
/// Count of active workflow executions
|
|
active_workflows: Arc<AtomicUsize>,
|
|
_phantom: PhantomData<D>,
|
|
}
|
|
|
|
impl<D: WorkflowData> WorkflowEngine<D, InMemoryStore<D>> {
|
|
pub fn new() -> Self {
|
|
Self::with_store(InMemoryStore::new())
|
|
}
|
|
}
|
|
|
|
impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
|
|
/// 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);
|
|
}
|
|
|
|
/// 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<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
|
|
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<D>,
|
|
) -> 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<EventBus> {
|
|
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<WorkflowInstanceId> {
|
|
// 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<WorkflowDefinition<D>>,
|
|
) -> 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).await? {
|
|
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;
|
|
})
|
|
.await?;
|
|
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
|
|
let mut tasks_launched = 0;
|
|
for step_idx in ready_step_indices {
|
|
let step = &definition.steps[step_idx];
|
|
tracker.write().running.insert(step.id.clone());
|
|
tasks_launched += 1;
|
|
|
|
let 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 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"
|
|
);
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
let should_wait = tasks_launched > 0 || !tracker.read().running.is_empty();
|
|
if should_wait {
|
|
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;
|
|
})
|
|
.await?;
|
|
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;
|
|
})
|
|
.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<D>,
|
|
definition: &WorkflowDefinition<D>,
|
|
) -> 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).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<WorkflowState<D>> {
|
|
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<String, String> {
|
|
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<AtomicUsize>,
|
|
}
|
|
|
|
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<D> + 'static> {
|
|
engine: &'a WorkflowEngine<D, S>,
|
|
committed: bool,
|
|
}
|
|
|
|
impl<'a, D: WorkflowData, S: StateStore<D> + 'static> StartGuard<'a, D, S> {
|
|
fn new(engine: &'a WorkflowEngine<D, S>) -> Self {
|
|
engine.active_workflows.fetch_add(1, Ordering::AcqRel);
|
|
Self {
|
|
engine,
|
|
committed: false,
|
|
}
|
|
}
|
|
|
|
fn commit(mut self) {
|
|
self.committed = true;
|
|
}
|
|
}
|
|
|
|
impl<D: WorkflowData, S: StateStore<D> + '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<WorkflowEngine>` rather
|
|
/// than cloning.
|
|
impl<D: WorkflowData, S: StateStore<D> + 'static> Clone for WorkflowEngine<D, S> {
|
|
fn clone(&self) -> Self {
|
|
self.clone_for_execution()
|
|
}
|
|
}
|
|
|
|
impl<D: WorkflowData> Default for WorkflowEngine<D, InMemoryStore<D>> {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl<D: WorkflowData, S: StateStore<D> + 'static> std::fmt::Debug for WorkflowEngine<D, S> {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("WorkflowEngine")
|
|
.field("definitions_count", &self.definitions.read().len())
|
|
.finish()
|
|
}
|
|
}
|