molframe_core/io/
structure_batch_pool.rs

1//! Reservation-preserving reuse for one structural batch source.
2
3use super::{StructureBatch, StructureBatchBuffer, StructureBatchBuilder, StructureBatchError};
4use crate::execution::BatchBufferPool;
5use crate::{Batch, BatchLease, ChunkDescriptor, ExecutionContext, MemoryReservation};
6use std::sync::Arc;
7
8/// One reusable columnar batch owned by a structural source.
9///
10/// Retained capacity never becomes invisible to the execution budget: the
11/// pool stores the live reservation beside the buffer and releases both when
12/// the source is dropped.
13#[derive(Clone, Debug)]
14pub struct StructureBatchPool {
15    inner: Arc<BatchBufferPool<StructureBatch>>,
16}
17
18impl StructureBatchPool {
19    /// Creates an empty one-slot pool.
20    #[must_use]
21    pub fn new() -> Self {
22        Self {
23            inner: Arc::new(BatchBufferPool::new()),
24        }
25    }
26
27    /// Acquires a fresh or recycled builder after charging its full working
28    /// allowance. A recycled allocation may exceed `reserve_bytes` when its
29    /// existing reservation still covers it and it remains within
30    /// `max_retained_bytes`. `None` means the context cannot currently make
31    /// progress.
32    ///
33    /// # Errors
34    ///
35    /// Rejects unrepresentable batch identity.
36    pub fn acquire(
37        &self,
38        descriptor: ChunkDescriptor,
39        reserve_bytes: usize,
40        max_retained_bytes: usize,
41        context: &ExecutionContext,
42    ) -> Result<Option<(StructureBatchBuilder, MemoryReservation)>, StructureBatchError> {
43        let capacity = descriptor.rows();
44        let recycled = self.inner.take().filter(|(batch, reservation)| {
45            reservation.belongs_to(context) && batch.can_reuse(capacity, max_retained_bytes)
46        });
47        if let Some((batch, mut reservation)) = recycled {
48            if reservation.bytes() < reserve_bytes
49                && reservation
50                    .try_grow(reserve_bytes - reservation.bytes())
51                    .is_err()
52            {
53                self.inner.put(batch, reservation)?;
54                return Ok(None);
55            }
56            let buffer = StructureBatchBuffer::new(batch);
57            let builder = StructureBatchBuilder::reuse(
58                buffer,
59                descriptor.dataset(),
60                descriptor.chunk(),
61                descriptor.logical_start(),
62            )?;
63            return Ok(Some((builder, reservation)));
64        }
65        let Ok(reservation) = context.try_reserve(reserve_bytes) else {
66            return Ok(None);
67        };
68        let builder = StructureBatchBuilder::new(
69            descriptor.dataset(),
70            descriptor.chunk(),
71            descriptor.logical_start(),
72            capacity,
73        )?;
74        Ok(Some((builder, reservation)))
75    }
76
77    /// Publishes a batch whose capacity returns to this pool with its lease.
78    ///
79    /// # Errors
80    ///
81    /// Rejects a batch larger than the producer's reservation.
82    pub fn lease(
83        &self,
84        buffer: StructureBatchBuffer,
85        reservation: MemoryReservation,
86    ) -> Result<BatchLease<StructureBatch>, StructureBatchError> {
87        BatchLease::try_from_reservation_recycling(
88            buffer.into_owned(),
89            reservation,
90            Arc::clone(&self.inner),
91        )
92        .map_err(StructureBatchError::from)
93    }
94
95    /// Keeps a rejected candidate available for a subsequent demand.
96    pub fn recycle(&self, buffer: StructureBatchBuffer, reservation: MemoryReservation) {
97        if buffer.batch().retained_bytes() > reservation.bytes() {
98            drop(buffer);
99            return;
100        }
101        // Capacity was checked above. Rejection still destroys the batch before
102        // releasing its charge, so no unaccounted idle allocation survives.
103        let _ = self.inner.put(buffer.into_owned(), reservation);
104    }
105}
106
107impl Default for StructureBatchPool {
108    fn default() -> Self {
109        Self::new()
110    }
111}
112
113#[cfg(test)]
114#[path = "structure_batch_pool_tests.rs"]
115mod tests;
molframe_core/io/structure_batch_pool.rs

molframe_core/io/
structure_batch_pool.rs
