From 2f2238c51721c61d083c0ae9e5562ec30c098863 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Wed, 19 Aug 2026 19:52:45 +0200 Subject: [PATCH 1/5] feat(rvm): add per-execution memory budgets Add opt-in live-memory budgets for run-to-completion RVM evaluations with typed Rust, FFI, and C# failures. Preserve the process-global limit as a separate safeguard and reject budgeted suspendable execution. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7ba89dec-e1cf-482a-9f10-c97b107ae6ef --- benches/rvm_benchmark.rs | 32 ++- bindings/csharp/README.md | 23 ++ .../Regorus.Tests/RvmMemoryBudgetTests.cs | 86 +++++++ bindings/csharp/Regorus/MemoryBudgetConfig.cs | 46 ++++ bindings/csharp/Regorus/NativeMethods.cs | 19 ++ .../RegorusMemoryBudgetExceededException.cs | 18 ++ bindings/csharp/Regorus/Rvm.cs | 32 +++ bindings/csharp/Regorus/StatusExtensions.cs | 1 + bindings/ffi/src/common.rs | 3 + bindings/ffi/src/limits.rs | 28 ++ bindings/ffi/src/rvm.rs | 168 +++++++++++- docs/limits/memory_budget.md | 46 ++++ mimalloc/src/lib.rs | 4 +- mimalloc/src/limits.rs | 25 ++ mimalloc/src/mimalloc.rs | 6 +- src/lib.rs | 2 +- src/rvm/vm/errors.rs | 8 + src/rvm/vm/execution.rs | 30 ++- src/rvm/vm/machine.rs | 147 ++++++++++- src/rvm/vm/rules.rs | 1 + src/rvm/vm/state.rs | 33 ++- src/utils/limits/memory.rs | 12 + src/utils/limits/mod.rs | 4 +- tests/memory_limits.rs | 242 +++++++++++++++++- 24 files changed, 962 insertions(+), 54 deletions(-) create mode 100644 bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs create mode 100644 bindings/csharp/Regorus/MemoryBudgetConfig.cs create mode 100644 bindings/csharp/Regorus/RegorusMemoryBudgetExceededException.cs create mode 100644 docs/limits/memory_budget.md diff --git a/benches/rvm_benchmark.rs b/benches/rvm_benchmark.rs index 7ac093401..4cbe76892 100644 --- a/benches/rvm_benchmark.rs +++ b/benches/rvm_benchmark.rs @@ -38,6 +38,8 @@ use std::hint::black_box; use std::num::NonZeroU32; +#[cfg(feature = "allocator-memory-limits")] +use std::num::NonZeroU64; use std::path::Path; use std::sync::Arc; use std::time::Duration; @@ -50,6 +52,8 @@ use regorus::languages::rego::compiler::Compiler; use regorus::rvm::program::Program; use regorus::rvm::vm::{ExecutionMode, RegoVM}; use regorus::utils::limits::ExecutionTimerConfig; +#[cfg(feature = "allocator-memory-limits")] +use regorus::MemoryBudgetConfig; use regorus::{Engine, Rc, Value}; // --------------------------------------------------------------------------- @@ -68,28 +72,39 @@ struct EvalConfig { name: &'static str, mode: ExecutionMode, limits: bool, + memory_budget: bool, } -const EVAL_CONFIGS: [EvalConfig; 4] = [ +const EVAL_CONFIGS: [EvalConfig; 5] = [ EvalConfig { name: "regular_no_limits", mode: ExecutionMode::RunToCompletion, limits: false, + memory_budget: false, + }, + EvalConfig { + name: "regular_memory_budget", + mode: ExecutionMode::RunToCompletion, + limits: false, + memory_budget: true, }, EvalConfig { name: "regular_with_limits", mode: ExecutionMode::RunToCompletion, limits: true, + memory_budget: true, }, EvalConfig { name: "suspendable_no_limits", mode: ExecutionMode::Suspendable, limits: false, + memory_budget: false, }, EvalConfig { name: "suspendable_with_limits", mode: ExecutionMode::Suspendable, limits: true, + memory_budget: false, }, ]; @@ -358,9 +373,9 @@ fn compile_all_programs() -> Vec { // Limit helpers // --------------------------------------------------------------------------- -/// Apply or remove production-style limits based on a boolean flag. -fn configure_limits(vm: &mut RegoVM, limits: bool) { - if limits { +/// Apply the limits selected for one benchmark configuration. +fn configure_limits(vm: &mut RegoVM, config: EvalConfig) { + if config.limits { #[cfg(feature = "allocator-memory-limits")] regorus::set_global_memory_limit(Some(MEMORY_LIMIT_BYTES)); vm.set_execution_timer_config(Some(ExecutionTimerConfig { @@ -374,6 +389,11 @@ fn configure_limits(vm: &mut RegoVM, limits: bool) { vm.set_execution_timer_config(None); vm.set_max_instructions(usize::MAX); } + + #[cfg(feature = "allocator-memory-limits")] + vm.set_memory_budget_config(config.memory_budget.then(|| MemoryBudgetConfig { + limit: NonZeroU64::new(MEMORY_LIMIT_BYTES).expect("non-zero memory budget"), + })); } // --------------------------------------------------------------------------- @@ -408,7 +428,7 @@ fn bench_cold(c: &mut Criterion) { vm.set_data(black_box(d.clone())).unwrap(); } vm.set_input(black_box(input.clone())); - configure_limits(&mut vm, config.limits); + configure_limits(&mut vm, config); black_box(vm.execute().unwrap()) }) }); @@ -445,7 +465,7 @@ fn bench_hot(c: &mut Criterion) { if let Some(ref d) = data { vm.set_data(d.clone()).unwrap(); } - configure_limits(&mut vm, config.limits); + configure_limits(&mut vm, config); // Warm up: fill register window pools, caches, etc. vm.set_input(inputs[0].clone()); diff --git a/bindings/csharp/README.md b/bindings/csharp/README.md index 0d5669bf7..01f7d775e 100644 --- a/bindings/csharp/README.md +++ b/bindings/csharp/README.md @@ -105,6 +105,29 @@ var result = vm.Execute(); Console.WriteLine($"allow: {result}"); ``` +### Per-execution memory budget + +RVM run-to-completion evaluation can use an optional additional live-memory budget. Each call to `Execute` or `ExecuteEntryPoint` starts with a fresh budget. Program compilation and data/input loading are not charged. + +```csharp +using var vm = new Rvm(); +vm.LoadProgram(program); +vm.SetDataJson(Data); +vm.SetInputJson(Input); +vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(16 * 1024 * 1024)); + +try +{ + var result = vm.Execute(); +} +catch (RegorusMemoryBudgetExceededException ex) +{ + Console.WriteLine(ex.Message); +} +``` + +The budget is cooperative and may overshoot between VM checks. It is not supported in suspendable execution mode. `ClearMemoryBudgetConfig` restores the previous unlimited per-execution behavior. The process-wide limit exposed by `MemoryLimits` remains a separate safeguard. + ## Azure RBAC Condition Evaluation Evaluate Azure RBAC condition expressions directly with a JSON evaluation context: diff --git a/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs new file mode 100644 index 000000000..5c2ed456c --- /dev/null +++ b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs @@ -0,0 +1,86 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Linq; +using System.Text.Json; +using Microsoft.VisualStudio.TestTools.UnitTesting; + +namespace Regorus.Tests; + +[TestClass] +public sealed class RvmMemoryBudgetTests +{ + private const string Policy = """ +package limits.memory +import rego.v1 + +large_array := json.unmarshal(data.large_json) +"""; + + private const string EntryPoint = "data.limits.memory.large_array"; + + [TestMethod] + public void Memory_budget_must_be_non_zero() + { + Assert.ThrowsException(() => new MemoryBudgetConfig(0)); + } + + [TestMethod] + public void Execute_exceeding_memory_budget_throws_typed_exception() + { + using var program = CreateProgram(); + using var vm = CreateRvm(program); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(1)); + + Assert.ThrowsException(() => vm.Execute()); + } + + [TestMethod] + public void Clearing_memory_budget_restores_unlimited_execution() + { + using var program = CreateProgram(); + using var vm = CreateRvm(program); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(1)); + Assert.ThrowsException(() => vm.Execute()); + + vm.ClearMemoryBudgetConfig(); + + var result = vm.Execute(); + Assert.IsFalse(string.IsNullOrWhiteSpace(result)); + } + + [TestMethod] + public void Suspendable_execution_rejects_memory_budget() + { + using var vm = new Rvm(); + vm.SetExecutionMode(ExecutionMode.Suspendable); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(1024)); + + var exception = Assert.ThrowsException(() => vm.Execute()); + StringAssert.Contains(exception.Message, "not supported for suspendable execution"); + } + + private static Program CreateProgram() + { + var modules = new[] { new PolicyModule("memory_budget.rego", Policy) }; + return Program.CompileFromModules(CreateData(), modules, new[] { EntryPoint }); + } + + private static Rvm CreateRvm(Program program) + { + var vm = new Rvm(); + vm.LoadProgram(program); + vm.SetDataJson(CreateData()); + return vm; + } + + private static string CreateData() + { + var values = Enumerable.Range(0, 200_000).ToArray(); + return JsonSerializer.Serialize(new + { + large_json = JsonSerializer.Serialize(values), + }); + } +} diff --git a/bindings/csharp/Regorus/MemoryBudgetConfig.cs b/bindings/csharp/Regorus/MemoryBudgetConfig.cs new file mode 100644 index 000000000..9df08a1c6 --- /dev/null +++ b/bindings/csharp/Regorus/MemoryBudgetConfig.cs @@ -0,0 +1,46 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; + +namespace Regorus +{ + /// + /// Configures the additional live-memory budget for one RVM execution. + /// + public readonly struct MemoryBudgetConfig + { + /// + /// Initializes a new instance of the struct. + /// + /// Maximum additional live bytes allowed during one execution. + /// Thrown when is zero. + public MemoryBudgetConfig(ulong limitBytes) + { + if (limitBytes == 0) + { + throw new ArgumentOutOfRangeException(nameof(limitBytes), "Memory budget must be non-zero."); + } + + LimitBytes = limitBytes; + } + + /// + /// Gets the maximum additional live bytes allowed during one execution. + /// + public ulong LimitBytes { get; } + + internal Regorus.Internal.RegorusMemoryBudgetConfig ToNative() + { + if (LimitBytes == 0) + { + throw new InvalidOperationException("Memory budget must be non-zero."); + } + + return new Regorus.Internal.RegorusMemoryBudgetConfig + { + limit_bytes = LimitBytes, + }; + } + } +} diff --git a/bindings/csharp/Regorus/NativeMethods.cs b/bindings/csharp/Regorus/NativeMethods.cs index 327d5b15e..d46c9bdbc 100644 --- a/bindings/csharp/Regorus/NativeMethods.cs +++ b/bindings/csharp/Regorus/NativeMethods.cs @@ -245,6 +245,12 @@ internal static unsafe partial class API /// [DllImport(LibraryName, EntryPoint = "regorus_rvm_set_execution_timer_config", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] internal static extern RegorusResult regorus_rvm_set_execution_timer_config(RegorusRvm* vm, [MarshalAs(UnmanagedType.I1)] bool has_config, RegorusExecutionTimerConfig config); + + /// + /// Set memory budget configuration. + /// + [DllImport(LibraryName, EntryPoint = "regorus_rvm_set_memory_budget_config", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] + internal static extern RegorusResult regorus_rvm_set_memory_budget_config(RegorusRvm* vm, [MarshalAs(UnmanagedType.I1)] bool has_config, RegorusMemoryBudgetConfig config); /// Add a policy. /// The policy is parsed into AST. /// See https://docs.rs/regorus/latest/regorus/struct.Engine.html#method.add_policy @@ -828,6 +834,10 @@ internal enum RegorusStatus : uint /// The engine remains poisoned because a previous panic was detected. /// Poisoned, + /// + /// An RVM execution exceeded its configured memory budget. + /// + MemoryBudgetExceeded, } /// @@ -883,6 +893,15 @@ internal struct RegorusExecutionTimerConfig public uint check_interval; } + /// + /// FFI representation of the RVM memory budget configuration. + /// + [StructLayout(LayoutKind.Sequential)] + internal struct RegorusMemoryBudgetConfig + { + public ulong limit_bytes; + } + /// /// FFI representation of the policy length configuration. /// diff --git a/bindings/csharp/Regorus/RegorusMemoryBudgetExceededException.cs b/bindings/csharp/Regorus/RegorusMemoryBudgetExceededException.cs new file mode 100644 index 000000000..79a9301e0 --- /dev/null +++ b/bindings/csharp/Regorus/RegorusMemoryBudgetExceededException.cs @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; + +namespace Regorus +{ + /// + /// The exception thrown when an RVM execution exceeds its configured memory budget. + /// + public sealed class RegorusMemoryBudgetExceededException : InvalidOperationException + { + internal RegorusMemoryBudgetExceededException(string message) + : base(message) + { + } + } +} diff --git a/bindings/csharp/Regorus/Rvm.cs b/bindings/csharp/Regorus/Rvm.cs index 45c2fda3c..884731e58 100644 --- a/bindings/csharp/Regorus/Rvm.cs +++ b/bindings/csharp/Regorus/Rvm.cs @@ -144,6 +144,38 @@ public void SetExecutionMode(ExecutionMode mode) SetExecutionMode((byte)mode); } + /// + /// Configure a fresh memory budget for every run-to-completion execution. + /// + /// Memory-budget configuration. + public void SetMemoryBudgetConfig(MemoryBudgetConfig config) + { + var nativeConfig = config.ToNative(); + UseHandle(vmPtr => + { + CheckAndDropResult(API.regorus_rvm_set_memory_budget_config( + (RegorusRvm*)vmPtr, + has_config: true, + nativeConfig)); + return 0; + }); + } + + /// + /// Clear the per-execution memory budget. + /// + public void ClearMemoryBudgetConfig() + { + UseHandle(vmPtr => + { + CheckAndDropResult(API.regorus_rvm_set_memory_budget_config( + (RegorusRvm*)vmPtr, + has_config: false, + default)); + return 0; + }); + } + /// /// Execute the program and return the JSON result. /// diff --git a/bindings/csharp/Regorus/StatusExtensions.cs b/bindings/csharp/Regorus/StatusExtensions.cs index 469bce06e..1a4ae2307 100644 --- a/bindings/csharp/Regorus/StatusExtensions.cs +++ b/bindings/csharp/Regorus/StatusExtensions.cs @@ -17,6 +17,7 @@ internal static Exception CreateException(this RegorusStatus status, string? mes { RegorusStatus.Panic => new InvalidOperationException($"Regorus engine panicked: {details}"), RegorusStatus.Poisoned => new InvalidOperationException($"Regorus engine is poisoned: {details}"), + RegorusStatus.MemoryBudgetExceeded => new RegorusMemoryBudgetExceededException(details), _ => new InvalidOperationException(details), }; } diff --git a/bindings/ffi/src/common.rs b/bindings/ffi/src/common.rs index bf8cca3d0..859083f64 100644 --- a/bindings/ffi/src/common.rs +++ b/bindings/ffi/src/common.rs @@ -43,6 +43,9 @@ pub enum RegorusStatus { /// The engine remains poisoned because a previous panic was detected. Poisoned, + + /// An RVM execution exceeded its configured memory budget. + MemoryBudgetExceeded, } /// Type of data contained in RegorusResult diff --git a/bindings/ffi/src/limits.rs b/bindings/ffi/src/limits.rs index dc98a4ee8..e684f5781 100644 --- a/bindings/ffi/src/limits.rs +++ b/bindings/ffi/src/limits.rs @@ -4,6 +4,8 @@ use crate::common::{to_regorus_result, RegorusResult, RegorusStatus}; use alloc::format; use anyhow::{anyhow, Result}; +#[cfg(feature = "allocator-memory-limits")] +use core::num::NonZeroU64; use core::num::{NonZeroU32, NonZeroUsize}; use core::time::Duration; use regorus::utils::limits::{self, ExecutionTimerConfig}; @@ -145,6 +147,23 @@ pub struct RegorusExecutionTimerConfig { pub check_interval: u32, } +/// FFI representation of [`regorus::MemoryBudgetConfig`]. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct RegorusMemoryBudgetConfig { + /// Maximum additional live bytes allowed during one execution. + pub limit_bytes: u64, +} + +#[cfg(feature = "allocator-memory-limits")] +impl RegorusMemoryBudgetConfig { + pub fn to_memory_budget_config(self) -> Result { + let limit = NonZeroU64::new(self.limit_bytes) + .ok_or_else(|| anyhow!("memory_budget.limit_bytes must be non-zero"))?; + Ok(regorus::MemoryBudgetConfig { limit }) + } +} + impl RegorusExecutionTimerConfig { pub fn to_execution_timer_config(self) -> Result { let check_interval = NonZeroU32::new(self.check_interval) @@ -235,6 +254,8 @@ pub extern "C" fn regorus_clear_cache() -> RegorusResult { #[cfg(test)] mod tests { + #[cfg(feature = "allocator-memory-limits")] + use super::RegorusMemoryBudgetConfig; use super::{ optional_u64_to_result, regorus_get_global_memory_limit, regorus_set_global_memory_limit, }; @@ -256,6 +277,13 @@ mod tests { assert_eq!(result.int_value, 0); } + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn memory_budget_must_be_non_zero() { + let config = RegorusMemoryBudgetConfig { limit_bytes: 0 }; + assert!(config.to_memory_budget_config().is_err()); + } + #[test] fn ffi_roundtrips_global_limit() { let limit = 456_u64; diff --git a/bindings/ffi/src/rvm.rs b/bindings/ffi/src/rvm.rs index 7603305c1..63c028c78 100644 --- a/bindings/ffi/src/rvm.rs +++ b/bindings/ffi/src/rvm.rs @@ -7,7 +7,7 @@ use crate::common::{ }; use crate::compile::RegorusPolicyModule; use crate::compiled_policy::RegorusCompiledPolicy; -use crate::limits::RegorusExecutionTimerConfig; +use crate::limits::{RegorusExecutionTimerConfig, RegorusMemoryBudgetConfig}; use crate::lock::{new_handle, try_read, try_write, Handle, ReadGuard, WriteGuard}; use crate::panic_guard::with_unwind_guard; use alloc::boxed::Box; @@ -23,7 +23,7 @@ use regorus::rvm::program::{ generate_assembly_listing, generate_tabular_assembly_listing, AssemblyListingConfig, DeserializationResult, Program, }; -use regorus::rvm::vm::{ExecutionMode, ExecutionState, RegoVM}; +use regorus::rvm::vm::{ExecutionMode, ExecutionState, RegoVM, VmError}; use regorus::PolicyModule; use regorus::Value; @@ -56,6 +56,19 @@ impl RegorusRvm { } } +fn to_rvm_string_result(output: Result) -> RegorusResult { + match output { + Ok(json) => RegorusResult::ok_string(json), + Err(err) => { + let status = match err.downcast_ref::() { + Some(VmError::MemoryBudgetExceeded { .. }) => RegorusStatus::MemoryBudgetExceeded, + _ => RegorusStatus::Error, + }; + RegorusResult::err_with_message(status, err.to_string()) + } + } +} + /// Drop a `RegorusProgram`. #[no_mangle] pub extern "C" fn regorus_program_drop(program: *mut RegorusProgram) { @@ -478,6 +491,53 @@ pub extern "C" fn regorus_rvm_set_execution_timer_config( }) } +/// Configure the per-VM memory budget for run-to-completion execution. +#[cfg(feature = "allocator-memory-limits")] +#[no_mangle] +pub extern "C" fn regorus_rvm_set_memory_budget_config( + vm: *mut RegorusRvm, + has_config: bool, + config: RegorusMemoryBudgetConfig, +) -> RegorusResult { + with_unwind_guard(|| { + let config = if has_config { + match config.to_memory_budget_config() { + Ok(config) => Some(config), + Err(err) => { + return RegorusResult::err_with_message( + RegorusStatus::InvalidArgument, + err.to_string(), + ) + } + } + } else { + None + }; + + to_regorus_result(|| -> Result<()> { + let vm = to_shared_ref(vm as *const RegorusRvm)?; + let mut guard = vm.try_write()?; + guard.set_memory_budget_config(config); + Ok(()) + }()) + }) +} + +/// Report that memory budgets are unavailable without allocator tracking. +#[cfg(not(feature = "allocator-memory-limits"))] +#[no_mangle] +pub extern "C" fn regorus_rvm_set_memory_budget_config( + _vm: *mut RegorusRvm, + _has_config: bool, + _config: RegorusMemoryBudgetConfig, +) -> RegorusResult { + RegorusResult::err_with_message( + RegorusStatus::InvalidArgument, + "regorus_rvm_set_memory_budget_config unavailable: regorus built without allocator-memory-limits feature" + .into(), + ) +} + /// Execute the program's main entry point. #[no_mangle] pub extern "C" fn regorus_rvm_execute(vm: *mut RegorusRvm) -> RegorusResult { @@ -489,10 +549,7 @@ pub extern "C" fn regorus_rvm_execute(vm: *mut RegorusRvm) -> RegorusResult { result.to_json_str() }(); - match output { - Ok(json) => RegorusResult::ok_string(json), - Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), - } + to_rvm_string_result(output) }) } @@ -511,10 +568,7 @@ pub extern "C" fn regorus_rvm_execute_entry_point_by_name( result.to_json_str() }(); - match output { - Ok(json) => RegorusResult::ok_string(json), - Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), - } + to_rvm_string_result(output) }) } @@ -532,10 +586,7 @@ pub extern "C" fn regorus_rvm_execute_entry_point_by_index( result.to_json_str() }(); - match output { - Ok(json) => RegorusResult::ok_string(json), - Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), - } + to_rvm_string_result(output) }) } @@ -584,6 +635,95 @@ pub extern "C" fn regorus_rvm_get_execution_state(vm: *mut RegorusRvm) -> Regoru }) } +#[cfg(all(test, feature = "allocator-memory-limits"))] +mod tests { + use super::{ + regorus_rvm_drop, regorus_rvm_execute, regorus_rvm_new, + regorus_rvm_set_memory_budget_config, RegorusRvm, + }; + use crate::common::{regorus_result_drop, RegorusStatus}; + use crate::limits::RegorusMemoryBudgetConfig; + use alloc::boxed::Box; + use core::num::NonZeroU64; + use regorus::languages::rego::compiler::Compiler; + use regorus::rvm::vm::RegoVM; + use regorus::{Engine, MemoryBudgetConfig, Rc, Value}; + + const POLICY: &str = r#" +package limits.memory +import rego.v1 + +copy := [value | some value in input] +"#; + + #[test] + fn ffi_memory_budget_setter_validates_and_clears_configuration() { + let vm = regorus_rvm_new(); + + let result = regorus_rvm_set_memory_budget_config( + vm, + true, + RegorusMemoryBudgetConfig { limit_bytes: 0 }, + ); + assert!(matches!(result.status, RegorusStatus::InvalidArgument)); + regorus_result_drop(result); + + let result = regorus_rvm_set_memory_budget_config( + vm, + true, + RegorusMemoryBudgetConfig { limit_bytes: 1024 }, + ); + assert!(matches!(result.status, RegorusStatus::Ok)); + regorus_result_drop(result); + + let result = regorus_rvm_set_memory_budget_config( + vm, + false, + RegorusMemoryBudgetConfig { limit_bytes: 0 }, + ); + assert!(matches!(result.status, RegorusStatus::Ok)); + regorus_result_drop(result); + + regorus_rvm_drop(vm); + } + + #[test] + fn ffi_execution_reports_memory_budget_status() { + let entrypoint = Rc::from("data.limits.memory.copy"); + let mut engine = Engine::new(); + engine + .add_policy("memory_budget.rego".into(), POLICY.into()) + .expect("add policy"); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy"); + let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) + .expect("compile VM program"); + + let mut vm = RegoVM::new(); + vm.load_program(program); + vm.set_input( + Value::from_json_str(&format!( + "[{}]", + (0..50_000) + .map(|value| value.to_string()) + .collect::>() + .join(",") + )) + .expect("parse input"), + ); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1).expect("non-zero budget"), + })); + + let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); + let result = regorus_rvm_execute(vm); + assert!(matches!(result.status, RegorusStatus::MemoryBudgetExceeded)); + regorus_result_drop(result); + regorus_rvm_drop(vm); + } +} + fn convert_c_entry_points( entry_points: *const *const c_char, entry_points_len: usize, diff --git a/docs/limits/memory_budget.md b/docs/limits/memory_budget.md new file mode 100644 index 000000000..e3d71f83b --- /dev/null +++ b/docs/limits/memory_budget.md @@ -0,0 +1,46 @@ +# RVM memory budgets + +RVM run-to-completion execution supports an optional memory budget when Regorus is built with the `allocator-memory-limits` feature. + +The budget limits additional live bytes on the execution thread. Regorus captures a baseline when execution starts and compares later live-byte samples with that baseline. Every call to `execute`, `execute_entry_point_by_name`, or `execute_entry_point_by_index` starts with a fresh budget. + +```rust +use core::num::NonZeroU64; +use regorus::rvm::vm::RegoVM; +use regorus::MemoryBudgetConfig; + +let mut vm = RegoVM::new(); +vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(16 * 1024 * 1024).expect("non-zero budget"), +})); +``` + +No configured budget preserves existing RVM behavior. A zero-byte budget is not representable in Rust and is rejected by language bindings. + +## Included work + +The budget starts when RVM execution begins. Allocations retained by rule evaluation and its result count against the budget. + +Program compilation, program loading, data loading, input loading, and context loading happen before the execution baseline and are not charged. + +## Enforcement + +Regorus checks the budget cooperatively during VM dispatch and once before returning a successful result. Enforcement can overshoot between checks. A short-lived allocation created and freed inside one instruction may not be observed. + +Exhaustion returns `VmError::MemoryBudgetExceeded`, including: + +- additional live-byte usage observed for the evaluation +- configured budget +- VM program counter + +The C FFI reports `RegorusStatus::MemoryBudgetExceeded`. The C# binding throws `RegorusMemoryBudgetExceededException`. + +## Execution modes + +The first implementation supports run-to-completion execution only. Configuring a budget and executing in suspendable mode returns `VmError::MemoryBudgetUnsupportedInSuspendableExecution`. + +Suspendable execution may resume on another thread. A thread-local baseline cannot safely span that migration without evaluation-owned allocation attribution. + +## Process-global limit + +The existing process-global memory limit remains separate. It protects the process as a whole and is not an isolation mechanism for individual evaluations. When both controls are configured, the per-evaluation budget is checked first. diff --git a/mimalloc/src/lib.rs b/mimalloc/src/lib.rs index b12204254..9b2479624 100644 --- a/mimalloc/src/lib.rs +++ b/mimalloc/src/lib.rs @@ -7,8 +7,8 @@ pub mod mimalloc; #[cfg(feature = "allocator-memory-limits")] #[cfg(not(any(target_family = "wasm", miri)))] pub use mimalloc::{ - allocation_stats_snapshot, current_thread_allocation_stats, global_allocation_stats_snapshot, - GlobalAllocationStats, ThreadAllocationStats, + allocation_stats_snapshot, current_thread_allocation_stats, current_thread_live_bytes, + global_allocation_stats_snapshot, GlobalAllocationStats, ThreadAllocationStats, }; #[cfg(feature = "allocator-memory-limits")] diff --git a/mimalloc/src/limits.rs b/mimalloc/src/limits.rs index 8d5d98ec9..df1588dff 100644 --- a/mimalloc/src/limits.rs +++ b/mimalloc/src/limits.rs @@ -211,6 +211,11 @@ pub fn current_thread_allocation_stats() -> ThreadAllocationStats { allocation_stats_snapshot().1 } +/// Return the current thread's live-byte count without publishing counters globally. +pub fn current_thread_live_bytes() -> i64 { + THREAD_COUNTERS.with(|counters| counters.allocated.get()) +} + /// Return the unflushed allocation delta for the current thread. pub fn thread_allocation_pending_delta() -> i64 { THREAD_COUNTERS.with(|counters| counters.pending_delta()) @@ -265,3 +270,23 @@ pub fn thread_flush_threshold() -> Option { let value = THREAD_FLUSH_THRESHOLD.load(Ordering::Relaxed); (value > 0).then_some(value as u64) } + +#[cfg(test)] +mod tests { + use super::{current_thread_live_bytes, record_alloc, record_free}; + + #[test] + fn current_thread_live_bytes_tracks_usage_without_a_snapshot() { + const SIZE: usize = 4096; + let before = current_thread_live_bytes(); + + record_alloc(SIZE); + assert_eq!( + current_thread_live_bytes(), + before.saturating_add(SIZE as i64) + ); + + record_free(SIZE); + assert_eq!(current_thread_live_bytes(), before); + } +} diff --git a/mimalloc/src/mimalloc.rs b/mimalloc/src/mimalloc.rs index 74f42e441..b6cda90ac 100644 --- a/mimalloc/src/mimalloc.rs +++ b/mimalloc/src/mimalloc.rs @@ -9,9 +9,9 @@ use mimalloc_sys::{ #[cfg(feature = "allocator-memory-limits")] pub use crate::limits::{ - allocation_stats_snapshot, current_thread_allocation_stats, flush_thread_counters, - global_allocation_stats_snapshot, record_alloc, record_free, set_thread_flush_threshold, - GlobalAllocationStats, ThreadAllocationStats, + allocation_stats_snapshot, current_thread_allocation_stats, current_thread_live_bytes, + flush_thread_counters, global_allocation_stats_snapshot, record_alloc, record_free, + set_thread_flush_threshold, GlobalAllocationStats, ThreadAllocationStats, }; pub struct Mimalloc; diff --git a/src/lib.rs b/src/lib.rs index 953d8bbe1..aabe7e48f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -175,7 +175,7 @@ pub use utils::limits::PolicyLengthConfig; pub use utils::limits::{ check_global_memory_limit, enforce_memory_limit, flush_thread_memory_counters, global_memory_limit, set_global_memory_limit, set_thread_flush_threshold_override, - thread_memory_flush_threshold, + thread_memory_flush_threshold, MemoryBudgetConfig, }; pub use value::Value; diff --git a/src/rvm/vm/errors.rs b/src/rvm/vm/errors.rs index abe7beb6f..cdbd753f7 100644 --- a/src/rvm/vm/errors.rs +++ b/src/rvm/vm/errors.rs @@ -28,6 +28,14 @@ pub enum VmError { #[error("Execution exceeded memory limit (usage={usage} bytes, limit={limit} bytes, pc={pc})")] MemoryLimitExceeded { usage: u64, limit: u64, pc: usize }, + #[error( + "Execution exceeded memory budget (usage={usage} bytes, budget={budget} bytes, pc={pc})" + )] + MemoryBudgetExceeded { usage: u64, budget: u64, pc: usize }, + + #[error("Memory budgets are not supported for suspendable execution (pc={pc})")] + MemoryBudgetUnsupportedInSuspendableExecution { pc: usize }, + #[error("Compiled regex exceeded size limit ({limit} bytes, pc={pc})")] RegexSizeLimitExceeded { limit: usize, pc: usize }, diff --git a/src/rvm/vm/execution.rs b/src/rvm/vm/execution.rs index 22ff6def1..0b50bdbac 100644 --- a/src/rvm/vm/execution.rs +++ b/src/rvm/vm/execution.rs @@ -16,6 +16,7 @@ use super::machine::RegoVM; impl RegoVM { pub fn execute(&mut self) -> Result { + self.ensure_memory_budget_execution_mode()?; match self.execution_mode { ExecutionMode::RunToCompletion => self.execute_run_to_completion(), ExecutionMode::Suspendable => self.execute_suspendable(), @@ -23,6 +24,7 @@ impl RegoVM { } pub fn execute_entry_point_by_index(&mut self, index: usize) -> Result { + self.ensure_memory_budget_execution_mode()?; let (entry_point_name, entry_point_pc) = { let (name, &pc) = self.program.entry_points.get_index(index).ok_or( VmError::InvalidEntryPointIndex { @@ -44,7 +46,7 @@ impl RegoVM { match self.execution_mode { ExecutionMode::RunToCompletion => { - self.reset_execution_state(); + self.reset_run_to_completion_state(); self.reset_execution_timer_state(); self.validate_vm_state()?; @@ -56,7 +58,11 @@ impl RegoVM { } })?; - self.jump_to(entry_point_pc_u32) + let result = self + .jump_to(entry_point_pc_u32) + .map_err(|err| self.apply_memory_budget_precedence(err))?; + self.check_memory_budget_now()?; + Ok(result) } ExecutionMode::Suspendable => { self.reset_execution_state(); @@ -69,6 +75,7 @@ impl RegoVM { } pub fn execute_entry_point_by_name(&mut self, name: &str) -> Result { + self.ensure_memory_budget_execution_mode()?; let entry_point_pc = self.program .get_entry_point(name) @@ -88,7 +95,7 @@ impl RegoVM { match self.execution_mode { ExecutionMode::RunToCompletion => { - self.reset_execution_state(); + self.reset_run_to_completion_state(); self.reset_execution_timer_state(); self.validate_vm_state()?; @@ -100,7 +107,11 @@ impl RegoVM { } })?; - self.jump_to(entry_point_pc_u32) + let result = self + .jump_to(entry_point_pc_u32) + .map_err(|err| self.apply_memory_budget_precedence(err))?; + self.check_memory_budget_now()?; + Ok(result) } ExecutionMode::Suspendable => { self.reset_execution_state(); @@ -163,10 +174,17 @@ impl RegoVM { } fn execute_run_to_completion(&mut self) -> Result { - self.reset_execution_state(); + self.reset_run_to_completion_state(); self.reset_execution_timer_state(); self.execution_state = ExecutionState::Running; - match self.jump_to(0_u32) { + let result = self + .jump_to(0_u32) + .map_err(|err| self.apply_memory_budget_precedence(err)) + .and_then(|value| { + self.check_memory_budget_now()?; + Ok(value) + }); + match result { Ok(value) => { self.execution_state = ExecutionState::Completed { result: value.clone(), diff --git a/src/rvm/vm/machine.rs b/src/rvm/vm/machine.rs index 48a078607..2ffc7d9ca 100644 --- a/src/rvm/vm/machine.rs +++ b/src/rvm/vm/machine.rs @@ -4,6 +4,8 @@ use crate::rvm::program::Program; #[cfg(all(feature = "allocator-memory-limits", not(miri)))] use crate::utils::limits; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +use crate::utils::limits::MemoryBudgetConfig; use crate::utils::limits::{ fallback_execution_timer_config, monotonic_now, ExecutionTimer, ExecutionTimerConfig, LimitError, @@ -25,6 +27,9 @@ use super::execution_model::{ BreakpointSet, ExecutionMode, ExecutionStack, ExecutionState, SuspendReason, }; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +const MEMORY_BUDGET_CHECK_STRIDE: u32 = 16; + /// The Rego Virtual Machine #[derive(Debug)] pub struct RegoVM { @@ -131,6 +136,18 @@ pub struct RegoVM { /// Elapsed wall-clock time recorded when the VM entered a suspended state pub(super) execution_timer_elapsed_at_suspend: Option, + /// Optional additional live-memory budget for each run-to-completion execution + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) memory_budget_config: Option, + + /// Current-thread live-byte baseline captured at execution start + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) memory_budget_baseline: i64, + + /// Number of dispatch checks since the last memory-budget sample + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) memory_budget_ticks: u32, + /// Cached dummy span for builtin calls (avoids Source::from_contents per call) pub(super) dummy_span: Option, @@ -197,6 +214,12 @@ impl RegoVM { execution_timer_config: None, execution_timer: ExecutionTimer::new(fallback_timer), execution_timer_elapsed_at_suspend: None, + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + memory_budget_config: None, + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + memory_budget_baseline: 0, + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + memory_budget_ticks: 0, dummy_span: None, dummy_exprs: Vec::new(), cached_builtin_args: Vec::new(), @@ -402,6 +425,105 @@ impl RegoVM { self.execution_timer_config } + /// Configure a fresh memory budget for every run-to-completion execution. + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub fn set_memory_budget_config(&mut self, config: Option) { + self.memory_budget_config = config; + self.reset_memory_budget_state(); + } + + /// Return the configured per-execution memory budget. + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub const fn memory_budget_config(&self) -> Option { + self.memory_budget_config + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) fn reset_memory_budget_state(&mut self) { + self.memory_budget_baseline = if self.memory_budget_config.is_some() { + limits::current_thread_live_bytes() + } else { + 0 + }; + self.memory_budget_ticks = 0; + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) const fn ensure_memory_budget_execution_mode(&self) -> Result<()> { + if self.memory_budget_config.is_some() + && matches!(self.execution_mode, ExecutionMode::Suspendable) + { + return Err(VmError::MemoryBudgetUnsupportedInSuspendableExecution { pc: self.pc }); + } + + Ok(()) + } + + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + #[allow(clippy::unused_self)] + pub(super) const fn ensure_memory_budget_execution_mode(&self) -> Result<()> { + Ok(()) + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + fn check_memory_budget_if_needed(&mut self) -> Result<()> { + if self.memory_budget_config.is_none() { + self.memory_budget_ticks = 0; + return Ok(()); + } + + let next = self.memory_budget_ticks.saturating_add(1); + if next < MEMORY_BUDGET_CHECK_STRIDE { + self.memory_budget_ticks = next; + return Ok(()); + } + + self.memory_budget_ticks = 0; + self.check_memory_budget_now() + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) fn check_memory_budget_now(&self) -> Result<()> { + let Some(config) = self.memory_budget_config else { + return Ok(()); + }; + + let current = limits::current_thread_live_bytes(); + let usage = + u64::try_from(current.saturating_sub(self.memory_budget_baseline)).unwrap_or_default(); + let budget = config.limit.get(); + if usage > budget { + return Err(VmError::MemoryBudgetExceeded { + usage, + budget, + pc: self.pc, + }); + } + + Ok(()) + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) fn apply_memory_budget_precedence(&self, err: VmError) -> VmError { + if matches!(err, VmError::MemoryLimitExceeded { .. }) { + self.check_memory_budget_now().err().unwrap_or(err) + } else { + err + } + } + + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + #[allow(clippy::unused_self)] + pub(super) const fn check_memory_budget_now(&self) -> Result<()> { + Ok(()) + } + + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + #[allow(clippy::unused_self)] + pub(super) fn apply_memory_budget_precedence(&self, err: VmError) -> VmError { + err + } + pub(super) fn reset_execution_timer_state(&mut self) { let config = self.effective_execution_timer_config(); self.execution_timer = ExecutionTimer::new(config); @@ -541,17 +663,20 @@ impl RegoVM { #[cfg(all(feature = "allocator-memory-limits", not(miri)))] pub(super) fn memory_check(&mut self) -> Result<()> { - limits::check_memory_limit_if_needed().map_err(|err| match err { - LimitError::MemoryLimitExceeded { usage, limit } => VmError::MemoryLimitExceeded { - usage, - limit, - pc: self.pc, - }, - other => VmError::Internal { - message: format!("unexpected limit error: {other}"), - pc: self.pc, - }, - }) + self.check_memory_budget_if_needed()?; + limits::check_memory_limit_if_needed() + .map_err(|err| match err { + LimitError::MemoryLimitExceeded { usage, limit } => VmError::MemoryLimitExceeded { + usage, + limit, + pc: self.pc, + }, + other => VmError::Internal { + message: format!("unexpected limit error: {other}"), + pc: self.pc, + }, + }) + .map_err(|err| self.apply_memory_budget_precedence(err)) } #[cfg(any(miri, not(feature = "allocator-memory-limits")))] diff --git a/src/rvm/vm/rules.rs b/src/rvm/vm/rules.rs index 863a176e3..7921ca38b 100644 --- a/src/rvm/vm/rules.rs +++ b/src/rvm/vm/rules.rs @@ -24,6 +24,7 @@ impl RegoVM { err, VmError::TimeLimitExceeded { .. } | VmError::MemoryLimitExceeded { .. } + | VmError::MemoryBudgetExceeded { .. } | VmError::RegexSizeLimitExceeded { .. } | VmError::InstructionLimitExceeded { .. } ) diff --git a/src/rvm/vm/state.rs b/src/rvm/vm/state.rs index 28fb05b4c..15eff1ff4 100644 --- a/src/rvm/vm/state.rs +++ b/src/rvm/vm/state.rs @@ -11,30 +11,45 @@ use super::machine::RegoVM; impl RegoVM { /// Reset all execution state and return objects to pools for reuse pub(super) fn reset_execution_state(&mut self) { + self.release_previous_execution_state(); + self.initialize_execution_state(); + } + + /// Release values retained by the previous execution before capturing a new memory baseline. + pub(super) fn reset_run_to_completion_state(&mut self) { + self.release_previous_execution_state(); + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + self.reset_memory_budget_state(); + self.initialize_execution_state(); + } + + fn release_previous_execution_state(&mut self) { + self.evaluated = Value::Undefined; + self.execution_state = ExecutionState::Ready; + self.execution_stack.clear(); + self.return_to_pools(); + self.rule_cache.clear(); + self.registers.clear(); + self.builtins_cache.clear(); + self.cached_builtin_args.clear(); + } + + fn initialize_execution_state(&mut self) { // Reset basic execution state self.executed_instructions = 0; self.pc = 0; self.evaluated = Value::new_object(); self.cache_hits = 0; - // Reset suspendable execution state - self.execution_stack.clear(); self.execution_state = ExecutionState::Ready; - // Return objects to pools and clear stacks - self.return_to_pools(); - // Reset rule cache self.rule_cache = alloc::vec![(false, Value::Undefined); self.program.rule_infos.len()]; // Reset registers to clean state - self.registers.clear(); self.registers .resize(self.base_register_count, Value::Undefined); - // Builtin cache entries only live for a single execution - self.builtins_cache.clear(); - // Postcondition: every stack/cache that `reset_execution_state` touches // must be in its documented "clean" shape. This catches accidental // omissions in future edits to this function. diff --git a/src/utils/limits/memory.rs b/src/utils/limits/memory.rs index 4fe282013..148a27a50 100644 --- a/src/utils/limits/memory.rs +++ b/src/utils/limits/memory.rs @@ -3,6 +3,7 @@ use super::error::LimitError; use core::cell::Cell; +use core::num::NonZeroU64; use core::sync::atomic::{AtomicU64, Ordering}; use std::thread_local; @@ -14,6 +15,13 @@ const MEMORY_CHECK_STRIDE: u32 = 16; // catch short bursts before they exceed typical entry budgets while still amortizing the atomic. const MEMORY_CHECK_DELTA_BYTES: u64 = 32 * 1024; +/// Configuration for a fresh RVM memory budget. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MemoryBudgetConfig { + /// Maximum additional live bytes allowed during one execution. + pub limit: NonZeroU64, +} + thread_local! { // Per-thread stride counter used to amortize global memory checks. static MEMORY_CHECK_TICKS: Cell = const { Cell::new(0) }; @@ -211,3 +219,7 @@ pub fn global_memory_limit() -> Option { let limit = GLOBAL_MEMORY_LIMIT.load(Ordering::Relaxed); (limit != u64::MAX).then_some(limit) } + +pub fn current_thread_live_bytes() -> i64 { + mimalloc::current_thread_live_bytes() +} diff --git a/src/utils/limits/mod.rs b/src/utils/limits/mod.rs index 69d9c1fe3..0988b699b 100644 --- a/src/utils/limits/mod.rs +++ b/src/utils/limits/mod.rs @@ -14,12 +14,14 @@ mod time; #[allow(unused_imports)] pub use error::LimitError; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +pub(crate) use memory::current_thread_live_bytes; #[allow(unused_imports)] #[cfg(all(feature = "allocator-memory-limits", not(miri)))] pub use memory::{ check_global_memory_limit, enforce_memory_limit, flush_thread_memory_counters, global_memory_limit, set_global_memory_limit, set_thread_flush_threshold_override, - thread_memory_flush_threshold, + thread_memory_flush_threshold, MemoryBudgetConfig, }; #[allow(unused_imports)] diff --git a/tests/memory_limits.rs b/tests/memory_limits.rs index e80d05c21..5942707b6 100644 --- a/tests/memory_limits.rs +++ b/tests/memory_limits.rs @@ -1,9 +1,15 @@ #![cfg(all(feature = "mimalloc", feature = "allocator-memory-limits", not(miri)))] +#[cfg(feature = "rvm")] +use std::num::NonZeroU64; +#[cfg(feature = "rvm")] +use std::sync::{Arc, Barrier}; use std::sync::{Mutex, OnceLock}; use anyhow::Error; use mimalloc::global_allocation_stats_snapshot; +#[cfg(feature = "rvm")] +use regorus::MemoryBudgetConfig; use regorus::{set_global_memory_limit, Engine, LimitError, Value}; #[cfg(feature = "rvm")] @@ -139,7 +145,7 @@ fn vm_memory_limit_on_entry() { .expect("compile VM program"); let mut vm = RegoVM::new(); - vm.load_program(program); + vm.load_program(program.clone()); vm.set_data(engine.get_data()).expect("set data"); vm.set_input(Value::Undefined); @@ -283,3 +289,237 @@ fn add_data_conflict_is_atomic_on_allocator_build() { Value::from_json_str(r#"{ "a": { "z": 1 } }"#).expect("valid JSON") ); } + +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_is_enforced_per_execution() { + let _guard = LimitGuard::lock(); + let mut engine = new_engine_with_module(LARGE_PARSE_MODULE); + let large_data = large_json_data(200_000); + engine.add_data(large_data).expect("add large JSON data"); + + let entrypoint = Rc::from("data.limit.large_array"); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy for VM"); + let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) + .expect("compile VM program"); + + let mut vm = RegoVM::new(); + vm.load_program(program); + vm.set_data(engine.get_data()).expect("set data"); + vm.set_input(Value::Undefined); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1).expect("non-zero budget"), + })); + + match vm.execute() { + Err(VmError::MemoryBudgetExceeded { .. }) => {} + Err(other) => panic!("expected VM memory budget error, got {other}"), + Ok(value) => panic!("expected VM memory budget error, got value {value:?}"), + } +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_takes_precedence_over_global_limit() { + let mut guard = LimitGuard::lock(); + let mut engine = new_engine_with_module(SIMPLE_MODULE); + let entrypoint = Rc::from("data.limit.allow"); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy for VM"); + let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) + .expect("compile VM program"); + + let mut vm = RegoVM::new(); + vm.load_program(program.clone()); + vm.set_data(engine.get_data()).expect("set data"); + vm.set_input(Value::Undefined); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1).expect("non-zero budget"), + })); + guard.set_below_current_usage(); + + assert!(matches!( + vm.execute(), + Err(VmError::MemoryBudgetExceeded { .. }) + )); + + set_global_memory_limit(None); + let mut vm = RegoVM::new(); + vm.load_program(program.clone()); + vm.set_data(engine.get_data()).expect("set data"); + vm.set_input(Value::Undefined); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1).expect("non-zero budget"), + })); + guard.set_below_current_usage(); + + assert!(matches!( + vm.execute_entry_point_by_name("data.limit.allow"), + Err(VmError::MemoryBudgetExceeded { .. }) + )); + + set_global_memory_limit(None); + let mut vm = RegoVM::new(); + vm.load_program(program); + vm.set_data(engine.get_data()).expect("set data"); + vm.set_input(Value::Undefined); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1).expect("non-zero budget"), + })); + guard.set_below_current_usage(); + + assert!(matches!( + vm.execute_entry_point_by_index(0), + Err(VmError::MemoryBudgetExceeded { .. }) + )); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_is_fresh_for_each_execution() { + let _guard = LimitGuard::lock(); + let mut engine = new_engine_with_module(SIMPLE_MODULE); + let entrypoint = Rc::from("data.limit.allow"); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy for VM"); + let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) + .expect("compile VM program"); + + let mut vm = RegoVM::new(); + vm.load_program(program); + vm.set_data(engine.get_data()).expect("set data"); + vm.set_input(Value::Undefined); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1024 * 1024).expect("non-zero budget"), + })); + + assert_eq!(vm.execute().expect("first execution"), Value::Bool(true)); + assert_eq!(vm.execute().expect("second execution"), Value::Bool(true)); + assert_eq!( + vm.execute_entry_point_by_name("data.limit.allow") + .expect("named entry point"), + Value::Bool(true) + ); + assert_eq!( + vm.execute_entry_point_by_index(0) + .expect("indexed entry point"), + Value::Bool(true) + ); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_does_not_receive_credit_from_previous_results() { + let _guard = LimitGuard::lock(); + let mut engine = new_engine_with_module(LARGE_PARSE_MODULE); + let large_data = large_json_data(50_000); + engine.add_data(large_data).expect("add large JSON data"); + + let entrypoint = Rc::from("data.limit.large_array"); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy for VM"); + let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) + .expect("compile VM program"); + + let mut vm = RegoVM::new(); + vm.load_program(program); + vm.set_data(engine.get_data()).expect("set data"); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(256 * 1024 * 1024).expect("non-zero budget"), + })); + assert!(matches!( + vm.execute().expect("first execution"), + Value::Array(_) + )); + + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1).expect("non-zero budget"), + })); + + assert!(matches!( + vm.execute(), + Err(VmError::MemoryBudgetExceeded { .. }) + )); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_rejects_suspendable_execution() { + let _guard = LimitGuard::lock(); + let mut vm = RegoVM::new(); + vm.set_execution_mode(regorus::rvm::vm::ExecutionMode::Suspendable); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1024).expect("non-zero budget"), + })); + + match vm.execute() { + Err(VmError::MemoryBudgetUnsupportedInSuspendableExecution { .. }) => {} + Err(other) => panic!("expected unsupported memory budget error, got {other}"), + Ok(value) => panic!("expected unsupported memory budget error, got value {value:?}"), + } +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budgets_are_independent_across_threads() { + let _guard = LimitGuard::lock(); + let mut engine = new_engine_with_module(LARGE_PARSE_MODULE); + let large_data = large_json_data(50_000); + engine + .add_data(large_data.clone()) + .expect("add large JSON data"); + + let entrypoint = Rc::from("data.limit.large_array"); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy for VM"); + let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) + .expect("compile VM program"); + let barrier = Arc::new(Barrier::new(2)); + + std::thread::scope(|scope| { + let constrained_program = program.clone(); + let constrained_data = large_data.clone(); + let constrained_barrier = barrier.clone(); + let constrained = scope.spawn(move || { + let mut vm = RegoVM::new(); + vm.load_program(constrained_program); + vm.set_data(constrained_data).expect("set constrained data"); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1).expect("non-zero budget"), + })); + constrained_barrier.wait(); + vm.execute() + }); + + let relaxed_program = program.clone(); + let relaxed_barrier = barrier.clone(); + let relaxed = scope.spawn(move || { + let mut vm = RegoVM::new(); + vm.load_program(relaxed_program); + vm.set_data(large_data).expect("set relaxed data"); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(256 * 1024 * 1024).expect("non-zero budget"), + })); + relaxed_barrier.wait(); + vm.execute() + }); + + assert!(matches!( + constrained.join().expect("constrained thread"), + Err(VmError::MemoryBudgetExceeded { .. }) + )); + assert!(matches!( + relaxed + .join() + .expect("relaxed thread") + .expect("relaxed execution"), + Value::Array(_) + )); + }); +} From fbf805cf601a83cedf924fc7857a7a5488efae57 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Thu, 20 Aug 2026 15:08:43 +0200 Subject: [PATCH 2/5] fix(rvm): address memory budget review feedback Tighten run-to-completion accounting, reject budgeted resume, include FFI serialization, and expose typed binding failures. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7ba89dec-e1cf-482a-9f10-c97b107ae6ef --- CHANGELOG.md | 5 +- bindings/csharp/API.md | 21 ++ .../Regorus.Tests/RvmMemoryBudgetTests.cs | 18 +- bindings/csharp/Regorus/MemoryBudgetConfig.cs | 2 +- bindings/csharp/Regorus/NativeMethods.cs | 6 + ...RegorusMemoryBudgetUnsupportedException.cs | 18 ++ bindings/csharp/Regorus/StatusExtensions.cs | 1 + bindings/ffi/CHANGELOG.md | 4 + bindings/ffi/src/common.rs | 3 + bindings/ffi/src/limits.rs | 38 +-- bindings/ffi/src/rvm.rs | 222 +++++++++++++---- docs/limits/memory_budget.md | 10 +- src/lib.rs | 1 + src/rvm/vm/execution.rs | 8 +- src/rvm/vm/machine.rs | 139 ++++++++--- src/utils/limits/memory.rs | 1 + src/utils/limits/mod.rs | 1 + tests/memory_limits.rs | 233 ++++++++++++------ 18 files changed, 546 insertions(+), 185 deletions(-) create mode 100644 bindings/csharp/Regorus/RegorusMemoryBudgetUnsupportedException.cs diff --git a/CHANGELOG.md b/CHANGELOG.md index fc09e5b43..77d09910b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- *(rvm)* add opt-in per-execution memory budgets for run-to-completion evaluation, including typed Rust and binding errors ([#792](https://github.com/microsoft/regorus/pull/792)) + ## [0.11.0](https://github.com/microsoft/regorus/compare/regorus-v0.10.1...regorus-v0.11.0) - 2026-07-21 ### Added @@ -634,4 +638,3 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - LICENSE committed - CODE_OF_CONDUCT.md committed - Initial commit - diff --git a/bindings/csharp/API.md b/bindings/csharp/API.md index a8b0c83fa..a9d4b2406 100644 --- a/bindings/csharp/API.md +++ b/bindings/csharp/API.md @@ -81,6 +81,27 @@ The Regorus C# bindings provide a modern, thread-safe API for compiling and eval - **Thread Safety**: All operations are thread-safe without external synchronization - **Registry Management**: Centralized management of targets and schemas - **Policy Introspection**: Rich metadata about compiled policies +- **RVM Memory Budgets**: Optional per-execution live-memory limits for run-to-completion evaluation + +## RVM Memory Budgets + +`Rvm.SetMemoryBudgetConfig` configures a fresh non-zero budget for each `Execute` or `ExecuteEntryPoint` call. `Rvm.ClearMemoryBudgetConfig` restores unlimited execution. Compilation and data, input, and context loading are excluded; execution setup, evaluation, and result serialization are included. + +```csharp +using var vm = new Rvm(); +vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(16UL * 1024 * 1024)); + +try +{ + var result = vm.Execute(); +} +catch (RegorusMemoryBudgetExceededException) +{ + // The execution exceeded its configured budget. +} +``` + +Memory budgets require a native library built with allocator memory tracking and are supported only for run-to-completion execution. `RegorusMemoryBudgetUnsupportedException` is thrown if a configured budget is used to start or resume suspendable execution. Enforcement is cooperative, so one instruction can overshoot before the next checkpoint. ## Core Classes diff --git a/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs index 5c2ed456c..589286951 100644 --- a/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs +++ b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs @@ -11,6 +11,8 @@ namespace Regorus.Tests; [TestClass] public sealed class RvmMemoryBudgetTests { + private const ulong TightMemoryBudgetBytes = 64 * 1024; + private const string Policy = """ package limits.memory import rego.v1 @@ -24,6 +26,9 @@ import rego.v1 public void Memory_budget_must_be_non_zero() { Assert.ThrowsException(() => new MemoryBudgetConfig(0)); + + using var vm = new Rvm(); + Assert.ThrowsException(() => vm.SetMemoryBudgetConfig(default)); } [TestMethod] @@ -31,9 +36,9 @@ public void Execute_exceeding_memory_budget_throws_typed_exception() { using var program = CreateProgram(); using var vm = CreateRvm(program); - vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(1)); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); - Assert.ThrowsException(() => vm.Execute()); + Assert.ThrowsException(() => vm.ExecuteEntryPoint(EntryPoint)); } [TestMethod] @@ -41,12 +46,12 @@ public void Clearing_memory_budget_restores_unlimited_execution() { using var program = CreateProgram(); using var vm = CreateRvm(program); - vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(1)); - Assert.ThrowsException(() => vm.Execute()); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); + Assert.ThrowsException(() => vm.ExecuteEntryPoint(EntryPoint)); vm.ClearMemoryBudgetConfig(); - var result = vm.Execute(); + var result = vm.ExecuteEntryPoint(EntryPoint); Assert.IsFalse(string.IsNullOrWhiteSpace(result)); } @@ -57,8 +62,7 @@ public void Suspendable_execution_rejects_memory_budget() vm.SetExecutionMode(ExecutionMode.Suspendable); vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(1024)); - var exception = Assert.ThrowsException(() => vm.Execute()); - StringAssert.Contains(exception.Message, "not supported for suspendable execution"); + Assert.ThrowsException(() => vm.Execute()); } private static Program CreateProgram() diff --git a/bindings/csharp/Regorus/MemoryBudgetConfig.cs b/bindings/csharp/Regorus/MemoryBudgetConfig.cs index 9df08a1c6..a77d5ed10 100644 --- a/bindings/csharp/Regorus/MemoryBudgetConfig.cs +++ b/bindings/csharp/Regorus/MemoryBudgetConfig.cs @@ -34,7 +34,7 @@ internal Regorus.Internal.RegorusMemoryBudgetConfig ToNative() { if (LimitBytes == 0) { - throw new InvalidOperationException("Memory budget must be non-zero."); + throw new ArgumentOutOfRangeException(nameof(LimitBytes), "Memory budget must be non-zero."); } return new Regorus.Internal.RegorusMemoryBudgetConfig diff --git a/bindings/csharp/Regorus/NativeMethods.cs b/bindings/csharp/Regorus/NativeMethods.cs index d46c9bdbc..479c9b7c9 100644 --- a/bindings/csharp/Regorus/NativeMethods.cs +++ b/bindings/csharp/Regorus/NativeMethods.cs @@ -251,6 +251,8 @@ internal static unsafe partial class API /// [DllImport(LibraryName, EntryPoint = "regorus_rvm_set_memory_budget_config", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] internal static extern RegorusResult regorus_rvm_set_memory_budget_config(RegorusRvm* vm, [MarshalAs(UnmanagedType.I1)] bool has_config, RegorusMemoryBudgetConfig config); + + /// /// Add a policy. /// The policy is parsed into AST. /// See https://docs.rs/regorus/latest/regorus/struct.Engine.html#method.add_policy @@ -838,6 +840,10 @@ internal enum RegorusStatus : uint /// An RVM execution exceeded its configured memory budget. /// MemoryBudgetExceeded, + /// + /// An RVM memory budget was used with suspendable execution. + /// + MemoryBudgetUnsupportedInSuspendableExecution, } /// diff --git a/bindings/csharp/Regorus/RegorusMemoryBudgetUnsupportedException.cs b/bindings/csharp/Regorus/RegorusMemoryBudgetUnsupportedException.cs new file mode 100644 index 000000000..b1b36d816 --- /dev/null +++ b/bindings/csharp/Regorus/RegorusMemoryBudgetUnsupportedException.cs @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; + +namespace Regorus +{ + /// + /// The exception thrown when an RVM memory budget is used with suspendable execution. + /// + public sealed class RegorusMemoryBudgetUnsupportedException : InvalidOperationException + { + internal RegorusMemoryBudgetUnsupportedException(string message) + : base(message) + { + } + } +} diff --git a/bindings/csharp/Regorus/StatusExtensions.cs b/bindings/csharp/Regorus/StatusExtensions.cs index 1a4ae2307..ffbd284c8 100644 --- a/bindings/csharp/Regorus/StatusExtensions.cs +++ b/bindings/csharp/Regorus/StatusExtensions.cs @@ -18,6 +18,7 @@ internal static Exception CreateException(this RegorusStatus status, string? mes RegorusStatus.Panic => new InvalidOperationException($"Regorus engine panicked: {details}"), RegorusStatus.Poisoned => new InvalidOperationException($"Regorus engine is poisoned: {details}"), RegorusStatus.MemoryBudgetExceeded => new RegorusMemoryBudgetExceededException(details), + RegorusStatus.MemoryBudgetUnsupportedInSuspendableExecution => new RegorusMemoryBudgetUnsupportedException(details), _ => new InvalidOperationException(details), }; } diff --git a/bindings/ffi/CHANGELOG.md b/bindings/ffi/CHANGELOG.md index d9c001215..f1b156b0d 100644 --- a/bindings/ffi/CHANGELOG.md +++ b/bindings/ffi/CHANGELOG.md @@ -6,6 +6,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Add `regorus_rvm_set_memory_budget_config`, `RegorusMemoryBudgetConfig`, and appended statuses for memory-budget exhaustion and unsupported suspendable execution. + ## [0.1.0](https://github.com/microsoft/regorus/releases/tag/regorus-ffi-v0.1.0) - 2024-02-08 ### Other diff --git a/bindings/ffi/src/common.rs b/bindings/ffi/src/common.rs index 859083f64..cf0b19c42 100644 --- a/bindings/ffi/src/common.rs +++ b/bindings/ffi/src/common.rs @@ -46,6 +46,9 @@ pub enum RegorusStatus { /// An RVM execution exceeded its configured memory budget. MemoryBudgetExceeded, + + /// An RVM memory budget was used with suspendable execution. + MemoryBudgetUnsupportedInSuspendableExecution, } /// Type of data contained in RegorusResult diff --git a/bindings/ffi/src/limits.rs b/bindings/ffi/src/limits.rs index e684f5781..c69767981 100644 --- a/bindings/ffi/src/limits.rs +++ b/bindings/ffi/src/limits.rs @@ -4,13 +4,13 @@ use crate::common::{to_regorus_result, RegorusResult, RegorusStatus}; use alloc::format; use anyhow::{anyhow, Result}; -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] use core::num::NonZeroU64; use core::num::{NonZeroU32, NonZeroUsize}; use core::time::Duration; use regorus::utils::limits::{self, ExecutionTimerConfig}; -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] fn some_or_none(flag: bool, value: u64) -> Option { if flag { Some(value) @@ -19,7 +19,7 @@ fn some_or_none(flag: bool, value: u64) -> Option { } } -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] fn optional_u64_to_result(value: Option) -> RegorusResult { match value { Some(bytes) => { @@ -45,32 +45,32 @@ fn optional_u64_to_result(value: Option) -> RegorusResult { } } -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[no_mangle] pub extern "C" fn regorus_set_global_memory_limit(limit: u64, has_limit: bool) -> RegorusResult { ::regorus::set_global_memory_limit(some_or_none(has_limit, limit)); RegorusResult::ok_void() } -#[cfg(not(feature = "allocator-memory-limits"))] +#[cfg(any(not(feature = "allocator-memory-limits"), miri))] #[no_mangle] pub extern "C" fn regorus_set_global_memory_limit(_limit: u64, _has_limit: bool) -> RegorusResult { feature_disabled("regorus_set_global_memory_limit") } -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[no_mangle] pub extern "C" fn regorus_get_global_memory_limit() -> RegorusResult { optional_u64_to_result(::regorus::global_memory_limit()) } -#[cfg(not(feature = "allocator-memory-limits"))] +#[cfg(any(not(feature = "allocator-memory-limits"), miri))] #[no_mangle] pub extern "C" fn regorus_get_global_memory_limit() -> RegorusResult { feature_disabled("regorus_get_global_memory_limit") } -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[no_mangle] pub extern "C" fn regorus_check_global_memory_limit() -> RegorusResult { match ::regorus::check_global_memory_limit() { @@ -79,26 +79,26 @@ pub extern "C" fn regorus_check_global_memory_limit() -> RegorusResult { } } -#[cfg(not(feature = "allocator-memory-limits"))] +#[cfg(any(not(feature = "allocator-memory-limits"), miri))] #[no_mangle] pub extern "C" fn regorus_check_global_memory_limit() -> RegorusResult { feature_disabled("regorus_check_global_memory_limit") } -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[no_mangle] pub extern "C" fn regorus_flush_thread_memory_counters() -> RegorusResult { ::regorus::flush_thread_memory_counters(); RegorusResult::ok_void() } -#[cfg(not(feature = "allocator-memory-limits"))] +#[cfg(any(not(feature = "allocator-memory-limits"), miri))] #[no_mangle] pub extern "C" fn regorus_flush_thread_memory_counters() -> RegorusResult { feature_disabled("regorus_flush_thread_memory_counters") } -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[no_mangle] pub extern "C" fn regorus_set_thread_flush_threshold_override( bytes: u64, @@ -108,7 +108,7 @@ pub extern "C" fn regorus_set_thread_flush_threshold_override( RegorusResult::ok_void() } -#[cfg(not(feature = "allocator-memory-limits"))] +#[cfg(any(not(feature = "allocator-memory-limits"), miri))] #[no_mangle] pub extern "C" fn regorus_set_thread_flush_threshold_override( _bytes: u64, @@ -117,23 +117,23 @@ pub extern "C" fn regorus_set_thread_flush_threshold_override( feature_disabled("regorus_set_thread_flush_threshold_override") } -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[no_mangle] pub extern "C" fn regorus_get_thread_memory_flush_threshold() -> RegorusResult { optional_u64_to_result(::regorus::thread_memory_flush_threshold()) } -#[cfg(not(feature = "allocator-memory-limits"))] +#[cfg(any(not(feature = "allocator-memory-limits"), miri))] #[no_mangle] pub extern "C" fn regorus_get_thread_memory_flush_threshold() -> RegorusResult { feature_disabled("regorus_get_thread_memory_flush_threshold") } -#[cfg(not(feature = "allocator-memory-limits"))] +#[cfg(any(not(feature = "allocator-memory-limits"), miri))] fn feature_disabled(function: &str) -> RegorusResult { RegorusResult::err_with_message( RegorusStatus::InvalidArgument, - format!("{function} unavailable: regorus built without allocator-memory-limits feature"), + format!("{function} unavailable: allocator memory tracking is disabled"), ) } @@ -155,7 +155,7 @@ pub struct RegorusMemoryBudgetConfig { pub limit_bytes: u64, } -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] impl RegorusMemoryBudgetConfig { pub fn to_memory_budget_config(self) -> Result { let limit = NonZeroU64::new(self.limit_bytes) @@ -252,7 +252,7 @@ pub extern "C" fn regorus_clear_cache() -> RegorusResult { RegorusResult::ok_void() } -#[cfg(test)] +#[cfg(all(test, not(miri)))] mod tests { #[cfg(feature = "allocator-memory-limits")] use super::RegorusMemoryBudgetConfig; diff --git a/bindings/ffi/src/rvm.rs b/bindings/ffi/src/rvm.rs index 63c028c78..da21574e9 100644 --- a/bindings/ffi/src/rvm.rs +++ b/bindings/ffi/src/rvm.rs @@ -5,6 +5,8 @@ use crate::common::{ from_c_str, to_ref, to_regorus_result, to_shared_ref, RegorusBuffer, RegorusResult, RegorusStatus, }; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +use crate::common::regorus_result_drop; use crate::compile::RegorusPolicyModule; use crate::compiled_policy::RegorusCompiledPolicy; use crate::limits::{RegorusExecutionTimerConfig, RegorusMemoryBudgetConfig}; @@ -59,13 +61,44 @@ impl RegorusRvm { fn to_rvm_string_result(output: Result) -> RegorusResult { match output { Ok(json) => RegorusResult::ok_string(json), - Err(err) => { - let status = match err.downcast_ref::() { - Some(VmError::MemoryBudgetExceeded { .. }) => RegorusStatus::MemoryBudgetExceeded, - _ => RegorusStatus::Error, - }; - RegorusResult::err_with_message(status, err.to_string()) + Err(err) => to_rvm_error_result(err), + } +} + +fn to_rvm_error_result(err: anyhow::Error) -> RegorusResult { + let status = match err.downcast_ref::() { + Some(VmError::MemoryBudgetExceeded { .. }) => RegorusStatus::MemoryBudgetExceeded, + Some(VmError::MemoryBudgetUnsupportedInSuspendableExecution { .. }) => { + RegorusStatus::MemoryBudgetUnsupportedInSuspendableExecution } + _ => RegorusStatus::Error, + }; + RegorusResult::err_with_message(status, err.to_string()) +} + +fn execute_to_rvm_result(vm: *mut RegorusRvm, execute: F) -> RegorusResult +where + F: FnOnce(&mut RegoVM) -> core::result::Result, +{ + let output = || -> Result { + let vm = to_shared_ref(vm as *const RegorusRvm)?; + let mut guard = vm.try_write()?; + let value = execute(&mut guard)?; + let json = value.to_json_str()?; + let result = RegorusResult::ok_string(json); + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + if let Err(err) = guard.check_memory_budget() { + regorus_result_drop(result); + return Err(err.into()); + } + + Ok(result) + }(); + + match output { + Ok(result) => result, + Err(err) => to_rvm_error_result(err), } } @@ -492,7 +525,7 @@ pub extern "C" fn regorus_rvm_set_execution_timer_config( } /// Configure the per-VM memory budget for run-to-completion execution. -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[no_mangle] pub extern "C" fn regorus_rvm_set_memory_budget_config( vm: *mut RegorusRvm, @@ -524,7 +557,7 @@ pub extern "C" fn regorus_rvm_set_memory_budget_config( } /// Report that memory budgets are unavailable without allocator tracking. -#[cfg(not(feature = "allocator-memory-limits"))] +#[cfg(any(not(feature = "allocator-memory-limits"), miri))] #[no_mangle] pub extern "C" fn regorus_rvm_set_memory_budget_config( _vm: *mut RegorusRvm, @@ -533,7 +566,7 @@ pub extern "C" fn regorus_rvm_set_memory_budget_config( ) -> RegorusResult { RegorusResult::err_with_message( RegorusStatus::InvalidArgument, - "regorus_rvm_set_memory_budget_config unavailable: regorus built without allocator-memory-limits feature" + "regorus_rvm_set_memory_budget_config unavailable: allocator memory tracking is disabled" .into(), ) } @@ -541,16 +574,7 @@ pub extern "C" fn regorus_rvm_set_memory_budget_config( /// Execute the program's main entry point. #[no_mangle] pub extern "C" fn regorus_rvm_execute(vm: *mut RegorusRvm) -> RegorusResult { - with_unwind_guard(|| { - let output = || -> Result { - let vm = to_shared_ref(vm as *const RegorusRvm)?; - let mut guard = vm.try_write()?; - let result = guard.execute()?; - result.to_json_str() - }(); - - to_rvm_string_result(output) - }) + with_unwind_guard(|| execute_to_rvm_result(vm, RegoVM::execute)) } /// Execute a named entry point. @@ -560,15 +584,11 @@ pub extern "C" fn regorus_rvm_execute_entry_point_by_name( entry_point: *const c_char, ) -> RegorusResult { with_unwind_guard(|| { - let output = || -> Result { - let vm = to_shared_ref(vm as *const RegorusRvm)?; - let mut guard = vm.try_write()?; - let name = from_c_str(entry_point)?; - let result = guard.execute_entry_point_by_name(&name)?; - result.to_json_str() - }(); - - to_rvm_string_result(output) + let name = match from_c_str(entry_point) { + Ok(name) => name, + Err(err) => return to_rvm_error_result(err), + }; + execute_to_rvm_result(vm, |guard| guard.execute_entry_point_by_name(&name)) }) } @@ -579,14 +599,7 @@ pub extern "C" fn regorus_rvm_execute_entry_point_by_index( index: usize, ) -> RegorusResult { with_unwind_guard(|| { - let output = || -> Result { - let vm = to_shared_ref(vm as *const RegorusRvm)?; - let mut guard = vm.try_write()?; - let result = guard.execute_entry_point_by_index(index)?; - result.to_json_str() - }(); - - to_rvm_string_result(output) + execute_to_rvm_result(vm, |guard| guard.execute_entry_point_by_index(index)) }) } @@ -610,10 +623,7 @@ pub extern "C" fn regorus_rvm_resume( result.to_json_str() }(); - match output { - Ok(json) => RegorusResult::ok_string(json), - Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), - } + to_rvm_string_result(output) }) } @@ -635,18 +645,25 @@ pub extern "C" fn regorus_rvm_get_execution_state(vm: *mut RegorusRvm) -> Regoru }) } -#[cfg(all(test, feature = "allocator-memory-limits"))] +#[cfg(all(test, feature = "allocator-memory-limits", not(miri)))] mod tests { use super::{ - regorus_rvm_drop, regorus_rvm_execute, regorus_rvm_new, + regorus_rvm_drop, regorus_rvm_execute, regorus_rvm_execute_entry_point_by_index, + regorus_rvm_execute_entry_point_by_name, regorus_rvm_new, regorus_rvm_resume, regorus_rvm_set_memory_budget_config, RegorusRvm, }; use crate::common::{regorus_result_drop, RegorusStatus}; use crate::limits::RegorusMemoryBudgetConfig; use alloc::boxed::Box; - use core::num::NonZeroU64; + use alloc::ffi::CString; + use alloc::string::ToString; + use alloc::sync::Arc; + use alloc::vec; + use core::ptr; use regorus::languages::rego::compiler::Compiler; - use regorus::rvm::vm::RegoVM; + use regorus::rvm::instructions::Instruction; + use regorus::rvm::program::Program; + use regorus::rvm::vm::{ExecutionMode, RegoVM}; use regorus::{Engine, MemoryBudgetConfig, Rc, Value}; const POLICY: &str = r#" @@ -656,6 +673,57 @@ import rego.v1 copy := [value | some value in input] "#; + const TIGHT_MEMORY_BUDGET_BYTES: u64 = 64 * 1024; + + fn memory_budget(limit: u64) -> MemoryBudgetConfig { + MemoryBudgetConfig { + limit: core::num::NonZeroU64::new(limit).expect("non-zero budget"), + } + } + + fn host_await_program() -> Arc { + 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 = vec![Value::from("id"), Value::from(1)]; + program.instructions = 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 = vec![None; program.instructions.len()]; + Arc::new(program) + } + + fn preloaded_result_program() -> Arc { + let mut program = Program::new(); + program.dispatch_window_size = 1; + program.max_rule_window_size = 1; + program.entry_points.insert("main".to_string(), 0); + program.literals = vec![Value::from("x".repeat(2 * 1024 * 1024))]; + program.instructions = vec![ + Instruction::Load { + dest: 0, + literal_idx: 0, + }, + Instruction::Return { value: 0 }, + ]; + program.instruction_spans = vec![None; program.instructions.len()]; + Arc::new(program) + } + #[test] fn ffi_memory_budget_setter_validates_and_clears_configuration() { let vm = regorus_rvm_new(); @@ -712,16 +780,74 @@ copy := [value | some value in input] )) .expect("parse input"), ); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1).expect("non-zero budget"), - })); + vm.set_memory_budget_config(Some(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); - let result = regorus_rvm_execute(vm); + let result = regorus_rvm_execute_entry_point_by_index(vm, 0); assert!(matches!(result.status, RegorusStatus::MemoryBudgetExceeded)); + assert!(result.output.is_null()); + regorus_result_drop(result); + regorus_rvm_drop(vm); + } + + #[test] + fn ffi_result_serialization_is_included_in_memory_budget() { + let mut vm = RegoVM::new(); + vm.load_program(preloaded_result_program()); + vm.set_memory_budget_config(Some(memory_budget(512 * 1024))); + assert!(vm.execute().is_ok(), "core execution should fit the budget"); + + let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); + let entrypoint = CString::new("main").expect("entry point CString"); + let results = [ + regorus_rvm_execute(vm), + regorus_rvm_execute_entry_point_by_name(vm, entrypoint.as_ptr()), + regorus_rvm_execute_entry_point_by_index(vm, 0), + ]; + + for result in results { + assert!(matches!(result.status, RegorusStatus::MemoryBudgetExceeded)); + assert!(result.output.is_null()); + regorus_result_drop(result); + } + regorus_rvm_drop(vm); + } + + #[test] + fn ffi_resume_reports_unsupported_memory_budget_status() { + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(host_await_program()); + vm.execute().expect("suspend execution"); + + let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); + let set_result = regorus_rvm_set_memory_budget_config( + vm, + true, + RegorusMemoryBudgetConfig { + limit_bytes: 1024 * 1024, + }, + ); + assert!(matches!(set_result.status, RegorusStatus::Ok)); + regorus_result_drop(set_result); + + let result = regorus_rvm_resume(vm, ptr::null(), false); + assert!(matches!( + result.status, + RegorusStatus::MemoryBudgetUnsupportedInSuspendableExecution + )); regorus_result_drop(result); regorus_rvm_drop(vm); } + + #[test] + fn memory_budget_status_values_are_appended() { + assert_eq!(RegorusStatus::MemoryBudgetExceeded as u32, 10); + assert_eq!( + RegorusStatus::MemoryBudgetUnsupportedInSuspendableExecution as u32, + 11 + ); + } } fn convert_c_entry_points( diff --git a/docs/limits/memory_budget.md b/docs/limits/memory_budget.md index e3d71f83b..5acd3cd8b 100644 --- a/docs/limits/memory_budget.md +++ b/docs/limits/memory_budget.md @@ -19,13 +19,17 @@ No configured budget preserves existing RVM behavior. A zero-byte budget is not ## Included work -The budget starts when RVM execution begins. Allocations retained by rule evaluation and its result count against the budget. +The budget starts when RVM execution begins. Fresh execution-state initialization, rule evaluation, and allocations retained by the result count against the budget. Program compilation, program loading, data loading, input loading, and context loading happen before the execution baseline and are not charged. +The C FFI and C# bindings also check the budget after result JSON serialization and native string marshaling, before returning success. + ## Enforcement -Regorus checks the budget cooperatively during VM dispatch and once before returning a successful result. Enforcement can overshoot between checks. A short-lived allocation created and freed inside one instruction may not be observed. +Regorus checks the budget at every VM memory checkpoint and once before returning a successful result. This is cooperative enforcement, not an allocation-time hard cap: a single instruction or builtin can overshoot the budget by an unbounded amount before the next checkpoint. A short-lived allocation created and freed entirely inside one instruction may not be observed. + +Accounting uses the execution thread's live-byte counter rather than allocation ownership. When a sample falls below the baseline, Regorus lowers the baseline so an already-observed foreign free does not grant credit to later work. A foreign free can still offset evaluation allocations when both occur between samples because the allocator does not retain execution ownership for each allocation. The control therefore bounds observed additional live bytes on the execution thread, not memory attributed to an execution across threads. Exhaustion returns `VmError::MemoryBudgetExceeded`, including: @@ -37,7 +41,7 @@ The C FFI reports `RegorusStatus::MemoryBudgetExceeded`. The C# binding throws ` ## Execution modes -The first implementation supports run-to-completion execution only. Configuring a budget and executing in suspendable mode returns `VmError::MemoryBudgetUnsupportedInSuspendableExecution`. +The first implementation supports run-to-completion execution only. Configuring a budget and starting or resuming suspendable execution returns `VmError::MemoryBudgetUnsupportedInSuspendableExecution`. The FFI reports `RegorusStatus::MemoryBudgetUnsupportedInSuspendableExecution`, and C# throws `RegorusMemoryBudgetUnsupportedException`. Suspendable execution may resume on another thread. A thread-local baseline cannot safely span that migration without evaluation-owned allocation attribution. diff --git a/src/lib.rs b/src/lib.rs index aabe7e48f..de618e05f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -172,6 +172,7 @@ pub use policy_info::PolicyInfo; pub use utils::limits::LimitError; pub use utils::limits::PolicyLengthConfig; #[cfg(all(feature = "allocator-memory-limits", not(miri)))] +#[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] pub use utils::limits::{ check_global_memory_limit, enforce_memory_limit, flush_thread_memory_counters, global_memory_limit, set_global_memory_limit, set_thread_flush_threshold_override, diff --git a/src/rvm/vm/execution.rs b/src/rvm/vm/execution.rs index 0b50bdbac..a6b0301bc 100644 --- a/src/rvm/vm/execution.rs +++ b/src/rvm/vm/execution.rs @@ -61,7 +61,7 @@ impl RegoVM { let result = self .jump_to(entry_point_pc_u32) .map_err(|err| self.apply_memory_budget_precedence(err))?; - self.check_memory_budget_now()?; + self.check_memory_budget()?; Ok(result) } ExecutionMode::Suspendable => { @@ -110,7 +110,7 @@ impl RegoVM { let result = self .jump_to(entry_point_pc_u32) .map_err(|err| self.apply_memory_budget_precedence(err))?; - self.check_memory_budget_now()?; + self.check_memory_budget()?; Ok(result) } ExecutionMode::Suspendable => { @@ -181,7 +181,7 @@ impl RegoVM { .jump_to(0_u32) .map_err(|err| self.apply_memory_budget_precedence(err)) .and_then(|value| { - self.check_memory_budget_now()?; + self.check_memory_budget()?; Ok(value) }); match result { @@ -228,6 +228,8 @@ impl RegoVM { } pub fn resume(&mut self, resume_value: Option) -> Result { + self.ensure_memory_budget_resume_supported()?; + // Precondition is enforced below by returning `VmError::InvalidResumeState` // for any non-`Suspended` state. A `debug_assert!` here would diverge // debug vs release behavior and, when invoked via FFI, would trip the diff --git a/src/rvm/vm/machine.rs b/src/rvm/vm/machine.rs index 2ffc7d9ca..f8d79bb47 100644 --- a/src/rvm/vm/machine.rs +++ b/src/rvm/vm/machine.rs @@ -27,9 +27,6 @@ use super::execution_model::{ BreakpointSet, ExecutionMode, ExecutionStack, ExecutionState, SuspendReason, }; -#[cfg(all(feature = "allocator-memory-limits", not(miri)))] -const MEMORY_BUDGET_CHECK_STRIDE: u32 = 16; - /// The Rego Virtual Machine #[derive(Debug)] pub struct RegoVM { @@ -144,9 +141,9 @@ pub struct RegoVM { #[cfg(all(feature = "allocator-memory-limits", not(miri)))] pub(super) memory_budget_baseline: i64, - /// Number of dispatch checks since the last memory-budget sample + /// Whether the current baseline belongs to an active run-to-completion execution #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - pub(super) memory_budget_ticks: u32, + pub(super) memory_budget_active: bool, /// Cached dummy span for builtin calls (avoids Source::from_contents per call) pub(super) dummy_span: Option, @@ -219,7 +216,7 @@ impl RegoVM { #[cfg(all(feature = "allocator-memory-limits", not(miri)))] memory_budget_baseline: 0, #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - memory_budget_ticks: 0, + memory_budget_active: false, dummy_span: None, dummy_exprs: Vec::new(), cached_builtin_args: Vec::new(), @@ -427,13 +424,16 @@ impl RegoVM { /// Configure a fresh memory budget for every run-to-completion execution. #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - pub fn set_memory_budget_config(&mut self, config: Option) { + #[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] + pub const fn set_memory_budget_config(&mut self, config: Option) { self.memory_budget_config = config; - self.reset_memory_budget_state(); + self.memory_budget_baseline = 0; + self.memory_budget_active = false; } /// Return the configured per-execution memory budget. #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + #[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] pub const fn memory_budget_config(&self) -> Option { self.memory_budget_config } @@ -445,7 +445,7 @@ impl RegoVM { } else { 0 }; - self.memory_budget_ticks = 0; + self.memory_budget_active = self.memory_budget_config.is_some(); } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] @@ -459,38 +459,46 @@ impl RegoVM { Ok(()) } + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) const fn ensure_memory_budget_resume_supported(&self) -> Result<()> { + if self.memory_budget_config.is_some() { + return Err(VmError::MemoryBudgetUnsupportedInSuspendableExecution { pc: self.pc }); + } + + Ok(()) + } + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] #[allow(clippy::unused_self)] pub(super) const fn ensure_memory_budget_execution_mode(&self) -> Result<()> { Ok(()) } - #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - fn check_memory_budget_if_needed(&mut self) -> Result<()> { - if self.memory_budget_config.is_none() { - self.memory_budget_ticks = 0; - return Ok(()); - } - - let next = self.memory_budget_ticks.saturating_add(1); - if next < MEMORY_BUDGET_CHECK_STRIDE { - self.memory_budget_ticks = next; - return Ok(()); - } - - self.memory_budget_ticks = 0; - self.check_memory_budget_now() + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + #[allow(clippy::unused_self)] + pub(super) const fn ensure_memory_budget_resume_supported(&self) -> Result<()> { + Ok(()) } + /// Check the configured budget against the latest run-to-completion execution baseline. + /// + /// Bindings can call this on the execution thread after result serialization so their + /// marshaling allocations are included before returning success. #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - pub(super) fn check_memory_budget_now(&self) -> Result<()> { - let Some(config) = self.memory_budget_config else { + #[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] + pub fn check_memory_budget(&mut self) -> Result<()> { + let Some(config) = self + .memory_budget_config + .filter(|_| self.memory_budget_active) + else { return Ok(()); }; let current = limits::current_thread_live_bytes(); - let usage = - u64::try_from(current.saturating_sub(self.memory_budget_baseline)).unwrap_or_default(); + self.memory_budget_baseline = self.memory_budget_baseline.min(current); + let usage = current + .saturating_sub(self.memory_budget_baseline) + .unsigned_abs(); let budget = config.limit.get(); if usage > budget { return Err(VmError::MemoryBudgetExceeded { @@ -504,9 +512,9 @@ impl RegoVM { } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - pub(super) fn apply_memory_budget_precedence(&self, err: VmError) -> VmError { + pub(super) fn apply_memory_budget_precedence(&mut self, err: VmError) -> VmError { if matches!(err, VmError::MemoryLimitExceeded { .. }) { - self.check_memory_budget_now().err().unwrap_or(err) + self.check_memory_budget().err().unwrap_or(err) } else { err } @@ -514,13 +522,13 @@ impl RegoVM { #[cfg(any(miri, not(feature = "allocator-memory-limits")))] #[allow(clippy::unused_self)] - pub(super) const fn check_memory_budget_now(&self) -> Result<()> { + pub(super) const fn check_memory_budget(&mut self) -> Result<()> { Ok(()) } #[cfg(any(miri, not(feature = "allocator-memory-limits")))] #[allow(clippy::unused_self)] - pub(super) fn apply_memory_budget_precedence(&self, err: VmError) -> VmError { + pub(super) fn apply_memory_budget_precedence(&mut self, err: VmError) -> VmError { err } @@ -663,7 +671,7 @@ impl RegoVM { #[cfg(all(feature = "allocator-memory-limits", not(miri)))] pub(super) fn memory_check(&mut self) -> Result<()> { - self.check_memory_budget_if_needed()?; + self.check_memory_budget()?; limits::check_memory_limit_if_needed() .map_err(|err| match err { LimitError::MemoryLimitExceeded { usage, limit } => VmError::MemoryLimitExceeded { @@ -725,3 +733,68 @@ impl RegoVM { Ok(()) } } + +#[cfg(all(test, feature = "allocator-memory-limits", not(miri)))] +mod memory_budget_tests { + use super::RegoVM; + use super::VmError; + use crate::MemoryBudgetConfig; + use alloc::vec; + use core::num::NonZeroU64; + + #[test] + fn foreign_free_observed_before_allocation_does_not_grant_budget_credit() -> anyhow::Result<()> + { + const FOREIGN_ALLOCATION_BYTES: usize = 512 * 1024; + const LOCAL_ALLOCATION_BYTES: usize = 256 * 1024; + const BUDGET_BYTES: u64 = 128 * 1024; + + let foreign_allocation = + std::thread::spawn(|| vec![0_u8; FOREIGN_ALLOCATION_BYTES].into_boxed_slice()) + .join() + .map_err(|_| anyhow::anyhow!("allocation thread panicked"))?; + + let mut vm = RegoVM::new(); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(BUDGET_BYTES).unwrap_or(NonZeroU64::MIN), + })); + vm.reset_memory_budget_state(); + + drop(foreign_allocation); + vm.check_memory_budget()?; + + let local_allocation = vec![0_u8; LOCAL_ALLOCATION_BYTES]; + core::hint::black_box(&local_allocation); + + match vm.check_memory_budget() { + Err(VmError::MemoryBudgetExceeded { .. }) => Ok(()), + Err(err) => Err(anyhow::anyhow!("unexpected memory budget error: {err}")), + Ok(()) => Err(anyhow::anyhow!("expected memory budget exhaustion")), + } + } + + #[test] + fn memory_budget_error_takes_precedence_over_global_limit_error() { + const ALLOCATION_BYTES: usize = 256 * 1024; + const ALLOCATION_BYTES_U64: u64 = 256 * 1024; + const BUDGET_BYTES: u64 = 128 * 1024; + + let mut vm = RegoVM::new(); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(BUDGET_BYTES).unwrap_or(NonZeroU64::MIN), + })); + vm.reset_memory_budget_state(); + + let allocation = vec![0_u8; ALLOCATION_BYTES]; + core::hint::black_box(&allocation); + + assert!(matches!( + vm.apply_memory_budget_precedence(VmError::MemoryLimitExceeded { + usage: ALLOCATION_BYTES_U64, + limit: BUDGET_BYTES, + pc: 0, + }), + VmError::MemoryBudgetExceeded { .. } + )); + } +} diff --git a/src/utils/limits/memory.rs b/src/utils/limits/memory.rs index 148a27a50..fac97ce2d 100644 --- a/src/utils/limits/memory.rs +++ b/src/utils/limits/memory.rs @@ -17,6 +17,7 @@ const MEMORY_CHECK_DELTA_BYTES: u64 = 32 * 1024; /// Configuration for a fresh RVM memory budget. #[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] pub struct MemoryBudgetConfig { /// Maximum additional live bytes allowed during one execution. pub limit: NonZeroU64, diff --git a/src/utils/limits/mod.rs b/src/utils/limits/mod.rs index 0988b699b..c8f4e7e52 100644 --- a/src/utils/limits/mod.rs +++ b/src/utils/limits/mod.rs @@ -18,6 +18,7 @@ pub use error::LimitError; pub(crate) use memory::current_thread_live_bytes; #[allow(unused_imports)] #[cfg(all(feature = "allocator-memory-limits", not(miri)))] +#[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] pub use memory::{ check_global_memory_limit, enforce_memory_limit, flush_thread_memory_counters, global_memory_limit, set_global_memory_limit, set_thread_flush_threshold_override, diff --git a/tests/memory_limits.rs b/tests/memory_limits.rs index 5942707b6..4afa09f22 100644 --- a/tests/memory_limits.rs +++ b/tests/memory_limits.rs @@ -15,9 +15,11 @@ use regorus::{set_global_memory_limit, Engine, LimitError, Value}; #[cfg(feature = "rvm")] use regorus::languages::rego::compiler::Compiler; #[cfg(feature = "rvm")] -use regorus::rvm::vm::RegoVM; +use regorus::rvm::instructions::Instruction; #[cfg(feature = "rvm")] -use regorus::rvm::vm::VmError; +use regorus::rvm::program::Program; +#[cfg(feature = "rvm")] +use regorus::rvm::vm::{ExecutionMode, ExecutionState, RegoVM, VmError}; #[cfg(feature = "rvm")] use regorus::Rc; @@ -78,6 +80,12 @@ package limit large_array := json.unmarshal(data.limit.large_json) "#; +#[cfg(feature = "rvm")] +const TIGHT_MEMORY_BUDGET_BYTES: u64 = 64 * 1024; + +#[cfg(feature = "rvm")] +const RELAXED_MEMORY_BUDGET_BYTES: u64 = 256 * 1024 * 1024; + #[cfg(feature = "jsonpatch")] const JSON_PATCH_MODULE: &str = r#" package limit @@ -121,6 +129,40 @@ fn new_engine_with_module(module: &str) -> Engine { engine } +#[cfg(feature = "rvm")] +fn memory_budget(limit: u64) -> MemoryBudgetConfig { + MemoryBudgetConfig { + limit: NonZeroU64::new(limit).expect("non-zero budget"), + } +} + +#[cfg(feature = "rvm")] +fn host_await_program() -> Arc { + 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 = vec![Value::from("id"), Value::from(1)]; + program.instructions = 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 = vec![None; program.instructions.len()]; + Arc::new(program) +} + #[test] fn interpreter_memory_limit_on_entry() { let mut guard = LimitGuard::lock(); @@ -305,27 +347,57 @@ fn vm_memory_budget_is_enforced_per_execution() { let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) .expect("compile VM program"); + let instruction_count = program.instructions.len(); + + let mut relaxed_vm = RegoVM::new(); + relaxed_vm.load_program(program.clone()); + relaxed_vm + .set_data(engine.get_data()) + .expect("set relaxed data"); + relaxed_vm.set_input(Value::Undefined); + relaxed_vm.set_memory_budget_config(Some(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); + match relaxed_vm + .execute_entry_point_by_name(entrypoint.as_ref()) + .expect("relaxed execution") + { + Value::Array(values) => assert_eq!(values.len(), 200_000), + value => panic!("expected large array, got {value:?}"), + } + let mut vm = RegoVM::new(); vm.load_program(program); - vm.set_data(engine.get_data()).expect("set data"); + vm.set_data(engine.get_data()) + .expect("set constrained data"); vm.set_input(Value::Undefined); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1).expect("non-zero budget"), - })); + vm.set_memory_budget_config(Some(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); - match vm.execute() { - Err(VmError::MemoryBudgetExceeded { .. }) => {} + match vm.execute_entry_point_by_name(entrypoint.as_ref()) { + Err(VmError::MemoryBudgetExceeded { usage, budget, pc }) => { + assert!(usage > budget); + assert_eq!(budget, TIGHT_MEMORY_BUDGET_BYTES); + assert!(pc < instruction_count); + } Err(other) => panic!("expected VM memory budget error, got {other}"), Ok(value) => panic!("expected VM memory budget error, got value {value:?}"), } + + vm.set_memory_budget_config(None); + assert!(matches!( + vm.execute_entry_point_by_name(entrypoint.as_ref()) + .expect("execution after clearing budget"), + Value::Array(_) + )); } #[cfg(feature = "rvm")] #[test] -fn vm_memory_budget_takes_precedence_over_global_limit() { +fn vm_memory_budget_does_not_mask_global_limit() { let mut guard = LimitGuard::lock(); - let mut engine = new_engine_with_module(SIMPLE_MODULE); - let entrypoint = Rc::from("data.limit.allow"); + let mut engine = new_engine_with_module(LARGE_PARSE_MODULE); + engine + .add_data(large_json_data(200_000)) + .expect("add large JSON data"); + let entrypoint = Rc::from("data.limit.large_array"); let compiled = engine .compile_with_entrypoint(&entrypoint) .expect("compile policy for VM"); @@ -333,48 +405,46 @@ fn vm_memory_budget_takes_precedence_over_global_limit() { .expect("compile VM program"); let mut vm = RegoVM::new(); - vm.load_program(program.clone()); - vm.set_data(engine.get_data()).expect("set data"); - vm.set_input(Value::Undefined); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1).expect("non-zero budget"), - })); - guard.set_below_current_usage(); - - assert!(matches!( - vm.execute(), - Err(VmError::MemoryBudgetExceeded { .. }) - )); - - set_global_memory_limit(None); - let mut vm = RegoVM::new(); - vm.load_program(program.clone()); + vm.load_program(program); vm.set_data(engine.get_data()).expect("set data"); vm.set_input(Value::Undefined); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1).expect("non-zero budget"), - })); + vm.set_memory_budget_config(Some(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); guard.set_below_current_usage(); assert!(matches!( - vm.execute_entry_point_by_name("data.limit.allow"), - Err(VmError::MemoryBudgetExceeded { .. }) + vm.execute_entry_point_by_name(entrypoint.as_ref()), + Err(VmError::MemoryLimitExceeded { .. }) )); +} - set_global_memory_limit(None); - let mut vm = RegoVM::new(); - vm.load_program(program); - vm.set_data(engine.get_data()).expect("set data"); - vm.set_input(Value::Undefined); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1).expect("non-zero budget"), - })); - guard.set_below_current_usage(); +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_is_enforced_for_named_and_indexed_entry_points() { + let _guard = LimitGuard::lock(); + let mut engine = new_engine_with_module(LARGE_PARSE_MODULE); + engine + .add_data(large_json_data(200_000)) + .expect("add large JSON data"); + let entrypoint = Rc::from("data.limit.large_array"); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy for VM"); + let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) + .expect("compile VM program"); - assert!(matches!( - vm.execute_entry_point_by_index(0), - Err(VmError::MemoryBudgetExceeded { .. }) - )); + for execute_by_name in [true, false] { + let mut vm = RegoVM::new(); + vm.load_program(program.clone()); + vm.set_data(engine.get_data()).expect("set data"); + vm.set_memory_budget_config(Some(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); + + let result = if execute_by_name { + vm.execute_entry_point_by_name(entrypoint.as_ref()) + } else { + vm.execute_entry_point_by_index(0) + }; + assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); + } } #[cfg(feature = "rvm")] @@ -393,11 +463,13 @@ fn vm_memory_budget_is_fresh_for_each_execution() { vm.load_program(program); vm.set_data(engine.get_data()).expect("set data"); vm.set_input(Value::Undefined); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1024 * 1024).expect("non-zero budget"), - })); + vm.set_memory_budget_config(Some(memory_budget(1024 * 1024))); assert_eq!(vm.execute().expect("first execution"), Value::Bool(true)); + + let allocation_between_executions = vec![0_u8; 4 * 1024 * 1024]; + core::hint::black_box(&allocation_between_executions); + assert_eq!(vm.execute().expect("second execution"), Value::Bool(true)); assert_eq!( vm.execute_entry_point_by_name("data.limit.allow") @@ -429,17 +501,13 @@ fn vm_memory_budget_does_not_receive_credit_from_previous_results() { let mut vm = RegoVM::new(); vm.load_program(program); vm.set_data(engine.get_data()).expect("set data"); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(256 * 1024 * 1024).expect("non-zero budget"), - })); + vm.set_memory_budget_config(Some(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); assert!(matches!( vm.execute().expect("first execution"), Value::Array(_) )); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1).expect("non-zero budget"), - })); + vm.set_memory_budget_config(Some(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); assert!(matches!( vm.execute(), @@ -453,9 +521,7 @@ fn vm_memory_budget_rejects_suspendable_execution() { let _guard = LimitGuard::lock(); let mut vm = RegoVM::new(); vm.set_execution_mode(regorus::rvm::vm::ExecutionMode::Suspendable); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1024).expect("non-zero budget"), - })); + vm.set_memory_budget_config(Some(memory_budget(1024))); match vm.execute() { Err(VmError::MemoryBudgetUnsupportedInSuspendableExecution { .. }) => {} @@ -464,6 +530,36 @@ fn vm_memory_budget_rejects_suspendable_execution() { } } +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_rejects_resume_after_suspension() { + let _guard = LimitGuard::lock(); + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(host_await_program()); + + vm.execute().expect("suspend execution"); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + + vm.set_memory_budget_config(Some(memory_budget(1024 * 1024))); + vm.set_execution_mode(ExecutionMode::RunToCompletion); + + assert!(matches!( + vm.resume(Some(Value::from(42))), + Err(VmError::MemoryBudgetUnsupportedInSuspendableExecution { .. }) + )); + + vm.set_memory_budget_config(None); + assert_eq!( + vm.resume(Some(Value::from(42))) + .expect("resume after clearing budget"), + Value::from(42) + ); +} + #[cfg(feature = "rvm")] #[test] fn vm_memory_budgets_are_independent_across_threads() { @@ -490,9 +586,7 @@ fn vm_memory_budgets_are_independent_across_threads() { let mut vm = RegoVM::new(); vm.load_program(constrained_program); vm.set_data(constrained_data).expect("set constrained data"); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1).expect("non-zero budget"), - })); + vm.set_memory_budget_config(Some(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); constrained_barrier.wait(); vm.execute() }); @@ -503,9 +597,7 @@ fn vm_memory_budgets_are_independent_across_threads() { let mut vm = RegoVM::new(); vm.load_program(relaxed_program); vm.set_data(large_data).expect("set relaxed data"); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(256 * 1024 * 1024).expect("non-zero budget"), - })); + vm.set_memory_budget_config(Some(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); relaxed_barrier.wait(); vm.execute() }); @@ -514,12 +606,13 @@ fn vm_memory_budgets_are_independent_across_threads() { constrained.join().expect("constrained thread"), Err(VmError::MemoryBudgetExceeded { .. }) )); - assert!(matches!( - relaxed - .join() - .expect("relaxed thread") - .expect("relaxed execution"), - Value::Array(_) - )); + match relaxed + .join() + .expect("relaxed thread") + .expect("relaxed execution") + { + Value::Array(values) => assert_eq!(values.len(), 50_000), + value => panic!("expected relaxed array, got {value:?}"), + } }); } From e61aa28b719091e433a00e6ce5d89213fabf9c45 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Thu, 20 Aug 2026 16:33:15 +0200 Subject: [PATCH 3/5] fix(ffi): format conditional import Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7ba89dec-e1cf-482a-9f10-c97b107ae6ef --- bindings/ffi/src/rvm.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bindings/ffi/src/rvm.rs b/bindings/ffi/src/rvm.rs index da21574e9..20699c8c2 100644 --- a/bindings/ffi/src/rvm.rs +++ b/bindings/ffi/src/rvm.rs @@ -1,12 +1,12 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +use crate::common::regorus_result_drop; use crate::common::{ from_c_str, to_ref, to_regorus_result, to_shared_ref, RegorusBuffer, RegorusResult, RegorusStatus, }; -#[cfg(all(feature = "allocator-memory-limits", not(miri)))] -use crate::common::regorus_result_drop; use crate::compile::RegorusPolicyModule; use crate::compiled_policy::RegorusCompiledPolicy; use crate::limits::{RegorusExecutionTimerConfig, RegorusMemoryBudgetConfig}; From 3ac1c1a9be66e5d266e6a6e9414a4d31584ac8d1 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Mon, 24 Aug 2026 20:55:11 +0200 Subject: [PATCH 4/5] fix(rvm): include evaluation data in memory budgets Add synchronous scoped accounting across evaluation-specific data, execution, and native result production. Expose matching one-call FFI and C# APIs with terminal cleanup and reuse guarantees. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 1 + benches/rvm_benchmark.rs | 6 +- bindings/csharp/API.md | 28 +- bindings/csharp/README.md | 14 +- .../Regorus.Tests/RvmMemoryBudgetTests.cs | 189 ++++++++ bindings/csharp/Regorus/NativeMethods.cs | 18 + bindings/csharp/Regorus/Rvm.cs | 51 +++ bindings/ffi/CHANGELOG.md | 4 + bindings/ffi/src/engine.rs | 1 + bindings/ffi/src/limits.rs | 2 +- bindings/ffi/src/panic_guard.rs | 131 ++++++ bindings/ffi/src/rvm.rs | 427 ++++++++++++++++-- docs/limits/memory_budget.md | 36 +- src/rvm/vm/errors.rs | 3 + src/rvm/vm/execution.rs | 103 +++-- src/rvm/vm/machine.rs | 160 ++++++- src/rvm/vm/state.rs | 400 +++++++++++++++- tests/memory_limits.rs | 124 ++++- 18 files changed, 1590 insertions(+), 108 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 77d09910b..2b03aad49 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - *(rvm)* add opt-in per-execution memory budgets for run-to-completion evaluation, including typed Rust and binding errors ([#792](https://github.com/microsoft/regorus/pull/792)) +- *(rvm,ffi,csharp)* add scoped evaluation-memory budgeting via `RegoVM::with_evaluation_memory_budget` and one-call data execution helpers for charging evaluation-specific data with execution ([#792](https://github.com/microsoft/regorus/pull/792)) ## [0.11.0](https://github.com/microsoft/regorus/compare/regorus-v0.10.1...regorus-v0.11.0) - 2026-07-21 diff --git a/benches/rvm_benchmark.rs b/benches/rvm_benchmark.rs index 4cbe76892..cbcd1e610 100644 --- a/benches/rvm_benchmark.rs +++ b/benches/rvm_benchmark.rs @@ -38,7 +38,7 @@ use std::hint::black_box; use std::num::NonZeroU32; -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] use std::num::NonZeroU64; use std::path::Path; use std::sync::Arc; @@ -52,7 +52,7 @@ use regorus::languages::rego::compiler::Compiler; use regorus::rvm::program::Program; use regorus::rvm::vm::{ExecutionMode, RegoVM}; use regorus::utils::limits::ExecutionTimerConfig; -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] use regorus::MemoryBudgetConfig; use regorus::{Engine, Rc, Value}; @@ -390,7 +390,7 @@ fn configure_limits(vm: &mut RegoVM, config: EvalConfig) { vm.set_max_instructions(usize::MAX); } - #[cfg(feature = "allocator-memory-limits")] + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] vm.set_memory_budget_config(config.memory_budget.then(|| MemoryBudgetConfig { limit: NonZeroU64::new(MEMORY_LIMIT_BYTES).expect("non-zero memory budget"), })); diff --git a/bindings/csharp/API.md b/bindings/csharp/API.md index a9d4b2406..a7b27833f 100644 --- a/bindings/csharp/API.md +++ b/bindings/csharp/API.md @@ -85,7 +85,27 @@ The Regorus C# bindings provide a modern, thread-safe API for compiling and eval ## RVM Memory Budgets -`Rvm.SetMemoryBudgetConfig` configures a fresh non-zero budget for each `Execute` or `ExecuteEntryPoint` call. `Rvm.ClearMemoryBudgetConfig` restores unlimited execution. Compilation and data, input, and context loading are excluded; execution setup, evaluation, and result serialization are included. +`Rvm.SetMemoryBudgetConfig` configures a fresh non-zero budget for each run-to-completion execution. `Rvm.ClearMemoryBudgetConfig` restores unlimited execution. + +```csharp +public readonly struct MemoryBudgetConfig +{ + public MemoryBudgetConfig(ulong limitBytes); + public ulong LimitBytes { get; } +} + +public sealed class Rvm : IDisposable +{ + public void SetMemoryBudgetConfig(MemoryBudgetConfig config); + public void ClearMemoryBudgetConfig(); + + public string? ExecuteWithDataJson(string dataJson); + public string? ExecuteEntryPointWithDataJson(string entryPoint, string dataJson); + public string? ExecuteEntryPointWithDataJson(ulong index, string dataJson); +} +``` + +Ordinary `Execute` and `ExecuteEntryPoint` calls use the existing VM data and start a fresh budget for execution. Program compilation, program loading, and prior `SetDataJson`, `SetInputJson`, and `SetContextJson` calls are excluded; use this path for static or preloaded data. ```csharp using var vm = new Rvm(); @@ -101,7 +121,11 @@ catch (RegorusMemoryBudgetExceededException) } ``` -Memory budgets require a native library built with allocator memory tracking and are supported only for run-to-completion execution. `RegorusMemoryBudgetUnsupportedException` is thrown if a configured budget is used to start or resume suspendable execution. Enforcement is cooperative, so one instruction can overshoot before the next checkpoint. +Use `ExecuteWithDataJson`, `ExecuteEntryPointWithDataJson(string, string)`, or `ExecuteEntryPointWithDataJson(ulong, string)` when evaluation-specific data should be charged with execution. These methods make one native call that includes native JSON parsing/storage for `dataJson`, execution, native JSON serialization, and native C-string allocation. Managed UTF-8 decoding and the managed C# `string` allocation after the native call returns are not charged. + +If scoped data replacement fails, the VM keeps its previous data. The previous and provisional native data can coexist transiently, so both count toward the peak live bytes observed by the budget. + +Memory budgets require a native library built with allocator memory tracking and are supported only for run-to-completion execution. `RegorusMemoryBudgetExceededException` is thrown when a budget is exceeded. `RegorusMemoryBudgetUnsupportedException` is thrown if a configured budget is used to start or resume suspendable execution. Enforcement is cooperative, so one instruction can overshoot before the next checkpoint. Same-thread baseline ratcheting can make the effective limit stricter after unrelated frees are observed; those frees are never credited back. Public multi-call begin/end scopes are intentionally absent because allocator counters are thread-local and abandoned or cross-thread scopes would be unsafe. ## Core Classes diff --git a/bindings/csharp/README.md b/bindings/csharp/README.md index 01f7d775e..36c67def1 100644 --- a/bindings/csharp/README.md +++ b/bindings/csharp/README.md @@ -107,7 +107,7 @@ Console.WriteLine($"allow: {result}"); ### Per-execution memory budget -RVM run-to-completion evaluation can use an optional additional live-memory budget. Each call to `Execute` or `ExecuteEntryPoint` starts with a fresh budget. Program compilation and data/input loading are not charged. +RVM run-to-completion evaluation can use an optional additional live-memory budget. Each ordinary `Execute` or `ExecuteEntryPoint` call starts with a fresh budget for execution; program compilation, program loading, and prior `SetDataJson`, `SetInputJson`, and `SetContextJson` calls are not charged. ```csharp using var vm = new Rvm(); @@ -126,7 +126,17 @@ catch (RegorusMemoryBudgetExceededException ex) } ``` -The budget is cooperative and may overshoot between VM checks. It is not supported in suspendable execution mode. `ClearMemoryBudgetConfig` restores the previous unlimited per-execution behavior. The process-wide limit exposed by `MemoryLimits` remains a separate safeguard. +Use the one-call `Execute*WithDataJson` methods when evaluation-specific data should be charged in the same native budget scope as execution: + +```csharp +var result = vm.ExecuteWithDataJson(Data); +var named = vm.ExecuteEntryPointWithDataJson("data.demo.allow", Data); +var indexed = vm.ExecuteEntryPointWithDataJson(0UL, Data); +``` + +Those methods include native JSON parsing/storage for `dataJson`, execution, native result JSON serialization, and native C-string allocation. Managed UTF-8 decoding and C# `string` allocation after the native call returns are not charged. If scoped data replacement fails, the previous VM data is preserved, but the previous and provisional data may coexist transiently and count toward peak live bytes. + +The budget is cooperative and may overshoot between VM checks. Same-thread allocation-counter baseline ratcheting can make the effective limit stricter after unrelated frees are observed; those frees are not credited back. Budgets are not supported in suspendable execution mode. `ClearMemoryBudgetConfig` restores the previous unlimited per-execution behavior. Public multi-call begin/end scopes are intentionally absent because allocator counters are thread-local. The process-wide limit exposed by `MemoryLimits` remains a separate safeguard. ## Azure RBAC Condition Evaluation diff --git a/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs index 589286951..09c6f99a5 100644 --- a/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs +++ b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs @@ -22,6 +22,25 @@ import rego.v1 private const string EntryPoint = "data.limits.memory.large_array"; + private const string PreloadedResultPolicy = """ +package limits.memory + +large_string := data.large_string +"""; + + private const string PreloadedResultEntryPoint = "data.limits.memory.large_string"; + + private const string DataResultPolicy = """ +package limits.memory + +main := data.value +named := data.value +"""; + + private const string MainEntryPoint = "data.limits.memory.main"; + + private const string NamedEntryPoint = "data.limits.memory.named"; + [TestMethod] public void Memory_budget_must_be_non_zero() { @@ -55,6 +74,144 @@ public void Clearing_memory_budget_restores_unlimited_execution() Assert.IsFalse(string.IsNullOrWhiteSpace(result)); } + [TestMethod] + public void Serialization_budget_failure_leaves_error_state() + { + var data = JsonSerializer.Serialize(new + { + large_string = new string('x', 2 * 1024 * 1024), + }); + var modules = new[] { new PolicyModule("memory_budget.rego", PreloadedResultPolicy) }; + using var program = Program.CompileFromModules(data, modules, new[] { PreloadedResultEntryPoint }); + using var vm = new Rvm(); + vm.LoadProgram(program); + vm.SetDataJson(data); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(512 * 1024)); + + Assert.ThrowsException( + () => vm.ExecuteEntryPoint(PreloadedResultEntryPoint)); + + var state = vm.GetExecutionState(); + Assert.IsNotNull(state); + StringAssert.Contains(state, "Error { error: MemoryBudgetExceeded"); + } + + [TestMethod] + public void Execute_with_data_json_runs_main_entry_point() + { + using var program = CreateDataResultProgram(); + using var vm = CreateRvm(program); + + var result = vm.ExecuteWithDataJson(CreateValueData("one-call")); + + Assert.AreEqual("\"one-call\"", result); + } + + [TestMethod] + public void Execute_entry_point_with_data_json_runs_named_and_indexed_entry_points() + { + using var program = CreateDataResultProgram(); + using var vm = CreateRvm(program); + + Assert.AreEqual("\"named\"", vm.ExecuteEntryPointWithDataJson(NamedEntryPoint, CreateValueData("named"))); + Assert.AreEqual("\"indexed\"", vm.ExecuteEntryPointWithDataJson(0, CreateValueData("indexed"))); + } + + [TestMethod] + public void Execute_with_data_json_exceeding_data_budget_throws_typed_exception() + { + using var program = CreateDataResultProgram(); + using var vm = CreateRvm(program); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); + + Assert.ThrowsException( + () => vm.ExecuteWithDataJson(CreateOversizedValueData())); + } + + [TestMethod] + public void Execute_with_data_json_failed_replacement_preserves_previous_vm_data() + { + using var program = CreateDataResultProgram(); + using var vm = CreateRvm(program); + vm.SetDataJson(CreateValueData("previous")); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); + + Assert.ThrowsException( + () => vm.ExecuteWithDataJson(CreateOversizedValueData())); + + var state = vm.GetExecutionState(); + Assert.IsNotNull(state); + StringAssert.Contains(state, "Error { error: MemoryBudgetExceeded"); + + vm.ClearMemoryBudgetConfig(); + + Assert.AreEqual("\"previous\"", vm.ExecuteEntryPoint(MainEntryPoint)); + } + + [TestMethod] + public void Execute_with_malformed_data_json_leaves_error_state_and_preserves_previous_vm_data() + { + using var program = CreateDataResultProgram(); + using var vm = CreateRvm(program); + vm.SetDataJson(CreateValueData("previous")); + + Assert.ThrowsException(() => vm.ExecuteWithDataJson("{")); + + var state = vm.GetExecutionState(); + Assert.IsNotNull(state); + StringAssert.Contains(state, "Error { error:"); + + Assert.AreEqual("\"previous\"", vm.ExecuteEntryPoint(MainEntryPoint)); + } + + [TestMethod] + public void Execute_with_data_json_reuses_vm_with_a_fresh_budget_after_success_and_failure() + { + using var program = CreateDataResultProgram(); + using var vm = CreateRvm(program); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); + + Assert.AreEqual("\"first\"", vm.ExecuteWithDataJson(CreateValueData("first"))); + Assert.ThrowsException( + () => vm.ExecuteWithDataJson(CreateOversizedValueData())); + Assert.AreEqual("\"second\"", vm.ExecuteWithDataJson(CreateValueData("second"))); + } + + [TestMethod] + public void Execute_with_data_json_uses_unlimited_default_and_cleared_budget() + { + using var program = CreateDataResultProgram(); + using var vm = CreateRvm(program); + + Assert.AreEqual("\"default\"", vm.ExecuteWithDataJson(CreateValueData("default"))); + + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); + Assert.ThrowsException( + () => vm.ExecuteWithDataJson(CreateOversizedValueData())); + + vm.ClearMemoryBudgetConfig(); + + var result = vm.ExecuteWithDataJson(CreateOversizedValueData()); + Assert.IsFalse(string.IsNullOrWhiteSpace(result)); + } + + [TestMethod] + public void Execute_with_data_json_counts_native_result_serialization_but_not_managed_string_copy() + { + using var program = CreateLargeResultProgram(); + using var vm = CreateRvm(program); + vm.SetInputJson(CreateLargeResultInput()); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(512 * 1024)); + + Assert.ThrowsException(() => vm.ExecuteWithDataJson("{}")); + + vm.ClearMemoryBudgetConfig(); + + var result = vm.ExecuteWithDataJson("{}"); + Assert.IsNotNull(result); + Assert.AreEqual((2 * 1024 * 1024) + 2, result!.Length); + } + [TestMethod] public void Suspendable_execution_rejects_memory_budget() { @@ -65,6 +222,38 @@ public void Suspendable_execution_rejects_memory_budget() Assert.ThrowsException(() => vm.Execute()); } + private static Program CreateDataResultProgram() + { + var modules = new[] { new PolicyModule("data_result.rego", DataResultPolicy) }; + return Program.CompileFromModules("{}", modules, new[] { MainEntryPoint, NamedEntryPoint }); + } + + private static Program CreateLargeResultProgram() + { + const string policy = """ +package limits.memory + +large_string := input.large_string +"""; + var modules = new[] { new PolicyModule("large_result.rego", policy) }; + return Program.CompileFromModules("{}", modules, new[] { "data.limits.memory.large_string" }); + } + + private static string CreateLargeResultInput() + { + return JsonSerializer.Serialize(new { large_string = new string('x', 2 * 1024 * 1024) }); + } + + private static string CreateValueData(string value) + { + return JsonSerializer.Serialize(new { value }); + } + + private static string CreateOversizedValueData() + { + return CreateValueData(new string('x', 2 * 1024 * 1024)); + } + private static Program CreateProgram() { var modules = new[] { new PolicyModule("memory_budget.rego", Policy) }; diff --git a/bindings/csharp/Regorus/NativeMethods.cs b/bindings/csharp/Regorus/NativeMethods.cs index 479c9b7c9..b07f82f96 100644 --- a/bindings/csharp/Regorus/NativeMethods.cs +++ b/bindings/csharp/Regorus/NativeMethods.cs @@ -204,6 +204,24 @@ internal static unsafe partial class API [DllImport(LibraryName, EntryPoint = "regorus_rvm_execute_entry_point_by_index", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] internal static extern RegorusResult regorus_rvm_execute_entry_point_by_index(RegorusRvm* vm, UIntPtr index); + /// + /// Set the data document and execute the program in one native call. + /// + [DllImport(LibraryName, EntryPoint = "regorus_rvm_execute_with_data", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] + internal static extern RegorusResult regorus_rvm_execute_with_data(RegorusRvm* vm, byte* data_json); + + /// + /// Set the data document and execute a named entry point in one native call. + /// + [DllImport(LibraryName, EntryPoint = "regorus_rvm_execute_entry_point_by_name_with_data", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] + internal static extern RegorusResult regorus_rvm_execute_entry_point_by_name_with_data(RegorusRvm* vm, byte* entry_point, byte* data_json); + + /// + /// Set the data document and execute an entry point by index in one native call. + /// + [DllImport(LibraryName, EntryPoint = "regorus_rvm_execute_entry_point_by_index_with_data", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] + internal static extern RegorusResult regorus_rvm_execute_entry_point_by_index_with_data(RegorusRvm* vm, UIntPtr index, byte* data_json); + /// /// Resume execution. /// diff --git a/bindings/csharp/Regorus/Rvm.cs b/bindings/csharp/Regorus/Rvm.cs index 884731e58..e6e0c3e29 100644 --- a/bindings/csharp/Regorus/Rvm.cs +++ b/bindings/csharp/Regorus/Rvm.cs @@ -212,6 +212,57 @@ public void ClearMemoryBudgetConfig() }); } + /// + /// Set evaluation-specific data JSON and execute the program in one native call. + /// + public string? ExecuteWithDataJson(string dataJson) + { + return Utf8Marshaller.WithUtf8(dataJson, dataPtr => + { + return UseHandle(vmPtr => + { + return CheckAndDropResult(API.regorus_rvm_execute_with_data((RegorusRvm*)vmPtr, (byte*)dataPtr)); + }); + }); + } + + /// + /// Set evaluation-specific data JSON and execute a named entry point in one native call. + /// + public string? ExecuteEntryPointWithDataJson(string entryPoint, string dataJson) + { + return Utf8Marshaller.WithUtf8(entryPoint, entryPtr => + { + return Utf8Marshaller.WithUtf8(dataJson, dataPtr => + { + return UseHandle(vmPtr => + { + return CheckAndDropResult(API.regorus_rvm_execute_entry_point_by_name_with_data( + (RegorusRvm*)vmPtr, + (byte*)entryPtr, + (byte*)dataPtr)); + }); + }); + }); + } + + /// + /// Set evaluation-specific data JSON and execute an entry point by index in one native call. + /// + public string? ExecuteEntryPointWithDataJson(ulong index, string dataJson) + { + return Utf8Marshaller.WithUtf8(dataJson, dataPtr => + { + return UseHandle(vmPtr => + { + return CheckAndDropResult(API.regorus_rvm_execute_entry_point_by_index_with_data( + (RegorusRvm*)vmPtr, + (UIntPtr)index, + (byte*)dataPtr)); + }); + }); + } + /// /// Resume execution with an optional value. /// diff --git a/bindings/ffi/CHANGELOG.md b/bindings/ffi/CHANGELOG.md index f1b156b0d..04ab70683 100644 --- a/bindings/ffi/CHANGELOG.md +++ b/bindings/ffi/CHANGELOG.md @@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Add `regorus_rvm_set_memory_budget_config`, `RegorusMemoryBudgetConfig`, and appended statuses for memory-budget exhaustion and unsupported suspendable execution. +- Add one-call RVM data execution APIs: `regorus_rvm_execute_with_data`, + `regorus_rvm_execute_entry_point_by_name_with_data`, and + `regorus_rvm_execute_entry_point_by_index_with_data`. Their memory-budget scope covers data + preparation, execution, and native result production. ## [0.1.0](https://github.com/microsoft/regorus/releases/tag/regorus-ffi-v0.1.0) - 2024-02-08 diff --git a/bindings/ffi/src/engine.rs b/bindings/ffi/src/engine.rs index 3ec469cc0..52016cde5 100644 --- a/bindings/ffi/src/engine.rs +++ b/bindings/ffi/src/engine.rs @@ -92,6 +92,7 @@ mod panic_tests { #[test] fn catches_extension_panics_and_marks_poison() { + let _poison_test_lock = crate::panic_guard::lock_poison_test_state(); reset_poison(); let engine_ptr = regorus_engine_new(); diff --git a/bindings/ffi/src/limits.rs b/bindings/ffi/src/limits.rs index c69767981..43459e02f 100644 --- a/bindings/ffi/src/limits.rs +++ b/bindings/ffi/src/limits.rs @@ -252,7 +252,7 @@ pub extern "C" fn regorus_clear_cache() -> RegorusResult { RegorusResult::ok_void() } -#[cfg(all(test, not(miri)))] +#[cfg(all(test, feature = "allocator-memory-limits", not(miri)))] mod tests { #[cfg(feature = "allocator-memory-limits")] use super::RegorusMemoryBudgetConfig; diff --git a/bindings/ffi/src/panic_guard.rs b/bindings/ffi/src/panic_guard.rs index 0293e8f94..f5012a760 100644 --- a/bindings/ffi/src/panic_guard.rs +++ b/bindings/ffi/src/panic_guard.rs @@ -70,6 +70,82 @@ impl Drop for PanicHookGuard { static POISONED: AtomicBool = AtomicBool::new(false); +#[cfg(all(test, feature = "std"))] +static POISON_TEST_GATE: std::sync::RwLock<()> = std::sync::RwLock::new(()); + +#[cfg(all(test, feature = "std"))] +std::thread_local! { + static POISON_TEST_GATE_DEPTH: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +#[cfg(all(test, feature = "std"))] +pub(crate) struct PoisonTestLock { + lock: Option>, +} + +#[cfg(all(test, feature = "std"))] +impl Drop for PoisonTestLock { + fn drop(&mut self) { + self.lock.take(); + decrement_poison_test_gate_depth(); + } +} + +#[cfg(all(test, feature = "std"))] +pub(crate) fn lock_poison_test_state() -> PoisonTestLock { + debug_assert!( + !poison_test_gate_is_held(), + "the poison test gate cannot be upgraded from a nested guarded call" + ); + let lock = POISON_TEST_GATE + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + increment_poison_test_gate_depth(); + PoisonTestLock { lock: Some(lock) } +} + +#[cfg(all(test, feature = "std"))] +struct PoisonTestCallLock { + lock: Option>, +} + +#[cfg(all(test, feature = "std"))] +impl Drop for PoisonTestCallLock { + fn drop(&mut self) { + self.lock.take(); + decrement_poison_test_gate_depth(); + } +} + +#[cfg(all(test, feature = "std"))] +fn poison_test_call_lock() -> PoisonTestCallLock { + let nested = poison_test_gate_is_held(); + increment_poison_test_gate_depth(); + if nested { + PoisonTestCallLock { lock: None } + } else { + let lock = POISON_TEST_GATE + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + PoisonTestCallLock { lock: Some(lock) } + } +} + +#[cfg(all(test, feature = "std"))] +fn poison_test_gate_is_held() -> bool { + POISON_TEST_GATE_DEPTH.with(|depth| depth.get() > 0) +} + +#[cfg(all(test, feature = "std"))] +fn increment_poison_test_gate_depth() { + POISON_TEST_GATE_DEPTH.with(|depth| depth.set(depth.get().saturating_add(1))); +} + +#[cfg(all(test, feature = "std"))] +fn decrement_poison_test_gate_depth() { + POISON_TEST_GATE_DEPTH.with(|depth| depth.set(depth.get().saturating_sub(1))); +} + /// Result of attempting to run `f` while guarding against unwinding. pub(crate) enum GuardResult { /// Closure completed successfully. @@ -82,6 +158,9 @@ pub(crate) fn with_unwind_guard(f: F) -> RegorusResult where F: FnOnce() -> RegorusResult, { + #[cfg(all(test, feature = "std"))] + let _poison_test_call_lock = poison_test_call_lock(); + if is_poisoned() { return poisoned_result(); } @@ -157,3 +236,55 @@ pub(crate) fn is_poisoned() -> bool { pub(crate) fn reset_poison() { POISONED.store(false, Ordering::Release); } + +#[cfg(all(test, feature = "std"))] +mod tests { + use super::with_unwind_guard; + use crate::common::{regorus_result_drop, RegorusResult, RegorusStatus}; + use std::sync::mpsc; + use std::time::Duration; + + #[test] + fn ordinary_guarded_calls_are_not_serialized_by_the_poison_test_gate() { + let (first_entered_sender, first_entered_receiver) = mpsc::channel(); + let (second_entered_sender, second_entered_receiver) = mpsc::channel(); + + let first = std::thread::spawn(move || { + let result = with_unwind_guard(|| { + let _ = first_entered_sender.send(()); + let second_entered = second_entered_receiver + .recv_timeout(Duration::from_secs(1)) + .is_ok(); + RegorusResult::ok_bool(second_entered) + }); + let success = matches!(result.status, RegorusStatus::Ok) && result.bool_value; + regorus_result_drop(result); + success + }); + + assert!( + first_entered_receiver + .recv_timeout(Duration::from_secs(1)) + .is_ok(), + "first guarded call did not start" + ); + + let second = std::thread::spawn(move || { + let result = with_unwind_guard(|| { + let _ = second_entered_sender.send(()); + RegorusResult::ok_void() + }); + let success = matches!(result.status, RegorusStatus::Ok); + regorus_result_drop(result); + success + }); + + let first_success = first.join().unwrap_or(false); + let second_success = second.join().unwrap_or(false); + assert!(second_success, "second guarded call did not complete"); + assert!( + first_success, + "ordinary guarded calls were serialized by the poison test gate" + ); + } +} diff --git a/bindings/ffi/src/rvm.rs b/bindings/ffi/src/rvm.rs index 20699c8c2..f0a6f2a7a 100644 --- a/bindings/ffi/src/rvm.rs +++ b/bindings/ffi/src/rvm.rs @@ -76,24 +76,122 @@ fn to_rvm_error_result(err: anyhow::Error) -> RegorusResult { RegorusResult::err_with_message(status, err.to_string()) } -fn execute_to_rvm_result(vm: *mut RegorusRvm, execute: F) -> RegorusResult -where - F: FnOnce(&mut RegoVM) -> core::result::Result, -{ +enum RvmExecution { + Main, + Named(String), + Indexed(usize), +} + +/// Owns a provisional native result until the surrounding budget scope succeeds. +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +struct ProvisionalRvmResult { + result: RegorusResult, + owns_result: bool, +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +impl ProvisionalRvmResult { + fn new(result: RegorusResult) -> Self { + Self { + result, + owns_result: true, + } + } + + fn into_result(mut self) -> RegorusResult { + self.owns_result = false; + core::mem::replace(&mut self.result, RegorusResult::ok_void()) + } +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +impl Drop for ProvisionalRvmResult { + fn drop(&mut self) { + if self.owns_result { + #[cfg(all(test, feature = "std"))] + PROVISIONAL_RVM_RESULT_DROP_COUNT.with(|count| { + count.set(count.get().saturating_add(1)); + }); + regorus_result_drop(core::mem::replace( + &mut self.result, + RegorusResult::ok_void(), + )); + } + } +} + +#[cfg(all(test, feature = "allocator-memory-limits", not(miri), feature = "std"))] +std::thread_local! { + static PROVISIONAL_RVM_RESULT_DROP_COUNT: std::cell::Cell = + const { std::cell::Cell::new(0) }; +} + +#[cfg(all(test, feature = "allocator-memory-limits", not(miri), feature = "std"))] +fn reset_provisional_rvm_result_drop_count() { + PROVISIONAL_RVM_RESULT_DROP_COUNT.with(|count| count.set(0)); +} + +#[cfg(all(test, feature = "allocator-memory-limits", not(miri), feature = "std"))] +fn provisional_rvm_result_drop_count() -> usize { + PROVISIONAL_RVM_RESULT_DROP_COUNT.with(std::cell::Cell::get) +} + +fn prepare_data( + vm: &mut RegoVM, + data_json: Option<*const c_char>, +) -> core::result::Result<(), VmError> { + if let Some(data_json) = data_json { + let data = Value::from_json_str(&from_c_str(data_json)?).map_err(VmError::from)?; + vm.set_data(data)?; + } + + Ok(()) +} + +fn execute_rvm(vm: &mut RegoVM, execution: RvmExecution) -> core::result::Result { + match execution { + RvmExecution::Main => vm.execute(), + RvmExecution::Named(entry_point) => vm.execute_entry_point_by_name(&entry_point), + RvmExecution::Indexed(index) => vm.execute_entry_point_by_index(index), + } +} + +fn execute_to_rvm_result( + vm: *mut RegorusRvm, + data_json: Option<*const c_char>, + execution: RvmExecution, +) -> RegorusResult { let output = || -> Result { let vm = to_shared_ref(vm as *const RegorusRvm)?; let mut guard = vm.try_write()?; - let value = execute(&mut guard)?; - let json = value.to_json_str()?; - let result = RegorusResult::ok_string(json); #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - if let Err(err) = guard.check_memory_budget() { - regorus_result_drop(result); - return Err(err.into()); + { + let result = guard.with_evaluation_memory_budget(|guard| { + prepare_data(guard, data_json)?; + let value = execute_rvm(guard, execution)?; + let json = value.to_json_str().map_err(VmError::from)?; + let result = ProvisionalRvmResult::new(RegorusResult::ok_string(json)); + + // This remains separate from the scope's final check so the + // provisional output is dropped before an error crosses FFI. + guard.check_memory_budget()?; + Ok(result) + }); + + match result { + Ok(result) => Ok(result.into_result()), + Err(err) => Err(anyhow::Error::from(guard.fail_run_to_completion(err))), + } } - Ok(result) + #[cfg(any(not(feature = "allocator-memory-limits"), miri))] + { + prepare_data(&mut guard, data_json)?; + let value = execute_rvm(&mut guard, execution)?; + let json = value.to_json_str()?; + Ok(RegorusResult::ok_string(json)) + } }(); match output { @@ -574,7 +672,7 @@ pub extern "C" fn regorus_rvm_set_memory_budget_config( /// Execute the program's main entry point. #[no_mangle] pub extern "C" fn regorus_rvm_execute(vm: *mut RegorusRvm) -> RegorusResult { - with_unwind_guard(|| execute_to_rvm_result(vm, RegoVM::execute)) + with_unwind_guard(|| execute_to_rvm_result(vm, None, RvmExecution::Main)) } /// Execute a named entry point. @@ -584,11 +682,11 @@ pub extern "C" fn regorus_rvm_execute_entry_point_by_name( entry_point: *const c_char, ) -> RegorusResult { with_unwind_guard(|| { - let name = match from_c_str(entry_point) { - Ok(name) => name, + let entry_point = match from_c_str(entry_point) { + Ok(entry_point) => entry_point, Err(err) => return to_rvm_error_result(err), }; - execute_to_rvm_result(vm, |guard| guard.execute_entry_point_by_name(&name)) + execute_to_rvm_result(vm, None, RvmExecution::Named(entry_point)) }) } @@ -597,12 +695,45 @@ pub extern "C" fn regorus_rvm_execute_entry_point_by_name( pub extern "C" fn regorus_rvm_execute_entry_point_by_index( vm: *mut RegorusRvm, index: usize, +) -> RegorusResult { + with_unwind_guard(|| execute_to_rvm_result(vm, None, RvmExecution::Indexed(index))) +} + +/// Set the VM data from JSON and execute the program's main entry point in one call. +#[no_mangle] +pub extern "C" fn regorus_rvm_execute_with_data( + vm: *mut RegorusRvm, + data_json: *const c_char, +) -> RegorusResult { + with_unwind_guard(|| execute_to_rvm_result(vm, Some(data_json), RvmExecution::Main)) +} + +/// Set the VM data from JSON and execute a named entry point in one call. +#[no_mangle] +pub extern "C" fn regorus_rvm_execute_entry_point_by_name_with_data( + vm: *mut RegorusRvm, + entry_point: *const c_char, + data_json: *const c_char, ) -> RegorusResult { with_unwind_guard(|| { - execute_to_rvm_result(vm, |guard| guard.execute_entry_point_by_index(index)) + let entry_point = match from_c_str(entry_point) { + Ok(entry_point) => entry_point, + Err(err) => return to_rvm_error_result(err), + }; + execute_to_rvm_result(vm, Some(data_json), RvmExecution::Named(entry_point)) }) } +/// Set the VM data from JSON and execute an entry point by index in one call. +#[no_mangle] +pub extern "C" fn regorus_rvm_execute_entry_point_by_index_with_data( + vm: *mut RegorusRvm, + index: usize, + data_json: *const c_char, +) -> RegorusResult { + with_unwind_guard(|| execute_to_rvm_result(vm, Some(data_json), RvmExecution::Indexed(index))) +} + /// Resume execution for suspendable runs. #[no_mangle] pub extern "C" fn regorus_rvm_resume( @@ -649,16 +780,19 @@ pub extern "C" fn regorus_rvm_get_execution_state(vm: *mut RegorusRvm) -> Regoru mod tests { use super::{ regorus_rvm_drop, regorus_rvm_execute, regorus_rvm_execute_entry_point_by_index, - regorus_rvm_execute_entry_point_by_name, regorus_rvm_new, regorus_rvm_resume, - regorus_rvm_set_memory_budget_config, RegorusRvm, + regorus_rvm_execute_entry_point_by_index_with_data, + regorus_rvm_execute_entry_point_by_name, regorus_rvm_execute_entry_point_by_name_with_data, + regorus_rvm_execute_with_data, regorus_rvm_get_execution_state, regorus_rvm_new, + regorus_rvm_resume, regorus_rvm_set_data, regorus_rvm_set_memory_budget_config, RegorusRvm, }; - use crate::common::{regorus_result_drop, RegorusStatus}; + use crate::common::{regorus_result_drop, RegorusResult, RegorusStatus}; use crate::limits::RegorusMemoryBudgetConfig; use alloc::boxed::Box; use alloc::ffi::CString; use alloc::string::ToString; use alloc::sync::Arc; use alloc::vec; + use core::ffi::CStr; use core::ptr; use regorus::languages::rego::compiler::Compiler; use regorus::rvm::instructions::Instruction; @@ -681,6 +815,31 @@ copy := [value | some value in input] } } + fn assert_memory_budget_failure_state(vm: *mut RegorusRvm, result: RegorusResult) { + assert!(matches!(result.status, RegorusStatus::MemoryBudgetExceeded)); + assert!(result.output.is_null()); + regorus_result_drop(result); + + assert_execution_state(vm, "Error { error: MemoryBudgetExceeded"); + } + + fn assert_execution_state_error(vm: *mut RegorusRvm) { + assert_execution_state(vm, "Error { error:"); + } + + fn assert_execution_state(vm: *mut RegorusRvm, expected_prefix: &str) { + let state = regorus_rvm_get_execution_state(vm); + assert!(matches!(state.status, RegorusStatus::Ok)); + let state_text = unsafe { CStr::from_ptr(state.output) } + .to_str() + .expect("execution state UTF-8"); + assert!( + state_text.starts_with(expected_prefix), + "unexpected execution state: {state_text}" + ); + regorus_result_drop(state); + } + fn host_await_program() -> Arc { let mut program = Program::new(); program.dispatch_window_size = 3; @@ -724,6 +883,67 @@ copy := [value | some value in input] Arc::new(program) } + fn data_result_program() -> Arc { + let mut program = Program::new(); + program.dispatch_window_size = 1; + program.max_rule_window_size = 1; + program.entry_points.insert("main".to_string(), 0); + program.entry_points.insert("named".to_string(), 0); + program.instructions = vec![ + Instruction::LoadData { dest: 0 }, + Instruction::Return { value: 0 }, + ]; + program.instruction_spans = vec![None; program.instructions.len()]; + Arc::new(program) + } + + fn result_json(result: RegorusResult) -> String { + if !matches!(result.status, RegorusStatus::Ok) { + let message = if result.error_message.is_null() { + "" + } else { + unsafe { CStr::from_ptr(result.error_message) } + .to_str() + .unwrap_or("") + }; + panic!( + "unexpected RVM result status {:?}: {message}", + result.status + ); + } + let output = unsafe { CStr::from_ptr(result.output) } + .to_str() + .expect("result UTF-8") + .to_string(); + regorus_result_drop(result); + output + } + + fn assert_result_json_eq(result: RegorusResult, expected: &str) { + assert_eq!( + Value::from_json_str(&result_json(result)).expect("result JSON"), + Value::from_json_str(expected).expect("expected JSON") + ); + } + + #[cfg(feature = "std")] + #[test] + fn provisional_rvm_result_transfers_or_drops_ownership() { + super::reset_provisional_rvm_result_drop_count(); + { + let _provisional = + super::ProvisionalRvmResult::new(RegorusResult::ok_string("dropped".into())); + } + assert_eq!(super::provisional_rvm_result_drop_count(), 1); + + super::reset_provisional_rvm_result_drop_count(); + let result = + super::ProvisionalRvmResult::new(RegorusResult::ok_string("transferred".into())) + .into_result(); + assert_eq!(super::provisional_rvm_result_drop_count(), 0); + regorus_result_drop(result); + } + #[test] fn ffi_memory_budget_setter_validates_and_clears_configuration() { let vm = regorus_rvm_new(); @@ -799,17 +1019,148 @@ copy := [value | some value in input] let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); let entrypoint = CString::new("main").expect("entry point CString"); - let results = [ - regorus_rvm_execute(vm), + assert_memory_budget_failure_state(vm, regorus_rvm_execute(vm)); + assert_memory_budget_failure_state( + vm, regorus_rvm_execute_entry_point_by_name(vm, entrypoint.as_ptr()), - regorus_rvm_execute_entry_point_by_index(vm, 0), - ]; + ); + assert_memory_budget_failure_state(vm, regorus_rvm_execute_entry_point_by_index(vm, 0)); + regorus_rvm_drop(vm); + } - for result in results { - assert!(matches!(result.status, RegorusStatus::MemoryBudgetExceeded)); - assert!(result.output.is_null()); - regorus_result_drop(result); - } + #[test] + fn ffi_execute_with_data_runs_main_named_and_indexed_entry_points() { + let data = CString::new(r#"{"value":"one-call"}"#).expect("data CString"); + let entrypoint = CString::new("named").expect("entry point CString"); + + let mut main_vm = RegoVM::new(); + main_vm.load_program(data_result_program()); + let main_vm = Box::into_raw(Box::new(RegorusRvm::new(main_vm))); + assert_result_json_eq( + regorus_rvm_execute_with_data(main_vm, data.as_ptr()), + r#"{"value":"one-call"}"#, + ); + assert_result_json_eq( + regorus_rvm_execute_with_data(main_vm, data.as_ptr()), + r#"{"value":"one-call"}"#, + ); + regorus_rvm_drop(main_vm); + + let mut named_vm = RegoVM::new(); + named_vm.load_program(data_result_program()); + let named_vm = Box::into_raw(Box::new(RegorusRvm::new(named_vm))); + assert_result_json_eq( + regorus_rvm_execute_entry_point_by_name_with_data( + named_vm, + entrypoint.as_ptr(), + data.as_ptr(), + ), + r#"{"value":"one-call"}"#, + ); + assert_result_json_eq( + regorus_rvm_execute_entry_point_by_name_with_data( + named_vm, + entrypoint.as_ptr(), + data.as_ptr(), + ), + r#"{"value":"one-call"}"#, + ); + regorus_rvm_drop(named_vm); + + let mut indexed_vm = RegoVM::new(); + indexed_vm.load_program(data_result_program()); + let indexed_vm = Box::into_raw(Box::new(RegorusRvm::new(indexed_vm))); + assert_result_json_eq( + regorus_rvm_execute_entry_point_by_index_with_data(indexed_vm, 0, data.as_ptr()), + r#"{"value":"one-call"}"#, + ); + assert_result_json_eq( + regorus_rvm_execute_entry_point_by_index_with_data(indexed_vm, 0, data.as_ptr()), + r#"{"value":"one-call"}"#, + ); + regorus_rvm_drop(indexed_vm); + } + + #[test] + fn ffi_execute_with_data_rejects_oversized_replacement_without_losing_previous_data() { + let mut vm = RegoVM::new(); + vm.load_program(data_result_program()); + let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); + let old_data = CString::new(r#"{"previous":true}"#).expect("old data CString"); + let set_data = regorus_rvm_set_data(vm, old_data.as_ptr()); + assert!(matches!(set_data.status, RegorusStatus::Ok)); + regorus_result_drop(set_data); + + let set_budget = regorus_rvm_set_memory_budget_config( + vm, + true, + RegorusMemoryBudgetConfig { + limit_bytes: TIGHT_MEMORY_BUDGET_BYTES, + }, + ); + assert!(matches!(set_budget.status, RegorusStatus::Ok)); + regorus_result_drop(set_budget); + + let oversized_data = CString::new(format!( + r#"{{"replacement":"{}"}}"#, + "x".repeat(2 * 1024 * 1024) + )) + .expect("oversized data CString"); + let rejected = regorus_rvm_execute_with_data(vm, oversized_data.as_ptr()); + assert_memory_budget_failure_state(vm, rejected); + + let clear_budget = regorus_rvm_set_memory_budget_config( + vm, + false, + RegorusMemoryBudgetConfig { limit_bytes: 0 }, + ); + assert!(matches!(clear_budget.status, RegorusStatus::Ok)); + regorus_result_drop(clear_budget); + assert_result_json_eq(regorus_rvm_execute(vm), r#"{"previous":true}"#); + regorus_rvm_drop(vm); + } + + #[test] + fn ffi_malformed_scoped_data_terminalizes_and_preserves_previous_data() { + let mut vm = RegoVM::new(); + vm.load_program(data_result_program()); + let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); + let old_data = CString::new(r#"{"previous":true}"#).expect("old data CString"); + let set_data = regorus_rvm_set_data(vm, old_data.as_ptr()); + assert!(matches!(set_data.status, RegorusStatus::Ok)); + regorus_result_drop(set_data); + + let malformed_data = CString::new("{").expect("malformed data CString"); + let rejected = regorus_rvm_execute_with_data(vm, malformed_data.as_ptr()); + assert!(matches!(rejected.status, RegorusStatus::Error)); + assert!(rejected.output.is_null()); + regorus_result_drop(rejected); + assert_execution_state_error(vm); + + assert_result_json_eq(regorus_rvm_execute(vm), r#"{"previous":true}"#); + regorus_rvm_drop(vm); + } + + #[test] + fn ffi_execute_with_data_counts_native_result_production_and_leaves_vm_reusable() { + let mut vm = RegoVM::new(); + vm.load_program(preloaded_result_program()); + vm.set_memory_budget_config(Some(memory_budget(512 * 1024))); + let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); + let data = CString::new("{}").expect("data CString"); + + assert_memory_budget_failure_state(vm, regorus_rvm_execute_with_data(vm, data.as_ptr())); + + let clear_budget = regorus_rvm_set_memory_budget_config( + vm, + false, + RegorusMemoryBudgetConfig { limit_bytes: 0 }, + ); + assert!(matches!(clear_budget.status, RegorusStatus::Ok)); + regorus_result_drop(clear_budget); + let result = regorus_rvm_execute_with_data(vm, data.as_ptr()); + assert!(matches!(result.status, RegorusStatus::Ok)); + regorus_result_drop(result); regorus_rvm_drop(vm); } @@ -850,6 +1201,26 @@ copy := [value | some value in input] } } +#[cfg(all(test, any(not(feature = "allocator-memory-limits"), miri)))] +mod unsupported_memory_budget_tests { + use super::{regorus_rvm_drop, regorus_rvm_new, regorus_rvm_set_memory_budget_config}; + use crate::common::{regorus_result_drop, RegorusStatus}; + use crate::limits::RegorusMemoryBudgetConfig; + + #[test] + fn ffi_memory_budget_setter_is_unsupported_without_allocator_tracking() { + let vm = regorus_rvm_new(); + let result = regorus_rvm_set_memory_budget_config( + vm, + true, + RegorusMemoryBudgetConfig { limit_bytes: 1024 }, + ); + assert!(matches!(result.status, RegorusStatus::InvalidArgument)); + regorus_result_drop(result); + regorus_rvm_drop(vm); + } +} + fn convert_c_entry_points( entry_points: *const *const c_char, entry_points_len: usize, diff --git a/docs/limits/memory_budget.md b/docs/limits/memory_budget.md index 5acd3cd8b..3a8ecc060 100644 --- a/docs/limits/memory_budget.md +++ b/docs/limits/memory_budget.md @@ -1,8 +1,8 @@ # RVM memory budgets -RVM run-to-completion execution supports an optional memory budget when Regorus is built with the `allocator-memory-limits` feature. +RVM run-to-completion evaluation supports an optional memory budget when Regorus is built with the `allocator-memory-limits` feature. -The budget limits additional live bytes on the execution thread. Regorus captures a baseline when execution starts and compares later live-byte samples with that baseline. Every call to `execute`, `execute_entry_point_by_name`, or `execute_entry_point_by_index` starts with a fresh budget. +The budget limits additional live bytes on the execution thread. Regorus captures a baseline when a budgeted scope starts and compares later live-byte samples with that baseline. Each ordinary Rust `execute`, `execute_entry_point_by_name`, or `execute_entry_point_by_index` call starts with a fresh execution-only budget. ```rust use core::num::NonZeroU64; @@ -17,19 +17,39 @@ vm.set_memory_budget_config(Some(MemoryBudgetConfig { No configured budget preserves existing RVM behavior. A zero-byte budget is not representable in Rust and is rejected by language bindings. -## Included work +## Budget scopes -The budget starts when RVM execution begins. Fresh execution-state initialization, rule evaluation, and allocations retained by the result count against the budget. +The ordinary `execute*` APIs start their budget when RVM execution begins. Fresh execution-state initialization, rule evaluation, and allocations retained by the result count against the budget. Program compilation, program loading, data loading, input loading, and context loading happen before the execution baseline and are not charged. -The C FFI and C# bindings also check the budget after result JSON serialization and native string marshaling, before returning success. +For callers that need evaluation-specific data to be charged with the evaluation, Rust exposes `RegoVM::with_evaluation_memory_budget`. This is a synchronous same-thread scope for allocations created after the closure begins. It can include evaluation-specific data construction and assignment, execution, and caller post-processing performed inside the closure. Nested scopes are rejected, and the scope is deactivated before the method returns on success, error, or unwinding. + +If a scope begins without a configured budget, configuring or re-enabling one inside the scope captures a fresh current-thread baseline at that transition. Changing one configured threshold to another preserves the existing scope baseline. + +The C FFI exposes one-call data execution APIs: + +- `regorus_rvm_execute_with_data` +- `regorus_rvm_execute_entry_point_by_name_with_data` +- `regorus_rvm_execute_entry_point_by_index_with_data` + +The C# binding exposes matching methods: + +- `Rvm.ExecuteWithDataJson(string dataJson)` +- `Rvm.ExecuteEntryPointWithDataJson(string entryPoint, string dataJson)` +- `Rvm.ExecuteEntryPointWithDataJson(ulong index, string dataJson)` + +These APIs run in one native call and use one budget scope covering native JSON parsing and storage for the supplied data, execution, native JSON serialization of the result, and native C-string allocation. Program loading, compilation, input loading, context loading, and managed C# UTF-8 decoding and `string` allocation remain outside the scope. + +Existing `set_data` / C# `SetDataJson` before execution remains the static/preloaded data path. That work happens before the ordinary execution budget and is excluded from it. + +If scoped data replacement fails, the previous VM data is preserved. The old data and provisional replacement can coexist transiently, so both contribute to the observed peak live bytes. ## Enforcement Regorus checks the budget at every VM memory checkpoint and once before returning a successful result. This is cooperative enforcement, not an allocation-time hard cap: a single instruction or builtin can overshoot the budget by an unbounded amount before the next checkpoint. A short-lived allocation created and freed entirely inside one instruction may not be observed. -Accounting uses the execution thread's live-byte counter rather than allocation ownership. When a sample falls below the baseline, Regorus lowers the baseline so an already-observed foreign free does not grant credit to later work. A foreign free can still offset evaluation allocations when both occur between samples because the allocator does not retain execution ownership for each allocation. The control therefore bounds observed additional live bytes on the execution thread, not memory attributed to an execution across threads. +Accounting uses the execution thread's live-byte counter rather than allocation ownership. When a sample falls below the baseline, Regorus lowers the baseline so an already-observed foreign free does not grant credit to later work. This downward baseline ratchet can make the effective limit stricter than the configured limit when unrelated same-thread frees are observed; it never grants those foreign frees back. A foreign free can still offset evaluation allocations when both occur between samples because the allocator does not retain execution ownership for each allocation. The control therefore bounds observed additional live bytes on the execution thread, not memory attributed to an execution across threads. Exhaustion returns `VmError::MemoryBudgetExceeded`, including: @@ -37,7 +57,7 @@ Exhaustion returns `VmError::MemoryBudgetExceeded`, including: - configured budget - VM program counter -The C FFI reports `RegorusStatus::MemoryBudgetExceeded`. The C# binding throws `RegorusMemoryBudgetExceededException`. +The VM transitions to `ExecutionState::Error` and releases values retained by the failed execution. The C FFI reports `RegorusStatus::MemoryBudgetExceeded`, including when native result serialization exceeds the budget. The C# binding throws `RegorusMemoryBudgetExceededException`. Every success or failure terminal path deactivates its owned scope, and reused VMs get a fresh budget for the next execution. ## Execution modes @@ -45,6 +65,8 @@ The first implementation supports run-to-completion execution only. Configuring Suspendable execution may resume on another thread. A thread-local baseline cannot safely span that migration without evaluation-owned allocation attribution. +Public multi-call begin/end memory-budget scopes are intentionally absent because allocation counters are thread-local and abandoned or cross-thread scopes would be unsafe. Rust, C FFI, and C# are supported by this PR; other bindings require follow-up work. + ## Process-global limit The existing process-global memory limit remains separate. It protects the process as a whole and is not an isolation mechanism for individual evaluations. When both controls are configured, the per-evaluation budget is checked first. diff --git a/src/rvm/vm/errors.rs b/src/rvm/vm/errors.rs index cdbd753f7..3187ea7d1 100644 --- a/src/rvm/vm/errors.rs +++ b/src/rvm/vm/errors.rs @@ -312,6 +312,9 @@ pub enum VmError { #[error("Internal VM error: {message} (pc={pc})")] Internal { message: String, pc: usize }, + + #[error("An evaluation memory budget scope is already active (pc={pc})")] + MemoryBudgetScopeAlreadyActive { pc: usize }, } impl From for VmError { diff --git a/src/rvm/vm/execution.rs b/src/rvm/vm/execution.rs index a6b0301bc..07f1fc5b0 100644 --- a/src/rvm/vm/execution.rs +++ b/src/rvm/vm/execution.rs @@ -46,23 +46,31 @@ impl RegoVM { match self.execution_mode { ExecutionMode::RunToCompletion => { - self.reset_run_to_completion_state(); - self.reset_execution_timer_state(); + let result = (|| { + self.reset_run_to_completion_state()?; + self.reset_execution_timer_state(); + self.validate_vm_state()?; + let entry_point_pc_u32 = u32::try_from(entry_point_pc).map_err(|_| { + VmError::EntryPointPcOutOfBounds { + pc: entry_point_pc, + instruction_count: self.program.instructions.len(), + entry_point: entry_point_name.clone(), + } + })?; - self.validate_vm_state()?; - let entry_point_pc_u32 = u32::try_from(entry_point_pc).map_err(|_| { - VmError::EntryPointPcOutOfBounds { - pc: entry_point_pc, - instruction_count: self.program.instructions.len(), - entry_point: entry_point_name.clone(), + let result = self + .jump_to(entry_point_pc_u32) + .map_err(|err| self.apply_memory_budget_precedence(err))?; + self.check_memory_budget()?; + Ok(result) + })(); + match result { + Ok(value) => { + self.finish_implicit_memory_budget_execution(); + Ok(value) } - })?; - - let result = self - .jump_to(entry_point_pc_u32) - .map_err(|err| self.apply_memory_budget_precedence(err))?; - self.check_memory_budget()?; - Ok(result) + Err(err) => Err(self.fail_run_to_completion(err)), + } } ExecutionMode::Suspendable => { self.reset_execution_state(); @@ -95,23 +103,31 @@ impl RegoVM { match self.execution_mode { ExecutionMode::RunToCompletion => { - self.reset_run_to_completion_state(); - self.reset_execution_timer_state(); + let result = (|| { + self.reset_run_to_completion_state()?; + self.reset_execution_timer_state(); + self.validate_vm_state()?; + let entry_point_pc_u32 = u32::try_from(entry_point_pc).map_err(|_| { + VmError::EntryPointPcOutOfBounds { + pc: entry_point_pc, + instruction_count: self.program.instructions.len(), + entry_point: String::from(name), + } + })?; - self.validate_vm_state()?; - let entry_point_pc_u32 = u32::try_from(entry_point_pc).map_err(|_| { - VmError::EntryPointPcOutOfBounds { - pc: entry_point_pc, - instruction_count: self.program.instructions.len(), - entry_point: String::from(name), + let result = self + .jump_to(entry_point_pc_u32) + .map_err(|err| self.apply_memory_budget_precedence(err))?; + self.check_memory_budget()?; + Ok(result) + })(); + match result { + Ok(value) => { + self.finish_implicit_memory_budget_execution(); + Ok(value) } - })?; - - let result = self - .jump_to(entry_point_pc_u32) - .map_err(|err| self.apply_memory_budget_precedence(err))?; - self.check_memory_budget()?; - Ok(result) + Err(err) => Err(self.fail_run_to_completion(err)), + } } ExecutionMode::Suspendable => { self.reset_execution_state(); @@ -174,27 +190,26 @@ impl RegoVM { } fn execute_run_to_completion(&mut self) -> Result { - self.reset_run_to_completion_state(); - self.reset_execution_timer_state(); - self.execution_state = ExecutionState::Running; - let result = self - .jump_to(0_u32) - .map_err(|err| self.apply_memory_budget_precedence(err)) - .and_then(|value| { - self.check_memory_budget()?; - Ok(value) - }); + let result = (|| { + self.reset_run_to_completion_state()?; + self.reset_execution_timer_state(); + self.execution_state = ExecutionState::Running; + self.jump_to(0_u32) + .map_err(|err| self.apply_memory_budget_precedence(err)) + .and_then(|value| { + self.check_memory_budget()?; + Ok(value) + }) + })(); match result { Ok(value) => { self.execution_state = ExecutionState::Completed { result: value.clone(), }; + self.finish_implicit_memory_budget_execution(); Ok(value) } - Err(err) => { - self.execution_state = ExecutionState::Error { error: err.clone() }; - Err(err) - } + Err(err) => Err(self.fail_run_to_completion(err)), } } diff --git a/src/rvm/vm/machine.rs b/src/rvm/vm/machine.rs index f8d79bb47..b1630d704 100644 --- a/src/rvm/vm/machine.rs +++ b/src/rvm/vm/machine.rs @@ -27,6 +27,14 @@ use super::execution_model::{ BreakpointSet, ExecutionMode, ExecutionStack, ExecutionState, SuspendReason, }; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum MemoryBudgetLifecycle { + Inactive, + ImplicitExecution, + ExternalScope, +} + /// The Rego Virtual Machine #[derive(Debug)] pub struct RegoVM { @@ -141,9 +149,9 @@ pub struct RegoVM { #[cfg(all(feature = "allocator-memory-limits", not(miri)))] pub(super) memory_budget_baseline: i64, - /// Whether the current baseline belongs to an active run-to-completion execution + /// Owner of the current memory-budget baseline. #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - pub(super) memory_budget_active: bool, + pub(super) memory_budget_lifecycle: MemoryBudgetLifecycle, /// Cached dummy span for builtin calls (avoids Source::from_contents per call) pub(super) dummy_span: Option, @@ -169,6 +177,25 @@ pub struct RegoVM { pub(super) metadata_value: Value, } +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +struct EvaluationMemoryBudgetScope<'a> { + vm: &'a mut RegoVM, +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +impl EvaluationMemoryBudgetScope<'_> { + const fn vm(&mut self) -> &mut RegoVM { + self.vm + } +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +impl Drop for EvaluationMemoryBudgetScope<'_> { + fn drop(&mut self) { + self.vm.finish_external_memory_budget_scope(); + } +} + impl Default for RegoVM { fn default() -> Self { Self::new() @@ -216,7 +243,7 @@ impl RegoVM { #[cfg(all(feature = "allocator-memory-limits", not(miri)))] memory_budget_baseline: 0, #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - memory_budget_active: false, + memory_budget_lifecycle: MemoryBudgetLifecycle::Inactive, dummy_span: None, dummy_exprs: Vec::new(), cached_builtin_args: Vec::new(), @@ -284,6 +311,16 @@ impl RegoVM { // Check for conflicts between rule tree and data self.program.check_rule_data_conflicts(&data)?; + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::ExternalScope + ) { + // Keep both the installed and provisional data alive until after the + // checkpoint so rejected replacements are atomic at their peak usage. + self.check_memory_budget()?; + } + self.data = data; Ok(()) } @@ -425,10 +462,24 @@ impl RegoVM { /// Configure a fresh memory budget for every run-to-completion execution. #[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] - pub const fn set_memory_budget_config(&mut self, config: Option) { + pub fn set_memory_budget_config(&mut self, config: Option) { + let was_configured = self.memory_budget_config.is_some(); self.memory_budget_config = config; + + if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::ExternalScope + ) { + if !was_configured && self.memory_budget_config.is_some() { + self.memory_budget_baseline = limits::current_thread_live_bytes(); + } else if self.memory_budget_config.is_none() { + self.memory_budget_baseline = 0; + } + return; + } + self.memory_budget_baseline = 0; - self.memory_budget_active = false; + self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; } /// Return the configured per-execution memory budget. @@ -440,12 +491,90 @@ impl RegoVM { #[cfg(all(feature = "allocator-memory-limits", not(miri)))] pub(super) fn reset_memory_budget_state(&mut self) { + if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::ExternalScope + ) { + return; + } + + self.memory_budget_baseline = if self.memory_budget_config.is_some() { + limits::current_thread_live_bytes() + } else { + 0 + }; + self.memory_budget_lifecycle = if self.memory_budget_config.is_some() { + MemoryBudgetLifecycle::ImplicitExecution + } else { + MemoryBudgetLifecycle::Inactive + }; + } + + /// Run a synchronous evaluation operation under one memory-budget baseline. + /// + /// When a memory budget is configured, data assigned with [`Self::set_data`], + /// VM execution, and allocations retained through the closure share the baseline + /// captured on entry. If a scope starts unconfigured, enabling or re-enabling a + /// budget captures a fresh baseline at that transition; changing one configured + /// threshold to another preserves the existing baseline. Nested scopes return + /// [`VmError::MemoryBudgetScopeAlreadyActive`]. + /// + /// The closure must complete synchronously on the calling thread. Its scope is + /// deactivated before this method returns, including when the closure returns an error + /// or unwinds. + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + #[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] + pub fn with_evaluation_memory_budget(&mut self, operation: F) -> Result + where + F: FnOnce(&mut Self) -> Result, + { + if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::ExternalScope + ) { + return Err(VmError::MemoryBudgetScopeAlreadyActive { pc: self.pc }); + } + self.memory_budget_baseline = if self.memory_budget_config.is_some() { limits::current_thread_live_bytes() } else { 0 }; - self.memory_budget_active = self.memory_budget_config.is_some(); + self.memory_budget_lifecycle = MemoryBudgetLifecycle::ExternalScope; + + let mut scope = EvaluationMemoryBudgetScope { vm: self }; + let result = operation(scope.vm()).and_then(|value| { + scope.vm().check_memory_budget()?; + Ok(value) + }); + + result + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) const fn finish_implicit_memory_budget_execution(&mut self) { + if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::ImplicitExecution + ) { + self.memory_budget_baseline = 0; + self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; + } + } + + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + #[allow(clippy::unused_self)] + pub(super) const fn finish_implicit_memory_budget_execution(&mut self) {} + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + const fn finish_external_memory_budget_scope(&mut self) { + if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::ExternalScope + ) { + self.memory_budget_baseline = 0; + self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; + } } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] @@ -480,17 +609,22 @@ impl RegoVM { Ok(()) } - /// Check the configured budget against the latest run-to-completion execution baseline. + /// Check the configured budget against the active evaluation-memory baseline. /// - /// Bindings can call this on the execution thread after result serialization so their - /// marshaling allocations are included before returning success. + /// Within [`Self::with_evaluation_memory_budget`], call this after scoped data + /// preparation or native result post-processing to include those allocations before + /// returning success. Ordinary `execute*` calls checkpoint their execution budget and + /// deactivate it on successful return; bindings that need to charge serialization must + /// wrap execution and serialization in an external evaluation-memory scope. #[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] pub fn check_memory_budget(&mut self) -> Result<()> { - let Some(config) = self - .memory_budget_config - .filter(|_| self.memory_budget_active) - else { + let Some(config) = self.memory_budget_config.filter(|_| { + !matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + ) + }) else { return Ok(()); }; diff --git a/src/rvm/vm/state.rs b/src/rvm/vm/state.rs index 15eff1ff4..1c9bb1ad8 100644 --- a/src/rvm/vm/state.rs +++ b/src/rvm/vm/state.rs @@ -16,11 +16,21 @@ impl RegoVM { } /// Release values retained by the previous execution before capturing a new memory baseline. - pub(super) fn reset_run_to_completion_state(&mut self) { + pub(super) fn reset_run_to_completion_state(&mut self) -> Result<()> { self.release_previous_execution_state(); #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - self.reset_memory_budget_state(); + if matches!( + self.memory_budget_lifecycle, + super::machine::MemoryBudgetLifecycle::ExternalScope + ) { + // An external scope owns the baseline. Observe the post-release trough + // before initialization so fresh state allocations are charged to it. + self.check_memory_budget()?; + } else { + self.reset_memory_budget_state(); + } self.initialize_execution_state(); + Ok(()) } fn release_previous_execution_state(&mut self) { @@ -34,6 +44,23 @@ impl RegoVM { self.cached_builtin_args.clear(); } + /// Release values retained by a failed run-to-completion execution and record its error. + /// + /// This is public only so binding crates can apply the same terminal transition when + /// post-execution marshaling exceeds the execution budget. + #[doc(hidden)] + pub fn fail_run_to_completion(&mut self, error: VmError) -> VmError { + self.release_previous_execution_state(); + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + { + self.finish_implicit_memory_budget_execution(); + } + self.execution_state = ExecutionState::Error { + error: error.clone(), + }; + error + } + fn initialize_execution_state(&mut self) { // Reset basic execution state self.executed_instructions = 0; @@ -236,3 +263,372 @@ impl RegoVM { Ok(()) } } + +#[cfg(test)] +mod tests { + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + use super::super::machine::MemoryBudgetLifecycle; + use super::{ExecutionState, RegoVM, Value, VmError}; + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + use crate::rvm::program::{RuleInfo, RuleType}; + use crate::rvm::{instructions::Instruction, program::Program}; + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + use crate::MemoryBudgetConfig; + use alloc::sync::Arc; + use alloc::vec; + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + use core::num::NonZeroU64; + + #[test] + fn failed_run_to_completion_releases_retained_values() { + let mut vm = RegoVM::new(); + let retained = Value::from("retained"); + vm.evaluated = retained.clone(); + vm.registers = vec![retained.clone()]; + vm.rule_cache = vec![(true, retained.clone())]; + vm.register_stack.push(vec![retained.clone()]); + vm.cached_builtin_args = vec![retained.clone()]; + vm.execution_state = ExecutionState::Completed { result: retained }; + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + { + vm.memory_budget_lifecycle = MemoryBudgetLifecycle::ImplicitExecution; + } + + let error = VmError::MemoryBudgetExceeded { + usage: 2, + budget: 1, + pc: 3, + }; + + assert_eq!(vm.fail_run_to_completion(error.clone()), error); + assert!(matches!(vm.evaluated, Value::Undefined)); + assert!(vm.registers.is_empty()); + assert!(vm.rule_cache.is_empty()); + assert!(vm.register_stack.is_empty()); + assert!(vm + .register_window_pool + .iter() + .all(alloc::vec::Vec::is_empty)); + assert!(vm.cached_builtin_args.is_empty()); + assert_eq!( + vm.execution_state, + ExecutionState::Error { + error: error.clone() + } + ); + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + #[allow(clippy::expect_used)] + #[test] + fn successful_ordinary_executions_finish_their_implicit_memory_budget() { + let mut program = Program::new(); + program.entry_points.insert("main".into(), 0); + program.instructions = vec![Instruction::Return { value: 0 }]; + program.instruction_spans = vec![None]; + + let mut vm = RegoVM::new(); + vm.load_program(Arc::new(program)); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1024 * 1024).unwrap_or(NonZeroU64::MIN), + })); + + assert_eq!( + vm.execute().expect("main execution succeeds"), + Value::Undefined + ); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + + assert_eq!( + vm.execute_entry_point_by_name("main") + .expect("named execution succeeds"), + Value::Undefined + ); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + + assert_eq!( + vm.execute_entry_point_by_index(0) + .expect("indexed execution succeeds"), + Value::Undefined + ); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + + vm.with_evaluation_memory_budget(|vm| { + assert_eq!( + vm.execute().expect("main scope execution"), + Value::Undefined + ); + assert_eq!( + vm.execute_entry_point_by_name("main") + .expect("named scope execution"), + Value::Undefined + ); + assert_eq!( + vm.execute_entry_point_by_index(0) + .expect("indexed scope execution"), + Value::Undefined + ); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::ExternalScope + )); + Ok(()) + }) + .expect("external scope succeeds"); + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + #[allow(clippy::expect_used)] + #[test] + fn evaluation_memory_budget_scope_rejects_nesting_and_cleans_up_terminal_paths() { + let mut vm = RegoVM::new(); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1024 * 1024).unwrap_or(NonZeroU64::MIN), + })); + + vm.with_evaluation_memory_budget(|vm| { + assert!(matches!( + vm.with_evaluation_memory_budget(|_| Ok(())), + Err(VmError::MemoryBudgetScopeAlreadyActive { .. }) + )); + assert_eq!(vm.execute()?, Value::Undefined); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::ExternalScope + )); + Ok(()) + }) + .expect("outer scope succeeds"); + + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + + assert!(matches!( + vm.with_evaluation_memory_budget(|vm| { + let result = vm.execute_entry_point_by_index(0); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::ExternalScope + )); + result + }), + Err(VmError::InvalidEntryPointIndex { .. }) + )); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + + let mut error_program = Program::new(); + error_program.instructions = vec![Instruction::Return { value: 0 }]; + error_program.instruction_spans = vec![None]; + vm.load_program(Arc::new(error_program)); + vm.set_max_instructions(0); + assert!(matches!( + vm.with_evaluation_memory_budget(|vm| { + let result = vm.execute(); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::ExternalScope + )); + result + }), + Err(VmError::InstructionLimitExceeded { .. }) + )); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + vm.set_max_instructions(25_000); + + vm.set_execution_mode(super::super::execution_model::ExecutionMode::Suspendable); + assert!(matches!( + vm.with_evaluation_memory_budget(|vm| { + let result = vm.execute(); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::ExternalScope + )); + result + }), + Err(VmError::MemoryBudgetUnsupportedInSuspendableExecution { .. }) + )); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + vm.set_execution_mode(super::super::execution_model::ExecutionMode::RunToCompletion); + + let closure_error = VmError::Internal { + message: "closure failure".into(), + pc: 0, + }; + assert_eq!( + vm.with_evaluation_memory_budget(|_| Err::<(), _>(closure_error.clone())), + Err(closure_error) + ); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + assert_eq!( + vm.with_evaluation_memory_budget(|vm| vm.execute()) + .expect("fresh scope after closure failure"), + Value::Undefined + ); + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + #[allow(clippy::panic)] + #[test] + fn evaluation_memory_budget_scope_cleans_up_during_unwind() { + let mut vm = RegoVM::new(); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1024 * 1024).unwrap_or(NonZeroU64::MIN), + })); + + let unwind = std::panic::catch_unwind(core::panic::AssertUnwindSafe(|| { + let _ = vm.with_evaluation_memory_budget(|_| -> core::result::Result<(), VmError> { + panic!("scope panic for cleanup regression"); + }); + })); + + assert!(unwind.is_err()); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + assert_eq!( + vm.with_evaluation_memory_budget(|vm| vm.execute()), + Ok(Value::Undefined) + ); + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + #[test] + fn scoped_budget_configuration_captures_a_fresh_baseline() { + let tiny_budget = MemoryBudgetConfig { + limit: NonZeroU64::new(1).unwrap_or(NonZeroU64::MIN), + }; + let relaxed_budget = MemoryBudgetConfig { + limit: NonZeroU64::new(1024 * 1024).unwrap_or(NonZeroU64::MIN), + }; + let tightened_budget = MemoryBudgetConfig { + limit: NonZeroU64::new(64 * 1024).unwrap_or(NonZeroU64::MIN), + }; + + let mut unconfigured_vm = RegoVM::new(); + let configured_from_none = unconfigured_vm.with_evaluation_memory_budget(|vm| { + vm.set_memory_budget_config(Some(tiny_budget)); + vm.check_memory_budget()?; + + vm.set_memory_budget_config(None); + let allocation = alloc::vec![0_u8; 128 * 1024]; + core::hint::black_box(&allocation); + + vm.set_memory_budget_config(Some(tiny_budget)); + vm.check_memory_budget() + }); + assert!(matches!(configured_from_none, Ok(()))); + + let mut configured_vm = RegoVM::new(); + configured_vm.set_memory_budget_config(Some(relaxed_budget)); + let reconfigured = configured_vm.with_evaluation_memory_budget(|vm| { + let allocation = alloc::vec![0_u8; 128 * 1024]; + core::hint::black_box(&allocation); + + vm.set_memory_budget_config(Some(tightened_budget)); + vm.check_memory_budget() + }); + assert!(matches!( + reconfigured, + Err(VmError::MemoryBudgetExceeded { .. }) + )); + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + #[test] + fn external_scope_counts_fresh_initialization_after_releasing_previous_state() { + let rule_info = RuleInfo::new( + "unused".into(), + RuleType::Complete, + crate::Rc::new(alloc::vec![]), + 0, + 0, + ); + let mut program = Program::new(); + program.dispatch_window_size = 1; + program.max_rule_window_size = 1; + program.entry_points.insert("main".into(), 0); + program.rule_infos = alloc::vec![rule_info; Program::MAX_RULES]; + program.instructions = alloc::vec![Instruction::Return { value: 0 }]; + program.instruction_spans = alloc::vec![None]; + + let mut vm = RegoVM::new(); + vm.load_program(Arc::new(program)); + // This models a reused VM whose completed result is released before a fresh + // execution-state allocation. Keeping the previous cache empty makes the + // fresh rule-cache allocation observable at the next VM checkpoint. + vm.rule_cache = alloc::vec![]; + let retained = Value::from( + (0..(Program::MAX_RULES * 4)) + .map(Value::from) + .collect::>(), + ); + vm.evaluated = retained.clone(); + vm.execution_state = ExecutionState::Completed { result: retained }; + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(32 * 1024).unwrap_or(NonZeroU64::MIN), + })); + + let result = vm.with_evaluation_memory_budget(|vm| vm.execute()); + + assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); + assert!(matches!( + vm.execution_state, + ExecutionState::Error { + error: VmError::MemoryBudgetExceeded { .. } + } + )); + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + #[test] + fn scoped_set_data_is_atomic_when_the_budget_is_exhausted() { + let mut vm = RegoVM::new(); + let original = Value::from("original"); + vm.data = original.clone(); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1024).unwrap_or(NonZeroU64::MIN), + })); + + let result = vm.with_evaluation_memory_budget(|vm| { + let candidate = + Value::from((0..16_384).map(Value::from).collect::>()); + vm.set_data(candidate) + }); + + assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); + assert_eq!(vm.data, original); + assert!(matches!( + vm.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + )); + } +} diff --git a/tests/memory_limits.rs b/tests/memory_limits.rs index 4afa09f22..1e69d1e70 100644 --- a/tests/memory_limits.rs +++ b/tests/memory_limits.rs @@ -438,12 +438,20 @@ fn vm_memory_budget_is_enforced_for_named_and_indexed_entry_points() { vm.set_data(engine.get_data()).expect("set data"); vm.set_memory_budget_config(Some(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); - let result = if execute_by_name { - vm.execute_entry_point_by_name(entrypoint.as_ref()) - } else { - vm.execute_entry_point_by_index(0) - }; - assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); + for _ in 0..2 { + let result = if execute_by_name { + vm.execute_entry_point_by_name(entrypoint.as_ref()) + } else { + vm.execute_entry_point_by_index(0) + }; + assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); + assert!(matches!( + vm.execution_state(), + ExecutionState::Error { + error: VmError::MemoryBudgetExceeded { .. } + } + )); + } } } @@ -483,6 +491,110 @@ fn vm_memory_budget_is_fresh_for_each_execution() { ); } +#[cfg(feature = "rvm")] +#[test] +fn vm_evaluation_memory_budget_scope_preserves_its_baseline_across_execution() { + let _guard = LimitGuard::lock(); + let mut vm = RegoVM::new(); + vm.set_memory_budget_config(Some(memory_budget(64 * 1024))); + + let result = vm.with_evaluation_memory_budget(|vm| { + let allocation = vec![0_u8; 128 * 1024]; + core::hint::black_box(&allocation); + vm.execute() + }); + + assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); + + // A completed scope must not leave a stale baseline on a reusable VM. + assert_eq!( + vm.with_evaluation_memory_budget(|vm| vm.execute()) + .expect("fresh scope after exhaustion"), + Value::Undefined + ); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_evaluation_memory_budget_charges_scoped_data_and_execution_cumulatively() { + let _guard = LimitGuard::lock(); + const BUDGET_BYTES: u64 = 832 * 1024; + + let mut engine = new_engine_with_module(LARGE_PARSE_MODULE); + let entrypoint = Rc::from("data.limit.large_array"); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy for VM"); + let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) + .expect("compile VM program"); + + let mut ordinary_vm = RegoVM::new(); + ordinary_vm.load_program(program.clone()); + ordinary_vm + .set_data(large_json_data(20_000)) + .expect("set static data"); + ordinary_vm.set_input(Value::Undefined); + ordinary_vm.set_memory_budget_config(Some(memory_budget(BUDGET_BYTES))); + assert!(matches!( + ordinary_vm + .execute_entry_point_by_name(entrypoint.as_ref()) + .expect("execution-only budget permits parsing"), + Value::Array(_) + )); + + let mut vm = RegoVM::new(); + vm.load_program(program); + vm.set_input(Value::Undefined); + vm.set_memory_budget_config(Some(memory_budget(BUDGET_BYTES))); + + let result = vm.with_evaluation_memory_budget(|vm| { + vm.set_data(large_json_data(20_000))?; + vm.execute_entry_point_by_name(entrypoint.as_ref()) + }); + + assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_evaluation_memory_budget_scope_without_configuration_is_behavioral() { + let mut vm = RegoVM::new(); + + vm.with_evaluation_memory_budget(|vm| { + assert!(matches!( + vm.with_evaluation_memory_budget(|_| Ok(())), + Err(VmError::MemoryBudgetScopeAlreadyActive { .. }) + )); + assert_eq!(vm.execute()?, Value::Undefined); + Ok(()) + }) + .expect("unconfigured scope succeeds"); + + assert_eq!( + vm.with_evaluation_memory_budget(|vm| vm.execute()) + .expect("unconfigured scope is reusable"), + Value::Undefined + ); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_excludes_static_data_before_ordinary_execution() { + let _guard = LimitGuard::lock(); + let mut vm = RegoVM::new(); + vm.set_data(Value::from( + (0..100_000).map(Value::from).collect::>(), + )) + .expect("set static data"); + vm.set_memory_budget_config(Some(memory_budget(16 * 1024))); + + assert_eq!( + vm.execute() + .expect("static data is outside ordinary budget"), + Value::Undefined + ); +} + #[cfg(feature = "rvm")] #[test] fn vm_memory_budget_does_not_receive_credit_from_previous_results() { From e5da64781f6be9116cfa3bebf40456e925ccc820 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Wed, 26 Aug 2026 10:56:22 +0200 Subject: [PATCH 5/5] fix(rvm): restore execute-only memory budgets Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 1 - benches/rvm_benchmark.rs | 9 +- bindings/csharp/API.md | 12 +- bindings/csharp/README.md | 14 +- .../Regorus.Tests/RvmMemoryBudgetTests.cs | 159 -------- bindings/csharp/Regorus/NativeMethods.cs | 18 - bindings/csharp/Regorus/Rvm.cs | 51 --- bindings/ffi/CHANGELOG.md | 9 +- bindings/ffi/src/common.rs | 13 + bindings/ffi/src/rvm.rs | 379 ++---------------- docs/limits/memory_budget.md | 34 +- src/rvm/vm/errors.rs | 3 - src/rvm/vm/machine.rs | 260 ++++++++---- src/rvm/vm/state.rs | 249 +----------- tests/memory_limits.rs | 86 ---- 15 files changed, 265 insertions(+), 1032 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b03aad49..77d09910b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - *(rvm)* add opt-in per-execution memory budgets for run-to-completion evaluation, including typed Rust and binding errors ([#792](https://github.com/microsoft/regorus/pull/792)) -- *(rvm,ffi,csharp)* add scoped evaluation-memory budgeting via `RegoVM::with_evaluation_memory_budget` and one-call data execution helpers for charging evaluation-specific data with execution ([#792](https://github.com/microsoft/regorus/pull/792)) ## [0.11.0](https://github.com/microsoft/regorus/compare/regorus-v0.10.1...regorus-v0.11.0) - 2026-07-21 diff --git a/benches/rvm_benchmark.rs b/benches/rvm_benchmark.rs index cbcd1e610..1d12f16de 100644 --- a/benches/rvm_benchmark.rs +++ b/benches/rvm_benchmark.rs @@ -61,7 +61,7 @@ use regorus::{Engine, Rc, Value}; // hot path (memory_check, execution_timer_tick, instruction-limit compare). // --------------------------------------------------------------------------- -#[cfg(feature = "allocator-memory-limits")] +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] const MEMORY_LIMIT_BYTES: u64 = 256 * 1024 * 1024; const TIME_LIMIT: Duration = Duration::from_secs(30); const TIMER_CHECK_INTERVAL: NonZeroU32 = NonZeroU32::new(16).unwrap(); @@ -375,8 +375,11 @@ fn compile_all_programs() -> Vec { /// Apply the limits selected for one benchmark configuration. fn configure_limits(vm: &mut RegoVM, config: EvalConfig) { + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + let _ = config.memory_budget; + if config.limits { - #[cfg(feature = "allocator-memory-limits")] + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] regorus::set_global_memory_limit(Some(MEMORY_LIMIT_BYTES)); vm.set_execution_timer_config(Some(ExecutionTimerConfig { limit: TIME_LIMIT, @@ -384,7 +387,7 @@ fn configure_limits(vm: &mut RegoVM, config: EvalConfig) { })); vm.set_max_instructions(INSTRUCTION_LIMIT); } else { - #[cfg(feature = "allocator-memory-limits")] + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] regorus::set_global_memory_limit(None); vm.set_execution_timer_config(None); vm.set_max_instructions(usize::MAX); diff --git a/bindings/csharp/API.md b/bindings/csharp/API.md index a7b27833f..8c3b62b2b 100644 --- a/bindings/csharp/API.md +++ b/bindings/csharp/API.md @@ -98,14 +98,10 @@ public sealed class Rvm : IDisposable { public void SetMemoryBudgetConfig(MemoryBudgetConfig config); public void ClearMemoryBudgetConfig(); - - public string? ExecuteWithDataJson(string dataJson); - public string? ExecuteEntryPointWithDataJson(string entryPoint, string dataJson); - public string? ExecuteEntryPointWithDataJson(ulong index, string dataJson); } ``` -Ordinary `Execute` and `ExecuteEntryPoint` calls use the existing VM data and start a fresh budget for execution. Program compilation, program loading, and prior `SetDataJson`, `SetInputJson`, and `SetContextJson` calls are excluded; use this path for static or preloaded data. +Ordinary `Execute` and `ExecuteEntryPoint` calls use the existing VM data and start a fresh budget for execution. Program compilation, program loading, and prior `SetDataJson`, `SetInputJson`, and `SetContextJson` calls occur before and outside the budget; use this path for static or preloaded data. ```csharp using var vm = new Rvm(); @@ -121,11 +117,9 @@ catch (RegorusMemoryBudgetExceededException) } ``` -Use `ExecuteWithDataJson`, `ExecuteEntryPointWithDataJson(string, string)`, or `ExecuteEntryPointWithDataJson(ulong, string)` when evaluation-specific data should be charged with execution. These methods make one native call that includes native JSON parsing/storage for `dataJson`, execution, native JSON serialization, and native C-string allocation. Managed UTF-8 decoding and the managed C# `string` allocation after the native call returns are not charged. - -If scoped data replacement fails, the VM keeps its previous data. The previous and provisional native data can coexist transiently, so both count toward the peak live bytes observed by the budget. +Native result JSON serialization and C-string allocation are included before an `Execute` or `ExecuteEntryPoint` call returns. Managed UTF-8 decoding and the managed C# `string` allocation after the native call returns are excluded. -Memory budgets require a native library built with allocator memory tracking and are supported only for run-to-completion execution. `RegorusMemoryBudgetExceededException` is thrown when a budget is exceeded. `RegorusMemoryBudgetUnsupportedException` is thrown if a configured budget is used to start or resume suspendable execution. Enforcement is cooperative, so one instruction can overshoot before the next checkpoint. Same-thread baseline ratcheting can make the effective limit stricter after unrelated frees are observed; those frees are never credited back. Public multi-call begin/end scopes are intentionally absent because allocator counters are thread-local and abandoned or cross-thread scopes would be unsafe. +Memory budgets require a native library built with allocator memory tracking and are supported only for run-to-completion execution. `RegorusMemoryBudgetExceededException` is thrown when a budget is exceeded. `RegorusMemoryBudgetUnsupportedException` is thrown if a configured budget is used to start or resume suspendable execution. Enforcement is cooperative, so one instruction can overshoot before the next checkpoint. Same-thread baseline ratcheting can make the effective limit stricter after unrelated frees are observed; those frees are never credited back. Failed terminal execution clears retained state, and a reused VM gets a fresh budget. Public multi-call begin/end scopes are intentionally absent because allocator counters are thread-local and abandoned or cross-thread scopes would be unsafe. ## Core Classes diff --git a/bindings/csharp/README.md b/bindings/csharp/README.md index 36c67def1..55807176a 100644 --- a/bindings/csharp/README.md +++ b/bindings/csharp/README.md @@ -107,7 +107,7 @@ Console.WriteLine($"allow: {result}"); ### Per-execution memory budget -RVM run-to-completion evaluation can use an optional additional live-memory budget. Each ordinary `Execute` or `ExecuteEntryPoint` call starts with a fresh budget for execution; program compilation, program loading, and prior `SetDataJson`, `SetInputJson`, and `SetContextJson` calls are not charged. +RVM run-to-completion evaluation can use an optional additional live-memory budget. Each ordinary `Execute` or `ExecuteEntryPoint` call starts with a fresh budget for execution; program compilation, program loading, and prior `SetDataJson`, `SetInputJson`, and `SetContextJson` calls occur before and outside that budget. ```csharp using var vm = new Rvm(); @@ -126,17 +126,9 @@ catch (RegorusMemoryBudgetExceededException ex) } ``` -Use the one-call `Execute*WithDataJson` methods when evaluation-specific data should be charged in the same native budget scope as execution: +The native execution budget includes result JSON serialization and C-string allocation before the native call returns. Managed UTF-8 decoding and C# `string` allocation after that return are not charged. -```csharp -var result = vm.ExecuteWithDataJson(Data); -var named = vm.ExecuteEntryPointWithDataJson("data.demo.allow", Data); -var indexed = vm.ExecuteEntryPointWithDataJson(0UL, Data); -``` - -Those methods include native JSON parsing/storage for `dataJson`, execution, native result JSON serialization, and native C-string allocation. Managed UTF-8 decoding and C# `string` allocation after the native call returns are not charged. If scoped data replacement fails, the previous VM data is preserved, but the previous and provisional data may coexist transiently and count toward peak live bytes. - -The budget is cooperative and may overshoot between VM checks. Same-thread allocation-counter baseline ratcheting can make the effective limit stricter after unrelated frees are observed; those frees are not credited back. Budgets are not supported in suspendable execution mode. `ClearMemoryBudgetConfig` restores the previous unlimited per-execution behavior. Public multi-call begin/end scopes are intentionally absent because allocator counters are thread-local. The process-wide limit exposed by `MemoryLimits` remains a separate safeguard. +The budget is cooperative and may overshoot between VM checks. Same-thread allocation-counter baseline ratcheting can make the effective limit stricter after unrelated frees are observed; those frees are not credited back. Budgets are not supported in suspendable execution mode. `ClearMemoryBudgetConfig` restores the previous unlimited per-execution behavior. Public multi-call begin/end scopes are intentionally absent because allocator counters are thread-local. Failed terminal execution clears retained state, and a reused VM starts a fresh budget. The process-wide limit exposed by `MemoryLimits` remains a separate safeguard. ## Azure RBAC Condition Evaluation diff --git a/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs index 09c6f99a5..34c34ce72 100644 --- a/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs +++ b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs @@ -30,17 +30,6 @@ package limits.memory private const string PreloadedResultEntryPoint = "data.limits.memory.large_string"; - private const string DataResultPolicy = """ -package limits.memory - -main := data.value -named := data.value -"""; - - private const string MainEntryPoint = "data.limits.memory.main"; - - private const string NamedEntryPoint = "data.limits.memory.named"; - [TestMethod] public void Memory_budget_must_be_non_zero() { @@ -96,122 +85,6 @@ public void Serialization_budget_failure_leaves_error_state() StringAssert.Contains(state, "Error { error: MemoryBudgetExceeded"); } - [TestMethod] - public void Execute_with_data_json_runs_main_entry_point() - { - using var program = CreateDataResultProgram(); - using var vm = CreateRvm(program); - - var result = vm.ExecuteWithDataJson(CreateValueData("one-call")); - - Assert.AreEqual("\"one-call\"", result); - } - - [TestMethod] - public void Execute_entry_point_with_data_json_runs_named_and_indexed_entry_points() - { - using var program = CreateDataResultProgram(); - using var vm = CreateRvm(program); - - Assert.AreEqual("\"named\"", vm.ExecuteEntryPointWithDataJson(NamedEntryPoint, CreateValueData("named"))); - Assert.AreEqual("\"indexed\"", vm.ExecuteEntryPointWithDataJson(0, CreateValueData("indexed"))); - } - - [TestMethod] - public void Execute_with_data_json_exceeding_data_budget_throws_typed_exception() - { - using var program = CreateDataResultProgram(); - using var vm = CreateRvm(program); - vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); - - Assert.ThrowsException( - () => vm.ExecuteWithDataJson(CreateOversizedValueData())); - } - - [TestMethod] - public void Execute_with_data_json_failed_replacement_preserves_previous_vm_data() - { - using var program = CreateDataResultProgram(); - using var vm = CreateRvm(program); - vm.SetDataJson(CreateValueData("previous")); - vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); - - Assert.ThrowsException( - () => vm.ExecuteWithDataJson(CreateOversizedValueData())); - - var state = vm.GetExecutionState(); - Assert.IsNotNull(state); - StringAssert.Contains(state, "Error { error: MemoryBudgetExceeded"); - - vm.ClearMemoryBudgetConfig(); - - Assert.AreEqual("\"previous\"", vm.ExecuteEntryPoint(MainEntryPoint)); - } - - [TestMethod] - public void Execute_with_malformed_data_json_leaves_error_state_and_preserves_previous_vm_data() - { - using var program = CreateDataResultProgram(); - using var vm = CreateRvm(program); - vm.SetDataJson(CreateValueData("previous")); - - Assert.ThrowsException(() => vm.ExecuteWithDataJson("{")); - - var state = vm.GetExecutionState(); - Assert.IsNotNull(state); - StringAssert.Contains(state, "Error { error:"); - - Assert.AreEqual("\"previous\"", vm.ExecuteEntryPoint(MainEntryPoint)); - } - - [TestMethod] - public void Execute_with_data_json_reuses_vm_with_a_fresh_budget_after_success_and_failure() - { - using var program = CreateDataResultProgram(); - using var vm = CreateRvm(program); - vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); - - Assert.AreEqual("\"first\"", vm.ExecuteWithDataJson(CreateValueData("first"))); - Assert.ThrowsException( - () => vm.ExecuteWithDataJson(CreateOversizedValueData())); - Assert.AreEqual("\"second\"", vm.ExecuteWithDataJson(CreateValueData("second"))); - } - - [TestMethod] - public void Execute_with_data_json_uses_unlimited_default_and_cleared_budget() - { - using var program = CreateDataResultProgram(); - using var vm = CreateRvm(program); - - Assert.AreEqual("\"default\"", vm.ExecuteWithDataJson(CreateValueData("default"))); - - vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); - Assert.ThrowsException( - () => vm.ExecuteWithDataJson(CreateOversizedValueData())); - - vm.ClearMemoryBudgetConfig(); - - var result = vm.ExecuteWithDataJson(CreateOversizedValueData()); - Assert.IsFalse(string.IsNullOrWhiteSpace(result)); - } - - [TestMethod] - public void Execute_with_data_json_counts_native_result_serialization_but_not_managed_string_copy() - { - using var program = CreateLargeResultProgram(); - using var vm = CreateRvm(program); - vm.SetInputJson(CreateLargeResultInput()); - vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(512 * 1024)); - - Assert.ThrowsException(() => vm.ExecuteWithDataJson("{}")); - - vm.ClearMemoryBudgetConfig(); - - var result = vm.ExecuteWithDataJson("{}"); - Assert.IsNotNull(result); - Assert.AreEqual((2 * 1024 * 1024) + 2, result!.Length); - } - [TestMethod] public void Suspendable_execution_rejects_memory_budget() { @@ -222,38 +95,6 @@ public void Suspendable_execution_rejects_memory_budget() Assert.ThrowsException(() => vm.Execute()); } - private static Program CreateDataResultProgram() - { - var modules = new[] { new PolicyModule("data_result.rego", DataResultPolicy) }; - return Program.CompileFromModules("{}", modules, new[] { MainEntryPoint, NamedEntryPoint }); - } - - private static Program CreateLargeResultProgram() - { - const string policy = """ -package limits.memory - -large_string := input.large_string -"""; - var modules = new[] { new PolicyModule("large_result.rego", policy) }; - return Program.CompileFromModules("{}", modules, new[] { "data.limits.memory.large_string" }); - } - - private static string CreateLargeResultInput() - { - return JsonSerializer.Serialize(new { large_string = new string('x', 2 * 1024 * 1024) }); - } - - private static string CreateValueData(string value) - { - return JsonSerializer.Serialize(new { value }); - } - - private static string CreateOversizedValueData() - { - return CreateValueData(new string('x', 2 * 1024 * 1024)); - } - private static Program CreateProgram() { var modules = new[] { new PolicyModule("memory_budget.rego", Policy) }; diff --git a/bindings/csharp/Regorus/NativeMethods.cs b/bindings/csharp/Regorus/NativeMethods.cs index b07f82f96..479c9b7c9 100644 --- a/bindings/csharp/Regorus/NativeMethods.cs +++ b/bindings/csharp/Regorus/NativeMethods.cs @@ -204,24 +204,6 @@ internal static unsafe partial class API [DllImport(LibraryName, EntryPoint = "regorus_rvm_execute_entry_point_by_index", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] internal static extern RegorusResult regorus_rvm_execute_entry_point_by_index(RegorusRvm* vm, UIntPtr index); - /// - /// Set the data document and execute the program in one native call. - /// - [DllImport(LibraryName, EntryPoint = "regorus_rvm_execute_with_data", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] - internal static extern RegorusResult regorus_rvm_execute_with_data(RegorusRvm* vm, byte* data_json); - - /// - /// Set the data document and execute a named entry point in one native call. - /// - [DllImport(LibraryName, EntryPoint = "regorus_rvm_execute_entry_point_by_name_with_data", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] - internal static extern RegorusResult regorus_rvm_execute_entry_point_by_name_with_data(RegorusRvm* vm, byte* entry_point, byte* data_json); - - /// - /// Set the data document and execute an entry point by index in one native call. - /// - [DllImport(LibraryName, EntryPoint = "regorus_rvm_execute_entry_point_by_index_with_data", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)] - internal static extern RegorusResult regorus_rvm_execute_entry_point_by_index_with_data(RegorusRvm* vm, UIntPtr index, byte* data_json); - /// /// Resume execution. /// diff --git a/bindings/csharp/Regorus/Rvm.cs b/bindings/csharp/Regorus/Rvm.cs index e6e0c3e29..884731e58 100644 --- a/bindings/csharp/Regorus/Rvm.cs +++ b/bindings/csharp/Regorus/Rvm.cs @@ -212,57 +212,6 @@ public void ClearMemoryBudgetConfig() }); } - /// - /// Set evaluation-specific data JSON and execute the program in one native call. - /// - public string? ExecuteWithDataJson(string dataJson) - { - return Utf8Marshaller.WithUtf8(dataJson, dataPtr => - { - return UseHandle(vmPtr => - { - return CheckAndDropResult(API.regorus_rvm_execute_with_data((RegorusRvm*)vmPtr, (byte*)dataPtr)); - }); - }); - } - - /// - /// Set evaluation-specific data JSON and execute a named entry point in one native call. - /// - public string? ExecuteEntryPointWithDataJson(string entryPoint, string dataJson) - { - return Utf8Marshaller.WithUtf8(entryPoint, entryPtr => - { - return Utf8Marshaller.WithUtf8(dataJson, dataPtr => - { - return UseHandle(vmPtr => - { - return CheckAndDropResult(API.regorus_rvm_execute_entry_point_by_name_with_data( - (RegorusRvm*)vmPtr, - (byte*)entryPtr, - (byte*)dataPtr)); - }); - }); - }); - } - - /// - /// Set evaluation-specific data JSON and execute an entry point by index in one native call. - /// - public string? ExecuteEntryPointWithDataJson(ulong index, string dataJson) - { - return Utf8Marshaller.WithUtf8(dataJson, dataPtr => - { - return UseHandle(vmPtr => - { - return CheckAndDropResult(API.regorus_rvm_execute_entry_point_by_index_with_data( - (RegorusRvm*)vmPtr, - (UIntPtr)index, - (byte*)dataPtr)); - }); - }); - } - /// /// Resume execution with an optional value. /// diff --git a/bindings/ffi/CHANGELOG.md b/bindings/ffi/CHANGELOG.md index 04ab70683..a100511db 100644 --- a/bindings/ffi/CHANGELOG.md +++ b/bindings/ffi/CHANGELOG.md @@ -8,11 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add `regorus_rvm_set_memory_budget_config`, `RegorusMemoryBudgetConfig`, and appended statuses for memory-budget exhaustion and unsupported suspendable execution. -- Add one-call RVM data execution APIs: `regorus_rvm_execute_with_data`, - `regorus_rvm_execute_entry_point_by_name_with_data`, and - `regorus_rvm_execute_entry_point_by_index_with_data`. Their memory-budget scope covers data - preparation, execution, and native result production. +- Add execute-only RVM memory-budget configuration through + `regorus_rvm_set_memory_budget_config`, `RegorusMemoryBudgetConfig`, and appended statuses for + exhaustion and unsupported suspendable execution. Native execute result serialization and C-string + allocation are included in the configured budget. ## [0.1.0](https://github.com/microsoft/regorus/releases/tag/regorus-ffi-v0.1.0) - 2024-02-08 diff --git a/bindings/ffi/src/common.rs b/bindings/ffi/src/common.rs index cf0b19c42..53d37648d 100644 --- a/bindings/ffi/src/common.rs +++ b/bindings/ffi/src/common.rs @@ -139,6 +139,19 @@ impl RegorusResult { } } + /// Create a successful result from an already allocated C string. + pub(crate) fn ok_c_string(output: CString) -> Self { + Self { + status: RegorusStatus::Ok, + data_type: RegorusDataType::String, + output: output.into_raw(), + bool_value: false, + int_value: 0, + pointer_value: ptr::null_mut(), + error_message: ptr::null_mut(), + } + } + /// Create a successful result with boolean value. #[allow(unused)] pub(crate) fn ok_bool(value: bool) -> Self { diff --git a/bindings/ffi/src/rvm.rs b/bindings/ffi/src/rvm.rs index f0a6f2a7a..bc14da916 100644 --- a/bindings/ffi/src/rvm.rs +++ b/bindings/ffi/src/rvm.rs @@ -1,8 +1,6 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -#[cfg(all(feature = "allocator-memory-limits", not(miri)))] -use crate::common::regorus_result_drop; use crate::common::{ from_c_str, to_ref, to_regorus_result, to_shared_ref, RegorusBuffer, RegorusResult, RegorusStatus, @@ -82,116 +80,21 @@ enum RvmExecution { Indexed(usize), } -/// Owns a provisional native result until the surrounding budget scope succeeds. -#[cfg(all(feature = "allocator-memory-limits", not(miri)))] -struct ProvisionalRvmResult { - result: RegorusResult, - owns_result: bool, -} - -#[cfg(all(feature = "allocator-memory-limits", not(miri)))] -impl ProvisionalRvmResult { - fn new(result: RegorusResult) -> Self { - Self { - result, - owns_result: true, - } - } - - fn into_result(mut self) -> RegorusResult { - self.owns_result = false; - core::mem::replace(&mut self.result, RegorusResult::ok_void()) - } -} - -#[cfg(all(feature = "allocator-memory-limits", not(miri)))] -impl Drop for ProvisionalRvmResult { - fn drop(&mut self) { - if self.owns_result { - #[cfg(all(test, feature = "std"))] - PROVISIONAL_RVM_RESULT_DROP_COUNT.with(|count| { - count.set(count.get().saturating_add(1)); - }); - regorus_result_drop(core::mem::replace( - &mut self.result, - RegorusResult::ok_void(), - )); - } - } -} - -#[cfg(all(test, feature = "allocator-memory-limits", not(miri), feature = "std"))] -std::thread_local! { - static PROVISIONAL_RVM_RESULT_DROP_COUNT: std::cell::Cell = - const { std::cell::Cell::new(0) }; -} - -#[cfg(all(test, feature = "allocator-memory-limits", not(miri), feature = "std"))] -fn reset_provisional_rvm_result_drop_count() { - PROVISIONAL_RVM_RESULT_DROP_COUNT.with(|count| count.set(0)); -} - -#[cfg(all(test, feature = "allocator-memory-limits", not(miri), feature = "std"))] -fn provisional_rvm_result_drop_count() -> usize { - PROVISIONAL_RVM_RESULT_DROP_COUNT.with(std::cell::Cell::get) -} - -fn prepare_data( - vm: &mut RegoVM, - data_json: Option<*const c_char>, -) -> core::result::Result<(), VmError> { - if let Some(data_json) = data_json { - let data = Value::from_json_str(&from_c_str(data_json)?).map_err(VmError::from)?; - vm.set_data(data)?; - } - - Ok(()) -} - -fn execute_rvm(vm: &mut RegoVM, execution: RvmExecution) -> core::result::Result { - match execution { - RvmExecution::Main => vm.execute(), - RvmExecution::Named(entry_point) => vm.execute_entry_point_by_name(&entry_point), - RvmExecution::Indexed(index) => vm.execute_entry_point_by_index(index), - } -} - -fn execute_to_rvm_result( - vm: *mut RegorusRvm, - data_json: Option<*const c_char>, - execution: RvmExecution, -) -> RegorusResult { +fn execute_to_rvm_result(vm: *mut RegorusRvm, execution: RvmExecution) -> RegorusResult { let output = || -> Result { let vm = to_shared_ref(vm as *const RegorusRvm)?; let mut guard = vm.try_write()?; - #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - { - let result = guard.with_evaluation_memory_budget(|guard| { - prepare_data(guard, data_json)?; - let value = execute_rvm(guard, execution)?; - let json = value.to_json_str().map_err(VmError::from)?; - let result = ProvisionalRvmResult::new(RegorusResult::ok_string(json)); - - // This remains separate from the scope's final check so the - // provisional output is dropped before an error crosses FFI. - guard.check_memory_budget()?; - Ok(result) - }); - - match result { - Ok(result) => Ok(result.into_result()), - Err(err) => Err(anyhow::Error::from(guard.fail_run_to_completion(err))), + let output = match execution { + RvmExecution::Main => guard.execute_to_c_string_for_ffi()?, + RvmExecution::Named(entry_point) => { + guard.execute_entry_point_by_name_to_c_string_for_ffi(&entry_point)? } - } - - #[cfg(any(not(feature = "allocator-memory-limits"), miri))] - { - prepare_data(&mut guard, data_json)?; - let value = execute_rvm(&mut guard, execution)?; - let json = value.to_json_str()?; - Ok(RegorusResult::ok_string(json)) - } + RvmExecution::Indexed(index) => { + guard.execute_entry_point_by_index_to_c_string_for_ffi(index)? + } + }; + Ok(RegorusResult::ok_c_string(output)) }(); match output { @@ -672,7 +575,7 @@ pub extern "C" fn regorus_rvm_set_memory_budget_config( /// Execute the program's main entry point. #[no_mangle] pub extern "C" fn regorus_rvm_execute(vm: *mut RegorusRvm) -> RegorusResult { - with_unwind_guard(|| execute_to_rvm_result(vm, None, RvmExecution::Main)) + with_unwind_guard(|| execute_to_rvm_result(vm, RvmExecution::Main)) } /// Execute a named entry point. @@ -686,7 +589,7 @@ pub extern "C" fn regorus_rvm_execute_entry_point_by_name( Ok(entry_point) => entry_point, Err(err) => return to_rvm_error_result(err), }; - execute_to_rvm_result(vm, None, RvmExecution::Named(entry_point)) + execute_to_rvm_result(vm, RvmExecution::Named(entry_point)) }) } @@ -696,42 +599,7 @@ pub extern "C" fn regorus_rvm_execute_entry_point_by_index( vm: *mut RegorusRvm, index: usize, ) -> RegorusResult { - with_unwind_guard(|| execute_to_rvm_result(vm, None, RvmExecution::Indexed(index))) -} - -/// Set the VM data from JSON and execute the program's main entry point in one call. -#[no_mangle] -pub extern "C" fn regorus_rvm_execute_with_data( - vm: *mut RegorusRvm, - data_json: *const c_char, -) -> RegorusResult { - with_unwind_guard(|| execute_to_rvm_result(vm, Some(data_json), RvmExecution::Main)) -} - -/// Set the VM data from JSON and execute a named entry point in one call. -#[no_mangle] -pub extern "C" fn regorus_rvm_execute_entry_point_by_name_with_data( - vm: *mut RegorusRvm, - entry_point: *const c_char, - data_json: *const c_char, -) -> RegorusResult { - with_unwind_guard(|| { - let entry_point = match from_c_str(entry_point) { - Ok(entry_point) => entry_point, - Err(err) => return to_rvm_error_result(err), - }; - execute_to_rvm_result(vm, Some(data_json), RvmExecution::Named(entry_point)) - }) -} - -/// Set the VM data from JSON and execute an entry point by index in one call. -#[no_mangle] -pub extern "C" fn regorus_rvm_execute_entry_point_by_index_with_data( - vm: *mut RegorusRvm, - index: usize, - data_json: *const c_char, -) -> RegorusResult { - with_unwind_guard(|| execute_to_rvm_result(vm, Some(data_json), RvmExecution::Indexed(index))) + with_unwind_guard(|| execute_to_rvm_result(vm, RvmExecution::Indexed(index))) } /// Resume execution for suspendable runs. @@ -780,9 +648,7 @@ pub extern "C" fn regorus_rvm_get_execution_state(vm: *mut RegorusRvm) -> Regoru mod tests { use super::{ regorus_rvm_drop, regorus_rvm_execute, regorus_rvm_execute_entry_point_by_index, - regorus_rvm_execute_entry_point_by_index_with_data, - regorus_rvm_execute_entry_point_by_name, regorus_rvm_execute_entry_point_by_name_with_data, - regorus_rvm_execute_with_data, regorus_rvm_get_execution_state, regorus_rvm_new, + regorus_rvm_execute_entry_point_by_name, regorus_rvm_get_execution_state, regorus_rvm_new, regorus_rvm_resume, regorus_rvm_set_data, regorus_rvm_set_memory_budget_config, RegorusRvm, }; use crate::common::{regorus_result_drop, RegorusResult, RegorusStatus}; @@ -823,10 +689,6 @@ copy := [value | some value in input] assert_execution_state(vm, "Error { error: MemoryBudgetExceeded"); } - fn assert_execution_state_error(vm: *mut RegorusRvm) { - assert_execution_state(vm, "Error { error:"); - } - fn assert_execution_state(vm: *mut RegorusRvm, expected_prefix: &str) { let state = regorus_rvm_get_execution_state(vm); assert!(matches!(state.status, RegorusStatus::Ok)); @@ -883,67 +745,6 @@ copy := [value | some value in input] Arc::new(program) } - fn data_result_program() -> Arc { - let mut program = Program::new(); - program.dispatch_window_size = 1; - program.max_rule_window_size = 1; - program.entry_points.insert("main".to_string(), 0); - program.entry_points.insert("named".to_string(), 0); - program.instructions = vec![ - Instruction::LoadData { dest: 0 }, - Instruction::Return { value: 0 }, - ]; - program.instruction_spans = vec![None; program.instructions.len()]; - Arc::new(program) - } - - fn result_json(result: RegorusResult) -> String { - if !matches!(result.status, RegorusStatus::Ok) { - let message = if result.error_message.is_null() { - "" - } else { - unsafe { CStr::from_ptr(result.error_message) } - .to_str() - .unwrap_or("") - }; - panic!( - "unexpected RVM result status {:?}: {message}", - result.status - ); - } - let output = unsafe { CStr::from_ptr(result.output) } - .to_str() - .expect("result UTF-8") - .to_string(); - regorus_result_drop(result); - output - } - - fn assert_result_json_eq(result: RegorusResult, expected: &str) { - assert_eq!( - Value::from_json_str(&result_json(result)).expect("result JSON"), - Value::from_json_str(expected).expect("expected JSON") - ); - } - - #[cfg(feature = "std")] - #[test] - fn provisional_rvm_result_transfers_or_drops_ownership() { - super::reset_provisional_rvm_result_drop_count(); - { - let _provisional = - super::ProvisionalRvmResult::new(RegorusResult::ok_string("dropped".into())); - } - assert_eq!(super::provisional_rvm_result_drop_count(), 1); - - super::reset_provisional_rvm_result_drop_count(); - let result = - super::ProvisionalRvmResult::new(RegorusResult::ok_string("transferred".into())) - .into_result(); - assert_eq!(super::provisional_rvm_result_drop_count(), 0); - regorus_result_drop(result); - } - #[test] fn ffi_memory_budget_setter_validates_and_clears_configuration() { let vm = regorus_rvm_new(); @@ -975,6 +776,31 @@ copy := [value | some value in input] regorus_rvm_drop(vm); } + #[test] + fn ffi_preloaded_data_is_outside_the_execution_budget() { + let vm = regorus_rvm_new(); + let set_budget = regorus_rvm_set_memory_budget_config( + vm, + true, + RegorusMemoryBudgetConfig { + limit_bytes: 16 * 1024, + }, + ); + assert!(matches!(set_budget.status, RegorusStatus::Ok)); + regorus_result_drop(set_budget); + + let data = CString::new(format!(r#"{{"value":"{}"}}"#, "x".repeat(2 * 1024 * 1024))) + .expect("preloaded data CString"); + let set_data = regorus_rvm_set_data(vm, data.as_ptr()); + assert!(matches!(set_data.status, RegorusStatus::Ok)); + regorus_result_drop(set_data); + + let result = regorus_rvm_execute(vm); + assert!(matches!(result.status, RegorusStatus::Ok)); + regorus_result_drop(result); + regorus_rvm_drop(vm); + } + #[test] fn ffi_execution_reports_memory_budget_status() { let entrypoint = Rc::from("data.limits.memory.copy"); @@ -1025,131 +851,6 @@ copy := [value | some value in input] regorus_rvm_execute_entry_point_by_name(vm, entrypoint.as_ptr()), ); assert_memory_budget_failure_state(vm, regorus_rvm_execute_entry_point_by_index(vm, 0)); - regorus_rvm_drop(vm); - } - - #[test] - fn ffi_execute_with_data_runs_main_named_and_indexed_entry_points() { - let data = CString::new(r#"{"value":"one-call"}"#).expect("data CString"); - let entrypoint = CString::new("named").expect("entry point CString"); - - let mut main_vm = RegoVM::new(); - main_vm.load_program(data_result_program()); - let main_vm = Box::into_raw(Box::new(RegorusRvm::new(main_vm))); - assert_result_json_eq( - regorus_rvm_execute_with_data(main_vm, data.as_ptr()), - r#"{"value":"one-call"}"#, - ); - assert_result_json_eq( - regorus_rvm_execute_with_data(main_vm, data.as_ptr()), - r#"{"value":"one-call"}"#, - ); - regorus_rvm_drop(main_vm); - - let mut named_vm = RegoVM::new(); - named_vm.load_program(data_result_program()); - let named_vm = Box::into_raw(Box::new(RegorusRvm::new(named_vm))); - assert_result_json_eq( - regorus_rvm_execute_entry_point_by_name_with_data( - named_vm, - entrypoint.as_ptr(), - data.as_ptr(), - ), - r#"{"value":"one-call"}"#, - ); - assert_result_json_eq( - regorus_rvm_execute_entry_point_by_name_with_data( - named_vm, - entrypoint.as_ptr(), - data.as_ptr(), - ), - r#"{"value":"one-call"}"#, - ); - regorus_rvm_drop(named_vm); - - let mut indexed_vm = RegoVM::new(); - indexed_vm.load_program(data_result_program()); - let indexed_vm = Box::into_raw(Box::new(RegorusRvm::new(indexed_vm))); - assert_result_json_eq( - regorus_rvm_execute_entry_point_by_index_with_data(indexed_vm, 0, data.as_ptr()), - r#"{"value":"one-call"}"#, - ); - assert_result_json_eq( - regorus_rvm_execute_entry_point_by_index_with_data(indexed_vm, 0, data.as_ptr()), - r#"{"value":"one-call"}"#, - ); - regorus_rvm_drop(indexed_vm); - } - - #[test] - fn ffi_execute_with_data_rejects_oversized_replacement_without_losing_previous_data() { - let mut vm = RegoVM::new(); - vm.load_program(data_result_program()); - let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); - let old_data = CString::new(r#"{"previous":true}"#).expect("old data CString"); - let set_data = regorus_rvm_set_data(vm, old_data.as_ptr()); - assert!(matches!(set_data.status, RegorusStatus::Ok)); - regorus_result_drop(set_data); - - let set_budget = regorus_rvm_set_memory_budget_config( - vm, - true, - RegorusMemoryBudgetConfig { - limit_bytes: TIGHT_MEMORY_BUDGET_BYTES, - }, - ); - assert!(matches!(set_budget.status, RegorusStatus::Ok)); - regorus_result_drop(set_budget); - - let oversized_data = CString::new(format!( - r#"{{"replacement":"{}"}}"#, - "x".repeat(2 * 1024 * 1024) - )) - .expect("oversized data CString"); - let rejected = regorus_rvm_execute_with_data(vm, oversized_data.as_ptr()); - assert_memory_budget_failure_state(vm, rejected); - - let clear_budget = regorus_rvm_set_memory_budget_config( - vm, - false, - RegorusMemoryBudgetConfig { limit_bytes: 0 }, - ); - assert!(matches!(clear_budget.status, RegorusStatus::Ok)); - regorus_result_drop(clear_budget); - assert_result_json_eq(regorus_rvm_execute(vm), r#"{"previous":true}"#); - regorus_rvm_drop(vm); - } - - #[test] - fn ffi_malformed_scoped_data_terminalizes_and_preserves_previous_data() { - let mut vm = RegoVM::new(); - vm.load_program(data_result_program()); - let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); - let old_data = CString::new(r#"{"previous":true}"#).expect("old data CString"); - let set_data = regorus_rvm_set_data(vm, old_data.as_ptr()); - assert!(matches!(set_data.status, RegorusStatus::Ok)); - regorus_result_drop(set_data); - - let malformed_data = CString::new("{").expect("malformed data CString"); - let rejected = regorus_rvm_execute_with_data(vm, malformed_data.as_ptr()); - assert!(matches!(rejected.status, RegorusStatus::Error)); - assert!(rejected.output.is_null()); - regorus_result_drop(rejected); - assert_execution_state_error(vm); - - assert_result_json_eq(regorus_rvm_execute(vm), r#"{"previous":true}"#); - regorus_rvm_drop(vm); - } - - #[test] - fn ffi_execute_with_data_counts_native_result_production_and_leaves_vm_reusable() { - let mut vm = RegoVM::new(); - vm.load_program(preloaded_result_program()); - vm.set_memory_budget_config(Some(memory_budget(512 * 1024))); - let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); - let data = CString::new("{}").expect("data CString"); - - assert_memory_budget_failure_state(vm, regorus_rvm_execute_with_data(vm, data.as_ptr())); let clear_budget = regorus_rvm_set_memory_budget_config( vm, @@ -1158,7 +859,7 @@ copy := [value | some value in input] ); assert!(matches!(clear_budget.status, RegorusStatus::Ok)); regorus_result_drop(clear_budget); - let result = regorus_rvm_execute_with_data(vm, data.as_ptr()); + let result = regorus_rvm_execute(vm); assert!(matches!(result.status, RegorusStatus::Ok)); regorus_result_drop(result); regorus_rvm_drop(vm); diff --git a/docs/limits/memory_budget.md b/docs/limits/memory_budget.md index 3a8ecc060..60d3d500f 100644 --- a/docs/limits/memory_budget.md +++ b/docs/limits/memory_budget.md @@ -2,7 +2,7 @@ RVM run-to-completion evaluation supports an optional memory budget when Regorus is built with the `allocator-memory-limits` feature. -The budget limits additional live bytes on the execution thread. Regorus captures a baseline when a budgeted scope starts and compares later live-byte samples with that baseline. Each ordinary Rust `execute`, `execute_entry_point_by_name`, or `execute_entry_point_by_index` call starts with a fresh execution-only budget. +The budget limits additional live bytes on the execution thread. Regorus captures a baseline when execution starts and compares later live-byte samples with that baseline. Each Rust `execute`, `execute_entry_point_by_name`, or `execute_entry_point_by_index` call starts with a fresh execution-only budget. ```rust use core::num::NonZeroU64; @@ -17,33 +17,15 @@ vm.set_memory_budget_config(Some(MemoryBudgetConfig { No configured budget preserves existing RVM behavior. A zero-byte budget is not representable in Rust and is rejected by language bindings. -## Budget scopes +## Included work -The ordinary `execute*` APIs start their budget when RVM execution begins. Fresh execution-state initialization, rule evaluation, and allocations retained by the result count against the budget. +The `execute*` APIs start their budget when RVM execution begins. Fresh execution-state initialization, rule evaluation, and allocations retained by the result count against the budget. -Program compilation, program loading, data loading, input loading, and context loading happen before the execution baseline and are not charged. +Program compilation, program loading, data loading, input loading, and context loading happen before and outside the execution baseline and are not charged. -For callers that need evaluation-specific data to be charged with the evaluation, Rust exposes `RegoVM::with_evaluation_memory_budget`. This is a synchronous same-thread scope for allocations created after the closure begins. It can include evaluation-specific data construction and assignment, execution, and caller post-processing performed inside the closure. Nested scopes are rejected, and the scope is deactivated before the method returns on success, error, or unwinding. +The C FFI keeps an internal execution window open through immediate native result JSON serialization and C-string allocation, then closes it on success, error, or unwinding. The C# binding receives that native string after the window has closed, so managed UTF-8 decoding and managed `string` allocation are excluded. -If a scope begins without a configured budget, configuring or re-enabling one inside the scope captures a fresh current-thread baseline at that transition. Changing one configured threshold to another preserves the existing scope baseline. - -The C FFI exposes one-call data execution APIs: - -- `regorus_rvm_execute_with_data` -- `regorus_rvm_execute_entry_point_by_name_with_data` -- `regorus_rvm_execute_entry_point_by_index_with_data` - -The C# binding exposes matching methods: - -- `Rvm.ExecuteWithDataJson(string dataJson)` -- `Rvm.ExecuteEntryPointWithDataJson(string entryPoint, string dataJson)` -- `Rvm.ExecuteEntryPointWithDataJson(ulong index, string dataJson)` - -These APIs run in one native call and use one budget scope covering native JSON parsing and storage for the supplied data, execution, native JSON serialization of the result, and native C-string allocation. Program loading, compilation, input loading, context loading, and managed C# UTF-8 decoding and `string` allocation remain outside the scope. - -Existing `set_data` / C# `SetDataJson` before execution remains the static/preloaded data path. That work happens before the ordinary execution budget and is excluded from it. - -If scoped data replacement fails, the previous VM data is preserved. The old data and provisional replacement can coexist transiently, so both contribute to the observed peak live bytes. +There is no public multi-call begin/end memory-budget scope. Public scopes could be abandoned or move across threads while allocator counters are thread-local. Rust, C FFI, and C# are supported by this API; other bindings require follow-up work. ## Enforcement @@ -57,7 +39,7 @@ Exhaustion returns `VmError::MemoryBudgetExceeded`, including: - configured budget - VM program counter -The VM transitions to `ExecutionState::Error` and releases values retained by the failed execution. The C FFI reports `RegorusStatus::MemoryBudgetExceeded`, including when native result serialization exceeds the budget. The C# binding throws `RegorusMemoryBudgetExceededException`. Every success or failure terminal path deactivates its owned scope, and reused VMs get a fresh budget for the next execution. +The VM transitions to `ExecutionState::Error` and releases values retained by a failed execution. The C FFI reports `RegorusStatus::MemoryBudgetExceeded`, including when native result serialization or C-string allocation exceeds the budget. The C# binding throws `RegorusMemoryBudgetExceededException`. Every terminal path clears its execution window, and reused VMs get a fresh budget for the next execution. ## Execution modes @@ -65,8 +47,6 @@ The first implementation supports run-to-completion execution only. Configuring Suspendable execution may resume on another thread. A thread-local baseline cannot safely span that migration without evaluation-owned allocation attribution. -Public multi-call begin/end memory-budget scopes are intentionally absent because allocation counters are thread-local and abandoned or cross-thread scopes would be unsafe. Rust, C FFI, and C# are supported by this PR; other bindings require follow-up work. - ## Process-global limit The existing process-global memory limit remains separate. It protects the process as a whole and is not an isolation mechanism for individual evaluations. When both controls are configured, the per-evaluation budget is checked first. diff --git a/src/rvm/vm/errors.rs b/src/rvm/vm/errors.rs index 3187ea7d1..cdbd753f7 100644 --- a/src/rvm/vm/errors.rs +++ b/src/rvm/vm/errors.rs @@ -312,9 +312,6 @@ pub enum VmError { #[error("Internal VM error: {message} (pc={pc})")] Internal { message: String, pc: usize }, - - #[error("An evaluation memory budget scope is already active (pc={pc})")] - MemoryBudgetScopeAlreadyActive { pc: usize }, } impl From for VmError { diff --git a/src/rvm/vm/machine.rs b/src/rvm/vm/machine.rs index b1630d704..3a51b6bef 100644 --- a/src/rvm/vm/machine.rs +++ b/src/rvm/vm/machine.rs @@ -13,6 +13,7 @@ use crate::utils::limits::{ use crate::value::Value; use crate::CompiledPolicy; use alloc::collections::{btree_map::Entry, BTreeMap, VecDeque}; +use alloc::ffi::CString; #[cfg(all(feature = "allocator-memory-limits", not(miri)))] use alloc::format; use alloc::string::String; @@ -32,7 +33,7 @@ use super::execution_model::{ pub(super) enum MemoryBudgetLifecycle { Inactive, ImplicitExecution, - ExternalScope, + FfiResultSerialization, } /// The Rego Virtual Machine @@ -178,21 +179,21 @@ pub struct RegoVM { } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] -struct EvaluationMemoryBudgetScope<'a> { +struct FfiResultSerializationBudget<'a> { vm: &'a mut RegoVM, } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] -impl EvaluationMemoryBudgetScope<'_> { +impl FfiResultSerializationBudget<'_> { const fn vm(&mut self) -> &mut RegoVM { self.vm } } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] -impl Drop for EvaluationMemoryBudgetScope<'_> { +impl Drop for FfiResultSerializationBudget<'_> { fn drop(&mut self) { - self.vm.finish_external_memory_budget_scope(); + self.vm.finish_ffi_result_serialization_memory_budget(); } } @@ -311,16 +312,6 @@ impl RegoVM { // Check for conflicts between rule tree and data self.program.check_rule_data_conflicts(&data)?; - #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - if matches!( - self.memory_budget_lifecycle, - MemoryBudgetLifecycle::ExternalScope - ) { - // Keep both the installed and provisional data alive until after the - // checkpoint so rejected replacements are atomic at their peak usage. - self.check_memory_budget()?; - } - self.data = data; Ok(()) } @@ -462,38 +453,17 @@ impl RegoVM { /// Configure a fresh memory budget for every run-to-completion execution. #[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] - pub fn set_memory_budget_config(&mut self, config: Option) { - let was_configured = self.memory_budget_config.is_some(); + pub const fn set_memory_budget_config(&mut self, config: Option) { self.memory_budget_config = config; - - if matches!( - self.memory_budget_lifecycle, - MemoryBudgetLifecycle::ExternalScope - ) { - if !was_configured && self.memory_budget_config.is_some() { - self.memory_budget_baseline = limits::current_thread_live_bytes(); - } else if self.memory_budget_config.is_none() { - self.memory_budget_baseline = 0; - } - return; - } - self.memory_budget_baseline = 0; self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; } - /// Return the configured per-execution memory budget. - #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - #[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] - pub const fn memory_budget_config(&self) -> Option { - self.memory_budget_config - } - #[cfg(all(feature = "allocator-memory-limits", not(miri)))] pub(super) fn reset_memory_budget_state(&mut self) { if matches!( self.memory_budget_lifecycle, - MemoryBudgetLifecycle::ExternalScope + MemoryBudgetLifecycle::FfiResultSerialization ) { return; } @@ -510,70 +480,117 @@ impl RegoVM { }; } - /// Run a synchronous evaluation operation under one memory-budget baseline. - /// - /// When a memory budget is configured, data assigned with [`Self::set_data`], - /// VM execution, and allocations retained through the closure share the baseline - /// captured on entry. If a scope starts unconfigured, enabling or re-enabling a - /// budget captures a fresh baseline at that transition; changing one configured - /// threshold to another preserves the existing baseline. Nested scopes return - /// [`VmError::MemoryBudgetScopeAlreadyActive`]. - /// - /// The closure must complete synchronously on the calling thread. Its scope is - /// deactivated before this method returns, including when the closure returns an error - /// or unwinds. #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - #[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] - pub fn with_evaluation_memory_budget(&mut self, operation: F) -> Result - where - F: FnOnce(&mut Self) -> Result, - { + pub(super) const fn finish_implicit_memory_budget_execution(&mut self) { if matches!( self.memory_budget_lifecycle, - MemoryBudgetLifecycle::ExternalScope + MemoryBudgetLifecycle::ImplicitExecution ) { - return Err(VmError::MemoryBudgetScopeAlreadyActive { pc: self.pc }); + self.memory_budget_baseline = 0; + self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; } + } + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + #[allow(clippy::unused_self)] + pub(super) const fn finish_implicit_memory_budget_execution(&mut self) {} + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + fn begin_ffi_result_serialization_memory_budget(&mut self) { self.memory_budget_baseline = if self.memory_budget_config.is_some() { limits::current_thread_live_bytes() } else { 0 }; - self.memory_budget_lifecycle = MemoryBudgetLifecycle::ExternalScope; - - let mut scope = EvaluationMemoryBudgetScope { vm: self }; - let result = operation(scope.vm()).and_then(|value| { - scope.vm().check_memory_budget()?; - Ok(value) - }); - - result + self.memory_budget_lifecycle = MemoryBudgetLifecycle::FfiResultSerialization; } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - pub(super) const fn finish_implicit_memory_budget_execution(&mut self) { + const fn finish_ffi_result_serialization_memory_budget(&mut self) { if matches!( self.memory_budget_lifecycle, - MemoryBudgetLifecycle::ImplicitExecution + MemoryBudgetLifecycle::FfiResultSerialization ) { self.memory_budget_baseline = 0; self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; } } - #[cfg(any(miri, not(feature = "allocator-memory-limits")))] - #[allow(clippy::unused_self)] - pub(super) const fn finish_implicit_memory_budget_execution(&mut self) {} + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) const fn finish_active_memory_budget_execution(&mut self) { + self.memory_budget_baseline = 0; + self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; + } + + /// Execute and serialize a main entry point for the native FFI. + /// + /// This is an internal binding hook, not a general-purpose budget scope. It starts the + /// execution budget after data, input, context, and program loading have completed, then + /// keeps that budget active only through immediate native JSON and C-string production. + #[doc(hidden)] + pub fn execute_to_c_string_for_ffi(&mut self) -> Result { + self.execute_to_c_string_for_ffi_with(Self::execute) + } + + /// Execute and serialize a named entry point for the native FFI. + #[doc(hidden)] + pub fn execute_entry_point_by_name_to_c_string_for_ffi( + &mut self, + name: &str, + ) -> Result { + self.execute_to_c_string_for_ffi_with(|vm| vm.execute_entry_point_by_name(name)) + } + + /// Execute and serialize an indexed entry point for the native FFI. + #[doc(hidden)] + pub fn execute_entry_point_by_index_to_c_string_for_ffi( + &mut self, + index: usize, + ) -> Result { + self.execute_to_c_string_for_ffi_with(|vm| vm.execute_entry_point_by_index(index)) + } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - const fn finish_external_memory_budget_scope(&mut self) { - if matches!( - self.memory_budget_lifecycle, - MemoryBudgetLifecycle::ExternalScope - ) { - self.memory_budget_baseline = 0; - self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; + fn execute_to_c_string_for_ffi_with(&mut self, execute: F) -> Result + where + F: FnOnce(&mut Self) -> Result, + { + self.begin_ffi_result_serialization_memory_budget(); + let mut budget = FfiResultSerializationBudget { vm: self }; + let output = (|| { + let value = execute(budget.vm())?; + let json = value.to_json_str().map_err(VmError::from)?; + let output = CString::new(json).map_err(|_| VmError::Internal { + message: String::from("RVM JSON result contained an interior NUL byte"), + pc: budget.vm().pc, + })?; + budget.vm().check_memory_budget()?; + Ok(output) + })(); + + match output { + Ok(output) => Ok(output), + Err(error) => Err(budget.vm().fail_run_to_completion(error)), + } + } + + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + fn execute_to_c_string_for_ffi_with(&mut self, execute: F) -> Result + where + F: FnOnce(&mut Self) -> Result, + { + let output = (|| { + let value = execute(self)?; + let json = value.to_json_str().map_err(VmError::from)?; + CString::new(json).map_err(|_| VmError::Internal { + message: String::from("RVM JSON result contained an interior NUL byte"), + pc: self.pc, + }) + })(); + + match output { + Ok(output) => Ok(output), + Err(error) => Err(self.fail_run_to_completion(error)), } } @@ -609,16 +626,9 @@ impl RegoVM { Ok(()) } - /// Check the configured budget against the active evaluation-memory baseline. - /// - /// Within [`Self::with_evaluation_memory_budget`], call this after scoped data - /// preparation or native result post-processing to include those allocations before - /// returning success. Ordinary `execute*` calls checkpoint their execution budget and - /// deactivate it on successful return; bindings that need to charge serialization must - /// wrap execution and serialization in an external evaluation-memory scope. + /// Check the configured budget against the active execution baseline. #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - #[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] - pub fn check_memory_budget(&mut self) -> Result<()> { + pub(super) fn check_memory_budget(&mut self) -> Result<()> { let Some(config) = self.memory_budget_config.filter(|_| { !matches!( self.memory_budget_lifecycle, @@ -931,4 +941,80 @@ mod memory_budget_tests { VmError::MemoryBudgetExceeded { .. } )); } + + #[allow(clippy::expect_used)] + #[test] + fn ffi_result_serialization_window_cleans_up_on_terminal_paths() { + let mut program = crate::rvm::program::Program::new(); + program.entry_points.insert("main".into(), 0); + program.instructions = + alloc::vec![crate::rvm::instructions::Instruction::Return { value: 0 }]; + program.instruction_spans = alloc::vec![None]; + + let mut vm = RegoVM::new(); + vm.load_program(alloc::sync::Arc::new(program)); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1024 * 1024).unwrap_or(NonZeroU64::MIN), + })); + + assert_eq!( + vm.execute_to_c_string_for_ffi() + .expect("main FFI serialization succeeds") + .as_bytes_with_nul(), + b"\"\"\0" + ); + assert_eq!( + vm.execute_entry_point_by_name_to_c_string_for_ffi("main") + .expect("named FFI serialization succeeds") + .as_bytes_with_nul(), + b"\"\"\0" + ); + assert_eq!( + vm.execute_entry_point_by_index_to_c_string_for_ffi(0) + .expect("indexed FFI serialization succeeds") + .as_bytes_with_nul(), + b"\"\"\0" + ); + assert!(matches!( + vm.memory_budget_lifecycle, + super::MemoryBudgetLifecycle::Inactive + )); + + vm.set_max_instructions(0); + assert!(matches!( + vm.execute_to_c_string_for_ffi(), + Err(VmError::InstructionLimitExceeded { .. }) + )); + assert!(matches!( + vm.execution_state, + super::super::execution_model::ExecutionState::Error { + error: VmError::InstructionLimitExceeded { .. } + } + )); + assert!(matches!( + vm.memory_budget_lifecycle, + super::MemoryBudgetLifecycle::Inactive + )); + } + + #[allow(clippy::panic)] + #[test] + fn ffi_result_serialization_window_deactivates_during_unwind() { + let mut vm = RegoVM::new(); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(1024 * 1024).unwrap_or(NonZeroU64::MIN), + })); + + let unwind = std::panic::catch_unwind(core::panic::AssertUnwindSafe(|| { + vm.begin_ffi_result_serialization_memory_budget(); + let _budget = super::FfiResultSerializationBudget { vm: &mut vm }; + panic!("FFI serialization cleanup regression"); + })); + + assert!(unwind.is_err()); + assert!(matches!( + vm.memory_budget_lifecycle, + super::MemoryBudgetLifecycle::Inactive + )); + } } diff --git a/src/rvm/vm/state.rs b/src/rvm/vm/state.rs index 1c9bb1ad8..afe559fda 100644 --- a/src/rvm/vm/state.rs +++ b/src/rvm/vm/state.rs @@ -21,10 +21,11 @@ impl RegoVM { #[cfg(all(feature = "allocator-memory-limits", not(miri)))] if matches!( self.memory_budget_lifecycle, - super::machine::MemoryBudgetLifecycle::ExternalScope + super::machine::MemoryBudgetLifecycle::FfiResultSerialization ) { - // An external scope owns the baseline. Observe the post-release trough - // before initialization so fresh state allocations are charged to it. + // The internal FFI serialization window owns the baseline. Observe the + // post-release trough before initialization so fresh state allocations are + // charged to that execution. self.check_memory_budget()?; } else { self.reset_memory_budget_state(); @@ -46,14 +47,11 @@ impl RegoVM { /// Release values retained by a failed run-to-completion execution and record its error. /// - /// This is public only so binding crates can apply the same terminal transition when - /// post-execution marshaling exceeds the execution budget. - #[doc(hidden)] - pub fn fail_run_to_completion(&mut self, error: VmError) -> VmError { + pub(super) fn fail_run_to_completion(&mut self, error: VmError) -> VmError { self.release_previous_execution_state(); #[cfg(all(feature = "allocator-memory-limits", not(miri)))] { - self.finish_implicit_memory_budget_execution(); + self.finish_active_memory_budget_execution(); } self.execution_state = ExecutionState::Error { error: error.clone(), @@ -271,9 +269,11 @@ mod tests { use super::{ExecutionState, RegoVM, Value, VmError}; #[cfg(all(feature = "allocator-memory-limits", not(miri)))] use crate::rvm::program::{RuleInfo, RuleType}; + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] use crate::rvm::{instructions::Instruction, program::Program}; #[cfg(all(feature = "allocator-memory-limits", not(miri)))] use crate::MemoryBudgetConfig; + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] use alloc::sync::Arc; use alloc::vec; #[cfg(all(feature = "allocator-memory-limits", not(miri)))] @@ -366,205 +366,11 @@ mod tests { vm.memory_budget_lifecycle, MemoryBudgetLifecycle::Inactive )); - - vm.with_evaluation_memory_budget(|vm| { - assert_eq!( - vm.execute().expect("main scope execution"), - Value::Undefined - ); - assert_eq!( - vm.execute_entry_point_by_name("main") - .expect("named scope execution"), - Value::Undefined - ); - assert_eq!( - vm.execute_entry_point_by_index(0) - .expect("indexed scope execution"), - Value::Undefined - ); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::ExternalScope - )); - Ok(()) - }) - .expect("external scope succeeds"); } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - #[allow(clippy::expect_used)] #[test] - fn evaluation_memory_budget_scope_rejects_nesting_and_cleans_up_terminal_paths() { - let mut vm = RegoVM::new(); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1024 * 1024).unwrap_or(NonZeroU64::MIN), - })); - - vm.with_evaluation_memory_budget(|vm| { - assert!(matches!( - vm.with_evaluation_memory_budget(|_| Ok(())), - Err(VmError::MemoryBudgetScopeAlreadyActive { .. }) - )); - assert_eq!(vm.execute()?, Value::Undefined); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::ExternalScope - )); - Ok(()) - }) - .expect("outer scope succeeds"); - - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::Inactive - )); - - assert!(matches!( - vm.with_evaluation_memory_budget(|vm| { - let result = vm.execute_entry_point_by_index(0); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::ExternalScope - )); - result - }), - Err(VmError::InvalidEntryPointIndex { .. }) - )); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::Inactive - )); - - let mut error_program = Program::new(); - error_program.instructions = vec![Instruction::Return { value: 0 }]; - error_program.instruction_spans = vec![None]; - vm.load_program(Arc::new(error_program)); - vm.set_max_instructions(0); - assert!(matches!( - vm.with_evaluation_memory_budget(|vm| { - let result = vm.execute(); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::ExternalScope - )); - result - }), - Err(VmError::InstructionLimitExceeded { .. }) - )); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::Inactive - )); - vm.set_max_instructions(25_000); - - vm.set_execution_mode(super::super::execution_model::ExecutionMode::Suspendable); - assert!(matches!( - vm.with_evaluation_memory_budget(|vm| { - let result = vm.execute(); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::ExternalScope - )); - result - }), - Err(VmError::MemoryBudgetUnsupportedInSuspendableExecution { .. }) - )); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::Inactive - )); - vm.set_execution_mode(super::super::execution_model::ExecutionMode::RunToCompletion); - - let closure_error = VmError::Internal { - message: "closure failure".into(), - pc: 0, - }; - assert_eq!( - vm.with_evaluation_memory_budget(|_| Err::<(), _>(closure_error.clone())), - Err(closure_error) - ); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::Inactive - )); - assert_eq!( - vm.with_evaluation_memory_budget(|vm| vm.execute()) - .expect("fresh scope after closure failure"), - Value::Undefined - ); - } - - #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - #[allow(clippy::panic)] - #[test] - fn evaluation_memory_budget_scope_cleans_up_during_unwind() { - let mut vm = RegoVM::new(); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1024 * 1024).unwrap_or(NonZeroU64::MIN), - })); - - let unwind = std::panic::catch_unwind(core::panic::AssertUnwindSafe(|| { - let _ = vm.with_evaluation_memory_budget(|_| -> core::result::Result<(), VmError> { - panic!("scope panic for cleanup regression"); - }); - })); - - assert!(unwind.is_err()); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::Inactive - )); - assert_eq!( - vm.with_evaluation_memory_budget(|vm| vm.execute()), - Ok(Value::Undefined) - ); - } - - #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - #[test] - fn scoped_budget_configuration_captures_a_fresh_baseline() { - let tiny_budget = MemoryBudgetConfig { - limit: NonZeroU64::new(1).unwrap_or(NonZeroU64::MIN), - }; - let relaxed_budget = MemoryBudgetConfig { - limit: NonZeroU64::new(1024 * 1024).unwrap_or(NonZeroU64::MIN), - }; - let tightened_budget = MemoryBudgetConfig { - limit: NonZeroU64::new(64 * 1024).unwrap_or(NonZeroU64::MIN), - }; - - let mut unconfigured_vm = RegoVM::new(); - let configured_from_none = unconfigured_vm.with_evaluation_memory_budget(|vm| { - vm.set_memory_budget_config(Some(tiny_budget)); - vm.check_memory_budget()?; - - vm.set_memory_budget_config(None); - let allocation = alloc::vec![0_u8; 128 * 1024]; - core::hint::black_box(&allocation); - - vm.set_memory_budget_config(Some(tiny_budget)); - vm.check_memory_budget() - }); - assert!(matches!(configured_from_none, Ok(()))); - - let mut configured_vm = RegoVM::new(); - configured_vm.set_memory_budget_config(Some(relaxed_budget)); - let reconfigured = configured_vm.with_evaluation_memory_budget(|vm| { - let allocation = alloc::vec![0_u8; 128 * 1024]; - core::hint::black_box(&allocation); - - vm.set_memory_budget_config(Some(tightened_budget)); - vm.check_memory_budget() - }); - assert!(matches!( - reconfigured, - Err(VmError::MemoryBudgetExceeded { .. }) - )); - } - - #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - #[test] - fn external_scope_counts_fresh_initialization_after_releasing_previous_state() { + fn ffi_result_serialization_counts_fresh_initialization_after_releasing_previous_state() { let rule_info = RuleInfo::new( "unused".into(), RuleType::Complete, @@ -582,9 +388,9 @@ mod tests { let mut vm = RegoVM::new(); vm.load_program(Arc::new(program)); - // This models a reused VM whose completed result is released before a fresh - // execution-state allocation. Keeping the previous cache empty makes the - // fresh rule-cache allocation observable at the next VM checkpoint. + // Model a reused VM whose completed result is released before fresh + // execution-state allocation. The FFI window must observe that trough, + // then charge the fresh rule-cache allocation. vm.rule_cache = alloc::vec![]; let retained = Value::from( (0..(Program::MAX_RULES * 4)) @@ -597,9 +403,10 @@ mod tests { limit: NonZeroU64::new(32 * 1024).unwrap_or(NonZeroU64::MIN), })); - let result = vm.with_evaluation_memory_budget(|vm| vm.execute()); - - assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); + assert!(matches!( + vm.execute_to_c_string_for_ffi(), + Err(VmError::MemoryBudgetExceeded { .. }) + )); assert!(matches!( vm.execution_state, ExecutionState::Error { @@ -607,28 +414,4 @@ mod tests { } )); } - - #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - #[test] - fn scoped_set_data_is_atomic_when_the_budget_is_exhausted() { - let mut vm = RegoVM::new(); - let original = Value::from("original"); - vm.data = original.clone(); - vm.set_memory_budget_config(Some(MemoryBudgetConfig { - limit: NonZeroU64::new(1024).unwrap_or(NonZeroU64::MIN), - })); - - let result = vm.with_evaluation_memory_budget(|vm| { - let candidate = - Value::from((0..16_384).map(Value::from).collect::>()); - vm.set_data(candidate) - }); - - assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); - assert_eq!(vm.data, original); - assert!(matches!( - vm.memory_budget_lifecycle, - MemoryBudgetLifecycle::Inactive - )); - } } diff --git a/tests/memory_limits.rs b/tests/memory_limits.rs index 1e69d1e70..cbb9bec05 100644 --- a/tests/memory_limits.rs +++ b/tests/memory_limits.rs @@ -491,92 +491,6 @@ fn vm_memory_budget_is_fresh_for_each_execution() { ); } -#[cfg(feature = "rvm")] -#[test] -fn vm_evaluation_memory_budget_scope_preserves_its_baseline_across_execution() { - let _guard = LimitGuard::lock(); - let mut vm = RegoVM::new(); - vm.set_memory_budget_config(Some(memory_budget(64 * 1024))); - - let result = vm.with_evaluation_memory_budget(|vm| { - let allocation = vec![0_u8; 128 * 1024]; - core::hint::black_box(&allocation); - vm.execute() - }); - - assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); - - // A completed scope must not leave a stale baseline on a reusable VM. - assert_eq!( - vm.with_evaluation_memory_budget(|vm| vm.execute()) - .expect("fresh scope after exhaustion"), - Value::Undefined - ); -} - -#[cfg(feature = "rvm")] -#[test] -fn vm_evaluation_memory_budget_charges_scoped_data_and_execution_cumulatively() { - let _guard = LimitGuard::lock(); - const BUDGET_BYTES: u64 = 832 * 1024; - - let mut engine = new_engine_with_module(LARGE_PARSE_MODULE); - let entrypoint = Rc::from("data.limit.large_array"); - let compiled = engine - .compile_with_entrypoint(&entrypoint) - .expect("compile policy for VM"); - let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) - .expect("compile VM program"); - - let mut ordinary_vm = RegoVM::new(); - ordinary_vm.load_program(program.clone()); - ordinary_vm - .set_data(large_json_data(20_000)) - .expect("set static data"); - ordinary_vm.set_input(Value::Undefined); - ordinary_vm.set_memory_budget_config(Some(memory_budget(BUDGET_BYTES))); - assert!(matches!( - ordinary_vm - .execute_entry_point_by_name(entrypoint.as_ref()) - .expect("execution-only budget permits parsing"), - Value::Array(_) - )); - - let mut vm = RegoVM::new(); - vm.load_program(program); - vm.set_input(Value::Undefined); - vm.set_memory_budget_config(Some(memory_budget(BUDGET_BYTES))); - - let result = vm.with_evaluation_memory_budget(|vm| { - vm.set_data(large_json_data(20_000))?; - vm.execute_entry_point_by_name(entrypoint.as_ref()) - }); - - assert!(matches!(result, Err(VmError::MemoryBudgetExceeded { .. }))); -} - -#[cfg(feature = "rvm")] -#[test] -fn vm_evaluation_memory_budget_scope_without_configuration_is_behavioral() { - let mut vm = RegoVM::new(); - - vm.with_evaluation_memory_budget(|vm| { - assert!(matches!( - vm.with_evaluation_memory_budget(|_| Ok(())), - Err(VmError::MemoryBudgetScopeAlreadyActive { .. }) - )); - assert_eq!(vm.execute()?, Value::Undefined); - Ok(()) - }) - .expect("unconfigured scope succeeds"); - - assert_eq!( - vm.with_evaluation_memory_budget(|vm| vm.execute()) - .expect("unconfigured scope is reusable"), - Value::Undefined - ); -} - #[cfg(feature = "rvm")] #[test] fn vm_memory_budget_excludes_static_data_before_ordinary_execution() {