krishiv_sql/
unspillable_headroom.rs

1//! A memory pool that keeps a slice of the budget out of reach of spillable
2//! consumers, so an operator that *cannot* spill can still make progress.
3//!
4//! ## The failure this exists for
5//!
6//! `FairSpillPool` treats its two consumer classes asymmetrically. A spillable
7//! consumer is capped at its fair share of `pool_size - unspillable`; an
8//! **unspillable** one gets whatever is left after *both* classes:
9//!
10//! ```text
11//! // datafusion-execution-54/src/memory_pool/pool.rs, FairSpillPool::try_grow
12//! false => {
13//!     let available = self.pool_size
14//!         .saturating_sub(state.unspillable + state.spillable);
15//! ```
16//!
17//! So N spillable consumers, each politely inside its own `pool/N` share, can
18//! between them occupy the entire pool — and the pool has no way to make any of
19//! them give it back. Spilling in DataFusion is driven by a consumer's *own*
20//! `try_grow` failing; there is no callback the pool can invoke to reclaim.
21//! The next unspillable consumer to ask for memory therefore gets zero, however
22//! little it wants.
23//!
24//! TPC-H q10 and q11 at SF100 died exactly there:
25//!
26//! ```text
27//! Failed to allocate additional 877.0 B for HashJoinInput with 0.0 B already
28//! allocated for this reservation - 0.0 B remain available for the total
29//! memory pool: fair(pool_size: 2.3 GB)
30//! ```
31//!
32//! 877 bytes, refused by a 2.3 GB pool. A hash join build side cannot spill
33//! (DataFusion has no spilling `HashJoinExec`), so these are the *small* joins
34//! [`crate::spillable_join::SpillableJoinSelection`] correctly declined to
35//! convert — they need very little, and there was nothing left to give.
36//!
37//! ## What this does
38//!
39//! Bounds the *total* spillable footprint at `pool_size - headroom`, leaving
40//! `headroom` that only unspillable consumers can occupy. Spillable consumers
41//! hit their ceiling earlier and spill, which is the behaviour they are built
42//! for and already exercise; unspillable ones keep a floor they can always draw
43//! on. The fair-share rule between spillable consumers is unchanged — that is
44//! still `FairSpillPool`'s job, and this delegates to it.
45//!
46//! This is a ceiling on spillers, not a reservation: when no unspillable
47//! consumer is running, the headroom simply goes unused, which costs a query
48//! that spills slightly earlier than it strictly had to. That is the trade —
49//! a little more spilling against queries that cannot run at all.
50
51use std::sync::{Arc, Mutex};
52
53use datafusion::execution::memory_pool::{MemoryPool, MemoryReservation};
54
55/// Fraction of the pool held back for consumers that cannot spill.
56///
57/// A quarter, because the joins that land here are by construction the ones
58/// under `SpillableJoinSelection`'s conversion threshold — anything larger has
59/// already become a (spillable) sort-merge join — and a handful of those fit
60/// comfortably in a quarter of any pool worth having. Too large a headroom
61/// makes every spilling query spill sooner for no benefit.
62pub const DEFAULT_UNSPILLABLE_HEADROOM_NUMERATOR: usize = 1;
63/// Denominator of [`DEFAULT_UNSPILLABLE_HEADROOM_NUMERATOR`].
64pub const DEFAULT_UNSPILLABLE_HEADROOM_DENOMINATOR: usize = 4;
65
66/// Environment override for the headroom, as a percentage of the pool.
67///
68/// `0` disables it and restores plain `FairSpillPool` behaviour.
69pub const UNSPILLABLE_HEADROOM_PERCENT_ENV: &str = "KRISHIV_UNSPILLABLE_HEADROOM_PERCENT";
70
71/// Headroom in bytes for a pool of `pool_size`, honouring the env override.
72#[must_use]
73pub fn headroom_bytes(pool_size: usize) -> usize {
74    let percent = std::env::var(UNSPILLABLE_HEADROOM_PERCENT_ENV)
75        .ok()
76        .and_then(|v| v.trim().parse::<usize>().ok())
77        .filter(|p| *p <= 100);
78    match percent {
79        Some(p) => pool_size / 100 * p,
80        None => {
81            pool_size / DEFAULT_UNSPILLABLE_HEADROOM_DENOMINATOR
82                * DEFAULT_UNSPILLABLE_HEADROOM_NUMERATOR
83        }
84    }
85}
86
87/// See the module docs.
88#[derive(Debug)]
89pub struct UnspillableHeadroomPool {
90    inner: Arc<dyn MemoryPool>,
91    /// Ceiling on the *total* bytes held by spillable consumers.
92    spillable_ceiling: usize,
93    /// Bytes currently held by spillable consumers, tracked here because the
94    /// inner pool does not expose the split.
95    spillable_used: Mutex<usize>,
96    pool_size: usize,
97}
98
99impl UnspillableHeadroomPool {
100    /// Wrap `inner` (a pool of `pool_size` bytes), holding `headroom` back from
101    /// spillable consumers.
102    ///
103    /// A `headroom` of 0, or one at least as large as the pool, disables the
104    /// ceiling — the wrapper then delegates everything unchanged rather than
105    /// bounding spillers to nothing, which would deadlock every spilling query.
106    #[must_use]
107    pub fn new(inner: Arc<dyn MemoryPool>, pool_size: usize, headroom: usize) -> Self {
108        let spillable_ceiling = if headroom == 0 || headroom >= pool_size {
109            pool_size
110        } else {
111            pool_size - headroom
112        };
113        Self {
114            inner,
115            spillable_ceiling,
116            spillable_used: Mutex::new(0),
117            pool_size,
118        }
119    }
120
121    /// The ceiling spillable consumers are held to, in bytes.
122    #[must_use]
123    pub fn spillable_ceiling(&self) -> usize {
124        self.spillable_ceiling
125    }
126
127    fn add_spillable(&self, additional: usize) {
128        if let Ok(mut used) = self.spillable_used.lock() {
129            *used = used.saturating_add(additional);
130        }
131    }
132
133    fn sub_spillable(&self, shrink: usize) {
134        if let Ok(mut used) = self.spillable_used.lock() {
135            *used = used.saturating_sub(shrink);
136        }
137    }
138}
139
140impl std::fmt::Display for UnspillableHeadroomPool {
141    /// Mirrors `FairSpillPool`'s form, with the ceiling — this string is what
142    /// an exhaustion error prints, and "fair(pool_size: 2.3 GB)" alone was not
143    /// enough to tell q10's failure apart from a genuinely full pool.
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        write!(
146            f,
147            "fair+unspillable-headroom(pool_size: {}, spillable_ceiling: {})",
148            human_bytes(self.pool_size),
149            human_bytes(self.spillable_ceiling)
150        )
151    }
152}
153
154impl MemoryPool for UnspillableHeadroomPool {
155    fn name(&self) -> &str {
156        "fair+unspillable-headroom"
157    }
158
159    fn register(&self, consumer: &datafusion::execution::memory_pool::MemoryConsumer) {
160        self.inner.register(consumer);
161    }
162
163    fn unregister(&self, consumer: &datafusion::execution::memory_pool::MemoryConsumer) {
164        self.inner.unregister(consumer);
165    }
166
167    fn grow(&self, reservation: &MemoryReservation, additional: usize) {
168        // Infallible by contract, so the ceiling cannot be enforced here — only
169        // tracked, so `try_grow` keeps seeing the truth.
170        if reservation.consumer().can_spill() {
171            self.add_spillable(additional);
172        }
173        self.inner.grow(reservation, additional);
174    }
175
176    fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
177        if reservation.consumer().can_spill() {
178            self.sub_spillable(shrink);
179        }
180        self.inner.shrink(reservation, shrink);
181    }
182
183    fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> datafusion::error::Result<()> {
184        if !reservation.consumer().can_spill() {
185            return self.inner.try_grow(reservation, additional);
186        }
187        // Hold the spillable counter across the delegated call so a concurrent
188        // spiller cannot slip past the ceiling between the check and the grow.
189        // Lock order is always ours-then-inner's, so this cannot deadlock.
190        let Ok(mut used) = self.spillable_used.lock() else {
191            return self.inner.try_grow(reservation, additional);
192        };
193        let requested = used.saturating_add(additional);
194        if requested > self.spillable_ceiling {
195            return Err(datafusion::error::DataFusionError::ResourcesExhausted(format!(
196                "spillable consumers are capped at {} of the {} pool so that operators \
197                 which cannot spill (hash join build sides) keep a usable floor; \
198                 '{}' asked for {additional} more with {} already held across all \
199                 spillable consumers. This consumer should spill. Set {}=0 to \
200                 restore unbounded fair-share behaviour.",
201                human_bytes(self.spillable_ceiling),
202                human_bytes(self.pool_size),
203                reservation.consumer().name(),
204                human_bytes(*used),
205                UNSPILLABLE_HEADROOM_PERCENT_ENV,
206            )));
207        }
208        self.inner.try_grow(reservation, additional)?;
209        *used = requested;
210        Ok(())
211    }
212
213    fn reserved(&self) -> usize {
214        self.inner.reserved()
215    }
216}
217
218fn human_bytes(bytes: usize) -> String {
219    const MIB: usize = 1024 * 1024;
220    if bytes >= MIB {
221        format!("{:.1} MiB", bytes as f64 / MIB as f64)
222    } else {
223        format!("{bytes} B")
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230    use datafusion::execution::memory_pool::{FairSpillPool, MemoryConsumer};
231
232    fn pool(size: usize, headroom: usize) -> Arc<dyn MemoryPool> {
233        Arc::new(UnspillableHeadroomPool::new(
234            Arc::new(FairSpillPool::new(size)),
235            size,
236            headroom,
237        ))
238    }
239
240    /// The q10/q11 failure, reproduced against the unwrapped pool and fixed by
241    /// the wrapper.
242    ///
243    /// A spillable consumer takes everything `FairSpillPool` will give it —
244    /// which is the whole pool when it is the only spiller — and a hash join
245    /// build side is then refused a few hundred bytes.
246    #[test]
247    fn a_spiller_cannot_starve_a_consumer_that_cannot_spill() {
248        const SIZE: usize = 1024 * 1024;
249
250        // Without headroom: the spiller takes the pool and the join gets nothing.
251        let bare: Arc<dyn MemoryPool> = Arc::new(FairSpillPool::new(SIZE));
252        let spiller = MemoryConsumer::new("ShuffleWriteBuffer")
253            .with_can_spill(true)
254            .register(&bare);
255        spiller.try_grow(SIZE).expect("the only spiller may take it all");
256        let join = MemoryConsumer::new("HashJoinInput").register(&bare);
257        let error = join
258            .try_grow(877)
259            .expect_err("this is the q10/q11 failure and it must reproduce");
260        assert!(
261            error.to_string().contains("HashJoinInput"),
262            "got: {error}"
263        );
264
265        // With headroom: the spiller is capped, and the join is served.
266        let guarded = pool(SIZE, SIZE / 4);
267        let spiller = MemoryConsumer::new("ShuffleWriteBuffer")
268            .with_can_spill(true)
269            .register(&guarded);
270        let error = spiller
271            .try_grow(SIZE)
272            .expect_err("a spiller must not be able to take the whole pool");
273        assert!(
274            error.to_string().contains("cannot spill"),
275            "the refusal must say why, got: {error}"
276        );
277        spiller
278            .try_grow(SIZE / 4 * 3)
279            .expect("up to the ceiling is still allowed");
280        let join = MemoryConsumer::new("HashJoinInput").register(&guarded);
281        join.try_grow(877)
282            .expect("the headroom exists precisely for this");
283    }
284
285    /// Several spillers, each inside its own fair share, must not sum past the
286    /// ceiling — the sum is what starved the join, not any single consumer.
287    #[test]
288    fn the_ceiling_bounds_spillers_in_aggregate_not_individually() {
289        const SIZE: usize = 1024 * 1024;
290        let guarded = pool(SIZE, SIZE / 4);
291        let mut held = Vec::new();
292        for i in 0..4 {
293            let c = MemoryConsumer::new(format!("spiller{i}"))
294                .with_can_spill(true)
295                .register(&guarded);
296            // A quarter each: individually fine, collectively over the ceiling.
297            if i < 3 {
298                c.try_grow(SIZE / 4).expect("within the ceiling");
299            } else {
300                c.try_grow(SIZE / 4)
301                    .expect_err("the fourth quarter crosses the ceiling");
302            }
303            held.push(c);
304        }
305        let join = MemoryConsumer::new("HashJoinInput").register(&guarded);
306        join.try_grow(SIZE / 8).expect("headroom is intact");
307    }
308
309    /// Shrinking returns capacity to the spillable budget; a consumer that
310    /// spilled must be able to grow again afterwards.
311    #[test]
312    fn shrinking_returns_capacity_to_the_spillable_budget() {
313        const SIZE: usize = 1024 * 1024;
314        let guarded = pool(SIZE, SIZE / 4);
315        let spiller = MemoryConsumer::new("s")
316            .with_can_spill(true)
317            .register(&guarded);
318        spiller.try_grow(SIZE / 4 * 3).expect("fills the ceiling");
319        spiller
320            .try_grow(1)
321            .expect_err("nothing left under the ceiling");
322        spiller.shrink(SIZE / 2); // it spilled
323        spiller
324            .try_grow(SIZE / 4)
325            .expect("capacity came back after spilling");
326    }
327
328    /// **Every** bounded `EngineMemory` must install the guard — most of all
329    /// `Private`, which is what the executor uses.
330    ///
331    /// The first version of this fix wrapped only `EngineMemory::shared_pool`.
332    /// Every executor task engine is built with `EngineMemory::Private`
333    /// (krishiv-executor `task_sql_engine`), so the protection was absent from
334    /// the one deployment q10 and q11 fail in — the code shipped, the tests
335    /// passed, and nothing on the cluster changed. Asserting behaviour through
336    /// the real constructor is what makes that visible.
337    #[test]
338    fn both_bounded_engine_memories_install_the_guard() {
339        const SIZE: usize = 1024 * 1024;
340        for (label, pool) in [
341            ("Private", crate::EngineMemory::Private(SIZE).pool()),
342            ("Shared", Some(crate::EngineMemory::shared_pool(SIZE))),
343        ] {
344            let pool = pool.unwrap_or_else(|| panic!("{label} must install a pool"));
345            assert_eq!(
346                pool.name(),
347                "fair+unspillable-headroom",
348                "{label} installed an unguarded pool"
349            );
350            // Behaviour, not just the name: a lone spiller must be refused the
351            // whole budget so a hash join build side keeps a floor.
352            let spiller = MemoryConsumer::new("s").with_can_spill(true).register(&pool);
353            assert!(
354                spiller.try_grow(SIZE).is_err(),
355                "{label}: a lone spiller took the entire pool, so the guard is absent"
356            );
357            spiller
358                .try_grow(SIZE / 4 * 3)
359                .unwrap_or_else(|e| panic!("{label}: the ceiling itself must be reachable — {e}"));
360            let join = MemoryConsumer::new("HashJoinInput").register(&pool);
361            join.try_grow(877)
362                .unwrap_or_else(|e| panic!("{label}: headroom absent — {e}"));
363        }
364    }
365
366    /// Zero headroom is the documented escape hatch and must behave exactly
367    /// like the pool it wraps.
368    #[test]
369    fn zero_headroom_delegates_unchanged() {
370        const SIZE: usize = 1024 * 1024;
371        let guarded = pool(SIZE, 0);
372        let spiller = MemoryConsumer::new("s")
373            .with_can_spill(true)
374            .register(&guarded);
375        spiller
376            .try_grow(SIZE)
377            .expect("with no headroom a lone spiller may still take everything");
378    }
379
380    /// A headroom larger than the pool must not bound spillers to nothing.
381    #[test]
382    fn absurd_headroom_does_not_deadlock_every_spiller() {
383        const SIZE: usize = 1024 * 1024;
384        let guarded = pool(SIZE, SIZE * 4);
385        let spiller = MemoryConsumer::new("s")
386            .with_can_spill(true)
387            .register(&guarded);
388        spiller.try_grow(SIZE).expect("ceiling disabled, not zeroed");
389    }
390}
krishiv_sql/unspillable_headroom.rs

krishiv_sql/
unspillable_headroom.rs
