Skip to content

Instantly share code, notes, and snippets.

@thomcc
Created Apr 2, 2021
Embed
What would you like to do?
//! This is a (hopefully) correct implementation of a lockfree stack.
//!
//! Doing this without a memory reclamation implementation requires pointer
//! tagging, which is slightly architecture specific. This supports aarch64 and
//! x86_64 so long as 5-level page tables aren't in use (rare to the point of
//! nonexistence), and arm64e (or other hardware pointer authentication) is not
//! enabled, which is not currently supported by Rust. On x86_64/aarch64 where
//! this doesn't hold, it will panic.
//!
//! That said, while this implementation is believed to be correct, it could
//! have bugs so general use externally is somewhat discouraged — (it certainly
//! has inefficiencies, for example).
use core::marker::PhantomData;
use core::sync::atomic::{AtomicPtr, AtomicU64, Ordering::*};
type AtomicTaggedPtr = AtomicU64;
#[cfg(target_pointer_width = "64")]
const COUNTER_BITS: usize = 16;
#[cfg(target_pointer_width = "32")]
const COUNTER_BITS: usize = 32;
const PTR_MASK: u64 = (!0u64) >> COUNTER_BITS;
const COUNTER_MASK: u64 = !PTR_MASK;
const COUNTER_INC: u64 = PTR_MASK + 1;
pub struct Stack<T> {
head: AtomicTaggedPtr,
_boo: PhantomData<T>,
}
struct Node<T> {
data: T,
next: AtomicPtr<Node<T>>,
}
impl<T> Drop for Stack<T> {
fn drop(&mut self) {
while self.pop().is_some() {}
}
}
impl<T> Stack<T> {
pub fn new() -> Self {
Self {
head: AtomicTaggedPtr::new(0),
_boo: PhantomData,
}
}
pub fn push(&self, data: T) {
let n = Box::into_raw(Box::new(Node {
next: AtomicPtr::new(core::ptr::null_mut()),
data,
}));
assert_eq!((n as u64) & COUNTER_MASK, 0, "ptr: {:p}", n);
let mut cmp = self.head.load(Relaxed);
loop {
let counter = (cmp & COUNTER_MASK).wrapping_add(COUNTER_INC);
assert_eq!(counter & PTR_MASK, 0);
let swap = (n as u64) | counter;
// Safety: We fully own `n` still, so this is fine
unsafe {
(*n).next = AtomicPtr::new((cmp & PTR_MASK) as usize as *mut Nod
e<T>);
}
match self.head.compare_exchange_weak(cmp, swap, Release, Relaxed) {
Ok(_) => break,
Err(new) => cmp = new,
}
}
}
pub fn pop(&self) -> Option<T> {
let mut cmp = self.head.load(Acquire);
let cur = loop {
let cur = (cmp & PTR_MASK) as usize as *mut Node<T>;
if cur.is_null() {
return None;
}
let next = unsafe { (*cur).next.load(Relaxed) };
let counter = cmp & COUNTER_MASK;
let swap = (next as u64) | counter;
match self.head.compare_exchange_weak(cmp, swap, Acquire, Relaxed) {
Ok(_) => break cur,
Err(h) => cmp = h,
}
};
assert!(!cur.is_null());
unsafe {
let data = core::ptr::read(&(*cur).data);
std::alloc::dealloc(cur.cast::<u8>(), std::alloc::Layout::new::<Node
<T>>());
Some(data)
}
}
}
// send+sync for sendable data.
unsafe impl<T> Send for Stack<T> where T: Send + 'static {}
unsafe impl<T> Sync for Stack<T> where T: Send + 'static {}
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
fn main() {
cobb::run_test(cobb::TestCfg::<Stack<usize>> {
threads: if cfg!(miri) { 8 } else { 16 },
iterations: if cfg!(miri) { 100 } else { 1000 },
sub_iterations: if cfg!(miri) { 10 } else { 20 },
setup: || Stack::new(),
test: |stk, tctx| {
stk.push(tctx.thread_index());
let _ = stk.pop();
},
..Default::default()
});
}
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
fn main() {
panic!("Sorry, this architecture is unsupported");
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment