mz_ore/pool/

file.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.

//! The file-backed extent store: a userspace extent allocator over one
//! anonymous file per extent class in a scratch directory.
//!
//! **Files.** Each class of the ladder the caller passes to
//! [`FileStore::open`] gets one file, opened `O_TMPFILE` so it has no name
//! and the kernel frees its blocks when the process exits, crash included.
//! Where the filesystem refuses `O_TMPFILE`, the store creates a uniquely
//! named file and unlinks it immediately. A slot `(class, index)` lives at
//! byte offset `index * class_size` of its class's file.
//!
//! **Block ownership.** Each class allocates slot indices with a
//! [`SlotAllocator`] behind its own mutex. Its warm side holds free slots
//! whose blocks are still allocated, its cold side free slots whose blocks
//! were punched, and slots above its high-water mark were never touched.
//! Freeing a slot pushes it warm with no filesystem call, so slot reuse
//! performs no metadata operation. A slot taken cold or from the high-water
//! mark is `fallocate`d before it is handed out, which turns `ENOSPC` into
//! an allocation failure before any data is in flight.
//!
//! **Capacity.** Every slot with allocated blocks, in use or warm, is
//! charged at its class size against the store's capacity. When an
//! allocation does not fit, the store punches holes over warm slots of the
//! other class holding the most warm bytes, moving them cold, until it
//! fits. Independently, [`FileStore::release_warm`] punches warm slots on
//! its caller's schedule, so the bytes the store holds follow its live data
//! and space freed in one class is available to any other. A slot it punches
//! pays a `fallocate` when next allocated.
//!
//! **Alignment.** Class sizes are page multiples, and every transfer covers
//! whole pages from a page-aligned buffer, which satisfies `O_DIRECT` on
//! devices with 512-byte and 4 KiB logical blocks. The store owns that
//! rounding: [`FileStore::write`] and [`FileStore::read`] return the length
//! they transferred. `open` probes one page-aligned write and read. A probe
//! failing with `EINVAL` means the device needs a larger alignment or the
//! filesystem lacks direct I/O, and the store reopens its files buffered:
//! each write is then followed by a synchronous writeback and a page-cache
//! drop over its range. Memory-backed filesystems are rejected by type,
//! since tmpfs accepts the probe and would hold extents in RAM.
//!
//! **Errors.** Running out of capacity, or `ENOSPC` from `fallocate`, fails
//! the allocation with [`AllocError::Full`]. An `ENOSPC` from a write fails
//! it with [`WriteError::Full`] after the store punches and takes back the
//! slot. A short direct write is retried whole first, and its final error
//! is classified the same way. Either `ENOSPC` additionally lowers the
//! capacity to the bytes allocated at that moment. The volume can be shared
//! with other writers, so [`FileStore::refresh_capacity`] re-derives the
//! capacity from the volume's free space, and room another writer frees
//! becomes usable again. Any other write or `fallocate` error disables
//! writes for the store's lifetime. Data
//! already written stays readable. A read error, short read, or checksum
//! mismatch panics: there is no correct value to return, and the pool's
//! contents are recreatable.

use std::alloc::Layout;
use std::fs::File;
use std::io;
use std::path::{Path, PathBuf};
use std::ptr::NonNull;
use std::sync::Mutex;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::Instant;

use itertools::Itertools;

use crate::cast::CastFrom;
use crate::pool::region::{SlotAllocator, page_size};

/// Where the store's I/O goes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum IoMode {
    /// `O_DIRECT`: transfers bypass the page cache.
    Direct,
    /// Page-cache I/O, written back and dropped after each write.
    Buffered,
}

/// A slot in the store: `index` in the file of extent class `class`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct FileSlot {
    pub(crate) class: usize,
    pub(crate) index: u32,
}

/// Why an allocation failed.
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum AllocError {
    /// The store's capacity, or the filesystem, has no room for the slot.
    Full,
    /// An earlier I/O error disabled writes for the store's lifetime.
    WritesDisabled,
}

/// Why a write failed.
#[derive(Debug)]
pub(crate) enum WriteError {
    /// The filesystem had no room for the slot's data. The store took the
    /// slot back, so the caller must not free it.
    Full,
    /// Writes are disabled for the store's lifetime, by this write's I/O
    /// error or an earlier one. The caller still owns the slot and must free
    /// it.
    Disabled,
}

