use super::graph::WorkflowGraph;
use super::node::{NodeType, WorkflowNode};
use super::state::{
ExecutionRecord, NodeExecutionRecord, NodeResult, NodeStatus, WorkflowContext, WorkflowStatus,
WorkflowValue,
};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::{RwLock, Semaphore, mpsc, oneshot};
use tracing::{error, info, warn};
#[derive(Debug, Clone)]
pub struct ExecutorConfig {
pub max_parallelism: usize,
pub stop_on_failure: bool,
pub enable_checkpoints: bool,
pub checkpoint_interval: usize,
pub execution_timeout_ms: Option<u64>,
}
impl Default for ExecutorConfig {
fn default() -> Self {
Self {
max_parallelism: 10,
stop_on_failure: true,
enable_checkpoints: true,
checkpoint_interval: 5,
execution_timeout_ms: None,
}
}
}
#[derive(Debug, Clone)]
pub enum ExecutionEvent {
WorkflowStarted {
workflow_id: String,
execution_id: String,
},
WorkflowCompleted {
workflow_id: String,
execution_id: String,
status: WorkflowStatus,
},
NodeStarted { node_id: String },
NodeCompleted { node_id: String, result: NodeResult },
NodeFailed { node_id: String, error: String },
CheckpointCreated { label: String },
ExternalEvent {
event_type: String,
data: WorkflowValue,
},
}
pub struct WorkflowExecutor {
config: ExecutorConfig,
event_tx: Option<mpsc::Sender<ExecutionEvent>>,
sub_workflows: Arc<RwLock<HashMap<String, Arc<WorkflowGraph>>>>,
event_waiters: Arc<RwLock<HashMap<String, Vec<oneshot::Sender<WorkflowValue>>>>>,
semaphore: Arc<Semaphore>,
}
impl WorkflowExecutor {
pub fn new(config: ExecutorConfig) -> Self {
let semaphore = Arc::new(Semaphore::new(config.max_parallelism));
Self {
config,
event_tx: None,
sub_workflows: Arc::new(RwLock::new(HashMap::new())),
event_waiters: Arc::new(RwLock::new(HashMap::new())),
semaphore,
}
}
pub fn with_event_sender(mut self, tx: mpsc::Sender<ExecutionEvent>) -> Self {
self.event_tx = Some(tx);
self
}
pub async fn register_sub_workflow(&self, id: &str, graph: WorkflowGraph) {
let mut workflows = self.sub_workflows.write().await;
workflows.insert(id.to_string(), Arc::new(graph));
}
async fn emit_event(&self, event: ExecutionEvent) {
if let Some(ref tx) = self.event_tx {
let _ = tx.send(event).await;
}
}
pub async fn send_external_event(&self, event_type: &str, data: WorkflowValue) {
let mut waiters = self.event_waiters.write().await;
if let Some(senders) = waiters.remove(event_type) {
for sender in senders {
let _ = sender.send(data.clone());
}
}
}
pub async fn execute(
&self,
graph: &WorkflowGraph,
input: WorkflowValue,
) -> Result<ExecutionRecord, String> {
let start_time = Instant::now();
let ctx = WorkflowContext::new(&graph.id);
ctx.set_input(input.clone()).await;
self.emit_event(ExecutionEvent::WorkflowStarted {
workflow_id: graph.id.clone(),
execution_id: ctx.execution_id.clone(),
})
.await;
info!(
"Starting workflow execution: {} ({})",
graph.name, ctx.execution_id
);
if let Err(errors) = graph.validate() {
let error_msg = errors.join("; ");
error!("Workflow validation failed: {}", error_msg);
return Err(error_msg);
}
let start_node_id = graph
.start_node()
.ok_or_else(|| "No start node".to_string())?;
let mut execution_record = ExecutionRecord {
execution_id: ctx.execution_id.clone(),
workflow_id: graph.id.clone(),
started_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64,
ended_at: None,
status: WorkflowStatus::Running,
node_records: Vec::new(),
};
let result = self
.execute_from_node(graph, &ctx, start_node_id, input, &mut execution_record)
.await;
let duration = start_time.elapsed();
execution_record.ended_at = Some(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64,
);
match result {
Ok(_) => {
execution_record.status = WorkflowStatus::Completed;
info!("Workflow {} completed in {:?}", graph.name, duration);
}
Err(ref e) => {
execution_record.status = WorkflowStatus::Failed(e.clone());
error!("Workflow {} failed: {}", graph.name, e);
}
}
self.emit_event(ExecutionEvent::WorkflowCompleted {
workflow_id: graph.id.clone(),
execution_id: ctx.execution_id.clone(),
status: execution_record.status.clone(),
})
.await;
Ok(execution_record)
}
async fn execute_from_node(
&self,
graph: &WorkflowGraph,
ctx: &WorkflowContext,
start_node_id: &str,
initial_input: WorkflowValue,
record: &mut ExecutionRecord,
) -> Result<WorkflowValue, String> {
let mut current_node_id = start_node_id.to_string();
let mut current_input = initial_input;
loop {
let node = graph
.get_node(¤t_node_id)
.ok_or_else(|| format!("Node {} not found", current_node_id))?;
ctx.set_node_status(¤t_node_id, NodeStatus::Running)
.await;
self.emit_event(ExecutionEvent::NodeStarted {
node_id: current_node_id.clone(),
})
.await;
let start_time = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
let result = match node.node_type() {
NodeType::Parallel => {
self.execute_parallel(graph, ctx, node, current_input.clone(), record)
.await
}
NodeType::Join => self.execute_join(graph, ctx, node, record).await,
NodeType::SubWorkflow => {
self.execute_sub_workflow(graph, ctx, node, current_input.clone(), record)
.await
}
NodeType::Wait => self.execute_wait(ctx, node, current_input.clone()).await,
_ => {
let result = node.execute(ctx, current_input.clone()).await;
ctx.set_node_output(¤t_node_id, result.output.clone())
.await;
ctx.set_node_status(¤t_node_id, result.status.clone())
.await;
self.emit_event(ExecutionEvent::NodeCompleted {
node_id: current_node_id.clone(),
result: result.clone(),
})
.await;
if result.status.is_success() {
Ok(result.output)
} else {
Err(result.error.unwrap_or_else(|| "Unknown error".to_string()))
}
}
};
let end_time = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
record.node_records.push(NodeExecutionRecord {
node_id: current_node_id.clone(),
started_at: start_time,
ended_at: end_time,
status: ctx
.get_node_status(¤t_node_id)
.await
.unwrap_or(NodeStatus::Pending),
retry_count: 0,
});
if self.config.enable_checkpoints
&& record
.node_records
.len()
.is_multiple_of(self.config.checkpoint_interval)
{
let label = format!("auto_checkpoint_{}", record.node_records.len());
ctx.create_checkpoint(&label).await;
self.emit_event(ExecutionEvent::CheckpointCreated { label })
.await;
}
match result {
Ok(output) => {
let next = self.determine_next_node(graph, node, &output).await;
match next {
Some(next_node_id) => {
current_node_id = next_node_id;
current_input = output;
}
None => {
return Ok(output);
}
}
}
Err(e) => {
if let Some(error_handler) = graph.get_error_handler(¤t_node_id) {
warn!(
"Node {} failed, executing error handler: {}",
current_node_id, error_handler
);
let error_input = WorkflowValue::Map({
let mut m = HashMap::new();
m.insert("error".to_string(), WorkflowValue::String(e.clone()));
m.insert(
"node_id".to_string(),
WorkflowValue::String(current_node_id.clone()),
);
m
});
current_node_id = error_handler.to_string();
current_input = error_input;
} else if self.config.stop_on_failure {
return Err(e);
} else {
warn!("Node {} failed but continuing: {}", current_node_id, e);
if let Some(next_node_id) = graph.get_next_node(¤t_node_id, None) {
current_node_id = next_node_id.to_string();
current_input = WorkflowValue::Null;
} else {
return Err(e);
}
}
}
}
}
}
async fn determine_next_node(
&self,
graph: &WorkflowGraph,
node: &WorkflowNode,
output: &WorkflowValue,
) -> Option<String> {
let node_id = node.id();
match node.node_type() {
NodeType::Condition => {
let condition = output.as_str().unwrap_or("false");
graph
.get_next_node(node_id, Some(condition))
.map(|s| s.to_string())
}
NodeType::End => {
None
}
_ => {
graph.get_next_node(node_id, None).map(|s| s.to_string())
}
}
}
async fn execute_parallel(
&self,
graph: &WorkflowGraph,
ctx: &WorkflowContext,
node: &WorkflowNode,
input: WorkflowValue,
record: &mut ExecutionRecord,
) -> Result<WorkflowValue, String> {
let branches = node.parallel_branches();
if branches.is_empty() {
let edges = graph.get_outgoing_edges(node.id());
let branch_ids: Vec<String> = edges.iter().map(|e| e.to.clone()).collect();
if branch_ids.is_empty() {
return Ok(input);
}
return self
.execute_branches_parallel(graph, ctx, &branch_ids, input, record)
.await;
}
self.execute_branches_parallel(graph, ctx, branches, input, record)
.await
}
async fn execute_branches_parallel(
&self,
graph: &WorkflowGraph,
ctx: &WorkflowContext,
branches: &[String],
input: WorkflowValue,
_record: &mut ExecutionRecord,
) -> Result<WorkflowValue, String> {
let mut results = HashMap::new();
let mut errors = Vec::new();
for branch_id in branches {
if let Some(node) = graph.get_node(branch_id) {
let result = node.execute(ctx, input.clone()).await;
ctx.set_node_output(branch_id, result.output.clone()).await;
ctx.set_node_status(branch_id, result.status.clone()).await;
if result.status.is_success() {
results.insert(branch_id.clone(), result.output);
} else {
errors.push(format!(
"{}: {}",
branch_id,
result.error.unwrap_or_else(|| "Unknown error".to_string())
));
}
} else {
errors.push(format!("Node {} not found", branch_id));
}
}
if !errors.is_empty() && self.config.stop_on_failure {
return Err(errors.join("; "));
}
Ok(WorkflowValue::Map(results))
}
async fn execute_join(
&self,
_graph: &WorkflowGraph,
ctx: &WorkflowContext,
node: &WorkflowNode,
_record: &mut ExecutionRecord,
) -> Result<WorkflowValue, String> {
let wait_for = node.join_nodes();
let mut all_completed = false;
let mut attempts = 0;
const MAX_ATTEMPTS: u32 = 1000;
while !all_completed && attempts < MAX_ATTEMPTS {
all_completed = true;
for node_id in wait_for {
match ctx.get_node_status(node_id).await {
Some(status) if status.is_terminal() => {}
_ => {
all_completed = false;
break;
}
}
}
if !all_completed {
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
attempts += 1;
}
}
if !all_completed {
return Err("Join timeout waiting for nodes".to_string());
}
let outputs = ctx
.get_node_outputs(&wait_for.iter().map(|s| s.as_str()).collect::<Vec<_>>())
.await;
let result = node.execute(ctx, WorkflowValue::Map(outputs)).await;
ctx.set_node_output(node.id(), result.output.clone()).await;
ctx.set_node_status(node.id(), result.status.clone()).await;
if result.status.is_success() {
Ok(result.output)
} else {
Err(result.error.unwrap_or_else(|| "Join failed".to_string()))
}
}
async fn execute_sub_workflow(
&self,
_graph: &WorkflowGraph,
ctx: &WorkflowContext,
node: &WorkflowNode,
input: WorkflowValue,
_record: &mut ExecutionRecord,
) -> Result<WorkflowValue, String> {
let sub_workflow_id = node
.sub_workflow_id()
.ok_or_else(|| "No sub-workflow specified".to_string())?;
let workflows = self.sub_workflows.read().await;
let sub_graph = workflows
.get(sub_workflow_id)
.ok_or_else(|| format!("Sub-workflow {} not found", sub_workflow_id))?
.clone();
drop(workflows);
info!("Executing sub-workflow: {}", sub_workflow_id);
let _sub_record = self.execute_parallel_workflow(&sub_graph, input).await?;
let output = if let Some(end_node) = sub_graph.end_nodes().first() {
ctx.get_node_output(end_node)
.await
.unwrap_or(WorkflowValue::Null)
} else {
WorkflowValue::Null
};
ctx.set_node_output(node.id(), output.clone()).await;
ctx.set_node_status(node.id(), NodeStatus::Completed).await;
Ok(output)
}
async fn execute_wait(
&self,
ctx: &WorkflowContext,
node: &WorkflowNode,
_input: WorkflowValue,
) -> Result<WorkflowValue, String> {
let event_type = node
.wait_event_type()
.ok_or_else(|| "No event type specified".to_string())?;
info!("Waiting for event: {}", event_type);
let (tx, rx) = oneshot::channel();
{
let mut waiters = self.event_waiters.write().await;
waiters.entry(event_type.to_string()).or_default().push(tx);
}
let timeout = node.config.timeout.execution_timeout_ms;
let result = if timeout > 0 {
tokio::time::timeout(std::time::Duration::from_millis(timeout), rx)
.await
.map_err(|_| "Wait timeout".to_string())?
.map_err(|_| "Wait cancelled".to_string())?
} else {
rx.await.map_err(|_| "Wait cancelled".to_string())?
};
ctx.set_node_output(node.id(), result.clone()).await;
ctx.set_node_status(node.id(), NodeStatus::Completed).await;
Ok(result)
}
pub async fn execute_parallel_workflow(
&self,
graph: &WorkflowGraph,
input: WorkflowValue,
) -> Result<ExecutionRecord, String> {
let ctx = WorkflowContext::new(&graph.id);
ctx.set_input(input.clone()).await;
let start_time = Instant::now();
info!(
"Starting layered workflow execution: {} ({})",
graph.name, ctx.execution_id
);
let groups = graph.get_parallel_groups();
let mut execution_record = ExecutionRecord {
execution_id: ctx.execution_id.clone(),
workflow_id: graph.id.clone(),
started_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64,
ended_at: None,
status: WorkflowStatus::Running,
node_records: Vec::new(),
};
for group in groups {
for node_id in group {
let node_start_time = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
let result = if let Some(node) = graph.get_node(&node_id) {
let predecessors = graph.get_predecessors(&node_id);
let node_input = if predecessors.is_empty() {
ctx.get_input().await
} else if predecessors.len() == 1 {
ctx.get_node_output(predecessors[0])
.await
.unwrap_or(WorkflowValue::Null)
} else {
let outputs = ctx.get_node_outputs(&predecessors).await;
WorkflowValue::Map(outputs)
};
let result = node.execute(&ctx, node_input).await;
ctx.set_node_output(&node_id, result.output.clone()).await;
ctx.set_node_status(&node_id, result.status.clone()).await;
result
} else {
NodeResult::failed(&node_id, "Node not found", 0)
};
let node_end_time = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
let record_entry = NodeExecutionRecord {
node_id: node_id.clone(),
started_at: node_start_time,
ended_at: node_end_time,
status: result.status.clone(),
retry_count: result.retry_count,
};
execution_record.node_records.push(record_entry);
if !result.status.is_success() && self.config.stop_on_failure {
execution_record.status = WorkflowStatus::Failed(
result.error.unwrap_or_else(|| "Unknown error".to_string()),
);
execution_record.ended_at = Some(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64,
);
return Ok(execution_record);
}
}
}
let duration = start_time.elapsed();
execution_record.status = WorkflowStatus::Completed;
execution_record.ended_at = Some(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64,
);
info!(
"Layered workflow {} completed in {:?}",
graph.name, duration
);
Ok(execution_record)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_simple_workflow_execution() {
let mut graph = WorkflowGraph::new("test", "Simple Workflow");
graph.add_node(WorkflowNode::start("start"));
graph.add_node(WorkflowNode::task(
"double",
"Double",
|_ctx, input| async move {
let value = input.as_i64().unwrap_or(0);
Ok(WorkflowValue::Int(value * 2))
},
));
graph.add_node(WorkflowNode::task(
"add_ten",
"Add Ten",
|_ctx, input| async move {
let value = input.as_i64().unwrap_or(0);
Ok(WorkflowValue::Int(value + 10))
},
));
graph.add_node(WorkflowNode::end("end"));
graph.connect("start", "double");
graph.connect("double", "add_ten");
graph.connect("add_ten", "end");
let executor = WorkflowExecutor::new(ExecutorConfig::default());
let result = executor
.execute(&graph, WorkflowValue::Int(5))
.await
.unwrap();
assert!(matches!(result.status, WorkflowStatus::Completed));
}
#[tokio::test]
async fn test_conditional_workflow() {
let mut graph = WorkflowGraph::new("test", "Conditional Workflow");
graph.add_node(WorkflowNode::start("start"));
graph.add_node(WorkflowNode::condition(
"check",
"Check Value",
|_ctx, input| async move { input.as_i64().unwrap_or(0) > 10 },
));
graph.add_node(WorkflowNode::task(
"high",
"High Path",
|_ctx, _input| async move { Ok(WorkflowValue::String("high".to_string())) },
));
graph.add_node(WorkflowNode::task(
"low",
"Low Path",
|_ctx, _input| async move { Ok(WorkflowValue::String("low".to_string())) },
));
graph.add_node(WorkflowNode::end("end"));
graph.connect("start", "check");
graph.connect_conditional("check", "high", "true");
graph.connect_conditional("check", "low", "false");
graph.connect("high", "end");
graph.connect("low", "end");
let executor = WorkflowExecutor::new(ExecutorConfig::default());
let result = executor
.execute(&graph, WorkflowValue::Int(20))
.await
.unwrap();
assert!(matches!(result.status, WorkflowStatus::Completed));
}
}