Learn Rust Series (#68) - Building a Lock-Free Stack (Treiber)

scipio(71)

Published in

StemSocial

Words

3163

Reading

15 min

Listen

Play

Learn Rust Series (#68) - Building a Lock-Free Stack (Treiber)

What will I learn

You will build a real lock-free stack, the classic Treiber stack, using only atomics;
how AtomicPtr holds the head pointer and how a CAS loop implements push and pop with no lock;
how Box::into_raw and Box::from_raw hand ownership to a raw pointer and reclaim it exactly once;
why push/pop take &self, so scoped threads can share the stack without Arc or a Mutex;
what the ABA problem is, and why production lock-free code needs safe memory reclamation.

Requirements

A working modern computer running macOS, Windows or Ubuntu;
An installed Rust toolchain (via rustup, from rustup.rs);
The previous sixty-seven episodes, especially atomics, memory ordering, and raw pointers;
The ambition to learn systems programming from the ground up.

Difficulty

Advanced

Curriculum (of the `Learn Rust Series`):

Learn Rust Series (#68) - Building a Lock-Free Stack (Treiber)

Time to put the atomics to work and build something real: a lock-free stack. The Treiber stack, from a 1986 IBM paper by R. Kent Treiber, is the "hello world" of lock-free data structures, and it is genuinely beautiful once you have compare-and-swap in your toolbox. Instead of a lock, the head pointer is an AtomicPtr, and both push and pop are CAS loops that optimistically try to swing the head and simply retry if another thread got there first. There is no blocking anywhere -- no thread ever waits on another to release something -- and yet the structure stays correct under any amount of concurrent hammering. Building one teaches you exactly how the lock-free world thinks, sharp edges very much included ;-)

I want to be honest up front, the way I was last episode: this is the hardest data structure we have built so far, and it is the first one where a single misplaced ordering flag is the difference between "works" and "silently corrupts memory". So we are going to build up to it in small, understandable steps -- a lock-based version for contrast, then the raw-ownership plumbing, then the head pointer, and only then the whole thing -- and I will keep pointing back to episode 66 (atomics) and episode 67 (memory ordering) as we go, because this is where all of that theory finally gets spent.

Solutions to Episode 67 Exercises

Episode 67 was memory ordering -- Relaxed, Acquire, Release, and SeqCst -- so all three solutions live in std::sync::atomic.

Exercise 1 asked for a Release/Acquire flag that publishes a value stored just before the flag is set: one thread writes a number with a Relaxed store, then flips an AtomicBool with Release; another spins on an Acquire load and prints the number once the flag is seen:

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::thread;

fn main() {
    let ready = Arc::new(AtomicBool::new(false));
    let data = Arc::new(AtomicU32::new(0));
    let (r, d) = (Arc::clone(&ready), Arc::clone(&data));
    let t = thread::spawn(move || {
        d.store(7, Ordering::Relaxed);     // write the payload first
        r.store(true, Ordering::Release);  // publish -- no earlier write may cross this
    });
    while !ready.load(Ordering::Acquire) {} // once we see true, the 7 is visible too
    t.join().unwrap();
    println!("{}", data.load(Ordering::Relaxed)); // 7, guaranteed
}

The Release store forbids the data = 7 write from sliding past it, the Acquire load forbids the data read from sliding before the flag check, and together they form the happens-before edge that drags the payload across. This exact handshake is the machinery we lean on in the stack below.

Exercise 2 wanted a plain counter incremented from several threads with fetch_add(1, Ordering::Relaxed), joined, totalled, with a comment on why Relaxed is safe here:

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::thread;

fn main() {
    let count = Arc::new(AtomicU64::new(0));
    let mut handles = Vec::new();
    for _ in 0..4 {
        let c = Arc::clone(&count);
        handles.push(thread::spawn(move || {
            for _ in 0..250 {
                // Relaxed is safe: no other memory hangs off this value,
                // we only need each +1 to be atomic, not ordered vs other data.
                c.fetch_add(1, Ordering::Relaxed);
            }
        }));
    }
    for h in handles { h.join().unwrap(); }
    println!("{}", count.load(Ordering::Relaxed)); // 1000, exactly
}

Relaxed is correct precisely because no thread ever uses the counter's value to decide it is safe to touch some other piece of memory. It is a standalone statistic, so atomicity is all we need and we pay for nothing more.

Exercise 3 asked for a SeqCst one-shot claim that several threads race for, with a comment on why a weaker ordering would leave the outcome unclear:

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread;

fn main() {
    let claimed = Arc::new(AtomicBool::new(false));
    let mut handles = Vec::new();
    for id in 0..8 {
        let c = Arc::clone(&claimed);
        handles.push(thread::spawn(move || {
            // SeqCst so every thread agrees on a single global order of who claimed first.
            if c.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst).is_ok() {
                println!("thread {id} claimed it");
            }
        }));
    }
    for h in handles { h.join().unwrap(); }
}

Exactly one thread finds false, installs true, and wins; every other thread now finds true and stays silent. SeqCst gives all threads one agreed-upon timeline, which is exactly what a "who was first" question needs. Right -- homework cleared, now the stack.

First, the lock-based version for contrast

To appreciate the lock-free version, here is the boring-but-correct one we could have written back in episode 65: a Vec behind a Mutex. It works perfectly, and honestly for most programs it is the right answer. But every push and pop has to take the single lock, so threads serialise -- only one can touch the stack at a time, and the rest queue up behind it:

use std::sync::Mutex;

struct LockedStack<T> {
    items: Mutex<Vec<T>>,
}

impl<T> LockedStack<T> {
    fn new() -> Self { LockedStack { items: Mutex::new(Vec::new()) } }
    fn push(&self, v: T) { self.items.lock().unwrap().push(v); }
    fn pop(&self) -> Option<T> { self.items.lock().unwrap().pop() }
}

fn main() {
    let s = LockedStack::new();
    s.push(1);
    s.push(2);
    println!("{:?} {:?}", s.pop(), s.pop()); // Some(2) Some(1)
}

The lock-free version we are about to build removes that serialisation entirely. No thread ever blocks another. When two threads collide, the loser does not wait -- it simply notices its attempt failed and tries again immediately with fresh information. That is the whole philosophical difference: a lock makes threads take turns, whereas a lock-free structure lets them all charge ahead and reconciles the collisions after the fact. Having said that, "lock-free" does not mean "free" -- it means the cost moves from waiting to retrying, and it means you take on the memory-management job the lock used to hide.

The building block: raw ownership

A lock-free structure cannot store a Box directly in the shared pointer, because an atomic must hold a plain machine address -- a *mut T -- not a smart pointer with a destructor. So we need a way to turn an owned heap allocation into a bare address, hand that address around between threads, and later turn it back into something that will free itself. That is exactly what Box::into_raw and Box::from_raw are for. into_raw releases the allocation as a raw pointer and gives up ownership (so the Box will no longer free it), and from_raw reclaims that same pointer as a Box and takes ownership back (so it will free it, once):

fn main() {
    let raw: *mut String = Box::into_raw(Box::new(String::from("heap")));
    // The address could now live in an AtomicPtr, be handed between threads, etc.
    // Nothing frees it while it is "just an address" -- that is the danger and the point.
    let back: Box<String> = unsafe { Box::from_raw(raw) }; // reclaim ownership exactly once
    println!("{}", *back); // heap
}

The word "exactly once" is doing enormous work in that comment, and it is the single most important rule of this whole episode. Every allocation that leaves through into_raw must come back through from_raw precisely one time. Call from_raw twice on the same address and you get a double-free (undefined behaviour, the classic memory-corruption bug). Never call it at all and you leak the allocation forever. This is the accounting the Mutex version did for us behind the scenes with Vec, and now it is our responsibility to keep the books straight. Remember Box::into_raw and Box::from_raw from episode 41 when we built our own smart pointer -- same tools, higher stakes.

The head pointer as an AtomicPtr

The stack's head is an AtomicPtr<Node<T>>. An AtomicPtr is just an atomic wrapper around a raw pointer -- it holds an address and lets many threads load, store, and compare-and-swap it without a data race. Crucially, AtomicPtr is Sync unconditionally, which is precisely what will let a bare &Stack be shared across threads. Here is how load, store, and null behave on one in isolation, before we wire it into anything:

use std::sync::atomic::{AtomicPtr, Ordering};
use std::ptr;

fn main() {
    let mut value = 42;
    let p = AtomicPtr::new(&mut value as *mut i32);
    let loaded = p.load(Ordering::Acquire);       // read the current address
    println!("{}", unsafe { *loaded }); // 42     // deref it (unsafe: we vouch it is valid)
    p.store(ptr::null_mut(), Ordering::Release);   // an empty stack is a null head
    println!("null now? {}", p.load(Ordering::Acquire).is_null()); // true
}

Notice the orderings already: I load with Acquire and store with Release, not Relaxed. That is not decoration. When the head points at a node, following that pointer means reading the node's fields -- so the load that publishes "here is the new head" must be an Acquire, and the store that installs it must be a Release, so that the node's contents (written before the store) are guaranteed visible to whoever reads the head afterwards. This is the episode-67 handshake applied to a data structure: the head pointer is the flag, and the node it points at is the payload being published.

The Treiber stack

Now the whole thing. Each Node holds a value and a raw next pointer to the node beneath it. push allocates a node, points its next at the current head, and CAS-swaps the head to it. pop reads the head, reads the head's next, and CAS-swaps the head down to next, then reclaims the old head. Both operations loop: they read, prepare, attempt the swap, and retry from scratch if the head moved underneath them. And because we now own the raw nodes, Drop has to walk the list and free whatever is left:

use std::sync::atomic::{AtomicPtr, Ordering};
use std::ptr;

struct Node<T> {
    value: T,
    next: *mut Node<T>,
}

pub struct Stack<T> {
    head: AtomicPtr<Node<T>>,
}

impl<T> Stack<T> {
    pub fn new() -> Self {
        Stack { head: AtomicPtr::new(ptr::null_mut()) }
    }

    pub fn push(&self, value: T) {
        let node = Box::into_raw(Box::new(Node { value, next: ptr::null_mut() }));
        loop {
            let head = self.head.load(Ordering::Acquire);
            unsafe { (*node).next = head; } // our new node points at the old top
            // swing head to our node; if another thread moved it first, retry
            if self.head
                .compare_exchange(head, node, Ordering::Release, Ordering::Acquire)
                .is_ok()
            {
                break;
            }
        }
    }

    pub fn pop(&self) -> Option<T> {
        loop {
            let head = self.head.load(Ordering::Acquire);
            if head.is_null() {
                return None; // empty stack
            }
            let next = unsafe { (*head).next };
            // try to move head down to next; retry if someone beat us to it
            if self.head
                .compare_exchange(head, next, Ordering::Release, Ordering::Acquire)
                .is_ok()
            {
                let boxed = unsafe { Box::from_raw(head) }; // reclaim and free the popped node
                return Some(boxed.value);
            }
        }
    }
}

impl<T> Drop for Stack<T> {
    fn drop(&mut self) {
        while self.pop().is_some() {} // free every remaining node
    }
}

fn main() {
    let stack = Stack::new();
    stack.push(1);
    stack.push(2);
    stack.push(3);
    println!("{:?}", stack.pop()); // Some(3), last in first out
    println!("{:?}", stack.pop()); // Some(2)
    println!("{:?}", stack.pop()); // Some(1)
    println!("{:?}", stack.pop()); // None
}

Let us walk push slowly, because the CAS loop is the heart of everything. First we allocate the node and get a raw pointer to it. Then, inside the loop, we read the current head. We set our node's next to that head, so our node now sits logically on top of the existing stack. Then comes the pivotal line: compare_exchange(head, node, ...) says "if the head is still the value I just read, replace it with my node; otherwise tell me it changed". If it succeeds, we are done -- our node is the new top. If it fails, some other thread pushed or popped in the microseconds since our read, so the head we captured is stale. compare_exchange hands us back the new head via its Err, the loop goes round again, we re-point our node's next at the fresh head, and we retry. We never corrupt anything, because the swap only ever lands when the world is exactly as we last saw it.

pop is the mirror image, with one extra subtlety. We read the head; if it is null the stack is empty and we return None. Otherwise we read the head node's next, and CAS the head down to that next. On success, we are the unique thread that unlinked this particular node -- so we, and only we, are allowed to from_raw it and free it. That uniqueness is what keeps the "exactly once" rule intact: the winning CAS is the moment ownership of the popped node transfers cleanly to us. The Drop impl reuses pop in a loop to drain and free anything still on the stack when it goes out of scope, so we do not leak.

One thing worth flagging: this stack is correct, but it is not the last word, and I will show you the flaw shortly. For a single-threaded sanity check, though, the main above behaves exactly like any stack -- push 1, 2, 3, pop them back as 3, 2, 1, then None. Last in, first out, no surprises.

Sharing it across threads, lock-free

Here is the payoff. Because push and pop both take &self (a shared reference, never &mut self), and because AtomicPtr is Sync, the whole Stack is Sync automatically. That means we can hand out a plain &stack to many threads at once. Combined with thread::scope from episode 62 -- which lets threads borrow local data without Arc because the scope guarantees they finish before the borrow ends -- we get concurrent pushes with no lock and no reference counting anywhere in sight:

use std::sync::atomic::{AtomicPtr, Ordering};
use std::ptr;
use std::thread;

struct Node { value: i32, next: *mut Node }
struct IntStack { head: AtomicPtr<Node> }

impl IntStack {
    fn new() -> Self { IntStack { head: AtomicPtr::new(ptr::null_mut()) } }
    fn push(&self, value: i32) {
        let node = Box::into_raw(Box::new(Node { value, next: ptr::null_mut() }));
        loop {
            let head = self.head.load(Ordering::Acquire);
            unsafe { (*node).next = head; }
            if self.head.compare_exchange(head, node, Ordering::Release, Ordering::Acquire).is_ok() {
                break;
            }
        }
    }
    fn pop(&self) -> Option<i32> {
        loop {
            let head = self.head.load(Ordering::Acquire);
            if head.is_null() { return None; }
            let next = unsafe { (*head).next };
            if self.head.compare_exchange(head, next, Ordering::Release, Ordering::Acquire).is_ok() {
                return Some(unsafe { Box::from_raw(head) }.value);
            }
        }
    }
}

fn main() {
    let stack = IntStack::new();
    thread::scope(|s| {
        let stack_ref = &stack;
        for i in 0..8 {
            s.spawn(move || stack_ref.push(i)); // 8 threads push concurrently, lock-free
        }
    });
    let mut count = 0;
    while stack.pop().is_some() { count += 1; }
    println!("popped {count} items after concurrent pushes"); // popped 8 items
}

Eight threads push into the same stack at the same time, every one of them going through the CAS loop, colliding, retrying, and eventually all landing their node. When the scope ends and every pusher has joined, the main thread drains the stack and finds all 8 items present -- none lost, none duplicated. This is the thing a lock would have forced into single file, now happening genuinely in parallel. It is a small program, but it is doing something that would make a C programmer sweat, and Rust let us express the shared-across-threads part with an ordinary &.

The catch: the ABA problem

I promised a flaw, and here it is -- a real, famous one that has bitten production systems. Picture this interleaving in pop. Thread 1 reads the head, sees it points at node A, and reads A's next (say it points at node B). Now thread 1 gets suspended by the scheduler for a moment. In that window, thread 2 pops A (freeing it), pops B, and then pushes a brand new node -- and the allocator, being thrifty, happens to reuse A's old address for it. The head now points at "A" again, but it is a different node that merely lives at the same address. Thread 1 wakes up, does its compare_exchange(A, B, ...), sees the head is "still A", concludes nothing changed, and swaps in B -- a node that was already popped and freed. The stack is now corrupt, pointing into reclaimed memory. That is the ABA problem: a value went A -> B -> A, and a naive compare-and-swap cannot tell the difference between "never changed" and "changed and changed back":

fn main() {
    // A location reads A, changes to B, then back to A. A naive compare_exchange
    // sees "still A" and wrongly concludes nothing happened in between.
    let (a1, b, a2) = ("A", "B", "A");
    let looks_unchanged = a1 == a2; // true...
    println!("looks unchanged? {looks_unchanged}, yet it went {a1}->{b}->{a2}"); // true
}

The root cause is memory reuse: the CAS compares addresses, and addresses get recycled. Real lock-free stacks defeat this in one of two ways. The first is a tagged pointer -- you glue a version counter next to the address and CAS both together, so even if the address returns to A, the counter has advanced and the CAS correctly fails. The second, and by far the more common approach in Rust, is safe memory reclamation: you simply do not free a popped node until you can prove no other thread could still be holding a stale pointer to it. The go-to crate for that is crossbeam-epoch, which uses epoch-based reclamation to defer frees until every thread has moved past the danger window. There is also hazard-pointer machinery in crates like haphazard. Either way, the lesson of Phase 5's hardest episode is this: lock-free algorithms are writable by mortals, but lock-free memory reclamation is the dragon, and in production you reach for a vetted crate rather than rolling your own.

So is our stack useless? Not at all -- it is correct as long as popped memory is not reused while another thread still holds a stale pointer to it, which is a genuinely strong caveat but not an impossible one to satisfy in constrained settings. For learning, it is perfect: it shows you the real shape of lock-free code, the CAS-loop-and-retry pattern, and exactly where the difficulty hides. Just do not ship it as-is under high contention without a reclamation scheme, and you will have understood the honest state of the art. Next episode we take on the queue, which is harder still -- two ends to coordinate instead of one -- and the ABA dragon gets even hungrier there.

Bedankt en tot de volgende keer!

Building a Treiber stack is one of those exercises that rewires how you think about concurrency. Once you have written a CAS loop by hand and watched eight threads reconcile their collisions with no lock in sight, the whole lock-free world stops being mysterious and starts being engineering -- hard engineering with a memory-reclamation catch, but engineering none the less. Try the exercises below, break the thing on purpose, and I will see you in the next one ;-)

Exercises

Add an is_empty(&self) -> bool method to Stack<T> that returns whether the head pointer is currently null, using a single Acquire load. Note in a comment why the answer can be stale the instant it is returned under concurrency.
Add a peek(&self) -> Option<i32> to the IntStack that returns a copy of the top value without popping it, and explain in a comment why reading through the head pointer is only safe here under low contention (hint: the ABA and use-after-free windows).
Stress-test the IntStack: push from four scoped threads (each pushing a range of numbers), join them, then pop everything on the main thread and count the items -- assert the count equals the total number pushed.

@scipio