/// Number of read-latency buckets. Bucket `i < READ_LATENCY_BUCKETS - 1`
/// counts reads that took less than `32 µs << i` and, for `i > 0`, at least
/// `16 µs << i`. The last bucket counts reads of 65.536 ms and more.
pub const READ_LATENCY_BUCKETS: usize = 13;

/// A snapshot of the store's counters.
#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct FileStoreStats {
    pub(crate) writes: u64,
    pub(crate) reads: u64,
    /// Bytes transferred by reads, page-rounded.
    pub(crate) read_bytes: u64,
    pub(crate) read_latency_buckets: [u64; READ_LATENCY_BUCKETS],
    /// I/O errors that disabled writes.
    pub(crate) write_errors: u64,
    pub(crate) holes_punched_bytes: u64,
    /// Bytes of slots with allocated blocks, in use or warm: what the store
    /// holds on the filesystem.
    pub(crate) allocated_bytes: u64,
}

/// Attempts a direct write makes at a slot before a short write counts as
/// failed.
const DIRECT_WRITE_ATTEMPTS: usize = 3;

/// The checksum the store verifies on every read.
pub(crate) fn crc(bytes: &[u8]) -> u32 {
    crc32fast::hash(bytes)
}

/// The bucket of [`FileStoreStats::read_latency_buckets`] for a read that
/// took `micros` microseconds.
fn read_latency_bucket(micros: u64) -> usize {
    // `floor(log2(micros / 16))`, with everything below 32 µs in bucket 0.
    let log2 = (micros / 16).checked_ilog2().unwrap_or(0);
    usize::cast_from(log2).min(READ_LATENCY_BUCKETS - 1)
}

/// One extent class: its file and slot allocator.
#[derive(Debug)]
struct Class {
    size: usize,
    file: File,
    slots: Mutex<SlotAllocator>,
    /// Bytes on the warm side of `slots`, updated under its mutex and read
    /// without it.
    warm_bytes: AtomicU64,
}

impl Class {
    fn slots(&self) -> std::sync::MutexGuard<'_, SlotAllocator> {
        self.slots.lock().expect("file store allocator poisoned")
    }
}

#[derive(Debug, Default)]
struct Counters {
    writes: AtomicU64,
    reads: AtomicU64,
    read_bytes: AtomicU64,
    read_latency_buckets: [AtomicU64; READ_LATENCY_BUCKETS],
    write_errors: AtomicU64,
    holes_punched_bytes: AtomicU64,
}

/// The file-backed extent store. See the module documentation.
#[derive(Debug)]
pub(crate) struct FileStore {
    classes: Vec<Class>,
    io_mode: IoMode,
    page: usize,
    capacity: AtomicU64,
    /// The capacity `open` set. Re-derivation never raises the capacity
    /// above it.
    capacity_limit: u64,
    dir: PathBuf,
    /// Bytes of slots with allocated blocks, in use or warm. Exceeds
    /// `capacity` after `ENOSPC` or a re-derivation lowers the capacity,
    /// until slots are punched.
    allocated: AtomicU64,
    /// Bytes of slots taken off a warm list to be punched and not yet put
    /// back. They are still charged to `allocated`.
    punching: AtomicU64,
    writes_disabled: AtomicBool,
    enospc_seen: AtomicBool,
    counters: Counters,
}

/// The bytes a store may allocate on a volume with `available` bytes free
/// to unprivileged writers: those bytes minus the larger of 1 GiB and 2% of
/// them.
fn usable_bytes(available: u64) -> u64 {
    // Headroom for filesystem metadata of the preallocated extents and for
    // incidental writers sharing the volume.
    available.saturating_sub(std::cmp::max(1 << 30, available / 50))
}

/// The capacity of a new store on a volume with `available` bytes free, per
/// [`usable_bytes`]. `None` if that cannot hold one slot of `largest_class`
/// bytes, the largest extent class.
fn derived_capacity(available: u64, largest_class: usize) -> Option<u64> {
    let capacity = usable_bytes(available);
    (capacity >= u64::cast_from(largest_class)).then_some(capacity)
}

