Building a Lock-Free Multi-Producer Single-Consumer (MPSC) Ring Buffer in Rust with Atomic Memory Orderings — code snippets

snippet-1.txt

use std::cell::UnsafeCell;
use std::mem::MaybeUninit;
use std::sync::atomic::AtomicUsize;

#[repr(align(64))]
struct Cell<T> {
    // Monotonically increasing sequence to synchronize producers and consumer
    sequence: AtomicUsize,
    // The actual storage slot for the message, using UnsafeCell to allow interior mutability
    value: UnsafeCell<MaybeUninit<T>>,
}

// Ensure cache line size (64 bytes) to prevent false sharing
const CACHE_LINE_SIZE: usize = 64;

#[repr(align(64))]
pub struct MpscRingBuffer<T> {
    // Bounded buffer of cells. Size must be a power of two.
    buffer: Box<[Cell<T>]>,
    mask: usize,
    
    // Head index, modified only by the single consumer thread
    head: CacheAlignedUsize,
    
    // Tail index, competed for by multiple producer threads
    tail: CacheAlignedUsize,
}

#[repr(align(64))]
struct CacheAlignedUsize(AtomicUsize);

snippet-2.txt

impl<T> MpscRingBuffer<T> {
    pub fn new(capacity: usize) -> Self {
        assert!(capacity >= 2, "Capacity must be at least 2");
        assert!(capacity.is_power_of_two(), "Capacity must be a power of two");

        let mut buffer = Vec::with_capacity(capacity);
        for i in 0..capacity {
            buffer.push(Cell {
                sequence: AtomicUsize::new(i),
                value: UnsafeCell::new(MaybeUninit::uninit()),
            });
        }

        Self {
            buffer: buffer.into_boxed_slice(),
            mask: capacity - 1,
            head: CacheAlignedUsize(AtomicUsize::new(0)),
            tail: CacheAlignedUsize(AtomicUsize::new(0)),
        }
    }
}

snippet-3.txt

use std::sync::atomic::Ordering;

impl<T> MpscRingBuffer<T> {
    pub fn try_push(&self, value: T) -> Result<(), T> {
        let mut pos = self.tail.0.load(Ordering::Relaxed);
        loop {
            let cell = &self.buffer[pos & self.mask];
            let seq = cell.sequence.load(Ordering::Acquire);
            let diff = seq as isize - pos as isize;

            if diff == 0 {
                // The slot is empty and ready for this write cycle.
                // Attempt to claim the slot by advancing the tail pointer.
                match self.tail.0.compare_exchange_weak(
                    pos,
                    pos + 1,
                    Ordering::Relaxed,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => {
                        // Safety: We have claimed this slot uniquely. No other thread
                        // can write here. The consumer won't read it until we update the sequence.
                        unsafe {
                            *cell.value.get() = MaybeUninit::new(value);
                        }
                        // Signal to the consumer that writing is complete
                        cell.sequence.store(pos + 1, Ordering::Release);
                        return Ok(());
                    }
                    Err(actual) => {
                        // CAS failed, reload pos and try again
                        pos = actual;
                    }
                }
            } else if diff < 0 {
                // The cell has not been read by the consumer yet (queue is full)
                return Err(value);
            } else {
                // The tail has already been advanced by another thread.
                // We must reload the tail to check the next slot.
                pos = self.tail.0.load(Ordering::Relaxed);
            }
        }
    }
}

snippet-4.txt

impl<T> MpscRingBuffer<T> {
    pub fn try_pop(&self) -> Option<T> {
        // Only a single consumer thread calls this. 
        // Thus, we can load head with Relaxed ordering.
        let head = self.head.0.load(Ordering::Relaxed);
        let cell = &self.buffer[head & self.mask];
        let seq = cell.sequence.load(Ordering::Acquire);
        let diff = seq as isize - (head + 1) as isize;

        if diff == 0 {
            // The slot has been written by a producer and is ready to read.
            // Move head forward to claim the slot. We can use Relaxed ordering here 
            // since this thread is the sole writer to the head index.
            self.head.0.store(head + 1, Ordering::Relaxed);

            // Safety: We verified that the slot is written (seq == head + 1).
            // We read the value and update the sequence to signal producers.
            let value = unsafe {
                let ptr = cell.value.get();
                ptr.read().assume_init()
            };

            // Set the sequence to the index it will have on the next wrap-around.
            // Ordering::Release ensures that the read of the data completes 
            // before the slot is marked as free.
            cell.sequence.store(head + self.mask + 1, Ordering::Release);
            Some(value)
        } else {
            // Either the slot is empty (diff < 0) or the writer hasn't finished writing (diff > 0)
            None
        }
    }
}

