use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
#[derive(Debug)]
pub struct BatchIdGen {
next: AtomicU64,
}
impl BatchIdGen {
pub const fn new() -> Self {
Self {
next: AtomicU64::new(1),
}
}
pub fn next(&self) -> u64 {
self.next.fetch_add(1, Ordering::Relaxed)
}
}
impl Default for BatchIdGen {
fn default() -> Self {
Self::new()
}
}
pub(crate) fn take_next(counter: &AtomicU64) -> Option<u64> {
let mut current = counter.load(Ordering::Relaxed);
loop {
let advanced = current.checked_add(1)?;
match counter.compare_exchange_weak(current, advanced, Ordering::Relaxed, Ordering::Relaxed)
{
Ok(previous) => return Some(previous),
Err(observed) => current = observed,
}
}
}
#[derive(Debug)]
pub struct PendingQueue<T> {
items: Mutex<Vec<T>>,
}
impl<T> Default for PendingQueue<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> PendingQueue<T> {
pub const fn new() -> Self {
Self {
items: Mutex::new(Vec::new()),
}
}
pub fn push(&self, item: T) {
if let Ok(mut g) = self.items.lock() {
g.push(item);
}
}
pub fn drain(&self) -> Vec<T> {
match self.items.lock() {
Ok(mut g) => std::mem::take(&mut *g),
Err(_) => Vec::new(),
}
}
pub fn requeue(&self, item: T) {
if let Ok(mut g) = self.items.lock() {
g.insert(0, item);
}
}
pub fn retain<F>(&self, predicate: F)
where
F: FnMut(&T) -> bool,
{
if let Ok(mut g) = self.items.lock() {
g.retain(predicate);
}
}
pub fn with_first_mut<P, F, R>(&self, mut predicate: P, f: F) -> Option<R>
where
P: FnMut(&T) -> bool,
F: FnOnce(&mut T) -> R,
{
let mut g = self.items.lock().ok()?;
let entry = g.iter_mut().find(|t| predicate(t))?;
Some(f(entry))
}
pub fn len(&self) -> usize {
self.items.lock().map(|g| g.len()).unwrap_or(0)
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn batch_id_gen_is_monotonic_and_nonzero() {
let g = BatchIdGen::new();
let a = g.next();
let b = g.next();
let c = g.next();
assert!(a > 0 && b > a && c > b);
}
#[test]
fn take_next_stops_at_the_counter_limit() {
let counter = AtomicU64::new(u64::MAX - 1);
assert_eq!(take_next(&counter), Some(u64::MAX - 1));
assert_eq!(take_next(&counter), None);
assert_eq!(counter.load(Ordering::Relaxed), u64::MAX);
}
#[test]
fn push_drain_round_trip() {
let q = PendingQueue::<u32>::new();
q.push(1);
q.push(2);
let out = q.drain();
assert_eq!(out, vec![1, 2]);
assert!(q.drain().is_empty());
}
#[test]
fn requeue_inserts_at_head() {
let q = PendingQueue::<u32>::new();
q.push(1);
q.requeue(99);
assert_eq!(q.drain(), vec![99, 1]);
}
#[test]
fn retain_removes_matches() {
let q = PendingQueue::<u32>::new();
for i in 0..5 {
q.push(i);
}
q.retain(|x| *x % 2 == 0);
assert_eq!(q.drain(), vec![0, 2, 4]);
}
#[test]
fn with_first_mut_targets_match_only() {
let q = PendingQueue::<(u32, u32)>::new();
q.push((1, 10));
q.push((2, 20));
let touched = q.with_first_mut(|x| x.0 == 2, |x| x.1 += 5);
assert_eq!(touched, Some(()));
assert_eq!(q.drain(), vec![(1, 10), (2, 25)]);
}
}