tokio/sync/mpsc/

list.rs

//! A concurrent, lock-free, FIFO list.

use crate::loom::sync::atomic::{AtomicPtr, AtomicUsize};
use crate::loom::thread;
use crate::sync::mpsc::block::{self, Block};

use std::fmt;
use std::ptr::NonNull;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release};

/// List queue transmit handle.
pub(crate) struct Tx<T> {
    /// Tail in the `Block` mpmc list.
    block_tail: AtomicPtr<Block<T>>,

    /// Position to push the next message. This references a block and offset
    /// into the block.
    tail_position: AtomicUsize,
}

/// List queue receive handle
pub(crate) struct Rx<T> {
    /// Pointer to the block being processed.
    head: NonNull<Block<T>>,

    /// Next slot index to process.
    index: usize,

    /// Pointer to the next block pending release.
    free_head: NonNull<Block<T>>,
}

/// Return value of `Rx::try_pop`.
pub(crate) enum TryPopResult<T> {
    /// Successfully popped a value.
    Ok(T),
    /// The channel is empty.
    ///
    /// Note that `list.rs` only tracks the close state set by senders. If the
    /// channel is closed by `Rx::close()`, then `TryPopResult::Empty` is still
    /// returned, and the close state needs to be handled by `chan.rs`.
    Empty,
    /// The channel is empty and closed.
    ///
    /// Returned when the send half is closed (all senders dropped).
    Closed,
    /// The channel is not empty, but the first value is being written.
    Busy,
}

pub(crate) fn channel<T>() -> (Tx<T>, Rx<T>) {
    channel_from_index(0)
}

pub(crate) fn channel_from_index<T>(start_index: usize) -> (Tx<T>, Rx<T>) {
    debug_assert_eq!(block::offset(start_index), 0);

    // Create the initial block shared between the tx and rx halves.
    let initial_block = Block::new(start_index);
    let initial_block_ptr = Box::into_raw(initial_block);

    let tx = Tx {
        block_tail: AtomicPtr::new(initial_block_ptr),
        tail_position: AtomicUsize::new(start_index),
    };

    let head = NonNull::new(initial_block_ptr).unwrap();

    let rx = Rx {
        head,
        index: start_index,
        free_head: head,
    };

    (tx, rx)
}

impl<T> Tx<T> {
    /// Pushes a value into the list.
    pub(crate) fn push(&self, value: T) {
        // First, claim a slot for the value. `Acquire` is used here to
        // synchronize with the `fetch_add` in `reclaim_blocks`.
        let slot_index = self.tail_position.fetch_add(1, Acquire);

        // Load the current block and write the value
        let block = self.find_block(slot_index);

        unsafe {
            // Write the value to the block
            block.as_ref().write(slot_index, value);
        }
    }

    /// Closes the send half of the list.
    ///
    /// Similar process as pushing a value, but instead of writing the value &
    /// setting the ready flag, the `TX_CLOSED` flag is set on the block.
    pub(crate) fn close(&self) {
        // First, claim a slot for the value. This is the last slot that will be
        // claimed.
        let slot_index = self.tail_position.fetch_add(1, Acquire);

        let block = self.find_block(slot_index);

        unsafe { block.as_ref().tx_close() }
    }

