spatialrust_records/
bounded.rs

1//! Backward-compatible bounded record streams with leased chunk ownership.
2
3use std::sync::{
4    atomic::{AtomicU64, Ordering},
5    mpsc::{sync_channel, Receiver},
6    Arc, Mutex,
7};
8use std::thread::{self, JoinHandle};
9
10use spatialrust_core::{FieldSemantic, PointBuffer, PointBufferSet, PointCloud, SpatialMetadata};
11
12use crate::{
13    CancellationToken, MemoryReservation, MemoryTracker, RecordsError, RecordsResult,
14    SchemaDescriptor, SpatialRecord, SpatialRecordSink, SpatialRecordSource, StreamOptions,
15};
16
17/// Stable location of one chunk in its source stream.
18#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
19pub struct ChunkIdentity {
20    /// Zero-based chunk sequence.
21    pub sequence: u64,
22    /// Global point offset of the first point in the chunk.
23    pub point_offset: u64,
24}
25
26/// Finite axis-aligned bounds carried by a record chunk.
27#[derive(Clone, Copy, Debug, PartialEq)]
28pub struct RecordBounds3 {
29    /// Inclusive minimum XYZ coordinates.
30    pub min: [f64; 3],
31    /// Inclusive maximum XYZ coordinates.
32    pub max: [f64; 3],
33}
34
35impl RecordBounds3 {
36    /// Creates validated finite, ordered bounds.
37    pub fn new(min: [f64; 3], max: [f64; 3]) -> RecordsResult<Self> {
38        if min.into_iter().chain(max).any(|component| !component.is_finite())
39            || (0..3).any(|axis| min[axis] > max[axis])
40        {
41            return Err(RecordsError::InvalidChunk(
42                "chunk bounds must be finite and ordered".into(),
43            ));
44        }
45        Ok(Self { min, max })
46    }
47}
48
49type BufferPool = Arc<Mutex<Vec<PointBufferSet>>>;
50
51/// One owned record plus the memory lease that keeps it inside the hard limit.
52///
53/// The record is borrowed from this wrapper and cannot outlive its reservation.
54/// Source-specific buffers may be returned to a recycling pool when the chunk
55/// is dropped.
56#[derive(Debug)]
57pub struct SpatialRecordChunk {
58    identity: ChunkIdentity,
59    bounds: Option<RecordBounds3>,
60    tracked_bytes: u64,
61    record: Option<SpatialRecord>,
62    reservation: MemoryReservation,
63    recycle_pool: Option<BufferPool>,
64}
65
66impl SpatialRecordChunk {
67    fn new(
68        identity: ChunkIdentity,
69        bounds: Option<RecordBounds3>,
70        tracked_bytes: u64,
71        record: SpatialRecord,
72        reservation: MemoryReservation,
73    ) -> Self {
74        Self {
75            identity,
76            bounds,
77            tracked_bytes,
78            record: Some(record),
79            reservation,
80            recycle_pool: None,
81        }
82    }
83
84    fn with_recycle_pool(mut self, pool: BufferPool) -> Self {
85        self.recycle_pool = Some(pool);
86        self
87    }
88
89    /// Creates a chunk from a reservation acquired before source allocation.
90    ///
91    /// The reservation is shrunk to the record's allocated column capacity.
92    /// Construction fails if the source allocated more than it reserved.
93    pub fn try_from_reserved(
94        identity: ChunkIdentity,
95        record: SpatialRecord,
96        mut reservation: MemoryReservation,
97    ) -> RecordsResult<Self> {
98        let tracked_bytes = record_storage_bytes(&record)?;
99        if tracked_bytes > reservation.bytes() {
100            return Err(RecordsError::InvalidChunk(format!(
101                "record column capacity {tracked_bytes} exceeds reservation {}",
102                reservation.bytes()
103            )));
104        }
105        reservation.shrink_to(tracked_bytes)?;
106        let bounds = record_bounds(&record);
107        Ok(Self::new(identity, bounds, tracked_bytes, record, reservation))
108    }
109
110    /// Returns the deterministic source identity.
111    #[must_use]
112    pub const fn identity(&self) -> ChunkIdentity {
113        self.identity
114    }
115
116    /// Returns optional finite bounds.
117    #[must_use]
118    pub const fn bounds(&self) -> Option<RecordBounds3> {
119        self.bounds
120    }
121
122    /// Returns explicitly tracked column-capacity bytes for this chunk.
123    #[must_use]
124    pub const fn tracked_bytes(&self) -> u64 {
125        self.tracked_bytes
126    }
127
128    /// Borrows the owned record while its memory lease remains alive.
129    #[must_use]
130    pub fn record(&self) -> &SpatialRecord {
131        self.record.as_ref().expect("record is present until chunk drop")
132    }
133
134    /// Returns the exact held reservation.
135    #[must_use]
136    pub const fn reservation(&self) -> &MemoryReservation {
137        &self.reservation
138    }
139}
140
141impl Drop for SpatialRecordChunk {
142    fn drop(&mut self) {
143        let Some(pool) = &self.recycle_pool else {
144            return;
145        };
146        let Some(record) = self.record.take() else {
147            return;
148        };
149        let (_, mut buffers, _) = record.into_cloud().into_parts();
150        clear_buffers(&mut buffers);
151        lock_pool(pool).push(buffers);
152    }
153}
154
155/// Pull-based bounded source. Existing [`SpatialRecordSource`] remains unchanged.
156pub trait BoundedSpatialRecordSource {
157    /// Schema shared by every emitted record.
158    fn schema(&self) -> &SchemaDescriptor;
159    /// Validated execution options.
160    fn options(&self) -> &StreamOptions;
161    /// Exact tracker shared by emitted chunk leases.
162    fn memory_tracker(&self) -> &MemoryTracker;
163    /// Cooperative cancellation state checked at chunk boundaries.
164    fn cancellation_token(&self) -> CancellationToken;
165    /// Conservative maximum resident bytes for one emitted chunk.
166    fn max_chunk_bytes(&self) -> u64;
167    /// Returns the next leased chunk, or `None` after an explicit end.
168    fn next_chunk(&mut self) -> Option<RecordsResult<SpatialRecordChunk>>;
169}
170
171impl<S: BoundedSpatialRecordSource + ?Sized> BoundedSpatialRecordSource for Box<S> {
172    fn schema(&self) -> &SchemaDescriptor {
173        (**self).schema()
174    }
175
176    fn options(&self) -> &StreamOptions {
177        (**self).options()
178    }
179
180    fn memory_tracker(&self) -> &MemoryTracker {
181        (**self).memory_tracker()
182    }
183
184    fn cancellation_token(&self) -> CancellationToken {
185        (**self).cancellation_token()
186    }
187
188    fn max_chunk_bytes(&self) -> u64 {
189        (**self).max_chunk_bytes()
190    }
191
192    fn next_chunk(&mut self) -> Option<RecordsResult<SpatialRecordChunk>> {
193        (**self).next_chunk()
194    }
195}
196
197/// Push-based bounded sink that cannot retain a record beyond the chunk lease.
198pub trait BoundedSpatialRecordSink {
199    /// Writes one chunk synchronously.
200    fn write_chunk(&mut self, chunk: &SpatialRecordChunk) -> RecordsResult<()>;
201
202    /// Finalizes the sink.
203    fn finish(&mut self) -> RecordsResult<()> {
204        Ok(())
205    }
206}
207
208impl<S: BoundedSpatialRecordSink + ?Sized> BoundedSpatialRecordSink for Box<S> {
209    fn write_chunk(&mut self, chunk: &SpatialRecordChunk) -> RecordsResult<()> {
210        (**self).write_chunk(chunk)
211    }
212
213    fn finish(&mut self) -> RecordsResult<()> {
214        (**self).finish()
215    }
216}
217
218/// Adapts an existing record source without changing its public trait.
219pub struct LegacyBoundedSource<S> {
220    source: S,
221    schema: SchemaDescriptor,
222    options: StreamOptions,
223    tracker: MemoryTracker,
224    cancellation: CancellationToken,
225    next_sequence: u64,
226    next_point_offset: u64,
227    max_chunk_bytes: u64,
228}
229
230impl<S: SpatialRecordSource> LegacyBoundedSource<S> {
231    /// Wraps `source`, reserving its declared maximum before each pull.
232    pub fn try_new(
233        source: S,
234        options: StreamOptions,
235        cancellation: CancellationToken,
236    ) -> RecordsResult<Self> {
237        let schema = source.schema().clone();
238        let max_chunk_bytes = max_storage_bytes(&schema, options.chunk_points())?;
239        if max_chunk_bytes > options.memory_budget().limit_bytes() {
240            return Err(RecordsError::MemoryBudgetExceeded {
241                requested: max_chunk_bytes,
242                current: 0,
243                limit: options.memory_budget().limit_bytes(),
244            });
245        }
246        let tracker = MemoryTracker::new(options.memory_budget());
247        Ok(Self {
248            source,
249            schema,
250            options,
251            tracker,
252            cancellation,
253            next_sequence: 0,
254            next_point_offset: 0,
255            max_chunk_bytes,
256        })
257    }
258
259    /// Consumes the adapter and returns the legacy source.
260    #[must_use]
261    pub fn into_inner(self) -> S {
262        self.source
263    }
264}
265
266impl<S: SpatialRecordSource> BoundedSpatialRecordSource for LegacyBoundedSource<S> {
267    fn schema(&self) -> &SchemaDescriptor {
268        &self.schema
269    }
270
271    fn options(&self) -> &StreamOptions {
272        &self.options
273    }
274
275    fn memory_tracker(&self) -> &MemoryTracker {
276        &self.tracker
277    }
278
279    fn cancellation_token(&self) -> CancellationToken {
280        self.cancellation.clone()
281    }
282
283    fn max_chunk_bytes(&self) -> u64 {
284        self.max_chunk_bytes
285    }
286
287    fn next_chunk(&mut self) -> Option<RecordsResult<SpatialRecordChunk>> {
288        if let Err(error) = self.cancellation.check() {
289            return Some(Err(error));
290        }
291        let mut reservation = match self.tracker.try_reserve(self.max_chunk_bytes) {
292            Ok(reservation) => reservation,
293            Err(error) => return Some(Err(error)),
294        };
295        let record = match self.source.next_record()? {
296            Ok(record) => record,
297            Err(error) => return Some(Err(error)),
298        };
299        if record.cloud().is_empty() {
300            return Some(Err(RecordsError::InvalidChunk(
301                "sources must not emit empty chunks".into(),
302            )));
303        }
304        if record.cloud().len() > self.options.chunk_points() {
305            return Some(Err(RecordsError::InvalidChunk(format!(
306                "source emitted {} points above declared chunk limit {}",
307                record.cloud().len(),
308                self.options.chunk_points()
309            ))));
310        }
311        let tracked_bytes = match record_storage_bytes(&record) {
312            Ok(bytes) => bytes,
313            Err(error) => return Some(Err(error)),
314        };
315        if tracked_bytes > self.max_chunk_bytes {
316            return Some(Err(RecordsError::InvalidChunk(format!(
317                "chunk storage {tracked_bytes} exceeds reserved maximum {}",
318                self.max_chunk_bytes
319            ))));
320        }
321        if let Err(error) = reservation.shrink_to(tracked_bytes) {
322            return Some(Err(error));
323        }
324        let point_count = match u64::try_from(record.cloud().len()) {
325            Ok(value) => value,
326            Err(_) => {
327                return Some(Err(RecordsError::InvalidChunk(
328                    "chunk point count does not fit u64".into(),
329                )))
330            }
331        };
332        let identity =
333            ChunkIdentity { sequence: self.next_sequence, point_offset: self.next_point_offset };
334        self.next_sequence = match self.next_sequence.checked_add(1) {
335            Some(value) => value,
336            None => return Some(Err(RecordsError::ReceiptOverflow("chunk sequence".into()))),
337        };
338        self.next_point_offset = match self.next_point_offset.checked_add(point_count) {
339            Some(value) => value,
340            None => return Some(Err(RecordsError::ReceiptOverflow("point offset".into()))),
341        };
342        let bounds = record_bounds(&record);
343        Some(Ok(SpatialRecordChunk::new(identity, bounds, tracked_bytes, record, reservation)))
344    }
345}
346
347/// Adapts an existing synchronous sink to leased record chunks.
348pub struct LegacyBoundedSink<S> {
349    sink: S,
350}
351
352impl<S> LegacyBoundedSink<S> {
353    /// Creates a synchronous bounded sink adapter.
354    #[must_use]
355    pub const fn new(sink: S) -> Self {
356        Self { sink }
357    }
358
359    /// Consumes the adapter and returns its sink.
360    #[must_use]
361    pub fn into_inner(self) -> S {
362        self.sink
363    }
364}
365
366impl<S: SpatialRecordSink> BoundedSpatialRecordSink for LegacyBoundedSink<S> {
367    fn write_chunk(&mut self, chunk: &SpatialRecordChunk) -> RecordsResult<()> {
368        self.sink.write_record(chunk.record())
369    }
370
371    fn finish(&mut self) -> RecordsResult<()> {
372        self.sink.finish()
373    }
374}
375
376enum PrefetchMessage {
377    Item(Box<RecordsResult<SpatialRecordChunk>>),
378    End,
379}
380
381/// Single-worker deterministic prefetch with count and memory backpressure.
382pub struct PrefetchRecordSource {
383    schema: SchemaDescriptor,
384    options: StreamOptions,
385    tracker: MemoryTracker,
386    cancellation: CancellationToken,
387    max_chunk_bytes: u64,
388    receiver: Option<Receiver<PrefetchMessage>>,
389    worker: Option<JoinHandle<()>>,
390    finished: bool,
391    next_expected_sequence: u64,
392    next_expected_point_offset: u64,
393}
394
395impl PrefetchRecordSource {
396    /// Starts prefetch using `source.options().prefetch_chunks()` as queue capacity.
397    pub fn try_new<S>(mut source: S) -> RecordsResult<Self>
398    where
399        S: BoundedSpatialRecordSource + Send + 'static,
400    {
401        let options = source.options().clone();
402        let capacity = options.prefetch_chunks();
403        if capacity == 0 {
404            return Err(RecordsError::InvalidConfiguration(
405                "prefetch source requires prefetch_chunks > 0".into(),
406            ));
407        }
408        let max_chunk_bytes = source.max_chunk_bytes();
409        // Besides the bounded queue, one lease may be held by the consumer and
410        // one by the producer while `send` waits for queue capacity.
411        let concurrent_chunks =
412            u64::try_from(capacity).ok().and_then(|value| value.checked_add(2)).ok_or_else(
413                || RecordsError::InvalidConfiguration("prefetch capacity overflow".into()),
414            )?;
415        let required = max_chunk_bytes.checked_mul(concurrent_chunks).ok_or_else(|| {
416            RecordsError::InvalidConfiguration("prefetch memory requirement overflow".into())
417        })?;
418        if required > options.memory_budget().limit_bytes() {
419            return Err(RecordsError::MemoryBudgetExceeded {
420                requested: required,
421                current: 0,
422                limit: options.memory_budget().limit_bytes(),
423            });
424        }
425
426        let schema = source.schema().clone();
427        let tracker = source.memory_tracker().clone();
428        let cancellation = source.cancellation_token();
429        let worker_cancellation = cancellation.clone();
430        let (sender, receiver) = sync_channel(capacity);
431        let worker = thread::spawn(move || loop {
432            if worker_cancellation.is_cancelled() {
433                break;
434            }
435            match source.next_chunk() {
436                Some(result) => {
437                    let stop = result.is_err();
438                    if sender.send(PrefetchMessage::Item(Box::new(result))).is_err() || stop {
439                        break;
440                    }
441                }
442                None => {
443                    let _ = sender.send(PrefetchMessage::End);
444                    break;
445                }
446            }
447        });
448        Ok(Self {
449            schema,
450            options,
451            tracker,
452            cancellation,
453            max_chunk_bytes,
454            receiver: Some(receiver),
455            worker: Some(worker),
456            finished: false,
457            next_expected_sequence: 0,
458            next_expected_point_offset: 0,
459        })
460    }
461}
462
463impl BoundedSpatialRecordSource for PrefetchRecordSource {
464    fn schema(&self) -> &SchemaDescriptor {
465        &self.schema
466    }
467
468    fn options(&self) -> &StreamOptions {
469        &self.options
470    }
471
472    fn memory_tracker(&self) -> &MemoryTracker {
473        &self.tracker
474    }
475
476    fn cancellation_token(&self) -> CancellationToken {
477        self.cancellation.clone()
478    }
479
480    fn max_chunk_bytes(&self) -> u64 {
481        self.max_chunk_bytes
482    }
483
484    fn next_chunk(&mut self) -> Option<RecordsResult<SpatialRecordChunk>> {
485        if self.finished {
486            return None;
487        }
488        let Some(receiver) = &self.receiver else {
489            self.finished = true;
490            return Some(Err(RecordsError::StreamClosed));
491        };
492        match receiver.recv() {
493            Ok(PrefetchMessage::Item(result)) => match *result {
494                Ok(chunk) => {
495                    let identity = chunk.identity();
496                    if identity.sequence != self.next_expected_sequence
497                        || identity.point_offset != self.next_expected_point_offset
498                    {
499                        self.cancellation.cancel();
500                        return Some(Err(RecordsError::InvalidChunk(format!(
501                            "non-contiguous chunk identity {:?}, expected sequence {} offset {}",
502                            identity, self.next_expected_sequence, self.next_expected_point_offset
503                        ))));
504                    }
505                    let point_count = match u64::try_from(chunk.record().cloud().len()) {
506                        Ok(value) => value,
507                        Err(_) => {
508                            self.cancellation.cancel();
509                            return Some(Err(RecordsError::InvalidChunk(
510                                "chunk point count does not fit u64".into(),
511                            )));
512                        }
513                    };
514                    self.next_expected_sequence = match self.next_expected_sequence.checked_add(1) {
515                        Some(value) => value,
516                        None => {
517                            self.cancellation.cancel();
518                            return Some(Err(RecordsError::ReceiptOverflow(
519                                "prefetch chunk sequence".into(),
520                            )));
521                        }
522                    };
523                    self.next_expected_point_offset =
524                        match self.next_expected_point_offset.checked_add(point_count) {
525                            Some(value) => value,
526                            None => {
527                                self.cancellation.cancel();
528                                return Some(Err(RecordsError::ReceiptOverflow(
529                                    "prefetch point offset".into(),
530                                )));
531                            }
532                        };
533                    Some(Ok(chunk))
534                }
535                Err(error) => Some(Err(error)),
536            },
537            Ok(PrefetchMessage::End) => {
538                self.finished = true;
539                None
540            }
541            Err(_) => {
542                self.finished = true;
543                Some(Err(RecordsError::StreamClosed))
544            }
545        }
546    }
547}
548
549impl Drop for PrefetchRecordSource {
550    fn drop(&mut self) {
551        self.cancellation.cancel();
552        self.receiver.take();
553        if let Some(worker) = self.worker.take() {
554            let _ = worker.join();
555        }
556    }
557}
558
559/// In-memory source used to verify real buffer recycling and lease accounting.
560///
561/// The immutable backing cloud represents external source storage. Emitted
562/// chunk buffers are separately reserved and returned to the pool on drop.
563pub struct RecyclingMemoryChunkSource {
564    schema: SchemaDescriptor,
565    cloud: PointCloud,
566    options: StreamOptions,
567    tracker: MemoryTracker,
568    cancellation: CancellationToken,
569    next_sequence: u64,
570    next_point_offset: u64,
571    offset: usize,
572    max_chunk_bytes: u64,
573    pool: BufferPool,
574    buffer_set_allocations: Arc<AtomicU64>,
575}
576
577impl RecyclingMemoryChunkSource {
578    /// Creates a recycling source over an immutable backing cloud.
579    pub fn try_new(
580        schema: SchemaDescriptor,
581        cloud: PointCloud,
582        options: StreamOptions,
583        cancellation: CancellationToken,
584    ) -> RecordsResult<Self> {
585        if cloud.schema() != schema.point_schema() {
586            return Err(RecordsError::SchemaMismatch(
587                "recycling source cloud schema must match descriptor".into(),
588            ));
589        }
590        cloud.validate()?;
591        let max_chunk_bytes = max_storage_bytes(&schema, options.chunk_points())?;
592        if max_chunk_bytes > options.memory_budget().limit_bytes() {
593            return Err(RecordsError::MemoryBudgetExceeded {
594                requested: max_chunk_bytes,
595                current: 0,
596                limit: options.memory_budget().limit_bytes(),
597            });
598        }
599        let tracker = MemoryTracker::new(options.memory_budget());
600        Ok(Self {
601            schema,
602            cloud,
603            options,
604            tracker,
605            cancellation,
606            next_sequence: 0,
607            next_point_offset: 0,
608            offset: 0,
609            max_chunk_bytes,
610            pool: Arc::new(Mutex::new(Vec::new())),
611            buffer_set_allocations: Arc::new(AtomicU64::new(0)),
612        })
613    }
614
615    /// Returns buffer-set allocations performed because the pool was empty.
616    #[must_use]
617    pub fn buffer_set_allocations(&self) -> u64 {
618        self.buffer_set_allocations.load(Ordering::Acquire)
619    }
620
621    /// Returns recycled buffer sets currently available to the source.
622    #[must_use]
623    pub fn pooled_buffer_sets(&self) -> usize {
624        lock_pool(&self.pool).len()
625    }
626
627    fn take_buffers(&self) -> PointBufferSet {
628        if let Some(buffers) = lock_pool(&self.pool).pop() {
629            return buffers;
630        }
631        self.buffer_set_allocations.fetch_add(1, Ordering::AcqRel);
632        let mut buffers = PointBufferSet::new();
633        for field in self.schema.point_schema().fields() {
634            buffers.insert(
635                field.name.clone(),
636                PointBuffer::with_capacity(field.dtype, self.options.chunk_points()),
637            );
638        }
639        buffers
640    }
641}
642
643impl BoundedSpatialRecordSource for RecyclingMemoryChunkSource {
644    fn schema(&self) -> &SchemaDescriptor {
645        &self.schema
646    }
647
648    fn options(&self) -> &StreamOptions {
649        &self.options
650    }
651
652    fn memory_tracker(&self) -> &MemoryTracker {
653        &self.tracker
654    }
655
656    fn cancellation_token(&self) -> CancellationToken {
657        self.cancellation.clone()
658    }
659
660    fn max_chunk_bytes(&self) -> u64 {
661        self.max_chunk_bytes
662    }
663
664    fn next_chunk(&mut self) -> Option<RecordsResult<SpatialRecordChunk>> {
665        if let Err(error) = self.cancellation.check() {
666            return Some(Err(error));
667        }
668        if self.offset >= self.cloud.len() {
669            return None;
670        }
671        let end = (self.offset + self.options.chunk_points()).min(self.cloud.len());
672        let range = self.offset..end;
673        let point_count = end - self.offset;
674        // Recycled vectors retain full chunk capacity even for the final short
675        // chunk, so the lease keeps the full reusable allocation reserved.
676        let tracked_bytes = self.max_chunk_bytes;
677        let reservation = match self.tracker.try_reserve(tracked_bytes) {
678            Ok(reservation) => reservation,
679            Err(error) => return Some(Err(error)),
680        };
681        let mut buffers = self.take_buffers();
682        for field in self.schema.point_schema().fields() {
683            let source = match self.cloud.field(&field.name) {
684                Ok(buffer) => buffer,
685                Err(error) => return Some(Err(error.into())),
686            };
687            let Some(destination) = buffers.get_mut(&field.name) else {
688                return Some(Err(RecordsError::SchemaMismatch(format!(
689                    "recycled buffer set is missing `{}`",
690                    field.name
691                ))));
692            };
693            if let Err(error) = copy_buffer_range(destination, source, range.clone()) {
694                return Some(Err(error));
695            }
696        }
697        let metadata: SpatialMetadata = self.cloud.metadata().clone();
698        let cloud =
699            match PointCloud::try_from_parts(self.schema.point_schema().clone(), buffers, metadata)
700            {
701                Ok(cloud) => cloud,
702                Err(error) => return Some(Err(error.into())),
703            };
704        let record = match SpatialRecord::try_new(self.schema.clone(), cloud) {
705            Ok(record) => record,
706            Err(error) => return Some(Err(error)),
707        };
708        let identity =
709            ChunkIdentity { sequence: self.next_sequence, point_offset: self.next_point_offset };
710        self.next_sequence = match self.next_sequence.checked_add(1) {
711            Some(value) => value,
712            None => return Some(Err(RecordsError::ReceiptOverflow("chunk sequence".into()))),
713        };
714        let point_count = match u64::try_from(point_count) {
715            Ok(value) => value,
716            Err(_) => {
717                return Some(Err(RecordsError::InvalidChunk(
718                    "chunk point count does not fit u64".into(),
719                )))
720            }
721        };
722        self.next_point_offset = match self.next_point_offset.checked_add(point_count) {
723            Some(value) => value,
724            None => return Some(Err(RecordsError::ReceiptOverflow("point offset".into()))),
725        };
726        self.offset = end;
727        let bounds = record_bounds(&record);
728        Some(Ok(SpatialRecordChunk::new(identity, bounds, tracked_bytes, record, reservation)
729            .with_recycle_pool(self.pool.clone())))
730    }
731}
732
733/// Returns the allocated scalar capacity bytes owned by a record.
734pub fn record_storage_bytes(record: &SpatialRecord) -> RecordsResult<u64> {
735    let mut total = 0_u64;
736    for field in record.schema().point_schema().fields() {
737        let buffer = record.cloud().field(&field.name)?;
738        let bytes = buffer_capacity(buffer)
739            .and_then(|len| len.checked_mul(buffer_scalar_bytes(buffer)))
740            .ok_or_else(|| RecordsError::InvalidChunk("record storage byte overflow".into()))?;
741        total = total
742            .checked_add(bytes)
743            .ok_or_else(|| RecordsError::InvalidChunk("record storage byte overflow".into()))?;
744    }
745    Ok(total)
746}
747
748fn max_storage_bytes(schema: &SchemaDescriptor, point_count: usize) -> RecordsResult<u64> {
749    let point_count = u64::try_from(point_count)
750        .map_err(|_| RecordsError::InvalidConfiguration("chunk size does not fit u64".into()))?;
751    let mut bytes_per_point = 0_u64;
752    for field in schema.point_schema().fields() {
753        let scalar_bytes = match field.dtype {
754            spatialrust_core::DType::F32 | spatialrust_core::DType::F16 => 4,
755            spatialrust_core::DType::F64 => 8,
756            spatialrust_core::DType::U8 => 1,
757            spatialrust_core::DType::U16 => 2,
758            spatialrust_core::DType::U32 | spatialrust_core::DType::I32 => 4,
759        };
760        bytes_per_point = bytes_per_point.checked_add(scalar_bytes).ok_or_else(|| {
761            RecordsError::InvalidConfiguration("schema storage byte overflow".into())
762        })?;
763    }
764    bytes_per_point
765        .checked_mul(point_count)
766        .ok_or_else(|| RecordsError::InvalidConfiguration("chunk storage byte overflow".into()))
767}
768
769fn buffer_scalar_bytes(buffer: &PointBuffer) -> u64 {
770    match buffer {
771        PointBuffer::F32(_) | PointBuffer::U32(_) | PointBuffer::I32(_) => 4,
772        PointBuffer::F64(_) => 8,
773        PointBuffer::U8(_) => 1,
774        PointBuffer::U16(_) => 2,
775    }
776}
777
778fn buffer_capacity(buffer: &PointBuffer) -> Option<u64> {
779    let capacity = match buffer {
780        PointBuffer::F32(values) => values.capacity(),
781        PointBuffer::F64(values) => values.capacity(),
782        PointBuffer::U8(values) => values.capacity(),
783        PointBuffer::U16(values) => values.capacity(),
784        PointBuffer::U32(values) => values.capacity(),
785        PointBuffer::I32(values) => values.capacity(),
786    };
787    u64::try_from(capacity).ok()
788}
789
790fn record_bounds(record: &SpatialRecord) -> Option<RecordBounds3> {
791    let schema = record.schema().point_schema();
792    let x = record
793        .cloud()
794        .field(&schema.find_semantic(FieldSemantic::PositionX)?.name)
795        .ok()?
796        .as_f32()
797        .ok()?;
798    let y = record
799        .cloud()
800        .field(&schema.find_semantic(FieldSemantic::PositionY)?.name)
801        .ok()?
802        .as_f32()
803        .ok()?;
804    let z = record
805        .cloud()
806        .field(&schema.find_semantic(FieldSemantic::PositionZ)?.name)
807        .ok()?
808        .as_f32()
809        .ok()?;
810    let mut min = [f64::INFINITY; 3];
811    let mut max = [f64::NEG_INFINITY; 3];
812    let mut found = false;
813    for ((&x, &y), &z) in x.iter().zip(y).zip(z) {
814        let point = [f64::from(x), f64::from(y), f64::from(z)];
815        if point.iter().any(|value| !value.is_finite()) {
816            continue;
817        }
818        found = true;
819        for axis in 0..3 {
820            min[axis] = min[axis].min(point[axis]);
821            max[axis] = max[axis].max(point[axis]);
822        }
823    }
824    found.then_some(RecordBounds3 { min, max })
825}
826
827fn clear_buffers(buffers: &mut PointBufferSet) {
828    for (_, buffer) in buffers.iter_mut() {
829        clear_buffer(buffer);
830    }
831}
832
833fn clear_buffer(buffer: &mut PointBuffer) {
834    match buffer {
835        PointBuffer::F32(values) => values.clear(),
836        PointBuffer::F64(values) => values.clear(),
837        PointBuffer::U8(values) => values.clear(),
838        PointBuffer::U16(values) => values.clear(),
839        PointBuffer::U32(values) => values.clear(),
840        PointBuffer::I32(values) => values.clear(),
841    }
842}
843
844fn copy_buffer_range(
845    destination: &mut PointBuffer,
846    source: &PointBuffer,
847    range: std::ops::Range<usize>,
848) -> RecordsResult<()> {
849    clear_buffer(destination);
850    match (destination, source) {
851        (PointBuffer::F32(dst), PointBuffer::F32(src)) => {
852            dst.extend_from_slice(&src[range]);
853        }
854        (PointBuffer::F64(dst), PointBuffer::F64(src)) => {
855            dst.extend_from_slice(&src[range]);
856        }
857        (PointBuffer::U8(dst), PointBuffer::U8(src)) => {
858            dst.extend_from_slice(&src[range]);
859        }
860        (PointBuffer::U16(dst), PointBuffer::U16(src)) => {
861            dst.extend_from_slice(&src[range]);
862        }
863        (PointBuffer::U32(dst), PointBuffer::U32(src)) => {
864            dst.extend_from_slice(&src[range]);
865        }
866        (PointBuffer::I32(dst), PointBuffer::I32(src)) => {
867            dst.extend_from_slice(&src[range]);
868        }
869        (dst, src) => {
870            return Err(RecordsError::SchemaMismatch(format!(
871                "cannot recycle {:?} storage for {:?}",
872                dst.dtype(),
873                src.dtype()
874            )));
875        }
876    }
877    Ok(())
878}
879
880fn lock_pool(pool: &BufferPool) -> std::sync::MutexGuard<'_, Vec<PointBufferSet>> {
881    pool.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
882}
883
884#[cfg(test)]
885mod tests {
886    use super::{
887        record_storage_bytes, BoundedSpatialRecordSink, BoundedSpatialRecordSource,
888        LegacyBoundedSink, LegacyBoundedSource, PrefetchRecordSource, RecyclingMemoryChunkSource,
889    };
890    use crate::{
891        CancellationToken, MemoryBudget, MemoryChunkSink, MemoryChunkSource, SchemaDescriptor,
892        SchemaVersion, StreamOptions,
893    };
894    use spatialrust_core::{
895        PointBuffer, PointBufferSet, PointCloud, PointCloudBuilder, SpatialMetadata,
896        StandardSchemas,
897    };
898
899    fn cloud(points: usize) -> spatialrust_core::PointCloud {
900        let mut builder = PointCloudBuilder::xyz();
901        for index in 0..points {
902            builder.push_point([index as f32, 1.0, -1.0]).unwrap();
903        }
904        builder.build().unwrap()
905    }
906
907    fn schema() -> SchemaDescriptor {
908        SchemaDescriptor::try_new("point", SchemaVersion::new(1, 0), StandardSchemas::point_xyz())
909            .unwrap()
910    }
911
912    #[test]
913    fn storage_accounting_uses_vector_capacity_not_only_length() {
914        let mut buffers = PointBufferSet::new();
915        for name in ["x", "y", "z"] {
916            let mut values = Vec::with_capacity(10);
917            values.push(1.0);
918            buffers.insert(name, PointBuffer::from_f32(values));
919        }
920        let cloud = PointCloud::try_from_parts(
921            StandardSchemas::point_xyz(),
922            buffers,
923            SpatialMetadata::default(),
924        )
925        .unwrap();
926        let record =
927            crate::SpatialRecord::try_from_cloud("point", SchemaVersion::new(1, 0), cloud).unwrap();
928        assert_eq!(record_storage_bytes(&record).unwrap(), 120);
929    }
930
931    #[test]
932    fn legacy_source_gets_identity_bounds_and_drop_scoped_memory() {
933        let legacy = MemoryChunkSource::try_new(schema(), cloud(5), 2).unwrap();
934        let options = StreamOptions::new(2, MemoryBudget::new(24).unwrap()).unwrap();
935        let mut source =
936            LegacyBoundedSource::try_new(legacy, options, CancellationToken::default()).unwrap();
937
938        let first = source.next_chunk().unwrap().unwrap();
939        assert_eq!(first.identity().sequence, 0);
940        assert_eq!(first.identity().point_offset, 0);
941        assert_eq!(first.bounds().unwrap().min, [0.0, 1.0, -1.0]);
942        assert_eq!(source.memory_tracker().snapshot().current_bytes, 24);
943        drop(first);
944        assert_eq!(source.memory_tracker().snapshot().current_bytes, 0);
945
946        let second = source.next_chunk().unwrap().unwrap();
947        assert_eq!(second.identity().point_offset, 2);
948    }
949
950    #[test]
951    fn cancellation_stops_before_pulling_another_chunk() {
952        let token = CancellationToken::default();
953        let legacy = MemoryChunkSource::try_new(schema(), cloud(2), 2).unwrap();
954        let options = StreamOptions::new(2, MemoryBudget::new(24).unwrap()).unwrap();
955        let mut source = LegacyBoundedSource::try_new(legacy, options, token.clone()).unwrap();
956        token.cancel();
957        assert!(source.next_chunk().unwrap().is_err());
958    }
959
960    #[test]
961    fn recycling_source_reuses_one_buffer_set_in_steady_state() {
962        let options = StreamOptions::new(2, MemoryBudget::new(24).unwrap()).unwrap();
963        let mut source = RecyclingMemoryChunkSource::try_new(
964            schema(),
965            cloud(5),
966            options,
967            CancellationToken::default(),
968        )
969        .unwrap();
970        for expected_sequence in 0..3 {
971            let chunk = source.next_chunk().unwrap().unwrap();
972            assert_eq!(chunk.identity().sequence, expected_sequence);
973            drop(chunk);
974            assert_eq!(source.pooled_buffer_sets(), 1);
975        }
976        assert_eq!(source.buffer_set_allocations(), 1);
977        assert!(source.next_chunk().is_none());
978    }
979
980    #[test]
981    fn prefetch_preserves_order_and_stays_within_budget() {
982        let legacy = MemoryChunkSource::try_new(schema(), cloud(6), 2).unwrap();
983        let options =
984            StreamOptions::new(2, MemoryBudget::new(96).unwrap()).unwrap().with_prefetch_chunks(2);
985        let bounded =
986            LegacyBoundedSource::try_new(legacy, options, CancellationToken::default()).unwrap();
987        let mut source = PrefetchRecordSource::try_new(bounded).unwrap();
988        for expected in 0..3 {
989            let chunk = source.next_chunk().unwrap().unwrap();
990            assert_eq!(chunk.identity().sequence, expected);
991        }
992        assert!(source.next_chunk().is_none());
993        assert!(source.memory_tracker().snapshot().peak_bytes <= 96);
994    }
995
996    #[test]
997    fn prefetch_rejects_capacity_that_cannot_fit_budget() {
998        let legacy = MemoryChunkSource::try_new(schema(), cloud(6), 2).unwrap();
999        let options =
1000            StreamOptions::new(2, MemoryBudget::new(72).unwrap()).unwrap().with_prefetch_chunks(2);
1001        let bounded =
1002            LegacyBoundedSource::try_new(legacy, options, CancellationToken::default()).unwrap();
1003        assert!(PrefetchRecordSource::try_new(bounded).is_err());
1004    }
1005
1006    #[test]
1007    fn bounded_sink_bridges_existing_sink_synchronously() {
1008        let legacy = MemoryChunkSource::try_new(schema(), cloud(4), 2).unwrap();
1009        let options = StreamOptions::new(2, MemoryBudget::new(24).unwrap()).unwrap();
1010        let mut source =
1011            LegacyBoundedSource::try_new(legacy, options, CancellationToken::default()).unwrap();
1012        let mut sink = LegacyBoundedSink::new(MemoryChunkSink::new());
1013        while let Some(chunk) = source.next_chunk() {
1014            sink.write_chunk(&chunk.unwrap()).unwrap();
1015        }
1016        BoundedSpatialRecordSink::finish(&mut sink).unwrap();
1017        let record = sink.into_inner().into_record().unwrap().unwrap();
1018        assert_eq!(record.cloud().len(), 4);
1019    }
1020}
spatialrust_records/bounded.rs

spatialrust_records/
bounded.rs