snippet-5.txt

impl<T> Drop for MpscRingBuffer<T> {
    fn drop(&mut self) {
        let mut head = self.head.0.load(Ordering::Relaxed);
        let tail = self.tail.0.load(Ordering::Relaxed);
        
        while head < tail {
            let cell = &self.buffer[head & self.mask];
            let seq = cell.sequence.load(Ordering::Relaxed);
            
            // If the slot has been written (sequence == head + 1), drop it.
            if seq as isize - (head + 1) as isize == 0 {
                unsafe {
                    let ptr = cell.value.get();
                    ptr.read().assume_init(); // Reading out the value drops it
                }
            }
            head += 1;
        }
    }
}

snippet-6.txt

// Safety: UnsafeCell prevents automatic Sync implementation. We must implement it manually.
// - Sync is safe because access to cell.value is synchronized using cell.sequence.
// - Send is safe because the queue permits transferring ownership of T between threads.
unsafe impl<T: Send> Send for MpscRingBuffer<T> {}
unsafe impl<T: Send> Sync for MpscRingBuffer<T> {}

snippet-7.txt

use std::sync::Arc;
use std::thread;
use std::time::Instant;

fn main() {
    let cap = 65536; // 2^16
    let queue = Arc::new(MpscRingBuffer::new(cap));
    let num_producers = 4;
    let iterations = 1_000_000;

    let mut handles = Vec::new();

    // Spawn Producers
    for thread_id in 0..num_producers {
        let q = Arc::clone(&queue);
        let handle = thread::spawn(move || {
            // Pin producer thread to a specific CPU core to prevent migration
            if let Some(core_ids) = core_affinity::get_core_ids() {
                if let Some(core_id) = core_ids.get(thread_id % core_ids.len()) {
                    core_affinity::set_for_current(*core_id);
                }
            }

            for i in 0..iterations {
                let mut val = i;
                while let Err(returned) = q.try_push(val) {
                    val = returned;
                    std::thread::yield_now();
                }
            }
        });
        handles.push(handle);
    }

    // Spawn Consumer
    let q = Arc::clone(&queue);
    let consumer_handle = thread::spawn(move || {
        // Pin consumer thread to a dedicated core
        if let Some(core_ids) = core_affinity::get_core_ids() {
            if let Some(core_id) = core_ids.get(num_producers % core_ids.len()) {
                core_affinity::set_for_current(*core_id);
            }
        }

        let start = Instant::now();
        let total_expected = num_producers * iterations;
        let mut count = 0;

        while count < total_expected {
            if q.try_pop().is_some() {
                count += 1;
            } else {
                std::thread::yield_now();
            }
        }

        let duration = start.elapsed();
        let rate = (total_expected as f64) / duration.as_secs_f64();
        println!(
            "Processed {} messages in {:?}. Throughput: {:.2} ops/sec",
            total_expected, duration, rate
        );
    });

    for h in handles {
        h.join().unwrap();
    }
    consumer_handle.join().unwrap();
}

snippet-8.txt

// An optimized spinning function using std::hint::spin_loop to minimize CPU power
// consumption and prevent pipeline stalls due to speculative execution.
fn spin_backoff(counter: &mut usize) {
    if *counter < 10 {
        // Issue a PAUSE instruction to the CPU
        std::hint::spin_loop();
    } else if *counter < 20 {
        // Yield execution back to the OS scheduler
        std::thread::yield_now();
    } else {
        // Micro-sleep for 1 microsecond to prevent severe thread starvation
        std::thread::sleep(std::time::Duration::from_micros(1));
    }
    *counter += 1;
}