mz_timely_util/

pool_config.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License in the LICENSE file at the
// root of this repository, or online at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Process-wide installation and configuration point for the buffer pool
//! that backs chunk spilling.
//!
//! Chunk spill consumers, arriving with the chunk-batcher work, resolve their
//! backing through [`active_pool`] at each spill decision, so live
//! reconfiguration takes effect on the next call.
//! The pool itself ([`mz_ore::pool::Pool`]) is a process singleton: live
//! chunk handles keep their `Arc` into it, so replacing it on reconfigure
//! would split residency accounting across two budgets. Operator-driven
//! tunes instead retune the one instance in place via [`apply_pool_config`].
//!
//! See `doc/developer/design/20260610_buffer_managed_state.md` for the
//! buffer-pool design.

#![deny(missing_docs)]

use mz_ore::cast::CastLossy;

pub mod metrics;

use mz_ore::cast::CastFrom;
pub use mz_ore::pool::ExtentBackend;

/// Process-wide buffer pool shared by every spill consumer in the process.
///
/// Construction reserves virtual address space only: 1 TiB per chunk size
/// class plus 1 TiB per extent-arena class, tens of TiB in total. Physical
/// memory is paid per resident chunk. On the rare platforms or
/// configurations where the reservation fails, the pool is permanently
/// unavailable for this process and [`apply_pool_config`] reports that by
/// returning `false`.
static GLOBAL_POOL: std::sync::OnceLock<Option<mz_ore::pool::Pool>> = std::sync::OnceLock::new();

/// Whether [`apply_pool_config`] has installed the pool as the process's
/// spill mechanism. [`active_pool`] reads this so consumers stay inert until
/// the first config apply budgets the pool.
static POOL_MODE: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);

/// Whether the pool was asked for a file backend and installed on swap.
static FILE_FALLBACK: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);

/// Whether [`apply_pool_config`] was asked for a file backend and installed
/// the pool on swap instead. Covers installation only: a store that later
/// disables writes on an I/O error reports through
/// `mz_column_pool_extent_file_write_errors_total`.
pub fn file_backend_fell_back() -> bool {
    FILE_FALLBACK.load(std::sync::atomic::Ordering::Relaxed)
}

/// Returns the process-wide buffer pool only if something already
/// initialized it; never triggers the virtual reservation itself. Metrics
/// scrapes read through this so that observing a process (which may never
/// install the pool) does not mmap the pool's address space as a side
/// effect.
pub fn global_pool_peek() -> Option<mz_ore::pool::Pool> {
    GLOBAL_POOL.get().cloned().flatten()
}

/// Returns the pool only when [`apply_pool_config`] has installed and
/// budgeted it, and never initializes anything itself. Chunk spill consumers
/// resolve their backing through this at each spill decision, so an
/// unconfigured or gated-off process keeps every chunk resident.
pub fn active_pool() -> Option<mz_ore::pool::Pool> {
    if POOL_MODE.load(std::sync::atomic::Ordering::Relaxed) {
        global_pool_peek()
    } else {
        None
    }
}

/// Builds the pool for `backend`, creating a file backend's directory. A
/// file store that cannot be created or opened falls back to swap, which
/// needs no scratch space. This is the only place that decides the fallback.
fn build_pool(backend: ExtentBackend) -> Option<mz_ore::pool::Pool> {
    // Logged with the installed backend, so a swap pool tells "file never
    // requested" apart from a fallback.
    let requested_file = matches!(backend, ExtentBackend::File { .. });
    let built = match &backend {
        ExtentBackend::Swap => mz_ore::pool::Pool::new(),
        ExtentBackend::File { dir, .. } => std::fs::create_dir_all(dir)
            .and_then(|()| mz_ore::pool::Pool::with_backend(backend.clone()))
            .or_else(|err| {
                // The error can also come from the pool's own reservation,
                // which the swap attempt then reports again.
                tracing::warn!(%err, ?dir, "buffer pool with a file store failed, trying swap");
                let pool = mz_ore::pool::Pool::new();
                if pool.is_ok() {
                    FILE_FALLBACK.store(true, std::sync::atomic::Ordering::Relaxed);
                }
                pool
            }),
    };
    match built {
        Ok(pool) => {
            tracing::info!(
                requested_file,
                backend = ?pool.backend_kind(),
                "buffer pool installed",
            );
            Some(pool)
        }
        Err(err) => {
            tracing::warn!(
                %err,
                "buffer pool reservation failed; pool spilling unavailable",
            );
            None
        }
    }
}

