// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    counter::Counter,
    event::builder::SlowStartExitCause,
    random,
    recovery::{
        congestion_controller::{self, CongestionController, Publisher},
        cubic::{FastRetransmission::*, State::*},
        hybrid_slow_start::HybridSlowStart,
        pacing::Pacer,
        RttEstimator,
    },
    time::Timestamp,
};
use core::{
    cmp::{max, min},
    time::Duration,
};
#[cfg(not(feature = "std"))]
use num_traits::Float as _;

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.3
//#                 New Path or      +------------+
//#            persistent congestion |   Slow     |
//#        (O)---------------------->|   Start    |
//#                                  +------------+
//#                                        |
//#                                Loss or |
//#                        ECN-CE increase |
//#                                        v
//# +------------+     Loss or       +------------+
//# | Congestion |  ECN-CE increase  |  Recovery  |
//# | Avoidance  |------------------>|   Period   |
//# +------------+                   +------------+
//#           ^                            |
//#           |                            |
//#          +----------------------------+
//#              Acknowledgment of packet
//#                sent during recovery
// This implementation uses Hybrid Slow Start, which allows for
// Slow Start to exit directly to Congestion Avoidance.
#[derive(Clone, Debug, PartialEq, Eq)]
enum State {
    SlowStart,
    Recovery(Timestamp, FastRetransmission),
    CongestionAvoidance(CongestionAvoidanceTiming),
}

impl State {
    /// Returns State::CongestionAvoidance initialized with the given `start_time`
    fn congestion_avoidance(start_time: Timestamp) -> Self {
        Self::CongestionAvoidance(CongestionAvoidanceTiming {
            start_time,
            window_increase_time: start_time,
            app_limited_time: None,
        })
    }

    /// Called when app limited after sending has completed for a round and an ACK has been received.
    fn on_app_limited(&mut self, timestamp: Timestamp) {
        if let CongestionAvoidance(ref mut timing) = self {
            debug_assert!(
                timing
                    .app_limited_time
                    .map_or(true, |app_limited_time| timestamp >= app_limited_time),
                "timestamp must be monotonically increasing"
            );
            debug_assert!(
                timestamp >= timing.window_increase_time,
                "timestamp must not be before the window was last increased"
            );

            timing.app_limited_time = Some(timestamp);
        }
    }

    /// Returns true if the state is `SlowStart`
    fn is_slow_start(&self) -> bool {
        matches!(self, SlowStart)
    }
}

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.2
//# If the congestion window is reduced immediately, a
//# single packet can be sent prior to reduction.  This speeds up loss
//# recovery if the data in the lost packet is retransmitted and is
//# similar to TCP as described in Section 5 of [RFC6675].
#[derive(Clone, Debug, PartialEq, Eq)]
enum FastRetransmission {
    Idle,
    RequiresTransmission,
}

#[derive(Clone, Debug, PartialEq, Eq)]
struct CongestionAvoidanceTiming {
    // The time the congestion avoidance state was entered
    start_time: Timestamp,
    // The time the congestion window was last increased
    window_increase_time: Timestamp,
    // The last time the current congestion window was underutilized
    app_limited_time: Option<Timestamp>,
}

impl CongestionAvoidanceTiming {
    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.1
    //# t is the elapsed time from the beginning of the current congestion avoidance
    fn t(&self, timestamp: Timestamp) -> Duration {
        timestamp - self.start_time
    }

    /// Called when the congestion window is being increased.
    ///
    /// Adjusts the start time by the period from the last window increase to the app limited time
    /// to avoid counting the app limited time period in W_cubic.
    fn on_window_increase(&mut self, timestamp: Timestamp) {
        if let Some(app_limited_time) = self.app_limited_time.take() {
            //= https://www.rfc-editor.org/rfc/rfc8312#section-5.8
            //# CUBIC does not raise its congestion window size if the flow is
            //# currently limited by the application instead of the congestion
            //# window.  In case of long periods when cwnd has not been updated due
            //# to the application rate limit, such as idle periods, t in Eq. 1 MUST
            //# NOT include these periods; otherwise, W_cubic(t) might be very high
            //# after restarting from these periods.

            // Adjust the congestion avoidance start time by the app limited duration
            self.start_time += app_limited_time - self.window_increase_time;
        }

        self.window_increase_time = timestamp;
    }
}

/// A congestion controller that implements "CUBIC for Fast Long-Distance Networks"
/// as specified in <https://tools.ietf.org/html/rfc8312>. The Hybrid Slow Start algorithm
/// is used for determining the slow start threshold.
#[derive(Clone, Debug)]
pub struct CubicCongestionController {
    cubic: Cubic,
    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.8
    //# CUBIC MUST employ a slow-start algorithm, when the cwnd is no more
    //# than ssthresh.

    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.8
    //# Among the slow-start algorithms, CUBIC MAY choose the
    //# standard TCP slow start [RFC5681] in general networks, or the limited
    //# slow start [RFC3742] or hybrid slow start [HR08] for fast and long-
    //# distance networks.
    slow_start: HybridSlowStart,
    pacer: Pacer,
    max_datagram_size: u16,
    congestion_window: f32,
    state: State,
    //= https://www.rfc-editor.org/rfc/rfc9002#appendix-B.2
    //# The sum of the size in bytes of all sent packets
    //# that contain at least one ack-eliciting or PADDING frame and have
    //# not been acknowledged or declared lost.  The size does not include
    //# IP or UDP overhead, but does include the QUIC header and
    //# Authenticated Encryption with Associated Data (AEAD) overhead.
    //# Packets only containing ACK frames do not count toward
    //# bytes_in_flight to ensure congestion control does not impede
    //# congestion feedback.

    // ACK + PADDING packets do not contribute to bytes_in_flight
    // See https://github.com/aws/s2n-quic/pull/1514
    bytes_in_flight: BytesInFlight,
    time_of_last_sent_packet: Option<Timestamp>,
    under_utilized: bool,
    // The highest number of bytes in flight seen when an ACK was received,
    // since the last congestion event.
    bytes_in_flight_hi: BytesInFlight,
}

type BytesInFlight = Counter<u32>;

impl CongestionController for CubicCongestionController {
    type PacketInfo = ();

    #[inline]
    fn congestion_window(&self) -> u32 {
        self.congestion_window as u32
    }

    #[inline]
    fn bytes_in_flight(&self) -> u32 {
        *self.bytes_in_flight
    }

    #[inline]
    fn is_congestion_limited(&self) -> bool {
        let available_congestion_window = self
            .congestion_window()
            .saturating_sub(*self.bytes_in_flight);
        available_congestion_window < self.max_datagram_size as u32
    }

    #[inline]
    fn requires_fast_retransmission(&self) -> bool {
        matches!(self.state, Recovery(_, RequiresTransmission))
    }

    #[inline]
    fn on_packet_sent<Pub: Publisher>(
        &mut self,
        time_sent: Timestamp,
        bytes_sent: usize,
        app_limited: Option<bool>,
        rtt_estimator: &RttEstimator,
        publisher: &mut Pub,
    ) {
        if bytes_sent == 0 {
            // Packet was not congestion controlled
            return;
        }

        self.bytes_in_flight
            .try_add(bytes_sent)
            .expect("bytes sent should not exceed u32::MAX");

        if let Some(app_limited) = app_limited {
            // We check both the given `app_limited` value and is_congestion_window_under_utilized()
            // as is_congestion_window_under_utilized() is more lenient with respect to the utilization
            // of the congestion window than the app_limited check. is_congestion_window_under_utilized()
            // returns true if there are more than 3 MTU's of space left in the cwnd, or less than
            // half the cwnd is utilized in slow start.
            self.under_utilized = app_limited && self.is_congestion_window_under_utilized();
        } else {
            // We don't externally determine `app_limited` in the Initial and Handshake packet
            // spaces, so set under_utilized based on is_congestion_window_under_utilized alone
            self.under_utilized = self.is_congestion_window_under_utilized();
        }

        if let Recovery(recovery_start_time, RequiresTransmission) = self.state {
            // A packet has been sent since we entered recovery (fast retransmission)
            // so flip the state back to idle.
            self.state = Recovery(recovery_start_time, Idle);
        }

        self.time_of_last_sent_packet = Some(time_sent);

        self.pacer.on_packet_sent(
            time_sent,
            bytes_sent,
            rtt_estimator,
            self.congestion_window(),
            self.max_datagram_size,
            self.state.is_slow_start(),
            publisher,
        );
    }

    #[inline]
    fn on_rtt_update<Pub: Publisher>(
        &mut self,
        time_sent: Timestamp,
        now: Timestamp,
        rtt_estimator: &RttEstimator,
        publisher: &mut Pub,
    ) {
        // Update the Slow Start algorithm each time the RTT
        // estimate is updated to find the slow start threshold.
        self.slow_start.on_rtt_update(
            self.congestion_window,
            time_sent,
            self.time_of_last_sent_packet
                .expect("At least one packet must be sent to update RTT"),
            rtt_estimator.latest_rtt(),
        );

        if self.state.is_slow_start() && self.congestion_window >= self.slow_start.threshold {
            publisher.on_slow_start_exited(SlowStartExitCause::Rtt, self.congestion_window());
            //= https://www.rfc-editor.org/rfc/rfc8312#section-4.8
            //# In the case when CUBIC runs the hybrid slow start [HR08], it may exit
            //# the first slow start without incurring any packet loss and thus W_max
            //# is undefined.  In this special case, CUBIC switches to congestion
            //# avoidance and increases its congestion window size using Eq. 1, where
            //# t is the elapsed time since the beginning of the current congestion
            //# avoidance, K is set to 0, and W_max is set to the congestion window
            //# size at the beginning of the current congestion avoidance.
            self.state = State::congestion_avoidance(now);
            self.cubic.on_slow_start_exit(self.congestion_window);
        }
    }

