From e3ba05e8f6a4d4fdfce71f59d31fbbd791e9f71b Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Tue, 6 Oct 2026 12:05:25 +0200 Subject: [PATCH 01/15] feat(rvm): support suspendable memory budgets Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: d6592a3a-3255-447c-a539-00a799431708 --- CHANGELOG.md | 1 + benches/rvm_benchmark.rs | 8 +- bindings/csharp/API.md | 12 +- bindings/csharp/README.md | 8 +- .../Regorus.Tests/RvmMemoryBudgetTests.cs | 172 ++++- bindings/csharp/Regorus/NativeMethods.cs | 2 +- ...RegorusMemoryBudgetUnsupportedException.cs | 2 +- bindings/csharp/Regorus/Rvm.cs | 2 +- bindings/ffi/CHANGELOG.md | 11 +- bindings/ffi/src/common.rs | 3 +- bindings/ffi/src/lib.rs | 3 + bindings/ffi/src/limits.rs | 3 + bindings/ffi/src/rvm.rs | 631 ++++++++++++--- docs/limits/memory_budget.md | 127 +++- docs/rvm/vm-runtime.md | 14 +- mimalloc/src/limits.rs | 186 ++++- mimalloc/src/mimalloc.rs | 719 +++++++++++++++++- src/rvm/vm/execution.rs | 328 +++++--- src/rvm/vm/machine.rs | 545 ++++++++++++- src/rvm/vm/state.rs | 54 +- src/utils/limits/memory.rs | 2 + src/utils/limits/mod.rs | 28 +- src/value/mod.rs | 132 +++- src/value/tests.rs | 53 ++ tests/memory_limits.rs | 531 ++++++++++++- 25 files changed, 3173 insertions(+), 404 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b2fe012ee..51500090c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - Rule-conflict diagnostics now retain the current source envelope and identify the previous rule by file, line, and column without nesting a second diagnostic; literal percent, CR, and LF characters in the previous file label are percent-encoded. RVM conflict result semantics are unchanged. +- Suspendable RVM budgets remain immutable across host-await resumes, include native resume parsing and immediate result serialization, and leave malformed JSON syntax/EOF retryable; independent unbudgeted VM calls do not inherit a caller's account, and run-to-completion thread-baseline behavior remains unchanged. ## [0.12.0](https://github.com/microsoft/regorus/compare/regorus-v0.11.0...regorus-v0.12.0) - 2026-09-01 diff --git a/benches/rvm_benchmark.rs b/benches/rvm_benchmark.rs index 1d12f16de..a18e743c9 100644 --- a/benches/rvm_benchmark.rs +++ b/benches/rvm_benchmark.rs @@ -75,7 +75,7 @@ struct EvalConfig { memory_budget: bool, } -const EVAL_CONFIGS: [EvalConfig; 5] = [ +const EVAL_CONFIGS: [EvalConfig; 6] = [ EvalConfig { name: "regular_no_limits", mode: ExecutionMode::RunToCompletion, @@ -106,6 +106,12 @@ const EVAL_CONFIGS: [EvalConfig; 5] = [ limits: true, memory_budget: false, }, + EvalConfig { + name: "suspendable_memory_budget", + mode: ExecutionMode::Suspendable, + limits: false, + memory_budget: true, + }, ]; // --------------------------------------------------------------------------- diff --git a/bindings/csharp/API.md b/bindings/csharp/API.md index ae1cd7e1b..afc073fb7 100644 --- a/bindings/csharp/API.md +++ b/bindings/csharp/API.md @@ -85,7 +85,7 @@ The Regorus C# bindings provide a modern, thread-safe API for compiling and eval ## RVM Memory Budgets -`Rvm.SetMemoryBudgetConfig` configures a fresh non-zero budget for each run-to-completion execution. `Rvm.ClearMemoryBudgetConfig` restores unlimited execution. +`Rvm.SetMemoryBudgetConfig` configures a fresh non-zero budget for the next initial RVM execution in either run-to-completion or suspendable mode. `Rvm.ClearMemoryBudgetConfig` affects the next initial execution; it does not change an active suspendable continuation. ```csharp public readonly struct MemoryBudgetConfig @@ -101,7 +101,7 @@ public sealed class Rvm : IDisposable } ``` -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. +`Execute` and `ExecuteEntryPoint` start a fresh budget for each accepted initial execution. A suspendable execution retains the same budget across `Resume` calls, including resumes on another thread. 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(); @@ -117,11 +117,13 @@ catch (RegorusMemoryBudgetExceededException) } ``` -Native result JSON serialization and Rust `CString` 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. +Native resume C-string copying and JSON parsing, VM resume work, result JSON serialization, and Rust `CString` allocation are included before an `Execute`, `ExecuteEntryPoint`, or `Resume` 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. +Memory budgets require a native library built with allocator memory tracking. Changing or clearing the configuration while a suspendable execution is active affects only the next initial execution. `RegorusMemoryBudgetExceededException` is thrown when a budget is exceeded. Native status 11 and `RegorusMemoryBudgetUnsupportedException` remain reserved for compatibility and are not used to reject suspendable execution. -Enforcement samples observed live bytes at checkpoints; it is not an allocation-time peak-memory cap. An instruction, builtin, native serialization, or `CString` allocation can overshoot before it is rejected. The reported native usage is a change in the execution thread's live-byte counter, not exact memory owned by the query. Synchronous host work affects the sample, cross-thread frees can temporarily skew it, and downward baseline ratcheting does not restore lost headroom. +Enforcement samples currently live requested bytes at checkpoints; it is not an allocation-time peak-memory cap. An instruction, builtin, native input parse, serialization, or `CString` allocation can overshoot before it is rejected. Cross-thread frees debit the allocation's originating execution. Synchronous host work on the execution thread is included; managed allocations, host I/O, getter serialization, and other allocators are excluded. + +The following warm-VM baseline note applies to run-to-completion accounting. Suspendable budgets instead charge allocations by their originating execution, so retained pre-execution capacity is not charged unless it is reallocated during the run. A reused VM gets a fresh baseline, but capacities and pools retained before that baseline can make a warm execution allocate differently from a fresh one. Failed terminal execution clears retained state. Public multi-call begin/end scopes are intentionally absent because allocator counters are thread-local and abandoned or cross-thread scopes would be unsafe. The [detailed RVM memory budget documentation](../../docs/limits/memory_budget.md) is authoritative for accounting behavior. diff --git a/bindings/csharp/README.md b/bindings/csharp/README.md index ee70bd5a3..ebdf6772b 100644 --- a/bindings/csharp/README.md +++ b/bindings/csharp/README.md @@ -124,7 +124,7 @@ of being truncated. ### 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. +RVM run-to-completion and suspendable evaluation can use an optional live-memory budget. Each accepted initial `Execute` or `ExecuteEntryPoint` call snapshots the configured budget; `Resume` continues the same budget, including when it runs on a different thread. Program compilation, program loading, and prior `SetDataJson`, `SetInputJson`, and `SetContextJson` calls occur before and outside that budget. Changing or clearing the configuration affects only the next initial execution. ```csharp using var vm = new Rvm(); @@ -143,13 +143,13 @@ catch (RegorusMemoryBudgetExceededException ex) } ``` -The native execution budget includes result JSON serialization and Rust `CString` allocation before the native call returns. Managed UTF-8 decoding and C# `string` allocation after that return are not charged. +The native execution budget includes resume C-string copying and JSON parsing, VM resume work, result JSON serialization, and Rust `CString` allocation before the native call returns. Managed UTF-8 decoding and C# `string` allocation after that return are not charged. This is a cooperative observed-live-bytes budget, not an allocation-time peak-memory cap. A single instruction, builtin, native serialization, or `CString` allocation may exceed the limit before the next checkpoint; configure headroom for that overshoot. -Accounting observes the execution thread rather than allocation ownership. Synchronous host work on that thread contributes to usage, cross-thread frees can temporarily skew observations, and a downward baseline ratchet can permanently remove headroom during an execution. A reused VM receives a new baseline, but retained capacities and pools can make it allocate differently from a fresh VM. See the [RVM memory budget documentation](../../docs/limits/memory_budget.md) for the detailed accounting model. +Suspendable accounting tracks live requested bytes by allocation origin. The execution account is selected only during a synchronous segment; each block retains its account until freed, even on another thread. Successful reallocation charges the full replacement size to the currently selected execution. Run-to-completion accounting keeps its existing execution-thread baseline and downward ratchet. See the [RVM memory budget documentation](../../docs/limits/memory_budget.md) for the detailed accounting model. -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. The process-wide limit exposed by `MemoryLimits` remains a separate safeguard. +Terminal budget or runtime failures invalidate retained results and clear execution state. Missing/unexpected resume input and malformed JSON are retryable while a valid continuation remains. Native status 11 and `RegorusMemoryBudgetUnsupportedException` remain reserved for compatibility and are not used to reject suspendable execution. The process-wide limit exposed by `MemoryLimits` remains a separate safeguard. ## RVM with Registered Host-Await Builtins diff --git a/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs index 34c34ce72..17e7580d4 100644 --- a/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs +++ b/bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs @@ -2,8 +2,12 @@ // Licensed under the MIT License. using System; +using System.Diagnostics; +using System.IO; using System.Linq; +using System.Security.Cryptography; using System.Text.Json; +using System.Threading; using Microsoft.VisualStudio.TestTools.UnitTesting; namespace Regorus.Tests; @@ -11,6 +15,8 @@ namespace Regorus.Tests; [TestClass] public sealed class RvmMemoryBudgetTests { + public TestContext TestContext { get; set; } = null!; + private const ulong TightMemoryBudgetBytes = 64 * 1024; private const string Policy = """ @@ -30,6 +36,27 @@ package limits.memory private const string PreloadedResultEntryPoint = "data.limits.memory.large_string"; + private const string SuspendableHostAwaitPolicy = """ +package limits.memory +import rego.v1 + +result := count(__builtin_host_await( + __builtin_host_await(input.value, "first"), + "second" +)) +"""; + + private const string SuspendableHostAwaitEntryPoint = "data.limits.memory.result"; + private const ulong SuspendableMemoryBudgetBytes = 768 * 1024; + + private const string NativeModulePolicy = """ +package limits.native + +result := 1 +"""; + + private const string NativeModuleEntryPoint = "data.limits.native.result"; + [TestMethod] public void Memory_budget_must_be_non_zero() { @@ -86,13 +113,152 @@ public void Serialization_budget_failure_leaves_error_state() } [TestMethod] - public void Suspendable_execution_rejects_memory_budget() + public void Suspendable_budget_spans_compiled_host_await_thread_migration() { + var modules = new[] { new PolicyModule("memory_budget.rego", SuspendableHostAwaitPolicy) }; + using var program = Program.CompileFromModules( + """{"value":"request"}""", + modules, + new[] { SuspendableHostAwaitEntryPoint }); using var vm = new Rvm(); vm.SetExecutionMode(ExecutionMode.Suspendable); - vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(1024)); + vm.LoadProgram(program); + vm.SetInputJson("""{"value":"request"}"""); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(SuspendableMemoryBudgetBytes)); + + Exception? executeError = null; + Exception? firstResumeError = null; + Exception? secondResumeError = null; + string? finalResult = null; + using var executeFinished = new ManualResetEventSlim(); + using var firstResumeFinished = new ManualResetEventSlim(); + + var executeThread = new Thread(() => + { + try + { + vm.Execute(); + } + catch (Exception ex) + { + executeError = ex; + } + finally + { + executeFinished.Set(); + } + }); + + var firstResumeThread = new Thread(() => + { + try + { + if (!executeFinished.Wait(TimeSpan.FromSeconds(30))) + { + throw new TimeoutException("Execute did not finish."); + } + + vm.Resume("\"first\""); + } + catch (Exception ex) + { + firstResumeError = ex; + } + finally + { + firstResumeFinished.Set(); + } + }); + + var secondResumeThread = new Thread(() => + { + try + { + if (!firstResumeFinished.Wait(TimeSpan.FromSeconds(30))) + { + throw new TimeoutException("First resume did not finish."); + } + + finalResult = vm.Resume("\"second\""); + } + catch (Exception ex) + { + secondResumeError = ex; + } + }); + + executeThread.Start(); + firstResumeThread.Start(); + secondResumeThread.Start(); + executeThread.Join(); + firstResumeThread.Join(); + secondResumeThread.Join(); + + Assert.IsNull(executeError, executeError?.ToString()); + Assert.IsNull(firstResumeError, firstResumeError?.ToString()); + Assert.IsNull(secondResumeError, secondResumeError?.ToString()); + Assert.AreEqual("6", finalResult); + } + + [TestMethod] + public void Suspendable_resume_parse_exhaustion_throws_typed_exception() + { + var modules = new[] { new PolicyModule("memory_budget.rego", SuspendableHostAwaitPolicy) }; + using var program = Program.CompileFromModules( + """{"value":"request"}""", + modules, + new[] { SuspendableHostAwaitEntryPoint }); + using var vm = new Rvm(); + vm.SetExecutionMode(ExecutionMode.Suspendable); + vm.LoadProgram(program); + vm.SetInputJson("""{"value":"request"}"""); + vm.SetMemoryBudgetConfig(new MemoryBudgetConfig(SuspendableMemoryBudgetBytes)); + + vm.Execute(); + vm.ClearMemoryBudgetConfig(); + vm.SetExecutionMode(ExecutionMode.RunToCompletion); + vm.Resume(JsonSerializer.Serialize(new string('x', 450 * 1024))); + + Assert.ThrowsException( + () => vm.Resume(JsonSerializer.Serialize(new string('y', 450 * 1024)))); + StringAssert.Contains(vm.GetExecutionState()!, "Error { error: MemoryBudgetExceeded"); + } + + [TestMethod] + public void Package_native_module_is_identified_after_native_execution() + { + if (!OperatingSystem.IsWindows()) + { + Assert.Inconclusive("The native module identity assertion requires a Windows DLL."); + return; + } + + var modules = new[] { new PolicyModule("native_module.rego", NativeModulePolicy) }; + using var program = Program.CompileFromModules( + "{}", + modules, + new[] { NativeModuleEntryPoint }); + using var vm = new Rvm(); + vm.LoadProgram(program); + + Assert.AreEqual("1", vm.ExecuteEntryPoint(NativeModuleEntryPoint)); + + using var process = Process.GetCurrentProcess(); + var nativeModule = process.Modules + .Cast() + .FirstOrDefault(module => + string.Equals( + module.ModuleName, + "regorus_ffi.dll", + StringComparison.OrdinalIgnoreCase)); + Assert.IsNotNull(nativeModule); - Assert.ThrowsException(() => vm.Execute()); + var modulePath = nativeModule!.FileName; + Assert.IsTrue(Path.IsPathFullyQualified(modulePath), modulePath); + using var moduleStream = File.OpenRead(modulePath); + var moduleHash = Convert.ToHexString(SHA256.HashData(moduleStream)); + TestContext.WriteLine($"REGORUS_NATIVE_MODULE_PATH={modulePath}"); + TestContext.WriteLine($"REGORUS_NATIVE_MODULE_SHA256={moduleHash}"); } private static Program CreateProgram() diff --git a/bindings/csharp/Regorus/NativeMethods.cs b/bindings/csharp/Regorus/NativeMethods.cs index 6859081d4..c7fbb1287 100644 --- a/bindings/csharp/Regorus/NativeMethods.cs +++ b/bindings/csharp/Regorus/NativeMethods.cs @@ -865,7 +865,7 @@ internal enum RegorusStatus : uint /// MemoryBudgetExceeded, /// - /// An RVM memory budget was used with suspendable execution. + /// Reserved for native ABI compatibility. /// MemoryBudgetUnsupportedInSuspendableExecution, } diff --git a/bindings/csharp/Regorus/RegorusMemoryBudgetUnsupportedException.cs b/bindings/csharp/Regorus/RegorusMemoryBudgetUnsupportedException.cs index b1b36d816..2853a221e 100644 --- a/bindings/csharp/Regorus/RegorusMemoryBudgetUnsupportedException.cs +++ b/bindings/csharp/Regorus/RegorusMemoryBudgetUnsupportedException.cs @@ -6,7 +6,7 @@ namespace Regorus { /// - /// The exception thrown when an RVM memory budget is used with suspendable execution. + /// Compatibility exception for the reserved native status 11. /// public sealed class RegorusMemoryBudgetUnsupportedException : InvalidOperationException { diff --git a/bindings/csharp/Regorus/Rvm.cs b/bindings/csharp/Regorus/Rvm.cs index c1776c63c..00c3add54 100644 --- a/bindings/csharp/Regorus/Rvm.cs +++ b/bindings/csharp/Regorus/Rvm.cs @@ -176,7 +176,7 @@ public void SetMaxInstructions(ulong maxInstructions) } /// - /// Configure a fresh memory budget for every run-to-completion execution. + /// Configure the memory budget for the next initial RVM execution. /// /// Memory-budget configuration. public void SetMemoryBudgetConfig(MemoryBudgetConfig config) diff --git a/bindings/ffi/CHANGELOG.md b/bindings/ffi/CHANGELOG.md index a100511db..a79836f8c 100644 --- a/bindings/ffi/CHANGELOG.md +++ b/bindings/ffi/CHANGELOG.md @@ -8,10 +8,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### 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. +- Extend per-VM memory budgets configured through `regorus_rvm_set_memory_budget_config` to + suspendable RVM executions. One requested-byte account and its initial budget span all resumes; + the account outlives the VM while attributed blocks remain live. Native resume JSON parsing plus + immediate result JSON/C-string production are included, and malformed JSON syntax/EOF remains + retryable without changing the active budget. Run-to-completion thread-baseline behavior is + unchanged. Status 10 reports budget exhaustion; status 11 remains reserved. Existing C ABI + signatures, C# budget APIs, and `RegorusMemoryBudgetExceededException` mapping remain unchanged. ## [0.1.0](https://github.com/microsoft/regorus/releases/tag/regorus-ffi-v0.1.0) - 2024-02-08 diff --git a/bindings/ffi/src/common.rs b/bindings/ffi/src/common.rs index 53d37648d..8a259fbe9 100644 --- a/bindings/ffi/src/common.rs +++ b/bindings/ffi/src/common.rs @@ -47,7 +47,8 @@ pub enum RegorusStatus { /// An RVM execution exceeded its configured memory budget. MemoryBudgetExceeded, - /// An RVM memory budget was used with suspendable execution. + /// Reserved for ABI compatibility. + #[allow(dead_code)] MemoryBudgetUnsupportedInSuspendableExecution, } diff --git a/bindings/ffi/src/lib.rs b/bindings/ffi/src/lib.rs index fa97f632a..c20ac4120 100644 --- a/bindings/ffi/src/lib.rs +++ b/bindings/ffi/src/lib.rs @@ -5,6 +5,9 @@ extern crate alloc; +#[cfg(all(test, feature = "allocator-memory-limits", not(miri)))] +pub(crate) static TEST_GLOBAL_MEMORY_LIMIT_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + mod alias_registry; mod allocator; mod common; diff --git a/bindings/ffi/src/limits.rs b/bindings/ffi/src/limits.rs index 43459e02f..a0c2a4ffe 100644 --- a/bindings/ffi/src/limits.rs +++ b/bindings/ffi/src/limits.rs @@ -286,6 +286,9 @@ mod tests { #[test] fn ffi_roundtrips_global_limit() { + let _lock = crate::TEST_GLOBAL_MEMORY_LIMIT_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); let limit = 456_u64; let result = regorus_set_global_memory_limit(limit, true); assert!(matches!(result.status, RegorusStatus::Ok)); diff --git a/bindings/ffi/src/rvm.rs b/bindings/ffi/src/rvm.rs index 814d18d4f..2c8844ade 100644 --- a/bindings/ffi/src/rvm.rs +++ b/bindings/ffi/src/rvm.rs @@ -24,7 +24,9 @@ use regorus::rvm::program::{ DeserializationResult, Program, }; use regorus::rvm::vm::{ExecutionMode, ExecutionState, RegoVM, VmError}; -use regorus::utils::limits::check_memory_limit_if_needed; +use regorus::utils::limits::{ + check_memory_limit_if_needed, with_unowned_memory_budget_scope_for_ffi, +}; use regorus::PolicyModule; use regorus::Value; @@ -57,19 +59,9 @@ 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()) @@ -647,37 +639,44 @@ pub extern "C" fn regorus_rvm_resume( has_value: bool, ) -> RegorusResult { with_unwind_guard(|| { - let output = || -> Result { + let output = || -> Result { let vm = to_shared_ref(vm as *const RegorusRvm)?; let mut guard = vm.try_write()?; - let value = if has_value { - Some(Value::from_json_str(&from_c_str(resume_value_json)?)?) - } else { - None - }; - let result = guard.resume(value)?; - result.to_json_str() + let output = guard.resume_to_c_string_for_ffi(|| { + if has_value { + let resume_json = from_c_str(resume_value_json).map_err(VmError::from)?; + Ok(Some(resume_json)) + } else { + Ok(None) + } + })?; + Ok(RegorusResult::ok_c_string(output)) }(); - to_rvm_string_result(output) + match output { + Ok(result) => result, + Err(err) => to_rvm_error_result(err), + } }) } /// Get the current execution state of the VM. #[no_mangle] pub extern "C" fn regorus_rvm_get_execution_state(vm: *mut RegorusRvm) -> RegorusResult { - with_unwind_guard(|| { - let output = || -> Result { - let vm = to_shared_ref(vm as *const RegorusRvm)?; - let guard = vm.try_read()?; - let state: ExecutionState = guard.execution_state().clone(); - Ok(format!("{:?}", state)) - }(); - - match output { - Ok(json) => RegorusResult::ok_string(json), - Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), - } + with_unowned_memory_budget_scope_for_ffi(|| { + with_unwind_guard(|| { + let output = || -> Result { + let vm = to_shared_ref(vm as *const RegorusRvm)?; + let guard = vm.try_read()?; + let state: ExecutionState = guard.execution_state().clone(); + Ok(format!("{:?}", state)) + }(); + + match output { + Ok(json) => RegorusResult::ok_string(json), + Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), + } + }) }) } @@ -685,22 +684,25 @@ pub extern "C" fn regorus_rvm_get_execution_state(vm: *mut RegorusRvm) -> Regoru mod tests { use super::{ regorus_rvm_drop, regorus_rvm_execute, regorus_rvm_execute_entry_point_by_index, - regorus_rvm_execute_entry_point_by_name, regorus_rvm_get_execution_state, regorus_rvm_new, - regorus_rvm_resume, regorus_rvm_set_data, regorus_rvm_set_memory_budget_config, RegorusRvm, + regorus_rvm_execute_entry_point_by_name, regorus_rvm_get_execution_state, + regorus_rvm_get_host_await_argument, regorus_rvm_get_host_await_identifier, + regorus_rvm_new, regorus_rvm_resume, regorus_rvm_set_data, regorus_rvm_set_execution_mode, + 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::String; use alloc::string::ToString; use alloc::sync::Arc; use alloc::vec; + use alloc::vec::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::rvm::vm::{ExecutionMode, RegoVM, VmError}; use regorus::{Engine, MemoryBudgetConfig, Rc, Value}; const POLICY: &str = r#" @@ -708,9 +710,29 @@ package limits.memory import rego.v1 copy := [value | some value in input] +"#; + + const SUSPENDABLE_HOST_AWAIT_POLICY: &str = r#" +package limits.memory +import rego.v1 + +result := count(__builtin_host_await( + __builtin_host_await(input.value, "first"), + "second" +)) +"#; + + const SUSPENDABLE_REPEAT_RESULT_POLICY: &str = r#" +package limits.memory +import rego.v1 + +value := __builtin_host_await(input.value, "first") +result := [value, value, value, value, value, value, value, value] "#; const TIGHT_MEMORY_BUDGET_BYTES: u64 = 64 * 1024; + const NESTED_GETTER_MEMORY_BUDGET_BYTES: u64 = 8 * 1024; + const NESTED_GETTER_REPETITIONS: usize = 2048; fn memory_budget(limit: u64) -> MemoryBudgetConfig { MemoryBudgetConfig { @@ -739,30 +761,96 @@ copy := [value | some value in input] 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 compiled_suspendable_host_await_program() -> Arc { + let entrypoint = Rc::from("data.limits.memory.result"); + let mut engine = Engine::new(); + engine + .add_policy( + "memory_budget.rego".into(), + SUSPENDABLE_HOST_AWAIT_POLICY.into(), + ) + .expect("add policy"); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy"); + Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) + .expect("compile VM program") + } + + fn compiled_suspendable_repeat_result_program() -> Arc { + let entrypoint = Rc::from("data.limits.memory.result"); + let mut engine = Engine::new(); + engine + .add_policy( + "memory_budget.rego".into(), + SUSPENDABLE_REPEAT_RESULT_POLICY.into(), + ) + .expect("add policy"); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy"); + Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]) + .expect("compile VM program") + } + + fn suspendable_host_await_vm(program: &Arc, budget: Option) -> RegoVM { + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(Arc::clone(program)); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + if let Some(limit) = budget { + vm.set_memory_budget_config(Some(memory_budget(limit))); + } + vm + } + + fn assert_getter_scope_restored( + program: &Arc, + child_vm: *mut RegorusRvm, + getter: extern "C" fn(*mut RegorusRvm) -> RegorusResult, + expected_output: &str, + expected_output_is_prefix: bool, + ) { + let mut outer_vm = + suspendable_host_await_vm(program, Some(NESTED_GETTER_MEMORY_BUDGET_BYTES)); + outer_vm + .execute() + .expect("outer VM should suspend at HostAwait within its budget"); + + let mut getter_result = None; + let mut parent_probe: Option> = None; + let result = outer_vm.resume_to_c_string_for_ffi(|| { + getter_result = Some(getter(child_vm)); + parent_probe = Some(vec![0u8; NESTED_GETTER_MEMORY_BUDGET_BYTES as usize]); + core::hint::black_box(parent_probe.as_ref()); + Err(VmError::Internal { + message: String::from("stop after the post-getter parent allocation"), + pc: 0, + }) + }); + + let is_budget_error = match &result { + Err(VmError::MemoryBudgetExceeded { .. }) => true, + Ok(_) => false, + Err(_) => false, + }; + assert!( + is_budget_error, + "post-getter parent allocation should be charged before result serialization: {result:?}" + ); + + let getter_result = getter_result.expect("getter should be invoked"); + assert!(matches!(getter_result.status, RegorusStatus::Ok)); + assert!(!getter_result.output.is_null()); + let output = unsafe { CStr::from_ptr(getter_result.output) } + .to_str() + .expect("getter output UTF-8"); + if expected_output_is_prefix { + assert!(output.starts_with(expected_output)); + } else { + assert_eq!(output, expected_output); + } + regorus_result_drop(getter_result); } fn preloaded_result_program() -> Arc { @@ -903,29 +991,364 @@ copy := [value | some value in input] } #[test] - fn ffi_resume_reports_unsupported_memory_budget_status() { + fn ffi_suspendable_budget_spans_host_await_resumes_and_charges_json_input() { let mut vm = RegoVM::new(); vm.set_execution_mode(ExecutionMode::Suspendable); - vm.load_program(host_await_program()); - vm.execute().expect("suspend execution"); - + vm.load_program(compiled_suspendable_host_await_program()); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(768 * 1024))); let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); - let set_result = regorus_rvm_set_memory_budget_config( + + let mode = regorus_rvm_set_execution_mode(vm, 1); + assert!(matches!(mode.status, RegorusStatus::Ok)); + regorus_result_drop(mode); + let execute = regorus_rvm_execute(vm); + assert!(matches!(execute.status, RegorusStatus::Ok)); + regorus_result_drop(execute); + + let mode = regorus_rvm_set_execution_mode(vm, 0); + assert!(matches!(mode.status, RegorusStatus::Ok)); + regorus_result_drop(mode); + 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 first_resume = + CString::new(format!("\"{}\"", "x".repeat(450 * 1024))).expect("first resume JSON"); + let first = regorus_rvm_resume(vm, first_resume.as_ptr(), true); + assert!(matches!(first.status, RegorusStatus::Ok)); + regorus_result_drop(first); + + let second_resume = + CString::new(format!("\"{}\"", "y".repeat(450 * 1024))).expect("second resume JSON"); + assert_memory_budget_failure_state( vm, + regorus_rvm_resume(vm, second_resume.as_ptr(), true), + ); + regorus_rvm_drop(vm); + } + + #[test] + fn ffi_suspendable_resume_validation_errors_preserve_host_await_continuation() { + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(compiled_suspendable_host_await_program()); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(2 * 1024 * 1024))); + let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); + + let execute = regorus_rvm_execute(vm); + assert!(matches!(execute.status, RegorusStatus::Ok)); + regorus_result_drop(execute); + assert_execution_state(vm, "Suspended {"); + + let missing = regorus_rvm_resume(vm, core::ptr::null(), false); + assert!(!matches!(missing.status, RegorusStatus::Ok)); + regorus_result_drop(missing); + assert_execution_state(vm, "Suspended {"); + + let malformed_json = CString::new("{").expect("malformed JSON CString"); + let malformed = regorus_rvm_resume(vm, malformed_json.as_ptr(), true); + assert!(!matches!(malformed.status, RegorusStatus::Ok)); + regorus_result_drop(malformed); + assert_execution_state(vm, "Suspended {"); + + let first_resume = CString::new("\"first\"").expect("first resume JSON"); + let first = regorus_rvm_resume(vm, first_resume.as_ptr(), true); + assert!(matches!(first.status, RegorusStatus::Ok)); + regorus_result_drop(first); + assert_execution_state(vm, "Suspended {"); + + let second_resume = CString::new("\"second\"").expect("second resume JSON"); + let second = regorus_rvm_resume(vm, second_resume.as_ptr(), true); + assert!(matches!(second.status, RegorusStatus::Ok)); + regorus_result_drop(second); + assert_execution_state(vm, "Completed {"); + + regorus_rvm_drop(vm); + } + + #[test] + fn correction_ffi_malformed_resume_is_retryable_under_global_memory_pressure() { + let _lock = crate::TEST_GLOBAL_MEMORY_LIMIT_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + struct ResetGlobalMemoryLimit; + impl Drop for ResetGlobalMemoryLimit { + fn drop(&mut self) { + regorus::set_global_memory_limit(None); + } + } + + let _reset_global_limit = ResetGlobalMemoryLimit; + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(compiled_suspendable_host_await_program()); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(2 * 1024 * 1024))); + let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); + + let execute = regorus_rvm_execute(vm); + assert!(matches!(execute.status, RegorusStatus::Ok)); + regorus_result_drop(execute); + assert_execution_state(vm, "Suspended {"); + + regorus::set_global_memory_limit(Some(1)); + let malformed_json = CString::new("{").expect("malformed JSON CString"); + let malformed = regorus_rvm_resume(vm, malformed_json.as_ptr(), true); + assert!(matches!(malformed.status, RegorusStatus::Error)); + regorus_result_drop(malformed); + regorus::set_global_memory_limit(None); + + assert_execution_state(vm, "Suspended {"); + let valid_resume = CString::new("\"response\"").expect("valid resume JSON"); + let resumed = regorus_rvm_resume(vm, valid_resume.as_ptr(), true); + assert!(matches!(resumed.status, RegorusStatus::Ok)); + regorus_result_drop(resumed); + assert_execution_state(vm, "Suspended {"); + let second_resume = CString::new("\"done\"").expect("second resume JSON"); + let completed = regorus_rvm_resume(vm, second_resume.as_ptr(), true); + assert!(matches!(completed.status, RegorusStatus::Ok)); + regorus_result_drop(completed); + assert_execution_state(vm, "Completed {"); + regorus_rvm_drop(vm); + } + + #[test] + fn ffi_valid_resume_global_limit_failure_terminalizes_with_generic_error_status() { + let _lock = crate::TEST_GLOBAL_MEMORY_LIMIT_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + struct ResetGlobalMemoryLimit; + impl Drop for ResetGlobalMemoryLimit { + fn drop(&mut self) { + regorus::set_global_memory_limit(None); + } + } + + let _reset_global_limit = ResetGlobalMemoryLimit; + regorus::set_global_memory_limit(None); + regorus::utils::limits::check_memory_limit_if_needed() + .expect("disabled global limit resets check state"); + + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(compiled_suspendable_host_await_program()); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(128 * 1024 * 1024))); + let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); + + let execute = regorus_rvm_execute(vm); + assert!(matches!(execute.status, RegorusStatus::Ok)); + regorus_result_drop(execute); + assert_execution_state(vm, "Suspended {"); + + let mut json = String::with_capacity(2 + 2 * 4096); + json.push('['); + for index in 0..4096 { + if index > 0 { + json.push(','); + } + json.push('0'); + } + json.push(']'); + let json = CString::new(json).expect("valid numeric array JSON"); + + regorus::set_global_memory_limit(Some(1)); + let baseline_usage = match regorus::utils::limits::check_global_memory_limit() { + Err(regorus::LimitError::MemoryLimitExceeded { usage, limit: 1 }) => usage, + other => panic!("expected to observe current usage above one byte, got {other:?}"), + }; + let global_limit = baseline_usage.saturating_add(32 * 1024); + regorus::set_global_memory_limit(Some(global_limit)); + + let failure = regorus_rvm_resume(vm, json.as_ptr(), true); + let post_parse_check = regorus::utils::limits::check_global_memory_limit(); + println!( + "FFI resume conversion RED: baseline={baseline_usage}, limit={global_limit}, post_parse_check={post_parse_check:?}, status={:?}", + failure.status + ); + assert!(matches!(failure.status, RegorusStatus::Error)); + assert!(failure.output.is_null()); + regorus_result_drop(failure); + assert!( + post_parse_check.is_ok(), + "partial Value allocations should be gone while the input CString remains" + ); + assert_execution_state(vm, "Error { error: MemoryLimitExceeded"); + + let rejected_json = CString::new("\"rejected\"").expect("rejected resume JSON"); + let rejected = regorus_rvm_resume(vm, rejected_json.as_ptr(), true); + assert!(matches!(rejected.status, RegorusStatus::Error)); + assert!(rejected.output.is_null()); + regorus_result_drop(rejected); + assert_execution_state(vm, "Error { error: MemoryLimitExceeded"); + + regorus::set_global_memory_limit(None); + regorus::utils::limits::check_memory_limit_if_needed() + .expect("clearing the global limit resets check state"); + let restarted = regorus_rvm_execute(vm); + assert!(matches!(restarted.status, RegorusStatus::Ok)); + regorus_result_drop(restarted); + assert_execution_state(vm, "Suspended {"); + let first_resume = CString::new("\"independent\"").expect("first resume JSON"); + let first = regorus_rvm_resume(vm, first_resume.as_ptr(), true); + assert!(matches!(first.status, RegorusStatus::Ok)); + regorus_result_drop(first); + let second_resume = CString::new("\"done\"").expect("second resume JSON"); + let second = regorus_rvm_resume(vm, second_resume.as_ptr(), true); + assert!(matches!(second.status, RegorusStatus::Ok)); + regorus_result_drop(second); + assert_execution_state(vm, "Completed {"); + + regorus_rvm_drop(vm); + } + + #[test] + fn ffi_read_only_getters_exclude_nested_output_and_restore_parent_budget_scope() { + let program = compiled_suspendable_host_await_program(); + let mut child_vm = suspendable_host_await_vm(&program, None); + child_vm + .execute() + .expect("child VM should suspend at HostAwait"); + let child_vm = Box::into_raw(Box::new(RegorusRvm::new(child_vm))); + + let mut outer_vm = + suspendable_host_await_vm(&program, Some(NESTED_GETTER_MEMORY_BUDGET_BYTES)); + outer_vm + .execute() + .expect("outer VM should suspend at HostAwait within its budget"); + + let mut getter_results = Vec::with_capacity(NESTED_GETTER_REPETITIONS * 3); + let mut parent_allocation = None; + let resumed = outer_vm.resume_to_c_string_for_ffi(|| { + for _ in 0..NESTED_GETTER_REPETITIONS { + getter_results.push(regorus_rvm_get_execution_state(child_vm)); + getter_results.push(regorus_rvm_get_host_await_argument(child_vm)); + getter_results.push(regorus_rvm_get_host_await_identifier(child_vm)); + } + parent_allocation = Some(vec![0u8; 64]); + core::hint::black_box(parent_allocation.as_ref()); + Ok(Some(String::from("\"response\""))) + }); + + assert!( + resumed.is_ok(), + "getter serialization should be excluded from the unrelated outer budget: {resumed:?}" + ); + assert!(!resumed + .expect("outer resume should succeed") + .as_bytes() + .is_empty()); + assert!( + alloc::format!("{:?}", outer_vm.execution_state()).starts_with("Suspended {"), + "outer continuation should remain suspended after its first resume" + ); + assert_eq!( + getter_results.len(), + NESTED_GETTER_REPETITIONS * 3, + "all three getters should run for every retained output" + ); + for result in &getter_results { + assert!(matches!(result.status, RegorusStatus::Ok)); + assert!(!result.output.is_null()); + } + + let mut retained_c_string_bytes = [0usize; 3]; + for getter_results in getter_results.chunks_exact(3) { + for (total, result) in retained_c_string_bytes.iter_mut().zip(getter_results) { + let bytes_with_nul = unsafe { CStr::from_ptr(result.output) } + .to_bytes_with_nul() + .len(); + *total = total + .checked_add(bytes_with_nul) + .expect("retained getter output byte count"); + } + } + for (getter, bytes) in [ + ("execution state", retained_c_string_bytes[0]), + ("HostAwait argument", retained_c_string_bytes[1]), + ("HostAwait identifier", retained_c_string_bytes[2]), + ] { + assert!( + bytes > NESTED_GETTER_MEMORY_BUDGET_BYTES as usize, + "{getter} CString bytes {bytes} must exceed the outer budget of \ + {NESTED_GETTER_MEMORY_BUDGET_BYTES}" + ); + } + + let mut first_outputs = getter_results.iter(); + let state = first_outputs.next().expect("execution-state output"); + let argument = first_outputs.next().expect("HostAwait argument output"); + let identifier = first_outputs.next().expect("HostAwait identifier output"); + assert!(unsafe { CStr::from_ptr(state.output) } + .to_str() + .expect("execution-state UTF-8") + .starts_with("Suspended {")); + assert_eq!( + unsafe { CStr::from_ptr(argument.output) } + .to_str() + .expect("HostAwait argument UTF-8"), + "\"request\"" + ); + assert_eq!( + unsafe { CStr::from_ptr(identifier.output) } + .to_str() + .expect("HostAwait identifier UTF-8"), + "first" + ); + + for result in getter_results.drain(..) { + regorus_result_drop(result); + } + + assert_getter_scope_restored( + &program, + child_vm, + regorus_rvm_get_execution_state, + "Suspended {", true, - RegorusMemoryBudgetConfig { - limit_bytes: 1024 * 1024, - }, ); - assert!(matches!(set_result.status, RegorusStatus::Ok)); - regorus_result_drop(set_result); + assert_getter_scope_restored( + &program, + child_vm, + regorus_rvm_get_host_await_argument, + "\"request\"", + false, + ); + assert_getter_scope_restored( + &program, + child_vm, + regorus_rvm_get_host_await_identifier, + "first", + false, + ); + + regorus_rvm_drop(child_vm); + } + + #[test] + fn ffi_suspendable_result_serialization_uses_active_memory_budget() { + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(compiled_suspendable_repeat_result_program()); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(512 * 1024))); + let vm = Box::into_raw(Box::new(RegorusRvm::new(vm))); + + let execute = regorus_rvm_execute(vm); + assert!(matches!(execute.status, RegorusStatus::Ok)); + regorus_result_drop(execute); + assert_execution_state(vm, "Suspended {"); + + let large_resume = + CString::new(format!("\"{}\"", "x".repeat(128 * 1024))).expect("large resume JSON"); + assert_memory_budget_failure_state(vm, regorus_rvm_resume(vm, large_resume.as_ptr(), true)); - let result = regorus_rvm_resume(vm, ptr::null(), false); - assert!(matches!( - result.status, - RegorusStatus::MemoryBudgetUnsupportedInSuspendableExecution - )); - regorus_result_drop(result); regorus_rvm_drop(vm); } @@ -1243,21 +1666,23 @@ pub extern "C" fn regorus_rvm_set_host_await_responses( /// or None if the VM is not in a HostAwait-suspended state. #[no_mangle] pub extern "C" fn regorus_rvm_get_host_await_argument(vm: *mut RegorusRvm) -> RegorusResult { - with_unwind_guard(|| { - let output = || -> Result> { - let vm = to_shared_ref(vm as *const RegorusRvm)?; - let guard = vm.try_read()?; - match guard.get_host_await_argument() { - Some(arg) => Ok(Some(arg.to_json_str()?)), - None => Ok(None), - } - }(); + with_unowned_memory_budget_scope_for_ffi(|| { + with_unwind_guard(|| { + let output = || -> Result> { + let vm = to_shared_ref(vm as *const RegorusRvm)?; + let guard = vm.try_read()?; + match guard.get_host_await_argument() { + Some(arg) => Ok(Some(arg.to_json_str()?)), + None => Ok(None), + } + }(); - match output { - Ok(Some(json)) => RegorusResult::ok_string(json), - Ok(None) => RegorusResult::ok_void(), - Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), - } + match output { + Ok(Some(json)) => RegorusResult::ok_string(json), + Ok(None) => RegorusResult::ok_void(), + Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), + } + }) }) } @@ -1296,18 +1721,20 @@ fn host_await_identifier_string(identifier: Option<&Value>) -> Result RegorusResult { - with_unwind_guard(|| { - let output = || -> Result> { - let vm = to_shared_ref(vm as *const RegorusRvm)?; - let guard = vm.try_read()?; - host_await_identifier_string(guard.get_host_await_identifier()) - }(); - - match output { - Ok(Some(identifier)) => RegorusResult::ok_string(identifier), - Ok(None) => RegorusResult::ok_void(), - Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), - } + with_unowned_memory_budget_scope_for_ffi(|| { + with_unwind_guard(|| { + let output = || -> Result> { + let vm = to_shared_ref(vm as *const RegorusRvm)?; + let guard = vm.try_read()?; + host_await_identifier_string(guard.get_host_await_identifier()) + }(); + + match output { + Ok(Some(identifier)) => RegorusResult::ok_string(identifier), + Ok(None) => RegorusResult::ok_void(), + Err(err) => RegorusResult::err_with_message(RegorusStatus::Error, err.to_string()), + } + }) }) } diff --git a/docs/limits/memory_budget.md b/docs/limits/memory_budget.md index ad73e3377..47632a23c 100644 --- a/docs/limits/memory_budget.md +++ b/docs/limits/memory_budget.md @@ -1,8 +1,12 @@ # RVM memory budgets -RVM run-to-completion evaluation supports an optional memory budget when Regorus is built with the `allocator-memory-limits` feature. +RVM run-to-completion and suspendable execution support 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. +Each accepted initial `execute` or entry-point execution snapshots the configured +budget. `set_memory_budget_config` and `set_execution_mode` configure the next +initial execution; changing or clearing either setting while a suspendable run +is suspended does not change that run. Resuming never starts a new budget. ```rust use core::num::NonZeroU64; @@ -15,42 +19,83 @@ vm.set_memory_budget_config(Some(MemoryBudgetConfig { })); ``` -No configured budget preserves existing RVM behavior. A zero-byte budget is not representable in Rust and is rejected by language bindings. - -## Included work - -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 `CString` 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 samples memory at VM instruction checkpoints and once before returning a successful result. The C FFI also samples after native result JSON serialization and `CString` allocation. Enforcement is cooperative and checkpoint-based, not an allocation-time peak-memory hard cap. One instruction, builtin, result serialization, or `CString` allocation can temporarily overshoot the configured limit before the next sample. For example, a builtin can allocate far more than its remaining headroom and be rejected only after it returns; an allocation created and freed entirely between samples may not be observed at all. Callers should configure enough headroom for this overshoot. - -Accounting uses the execution thread's live-byte counter rather than allocation ownership. Allocations and frees performed by synchronous host callbacks or builtins on that thread affect the observation. Objects allocated on one thread and freed on another can temporarily skew thread-level observations while allocator counters are reconciled. The control therefore bounds observed additional live bytes on the execution thread, not memory owned by a query or attributed across threads. - -When a sampled live-byte count falls below the current baseline, Regorus lowers the baseline so an observed same-thread free does not grant headroom to later allocations. This downward ratchet is never restored during the execution and can make the effective budget stricter than configured after unrelated or legitimate frees. A foreign free can still offset evaluation allocations when both occur between samples. - -A fresh budget means a new baseline is captured for each execution, not that the VM is returned to a newly constructed state. Reused VMs retain capacities and pools that are already live before the baseline. An identical policy and input can therefore allocate differently, and may have a different budget outcome, on a warm VM than on a fresh VM. - -Exhaustion returns `VmError::MemoryBudgetExceeded`, including: - -- `usage`, the observed execution-thread live-byte increase above the ratcheted baseline; this is diagnostic thread-level change, not exact query-owned memory -- 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 `CString` 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. +A zero-byte budget is not representable in Rust and is rejected by language +bindings. With no configured budget, existing RVM behavior is unchanged. + +## Accounting scope + +Run-to-completion execution retains its existing thread-baseline accounting: +sampled live bytes are observed on the execution thread, with the baseline +ratcheted downward when usage falls. Cross-thread frees can therefore skew its +observations. Suspendable execution instead snapshots one thread-safe account +for the accepted execution. It tracks currently live requested bytes allocated +while that execution is selected, including retained VM frames, registers, +caches, continuations, and allocations made synchronously by builtins or +callbacks. Allocator metadata and allocator-rounded capacity are not charged. + +The account is selected only during each synchronous execution segment. Each +allocation records its requested size and account; freeing the block debits that +same account on any thread, including after the VM or execution has ended. +Nested and concurrent executions therefore have independent accounts, and a +resume may run on a different thread without moving or resetting its budget. +The account remains alive while any attributed allocation remains live. +An independent unbudgeted VM entry, resume, or immediate native conversion +temporarily selects no execution account, even when nested inside a budgeted +callback; the previous selector is restored on return or unwind. This selector +is separate from run-to-completion thread-baseline accounting. + +A successful `realloc` is treated as replacing the old block: its old account is +debited and the full new requested size is charged to the account selected for +the reallocation, or to no execution if none is selected. A failed `realloc` +preserves the old block and its attribution. Allocations made before execution +are not charged unless a successful reallocation replaces them during that +execution. + +Program compilation and loading, data/input/context loading, host I/O, cache +getter serialization, managed C# encoding/decoding, other allocators, and +detached work are outside the budget. For native FFI execute and resume calls, +the C-string copy and JSON parsing of resume input, VM resume work, immediate +result JSON serialization, and native `CString` allocation are included. +Managed C# string allocation after the native call returns is excluded. + +The process-global requested-byte limit remains independent and its counters +are not changed by per-execution attribution. At shared checkpoints the +per-execution budget takes precedence over a global memory-limit error. + +## Enforcement and lifecycle + +Enforcement is cooperative, not an allocation-time peak-memory cap. Regorus +checks before instruction dispatch, after resume input conversion, before +returning from a suspension point, after terminal execution, and after native +result serialization. A single instruction, builtin, callback, input parse, or +serialization operation can temporarily overshoot before the next check. +Usage equal to the configured limit succeeds; usage greater than it fails. + +Exhaustion returns `VmError::MemoryBudgetExceeded`, including the observed live +requested bytes, configured budget, and VM program counter. A terminal failure +of a budgeted suspendable execution invalidates retained and completed +results, releases execution state, and records `ExecutionState::Error`. +Unbudgeted initial failures also record `Error` while preserving legacy +register contents. Missing or unexpected resume values and native JSON +syntax/EOF errors are retryable and preserve a valid suspended continuation. +A later global or per-execution counter observation does not turn malformed +syntax into a terminal resource failure; actual typed resource failures during +conversion remain terminal. + +The C FFI reports `RegorusStatus::MemoryBudgetExceeded` (status 10), including +when native resume parsing or immediate result serialization exceeds the +budget. Status 11 remains reserved for ABI compatibility. C# throws +`RegorusMemoryBudgetExceededException`. The FFI budget window closes on success, +terminal error, or unwinding; allocations already returned to a host remain +attributed until freed. + +Loading a new program while a budgeted suspendable execution is active abandons +that continuation and retires its budget. Starting another initial execution +also abandons any prior continuation. Terminal success retires enforcement, +but the allocation account remains available to debit outstanding exported +blocks. Run-to-completion accounting and ordinary no-budget execution retain +their existing behavior. + +There is no public multi-call begin/end memory-budget scope; the RVM owns the +account and retires enforcement with its execution lifecycle. Rust, C FFI, and +C# are supported; other bindings require follow-up work. diff --git a/docs/rvm/vm-runtime.md b/docs/rvm/vm-runtime.md index efc895353..5ce0d8283 100644 --- a/docs/rvm/vm-runtime.md +++ b/docs/rvm/vm-runtime.md @@ -34,6 +34,8 @@ fields to responsibilities. | `step_mode` | Enables single-step suspension after each instruction. | `execution.rs` | | `host_await_responses` | Pre-scripted responses keyed by identifier (run-to-completion). | `dispatch.rs` | | `execution_mode` | `RunToCompletion` or `Suspendable`. | `execution.rs` | +| `next_execution_mode` | Mode selected for the next initial execution. | `machine.rs` | +| memory-budget account | Per-execution requested-byte owner retained across suspendable segments. | `machine.rs`, allocator shim | | `frame_pc_overridden` | Tracks manual PC updates inside frames. | `execution.rs`, `loops.rs`, `comprehension.rs` | | `strict_builtin_errors` | Configures builtin failure handling (error vs `undefined`). | `machine.rs`, `arithmetic.rs`, `dispatch.rs` | @@ -102,6 +104,17 @@ fields to responsibilities. `resume(resume_value)` to continue. - Completion: when `execution_stack` becomes empty the VM sets `ExecutionState::Completed { result }`. +- With `allocator-memory-limits`, one requested-byte account is snapshotted for + the initial execution and retained across suspensions. The account is selected + only while a synchronous segment runs; allocation headers preserve ownership + for frees on other threads. Updating the configured budget or execution mode + affects only the next initial execution. Native FFI resume parsing and + immediate result serialization remain within the active budget. +- An independent unbudgeted VM entry, resume, or native conversion nested inside + a budgeted callback temporarily masks the caller's account and restores it on + return or unwind. This does not alter run-to-completion thread-baseline + accounting. Initial suspendable failures publish `ExecutionState::Error`; + unbudgeted failures retain the legacy register contents. `execution_model.rs` defines the frame types and state machine: @@ -277,4 +290,3 @@ internal runtime field and is not a public C# getter. combines nested loops, comprehensions, function calls, and host awaits. --- - diff --git a/mimalloc/src/limits.rs b/mimalloc/src/limits.rs index df1588dff..9e2303088 100644 --- a/mimalloc/src/limits.rs +++ b/mimalloc/src/limits.rs @@ -7,8 +7,12 @@ //! Otherwise the functions become no-ops so the rest of the crate can call //! them unconditionally. +use core::alloc::Layout; use core::cell::Cell; -use core::sync::atomic::{AtomicI64, Ordering}; +use core::ffi::c_void; +use core::marker::PhantomData; +use core::sync::atomic::{AtomicI64, AtomicU64, AtomicUsize, Ordering}; +use mimalloc_sys::{mi_free, mi_malloc_aligned}; use std::thread_local; // Total memory allocated through the global allocator. @@ -21,6 +25,9 @@ const DEFAULT_THREAD_FLUSH_THRESHOLD: i64 = 1024 * 1024; static THREAD_FLUSH_THRESHOLD: AtomicI64 = AtomicI64::new(DEFAULT_THREAD_FLUSH_THRESHOLD); thread_local! { + // The account is selected only for the current synchronous execution segment. + static CURRENT_MEMORY_BUDGET_ACCOUNT: Cell<*mut MemoryBudgetAccountState> = + const { Cell::new(core::ptr::null_mut()) }; // Each thread tracks its own allocations and deallocations. Periodically the thread's // stats are flushed into the global counters. Such a design minimizes contention on the // global atomic variables on each allocation/free operation. @@ -29,6 +36,183 @@ thread_local! { static THREAD_FLUSHED_SINCE_CHECK: Cell = const { Cell::new(false) }; } +#[derive(Debug)] +struct MemoryBudgetAccountState { + references: AtomicUsize, + live_requested_bytes: AtomicU64, +} + +/// Thread-safe requested-byte accounting for one suspendable execution. +/// +/// The account state is allocated directly from mimalloc so creating an account cannot recurse +/// through the Rust global allocator whose blocks it tracks. +#[derive(Debug)] +pub struct MemoryBudgetAccount { + state: core::ptr::NonNull, +} + +// SAFETY: the shared state contains only atomics and is kept alive by the intrusive reference +// count held by this handle and by every allocation header that records it. +unsafe impl Send for MemoryBudgetAccount {} +unsafe impl Sync for MemoryBudgetAccount {} + +impl MemoryBudgetAccount { + /// Create an empty account backed by non-recursive mimalloc storage. + pub fn new() -> Self { + let layout = Layout::new::(); + let pointer = unsafe { + mi_malloc_aligned(layout.size(), layout.align()).cast::() + }; + if pointer.is_null() { + std::alloc::handle_alloc_error(layout); + } + + unsafe { + pointer.write(MemoryBudgetAccountState { + references: AtomicUsize::new(1), + live_requested_bytes: AtomicU64::new(0), + }); + } + + Self { + state: unsafe { core::ptr::NonNull::new_unchecked(pointer) }, + } + } + + /// Return the current live requested bytes attributed to this execution. + pub fn live_bytes(&self) -> u64 { + unsafe { + self.state + .as_ref() + .live_requested_bytes + .load(Ordering::Relaxed) + } + } + + #[cfg(test)] + pub(crate) fn reference_count_for_test(&self) -> usize { + unsafe { self.state.as_ref().references.load(Ordering::Relaxed) } + } + + /// Select this account for the duration of a synchronous execution segment. + pub fn with_scope(&self, operation: impl FnOnce() -> R) -> R { + let previous = + CURRENT_MEMORY_BUDGET_ACCOUNT.with(|selected| selected.replace(self.state.as_ptr())); + let _scope = MemoryBudgetAccountScope { + previous, + _account: PhantomData, + }; + operation() + } +} + +/// Select no execution account for the duration of an independent synchronous operation. +#[doc(hidden)] +pub fn with_unowned_memory_budget_scope(operation: impl FnOnce() -> R) -> R { + let previous = + CURRENT_MEMORY_BUDGET_ACCOUNT.with(|selected| selected.replace(core::ptr::null_mut())); + let _scope = UnownedMemoryBudgetAccountScope { previous }; + operation() +} + +impl Clone for MemoryBudgetAccount { + fn clone(&self) -> Self { + unsafe { retain_memory_budget_account(self.state.as_ptr().cast()) }; + Self { state: self.state } + } +} + +impl Drop for MemoryBudgetAccount { + fn drop(&mut self) { + unsafe { release_memory_budget_account(self.state.as_ptr().cast()) }; + } +} + +/// Restores the previously selected account when an execution segment ends or unwinds. +struct MemoryBudgetAccountScope<'a> { + previous: *mut MemoryBudgetAccountState, + _account: PhantomData<&'a MemoryBudgetAccount>, +} + +impl Drop for MemoryBudgetAccountScope<'_> { + fn drop(&mut self) { + CURRENT_MEMORY_BUDGET_ACCOUNT.with(|selected| selected.set(self.previous)); + } +} + +struct UnownedMemoryBudgetAccountScope { + previous: *mut MemoryBudgetAccountState, +} + +impl Drop for UnownedMemoryBudgetAccountScope { + fn drop(&mut self) { + CURRENT_MEMORY_BUDGET_ACCOUNT.with(|selected| selected.set(self.previous)); + } +} + +/// Return the account selected for allocations on this thread. +pub(crate) fn current_memory_budget_account() -> *mut c_void { + CURRENT_MEMORY_BUDGET_ACCOUNT + .with(Cell::get) + .cast::() +} + +/// Retain the account reference stored in an allocation header. +pub(crate) unsafe fn retain_memory_budget_account(account: *mut c_void) { + if !account.is_null() { + let _ = (*account.cast::()) + .references + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |references| { + Some(references.saturating_add(1)) + }); + } +} + +/// Release the account reference stored in an allocation header. +pub(crate) unsafe fn release_memory_budget_account(account: *mut c_void) { + if account.is_null() { + return; + } + + let state = account.cast::(); + let Ok(references) = + (*state) + .references + .fetch_update(Ordering::Release, Ordering::Relaxed, |references| { + references.checked_sub(1) + }) + else { + return; + }; + if references == 1 { + core::sync::atomic::fence(Ordering::Acquire); + core::ptr::drop_in_place(state); + mi_free(state.cast::()); + } +} + +/// Add requested bytes to the account stored in an allocation header. +pub(crate) unsafe fn record_memory_budget_alloc(account: *mut c_void, size: usize) { + if !account.is_null() { + let _ = (*account.cast::()) + .live_requested_bytes + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |live_bytes| { + Some(live_bytes.saturating_add(u64::try_from(size).unwrap_or(u64::MAX))) + }); + } +} + +/// Subtract requested bytes from the account stored in an allocation header. +pub(crate) unsafe fn record_memory_budget_free(account: *mut c_void, size: usize) { + if !account.is_null() { + let _ = (*account.cast::()) + .live_requested_bytes + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |live_bytes| { + Some(live_bytes.saturating_sub(u64::try_from(size).unwrap_or(u64::MAX))) + }); + } +} + /// Process-wide view of allocator usage at the moment of sampling. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct GlobalAllocationStats { diff --git a/mimalloc/src/mimalloc.rs b/mimalloc/src/mimalloc.rs index b6cda90ac..fcd215ac3 100644 --- a/mimalloc/src/mimalloc.rs +++ b/mimalloc/src/mimalloc.rs @@ -3,6 +3,10 @@ use core::alloc::{GlobalAlloc, Layout}; use core::ffi::c_void; +#[cfg(feature = "allocator-memory-limits")] +use core::mem::{align_of, size_of}; +#[cfg(feature = "allocator-memory-limits")] +use core::ptr; use mimalloc_sys::{ mi_free, mi_malloc_aligned, mi_realloc_aligned, mi_zalloc_aligned, MI_ALIGNMENT_MAX, }; @@ -11,62 +15,229 @@ use mimalloc_sys::{ pub use crate::limits::{ 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, + set_thread_flush_threshold, GlobalAllocationStats, MemoryBudgetAccount, ThreadAllocationStats, +}; + +#[cfg(feature = "allocator-memory-limits")] +use crate::limits::{ + current_memory_budget_account, record_memory_budget_alloc, record_memory_budget_free, + release_memory_budget_account, retain_memory_budget_account, }; + +#[cfg(feature = "allocator-memory-limits")] +#[repr(C)] +#[derive(Clone, Copy)] +struct AllocationHeader { + requested_size: usize, + account: *mut c_void, +} + +#[cfg(feature = "allocator-memory-limits")] +fn allocation_header_padding(align: usize) -> usize { + // Layout alignments are powers of two; this rounds the prefix without overflowing. + size_of::().wrapping_neg() & align.saturating_sub(1) +} + +#[cfg(feature = "allocator-memory-limits")] +fn allocation_header_layout(layout: Layout) -> Option<(usize, usize, usize)> { + let header_size = size_of::(); + let align = layout.align(); + let Some(offset) = header_size.checked_add(allocation_header_padding(align)) else { + return None; + }; + let Some(raw_size) = offset.checked_add(layout.size()) else { + return None; + }; + let raw_align = align.max(align_of::()); + Some((offset, raw_align, raw_size)) +} + +#[cfg(feature = "allocator-memory-limits")] +unsafe fn allocate_with_account(layout: Layout, account: *mut c_void, zeroed: bool) -> *mut u8 { + let Some((offset, raw_align, raw_size)) = allocation_header_layout(layout) else { + return ptr::null_mut(); + }; + let base = if zeroed { + mi_zalloc_aligned(raw_size, raw_align) + } else { + mi_malloc_aligned(raw_size, raw_align) + }; + if base.is_null() { + return ptr::null_mut(); + } + + let user = base.cast::().add(offset); + let header = user + .sub(size_of::()) + .cast::(); + retain_memory_budget_account(account); + header.write(AllocationHeader { + requested_size: layout.size(), + account, + }); + user +} + +#[cfg(feature = "allocator-memory-limits")] +unsafe fn allocation_header(ptr: *mut u8) -> *mut AllocationHeader { + ptr.sub(size_of::()) + .cast::() +} + +#[cfg(feature = "allocator-memory-limits")] +unsafe fn allocation_base(ptr: *mut u8, align: usize) -> *mut c_void { + allocation_header(ptr) + .cast::() + .sub(allocation_header_padding(align)) + .cast::() +} + pub struct Mimalloc; unsafe impl GlobalAlloc for Mimalloc { #[inline] unsafe fn alloc(&self, layout: Layout) -> *mut u8 { debug_assert!(layout.align() < MI_ALIGNMENT_MAX); - let ptr = mi_malloc_aligned(layout.size(), layout.align()).cast::(); - #[cfg(feature = "allocator-memory-limits")] - if !ptr.is_null() { - record_alloc(layout.size()); - } + let ptr = { + let account = current_memory_budget_account(); + let ptr = allocate_with_account(layout, account, false); + if !ptr.is_null() { + record_memory_budget_alloc(account, layout.size()); + record_alloc(layout.size()); + } + ptr + }; + + #[cfg(not(feature = "allocator-memory-limits"))] + let ptr = mi_malloc_aligned(layout.size(), layout.align()).cast::(); ptr } #[inline] unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) { #[cfg(feature = "allocator-memory-limits")] - record_free(_layout.size()); - mi_free(ptr.cast::()); + { + let header = allocation_header(ptr).read(); + let base = allocation_base(ptr, _layout.align()); + record_memory_budget_free(header.account, header.requested_size); + release_memory_budget_account(header.account); + record_free(header.requested_size); + mi_free(base); + } + + #[cfg(not(feature = "allocator-memory-limits"))] + { + mi_free(ptr.cast::()); + } } #[inline] unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { debug_assert!(layout.align() < MI_ALIGNMENT_MAX); - let ptr = mi_zalloc_aligned(layout.size(), layout.align()).cast::(); - #[cfg(feature = "allocator-memory-limits")] - if !ptr.is_null() { - record_alloc(layout.size()); - } + let ptr = { + let account = current_memory_budget_account(); + let ptr = allocate_with_account(layout, account, true); + if !ptr.is_null() { + record_memory_budget_alloc(account, layout.size()); + record_alloc(layout.size()); + } + ptr + }; + + #[cfg(not(feature = "allocator-memory-limits"))] + let ptr = mi_zalloc_aligned(layout.size(), layout.align()).cast::(); ptr } #[inline] unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { debug_assert!(layout.align() < MI_ALIGNMENT_MAX); - let result = - mi_realloc_aligned(ptr.cast::(), new_size, layout.align()).cast::(); - #[cfg(feature = "allocator-memory-limits")] - if !result.is_null() { - record_free(layout.size()); + let result = { + let Ok(new_layout) = Layout::from_size_align(new_size, layout.align()) else { + return ptr::null_mut(); + }; + let old_header = allocation_header(ptr).read(); + let old_base = allocation_base(ptr, layout.align()); + let Some((offset, raw_align, raw_size)) = allocation_header_layout(new_layout) else { + return ptr::null_mut(); + }; + let account = current_memory_budget_account(); + retain_memory_budget_account(account); + #[cfg(test)] + let base = if tests::take_next_native_realloc_failure_for_test() { + ptr::null_mut() + } else { + mi_realloc_aligned(old_base, raw_size, raw_align) + }; + #[cfg(not(test))] + let base = mi_realloc_aligned(old_base, raw_size, raw_align); + if base.is_null() { + release_memory_budget_account(account); + return ptr::null_mut(); + } + + let result = base.cast::().add(offset); + allocation_header(result).write(AllocationHeader { + requested_size: new_size, + account, + }); + record_memory_budget_free(old_header.account, old_header.requested_size); + record_memory_budget_alloc(account, new_size); + record_free(old_header.requested_size); record_alloc(new_size); - } + release_memory_budget_account(old_header.account); + result + }; + + #[cfg(not(feature = "allocator-memory-limits"))] + let result = + mi_realloc_aligned(ptr.cast::(), new_size, layout.align()).cast::(); result } } #[cfg(test)] mod tests { + #[cfg(feature = "allocator-memory-limits")] + use super::MemoryBudgetAccount; use super::*; use std::error::Error; + #[cfg(feature = "allocator-memory-limits")] + std::thread_local! { + static FAIL_NEXT_NATIVE_REALLOC_FOR_TEST: core::cell::Cell = + const { core::cell::Cell::new(false) }; + } + + #[cfg(feature = "allocator-memory-limits")] + struct FailNextNativeReallocGuard { + previous: bool, + } + + #[cfg(feature = "allocator-memory-limits")] + impl FailNextNativeReallocGuard { + fn arm() -> Self { + let previous = + FAIL_NEXT_NATIVE_REALLOC_FOR_TEST.with(|should_fail| should_fail.replace(true)); + Self { previous } + } + } + + #[cfg(feature = "allocator-memory-limits")] + impl Drop for FailNextNativeReallocGuard { + fn drop(&mut self) { + FAIL_NEXT_NATIVE_REALLOC_FOR_TEST.with(|should_fail| should_fail.set(self.previous)); + } + } + + #[cfg(feature = "allocator-memory-limits")] + pub(super) fn take_next_native_realloc_failure_for_test() -> bool { + FAIL_NEXT_NATIVE_REALLOC_FOR_TEST.with(|should_fail| should_fail.replace(false)) + } + #[test] fn memory_can_be_allocated_and_freed() -> Result<(), Box> { let layout = Layout::from_size_align(8, 8)?; @@ -93,6 +264,69 @@ mod tests { Ok(()) } + #[cfg(all(feature = "allocator-memory-limits", target_pointer_width = "64"))] + #[test] + fn two_word_header_uses_expected_alignment_eight_geometry() -> Result<(), Box> { + let layout = Layout::from_size_align(1, 8)?; + let native_request = allocation_header_layout(layout); + let allocator = Mimalloc; + let ptr = unsafe { allocator.alloc(layout) }; + assert!(!ptr.is_null()); + let user_to_header = ptr as usize - unsafe { allocation_header(ptr) } as usize; + unsafe { allocator.dealloc(ptr, layout) }; + + println!( + "geometry discriminator observed native_request={native_request:?}, user_to_header={user_to_header}; expected native_request=Some((16, 8, 17)), user_to_header=16" + ); + assert_eq!( + user_to_header, 16, + "alignment-8 user pointer displacement should be the two-word header size" + ); + assert_eq!( + native_request, + Some((16, 8, 17)), + "alignment-8 raw native allocation should request 16-byte prefix plus 1 user byte" + ); + Ok(()) + } + + #[cfg(all(feature = "allocator-memory-limits", target_pointer_width = "64"))] + #[test] + fn header_geometry_preserves_zeroed_aligned_reallocations() -> Result<(), Box> { + let allocator = Mimalloc; + let cases = [ + (1, (16, 8, 33)), + (8, (16, 8, 33)), + (16, (16, 16, 33)), + (64, (64, 64, 81)), + (256, (256, 256, 273)), + ]; + + for (align, expected_native_request) in cases { + let old_layout = Layout::from_size_align(17, align)?; + assert_eq!( + allocation_header_layout(old_layout), + Some(expected_native_request), + "unexpected native geometry for alignment {align}" + ); + + let ptr = unsafe { allocator.alloc_zeroed(old_layout) }; + assert!(!ptr.is_null()); + assert_eq!((ptr as usize) % align, 0); + assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0 })); + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + + let new_size = 33; + let ptr = unsafe { allocator.realloc(ptr, old_layout, new_size) }; + assert!(!ptr.is_null()); + assert_eq!((ptr as usize) % align, 0); + assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0x5a })); + + unsafe { allocator.dealloc(ptr, Layout::from_size_align(new_size, align)?) }; + } + Ok(()) + } + #[test] fn large_chunks_of_memory_can_be_allocated_and_freed() -> Result<(), Box> { let layout = Layout::from_size_align(2 * 1024 * 1024 * 1024, 8)?; @@ -130,4 +364,451 @@ mod tests { assert!(after_global >= before_global); Ok(()) } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn execution_budget_tracks_requested_bytes_and_origin_free_across_threads( + ) -> Result<(), Box> { + let allocator = Mimalloc; + let outer = MemoryBudgetAccount::new(); + let inner = MemoryBudgetAccount::new(); + let unowned_layout = Layout::from_size_align(11, 8)?; + let outer_layout = Layout::from_size_align(13, 8)?; + let inner_layout = Layout::from_size_align(19, 64)?; + let after_inner_layout = Layout::from_size_align(7, 8)?; + + let unowned = unsafe { allocator.alloc(unowned_layout) }; + assert!(!unowned.is_null()); + let (outer_ptr, after_inner_ptr, inner_ptr) = outer.with_scope(|| { + let outer_ptr = unsafe { allocator.alloc(outer_layout) }; + assert!(!outer_ptr.is_null()); + assert_eq!(outer.live_bytes(), 13); + assert_eq!(inner.live_bytes(), 0); + + let inner_ptr = inner.with_scope(|| unsafe { allocator.alloc_zeroed(inner_layout) }); + assert!(!inner_ptr.is_null()); + + assert_eq!(outer.live_bytes(), 13); + assert_eq!(inner.live_bytes(), 19); + assert_eq!((inner_ptr as usize) % 64, 0); + assert!((0..19).all(|offset| unsafe { *inner_ptr.add(offset) == 0 })); + + let after_inner_ptr = unsafe { allocator.alloc(after_inner_layout) }; + assert!(!after_inner_ptr.is_null()); + assert_eq!(outer.live_bytes(), 20); + unsafe { allocator.dealloc(unowned, unowned_layout) }; + assert_eq!(outer.live_bytes(), 20); + (outer_ptr, after_inner_ptr, inner_ptr) + }); + + let inner_address = inner_ptr as usize; + std::thread::spawn(move || unsafe { + Mimalloc.dealloc(inner_address as *mut u8, inner_layout); + }) + .join() + .expect("cross-thread free"); + assert_eq!(inner.live_bytes(), 0); + + unsafe { + allocator.dealloc(outer_ptr, outer_layout); + allocator.dealloc(after_inner_ptr, after_inner_layout); + } + assert_eq!(outer.live_bytes(), 0); + Ok(()) + } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn execution_budget_account_tracks_exact_boundary_and_alias_free() -> Result<(), Box> + { + let allocator = Mimalloc; + let account = MemoryBudgetAccount::new(); + let limit = 32_u64; + let layout = Layout::from_size_align(limit as usize, 8)?; + let extra_layout = Layout::from_size_align(1, 8)?; + + let ptr = account.with_scope(|| { + let ptr = unsafe { allocator.alloc(layout) }; + assert!(!ptr.is_null()); + assert_eq!(account.live_bytes(), limit); + assert!(account.live_bytes() <= limit); + + let extra = unsafe { allocator.alloc(extra_layout) }; + assert!(!extra.is_null()); + assert_eq!(account.live_bytes(), limit + 1); + assert!(account.live_bytes() > limit); + unsafe { allocator.dealloc(extra, extra_layout) }; + assert_eq!(account.live_bytes(), limit); + ptr + }); + + let alias = ptr; + unsafe { allocator.dealloc(alias, layout) }; + assert_eq!(account.live_bytes(), 0); + Ok(()) + } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn successful_realloc_transfers_origin_and_failed_realloc_preserves_it( + ) -> Result<(), Box> { + let allocator = Mimalloc; + let first = MemoryBudgetAccount::new(); + let second = MemoryBudgetAccount::new(); + let old_layout = Layout::from_size_align(16, 16)?; + let new_layout = Layout::from_size_align(64, 16)?; + + let mut ptr = first.with_scope(|| { + let ptr = unsafe { allocator.alloc(old_layout) }; + assert!(!ptr.is_null()); + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + ptr + }); + assert_eq!(first.live_bytes(), 16); + + let global_before_failure = global_allocation_stats_snapshot().allocated; + let failed = + second.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, usize::MAX) }); + assert!(failed.is_null()); + assert_eq!(first.live_bytes(), 16); + assert_eq!(second.live_bytes(), 0); + assert_eq!( + global_allocation_stats_snapshot().allocated, + global_before_failure + ); + assert!((0..16).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + + ptr = + second.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, new_layout.size()) }); + assert!(!ptr.is_null()); + assert_eq!(first.live_bytes(), 0); + assert_eq!(second.live_bytes(), 64); + assert_eq!((ptr as usize) % new_layout.align(), 0); + assert!((0..16).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + + unsafe { allocator.dealloc(ptr, new_layout) }; + assert_eq!(second.live_bytes(), 0); + Ok(()) + } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn injected_native_realloc_failure_preserves_origin_and_is_one_shot( + ) -> Result<(), Box> { + let allocator = Mimalloc; + let first = MemoryBudgetAccount::new(); + let current = MemoryBudgetAccount::new(); + let old_layout = Layout::from_size_align(16, 16)?; + let new_layout = Layout::from_size_align(64, 16)?; + let global_before = global_allocation_stats_snapshot().allocated; + + let ptr = first.with_scope(|| { + let ptr = unsafe { allocator.alloc(old_layout) }; + assert!(!ptr.is_null()); + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + ptr + }); + assert_eq!(first.live_bytes(), old_layout.size() as u64); + assert_eq!(current.live_bytes(), 0); + assert_eq!(first.reference_count_for_test(), 2); + assert_eq!(current.reference_count_for_test(), 1); + + let failure = FailNextNativeReallocGuard::arm(); + let result = + current.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, new_layout.size()) }); + drop(failure); + if !result.is_null() { + unsafe { allocator.dealloc(result, new_layout) }; + } + assert!( + result.is_null(), + "the fail-next test hook should force the native null-result path" + ); + assert_eq!(first.live_bytes(), old_layout.size() as u64); + assert_eq!(current.live_bytes(), 0); + assert_eq!(first.reference_count_for_test(), 2); + assert_eq!(current.reference_count_for_test(), 1); + assert_eq!( + global_allocation_stats_snapshot().allocated, + global_before + old_layout.size() as u64 + ); + let header = unsafe { allocation_header(ptr).read() }; + assert_eq!(header.requested_size, old_layout.size()); + assert!(!header.account.is_null()); + assert!((0..old_layout.size()).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + + drop(first); + unsafe { allocator.dealloc(ptr, old_layout) }; + + let next = current.with_scope(|| { + let ptr = unsafe { allocator.alloc(old_layout) }; + assert!(!ptr.is_null()); + unsafe { ptr.write_bytes(0x6b, old_layout.size()) }; + ptr + }); + let next = current + .with_scope(|| unsafe { allocator.realloc(next, old_layout, new_layout.size()) }); + assert!(!next.is_null()); + assert_eq!(current.live_bytes(), new_layout.size() as u64); + assert!((0..old_layout.size()).all(|offset| unsafe { *next.add(offset) == 0x6b })); + unsafe { allocator.dealloc(next, new_layout) }; + assert_eq!(current.live_bytes(), 0); + assert_eq!(global_allocation_stats_snapshot().allocated, global_before); + Ok(()) + } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn realloc_updates_requested_bytes_for_each_origin_transition() -> Result<(), Box> { + let allocator = Mimalloc; + let first = MemoryBudgetAccount::new(); + let second = MemoryBudgetAccount::new(); + let third = MemoryBudgetAccount::new(); + let mut layout = Layout::from_size_align(23, 64)?; + let mut ptr = first.with_scope(|| { + let ptr = unsafe { allocator.alloc(layout) }; + assert!(!ptr.is_null()); + unsafe { ptr.write_bytes(0x5a, layout.size()) }; + ptr + }); + + let transitions: [(Option<&MemoryBudgetAccount>, usize); 6] = [ + (Some(&first), 47), + (Some(&second), 31), + (None, 63), + (None, 55), + (Some(&third), 77), + (Some(&third), 51), + ]; + for (next_owner, new_size) in transitions { + let global_before = global_allocation_stats_snapshot().allocated; + ptr = if let Some(next_owner) = next_owner { + next_owner.with_scope(|| unsafe { allocator.realloc(ptr, layout, new_size) }) + } else { + unsafe { allocator.realloc(ptr, layout, new_size) } + }; + assert!(!ptr.is_null()); + assert_eq!((ptr as usize) % layout.align(), 0); + assert!((0..23).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + assert_eq!( + global_allocation_stats_snapshot().allocated, + global_before - layout.size() as u64 + new_size as u64 + ); + + layout = Layout::from_size_align(new_size, layout.align())?; + assert_eq!( + first.live_bytes(), + if next_owner.map_or(false, |owner| core::ptr::eq(owner, &first)) { + new_size as u64 + } else { + 0 + } + ); + assert_eq!( + second.live_bytes(), + if next_owner.map_or(false, |owner| core::ptr::eq(owner, &second)) { + new_size as u64 + } else { + 0 + } + ); + assert_eq!( + third.live_bytes(), + if next_owner.map_or(false, |owner| core::ptr::eq(owner, &third)) { + new_size as u64 + } else { + 0 + } + ); + } + + unsafe { allocator.dealloc(ptr, layout) }; + assert_eq!(third.live_bytes(), 0); + Ok(()) + } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn realloc_transfers_origin_across_threads_and_free_debits_it() -> Result<(), Box> { + let first = MemoryBudgetAccount::new(); + let second = MemoryBudgetAccount::new(); + let old_layout = Layout::from_size_align(17, 64)?; + let new_layout = Layout::from_size_align(257, 64)?; + let ptr = first.with_scope(|| { + let ptr = unsafe { Mimalloc.alloc(old_layout) }; + assert!(!ptr.is_null()); + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + ptr + }); + + let address = ptr as usize; + let worker_account = second.clone(); + let (sender, receiver) = std::sync::mpsc::channel(); + let worker = std::thread::spawn(move || { + let ptr = address as *mut u8; + let ptr = worker_account + .with_scope(|| unsafe { Mimalloc.realloc(ptr, old_layout, new_layout.size()) }); + sender.send(ptr as usize).expect("send reallocated pointer"); + }); + worker.join().expect("cross-thread realloc"); + let ptr = receiver.recv().expect("receive reallocated pointer") as *mut u8; + + assert!(!ptr.is_null()); + assert_eq!(first.live_bytes(), 0); + assert_eq!(second.live_bytes(), new_layout.size() as u64); + assert_eq!((ptr as usize) % new_layout.align(), 0); + assert!((0..old_layout.size()).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + + let address = ptr as usize; + let (sender, receiver) = std::sync::mpsc::channel(); + let freer = std::thread::spawn(move || { + unsafe { Mimalloc.dealloc(address as *mut u8, new_layout) }; + sender.send(()).expect("send free completion"); + }); + receiver.recv().expect("receive free completion"); + freer.join().expect("cross-thread free"); + assert_eq!(second.live_bytes(), 0); + Ok(()) + } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn native_aligned_realloc_reuses_a_same_bin_shrink() -> Result<(), Box> { + let allocator = Mimalloc; + let account = MemoryBudgetAccount::new(); + let old_layout = Layout::from_size_align(32, 8)?; + let ptr = account.with_scope(|| { + let ptr = unsafe { allocator.alloc(old_layout) }; + assert!(!ptr.is_null()); + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + ptr + }); + + // This characterizes vendored mimalloc's same-bin shrink, not GlobalAlloc pointer stability. + let original = ptr; + let ptr = account.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, 31) }); + assert!(!ptr.is_null()); + assert_eq!(ptr, original); + assert_eq!(account.live_bytes(), 31); + assert!((0..31).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + + unsafe { allocator.dealloc(ptr, Layout::from_size_align(31, 8)?) }; + assert_eq!(account.live_bytes(), 0); + Ok(()) + } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn native_aligned_realloc_moves_for_small_to_large_growth() -> Result<(), Box> { + let allocator = Mimalloc; + let account = MemoryBudgetAccount::new(); + let old_layout = Layout::from_size_align(17, 16)?; + let new_layout = Layout::from_size_align(1024, 16)?; + let ptr = account.with_scope(|| { + let ptr = unsafe { allocator.alloc(old_layout) }; + assert!(!ptr.is_null()); + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + ptr + }); + let original = ptr; + + let ptr = + account.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, new_layout.size()) }); + assert!(!ptr.is_null()); + assert_ne!( + ptr, original, + "vendored aligned realloc must move when this growth needs a larger native block" + ); + assert_eq!(account.live_bytes(), new_layout.size() as u64); + assert_eq!((ptr as usize) % new_layout.align(), 0); + assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0x5a })); + + unsafe { allocator.dealloc(ptr, new_layout) }; + assert_eq!(account.live_bytes(), 0); + Ok(()) + } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn successful_realloc_to_unowned_debits_the_origin_account() -> Result<(), Box> { + let allocator = Mimalloc; + let account = MemoryBudgetAccount::new(); + let old_layout = Layout::from_size_align(16, 16)?; + let new_layout = Layout::from_size_align(64, 16)?; + + let ptr = account.with_scope(|| { + let ptr = unsafe { allocator.alloc(old_layout) }; + assert!(!ptr.is_null()); + ptr + }); + assert_eq!(account.live_bytes(), 16); + + let ptr = unsafe { allocator.realloc(ptr, old_layout, new_layout.size()) }; + assert!(!ptr.is_null()); + assert_eq!(account.live_bytes(), 0); + assert_eq!((ptr as usize) % new_layout.align(), 0); + unsafe { allocator.dealloc(ptr, new_layout) }; + Ok(()) + } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn execution_budget_account_outlives_its_vm_handle_while_blocks_are_live( + ) -> Result<(), Box> { + let layout = Layout::from_size_align(32, 8)?; + let ptr = { + let account = MemoryBudgetAccount::new(); + account.with_scope(|| { + let ptr = unsafe { Mimalloc.alloc(layout) }; + assert!(!ptr.is_null()); + ptr + }) + }; + + unsafe { Mimalloc.dealloc(ptr, layout) }; + Ok(()) + } + + #[cfg(feature = "allocator-memory-limits")] + #[test] + fn execution_budget_scope_restores_nested_accounts_after_unwind() -> Result<(), Box> + { + let allocator = Mimalloc; + let outer = MemoryBudgetAccount::new(); + let inner = MemoryBudgetAccount::new(); + let layout = Layout::from_size_align(8, 8)?; + + let outer_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + outer.with_scope(|| { + let outer_ptr = unsafe { allocator.alloc(layout) }; + assert!(!outer_ptr.is_null()); + assert_eq!(outer.live_bytes(), 8); + + let inner_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + inner.with_scope(|| { + let inner_ptr = unsafe { allocator.alloc(layout) }; + assert!(!inner_ptr.is_null()); + assert_eq!(inner.live_bytes(), 8); + unsafe { allocator.dealloc(inner_ptr, layout) }; + panic!("inner scope unwind"); + }); + })); + assert!(inner_result.is_err()); + assert_eq!(inner.live_bytes(), 0); + assert_eq!(outer.live_bytes(), 8); + unsafe { allocator.dealloc(outer_ptr, layout) }; + panic!("outer scope unwind"); + }); + })); + + assert!(outer_result.is_err()); + assert_eq!(outer.live_bytes(), 0); + assert_eq!(inner.live_bytes(), 0); + let unowned = unsafe { allocator.alloc(layout) }; + assert!(!unowned.is_null()); + assert_eq!(outer.live_bytes(), 0); + assert_eq!(inner.live_bytes(), 0); + unsafe { allocator.dealloc(unowned, layout) }; + Ok(()) + } } diff --git a/src/rvm/vm/execution.rs b/src/rvm/vm/execution.rs index 463f17d8d..740ec236f 100644 --- a/src/rvm/vm/execution.rs +++ b/src/rvm/vm/execution.rs @@ -16,127 +16,124 @@ 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(), - } + crate::utils::limits::without_memory_budget_scope(|| { + self.ensure_memory_budget_execution_mode()?; + self.execution_mode = self.next_execution_mode; + match self.execution_mode { + ExecutionMode::RunToCompletion => self.execute_run_to_completion(), + ExecutionMode::Suspendable => self.execute_suspendable(), + } + }) } 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 { - index, - max_index: self.program.entry_points.len().saturating_sub(1), - pc: self.pc, - }, - )?; - (name.clone(), pc) - }; + crate::utils::limits::without_memory_budget_scope(|| { + 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 { + index, + max_index: self.program.entry_points.len().saturating_sub(1), + pc: self.pc, + }, + )?; + (name.clone(), pc) + }; - if entry_point_pc >= self.program.instructions.len() { - return Err(VmError::EntryPointPcOutOfBounds { - pc: entry_point_pc, - instruction_count: self.program.instructions.len(), - entry_point: entry_point_name, - }); - } + if entry_point_pc >= self.program.instructions.len() { + return Err(VmError::EntryPointPcOutOfBounds { + pc: entry_point_pc, + instruction_count: self.program.instructions.len(), + entry_point: entry_point_name, + }); + } - match self.execution_mode { - ExecutionMode::RunToCompletion => { - 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.execution_mode = self.next_execution_mode; + match self.execution_mode { + ExecutionMode::RunToCompletion => { + 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(), + } + })?; + + let result = self + .jump_to(entry_point_pc_u32) + .map_err(|err| self.apply_memory_budget_precedence(err))?; + self.check_memory_budget()?; + Ok(result) + })(); + match result { + Ok(value) => { + self.finish_implicit_memory_budget_execution(); + Ok(value) } - })?; - - let result = self - .jump_to(entry_point_pc_u32) - .map_err(|err| self.apply_memory_budget_precedence(err))?; - self.check_memory_budget()?; - Ok(result) - })(); - match result { - Ok(value) => { - self.finish_implicit_memory_budget_execution(); - Ok(value) + Err(err) => Err(self.fail_run_to_completion(err)), } - Err(err) => Err(self.fail_run_to_completion(err)), } + ExecutionMode::Suspendable => self.execute_suspendable_entry(entry_point_pc), } - ExecutionMode::Suspendable => { - self.reset_execution_state(); - self.reset_execution_timer_state(); - - self.validate_vm_state()?; - self.execute_suspendable_entry(entry_point_pc) - } - } + }) } 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) - .ok_or_else(|| VmError::EntryPointNotFound { - name: String::from(name), - available: self.program.entry_points.keys().cloned().collect(), - pc: self.pc, - })?; + crate::utils::limits::without_memory_budget_scope(|| { + self.ensure_memory_budget_execution_mode()?; + let entry_point_pc = + self.program + .get_entry_point(name) + .ok_or_else(|| VmError::EntryPointNotFound { + name: String::from(name), + available: self.program.entry_points.keys().cloned().collect(), + pc: self.pc, + })?; - if entry_point_pc >= self.program.instructions.len() { - return Err(VmError::EntryPointPcOutOfBounds { - pc: entry_point_pc, - instruction_count: self.program.instructions.len(), - entry_point: String::from(name), - }); - } + if entry_point_pc >= self.program.instructions.len() { + return Err(VmError::EntryPointPcOutOfBounds { + pc: entry_point_pc, + instruction_count: self.program.instructions.len(), + entry_point: String::from(name), + }); + } - match self.execution_mode { - ExecutionMode::RunToCompletion => { - 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.execution_mode = self.next_execution_mode; + match self.execution_mode { + ExecutionMode::RunToCompletion => { + 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), + } + })?; + + let result = self + .jump_to(entry_point_pc_u32) + .map_err(|err| self.apply_memory_budget_precedence(err))?; + self.check_memory_budget()?; + Ok(result) + })(); + match result { + Ok(value) => { + self.finish_implicit_memory_budget_execution(); + Ok(value) } - })?; - - let result = self - .jump_to(entry_point_pc_u32) - .map_err(|err| self.apply_memory_budget_precedence(err))?; - self.check_memory_budget()?; - Ok(result) - })(); - match result { - Ok(value) => { - self.finish_implicit_memory_budget_execution(); - Ok(value) + Err(err) => Err(self.fail_run_to_completion(err)), } - Err(err) => Err(self.fail_run_to_completion(err)), } + ExecutionMode::Suspendable => self.execute_suspendable_entry(entry_point_pc), } - ExecutionMode::Suspendable => { - self.reset_execution_state(); - self.reset_execution_timer_state(); - - self.validate_vm_state()?; - self.execute_suspendable_entry(entry_point_pc) - } - } + }) } pub(super) fn jump_to(&mut self, target: u32) -> Result { @@ -214,37 +211,119 @@ impl RegoVM { } fn execute_suspendable(&mut self) -> Result { - self.reset_execution_state(); - self.reset_execution_timer_state(); - self.execution_state = ExecutionState::Running; - match self.run_stackless_from(0) { - Ok(result) => Ok(result), - Err(err) => { - self.execution_state = ExecutionState::Error { error: err.clone() }; - Err(err) - } - } + self.execute_new_suspendable_execution(0, false) } fn execute_suspendable_entry(&mut self, entry_point_pc: usize) -> Result { - // Precondition: callers (execute_entry_point_by_{index,name}) reset the - // VM before invoking this method, so the VM must be in a clean state. - #[cfg(debug_assertions)] - self.debug_assert_state_is_clean(); - self.execution_state = ExecutionState::Running; - self.reset_execution_timer_state(); - match self.run_stackless_from(entry_point_pc) { - Ok(result) => Ok(result), - Err(err) => { - self.execution_state = ExecutionState::Error { error: err.clone() }; - Err(err) + self.execute_new_suspendable_execution(entry_point_pc, true) + } + + fn execute_new_suspendable_execution( + &mut self, + start_pc: usize, + validate_vm_state: bool, + ) -> Result { + self.release_previous_execution_state(); + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + self.begin_suspendable_memory_budget_execution(); + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + let budget_account = self.suspendable_memory_budget_account.clone(); + let mut run = || { + self.initialize_execution_state(); + #[cfg(debug_assertions)] + self.debug_assert_state_is_clean(); + self.execution_state = ExecutionState::Running; + self.reset_execution_timer_state(); + + self.pc = start_pc; + let result = (|| { + if validate_vm_state { + self.validate_vm_state()?; + } + self.check_memory_budget() + })() + .and_then(|()| self.run_stackless_from(start_pc)) + .and_then(|value| { + self.check_memory_budget()?; + Ok(value) + }); + match result { + Ok(value) => { + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + if !self.ffi_memory_budget_scope_active + && !matches!(self.execution_state, ExecutionState::Suspended { .. }) + { + self.finish_active_memory_budget_execution(); + } + Ok(value) + } + Err(err) => { + let err = self.apply_memory_budget_precedence(err); + Err(self.fail_initial_suspendable_execution(err)) + } } + }; + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + if let Some(account) = budget_account.as_ref() { + account.with_scope(run) + } else { + run() } + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + run() } pub fn resume(&mut self, resume_value: Option) -> Result { self.ensure_memory_budget_resume_supported()?; + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + let budget_account = self.suspendable_memory_budget_account.clone(); + + let resume = || { + let result = self + .resume_suspendable_segment(resume_value) + .and_then(|value| { + self.check_memory_budget()?; + Ok(value) + }); + match result { + Ok(value) => { + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + if !self.ffi_memory_budget_scope_active + && !matches!(self.execution_state, ExecutionState::Suspended { .. }) + { + self.finish_active_memory_budget_execution(); + } + Ok(value) + } + Err(err) + if matches!( + err, + VmError::MissingResumeValue { .. } + | VmError::UnexpectedResumeValue { .. } + | VmError::InvalidResumeState { .. } + ) => + { + Err(err) + } + Err(err) => { + let err = self.apply_memory_budget_precedence(err); + Err(self.fail_suspendable_execution(err)) + } + } + }; + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + if let Some(account) = budget_account.as_ref() { + account.with_scope(resume) + } else { + crate::utils::limits::without_memory_budget_scope(resume) + } + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + crate::utils::limits::without_memory_budget_scope(resume) + } + fn resume_suspendable_segment(&mut self, resume_value: Option) -> Result { // 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 @@ -286,6 +365,7 @@ impl RegoVM { if let Some(slot) = self.registers.get_mut(dest_index) { *slot = value; } + self.check_memory_budget()?; } other_reason => { if resume_value.is_some() { diff --git a/src/rvm/vm/machine.rs b/src/rvm/vm/machine.rs index a887d8364..7cac372ea 100644 --- a/src/rvm/vm/machine.rs +++ b/src/rvm/vm/machine.rs @@ -5,11 +5,15 @@ 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::MemoryBudgetAccount; +#[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, }; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +use crate::value::ResumeJsonError; use crate::value::Value; use crate::CompiledPolicy; use alloc::collections::{btree_map::Entry, BTreeMap, VecDeque}; @@ -34,6 +38,7 @@ pub(super) enum MemoryBudgetLifecycle { Inactive, ImplicitExecution, FfiResultSerialization, + SuspendableExecution, } /// The Rego Virtual Machine @@ -95,6 +100,12 @@ pub struct RegoVM { #[cfg(test)] pub(super) memory_check_count: usize, + #[cfg(all(test, feature = "allocator-memory-limits", not(miri)))] + last_memory_budget_usage_for_test: Option, + + #[cfg(all(test, feature = "allocator-memory-limits", not(miri)))] + ffi_output_start_usage_for_test: Option, + /// Cache for evaluated paths in virtual data document lookup /// Structure: evaluated[path_component1][path_component2]...[Undefined] = result_value pub(super) evaluated: Value, @@ -120,6 +131,9 @@ pub struct RegoVM { /// Current execution mode (run-to-completion vs suspendable) pub(super) execution_mode: ExecutionMode, + /// Execution mode selected for the next initial execution. + pub(super) next_execution_mode: ExecutionMode, + /// Tracks whether the current top-of-stack frame PC was explicitly set by an instruction pub(super) frame_pc_overridden: bool, @@ -145,10 +159,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 + /// Optional budget selected for the next initial execution. #[cfg(all(feature = "allocator-memory-limits", not(miri)))] pub(super) memory_budget_config: Option, + /// Budget snapshotted for the currently active execution. + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) active_memory_budget_config: Option, + + /// Requested-byte account retained across suspendable execution segments. + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) suspendable_memory_budget_account: Option, + /// Current-thread live-byte baseline captured at execution start #[cfg(all(feature = "allocator-memory-limits", not(miri)))] pub(super) memory_budget_baseline: i64, @@ -157,6 +179,10 @@ pub struct RegoVM { #[cfg(all(feature = "allocator-memory-limits", not(miri)))] pub(super) memory_budget_lifecycle: MemoryBudgetLifecycle, + /// Keeps an FFI segment active through immediate native result serialization. + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) ffi_memory_budget_scope_active: bool, + /// Cached dummy span for builtin calls (avoids Source::from_contents per call) pub(super) dummy_span: Option, @@ -230,6 +256,10 @@ impl RegoVM { executed_instructions: 0, #[cfg(test)] memory_check_count: 0, + #[cfg(all(test, feature = "allocator-memory-limits", not(miri)))] + last_memory_budget_usage_for_test: None, + #[cfg(all(test, feature = "allocator-memory-limits", not(miri)))] + ffi_output_start_usage_for_test: None, evaluated: Value::new_object(), // Initialize evaluation cache cache_hits: 0, // Initialize cache hit counter execution_stack: ExecutionStack::new(), @@ -238,6 +268,7 @@ impl RegoVM { step_mode: false, host_await_responses: BTreeMap::new(), execution_mode: ExecutionMode::RunToCompletion, + next_execution_mode: ExecutionMode::RunToCompletion, frame_pc_overridden: false, strict_builtin_errors: false, builtins_cache: BTreeMap::new(), @@ -247,9 +278,15 @@ impl RegoVM { #[cfg(all(feature = "allocator-memory-limits", not(miri)))] memory_budget_config: None, #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + active_memory_budget_config: None, + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + suspendable_memory_budget_account: 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, + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + ffi_memory_budget_scope_active: false, dummy_span: None, dummy_exprs: Vec::new(), cached_builtin_args: Vec::new(), @@ -267,6 +304,15 @@ impl RegoVM { /// Load a complete program for execution pub fn load_program(&mut self, program: Arc) { + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::SuspendableExecution + ) { + self.release_previous_execution_state(); + self.finish_active_memory_budget_execution(); + } + self.program = program.clone(); // Use the dispatch window size from the program for initial register allocation @@ -368,7 +414,7 @@ impl RegoVM { /// Set the execution mode for the VM pub const fn set_execution_mode(&mut self, mode: ExecutionMode) { - self.execution_mode = mode; + self.next_execution_mode = mode; } /// Configure whether builtin operations should raise errors strictly @@ -438,9 +484,9 @@ impl RegoVM { Ok(response) } - /// Get the current execution mode + /// Get the execution mode selected for the next initial execution. pub const fn get_execution_mode(&self) -> ExecutionMode { - self.execution_mode + self.next_execution_mode } /// Configure the execution timer to use the supplied configuration, or fall back to the global @@ -455,13 +501,25 @@ impl RegoVM { self.execution_timer_config } - /// Configure a fresh memory budget for every run-to-completion execution. + /// Configure a fresh memory budget for the next initial 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; + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) fn begin_suspendable_memory_budget_execution(&mut self) { + self.active_memory_budget_config = self.memory_budget_config; self.memory_budget_baseline = 0; - self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; + self.suspendable_memory_budget_account = self + .active_memory_budget_config + .map(|_| MemoryBudgetAccount::new()); + self.memory_budget_lifecycle = if self.suspendable_memory_budget_account.is_some() { + MemoryBudgetLifecycle::SuspendableExecution + } else { + MemoryBudgetLifecycle::Inactive + }; } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] @@ -473,12 +531,16 @@ impl RegoVM { return; } - self.memory_budget_baseline = if self.memory_budget_config.is_some() { + self.suspendable_memory_budget_account = None; + self.active_memory_budget_config = self.memory_budget_config; + self.memory_budget_baseline = if self.active_memory_budget_config.is_some() { limits::current_thread_live_bytes() } else { 0 }; - self.memory_budget_lifecycle = if self.memory_budget_config.is_some() { + self.memory_budget_lifecycle = if self.ffi_memory_budget_scope_active { + MemoryBudgetLifecycle::FfiResultSerialization + } else if self.active_memory_budget_config.is_some() { MemoryBudgetLifecycle::ImplicitExecution } else { MemoryBudgetLifecycle::Inactive @@ -492,6 +554,7 @@ impl RegoVM { MemoryBudgetLifecycle::ImplicitExecution ) { self.memory_budget_baseline = 0; + self.active_memory_budget_config = None; self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; } } @@ -502,7 +565,18 @@ impl RegoVM { #[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() { + self.ffi_memory_budget_scope_active = true; + if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::SuspendableExecution + ) || matches!(self.next_execution_mode, ExecutionMode::Suspendable) + { + return; + } + + self.active_memory_budget_config = self.memory_budget_config; + self.suspendable_memory_budget_account = None; + self.memory_budget_baseline = if self.active_memory_budget_config.is_some() { limits::current_thread_live_bytes() } else { 0 @@ -511,19 +585,27 @@ impl RegoVM { } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - const fn finish_ffi_result_serialization_memory_budget(&mut self) { + fn finish_ffi_result_serialization_memory_budget(&mut self) { + self.ffi_memory_budget_scope_active = false; if matches!( self.memory_budget_lifecycle, MemoryBudgetLifecycle::FfiResultSerialization ) { - self.memory_budget_baseline = 0; - self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; + self.finish_active_memory_budget_execution(); + } else if matches!( + self.memory_budget_lifecycle, + MemoryBudgetLifecycle::SuspendableExecution + ) && !matches!(self.execution_state, ExecutionState::Suspended { .. }) + { + self.finish_active_memory_budget_execution(); } } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - pub(super) const fn finish_active_memory_budget_execution(&mut self) { + pub(super) fn finish_active_memory_budget_execution(&mut self) { self.memory_budget_baseline = 0; + self.active_memory_budget_config = None; + self.suspendable_memory_budget_account = None; self.memory_budget_lifecycle = MemoryBudgetLifecycle::Inactive; } @@ -555,6 +637,112 @@ impl RegoVM { self.execute_to_c_string_for_ffi_with(|vm| vm.execute_entry_point_by_index(index)) } + /// Parse a native resume value, resume the VM, and serialize the result under one budget. + #[doc(hidden)] + pub fn resume_to_c_string_for_ffi(&mut self, parse_resume_json: F) -> Result + where + F: FnOnce() -> Result>, + { + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + { + self.begin_ffi_result_serialization_memory_budget(); + let mut budget = FfiResultSerializationBudget { vm: self }; + let account = budget.vm().suspendable_memory_budget_account.clone(); + + let process = || { + let resume_json = match parse_resume_json() { + Ok(value) => value, + Err(error) => { + if let Err(budget_error) = budget.vm().check_memory_budget() { + return Err(budget.vm().fail_suspendable_execution(budget_error)); + } + let error = budget.vm().apply_memory_budget_precedence(error); + if matches!( + error, + VmError::MemoryBudgetExceeded { .. } + | VmError::MemoryLimitExceeded { .. } + ) { + return Err(budget.vm().fail_suspendable_execution(error)); + } + return Err(error); + } + }; + let resume_value = match resume_json { + Some(json) => match Value::from_json_str_for_resume(&json) { + Ok(value) => Some(value), + Err(ResumeJsonError::Malformed(error)) => { + return Err(VmError::from(error)); + } + Err(ResumeJsonError::Other(error)) => { + let error = VmError::from(error); + if let Err(budget_error) = budget.vm().check_memory_budget() { + return Err(budget.vm().fail_suspendable_execution(budget_error)); + } + let error = budget.vm().apply_memory_budget_precedence(error); + if matches!( + error, + VmError::MemoryBudgetExceeded { .. } + | VmError::MemoryLimitExceeded { .. } + ) { + return Err(budget.vm().fail_suspendable_execution(error)); + } + return Err(error); + } + }, + None => None, + }; + + let value = budget.vm().resume(resume_value)?; + let output = (|| { + #[cfg(test)] + { + budget.vm().ffi_output_start_usage_for_test = budget + .vm() + .suspendable_memory_budget_account + .as_ref() + .map(|account| account.live_bytes()); + } + 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) => { + let error = budget.vm().apply_memory_budget_precedence(error); + Err(budget.vm().fail_suspendable_execution(error)) + } + } + }; + + if let Some(account) = account.as_ref() { + account.with_scope(process) + } else { + crate::utils::limits::without_memory_budget_scope(process) + } + } + + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + { + crate::utils::limits::without_memory_budget_scope(|| { + let resume_value = parse_resume_json()? + .map(|json| Value::from_json_str(&json).map_err(VmError::from)) + .transpose()?; + let value = self.resume(resume_value)?; + 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, + }) + }) + } + } + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] fn execute_to_c_string_for_ffi_with(&mut self, execute: F) -> Result where @@ -564,13 +752,26 @@ impl RegoVM { 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) + let account = budget.vm().suspendable_memory_budget_account.clone(); + let serialize = || { + #[cfg(test)] + { + budget.vm().ffi_output_start_usage_for_test = + account.as_ref().map(|account| account.live_bytes()); + } + 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) + }; + if let Some(account) = account.as_ref() { + account.with_scope(serialize) + } else { + crate::utils::limits::without_memory_budget_scope(serialize) + } })(); match output { @@ -601,21 +802,11 @@ impl RegoVM { #[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(()) } @@ -634,20 +825,33 @@ impl RegoVM { /// 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 { + let Some(config) = self.active_memory_budget_config 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 usage = match self.memory_budget_lifecycle { + MemoryBudgetLifecycle::Inactive => return Ok(()), + MemoryBudgetLifecycle::SuspendableExecution => self + .suspendable_memory_budget_account + .as_ref() + .ok_or_else(|| VmError::Internal { + message: String::from("active suspendable memory budget has no account"), + pc: self.pc, + })? + .live_bytes(), + MemoryBudgetLifecycle::ImplicitExecution + | MemoryBudgetLifecycle::FfiResultSerialization => { + let current = limits::current_thread_live_bytes(); + self.memory_budget_baseline = self.memory_budget_baseline.min(current); + current + .saturating_sub(self.memory_budget_baseline) + .unsigned_abs() + } + }; + #[cfg(test)] + { + self.last_memory_budget_usage_for_test = Some(usage); + } let budget = config.limit.get(); if usage > budget { return Err(VmError::MemoryBudgetExceeded { @@ -919,10 +1123,269 @@ impl RegoVM { mod memory_budget_tests { use super::RegoVM; use super::VmError; + use crate::languages::rego::compiler::Compiler; + use crate::rvm::vm::ExecutionMode; use crate::MemoryBudgetConfig; + use crate::{Engine, Rc, Value}; + use alloc::ffi::CString; + use alloc::string::String; + use alloc::sync::Arc; use alloc::vec; use core::num::NonZeroU64; + const NATIVE_RESUME_OUTPUT_BYTES: usize = 128 * 1024; + const NATIVE_OUTPUT_TEST_CAP_BYTES: u64 = 1024 * 1024; + const NATIVE_RESUME_OUTPUT_POLICY: &str = r#" +package limit +import rego.v1 + +result := sprintf("%s%s", [ + __builtin_host_await(input.value, "first"), + data.limit.large_text +]) +"#; + + fn native_resume_output_vm( + program: &Arc, + data: &Value, + budget: u64, + ) -> RegoVM { + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(program.clone()); + vm.set_data(data.clone()).expect("set fixture data"); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(budget).unwrap_or(NonZeroU64::MIN), + })); + vm + } + + fn compile_native_resume_output_fixture() -> (Arc, Value) { + let data_json = alloc::format!( + r#"{{"limit":{{"large_text":"{}"}}}}"#, + "x".repeat(NATIVE_RESUME_OUTPUT_BYTES) + ); + let data = Value::from_json_str(&data_json).expect("valid fixture data"); + + let mut engine = Engine::new(); + engine + .add_policy( + "memory_budget.rego".into(), + NATIVE_RESUME_OUTPUT_POLICY.into(), + ) + .expect("add fixture policy"); + let entry_point = Rc::from("data.limit.result"); + let compiled = engine + .compile_with_entrypoint(&entry_point) + .expect("compile fixture policy"); + let program = Compiler::compile_from_policy(&compiled, &[entry_point.as_ref()]) + .expect("compile fixture VM program"); + + (program, data) + } + + fn resume_to_large_native_output( + mut vm: RegoVM, + ) -> (RegoVM, core::result::Result) { + let initial = vm + .execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result") + .expect("initial HostAwait output"); + assert_eq!(initial.as_bytes_with_nul(), b"\"\"\0"); + assert!(matches!( + vm.execution_state, + super::super::execution_model::ExecutionState::Suspended { .. } + )); + + let response = String::from("\"ok\""); + let output = vm.resume_to_c_string_for_ffi(move || Ok(Some(response))); + (vm, output) + } + + #[allow(clippy::expect_used)] + #[test] + fn native_resume_output_allows_exact_budget_and_rejects_one_byte_over() { + let (program, data) = compile_native_resume_output_fixture(); + let vm = native_resume_output_vm(&program, &data, NATIVE_OUTPUT_TEST_CAP_BYTES); + let (vm, output) = resume_to_large_native_output(vm); + let output = output.expect("budgeted native resume output succeeds"); + assert_eq!( + output.as_bytes().len(), + NATIVE_RESUME_OUTPUT_BYTES.saturating_add(4) + ); + let expected_output = output.as_bytes().to_vec(); + let sampled_usage = vm + .last_memory_budget_usage_for_test + .expect("final native output check should record its live usage"); + assert!(sampled_usage >= NATIVE_RESUME_OUTPUT_BYTES as u64); + assert!(sampled_usage <= NATIVE_OUTPUT_TEST_CAP_BYTES); + std::println!( + "native_resume_final_output_check_usage={sampled_usage}; output_json_bytes={}", + output.as_bytes().len() + ); + + let exact_vm = native_resume_output_vm(&program, &data, sampled_usage); + let (exact_vm, exact_output) = resume_to_large_native_output(exact_vm); + let exact_output = exact_output.expect("usage equal to budget succeeds"); + assert_eq!(exact_output.as_bytes(), expected_output); + assert_eq!( + exact_vm.last_memory_budget_usage_for_test, + Some(sampled_usage) + ); + + let one_byte_over_budget = sampled_usage + .checked_sub(1) + .expect("fixture usage must exceed zero"); + assert_eq!(sampled_usage, one_byte_over_budget.saturating_add(1)); + let over_budget_vm = native_resume_output_vm(&program, &data, one_byte_over_budget); + let (mut over_budget_vm, over_budget_output) = + resume_to_large_native_output(over_budget_vm); + assert!(matches!( + over_budget_output, + Err(VmError::MemoryBudgetExceeded { usage, budget, .. }) + if usage == sampled_usage && budget == one_byte_over_budget + )); + assert!(matches!( + &over_budget_vm.execution_state, + super::super::execution_model::ExecutionState::Error { + error: VmError::MemoryBudgetExceeded { usage, budget, .. } + } if *usage == sampled_usage && *budget == one_byte_over_budget + )); + assert!(over_budget_vm.execution_stack.is_empty()); + assert!(over_budget_vm.host_await_responses.is_empty()); + assert!(over_budget_vm.suspendable_memory_budget_account.is_none()); + assert!(matches!( + over_budget_vm.memory_budget_lifecycle, + super::MemoryBudgetLifecycle::Inactive + )); + assert!(!matches!( + &over_budget_vm.execution_state, + super::super::execution_model::ExecutionState::Completed { .. } + )); + + let further_resume = + over_budget_vm.resume_to_c_string_for_ffi(|| Ok(Some(String::from("\"retry\"")))); + assert!(further_resume.is_err()); + assert!(matches!( + &over_budget_vm.execution_state, + super::super::execution_model::ExecutionState::Error { .. } + )); + + let plus_one_vm = native_resume_output_vm(&program, &data, sampled_usage.saturating_add(1)); + let (_, plus_one_output) = resume_to_large_native_output(plus_one_vm); + assert_eq!( + plus_one_output + .expect("budget one byte above usage succeeds") + .as_bytes(), + expected_output + ); + } + + #[allow(clippy::expect_used)] + #[test] + fn large_native_resume_output_fails_at_the_final_check_and_allows_reuse() { + let (program, data) = compile_native_resume_output_fixture(); + let calibration_vm = native_resume_output_vm(&program, &data, NATIVE_OUTPUT_TEST_CAP_BYTES); + let (calibration_vm, calibration_output) = resume_to_large_native_output(calibration_vm); + calibration_output.expect("calibration output succeeds under large cap"); + let output_start_usage = calibration_vm + .ffi_output_start_usage_for_test + .expect("resume records live usage before native output construction"); + let final_usage = calibration_vm + .last_memory_budget_usage_for_test + .expect("calibration records final native output check usage"); + assert!(output_start_usage < final_usage); + + let mut vm = native_resume_output_vm(&program, &data, output_start_usage); + let initial = vm + .execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result") + .expect("initial HostAwait output"); + assert_eq!(initial.as_bytes_with_nul(), b"\"\"\0"); + assert!(matches!( + vm.execution_state, + super::super::execution_model::ExecutionState::Suspended { .. } + )); + let before_resume_usage = vm + .suspendable_memory_budget_account + .as_ref() + .expect("suspended execution has an account") + .live_bytes(); + assert!(before_resume_usage < output_start_usage); + + let response = String::from("\"ok\""); + let error = vm + .resume_to_c_string_for_ffi(move || Ok(Some(response))) + .expect_err("native JSON and C-string production exceeds the cap"); + let observed_output_start = vm + .ffi_output_start_usage_for_test + .expect("native output stage was reached"); + let sampled_usage = vm + .last_memory_budget_usage_for_test + .expect("final output check should record live usage"); + assert_eq!(observed_output_start, output_start_usage); + assert!(sampled_usage > output_start_usage); + assert!(matches!( + error, + VmError::MemoryBudgetExceeded { usage, budget, .. } + if usage == sampled_usage && budget == output_start_usage + )); + assert!(matches!( + &vm.execution_state, + super::super::execution_model::ExecutionState::Error { + error: VmError::MemoryBudgetExceeded { usage, budget, .. } + } if *usage == sampled_usage && *budget == output_start_usage + )); + assert!(vm.execution_stack.is_empty()); + assert!(vm.host_await_responses.is_empty()); + assert!(vm.suspendable_memory_budget_account.is_none()); + assert!(matches!( + vm.memory_budget_lifecycle, + super::MemoryBudgetLifecycle::Inactive + )); + let further_resume = vm.resume_to_c_string_for_ffi(|| Ok(Some(String::from("\"retry\"")))); + assert!(further_resume.is_err()); + assert!(matches!( + &vm.execution_state, + super::super::execution_model::ExecutionState::Error { .. } + )); + + vm.set_memory_budget_config(Some(MemoryBudgetConfig { + limit: NonZeroU64::new(NATIVE_OUTPUT_TEST_CAP_BYTES).unwrap_or(NonZeroU64::MIN), + })); + let (_, reused_output) = resume_to_large_native_output(vm); + assert_eq!( + reused_output + .expect("independent execution completes under its own budget") + .as_bytes() + .len(), + NATIVE_RESUME_OUTPUT_BYTES.saturating_add(4) + ); + } + + #[allow(clippy::expect_used)] + #[test] + fn one_byte_native_resume_budget_fails_and_terminalizes() { + let (program, data) = compile_native_resume_output_fixture(); + let mut vm = native_resume_output_vm(&program, &data, 1); + let error = vm + .execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result") + .expect_err("one-byte budget must not reach a valid HostAwait continuation"); + assert!(matches!(error, VmError::MemoryBudgetExceeded { .. })); + assert!(matches!( + &vm.execution_state, + super::super::execution_model::ExecutionState::Error { + error: VmError::MemoryBudgetExceeded { .. } + } + )); + assert!(vm.execution_stack.is_empty()); + assert!(vm.host_await_responses.is_empty()); + assert!(vm.suspendable_memory_budget_account.is_none()); + assert!(matches!( + vm.memory_budget_lifecycle, + super::MemoryBudgetLifecycle::Inactive + )); + } + #[test] fn foreign_free_observed_before_allocation_does_not_grant_budget_credit() -> anyhow::Result<()> { diff --git a/src/rvm/vm/state.rs b/src/rvm/vm/state.rs index afe559fda..507972ea1 100644 --- a/src/rvm/vm/state.rs +++ b/src/rvm/vm/state.rs @@ -9,12 +9,6 @@ use super::execution_model::ExecutionState; 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(); @@ -34,7 +28,7 @@ impl RegoVM { Ok(()) } - fn release_previous_execution_state(&mut self) { + pub(super) fn release_previous_execution_state(&mut self) { self.evaluated = Value::Undefined; self.execution_state = ExecutionState::Ready; self.execution_stack.clear(); @@ -59,7 +53,51 @@ impl RegoVM { error } - fn initialize_execution_state(&mut self) { + /// Publish an initial suspendable execution error without discarding legacy unbudgeted state. + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) fn fail_initial_suspendable_execution(&mut self, error: VmError) -> VmError { + if self.active_memory_budget_config.is_some() { + return self.fail_suspendable_execution(error); + } + + self.execution_state = ExecutionState::Error { + error: error.clone(), + }; + error + } + + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + pub(super) fn fail_initial_suspendable_execution(&mut self, error: VmError) -> VmError { + self.execution_state = ExecutionState::Error { + error: error.clone(), + }; + error + } + + /// Release state after a failed budgeted suspendable execution. + /// + /// Unbudgeted executions preserve their legacy error state. + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(super) fn fail_suspendable_execution(&mut self, error: VmError) -> VmError { + if self.active_memory_budget_config.is_none() { + return error; + } + + self.release_previous_execution_state(); + self.finish_active_memory_budget_execution(); + self.execution_state = ExecutionState::Error { + error: error.clone(), + }; + error + } + + #[cfg(any(miri, not(feature = "allocator-memory-limits")))] + #[allow(clippy::unused_self)] + pub(super) fn fail_suspendable_execution(&mut self, error: VmError) -> VmError { + error + } + + pub(super) fn initialize_execution_state(&mut self) { // Reset basic execution state self.executed_instructions = 0; self.pc = 0; diff --git a/src/utils/limits/memory.rs b/src/utils/limits/memory.rs index fac97ce2d..250f30b58 100644 --- a/src/utils/limits/memory.rs +++ b/src/utils/limits/memory.rs @@ -7,6 +7,8 @@ use core::num::NonZeroU64; use core::sync::atomic::{AtomicU64, Ordering}; use std::thread_local; +pub(crate) use mimalloc::limits::MemoryBudgetAccount; + static GLOBAL_MEMORY_LIMIT: AtomicU64 = AtomicU64::new(u64::MAX); // Maximum iteration count before forcing a global memory check. diff --git a/src/utils/limits/mod.rs b/src/utils/limits/mod.rs index c8f4e7e52..450aafddd 100644 --- a/src/utils/limits/mod.rs +++ b/src/utils/limits/mod.rs @@ -14,8 +14,6 @@ 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")))] @@ -24,6 +22,8 @@ pub use memory::{ global_memory_limit, set_global_memory_limit, set_thread_flush_threshold_override, thread_memory_flush_threshold, MemoryBudgetConfig, }; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +pub(crate) use memory::{current_thread_live_bytes, MemoryBudgetAccount}; #[allow(unused_imports)] pub use time::{ @@ -34,6 +34,30 @@ pub use time::{ pub use length::PolicyLengthConfig; pub(crate) use length::{DEFAULT_MAX_COL, DEFAULT_MAX_FILE_BYTES, DEFAULT_MAX_LINES}; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +#[inline] +pub(crate) fn without_memory_budget_scope(operation: impl FnOnce() -> R) -> R { + mimalloc::limits::with_unowned_memory_budget_scope(operation) +} + +#[cfg(any(miri, not(feature = "allocator-memory-limits")))] +#[inline] +pub(crate) fn without_memory_budget_scope(operation: impl FnOnce() -> R) -> R { + operation() +} + +/// Run a binding-side read-only operation without attributing its allocations +/// to an active VM execution budget. +/// +/// This doc-hidden bridge is needed by the separately compiled FFI crate, +/// which cannot call the crate-private selector helper. Process-global +/// allocation accounting remains enabled. +#[doc(hidden)] +#[inline] +pub fn with_unowned_memory_budget_scope_for_ffi(operation: impl FnOnce() -> R) -> R { + without_memory_budget_scope(operation) +} + #[cfg(test)] pub use time::acquire_limits_test_lock; diff --git a/src/value/mod.rs b/src/value/mod.rs index 3c24bbf0a..7f090edc1 100644 --- a/src/value/mod.rs +++ b/src/value/mod.rs @@ -37,7 +37,11 @@ use crate::number::Number; use alloc::collections::{BTreeMap, BTreeSet}; use alloc::vec::Vec; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +use core::cell::Cell; use core::fmt; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +use core::marker::PhantomData; use core::ops; use core::convert::AsRef; @@ -99,7 +103,80 @@ fn enforce_limit_anyhow() -> Result<()> { #[inline] fn enforce_limit_for() -> core::result::Result<(), E> { - crate::utils::limits::check_memory_limit_if_needed().map_err(|err| E::custom(err.to_string())) + crate::utils::limits::check_memory_limit_if_needed().map_err(|err| { + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + record_resume_deserialization_limit_error(err); + E::custom(err.to_string()) + }) +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +#[derive(Clone, Copy)] +struct ResumeDeserializationLimitErrorCaptureState { + active: bool, + error: Option, +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +std::thread_local! { + static RESUME_DESERIALIZATION_LIMIT_ERROR_CAPTURE: Cell = + const { + Cell::new(ResumeDeserializationLimitErrorCaptureState { + active: false, + error: None, + }) + }; +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +struct ResumeDeserializationLimitErrorCapture { + previous: ResumeDeserializationLimitErrorCaptureState, + _thread_bound: PhantomData>, +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +impl ResumeDeserializationLimitErrorCapture { + fn new() -> Self { + let previous = RESUME_DESERIALIZATION_LIMIT_ERROR_CAPTURE.with(|capture| { + let previous = capture.get(); + capture.set(ResumeDeserializationLimitErrorCaptureState { + active: true, + error: None, + }); + previous + }); + Self { + previous, + _thread_bound: PhantomData, + } + } + + fn take(&self) -> Option { + RESUME_DESERIALIZATION_LIMIT_ERROR_CAPTURE.with(|capture| { + let mut state = capture.get(); + let error = state.error.take(); + capture.set(state); + error + }) + } +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +impl Drop for ResumeDeserializationLimitErrorCapture { + fn drop(&mut self) { + RESUME_DESERIALIZATION_LIMIT_ERROR_CAPTURE.with(|capture| capture.set(self.previous)); + } +} + +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +fn record_resume_deserialization_limit_error(error: crate::utils::limits::LimitError) { + RESUME_DESERIALIZATION_LIMIT_ERROR_CAPTURE.with(|capture| { + let mut state = capture.get(); + if state.active && state.error.is_none() { + state.error = Some(error); + capture.set(state); + } + }); } #[doc(hidden)] @@ -351,6 +428,19 @@ impl Value { } } +pub(crate) enum ResumeJsonError { + Malformed(anyhow::Error), + Other(anyhow::Error), +} + +impl ResumeJsonError { + fn into_anyhow(self) -> anyhow::Error { + match self { + Self::Malformed(error) | Self::Other(error) => error, + } + } +} + impl Value { /// Deserialize a [`Value`] from JSON. /// ``` @@ -382,25 +472,59 @@ impl Value { /// # } /// ``` pub fn from_json_str(json: &str) -> Result { + Self::parse_json_str(json, false).map_err(ResumeJsonError::into_anyhow) + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + pub(crate) fn from_json_str_for_resume( + json: &str, + ) -> core::result::Result { + Self::parse_json_str(json, true) + } + + fn parse_json_str( + json: &str, + keep_malformed_retryable: bool, + ) -> core::result::Result { // Intern object keys for the duration of the parse: a homogeneous array // of N objects sharing K keys allocates K key strings instead of N * K. #[cfg(feature = "std")] let _guard = interning::InternGuard::install(); + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + let limit_error_capture = + keep_malformed_retryable.then(ResumeDeserializationLimitErrorCapture::new); match serde_json::from_str::(json) { Ok(value) => Ok(value), Err(err) => { + if keep_malformed_retryable + && matches!( + err.classify(), + serde_json::error::Category::Syntax | serde_json::error::Category::Eof + ) + { + return Err(ResumeJsonError::Malformed(anyhow!(err))); + } + + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + if let Some(limit_error) = limit_error_capture + .as_ref() + .and_then(ResumeDeserializationLimitErrorCapture::take) + { + return Err(ResumeJsonError::Other(anyhow!(limit_error))); + } + #[cfg(all(feature = "allocator-memory-limits", not(miri)))] { // Re-validate allocator limits when serde parsing fails to surface LimitError. match crate::utils::limits::check_global_memory_limit() { - Err(limit_err) => Err(anyhow!(limit_err)), - Ok(_) => Err(anyhow!(err)), + Err(limit_err) => Err(ResumeJsonError::Other(anyhow!(limit_err))), + Ok(_) => Err(ResumeJsonError::Other(anyhow!(err))), } } #[cfg(any(miri, not(feature = "allocator-memory-limits")))] { - Err(anyhow!(err)) + Err(ResumeJsonError::Other(anyhow!(err))) } } } diff --git a/src/value/tests.rs b/src/value/tests.rs index 2465951fc..cb755b597 100644 --- a/src/value/tests.rs +++ b/src/value/tests.rs @@ -21,6 +21,59 @@ use alloc::vec::Vec; use super::{Object, Set}; use crate::value::Value; +#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +#[test] +fn resume_deserialization_limit_error_capture_is_nested_thread_local_and_unwind_safe() { + use super::{ + record_resume_deserialization_limit_error, ResumeDeserializationLimitErrorCapture, + }; + use crate::utils::limits::LimitError; + + let outer_error = LimitError::MemoryLimitExceeded { + usage: 256, + limit: 128, + }; + let nested_error = LimitError::MemoryLimitExceeded { + usage: 512, + limit: 256, + }; + let thread_error = LimitError::MemoryLimitExceeded { + usage: 1024, + limit: 512, + }; + let outer_capture = ResumeDeserializationLimitErrorCapture::new(); + record_resume_deserialization_limit_error(outer_error); + + { + let nested_capture = ResumeDeserializationLimitErrorCapture::new(); + assert_eq!(nested_capture.take(), None); + record_resume_deserialization_limit_error(nested_error); + assert_eq!(nested_capture.take(), Some(nested_error)); + } + + let thread_result = std::thread::spawn(move || { + let thread_capture = ResumeDeserializationLimitErrorCapture::new(); + assert_eq!(thread_capture.take(), None); + record_resume_deserialization_limit_error(thread_error); + thread_capture.take() + }) + .join() + .expect("thread-local capture test thread should complete"); + assert_eq!(thread_result, Some(thread_error)); + + let unwind_result = std::panic::catch_unwind(|| { + let _unwinding_capture = ResumeDeserializationLimitErrorCapture::new(); + record_resume_deserialization_limit_error(nested_error); + panic!("exercise RAII restoration during unwind"); + }); + assert!(unwind_result.is_err()); + assert_eq!(outer_capture.take(), Some(outer_error)); + + drop(outer_capture); + let subsequent_capture = ResumeDeserializationLimitErrorCapture::new(); + assert_eq!(subsequent_capture.take(), None); +} + fn val(i: u64) -> Value { Value::from(i) } diff --git a/tests/memory_limits.rs b/tests/memory_limits.rs index 7c0dd758a..24382c5e6 100644 --- a/tests/memory_limits.rs +++ b/tests/memory_limits.rs @@ -86,6 +86,25 @@ package limit quoted := sprintf("%q", [input]) "#; +#[cfg(feature = "rvm")] +const SUSPENDABLE_HOST_AWAIT_MODULE: &str = r#" +package limit +import rego.v1 + +result := __builtin_host_await(input.value, "first") +"#; + +#[cfg(feature = "rvm")] +const SUSPENDABLE_RESUME_BUDGET_MODULE: &str = r#" +package limit +import rego.v1 + +result := count(__builtin_host_await([ + __builtin_host_await(input.value, "first"), + json.unmarshal(data.limit.large_json) +], "second")) +"#; + #[cfg(feature = "rvm")] const TIGHT_MEMORY_BUDGET_BYTES: u64 = 64 * 1024; @@ -143,30 +162,13 @@ fn memory_budget(limit: u64) -> MemoryBudgetConfig { } #[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) +fn compile_memory_budget_module(module: &str, entrypoint: &str) -> Arc { + let mut engine = new_engine_with_module(module); + let entrypoint = Rc::from(entrypoint); + let compiled = engine + .compile_with_entrypoint(&entrypoint) + .expect("compile policy for VM"); + Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()]).expect("compile VM program") } #[test] @@ -564,46 +566,495 @@ fn vm_memory_budget_does_not_receive_credit_from_previous_results() { #[cfg(feature = "rvm")] #[test] -fn vm_memory_budget_rejects_suspendable_execution() { +fn vm_memory_budget_accepts_compiled_suspendable_host_await() { 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))); + vm.load_program(compile_memory_budget_module( + SUSPENDABLE_HOST_AWAIT_MODULE, + "data.limit.result", + )); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(1024 * 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:?}"), + assert_eq!( + vm.execute_entry_point_by_name("data.limit.result") + .expect("budgeted suspendable execute"), + Value::Undefined + ); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + assert_eq!( + vm.resume(Some(Value::from("response"))) + .expect("budgeted suspendable resume"), + Value::from("response") + ); +} + +#[cfg(feature = "rvm")] +#[test] +fn correction_unbudgeted_vm_entries_do_not_inherit_parent_account() { + let _guard = LimitGuard::lock(); + let program = compile_memory_budget_module(SPRINTF_QUOTE_MODULE, "data.limit.quoted"); + let outer = mimalloc::limits::MemoryBudgetAccount::new(); + + for mode in [ExecutionMode::RunToCompletion, ExecutionMode::Suspendable] { + for entry_point in ["default", "name", "index"] { + let mut vm = RegoVM::new(); + vm.set_execution_mode(mode); + vm.load_program(program.clone()); + vm.set_input(Value::from("x".repeat(512))); + + let result = outer.with_scope(|| { + let result = match entry_point { + "default" => vm.execute(), + "name" => vm.execute_entry_point_by_name("data.limit.quoted"), + _ => vm.execute_entry_point_by_index(0), + } + .expect("unbudgeted child execution succeeds"); + assert!(matches!( + &result, + Value::String(value) if value.len() > 512 + )); + assert_eq!( + outer.live_bytes(), + 0, + "{mode:?} {entry_point} child allocations must remain unowned" + ); + + let parent_probe = vec![0_u8; 64]; + core::hint::black_box(&parent_probe); + assert_eq!(outer.live_bytes(), 64); + drop(parent_probe); + assert_eq!(outer.live_bytes(), 0); + result + }); + + assert!(matches!(&result, Value::String(value) if value.len() > 512)); + } } } #[cfg(feature = "rvm")] #[test] -fn vm_memory_budget_rejects_resume_after_suspension() { +fn correction_unbudgeted_native_host_await_conversion_restores_parent_account() { + let _guard = LimitGuard::lock(); + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(compile_memory_budget_module( + SUSPENDABLE_HOST_AWAIT_MODULE, + "data.limit.result", + )); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + let outer = mimalloc::limits::MemoryBudgetAccount::new(); + + let (initial, resumed) = outer.with_scope(|| { + let initial = vm + .execute_to_c_string_for_ffi() + .expect("native execution suspends"); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + assert_eq!(outer.live_bytes(), 0); + + let resumed = vm + .resume_to_c_string_for_ffi(|| Ok(Some("\"response\"".into()))) + .expect("native resume and conversion succeed"); + assert_eq!(resumed.as_bytes(), b"\"response\""); + assert_eq!(outer.live_bytes(), 0); + + let parent_probe = vec![0_u8; 64]; + core::hint::black_box(&parent_probe); + assert_eq!(outer.live_bytes(), 64); + drop(parent_probe); + assert_eq!(outer.live_bytes(), 0); + (initial, resumed) + }); + + assert!(!initial.as_bytes().is_empty()); + assert_eq!(resumed.as_bytes(), b"\"response\""); +} + +#[cfg(feature = "rvm")] +#[test] +fn correction_malformed_native_resume_preserves_over_budget_continuation() { let _guard = LimitGuard::lock(); let mut vm = RegoVM::new(); vm.set_execution_mode(ExecutionMode::Suspendable); - vm.load_program(host_await_program()); + vm.load_program(compile_memory_budget_module( + SUSPENDABLE_HOST_AWAIT_MODULE, + "data.limit.result", + )); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(1024 * 1024))); - vm.execute().expect("suspend execution"); + assert_eq!( + vm.execute_entry_point_by_name("data.limit.result") + .expect("initial execution suspends"), + Value::Undefined + ); assert!(matches!( vm.execution_state(), ExecutionState::Suspended { .. } )); - vm.set_memory_budget_config(Some(memory_budget(1024 * 1024))); - vm.set_execution_mode(ExecutionMode::RunToCompletion); + let mut retained_probe = Vec::new(); + let malformed = vm.resume_to_c_string_for_ffi(|| { + retained_probe.resize(1024 * 1024 + 1, 0); + Ok(Some("{".into())) + }); + assert!( + matches!(malformed, Err(VmError::ArithmeticError { .. })), + "unexpected malformed input error: {malformed:?}" + ); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + + drop(retained_probe); + assert_eq!( + vm.resume_to_c_string_for_ffi(|| Ok(Some("\"response\"".into()))) + .expect("valid resume after malformed input") + .as_bytes(), + b"\"response\"" + ); +} +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_limit_during_valid_native_resume_conversion_preserves_typed_error() { + let mut guard = LimitGuard::lock(); + let program = compile_memory_budget_module(SUSPENDABLE_HOST_AWAIT_MODULE, "data.limit.result"); + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(program.clone()); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); + + assert_eq!( + vm.execute_entry_point_by_name("data.limit.result") + .expect("initial execution suspends"), + Value::Undefined + ); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + + let mut json = String::with_capacity(2 + 2 * 4096); + json.push('['); + for index in 0..4096 { + if index > 0 { + json.push(','); + } + json.push('0'); + } + json.push(']'); + let retained_json = json.clone(); + + regorus::utils::limits::check_memory_limit_if_needed() + .expect("disabled global limit resets this thread's check state"); + let baseline_usage = global_allocation_stats_snapshot().allocated as u64; + let global_limit = baseline_usage.saturating_add(32 * 1024); + guard.set_absolute_limit(global_limit); + + let result = vm.resume_to_c_string_for_ffi(move || Ok(Some(json))); + let post_parse_usage = global_allocation_stats_snapshot().allocated as u64; + println!( + "resume conversion RED: baseline={baseline_usage}, limit={global_limit}, post_parse={post_parse_usage}, error={result:?}, state={:?}", + vm.execution_state() + ); + core::hint::black_box(retained_json.as_str()); + assert!( + post_parse_usage < global_limit, + "partial Value allocations should be gone while the retained JSON remains" + ); + + let error = result.expect_err("valid array parsing should hit the global limit"); + assert!( + matches!( + error, + VmError::MemoryLimitExceeded { + usage, + limit: observed_limit, + .. + } if observed_limit == global_limit && usage > global_limit + ), + "expected the original over-limit error, got {error:?}" + ); + assert!(matches!( + vm.execution_state(), + ExecutionState::Error { + error: VmError::MemoryLimitExceeded { .. } + } + )); + assert!(vm.get_registers().is_empty()); + assert!(vm.get_host_await_argument().is_none()); + assert!(vm.get_host_await_identifier().is_none()); + assert!(matches!( + vm.resume(Some(Value::from("rejected"))), + Err(VmError::InvalidResumeState { .. }) + )); + + regorus::set_global_memory_limit(None); + regorus::utils::limits::check_memory_limit_if_needed() + .expect("clearing the global limit resets this thread's check state"); + assert_eq!( + vm.execute_entry_point_by_name("data.limit.result") + .expect("same VM can start a fresh execution"), + Value::Undefined + ); + assert_eq!( + vm.resume_to_c_string_for_ffi(|| Ok(Some(String::from("\"independent\"")))) + .expect("fresh execution can resume") + .as_bytes(), + b"\"independent\"" + ); +} + +#[cfg(feature = "rvm")] +#[test] +fn correction_unbudgeted_initial_dispatch_limit_publishes_error() { + let mut guard = LimitGuard::lock(); + let mut program = Program::new(); + for _ in 0..31 { + program.add_instruction(Instruction::LoadTrue { dest: 0 }, None); + } + program.add_instruction(Instruction::Return { value: 0 }, None); + + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(Arc::new(program)); + regorus::utils::limits::check_memory_limit_if_needed() + .expect("reset disabled global check state"); + guard.set_below_current_usage(); + + assert!(matches!( + vm.execute(), + Err(VmError::MemoryLimitExceeded { limit: 1, .. }) + )); assert!(matches!( - vm.resume(Some(Value::from(42))), - Err(VmError::MemoryBudgetUnsupportedInSuspendableExecution { .. }) + vm.execution_state(), + ExecutionState::Error { + error: VmError::MemoryLimitExceeded { limit: 1, .. } + } + )); + assert_eq!(vm.get_registers().first(), Some(&Value::Bool(true))); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_is_enforced_during_compiled_host_await_resume() { + let _guard = LimitGuard::lock(); + let mut engine = new_engine_with_module(SUSPENDABLE_RESUME_BUDGET_MODULE); + engine + .add_data(large_json_data(20_000)) + .expect("add bounded resume data"); + + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(compile_memory_budget_module( + SUSPENDABLE_RESUME_BUDGET_MODULE, + "data.limit.result", + )); + vm.set_data(engine.get_data()).expect("set data"); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); + + assert_eq!( + vm.execute_entry_point_by_name("data.limit.result") + .expect("initial execution suspends"), + Value::Undefined + ); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + + assert!(matches!( + vm.resume(Some(Value::from("response"))), + Err(VmError::MemoryBudgetExceeded { .. }) + )); + assert!(matches!( + vm.execution_state(), + ExecutionState::Error { + error: VmError::MemoryBudgetExceeded { .. } + } + )); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_configuration_is_latched_until_suspendable_execution_finishes() { + let _guard = LimitGuard::lock(); + let mut engine = new_engine_with_module(SUSPENDABLE_RESUME_BUDGET_MODULE); + engine + .add_data(large_json_data(20_000)) + .expect("add bounded resume data"); + + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(compile_memory_budget_module( + SUSPENDABLE_RESUME_BUDGET_MODULE, + "data.limit.result", + )); + vm.set_data(engine.get_data()).expect("set data"); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(TIGHT_MEMORY_BUDGET_BYTES))); + + assert_eq!( + vm.execute_entry_point_by_name("data.limit.result") + .expect("initial execution suspends"), + Value::Undefined + ); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } )); vm.set_memory_budget_config(None); + assert!(matches!( + vm.resume(Some(Value::from("response"))), + Err(VmError::MemoryBudgetExceeded { .. }) + )); + assert!(matches!( + vm.execution_state(), + ExecutionState::Error { + error: VmError::MemoryBudgetExceeded { .. } + } + )); + + assert_eq!( + vm.execute_entry_point_by_name("data.limit.result") + .expect("next execution uses the cleared budget"), + Value::Undefined + ); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_memory_budget_load_program_abandons_suspended_continuation() { + let _guard = LimitGuard::lock(); + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(compile_memory_budget_module( + SUSPENDABLE_HOST_AWAIT_MODULE, + "data.limit.result", + )); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); + assert_eq!( - vm.resume(Some(Value::from(42))) - .expect("resume after clearing budget"), - Value::from(42) + vm.execute_entry_point_by_name("data.limit.result") + .expect("initial execution suspends"), + Value::Undefined + ); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + + vm.load_program(compile_memory_budget_module( + SIMPLE_MODULE, + "data.limit.allow", + )); + assert!(matches!( + vm.resume(Some(Value::from("stale continuation"))), + Err(VmError::InvalidResumeState { .. }) + )); + assert_eq!( + vm.execute_entry_point_by_name("data.limit.allow") + .expect("new program can execute after abandonment"), + Value::Bool(true) + ); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_suspendable_missing_resume_value_preserves_host_await_continuation() { + let _guard = LimitGuard::lock(); + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.load_program(compile_memory_budget_module( + SUSPENDABLE_HOST_AWAIT_MODULE, + "data.limit.result", + )); + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_memory_budget_config(Some(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); + + assert_eq!( + vm.execute_entry_point_by_name("data.limit.result") + .expect("initial execution suspends"), + Value::Undefined + ); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + + assert!(matches!( + vm.resume(None), + Err(VmError::MissingResumeValue { .. }) + )); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + assert_eq!( + vm.resume(Some(Value::from("response"))) + .expect("continuation remains resumable"), + Value::from("response") + ); +} + +#[cfg(feature = "rvm")] +#[test] +fn vm_suspendable_unexpected_resume_value_preserves_step_continuation() { + let _guard = LimitGuard::lock(); + let mut vm = RegoVM::new(); + vm.set_execution_mode(ExecutionMode::Suspendable); + vm.set_step_mode(true); + vm.load_program(compile_memory_budget_module( + SIMPLE_MODULE, + "data.limit.allow", + )); + vm.set_memory_budget_config(Some(memory_budget(RELAXED_MEMORY_BUDGET_BYTES))); + + vm.execute_entry_point_by_name("data.limit.allow") + .expect("first step"); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + + assert!(matches!( + vm.resume(Some(Value::from("unexpected"))), + Err(VmError::UnexpectedResumeValue { .. }) + )); + assert!(matches!( + vm.execution_state(), + ExecutionState::Suspended { .. } + )); + + for _ in 0..32 { + if !matches!(vm.execution_state(), ExecutionState::Suspended { .. }) { + break; + } + vm.resume(None).expect("resume after rejected input"); + } + assert_eq!( + vm.execution_state(), + &ExecutionState::Completed { + result: Value::Bool(true) + } ); } From 0e5bdb327d42d23a95ebdbd3b571b40dcc4c0e4c Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Thu, 8 Oct 2026 14:53:14 +0200 Subject: [PATCH 02/15] fix(ci): retain failed Clippy diagnostics Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .github/workflows/rust-clippy.yml | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/.github/workflows/rust-clippy.yml b/.github/workflows/rust-clippy.yml index 83b641f1c..9009ce4a7 100644 --- a/.github/workflows/rust-clippy.yml +++ b/.github/workflows/rust-clippy.yml @@ -47,9 +47,23 @@ jobs: - name: Fetch run: cargo fetch --locked + - name: Clear cached Clippy diagnostics + run: rm -f target/xtask/clippy.json + - name: Run rust-clippy + id: clippy run: cargo xtask clippy --sarif rust-clippy-results.sarif + - name: Expose raw Clippy diagnostics on failure + if: ${{ failure() && steps.clippy.outcome == 'failure' }} + shell: bash + run: | + if [ -s target/xtask/clippy.json ]; then + cat target/xtask/clippy.json + else + echo "::warning::Clippy failed before producing a nonempty target/xtask/clippy.json" + fi + - name: Upload analysis results to GitHub if: ${{ hashFiles('rust-clippy-results.sarif') != '' }} uses: github/codeql-action/upload-sarif@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v3.29.11 From 6bdca3ecc9ffac95226d07315016467aed295538 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Thu, 8 Oct 2026 17:43:28 +0200 Subject: [PATCH 03/15] fix: repair allocator lints and RVM feature gates Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .github/workflows/tests-debug.yml | 5 +++++ mimalloc/src/limits.rs | 6 ++++++ mimalloc/src/mimalloc.rs | 10 +++------- src/utils/limits/memory.rs | 1 + src/utils/limits/mod.rs | 2 +- src/value/mod.rs | 2 +- 6 files changed, 17 insertions(+), 9 deletions(-) diff --git a/.github/workflows/tests-debug.yml b/.github/workflows/tests-debug.yml index 071417bc6..6295a8c72 100644 --- a/.github/workflows/tests-debug.yml +++ b/.github/workflows/tests-debug.yml @@ -29,5 +29,10 @@ jobs: shared-key: ${{ runner.os }}-regorus - name: Fetch dependencies run: cargo fetch --locked + - name: Check allocator features and lifecycle tests + run: | + cargo check --frozen --no-default-features --features std,arc,full-opa,rvm,allocator-memory-limits + cargo test --manifest-path mimalloc/Cargo.toml --features allocator-memory-limits --frozen execution_budget_ + cargo test --manifest-path mimalloc/Cargo.toml --features allocator-memory-limits --frozen header_ - name: Run debug CI suite run: cargo xtask ci-debug --frozen diff --git a/mimalloc/src/limits.rs b/mimalloc/src/limits.rs index 9e2303088..57ec9a1a9 100644 --- a/mimalloc/src/limits.rs +++ b/mimalloc/src/limits.rs @@ -106,6 +106,12 @@ impl MemoryBudgetAccount { } } +impl Default for MemoryBudgetAccount { + fn default() -> Self { + Self::new() + } +} + /// Select no execution account for the duration of an independent synchronous operation. #[doc(hidden)] pub fn with_unowned_memory_budget_scope(operation: impl FnOnce() -> R) -> R { diff --git a/mimalloc/src/mimalloc.rs b/mimalloc/src/mimalloc.rs index fcd215ac3..27ca1bf88 100644 --- a/mimalloc/src/mimalloc.rs +++ b/mimalloc/src/mimalloc.rs @@ -42,12 +42,8 @@ fn allocation_header_padding(align: usize) -> usize { fn allocation_header_layout(layout: Layout) -> Option<(usize, usize, usize)> { let header_size = size_of::(); let align = layout.align(); - let Some(offset) = header_size.checked_add(allocation_header_padding(align)) else { - return None; - }; - let Some(raw_size) = offset.checked_add(layout.size()) else { - return None; - }; + let offset = header_size.checked_add(allocation_header_padding(align))?; + let raw_size = offset.checked_add(layout.size())?; let raw_align = align.max(align_of::()); Some((offset, raw_align, raw_size)) } @@ -371,7 +367,7 @@ mod tests { ) -> Result<(), Box> { let allocator = Mimalloc; let outer = MemoryBudgetAccount::new(); - let inner = MemoryBudgetAccount::new(); + let inner = MemoryBudgetAccount::default(); let unowned_layout = Layout::from_size_align(11, 8)?; let outer_layout = Layout::from_size_align(13, 8)?; let inner_layout = Layout::from_size_align(19, 64)?; diff --git a/src/utils/limits/memory.rs b/src/utils/limits/memory.rs index 250f30b58..3d64540ad 100644 --- a/src/utils/limits/memory.rs +++ b/src/utils/limits/memory.rs @@ -7,6 +7,7 @@ use core::num::NonZeroU64; use core::sync::atomic::{AtomicU64, Ordering}; use std::thread_local; +#[cfg(feature = "rvm")] pub(crate) use mimalloc::limits::MemoryBudgetAccount; static GLOBAL_MEMORY_LIMIT: AtomicU64 = AtomicU64::new(u64::MAX); diff --git a/src/utils/limits/mod.rs b/src/utils/limits/mod.rs index 450aafddd..be19985a4 100644 --- a/src/utils/limits/mod.rs +++ b/src/utils/limits/mod.rs @@ -22,7 +22,7 @@ pub use memory::{ global_memory_limit, set_global_memory_limit, set_thread_flush_threshold_override, thread_memory_flush_threshold, MemoryBudgetConfig, }; -#[cfg(all(feature = "allocator-memory-limits", not(miri)))] +#[cfg(all(feature = "rvm", feature = "allocator-memory-limits", not(miri)))] pub(crate) use memory::{current_thread_live_bytes, MemoryBudgetAccount}; #[allow(unused_imports)] diff --git a/src/value/mod.rs b/src/value/mod.rs index 7f090edc1..8f1a330b3 100644 --- a/src/value/mod.rs +++ b/src/value/mod.rs @@ -475,7 +475,7 @@ impl Value { Self::parse_json_str(json, false).map_err(ResumeJsonError::into_anyhow) } - #[cfg(all(feature = "allocator-memory-limits", not(miri)))] + #[cfg(all(feature = "rvm", feature = "allocator-memory-limits", not(miri)))] pub(crate) fn from_json_str_for_resume( json: &str, ) -> core::result::Result { From 21219c96f5bb1de86eb1bbd90c2d46ddaab1552b Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Thu, 8 Oct 2026 18:24:58 +0200 Subject: [PATCH 04/15] fix: resolve remaining memory budget lints Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .github/workflows/tests-debug.yml | 2 + src/rvm/vm/execution.rs | 8 +-- src/rvm/vm/machine.rs | 104 +++++++++++++++--------------- src/utils/limits/memory.rs | 2 +- src/value/mod.rs | 2 +- src/value/tests.rs | 27 ++++++-- 6 files changed, 82 insertions(+), 63 deletions(-) diff --git a/.github/workflows/tests-debug.yml b/.github/workflows/tests-debug.yml index 6295a8c72..28f58966e 100644 --- a/.github/workflows/tests-debug.yml +++ b/.github/workflows/tests-debug.yml @@ -34,5 +34,7 @@ jobs: cargo check --frozen --no-default-features --features std,arc,full-opa,rvm,allocator-memory-limits cargo test --manifest-path mimalloc/Cargo.toml --features allocator-memory-limits --frozen execution_budget_ cargo test --manifest-path mimalloc/Cargo.toml --features allocator-memory-limits --frozen header_ + cargo test --frozen --no-default-features --features std,arc,full-opa,rvm,allocator-memory-limits --lib native_resume_ + cargo test --frozen --no-default-features --features std,arc,full-opa,rvm,allocator-memory-limits --lib resume_deserialization_limit_error_capture_ - name: Run debug CI suite run: cargo xtask ci-debug --frozen diff --git a/src/rvm/vm/execution.rs b/src/rvm/vm/execution.rs index 740ec236f..86c92711c 100644 --- a/src/rvm/vm/execution.rs +++ b/src/rvm/vm/execution.rs @@ -17,7 +17,7 @@ use super::machine::RegoVM; impl RegoVM { pub fn execute(&mut self) -> Result { crate::utils::limits::without_memory_budget_scope(|| { - self.ensure_memory_budget_execution_mode()?; + Self::ensure_memory_budget_execution_mode()?; self.execution_mode = self.next_execution_mode; match self.execution_mode { ExecutionMode::RunToCompletion => self.execute_run_to_completion(), @@ -28,7 +28,7 @@ impl RegoVM { pub fn execute_entry_point_by_index(&mut self, index: usize) -> Result { crate::utils::limits::without_memory_budget_scope(|| { - self.ensure_memory_budget_execution_mode()?; + 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 { @@ -84,7 +84,7 @@ impl RegoVM { pub fn execute_entry_point_by_name(&mut self, name: &str) -> Result { crate::utils::limits::without_memory_budget_scope(|| { - self.ensure_memory_budget_execution_mode()?; + Self::ensure_memory_budget_execution_mode()?; let entry_point_pc = self.program .get_entry_point(name) @@ -275,7 +275,7 @@ impl RegoVM { } pub fn resume(&mut self, resume_value: Option) -> Result { - self.ensure_memory_budget_resume_supported()?; + Self::ensure_memory_budget_resume_supported()?; #[cfg(all(feature = "allocator-memory-limits", not(miri)))] let budget_account = self.suspendable_memory_budget_account.clone(); diff --git a/src/rvm/vm/machine.rs b/src/rvm/vm/machine.rs index 7cac372ea..771c53646 100644 --- a/src/rvm/vm/machine.rs +++ b/src/rvm/vm/machine.rs @@ -590,12 +590,10 @@ impl RegoVM { if matches!( self.memory_budget_lifecycle, MemoryBudgetLifecycle::FfiResultSerialization - ) { - self.finish_active_memory_budget_execution(); - } else if matches!( + ) || (matches!( self.memory_budget_lifecycle, MemoryBudgetLifecycle::SuspendableExecution - ) && !matches!(self.execution_state, ExecutionState::Suspended { .. }) + ) && !matches!(self.execution_state, ExecutionState::Suspended { .. })) { self.finish_active_memory_budget_execution(); } @@ -700,7 +698,7 @@ impl RegoVM { .vm() .suspendable_memory_budget_account .as_ref() - .map(|account| account.live_bytes()); + .map(|sampled_account| sampled_account.live_bytes()); } let json = value.to_json_str().map_err(VmError::from)?; let output = CString::new(json).map_err(|_| VmError::Internal { @@ -801,24 +799,22 @@ impl RegoVM { } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - pub(super) const fn ensure_memory_budget_execution_mode(&self) -> Result<()> { + pub(super) const fn ensure_memory_budget_execution_mode() -> Result<()> { Ok(()) } #[cfg(all(feature = "allocator-memory-limits", not(miri)))] - pub(super) const fn ensure_memory_budget_resume_supported(&self) -> Result<()> { + pub(super) const fn ensure_memory_budget_resume_supported() -> Result<()> { Ok(()) } #[cfg(any(miri, not(feature = "allocator-memory-limits")))] - #[allow(clippy::unused_self)] - pub(super) const fn ensure_memory_budget_execution_mode(&self) -> Result<()> { + pub(super) const fn ensure_memory_budget_execution_mode() -> 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<()> { + pub(super) const fn ensure_memory_budget_resume_supported() -> Result<()> { Ok(()) } @@ -1149,48 +1145,44 @@ result := sprintf("%s%s", [ program: &Arc, data: &Value, budget: u64, - ) -> RegoVM { + ) -> anyhow::Result { let mut vm = RegoVM::new(); vm.set_execution_mode(ExecutionMode::Suspendable); vm.load_program(program.clone()); - vm.set_data(data.clone()).expect("set fixture data"); - vm.set_input(Value::from_json_str(r#"{"value":"request"}"#).expect("valid input JSON")); + vm.set_data(data.clone())?; + vm.set_input(Value::from_json_str(r#"{"value":"request"}"#)?); vm.set_memory_budget_config(Some(MemoryBudgetConfig { limit: NonZeroU64::new(budget).unwrap_or(NonZeroU64::MIN), })); - vm + Ok(vm) } - fn compile_native_resume_output_fixture() -> (Arc, Value) { + fn compile_native_resume_output_fixture( + ) -> anyhow::Result<(Arc, Value)> { let data_json = alloc::format!( r#"{{"limit":{{"large_text":"{}"}}}}"#, "x".repeat(NATIVE_RESUME_OUTPUT_BYTES) ); - let data = Value::from_json_str(&data_json).expect("valid fixture data"); + let data = Value::from_json_str(&data_json)?; let mut engine = Engine::new(); engine .add_policy( "memory_budget.rego".into(), NATIVE_RESUME_OUTPUT_POLICY.into(), - ) - .expect("add fixture policy"); + )?; let entry_point = Rc::from("data.limit.result"); - let compiled = engine - .compile_with_entrypoint(&entry_point) - .expect("compile fixture policy"); - let program = Compiler::compile_from_policy(&compiled, &[entry_point.as_ref()]) - .expect("compile fixture VM program"); + let compiled = engine.compile_with_entrypoint(&entry_point)?; + let program = Compiler::compile_from_policy(&compiled, &[entry_point.as_ref()])?; - (program, data) + Ok((program, data)) } fn resume_to_large_native_output( mut vm: RegoVM, - ) -> (RegoVM, core::result::Result) { + ) -> anyhow::Result<(RegoVM, core::result::Result)> { let initial = vm - .execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result") - .expect("initial HostAwait output"); + .execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result")?; assert_eq!(initial.as_bytes_with_nul(), b"\"\"\0"); assert!(matches!( vm.execution_state, @@ -1199,15 +1191,15 @@ result := sprintf("%s%s", [ let response = String::from("\"ok\""); let output = vm.resume_to_c_string_for_ffi(move || Ok(Some(response))); - (vm, output) + Ok((vm, output)) } #[allow(clippy::expect_used)] #[test] - fn native_resume_output_allows_exact_budget_and_rejects_one_byte_over() { - let (program, data) = compile_native_resume_output_fixture(); - let vm = native_resume_output_vm(&program, &data, NATIVE_OUTPUT_TEST_CAP_BYTES); - let (vm, output) = resume_to_large_native_output(vm); + fn native_resume_output_allows_exact_budget_and_rejects_one_byte_over() -> anyhow::Result<()> { + let (program, data) = compile_native_resume_output_fixture()?; + let vm = native_resume_output_vm(&program, &data, NATIVE_OUTPUT_TEST_CAP_BYTES)?; + let (vm, output) = resume_to_large_native_output(vm)?; let output = output.expect("budgeted native resume output succeeds"); assert_eq!( output.as_bytes().len(), @@ -1217,15 +1209,15 @@ result := sprintf("%s%s", [ let sampled_usage = vm .last_memory_budget_usage_for_test .expect("final native output check should record its live usage"); - assert!(sampled_usage >= NATIVE_RESUME_OUTPUT_BYTES as u64); + assert!(sampled_usage >= u64::try_from(NATIVE_RESUME_OUTPUT_BYTES)?); assert!(sampled_usage <= NATIVE_OUTPUT_TEST_CAP_BYTES); std::println!( "native_resume_final_output_check_usage={sampled_usage}; output_json_bytes={}", output.as_bytes().len() ); - let exact_vm = native_resume_output_vm(&program, &data, sampled_usage); - let (exact_vm, exact_output) = resume_to_large_native_output(exact_vm); + let exact_vm = native_resume_output_vm(&program, &data, sampled_usage)?; + let (exact_vm, exact_output) = resume_to_large_native_output(exact_vm)?; let exact_output = exact_output.expect("usage equal to budget succeeds"); assert_eq!(exact_output.as_bytes(), expected_output); assert_eq!( @@ -1237,9 +1229,9 @@ result := sprintf("%s%s", [ .checked_sub(1) .expect("fixture usage must exceed zero"); assert_eq!(sampled_usage, one_byte_over_budget.saturating_add(1)); - let over_budget_vm = native_resume_output_vm(&program, &data, one_byte_over_budget); + let over_budget_vm = native_resume_output_vm(&program, &data, one_byte_over_budget)?; let (mut over_budget_vm, over_budget_output) = - resume_to_large_native_output(over_budget_vm); + resume_to_large_native_output(over_budget_vm)?; assert!(matches!( over_budget_output, Err(VmError::MemoryBudgetExceeded { usage, budget, .. }) @@ -1265,28 +1257,34 @@ result := sprintf("%s%s", [ let further_resume = over_budget_vm.resume_to_c_string_for_ffi(|| Ok(Some(String::from("\"retry\"")))); - assert!(further_resume.is_err()); + further_resume + .err() + .ok_or_else(|| anyhow::anyhow!("further resume unexpectedly succeeded"))?; assert!(matches!( &over_budget_vm.execution_state, super::super::execution_model::ExecutionState::Error { .. } )); - let plus_one_vm = native_resume_output_vm(&program, &data, sampled_usage.saturating_add(1)); - let (_, plus_one_output) = resume_to_large_native_output(plus_one_vm); + let plus_one_vm = + native_resume_output_vm(&program, &data, sampled_usage.saturating_add(1))?; + let (_, plus_one_output) = resume_to_large_native_output(plus_one_vm)?; assert_eq!( plus_one_output .expect("budget one byte above usage succeeds") .as_bytes(), expected_output ); + Ok(()) } #[allow(clippy::expect_used)] #[test] - fn large_native_resume_output_fails_at_the_final_check_and_allows_reuse() { - let (program, data) = compile_native_resume_output_fixture(); - let calibration_vm = native_resume_output_vm(&program, &data, NATIVE_OUTPUT_TEST_CAP_BYTES); - let (calibration_vm, calibration_output) = resume_to_large_native_output(calibration_vm); + fn large_native_resume_output_fails_at_the_final_check_and_allows_reuse( + ) -> anyhow::Result<()> { + let (program, data) = compile_native_resume_output_fixture()?; + let calibration_vm = + native_resume_output_vm(&program, &data, NATIVE_OUTPUT_TEST_CAP_BYTES)?; + let (calibration_vm, calibration_output) = resume_to_large_native_output(calibration_vm)?; calibration_output.expect("calibration output succeeds under large cap"); let output_start_usage = calibration_vm .ffi_output_start_usage_for_test @@ -1296,7 +1294,7 @@ result := sprintf("%s%s", [ .expect("calibration records final native output check usage"); assert!(output_start_usage < final_usage); - let mut vm = native_resume_output_vm(&program, &data, output_start_usage); + let mut vm = native_resume_output_vm(&program, &data, output_start_usage)?; let initial = vm .execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result") .expect("initial HostAwait output"); @@ -1343,7 +1341,9 @@ result := sprintf("%s%s", [ super::MemoryBudgetLifecycle::Inactive )); let further_resume = vm.resume_to_c_string_for_ffi(|| Ok(Some(String::from("\"retry\"")))); - assert!(further_resume.is_err()); + further_resume + .err() + .ok_or_else(|| anyhow::anyhow!("further resume unexpectedly succeeded"))?; assert!(matches!( &vm.execution_state, super::super::execution_model::ExecutionState::Error { .. } @@ -1352,7 +1352,7 @@ result := sprintf("%s%s", [ vm.set_memory_budget_config(Some(MemoryBudgetConfig { limit: NonZeroU64::new(NATIVE_OUTPUT_TEST_CAP_BYTES).unwrap_or(NonZeroU64::MIN), })); - let (_, reused_output) = resume_to_large_native_output(vm); + let (_, reused_output) = resume_to_large_native_output(vm)?; assert_eq!( reused_output .expect("independent execution completes under its own budget") @@ -1360,13 +1360,14 @@ result := sprintf("%s%s", [ .len(), NATIVE_RESUME_OUTPUT_BYTES.saturating_add(4) ); + Ok(()) } #[allow(clippy::expect_used)] #[test] - fn one_byte_native_resume_budget_fails_and_terminalizes() { - let (program, data) = compile_native_resume_output_fixture(); - let mut vm = native_resume_output_vm(&program, &data, 1); + fn one_byte_native_resume_budget_fails_and_terminalizes() -> anyhow::Result<()> { + let (program, data) = compile_native_resume_output_fixture()?; + let mut vm = native_resume_output_vm(&program, &data, 1)?; let error = vm .execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result") .expect_err("one-byte budget must not reach a valid HostAwait continuation"); @@ -1384,6 +1385,7 @@ result := sprintf("%s%s", [ vm.memory_budget_lifecycle, super::MemoryBudgetLifecycle::Inactive )); + Ok(()) } #[test] diff --git a/src/utils/limits/memory.rs b/src/utils/limits/memory.rs index 3d64540ad..019e21193 100644 --- a/src/utils/limits/memory.rs +++ b/src/utils/limits/memory.rs @@ -8,7 +8,7 @@ use core::sync::atomic::{AtomicU64, Ordering}; use std::thread_local; #[cfg(feature = "rvm")] -pub(crate) use mimalloc::limits::MemoryBudgetAccount; +pub use mimalloc::limits::MemoryBudgetAccount; static GLOBAL_MEMORY_LIMIT: AtomicU64 = AtomicU64::new(u64::MAX); diff --git a/src/value/mod.rs b/src/value/mod.rs index 8f1a330b3..4c1b63465 100644 --- a/src/value/mod.rs +++ b/src/value/mod.rs @@ -151,7 +151,7 @@ impl ResumeDeserializationLimitErrorCapture { } } - fn take(&self) -> Option { + fn take(_capture: &Self) -> Option { RESUME_DESERIALIZATION_LIMIT_ERROR_CAPTURE.with(|capture| { let mut state = capture.get(); let error = state.error.take(); diff --git a/src/value/tests.rs b/src/value/tests.rs index cb755b597..215638ebc 100644 --- a/src/value/tests.rs +++ b/src/value/tests.rs @@ -46,16 +46,25 @@ fn resume_deserialization_limit_error_capture_is_nested_thread_local_and_unwind_ { let nested_capture = ResumeDeserializationLimitErrorCapture::new(); - assert_eq!(nested_capture.take(), None); + assert_eq!( + ResumeDeserializationLimitErrorCapture::take(&nested_capture), + None + ); record_resume_deserialization_limit_error(nested_error); - assert_eq!(nested_capture.take(), Some(nested_error)); + assert_eq!( + ResumeDeserializationLimitErrorCapture::take(&nested_capture), + Some(nested_error) + ); } let thread_result = std::thread::spawn(move || { let thread_capture = ResumeDeserializationLimitErrorCapture::new(); - assert_eq!(thread_capture.take(), None); + assert_eq!( + ResumeDeserializationLimitErrorCapture::take(&thread_capture), + None + ); record_resume_deserialization_limit_error(thread_error); - thread_capture.take() + ResumeDeserializationLimitErrorCapture::take(&thread_capture) }) .join() .expect("thread-local capture test thread should complete"); @@ -67,11 +76,17 @@ fn resume_deserialization_limit_error_capture_is_nested_thread_local_and_unwind_ panic!("exercise RAII restoration during unwind"); }); assert!(unwind_result.is_err()); - assert_eq!(outer_capture.take(), Some(outer_error)); + assert_eq!( + ResumeDeserializationLimitErrorCapture::take(&outer_capture), + Some(outer_error) + ); drop(outer_capture); let subsequent_capture = ResumeDeserializationLimitErrorCapture::new(); - assert_eq!(subsequent_capture.take(), None); + assert_eq!( + ResumeDeserializationLimitErrorCapture::take(&subsequent_capture), + None + ); } fn val(i: u64) -> Value { From 385e5c5497279522e6e1a2f49f092d643d8b5f88 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Thu, 8 Oct 2026 21:17:26 +0200 Subject: [PATCH 05/15] fix(test): use fallible memory budget assertions Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/rvm/vm/machine.rs | 104 ++++++++++++++++++++---------------------- 1 file changed, 49 insertions(+), 55 deletions(-) diff --git a/src/rvm/vm/machine.rs b/src/rvm/vm/machine.rs index 771c53646..e75e2f883 100644 --- a/src/rvm/vm/machine.rs +++ b/src/rvm/vm/machine.rs @@ -1166,11 +1166,10 @@ result := sprintf("%s%s", [ let data = Value::from_json_str(&data_json)?; let mut engine = Engine::new(); - engine - .add_policy( - "memory_budget.rego".into(), - NATIVE_RESUME_OUTPUT_POLICY.into(), - )?; + engine.add_policy( + "memory_budget.rego".into(), + NATIVE_RESUME_OUTPUT_POLICY.into(), + )?; let entry_point = Rc::from("data.limit.result"); let compiled = engine.compile_with_entrypoint(&entry_point)?; let program = Compiler::compile_from_policy(&compiled, &[entry_point.as_ref()])?; @@ -1181,10 +1180,9 @@ result := sprintf("%s%s", [ fn resume_to_large_native_output( mut vm: RegoVM, ) -> anyhow::Result<(RegoVM, core::result::Result)> { - let initial = vm - .execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result")?; - assert_eq!(initial.as_bytes_with_nul(), b"\"\"\0"); - assert!(matches!( + let initial = vm.execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result")?; + anyhow::ensure!(initial.as_bytes_with_nul() == b"\"\"\0"); + anyhow::ensure!(matches!( vm.execution_state, super::super::execution_model::ExecutionState::Suspended { .. } )); @@ -1201,16 +1199,13 @@ result := sprintf("%s%s", [ let vm = native_resume_output_vm(&program, &data, NATIVE_OUTPUT_TEST_CAP_BYTES)?; let (vm, output) = resume_to_large_native_output(vm)?; let output = output.expect("budgeted native resume output succeeds"); - assert_eq!( - output.as_bytes().len(), - NATIVE_RESUME_OUTPUT_BYTES.saturating_add(4) - ); + anyhow::ensure!(output.as_bytes().len() == NATIVE_RESUME_OUTPUT_BYTES.saturating_add(4)); let expected_output = output.as_bytes().to_vec(); let sampled_usage = vm .last_memory_budget_usage_for_test .expect("final native output check should record its live usage"); - assert!(sampled_usage >= u64::try_from(NATIVE_RESUME_OUTPUT_BYTES)?); - assert!(sampled_usage <= NATIVE_OUTPUT_TEST_CAP_BYTES); + anyhow::ensure!(sampled_usage >= u64::try_from(NATIVE_RESUME_OUTPUT_BYTES)?); + anyhow::ensure!(sampled_usage <= NATIVE_OUTPUT_TEST_CAP_BYTES); std::println!( "native_resume_final_output_check_usage={sampled_usage}; output_json_bytes={}", output.as_bytes().len() @@ -1219,38 +1214,37 @@ result := sprintf("%s%s", [ let exact_vm = native_resume_output_vm(&program, &data, sampled_usage)?; let (exact_vm, exact_output) = resume_to_large_native_output(exact_vm)?; let exact_output = exact_output.expect("usage equal to budget succeeds"); - assert_eq!(exact_output.as_bytes(), expected_output); - assert_eq!( - exact_vm.last_memory_budget_usage_for_test, - Some(sampled_usage) + anyhow::ensure!(exact_output.as_bytes() == expected_output); + anyhow::ensure!( + exact_vm.last_memory_budget_usage_for_test == Some(sampled_usage) ); let one_byte_over_budget = sampled_usage .checked_sub(1) .expect("fixture usage must exceed zero"); - assert_eq!(sampled_usage, one_byte_over_budget.saturating_add(1)); + anyhow::ensure!(sampled_usage == one_byte_over_budget.saturating_add(1)); let over_budget_vm = native_resume_output_vm(&program, &data, one_byte_over_budget)?; let (mut over_budget_vm, over_budget_output) = resume_to_large_native_output(over_budget_vm)?; - assert!(matches!( + anyhow::ensure!(matches!( over_budget_output, Err(VmError::MemoryBudgetExceeded { usage, budget, .. }) if usage == sampled_usage && budget == one_byte_over_budget )); - assert!(matches!( + anyhow::ensure!(matches!( &over_budget_vm.execution_state, super::super::execution_model::ExecutionState::Error { error: VmError::MemoryBudgetExceeded { usage, budget, .. } } if *usage == sampled_usage && *budget == one_byte_over_budget )); - assert!(over_budget_vm.execution_stack.is_empty()); - assert!(over_budget_vm.host_await_responses.is_empty()); - assert!(over_budget_vm.suspendable_memory_budget_account.is_none()); - assert!(matches!( + anyhow::ensure!(over_budget_vm.execution_stack.is_empty()); + anyhow::ensure!(over_budget_vm.host_await_responses.is_empty()); + anyhow::ensure!(over_budget_vm.suspendable_memory_budget_account.is_none()); + anyhow::ensure!(matches!( over_budget_vm.memory_budget_lifecycle, super::MemoryBudgetLifecycle::Inactive )); - assert!(!matches!( + anyhow::ensure!(!matches!( &over_budget_vm.execution_state, super::super::execution_model::ExecutionState::Completed { .. } )); @@ -1260,7 +1254,7 @@ result := sprintf("%s%s", [ further_resume .err() .ok_or_else(|| anyhow::anyhow!("further resume unexpectedly succeeded"))?; - assert!(matches!( + anyhow::ensure!(matches!( &over_budget_vm.execution_state, super::super::execution_model::ExecutionState::Error { .. } )); @@ -1268,19 +1262,19 @@ result := sprintf("%s%s", [ let plus_one_vm = native_resume_output_vm(&program, &data, sampled_usage.saturating_add(1))?; let (_, plus_one_output) = resume_to_large_native_output(plus_one_vm)?; - assert_eq!( + anyhow::ensure!( plus_one_output .expect("budget one byte above usage succeeds") - .as_bytes(), - expected_output + .as_bytes() + == expected_output ); Ok(()) } #[allow(clippy::expect_used)] #[test] - fn large_native_resume_output_fails_at_the_final_check_and_allows_reuse( - ) -> anyhow::Result<()> { + fn large_native_resume_output_fails_at_the_final_check_and_allows_reuse() -> anyhow::Result<()> + { let (program, data) = compile_native_resume_output_fixture()?; let calibration_vm = native_resume_output_vm(&program, &data, NATIVE_OUTPUT_TEST_CAP_BYTES)?; @@ -1292,14 +1286,14 @@ result := sprintf("%s%s", [ let final_usage = calibration_vm .last_memory_budget_usage_for_test .expect("calibration records final native output check usage"); - assert!(output_start_usage < final_usage); + anyhow::ensure!(output_start_usage < final_usage); let mut vm = native_resume_output_vm(&program, &data, output_start_usage)?; let initial = vm .execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result") .expect("initial HostAwait output"); - assert_eq!(initial.as_bytes_with_nul(), b"\"\"\0"); - assert!(matches!( + anyhow::ensure!(initial.as_bytes_with_nul() == b"\"\"\0"); + anyhow::ensure!(matches!( vm.execution_state, super::super::execution_model::ExecutionState::Suspended { .. } )); @@ -1308,7 +1302,7 @@ result := sprintf("%s%s", [ .as_ref() .expect("suspended execution has an account") .live_bytes(); - assert!(before_resume_usage < output_start_usage); + anyhow::ensure!(before_resume_usage < output_start_usage); let response = String::from("\"ok\""); let error = vm @@ -1320,23 +1314,23 @@ result := sprintf("%s%s", [ let sampled_usage = vm .last_memory_budget_usage_for_test .expect("final output check should record live usage"); - assert_eq!(observed_output_start, output_start_usage); - assert!(sampled_usage > output_start_usage); - assert!(matches!( + anyhow::ensure!(observed_output_start == output_start_usage); + anyhow::ensure!(sampled_usage > output_start_usage); + anyhow::ensure!(matches!( error, VmError::MemoryBudgetExceeded { usage, budget, .. } if usage == sampled_usage && budget == output_start_usage )); - assert!(matches!( + anyhow::ensure!(matches!( &vm.execution_state, super::super::execution_model::ExecutionState::Error { error: VmError::MemoryBudgetExceeded { usage, budget, .. } } if *usage == sampled_usage && *budget == output_start_usage )); - assert!(vm.execution_stack.is_empty()); - assert!(vm.host_await_responses.is_empty()); - assert!(vm.suspendable_memory_budget_account.is_none()); - assert!(matches!( + anyhow::ensure!(vm.execution_stack.is_empty()); + anyhow::ensure!(vm.host_await_responses.is_empty()); + anyhow::ensure!(vm.suspendable_memory_budget_account.is_none()); + anyhow::ensure!(matches!( vm.memory_budget_lifecycle, super::MemoryBudgetLifecycle::Inactive )); @@ -1344,7 +1338,7 @@ result := sprintf("%s%s", [ further_resume .err() .ok_or_else(|| anyhow::anyhow!("further resume unexpectedly succeeded"))?; - assert!(matches!( + anyhow::ensure!(matches!( &vm.execution_state, super::super::execution_model::ExecutionState::Error { .. } )); @@ -1353,12 +1347,12 @@ result := sprintf("%s%s", [ limit: NonZeroU64::new(NATIVE_OUTPUT_TEST_CAP_BYTES).unwrap_or(NonZeroU64::MIN), })); let (_, reused_output) = resume_to_large_native_output(vm)?; - assert_eq!( + anyhow::ensure!( reused_output .expect("independent execution completes under its own budget") .as_bytes() - .len(), - NATIVE_RESUME_OUTPUT_BYTES.saturating_add(4) + .len() + == NATIVE_RESUME_OUTPUT_BYTES.saturating_add(4) ); Ok(()) } @@ -1371,17 +1365,17 @@ result := sprintf("%s%s", [ let error = vm .execute_entry_point_by_name_to_c_string_for_ffi("data.limit.result") .expect_err("one-byte budget must not reach a valid HostAwait continuation"); - assert!(matches!(error, VmError::MemoryBudgetExceeded { .. })); - assert!(matches!( + anyhow::ensure!(matches!(error, VmError::MemoryBudgetExceeded { .. })); + anyhow::ensure!(matches!( &vm.execution_state, super::super::execution_model::ExecutionState::Error { error: VmError::MemoryBudgetExceeded { .. } } )); - assert!(vm.execution_stack.is_empty()); - assert!(vm.host_await_responses.is_empty()); - assert!(vm.suspendable_memory_budget_account.is_none()); - assert!(matches!( + anyhow::ensure!(vm.execution_stack.is_empty()); + anyhow::ensure!(vm.host_await_responses.is_empty()); + anyhow::ensure!(vm.suspendable_memory_budget_account.is_none()); + anyhow::ensure!(matches!( vm.memory_budget_lifecycle, super::MemoryBudgetLifecycle::Inactive )); From fd6d32e5c4d438980a2a763f159cb0a74d4a7032 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Thu, 8 Oct 2026 23:06:12 +0200 Subject: [PATCH 06/15] style(test): format native resume budget assertion Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/rvm/vm/machine.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/rvm/vm/machine.rs b/src/rvm/vm/machine.rs index e75e2f883..494934a00 100644 --- a/src/rvm/vm/machine.rs +++ b/src/rvm/vm/machine.rs @@ -1215,9 +1215,7 @@ result := sprintf("%s%s", [ let (exact_vm, exact_output) = resume_to_large_native_output(exact_vm)?; let exact_output = exact_output.expect("usage equal to budget succeeds"); anyhow::ensure!(exact_output.as_bytes() == expected_output); - anyhow::ensure!( - exact_vm.last_memory_budget_usage_for_test == Some(sampled_usage) - ); + anyhow::ensure!(exact_vm.last_memory_budget_usage_for_test == Some(sampled_usage)); let one_byte_over_budget = sampled_usage .checked_sub(1) From 13a66d8be63b0540092f9a4f07eb00a9af619e68 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Fri, 9 Oct 2026 00:15:13 +0200 Subject: [PATCH 07/15] fix(test): serialize FFI test harness Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- xtask/src/tasks/bindings/ffi.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/xtask/src/tasks/bindings/ffi.rs b/xtask/src/tasks/bindings/ffi.rs index f09531df1..725863df5 100644 --- a/xtask/src/tasks/bindings/ffi.rs +++ b/xtask/src/tasks/bindings/ffi.rs @@ -161,9 +161,12 @@ impl TestFfiCommand { } test_args.push(OsString::from("--features")); test_args.push(OsString::from("contention_checks")); + // Global memory-limit tests share process-wide allocator counters. + test_args.push(OsString::from("--")); + test_args.push(OsString::from("--test-threads=1")); run_cargo_step( &ffi_dir, - "cargo test --features contention_checks (bindings/ffi)", + "cargo test --features contention_checks -- --test-threads=1 (bindings/ffi)", test_args, )?; From f018788bb873182b172bfdc8733fd4c18708e494 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Fri, 9 Oct 2026 10:42:07 +0200 Subject: [PATCH 08/15] ci: retain Rust CodeQL SARIF artifacts Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .github/workflows/codeql.yml | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 7d684703a..a6cdbf264 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -188,6 +188,16 @@ jobs: run: cargo xtask build-wasm --release - name: Perform CodeQL Analysis + id: analyze uses: github/codeql-action/analyze@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 with: - category: "/language:${{matrix.language}}" \ No newline at end of file + category: "/language:${{matrix.language}}" + + - name: Retain Rust SARIF + if: matrix.language == 'rust' + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: codeql-rust-sarif-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ steps.analyze.outputs.sarif-output }}/*.sarif + if-no-files-found: error + retention-days: 1 \ No newline at end of file From 59542065c0cdd7310f63ac85e9d899e482c86ee0 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Fri, 9 Oct 2026 16:45:16 +0200 Subject: [PATCH 09/15] ci: export fifth Rust CodeQL path Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .github/workflows/codeql.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index e00588a03..e8e483aa1 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -190,6 +190,8 @@ jobs: - name: Perform CodeQL Analysis id: analyze uses: github/codeql-action/analyze@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4.38.2 + env: + CODEQL_ACTION_EXTRA_OPTIONS: ${{ matrix.language == 'rust' && '{"database":{"interpret-results":["--max-paths=5"]}}' || '{}' }} with: category: "/language:${{matrix.language}}" From e97324677abdf18d52079d50925d93f9e4a524af Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Fri, 9 Oct 2026 19:35:23 +0200 Subject: [PATCH 10/15] ci: characterize explicit Rust null guard Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .github/workflows/codeql.yml | 6 ++++++ mimalloc/src/mimalloc.rs | 11 +++++++---- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index e8e483aa1..9be069615 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -132,6 +132,12 @@ jobs: sudo apt-get update sudo apt-get install -y build-essential cmake + - name: Check Rust guard diagnostic preservation + if: matrix.language == 'rust' + run: | + cargo test --locked -p regorus-mimalloc --features allocator-memory-limits mimalloc::tests::header_geometry_preserves_zeroed_aligned_reallocations -- --exact + cargo clippy --locked -p regorus-mimalloc --all-targets --features allocator-memory-limits -- -D warnings + - name: Install Python build dependencies if: matrix.language == 'python' working-directory: ${{ matrix.working-directory }} diff --git a/mimalloc/src/mimalloc.rs b/mimalloc/src/mimalloc.rs index 27ca1bf88..863225213 100644 --- a/mimalloc/src/mimalloc.rs +++ b/mimalloc/src/mimalloc.rs @@ -307,10 +307,13 @@ mod tests { ); let ptr = unsafe { allocator.alloc_zeroed(old_layout) }; - assert!(!ptr.is_null()); - assert_eq!((ptr as usize) % align, 0); - assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0 })); - unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + assert_eq!((ptr as usize) % align, 0); + assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0 })); + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + } let new_size = 33; let ptr = unsafe { allocator.realloc(ptr, old_layout, new_size) }; From 3a17a9fd3cb50bd4ad375ab6cb8d31e166d41936 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Sat, 10 Oct 2026 19:24:22 +0200 Subject: [PATCH 11/15] test: establish integrated allocator guard control Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mimalloc/src/mimalloc.rs | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/mimalloc/src/mimalloc.rs b/mimalloc/src/mimalloc.rs index 863225213..27ca1bf88 100644 --- a/mimalloc/src/mimalloc.rs +++ b/mimalloc/src/mimalloc.rs @@ -307,13 +307,10 @@ mod tests { ); let ptr = unsafe { allocator.alloc_zeroed(old_layout) }; - if ptr.is_null() { - assert!(!ptr.is_null()); - } else { - assert_eq!((ptr as usize) % align, 0); - assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0 })); - unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; - } + assert!(!ptr.is_null()); + assert_eq!((ptr as usize) % align, 0); + assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0 })); + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; let new_size = 33; let ptr = unsafe { allocator.realloc(ptr, old_layout, new_size) }; From 362f512377e40514c846805608fd0d38b4760493 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Sat, 10 Oct 2026 22:23:44 +0200 Subject: [PATCH 12/15] test: simplify allocator ownership assertions Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mimalloc/src/mimalloc.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mimalloc/src/mimalloc.rs b/mimalloc/src/mimalloc.rs index 27ca1bf88..1fbf76697 100644 --- a/mimalloc/src/mimalloc.rs +++ b/mimalloc/src/mimalloc.rs @@ -594,7 +594,7 @@ mod tests { layout = Layout::from_size_align(new_size, layout.align())?; assert_eq!( first.live_bytes(), - if next_owner.map_or(false, |owner| core::ptr::eq(owner, &first)) { + if next_owner.is_some_and(|owner| core::ptr::eq(owner, &first)) { new_size as u64 } else { 0 @@ -602,7 +602,7 @@ mod tests { ); assert_eq!( second.live_bytes(), - if next_owner.map_or(false, |owner| core::ptr::eq(owner, &second)) { + if next_owner.is_some_and(|owner| core::ptr::eq(owner, &second)) { new_size as u64 } else { 0 @@ -610,7 +610,7 @@ mod tests { ); assert_eq!( third.live_bytes(), - if next_owner.map_or(false, |owner| core::ptr::eq(owner, &third)) { + if next_owner.is_some_and(|owner| core::ptr::eq(owner, &third)) { new_size as u64 } else { 0 From 16afa478230c00af144b208db0b1d2928dddde38 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Sat, 10 Oct 2026 23:31:21 +0200 Subject: [PATCH 13/15] test: expose allocator fixture null guard Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mimalloc/src/mimalloc.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/mimalloc/src/mimalloc.rs b/mimalloc/src/mimalloc.rs index 1fbf76697..97b76f4c2 100644 --- a/mimalloc/src/mimalloc.rs +++ b/mimalloc/src/mimalloc.rs @@ -307,10 +307,13 @@ mod tests { ); let ptr = unsafe { allocator.alloc_zeroed(old_layout) }; - assert!(!ptr.is_null()); - assert_eq!((ptr as usize) % align, 0); - assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0 })); - unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + assert_eq!((ptr as usize) % align, 0); + assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0 })); + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + } let new_size = 33; let ptr = unsafe { allocator.realloc(ptr, old_layout, new_size) }; From 1ab0820ae73ad474c636843a5a96e1f04cd21fc8 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Mon, 12 Oct 2026 12:11:16 +0200 Subject: [PATCH 14/15] test: guard allocator fixture pointer lifetimes Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mimalloc/src/mimalloc.rs | 444 +++++++++++++++++++++++---------------- 1 file changed, 267 insertions(+), 177 deletions(-) diff --git a/mimalloc/src/mimalloc.rs b/mimalloc/src/mimalloc.rs index 97b76f4c2..d2e598c7b 100644 --- a/mimalloc/src/mimalloc.rs +++ b/mimalloc/src/mimalloc.rs @@ -241,8 +241,11 @@ mod tests { unsafe { let ptr = alloc.alloc(layout); - assert!(!ptr.cast::().is_null()); - alloc.dealloc(ptr, layout); + if ptr.cast::().is_null() { + assert!(!ptr.cast::().is_null()); + } else { + alloc.dealloc(ptr, layout); + } } Ok(()) } @@ -254,8 +257,11 @@ mod tests { unsafe { let ptr = alloc.alloc_zeroed(layout); - assert!(!ptr.cast::().is_null()); - alloc.dealloc(ptr, layout); + if ptr.cast::().is_null() { + assert!(!ptr.cast::().is_null()); + } else { + alloc.dealloc(ptr, layout); + } } Ok(()) } @@ -313,15 +319,18 @@ mod tests { assert_eq!((ptr as usize) % align, 0); assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0 })); unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; - } - let new_size = 33; - let ptr = unsafe { allocator.realloc(ptr, old_layout, new_size) }; - assert!(!ptr.is_null()); - assert_eq!((ptr as usize) % align, 0); - assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0x5a })); + let new_size = 33; + let ptr = unsafe { allocator.realloc(ptr, old_layout, new_size) }; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + assert_eq!((ptr as usize) % align, 0); + assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0x5a })); - unsafe { allocator.dealloc(ptr, Layout::from_size_align(new_size, align)?) }; + unsafe { allocator.dealloc(ptr, Layout::from_size_align(new_size, align)?) }; + } + } } Ok(()) } @@ -459,34 +468,44 @@ mod tests { let mut ptr = first.with_scope(|| { let ptr = unsafe { allocator.alloc(old_layout) }; - assert!(!ptr.is_null()); - unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + } ptr }); - assert_eq!(first.live_bytes(), 16); - - let global_before_failure = global_allocation_stats_snapshot().allocated; - let failed = - second.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, usize::MAX) }); - assert!(failed.is_null()); - assert_eq!(first.live_bytes(), 16); - assert_eq!(second.live_bytes(), 0); - assert_eq!( - global_allocation_stats_snapshot().allocated, - global_before_failure - ); - assert!((0..16).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + assert_eq!(first.live_bytes(), 16); + + let global_before_failure = global_allocation_stats_snapshot().allocated; + let failed = + second.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, usize::MAX) }); + assert!(failed.is_null()); + assert_eq!(first.live_bytes(), 16); + assert_eq!(second.live_bytes(), 0); + assert_eq!( + global_allocation_stats_snapshot().allocated, + global_before_failure + ); + assert!((0..16).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); - ptr = - second.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, new_layout.size()) }); - assert!(!ptr.is_null()); - assert_eq!(first.live_bytes(), 0); - assert_eq!(second.live_bytes(), 64); - assert_eq!((ptr as usize) % new_layout.align(), 0); - assert!((0..16).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + ptr = second + .with_scope(|| unsafe { allocator.realloc(ptr, old_layout, new_layout.size()) }); + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + assert_eq!(first.live_bytes(), 0); + assert_eq!(second.live_bytes(), 64); + assert_eq!((ptr as usize) % new_layout.align(), 0); + assert!((0..16).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); - unsafe { allocator.dealloc(ptr, new_layout) }; - assert_eq!(second.live_bytes(), 0); + unsafe { allocator.dealloc(ptr, new_layout) }; + assert_eq!(second.live_bytes(), 0); + } + } Ok(()) } @@ -503,56 +522,74 @@ mod tests { let ptr = first.with_scope(|| { let ptr = unsafe { allocator.alloc(old_layout) }; - assert!(!ptr.is_null()); - unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + } ptr }); - assert_eq!(first.live_bytes(), old_layout.size() as u64); - assert_eq!(current.live_bytes(), 0); - assert_eq!(first.reference_count_for_test(), 2); - assert_eq!(current.reference_count_for_test(), 1); - - let failure = FailNextNativeReallocGuard::arm(); - let result = - current.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, new_layout.size()) }); - drop(failure); - if !result.is_null() { - unsafe { allocator.dealloc(result, new_layout) }; - } - assert!( - result.is_null(), - "the fail-next test hook should force the native null-result path" - ); - assert_eq!(first.live_bytes(), old_layout.size() as u64); - assert_eq!(current.live_bytes(), 0); - assert_eq!(first.reference_count_for_test(), 2); - assert_eq!(current.reference_count_for_test(), 1); - assert_eq!( - global_allocation_stats_snapshot().allocated, - global_before + old_layout.size() as u64 - ); - let header = unsafe { allocation_header(ptr).read() }; - assert_eq!(header.requested_size, old_layout.size()); - assert!(!header.account.is_null()); - assert!((0..old_layout.size()).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); - - drop(first); - unsafe { allocator.dealloc(ptr, old_layout) }; - - let next = current.with_scope(|| { - let ptr = unsafe { allocator.alloc(old_layout) }; + if ptr.is_null() { assert!(!ptr.is_null()); - unsafe { ptr.write_bytes(0x6b, old_layout.size()) }; - ptr - }); - let next = current - .with_scope(|| unsafe { allocator.realloc(next, old_layout, new_layout.size()) }); - assert!(!next.is_null()); - assert_eq!(current.live_bytes(), new_layout.size() as u64); - assert!((0..old_layout.size()).all(|offset| unsafe { *next.add(offset) == 0x6b })); - unsafe { allocator.dealloc(next, new_layout) }; - assert_eq!(current.live_bytes(), 0); - assert_eq!(global_allocation_stats_snapshot().allocated, global_before); + } else { + assert_eq!(first.live_bytes(), old_layout.size() as u64); + assert_eq!(current.live_bytes(), 0); + assert_eq!(first.reference_count_for_test(), 2); + assert_eq!(current.reference_count_for_test(), 1); + + let failure = FailNextNativeReallocGuard::arm(); + let result = current + .with_scope(|| unsafe { allocator.realloc(ptr, old_layout, new_layout.size()) }); + drop(failure); + if !result.is_null() { + unsafe { allocator.dealloc(result, new_layout) }; + } + assert!( + result.is_null(), + "the fail-next test hook should force the native null-result path" + ); + assert_eq!(first.live_bytes(), old_layout.size() as u64); + assert_eq!(current.live_bytes(), 0); + assert_eq!(first.reference_count_for_test(), 2); + assert_eq!(current.reference_count_for_test(), 1); + assert_eq!( + global_allocation_stats_snapshot().allocated, + global_before + old_layout.size() as u64 + ); + let header = unsafe { allocation_header(ptr).read() }; + assert_eq!(header.requested_size, old_layout.size()); + assert!(!header.account.is_null()); + assert!((0..old_layout.size()).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + + drop(first); + unsafe { allocator.dealloc(ptr, old_layout) }; + + let next = current.with_scope(|| { + let ptr = unsafe { allocator.alloc(old_layout) }; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + unsafe { ptr.write_bytes(0x6b, old_layout.size()) }; + } + ptr + }); + if next.is_null() { + assert!(!next.is_null()); + } else { + let next = current.with_scope(|| unsafe { + allocator.realloc(next, old_layout, new_layout.size()) + }); + if next.is_null() { + assert!(!next.is_null()); + } else { + assert_eq!(current.live_bytes(), new_layout.size() as u64); + assert!((0..old_layout.size()).all(|offset| unsafe { *next.add(offset) == 0x6b })); + unsafe { allocator.dealloc(next, new_layout) }; + assert_eq!(current.live_bytes(), 0); + assert_eq!(global_allocation_stats_snapshot().allocated, global_before); + } + } + } Ok(()) } @@ -566,8 +603,11 @@ mod tests { let mut layout = Layout::from_size_align(23, 64)?; let mut ptr = first.with_scope(|| { let ptr = unsafe { allocator.alloc(layout) }; - assert!(!ptr.is_null()); - unsafe { ptr.write_bytes(0x5a, layout.size()) }; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + unsafe { ptr.write_bytes(0x5a, layout.size()) }; + } ptr }); @@ -581,48 +621,59 @@ mod tests { ]; for (next_owner, new_size) in transitions { let global_before = global_allocation_stats_snapshot().allocated; - ptr = if let Some(next_owner) = next_owner { - next_owner.with_scope(|| unsafe { allocator.realloc(ptr, layout, new_size) }) + if ptr.is_null() { + assert!(!ptr.is_null()); } else { - unsafe { allocator.realloc(ptr, layout, new_size) } - }; - assert!(!ptr.is_null()); - assert_eq!((ptr as usize) % layout.align(), 0); - assert!((0..23).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); - assert_eq!( - global_allocation_stats_snapshot().allocated, - global_before - layout.size() as u64 + new_size as u64 - ); - - layout = Layout::from_size_align(new_size, layout.align())?; - assert_eq!( - first.live_bytes(), - if next_owner.is_some_and(|owner| core::ptr::eq(owner, &first)) { - new_size as u64 + ptr = if let Some(next_owner) = next_owner { + next_owner.with_scope(|| unsafe { allocator.realloc(ptr, layout, new_size) }) } else { - 0 - } - ); - assert_eq!( - second.live_bytes(), - if next_owner.is_some_and(|owner| core::ptr::eq(owner, &second)) { - new_size as u64 + unsafe { allocator.realloc(ptr, layout, new_size) } + }; + if ptr.is_null() { + assert!(!ptr.is_null()); } else { - 0 + assert_eq!((ptr as usize) % layout.align(), 0); + assert!((0..23).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + assert_eq!( + global_allocation_stats_snapshot().allocated, + global_before - layout.size() as u64 + new_size as u64 + ); + + layout = Layout::from_size_align(new_size, layout.align())?; + assert_eq!( + first.live_bytes(), + if next_owner.is_some_and(|owner| core::ptr::eq(owner, &first)) { + new_size as u64 + } else { + 0 + } + ); + assert_eq!( + second.live_bytes(), + if next_owner.is_some_and(|owner| core::ptr::eq(owner, &second)) { + new_size as u64 + } else { + 0 + } + ); + assert_eq!( + third.live_bytes(), + if next_owner.is_some_and(|owner| core::ptr::eq(owner, &third)) { + new_size as u64 + } else { + 0 + } + ); } - ); - assert_eq!( - third.live_bytes(), - if next_owner.is_some_and(|owner| core::ptr::eq(owner, &third)) { - new_size as u64 - } else { - 0 - } - ); + } } - unsafe { allocator.dealloc(ptr, layout) }; - assert_eq!(third.live_bytes(), 0); + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + unsafe { allocator.dealloc(ptr, layout) }; + assert_eq!(third.live_bytes(), 0); + } Ok(()) } @@ -635,38 +686,57 @@ mod tests { let new_layout = Layout::from_size_align(257, 64)?; let ptr = first.with_scope(|| { let ptr = unsafe { Mimalloc.alloc(old_layout) }; - assert!(!ptr.is_null()); - unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + } ptr }); + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + let address = ptr as usize; + let worker_account = second.clone(); + let (sender, receiver) = std::sync::mpsc::channel(); + let worker = std::thread::spawn(move || { + let ptr = address as *mut u8; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + let ptr = worker_account.with_scope(|| unsafe { + Mimalloc.realloc(ptr, old_layout, new_layout.size()) + }); + sender.send(ptr as usize).expect("send reallocated pointer"); + } + }); + worker.join().expect("cross-thread realloc"); + let ptr = receiver.recv().expect("receive reallocated pointer") as *mut u8; - let address = ptr as usize; - let worker_account = second.clone(); - let (sender, receiver) = std::sync::mpsc::channel(); - let worker = std::thread::spawn(move || { - let ptr = address as *mut u8; - let ptr = worker_account - .with_scope(|| unsafe { Mimalloc.realloc(ptr, old_layout, new_layout.size()) }); - sender.send(ptr as usize).expect("send reallocated pointer"); - }); - worker.join().expect("cross-thread realloc"); - let ptr = receiver.recv().expect("receive reallocated pointer") as *mut u8; - - assert!(!ptr.is_null()); - assert_eq!(first.live_bytes(), 0); - assert_eq!(second.live_bytes(), new_layout.size() as u64); - assert_eq!((ptr as usize) % new_layout.align(), 0); - assert!((0..old_layout.size()).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); - - let address = ptr as usize; - let (sender, receiver) = std::sync::mpsc::channel(); - let freer = std::thread::spawn(move || { - unsafe { Mimalloc.dealloc(address as *mut u8, new_layout) }; - sender.send(()).expect("send free completion"); - }); - receiver.recv().expect("receive free completion"); - freer.join().expect("cross-thread free"); - assert_eq!(second.live_bytes(), 0); + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + assert_eq!(first.live_bytes(), 0); + assert_eq!(second.live_bytes(), new_layout.size() as u64); + assert_eq!((ptr as usize) % new_layout.align(), 0); + assert!((0..old_layout.size()).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + + let address = ptr as usize; + let (sender, receiver) = std::sync::mpsc::channel(); + let freer = std::thread::spawn(move || { + let ptr = address as *mut u8; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + unsafe { Mimalloc.dealloc(ptr, new_layout) }; + } + sender.send(()).expect("send free completion"); + }); + receiver.recv().expect("receive free completion"); + freer.join().expect("cross-thread free"); + assert_eq!(second.live_bytes(), 0); + } + } Ok(()) } @@ -678,21 +748,31 @@ mod tests { let old_layout = Layout::from_size_align(32, 8)?; let ptr = account.with_scope(|| { let ptr = unsafe { allocator.alloc(old_layout) }; - assert!(!ptr.is_null()); - unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + } ptr }); - // This characterizes vendored mimalloc's same-bin shrink, not GlobalAlloc pointer stability. - let original = ptr; - let ptr = account.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, 31) }); - assert!(!ptr.is_null()); - assert_eq!(ptr, original); - assert_eq!(account.live_bytes(), 31); - assert!((0..31).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + // This characterizes vendored mimalloc's same-bin shrink, not GlobalAlloc pointer stability. + let original = ptr; + let ptr = account.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, 31) }); + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + assert_eq!(ptr, original); + assert_eq!(account.live_bytes(), 31); + assert!((0..31).all(|offset| unsafe { *ptr.add(offset) == 0x5a })); - unsafe { allocator.dealloc(ptr, Layout::from_size_align(31, 8)?) }; - assert_eq!(account.live_bytes(), 0); + unsafe { allocator.dealloc(ptr, Layout::from_size_align(31, 8)?) }; + assert_eq!(account.live_bytes(), 0); + } + } Ok(()) } @@ -705,25 +785,35 @@ mod tests { let new_layout = Layout::from_size_align(1024, 16)?; let ptr = account.with_scope(|| { let ptr = unsafe { allocator.alloc(old_layout) }; - assert!(!ptr.is_null()); - unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + unsafe { ptr.write_bytes(0x5a, old_layout.size()) }; + } ptr }); - let original = ptr; - - let ptr = - account.with_scope(|| unsafe { allocator.realloc(ptr, old_layout, new_layout.size()) }); - assert!(!ptr.is_null()); - assert_ne!( - ptr, original, - "vendored aligned realloc must move when this growth needs a larger native block" - ); - assert_eq!(account.live_bytes(), new_layout.size() as u64); - assert_eq!((ptr as usize) % new_layout.align(), 0); - assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0x5a })); - - unsafe { allocator.dealloc(ptr, new_layout) }; - assert_eq!(account.live_bytes(), 0); + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + let original = ptr; + let ptr = account.with_scope(|| unsafe { + allocator.realloc(ptr, old_layout, new_layout.size()) + }); + if ptr.is_null() { + assert!(!ptr.is_null()); + } else { + assert_ne!( + ptr, original, + "vendored aligned realloc must move when this growth needs a larger native block" + ); + assert_eq!(account.live_bytes(), new_layout.size() as u64); + assert_eq!((ptr as usize) % new_layout.align(), 0); + assert!((0..old_layout.size()).all(|index| unsafe { *ptr.add(index) == 0x5a })); + + unsafe { allocator.dealloc(ptr, new_layout) }; + assert_eq!(account.live_bytes(), 0); + } + } Ok(()) } From 61fc3b39eb413e7a2bf59b4f98566d88055117b9 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Mon, 12 Oct 2026 13:04:08 +0200 Subject: [PATCH 15/15] style: format allocator fixture guards Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mimalloc/src/mimalloc.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/mimalloc/src/mimalloc.rs b/mimalloc/src/mimalloc.rs index d2e598c7b..bfcd2d235 100644 --- a/mimalloc/src/mimalloc.rs +++ b/mimalloc/src/mimalloc.rs @@ -583,7 +583,9 @@ mod tests { assert!(!next.is_null()); } else { assert_eq!(current.live_bytes(), new_layout.size() as u64); - assert!((0..old_layout.size()).all(|offset| unsafe { *next.add(offset) == 0x6b })); + assert!( + (0..old_layout.size()).all(|offset| unsafe { *next.add(offset) == 0x6b }) + ); unsafe { allocator.dealloc(next, new_layout) }; assert_eq!(current.live_bytes(), 0); assert_eq!(global_allocation_stats_snapshot().allocated, global_before); @@ -796,9 +798,8 @@ mod tests { assert!(!ptr.is_null()); } else { let original = ptr; - let ptr = account.with_scope(|| unsafe { - allocator.realloc(ptr, old_layout, new_layout.size()) - }); + let ptr = account + .with_scope(|| unsafe { allocator.realloc(ptr, old_layout, new_layout.size()) }); if ptr.is_null() { assert!(!ptr.is_null()); } else {