mirror of
https://github.com/microsoft/regorus.git
synced 2026-08-05 02:16:11 +00:00
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>
This commit is contained in:
Anand Krishnamoorthi
committed by
GitHub
GitHub
parent
80686d6ed1
commit
fd59bb5a91
68
src/utils/limits/error.rs
Normal file
68
src/utils/limits/error.rs
Normal file
@@ -0,0 +1,68 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#![allow(dead_code)]
|
||||
|
||||
use core::fmt;
|
||||
use core::time::Duration;
|
||||
|
||||
/// Errors reported when execution time or memory ceilings are enforced.
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub enum LimitError {
|
||||
/// Reported when the execution timer observes elapsed time beyond the configured limit.
|
||||
TimeLimitExceeded {
|
||||
/// Elapsed work duration when the threshold was exceeded.
|
||||
elapsed: Duration,
|
||||
/// Configured time limit.
|
||||
limit: Duration,
|
||||
},
|
||||
/// Reported when the memory tracker estimates usage beyond the configured limit.
|
||||
MemoryLimitExceeded {
|
||||
/// Estimated bytes in use when the limit was detected.
|
||||
usage: u64,
|
||||
/// Configured memory ceiling in bytes.
|
||||
limit: u64,
|
||||
},
|
||||
}
|
||||
|
||||
impl fmt::Debug for LimitError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Self::TimeLimitExceeded { elapsed, limit } => f
|
||||
.debug_struct("TimeLimitExceeded")
|
||||
.field("elapsed", elapsed)
|
||||
.field("limit", limit)
|
||||
.finish(),
|
||||
Self::MemoryLimitExceeded { usage, limit } => f
|
||||
.debug_struct("MemoryLimitExceeded")
|
||||
.field("usage", usage)
|
||||
.field("limit", limit)
|
||||
.finish(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for LimitError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Self::TimeLimitExceeded { elapsed, limit } => {
|
||||
let elapsed_ns = elapsed.as_nanos();
|
||||
let limit_ns = limit.as_nanos();
|
||||
write!(
|
||||
f,
|
||||
"execution exceeded time limit (elapsed={}ns, limit={}ns)",
|
||||
elapsed_ns, limit_ns
|
||||
)
|
||||
}
|
||||
Self::MemoryLimitExceeded { usage, limit } => {
|
||||
write!(
|
||||
f,
|
||||
"execution exceeded memory limit (usage={} bytes, limit={} bytes)",
|
||||
usage, limit
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl core::error::Error for LimitError {}
|
||||
213
src/utils/limits/memory.rs
Normal file
213
src/utils/limits/memory.rs
Normal file
@@ -0,0 +1,213 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
use super::error::LimitError;
|
||||
use core::cell::Cell;
|
||||
use core::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::thread_local;
|
||||
|
||||
static GLOBAL_MEMORY_LIMIT: AtomicU64 = AtomicU64::new(u64::MAX);
|
||||
|
||||
// Maximum iteration count before forcing a global memory check.
|
||||
const MEMORY_CHECK_STRIDE: u32 = 16;
|
||||
// Pending per-thread allocator delta (bytes) that triggers a memory check. Chosen at 32 KiB to
|
||||
// catch short bursts before they exceed typical entry budgets while still amortizing the atomic.
|
||||
const MEMORY_CHECK_DELTA_BYTES: u64 = 32 * 1024;
|
||||
|
||||
thread_local! {
|
||||
// Per-thread stride counter used to amortize global memory checks.
|
||||
static MEMORY_CHECK_TICKS: Cell<u32> = const { Cell::new(0) };
|
||||
}
|
||||
|
||||
/// Sets the global memory limit in bytes; `None` disables enforcement.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```rust
|
||||
/// #![cfg(feature = "allocator-memory-limits")]
|
||||
/// use regorus::{set_global_memory_limit, Engine, LimitError, Value};
|
||||
///
|
||||
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// const LIMIT: u64 = 32 * 1024;
|
||||
/// set_global_memory_limit(Some(LIMIT));
|
||||
///
|
||||
/// let mut engine = Engine::new();
|
||||
/// engine.add_policy(
|
||||
/// "limit.rego".to_string(),
|
||||
/// "package limit\nallow if input.blob != \"\"".to_string(),
|
||||
/// )?;
|
||||
///
|
||||
/// let payload = format!("{{\"blob\":\"{}\"}}", "x".repeat(128 * 1024));
|
||||
///
|
||||
/// // Prepare input while the limit is relaxed so parsing succeeds on constrained builds.
|
||||
/// set_global_memory_limit(Some(LIMIT * 32));
|
||||
/// engine.set_input(Value::from_json_str(&payload)?);
|
||||
///
|
||||
/// // Tighten the budget and observe evaluation fail fast on the smaller limit.
|
||||
/// set_global_memory_limit(Some(LIMIT));
|
||||
/// let err = engine
|
||||
/// .eval_query("data.limit.allow".to_string(), false)
|
||||
/// .unwrap_err();
|
||||
/// let limit = err.downcast_ref::<LimitError>().copied();
|
||||
/// assert!(matches!(
|
||||
/// limit,
|
||||
/// Some(LimitError::MemoryLimitExceeded { limit: LIMIT, .. })
|
||||
/// ));
|
||||
///
|
||||
/// // Raise the ceiling and the same evaluation succeeds.
|
||||
/// set_global_memory_limit(Some(LIMIT * 32));
|
||||
/// let result = engine.eval_query("data.limit.allow".to_string(), false)?;
|
||||
/// assert_eq!(result.result.len(), 1);
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub fn set_global_memory_limit(limit: Option<u64>) {
|
||||
let value = limit.unwrap_or(u64::MAX);
|
||||
GLOBAL_MEMORY_LIMIT.store(value, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Flushes this thread's Regorus allocation counters into the global aggregates.
|
||||
///
|
||||
/// Regorus batches per-thread deltas to keep hot paths uncontended. Automatic publication occurs when
|
||||
/// [`check_global_memory_limit`] runs, when the configured threshold (see
|
||||
/// [`thread_memory_flush_threshold`]) is exceeded, or when allocation statistics are queried.
|
||||
/// Workloads that burst and then idle—or threads that are about to terminate—can call this helper to
|
||||
/// push their pending deltas immediately and avoid dropping buffered usage.
|
||||
///
|
||||
/// ```rust
|
||||
/// #![cfg(feature = "allocator-memory-limits")]
|
||||
/// use regorus::flush_thread_memory_counters;
|
||||
///
|
||||
/// flush_thread_memory_counters();
|
||||
/// ```
|
||||
///
|
||||
/// Flush before a worker thread terminates:
|
||||
/// ```no_run
|
||||
/// # use regorus::flush_thread_memory_counters;
|
||||
/// std::thread::spawn(|| {
|
||||
/// // ... do work ...
|
||||
/// flush_thread_memory_counters(); // publish before exiting
|
||||
/// });
|
||||
/// ```
|
||||
pub fn flush_thread_memory_counters() {
|
||||
mimalloc::limits::flush_thread_counters();
|
||||
}
|
||||
|
||||
/// Configures the per-thread flush threshold in bytes.
|
||||
///
|
||||
/// Each thread buffers allocation deltas locally and publishes them automatically once the absolute
|
||||
/// difference since the last flush exceeds this threshold. Passing `None` restores the default of
|
||||
/// 1 MiB. Setting the threshold to zero disables automatic flushing, requiring manual calls to
|
||||
/// [`flush_thread_memory_counters()`]. Values larger than [`i64::MAX`] are saturated.
|
||||
pub fn set_thread_flush_threshold_override(bytes: Option<u64>) {
|
||||
mimalloc::limits::set_thread_flush_threshold(bytes);
|
||||
}
|
||||
|
||||
/// Returns the current per-thread flush threshold in bytes, if automatic flushing is enabled.
|
||||
///
|
||||
/// When the threshold is disabled (zero or negative), `None` is returned. Otherwise the value
|
||||
/// represents the absolute delta that will trigger an automatic flush.
|
||||
///
|
||||
/// ```rust
|
||||
/// #![cfg(feature = "allocator-memory-limits")]
|
||||
/// use regorus::{set_thread_flush_threshold_override, thread_memory_flush_threshold};
|
||||
///
|
||||
/// set_thread_flush_threshold_override(Some(256 * 1024));
|
||||
/// assert_eq!(thread_memory_flush_threshold(), Some(256 * 1024));
|
||||
///
|
||||
/// set_thread_flush_threshold_override(Some(0));
|
||||
/// assert_eq!(thread_memory_flush_threshold(), None);
|
||||
/// ```
|
||||
pub fn thread_memory_flush_threshold() -> Option<u64> {
|
||||
mimalloc::limits::thread_flush_threshold()
|
||||
}
|
||||
|
||||
/// Validates the current allocation usage against the configured global limit.
|
||||
///
|
||||
/// This helper consults the global memory limit and returns [`LimitError::MemoryLimitExceeded`]
|
||||
/// if tracked usage exceeds the configured ceiling.
|
||||
///
|
||||
/// ```rust
|
||||
/// #![cfg(feature = "allocator-memory-limits")]
|
||||
/// use regorus::{check_global_memory_limit, set_global_memory_limit, LimitError};
|
||||
///
|
||||
/// const LIMIT: u64 = 32 * 1024;
|
||||
/// set_global_memory_limit(Some(LIMIT));
|
||||
/// let _buffer = vec![0u8; 128 * 1024];
|
||||
///
|
||||
/// let outcome = check_global_memory_limit();
|
||||
/// assert!(matches!(
|
||||
/// outcome,
|
||||
/// Err(LimitError::MemoryLimitExceeded { limit: LIMIT, .. })
|
||||
/// ));
|
||||
///
|
||||
/// set_global_memory_limit(Some(LIMIT * 32));
|
||||
/// check_global_memory_limit().unwrap();
|
||||
/// ```
|
||||
pub fn check_global_memory_limit() -> Result<(), LimitError> {
|
||||
if let Some(limit) = global_memory_limit() {
|
||||
let (stats, _) = mimalloc::allocation_stats_snapshot();
|
||||
let usage = stats.allocated;
|
||||
if usage > limit {
|
||||
return Err(LimitError::MemoryLimitExceeded { usage, limit });
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Enforces the currently configured memory ceiling, if any.
|
||||
#[inline]
|
||||
pub fn enforce_memory_limit() -> Result<(), LimitError> {
|
||||
check_global_memory_limit()
|
||||
}
|
||||
|
||||
/// Performs a throttled global memory check, combining a lightweight stride counter
|
||||
/// with the allocator's pending per-thread delta to avoid excessive atomics on hot paths.
|
||||
/// The global limit is only evaluated when the stride expires or the pending delta crosses
|
||||
/// the configured watermark.
|
||||
pub(super) fn check_memory_limit_if_needed() -> Result<(), LimitError> {
|
||||
if global_memory_limit().is_none() {
|
||||
// Reset state when enforcement is disabled to avoid stale counters.
|
||||
MEMORY_CHECK_TICKS.with(|ticks| ticks.set(0));
|
||||
let _ = mimalloc::limits::take_thread_flushed_since_check_flag();
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Inspect unflushed per-thread allocator usage to catch large bursts early.
|
||||
let pending_delta = mimalloc::limits::thread_allocation_pending_delta().unsigned_abs() as u64;
|
||||
let flush_hint = mimalloc::limits::take_thread_flushed_since_check_flag();
|
||||
|
||||
MEMORY_CHECK_TICKS.with(|ticks| {
|
||||
let next = ticks.get().wrapping_add(1);
|
||||
if flush_hint || next >= MEMORY_CHECK_STRIDE || pending_delta >= MEMORY_CHECK_DELTA_BYTES {
|
||||
// Publish usage when either the stride or delta threshold is hit.
|
||||
ticks.set(0);
|
||||
check_global_memory_limit()
|
||||
} else {
|
||||
ticks.set(next);
|
||||
Ok(())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the currently-configured global memory limit in bytes, if any.
|
||||
///
|
||||
/// When [`set_global_memory_limit`] is called with `Some(value)`, that value is reported here.
|
||||
/// Passing `None` to [`set_global_memory_limit`] removes the limit, causing this function to return
|
||||
/// `None`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```rust
|
||||
/// #![cfg(feature = "allocator-memory-limits")]
|
||||
/// use regorus::{global_memory_limit, set_global_memory_limit};
|
||||
/// set_global_memory_limit(Some(123));
|
||||
/// assert_eq!(global_memory_limit(), Some(123));
|
||||
///
|
||||
/// set_global_memory_limit(None);
|
||||
/// assert_eq!(global_memory_limit(), None);
|
||||
/// ```
|
||||
pub fn global_memory_limit() -> Option<u64> {
|
||||
let limit = GLOBAL_MEMORY_LIMIT.load(Ordering::Relaxed);
|
||||
(limit != u64::MAX).then_some(limit)
|
||||
}
|
||||
37
src/utils/limits/mod.rs
Normal file
37
src/utils/limits/mod.rs
Normal file
@@ -0,0 +1,37 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
//! Helpers for cooperative execution time and memory limits.
|
||||
|
||||
#![allow(dead_code)]
|
||||
|
||||
mod error;
|
||||
#[cfg(feature = "allocator-memory-limits")]
|
||||
mod memory;
|
||||
#[allow(unused_imports)]
|
||||
pub use error::LimitError;
|
||||
#[allow(unused_imports)]
|
||||
#[cfg(feature = "allocator-memory-limits")]
|
||||
pub use memory::{
|
||||
check_global_memory_limit, enforce_memory_limit, flush_thread_memory_counters,
|
||||
global_memory_limit, set_global_memory_limit, set_thread_flush_threshold_override,
|
||||
thread_memory_flush_threshold,
|
||||
};
|
||||
|
||||
#[cfg(feature = "allocator-memory-limits")]
|
||||
#[inline]
|
||||
pub fn check_memory_limit_if_needed() -> core::result::Result<(), LimitError> {
|
||||
memory::check_memory_limit_if_needed()
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "allocator-memory-limits"))]
|
||||
#[inline]
|
||||
pub fn enforce_memory_limit() -> core::result::Result<(), LimitError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "allocator-memory-limits"))]
|
||||
#[inline]
|
||||
pub fn check_memory_limit_if_needed() -> core::result::Result<(), LimitError> {
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user