Quelle  tracking.rs

// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Tracking of received packets and generating ACKs thereof.

use std::{
    cmp::min,
    collections::VecDeque,
    fmt::{self, Display, Formatter},
    time::{Duration, Instant},
};

use enum_map::{Enum, EnumMap};
use enumset::{EnumSet, EnumSetType};
use log::{Level, log_enabled};
use neqo_common::{Buffer, Ecn, MAX_VARINT, qdebug, qtrace, qwarn};
use nss::Epoch;
use smallvec::SmallVec;
use strum::{Display, EnumIter};

use crate::{
    Error, Res, Stats, ecn,
    frame::{FrameEncoder as _, FrameType},
    packet,
    recovery::{self},
    stats::FrameStats,
};

#[derive(Debug, PartialOrd, Ord, EnumSetType, Enum, EnumIter, Display)]
pub enum PacketNumberSpace {
    #[strum(to_string = "in")]
    Initial,
    #[strum(to_string = "hs")]
    Handshake,
    #[strum(to_string = "ap")]
    ApplicationData,
}

impl From<Epoch> for PacketNumberSpace {
    fn from(epoch: Epoch) -> Self {
        match epoch {
            Epoch::Initial => Self::Initial,
            Epoch::Handshake => Self::Handshake,
            Epoch::ApplicationData | Epoch::ZeroRtt => Self::ApplicationData,
        }
    }
}

impl From<PacketNumberSpace> for Epoch {
    fn from(val: PacketNumberSpace) -> Self {
        match val {
            PacketNumberSpace::Initial => Self::Initial,
            PacketNumberSpace::Handshake => Self::Handshake,
            PacketNumberSpace::ApplicationData => Self::ApplicationData,
        }
    }
}

#[expect(clippy::fallible_impl_from, reason = "OK here.")]
impl From<packet::Type> for PacketNumberSpace {
    fn from(pt: packet::Type) -> Self {
        match pt {
            packet::Type::Initial => Self::Initial,
            packet::Type::Handshake => Self::Handshake,
            packet::Type::ZeroRtt | packet::Type::Short => Self::ApplicationData,
            _ => panic!("Attempted to get space from wrong packet type"),
        }
    }
}

pub type PacketNumberSpaceSet = EnumSet<PacketNumberSpace>;

/// `InsertionResult` tracks whether something was inserted for `PacketRange::add()`.
pub enum InsertionResult {
    Largest,
    Smallest,
    NotInserted,
}

#[derive(Clone, Debug, Default)]
pub struct PacketRange {
    largest: packet::Number,
    smallest: packet::Number,
    ack_needed: bool,
}

impl PacketRange {
    /// Make a single packet range.
    pub const fn new(pn: packet::Number) -> Self {
        Self {
            largest: pn,
            smallest: pn,
            ack_needed: true,
        }
    }

    /// Get the number of acknowledged packets in the range.
    pub const fn len(&self) -> u64 {
        self.largest - self.smallest + 1
    }

    /// Returns whether this needs to be sent.
    pub const fn ack_needed(&self) -> bool {
        self.ack_needed
    }

    /// Return whether the given number is in the range.
    pub const fn contains(&self, pn: packet::Number) -> bool {
        (pn >= self.smallest) && (pn <= self.largest)
    }

    /// Maybe add a packet number to the range.  Returns true if it was added
    /// at the small end (which indicates that this might need merging with a
    /// preceding range).
    pub fn add(&mut self, pn: packet::Number) -> InsertionResult {
        assert!(!self.contains(pn));
        // Only insert if this is adjacent the current range.
        if (self.largest + 1) == pn {
            qtrace!("[{self}] Adding largest {pn}");
            self.largest += 1;
            self.ack_needed = true;
            InsertionResult::Largest
        } else if self.smallest == (pn + 1) {
            qtrace!("[{self}] Adding smallest {pn}");
            self.smallest -= 1;
            self.ack_needed = true;
            InsertionResult::Smallest
        } else {
            InsertionResult::NotInserted
        }
    }

    /// Maybe merge a higher-numbered range into this.
    fn merge_larger(&mut self, other: &Self) {
        qdebug!("[{self}] Merging {other}");
        // This only works if they are immediately adjacent.
        assert_eq!(self.largest + 1, other.smallest);

        self.largest = other.largest;
        self.ack_needed = self.ack_needed || other.ack_needed;
    }

    /// When a packet containing the range `other` is acknowledged,
    /// clear the `ack_needed` attribute on this.
    /// Requires that other is equal to this, or a larger range.
    pub const fn acknowledged(&mut self, other: &Self) {
        if (other.smallest <= self.smallest) && (other.largest >= self.largest) {
            self.ack_needed = false;
        }
    }
}

impl Display for PacketRange {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "{}->{}", self.largest, self.smallest)
    }
}

/// The default maximum ACK delay we use locally and advertise to the remote.
pub const DEFAULT_LOCAL_ACK_DELAY: Duration = Duration::from_millis(20);
/// The default maximum ACK delay we assume the remote uses.
///
/// > If this value is absent, a default of 25 milliseconds is assumed.
///
/// <https://datatracker.ietf.org/doc/html/rfc9000#section-18.2>
pub const DEFAULT_REMOTE_ACK_DELAY: Duration = Duration::from_millis(25);
/// The default number of in-order packets we will receive after
/// largest acknowledged without sending an immediate acknowledgment.
pub const DEFAULT_ACK_PACKET_TOLERANCE: packet::Number = 1;
const MAX_TRACKED_RANGES: usize = 32;
const MAX_ACKS_PER_FRAME: usize = 32;

