hurray_io/stream/reader.rs
1use tokio::io::{AsyncRead, AsyncReadExt};
2
3use crate::stream::sans_io::{CompositeAssembler, ReadStep, StreamReaderCore};
4// `pub use`, not `use`: 0.1.0 published these at `stream::reader::*` as well as at
5// `stream::*`, because this module is public and defined them. They belong to the core
6// now, and the published paths keep resolving. `hurray_io::stream::*` is canonical.
7pub use crate::stream::types::{
8 StreamComposite, StreamItem, StreamReaderOptions, StreamTensor, DEFAULT_MAX_COMPOSITE_DEPTH,
9 DEFAULT_MAX_DESCRIPTOR_BYTES,
10};
11use crate::{Error, Result};
12
13/// Reads tensors from an async source in the Hurray streaming wire format.
14///
15/// Call [`next_tensor`][StreamReader::next_tensor] in a loop until it returns
16/// `Ok(None)` (clean EOF).
17///
18/// This is a driver, not a decoder: [`StreamReaderCore`] holds the wire format and the
19/// read position, and this type only moves bytes into the destination the core asks for
20/// (ADR-041 § 7).
21///
22/// # Examples
23///
24/// ```no_run
25/// # #[tokio::main]
26/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
27/// use hurray_io::stream::StreamReader;
28///
29/// let wire: &[u8] = &[]; // replace with actual data
30/// let mut reader = StreamReader::new(wire);
31/// while let Some(tensor) = reader.next_tensor().await? {
32/// println!(
33/// "tensor with {} buffer(s)",
34/// tensor.buffers.len()
35/// );
36/// }
37/// # Ok(())
38/// # }
39/// ```
40pub struct StreamReader<R> {
41 inner: R,
42 core: StreamReaderCore,
43 // Held across calls, not created per call: between successful `next_item` calls it is
44 // provably empty (a call returns only once a *top-level* item closes), so this changes
45 // nothing observable today — but a half-assembled tree is exactly what has to survive
46 // a suspension when the readiness-driven driver arrives (ADR-041 § 6a).
47 assembler: CompositeAssembler,
48}
49
50impl<R: AsyncRead + Unpin> StreamReader<R> {
51 /// Creates a reader with default options.
52 pub fn new(inner: R) -> Self {
53 Self::with_options(inner, StreamReaderOptions::default())
54 }
55
56 /// Creates a reader that enforces `ProducerSynced` on every buffer.
57 ///
58 /// Use when the stream was produced by a remote machine and GPU/semaphore
59 /// sync primitives are not meaningful locally.
60 pub fn cross_machine(inner: R) -> Self {
61 Self::with_options(
62 inner,
63 StreamReaderOptions {
64 enforce_cross_machine_sync: true,
65 ..Default::default()
66 },
67 )
68 }
69
70 /// Creates a reader with custom options.
71 pub fn with_options(inner: R, options: StreamReaderOptions) -> Self {
72 Self {
73 inner,
74 assembler: CompositeAssembler::new(options.max_composite_depth),
75 core: StreamReaderCore::with_options(options),
76 }
77 }
78
79 /// Reads the next tensor from the stream.
80 ///
81 /// Returns `Ok(None)` on a clean EOF (no bytes remaining before a descriptor starts).
82 ///
83 /// Oblivious to composition: a composite head comes back like any other tensor, with
84 /// no buffers, followed by its members in wire order. Use
85 /// [`next_item`][StreamReader::next_item] to have them assembled.
86 ///
87 /// # Errors
88 ///
89 /// - [`Error::UnexpectedEof`] — stream ended mid-descriptor or mid-buffer
90 /// - [`Error::InvalidHeader`] — malformed descriptor prefix
91 /// - [`Error::FrameTooLarge`] — descriptor or buffer exceeds configured limit
92 /// - [`Error::InvalidCrossMachineSyncMode`] — cross-machine mode and a non-`ProducerSynced` buffer
93 /// - [`Error::UnallocatableBuffer`] — a buffer's declared size and alignment cannot be allocated
94 /// - [`Error::Core`] — descriptor decode failed
95 /// - [`Error::Io`] — underlying read error
96 pub async fn next_tensor(&mut self) -> Result<Option<StreamTensor>> {
97 loop {
98 match self.core.step()? {
99 ReadStep::Input { .. } => self.fill().await?,
100 // `take_tensor` yields `Some` exactly when `step` reported `Tensor`.
101 ReadStep::Tensor => return Ok(Some(self.take_tensor()?)),
102 ReadStep::Finished => return Ok(None),
103 }
104 }
105 }
106
107 /// Reads the next item, assembling composites (head + members) into a
108 /// [`StreamItem::Composite`] and validating them.
109 ///
110 /// When the next descriptor on the wire is a composite head (`layout_tag = 0x0B`), this
111 /// reads its declared `member_count` members — each of which may itself be a composite
112 /// (nesting) — validates the group with
113 /// [`CompositeValidator`][hurray_core::CompositeValidator] (member count, plus
114 /// partition coverage / overlay ordering), and returns them together. Otherwise it
115 /// behaves like [`next_tensor`][StreamReader::next_tensor], returning a
116 /// [`StreamItem::Tensor`]. Returns `Ok(None)` on a clean EOF.
117 ///
118 /// # Examples
119 ///
120 /// ```no_run
121 /// # #[tokio::main]
122 /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
123 /// use hurray_io::stream::{StreamItem, StreamReader};
124 ///
125 /// let wire: &[u8] = &[]; // replace with actual data
126 /// let mut reader = StreamReader::new(wire);
127 /// while let Some(item) = reader.next_item().await? {
128 /// match item {
129 /// StreamItem::Tensor(t) => println!("tensor, {} buffer(s)", t.buffers.len()),
130 /// StreamItem::Composite(c) => println!("composite, {} member(s)", c.members.len()),
131 /// }
132 /// }
133 /// # Ok(())
134 /// # }
135 /// ```
136 ///
137 /// # Errors
138 ///
139 /// In addition to the errors of [`next_tensor`][StreamReader::next_tensor]:
140 ///
141 /// - [`Error::TornComposite`] — the stream ended before the head's `member_count`
142 /// members were read
143 /// - [`Error::CompositeNestingTooDeep`] — nesting exceeded `max_composite_depth`
144 /// - [`Error::Core`] — composite validation failed (e.g. partition does not cover the
145 /// index space, overlay ordering, member-count mismatch)
146 pub async fn next_item(&mut self) -> Result<Option<StreamItem>> {
147 loop {
148 match self.core.step()? {
149 ReadStep::Input { .. } => self.fill().await?,
150 ReadStep::Tensor => {
151 let tensor = self.take_tensor()?;
152 if let Some(item) = self.assembler.push(tensor)? {
153 return Ok(Some(item));
154 }
155 }
156 ReadStep::Finished => {
157 self.assembler.finish()?;
158 return Ok(None);
159 }
160 }
161 }
162 }
163
164 /// Answers one [`ReadStep::Input`]: read into the core's destination, report the count.
165 ///
166 /// The transfer and the report are one step with nothing suspended between them, which
167 /// is the shape ADR-041 § 6a requires of every driver.
168 async fn fill(&mut self) -> Result<()> {
169 let n = loop {
170 // Disjoint field borrows: the destination comes from `core`, the bytes from
171 // `inner`.
172 match self.inner.read(self.core.input_dst()).await {
173 Ok(n) => break n,
174 // `read_exact`, which this replaces, absorbed this internally.
175 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
176 // A read that errors moved no bytes, so the count is known to be zero and
177 // the core's position survives — no need to declare it lost (§ 6).
178 Err(e) => return Err(Error::Io(e)),
179 }
180 };
181 // `n == 0` is end of input: the core decides whether that is clean or torn.
182 self.core.report_filled(n)
183 }
184
185 /// Collects the tensor the core just announced.
186 fn take_tensor(&mut self) -> Result<StreamTensor> {
187 self.core.take_tensor().ok_or(Error::CoreFailed)
188 }
189
190 /// Returns the underlying reader.
191 pub fn into_inner(self) -> R {
192 self.inner
193 }
194}