#![allow(deprecated)]
use core::cell::{Cell, RefCell};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use crate::concurrent::{EpochHandle, EpochRegion, ShardedHandle};
struct ShardedInner<T> {
shards: Box<[EpochRegion<T>]>,
next_shard: AtomicUsize,
}
struct ErasedGuard {
tokens: Arc<[AtomicBool]>,
shard: Option<u16>,
}
impl Drop for ErasedGuard {
fn drop(&mut self) {
if let Some(id) = self.shard {
if let Some(occupied) = self.tokens.get(usize::from(id)) {
occupied.store(false, Ordering::Release);
}
}
}
}
thread_local! {
static MY_SHARD: Cell<Option<u16>> = const { Cell::new(None) };
}
thread_local! {
static ERASED_GUARD: RefCell<Option<ErasedGuard>> = const { RefCell::new(None) };
}
const DEFAULT_CAP_PER_SHARD: usize = 1024;
const MAX_SHARDS: usize = u16::MAX as usize;
#[deprecated(
since = "0.1.0",
note = "concurrent regions are legacy/research-tier; use the production allocator stack (`alloc-xthread`) for cross-thread allocation needs"
)]
pub struct ShardedRegion<T> {
inner: Arc<ShardedInner<T>>,
tokens: Arc<[AtomicBool]>,
}
impl<T> ShardedRegion<T> {
#[must_use]
pub fn with_shards(n: usize, cap_per_shard: usize) -> Self {
assert!(n > 0, "ShardedRegion::with_shards: n must be > 0");
assert!(
n <= MAX_SHARDS,
"ShardedRegion::with_shards: n={n} exceeds the u16 shard-id space ({MAX_SHARDS})"
);
let shards: Vec<EpochRegion<T>> = (0..n)
.map(|_| EpochRegion::with_capacity(cap_per_shard))
.collect();
let tokens: Vec<AtomicBool> = (0..n).map(|_| AtomicBool::new(false)).collect();
Self {
inner: Arc::new(ShardedInner {
shards: shards.into_boxed_slice(),
next_shard: AtomicUsize::new(0),
}),
tokens: Arc::from(tokens.into_boxed_slice()),
}
}
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn shard_count(&self) -> usize {
self.inner.shards.len()
}
#[must_use]
pub fn len(&self) -> usize {
self.inner.shards.iter().map(EpochRegion::len).sum()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.inner.shards.iter().all(EpochRegion::is_empty)
}
fn my_shard(&self) -> Option<u16> {
MY_SHARD.with(|cell| cell.get())
}
fn claim_or_get_shard(&self) -> u16 {
let n = self.inner.shards.len();
if let Some(id) = self.my_shard() {
if usize::from(id) < n {
return id;
}
}
let mut claimed_exclusively: Option<u16> = None;
for (i, occupied) in self.tokens.iter().enumerate() {
if occupied
.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_ok()
{
claimed_exclusively =
Some(u16::try_from(i).expect("shard index fits u16: n <= u16::MAX"));
break;
}
}
let id = claimed_exclusively.unwrap_or_else(|| {
let ticket = self.inner.next_shard.fetch_add(1, Ordering::Relaxed);
u16::try_from(ticket % n)
.expect("shard id fits u16: ticket%n where n<=u16::MAX cannot exceed u16::MAX")
});
MY_SHARD.with(|cell| cell.set(Some(id)));
ERASED_GUARD.with(|slot| {
let mut slot = slot.borrow_mut();
if slot.is_some() {
return;
}
*slot = Some(ErasedGuard {
tokens: Arc::clone(&self.tokens),
shard: claimed_exclusively,
});
});
id
}
pub fn insert(&self, value: T) -> Result<ShardedHandle<T>, T> {
let shard = self.claim_or_get_shard();
match self.inner.shards[usize::from(shard)].insert(value) {
Ok(inner) => Ok(ShardedHandle::new(shard, inner)),
Err(value) => Err(value),
}
}
pub fn get_with<R>(&self, handle: ShardedHandle<T>, f: impl FnOnce(&T) -> R) -> Option<R> {
let shard = self.inner.shards.get(usize::from(handle.shard))?;
shard.get_with(handle.inner, f)
}
pub fn get_cloned(&self, handle: ShardedHandle<T>) -> Option<T>
where
T: Clone,
{
self.get_with(handle, T::clone)
}
pub fn remove(&self, handle: ShardedHandle<T>) -> bool {
let Some(shard) = self.inner.shards.get(usize::from(handle.shard)) else {
return false;
};
let mine = self.my_shard() == Some(handle.shard);
if mine {
shard.remove(handle.inner)
} else {
shard.remote_evict(handle.inner)
}
}
#[must_use]
pub fn bind_current_thread_to_shard(&self, shard: u16) -> bool {
if usize::from(shard) >= self.inner.shards.len() {
return false;
}
let claimed_exclusively = self.tokens[usize::from(shard)]
.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_ok();
MY_SHARD.with(|cell| cell.set(Some(shard)));
ERASED_GUARD.with(|slot| {
let mut slot = slot.borrow_mut();
if slot.is_none() {
*slot = Some(ErasedGuard {
tokens: Arc::clone(&self.tokens),
shard: claimed_exclusively.then_some(shard),
});
}
});
true
}
#[doc(hidden)]
pub fn _reset_my_shard_binding_for_tests() {
MY_SHARD.with(|cell| cell.set(None));
}
}
impl<T> Default for ShardedRegion<T> {
fn default() -> Self {
let n = std::thread::available_parallelism()
.map(|p| p.get())
.unwrap_or(1)
.min(MAX_SHARDS);
Self::with_shards(n.max(1), DEFAULT_CAP_PER_SHARD)
}
}
#[allow(clippy::unused_self)]
impl<T> ShardedRegion<T> {
#[must_use]
pub fn split_handle(handle: ShardedHandle<T>) -> (u16, EpochHandle<T>) {
(handle.shard, handle.inner)
}
}