rstango 0.1.0

//! Our goal here is to create a synchronization queue for 

use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
use std::cell::UnsafeCell;
use std::mem::MaybeUninit;
use std::ptr::NonNull;
use std::pin::Pin;
use std::sync::Arc;

use linux_futex::Futex;


#[derive(Debug)]
pub struct Queue<T> {
    inner: Pin<Arc<QueueInner<T>>>,
}

#[derive(Debug)]
pub struct WorkerQueue<T> {
    /// A reference to the inner state of the queue.
    queue: Pin<Arc<QueueInner<T>>>,
    /// The index of this queue worker.
    index: u8,
    /// A pointer to the worker inner state, this pointer is only available while the 
    /// queue's inner state is hold in this structure. Theoretically, we could just 
    /// resolve this pointer by using queue.workers[index], but this would imply some
    /// inderection.
    inner: NonNull<WorkerQueueInner<T>>,
}

/// SAFETY: Implementing Send because the shared pointer removed it, because this pointer
/// will be valid while the queue's inner state is valid, and it is stored with it.
unsafe impl<T> Send for WorkerQueue<T> { }

#[derive(Debug)]
struct QueueInner<T> {
    /// An array of workers, the fixed size of 32 comes from the size of our mask.
    workers: [WorkerQueueInner<T>; 64],
    /// The number of workers currently initialized, we are not relying on this value to
    /// know the actual number of initialized workers, but rather on the mask below. Thus
    /// this value can be accessed with Relaxed ordering.
    workers_count: AtomicU8,
    /// This mask contains a bit packed value which represent the workers that are 
    /// currently available for work. This mask contains, at bit N, if the N-th worker
    /// is currently available for receiving work.
    workers_mask: AtomicU64,
}

#[derive(Debug)]
#[repr(align(64))]
struct WorkerQueueInner<T> {
    value: UnsafeCell<MaybeUninit<T>>,
    state: Futex<linux_futex::Private>,
}

// SAFETY: The cell is atomically guarded by the state.
unsafe impl<T> Sync for WorkerQueueInner<T> { }

/// State for when the worker has no store value.
const STATE_UNINIT: u32 = 0;
/// State for when the sender has set the value into the worker.
const STATE_SET: u32 = 1;

impl<T> Queue<T> {

    pub fn new() -> Self {
        Self {
            inner: Arc::pin(QueueInner { 
                workers: [ const { WorkerQueueInner { 
                    value: UnsafeCell::new(MaybeUninit::uninit()),
                    state: Futex::new(STATE_UNINIT),
                } }; 64],
                workers_count: AtomicU8::new(0), 
                workers_mask: AtomicU64::new(0),
            }),
        }
    }

    /// Allocate a new worker.
    pub fn worker(&self) -> Option<WorkerQueue<T>> {

        let index = self.inner.workers_count.fetch_add(1, Ordering::Relaxed);
        if index as usize >= self.inner.workers.len() {
            return None;
        }
        
        let inner = &raw const self.inner.workers[index as usize];
        let inner = NonNull::new(inner.cast_mut()).unwrap();

        Some(WorkerQueue {
            queue: self.inner.clone(),
            index,
            inner,
        })

    }

    /// Send a value if possible.
    pub fn send(&self, value: T) -> Result<(), T> {

        // In this loop, we try to acquire a worker index to push our value. By using the
        // trailing_zeros on the mask, we ensure that we try to use the first workers 
        // before using the last ones. This allows us avoid splitting work on many cores
        // (and thus nullifying L1 cache usage, at least).
        let index = loop {
            let mask = self.inner.workers_mask.load(Ordering::Relaxed);
            if mask == 0 {
                return Err(value);
            }
            let index = mask.trailing_zeros();
            // We are using the Acquire on the loading of the value, thus ensuring that
            // if we are successful, the operations below will effectively see a valid
            // and initialized thread.
            if self.inner.workers_mask.compare_exchange_weak(mask, mask & !(1 << index), Ordering::Relaxed, Ordering::Relaxed).is_ok() {
                break index;
            }
        };

        let worker = &self.inner.workers[index as usize];

        // SAFETY: We know that we are the only one to access this worker, and the worker
        // has set its index in the mask only after it is sure that it will not read from
        // this value, not until we change the state to SET.
        unsafe {
            worker.value.get().write(MaybeUninit::new(value));
        }

        // Use the release order, so when the worker will read this SET state it will be
        // sure that the value has been fully flushed.
        worker.state.value.store(STATE_SET, Ordering::Release);
        worker.state.wake(1);

        Ok(())

    }

}

impl<T> Clone for Queue<T> {
    fn clone(&self) -> Self {
        Self { inner: self.inner.clone() }
    }
}

impl<T> WorkerQueue<T> {

    /// Receive a value from this worker.
    pub fn recv(&mut self) -> T {

        // SAFETY: This pointer is valid as long as the internal Arc<QueueInner> is alive,
        // and because it's guaranteed to at least live for the duration of this call, we
        // are ok. We also only get shared references of this worker's inner, never 
        // mutably.
        let inner = unsafe { self.inner.as_ref() };

        // Mark this worker as available for work.
        // Note that the state should've been reset to UNINIT by the last recv call.
        // We are using the Release ordering, so we are sure that the thread as been set
        // in the shared place!
        self.queue.workers_mask.fetch_or(1 << self.index, Ordering::Relaxed);

        let mut cycles = 0u32;

        loop {
            
            let state = inner.state.value.load(Ordering::Acquire);
            if state == STATE_SET {
                break;
            }

            cycles += 1;
            if cycles <= 1000 {
                for _ in 0..8 {
                    std::hint::spin_loop();
                }
            } else if inner.state.wait(state).is_err() {
                // If the value is wrong, it could only been changed to SET, so we break!
                break;
            }

        }

        // SAFETY: We are absolutely sure that we are now the only thread to access it,
        // because the sender that sent the value is no longer using it.
        unsafe {
            let val = inner.value.get().replace(MaybeUninit::uninit()).assume_init();
            inner.state.value.store(STATE_UNINIT, Ordering::Release);
            val
        }

    }

}