alopex_dataframe/physical/
budget.rs

1//! Resource accounting and terminal-state contracts for bounded execution.
2//!
3//! A reservation must be acquired before this execution layer takes ownership
4//! of an allocation.  The reservation is RAII-backed so every successful
5//! reservation is released exactly once, including on error paths.
6
7use std::num::NonZeroUsize;
8use std::sync::{Arc, Mutex};
9
10use crate::{DataFrameError, Result};
11
12/// Options shared by bounded materialization and streaming execution.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub struct StreamOptions {
15    /// Maximum bytes owned by the execution pipeline at one time.
16    pub memory_limit_bytes: u64,
17    /// Maximum number of pipeline-owned batches at one time.
18    pub max_in_flight_batches: NonZeroUsize,
19    /// Requested source row cap. It never relaxes `memory_limit_bytes`.
20    pub batch_rows: NonZeroUsize,
21}
22
23impl StreamOptions {
24    /// Construct validated bounded-execution options.
25    pub fn new(
26        memory_limit_bytes: u64,
27        max_in_flight_batches: NonZeroUsize,
28        batch_rows: NonZeroUsize,
29    ) -> Self {
30        Self {
31            memory_limit_bytes,
32            max_in_flight_batches,
33            batch_rows,
34        }
35    }
36}
37
38impl Default for StreamOptions {
39    fn default() -> Self {
40        Self {
41            memory_limit_bytes: 64 * 1024 * 1024,
42            max_in_flight_batches: NonZeroUsize::MIN,
43            batch_rows: NonZeroUsize::new(8_192).expect("8,192 is non-zero"),
44        }
45    }
46}
47
48/// The allocation domain that owns a reservation.
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum ResourceScope {
51    /// Raw input bytes held by a source reader.
52    Source,
53    /// Decode buffers and Arrow arrays constructed by a source reader.
54    Decode,
55    /// A physical operator's intermediate state.
56    Operator,
57    /// Expression evaluation and memoization state.
58    Expression,
59    /// A batch waiting to be transferred to the consumer.
60    Output,
61    /// Batches retained while constructing a bounded materialized result.
62    MaterializedOutput,
63    /// Temporary spill state, if a bounded executor supports it.
64    Spill,
65}
66
67impl ResourceScope {
68    /// Stable scope name for diagnostics and telemetry.
69    pub const fn as_str(self) -> &'static str {
70        match self {
71            Self::Source => "source",
72            Self::Decode => "decode",
73            Self::Operator => "operator",
74            Self::Expression => "expression",
75            Self::Output => "output",
76            Self::MaterializedOutput => "materialized_output",
77            Self::Spill => "spill",
78        }
79    }
80}
81
82/// Snapshot of currently pipeline-owned resources.
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
84pub struct ResourceUsage {
85    /// Bytes currently reserved by the execution pipeline.
86    pub reserved_bytes: u64,
87    /// Batches currently reserved by the execution pipeline.
88    pub reserved_batches: usize,
89}
90
91#[derive(Debug)]
92struct BudgetState {
93    usage: ResourceUsage,
94}
95
96/// Shared, allocation-before-ownership resource budget.
97#[derive(Debug, Clone)]
98pub struct ResourceBudget {
99    memory_limit_bytes: u64,
100    max_in_flight_batches: usize,
101    state: Arc<Mutex<BudgetState>>,
102}
103
104impl ResourceBudget {
105    /// Create a budget from public bounded-execution options.
106    pub fn from_options(options: StreamOptions) -> Self {
107        Self::new(options.memory_limit_bytes, options.max_in_flight_batches)
108    }
109
110    /// Create a budget with a byte limit and a positive in-flight batch limit.
111    pub fn new(memory_limit_bytes: u64, max_in_flight_batches: NonZeroUsize) -> Self {
112        Self {
113            memory_limit_bytes,
114            max_in_flight_batches: max_in_flight_batches.get(),
115            state: Arc::new(Mutex::new(BudgetState {
116                usage: ResourceUsage {
117                    reserved_bytes: 0,
118                    reserved_batches: 0,
119                },
120            })),
121        }
122    }
123
124    /// Reserve byte ownership before allocating or retaining data.
125    pub fn reserve(&self, scope: ResourceScope, bytes: u64) -> Result<ResourceReservation> {
126        self.reserve_inner(scope, bytes, 0)
127    }
128
129    /// Reserve ownership for one decoded or retained batch before publishing it.
130    pub fn reserve_batch(&self, scope: ResourceScope, bytes: u64) -> Result<ResourceReservation> {
131        self.reserve_inner(scope, bytes, 1)
132    }
133
134    /// Return the configured byte limit.
135    pub const fn memory_limit_bytes(&self) -> u64 {
136        self.memory_limit_bytes
137    }
138
139    /// Return the configured pipeline-owned batch limit.
140    pub const fn max_in_flight_batches(&self) -> usize {
141        self.max_in_flight_batches
142    }
143
144    /// Obtain an accounting snapshot for tests, diagnostics, and stream status.
145    pub fn usage(&self) -> ResourceUsage {
146        self.state
147            .lock()
148            .expect("resource budget mutex poisoned")
149            .usage
150    }
151
152    fn reserve_inner(
153        &self,
154        scope: ResourceScope,
155        bytes: u64,
156        batches: usize,
157    ) -> Result<ResourceReservation> {
158        let mut state = self.state.lock().expect("resource budget mutex poisoned");
159        let observed_bytes = state.usage.reserved_bytes.saturating_add(bytes);
160        let observed_batches = state.usage.reserved_batches.saturating_add(batches);
161
162        if observed_bytes > self.memory_limit_bytes || observed_batches > self.max_in_flight_batches
163        {
164            return Err(DataFrameError::resource_limit_exceeded(
165                self.memory_limit_bytes,
166                observed_bytes,
167                self.max_in_flight_batches,
168                observed_batches,
169                scope,
170            ));
171        }
172
173        state.usage = ResourceUsage {
174            reserved_bytes: observed_bytes,
175            reserved_batches: observed_batches,
176        };
177
178        Ok(ResourceReservation {
179            budget: self.clone(),
180            bytes,
181            batches,
182            released: false,
183        })
184    }
185
186    fn release(&self, bytes: u64, batches: usize) {
187        let mut state = self.state.lock().expect("resource budget mutex poisoned");
188        state.usage.reserved_bytes = state.usage.reserved_bytes.saturating_sub(bytes);
189        state.usage.reserved_batches = state.usage.reserved_batches.saturating_sub(batches);
190    }
191}
192
193/// A single ownership reservation released exactly once on `Drop` or `release`.
194#[derive(Debug)]
195pub struct ResourceReservation {
196    budget: ResourceBudget,
197    bytes: u64,
198    batches: usize,
199    released: bool,
200}
201
202impl ResourceReservation {
203    /// Convert a byte-only reservation into the single batch ownership slot for a published
204    /// result. The caller must first release the source batch slot after the source batch is no
205    /// longer needed, so a `max_in_flight_batches` limit of one remains usable for transforms.
206    pub fn promote_to_batch(&mut self, scope: ResourceScope) -> Result<()> {
207        if self.released || self.batches != 0 {
208            return Ok(());
209        }
210
211        let mut state = self
212            .budget
213            .state
214            .lock()
215            .expect("resource budget mutex poisoned");
216        let observed_batches = state.usage.reserved_batches.saturating_add(1);
217        if observed_batches > self.budget.max_in_flight_batches {
218            return Err(DataFrameError::resource_limit_exceeded(
219                self.budget.memory_limit_bytes,
220                state.usage.reserved_bytes,
221                self.budget.max_in_flight_batches,
222                observed_batches,
223                scope,
224            ));
225        }
226        state.usage.reserved_batches = observed_batches;
227        self.batches = 1;
228        Ok(())
229    }
230
231    /// Release this reservation now. Repeated calls are no-ops.
232    pub fn release(&mut self) {
233        if !self.released {
234            self.budget.release(self.bytes, self.batches);
235            self.released = true;
236        }
237    }
238
239    /// Retain only `bytes` of this reservation, returning any excess immediately.
240    ///
241    /// This is used for bounded metadata probing: reserve before an opaque parser allocates,
242    /// then retain its measured metadata footprint rather than pessimistically pinning the full
243    /// execution budget for the lifetime of a source.
244    pub fn shrink_to(&mut self, bytes: u64) {
245        let bytes = bytes.min(self.bytes);
246        if !self.released && bytes < self.bytes {
247            self.budget.release(self.bytes - bytes, 0);
248            self.bytes = bytes;
249        }
250    }
251
252    /// Return the bytes covered by this reservation.
253    pub const fn bytes(&self) -> u64 {
254        self.bytes
255    }
256}
257
258impl Drop for ResourceReservation {
259    fn drop(&mut self) {
260        self.release();
261    }
262}
263
264/// Stable classes for stream terminal failures.
265#[derive(Debug, Clone, Copy, PartialEq, Eq)]
266pub enum StreamFailureClass {
267    /// Plan or source has no streaming implementation.
268    Unsupported,
269    /// A configured resource bound was exceeded.
270    ResourceLimit,
271    /// The source could not be read.
272    Source,
273    /// Source content could not be decoded.
274    Decode,
275    /// The source or result schema is invalid.
276    Schema,
277    /// The caller cancelled execution.
278    Cancelled,
279    /// The caller closed execution before normal completion.
280    Closed,
281    /// An internal execution failure occurred.
282    Internal,
283}
284
285impl StreamFailureClass {
286    /// Stable diagnostic classification string.
287    pub const fn as_str(self) -> &'static str {
288        match self {
289            Self::Unsupported => "unsupported",
290            Self::ResourceLimit => "resource_limit",
291            Self::Source => "source",
292            Self::Decode => "decode",
293            Self::Schema => "schema",
294            Self::Cancelled => "cancelled",
295            Self::Closed => "closed",
296            Self::Internal => "internal",
297        }
298    }
299}
300
301/// Repeatable, structured terminal failure metadata.
302#[derive(Debug, Clone, PartialEq, Eq)]
303pub struct StreamFailure {
304    /// Stable diagnostic code; callers must not parse error display text.
305    pub code: &'static str,
306    /// Broad failure category.
307    pub classification: StreamFailureClass,
308}
309
310impl StreamFailure {
311    /// Construct structured terminal failure metadata.
312    pub const fn new(code: &'static str, classification: StreamFailureClass) -> Self {
313        Self {
314            code,
315            classification,
316        }
317    }
318}
319
320/// One-way lifecycle state for a `DataFrameStream`.
321#[derive(Debug, Clone, PartialEq, Eq)]
322pub enum StreamTerminal {
323    /// The stream may yield its next batch.
324    Open,
325    /// The source ended normally; later reads return end-of-stream.
326    Exhausted,
327    /// The consumer closed the stream before exhaustion.
328    Closed,
329    /// The consumer cancelled the stream before exhaustion.
330    Cancelled,
331    /// The stream failed; later reads reproduce this structured failure.
332    Failed(StreamFailure),
333}
334
335impl StreamTerminal {
336    /// Return the corresponding repeatable error, if this terminal state cannot yield.
337    pub fn as_error(&self) -> Option<DataFrameError> {
338        match self {
339            Self::Open | Self::Exhausted => None,
340            Self::Closed => Some(DataFrameError::stream_closed()),
341            Self::Cancelled => Some(DataFrameError::stream_cancelled()),
342            Self::Failed(failure) => Some(DataFrameError::stream_failed(
343                failure.code,
344                failure.classification,
345            )),
346        }
347    }
348}
349
350/// Shared terminal state with idempotent transition semantics.
351#[derive(Debug, Clone)]
352pub struct StreamTerminalState {
353    terminal: Arc<Mutex<StreamTerminal>>,
354}
355
356impl Default for StreamTerminalState {
357    fn default() -> Self {
358        Self::new()
359    }
360}
361
362impl StreamTerminalState {
363    /// Start in the open state.
364    pub fn new() -> Self {
365        Self {
366            terminal: Arc::new(Mutex::new(StreamTerminal::Open)),
367        }
368    }
369
370    /// Return the current terminal state.
371    pub fn status(&self) -> StreamTerminal {
372        self.terminal
373            .lock()
374            .expect("stream terminal mutex poisoned")
375            .clone()
376    }
377
378    /// Move from `Open` to a terminal state once; later transitions retain the first result.
379    pub fn finish(&self, requested: StreamTerminal) -> StreamTerminal {
380        debug_assert!(!matches!(requested, StreamTerminal::Open));
381        let mut terminal = self
382            .terminal
383            .lock()
384            .expect("stream terminal mutex poisoned");
385        if matches!(*terminal, StreamTerminal::Open) {
386            *terminal = requested;
387        }
388        terminal.clone()
389    }
390
391    /// Verify that another batch may be requested from this stream.
392    pub fn ensure_open(&self) -> Result<()> {
393        let terminal = self.status();
394        match terminal {
395            StreamTerminal::Open => Ok(()),
396            StreamTerminal::Exhausted => Ok(()),
397            _ => Err(terminal
398                .as_error()
399                .expect("non-open non-exhausted terminals have errors")),
400        }
401    }
402}
403
404#[cfg(test)]
405mod tests {
406    use std::num::NonZeroUsize;
407
408    use super::{
409        ResourceBudget, ResourceScope, StreamFailure, StreamFailureClass, StreamTerminal,
410        StreamTerminalState,
411    };
412    use crate::DataFrameError;
413
414    #[test]
415    fn reservation_accounts_bytes_and_batches_and_releases_once() {
416        let budget = ResourceBudget::new(16, NonZeroUsize::new(1).unwrap());
417        let mut reservation = budget.reserve_batch(ResourceScope::Decode, 12).unwrap();
418        assert_eq!(budget.usage().reserved_bytes, 12);
419        assert_eq!(budget.usage().reserved_batches, 1);
420
421        reservation.release();
422        reservation.release();
423        assert_eq!(budget.usage().reserved_bytes, 0);
424        assert_eq!(budget.usage().reserved_batches, 0);
425    }
426
427    #[test]
428    fn reservation_rejects_before_mutating_usage() {
429        let budget = ResourceBudget::new(16, NonZeroUsize::new(1).unwrap());
430        let err = budget.reserve_batch(ResourceScope::Output, 17).unwrap_err();
431        assert!(matches!(err, DataFrameError::ResourceLimitExceeded { .. }));
432        assert_eq!(budget.usage().reserved_bytes, 0);
433        assert_eq!(budget.usage().reserved_batches, 0);
434    }
435
436    #[test]
437    fn terminal_keeps_the_first_non_open_outcome() {
438        let state = StreamTerminalState::new();
439        let failure = StreamFailure::new("decode_failed", StreamFailureClass::Decode);
440        assert_eq!(
441            state.finish(StreamTerminal::Failed(failure.clone())),
442            StreamTerminal::Failed(failure)
443        );
444        assert_eq!(state.finish(StreamTerminal::Cancelled), state.status());
445        assert!(matches!(
446            state.ensure_open(),
447            Err(DataFrameError::StreamFailed {
448                code: "decode_failed",
449                ..
450            })
451        ));
452    }
453}
alopex_dataframe/physical/budget.rs

alopex_dataframe/physical/
budget.rs
