This is a follow-up to my previous post, where we reproduced the ABA problem in a lock-free stack by manually sleeping the threads to achieve the desired scheduling. This is, however, an undeniably unviable strategy for testing and finding bugs in real life. To quote the docs, Loom offers a way to “run tests many times, permuting the possible concurrent executions of each test according to what constitutes valid executions under the C11 memory model.”

My goal for this post is to explore the usage of loom and apply it to our broken lock-free stack, to see if it can detect the aba bug. It’s my first time working with this tool, so it should be exciting!

The code can be found on my github!

ABA Scenario Recap

Here’s a diagram for a quick recap of the ABA scenario:

Loom Setup

Shim For Using Loom Primitives

The way loom tests our code essentially is by replacing the std primitives with its own:

What we want is to be able to use the std types normally and only replace them when compiling with the loom flag.

So, for running the loom tests, we pass RUSTFLAGS="--cfg loom".

In order to toggle the loom types, we can add a shim to our lib.rs file:

#[cfg(not(loom))]
pub use std::sync::atomic::{AtomicPtr, Ordering};

#[cfg(loom)]
pub use loom::sync::atomic::{AtomicPtr, Ordering};

I’m just demonstrating the usecase, but we’d need to do this for all of types we want to replace (check out the complete lib.rs for this demo).

Cargo.toml

Then, we want to add loom to our project, but gated behind a config flag, so that it doesn’t bloat our production builds. I added this my Cargo.toml:

// tells cargo to only include when compiling with --cfg loom
[target.'cfg(loom)'.dependencies]
loom = "0.7"

// registers loom as custom configuration flag (gets rid of compiler warnings)
[lints.rust]
unexpected_cfgs = { level = "warn", check-cfg = ['cfg(loom)'] }

Wrapping Raw Pointers In UnsafeCell

This is something I didn’t realize until later and didn’t understand why the loom tests passed for the broken stack ;)

We need to wrap the fields that get unsafely read or written to across threads in loom::cell::UnsafeCell. This gives loom visibility over the memory that it needs to track for conflicting accesses.

