use std::path::Path;
use serde::Serialize;
use crate::config::{EnvVar, GraphMemoryConfig, read_env_var, workspace_config};
pub const DEFAULT_SNAPSHOT_CAP_MB: u64 = 250;
pub const DEFAULT_PER_ALGORITHM_CAP_MB: u64 = 100;
pub const DEFAULT_DEGRADED_BELOW_PCT: u8 = 80;
pub const PER_NODE_BYTES: u64 = 32;
pub const PER_EDGE_BYTES: u64 = 96;
pub const LARGE_GRAPH_UNCACHED_CODE: &str = "large_graph_uncached";
pub const UNEXPECTED_GROWTH_CODE: &str = "unexpected_growth";
pub const MEMORY_PRESSURE_CODE: &str = "memory_pressure";
pub const ALGORITHM_MEMORY_CAP_CODE: &str = "algorithm_memory_cap";
pub const APPROACHING_CAP_CODE: &str = "snapshot_approaching_cap";
const DEFAULT_GROWTH_MULTIPLIER_BASIS_POINTS: u32 = 15_000;
const BYTES_PER_MIB: u64 = 1024 * 1024;
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct MemoryBudgetPolicy {
pub snapshot_cap_bytes: u64,
pub per_algorithm_cap_bytes: u64,
pub degraded_below_pct: u8,
pub growth_multiplier_basis_points: u32,
}
impl MemoryBudgetPolicy {
#[must_use]
pub const fn defaults() -> Self {
Self {
snapshot_cap_bytes: DEFAULT_SNAPSHOT_CAP_MB * 1024 * 1024,
per_algorithm_cap_bytes: DEFAULT_PER_ALGORITHM_CAP_MB * 1024 * 1024,
degraded_below_pct: DEFAULT_DEGRADED_BELOW_PCT,
growth_multiplier_basis_points: DEFAULT_GROWTH_MULTIPLIER_BASIS_POINTS,
}
}
#[must_use]
pub fn from_config(memory: &GraphMemoryConfig, env: MemoryBudgetEnv) -> Self {
let defaults = Self::defaults();
Self {
snapshot_cap_bytes: mb_to_bytes(
env.snapshot_cap_mb
.or(memory.snapshot_cap_mb)
.unwrap_or(DEFAULT_SNAPSHOT_CAP_MB),
),
per_algorithm_cap_bytes: mb_to_bytes(
env.per_algorithm_cap_mb
.or(memory.per_algorithm_cap_mb)
.unwrap_or(DEFAULT_PER_ALGORITHM_CAP_MB),
),
degraded_below_pct: clamp_percent_u8(
env.degraded_below_pct
.or(memory.degraded_below_pct)
.unwrap_or(u64::from(DEFAULT_DEGRADED_BELOW_PCT)),
),
growth_multiplier_basis_points: saturating_u32(
env.growth_multiplier_basis_points
.or(memory.growth_multiplier_basis_points)
.unwrap_or(u64::from(DEFAULT_GROWTH_MULTIPLIER_BASIS_POINTS)),
)
.max(1),
}
.with_zero_cap_defaults(defaults)
}
fn with_zero_cap_defaults(mut self, defaults: Self) -> Self {
if self.snapshot_cap_bytes == 0 {
self.snapshot_cap_bytes = defaults.snapshot_cap_bytes;
}
if self.per_algorithm_cap_bytes == 0 {
self.per_algorithm_cap_bytes = defaults.per_algorithm_cap_bytes;
}
self
}
}
impl Default for MemoryBudgetPolicy {
fn default() -> Self {
Self::defaults()
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct MemoryBudgetEnv {
pub snapshot_cap_mb: Option<u64>,
pub per_algorithm_cap_mb: Option<u64>,
pub degraded_below_pct: Option<u64>,
pub growth_multiplier_basis_points: Option<u64>,
}
impl MemoryBudgetEnv {
#[must_use]
pub fn current() -> Self {
Self {
snapshot_cap_mb: read_env_u64(EnvVar::GraphMemorySnapshotCapMb),
per_algorithm_cap_mb: read_env_u64(EnvVar::GraphMemoryPerAlgorithmCapMb),
degraded_below_pct: read_env_u64(EnvVar::GraphMemoryDegradedBelowPct),
growth_multiplier_basis_points: read_env_u64(
EnvVar::GraphMemoryGrowthMultiplierBasisPoints,
),
}
}
}
#[must_use]
pub fn workspace_memory_budget_policy(workspace_path: &Path) -> MemoryBudgetPolicy {
let memory = workspace_config(workspace_path)
.map(|config| config.graph.memory)
.unwrap_or_default();
MemoryBudgetPolicy::from_config(&memory, MemoryBudgetEnv::current())
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct MemoryBudgetRefusal {
pub code: &'static str,
pub severity: &'static str,
pub message: &'static str,
pub repair: &'static str,
pub observed_bytes: u64,
pub limit_bytes: u64,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum SnapshotAdmissionDecision {
Admit {
estimate_bytes: u64,
headroom_bytes: u64,
approaching_cap: bool,
},
Refuse(MemoryBudgetRefusal),
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum InBuildGrowthDecision {
Continue {
observed_bytes: u64,
allowed_bytes: u64,
},
Abort(MemoryBudgetRefusal),
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum AlgorithmAdmissionDecision {
Admit { combined_bytes: u64 },
Refuse(MemoryBudgetRefusal),
}
#[must_use]
pub fn estimate_snapshot_bytes(node_count: usize, edge_count: usize) -> u64 {
let nodes = u64::try_from(node_count).unwrap_or(u64::MAX);
let edges = u64::try_from(edge_count).unwrap_or(u64::MAX);
let node_bytes = nodes.saturating_mul(PER_NODE_BYTES);
let edge_bytes = edges.saturating_mul(PER_EDGE_BYTES);
node_bytes.saturating_add(edge_bytes)
}
#[must_use]
pub fn check_snapshot_admission(
estimate_bytes: u64,
policy: &MemoryBudgetPolicy,
) -> SnapshotAdmissionDecision {
if estimate_bytes > policy.snapshot_cap_bytes {
return SnapshotAdmissionDecision::Refuse(MemoryBudgetRefusal {
code: LARGE_GRAPH_UNCACHED_CODE,
severity: "high",
message: "graph snapshot estimate exceeds the configured cap; build skipped",
repair: "raise graph.memory.snapshot_cap_mb or shrink the workspace",
observed_bytes: estimate_bytes,
limit_bytes: policy.snapshot_cap_bytes,
});
}
let headroom_bytes = policy.snapshot_cap_bytes.saturating_sub(estimate_bytes);
let advisory_threshold = scale_by_basis_points(
policy.snapshot_cap_bytes,
u32::from(policy.degraded_below_pct) * 100,
);
SnapshotAdmissionDecision::Admit {
estimate_bytes,
headroom_bytes,
approaching_cap: estimate_bytes >= advisory_threshold,
}
}
#[must_use]
pub fn check_in_build_growth(
pre_build_estimate: u64,
observed_bytes: u64,
policy: &MemoryBudgetPolicy,
) -> InBuildGrowthDecision {
let allowed = scale_by_basis_points(pre_build_estimate, policy.growth_multiplier_basis_points);
if observed_bytes > allowed {
InBuildGrowthDecision::Abort(MemoryBudgetRefusal {
code: UNEXPECTED_GROWTH_CODE,
severity: "warning",
message: "in-build allocation grew past the tripwire; aborted and rolled back",
repair: "shrink the workspace, rerun snapshot refresh, or raise the growth multiplier",
observed_bytes,
limit_bytes: allowed,
})
} else {
InBuildGrowthDecision::Continue {
observed_bytes,
allowed_bytes: allowed,
}
}
}
#[must_use]
pub fn check_algorithm_admission(
active_resident_bytes: u64,
requested_bytes: u64,
policy: &MemoryBudgetPolicy,
) -> AlgorithmAdmissionDecision {
if requested_bytes > policy.per_algorithm_cap_bytes {
return AlgorithmAdmissionDecision::Refuse(MemoryBudgetRefusal {
code: ALGORITHM_MEMORY_CAP_CODE,
severity: "warning",
message: "requested algorithm working-set exceeds the per-algorithm cap; refused before allocating",
repair: "raise graph.memory.per_algorithm_cap_mb or use a sampled algorithm",
observed_bytes: requested_bytes,
limit_bytes: policy.per_algorithm_cap_bytes,
});
}
let combined_bytes = active_resident_bytes.saturating_add(requested_bytes);
if combined_bytes > policy.snapshot_cap_bytes {
return AlgorithmAdmissionDecision::Refuse(MemoryBudgetRefusal {
code: MEMORY_PRESSURE_CODE,
severity: "warning",
message: "active snapshots plus the requested algorithm would exceed the snapshot cap; refused before allocating",
repair: "wait for an active snapshot to release, raise graph.memory.snapshot_cap_mb, or skip the algorithm",
observed_bytes: combined_bytes,
limit_bytes: policy.snapshot_cap_bytes,
});
}
AlgorithmAdmissionDecision::Admit { combined_bytes }
}
fn scale_by_basis_points(bytes: u64, basis_points: u32) -> u64 {
let basis_points = u64::from(basis_points);
bytes
.saturating_mul(basis_points)
.checked_div(10_000)
.unwrap_or(u64::MAX)
}
fn mb_to_bytes(megabytes: u64) -> u64 {
megabytes.saturating_mul(BYTES_PER_MIB)
}
fn clamp_percent_u8(value: u64) -> u8 {
u8::try_from(value.min(100)).unwrap_or(100)
}
fn saturating_u32(value: u64) -> u32 {
u32::try_from(value).unwrap_or(u32::MAX)
}
fn read_env_u64(var: EnvVar) -> Option<u64> {
read_env_var(var).and_then(|raw| raw.parse::<u64>().ok())
}
#[cfg(test)]
mod tests {
use super::*;
fn mb(megabytes: u64) -> u64 {
megabytes * 1024 * 1024
}
#[test]
fn policy_reads_graph_memory_config_and_env_overrides() {
let config = GraphMemoryConfig {
snapshot_cap_mb: Some(64),
per_algorithm_cap_mb: Some(16),
degraded_below_pct: Some(75),
growth_multiplier_basis_points: Some(12_500),
};
let env = MemoryBudgetEnv {
snapshot_cap_mb: Some(32),
per_algorithm_cap_mb: None,
degraded_below_pct: Some(90),
growth_multiplier_basis_points: None,
};
let policy = MemoryBudgetPolicy::from_config(&config, env);
assert_eq!(policy.snapshot_cap_bytes, mb(32));
assert_eq!(policy.per_algorithm_cap_bytes, mb(16));
assert_eq!(policy.degraded_below_pct, 90);
assert_eq!(policy.growth_multiplier_basis_points, 12_500);
}
#[test]
fn estimate_uses_documented_per_node_and_per_edge_constants() {
let estimate = estimate_snapshot_bytes(10, 25);
assert_eq!(estimate, 10 * PER_NODE_BYTES + 25 * PER_EDGE_BYTES);
}
#[test]
fn estimate_saturates_at_u64_max_on_extreme_inputs() {
let extreme = estimate_snapshot_bytes(usize::MAX, usize::MAX);
assert_eq!(extreme, u64::MAX);
}
#[test]
fn snapshot_under_cap_admits_with_headroom() {
let policy = MemoryBudgetPolicy::defaults();
let decision = check_snapshot_admission(mb(50), &policy);
let SnapshotAdmissionDecision::Admit {
estimate_bytes,
headroom_bytes,
approaching_cap,
} = decision
else {
panic!("expected Admit");
};
assert_eq!(estimate_bytes, mb(50));
assert_eq!(headroom_bytes, mb(200));
assert!(!approaching_cap);
}
#[test]
fn snapshot_past_threshold_admits_with_approaching_cap_flag() {
let policy = MemoryBudgetPolicy::defaults();
let decision = check_snapshot_admission(mb(200), &policy);
let SnapshotAdmissionDecision::Admit {
approaching_cap, ..
} = decision
else {
panic!("expected Admit");
};
assert!(approaching_cap, "200 MB hits the 80% threshold");
}
#[test]
fn snapshot_over_cap_refuses_with_large_graph_uncached_code() {
let policy = MemoryBudgetPolicy::defaults();
let decision = check_snapshot_admission(mb(300), &policy);
let SnapshotAdmissionDecision::Refuse(refusal) = decision else {
panic!("expected Refuse");
};
assert_eq!(refusal.code, LARGE_GRAPH_UNCACHED_CODE);
assert_eq!(refusal.severity, "high");
assert_eq!(refusal.observed_bytes, mb(300));
assert_eq!(refusal.limit_bytes, mb(250));
}
#[test]
fn in_build_growth_within_tripwire_continues() {
let policy = MemoryBudgetPolicy::defaults();
let decision = check_in_build_growth(mb(100), mb(140), &policy);
let InBuildGrowthDecision::Continue {
observed_bytes,
allowed_bytes,
} = decision
else {
panic!("expected Continue");
};
assert_eq!(observed_bytes, mb(140));
assert_eq!(allowed_bytes, mb(150));
}
#[test]
fn in_build_growth_past_tripwire_aborts_with_unexpected_growth_code() {
let policy = MemoryBudgetPolicy::defaults();
let decision = check_in_build_growth(mb(100), mb(160), &policy);
let InBuildGrowthDecision::Abort(refusal) = decision else {
panic!("expected Abort");
};
assert_eq!(refusal.code, UNEXPECTED_GROWTH_CODE);
assert_eq!(refusal.observed_bytes, mb(160));
assert_eq!(refusal.limit_bytes, mb(150));
}
#[test]
fn algorithm_admission_refuses_per_algorithm_cap_before_combined_check() {
let policy = MemoryBudgetPolicy::defaults();
let decision = check_algorithm_admission(0, mb(150), &policy);
let AlgorithmAdmissionDecision::Refuse(refusal) = decision else {
panic!("expected Refuse");
};
assert_eq!(refusal.code, ALGORITHM_MEMORY_CAP_CODE);
assert_eq!(refusal.observed_bytes, mb(150));
assert_eq!(refusal.limit_bytes, mb(100));
}
#[test]
fn algorithm_admission_refuses_combined_pressure_when_total_breaches_cap() {
let policy = MemoryBudgetPolicy::defaults();
let decision = check_algorithm_admission(mb(200), mb(80), &policy);
let AlgorithmAdmissionDecision::Refuse(refusal) = decision else {
panic!("expected Refuse");
};
assert_eq!(refusal.code, MEMORY_PRESSURE_CODE);
assert_eq!(refusal.observed_bytes, mb(280));
assert_eq!(refusal.limit_bytes, mb(250));
}
#[test]
fn algorithm_admission_admits_when_within_all_caps() {
let policy = MemoryBudgetPolicy::defaults();
let decision = check_algorithm_admission(mb(80), mb(50), &policy);
let AlgorithmAdmissionDecision::Admit { combined_bytes } = decision else {
panic!("expected Admit");
};
assert_eq!(combined_bytes, mb(130));
}
#[test]
fn one_hundred_k_memory_estimate_triggers_graceful_refusal_not_oom() {
let policy = MemoryBudgetPolicy::defaults();
let estimate = estimate_snapshot_bytes(100_000, 500_000);
let SnapshotAdmissionDecision::Admit { .. } = check_snapshot_admission(estimate, &policy)
else {
panic!("light density should still admit");
};
let dense_estimate = estimate_snapshot_bytes(100_000, 3_000_000);
let SnapshotAdmissionDecision::Refuse(refusal) =
check_snapshot_admission(dense_estimate, &policy)
else {
panic!("dense 100k fixture must refuse");
};
assert_eq!(refusal.code, LARGE_GRAPH_UNCACHED_CODE);
}
#[test]
fn every_refusal_variant_carries_a_non_empty_repair_hint() {
let policy = MemoryBudgetPolicy::defaults();
let snapshot_refusal = match check_snapshot_admission(mb(1000), &policy) {
SnapshotAdmissionDecision::Refuse(r) => r,
other => panic!("expected Refuse, got {other:?}"),
};
let growth_refusal = match check_in_build_growth(mb(100), mb(500), &policy) {
InBuildGrowthDecision::Abort(r) => r,
other => panic!("expected Abort, got {other:?}"),
};
let algo_cap_refusal = match check_algorithm_admission(0, mb(500), &policy) {
AlgorithmAdmissionDecision::Refuse(r) => r,
other => panic!("expected Refuse, got {other:?}"),
};
let pressure_refusal = match check_algorithm_admission(mb(240), mb(20), &policy) {
AlgorithmAdmissionDecision::Refuse(r) => r,
other => panic!("expected Refuse, got {other:?}"),
};
for refusal in [
snapshot_refusal,
growth_refusal,
algo_cap_refusal,
pressure_refusal,
] {
assert!(!refusal.repair.is_empty(), "repair hint must be present");
assert!(!refusal.message.is_empty(), "message must be present");
assert!(refusal.observed_bytes > 0);
assert!(refusal.limit_bytes > 0);
}
}
#[test]
fn decisions_serialize_byte_stable_across_runs() {
let policy = MemoryBudgetPolicy::defaults();
let first = check_snapshot_admission(mb(50), &policy);
let second = check_snapshot_admission(mb(50), &policy);
let json_first = serde_json::to_string(&first).expect("serialize");
let json_second = serde_json::to_string(&second).expect("serialize");
assert_eq!(json_first, json_second);
}
}