Skip to content

Instantly share code, notes, and snippets.

@vtta
Created May 17, 2023 02:13
Show Gist options
  • Save vtta/48469e01dc0e3d713caef66506125b47 to your computer and use it in GitHub Desktop.
Save vtta/48469e01dc0e3d713caef66506125b47 to your computer and use it in GitHub Desktop.
A simple lock free list implementation to be use in the linux kernel based on [this](https://www.microsoft.com/en-us/research/wp-content/uploads/2001/10/2001-disc.pdf) microsoft paper
use core::{
fmt, ptr,
sync::atomic::{AtomicPtr, Ordering::SeqCst},
};
use kernel::prelude::*;
#[derive(Debug)]
pub struct Node<T> {
value: Option<T>,
next: AtomicPtr<Node<T>>,
}
impl<T> Default for Node<T> {
fn default() -> Self {
Self {
value: None,
next: AtomicPtr::new(ptr::null_mut()),
}
}
}
impl<T> Node<T> {
pub fn new(v: T) -> Self {
Self {
value: Some(v),
next: AtomicPtr::new(ptr::null_mut()),
}
}
pub fn link(&self, next: *mut Self) {
self.next.store(next, SeqCst)
}
pub fn put(&mut self, v: T) {
self.value.replace(v);
}
pub fn next(&self) -> *mut Self {
// SAFETY: we know LListNode<T>'s alignment is at least 8
(self.next.load(SeqCst) as usize & !0x1usize) as *mut Self
}
pub fn get(&self) -> Option<&T> {
self.value.as_ref()
}
pub fn get_mut(&mut self) -> Option<&mut T> {
self.value.as_mut()
}
pub fn try_swap_next(&self, old: *mut Self, new: *mut Self) -> bool {
self.next.swap(new, SeqCst) == old
}
pub fn mark(p: *mut Self) -> *mut Self {
(p as usize | 0x1) as *mut Self
}
pub fn is_marked(p: *mut Self) -> bool {
(p as usize & 0x1) != 0
}
}
macro_rules! as_ref {
($e:expr) => {
unsafe { $e.as_ref() }.expect("corrupted list node")
};
}
macro_rules! next {
($e:expr) => {
as_ref!($e).next()
};
}
#[derive(Debug)]
pub struct LList<T> {
// both head and tail are sentinels
head: AtomicPtr<Node<T>>,
tail: AtomicPtr<Node<T>>,
}
impl<T: Ord + Copy + fmt::Debug> fmt::Display for LList<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut curr = next!(self.head());
while !self.is_tail(curr) {
let next = next!(curr);
write!(f, "->{:?}", as_ref!(curr).get())?;
curr = if Node::is_marked(next) {
// curr must have at least two following nodes: <marked> <tail>
next!(next)
} else {
next
};
}
Ok(())
}
}
impl<T> Default for LList<T> {
fn default() -> Self {
let tail = Box::try_new(Node::default()).expect("cannot create list tail");
let head = Box::try_new(Node::default()).expect("cannot create list head");
let tail = AtomicPtr::new(Box::leak(tail));
let head = AtomicPtr::new(Box::leak(head));
as_ref!(head.load(SeqCst)).link(tail.load(SeqCst));
Self { head, tail }
}
}
// We require Copy here to avoid memory leak issue
impl<T: Ord + Copy> LList<T> {
fn head(&self) -> *mut Node<T> {
self.head.load(SeqCst)
}
fn tail(&self) -> *mut Node<T> {
self.tail.load(SeqCst)
}
fn is_head(&self, node: *mut Node<T>) -> bool {
self.head() == node
}
fn is_tail(&self, node: *mut Node<T>) -> bool {
self.tail() == node
}
fn _search(&self, v: &T) -> (*mut Node<T>, *mut Node<T>) {
// left should be the last node with value less than v
let mut left = self.head();
// head must have at least one following node
let mut left_next = next!(left);
'again: loop {
// 1. find the last not marked node with value smaller than the given one
let mut curr = next!(self.head());
while !self.is_tail(curr) && as_ref!(curr).get() < Some(v) {
// curr is not tail, which means we can call next()
let next = next!(curr);
left = curr;
left_next = next;
curr = if Node::is_marked(next) {
// curr must have at least two following nodes: <marked> <tail>
next!(next)
} else {
next
};
}
let right = curr;
// 2. check if no marked nodes lie in between
if left_next == right {
if !self.is_tail(right) && Node::is_marked(next!(right)) {
// we should not see consecutive marked nodes
continue 'again;
} else {
return (left, right);
}
}
// 3. remove any marked node
if as_ref!(left).try_swap_next(left_next, right) {
// TODO: we should drop the deleted node
// However, this node might still be used by others
// let b = unsafe { Box::from_raw(left_next) };
if !self.is_tail(right) && Node::is_marked(next!(right)) {
continue 'again;
} else {
return (left, right);
}
}
}
}
fn _find(&self, v: &T) -> (bool, *mut Node<T>, *mut Node<T>) {
let (l, r) = self._search(v);
let exists = !self.is_tail(r)
&& !r.is_null()
&& as_ref!(r)
.get()
.and_then(|x| (x == v).then_some(()))
.is_some();
(exists, l, r)
}
pub fn find(&self, v: &T) -> bool {
self._find(v).0
}
pub fn get(&self, v: &T) -> Option<&T> {
let (exists, _, r) = self._find(v);
if exists {
as_ref!(r).value.as_ref()
} else {
None
}
}
pub fn insert(&self, v: T) -> Result<bool> {
let node = AtomicPtr::new(Box::leak(Box::try_new(Node::new(v))?));
let n = node.load(SeqCst);
let v = as_ref!(n)
.get()
.expect("failure accessing just created node");
loop {
let (exists, l, r) = self._find(v);
if exists {
// release the memory
unsafe { Box::from_raw(node.load(SeqCst)) };
return Ok(false);
} else {
as_ref!(n).link(r);
if as_ref!(l).try_swap_next(r, n) {
return Ok(true);
}
}
}
}
pub fn delete(&self, v: &T) -> bool {
loop {
let (exists, left, right) = self._find(v);
if exists {
let right_next = next!(right);
if !Node::is_marked(right_next)
&& as_ref!(right).try_swap_next(right_next, Node::mark(right_next))
{
if as_ref!(left).try_swap_next(right, right_next) {
self._find(v);
}
return true;
}
} else {
return false;
}
}
}
pub fn empty(&self) -> bool {
self.is_tail(next!(self.head()))
}
}
@vtta
Copy link
Author

vtta commented May 17, 2023

Lesson learned: The challenge of designing a lock-free data structure is how to do reclamation.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment