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

use crate::{
    frame::ack_elicitation::AckElicitation, inet::ExplicitCongestionNotification, path,
    time::Timestamp, transmission,
};
use core::convert::TryInto;

//= https://www.rfc-editor.org/rfc/rfc9002#appendix-A.1

//= https://www.rfc-editor.org/rfc/rfc9002#appendix-A.1.1

#[cfg(feature = "alloc")]
pub type SentPackets<PacketInfo> = crate::packet::number::Map<SentPacketInfo<PacketInfo>>;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub struct SentPacketInfo<PacketInfo> {
    /// Indicates whether the packet counts towards bytes in flight
    pub congestion_controlled: bool,
    /// The number of bytes sent in the packet, not including UDP or IP overhead,
    /// but including QUIC framing overhead
    pub sent_bytes: u16,
    /// The time the packet was sent
    pub time_sent: Timestamp,
    /// Indicates whether a packet is ack-eliciting
    pub ack_elicitation: AckElicitation,
    /// The ID of the Path the packet was sent on
    pub path_id: path::Id,
    /// The ECN marker (if any) sent on the datagram that contained this packet
    pub ecn: ExplicitCongestionNotification,
    /// Indicates if the packet was part of a probe transmission
    pub transmission_mode: transmission::Mode,
    /// Additional packet metadata dictated by the congestion controller
    pub cc_packet_info: PacketInfo,
}

impl<PacketInfo> SentPacketInfo<PacketInfo> {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        congestion_controlled: bool,
        sent_bytes: usize,
        time_sent: Timestamp,
        ack_elicitation: AckElicitation,
        path_id: path::Id,
        ecn: ExplicitCongestionNotification,
        transmission_mode: transmission::Mode,
        cc_packet_info: PacketInfo,
    ) -> Self {
        debug_assert_eq!(
            sent_bytes > 0,
            congestion_controlled,
            "sent bytes should be zero for packets that are not congestion controlled"
        );

        SentPacketInfo {
            congestion_controlled,
            sent_bytes: sent_bytes
                .try_into()
                .expect("sent_bytes exceeds max UDP payload size"),
            time_sent,
            ack_elicitation,
            path_id,
            ecn,
            transmission_mode,
            cc_packet_info,
        }
    }
}

#[cfg(test)]
mod test {
    use crate::{
        frame::ack_elicitation::AckElicitation,
        inet::ExplicitCongestionNotification,
        path,
        recovery::SentPacketInfo,
        time::{Clock, NoopClock},
        transmission,
    };

    #[test]
    #[should_panic]
    fn too_large_packet() {
        SentPacketInfo::new(
            true,
            u16::MAX as usize + 1,
            NoopClock.get_time(),
            AckElicitation::Eliciting,
            unsafe { path::Id::new(0) },
            ExplicitCongestionNotification::default(),
            transmission::Mode::Normal,
            (),
        );
    }

    #[test]
    #[cfg_attr(miri, ignore)] // snapshot tests don't work on miri
    fn sent_packet_info_size_test() {
        insta::assert_debug_snapshot!(
            stringify!(sent_packet_info_size_test),
            core::mem::size_of::<SentPacketInfo<()>>()
        );

        assert_eq!(
            core::mem::size_of::<Option<SentPacketInfo<()>>>(),
            core::mem::size_of::<SentPacketInfo<()>>()
        );
    }
}


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

use crate::{
    event,
    packet::number::PacketNumberSpace,
    path,
    path::MINIMUM_MTU,
    random,
    recovery::{
        congestion_controller::PathPublisher, CongestionController, CubicCongestionController,
        RttEstimator,
    },
    time::{Clock, NoopClock, Timestamp},
};
use core::{fmt, ops::Range, time::Duration};
use insta::assert_debug_snapshot;
use plotters::prelude::*;
use std::{
    env,
    path::{Path, PathBuf},
};

const CHART_DIMENSIONS: (u32, u32) = (1024, 768);

fn type_name<T>() -> &'static str {
    core::any::type_name::<T>().split("::").last().unwrap()
}

// These simulations are too slow for Miri
#[test]
#[cfg_attr(miri, ignore)]
fn slow_start_unlimited_test() {
    let cc = CubicCongestionController::new(MINIMUM_MTU);

    slow_start_unlimited(cc, 12).finish();
}

#[test]
#[cfg_attr(miri, ignore)]
fn loss_at_3mb_test() {
    let cc = CubicCongestionController::new(MINIMUM_MTU);

    loss_at_3mb(cc, 135).finish();
}

#[test]
#[cfg_attr(miri, ignore)]
fn app_limited_1mb_test() {
    let cc = CubicCongestionController::new(MINIMUM_MTU);

    app_limited_1mb(cc, 120).finish();
}

#[test]
#[cfg_attr(miri, ignore)]
fn minimum_window_test() {
    let cc = CubicCongestionController::new(MINIMUM_MTU);

    minimum_window(cc, 10).finish();
}

#[test]
#[cfg_attr(miri, ignore)]
fn loss_at_3mb_and_2_75mb_test() {
    let cc = CubicCongestionController::new(MINIMUM_MTU);

    loss_at_3mb_and_2_75mb(cc, 120).finish();
}

#[derive(Debug)]
struct Simulation {
    name: &'static str,
    description: &'static str,
    cc: &'static str,
    rounds: Vec<Round>,
}

struct Round {
    number: usize,
    cwnd: u32,
}

impl fmt::Debug for Round {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Print out packet count rather than bytes to normalize any differences between
        // platforms.
        //
        // Floating point numbers are not guaranteed to be deterministic between CPU
        // architectures, compilers, or libc implementations.
        let packets = self.cwnd / MINIMUM_MTU as u32;

        write!(f, "{:>3}: pkts: {}", self.number, packets)
    }
}

impl Simulation {
    fn finish(&self) {
        if let Ok(dir) = env::var("RECOVERY_SIM_DIR") {
            let mut path = PathBuf::new();
            path.push(dir);
            path.push(self.filename());
            path.set_extension("svg");
            self.plot(&path);
        } else {
            self.assert_snapshot();
        }
    }

    fn plot<T: AsRef<Path> + ?Sized>(&self, path: &T) {
        let root_area = SVGBackend::new(path, CHART_DIMENSIONS).into_drawing_area();
        root_area.fill(&WHITE).expect("Could not fill chart");
        root_area
            .titled(&self.name(), ("sans-serif", 40))
            .expect("Could not add title");

        let mut ctx = ChartBuilder::on(&root_area)
            .set_label_area_size(LabelAreaPosition::Left, 120)
            .set_label_area_size(LabelAreaPosition::Bottom, 60)
            .margin(20)
            .margin_top(40)
            .caption(self.description, ("sans-serif", 20))
            .build_cartesian_2d(self.x_spec(), self.y_spec())
            .expect("Could not build chart");

        ctx.configure_mesh()
            .x_desc("Transmission Round")
            .label_style(("sans-serif", 20))
            .y_desc("Congestion window size (bytes)")
            .draw()
            .expect("Could not configure mesh");

        ctx.draw_series(LineSeries::new(
            self.rounds.iter().map(|x| (x.number as i32, x.cwnd as i32)),
            GREEN,
        ))
        .expect("Could not draw series");
    }

    fn x_spec(&self) -> Range<i32> {
        0..(self.rounds.len() as i32 + 1)
    }

    fn y_spec(&self) -> Range<i32> {
        let mut max = self.rounds.iter().map(|r| r.cwnd as i32).max().unwrap_or(0);

        // Add a 5% buffer
        max = (max as f32 * 1.05) as i32;

        0..max
    }

    fn assert_snapshot(&self) {
        assert_debug_snapshot!(self.filename(), self);
    }

    fn name(&self) -> String {
        let mut name = String::new();
        name.push_str(self.name);
        name.push_str(" - ");
        name.push_str(self.cc.split("::").last().unwrap());
        name
    }

    fn filename(&self) -> String {
        self.name().replace('.', "_").split_whitespace().collect()
    }
}

/// Simulates a network with no congestion experienced
fn slow_start_unlimited<CC: CongestionController>(
    mut congestion_controller: CC,
    num_rounds: usize,
) -> Simulation {
    Simulation {
        name: "Slow Start Unlimited",
        description: "Full congestion window utilization with no congestion experienced",
        cc: type_name::<CC>(),
        rounds: simulate_constant_rtt(&mut congestion_controller, &[], None, num_rounds),
    }
}

/// Simulates a network that experienced loss at a 3MB congestion window
fn loss_at_3mb<CC: CongestionController>(
    mut congestion_controller: CC,
    num_rounds: usize,
) -> Simulation {
    Simulation {
        name: "Loss at 3MB",
        description: "Full congestion window utilization with loss encountered at ~3MB",
        cc: type_name::<CC>(),
        rounds: simulate_constant_rtt(&mut congestion_controller, &[3_000_000], None, num_rounds),
    }
}

/// Simulates a network that experiences loss at a 750KB congestion window with the application
/// sending at most 1MB of data per round.
fn app_limited_1mb<CC: CongestionController>(
    mut congestion_controller: CC,
    num_rounds: usize,
) -> Simulation {
    const APP_LIMIT_BYTES: usize = 1_000_000;

    Simulation {
        name: "App Limited 1MB",
        description: "App limited to 1MB per round with loss encountered at ~750KB",
        cc: type_name::<CC>(),
        rounds: simulate_constant_rtt(
            &mut congestion_controller,
            &[750_000],
            Some(APP_LIMIT_BYTES),
            num_rounds,
        ),
    }
}

/// Simulates a network starting from the minimum window size with no further congestion
fn minimum_window<CC: CongestionController>(
    mut congestion_controller: CC,
    num_rounds: usize,
) -> Simulation {
    let time_zero = NoopClock.get_time();
    let rtt_estimator = RttEstimator::default();
    let random = &mut random::testing::Generator::default();
    let mut publisher = event::testing::Publisher::no_snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());

    let packet_info = congestion_controller.on_packet_sent(
        time_zero,
        MINIMUM_MTU as usize,
        None,
        &rtt_estimator,
        &mut publisher,
    );
    // Experience persistent congestion to drop to the minimum window
    congestion_controller.on_packet_lost(
        MINIMUM_MTU as u32,
        packet_info,
        true,
        false,
        random,
        time_zero,
        &mut publisher,
    );
    let packet_info = congestion_controller.on_packet_sent(
        time_zero,
        MINIMUM_MTU as usize,
        None,
        &rtt_estimator,
        &mut publisher,
    );
    // Lose a packet to exit slow start
    congestion_controller.on_packet_lost(
        MINIMUM_MTU as u32,
        packet_info,
        false,
        false,
        random,
        time_zero,
        &mut publisher,
    );

    Simulation {
        name: "Minimum Window",
        description: "Full congestion window utilization after starting from the minimum window",
        cc: type_name::<CC>(),
        rounds: simulate_constant_rtt(&mut congestion_controller, &[], None, num_rounds),
    }
}

/// Simulates a network that experienced loss at a 3MB congestion window and then at a 2.75MB window
/// With Cubic this will exercise the fast convergence algorithm
fn loss_at_3mb_and_2_75mb<CC: CongestionController>(
    mut congestion_controller: CC,
    num_rounds: usize,
) -> Simulation {
    Simulation {
        name: "Loss at 3MB and 2.75MB",
        description: "Loss encountered at ~3MB and ~2.75MB",
        cc: type_name::<CC>(),
        rounds: simulate_constant_rtt(
            &mut congestion_controller,
            &[3_000_000, 2_750_000],
            None,
            num_rounds,
        ),
    }
}

