cobs_io_async/
connection.rs

1use core::mem;
2use core::task::{Context, Poll, Waker};
3
4use crate::error::{ConfigError, ConnectionError};
5use crate::protocol::{self, ACK_FRAME_LEN, Packet};
6use crate::sync;
7use crate::timer::Timer;
8use crate::transport::Transport;
9
10pub(crate) const STEP_BUDGET: usize = 16;
11
12#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
13#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
14#[cfg_attr(feature = "defmt", derive(defmt::Format))]
15/// Parameters of a [`Reliable`] connection.
16///
17/// Create one with [`Config::new`] and override fields as needed.
18pub struct Config {
19    /// Identifier of the current session.
20    ///
21    /// Both peers must use the same value, and it should be fresh for each
22    /// session. Packets from a different session are ignored. This is not
23    /// authentication.
24    pub session: u64,
25
26    /// Time to wait for an acknowledgment before retransmitting, in
27    /// [`Timer`] ticks.
28    ///
29    /// The wait begins once the packet has been written and the underlying
30    /// transport flushed.
31    pub retransmit_timeout: u64,
32
33    /// Number of retransmissions attempted after the initial transmission.
34    ///
35    /// When the timeout elapses with none left, the connection fails with
36    /// [`ConnectionError::Timeout`]. Zero fails at the first timeout.
37    pub max_retries: u32,
38}
39
40impl Config {
41    /// Returns a configuration for `session` with a retransmission timeout
42    /// of 1000 ticks and 8 retries.
43    pub const fn new(session: u64) -> Self {
44        Self {
45            session,
46            retransmit_timeout: 1_000,
47            max_retries: 8,
48        }
49    }
50}
51
52#[derive(Clone, Copy, PartialEq, Eq)]
53enum Tx {
54    Idle,
55    Queued,
56    Sending,
57    Waiting,
58}
59
60#[derive(Clone, Copy, PartialEq, Eq)]
61enum Frame {
62    Ack,
63    Data,
64}
65
66#[derive(Clone, Copy, PartialEq, Eq)]
67enum Out {
68    Idle,
69    Writing(Frame, usize),
70    Flushing(Frame),
71}
72
73#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
74enum Work {
75    Blocked,
76    Progress,
77    Yield,
78}
79
80struct Budget(usize);
81
82impl Budget {
83    const fn new() -> Self {
84        Self(STEP_BUDGET)
85    }
86
87    fn spend(&mut self) -> bool {
88        match self.0.checked_sub(1) {
89            Some(left) => {
90                self.0 = left;
91                true
92            }
93            None => false,
94        }
95    }
96}
97
98#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
99/// An operation of a [`Reliable`] connection that can be pending.
100///
101/// Identifies the operation abandoned in
102/// [`Reliable::cancel_pending`]. The connection stores one waiter per
103/// operation.
104pub enum PendingOperation {
105    /// [`Transport::poll_read`].
106    Read,
107    /// [`Transport::poll_write`].
108    Write,
109    /// [`Transport::poll_flush`].
110    Flush,
111    /// [`Reliable::poll_progress`].
112    Progress,
113}
114
115#[derive(Clone, Copy, PartialEq, Eq)]
116enum Op {
117    Read,
118    Write,
119    Flush,
120    Progress,
121}
122
123const OPS: [Op; 4] = [Op::Read, Op::Write, Op::Flush, Op::Progress];
124
125#[derive(Default)]
126struct Waiters([Option<Waker>; 4]);
127
128impl Waiters {
129    fn slot(&mut self, op: Op) -> &mut Option<Waker> {
130        &mut self.0[op as usize]
131    }
132
133    fn register(&mut self, op: Op, waker: &Waker) {
134        let slot = self.slot(op);
135        match slot {
136            Some(held) if held.will_wake(waker) => {}
137            _ => *slot = Some(waker.clone()),
138        }
139    }
140
141    fn clear(&mut self, op: Op) {
142        *self.slot(op) = None;
143    }
144
145    fn wake(&mut self, op: Op) {
146        if let Some(waker) = self.slot(op).take() {
147            waker.wake();
148        }
149    }
150}
151
152enum Health<E> {
153    Ok,
154    Failed(ConnectionError<E>),
155    Reported,
156}
157
158/// A reliable, ordered byte stream over an unreliable [`Transport`].
159///
160/// Wraps the link `T` and the clock `C`, and implements [`Transport`] itself.
161/// Application bytes are divided into packets, framed with COBS, protected by
162/// a checksum, acknowledged and retransmitted. See the [crate
163/// documentation](crate) for the protocol, semantics and limitations.
164///
165/// # Capacity
166///
167/// `N` is the size in bytes of each internal frame buffer, which bounds a
168/// complete encoded packet including both delimiters. The largest payload per
169/// packet is [`MAX_PAYLOAD`](Self::MAX_PAYLOAD), which accounts for the
170/// 17 bytes of header and checksum and for COBS overhead; with the default
171/// `N` of 256 it is 236. [`new`](Self::new) rejects capacities too small for
172/// one byte of payload (any `N` below 21). The connection holds four buffers
173/// of `N` bytes, plus a small ACK frame buffer.
174///
175/// # Driving the connection
176///
177/// The connection makes progress only while polled. Each poll method drives
178/// both directions: it finishes partially written frames, sends pending
179/// acknowledgments between frames, processes input, checks the retransmission
180/// deadline and starts retransmissions. Call
181/// [`poll_progress`](Self::poll_progress) when no read, write or flush is
182/// outstanding.
183///
184/// # Cancellation
185///
186/// Dropping a future that awaits a poll method rolls back no protocol state
187/// and corrupts no accepted data. It does not, however, notify the
188/// connection. When tasks poll the operations separately, the transport and
189/// timer wake only the task that polled them last, so dropping that future
190/// can leave the others without a wakeup. An adapter that supports canceling
191/// pending operations must call
192/// [`cancel_pending`](Self::cancel_pending) for the dropped operation. The
193/// connection does not detect dropped futures.
194///
195/// # Failure
196///
197/// After a transport error, EOF, retry exhaustion or sequence exhaustion, the
198/// connection is permanently failed. The first poll to observe the failure
199/// returns the specific [`ConnectionError`]; later polls return
200/// [`ConnectionError::Failed`]. Payload already received can still be read
201/// first. Use [`into_parts`](Self::into_parts) to recover the transport and
202/// timer.
203pub struct Reliable<T: Transport, C: Timer, const N: usize = 256> {
204    transport: T,
205    timer: C,
206    config: Config,
207    health: Health<T::Error>,
208    tx: Tx,
209    tx_frame: [u8; N],
210    tx_len: usize,
211    tx_next: u64,
212    tx_deadline: u64,
213    tx_retries: u32,
214    final_input: Option<usize>,
215    out: Out,
216    ack_pending: Option<u32>,
217    ack_frame: [u8; ACK_FRAME_LEN],
218    ack_len: usize,
219    rx_buf: [u8; N],
220    rx_len: usize,
221    rx_discarding: bool,
222    rx_payload: [u8; N],
223    rx_start: usize,
224    rx_end: usize,
225    rx_next: u64,
226    scratch: [u8; N],
227    waiters: Waiters,
228}
229
230impl<T: Transport, C: Timer, const N: usize> Reliable<T, C, N> {
231    /// The largest number of application bytes carried by one packet.
232    ///
233    /// A single `poll_write` accepts at most this many bytes.
234    pub const MAX_PAYLOAD: usize = protocol::max_payload(N);
235
236    /// Bytes read after a deadline expires while looking for a late
237    /// acknowledgment: the tail of a partially received frame plus one ACK.
238    const FINAL_INPUT: usize = N + ACK_FRAME_LEN;
239
240    /// Creates a connection over `transport`, using `timer` for deadlines.
241    ///
242    /// # Errors
243    ///
244    /// Returns [`ConfigError::BufferTooSmall`] if `N` cannot hold a packet
245    /// carrying at least one payload byte.
246    pub fn new(transport: T, timer: C, config: Config) -> Result<Self, ConfigError> {
247        if Self::MAX_PAYLOAD == 0 {
248            return Err(ConfigError::BufferTooSmall);
249        }
250        Ok(Self {
251            transport,
252            timer,
253            config,
254            health: Health::Ok,
255            tx: Tx::Idle,
256            tx_frame: [0; N],
257            tx_len: 0,
258            tx_next: 0,
259            tx_deadline: 0,
260            tx_retries: 0,
261            final_input: None,
262            out: Out::Idle,
263            ack_pending: None,
264            ack_frame: [0; ACK_FRAME_LEN],
265            ack_len: 0,
266            rx_buf: [0; N],
267            rx_len: 0,
268            rx_discarding: false,
269            rx_payload: [0; N],
270            rx_start: 0,
271            rx_end: 0,
272            rx_next: 0,
273            scratch: [0; N],
274            waiters: Waiters::default(),
275        })
276    }
277
278    /// Consumes the connection and returns the transport and timer.
279    ///
280    /// Buffered data and protocol state are discarded.
281    pub fn into_parts(self) -> (T, C) {
282        (self.transport, self.timer)
283    }
284
285    /// Drives the connection without reading, writing or flushing.
286    ///
287    /// Use it to keep acknowledgments, retransmissions and incoming data
288    /// flowing while the application has no operation pending. Incoming
289    /// payload is retained for later reads, and further data is not
290    /// acknowledged until the application reads it.
291    ///
292    /// Returns `Pending` for as long as the connection is healthy, with
293    /// wakeups registered. Returns `Ready` with the error once the
294    /// connection has failed.
295    pub fn poll_progress(&mut self, cx: &mut Context<'_>) -> Poll<ConnectionError<T::Error>> {
296        self.drive(cx, &mut Budget::new());
297        let poll = match self.failure() {
298            Some(error) => Poll::Ready(error),
299            None => Poll::Pending,
300        };
301        self.finish(Op::Progress, cx, poll)
302    }
303
304    /// Abandons a pending `operation` and hands its wakeups to the others.
305    ///
306    /// The underlying [`Transport`] and [`Timer`] remember only the waker of
307    /// their latest poller. If that poller is dropped while pending, the
308    /// registration is lost and the connection cannot detect it, so the
309    /// surviving operations may never be woken. This method clears the stored
310    /// waiter of `operation` and wakes and removes every other stored waiter.
311    /// Each woken task polls again and re-registers with the transport and
312    /// timer.
313    ///
314    /// Protocol state is unchanged: an accepted packet, a partially written
315    /// frame, retained incoming payload, pending acknowledgments and
316    /// retransmission deadlines and retry counts all survive. Nothing is
317    /// sent, and no task is woken if no waiter is stored. Waiters are woken
318    /// once per call, so calling it repeatedly does not wake a task again
319    /// until it has polled and registered again.
320    ///
321    /// # Serialization
322    ///
323    /// Call it with exclusive access to the connection, after the last poll
324    /// of the abandoned operation and before the operation is forgotten. The
325    /// connection stores one waiter per operation, and does not distinguish
326    /// between futures awaiting the same one. Canceling an older future of an
327    /// operation while a newer one is pending clears the newer one's waiter,
328    /// so the adapter must prevent that or track the futures itself.
329    pub fn cancel_pending(&mut self, operation: PendingOperation) {
330        let canceled = match operation {
331            PendingOperation::Read => Op::Read,
332            PendingOperation::Write => Op::Write,
333            PendingOperation::Flush => Op::Flush,
334            PendingOperation::Progress => Op::Progress,
335        };
336        self.waiters.clear(canceled);
337        for op in OPS {
338            if op != canceled {
339                self.waiters.wake(op);
340            }
341        }
342    }
343
344    #[cfg(test)]
345    pub(crate) fn set_sequences(&mut self, tx_next: u64, rx_next: u64) {
346        self.tx_next = tx_next;
347        self.rx_next = rx_next;
348    }
349
350    fn fail(&mut self, error: ConnectionError<T::Error>) {
351        self.final_input = None;
352        if matches!(self.health, Health::Ok) {
353            self.health = Health::Failed(error);
354        }
355    }
356
357    fn failure(&mut self) -> Option<ConnectionError<T::Error>> {
358        match mem::replace(&mut self.health, Health::Reported) {
359            Health::Ok => {
360                self.health = Health::Ok;
361                None
362            }
363            Health::Failed(error) => Some(error),
364            Health::Reported => Some(ConnectionError::Failed),
365        }
366    }
367
368    fn data_out_busy(&self) -> bool {
369        matches!(
370            self.out,
371            Out::Writing(Frame::Data, _) | Out::Flushing(Frame::Data)
372        )
373    }
374
375    fn queue_packet(
376        &mut self,
377        cx: &mut Context<'_>,
378        buf: &[u8],
379    ) -> Poll<Result<usize, ConnectionError<T::Error>>> {
380        let mut budget = Budget::new();
381        let yielded = self.drive(cx, &mut budget) == Work::Yield;
382        let busy = !self.write_ready();
383        if !busy && self.tx_next > u64::from(u32::MAX) {
384            self.fail(ConnectionError::SequenceExhausted);
385        }
386        if let Some(error) = self.failure() {
387            return Poll::Ready(Err(error));
388        }
389        if busy {
390            return Poll::Pending;
391        }
392        let len = buf.len().min(Self::MAX_PAYLOAD);
393        self.tx_len = protocol::encode_data(
394            self.config.session,
395            self.tx_next as u32,
396            &buf[..len],
397            &mut self.scratch,
398            &mut self.tx_frame,
399        );
400        self.tx_next += 1;
401        self.tx_retries = 0;
402        self.tx = Tx::Queued;
403        if !yielded {
404            self.drive(cx, &mut budget);
405        }
406        Poll::Ready(Ok(len))
407    }
408
409    fn write_ready(&self) -> bool {
410        self.tx == Tx::Idle && !self.data_out_busy()
411    }
412
413    fn flush_ready(&self) -> bool {
414        self.tx == Tx::Idle && self.out == Out::Idle && self.ack_pending.is_none()
415    }
416
417    fn finish<R>(&mut self, op: Op, cx: &mut Context<'_>, poll: Poll<R>) -> Poll<R> {
418        let completed = poll.is_ready();
419        if completed {
420            self.waiters.clear(op);
421        } else {
422            self.waiters.register(op, cx.waker());
423        }
424        self.notify(op, completed);
425        poll
426    }
427
428    /// Wakes the retained waiters of the other operations whose condition now
429    /// holds. A completed operation, or a failure, wakes all of them: the
430    /// transport and timer remember only their latest poller, which may be a
431    /// task that is no longer waiting.
432    fn notify(&mut self, done: Op, completed: bool) {
433        let everyone = completed || !matches!(self.health, Health::Ok);
434        for op in OPS {
435            if op == done {
436                continue;
437            }
438            let due = everyone
439                || match op {
440                    Op::Read => self.rx_start < self.rx_end,
441                    Op::Write => self.write_ready(),
442                    Op::Flush => self.flush_ready(),
443                    Op::Progress => false,
444                };
445            if due {
446                self.waiters.wake(op);
447            }
448        }
449    }
450
451    fn drive(&mut self, cx: &mut Context<'_>, budget: &mut Budget) -> Work {
452        if !matches!(self.health, Health::Ok) {
453            return Work::Blocked;
454        }
455        loop {
456            match self.step(cx, budget) {
457                Ok(Work::Progress) => {}
458                Ok(Work::Blocked) => return Work::Blocked,
459                Ok(Work::Yield) => {
460                    cx.waker().wake_by_ref();
461                    return Work::Yield;
462                }
463                Err(error) => {
464                    self.fail(error);
465                    return Work::Blocked;
466                }
467            }
468        }
469    }
470
471    fn step(
472        &mut self,
473        cx: &mut Context<'_>,
474        budget: &mut Budget,
475    ) -> Result<Work, ConnectionError<T::Error>> {
476        let output = self.poll_output(cx, budget)?;
477        let input = self.poll_input(cx, budget)?;
478        let timer = self.poll_timer(cx, input)?;
479        Ok(output.max(input).max(timer))
480    }
481
482    fn poll_output(
483        &mut self,
484        cx: &mut Context<'_>,
485        budget: &mut Budget,
486    ) -> Result<Work, ConnectionError<T::Error>> {
487        if self.out == Out::Idle {
488            if let Some(seq) = self.ack_pending.take() {
489                self.ack_len = protocol::encode_ack(self.config.session, seq, &mut self.ack_frame);
490                self.out = Out::Writing(Frame::Ack, 0);
491            } else if self.tx == Tx::Queued {
492                self.tx = Tx::Sending;
493                self.out = Out::Writing(Frame::Data, 0);
494            } else {
495                return Ok(Work::Blocked);
496            }
497        }
498        if !budget.spend() {
499            return Ok(Work::Yield);
500        }
501        match self.out {
502            Out::Idle => Ok(Work::Blocked),
503            Out::Writing(frame, offset) => {
504                let (pending, len) = match frame {
505                    Frame::Ack => (&self.ack_frame[offset..self.ack_len], self.ack_len),
506                    Frame::Data => (&self.tx_frame[offset..self.tx_len], self.tx_len),
507                };
508                let written = match self.transport.poll_write(cx, pending) {
509                    Poll::Pending => return Ok(Work::Blocked),
510                    Poll::Ready(Err(error)) => return Err(ConnectionError::Transport(error)),
511                    Poll::Ready(Ok(0)) => return Err(ConnectionError::WriteZero),
512                    Poll::Ready(Ok(written)) => written.min(pending.len()),
513                };
514                let offset = offset + written;
515                self.out = if offset == len {
516                    Out::Flushing(frame)
517                } else {
518                    Out::Writing(frame, offset)
519                };
520                Ok(Work::Progress)
521            }
522            Out::Flushing(frame) => match self.transport.poll_flush(cx) {
523                Poll::Pending => Ok(Work::Blocked),
524                Poll::Ready(Err(error)) => Err(ConnectionError::Transport(error)),
525                Poll::Ready(Ok(())) => {
526                    if frame == Frame::Data && self.tx == Tx::Sending {
527                        self.tx = Tx::Waiting;
528                        self.final_input = None;
529                        self.tx_deadline = self
530                            .timer
531                            .now()
532                            .saturating_add(self.config.retransmit_timeout);
533                    }
534                    self.out = Out::Idle;
535                    Ok(Work::Progress)
536                }
537            },
538        }
539    }
540
541    fn poll_input(
542        &mut self,
543        cx: &mut Context<'_>,
544        budget: &mut Budget,
545    ) -> Result<Work, ConnectionError<T::Error>> {
546        if !budget.spend() {
547            return Ok(Work::Yield);
548        }
549        let free = &mut self.rx_buf[self.rx_len..];
550        let read = match self.transport.poll_read(cx, free) {
551            Poll::Pending => return Ok(Work::Blocked),
552            Poll::Ready(Err(error)) => return Err(ConnectionError::Transport(error)),
553            Poll::Ready(Ok(0)) => return Err(ConnectionError::UnexpectedEof),
554            Poll::Ready(Ok(read)) => read.min(free.len()),
555        };
556        self.rx_len += read;
557        if let Some(left) = self.final_input.as_mut() {
558            *left = left.saturating_sub(read);
559        }
560        self.process_input();
561        Ok(Work::Progress)
562    }
563
564    /// Expiry of the deadline does not fail the packet at once: ready input
565    /// gets a final allowance, since the acknowledgment may already be queued
566    /// behind other bytes. The packet is retransmitted, or the connection
567    /// times out, once the input blocks or the allowance is spent.
568    fn poll_timer(
569        &mut self,
570        cx: &mut Context<'_>,
571        input: Work,
572    ) -> Result<Work, ConnectionError<T::Error>> {
573        if self.tx != Tx::Waiting {
574            return Ok(Work::Blocked);
575        }
576        match self.final_input {
577            None => {
578                if self.timer.poll_deadline(cx, self.tx_deadline).is_pending() {
579                    return Ok(Work::Blocked);
580                }
581                self.final_input = Some(Self::FINAL_INPUT);
582                Ok(Work::Progress)
583            }
584            Some(left) if input == Work::Blocked || left == 0 => {
585                self.final_input = None;
586                if self.tx_retries >= self.config.max_retries {
587                    return Err(ConnectionError::Timeout);
588                }
589                self.tx_retries += 1;
590                self.tx = Tx::Queued;
591                Ok(Work::Progress)
592            }
593            Some(_) => Ok(Work::Blocked),
594        }
595    }
596
597    fn process_input(&mut self) {
598        loop {
599            let Some(end) = self.rx_buf[..self.rx_len]
600                .iter()
601                .position(|&byte| byte == 0)
602            else {
603                if self.rx_discarding || self.rx_len == N {
604                    self.rx_discarding = true;
605                    self.rx_len = 0;
606                }
607                return;
608            };
609            if self.rx_discarding {
610                self.rx_discarding = false;
611            } else if end > 0 {
612                self.process_frame(end);
613            }
614            self.rx_buf.copy_within(end + 1..self.rx_len, 0);
615            self.rx_len -= end + 1;
616        }
617    }
618
619    fn process_frame(&mut self, end: usize) {
620        let Ok(decoded) = sync::decode_to_slice(&self.rx_buf[..=end], &mut self.scratch) else {
621            return;
622        };
623        match protocol::parse(&self.scratch[..decoded.len]) {
624            Some(Packet::Ack { session, seq }) => {
625                if session == self.config.session
626                    && self.tx != Tx::Idle
627                    && u64::from(seq) + 1 == self.tx_next
628                {
629                    self.tx = Tx::Idle;
630                    self.tx_retries = 0;
631                    self.final_input = None;
632                }
633            }
634            Some(Packet::Data {
635                session,
636                seq,
637                payload,
638            }) => {
639                if session != self.config.session || payload.len() > Self::MAX_PAYLOAD {
640                    return;
641                }
642                if u64::from(seq) == self.rx_next {
643                    if self.rx_start == self.rx_end {
644                        self.rx_payload[..payload.len()].copy_from_slice(payload);
645                        self.rx_start = 0;
646                        self.rx_end = payload.len();
647                        self.rx_next += 1;
648                        self.ack_pending = Some(seq);
649                    }
650                } else if self.rx_next > 0
651                    && u64::from(seq) == self.rx_next - 1
652                    && self.ack_pending.is_none()
653                {
654                    self.ack_pending = Some(seq);
655                }
656            }
657            None => {}
658        }
659    }
660}
661
662impl<T: Transport, C: Timer, const N: usize> Transport for Reliable<T, C, N> {
663    type Error = ConnectionError<T::Error>;
664
665    fn poll_read(
666        &mut self,
667        cx: &mut Context<'_>,
668        buf: &mut [u8],
669    ) -> Poll<Result<usize, Self::Error>> {
670        if buf.is_empty() {
671            return Poll::Ready(Ok(0));
672        }
673        self.drive(cx, &mut Budget::new());
674        let poll = if self.rx_start < self.rx_end {
675            let len = buf.len().min(self.rx_end - self.rx_start);
676            buf[..len].copy_from_slice(&self.rx_payload[self.rx_start..self.rx_start + len]);
677            self.rx_start += len;
678            Poll::Ready(Ok(len))
679        } else {
680            match self.failure() {
681                Some(error) => Poll::Ready(Err(error)),
682                None => Poll::Pending,
683            }
684        };
685        self.finish(Op::Read, cx, poll)
686    }
687
688    fn poll_write(&mut self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, Self::Error>> {
689        if buf.is_empty() {
690            return Poll::Ready(Ok(0));
691        }
692        let poll = self.queue_packet(cx, buf);
693        self.finish(Op::Write, cx, poll)
694    }
695
696    fn poll_flush(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
697        self.drive(cx, &mut Budget::new());
698        let poll = if let Some(error) = self.failure() {
699            Poll::Ready(Err(error))
700        } else if self.flush_ready() {
701            Poll::Ready(Ok(()))
702        } else {
703            Poll::Pending
704        };
705        self.finish(Op::Flush, cx, poll)
706    }
707}
cobs_io_async/connection.rs

cobs_io_async/
connection.rs