    fn find_block(&self, slot_index: usize) -> NonNull<Block<T>> {
        // The start index of the block that contains `index`.
        let start_index = block::start_index(slot_index);

        // The index offset into the block
        let offset = block::offset(slot_index);

        // Load the current head of the block
        let mut block_ptr = self.block_tail.load(Acquire);

        let block = unsafe { &*block_ptr };

        // Calculate the distance between the tail ptr and the target block
        let distance = block.distance(start_index);

        // Decide if this call to `find_block` should attempt to update the
        // `block_tail` pointer.
        //
        // Updating `block_tail` is not always performed in order to reduce
        // contention.
        //
        // When set, as the routine walks the linked list, it attempts to update
        // `block_tail`. If the update cannot be performed, `try_updating_tail`
        // is unset.
        let mut try_updating_tail = distance > offset;

        // Walk the linked list of blocks until the block with `start_index` is
        // found.
        loop {
            let block = unsafe { &(*block_ptr) };

            if block.is_at_index(start_index) {
                return unsafe { NonNull::new_unchecked(block_ptr) };
            }

            let next_block = block
                .load_next(Acquire)
                // There is no allocated next block, grow the linked list.
                .unwrap_or_else(|| block.grow());

            // If the block is **not** final, then the tail pointer cannot be
            // advanced any more.
            try_updating_tail &= block.is_final();

            if try_updating_tail {
                // Advancing `block_tail` must happen when walking the linked
                // list. `block_tail` may not advance passed any blocks that are
                // not "final". At the point a block is finalized, it is unknown
                // if there are any prior blocks that are unfinalized, which
                // makes it impossible to advance `block_tail`.
                //
                // While walking the linked list, `block_tail` can be advanced
                // as long as finalized blocks are traversed.
                //
                // Release ordering is used to ensure that any subsequent reads
                // are able to see the memory pointed to by `block_tail`.
                //
                // Acquire is not needed as any "actual" value is not accessed.
                // At this point, the linked list is walked to acquire blocks.
                if self
                    .block_tail
                    .compare_exchange(block_ptr, next_block.as_ptr(), Release, Relaxed)
                    .is_ok()
                {
                    // Synchronize with any senders
                    let tail_position = self.tail_position.fetch_add(0, Release);

                    unsafe {
                        block.tx_release(tail_position);
                    }
                } else {
                    // A concurrent sender is also working on advancing
                    // `block_tail` and this thread is falling behind.
                    //
                    // Stop trying to advance the tail pointer
                    try_updating_tail = false;
                }
            }

            block_ptr = next_block.as_ptr();

            thread::yield_now();
        }
    }

    /// # Safety
    ///
    /// Behavior is undefined if any of the following conditions are violated:
    ///
    /// - The `block` was created by [`Box::into_raw`].
    /// - The `block` is not currently part of any linked list.
    /// - The `block` is a valid pointer to a [`Block<T>`].
    pub(crate) unsafe fn reclaim_block(&self, mut block: NonNull<Block<T>>) {
        // The block has been removed from the linked list and ownership
        // is reclaimed.
        //
        // Before dropping the block, see if it can be reused by
        // inserting it back at the end of the linked list.
        //
        // First, reset the data
        //
        // Safety: caller guarantees the block is valid and not in any list.
        unsafe {
            block.as_mut().reclaim();
        }

        let mut reused = false;

        // Attempt to insert the block at the end
        //
        // Walk at most three times
        let curr_ptr = self.block_tail.load(Acquire);

        // The pointer can never be null
        debug_assert!(!curr_ptr.is_null());

        // Safety: curr_ptr is never null.
        let mut curr = unsafe { NonNull::new_unchecked(curr_ptr) };

        // TODO: Unify this logic with Block::grow
        for _ in 0..3 {
            match unsafe { curr.as_ref().try_push(&mut block, AcqRel, Acquire) } {
                Ok(()) => {
                    reused = true;
                    break;
                }
                Err(next) => {
                    curr = next;
                }
            }
        }

        if !reused {
            // Safety:
            //
            // 1. Caller guarantees the block is valid and not in any list.
            // 2. The block was created by `Box::into_raw`.
            let _ = unsafe { Box::from_raw(block.as_ptr()) };
        }
    }
}

impl<T> fmt::Debug for Tx<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("Tx")
            .field("block_tail", &self.block_tail.load(Relaxed))
            .field("tail_position", &self.tail_position.load(Relaxed))
            .finish()
    }
}

impl<T> Rx<T> {
    pub(crate) fn is_empty(&self, tx: &Tx<T>) -> bool {
        let block = unsafe { self.head.as_ref() };
        if block.has_value(self.index) {
            return false;
        }

        // It is possible that a block has no value "now" but the list is still not empty.
        // To be sure, it is necessary to check the length of the list.
        self.len(tx) == 0
    }

