use alloc::boxed::Box;
use core::{marker::PhantomData, ptr::null_mut};
use crossbeam_utils::CachePadded;
use crate::{
BoundedCollection,
sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering},
utils::Backoff,
};
#[derive(Debug)]
pub struct Resizable<Q> {
cores: [AtomicPtr<Q>; 2],
push_epoch: CachePadded<AtomicUsize>,
pop_epoch: CachePadded<AtomicUsize>,
active_pushes: CachePadded<[AtomicUsize; 2]>,
active_reads: CachePadded<[AtomicUsize; 2]>,
is_resizing: AtomicBool,
_marker: PhantomData<Box<Q>>,
}
impl<Q> Resizable<Q> {
pub fn from_parts(left: Q, right: Q) -> Self {
Self {
cores: [
AtomicPtr::new(Box::into_raw(Box::new(left))),
AtomicPtr::new(Box::into_raw(Box::new(right))),
],
active_pushes: [AtomicUsize::new(0), AtomicUsize::new(0)].into(),
active_reads: [AtomicUsize::new(0), AtomicUsize::new(0)].into(),
push_epoch: AtomicUsize::new(0).into(),
pop_epoch: AtomicUsize::new(0).into(),
is_resizing: AtomicBool::new(false),
_marker: PhantomData,
}
}
}
impl<Q> Resizable<Q>
where
Q: BoundedCollection,
{
#[track_caller]
pub fn with_capacity(size: usize) -> Resizable<Q> {
Self::from_parts(Q::with_capacity(size), Q::with_capacity(1))
}
#[track_caller]
pub fn new() -> Self {
Self::from_parts(Q::with_capacity(1), Q::with_capacity(1))
}
}
impl<Q> Default for Resizable<Q>
where
Q: Default,
{
fn default() -> Self {
Self::from_parts(Q::default(), Q::default())
}
}
impl<Q> Drop for Resizable<Q> {
fn drop(&mut self) {
let left = self.cores[0].swap(null_mut(), Ordering::Acquire);
if !left.is_null() {
_ = unsafe { Box::from_raw(left) };
}
let right = self.cores[1].swap(null_mut(), Ordering::Acquire);
if !right.is_null() {
_ = unsafe { Box::from_raw(right) };
}
}
}
impl<Q> Resizable<Q>
where
Q: BoundedCollection,
{
pub fn resize(&self, size: usize) -> bool {
if size == 0 {
return false;
}
#[cfg(loom)]
crate::sync::atomic::fence(Ordering::SeqCst);
let push_epoch = self.push_epoch.load(Ordering::SeqCst);
let pop_epoch = self.pop_epoch.load(Ordering::SeqCst);
if pop_epoch != push_epoch {
return false;
}
let old_idx = (push_epoch + 1) % 2;
#[cfg(loom)]
crate::sync::atomic::fence(Ordering::SeqCst);
if self.active_reads[old_idx].load(Ordering::SeqCst) != 0
|| self.active_pushes[old_idx].load(Ordering::SeqCst) != 0
{
return false;
}
if self.is_resizing.swap(true, Ordering::AcqRel) {
return false;
}
#[cfg(loom)]
crate::sync::atomic::fence(Ordering::SeqCst);
if self.push_epoch.load(Ordering::SeqCst) != push_epoch {
self.is_resizing.store(false, Ordering::Release);
return false;
}
let new_queue = Box::into_raw(Box::new(Q::with_capacity(size)));
let old_queue = self.cores[old_idx].swap(new_queue, Ordering::AcqRel);
self.push_epoch.fetch_add(1, Ordering::Release);
let q = unsafe { Box::from_raw(old_queue) };
#[cfg(not(any(loom, shuttle)))]
debug_assert!(q.try_pop().is_none());
#[cfg(any(loom, shuttle))]
assert!(q.try_pop().is_none());
self.is_resizing.store(false, Ordering::Release);
true
}
}
impl<Q> Resizable<Q> {
fn get_queue(&self, epoch: usize) -> &Q {
let queue = self.cores[epoch % 2].load(Ordering::Acquire);
unsafe { &*queue }
}
fn register_push(&self, target_epoch: usize) -> bool {
#[cfg(loom)]
crate::sync::atomic::fence(Ordering::SeqCst);
self.active_reads[target_epoch % 2].fetch_add(1, Ordering::SeqCst);
#[cfg(loom)]
crate::sync::atomic::fence(Ordering::SeqCst);
let current_push = self.push_epoch.load(Ordering::SeqCst);
if target_epoch != current_push {
self.deregister_reader(target_epoch);
return false;
}
true
}
fn register_pop(&self, target_epoch: usize) -> bool {
#[cfg(loom)]
crate::sync::atomic::fence(Ordering::SeqCst);
self.active_reads[target_epoch % 2].fetch_add(1, Ordering::SeqCst);
#[cfg(loom)]
crate::sync::atomic::fence(Ordering::SeqCst);
let current_pop = self.pop_epoch.load(Ordering::SeqCst);
if target_epoch != current_pop {
self.deregister_reader(target_epoch);
return false;
}
true
}
fn deregister_reader(&self, epoch: usize) {
self.active_reads[epoch % 2].fetch_sub(1, Ordering::Release);
}
}
impl<Q> Resizable<Q>
where
Q: BoundedCollection,
{
fn try_pop_from(
&self,
epoch: usize,
registration: impl Fn(&Self, usize) -> bool,
) -> Result<Option<Q::Item>, ()> {
if !registration(self, epoch) {
return Err(());
}
let item = self.get_queue(epoch).try_pop();
self.deregister_reader(epoch);
Ok(item)
}
}
impl<Q> BoundedCollection for Resizable<Q>
where
Q: BoundedCollection,
{
type Item = Q::Item;
fn try_push(&self, item: Self::Item) -> Result<(), Self::Item> {
let mut backoff = Backoff::new();
loop {
let push_epoch = self.push_epoch.load(Ordering::Acquire);
self.active_pushes[push_epoch % 2].fetch_add(1, Ordering::SeqCst);
#[cfg(loom)]
crate::sync::atomic::fence(Ordering::SeqCst);
if self.push_epoch.load(Ordering::SeqCst) == push_epoch {
let r = self.get_queue(push_epoch).try_push(item);
self.active_pushes[push_epoch % 2].fetch_sub(1, Ordering::Release);
return r;
}
self.active_pushes[push_epoch % 2].fetch_sub(1, Ordering::Release);
backoff.backoff();
}
}
fn try_pop(&self) -> Option<Self::Item> {
#[cfg(any(shuttle, loom))]
let mut backoff = Backoff::new();
for _ in 0..2 {
loop {
let push_epoch = self.push_epoch.load(Ordering::Acquire);
let pop_epoch = self.pop_epoch.load(Ordering::Acquire);
if pop_epoch != push_epoch {
let Ok(item) = self.try_pop_from(pop_epoch, Self::register_pop) else {
#[cfg(any(shuttle, loom))]
backoff.backoff();
continue;
};
if item.is_some() {
return item;
}
if self.active_pushes[pop_epoch % 2].load(Ordering::Acquire) == 0 {
let Ok(item) = self.try_pop_from(pop_epoch, Self::register_pop) else {
#[cfg(any(shuttle, loom))]
backoff.backoff();
continue;
};
if item.is_some() {
return item;
}
#[cfg(loom)]
crate::sync::atomic::fence(Ordering::SeqCst);
_ = self.pop_epoch.compare_exchange_weak(
pop_epoch,
pop_epoch + 1,
Ordering::SeqCst,
Ordering::Relaxed,
);
#[cfg(any(shuttle, loom))]
backoff.backoff();
continue;
}
}
let Ok(item) = self.try_pop_from(push_epoch, Self::register_push) else {
#[cfg(any(shuttle, loom))]
backoff.backoff();
continue;
};
if item.is_none() && push_epoch != self.push_epoch.load(Ordering::Acquire) {
#[cfg(any(shuttle, loom))]
backoff.backoff();
continue;
}
if item.is_some() || push_epoch == pop_epoch {
return item;
}
#[cfg(any(loom, shuttle))]
backoff.backoff();
break;
}
}
None
}
fn capacity(&self) -> usize {
loop {
let push_epoch = self.push_epoch.load(Ordering::Acquire);
if !self.register_push(push_epoch) {
continue;
}
let cap = self.get_queue(push_epoch).capacity();
self.deregister_reader(push_epoch);
return cap;
}
}
fn len(&self) -> usize {
loop {
let push_epoch = self.push_epoch.load(Ordering::Acquire);
if !self.register_push(push_epoch) {
continue;
}
let pop_epoch = self.pop_epoch.load(Ordering::Acquire);
let pop_len = if pop_epoch != push_epoch {
if !self.register_pop(pop_epoch) {
self.deregister_reader(push_epoch);
continue;
}
let pop_len = self.get_queue(pop_epoch).len();
self.deregister_reader(pop_epoch);
pop_len
} else {
0
};
let len = self.get_queue(push_epoch).len() + pop_len;
self.deregister_reader(push_epoch);
return len;
}
}
fn is_empty(&self) -> bool {
loop {
let push_epoch = self.push_epoch.load(Ordering::Acquire);
if !self.register_push(push_epoch) {
continue;
}
let pop_epoch = self.pop_epoch.load(Ordering::Acquire);
let pop_is_empty = if pop_epoch != push_epoch {
if !self.register_pop(pop_epoch) {
self.deregister_reader(push_epoch);
continue;
}
let pop_is_empty = self.get_queue(pop_epoch).is_empty();
self.deregister_reader(pop_epoch);
pop_is_empty
} else {
true
};
let is_empty = self.get_queue(push_epoch).is_empty() && pop_is_empty;
self.deregister_reader(push_epoch);
return is_empty;
}
}
fn is_full(&self) -> bool {
loop {
let push_epoch = self.push_epoch.load(Ordering::Acquire);
if !self.register_push(push_epoch) {
continue;
}
let is_full = self.get_queue(push_epoch).is_full();
self.deregister_reader(push_epoch);
return is_full;
}
}
fn with_capacity(capacity: usize) -> Self {
Resizable::with_capacity(capacity)
}
}
#[cfg(not(any(loom, shuttle, echeneis)))]
impl<Q> Resizable<Q> {
pub fn into_parts(self) -> [Box<Q>; 2] {
let left = self.cores[0].swap(null_mut(), Ordering::Acquire);
let left = unsafe { Box::from_raw(left) };
let right = self.cores[1].swap(null_mut(), Ordering::Acquire);
let right = unsafe { Box::from_raw(right) };
[left, right]
}
pub fn parts_mut(&mut self) -> [&mut Q; 2] {
let left = unsafe { &mut **self.cores[0].get_mut() };
let right = unsafe { &mut **self.cores[1].get_mut() };
[left, right]
}
pub fn into_current(mut self) -> Box<Q> {
let push_epoch = *self.push_epoch.get_mut();
let parts = self.into_parts();
parts.into_iter().nth(push_epoch % 2).unwrap()
}
pub fn current_mut(&mut self) -> &mut Q {
unsafe { &mut **self.cores[*self.push_epoch.get_mut() % 2].get_mut() }
}
}
#[cfg(not(any(loom, shuttle, echeneis)))]
impl<Q> Resizable<Q>
where
Q: BoundedCollection,
{
pub fn pop_mut(&mut self) -> Option<Q::Item> {
let pop_idx = *self.pop_epoch.get_mut() % 2;
let push_idx = *self.push_epoch.get_mut() % 2;
if pop_idx == push_idx {
self.current_mut().try_pop()
} else {
let parts = self.parts_mut();
parts[pop_idx].try_pop().or(parts[push_idx].try_pop())
}
}
pub fn push_mut(&mut self, item: Q::Item) -> Result<(), Q::Item> {
self.current_mut().try_push(item)
}
pub fn clear(&mut self) {
for core in self.parts_mut() {
while core.try_pop().is_some() {}
}
}
pub fn migrate(&mut self) {
let pop_epoch = *self.pop_epoch.get_mut();
let push_epoch = *self.push_epoch.get_mut();
if pop_epoch == push_epoch {
return;
}
let parts = self.parts_mut();
let pop_queue = &parts[pop_epoch % 2];
let push_queue = &parts[push_epoch % 2];
let total_items = pop_queue.len() + push_queue.len();
let empty_collection = Q::with_capacity(total_items);
while let Some(item) = pop_queue.try_pop() {
_ = empty_collection.try_push(item);
}
while let Some(item) = push_queue.try_pop() {
_ = empty_collection.try_push(item);
}
let new_queue = Box::into_raw(Box::new(empty_collection));
let old_queue = self.cores[pop_epoch % 2].swap(new_queue, Ordering::Relaxed);
self.push_epoch.fetch_add(1, Ordering::Relaxed);
self.pop_epoch.store(push_epoch + 1, Ordering::Relaxed);
let _q = unsafe { Box::from_raw(old_queue) };
}
pub fn drain(&mut self) -> Drain<'_, Q, impl Iterator<Item = &mut Q>> {
let cores = self.parts_mut();
Drain {
cores: cores.into_iter(),
next_up: None,
}
}
}
#[cfg(not(any(loom, shuttle, echeneis)))]
impl<Q: BoundedCollection> Extend<Q::Item> for Resizable<Q> {
fn extend<I: IntoIterator<Item = Q::Item>>(&mut self, iter: I) {
for item in iter {
if let Err(item) = self.push_mut(item) {
self.migrate();
let cap = self.capacity();
debug_assert!(
self.resize(cap * 2),
"resize is promised to succeed after a call to migrate"
);
_ = self.push_mut(item);
}
}
}
}
pub struct Drain<'a, Q: BoundedCollection, I: Iterator<Item = &'a mut Q>> {
cores: I,
next_up: Option<&'a mut Q>,
}
impl<'a, Q: BoundedCollection, I: Iterator<Item = &'a mut Q>> Iterator for Drain<'a, Q, I> {
type Item = Q::Item;
fn next(&mut self) -> Option<Self::Item> {
if let Some(current) = &mut self.next_up {
if let Some(item) = current.try_pop() {
return Some(item);
}
}
for next in self.cores.by_ref() {
if let Some(item) = next.try_pop() {
self.next_up.replace(next);
return Some(item);
}
}
None
}
}
impl<'a, Q: BoundedCollection, I: Iterator<Item = &'a mut Q>> Drop for Drain<'a, Q, I> {
fn drop(&mut self) {
for _ in self {}
}
}
pub struct IntoIter<Q: BoundedCollection> {
old: Option<Box<Q>>,
new: Option<Box<Q>>,
}
impl<Q: BoundedCollection> Iterator for IntoIter<Q> {
type Item = Q::Item;
fn next(&mut self) -> Option<Self::Item> {
if let Some(old_q) = &mut self.old {
if let Some(item) = old_q.try_pop() {
return Some(item);
}
self.old = None;
}
self.new.as_mut().and_then(|q| q.try_pop())
}
}
#[cfg(not(any(loom, shuttle, echeneis)))]
impl<Q: BoundedCollection> IntoIterator for Resizable<Q> {
type IntoIter = IntoIter<Q>;
type Item = Q::Item;
fn into_iter(mut self) -> Self::IntoIter {
let push_epoch = *self.push_epoch.get_mut();
let [left, right] = self.into_parts();
#[allow(clippy::manual_is_multiple_of)]
let (old, new) = if push_epoch % 2 == 0 {
(right, left)
} else {
(left, right)
};
IntoIter {
old: Some(old),
new: Some(new),
}
}
}