/// A structure that tracks what was included in an ACK.
#[derive(Debug, Clone)]
pub struct AckToken {
    space: PacketNumberSpace,
    ranges: Box<[PacketRange]>,
}

impl AckToken {
    /// Get the space for this token.
    pub const fn space(&self) -> PacketNumberSpace {
        self.space
    }
}

/// A structure that tracks what packets have been received,
/// and what needs acknowledgement for a packet number space.
#[derive(Debug)]
pub struct RecvdPackets {
    space: PacketNumberSpace,
    ranges: VecDeque<PacketRange>,
    /// The packet number of the lowest number packet that we are tracking.
    min_tracked: packet::Number,
    /// The time we got the largest acknowledged.
    largest_pn_time: Option<Instant>,
    /// The time that we should be sending an ACK.
    ack_time: Option<Instant>,
    /// The time we last sent an ACK.
    last_ack_time: Option<Instant>,
    /// The current ACK frequency sequence number.
    ack_frequency_seqno: u64,
    /// The time to delay after receiving the first packet that is
    /// not immediately acknowledged.
    ack_delay: Duration,
    /// The number of ack-eliciting packets that have been received, but
    /// not acknowledged.
    unacknowledged_count: packet::Number,
    /// The number of contiguous packets that can be received without
    /// acknowledging immediately.
    unacknowledged_tolerance: packet::Number,
    /// Whether we are ignoring packets that arrive out of order
    /// for the purposes of generating immediate acknowledgment.
    ignore_order: bool,
    // The counts of different ECN marks that have been received.
    ecn_count: ecn::Count,
}

impl RecvdPackets {
    /// Make a new `RecvdPackets` for the indicated packet number space.
    pub fn new(space: PacketNumberSpace) -> Self {
        Self {
            space,
            ranges: VecDeque::new(),
            min_tracked: 0,
            largest_pn_time: None,
            ack_time: None,
            last_ack_time: None,
            ack_frequency_seqno: 0,
            ack_delay: DEFAULT_LOCAL_ACK_DELAY,
            unacknowledged_count: 0,
            unacknowledged_tolerance: if space == PacketNumberSpace::ApplicationData {
                DEFAULT_ACK_PACKET_TOLERANCE
            } else {
                // ACK more aggressively
                0
            },
            ignore_order: false,
            ecn_count: ecn::Count::default(),
        }
    }

    /// Get the ECN counts.
    pub const fn ecn_marks(&mut self) -> &mut ecn::Count {
        &mut self.ecn_count
    }

    /// Get the time at which the next ACK should be sent.
    pub const fn ack_time(&self) -> Option<Instant> {
        self.ack_time
    }

    /// Update acknowledgment delay parameters.
    pub const fn ack_freq(
        &mut self,
        seqno: u64,
        tolerance: packet::Number,
        delay: Duration,
        ignore_order: bool,
    ) {
        // Yes, this means that we will overwrite values if a sequence number is
        // reused, but that is better than using an `Option<packet::Number>`
        // when it will always be `Some`.
        if seqno >= self.ack_frequency_seqno {
            self.ack_frequency_seqno = seqno;
            self.unacknowledged_tolerance = tolerance;
            self.ack_delay = delay;
            self.ignore_order = ignore_order;
        }
    }

    /// Returns true if an ACK frame should be sent now.
    fn ack_now(&self, now: Instant, rtt: Duration) -> bool {
        // If ack_time is Some, then we have something to acknowledge.
        // In that case, either ack because `now >= ack_time`, or
        // because it is more than an RTT since the last time we sent an ack.
        self.ack_time.is_some_and(|next| {
            next <= now || self.last_ack_time.is_some_and(|last| last + rtt <= now)
        })
    }

    // A simple addition of a packet number to the tracked set.
    // This doesn't do a binary search on the assumption that
    // new packets will generally be added to the start of the list.
    fn add(&mut self, pn: packet::Number) -> Res<()> {
        for i in 0..self.ranges.len() {
            match self.ranges[i].add(pn) {
                InsertionResult::Largest => return Ok(()),
                InsertionResult::Smallest => {
                    // If this was the smallest, it might have filled a gap.
                    let nxt = i + 1;
                    if (nxt < self.ranges.len()) && (pn - 1 == self.ranges[nxt].largest) {
                        let larger = self.ranges.remove(i).ok_or(Error::Internal)?;
                        self.ranges[i].merge_larger(&larger);
                    }
                    return Ok(());
                }
                InsertionResult::NotInserted => {
                    if self.ranges[i].largest < pn {
                        self.ranges.insert(i, PacketRange::new(pn));
                        return Ok(());
                    }
                }
            }
        }
        self.ranges.push_back(PacketRange::new(pn));
        Ok(())
    }

