863 lines
24 KiB
Rust
863 lines
24 KiB
Rust
//! DAG (Directed Acyclic Graph) scheduling system with Airflow-like capabilities
|
|
//!
|
|
//! Provides comprehensive task orchestration including:
|
|
//! - Task dependency resolution with topological sorting
|
|
//! - Parallel execution of independent tasks
|
|
//! - Retry mechanisms with exponential backoff
|
|
//! - Conditional task execution based on upstream results
|
|
//! - Dynamic DAG generation and modification
|
|
//! - Resource allocation and priority scheduling
|
|
|
|
use chrono::{DateTime, Utc};
|
|
use petgraph::algo::{is_cyclic_directed, toposort};
|
|
use petgraph::{Directed, Graph};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
use tokio::sync::RwLock;
|
|
use tracing::{debug, instrument, warn};
|
|
use uuid::Uuid;
|
|
|
|
use crate::{
|
|
EtlError, Result,
|
|
connectors::{DataSink, DataSource},
|
|
monitoring::EtlMetrics,
|
|
state::StateManager,
|
|
transform::Transformation,
|
|
};
|
|
|
|
/// Configuration for DAG scheduling
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct DagConfig {
|
|
/// Maximum number of task retries
|
|
pub max_retries: u32,
|
|
/// Base retry delay in milliseconds
|
|
pub base_retry_delay_ms: u64,
|
|
/// Maximum retry delay in milliseconds
|
|
pub max_retry_delay_ms: u64,
|
|
/// Task execution timeout in milliseconds
|
|
pub task_timeout_ms: u64,
|
|
/// Enable dynamic DAG modifications
|
|
pub allow_dynamic_dag: bool,
|
|
/// Maximum DAG depth
|
|
pub max_dag_depth: usize,
|
|
/// Enable priority-based scheduling
|
|
pub priority_scheduling: bool,
|
|
}
|
|
|
|
impl Default for DagConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_retries: 3,
|
|
base_retry_delay_ms: 1000,
|
|
max_retry_delay_ms: 60000,
|
|
task_timeout_ms: 300000, // 5 minutes
|
|
allow_dynamic_dag: true,
|
|
max_dag_depth: 50,
|
|
priority_scheduling: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// DAG scheduler responsible for task orchestration
|
|
pub struct DagScheduler {
|
|
/// Configuration
|
|
config: DagConfig,
|
|
/// State manager for persistence
|
|
state_manager: Arc<StateManager>,
|
|
/// Metrics collector
|
|
metrics: Arc<EtlMetrics>,
|
|
/// Active DAG executions
|
|
active_executions: Arc<RwLock<HashMap<Uuid, DagExecution>>>,
|
|
}
|
|
|
|
/// Represents a DAG of tasks with dependencies
|
|
#[derive(Debug, Clone)]
|
|
pub struct TaskGraph {
|
|
/// Tasks in the graph
|
|
pub tasks: HashMap<String, Task>,
|
|
/// Task dependencies (`task_id` -> list of dependencies)
|
|
pub dependencies: HashMap<String, Vec<String>>,
|
|
/// Graph metadata
|
|
pub metadata: GraphMetadata,
|
|
}
|
|
|
|
/// Metadata for a task graph
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct GraphMetadata {
|
|
/// Graph identifier
|
|
pub graph_id: String,
|
|
/// Graph name
|
|
pub name: String,
|
|
/// Description
|
|
pub description: Option<String>,
|
|
/// Tags for categorization
|
|
pub tags: Vec<String>,
|
|
/// Creation timestamp
|
|
pub created_at: DateTime<Utc>,
|
|
/// Last modified timestamp
|
|
pub modified_at: DateTime<Utc>,
|
|
/// Version
|
|
pub version: String,
|
|
}
|
|
|
|
/// Individual task in a DAG
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Task {
|
|
/// Task identifier
|
|
pub id: String,
|
|
/// Task name
|
|
pub name: String,
|
|
/// Task description
|
|
pub description: Option<String>,
|
|
/// Task type
|
|
pub task_type: TaskType,
|
|
/// Data source configuration
|
|
pub source: Option<DataSource>,
|
|
/// Data transformation configuration
|
|
pub transformation: Option<Transformation>,
|
|
/// Data sink configuration
|
|
pub sink: Option<DataSink>,
|
|
/// Task dependencies
|
|
pub dependencies: Vec<String>,
|
|
/// Task priority (higher number = higher priority)
|
|
pub priority: i32,
|
|
/// Retry configuration
|
|
pub retry_config: RetryConfig,
|
|
/// Resource requirements
|
|
pub resources: ResourceRequirements,
|
|
/// Conditional execution configuration
|
|
pub condition: Option<TaskCondition>,
|
|
/// Task timeout override
|
|
pub timeout_ms: Option<u64>,
|
|
/// Task metadata
|
|
pub metadata: HashMap<String, String>,
|
|
/// Task configuration parameters
|
|
pub config: HashMap<String, serde_json::Value>,
|
|
}
|
|
|
|
/// Task priority levels
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Default)]
|
|
pub enum TaskPriority {
|
|
Low = 1,
|
|
#[default]
|
|
Medium = 2,
|
|
High = 3,
|
|
Critical = 4,
|
|
}
|
|
|
|
/// Types of tasks that can be executed
|
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
|
pub enum TaskType {
|
|
/// Extract data from a source
|
|
Extract,
|
|
/// Transform data using specified transformations
|
|
Transform,
|
|
/// Load data to a destination
|
|
Load,
|
|
/// Combined ETL operation
|
|
Etl,
|
|
/// Data quality check
|
|
QualityCheck,
|
|
/// Custom task with user-defined logic
|
|
Custom(String),
|
|
/// Conditional branching task
|
|
Branch,
|
|
/// Wait/delay task
|
|
Wait(Duration),
|
|
}
|
|
|
|
impl std::fmt::Display for TaskType {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
Self::Extract => write!(f, "Extract"),
|
|
Self::Transform => write!(f, "Transform"),
|
|
Self::Load => write!(f, "Load"),
|
|
Self::Etl => write!(f, "ETL"),
|
|
Self::QualityCheck => write!(f, "QualityCheck"),
|
|
Self::Custom(name) => write!(f, "Custom({name})"),
|
|
Self::Branch => write!(f, "Branch"),
|
|
Self::Wait(duration) => write!(f, "Wait({duration:?})"),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl TaskType {
|
|
#[must_use]
|
|
pub fn as_str(&self) -> &str {
|
|
match self {
|
|
Self::Extract => "extract",
|
|
Self::Transform => "transform",
|
|
Self::Load => "load",
|
|
Self::Etl => "etl",
|
|
Self::QualityCheck => "quality_check",
|
|
Self::Custom(_) => "custom",
|
|
Self::Branch => "branch",
|
|
Self::Wait(_) => "wait",
|
|
}
|
|
}
|
|
|
|
pub fn from_str(s: &str) -> Result<Self> {
|
|
match s.to_lowercase().as_str() {
|
|
"extract" => Ok(Self::Extract),
|
|
"transform" => Ok(Self::Transform),
|
|
"load" => Ok(Self::Load),
|
|
"etl" => Ok(Self::Etl),
|
|
"quality_check" => Ok(Self::QualityCheck),
|
|
"branch" => Ok(Self::Branch),
|
|
_ if s.starts_with("custom:") => {
|
|
let name = s.strip_prefix("custom:").unwrap_or(s);
|
|
Ok(Self::Custom(name.to_string()))
|
|
}
|
|
_ => Err(EtlError::validation(format!("Unknown task type: {s}"))),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Task {
|
|
#[must_use]
|
|
pub fn new(id: String, task_type: TaskType) -> Self {
|
|
Self {
|
|
id,
|
|
name: String::new(),
|
|
description: None,
|
|
task_type,
|
|
source: None,
|
|
transformation: None,
|
|
sink: None,
|
|
dependencies: Vec::new(),
|
|
retry_config: RetryConfig::default(),
|
|
priority: 0,
|
|
resources: ResourceRequirements::default(),
|
|
condition: None,
|
|
timeout_ms: None,
|
|
metadata: HashMap::new(),
|
|
config: HashMap::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Retry configuration for tasks
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct RetryConfig {
|
|
/// Maximum number of retries
|
|
pub max_retries: u32,
|
|
/// Retry strategy
|
|
pub strategy: RetryStrategy,
|
|
/// Base delay between retries
|
|
pub base_delay_ms: u64,
|
|
/// Maximum delay between retries
|
|
pub max_delay_ms: u64,
|
|
/// Exponential backoff multiplier
|
|
pub backoff_multiplier: f64,
|
|
}
|
|
|
|
/// Retry strategies for failed tasks
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum RetryStrategy {
|
|
/// Fixed delay between retries
|
|
Fixed,
|
|
/// Exponential backoff with jitter
|
|
ExponentialBackoff,
|
|
/// Linear backoff
|
|
LinearBackoff,
|
|
/// Custom retry logic
|
|
Custom(String),
|
|
}
|
|
|
|
/// Resource requirements for task execution
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ResourceRequirements {
|
|
/// CPU cores required
|
|
pub cpu_cores: f32,
|
|
/// Memory required in MB
|
|
pub memory_mb: u64,
|
|
/// Disk space required in MB
|
|
pub disk_mb: u64,
|
|
/// GPU resources required
|
|
pub gpu_memory_mb: Option<u64>,
|
|
/// Network bandwidth required in Mbps
|
|
pub network_mbps: Option<u64>,
|
|
}
|
|
|
|
/// Condition for task execution
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TaskCondition {
|
|
/// Condition type
|
|
pub condition_type: ConditionType,
|
|
/// Condition expression
|
|
pub expression: String,
|
|
/// Task to execute if condition is true
|
|
pub on_true: Option<String>,
|
|
/// Task to execute if condition is false
|
|
pub on_false: Option<String>,
|
|
}
|
|
|
|
/// Types of task conditions
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum ConditionType {
|
|
/// Condition based on upstream task status
|
|
UpstreamStatus,
|
|
/// Condition based on data availability
|
|
DataAvailability,
|
|
/// Custom condition logic
|
|
Custom(String),
|
|
/// Time-based condition
|
|
TimeBased,
|
|
/// Resource availability condition
|
|
ResourceAvailability,
|
|
}
|
|
|
|
/// Status of a task execution
|
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
|
pub enum TaskStatus {
|
|
/// Task is pending execution
|
|
Pending,
|
|
/// Task is currently running
|
|
Running,
|
|
/// Task completed successfully
|
|
Success,
|
|
/// Task failed
|
|
Failed,
|
|
/// Task was skipped due to conditions
|
|
Skipped,
|
|
/// Task was cancelled
|
|
Cancelled,
|
|
/// Task timed out
|
|
TimedOut,
|
|
/// Task is waiting for retry
|
|
Retrying,
|
|
}
|
|
|
|
/// Result of task execution
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TaskResult {
|
|
/// Task identifier
|
|
pub task_id: String,
|
|
/// Execution status
|
|
pub status: TaskStatus,
|
|
/// Start time
|
|
pub start_time: DateTime<Utc>,
|
|
/// End time
|
|
pub end_time: Option<DateTime<Utc>>,
|
|
/// Records processed
|
|
pub records_processed: u64,
|
|
/// Bytes processed
|
|
pub bytes_processed: u64,
|
|
/// Error message if failed
|
|
pub error_message: Option<String>,
|
|
/// Output data location
|
|
pub output_location: Option<String>,
|
|
/// Task metrics
|
|
pub metrics: TaskMetrics,
|
|
}
|
|
|
|
/// Metrics for task execution
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TaskMetrics {
|
|
/// Execution duration in milliseconds
|
|
pub duration_ms: u64,
|
|
/// Peak memory usage in bytes
|
|
pub peak_memory_bytes: u64,
|
|
/// CPU time used in milliseconds
|
|
pub cpu_time_ms: u64,
|
|
/// I/O operations count
|
|
pub io_operations: u64,
|
|
/// Network bytes transferred
|
|
pub network_bytes: u64,
|
|
}
|
|
|
|
/// Execution context for a DAG
|
|
#[derive(Debug)]
|
|
pub struct DagExecution {
|
|
/// Execution identifier
|
|
pub execution_id: Uuid,
|
|
/// Task graph being executed
|
|
pub graph: TaskGraph,
|
|
/// Execution plan
|
|
pub plan: ExecutionPlan,
|
|
/// Task statuses
|
|
pub task_status: HashMap<String, TaskStatus>,
|
|
/// Task results
|
|
pub task_results: HashMap<String, TaskResult>,
|
|
/// Execution start time
|
|
pub start_time: Instant,
|
|
/// Current execution stage
|
|
pub current_stage: usize,
|
|
}
|
|
|
|
/// Execution plan for a DAG
|
|
#[derive(Debug, Clone)]
|
|
pub struct ExecutionPlan {
|
|
/// Execution stages (tasks that can run in parallel)
|
|
pub stages: Vec<ExecutionStage>,
|
|
/// Total estimated duration
|
|
pub estimated_duration: Duration,
|
|
/// Resource requirements per stage
|
|
pub stage_resources: Vec<ResourceRequirements>,
|
|
}
|
|
|
|
/// A stage of execution containing tasks that can run in parallel
|
|
#[derive(Debug, Clone)]
|
|
pub struct ExecutionStage {
|
|
/// Stage identifier
|
|
pub stage_id: usize,
|
|
/// Tasks in this stage
|
|
pub tasks: Vec<Task>,
|
|
/// Stage dependencies (previous stages that must complete)
|
|
pub dependencies: Vec<usize>,
|
|
}
|
|
|
|
/// Context for task execution
|
|
#[derive(Debug)]
|
|
pub struct ExecutionContext {
|
|
/// Task being executed
|
|
pub task: Task,
|
|
/// Execution start time
|
|
pub start_time: Instant,
|
|
/// Retry attempt number
|
|
pub retry_attempt: u32,
|
|
/// Execution metadata
|
|
pub metadata: HashMap<String, String>,
|
|
}
|
|
|
|
/// Builder for creating tasks
|
|
pub struct TaskBuilder {
|
|
task: Task,
|
|
}
|
|
|
|
impl DagScheduler {
|
|
/// Create a new DAG scheduler
|
|
pub async fn new(
|
|
config: DagConfig,
|
|
state_manager: Arc<StateManager>,
|
|
metrics: Arc<EtlMetrics>,
|
|
) -> Result<Self> {
|
|
Ok(Self {
|
|
config,
|
|
state_manager,
|
|
metrics,
|
|
active_executions: Arc::new(RwLock::new(HashMap::new())),
|
|
})
|
|
}
|
|
|
|
/// Plan the execution of a DAG
|
|
#[instrument(skip(self, graph))]
|
|
pub async fn plan_execution(&self, graph: &TaskGraph) -> Result<ExecutionPlan> {
|
|
debug!(
|
|
"Planning DAG execution for graph: {}",
|
|
graph.metadata.graph_id
|
|
);
|
|
|
|
// Validate the graph
|
|
graph.validate()?;
|
|
|
|
// Create execution stages using topological sort
|
|
let stages = self.create_execution_stages(graph)?;
|
|
|
|
// Estimate execution duration
|
|
let estimated_duration = self.estimate_execution_duration(&stages);
|
|
|
|
// Calculate resource requirements
|
|
let stage_resources = self.calculate_stage_resources(&stages);
|
|
|
|
Ok(ExecutionPlan {
|
|
stages,
|
|
estimated_duration,
|
|
stage_resources,
|
|
})
|
|
}
|
|
|
|
/// Create execution stages from task graph
|
|
fn create_execution_stages(&self, graph: &TaskGraph) -> Result<Vec<ExecutionStage>> {
|
|
// Build petgraph from task dependencies
|
|
let mut pg = Graph::<String, (), Directed>::new();
|
|
let mut node_indices = HashMap::new();
|
|
|
|
// Add nodes
|
|
for task_id in graph.tasks.keys() {
|
|
let idx = pg.add_node(task_id.clone());
|
|
node_indices.insert(task_id.clone(), idx);
|
|
}
|
|
|
|
// Add edges for dependencies
|
|
for (task_id, deps) in &graph.dependencies {
|
|
let task_idx = node_indices[task_id];
|
|
for dep in deps {
|
|
if let Some(&dep_idx) = node_indices.get(dep) {
|
|
pg.add_edge(dep_idx, task_idx, ());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check for cycles
|
|
if is_cyclic_directed(&pg) {
|
|
return Err(EtlError::DagValidation("Graph contains cycles".to_string()));
|
|
}
|
|
|
|
// Perform topological sort
|
|
let topo_sort = toposort(&pg, None).map_err(|_| {
|
|
EtlError::DagValidation("Failed to sort graph topologically".to_string())
|
|
})?;
|
|
|
|
// Group tasks into stages based on dependency levels
|
|
let mut stages = Vec::new();
|
|
let mut task_levels = HashMap::new();
|
|
|
|
// Calculate dependency levels
|
|
for node_idx in topo_sort {
|
|
let task_id = &pg[node_idx];
|
|
let task = &graph.tasks[task_id];
|
|
|
|
let level = if task.dependencies.is_empty() {
|
|
0
|
|
} else {
|
|
task.dependencies
|
|
.iter()
|
|
.map(|dep| task_levels.get(dep).unwrap_or(&0) + 1)
|
|
.max()
|
|
.unwrap_or(0)
|
|
};
|
|
|
|
task_levels.insert(task_id.clone(), level);
|
|
}
|
|
|
|
// Group tasks by level into stages
|
|
let max_level = task_levels.values().max().unwrap_or(&0);
|
|
for stage_id in 0..=*max_level {
|
|
let stage_tasks: Vec<Task> = task_levels
|
|
.iter()
|
|
.filter(|(_, level)| **level == stage_id)
|
|
.map(|(task_id, _)| graph.tasks[task_id].clone())
|
|
.collect();
|
|
|
|
if !stage_tasks.is_empty() {
|
|
stages.push(ExecutionStage {
|
|
stage_id,
|
|
tasks: stage_tasks,
|
|
dependencies: if stage_id > 0 {
|
|
vec![stage_id - 1]
|
|
} else {
|
|
vec![]
|
|
},
|
|
});
|
|
}
|
|
}
|
|
|
|
Ok(stages)
|
|
}
|
|
|
|
/// Estimate total execution duration
|
|
fn estimate_execution_duration(&self, stages: &[ExecutionStage]) -> Duration {
|
|
let mut total_duration = Duration::ZERO;
|
|
|
|
for stage in stages {
|
|
// For parallel execution, use the longest task in the stage
|
|
let stage_duration = stage
|
|
.tasks
|
|
.iter()
|
|
.map(|task| {
|
|
Duration::from_millis(task.timeout_ms.unwrap_or(self.config.task_timeout_ms))
|
|
})
|
|
.max()
|
|
.unwrap_or(Duration::ZERO);
|
|
|
|
total_duration += stage_duration;
|
|
}
|
|
|
|
total_duration
|
|
}
|
|
|
|
/// Calculate resource requirements for each stage
|
|
fn calculate_stage_resources(&self, stages: &[ExecutionStage]) -> Vec<ResourceRequirements> {
|
|
stages
|
|
.iter()
|
|
.map(|stage| {
|
|
// Sum resources for all tasks in the stage (parallel execution)
|
|
let total_cpu: f32 = stage.tasks.iter().map(|t| t.resources.cpu_cores).sum();
|
|
let total_memory: u64 = stage.tasks.iter().map(|t| t.resources.memory_mb).sum();
|
|
let total_disk: u64 = stage.tasks.iter().map(|t| t.resources.disk_mb).sum();
|
|
|
|
ResourceRequirements {
|
|
cpu_cores: total_cpu,
|
|
memory_mb: total_memory,
|
|
disk_mb: total_disk,
|
|
gpu_memory_mb: None,
|
|
network_mbps: None,
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
impl TaskGraph {
|
|
/// Create a new empty task graph
|
|
#[must_use]
|
|
pub fn new() -> Self {
|
|
Self {
|
|
tasks: HashMap::new(),
|
|
dependencies: HashMap::new(),
|
|
metadata: GraphMetadata {
|
|
graph_id: Uuid::new_v4().to_string(),
|
|
name: "unnamed_graph".to_string(),
|
|
description: None,
|
|
tags: Vec::new(),
|
|
created_at: Utc::now(),
|
|
modified_at: Utc::now(),
|
|
version: "1.0.0".to_string(),
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Add a task to the graph
|
|
pub fn add_task(&mut self, task: Task) {
|
|
let deps = task.dependencies.clone();
|
|
self.dependencies.insert(task.id.clone(), deps);
|
|
self.tasks.insert(task.id.clone(), task);
|
|
self.metadata.modified_at = Utc::now();
|
|
}
|
|
|
|
/// Remove a task from the graph
|
|
pub fn remove_task(&mut self, task_id: &str) -> Option<Task> {
|
|
self.dependencies.remove(task_id);
|
|
|
|
// Remove this task from other tasks' dependencies
|
|
for deps in self.dependencies.values_mut() {
|
|
deps.retain(|dep| dep != task_id);
|
|
}
|
|
|
|
self.metadata.modified_at = Utc::now();
|
|
self.tasks.remove(task_id)
|
|
}
|
|
|
|
/// Add a dependency between tasks
|
|
pub fn add_dependency(&mut self, task_id: &str, dependency: &str) -> Result<()> {
|
|
if !self.tasks.contains_key(task_id) {
|
|
return Err(EtlError::DagValidation(format!("Task {task_id} not found")));
|
|
}
|
|
|
|
if !self.tasks.contains_key(dependency) {
|
|
return Err(EtlError::DagValidation(format!(
|
|
"Dependency task {dependency} not found"
|
|
)));
|
|
}
|
|
|
|
self.dependencies
|
|
.entry(task_id.to_string())
|
|
.or_default()
|
|
.push(dependency.to_string());
|
|
|
|
// Update the task's dependencies as well
|
|
if let Some(task) = self.tasks.get_mut(task_id)
|
|
&& !task.dependencies.contains(&dependency.to_string())
|
|
{
|
|
task.dependencies.push(dependency.to_string());
|
|
}
|
|
|
|
self.metadata.modified_at = Utc::now();
|
|
|
|
// Check for cycles after adding dependency
|
|
self.validate_no_cycles()
|
|
}
|
|
|
|
/// Validate the graph for correctness
|
|
pub fn validate(&self) -> Result<()> {
|
|
// Check for missing dependencies
|
|
for (task_id, deps) in &self.dependencies {
|
|
for dep in deps {
|
|
if !self.tasks.contains_key(dep) {
|
|
return Err(EtlError::DagValidation(format!(
|
|
"Task {task_id} depends on non-existent task {dep}"
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check for cycles
|
|
self.validate_no_cycles()?;
|
|
|
|
// Check for orphaned tasks (tasks with no dependencies and no dependents)
|
|
let mut has_dependents = HashSet::new();
|
|
for deps in self.dependencies.values() {
|
|
for dep in deps {
|
|
has_dependents.insert(dep.clone());
|
|
}
|
|
}
|
|
|
|
for task_id in self.tasks.keys() {
|
|
if self
|
|
.dependencies
|
|
.get(task_id)
|
|
.is_none_or(std::vec::Vec::is_empty)
|
|
&& !has_dependents.contains(task_id)
|
|
{
|
|
warn!(
|
|
"Task {} is orphaned (no dependencies or dependents)",
|
|
task_id
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Validate that the graph contains no cycles
|
|
fn validate_no_cycles(&self) -> Result<()> {
|
|
let mut visited = HashSet::new();
|
|
let mut rec_stack = HashSet::new();
|
|
|
|
for task_id in self.tasks.keys() {
|
|
if !visited.contains(task_id)
|
|
&& self.has_cycle_dfs(task_id, &mut visited, &mut rec_stack)
|
|
{
|
|
return Err(EtlError::DagValidation("Graph contains cycles".to_string()));
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// DFS cycle detection helper
|
|
fn has_cycle_dfs(
|
|
&self,
|
|
task_id: &str,
|
|
visited: &mut HashSet<String>,
|
|
rec_stack: &mut HashSet<String>,
|
|
) -> bool {
|
|
visited.insert(task_id.to_string());
|
|
rec_stack.insert(task_id.to_string());
|
|
|
|
if let Some(deps) = self.dependencies.get(task_id) {
|
|
for dep in deps {
|
|
if !visited.contains(dep) {
|
|
if self.has_cycle_dfs(dep, visited, rec_stack) {
|
|
return true;
|
|
}
|
|
} else if rec_stack.contains(dep) {
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
|
|
rec_stack.remove(task_id);
|
|
false
|
|
}
|
|
}
|
|
|
|
impl TaskBuilder {
|
|
/// Create a new task builder
|
|
#[must_use]
|
|
pub fn new(id: &str) -> Self {
|
|
Self {
|
|
task: Task::new(id.to_string(), TaskType::Transform),
|
|
}
|
|
}
|
|
|
|
/// Set task name
|
|
#[must_use]
|
|
pub fn name(mut self, name: &str) -> Self {
|
|
self.task.name = name.to_string();
|
|
self
|
|
}
|
|
|
|
/// Set task description
|
|
#[must_use]
|
|
pub fn description(mut self, description: &str) -> Self {
|
|
self.task.description = Some(description.to_string());
|
|
self
|
|
}
|
|
|
|
/// Set task type
|
|
#[must_use]
|
|
pub fn task_type(mut self, task_type: TaskType) -> Self {
|
|
self.task.task_type = task_type;
|
|
self
|
|
}
|
|
|
|
/// Add a dependency
|
|
#[must_use]
|
|
pub fn depends_on(mut self, dependency: &str) -> Self {
|
|
self.task.dependencies.push(dependency.to_string());
|
|
self
|
|
}
|
|
|
|
/// Set data source
|
|
#[must_use]
|
|
pub fn with_source(mut self, source: DataSource) -> Self {
|
|
self.task.source = Some(source);
|
|
self
|
|
}
|
|
|
|
/// Set transformation
|
|
#[must_use]
|
|
pub fn with_transformation(mut self, transformation: Transformation) -> Self {
|
|
self.task.transformation = Some(transformation);
|
|
self
|
|
}
|
|
|
|
/// Set data sink
|
|
#[must_use]
|
|
pub fn with_sink(mut self, sink: DataSink) -> Self {
|
|
self.task.sink = Some(sink);
|
|
self
|
|
}
|
|
|
|
/// Set task priority
|
|
#[must_use]
|
|
pub fn priority(mut self, priority: i32) -> Self {
|
|
self.task.priority = priority;
|
|
self
|
|
}
|
|
|
|
/// Set retry configuration
|
|
#[must_use]
|
|
pub fn retry_config(mut self, config: RetryConfig) -> Self {
|
|
self.task.retry_config = config;
|
|
self
|
|
}
|
|
|
|
/// Set resource requirements
|
|
#[must_use]
|
|
pub fn resources(mut self, resources: ResourceRequirements) -> Self {
|
|
self.task.resources = resources;
|
|
self
|
|
}
|
|
|
|
/// Build the task
|
|
#[must_use]
|
|
pub fn build(self) -> Task {
|
|
self.task
|
|
}
|
|
}
|
|
|
|
impl Default for RetryConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_retries: 3,
|
|
strategy: RetryStrategy::ExponentialBackoff,
|
|
base_delay_ms: 1000,
|
|
max_delay_ms: 60000,
|
|
backoff_multiplier: 2.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for ResourceRequirements {
|
|
fn default() -> Self {
|
|
Self {
|
|
cpu_cores: 1.0,
|
|
memory_mb: 512,
|
|
disk_mb: 1024,
|
|
gpu_memory_mb: None,
|
|
network_mbps: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for TaskGraph {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|