/// Simulate the given number of rounds with drops occurring at the given congestion window sizes
/// and limited to the given app limit
fn simulate_constant_rtt<CC: CongestionController>(
    congestion_controller: &mut CC,
    drops: &[u32],
    app_limit: Option<usize>,
    num_rounds: usize,
) -> Vec<Round> {
    let time_zero = NoopClock.get_time();
    let mut rtt_estimator = RttEstimator::default();
    let random = &mut random::testing::Generator::default();
    let mut publisher = event::testing::Publisher::no_snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());

    // Update the rtt with 200 ms
    rtt_estimator.update_rtt(
        Duration::from_millis(0),
        Duration::from_millis(200),
        time_zero,
        true,
        PacketNumberSpace::ApplicationData,
    );
    let mut round_start = NoopClock.get_time() + Duration::from_millis(1);
    let mut rounds = Vec::with_capacity(num_rounds);
    let mut drop_index = 0;

    for round in 0..num_rounds {
        rounds.push(Round {
            number: round,
            cwnd: congestion_controller.congestion_window(),
        });

        round_start += Duration::from_millis(200);

        if drop_index < drops.len()
            && congestion_controller.congestion_window() >= drops[drop_index]
        {
            // Drop a packet
            let packet_info = congestion_controller.on_packet_sent(
                round_start,
                MINIMUM_MTU as usize,
                None,
                &rtt_estimator,
                &mut publisher,
            );
            congestion_controller.on_packet_lost(
                MINIMUM_MTU as u32,
                packet_info,
                false,
                false,
                random,
                round_start,
                &mut publisher,
            );
            drop_index += 1;
        } else {
            let send_bytes = (congestion_controller.congestion_window() as usize)
                .min(app_limit.unwrap_or(usize::MAX));

            // Send and ack the full congestion window
            send_and_ack(
                congestion_controller,
                &rtt_estimator,
                round_start,
                send_bytes,
            );
        }
    }

    rounds
}

/// Send and acknowledge the given amount of bytes using the given congestion controller
fn send_and_ack<CC: CongestionController>(
    congestion_controller: &mut CC,
    rtt_estimator: &RttEstimator,
    timestamp: Timestamp,
    bytes: usize,
) {
    let random = &mut random::testing::Generator::default();
    let mut tx_remaining = bytes;
    let mut rx_remaining = 0;
    let mut now = timestamp;
    let ack_receive_time = now + rtt_estimator.min_rtt();
    // Allow acks to start being received after this time, to simulate
    // acks arriving while sending is paused by the pacer.
    let earliest_ack_receive_time = ack_receive_time - Duration::from_millis(50);
    let sending_full_cwnd = bytes as u32 == congestion_controller.congestion_window();
    let mut publisher = event::testing::Publisher::no_snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());

    let mut packet_info = None;

    while tx_remaining > 0 || rx_remaining > 0 {
        while tx_remaining > 0 {
            if let Some(edt) = congestion_controller.earliest_departure_time() {
                if !edt.has_elapsed(now) {
                    // We are blocked by the pacer, stop sending and fast forward to the earliest departure time
                    now = edt;
                    break;
                }
            }

            let bytes_sent = tx_remaining.min(MINIMUM_MTU as usize);
            let app_limited = tx_remaining - bytes_sent == 0 && !sending_full_cwnd;

            packet_info = Some(congestion_controller.on_packet_sent(
                now,
                bytes_sent,
                Some(app_limited),
                rtt_estimator,
                &mut publisher,
            ));
            tx_remaining -= bytes_sent;
            rx_remaining += bytes_sent;
        }

        if tx_remaining == 0 {
            // Nothing left to send, so fast forward to when we receive acks for everything
            now = ack_receive_time;
        }

        while now >= earliest_ack_receive_time && rx_remaining > 0 {
            let bytes_acked = rx_remaining.min(MINIMUM_MTU as usize);

            congestion_controller.on_ack(
                now,
                bytes_acked,
                packet_info.unwrap(),
                rtt_estimator,
                random,
                now,
                &mut publisher,
            );
            rx_remaining -= bytes_acked;
        }
    }
}


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

use crate::{event, event::IntoEvent, recovery::congestion_controller::Publisher, time::Timestamp};
use core::{
    cmp::{max, Ordering},
    time::Duration,
};
use num_rational::Ratio;
use num_traits::Inv;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
/// Bandwidth-related data tracked for each sent packet
pub struct PacketInfo {
    /// [Estimator::delivered_bytes] at the time this packet was sent.
    pub delivered_bytes: u64,
    /// `Estimator::delivered_time` at the time this packet was sent.
    pub delivered_time: Timestamp,
    /// `Estimator::lost_bytes` at the time this packet was sent.
    pub lost_bytes: u64,
    /// `Estimator::ecn_ce_count` at the time this packet was sent.
    pub ecn_ce_count: u64,
    /// `Estimator::first_sent_time` at the time this packet was sent.
    pub first_sent_time: Timestamp,
    /// The volume of data that was estimated to be in flight at the
    /// time of the transmission of this packet
    pub bytes_in_flight: u32,
    /// Whether the path send rate was limited by the application rather
    /// than congestion control at the time this packet was sent
    pub is_app_limited: bool,
}

/// Represents a rate at which data is transferred
///
/// While bandwidth is typically thought of as an amount of data over a fixed
/// amount of time (bytes per second, for example), in this case we internally
/// represent bandwidth as the inverse: an amount of time to send a fixed amount
/// of data (nanoseconds per kibibyte or 1024 bytes, in this case). This allows for
/// some of the math operations needed on `Bandwidth` to avoid division, while
/// reducing the likelihood of panicking due to overflow.
///
/// The maximum (non-infinite) value that can be represented is ~1 TB/second.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct Bandwidth {
    nanos_per_kibibyte: u64,
}

// 2^10 = 1024 bytes in kibibyte
const KIBIBYTE_SHIFT: u8 = 10;

impl Bandwidth {
    pub const ZERO: Bandwidth = Bandwidth {
        nanos_per_kibibyte: u64::MAX,
    };

    pub const INFINITY: Bandwidth = Bandwidth {
        nanos_per_kibibyte: 0,
    };

    /// Constructs a new `Bandwidth` with the given bytes per interval
    pub const fn new(bytes: u64, interval: Duration) -> Self {
        let interval = (interval.as_nanos() as u64) << KIBIBYTE_SHIFT;
        if interval == 0 || bytes == 0 {
            Bandwidth::ZERO
        } else {
            Self {
                nanos_per_kibibyte: interval / bytes,
            }
        }
    }

    /// Represents the bandwidth as bytes per second
    pub fn as_bytes_per_second(&self) -> u64 {
        const ONE_SECOND_IN_NANOS: u64 = Duration::from_secs(1).as_nanos() as u64;

        if *self == Bandwidth::INFINITY {
            return u64::MAX;
        }

        (ONE_SECOND_IN_NANOS << KIBIBYTE_SHIFT) / self.nanos_per_kibibyte
    }
}

impl Default for Bandwidth {
    fn default() -> Self {
        Bandwidth::ZERO
    }
}

impl core::cmp::PartialOrd for Bandwidth {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl core::cmp::Ord for Bandwidth {
    fn cmp(&self, other: &Self) -> Ordering {
        // The higher the nanos_per_kibibyte, the lower the bandwidth,
        // so reverse the ordering when comparing
        self.nanos_per_kibibyte
            .cmp(&other.nanos_per_kibibyte)
            .reverse()
    }
}

impl core::ops::Mul<Ratio<u64>> for Bandwidth {
    type Output = Bandwidth;

    fn mul(self, rhs: Ratio<u64>) -> Self::Output {
        if self == Bandwidth::ZERO {
            return Bandwidth::ZERO;
        }

        Bandwidth {
            // Since `Bandwidth` is represented as time/byte and not bytes/time, we should divide
            // by the given ratio to result in a higher nanos_per_kibibyte value (lower bandwidth).
            // To avoid division, we can multiply by the inverse of the ratio instead
            nanos_per_kibibyte: (rhs.inv() * self.nanos_per_kibibyte).to_integer(),
        }
    }
}

impl core::ops::Mul<Duration> for Bandwidth {
    type Output = u64;

    fn mul(self, rhs: Duration) -> Self::Output {
        if self == Bandwidth::INFINITY {
            return u64::MAX;
        } else if rhs.is_zero() {
            return 0;
        }

        let interval = (rhs.as_nanos() as u64) << KIBIBYTE_SHIFT;

        if interval == 0 {
            return u64::MAX;
        }

        interval / self.nanos_per_kibibyte
    }
}

/// Divides a count of bytes represented as a u64 by the given `Bandwidth`
///
/// Since `Bandwidth` is a rate of bytes over a time period, this division
/// results in a `Duration` being returned, representing how long a path
/// with the given `Bandwidth` would take to transmit the given number of
/// bytes.
impl core::ops::Div<Bandwidth> for u64 {
    type Output = Duration;

    fn div(self, rhs: Bandwidth) -> Self::Output {
        #[allow(clippy::suspicious_arithmetic_impl)]
        Duration::from_nanos(rhs.nanos_per_kibibyte.saturating_mul(self) >> KIBIBYTE_SHIFT)
    }
}

#[derive(Clone, Copy, Debug, Default)]
/// A bandwidth delivery rate estimate with associated metadata
pub struct RateSample {
    /// The length of the sampling interval
    pub interval: Duration,
    /// The amount of data in bytes marked as delivered over the sampling interval
    pub delivered_bytes: u64,
    /// The amount of data in bytes marked as lost over the sampling interval
    pub lost_bytes: u64,
    /// The number of packets marked as explicit congestion experienced over the sampling interval
    pub ecn_ce_count: u64,
    /// [PacketInfo::is_app_limited] from the most recent acknowledged packet
    pub is_app_limited: bool,
    /// [PacketInfo::delivered_bytes] from the most recent acknowledged packet
    pub prior_delivered_bytes: u64,
    /// [PacketInfo::bytes_in_flight] from the most recent acknowledged packet
    pub bytes_in_flight: u32,
    /// [PacketInfo::lost_bytes] from the most recent acknowledged packet
    pub prior_lost_bytes: u64,
    /// [PacketInfo::ecn_ce_count] from the most recent acknowledged packet
    pub prior_ecn_ce_count: u64,
}

impl RateSample {
    /// Updates the rate sample with the most recent acknowledged packet
    fn on_ack(&mut self, packet_info: PacketInfo) {
        debug_assert!(
            packet_info.delivered_bytes >= self.prior_delivered_bytes,
            "on_ack should only be called with the newest acknowledged packet"
        );

        self.is_app_limited = packet_info.is_app_limited;
        self.prior_delivered_bytes = packet_info.delivered_bytes;
        self.prior_lost_bytes = packet_info.lost_bytes;
        self.prior_ecn_ce_count = packet_info.ecn_ce_count;
        self.bytes_in_flight = packet_info.bytes_in_flight;
    }