/// Applies a buffer-pool configuration, installing the pool as the process's
/// spill mechanism. This is the only path that constructs the pool. A process
/// that never calls it, because the spill gate is off or it is unconfigured,
/// never reserves the pool's address space or spawns its spill threads.
/// Returns `false` (and changes nothing) if the pool is unavailable because
/// its virtual reservation failed.
///
/// The first call installs the pool, and its `backend` fixes the extent store
/// for the process. Later calls only retune the installed pool, so a changed
/// backend, or a flag or scratch directory that feeds it, has no effect.
///
/// On success the pool becomes reachable through [`active_pool`] and its
/// resident budget is retuned in place so live handles stay coherent.
pub fn apply_pool_config(cfg: PoolPagerConfig, backend: ExtentBackend) -> bool {
    let applied = apply_to(&GLOBAL_POOL, cfg, backend);
    if applied {
        POOL_MODE.store(true, std::sync::atomic::Ordering::Relaxed);
    }
    applied
}

/// [`apply_pool_config`] over an explicit slot, so tests can observe the
/// install-once behavior without racing on the process-wide pool.
fn apply_to(
    slot: &std::sync::OnceLock<Option<mz_ore::pool::Pool>>,
    cfg: PoolPagerConfig,
    backend: ExtentBackend,
) -> bool {
    let Some(pool) = slot.get_or_init(|| build_pool(backend)) else {
        return false;
    };
    pool.set_budget(cfg.budget_bytes);
    pool.set_rss_target(cfg.rss_target_bytes);
    pool.set_spill_threads(cfg.spill_threads);
    pool.set_eager_backing(cfg.eager_backing);
    pool.set_file_release(
        cfg.file_release_interval,
        u64::cast_from(cfg.file_release_bytes),
    );
    true
}

/// Inputs to [`apply_pool_config`]. All sizes are absolute bytes; fractions
/// are resolved by the caller against *physical RAM* (see
/// `mz_ore::memory::physical_memory_bytes`), never against an announced
/// limit that may include swap.
#[derive(Clone, Copy, Debug)]
pub struct PoolPagerConfig {
    /// Resident-bytes budget for uncompressed slots.
    pub budget_bytes: usize,
    /// Spill threads for off-worker eviction I/O (spawn-once).
    pub spill_threads: usize,
    /// Whether idle spill threads eagerly compress chunks to
    /// `BackedResident` ahead of pressure.
    pub eager_backing: bool,
    /// Ceiling on the pool's total RSS; the compressed-resident tier is the
    /// headroom above the budget and warm cap. Zero collapses the tier.
    pub rss_target_bytes: usize,
    /// How often spill threads release free file store slots.
    pub file_release_interval: std::time::Duration,
    /// Bytes of free file store slots released per extent class and
    /// interval. Zero disables the release.
    pub file_release_bytes: usize,
}