    fn trim_ranges(&mut self, stats: &mut Stats) -> Res<()> {
        // Limit the number of ranges that are tracked to MAX_TRACKED_RANGES.
        if self.ranges.len() > MAX_TRACKED_RANGES {
            let oldest = self.ranges.pop_back().ok_or(Error::Internal)?;
            if oldest.ack_needed {
                qwarn!("[{self}] Dropping unacknowledged ACK range: {oldest}");
                stats.unacked_range_dropped += 1;
            } else {
                qdebug!("[{self}] Drop ACK range: {oldest}");
            }
            self.min_tracked = oldest.largest + 1;
        }
        Ok(())
    }

    /// Add the packet to the tracked set.
    /// Return true if the packet was the largest received so far.
    pub fn set_received(
        &mut self,
        now: Instant,
        pn: packet::Number,
        ack_eliciting: bool,
        stats: &mut Stats,
    ) -> Res<bool> {
        let next_in_order_pn = self.ranges.front().map_or(0, |r| r.largest + 1);
        qtrace!("[{self}] received {pn}, next: {next_in_order_pn}");

        self.add(pn)?;
        self.trim_ranges(stats)?;

        // The new addition was the largest, so update the time we use for calculating ACK delay.
        let largest = if pn >= next_in_order_pn {
            self.largest_pn_time = Some(now);
            true
        } else {
            false
        };

        if ack_eliciting {
            self.unacknowledged_count += 1;

            let immediate_ack = self.space != PacketNumberSpace::ApplicationData
                || (pn != next_in_order_pn && !self.ignore_order)
                || self.unacknowledged_count > self.unacknowledged_tolerance;

            let ack_time = if immediate_ack {
                now
            } else {
                // Note that `ack_delay` can change and that won't take effect if
                // we are waiting on the previous delay timer.
                // If ACK delay increases, we might send an ACK a bit early;
                // if ACK delay decreases, we might send an ACK a bit later.
                // We could use min() here, but change is rare and the size
                // of the change is very small.
                self.ack_time.unwrap_or_else(|| now + self.ack_delay)
            };
            qdebug!("[{self}] Set ACK timer to {ack_time:?}");
            self.ack_time = Some(ack_time);
        }
        Ok(largest)
    }

    /// If we just received a PING frame, we should immediately acknowledge.
    pub fn immediate_ack(&mut self, now: Instant) {
        self.ack_time = Some(now);
        qdebug!("[{self}] immediate_ack at {now:?}");
    }

    /// Check if the packet is a duplicate.
    pub fn is_duplicate(&self, pn: packet::Number) -> bool {
        if pn < self.min_tracked {
            return true;
        }
        self.ranges
            .iter()
            .take_while(|r| pn <= r.largest)
            .any(|r| r.contains(pn))
    }

    /// Mark the given range as having been acknowledged.
    pub fn acknowledged(&mut self, acked: &[PacketRange]) {
        let mut range_iter = self.ranges.iter_mut();
        let mut cur = range_iter.next().expect("should have at least one range");
        for ack in acked {
            while cur.smallest > ack.largest {
                let Some(next) = range_iter.next() else {
                    return;
                };
                cur = next;
            }
            cur.acknowledged(ack);
        }
    }

    /// Length of the worst possible ACK frame, assuming only one range and ECN counts.
    /// Note that this assumes one byte for the type and count of extra ranges.
    pub const USEFUL_ACK_LEN: usize = 1 + 8 + 8 + 1 + 8 + 3 * 8;

    /// Generate an ACK frame for this packet number space.
    ///
    /// Unlike other frame generators this doesn't modify the underlying instance
    /// to track what has been sent. This only clears the delayed ACK timer.
    ///
    /// When sending ACKs, we want to always send the most recent ranges,
    /// even if they have been sent in other packets.
    ///
    /// We don't send ranges that have been acknowledged, but they still need
    /// to be tracked so that duplicates can be detected.
    fn write_frame<B: Buffer>(
        &mut self,
        now: Instant,
        rtt: Duration,
        builder: &mut packet::Builder<B>,
        tokens: &mut recovery::Tokens,
        stats: &mut FrameStats,
    ) {
        // Check that we aren't delaying ACKs.
        if !self.ack_now(now, rtt) {
            return;
        }

        // Drop extra ACK ranges to fit the available space.  Do this based on
        // a worst-case estimate of frame size for simplicity.
        //
        // When congestion limited, ACK-only packets are 255 bytes at most
        // (`recovery::ACK_ONLY_SIZE_LIMIT - 1`).  This results in limiting the
        // ranges to 13 here.
        let max_ranges = if let Some(avail) = builder.remaining().checked_sub(Self::USEFUL_ACK_LEN)
        {
            // Apply a hard maximum to keep plenty of space for other stuff.
            min(1 + (avail / 16), MAX_ACKS_PER_FRAME)
        } else {
            return;
        };

        let ranges = self
            .ranges
            .iter()
            .filter(|r| r.ack_needed())
            .take(max_ranges)
            .cloned()
            .collect::<SmallVec<[_; MAX_TRACKED_RANGES]>>();
        if ranges.is_empty() {
            return;
        }

        let mut iter = ranges.iter();
        let Some(first) = iter.next() else { return };
        stats.largest_acknowledged = first.largest;
        stats.ack += 1;

        let Some(largest_pn_time) = self.largest_pn_time else {
            return;
        };
        let elapsed = now.duration_since(largest_pn_time);
        // We use the default exponent, so delay is in multiples of 8 microseconds.
        let ack_delay = u64::try_from(elapsed.as_micros() / 8).unwrap_or(u64::MAX);
        let ack_delay = min(MAX_VARINT, ack_delay);
        let Ok(extra_ranges) = u64::try_from(ranges.len() - 1) else {
            return;
        };

        builder.encode_frame(
            if self.ecn_count.is_some() {
                FrameType::AckEcn
            } else {
                FrameType::Ack
            },
            |b| {
                b.encode_varint(first.largest);
                b.encode_varint(ack_delay);
                b.encode_varint(extra_ranges); // extra ranges
                b.encode_varint(first.len() - 1); // first range

                let mut last = first.smallest;
                for r in iter {
                    // The difference must be at least 2 because 0-length gaps,
                    // (difference 1) are illegal.
                    b.encode_varint(last - r.largest - 2); // Gap
                    b.encode_varint(r.len() - 1); // Range
                    last = r.smallest;
                }

                if self.ecn_count.is_some() {
                    b.encode_varint(self.ecn_count[Ecn::Ect0]);
                    b.encode_varint(self.ecn_count[Ecn::Ect1]);
                    b.encode_varint(self.ecn_count[Ecn::Ce]);
                }
            },
        );

        // We've sent an ACK, reset the timer.
        self.ack_time = None;
        self.last_ack_time = Some(now);
        self.unacknowledged_count = 0;

        tokens.push(recovery::Token::Ack(AckToken {
            space: self.space,
            ranges: ranges.into_boxed_slice(),
        }));
    }
}