    #[inline]
    fn on_ack<Pub: Publisher>(
        &mut self,
        newest_acked_time_sent: Timestamp,
        bytes_acknowledged: usize,
        _newest_acked_packet_info: Self::PacketInfo,
        rtt_estimator: &RttEstimator,
        _random_generator: &mut dyn random::Generator,
        ack_receive_time: Timestamp,
        publisher: &mut Pub,
    ) {
        self.bytes_in_flight_hi = self.bytes_in_flight_hi.max(self.bytes_in_flight);
        self.bytes_in_flight
            .try_sub(bytes_acknowledged)
            .expect("bytes_acknowledged should not exceed u32::MAX");

        if self.under_utilized {
            self.state.on_app_limited(ack_receive_time);

            //= https://www.rfc-editor.org/rfc/rfc9002#section-7.8
            //# When bytes in flight is smaller than the congestion window and
            //# sending is not pacing limited, the congestion window is
            //# underutilized.  This can happen due to insufficient application data
            //# or flow control limits.  When this occurs, the congestion window
            //# SHOULD NOT be increased in either slow start or congestion avoidance.
            return;
        }

        // Check if this ack causes the controller to exit recovery
        if let State::Recovery(recovery_start_time, _) = self.state {
            if newest_acked_time_sent > recovery_start_time {
                //= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.2
                //# A recovery period ends and the sender enters congestion avoidance
                //# when a packet sent during the recovery period is acknowledged.
                self.state = State::congestion_avoidance(ack_receive_time)
            }
        };

        // The congestion window may continue to grow while app-limited due to pacing
        // interrupting sending while momentarily not app-limited. To avoid the congestion
        // window growing too far beyond bytes in flight, we limit the maximum cwnd to
        // the highest bytes_in_flight value seen since the last congestion event multiplied
        // by a multiplier depending on the current state.
        const SLOW_START_MAX_CWND_MULTIPLIER: f32 = 2.0;
        const MAX_CWND_MULTIPLIER: f32 = 1.5;
        let max_cwnd = match self.state {
            SlowStart => *self.bytes_in_flight_hi as f32 * SLOW_START_MAX_CWND_MULTIPLIER,
            Recovery(_, _) => self.congestion_window,
            CongestionAvoidance(_) => *self.bytes_in_flight_hi as f32 * MAX_CWND_MULTIPLIER,
        }
        .max(self.cubic.minimum_window());

        if self.congestion_window >= max_cwnd {
            // The window is already larger than the max, so we can return early
            return;
        }

        match self.state {
            SlowStart => {
                //= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.1
                //# While a sender is in slow start, the congestion window increases by
                //# the number of bytes acknowledged when each acknowledgment is
                //# processed.  This results in exponential growth of the congestion
                //# window.
                self.congestion_window = (self.congestion_window
                    + self.slow_start.cwnd_increment(bytes_acknowledged))
                .min(max_cwnd);

                if self.congestion_window >= self.slow_start.threshold {
                    // The congestion window has exceeded a previously determined slow start threshold
                    // so transition to congestion avoidance and notify cubic of the slow start exit
                    publisher
                        .on_slow_start_exited(SlowStartExitCause::Other, self.congestion_window());
                    self.state = State::congestion_avoidance(ack_receive_time);
                    self.cubic.on_slow_start_exit(self.congestion_window);
                }
            }
            Recovery(_, _) => {
                // Don't increase the congestion window while in recovery
            }
            CongestionAvoidance(ref mut timing) => {
                timing.on_window_increase(ack_receive_time);

                //= https://www.rfc-editor.org/rfc/rfc8312#section-4.1
                //# t is the elapsed time from the beginning of the current congestion avoidance
                let t = timing.t(ack_receive_time);

                //= https://www.rfc-editor.org/rfc/rfc8312#section-4.1
                //# RTT is the weighted average RTT
                // TODO: Linux Kernel Cubic implementation uses min RTT, possibly
                //      because it is more stable than smoothed_rtt. Other implementations
                //      have followed Linux's choice, so we will as well. The end result is a more
                //      conservative rate of increase of the congestion window. This requires
                //      investigation and testing to evaluate if smoothed_rtt would be a better input.
                let rtt = rtt_estimator.min_rtt();

                self.congestion_avoidance(t, rtt, bytes_acknowledged, max_cwnd);
            }
        };

        debug_assert!(self.congestion_window >= self.cubic.minimum_window());
    }

    #[inline]
    fn on_packet_lost<Pub: Publisher>(
        &mut self,
        lost_bytes: u32,
        _packet_info: Self::PacketInfo,
        persistent_congestion: bool,
        _new_loss_burst: bool,
        _random_generator: &mut dyn random::Generator,
        timestamp: Timestamp,
        publisher: &mut Pub,
    ) {
        debug_assert!(lost_bytes > 0);

        self.bytes_in_flight -= lost_bytes;

        if self.state.is_slow_start() && !persistent_congestion {
            publisher
                .on_slow_start_exited(SlowStartExitCause::PacketLoss, self.congestion_window());
        }

        self.on_congestion_event(timestamp);

        //= https://www.rfc-editor.org/rfc/rfc9002#section-7.6.2
        //# When persistent congestion is declared, the sender's congestion
        //# window MUST be reduced to the minimum congestion window
        //# (kMinimumWindow), similar to a TCP sender's response on an RTO
        //# [RFC5681].
        if persistent_congestion {
            self.congestion_window = self.cubic.minimum_window();
            self.state = State::SlowStart;
            self.cubic.reset();
        }
    }

    #[inline]
    fn on_explicit_congestion<Pub: Publisher>(
        &mut self,
        _ce_count: u64,
        event_time: Timestamp,
        publisher: &mut Pub,
    ) {
        if self.state.is_slow_start() {
            publisher.on_slow_start_exited(SlowStartExitCause::Ecn, self.congestion_window());
        }

        //= https://www.rfc-editor.org/rfc/rfc9002#section-7.1
        //# If a path has been validated to support Explicit Congestion
        //# Notification (ECN) [RFC3168] [RFC8311], QUIC treats a Congestion
        //# Experienced (CE) codepoint in the IP header as a signal of
        //# congestion.
        self.on_congestion_event(event_time);
    }

    //= https://www.rfc-editor.org/rfc/rfc8899#section-3
    //= type=exception
    //= reason=See https://github.com/aws/s2n-quic/issues/959
    //# An update to the PLPMTU (or MPS) MUST NOT increase the congestion
    //# window measured in bytes [RFC4821].

    //= https://www.rfc-editor.org/rfc/rfc8899#section-3
    //# A PL that maintains the congestion window in terms of a limit to
    //# the number of outstanding fixed-size packets SHOULD adapt this
    //# limit to compensate for the size of the actual packets.

    //= https://www.rfc-editor.org/rfc/rfc9002#section-7.2
    //= type=exception
    //= reason=The maximum datagram size remains at the minimum (1200 bytes) during the handshake
    //# If the maximum datagram size is decreased in order to complete the
    //# handshake, the congestion window SHOULD be set to the new initial
    //# congestion window.
    #[inline]
    fn on_mtu_update<Pub: Publisher>(&mut self, max_datagram_size: u16, _publisher: &mut Pub) {
        let old_max_datagram_size = self.max_datagram_size;
        self.max_datagram_size = max_datagram_size;
        self.cubic.max_datagram_size = max_datagram_size;

        self.congestion_window =
            (self.congestion_window / old_max_datagram_size as f32) * max_datagram_size as f32;
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-6.4
    //# When Initial and Handshake packet protection keys are discarded (see
    //# Section 4.9 of [QUIC-TLS]), all packets that were sent with those
    //# keys can no longer be acknowledged because their acknowledgments
    //# cannot be processed.  The sender MUST discard all recovery state
    //# associated with those packets and MUST remove them from the count of
    //# bytes in flight.
    #[inline]
    fn on_packet_discarded<Pub: Publisher>(&mut self, bytes_sent: usize, _publisher: &mut Pub) {
        self.bytes_in_flight
            .try_sub(bytes_sent)
            .expect("bytes sent should not exceed u32::MAX");

        if let Recovery(recovery_start_time, RequiresTransmission) = self.state {
            // If any of the discarded packets were lost, they will no longer be retransmitted
            // so flip the Recovery status back to Idle so it is not waiting for a
            // retransmission that may never come.
            self.state = Recovery(recovery_start_time, Idle);
        }
    }

    #[inline]
    fn earliest_departure_time(&self) -> Option<Timestamp> {
        self.pacer.earliest_departure_time()
    }
}

impl CubicCongestionController {
    // max_datagram_size is the current max_datagram_size, and is
    // expected to be 1200 when the congestion controller is created.
    pub fn new(max_datagram_size: u16) -> Self {
        Self {
            cubic: Cubic::new(max_datagram_size),
            slow_start: HybridSlowStart::new(max_datagram_size),
            pacer: Pacer::default(),
            max_datagram_size,
            congestion_window: CubicCongestionController::initial_window(max_datagram_size) as f32,
            state: SlowStart,
            bytes_in_flight: Counter::new(0),
            time_of_last_sent_packet: None,
            under_utilized: true,
            bytes_in_flight_hi: Counter::new(0),
        }
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-7.2
    //# Endpoints SHOULD use an initial congestion
    //# window of ten times the maximum datagram size (max_datagram_size),
    //# while limiting the window to the larger of 14,720 bytes or twice the
    //# maximum datagram size.

    //= https://www.rfc-editor.org/rfc/rfc9002#section-7.2
    //# If the maximum datagram size changes during the connection, the
    //# initial congestion window SHOULD be recalculated with the new size.
    #[inline]
    fn initial_window(max_datagram_size: u16) -> u32 {
        const INITIAL_WINDOW_LIMIT: u32 = 14720;
        min(
            10 * max_datagram_size as u32,
            max(INITIAL_WINDOW_LIMIT, 2 * max_datagram_size as u32),
        )
    }

    #[inline]
    fn congestion_avoidance(
        &mut self,
        t: Duration,
        rtt: Duration,
        sent_bytes: usize,
        max_cwnd: f32,
    ) {
        let w_cubic = self.cubic.w_cubic(t);
        let w_est = self.cubic.w_est(t, rtt);
        // limit the window increase to half the acked bytes
        // as the Linux implementation of Cubic does.
        let max_cwnd = (self.congestion_window + sent_bytes as f32 / 2.0).min(max_cwnd);

        if w_cubic < w_est {
            // TCP-Friendly Region
            //= https://www.rfc-editor.org/rfc/rfc8312#section-4.2
            //# When receiving an ACK in congestion avoidance (cwnd could be greater than
            //# or less than W_max), CUBIC checks whether W_cubic(t) is less than
            //# W_est(t).  If so, CUBIC is in the TCP-friendly region and cwnd SHOULD
            //# be set to W_est(t) at each reception of an ACK.
            self.congestion_window = self.packets_to_bytes(w_est).min(max_cwnd);
        } else {
            //= https://www.rfc-editor.org/rfc/rfc8312#section-4.1
            //# Upon receiving an ACK during congestion avoidance, CUBIC computes the
            //# window increase rate during the next RTT period using Eq. 1.  It sets
            //# W_cubic(t+RTT) as the candidate target value of the congestion
            //# window

            // Concave Region
            //= https://www.rfc-editor.org/rfc/rfc8312#section-4.3
            //# When receiving an ACK in congestion avoidance, if CUBIC is not in the
            //# TCP-friendly region and cwnd is less than W_max, then CUBIC is in the
            //# concave region.  In this region, cwnd MUST be incremented by
            //# (W_cubic(t+RTT) - cwnd)/cwnd for each received ACK, where
            //# W_cubic(t+RTT) is calculated using Eq. 1.

            // Convex Region
            //# https://www.rfc-editor.org/rfc/rfc8312#section-4.4
            //# When receiving an ACK in congestion avoidance, if CUBIC is not in the
            //# TCP-friendly region and cwnd is larger than or equal to W_max, then
            //# CUBIC is in the convex region.

            //= https://www.rfc-editor.org/rfc/rfc8312#section-4.4
            //# In this region, cwnd MUST be incremented by
            //# (W_cubic(t+RTT) - cwnd)/cwnd for each received ACK, where
            //# W_cubic(t+RTT) is calculated using Eq. 1.

            // The congestion window is adjusted in the same way in the convex and concave regions.
            // A target congestion window is calculated for where the congestion window should be
            // by the end of one RTT. That target is used for calculating the required rate of increase
            // based on where the congestion window currently is. Assuming a full congestion window's
            // worth of packets will be sent and acknowledged within that RTT, we increase the
            // congestion window by this increment for each acknowledgement. As long as all packets
            // are sent and acknowledged by the end of the RTT, the congestion window will reach the
            // target size. Otherwise it will be smaller, reflecting that the network latency is
            // higher than needed to achieve the target window, and thus a smaller congestion window
            // is appropriate.
            let target_congestion_window = self.packets_to_bytes(self.cubic.w_cubic(t + rtt));

            // Decreases in the RTT estimate can cause the congestion window to get ahead of the
            // target. In the case where the congestion window has already exceeded the target,
            // we return without any further adjustment to the window.
            if self.congestion_window >= target_congestion_window {
                return;
            }

            let window_increase_rate =
                (target_congestion_window - self.congestion_window) / self.congestion_window;
            let window_increment = self.packets_to_bytes(window_increase_rate);

            self.congestion_window = (self.congestion_window + window_increment).min(max_cwnd);
        }
    }

    #[inline]
    fn on_congestion_event(&mut self, event_time: Timestamp) {
        // Reset bytes_in_flight_hi
        self.bytes_in_flight_hi = BytesInFlight::new(0);

        // No reaction if already in a recovery period.
        if matches!(self.state, Recovery(_, _)) {
            return;
        }

        // Enter recovery period.

        //= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.1
        //# The sender MUST exit slow start and enter a recovery period when a
        //# packet is lost or when the ECN-CE count reported by its peer
        //# increases.

        //= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.2
        //# If the congestion window is reduced immediately, a
        //# single packet can be sent prior to reduction.  This speeds up loss
        //# recovery if the data in the lost packet is retransmitted and is
        //# similar to TCP as described in Section 5 of [RFC6675].
        self.state = Recovery(event_time, RequiresTransmission);

        //= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.2
        //# Implementations MAY reduce the congestion window immediately upon
        //# entering a recovery period or use other mechanisms, such as
        //# Proportional Rate Reduction [PRR], to reduce the congestion window
        //# more gradually.

        //= https://www.rfc-editor.org/rfc/rfc9002#section-7.2
        //# The minimum congestion window is the smallest value the congestion
        //# window can attain in response to loss, an increase in the peer-
        //# reported ECN-CE count, or persistent congestion.
        self.congestion_window = self.cubic.multiplicative_decrease(self.congestion_window);

        // Update Hybrid Slow Start with the decreased congestion window.
        self.slow_start.on_congestion_event(self.congestion_window);
    }

    #[inline]
    fn packets_to_bytes(&self, cwnd: f32) -> f32 {
        cwnd * self.max_datagram_size as f32
    }

    /// Returns true if the congestion window is under utilized and should not grow larger
    /// without further evidence of the stability of the current window.
    #[inline]
    fn is_congestion_window_under_utilized(&self) -> bool {
        // This value is based on kMaxBurstBytes from Chromium
        // https://source.chromium.org/chromium/chromium/src/+/master:net/third_party/quiche/src/quic/core/congestion_control/tcp_cubic_sender_bytes.cc;l=23;drc=f803516d2656ed829e54b2e819731763ca6cf4d9
        const MAX_BURST_MULTIPLIER: u32 = 3;

        if self.is_congestion_limited() {
            return false;
        }

        // In slow start, allow the congestion window to increase as long as half of it is
        // being used. This allows for the window to increase rapidly.
        if self.state.is_slow_start() && self.bytes_in_flight >= self.congestion_window() / 2 {
            return false;
        }

        // Otherwise allow the window to increase while MAX_BURST_MULTIPLIER packets are available
        // in the window.
        let available_congestion_window = self
            .congestion_window()
            .saturating_sub(*self.bytes_in_flight);
        available_congestion_window > self.max_datagram_size as u32 * MAX_BURST_MULTIPLIER
    }
}

/// Core functions of "CUBIC for Fast Long-Distance Networks" as specified in
/// https://tools.ietf.org/html/rfc8312. The unit of all window sizes is in
/// packets of size max_datagram_size to maintain alignment with the specification.
/// Thus, window sizes should be converted to bytes before applying to the
/// congestion window in the congestion controller.
#[derive(Clone, Debug)]
struct Cubic {
    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.1
    //# W_max is the window size just before the window is
    //# reduced in the last congestion event.
    w_max: f32,
    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.6
    //# a flow remembers the last value of W_max before it
    //# updates W_max for the current congestion event.
    //# Let us call the last value of W_max to be W_last_max.
    w_last_max: f32,
    // k is the time until we expect to reach w_max
    k: Duration,
    max_datagram_size: u16,
}

//= https://www.rfc-editor.org/rfc/rfc8312#section-5.1
//# Based on these observations and our experiments, we find C=0.4
//# gives a good balance between TCP-friendliness and aggressiveness
//# of window increase.  Therefore, C SHOULD be set to 0.4.
const C: f32 = 0.4;

//= https://www.rfc-editor.org/rfc/rfc8312#section-4.5
//# Parameter beta_cubic SHOULD be set to 0.7.
const BETA_CUBIC: f32 = 0.7;

impl Cubic {
    pub fn new(max_datagram_size: u16) -> Self {
        Cubic {
            w_max: 0.0,
            w_last_max: 0.0,
            k: Duration::ZERO,
            max_datagram_size,
        }
    }