    /// Gets the delivery rate of this rate sample
    pub fn delivery_rate(&self) -> Bandwidth {
        Bandwidth::new(self.delivered_bytes, self.interval)
    }
}

impl IntoEvent<event::builder::RateSample> for RateSample {
    #[inline]
    fn into_event(self) -> event::builder::RateSample {
        event::builder::RateSample {
            interval: self.interval.into_event(),
            delivered_bytes: self.delivered_bytes.into_event(),
            lost_bytes: self.lost_bytes.into_event(),
            ecn_ce_count: self.ecn_ce_count.into_event(),
            is_app_limited: self.is_app_limited.into_event(),
            prior_delivered_bytes: self.prior_delivered_bytes.into_event(),
            bytes_in_flight: self.bytes_in_flight.into_event(),
            prior_lost_bytes: self.prior_lost_bytes.into_event(),
            prior_ecn_ce_count: self.prior_ecn_ce_count.into_event(),
            delivery_rate_bytes_per_second: self.delivery_rate().as_bytes_per_second().into_event(),
        }
    }
}

/// Bandwidth estimator as defined in [Delivery Rate Estimation](https://datatracker.ietf.org/doc/draft-cheng-iccrg-delivery-rate-estimation/)
/// and [BBR Congestion Control](https://datatracker.ietf.org/doc/draft-cardwell-iccrg-bbr-congestion-control/).
#[derive(Clone, Debug, Default)]
pub struct Estimator {
    //= https://tools.ietf.org/id/draft-cheng-iccrg-delivery-rate-estimation-02#2.2
    //# The amount of data delivered MAY be tracked in units of either octets or packets.
    /// The total amount of data in bytes delivered so far over the lifetime of the path, not including
    /// non-congestion-controlled packets such as pure ACK packets.
    delivered_bytes: u64,
    /// The timestamp when delivered_bytes was last updated, or if the connection
    /// was recently idle, the send time of the first packet sent after resuming from idle.
    delivered_time: Option<Timestamp>,
    /// The total amount of data in bytes declared lost so far over the lifetime of the path, not including
    /// non-congestion-controlled packets such as pure ACK packets.
    lost_bytes: u64,
    /// The total amount of explicit congestion experienced packets over the lifetime of the path
    ecn_ce_count: u64,
    /// The send time of the packet that was most recently marked as delivered, or if the connection
    /// was recently idle, the send time of the first packet sent after resuming from idle.
    first_sent_time: Option<Timestamp>,
    /// The `delivered_bytes` that marks the end of the current application-limited period, or
    /// `None` if the connection is not currently application-limited.
    app_limited_delivered_bytes: Option<u64>,
    rate_sample: RateSample,
}

impl Estimator {
    /// Gets the total amount of data in bytes delivered so far over the lifetime of the path, not including
    /// non-congestion-controlled packets such as pure ACK packets.
    pub fn delivered_bytes(&self) -> u64 {
        self.delivered_bytes
    }

    /// Gets the total amount of data in bytes lost so far over the lifetime of the path
    pub fn lost_bytes(&self) -> u64 {
        self.lost_bytes
    }

    /// Gets the latest [RateSample]
    pub fn rate_sample(&self) -> RateSample {
        self.rate_sample
    }

    /// Returns true if the path is currently in an application-limited period
    pub fn is_app_limited(&self) -> bool {
        self.app_limited_delivered_bytes.is_some()
    }

    /// Called when a packet is transmitted
    pub fn on_packet_sent(
        &mut self,
        prior_bytes_in_flight: u32,
        sent_bytes: usize,
        app_limited: Option<bool>,
        now: Timestamp,
    ) -> PacketInfo {
        //= https://tools.ietf.org/id/draft-cheng-iccrg-delivery-rate-estimation-02#3.2
        //# If there are no packets in flight yet, then we can start the delivery rate interval
        //# at the current time, since we know that any ACKs after now indicate that the network
        //# was able to deliver those packets completely in the sampling interval between now
        //# and the next ACK.

        //= https://tools.ietf.org/id/draft-cheng-iccrg-delivery-rate-estimation-02#3.2
        //# Upon each packet transmission, the sender executes the following steps:
        //#
        //# SendPacket(Packet P):
        //#   if (SND.NXT == SND.UNA)  /* no packets in flight yet? */
        //#     C.first_sent_time  = C.delivered_time = Now()
        //#   P.first_sent_time = C.first_sent_time
        //#   P.delivered_time  = C.delivered_time
        //#   P.delivered       = C.delivered
        //#   P.is_app_limited  = (C.app_limited != 0)
        if prior_bytes_in_flight == 0 {
            self.first_sent_time = Some(now);
            self.delivered_time = Some(now);
        }

        let bytes_in_flight = prior_bytes_in_flight.saturating_add(sent_bytes as u32);

        if app_limited.unwrap_or(true) {
            self.on_app_limited(bytes_in_flight);
        }

        PacketInfo {
            delivered_bytes: self.delivered_bytes,
            delivered_time: self
                .delivered_time
                .expect("initialized on first sent packet"),
            lost_bytes: self.lost_bytes,
            ecn_ce_count: self.ecn_ce_count,
            first_sent_time: self
                .first_sent_time
                .expect("initialized on first sent packet"),
            bytes_in_flight,
            is_app_limited: self.app_limited_delivered_bytes.is_some(),
        }
    }

    //= https://tools.ietf.org/id/draft-cheng-iccrg-delivery-rate-estimation-02#3.3
    //# For each packet that was newly SACKed or ACKed, UpdateRateSample() updates the
    //# rate sample based on a snapshot of connection delivery information from the time
    //# at which the packet was last transmitted.
    /// Called for each acknowledgement of one or more packets
    pub fn on_ack<Pub: Publisher>(
        &mut self,
        bytes_acknowledged: usize,
        newest_acked_time_sent: Timestamp,
        newest_acked_packet_info: PacketInfo,
        now: Timestamp,
        publisher: &mut Pub,
    ) {
        self.delivered_bytes += bytes_acknowledged as u64;
        self.delivered_time = Some(now);

        let is_app_limited_period_over =
            |app_limited_bytes| self.delivered_bytes > app_limited_bytes;
        if self
            .app_limited_delivered_bytes
            .map_or(false, is_app_limited_period_over)
        {
            // Clear app-limited field if bubble is ACKed and gone
            self.app_limited_delivered_bytes = None;
        }

        //= https://tools.ietf.org/id/draft-cheng-iccrg-delivery-rate-estimation-02#3.3
        //# UpdateRateSample() is invoked multiple times when a stretched ACK acknowledges
        //# multiple data packets. In this case we use the information from the most recently
        //# sent packet, i.e., the packet with the highest "P.delivered" value.
        if self.rate_sample.prior_delivered_bytes == 0
            || newest_acked_packet_info.delivered_bytes > self.rate_sample.prior_delivered_bytes
        {
            // Update info using the newest packet
            self.rate_sample.on_ack(newest_acked_packet_info);
            self.first_sent_time = Some(newest_acked_time_sent);

            let send_elapsed = newest_acked_time_sent - newest_acked_packet_info.first_sent_time;
            let ack_elapsed = now - newest_acked_packet_info.delivered_time;

            //= https://tools.ietf.org/id/draft-cheng-iccrg-delivery-rate-estimation-02#2.2.4
            //# Since it is physically impossible to have data delivered faster than it is sent
            //# in a sustained fashion, when the estimator notices that the ack_rate for a flight
            //# is faster than the send rate for the flight, it filters out the implausible ack_rate
            //# by capping the delivery rate sample to be no higher than the send rate.
            self.rate_sample.interval = max(send_elapsed, ack_elapsed);
        }

        self.rate_sample.delivered_bytes =
            self.delivered_bytes - self.rate_sample.prior_delivered_bytes;
        // Lost bytes and ECN CE Count are updated here as well as in `on_loss` and `on_explicit_congestion`
        // so the values are up to date even when no loss or ECN CE markings are received.
        self.rate_sample.lost_bytes = self.lost_bytes - self.rate_sample.prior_lost_bytes;
        self.rate_sample.ecn_ce_count = self.ecn_ce_count - self.rate_sample.prior_ecn_ce_count;

        publisher.on_delivery_rate_sampled(self.rate_sample);
    }

    /// Called when packets are declared lost
    #[inline]
    pub fn on_loss(&mut self, lost_bytes: usize) {
        self.lost_bytes += lost_bytes as u64;
        self.rate_sample.lost_bytes = self.lost_bytes - self.rate_sample.prior_lost_bytes;

        // Move the app-limited period earlier as the lost bytes will not be delivered
        if let Some(ref mut app_limited_delivered_bytes) = self.app_limited_delivered_bytes {
            *app_limited_delivered_bytes =
                app_limited_delivered_bytes.saturating_sub(lost_bytes as u64)
        }
    }

    /// Called when packets are discarded
    #[inline]
    pub fn on_packet_discarded(&mut self, bytes_sent: usize) {
        // Move the app-limited period earlier as the discarded bytes will not be delivered
        if let Some(ref mut app_limited_delivered_bytes) = self.app_limited_delivered_bytes {
            *app_limited_delivered_bytes =
                app_limited_delivered_bytes.saturating_sub(bytes_sent as u64)
        }
    }

    /// Called each time explicit congestion is recorded
    #[inline]
    pub fn on_explicit_congestion(&mut self, ce_count: u64) {
        self.ecn_ce_count += ce_count;
        self.rate_sample.ecn_ce_count = self.ecn_ce_count - self.rate_sample.prior_ecn_ce_count;
    }

    /// Mark the path as app limited until the given `bytes_in_flight` are acknowledged
    pub fn on_app_limited(&mut self, bytes_in_flight: u32) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.3.4.4
        //# MarkConnectionAppLimited():
        //#   C.app_limited =
        //#     (C.delivered + packets_in_flight) ? : 1
        self.app_limited_delivered_bytes = Some(self.delivered_bytes + bytes_in_flight as u64);
    }

    #[cfg(test)]
    pub fn set_rate_sample_for_test(&mut self, rate_sample: RateSample) {
        self.rate_sample = rate_sample
    }

    #[cfg(test)]
    pub fn set_delivered_bytes_for_test(&mut self, delivered_bytes: u64) {
        self.delivered_bytes = delivered_bytes
    }
}

#[cfg(test)]
mod tests;


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

pub use estimator::*;
mod estimator;


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

use super::*;
use crate::{
    event, path,
    recovery::congestion_controller::PathPublisher,
    time::{Clock, NoopClock},
};

#[test]
fn bandwidth() {
    let result = Bandwidth::new(1000, Duration::from_secs(10));

    // 10 seconds = 10^10 nanoseconds, nanos_per_byte = 10^10 / 1000 = 10_000_000
    // nanos_per_kibibyte = 10_000_000 * 1024 =
    assert_eq!(10_240_000_000, result.nanos_per_kibibyte);
}

#[test]
fn bandwidth_saturating() {
    let result = Bandwidth::new(u64::MAX, Duration::from_secs(1));

    assert_eq!(Bandwidth::INFINITY, result);
}

#[test]
fn bandwidth_zero_interval() {
    let result = Bandwidth::new(500, Duration::ZERO);

    assert_eq!(Bandwidth::ZERO, result);
}
/// Test that the maximum supported interval (~104 days) does not panic
#[test]
fn bandwidth_max_interval() {
    let result = Bandwidth::new(1, Duration::from_nanos(1 << 53));

    assert!(result > Bandwidth::ZERO);
}

/// Test that a value larger than the maximum supported interval does not panic
#[test]
fn bandwidth_interval_too_long() {
    let result = Bandwidth::new(1, Duration::from_nanos(1 << 54));

    assert_eq!(Bandwidth::ZERO, result);
}

#[test]
fn bandwidth_mul_ratio() {
    let bandwidth = Bandwidth::new(4000, Duration::from_secs(1));

    let result = bandwidth * Ratio::new(3, 4);

    assert_eq!(result, Bandwidth::new(3000, Duration::from_secs(1)));
}

#[test]
fn bandwidth_zero_mul_ratio() {
    assert_eq!(Bandwidth::ZERO, Bandwidth::ZERO * Ratio::new(3, 7));
}