impl FileStore {
    /// Opens one anonymous file per entry of `classes` in `dir`. `classes`
    /// is the ladder of slot sizes, ascending page multiples. `capacity_bytes`
    /// `None` derives capacity from the volume per [`derived_capacity`].
    ///
    /// Fails with [`io::ErrorKind::Unsupported`] when `dir` is on a
    /// memory-backed filesystem or the platform is not Linux, and with
    /// [`io::ErrorKind::StorageFull`] when a derived capacity cannot hold
    /// one slot of the largest class.
    pub(crate) fn open(
        dir: &Path,
        capacity_bytes: Option<u64>,
        classes: &[usize],
    ) -> io::Result<FileStore> {
        if sys::is_memory_backed(dir)? {
            return Err(io::Error::new(
                io::ErrorKind::Unsupported,
                "scratch directory is memory-backed",
            ));
        }
        let page = page_size();
        assert!(
            !classes.is_empty()
                && classes.iter().all(|&size| size > 0 && size % page == 0)
                && classes.is_sorted(),
            "classes must be ascending page multiples: {classes:?}",
        );
        let capacity = match capacity_bytes {
            Some(capacity) => capacity,
            None => {
                let available = sys::available_bytes(dir)?;
                let largest = *classes.last().expect("classes is not empty");
                derived_capacity(available, largest).ok_or_else(|| {
                    io::Error::new(
                        io::ErrorKind::StorageFull,
                        format!(
                            "scratch volume has {available} bytes available, too few for \
                             the file store"
                        ),
                    )
                })?
            }
        };
        let (files, io_mode) = match open_files(dir, classes.len(), page, IoMode::Direct) {
            Ok(files) => (files, IoMode::Direct),
            Err(err) if err.raw_os_error() == Some(libc::EINVAL) => {
                tracing::warn!(
                    dir = %dir.display(),
                    "pool file store: direct I/O unsupported ({err}), using buffered I/O",
                );
                (
                    open_files(dir, classes.len(), page, IoMode::Buffered)?,
                    IoMode::Buffered,
                )
            }
            Err(err) => return Err(err),
        };
        let classes = classes
            .iter()
            .copied()
            .zip_eq(files)
            .map(|(size, file)| Class {
                size,
                file,
                slots: Mutex::new(SlotAllocator::new(u32::MAX)),
                warm_bytes: AtomicU64::new(0),
            })
            .collect();
        Ok(FileStore {
            classes,
            io_mode,
            page,
            capacity: AtomicU64::new(capacity),
            capacity_limit: capacity,
            dir: dir.to_path_buf(),
            allocated: AtomicU64::new(0),
            punching: AtomicU64::new(0),
            writes_disabled: AtomicBool::new(false),
            enospc_seen: AtomicBool::new(false),
            counters: Counters::default(),
        })
    }

    pub(crate) fn io_mode(&self) -> IoMode {
        self.io_mode
    }

    /// The class for a stored payload of `comp_len` bytes, `None` if none fits.
    pub(crate) fn class_for(&self, comp_len: usize) -> Option<usize> {
        self.classes.iter().position(|c| c.size >= comp_len)
    }

    pub(crate) fn class_size(&self, class: usize) -> usize {
        self.classes[class].size
    }

    /// Reserves a slot with allocated blocks.
    pub(crate) fn alloc(&self, class: usize) -> Result<FileSlot, AllocError> {
        if self.writes_disabled() {
            return Err(AllocError::WritesDisabled);
        }
        let c = &self.classes[class];
        let size = u64::cast_from(c.size);
        let allocated = {
            let mut slots = c.slots();
            let allocated = slots.alloc();
            if let Some((_, true)) = allocated {
                c.warm_bytes.fetch_sub(size, Ordering::Relaxed);
            }
            allocated
        };
        let Some((index, warm)) = allocated else {
            // Every `u32` index of the class is in use: `u32::MAX` slots,
            // 16 TiB at a one-page class on a 4 KiB page.
            return Err(AllocError::Full);
        };
        let slot = FileSlot { class, index };
        if warm {
            return Ok(slot);
        }
        // A cold or never-touched slot has no blocks. A concurrent
        // allocation can take the room a punch made, so punch again until
        // the reservation succeeds or nothing is left to punch.
        loop {
            if self.reserve(size) {
                break;
            }
            if self.punch_for(class, size) {
                continue;
            }
            // Slots another thread is punching, including this class's warm
            // slots, stay charged until each punch returns. Failing now
            // would report `Full` with room about to free.
            if self.punching.load(Ordering::Acquire) > 0 {
                std::thread::yield_now();
                continue;
            }
            // A punch that finished after the failed reservation released
            // its bytes before it cleared `punching`.
            if self.reserve(size) {
                break;
            }
            c.slots().free(index, false);
            return Err(AllocError::Full);
        }
        match sys::fallocate(&c.file, self.offset(slot), c.size) {
            Ok(()) => Ok(slot),
            Err(err) => {
                // A failed `fallocate` can leave part of the range
                // allocated. Punching returns those blocks. If the punch
                // fails too, the slot goes warm and stays charged, so its
                // blocks stay counted.
                match sys::punch_hole(&c.file, self.offset(slot), c.size) {
                    Ok(()) => {
                        self.allocated.fetch_sub(size, Ordering::Relaxed);
                        self.counters
                            .holes_punched_bytes
                            .fetch_add(size, Ordering::Relaxed);
                        c.slots().free(index, false);
                    }
                    Err(punch_err) => {
                        tracing::warn!(
                            "pool file store: punching a slot after a failed fallocate failed: \
                             {punch_err}"
                        );
                        let mut slots = c.slots();
                        slots.free(index, true);
                        c.warm_bytes.fetch_add(size, Ordering::Relaxed);
                    }
                }
                if err.raw_os_error() == Some(libc::ENOSPC) {
                    self.lower_capacity("fallocate");
                    Err(AllocError::Full)
                } else {
                    self.disable_writes("fallocate", &err);
                    Err(AllocError::WritesDisabled)
                }
            }
        }
    }