impl Display for RecvdPackets {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "Recvd-{}", self.space)
    }
}

pub struct AckTracker {
    spaces: EnumMap<PacketNumberSpace, Option<RecvdPackets>>,
}

impl AckTracker {
    pub fn drop_space(&mut self, space: PacketNumberSpace) {
        assert_ne!(
            space,
            PacketNumberSpace::ApplicationData,
            "discarding application space"
        );
        if space == PacketNumberSpace::Handshake {
            assert!(self.spaces[PacketNumberSpace::Initial].is_none());
        }
        self.spaces[space].take();
    }

    pub fn get_mut(&mut self, space: PacketNumberSpace) -> Option<&style='color:red'>mut RecvdPackets> {
        self.spaces[space].as_mut()
    }

    pub fn ack_freq(
        &mut self,
        seqno: u64,
        tolerance: packet::Number,
        delay: Duration,
        ignore_order: bool,
    ) {
        // Only ApplicationData ever delays ACK.
        if let Some(space) = self.get_mut(PacketNumberSpace::ApplicationData) {
            space.ack_freq(seqno, tolerance, delay, ignore_order);
        }
    }

    /// Force an ACK to be generated immediately.
    pub fn immediate_ack(&mut self, space: PacketNumberSpace, now: Instant) {
        if let Some(space) = self.get_mut(space) {
            space.immediate_ack(now);
        }
    }

    /// Determine the earliest time that an ACK might be needed.
    pub fn ack_time(&self, now: Instant) -> Option<Instant> {
        if log_enabled!(Level::Trace) {
            for (space, recvd) in &self.spaces {
                if let Some(recvd) = recvd {
                    qtrace!("ack_time for {space} = {:?}", recvd.ack_time());
                }
            }
        }
        if self.spaces[PacketNumberSpace::Initial].is_none()
            && self.spaces[PacketNumberSpace::Handshake].is_none()
            && let Some(recvd) = &self.spaces[PacketNumberSpace::ApplicationData]
        {
            return recvd.ack_time();
        }

        // Ignore any time that is in the past relative to `now`.
        // That is something of a hack, but there are cases where we can't send ACK
        // frames for all spaces, which can mean that one space is stuck in the past.
        // That isn't a problem because we guarantee that earlier spaces will always
        // be able to send ACK frames.
        self.spaces
            .values()
            .flatten()
            .filter_map(|recvd| recvd.ack_time().filter(|t| *t > now))
            .min()
    }

    pub fn acked(&mut self, token: &AckToken) {
        if let Some(space) = self.get_mut(token.space) {
            space.acknowledged(&token.ranges);
        }
    }

    pub(crate) fn write_frame<B: Buffer>(
        &mut self,
        pn_space: PacketNumberSpace,
        now: Instant,
        rtt: Duration,
        builder: &mut packet::Builder<B>,
        tokens: &mut recovery::Tokens,
        stats: &mut FrameStats,
    ) {
        if let Some(space) = self.get_mut(pn_space) {
            space.write_frame(now, rtt, builder, tokens, stats);
        }
    }
}

impl Default for AckTracker {
    fn default() -> Self {
        Self {
            spaces: EnumMap::from_array([
                Some(RecvdPackets::new(PacketNumberSpace::Initial)),
                Some(RecvdPackets::new(PacketNumberSpace::Handshake)),
                Some(RecvdPackets::new(PacketNumberSpace::ApplicationData)),
            ]),
        }
    }
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
mod tests {
    use std::collections::HashSet;

    use neqo_common::{Decoder, Encoder};
    use test_fixture::now;

    use super::{
        AckTracker, Duration, Instant, MAX_TRACKED_RANGES, PacketNumberSpace, PacketRange,
        RecvdPackets,
    };
    use crate::{
        Stats,
        frame::Frame,
        packet,
        recovery::{self},
        stats::FrameStats,
    };

