linkrs_query/executor/streaming/spill/
quota.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2use std::sync::Arc;
3
4use linkrs_core::error::QueryError;
5
6// ── DiskQuota (separate from memory budget) ──────────────────────────────────
7
8/// Tracks disk space used by spill operations for a single query.
9///
10/// Disk quota is independent of the memory budget.  Exceeding disk quota
11/// produces a structured error rather than a silent spill-to-nowhere.
12#[derive(Debug, Clone)]
13pub struct DiskQuota {
14    max_bytes: u64,
15    used: Arc<AtomicU64>,
16}
17
18impl DiskQuota {
19    /// Create a quota with the given byte limit.  `0` = unlimited.
20    pub fn new(max_bytes: u64) -> Self {
21        Self {
22            max_bytes,
23            used: Arc::new(AtomicU64::new(0)),
24        }
25    }
26
27    /// Default quota: 2 GiB per query.
28    pub fn default_quota() -> Self {
29        Self::new(2 * 1024 * 1024 * 1024)
30    }
31
32    /// Try to reserve `bytes` of disk space.  Returns an error when the
33    /// quota would be exceeded.
34    pub fn try_reserve(&self, bytes: u64) -> Result<(), QueryError> {
35        if self.max_bytes == 0 {
36            // unlimited
37            let _ = self.used.fetch_add(bytes, Ordering::Relaxed);
38            return Ok(());
39        }
40        let mut prev = self.used.load(Ordering::Relaxed);
41        loop {
42            let total = prev.checked_add(bytes).ok_or_else(|| {
43                QueryError::execution(format!(
44                    "Disk quota overflow: request {} bytes overflows u64",
45                    bytes,
46                ))
47            })?;
48            if total > self.max_bytes {
49                return Err(QueryError::execution(format!(
50                    "Disk quota exceeded: request {} bytes, total {} > quota {} bytes",
51                    bytes, total, self.max_bytes,
52                )));
53            }
54            match self
55                .used
56                .compare_exchange_weak(prev, total, Ordering::Relaxed, Ordering::Relaxed)
57            {
58                Ok(_) => return Ok(()),
59                Err(current) => prev = current,
60            }
61        }
62    }
63
64    /// Release `bytes` of disk space.
65    pub fn release(&self, bytes: u64) {
66        self.used.fetch_sub(bytes, Ordering::Relaxed);
67    }
68
69    /// Current used disk space.
70    pub fn current(&self) -> u64 {
71        self.used.load(Ordering::Relaxed)
72    }
73
74    /// Maximum allowed disk space.
75    pub fn max(&self) -> u64 {
76        self.max_bytes
77    }
78}
linkrs_query/executor/streaming/spill/quota.rs

linkrs_query/executor/streaming/spill/
quota.rs
