mpsc

pub(crate) struct Block<T> {
    /// The header fields.
    header: BlockHeader<T>,

    /// Array containing values pushed into the block. Values are stored in a
    /// continuous array in order to improve cache line behavior when reading.
    /// The values must be manually dropped.
    values: Values<T>,
}

struct BlockHeader<T> {
    /// The start index of this block.
    ///
    /// Slots in this block have indices in `start_index .. start_index + BLOCK_CAP`.
    start_index: usize,

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

    /// Bitfield tracking slots that are ready to have their values consumed.
    ready_slots: AtomicUsize,

    /// The observed `tail_position` value *after* the block has been passed by
    /// `block_tail`.
    observed_tail_position: UnsafeCell<usize>,
}

/// Flag tracking that a block has gone through the sender's release routine.
///
/// When this is set, the receiver may consider freeing the block.
const RELEASED: usize = 1 << BLOCK_CAP;

/// Flag tracking all senders dropped.
///
/// When this flag is set, the send half of the channel has closed.
const TX_CLOSED: usize = RELEASED << 1;

/// Mask covering all bits used to track slot readiness.
const READY_MASK: usize = RELEASED - 1;

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,
}

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);
    }
}

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);
    }
}

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>>,
}

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);
// ...
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 {
    // ...
    let _ = unsafe { Box::from_raw(block.as_ptr()) };
}

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use tokio::sync::mpsc;

struct Counting;
static BIG: AtomicUsize = AtomicUsize::new(0);

// count only allocations big enough to be a Block<[u8; 64]> (32 * 64 bytes + header)
unsafe impl GlobalAlloc for Counting {
  unsafe fn alloc(&self, l: Layout) -> *mut u8 {
    if l.size() >= 32 * 64 {
      BIG.fetch_add(1, Relaxed);
    }
    unsafe { System.alloc(l) }
  }
  unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
    unsafe { System.dealloc(p, l) }
  }
}

#[global_allocator]
static A: Counting = Counting;

fn main() {
  println!("size_of Block for [u8;64]: ~{} bytes", 32 * 64);
  let (tx, mut rx) = mpsc::unbounded_channel::<[u8; 64]>();
  let base = BIG.load(Relaxed);

  // 1. ping-pong: send one, receive one
  for _ in 0..10_000 {
    tx.send([0; 64]).unwrap();
    rx.try_recv().unwrap();
  }
  println!("ping-pong 10000 msgs: blocks allocated = {}", BIG.load(Relaxed) - base);

  // 2. burst: 1000 queued, then drained; three rounds
  for round in 1..=3 {
    let before = BIG.load(Relaxed);
    for _ in 0..1000 {
      tx.send([0; 64]).unwrap();
    }
    while rx.try_recv().is_ok() {}
    println!("burst round {round} (1000 msgs): blocks allocated = {}", BIG.load(Relaxed) - before);
  }
}

size_of Block for [u8;64]: ~2048 bytes
ping-pong 10000 msgs: blocks allocated = 1
burst round 1 (1000 msgs): blocks allocated = 30
burst round 2 (1000 msgs): blocks allocated = 28
burst round 3 (1000 msgs): blocks allocated = 29

pub(super) struct Chan<T, S> {
    /// Handle to the push half of the lock-free list.
    tx: CachePadded<list::Tx<T>>,

    /// Receiver waker. Notified when a value is pushed into the channel.
    rx_waker: CachePadded<AtomicWaker>,

    /// Notifies all tasks listening for the receiver being dropped.
    notify_rx_closed: Notify,

    /// Coordinates access to channel's capacity.
    semaphore: S,

    /// Tracks the number of outstanding sender handles.
    ///
    /// When this drops to zero, the send half of the channel is closed.
    tx_count: AtomicUsize,

    /// Tracks the number of outstanding weak sender handles.
    tx_weak_count: AtomicUsize,

    /// Only accessed by `Rx` handle.
    rx_fields: UnsafeCell<RxFields<T>>,
}

pub(crate) trait Semaphore {
    fn is_idle(&self) -> bool;

    fn add_permit(&self);

    fn add_permits(&self, n: usize);

    fn close(&self);

    fn is_closed(&self) -> bool;
}

use std::time::Duration;
use tokio::sync::mpsc;
use tokio::sync::mpsc::error::TrySendError;
use tokio::task::yield_now;

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let (tx, mut rx) = mpsc::channel::<String>(2);
  tx.send("m0".into()).await.unwrap();
  tx.send("m1".into()).await.unwrap();
  println!("capacity after 2 sends: {}", tx.capacity());
  match tx.try_send("x".into()) {
    Err(TrySendError::Full(v)) => println!("try_send -> Full({v})"),
    other => println!("try_send -> {other:?}"),
  }

  let mut hs = Vec::new();
  for name in ["s1", "s2", "s3"] {
    let tx = tx.clone();
    hs.push(tokio::spawn(async move {
      println!("{name}: send start");
      tx.send(name.to_string()).await.unwrap();
      println!("{name}: send done");
    }));
    yield_now().await;
  }

  for _ in 0..5 {
    let v = rx.recv().await.unwrap();
    println!("recv {v}");
    yield_now().await;
  }
  for h in hs {
    h.await.unwrap();
  }

  drop(rx);
  match tx.try_send("y".into()) {
    Err(TrySendError::Closed(v)) => println!("after drop(rx): try_send -> Closed({v})"),
    other => println!("after drop(rx): try_send -> {other:?}"),
  }
  let r = tokio::time::timeout(Duration::from_millis(10), tx.send("z".into())).await;
  println!("after drop(rx): send -> {:?}", r.map(|x| x.map_err(|e| e.0)));
}

capacity after 2 sends: 0
try_send -> Full(x)
s1: send start
s2: send start
s3: send start
recv m0
s1: send done
recv m1
s2: send done
recv s1
s3: send done
recv s2
recv s3
after drop(rx): try_send -> Closed(y)
after drop(rx): send -> Ok(Err("z"))

pub(crate) fn recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
    use super::block::Read;

    ready!(crate::trace::trace_leaf(cx));

    // Keep track of task budget
    let coop = ready!(crate::task::coop::poll_proceed(cx));

    self.inner.rx_fields.with_mut(|rx_fields_ptr| {
        let rx_fields = unsafe { &mut *rx_fields_ptr };

        macro_rules! try_recv {
            () => {
                match rx_fields.list.pop(&self.inner.tx) {
                    Some(Read::Value(value)) => {
                        self.inner.semaphore.add_permit();
                        coop.made_progress();
                        return Ready(Some(value));
                    }
                    Some(Read::Closed) => {
                        // ...
                        debug_assert!(self.inner.semaphore.is_idle());
                        coop.made_progress();
                        return Ready(None);
                    }
                    None => {} // fall through
                }
            };
        }

        try_recv!();

        self.inner.rx_waker.register_by_ref(cx.waker());

        // It is possible that a value was pushed between attempting to read
        // and registering the task, so we have to check the channel a
        // second time here.
        try_recv!();

        if rx_fields.rx_closed && self.inner.semaphore.is_idle() {
            coop.made_progress();
            Ready(None)
        } else {
            Pending
        }
    })
}

A: Err("Full(..)")
B: Ok(())
C: capacity = 0
D: Some(2) Some(3)
E: capacity = 2