#[deny(unsafe_code)]
#[cfg(feature = "std")]
use std::time::{Duration, Instant};
use core::cell::RefCell;
use core::fmt;
use core::sync::atomic;
#[cfg(feature = "random")]
use core::sync::atomic::{AtomicUsize, Ordering};
#[cfg(feature = "random")]
use rand::{rngs::SmallRng, Rng, SeedableRng};
#[derive(Clone)]
pub struct BackOff {
strategy: RefCell<Strategy>,
}
impl Default for BackOff {
#[inline]
fn default() -> Self {
Self::new()
}
}
impl BackOff {
#[inline]
pub const fn new() -> Self {
Self { strategy: RefCell::new(Strategy::constant()) }
}
#[inline(never)]
pub fn spin_once() {
atomic::spin_loop_hint();
}
#[inline]
pub fn reset(&self) {
self.strategy.borrow_mut().reset();
}
#[inline]
pub fn spin(&self) {
let steps = self.strategy.borrow_mut().exponential_backoff();
for _ in 0..steps {
Self::spin_once();
}
}
#[inline]
pub fn advise_yield(&self) -> bool {
self.strategy.borrow().advise_yield()
}
}
#[cfg(feature = "random")]
impl BackOff {
pub fn random() -> Self {
Self { strategy: RefCell::new(Strategy::random()) }
}
pub fn random_with_seed(seed: u64) -> Self {
Self { strategy: RefCell::new(Strategy::random_with_seed(seed)) }
}
}
#[cfg(feature = "std")]
impl BackOff {
pub fn spin_for(dur: Duration) {
let now = Instant::now();
let end = now + dur;
while Instant::now() < end {
Self::spin_once();
}
}
#[inline]
pub fn yield_now() {
std::thread::yield_now();
}
}
impl fmt::Debug for BackOff {
#[inline]
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("BackOff").field("advise_yield", &self.advise_yield()).finish()
}
}
impl fmt::Display for BackOff {
#[inline]
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "advise yield: {}", self.advise_yield())
}
}
#[derive(Clone, Debug)]
enum Strategy {
Const {
pow: u32,
},
#[cfg(feature = "random")]
Random {
pow: u32,
rng: SmallRng,
},
}
impl Strategy {
const INIT_POW: u32 = 1;
const SPIN_LIMIT_POW: u32 = 7;
#[inline]
const fn constant() -> Self {
Strategy::Const { pow: Self::INIT_POW }
}
#[inline]
fn exponential_backoff(&mut self) -> u32 {
match self {
Strategy::Const { pow } => {
let steps = 1 << *pow;
if *pow < Self::SPIN_LIMIT_POW {
*pow += 1;
}
steps
}
#[cfg(feature = "random")]
Strategy::Random { pow, rng } => {
let low = 1 << (*pow - 1);
let high = 1 << *pow;
if *pow < Self::SPIN_LIMIT_POW {
*pow += 1;
}
rng.gen_range(low, high)
}
}
}
#[inline]
fn reset(&mut self) {
let pow = match self {
Strategy::Const { pow } => pow,
#[cfg(feature = "random")]
Strategy::Random { pow, .. } => pow,
};
*pow = Self::INIT_POW;
}
#[inline]
fn advise_yield(&self) -> bool {
let pow = match self {
Strategy::Const { pow } => *pow,
#[cfg(feature = "random")]
Strategy::Random { pow, .. } => *pow,
};
pow == Self::SPIN_LIMIT_POW
}
}
#[cfg(feature = "random")]
impl Strategy {
#[inline]
fn random() -> Self {
#[cfg(target_pointer_width = "32")]
const INIT_SEED: usize = 0x608c_dbfc;
#[cfg(target_pointer_width = "64")]
const INIT_SEED: usize = 0xd1dc_dceb_2fb4_70f3;
const SEED_INCREMENT: usize = 51;
static GLOBAL_SEED: AtomicUsize = AtomicUsize::new(INIT_SEED);
let seed = GLOBAL_SEED.fetch_add(SEED_INCREMENT, Ordering::Relaxed) as u64;
Strategy::Random { pow: Self::INIT_POW, rng: SmallRng::seed_from_u64(seed) }
}
#[inline]
fn random_with_seed(seed: u64) -> Self {
Strategy::Random { pow: Self::INIT_POW, rng: SmallRng::seed_from_u64(seed) }
}
}
#[cfg(test)]
mod tests {
use super::{BackOff, Strategy};
#[test]
fn spin_full_const() {
let backoff = BackOff::new();
let mut steps = 1;
while !backoff.advise_yield() {
backoff.spin();
steps += 1;
}
assert_eq!(steps, Strategy::SPIN_LIMIT_POW);
}
#[cfg(feature = "random")]
#[test]
fn spin_full_random() {
let backoff = BackOff::random();
let mut steps = 1;
while !backoff.advise_yield() {
backoff.spin();
steps += 1;
}
assert_eq!(steps, Strategy::SPIN_LIMIT_POW);
}
}