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/benches/rvm_benchmark.rs b/benches/rvm_benchmark.rs index 7ac093401..1d12f16de 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(all(feature = "allocator-memory-limits", not(miri)))] +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(all(feature = "allocator-memory-limits", not(miri)))] +use regorus::MemoryBudgetConfig; use regorus::{Engine, Rc, Value}; // --------------------------------------------------------------------------- @@ -57,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(); @@ -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,10 +373,13 @@ 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 { - #[cfg(feature = "allocator-memory-limits")] +/// 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(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, @@ -369,11 +387,16 @@ fn configure_limits(vm: &mut RegoVM, limits: bool) { })); 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); } + + #[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"), + })); } // --------------------------------------------------------------------------- @@ -408,7 +431,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 +468,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/API.md b/bindings/csharp/API.md index a8b0c83fa..8c3b62b2b 100644 --- a/bindings/csharp/API.md +++ b/bindings/csharp/API.md @@ -81,6 +81,45 @@ 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 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(); +} +``` + +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(); +vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(16UL * 1024 * 1024)); + +try +{ + var result = vm.Execute(); +} +catch (RegorusMemoryBudgetExceededException) +{ + // The execution exceeded its configured 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. 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 0d5669bf7..55807176a 100644 --- a/bindings/csharp/README.md +++ b/bindings/csharp/README.md @@ -105,6 +105,31 @@ 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 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(); +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 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. + +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 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..34c34ce72 --- /dev/null +++ b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs @@ -0,0 +1,120 @@ +// 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 ulong TightMemoryBudgetBytes = 64 * 1024; + + 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"; + + private const string PreloadedResultPolicy = """ +package limits.memory + +large_string := data.large_string +"""; + + private const string PreloadedResultEntryPoint = "data.limits.memory.large_string"; + + [TestMethod] + public void Memory_budget_must_be_non_zero() + { + Assert.ThrowsException(() => new MemoryBudgetConfig(0)); + + using var vm = new Rvm(); + Assert.ThrowsException(() => vm.SetMemoryBudgetConfig(default)); + } + + [TestMethod] + public void Execute_exceeding_memory_budget_throws_typed_exception() + { + using var program = CreateProgram(); + using var vm = CreateRvm(program); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); + + Assert.ThrowsException(() => vm.ExecuteEntryPoint(EntryPoint)); + } + + [TestMethod] + public void Clearing_memory_budget_restores_unlimited_execution() + { + using var program = CreateProgram(); + using var vm = CreateRvm(program); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(TightMemoryBudgetBytes)); + Assert.ThrowsException(() => vm.ExecuteEntryPoint(EntryPoint)); + + vm.ClearMemoryBudgetConfig(); + + var result = vm.ExecuteEntryPoint(EntryPoint); + 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 Suspendable_execution_rejects_memory_budget() + { + using var vm = new Rvm(); + vm.SetExecutionMode(ExecutionMode.Suspendable); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(1024)); + + Assert.ThrowsException(() => vm.Execute()); + } + + 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..a77d5ed10 --- /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 ArgumentOutOfRangeException(nameof(LimitBytes), "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..479c9b7c9 100644 --- a/bindings/csharp/Regorus/NativeMethods.cs +++ b/bindings/csharp/Regorus/NativeMethods.cs @@ -245,6 +245,14 @@ 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 +836,14 @@ internal enum RegorusStatus : uint /// The engine remains poisoned because a previous panic was detected. /// Poisoned, + /// + /// An RVM execution exceeded its configured memory budget. + /// + MemoryBudgetExceeded, + /// + /// An RVM memory budget was used with suspendable execution. + /// + MemoryBudgetUnsupportedInSuspendableExecution, } /// @@ -883,6 +899,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/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/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..ffbd284c8 100644 --- a/bindings/csharp/Regorus/StatusExtensions.cs +++ b/bindings/csharp/Regorus/StatusExtensions.cs @@ -17,6 +17,8 @@ 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..a100511db 100644 --- a/bindings/ffi/CHANGELOG.md +++ b/bindings/ffi/CHANGELOG.md @@ -6,6 +6,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- 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 ### Other diff --git a/bindings/ffi/src/common.rs b/bindings/ffi/src/common.rs index bf8cca3d0..53d37648d 100644 --- a/bindings/ffi/src/common.rs +++ b/bindings/ffi/src/common.rs @@ -43,6 +43,12 @@ pub enum RegorusStatus { /// The engine remains poisoned because a previous panic was detected. Poisoned, + + /// 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 @@ -133,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/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 dc98a4ee8..43459e02f 100644 --- a/bindings/ffi/src/limits.rs +++ b/bindings/ffi/src/limits.rs @@ -4,11 +4,13 @@ use crate::common::{to_regorus_result, RegorusResult, RegorusStatus}; use alloc::format; use anyhow::{anyhow, Result}; +#[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) @@ -17,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) => { @@ -43,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() { @@ -77,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, @@ -106,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, @@ -115,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"), ) } @@ -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(all(feature = "allocator-memory-limits", not(miri)))] +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) @@ -233,8 +252,10 @@ pub extern "C" fn regorus_clear_cache() -> RegorusResult { RegorusResult::ok_void() } -#[cfg(test)] +#[cfg(all(test, feature = "allocator-memory-limits", not(miri)))] 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/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 7603305c1..bc14da916 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,53 @@ impl RegorusRvm { } } +fn to_rvm_string_result(output: Result) -> RegorusResult { + match output { + Ok(json) => RegorusResult::ok_string(json), + 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()) +} + +enum RvmExecution { + Main, + Named(String), + Indexed(usize), +} + +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()?; + + 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)? + } + RvmExecution::Indexed(index) => { + guard.execute_entry_point_by_index_to_c_string_for_ffi(index)? + } + }; + Ok(RegorusResult::ok_c_string(output)) + }(); + + match output { + Ok(result) => result, + Err(err) => to_rvm_error_result(err), + } +} + /// Drop a `RegorusProgram`. #[no_mangle] pub extern "C" fn regorus_program_drop(program: *mut RegorusProgram) { @@ -478,24 +525,59 @@ pub extern "C" fn regorus_rvm_set_execution_timer_config( }) } -/// Execute the program's main entry point. +/// Configure the per-VM memory budget for run-to-completion execution. +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] #[no_mangle] -pub extern "C" fn regorus_rvm_execute(vm: *mut RegorusRvm) -> RegorusResult { +pub extern "C" fn regorus_rvm_set_memory_budget_config( + vm: *mut RegorusRvm, + has_config: bool, + config: RegorusMemoryBudgetConfig, +) -> RegorusResult { with_unwind_guard(|| { - let output = || -> Result { + 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()?; - let result = guard.execute()?; - result.to_json_str() - }(); - - match output { - Ok(json) => RegorusResult::ok_string(json), - Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), - } + guard.set_memory_budget_config(config); + Ok(()) + }()) }) } +/// Report that memory budgets are unavailable without allocator tracking. +#[cfg(any(not(feature = "allocator-memory-limits"), miri))] +#[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: allocator memory tracking is disabled" + .into(), + ) +} + +/// 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, RvmExecution::Main)) +} + /// Execute a named entry point. #[no_mangle] pub extern "C" fn regorus_rvm_execute_entry_point_by_name( @@ -503,18 +585,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() - }(); - - match output { - Ok(json) => RegorusResult::ok_string(json), - Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), - } + 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, RvmExecution::Named(entry_point)) }) } @@ -524,19 +599,7 @@ pub extern "C" fn regorus_rvm_execute_entry_point_by_index( vm: *mut RegorusRvm, 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() - }(); - - match output { - Ok(json) => RegorusResult::ok_string(json), - Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), - } - }) + with_unwind_guard(|| execute_to_rvm_result(vm, RvmExecution::Indexed(index))) } /// Resume execution for suspendable runs. @@ -559,10 +622,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) }) } @@ -584,6 +644,284 @@ pub extern "C" fn regorus_rvm_get_execution_state(vm: *mut RegorusRvm) -> Regoru }) } +#[cfg(all(test, feature = "allocator-memory-limits", not(miri)))] +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_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}; + 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; + use regorus::rvm::program::Program; + use regorus::rvm::vm::{ExecutionMode, RegoVM}; + use regorus::{Engine, MemoryBudgetConfig, Rc, Value}; + + const POLICY: &str = r#" +package limits.memory +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 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(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; + 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(); + + 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_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"); + 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(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); + + let vm = Box::into_raw(Box::new(RegorusRvm::new(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"); + 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()), + ); + assert_memory_budget_failure_state(vm, regorus_rvm_execute_entry_point_by_index(vm, 0)); + + 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(vm); + assert!(matches!(result.status, RegorusStatus::Ok)); + 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 + ); + } +} + +#[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 new file mode 100644 index 000000000..60d3d500f --- /dev/null +++ b/docs/limits/memory_budget.md @@ -0,0 +1,52 @@ +# RVM memory budgets + +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. 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; +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 `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 and outside the execution baseline and are not charged. + +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. + +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 + +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. 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: + +- additional live-byte usage observed for the evaluation +- configured budget +- VM program counter + +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 + +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. + +## 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..de618e05f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -172,10 +172,11 @@ 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, - 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..07f1fc5b0 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,19 +46,31 @@ impl RegoVM { match self.execution_mode { ExecutionMode::RunToCompletion => { - self.reset_execution_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) } - })?; - - self.jump_to(entry_point_pc_u32) + Err(err) => Err(self.fail_run_to_completion(err)), + } } ExecutionMode::Suspendable => { self.reset_execution_state(); @@ -69,6 +83,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,19 +103,31 @@ impl RegoVM { match self.execution_mode { ExecutionMode::RunToCompletion => { - self.reset_execution_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) } - })?; - - self.jump_to(entry_point_pc_u32) + Err(err) => Err(self.fail_run_to_completion(err)), + } } ExecutionMode::Suspendable => { self.reset_execution_state(); @@ -163,20 +190,26 @@ impl RegoVM { } fn execute_run_to_completion(&mut self) -> Result { - self.reset_execution_state(); - self.reset_execution_timer_state(); - self.execution_state = ExecutionState::Running; - match self.jump_to(0_u32) { + 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)), } } @@ -210,6 +243,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 48a078607..3a51b6bef 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, @@ -11,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; @@ -25,6 +28,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, + FfiResultSerialization, +} + /// The Rego Virtual Machine #[derive(Debug)] pub struct RegoVM { @@ -131,6 +142,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, + + /// Owner of the current memory-budget baseline. + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) memory_budget_lifecycle: MemoryBudgetLifecycle, + /// Cached dummy span for builtin calls (avoids Source::from_contents per call) pub(super) dummy_span: Option, @@ -155,6 +178,25 @@ pub struct RegoVM { pub(super) metadata_value: Value, } +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +struct FfiResultSerializationBudget<'a> { + vm: &'a mut RegoVM, +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +impl FfiResultSerializationBudget<'_> { + const fn vm(&mut self) -> &mut RegoVM { + self.vm + } +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +impl Drop for FfiResultSerializationBudget<'_> { + fn drop(&mut self) { + self.vm.finish_ffi_result_serialization_memory_budget(); + } +} + impl Default for RegoVM { fn default() -> Self { Self::new() @@ -197,6 +239,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_lifecycle: MemoryBudgetLifecycle::Inactive, dummy_span: None, dummy_exprs: Vec::new(), cached_builtin_args: Vec::new(), @@ -402,6 +450,232 @@ 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)))] + #[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.memory_budget_baseline = 0; + self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) fn reset_memory_budget_state(&mut self) { + if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::FfiResultSerialization + ) { + 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 + }; + } + + #[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)))] + 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::FfiResultSerialization; + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + const fn finish_ffi_result_serialization_memory_budget(&mut self) { + if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::FfiResultSerialization + ) { + self.memory_budget_baseline = 0; + self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; + } + } + + #[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)))] + 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)), + } + } + + #[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(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(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 active execution baseline. + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) fn check_memory_budget(&mut self) -> Result<()> { + let Some(config) = self.memory_budget_config.filter(|_| { + !matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::Inactive + ) + }) else { + return Ok(()); + }; + + let current = limits::current_thread_live_bytes(); + 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 { + usage, + budget, + pc: self.pc, + }); + } + + Ok(()) + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) fn apply_memory_budget_precedence(&mut self, err: VmError) -> VmError { + if matches!(err, VmError::MemoryLimitExceeded { .. }) { + self.check_memory_budget().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(&mut self) -> Result<()> { + Ok(()) + } + + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + #[allow(clippy::unused_self)] + pub(super) fn apply_memory_budget_precedence(&mut 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 +815,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()?; + 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")))] @@ -600,3 +877,144 @@ 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 { .. } + )); + } + + #[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/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..afe559fda 100644 --- a/src/rvm/vm/state.rs +++ b/src/rvm/vm/state.rs @@ -11,30 +11,70 @@ 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) -> Result<()> { + self.release_previous_execution_state(); + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + if matches!( + self.memory_budget_lifecycle, + super::machine::MemoryBudgetLifecycle::FfiResultSerialization + ) { + // 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(); + } + self.initialize_execution_state(); + Ok(()) + } + + 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(); + } + + /// Release values retained by a failed run-to-completion execution and record its error. + /// + 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_active_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; 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. @@ -221,3 +261,157 @@ 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}; + #[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)))] + 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 + )); + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + #[test] + fn ffi_result_serialization_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)); + // 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)) + .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), + })); + + assert!(matches!( + vm.execute_to_c_string_for_ffi(), + Err(VmError::MemoryBudgetExceeded { .. }) + )); + assert!(matches!( + vm.execution_state, + ExecutionState::Error { + error: VmError::MemoryBudgetExceeded { .. } + } + )); + } +} diff --git a/src/utils/limits/memory.rs b/src/utils/limits/memory.rs index 4fe282013..fac97ce2d 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,14 @@ 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)] +#[cfg_attr(docsrs, doc(cfg(feature = "allocator-memory-limits")))] +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 +220,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..c8f4e7e52 100644 --- a/src/utils/limits/mod.rs +++ b/src/utils/limits/mod.rs @@ -14,12 +14,15 @@ 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)))] +#[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, - 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..cbb9bec05 100644 --- a/tests/memory_limits.rs +++ b/tests/memory_limits.rs @@ -1,17 +1,25 @@ #![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")] 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; @@ -72,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 @@ -115,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(); @@ -139,7 +187,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 +331,314 @@ 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 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 constrained data"); + vm.set_input(Value::Undefined); + vm.set_memory_budget_config(Some(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); + + 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_does_not_mask_global_limit() { + let mut 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"); + + 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(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); + guard.set_below_current_usage(); + + assert!(matches!( + vm.execute_entry_point_by_name(entrypoint.as_ref()), + Err(VmError::MemoryLimitExceeded { .. }) + )); +} + +#[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"); + + 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))); + + 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 { .. } + } + )); + } + } +} + +#[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(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") + .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_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() { + 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(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); + assert!(matches!( + vm.execute().expect("first execution"), + Value::Array(_) + )); + + vm.set_memory_budget_config(Some(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); + + 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(memory_budget(1024))); + + 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_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() { + 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(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); + 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(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); + relaxed_barrier.wait(); + vm.execute() + }); + + assert!(matches!( + constrained.join().expect("constrained thread"), + Err(VmError::MemoryBudgetExceeded { .. }) + )); + 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:?}"), + } + }); +}