Add per-execution memory budgets to RVM

[Maksym (maksym-mishchenko)]

AB#3522638

A process-global allocator limit cannot isolate individual policy evaluations and may cause unrelated requests to fail. This adds an optional memory budget to each run-to-completion RVM execution.

The budget measures additional current-thread live bytes above a fresh execution baseline. Previous execution-owned values are released before that baseline is captured. Program compilation and data, input, and context loading remain outside the budget.

Budget exhaustion has additive typed Rust, FFI, and C# errors and takes precedence over the process-global limit. Budgeted suspendable execution is rejected because its thread-local accounting cannot safely span thread migration.

The opt-in controls are RegoVM::set_memory_budget_config and C# Rvm.SetMemoryBudgetConfig; existing behavior is unchanged when no budget is configured.

[Copilot]

Pull request overview

Adds opt-in, per-execution memory budgets to isolate RVM evaluations using thread-local allocator accounting.

Changes:

• Enforces fresh memory budgets for run-to-completion RVM execution.
• Adds typed Rust, FFI, and C# errors and configuration APIs.
• Adds documentation, tests, and benchmark coverage.

Reviewed changes

Copilot reviewed 24 out of 24 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/memory_limits.rs	Tests enforcement, precedence, reset, and threading.
src/utils/limits/mod.rs	Exports budget configuration and counters.
src/utils/limits/memory.rs	Defines memory budget configuration.
src/rvm/vm/state.rs	Resets state before capturing baselines.
src/rvm/vm/rules.rs	Treats budget exhaustion as fatal.
src/rvm/vm/machine.rs	Implements budget accounting and checks.
src/rvm/vm/execution.rs	Integrates budgets into execution entry points.
src/rvm/vm/errors.rs	Adds typed budget errors.
src/lib.rs	Exposes the Rust configuration API.
mimalloc/src/mimalloc.rs	Re-exports thread live-byte accounting.
mimalloc/src/limits.rs	Implements and tests live-byte sampling.
mimalloc/src/lib.rs	Exposes allocator accounting publicly.
docs/limits/memory_budget.md	Documents behavior and limitations.
bindings/ffi/src/rvm.rs	Adds FFI configuration and status mapping.
bindings/ffi/src/limits.rs	Defines FFI budget configuration.
bindings/ffi/src/common.rs	Adds the FFI exhaustion status.
bindings/csharp/Regorus/StatusExtensions.cs	Maps exhaustion to a typed exception.
bindings/csharp/Regorus/Rvm.cs	Adds budget configuration methods.
bindings/csharp/Regorus/RegorusMemoryBudgetExceededException.cs	Defines the typed exception.
bindings/csharp/Regorus/NativeMethods.cs	Adds native declarations and types.
bindings/csharp/Regorus/MemoryBudgetConfig.cs	Defines validated C# configuration.
bindings/csharp/Regorus.Tests/RvmMemoryBudgetTests.cs	Tests the C# API.
bindings/csharp/README.md	Documents C# usage.
benches/rvm_benchmark.rs	Benchmarks budget overhead.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Anand Krishnamoorthi (anakrish)]

Maksym (@maksym-mishchenko) Thanks for doing this very useful feature! Overall looks good to me. Copilot reviews found some interesting cases that are worth addressing.

[Maksym (maksym-mishchenko)]

Maksym (Maksym (@maksym-mishchenko)) Thanks for doing this very useful feature! Overall looks good to me. Copilot reviews found some interesting cases that are worth addressing.

Hi Anand Krishnamoorthi (@anakrish), thanks for the thorough review. I addressed comments, I kept the two API suggestions unchanged for the reasons explained in their threads. Could you please take another look when you have time?

[Anand Krishnamoorthi (anakrish)]

Low: These benchmark imports/calls use #[cfg(feature = "allocator-memory-limits")], but the core MemoryBudgetConfig export and RegoVM::set_memory_budget_config are gated by all(feature = "allocator-memory-limits", not(miri)). A bench build selected under Miri can therefore fail to compile.

Suggested change:

#[cfg(all(feature = "allocator-memory-limits", not(miri)))]
use std::num::NonZeroU64;

#[cfg(all(feature = "allocator-memory-limits", not(miri)))]
use regorus::MemoryBudgetConfig;

