dataflow_rs/run/
budget.rs

1//! Run-level condition error policy, task/expression budgets, and host
2//! accounting hooks.
3
4use crate::base::{BudgetCause, Error, ErrorKind, ExprCause, Result};
5use crate::expr::{Expr, ExprEngine, FromValue};
6use bumpalo::Bump;
7use datalogic_rs::Logic;
8use datavalue::DataValue;
9use serde::{Deserialize, Serialize};
10use std::fmt;
11use std::sync::atomic::{AtomicU64, Ordering};
12
13/// How condition evaluation failures are handled during a run.
14///
15/// Configured on [`RunOptions`](super::RunOptions::with_condition_error_policy)
16/// or as an engine-wide default on [`EngineConfig`](super::EngineConfig::with_condition_error_policy).
17#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub enum ConditionErrorPolicy {
20    /// Compatibility behavior: workflow, group, and task condition evaluation
21    /// errors skip the gated step, and `filter` ordinary evaluation errors
22    /// follow the task's `on_reject` policy ([`ErrorKind::Budget`] in `filter`
23    /// still halts the run).
24    #[default]
25    Skip,
26    /// Fail-closed: any condition evaluation error (workflow, group, task, or
27    /// `filter`, whether [`ErrorKind::Eval`] or [`ErrorKind::Budget`]) halts
28    /// the run immediately as [`RunStatus::Stopped`](super::RunStatus::Stopped)
29    /// with no further task dispatches.
30    Stop,
31}
32
33/// The kind of step about to be dispatched when charging a task budget or
34/// consulting a [`RunAccountant`].
35#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize)]
36#[serde(tag = "kind", rename_all = "snake_case")]
37pub enum DispatchKind<'a> {
38    /// An inline built-in task (`map`, `validation`, `filter`, `log`, etc.).
39    InlineBuiltin {
40        /// Workflow ID.
41        workflow_id: &'a str,
42        /// Task ID.
43        task_id: &'a str,
44        /// Built-in function name.
45        function: &'a str,
46    },
47    /// A host-registered async function (`Call`, `Parse`, or `Publish`).
48    HostTask {
49        /// Workflow ID.
50        workflow_id: &'a str,
51        /// Task ID.
52        task_id: &'a str,
53        /// Registered function name.
54        function: &'a str,
55    },
56    /// One sweep of a workflow loop.
57    LoopStep {
58        /// Workflow ID.
59        workflow_id: &'a str,
60        /// Zero-based sweep index (`0, 1, 2, ...`).
61        iteration: u64,
62    },
63    /// One element of a `for_each` fan-out.
64    ForEachElement {
65        /// Workflow ID.
66        workflow_id: &'a str,
67        /// Task ID.
68        task_id: &'a str,
69        /// Zero-based element index.
70        index: usize,
71    },
72}
73
74/// RAII guard returned by [`RunAccountant::before_task_dispatch`].
75///
76/// Holds an optional drop callback so host accountants can track in-flight
77/// concurrency or release per-dispatch leases when a step finishes or is
78/// cancelled (including when a `for_each` branch is dropped on fail-fast).
79pub struct OccupancyGuard<'a> {
80    on_drop: Option<Box<dyn FnOnce() + Send + 'a>>,
81}
82
83impl<'a> OccupancyGuard<'a> {
84    /// A guard that does nothing when dropped.
85    pub const fn noop() -> Self {
86        Self { on_drop: None }
87    }
88
89    /// A guard that runs `on_drop` once when dropped.
90    pub fn new(on_drop: impl FnOnce() + Send + 'a) -> Self {
91        Self {
92            on_drop: Some(Box::new(on_drop)),
93        }
94    }
95}
96
97impl Default for OccupancyGuard<'_> {
98    fn default() -> Self {
99        Self::noop()
100    }
101}
102
103impl fmt::Debug for OccupancyGuard<'_> {
104    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
105        f.debug_struct("OccupancyGuard")
106            .field("active", &self.on_drop.is_some())
107            .finish()
108    }
109}
110
111impl Drop for OccupancyGuard<'_> {
112    fn drop(&mut self) {
113        if let Some(f) = self.on_drop.take() {
114            f();
115        }
116    }
117}
118
119/// Host-owned accounting and admission hook for a single run.
120///
121/// Implementations must be thread-safe (`Send + Sync`) because `for_each`
122/// branches may invoke hooks concurrently across tasks.
123pub trait RunAccountant: Send + Sync {
124    /// Called immediately before dispatching an inline built-in, host task,
125    /// loop sweep, or `for_each` element.
126    ///
127    /// Returning `Err(e)` aborts the dispatch before any side effects occur;
128    /// the error is propagated as the step failure (typically [`ErrorKind::Budget`]).
129    ///
130    /// # Errors
131    ///
132    /// Returns an [`Error`] if the accountant rejects the dispatch.
133    fn before_task_dispatch(&self, kind: DispatchKind<'_>) -> Result<OccupancyGuard<'_>> {
134        let _ = kind;
135        Ok(OccupancyGuard::noop())
136    }
137
138    /// Called after each metered expression evaluation (both `Ok` and `Err`)
139    /// with the number of datalogic operations consumed.
140    fn on_expr_evaluated(&self, ops: u64) {
141        let _ = ops;
142    }
143
144    /// Called when a host function charges domain-specific work units via
145    /// [`Ctx::accountant`](crate::function::Ctx::accountant).
146    ///
147    /// # Errors
148    ///
149    /// Returns an [`Error`] if the custom charge exceeds host limits.
150    fn charge_custom(&self, kind: &str, units: u64) -> Result<()> {
151        let _ = (kind, units);
152        Ok(())
153    }
154}
155
156/// Shared per-run task and expression op budget tracker.
157pub(crate) struct TaskBudget<'a> {
158    max_tasks: Option<u64>,
159    tasks_used: AtomicU64,
160    max_expr_ops: Option<u64>,
161    expr_ops_used: AtomicU64,
162    accountant: Option<&'a dyn RunAccountant>,
163}
164
165impl fmt::Debug for TaskBudget<'_> {
166    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
167        f.debug_struct("TaskBudget")
168            .field("max_tasks", &self.max_tasks)
169            .field("tasks_used", &self.tasks_used.load(Ordering::Relaxed))
170            .field("max_expr_ops", &self.max_expr_ops)
171            .field("expr_ops_used", &self.expr_ops_used.load(Ordering::Relaxed))
172            .field("has_accountant", &self.accountant.is_some())
173            .finish()
174    }
175}
176
177impl<'a> TaskBudget<'a> {
178    /// Creates a new per-run budget tracker.
179    pub(crate) const fn new(
180        max_tasks: Option<u64>,
181        max_expr_ops: Option<u64>,
182        accountant: Option<&'a dyn RunAccountant>,
183    ) -> Self {
184        Self {
185            max_tasks,
186            tasks_used: AtomicU64::new(0),
187            max_expr_ops,
188            expr_ops_used: AtomicU64::new(0),
189            accountant,
190        }
191    }
192
193    /// An unlimited tracker with no accountant, used by standalone test helpers.
194    pub(crate) const fn unlimited() -> Self {
195        Self::new(None, None, None)
196    }
197
198    /// The host accountant attached to this run, if any.
199    pub(crate) const fn accountant(&self) -> Option<&'a dyn RunAccountant> {
200        self.accountant
201    }
202
203    /// Whether any run-level budget or accountant is active on this run.
204    pub(crate) const fn is_run_limited(&self) -> bool {
205        self.max_tasks.is_some() || self.max_expr_ops.is_some() || self.accountant.is_some()
206    }
207
208    /// Atomically charges one dispatch against `max_tasks` and consults the
209    /// [`RunAccountant`], returning its [`OccupancyGuard`] if admitted.
210    pub(crate) fn charge_dispatch(&self, kind: DispatchKind<'_>) -> Result<OccupancyGuard<'_>> {
211        if let Some(limit) = self.max_tasks {
212            let admitted =
213                self.tasks_used
214                    .try_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
215                        (used < limit).then_some(used + 1)
216                    });
217            if admitted.is_err() {
218                let used = limit.saturating_add(1);
219                return Err(Error::new(
220                    ErrorKind::Budget,
221                    format!("run task budget exceeded ({used} > {limit} max tasks)"),
222                )
223                .with_expr_cause(ExprCause::budget_exceeded(used, limit, None)));
224            }
225        }
226        match self.accountant {
227            Some(acc) => acc.before_task_dispatch(kind),
228            None => Ok(OccupancyGuard::noop()),
229        }
230    }
231
232    /// Evaluates `expr` against `root` in `arena`, charging the run's
233    /// `expr_ops_budget` and notifying the [`RunAccountant`] on both `Ok` and `Err`.
234    pub(crate) fn eval_view<'e, T: FromValue>(
235        &self,
236        exprs: &ExprEngine,
237        expr: &'e Expr<T>,
238        root: &DataValue<'_>,
239        arena: &Bump,
240    ) -> Result<T::Out<'e>> {
241        if let Some(c) = expr.constant() {
242            return Ok(T::lend(c));
243        }
244        let ops_limit = self.remaining_ops_limit(exprs);
245        let (mut res, ops) = exprs.eval_view_metered(expr, root, arena, ops_limit);
246        self.reconcile_ops(ops, &mut res);
247        res
248    }
249
250    /// Evaluates `logic` against `root` in `arena`, charging the run's
251    /// `expr_ops_budget` and notifying the [`RunAccountant`] on both `Ok` and `Err`.
252    pub(crate) fn eval_in<'v>(
253        &self,
254        exprs: &ExprEngine,
255        logic: &'v Logic,
256        reads_secret: bool,
257        root: &'v DataValue<'v>,
258        arena: &'v Bump,
259    ) -> Result<(&'v DataValue<'v>, Option<u64>)> {
260        let ops_limit = self.remaining_ops_limit(exprs);
261        let (mut res, ops) = exprs.eval_in_metered(logic, reads_secret, root, arena, ops_limit);
262        self.reconcile_ops(ops, &mut res);
263        res.map(|v| (v, ops))
264    }
265
266    /// Computes the per-evaluation op ceiling to pass to `datalogic-rs` when a
267    /// run-level `expr_ops_budget` is active.
268    fn remaining_ops_limit(&self, exprs: &ExprEngine) -> Option<u64> {
269        let limit = self.max_expr_ops?;
270        let used = self.expr_ops_used.load(Ordering::Relaxed);
271        let remaining = limit.saturating_sub(used);
272        #[cfg(feature = "budget")]
273        {
274            let engine_limit = exprs.logic().resolve_ops_budget(None);
275            if engine_limit < remaining {
276                return Some(engine_limit);
277            }
278        }
279        let _ = exprs;
280        Some(remaining)
281    }
282
283    /// Charges `ops` against `expr_ops_used`, notifies `accountant`, and
284    /// upgrades or rewrites `res` if the run-level `expr_ops_budget` was
285    /// exceeded.
286    fn reconcile_ops<T>(&self, ops: Option<u64>, res: &mut Result<T>) {
287        let Some(ops) = ops else {
288            return;
289        };
290        if let Some(acc) = self.accountant {
291            acc.on_expr_evaluated(ops);
292        }
293        let Some(limit) = self.max_expr_ops else {
294            return;
295        };
296        let prev = self.expr_ops_used.fetch_add(ops, Ordering::Relaxed);
297        let total = prev.saturating_add(ops);
298        if total <= limit {
299            return;
300        }
301        match res {
302            Err(err) if *err.kind() == ErrorKind::Budget => {
303                // `datalogic-rs` already halted at the exact AST node; update
304                // the budget numbers on `ExprCause` to reflect the run-level
305                // cumulative total and limit when the run ceiling was tighter.
306                if let Some(mut cause) = err.expr_cause().cloned()
307                    && cause.budget.is_none_or(|b| b.limit <= limit)
308                {
309                    cause.budget = Some(BudgetCause { used: total, limit });
310                    cause.ops = Some(ops);
311                    let updated = mem_take_error(err).with_expr_cause(cause);
312                    *err = updated;
313                }
314            }
315            _ => {
316                *res = Err(Error::new(
317                    ErrorKind::Budget,
318                    format!("run expression ops budget exceeded ({total} > {limit} max ops)"),
319                )
320                .with_expr_cause(ExprCause::budget_exceeded(
321                    total,
322                    limit,
323                    Some(ops),
324                )));
325            }
326        }
327    }
328}
329
330fn mem_take_error(err: &mut Error) -> Error {
331    std::mem::replace(err, Error::new(ErrorKind::Eval, ""))
332}
dataflow_rs/run/budget.rs

dataflow_rs/run/
budget.rs