    const RTT: Duration = Duration::from_millis(100);

    fn test_ack_range(pns: &[packet::Number], nranges: usize) {
        let mut rp = RecvdPackets::new(PacketNumberSpace::Initial); // Any space will do.
        assert_eq!(rp.to_string(), "Recvd-in");
        let mut packets = HashSet::new();

        for pn in pns {
            rp.set_received(now(), *pn, true, &mut Stats::default())
                .unwrap();
            packets.insert(*pn);
        }

        assert_eq!(rp.ranges.len(), nranges);

        // Check that all these packets will be detected as duplicates.
        for pn in pns {
            assert!(rp.is_duplicate(*pn));
        }

        // Check that the ranges decrease monotonically and don't overlap.
        let mut iter = rp.ranges.iter();
        let mut last = iter.next().expect("should have at least one");
        for n in iter {
            assert!(n.largest + 1 < last.smallest);
            last = n;
        }

        // Check that the ranges include the right values.
        let mut in_ranges = HashSet::new();
        for range in &rp.ranges {
            for included in range.smallest..=range.largest {
                in_ranges.insert(included);
            }
        }
        assert_eq!(packets, in_ranges);
    }

    #[test]
    fn pn0() {
        test_ack_range(&[0], 1);
    }

    #[test]
    fn pn1() {
        test_ack_range(&[1], 1);
    }

    #[test]
    fn two_ranges() {
        test_ack_range(&[0, 1, 2, 5, 6, 7], 2);
    }

    #[test]
    fn fill_in_range() {
        test_ack_range(&[0, 1, 2, 5, 6, 7, 3, 4], 1);
    }

    #[test]
    fn too_many_ranges() {
        let mut rp = RecvdPackets::new(PacketNumberSpace::Initial); // Any space will do.
        let mut stats = Stats::default();

        // This will add one too many disjoint ranges.
        for i in 0..=MAX_TRACKED_RANGES {
            rp.set_received(now(), (i * 2) as u64, true, &mut stats)
                .unwrap();
        }

        assert_eq!(rp.ranges.len(), MAX_TRACKED_RANGES);
        assert_eq!(rp.ranges.back().unwrap().largest, 2);

        // Even though the range was dropped, we still consider it a duplicate.
        assert!(rp.is_duplicate(0));
        assert!(!rp.is_duplicate(1));
        assert!(rp.is_duplicate(2));
    }

    #[test]
    fn ack_delay() {
        const COUNT: packet::Number = 9;
        const DELAY: Duration = Duration::from_millis(7);
        let mut stats = Stats::default();
        // Only application data packets are delayed.
        let mut rp = RecvdPackets::new(PacketNumberSpace::ApplicationData);
        assert!(rp.ack_time().is_none());
        assert!(!rp.ack_now(now(), RTT));

        rp.ack_freq(0, COUNT, DELAY, false);

        // Some packets won't cause an ACK to be needed.
        for i in 0..COUNT {
            rp.set_received(now(), i, true, &mut stats).unwrap();
            assert_eq!(Some(now() + DELAY), rp.ack_time());
            assert!(!rp.ack_now(now(), RTT));
            assert!(rp.ack_now(now() + DELAY, RTT));
        }

        // Exceeding COUNT will move the ACK time to now.
        rp.set_received(now(), COUNT, true, &mut stats).unwrap();
        assert_eq!(Some(now()), rp.ack_time());
        assert!(rp.ack_now(now(), RTT));
    }

    #[test]
    fn no_ack_delay() {
        let mut stats = Stats::default();
        for space in &[PacketNumberSpace::Initial, PacketNumberSpace::Handshake] {
            let mut rp = RecvdPackets::new(*space);
            assert!(rp.ack_time().is_none());
            assert!(!rp.ack_now(now(), RTT));

            // Any packet in these spaces is acknowledged straight away.
            rp.set_received(now(), 0, true, &mut stats).unwrap();
            assert_eq!(Some(now()), rp.ack_time());
            assert!(rp.ack_now(now(), RTT));
        }
    }

    #[test]
    fn ooo_no_ack_delay_new() {
        let mut stats = Stats::default();
        let mut rp = RecvdPackets::new(PacketNumberSpace::ApplicationData);
        assert!(rp.ack_time().is_none());
        assert!(!rp.ack_now(now(), RTT));

        // Anything other than packet 0 is acknowledged immediately.
        rp.set_received(now(), 1, true, &mut stats).unwrap();
        assert_eq!(Some(now()), rp.ack_time());
        assert!(rp.ack_now(now(), RTT));
    }

    fn write_frame_at(rp: &mut RecvdPackets, now: Instant) {
        let mut builder =
            packet::Builder::short(Encoder::default(), false, None::<&[u8]>, packet::LIMIT);
        let mut stats = FrameStats::default();
        let mut tokens = recovery::Tokens::new();
        rp.write_frame(now, RTT, &mut builder, &mut tokens, &mut stats);
        assert!(!tokens.is_empty());
        assert_eq!(stats.ack, 1);
    }

    fn write_frame(rp: &mut RecvdPackets) {
        write_frame_at(rp, now());
    }

