Add optional execution-cancellation breakpoint

Cargo.toml

@@ -56,6 +56,13 @@ lua = ["dep:mlua"]
 # Property-based testing support via proptest.
 proptest = ["dep:proptest", "dep:proptest-derive"]
 
+# Enables checking a caller-supplied cancellation flag on every expression resolution,
+# via `Runtime::set_cancellation_flag`. A no-op unless a flag is actually registered;
+# intended for embedders that need a hard safety net against runaway scripts (set the
+# flag from another thread, e.g. after your own timeout elapses, and the program panics
+# the next time it's observed).
+execution_cancellation = ["compiler"]
+
 # Enables environment variable access functions (e.g. `get_env_var`).
 enable_env_functions = []
 

src/compiler/context.rs

@@ -49,4 +49,20 @@ impl<'a> Context<'a> {
     pub fn timezone(&self) -> &TimeZone {
         self.timezone
     }
+
+    /// Checks whether the program has been cancelled via a caller-supplied
+    /// flag, panicking if so. Called on every expression resolution — it is
+    /// used for cancellation only, hence the name.
+    ///
+    /// Only compiled in when the `execution_cancellation` feature is
+    /// enabled; the call site in [`Expr::resolve`](super::expression::Expr)
+    /// is `#[cfg]`-gated the same way, so this is a true no-op otherwise. A
+    /// flag must also be registered via
+    /// [`RuntimeState::set_cancellation_flag`](super::state::RuntimeState::set_cancellation_flag)
+    /// for the check to do anything.
+    #[cfg(feature = "execution_cancellation")]
+    #[inline]
+    pub(crate) fn cancel_breakpoint(&mut self) {
+        self.state.check_cancellation();
+    }
 }

src/compiler/expression.rs

@@ -234,6 +234,9 @@ impl Expression for Expr {
             Return, Unary, Variable,
         };
 
+        #[cfg(feature = "execution_cancellation")]
+        ctx.cancel_breakpoint();
+
         match self {
             Literal(v) => v.resolve(ctx),
             Container(v) => v.resolve(ctx),

src/compiler/runtime.rs

@@ -1,3 +1,7 @@
+#[cfg(feature = "execution_cancellation")]
+use std::sync::Arc;
+#[cfg(feature = "execution_cancellation")]
+use std::sync::atomic::AtomicBool;
 use std::{error::Error, fmt};
 
 use crate::path::OwnedTargetPath;
@@ -127,3 +131,57 @@ impl Runtime {
         }
     }
 }
+
+#[cfg(feature = "execution_cancellation")]
+impl Runtime {
+    /// Registers a cancellation flag [`Runtime::resolve`] will check on
+    /// every expression resolution. Set it to `true` from any thread —
+    /// after your own timeout elapses, on a client disconnect, on shutdown,
+    /// whatever your cancellation source is — to abort a running program;
+    /// it panics rather than returning a [`RuntimeResult`].
+    pub fn set_cancellation_flag(&mut self, flag: Arc<AtomicBool>) {
+        self.state.set_cancellation_flag(flag);
+    }
+}
+
+#[cfg(all(test, feature = "execution_cancellation", feature = "stdlib"))]
+mod execution_cancellation_tests {
+    use std::collections::BTreeMap;
+    use std::sync::Arc;
+    use std::sync::atomic::AtomicBool;
+
+    use super::{Runtime, TimeZone};
+    use crate::compiler::state::RuntimeState;
+    use crate::value::Value;
+
+    #[test]
+    fn for_each_loop_panics_when_already_cancelled() {
+        let source = r"
+            count = 0
+            for_each(array!(.items)) -> |_index, _value| {
+                count = count + 1
+            }
+            count
+        ";
+
+        let program = crate::compiler::compile(source, &crate::stdlib::all())
+            .expect("program should compile")
+            .program;
+
+        let mut target: Value =
+            BTreeMap::from([("items".into(), Value::Array(vec![Value::from(1); 5_000]))]).into();
+
+        let flag = Arc::new(AtomicBool::new(true));
+        let mut runtime = Runtime::new(RuntimeState::default());
+        runtime.set_cancellation_flag(flag);
+
+        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+            runtime.resolve(&mut target, &program, &TimeZone::default())
+        }));
+
+        assert!(
+            result.is_err(),
+            "expected the for_each loop to be cancelled before it started"
+        );
+    }
+}

src/compiler/state.rs

@@ -1,6 +1,10 @@
 use crate::path::PathPrefix;
 use crate::value::{Kind, Value};
 use std::collections::{HashMap, hash_map::Entry};
+#[cfg(feature = "execution_cancellation")]
+use std::sync::Arc;
+#[cfg(feature = "execution_cancellation")]
+use std::sync::atomic::{AtomicBool, Ordering};
 
 use super::{TypeDef, parser::ast::Ident, type_def::Details, value::Collection};
 
@@ -176,6 +180,11 @@ impl ExternalEnv {
 pub struct RuntimeState {
     /// The [`Value`] stored in each variable.
     variables: HashMap<Ident, Value>,
+
+    /// An optional flag the running program is cancelled through.
+    /// See [`RuntimeState::set_cancellation_flag`].
+    #[cfg(feature = "execution_cancellation")]
+    cancel_flag: Option<Arc<AtomicBool>>,
 }
 
 impl RuntimeState {
@@ -215,3 +224,63 @@ impl RuntimeState {
         }
     }
 }
+
+#[cfg(feature = "execution_cancellation")]
+impl RuntimeState {
+    /// Registers a flag the running program will check on every expression
+    /// resolution. Set it to `true` from any thread — after your own
+    /// timeout elapses, on a client disconnect, on shutdown, whatever your
+    /// cancellation source is — and the program panics the next time it's
+    /// observed, rather than running to completion.
+    ///
+    /// This is a hard safety net against runaway scripts, not a regular
+    /// control-flow mechanism: cancellation panics rather than returning a
+    /// `Terminate` error.
+    pub fn set_cancellation_flag(&mut self, flag: Arc<AtomicBool>) {
+        self.cancel_flag = Some(flag);
+    }
+
+    /// Removes any previously registered cancellation flag.
+    pub fn clear_cancellation_flag(&mut self) {
+        self.cancel_flag = None;
+    }
+
+    pub(crate) fn check_cancellation(&self) {
+        if let Some(flag) = &self.cancel_flag {
+            assert!(
+                !flag.load(Ordering::Relaxed),
+                "VRL program execution was cancelled"
+            );
+        }
+    }
+}
+
+#[cfg(all(test, feature = "execution_cancellation"))]
+mod execution_cancellation_tests {
+    use super::RuntimeState;
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicBool, Ordering};
+
+    #[test]
+    fn panics_once_the_flag_is_set() {
+        let flag = Arc::new(AtomicBool::new(false));
+        let mut state = RuntimeState::default();
+        state.set_cancellation_flag(Arc::clone(&flag));
+
+        state.check_cancellation();
+
+        flag.store(true, Ordering::Relaxed);
+
+        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+            state.check_cancellation();
+        }));
+
+        assert!(result.is_err(), "expected check_cancellation to panic");
+    }
+
+    #[test]
+    fn no_flag_configured_never_panics() {
+        let state = RuntimeState::default();
+        state.check_cancellation();
+    }
+}