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:
committed by
GitHub
parent
80686d6ed1
commit
fd59bb5a91
+45
-4
@@ -22,7 +22,7 @@ use core::convert::AsRef;
|
||||
use core::str::FromStr;
|
||||
|
||||
use anyhow::{anyhow, bail, Result};
|
||||
use serde::de::{self, Deserializer, MapAccess, SeqAccess, Visitor};
|
||||
use serde::de::{self, Deserializer, Error as DeError, MapAccess, SeqAccess, Visitor};
|
||||
use serde::ser::{SerializeMap, Serializer};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
@@ -70,6 +70,16 @@ pub enum Value {
|
||||
Undefined,
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn enforce_limit_anyhow() -> Result<()> {
|
||||
crate::utils::limits::check_memory_limit_if_needed().map_err(|err| anyhow!(err))
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn enforce_limit_for<E: DeError>() -> core::result::Result<(), E> {
|
||||
crate::utils::limits::check_memory_limit_if_needed().map_err(|err| E::custom(err.to_string()))
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
impl Serialize for Value {
|
||||
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
||||
@@ -199,6 +209,8 @@ impl<'de> Visitor<'de> for ValueVisitor {
|
||||
let mut arr = vec![];
|
||||
while let Some(v) = visitor.next_element()? {
|
||||
arr.push(v);
|
||||
// Enforce allocator limit while expanding a deserialized array.
|
||||
enforce_limit_for::<V::Error>()?;
|
||||
}
|
||||
Ok(Value::from(arr))
|
||||
}
|
||||
@@ -218,8 +230,12 @@ impl<'de> Visitor<'de> for ValueVisitor {
|
||||
}
|
||||
let mut map = BTreeMap::new();
|
||||
map.insert(key, value);
|
||||
// Enforce allocator limit while expanding a deserialized object.
|
||||
enforce_limit_for::<V::Error>()?;
|
||||
while let Some((key, value)) = visitor.next_entry()? {
|
||||
map.insert(key, value);
|
||||
// Enforce allocator limit while expanding a deserialized object.
|
||||
enforce_limit_for::<V::Error>()?;
|
||||
}
|
||||
Ok(Value::from(map))
|
||||
} else {
|
||||
@@ -334,7 +350,24 @@ impl Value {
|
||||
/// # }
|
||||
/// ```
|
||||
pub fn from_json_str(json: &str) -> Result<Value> {
|
||||
serde_json::from_str(json).map_err(anyhow::Error::msg)
|
||||
match serde_json::from_str::<Value>(json) {
|
||||
Ok(value) => Ok(value),
|
||||
Err(err) => {
|
||||
#[cfg(feature = "allocator-memory-limits")]
|
||||
{
|
||||
// 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)),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "allocator-memory-limits"))]
|
||||
{
|
||||
Err(anyhow!(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Deserialize a [`Value`] from a file containing JSON.
|
||||
@@ -1305,6 +1338,8 @@ impl Value {
|
||||
if let Value::Object(map) = self {
|
||||
if map.get(&key).is_none() {
|
||||
Rc::make_mut(map).insert(key.clone(), Value::Undefined);
|
||||
// Enforce allocator limit while creating nested object entries.
|
||||
enforce_limit_anyhow()?;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1330,7 +1365,9 @@ impl Value {
|
||||
match (self, &mut new) {
|
||||
(v @ Value::Undefined, _) => *v = new,
|
||||
(Value::Set(ref mut set), Value::Set(new)) => {
|
||||
Rc::make_mut(set).append(Rc::make_mut(new))
|
||||
Rc::make_mut(set).append(Rc::make_mut(new));
|
||||
// Enforce allocator limit after merging set entries.
|
||||
enforce_limit_anyhow()?;
|
||||
}
|
||||
(Value::Object(map), Value::Object(new)) => {
|
||||
for (k, v) in new.iter() {
|
||||
@@ -1343,7 +1380,11 @@ impl Value {
|
||||
serde_json::to_string_pretty(&v).map_err(anyhow::Error::msg)?,
|
||||
)
|
||||
}
|
||||
_ => Rc::make_mut(map).insert(k.clone(), v.clone()),
|
||||
_ => {
|
||||
Rc::make_mut(map).insert(k.clone(), v.clone());
|
||||
// Enforce allocator limit after merging object entries.
|
||||
enforce_limit_anyhow()?;
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user