    #[test]
    fn ooo_no_ack_delay_fill() {
        let mut stats = Stats::default();
        let mut rp = RecvdPackets::new(PacketNumberSpace::ApplicationData);
        rp.set_received(now(), 1, true, &mut stats).unwrap();
        write_frame(&mut rp);

        // Filling in behind the largest acknowledged causes immediate ACK.
        rp.set_received(now(), 0, true, &mut stats).unwrap();
        write_frame(&mut rp);

        // Receiving the next packet won't elicit an ACK.
        rp.set_received(now(), 2, true, &mut stats).unwrap();
        assert!(!rp.ack_now(now(), RTT));
    }

    #[test]
    fn immediate_ack_after_rtt() {
        let mut stats = Stats::default();
        let mut rp = RecvdPackets::new(PacketNumberSpace::ApplicationData);
        rp.set_received(now(), 1, true, &mut stats).unwrap();
        write_frame(&mut rp);

        // Filling in behind the largest acknowledged causes immediate ACK.
        rp.set_received(now(), 0, true, &mut stats).unwrap();
        write_frame(&mut rp);

        // A new packet ordinarily doesn't result in an ACK, but this time it does.
        rp.set_received(now() + RTT, 2, true, &mut stats).unwrap();
        write_frame_at(&mut rp, now() + RTT);
    }

    #[test]
    fn ooo_no_ack_delay_threshold_new() {
        let mut stats = Stats::default();
        let mut rp = RecvdPackets::new(PacketNumberSpace::ApplicationData);

        // Set tolerance to 2 and then it takes three packets.
        rp.ack_freq(0, 2, Duration::from_millis(10), true);

        rp.set_received(now(), 1, true, &mut stats).unwrap();
        assert_ne!(Some(now()), rp.ack_time());
        rp.set_received(now(), 2, true, &mut stats).unwrap();
        assert_ne!(Some(now()), rp.ack_time());
        rp.set_received(now(), 3, true, &mut stats).unwrap();
        assert_eq!(Some(now()), rp.ack_time());
    }

    #[test]
    fn ooo_no_ack_delay_threshold_gap() {
        let mut stats = Stats::default();
        let mut rp = RecvdPackets::new(PacketNumberSpace::ApplicationData);
        rp.set_received(now(), 1, true, &mut stats).unwrap();
        write_frame(&mut rp);

        // Set tolerance to 2 and then it takes three packets.
        rp.ack_freq(0, 2, Duration::from_millis(10), true);

        rp.set_received(now(), 3, true, &mut stats).unwrap();
        assert_ne!(Some(now()), rp.ack_time());
        rp.set_received(now(), 4, true, &mut stats).unwrap();
        assert_ne!(Some(now()), rp.ack_time());
        rp.set_received(now(), 5, true, &mut stats).unwrap();
        assert_eq!(Some(now()), rp.ack_time());
    }

    /// Test that an in-order packet that is not ack-eliciting doesn't
    /// increase the number of packets needed to cause an ACK.
    #[test]
    fn non_ack_eliciting_skip() {
        let mut stats = Stats::default();
        let mut rp = RecvdPackets::new(PacketNumberSpace::ApplicationData);
        rp.ack_freq(0, 1, Duration::from_millis(10), true);

        // This should be ignored.
        rp.set_received(now(), 0, false, &mut stats).unwrap();
        assert_ne!(Some(now()), rp.ack_time());
        // Skip 1 (it has no effect).
        rp.set_received(now(), 2, true, &mut stats).unwrap();
        assert_ne!(Some(now()), rp.ack_time());
        rp.set_received(now(), 3, true, &mut stats).unwrap();
        assert_eq!(Some(now()), rp.ack_time());
    }

    /// If a packet that is not ack-eliciting is reordered, that's fine too.
    #[test]
    fn non_ack_eliciting_reorder() {
        let mut stats = Stats::default();
        let mut rp = RecvdPackets::new(PacketNumberSpace::ApplicationData);
        rp.ack_freq(0, 1, Duration::from_millis(10), false);

        // These are out of order, but they are not ack-eliciting.
        rp.set_received(now(), 1, false, &mut stats).unwrap();
        assert_ne!(Some(now()), rp.ack_time());
        rp.set_received(now(), 0, false, &mut stats).unwrap();
        assert_ne!(Some(now()), rp.ack_time());

        // These are in order.
        rp.set_received(now(), 2, true, &mut stats).unwrap();
        assert_ne!(Some(now()), rp.ack_time());
        rp.set_received(now(), 3, true, &mut stats).unwrap();
        assert_eq!(Some(now()), rp.ack_time());
    }

    #[test]
    fn aggregate_ack_time() {
        const DELAY: Duration = Duration::from_millis(17);
        let mut stats = Stats::default();
        let mut tracker = AckTracker::default();
        tracker.ack_freq(0, 1, DELAY, false);
        // This packet won't trigger an ACK.
        tracker
            .get_mut(PacketNumberSpace::Handshake)
            .unwrap()
            .set_received(now(), 0, false, &mut stats)
            .unwrap();
        assert_eq!(None, tracker.ack_time(now()));

        // This should be delayed.
        tracker
            .get_mut(PacketNumberSpace::ApplicationData)
            .unwrap()
            .set_received(now(), 0, true, &mut stats)
            .unwrap();
        assert_eq!(Some(now() + DELAY), tracker.ack_time(now()));

        // This should move the time forward.
        let later = now() + (DELAY / 2);
        tracker
            .get_mut(PacketNumberSpace::Initial)
            .unwrap()
            .set_received(later, 0, true, &mut stats)
            .unwrap();
        assert_eq!(Some(later), tracker.ack_time(now()));
    }

