#![no_std]
#[cfg(feature = "std")]
extern crate std;
use core::{hint::spin_loop, sync::atomic::Ordering};
pub trait BackoffStrategy: Default + Send + Sync + 'static {
const BACKOFF: bool = true;
fn backoff(&mut self) -> RetryStrategy;
#[inline]
fn will_reload(&self) -> bool {
false
}
#[inline]
fn backoff_reload<T: PartialEq, F: FnMut() -> T>(
&mut self,
mut current: T,
mut reload: F,
) -> T {
loop {
match self.backoff() {
RetryStrategy::NoReload => return current,
RetryStrategy::Reload => return reload(),
RetryStrategy::ReloadUntilUnchanged => {
let reloaded = reload();
if reloaded == current {
return current;
}
current = reloaded;
}
}
}
}
#[inline]
fn backoff_until<C: BackoffUntilCondition, F: FnMut() -> C>(&mut self, mut f: F) -> C::Result {
loop {
if let Some(res) = f().into_result() {
return res;
}
self.backoff();
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RetryStrategy {
NoReload,
Reload,
ReloadUntilUnchanged,
}
pub trait BackoffUntilCondition {
type Result;
fn into_result(self) -> Option<Self::Result>;
}
impl BackoffUntilCondition for bool {
type Result = ();
#[inline]
fn into_result(self) -> Option<Self::Result> {
if self { Some(()) } else { None }
}
}
impl<T> BackoffUntilCondition for Option<T> {
type Result = T;
#[inline]
fn into_result(self) -> Option<Self::Result> {
self
}
}
#[derive(Debug, Default)]
pub struct NoBackoff;
impl BackoffStrategy for NoBackoff {
const BACKOFF: bool = false;
#[inline]
fn backoff(&mut self) -> RetryStrategy {
RetryStrategy::NoReload
}
}
impl BoundedBackoffStrategy for NoBackoff {
#[inline]
fn is_completed(&self) -> bool {
true
}
}
pub trait BoundedBackoffStrategy: BackoffStrategy {
fn is_completed(&self) -> bool;
#[inline]
fn try_backoff_until<C: BackoffUntilCondition, F: FnMut() -> C>(
&mut self,
mut f: F,
) -> Option<C::Result> {
loop {
if let Some(res) = f().into_result() {
return Some(res);
}
if self.is_completed() {
return None;
}
self.backoff();
}
}
}
#[derive(Debug, Default)]
pub struct BackoffLimit<S, const LIMIT: usize> {
strategy: S,
iter: usize,
}
impl<S: BackoffStrategy, const LIMIT: usize> BackoffLimit<S, LIMIT> {
pub fn new(strategy: S) -> Self {
Self { strategy, iter: 0 }
}
}
impl<S: BackoffStrategy, const LIMIT: usize> BackoffStrategy for BackoffLimit<S, LIMIT> {
const BACKOFF: bool = S::BACKOFF;
#[inline]
fn backoff(&mut self) -> RetryStrategy {
let retry = self.strategy.backoff();
self.iter = self.iter.saturating_add(1);
retry
}
#[inline]
fn will_reload(&self) -> bool {
self.strategy.will_reload()
}
}
impl<S: BackoffStrategy, const LIMIT: usize> BoundedBackoffStrategy for BackoffLimit<S, LIMIT> {
#[inline]
fn is_completed(&self) -> bool {
self.iter >= LIMIT
}
}
#[derive(Debug, Default)]
pub struct SpinBackoff;
impl BackoffStrategy for SpinBackoff {
#[inline]
fn backoff(&mut self) -> RetryStrategy {
spin_loop();
RetryStrategy::Reload
}
#[inline]
fn will_reload(&self) -> bool {
true
}
}
#[derive(Debug, Default)]
pub struct ExponentialBackoff<
const SPIN_LIMIT: usize,
const UNTIL_UNCHANGED_LIMIT: usize = 0,
const YIELD_AFTER: usize = { usize::MAX },
> {
iter: usize,
}
impl<const SPIN_LIMIT: usize, const UNTIL_UNCHANGED_LIMIT: usize, const YIELD_AFTER: usize>
ExponentialBackoff<SPIN_LIMIT, UNTIL_UNCHANGED_LIMIT, YIELD_AFTER>
{
const ASSERT_SPIN_LIMIT: () = assert!(
SPIN_LIMIT < usize::BITS as usize,
"SPIN_LIMIT must be lower than usize::BITS"
);
pub fn starts_at(iter: usize) -> Self {
ExponentialBackoff { iter }
}
pub fn iter_count(&self) -> usize {
self.iter
}
}
impl<const SPIN_LIMIT: usize, const UNTIL_UNCHANGED_LIMIT: usize, const YIELD_AFTER: usize>
BackoffStrategy for ExponentialBackoff<SPIN_LIMIT, UNTIL_UNCHANGED_LIMIT, YIELD_AFTER>
{
#[inline]
fn backoff(&mut self) -> RetryStrategy {
let () = Self::ASSERT_SPIN_LIMIT;
if cfg!(feature = "std") && self.iter >= YIELD_AFTER {
#[cfg(feature = "std")]
std::thread::yield_now();
} else {
for _ in 0..1usize << self.iter.min(SPIN_LIMIT) {
spin_loop();
}
}
self.iter = self.iter.saturating_add(1);
if self.iter <= UNTIL_UNCHANGED_LIMIT {
RetryStrategy::ReloadUntilUnchanged
} else {
RetryStrategy::Reload
}
}
#[inline]
fn will_reload(&self) -> bool {
true
}
}
#[derive(Debug, Default)]
pub struct BackoffState<S> {
strategy: S,
enabled: bool,
}
impl<S: BackoffStrategy> BackoffState<S> {
pub fn new(strategy: S) -> Self {
Self {
strategy,
enabled: false,
}
}
pub fn enable(mut self) -> Self {
self.enabled = true;
self
}
#[inline]
pub fn backoff_reload<T: PartialEq, F: FnOnce() -> T>(
&mut self,
current: &mut T,
reload: F,
) -> bool {
if !self.enabled {
self.enabled = true;
return false;
}
let retry = self.strategy.backoff();
if retry == RetryStrategy::NoReload {
return false;
}
let reloaded = reload();
if reloaded == *current {
return false;
}
*current = reloaded;
self.enabled = retry == RetryStrategy::ReloadUntilUnchanged;
true
}
}
impl<S: BackoffStrategy> From<S> for BackoffState<S> {
fn from(strategy: S) -> Self {
Self::new(strategy)
}
}
pub trait Atomic {
type Value: Copy + PartialEq;
fn load(&self, ordering: Ordering) -> Self::Value;
fn compare_exchange_weak(
&self,
current: Self::Value,
new: Self::Value,
success: Ordering,
failure: Ordering,
) -> Result<Self::Value, Self::Value>;
}
pub trait AtomicWithBackoffExt: Atomic {
fn try_update_with_backoff<S, F>(
&self,
set_order: Ordering,
fetch_order: Ordering,
mut f: F,
strategy: S,
) -> Result<Self::Value, Self::Value>
where
S: BackoffStrategy,
F: FnMut(Self::Value) -> Option<Self::Value>,
{
let mut backoff = BackoffState::new(strategy);
let mut current = self.load(fetch_order);
loop {
let new = f(current).ok_or(current)?;
if backoff.backoff_reload(&mut current, || self.load(fetch_order)) {
continue;
}
match self.compare_exchange_weak(current, new, set_order, fetch_order) {
Ok(x) => return Ok(x),
Err(cur) => current = cur,
}
}
}
fn update_with_backoff<S, F>(
&self,
set_order: Ordering,
fetch_order: Ordering,
mut f: F,
mut strategy: S,
) -> Self::Value
where
S: BackoffStrategy,
F: FnMut(Self::Value) -> Self::Value,
{
let mut current = self.load(fetch_order);
loop {
let failure_order = if strategy.will_reload() {
Ordering::Relaxed
} else {
fetch_order
};
match self.compare_exchange_weak(current, f(current), set_order, failure_order) {
Ok(x) => return x,
Err(cur) => current = strategy.backoff_reload(cur, || self.load(fetch_order)),
}
}
}
}
impl<T: Atomic> AtomicWithBackoffExt for T {}
macro_rules! impl_atomic {
($($($atomic:ident)::+ $(<$t:ident>)? => $value:ty,)*) => {$(
impl$(<$t>)? Atomic for $($atomic)::+$(<$t>)? {
type Value = $value;
#[inline(always)]
fn load(&self, ordering: Ordering) -> Self::Value {
self.load(ordering)
}
#[inline(always)]
fn compare_exchange_weak(
&self,
current: Self::Value,
new: Self::Value,
success: Ordering,
failure: Ordering,
) -> Result<Self::Value, Self::Value> {
self.compare_exchange_weak(current, new, success, failure)
}
}
)*};
}
macro_rules! impl_core_atomic {
($($size:literal: $atomic:ident $(<$t:ident>)? => $value:ty,)*) => {$(
#[cfg(target_has_atomic = $size)]
impl_atomic!(core::sync::atomic::$atomic $(<$t>)? => $value,);
)*};
}
impl_core_atomic! {
"8": AtomicBool => bool,
"8": AtomicI8 => i8,
"8": AtomicU8 => u8,
"16": AtomicI16 => i16,
"16": AtomicU16 => u16,
"32": AtomicI32 => i32,
"32": AtomicU32 => u32,
"64": AtomicI64 => i64,
"64": AtomicU64 => u64,
"ptr": AtomicIsize => isize,
"ptr": AtomicUsize => usize,
"ptr": AtomicPtr<T> => *mut T,
}
#[cfg(feature = "portable-atomic")]
macro_rules! impl_portable_atomic {
($($cfg:ident: $atomic:ident $(<$t:ident>)? => $value:ty,)*) => {
portable_atomic::cfg_has_atomic_cas! {$(
portable_atomic::$cfg! {
impl_atomic!(portable_atomic::$atomic $(<$t>)? => $value,);
}
)*}
};
}
#[cfg(feature = "portable-atomic")]
impl_portable_atomic! {
cfg_has_atomic_8: AtomicBool => bool,
cfg_has_atomic_8: AtomicI8 => i8,
cfg_has_atomic_8: AtomicU8 => u8,
cfg_has_atomic_16: AtomicI16 => i16,
cfg_has_atomic_16: AtomicU16 => u16,
cfg_has_atomic_32: AtomicI32 => i32,
cfg_has_atomic_32: AtomicU32 => u32,
cfg_has_atomic_64: AtomicI64 => i64,
cfg_has_atomic_64: AtomicU64 => u64,
cfg_has_atomic_128: AtomicI128 => i128,
cfg_has_atomic_128: AtomicU128 => u128,
cfg_has_atomic_ptr: AtomicIsize => isize,
cfg_has_atomic_ptr: AtomicUsize => usize,
cfg_has_atomic_ptr: AtomicPtr<T> => *mut T,
}
#[cfg(loom)]
macro_rules! impl_loom_atomic {
($($($width:literal:)? $atomic:ident $(<$t:ident>)? => $value:ty,)*) => {$(
$(#[cfg(target_pointer_width = $width)])?
impl_atomic!(loom::sync::atomic::$atomic $(<$t>)? => $value,);
)*};
}
#[cfg(loom)]
impl_loom_atomic! {
AtomicBool => bool,
AtomicI8 => i8,
AtomicU8 => u8,
AtomicI16 => i16,
AtomicU16 => u16,
AtomicI32 => i32,
AtomicU32 => u32,
"64": AtomicI64 => i64,
"64": AtomicU64 => u64,
AtomicIsize => isize,
AtomicUsize => usize,
AtomicPtr<T> => *mut T,
}
#[cfg(test)]
mod tests {
extern crate std;
use core::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::{sync::Arc, thread};
use crate::{AtomicWithBackoffExt, BackoffLimit, BoundedBackoffStrategy, NoBackoff};
fn increment_twice<R: Send + 'static>(
f: impl Fn(&AtomicUsize) -> R + Copy + Send + 'static,
) -> [R; 2] {
let atomic = Arc::new(AtomicUsize::new(0));
let spawn = || {
let atomic = atomic.clone();
thread::spawn(move || f(&atomic))
};
let (t1, t2) = (spawn(), spawn());
[t1.join().unwrap(), t2.join().unwrap()]
}
#[test]
fn update() {
let results = increment_twice(|atomic| {
atomic.update_with_backoff(Relaxed, Relaxed, |x| x + 1, NoBackoff)
});
assert!(results == [0, 1] || results == [1, 0], "{:?}", results);
}
#[test]
fn try_update() {
let results = increment_twice(|atomic| {
let incr = |x| if x != 1 { Some(x + 1) } else { None };
atomic.try_update_with_backoff(Relaxed, Relaxed, incr, NoBackoff)
});
assert!(
results == [Ok(0), Err(1)] || results == [Err(1), Ok(0)],
"{:?}",
results
);
}
#[test]
fn backoff_limit() {
let mut backoff = BackoffLimit::<NoBackoff, 2>::default();
let mut calls = 0;
let res = backoff.try_backoff_until(|| {
calls += 1;
false
});
assert_eq!(res, None);
assert_eq!(calls, 3);
assert!(backoff.is_completed());
assert_eq!(backoff.try_backoff_until(|| Some(42)), Some(42));
assert!(NoBackoff.is_completed());
}
}