    /// Reset to the original state
    #[inline]
    pub fn reset(&mut self) {
        self.w_max = 0.0;
        self.w_last_max = 0.0;
        self.k = Duration::ZERO;
    }

    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.1
    //# CUBIC uses the following window increase function:
    //#
    //#    W_cubic(t) = C*(t-K)^3 + W_max (Eq. 1)
    //#
    //# where C is a constant fixed to determine the aggressiveness of window
    //# increase in high BDP networks, t is the elapsed time from the
    //# beginning of the current congestion avoidance, and K is the time
    //# period that the above function takes to increase the current window
    //# size to W_max if there are no further congestion events and is
    //# calculated using the following equation:
    //#
    //#    K = cubic_root(W_max*(1-beta_cubic)/C) (Eq. 2)
    //#
    //# where beta_cubic is the CUBIC multiplication decrease factor
    #[inline]
    fn w_cubic(&self, t: Duration) -> f32 {
        C * (t.as_secs_f32() - self.k.as_secs_f32()).powi(3) + self.w_max
    }

    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.2
    //# W_est(t) = W_max*beta_cubic +
    //               [3*(1-beta_cubic)/(1+beta_cubic)] * (t/RTT) (Eq. 4)
    #[inline]
    fn w_est(&self, t: Duration, rtt: Duration) -> f32 {
        self.w_max.mul_add(
            BETA_CUBIC,
            (3.0 * (1.0 - BETA_CUBIC) / (1.0 + BETA_CUBIC)) * (t.as_secs_f32() / rtt.as_secs_f32()),
        )
    }

    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.5
    //# When a packet loss is detected by duplicate ACKs or a network
    //# congestion is detected by ECN-Echo ACKs, CUBIC updates its W_max,
    //# cwnd, and ssthresh as follows.  Parameter beta_cubic SHOULD be set to
    //# 0.7.
    //#
    //#    W_max = cwnd;                 // save window size before reduction
    //#    ssthresh = cwnd * beta_cubic; // new slow-start threshold
    //#    ssthresh = max(ssthresh, 2);  // threshold is at least 2 MSS
    //#    cwnd = cwnd * beta_cubic;     // window reduction
    // This does not change the units of the congestion window
    #[inline]
    fn multiplicative_decrease(&mut self, cwnd: f32) -> f32 {
        self.w_max = self.bytes_to_packets(cwnd);

        //= https://www.rfc-editor.org/rfc/rfc8312#section-4.6
        //# To speed up this bandwidth release by
        //# existing flows, the following mechanism called "fast convergence"
        //# SHOULD be implemented.

        //= https://www.rfc-editor.org/rfc/rfc8312#section-4.6
        //# With fast convergence, when a congestion event occurs, before the
        //# window reduction of the congestion window, a flow remembers the last
        //# value of W_max before it updates W_max for the current congestion
        //# event.  Let us call the last value of W_max to be W_last_max.
        //#
        //#    if (W_max < W_last_max){ // should we make room for others
        //#       W_last_max = W_max;             // remember the last W_max
        //#       W_max = W_max*(1.0+beta_cubic)/2.0; // further reduce W_max
        //#    } else {
        //#       W_last_max = W_max              // remember the last W_max
        //#    }
        //#
        //# At a congestion event, if the current value of W_max is less than
        //# W_last_max, this indicates that the saturation point experienced by
        //# this flow is getting reduced because of the change in available
        //# bandwidth.  Then we allow this flow to release more bandwidth by
        //# reducing W_max further.  This action effectively lengthens the time
        //# for this flow to increase its congestion window because the reduced
        //# W_max forces the flow to have the plateau earlier.  This allows more
        //# time for the new flow to catch up to its congestion window size.
        let w_max = self.w_max;
        if w_max < self.w_last_max {
            self.w_max = (w_max * (1.0 + BETA_CUBIC) / 2.0)
                .max(self.bytes_to_packets(self.minimum_window()));
        }
        self.w_last_max = w_max;

        let cwnd_start = (cwnd * BETA_CUBIC).max(self.minimum_window());

        //= https://tools.ietf.org/id/draft-eggert-tcpm-rfc8312bis-01#4.2
        //# _K_ is the time period that the above
        //# function takes to increase the congestion window size at the
        //# beginning of the current congestion avoidance stage to _W_(max)_ if
        //# there are no further congestion events and is calculated using the
        //# following equation:
        //#
        //#                                ________________
        //#                               /W    - cwnd
        //#                           3  /  max       start
        //#                       K = | /  ----------------
        //#                           |/           C
        //#
        //#                                Figure 2
        //#
        //# where _cwnd_(start)_ is the congestion window at the beginning of the
        //# current congestion avoidance stage.
        self.k =
            Duration::from_secs_f32(((self.w_max - self.bytes_to_packets(cwnd_start)) / C).cbrt());

        cwnd_start
    }

    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.8
    //# In the case when CUBIC runs the hybrid slow start [HR08], it may exit
    //# the first slow start without incurring any packet loss and thus W_max
    //# is undefined.  In this special case, CUBIC switches to congestion
    //# avoidance and increases its congestion window size using Eq. 1, where
    //# t is the elapsed time since the beginning of the current congestion
    //# avoidance, K is set to 0, and W_max is set to the congestion window
    //# size at the beginning of the current congestion avoidance.
    #[inline]
    fn on_slow_start_exit(&mut self, cwnd: f32) {
        self.w_max = self.bytes_to_packets(cwnd);

        // We are currently at the w_max, so set k to zero indicating zero
        // seconds to reach the max
        self.k = Duration::from_secs(0);
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-7.2
    //# The minimum congestion window is the smallest value the congestion
    //# window can attain in response to loss, an increase in the peer-
    //# reported ECN-CE count, or persistent congestion.  The RECOMMENDED
    //# value is 2 * max_datagram_size.
    #[inline]
    fn minimum_window(&self) -> f32 {
        2.0 * self.max_datagram_size as f32
    }

    #[inline]
    fn bytes_to_packets(&self, bytes: f32) -> f32 {
        bytes / self.max_datagram_size as f32
    }
}

#[non_exhaustive]
#[derive(Debug, Default)]
pub struct Endpoint {}

impl congestion_controller::Endpoint for Endpoint {
    type CongestionController = CubicCongestionController;