    #[test]
    #[should_panic(expected = "discarding application space")]
    fn drop_app() {
        let mut tracker = AckTracker::default();
        tracker.drop_space(PacketNumberSpace::ApplicationData);
    }

    #[test]
    fn drop_spaces() {
        let mut stats = Stats::default();
        let mut tracker = AckTracker::default();
        let mut builder =
            packet::Builder::short(Encoder::default(), false, None::<&[u8]>, packet::LIMIT);
        tracker
            .get_mut(PacketNumberSpace::Initial)
            .unwrap()
            .set_received(now(), 0, true, &mut stats)
            .unwrap();
        // The reference time for `ack_time` has to be in the past or we filter out the timer.
        assert!(
            tracker
                .ack_time(now().checked_sub(Duration::from_millis(1)).unwrap())
                .is_some()
        );

        let mut tokens = recovery::Tokens::new();
        let mut frame_stats = FrameStats::default();
        tracker.write_frame(
            PacketNumberSpace::Initial,
            now(),
            RTT,
            &mut builder,
            &mut tokens,
            &mut frame_stats,
        );
        assert_eq!(frame_stats.ack, 1);

        // Mark another packet as received so we have cause to send another ACK in that space.
        tracker
            .get_mut(PacketNumberSpace::Initial)
            .unwrap()
            .set_received(now(), 1, true, &mut stats)
            .unwrap();
        assert!(
            tracker
                .ack_time(now().checked_sub(Duration::from_millis(1)).unwrap())
                .is_some()
        );

        // Now drop that space.
        tracker.drop_space(PacketNumberSpace::Initial);

        assert!(tracker.get_mut(PacketNumberSpace::Initial).is_none());
        assert!(
            tracker
                .ack_time(now().checked_sub(Duration::from_millis(1)).unwrap())
                .is_none()
        );
        tracker.write_frame(
            PacketNumberSpace::Initial,
            now(),
            RTT,
            &mut builder,
            &mut tokens,
            &mut frame_stats,
        );
        assert_eq!(frame_stats.ack, 1);
        if let recovery::Token::Ack(tok) = &tokens[0] {
            tracker.acked(tok); // Should be a noop.
        } else {
            panic!("not an ACK token");
        }
    }

    /// ACK delay encodes elapsed microseconds divided by 8 (the default ACK delay exponent).
    #[test]
    fn ack_delay_encoding() {
        let t = now();
        // 16µs → ack_delay = 16/8 = 2.
        let elapsed = Duration::from_micros(16);

        let mut tracker = AckTracker::default();
        tracker
            .get_mut(PacketNumberSpace::Initial)
            .unwrap()
            .set_received(t, 0, true, &mut Stats::default())
            .unwrap();

        let mut builder =
            packet::Builder::short(Encoder::default(), false, None::<&[u8]>, packet::LIMIT);
        let mut stats = FrameStats::default();
        tracker.write_frame(
            PacketNumberSpace::Initial,
            t + elapsed,
            RTT,
            &mut builder,
            &mut recovery::Tokens::new(),
            &mut stats,
        );
        assert_eq!(stats.ack, 1, "ACK frame should have been written");

        // Decode the ACK frame from the builder output and check ack_delay.
        let enc: Encoder = builder.into();
        let bytes = Vec::from(enc);
        // Skip the 1-byte packet header; remainder is the ACK frame.
        let mut dec = Decoder::from(&bytes[1..]);
        let frame = Frame::decode(&mut dec).unwrap();
        let Frame::Ack { ack_delay, .. } = frame else {
            panic!("expected ACK frame, got {frame:?}");
        };
        assert_eq!(ack_delay, 2, "ack_delay must be 16\u{b5}s / 8 = 2");
    }

    #[test]
    fn no_room_for_ack() {
        let mut tracker = AckTracker::default();
        tracker
            .get_mut(PacketNumberSpace::Initial)
            .unwrap()
            .set_received(now(), 0, true, &mut Stats::default())
            .unwrap();
        assert!(
            tracker
                .ack_time(now().checked_sub(Duration::from_millis(1)).unwrap())
                .is_some()
        );

        let mut builder =
            packet::Builder::short(Encoder::default(), false, None::<&[u8]>, packet::LIMIT);
        builder.set_limit(10);

        let mut stats = FrameStats::default();
        tracker.write_frame(
            PacketNumberSpace::Initial,
            now(),
            RTT,
            &mut builder,
            &mut recovery::Tokens::new(),
            &mut stats,
        );
        assert_eq!(stats.ack, 0);
        assert_eq!(builder.len(), 1); // Only the short packet header has been added.
    }

