batch_semaphore

pub(crate) struct Semaphore {
    waiters: Mutex<Waitlist>,
    /// The current number of available permits in the semaphore.
    permits: AtomicUsize,
    // ...
}

struct Waitlist {
    queue: LinkedList<Waiter>,
    closed: bool,
}

pub(crate) const MAX_PERMITS: usize = usize::MAX >> 3;
const CLOSED: usize = 1;
// The least-significant bit in the number of permits is reserved to use
// as a flag indicating that the semaphore has been closed. Consequently
// PERMIT_SHIFT is used to leave that bit for that purpose.
const PERMIT_SHIFT: usize = 1;

pub(crate) struct Acquire<'a> {
    node: Waiter,
    semaphore: &'a Semaphore,
    num_permits: usize,
    queued: bool,
}

struct Waiter {
    state: AtomicUsize,
    waker: UnsafeCell<Option<Waker>>,
    pointers: linked_list::Pointers<Waiter>,
    // ...
    _p: PhantomPinned,
}

let mut curr = self.permits.load(Acquire);
let mut waiters = loop {
    // Has the semaphore closed?
    if curr & Self::CLOSED > 0 {
        return Poll::Ready(Err(AcquireError::closed()));
    }

    let mut remaining = 0;
    let total = curr
        .checked_add(acquired)
        .expect("number of permits must not overflow");
    let (next, acq) = if total >= needed {
        let next = curr - (needed - acquired);
        (next, needed >> Self::PERMIT_SHIFT)
    } else {
        remaining = (needed - acquired) - curr;
        (0, curr >> Self::PERMIT_SHIFT)
    };

    if remaining > 0 && lock.is_none() {
        // ...
        lock = Some(self.waiters.lock());
    }

    match self.permits.compare_exchange(curr, next, AcqRel, Acquire) {
        Ok(_) => {
            acquired += acq;
            if remaining == 0 {
                if !*queued {
                    // ...
                    node.state.store(0, Release);
                    *queued = true;
                    // ...
                    return Poll::Ready(Ok(()));
                } else if lock.is_none() {
                    break self.waiters.lock();
                }
            }
            break lock.expect("lock must be acquired before waiting");
        }
        Err(actual) => curr = actual,
    }
};

let (should_remove, _assigned) = node.assign_permits(&mut acquired);
// ...
if should_remove {
    self.add_permits_locked(acquired, waiters);
    return Poll::Ready(Ok(()));
}
// ...
// If the waiter is not already in the wait queue, enqueue it.
if !was_queued {
    // ...
    waiters.queue.push_front(node);
}

while rem > 0 {
    let mut waiters = lock.take().unwrap_or_else(|| self.waiters.lock());
    'inner: while wakers.can_push() {
        // Was the waiter assigned enough permits to wake it?
        let _assigned = match waiters.queue.last() {
            Some(waiter) => {
                let (should_remove, assigned) = waiter.assign_permits(&mut rem);
                if !should_remove {
                    // ...
                    break 'inner;
                }
                assigned
            }
            None => {
                is_empty = true;
                // If we assigned permits to all the waiters in the queue, and there are
                // still permits left over, assign them back to the semaphore.
                break 'inner;
            }
        };
        let mut waiter = waiters.queue.pop_back().unwrap();
        if let Some(waker) =
            unsafe { waiter.as_mut().waker.with_mut(|waker| (*waker).take()) }
        {
            wakers.push(waker);
        }
        // ...
    }

    if rem > 0 && is_empty {
        // ...
        let prev = self.permits.fetch_add(rem << Self::PERMIT_SHIFT, Release);
        // ...
        rem = 0;
    }

    drop(waiters); // release the lock

    wakers.wake_all();
}