/// Spill threads for a process with `workers` timely workers at `fraction`.
///
/// `floor(workers * fraction)` rounded down to a power of two, and at least
/// one for any positive fraction. A fraction that is not positive yields
/// zero, which evicts inline on the calling thread.
pub fn spill_threads_for_workers(workers: usize, fraction: f64) -> usize {
    if fraction.is_nan() || fraction <= 0.0 {
        return 0;
    }
    let scaled = usize::cast_lossy((f64::cast_lossy(workers) * fraction).floor());
    match scaled.checked_ilog2() {
        Some(log) => 1 << log,
        None => 1,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[mz_ore::test]
    fn spill_threads_for_workers_steps_by_powers_of_two() {
        let at_half = [
            (1, 1),
            (2, 1),
            (4, 2),
            (6, 2),
            (8, 4),
            (12, 4),
            (16, 8),
            (24, 8),
            (31, 8),
            (62, 16),
        ];
        for (workers, threads) in at_half {
            assert_eq!(
                spill_threads_for_workers(workers, 0.5),
                threads,
                "{workers} workers"
            );
        }
        assert_eq!(spill_threads_for_workers(62, 0.0), 0);
        assert_eq!(spill_threads_for_workers(62, -1.0), 0);
        assert_eq!(spill_threads_for_workers(62, f64::NAN), 0);
        assert_eq!(spill_threads_for_workers(1, 0.01), 1);
    }

    const CONFIG: PoolPagerConfig = PoolPagerConfig {
        budget_bytes: 1 << 30,
        spill_threads: 0,
        eager_backing: false,
        rss_target_bytes: 0,
        file_release_interval: std::time::Duration::ZERO,
        file_release_bytes: 0,
    };

    /// `apply_pool_config` installs the pool: `active_pool` resolves to it
    /// afterwards. Mutates process-global state, so it is the only test in
    /// this module that observes `POOL_MODE`.
    #[mz_ore::test]
    #[cfg_attr(miri, ignore)] // unsupported operation: foreign function calls (mmap, madvise)
    fn apply_pool_config_installs_pool() {
        let ok = apply_pool_config(CONFIG, ExtentBackend::Swap);
        assert!(ok, "pool reservation expected to succeed in tests");
        assert!(active_pool().is_some());
        assert!(global_pool_peek().is_some());
    }

    /// The first apply fixes the backend, and later applies never change
    /// it or touch the directory they name. Uses its own slot because the
    /// process-wide pool is installed by whichever test runs first.
    #[mz_ore::test]
    #[cfg_attr(miri, ignore)] // unsupported operation: foreign function calls (mmap, madvise)
    fn first_apply_fixes_backend() {
        let slot = std::sync::OnceLock::new();
        let ok = apply_to(&slot, CONFIG, ExtentBackend::Swap);
        assert!(ok, "pool reservation expected to succeed in tests");

        let dir = tempfile::tempdir().expect("tempdir");
        let pool_dir = dir.path().join("pool");
        let ok = apply_to(
            &slot,
            CONFIG,
            ExtentBackend::File {
                dir: pool_dir.clone(),
                capacity_bytes: None,
            },
        );
        assert!(ok);
        assert!(!pool_dir.exists(), "a later apply creates no directory");
        let pool = slot.get().cloned().flatten().expect("installed");
        assert_eq!(
            pool.backend_kind(),
            mz_ore::pool::BackendKind::Swap,
            "the backend is fixed by the first application",
        );
    }

    /// A file backend whose directory cannot be created installs a swap
    /// pool.
    #[mz_ore::test]
    #[cfg_attr(miri, ignore)] // unsupported operation: foreign function calls (mmap, madvise)
    fn unusable_file_backend_falls_back_to_swap() {
        let file = tempfile::NamedTempFile::new().expect("tempfile");
        let slot = std::sync::OnceLock::new();
        let ok = apply_to(
            &slot,
            CONFIG,
            ExtentBackend::File {
                dir: file.path().join("pool"),
                capacity_bytes: None,
            },
        );
        assert!(ok, "the fallback installs a pool");
        let pool = slot.get().cloned().flatten().expect("installed");
        assert_eq!(pool.backend_kind(), mz_ore::pool::BackendKind::Swap);
        assert!(file_backend_fell_back(), "the fallback is reported");
    }
}