spill

server-rs/src/data/spill.rs

@@ -0,0 +1,330 @@
+//! Optional disk-backed storage for the large flat arrays in [`PropertyData`].
+//!
+//! In production every property array is held as an owned `Vec` in RAM — fastest,
+//! but the feature matrix plus the flat address-search arrays are ~4GB, which a
+//! memory-constrained dev box can't hold reliably. When a spill directory is
+//! configured (the `--spill-dir` dev flag), each large array is instead written
+//! to an anonymous file in that directory and memory-mapped read-only. The mapped
+//! pages are file-backed and clean, so under memory pressure the kernel evicts
+//! them (re-faulting from disk on next touch) rather than the process being
+//! OOM-killed — the same "let the kernel page it" trade-off the PMTiles reader
+//! already makes. The backing files are unlinked immediately after creation, so
+//! they leave nothing on disk and are reclaimed when the map drops at shutdown.
+//!
+//! [`SpillVec<T>`] dereferences to `[T]`, so every read site is byte-identical
+//! whether the data lives on the heap or in a map. With no `--spill-dir` (prod),
+//! the `Owned` variant is a thin wrapper over `Vec<T>` and the hot paths are
+//! unchanged.
+
+use std::fs::{File, OpenOptions};
+use std::marker::PhantomData;
+use std::path::Path;
+use std::sync::atomic::{AtomicU64, Ordering};
+
+use anyhow::Context;
+use memmap2::{Mmap, MmapMut};
+
+/// Marker for types that may back a [`SpillVec`]: fixed-size `Copy` values that
+/// round-trip losslessly through their raw bytes.
+///
+/// # Safety
+/// Implementors must be `Copy`, contain no padding, and accept — for any value we
+/// could have stored — the exact bytes we wrote when reinterpreted as `Self`. We
+/// only ever read back bytes produced from genuine `Self` values, so niche types
+/// such as `lasso::Spur` (a `NonZeroU32`) are sound here: the bytes always come
+/// from real, non-zero interner keys and are never reinterpreted from arbitrary
+/// input.
+pub unsafe trait SpillElem: Copy + 'static {}
+
+// SAFETY: plain integers have no padding and every bit pattern is a valid value.
+unsafe impl SpillElem for u8 {}
+unsafe impl SpillElem for u16 {}
+unsafe impl SpillElem for u32 {}
+unsafe impl SpillElem for f32 {}
+// SAFETY: `lasso::Spur` is a transparent wrapper over `NonZeroU32` (size/align 4).
+// We only store keys returned by an interner, which are never zero, so the bytes
+// we map back never form the one invalid (all-zero) pattern.
+unsafe impl SpillElem for lasso::Spur {}
+
+/// Process-unique suffix so concurrently-created spill files never collide before
+/// they are unlinked.
+static NEXT_SPILL_ID: AtomicU64 = AtomicU64::new(0);
+
+/// A read-only slice of `T` backed by either an owned heap `Vec` (production) or a
+/// memory-mapped spill file (dev). Dereferences to `[T]`.
+pub enum SpillVec<T: SpillElem> {
+    Owned(Vec<T>),
+    Mapped {
+        map: Mmap,
+        len: usize,
+        _marker: PhantomData<T>,
+    },
+}
+
+impl<T: SpillElem> SpillVec<T> {
+    /// Keep `values` on the heap (production behaviour, and the only path used in
+    /// tests).
+    pub fn owned(values: Vec<T>) -> Self {
+        SpillVec::Owned(values)
+    }
+
+    /// Spill `values` to `dir` and memory-map it when `dir` is `Some`; otherwise
+    /// keep it on the heap. `label` names the backing file for diagnostics.
+    pub fn maybe_spill(values: Vec<T>, dir: Option<&Path>, label: &str) -> anyhow::Result<Self> {
+        match dir {
+            None => Ok(SpillVec::owned(values)),
+            Some(dir) => spill_vec(values, dir, label),
+        }
+    }
+
+    /// Borrow the contents as a plain slice. (`SpillVec` also derefs to `[T]`, so
+    /// `&spill_vec` coerces to `&[T]`; this is for the rare spot where coercion
+    /// doesn't fire, e.g. building a tuple.)
+    pub fn as_slice(&self) -> &[T] {
+        self
+    }
+}
+
+impl SpillVec<u16> {
+    /// Build a `u16` array of `len` elements by filling it in place — on the heap
+    /// when `dir` is `None`, or directly inside an mmap-backed spill file when
+    /// `Some`, so the (large) buffer never simultaneously exists on the heap. Every
+    /// element is initialized to `default` first; `fill` then overwrites the cells
+    /// it owns. Used for the feature matrix, whose 3GB heap copy we want to avoid
+    /// when spilling.
+    pub fn build_u16(
+        len: usize,
+        default: u16,
+        dir: Option<&Path>,
+        label: &str,
+        fill: impl FnOnce(&mut [u16]),
+    ) -> anyhow::Result<Self> {
+        match dir {
+            Some(dir) if len > 0 => {
+                let byte_len = len * std::mem::size_of::<u16>();
+                let file = anon_file(dir, label)?;
+                allocate_spill_file(&file, byte_len, label)?;
+                // SAFETY: `file` is a freshly-created, exclusively-owned regular
+                // file sized to exactly `byte_len`; no other mapping aliases it.
+                let mut map = unsafe { MmapMut::map_mut(&file) }
+                    .with_context(|| format!("mapping spill file for '{label}'"))?;
+                // SAFETY: an mmap base is page-aligned (so aligned for `u16`) and
+                // the mapping is exactly `len * 2` bytes, valid to view as `len`
+                // `u16`s.
+                let slice =
+                    unsafe { std::slice::from_raw_parts_mut(map.as_mut_ptr().cast::<u16>(), len) };
+                slice.fill(default);
+                fill(slice);
+                let map = map
+                    .make_read_only()
+                    .with_context(|| format!("sealing spill file for '{label}'"))?;
+                Ok(SpillVec::Mapped {
+                    map,
+                    len,
+                    _marker: PhantomData,
+                })
+            }
+            _ => {
+                let mut values = vec![default; len];
+                fill(&mut values);
+                Ok(SpillVec::owned(values))
+            }
+        }
+    }
+}
+
+impl<T: SpillElem> std::ops::Deref for SpillVec<T> {
+    type Target = [T];
+
+    fn deref(&self) -> &[T] {
+        match self {
+            SpillVec::Owned(values) => values.as_slice(),
+            SpillVec::Mapped { map, len, .. } => {
+                let bytes: &[u8] = map;
+                debug_assert_eq!(bytes.len(), len * std::mem::size_of::<T>());
+                debug_assert_eq!(bytes.as_ptr().align_offset(std::mem::align_of::<T>()), 0);
+                // SAFETY: the mapping holds exactly `len * size_of::<T>()` bytes,
+                // all written from a `&[T]`; its base is page-aligned (hence aligned
+                // for `T`); `T: SpillElem` guarantees those bytes reinterpret as
+                // valid `T`s; and the read-only map outlives the borrow of `self`.
+                unsafe { std::slice::from_raw_parts(bytes.as_ptr().cast::<T>(), *len) }
+            }
+        }
+    }
+}
+
+fn spill_vec<T: SpillElem>(values: Vec<T>, dir: &Path, label: &str) -> anyhow::Result<SpillVec<T>> {
+    let len = values.len();
+    let byte_len = len * std::mem::size_of::<T>();
+    // An empty mapping is invalid; nothing to spill, so keep the empty Vec.
+    if byte_len == 0 {
+        return Ok(SpillVec::owned(values));
+    }
+
+    let file = anon_file(dir, label)?;
+    allocate_spill_file(&file, byte_len, label)?;
+    // SAFETY: `file` is a freshly-created, exclusively-owned regular file sized to
+    // exactly `byte_len`; no other mapping aliases it.
+    let mut map = unsafe { MmapMut::map_mut(&file) }
+        .with_context(|| format!("mapping spill file for '{label}' ({byte_len} bytes)"))?;
+    // SAFETY: `values` is a live `Vec<T>` of `Copy`, padding-free elements, so its
+    // store is `byte_len` initialized bytes valid to read as `u8`.
+    let src = unsafe { std::slice::from_raw_parts(values.as_ptr().cast::<u8>(), byte_len) };
+    map.copy_from_slice(src);
+    // The bytes now live in the file's page cache; release the heap copy.
+    drop(values);
+    let map = map
+        .make_read_only()
+        .with_context(|| format!("sealing spill file for '{label}'"))?;
+    Ok(SpillVec::Mapped {
+        map,
+        len,
+        _marker: PhantomData,
+    })
+}
+
+/// Size a freshly-created spill file to `byte_len` bytes, reserving the disk
+/// blocks up front. We write the array through a mutable mmap, and writing to a
+/// dirty mmap page the kernel can't back at writeback raises `SIGBUS` — so a
+/// plain sparse `set_len` would turn an out-of-space dev disk into a crash. On
+/// Linux `posix_fallocate` allocates the blocks now, surfacing the shortfall as a
+/// clean `ENOSPC` error here instead. Elsewhere we fall back to `set_len`.
+fn allocate_spill_file(file: &File, byte_len: usize, label: &str) -> anyhow::Result<()> {
+    #[cfg(target_os = "linux")]
+    {
+        use std::os::fd::AsRawFd;
+        // `posix_fallocate` returns an errno value directly (0 on success) and,
+        // unlike most libc calls, does not set the global errno.
+        let ret = unsafe { libc::posix_fallocate(file.as_raw_fd(), 0, byte_len as libc::off_t) };
+        if ret != 0 {
+            return Err(std::io::Error::from_raw_os_error(ret)).with_context(|| {
+                format!("reserving {byte_len} bytes of disk for spill file '{label}'")
+            });
+        }
+        Ok(())
+    }
+    #[cfg(not(target_os = "linux"))]
+    {
+        file.set_len(byte_len as u64)
+            .with_context(|| format!("sizing spill file for '{label}'"))
+    }
+}
+
+/// Create a spill file in `dir` and immediately unlink it. On Unix the open
+/// descriptor (and the mapping built from it) keep the inode alive, so the file is
+/// invisible in the directory and its blocks are reclaimed automatically when the
+/// map is dropped — no cleanup, no leftovers across runs.
+fn anon_file(dir: &Path, label: &str) -> anyhow::Result<File> {
+    std::fs::create_dir_all(dir)
+        .with_context(|| format!("creating spill directory {}", dir.display()))?;
+    let unique = NEXT_SPILL_ID.fetch_add(1, Ordering::Relaxed);
+    let path = dir.join(format!(
+        ".ppc-spill-{}-{label}-{unique}.bin",
+        std::process::id()
+    ));
+    let file = OpenOptions::new()
+        .read(true)
+        .write(true)
+        .create_new(true)
+        .open(&path)
+        .with_context(|| format!("creating spill file {}", path.display()))?;
+    #[cfg(unix)]
+    let _ = std::fs::remove_file(&path);
+    Ok(file)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A scratch directory under the system temp dir; cleaned up on drop.
+    struct TempDir(std::path::PathBuf);
+
+    impl TempDir {
+        fn new(tag: &str) -> Self {
+            let id = NEXT_SPILL_ID.fetch_add(1, Ordering::Relaxed);
+            let dir = std::env::temp_dir()
+                .join(format!("ppc-spill-test-{}-{tag}-{id}", std::process::id()));
+            std::fs::create_dir_all(&dir).expect("create temp dir");
+            TempDir(dir)
+        }
+        fn path(&self) -> &Path {
+            &self.0
+        }
+    }
+
+    impl Drop for TempDir {
+        fn drop(&mut self) {
+            let _ = std::fs::remove_dir_all(&self.0);
+        }
+    }
+
+    #[test]
+    fn maybe_spill_roundtrips_identically_owned_and_mapped() {
+        let values: Vec<u32> = (0..10_000u32)
+            .map(|n| n.wrapping_mul(2_654_435_761))
+            .collect();
+
+        let owned = SpillVec::maybe_spill(values.clone(), None, "u32_owned").unwrap();
+        assert!(matches!(owned, SpillVec::Owned(_)));
+        assert_eq!(&*owned, values.as_slice());
+
+        let dir = TempDir::new("u32");
+        let mapped = SpillVec::maybe_spill(values.clone(), Some(dir.path()), "u32_mapped").unwrap();
+        assert!(matches!(mapped, SpillVec::Mapped { .. }));
+        // The mmap-backed slice must be byte-identical to the original Vec.
+        assert_eq!(&*mapped, values.as_slice());
+    }
+
+    #[test]
+    fn build_u16_fills_in_place_for_both_backings() {
+        let len = 4096usize;
+        let fill = |slice: &mut [u16]| {
+            for (idx, cell) in slice.iter_mut().enumerate() {
+                if idx % 3 == 0 {
+                    *cell = idx as u16; // leave the rest at `default`
+                }
+            }
+        };
+
+        let owned = SpillVec::build_u16(len, 7, None, "u16_owned", fill).unwrap();
+        let dir = TempDir::new("u16");
+        let mapped = SpillVec::build_u16(len, 7, Some(dir.path()), "u16_mapped", fill).unwrap();
+
+        assert_eq!(owned.len(), len);
+        assert_eq!(mapped.len(), len);
+        // Identical contents, and untouched cells read back as `default` (7).
+        assert_eq!(&*owned, &*mapped);
+        assert_eq!(owned[0], 0);
+        assert_eq!(owned[1], 7);
+        assert_eq!(owned[3], 3);
+    }
+
+    #[test]
+    fn spur_keys_survive_the_mmap_roundtrip() {
+        // Spur is a NonZeroU32 niche type — exercises the SpillElem soundness claim.
+        let mut rodeo = lasso::Rodeo::default();
+        let keys: Vec<lasso::Spur> = (0..2000)
+            .map(|n| rodeo.get_or_intern(format!("token-{n}")))
+            .collect();
+
+        let dir = TempDir::new("spur");
+        let mapped = SpillVec::maybe_spill(keys.clone(), Some(dir.path()), "spurs").unwrap();
+        assert!(matches!(mapped, SpillVec::Mapped { .. }));
+        assert_eq!(&*mapped, keys.as_slice());
+        // Keys must still resolve back to their original strings through the interner.
+        let reader = rodeo.into_resolver();
+        for (idx, &key) in mapped.iter().enumerate() {
+            assert_eq!(reader.resolve(&key), format!("token-{idx}"));
+        }
+    }
+
+    #[test]
+    fn empty_input_stays_owned() {
+        let dir = TempDir::new("empty");
+        let mapped = SpillVec::maybe_spill(Vec::<u32>::new(), Some(dir.path()), "empty").unwrap();
+        // A zero-length mmap is invalid, so empties fall back to an owned Vec.
+        assert!(matches!(mapped, SpillVec::Owned(_)));
+        assert!(mapped.is_empty());
+    }
+}