#[test]
fn bandwidth_mul_duration() {
    let bandwidth = Bandwidth::new(7000, Duration::from_secs(2));

    let result = bandwidth * Duration::from_secs(10);

    assert_eq!(result, 35000);
}

#[test]
fn bandwidth_mul_duration_too_long() {
    let bandwidth = Bandwidth::new(1, Duration::from_nanos(1));

    let result = bandwidth * Duration::from_nanos(1 << 53);

    assert!(result < u64::MAX);

    let result = bandwidth * Duration::from_nanos(1 << 54);

    assert_eq!(result, u64::MAX);
}

#[test]
fn bandwidth_mul_saturation() {
    let bandwidth = Bandwidth::INFINITY;

    let result = bandwidth * Duration::from_secs(10);

    assert_eq!(result, u64::MAX);
}

#[test]
fn u64_div_bandwidth() {
    let bandwidth = Bandwidth::new(10_000, Duration::from_secs(1));
    let bytes = 200_000;
    assert_eq!(bytes / bandwidth, Duration::from_secs(20));

    let bandwidth = Bandwidth::new(10_000, Duration::from_secs(1));
    let bytes = 2_000;
    assert_eq!(bytes / bandwidth, Duration::from_millis(200));
}

#[test]
fn bandwidth_ordering() {
    let low = Bandwidth::new(10_000, Duration::from_secs(1));
    let high = Bandwidth::new(20_000, Duration::from_secs(1));

    assert!(high > low);
    assert_eq!(high, low.max(high));
}

// first_sent_time and delivered_time typically hold values from recently acknowledged packets. However,
// when  no packet has been sent yet, or there are no packets currently in flight, these values are initialized
// with the time when a packet is sent. This test confirms first_sent_time and delivered_time are
// initialized properly on the first packet sent, and on the first packet sent after an idle period.
#[test]
fn on_packet_sent_timestamp_initialization() {
    let t0 = NoopClock.get_time();
    let mut bw_estimator = Estimator::default();

    // Test that first_sent_time and delivered_time are updated on the first sent packet
    let packet_info = bw_estimator.on_packet_sent(0, 0, None, t0);
    assert_eq!(t0, packet_info.first_sent_time);
    assert_eq!(t0, packet_info.delivered_time);
    assert_eq!(Some(t0), bw_estimator.first_sent_time);
    assert_eq!(Some(t0), bw_estimator.delivered_time);

    // Test that first_sent_time and delivered_time are not updated if packets are in flight
    let t1 = t0 + Duration::from_secs(1);
    let packet_info = bw_estimator.on_packet_sent(1500, 0, None, t1);
    assert_eq!(t0, packet_info.first_sent_time);
    assert_eq!(t0, packet_info.delivered_time);
    assert_eq!(Some(t0), bw_estimator.first_sent_time);
    assert_eq!(Some(t0), bw_estimator.delivered_time);

    // Test that first_sent_time and delivered_time are updated after an idle period
    let t2 = t0 + Duration::from_secs(2);
    let packet_info = bw_estimator.on_packet_sent(0, 0, None, t2);
    assert_eq!(t2, packet_info.first_sent_time);
    assert_eq!(t2, packet_info.delivered_time);
    assert_eq!(Some(t2), bw_estimator.first_sent_time);
    assert_eq!(Some(t2), bw_estimator.delivered_time);
}

#[test]
fn on_packet_sent() {
    let first_sent_time = NoopClock.get_time();
    let delivered_time = first_sent_time + Duration::from_secs(1);
    let mut bw_estimator = Estimator {
        delivered_bytes: 15000,
        delivered_time: Some(delivered_time),
        lost_bytes: 100,
        ecn_ce_count: 5,
        first_sent_time: Some(first_sent_time),
        app_limited_delivered_bytes: None,
        rate_sample: Default::default(),
    };

    let packet_info = bw_estimator.on_packet_sent(500, 100, Some(true), first_sent_time);
    assert_eq!(first_sent_time, packet_info.first_sent_time);
    assert_eq!(delivered_time, packet_info.delivered_time);
    assert_eq!(15000, packet_info.delivered_bytes);
    assert_eq!(100, packet_info.lost_bytes);
    assert_eq!(5, packet_info.ecn_ce_count);
    assert!(packet_info.is_app_limited);
    assert_eq!(500 + 100, packet_info.bytes_in_flight);
    assert_eq!(
        Some(500 + 15000 + 100),
        bw_estimator.app_limited_delivered_bytes
    );
}

#[test]
fn app_limited() {
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let first_sent_time = NoopClock.get_time();
    let delivered_time = first_sent_time + Duration::from_secs(1);
    let mut bw_estimator = Estimator {
        delivered_bytes: 15000,
        delivered_time: Some(delivered_time),
        lost_bytes: 100,
        ecn_ce_count: 5,
        first_sent_time: Some(first_sent_time),
        app_limited_delivered_bytes: None,
        rate_sample: Default::default(),
    };

    // Packet is sent while app-limited, starting the app limited period
    let packet_info = bw_estimator.on_packet_sent(1500, 100, Some(true), first_sent_time);
    assert!(packet_info.is_app_limited);
    assert_eq!(
        Some(1500 + 15000 + 100),
        bw_estimator.app_limited_delivered_bytes
    );

    // Packet is sent while not app-limited, but the app limited continues until all previous bytes in flight have been acknowledged
    let packet_info = bw_estimator.on_packet_sent(500, 100, Some(false), first_sent_time);
    assert!(packet_info.is_app_limited);
    assert_eq!(
        Some(1500 + 15000 + 100),
        bw_estimator.app_limited_delivered_bytes
    );

    // Packet is sent while app-limited is not determined, this should default to app-limited
    let packet_info = bw_estimator.on_packet_sent(2500, 100, None, first_sent_time);
    assert!(packet_info.is_app_limited);
    assert_eq!(
        Some(2500 + 15000 + 100),
        bw_estimator.app_limited_delivered_bytes
    );

    let packet_info = PacketInfo {
        delivered_bytes: bw_estimator.delivered_bytes,
        delivered_time,
        lost_bytes: 0,
        ecn_ce_count: 0,
        first_sent_time,
        bytes_in_flight: 1500,
        is_app_limited: false,
    };

    // Acknowledge all the bytes that were inflight when the app-limited period began
    bw_estimator.on_ack(
        2600,
        delivered_time,
        packet_info,
        delivered_time,
        &mut publisher,
    );
    // Still app_limited, since we need bytes to be acknowledged after the app limited period
    assert_eq!(
        Some(2500 + 100 + 15000),
        bw_estimator.app_limited_delivered_bytes
    );

    // Acknowledge one more byte
    bw_estimator.on_ack(
        1,
        delivered_time,
        packet_info,
        delivered_time,
        &mut publisher,
    );
    // Now the app limited period is over
    assert_eq!(None, bw_estimator.app_limited_delivered_bytes);
}

#[test]
fn on_packet_ack_rate_sample() {
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let t0 = NoopClock.get_time() + Duration::from_secs(60);
    let t1 = t0 + Duration::from_secs(1);
    let t2 = t0 + Duration::from_secs(2);
    let mut bw_estimator = Estimator::default();

    // Send three packets. In between each send, other packets were acknowledged, lost or had ECN CE,
    // and thus the delivered_bytes amount, lost_bytes amount, and ecn ce count is increased.
    let packet_1 = bw_estimator.on_packet_sent(0, 100, Some(false), t0);
    bw_estimator.delivered_bytes = 100000;
    bw_estimator.lost_bytes = 100;
    bw_estimator.ecn_ce_count = 5;
    let packet_2 = bw_estimator.on_packet_sent(1500, 100, Some(true), t1);
    bw_estimator.delivered_bytes = 200000;
    bw_estimator.lost_bytes = 150;
    bw_estimator.ecn_ce_count = 15;
    let packet_3 = bw_estimator.on_packet_sent(3000, 100, Some(false), t2);

    let now = t0 + Duration::from_secs(10);
    let delivered_bytes = bw_estimator.delivered_bytes;
    bw_estimator.on_ack(1500, t0, packet_1, now, &mut publisher);

    assert_eq!(bw_estimator.delivered_bytes, delivered_bytes + 1500);
    assert_eq!(bw_estimator.delivered_time, Some(now));

    // Rate sample is updated since this is the first packet delivered
    assert_eq!(
        packet_1.delivered_bytes,
        bw_estimator.rate_sample.prior_delivered_bytes
    );
    assert_eq!(
        packet_1.lost_bytes,
        bw_estimator.rate_sample.prior_lost_bytes
    );
    assert_eq!(
        packet_1.is_app_limited,
        bw_estimator.rate_sample.is_app_limited
    );
    assert_eq!(
        packet_1.bytes_in_flight,
        bw_estimator.rate_sample.bytes_in_flight
    );
    assert_eq!(Some(t0), bw_estimator.first_sent_time);
    assert_eq!(now - t0, bw_estimator.rate_sample.interval);

    // Delivered bytes, lost bytes, and ECN CE counts reflect the total lifetime counts,
    // since there was no prior delivered packets
    assert_eq!(
        bw_estimator.delivered_bytes,
        bw_estimator.rate_sample.delivered_bytes
    );
    assert_eq!(bw_estimator.lost_bytes, bw_estimator.rate_sample.lost_bytes);
    assert_eq!(
        bw_estimator.ecn_ce_count,
        bw_estimator.rate_sample.ecn_ce_count
    );

    // Ack a newer packet
    let now = now + Duration::from_secs(1);
    let delivered_bytes = bw_estimator.delivered_bytes;
    bw_estimator.on_ack(1500, t2, packet_3, now, &mut publisher);

    assert_eq!(bw_estimator.delivered_bytes, delivered_bytes + 1500);
    assert_eq!(bw_estimator.delivered_time, Some(now));

    // Rate sample is updated since this packet is newer (has a higher delivered_bytes)
    assert!(packet_3.delivered_bytes > packet_1.delivered_bytes);
    assert_eq!(
        packet_3.delivered_bytes,
        bw_estimator.rate_sample.prior_delivered_bytes
    );
    assert_eq!(
        packet_3.lost_bytes,
        bw_estimator.rate_sample.prior_lost_bytes
    );
    assert_eq!(
        packet_3.is_app_limited,
        bw_estimator.rate_sample.is_app_limited
    );
    assert_eq!(
        packet_3.bytes_in_flight,
        bw_estimator.rate_sample.bytes_in_flight
    );
    assert_eq!(Some(t2), bw_estimator.first_sent_time);
    assert_eq!(now - t0, bw_estimator.rate_sample.interval);

    // Now the delivered_bytes in the rate sample should reflect the current lifetime delivered bytes (200000 + 1500 + 1500)
    // minus the delivered bytes value at the time this packet was transmitted (200000)
    assert_eq!(3000, bw_estimator.rate_sample.delivered_bytes);
    // Lost bytes and ECN CE count are 0 since no additional loss or ECN CE was received since this packet was transmitted
    assert_eq!(0, bw_estimator.rate_sample.lost_bytes);
    assert_eq!(0, bw_estimator.rate_sample.ecn_ce_count);

    // Ack an older packet
    let now = now + Duration::from_secs(1);
    let delivered_bytes = bw_estimator.delivered_bytes;
    bw_estimator.on_ack(1500, t1, packet_2, now, &mut publisher);

    assert_eq!(bw_estimator.delivered_bytes, delivered_bytes + 1500);
    assert_eq!(bw_estimator.delivered_time, Some(now));

    //= https://tools.ietf.org/id/draft-cheng-iccrg-delivery-rate-estimation-02#3.3
    //= type=test
    //# UpdateRateSample() is invoked multiple times when a stretched ACK acknowledges
    //# multiple data packets. In this case we use the information from the most recently
    //# sent packet, i.e., the packet with the highest "P.delivered" value.

    // Rate sample is not updated since this packet is older than the current sample
    assert_eq!(
        packet_3.delivered_bytes,
        bw_estimator.rate_sample.prior_delivered_bytes
    );
    assert_eq!(
        packet_3.lost_bytes,
        bw_estimator.rate_sample.prior_lost_bytes
    );
    assert_eq!(
        packet_3.is_app_limited,
        bw_estimator.rate_sample.is_app_limited
    );
    assert_eq!(
        packet_3.bytes_in_flight,
        bw_estimator.rate_sample.bytes_in_flight
    );
    assert_eq!(Some(t2), bw_estimator.first_sent_time);

    // Delivered bytes is increased to include packet 2, as it is within the sampling interval
    assert_eq!(4500, bw_estimator.rate_sample.delivered_bytes);
    assert_eq!(0, bw_estimator.rate_sample.lost_bytes);
    assert_eq!(0, bw_estimator.rate_sample.ecn_ce_count);
}

