hopr_protocol_session/socket/
ack_state.rs

1//! This module defines the [`SocketState`] that turns [`SessionSocket`](super::SessionSocket) into
2//! a reliable socket, with segment/frame retransmission and frame acknowledgements.
3
4use std::{
5    sync::atomic::AtomicBool,
6    time::{Duration, Instant},
7};
8
9use futures::{FutureExt, StreamExt};
10use futures_time::stream::StreamExt as TimeStreamExt;
11use hopr_utils::network_types::crossfire_sink::{CrossfireSink, bounded_sink_channel};
12use tracing::Instrument;
13
14use crate::{
15    errors::SessionError,
16    processing::types::FrameInspector,
17    protocol::{FrameAcknowledgements, FrameId, Segment, SegmentId, SegmentRequest, SeqIndicator, SessionMessage},
18    socket::{SocketState, state::SocketComponents},
19    utils::{
20        RetriedFrameId, RingBufferProducer, RingBufferView, next_deadline_with_backoff, searchable_ringbuffer,
21        skip_queue::{Skip, SkipDelaySender, skip_delay_channel},
22    },
23};
24
25/// Indicates the acknowledgement mode of a [stateful](AcknowledgementState) Session socket.
26#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
27pub enum AcknowledgementMode {
28    /// Partial frames are acknowledged, leading to receiver-driven retransmission requests.
29    /// The frame sender is never going to retransmit the entire frame.
30    Partial,
31    /// Only full frames are acknowledged, leading to the full-frame retransmission if
32    /// no acknowledgement is received by the frame sender.
33    Full,
34    /// Both partial and full acknowledgements are sent by the receiver.
35    ///
36    /// If a frame is partially acknowledged first, only receiver-driven retransmission requests follow
37    /// (as with [`AcknowledgementMode::Partial`].
38    ///
39    /// If the frame sender receives no acknowledgement (partial nor full), it retransmits
40    /// the entire frame (as in [`AcknowledgementMode::Full`]).
41    #[default]
42    Both,
43}
44
45impl AcknowledgementMode {
46    /// Indicates if `self` is [`AcknowledgementMode::Partial`] or [`AcknowledgementMode::Both`].
47    #[inline]
48    fn is_partial_ack_enabled(&self) -> bool {
49        matches!(self, Self::Partial | Self::Both)
50    }
51
52    /// Indicates if `self` is [`AcknowledgementMode::Full`] or [`AcknowledgementMode::Both`].
53    #[inline]
54    fn is_full_ack_enabled(&self) -> bool {
55        matches!(self, Self::Full | Self::Both)
56    }
57}
58
59/// Configuration object of the [`AcknowledgementState`].
60#[derive(Debug, Clone, Copy, PartialEq, smart_default::SmartDefault)]
61pub struct AcknowledgementStateConfig {
62    /// Mode of frame acknowledgement.
63    ///
64    /// Default is [`AcknowledgementMode::Both`]
65    pub mode: AcknowledgementMode,
66
67    /// The expected (average) latency of a packet (= single frame segment).
68    ///
69    /// Default is 20 ms
70    #[default(Duration::from_millis(20))]
71    pub expected_packet_latency: Duration,
72
73    /// Backoff base applied for segment or frame retransmissions.
74    ///
75    /// Default is 1.2
76    #[default(1.2)]
77    pub backoff_base: f64,
78
79    /// The maximum number of receiver-driven segment retransmission requests.
80    ///
81    /// Default is 3
82    #[default(3)]
83    pub max_incoming_frame_retries: usize,
84
85    /// The maximum number of sender-driven full-frame retransmissions.
86    ///
87    /// Default is 3
88    #[default(3)]
89    pub max_outgoing_frame_retries: usize,
90
91    /// Delay between acknowledgement batches.
92    ///
93    /// Default is 50 ms
94    #[default(Duration::from_millis(50))]
95    pub acknowledgement_delay: Duration,
96
97    /// Number of segments to hold back for retransmission upon other party's request.
98    /// Minimum is 1024.
99    ///
100    /// Default is 16 384.
101    #[default(16384)]
102    pub lookbehind_segments: usize,
103}
104
105impl AcknowledgementStateConfig {
106    fn normalize(self) -> AcknowledgementStateConfig {
107        Self {
108            mode: self.mode,
109            expected_packet_latency: self.expected_packet_latency.max(Duration::from_millis(1)),
110            backoff_base: self.backoff_base.max(1.0),
111            max_incoming_frame_retries: self.max_incoming_frame_retries,
112            max_outgoing_frame_retries: self.max_outgoing_frame_retries,
113            acknowledgement_delay: self.acknowledgement_delay.max(Duration::from_millis(1)),
114            lookbehind_segments: self.lookbehind_segments.max(1024),
115        }
116    }
117}
118
119#[derive(Clone)]
120struct AcknowledgementStateContext<const C: usize> {
121    rb_tx: RingBufferProducer<Segment>,
122    rb_rx: RingBufferView<Segment>,
123    incoming_frame_retries_tx: SkipDelaySender<RetriedFrameId>,
124    outgoing_frame_retries_tx: SkipDelaySender<RetriedFrameId>,
125    ack_tx: CrossfireSink<FrameId>,
126    inspector: FrameInspector,
127    ctl_tx: CrossfireSink<SessionMessage<C>>,
128}
129
130#[cfg_attr(doc, aquamarine::aquamarine)]
131/// Represents a Session socket state is able to process acknowledgements.
132///
133/// # Retransmission driven by the Receiver
134/// ```mermaid
135/// sequenceDiagram
136///     Note over Sender,Receiver: Frame 1
137///     rect rgb(191, 223, 255)
138///     Note left of Sender: Frame 1 in buffer
139///     Sender->>Receiver: Segment 1/3 of Frame 1
140///     Sender->>Receiver: Segment 2/3 of Frame 1
141///     Sender--xReceiver: Segment 3/3 of Frame 1
142///     Note right of Receiver: RTO_BASE_RECEIVER elapsed
143///     Receiver->>Sender: Request Segment 3 of Frame 1
144///     Sender->>Receiver: Segment 3/3 of Frame 1
145///     Receiver->>Sender: Acknowledge Frame 1
146///     Note left of Sender: Frame 1 dropped from buffer
147///     end
148///     Note over Sender,Receiver: Frame 1 delivered
149/// ```
150///
151/// # Retransmission driven by the Sender
152/// ```mermaid
153/// sequenceDiagram
154///     Note over Sender,Receiver: Frame 1
155///     rect rgb(191, 223, 255)
156///     Note left of Sender: Frame 1 in buffer
157///     Sender->>Receiver: Segment 1/3 of Frame 1
158///     Sender->>Receiver: Segment 2/3 of Frame 1
159///     Sender--xReceiver: Segment 3/3 of Frame 1
160///     Note right of Receiver: RTO_BASE_RECEIVER elapsed
161///     Receiver--xSender: Request Segment 3 of Frame 1
162///     Note left of Sender: RTO_BASE_SENDER elapsed
163///     Sender->>Receiver: Segment 1/3 of Frame 1
164///     Sender->>Receiver: Segment 2/3 of Frame 1
165///     Sender->>Receiver: Segment 3/3 of Frame 1
166///     Receiver->>Sender: Acknowledge Frame 1
167///     Note left of Sender: Frame 1 dropped from buffer
168///     end
169///     Note over Sender,Receiver: Frame 1 delivered
170/// ```
171///
172/// # Sender-Receiver retransmission handover
173///
174/// ```mermaid
175///    sequenceDiagram
176///     Note over Sender,Receiver: Frame 1
177///     rect rgb(191, 223, 255)
178///     Note left of Sender: Frame 1 in buffer
179///     Sender->>Receiver: Segment 1/3 of Frame 1
180///     Sender--xReceiver: Segment 2/3 of Frame 1
181///     Sender--xReceiver: Segment 3/3 of Frame 1
182///     Note right of Receiver: RTO_BASE_RECEIVER elapsed
183///     Receiver->>Sender: Request Segments 2,3 of Frame 1
184///     Note left of Sender: RTO_BASE_SENDER cancelled
185///     Sender->>Receiver: Segment 2/3 of Frame 1
186///     Sender--xReceiver: Segment 3/3 of Frame 1
187///     Note right of Receiver: RTO_BASE_RECEIVER elapsed
188///     Receiver--xSender: Request Segments 3 of Frame 1
189///     Note right of Receiver: RTO_BASE_RECEIVER elapsed
190///     Receiver->>Sender: Request Segments 3 of Frame 1
191///     Sender->>Receiver: Segment 3/3 of Frame 1
192///     Receiver->>Sender: Acknowledge Frame 1
193///     Note left of Sender: Frame 1 dropped from buffer
194///     end
195///     Note over Sender,Receiver: Frame 1 delivered
196/// ```
197///
198/// # Retransmission failure
199///
200/// ```mermaid
201///    sequenceDiagram
202///     Note over Sender,Receiver: Frame 1
203///     rect rgb(191, 223, 255)
204///     Note left of Sender: Frame 1 in buffer
205///     Sender->>Receiver: Segment 1/3 of Frame 1
206///     Sender->>Receiver: Segment 2/3 of Frame 1
207///     Sender--xReceiver: Segment 3/3 of Frame 1
208///     Note right of Receiver: RTO_BASE_RECEIVER elapsed
209///     Receiver--xSender: Request Segment 3 of Frame 1
210///     Note left of Sender: RTO_BASE_SENDER elapsed
211///     Sender--xReceiver: Segment 1/3 of Frame 1
212///     Sender--xReceiver: Segment 2/3 of Frame 1
213///     Sender--xReceiver: Segment 3/3 of Frame 1
214///     Note left of Sender: FRAME_MAX_AGE elapsed<br/>Frame 1 dropped from buffer
215///     Note right of Receiver: FRAME_MAX_AGE elapsed<br/>Frame 1 dropped from buffer
216///     end
217///     Note over Sender,Receiver: Frame 1 never delivered
218/// ```
219#[derive(Clone)]
220pub struct AcknowledgementState<const C: usize> {
221    id: String,
222    cfg: AcknowledgementStateConfig,
223    context: Option<AcknowledgementStateContext<C>>,
224    started: std::sync::Arc<AtomicBool>,
225}
226
227impl<const C: usize> AcknowledgementState<C> {
228    pub fn new<I: std::fmt::Display>(session_id: I, cfg: AcknowledgementStateConfig) -> Self {
229        Self {
230            id: session_id.to_string(),
231            cfg: cfg.normalize(),
232            context: Default::default(),
233            started: std::sync::Arc::new(AtomicBool::new(false)),
234        }
235    }
236}
237
238impl<const C: usize> SocketState<C> for AcknowledgementState<C> {
239    fn session_id(&self) -> &str {
240        &self.id
241    }
242
243    #[tracing::instrument(name = "AcknowledgementState", skip(self, socket_components), fields(session_id = self.id))]
244    fn run(&mut self, socket_components: SocketComponents<C>) -> Result<(), SessionError> {
245        if self.started.load(std::sync::atomic::Ordering::Relaxed) && self.context.is_some() {
246            return Err(SessionError::InvalidState("state is already running".into()));
247        }
248
249        let (incoming_frame_retries_tx, incoming_frame_retries_rx) = skip_delay_channel();
250        let (outgoing_frame_retries_tx, outgoing_frame_retries_rx) = skip_delay_channel();
251        let (rb_tx, rb_rx) = searchable_ringbuffer(self.cfg.lookbehind_segments);
252
253        // Full frame acknowledgements get a special channel with fixed capacity
254        let (ack_tx, ack_rx) = bounded_sink_channel::<FrameId>(2 * self.cfg.lookbehind_segments);
255
256        let context = self.context.insert(AcknowledgementStateContext {
257            rb_tx,
258            rb_rx,
259            incoming_frame_retries_tx,
260            outgoing_frame_retries_tx,
261            ack_tx,
262            ctl_tx: socket_components.ctl_tx,
263            inspector: socket_components
264                .inspector
265                .ok_or(SessionError::InvalidState("inspector is not available".into()))?,
266        });
267
268        if self.cfg.mode.is_partial_ack_enabled() {
269            // For partially received frames incomplete for too long,
270            // missing segments will be asked for retransmission
271            let mut incoming_frame_retries_tx_clone = context.incoming_frame_retries_tx.clone();
272            let ctl_tx_clone = context.ctl_tx.clone();
273            let frame_inspector_clone = context.inspector.clone();
274            let cfg = self.cfg;
275            hopr_utils::runtime::prelude::spawn(incoming_frame_retries_rx
276                .filter_map(move |rf| {
277                    let frame_id = rf.frame_id;
278                    let missing_segments = frame_inspector_clone.missing_segments(&frame_id).unwrap_or_default();
279                    if !missing_segments.is_empty() {
280                        // Find out if we need to subscribe for further retries of this Frame
281                        if let Some(next) = rf.next() {
282                            // Register the next retry if still possible
283                            let retry_at = next_deadline_with_backoff(next.retry_count, cfg.backoff_base, cfg.expected_packet_latency);
284                            if let Err(error) = incoming_frame_retries_tx_clone.send_one((next, retry_at)) {
285                                tracing::error!(frame_id, %error, "failed to register next resend of incoming frame");
286                            } else {
287                                tracing::debug!(frame_id, retry_in = ?retry_at.saturating_duration_since(Instant::now()), "next resend request of incoming frame segments");
288                            }
289                        } else {
290                            tracing::debug!(frame_id, "last request of incoming frame segments");
291                        }
292
293                        futures::future::ready(Some((frame_id, missing_segments)))
294                    } else {
295                        tracing::debug!(frame_id, "no more missing segments in frame");
296                        futures::future::ready(None)
297                    }
298                })
299                .ready_chunks(SegmentRequest::<C>::MAX_ENTRIES)
300                .inspect(|r| tracing::trace!(req = ?r, "requesting segments resend"))
301                .map(|a| Ok(SessionMessage::<C>::Request(a.into_iter().collect())))
302                .forward(ctl_tx_clone)
303                .map(move |res| match res {
304                    Ok(_) => tracing::debug!("incoming frame resends processing done"),
305                    Err(error) => tracing::error!(%error, "error while processing incoming frame resends")
306                })
307                .instrument(tracing::debug_span!("incoming_frame_retries_sender"))
308            );
309        }
310
311        // Send out Frame Acknowledgements chunked as Control messages
312        let ctl_tx_clone = context.ctl_tx.clone();
313        let ack_delay = self.cfg.acknowledgement_delay;
314        hopr_utils::runtime::prelude::spawn(
315            ack_rx
316                .buffer(futures_time::time::Duration::from(ack_delay))
317                .flat_map(|acks| futures::stream::iter(FrameAcknowledgements::<C>::new_multiple(acks)))
318                .filter(|acks| futures::future::ready(!acks.is_empty()))
319                .inspect(|acks| tracing::trace!(?acks, "acknowledgements sent"))
320                .map(|acks| Ok(SessionMessage::<C>::Acknowledge(acks)))
321                .forward(ctl_tx_clone)
322                .map(move |res| match res {
323                    Ok(_) => tracing::debug!("acknowledgement forwarding done"),
324                    Err(error) => tracing::debug!(%error, "acknowledgement forwarding failed"),
325                })
326                .instrument(tracing::debug_span!("acknowledgement_sender")),
327        );
328
329        // Resend outgoing frame Segments if they were not (partially or fully) acknowledged
330        let mut outgoing_frame_retries_tx_clone = context.outgoing_frame_retries_tx.clone();
331        let ctl_tx_clone = context.ctl_tx.clone();
332        let rb_rx_clone = context.rb_rx.clone();
333        let cfg = self.cfg;
334        hopr_utils::runtime::prelude::spawn(
335            outgoing_frame_retries_rx
336                .map(move |rf: RetriedFrameId| {
337                    // Find out if the frame can be retried again in the future
338                    let frame_id = rf.frame_id;
339                    if let Some(next) = rf.next() {
340                        // Register the next retry if still possible
341                        let retry_at =
342                            next_deadline_with_backoff(next.retry_count, cfg.backoff_base, cfg.expected_packet_latency);
343                        if let Err(error) = outgoing_frame_retries_tx_clone.send_one((next, retry_at)) {
344                            tracing::error!(frame_id, %error, "failed to register next retry of frame");
345                        } else {
346                            tracing::debug!(frame_id, retry_in = ?retry_at.saturating_duration_since(Instant::now()), "next resend of outgoing frame");
347                        }
348                    } else {
349                        tracing::debug!(frame_id, "last outgoing retry of frame");
350                    }
351                    tracing::trace!(frame_id, "going to re-send entire frame");
352                    frame_id
353                })
354                .flat_map(move |frame_id| {
355                    // Find out all the segments of that frame to be retransmitted
356                    futures::stream::iter(
357                        rb_rx_clone
358                            .find(|s: &Segment| s.id().0 == frame_id)
359                            .into_iter()
360                            .inspect(|s| tracing::trace!(seg_id = %s.id(), "segment retransmit"))
361                            .map(|s| Ok(SessionMessage::<C>::Segment(s))),
362                    )
363                })
364                .forward(ctl_tx_clone) // Retransmit all the segments
365                .map(move |res| match res {
366                    Ok(_) => tracing::debug!("outgoing frame retries processing done"),
367                    Err(error) => tracing::error!(%error, "error while processing outgoing frame retries"),
368                })
369                .instrument(tracing::debug_span!("outgoing_frame_retries_sender")),
370        );
371
372        tracing::debug!("acknowledgement state has been started");
373        self.started.store(true, std::sync::atomic::Ordering::Relaxed);
374
375        Ok(())
376    }
377
378    #[tracing::instrument(name = "AcknowledgementState::stop", skip(self), fields(session_id = self.id))]
379    fn stop(&mut self) -> Result<(), SessionError> {
380        if let Some(mut ctx) = self.context.take() {
381            ctx.outgoing_frame_retries_tx.force_close();
382            ctx.incoming_frame_retries_tx.force_close();
383            // ack_tx and ctl_tx close when their last clone is dropped;
384            // dropping ctx here drops the originals, spawned-task clones drop when those tasks complete.
385
386            self.started.store(false, std::sync::atomic::Ordering::Relaxed);
387            tracing::debug!("state has been stopped");
388        } else {
389            tracing::warn!("cannot be stopped, because it is not running");
390        }
391
392        Ok(())
393    }
394
395    #[tracing::instrument(name = "AcknowledgementState::incoming_segment", skip(self), fields(session_id = self.id, frame_id = seg_id.0))]
396    fn incoming_segment(&mut self, seg_id: &SegmentId, _ind: SeqIndicator) -> Result<(), SessionError> {
397        tracing::trace!("segment received");
398
399        let ctx = self
400            .started
401            .load(std::sync::atomic::Ordering::Relaxed)
402            .then_some(self.context.as_mut())
403            .flatten()
404            .ok_or(SessionError::StateNotRunning)?;
405
406        // Register future requesting of segments for this frame
407        if self.cfg.mode.is_partial_ack_enabled() {
408            // Every incoming segment of this frame will move the deadline further
409            // into the future.
410            if let Err(error) = ctx.incoming_frame_retries_tx.send_one((
411                RetriedFrameId::with_retries(seg_id.0, self.cfg.max_incoming_frame_retries),
412                self.cfg.expected_packet_latency, // RTO_BASE_RECEIVER - when we expect the next segment to arrive
413            )) {
414                tracing::error!(%error, "failed to register incoming retry for frame");
415            }
416        }
417        Ok(())
418    }
419
420    #[tracing::instrument(name = "AcknowledgementState::incoming_retransmission_request", skip(self, request), fields(session_id = self.id))]
421    fn incoming_retransmission_request(&mut self, request: SegmentRequest<C>) -> Result<(), SessionError> {
422        // The state will respond to segment retransmission requests even
423        // if it has this feature disabled in the config.
424        tracing::trace!(count = request.len(), "segment retransmission requested");
425
426        let ctx = self
427            .started
428            .load(std::sync::atomic::Ordering::Relaxed)
429            .then_some(self.context.as_mut())
430            .flatten()
431            .ok_or(SessionError::StateNotRunning)?;
432
433        let (mut missing_seg_ids, mut missing_frame_ids): (Vec<_>, Vec<_>) =
434            request.into_iter().map(|s| (s, s.0)).unzip();
435
436        // Perform a single find to lock the RB only once
437        let segments = ctx.rb_rx.find(|s| {
438            // SegmentIds are guaranteed to be sorted, so we can use binary search
439            if let Ok(i) = missing_seg_ids.binary_search(&s.id()) {
440                missing_seg_ids.remove(i);
441                true
442            } else {
443                false
444            }
445        });
446
447        tracing::trace!(
448            found = segments.len(),
449            requested = missing_frame_ids.len(),
450            "found matching segments to be retransmitted"
451        );
452
453        // Partially acknowledged frames will not need to be fully resent in the future.
454        // Cancel all partially acknowledged frame resends.
455        if self.cfg.mode.is_full_ack_enabled() {
456            // Since the FrameIds are guaranteed to be sorted, we can simply dedup them.
457            missing_frame_ids.dedup();
458
459            if let Err(error) = ctx.outgoing_frame_retries_tx.send_many(
460                missing_frame_ids
461                    .into_iter()
462                    .map(|frame_id| (RetriedFrameId::no_retries(frame_id), Skip).into()),
463            ) {
464                tracing::error!(%error, "failed to cancel frame resend of partially acknowledged frames");
465            }
466        }
467
468        // Resend the segments via the Control Stream
469        segments
470            .into_iter()
471            .try_for_each(|s| {
472                tracing::trace!(seg_id = %s.id(), "retransmit segment on request");
473                ctx.ctl_tx.try_send(SessionMessage::Segment(s))
474            })
475            .map_err(|e| SessionError::ProcessingError(e.to_string()))
476    }
477
478    #[tracing::instrument(name = "AcknowledgementState::incoming_acknowledged_frames", skip(self), fields(session_id = self.id))]
479    fn incoming_acknowledged_frames(&mut self, ack: FrameAcknowledgements<C>) -> Result<(), SessionError> {
480        tracing::trace!(count = ack.len(), "frame acknowledgements received");
481
482        let ctx = self
483            .started
484            .load(std::sync::atomic::Ordering::Relaxed)
485            .then_some(self.context.as_mut())
486            .flatten()
487            .ok_or(SessionError::StateNotRunning)?;
488
489        // Frame acknowledged, we will not need to resend it
490        if self.cfg.mode.is_full_ack_enabled()
491            && let Err(error) = ctx.outgoing_frame_retries_tx.send_many(
492                ack.into_iter()
493                    .inspect(|frame_id| tracing::trace!(frame_id, "frame acknowledged"))
494                    .map(|frame_id| (RetriedFrameId::no_retries(frame_id), Skip).into()),
495            )
496        {
497            tracing::error!(%error, "failed to cancel frame resend");
498        }
499
500        Ok(())
501    }
502
503    #[tracing::instrument(name = "AcknowledgementState::frame_complete", skip(self), fields(session_id = self.id))]
504    fn frame_complete(&mut self, frame_id: FrameId) -> Result<(), SessionError> {
505        tracing::trace!("frame complete");
506
507        let ctx = self
508            .started
509            .load(std::sync::atomic::Ordering::Relaxed)
510            .then_some(self.context.as_mut())
511            .flatten()
512            .ok_or(SessionError::StateNotRunning)?;
513
514        // Since the frame has been completed, push its ID into the acknowledgement queue
515        if let Err(error) = ctx.ack_tx.try_send(frame_id) {
516            tracing::error!(%error, "failed to acknowledge frame");
517        }
518
519        if self.cfg.mode.is_partial_ack_enabled() {
520            // No more requesting of segment retransmissions from frames that were completed
521            if let Err(error) = ctx
522                .incoming_frame_retries_tx
523                .send_one((RetriedFrameId::no_retries(frame_id), Skip))
524            {
525                tracing::error!(%error, "failed to cancel retry of acknowledged frame");
526            }
527        }
528
529        Ok(())
530    }
531
532    #[tracing::instrument(name = "AcknowledgementState::frame_emitted", skip(self), fields(session_id = self.id))]
533    fn frame_emitted(&mut self, id: FrameId) -> Result<(), SessionError> {
534        tracing::trace!("frame emitted");
535        let _ = self
536            .started
537            .load(std::sync::atomic::Ordering::Relaxed)
538            .then_some(self.context.as_mut())
539            .flatten()
540            .ok_or(SessionError::StateNotRunning)?;
541        Ok(())
542    }
543
544    #[tracing::instrument(name = "AcknowledgementState::frame_discarded", skip(self), fields(session_id = self.id))]
545    fn frame_discarded(&mut self, frame_id: FrameId) -> Result<(), SessionError> {
546        tracing::trace!("frame discarded");
547
548        let ctx = self
549            .started
550            .load(std::sync::atomic::Ordering::Relaxed)
551            .then_some(self.context.as_mut())
552            .flatten()
553            .ok_or(SessionError::StateNotRunning)?;
554
555        if self.cfg.mode.is_partial_ack_enabled() {
556            // No more requesting of segment retransmissions from frames that were discarded
557            if let Err(error) = ctx
558                .incoming_frame_retries_tx
559                .send_one((RetriedFrameId::no_retries(frame_id), Skip))
560            {
561                tracing::error!(%error, "failed to cancel retry of acknowledged frame");
562            }
563        }
564
565        Ok(())
566    }
567
568    #[tracing::instrument(name = "AcknowledgementState::segment_sent", skip(self, segment), fields(session_id = self.id, frame_id = segment.frame_id, seq_idx = segment.seq_idx))]
569    fn segment_sent(&mut self, segment: &Segment) -> Result<(), SessionError> {
570        tracing::trace!("segment sent");
571
572        let ctx = self
573            .started
574            .load(std::sync::atomic::Ordering::Relaxed)
575            .then_some(self.context.as_mut())
576            .flatten()
577            .ok_or(SessionError::StateNotRunning)?;
578
579        // Since segments are re-sent via Control stream, they are not later fed again
580        // into the ring buffer.
581        ctx.rb_tx.push(segment.clone());
582
583        // When the last segment of a frame has been sent,
584        // add it to outgoing retries (if the full ack mode is enabled).
585        if segment.is_last() && self.cfg.mode.is_full_ack_enabled() {
586            tracing::trace!("last segment of frame sent");
587
588            if let Err(error) = ctx.outgoing_frame_retries_tx.send_one((
589                RetriedFrameId::with_retries(segment.frame_id, self.cfg.max_outgoing_frame_retries),
590                // The whole frame should be delivered and acknowledged
591                // once all its segments (seq_len) are sent,
592                // and the acknowledgement also comes back to us.
593                // Therefore, RTO_BASE_SENDER = latency * (seq_len + 1)
594                self.cfg.expected_packet_latency * (segment.seq_flags.seq_len() + 1) as u32,
595            )) {
596                tracing::error!(%error, "failed to insert outgoing retry of a frame");
597            }
598        }
599
600        Ok(())
601    }
602}
603
604#[cfg(test)]
605mod tests {
606    use anyhow::Context;
607
608    use super::*;
609    use crate::{
610        processing::types::{FrameBuilder, FrameDashMap, FrameMap},
611        protocol::SeqNum,
612        utils::segment,
613    };
614
615    const FRAME_SIZE: usize = 1500;
616
617    const MTU: usize = 1000;
618
619    #[test_log::test(tokio::test)]
620    async fn ack_state_sender_must_acknowledge_completed_frames() -> anyhow::Result<()> {
621        let cfg = AcknowledgementStateConfig {
622            expected_packet_latency: Duration::from_millis(10),
623            acknowledgement_delay: Duration::from_millis(2),
624            ..Default::default()
625        };
626
627        let inspector = FrameInspector(FrameDashMap::with_capacity(10));
628        let (ctl_tx, ctl_rx) = bounded_sink_channel::<SessionMessage<MTU>>(1024);
629
630        let mut state = AcknowledgementState::<MTU>::new("test", cfg);
631        state.run(SocketComponents {
632            inspector: inspector.into(),
633            ctl_tx,
634        })?;
635
636        let acked_frame_ids = [1, 2, 3];
637
638        for &frame_id in &acked_frame_ids {
639            state.frame_complete(frame_id)?;
640        }
641
642        tokio::time::sleep(cfg.acknowledgement_delay * 2).await;
643
644        state.stop()?;
645
646        let ctl_msgs = tokio::time::timeout(Duration::from_millis(100), ctl_rx.collect::<Vec<_>>())
647            .await
648            .context("timeout receiving Control messages")?;
649
650        assert_eq!(1, ctl_msgs.len());
651
652        assert_eq!(
653            ctl_msgs[0],
654            SessionMessage::Acknowledge(acked_frame_ids.to_vec().try_into()?)
655        );
656
657        Ok(())
658    }
659
660    #[parameterized::parameterized(num_frames = { 1, 2, 3 })]
661    #[parameterized_macro(test_log::test(tokio::test))]
662    async fn ack_state_sender_must_resend_unacknowledged_frames(num_frames: usize) -> anyhow::Result<()> {
663        const NUM_RETRIES: usize = 2;
664
665        let cfg = AcknowledgementStateConfig {
666            mode: AcknowledgementMode::Full,
667            expected_packet_latency: Duration::from_millis(2),
668            max_outgoing_frame_retries: NUM_RETRIES,
669            ..Default::default()
670        };
671
672        let inspector = FrameInspector(FrameDashMap::with_capacity(10));
673        let (ctl_tx, ctl_rx) = bounded_sink_channel::<SessionMessage<MTU>>(1024);
674
675        let mut state = AcknowledgementState::<MTU>::new("test", cfg);
676        state.run(SocketComponents {
677            inspector: inspector.into(),
678            ctl_tx,
679        })?;
680
681        let mut expected_frame_segments = Vec::new();
682        let num_segments_in_frame = FRAME_SIZE / MTU + 1;
683        for i in 1..=num_frames {
684            let expected_segments = segment(
685                hopr_types::crypto_random::random_bytes::<FRAME_SIZE>(),
686                MTU,
687                i as FrameId,
688            )?;
689            for segment in &expected_segments {
690                state.segment_sent(segment)?;
691            }
692            expected_frame_segments.push(expected_segments);
693        }
694
695        let expected_frame_delivery = cfg.expected_packet_latency * (num_segments_in_frame + 1) as u32;
696        tokio::time::sleep(2 * expected_frame_delivery).await;
697        state.stop()?;
698
699        let ctl_msg = tokio::time::timeout(Duration::from_millis(100), ctl_rx.collect::<Vec<_>>())
700            .await
701            .context("timeout receiving Control message")?;
702
703        let retransmitted_segments = ctl_msg
704            .into_iter()
705            .map(|m| m.try_as_segment().ok_or(anyhow::anyhow!("must be segment")))
706            .collect::<Result<Vec<_>, _>>()?;
707
708        assert_eq!(
709            NUM_RETRIES * num_segments_in_frame * num_frames,
710            retransmitted_segments.len()
711        );
712
713        let total_segments = expected_frame_segments.iter().map(|m| m.len()).sum::<usize>();
714        let expected_segments = expected_frame_segments
715            .into_iter()
716            .flatten()
717            .cycle()
718            .take(total_segments * NUM_RETRIES)
719            .collect::<Vec<_>>();
720        assert_eq!(expected_segments, retransmitted_segments);
721
722        Ok(())
723    }
724
725    #[test_log::test(tokio::test)]
726    async fn ack_state_sender_must_not_resend_unacknowledged_frame_when_full_resend_disabled() -> anyhow::Result<()> {
727        const NUM_RETRIES: usize = 2;
728
729        let cfg = AcknowledgementStateConfig {
730            mode: AcknowledgementMode::Partial,
731            expected_packet_latency: Duration::from_millis(2),
732            max_outgoing_frame_retries: NUM_RETRIES,
733            ..Default::default()
734        };
735
736        let inspector = FrameInspector(FrameDashMap::with_capacity(10));
737        let (ctl_tx, ctl_rx) = bounded_sink_channel::<SessionMessage<MTU>>(1024);
738
739        let mut state = AcknowledgementState::<MTU>::new("test", cfg);
740        state.run(SocketComponents {
741            inspector: inspector.into(),
742            ctl_tx,
743        })?;
744
745        let expected_segments = segment(hopr_types::crypto_random::random_bytes::<FRAME_SIZE>(), MTU, 1)?;
746        for segment in &expected_segments {
747            state.segment_sent(segment)?;
748        }
749
750        let expected_frame_delivery = cfg.expected_packet_latency * (expected_segments.len() + 1) as u32;
751        tokio::time::sleep(2 * expected_frame_delivery).await;
752        state.stop()?;
753
754        // No retransmission should be sent because it is disabled
755        assert!(ctl_rx.collect::<Vec<_>>().await.is_empty());
756
757        Ok(())
758    }
759
760    #[tokio::test]
761    async fn ack_state_sender_must_not_resend_acknowledged_frame() -> anyhow::Result<()> {
762        let cfg = AcknowledgementStateConfig {
763            mode: AcknowledgementMode::Full,
764            expected_packet_latency: Duration::from_millis(2),
765            max_outgoing_frame_retries: 1,
766            ..Default::default()
767        };
768
769        let inspector = FrameInspector(FrameDashMap::with_capacity(10));
770        let (ctl_tx, ctl_rx) = bounded_sink_channel::<SessionMessage<MTU>>(1024);
771
772        let mut state = AcknowledgementState::<MTU>::new("test", cfg);
773        state.run(SocketComponents {
774            inspector: inspector.into(),
775            ctl_tx,
776        })?;
777
778        let expected_segments = segment(hopr_types::crypto_random::random_bytes::<{ FRAME_SIZE * 2 }>(), MTU, 1)?;
779        for segment in &expected_segments {
780            state.segment_sent(segment)?;
781        }
782
783        // Acknowledge the frame
784        state.incoming_acknowledged_frames(vec![1].try_into()?)?;
785
786        tokio::time::sleep(10 * cfg.expected_packet_latency).await;
787
788        state.stop()?;
789
790        // No retransmission should be sent because the frame was already acknowledged.
791        assert!(ctl_rx.collect::<Vec<_>>().await.is_empty());
792
793        Ok(())
794    }
795
796    #[test_log::test(tokio::test)]
797    async fn ack_state_sender_must_not_resend_entire_frame_when_already_partially_acknowledged() -> anyhow::Result<()> {
798        let cfg = AcknowledgementStateConfig {
799            mode: AcknowledgementMode::Full,
800            expected_packet_latency: Duration::from_millis(2),
801            max_outgoing_frame_retries: 1,
802            ..Default::default()
803        };
804
805        let inspector = FrameInspector(FrameDashMap::with_capacity(10));
806        let (ctl_tx, ctl_rx) = bounded_sink_channel::<SessionMessage<MTU>>(1024);
807
808        let mut state = AcknowledgementState::<MTU>::new("test", cfg);
809        state.run(SocketComponents {
810            inspector: inspector.into(),
811            ctl_tx,
812        })?;
813
814        let expected_segments = segment(hopr_types::crypto_random::random_bytes::<{ FRAME_SIZE * 2 }>(), MTU, 1)?;
815
816        // Load segments into the ring buffer
817        for segment in &expected_segments {
818            state.segment_sent(segment)?;
819        }
820
821        tokio::time::sleep(cfg.expected_packet_latency).await;
822
823        // Partially acknowledge the frame (report the first segment as missing)
824        state.incoming_retransmission_request(SegmentRequest::from_iter([(1, [0b10000000].into())]))?;
825
826        state.stop()?;
827
828        // Only segment 1 must be retransmitted
829        let ctl_msgs = ctl_rx.collect::<Vec<_>>().await;
830        assert_eq!(1, ctl_msgs.len());
831        assert_eq!(ctl_msgs[0], SessionMessage::Segment(expected_segments[0].clone()));
832
833        Ok(())
834    }
835
836    #[test_log::test(tokio::test)]
837    async fn ack_state_sender_must_retransmit_segments_when_requested() -> anyhow::Result<()> {
838        let cfg = AcknowledgementStateConfig {
839            mode: AcknowledgementMode::Full,
840            expected_packet_latency: Duration::from_millis(2),
841            max_outgoing_frame_retries: 1,
842            ..Default::default()
843        };
844
845        let inspector = FrameInspector(FrameDashMap::with_capacity(10));
846        let (ctl_tx, ctl_rx) = bounded_sink_channel::<SessionMessage<MTU>>(1024);
847
848        let mut state = AcknowledgementState::<MTU>::new("test", cfg);
849        state.run(SocketComponents {
850            inspector: inspector.into(),
851            ctl_tx,
852        })?;
853
854        let expected_segments_1 = segment(hopr_types::crypto_random::random_bytes::<{ FRAME_SIZE * 2 }>(), MTU, 1)?;
855        // Load frame 1 segments into the ring buffer
856        for segment in &expected_segments_1 {
857            state.segment_sent(segment)?;
858        }
859
860        let expected_segments_2 = segment(hopr_types::crypto_random::random_bytes::<{ FRAME_SIZE * 2 }>(), MTU, 2)?;
861        // Load frame 2 segments into the ring buffer
862        for segment in &expected_segments_2 {
863            state.segment_sent(segment)?;
864        }
865
866        tokio::time::sleep(cfg.expected_packet_latency).await;
867
868        // Request different segments to be retransmitted
869        state.incoming_retransmission_request(SegmentRequest::from_iter([
870            (1, [0b11100000].into()),
871            (2, [0b11100000].into()),
872        ]))?;
873        tokio::time::sleep(cfg.expected_packet_latency).await;
874
875        state.incoming_retransmission_request(SegmentRequest::from_iter([(2, [0b11000000].into())]))?;
876        tokio::time::sleep(cfg.expected_packet_latency).await;
877
878        state.incoming_retransmission_request(SegmentRequest::from_iter([(2, [0b01000000].into())]))?;
879        tokio::time::sleep(cfg.expected_packet_latency).await;
880
881        state.stop()?;
882
883        let ctl_msgs = ctl_rx.collect::<Vec<_>>().await;
884
885        assert_eq!(9, ctl_msgs.len());
886        // Request 1 - frame 1
887        assert_eq!(ctl_msgs[0], SessionMessage::Segment(expected_segments_1[0].clone()));
888        assert_eq!(ctl_msgs[1], SessionMessage::Segment(expected_segments_1[1].clone()));
889        assert_eq!(ctl_msgs[2], SessionMessage::Segment(expected_segments_1[2].clone()));
890        // Request 1 - frame 2
891        assert_eq!(ctl_msgs[3], SessionMessage::Segment(expected_segments_2[0].clone()));
892        assert_eq!(ctl_msgs[4], SessionMessage::Segment(expected_segments_2[1].clone()));
893        assert_eq!(ctl_msgs[5], SessionMessage::Segment(expected_segments_2[2].clone()));
894
895        // Request 2 - frame 2
896        assert_eq!(ctl_msgs[6], SessionMessage::Segment(expected_segments_2[0].clone()));
897        assert_eq!(ctl_msgs[7], SessionMessage::Segment(expected_segments_2[1].clone()));
898
899        // Request 3 - frame 2
900        assert_eq!(ctl_msgs[8], SessionMessage::Segment(expected_segments_2[1].clone()));
901
902        Ok(())
903    }
904
905    #[tokio::test]
906    async fn ack_state_receiver_must_request_missing_frames_when_partial_acks_are_enabled() -> anyhow::Result<()> {
907        let cfg = AcknowledgementStateConfig {
908            mode: AcknowledgementMode::Partial,
909            expected_packet_latency: Duration::from_millis(2),
910            max_incoming_frame_retries: 1,
911            ..Default::default()
912        };
913
914        let mut inspector = FrameInspector(FrameDashMap::with_capacity(10));
915        let (ctl_tx, ctl_rx) = bounded_sink_channel::<SessionMessage<MTU>>(1024);
916
917        let segments = segment(hopr_types::crypto_random::random_bytes::<FRAME_SIZE>(), MTU, 1)?;
918
919        inspector
920            .0
921            .entry(1)
922            .try_as_vacant()
923            .ok_or(anyhow::anyhow!("frame 1 must be vacant"))?
924            .insert(FrameBuilder::from(segments[0].clone()));
925
926        let mut state = AcknowledgementState::<MTU>::new("test", cfg);
927        state.run(SocketComponents {
928            inspector: inspector.into(),
929            ctl_tx,
930        })?;
931
932        state.incoming_segment(&segments[0].id(), (segments.len() as SeqNum).try_into()?)?;
933
934        tokio::time::sleep(cfg.expected_packet_latency * 2).await;
935
936        state.stop()?;
937
938        let ctl_msgs = tokio::time::timeout(Duration::from_millis(100), ctl_rx.collect::<Vec<_>>())
939            .await
940            .context("timeout receiving Control messages")?;
941
942        assert_eq!(1, ctl_msgs.len());
943        assert_eq!(
944            ctl_msgs[0],
945            SessionMessage::Request(SegmentRequest::from_iter([(1, [0b01000000].into())]))
946        );
947
948        Ok(())
949    }
950
951    #[tokio::test]
952    async fn ack_state_receiver_must_not_request_missing_frames_when_partial_acks_are_disabled() -> anyhow::Result<()> {
953        let cfg = AcknowledgementStateConfig {
954            mode: AcknowledgementMode::Full,
955            expected_packet_latency: Duration::from_millis(2),
956            max_incoming_frame_retries: 1,
957            ..Default::default()
958        };
959
960        let mut inspector = FrameInspector(FrameDashMap::with_capacity(10));
961        let (ctl_tx, ctl_rx) = bounded_sink_channel::<SessionMessage<MTU>>(1024);
962
963        let segments = segment(hopr_types::crypto_random::random_bytes::<FRAME_SIZE>(), MTU, 1)?;
964
965        inspector
966            .0
967            .entry(1)
968            .try_as_vacant()
969            .ok_or(anyhow::anyhow!("frame 1 must be vacant"))?
970            .insert(FrameBuilder::from(segments[0].clone()));
971
972        let mut state = AcknowledgementState::<MTU>::new("test", cfg);
973        state.run(SocketComponents {
974            inspector: inspector.into(),
975            ctl_tx,
976        })?;
977
978        state.incoming_segment(&segments[0].id(), (segments.len() as SeqNum).try_into()?)?;
979
980        tokio::time::sleep(cfg.expected_packet_latency * 2).await;
981
982        state.stop()?;
983
984        let ctl_msgs = tokio::time::timeout(Duration::from_millis(100), ctl_rx.collect::<Vec<_>>())
985            .await
986            .context("timeout receiving Control messages")?;
987
988        assert!(ctl_msgs.iter().all(|m| !matches!(m, SessionMessage::Request(_))));
989
990        Ok(())
991    }
992
993    #[tokio::test]
994    async fn ack_state_receiver_must_continue_requesting_missing_frames_when_frame_not_completed() -> anyhow::Result<()>
995    {
996        let cfg = AcknowledgementStateConfig {
997            mode: AcknowledgementMode::Partial,
998            expected_packet_latency: Duration::from_millis(2),
999            max_incoming_frame_retries: 3,
1000            ..Default::default()
1001        };
1002
1003        let mut inspector = FrameInspector(FrameDashMap::with_capacity(10));
1004        let (ctl_tx, ctl_rx) = bounded_sink_channel::<SessionMessage<MTU>>(1024);
1005
1006        let segments = segment(hopr_types::crypto_random::random_bytes::<{ 2 * FRAME_SIZE }>(), MTU, 1)?;
1007
1008        inspector
1009            .0
1010            .entry(1)
1011            .try_as_vacant()
1012            .ok_or(anyhow::anyhow!("frame 1 must be vacant"))?
1013            .insert(FrameBuilder::from(segments[0].clone()));
1014
1015        let mut state = AcknowledgementState::<MTU>::new("test", cfg);
1016        state.run(SocketComponents {
1017            inspector: inspector.clone().into(),
1018            ctl_tx,
1019        })?;
1020
1021        state.incoming_segment(&segments[0].id(), (segments.len() as SeqNum).try_into()?)?;
1022
1023        tokio::time::sleep(cfg.expected_packet_latency * 2).await;
1024
1025        inspector
1026            .0
1027            .entry(1)
1028            .try_as_occupied()
1029            .ok_or(anyhow::anyhow!("frame 1 must be occupied"))?
1030            .get_mut()
1031            .add_segment(segments[1].clone())?;
1032
1033        state.incoming_segment(&segments[1].id(), (segments.len() as SeqNum).try_into()?)?;
1034
1035        tokio::time::sleep(cfg.expected_packet_latency * 2).await;
1036
1037        state.stop()?;
1038
1039        let ctl_msgs = tokio::time::timeout(Duration::from_millis(100), ctl_rx.collect::<Vec<_>>())
1040            .await
1041            .context("timeout receiving Control messages")?;
1042
1043        assert_eq!(2, ctl_msgs.len());
1044        assert_eq!(
1045            ctl_msgs[0],
1046            SessionMessage::Request(SegmentRequest::from_iter([(1, [0b01100000].into())]))
1047        );
1048
1049        assert_eq!(
1050            ctl_msgs[1],
1051            SessionMessage::Request(SegmentRequest::from_iter([(1, [0b00100000].into())]))
1052        );
1053
1054        Ok(())
1055    }
1056
1057    #[tokio::test]
1058    async fn ack_state_receiver_must_continue_requesting_missing_frames_and_acknowledge_once_complete()
1059    -> anyhow::Result<()> {
1060        let cfg = AcknowledgementStateConfig {
1061            mode: AcknowledgementMode::Partial,
1062            expected_packet_latency: Duration::from_millis(2),
1063            max_incoming_frame_retries: 3,
1064            acknowledgement_delay: Duration::from_millis(5),
1065            ..Default::default()
1066        };
1067
1068        let mut inspector = FrameInspector(FrameDashMap::with_capacity(10));
1069        let (ctl_tx, ctl_rx) = bounded_sink_channel::<SessionMessage<MTU>>(1024);
1070
1071        let segments = segment(hopr_types::crypto_random::random_bytes::<{ 2 * FRAME_SIZE }>(), MTU, 1)?;
1072
1073        // Segment 1
1074        inspector
1075            .0
1076            .entry(1)
1077            .try_as_vacant()
1078            .ok_or(anyhow::anyhow!("frame 1 must be vacant"))?
1079            .insert(FrameBuilder::from(segments[0].clone()));
1080
1081        let mut state = AcknowledgementState::<MTU>::new("test", cfg);
1082        state.run(SocketComponents {
1083            inspector: inspector.clone().into(),
1084            ctl_tx,
1085        })?;
1086
1087        // Segment 2
1088        state.incoming_segment(&segments[0].id(), (segments.len() as SeqNum).try_into()?)?;
1089
1090        tokio::time::sleep(cfg.expected_packet_latency * 2).await;
1091
1092        inspector
1093            .0
1094            .entry(1)
1095            .try_as_occupied()
1096            .ok_or(anyhow::anyhow!("frame 1 must be occupied"))?
1097            .get_mut()
1098            .add_segment(segments[1].clone())?;
1099
1100        state.incoming_segment(&segments[1].id(), (segments.len() as SeqNum).try_into()?)?;
1101
1102        tokio::time::sleep(cfg.expected_packet_latency * 2).await;
1103
1104        // Segment 3
1105        inspector
1106            .0
1107            .entry(1)
1108            .try_as_occupied()
1109            .ok_or(anyhow::anyhow!("frame 1 must be occupied"))?
1110            .get_mut()
1111            .add_segment(segments[2].clone())?;
1112
1113        state.incoming_segment(&segments[2].id(), (segments.len() as SeqNum).try_into()?)?;
1114        state.frame_complete(1)?;
1115
1116        tokio::time::sleep(cfg.acknowledgement_delay * 2).await;
1117
1118        state.stop()?;
1119
1120        let ctl_msgs = tokio::time::timeout(Duration::from_millis(100), ctl_rx.collect::<Vec<_>>())
1121            .await
1122            .context("timeout receiving Control messages")?;
1123
1124        assert_eq!(3, ctl_msgs.len());
1125        assert_eq!(
1126            ctl_msgs[0],
1127            SessionMessage::Request(SegmentRequest::from_iter([(1, [0b01100000].into())]))
1128        );
1129
1130        assert_eq!(
1131            ctl_msgs[1],
1132            SessionMessage::Request(SegmentRequest::from_iter([(1, [0b00100000].into())]))
1133        );
1134
1135        assert_eq!(ctl_msgs[2], SessionMessage::Acknowledge(vec![1].try_into()?));
1136
1137        Ok(())
1138    }
1139}
hopr_protocol_session/socket/ack_state.rs

hopr_protocol_session/socket/
ack_state.rs
