1
\$\begingroup\$

I've created a lock-free lamport queue in rust. My code should be clear enough that you won't need to understand a lamport queue before reading it, but if you'd like a reference, this paper give a good outline of the design.

It is a multi-threaded queue that requires no locking on accesses, although there can only be one reader and one writer. I enforce this constraint by creating wrapper types in LamportQueueReader and LamportQueueWriter, which I hope act as 0 cost abstractions.

I'm confident in the correctness of my code, as long as I'm not invoking undefined behavior (UB), so I'd love feedback on any potential UB that I'm invoking. This is my first rust project, so I'd also like feedback on general design, naming, code structure, etc.

A couple specific points I'd like some feedback on are:

  • Is my use of Deref in src/lamport_queue.rs considered bad practice? Is there a better way to handle the typing there?
  • Does any my queue code invoke UB? Could I eliminate the volatile reads at the starts of my push/pop functions?
  • Is my approach to threading + timing in src/srsw_queue_verifier.rs sane? Am I over complicating things there?

Any general feedback on performance, code style, or design is also welcome!

Here is the github repository, although I've reproduced the code below for convenience. The bulk of it is in src/lamport_queue.rs but I'm also open to feedback on src/srsw_queue_verifier.rs (which is a loose framework for testing my queue implementation) or src/main.rs which compares the performance of my queue to mpsc::sync_channel, a rust-provided structure for cross thread communication.

src/lamport_queue.rs

use std::default::Default;
use std::ops::Deref;
use std::ptr;
use std::sync::{atomic, Arc};

pub struct LamportQueue<D: Default + Clone> {
    elements: Box<[D]>,
    head: usize,
    tail: usize,
    open: bool,
}

// Hoping that the Reader and Writer are effectively `newtype`
// and therefore 0 cost
pub struct LamportQueueReader<D: Default + Clone> {
    q: Arc<LamportQueue<D>>,
}

pub struct LamportQueueWriter<D: Default + Clone> {
    q: Arc<LamportQueue<D>>,
}

// Most of the methods on Reader and Writer are equally valid on either
// and simply read the underlying Queue, so we use Reader/Writer
// as smart pointers and use Deref to make them callable.
impl<D: Default + Clone> LamportQueue<D> {
    // Allocate the queue and return handles to enqueue/dequeue
    // Use owned handles to ensure there is only one reader and one writer
    pub fn new(capacity: usize) -> (LamportQueueReader<D>, LamportQueueWriter<D>) {
        let q = Arc::new(Self {
            elements: vec![D::default(); capacity].into_boxed_slice(),
            head: 0,
            tail: 0,
            open: true,
        });
        (
            LamportQueueReader { q: Arc::clone(&q) },
            LamportQueueWriter { q: Arc::clone(&q) },
        )
    }

    pub fn len(&self) -> usize {
        unsafe { ptr::read_volatile(&self.tail) - ptr::read_volatile(&self.head) }
    }

    pub fn capacity(&self) -> usize {
        // Don't need a volatile read here because it should be constant
        self.elements.len()
    }

    pub fn closed(&self) -> bool {
        unsafe { !ptr::read_volatile(&self.open) }
    }

    // This really should be &mut and should live in Reader or Writer,
    // because you should need to own either the reader or the writer
    // to close the queue.
    // I've decided to put it here instead and leave it as a shared
    // reference (`&`) rather than exclusive (`&mut`) is because it's
    // technically an OK thing to do, and it avoids duplicating the code.
    // Anyways, I don't really like this, but it's fine I guess :/
    pub fn close(&self) {
        unsafe { ptr::write_volatile(ptr::addr_of!(self.open) as *mut bool, false) }
    }
}

impl<D: Default + Clone> Deref for LamportQueueReader<D> {
    type Target = LamportQueue<D>;

    fn deref(&self) -> &Self::Target {
        &self.q
    }
}

impl<D: Default + Clone> Deref for LamportQueueWriter<D> {
    type Target = LamportQueue<D>;

    fn deref(&self) -> &Self::Target {
        &self.q
    }
}

// Actual push/pop implimentations
impl<D: Default + Clone> LamportQueueWriter<D> {
    pub fn push(&mut self, d: D) -> bool {
        // In theory, all reads to `self.tail` can be done without volatile, since
        // we can guarantee that we are the only writer to it. However, we only have
        // an immutable reference to `self.tail` and we use a volatile write to
        // mutate it to ensure the Writer sees our write, so I'm unsure how the
        // compiler would elide accesses across the volatile write.
        // 
        // So, I've left in a single volatile read to avoid the potential for compiler
        // access elision.
        let tail = unsafe { ptr::read_volatile(&self.tail) };

        // Ensure that the list is not full
        // It is not a problem if head returns an artifically low number because
        // the caller can simply retry the operation
        if (tail - unsafe { ptr::read_volatile(&self.head) }) == self.elements.len() {
            return false;
        }
        unsafe { ptr::write_volatile(ptr::addr_of!(self.elements[tail % self.elements.len()]) as *mut D, d) };

        // While rust's volatile guarentees that the compiler won't reorder memory accesses,
        // the fence is still necessary to avoid hardware reordering
        atomic::fence(atomic::Ordering::SeqCst);

        unsafe { ptr::write_volatile(ptr::addr_of!(self.tail) as *mut usize, tail + 1) };

        true
    }
}