//= https://tools.ietf.org/id/draft-cheng-iccrg-delivery-rate-estimation-02#2.2.4
//= type=test
//# Since it is physically impossible to have data delivered faster than it is sent
//# in a sustained fashion, when the estimator notices that the ack_rate for a flight
//# is faster than the send rate for the flight, it filters out the implausible ack_rate
//# by capping the delivery rate sample to be no higher than the send rate.
#[test]
fn on_packet_ack_implausible_ack_rate() {
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let t0 = NoopClock.get_time();
    let mut bw_estimator = Estimator::default();

    // A packet is sent and acknowledged 4 seconds later
    let packet_info = bw_estimator.on_packet_sent(0, 100, Some(false), t0);
    let t4 = t0 + Duration::from_secs(4);
    bw_estimator.on_ack(1500, t0, packet_info, t4, &mut publisher);

    // A packet is sent and acknowledged 1 second later
    let t5 = t0 + Duration::from_secs(5);
    let packet_info = bw_estimator.on_packet_sent(1500, 100, Some(false), t5);
    let now = t0 + Duration::from_secs(6);
    bw_estimator.on_ack(
        1500,
        t0 + Duration::from_secs(5),
        packet_info,
        now,
        &mut publisher,
    );

    let send_elapsed = t5 - packet_info.first_sent_time;
    let ack_elapsed = now - packet_info.delivered_time;

    // The amount of time to send the packets was greater than the time to
    // acknowledge them, indicating the ack rate would be implausible.
    assert!(send_elapsed > ack_elapsed);

    // The rate sample interval is based on the send_elapsed time since the
    // ack_elapsed time was implausible
    assert_eq!(send_elapsed, bw_estimator.rate_sample.interval);
}

#[test]
fn on_packet_loss() {
    let mut bw_estimator = Estimator::default();

    assert_eq!(0, bw_estimator.lost_bytes);
    assert_eq!(0, bw_estimator.rate_sample.lost_bytes);

    bw_estimator.on_loss(500);

    assert_eq!(500, bw_estimator.lost_bytes);
    assert_eq!(500, bw_estimator.rate_sample.lost_bytes);

    bw_estimator.on_loss(250);

    assert_eq!(750, bw_estimator.lost_bytes);
    assert_eq!(750, bw_estimator.rate_sample.lost_bytes);

    bw_estimator.on_app_limited(1000);
    bw_estimator.on_loss(250);
    assert_eq!(Some(750), bw_estimator.app_limited_delivered_bytes);

    bw_estimator.on_loss(1000);
    assert_eq!(Some(0), bw_estimator.app_limited_delivered_bytes);
}

#[test]
fn on_packet_discarded() {
    let mut bw_estimator = Estimator::default();

    bw_estimator.on_app_limited(1000);
    bw_estimator.on_packet_discarded(250);
    assert_eq!(Some(750), bw_estimator.app_limited_delivered_bytes);

    bw_estimator.on_packet_discarded(1000);
    assert_eq!(Some(0), bw_estimator.app_limited_delivered_bytes);
}

#[test]
fn on_explicit_congestion() {
    let mut bw_estimator = Estimator::default();

    assert_eq!(0, bw_estimator.ecn_ce_count);
    assert_eq!(0, bw_estimator.rate_sample.ecn_ce_count);

    bw_estimator.on_explicit_congestion(5);

    assert_eq!(5, bw_estimator.ecn_ce_count);
    assert_eq!(5, bw_estimator.rate_sample.ecn_ce_count);

    bw_estimator.on_explicit_congestion(3);

    assert_eq!(8, bw_estimator.ecn_ce_count);
    assert_eq!(8, bw_estimator.rate_sample.ecn_ce_count);
}

#[test]
fn as_bytes_per_second() {
    let bandwidth = Bandwidth::new(10_000, Duration::from_secs(1));

    assert_eq!(10_000, bandwidth.as_bytes_per_second());
    assert_eq!(0, Bandwidth::ZERO.as_bytes_per_second());
    assert_eq!(u64::MAX, Bandwidth::INFINITY.as_bytes_per_second());
}


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

use crate::{
    counter::{Counter, Saturating},
    recovery::{
        bandwidth::{Bandwidth, PacketInfo, RateSample},
        bbr::{data_rate, data_volume, round, BbrCongestionController},
    },
};

#[derive(Clone, Debug, Default)]
pub(crate) struct State {
    /// Tracks round trips for ensuring BBR reacts to congestion only once per round
    loss_round_counter: round::Counter,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.10
    //# a 1-round-trip max of delivered bandwidth (rs.delivery_rate)
    bw_latest: Bandwidth,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.10
    //# a 1-round-trip max of delivered volume of data (rs.delivered)
    inflight_latest: u64,
    /// The number of bursts of loss encountered during the current round trip
    loss_bursts_in_round: Counter<u8, Saturating>,
    /// True if packets were marked with ECN CE at any point in the current round trip
    ecn_in_round: bool,
}

impl State {
    /// Updates the congestion state based on the latest delivery signals
    ///
    /// Called near the start of ACK processing
    #[allow(clippy::too_many_arguments)]
    pub(super) fn update(
        &mut self,
        packet_info: PacketInfo,
        rate_sample: RateSample,
        delivered_bytes: u64,
        data_rate_model: &mut data_rate::Model,
        data_volume_model: &mut data_volume::Model,
        is_probing_for_bandwidth: bool,
        cwnd: u32,
        ecn_alpha: f64,
    ) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
        //# BBRUpdateLatestDeliverySignals():
        //#   BBR.loss_round_start = 0
        //#   BBR.bw_latest       = max(BBR.bw_latest,       rs.delivery_rate)
        //#   BBR.inflight_latest = max(BBR.inflight_latest, rs.delivered)
        //#   if (rs.prior_delivered >= BBR.loss_round_delivered)
        //#     BBR.loss_round_delivered = C.delivered
        //#     BBR.loss_round_start = 1

        self.loss_round_counter.on_ack(packet_info, delivered_bytes);
        self.bw_latest = self.bw_latest.max(rate_sample.delivery_rate());
        self.inflight_latest = self.inflight_latest.max(rate_sample.delivered_bytes);

        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
        //# BBRUpdateCongestionSignals():
        //#   BBRUpdateMaxBw()
        //#   if (rs.losses > 0)
        //#     BBR.loss_in_round = 1
        //#   if (!BBR.loss_round_start)
        //#     return  /* wait until end of round trip */
        //#   BBRAdaptLowerBoundsFromCongestion()
        //#   BBR.loss_in_round = 0

        data_rate_model.update_max_bw(rate_sample);

        if rate_sample.ecn_ce_count > 0 {
            self.ecn_in_round = true;
        }

        if self.loss_round_counter.round_start() && !is_probing_for_bandwidth {
            //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
            //# When not explicitly accelerating to probe for bandwidth (Drain, ProbeRTT,
            //# ProbeBW_DOWN, ProbeBW_CRUISE), BBR responds to loss by slowing down to some extent.

            //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
            //# BBRAdaptLowerBoundsFromCongestion():
            //#   if (BBRIsProbingBW())
            //#     return
            //#   if (BBR.loss_in_round())
            //#     BBRInitLowerBounds()
            //#     BBRLossLowerBounds()

            if self.loss_in_round() {
                // The following update_lower_bound method combines the functionality of
                // BBRInitLowerBounds() and BBRLossLowerBounds()
                data_rate_model.update_lower_bound(self.bw_latest);
            }

            // Update inflight_lo to the lower of the ECN and the Loss based values
            // if there is loss or ECN in the round
            data_volume_model.update_lower_bound(
                cwnd,
                self.inflight_latest,
                self.loss_in_round(),
                self.ecn_in_round,
                ecn_alpha,
            );
        }
    }

    /// Initializes the congestion state for the next round
    ///
    /// Called near the end of ACK processing
    pub(super) fn advance(&mut self, rate_sample: RateSample) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
        //# BBRAdvanceLatestDeliverySignals():
        //#   if (BBR.loss_round_start)
        //#     BBR.bw_latest       = rs.delivery_rate
        //#     BBR.inflight_latest = rs.delivered
        if self.loss_round_counter.round_start() {
            self.bw_latest = rate_sample.delivery_rate();
            self.inflight_latest = rate_sample.delivered_bytes;
            self.loss_bursts_in_round = Default::default();
            self.ecn_in_round = false;
        }
    }

    /// Resets the congestion signals
    pub(super) fn reset(&mut self) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
        //# BBRResetCongestionSignals():
        //#   BBR.loss_in_round = 0
        //#   BBR.bw_latest = 0
        //#   BBR.inflight_latest = 0
        self.loss_bursts_in_round = Default::default();
        self.ecn_in_round = false;
        self.bw_latest = Bandwidth::ZERO;
        self.inflight_latest = 0;
    }

    #[inline]
    pub(super) fn on_packet_lost(&mut self, delivered_bytes: u64, new_lost_burst: bool) {
        if !self.loss_in_round() {
            self.loss_round_counter.set_round_end(delivered_bytes);
        }

        if new_lost_burst {
            self.loss_bursts_in_round += 1;
        }
    }

    #[inline]
    pub(super) fn on_explicit_congestion(&mut self) {
        self.ecn_in_round = true;
    }

    #[inline]
    /// Returns true if this is the beginning of a new loss round
    pub(super) fn loss_round_start(&self) -> bool {
        self.loss_round_counter.round_start()
    }

    #[inline]
    /// Returns true if there was loss in the current round
    pub(super) fn loss_in_round(&self) -> bool {
        self.loss_bursts_in_round > 0
    }

    #[inline]
    /// Returns the number of loss busts in the current round
    pub(super) fn loss_bursts_in_round(&self) -> u8 {
        *self.loss_bursts_in_round
    }

    #[inline]
    /// Returns true if there were ECN CE marked packets in the current round
    pub(super) fn ecn_in_round(&self) -> bool {
        self.ecn_in_round
    }
}

/// Methods related to Congestion state
impl BbrCongestionController {
    /// Updates delivery and congestion signals according to
    /// BBRUpdateLatestDeliverySignals() and BBRUpdateCongestionSignals()
    #[inline]
    pub(super) fn update_latest_signals(&mut self, packet_info: PacketInfo) {
        self.congestion_state.update(
            packet_info,
            self.bw_estimator.rate_sample(),
            self.bw_estimator.delivered_bytes(),
            &mut self.data_rate_model,
            &mut self.data_volume_model,
            self.state.is_probing_for_bandwidth(),
            self.cwnd,
            self.ecn_state.alpha(),
        );
    }
}