    // Guaranteed to return true if `slot_index` is the fake message sent on channel close.
    // Guaranteed to return false if `slot_index` is a fully sent message.
    //
    // For messages that are partially sent, may return either true or false.
    fn is_maybe_closed(&self, tx: &Tx<T>, slot_index: usize) -> bool {
        let start_index = block::start_index(slot_index);

        let tail = tx.block_tail.load(Acquire);
        // SAFETY: Only the receiver frees blocks, so since we are the receiver, this will not be
        // freed right now.
        let tail_ref = unsafe { &*tail };
        if tail_ref.is_at_index(start_index) {
            return !tail_ref.has_value(slot_index);
        }

        // This method is optimized for checking whether the last value is present, so most of the
        // time it is in `block_tail`. However, this isn't always the case since it's possible
        // that the list was grown with an empty block, in which case `block_tail` points one block
        // too far. To handle this case, we walk the list from the head.
        let mut block_ptr = Some(self.head);

        while let Some(block) = block_ptr {
            // SAFETY: Only the receiver frees blocks, so since we are the receiver, this will not
            // be freed right now.
            let block_ref = unsafe { block.as_ref() };
            if block_ref.is_at_index(start_index) {
                return !block_ref.has_value(slot_index);
            }
            block_ptr = block_ref.load_next(Acquire);
        }
        true
    }

    pub(crate) fn len(&self, tx: &Tx<T>) -> usize {
        let tail_position = tx.tail_position.load(Acquire);
        let mut len = tail_position.wrapping_sub(self.index);
        debug_assert!(0 <= len as isize);
        if len == 0 {
            return 0;
        }
        // There are messages present in the queue. However, it's possible that the last message is
        // a fake "closed" message that we do not wish to count. To avoid counting it, we do not
        // count the last message if the ready bit is unset.
        //
        // Note that it is also possible for the ready bit to be unset on a normal message, but
        // this happens only if that message is currently being sent *right now* in parallel on
        // another thread. That is okay because it is optional to count messages that are currently
        // being sent.
        if self.is_maybe_closed(tx, tail_position.wrapping_sub(1)) {
            len -= 1;
        }
        len
    }

    /// Pops the next value off the queue.
    pub(crate) fn pop(&mut self, tx: &Tx<T>) -> Option<block::Read<T>> {
        // Advance `head`, if needed
        if !self.try_advancing_head() {
            return None;
        }

        self.reclaim_blocks(tx);

        unsafe {
            let block = self.head.as_ref();

            let ret = block.read(self.index);

            if let Some(block::Read::Value(..)) = ret {
                self.index = self.index.wrapping_add(1);
            }

            ret
        }
    }

    /// Pops the next value off the queue, detecting whether the block
    /// is busy or empty on failure.
    ///
    /// This function exists because `Rx::pop` can return `None` even if the
    /// channel's queue contains a message that has been completely written.
    /// This can happen if the fully delivered message is behind another message
    /// that is in the middle of being written to the block, since the channel
    /// can't return the messages out of order.
    pub(crate) fn try_pop(&mut self, tx: &Tx<T>) -> TryPopResult<T> {
        let tail_position = tx.tail_position.load(Acquire);
        let result = self.pop(tx);

        match result {
            Some(block::Read::Value(t)) => TryPopResult::Ok(t),
            Some(block::Read::Closed) => TryPopResult::Closed,
            None if tail_position == self.index => TryPopResult::Empty,
            None => TryPopResult::Busy,
        }
    }

    /// Tries advancing the block pointer to the block referenced by `self.index`.
    ///
    /// Returns `true` if successful, `false` if there is no next block to load.
    fn try_advancing_head(&mut self) -> bool {
        let block_index = block::start_index(self.index);

        loop {
            let next_block = {
                let block = unsafe { self.head.as_ref() };

                if block.is_at_index(block_index) {
                    return true;
                }

                block.load_next(Acquire)
            };

            let next_block = match next_block {
                Some(next_block) => next_block,
                None => {
                    return false;
                }
            };

            self.head = next_block;

            thread::yield_now();
        }
    }

    fn reclaim_blocks(&mut self, tx: &Tx<T>) {
        while self.free_head != self.head {
            unsafe {
                // Get a handle to the block that will be freed and update
                // `free_head` to point to the next block.
                let block = self.free_head;

                let observed_tail_position = block.as_ref().observed_tail_position();

                let required_index = match observed_tail_position {
                    Some(i) => i,
                    None => return,
                };

                if required_index.wrapping_sub(self.index) as isize > 0 {
                    return;
                }

                // We may read the next pointer with `Relaxed` ordering as it is
                // guaranteed that the `reclaim_blocks` routine trails the `recv`
                // routine. Any memory accessed by `reclaim_blocks` has already
                // been acquired by `recv`.
                let next_block = block.as_ref().load_next(Relaxed);

                // Update the free list head
                self.free_head = next_block.unwrap();

                // Push the emptied block onto the back of the queue, making it
                // available to senders.
                tx.reclaim_block(block);
            }

            thread::yield_now();
        }
    }

