use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Once;
use super::engine::page_cache::MIN_CACHE_CAPACITY;
use super::memory_budget::MemoryBudget;
use super::profile::DeployProfile;
pub const PAGE_CACHE_PAGE_SIZE_BYTES: u64 = reddb_file::PAGED_PAGE_SIZE as u64;
pub const BASIS_POINTS_PER_WHOLE: u32 = 10_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum MemoryPool {
PageCache,
BlobCacheL1,
SegmentArena,
IndexMemory,
WalBuffers,
}
pub const MEMORY_POOLS: [MemoryPool; MEMORY_POOL_COUNT] = [
MemoryPool::PageCache,
MemoryPool::BlobCacheL1,
MemoryPool::SegmentArena,
MemoryPool::IndexMemory,
MemoryPool::WalBuffers,
];
pub const MEMORY_POOL_COUNT: usize = 5;
impl MemoryPool {
pub const fn as_str(self) -> &'static str {
match self {
Self::PageCache => "page_cache",
Self::BlobCacheL1 => "blob_cache_l1",
Self::SegmentArena => "segment_arena",
Self::IndexMemory => "index_memory",
Self::WalBuffers => "wal_buffers",
}
}
pub const fn index(self) -> usize {
match self {
Self::PageCache => 0,
Self::BlobCacheL1 => 1,
Self::SegmentArena => 2,
Self::IndexMemory => 3,
Self::WalBuffers => 4,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BudgetSharePolicy {
basis_points: [u32; MEMORY_POOL_COUNT],
}
const DEFAULT_POLICY: BudgetSharePolicy = BudgetSharePolicy {
basis_points: [
2_000, 1_500, 4_000, 1_000, 500, ],
};
const SERVERLESS_POLICY: BudgetSharePolicy = BudgetSharePolicy {
basis_points: [
1_500, 1_000, 4_000, 1_000, 500, ],
};
impl BudgetSharePolicy {
pub const fn for_profile(profile: DeployProfile) -> Self {
match profile {
DeployProfile::Serverless => SERVERLESS_POLICY,
DeployProfile::Embedded | DeployProfile::PrimaryReplica | DeployProfile::Cluster => {
DEFAULT_POLICY
}
}
}
pub const fn basis_points(&self, pool: MemoryPool) -> u32 {
self.basis_points[pool.index()]
}
pub fn total_basis_points(&self) -> u32 {
self.basis_points.iter().sum()
}
pub fn reserve_basis_points(&self) -> u32 {
BASIS_POINTS_PER_WHOLE - self.total_basis_points()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BudgetShares {
budget_bytes: u64,
policy: BudgetSharePolicy,
share_bytes: [u64; MEMORY_POOL_COUNT],
}
impl BudgetShares {
pub fn resolve(budget: MemoryBudget, profile: DeployProfile) -> Self {
let policy = BudgetSharePolicy::for_profile(profile);
assert!(
policy.total_basis_points() > 0,
"invariant: budget share policy must claim some of the budget"
);
assert!(
policy.total_basis_points() <= BASIS_POINTS_PER_WHOLE,
"invariant: Σ(share fractions) = {} bp exceeds the whole budget ({BASIS_POINTS_PER_WHOLE} bp)",
policy.total_basis_points()
);
let budget_bytes = budget.resolved_bytes;
assert!(
budget_bytes > 0,
"invariant: the resolved budget is strictly positive (ADR 0073 §1: no unlimited mode)"
);
let per_basis_point = budget_bytes / u64::from(BASIS_POINTS_PER_WHOLE);
let mut share_bytes = [0_u64; MEMORY_POOL_COUNT];
for pool in MEMORY_POOLS {
share_bytes[pool.index()] = per_basis_point * u64::from(policy.basis_points(pool));
}
let shares = Self {
budget_bytes,
policy,
share_bytes,
};
assert!(
shares.total_share_bytes() <= budget_bytes,
"invariant: Σ(shares) = {} exceeds budget {budget_bytes}",
shares.total_share_bytes()
);
assert!(
shares.total_share_bytes() > 0 || budget_bytes < u64::from(BASIS_POINTS_PER_WHOLE),
"invariant: a budget of {budget_bytes} bytes must reach the pools"
);
shares
}
pub fn budget_bytes(&self) -> u64 {
self.budget_bytes
}
pub fn policy(&self) -> BudgetSharePolicy {
self.policy
}
pub fn share_bytes(&self, pool: MemoryPool) -> u64 {
self.share_bytes[pool.index()]
}
pub fn total_share_bytes(&self) -> u64 {
self.share_bytes.iter().sum()
}
pub fn page_cache_slots(&self) -> usize {
let slots = self.share_bytes(MemoryPool::PageCache) / PAGE_CACHE_PAGE_SIZE_BYTES;
usize::try_from(slots)
.unwrap_or(usize::MAX)
.max(MIN_CACHE_CAPACITY)
}
pub fn blob_cache_l1_bytes(&self) -> usize {
usize::try_from(self.share_bytes(MemoryPool::BlobCacheL1)).unwrap_or(usize::MAX)
}
pub fn log_once(&self) {
static LOGGED: Once = Once::new();
LOGGED.call_once(|| {
for pool in MEMORY_POOLS {
tracing::info!(
pool = pool.as_str(),
share_bytes = self.share_bytes(pool),
basis_points = self.policy.basis_points(pool),
budget_bytes = self.budget_bytes,
"memory budget share assigned"
);
}
tracing::info!(
total_share_bytes = self.total_share_bytes(),
reserve_basis_points = self.policy.reserve_basis_points(),
budget_bytes = self.budget_bytes,
"memory budget shares resolved"
);
});
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PoolUsage {
pub pool: MemoryPool,
pub share_bytes: u64,
pub used_bytes: u64,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct MemoryEnforcementSnapshot {
pub admissions_denied: u64,
pub pressure_reclamations_triggered: u64,
pub pressure_bytes_reclaimed: u64,
pub high_water_bytes: u64,
}
#[derive(Debug)]
pub struct MemoryAccounting {
budget: MemoryBudget,
shares: BudgetShares,
used_bytes: [AtomicU64; MEMORY_POOL_COUNT],
admissions_denied: AtomicU64,
pressure_reclamations_triggered: AtomicU64,
pressure_bytes_reclaimed: AtomicU64,
high_water_bytes: AtomicU64,
}
impl MemoryAccounting {
pub fn new(budget: MemoryBudget, profile: DeployProfile) -> Self {
Self::from_shares(budget, BudgetShares::resolve(budget, profile))
}
pub fn from_shares(budget: MemoryBudget, shares: BudgetShares) -> Self {
assert_eq!(
budget.resolved_bytes, shares.budget_bytes,
"invariant: accounting shares must divide this process's budget"
);
Self {
budget,
shares,
used_bytes: std::array::from_fn(|_| AtomicU64::new(0)),
admissions_denied: AtomicU64::new(0),
pressure_reclamations_triggered: AtomicU64::new(0),
pressure_bytes_reclaimed: AtomicU64::new(0),
high_water_bytes: AtomicU64::new(0),
}
}
pub fn budget(&self) -> MemoryBudget {
self.budget
}
pub fn shares(&self) -> BudgetShares {
self.shares
}
pub fn report(&self, pool: MemoryPool, bytes: u64) {
self.used_bytes[pool.index()].store(bytes, Ordering::Relaxed);
}
pub fn charge(&self, pool: MemoryPool, bytes: u64) {
self.used_bytes[pool.index()].fetch_add(bytes, Ordering::Relaxed);
}
pub fn release(&self, pool: MemoryPool, bytes: u64) {
let slot = &self.used_bytes[pool.index()];
let _ = slot.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
Some(used.saturating_sub(bytes))
});
}
pub fn used_bytes(&self, pool: MemoryPool) -> u64 {
self.used_bytes[pool.index()].load(Ordering::Relaxed)
}
pub fn share_bytes(&self, pool: MemoryPool) -> u64 {
self.shares.share_bytes(pool)
}
pub fn total_used_bytes(&self) -> u64 {
MEMORY_POOLS
.iter()
.map(|pool| self.used_bytes(*pool))
.fold(0, u64::saturating_add)
}
pub fn total_share_bytes(&self) -> u64 {
self.shares.total_share_bytes()
}
pub fn snapshot(&self) -> [PoolUsage; MEMORY_POOL_COUNT] {
MEMORY_POOLS.map(|pool| PoolUsage {
pool,
share_bytes: self.share_bytes(pool),
used_bytes: self.used_bytes(pool),
})
}
pub fn observe_total_used(&self, total: u64) {
let mut current = self.high_water_bytes.load(Ordering::Relaxed);
while total > current {
match self.high_water_bytes.compare_exchange_weak(
current,
total,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => break,
Err(observed) => current = observed,
}
}
}
pub fn record_admission_denied(&self) {
self.admissions_denied.fetch_add(1, Ordering::Relaxed);
}
pub fn record_pressure_reclamation(&self, reclaimed_bytes: u64) {
self.pressure_reclamations_triggered
.fetch_add(1, Ordering::Relaxed);
self.pressure_bytes_reclaimed
.fetch_add(reclaimed_bytes, Ordering::Relaxed);
}
pub fn enforcement_snapshot(&self) -> MemoryEnforcementSnapshot {
MemoryEnforcementSnapshot {
admissions_denied: self.admissions_denied.load(Ordering::Relaxed),
pressure_reclamations_triggered: self
.pressure_reclamations_triggered
.load(Ordering::Relaxed),
pressure_bytes_reclaimed: self.pressure_bytes_reclaimed.load(Ordering::Relaxed),
high_water_bytes: self.high_water_bytes.load(Ordering::Relaxed),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::memory_budget::MemoryBudgetSource;
use proptest::prelude::*;
const GIB: u64 = 1 << 30;
const MIB: u64 = 1 << 20;
const PROFILES: [DeployProfile; 4] = [
DeployProfile::Embedded,
DeployProfile::Serverless,
DeployProfile::PrimaryReplica,
DeployProfile::Cluster,
];
fn budget(bytes: u64) -> MemoryBudget {
MemoryBudget {
resolved_bytes: bytes,
source: MemoryBudgetSource::Config,
}
}
#[test]
fn every_policy_claims_at_most_the_whole_budget() {
for profile in PROFILES {
let policy = BudgetSharePolicy::for_profile(profile);
assert!(policy.total_basis_points() > 0, "{profile:?}");
assert!(
policy.total_basis_points() <= BASIS_POINTS_PER_WHOLE,
"{profile:?} claims {} bp",
policy.total_basis_points()
);
assert!(
policy.reserve_basis_points() > 0,
"{profile:?} must leave an unpooled reserve for plan caches and slack"
);
}
}
#[test]
fn sum_of_shares_never_exceeds_the_budget() {
let budgets = [
1,
2,
9_999,
10_000,
10_001,
64 * MIB,
256 * MIB,
512 * MIB,
GIB,
2 * GIB,
11 * GIB,
1 << 40,
1 << 50,
u64::MAX / 2,
u64::MAX,
];
for profile in PROFILES {
for bytes in budgets {
let shares = BudgetShares::resolve(budget(bytes), profile);
assert!(
shares.total_share_bytes() <= bytes,
"{profile:?} budget {bytes}: Σ(shares) = {}",
shares.total_share_bytes()
);
if bytes >= u64::from(BASIS_POINTS_PER_WHOLE) {
for pool in MEMORY_POOLS {
assert!(
shares.share_bytes(pool) > 0,
"{profile:?} budget {bytes}: {} starved",
pool.as_str()
);
}
}
}
}
}
#[test]
fn sum_of_shares_never_exceeds_the_budget_across_a_dense_sweep() {
for profile in PROFILES {
for step in 0..2_000_u64 {
let bytes = 1 + step * 7_919; let shares = BudgetShares::resolve(budget(bytes), profile);
assert!(
shares.total_share_bytes() <= bytes,
"{profile:?} budget {bytes}: Σ(shares) = {}",
shares.total_share_bytes()
);
}
}
}
#[test]
fn doubling_the_budget_doubles_the_page_cache_slots_and_the_blob_l1_ceiling() {
for profile in PROFILES {
let base = BudgetShares::resolve(budget(GIB), profile);
let doubled = BudgetShares::resolve(budget(2 * GIB), profile);
assert_eq!(
doubled.page_cache_slots(),
base.page_cache_slots() * 2,
"{profile:?} page cache slots must scale with the budget"
);
assert_eq!(
doubled.blob_cache_l1_bytes(),
base.blob_cache_l1_bytes() * 2,
"{profile:?} blob L1 ceiling must scale with the budget"
);
assert_eq!(
doubled.share_bytes(MemoryPool::SegmentArena),
base.share_bytes(MemoryPool::SegmentArena) * 2,
);
}
}
#[test]
fn page_cache_slots_are_the_share_divided_by_the_fixed_page_size() {
let shares = BudgetShares::resolve(budget(10 * GIB), DeployProfile::Embedded);
let expected = shares.share_bytes(MemoryPool::PageCache) / PAGE_CACHE_PAGE_SIZE_BYTES;
assert_eq!(shares.page_cache_slots() as u64, expected);
assert_eq!(PAGE_CACHE_PAGE_SIZE_BYTES, 16 * 1024);
}
#[test]
fn a_budget_too_small_to_fill_a_slot_still_yields_a_usable_page_cache() {
let shares = BudgetShares::resolve(budget(1), DeployProfile::Serverless);
assert_eq!(shares.total_share_bytes(), 0, "nothing reaches the pools");
assert_eq!(shares.page_cache_slots(), MIN_CACHE_CAPACITY);
assert_eq!(shares.blob_cache_l1_bytes(), 0);
}
#[test]
fn serverless_shrinks_the_caches_relative_to_the_server_profiles() {
let serverless = BudgetShares::resolve(budget(GIB), DeployProfile::Serverless);
let embedded = BudgetShares::resolve(budget(GIB), DeployProfile::Embedded);
assert!(serverless.page_cache_slots() < embedded.page_cache_slots());
assert!(serverless.blob_cache_l1_bytes() < embedded.blob_cache_l1_bytes());
assert!(serverless.total_share_bytes() < embedded.total_share_bytes());
assert_eq!(
serverless.share_bytes(MemoryPool::SegmentArena),
embedded.share_bytes(MemoryPool::SegmentArena),
"the RAM-resident store keeps its share on every profile"
);
}
#[test]
fn the_serverless_default_budget_produces_a_bounded_blob_l1() {
let shares = BudgetShares::resolve(
budget(super::super::memory_budget::SERVERLESS_PROFILE_BUDGET_BYTES),
DeployProfile::Serverless,
);
assert!(
shares.blob_cache_l1_bytes() < crate::storage::cache::blob::DEFAULT_BLOB_L1_BYTES_MAX,
"serverless L1 = {} bytes",
shares.blob_cache_l1_bytes()
);
}
#[test]
fn pool_labels_are_stable_and_unique() {
let labels: Vec<&str> = MEMORY_POOLS.iter().map(|pool| pool.as_str()).collect();
assert_eq!(
labels,
vec![
"page_cache",
"blob_cache_l1",
"segment_arena",
"index_memory",
"wal_buffers"
]
);
let unique: std::collections::HashSet<_> = labels.iter().collect();
assert_eq!(unique.len(), MEMORY_POOL_COUNT);
for (position, pool) in MEMORY_POOLS.iter().enumerate() {
assert_eq!(pool.index(), position, "{} index", pool.as_str());
}
}
#[test]
fn accounting_charges_releases_and_totals_across_pools() {
let accounting = MemoryAccounting::new(budget(GIB), DeployProfile::Embedded);
assert_eq!(accounting.total_used_bytes(), 0);
accounting.charge(MemoryPool::SegmentArena, 4_096);
accounting.charge(MemoryPool::SegmentArena, 1_024);
accounting.charge(MemoryPool::WalBuffers, 64);
assert_eq!(accounting.used_bytes(MemoryPool::SegmentArena), 5_120);
assert_eq!(accounting.total_used_bytes(), 5_184);
accounting.release(MemoryPool::SegmentArena, 1_024);
assert_eq!(accounting.used_bytes(MemoryPool::SegmentArena), 4_096);
accounting.report(MemoryPool::SegmentArena, 42);
assert_eq!(accounting.used_bytes(MemoryPool::SegmentArena), 42);
assert_eq!(accounting.total_used_bytes(), 106);
}
#[test]
fn releasing_more_than_charged_saturates_instead_of_wrapping() {
let accounting = MemoryAccounting::new(budget(GIB), DeployProfile::Embedded);
accounting.charge(MemoryPool::IndexMemory, 10);
accounting.release(MemoryPool::IndexMemory, 1_000);
assert_eq!(accounting.used_bytes(MemoryPool::IndexMemory), 0);
}
#[test]
fn a_snapshot_carries_every_pool_with_its_share_and_usage() {
let accounting = MemoryAccounting::new(budget(4 * GIB), DeployProfile::Cluster);
accounting.report(MemoryPool::PageCache, 777);
let snapshot = accounting.snapshot();
assert_eq!(snapshot.len(), MEMORY_POOL_COUNT);
assert_eq!(snapshot[0].pool, MemoryPool::PageCache);
assert_eq!(snapshot[0].used_bytes, 777);
assert_eq!(
snapshot[0].share_bytes,
accounting.share_bytes(MemoryPool::PageCache)
);
assert_eq!(
snapshot.iter().map(|row| row.share_bytes).sum::<u64>(),
accounting.total_share_bytes()
);
assert!(accounting.total_share_bytes() <= accounting.budget().resolved_bytes);
}
proptest::proptest! {
#[test]
fn randomized_accounting_sequences_never_underflow_or_exceed_when_admitted(
ops in proptest::collection::vec((0usize..MEMORY_POOL_COUNT, 0u8..3, 0u64..4096), 1..256)
) {
let accounting = MemoryAccounting::new(budget(64 * MIB), DeployProfile::Embedded);
let budget = accounting.budget().resolved_bytes;
for (pool_idx, op, bytes) in ops {
let pool = MEMORY_POOLS[pool_idx];
match op {
0 => {
if accounting.total_used_bytes().saturating_add(bytes) <= budget {
accounting.charge(pool, bytes);
}
}
1 => accounting.release(pool, bytes),
_ => accounting.report(pool, accounting.used_bytes(pool).saturating_sub(bytes)),
}
prop_assert!(accounting.used_bytes(pool) <= budget);
prop_assert!(accounting.total_used_bytes() <= budget);
}
}
}
}