spatialrust_pipeline/
workflow.rs

1//! Type-erased, metered bounded-memory streaming workflows.
2
3use std::sync::{Arc, Mutex, MutexGuard};
4use std::time::Instant;
5
6use spatialrust_math::Mat4;
7use spatialrust_records::{
8    BoundedSpatialRecordSink, BoundedSpatialRecordSource, CancellationToken, MemoryTracker,
9    RecordsError, RecordsResult, SchemaDescriptor, SpatialRecordChunk, StreamOptions,
10    StreamingReceipt,
11};
12
13use crate::{ChunkMapSource, StreamingVoxelConfig, StreamingVoxelSource};
14
15type ReceiptState = Arc<Mutex<StreamingReceipt>>;
16
17/// Composable bounded-memory point-cloud stream.
18pub struct StreamingPipeline {
19    source: Box<dyn BoundedSpatialRecordSource>,
20    receipt: ReceiptState,
21}
22
23impl StreamingPipeline {
24    /// Starts a workflow and meters chunks read from the original source.
25    pub fn new(
26        source: impl BoundedSpatialRecordSource + 'static,
27        source_id: impl Into<String>,
28    ) -> RecordsResult<Self> {
29        let receipt = Arc::new(Mutex::new(StreamingReceipt::new(source_id)?));
30        let source = MeteredInputSource { source, receipt: receipt.clone() };
31        Ok(Self { source: Box::new(source), receipt })
32    }
33
34    /// Returns the output schema.
35    #[must_use]
36    pub fn schema(&self) -> &SchemaDescriptor {
37        self.source.schema()
38    }
39
40    /// Returns the shared cooperative cancellation token.
41    #[must_use]
42    pub fn cancellation_token(&self) -> CancellationToken {
43        self.source.cancellation_token()
44    }
45
46    /// Adds an inclusive axis-aligned crop.
47    pub fn crop(self, min: [f32; 3], max: [f32; 3], invert: bool) -> RecordsResult<Self> {
48        let source = ChunkMapSource::crop(self.source, min, max, invert)?;
49        Ok(Self { source: Box::new(source), receipt: self.receipt })
50    }
51
52    /// Adds an affine position/normal transform.
53    pub fn transform(self, transform: Mat4<f32>) -> RecordsResult<Self> {
54        let source = ChunkMapSource::transform(self.source, transform)?;
55        Ok(Self { source: Box::new(source), receipt: self.receipt })
56    }
57
58    /// Adds deterministic global voxel aggregation backed by bounded spool storage.
59    pub fn voxel(self, config: StreamingVoxelConfig) -> RecordsResult<Self> {
60        let started = Instant::now();
61        let source = StreamingVoxelSource::try_build(self.source, config)?;
62        let spill_bytes = source.spool_bytes();
63        {
64            let mut receipt = lock_receipt(&self.receipt)?;
65            receipt.record_spill(spill_bytes)?;
66            receipt.record_phase("voxel", elapsed_ns(started), spill_bytes)?;
67            receipt.capture_memory(source.memory_tracker());
68        }
69        Ok(Self { source: Box::new(source), receipt: self.receipt })
70    }
71
72    /// Drains the workflow into a synchronous bounded sink and returns its receipt.
73    pub fn run_to_sink(
74        self,
75        sink: &mut dyn BoundedSpatialRecordSink,
76    ) -> RecordsResult<StreamingReceipt> {
77        let mut stream = self.into_iter();
78        for chunk in stream.by_ref() {
79            let chunk = chunk?;
80            sink.write_chunk(&chunk)?;
81        }
82        sink.finish()?;
83        stream.receipt()
84    }
85}
86
87impl IntoIterator for StreamingPipeline {
88    type Item = RecordsResult<SpatialRecordChunk>;
89    type IntoIter = StreamingPipelineIter;
90
91    fn into_iter(self) -> Self::IntoIter {
92        let tracker = self.source.memory_tracker().clone();
93        StreamingPipelineIter {
94            source: self.source,
95            receipt: self.receipt,
96            tracker,
97            completed: false,
98        }
99    }
100}
101
102/// Pull iterator that meters final output and exposes a live receipt snapshot.
103pub struct StreamingPipelineIter {
104    source: Box<dyn BoundedSpatialRecordSource>,
105    receipt: ReceiptState,
106    tracker: MemoryTracker,
107    completed: bool,
108}
109
110impl StreamingPipelineIter {
111    /// Returns the shared cooperative cancellation token.
112    #[must_use]
113    pub fn cancellation_token(&self) -> CancellationToken {
114        self.source.cancellation_token()
115    }
116
117    /// Returns the output record schema.
118    #[must_use]
119    pub fn schema(&self) -> &SchemaDescriptor {
120        self.source.schema()
121    }
122
123    /// Clones the receipt as observed at the latest completed chunk boundary.
124    pub fn receipt(&self) -> RecordsResult<StreamingReceipt> {
125        let mut receipt = lock_receipt(&self.receipt)?;
126        receipt.capture_memory(&self.tracker);
127        Ok(receipt.clone())
128    }
129}
130
131impl Iterator for StreamingPipelineIter {
132    type Item = RecordsResult<SpatialRecordChunk>;
133
134    fn next(&mut self) -> Option<Self::Item> {
135        if self.completed {
136            return None;
137        }
138        let next = self.source.next_chunk();
139        match next {
140            Some(Ok(chunk)) => {
141                let points = match u64::try_from(chunk.record().cloud().len()) {
142                    Ok(points) => points,
143                    Err(_) => {
144                        self.completed = true;
145                        return Some(Err(RecordsError::ReceiptOverflow(
146                            "pipeline output point count".into(),
147                        )));
148                    }
149                };
150                if let Err(error) = lock_receipt(&self.receipt).and_then(|mut receipt| {
151                    receipt.record_output_chunk(points, chunk.tracked_bytes())
152                }) {
153                    self.completed = true;
154                    return Some(Err(error));
155                }
156                Some(Ok(chunk))
157            }
158            Some(Err(error)) => {
159                self.completed = true;
160                Some(Err(error))
161            }
162            None => {
163                self.completed = true;
164                if let Ok(mut receipt) = lock_receipt(&self.receipt) {
165                    receipt.capture_memory(&self.tracker);
166                }
167                None
168            }
169        }
170    }
171}
172
173struct MeteredInputSource<S> {
174    source: S,
175    receipt: ReceiptState,
176}
177
178impl<S: BoundedSpatialRecordSource> BoundedSpatialRecordSource for MeteredInputSource<S> {
179    fn schema(&self) -> &SchemaDescriptor {
180        self.source.schema()
181    }
182
183    fn options(&self) -> &StreamOptions {
184        self.source.options()
185    }
186
187    fn memory_tracker(&self) -> &MemoryTracker {
188        self.source.memory_tracker()
189    }
190
191    fn cancellation_token(&self) -> CancellationToken {
192        self.source.cancellation_token()
193    }
194
195    fn max_chunk_bytes(&self) -> u64 {
196        self.source.max_chunk_bytes()
197    }
198
199    fn next_chunk(&mut self) -> Option<RecordsResult<SpatialRecordChunk>> {
200        match self.source.next_chunk()? {
201            Ok(chunk) => {
202                let points = match u64::try_from(chunk.record().cloud().len()) {
203                    Ok(points) => points,
204                    Err(_) => {
205                        return Some(Err(RecordsError::ReceiptOverflow(
206                            "pipeline input point count".into(),
207                        )));
208                    }
209                };
210                match lock_receipt(&self.receipt).and_then(|mut receipt| {
211                    receipt.record_input_chunk(points, chunk.tracked_bytes())
212                }) {
213                    Ok(()) => Some(Ok(chunk)),
214                    Err(error) => Some(Err(error)),
215                }
216            }
217            Err(error) => Some(Err(error)),
218        }
219    }
220}
221
222fn lock_receipt(receipt: &ReceiptState) -> RecordsResult<MutexGuard<'_, StreamingReceipt>> {
223    receipt
224        .lock()
225        .map_err(|_| RecordsError::InvalidReceipt("streaming receipt lock poisoned".into()))
226}
227
228fn elapsed_ns(started: Instant) -> u64 {
229    u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX)
230}
231
232#[cfg(test)]
233mod tests {
234    use super::StreamingPipeline;
235    use spatialrust_core::{HasPositions3, PointCloudBuilder, StandardSchemas};
236    use spatialrust_records::{
237        CancellationToken, MemoryBudget, RecyclingMemoryChunkSource, SchemaDescriptor,
238        SchemaVersion, StreamOptions,
239    };
240
241    fn source() -> RecyclingMemoryChunkSource {
242        let mut builder = PointCloudBuilder::xyz();
243        for point in [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]] {
244            builder.push_point(point).unwrap();
245        }
246        let schema = SchemaDescriptor::try_new(
247            "workflow.xyz",
248            SchemaVersion::new(1, 0),
249            StandardSchemas::point_xyz(),
250        )
251        .unwrap();
252        let options = StreamOptions::new(2, MemoryBudget::new(4096).unwrap()).unwrap();
253        RecyclingMemoryChunkSource::try_new(
254            schema,
255            builder.build().unwrap(),
256            options,
257            CancellationToken::default(),
258        )
259        .unwrap()
260    }
261
262    #[test]
263    fn meters_input_and_output_around_composable_crop() {
264        let pipeline = StreamingPipeline::new(source(), "memory")
265            .unwrap()
266            .crop([0.5, -1.0, -1.0], [2.0, 1.0, 1.0], false)
267            .unwrap();
268        let mut stream = pipeline.into_iter();
269        let mut x = Vec::new();
270        for chunk in stream.by_ref() {
271            let chunk = chunk.unwrap();
272            x.extend_from_slice(chunk.record().cloud().positions3().unwrap().0);
273        }
274        assert_eq!(x, [1.0, 2.0]);
275        let receipt = stream.receipt().unwrap();
276        assert_eq!(receipt.input_points(), 3);
277        assert_eq!(receipt.output_points(), 2);
278        assert_eq!(receipt.chunks_read(), 2);
279        assert_eq!(receipt.chunks_written(), 2);
280    }
281
282    #[test]
283    fn iterator_observes_shared_cancellation() {
284        let pipeline = StreamingPipeline::new(source(), "memory").unwrap();
285        let token = pipeline.cancellation_token();
286        token.cancel();
287        let mut stream = pipeline.into_iter();
288        assert!(stream.next().unwrap().is_err());
289        assert!(stream.next().is_none());
290    }
291}
spatialrust_pipeline/workflow.rs

spatialrust_pipeline/
workflow.rs
