Crossbeam的无锁并发Channel解析

/// Unbounded channel implemented as a linked list.
///
/// Each message sent into the channel is assigned a sequence number, i.e. an index. Indices are
/// represented as numbers of type `usize` and wrap on overflow.
///
/// Consecutive messages are grouped into blocks in order to put less pressure on the allocator and
/// improve cache efficiency.
pub struct Channel<T> {
    /// The head of the channel.
    head: CachePadded<Position<T>>,

    /// The tail of the channel.
    tail: CachePadded<Position<T>>,

    /// Receivers waiting while the channel is empty and not disconnected.
    receivers: SyncWaker,

    /// Indicates that dropping a `Channel<T>` may drop messages of type `T`.
    _marker: PhantomData<T>,
}

/// A position in a channel.
#[derive(Debug)]
struct Position<T> {
    /// The index in the channel.
    index: AtomicUsize,

    /// The block in the linked list.
    block: AtomicPtr<Block<T>>,
}

// Each block covers one "lap" of indices.
const LAP: usize = 32;
// The maximum number of messages a block can hold.
const BLOCK_CAP: usize = LAP - 1;

// ..

/// A block in a linked list.
///
/// Each block in the list can hold up to `BLOCK_CAP` messages.
struct Block<T> {
    /// The next block in the linked list.
    next: AtomicPtr<Block<T>>,

    /// Slots for messages.
    slots: [Slot<T>; BLOCK_CAP],
}

/// A slot in a block.
struct Slot<T> {
    /// The message.
    msg: UnsafeCell<MaybeUninit<T>>,

    /// The state of the slot.
    state: AtomicUsize,
}

/// Sends a message into the channel.
pub fn send(&self, msg: T, _deadline: Option<Instant>) -> Result<(), SendTimeoutError<T>> {
	let token = &mut Token::default();
	assert!(self.start_send(token));
	unsafe {
		self.write(token, msg)
			.map_err(SendTimeoutError::Disconnected)
	}
}

/// The token type for the list flavor.
#[derive(Debug)]
pub struct ListToken {
    /// The block of slots.
    block: *const u8,

    /// The offset into the block.
    offset: usize,
}

// Calculate the offset of the index into the block.
let offset = (tail >> SHIFT) % LAP;

// If we reached the end of the block, wait until the next one is installed.
if offset == BLOCK_CAP {
	backoff.snooze();
	tail = self.tail.index.load(Ordering::Acquire);
	block = self.tail.block.load(Ordering::Acquire);
	continue;
}

#[inline]
pub fn snooze(&self) {
	if self.step.get() <= SPIN_LIMIT {
		for _ in 0..1 << self.step.get() {
			atomic::spin_loop_hint();
		}
	} else {
		#[cfg(not(feature = "std"))]
		for _ in 0..1 << self.step.get() {
			atomic::spin_loop_hint();
		}

		#[cfg(feature = "std")]
		::std::thread::yield_now();
	}

	if self.step.get() <= YIELD_LIMIT {
		self.step.set(self.step.get() + 1);
	}
}

// If we're going to have to install the next block, allocate it in advance in order to
// make the wait for other threads as short as possible.
if offset + 1 == BLOCK_CAP && next_block.is_none() {
	next_block = Some(Box::new(Block::<T>::new()));
}

// If this is the first message to be sent into the channel, we need to allocate the
// first block and install it.
if block.is_null() {
	let new = Box::into_raw(Box::new(Block::<T>::new()));

	if self
		.tail
		.block
		.compare_and_swap(block, new, Ordering::Release)
		== block
	{
		self.head.block.store(new, Ordering::Release);
		block = new;
	} else {
		next_block = unsafe { Some(Box::from_raw(new)) };
		tail = self.tail.index.load(Ordering::Acquire);
		block = self.tail.block.load(Ordering::Acquire);
		continue;
	}
}

/// Creates a new unbounded channel.
pub fn new() -> Self {
	Channel {
		head: CachePadded::new(Position {
			block: AtomicPtr::new(ptr::null_mut()),
			index: AtomicUsize::new(0),
		}),
		tail: CachePadded::new(Position {
			block: AtomicPtr::new(ptr::null_mut()),
			index: AtomicUsize::new(0),
		}),
		receivers: SyncWaker::new(),
		_marker: PhantomData,
	}
}

let new_tail = tail + (1 << SHIFT);

// Try advancing the tail forward.
match self.tail.index.compare_exchange_weak(
	tail,
	new_tail,
	Ordering::SeqCst,
	Ordering::Acquire,
) {
	Ok(_) => unsafe {
		// If we've reached the end of the block, install the next one.
		if offset + 1 == BLOCK_CAP {
			let next_block = Box::into_raw(next_block.unwrap());
			self.tail.block.store(next_block, Ordering::Release);
			self.tail.index.fetch_add(1 << SHIFT, Ordering::Release);
			(*block).next.store(next_block, Ordering::Release);
		}

		token.list.block = block as *const u8;
		token.list.offset = offset;
		return true;
	},
	Err(t) => {
		tail = t;
		block = self.tail.block.load(Ordering::Acquire);
		backoff.spin();
	}
}

