query_runtime/runtime/
cache_fill.rs

1//! Service-wide lakehouse cache-fill resource policy and gates.
2
3use std::{
4    sync::{Arc, Condvar, Mutex, PoisonError},
5    time::{Duration, Instant},
6};
7
8use query_planning::error::{QueryError, QueryErrorCode, Result};
9
10use super::ExecutionCancellation;
11
12/// Service-wide resources used only by complete-object lakehouse cache fills.
13#[derive(Debug, Clone, PartialEq, Eq, Hash)]
14pub struct LakehouseCacheFillConfig {
15    /// Maximum complete-object fill bytes held concurrently.
16    /// `None` uses one quarter of detected host RAM, capped at 16 GiB,
17    /// with a 1 GiB fallback when RAM detection is unavailable.
18    pub budget_bytes: Option<u64>,
19    /// Timeout for blocking byte-budget and operation-gate acquisition.
20    pub acquire_timeout: Duration,
21    /// Max concurrent complete-object cache maintenance operations.
22    pub max_concurrent_ops: usize,
23}
24
25impl LakehouseCacheFillConfig {
26    pub(crate) const DEFAULT_BUDGET_BYTES: u64 = 1024 * 1024 * 1024;
27    pub(crate) const DEFAULT_BUDGET_CAP_BYTES: u64 = 16 * 1024 * 1024 * 1024;
28    /// Default wait bound for cache-fill byte and operation permits, 30 seconds.
29    pub const DEFAULT_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(30);
30    pub(crate) const DEFAULT_MAX_CONCURRENT_OPS: usize = 4;
31
32    /// Validates the service-wide fill policy.
33    /// Returns `InvalidConfiguration` for a zero explicit byte budget, zero
34    /// acquisition timeout or zero concurrent operations.
35    pub fn validate(&self) -> Result<()> {
36        if let Some(budget_bytes) = self.budget_bytes {
37            validate_budget_bytes(budget_bytes)?;
38        }
39        if self.acquire_timeout.is_zero() {
40            return Err(QueryError::new(
41                QueryErrorCode::InvalidConfiguration,
42                "query_engine.lakehouse_cache.fill_config.validate",
43                "lakehouse cache fill acquire timeout must be greater than zero",
44            )
45            .with_fact("field", "lakehouse_cache_fill_acquire_timeout")
46            .with_fact("reason", "zero"));
47        }
48        if self.max_concurrent_ops == 0 {
49            return Err(QueryError::new(
50                QueryErrorCode::InvalidConfiguration,
51                "query_engine.lakehouse_cache.fill_config.validate",
52                "lakehouse cache fill max concurrent operations must be greater than zero",
53            )
54            .with_fact("field", "lakehouse_cache_fill_max_concurrent_ops")
55            .with_fact("reason", "zero"));
56        }
57        Ok(())
58    }
59}
60
61impl Default for LakehouseCacheFillConfig {
62    fn default() -> Self {
63        Self {
64            budget_bytes: None,
65            acquire_timeout: Self::DEFAULT_ACQUIRE_TIMEOUT,
66            max_concurrent_ops: Self::DEFAULT_MAX_CONCURRENT_OPS,
67        }
68    }
69}
70
71pub(crate) fn default_lakehouse_cache_fill_budget_bytes() -> u64 {
72    default_lakehouse_cache_fill_budget_bytes_from_physical_memory(detect_physical_memory_bytes())
73}
74
75pub(super) fn default_lakehouse_cache_fill_budget_bytes_from_physical_memory(
76    physical_memory_bytes: Option<u64>,
77) -> u64 {
78    physical_memory_bytes
79        .and_then(|bytes| bytes.checked_div(4))
80        .map(|bytes| bytes.min(LakehouseCacheFillConfig::DEFAULT_BUDGET_CAP_BYTES))
81        .filter(|bytes| *bytes > 0)
82        .unwrap_or(LakehouseCacheFillConfig::DEFAULT_BUDGET_BYTES)
83}
84
85fn detect_physical_memory_bytes() -> Option<u64> {
86    let meminfo = std::fs::read_to_string("/proc/meminfo").ok()?;
87    let line = meminfo.lines().find(|line| line.starts_with("MemTotal:"))?;
88    let kib = line.split_whitespace().nth(1)?.parse::<u64>().ok()?;
89    kib.checked_mul(1024)
90}
91
92#[derive(Debug)]
93struct CacheFillByteBudgetInner {
94    budget_bytes: u64,
95    state: Mutex<CacheFillByteBudgetState>,
96    available: Condvar,
97}
98
99#[derive(Debug, Default)]
100struct CacheFillByteBudgetState {
101    in_flight_bytes: u64,
102}
103
104/// Service-wide byte budget for complete-object lakehouse cache maintenance.
105#[derive(Debug, Clone)]
106pub struct CacheFillByteBudget {
107    inner: Arc<CacheFillByteBudgetInner>,
108}
109
110impl CacheFillByteBudget {
111    pub fn new(budget_bytes: u64) -> Result<Self> {
112        validate_budget_bytes(budget_bytes)?;
113        Ok(Self {
114            inner: Arc::new(CacheFillByteBudgetInner {
115                budget_bytes,
116                state: Mutex::new(CacheFillByteBudgetState::default()),
117                available: Condvar::new(),
118            }),
119        })
120    }
121
122    pub fn budget_bytes(&self) -> u64 {
123        self.inner.budget_bytes
124    }
125
126    #[cfg(test)]
127    pub(crate) fn try_acquire(&self, bytes: u64) -> Result<Option<CacheFillBytePermit>> {
128        validate_request_bytes(bytes, self.budget_bytes())?;
129        let mut state = self.lock_state("query_engine.lakehouse_cache.fill_budget.try_acquire")?;
130        if !bytes_fit(state.in_flight_bytes, bytes, self.budget_bytes()) {
131            return Ok(None);
132        }
133        state.in_flight_bytes = state.in_flight_bytes.saturating_add(bytes);
134        Ok(Some(CacheFillBytePermit { budget: self.clone(), bytes }))
135    }
136
137    pub fn acquire_with_timeout(
138        &self,
139        cancellation: &ExecutionCancellation,
140        bytes: u64,
141        timeout: Duration,
142        operation: &'static str,
143    ) -> Result<Option<CacheFillBytePermit>> {
144        validate_request_bytes(bytes, self.budget_bytes())?;
145        let start = Instant::now();
146        let mut state = self.lock_state(operation)?;
147        loop {
148            cancellation.check(operation)?;
149            if bytes_fit(state.in_flight_bytes, bytes, self.budget_bytes()) {
150                state.in_flight_bytes = state.in_flight_bytes.saturating_add(bytes);
151                return Ok(Some(CacheFillBytePermit { budget: self.clone(), bytes }));
152            }
153            let Some(remaining) = timeout.checked_sub(start.elapsed()) else {
154                return Ok(None);
155            };
156            let wait = remaining.min(Duration::from_millis(100));
157            let (next_state, _) = self.inner.available.wait_timeout(state, wait).map_err(|_| {
158                QueryError::new(
159                    QueryErrorCode::Internal,
160                    operation,
161                    "lakehouse cache fill byte-budget mutex was poisoned while waiting",
162                )
163            })?;
164            state = next_state;
165        }
166    }
167
168    fn try_grow_permit(&self, permit: &mut CacheFillBytePermit, new_bytes: u64) -> Result<bool> {
169        if new_bytes <= permit.bytes {
170            return Ok(true);
171        }
172        validate_request_bytes(new_bytes, self.budget_bytes())?;
173        let delta = new_bytes - permit.bytes;
174        let mut state = self.lock_state("query_engine.lakehouse_cache.fill_budget.grow")?;
175        // Growth never waits: two fills that each hold part of the budget and
176        // wait for the remainder deadlock while holding their shares and the
177        // object-store worker. A fill that cannot grow abandons its cache
178        // write instead; the reader still gets its data.
179        if !bytes_fit(state.in_flight_bytes, delta, self.budget_bytes()) {
180            return Ok(false);
181        }
182        state.in_flight_bytes = state.in_flight_bytes.saturating_add(delta);
183        permit.bytes = new_bytes;
184        Ok(true)
185    }
186
187    fn release(&self, bytes: u64) {
188        // The state is one integer; skipping the release on a poisoned lock would leak these bytes for good.
189        let mut state = self.inner.state.lock().unwrap_or_else(PoisonError::into_inner);
190        state.in_flight_bytes = state.in_flight_bytes.saturating_sub(bytes);
191        self.inner.available.notify_all();
192    }
193
194    fn lock_state(&self, operation: &'static str) -> Result<std::sync::MutexGuard<'_, CacheFillByteBudgetState>> {
195        self.inner.state.lock().map_err(|_| {
196            QueryError::new(QueryErrorCode::Internal, operation, "lakehouse cache fill byte-budget mutex was poisoned")
197        })
198    }
199}
200
201impl Default for CacheFillByteBudget {
202    fn default() -> Self {
203        Self::new(default_lakehouse_cache_fill_budget_bytes()).expect("default cache fill byte budget is non-zero")
204    }
205}
206
207/// RAII permit for bytes charged to a cache-fill byte budget.
208#[derive(Debug)]
209pub struct CacheFillBytePermit {
210    budget: CacheFillByteBudget,
211    bytes: u64,
212}
213
214impl CacheFillBytePermit {
215    pub fn bytes(&self) -> u64 {
216        self.bytes
217    }
218
219    /// Grows this permit without waiting. Returns `false` when the budget
220    /// cannot cover the growth; the caller abandons its cache write and the
221    /// reader still gets its data.
222    pub fn try_grow(&mut self, new_bytes: u64) -> Result<bool> {
223        let budget = self.budget.clone();
224        budget.try_grow_permit(self, new_bytes)
225    }
226}
227
228impl Drop for CacheFillBytePermit {
229    fn drop(&mut self) {
230        self.budget.release(self.bytes);
231    }
232}
233
234#[derive(Debug)]
235struct CacheFillConcurrencyGateInner {
236    max_concurrent_ops: usize,
237    state: Mutex<CacheFillConcurrencyGateState>,
238    available: Condvar,
239}
240
241#[derive(Debug, Default)]
242struct CacheFillConcurrencyGateState {
243    in_flight_ops: usize,
244}
245
246#[derive(Debug, Clone)]
247pub(super) struct CacheFillConcurrencyGate {
248    inner: Arc<CacheFillConcurrencyGateInner>,
249}
250
251impl CacheFillConcurrencyGate {
252    pub(super) fn new(max_concurrent_ops: usize) -> Result<Self> {
253        let config = LakehouseCacheFillConfig { max_concurrent_ops, ..Default::default() };
254        config.validate()?;
255        Ok(Self {
256            inner: Arc::new(CacheFillConcurrencyGateInner {
257                max_concurrent_ops,
258                state: Mutex::new(CacheFillConcurrencyGateState::default()),
259                available: Condvar::new(),
260            }),
261        })
262    }
263
264    pub(super) fn max_concurrent_ops(&self) -> usize {
265        self.inner.max_concurrent_ops
266    }
267
268    pub(super) fn acquire_with_timeout(
269        &self,
270        cancellation: &ExecutionCancellation,
271        timeout: Duration,
272        operation: &'static str,
273    ) -> Result<Option<CacheFillOperationPermit>> {
274        let start = Instant::now();
275        let mut state = self.lock_state(operation)?;
276        loop {
277            cancellation.check(operation)?;
278            if state.in_flight_ops < self.max_concurrent_ops() {
279                state.in_flight_ops = state.in_flight_ops.saturating_add(1);
280                return Ok(Some(CacheFillOperationPermit { gate: self.clone() }));
281            }
282            let Some(remaining) = timeout.checked_sub(start.elapsed()) else {
283                return Ok(None);
284            };
285            let wait = remaining.min(Duration::from_millis(100));
286            let (next_state, _) = self.inner.available.wait_timeout(state, wait).map_err(|_| {
287                QueryError::new(
288                    QueryErrorCode::Internal,
289                    operation,
290                    "lakehouse cache fill concurrency mutex was poisoned while waiting",
291                )
292            })?;
293            state = next_state;
294        }
295    }
296
297    fn release(&self) {
298        let mut state = self.inner.state.lock().unwrap_or_else(PoisonError::into_inner);
299        state.in_flight_ops = state.in_flight_ops.saturating_sub(1);
300        self.inner.available.notify_one();
301    }
302
303    fn lock_state(&self, operation: &'static str) -> Result<std::sync::MutexGuard<'_, CacheFillConcurrencyGateState>> {
304        self.inner.state.lock().map_err(|_| {
305            QueryError::new(QueryErrorCode::Internal, operation, "lakehouse cache fill concurrency mutex was poisoned")
306        })
307    }
308}
309
310pub struct CacheFillOperationPermit {
311    gate: CacheFillConcurrencyGate,
312}
313
314impl Drop for CacheFillOperationPermit {
315    fn drop(&mut self) {
316        self.gate.release();
317    }
318}
319
320fn bytes_fit(in_flight_bytes: u64, requested_bytes: u64, budget_bytes: u64) -> bool {
321    in_flight_bytes.saturating_add(requested_bytes) <= budget_bytes
322}
323
324fn validate_budget_bytes(budget_bytes: u64) -> Result<()> {
325    if budget_bytes == 0 {
326        return Err(QueryError::new(
327            QueryErrorCode::InvalidConfiguration,
328            "query_engine.lakehouse_cache.fill_budget.validate",
329            "lakehouse cache fill budget bytes must be greater than zero",
330        )
331        .with_fact("field", "lakehouse_cache_fill_budget_bytes")
332        .with_fact("reason", "zero"));
333    }
334    Ok(())
335}
336
337fn validate_request_bytes(requested_bytes: u64, budget_bytes: u64) -> Result<()> {
338    if requested_bytes == 0 {
339        return Err(QueryError::new(
340            QueryErrorCode::InvalidPlan,
341            "query_engine.lakehouse_cache.fill_budget.request",
342            "lakehouse cache fill request bytes must be greater than zero",
343        )
344        .with_fact("field", "lakehouse_cache_fill_request_bytes")
345        .with_fact("reason", "zero"));
346    }
347    if requested_bytes > budget_bytes {
348        return Err(QueryError::new(
349            QueryErrorCode::MemoryReservationFailed,
350            "query_engine.lakehouse_cache.fill_budget.request",
351            "lakehouse cache fill exceeds its configured byte budget",
352        )
353        .with_fact("requested_bytes", requested_bytes)
354        .with_fact("budget_bytes", budget_bytes));
355    }
356    Ok(())
357}
358
359#[cfg(test)]
360mod tests {
361    use super::*;
362
363    #[test]
364    fn cache_fill_byte_budget_try_acquire_does_not_wait_when_capacity_is_held() {
365        let budget = CacheFillByteBudget::new(10).unwrap();
366        let held = budget.try_acquire(10).unwrap().unwrap();
367        assert!(budget.try_acquire(1).unwrap().is_none());
368        drop(held);
369        assert!(budget.try_acquire(1).unwrap().is_some());
370    }
371
372    #[test]
373    fn permits_return_their_share_after_the_state_mutex_is_poisoned() {
374        fn poison<T: Send>(mutex: &Mutex<T>) {
375            std::thread::scope(|scope| {
376                scope
377                    .spawn(|| {
378                        let _state = mutex.lock().unwrap();
379                        panic!("poison the state mutex");
380                    })
381                    .join()
382                    .expect_err("the poisoning thread panics while holding the lock");
383            });
384        }
385
386        let budget = CacheFillByteBudget::new(10).unwrap();
387        let gate = CacheFillConcurrencyGate::new(1).unwrap();
388        let bytes = budget.try_acquire(10).unwrap().unwrap();
389        let operation =
390            gate.acquire_with_timeout(&ExecutionCancellation::new(), Duration::from_secs(1), "test").unwrap().unwrap();
391        poison(&budget.inner.state);
392        poison(&gate.inner.state);
393
394        drop(bytes);
395        drop(operation);
396
397        let in_flight_bytes = budget.inner.state.lock().unwrap_or_else(PoisonError::into_inner).in_flight_bytes;
398        let in_flight_ops = gate.inner.state.lock().unwrap_or_else(PoisonError::into_inner).in_flight_ops;
399        assert_eq!(in_flight_bytes, 0, "the byte budget is returned");
400        assert_eq!(in_flight_ops, 0, "the concurrency slot is returned");
401    }
402}
query_runtime/runtime/cache_fill.rs

query_runtime/runtime/
cache_fill.rs
