Skip to main content

dre_core/
engine.rs

1//! The run engine's parts around a Binding's run ([`crate::run`]):
2//!
3//! - the **plan**: which Bindings run, with every profile's entry checked ([`crate::run::plan`]);
4//! - the **executor**: runs a plan on a worker thread, cancellable through a [`CancelToken`],
5//!   reporting what happens as [`RunEvent`]s on a channel. The terminal UI, `--log-format json`,
6//!   the log file and `run_results.json` are consumers; whoever drives the run decides how
7//!   events look ([`dispatch`] replays them into a [`Ui`]);
8//! - the **[`RunStore`]**: where a run's files go in the target path.
9
10use std::path::{Path, PathBuf};
11use std::sync::Arc;
12use std::sync::atomic::{AtomicU8, Ordering};
13use std::sync::mpsc::Sender;
14use std::time::Duration;
15
16use dre_protocol::host::LogSink;
17use serde_json::{Map as JsonMap, Value as Json};
18
19use crate::run::{BindingOutcome, BindingPlan, Level, RunTargets, ShownMessage, Ui};
20
21/// Something that happened during a run.
22#[derive(Debug, Clone)]
23pub enum RunEvent {
24    /// The run's target and each used profile's entry, once they're all known to exist.
25    Targets(RunTargets),
26    /// The number of Bindings about to run.
27    Plan(usize),
28    BindingStart {
29        report: String,
30        set: Option<String>,
31    },
32    /// The running Binding's log file (`runs/<run id>/dre.log`), once its run folder exists.
33    BindingLog(PathBuf),
34    /// The schedule a Binding runs under (if any) and every var it uses.
35    BindingVars {
36        schedule: Option<String>,
37        schedule_vars: Option<JsonMap<String, Json>>,
38        vars: JsonMap<String, Json>,
39    },
40    /// One step inside the current Binding.
41    Step {
42        level: Level,
43        verb: String,
44        detail: String,
45        elapsed: Option<Duration>,
46    },
47    Warn(String),
48    BindingEnd(BindingOutcome),
49    /// A Binding compiled: what it would do.
50    Compiled(BindingPlan),
51    /// `--preview`: a rendered message, shown instead of sent.
52    Message(ShownMessage),
53    /// A line a plugin wrote to stderr: `(plugin, line)`.
54    PluginLog(String, String),
55    /// The full text of a statement sent to a source: `(label, sql)`.
56    Sql(String, String),
57    /// Bindings running at once (`threads:`): one Binding started or finished. Sent as it
58    /// happens; the Binding's own events follow later, together, as a [`RunEvent::Block`].
59    Running {
60        label: String,
61        running: bool,
62    },
63    /// One Binding's events, in order, when Bindings run at once: shown together so their lines
64    /// don't interleave.
65    Block(Vec<RunEvent>),
66}
67
68/// Where a running Binding sends its events. Cheap to clone, and `Send`: the executor's worker
69/// and plugins' stderr readers hold one each.
70#[derive(Clone)]
71pub struct Events {
72    tx: Sender<RunEvent>,
73    /// When Bindings run at once, a Binding's events are held here and sent as one block.
74    held: Option<std::sync::Arc<std::sync::Mutex<Vec<RunEvent>>>>,
75}
76
77impl Events {
78    pub fn new(sender: Sender<RunEvent>) -> Events {
79        Events {
80            tx: sender,
81            held: None,
82        }
83    }
84
85    /// Events that are held until [`Events::flush`], for one of several Bindings running at once.
86    pub fn held(&self) -> Events {
87        Events {
88            tx: self.tx.clone(),
89            held: Some(Default::default()),
90        }
91    }
92
93    /// Send an event. Nobody listening (the run was abandoned) isn't an error.
94    pub fn emit(&self, e: RunEvent) {
95        match &self.held {
96            Some(h) => h.lock().unwrap().push(e),
97            None => {
98                let _ = self.tx.send(e);
99            }
100        }
101    }
102
103    /// Send an event now, even from held events.
104    pub fn emit_now(&self, e: RunEvent) {
105        let _ = self.tx.send(e);
106    }
107
108    /// Send the held events as one block.
109    pub fn flush(&self) {
110        if let Some(h) = &self.held {
111            let events = std::mem::take(&mut *h.lock().unwrap());
112            if !events.is_empty() {
113                let _ = self.tx.send(RunEvent::Block(events));
114            }
115        }
116    }
117
118    pub fn step(&self, level: Level, verb: &str, detail: &str, elapsed: Option<Duration>) {
119        self.emit(RunEvent::Step {
120            level,
121            verb: verb.to_string(),
122            detail: detail.to_string(),
123            elapsed,
124        });
125    }
126
127    pub fn warn(&self, msg: &str) {
128        self.emit(RunEvent::Warn(msg.to_string()));
129    }
130
131    pub fn compiled(&self, plan: &BindingPlan) {
132        self.emit(RunEvent::Compiled(plan.clone()));
133    }
134
135    pub fn message(&self, message: &ShownMessage) {
136        self.emit(RunEvent::Message(message.clone()));
137    }
138
139    pub fn binding_vars(
140        &self,
141        schedule: Option<&str>,
142        schedule_vars: Option<&JsonMap<String, Json>>,
143        vars: &JsonMap<String, Json>,
144    ) {
145        self.emit(RunEvent::BindingVars {
146            schedule: schedule.map(str::to_string),
147            schedule_vars: schedule_vars.cloned(),
148            vars: vars.clone(),
149        });
150    }
151
152    /// Where plugin stderr goes: into the event stream.
153    pub fn plugin_log(&self) -> LogSink {
154        let events = self.clone();
155        Arc::new(move |plugin: &str, line: &str| events.emit(RunEvent::PluginLog(plugin.into(), line.into())))
156    }
157
158    /// Where every statement's full text goes: into the event stream.
159    pub fn sql_log(&self) -> LogSink {
160        let events = self.clone();
161        Arc::new(move |label: &str, sql: &str| events.emit(RunEvent::Sql(label.into(), sql.into())))
162    }
163}
164
165/// Replay an event into a [`Ui`]. `plugin_log` and `sql_log` are the UI's own sinks, taken once
166/// before the run.
167pub fn dispatch(ui: &mut dyn Ui, e: RunEvent, plugin_log: &LogSink, sql_log: &LogSink) {
168    match e {
169        RunEvent::Targets(t) => ui.targets(&t),
170        RunEvent::Plan(n) => ui.plan(n),
171        RunEvent::BindingStart { report, set } => ui.binding_start(&report, set.as_deref()),
172        RunEvent::BindingLog(path) => ui.binding_log(&path),
173        RunEvent::BindingVars {
174            schedule,
175            schedule_vars,
176            vars,
177        } => ui.binding_vars(schedule.as_deref(), schedule_vars.as_ref(), &vars),
178        RunEvent::Step {
179            level,
180            verb,
181            detail,
182            elapsed,
183        } => ui.step(level, &verb, &detail, elapsed),
184        RunEvent::Warn(msg) => ui.warn(&msg),
185        RunEvent::BindingEnd(o) => ui.binding_end(&o),
186        RunEvent::Compiled(p) => ui.compiled(&p),
187        RunEvent::Message(m) => ui.message(&m),
188        RunEvent::PluginLog(plugin, line) => plugin_log(&plugin, &line),
189        RunEvent::Sql(label, sql) => sql_log(&label, &sql),
190        RunEvent::Running { label, running } => ui.running(&label, running),
191        RunEvent::Block(events) => {
192            for e in events {
193                dispatch(ui, e, plugin_log, sql_log);
194            }
195        }
196    }
197}
198
199/// Why a run was asked to stop.
200#[derive(Debug, Clone, Copy, PartialEq, Eq)]
201pub enum CancelReason {
202    /// Ctrl-C (SIGINT).
203    Interrupt,
204    /// SIGTERM, or on Windows Ctrl-Break, closing the console, logging off or shutting down.
205    Terminate,
206    /// The run took longer than its timeout (`--timeout`, `DRE_RUN_TIMEOUT`, `flags: run_timeout`).
207    Timeout,
208}
209
210impl CancelReason {
211    /// The conventional exit code: 130 after an interrupt, 143 after a termination, 124 after a
212    /// timeout (as `timeout(1)`).
213    pub fn exit_code(self) -> u8 {
214        match self {
215            CancelReason::Interrupt => 130,
216            CancelReason::Terminate => 143,
217            CancelReason::Timeout => 124,
218        }
219    }
220
221    fn code(self) -> u8 {
222        match self {
223            CancelReason::Interrupt => 1,
224            CancelReason::Terminate => 2,
225            CancelReason::Timeout => 3,
226        }
227    }
228}
229
230/// Asks a run to stop. Clones share the request; the first reason given sticks.
231#[derive(Debug, Clone, Default)]
232pub struct CancelToken(Arc<AtomicU8>);
233
234impl CancelToken {
235    pub fn new() -> CancelToken {
236        CancelToken::default()
237    }
238
239    /// Cancel as an interrupt (Ctrl-C).
240    pub fn cancel(&self) {
241        self.cancel_with(CancelReason::Interrupt);
242    }
243
244    pub fn cancel_with(&self, reason: CancelReason) {
245        let _ = self
246            .0
247            .compare_exchange(0, reason.code(), Ordering::SeqCst, Ordering::SeqCst);
248    }
249
250    pub fn is_cancelled(&self) -> bool {
251        self.reason().is_some()
252    }
253
254    pub fn reason(&self) -> Option<CancelReason> {
255        match self.0.load(Ordering::SeqCst) {
256            1 => Some(CancelReason::Interrupt),
257            2 => Some(CancelReason::Terminate),
258            3 => Some(CancelReason::Timeout),
259            _ => None,
260        }
261    }
262}
263
264/// Where a run's files go in the target path, for one Binding (`<report>/<binding>`):
265/// `compiled/` (rendered SQL), `run/` (outputs and `run_results.json`) and `schema/` (the schema
266/// snapshot drift is checked against).
267#[derive(Debug, Clone)]
268pub struct RunStore {
269    target: PathBuf,
270    /// Bindings running at once: each destination path a Binding delivers to, and which Binding,
271    /// so two never write the same file.
272    claims: Option<std::sync::Arc<std::sync::Mutex<std::collections::HashMap<String, String>>>>,
273}
274
275impl RunStore {
276    pub fn new(target: &Path) -> RunStore {
277        RunStore {
278            target: target.to_path_buf(),
279            claims: None,
280        }
281    }
282
283    /// A store for Bindings running at once, which refuses a destination path two of them use.
284    pub fn concurrent(target: &Path) -> RunStore {
285        RunStore {
286            target: target.to_path_buf(),
287            claims: Some(Default::default()),
288        }
289    }
290
291    /// Claim `profile`'s `path` for `binding`; the Binding that already has it, if another.
292    pub fn claim(&self, profile: &str, path: &str, binding: &str) -> Option<String> {
293        let mut claims = self.claims.as_ref()?.lock().unwrap();
294        let owner = claims
295            .entry(format!("{profile}\u{0}{path}"))
296            .or_insert_with(|| binding.to_string());
297        (owner != binding).then(|| owner.clone())
298    }
299
300    fn binding(&self, kind: &str, report: &str, binding: &str) -> PathBuf {
301        self.target.join(kind).join(report).join(binding)
302    }
303
304    /// Where a Binding's rendered SQL goes.
305    pub fn compiled_dir(&self, report: &str, binding: &str) -> PathBuf {
306        self.binding("compiled", report, binding)
307    }
308
309    /// Where a Binding's outputs and `run_results.json` go.
310    pub fn run_dir(&self, report: &str, binding: &str) -> PathBuf {
311        self.binding("run", report, binding)
312    }
313
314    /// Where a Binding's schema snapshot goes.
315    pub fn schema_dir(&self, report: &str, binding: &str) -> PathBuf {
316        self.binding("schema", report, binding)
317    }
318
319    /// Write a Binding's `run_results.json` into its run folder (secrets masked).
320    pub fn write_results(
321        &self,
322        run_dir: &Path,
323        results: &crate::run_results::RunResults,
324    ) -> std::io::Result<()> {
325        std::fs::create_dir_all(run_dir)?;
326        crate::fs::write_atomic(
327            &run_dir.join("run_results.json"),
328            (crate::secrets::to_json_pretty(results)? + "\n").as_bytes(),
329        )
330    }
331}
332
333#[cfg(test)]
334mod tests {
335    use super::*;
336
337    #[test]
338    fn a_cancel_is_seen_by_every_clone() {
339        let a = CancelToken::new();
340        let b = a.clone();
341        assert!(!b.is_cancelled());
342        a.cancel_with(CancelReason::Terminate);
343        assert!(b.is_cancelled());
344        // The first reason sticks.
345        b.cancel();
346        assert_eq!(a.reason(), Some(CancelReason::Terminate));
347        assert_eq!(CancelReason::Terminate.exit_code(), 143);
348    }
349
350    #[test]
351    fn the_store_lays_out_a_bindings_folders() {
352        let s = RunStore::new(Path::new("/t"));
353        assert_eq!(s.run_dir("daily", "default"), Path::new("/t/run/daily/default"));
354        assert_eq!(s.compiled_dir("daily", "eu"), Path::new("/t/compiled/daily/eu"));
355        assert_eq!(s.schema_dir("daily", "eu"), Path::new("/t/schema/daily/eu"));
356    }
357
358    #[test]
359    fn events_arrive_in_order_and_log_sinks_feed_the_stream() {
360        let (tx, rx) = std::sync::mpsc::channel();
361        let events = Events::new(tx);
362        events.warn("one");
363        events.sql_log()("q.sql:1", "select 1");
364        events.plugin_log()("duckdb", "hello");
365        drop(events);
366        let got: Vec<String> = rx.iter().map(|e| format!("{e:?}")).collect();
367        assert_eq!(got.len(), 3);
368        assert!(got[0].starts_with("Warn") && got[1].starts_with("Sql") && got[2].starts_with("PluginLog"));
369    }
370}