    /// Returns a slot. Its blocks stay allocated. Must not be called while a
    /// read or write of the slot is in flight, since the slot can be punched
    /// or handed out again at once.
    pub(crate) fn free(&self, slot: FileSlot) {
        let c = &self.classes[slot.class];
        let mut slots = c.slots();
        slots.free(slot.index, true);
        c.warm_bytes
            .fetch_add(u64::cast_from(c.size), Ordering::Relaxed);
    }

    /// Writes `buf[..round_up(len, page)]` to `slot` and returns the rounded
    /// length. `buf` must be at least that long, and page aligned in
    /// [`IoMode::Direct`]. Fails without I/O once writes are disabled. On
    /// `ENOSPC`, takes the slot back and fails with [`WriteError::Full`]. On
    /// any other I/O error, disables further writes for the store's lifetime.
    pub(crate) fn write(
        &self,
        slot: FileSlot,
        buf: &[u8],
        len: usize,
    ) -> Result<usize, WriteError> {
        if self.writes_disabled() {
            return Err(WriteError::Disabled);
        }
        let c = &self.classes[slot.class];
        let n = len.next_multiple_of(self.page);
        assert!(
            n <= c.size,
            "write of {len} bytes exceeds class size {}",
            c.size
        );
        assert!(
            n <= buf.len(),
            "write of {n} bytes from a {}-byte buffer",
            buf.len()
        );
        if self.io_mode == IoMode::Direct {
            assert_eq!(
                buf.as_ptr().addr() % self.page,
                0,
                "write buffer is not page aligned"
            );
        }
        let offset = self.offset(slot);
        let result =
            write_all(&c.file, &buf[..n], offset, self.io_mode).and_then(|()| match self.io_mode {
                IoMode::Direct => Ok(()),
                IoMode::Buffered => sys::writeback_and_drop(&c.file, offset, n),
            });
        match result {
            Ok(()) => {
                self.counters.writes.fetch_add(1, Ordering::Relaxed);
                Ok(n)
            }
            Err(err) if err.raw_os_error() == Some(libc::ENOSPC) => {
                self.reclaim_after_enospc(slot);
                Err(WriteError::Full)
            }
            Err(err) => {
                self.disable_writes("write", &err);
                Err(WriteError::Disabled)
            }
        }
    }

    /// Takes back `slot` after a write to it hit `ENOSPC`, and lowers the
    /// capacity.
    fn reclaim_after_enospc(&self, slot: FileSlot) {
        let c = &self.classes[slot.class];
        let size = u64::cast_from(c.size);
        // A filesystem that allocates on overwrite (copy-on-write, or thin
        // provisioning) can fail the write despite the `fallocate`. Its
        // blocks are then in an unknown state, and a warm slot would hand
        // the same failure to the next demotion of this class. Punching
        // returns the blocks and makes the slot cold. If the punch fails,
        // the slot goes on no free list and stays charged, so it is never
        // reused and its blocks stay counted.
        match sys::punch_hole(&c.file, self.offset(slot), c.size) {
            Ok(()) => {
                c.slots().free(slot.index, false);
                self.allocated.fetch_sub(size, Ordering::Relaxed);
                self.counters
                    .holes_punched_bytes
                    .fetch_add(size, Ordering::Relaxed);
            }
            Err(err) => {
                tracing::warn!(
                    "pool file store: punching a slot after ENOSPC failed, leaking it: {err}"
                );
            }
        }
        self.lower_capacity("write");
    }

