dag-executor 0.1.0

A production-ready DAG executor with state management and advanced patterns
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

//! Integration tests for complete workflows and task types.

use std::sync::Arc;
use std::time::Duration;

use crate::fixtures::ok_task;
use dag_executor::context::Context;
use dag_executor::prelude::*;

fn executor() -> DagExecutor {
    DagExecutor::builder()
        .persist(false)
        .concurrency(16)
        .build()
}

#[tokio::test]
async fn runs_a_linear_workflow() {
    let mut dag = Dag::new();
    dag.add_task(ok_task("a", &[])).unwrap();
    dag.add_task(ok_task("b", &["a"])).unwrap();
    dag.add_task(ok_task("c", &["b"])).unwrap();

    let report = executor().run(dag).await.unwrap();
    assert!(report.is_success());
    assert_eq!(report.count_in(TaskState::Completed), 3);
}

#[tokio::test]
async fn runs_diamond_with_fan_out_in() {
    let mut dag = Dag::new();
    dag.add_task(ok_task("start", &[])).unwrap();
    for task in patterns::fan_out_in(
        "stage",
        16,
        Some("start"),
        |_ctx, i| async move { Ok(serde_json::json!({ "value": i })) },
        |_ctx, results| async move {
            let sum: u64 = results
                .iter()
                .filter_map(|r| r.get("value").and_then(|v| v.as_u64()))
                .sum();
            Ok(serde_json::json!(sum))
        },
    ) {
        dag.add_task(task).unwrap();
    }

    let report = executor().run(dag).await.unwrap();
    assert!(report.is_success());
    // Sum of 0..16 == 120.
    assert_eq!(
        report.records["stage.aggregate"].output,
        Some(serde_json::json!(120))
    );
}

#[tokio::test]
async fn conditional_branch_drives_downstream() {
    let mut dag = Dag::new();
    dag.add_task(Arc::new(ConditionalTask::new(
        "cond",
        "yes",
        "no",
        |_ctx| async { Ok(true) },
    )))
    .unwrap();
    dag.add_task(Arc::new(
        BasicTask::new("after", |ctx: Arc<Context>| async move {
            Ok(ctx.get("cond.branch").unwrap_or(serde_json::Value::Null))
        })
        .with_deps(["cond"]),
    ))
    .unwrap();

    let report = executor().run(dag).await.unwrap();
    assert!(report.is_success());
    assert_eq!(
        report.records["after"].output,
        Some(serde_json::json!("yes"))
    );
}

#[tokio::test]
async fn loop_task_breaks_early() {
    let mut dag = Dag::new();
    dag.add_task(Arc::new(
        LoopTask::new(
            "loop",
            10,
            |_ctx, i| async move { Ok(serde_json::json!(i)) },
        )
        .with_break(|out| out.as_u64() == Some(2)),
    ))
    .unwrap();

    let report = executor().run(dag).await.unwrap();
    let out = report.records["loop"].output.clone().unwrap();
    assert_eq!(out["iterations"], serde_json::json!(3)); // i = 0,1,2
    assert_eq!(out["broke_early"], serde_json::json!(true));
}

#[tokio::test]
async fn event_driven_task_waits_for_signal() {
    let mut dag = Dag::new();
    dag.add_task(Arc::new(
        EventDrivenTask::new("waiter", "go", |_ctx| async {
            Ok(serde_json::json!("woke up"))
        })
        .with_timeout(Duration::from_secs(5)),
    ))
    .unwrap();
    dag.add_task(Arc::new(BasicTask::new(
        "emitter",
        |ctx: Arc<Context>| async move {
            // Give the waiter time to register before emitting.
            tokio::time::sleep(Duration::from_millis(50)).await;
            ctx.emit("go");
            Ok(serde_json::Value::Null)
        },
    )))
    .unwrap();

    let report = executor().run(dag).await.unwrap();
    assert!(report.is_success());
    assert_eq!(
        report.records["waiter"].output,
        Some(serde_json::json!("woke up"))
    );
}