    fn new_congestion_controller(
        &mut self,
        path_info: congestion_controller::PathInfo,
    ) -> Self::CongestionController {
        CubicCongestionController::new(path_info.max_datagram_size)
    }
}

#[cfg(test)]
mod tests;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::time::Timestamp;
use core::time::Duration;
#[cfg(not(feature = "std"))]
use num_traits::Float as _;

/// An implementation of the Hybrid Slow Start algorithm described in
/// "Hybrid Slow Start for High-Bandwidth and Long-Distance Networks"
/// https://pdfs.semanticscholar.org/25e9/ef3f03315782c7f1cbcd31b587857adae7d1.pdf
/// Most of the constants have been updated since this paper. This constants used in
/// this implementation are based on https://github.com/torvalds/linux/blob/net/ipv4/tcp_cubic.c
#[derive(Clone, Debug)]
pub struct HybridSlowStart {
    sample_count: usize,
    last_min_rtt: Option<Duration>,
    cur_min_rtt: Option<Duration>,
    pub(super) threshold: f32,
    max_datagram_size: u16,
    rtt_round_end_time: Option<Timestamp>,
    use_hystart_plus_plus: bool,
    ss_growth_divisor: f32,
    css_count: usize,
    css_baseline_min_rtt: Duration,
    css_threshold: f32,
}

/// Minimum slow start threshold in multiples of the max_datagram_size.
/// Defined as "hystart_low_window" in tcp_cubic.c
const LOW_SSTHRESH: f32 = 16.0;
/// Factor for dividing the RTT to determine the threshold. Defined in tcp_cubic.c (not a constant)
//= https://tools.ietf.org/id/draft-ietf-tcpm-hystartplusplus-04.txt#section-4.2
//#   o  RttThresh = clamp(MIN_RTT_THRESH, lastRoundMinRTT / 8, MAX_RTT_THRESH)
const THRESHOLD_DIVIDEND: u32 = 8;

//= https://tools.ietf.org/id/draft-ietf-tcpm-hystartplusplus-04.txt#section-4.3
//# It is RECOMMENDED that a HyStart++ implementation use the following
//# constants:
//#
//# *  MIN_RTT_THRESH = 4 msec
//#
//# *  MAX_RTT_THRESH = 16 msec
//#
//# *  N_RTT_SAMPLE = 8
//#
//# *  CSS_GROWTH_DIVISOR = 4
//#
//# *  CSS_ROUNDS = 5

/// Number of samples required before determining the slow start threshold.
/// Defined as "HYSTART_MIN_SAMPLES" in tcp_cubic.c
const N_SAMPLING: usize = 8;
/// Minimum increase in delay to consider. Defined as"HYSTART_DELAY_MIN" in tcp_cubic.c
const MIN_DELAY_THRESHOLD: Duration = Duration::from_millis(4);
/// Maximum increase in delay to consider. Defined as"HYSTART_DELAY_MAX" in tcp_cubic.c
const MAX_DELAY_THRESHOLD: Duration = Duration::from_millis(16);
/// Growth devisor for CSS phase
const CSS_GROWTH_DIVISOR: f32 = 4.0;
/// Maximum rounds for CSS phase
const CSS_ROUNDS: usize = 5;
/// environment variable for using hystart++
#[cfg(feature = "std")]
const USE_HYSTART_PLUS_PLUS: &str = "S2N_UNSTABLE_USE_HYSTART_PP";

impl HybridSlowStart {
    /// Constructs a new `HybridSlowStart`. `max_datagram_size` is used for determining
    /// the minimum slow start threshold.
    pub fn new(max_datagram_size: u16) -> Self {
        Self {
            sample_count: 0,
            last_min_rtt: None,
            cur_min_rtt: None,
            //= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.1
            //# A sender begins in slow start because the slow start threshold
            //# is initialized to an infinite value.
            threshold: f32::MAX,
            max_datagram_size,
            rtt_round_end_time: None,
            use_hystart_plus_plus: Self::use_hystart_parameter(),
            ss_growth_divisor: 1.0,
            css_count: 0,
            css_baseline_min_rtt: Duration::ZERO,
            css_threshold: f32::MAX,
        }
    }

    /// Called each time the round trip time estimate is
    /// updated. The algorithm detects if the min RTT over
    /// a number of samples has increased since the last
    /// round of samples and if so will set the slow start
    /// threshold.
    pub fn on_rtt_update(
        &mut self,
        congestion_window: f32,
        time_sent: Timestamp,
        time_of_last_sent_packet: Timestamp,
        rtt: Duration,
    ) {
        let ss_threshold_found = self.threshold < f32::MAX;
        if congestion_window >= self.threshold || (self.use_hystart_plus_plus && ss_threshold_found)
        {
            //= https://tools.ietf.org/id/draft-ietf-tcpm-hystartplusplus-04.txt#section-4.3
            //# An implementation SHOULD use HyStart++ only for the initial slow
            //# start (when ssthresh is at its initial value of arbitrarily high per
            //# [RFC5681]) and fall back to using traditional slow start for the
            //# remainder of the connection lifetime.
            return;
        }

        // An RTT round is over when a packet that was sent after the packet that
        // started the RTT round (or the starting packet itself) is acknowledged
        let rtt_round_is_over = self
            .rtt_round_end_time
            .map_or(true, |end_time| time_sent >= end_time);

        if rtt_round_is_over {
            // Start a new round and save the previous min RTT
            self.last_min_rtt = self.cur_min_rtt;
            self.cur_min_rtt = None;
            self.sample_count = 0;
            // End this round when packets sent after the current last sent packet
            // start getting acknowledged.
            self.rtt_round_end_time = Some(time_of_last_sent_packet);
        }

        if self.sample_count < N_SAMPLING {
            // Sample the delay, saving the minimum
            self.cur_min_rtt = Some(rtt.min(self.cur_min_rtt.unwrap_or(rtt)));
        }

        self.sample_count += 1;

        // We've gathered enough samples and there have been at least 2 RTT rounds
        // to compare, so check if the delay has increased between the rounds
        if let (N_SAMPLING, Some(last_min_rtt), Some(cur_min_rtt)) =
            (self.sample_count, self.last_min_rtt, self.cur_min_rtt)
        {
            if congestion_window >= self.css_threshold {
                self.css_count += 1;
                if cur_min_rtt < self.css_baseline_min_rtt {
                    // resume slow start
                    self.css_threshold = self.threshold;
                    self.ss_growth_divisor = 1.0;
                    self.css_count = 0;
                }
                if self.css_count >= CSS_ROUNDS {
                    // exit slow start phase
                    self.threshold = congestion_window;
                    self.css_threshold = f32::MAX;
                    self.ss_growth_divisor = 1.0;
                }
            } else {
                let threshold = last_min_rtt / THRESHOLD_DIVIDEND;
                // Clamp n to the min and max thresholds
                let threshold = threshold.min(MAX_DELAY_THRESHOLD).max(MIN_DELAY_THRESHOLD);
                let delay_increase_is_over_threshold = cur_min_rtt >= last_min_rtt + threshold;
                let congestion_window_is_above_minimum = congestion_window >= self.low_ssthresh();

                if self.use_hystart_plus_plus {
                    // if delay is beyond threshold, go into css phase
                    if delay_increase_is_over_threshold {
                        self.css_threshold = congestion_window;
                        self.css_baseline_min_rtt = cur_min_rtt;
                        self.ss_growth_divisor = CSS_GROWTH_DIVISOR;
                        self.css_count = 0;
                    }
                } else if delay_increase_is_over_threshold && congestion_window_is_above_minimum {
                    self.threshold = congestion_window;
                }
            }
        }
    }

    /// return cwnd increment during slow start phase
    /// should be called from on_packet_ack
    pub fn cwnd_increment(&self, sent_bytes: usize) -> f32 {
        if cfg!(debug_assertions) && !self.use_hystart_plus_plus {
            assert!((self.ss_growth_divisor - 1.0).abs() < f32::EPSILON);
        }
        (sent_bytes as f32) / self.ss_growth_divisor
    }

    /// Called when a congestion event is experienced. Sets the
    /// slow start threshold to the minimum of the Hybrid Slow Start threshold
    /// and the given congestion window. This will ensure we exit slow start
    /// early enough to avoid further congestion.
    pub fn on_congestion_event(&mut self, ssthresh: f32) {
        self.threshold = self.threshold.min(ssthresh).max(self.low_ssthresh());
        self.ss_growth_divisor = 1.0;
        self.css_threshold = f32::MAX;
    }

    fn low_ssthresh(&self) -> f32 {
        LOW_SSTHRESH * self.max_datagram_size as f32
    }

    #[cfg(feature = "std")]
    fn use_hystart_parameter() -> bool {
        use once_cell::sync::OnceCell;
        static USE_HYSTART_PP: OnceCell<bool> = OnceCell::new();
        *USE_HYSTART_PP.get_or_init(|| std::env::var(USE_HYSTART_PLUS_PLUS).is_ok())
    }

    #[cfg(not(feature = "std"))]
    fn use_hystart_parameter() -> bool {
        false
    }
}

#[cfg(test)]
mod test {
    use crate::{
        assert_delta,
        recovery::hybrid_slow_start::HybridSlowStart,
        time::{Clock, NoopClock},
    };
    use core::time::Duration;

    #[test]
    fn on_congestion_event() {
        let mut slow_start = HybridSlowStart::new(10);
        slow_start.threshold = 501.0;

        // Setting a threshold lower than the current threshold
        // will override the current threshold
        slow_start.on_congestion_event(500.0);
        assert_delta!(slow_start.threshold, 500.0, 0.001);

        slow_start.threshold = 501.0;

        // Setting a threshold higher than the current threshold
        // keeps the current threshold intact
        slow_start.on_congestion_event(502.0);
        assert_delta!(slow_start.threshold, 501.0, 0.001);

        slow_start.threshold = 501.0;

        // Setting a threshold lower than low_ssthresh
        // sets the threshold to low_ssthresh
        slow_start.on_congestion_event(slow_start.low_ssthresh() - 1.0);
        assert_delta!(slow_start.threshold, slow_start.low_ssthresh(), 0.001);
    }

    #[test]
    fn on_rtt_update_above_threshold() {
        let mut slow_start = HybridSlowStart::new(10);
        let time_zero = NoopClock.get_time();
        slow_start.threshold = 500.0;

        assert_eq!(slow_start.sample_count, 0);
        slow_start.on_rtt_update(750.0, time_zero, time_zero, Duration::from_secs(1));

        assert_delta!(slow_start.threshold, 500.0, 0.001);
        assert_eq!(slow_start.sample_count, 0);
    }

    #[test]
    fn on_rtt_update() {
        let mut slow_start = HybridSlowStart::new(10);

        assert_eq!(slow_start.sample_count, 0);

        let time_zero = NoopClock.get_time() + Duration::from_secs(10);

        // -- Round 1 --

        // t=0-9: Send packet #1-10
        let time_of_last_sent_packet = time_zero + Duration::from_millis(9);

        // t=10: Acknowledge packets #1-7 all with RTT 200
        for i in 0..=6 {
            slow_start.on_rtt_update(
                1000.0,
                time_zero + Duration::from_millis(i),
                time_of_last_sent_packet,
                Duration::from_millis(200),
            );
            assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(200)));
        }