    /// Lowers the capacity to the bytes allocated now, after `op` hit
    /// `ENOSPC`.
    fn lower_capacity(&self, op: &str) {
        // The headroom was too small for this filesystem. The bytes
        // allocated now are evidently what fits.
        let allocated = self.allocated.load(Ordering::Relaxed);
        self.capacity.fetch_min(allocated, Ordering::Relaxed);
        if !self.enospc_seen.swap(true, Ordering::Relaxed) {
            tracing::warn!(
                capacity = allocated,
                "pool file store: {op} hit ENOSPC, lowering capacity",
            );
        }
    }

    /// Reads the slot's first `len` stored bytes into `dst` (resized to
    /// `len`), verifying `crc`, and returns the length transferred, `len`
    /// rounded up to a page. Panics on I/O error, short read, or checksum
    /// mismatch.
    pub(crate) fn read(&self, slot: FileSlot, len: usize, crc: u32, dst: &mut AlignedBuf) -> usize {
        let c = &self.classes[slot.class];
        let n = len.next_multiple_of(self.page);
        assert!(
            n <= c.size,
            "read of {len} bytes exceeds class size {}",
            c.size
        );
        let offset = self.offset(slot);
        dst.resize(n);
        let start = Instant::now();
        let mut done = 0;
        while done < n {
            match sys::pread(
                &c.file,
                &mut dst.as_mut_slice()[done..],
                offset + u64::cast_from(done),
            ) {
                Ok(0) => panic!(
                    "pool file store: short read: class {}, index {}, offset {offset}, \
                     {done} of {n} bytes",
                    slot.class, slot.index,
                ),
                // A direct read resumed at `offset + done` would be
                // unaligned unless `done` is, and fail `EINVAL`.
                Ok(k) if self.io_mode == IoMode::Direct && done + k < n => panic!(
                    "pool file store: short read: class {}, index {}, offset {offset}, \
                     {} of {n} bytes",
                    slot.class,
                    slot.index,
                    done + k,
                ),
                Ok(k) => done += k,
                Err(err) if err.kind() == io::ErrorKind::Interrupted => {}
                Err(err) => panic!(
                    "pool file store: read failed: class {}, index {}, offset {offset}: {err}",
                    slot.class, slot.index,
                ),
            }
        }
        let micros = u64::try_from(start.elapsed().as_micros()).unwrap_or(u64::MAX);
        if self.io_mode == IoMode::Buffered {
            // The read left the range in the page cache, charged to the
            // process's cgroup. A failed advice leaves the read correct.
            let _ = sys::drop_cache(&c.file, offset, n);
        }
        dst.truncate(len);
        let observed = self::crc(dst.as_slice());
        assert!(
            observed == crc,
            "pool file store: checksum mismatch: class {}, index {}, offset {offset}, \
             expected {crc:#010x}, observed {observed:#010x}",
            slot.class,
            slot.index,
        );
        self.counters.reads.fetch_add(1, Ordering::Relaxed);
        self.counters
            .read_bytes
            .fetch_add(u64::cast_from(n), Ordering::Relaxed);
        self.counters.read_latency_buckets[read_latency_bucket(micros)]
            .fetch_add(1, Ordering::Relaxed);
        n
    }

    pub(crate) fn capacity_bytes(&self) -> u64 {
        self.capacity.load(Ordering::Relaxed)
    }

    pub(crate) fn writes_disabled(&self) -> bool {
        self.writes_disabled.load(Ordering::Relaxed)
    }

    /// Whether an `alloc` of `class` could currently succeed without I/O
    /// errors: writes enabled and either a warm free slot, capacity room, or
    /// room after punching other classes' warm slots exists. Takes no lock.
    pub(crate) fn can_alloc(&self, class: usize) -> bool {
        if self.writes_disabled() {
            return false;
        }
        let c = &self.classes[class];
        let size = u64::cast_from(c.size);
        if c.warm_bytes.load(Ordering::Relaxed) > 0 || self.fits(size) {
            return true;
        }
        let punchable: u64 = self
            .classes
            .iter()
            .enumerate()
            .filter(|&(i, _)| i != class)
            .map(|(_, c)| c.warm_bytes.load(Ordering::Relaxed))
            .sum::<u64>()
            .saturating_add(self.punching.load(Ordering::Relaxed));
        self.allocated
            .load(Ordering::Relaxed)
            .saturating_sub(punchable)
            .saturating_add(size)
            <= self.capacity_bytes()
    }

