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 optional per-execution memory budgets for run-to-completion RVM evaluations.

Enforcement is cooperative and checkpoint-based. The budget observes current-thread live bytes relative to a fresh execution baseline; it is not an allocation-time peak cap, so one instruction, builtin, native serialization step, or CString allocation may temporarily overshoot. Reported usage is a diagnostic thread-level change rather than exact query-owned memory. Synchronous callbacks on the execution thread affect accounting, cross-thread frees may temporarily skew observations, and a lower live-byte sample ratchets the baseline downward without restoring lost headroom.

A reused VM captures a fresh baseline after prior execution state is released, but retained capacities and pools precede that baseline, so warm and fresh VMs may allocate differently. Program compilation and loading data, input, or context remain outside the execution budget. Native result serialization and CString construction are included; managed C# UTF-8 decoding and managed-string allocation are excluded.

Budget failures return additive typed Rust, FFI, and C# errors, invalidate completed results, release retained execution state, and take precedence over the process-global allocator limit. Budgeted suspendable execution is rejected because thread-local accounting cannot safely span host-await thread migration; host-await accounting is separate follow-up work.

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.

[Anand Krishnamoorthi (anakrish)]

I found a performance issue in the new memory-budget enforcement path: memory_check() is invoked twice per VM instruction in run-to-completion execution. src/rvm/vm/execution.rs calls self.memory_check() before dispatch, and src/rvm/vm/dispatch.rs calls it again at the start of execute_instruction(). The same duplicate exists on the suspendable dispatch path.

With allocator limits enabled, each call can sample the budget and run the throttled global-limit check, so this adds avoidable hot-path overhead and makes checkpoint cadence effectively 2x what the execution loop suggests.

Suggested fix: retain the check at one dispatch boundary (preferably the outer execution loop, unless a specific instruction-level invariant requires the inner check), then add a test-only check counter or benchmark assertion to verify one check per dispatched instruction. Please also compare the benchmark with one vs. two checks so the overhead is measurable.

[Anand Krishnamoorthi (anakrish)]

Documentation suggestion — cooperative enforcement and overshoot:

Please make the hard-cap distinction especially explicit in docs/limits/memory_budget.md and the C# README/API docs. The budget is checked at VM checkpoints and after native result serialization; a single instruction, builtin, to_json_str(), or CString allocation can temporarily exceed the configured limit before the next check. This is a cooperative observed-live-bytes budget, not an allocation-time peak-memory cap. A short example of this behavior would help callers choose an appropriate headroom.

[Anand Krishnamoorthi (anakrish)]

Documentation suggestion — VM reuse and warm pools:

Please document that budget outcomes can depend on prior executions when the same RegoVM is reused. Execution state is reset and the baseline is fresh, but retained VM pools/capacities are already live outside the new baseline, so an identical policy/input may have different allocation behavior on a warm VM than on a fresh VM. This is an important qualification to the phrase “fresh per-execution budget.”

[Anand Krishnamoorthi (anakrish)]

Documentation suggestion — baseline ratcheting and the usage diagnostic:

Please clarify that the baseline is lowered when sampled same-thread live bytes fall below the prior baseline, and that this lost headroom is not restored. As a result, the effective budget can become stricter after unrelated or legitimate same-thread frees, and VmError::MemoryBudgetExceeded { usage, .. } should not be interpreted as exact memory owned by the policy evaluation or as a net-allocation measurement.

[Anand Krishnamoorthi (anakrish)]

Documentation suggestion — host callbacks and cross-thread frees:

Please call out that accounting is based on the current thread’s live-byte counter, not allocation ownership. Allocations/frees performed by synchronous custom or host builtins on the execution thread are therefore observed as part of the budget; objects allocated on one thread and freed on another can temporarily skew per-thread/global observations until counters are published. This would help embedding applications avoid assuming strict query-owned attribution or thread-independent accounting.