        // This round will end when packet #10 is acknowledged
        assert_eq!(
            slow_start.rtt_round_end_time,
            Some(time_of_last_sent_packet)
        );

        // t=11: Acknowledge packet #8 with RTT 100
        slow_start.on_rtt_update(
            1000.0,
            time_zero + Duration::from_millis(7),
            time_of_last_sent_packet,
            Duration::from_millis(100),
        );

        // The current minimum should now be 100
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(100)));

        // t=11: Acknowledge packet #9 with RTT 50
        slow_start.on_rtt_update(
            1000.0,
            time_zero + Duration::from_millis(8),
            time_of_last_sent_packet,
            Duration::from_millis(50),
        );

        // The current minimum is still 100 because we've already collected N_SAMPLING samples
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(100)));

        // -- Round 2 --

        // t=20-29: Send packet #11-20
        let time_of_last_sent_packet = time_zero + Duration::from_millis(29);

        // t=30: Acknowledge packet #10 with RTT 400, ending the first round and starting the second
        slow_start.on_rtt_update(
            1000.0,
            time_zero + Duration::from_millis(9),
            time_of_last_sent_packet,
            Duration::from_millis(400),
        );

        // The last minimum is saved in last_min_rtt
        assert_eq!(slow_start.last_min_rtt, Some(Duration::from_millis(100)));
        // The current minimum is now 400 because we've started a new round
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(400)));
        // This round will end when packet #20 is acknowledged
        assert_eq!(
            slow_start.rtt_round_end_time,
            Some(time_of_last_sent_packet)
        );

        // t=31: Acknowledge packets #11-16 all with RTT 500
        for i in 20..=25 {
            slow_start.on_rtt_update(
                1000.0,
                time_zero + Duration::from_millis(i),
                time_of_last_sent_packet,
                Duration::from_millis(500),
            );
            assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(400)));
        }

        // t=32: Acknowledge packet #17, the 8th sample in Round 2
        slow_start.on_rtt_update(
            2000.0,
            time_zero + Duration::from_millis(27),
            time_of_last_sent_packet,
            Duration::from_millis(112),
        );

        assert_eq!(slow_start.last_min_rtt, Some(Duration::from_millis(100)));
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(112)));
        // The last ack was the 8th sample, but since the min rtt increased below the threshold,
        // the slow start threshold remains the same
        assert_delta!(slow_start.threshold, f32::MAX, 0.001);

        // -- Round 3 --

        // t=40-49: Send packet #21-30
        let time_of_last_sent_packet = time_zero + Duration::from_millis(49);

        // t=50: Acknowledge packets 21-27
        for i in 40..=46 {
            slow_start.on_rtt_update(
                1000.0,
                time_zero + Duration::from_millis(i),
                time_of_last_sent_packet,
                Duration::from_millis(500),
            );
            assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(500)));
        }
        // t=51: Acknowledge packet 28, the 8th sample in Round 3
        slow_start.on_rtt_update(
            5000.0,
            time_zero + Duration::from_millis(38),
            time_of_last_sent_packet,
            Duration::from_millis(126),
        );

        assert_eq!(slow_start.last_min_rtt, Some(Duration::from_millis(112)));
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(126)));
        // The last ack was the 8th sample, and since the min rtt increased above the threshold,
        // the slow start threshold was set to the current congestion window
        assert_delta!(slow_start.threshold, 5000.0, 0.001);
    }

    #[test]
    fn on_rtt_update_with_hystartplus_1() {
        let mut slow_start = HybridSlowStart::new(10);
        // use hystart++
        slow_start.use_hystart_plus_plus = true;

        assert_eq!(slow_start.sample_count, 0);

        let time_zero = NoopClock.get_time() + Duration::from_secs(10);

        // -- Round 1 --

        // t=0-9: Send packet #1-10
        let time_of_last_sent_packet = time_zero + Duration::from_millis(9);

        // t=10: Acknowledge packets #1-7 all with RTT 200
        for i in 0..=6 {
            slow_start.on_rtt_update(
                1000.0,
                time_zero + Duration::from_millis(i),
                time_of_last_sent_packet,
                Duration::from_millis(200),
            );
            assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(200)));
        }

        // This round will end when packet #10 is acknowledged
        assert_eq!(
            slow_start.rtt_round_end_time,
            Some(time_of_last_sent_packet)
        );

        // t=11: Acknowledge packet #8 with RTT 100
        slow_start.on_rtt_update(
            1000.0,
            time_zero + Duration::from_millis(7),
            time_of_last_sent_packet,
            Duration::from_millis(100),
        );

        // The current minimum should now be 100
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(100)));

        // t=11: Acknowledge packet #9 with RTT 50
        slow_start.on_rtt_update(
            1000.0,
            time_zero + Duration::from_millis(8),
            time_of_last_sent_packet,
            Duration::from_millis(50),
        );

        // The current minimum is still 100 because we've already collected N_SAMPLING samples
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(100)));

        // -- Round 2 --

        // t=20-29: Send packet #11-20
        let time_of_last_sent_packet = time_zero + Duration::from_millis(29);

        // t=30: Acknowledge packet #10 with RTT 400, ending the first round and starting the second
        slow_start.on_rtt_update(
            1000.0,
            time_zero + Duration::from_millis(9),
            time_of_last_sent_packet,
            Duration::from_millis(400),
        );

        // The last minimum is saved in last_min_rtt
        assert_eq!(slow_start.last_min_rtt, Some(Duration::from_millis(100)));
        // The current minimum is now 400 because we've started a new round
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(400)));
        // This round will end when packet #20 is acknowledged
        assert_eq!(
            slow_start.rtt_round_end_time,
            Some(time_of_last_sent_packet)
        );

        // t=31: Acknowledge packets #11-16 all with RTT 500
        for i in 20..=25 {
            slow_start.on_rtt_update(
                1000.0,
                time_zero + Duration::from_millis(i),
                time_of_last_sent_packet,
                Duration::from_millis(500),
            );
            assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(400)));
        }

        // t=32: Acknowledge packet #17, the 8th sample in Round 2
        slow_start.on_rtt_update(
            2000.0,
            time_zero + Duration::from_millis(27),
            time_of_last_sent_packet,
            Duration::from_millis(112),
        );

        assert_eq!(slow_start.last_min_rtt, Some(Duration::from_millis(100)));
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(112)));
        // The last ack was the 8th sample, but since the min rtt increased below the threshold,
        // the slow start threshold remains the same
        assert_delta!(slow_start.threshold, f32::MAX, 0.001);

        // -- Round 3 --

        // t=40-49: Send packet #21-30
        let time_of_last_sent_packet = time_zero + Duration::from_millis(49);

        // t=50: Acknowledge packets 21-27
        for i in 40..=46 {
            slow_start.on_rtt_update(
                1000.0,
                time_zero + Duration::from_millis(i),
                time_of_last_sent_packet,
                Duration::from_millis(500),
            );
            assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(500)));
        }
        // t=51: Acknowledge packet 28, the 8th sample in Round 3
        slow_start.on_rtt_update(
            5000.0,
            time_zero + Duration::from_millis(38),
            time_of_last_sent_packet,
            Duration::from_millis(126),
        );

        assert_eq!(slow_start.last_min_rtt, Some(Duration::from_millis(112)));
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(126)));
        // The last ack was the 8th sample, and since the min rtt increased above the threshold,
        // the slow start cssthreshold was set to the current congestion window
        // Also, cssbaseline is cur_min_rtt and ss_growth_divisor is 4.0
        assert_delta!(slow_start.css_threshold, 5000.0, 0.001);
        assert_eq!(slow_start.css_baseline_min_rtt, Duration::from_millis(126));
        assert_delta!(slow_start.ss_growth_divisor, 4.0, 0.001);

        // -- Round 4 --
        // t=50-59: Send packet #31-40
        // first round for CSS
        let time_of_last_sent_packet = time_zero + Duration::from_millis(69);

        // t=60: Acknowledge packets 31-37
        for i in 60..=66 {
            slow_start.on_rtt_update(
                1000.0,
                time_zero + Duration::from_millis(i),
                time_of_last_sent_packet,
                Duration::from_millis(500),
            );
            assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(500)));
        }

        // t=51: Acknowledge packet 38, the 8th sample in Round 4
        slow_start.on_rtt_update(
            5000.0,
            time_zero + Duration::from_millis(38),
            time_of_last_sent_packet,
            Duration::from_millis(130),
        );

        let cwnd_increment = slow_start.cwnd_increment(1000);

        assert_delta!(cwnd_increment, 250.0, 0.001);
        assert_eq!(slow_start.last_min_rtt, Some(Duration::from_millis(126)));
        assert_eq!(slow_start.css_baseline_min_rtt, Duration::from_millis(126));
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(130)));
        // because the cur_min_rtt is bigger than css_baseline_min_rtt, CSS phase should continue
        assert_eq!(slow_start.css_count, 1);
    }

    #[test]
    fn on_rtt_update_with_hystartplus_2() {
        let mut slow_start = HybridSlowStart::new(10);
        // use hystart++
        slow_start.use_hystart_plus_plus = true;

        // emulate Round 1 and Round 2
        let time_zero = NoopClock.get_time() + Duration::from_secs(10);
        slow_start.cur_min_rtt = Some(Duration::from_millis(112));

        // -- Round 3 --

        // t=40-49: Send packet #21-30
        let time_of_last_sent_packet = time_zero + Duration::from_millis(49);

        // t=50: Acknowledge packets 21-27
        for i in 40..=46 {
            slow_start.on_rtt_update(
                1000.0,
                time_zero + Duration::from_millis(i),
                time_of_last_sent_packet,
                Duration::from_millis(500),
            );
            assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(500)));
        }
        // t=51: Acknowledge packet 28, the 8th sample in Round 3
        slow_start.on_rtt_update(
            5000.0,
            time_zero + Duration::from_millis(38),
            time_of_last_sent_packet,
            Duration::from_millis(126),
        );

        //lassert_eq!(slow_start.last_min_rtt, Some(Duration::from_millis(112)));
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(126)));
        // The last ack was the 8th sample, and since the min rtt increased above the threshold,
        // the slow start cssthreshold was set to the current congestion window
        // Also, cssbaseline is cur_min_rtt and ss_growth_divisor is 4.0
        assert_delta!(slow_start.css_threshold, 5000.0, 0.001);
        assert_eq!(slow_start.css_baseline_min_rtt, Duration::from_millis(126));
        assert_delta!(slow_start.ss_growth_divisor, 4.0, 0.001);

        // -- Round 4 --
        // t=50-59: Send packet #31-40
        // first round for CSS
        let time_of_last_sent_packet = time_zero + Duration::from_millis(69);

        // t=60: Acknowledge packets 31-37
        for i in 60..=66 {
            slow_start.on_rtt_update(
                1000.0,
                time_zero + Duration::from_millis(i),
                time_of_last_sent_packet,
                Duration::from_millis(500),
            );
            assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(500)));
        }

        // t=51: Acknowledge packet 38, the 8th sample in Round 4
        slow_start.on_rtt_update(
            5000.0,
            time_zero + Duration::from_millis(38),
            time_of_last_sent_packet,
            Duration::from_millis(125),
        );

        assert_eq!(slow_start.last_min_rtt, Some(Duration::from_millis(126)));
        assert_eq!(slow_start.css_baseline_min_rtt, Duration::from_millis(126));
        assert_eq!(slow_start.cur_min_rtt, Some(Duration::from_millis(125)));
        // because the cur_min_rtt is smaller than css_baseline_min_rtt, should resume slow start
        assert_delta!(slow_start.ss_growth_divisor, 1.0, 0.001);
        assert_delta!(slow_start.css_threshold, f32::MAX, 0.001);
        assert_eq!(slow_start.css_count, 0);
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