#[cfg(test)]
pub mod testing {
    use crate::{
        recovery::{
            bandwidth::{Bandwidth, PacketInfo, RateSample},
            bbr::{congestion, data_rate, data_volume},
        },
        time::{Clock, NoopClock},
    };
    use std::time::Duration;

    /// Asserts that the given `congestion::State` has been reset
    pub(crate) fn assert_reset(state: congestion::State) {
        assert!(!state.loss_in_round());
        assert!(!state.ecn_in_round);
        assert_eq!(Bandwidth::ZERO, state.bw_latest);
        assert_eq!(0, state.inflight_latest);
    }

    /// Return congestion::State updated with data
    pub(crate) fn test_state() -> congestion::State {
        let mut state = congestion::State::default();

        let now = NoopClock.get_time();
        let packet_info = PacketInfo {
            delivered_bytes: 100,
            delivered_time: now,
            lost_bytes: 0,
            ecn_ce_count: 0,
            first_sent_time: now,
            bytes_in_flight: 0,
            is_app_limited: false,
        };
        let rate_sample = RateSample {
            interval: Duration::from_millis(10),
            delivered_bytes: 100,
            lost_bytes: 50,
            ..Default::default()
        };
        let mut data_rate_model = data_rate::Model::new();
        let mut data_volume_model = data_volume::Model::new();

        state.update(
            packet_info,
            rate_sample,
            500,
            &mut data_rate_model,
            &mut data_volume_model,
            false,
            100,
            1.0,
        );

        state
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::time::{Clock, NoopClock};
    use core::time::Duration;

    #[test]
    fn update() {
        let mut state = State::default();

        let now = NoopClock.get_time();
        let mut packet_info = PacketInfo {
            delivered_bytes: 100,
            delivered_time: now,
            lost_bytes: 0,
            ecn_ce_count: 0,
            first_sent_time: now,
            bytes_in_flight: 0,
            is_app_limited: false,
        };
        let rate_sample = RateSample {
            interval: Duration::from_millis(10),
            delivered_bytes: 100,
            lost_bytes: 50,
            ..Default::default()
        };
        let mut data_rate_model = data_rate::Model::new();
        let mut data_volume_model = data_volume::Model::new();

        state.on_packet_lost(100, true);
        state.update(
            packet_info,
            rate_sample,
            500,
            &mut data_rate_model,
            &mut data_volume_model,
            false,
            100,
            1.0,
        );

        assert!(state.loss_round_counter.round_start());
        assert_eq!(rate_sample.delivery_rate(), state.bw_latest);
        assert_eq!(rate_sample.delivered_bytes, state.inflight_latest);
        assert_eq!(rate_sample.delivery_rate(), data_rate_model.max_bw());
        // Since there was loss in the round, the lower bounds are updated
        assert_eq!(rate_sample.delivery_rate(), data_rate_model.bw_lo());
        assert_eq!(rate_sample.delivered_bytes, data_volume_model.inflight_lo());

        state.advance(rate_sample);
        // Loss and ecn in round are reset
        assert!(!state.loss_in_round());
        assert!(!state.ecn_in_round);

        packet_info.delivered_bytes = 400;

        let new_rate_sample = RateSample {
            interval: Duration::from_millis(10),
            delivered_bytes: 1000,
            lost_bytes: 50,
            ecn_ce_count: 5,
            ..Default::default()
        };

        state.on_packet_lost(500, true);
        state.update(
            packet_info,
            new_rate_sample,
            500,
            &mut data_rate_model,
            &mut data_volume_model,
            false,
            100,
            1.0,
        );

        assert!(!state.loss_round_counter.round_start());
        assert_eq!(new_rate_sample.delivery_rate(), state.bw_latest);
        assert_eq!(new_rate_sample.delivered_bytes, state.inflight_latest);
        assert_eq!(new_rate_sample.delivery_rate(), data_rate_model.max_bw());
        // It is not the start of a round, so lower bounds are not updated
        assert_eq!(rate_sample.delivery_rate(), data_rate_model.bw_lo());
        assert_eq!(rate_sample.delivered_bytes, data_volume_model.inflight_lo());
        assert!(state.loss_in_round());
        assert!(state.ecn_in_round);

        // This packet ends the rounds
        packet_info.delivered_bytes = 500;

        state.update(
            packet_info,
            new_rate_sample,
            500,
            &mut data_rate_model,
            &mut data_volume_model,
            true, // we are probing bw, so lower bounds should not update
            100,
            1.0,
        );

        assert!(state.loss_round_counter.round_start());
        // we are probing bw, so lower bounds should not update
        assert_eq!(rate_sample.delivery_rate(), data_rate_model.bw_lo());
        assert_eq!(rate_sample.delivered_bytes, data_volume_model.inflight_lo());
        // loss and ecn in round are still reset though
        state.advance(rate_sample);
        assert!(!state.loss_in_round());
        assert!(!state.ecn_in_round);
    }

    #[test]
    fn advance() {
        let mut state = State::default();

        let now = NoopClock.get_time();
        let packet_info = PacketInfo {
            delivered_bytes: 100,
            delivered_time: now,
            lost_bytes: 0,
            ecn_ce_count: 0,
            first_sent_time: now,
            bytes_in_flight: 0,
            is_app_limited: false,
        };
        let mut rate_sample = RateSample {
            interval: Duration::from_millis(10),
            delivered_bytes: 100,
            ..Default::default()
        };

        state.on_packet_lost(100, true);
        state.on_packet_lost(100, true);
        state.on_explicit_congestion();
        state.update(
            packet_info,
            rate_sample,
            100,
            &mut data_rate::Model::new(),
            &mut data_volume::Model::new(),
            false,
            100,
            1.0,
        );

        assert!(state.loss_round_counter.round_start());
        assert_eq!(rate_sample.delivery_rate(), state.bw_latest);
        assert_eq!(rate_sample.delivered_bytes, state.inflight_latest);
        assert_eq!(2, state.loss_bursts_in_round());
        assert!(state.ecn_in_round);

        rate_sample.delivered_bytes = 500;
        state.advance(rate_sample);

        assert_eq!(rate_sample.delivery_rate(), state.bw_latest);
        assert_eq!(rate_sample.delivered_bytes, state.inflight_latest);
        assert_eq!(0, state.loss_bursts_in_round());
        assert!(!state.ecn_in_round);
    }

    #[test]
    fn reset() {
        let mut state = State {
            loss_round_counter: Default::default(),
            loss_bursts_in_round: Counter::new(10),
            inflight_latest: 100,
            bw_latest: Bandwidth::INFINITY,
            ecn_in_round: true,
        };

        state.reset();

        assert!(!state.loss_in_round());
        assert!(!state.ecn_in_round);
        assert_eq!(Bandwidth::ZERO, state.bw_latest);
        assert_eq!(0, state.inflight_latest);
    }
}


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

use crate::recovery::{
    bandwidth::{Bandwidth, RateSample},
    bbr::{windowed_filter::WindowedMaxFilter, BETA},
};

//= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.9.1
//# The data rate model parameters together estimate both the sending rate required to reach the
//# full bandwidth available to the flow (BBR.max_bw), and the maximum pacing rate control parameter
//# that is consistent with the queue pressure objective (BBR.bw).

#[derive(Clone, Debug)]
pub(crate) struct Model {
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.9.1
    //# The windowed maximum recent bandwidth sample - obtained using the BBR delivery rate sampling
    //# algorithm [draft-cheng-iccrg-delivery-rate-estimation] - measured during the current or
    //# previous bandwidth probing cycle (or during Startup, if the flow is still in that state).
    max_bw_filter: WindowedMaxFilter<Bandwidth, core::num::Wrapping<u8>, core::num::Wrapping<u8>>,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.9.1
    //# The long-term maximum sending bandwidth that the algorithm estimates will produce acceptable
    //# queue pressure, based on signals in the current or previous bandwidth probing cycle, as
    //# measured by loss.
    bw_hi: Bandwidth,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.9.1
    //# The short-term maximum sending bandwidth that the algorithm estimates is safe for matching
    //# the current network path delivery rate, based on any loss signals in the current bandwidth
    //# probing cycle. This is generally lower than max_bw or bw_hi (thus the name).
    bw_lo: Bandwidth,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.9.1
    //# The maximum sending bandwidth that the algorithm estimates is appropriate for matching the
    //# current network path delivery rate, given all available signals in the model, at any time
    //# scale. It is the min() of max_bw, bw_hi, and bw_lo.
    bw: Bandwidth,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.11
    //# The virtual time used by the BBR.max_bw filter window.
    cycle_count: core::num::Wrapping<u8>,
}

impl Model {
    /// Constructs a new `data_rate::Model`
    pub fn new() -> Self {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.11
        //# The filter window length for BBR.MaxBwFilter = 2 (representing up to 2 ProbeBW cycles,
        //# the current cycle and the previous full cycle)
        const MAX_BW_FILTER_LEN: core::num::Wrapping<u8> = core::num::Wrapping(2);

        Self {
            max_bw_filter: WindowedMaxFilter::new(MAX_BW_FILTER_LEN),
            bw_hi: Bandwidth::INFINITY,
            bw_lo: Bandwidth::INFINITY,
            bw: Bandwidth::ZERO,
            cycle_count: Default::default(),
        }
    }

    /// The windowed maximum recent bandwidth sample
    pub fn max_bw(&self) -> Bandwidth {
        self.max_bw_filter.value().unwrap_or(Bandwidth::ZERO)
    }

    /// The long-term maximum sending bandwidth that the algorithm estimates
    /// will produce acceptable queue pressure
    pub fn bw_hi(&self) -> Bandwidth {
        self.bw_hi
    }

    /// The short-term maximum sending bandwidth that the algorithm estimates
    /// is safe for matching the current network path delivery rate
    pub fn bw_lo(&self) -> Bandwidth {
        self.bw_lo
    }

    /// The maximum sending bandwidth that the algorithm estimates is appropriate for
    /// matching the current network path delivery rate
    pub fn bw(&self) -> Bandwidth {
        self.bw
    }

    /// Increments the virtual time tracked for counting cyclical progression through ProbeBW cycles
    pub fn advance_max_bw_filter(&mut self) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.2.5
        //# BBRAdvanceMaxBwFilter():
        //#   BBR.cycle_count++
        self.cycle_count += core::num::Wrapping(1)
    }

