laminar_core/streaming/
arrow_admission.rs

1//! Retained Arrow storage admission for in-process source queues.
2
3use std::sync::Arc;
4
5use arrow::array::{ArrayRef, RecordBatch};
6use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError};
7
8use super::StreamingError;
9
10/// Default Arrow-byte limit per source for batches its sink has not yet received (64 MiB).
11pub const DEFAULT_SOURCE_MAX_QUEUED_BYTES: usize = 64 * 1024 * 1024;
12
13/// Largest byte limit supported by source semaphore reservations.
14pub const MAX_SOURCE_QUEUED_BYTES: usize = if Semaphore::MAX_PERMITS < u32::MAX as usize {
15    Semaphore::MAX_PERMITS
16} else {
17    u32::MAX as usize
18};
19
20/// Validates an in-process source's Arrow-byte limit.
21///
22/// # Errors
23/// Returns [`StreamingError::InvalidConfig`] for zero or values above
24/// [`MAX_SOURCE_QUEUED_BYTES`].
25pub fn validate_source_max_queued_bytes(bytes: usize) -> Result<(), StreamingError> {
26    if bytes == 0 || bytes > MAX_SOURCE_QUEUED_BYTES {
27        return Err(StreamingError::InvalidConfig(
28            "max_queued_bytes must be nonzero and at most MAX_SOURCE_QUEUED_BYTES".into(),
29        ));
30    }
31    Ok(())
32}
33
34/// Arrow-reported retained array storage plus fixed batch/column charges.
35///
36/// Slices, nested arrays and views include their backing storage. Aliases are charged
37/// independently. Schema metadata, allocator overhead and caller-owned copies are excluded.
38/// Arithmetic overflow saturates so admission fails closed.
39#[must_use]
40pub fn retained_arrow_bytes(batch: &RecordBatch) -> usize {
41    batch
42        .columns()
43        .iter()
44        .fold(std::mem::size_of::<RecordBatch>(), |total, column| {
45            total
46                .saturating_add(column.get_array_memory_size())
47                .saturating_add(std::mem::size_of::<ArrayRef>())
48        })
49}
50
51pub(crate) struct QueuedArrowBatch {
52    batch: RecordBatch,
53    _permit: OwnedSemaphorePermit,
54}
55
56impl QueuedArrowBatch {
57    pub(crate) fn admit(
58        batch: RecordBatch,
59        budget: &Arc<Semaphore>,
60        limit: usize,
61    ) -> Result<Self, StreamingError> {
62        validate_source_max_queued_bytes(limit)?;
63        let bytes = retained_arrow_bytes(&batch);
64        let permits = u32::try_from(bytes)
65            .ok()
66            .filter(|_| bytes <= limit)
67            .ok_or(StreamingError::BatchTooLarge { bytes, limit })?;
68        let permit =
69            Arc::clone(budget)
70                .try_acquire_many_owned(permits)
71                .map_err(|error| match error {
72                    TryAcquireError::NoPermits => StreamingError::ChannelFull,
73                    TryAcquireError::Closed => StreamingError::Disconnected,
74                })?;
75        Ok(Self {
76            batch,
77            _permit: permit,
78        })
79    }
80
81    pub(crate) fn into_batch(self) -> RecordBatch {
82        self.batch
83    }
84}
laminar_core/streaming/arrow_admission.rs

laminar_core/streaming/
arrow_admission.rs
