differential_dataflow/columnar/trace/spill.rs
1//! Per-worker spill control for the columnar `Chunk` trace.
2//!
3//! [`ColChunk`](super::ColChunk) is either resident (a trie in
4//! memory) or paged (resident bounds + a byte handle to fetch the trie back).
5//! `Chunk::settle` is where committed chunks may be **paged out**: it calls
6//! [`try_page`], which consults a per-worker [`SpillState`] and, when over budget,
7//! serializes the trie and hands the bytes to a pluggable [`BytesStore`]. Reads
8//! ([`Chunk::merge`](crate::trace::chunk::Chunk::merge) etc., and the cursor) fetch paged chunks back on demand.
9//!
10//! The backend is pluggable: a worker calls [`install`] with its own
11//! [`BytesStore`] (e.g. file-backed). With nothing installed, no chunk is ever
12//! paged and `ColChunk` stays resident — zero overhead beyond a thread-local peek.
13//!
14//! # Known limitation
15//!
16//! `settle` sees only its **local** committed output, not the whole batcher
17//! queue (timely's `MergeBatcher` does not expose it, and the shared `settle` is
18//! generic so it can't notify a columnar-specific accountant per consumed chunk).
19//! So the eviction policy is a simple per-worker **high-water mark**: keep the
20//! first `budget_records` worth of committed chunks resident, page the rest. It
21//! bounds the resident set and round-trips through the store correctly, but it
22//! favours keeping older data and can't reclaim budget as data leaves — coarser
23//! than the exact head-reserve-over-the-queue policy the old bespoke batcher ran.
24
25use std::cell::RefCell;
26use std::sync::Arc;
27use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
28
29use crate::columnar::layout::ColumnarUpdate as Update;
30use crate::columnar::updates::{Updates, UpdatesTyped};
31
32/// Append a chunk's bytes to backing storage, returning a handle to read them back.
33pub trait BytesStore {
34 /// Persist `bytes`, returning a source that can reproduce them.
35 fn store(&mut self, bytes: &[u8]) -> Box<dyn BytesSource>;
36}
37
38/// A handle to bytes previously written via a [`BytesStore`].
39pub trait BytesSource {
40 /// Read the bytes back. Panics on failure (the trace cannot make progress without them).
41 fn load(&self) -> Vec<u8>;
42}
43
44/// Cumulative spill counters, shared (via `Arc`) so a run can sum across workers.
45#[derive(Default)]
46pub struct SpillStats {
47 /// Chunks paged out.
48 pub spilled_chunks: AtomicUsize,
49 /// Records (updates) in paged-out chunks.
50 pub spilled_records: AtomicUsize,
51 /// Chunks fetched back in.
52 pub fetched_chunks: AtomicUsize,
53 /// Serialized bytes written.
54 pub bytes_written: AtomicUsize,
55}
56
57/// Per-worker spill state: the budget, the running resident estimate, the store,
58/// and counters.
59struct SpillState {
60 /// Keep up to this many records resident before paging committed chunks.
61 budget_records: usize,
62 /// Records kept resident so far — a monotonic high-water mark. Once it reaches
63 /// `budget_records`, every further committed chunk is paged (see the module's
64 /// "Known limitation").
65 resident_records: usize,
66 /// Pluggable backing store.
67 store: Box<dyn BytesStore>,
68 /// Counters, shared with the installer.
69 stats: Arc<SpillStats>,
70}
71
72thread_local! {
73 static STATE: RefCell<Option<SpillState>> = const { RefCell::new(None) };
74}
75
76/// Install spill on this worker: page committed chunks once resident records
77/// exceed `budget_records`, writing to `store`. `stats` receives the counters.
78pub fn install(budget_records: usize, store: Box<dyn BytesStore>, stats: Arc<SpillStats>) {
79 STATE.with(|s| *s.borrow_mut() = Some(SpillState { budget_records, resident_records: 0, store, stats }));
80}
81
82/// Remove this worker's spill state (so later traces stay resident).
83pub fn uninstall() {
84 STATE.with(|s| *s.borrow_mut() = None);
85}
86
87/// Whether a spiller is installed on this worker (a cheap peek `settle` makes
88/// before computing any per-chunk paging metadata).
89pub(crate) fn active() -> bool {
90 STATE.with(|s| s.borrow().is_some())
91}
92
93/// Try to page `updates` out. On success returns a [`BytesSource`] handle to fetch
94/// it back; otherwise returns the trie to keep resident.
95///
96/// Policy: a **high-water mark** — keep the first `budget_records` worth of
97/// committed chunks resident, page everything after. Simple and robust (the
98/// counter only rises, so paging is deterministic once the budget is reached);
99/// it does not reclaim budget when data later leaves the batcher, since
100/// `settle` only sees its local output (see the module-level limitation).
101pub(crate) fn try_page<U: Update>(updates: UpdatesTyped<U>) -> Result<Box<dyn BytesSource>, UpdatesTyped<U>> {
102 STATE.with(|s| {
103 let mut guard = s.borrow_mut();
104 let Some(state) = guard.as_mut() else { return Err(updates) };
105 let records = updates.len();
106 if state.resident_records + records <= state.budget_records {
107 state.resident_records += records;
108 return Err(updates);
109 }
110 // Over the high-water mark: serialize and page out.
111 let mut buf = Vec::new();
112 Updates::<U>::from(updates).write_to(&mut buf);
113 state.stats.spilled_chunks.fetch_add(1, Relaxed);
114 state.stats.spilled_records.fetch_add(records, Relaxed);
115 state.stats.bytes_written.fetch_add(buf.len(), Relaxed);
116 Ok(state.store.store(&buf))
117 })
118}
119
120/// A paged chunk was fetched back in; bump the counter (for stats).
121pub(crate) fn note_fetched() {
122 STATE.with(|s| {
123 if let Some(state) = s.borrow_mut().as_mut() {
124 state.stats.fetched_chunks.fetch_add(1, Relaxed);
125 }
126 });
127}
128
129/// Reconstruct a trie from bytes a [`BytesSource`] produced (the inverse of
130/// `try_page`'s serialization).
131pub(crate) fn decode<U: Update>(source: &dyn BytesSource) -> UpdatesTyped<U> {
132 let bytes = timely::bytes::arc::BytesMut::from(source.load()).freeze();
133 Updates::<U>::read_from(bytes).into_typed()
134}