hurray_ffi/stream/reader.rs
1//! Blocking stream reads over a caller-supplied transport (ADR-041 §§ 7, 8).
2//!
3//! [`HurrayReader`] is a driver, not a parser: the decoding is `hurray-io`'s sans-io
4//! `StreamReaderCore`, and this module only moves bytes between the caller's transport and
5//! the destination the core hands it. Which is why there is no runtime here — the core's
6//! progress API never blocks, so the blocking happens entirely inside the caller's own
7//! read callback.
8
9use hurray_io::stream::{CompositeAssembler, ReadStep, StreamReaderCore, StreamReaderOptions};
10
11use crate::{
12 item::{item_handle, HurrayItem},
13 panic::{catch, null_check},
14 status::{
15 status_from_io_error, HurrayStatus, HURRAY_END_OF_STREAM, HURRAY_ERR_INTERNAL,
16 HURRAY_ERR_NULL_POINTER, HURRAY_OK,
17 },
18 stream::{transport::Transport, HurrayTransport},
19 tensor::{tensor_handle, HurrayTensor},
20};
21
22// ── Options ───────────────────────────────────────────────────────────────────
23
24/// Resource limits and transport assumptions for a reader.
25///
26/// Pass null to [`hurray_reader_new`] for the defaults. The fields map one-to-one onto
27/// `hurray-io`'s `StreamReaderOptions`; a reader over an untrusted socket SHOULD set both
28/// limits, because a descriptor's declared sizes are attacker-controlled until they are
29/// bounded.
30#[repr(C)]
31pub struct HurrayReaderOptions {
32 /// Maximum byte length of a single descriptor. Default: 16 MiB.
33 pub max_descriptor_bytes: u64,
34 /// Maximum byte length of a single buffer. Default: unbounded.
35 pub max_buffer_bytes: u64,
36 /// Reject buffers whose sync mode is not `ProducerSynced`. Set this when the stream
37 /// crosses a machine boundary, where a device-local synchronization handle means
38 /// nothing.
39 pub enforce_cross_machine_sync: bool,
40 /// Maximum composite nesting depth accepted by [`hurray_reader_next_item`].
41 /// Default: 64. Pass `0` for that default.
42 ///
43 /// Bounds the assembler, not the byte machine: the core reassembles nothing, so
44 /// [`hurray_reader_next_tensor`] is unaffected by this field.
45 pub max_composite_depth: usize,
46}
47
48impl HurrayReaderOptions {
49 /// Translates the C struct into `hurray-io`'s options.
50 fn to_io(&self) -> StreamReaderOptions {
51 let defaults = StreamReaderOptions::default();
52 StreamReaderOptions {
53 max_descriptor_bytes: self.max_descriptor_bytes,
54 max_buffer_bytes: self.max_buffer_bytes,
55 enforce_cross_machine_sync: self.enforce_cross_machine_sync,
56 // Zero is the request for the default rather than a limit of zero: a depth of
57 // zero would refuse every composite, which no caller means by `{0}`.
58 max_composite_depth: if self.max_composite_depth == 0 {
59 defaults.max_composite_depth
60 } else {
61 self.max_composite_depth
62 },
63 }
64 }
65}
66
67// ── HurrayReader ──────────────────────────────────────────────────────────────
68
69/// Opaque handle to a blocking stream reader.
70///
71/// Not `#[repr(C)]`; the layout is an implementation detail. Created by
72/// [`hurray_reader_new`] and released exactly once by [`hurray_reader_destroy`].
73///
74/// One handle, one thread at a time.
75pub struct HurrayReader {
76 core: StreamReaderCore,
77 transport: Transport,
78 /// Held across calls, not created per call: a composite spans several tensors, so a
79 /// half-assembled one has to survive between them.
80 assembler: CompositeAssembler,
81}
82
83// ── Construction / destruction ────────────────────────────────────────────────
84
85/// Creates a reader over `transport`.
86///
87/// The vtable is copied, so the caller MAY free the [`HurrayTransport`] struct immediately;
88/// its `ctx` must stay valid until the reader is destroyed, at which point the vtable's
89/// release callback is invoked.
90///
91/// # Arguments
92///
93/// - `transport` — Non-null; its `read` callback MUST be non-null.
94/// - `options` — MAY be null for [`HurrayReaderOptions`]' documented defaults.
95/// - `out_handle` — Non-null pointer to a `*mut HurrayReader`; set on success.
96///
97/// # Safety
98///
99/// - `transport` MUST point to a readable [`HurrayTransport`] for the duration of the call.
100/// - `options`, if non-null, MUST point to a readable [`HurrayReaderOptions`].
101/// - `out_handle` MUST be a valid, non-null, writable pointer.
102#[no_mangle]
103pub unsafe extern "C" fn hurray_reader_new(
104 transport: *const HurrayTransport,
105 options: *const HurrayReaderOptions,
106 out_handle: *mut *mut HurrayReader,
107) -> HurrayStatus {
108 catch(|| {
109 // options may be null: that is the request for defaults.
110 null_check!(transport, out_handle);
111
112 // SAFETY: caller guarantees transport is readable for this call.
113 let transport = Transport::from_c(transport);
114 if !transport.can_read() {
115 return HURRAY_ERR_NULL_POINTER;
116 }
117
118 let io_options = if options.is_null() {
119 StreamReaderOptions::default()
120 } else {
121 // SAFETY: non-null and caller-guaranteed readable.
122 (*options).to_io()
123 };
124
125 // The core ignores `max_composite_depth` — it bounds the assembler (ADR-041 § 10).
126 let max_composite_depth = io_options.max_composite_depth;
127 let reader = HurrayReader {
128 core: StreamReaderCore::with_options(io_options),
129 transport,
130 assembler: CompositeAssembler::new(max_composite_depth),
131 };
132 // SAFETY: Box::into_raw transfers ownership to the caller, who must call
133 // hurray_reader_destroy exactly once.
134 *out_handle = Box::into_raw(Box::new(reader));
135 HURRAY_OK
136 })
137}
138
139/// Destroys a reader, releases its transport context, and writes null through `handle`.
140///
141/// Tensors already produced by [`hurray_reader_next_tensor`] are unaffected: each owns its
142/// own buffers and outlives the reader.
143///
144/// # Safety
145///
146/// - `handle` MUST be a valid, non-null pointer to a `*mut HurrayReader` from
147/// [`hurray_reader_new`].
148/// - That reader MUST NOT have been destroyed already.
149#[no_mangle]
150pub unsafe extern "C" fn hurray_reader_destroy(handle: *mut *mut HurrayReader) -> HurrayStatus {
151 catch(|| {
152 null_check!(handle);
153 let reader = *handle;
154 if reader.is_null() {
155 return HURRAY_ERR_INTERNAL;
156 }
157 // SAFETY: created by hurray_reader_new via Box::into_raw; the caller guarantees
158 // this is the first and only destroy call. Dropping it runs the transport's
159 // release callback.
160 drop(Box::from_raw(reader));
161 *handle = std::ptr::null_mut();
162 HURRAY_OK
163 })
164}
165
166// ── Reading ───────────────────────────────────────────────────────────────────
167
168/// Reads the next tensor, blocking in the transport's read callback as needed.
169///
170/// Returns:
171///
172/// - [`HURRAY_OK`] — `*out_tensor` is a new [`HurrayTensor`], to be released with
173/// [`hurray_tensor_destroy`][crate::hurray_tensor_destroy].
174/// - [`HURRAY_END_OF_STREAM`] — input ended cleanly at a tensor boundary. `*out_tensor` is
175/// untouched. Positive, because this is how a stream ends, not a failure.
176/// - [`HURRAY_ERR_UNEXPECTED_EOF`][crate::HURRAY_ERR_UNEXPECTED_EOF] — input ended inside
177/// a descriptor or a buffer.
178/// - Whatever status the read callback returned, if it failed. The reader's position is
179/// intact in that case — the transport contract says a failing callback moved no bytes —
180/// so **the same call MAY be retried** and will resume mid-tensor rather than restart it.
181///
182/// A composite head arrives as the buffer-less tensor it is on the wire, and its members
183/// as the tensors that follow it. Use [`hurray_reader_next_item`] to have them assembled
184/// into one walkable item instead.
185///
186/// Every buffer in the returned tensor points at memory the read core allocated at the
187/// alignment its buffer handle *declares*, and the transport wrote into it directly — so a
188/// SIMD load on it is sound, and no copy of the payload happened after that write.
189///
190/// # Safety
191///
192/// - `handle` MUST point to a live [`HurrayReader`].
193/// - `out_tensor` MUST be a valid, non-null, writable pointer.
194#[no_mangle]
195pub unsafe extern "C" fn hurray_reader_next_tensor(
196 handle: *mut HurrayReader,
197 out_tensor: *mut *mut HurrayTensor,
198) -> HurrayStatus {
199 catch(|| {
200 null_check!(handle, out_tensor);
201 // SAFETY: handle is non-null and points to a live HurrayReader, which the C ABI's
202 // one-handle-one-thread rule says nothing else is touching.
203 let reader = &mut *handle;
204
205 loop {
206 match reader.core.step() {
207 Err(e) => return status_from_io_error(&e),
208 Ok(ReadStep::Finished) => return HURRAY_END_OF_STREAM,
209 Ok(ReadStep::Tensor) => {
210 let Some(tensor) = reader.core.take_tensor() else {
211 // step() said Tensor; the core cannot then have none.
212 return HURRAY_ERR_INTERNAL;
213 };
214 *out_tensor = Box::into_raw(tensor_handle(tensor.descriptor, tensor.buffers));
215 return HURRAY_OK;
216 }
217 Ok(ReadStep::Input { .. }) => {
218 // Disjoint fields: the core lends out its destination while the
219 // transport fills it.
220 let dst = reader.core.input_dst();
221 let filled = match reader.transport.read(dst) {
222 Ok(n) => n,
223 // No bytes moved, so no report is owed and the position still
224 // matches the transport's: § 6's LostReadCount is unreachable from
225 // here, and the call is retryable.
226 Err(status) => return status,
227 };
228 if let Err(e) = reader.core.report_filled(filled) {
229 return status_from_io_error(&e);
230 }
231 }
232 }
233 }
234 })
235}
236
237/// Reads the next *item*: one tensor, or a composite and everything under it.
238///
239/// The assembled view of the same bytes [`hurray_reader_next_tensor`] reads flat. A
240/// composite head and its members arrive as one [`HurrayItem`] whose nodes a caller walks
241/// in preorder — see [`crate::item`] for the shape, and `hurray_item_node` for the walk.
242///
243/// Returns:
244///
245/// - [`HURRAY_OK`] — `*out_item` is a new [`HurrayItem`], to be released with
246/// [`hurray_item_destroy`][crate::hurray_item_destroy].
247/// - [`HURRAY_END_OF_STREAM`] — input ended cleanly at an item boundary. `*out_item` is
248/// untouched.
249/// - [`HURRAY_ERR_INVALID_STREAM`][crate::HURRAY_ERR_INVALID_STREAM] — the stream ended
250/// with a composite head still owed members, or a group failed validation (member
251/// count, partition coverage, overlay ordering).
252/// - [`HURRAY_ERR_LIMIT_EXCEEDED`][crate::HURRAY_ERR_LIMIT_EXCEEDED] — nesting exceeded
253/// `max_composite_depth`.
254/// - Whatever status the read callback returned, if it failed. The reader's position is
255/// intact, so the same call MAY be retried and resumes mid-tensor.
256///
257/// Both views come from one core, and the assembly happens on this side of it. So mixing
258/// the two calls on one reader is allowed but rarely meant: a
259/// [`hurray_reader_next_tensor`] issued while a composite head is open takes the member
260/// that head was waiting for, and the group stays open for the rest of the stream.
261///
262/// # Safety
263///
264/// - `handle` MUST point to a live [`HurrayReader`].
265/// - `out_item` MUST be a valid, non-null, writable pointer.
266#[no_mangle]
267pub unsafe extern "C" fn hurray_reader_next_item(
268 handle: *mut HurrayReader,
269 out_item: *mut *mut HurrayItem,
270) -> HurrayStatus {
271 catch(|| {
272 null_check!(handle, out_item);
273 // SAFETY: handle is non-null and points to a live HurrayReader, which the C ABI's
274 // one-handle-one-thread rule says nothing else is touching.
275 let reader = &mut *handle;
276
277 loop {
278 match reader.core.step() {
279 Err(e) => return status_from_io_error(&e),
280 Ok(ReadStep::Finished) => {
281 // A head still owed members at end of input is a torn composite, not a
282 // clean end of stream.
283 if let Err(e) = reader.assembler.finish() {
284 return status_from_io_error(&e);
285 }
286 return HURRAY_END_OF_STREAM;
287 }
288 Ok(ReadStep::Tensor) => {
289 let Some(tensor) = reader.core.take_tensor() else {
290 // step() said Tensor; the core cannot then have none.
291 return HURRAY_ERR_INTERNAL;
292 };
293 match reader.assembler.push(tensor) {
294 Err(e) => return status_from_io_error(&e),
295 // Still filling a composite: nothing to hand back yet.
296 Ok(None) => continue,
297 Ok(Some(item)) => {
298 *out_item = Box::into_raw(item_handle(item));
299 return HURRAY_OK;
300 }
301 }
302 }
303 Ok(ReadStep::Input { .. }) => {
304 // Disjoint fields: the core lends out its destination while the
305 // transport fills it.
306 let dst = reader.core.input_dst();
307 let filled = match reader.transport.read(dst) {
308 Ok(n) => n,
309 // No bytes moved, so no report is owed and the position still
310 // matches the transport's: the call is retryable.
311 Err(status) => return status,
312 };
313 if let Err(e) = reader.core.report_filled(filled) {
314 return status_from_io_error(&e);
315 }
316 }
317 }
318 }
319 })
320}