    /// Whether [`FileStore::can_alloc`] holds for some class.
    pub(crate) fn can_alloc_any(&self) -> bool {
        (0..self.classes.len()).any(|class| self.can_alloc(class))
    }

    /// Punches warm slots of every class, up to `max_bytes` per class
    /// rounded up to a whole slot, and returns the bytes punched.
    pub(crate) fn release_warm(&self, max_bytes: u64) -> u64 {
        (0..self.classes.len())
            .map(|class| self.punch_warm(class, max_bytes).0)
            .sum()
    }

    /// Sets the capacity to the bytes allocated plus the volume's
    /// [`usable_bytes`], at most the capacity `open` set. Keeps the capacity
    /// if the volume cannot be queried.
    pub(crate) fn refresh_capacity(&self) {
        let Ok(available) = sys::available_bytes(&self.dir) else {
            return;
        };
        // The store's own blocks are absent from `available`, so they are
        // added back. A concurrent `ENOSPC` can lower the capacity between
        // the load and the store, and is then overwritten by a value derived
        // from free space observed just before it, which the next `ENOSPC`
        // lowers again.
        let capacity = self
            .allocated
            .load(Ordering::Relaxed)
            .saturating_add(usable_bytes(available))
            .min(self.capacity_limit);
        self.capacity.store(capacity, Ordering::Relaxed);
    }

    pub(crate) fn stats(&self) -> FileStoreStats {
        let c = &self.counters;
        FileStoreStats {
            writes: c.writes.load(Ordering::Relaxed),
            reads: c.reads.load(Ordering::Relaxed),
            read_bytes: c.read_bytes.load(Ordering::Relaxed),
            read_latency_buckets: std::array::from_fn(|i| {
                c.read_latency_buckets[i].load(Ordering::Relaxed)
            }),
            write_errors: c.write_errors.load(Ordering::Relaxed),
            holes_punched_bytes: c.holes_punched_bytes.load(Ordering::Relaxed),
            allocated_bytes: self.allocated.load(Ordering::Relaxed),
        }
    }

    fn offset(&self, slot: FileSlot) -> u64 {
        u64::cast_from(slot.index) * u64::cast_from(self.classes[slot.class].size)
    }

    fn fits(&self, size: u64) -> bool {
        self.allocated.load(Ordering::Relaxed).saturating_add(size) <= self.capacity_bytes()
    }

    /// Charges `size` bytes against capacity, or returns `false` if they do
    /// not fit.
    fn reserve(&self, size: u64) -> bool {
        self.allocated
            .try_update(Ordering::Relaxed, Ordering::Relaxed, |allocated| {
                allocated
                    .checked_add(size)
                    .filter(|&n| n <= self.capacity_bytes())
            })
            .is_ok()
    }

    /// Punches warm slots of classes other than `class` until `size` more
    /// bytes fit or no other class has a warm slot, and returns whether they
    /// fit.
    fn punch_for(&self, class: usize, size: u64) -> bool {
        while !self.fits(size) {
            // Choose from the lock-free warm-bytes snapshot: holding two
            // class mutexes at once could deadlock against a concurrent
            // punch.
            let victim = self
                .classes
                .iter()
                .enumerate()
                .filter(|&(i, _)| i != class)
                .map(|(i, c)| (c.warm_bytes.load(Ordering::Relaxed), i))
                .filter(|&(warm_bytes, _)| warm_bytes > 0)
                .max();
            let Some((_, victim)) = victim else {
                return false;
            };
            let shortfall = self
                .allocated
                .load(Ordering::Relaxed)
                .saturating_add(size)
                .saturating_sub(self.capacity_bytes());
            let (_, failed) = self.punch_warm(victim, shortfall);
            if failed {
                return false;
            }
        }
        true
    }