use std::sync::Arc;
use tokio::sync::Semaphore;
use tokio::task::yield_now;

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let sem = Arc::new(Semaphore::new(3));
  // main holds all 3 permits
  let held = sem.clone().acquire_many_owned(3).await.unwrap();

  let mut handles = Vec::new();
  // A asks for 3 first, then B, C, D ask for 1 each
  for (name, n) in [("A", 3u32), ("B", 1), ("C", 1), ("D", 1)] {
    let sem = sem.clone();
    handles.push(tokio::spawn(async move {
      let p = sem.acquire_many(n).await.unwrap();
      println!("{name} got {n}");
      yield_now().await;
      drop(p);
    }));
    yield_now().await; // let it enqueue before the next one
  }

  let mut held = held;
  for i in 1..=3 {
    // give back one permit at a time
    let one = held.split(1).unwrap();
    drop(one);
    println!(
      "released {i}: available={} try_acquire(1)={:?}",
      sem.available_permits(),
      sem.try_acquire().map(|_| ()).map_err(|e| e.to_string())
    );
    yield_now().await;
  }
  for h in handles {
    h.await.unwrap();
  }
  println!("end: available={}", sem.available_permits());
}

released 1: available=0 try_acquire(1)=Err("no permits available")
released 2: available=0 try_acquire(1)=Err("no permits available")
released 3: available=0 try_acquire(1)=Err("no permits available")
A got 3
B got 1
C got 1
D got 1
end: available=3

impl Drop for Acquire<'_> {
    fn drop(&mut self) {
        // If the future is completed, there is no node in the wait list, so we
        // can skip acquiring the lock.
        if !self.queued {
            return;
        }

        // This is where we ensure safety. The future is being dropped,
        // which means we must ensure that the waiter entry is no longer stored
        // in the linked list.
        let mut waiters = self.semaphore.waiters.lock();

        // remove the entry from the list
        let node = NonNull::from(&mut self.node);
        // Safety: we have locked the wait list.
        unsafe { waiters.queue.remove(node) };

        let acquired_permits = self.num_permits - self.node.state.load(Acquire);
        if acquired_permits > 0 {
            self.semaphore.add_permits_locked(acquired_permits, waiters);
        }
    }
}

let result = match semaphore.poll_acquire(cx, needed, node, queued) {
    Poll::Pending => Poll::Pending,
    Poll::Ready(r) => {
        r?;
        // ...
        let coop = ready!(crate::task::coop::poll_proceed(cx));

        coop.made_progress();
        Poll::Ready(Ok(()))
    }
};

use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::task::yield_now;

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let m = Arc::new(Mutex::new(Vec::new()));
  let guard = m.lock().await;

  let mut handles = Vec::new();
  for id in [3, 1, 4, 0, 2] {
    let m = m.clone();
    handles.push(tokio::spawn(async move {
      m.lock().await.push(id);
    }));
    yield_now().await; // id is queued before the next task is spawned
  }

  drop(guard);
  // the permit was handed to the first waiter, not put back
  println!("after unlock: try_lock ok = {}", m.try_lock().is_ok());

  for h in handles {
    h.await.unwrap();
  }
  println!("lock order: {:?}", *m.lock().await);
}

after unlock: try_lock ok = false
lock order: [3, 1, 4, 0, 2]

#[cfg(not(loom))]
const MAX_READS: u32 = u32::MAX >> 3;

use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::task::yield_now;

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let l = Arc::new(RwLock::new(0));
  let r1 = l.read().await;

  let lw = l.clone();
  let w = tokio::spawn(async move {
    *lw.write().await += 1;
    println!("writer done");
  });
  yield_now().await; // writer is now queued

  println!("try_read while writer waits: ok = {}", l.try_read().is_ok());
  let lr = l.clone();
  let r2 = tokio::spawn(async move {
    let v = *lr.read().await;
    println!("reader 2 sees {v}");
  });
  yield_now().await;

  drop(r1);
  println!("dropped reader 1");
  w.await.unwrap();
  r2.await.unwrap();
}

try_read while writer waits: ok = false
dropped reader 1
writer done
reader 2 sees 1

1: available = 0
B got 1
2: A got its 3 = false
3: available = 0
4: available = 1
5: available = 2