/// Writes a message into the channel.
pub unsafe fn write(&self, token: &mut Token, msg: T) -> Result<(), T> {
	// If there is no slot, the channel is disconnected.
	if token.list.block.is_null() {
		return Err(msg);
	}

	// Write the message into the slot.
	let block = token.list.block as *mut Block<T>;
	let offset = token.list.offset;
	let slot = (*block).slots.get_unchecked(offset);
	slot.msg.get().write(MaybeUninit::new(msg));
	slot.state.fetch_or(WRITE, Ordering::Release);

	// Wake a sleeping receiver.
	self.receivers.notify();
	Ok(())
}

/// Receives a message from the channel.
pub fn recv(&self, deadline: Option<Instant>) -> Result<T, RecvTimeoutError> {
	let token = &mut Token::default();
	loop {
		// Try receiving a message several times.
		let backoff = Backoff::new();
		loop {
			if self.start_recv(token) {
				unsafe {
					return self.read(token).map_err(|_| RecvTimeoutError::Disconnected);
				}
			}

			if backoff.is_completed() {
				break;
			} else {
				backoff.snooze();
			}
		}

		if let Some(d) = deadline {
			if Instant::now() >= d {
				return Err(RecvTimeoutError::Timeout);
			}
		}

		// Prepare for blocking until a sender wakes us up.
		Context::with(|cx| {
			let oper = Operation::hook(token);
			self.receivers.register(oper, cx);

			// Has the channel become ready just now?
			if !self.is_empty() || self.is_disconnected() {
				let _ = cx.try_select(Selected::Aborted);
			}

			// Block the current thread.
			let sel = cx.wait_until(deadline);

			match sel {
				Selected::Waiting => unreachable!(),
				Selected::Aborted | Selected::Disconnected => {
					self.receivers.unregister(oper).unwrap();
					// If the channel was disconnected, we still have to check for remaining
					// messages.
				}
				Selected::Operation(_) => {}
			}
		});
	}
}

if new_head & MARK_BIT == 0 {
	atomic::fence(Ordering::SeqCst);
	let tail = self.tail.index.load(Ordering::Relaxed);

	// If the tail equals the head, that means the channel is empty.
	if head >> SHIFT == tail >> SHIFT {
		// If the channel is disconnected...
		if tail & MARK_BIT != 0 {
			// ...then receive an error.
			token.list.block = ptr::null();
			return true;
		} else {
			// Otherwise, the receive operation is not ready.
			return false;
		}
	}

	// If head and tail are not in the same block, set `MARK_BIT` in head.
	if (head >> SHIFT) / LAP != (tail >> SHIFT) / LAP {
		new_head |= MARK_BIT;
	}
}

// The block can be null here only if the first message is being sent into the channel.
// In that case, just wait until it gets initialized.
if block.is_null() {
	backoff.snooze();
	head = self.head.index.load(Ordering::Acquire);
	block = self.head.block.load(Ordering::Acquire);
	continue;
}

// Try moving the head index forward.
match self.head.index.compare_exchange_weak(
	head,
	new_head,
	Ordering::SeqCst,
	Ordering::Acquire,
) {
	Ok(_) => unsafe {
		// If we've reached the end of the block, move to the next one.
		if offset + 1 == BLOCK_CAP {
			let next = (*block).wait_next();
			let mut next_index = (new_head & !MARK_BIT).wrapping_add(1 << SHIFT);
			if !(*next).next.load(Ordering::Relaxed).is_null() {
				next_index |= MARK_BIT;
			}

			self.head.block.store(next, Ordering::Release);
			self.head.index.store(next_index, Ordering::Release);
		}

		token.list.block = block as *const u8;
		token.list.offset = offset;
		return true;
	},
	Err(h) => {
		head = h;
		block = self.head.block.load(Ordering::Acquire);
		backoff.spin();
	}
}

/// Reads a message from the channel.
pub unsafe fn read(&self, token: &mut Token) -> Result<T, ()> {
	if token.list.block.is_null() {
		// The channel is disconnected.
		return Err(());
	}

	// Read the message.
	let block = token.list.block as *mut Block<T>;
	let offset = token.list.offset;
	let slot = (*block).slots.get_unchecked(offset);
	slot.wait_write();
	let msg = slot.msg.get().read().assume_init();

	// Destroy the block if we've reached the end, or if another thread wanted to destroy but
	// couldn't because we were busy reading from the slot.
	if offset + 1 == BLOCK_CAP {
		Block::destroy(block, 0);
	} else if slot.state.fetch_or(READ, Ordering::AcqRel) & DESTROY != 0 {
		Block::destroy(block, offset + 1);
	}

	Ok(msg)
}

