1use std::sync::atomic::{AtomicU64, Ordering};
27use std::sync::Once;
28
29use super::engine::page_cache::MIN_CACHE_CAPACITY;
30use super::memory_budget::MemoryBudget;
31use super::profile::DeployProfile;
32
33pub const PAGE_CACHE_PAGE_SIZE_BYTES: u64 = reddb_file::PAGED_PAGE_SIZE as u64;
36
37pub const BASIS_POINTS_PER_WHOLE: u32 = 10_000;
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
49pub enum MemoryPool {
50 PageCache,
52 BlobCacheL1,
54 SegmentArena,
56 IndexMemory,
58 WalBuffers,
60}
61
62pub const MEMORY_POOLS: [MemoryPool; MEMORY_POOL_COUNT] = [
64 MemoryPool::PageCache,
65 MemoryPool::BlobCacheL1,
66 MemoryPool::SegmentArena,
67 MemoryPool::IndexMemory,
68 MemoryPool::WalBuffers,
69];
70
71pub const MEMORY_POOL_COUNT: usize = 5;
73
74impl MemoryPool {
75 pub const fn as_str(self) -> &'static str {
77 match self {
78 Self::PageCache => "page_cache",
79 Self::BlobCacheL1 => "blob_cache_l1",
80 Self::SegmentArena => "segment_arena",
81 Self::IndexMemory => "index_memory",
82 Self::WalBuffers => "wal_buffers",
83 }
84 }
85
86 pub const fn index(self) -> usize {
88 match self {
89 Self::PageCache => 0,
90 Self::BlobCacheL1 => 1,
91 Self::SegmentArena => 2,
92 Self::IndexMemory => 3,
93 Self::WalBuffers => 4,
94 }
95 }
96}
97
98#[derive(Debug, Clone, Copy, PartialEq, Eq)]
104pub struct BudgetSharePolicy {
105 basis_points: [u32; MEMORY_POOL_COUNT],
106}
107
108const DEFAULT_POLICY: BudgetSharePolicy = BudgetSharePolicy {
113 basis_points: [
114 2_000, 1_500, 4_000, 1_000, 500, ],
120};
121
122const SERVERLESS_POLICY: BudgetSharePolicy = BudgetSharePolicy {
126 basis_points: [
127 1_500, 1_000, 4_000, 1_000, 500, ],
133};
134
135impl BudgetSharePolicy {
136 pub const fn for_profile(profile: DeployProfile) -> Self {
138 match profile {
139 DeployProfile::Serverless => SERVERLESS_POLICY,
140 DeployProfile::Embedded | DeployProfile::PrimaryReplica | DeployProfile::Cluster => {
141 DEFAULT_POLICY
142 }
143 }
144 }
145
146 pub const fn basis_points(&self, pool: MemoryPool) -> u32 {
148 self.basis_points[pool.index()]
149 }
150
151 pub fn total_basis_points(&self) -> u32 {
153 self.basis_points.iter().sum()
154 }
155
156 pub fn reserve_basis_points(&self) -> u32 {
159 BASIS_POINTS_PER_WHOLE - self.total_basis_points()
160 }
161}
162
163#[derive(Debug, Clone, Copy, PartialEq, Eq)]
165pub struct BudgetShares {
166 budget_bytes: u64,
167 policy: BudgetSharePolicy,
168 share_bytes: [u64; MEMORY_POOL_COUNT],
169}
170
171impl BudgetShares {
172 pub fn resolve(budget: MemoryBudget, profile: DeployProfile) -> Self {
180 let policy = BudgetSharePolicy::for_profile(profile);
181
182 assert!(
187 policy.total_basis_points() > 0,
188 "invariant: budget share policy must claim some of the budget"
189 );
190 assert!(
191 policy.total_basis_points() <= BASIS_POINTS_PER_WHOLE,
192 "invariant: Σ(share fractions) = {} bp exceeds the whole budget ({BASIS_POINTS_PER_WHOLE} bp)",
193 policy.total_basis_points()
194 );
195
196 let budget_bytes = budget.resolved_bytes;
197 assert!(
198 budget_bytes > 0,
199 "invariant: the resolved budget is strictly positive (ADR 0073 §1: no unlimited mode)"
200 );
201
202 let per_basis_point = budget_bytes / u64::from(BASIS_POINTS_PER_WHOLE);
203 let mut share_bytes = [0_u64; MEMORY_POOL_COUNT];
204 for pool in MEMORY_POOLS {
205 share_bytes[pool.index()] = per_basis_point * u64::from(policy.basis_points(pool));
206 }
207
208 let shares = Self {
209 budget_bytes,
210 policy,
211 share_bytes,
212 };
213
214 assert!(
218 shares.total_share_bytes() <= budget_bytes,
219 "invariant: Σ(shares) = {} exceeds budget {budget_bytes}",
220 shares.total_share_bytes()
221 );
222 assert!(
223 shares.total_share_bytes() > 0 || budget_bytes < u64::from(BASIS_POINTS_PER_WHOLE),
224 "invariant: a budget of {budget_bytes} bytes must reach the pools"
225 );
226
227 shares
228 }
229
230 pub fn budget_bytes(&self) -> u64 {
232 self.budget_bytes
233 }
234
235 pub fn policy(&self) -> BudgetSharePolicy {
237 self.policy
238 }
239
240 pub fn share_bytes(&self, pool: MemoryPool) -> u64 {
242 self.share_bytes[pool.index()]
243 }
244
245 pub fn total_share_bytes(&self) -> u64 {
247 self.share_bytes.iter().sum()
248 }
249
250 pub fn page_cache_slots(&self) -> usize {
257 let slots = self.share_bytes(MemoryPool::PageCache) / PAGE_CACHE_PAGE_SIZE_BYTES;
258 usize::try_from(slots)
259 .unwrap_or(usize::MAX)
260 .max(MIN_CACHE_CAPACITY)
261 }
262
263 pub fn blob_cache_l1_bytes(&self) -> usize {
265 usize::try_from(self.share_bytes(MemoryPool::BlobCacheL1)).unwrap_or(usize::MAX)
266 }
267
268 pub fn log_once(&self) {
272 static LOGGED: Once = Once::new();
273 LOGGED.call_once(|| {
274 for pool in MEMORY_POOLS {
275 tracing::info!(
276 pool = pool.as_str(),
277 share_bytes = self.share_bytes(pool),
278 basis_points = self.policy.basis_points(pool),
279 budget_bytes = self.budget_bytes,
280 "memory budget share assigned"
281 );
282 }
283 tracing::info!(
284 total_share_bytes = self.total_share_bytes(),
285 reserve_basis_points = self.policy.reserve_basis_points(),
286 budget_bytes = self.budget_bytes,
287 "memory budget shares resolved"
288 );
289 });
290 }
291}
292
293#[derive(Debug, Clone, Copy, PartialEq, Eq)]
295pub struct PoolUsage {
296 pub pool: MemoryPool,
297 pub share_bytes: u64,
298 pub used_bytes: u64,
299}
300
301#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
303pub struct MemoryEnforcementSnapshot {
304 pub admissions_denied: u64,
305 pub pressure_reclamations_triggered: u64,
306 pub pressure_bytes_reclaimed: u64,
307 pub high_water_bytes: u64,
308}
309
310#[derive(Debug)]
317pub struct MemoryAccounting {
318 budget: MemoryBudget,
319 shares: BudgetShares,
320 used_bytes: [AtomicU64; MEMORY_POOL_COUNT],
321 admissions_denied: AtomicU64,
322 pressure_reclamations_triggered: AtomicU64,
323 pressure_bytes_reclaimed: AtomicU64,
324 high_water_bytes: AtomicU64,
325}
326
327impl MemoryAccounting {
328 pub fn new(budget: MemoryBudget, profile: DeployProfile) -> Self {
330 Self::from_shares(budget, BudgetShares::resolve(budget, profile))
331 }
332
333 pub fn from_shares(budget: MemoryBudget, shares: BudgetShares) -> Self {
335 assert_eq!(
336 budget.resolved_bytes, shares.budget_bytes,
337 "invariant: accounting shares must divide this process's budget"
338 );
339 Self {
340 budget,
341 shares,
342 used_bytes: std::array::from_fn(|_| AtomicU64::new(0)),
343 admissions_denied: AtomicU64::new(0),
344 pressure_reclamations_triggered: AtomicU64::new(0),
345 pressure_bytes_reclaimed: AtomicU64::new(0),
346 high_water_bytes: AtomicU64::new(0),
347 }
348 }
349
350 pub fn budget(&self) -> MemoryBudget {
352 self.budget
353 }
354
355 pub fn shares(&self) -> BudgetShares {
357 self.shares
358 }
359
360 pub fn report(&self, pool: MemoryPool, bytes: u64) {
363 self.used_bytes[pool.index()].store(bytes, Ordering::Relaxed);
364 }
365
366 pub fn charge(&self, pool: MemoryPool, bytes: u64) {
368 self.used_bytes[pool.index()].fetch_add(bytes, Ordering::Relaxed);
369 }
370
371 pub fn release(&self, pool: MemoryPool, bytes: u64) {
377 let slot = &self.used_bytes[pool.index()];
378 let _ = slot.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
379 Some(used.saturating_sub(bytes))
380 });
381 }
382
383 pub fn used_bytes(&self, pool: MemoryPool) -> u64 {
385 self.used_bytes[pool.index()].load(Ordering::Relaxed)
386 }
387
388 pub fn share_bytes(&self, pool: MemoryPool) -> u64 {
390 self.shares.share_bytes(pool)
391 }
392
393 pub fn total_used_bytes(&self) -> u64 {
395 MEMORY_POOLS
396 .iter()
397 .map(|pool| self.used_bytes(*pool))
398 .fold(0, u64::saturating_add)
399 }
400
401 pub fn total_share_bytes(&self) -> u64 {
403 self.shares.total_share_bytes()
404 }
405
406 pub fn snapshot(&self) -> [PoolUsage; MEMORY_POOL_COUNT] {
408 MEMORY_POOLS.map(|pool| PoolUsage {
409 pool,
410 share_bytes: self.share_bytes(pool),
411 used_bytes: self.used_bytes(pool),
412 })
413 }
414
415 pub fn observe_total_used(&self, total: u64) {
417 let mut current = self.high_water_bytes.load(Ordering::Relaxed);
418 while total > current {
419 match self.high_water_bytes.compare_exchange_weak(
420 current,
421 total,
422 Ordering::Relaxed,
423 Ordering::Relaxed,
424 ) {
425 Ok(_) => break,
426 Err(observed) => current = observed,
427 }
428 }
429 }
430
431 pub fn record_admission_denied(&self) {
432 self.admissions_denied.fetch_add(1, Ordering::Relaxed);
433 }
434
435 pub fn record_pressure_reclamation(&self, reclaimed_bytes: u64) {
436 self.pressure_reclamations_triggered
437 .fetch_add(1, Ordering::Relaxed);
438 self.pressure_bytes_reclaimed
439 .fetch_add(reclaimed_bytes, Ordering::Relaxed);
440 }
441
442 pub fn enforcement_snapshot(&self) -> MemoryEnforcementSnapshot {
443 MemoryEnforcementSnapshot {
444 admissions_denied: self.admissions_denied.load(Ordering::Relaxed),
445 pressure_reclamations_triggered: self
446 .pressure_reclamations_triggered
447 .load(Ordering::Relaxed),
448 pressure_bytes_reclaimed: self.pressure_bytes_reclaimed.load(Ordering::Relaxed),
449 high_water_bytes: self.high_water_bytes.load(Ordering::Relaxed),
450 }
451 }
452}
453
454#[cfg(test)]
455mod tests {
456 use super::*;
457 use crate::storage::memory_budget::MemoryBudgetSource;
458 use proptest::prelude::*;
459
460 const GIB: u64 = 1 << 30;
461 const MIB: u64 = 1 << 20;
462
463 const PROFILES: [DeployProfile; 4] = [
464 DeployProfile::Embedded,
465 DeployProfile::Serverless,
466 DeployProfile::PrimaryReplica,
467 DeployProfile::Cluster,
468 ];
469
470 fn budget(bytes: u64) -> MemoryBudget {
471 MemoryBudget {
472 resolved_bytes: bytes,
473 source: MemoryBudgetSource::Config,
474 }
475 }
476
477 #[test]
478 fn every_policy_claims_at_most_the_whole_budget() {
479 for profile in PROFILES {
480 let policy = BudgetSharePolicy::for_profile(profile);
481 assert!(policy.total_basis_points() > 0, "{profile:?}");
482 assert!(
483 policy.total_basis_points() <= BASIS_POINTS_PER_WHOLE,
484 "{profile:?} claims {} bp",
485 policy.total_basis_points()
486 );
487 assert!(
488 policy.reserve_basis_points() > 0,
489 "{profile:?} must leave an unpooled reserve for plan caches and slack"
490 );
491 }
492 }
493
494 #[test]
497 fn sum_of_shares_never_exceeds_the_budget() {
498 let budgets = [
499 1,
500 2,
501 9_999,
502 10_000,
503 10_001,
504 64 * MIB,
505 256 * MIB,
506 512 * MIB,
507 GIB,
508 2 * GIB,
509 11 * GIB,
510 1 << 40,
511 1 << 50,
512 u64::MAX / 2,
513 u64::MAX,
514 ];
515
516 for profile in PROFILES {
517 for bytes in budgets {
518 let shares = BudgetShares::resolve(budget(bytes), profile);
519 assert!(
520 shares.total_share_bytes() <= bytes,
521 "{profile:?} budget {bytes}: Σ(shares) = {}",
522 shares.total_share_bytes()
523 );
524 if bytes >= u64::from(BASIS_POINTS_PER_WHOLE) {
526 for pool in MEMORY_POOLS {
527 assert!(
528 shares.share_bytes(pool) > 0,
529 "{profile:?} budget {bytes}: {} starved",
530 pool.as_str()
531 );
532 }
533 }
534 }
535 }
536 }
537
538 #[test]
542 fn sum_of_shares_never_exceeds_the_budget_across_a_dense_sweep() {
543 for profile in PROFILES {
544 for step in 0..2_000_u64 {
545 let bytes = 1 + step * 7_919; let shares = BudgetShares::resolve(budget(bytes), profile);
547 assert!(
548 shares.total_share_bytes() <= bytes,
549 "{profile:?} budget {bytes}: Σ(shares) = {}",
550 shares.total_share_bytes()
551 );
552 }
553 }
554 }
555
556 #[test]
557 fn doubling_the_budget_doubles_the_page_cache_slots_and_the_blob_l1_ceiling() {
558 for profile in PROFILES {
559 let base = BudgetShares::resolve(budget(GIB), profile);
560 let doubled = BudgetShares::resolve(budget(2 * GIB), profile);
561
562 assert_eq!(
563 doubled.page_cache_slots(),
564 base.page_cache_slots() * 2,
565 "{profile:?} page cache slots must scale with the budget"
566 );
567 assert_eq!(
568 doubled.blob_cache_l1_bytes(),
569 base.blob_cache_l1_bytes() * 2,
570 "{profile:?} blob L1 ceiling must scale with the budget"
571 );
572 assert_eq!(
573 doubled.share_bytes(MemoryPool::SegmentArena),
574 base.share_bytes(MemoryPool::SegmentArena) * 2,
575 );
576 }
577 }
578
579 #[test]
580 fn page_cache_slots_are_the_share_divided_by_the_fixed_page_size() {
581 let shares = BudgetShares::resolve(budget(10 * GIB), DeployProfile::Embedded);
582 let expected = shares.share_bytes(MemoryPool::PageCache) / PAGE_CACHE_PAGE_SIZE_BYTES;
583 assert_eq!(shares.page_cache_slots() as u64, expected);
584 assert_eq!(PAGE_CACHE_PAGE_SIZE_BYTES, 16 * 1024);
585 }
586
587 #[test]
588 fn a_budget_too_small_to_fill_a_slot_still_yields_a_usable_page_cache() {
589 let shares = BudgetShares::resolve(budget(1), DeployProfile::Serverless);
590 assert_eq!(shares.total_share_bytes(), 0, "nothing reaches the pools");
591 assert_eq!(shares.page_cache_slots(), MIN_CACHE_CAPACITY);
592 assert_eq!(shares.blob_cache_l1_bytes(), 0);
593 }
594
595 #[test]
596 fn serverless_shrinks_the_caches_relative_to_the_server_profiles() {
597 let serverless = BudgetShares::resolve(budget(GIB), DeployProfile::Serverless);
598 let embedded = BudgetShares::resolve(budget(GIB), DeployProfile::Embedded);
599
600 assert!(serverless.page_cache_slots() < embedded.page_cache_slots());
601 assert!(serverless.blob_cache_l1_bytes() < embedded.blob_cache_l1_bytes());
602 assert!(serverless.total_share_bytes() < embedded.total_share_bytes());
603 assert_eq!(
604 serverless.share_bytes(MemoryPool::SegmentArena),
605 embedded.share_bytes(MemoryPool::SegmentArena),
606 "the RAM-resident store keeps its share on every profile"
607 );
608 }
609
610 #[test]
611 fn the_serverless_default_budget_produces_a_bounded_blob_l1() {
612 let shares = BudgetShares::resolve(
613 budget(super::super::memory_budget::SERVERLESS_PROFILE_BUDGET_BYTES),
614 DeployProfile::Serverless,
615 );
616 assert!(
619 shares.blob_cache_l1_bytes() < crate::storage::cache::blob::DEFAULT_BLOB_L1_BYTES_MAX,
620 "serverless L1 = {} bytes",
621 shares.blob_cache_l1_bytes()
622 );
623 }
624
625 #[test]
626 fn pool_labels_are_stable_and_unique() {
627 let labels: Vec<&str> = MEMORY_POOLS.iter().map(|pool| pool.as_str()).collect();
628 assert_eq!(
629 labels,
630 vec![
631 "page_cache",
632 "blob_cache_l1",
633 "segment_arena",
634 "index_memory",
635 "wal_buffers"
636 ]
637 );
638 let unique: std::collections::HashSet<_> = labels.iter().collect();
639 assert_eq!(unique.len(), MEMORY_POOL_COUNT);
640
641 for (position, pool) in MEMORY_POOLS.iter().enumerate() {
642 assert_eq!(pool.index(), position, "{} index", pool.as_str());
643 }
644 }
645
646 #[test]
647 fn accounting_charges_releases_and_totals_across_pools() {
648 let accounting = MemoryAccounting::new(budget(GIB), DeployProfile::Embedded);
649 assert_eq!(accounting.total_used_bytes(), 0);
650
651 accounting.charge(MemoryPool::SegmentArena, 4_096);
652 accounting.charge(MemoryPool::SegmentArena, 1_024);
653 accounting.charge(MemoryPool::WalBuffers, 64);
654 assert_eq!(accounting.used_bytes(MemoryPool::SegmentArena), 5_120);
655 assert_eq!(accounting.total_used_bytes(), 5_184);
656
657 accounting.release(MemoryPool::SegmentArena, 1_024);
658 assert_eq!(accounting.used_bytes(MemoryPool::SegmentArena), 4_096);
659
660 accounting.report(MemoryPool::SegmentArena, 42);
661 assert_eq!(accounting.used_bytes(MemoryPool::SegmentArena), 42);
662 assert_eq!(accounting.total_used_bytes(), 106);
663 }
664
665 #[test]
666 fn releasing_more_than_charged_saturates_instead_of_wrapping() {
667 let accounting = MemoryAccounting::new(budget(GIB), DeployProfile::Embedded);
668 accounting.charge(MemoryPool::IndexMemory, 10);
669 accounting.release(MemoryPool::IndexMemory, 1_000);
670 assert_eq!(accounting.used_bytes(MemoryPool::IndexMemory), 0);
671 }
672
673 #[test]
674 fn a_snapshot_carries_every_pool_with_its_share_and_usage() {
675 let accounting = MemoryAccounting::new(budget(4 * GIB), DeployProfile::Cluster);
676 accounting.report(MemoryPool::PageCache, 777);
677
678 let snapshot = accounting.snapshot();
679 assert_eq!(snapshot.len(), MEMORY_POOL_COUNT);
680 assert_eq!(snapshot[0].pool, MemoryPool::PageCache);
681 assert_eq!(snapshot[0].used_bytes, 777);
682 assert_eq!(
683 snapshot[0].share_bytes,
684 accounting.share_bytes(MemoryPool::PageCache)
685 );
686 assert_eq!(
687 snapshot.iter().map(|row| row.share_bytes).sum::<u64>(),
688 accounting.total_share_bytes()
689 );
690 assert!(accounting.total_share_bytes() <= accounting.budget().resolved_bytes);
691 }
692
693 proptest::proptest! {
694 #[test]
695 fn randomized_accounting_sequences_never_underflow_or_exceed_when_admitted(
696 ops in proptest::collection::vec((0usize..MEMORY_POOL_COUNT, 0u8..3, 0u64..4096), 1..256)
697 ) {
698 let accounting = MemoryAccounting::new(budget(64 * MIB), DeployProfile::Embedded);
699 let budget = accounting.budget().resolved_bytes;
700
701 for (pool_idx, op, bytes) in ops {
702 let pool = MEMORY_POOLS[pool_idx];
703 match op {
704 0 => {
705 if accounting.total_used_bytes().saturating_add(bytes) <= budget {
706 accounting.charge(pool, bytes);
707 }
708 }
709 1 => accounting.release(pool, bytes),
710 _ => accounting.report(pool, accounting.used_bytes(pool).saturating_sub(bytes)),
711 }
712
713 prop_assert!(accounting.used_bytes(pool) <= budget);
714 prop_assert!(accounting.total_used_bytes() <= budget);
715 }
716 }
717 }
718}