fission_core/
data_stream.rs

1//! Runtime-owned streams for large binary data.
2//!
3//! Large data should not travel through action payloads or capability results as
4//! `Vec<u8>`. Host capabilities register streams here, then reducers pass the
5//! returned [`DataStreamId`] to jobs or services that consume the stream
6//! asynchronously.
7
8use bytes::{Bytes, BytesMut};
9use futures_core::Stream;
10use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12use std::fmt;
13use std::pin::Pin;
14use std::sync::atomic::{AtomicU64, Ordering};
15use std::sync::{Arc, Mutex};
16use std::task::{Context, Poll};
17
18/// Unique identifier for a runtime-owned binary stream.
19#[derive(
20    Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
21)]
22pub struct DataStreamId(pub u64);
23
24/// Machine-readable class of stream failure.
25#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
26pub enum FissionDataStreamErrorKind {
27    /// The stream id is unknown.
28    NotFound,
29    /// The stream was already opened by a previous consumer.
30    AlreadyConsumed,
31    /// The host or user cancelled the stream.
32    Cancelled,
33    /// The host denied access to the data.
34    PermissionDenied,
35    /// The host cannot provide this stream operation.
36    Unsupported,
37    /// I/O failed while reading the stream.
38    Io,
39    /// The data was malformed or could not be decoded by the producer.
40    InvalidData,
41    /// Any other host-specific failure.
42    Other,
43}
44
45/// Error item yielded by [`FissionDataStream`].
46#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
47pub struct FissionDataStreamError {
48    /// Stable error category.
49    pub kind: FissionDataStreamErrorKind,
50    /// Human-readable diagnostic message.
51    pub message: String,
52}
53
54impl FissionDataStreamError {
55    /// Creates a stream error.
56    pub fn new(kind: FissionDataStreamErrorKind, message: impl Into<String>) -> Self {
57        Self {
58            kind,
59            message: message.into(),
60        }
61    }
62
63    /// Creates a not-found stream error.
64    pub fn not_found(id: DataStreamId) -> Self {
65        Self::new(
66            FissionDataStreamErrorKind::NotFound,
67            format!("data stream {} was not found", id.0),
68        )
69    }
70
71    /// Creates an already-consumed stream error.
72    pub fn already_consumed(id: DataStreamId) -> Self {
73        Self::new(
74            FissionDataStreamErrorKind::AlreadyConsumed,
75            format!("data stream {} was already consumed", id.0),
76        )
77    }
78}
79
80impl std::error::Error for FissionDataStreamError {}
81
82impl fmt::Display for FissionDataStreamError {
83    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
84        write!(f, "{:?}: {}", self.kind, self.message)
85    }
86}
87
88impl From<std::io::Error> for FissionDataStreamError {
89    fn from(error: std::io::Error) -> Self {
90        Self::new(FissionDataStreamErrorKind::Io, error.to_string())
91    }
92}
93
94/// Fission's framework-wide stream trait for large binary payloads.
95///
96/// It intentionally builds on `futures_core::Stream` because Rust does not
97/// currently provide a stable `std::stream::Stream`.
98pub trait FissionDataStream:
99    Stream<Item = Result<Bytes, FissionDataStreamError>> + Send + 'static
100{
101}
102
103impl<T> FissionDataStream for T where
104    T: Stream<Item = Result<Bytes, FissionDataStreamError>> + Send + 'static
105{
106}
107
108/// Boxed dynamic Fission data stream.
109pub type BoxFissionDataStream = Pin<Box<dyn FissionDataStream>>;
110
111/// Host/runtime registry for one-shot data streams.
112#[derive(Clone, Default)]
113pub struct DataStreamRegistry {
114    inner: Arc<DataStreamRegistryInner>,
115}
116
117#[derive(Default)]
118struct DataStreamRegistryInner {
119    next_id: AtomicU64,
120    streams: Mutex<HashMap<DataStreamId, Option<BoxFissionDataStream>>>,
121}
122
123impl fmt::Debug for DataStreamRegistry {
124    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
125        let len = self
126            .inner
127            .streams
128            .lock()
129            .map(|streams| streams.len())
130            .unwrap_or_default();
131        f.debug_struct("DataStreamRegistry")
132            .field("streams", &len)
133            .finish()
134    }
135}
136
137impl DataStreamRegistry {
138    /// Creates an empty stream registry.
139    pub fn new() -> Self {
140        Self::default()
141    }
142
143    /// Registers a one-shot stream and returns its public id.
144    pub fn register(&self, stream: BoxFissionDataStream) -> DataStreamId {
145        let id = DataStreamId(self.inner.next_id.fetch_add(1, Ordering::Relaxed) + 1);
146        self.inner
147            .streams
148            .lock()
149            .expect("data stream registry lock poisoned")
150            .insert(id, Some(stream));
151        id
152    }
153
154    /// Opens the stream for consumption.
155    ///
156    /// Streams are single-consumer by default. Opening removes the stream from
157    /// the registry slot so subsequent opens fail deterministically.
158    pub fn open(&self, id: DataStreamId) -> Result<BoxFissionDataStream, FissionDataStreamError> {
159        let mut streams = self.inner.streams.lock().map_err(|_| {
160            FissionDataStreamError::new(
161                FissionDataStreamErrorKind::Other,
162                "data stream registry lock was poisoned",
163            )
164        })?;
165        match streams.get_mut(&id) {
166            Some(slot) => slot
167                .take()
168                .ok_or_else(|| FissionDataStreamError::already_consumed(id)),
169            None => Err(FissionDataStreamError::not_found(id)),
170        }
171    }
172
173    /// Releases a stream without consuming it.
174    pub fn release(&self, id: DataStreamId) -> bool {
175        self.inner
176            .streams
177            .lock()
178            .map(|mut streams| streams.remove(&id).is_some())
179            .unwrap_or(false)
180    }
181}
182
183struct SingleChunkDataStream {
184    chunk: Option<Result<Bytes, FissionDataStreamError>>,
185}
186
187impl Stream for SingleChunkDataStream {
188    type Item = Result<Bytes, FissionDataStreamError>;
189
190    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
191        Poll::Ready(self.chunk.take())
192    }
193}
194
195/// Creates a stream that yields a single byte chunk.
196pub fn single_chunk_data_stream(bytes: impl Into<Bytes>) -> BoxFissionDataStream {
197    Box::pin(SingleChunkDataStream {
198        chunk: Some(Ok(bytes.into())),
199    })
200}
201
202/// Creates an empty stream.
203pub fn empty_data_stream() -> BoxFissionDataStream {
204    Box::pin(SingleChunkDataStream { chunk: None })
205}
206
207/// Collects a data stream into one contiguous byte buffer.
208///
209/// Use this only at shell boundaries that must interoperate with platform APIs
210/// requiring contiguous memory. Application-facing large-data flows should keep
211/// streaming chunks instead.
212pub async fn collect_data_stream(
213    mut stream: BoxFissionDataStream,
214) -> Result<Bytes, FissionDataStreamError> {
215    let mut out = BytesMut::new();
216    while let Some(chunk) = std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await {
217        out.extend_from_slice(&chunk?);
218    }
219    Ok(out.freeze())
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225    use std::future::Future;
226    use std::task::{RawWaker, RawWakerVTable, Waker};
227
228    fn noop_waker() -> Waker {
229        unsafe fn clone(data: *const ()) -> RawWaker {
230            RawWaker::new(data, &VTABLE)
231        }
232        unsafe fn wake(_data: *const ()) {}
233        unsafe fn wake_by_ref(_data: *const ()) {}
234        unsafe fn drop(_data: *const ()) {}
235
236        static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);
237        unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) }
238    }
239
240    fn block_on_ready<F: Future>(future: F) -> F::Output {
241        let waker = noop_waker();
242        let mut cx = Context::from_waker(&waker);
243        let mut future = Box::pin(future);
244        match future.as_mut().poll(&mut cx) {
245            Poll::Ready(output) => output,
246            Poll::Pending => panic!("test future unexpectedly pending"),
247        }
248    }
249
250    #[test]
251    fn registry_opens_stream_once() {
252        let registry = DataStreamRegistry::new();
253        let id = registry.register(single_chunk_data_stream(Bytes::from_static(b"abc")));
254
255        let bytes = block_on_ready(collect_data_stream(registry.open(id).unwrap())).unwrap();
256        assert_eq!(bytes.as_ref(), b"abc");
257
258        match registry.open(id) {
259            Ok(_) => panic!("stream should be consumed"),
260            Err(error) => assert_eq!(error.kind, FissionDataStreamErrorKind::AlreadyConsumed),
261        }
262    }
263
264    #[test]
265    fn registry_release_removes_unopened_stream() {
266        let registry = DataStreamRegistry::new();
267        let id = registry.register(empty_data_stream());
268
269        assert!(registry.release(id));
270        match registry.open(id) {
271            Ok(_) => panic!("released stream should be removed"),
272            Err(error) => assert_eq!(error.kind, FissionDataStreamErrorKind::NotFound),
273        }
274    }
275
276    #[test]
277    fn cloned_registries_share_streams() {
278        let registry = DataStreamRegistry::new();
279        let clone = registry.clone();
280        let id = registry.register(single_chunk_data_stream(Bytes::from_static(b"shared")));
281
282        let bytes = block_on_ready(collect_data_stream(clone.open(id).unwrap())).unwrap();
283        assert_eq!(bytes.as_ref(), b"shared");
284    }
285}
fission_core/data_stream.rs

fission_core/
data_stream.rs
