1use std::sync::Once;
25
26use super::profile::DeployProfile;
27
28pub const MEMORY_BUDGET_ENV: &str = "REDDB_MEMORY_BUDGET";
32
33pub const CGROUP_V2_MEMORY_MAX_PATH: &str = "/sys/fs/cgroup/memory.max";
36
37pub const CGROUP_V1_MEMORY_LIMIT_PATH: &str = "/sys/fs/cgroup/memory/memory.limit_in_bytes";
40
41pub const PHYSICAL_RAM_BUDGET_PERCENT: u64 = 70;
50
51pub const SERVERLESS_PROFILE_BUDGET_BYTES: u64 = 256 * 1024 * 1024;
56
57pub const MINIMUM_DETECTED_BUDGET_BYTES: u64 = 64 * 1024 * 1024;
65
66const CGROUP_V1_UNLIMITED_SENTINEL: u64 = 0x7FFF_FFFF_FFFF_F000;
71
72#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
74pub enum MemoryBudgetSource {
75 Config,
77 ProfileDefault,
79 CgroupV2,
81 CgroupV1,
83 PhysicalFraction,
85}
86
87impl MemoryBudgetSource {
88 pub const fn as_str(self) -> &'static str {
90 match self {
91 Self::Config => "config",
92 Self::ProfileDefault => "profile-default",
93 Self::CgroupV2 => "cgroup-v2",
94 Self::CgroupV1 => "cgroup-v1",
95 Self::PhysicalFraction => "physical-fraction",
96 }
97 }
98}
99
100#[derive(Debug, Clone, Copy, PartialEq, Eq)]
102pub struct MemoryBudget {
103 pub resolved_bytes: u64,
104 pub source: MemoryBudgetSource,
105}
106
107#[derive(Debug, Clone, PartialEq, Eq)]
110pub struct InvalidMemoryBudget {
111 pub raw: String,
112}
113
114impl std::fmt::Display for InvalidMemoryBudget {
115 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116 write!(
117 f,
118 "invalid memory budget `{}`: expected a positive byte count with an optional \
119 unit suffix (`536870912`, `512MiB`, `2GiB`, `2GB`); RedDB has no unlimited \
120 mode — omit the setting to use the detected default",
121 self.raw
122 )
123 }
124}
125
126impl std::error::Error for InvalidMemoryBudget {}
127
128#[derive(Debug, Clone, Default, PartialEq, Eq)]
131pub struct MemoryBudgetInputs {
132 pub configured: Option<String>,
134 pub profile: Option<DeployProfile>,
136 pub cgroup_v2_memory_max: Option<String>,
138 pub cgroup_v1_memory_limit: Option<String>,
140 pub physical_ram_bytes: u64,
142}
143
144pub fn resolve(inputs: &MemoryBudgetInputs) -> Result<MemoryBudget, InvalidMemoryBudget> {
150 if let Some(raw) = inputs.configured.as_deref() {
151 let resolved_bytes = parse_budget_bytes(raw).ok_or_else(|| InvalidMemoryBudget {
152 raw: raw.to_string(),
153 })?;
154 return Ok(MemoryBudget {
155 resolved_bytes,
156 source: MemoryBudgetSource::Config,
157 });
158 }
159
160 if let Some(resolved_bytes) = inputs.profile.and_then(profile_default_bytes) {
161 return Ok(MemoryBudget {
162 resolved_bytes,
163 source: MemoryBudgetSource::ProfileDefault,
164 });
165 }
166
167 if let Some(resolved_bytes) = inputs
168 .cgroup_v2_memory_max
169 .as_deref()
170 .and_then(parse_cgroup_v2_limit)
171 {
172 return Ok(MemoryBudget {
173 resolved_bytes,
174 source: MemoryBudgetSource::CgroupV2,
175 });
176 }
177
178 if let Some(resolved_bytes) = inputs
179 .cgroup_v1_memory_limit
180 .as_deref()
181 .and_then(parse_cgroup_v1_limit)
182 {
183 return Ok(MemoryBudget {
184 resolved_bytes,
185 source: MemoryBudgetSource::CgroupV1,
186 });
187 }
188
189 Ok(MemoryBudget {
190 resolved_bytes: physical_fraction_bytes(inputs.physical_ram_bytes),
191 source: MemoryBudgetSource::PhysicalFraction,
192 })
193}
194
195pub fn host_inputs(profile: DeployProfile, configured_bytes: Option<u64>) -> MemoryBudgetInputs {
198 let configured = configured_bytes
201 .map(|bytes| bytes.to_string())
202 .or_else(|| std::env::var(MEMORY_BUDGET_ENV).ok())
203 .filter(|raw| !raw.trim().is_empty());
204
205 let (cgroup_v2, cgroup_v1) = crate::utils::cgroup_memory::limits();
206 MemoryBudgetInputs {
207 configured,
208 profile: Some(profile),
209 cgroup_v2_memory_max: cgroup_v2.map(|bytes| bytes.to_string()),
210 cgroup_v1_memory_limit: cgroup_v1.map(|bytes| bytes.to_string()),
211 physical_ram_bytes: physical_ram_bytes(),
212 }
213}
214
215pub fn resolve_for_boot(
218 profile: DeployProfile,
219 configured_bytes: Option<u64>,
220) -> Result<MemoryBudget, InvalidMemoryBudget> {
221 resolve(&host_inputs(profile, configured_bytes))
222}
223
224pub fn log_resolved_once(budget: &MemoryBudget) {
227 static LOGGED: Once = Once::new();
228 LOGGED.call_once(|| {
229 tracing::info!(
230 resolved_bytes = budget.resolved_bytes,
231 source = budget.source.as_str(),
232 "memory budget resolved"
233 );
234 });
235}
236
237fn profile_default_bytes(profile: DeployProfile) -> Option<u64> {
239 match profile {
240 DeployProfile::Serverless => Some(SERVERLESS_PROFILE_BUDGET_BYTES),
241 DeployProfile::Embedded | DeployProfile::PrimaryReplica | DeployProfile::Cluster => None,
242 }
243}
244
245fn parse_budget_bytes(raw: &str) -> Option<u64> {
249 let trimmed = raw.trim();
250 let digits_end = trimmed
251 .char_indices()
252 .find(|(_, ch)| !ch.is_ascii_digit())
253 .map_or(trimmed.len(), |(idx, _)| idx);
254 if digits_end == 0 {
255 return None;
256 }
257
258 let value = trimmed[..digits_end].parse::<u64>().ok()?;
259 let multiplier = match trimmed[digits_end..].trim().to_ascii_lowercase().as_str() {
260 "" | "b" => 1,
261 "kb" => 1_000,
262 "mb" => 1_000_000,
263 "gb" => 1_000_000_000,
264 "tb" => 1_000_000_000_000,
265 "kib" => 1 << 10,
266 "mib" => 1 << 20,
267 "gib" => 1 << 30,
268 "tib" => 1 << 40,
269 _ => return None,
270 };
271
272 let bytes = value.checked_mul(multiplier)?;
273 (bytes > 0).then_some(bytes)
274}
275
276fn parse_cgroup_v2_limit(raw: &str) -> Option<u64> {
278 let trimmed = raw.trim();
279 if trimmed.eq_ignore_ascii_case("max") {
280 return None;
281 }
282 trimmed.parse::<u64>().ok().filter(|bytes| *bytes > 0)
283}
284
285fn parse_cgroup_v1_limit(raw: &str) -> Option<u64> {
287 raw.trim()
288 .parse::<u64>()
289 .ok()
290 .filter(|bytes| *bytes > 0 && *bytes < CGROUP_V1_UNLIMITED_SENTINEL)
291}
292
293fn physical_fraction_bytes(physical_ram_bytes: u64) -> u64 {
295 (physical_ram_bytes / 100)
298 .saturating_mul(PHYSICAL_RAM_BUDGET_PERCENT)
299 .max(MINIMUM_DETECTED_BUDGET_BYTES)
300}
301
302#[cfg(target_os = "linux")]
303fn physical_ram_bytes() -> u64 {
304 std::fs::read_to_string("/proc/meminfo")
305 .ok()
306 .and_then(|meminfo| {
307 meminfo
308 .lines()
309 .find(|line| line.starts_with("MemTotal:"))
310 .and_then(|line| line.split_whitespace().nth(1))
311 .and_then(|kb| kb.parse::<u64>().ok())
312 })
313 .map_or(0, |kb| kb.saturating_mul(1024))
314}
315
316#[cfg(not(target_os = "linux"))]
317fn physical_ram_bytes() -> u64 {
318 0
319}
320
321#[cfg(test)]
322mod tests {
323 use super::*;
324
325 const GIB: u64 = 1 << 30;
326
327 fn inputs() -> MemoryBudgetInputs {
328 MemoryBudgetInputs {
329 profile: Some(DeployProfile::Embedded),
330 physical_ram_bytes: 16 * GIB,
331 ..MemoryBudgetInputs::default()
332 }
333 }
334
335 #[test]
336 fn explicit_configuration_wins_over_every_other_tier() {
337 let resolved = resolve(&MemoryBudgetInputs {
338 configured: Some("512MiB".to_string()),
339 profile: Some(DeployProfile::Serverless),
340 cgroup_v2_memory_max: Some("1073741824".to_string()),
341 cgroup_v1_memory_limit: Some("2147483648".to_string()),
342 ..inputs()
343 })
344 .expect("configured budget resolves");
345
346 assert_eq!(resolved.source, MemoryBudgetSource::Config);
347 assert_eq!(resolved.resolved_bytes, 512 * 1024 * 1024);
348 }
349
350 #[test]
351 fn configured_budget_accepts_bare_bytes_and_unit_suffixes() {
352 let cases = [
353 ("536870912", 536_870_912),
354 (" 536870912 ", 536_870_912),
355 ("1024b", 1024),
356 ("2KiB", 2048),
357 ("512MiB", 512 * 1024 * 1024),
358 ("2GiB", 2 * GIB),
359 ("1TiB", 1 << 40),
360 ("2GB", 2_000_000_000),
361 ("500mb", 500_000_000),
362 ];
363
364 for (raw, expected) in cases {
365 assert_eq!(parse_budget_bytes(raw), Some(expected), "parsing {raw}");
366 }
367 }
368
369 #[test]
370 fn nonsensical_configured_budget_is_rejected_not_replaced() {
371 for raw in [
372 "0",
373 "0MiB",
374 "-1",
375 "unlimited",
376 "max",
377 "",
378 " ",
379 "12 apples",
380 "MiB",
381 "1.5GiB",
382 ] {
383 let err = resolve(&MemoryBudgetInputs {
384 configured: Some(raw.to_string()),
385 ..inputs()
386 })
387 .expect_err("nonsensical configured budget must fail boot");
388
389 let message = err.to_string();
390 assert!(
391 message.contains("expected a positive byte count"),
392 "{message}"
393 );
394 assert!(
395 message.contains("512MiB"),
396 "error names the valid form: {message}"
397 );
398 assert!(message.contains("no unlimited mode"), "{message}");
399 }
400 }
401
402 #[test]
403 fn serverless_profile_ships_a_strict_default() {
404 let resolved = resolve(&MemoryBudgetInputs {
405 profile: Some(DeployProfile::Serverless),
406 cgroup_v2_memory_max: Some("34359738368".to_string()),
407 ..inputs()
408 })
409 .expect("profile default resolves");
410
411 assert_eq!(resolved.source, MemoryBudgetSource::ProfileDefault);
412 assert_eq!(resolved.resolved_bytes, SERVERLESS_PROFILE_BUDGET_BYTES);
413 }
414
415 #[test]
416 fn detection_profiles_fall_through_to_the_cgroup() {
417 for profile in [
418 DeployProfile::Embedded,
419 DeployProfile::PrimaryReplica,
420 DeployProfile::Cluster,
421 ] {
422 let resolved = resolve(&MemoryBudgetInputs {
423 profile: Some(profile),
424 cgroup_v2_memory_max: Some("1073741824\n".to_string()),
425 ..inputs()
426 })
427 .expect("cgroup budget resolves");
428
429 assert_eq!(resolved.source, MemoryBudgetSource::CgroupV2);
430 assert_eq!(resolved.resolved_bytes, GIB);
431 }
432 }
433
434 #[test]
435 fn cgroup_v2_max_falls_through_to_v1_then_physical() {
436 let to_v1 = resolve(&MemoryBudgetInputs {
437 cgroup_v2_memory_max: Some("max\n".to_string()),
438 cgroup_v1_memory_limit: Some("2147483648".to_string()),
439 ..inputs()
440 })
441 .expect("v1 fallback resolves");
442 assert_eq!(to_v1.source, MemoryBudgetSource::CgroupV1);
443 assert_eq!(to_v1.resolved_bytes, 2 * GIB);
444
445 let to_physical = resolve(&MemoryBudgetInputs {
446 cgroup_v2_memory_max: Some("MAX".to_string()),
447 ..inputs()
448 })
449 .expect("physical fallback resolves");
450 assert_eq!(to_physical.source, MemoryBudgetSource::PhysicalFraction);
451 }
452
453 #[test]
454 fn cgroup_v1_unlimited_sentinel_falls_through_to_physical() {
455 let resolved = resolve(&MemoryBudgetInputs {
456 cgroup_v1_memory_limit: Some(CGROUP_V1_UNLIMITED_SENTINEL.to_string()),
457 ..inputs()
458 })
459 .expect("sentinel falls through");
460
461 assert_eq!(resolved.source, MemoryBudgetSource::PhysicalFraction);
462 assert_eq!(resolved.resolved_bytes, 16 * GIB / 100 * 70);
463 }
464
465 #[test]
466 fn physical_fraction_is_the_last_tier_and_never_zero() {
467 let detected = resolve(&inputs()).expect("physical fraction resolves");
468 assert_eq!(detected.source, MemoryBudgetSource::PhysicalFraction);
469 assert_eq!(detected.resolved_bytes, 16 * GIB / 100 * 70);
470
471 let undetectable = resolve(&MemoryBudgetInputs {
473 physical_ram_bytes: 0,
474 ..inputs()
475 })
476 .expect("undetectable physical RAM resolves");
477 assert_eq!(undetectable.source, MemoryBudgetSource::PhysicalFraction);
478 assert_eq!(undetectable.resolved_bytes, MINIMUM_DETECTED_BUDGET_BYTES);
479 }
480
481 #[test]
482 fn every_resolved_budget_is_strictly_positive() {
483 let profiles = [
484 None,
485 Some(DeployProfile::Embedded),
486 Some(DeployProfile::Serverless),
487 Some(DeployProfile::PrimaryReplica),
488 Some(DeployProfile::Cluster),
489 ];
490 let cgroup_v2 = [None, Some("max"), Some("0"), Some("1"), Some("garbage")];
491 let cgroup_v1 = [None, Some("0"), Some("1"), Some("9223372036854771712")];
492 let physical = [0, 1, 1024, 16 * GIB, u64::MAX];
493 let configured = [None, Some("1"), Some("1TiB")];
494
495 for profile in profiles {
496 for v2 in cgroup_v2 {
497 for v1 in cgroup_v1 {
498 for ram in physical {
499 for config in configured {
500 let resolved = resolve(&MemoryBudgetInputs {
501 configured: config.map(str::to_string),
502 profile,
503 cgroup_v2_memory_max: v2.map(str::to_string),
504 cgroup_v1_memory_limit: v1.map(str::to_string),
505 physical_ram_bytes: ram,
506 })
507 .expect("valid inputs always resolve");
508 assert!(resolved.resolved_bytes > 0);
509 }
510 }
511 }
512 }
513 }
514 }
515
516 #[test]
517 fn source_labels_are_the_documented_tier_names() {
518 assert_eq!(MemoryBudgetSource::Config.as_str(), "config");
519 assert_eq!(
520 MemoryBudgetSource::ProfileDefault.as_str(),
521 "profile-default"
522 );
523 assert_eq!(MemoryBudgetSource::CgroupV2.as_str(), "cgroup-v2");
524 assert_eq!(MemoryBudgetSource::CgroupV1.as_str(), "cgroup-v1");
525 assert_eq!(
526 MemoryBudgetSource::PhysicalFraction.as_str(),
527 "physical-fraction"
528 );
529 }
530}