    /// Punches warm slots of `class` until at least `want` bytes are punched
    /// or none is left, moving them cold. Returns the bytes punched and
    /// whether a punch failed, which leaves the failed slot and the ones not
    /// yet punched warm.
    fn punch_warm(&self, class: usize, want: u64) -> (u64, bool) {
        let c = &self.classes[class];
        let class_size = u64::cast_from(c.size);
        // Take the victims off the warm list under the lock and punch them
        // after releasing it. `free` takes this mutex on workers holding a
        // chunk's state lock, so it must not wait out the punches. A popped
        // slot is on no free list, so nothing can allocate it meanwhile.
        let popped: Vec<u32> = {
            let mut slots = c.slots();
            let mut popped = Vec::new();
            while u64::cast_from(popped.len()).saturating_mul(class_size) < want {
                let Some(index) = slots.pop_oldest_warm() else {
                    break;
                };
                c.warm_bytes.fetch_sub(class_size, Ordering::Relaxed);
                popped.push(index);
            }
            // Under the lock, so an `alloc` never sees the slots neither
            // warm nor punching.
            self.punching.fetch_add(
                u64::cast_from(popped.len()).saturating_mul(class_size),
                Ordering::Relaxed,
            );
            popped
        };
        let mut punched = 0;
        for &index in &popped {
            let slot = FileSlot { class, index };
            if let Err(err) = sys::punch_hole(&c.file, self.offset(slot), c.size) {
                tracing::warn!("pool file store: punching a free slot failed: {err}");
                break;
            }
            punched += 1;
            self.allocated.fetch_sub(class_size, Ordering::Relaxed);
            // Release orders the `allocated` decrement before it for
            // `alloc`, which reads `punching` with acquire.
            self.punching.fetch_sub(class_size, Ordering::Release);
            self.counters
                .holes_punched_bytes
                .fetch_add(class_size, Ordering::Relaxed);
        }
        let (cold, warm) = popped.split_at(punched);
        let mut slots = c.slots();
        for &index in cold {
            slots.free(index, false);
        }
        // Back to the old end, in their original order, so the next release
        // retries them first and `alloc` does not prefer them.
        for &index in warm.iter().rev() {
            slots.return_oldest_warm(index);
            c.warm_bytes.fetch_add(class_size, Ordering::Relaxed);
        }
        self.punching.fetch_sub(
            u64::cast_from(warm.len()).saturating_mul(class_size),
            Ordering::Release,
        );
        (
            u64::cast_from(punched).saturating_mul(class_size),
            !warm.is_empty(),
        )
    }

    fn disable_writes(&self, op: &str, err: &io::Error) {
        self.counters.write_errors.fetch_add(1, Ordering::Relaxed);
        if !self.writes_disabled.swap(true, Ordering::Relaxed) {
            tracing::warn!("pool file store: {op} failed, disabling writes: {err}");
        }
    }
}

/// Opens `classes` anonymous files for `io_mode` and probes the first. In
/// [`IoMode::Direct`], an `EINVAL` from the open or the probe means direct
/// I/O is unusable here.
fn open_files(dir: &Path, classes: usize, page: usize, io_mode: IoMode) -> io::Result<Vec<File>> {
    let direct = io_mode == IoMode::Direct;
    let files = (0..classes)
        .map(|class| sys::open_anonymous(dir, class, direct))
        .collect::<io::Result<Vec<_>>>()?;
    probe(&files[0], page, io_mode)?;
    Ok(files)
}

/// Allocates, writes, reads back, and punches one page at offset 0.
fn probe(file: &File, page: usize, io_mode: IoMode) -> io::Result<()> {
    let pattern: Vec<u8> = (0..page)
        .map(|i| u8::try_from(i % 251).expect("fits"))
        .collect();
    let src = AlignedBuf::from_bytes(&pattern);
    sys::fallocate(file, 0, page)?;
    write_all(file, src.as_slice(), 0, io_mode)?;
    let mut dst = AlignedBuf::new();
    dst.resize(page);
    let mut done = 0;
    while done < page {
        match sys::pread(file, &mut dst.as_mut_slice()[done..], u64::cast_from(done)) {
            Ok(0) => return Err(io::Error::from(io::ErrorKind::UnexpectedEof)),
            Ok(k) => done += k,
            Err(err) if err.kind() == io::ErrorKind::Interrupted => {}
            Err(err) => return Err(err),
        }
    }
    if dst.as_slice() != src.as_slice() {
        return Err(io::Error::other("probe read back different bytes"));
    }
    sys::punch_hole(file, 0, page)
}

