swale 0.2.0

A scheduled asset-graph orchestrator on object storage.
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
131
132
133
134
135

//! The pools: one [`WorkflowRuntime`] per pool, all over one queue and one
//! store.
//!
//! Every task instance runs on the runtime of its node's pool. The queue of
//! the pool `name` is `swale-pool-{name}`, and its memos are at `memos/{name}`
//! within the store prefix. The terminal hook of each runtime is the
//! [`RecordHook`], which writes the node's record and enqueues the event for
//! the scheduler.

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

use taquba::Queue;
use taquba::object_store::ObjectStore;
use taquba_workflow::{RunnerHandle, WorkflowRuntime};
use tokio_util::sync::CancellationToken;

use crate::dispatch::Dispatch;
use crate::hook::RecordHook;
use crate::operator::OperatorSet;
use crate::store::store_path;

/// The runtime of a pool.
pub type PoolRuntime = WorkflowRuntime<Dispatch, RecordHook>;

/// One [`WorkflowRuntime`] per pool, all over one queue and one store.
pub struct Pools {
    runtimes: HashMap<String, PoolRuntime>,
}

/// Builds a [`Pools`].
pub struct PoolsBuilder {
    queue: Arc<Queue>,
    store: Arc<dyn ObjectStore>,
    dispatch: Dispatch,
    hook: RecordHook,
    poll_interval: Duration,
    memo_retention: Duration,
    store_prefix: String,
    pools: Vec<(String, usize)>,
}

impl PoolsBuilder {
    /// Adds the pool `name` with `max_concurrent_steps` steps at a time. Its
    /// queue is `swale-pool-{name}`, and its memos are at `memos/{name}`
    /// within the store prefix.
    pub fn pool(mut self, name: impl Into<String>, max_concurrent_steps: usize) -> Self {
        self.pools.push((name.into(), max_concurrent_steps));
        self
    }

    /// The poll interval of every pool's step worker.
    pub fn poll_interval(mut self, interval: Duration) -> Self {
        self.poll_interval = interval;
        self
    }

    /// The time the memos and the run result record of a terminated task
    /// instance are kept, seven days by default. The scheduler does not read
    /// them, so the retention determines how long an operator can inspect
    /// them.
    pub fn memo_retention(mut self, retention: Duration) -> Self {
        self.memo_retention = retention;
        self
    }

    /// The path within the store that every pool writes its memos under,
    /// for a store whose queue is opened at a prefix. Empty by default.
    pub fn store_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.store_prefix = prefix.into();
        self
    }

    /// Builds the runtimes.
    pub fn build(self) -> Pools {
        let runtimes = self
            .pools
            .into_iter()
            .map(|(name, concurrency)| {
                let runtime = WorkflowRuntime::builder(
                    self.queue.clone(),
                    self.store.clone(),
                    self.dispatch.clone(),
                    self.hook.clone(),
                )
                .queue_name(format!("swale-pool-{name}"))
                .memo_prefix(store_path(&self.store_prefix, &format!("memos/{name}")))
                .max_concurrent_steps(concurrency)
                .poll_interval(self.poll_interval)
                .memo_retention(self.memo_retention)
                .build();
                (name, runtime)
            })
            .collect();
        Pools { runtimes }
    }
}

impl Pools {
    /// Starts building pools over `queue` and `store`, with `operators` as
    /// the dispatch and `hook` as the terminal hook of every pool.
    pub fn builder(
        queue: Arc<Queue>,
        store: Arc<dyn ObjectStore>,
        operators: Arc<OperatorSet>,
        hook: RecordHook,
    ) -> PoolsBuilder {
        let dispatch = Dispatch::new(operators, queue.clock());
        PoolsBuilder {
            queue,
            store,
            dispatch,
            hook,
            poll_interval: Duration::from_millis(250),
            memo_retention: Duration::from_secs(7 * 86_400),
            store_prefix: String::new(),
            pools: Vec::new(),
        }
    }

    /// The runtime of the pool `name`.
    pub fn runtime(&self, name: &str) -> Option<&PoolRuntime> {
        self.runtimes.get(name)
    }

    /// Spawns the step worker of every pool. Each stops when `shutdown` is
    /// cancelled.
    pub fn spawn(&self, shutdown: &CancellationToken) -> Vec<RunnerHandle> {
        self.runtimes
            .values()
            .map(|runtime| runtime.spawn(shutdown.clone().cancelled_owned()))
            .collect()
    }
}