hurray_io/stream/sans_io/write.rs
1//! The sequential write half of the sans-io stream core.
2//!
3//! Where the read core *owns* its destination, this one owns nothing: it hands the driver
4//! the caller's own payload slices, one at a time, and never assembles them (ADR-041 § 3).
5
6use hurray_core::{CompositeValidator, SyncMode, TensorDescriptor};
7
8use super::Unusable;
9use crate::stream::types::CompositeNode;
10use crate::{Error, Result};
11
12/// What the core has for its driver.
13///
14/// One request outstanding at a time, answered with
15/// [`report_written`][WriteOp::report_written] (ADR-041 § 4).
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum WriteStep {
18 /// Write [`output_src`][WriteOp::output_src] — `len` bytes — then report how many the
19 /// sink accepted. Fewer is normal; `0` is never progress.
20 Output {
21 /// Bytes remaining for the current slice. Never `0`.
22 len: usize,
23 },
24 /// Every byte of the submitted item is written. The borrowed payloads are released:
25 /// the core holds no pointer into them from here on (ADR-041 § 3).
26 ItemWritten,
27}
28
29/// Encodes the Hurray stream wire format into bytes a driver writes.
30///
31/// A plain tensor or a whole composite is submitted as one operation; the core emits the
32/// descriptor, then each payload in turn, and says when the item is written. Validation
33/// happens on submission, so a torn or invalid item never reaches a sink.
34///
35/// # Examples
36///
37/// ```
38/// use hurray_core::{
39/// BufferHandle, DeviceTag, ElementType, LayoutDescriptor, Shape, SyncMode,
40/// TensorDescriptor, MIN_BUFFER_ALIGNMENT,
41/// };
42/// use hurray_io::stream::{StreamWriterCore, WriteStep};
43///
44/// let handle =
45/// BufferHandle::new(8, MIN_BUFFER_ALIGNMENT, DeviceTag::Cpu, SyncMode::ProducerSynced)?;
46/// let descriptor = TensorDescriptor::new(
47/// 1, 0, ElementType::Float64, Shape::new(vec![1u64])?, 0,
48/// LayoutDescriptor::RowMajor, vec![handle], None, None, None, None,
49/// )?;
50/// let payload = [1u8; 8];
51/// // Bound, not a temporary: the operation borrows the payloads until the item is
52/// // written, which is ADR-041 § 3's lifetime rule showing up in the type system.
53/// let buffers: [&[u8]; 1] = [&payload];
54///
55/// let mut sink = Vec::new();
56/// let mut core = StreamWriterCore::new();
57/// let mut op = core.tensor(&descriptor, &buffers)?;
58/// loop {
59/// match op.step()? {
60/// WriteStep::Output { .. } => {
61/// // A sink that accepts 3 bytes at a time: a short write is normal.
62/// let src = op.output_src();
63/// let n = src.len().min(3);
64/// sink.extend_from_slice(&src[..n]);
65/// op.report_written(n)?;
66/// }
67/// WriteStep::ItemWritten => break,
68/// }
69/// }
70/// assert_eq!(sink.len(), descriptor.encode()?.len() + payload.len());
71/// # Ok::<(), Box<dyn std::error::Error>>(())
72/// ```
73pub struct StreamWriterCore {
74 enforce_cross_machine_sync: bool,
75 /// Set once the core may no longer be used; survives the [`WriteOp`] that set it.
76 unusable: Option<Unusable>,
77 /// Set when the item in flight has had a byte accepted, cleared when it completes
78 /// (§ 6b). Carried on the core, not the op, because a driver that returns an error and
79 /// drops the op reports nothing — so the poison cannot depend on being told.
80 torn: bool,
81}
82
83impl Default for StreamWriterCore {
84 fn default() -> Self {
85 Self::new()
86 }
87}
88
89impl StreamWriterCore {
90 /// Creates a core that accepts any valid [`SyncMode`].
91 ///
92 /// # Examples
93 ///
94 /// See the type-level example.
95 pub fn new() -> Self {
96 Self {
97 enforce_cross_machine_sync: false,
98 unusable: None,
99 torn: false,
100 }
101 }
102
103 /// Creates a core that rejects buffers whose `sync_mode` is not
104 /// [`SyncMode::ProducerSynced`], for a stream that crosses machine boundaries.
105 ///
106 /// # Examples
107 ///
108 /// ```
109 /// use hurray_io::stream::StreamWriterCore;
110 ///
111 /// let core = StreamWriterCore::cross_machine();
112 /// assert!(!core.is_unusable());
113 /// ```
114 pub fn cross_machine() -> Self {
115 Self {
116 enforce_cross_machine_sync: true,
117 unusable: None,
118 torn: false,
119 }
120 }
121
122 /// Submits one tensor: its descriptor followed by one payload slice per buffer.
123 ///
124 /// The payloads stay borrowed for the operation's lifetime and are released when it
125 /// reports [`WriteStep::ItemWritten`] — which is also the point a binding may drop the
126 /// reference pinning them (ADR-041 § 3).
127 ///
128 /// # Errors
129 ///
130 /// - [`Error::MultiBufferLengthMismatch`] — `buffers.len()` ≠ `descriptor.buffers.len()`
131 /// - [`Error::BufferSizeMismatch`] — a slice's length ≠ its handle's `byte_size`
132 /// - [`Error::InvalidCrossMachineSyncMode`] — cross-machine mode and a buffer that is
133 /// not `ProducerSynced`
134 /// - [`Error::Core`] — the descriptor does not encode
135 /// - [`Error::CoreUnusable`] — a previous write lost its byte count, so this core and
136 /// its sink are poisoned
137 ///
138 /// # Examples
139 ///
140 /// See the type-level example.
141 pub fn tensor<'a>(
142 &'a mut self,
143 descriptor: &'a TensorDescriptor,
144 buffers: &'a [&'a [u8]],
145 ) -> Result<WriteOp<'a>> {
146 if let Some(reason) = self.retired() {
147 return Err(Error::CoreUnusable(reason));
148 }
149 self.validate_tensor(descriptor, buffers)?;
150 let staging = descriptor.encode()?;
151 Ok(WriteOp {
152 core: self,
153 staging,
154 payloads: buffers,
155 state: OpState::Descriptor,
156 written: 0,
157 emitted_bytes: 0,
158 stack: Vec::new(),
159 })
160 }
161
162 /// Submits a composite: its head, then every member's descriptor and payloads, in
163 /// order (ADR-027 § Binding).
164 ///
165 /// The whole tree is validated before a byte is emitted, at every nesting level.
166 ///
167 /// # Errors
168 ///
169 /// The errors of [`tensor`][Self::tensor] for each member, plus [`Error::Core`] when
170 /// the head is not a valid composite head or the group fails validation (member count,
171 /// partition coverage, overlay ordering).
172 ///
173 /// # Examples
174 ///
175 /// ```
176 /// use hurray_core::{
177 /// layout::{CompositeLayout, CompositionRule},
178 /// BufferHandle, DeviceTag, ElementType, LayoutDescriptor, Shape, SyncMode,
179 /// TensorDescriptor, MIN_BUFFER_ALIGNMENT,
180 /// };
181 /// use hurray_io::stream::{CompositeNode, StreamWriterCore, WriteStep};
182 ///
183 /// let head = TensorDescriptor::new(
184 /// 1, 0, ElementType::Float32, Shape::new(vec![1u64])?, 0,
185 /// LayoutDescriptor::Composite(CompositeLayout::new(CompositionRule::Group, 1)?),
186 /// vec![], None, None, None, None,
187 /// )?;
188 /// let handle =
189 /// BufferHandle::new(4, MIN_BUFFER_ALIGNMENT, DeviceTag::Cpu, SyncMode::ProducerSynced)?;
190 /// let member = TensorDescriptor::new(
191 /// 1, 0, ElementType::Float32, Shape::new(vec![1u64])?, 0,
192 /// LayoutDescriptor::RowMajor, vec![handle], None, None, None, None,
193 /// )?;
194 /// let data = [7u8; 4];
195 /// let buffers: [&[u8]; 1] = [&data];
196 /// let members = [CompositeNode::Tensor { descriptor: &member, buffers: &buffers }];
197 ///
198 /// let mut core = StreamWriterCore::new();
199 /// let mut op = core.composite(&head, &members)?;
200 /// let mut emitted = 0usize;
201 /// while let WriteStep::Output { len } = op.step()? {
202 /// emitted += len;
203 /// op.report_written(len)?;
204 /// }
205 /// assert_eq!(emitted, head.encode()?.len() + member.encode()?.len() + data.len());
206 /// # Ok::<(), Box<dyn std::error::Error>>(())
207 /// ```
208 pub fn composite<'a>(
209 &'a mut self,
210 head: &'a TensorDescriptor,
211 members: &'a [CompositeNode<'a>],
212 ) -> Result<WriteOp<'a>> {
213 if let Some(reason) = self.retired() {
214 return Err(Error::CoreUnusable(reason));
215 }
216 self.validate_composite(head, members)?;
217 // The head owns no data buffers (composite-head invariant).
218 let staging = head.encode()?;
219 Ok(WriteOp {
220 core: self,
221 staging,
222 payloads: &[],
223 state: OpState::Descriptor,
224 written: 0,
225 emitted_bytes: 0,
226 stack: vec![Frame { members, next: 0 }],
227 })
228 }
229
230 /// Whether this core — and the sink behind it — has been poisoned.
231 ///
232 /// # Examples
233 ///
234 /// See [`WriteOp::mark_unusable`].
235 pub fn is_unusable(&self) -> bool {
236 self.retired().is_some()
237 }
238
239 /// The reason this core is retired, if it is.
240 ///
241 /// `unusable` wins over `torn`: a lost count is strictly worse than a known-incomplete
242 /// item. Read through here at every site so the precedence cannot drift.
243 fn retired(&self) -> Option<Unusable> {
244 self.unusable
245 .or(self.torn.then_some(Unusable::IncompleteOutput))
246 }
247
248 /// Checks one tensor's payload slices against its buffer table, writing nothing.
249 fn validate_tensor(&self, descriptor: &TensorDescriptor, buffers: &[&[u8]]) -> Result<()> {
250 if buffers.len() != descriptor.buffers.len() {
251 return Err(Error::MultiBufferLengthMismatch {
252 declared: descriptor.buffers.len(),
253 actual: buffers.len(),
254 });
255 }
256 for (i, (handle, buf)) in descriptor.buffers.iter().zip(buffers).enumerate() {
257 let declared = handle.byte_size();
258 let actual = buf.len() as u64;
259 if declared != actual {
260 return Err(Error::BufferSizeMismatch {
261 index: i,
262 declared,
263 actual,
264 });
265 }
266 if self.enforce_cross_machine_sync && handle.sync_mode() != SyncMode::ProducerSynced {
267 return Err(Error::InvalidCrossMachineSyncMode {
268 index: i,
269 actual: handle.sync_mode().to_byte(),
270 });
271 }
272 }
273 Ok(())
274 }
275
276 /// Validates a composite and every composite nested inside it, writing nothing.
277 fn validate_composite(
278 &self,
279 head: &TensorDescriptor,
280 members: &[CompositeNode<'_>],
281 ) -> Result<()> {
282 let mut validator = CompositeValidator::new(head)?;
283 for member in members {
284 validator.push_member(member.descriptor())?;
285 }
286 validator.finish()?;
287
288 for member in members {
289 match member {
290 CompositeNode::Tensor {
291 descriptor,
292 buffers,
293 } => self.validate_tensor(descriptor, buffers)?,
294 CompositeNode::Composite { head, members } => {
295 self.validate_composite(head, members)?
296 }
297 }
298 }
299 Ok(())
300 }
301}
302
303/// Composite members still to emit at one nesting level.
304struct Frame<'a> {
305 members: &'a [CompositeNode<'a>],
306 next: usize,
307}
308
309/// What part of the current node is being emitted.
310#[derive(Debug, Clone, Copy)]
311enum OpState {
312 /// Emitting the staged descriptor bytes.
313 Descriptor,
314 /// Emitting payload `index` of the current tensor.
315 Payload { index: usize },
316 /// Every byte of the item is out.
317 Done,
318 /// A previous call ended this operation.
319 Failed,
320}
321
322/// One item being written: a tensor or a whole composite.
323///
324/// Borrowing both the core and the caller's payloads for its lifetime is what enforces
325/// ADR-041 § 3's ownership rule in the type system — the core cannot hold a payload
326/// pointer after the operation ends, because the operation is what holds them.
327///
328/// # Examples
329///
330/// See [`StreamWriterCore`].
331pub struct WriteOp<'a> {
332 core: &'a mut StreamWriterCore,
333 /// Encoded descriptor bytes of the node being emitted.
334 staging: Vec<u8>,
335 /// Payload slices of the node being emitted, in buffer order.
336 payloads: &'a [&'a [u8]],
337 state: OpState,
338 /// Bytes of the current slice already written.
339 written: usize,
340 /// Bytes of *this item* already accepted by the sink — per item, not per slice, so
341 /// `begin_payload`/`next_node` must not reset it. The § 6b predicate.
342 emitted_bytes: u64,
343 /// Members awaiting emission, outermost first: ADR-041 § 10's explicit stack, in place
344 /// of the recursion `write_composite_unchecked` used.
345 stack: Vec<Frame<'a>>,
346}
347
348impl WriteOp<'_> {
349 /// Advances the operation and returns what it needs next. Never blocks.
350 ///
351 /// # Errors
352 ///
353 /// - [`Error::Core`] — a member descriptor does not encode
354 /// - [`Error::CoreFailed`] — a previous call ended this operation
355 /// - [`Error::CoreUnusable`] — the core was poisoned mid-item
356 ///
357 /// # Examples
358 ///
359 /// See [`StreamWriterCore`].
360 pub fn step(&mut self) -> Result<WriteStep> {
361 // `unusable`, not `retired()`: `torn` is set by *this* operation's first byte, and
362 // this operation is the one that can still clear it by finishing the item. `torn`
363 // only means damage once the op is gone and new work is submitted.
364 if let Some(reason) = self.core.unusable {
365 return Err(Error::CoreUnusable(reason));
366 }
367 loop {
368 match self.state {
369 OpState::Failed => return Err(Error::CoreFailed),
370 OpState::Done => {
371 // The item is whole, so the output is conforming again (§ 6b). Nothing
372 // to commit: the stream core holds no names and no index.
373 self.core.torn = false;
374 return Ok(WriteStep::ItemWritten);
375 }
376 OpState::Descriptor => {
377 if self.written < self.staging.len() {
378 return Ok(WriteStep::Output {
379 len: self.staging.len() - self.written,
380 });
381 }
382 self.begin_payload(0)?;
383 }
384 OpState::Payload { index } => {
385 let remaining = self.remaining();
386 if remaining > 0 {
387 return Ok(WriteStep::Output { len: remaining });
388 }
389 self.begin_payload(index + 1)?;
390 }
391 }
392 }
393 }
394
395 /// The bytes to write for the outstanding [`WriteStep::Output`] request.
396 ///
397 /// Either the staged descriptor or one of the caller's payload slices — never a copy
398 /// of them, and never a concatenation (ADR-041 § 3).
399 ///
400 /// # Examples
401 ///
402 /// See [`StreamWriterCore`].
403 pub fn output_src(&self) -> &[u8] {
404 match self.state {
405 OpState::Descriptor => &self.staging[self.written.min(self.staging.len())..],
406 OpState::Payload { index } => match self.payloads.get(index) {
407 Some(payload) => &payload[self.written.min(payload.len())..],
408 None => &[],
409 },
410 _ => &[],
411 }
412 }
413
414 /// Reports that the sink accepted `n` bytes of [`output_src`][Self::output_src].
415 ///
416 /// A short write re-offers the remainder.
417 ///
418 /// # Errors
419 ///
420 /// - [`Error::ZeroLengthWrite`] — `n == 0`, which § 4 forbids treating as progress. A
421 /// sink that cannot accept bytes is reported with
422 /// [`report_sink_closed`][Self::report_sink_closed] instead.
423 /// - [`Error::InvalidProgressReport`] — more bytes reported than were offered. An
424 /// unknown number of payload bytes is then on the wire and cannot be un-sent, so
425 /// this poisons the core and its sink (ADR-041 § 6a).
426 ///
427 /// # Examples
428 ///
429 /// See [`StreamWriterCore`].
430 pub fn report_written(&mut self, n: usize) -> Result<()> {
431 // See `step`: the in-flight item's own `torn` is not a bar to continuing it.
432 if let Some(reason) = self.core.unusable {
433 return Err(Error::CoreUnusable(reason));
434 }
435 if matches!(self.state, OpState::Failed) {
436 return Err(Error::CoreFailed);
437 }
438 if n == 0 {
439 // § 4: a zero-length write is not success. Position is still known, so the
440 // operation stays usable — the driver may retry or declare the sink closed.
441 return Err(Error::ZeroLengthWrite);
442 }
443 let offered = self.remaining();
444 if n > offered {
445 self.core.unusable = Some(Unusable::LostWriteCount);
446 self.state = OpState::Failed;
447 return Err(Error::InvalidProgressReport {
448 reported: n,
449 offered,
450 });
451 }
452 if self.emitted_bytes == 0 {
453 // Poison at the first accepted byte, not in `Drop`: `Drop` can be skipped by
454 // `forget`, delayed by a live op, or run while unwinding. The transient state
455 // is unobservable anyway — the op holds `&mut core` for its whole life.
456 self.core.torn = true;
457 }
458 self.emitted_bytes += n as u64;
459 self.written += n;
460 Ok(())
461 }
462
463 /// Bytes of this item already accepted by the sink, padding included.
464 ///
465 /// Zero means nothing of the item has left the driver, which is exactly when the item
466 /// may be resubmitted after a failure (§ 6b).
467 ///
468 /// # Examples
469 ///
470 /// ```
471 /// use hurray_core::{
472 /// BufferHandle, DeviceTag, ElementType, LayoutDescriptor, Shape, SyncMode,
473 /// TensorDescriptor, MIN_BUFFER_ALIGNMENT,
474 /// };
475 /// use hurray_io::stream::StreamWriterCore;
476 ///
477 /// let handle =
478 /// BufferHandle::new(8, MIN_BUFFER_ALIGNMENT, DeviceTag::Cpu, SyncMode::ProducerSynced)?;
479 /// let descriptor = TensorDescriptor::new(
480 /// 1, 0, ElementType::Float64, Shape::new(vec![1u64])?, 0,
481 /// LayoutDescriptor::RowMajor, vec![handle], None, None, None, None,
482 /// )?;
483 /// let payload = [1u8; 8];
484 /// let buffers: [&[u8]; 1] = [&payload];
485 ///
486 /// let mut core = StreamWriterCore::new();
487 /// let mut op = core.tensor(&descriptor, &buffers)?;
488 /// // Stepping offers bytes; only a report counts them.
489 /// op.step()?;
490 /// assert_eq!(op.emitted(), 0);
491 ///
492 /// op.report_written(3)?;
493 /// assert_eq!(op.emitted(), 3);
494 /// # Ok::<(), Box<dyn std::error::Error>>(())
495 /// ```
496 pub fn emitted(&self) -> u64 {
497 self.emitted_bytes
498 }
499
500 /// Declares the sink closed while this item is in flight.
501 ///
502 /// Returns the error to propagate. A partially written item is already on the wire and
503 /// a stream cannot be un-sent, so this poisons the core as well: a fresh writer over
504 /// the same connection is not a recovery (ADR-041 § 6a).
505 ///
506 /// # Examples
507 ///
508 /// See [`mark_unusable`][Self::mark_unusable].
509 pub fn report_sink_closed(&mut self) -> Error {
510 // Not `LostWriteCount`: the count is known exactly. What is lost is the ability to
511 // finish the item, which is a different claim (§ 6b).
512 self.core.unusable = Some(Unusable::IncompleteOutput);
513 self.state = OpState::Failed;
514 Error::SinkClosed
515 }
516
517 /// Marks the core unusable: a write moved bytes whose count was lost.
518 ///
519 /// The read side's equivalent damages only the core's own position; here the bytes are
520 /// already on the wire, so the sink is poisoned too (ADR-041 § 6a).
521 ///
522 /// # Examples
523 ///
524 /// ```
525 /// # use hurray_core::{
526 /// # BufferHandle, DeviceTag, ElementType, LayoutDescriptor, Shape, SyncMode,
527 /// # TensorDescriptor, MIN_BUFFER_ALIGNMENT,
528 /// # };
529 /// use hurray_io::stream::StreamWriterCore;
530 ///
531 /// # let handle =
532 /// # BufferHandle::new(4, MIN_BUFFER_ALIGNMENT, DeviceTag::Cpu, SyncMode::ProducerSynced)?;
533 /// # let descriptor = TensorDescriptor::new(
534 /// # 1, 0, ElementType::Float32, Shape::new(vec![1u64])?, 0,
535 /// # LayoutDescriptor::RowMajor, vec![handle], None, None, None, None,
536 /// # )?;
537 /// # let payload = [0u8; 4];
538 /// # let buffers: [&[u8]; 1] = [&payload];
539 /// let mut core = StreamWriterCore::new();
540 /// let mut op = core.tensor(&descriptor, &buffers)?;
541 /// op.mark_unusable();
542 /// drop(op);
543 ///
544 /// // The poison outlives the operation: this sink is not reusable.
545 /// assert!(core.is_unusable());
546 /// assert!(core.tensor(&descriptor, &buffers).is_err());
547 /// # Ok::<(), Box<dyn std::error::Error>>(())
548 /// ```
549 pub fn mark_unusable(&mut self) {
550 self.core.unusable = Some(Unusable::LostWriteCount);
551 self.state = OpState::Failed;
552 }
553
554 /// Bytes still to write for the current slice.
555 fn remaining(&self) -> usize {
556 match self.state {
557 OpState::Descriptor => self.staging.len().saturating_sub(self.written),
558 OpState::Payload { index } => match self.payloads.get(index) {
559 Some(payload) => payload.len().saturating_sub(self.written),
560 None => 0,
561 },
562 _ => 0,
563 }
564 }
565
566 /// Moves to payload `index` of the current node, or past it to the next node.
567 fn begin_payload(&mut self, index: usize) -> Result<()> {
568 self.written = 0;
569 if index < self.payloads.len() {
570 self.state = OpState::Payload { index };
571 return Ok(());
572 }
573 self.next_node()
574 }
575
576 /// Takes the next composite member off the stack, or finishes the item.
577 fn next_node(&mut self) -> Result<()> {
578 loop {
579 let Some(frame) = self.stack.last_mut() else {
580 self.state = OpState::Done;
581 self.payloads = &[];
582 return Ok(());
583 };
584 let Some(node) = frame.members.get(frame.next) else {
585 self.stack.pop();
586 continue;
587 };
588 frame.next += 1;
589
590 match node {
591 CompositeNode::Tensor {
592 descriptor,
593 buffers,
594 } => {
595 self.staging = match descriptor.encode() {
596 Ok(staging) => staging,
597 Err(e) => {
598 self.state = OpState::Failed;
599 return Err(e.into());
600 }
601 };
602 self.payloads = buffers;
603 }
604 CompositeNode::Composite { head, members } => {
605 self.staging = match head.encode() {
606 Ok(staging) => staging,
607 Err(e) => {
608 self.state = OpState::Failed;
609 return Err(e.into());
610 }
611 };
612 self.payloads = &[];
613 self.stack.push(Frame { members, next: 0 });
614 }
615 }
616 self.written = 0;
617 self.state = OpState::Descriptor;
618 return Ok(());
619 }
620 }
621}