pub use congestion_controller::CongestionController;
pub use cubic::CubicCongestionController;
pub use rtt_estimator::*;
pub use sent_packets::*;

pub mod bandwidth;
pub mod bbr;
pub mod congestion_controller;
pub mod cubic;
mod hybrid_slow_start;
mod pacing;
mod rtt_estimator;
mod sent_packets;

#[cfg(test)]
mod simulation;

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.7
//# Senders SHOULD limit bursts to the initial congestion window; see
//# Section 7.2.

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.2
//# Endpoints SHOULD use an initial congestion
//# window of ten times the maximum datagram size (max_datagram_size),
//# while limiting the window to the larger of 14,720 bytes or twice the
//# maximum datagram size.

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.7
//= type=TODO
//= feature=Packet pacing
//= tracking-issue=1073
//# A sender with knowledge that the network path to the
//# receiver can absorb larger bursts MAY use a higher limit.
pub const MAX_BURST_PACKETS: u32 = 10;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    counter::{Counter, Saturating},
    recovery::{
        bandwidth::Bandwidth, congestion_controller::Publisher, RttEstimator, MAX_BURST_PACKETS,
    },
    time::{Duration, Timestamp},
};
use num_rational::Ratio;

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.7
//# Using a value for "N" that is small, but at least 1 (for example, 1.25) ensures
//# that variations in RTT do not result in underutilization of the congestion window.
const N: Ratio<u64> = Ratio::new_raw(5, 4); // 5/4 = 1.25

// In Slow Start, the congestion window grows rapidly, so there is a higher likelihood the congestion
// window may be underutilized due to pacing. To prevent that, we use a higher value for `N` while
// in slow start, as done in Linux:
// https://github.com/torvalds/linux/blob/fc02cb2b37fe2cbf1d3334b9f0f0eab9431766c4/net/ipv4/tcp_input.c#L905-L906
const SLOW_START_N: Ratio<u64> = Ratio::new_raw(2, 1); // 2/1 = 2.00

// Jim Roskind demonstrated the second packet sent on a path has a higher probability of loss due to
// network routers being busy setting up routing tables triggered by the first packet. Setting this
// value to a Duration greater than zero will introduce that delay into the second packet.
// See https://www.ietf.org/proceedings/88/slides/slides-88-tsvarea-10.pdf
// TODO: Determine an appropriate value for this that balances improvements to 2nd packet loss and delay
pub const INITIAL_INTERVAL: Duration = Duration::from_millis(0);

/// low RTT networks should not be using pacing since it'll take longer to wake up from
/// a timer than it would to deliver a packet
pub const MINIMUM_PACING_RTT: Duration = Duration::from_millis(2);

/// A packet pacer that returns departure times that evenly distribute bursts of packets over time
#[derive(Clone, Debug, Default)]
pub struct Pacer {
    // The capacity of the current departure time slot
    capacity: Counter<u32, Saturating>,
    // The time the next packet should be transmitted
    next_packet_departure_time: Option<Timestamp>,
}

impl Pacer {
    /// Called when each packet has been written
    #[allow(clippy::too_many_arguments)]
    #[inline]
    pub fn on_packet_sent<Pub: Publisher>(
        &mut self,
        now: Timestamp,
        bytes_sent: usize,
        rtt_estimator: &RttEstimator,
        congestion_window: u32,
        max_datagram_size: u16,
        slow_start: bool,
        publisher: &mut Pub,
    ) {
        if rtt_estimator.smoothed_rtt() < MINIMUM_PACING_RTT {
            return;
        }

        if self.capacity == 0 {
            if let Some(next_packet_departure_time) = self.next_packet_departure_time {
                let interval = Self::interval(
                    rtt_estimator.smoothed_rtt(),
                    congestion_window,
                    max_datagram_size,
                    slow_start,
                    publisher,
                );
                self.next_packet_departure_time =
                    Some((next_packet_departure_time + interval).max(now));
            } else {
                self.next_packet_departure_time = Some(now + INITIAL_INTERVAL);
            }
            self.capacity = Counter::new(MAX_BURST_PACKETS * max_datagram_size as u32);
        }

        self.capacity -= bytes_sent as u32;
    }

    /// Returns the earliest time that a packet may be transmitted.
    ///
    /// If the time is in the past or is `None`, the packet should be transmitted immediately.
    pub fn earliest_departure_time(&self) -> Option<Timestamp> {
        self.next_packet_departure_time
    }

    // Recalculate the interval between bursts of paced packets
    #[inline]
    fn interval<Pub: Publisher>(
        rtt: Duration,
        congestion_window: u32,
        max_datagram_size: u16,
        slow_start: bool,
        publisher: &mut Pub,
    ) -> Duration {
        debug_assert_ne!(congestion_window, 0);

        let n = if slow_start { SLOW_START_N } else { N };

        //= https://www.rfc-editor.org/rfc/rfc9002#section-7.7
        //# A perfectly paced sender spreads packets exactly evenly over time.
        //# For a window-based congestion controller, such as the one in this
        //# document, that rate can be computed by averaging the congestion
        //# window over the RTT. Expressed as a rate in units of bytes per time,
        //# where congestion_window is in bytes:
        //#
        //# rate = N * congestion_window / smoothed_rtt
        let pacing_rate = Bandwidth::new(congestion_window as u64, rtt) * n;

        // `MAX_BURST_PACKETS` is incorporated into the formula since we are trying to spread
        // bursts of packets evenly over time.
        let packet_size = MAX_BURST_PACKETS * max_datagram_size as u32;

        publisher.on_pacing_rate_updated(pacing_rate, packet_size, n);

        packet_size as u64 / pacing_rate
    }
}

#[cfg(test)]
mod tests;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    packet::number::PacketNumberSpace, time::Timestamp, transport::parameters::MaxAckDelay,
};
use core::{
    cmp::{max, min},
    time::Duration,
};

//= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.2
//# When no previous RTT is available, the initial RTT
//# SHOULD be set to 333 milliseconds.  This results in handshakes
//# starting with a PTO of 1 second, as recommended for TCP's initial
//# RTO; see Section 2 of [RFC6298].
pub const DEFAULT_INITIAL_RTT: Duration = Duration::from_millis(333);
const ZERO_DURATION: Duration = Duration::from_millis(0);

//= https://www.rfc-editor.org/rfc/rfc9002#section-6.1.2
//# The RECOMMENDED value of the timer granularity (kGranularity) is 1 millisecond.
pub const K_GRANULARITY: Duration = Duration::from_millis(1);

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.6.1
//# The RECOMMENDED value for kPersistentCongestionThreshold is 3, which
//# results in behavior that is approximately equivalent to a TCP sender
//# declaring an RTO after two TLPs.
const K_PERSISTENT_CONGESTION_THRESHOLD: u64 = 3;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct RttEstimator {
    /// Latest RTT sample
    latest_rtt: Duration,
    /// The minimum value observed over the lifetime of the connection
    min_rtt: Duration,
    /// An exponentially-weighted moving average
    smoothed_rtt: Duration,
    /// The variance in the observed RTT samples
    rttvar: Duration,
    /// The maximum amount of time by which the receiver intends to delay acknowledgments for
    /// packets in the ApplicationData packet number space. The actual ack_delay in a received
    /// ACK frame may be larger due to late timers, reordering, or lost ACK frames.
    max_ack_delay: Duration,
    /// The time that the first RTT sample was obtained
    first_rtt_sample: Option<Timestamp>,
}

impl Default for RttEstimator {
    fn default() -> Self {
        RttEstimator::new(Duration::ZERO)
    }
}

impl RttEstimator {
    /// Creates a new RTT Estimator with default initial values using the given `max_ack_delay`.
    pub fn new(max_ack_delay: Duration) -> Self {
        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
        //# Before any RTT samples are available for a new path or when the
        //# estimator is reset, the estimator is initialized using the initial RTT;
        //# see Section 6.2.2.
        //#
        //# smoothed_rtt and rttvar are initialized as follows, where kInitialRtt
        //# contains the initial RTT value:
        //
        //# smoothed_rtt = kInitialRtt
        //# rttvar = kInitialRtt / 2
        let smoothed_rtt = DEFAULT_INITIAL_RTT;
        let rttvar = DEFAULT_INITIAL_RTT / 2;

        Self {
            latest_rtt: DEFAULT_INITIAL_RTT,
            min_rtt: DEFAULT_INITIAL_RTT,
            smoothed_rtt,
            rttvar,
            max_ack_delay,
            first_rtt_sample: None,
        }
    }

    /// Gets the latest round trip time sample
    #[inline]
    pub fn latest_rtt(&self) -> Duration {
        self.latest_rtt
    }

    /// Gets the weighted average round trip time
    #[inline]
    pub fn smoothed_rtt(&self) -> Duration {
        self.smoothed_rtt
    }

    /// Gets the minimum round trip time
    #[inline]
    pub fn min_rtt(&self) -> Duration {
        self.min_rtt
    }

    /// Gets the variance in observed round trip time samples
    #[inline]
    pub fn rttvar(&self) -> Duration {
        self.rttvar
    }

    /// Gets the timestamp of the first RTT sample
    #[inline]
    pub fn first_rtt_sample(&self) -> Option<Timestamp> {
        self.first_rtt_sample
    }

