use llmosafe::ResourceGuard;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PressureZone {
Green,
Yellow,
Orange,
Red,
}
#[derive(Debug, Clone)]
pub struct CacheDirective {
pub zone: PressureZone,
pub pressure: u8,
pub evict_fraction: f64,
pub allow_new_entries: bool,
pub allow_mmap_pin: bool,
}
#[allow(clippy::missing_fields_in_debug)]
pub struct AdaptiveCachePolicy {
guard: ResourceGuard,
ceiling: usize,
}
impl AdaptiveCachePolicy {
#[must_use]
#[allow(clippy::cast_precision_loss, clippy::cast_sign_loss, clippy::cast_possible_truncation, clippy::as_conversions)]
pub fn new(ceiling_fraction: f64) -> Self {
let guard = ResourceGuard::auto(ceiling_fraction);
let ceiling = (ResourceGuard::system_memory_bytes() as f64 * ceiling_fraction) as usize;
Self { guard, ceiling }
}
#[must_use]
pub fn new_with_guard(guard: ResourceGuard, ceiling_bytes: usize) -> Self {
Self {
guard,
ceiling: ceiling_bytes,
}
}
#[must_use]
pub fn guard(&self) -> &ResourceGuard {
&self.guard
}
#[must_use]
pub fn directive(&self) -> CacheDirective {
let pressure = self.guard.pressure();
let (zone, evict_fraction, allow_new_entries, allow_mmap_pin) = match pressure {
0..=40 => (PressureZone::Green, 0.0_f64, true, true),
41..=70 => (PressureZone::Yellow, 0.1_f64, true, true),
71..=90 => (PressureZone::Orange, 0.5_f64, false, false),
_ => (PressureZone::Red, 1.0_f64, false, false),
};
CacheDirective {
zone,
pressure,
evict_fraction,
allow_new_entries,
allow_mmap_pin,
}
}
#[must_use]
pub fn pressure(&self) -> u8 {
self.guard.pressure()
}
#[must_use]
pub fn rss_bytes(&self) -> usize {
ResourceGuard::current_rss_bytes()
}
#[must_use]
pub fn ceiling_bytes(&self) -> usize {
self.ceiling
}
}
impl std::fmt::Debug for AdaptiveCachePolicy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AdaptiveCachePolicy")
.field("ceiling", &self.ceiling)
.finish_non_exhaustive()
}
}
impl PressureZone {
#[must_use]
pub fn from_pressure(pressure: u8) -> Self {
match pressure {
0..=40 => Self::Green,
41..=70 => Self::Yellow,
71..=90 => Self::Orange,
_ => Self::Red,
}
}
}
#[cfg(test)]
#[allow(clippy::as_conversions, clippy::unwrap_used, clippy::indexing_slicing)]
mod tests {
use super::*;
#[test]
fn zone_classification_all_zones() {
assert_eq!(PressureZone::from_pressure(0), PressureZone::Green);
assert_eq!(PressureZone::from_pressure(20), PressureZone::Green);
assert_eq!(PressureZone::from_pressure(40), PressureZone::Green);
assert_eq!(PressureZone::from_pressure(41), PressureZone::Yellow);
assert_eq!(PressureZone::from_pressure(55), PressureZone::Yellow);
assert_eq!(PressureZone::from_pressure(70), PressureZone::Yellow);
assert_eq!(PressureZone::from_pressure(71), PressureZone::Orange);
assert_eq!(PressureZone::from_pressure(80), PressureZone::Orange);
assert_eq!(PressureZone::from_pressure(90), PressureZone::Orange);
assert_eq!(PressureZone::from_pressure(91), PressureZone::Red);
assert_eq!(PressureZone::from_pressure(100), PressureZone::Red);
}
#[test]
fn directive_fields_per_zone() {
let policy = AdaptiveCachePolicy::new(0.75);
let green = CacheDirective {
zone: PressureZone::Green,
pressure: 25,
evict_fraction: 0.0,
allow_new_entries: true,
allow_mmap_pin: true,
};
let yellow = CacheDirective {
zone: PressureZone::Yellow,
pressure: 55,
evict_fraction: 0.1,
allow_new_entries: true,
allow_mmap_pin: true,
};
let orange = CacheDirective {
zone: PressureZone::Orange,
pressure: 80,
evict_fraction: 0.5,
allow_new_entries: false,
allow_mmap_pin: false,
};
let red = CacheDirective {
zone: PressureZone::Red,
pressure: 95,
evict_fraction: 1.0,
allow_new_entries: false,
allow_mmap_pin: false,
};
for (label, pressure, expected) in [
("green", 25, &green),
("yellow", 55, &yellow),
("orange", 80, &orange),
("red", 95, &red),
] {
let zone = PressureZone::from_pressure(pressure);
let (evict_fraction, allow_new_entries, allow_mmap_pin) = match zone {
PressureZone::Green => (0.0_f64, true, true),
PressureZone::Yellow => (0.1_f64, true, true),
PressureZone::Orange => (0.5_f64, false, false),
PressureZone::Red => (1.0_f64, false, false),
};
assert_eq!(
zone, expected.zone,
"{label}: zone mismatch for pressure {pressure}"
);
assert!(
(evict_fraction - expected.evict_fraction).abs() < f64::EPSILON,
"{label}: evict_fraction mismatch"
);
assert_eq!(
allow_new_entries, expected.allow_new_entries,
"{label}: allow_new_entries mismatch"
);
assert_eq!(
allow_mmap_pin, expected.allow_mmap_pin,
"{label}: allow_mmap_pin mismatch"
);
}
let directive = policy.directive();
assert!(directive.pressure <= 100);
assert_eq!(directive.zone, PressureZone::from_pressure(directive.pressure));
}
#[test]
#[expect(clippy::integer_division, reason = "intentional ceiling via floor(sys_mem * fraction)")]
fn ceiling_fraction_applied_correctly() {
let sys_mem = ResourceGuard::system_memory_bytes();
if sys_mem == 0 {
return;
}
let expected_ceiling = sys_mem / 2;
let policy = AdaptiveCachePolicy::new(0.5_f64);
assert_eq!(
policy.ceiling_bytes(),
expected_ceiling,
"ceiling_bytes should equal system_memory * 0.5"
);
let expected_75 = sys_mem.saturating_mul(3) / 4;
let policy_75 = AdaptiveCachePolicy::new(0.75_f64);
assert_eq!(
policy_75.ceiling_bytes(),
expected_75,
"ceiling_bytes with 0.75 fraction"
);
}
#[test]
fn pressure_returns_bounded_value() {
let policy = AdaptiveCachePolicy::new(0.75);
let pressure = policy.pressure();
assert!(
pressure <= 100,
"pressure {pressure} should be <= 100"
);
}
#[test]
fn rss_bytes_is_non_negative() {
let policy = AdaptiveCachePolicy::new(0.75);
let rss = policy.rss_bytes();
assert!(
rss <= ResourceGuard::system_memory_bytes(),
"RSS {rss} should not exceed system memory"
);
}
#[test]
fn debug_impl_works() {
let policy = AdaptiveCachePolicy::new(0.75);
let debug_str = format!("{policy:?}");
assert!(
debug_str.contains("AdaptiveCachePolicy"),
"Debug output should contain type name"
);
assert!(
debug_str.contains("ceiling"),
"Debug output should contain ceiling field"
);
}
#[test]
#[expect(clippy::integer_division, reason = "intentional ceiling computation")]
fn new_with_guard_shares_ceiling() {
let sys_mem = ResourceGuard::system_memory_bytes();
if sys_mem == 0 {
return;
}
let ceiling = sys_mem / 2;
let guard = ResourceGuard::new(ceiling);
let policy = AdaptiveCachePolicy::new_with_guard(guard, ceiling);
assert_eq!(
policy.ceiling_bytes(), ceiling,
"new_with_guard should use the provided ceiling"
);
assert_eq!(
policy.guard().pressure(),
ResourceGuard::new(ceiling).pressure(),
"guard pressure should match"
);
}
}