    /// Effectively `Drop` all the blocks. Should only be called once, when
    /// the list is dropping.
    pub(super) unsafe fn free_blocks(&mut self) {
        debug_assert_ne!(self.free_head, NonNull::dangling());

        let mut cur = Some(self.free_head);

        #[cfg(debug_assertions)]
        {
            // to trigger the debug assert above so as to catch that we
            // don't call `free_blocks` more than once.
            self.free_head = NonNull::dangling();
            self.head = NonNull::dangling();
        }

        while let Some(block) = cur {
            cur = unsafe { block.as_ref() }.load_next(Relaxed);
            drop(unsafe { Box::from_raw(block.as_ptr()) });
        }
    }
}

impl<T> fmt::Debug for Rx<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("Rx")
            .field("head", &self.head)
            .field("index", &self.index)
            .field("free_head", &self.free_head)
            .finish()
    }
}

#[cfg(all(test, not(loom)))]
mod tests {
    use crate::sync::mpsc::unbounded::unbounded_channel_from_index;
    use crate::sync::mpsc::BLOCK_CAP;

    #[cfg(all(target_family = "wasm", not(target_os = "wasi")))]
    use wasm_bindgen_test::wasm_bindgen_test as test;

    #[test]
    #[cfg(not(target_family = "wasm"))]
    fn wraparound() {
        use super::*;

        let (tx, mut rx) = channel_from_index(0usize.wrapping_sub(2 * BLOCK_CAP));
        let head = rx.free_head;

        for i in 0..BLOCK_CAP {
            tx.push(i);
            assert!(matches!(rx.pop(&tx), Some(block::Read::Value(v)) if v == i));
        }

        // Simulate a slow sender claiming the first slot of the second block
        // while still holding a pointer to the first block.
        let slow_slot = tx.tail_position.fetch_add(1, Acquire);
        let slow_tail = AtomicPtr::new(tx.block_tail.load(Acquire));

        // Fill the rest of the second block (wrapping `tail_position` to 0),
        // then push one item into the third block to retire the first block.
        for i in 1..BLOCK_CAP {
            tx.push(i);
        }
        assert_eq!(rx.len(&tx), BLOCK_CAP);
        tx.push(BLOCK_CAP);

        // Advancing `rx` to the second block must not reclaim the first block yet,
        // because `rx.index` has not reached the wrapped `required_index` (1).
        std::thread::scope(|s| {
            s.spawn(|| unsafe {
                let slow_tail = &*slow_tail.load(Relaxed);
                let slow_block = slow_tail.load_next(Acquire).unwrap();
                assert!(slow_block
                    .as_ref()
                    .is_at_index(block::start_index(slow_slot)));
            });
            assert!(rx.pop(&tx).is_none());
        });
        assert_eq!(rx.free_head, head);

        unsafe {
            let slow_tail = &*slow_tail.load(Relaxed);
            let slow_block = slow_tail.load_next(Acquire).unwrap();
            slow_block.as_ref().write(slow_slot, 0);
        }

        for i in 0..=BLOCK_CAP {
            assert!(matches!(rx.pop(&tx), Some(block::Read::Value(v)) if v == i));
        }
        unsafe { rx.free_blocks() };
    }

    #[test]
    fn wraparound_unbounded() {
        let (tx, mut rx) = unbounded_channel_from_index(0usize.wrapping_sub(2 * BLOCK_CAP));

        for i in 0..4 * BLOCK_CAP {
            tx.send(i).unwrap();
            assert_eq!(rx.len(), i + 1);
            assert!(!rx.is_empty());
        }

        for i in 0..4 * BLOCK_CAP {
            assert_eq!(rx.try_recv().unwrap(), i);
            assert_eq!(rx.len(), 4 * BLOCK_CAP - 1 - i);
        }
        assert!(rx.is_empty());
    }
}