    /// Updates `max_bw` with the given `rate_sample`
    pub fn update_max_bw(&mut self, rate_sample: RateSample) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.2.4
        //# BBRUpdateMaxBw()
        //#   BBRUpdateRound()
        //#   if (rs.delivery_rate >= BBR.max_bw || !rs.is_app_limited)
        //#       BBR.max_bw = update_windowed_max_filter(
        //#                     filter=BBR.MaxBwFilter,
        //#                     value=rs.delivery_rate,
        //#                     time=BBR.cycle_count,
        //#                     window_length=MaxBwFilterLen)

        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.2.3
        //# By default, the estimator discards application-limited samples, since by definition they
        //# reflect application limits.  However, the estimator does use application-limited samples
        //# if the measured delivery rate happens to be larger than the current BBR.max_bw estimate,
        //# since this indicates the current BBR.Max_bw estimate is too low.
        if rate_sample.delivery_rate() > self.max_bw() || !rate_sample.is_app_limited {
            self.max_bw_filter
                .update(rate_sample.delivery_rate(), self.cycle_count);
        }
    }

    /// Updates `bw_hi` with the given `bw`
    #[allow(dead_code)] // TODO: See note in probe_bw.rs about updating bw_hi
    pub fn update_upper_bound(&mut self, bw: Bandwidth) {
        self.bw_hi = bw
    }

    /// Updates `bw_lo` with the given `bw` if it exceeds the current `bw_lo` * `bbr::BETA`
    pub fn update_lower_bound(&mut self, bw: Bandwidth) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
        //# BBRInitLowerBounds():
        //#   if (BBR.bw_lo == Infinity)
        //#     BBR.bw_lo = BBR.max_bw
        if self.bw_lo == Bandwidth::INFINITY {
            self.bw_lo = self.max_bw()
        }

        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
        //# BBRLossLowerBounds():
        //#   BBR.bw_lo       = max(BBR.bw_latest,
        //#                         BBRBeta * BBR.bw_lo)
        self.bw_lo = bw.max(self.bw_lo * BETA);
    }

    /// Resets `bw_lo` to its initial value
    pub fn reset_lower_bound(&mut self) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
        //# BBRResetLowerBounds():
        //#   BBR.bw_lo       = Infinity
        self.bw_lo = Bandwidth::INFINITY
    }

    /// Bounds `bw` to min(`max_bw`, `bw_lo`, `bw_hi)
    pub fn bound_bw_for_model(&mut self) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
        //# BBRBoundBWForModel():
        //#   BBR.bw = min(BBR.max_bw, BBR.bw_lo, BBR.bw_hi)
        self.bw = self.max_bw().min(self.bw_lo()).min(self.bw_hi())
    }

    #[cfg(test)]
    pub fn cycle_count(&self) -> u8 {
        self.cycle_count.0
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[test]
    fn new() {
        let model = Model::new();

        assert_eq!(Bandwidth::ZERO, model.max_bw());
        assert_eq!(Bandwidth::ZERO, model.bw());
        assert_eq!(Bandwidth::INFINITY, model.bw_hi());
        assert_eq!(Bandwidth::INFINITY, model.bw_lo());
    }

    #[test]
    fn update_max_bw() {
        let mut model = Model::new();

        let mut rate_sample = RateSample {
            interval: Duration::from_millis(10),
            delivered_bytes: 100,
            is_app_limited: true,
            ..Default::default()
        };
        let bw = rate_sample.delivery_rate();

        // Sample is app limited, but we don't have any data yet, so we use the sample
        model.update_max_bw(rate_sample);
        assert_eq!(bw, model.max_bw());

        model.advance_max_bw_filter();

        rate_sample.is_app_limited = false;
        rate_sample.delivered_bytes = 50;

        // Sample is not app limited, so the sample is used. It is less than the max
        // though, so the max does not change
        model.update_max_bw(rate_sample);
        assert_eq!(bw, model.max_bw());

        rate_sample.delivered_bytes = 75;
        let bw = rate_sample.delivery_rate();

        model.advance_max_bw_filter();

        // The MaxBwFilter only tracks 2 values, so the previous max of 100 falls off and
        // the latest bw estimate becomes the max
        model.update_max_bw(rate_sample);
        assert_eq!(bw, model.max_bw());

        // Advance the max bw filter twice so existing values should fall off
        model.advance_max_bw_filter();
        model.advance_max_bw_filter();
        rate_sample.is_app_limited = true;
        rate_sample.delivered_bytes = 50;

        // The sample is app limited so the max stays the same
        model.update_max_bw(rate_sample);
        assert_eq!(bw, model.max_bw());
    }

    #[test]
    fn update_lower_bound() {
        let mut model = Model::new();

        let bw = Bandwidth::new(50, Duration::from_millis(10));

        let rate_sample = RateSample {
            interval: Duration::from_millis(10),
            delivered_bytes: 100,
            ..Default::default()
        };

        model.update_max_bw(rate_sample);
        model.update_lower_bound(bw);

        let max_bw = model.max_bw();

        // We didn't have a valid bw_lo value yet, and the given bw is lower than max_bw * BETA,
        // so bw_lo is set to max_bw * BETA
        assert_eq!(max_bw * BETA, model.bw_lo());

        let lower_bw = Bandwidth::new(50, Duration::from_millis(10));
        model.update_upper_bound(lower_bw);

        // The new sample is lower than bw_lo, so don't update bw_lo
        assert_eq!(max_bw * BETA, model.bw_lo());

        let higher_bw = Bandwidth::new(150, Duration::from_millis(10));
        model.update_lower_bound(higher_bw);

        // The new sample is higher than bw_lo, so update bw_lo
        assert_eq!(higher_bw, model.bw_lo());

        // Resetting the lower bound sets bw_lo to Bandwidth::MAX
        model.reset_lower_bound();
        assert_eq!(Bandwidth::INFINITY, model.bw_lo());
    }

    #[test]
    fn bound_bw_for_model() {
        let mut model = Model::new();

        let mut rate_sample = RateSample {
            interval: Duration::from_millis(10),
            delivered_bytes: 100,
            ..Default::default()
        };

        let high_bw = rate_sample.delivery_rate();
        let med_bw = Bandwidth::new(75, Duration::from_millis(10));
        let low_bw = Bandwidth::new(50, Duration::from_millis(10));

        model.update_max_bw(rate_sample);
        model.update_upper_bound(low_bw);
        model.update_lower_bound(med_bw);

        // bw is not updated yet
        assert_eq!(Bandwidth::ZERO, model.bw());

        model.bound_bw_for_model();

        assert_eq!(low_bw, model.bw());
        assert_eq!(
            model.max_bw().min(model.bw_lo()).min(model.bw_hi()),
            model.bw()
        );

        // bw_hi is no longer the min
        model.update_upper_bound(high_bw);
        model.bound_bw_for_model();

        assert_eq!(med_bw, model.bw());
        assert_eq!(
            model.max_bw().min(model.bw_lo()).min(model.bw_hi()),
            model.bw()
        );

        rate_sample.delivered_bytes = 10;
        let low_bw = rate_sample.delivery_rate();

        // bw_lo is no longer the min
        model.advance_max_bw_filter();
        model.advance_max_bw_filter();
        model.update_max_bw(rate_sample);
        model.bound_bw_for_model();

        assert_eq!(low_bw, model.bw());
        assert_eq!(
            model.max_bw().min(model.bw_lo()).min(model.bw_hi()),
            model.bw()
        );
    }
}


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

use crate::{
    recovery::{
        bandwidth::Bandwidth,
        bbr::{
            ecn::ECN_FACTOR,
            windowed_filter::{MinRttWindowedFilter, WindowedMaxFilter},
            BETA,
        },
    },
    time::Timestamp,
};
use core::time::Duration;

//= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.9.2
//# The data volume model parameters together estimate both the volume of in-flight data required to
//# reach the full bandwidth available to the flow (BBR.max_inflight), and the maximum volume of
//# data that is consistent with the queue pressure objective (cwnd).

#[derive(Clone, Debug)]
pub(crate) struct Model {
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.9.2
    //# The windowed minimum round-trip time sample measured over the last MinRTTFilterLen = 10 seconds.
    min_rtt_filter: MinRttWindowedFilter,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.9.2
    //# A volume of data that is the estimate of the recent degree of aggregation in the network path.
    extra_acked_filter: WindowedMaxFilter<u64, u64, u64>,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.12
    //# the start of the time interval for estimating the excess amount of data acknowledged due to aggregation effects.
    extra_acked_interval_start: Option<Timestamp>,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.12
    //# the volume of data marked as delivered since BBR.extra_acked_interval_start.
    extra_acked_delivered: u64,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.9.2
    //# Analogous to BBR.bw_hi, the long-term maximum volume of in-flight data that the algorithm
    //# estimates will produce acceptable queue pressure, based on signals in the current or
    //# previous bandwidth probing cycle, as measured by loss.
    inflight_hi: u64,
    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.9.2
    //# Analogous to BBR.bw_lo, the short-term maximum volume of in-flight data that the algorithm
    //# estimates is safe for matching the current network path delivery process, based on any loss
    //# signals in the current bandwidth probing cycle.
    inflight_lo: u64,
}

impl Model {
    /// Constructs a new `data_volume::Model`
    pub fn new() -> Self {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#2.12
        //# The window length of the BBR.ExtraACKedFilter max filter window:
        //# 10 (in units of packet-timed round trips).
        const EXTRA_ACKED_FILTER_LEN: u64 = 10;

        Self {
            min_rtt_filter: MinRttWindowedFilter::new(),
            extra_acked_filter: WindowedMaxFilter::new(EXTRA_ACKED_FILTER_LEN),
            extra_acked_interval_start: None,
            extra_acked_delivered: 0,
            inflight_hi: u64::MAX,
            inflight_lo: u64::MAX,
        }
    }

    /// The windowed maximum recent estimate in bytes of the degree of aggregation in the path
    pub fn extra_acked(&self) -> u64 {
        self.extra_acked_filter.value().unwrap_or(0)
    }

    /// The windowed minimum round trip time
    pub fn min_rtt(&self) -> Option<Duration> {
        self.min_rtt_filter.min_rtt()
    }

    /// The long-term maximum volume of in-flight data that the algorithm
    /// estimates will produce acceptable queue pressure
    pub fn inflight_hi(&self) -> u64 {
        self.inflight_hi
    }

    /// The short-term maximum volume of in-flight data that the algorithm
    /// estimates is safe for matching the current network path delivery process
    pub fn inflight_lo(&self) -> u64 {
        self.inflight_lo
    }

    /// True if the probe RTT has expired and is due for a refresh by entering the ProbeRTT state
    pub fn probe_rtt_expired(&self) -> bool {
        self.min_rtt_filter.probe_rtt_expired()
    }

    /// Overrides the last updated time for Min Probe RTT to ensure ProbeRTT is not entered until
    /// the next PROBE_RTT_INTERVAL.
    ///
    /// Called immediately after ProbeRTT period ends or after resuming from idle
    pub fn schedule_next_probe_rtt(&mut self, now: Timestamp) {
        self.min_rtt_filter.schedule_next_probe_rtt(now)
    }

    /// Update the min_rtt estimate with the given `rtt`
    pub fn update_min_rtt(&mut self, rtt: Duration, now: Timestamp) {
        self.min_rtt_filter.update(rtt, now)
    }

    /// Update the ack aggregation estimate
    pub fn update_ack_aggregation(
        &mut self,
        bw: Bandwidth,
        bytes_acknowledged: usize,
        cwnd: u32,
        round_count: u64,
        now: Timestamp,
    ) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.5
        //# BBRUpdateACKAggregation():
        //#   /* Find excess ACKed beyond expected amount over this interval */
        //#   interval = (Now() - BBR.extra_acked_interval_start)
        //#   expected_delivered = BBR.bw * interval
        //#   /* Reset interval if ACK rate is below expected rate: */
        //#   if (BBR.extra_acked_delivered <= expected_delivered)
        //#       BBR.extra_acked_delivered = 0
        //#       BBR.extra_acked_interval_start = Now()
        //#       expected_delivered = 0
        //#   BBR.extra_acked_delivered += rs.newly_acked
        //#   extra = BBR.extra_acked_delivered - expected_delivered
        //#   extra = min(extra, cwnd)
        //#   BBR.extra_acked =
        //#     update_windowed_max_filter(
        //#       filter=BBR.ExtraACKedFilter,
        //#       value=extra,
        //#       time=BBR.round_count,
        //#       window_length=BBRExtraAckedFilterLen)

