Skip to main content

crankshaft_engine/
lib.rs

1//! The engine that powers Crankshaft.
2
3use anyhow::Context;
4use anyhow::Result;
5use crankshaft_config::backend::Config;
6use crankshaft_events::Event;
7use indexmap::IndexMap;
8use tokio::sync::broadcast;
9use tokio_util::sync::CancellationToken;
10use tracing::debug;
11
12pub mod service;
13pub mod task;
14
15pub use task::Task;
16
17use crate::service::Runner;
18use crate::service::runner::Backend;
19use crate::service::runner::TaskHandle;
20
21/// The capacity for the events channel.
22///
23/// This is the number of events to buffer in the channel before receivers
24/// become lagged.
25///
26/// The value of `100` was chosen simply as a reasonable default.
27const EVENTS_CHANNEL_CAPACITY: usize = 100;
28
29/// A workflow execution engine.
30#[derive(Debug)]
31pub struct Engine {
32    /// The task runner(s).
33    runners: IndexMap<String, Runner>,
34    /// The events sender.
35    events: Option<broadcast::Sender<Event>>,
36    /// The monitor for the engine.
37    #[cfg(feature = "monitoring")]
38    monitor: Option<crankshaft_monitor::Monitor>,
39}
40
41impl Engine {
42    /// Constructs a new engine.
43    pub fn new() -> Self {
44        let (events_tx, _) = broadcast::channel(EVENTS_CHANNEL_CAPACITY);
45        Self {
46            runners: Default::default(),
47            events: Some(events_tx),
48            #[cfg(feature = "monitoring")]
49            monitor: None,
50        }
51    }
52
53    /// Constructs a new engine with monitoring enabled.
54    #[cfg(feature = "monitoring")]
55    pub async fn new_with_monitoring(addr: std::net::SocketAddr) -> Self {
56        let (events_tx, _) = broadcast::channel(EVENTS_CHANNEL_CAPACITY);
57        let monitor = crankshaft_monitor::Monitor::start(addr, events_tx.clone()).await;
58
59        Self {
60            runners: Default::default(),
61            events: Some(events_tx),
62            monitor: Some(monitor),
63        }
64    }
65
66    /// Adds a [`Backend`] to the engine.
67    pub async fn with(mut self, config: Config) -> Result<Self> {
68        let (name, kind, max_tasks, defaults) = config.into_parts();
69        let runner = Runner::initialize(kind, max_tasks, defaults, self.events.clone()).await?;
70        self.runners.insert(name, runner);
71        Ok(self)
72    }
73
74    /// Subscribes to the engine's events and returns a receiver.
75    ///
76    /// Returns an error if the engine has already been shut down.
77    pub fn subscribe(&self) -> Result<broadcast::Receiver<Event>> {
78        Ok(self
79            .events
80            .as_ref()
81            .context("engine has shut down")?
82            .subscribe())
83    }
84
85    /// Gets the names of the runners.
86    pub fn runners(&self) -> impl Iterator<Item = &str> {
87        self.runners.keys().map(|key| key.as_ref())
88    }
89
90    /// Shuts down the engine.
91    pub async fn shutdown(mut self) {
92        // Drop the events sender
93        self.events.take();
94
95        #[cfg(feature = "monitoring")]
96        if let Some(monitor) = self.monitor.take() {
97            monitor.stop().await;
98        }
99    }
100
101    /// Spawns a [`Task`] to be executed.
102    ///
103    /// The `cancellation` token can be used to gracefully cancel the task.
104    ///
105    /// A [`TaskHandle`] is returned, which contains a channel that can be
106    /// awaited for the result of the job.
107    pub async fn spawn(
108        &self,
109        name: impl AsRef<str>,
110        task: Task,
111        token: CancellationToken,
112    ) -> Result<TaskHandle> {
113        let name = name.as_ref();
114        let backend = self
115            .runners
116            .get(name)
117            .unwrap_or_else(|| panic!("backend not found: {name}"));
118
119        debug!(
120            "submitting job{job} to the `{name}` backend",
121            job = task
122                .name
123                .as_ref()
124                .map(|name| format!(" with name `{name}`"))
125                .unwrap_or_default(),
126        );
127
128        backend.spawn(task, token).await
129    }
130
131    /// Starts an instrumentation loop.
132    #[cfg(tokio_unstable)]
133    pub async fn start_instrument(delay_ms: u64) {
134        use tokio_metrics::RuntimeMonitor;
135        use tracing::info;
136
137        let handle = tokio::runtime::Handle::current();
138        let monitor = RuntimeMonitor::new(&handle);
139
140        tokio::spawn(async move {
141            for interval in monitor.intervals() {
142                info!("{:?}", interval.total_park_count);
143                tokio::time::sleep(Duration::from_millis(delay_ms)).await;
144            }
145        });
146    }
147}
148
149impl Default for Engine {
150    fn default() -> Self {
151        Self::new()
152    }
153}
154
155impl Drop for Engine {
156    fn drop(&mut self) {
157        // Drop the events sender before the monitor
158        self.events.take();
159    }
160}