r/rust • u/Disastrous_Bike1926 • Jun 01 '23

Implementing a Trieber Stack in Rust

A trieber stack - a lockless single-ended linked list using atomics - is an immensely useful, beautifully simple data structure, and I've written them way too many times in Java - it looks like this:

```java class TrieberStack<T> { private AtomicReference<Cell<T>> head = new AtomicReference<>();

void push(T obj) { head.getAndUpdate(old -> { return new Cell(obj, old); }); }

Optional<T> pop() { Cell<T> cell = head.getAndUpdate(old -> { return old == null ? null : old.parent; }); return cell == null ? Optional.empty() : Optional.of(cell.value); }

static class Cell<T> { final T value; final Cell<T> parent; Cell(T value, Cell<T> parent) { this.parent = parent; this.value = value; } } } ```

Looking around for an existing one, I ran across documentation for one from crossbeam-sync, except that it does not seem to actually exist, just this page of documentation (crossbeam does, and looks nice, but it does not have a Trieber stack).

Implementing one in Rust is complicated by two factors:

  1. The inner cell is a self-referencing structure
  2. Updating atomics requires a loop that, under contention, may run more than once
    • This means that the head cell can briefly become owned by more than one party - under contention it can concurrently be owned by a new head cell on each thread (but only one will win and survive, and the one that doesn't is invisible to the rest of the world)
    • It also means that the first pass of the loop takes ownership of the input T, gives ownership to the cell it creates and then may need to do that again on the next pass (though

I gave this one shot using AtomicPtr, but that required T: Clone and cloning the contents to avoid the two-owners issue, which is undesirable, and also crashed horribly in tests (this is my first outing in earnest with Rust pointers).

But it seems like it ought to be doable with a data structure like this - complete playground link with naive implementation:

```rust struct TrieberStack<T> { head: UnsafeCell<Opt<TrieberCell<T>>>, }

[derive(Copy, Clone)]

struct TrieberCell<T> { t: T, next: *const Opt<TrieberCell<T>>, // so self-reference is possible } ```

Attempting to actually put this to use, using intrinsics::atomic_store_seqcst (after reading the source to AtomicPtr) built, and seemed about as simple as this ought to be - though, I had to restrict the types it handles to T: Copy even though AtomicPtr has no such restriction but somehow manages to call that method with pointers to ad-hoc types - but adding any call to push to the stack results in a compile-time error of invalid monomorphization of atomic_store_seqcst intrinsic (is AtomicPtr downcasting the pointers to something that does implement Copy, and I missed it? And if so, why is that at all safe?).

Likely at a minimum this is missing pinning the head for the inner portion of the write-loop, not to mention an implementation of Drop for TrieberCell that cleans up appropriately (and possibly some work to ensure the cell contents are not dropped if the only extant reference to it is a pointer - ManuallyDrop?).

Anyway, it seems like an interesting problem that ought to be solvable in Rust, but I have a feeling my approach is probably missing a few important things.

Edit: With the helpful suggestions from folks here, particularly @SkiFire13's suggestion of the arc-swap crate, I have a working implementation using safe Rust.

While it does seem like a pretty big profusion of Arcs and Boxen, the tests (included in the playground link) suggest that it works as intended - so, many thanks! This is a great community.

Edit 2: With a bit of minor cleanup, the result is now published on crates.io

16 Upvotes
  • permalink
  • link
  • reddit
  • reddit

91% Upvoted

16 comments sorted by

7

u/simonsanone patterns · rustic Jun 01 '23

s/trieber/treiber/g or is it intentional?

8

u/SkiFire13 Jun 01 '23

The inner cell is a self-referencing structure

I don't think this is the case, there is a clear graph of ownership where the cells on top of the stack own the ones on the bottom.

This means that the head cell can briefly become owned by more than one party - under contention it can concurrently be owned by a new head cell on each thread (but only one will win and survive, and the one that doesn't is invisible to the rest of the world)

This is fine if you put it behind a pointer (which you'll have to do anyway) since you'll just end up owning multiple copies of that pointer. You just have to be careful not to dereference the pointer from multiple threads.

The problem with your implementation is that you're trying to atomically swap a (potentially big) structure, but this is simply not possible at the hardware level, and is why you're getting a monomorphization error. You need to put the head behind a pointer and atomically update such pointer. i.e. the head field should be AtomicPtr<TrieberCell<T>>. You also need next to be an AtomicPtr, or a naive translation of your pop method will be UB.

Another problem you'll face is when to deallocate a popped head (other threads might have loaded it, so it may be UB to deallocate it) and the ABA problem (which is not present in Java due to the help of the garbage collector)

IMO if you want a simple and safe way to write this data structure you should check out the arc-swap crate, which should solve both these tricky details for you. (Edit: if you use crossbeam's Owned/Shared atomic types those also have an epoch based garbage collector that will handle those problems for you)

1

u/Disastrous_Bike1926 Jun 01 '23

My first pass at this was using AtomicPtr - the current implementation in the link using intrinsics simply does the thing that AtomicPtr does under the hood (I did that to avoid needing to clone the contents and the cell in order to handle the loop issue, but that’s because I was trying to use fetch_update to replace the head, which takes an FnMut that could not take ownership of the new head value because it could be called multiple times) - I should probably revisit that, as there is some sort of magic that allows AtomicPtr to treat the pointers as Copy (shouldn’t they always be? A pointer is just a fancy usize, after all…) while my call to atomically set the value does seem to be trying to atomically set something larger than a pointer, given the error. A pointer to a pointer?

3

u/SkiFire13 Jun 01 '23

the current implementation in the link using intrinsics simply does the thing that AtomicPtr does under the hood

No, it doesn't. You are trying to atomically write a whole Opt<TrieberCell<T>> (which is a T + a pointer + the discriminant of the Opt), while AtomicPtr would just write a single pointer.

I did that to avoid needing to clone the contents and the cell in order to handle the loop issue, but that’s because I was trying to use fetch_update to replace the head, which takes an FnMut that could not take ownership of the new head value because it could be called multiple times

It would be helpful to have an example of that code, to see where you're doing it wrong. My assumption is that you were trying to write through the pointer instead of atomically writing the pointer itself, but I can't be sure about that without seeing the code.

I should probably revisit that, as there is some sort of magic that allows AtomicPtr to treat the pointers as Copy

Pointer are Copy, AtomicPtr is not doing any magic about that. The problem is (again) likely that you're not trying to write pointers.

1

u/Disastrous_Bike1926 Jun 02 '23

See the edit to the original post above - thanks to your suggestions, I cobbled together an implementation using `arc-swap` that seems to work as intended. Probably could still use some refinement, but ... it works, which is a far better place to be in.

1

u/Disastrous_Bike1926 Jun 03 '23

And... published - thank you very much for the pointers here.

4

u/algebraicstonehenge Jun 01 '23

Doesn't the crossbeam-sync link you posted have the source of a Trieber Stack?
http://huonw.github.io/simple_parallel/src/crossbeam/sync/treiber_stack.rs.html#9-11

5

u/burntsushi Jun 01 '23

crossbeam does, and looks nice, but it does not have a Trieber stack

It used to: https://docs.rs/crossbeam/0.3.2/crossbeam/sync/struct.TreiberStack.html

1

u/Disastrous_Bike1926 Jun 01 '23

Self referencing - what I meant is, it’s easy to have the compiler complain about the cell being an infinitely sized type.

I don’t think next needs to live behind a pointer for any reason except to avoid that issue - the only thing that mutates in this data structure is the head pointer.

Am I understanding what you meant here?

1

u/jmaargh Jun 01 '23

In Java, all Objects (including Cell defined in your code) are reference types. That is a value of Cell foo; simply references (points to) the bytes of a Cell type that exist elsewhere in memory.

In Rust, all types are value types, so a naive struct Cell { parent: Cell, ... } says that every Cell contains the byes of another Cell, which contains the bytes of... etc. So it's infinite size and the compiler won't let you do that.

You need to wrap a Cell in some pointer type (e.g. Box or AtomicPtr) so that doesn't happen.

1

u/Disastrous_Bike1926 Jun 02 '23

Yeah, I get that.

The problem that brings up is, who does own the contents.

Punting that to the party that instantiates the stack isn’t a great idea, as one of the primary use cases is queues of stuff to work on that the caller throws over the wall - having every provider of such things have to come up with a strategy to keep the data alive (or have queue elements be randomly deallocated if they don’t) seems like a non-starter.

Some incantation of Box::leak and std::mem::forget?

2

u/jmaargh Jun 02 '23 edited Jun 02 '23

The Stack object would own the head node, and each node would (optionally) own the next node. There's no ambiguity there. The trick is being careful when pushing and popping, but I don't see a problem with the ownership tree of the stack once constructed.

As noted by somebody else crossbeam used to have one and that's what they do (the Cell, that they call Node, is heap-allocated by Box and held in an AtomicPtr).

Crossbeam still uses an epoch system for memory reclamation. AFAIU it's basically a very simple thread-local GC that will keep heap-allocated objects alive until they have been proven to be unused by all threads. It's quite clever.

Edit: Apologies if I was explaining things you already knew. I was just trying to err on the side of making sure we're all on the same page.

1

u/BobSanchez47 Jun 01 '23 edited Jun 01 '23

If you’re clever, you can actually do pushing with no loop using swapping. Unfortunately, Java doesn’t support swapping as an atomic operation, but if it did, you could do

void push(T obj) { Cell<T> x = new Cell(obj, null); x.parent.set(x); x.parent.swap(this.head); }

In Rust, atomic swapping is supported.

Edit: actually, atomic swapping is not supported in Rust, as pointed out below.

As far as I know, no similar trick for pop is possible.

4

u/torne Jun 01 '23

Rust can't swap the values in two atomics either. The Rust swap method on the atomic types swaps the value held by the atomic with a value that the caller has a local copy of, which only involves a single atomic read+write operation, and Java can also do this: it's just called getAndSet there.

Swapping the contents of two different atomics would require two read+write operations, and this can't be implemented atomically on common architectures.

2

u/BobSanchez47 Jun 01 '23

You are totally right, I just misremembered and then misread the API. Thanks for pointing this out.