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

hawdb_executor/spill/
pool.rs
