feat(memory): Allocator-backed global memory limits (#544)

Policy evaluation at scale needs to be able to set memory limits
so that a bad policy does not hog memory or to ensure that
policy evaluation itself does not use too much memory which could
cause other components to suffer.

This PR introduces capability to set and enforce global memory limits.
It also lays the groundwork for enabling per evaluation limits in future.

Once a global memory limit is set, Regorus maintains per thread counters
to track memory activity (allocation, deallocation) of a thread.
These counters are periodically flushed to global memory counters.
Per thread counters avoid the contention that updating global counters
on each alloc/free would cause.

Policy evaluation periodically checks these counters and raises errors
if allocated memory has exceeded the configured limit.

Currently memory limit capability is exposed only to FFI and C#.

Also update mimalloc to v2.2.6

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

tests/memory_limits.rs

@@ -0,0 +1,187 @@
+#![cfg(all(feature = "mimalloc", feature = "allocator-memory-limits"))]
+
+use std::sync::{Mutex, OnceLock};
+
+use anyhow::Error;
+use mimalloc::global_allocation_stats_snapshot;
+use regorus::{set_global_memory_limit, Engine, LimitError, Value};
+
+#[cfg(feature = "rvm")]
+use regorus::languages::rego::compiler::Compiler;
+#[cfg(feature = "rvm")]
+use regorus::rvm::vm::RegoVM;
+#[cfg(feature = "rvm")]
+use regorus::rvm::vm::VmError;
+#[cfg(feature = "rvm")]
+use regorus::Rc;
+
+static LIMIT_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
+
+struct LimitGuard {
+    _guard: std::sync::MutexGuard<'static, ()>,
+}
+
+impl LimitGuard {
+    fn lock() -> Self {
+        let mutex = LIMIT_LOCK.get_or_init(|| Mutex::new(()));
+        let guard = mutex.lock().expect("limit mutex poisoned");
+        // Start with no global limit while the caller prepares state.
+        set_global_memory_limit(None);
+        Self { _guard: guard }
+    }
+
+    fn set_below_current_usage(&mut self) {
+        self.set_absolute_limit(1);
+    }
+
+    fn set_with_additional_budget(&mut self, budget: u64) {
+        self.set_with_usage_limit(|usage| usage.saturating_add(budget));
+    }
+
+    fn set_absolute_limit(&mut self, limit: u64) {
+        set_global_memory_limit(Some(limit));
+    }
+
+    fn set_with_usage_limit<F>(&mut self, calc: F)
+    where
+        F: FnOnce(u64) -> u64,
+    {
+        let usage = global_allocation_stats_snapshot().allocated as u64;
+        let limit = calc(usage);
+        self.set_absolute_limit(limit);
+    }
+}
+
+impl Drop for LimitGuard {
+    fn drop(&mut self) {
+        set_global_memory_limit(None);
+    }
+}
+
+const SIMPLE_MODULE: &str = r#"
+package limit
+
+allow if {
+    true
+}
+"#;
+
+const LARGE_PARSE_MODULE: &str = r#"
+package limit
+
+large_array := json.unmarshal(data.limit.large_json)
+"#;
+
+fn assert_memory_limit_error(err: &Error) {
+    match err.downcast_ref::<LimitError>() {
+        Some(LimitError::MemoryLimitExceeded { .. }) => {}
+        Some(other) => panic!("unexpected limit error variant: {other:?}"),
+        None => panic!("expected memory limit error, got: {err}"),
+    }
+}
+
+fn large_json_data(elements: usize) -> Value {
+    let mut payload = String::with_capacity(elements * 6);
+    payload.push('[');
+    for i in 0..elements {
+        if i > 0 {
+            payload.push(',');
+        }
+        payload.push_str(&i.to_string());
+    }
+    payload.push(']');
+
+    let json = serde_json::json!({
+        "limit": {
+            "large_json": payload,
+        }
+    });
+
+    Value::from_json_str(&json.to_string()).expect("valid JSON")
+}
+
+fn new_engine_with_module(module: &str) -> Engine {
+    let mut engine = Engine::new();
+    engine
+        .add_policy("limit.rego".to_string(), module.to_string())
+        .expect("add policy");
+    engine
+}
+
+#[test]
+fn interpreter_memory_limit_on_entry() {
+    let mut guard = LimitGuard::lock();
+    let mut engine = new_engine_with_module(SIMPLE_MODULE);
+    guard.set_below_current_usage();
+    let err = engine
+        .eval_query("data.limit.allow".to_string(), false)
+        .expect_err("expected interpreter memory limit error");
+    assert_memory_limit_error(&err);
+}
+
+#[cfg(feature = "rvm")]
+#[test]
+fn vm_memory_limit_on_entry() {
+    let mut guard = LimitGuard::lock();
+    let mut engine = new_engine_with_module(SIMPLE_MODULE);
+    let entrypoint = Rc::from("data.limit.allow");
+    let compiled = engine
+        .compile_with_entrypoint(&entrypoint)
+        .expect("compile policy for VM");
+    let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()])
+        .expect("compile VM program");
+
+    let mut vm = RegoVM::new();
+    vm.load_program(program);
+    vm.set_data(engine.get_data()).expect("set data");
+    vm.set_input(Value::Undefined);
+
+    guard.set_below_current_usage();
+    match vm.execute() {
+        Err(VmError::MemoryLimitExceeded { .. }) => {}
+        Err(other) => panic!("expected VM memory limit error, got {other}"),
+        Ok(value) => panic!("expected VM memory limit error, got value {value:?}"),
+    }
+}
+
+#[test]
+fn interpreter_memory_limit_during_large_allocation() {
+    let mut guard = LimitGuard::lock();
+    let mut engine = new_engine_with_module(LARGE_PARSE_MODULE);
+    let large_data = large_json_data(200_000);
+    engine.add_data(large_data).expect("add large JSON data");
+
+    guard.set_with_additional_budget(0);
+    let err = engine
+        .eval_rule("data.limit.large_array".to_string())
+        .expect_err("expected interpreter memory limit error while parsing");
+    assert_memory_limit_error(&err);
+}
+
+#[cfg(feature = "rvm")]
+#[test]
+fn vm_memory_limit_during_large_allocation() {
+    let mut guard = LimitGuard::lock();
+    let mut engine = new_engine_with_module(LARGE_PARSE_MODULE);
+    let large_data = large_json_data(200_000);
+    engine.add_data(large_data).expect("add large JSON data");
+
+    let entrypoint = Rc::from("data.limit.large_array");
+    let compiled = engine
+        .compile_with_entrypoint(&entrypoint)
+        .expect("compile policy for VM");
+    let program = Compiler::compile_from_policy(&compiled, &[entrypoint.as_ref()])
+        .expect("compile VM program");
+
+    let mut vm = RegoVM::new();
+    vm.load_program(program);
+    vm.set_data(engine.get_data()).expect("set data");
+    vm.set_input(Value::Undefined);
+
+    guard.set_with_additional_budget(0);
+    match vm.execute() {
+        Err(VmError::MemoryLimitExceeded { .. }) => {}
+        Err(other) => panic!("expected VM memory limit error, got {other}"),
+        Ok(value) => panic!("expected VM memory limit error, got value {value:?}"),
+    }
+}