        let mut expected_delivered = 0;

        if let Some(extra_acked_interval_start) = self.extra_acked_interval_start {
            // Find excess ACKed beyond expected amount over this interval
            let interval = now - extra_acked_interval_start;
            expected_delivered = bw * interval;
        }

        // Reset interval if ACK rate is below expected rate
        if self.extra_acked_delivered <= expected_delivered
            || self.extra_acked_interval_start.is_none()
        {
            self.extra_acked_delivered = 0;
            self.extra_acked_interval_start = Some(now);
            expected_delivered = 0;
        }

        self.extra_acked_delivered += bytes_acknowledged as u64;
        let extra = (self.extra_acked_delivered - expected_delivered).min(cwnd as u64);
        self.extra_acked_filter.update(extra, round_count);
    }

    /// Updates `inflight_hi` with the given `inflight_hi`
    pub fn update_upper_bound(&mut self, inflight_hi: u64) {
        self.inflight_hi = inflight_hi;
    }

    /// Updates `inflight_lo` if there is loss or ECN in the round
    pub fn update_lower_bound(
        &mut self,
        cwnd: u32,
        inflight_latest: u64,
        loss_in_round: bool,
        ecn_in_round: bool,
        ecn_alpha: f64,
    ) {
        if !loss_in_round && !ecn_in_round {
            return;
        }

        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
        //# if (BBR.inflight_lo == Infinity)
        //#     BBR.inflight_lo = cwnd
        if self.inflight_lo == u64::MAX {
            self.inflight_lo = cwnd as u64;
        }

        // Update inflight_lo to the lower of the values determined when loss_in_round or ecn_in_round
        // Based on https://github.com/google/bbr/blob/1a45fd4faf30229a3d3116de7bfe9d2f933d3562/net/ipv4/tcp_bbr2.c#L1618
        let ecn_inflight_lo = if ecn_in_round {
            ((1.0 - (ecn_alpha * ECN_FACTOR)) * self.inflight_lo as f64) as u64
        } else {
            u64::MAX
        };

        let loss_inflight_lo = if loss_in_round {
            //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
            //# BBR.inflight_lo = max(BBR.inflight_latest,
            //#                       BBRBeta * BBR.infligh_lo)
            inflight_latest.max((BETA * self.inflight_lo).to_integer())
        } else {
            u64::MAX
        };

        self.inflight_lo = loss_inflight_lo.min(ecn_inflight_lo);
    }

    /// Resets `inflight_lo` to its initial value
    pub fn reset_lower_bound(&mut self) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.5.6.3
        //# BBR.inflight_lo = Infinity
        self.inflight_lo = u64::MAX
    }

    /// Sets the `extra_acked_interval_start` to the given `timestamp`
    pub fn set_extra_acked_interval_start(&mut self, timestamp: Timestamp) {
        self.extra_acked_interval_start = Some(timestamp);
    }

    #[cfg(test)]
    pub fn set_extra_acked_for_test(&mut self, sample: u64, round_count: u64) {
        self.extra_acked_filter.update(sample, round_count);
    }

    #[cfg(test)]
    pub fn set_inflight_lo_for_test(&mut self, inflight_lo: u64) {
        self.inflight_lo = inflight_lo;
    }

    #[cfg(test)]
    pub fn extra_acked_interval_start(&self) -> Option<Timestamp> {
        self.extra_acked_interval_start
    }

    #[cfg(test)]
    pub fn next_probe_rtt(&self) -> Option<Timestamp> {
        self.min_rtt_filter.next_probe_rtt()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::time::{Clock, NoopClock};

    #[test]
    fn new() {
        let model = Model::new();

        assert_eq!(0, model.extra_acked());
        assert_eq!(None, model.min_rtt());
        assert_eq!(u64::MAX, model.inflight_hi());
        assert_eq!(u64::MAX, model.inflight_lo());
    }

    #[test]
    fn update_ack_aggregation() {
        let now = NoopClock.get_time();
        let mut model = Model::new();

        let now = now + Duration::from_millis(200);
        let bw = Bandwidth::new(1500, Duration::from_secs(1));

        // The first call to update_ack_aggregation starts a new ack aggregation epoch
        model.update_ack_aggregation(bw, 1600, 12000, 0, now);

        assert_eq!(1600, model.extra_acked());
        assert_eq!(Some(now), model.extra_acked_interval_start);
        assert_eq!(1600, model.extra_acked_delivered);

        let now = now + Duration::from_secs(1);

        model.update_ack_aggregation(bw, 1600, 12000, 1, now);

        // The BW sample indicates 1500 bytes should be ACKed over the interval, but instead 1600 were,
        // meaning the extra acked amount is 100 bytes. This is added to the initial 1600 extra acked
        // amount to arrive at 1700 bytes.
        assert_eq!(1700, model.extra_acked());

        let now = now + Duration::from_secs(1);

        // Even more extra data is acked, but since the cwnd is lower than the extra amount, that
        // value is used as the extra acked (1600 bytes). 1700 remains the max extra acked.
        model.update_ack_aggregation(bw, 1700, 1600, 2, now);
        assert_eq!(1700, model.extra_acked());
    }

    #[test]
    fn update_lower_bound() {
        let mut model = Model::new();

        model.update_lower_bound(1000, 100, true, false, 1.0);

        // We didn't have a valid inflight_lo value yet, and the given inflight_latest is lower than cwnd * BETA,
        // so inflight_lo is set to cwnd * BETA
        assert_eq!((BETA * 1000).to_integer(), model.inflight_lo());

        let inflight_latest = 1500;
        model.update_lower_bound(1000, inflight_latest, true, false, 1.0);

        // The new sample is higher than inflight_lo, so update inflight_lo
        assert_eq!(inflight_latest, model.inflight_lo());

        let ecn_alpha = 4.0 / 5.0;
        model.update_lower_bound(1000, inflight_latest, false, true, ecn_alpha);

        // There was ecn_in_round, so lower inflight_lo according to ecn_alpha
        // (1 - 4/5 * .33) * 1500 = 1104
        assert_eq!(1104, model.inflight_lo());

        // Resetting the lower bound sets inflight_lo to u64::MAX
        model.reset_lower_bound();
        assert_eq!(u64::MAX, model.inflight_lo());

        // There is both ECN and Loss in the round, but the Loss reduced value is lower, so use the Loss value
        model.update_lower_bound(1500, 100, true, true, ecn_alpha);
        assert_eq!(1050, model.inflight_lo());

        model.reset_lower_bound();

        // There is both ECN and Loss in the round, but the Loss reduced value is higher, so use the ECN value
        model.update_lower_bound(1500, 1200, true, true, ecn_alpha);
        assert_eq!(1104, model.inflight_lo());
    }
}


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

use crate::{
    random,
    recovery::{
        bbr::{startup, BbrCongestionController, State},
        congestion_controller::Publisher,
    },
    time::Timestamp,
};
use num_rational::Ratio;
use num_traits::One;

//= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.3.2
//# In Drain, BBR aims to quickly drain any queue created in Startup by switching to a
//# pacing_gain well below 1.0, until any estimated queue has been drained. It uses a
//# pacing_gain that is the inverse of the value used during Startup, chosen to try to
//# drain the queue in one round

// The wording above is somewhat ambiguous over whether the drain pacing_gain should be
// the inverse of the startup pacing_gain or startup cwnd_gain. However, the citation below
// makes it clear it is the inverse of the startup cwnd_gain. This is also supported
// by the following derivation:
// https://github.com/google/bbr/blob/master/Documentation/startup/gain/analysis/bbr_drain_gain.pdf

//= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.3.2
//#     BBR.pacing_gain = 1/BBRStartupCwndGain  /* pace slowly */
pub(crate) const PACING_GAIN: Ratio<u64> = Ratio::new_raw(1, 2);

//= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.3.2
//# BBREnterDrain():
//#     BBR.state = Drain
//#     BBR.pacing_gain = 1/BBRStartupCwndGain  /* pace slowly */
//#     BBR.cwnd_gain = BBRStartupCwndGain      /* maintain cwnd */
pub(crate) const CWND_GAIN: Ratio<u64> = startup::CWND_GAIN;

/// Methods related to the Drain state
impl BbrCongestionController {
    /// Enter the `Drain` state
    #[inline]
    pub(super) fn enter_drain<Pub: Publisher>(&mut self, publisher: &mut Pub) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.3.2
        //# BBREnterDrain():
        //#   BBR.state = Drain
        //#   BBR.pacing_gain = 1/BBRStartupCwndGain  /* pace slowly */
        //#   BBR.cwnd_gain = BBRStartupCwndGain      /* maintain cwnd */
        // pacing_gain and cwnd_gain are managed with the State enum

        // New BBR state requires updating the model
        self.try_fast_path = false;
        self.state.transition_to(State::Drain, publisher);
    }

    /// Checks if the `Drain` state is done and enters `ProbeBw` if so
    #[inline]
    pub(super) fn check_drain_done<Pub: Publisher>(
        &mut self,
        random_generator: &mut dyn random::Generator,
        now: Timestamp,
        publisher: &mut Pub,
    ) {
        //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.3.2
        //# BBRCheckDrain():
        //#   if (BBR.state == Drain and packets_in_flight <= BBRInflight(1.0))
        //#     BBREnterProbeBW()  /* BBR estimates the queue was drained */
        if self.state.is_drain()
            && self.bytes_in_flight <= self.inflight(self.data_rate_model.bw(), Ratio::one())
        {
            self.enter_probe_bw(false, random_generator, now, publisher);
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        counter::Counter,
        event, path,
        path::MINIMUM_MTU,
        random,
        recovery::{
            bandwidth::RateSample, bbr::BbrCongestionController,
            congestion_controller::PathPublisher,
        },
        time::{Clock, NoopClock},
    };
    use core::time::Duration;

    #[test]
    fn enter_drain() {
        let mut bbr = BbrCongestionController::new(MINIMUM_MTU);
        let mut publisher = event::testing::Publisher::snapshot();
        let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());

        bbr.enter_drain(&mut publisher);

        assert!(bbr.state.is_drain());
        assert!(!bbr.try_fast_path);
    }

    //= https://tools.ietf.org/id/draft-cardwell-iccrg-bbr-congestion-control-02#4.3.2
    //= type=test
    //# BBRCheckDrain():
    //#   if (BBR.state == Drain and packets_in_flight <= BBRInflight(1.0))
    //#     BBREnterProbeBW()  /* BBR estimates the queue was drained */
    #[test]
    fn check_drain_done() {
        let mut bbr = BbrCongestionController::new(MINIMUM_MTU);
        let now = NoopClock.get_time();
        let mut rng = random::testing::Generator::default();
        let mut publisher = event::testing::Publisher::snapshot();
        let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());

        // Not in drain yet
        bbr.check_drain_done(&mut rng, now, &mut publisher);
        assert!(bbr.state.is_startup());

        bbr.enter_drain(&mut publisher);
        bbr.bytes_in_flight = Counter::new(100);

        // bytes_in_flight > inflight
        bbr.check_drain_done(&mut rng, now, &mut publisher);
        assert!(!bbr.state.is_drain());

        let rate_sample = RateSample {
            delivered_bytes: 100_000,
            interval: Duration::from_millis(1),
            ..Default::default()
        };
        bbr.data_rate_model.update_max_bw(rate_sample);
        bbr.data_rate_model.bound_bw_for_model();

        // Now drain is done
        bbr.check_drain_done(&mut rng, now, &mut publisher);
        assert!(bbr.state.is_probing_bw());
    }
}