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(target_os = "windows")]
317fn physical_ram_bytes() -> u64 {
318 use windows_sys::Win32::System::SystemInformation::{GlobalMemoryStatusEx, MEMORYSTATUSEX};
319 let mut status = MEMORYSTATUSEX {
320 dwLength: u32::try_from(std::mem::size_of::<MEMORYSTATUSEX>())
321 .expect("Windows memory status size"),
322 ..Default::default()
323 };
324 if unsafe { GlobalMemoryStatusEx(&mut status) } == 0 {
327 return 0;
328 }
329 status.ullTotalPhys
330}
331
332#[cfg(not(any(target_os = "linux", target_os = "windows")))]
333fn physical_ram_bytes() -> u64 {
334 0
335}
336
337#[cfg(test)]
338mod tests {
339 use super::*;
340
341 const GIB: u64 = 1 << 30;
342
343 fn inputs() -> MemoryBudgetInputs {
344 MemoryBudgetInputs {
345 profile: Some(DeployProfile::Embedded),
346 physical_ram_bytes: 16 * GIB,
347 ..MemoryBudgetInputs::default()
348 }
349 }
350
351 #[test]
352 fn explicit_configuration_wins_over_every_other_tier() {
353 let resolved = resolve(&MemoryBudgetInputs {
354 configured: Some("512MiB".to_string()),
355 profile: Some(DeployProfile::Serverless),
356 cgroup_v2_memory_max: Some("1073741824".to_string()),
357 cgroup_v1_memory_limit: Some("2147483648".to_string()),
358 ..inputs()
359 })
360 .expect("configured budget resolves");
361
362 assert_eq!(resolved.source, MemoryBudgetSource::Config);
363 assert_eq!(resolved.resolved_bytes, 512 * 1024 * 1024);
364 }
365
366 #[test]
367 fn configured_budget_accepts_bare_bytes_and_unit_suffixes() {
368 let cases = [
369 ("536870912", 536_870_912),
370 (" 536870912 ", 536_870_912),
371 ("1024b", 1024),
372 ("2KiB", 2048),
373 ("512MiB", 512 * 1024 * 1024),
374 ("2GiB", 2 * GIB),
375 ("1TiB", 1 << 40),
376 ("2GB", 2_000_000_000),
377 ("500mb", 500_000_000),
378 ];
379
380 for (raw, expected) in cases {
381 assert_eq!(parse_budget_bytes(raw), Some(expected), "parsing {raw}");
382 }
383 }
384
385 #[test]
386 fn nonsensical_configured_budget_is_rejected_not_replaced() {
387 for raw in [
388 "0",
389 "0MiB",
390 "-1",
391 "unlimited",
392 "max",
393 "",
394 " ",
395 "12 apples",
396 "MiB",
397 "1.5GiB",
398 ] {
399 let err = resolve(&MemoryBudgetInputs {
400 configured: Some(raw.to_string()),
401 ..inputs()
402 })
403 .expect_err("nonsensical configured budget must fail boot");
404
405 let message = err.to_string();
406 assert!(
407 message.contains("expected a positive byte count"),
408 "{message}"
409 );
410 assert!(
411 message.contains("512MiB"),
412 "error names the valid form: {message}"
413 );
414 assert!(message.contains("no unlimited mode"), "{message}");
415 }
416 }
417
418 #[test]
419 fn serverless_profile_ships_a_strict_default() {
420 let resolved = resolve(&MemoryBudgetInputs {
421 profile: Some(DeployProfile::Serverless),
422 cgroup_v2_memory_max: Some("34359738368".to_string()),
423 ..inputs()
424 })
425 .expect("profile default resolves");
426
427 assert_eq!(resolved.source, MemoryBudgetSource::ProfileDefault);
428 assert_eq!(resolved.resolved_bytes, SERVERLESS_PROFILE_BUDGET_BYTES);
429 }
430
431 #[test]
432 fn detection_profiles_fall_through_to_the_cgroup() {
433 for profile in [
434 DeployProfile::Embedded,
435 DeployProfile::PrimaryReplica,
436 DeployProfile::Cluster,
437 ] {
438 let resolved = resolve(&MemoryBudgetInputs {
439 profile: Some(profile),
440 cgroup_v2_memory_max: Some("1073741824\n".to_string()),
441 ..inputs()
442 })
443 .expect("cgroup budget resolves");
444
445 assert_eq!(resolved.source, MemoryBudgetSource::CgroupV2);
446 assert_eq!(resolved.resolved_bytes, GIB);
447 }
448 }
449
450 #[test]
451 fn cgroup_v2_max_falls_through_to_v1_then_physical() {
452 let to_v1 = resolve(&MemoryBudgetInputs {
453 cgroup_v2_memory_max: Some("max\n".to_string()),
454 cgroup_v1_memory_limit: Some("2147483648".to_string()),
455 ..inputs()
456 })
457 .expect("v1 fallback resolves");
458 assert_eq!(to_v1.source, MemoryBudgetSource::CgroupV1);
459 assert_eq!(to_v1.resolved_bytes, 2 * GIB);
460
461 let to_physical = resolve(&MemoryBudgetInputs {
462 cgroup_v2_memory_max: Some("MAX".to_string()),
463 ..inputs()
464 })
465 .expect("physical fallback resolves");
466 assert_eq!(to_physical.source, MemoryBudgetSource::PhysicalFraction);
467 }
468
469 #[test]
470 fn cgroup_v1_unlimited_sentinel_falls_through_to_physical() {
471 let resolved = resolve(&MemoryBudgetInputs {
472 cgroup_v1_memory_limit: Some(CGROUP_V1_UNLIMITED_SENTINEL.to_string()),
473 ..inputs()
474 })
475 .expect("sentinel falls through");
476
477 assert_eq!(resolved.source, MemoryBudgetSource::PhysicalFraction);
478 assert_eq!(resolved.resolved_bytes, 16 * GIB / 100 * 70);
479 }
480
481 #[test]
482 fn physical_fraction_is_the_last_tier_and_never_zero() {
483 let detected = resolve(&inputs()).expect("physical fraction resolves");
484 assert_eq!(detected.source, MemoryBudgetSource::PhysicalFraction);
485 assert_eq!(detected.resolved_bytes, 16 * GIB / 100 * 70);
486
487 let undetectable = resolve(&MemoryBudgetInputs {
489 physical_ram_bytes: 0,
490 ..inputs()
491 })
492 .expect("undetectable physical RAM resolves");
493 assert_eq!(undetectable.source, MemoryBudgetSource::PhysicalFraction);
494 assert_eq!(undetectable.resolved_bytes, MINIMUM_DETECTED_BUDGET_BYTES);
495 }
496
497 #[test]
498 fn every_resolved_budget_is_strictly_positive() {
499 let profiles = [
500 None,
501 Some(DeployProfile::Embedded),
502 Some(DeployProfile::Serverless),
503 Some(DeployProfile::PrimaryReplica),
504 Some(DeployProfile::Cluster),
505 ];
506 let cgroup_v2 = [None, Some("max"), Some("0"), Some("1"), Some("garbage")];
507 let cgroup_v1 = [None, Some("0"), Some("1"), Some("9223372036854771712")];
508 let physical = [0, 1, 1024, 16 * GIB, u64::MAX];
509 let configured = [None, Some("1"), Some("1TiB")];
510
511 for profile in profiles {
512 for v2 in cgroup_v2 {
513 for v1 in cgroup_v1 {
514 for ram in physical {
515 for config in configured {
516 let resolved = resolve(&MemoryBudgetInputs {
517 configured: config.map(str::to_string),
518 profile,
519 cgroup_v2_memory_max: v2.map(str::to_string),
520 cgroup_v1_memory_limit: v1.map(str::to_string),
521 physical_ram_bytes: ram,
522 })
523 .expect("valid inputs always resolve");
524 assert!(resolved.resolved_bytes > 0);
525 }
526 }
527 }
528 }
529 }
530 }
531
532 #[test]
533 fn source_labels_are_the_documented_tier_names() {
534 assert_eq!(MemoryBudgetSource::Config.as_str(), "config");
535 assert_eq!(
536 MemoryBudgetSource::ProfileDefault.as_str(),
537 "profile-default"
538 );
539 assert_eq!(MemoryBudgetSource::CgroupV2.as_str(), "cgroup-v2");
540 assert_eq!(MemoryBudgetSource::CgroupV1.as_str(), "cgroup-v1");
541 assert_eq!(
542 MemoryBudgetSource::PhysicalFraction.as_str(),
543 "physical-fraction"
544 );
545 }
546}