/// Writes all of `buf` at `offset`, retrying on `EINTR` and short writes.
/// In [`IoMode::Direct`], a short write restarts the whole buffer at
/// `offset`, at most [`DIRECT_WRITE_ATTEMPTS`] times, and the final attempt's
/// error is returned.
fn write_all(file: &File, buf: &[u8], offset: u64, io_mode: IoMode) -> io::Result<()> {
    match io_mode {
        IoMode::Direct => {
            // Resuming at `offset + done` would be a direct write that is
            // aligned only if `done` is, which the kernel does not promise,
            // and an unaligned direct write fails `EINVAL`, which disables
            // writes. Rewriting the same bytes is idempotent, and when the
            // short write came from a full filesystem the retry returns the
            // `ENOSPC` the caller classifies.
            let mut attempts = 0;
            loop {
                match sys::pwrite(file, buf, offset) {
                    Ok(k) if k == buf.len() => return Ok(()),
                    Ok(_) => {
                        attempts += 1;
                        if attempts == DIRECT_WRITE_ATTEMPTS {
                            return Err(io::Error::from(io::ErrorKind::WriteZero));
                        }
                    }
                    Err(err) if err.kind() == io::ErrorKind::Interrupted => {}
                    Err(err) => return Err(err),
                }
            }
        }
        IoMode::Buffered => {
            let mut done = 0;
            while done < buf.len() {
                match sys::pwrite(file, &buf[done..], offset + u64::cast_from(done)) {
                    Ok(0) => return Err(io::Error::from(io::ErrorKind::WriteZero)),
                    Ok(k) => done += k,
                    Err(err) if err.kind() == io::ErrorKind::Interrupted => {}
                    Err(err) => return Err(err),
                }
            }
            Ok(())
        }
    }
}

/// Page-aligned, growable byte buffer for `O_DIRECT` reads.
#[derive(Debug)]
pub(crate) struct AlignedBuf {
    /// Valid for `cap` initialized bytes, or dangling when `cap == 0`.
    ptr: NonNull<u8>,
    len: usize,
    cap: usize,
}

// SAFETY: `AlignedBuf` exclusively owns its allocation, like `Vec<u8>`.
unsafe impl Send for AlignedBuf {}

impl AlignedBuf {
    pub(crate) const fn new() -> AlignedBuf {
        AlignedBuf {
            ptr: NonNull::dangling(),
            len: 0,
            cap: 0,
        }
    }

    pub(crate) fn as_slice(&self) -> &[u8] {
        // SAFETY: `ptr` is valid for `cap >= len` initialized bytes, or
        // dangling and well aligned with `len == 0`.
        unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
    }

    fn as_mut_slice(&mut self) -> &mut [u8] {
        // SAFETY: as in `as_slice`, and `&mut self` makes the borrow unique.
        unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
    }

    /// Releases the allocation if its capacity exceeds `max` bytes, which
    /// also empties the buffer.
    pub(crate) fn shrink_above(&mut self, max: usize) {
        if self.cap > max {
            self.release();
        }
    }

    /// Sets the length to `len`, growing capacity to a page multiple if
    /// needed. Contents are unspecified but initialized.
    fn resize(&mut self, len: usize) {
        if len > self.cap {
            self.release();
            let cap = len.next_multiple_of(page_size());
            let layout = Self::layout(cap);
            // SAFETY: `layout` has nonzero size because `len > 0`. Zeroing
            // keeps every byte up to `cap` initialized.
            let ptr = unsafe { std::alloc::alloc_zeroed(layout) };
            self.ptr = NonNull::new(ptr).unwrap_or_else(|| std::alloc::handle_alloc_error(layout));
            self.cap = cap;
        }
        self.len = len;
    }

    fn truncate(&mut self, len: usize) {
        self.len = self.len.min(len);
    }

    fn release(&mut self) {
        if self.cap > 0 {
            // SAFETY: `ptr` was allocated in `resize` with exactly this
            // layout, and no borrow outlives `&mut self`.
            unsafe { std::alloc::dealloc(self.ptr.as_ptr(), Self::layout(self.cap)) };
        }
        // NOTE: assigning `*self` would drop the old value and recurse
        // through `Drop`.
        self.ptr = NonNull::dangling();
        self.len = 0;
        self.cap = 0;
    }

    fn layout(cap: usize) -> Layout {
        Layout::from_size_align(cap, page_size()).expect("valid aligned buffer layout")
    }

    /// A buffer holding a copy of `bytes`.
    pub(crate) fn from_bytes(bytes: &[u8]) -> AlignedBuf {
        let mut buf = AlignedBuf::new();
        buf.resize(bytes.len());
        buf.as_mut_slice().copy_from_slice(bytes);
        buf
    }
}

impl Drop for AlignedBuf {
    fn drop(&mut self) {
        self.release();
    }
}

mod sys;

#[cfg(test)]
pub(crate) mod testing;

#[cfg(all(test, target_os = "linux"))]
mod tests;