rskit_dag/
dag.rs

1use std::collections::{HashMap, HashSet, VecDeque};
2use std::sync::Arc;
3
4use crate::node::DagNode;
5use parking_lot::Mutex;
6use rskit_errors::{AppError, AppResult, ErrorCode};
7use tokio::task::JoinSet;
8use tokio_util::sync::CancellationToken;
9
10/// Failure handling strategy for DAG execution.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12#[non_exhaustive]
13pub enum FailurePolicy {
14    /// Stop execution immediately when the first node fails.
15    FailFast,
16    /// Continue executing all nodes, even when upstream dependencies fail.
17    Continue,
18    /// Skip all downstream dependents of a failed node, but keep independent branches running.
19    SkipDependents,
20}
21
22struct NodeExecution {
23    node_id: String,
24    result: AppResult<serde_json::Value>,
25}
26
27/// A directed acyclic graph of executable nodes.
28///
29/// Nodes are executed in topological order with maximum parallelism:
30/// all nodes whose dependencies are satisfied run concurrently.
31/// Use [`with_max_parallelism`](Dag::with_max_parallelism) to limit
32/// how many nodes execute simultaneously.
33pub struct Dag {
34    nodes: HashMap<String, Arc<dyn DagNode>>,
35    /// Forward edges: `from_id` → downstream dependents
36    edges: HashMap<String, Vec<String>>,
37    /// Reverse edges: `to_id` → upstream dependencies
38    reverse_edges: HashMap<String, Vec<String>>,
39    /// Limit on concurrent node execution.
40    max_parallelism: usize,
41    /// Failure handling strategy.
42    failure_policy: FailurePolicy,
43}
44
45impl Default for Dag {
46    fn default() -> Self {
47        Self::new()
48    }
49}
50
51impl Dag {
52    /// Create an empty DAG.
53    #[must_use]
54    pub fn new() -> Self {
55        Self {
56            nodes: HashMap::new(),
57            edges: HashMap::new(),
58            reverse_edges: HashMap::new(),
59            max_parallelism: default_parallelism(),
60            failure_policy: FailurePolicy::FailFast,
61        }
62    }
63
64    /// Set the maximum number of nodes that can execute concurrently.
65    #[must_use]
66    pub fn with_max_parallelism(mut self, max: usize) -> Self {
67        self.max_parallelism = max.max(1);
68        self
69    }
70
71    /// Set the failure handling policy.
72    #[must_use]
73    pub const fn with_failure_policy(mut self, policy: FailurePolicy) -> Self {
74        self.failure_policy = policy;
75        self
76    }
77
78    /// Add a node to the DAG.
79    #[must_use]
80    pub fn add_node(mut self, node: impl DagNode) -> Self {
81        let id = node.id().to_owned();
82        self.nodes.insert(id.clone(), Arc::new(node));
83        self.edges.entry(id.clone()).or_default();
84        self.reverse_edges.entry(id).or_default();
85        self
86    }
87
88    /// Add a directed edge from one node to another.
89    pub fn add_edge(mut self, from: &str, to: &str) -> AppResult<Self> {
90        if !self.nodes.contains_key(from) {
91            return Err(AppError::new(
92                ErrorCode::InvalidInput,
93                format!("DAG node '{from}' not found"),
94            ));
95        }
96        if !self.nodes.contains_key(to) {
97            return Err(AppError::new(
98                ErrorCode::InvalidInput,
99                format!("DAG node '{to}' not found"),
100            ));
101        }
102
103        if self.path_exists(to, from) {
104            return Err(AppError::new(
105                ErrorCode::InvalidInput,
106                format!("DAG edge '{from}' -> '{to}' would create a cycle"),
107            ));
108        }
109
110        self.edges
111            .entry(from.to_owned())
112            .or_default()
113            .push(to.to_owned());
114        self.reverse_edges
115            .entry(to.to_owned())
116            .or_default()
117            .push(from.to_owned());
118
119        Ok(self)
120    }
121
122    /// Topological sort using Kahn's algorithm.
123    pub fn topological_sort(&self) -> AppResult<Vec<String>> {
124        let mut in_degree: HashMap<String, usize> = HashMap::new();
125        for id in self.nodes.keys() {
126            in_degree.entry(id.clone()).or_insert(0);
127        }
128        for (node_id, deps) in &self.reverse_edges {
129            *in_degree.entry(node_id.clone()).or_insert(0) = deps.len();
130        }
131
132        let mut queue: VecDeque<String> = VecDeque::new();
133        for (id, &deg) in &in_degree {
134            if deg == 0 {
135                queue.push_back(id.clone());
136            }
137        }
138
139        let mut sorted = Vec::with_capacity(self.nodes.len());
140        while let Some(id) = queue.pop_front() {
141            sorted.push(id.clone());
142            if let Some(dependents) = self.edges.get(&id) {
143                for dep in dependents {
144                    if let Some(deg) = in_degree.get_mut(dep) {
145                        *deg -= 1;
146                        if *deg == 0 {
147                            queue.push_back(dep.clone());
148                        }
149                    }
150                }
151            }
152        }
153
154        if sorted.len() != self.nodes.len() {
155            return Err(AppError::new(
156                ErrorCode::InvalidInput,
157                "DAG contains a cycle",
158            ));
159        }
160
161        Ok(sorted)
162    }
163
164    /// Execute all nodes respecting dependency order with maximum parallelism.
165    pub async fn execute(
166        &self,
167        cancel: CancellationToken,
168    ) -> AppResult<HashMap<String, serde_json::Value>> {
169        self.execute_with_inputs(HashMap::new(), cancel).await
170    }
171
172    /// Execute with initial inputs provided to root nodes.
173    pub async fn execute_with_inputs(
174        &self,
175        initial_inputs: HashMap<String, serde_json::Value>,
176        cancel: CancellationToken,
177    ) -> AppResult<HashMap<String, serde_json::Value>> {
178        let _ = self.topological_sort()?;
179
180        let outputs = Arc::new(Mutex::new(HashMap::<String, serde_json::Value>::new()));
181        let semaphore = Arc::new(tokio::sync::Semaphore::new(self.max_parallelism));
182
183        let mut remaining_in_degree: HashMap<String, usize> = self
184            .nodes
185            .keys()
186            .map(|id| {
187                (
188                    id.clone(),
189                    self.reverse_edges.get(id).map_or(0, std::vec::Vec::len),
190                )
191            })
192            .collect();
193
194        let mut join_set: JoinSet<NodeExecution> = JoinSet::new();
195        let mut pending = HashSet::new();
196        let mut completed = HashSet::new();
197        let mut failed = HashSet::new();
198        let mut skipped = HashSet::new();
199
200        for (id, degree) in &remaining_in_degree {
201            if *degree == 0 {
202                self.spawn_node(
203                    id.clone(),
204                    initial_inputs.clone(),
205                    cancel.clone(),
206                    Arc::clone(&semaphore),
207                    &mut join_set,
208                    &mut pending,
209                )?;
210            }
211        }
212
213        while let Some(joined) = join_set.join_next().await {
214            let execution = joined.map_err(|error| {
215                AppError::new(ErrorCode::Internal, format!("DAG task panicked: {error}"))
216            })?;
217            pending.remove(&execution.node_id);
218
219            match execution.result {
220                Ok(value) => {
221                    outputs.lock().insert(execution.node_id.clone(), value);
222                    completed.insert(execution.node_id.clone());
223                    self.schedule_dependents(
224                        &execution.node_id,
225                        &cancel,
226                        &initial_inputs,
227                        &outputs,
228                        &semaphore,
229                        &mut join_set,
230                        &mut remaining_in_degree,
231                        &mut pending,
232                        &completed,
233                        &failed,
234                        &skipped,
235                    )?;
236                }
237                Err(error) => match self.failure_policy {
238                    FailurePolicy::FailFast => return Err(error),
239                    FailurePolicy::Continue => {
240                        failed.insert(execution.node_id.clone());
241                        completed.insert(execution.node_id.clone());
242                        self.schedule_dependents(
243                            &execution.node_id,
244                            &cancel,
245                            &initial_inputs,
246                            &outputs,
247                            &semaphore,
248                            &mut join_set,
249                            &mut remaining_in_degree,
250                            &mut pending,
251                            &completed,
252                            &failed,
253                            &skipped,
254                        )?;
255                    }
256                    FailurePolicy::SkipDependents => {
257                        failed.insert(execution.node_id.clone());
258                        completed.insert(execution.node_id.clone());
259                        mark_skipped_dependents(
260                            &self.edges,
261                            &execution.node_id,
262                            &mut skipped,
263                            &pending,
264                        );
265                    }
266                },
267            }
268        }
269
270        Ok(outputs.lock().clone())
271    }
272
273    #[allow(clippy::too_many_arguments)]
274    fn schedule_dependents(
275        &self,
276        finished_id: &str,
277        cancel: &CancellationToken,
278        initial_inputs: &HashMap<String, serde_json::Value>,
279        outputs: &Arc<Mutex<HashMap<String, serde_json::Value>>>,
280        semaphore: &Arc<tokio::sync::Semaphore>,
281        join_set: &mut JoinSet<NodeExecution>,
282        remaining_in_degree: &mut HashMap<String, usize>,
283        pending: &mut HashSet<String>,
284        completed: &HashSet<String>,
285        failed: &HashSet<String>,
286        skipped: &HashSet<String>,
287    ) -> AppResult<()> {
288        if let Some(dependents) = self.edges.get(finished_id) {
289            for dependent_id in dependents {
290                if completed.contains(dependent_id)
291                    || pending.contains(dependent_id)
292                    || skipped.contains(dependent_id)
293                {
294                    continue;
295                }
296                if let Some(degree) = remaining_in_degree.get_mut(dependent_id) {
297                    *degree = degree.saturating_sub(1);
298                    if *degree == 0 {
299                        let inputs =
300                            self.collect_inputs(dependent_id, outputs, failed, initial_inputs);
301                        self.spawn_node(
302                            dependent_id.clone(),
303                            inputs,
304                            cancel.clone(),
305                            Arc::clone(semaphore),
306                            join_set,
307                            pending,
308                        )?;
309                    }
310                }
311            }
312        }
313        Ok(())
314    }
315
316    fn collect_inputs(
317        &self,
318        node_id: &str,
319        outputs: &Arc<Mutex<HashMap<String, serde_json::Value>>>,
320        failed: &HashSet<String>,
321        initial_inputs: &HashMap<String, serde_json::Value>,
322    ) -> HashMap<String, serde_json::Value> {
323        let reverse = self.reverse_edges.get(node_id).cloned().unwrap_or_default();
324        if reverse.is_empty() {
325            return initial_inputs.clone();
326        }
327
328        let output_guard = outputs.lock();
329        reverse
330            .iter()
331            .filter(|dependency| !failed.contains(*dependency))
332            .filter_map(|dependency| {
333                output_guard
334                    .get(dependency)
335                    .map(|value| (dependency.clone(), value.clone()))
336            })
337            .collect()
338    }
339
340    fn spawn_node(
341        &self,
342        node_id: String,
343        inputs: HashMap<String, serde_json::Value>,
344        cancel: CancellationToken,
345        semaphore: Arc<tokio::sync::Semaphore>,
346        join_set: &mut JoinSet<NodeExecution>,
347        pending: &mut HashSet<String>,
348    ) -> AppResult<()> {
349        let node = Arc::clone(self.nodes.get(&node_id).ok_or_else(|| {
350            AppError::new(
351                ErrorCode::Internal,
352                format!("DAG node '{node_id}' not found in node map"),
353            )
354        })?);
355        pending.insert(node_id.clone());
356        join_set.spawn(async move {
357            let permit_result = semaphore
358                .acquire()
359                .await
360                .map_err(|_| AppError::new(ErrorCode::Internal, "DAG semaphore closed"));
361
362            match permit_result {
363                Ok(_permit) => {
364                    tracing::debug!(node = %node_id, "executing DAG node");
365                    NodeExecution {
366                        node_id,
367                        result: node.execute(inputs, cancel).await,
368                    }
369                }
370                Err(error) => NodeExecution {
371                    node_id,
372                    result: Err(error),
373                },
374            }
375        });
376        Ok(())
377    }
378
379    fn path_exists(&self, from: &str, to: &str) -> bool {
380        let mut stack = vec![from.to_string()];
381        let mut visited = HashSet::new();
382        while let Some(node_id) = stack.pop() {
383            if node_id == to {
384                return true;
385            }
386            if !visited.insert(node_id.clone()) {
387                continue;
388            }
389            if let Some(children) = self.edges.get(&node_id) {
390                stack.extend(children.iter().cloned());
391            }
392        }
393        false
394    }
395}
396
397fn default_parallelism() -> usize {
398    std::thread::available_parallelism().map_or(1, std::num::NonZeroUsize::get)
399}
400
401fn mark_skipped_dependents(
402    edges: &HashMap<String, Vec<String>>,
403    failed_id: &str,
404    skipped: &mut HashSet<String>,
405    pending: &HashSet<String>,
406) {
407    let mut stack = edges.get(failed_id).cloned().unwrap_or_default();
408    while let Some(node_id) = stack.pop() {
409        if pending.contains(&node_id) || !skipped.insert(node_id.clone()) {
410            continue;
411        }
412        if let Some(children) = edges.get(&node_id) {
413            stack.extend(children.iter().cloned());
414        }
415    }
416}
417
418#[cfg(test)]
419mod tests {
420    use super::*;
421    use std::future::Future;
422    use std::pin::Pin;
423    use std::sync::atomic::{AtomicUsize, Ordering};
424
425    struct AddNode {
426        id: String,
427        value: i64,
428    }
429
430    impl DagNode for AddNode {
431        fn id(&self) -> &str {
432            &self.id
433        }
434
435        fn execute(
436            &self,
437            inputs: HashMap<String, serde_json::Value>,
438            _cancel: CancellationToken,
439        ) -> Pin<Box<dyn Future<Output = AppResult<serde_json::Value>> + Send + '_>> {
440            Box::pin(async move {
441                let sum: i64 = inputs
442                    .values()
443                    .filter_map(serde_json::Value::as_i64)
444                    .sum::<i64>()
445                    + self.value;
446                Ok(serde_json::json!(sum))
447            })
448        }
449    }
450
451    struct FailNode {
452        id: String,
453    }
454
455    impl DagNode for FailNode {
456        fn id(&self) -> &str {
457            &self.id
458        }
459
460        fn execute(
461            &self,
462            _inputs: HashMap<String, serde_json::Value>,
463            _cancel: CancellationToken,
464        ) -> Pin<Box<dyn Future<Output = AppResult<serde_json::Value>> + Send + '_>> {
465            Box::pin(async { Err(AppError::new(ErrorCode::Internal, "node failed")) })
466        }
467    }
468
469    struct CountingNode {
470        id: String,
471        counter: Arc<AtomicUsize>,
472        value: i64,
473    }
474
475    impl DagNode for CountingNode {
476        fn id(&self) -> &str {
477            &self.id
478        }
479
480        fn execute(
481            &self,
482            inputs: HashMap<String, serde_json::Value>,
483            _cancel: CancellationToken,
484        ) -> Pin<Box<dyn Future<Output = AppResult<serde_json::Value>> + Send + '_>> {
485            Box::pin(async move {
486                self.counter.fetch_add(1, Ordering::SeqCst);
487                let sum = inputs
488                    .values()
489                    .filter_map(serde_json::Value::as_i64)
490                    .sum::<i64>()
491                    + self.value;
492                Ok(serde_json::json!(sum))
493            })
494        }
495    }
496
497    #[tokio::test]
498    async fn linear_dag_executes_in_order() {
499        let dag = Dag::new()
500            .add_node(AddNode {
501                id: "a".into(),
502                value: 1,
503            })
504            .add_node(AddNode {
505                id: "b".into(),
506                value: 2,
507            })
508            .add_node(AddNode {
509                id: "c".into(),
510                value: 3,
511            });
512
513        let dag = dag.add_edge("a", "b").unwrap().add_edge("b", "c").unwrap();
514        let outputs = dag.execute(CancellationToken::new()).await.unwrap();
515
516        assert_eq!(outputs["a"], serde_json::json!(1));
517        assert_eq!(outputs["b"], serde_json::json!(3));
518        assert_eq!(outputs["c"], serde_json::json!(6));
519    }
520
521    #[tokio::test]
522    async fn diamond_dag_merges_inputs() {
523        let dag = Dag::new()
524            .add_node(AddNode {
525                id: "a".into(),
526                value: 10,
527            })
528            .add_node(AddNode {
529                id: "b".into(),
530                value: 1,
531            })
532            .add_node(AddNode {
533                id: "c".into(),
534                value: 2,
535            })
536            .add_node(AddNode {
537                id: "d".into(),
538                value: 0,
539            });
540
541        let dag = dag
542            .add_edge("a", "b")
543            .unwrap()
544            .add_edge("a", "c")
545            .unwrap()
546            .add_edge("b", "d")
547            .unwrap()
548            .add_edge("c", "d")
549            .unwrap();
550
551        let outputs = dag.execute(CancellationToken::new()).await.unwrap();
552
553        assert_eq!(outputs["a"], serde_json::json!(10));
554        assert_eq!(outputs["b"], serde_json::json!(11));
555        assert_eq!(outputs["c"], serde_json::json!(12));
556        assert_eq!(outputs["d"], serde_json::json!(23));
557    }
558
559    #[tokio::test]
560    async fn fail_fast_returns_first_error() {
561        let dag = Dag::new()
562            .add_node(FailNode { id: "a".into() })
563            .add_node(AddNode {
564                id: "b".into(),
565                value: 1,
566            })
567            .add_edge("a", "b")
568            .unwrap();
569
570        let result = dag.execute(CancellationToken::new()).await;
571        assert!(result.is_err());
572    }
573
574    #[tokio::test]
575    async fn continue_runs_dependents_with_partial_inputs() {
576        let dependent_runs = Arc::new(AtomicUsize::new(0));
577        let independent_runs = Arc::new(AtomicUsize::new(0));
578        let dag = Dag::new()
579            .with_failure_policy(FailurePolicy::Continue)
580            .add_node(FailNode { id: "a".into() })
581            .add_node(CountingNode {
582                id: "b".into(),
583                counter: dependent_runs.clone(),
584                value: 5,
585            })
586            .add_node(CountingNode {
587                id: "c".into(),
588                counter: independent_runs.clone(),
589                value: 9,
590            });
591
592        let dag = dag.add_edge("a", "b").unwrap();
593        let outputs = dag.execute(CancellationToken::new()).await.unwrap();
594
595        assert_eq!(dependent_runs.load(Ordering::SeqCst), 1);
596        assert_eq!(independent_runs.load(Ordering::SeqCst), 1);
597        assert_eq!(outputs["b"], serde_json::json!(5));
598        assert_eq!(outputs["c"], serde_json::json!(9));
599        assert!(!outputs.contains_key("a"));
600    }
601
602    #[tokio::test]
603    async fn skip_dependents_skips_failed_branch_only() {
604        let dependent_runs = Arc::new(AtomicUsize::new(0));
605        let independent_runs = Arc::new(AtomicUsize::new(0));
606        let dag = Dag::new()
607            .with_failure_policy(FailurePolicy::SkipDependents)
608            .add_node(FailNode { id: "a".into() })
609            .add_node(CountingNode {
610                id: "b".into(),
611                counter: dependent_runs.clone(),
612                value: 5,
613            })
614            .add_node(CountingNode {
615                id: "c".into(),
616                counter: independent_runs.clone(),
617                value: 9,
618            });
619
620        let dag = dag.add_edge("a", "b").unwrap();
621        let outputs = dag.execute(CancellationToken::new()).await.unwrap();
622
623        assert_eq!(dependent_runs.load(Ordering::SeqCst), 0);
624        assert_eq!(independent_runs.load(Ordering::SeqCst), 1);
625        assert!(!outputs.contains_key("b"));
626        assert_eq!(outputs["c"], serde_json::json!(9));
627    }
628
629    #[test]
630    fn cycle_detection() {
631        let dag = Dag::new()
632            .add_node(AddNode {
633                id: "a".into(),
634                value: 0,
635            })
636            .add_node(AddNode {
637                id: "b".into(),
638                value: 0,
639            });
640
641        let result = dag.add_edge("a", "b").unwrap().add_edge("b", "a");
642        assert!(result.is_err());
643    }
644}
rskit_dag/dag.rs

rskit_dag/
dag.rs