/// Sets the `DESTROY` bit in slots starting from `start` and destroys the block.
unsafe fn destroy(this: *mut Block<T>, start: usize) {
	// It is not necessary to set the `DESTROY` bit in the last slot because that slot has
	// begun destruction of the block.
	for i in start..BLOCK_CAP - 1 {
		let slot = (*this).slots.get_unchecked(i);

		// Mark the `DESTROY` bit if a thread is still using the slot.
		if slot.state.load(Ordering::Acquire) & READ == 0
			&& slot.state.fetch_or(DESTROY, Ordering::AcqRel) & READ == 0
		{
			// If a thread is still using the slot, it will continue destruction of the block.
			return;
		}
	}

	// No thread is using the block, now it is safe to destroy it.
	drop(Box::from_raw(this));
}

/// Reference counter internals.
struct Counter<C> {
    /// The number of senders associated with the channel.
    senders: AtomicUsize,

    /// The number of receivers associated with the channel.
    receivers: AtomicUsize,

    /// Set to `true` if the last sender or the last receiver reference deallocates the channel.
    destroy: AtomicBool,

    /// The internal channel.
    chan: C,
}

/// Wraps a channel into the reference counter.
pub fn new<C>(chan: C) -> (Sender<C>, Receiver<C>) {
    let counter = Box::into_raw(Box::new(Counter {
        senders: AtomicUsize::new(1),
        receivers: AtomicUsize::new(1),
        destroy: AtomicBool::new(false),
        chan,
    }));
    let s = Sender { counter };
    let r = Receiver { counter };
    (s, r)
}

impl<C> Sender<C> {
    /// Acquires another sender reference.
    pub fn acquire(&self) -> Sender<C> {
        let count = self.counter().senders.fetch_add(1, Ordering::Relaxed);

        // Cloning senders and calling `mem::forget` on the clones could potentially overflow the
        // counter. It's very difficult to recover sensibly from such degenerate scenarios so we
        // just abort when the count becomes very large.
        if count > isize::MAX as usize {
            process::abort();
        }

        Sender {
            counter: self.counter,
        }
    }
}

impl<C> Receiver<C> {
    /// Acquires another receiver reference.
    pub fn acquire(&self) -> Receiver<C> {
        let count = self.counter().receivers.fetch_add(1, Ordering::Relaxed);

        // Cloning receivers and calling `mem::forget` on the clones could potentially overflow the
        // counter. It's very difficult to recover sensibly from such degenerate scenarios so we
        // just abort when the count becomes very large.
        if count > isize::MAX as usize {
            process::abort();
        }

        Receiver {
            counter: self.counter,
        }
    }
}

impl<C> Sender<C> {
    /// Releases the sender reference.
    ///
    /// Function `disconnect` will be called if this is the last sender reference.
    pub unsafe fn release<F: FnOnce(&C) -> bool>(&self, disconnect: F) {
        if self.counter().senders.fetch_sub(1, Ordering::AcqRel) == 1 {
            disconnect(&self.counter().chan);

            if self.counter().destroy.swap(true, Ordering::AcqRel) {
                drop(Box::from_raw(self.counter));
            }
        }
    }
}

impl<C> Receiver<C> {
    /// Releases the receiver reference.
    ///
    /// Function `disconnect` will be called if this is the last receiver reference.
    pub unsafe fn release<F: FnOnce(&C) -> bool>(&self, disconnect: F) {
        if self.counter().receivers.fetch_sub(1, Ordering::AcqRel) == 1 {
            disconnect(&self.counter().chan);

            if self.counter().destroy.swap(true, Ordering::AcqRel) {
                drop(Box::from_raw(self.counter));
            }
        }
    }
}

impl<T> Drop for Channel<T> {
    fn drop(&mut self) {
        let mut head = self.head.index.load(Ordering::Relaxed);
        let mut tail = self.tail.index.load(Ordering::Relaxed);
        let mut block = self.head.block.load(Ordering::Relaxed);

        // Erase the lower bits.
        head &= !((1 << SHIFT) - 1);
        tail &= !((1 << SHIFT) - 1);

        unsafe {
            // Drop all messages between head and tail and deallocate the heap-allocated blocks.
            while head != tail {
                let offset = (head >> SHIFT) % LAP;

                if offset < BLOCK_CAP {
                    // Drop the message in the slot.
                    let slot = (*block).slots.get_unchecked(offset);
                    let p = &mut *slot.msg.get();
                    p.as_mut_ptr().drop_in_place();
                } else {
                    // Deallocate the block and move to the next one.
                    let next = (*block).next.load(Ordering::Relaxed);
                    drop(Box::from_raw(block));
                    block = next;
                }

                head = head.wrapping_add(1 << SHIFT);
            }

            // Deallocate the last remaining block.
            if !block.is_null() {
                drop(Box::from_raw(block));
            }
        }
    }
}