use std::sync::atomic::{AtomicU8, Ordering};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum BbrState {
Startup = 0,
Drain = 1,
ProbeBW = 2,
ProbeRTT = 3,
}
impl BbrState {
pub(crate) fn from_u8(value: u8) -> Option<Self> {
match value {
0 => Some(Self::Startup),
1 => Some(Self::Drain),
2 => Some(Self::ProbeBW),
3 => Some(Self::ProbeRTT),
_ => None,
}
}
pub(crate) fn is_probing_bandwidth(&self) -> bool {
matches!(self, Self::Startup)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum ProbeBwPhase {
Down = 0,
Cruise = 1,
Refill = 2,
Up = 3,
}
impl ProbeBwPhase {
pub(crate) fn from_u8(value: u8) -> Option<Self> {
match value {
0 => Some(Self::Down),
1 => Some(Self::Cruise),
2 => Some(Self::Refill),
3 => Some(Self::Up),
_ => None,
}
}
pub(crate) fn next(self) -> Self {
match self {
Self::Down => Self::Cruise,
Self::Cruise => Self::Refill,
Self::Refill => Self::Up,
Self::Up => Self::Down,
}
}
pub(crate) fn pacing_gain(self) -> f64 {
match self {
Self::Down => super::config::PROBE_BW_DOWN_PACING_GAIN,
Self::Cruise => super::config::PROBE_BW_CRUISE_PACING_GAIN,
Self::Refill => super::config::PROBE_BW_CRUISE_PACING_GAIN,
Self::Up => super::config::PROBE_BW_UP_PACING_GAIN,
}
}
}
pub(crate) struct AtomicBbrState(AtomicU8);
impl AtomicBbrState {
pub(crate) fn new(state: BbrState) -> Self {
Self(AtomicU8::new(state as u8))
}
pub(crate) fn load(&self) -> BbrState {
let value = self.0.load(Ordering::Acquire);
match BbrState::from_u8(value) {
Some(state) => state,
None => {
tracing::error!(
value,
"CRITICAL: Invalid BBR state value - possible memory corruption"
);
debug_assert!(false, "Invalid BBR state value: {}", value);
BbrState::ProbeBW
}
}
}
pub(crate) fn store(&self, state: BbrState) {
self.0.store(state as u8, Ordering::Release);
}
pub(crate) fn is_startup(&self) -> bool {
self.load() == BbrState::Startup
}
pub(crate) fn is_probe_rtt(&self) -> bool {
self.load() == BbrState::ProbeRTT
}
pub(crate) fn enter_startup(&self) {
self.store(BbrState::Startup);
}
pub(crate) fn enter_drain(&self) {
self.store(BbrState::Drain);
}
pub(crate) fn enter_probe_bw(&self) {
self.store(BbrState::ProbeBW);
}
pub(crate) fn enter_probe_rtt(&self) {
self.store(BbrState::ProbeRTT);
}
pub(crate) fn compare_exchange(
&self,
expected: BbrState,
new: BbrState,
) -> Result<BbrState, BbrState> {
self.0
.compare_exchange(
expected as u8,
new as u8,
Ordering::AcqRel,
Ordering::Acquire,
)
.map(|v| BbrState::from_u8(v).unwrap_or(BbrState::ProbeBW))
.map_err(|v| BbrState::from_u8(v).unwrap_or(BbrState::ProbeBW))
}
}
impl std::fmt::Debug for AtomicBbrState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "AtomicBbrState({:?})", self.load())
}
}
pub(crate) struct AtomicProbeBwPhase(AtomicU8);
impl AtomicProbeBwPhase {
pub(crate) fn new(phase: ProbeBwPhase) -> Self {
Self(AtomicU8::new(phase as u8))
}
pub(crate) fn load(&self) -> ProbeBwPhase {
let value = self.0.load(Ordering::Acquire);
ProbeBwPhase::from_u8(value).unwrap_or(ProbeBwPhase::Cruise)
}
pub(crate) fn store(&self, phase: ProbeBwPhase) {
self.0.store(phase as u8, Ordering::Release);
}
pub(crate) fn advance(&self) -> ProbeBwPhase {
let current = self.load();
let next = current.next();
self.store(next);
next
}
}
impl std::fmt::Debug for AtomicProbeBwPhase {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "AtomicProbeBwPhase({:?})", self.load())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bbr_state_values() {
assert_eq!(BbrState::Startup as u8, 0);
assert_eq!(BbrState::Drain as u8, 1);
assert_eq!(BbrState::ProbeBW as u8, 2);
assert_eq!(BbrState::ProbeRTT as u8, 3);
}
#[test]
fn test_probe_bw_phase_values() {
assert_eq!(ProbeBwPhase::Down as u8, 0);
assert_eq!(ProbeBwPhase::Cruise as u8, 1);
assert_eq!(ProbeBwPhase::Refill as u8, 2);
assert_eq!(ProbeBwPhase::Up as u8, 3);
}
#[test]
fn test_probe_bw_phase_cycle() {
assert_eq!(ProbeBwPhase::Down.next(), ProbeBwPhase::Cruise);
assert_eq!(ProbeBwPhase::Cruise.next(), ProbeBwPhase::Refill);
assert_eq!(ProbeBwPhase::Refill.next(), ProbeBwPhase::Up);
assert_eq!(ProbeBwPhase::Up.next(), ProbeBwPhase::Down);
}
#[test]
fn test_atomic_bbr_state() {
let state = AtomicBbrState::new(BbrState::Startup);
assert!(state.is_startup());
state.enter_drain();
assert_eq!(state.load(), BbrState::Drain);
state.enter_probe_bw();
assert_eq!(state.load(), BbrState::ProbeBW);
state.enter_probe_rtt();
assert!(state.is_probe_rtt());
}
#[test]
fn test_atomic_probe_bw_phase() {
let phase = AtomicProbeBwPhase::new(ProbeBwPhase::Down);
assert_eq!(phase.load(), ProbeBwPhase::Down);
let next = phase.advance();
assert_eq!(next, ProbeBwPhase::Cruise);
assert_eq!(phase.load(), ProbeBwPhase::Cruise);
}
}