hawdb_executor/spill/
pool.rs

1// Copyright 2026 Nowledge
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::ExecutionMemoryConfig;
16use hawdb_core::{
17    HawDBError, ResourceBudgetError, ResourceBudgetKind, ResourceBudgetScope, Result,
18};
19use hawdb_storage::file_descriptors::FileOpenContext;
20use std::collections::BTreeMap;
21use std::io::ErrorKind;
22use std::path::{Path, PathBuf};
23use std::sync::atomic::{AtomicU64, Ordering};
24use std::sync::{Arc, Condvar, Mutex, MutexGuard, OnceLock};
25use std::time::{Duration, SystemTime, UNIX_EPOCH};
26
27pub(super) const SPILL_FILE_PREFIX: &str = "hawdb-spill-v1-";
28pub(super) const SPILL_FILE_SUFFIX: &str = ".spill";
29static SPILL_POOLS: OnceLock<Mutex<BTreeMap<PathBuf, Arc<SharedSpillPool>>>> = OnceLock::new();
30static PROCESS_MARKER: OnceLock<String> = OnceLock::new();
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub struct SpillPoolSnapshot {
34    pub max_total_bytes: u64,
35    pub max_total_runs: usize,
36    pub min_free_bytes: u64,
37    pub free_space_probe_interval_bytes: u64,
38    pub free_space_probe_count: u64,
39    pub free_space_probe_failures: u64,
40    pub active_bytes: u64,
41    pub peak_active_bytes: u64,
42    pub pending_write_bytes: u64,
43    pub active_runs: usize,
44    pub peak_active_runs: usize,
45    pub orphan_files_removed: u64,
46    pub orphan_bytes_removed: u64,
47    pub orphan_cleanup_failures: u64,
48    pub run_delete_failures: u64,
49}
50
51#[derive(Debug, Clone, Copy)]
52struct SpillPoolLimits {
53    max_total_bytes: u64,
54    max_total_runs: usize,
55    min_free_bytes: u64,
56    free_space_probe_interval_bytes: u64,
57}
58
59impl SpillPoolLimits {
60    fn from_config(memory: &ExecutionMemoryConfig) -> Self {
61        Self {
62            max_total_bytes: memory.max_total_spill_bytes.get(),
63            max_total_runs: memory.max_total_spill_runs.get(),
64            min_free_bytes: memory.min_spill_free_bytes.get(),
65            free_space_probe_interval_bytes: memory.spill_free_space_probe_interval_bytes.get(),
66        }
67    }
68
69    fn tighten(&mut self, other: Self) {
70        self.max_total_bytes = self.max_total_bytes.min(other.max_total_bytes);
71        self.max_total_runs = self.max_total_runs.min(other.max_total_runs);
72        self.min_free_bytes = self.min_free_bytes.max(other.min_free_bytes);
73        self.free_space_probe_interval_bytes = self
74            .free_space_probe_interval_bytes
75            .min(other.free_space_probe_interval_bytes);
76    }
77}
78
79#[derive(Debug, Default)]
80struct OrphanCleanupStats {
81    files_removed: u64,
82    bytes_removed: u64,
83    failures: u64,
84}
85
86#[derive(Debug)]
87struct SpillPoolState {
88    limits: SpillPoolLimits,
89    active_bytes: u64,
90    peak_active_bytes: u64,
91    pending_write_bytes: u64,
92    last_probed_available_bytes: Option<u64>,
93    unreflected_reserved_bytes: u64,
94    bytes_reserved_since_probe: u64,
95    free_space_probe_in_progress: bool,
96    free_space_probe_count: u64,
97    free_space_probe_failures: u64,
98    active_runs: usize,
99    peak_active_runs: usize,
100    orphan_cleanup: OrphanCleanupStats,
101    run_delete_failures: u64,
102}
103
104#[derive(Debug)]
105struct SharedSpillPool {
106    directory: PathBuf,
107    state: Mutex<SpillPoolState>,
108    free_space_probe_ready: Condvar,
109    free_space_probe: Arc<dyn SpillSpaceProbe>,
110}
111
112trait SpillSpaceProbe: std::fmt::Debug + Send + Sync {
113    fn available_space(&self, directory: &Path) -> std::io::Result<u64>;
114}
115
116#[derive(Debug)]
117struct FileSystemSpillSpaceProbe;
118
119impl SpillSpaceProbe for FileSystemSpillSpaceProbe {
120    fn available_space(&self, directory: &Path) -> std::io::Result<u64> {
121        #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
122        {
123            fs2::available_space(directory)
124        }
125        #[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
126        {
127            let _ = directory;
128            Err(std::io::Error::new(
129                std::io::ErrorKind::Unsupported,
130                "disk spill is unavailable in browser WASM",
131            ))
132        }
133    }
134}
135
136#[derive(Debug, Clone)]
137pub(crate) struct SpillPool {
138    shared: Arc<SharedSpillPool>,
139    file_descriptors: FileOpenContext,
140}
141
142impl SpillPool {
143    pub(crate) fn open(memory: &ExecutionMemoryConfig) -> Result<Self> {
144        Self::open_with_context(memory, None)
145    }
146
147    pub(crate) fn open_with_context(
148        memory: &ExecutionMemoryConfig,
149        context: Option<&FileOpenContext>,
150    ) -> Result<Self> {
151        if cfg!(all(target_arch = "wasm32", target_os = "unknown")) {
152            return Err(HawDBError::Execution(
153                "disk spill is unavailable in browser WASM".into(),
154            ));
155        }
156        Self::open_admitted(memory, Arc::new(FileSystemSpillSpaceProbe), context)
157    }
158
159    #[cfg(test)]
160    fn open_with_probe(
161        memory: &ExecutionMemoryConfig,
162        free_space_probe: Arc<dyn SpillSpaceProbe>,
163    ) -> Result<Self> {
164        Self::open_admitted(memory, free_space_probe, None)
165    }
166
167    fn open_admitted(
168        memory: &ExecutionMemoryConfig,
169        free_space_probe: Arc<dyn SpillSpaceProbe>,
170        context: Option<&FileOpenContext>,
171    ) -> Result<Self> {
172        let file_descriptors = match context {
173            Some(context) => context.clone(),
174            None => FileOpenContext::for_path(&memory.spill_directory)?,
175        };
176        file_descriptors
177            .create_dir_all(&memory.spill_directory)
178            .map_err(|error| {
179                super::spill_io_error("create spill directory", &memory.spill_directory, error)
180            })?;
181        let directory = file_descriptors
182            .canonicalize(&memory.spill_directory)
183            .map_err(|error| {
184                super::spill_io_error("resolve spill directory", &memory.spill_directory, error)
185            })?;
186        let limits = SpillPoolLimits::from_config(memory);
187        let mut pools = lock_unpoisoned(SPILL_POOLS.get_or_init(Default::default));
188        if let Some(shared) = pools.get(&directory) {
189            lock_unpoisoned(&shared.state).limits.tighten(limits);
190            return Ok(Self {
191                shared: Arc::clone(shared),
192                file_descriptors,
193            });
194        }
195        let orphan_cleanup = cleanup_orphan_files(
196            &directory,
197            memory.spill_orphan_grace_period,
198            process_marker(),
199            &file_descriptors,
200        )?;
201        let shared = Arc::new(SharedSpillPool {
202            directory: directory.clone(),
203            state: Mutex::new(SpillPoolState {
204                limits,
205                active_bytes: 0,
206                peak_active_bytes: 0,
207                pending_write_bytes: 0,
208                last_probed_available_bytes: None,
209                unreflected_reserved_bytes: 0,
210                bytes_reserved_since_probe: 0,
211                free_space_probe_in_progress: false,
212                free_space_probe_count: 0,
213                free_space_probe_failures: 0,
214                active_runs: 0,
215                peak_active_runs: 0,
216                orphan_cleanup,
217                run_delete_failures: 0,
218            }),
219            free_space_probe_ready: Condvar::new(),
220            free_space_probe,
221        });
222        pools.insert(directory, Arc::clone(&shared));
223        Ok(Self {
224            shared,
225            file_descriptors,
226        })
227    }
228
229    pub(super) fn directory(&self) -> &Path {
230        &self.shared.directory
231    }
232
233    pub(super) fn file_descriptors(&self) -> &FileOpenContext {
234        &self.file_descriptors
235    }
236
237    pub(super) fn begin_run(&self, operator: &str) -> Result<()> {
238        let mut state = lock_unpoisoned(&self.shared.state);
239        if state.active_runs >= state.limits.max_total_runs {
240            return Err(HawDBError::from(ResourceBudgetError {
241                resource: ResourceBudgetKind::SpillRuns,
242                scope: ResourceBudgetScope::SharedPool,
243                owner: operator.to_string(),
244                memory_class: None,
245                requested: 1,
246                used: state.active_runs as u64,
247                budget: state.limits.max_total_runs as u64,
248                budget_option: "ExecutionMemoryConfig.max_total_spill_runs",
249                message: format!(
250                    "{operator} exceeded shared max_total_spill_runs {}",
251                    state.limits.max_total_runs
252                ),
253            }));
254        }
255        state.active_runs = state.active_runs.saturating_add(1);
256        state.peak_active_runs = state.peak_active_runs.max(state.active_runs);
257        Ok(())
258    }
259
260    pub(super) fn cancel_run(&self) {
261        let mut state = lock_unpoisoned(&self.shared.state);
262        state.active_runs = state.active_runs.saturating_sub(1);
263    }
264
265    pub(crate) fn reserve_bytes(
266        &self,
267        operator: &str,
268        bytes: u64,
269    ) -> Result<SpillWriteReservation> {
270        loop {
271            let mut state = lock_unpoisoned(&self.shared.state);
272            let next_active = state.active_bytes.checked_add(bytes);
273            if next_active.is_none_or(|next| next > state.limits.max_total_bytes) {
274                return Err(HawDBError::from(ResourceBudgetError {
275                    resource: ResourceBudgetKind::SpillBytes,
276                    scope: ResourceBudgetScope::SharedPool,
277                    owner: operator.to_string(),
278                    memory_class: None,
279                    requested: bytes,
280                    used: state.active_bytes,
281                    budget: state.limits.max_total_bytes,
282                    budget_option: "ExecutionMemoryConfig.max_total_spill_bytes",
283                    message: format!(
284                    "{operator} exceeded shared max_total_spill_bytes {} (next active total {})",
285                    state.limits.max_total_bytes,
286                    state.active_bytes.saturating_add(bytes)
287                    ),
288                }));
289            }
290            let next_active = next_active.expect("validated spill reservation fits u64");
291            let probe_required = state.last_probed_available_bytes.is_none()
292                || state.bytes_reserved_since_probe >= state.limits.free_space_probe_interval_bytes;
293            if probe_required {
294                if state.free_space_probe_in_progress {
295                    drop(wait_unpoisoned(&self.shared.free_space_probe_ready, state));
296                    continue;
297                }
298                state.free_space_probe_in_progress = true;
299                let pending_write_bytes = state.pending_write_bytes;
300                drop(state);
301
302                let probe_result = self
303                    .shared
304                    .free_space_probe
305                    .available_space(&self.shared.directory);
306
307                let mut state = lock_unpoisoned(&self.shared.state);
308                state.free_space_probe_in_progress = false;
309                state.free_space_probe_count = state.free_space_probe_count.saturating_add(1);
310                match probe_result {
311                    Ok(available) => {
312                        state.last_probed_available_bytes = Some(available);
313                        state.unreflected_reserved_bytes = pending_write_bytes;
314                        state.bytes_reserved_since_probe = 0;
315                        self.shared.free_space_probe_ready.notify_all();
316                    }
317                    Err(error) => {
318                        state.free_space_probe_failures =
319                            state.free_space_probe_failures.saturating_add(1);
320                        self.shared.free_space_probe_ready.notify_all();
321                        drop(state);
322                        return Err(HawDBError::Execution(format!(
323                            "failed to inspect free space for spill directory '{}': {error}",
324                            self.shared.directory.display()
325                        )));
326                    }
327                }
328                continue;
329            }
330            let Some(available) = state.last_probed_available_bytes else {
331                continue;
332            };
333            let required = state
334                .limits
335                .min_free_bytes
336                .checked_add(state.unreflected_reserved_bytes)
337                .and_then(|required| required.checked_add(bytes));
338            if required.is_none_or(|required| available < required) {
339                return Err(HawDBError::from(ResourceBudgetError {
340                    resource: ResourceBudgetKind::SpillBytes,
341                    scope: ResourceBudgetScope::FileSystem,
342                    owner: operator.to_string(),
343                    memory_class: None,
344                    requested: bytes,
345                    used: state.unreflected_reserved_bytes,
346                    budget: available.saturating_sub(state.limits.min_free_bytes),
347                    budget_option: "ExecutionMemoryConfig.min_spill_free_bytes",
348                    message: format!(
349                    "{operator} cannot preserve min_spill_free_bytes {}: filesystem has {available} bytes available and {bytes} bytes were requested",
350                    state.limits.min_free_bytes
351                    ),
352                }));
353            }
354            state.active_bytes = next_active;
355            state.peak_active_bytes = state.peak_active_bytes.max(next_active);
356            state.pending_write_bytes = state.pending_write_bytes.saturating_add(bytes);
357            state.unreflected_reserved_bytes =
358                state.unreflected_reserved_bytes.saturating_add(bytes);
359            state.bytes_reserved_since_probe =
360                state.bytes_reserved_since_probe.saturating_add(bytes);
361            drop(state);
362            return Ok(SpillWriteReservation {
363                pool: self.clone(),
364                bytes,
365                committed: false,
366            });
367        }
368    }
369
370    fn commit_write(&self, bytes: u64) {
371        let mut state = lock_unpoisoned(&self.shared.state);
372        state.pending_write_bytes = state.pending_write_bytes.saturating_sub(bytes);
373    }
374
375    fn rollback_write(&self, bytes: u64) {
376        let mut state = lock_unpoisoned(&self.shared.state);
377        state.pending_write_bytes = state.pending_write_bytes.saturating_sub(bytes);
378        state.active_bytes = state.active_bytes.saturating_sub(bytes);
379    }
380
381    fn finish_run(&self, bytes: u64, pending_bytes: u64, deleted: bool) {
382        let mut state = lock_unpoisoned(&self.shared.state);
383        state.pending_write_bytes = state.pending_write_bytes.saturating_sub(pending_bytes);
384        if deleted {
385            state.active_bytes = state.active_bytes.saturating_sub(bytes);
386            state.active_runs = state.active_runs.saturating_sub(1);
387        } else {
388            state.run_delete_failures = state.run_delete_failures.saturating_add(1);
389        }
390    }
391
392    fn snapshot(&self) -> SpillPoolSnapshot {
393        let state = lock_unpoisoned(&self.shared.state);
394        SpillPoolSnapshot {
395            max_total_bytes: state.limits.max_total_bytes,
396            max_total_runs: state.limits.max_total_runs,
397            min_free_bytes: state.limits.min_free_bytes,
398            free_space_probe_interval_bytes: state.limits.free_space_probe_interval_bytes,
399            free_space_probe_count: state.free_space_probe_count,
400            free_space_probe_failures: state.free_space_probe_failures,
401            active_bytes: state.active_bytes,
402            peak_active_bytes: state.peak_active_bytes,
403            pending_write_bytes: state.pending_write_bytes,
404            active_runs: state.active_runs,
405            peak_active_runs: state.peak_active_runs,
406            orphan_files_removed: state.orphan_cleanup.files_removed,
407            orphan_bytes_removed: state.orphan_cleanup.bytes_removed,
408            orphan_cleanup_failures: state.orphan_cleanup.failures,
409            run_delete_failures: state.run_delete_failures,
410        }
411    }
412}
413
414pub(crate) struct SpillWriteReservation {
415    pool: SpillPool,
416    bytes: u64,
417    committed: bool,
418}
419
420impl SpillWriteReservation {
421    pub(super) fn commit(mut self, lease: &RunLease) {
422        lease
423            .reserved_bytes
424            .fetch_add(self.bytes, Ordering::Relaxed);
425        lease
426            .pending_write_bytes
427            .fetch_add(self.bytes, Ordering::Relaxed);
428        self.committed = true;
429    }
430}
431
432impl Drop for SpillWriteReservation {
433    fn drop(&mut self) {
434        if !self.committed {
435            self.pool.rollback_write(self.bytes);
436        }
437    }
438}
439
440#[derive(Debug)]
441pub(super) struct RunLease {
442    path: PathBuf,
443    pool: SpillPool,
444    reserved_bytes: AtomicU64,
445    pending_write_bytes: AtomicU64,
446}
447
448impl RunLease {
449    pub(super) fn new(path: PathBuf, pool: SpillPool) -> Self {
450        Self {
451            path,
452            pool,
453            reserved_bytes: AtomicU64::new(0),
454            pending_write_bytes: AtomicU64::new(0),
455        }
456    }
457
458    pub(super) fn path(&self) -> &Path {
459        &self.path
460    }
461
462    pub(super) fn file_descriptors(&self) -> &FileOpenContext {
463        self.pool.file_descriptors()
464    }
465
466    pub(super) fn mark_flushed(&self) {
467        let pending_bytes = self.pending_write_bytes.swap(0, Ordering::Relaxed);
468        self.pool.commit_write(pending_bytes);
469    }
470}
471
472impl Drop for RunLease {
473    fn drop(&mut self) {
474        let deleted = match self.file_descriptors().remove_file(&self.path) {
475            Ok(()) => true,
476            Err(error) if error.kind() == ErrorKind::NotFound => true,
477            Err(_) => false,
478        };
479        self.pool.finish_run(
480            self.reserved_bytes.load(Ordering::Relaxed),
481            self.pending_write_bytes.load(Ordering::Relaxed),
482            deleted,
483        );
484    }
485}
486
487pub(crate) fn spill_pool_snapshot(memory: &ExecutionMemoryConfig) -> Result<SpillPoolSnapshot> {
488    Ok(SpillPool::open(memory)?.snapshot())
489}
490
491fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
492    mutex
493        .lock()
494        .unwrap_or_else(std::sync::PoisonError::into_inner)
495}
496
497fn wait_unpoisoned<'a, T>(condvar: &Condvar, guard: MutexGuard<'a, T>) -> MutexGuard<'a, T> {
498    condvar
499        .wait(guard)
500        .unwrap_or_else(std::sync::PoisonError::into_inner)
501}
502
503pub(super) fn process_marker() -> &'static str {
504    PROCESS_MARKER.get_or_init(|| {
505        let started = SystemTime::now()
506            .duration_since(UNIX_EPOCH)
507            .unwrap_or(Duration::ZERO)
508            .as_nanos();
509        format!("{}-{started}", std::process::id())
510    })
511}
512
513fn cleanup_orphan_files(
514    directory: &Path,
515    grace_period: Duration,
516    current_process_marker: &str,
517    file_descriptors: &FileOpenContext,
518) -> Result<OrphanCleanupStats> {
519    let mut stats = OrphanCleanupStats::default();
520    let current_prefix = format!("{SPILL_FILE_PREFIX}{current_process_marker}-");
521    let entries = file_descriptors
522        .read_dir(directory)
523        .map_err(|error| super::spill_io_error("inspect spill directory", directory, error))?;
524    for entry in entries {
525        let entry = match entry {
526            Ok(entry) => entry,
527            Err(_) => {
528                stats.failures = stats.failures.saturating_add(1);
529                continue;
530            }
531        };
532        let name = entry.file_name();
533        let Some(name) = name.to_str() else {
534            continue;
535        };
536        if !name.starts_with(SPILL_FILE_PREFIX)
537            || !name.ends_with(SPILL_FILE_SUFFIX)
538            || name.starts_with(&current_prefix)
539        {
540            continue;
541        }
542        let metadata = match entry.metadata() {
543            Ok(metadata) if metadata.is_file() => metadata,
544            Ok(_) => continue,
545            Err(_) => {
546                stats.failures = stats.failures.saturating_add(1);
547                continue;
548            }
549        };
550        let old_enough = metadata
551            .modified()
552            .ok()
553            .and_then(|modified| SystemTime::now().duration_since(modified).ok())
554            .is_some_and(|age| age >= grace_period);
555        if !old_enough {
556            continue;
557        }
558        match file_descriptors.remove_file(&entry.path()) {
559            Ok(()) => {
560                stats.files_removed = stats.files_removed.saturating_add(1);
561                stats.bytes_removed = stats.bytes_removed.saturating_add(metadata.len());
562            }
563            Err(_) => stats.failures = stats.failures.saturating_add(1),
564        }
565    }
566    Ok(stats)
567}
568
569#[cfg(test)]
570mod tests {
571    use super::*;
572    use std::num::NonZeroU64;
573    use std::sync::mpsc;
574    use std::thread;
575
576    static TEST_SEQUENCE: AtomicU64 = AtomicU64::new(0);
577
578    #[derive(Debug)]
579    struct AdjustableSpaceProbe {
580        available: AtomicU64,
581        calls: AtomicU64,
582    }
583
584    impl AdjustableSpaceProbe {
585        fn new(available: u64) -> Self {
586            Self {
587                available: AtomicU64::new(available),
588                calls: AtomicU64::new(0),
589            }
590        }
591    }
592
593    impl SpillSpaceProbe for AdjustableSpaceProbe {
594        fn available_space(&self, _directory: &Path) -> std::io::Result<u64> {
595            self.calls.fetch_add(1, Ordering::Relaxed);
596            Ok(self.available.load(Ordering::Relaxed))
597        }
598    }
599
600    #[derive(Debug, Default)]
601    struct BlockingSpaceProbe {
602        state: Mutex<(bool, bool)>,
603        ready: Condvar,
604        calls: AtomicU64,
605    }
606
607    impl BlockingSpaceProbe {
608        fn wait_until_started(&self) {
609            let mut state = lock_unpoisoned(&self.state);
610            while !state.0 {
611                state = wait_unpoisoned(&self.ready, state);
612            }
613        }
614
615        fn release(&self) {
616            let mut state = lock_unpoisoned(&self.state);
617            state.1 = true;
618            self.ready.notify_all();
619        }
620    }
621
622    impl SpillSpaceProbe for BlockingSpaceProbe {
623        fn available_space(&self, _directory: &Path) -> std::io::Result<u64> {
624            self.calls.fetch_add(1, Ordering::Relaxed);
625            let mut state = lock_unpoisoned(&self.state);
626            state.0 = true;
627            self.ready.notify_all();
628            while !state.1 {
629                state = wait_unpoisoned(&self.ready, state);
630            }
631            Ok(u64::MAX)
632        }
633    }
634
635    fn test_memory(name: &str, probe_interval_bytes: u64) -> ExecutionMemoryConfig {
636        let sequence = TEST_SEQUENCE.fetch_add(1, Ordering::Relaxed);
637        ExecutionMemoryConfig {
638            min_spill_free_bytes: NonZeroU64::new(1).unwrap(),
639            spill_free_space_probe_interval_bytes: NonZeroU64::new(probe_interval_bytes).unwrap(),
640            spill_directory: std::env::temp_dir().join(format!(
641                "hawdb-spill-pool-test-{}-{name}-{sequence}",
642                std::process::id()
643            )),
644            ..ExecutionMemoryConfig::default()
645        }
646    }
647
648    #[test]
649    fn free_space_probe_is_amortized_by_reserved_bytes() {
650        let memory = test_memory("watermark", 64);
651        let probe = Arc::new(AdjustableSpaceProbe::new(1024 * 1024));
652        let pool = SpillPool::open_with_probe(&memory, probe.clone()).unwrap();
653
654        for _ in 0..8 {
655            drop(pool.reserve_bytes("test", 8).unwrap());
656        }
657        let snapshot = pool.snapshot();
658        assert_eq!(snapshot.free_space_probe_interval_bytes, 64);
659        assert_eq!(snapshot.free_space_probe_count, 1);
660        assert_eq!(probe.calls.load(Ordering::Relaxed), 1);
661
662        drop(pool.reserve_bytes("test", 1).unwrap());
663        assert_eq!(pool.snapshot().free_space_probe_count, 2);
664        assert_eq!(probe.calls.load(Ordering::Relaxed), 2);
665
666        std::fs::remove_dir(&memory.spill_directory).unwrap();
667    }
668
669    #[test]
670    fn refreshed_free_space_still_enforces_the_reserve() {
671        let memory = test_memory("exhausted", 8);
672        let probe = Arc::new(AdjustableSpaceProbe::new(128));
673        let pool = SpillPool::open_with_probe(&memory, probe.clone()).unwrap();
674
675        drop(pool.reserve_bytes("test", 8).unwrap());
676        probe.available.store(8, Ordering::Relaxed);
677        let error = match pool.reserve_bytes("test", 8) {
678            Ok(_) => panic!("exhausted free space must reject a reservation"),
679            Err(error) => error,
680        };
681
682        assert!(error.to_string().contains("min_spill_free_bytes"));
683        let HawDBError::ResourceBudgetExceeded(details) = error else {
684            panic!("expected typed free-space rejection");
685        };
686        assert_eq!(details.scope, ResourceBudgetScope::FileSystem);
687        assert_eq!(details.resource, ResourceBudgetKind::SpillBytes);
688        assert_eq!((details.requested, details.used, details.budget), (8, 0, 7));
689        assert_eq!(
690            details.budget_option,
691            "ExecutionMemoryConfig.min_spill_free_bytes"
692        );
693        assert_eq!(pool.snapshot().active_bytes, 0);
694        assert_eq!(pool.snapshot().free_space_probe_count, 2);
695        assert_eq!(probe.calls.load(Ordering::Relaxed), 2);
696        std::fs::remove_dir(&memory.spill_directory).unwrap();
697    }
698
699    #[test]
700    fn concurrent_reservations_share_a_probe_without_holding_the_pool_mutex() {
701        let memory = test_memory("concurrent", 1024);
702        let probe = Arc::new(BlockingSpaceProbe::default());
703        let pool = Arc::new(SpillPool::open_with_probe(&memory, probe.clone()).unwrap());
704
705        let first_pool = Arc::clone(&pool);
706        let first = thread::spawn(move || first_pool.reserve_bytes("first", 8));
707        probe.wait_until_started();
708        let second_pool = Arc::clone(&pool);
709        let second = thread::spawn(move || second_pool.reserve_bytes("second", 8));
710
711        let snapshot_pool = Arc::clone(&pool);
712        let (snapshot_tx, snapshot_rx) = mpsc::channel();
713        let snapshot = thread::spawn(move || {
714            snapshot_pool.snapshot();
715            snapshot_tx.send(()).unwrap();
716        });
717        snapshot_rx
718            .recv_timeout(Duration::from_secs(1))
719            .expect("pool mutex must remain available during a free-space probe");
720
721        probe.release();
722        drop(first.join().unwrap().unwrap());
723        drop(second.join().unwrap().unwrap());
724        snapshot.join().unwrap();
725        assert_eq!(pool.snapshot().free_space_probe_count, 1);
726        assert_eq!(probe.calls.load(Ordering::Relaxed), 1);
727        std::fs::remove_dir(&memory.spill_directory).unwrap();
728    }
729
730    #[test]
731    fn shared_byte_overflow_refuses_without_probing_or_charging() {
732        let mut memory = test_memory("shared-overflow", 64);
733        memory.max_total_spill_bytes = NonZeroU64::MAX;
734        let probe = Arc::new(AdjustableSpaceProbe::new(u64::MAX));
735        let pool = SpillPool::open_with_probe(&memory, probe.clone()).unwrap();
736        // Model a nearly full shared pool without creating an exabyte file.
737        lock_unpoisoned(&pool.shared.state).active_bytes = u64::MAX - 8;
738        let before = pool.snapshot();
739        let error = pool.reserve_bytes("overflow", 16).err().unwrap();
740        let HawDBError::ResourceBudgetExceeded(details) = error else {
741            panic!("expected typed shared byte overflow");
742        };
743        assert_eq!(details.scope, ResourceBudgetScope::SharedPool);
744        assert_eq!(
745            (details.requested, details.used, details.budget),
746            (16, u64::MAX - 8, u64::MAX)
747        );
748        assert_eq!(pool.snapshot(), before);
749        assert_eq!(probe.calls.load(Ordering::Relaxed), 0);
750        std::fs::remove_dir(&memory.spill_directory).unwrap();
751    }
752
753    #[test]
754    fn free_space_reserve_overflow_refuses_at_maximum_available_bytes() {
755        let mut memory = test_memory("free-space-overflow", 64);
756        memory.max_total_spill_bytes = NonZeroU64::MAX;
757        let pool =
758            SpillPool::open_with_probe(&memory, Arc::new(AdjustableSpaceProbe::new(u64::MAX)))
759                .unwrap();
760        let error = pool.reserve_bytes("overflow", u64::MAX).err().unwrap();
761        let HawDBError::ResourceBudgetExceeded(details) = error else {
762            panic!("expected typed free-space overflow");
763        };
764        assert_eq!(details.scope, ResourceBudgetScope::FileSystem);
765        assert_eq!(
766            (details.requested, details.used, details.budget),
767            (u64::MAX, 0, u64::MAX - 1)
768        );
769        let snapshot = pool.snapshot();
770        assert_eq!(
771            (
772                snapshot.active_bytes,
773                snapshot.pending_write_bytes,
774                snapshot.peak_active_bytes
775            ),
776            (0, 0, 0)
777        );
778        std::fs::remove_dir(&memory.spill_directory).unwrap();
779    }
780}
hawdb_executor/spill/pool.rs

hawdb_executor/spill/
pool.rs
