kata-containers/regorus
1
0
Fork 0
mirror of https://github.com/microsoft/regorus.git synced 2026-08-05 02:16:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat!: add cooperative execution-time limits across engine, VM, and binding (#539)

Browse Source 
- Introduce ExecutionTimer/ExecutionTimerConfig to allow limiting evaluating time.
- To amortize time checking costs, checking interval can be configured via the notion of work units
- A global fallback time limit can be set to universally limit all evaluation in addition to engine level limit setting.
- Implement limnits in interpreter and RVM. In RVM, also handle suspend/resume so that time during pause is not counted.
- Add engine-level APIs to set/clear per-engine timer configuration and apply global fallback defaults.
- Surface execution-time limits through FFI and C# bindings
- Add C# tests and example usage to validate engine overrides, global fallback behavior, and compiled policy enforcement.
- Expand docs for execution-time limit
- Add interpreter YAML cases and VM unit tests for time-limit behavior and deterministic time sources.

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2026-01-28 05:58:03 +05:30
committed by GitHub
parent e68e852ee3
commit 394625d4bc
32 changed files with 2259 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
src/engine.rs
View File

@@ -8,11 +8,13 @@ use crate::interpreter::*;
use crate::lexer::*;
use crate::parser::*;
use crate::scheduler::*;
use crate::utils::{gather_functions, limits};
use crate::utils::gather_functions;
use crate::utils::limits::{self, fallback_execution_timer_config, ExecutionTimerConfig};
use crate::value::*;
use crate::*;
use crate::{Extension, QueryResults};
use crate::Rc;
use anyhow::{anyhow, bail, Result};
/// The Rego evaluation engine.
@@ -23,6 +25,7 @@ pub struct Engine {
interpreter: Interpreter,
prepared: bool,
rego_v1: bool,
execution_timer_config: Option<ExecutionTimerConfig>,
}
#[cfg(feature = "azure_policy")]
@@ -62,14 +65,27 @@ impl Default for Engine {
}
impl Engine {
fn effective_execution_timer_config(&self) -> Option<ExecutionTimerConfig> {
self.execution_timer_config
.or_else(fallback_execution_timer_config)
}
fn apply_effective_execution_timer_config(&mut self) {
let config = self.effective_execution_timer_config();
self.interpreter.set_execution_timer_config(config);
}
/// Create an instance of [Engine].
pub fn new() -> Self {
Self {
let mut engine = Self {
modules: Rc::new(vec![]),
interpreter: Interpreter::new(),
prepared: false,
rego_v1: true,
}
execution_timer_config: None,
};
engine.apply_effective_execution_timer_config();
engine
}
/// Enable rego v0.
@@ -101,6 +117,60 @@ impl Engine {
self.rego_v1 = !rego_v0;
}
/// Configure the execution timer.
///
/// Stores the supplied configuration and ensures the next evaluation is checked against those
/// limits. Engines start without a time limit and otherwise fall back to the global
/// configuration (if provided).
///
/// # Examples
///
/// ```
/// use std::num::NonZeroU32;
/// use std::time::Duration;
/// use regorus::utils::limits::ExecutionTimerConfig;
/// use regorus::Engine;
///
/// let mut engine = Engine::new();
/// let config = ExecutionTimerConfig {
/// limit: Duration::from_millis(10),
/// check_interval: NonZeroU32::new(1).unwrap(),
/// };
///
/// engine.set_execution_timer_config(config);
/// ```
pub fn set_execution_timer_config(&mut self, config: ExecutionTimerConfig) {
self.execution_timer_config = Some(config);
self.interpreter.set_execution_timer_config(Some(config));
}
/// Clear the engine-specific execution timer configuration, falling back to the global value.
///
/// # Examples
///
/// ```
/// use std::num::NonZeroU32;
/// use std::time::Duration;
/// use regorus::utils::limits::{
/// set_fallback_execution_timer_config,
/// ExecutionTimerConfig,
/// };
/// use regorus::Engine;
///
/// let mut engine = Engine::new();
/// let global = ExecutionTimerConfig {
/// limit: Duration::from_millis(5),
/// check_interval: NonZeroU32::new(1).unwrap(),
/// };
/// set_fallback_execution_timer_config(Some(global));
///
/// engine.clear_execution_timer_config();
/// ```
pub fn clear_execution_timer_config(&mut self) {
self.execution_timer_config = None;
self.apply_effective_execution_timer_config();
}
/// Add a policy.
///
/// The policy file will be parsed and converted to AST representation.
@@ -524,6 +594,7 @@ impl Engine {
#[cfg_attr(docsrs, doc(cfg(feature = "azure_policy")))]
pub fn compile_for_target(&mut self) -> Result<CompiledPolicy> {
self.prepare_for_eval(false, true)?;
self.apply_effective_execution_timer_config();
self.interpreter.clean_internal_evaluation_state();
self.interpreter.compile(None).map(CompiledPolicy::new)
}
@@ -648,6 +719,7 @@ impl Engine {
/// - [`crate::compile_policy_with_entrypoint`] for a higher-level convenience function
pub fn compile_with_entrypoint(&mut self, rule: &Rc<str>) -> Result<CompiledPolicy> {
self.prepare_for_eval(false, false)?;
self.apply_effective_execution_timer_config();
self.interpreter.clean_internal_evaluation_state();
self.interpreter
.compile(Some(rule.clone()))
@@ -696,6 +768,7 @@ impl Engine {
/// ```
pub fn eval_rule(&mut self, rule: String) -> Result<Value> {
self.prepare_for_eval(false, false)?;
self.apply_effective_execution_timer_config();
self.interpreter.clean_internal_evaluation_state();
self.interpreter.eval_rule_in_path(rule)
}
@@ -737,6 +810,7 @@ impl Engine {
/// ```
pub fn eval_query(&mut self, query: String, enable_tracing: bool) -> Result<QueryResults> {
self.prepare_for_eval(enable_tracing, false)?;
self.apply_effective_execution_timer_config();
self.interpreter.clean_internal_evaluation_state();
self.interpreter.create_rule_prefixes()?;
@@ -936,6 +1010,8 @@ impl Engine {
enable_tracing: bool,
) -> Result<QueryResults> {
self.eval_modules(enable_tracing)?;
// Restart the timer window for the user query after module evaluation.
self.apply_effective_execution_timer_config();
let (query_module, query_node, query_schedule) = self.make_query(query)?;
self.interpreter
@@ -1011,6 +1087,7 @@ impl Engine {
enable_tracing: bool,
) -> Result<Value> {
self.prepare_for_eval(enable_tracing, false)?;
self.apply_effective_execution_timer_config();
self.interpreter.clean_internal_evaluation_state();
self.interpreter.eval_rule(module, rule)?;
@@ -1021,6 +1098,7 @@ impl Engine {
#[doc(hidden)]
pub fn eval_modules(&mut self, enable_tracing: bool) -> Result<Value> {
self.prepare_for_eval(enable_tracing, false)?;
self.apply_effective_execution_timer_config();
self.interpreter.clean_internal_evaluation_state();
// Ensure that empty modules are created.
@@ -1440,11 +1518,14 @@ impl Engine {
compiled_policy: Rc<crate::compiled_policy::CompiledPolicyData>,
) -> Self {
let modules = compiled_policy.modules.clone();
Self {
let mut engine = Self {
modules,
interpreter: Interpreter::new_from_compiled_policy(compiled_policy),
rego_v1: true, // Value doesn't matter since this is used only for policy parsing
prepared: true,
}
execution_timer_config: None,
};
engine.apply_effective_execution_timer_config();
engine
}
}

77
src/interpreter.rs
View File

@@ -15,7 +15,11 @@ use crate::lexer::*;
use crate::lookup::Lookup;
use crate::parser::Parser;
use crate::scheduler::*;
use crate::utils::limits::{monotonic_now, ExecutionTimer, ExecutionTimerConfig};
#[cfg(feature = "std")]
use crate::utils::*;
#[cfg(not(feature = "std"))]
use crate::utils::{get_extra_arg, get_path_string, get_root_var, FunctionTable};
use crate::value::*;
use crate::*;
use crate::{Expression, Extension, Location, QueryResult, QueryResults};
@@ -23,6 +27,7 @@ use crate::{Expression, Extension, Location, QueryResult, QueryResults};
#[cfg(feature = "coverage")]
use crate::query::traversal::traverse;
use crate::Rc;
use alloc::collections::btree_map::Entry as BTreeMapEntry;
use alloc::collections::{BTreeMap, BTreeSet};
use anyhow::{anyhow, bail, Result};
@@ -96,6 +101,7 @@ pub struct Interpreter {
active_rules: Vec<Ref<Rule>>,
builtins_cache: BTreeMap<(&'static str, Vec<Value>), Value>,
no_rules_lookup: bool,
execution_timer: ExecutionTimer,
}
impl Default for Interpreter {
@@ -143,6 +149,7 @@ impl Clone for Interpreter {
query_module: None,
module: None,
no_rules_lookup: false,
execution_timer: ExecutionTimer::new(self.execution_timer.config()),
}
}
}
@@ -223,6 +230,7 @@ impl Interpreter {
gather_prints: false,
prints: Vec::default(),
execution_timer: ExecutionTimer::new(None),
}
}
@@ -266,9 +274,38 @@ impl Interpreter {
.data
.clone()
.unwrap_or_else(Value::new_object),
execution_timer: ExecutionTimer::new(None),
}
}
fn reset_execution_timer_state(&mut self) {
self.execution_timer.reset();
if self.execution_timer.limit().is_none() {
return;
}
if let Some(now) = monotonic_now() {
self.execution_timer.start(now);
}
}
fn execution_timer_tick(&mut self, work_units: u32) -> Result<()> {
if self.execution_timer.limit().is_none() {
return Ok(());
}
let Some(now) = monotonic_now() else {
return Ok(());
};
self.execution_timer.tick(work_units, now)?;
Ok(())
}
#[inline]
fn check_execution_time(&mut self) -> Result<()> {
self.execution_timer_tick(1)
}
fn compiled_policy_mut(&mut self) -> &mut CompiledPolicyData {
Rc::make_mut(&mut self.compiled_policy)
}
@@ -329,6 +366,11 @@ impl Interpreter {
self.compiled_policy_mut().strict_builtin_errors = b;
}
pub fn set_execution_timer_config(&mut self, config: Option<ExecutionTimerConfig>) {
self.execution_timer = ExecutionTimer::new(config);
self.reset_execution_timer_state();
}
pub fn set_input(&mut self, input: Value) {
self.input = input.clone();
// Update with_document["input"] too, in case if engine is being reused and was already prepared
@@ -357,6 +399,7 @@ impl Interpreter {
self.contexts = vec![];
self.rule_values.clear();
self.builtins_cache.clear();
self.reset_execution_timer_state();
}
#[cfg(feature = "allocator-memory-limits")]
@@ -366,7 +409,7 @@ impl Interpreter {
}
#[cfg(not(feature = "allocator-memory-limits"))]
fn memory_check(&mut self) -> Result<()> {
const fn memory_check(&mut self) -> Result<()> {
let _ = self; // quiet clippy::unused_self; retained for symmetry with VM path
Ok(())
}
@@ -640,6 +683,7 @@ impl Interpreter {
domain: &ExprRef,
query: &Ref<Query>,
) -> Result<bool> {
self.check_execution_time()?;
let domain = self.eval_expr(domain)?;
self.scopes.push(Scope::new());
@@ -700,6 +744,7 @@ impl Interpreter {
plan: &DestructuringPlan,
value: &Value,
) -> Result<Value> {
self.check_execution_time()?;
if value == &Value::Undefined {
return Ok(Value::Undefined);
}
@@ -799,6 +844,7 @@ impl Interpreter {
}
fn execute_assignment_plan(&mut self, plan: &AssignmentPlan) -> Result<Value> {
self.check_execution_time()?;
match plan {
AssignmentPlan::ColonEquals {
lhs_expr: _,
@@ -894,6 +940,7 @@ impl Interpreter {
collection: &ExprRef,
stmts: &[&LiteralStmt],
) -> Result<bool> {
self.check_execution_time()?;
let scope_saved = self.current_scope()?.clone();
let mut count: usize = 0;
@@ -1051,7 +1098,7 @@ impl Interpreter {
fn eval_stmt_impl(&mut self, stmt: &LiteralStmt, stmts: &[&LiteralStmt]) -> Result<bool> {
self.memory_check()?;
self.check_execution_time()?;
Ok(match &stmt.literal {
Literal::Expr { span, expr, .. } => {
let value = match expr.as_ref() {
@@ -1343,7 +1390,7 @@ impl Interpreter {
loops: &[HoistedLoop],
) -> Result<bool> {
self.memory_check()?;
self.check_execution_time()?;
if loops.is_empty() {
if let Some((first_stmt, tail_stmts)) = stmts.split_first() {
// Evaluate the current statement whose loop expressions have been hoisted.
@@ -1569,6 +1616,7 @@ impl Interpreter {
}
fn eval_rule_ref(&mut self, rule_refr: &ExprRef) -> Result<Vec<Value>> {
self.check_execution_time()?;
let mut comps = vec![];
let mut expr = rule_refr;
loop {
@@ -1711,6 +1759,7 @@ impl Interpreter {
}
fn eval_output_expr_in_loop(&mut self, loops: &[HoistedLoop]) -> Result<bool> {
self.check_execution_time()?;
if loops.is_empty() {
let (key_expr, output_expr) = self.get_exprs_from_context()?;
@@ -1944,6 +1993,7 @@ impl Interpreter {
}
fn eval_output_expr(&mut self) -> Result<bool> {
self.check_execution_time()?;
// Evaluate output expression after all the statements have been executed.
let (key_expr, output_expr) = self.get_exprs_from_context()?;
@@ -2074,6 +2124,7 @@ impl Interpreter {
}
fn eval_query(&mut self, query: &Ref<Query>) -> Result<bool> {
self.check_execution_time()?;
// Execute the query in a new scope
self.scopes.push(Scope::new());
let order_indices = {
@@ -2157,6 +2208,7 @@ impl Interpreter {
}
fn eval_array(&mut self, items: &Vec<ExprRef>) -> Result<Value> {
self.check_execution_time()?;
let mut array = Vec::new();
for item in items {
@@ -2172,6 +2224,7 @@ impl Interpreter {
}
fn eval_object(&mut self, fields: &Vec<(Span, ExprRef, ExprRef)>) -> Result<Value> {
self.check_execution_time()?;
let mut object = BTreeMap::new();
for (_, key, value) in fields {
@@ -2194,6 +2247,7 @@ impl Interpreter {
}
fn eval_set(&mut self, items: &Vec<ExprRef>) -> Result<Value> {
self.check_execution_time()?;
let mut set = BTreeSet::new();
for item in items {
@@ -2213,6 +2267,7 @@ impl Interpreter {
value: &ExprRef,
collection: &ExprRef,
) -> Result<Value> {
self.check_execution_time()?;
let value = self.eval_expr(value)?;
let collection = self.eval_expr(collection)?;
@@ -2250,6 +2305,7 @@ impl Interpreter {
}
fn eval_array_compr(&mut self, term: &ExprRef, query: &Ref<Query>) -> Result<Value> {
self.check_execution_time()?;
// Push new context
self.contexts.push(Context {
output_expr: Some(term.clone()),
@@ -2268,6 +2324,7 @@ impl Interpreter {
}
fn eval_set_compr(&mut self, term: &ExprRef, query: &Ref<Query>) -> Result<Value> {
self.check_execution_time()?;
// Push new context
self.contexts.push(Context {
output_expr: Some(term.clone()),
@@ -2290,6 +2347,7 @@ impl Interpreter {
value: &ExprRef,
query: &Ref<Query>,
) -> Result<Value> {
self.check_execution_time()?;
// Push new context
self.contexts.push(Context {
key_expr: Some(key.clone()),
@@ -2327,6 +2385,7 @@ impl Interpreter {
params: &[ExprRef],
args: Vec<Value>,
) -> Result<Value> {
self.check_execution_time()?;
// If any argument is undefined, then the call is undefined.
if args.iter().any(|a| a == &Value::Undefined) {
return Ok(Value::Undefined);
@@ -2467,6 +2526,7 @@ impl Interpreter {
fcn: &ExprRef,
params: &[ExprRef],
) -> Result<Value> {
self.check_execution_time()?;
// Return generated values of walk builtin.
if let Some(v) = self.get_loop_var_value(expr)? {
return Ok(v.clone());
@@ -2783,6 +2843,7 @@ impl Interpreter {
extra_arg: Option<ExprRef>,
allow_return_arg: bool,
) -> Result<Value> {
self.check_execution_time()?;
// TODO: global var check; interop with `some var`
if extra_arg.is_some() {
let (last_param, arg_prefix) = params
@@ -2832,6 +2893,7 @@ impl Interpreter {
}
fn ensure_module_evaluated(&mut self, path: String) -> Result<()> {
self.check_execution_time()?;
for module in self.compiled_policy.modules.clone().iter().cloned() {
if Some(&module) == self.module.as_ref() {
// Prevent cyclic evaluation.
@@ -2875,6 +2937,7 @@ impl Interpreter {
}
fn ensure_rule_evaluated(&mut self, path: String) -> Result<()> {
self.check_execution_time()?;
let mut matched = false;
if let Some(rules) = self.compiled_policy.rules.get(&path) {
matched = true;
@@ -3055,6 +3118,7 @@ impl Interpreter {
}
fn eval_expr(&mut self, expr: &ExprRef) -> Result<Value> {
self.check_execution_time()?;
#[cfg(feature = "coverage")]
if self.enable_coverage {
let span = expr.span();
@@ -3235,6 +3299,7 @@ impl Interpreter {
span: &Span,
bodies: &[RuleBody],
) -> Result<Value> {
self.check_execution_time()?;
let n_scopes = self.scopes.len();
let result = if bodies.is_empty() {
self.contexts.push(ctx.clone());
@@ -3496,6 +3561,7 @@ impl Interpreter {
}
pub fn eval_default_rule(&mut self, rule: &Ref<Rule>) -> Result<()> {
self.check_execution_time()?;
// Skip reprocessing rule.
if self.processed.contains(rule) {
return Ok(());
@@ -3569,6 +3635,7 @@ impl Interpreter {
/// Evaluate a default rule and return the resulting value for compiler consumers.
#[cfg(feature = "rvm")]
pub fn eval_default_rule_for_compiler(&mut self, rule_path: &str) -> Result<Value> {
self.check_execution_time()?;
self.input = Value::Undefined;
self.data = Value::Undefined;
self.ensure_loop_var_values_capacity();
@@ -3658,6 +3725,7 @@ impl Interpreter {
}
fn eval_rule_impl(&mut self, module: &Ref<Module>, rule: &Ref<Rule>) -> Result<()> {
self.check_execution_time()?;
match rule.as_ref() {
Rule::Spec {
span,
@@ -3744,6 +3812,7 @@ impl Interpreter {
}
pub fn eval_rule(&mut self, module: &Ref<Module>, rule: &Ref<Rule>) -> Result<()> {
self.check_execution_time()?;
// Set current module index
self.current_module_index = self.find_module_index(module);
@@ -3802,6 +3871,7 @@ impl Interpreter {
query_schedule: Schedule,
enable_tracing: bool,
) -> Result<QueryResults> {
self.check_execution_time()?;
self.traces = match enable_tracing {
true => Some(vec![]),
false => None,
@@ -4308,6 +4378,7 @@ impl Interpreter {
}
pub fn eval_rule_in_path(&mut self, path: String) -> Result<Value> {
self.check_execution_time()?;
if !self.compiled_policy.rule_paths.contains(&path) {
bail!("not a valid rule path");
}

2
src/lib.rs
View File

@@ -151,7 +151,7 @@ mod schema;
pub mod target;
#[cfg(any(test, all(feature = "yaml", feature = "std")))]
pub mod test_utils;
mod utils;
pub mod utils;
mod value;
#[cfg(feature = "azure_policy")]

280
src/rvm/tests/vm.rs
View File

@@ -18,8 +18,17 @@
#[cfg(test)]
mod tests {
use crate::rvm::program::Program;
use crate::rvm::tests::instruction_parser::{parse_instruction, parse_loop_mode};
use crate::rvm::tests::test_utils::test_round_trip_serialization;
#[cfg(any(test, not(feature = "std")))]
use crate::utils::limits::set_time_source;
use crate::utils::limits::{
acquire_limits_test_lock, fallback_execution_timer_config,
set_fallback_execution_timer_config, ExecutionTimerConfig, TimeSource,
};
use core::num::NonZeroU32;
use core::time::Duration;
#[derive(Debug, Clone, Deserialize, Serialize, Default)]
struct RuleInfoSpec {
rule_type: String,
@@ -48,11 +57,122 @@ mod tests {
use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::fs;
use std::sync::{Mutex, Once};
use test_generator::test_resources;
extern crate alloc;
extern crate std;
struct FallbackGuard(Option<ExecutionTimerConfig>);
impl Drop for FallbackGuard {
fn drop(&mut self) {
set_fallback_execution_timer_config(self.0);
}
}
fn install_fallback_config(config: Option<ExecutionTimerConfig>) -> FallbackGuard {
let previous = fallback_execution_timer_config();
set_fallback_execution_timer_config(config);
FallbackGuard(previous)
}
struct TimeSourceGuard {
previous_default: Duration,
previous_template: Vec<Duration>,
}
impl Drop for TimeSourceGuard {
fn drop(&mut self) {
let mut state = TIME_SOURCE_STATE
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
state.default_increment = self.previous_default;
state.template_increments = self.previous_template.clone();
state.reset_from_template();
}
}
fn configure_time_source(
increments: Vec<Duration>,
default_increment: Duration,
) -> TimeSourceGuard {
ensure_time_source_registered();
let mut state = TIME_SOURCE_STATE
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let guard = TimeSourceGuard {
previous_default: state.default_increment,
previous_template: state.template_increments.clone(),
};
state.default_increment = default_increment;
state.template_increments = increments;
state.reset_from_template();
guard
}
struct TestTimeSource;
struct TimeSourceState {
current: Duration,
started: bool,
default_increment: Duration,
increments: VecDeque<Duration>,
template_increments: Vec<Duration>,
}
impl TimeSourceState {
const fn new() -> Self {
Self {
current: Duration::ZERO,
started: false,
default_increment: Duration::from_millis(1),
increments: VecDeque::new(),
template_increments: Vec::new(),
}
}
fn reset_from_template(&mut self) {
self.current = Duration::ZERO;
self.started = false;
self.increments = VecDeque::from(self.template_increments.clone());
}
}
impl TimeSource for TestTimeSource {
fn now(&self) -> Option<Duration> {
let mut state = TIME_SOURCE_STATE
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if !state.started {
state.started = true;
return Some(state.current);
}
let increment = state
.increments
.pop_front()
.unwrap_or(state.default_increment);
state.current = state.current.saturating_add(increment);
Some(state.current)
}
}
static TEST_TIME_SOURCE: TestTimeSource = TestTimeSource;
static TIME_SOURCE_STATE: Mutex<TimeSourceState> = Mutex::new(TimeSourceState::new());
static TIME_SOURCE_ONCE: Once = Once::new();
fn ensure_time_source_registered() {
#[cfg(any(test, not(feature = "std")))]
TIME_SOURCE_ONCE.call_once(|| {
let _ = set_time_source(&TEST_TIME_SOURCE);
});
}
#[derive(Debug, Clone, Deserialize, Serialize)]
struct HostAwaitResponseSpec {
id: crate::Value,
@@ -62,6 +182,13 @@ mod tests {
values: Vec<crate::Value>,
}
fn default_vm_test_execution_timer_config() -> ExecutionTimerConfig {
ExecutionTimerConfig {
limit: Duration::from_secs(1),
check_interval: NonZeroU32::new(100).unwrap_or(NonZeroU32::MIN),
}
}
#[derive(Debug, Deserialize, Serialize)]
struct VmTestCase {
note: String,
@@ -644,6 +771,7 @@ mod tests {
vm.set_execution_mode(mode);
vm.set_step_mode(use_step_mode);
vm.set_strict_builtin_errors(strict);
vm.set_execution_timer_config(Some(default_vm_test_execution_timer_config()));
if let Some(data_value) = processed_data.clone() {
vm.set_data(data_value)?;
@@ -1001,6 +1129,158 @@ mod tests {
Ok(())
}
#[test]
fn vm_execution_time_limit_triggers_error() -> Result<()> {
use crate::rvm::instructions::Instruction;
use crate::utils::limits::acquire_limits_test_lock;
use core::num::NonZeroU32;
use core::time::Duration;
let _lock = acquire_limits_test_lock();
let config = ExecutionTimerConfig {
limit: Duration::from_nanos(1),
check_interval: NonZeroU32::new(1).unwrap(),
};
let _guard = install_fallback_config(Some(config));
let mut program = Program::new();
program.dispatch_window_size = 2;
program.max_rule_window_size = 2;
program.entry_points.insert("main".to_string(), 0);
const INSTRUCTION_COUNT: usize = 60_000;
program.instructions = (0..INSTRUCTION_COUNT)
.map(|_| Instruction::LoadNull { dest: 0 })
.collect();
program.instructions.push(Instruction::Return { value: 0 });
program.instruction_spans = alloc::vec![None; program.instructions.len()];
program.main_entry_point = 0;
let program = Arc::new(program);
let mut vm = RegoVM::new();
vm.set_max_instructions(usize::MAX);
vm.load_program(program);
let result = vm.execute();
assert!(
matches!(result, Err(VmError::TimeLimitExceeded { .. })),
"expected time limit error but got {result:?}"
);
Ok(())
}
#[test]
fn vm_execution_time_limit_override_allows_completion() -> Result<()> {
use crate::rvm::instructions::Instruction;
use crate::utils::limits::acquire_limits_test_lock;
use core::num::NonZeroU32;
use core::time::Duration;
let _lock = acquire_limits_test_lock();
let strict_config = ExecutionTimerConfig {
limit: Duration::from_nanos(1),
check_interval: NonZeroU32::new(1).unwrap(),
};
let _guard = install_fallback_config(Some(strict_config));
let mut program = Program::new();
program.dispatch_window_size = 2;
program.max_rule_window_size = 2;
program.entry_points.insert("main".to_string(), 0);
program.instructions = alloc::vec![
Instruction::LoadNull { dest: 0 },
Instruction::Return { value: 0 },
];
program.instruction_spans = alloc::vec![None; program.instructions.len()];
program.main_entry_point = 0;
let program = Arc::new(program);
let mut vm = RegoVM::new();
vm.load_program(program);
let relaxed_config = ExecutionTimerConfig {
limit: Duration::from_millis(10),
check_interval: NonZeroU32::new(1).unwrap(),
};
vm.set_execution_timer_config(Some(relaxed_config));
let result = vm.execute();
assert!(
result.is_ok(),
"expected successful execution, got {result:?}"
);
Ok(())
}
#[test]
fn vm_suspend_resume_excludes_suspended_time_from_limit() -> Result<()> {
use crate::rvm::instructions::Instruction;
let _lock = acquire_limits_test_lock();
let _guard = install_fallback_config(Some(ExecutionTimerConfig {
limit: Duration::from_millis(10),
check_interval: NonZeroU32::new(1).unwrap(),
}));
let _time_guard = configure_time_source(
alloc::vec![
Duration::from_millis(1),
Duration::from_millis(1),
Duration::from_millis(1),
Duration::from_millis(1),
Duration::from_millis(100),
Duration::from_millis(1),
],
Duration::from_millis(1),
);
let mut program = Program::new();
program.dispatch_window_size = 3;
program.max_rule_window_size = 3;
program.entry_points.insert("main".to_string(), 0);
program.literals = alloc::vec![Value::from("id"), Value::from(1)];
program.instructions = alloc::vec![
Instruction::Load {
dest: 0,
literal_idx: 0
},
Instruction::Load {
dest: 1,
literal_idx: 1
},
Instruction::HostAwait {
dest: 2,
arg: 1,
id: 0
},
Instruction::Return { value: 2 },
];
program.instruction_spans = alloc::vec![None; program.instructions.len()];
program.main_entry_point = 0;
let program = Arc::new(program);
let mut vm = RegoVM::new();
vm.set_execution_mode(ExecutionMode::Suspendable);
vm.load_program(program);
let _ = vm.execute()?;
match vm.execution_state() {
ExecutionState::Suspended { reason, .. } => {
assert!(matches!(reason, SuspendReason::HostAwait { .. }));
}
other => panic!("expected suspension, got {other:?}"),
}
let resumed = vm.resume(Some(Value::from(42)))?;
assert_eq!(resumed, Value::from(42));
Ok(())
}
#[test_resources("tests/rvm/vm/suites/*.yaml")]
fn run_vm_test_file(file: &str) {
run_vm_test_suite(file).unwrap()

10
src/rvm/vm/errors.rs
View File

@@ -5,6 +5,7 @@ use super::execution_model::SuspendReason;
use crate::value::Value;
use alloc::string::String;
use alloc::vec::Vec;
use core::time::Duration;
use thiserror::Error;
/// VM execution errors
@@ -17,7 +18,14 @@ pub enum VmError {
pc: usize,
},
#[error("Execution stopped: exceeded maximum memory limit of {limit} bytes with usage {usage} bytes (pc={pc})")]
#[error("Execution exceeded time limit (elapsed={elapsed:?}, limit={limit:?}, pc={pc})")]
TimeLimitExceeded {
elapsed: Duration,
limit: Duration,
pc: usize,
},
#[error("Execution exceeded memory limit (usage={usage} bytes, limit={limit} bytes, pc={pc})")]
MemoryLimitExceeded { usage: u64, limit: u64, pc: usize },
#[error("Literal index {index} out of bounds (pc={pc})")]

14
src/rvm/vm/execution.rs
View File

@@ -59,6 +59,7 @@ impl RegoVM {
match self.execution_mode {
ExecutionMode::RunToCompletion => {
self.reset_execution_state();
self.reset_execution_timer_state();
self.validate_vm_state()?;
let entry_point_pc_u32 = u32::try_from(entry_point_pc).map_err(|_| {
@@ -73,6 +74,7 @@ impl RegoVM {
}
ExecutionMode::Suspendable => {
self.reset_execution_state();
self.reset_execution_timer_state();
self.validate_vm_state()?;
self.execute_suspendable_entry(entry_point_pc)
@@ -101,6 +103,7 @@ impl RegoVM {
match self.execution_mode {
ExecutionMode::RunToCompletion => {
self.reset_execution_state();
self.reset_execution_timer_state();
self.validate_vm_state()?;
let entry_point_pc_u32 = u32::try_from(entry_point_pc).map_err(|_| {
@@ -115,6 +118,7 @@ impl RegoVM {
}
ExecutionMode::Suspendable => {
self.reset_execution_state();
self.reset_execution_timer_state();
self.validate_vm_state()?;
self.execute_suspendable_entry(entry_point_pc)
@@ -136,6 +140,7 @@ impl RegoVM {
});
}
self.execution_timer_tick(1)?;
self.executed_instructions = self.executed_instructions.saturating_add(1);
let instruction = program.instructions.get(self.pc).cloned().ok_or(
VmError::ProgramCounterOutOfBounds {
@@ -168,6 +173,7 @@ impl RegoVM {
fn execute_run_to_completion(&mut self) -> Result<Value> {
self.reset_execution_state();
self.reset_execution_timer_state();
self.execution_state = ExecutionState::Running;
match self.jump_to(0_u32) {
Ok(value) => {
@@ -185,6 +191,7 @@ impl RegoVM {
fn execute_suspendable(&mut self) -> Result<Value> {
self.reset_execution_state();
self.reset_execution_timer_state();
self.execution_state = ExecutionState::Running;
match self.run_stackless_from(0) {
Ok(result) => Ok(result),
@@ -197,6 +204,7 @@ impl RegoVM {
fn execute_suspendable_entry(&mut self, entry_point_pc: usize) -> Result<Value> {
self.execution_state = ExecutionState::Running;
self.reset_execution_timer_state();
match self.run_stackless_from(entry_point_pc) {
Ok(result) => Ok(result),
Err(err) => {
@@ -256,6 +264,7 @@ impl RegoVM {
}
self.execution_state = ExecutionState::Running;
self.restore_execution_timer_after_resume();
let program = self.program.clone();
self.run_stackless_loop(&program, &mut last_result)?;
@@ -376,6 +385,10 @@ impl RegoVM {
}
self.pc = frame_pc;
if let Err(err) = self.execution_timer_tick(1) {
self.execution_state = ExecutionState::Error { error: err.clone() };
return Err(err);
}
let instruction = program.instructions.get(self.pc).cloned().ok_or(
VmError::ProgramCounterOutOfBounds {
pc: self.pc,
@@ -460,6 +473,7 @@ impl RegoVM {
}
}
self.snapshot_execution_timer_on_suspend();
self.execution_state = ExecutionState::Suspended {
reason,
pc: self.pc,

109
src/rvm/vm/machine.rs
View File

@@ -3,15 +3,21 @@
use crate::rvm::program::Program;
#[cfg(feature = "allocator-memory-limits")]
use crate::utils::limits::{self, LimitError};
use crate::utils::limits;
use crate::utils::limits::{
fallback_execution_timer_config, monotonic_now, ExecutionTimer, ExecutionTimerConfig,
LimitError,
};
use crate::value::Value;
use crate::CompiledPolicy;
use alloc::collections::{btree_map::Entry, BTreeMap, VecDeque};
#[cfg(feature = "allocator-memory-limits")]
use alloc::format;
use alloc::string::String;
use alloc::sync::Arc;
use alloc::vec;
use alloc::vec::Vec;
use core::time::Duration;
use super::context::{CallRuleContext, ComprehensionContext, LoopContext};
use super::errors::{Result, VmError};
@@ -105,6 +111,15 @@ pub struct RegoVM {
/// Cache for builtin calls that must stay deterministic across a single evaluation
pub(super) builtins_cache: BTreeMap<(&'static str, Vec<Value>), Value>,
/// Optional override for the execution timer configuration
pub(super) execution_timer_config: Option<ExecutionTimerConfig>,
/// Cooperative execution timer used to enforce wall-clock limits
pub(super) execution_timer: ExecutionTimer,
/// Elapsed wall-clock time recorded when the VM entered a suspended state
pub(super) execution_timer_elapsed_at_suspend: Option<Duration>,
}
impl Default for RegoVM {
@@ -116,6 +131,8 @@ impl Default for RegoVM {
impl RegoVM {
/// Create a new virtual machine
pub fn new() -> Self {
let fallback_timer = fallback_execution_timer_config();
RegoVM {
registers: Vec::new(), // Start with no registers - will be resized when program is loaded
pc: 0,
@@ -143,6 +160,9 @@ impl RegoVM {
frame_pc_overridden: false,
strict_builtin_errors: false,
builtins_cache: BTreeMap::new(),
execution_timer_config: None,
execution_timer: ExecutionTimer::new(fallback_timer),
execution_timer_elapsed_at_suspend: None,
}
}
@@ -319,6 +339,93 @@ impl RegoVM {
self.execution_mode
}
/// Configure the execution timer to use the supplied configuration, or fall back to the global
/// default when `None` is provided.
pub fn set_execution_timer_config(&mut self, config: Option<ExecutionTimerConfig>) {
self.execution_timer_config = config;
self.reset_execution_timer_state();
}
/// Returns the currently configured execution timer, if any.
pub const fn execution_timer_config(&self) -> Option<ExecutionTimerConfig> {
self.execution_timer_config
}
pub(super) fn reset_execution_timer_state(&mut self) {
let config = self.effective_execution_timer_config();
self.execution_timer = ExecutionTimer::new(config);
self.execution_timer_elapsed_at_suspend = None;
if config.is_none() {
return;
}
if let Some(now) = monotonic_now() {
self.execution_timer.start(now);
}
}
fn effective_execution_timer_config(&self) -> Option<ExecutionTimerConfig> {
self.execution_timer_config
.or_else(fallback_execution_timer_config)
}
pub(super) fn execution_timer_tick(&mut self, work_units: u32) -> Result<()> {
if self.execution_timer.limit().is_none() {
return Ok(());
}
let Some(now) = monotonic_now() else {
return Ok(());
};
self.execution_timer
.tick(work_units, now)
.map_err(|err| match err {
LimitError::TimeLimitExceeded { elapsed, limit } => VmError::TimeLimitExceeded {
elapsed,
limit,
pc: self.pc,
},
LimitError::MemoryLimitExceeded { usage, limit } => VmError::MemoryLimitExceeded {
usage,
limit,
pc: self.pc,
},
})
}
pub(super) fn snapshot_execution_timer_on_suspend(&mut self) {
if self.execution_timer.config().is_none() {
self.execution_timer_elapsed_at_suspend = None;
return;
}
let Some(now) = monotonic_now() else {
self.execution_timer_elapsed_at_suspend = None;
return;
};
self.execution_timer_elapsed_at_suspend = self.execution_timer.elapsed(now);
}
pub(super) fn restore_execution_timer_after_resume(&mut self) {
if self.execution_timer.config().is_none() {
self.execution_timer_elapsed_at_suspend = None;
return;
}
let Some(elapsed) = self.execution_timer_elapsed_at_suspend.take() else {
return;
};
let Some(now) = monotonic_now() else {
return;
};
self.execution_timer.resume_from_elapsed(now, elapsed);
}
/// Get the current execution state of the VM
pub const fn execution_state(&self) -> &ExecutionState {
&self.execution_state

254
src/tests/interpreter/mod.rs
View File

@@ -15,11 +15,197 @@
use std::env;
use crate::test_utils::{check_output, ValueOrVec};
use crate::utils::limits::{
acquire_limits_test_lock, fallback_execution_timer_config, ExecutionTimerConfig,
};
use crate::*;
use anyhow::{bail, Result};
use core::num::NonZeroU32;
use core::time::Duration;
use serde::{Deserialize, Serialize};
use test_generator::test_resources;
use timer_test_support::{
apply_engine_timer, configure_time_source, reset_time_source, GlobalTimerGuard,
};
mod timer_test_support {
use super::{ExecutionTimerTestConfig, TimeSourceTestConfig};
#[cfg(any(test, not(feature = "std")))]
use crate::utils::limits::set_time_source;
use crate::utils::limits::{
fallback_execution_timer_config, set_fallback_execution_timer_config, ExecutionTimerConfig,
TimeSource,
};
use crate::Engine;
use anyhow::{anyhow, Result};
use core::num::NonZeroU32;
use core::time::Duration;
use std::collections::VecDeque;
use std::sync::{Mutex, Once};
use std::vec::Vec;
pub struct GlobalTimerGuard {
previous: Option<ExecutionTimerConfig>,
changed: bool,
}
impl GlobalTimerGuard {
pub fn apply(spec: Option<&ExecutionTimerTestConfig>) -> Result<Self> {
let previous = fallback_execution_timer_config();
let mut changed = false;
if let Some(config_spec) = spec {
if config_spec.disable.unwrap_or(false) {
set_fallback_execution_timer_config(None);
changed = true;
} else {
let config = config_from_spec(config_spec)?;
set_fallback_execution_timer_config(config);
changed = true;
}
}
Ok(Self { previous, changed })
}
}
impl Drop for GlobalTimerGuard {
fn drop(&mut self) {
if self.changed {
set_fallback_execution_timer_config(self.previous);
}
}
}
pub fn configure_time_source(spec: Option<&TimeSourceTestConfig>) {
ensure_time_source_registered();
let mut state = TIME_SOURCE_STATE
.lock()
.expect("time source mutex poisoned");
if let Some(cfg) = spec {
state.default_increment = cfg
.default_increment_ms
.map(Duration::from_millis)
.unwrap_or(DEFAULT_INCREMENT);
state.template_increments = cfg
.increments_ms
.iter()
.copied()
.map(Duration::from_millis)
.collect();
} else {
state.default_increment = DEFAULT_INCREMENT;
state.template_increments.clear();
}
state.reset_from_template();
}
pub fn reset_time_source() {
let mut state = TIME_SOURCE_STATE
.lock()
.expect("time source mutex poisoned");
state.reset_from_template();
}
pub fn apply_engine_timer(engine: &mut Engine, spec: &ExecutionTimerTestConfig) -> Result<()> {
if spec.disable.unwrap_or(false) {
engine.clear_execution_timer_config();
return Ok(());
}
match config_from_spec(spec)? {
Some(config) => engine.set_execution_timer_config(config),
None => engine.clear_execution_timer_config(),
}
Ok(())
}
const DEFAULT_INCREMENT: Duration = Duration::from_millis(1);
struct TestTimeSource;
struct TimeSourceState {
current: Duration,
started: bool,
default_increment: Duration,
increments: VecDeque<Duration>,
template_increments: Vec<Duration>,
}
impl TimeSourceState {
const fn new() -> Self {
Self {
current: Duration::ZERO,
started: false,
default_increment: DEFAULT_INCREMENT,
increments: VecDeque::new(),
template_increments: Vec::new(),
}
}
fn reset_from_template(&mut self) {
self.current = Duration::ZERO;
self.started = false;
self.increments = VecDeque::from(self.template_increments.clone());
}
}
impl TimeSource for TestTimeSource {
fn now(&self) -> Option<Duration> {
let mut state = TIME_SOURCE_STATE
.lock()
.expect("time source mutex poisoned");
if !state.started {
state.started = true;
return Some(state.current);
}
let increment = state
.increments
.pop_front()
.unwrap_or(state.default_increment);
state.current = state.current.saturating_add(increment);
Some(state.current)
}
}
static TEST_TIME_SOURCE: TestTimeSource = TestTimeSource;
static TIME_SOURCE_STATE: Mutex<TimeSourceState> = Mutex::new(TimeSourceState::new());
static TIME_SOURCE_ONCE: Once = Once::new();
fn ensure_time_source_registered() {
#[cfg(any(test, not(feature = "std")))]
TIME_SOURCE_ONCE.call_once(|| {
let _ = set_time_source(&TEST_TIME_SOURCE);
});
}
fn config_from_spec(spec: &ExecutionTimerTestConfig) -> Result<Option<ExecutionTimerConfig>> {
let limit_ms = match spec.limit_ms {
Some(value) => value,
None => return Ok(None),
};
let check_interval = spec
.check_interval
.map(|interval| {
NonZeroU32::new(interval)
.ok_or_else(|| anyhow!("execution_timer.check_interval must be non-zero"))
})
.transpose()? // Result<Option<NonZeroU32>>
.unwrap_or(NonZeroU32::MIN);
Ok(Some(ExecutionTimerConfig {
limit: Duration::from_millis(limit_ms),
check_interval,
}))
}
}
#[cfg(feature = "azure_policy")]
mod load_target_definitions {
@@ -158,6 +344,7 @@ fn push_query_results(query_results: QueryResults, results: &mut Vec<Value>) {
}
}
#[allow(clippy::too_many_arguments)]
pub fn eval_file(
regos: &[String],
data_opt: Option<Value>,
@@ -166,6 +353,7 @@ pub fn eval_file(
enable_tracing: bool,
strict: bool,
v0: bool,
execution_timer: Option<&ExecutionTimerTestConfig>,
) -> Result<(Vec<Value>, Vec<String>)> {
let mut engine: Engine = Engine::new();
engine.set_rego_v0(v0);
@@ -175,6 +363,15 @@ pub fn eval_file(
#[cfg(feature = "coverage")]
engine.set_enable_coverage(true);
let use_default_timer =
execution_timer.is_none() && fallback_execution_timer_config().is_none();
if let Some(spec) = execution_timer {
apply_engine_timer(&mut engine, spec)?;
} else if use_default_timer {
engine.set_execution_timer_config(default_engine_execution_timer_config());
}
let mut results = vec![];
let mut files = vec![];
@@ -199,10 +396,17 @@ pub fn eval_file(
}
let mut engine_full = engine.clone();
if let Some(spec) = execution_timer {
apply_engine_timer(&mut engine_full, spec)?;
} else if use_default_timer {
engine_full.set_execution_timer_config(default_engine_execution_timer_config());
}
if inputs.is_empty() {
// Now eval the query.
reset_time_source();
let r = engine.eval_query(query.to_string(), enable_tracing)?;
reset_time_source();
let r_full = engine_full.eval_query_and_all_rules(query.to_string(), enable_tracing)?;
if r != r_full {
std::println!(
@@ -220,7 +424,9 @@ pub fn eval_file(
engine_full.set_input(input);
// Now eval the query.
reset_time_source();
let r = engine.eval_query(query.to_string(), enable_tracing)?;
reset_time_source();
let r_full = engine_full.eval_query_and_all_rules(query.to_string(), enable_tracing)?;
if r != r_full {
std::println!(
@@ -239,6 +445,7 @@ pub fn eval_file(
}
#[cfg(feature = "azure_policy")]
#[allow(clippy::too_many_arguments)]
pub fn eval_file_with_rule_evaluation(
regos: &[String],
data_opt: Option<Value>,
@@ -247,6 +454,7 @@ pub fn eval_file_with_rule_evaluation(
_enable_tracing: bool,
strict: bool,
v0: bool,
execution_timer: Option<&ExecutionTimerTestConfig>,
) -> Result<(Vec<Value>, Vec<String>)> {
let mut engine: Engine = Engine::new();
engine.set_rego_v0(v0);
@@ -256,6 +464,15 @@ pub fn eval_file_with_rule_evaluation(
#[cfg(feature = "coverage")]
engine.set_enable_coverage(true);
let use_default_timer =
execution_timer.is_none() && fallback_execution_timer_config().is_none();
if let Some(spec) = execution_timer {
apply_engine_timer(&mut engine, spec)?;
} else if use_default_timer {
engine.set_execution_timer_config(default_engine_execution_timer_config());
}
let mut results = vec![];
let mut files = vec![];
@@ -288,7 +505,9 @@ pub fn eval_file_with_rule_evaluation(
for input in inputs {
engine.set_input(input.clone());
// Use eval_rule instead of eval_query for target tests
reset_time_source();
let r_engine = engine.eval_rule(query.to_string())?;
reset_time_source();
let r_compiled_policy = compiled_policy.eval_with_input(input)?;
assert_eq!(r_engine, r_compiled_policy);
results.push(r_engine);
@@ -297,6 +516,21 @@ pub fn eval_file_with_rule_evaluation(
Ok((results, engine.take_prints()?))
}
#[derive(Serialize, Deserialize, PartialEq, Debug, Default)]
#[serde(default)]
pub struct ExecutionTimerTestConfig {
limit_ms: Option<u64>,
check_interval: Option<u32>,
disable: Option<bool>,
}
#[derive(Serialize, Deserialize, PartialEq, Debug, Default)]
#[serde(default)]
pub struct TimeSourceTestConfig {
increments_ms: Vec<u64>,
default_increment_ms: Option<u64>,
}
#[derive(Serialize, Deserialize, PartialEq, Debug)]
struct TestCase {
data: Option<Value>,
@@ -315,18 +549,33 @@ struct TestCase {
want_error_code: Option<String>,
#[serde(default = "default_strict")]
strict: bool,
#[serde(default)]
execution_timer: Option<ExecutionTimerTestConfig>,
#[serde(default)]
global_execution_timer: Option<ExecutionTimerTestConfig>,
#[serde(default)]
time_source: Option<TimeSourceTestConfig>,
}
fn default_strict() -> bool {
true
}
fn default_engine_execution_timer_config() -> ExecutionTimerConfig {
ExecutionTimerConfig {
limit: Duration::from_secs(5),
check_interval: NonZeroU32::new(100).unwrap_or(NonZeroU32::MIN),
}
}
#[derive(Serialize, Deserialize, PartialEq, Debug)]
struct YamlTest {
cases: Vec<TestCase>,
}
fn yaml_test_impl(file: &str) -> Result<()> {
let _limits_lock = acquire_limits_test_lock();
let yaml_str = std::fs::read_to_string(file)?;
let test: YamlTest = serde_yaml::from_str(&yaml_str)?;
@@ -382,6 +631,9 @@ fn yaml_test_impl(file: &str) -> Result<()> {
continue;
}
let _timer_guard = GlobalTimerGuard::apply(case.global_execution_timer.as_ref())?;
configure_time_source(case.time_source.as_ref());
match (&case.want_result, &case.error) {
(Some(_), None) | (None, Some(_)) => (),
_ if case.no_result != Some(true) => {
@@ -405,6 +657,7 @@ fn yaml_test_impl(file: &str) -> Result<()> {
enable_tracing,
case.strict,
v0,
case.execution_timer.as_ref(),
)
}
#[cfg(not(feature = "azure_policy"))]
@@ -420,6 +673,7 @@ fn yaml_test_impl(file: &str) -> Result<()> {
enable_tracing,
case.strict,
v0,
case.execution_timer.as_ref(),
)
};

20
src/utils/limits/mod.rs
View File

@@ -8,8 +8,11 @@
mod error;
#[cfg(feature = "allocator-memory-limits")]
mod memory;
mod time;
#[allow(unused_imports)]
pub use error::LimitError;
#[allow(unused_imports)]
#[cfg(feature = "allocator-memory-limits")]
pub use memory::{
@@ -18,6 +21,19 @@ pub use memory::{
thread_memory_flush_threshold,
};
#[allow(unused_imports)]
pub use time::{
fallback_execution_timer_config, monotonic_now, set_fallback_execution_timer_config,
ExecutionTimer, ExecutionTimerConfig, TimeSource,
};
#[cfg(test)]
pub use time::acquire_limits_test_lock;
#[cfg(any(test, not(feature = "std")))]
#[allow(unused_imports)]
pub use time::{set_time_source, TimeSourceRegistrationError};
#[cfg(feature = "allocator-memory-limits")]
#[inline]
pub fn check_memory_limit_if_needed() -> core::result::Result<(), LimitError> {
@@ -26,12 +42,12 @@ pub fn check_memory_limit_if_needed() -> core::result::Result<(), LimitError> {
#[cfg(not(feature = "allocator-memory-limits"))]
#[inline]
pub fn enforce_memory_limit() -> core::result::Result<(), LimitError> {
pub const fn enforce_memory_limit() -> core::result::Result<(), LimitError> {
Ok(())
}
#[cfg(not(feature = "allocator-memory-limits"))]
#[inline]
pub fn check_memory_limit_if_needed() -> core::result::Result<(), LimitError> {
pub const fn check_memory_limit_if_needed() -> core::result::Result<(), LimitError> {
Ok(())
}

474
src/utils/limits/time.rs Normal file
View File

@@ -0,0 +1,474 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
/*
ExecutionTimer provides cooperative wall-clock enforcement for long-running
policy evaluations. The timer tracks three pieces of state:
- ExecutionTimerConfig, which holds the optional wall-clock budget and the
interval (in work units) between time checks.
- The monotonic start instant recorded via start(now), expressed as a
Duration from whatever time source the engine uses.
- An accumulator that counts work units so callers can amortize expensive
time queries; once the counter reaches the configured interval, tick()
performs a check and preserves any remainder.
The timer never calls into a clock directly. Instead, callers pass the
current monotonic Duration to start(), tick(), check_now(), or elapsed().
Helper monotonic_now() returns that Duration by selecting a TimeSource
implementation:
- On std builds we use StdTimeSource, which anchors a std::time::Instant via
OnceLock and reports elapsed() for stable, monotonic measurements.
- In tests and truly no_std builds we allow integrators to inject a global
&'static dyn TimeSource using set_time_source(). This override lives behind
a spin::Mutex<Option<...>> so the critical section stays small (just a
pointer read) while remaining usable in bare-metal environments.
With this design the interpreter can cheaply interleave work with periodic
limit checks. Std builds automatically use the Instant-backed source, while
embedded users configure both their ExecutionTimerConfig and a single global
time source without paying for per-interpreter callbacks or unsafe code.
*/
use core::num::NonZeroU32;
use core::time::Duration;
use spin::Mutex;
use super::LimitError;
#[cfg(test)]
use std::sync::{Mutex as StdMutex, MutexGuard as StdMutexGuard};
/// Public configuration for the cooperative execution time limiter.
///
/// The limiter reads this struct to determine how often it should check for wall-clock overruns and
/// what deadline to enforce. Engines without a configuration skip time checks; when a configuration
/// is present, it normally pairs a concrete deadline with a small [`NonZeroU32`] interval so
/// interpreter loops amortize their clock reads without skipping checks for long stretches of
/// repetitive work. The process-wide fallback installed via [`set_fallback_execution_timer_config`]
/// supplies this configuration when an engine lacks its own override.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ExecutionTimerConfig {
/// Maximum allowed wall-clock duration.
pub limit: Duration,
/// Number of work units between time checks (minimum 1).
pub check_interval: NonZeroU32,
}
/// Cooperative time-limit tracker shared across interpreter and VM loops.
#[derive(Debug)]
pub struct ExecutionTimer {
config: Option<ExecutionTimerConfig>,
start: Option<Duration>,
accumulated_units: u32,
last_elapsed: Duration,
}
/// Monotonic time provider.
pub trait TimeSource: Send + Sync {
/// Returns a non-decreasing duration since an arbitrary anchor.
fn now(&self) -> Option<Duration>;
}
#[cfg(feature = "std")]
#[derive(Debug)]
struct StdTimeSource;
#[cfg(feature = "std")]
impl StdTimeSource {
const fn new() -> Self {
Self
}
}
#[cfg(feature = "std")]
impl TimeSource for StdTimeSource {
fn now(&self) -> Option<Duration> {
use std::sync::OnceLock;
static ANCHOR: OnceLock<std::time::Instant> = OnceLock::new();
let anchor = ANCHOR.get_or_init(std::time::Instant::now);
Some(anchor.elapsed())
}
}
#[cfg(feature = "std")]
static STD_TIME_SOURCE: StdTimeSource = StdTimeSource::new();
#[cfg(any(test, not(feature = "std")))]
static TIME_SOURCE_OVERRIDE: Mutex<Option<&'static dyn TimeSource>> = Mutex::new(None);
#[cfg(any(test, not(feature = "std")))]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TimeSourceRegistrationError {
AlreadySet,
}
#[cfg(any(test, not(feature = "std")))]
impl core::fmt::Display for TimeSourceRegistrationError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::AlreadySet => f.write_str("time source already configured"),
}
}
}
#[cfg(any(test, not(feature = "std")))]
impl core::error::Error for TimeSourceRegistrationError {}
static FALLBACK_EXECUTION_TIMER_CONFIG: Mutex<Option<ExecutionTimerConfig>> = Mutex::new(None);
#[cfg(test)]
static LIMITS_TEST_LOCK: StdMutex<()> = StdMutex::new(());
#[cfg(test)]
pub fn acquire_limits_test_lock() -> StdMutexGuard<'static, ()> {
LIMITS_TEST_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
/// Returns the duration supplied by the chosen source for this build.
pub fn monotonic_now() -> Option<Duration> {
#[cfg(any(test, not(feature = "std")))]
// Spin mutex acquisition incurs only a few atomic ops; the critical section
// is a single pointer read, so uncontended overhead stays tiny.
if let Some(source) = {
let guard = TIME_SOURCE_OVERRIDE.lock();
*guard
} {
if let Some(duration) = source.now() {
return Some(duration);
}
}
#[cfg(feature = "std")]
{
STD_TIME_SOURCE.now()
}
#[cfg(not(feature = "std"))]
{
None
}
}
#[cfg(any(test, not(feature = "std")))]
pub fn set_time_source(source: &'static dyn TimeSource) -> Result<(), TimeSourceRegistrationError> {
let mut slot = TIME_SOURCE_OVERRIDE.lock();
if slot.is_some() {
Err(TimeSourceRegistrationError::AlreadySet)
} else {
*slot = Some(source);
Ok(())
}
}
/// Sets the process-wide fallback configuration for the execution time limiter. Engine instances can
/// override this fallback via [`Engine::set_execution_timer_config`](crate::Engine::set_execution_timer_config).
///
/// # Examples
///
/// ```
/// use std::num::NonZeroU32;
/// use std::time::Duration;
/// use regorus::utils::limits::{
/// fallback_execution_timer_config,
/// set_fallback_execution_timer_config,
/// ExecutionTimerConfig,
/// };
///
/// let config = ExecutionTimerConfig {
/// limit: Duration::from_secs(1),
/// check_interval: NonZeroU32::new(10).unwrap(),
/// };
/// set_fallback_execution_timer_config(Some(config));
/// assert_eq!(fallback_execution_timer_config(), Some(config));
/// ```
pub fn set_fallback_execution_timer_config(config: Option<ExecutionTimerConfig>) {
*FALLBACK_EXECUTION_TIMER_CONFIG.lock() = config;
}
/// Returns the process-wide fallback configuration for the execution time limiter, if any.
///
/// # Examples
///
/// ```
/// use regorus::utils::limits::fallback_execution_timer_config;
///
/// // By default no fallback execution timer is configured.
/// assert!(fallback_execution_timer_config().is_none());
/// ```
pub fn fallback_execution_timer_config() -> Option<ExecutionTimerConfig> {
let guard = FALLBACK_EXECUTION_TIMER_CONFIG.lock();
guard.as_ref().copied()
}
impl ExecutionTimer {
/// Construct a new timer with the provided configuration.
pub const fn new(config: Option<ExecutionTimerConfig>) -> Self {
Self {
config,
start: None,
accumulated_units: 0,
last_elapsed: Duration::ZERO,
}
}
/// Reset the timer state to its initial configuration without recording a start instant.
pub const fn reset(&mut self) {
self.start = None;
self.accumulated_units = 0;
self.last_elapsed = Duration::ZERO;
}
/// Reset any prior state and record the start instant.
pub const fn start(&mut self, now: Duration) {
self.start = Some(now);
self.accumulated_units = 0;
self.last_elapsed = Duration::ZERO;
}
/// Returns the timer configuration.
pub const fn config(&self) -> Option<ExecutionTimerConfig> {
self.config
}
/// Returns the configured limit.
pub const fn limit(&self) -> Option<Duration> {
match self.config {
Some(config) => Some(config.limit),
None => None,
}
}
/// Returns the last elapsed duration recorded by a check.
pub const fn last_elapsed(&self) -> Duration {
self.last_elapsed
}
/// Increment work units and run the periodic limit check when necessary.
pub fn tick(&mut self, work_units: u32, now: Duration) -> Result<(), LimitError> {
let Some(config) = self.config else {
return Ok(());
};
self.accumulated_units = self.accumulated_units.saturating_add(work_units);
if self.accumulated_units < config.check_interval.get() {
return Ok(());
}
// Preserve the remainder so that callers do not lose fractional work.
let interval = config.check_interval.get();
self.accumulated_units %= interval;
self.check_now(now)
}
/// Force an immediate check against the configured deadline.
pub fn check_now(&mut self, now: Duration) -> Result<(), LimitError> {
let Some(config) = self.config else {
return Ok(());
};
let Some(start) = self.start else {
return Ok(());
};
let elapsed = now.checked_sub(start).unwrap_or(Duration::ZERO);
self.last_elapsed = elapsed;
if elapsed > config.limit {
return Err(LimitError::TimeLimitExceeded {
elapsed,
limit: config.limit,
});
}
Ok(())
}
/// Compute elapsed time relative to the recorded start, if available.
pub fn elapsed(&self, now: Duration) -> Option<Duration> {
let start = self.start?;
Some(now.checked_sub(start).unwrap_or(Duration::ZERO))
}
/// Realign the timer start so that a previously consumed `elapsed` duration is preserved while
/// ignoring any wall-clock time that passed during a suspension window.
pub const fn resume_from_elapsed(&mut self, now: Duration, elapsed: Duration) {
if self.config.is_none() {
return;
}
self.start = Some(now.saturating_sub(elapsed));
self.last_elapsed = elapsed;
}
}
#[cfg(test)]
mod tests {
use super::*;
use core::num::NonZeroU32;
use core::sync::atomic::{AtomicU64, Ordering};
use core::time::Duration;
fn nz(value: u32) -> NonZeroU32 {
NonZeroU32::new(value).unwrap_or(NonZeroU32::MIN)
}
#[test]
fn tick_defers_checks_until_interval_is_reached() {
let mut timer = ExecutionTimer::new(Some(ExecutionTimerConfig {
limit: Duration::from_millis(100),
check_interval: nz(4),
}));
timer.start(Duration::from_millis(0));
for step in 1..4 {
let now = Duration::from_millis((step * 10) as u64);
let result = timer.tick(1, now);
assert_eq!(result, Ok(()), "tick before reaching interval must succeed");
assert_eq!(timer.last_elapsed(), Duration::ZERO);
}
let result = timer.tick(1, Duration::from_millis(40));
assert_eq!(result, Ok(()), "tick at interval boundary must succeed");
assert_eq!(timer.last_elapsed(), Duration::from_millis(40));
}
#[test]
fn check_now_reports_limit_exceeded() {
let mut timer = ExecutionTimer::new(Some(ExecutionTimerConfig {
limit: Duration::from_millis(25),
check_interval: nz(1),
}));
timer.start(Duration::from_millis(0));
assert_eq!(
timer.tick(1, Duration::from_millis(10)),
Ok(()),
"tick before limit breach must succeed"
);
let result = timer.check_now(Duration::from_millis(30));
assert!(matches!(&result, Err(LimitError::TimeLimitExceeded { .. })));
if let Err(LimitError::TimeLimitExceeded { elapsed, limit }) = result {
assert!(elapsed > limit);
assert_eq!(limit, Duration::from_millis(25));
}
}
#[test]
fn tick_reports_limit_exceeded() {
let mut timer = ExecutionTimer::new(Some(ExecutionTimerConfig {
limit: Duration::from_millis(30),
check_interval: nz(2),
}));
timer.start(Duration::from_millis(0));
assert_eq!(
timer.tick(1, Duration::from_millis(10)),
Ok(()),
"initial tick must succeed"
);
let result = timer.tick(1, Duration::from_millis(35));
assert!(matches!(&result, Err(LimitError::TimeLimitExceeded { .. })));
if let Err(LimitError::TimeLimitExceeded { elapsed, limit }) = result {
assert!(elapsed > limit);
assert_eq!(limit, Duration::from_millis(30));
assert_eq!(timer.last_elapsed(), elapsed);
}
}
#[test]
fn tick_before_start_is_noop() {
let mut timer = ExecutionTimer::new(Some(ExecutionTimerConfig {
limit: Duration::from_secs(1),
check_interval: nz(1),
}));
let result = timer.tick(1, Duration::from_millis(100));
assert_eq!(result, Ok(()), "tick before start should be ignored");
assert_eq!(timer.last_elapsed(), Duration::ZERO);
assert!(timer.elapsed(Duration::from_millis(200)).is_none());
}
#[test]
fn check_now_allows_elapsed_equal_to_limit() {
let mut timer = ExecutionTimer::new(Some(ExecutionTimerConfig {
limit: Duration::from_millis(50),
check_interval: nz(1),
}));
timer.start(Duration::from_millis(0));
assert_eq!(
timer.tick(1, Duration::from_millis(30)),
Ok(()),
"tick prior to equality check must succeed"
);
let result = timer.check_now(Duration::from_millis(50));
assert_eq!(result, Ok(()), "elapsed equal to limit must not fail");
assert_eq!(timer.last_elapsed(), Duration::from_millis(50));
}
#[test]
fn tick_is_noop_when_limit_disabled() {
let mut timer = ExecutionTimer::new(None);
timer.start(Duration::from_millis(0));
for step in 0..8 {
let now = Duration::from_millis((step + 1) as u64);
assert_eq!(
timer.tick(1, now),
Ok(()),
"ticks with disabled limit must succeed"
);
}
assert_eq!(timer.last_elapsed(), Duration::ZERO);
}
#[test]
fn check_now_is_noop_before_start() {
let mut timer = ExecutionTimer::new(None);
let result = timer.check_now(Duration::from_secs(1));
assert_eq!(result, Ok(()), "check before start must be ignored");
assert!(timer.elapsed(Duration::from_secs(2)).is_none());
}
#[test]
fn elapsed_reports_offset_from_start() {
let mut timer = ExecutionTimer::new(None);
timer.start(Duration::from_millis(5));
let elapsed = timer.elapsed(Duration::from_millis(20));
assert_eq!(elapsed, Some(Duration::from_millis(15)));
}
#[test]
fn monotonic_now_uses_override_when_present() {
static TEST_TIME: AtomicU64 = AtomicU64::new(0);
struct TestSource;
impl TimeSource for TestSource {
fn now(&self) -> Option<Duration> {
Some(Duration::from_nanos(TEST_TIME.load(Ordering::Relaxed)))
}
}
static SOURCE: TestSource = TestSource;
let _suite_guard = super::acquire_limits_test_lock();
let mut slot = super::TIME_SOURCE_OVERRIDE.lock();
let previous = (*slot).replace(&SOURCE);
drop(slot);
TEST_TIME.store(123_000_000, Ordering::Relaxed);
assert_eq!(monotonic_now(), Some(Duration::from_nanos(123_000_000)));
let mut slot = super::TIME_SOURCE_OVERRIDE.lock();
*slot = previous;
}
}