impl<D: Default + Clone> LamportQueueReader<D> {
    pub fn pop(&mut self) -> Option<D> {
        // After this read, the value of `tail` will always be >= `self.tail`
        let tail = unsafe { ptr::read_volatile(&self.tail) };

        // In theory, all reads to `self.head` can be done without volatile, since
        // we can guarantee that we are the only writer to it. However, we only have
        // an immutable reference to `self.head` and we use a volatile write to
        // mutate it to ensure the Writer sees our write, so I'm unsure how the
        // compiler would elide accesses across the volatile write.
        // 
        // So, I've left in a single volatile read to avoid the potential for compiler
        // access elision.
        let head = unsafe { ptr::read_volatile(&self.head) };

        // `tail > head` & `self.tail >= tail` => `self.tail > head`
        // => It is safe to read one element off of the head
        if tail - head == 0 {
            return None;
        }
        let idx = head % self.elements.len();
        let e = unsafe { ptr::read_volatile(ptr::addr_of!(self.elements[idx])) };

        // While rust's volatile guarentees that the compiler won't reorder memory accesses,
        // the fence is still necessary to avoid hardware reordering
        atomic::fence(atomic::Ordering::SeqCst);

        unsafe { ptr::write_volatile(ptr::addr_of!(self.head) as *mut usize, head + 1) };

        Some(e)
    }
}

src/srsw_queue_verifier.rs

use std::thread;
use std::thread::JoinHandle;
use std::mem;
use std::time;

// A potential refactor here is to use an Either<A, B> type to store send_list / send_handle
// but since this isn't perf critical and I don't know how I'd name those fields, I'm leaving
// it as is.
pub struct VerificationChecker<T> {
    send_list_size: usize,
    // These lists are options because they can be lent out to the sender/reciever
    // threads, but only once each. Once they've been lent, they're None
    send_list: Option<Box<[T]>>, // A list of random numbers to send
    recieve_list: Option<Vec<T>>, // A location to deposit random numbers once recieved
    sender_handle: Option<JoinHandle<Box<[T]>>>,
    reciever_handle: Option<JoinHandle<Vec<T>>>,
    start_time: Option<time::Instant>,
}

impl<T: 'static + Default + Copy + PartialEq + Send> VerificationChecker<T> {
    pub fn new<F: FnMut(&mut T)>(send_list_size: usize, f: F) -> Self {
        // Populate a boxed slice with random numbers generated by `f`
        let mut v: Box<[T]> = vec![T::default(); send_list_size].into_boxed_slice();
        v.iter_mut().for_each(f);

        Self {
            send_list_size,
            send_list: Some(v),
            recieve_list: Some(Vec::with_capacity(send_list_size)),
            sender_handle: None,
            reciever_handle: None,
            start_time: None,
        }
    }

    // Creates an equivilent but fresh VerificationChecker. Specifically, the list elements
    // that are sent will remain the same and be in the same order.
    //
    // We could share the list itself using Arc, however this could give later tests an advantage
    // over earlier tests by having the caches start hot.
    pub fn clone_send(&self) -> Self {
        Self {
            send_list_size: self.send_list_size,
            send_list: self.send_list.clone(),
            recieve_list: Some(Vec::with_capacity(self.send_list_size)),
            sender_handle: None,
            reciever_handle: None,
            start_time: None,
        }
    }

    // Wait for our test to terminate, conclude timing, print results
    //
    // A possible improvment here would be to seperate output printing/handling
    // from the test ending and timing.
    pub fn verify(&mut self) -> bool {
        // Wait for both threads to terminate
        // We can use unwrap here because if the threads can't be joined that's unrecoverable
        let send_list = self.sender_handle.take().expect("Cannot verify if a sender thread hasn't been started").join().unwrap();
        let recieve_list = self.reciever_handle.take().expect("Cannot verify if a reciever thread hasn't been started").join().unwrap();

        // End the benchmark
        let end_time = self.start_time.expect("Nothing to time, nothing has been started").elapsed();

        // Run correctness checks
        let same_length = send_list.len() == recieve_list.len();
        let mut numbers_match = true;
        for i in 0..self.send_list_size {
            if send_list[i] != recieve_list[i] {
                numbers_match = false;
                break;
            }
        }

        println!("Count is correct? {}", same_length);
        println!("Numbers match up? {}", numbers_match);
        println!("Elapsed time: {}ms", end_time.as_millis());

        same_length && numbers_match
    }

    // Start the enqueue thread (and timing)
    pub fn run_sender<F: 'static + FnOnce(&Box<[T]>) + Send>(&mut self, f: F) {
        // Do a little dance to move ownership of `self.send_list` into the closure
        // and replace it with None
        let mut sl = None;
        mem::swap(&mut self.send_list, &mut sl);
        let sl = sl.expect("Only one sender can be started for each verifier");

        // Start timing
        self.start_time = Some(time::Instant::now());

        // Spawn thread and store handle so we can wait later
        self.sender_handle = Some(thread::spawn(move || {
            f(&sl);
            sl
        }));
    }

    // Start the dequeue thread
    pub fn run_reciever<F: 'static + FnOnce(&mut Vec<T>) + Send>(&mut self, f: F) {
        // Do a little dance to move ownership of `self.recieve_list` into the closure
        // and replace it with None
        let mut rl = None;
        mem::swap(&mut self.recieve_list, &mut rl);
        let mut rl = rl.expect("Only one reciever can be started for each verifier");

        // Spawn thread and store handle so we can wait later
        self.reciever_handle = Some(thread::spawn(move || {
            f(&mut rl);
            rl
        }));
    }
}

