hurray_io/stream/sans_io/read.rs
1//! The sequential read half of the sans-io stream core.
2//!
3//! One state machine, driven one step at a time: it asks for bytes, is told how many
4//! arrived, and produces tensors. It never reads, never waits, and holds no transport.
5//!
6//! Tensors, not items: a composite head leaves here as the buffer-less tensor it is on the
7//! wire, and [`CompositeAssembler`][super::CompositeAssembler] puts the tree back together.
8
9use hurray_core::{BufferHandle, SyncMode, TensorDescriptor};
10
11use super::dest::{Destination, DestinationAlloc, DestinationRequest, SystemDestinations};
12use super::Unusable;
13use crate::stream::types::{StreamReaderOptions, StreamTensor};
14use crate::{Error, Result};
15
16/// magic(4) + version(1) + flags(1) + descriptor_length(4)
17const PREFIX_LEN: usize = 10;
18
19/// What the core needs next from its driver.
20///
21/// Exactly one request is outstanding at a time, and the driver answers it — with
22/// [`report_filled`][StreamReaderCore::report_filled] or
23/// [`take_tensor`][StreamReaderCore::take_tensor] — before the core produces another
24/// (ADR-041 § 4).
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum ReadStep {
27 /// Fill up to `len` bytes of [`input_dst`][StreamReaderCore::input_dst], then report
28 /// how many arrived. Filling fewer is normal; `0` means end of input.
29 Input {
30 /// Bytes remaining for the current target. Never `0`.
31 len: usize,
32 },
33 /// A tensor is complete. Collect it with
34 /// [`take_tensor`][StreamReaderCore::take_tensor].
35 Tensor,
36 /// End of input, reached cleanly at a tensor boundary.
37 Finished,
38}
39
40/// The tensor currently being assembled.
41struct Pending<D> {
42 descriptor: TensorDescriptor,
43 /// Buffers already filled and frozen, in declaration order.
44 buffers: Vec<bytes::Bytes>,
45 /// The destination for the buffer being filled.
46 dest: Option<D>,
47}
48
49/// Where the machine is. `Failed` and `Unusable` are distinct by design: a format error
50/// ends the stream with the position still known, while `Unusable` means a byte count was
51/// lost and the position is unrecoverable (ADR-041 § 6).
52#[derive(Debug, Clone, Copy)]
53enum State {
54 /// Filling the 10-byte frame prefix. End of input *here* is the clean one.
55 Prefix,
56 /// Filling the rest of the descriptor.
57 Descriptor,
58 /// Filling buffer `index` of the pending tensor.
59 Buffer { index: usize },
60 /// A tensor is waiting to be taken.
61 TensorReady,
62 /// End of input at a tensor boundary.
63 Finished,
64 /// A previous call returned an error; the stream cannot continue.
65 Failed,
66 /// Bytes moved and could not be accounted for.
67 Unusable(Unusable),
68}
69
70/// Decodes the Hurray stream wire format from bytes a driver supplies.
71///
72/// The core owns the destination every payload byte lands in, so the transport's write into
73/// it is the only copy (ADR-041 § 2), and that destination carries the alignment the buffer
74/// handle *declares* (§ 5). Position is kept across every short read, so a driver that
75/// stops — because its transport would block, or because its future was cancelled — can
76/// resume on the next call (§ 6).
77///
78/// # Examples
79///
80/// Driving it over a slice, with no runtime anywhere:
81///
82/// ```
83/// use hurray_core::{
84/// BufferHandle, DeviceTag, ElementType, LayoutDescriptor, Shape, SyncMode,
85/// TensorDescriptor, MIN_BUFFER_ALIGNMENT,
86/// };
87/// use hurray_io::stream::{ReadStep, StreamReaderCore, StreamWriterCore, WriteStep};
88///
89/// // A 64-byte float32 tensor, encoded with the write core.
90/// let handle =
91/// BufferHandle::new(64, MIN_BUFFER_ALIGNMENT, DeviceTag::Cpu, SyncMode::ProducerSynced)?;
92/// let descriptor = TensorDescriptor::new(
93/// 1, 0, ElementType::Float32, Shape::new(vec![4u64, 4])?, 0,
94/// LayoutDescriptor::RowMajor, vec![handle], None, None, None, None,
95/// )?;
96/// let payload: Vec<u8> = (0u8..64).collect();
97/// // Bound, not a temporary: the operation borrows the payloads until the item is
98/// // written, which is ADR-041 § 3's lifetime rule showing up in the type system.
99/// let buffers: [&[u8]; 1] = [&payload];
100///
101/// let mut wire = Vec::new();
102/// let mut writer = StreamWriterCore::new();
103/// let mut op = writer.tensor(&descriptor, &buffers)?;
104/// while let WriteStep::Output { .. } = op.step()? {
105/// let src = op.output_src();
106/// wire.extend_from_slice(src);
107/// let n = src.len();
108/// op.report_written(n)?;
109/// }
110///
111/// // Read it back, handing the core 7 bytes at a time to exercise short reads.
112/// let mut core = StreamReaderCore::new();
113/// let mut rest = &wire[..];
114/// let mut tensors = Vec::new();
115/// loop {
116/// match core.step()? {
117/// ReadStep::Input { len } => {
118/// let n = len.min(7).min(rest.len());
119/// core.input_dst()[..n].copy_from_slice(&rest[..n]);
120/// rest = &rest[n..];
121/// core.report_filled(n)?;
122/// }
123/// ReadStep::Tensor => tensors.extend(core.take_tensor()),
124/// ReadStep::Finished => break,
125/// }
126/// }
127///
128/// assert_eq!(tensors.len(), 1);
129/// let tensor = &tensors[0];
130/// assert_eq!(tensor.buffers[0].as_ref(), &payload[..]);
131/// // The destination carried the declared alignment, so a SIMD load is safe on it.
132/// assert_eq!(tensor.buffers[0].as_ptr() as usize % 64, 0);
133/// # Ok::<(), Box<dyn std::error::Error>>(())
134/// ```
135pub struct StreamReaderCore<A: DestinationAlloc = SystemDestinations> {
136 options: StreamReaderOptions,
137 destinations: A,
138 state: State,
139 /// Descriptor bytes, prefix included, for the frame being read.
140 staging: Vec<u8>,
141 /// Length of the target currently being filled: the prefix, the rest of the
142 /// descriptor, or one buffer.
143 target: usize,
144 /// Bytes of that target already filled.
145 filled: usize,
146 pending: Option<Pending<A::Dest>>,
147 tensor: Option<StreamTensor>,
148}
149
150impl Default for StreamReaderCore {
151 fn default() -> Self {
152 Self::new()
153 }
154}
155
156// The default allocator's constructors: `StreamReaderCore` with no type argument.
157impl StreamReaderCore<SystemDestinations> {
158 /// Creates a core with default options and the system allocator.
159 ///
160 /// # Examples
161 ///
162 /// See the type-level example.
163 pub fn new() -> Self {
164 Self::with_options(StreamReaderOptions::default())
165 }
166
167 /// Creates a core with custom options.
168 ///
169 /// `max_descriptor_bytes`, `max_buffer_bytes` and `enforce_cross_machine_sync` apply
170 /// here; `max_composite_depth` does not, because this core does not assemble
171 /// composites — it bounds
172 /// [`CompositeAssembler`][super::CompositeAssembler] instead.
173 ///
174 /// # Examples
175 ///
176 /// ```
177 /// use hurray_io::stream::{StreamReaderCore, StreamReaderOptions};
178 ///
179 /// let core = StreamReaderCore::with_options(StreamReaderOptions {
180 /// max_buffer_bytes: 1 << 20,
181 /// ..Default::default()
182 /// });
183 /// assert!(!core.is_unusable());
184 /// ```
185 pub fn with_options(options: StreamReaderOptions) -> Self {
186 Self::with_destinations(options, SystemDestinations)
187 }
188}
189
190impl<A: DestinationAlloc> StreamReaderCore<A> {
191 /// Creates a core whose read destinations come from `destinations`.
192 ///
193 /// This is the seam ADR-041 left open: a binding may supply the allocation, so a read
194 /// can land directly in memory its runtime owns, as long as the destination satisfies
195 /// the alignment the handle declares.
196 ///
197 /// # Examples
198 ///
199 /// ```
200 /// use hurray_io::stream::{
201 /// AlignedDestination, DestinationAlloc, DestinationRequest, StreamReaderCore,
202 /// StreamReaderOptions, SystemDestinations,
203 /// };
204 ///
205 /// /// Counts what the core asks for, then defers to the default allocator.
206 /// struct Counting {
207 /// seen: Vec<(u64, u32)>,
208 /// }
209 ///
210 /// impl DestinationAlloc for Counting {
211 /// type Dest = AlignedDestination;
212 ///
213 /// fn allocate(
214 /// &mut self,
215 /// request: DestinationRequest,
216 /// ) -> hurray_io::Result<Self::Dest> {
217 /// self.seen.push((request.byte_size, request.alignment));
218 /// SystemDestinations.allocate(request)
219 /// }
220 /// }
221 ///
222 /// let core = StreamReaderCore::with_destinations(
223 /// StreamReaderOptions::default(),
224 /// Counting { seen: Vec::new() },
225 /// );
226 /// assert!(!core.is_unusable());
227 /// ```
228 pub fn with_destinations(options: StreamReaderOptions, destinations: A) -> Self {
229 Self {
230 options,
231 destinations,
232 state: State::Prefix,
233 staging: vec![0u8; PREFIX_LEN],
234 target: PREFIX_LEN,
235 filled: 0,
236 pending: None,
237 tensor: None,
238 }
239 }
240
241 /// Advances the machine and returns what it needs next.
242 ///
243 /// Never blocks and never waits: every transition here is decode, allocate or
244 /// bookkeeping, which is what lets one core serve a blocking driver and an event loop
245 /// alike (ADR-041 § 4).
246 ///
247 /// # Errors
248 ///
249 /// - [`Error::InvalidHeader`], [`Error::Core`] — the frame or descriptor is malformed
250 /// - [`Error::FrameTooLarge`] — a descriptor or buffer exceeds the configured limit
251 /// - [`Error::InvalidCrossMachineSyncMode`] — cross-machine mode and a buffer that is
252 /// not `ProducerSynced`
253 /// - [`Error::UnallocatableBuffer`] — a destination the allocator cannot provide
254 /// - [`Error::CoreFailed`], [`Error::CoreUnusable`] — a previous call ended this core
255 ///
256 /// # Examples
257 ///
258 /// See the type-level example.
259 pub fn step(&mut self) -> Result<ReadStep> {
260 loop {
261 match self.state {
262 State::Unusable(reason) => return Err(Error::CoreUnusable(reason)),
263 State::Failed => return Err(Error::CoreFailed),
264 State::Finished => return Ok(ReadStep::Finished),
265 State::TensorReady => return Ok(ReadStep::Tensor),
266 State::Prefix => match self.remaining() {
267 0 => self.begin_descriptor()?,
268 len => return Ok(ReadStep::Input { len }),
269 },
270 State::Descriptor => match self.remaining() {
271 0 => self.decode_descriptor()?,
272 len => return Ok(ReadStep::Input { len }),
273 },
274 State::Buffer { index } => match self.remaining() {
275 0 => self.finish_buffer(index)?,
276 len => return Ok(ReadStep::Input { len }),
277 },
278 }
279 }
280 }
281
282 /// The destination for the outstanding [`ReadStep::Input`] request.
283 ///
284 /// Its address does not move while the request stands, so a driver may hold a pointer
285 /// into it across a suspension (ADR-041 § 6a). Empty when no request is outstanding.
286 ///
287 /// # Examples
288 ///
289 /// See the type-level example.
290 pub fn input_dst(&mut self) -> &mut [u8] {
291 match self.state {
292 State::Prefix | State::Descriptor => &mut self.staging[self.filled..self.target],
293 State::Buffer { .. } => {
294 let filled = self.filled;
295 match self.pending.as_mut().and_then(|p| p.dest.as_mut()) {
296 Some(dest) => &mut dest.as_mut_slice()[filled..],
297 None => &mut [],
298 }
299 }
300 _ => &mut [],
301 }
302 }
303
304 /// Reports that `n` bytes were filled into [`input_dst`][Self::input_dst].
305 ///
306 /// Deliberately the cheapest call in the API: it records a count and nothing else, so
307 /// a driver can pair the transfer and this report inside one synchronous callback,
308 /// where no cancellation can interleave (ADR-041 § 6a).
309 ///
310 /// A short fill is normal — the core keeps its position and re-offers the remainder.
311 /// `n == 0` means the transport reached end of input: clean at an item boundary, an
312 /// error anywhere else.
313 ///
314 /// # Errors
315 ///
316 /// - [`Error::UnexpectedEof`] — end of input mid-item
317 /// - [`Error::InvalidProgressReport`] — more bytes reported than were offered, which
318 /// also makes the core unusable: the position is lost
319 ///
320 /// # Examples
321 ///
322 /// See the type-level example.
323 pub fn report_filled(&mut self, n: usize) -> Result<()> {
324 match self.state {
325 State::Unusable(reason) => return Err(Error::CoreUnusable(reason)),
326 State::Failed => return Err(Error::CoreFailed),
327 State::Finished | State::TensorReady => {
328 return Err(Error::InvalidProgressReport {
329 reported: n,
330 offered: 0,
331 })
332 }
333 _ => {}
334 }
335
336 if n == 0 {
337 // § 4: only a zero-length read is end of input, and only at a tensor
338 // boundary is it the clean one.
339 if matches!(self.state, State::Prefix) && self.filled == 0 {
340 self.state = State::Finished;
341 return Ok(());
342 }
343 return Err(self.fail(Error::UnexpectedEof));
344 }
345
346 let offered = self.remaining();
347 if n > offered {
348 // Bytes moved that the core cannot place: its position no longer matches the
349 // transport's, and no inspection recovers it (§ 6 — 0x00 is valid data).
350 self.state = State::Unusable(Unusable::LostReadCount);
351 return Err(Error::InvalidProgressReport {
352 reported: n,
353 offered,
354 });
355 }
356 self.filled += n;
357 Ok(())
358 }
359
360 /// Takes the completed tensor announced by [`ReadStep::Tensor`].
361 ///
362 /// Returns `None` when no tensor is ready. The tensor is handed over only once every
363 /// one of its buffers is full, so a partially filled destination is never observable
364 /// (ADR-041 § 6).
365 ///
366 /// A composite head arrives here like any other tensor, carrying no buffers; putting
367 /// the tree back together is [`CompositeAssembler`][crate::stream::CompositeAssembler]'s
368 /// job, not the core's.
369 ///
370 /// # Examples
371 ///
372 /// See the type-level example.
373 pub fn take_tensor(&mut self) -> Option<StreamTensor> {
374 if !matches!(self.state, State::TensorReady) {
375 return None;
376 }
377 let tensor = self.tensor.take();
378 self.reset_to_prefix();
379 tensor
380 }
381
382 /// Marks the core unusable: a transport moved bytes whose count was lost.
383 ///
384 /// The position is unrecoverable, so every later call reports it by name rather than
385 /// letting the core diverge silently from the transport (ADR-041 § 6). A readiness-
386 /// driven driver never needs this; one built on a transfer-and-resolve primitive such
387 /// as `sock_recv_into` does, for the window between the transfer and the report.
388 ///
389 /// # Examples
390 ///
391 /// ```
392 /// use hurray_io::stream::StreamReaderCore;
393 ///
394 /// let mut core = StreamReaderCore::new();
395 /// core.mark_unusable();
396 /// assert!(core.is_unusable());
397 /// assert!(core.step().is_err());
398 /// ```
399 pub fn mark_unusable(&mut self) {
400 self.state = State::Unusable(Unusable::LostReadCount);
401 }
402
403 /// Whether the core has been marked unusable.
404 ///
405 /// # Examples
406 ///
407 /// See [`mark_unusable`][Self::mark_unusable].
408 pub fn is_unusable(&self) -> bool {
409 matches!(self.state, State::Unusable(_))
410 }
411
412 // ── transitions ───────────────────────────────────────────────────────────
413
414 /// Bytes still wanted for the current target.
415 fn remaining(&self) -> usize {
416 match self.state {
417 State::Prefix | State::Descriptor | State::Buffer { .. } => {
418 self.target.saturating_sub(self.filled)
419 }
420 _ => 0,
421 }
422 }
423
424 /// Records a terminal failure and hands the error back for propagation.
425 fn fail(&mut self, error: Error) -> Error {
426 self.state = State::Failed;
427 error
428 }
429
430 /// Prefix complete: size the staging buffer for the rest of the descriptor.
431 fn begin_descriptor(&mut self) -> Result<()> {
432 let descriptor_length = u32::from_le_bytes(
433 self.staging[6..10]
434 .try_into()
435 .unwrap_or([0xFF, 0xFF, 0xFF, 0xFF]),
436 ) as u64;
437
438 if descriptor_length < PREFIX_LEN as u64 {
439 let error = Error::InvalidHeader(format!(
440 "descriptor_length {descriptor_length} is below minimum {PREFIX_LEN}"
441 ));
442 return Err(self.fail(error));
443 }
444 let limit = self.options.max_descriptor_bytes;
445 let total = match usize::try_from(descriptor_length) {
446 Ok(total) if descriptor_length <= limit => total,
447 _ => {
448 let error = Error::FrameTooLarge {
449 kind: "descriptor",
450 value: descriptor_length,
451 limit,
452 };
453 return Err(self.fail(error));
454 }
455 };
456
457 self.staging.resize(total, 0);
458 self.target = total;
459 self.state = State::Descriptor;
460 Ok(())
461 }
462
463 /// Descriptor complete: decode it and start filling its buffers.
464 fn decode_descriptor(&mut self) -> Result<()> {
465 let descriptor = match TensorDescriptor::decode(&self.staging) {
466 Ok(descriptor) => descriptor,
467 Err(e) => return Err(self.fail(e.into())),
468 };
469
470 // A composite head needs no special case here: it is the one descriptor that
471 // declares no buffers, so it completes immediately and is handed over like any
472 // other tensor.
473 let buffer_count = descriptor.buffers.len();
474 self.pending = Some(Pending {
475 descriptor,
476 buffers: Vec::with_capacity(buffer_count),
477 dest: None,
478 });
479 if buffer_count == 0 {
480 return self.complete_tensor();
481 }
482 self.begin_buffer(0)
483 }
484
485 /// Allocates the destination for buffer `index` and enters `Buffer { index }`.
486 fn begin_buffer(&mut self, index: usize) -> Result<()> {
487 let handle = match self.buffer_handle(index) {
488 Some(handle) => handle,
489 None => return Err(self.fail(Error::CoreFailed)),
490 };
491
492 if self.options.enforce_cross_machine_sync && handle.sync_mode() != SyncMode::ProducerSynced
493 {
494 let error = Error::InvalidCrossMachineSyncMode {
495 index,
496 actual: handle.sync_mode().to_byte(),
497 };
498 return Err(self.fail(error));
499 }
500 if handle.byte_size() > self.options.max_buffer_bytes {
501 let error = Error::FrameTooLarge {
502 kind: "buffer",
503 value: handle.byte_size(),
504 limit: self.options.max_buffer_bytes,
505 };
506 return Err(self.fail(error));
507 }
508
509 let request = DestinationRequest {
510 index,
511 byte_size: handle.byte_size(),
512 alignment: handle.alignment(),
513 };
514 let len = match usize::try_from(request.byte_size) {
515 Ok(len) => len,
516 Err(_) => {
517 let error = Error::UnallocatableBuffer {
518 index,
519 byte_size: request.byte_size,
520 alignment: request.alignment,
521 };
522 return Err(self.fail(error));
523 }
524 };
525 let dest = match self.destinations.allocate(request) {
526 Ok(dest) => dest,
527 Err(e) => return Err(self.fail(e)),
528 };
529
530 match self.pending.as_mut() {
531 Some(pending) => pending.dest = Some(dest),
532 None => return Err(self.fail(Error::CoreFailed)),
533 }
534 self.target = len;
535 self.filled = 0;
536 self.state = State::Buffer { index };
537 Ok(())
538 }
539
540 /// Buffer `index` is full: freeze it, then move to the next one or finish the tensor.
541 fn finish_buffer(&mut self, index: usize) -> Result<()> {
542 let (frozen, buffer_count) = match self.pending.as_mut() {
543 Some(pending) => match pending.dest.take() {
544 Some(dest) => (dest.freeze(), pending.descriptor.buffers.len()),
545 None => return Err(self.fail(Error::CoreFailed)),
546 },
547 None => return Err(self.fail(Error::CoreFailed)),
548 };
549 match self.pending.as_mut() {
550 Some(pending) => pending.buffers.push(frozen),
551 None => return Err(self.fail(Error::CoreFailed)),
552 }
553
554 let next = index + 1;
555 if next < buffer_count {
556 self.begin_buffer(next)
557 } else {
558 self.complete_tensor()
559 }
560 }
561
562 /// Hands the pending tensor to the caller.
563 fn complete_tensor(&mut self) -> Result<()> {
564 let pending = match self.pending.take() {
565 Some(pending) => pending,
566 None => return Err(self.fail(Error::CoreFailed)),
567 };
568 self.tensor = Some(StreamTensor {
569 descriptor: pending.descriptor,
570 buffers: pending.buffers,
571 });
572 self.state = State::TensorReady;
573 Ok(())
574 }
575
576 /// Rewinds to the start of the next frame.
577 fn reset_to_prefix(&mut self) {
578 self.filled = 0;
579 self.staging.clear();
580 self.staging.resize(PREFIX_LEN, 0);
581 self.target = PREFIX_LEN;
582 self.state = State::Prefix;
583 }
584
585 fn buffer_handle(&self, index: usize) -> Option<BufferHandle> {
586 self.pending
587 .as_ref()?
588 .descriptor
589 .buffers
590 .get(index)
591 .copied()
592 }
593}