    /// Gets the max_ack_delay
    #[inline]
    pub fn max_ack_delay(&self) -> Duration {
        self.max_ack_delay
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
    //# The PTO period is the amount of time that a sender ought to wait for
    //# an acknowledgement of a sent packet.
    #[inline]
    pub fn pto_period(&self, pto_backoff: u32, space: PacketNumberSpace) -> Duration {
        // We operate on microseconds rather than `Duration` to improve efficiency.
        // See https://godbolt.org/z/osEd9rj9a

        //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
        //# When an ack-eliciting packet is transmitted, the sender schedules a
        //# timer for the PTO period as follows:
        //#
        //# PTO = smoothed_rtt + max(4*rttvar, kGranularity) + max_ack_delay
        let mut pto_period = self.smoothed_rtt().as_micros() as u64;

        //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
        //# The PTO period MUST be at least kGranularity, to avoid the timer
        //# expiring immediately.
        pto_period += max(
            self.rttvar_4x().as_micros() as u64,
            K_GRANULARITY.as_micros() as u64,
        );

        //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
        //# When the PTO is armed for Initial or Handshake packet number spaces,
        //# the max_ack_delay in the PTO period computation is set to 0, since
        //# the peer is expected to not delay these packets intentionally; see
        //# Section 13.2.1 of [QUIC-TRANSPORT].
        if space.is_application_data() {
            pto_period += self.max_ack_delay.as_micros() as u64;
        }

        //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
        //# Even when there are ack-eliciting packets in flight in multiple
        //# packet number spaces, the exponential increase in PTO occurs across
        //# all spaces to prevent excess load on the network.  For example, a
        //# timeout in the Initial packet number space doubles the length of
        //# the timeout in the Handshake packet number space.
        pto_period *= pto_backoff as u64;

        //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
        //# The PTO period is the amount of time that a sender ought to wait for
        //# an acknowledgement of a sent packet.
        Duration::from_micros(pto_period)
    }

    /// Sets the `max_ack_delay` value from the peer `MaxAckDelay` transport parameter
    pub fn on_max_ack_delay(&mut self, max_ack_delay: MaxAckDelay) {
        self.max_ack_delay = max_ack_delay.as_duration()
    }

    /// Updates the RTT estimate using the given `rtt_sample`
    #[inline]
    pub fn update_rtt(
        &mut self,
        mut ack_delay: Duration,
        rtt_sample: Duration,
        timestamp: Timestamp,
        is_handshake_confirmed: bool,
        space: PacketNumberSpace,
    ) {
        self.latest_rtt = rtt_sample.max(Duration::from_micros(1));

        if self.first_rtt_sample.is_none() {
            self.first_rtt_sample = Some(timestamp);
            //= https://www.rfc-editor.org/rfc/rfc9002#section-5.2
            //# min_rtt MUST be set to the latest_rtt on the first RTT sample.
            self.min_rtt = self.latest_rtt;
            //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
            //# On the first RTT sample after initialization, smoothed_rtt and rttvar
            //# are set as follows:
            //#
            //# smoothed_rtt = latest_rtt
            //# rttvar = latest_rtt / 2
            self.smoothed_rtt = self.latest_rtt;
            self.rttvar = self.latest_rtt / 2;
            return;
        }

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.2
        //# min_rtt MUST be set to the lesser of min_rtt and latest_rtt
        //# (Section 5.1) on all other samples.
        self.min_rtt = min(self.min_rtt, self.latest_rtt);

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
        //# when adjusting an RTT sample using peer-reported
        //# acknowledgment delays, an endpoint:
        //#
        //# *  MAY ignore the acknowledgment delay for Initial packets, since
        //#    these acknowledgments are not delayed by the peer (Section 13.2.1
        //#    of [QUIC-TRANSPORT]);
        if space.is_initial() {
            ack_delay = ZERO_DURATION;
        }

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
        //# To account for this, the endpoint SHOULD ignore
        //# max_ack_delay until the handshake is confirmed, as defined in
        //# Section 4.1.2 of [QUIC-TLS].

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
        //# *  SHOULD ignore the peer's max_ack_delay until the handshake is
        //#    confirmed;
        if is_handshake_confirmed {
            //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
            //# *  MUST use the lesser of the acknowledgement delay and the peer's
            //#    max_ack_delay after the handshake is confirmed; and
            ack_delay = min(ack_delay, self.max_ack_delay);
        }

        let mut adjusted_rtt = self.latest_rtt;

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
        //# *  MUST NOT subtract the acknowledgement delay from the RTT sample if
        //#    the resulting value is smaller than the min_rtt.
        if self.min_rtt + ack_delay < self.latest_rtt {
            adjusted_rtt -= ack_delay;
        } else if !is_handshake_confirmed {
            //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
            //# Therefore, prior to handshake
            //# confirmation, an endpoint MAY ignore RTT samples if adjusting the RTT
            //# sample for acknowledgement delay causes the sample to be less than
            //# the min_rtt.
            return;
        }

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
        //# On subsequent RTT samples, smoothed_rtt and rttvar evolve as follows:
        //#
        //# ack_delay = decoded acknowledgment delay from ACK frame
        //# if (handshake confirmed):
        //# ack_delay = min(ack_delay, max_ack_delay)
        //# adjusted_rtt = latest_rtt
        //# if (latest_rtt >= min_rtt + ack_delay):
        //#     adjusted_rtt = latest_rtt - ack_delay
        //# smoothed_rtt = 7/8 * smoothed_rtt + 1/8 * adjusted_rtt
        //# rttvar_sample = abs(smoothed_rtt - adjusted_rtt)
        //# rttvar = 3/4 * rttvar + 1/4 * rttvar_sample

        // this logic has been updated to follow the errata reported in https://www.rfc-editor.org/errata/eid7539
        let rttvar_sample = abs_difference(self.smoothed_rtt, adjusted_rtt);
        self.rttvar = weighted_average(self.rttvar, rttvar_sample, 4);
        self.smoothed_rtt = weighted_average(self.smoothed_rtt, adjusted_rtt, 8);
    }

    /// Calculates the persistent congestion threshold used for determining
    /// if persistent congestion is being encountered.
    #[inline]
    pub fn persistent_congestion_threshold(&self) -> Duration {
        // Since K_GRANULARITY is 1ms, we operate on milliseconds rather than `Duration` to improve efficiency.
        // See https://godbolt.org/z/4o71WPods

        //= https://www.rfc-editor.org/rfc/rfc9002#section-7.6.1
        //# The persistent congestion duration is computed as follows:
        //#
        //# (smoothed_rtt + max(4*rttvar, kGranularity) + max_ack_delay) *
        //#     kPersistentCongestionThreshold
        //#
        //# Unlike the PTO computation in Section 6.2, this duration includes the
        //# max_ack_delay irrespective of the packet number spaces in which
        //# losses are established.
        //#
        //# This duration allows a sender to send as many packets before
        //# establishing persistent congestion, including some in response to PTO
        //# expiration, as TCP does with Tail Loss Probes [RFC8985] and an RTO
        //# [RFC5681].
        Duration::from_millis(
            (self.smoothed_rtt.as_millis() as u64
                + max(
                    self.rttvar_4x().as_millis() as u64,
                    K_GRANULARITY.as_millis() as u64,
                )
                + self.max_ack_delay.as_millis() as u64)
                * K_PERSISTENT_CONGESTION_THRESHOLD,
        )
    }

    /// Allows min_rtt and smoothed_rtt to be overwritten on the next RTT sample
    /// after persistent congestion is established.
    #[inline]
    pub fn on_persistent_congestion(&mut self) {
        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.2
        //# Endpoints SHOULD set the min_rtt to the newest RTT sample after
        //# persistent congestion is established.
        self.first_rtt_sample = None;
    }

    #[inline]
    fn rttvar_4x(&self) -> Duration {
        // Operate on micros instead, as it's more efficient and we don't need the precision Duration gives
        Duration::from_micros(4 * self.rttvar.as_micros() as u64)
    }
}

#[inline]
fn abs_difference<T: core::ops::Sub + PartialOrd>(a: T, b: T) -> <T as core::ops::Sub>::Output {
    if a > b {
        a - b
    } else {
        b - a
    }
}

/// Optimized function for averaging two durations with a weight
/// See https://godbolt.org/z/65f9bYEcs
#[inline]
fn weighted_average(a: Duration, b: Duration, weight: u64) -> Duration {
    let mut a = a.as_nanos() as u64;
    // it's more accurate to multiply first but it risks overflow so we divide first
    a /= weight;
    a *= weight - 1;

    let mut b = b.as_nanos() as u64;
    b /= weight;

    Duration::from_nanos(a + b)
}

#[cfg(test)]
mod test {
    use crate::{
        packet::number::PacketNumberSpace,
        path::INITIAL_PTO_BACKOFF,
        recovery::{RttEstimator, DEFAULT_INITIAL_RTT, K_GRANULARITY},
        time::{Clock, Duration, NoopClock},
        transport::parameters::MaxAckDelay,
        varint::VarInt,
    };

    /// Test the initial values before any RTT samples
    #[test]
    fn initial_rtt_across_spaces() {
        let rtt_estimator = RttEstimator::new(Duration::from_millis(10));
        assert_eq!(rtt_estimator.min_rtt, DEFAULT_INITIAL_RTT);
        assert_eq!(rtt_estimator.latest_rtt(), DEFAULT_INITIAL_RTT);
        assert_eq!(rtt_estimator.smoothed_rtt(), DEFAULT_INITIAL_RTT);
        assert_eq!(rtt_estimator.rttvar(), DEFAULT_INITIAL_RTT / 2);
        assert_eq!(
            rtt_estimator.pto_period(INITIAL_PTO_BACKOFF, PacketNumberSpace::Initial),
            Duration::from_millis(999)
        );
        assert_eq!(
            rtt_estimator.pto_period(INITIAL_PTO_BACKOFF, PacketNumberSpace::Handshake),
            Duration::from_millis(999)
        );
        assert_eq!(
            rtt_estimator.pto_period(INITIAL_PTO_BACKOFF, PacketNumberSpace::ApplicationData),
            Duration::from_millis(1009)
        );
    }

