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