tinyflow-framework 0.1.0

Streaming runtime engine — checkpoint, chained task execution, and state store
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

//! Spawn [`ChainedStreamTask`](crate::runtime::chained_stream_task::ChainedStreamTask) from a
//! [`ChainJob`](crate::runtime::chain_job::ChainJob) (framework-only, no committer).

mod spawn;
mod validate;

use std::sync::Arc;

use crate::env::stream_env::StreamEnv;
use crate::runtime::checkpoint_manager::CheckpointManager;
use crate::runtime::job_runtime::{ChainJobRuntime, shutdown_chain_job};
use crate::runtime::task_lifecycle::report_task_failure;
use tinyflow_api::error::StreamResult;

pub use spawn::spawn_chained_tasks;

use spawn::spawn_chained_tasks as spawn_tasks;
use validate::resolve_job_state_for_execute;

/// Run the built operator chain until Ctrl+C, task failure, or checkpoint manager exit.
pub async fn execute(env: StreamEnv) -> StreamResult<()> {
    let (job_state, work_dir) = resolve_job_state_for_execute(&env).await?;
    let work_dir_ref = work_dir.as_deref();

    {
        let pending = std::mem::take(&mut *env.pending_job_state.lock().unwrap());
        for (namespace, key, value) in pending {
            job_state.put(&namespace, key.as_bytes(), &value).await?;
        }
    }

    let (ack_tx, ack_rx) = tokio::sync::mpsc::channel(64);
    let (error_tx, mut error_rx) = tokio::sync::mpsc::channel(64);
    let job_runtime = ChainJobRuntime::new();
    let mut task_id_counter = 0u32;

    let pending = spawn_tasks(
        &env,
        job_state,
        work_dir_ref,
        &job_runtime,
        &ack_tx,
        &error_tx,
        0,
        &mut task_id_counter,
    )
    .await?;

    for p in pending {
        job_runtime.register_join(p.spawn());
    }

    let cm_handle = if env.checkpoint_interval_ms > 0
        && !job_runtime.control_senders.lock().unwrap().is_empty()
    {
        let interval = env.checkpoint_interval_ms;
        log::info!(
            "[execute] checkpoint manager interval={}ms, {} tasks",
            interval,
            job_runtime.control_senders.lock().unwrap().len()
        );
        let mut cm =
            CheckpointManager::new(interval, ack_rx, Arc::clone(&job_runtime.control_senders));
        let cm_error_tx = error_tx.clone();
        Some(tokio::spawn(async move {
            if let Err(e) = cm.run().await {
                report_task_failure(&cm_error_tx, 0, "checkpoint-manager", e).await;
            }
        }))
    } else {
        drop(ack_rx);
        None
    };

    drop(ack_tx);

    let result = if let Some(handle) = cm_handle {
        tokio::select! {
            _ = tokio::signal::ctrl_c() => Ok(()),
            msg = error_rx.recv() => match msg {
                Some(e) => Err(e),
                None => Ok(()),
            },
            _ = handle => match error_rx.try_recv() {
                Ok(e) => Err(e),
                Err(_) => Ok(()),
            },
        }
    } else {
        tokio::select! {
            _ = tokio::signal::ctrl_c() => Ok(()),
            msg = error_rx.recv() => match msg {
                Some(e) => Err(e),
                None => Ok(()),
            },
        }
    };

    shutdown_chain_job(&job_runtime).await?;
    result
}