    /// Test a zero RTT value is treated as 1 Âµs
    #[test]
    fn zero_rtt_sample() {
        let mut rtt_estimator = RttEstimator::new(Duration::from_millis(10));
        let now = NoopClock.get_time();
        rtt_estimator.update_rtt(
            Duration::from_millis(10),
            Duration::from_millis(0),
            now,
            false,
            PacketNumberSpace::ApplicationData,
        );
        assert_eq!(rtt_estimator.min_rtt, Duration::from_micros(1));
        assert_eq!(rtt_estimator.latest_rtt(), Duration::from_micros(1));
        assert_eq!(rtt_estimator.first_rtt_sample(), Some(now));
        assert_eq!(
            rtt_estimator.pto_period(INITIAL_PTO_BACKOFF, PacketNumberSpace::Initial),
            Duration::from_micros(1001)
        );
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
    //= type=test
    //# *  MUST use the lesser of the acknowledgement delay and the peer's
    //#    max_ack_delay after the handshake is confirmed;
    #[test]
    fn max_ack_delay() {
        let mut rtt_estimator = RttEstimator::default();
        assert_eq!(Duration::ZERO, rtt_estimator.max_ack_delay());

        rtt_estimator.on_max_ack_delay(MaxAckDelay::new(VarInt::from_u8(10)).unwrap());
        assert_eq!(Duration::from_millis(10), rtt_estimator.max_ack_delay());

        let now = NoopClock.get_time();
        rtt_estimator.update_rtt(
            Duration::from_millis(0),
            Duration::from_millis(100),
            now,
            true,
            PacketNumberSpace::ApplicationData,
        );

        // Update when the handshake is confirmed
        rtt_estimator.update_rtt(
            Duration::from_millis(1000),
            Duration::from_millis(200),
            now,
            true,
            PacketNumberSpace::ApplicationData,
        );

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
        //= type=test
        //# *  MUST use the lesser of the acknowledgement delay and the peer's
        //# max_ack_delay after the handshake is confirmed; and
        assert_eq!(
            rtt_estimator.smoothed_rtt,
            7 * Duration::from_millis(100) / 8 + Duration::from_millis(200 - 10) / 8
        );
        assert_eq!(rtt_estimator.first_rtt_sample(), Some(now));

        let prev_smoothed_rtt = rtt_estimator.smoothed_rtt;

        // Update when the handshake is not confirmed
        rtt_estimator.update_rtt(
            Duration::from_millis(50),
            Duration::from_millis(200),
            now,
            false,
            PacketNumberSpace::ApplicationData,
        );

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
        //= type=test
        //# To account for this, the endpoint SHOULD ignore
        //# max_ack_delay until the handshake is confirmed, as defined in
        //# Section 4.1.2 of [QUIC-TLS].

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
        //= type=test
        //# *  SHOULD ignore the peer's max_ack_delay until the handshake is
        //# confirmed;
        assert_eq!(
            rtt_estimator.smoothed_rtt,
            7 * prev_smoothed_rtt / 8 + Duration::from_millis(200 - 50) / 8
        );
    }

    /// Test several rounds of RTT updates
    #[test]
    fn update_rtt() {
        let mut rtt_estimator = RttEstimator::new(Duration::from_millis(10));
        let now = NoopClock.get_time();
        let rtt_sample = Duration::from_millis(500);
        assert!(rtt_estimator.first_rtt_sample.is_none());
        rtt_estimator.update_rtt(
            Duration::from_millis(10),
            rtt_sample,
            now,
            true,
            PacketNumberSpace::ApplicationData,
        );

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.2
        //= type=test
        //# min_rtt MUST be set to the latest_rtt on the first RTT sample.
        assert_eq!(rtt_estimator.min_rtt, rtt_sample);
        assert_eq!(rtt_estimator.latest_rtt, rtt_sample);
        assert_eq!(rtt_estimator.smoothed_rtt, rtt_sample);
        assert_eq!(rtt_estimator.rttvar, rtt_sample / 2);
        assert_eq!(rtt_estimator.first_rtt_sample, Some(now));

        let prev_smoothed_rtt = rtt_estimator.smoothed_rtt;
        let rtt_sample = Duration::from_millis(800);
        let ack_delay = Duration::from_millis(10);

        rtt_estimator.update_rtt(
            ack_delay,
            rtt_sample,
            now + Duration::from_secs(1),
            true,
            PacketNumberSpace::ApplicationData,
        );

        let adjusted_rtt = rtt_sample - ack_delay;

        assert_eq!(rtt_estimator.min_rtt, prev_smoothed_rtt);
        assert_eq!(rtt_estimator.latest_rtt, rtt_sample);
        assert_eq!(
            rtt_estimator.smoothed_rtt,
            7 * prev_smoothed_rtt / 8 + adjusted_rtt / 8
        );
        assert_eq!(rtt_estimator.first_rtt_sample, Some(now));

        // This rtt_sample is a new minimum, so the ack_delay is not used for adjustment
        let prev_smoothed_rtt = rtt_estimator.smoothed_rtt;
        let rtt_sample = Duration::from_millis(200);
        let ack_delay = Duration::from_millis(10);

        rtt_estimator.update_rtt(
            ack_delay,
            rtt_sample,
            now + Duration::from_secs(2),
            true,
            PacketNumberSpace::ApplicationData,
        );

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.2
        //= type=test
        //# min_rtt MUST be set to the lesser of min_rtt and latest_rtt
        //# (Section 5.1) on all other samples.
        assert_eq!(rtt_estimator.min_rtt, rtt_sample);
        assert_eq!(rtt_estimator.latest_rtt, rtt_sample);
        assert_eq!(
            rtt_estimator.smoothed_rtt,
            7 * prev_smoothed_rtt / 8 + rtt_sample / 8
        );
        assert_eq!(rtt_estimator.first_rtt_sample, Some(now));
        assert_eq!(
            rtt_estimator.pto_period(INITIAL_PTO_BACKOFF, PacketNumberSpace::ApplicationData),
            Duration::from_micros(1620466)
        );
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
    //= type=test
    //# *  MUST NOT subtract the acknowledgement delay from the RTT sample if
    //#    the resulting value is smaller than the min_rtt.
    #[test]
    fn must_not_subtract_acknowledgement_delay_if_result_smaller_than_min_rtt() {
        let mut rtt_estimator = RttEstimator::new(Duration::from_millis(200));
        let now = NoopClock.get_time();

        rtt_estimator.min_rtt = Duration::from_millis(500);
        rtt_estimator.smoothed_rtt = Duration::from_millis(700);
        rtt_estimator.first_rtt_sample = Some(now);

        let rtt_sample = Duration::from_millis(600);
        let prev_smoothed_rtt = rtt_estimator.smoothed_rtt;

        rtt_estimator.update_rtt(
            Duration::from_millis(200),
            rtt_sample,
            now,
            true,
            PacketNumberSpace::ApplicationData,
        );

        assert_eq!(
            rtt_estimator.smoothed_rtt,
            7 * prev_smoothed_rtt / 8 + rtt_sample / 8
        );
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
    //= type=test
    //# Therefore, prior to handshake
    //# confirmation, an endpoint MAY ignore RTT samples if adjusting the RTT
    //# sample for acknowledgement delay causes the sample to be less than
    //# the min_rtt.
    #[test]
    fn prior_to_handshake_ignore_if_less_than_min_rtt() {
        let mut rtt_estimator = RttEstimator::new(Duration::from_millis(200));
        let now = NoopClock.get_time();
        let smoothed_rtt = Duration::from_millis(700);

        rtt_estimator.min_rtt = Duration::from_millis(500);
        rtt_estimator.smoothed_rtt = smoothed_rtt;
        rtt_estimator.first_rtt_sample = Some(now);

        let rtt_sample = Duration::from_millis(600);

        rtt_estimator.update_rtt(
            Duration::from_millis(200),
            rtt_sample,
            now,
            false,
            PacketNumberSpace::ApplicationData,
        );

        assert_eq!(rtt_estimator.smoothed_rtt, smoothed_rtt);
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-5.3
    //= type=test
    //# *  MAY ignore the acknowledgment delay for Initial packets, since
    //     these acknowledgments are not delayed by the peer (Section 13.2.1
    //     of [QUIC-TRANSPORT]);
    #[test]
    fn initial_space() {
        let mut rtt_estimator = RttEstimator::new(Duration::from_millis(10));
        let now = NoopClock.get_time();
        let rtt_sample = Duration::from_millis(500);
        rtt_estimator.update_rtt(
            Duration::from_millis(10),
            rtt_sample,
            now,
            true,
            PacketNumberSpace::Initial,
        );

        let prev_smoothed_rtt = rtt_estimator.smoothed_rtt;
        let rtt_sample = Duration::from_millis(1000);

        rtt_estimator.update_rtt(
            Duration::from_millis(100),
            rtt_sample,
            now,
            true,
            PacketNumberSpace::Initial,
        );

        assert_eq!(
            rtt_estimator.smoothed_rtt,
            7 * prev_smoothed_rtt / 8 + rtt_sample / 8
        );
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-7.6.1
    //= type=test
    //# The persistent congestion duration is computed as follows:
    //#
    //# (smoothed_rtt + max(4*rttvar, kGranularity) + max_ack_delay) *
    //#   kPersistentCongestionThreshold
    #[test]
    fn persistent_congestion_duration() {
        let max_ack_delay = Duration::from_millis(10);
        let mut rtt_estimator = RttEstimator::new(max_ack_delay);

        rtt_estimator.smoothed_rtt = Duration::from_millis(100);
        rtt_estimator.rttvar = Duration::from_millis(50);

        // persistent congestion period =
        // (smoothed_rtt + max(4*rttvar, kGranularity) + max_ack_delay) * kPersistentCongestionThreshold
        // = (100 + max(4*50, 1) + 10) * 3 = 930
        assert_eq!(
            Duration::from_millis(930),
            rtt_estimator.persistent_congestion_threshold()
        );

        rtt_estimator.rttvar = Duration::from_millis(0);

        //= https://www.rfc-editor.org/rfc/rfc9002#section-7.6.1
        //= type=test
        //# The RECOMMENDED value for kPersistentCongestionThreshold is 3, which
        //# results in behavior that is approximately equivalent to a TCP sender
        //# declaring an RTO after two TLPs.

        // persistent congestion period =
        // (smoothed_rtt + max(4*rttvar, kGranularity) + max_ack_delay) * kPersistentCongestionThreshold
        // = (100 + max(0, 1) + 10) * 3 = 333
        assert_eq!(
            Duration::from_millis(333),
            rtt_estimator.persistent_congestion_threshold()
        );
    }

    #[test]
    fn set_min_rtt_to_latest_sample_after_persistent_congestion() {
        let mut rtt_estimator = RttEstimator::new(Duration::from_millis(10));
        let now = NoopClock.get_time();
        let mut rtt_sample = Duration::from_millis(500);
        rtt_estimator.update_rtt(
            Duration::from_millis(10),
            rtt_sample,
            now,
            true,
            PacketNumberSpace::Initial,
        );

        assert_eq!(rtt_estimator.min_rtt(), rtt_sample);

        rtt_sample = Duration::from_millis(200);

        rtt_estimator.on_persistent_congestion();

        rtt_estimator.update_rtt(
            Duration::from_millis(10),
            rtt_sample,
            now,
            true,
            PacketNumberSpace::Initial,
        );

        //= https://www.rfc-editor.org/rfc/rfc9002#section-5.2
        //= type=test
        //# Endpoints SHOULD set the min_rtt to the newest RTT sample after
        //# persistent congestion is established.
        assert_eq!(rtt_estimator.min_rtt(), rtt_sample);
        assert_eq!(rtt_estimator.smoothed_rtt(), rtt_sample);
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
    //= type=test
    //# The PTO period MUST be at least kGranularity, to avoid the timer
    //# expiring immediately.
    #[test]
    fn pto_must_be_at_least_k_granularity() {
        let space = PacketNumberSpace::Handshake;
        let max_ack_delay = Duration::from_millis(0);
        let now = NoopClock.get_time();
        let mut rtt_estimator = RttEstimator::new(max_ack_delay);

        // Update RTT with the smallest possible sample
        rtt_estimator.update_rtt(
            Duration::from_millis(0),
            Duration::from_nanos(1),
            now,
            true,
            space,
        );

        let pto_period = rtt_estimator.pto_period(INITIAL_PTO_BACKOFF, space);
        assert!(pto_period >= K_GRANULARITY);

        // pto_period should have microsecond precision
        assert_eq!(pto_period, Duration::from_micros(1001))
    }

    #[test]
    #[cfg_attr(kani, kani::proof, kani::unwind(3), kani::solver(kissat))]
    #[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
    fn weighted_average_test() {
        bolero::check!()
            .with_type::<(u32, u32)>()
            .for_each(|(a, b)| {
                let a = Duration::from_nanos(*a as _);
                let b = Duration::from_nanos(*b as _);

                let weight = 8;

                // perform the unoptimized version
                let expected = ((weight - 1) * a) / weight + b / weight;
                let actual = super::weighted_average(a, b, weight as _);

                // assert that the unoptimized result matches the optimized to the nearest `weight` nanos
                assert!(
                    super::abs_difference(expected.as_nanos(), actual.as_nanos()) as u32 <= weight,
                    "expected: {:?}; actual: {:?}",
                    expected,
                    actual
                );
            })
    }
}