From 8d0ecdab1952020cdb9880114d9e18dfa14ab6c2 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Sat, 20 Jun 2026 11:03:04 +0100 Subject: [PATCH 1/2] Support spilling --- docker-compose.yml | 6 + server-rs/Cargo.lock | 1 + server-rs/Cargo.toml | 3 + server-rs/src/data.rs | 1 + server-rs/src/data/property/loading.rs | 115 ++++++++++++++------ server-rs/src/data/property/mod.rs | 46 ++++---- server-rs/src/main.rs | 18 ++- server-rs/src/routes/actual_listings.rs | 2 +- server-rs/src/routes/ai_filters/matching.rs | 2 +- server-rs/src/routes/export.rs | 2 +- server-rs/src/routes/filter_counts.rs | 2 +- server-rs/src/routes/hexagon_stats.rs | 4 +- server-rs/src/routes/hexagons.rs | 2 +- server-rs/src/routes/postcode_properties.rs | 2 +- server-rs/src/routes/postcode_stats.rs | 2 +- server-rs/src/routes/postcodes.rs | 2 +- server-rs/src/routes/properties.rs | 2 +- server-rs/src/routes/stats.rs | 3 +- 18 files changed, 150 insertions(+), 65 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index 3c4cf31..37468c1 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -33,6 +33,12 @@ services: # Fallback only — the binary uses jemalloc as its global allocator # (tuned via a baked-in malloc_conf). Caps glibc to 2 arenas. MALLOC_ARENA_MAX: "2" + # Dev only: spill the large property arrays (feature matrix + + # address-search index, ~4GB) to disk and memory-map them read-only + # instead of holding them in RAM, so a low-memory box can run the full + # dataset. Points at the `target` volume (real disk, not host-synced, + # not tmpfs). Comment out to keep everything resident (prod behaviour). + SPILL_DIR: /app/server-rs/target/spill POCKETBASE_URL: http://pocketbase:8090 POCKETBASE_ADMIN_EMAIL: *pb-email POCKETBASE_ADMIN_PASSWORD: *pb-password diff --git a/server-rs/Cargo.lock b/server-rs/Cargo.lock index db5ab54..bd39c8b 100644 --- a/server-rs/Cargo.lock +++ b/server-rs/Cargo.lock @@ -3887,6 +3887,7 @@ dependencies = [ "hmac 0.13.0", "lasso", "libc", + "memmap2", "metrics", "metrics-exporter-prometheus", "parking_lot", diff --git a/server-rs/Cargo.toml b/server-rs/Cargo.toml index 2d7d290..fa31288 100644 --- a/server-rs/Cargo.toml +++ b/server-rs/Cargo.toml @@ -27,6 +27,9 @@ urlencoding = "2" url = "2" rust_xlsxwriter = "0.94" pmtiles = { version = "0.23", features = ["mmap-async-tokio"] } +# Read-only mmap backing for the large property arrays under the `--spill-dir` +# dev flag (see data/spill.rs); lets a low-memory dev box page them from disk. +memmap2 = "0.9" rand = "0.10" hmac = "0.13" sha2 = "0.11" diff --git a/server-rs/src/data.rs b/server-rs/src/data.rs index 7d739f0..bf66cd6 100644 --- a/server-rs/src/data.rs +++ b/server-rs/src/data.rs @@ -4,6 +4,7 @@ mod places; mod poi; mod postcodes; mod property; +pub mod spill; pub mod travel_time; /// Apostrophe-like code points that should be elided (not treated as word breaks) when diff --git a/server-rs/src/data/property/loading.rs b/server-rs/src/data/property/loading.rs index 29d224d..0872ab2 100644 --- a/server-rs/src/data/property/loading.rs +++ b/server-rs/src/data/property/loading.rs @@ -19,6 +19,7 @@ use super::address_search::{ use super::poi_metrics::{PostcodePoiMetrics, NO_POI_METRIC_ROW}; use super::stats::{column_to_f32_vec, compute_feature_stats, FeatureStats, Histogram}; use super::{HistoricalPrice, PropertyData, RenovationEvent, TenureEvent}; +use crate::data::spill::SpillVec; const MISSING_COORDINATE_SAMPLE_LIMIT: usize = 10; const COUNTRY_COLUMN_CANDIDATES: &[&str] = &[ @@ -160,11 +161,25 @@ fn validate_no_england_rows_missing_coordinates( } impl PropertyData { - pub fn load(properties_path: &Path, postcode_features_path: &Path) -> anyhow::Result { - crate::data::run_polars_io(|| Self::load_inner(properties_path, postcode_features_path)) + /// Load the property data. When `spill` is `Some(dir)`, the large flat arrays + /// (the feature matrix and the address-search index) are written to anonymous + /// files in `dir` and memory-mapped read-only instead of held on the heap — + /// the `--spill-dir` dev flag, which lets a low-memory box page them from disk. + pub fn load( + properties_path: &Path, + postcode_features_path: &Path, + spill: Option<&Path>, + ) -> anyhow::Result { + crate::data::run_polars_io(|| { + Self::load_inner(properties_path, postcode_features_path, spill) + }) } - fn load_inner(properties_path: &Path, postcode_features_path: &Path) -> anyhow::Result { + fn load_inner( + properties_path: &Path, + postcode_features_path: &Path, + spill: Option<&Path>, + ) -> anyhow::Result { // Load postcode.parquet tracing::info!( "Loading postcode features from {:?}", @@ -1011,37 +1026,71 @@ impl PropertyData { // Transpose to row-major AND apply spatial permutation in one pass. // Combines numeric and enum features into a single feature_data array, quantized to u16. tracing::info!("Transposing to row-major layout (spatially sorted, quantized to u16)"); - let mut feature_data = vec![NAN_U16; row_count * num_features]; - feature_data - .par_chunks_mut(num_features) - .enumerate() - .for_each(|(new_row, row_slice)| { - let old_index = perm[new_row] as usize; - // Numeric features: quantize to u16 - for (feat_idx, col_vec) in numeric_col_major.iter().enumerate() { - let value = col_vec[old_index]; - row_slice[feat_idx] = if value.is_finite() { - let range = quant_range[feat_idx]; - if range > 0.0 { - let normalized = (value - quant_min[feat_idx]) / range; - (normalized * QUANT_SCALE).round().clamp(0.0, QUANT_SCALE) as u16 - } else { - 0 + // Built in place so that under `--spill-dir` the 3GB matrix lands straight + // in the mmap-backed file and never also exists on the heap. + let feature_data = SpillVec::build_u16( + row_count * num_features, + NAN_U16, + spill, + "feature_data", + |feature_data| { + feature_data + .par_chunks_mut(num_features) + .enumerate() + .for_each(|(new_row, row_slice)| { + let old_index = perm[new_row] as usize; + // Numeric features: quantize to u16 + for (feat_idx, col_vec) in numeric_col_major.iter().enumerate() { + let value = col_vec[old_index]; + row_slice[feat_idx] = if value.is_finite() { + let range = quant_range[feat_idx]; + if range > 0.0 { + let normalized = (value - quant_min[feat_idx]) / range; + (normalized * QUANT_SCALE).round().clamp(0.0, QUANT_SCALE) + as u16 + } else { + 0 + } + } else { + NAN_U16 + }; } - } else { - NAN_U16 - }; - } - // Enum features: store as u16 directly - for (enum_idx, (_, encoded)) in enum_col_major.iter().enumerate() { - let value = encoded[old_index]; - row_slice[num_numeric + enum_idx] = if value.is_finite() { - value as u16 - } else { - NAN_U16 - }; - } - }); + // Enum features: store as u16 directly + for (enum_idx, (_, encoded)) in enum_col_major.iter().enumerate() { + let value = encoded[old_index]; + row_slice[num_numeric + enum_idx] = if value.is_finite() { + value as u16 + } else { + NAN_U16 + }; + } + }); + }, + )?; + + // Spill the remaining large flat arrays (the address-search index) to disk + // when configured. The posting-list hashmaps and interners stay on the heap + // — they aren't flat arrays, and they're a small fraction of the footprint. + let address_buffer = + SpillVec::maybe_spill(address_buffer.into_bytes(), spill, "address_buffer")?; + let address_offsets = SpillVec::maybe_spill(address_offsets, spill, "address_offsets")?; + let address_lengths = SpillVec::maybe_spill(address_lengths, spill, "address_lengths")?; + let address_search_token_keys = SpillVec::maybe_spill( + address_search_token_keys, + spill, + "address_search_token_keys", + )?; + let address_search_token_offsets = SpillVec::maybe_spill( + address_search_token_offsets, + spill, + "address_search_token_offsets", + )?; + let address_search_token_lengths = SpillVec::maybe_spill( + address_search_token_lengths, + spill, + "address_search_token_lengths", + )?; + let postcode_keys = SpillVec::maybe_spill(postcode_keys, spill, "postcode_keys")?; tracing::info!("Data loading complete"); diff --git a/server-rs/src/data/property/mod.rs b/server-rs/src/data/property/mod.rs index 72a3b4d..8a3d20f 100644 --- a/server-rs/src/data/property/mod.rs +++ b/server-rs/src/data/property/mod.rs @@ -26,6 +26,7 @@ use rustc_hash::FxHashMap; use serde::Serialize; use crate::consts::NAN_U16; +use crate::data::spill::SpillVec; #[derive(Serialize, Clone)] pub struct RenovationEvent { @@ -61,7 +62,9 @@ pub struct PropertyData { /// Quantized to u16. NaN sentinel = u16::MAX (65535). /// Numeric features: encoded via (val - min) / range * 65534. /// Enum features: stored directly as u16 cast of the f32 index. - pub feature_data: Vec, + /// Heap-resident in prod; mmap-backed under `--spill-dir` (the largest array, + /// so the biggest dev-memory win). Derefs to `[u16]`, so reads are unchanged. + pub feature_data: SpillVec, /// Per-feature: range / QUANT_SCALE for fast decode. dequant_a: Vec, /// Per-feature: minimum value (offset for dequantization). @@ -72,15 +75,16 @@ pub struct PropertyData { pub poi_metrics: PostcodePoiMetrics, /// Unquantized last sale price used by the price-history chart. last_known_price_raw: Vec, - /// Contiguous buffer holding all address strings end-to-end. - address_buffer: String, + /// Contiguous buffer holding all address strings end-to-end (valid UTF-8 by + /// construction; spillable, so stored as bytes rather than `String`). + address_buffer: SpillVec, /// Byte offset into `address_buffer` where each row's address starts. - address_offsets: Vec, + address_offsets: SpillVec, /// Length in bytes of each row's address. - address_lengths: Vec, + address_lengths: SpillVec, /// Interned postcodes: reader is thread-safe, keys index into it. postcode_interner: lasso::RodeoReader, - postcode_keys: Vec, + postcode_keys: SpillVec, /// Rows for each postcode, keyed by the interned postcode key. postcode_row_index: FxHashMap>, /// Inverted index from address tokens to property rows. @@ -91,11 +95,11 @@ pub struct PropertyData { /// Resolve-only (no string->key reverse map): scoring only resolves keys. address_search_interner: lasso::RodeoResolver, /// Flat per-row normalized address-search token keys. - address_search_token_keys: Vec, + address_search_token_keys: SpillVec, /// Offset into `address_search_token_keys` for each row. - address_search_token_offsets: Vec, + address_search_token_offsets: SpillVec, /// Number of normalized address-search token keys for each row. - address_search_token_lengths: Vec, + address_search_token_lengths: SpillVec, /// For enum features: maps feature index to list of possible string values. /// Index in values list corresponds to the u16 value stored in feature_data. pub enum_values: rustc_hash::FxHashMap>, @@ -123,7 +127,11 @@ impl PropertyData { pub fn address(&self, row: usize) -> &str { let offset = self.address_offsets[row] as usize; let length = self.address_lengths[row] as usize; - &self.address_buffer[offset..offset + length] + let bytes = &self.address_buffer[offset..offset + length]; + // SAFETY: `address_buffer` is built by concatenating `&str` slices on UTF-8 + // char boundaries, so every recorded `[offset, offset + length)` range is + // valid UTF-8. + unsafe { std::str::from_utf8_unchecked(bytes) } } /// Get the postcode string for a given row. @@ -133,7 +141,7 @@ impl PropertyData { /// Get postcode components for field-level borrowing (avoids conflicting borrows with feature_data). pub fn postcode_parts(&self) -> (&lasso::RodeoReader, &[lasso::Spur]) { - (&self.postcode_interner, &self.postcode_keys) + (&self.postcode_interner, self.postcode_keys.as_slice()) } /// Property rows for a given postcode string, or empty if unknown. @@ -229,25 +237,25 @@ impl PropertyData { feature_names: Vec::new(), num_features: 0, num_numeric: 0, - feature_data: Vec::new(), + feature_data: SpillVec::owned(Vec::new()), dequant_a: Vec::new(), quant_min: Vec::new(), quant_range: Vec::new(), feature_stats: Vec::new(), poi_metrics: PostcodePoiMetrics::empty(0), last_known_price_raw: Vec::new(), - address_buffer: String::new(), - address_offsets: Vec::new(), - address_lengths: Vec::new(), + address_buffer: SpillVec::owned(Vec::new()), + address_offsets: SpillVec::owned(Vec::new()), + address_lengths: SpillVec::owned(Vec::new()), postcode_interner: lasso::Rodeo::default().into_reader(), - postcode_keys: Vec::new(), + postcode_keys: SpillVec::owned(Vec::new()), postcode_row_index: FxHashMap::default(), address_token_index: FxHashMap::default(), address_prefix_index: FxHashMap::default(), address_search_interner: lasso::Rodeo::default().into_resolver(), - address_search_token_keys: Vec::new(), - address_search_token_offsets: Vec::new(), - address_search_token_lengths: Vec::new(), + address_search_token_keys: SpillVec::owned(Vec::new()), + address_search_token_offsets: SpillVec::owned(Vec::new()), + address_search_token_lengths: SpillVec::owned(Vec::new()), enum_values: rustc_hash::FxHashMap::default(), enum_counts: rustc_hash::FxHashMap::default(), approx_build_date_bits: Vec::new(), diff --git a/server-rs/src/main.rs b/server-rs/src/main.rs index 525955a..c0e28c3 100644 --- a/server-rs/src/main.rs +++ b/server-rs/src/main.rs @@ -281,6 +281,14 @@ struct Cli { #[arg(long)] dist: Option, + /// Dev only: spill the large property arrays (feature matrix + address-search + /// index) to anonymous files in this directory and memory-map them read-only, + /// instead of holding them on the heap. Trades a little speed for a much + /// smaller resident set so a low-memory dev box can run the full dataset; the + /// mapped pages are file-backed and reclaimable under pressure. Omit in prod. + #[arg(long, env = "SPILL_DIR")] + spill_dir: Option, + /// URL of the screenshot service (e.g. http://screenshot:8002) #[arg(long, env = "SCREENSHOT_URL")] screenshot_url: String, @@ -465,7 +473,15 @@ async fn main() -> anyhow::Result<()> { cli.properties.display(), cli.postcode_features.display(), ); - let property_data = data::PropertyData::load(&cli.properties, &cli.postcode_features)?; + let spill_dir = cli.spill_dir.as_deref(); + if let Some(dir) = spill_dir { + info!( + "Spill-to-disk enabled: large property arrays will be memory-mapped from {}", + dir.display() + ); + } + let property_data = + data::PropertyData::load(&cli.properties, &cli.postcode_features, spill_dir)?; trim_allocator("property data load"); info!( rows = property_data.lat.len(), diff --git a/server-rs/src/routes/actual_listings.rs b/server-rs/src/routes/actual_listings.rs index 645abcc..738e70f 100644 --- a/server-rs/src/routes/actual_listings.rs +++ b/server-rs/src/routes/actual_listings.rs @@ -48,7 +48,7 @@ const KEEP_UNKNOWN_LISTING_FILTER_FEATURES: &[&str] = &[ "Construction year", "Interior height (m)", "Current energy rating", - "Potential energy rating" + "Potential energy rating", ]; const LISTING_BOUNDS_EPSILON_DEGREES: f64 = 0.00001; diff --git a/server-rs/src/routes/ai_filters/matching.rs b/server-rs/src/routes/ai_filters/matching.rs index 74686d3..984a52d 100644 --- a/server-rs/src/routes/ai_filters/matching.rs +++ b/server-rs/src/routes/ai_filters/matching.rs @@ -105,7 +105,7 @@ pub(super) fn count_matching_rows( .collect(); let has_travel = !travel_data.is_empty(); - let feature_data = &state.data.feature_data; + let feature_data: &[u16] = &state.data.feature_data; let num_features = state.data.num_features; let num_rows = state.data.lat.len(); let (pc_interner, pc_keys) = state.data.postcode_parts(); diff --git a/server-rs/src/routes/export.rs b/server-rs/src/routes/export.rs index 1d8491f..5551411 100644 --- a/server-rs/src/routes/export.rs +++ b/server-rs/src/routes/export.rs @@ -632,7 +632,7 @@ pub async fn get_export( let bytes = tokio::task::spawn_blocking(move || -> Result, String> { let t0 = std::time::Instant::now(); let num_features = state.data.num_features; - let feature_data = &state.data.feature_data; + let feature_data: &[u16] = &state.data.feature_data; let quant = state.data.quant_ref(); let feature_names = &state.data.feature_names; let enum_values = &state.data.enum_values; diff --git a/server-rs/src/routes/filter_counts.rs b/server-rs/src/routes/filter_counts.rs index 364fab4..a9d49e0 100644 --- a/server-rs/src/routes/filter_counts.rs +++ b/server-rs/src/routes/filter_counts.rs @@ -88,7 +88,7 @@ pub async fn get_filter_counts( let response = tokio::task::spawn_blocking(move || -> Result { let t0 = std::time::Instant::now(); let num_features = state.data.num_features; - let feature_data = &state.data.feature_data; + let feature_data: &[u16] = &state.data.feature_data; // Load travel time data let travel_data: Vec = travel_entries diff --git a/server-rs/src/routes/hexagon_stats.rs b/server-rs/src/routes/hexagon_stats.rs index cb79a61..289e280 100644 --- a/server-rs/src/routes/hexagon_stats.rs +++ b/server-rs/src/routes/hexagon_stats.rs @@ -226,7 +226,7 @@ pub(super) fn top_filter_exclusions( return Vec::new(); } - let feature_data = &data.feature_data; + let feature_data: &[u16] = &data.feature_data; let num_features = data.num_features; let quant = data.quant_ref(); let poi_quant = data.poi_metrics.quant_ref(); @@ -494,7 +494,7 @@ pub async fn get_hexagon_stats( .map_err(|err| format!("Invalid H3 resolution {}: {}", resolution, err))?; let need_parent = needs_parent(resolution); let num_features = state.data.num_features; - let feature_data = &state.data.feature_data; + let feature_data: &[u16] = &state.data.feature_data; let travel_data = load_travel_data(&state.travel_time_store, &travel_entries)?; let has_travel = !travel_entries.is_empty(); diff --git a/server-rs/src/routes/hexagons.rs b/server-rs/src/routes/hexagons.rs index 0824d72..d48bee2 100644 --- a/server-rs/src/routes/hexagons.rs +++ b/server-rs/src/routes/hexagons.rs @@ -332,7 +332,7 @@ pub async fn get_hexagons( .collect(); let num_features = state.data.num_features; - let feature_data = &state.data.feature_data; + let feature_data: &[u16] = &state.data.feature_data; let quant = state.data.quant_ref(); let (pc_interner, pc_keys) = state.data.postcode_parts(); let min_keys = &state.min_keys; diff --git a/server-rs/src/routes/postcode_properties.rs b/server-rs/src/routes/postcode_properties.rs index d6f08cf..352d351 100644 --- a/server-rs/src/routes/postcode_properties.rs +++ b/server-rs/src/routes/postcode_properties.rs @@ -108,7 +108,7 @@ pub async fn get_postcode_properties( let result = tokio::task::spawn_blocking(move || { let t0 = std::time::Instant::now(); let num_features = state.data.num_features; - let feature_data = &state.data.feature_data; + let feature_data: &[u16] = &state.data.feature_data; let feature_names = &state.data.feature_names; let feature_name_to_index = &state.feature_name_to_index; let enum_values = &state.data.enum_values; diff --git a/server-rs/src/routes/postcode_stats.rs b/server-rs/src/routes/postcode_stats.rs index b18536c..2bc406d 100644 --- a/server-rs/src/routes/postcode_stats.rs +++ b/server-rs/src/routes/postcode_stats.rs @@ -92,7 +92,7 @@ pub async fn get_postcode_stats( let response = tokio::task::spawn_blocking(move || { let start_time = std::time::Instant::now(); let num_features = state.data.num_features; - let feature_data = &state.data.feature_data; + let feature_data: &[u16] = &state.data.feature_data; let travel_data = load_travel_data(&state.travel_time_store, &travel_entries)?; let has_travel = !travel_entries.is_empty(); diff --git a/server-rs/src/routes/postcodes.rs b/server-rs/src/routes/postcodes.rs index 5203453..723bf7d 100644 --- a/server-rs/src/routes/postcodes.rs +++ b/server-rs/src/routes/postcodes.rs @@ -131,7 +131,7 @@ pub async fn get_postcodes( .collect(); let num_features = state.data.num_features; - let feature_data = &state.data.feature_data; + let feature_data: &[u16] = &state.data.feature_data; let quant = state.data.quant_ref(); let min_keys = &state.min_keys; let max_keys = &state.max_keys; diff --git a/server-rs/src/routes/properties.rs b/server-rs/src/routes/properties.rs index e9aa484..630aa6f 100644 --- a/server-rs/src/routes/properties.rs +++ b/server-rs/src/routes/properties.rs @@ -332,7 +332,7 @@ pub async fn get_hexagon_properties( .map_err(|err| format!("Invalid H3 resolution {}: {}", resolution, err))?; let need_parent = needs_parent(resolution); let num_features = state.data.num_features; - let feature_data = &state.data.feature_data; + let feature_data: &[u16] = &state.data.feature_data; let feature_names = &state.data.feature_names; let feature_name_to_index = &state.feature_name_to_index; let enum_values = &state.data.enum_values; diff --git a/server-rs/src/routes/stats.rs b/server-rs/src/routes/stats.rs index a2c0f59..8dad7c7 100644 --- a/server-rs/src/routes/stats.rs +++ b/server-rs/src/routes/stats.rs @@ -482,12 +482,13 @@ pub fn compute_poi_feature_stats( mod tests { use super::*; use crate::consts::NAN_U16; + use crate::data::spill::SpillVec; fn enum_data(values: &[u16]) -> PropertyData { let mut data = PropertyData::empty_for_tests(); data.num_features = 1; data.num_numeric = 0; // single enum feature at index 0 - data.feature_data = values.to_vec(); + data.feature_data = SpillVec::owned(values.to_vec()); data.enum_values .insert(0, vec!["Yes".to_string(), "No".to_string()]); data From 45882e7de2deee04b98540242f4150f9a6dc4f68 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Sat, 20 Jun 2026 11:03:07 +0100 Subject: [PATCH 2/2] spill --- server-rs/src/data/spill.rs | 330 ++++++++++++++++++++++++++++++++++++ 1 file changed, 330 insertions(+) create mode 100644 server-rs/src/data/spill.rs diff --git a/server-rs/src/data/spill.rs b/server-rs/src/data/spill.rs new file mode 100644 index 0000000..977031f --- /dev/null +++ b/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`] 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` 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 { + Owned(Vec), + Mapped { + map: Mmap, + len: usize, + _marker: PhantomData, + }, +} + +impl SpillVec { + /// Keep `values` on the heap (production behaviour, and the only path used in + /// tests). + pub fn owned(values: Vec) -> 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, dir: Option<&Path>, label: &str) -> anyhow::Result { + 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 { + /// 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 { + match dir { + Some(dir) if len > 0 => { + let byte_len = len * std::mem::size_of::(); + 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::(), 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 std::ops::Deref for SpillVec { + 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::()); + debug_assert_eq!(bytes.as_ptr().align_offset(std::mem::align_of::()), 0); + // SAFETY: the mapping holds exactly `len * size_of::()` 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::(), *len) } + } + } + } +} + +fn spill_vec(values: Vec, dir: &Path, label: &str) -> anyhow::Result> { + let len = values.len(); + let byte_len = len * std::mem::size_of::(); + // 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` 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::(), 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 { + 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 = (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 = (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::::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()); + } +}