#[cfg(all(feature = "allocator-memory-limits", not(miri)))]
vm.set_memory_budget_config(config.memory_budget.then(|| MemoryBudgetConfig {
    limit: NonZeroU64::new(MEMORY_LIMIT_BYTES).expect("non-zero memory budget"),
}));

Please use the same predicate for both imports and the configuration block.

[Anand Krishnamoorthi (anakrish)]

Medium: If this post-serialization budget check fails, the FFI returns MemoryBudgetExceeded after dropping only the provisional JSON result, but the VM has already stored the value as ExecutionState::Completed { result }. A caller can then call regorus_rvm_get_execution_state() and observe/re-serialize the oversized completed result despite the reported failure.

Suggested shape:

if let Err(err) = guard.check_memory_budget() {
    regorus_result_drop(result);
    guard.mark_execution_error(err.clone()); // clear retained result/state
    return Err(err.into());
}

Alternatively, move the final serialization check into a core-owned completion helper that transitions the VM to ExecutionState::Error and releases the retained result before returning. Please add a regression asserting the state is Error after this failure.

[Anand Krishnamoorthi (anakrish)]

Medium: When this final budget check fails, the Err arm records ExecutionState::Error but leaves the failed execution allocations (registers, rule_cache, evaluated, and pooled values) resident until the next execution or VM drop. The named/indexed entry-point paths and the FFI post-marshalling failure do not clean them up at all. Reusing a VM after repeated oversized failures can therefore accumulate roughly one failed result/state per cycle.

Suggested shape:

fn fail_execution(&mut self, err: VmError) -> VmError {
    self.release_previous_execution_state();
    self.execution_state = ExecutionState::Error { error: err.clone() };
    err
}

Use this helper in the run-to-completion Err arm, the named/indexed entry-point budget-error paths, and the FFI post-marshalling failure path. Preserve the error value while releasing the result/register/cache state, and add a regression that repeats budget failures on a reused VM and verifies memory/state are reset.

[Anand Krishnamoorthi (anakrish)]

Maksym (Maksym (Maksym (@maksym-mishchenko))) Thanks for doing this very useful feature! Overall looks good to me. Copilot reviews found some interesting cases that are worth addressing.

Hi Anand Krishnamoorthi (Anand Krishnamoorthi (@anakrish)), thanks for the thorough review. I addressed comments, I kept the two API suggestions unchanged for the reasons explained in their threads. Could you please take another look when you have time?

The rationale for both sounds good. Review rerun found 3 more comments worth addressing. Then it should be good to go.

[Mark Birger (kusha)]

Maksym (@maksym-mishchenko) I see the following drawbacks in this implementation:

• The budget is armed at execute*() time, not at set_memory_budget_config() time (state.rs:19-22 → machine.rs:442-449). set_memory_budget_config sets baseline = 0; active = false; the real baseline is captured on the next execution. So it is a per-execution budget, not a budget window a host can open and close.
• RegoVM::set_data / set_input / set_context / load_program all run before the baseline and are uncharged (stated in docs/limits/memory_budget.md). Via FFI, a host can set a budget, load a 500 MB data blob successfully, and only then have evaluation-on-top bounded. If the goal is per-request isolation, the hostile-input case is the one not covered.

Why haven't you pinned thread level baseline and used every allocation (delta gated) for comparison (just like a global limit)? Why have you chosen execute() calls?

Consider whether the intended primitive is a scope (begin(limit) / end() pinning a thread baseline, consulted by every allocating entry point — data load, compile, execute, serialize) rather than a per-execution budget. set_memory_budget_config is about to become public API with per-execution semantics baked in.

Anand Krishnamoorthi (@anakrish) any thoughts on the above ^?

---

Apart from the design:

Follow-ups on the fixes from the previous review round

• Amortization was removed rather than refined (machine.rs:673, dispatch.rs:28). Comment 3816390968 argued stride-16 sampling was too coarse and asked for a byte-delta trigger; the fix removed the stride entirely. check_memory_budget() now runs on every instruction, doing a THREAD_COUNTERS.with(...) TLS access per instruction whenever a budget is configured — while the global limit it takes precedence over is still amortized (MEMORY_CHECK_STRIDE = 16 + 32 KiB delta gate). The bench harness gained a regular_memory_budget config; could the numbers be posted?
• The .min() ratchet can make the budget stricter than configured (machine.rs:498). memory_budget_baseline = memory_budget_baseline.min(current) never recovers, so any free of pre-baseline memory observed on the execution thread permanently lowers the baseline. Fail-safe direction, but it produces non-deterministic MemoryBudgetExceeded below the configured limit. release_previous_execution_state() before baseline capture mitigates the main in-engine case. The doc says only "does not grant credit" — it should also state the budget can become stricter than configured.