src/main.rs

use rand::Rng;
use std::sync::atomic;
use std::sync::mpsc;
use clap::{App, Arg};

mod lamport_queue;
mod srsw_queue_verifier;

use lamport_queue::LamportQueue;
use lamport_queue::{LamportQueueReader, LamportQueueWriter};
use srsw_queue_verifier::VerificationChecker;

fn main() {
    let matches = App::new("tsm's lamport queue")
                      .version("0.0.1")
                      .author("thesecretmaster <thesecretmaster@dvtk.me>")
                      .about("Simple single reader / singler writer lamport queue performance testing (vs mpsc::sync_channel)")
                      .arg(Arg::with_name("queue length")
                           .short("l")
                           .long("queue-length")
                           .value_name("QUEUE LENGTH")
                           .help("Sets the queue length (both lamport queues and mpsc::sync_channel have queue lengths)")
                           .takes_value(true)
                           .required(true)
                           .default_value("3"))
                      .get_matches();

    let queue_len: usize = clap::value_t!(matches.value_of("queue length"), usize).expect("Positive integer is required for queue length");

    println!("Running tests with queue length of {}", queue_len);

    // Create the queue and prepare it for sharing
    let (reciever_handle, sender_handle) = LamportQueue::new(queue_len);
    let (tx, rx) = mpsc::sync_channel(queue_len);

    // Generate a list of random numbers to send down the queue
    let srsw_queue_verifier: VerificationChecker<usize> = VerificationChecker::new(100000, |_| rand::thread_rng().gen());
    let mpsc_verifier = srsw_queue_verifier.clone_send();

    atomic::compiler_fence(atomic::Ordering::SeqCst);

    test_lamport(srsw_queue_verifier, sender_handle, reciever_handle);

    // This probably does nothing, but it's here just in case :P
    atomic::compiler_fence(atomic::Ordering::SeqCst);

    test_mpsc(mpsc_verifier, tx, rx);
}

fn test_mpsc(mut mpsc_verifier: VerificationChecker<usize>, tx: mpsc::SyncSender<usize>, rx: mpsc::Receiver<usize>) {
    mpsc_verifier.run_sender(move |sl| {
        for i in sl.into_iter() {
            // Continually retry until the push is sucessful
            while tx.send(*i).is_err() {}
        }
        // Automatically closes when it's dropped
    });

    mpsc_verifier.run_reciever(move |rl| {
        // Continually pop off the queue until it's closed and empty
        loop {
            match rx.recv() {
                Err(_) => break,
                Ok(i) => rl.push(i),
            }
        }
    });

    println!("MPSC:");
    mpsc_verifier.verify();
}

fn test_lamport(mut srsw_queue_verifier: VerificationChecker<usize>, mut sender_handle: LamportQueueWriter<usize>, mut reciever_handle: LamportQueueReader<usize>) {
    srsw_queue_verifier.run_sender(move |sl| {
        for i in sl.into_iter() {
            // Continually retry until the push is sucessful
            while !sender_handle.push(*i) {}
        }
        // Close the loop to allow the reciever to terminate
        sender_handle.close();
    });

    srsw_queue_verifier.run_reciever(move |rl| {
        // Continually pop off the queue until it's closed and empty
        while !(reciever_handle.closed() && reciever_handle.len() == 0) {
            match reciever_handle.pop() {
                None => (),
                Some(i) => rl.push(i),
            }
        }
    });

    println!("SRSW:");
    srsw_queue_verifier.verify();
}
```
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.