Since the API for UnsafeCell in std and loom are a bit different, the (docs)[https://docs.rs/loom/latest/loom/#handling-loom-api-differences] recommends adding:

#[cfg(not(loom))]
#[derive(Debug)]
pub(crate) struct UnsafeCell<T>(std::cell::UnsafeCell<T>);

#[cfg(not(loom))]
impl<T> UnsafeCell<T> {
    pub(crate) fn new(data: T) -> UnsafeCell<T> {
        UnsafeCell(std::cell::UnsafeCell::new(data))
    }

    pub(crate) fn with<R>(&self, f: impl FnOnce(*const T) -> R) -> R {
        f(self.0.get())
    }

    pub(crate) fn with_mut<R>(&self, f: impl FnOnce(*mut T) -> R) -> R {
        f(self.0.get())
    }
}

Tldr; loom::cell::UnsafeCell uses .with and .with_mut to track reads and writes, whereas std::cell::UnsafeCell just uses .get, so we need to work around this a bit.

Node becomes:

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

For the push I’ve replaced

unsafe { (*new_head).next  = current_head };

with

unsafe {
    (*new_head).next.with_mut(|next_ptr| {
    *next_ptr = current_head;
    });
}

Let’s break this down:

(*new_head).next
│          │
│          └── retrieves next, which is of type UnsafeCell<*mut Node>
│
└── dereference
       navigates through the new_head raw pointer to get the actual Node struct


        .with_mut( |next_ptr| { ... } )
        │          │
        │          └── closure arg: next_ptr
        │                 the closure receives next_ptr, which is a pointer-to-the-pointer:
        │                 type: `*mut (*mut Node)`
        │                 think of this as: "a pointer targeting the inner pointer slot"
        │
        └── Calls .with_mut() on the UnsafeCell
               - cfg(loom): loom sees that a thread is writing
               - cfg(not(loom)): standard raw access
        
        
            *next_ptr = current_head;
            │         │
            │         └──  set the value inside the slot to point to current_head
            │
            └── dereference and write
                   *next_ptr dereferences the outer pointer to reach the inner pointer slot

For the pop, I need to replace

let new_head = unsafe { (*current_head).next };

with

let new_head = unsafe {
    (*current_head).next.with(|next_ptr| *next_ptr)
};

ABA Loom Test

We are now ready to write our loom test (separte from the normal #[test]s). We wrap everything in a loom::model. And… I had to rewrite this a bit because loom doesn’t have std::thread::scope. The reason for this, from what I could understand, is that std::thread::scope works with standard OS scheduling that loom cannot intercept.

With that out of the way, so the way this works is that loom runs the loom::model(|| { ... }) closure many times, each time with a different schedule and thread-interleaving permutation.

#[cfg(test)]
#[cfg(loom)]
mod loom_tests {
    use super::*;
    use crate::thread;
    use crate::Arc;
    use loom::model;

    #[test]
    fn aba_problem() {
        model(|| {
            let stack = Arc::new(Stack::<i32>::new());
            let s2 = stack.clone();
            stack.push(1);
            stack.push(2);
            stack.push(3);

            let t1 = thread::spawn(move || {
                stack.pop();
            });

            let t2 = thread::spawn(move || {
                s2.pop();
                s2.pop();
                s2.push(4);
            });

            t1.join().unwrap();
            t2.join().unwrap();
        });
    }
}

Running the test:

RUSTFLAGS="--cfg loom" RUST_BACKTRACE=1 cargo test --release aba_problem

Big Reveal **happy noises** (or maybe I celebrated too early)

running 1 test
test naive_lock_free_stack::loom_tests::aba_problem ... FAILED
[...]
(25841) thread `naive_lock_free_stack::loom_tests::aba_problem` panicked at loom-0.7.2/src/rt/object.rs:286:38:
index out of bounds: the len is 7 but the index is 34091991057


  
[ Click to expand full backtrace ]
  

     = note: stack backtrace:
             0: __rustc::rust_begin_unwind
             1: core::panicking::panic_fmt
             2: core::panicking::panic_bounds_check
             3: <scoped_tls::ScopedKey<core::cell::RefCell<loom::rt::scheduler::State>>>::with::<<loom::rt::scheduler::Scheduler>::with_state<<loom::rt::scheduler::Scheduler>::with_execution<loom::rt::synchronize<<loom::rt::cell::Cell>::start_write::{closure#0}, loom::rt::cell::Writing>::{closure#0}, loom::rt::cell::Writing>::{closure#0}, loom::rt::cell::Writing>::{closure#0}, loom::rt::cell::Writing>
             4: <loom::rt::cell::Cell>::start_write
             5: <[...]::naive_lock_free_stack::Stack<i32>>::pop
             6: <loom::rt::spawn<loom::thread::spawn_internal<visualizing_crossbeam_epoch::naive_lock_free_stack::loom_tests::aba_problem::{closure#0}::{closure#1}, ()>::{closure#0}>::{closure#1} as core::ops::function::FnOnce<()>>::call_once::{shim:vtable#0}
             7: <generator::stack::StackBox<<generator::gen_impl::GeneratorImpl<core::option::Option<alloc::boxed::Box<dyn core::ops::function::FnOnce<(), Output = ()>>>, ()>>::init_code<loom::rt::scheduler::spawn_thread::{closure#0}>::{closure#0}>>::call_once
             8: generator::detail::gen::gen_init_impl
             9: generator::detail::asm::gen_init
  

note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace.


failures:
    naive_lock_free_stack::loom_tests::aba_problem

test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

Soooo… the test fails with an index out of bounds. I looked into what this means, and in loom’s runtime engine, every atomic, UnsafeCell, thread and allocation is assigned a small object id. In this case, the valid ones are from 0 to 6, so this huge index basically doesn’t belong to any of the objects that loom is tracking.

However, this doesn’t yet prove that we’ve managed to hit the ABA path yet. After printing printing the addresses, it turns out it’s just a use-after-free case. I tried running it a few more times but no luck.

Essentially, what an use-after-free scenario could look like:

Soooo.. if we did want to reproduce the exact ABA scenario, we’d need to interfere with the timing ourselves, like in the previous post, which I believe doesn’t add any more value at this point.

Conclusion

While it is a bit dissapointing that we couldn’t catch the actual ABA bug, the why we couldn’t helped me better understand what loom actually does. Since it runs different permutations until it panics, the first time it encounters a typical use-after-free, it stops there. It did prove, however, that loom does detect that our stack is indeed broken, which is the most important thing.

Originally, I had written a section discussing the use-after-free bug on my previous post, which I removed because it added too much complexity to the structure of the article without much extra benefit. So, here, I used that diagram that I had created before, so that’s something :)

Thank you so much for tagging along! Any feedback is appreciated as I’m just dipping my toes in these tools.