New

• memory_budget_active is never cleared (machine.rs:448; only reset at machine.rs:431 and on the next execution). After a successful execute() the window stays live, so a Rust host calling the public check_memory_budget() later charges unrelated work to the last execution. Given it is pub and the doc invites bindings to call it post-serialization, either consume the window on read or document it as valid only immediately after execute, on the same thread.
• Binding parity — only C# is wired up. C++ (bindings/cpp/regorus.hpp), Go, Java, Python, Ruby, and WASM get the new RegorusStatus variants but no budget API. Fine as an increment, but worth an explicit note in the PR or a tracking issue.
• Doc overstatement — docs/limits/memory_budget.md says the check covers "native string marshaling" for C#. The managed-string copy happens after the FFI call returns and is on the managed heap, not the mimalloc heap, so it is not covered.

[Anand Krishnamoorthi (anakrish)]

• RegoVM::set_data / set_input / set_context / load_program all run before the baseline and are uncharged (stated in docs/limits/memory_budget.md). Via FFI, a host can set a budget, load a 500 MB data blob successfully, and only then have evaluation-on-top bounded. If the goal is per-request isolation, the hostile-input case is the one not covered.

Mark Birger (@kusha)
I can see rationale for both sides. If the data/input are not included in the evaluation budget then, as you showed, really large objects can escape the memory check.
On the other hand, it is easier for the caller to reason about/estimate the size of input and data (say from the length of the byte streams from which they are originally deserialized) and perform checks whereas it is harder for the caller to reason about a policy evaluation's memory consumption. This PR addresses the latter.

However, I do see that it would be nice to also limit input, data memory consumption.

Some challenges around implementation:

• Setting input, data and evaluation is not an atomic operation. One thread could set input/data and another could perform the evaluation. Per thread budgets don't play nicely with this.
• If input/data exist as C# objects as opposed to a byte stream, then they must anyways be serialized on the C# size..ie the 500MB limit would be exceeded even before reaching rust.

Some ideas

• Should we have individual limits around data/object sizes?
• Or for evaluation, should the base line start at earlier or set data/input and end at the end of evaluation? Ideally in case of FFI after the result has been serialized.

[Anand Krishnamoorthi (anakrish)]

Why haven't you pinned thread level baseline and used every allocation (delta gated) for comparison (just like a global limit)? Why have you chosen execute() calls?

Maybe this would be the simplest approach. Have a thread level limit on alive memory?
But if one thread is the producer of programs and input/data and other threads are merely users, then it may not bounded correctly.

[Anand Krishnamoorthi (anakrish)]

So it is a per-execution budget, not a budget window a host can open and close.

Yes. That was also my impression from the PR description and the implementation. An execution-level budget is the clearest semantic: it bounds the additional working memory used during one RVM execution and intentionally excludes objects such as input, data, and the program that may have been created before—and may outlive—that execution.

However, I do agree that large data/input may go undetected by this budget, even if they are eventually subject to configured global memory limit.

We should meet to discuss your scenario and see if the design can be generalized and the semantics made more clear/easy to reason about.

[Maksym (maksym-mishchenko)]

Mark Birger (@kusha) Anand Krishnamoorthi (@anakrish). I made another pass based on this discussion.

The existing execute* APIs still have a fresh budget for each execution. I added an opt-in path that starts before evaluation-specific JSON data is parsed and set, then stays active through execution and native result serialization.

I avoided a begin/end API across FFI calls because the accounting is thread-local, and separate calls may run on different OS threads. Rust uses a same-thread scope; FFI and C# do everything in one call.

Program loading and compilation are still excluded. I also left input and context out for now. Is covering data enough for the fetch scenario, or do you think input should be included before we approve the API?

The cleanup and documentation issues are fixed as well. The benchmark showed no measurable difference: 859.82 ns without the budget and 847.59 ns with it, with overlapping intervals.

[Anand Krishnamoorthi (anakrish)]

Maksym (@maksym-mishchenko) data is typically the data that doesn't share that often across individual policy evaluation. input is the actual input for a policy evaluation. IMO input would also need to be guarded.