use serde::{Deserialize, Serialize};
use super::{ByteBudgetConfig, Hysteresis};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SelfRegulationProfile {
#[default]
Throughput,
Balanced,
LowLatency,
}
impl SelfRegulationProfile {
#[must_use]
fn byte_budget_envelope(self) -> (u64, u64, u64, usize) {
match self {
Self::Throughput => (16 * 1024 * 1024, 128 * 1024 * 1024, 512 * 1024, 2000),
Self::Balanced => (8 * 1024 * 1024, 64 * 1024 * 1024, 256 * 1024, 1000),
Self::LowLatency => (1024 * 1024, 16 * 1024 * 1024, 128 * 1024, 500),
}
}
}
const fn default_enabled() -> bool {
true
}
fn default_pause_above() -> f64 {
0.80
}
fn default_resume_below() -> f64 {
0.65
}
fn default_target_rho() -> f64 {
0.7
}
fn default_md_factor() -> f64 {
0.5
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default)]
pub struct SelfRegulationConfig {
pub enabled: bool,
pub profile: SelfRegulationProfile,
pub pause_above: f64,
pub resume_below: f64,
pub target_rho: f64,
pub md_factor: f64,
}
impl Default for SelfRegulationConfig {
fn default() -> Self {
Self {
enabled: default_enabled(),
profile: SelfRegulationProfile::default(),
pause_above: default_pause_above(),
resume_below: default_resume_below(),
target_rho: default_target_rho(),
md_factor: default_md_factor(),
}
}
}
impl SelfRegulationConfig {
#[must_use]
pub fn from_cascade() -> Self {
#[cfg(feature = "config")]
{
if let Some(cfg) = crate::config::try_get()
&& let Ok(value) = cfg.unmarshal_key_registered::<Self>("self_regulation")
{
return value;
}
}
Self::default()
}
#[must_use]
pub fn hysteresis(&self) -> Hysteresis {
Hysteresis::new(self.pause_above, self.resume_below).unwrap_or_else(|e| {
tracing::warn!(
error = %e,
"invalid self_regulation hysteresis band; using defaults 0.80/0.65"
);
Hysteresis::new(default_pause_above(), default_resume_below())
.expect("default band is valid")
})
}
#[must_use]
pub fn byte_budget_config(&self) -> ByteBudgetConfig {
let (start_bytes, max_bytes, ai_step, record_cap) = self.profile.byte_budget_envelope();
ByteBudgetConfig {
start_bytes,
max_bytes,
ai_step,
record_cap,
target_rho: self.target_rho,
md_factor: self.md_factor,
..ByteBudgetConfig::default()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_is_enabled() {
let cfg = SelfRegulationConfig::default();
assert!(cfg.enabled, "governor is ON by default (opt-out)");
assert_eq!(cfg.profile, SelfRegulationProfile::Throughput);
}
#[test]
fn from_cascade_falls_back_to_default_on() {
let cfg = SelfRegulationConfig::from_cascade();
assert!(cfg.enabled);
}
#[test]
fn hysteresis_uses_config_band() {
let cfg = SelfRegulationConfig {
pause_above: 0.9,
resume_below: 0.5,
..Default::default()
};
let h = cfg.hysteresis();
assert!((h.pause_above - 0.9).abs() < 1e-9);
assert!((h.resume_below - 0.5).abs() < 1e-9);
}
#[test]
fn inverted_band_falls_back_to_defaults() {
let cfg = SelfRegulationConfig {
pause_above: 0.3,
resume_below: 0.8, ..Default::default()
};
let h = cfg.hysteresis();
assert!((h.pause_above - 0.80).abs() < 1e-9);
assert!((h.resume_below - 0.65).abs() < 1e-9);
}
#[test]
fn profile_sizes_the_byte_budget() {
let tp = SelfRegulationConfig {
profile: SelfRegulationProfile::Throughput,
..Default::default()
}
.byte_budget_config();
let ll = SelfRegulationConfig {
profile: SelfRegulationProfile::LowLatency,
..Default::default()
}
.byte_budget_config();
assert!(
tp.start_bytes > ll.start_bytes,
"throughput starts bigger than low-latency"
);
assert!(tp.max_bytes > ll.max_bytes);
}
#[cfg(feature = "config")]
#[test]
fn serde_roundtrip_and_disabled_parse() {
let yaml = "enabled: false\nprofile: low_latency\n";
let cfg: SelfRegulationConfig = serde_yaml_ng::from_str(yaml).unwrap();
assert!(!cfg.enabled);
assert_eq!(cfg.profile, SelfRegulationProfile::LowLatency);
assert!((cfg.pause_above - 0.80).abs() < 1e-9);
}
#[cfg(feature = "config")]
#[test]
fn profile_serialises_snake_case() {
let j = serde_json::to_string(&SelfRegulationProfile::LowLatency).unwrap();
assert_eq!(j, "\"low_latency\"");
let j = serde_json::to_string(&SelfRegulationProfile::Throughput).unwrap();
assert_eq!(j, "\"throughput\"");
let j = serde_json::to_string(&SelfRegulationProfile::Balanced).unwrap();
assert_eq!(j, "\"balanced\"");
}
}