    #[test]
    fn no_room_for_extra_range() {
        let mut stats = Stats::default();
        let mut tracker = AckTracker::default();
        tracker
            .get_mut(PacketNumberSpace::Initial)
            .unwrap()
            .set_received(now(), 0, true, &mut stats)
            .unwrap();
        tracker
            .get_mut(PacketNumberSpace::Initial)
            .unwrap()
            .set_received(now(), 2, true, &mut stats)
            .unwrap();
        assert!(
            tracker
                .ack_time(now().checked_sub(Duration::from_millis(1)).unwrap())
                .is_some()
        );

        let mut builder =
            packet::Builder::short(Encoder::default(), false, None::<&[u8]>, packet::LIMIT);
        // The code pessimistically assumes that each range needs 16 bytes to express.
        // So this won't be enough for a second range.
        builder.set_limit(RecvdPackets::USEFUL_ACK_LEN + 8);

        let mut stats = FrameStats::default();
        tracker.write_frame(
            PacketNumberSpace::Initial,
            now(),
            RTT,
            &mut builder,
            &mut recovery::Tokens::new(),
            &mut stats,
        );
        assert_eq!(stats.ack, 1);

        let mut dec = builder.as_decoder();
        dec.skip(1); // Skip the short header.
        let frame = Frame::decode(&mut dec).unwrap();
        if let Frame::Ack { ack_ranges, .. } = frame {
            assert_eq!(ack_ranges.len(), 0);
        } else {
            panic!("not an ACK!");
        }
    }

    #[test]
    fn ack_time_elapsed() {
        let mut tracker = AckTracker::default();

        // While we have multiple PN spaces, we ignore ACK timers from the past.
        // Send out of order to cause the delayed ack timer to be set to `now()`.
        tracker
            .get_mut(PacketNumberSpace::ApplicationData)
            .unwrap()
            .set_received(now(), 3, true, &mut Stats::default())
            .unwrap();
        assert!(tracker.ack_time(now() + Duration::from_millis(1)).is_none());

        // When we are reduced to one space, that filter is off.
        tracker.drop_space(PacketNumberSpace::Initial);
        tracker.drop_space(PacketNumberSpace::Handshake);
        assert_eq!(
            tracker.ack_time(now() + Duration::from_millis(1)),
            Some(now())
        );
    }

    #[test]
    fn from_packet_type() {
        assert_eq!(
            PacketNumberSpace::from(packet::Type::Initial),
            PacketNumberSpace::Initial
        );
        assert_eq!(
            PacketNumberSpace::from(packet::Type::Handshake),
            PacketNumberSpace::Handshake
        );
        assert_eq!(
            PacketNumberSpace::from(packet::Type::ZeroRtt),
            PacketNumberSpace::ApplicationData
        );
        assert_eq!(
            PacketNumberSpace::from(packet::Type::Short),
            PacketNumberSpace::ApplicationData
        );
        assert!(
            std::panic::catch_unwind(|| {
                PacketNumberSpace::from(packet::Type::VersionNegotiation)
            })
            .is_err()
        );
    }

    #[test]
    fn packet_range_acknowledged() {
        use super::PacketRange;
        let mut r = PacketRange::new(5);
        assert!(r.ack_needed());
        assert_eq!(r.to_string(), "5->5");
        assert_eq!(r.len(), 1);
        assert!(r.contains(5));
        assert!(!r.contains(4));
        r.acknowledged(&PacketRange {
            largest: 10,
            smallest: 0,
            ack_needed: false,
        });
        assert!(!r.ack_needed());

        // Test partial overlap: other.smallest <= self.smallest but other.largest < self.largest.
        // Should NOT clear ack_needed.
        let mut r2 = PacketRange::new(5);
        r2.add(6); // range is now 5..=6
        assert_eq!(r2.to_string(), "6->5");
        assert_eq!(r2.len(), 2);
        r2.acknowledged(&PacketRange {
            largest: 5,
            smallest: 0,
            ack_needed: false,
        });
        assert!(r2.ack_needed()); // Should still need ack
    }

    #[test]
    fn useful_ack_len() {
        // 1 (type) + 8 (largest) + 8 (delay) + 1 (count) + 8 (first range) + 24 (3 ECN counts)
        assert_eq!(RecvdPackets::USEFUL_ACK_LEN, 50);
    }

    #[test]
    fn trim_ranges_increments_stat() {
        // Each dropped range increments the stat by 1.
        let mut rp = RecvdPackets::new(PacketNumberSpace::Initial);
        let mut stats = Stats::default();
        // Fill with MAX_TRACKED_RANGES + 2 disjoint ack-eliciting packets.
        for i in 0..=(MAX_TRACKED_RANGES + 1) {
            rp.set_received(now(), (i * 2) as u64, true, &mut stats)
                .unwrap();
        }
        // Two ranges should have been dropped.
        assert_eq!(stats.unacked_range_dropped, 2);
    }

    #[test]
    fn acknowledged_tracks_duplicates() {
        let mut rp = RecvdPackets::new(PacketNumberSpace::ApplicationData);
        let mut stats = Stats::default();
        // Receive packets 0, 1, 2 — one contiguous range, ACK needed.
        for pn in 0u64..3 {
            rp.set_received(now(), pn, true, &mut stats).unwrap();
        }
        assert!(
            rp.ack_time().is_some(),
            "ACK should be needed before acknowledging"
        );

        // Simulate peer acknowledging our ACK by calling acknowledged.
        let acked = [PacketRange::new(0)]; // covers packet 0
        rp.acknowledged(&acked);
        // Ranges are still tracked for duplicate detection after acknowledgement.
        assert!(rp.is_duplicate(0));
        assert!(rp.is_duplicate(2));
    }
}