Skip to main content

reddb_server/storage/
memory_pools.rs

1//! Budget shares and the one shared accounting pool — ADR 0073 §2 made
2//! executable.
3//!
4//! The budget resolved at boot (`memory_budget`, ADR 0073 §1) is divided into
5//! **named shares** by a single allocation policy. The big memory consumers
6//! pre-size their structures from those shares and report live usage into
7//! **one shared accounting pool**.
8//!
9//! Per-subsystem caps *without* shared accounting are explicitly rejected by
10//! the ADR: each subsystem stays individually "within limits" while the sum
11//! kills the process (OOM by summation). One pool, one total.
12//!
13//! The invariant that makes the pool meaningful is
14//!
15//! > **Σ(shares) ≤ budget**
16//!
17//! It holds by construction: every share is `(budget / 10_000) * basis_points`
18//! and the policy's basis points sum to at most `10_000`. Boot asserts both
19//! halves of that statement rather than trusting the arithmetic.
20//!
21//! Reporting uses one atomic per pool. Runtime admission also shares a short
22//! ledger lock for sampled storage, outstanding mutations and owned query
23//! buffers. The ledger never holds storage locks or samples storage itself;
24//! callbacks can reserve output while the sampler waits for their read locks.
25
26use 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
33/// Fixed page size the page-cache share is divided by to obtain a slot count.
34/// Pages are fixed size → slots are fixed size → the arena is preallocated.
35pub const PAGE_CACHE_PAGE_SIZE_BYTES: u64 = reddb_file::PAGED_PAGE_SIZE as u64;
36
37/// Denominator of the policy fractions. Shares are expressed in basis points
38/// (parts per ten thousand) so the whole policy is integer arithmetic — no
39/// float rounding standing between the budget and the invariant.
40pub const BASIS_POINTS_PER_WHOLE: u32 = 10_000;
41
42/// The big memory consumers that receive a slice of the budget.
43///
44/// Only structures whose growth is `O(data)` or `O(traffic)` belong here.
45/// L2's *disk* extent is not memory and stays out; its RAM-resident metadata
46/// (B+ tree index, synopsis filters) is accounted under [`MemoryPool::BlobCacheL1`]
47/// because it is the blob cache's RAM footprint.
48#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
49pub enum MemoryPool {
50    /// Pager page cache — preallocated 16 KiB slots (ADR 0033 SIEVE cache).
51    PageCache,
52    /// Blob Cache RAM tier: L1 entries plus L2's RAM-resident metadata.
53    BlobCacheL1,
54    /// Unified segment arena — active/retained segments and in-flight consolidation.
55    SegmentArena,
56    /// Secondary-index memory — hash / bitmap / sorted / composite.
57    IndexMemory,
58    /// WAL group-commit queue and writer buffers.
59    WalBuffers,
60    /// Owned result and serialized query buffers, retained until their last owner drops.
61    QueryBuffers,
62}
63
64/// Every pool, in the order `red.stats` reports them.
65pub const MEMORY_POOLS: [MemoryPool; MEMORY_POOL_COUNT] = [
66    MemoryPool::PageCache,
67    MemoryPool::BlobCacheL1,
68    MemoryPool::SegmentArena,
69    MemoryPool::IndexMemory,
70    MemoryPool::WalBuffers,
71    MemoryPool::QueryBuffers,
72];
73
74/// Number of pools. Arrays are indexed by [`MemoryPool::index`].
75pub const MEMORY_POOL_COUNT: usize = 6;
76
77impl MemoryPool {
78    /// Stable label echoed by the boot log and the `red.stats` budget section.
79    pub const fn as_str(self) -> &'static str {
80        match self {
81            Self::PageCache => "page_cache",
82            Self::BlobCacheL1 => "blob_cache_l1",
83            Self::SegmentArena => "segment_arena",
84            Self::IndexMemory => "index_memory",
85            Self::WalBuffers => "wal_buffers",
86            Self::QueryBuffers => "query_buffers",
87        }
88    }
89
90    /// Dense index into the share and usage arrays.
91    pub const fn index(self) -> usize {
92        match self {
93            Self::PageCache => 0,
94            Self::BlobCacheL1 => 1,
95            Self::SegmentArena => 2,
96            Self::IndexMemory => 3,
97            Self::WalBuffers => 4,
98            Self::QueryBuffers => 5,
99        }
100    }
101}
102
103/// The single allocation policy: named fractions per pool, profile-adjustable.
104///
105/// No subsystem computes its own fraction. Adding a pool means adding a
106/// fraction here and taking it from another pool or from the reserve — which
107/// is the point: the tradeoff is visible in one place.
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub struct BudgetSharePolicy {
110    basis_points: [u32; MEMORY_POOL_COUNT],
111}
112
113/// Server-side profiles: the segment arena *is* the database (RAM-resident
114/// store), the page cache backs the paged tier and the B+ trees over it.
115/// The unnamed remainder (5%) absorbs plan caches, connection state and
116/// allocator slack — everything that is not a governed pool.
117const DEFAULT_POLICY: BudgetSharePolicy = BudgetSharePolicy {
118    basis_points: [
119        2_000, // page_cache
120        1_500, // blob_cache_l1
121        4_000, // segment_arena
122        1_000, // index_memory
123        500,   // wal_buffers
124        500,   // query_buffers
125    ],
126};
127
128/// Serverless: boundedness is a survival contract, not a tuning preference
129/// (ADR 0038 §1). A cold function instance has a small budget and a short
130/// life, so the caches shrink and the unnamed reserve doubles to 10%.
131const SERVERLESS_POLICY: BudgetSharePolicy = BudgetSharePolicy {
132    basis_points: [
133        1_500, // page_cache
134        1_000, // blob_cache_l1
135        4_000, // segment_arena
136        1_000, // index_memory
137        500,   // wal_buffers
138        1_000, // query_buffers
139    ],
140};
141
142impl BudgetSharePolicy {
143    /// The one policy for a deployment profile.
144    pub const fn for_profile(profile: DeployProfile) -> Self {
145        match profile {
146            DeployProfile::Serverless => SERVERLESS_POLICY,
147            DeployProfile::Embedded | DeployProfile::PrimaryReplica | DeployProfile::Cluster => {
148                DEFAULT_POLICY
149            }
150        }
151    }
152
153    /// This pool's fraction of the budget, in basis points.
154    pub const fn basis_points(&self, pool: MemoryPool) -> u32 {
155        self.basis_points[pool.index()]
156    }
157
158    /// Sum of every pool's fraction. Never exceeds [`BASIS_POINTS_PER_WHOLE`].
159    pub fn total_basis_points(&self) -> u32 {
160        self.basis_points.iter().sum()
161    }
162
163    /// The slice of the budget handed to no pool: plan caches, connection
164    /// state, allocator slack.
165    pub fn reserve_basis_points(&self) -> u32 {
166        BASIS_POINTS_PER_WHOLE - self.total_basis_points()
167    }
168}
169
170/// The resolved per-pool shares of one process budget.
171#[derive(Debug, Clone, Copy, PartialEq, Eq)]
172pub struct BudgetShares {
173    budget_bytes: u64,
174    policy: BudgetSharePolicy,
175    share_bytes: [u64; MEMORY_POOL_COUNT],
176}
177
178impl BudgetShares {
179    /// Divide a resolved budget among the pools.
180    ///
181    /// Integer division comes *first* (`budget / 10_000` then `* basis_points`)
182    /// for two reasons: `budget * basis_points` overflows `u64` for budgets
183    /// above ~1.8 EiB, and dividing first makes Σ(shares) ≤ budget an identity
184    /// rather than a rounding accident. The cost is at most 9_999 bytes of
185    /// unallocated remainder per pool.
186    pub fn resolve(budget: MemoryBudget, profile: DeployProfile) -> Self {
187        let policy = BudgetSharePolicy::for_profile(profile);
188
189        // Pair assertion: the policy claims a positive but not-more-than-whole
190        // fraction of the budget. Both halves, because a policy summing to 0
191        // is as broken as one summing to 12_000 and neither is caught by the
192        // other check.
193        assert!(
194            policy.total_basis_points() > 0,
195            "invariant: budget share policy must claim some of the budget"
196        );
197        assert!(
198            policy.total_basis_points() <= BASIS_POINTS_PER_WHOLE,
199            "invariant: Σ(share fractions) = {} bp exceeds the whole budget ({BASIS_POINTS_PER_WHOLE} bp)",
200            policy.total_basis_points()
201        );
202
203        let budget_bytes = budget.resolved_bytes;
204        assert!(
205            budget_bytes > 0,
206            "invariant: the resolved budget is strictly positive (ADR 0073 §1: no unlimited mode)"
207        );
208
209        let per_basis_point = budget_bytes / u64::from(BASIS_POINTS_PER_WHOLE);
210        let mut share_bytes = [0_u64; MEMORY_POOL_COUNT];
211        for pool in MEMORY_POOLS {
212            share_bytes[pool.index()] = per_basis_point * u64::from(policy.basis_points(pool));
213        }
214
215        let shares = Self {
216            budget_bytes,
217            policy,
218            share_bytes,
219        };
220
221        // The property this whole module exists to guarantee, asserted at boot
222        // rather than assumed. Its negative space too: the shares are not all
223        // zero unless the budget itself is smaller than a basis point.
224        assert!(
225            shares.total_share_bytes() <= budget_bytes,
226            "invariant: Σ(shares) = {} exceeds budget {budget_bytes}",
227            shares.total_share_bytes()
228        );
229        assert!(
230            shares.total_share_bytes() > 0 || budget_bytes < u64::from(BASIS_POINTS_PER_WHOLE),
231            "invariant: a budget of {budget_bytes} bytes must reach the pools"
232        );
233
234        shares
235    }
236
237    /// The budget these shares divide.
238    pub fn budget_bytes(&self) -> u64 {
239        self.budget_bytes
240    }
241
242    /// The policy that produced these shares.
243    pub fn policy(&self) -> BudgetSharePolicy {
244        self.policy
245    }
246
247    /// This pool's slice of the budget.
248    pub fn share_bytes(&self, pool: MemoryPool) -> u64 {
249        self.share_bytes[pool.index()]
250    }
251
252    /// Σ(shares). Never exceeds [`Self::budget_bytes`].
253    pub fn total_share_bytes(&self) -> u64 {
254        self.share_bytes.iter().sum()
255    }
256
257    /// Page-cache slot count: the share divided by the fixed page size.
258    ///
259    /// Clamped to the cache's structural minimum. A budget small enough to hit
260    /// the clamp buys a page cache slightly larger than its share; that is
261    /// visible in `red.stats` as `used_bytes > share_bytes` and, per ADR 0073
262    /// §4, is the enforcement slice's problem, not this one's.
263    pub fn page_cache_slots(&self) -> usize {
264        let slots = self.share_bytes(MemoryPool::PageCache) / PAGE_CACHE_PAGE_SIZE_BYTES;
265        usize::try_from(slots)
266            .unwrap_or(usize::MAX)
267            .max(MIN_CACHE_CAPACITY)
268    }
269
270    /// Blob Cache L1 byte ceiling: the RAM tier's share, verbatim.
271    pub fn blob_cache_l1_bytes(&self) -> usize {
272        usize::try_from(self.share_bytes(MemoryPool::BlobCacheL1)).unwrap_or(usize::MAX)
273    }
274
275    /// Emit one boot log line per pool naming its share. Guarded so a process
276    /// that opens several runtimes still logs exactly once, matching
277    /// `memory_budget::log_resolved_once`.
278    pub fn log_once(&self) {
279        static LOGGED: Once = Once::new();
280        LOGGED.call_once(|| {
281            for pool in MEMORY_POOLS {
282                tracing::info!(
283                    pool = pool.as_str(),
284                    share_bytes = self.share_bytes(pool),
285                    basis_points = self.policy.basis_points(pool),
286                    budget_bytes = self.budget_bytes,
287                    "memory budget share assigned"
288                );
289            }
290            tracing::info!(
291                total_share_bytes = self.total_share_bytes(),
292                reserve_basis_points = self.policy.reserve_basis_points(),
293                budget_bytes = self.budget_bytes,
294                "memory budget shares resolved"
295            );
296        });
297    }
298}
299
300/// One pool's row in the shared accounting.
301#[derive(Debug, Clone, Copy, PartialEq, Eq)]
302pub struct PoolUsage {
303    pub pool: MemoryPool,
304    pub share_bytes: u64,
305    pub used_bytes: u64,
306}
307
308/// Process-level enforcement counters surfaced by `red.stats`.
309#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
310pub struct MemoryEnforcementSnapshot {
311    pub admissions_denied: u64,
312    pub pressure_reclamations_triggered: u64,
313    pub pressure_bytes_reclaimed: u64,
314    pub high_water_bytes: u64,
315}
316
317/// The one shared accounting pool: fixed shares, live usage.
318///
319/// Usage is a relaxed atomic per pool. Reporting is a single store or
320/// fetch-add — no allocation, no lock, nothing a read hot path has to wait on.
321/// Relaxed is the right ordering because the counters are observability, not
322/// synchronisation: no other memory is published through them.
323#[derive(Debug)]
324pub struct MemoryAccounting {
325    budget: MemoryBudget,
326    shares: BudgetShares,
327    used_bytes: [AtomicU64; MEMORY_POOL_COUNT],
328    admissions_denied: AtomicU64,
329    pressure_reclamations_triggered: AtomicU64,
330    pressure_bytes_reclaimed: AtomicU64,
331    high_water_bytes: AtomicU64,
332    // Sampled storage plus outstanding mutation growth. This short ledger lock
333    // never covers sampling or storage access: read callbacks can admit output
334    // while a sampler waits for their segment lock without an ABBA deadlock.
335    admission_non_query_bytes: parking_lot::Mutex<u64>,
336}
337
338impl MemoryAccounting {
339    /// Build the pool from a resolved budget and the profile's policy.
340    pub fn new(budget: MemoryBudget, profile: DeployProfile) -> Self {
341        Self::from_shares(budget, BudgetShares::resolve(budget, profile))
342    }
343
344    /// Build the pool from shares already resolved (and already asserted).
345    pub fn from_shares(budget: MemoryBudget, shares: BudgetShares) -> Self {
346        assert_eq!(
347            budget.resolved_bytes, shares.budget_bytes,
348            "invariant: accounting shares must divide this process's budget"
349        );
350        Self {
351            budget,
352            shares,
353            used_bytes: std::array::from_fn(|_| AtomicU64::new(0)),
354            admissions_denied: AtomicU64::new(0),
355            pressure_reclamations_triggered: AtomicU64::new(0),
356            pressure_bytes_reclaimed: AtomicU64::new(0),
357            high_water_bytes: AtomicU64::new(0),
358            admission_non_query_bytes: parking_lot::Mutex::new(0),
359        }
360    }
361
362    /// The budget this process runs under.
363    pub fn budget(&self) -> MemoryBudget {
364        self.budget
365    }
366
367    /// The shares this pool accounts against.
368    pub fn shares(&self) -> BudgetShares {
369        self.shares
370    }
371
372    /// Overwrite a pool's live usage. The reporting shape for pools whose
373    /// footprint is a measured total rather than a running delta.
374    pub fn report(&self, pool: MemoryPool, bytes: u64) {
375        self.used_bytes[pool.index()].store(bytes, Ordering::Relaxed);
376    }
377
378    /// Charge `bytes` to a pool. One relaxed fetch-add.
379    pub fn charge(&self, pool: MemoryPool, bytes: u64) {
380        self.used_bytes[pool.index()].fetch_add(bytes, Ordering::Relaxed);
381    }
382
383    /// Return `bytes` to a pool, saturating at zero.
384    ///
385    /// Saturation rather than wrap: a double-release is a programmer error the
386    /// enforcement slice will assert on, but until then an accounting counter
387    /// underflowing to 18 exabytes would make `red.stats` lie spectacularly.
388    pub fn release(&self, pool: MemoryPool, bytes: u64) {
389        let slot = &self.used_bytes[pool.index()];
390        let _ = slot.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
391            Some(used.saturating_sub(bytes))
392        });
393    }
394
395    /// Publish a completed storage sample and its still-outstanding mutation
396    /// reservations. The runtime serializes samplers; owned query charges are
397    /// read under this ledger lock and can never be overwritten by a sample.
398    pub(crate) fn reconcile_non_query_admission(&self, bytes: u64) {
399        let mut non_query = self.admission_non_query_bytes.lock();
400        *non_query = bytes;
401        self.observe_total_used(bytes.saturating_add(self.used_bytes(MemoryPool::QueryBuffers)));
402    }
403
404    pub(crate) fn try_charge_query_buffer(&self, bytes: u64) -> bool {
405        let non_query = self.admission_non_query_bytes.lock();
406        let Some(projected) = non_query
407            .checked_add(self.used_bytes(MemoryPool::QueryBuffers))
408            .and_then(|used| used.checked_add(bytes))
409        else {
410            return false;
411        };
412        if projected > self.budget.resolved_bytes {
413            return false;
414        }
415        self.charge(MemoryPool::QueryBuffers, bytes);
416        self.observe_total_used(projected);
417        true
418    }
419
420    /// Hold growth until a later sample observes its completed publication.
421    pub(crate) fn try_reserve_mutation_growth(&self, bytes: u64) -> bool {
422        let mut non_query = self.admission_non_query_bytes.lock();
423        let Some(projected) = non_query
424            .checked_add(self.used_bytes(MemoryPool::QueryBuffers))
425            .and_then(|used| used.checked_add(bytes))
426        else {
427            return false;
428        };
429        if projected > self.budget.resolved_bytes {
430            return false;
431        }
432        *non_query += bytes;
433        self.observe_total_used(projected);
434        true
435    }
436
437    /// Owned buffers must be charged once and released once. A mismatch is an
438    /// ownership bug, not a sampled counter that can safely saturate.
439    pub(crate) fn release_query_buffer(&self, bytes: u64) {
440        self.used_bytes[MemoryPool::QueryBuffers.index()]
441            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
442                Some(used.checked_sub(bytes).expect("query buffer released once"))
443            })
444            .expect("query buffer counter update cannot be rejected");
445    }
446
447    /// This pool's live usage.
448    pub fn used_bytes(&self, pool: MemoryPool) -> u64 {
449        self.used_bytes[pool.index()].load(Ordering::Relaxed)
450    }
451
452    /// This pool's share of the budget.
453    pub fn share_bytes(&self, pool: MemoryPool) -> u64 {
454        self.shares.share_bytes(pool)
455    }
456
457    /// Σ(used) across every pool — the number the enforcement slice will gate on.
458    pub fn total_used_bytes(&self) -> u64 {
459        MEMORY_POOLS
460            .iter()
461            .map(|pool| self.used_bytes(*pool))
462            .fold(0, u64::saturating_add)
463    }
464
465    /// Σ(shares). Never exceeds the budget.
466    pub fn total_share_bytes(&self) -> u64 {
467        self.shares.total_share_bytes()
468    }
469
470    /// Every pool's share and usage, in [`MEMORY_POOLS`] order.
471    pub fn snapshot(&self) -> [PoolUsage; MEMORY_POOL_COUNT] {
472        MEMORY_POOLS.map(|pool| PoolUsage {
473            pool,
474            share_bytes: self.share_bytes(pool),
475            used_bytes: self.used_bytes(pool),
476        })
477    }
478
479    /// Record a sampled total as the shared accounting high-water mark.
480    pub fn observe_total_used(&self, total: u64) {
481        let mut current = self.high_water_bytes.load(Ordering::Relaxed);
482        while total > current {
483            match self.high_water_bytes.compare_exchange_weak(
484                current,
485                total,
486                Ordering::Relaxed,
487                Ordering::Relaxed,
488            ) {
489                Ok(_) => break,
490                Err(observed) => current = observed,
491            }
492        }
493    }
494
495    pub fn record_admission_denied(&self) {
496        self.admissions_denied.fetch_add(1, Ordering::Relaxed);
497    }
498
499    pub fn record_pressure_reclamation(&self, reclaimed_bytes: u64) {
500        self.pressure_reclamations_triggered
501            .fetch_add(1, Ordering::Relaxed);
502        self.pressure_bytes_reclaimed
503            .fetch_add(reclaimed_bytes, Ordering::Relaxed);
504    }
505
506    pub fn enforcement_snapshot(&self) -> MemoryEnforcementSnapshot {
507        MemoryEnforcementSnapshot {
508            admissions_denied: self.admissions_denied.load(Ordering::Relaxed),
509            pressure_reclamations_triggered: self
510                .pressure_reclamations_triggered
511                .load(Ordering::Relaxed),
512            pressure_bytes_reclaimed: self.pressure_bytes_reclaimed.load(Ordering::Relaxed),
513            high_water_bytes: self.high_water_bytes.load(Ordering::Relaxed),
514        }
515    }
516}
517
518#[cfg(test)]
519mod tests {
520    use super::*;
521    use crate::storage::memory_budget::MemoryBudgetSource;
522    use proptest::prelude::*;
523
524    const GIB: u64 = 1 << 30;
525    const MIB: u64 = 1 << 20;
526
527    const PROFILES: [DeployProfile; 4] = [
528        DeployProfile::Embedded,
529        DeployProfile::Serverless,
530        DeployProfile::PrimaryReplica,
531        DeployProfile::Cluster,
532    ];
533
534    fn budget(bytes: u64) -> MemoryBudget {
535        MemoryBudget {
536            resolved_bytes: bytes,
537            source: MemoryBudgetSource::Config,
538        }
539    }
540
541    #[test]
542    fn every_policy_claims_at_most_the_whole_budget() {
543        for profile in PROFILES {
544            let policy = BudgetSharePolicy::for_profile(profile);
545            assert!(policy.total_basis_points() > 0, "{profile:?}");
546            assert!(
547                policy.total_basis_points() <= BASIS_POINTS_PER_WHOLE,
548                "{profile:?} claims {} bp",
549                policy.total_basis_points()
550            );
551            assert!(
552                policy.reserve_basis_points() > 0,
553                "{profile:?} must leave an unpooled reserve for plan caches and slack"
554            );
555        }
556    }
557
558    /// The property: for any resolved budget across the valid range, and any
559    /// profile, Σ(shares) ≤ B. Boot asserts the same thing (see `resolve`).
560    #[test]
561    fn sum_of_shares_never_exceeds_the_budget() {
562        let budgets = [
563            1,
564            2,
565            9_999,
566            10_000,
567            10_001,
568            64 * MIB,
569            256 * MIB,
570            512 * MIB,
571            GIB,
572            2 * GIB,
573            11 * GIB,
574            1 << 40,
575            1 << 50,
576            u64::MAX / 2,
577            u64::MAX,
578        ];
579
580        for profile in PROFILES {
581            for bytes in budgets {
582                let shares = BudgetShares::resolve(budget(bytes), profile);
583                assert!(
584                    shares.total_share_bytes() <= bytes,
585                    "{profile:?} budget {bytes}: Σ(shares) = {}",
586                    shares.total_share_bytes()
587                );
588                // Positive space too: nothing above a basis point starves.
589                if bytes >= u64::from(BASIS_POINTS_PER_WHOLE) {
590                    for pool in MEMORY_POOLS {
591                        assert!(
592                            shares.share_bytes(pool) > 0,
593                            "{profile:?} budget {bytes}: {} starved",
594                            pool.as_str()
595                        );
596                    }
597                }
598            }
599        }
600    }
601
602    /// Exhaustive over a dense sweep — the same property, walked rather than
603    /// sampled, so an off-by-one in the basis-point arithmetic cannot hide
604    /// between the hand-picked budgets above.
605    #[test]
606    fn sum_of_shares_never_exceeds_the_budget_across_a_dense_sweep() {
607        for profile in PROFILES {
608            for step in 0..2_000_u64 {
609                let bytes = 1 + step * 7_919; // prime stride: hits every residue mod 10_000
610                let shares = BudgetShares::resolve(budget(bytes), profile);
611                assert!(
612                    shares.total_share_bytes() <= bytes,
613                    "{profile:?} budget {bytes}: Σ(shares) = {}",
614                    shares.total_share_bytes()
615                );
616            }
617        }
618    }
619
620    #[test]
621    fn doubling_the_budget_doubles_the_page_cache_slots_and_the_blob_l1_ceiling() {
622        for profile in PROFILES {
623            let base = BudgetShares::resolve(budget(GIB), profile);
624            let doubled = BudgetShares::resolve(budget(2 * GIB), profile);
625
626            assert_eq!(
627                doubled.page_cache_slots(),
628                base.page_cache_slots() * 2,
629                "{profile:?} page cache slots must scale with the budget"
630            );
631            assert_eq!(
632                doubled.blob_cache_l1_bytes(),
633                base.blob_cache_l1_bytes() * 2,
634                "{profile:?} blob L1 ceiling must scale with the budget"
635            );
636            assert_eq!(
637                doubled.share_bytes(MemoryPool::SegmentArena),
638                base.share_bytes(MemoryPool::SegmentArena) * 2,
639            );
640        }
641    }
642
643    #[test]
644    fn page_cache_slots_are_the_share_divided_by_the_fixed_page_size() {
645        let shares = BudgetShares::resolve(budget(10 * GIB), DeployProfile::Embedded);
646        let expected = shares.share_bytes(MemoryPool::PageCache) / PAGE_CACHE_PAGE_SIZE_BYTES;
647        assert_eq!(shares.page_cache_slots() as u64, expected);
648        assert_eq!(PAGE_CACHE_PAGE_SIZE_BYTES, 16 * 1024);
649    }
650
651    #[test]
652    fn a_budget_too_small_to_fill_a_slot_still_yields_a_usable_page_cache() {
653        let shares = BudgetShares::resolve(budget(1), DeployProfile::Serverless);
654        assert_eq!(shares.total_share_bytes(), 0, "nothing reaches the pools");
655        assert_eq!(shares.page_cache_slots(), MIN_CACHE_CAPACITY);
656        assert_eq!(shares.blob_cache_l1_bytes(), 0);
657    }
658
659    #[test]
660    fn serverless_shrinks_the_caches_relative_to_the_server_profiles() {
661        let serverless = BudgetShares::resolve(budget(GIB), DeployProfile::Serverless);
662        let embedded = BudgetShares::resolve(budget(GIB), DeployProfile::Embedded);
663
664        assert!(serverless.page_cache_slots() < embedded.page_cache_slots());
665        assert!(serverless.blob_cache_l1_bytes() < embedded.blob_cache_l1_bytes());
666        assert!(serverless.total_share_bytes() < embedded.total_share_bytes());
667        assert_eq!(
668            serverless.share_bytes(MemoryPool::SegmentArena),
669            embedded.share_bytes(MemoryPool::SegmentArena),
670            "the RAM-resident store keeps its share on every profile"
671        );
672    }
673
674    #[test]
675    fn the_serverless_default_budget_produces_a_bounded_blob_l1() {
676        let shares = BudgetShares::resolve(
677            budget(super::super::memory_budget::SERVERLESS_PROFILE_BUDGET_BYTES),
678            DeployProfile::Serverless,
679        );
680        // The old hardcoded 256 MiB L1 would have been the *entire* serverless
681        // budget on its own. That is the default this slice deletes.
682        assert!(
683            shares.blob_cache_l1_bytes() < crate::storage::cache::blob::DEFAULT_BLOB_L1_BYTES_MAX,
684            "serverless L1 = {} bytes",
685            shares.blob_cache_l1_bytes()
686        );
687    }
688
689    #[test]
690    fn pool_labels_are_stable_and_unique() {
691        let labels: Vec<&str> = MEMORY_POOLS.iter().map(|pool| pool.as_str()).collect();
692        assert_eq!(
693            labels,
694            vec![
695                "page_cache",
696                "blob_cache_l1",
697                "segment_arena",
698                "index_memory",
699                "wal_buffers",
700                "query_buffers"
701            ]
702        );
703        let unique: std::collections::HashSet<_> = labels.iter().collect();
704        assert_eq!(unique.len(), MEMORY_POOL_COUNT);
705
706        for (position, pool) in MEMORY_POOLS.iter().enumerate() {
707            assert_eq!(pool.index(), position, "{} index", pool.as_str());
708        }
709    }
710
711    #[test]
712    fn accounting_charges_releases_and_totals_across_pools() {
713        let accounting = MemoryAccounting::new(budget(GIB), DeployProfile::Embedded);
714        assert_eq!(accounting.total_used_bytes(), 0);
715
716        accounting.charge(MemoryPool::SegmentArena, 4_096);
717        accounting.charge(MemoryPool::SegmentArena, 1_024);
718        accounting.charge(MemoryPool::WalBuffers, 64);
719        assert_eq!(accounting.used_bytes(MemoryPool::SegmentArena), 5_120);
720        assert_eq!(accounting.total_used_bytes(), 5_184);
721
722        accounting.release(MemoryPool::SegmentArena, 1_024);
723        assert_eq!(accounting.used_bytes(MemoryPool::SegmentArena), 4_096);
724
725        accounting.report(MemoryPool::SegmentArena, 42);
726        assert_eq!(accounting.used_bytes(MemoryPool::SegmentArena), 42);
727        assert_eq!(accounting.total_used_bytes(), 106);
728    }
729
730    #[test]
731    fn releasing_more_than_charged_saturates_instead_of_wrapping() {
732        let accounting = MemoryAccounting::new(budget(GIB), DeployProfile::Embedded);
733        accounting.charge(MemoryPool::IndexMemory, 10);
734        accounting.release(MemoryPool::IndexMemory, 1_000);
735        assert_eq!(accounting.used_bytes(MemoryPool::IndexMemory), 0);
736    }
737
738    #[test]
739    fn a_snapshot_carries_every_pool_with_its_share_and_usage() {
740        let accounting = MemoryAccounting::new(budget(4 * GIB), DeployProfile::Cluster);
741        accounting.report(MemoryPool::PageCache, 777);
742
743        let snapshot = accounting.snapshot();
744        assert_eq!(snapshot.len(), MEMORY_POOL_COUNT);
745        assert_eq!(snapshot[0].pool, MemoryPool::PageCache);
746        assert_eq!(snapshot[0].used_bytes, 777);
747        assert_eq!(
748            snapshot[0].share_bytes,
749            accounting.share_bytes(MemoryPool::PageCache)
750        );
751        assert_eq!(
752            snapshot.iter().map(|row| row.share_bytes).sum::<u64>(),
753            accounting.total_share_bytes()
754        );
755        assert!(accounting.total_share_bytes() <= accounting.budget().resolved_bytes);
756    }
757
758    proptest::proptest! {
759        #[test]
760        fn randomized_accounting_sequences_never_underflow_or_exceed_when_admitted(
761            ops in proptest::collection::vec((0usize..MEMORY_POOL_COUNT, 0u8..3, 0u64..4096), 1..256)
762        ) {
763            let accounting = MemoryAccounting::new(budget(64 * MIB), DeployProfile::Embedded);
764            let budget = accounting.budget().resolved_bytes;
765
766            for (pool_idx, op, bytes) in ops {
767                let pool = MEMORY_POOLS[pool_idx];
768                match op {
769                    0 => {
770                        if accounting.total_used_bytes().saturating_add(bytes) <= budget {
771                            accounting.charge(pool, bytes);
772                        }
773                    }
774                    1 => accounting.release(pool, bytes),
775                    _ => accounting.report(pool, accounting.used_bytes(pool).saturating_sub(bytes)),
776                }
777
778                prop_assert!(accounting.used_bytes(pool) <= budget);
779                prop_assert!(accounting.total_used_bytes() <= budget);
780            }
781        }
782    }
783}