mz_timely_util/
pool_config.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Process-wide installation and configuration point for the buffer pool
17//! that backs chunk spilling.
18//!
19//! Chunk spill consumers, arriving with the chunk-batcher work, resolve their
20//! backing through [`active_pool`] at each spill decision, so live
21//! reconfiguration takes effect on the next call.
22//! The pool itself ([`mz_ore::pool::Pool`]) is a process singleton: live
23//! chunk handles keep their `Arc` into it, so replacing it on reconfigure
24//! would split residency accounting across two budgets. Operator-driven
25//! tunes instead retune the one instance in place via [`apply_pool_config`].
26//!
27//! See `doc/developer/design/20260610_buffer_managed_state.md` for the
28//! buffer-pool design.
29
30#![deny(missing_docs)]
31
32pub mod metrics;
33
34/// Process-wide buffer pool shared by every spill consumer in the process.
35///
36/// Construction reserves virtual address space only: 1 TiB per chunk size
37/// class plus 1 TiB per extent-arena class, tens of TiB in total. Physical
38/// memory is paid per resident chunk. On the rare platforms or
39/// configurations where the reservation fails, the pool is permanently
40/// unavailable for this process and [`apply_pool_config`] reports that by
41/// returning `false`.
42static GLOBAL_POOL: std::sync::OnceLock<Option<mz_ore::pool::Pool>> = std::sync::OnceLock::new();
43
44/// Whether [`apply_pool_config`] has installed the pool as the process's
45/// spill mechanism. [`active_pool`] reads this so consumers stay inert until
46/// the first config apply budgets the pool.
47static POOL_MODE: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
48
49/// Returns the process-wide buffer pool, initializing it on first call.
50/// `None` if the virtual reservation failed at first use.
51pub fn global_pool() -> Option<mz_ore::pool::Pool> {
52    GLOBAL_POOL
53        .get_or_init(|| match mz_ore::pool::Pool::new() {
54            Ok(pool) => Some(pool),
55            Err(err) => {
56                tracing::warn!(
57                    %err,
58                    "buffer pool reservation failed; pool spilling unavailable",
59                );
60                None
61            }
62        })
63        .clone()
64}
65
66/// Returns the process-wide buffer pool only if something already
67/// initialized it; never triggers the virtual reservation itself. Metrics
68/// scrapes read through this so that observing a process (which may never
69/// install the pool) does not mmap the pool's address space as a side
70/// effect.
71pub fn global_pool_peek() -> Option<mz_ore::pool::Pool> {
72    GLOBAL_POOL.get().cloned().flatten()
73}
74
75/// Returns the pool only when [`apply_pool_config`] has installed and
76/// budgeted it, and never initializes anything itself. Chunk spill consumers
77/// resolve their backing through this at each spill decision, so an
78/// unconfigured or gated-off process keeps every chunk resident.
79pub fn active_pool() -> Option<mz_ore::pool::Pool> {
80    if POOL_MODE.load(std::sync::atomic::Ordering::Relaxed) {
81        global_pool_peek()
82    } else {
83        None
84    }
85}
86
87/// Applies a buffer-pool configuration, installing the pool as the process's
88/// spill mechanism. This is the only path that constructs the pool. A process
89/// that never calls it, because the spill gate is off or it is unconfigured,
90/// never reserves the pool's address space or spawns its spill threads.
91/// Returns `false` (and changes nothing) if the pool is unavailable because
92/// its virtual reservation failed.
93///
94/// On success the pool becomes reachable through [`active_pool`] and its
95/// resident budget is retuned in place so live handles stay coherent.
96pub fn apply_pool_config(cfg: PoolPagerConfig) -> bool {
97    let Some(pool) = global_pool() else {
98        return false;
99    };
100    pool.set_budget(cfg.budget_bytes);
101    pool.set_rss_target(cfg.rss_target_bytes);
102    pool.set_spill_threads(cfg.spill_threads);
103    pool.set_eager_backing(cfg.eager_backing);
104    POOL_MODE.store(true, std::sync::atomic::Ordering::Relaxed);
105    true
106}
107
108/// Inputs to [`apply_pool_config`]. All sizes are absolute bytes; fractions
109/// are resolved by the caller against *physical RAM* (see
110/// `mz_ore::memory::physical_memory_bytes`), never against an announced
111/// limit that may include swap.
112#[derive(Clone, Copy, Debug)]
113pub struct PoolPagerConfig {
114    /// Resident-bytes budget for uncompressed slots.
115    pub budget_bytes: usize,
116    /// Spill threads for off-worker eviction I/O (spawn-once).
117    pub spill_threads: usize,
118    /// Whether idle spill threads eagerly compress chunks to
119    /// `BackedResident` ahead of pressure.
120    pub eager_backing: bool,
121    /// Ceiling on the pool's total RSS; the compressed-resident tier is the
122    /// headroom above the budget and warm cap. Zero collapses the tier.
123    pub rss_target_bytes: usize,
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    /// `apply_pool_config` installs the pool: `active_pool` resolves to it
131    /// afterwards, and the configured budget is visible on the instance.
132    /// Mutates process-global state, so it is the only test in this module
133    /// that observes `POOL_MODE`.
134    #[mz_ore::test]
135    #[cfg_attr(miri, ignore)] // unsupported operation: foreign function calls (mmap, madvise)
136    fn apply_pool_config_installs_pool() {
137        let ok = apply_pool_config(PoolPagerConfig {
138            budget_bytes: 1 << 30,
139            spill_threads: 0,
140            eager_backing: false,
141            rss_target_bytes: 0,
142        });
143        assert!(ok, "pool reservation expected to succeed in tests");
144        assert!(active_pool().is_some());
145        assert!(global_pool_peek().is_some());
146    }
147}
mz_timely_util/pool_config.rs

mz_timely_util/
pool_config.rs
