shuttle_engine/future/

batch_semaphore.rs

//! A counting semaphore supporting both async and sync operations.
use crate::runtime::execution::ExecutionState;
use crate::runtime::task::{clock::VectorClock, TaskId};
use crate::runtime::thread;
use crate::sync_types::{ResourceSignature, ResourceType};
use crate::{backtrace_enabled, current};
use std::cell::RefCell;
use std::collections::VecDeque;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::sync::Mutex;
use std::task::{Context, Poll, Waker};
use tracing::trace;

struct Waiter {
    /// The task waiting on this waiter's `Acquire`.
    ///
    /// Refreshed on every poll (like `waker`) rather than frozen at creation
    /// time. An `Acquire` future is not necessarily owned by the task that
    /// created it: it can be cached inside a longer-lived object and later
    /// polled by a different task (tokio's `poll_recv(&mut self, cx)` is the
    /// motivating example — the in-flight acquire lives in the `Receiver`, and
    /// a `Receiver` may be moved between tasks). The semaphore must unblock
    /// whoever is actually waiting now, so this follows the poller. This
    /// mirrors tokio's own `batch_semaphore`, which refreshes its waiter's
    /// `Waker` under a `will_wake` check.
    ///
    /// Stored as an atomic rather than a `Cell` to keep `Waiter` (and hence
    /// `Acquire`) `Sync`.
    task_id: AtomicUsize,
    num_permits: usize,
    is_queued: AtomicBool,
    has_permits: AtomicBool,
    /// Clock of the task that created this waiter. Note this is *not* refreshed
    /// when `task_id` is: it is only used to seed the causality of the acquired
    /// permits, and keeping the original enqueue clock is conservative (it can
    /// only add happens-before edges, never remove them).
    clock: VectorClock,
    waker: Mutex<Option<Waker>>,
}

// Implement debug in order to not output the `VectorClock`
impl fmt::Debug for Waiter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Waiter")
            .field("task_id", &self.task_id())
            .field("num_permits", &self.num_permits)
            .field("is_queued", &self.is_queued)
            .field("has_permits", &self.has_permits)
            .field("waker", &self.waker)
            .finish()
    }
}

impl Waiter {
    /// A `Waiter` is the part of an acquire that a *releasing* task can see and
    /// mutate, so it only needs to exist once an acquire actually blocks.
    ///
    /// `clock` is passed in rather than read from the ambient execution state,
    /// because it must be snapshotted when the `Acquire` was created, not when it
    /// later blocks: it feeds the happens-before edge recorded in
    /// `unblock_waiters_from_front`, and a scheduling point sits between those two
    /// moments. `task_id`, in contrast, tracks the current poller (see
    /// [`Waiter::task_id`]), so it is read here and refreshed on later polls.
    fn new(num_permits: usize, clock: VectorClock) -> Self {
        Self {
            task_id: AtomicUsize::new(ExecutionState::me().into()),
            num_permits,
            is_queued: AtomicBool::new(false),
            has_permits: AtomicBool::new(false),
            clock,
            waker: Mutex::new(None),
        }
    }

    /// The task currently waiting on this waiter. See [`Waiter::task_id`].
    fn task_id(&self) -> TaskId {
        TaskId::from(self.task_id.load(Ordering::SeqCst))
    }

    /// Point this waiter at the task that is polling it now, so that a later
    /// `release` unblocks the current poller rather than whoever polled first.
    fn set_task_id(&self, task_id: TaskId) {
        self.task_id.store(task_id.into(), Ordering::SeqCst);
    }
}

/// Number of permits (`num_available`) available to be acquired. The permits
/// are grouped into batches in the `permit_clocks` deque, such that batches
/// farther back correspond to later `release` calls. Each batch is a tuple
/// of the permits remaining in that batch and the clock of the event whence
/// the permits originate.
struct PermitsAvailable {
    // Invariant: the number of permits available is equal to the sum of the
    // batch sizes in the queue.
    num_available: usize,

    /// Batches of permits with associated clocks (corresponding to the
    /// `release` events that created them). This is an `Option` because the
    /// deque is lazily initialized; see `const_new`.
    permit_clocks: Option<VecDeque<(usize, VectorClock)>>,

    /// The clock of the last successful acquire event. Used for causal
    /// dependence in `try_acquire` failures.
    last_acquire: VectorClock,
}

// Implement debug in order to not output the `VectorClock`s
impl fmt::Debug for PermitsAvailable {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("PermitsAvailable")
            .field("num_available", &self.num_available)
            .finish()
    }
}

impl PermitsAvailable {
    fn new(num_permits: usize) -> Self {
        let mut permit_clocks = VecDeque::new();
        if num_permits > 0 {
            permit_clocks.push_back((num_permits, current::clock()));
        }
        Self {
            num_available: num_permits,
            permit_clocks: Some(permit_clocks),
            last_acquire: VectorClock::new(),
        }
    }

    const fn const_new(num_permits: usize) -> Self {
        // A `VecDeque` cannot be populated in a const fn, due to allocation.
        // Instead, we set `permit_clocks` to `None`, and initialize it lazily
        // when it is needed for the first time, to contain one batch of size
        // `num_permits`.
        Self {
            num_available: num_permits,
            permit_clocks: None,
            last_acquire: VectorClock::new(),
        }
    }

    fn available(&self) -> usize {
        self.num_available
    }

    fn init_permit_clocks(&mut self) {
        if self.permit_clocks.is_none() {
            let mut permit_clocks = VecDeque::new();
            if self.num_available > 0 {
                permit_clocks.push_back((self.num_available, VectorClock::new()));
            }
            self.permit_clocks = Some(permit_clocks);
        }
    }

    fn acquire(&mut self, mut num_permits: usize, acquire_clock: VectorClock) -> Result<VectorClock, TryAcquireError> {
        // Acquiring zero permits is always possible, and is not causally
        // dependent on any event.
        if num_permits == 0 {
            return Ok(VectorClock::new());
        }

        if num_permits <= self.num_available {
            self.init_permit_clocks();
            self.last_acquire.update(&acquire_clock);
            self.num_available -= num_permits;

            // Acquire `num_permits` from the available batches. This may
            // consume one or more batches from the queue. The resulting clock
            // is the join of all the batches used (fully or partially), since
            // the acquiry causally depends on the releases that created those
            // batches.
            let mut clock = VectorClock::new();
            let permit_clocks = self.permit_clocks.as_mut().unwrap();
            while let Some((batch_size, batch_clock)) = permit_clocks.front_mut() {
                clock.update(batch_clock);

                if num_permits < *batch_size {
                    // The current batch is larger than the number of permits
                    // requested: diminish batch, finish loop.
                    *batch_size -= num_permits;
                    num_permits = 0;
                } else {
                    // The current batch is fully consumed by the request.
                    // Remove it from the queue.
                    num_permits -= *batch_size;
                    permit_clocks.pop_front();
                }

                // Break early to avoid causally depending on the next batch.
                if num_permits == 0 {
                    break;
                }
            }

            assert_eq!(num_permits, 0);
            Ok(clock)
        } else {
            // There are not enough permits to fulfill the request.
            Err(TryAcquireError::NoPermits)
        }
    }

    fn release(&mut self, num_permits: usize, clock: VectorClock) {
        self.init_permit_clocks();
        self.num_available += num_permits;
        self.permit_clocks.as_mut().unwrap().push_back((num_permits, clock));
    }
}

/// Fairness mode for the semaphore. Determines which threads are woken when
/// permits are released.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Fairness {
    /// The semaphore is strictly fair, so earlier requesters always get
    /// priority over later ones.
    StrictlyFair,

    /// The semaphore makes no guarantees about fairness. In particular,
    /// a waiter can be starved by other threads.
    Unfair,
}

/// Where an acquire request sits relative to waiters that are already queued on
/// a [`Fairness::StrictlyFair`] semaphore. Ignored by an unfair semaphore, which
/// has no queue order to speak of.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Priority {
    /// The default: queue behind existing waiters, and do not take available
    /// permits while any waiter is queued.
    Back,

    /// Overtake every queued waiter: take available permits even when others are
    /// waiting, and if there still aren't enough, queue at the *front*.
    ///
    /// This is only correct for a requester that already holds permits of this
    /// semaphore and is escalating its own claim (see [`BatchSemaphore::upgrade`]).
    /// Such a request cannot be satisfied by making the queue wait its turn --
    /// queued waiters hold no permits, so they can never release what the
    /// requester is missing, and the requester will not release what it holds.
    /// Deadlock is avoided precisely by letting it overtake them.
    Front,
}

/// A counting semaphore which permits waiting on multiple permits at once,
/// and supports both asychronous and synchronous blocking operations.
#[derive(Debug)]
struct BatchSemaphoreState {
    id: Option<crate::annotations::ObjectId>,

    // Key invariants:
    //
    // (1) if `waiters` is nonempty and the head waiter is `H`,
    // then `H.num_permits > permits_available.available()`.  (In other words,
    // we are never in a state where there are enough permits available for the
    // first waiter.  This invariant is ensured by the `drop` handler below.)
    //
    // (2) W is in waiters iff W.is_queued
    //
    // (3) W.is_queued ==> !W.has_permits
    // Note: the converse is not true.  We can have !W.has_permits && !W.is_queued
    // when the Acquire is created but not yet polled.
    //
    // (4) closed ==> waiters.is_empty()
    waiters: VecDeque<Arc<Waiter>>,
    permits_available: PermitsAvailable,
    // TODO: should there be a clock for the close event?
    closed: bool,
}

impl BatchSemaphoreState {
    fn acquire_permits(
        &mut self,
        num_permits: usize,
        fairness: Fairness,
        priority: Priority,
    ) -> Result<(), TryAcquireError> {
        assert!(num_permits > 0);
        if self.closed {
            Err(TryAcquireError::Closed)
        } else if self.waiters.is_empty() || matches!(fairness, Fairness::Unfair) || priority == Priority::Front {
            // Permits here can be acquired in one of three scenarios:
            // - The waiter queue is empty; nobody else is waiting for permits,
            //   so if there are enough available, immediately succeed.
            // - The semaphore is operating in an unfair mode; the current
            //   thread is either requesting permits for the first time, or it
            //   was woken and selected by the scheduler. In either case, the
            //   thread may succeed, as long as there are enough permits.
            // - The request has `Priority::Front`, so it deliberately overtakes
            //   the queue (see `BatchSemaphore::upgrade`). Queued waiters hold
            //   no permits, so they cannot prevent this request from succeeding.

            let clock = self.permits_available.acquire(num_permits, current::clock())?;

            // If successful, the acquiry is causally dependent on the event
            // which released the acquired permits.
            ExecutionState::with(|s| {
                s.update_clock(&clock);
            });

            Ok(())
        } else {
            Err(TryAcquireError::NoPermits)
        }
    }

    fn unblock_waiters_from_front(&mut self) {
        while let Some(front) = self.waiters.front() {
            // A waiter whose task has already finished is stale: its `Acquire`
            // future was cancelled (e.g. a `select!` branch lost, or a
            // `poll_recv`-style API cached the `Acquire` inside a longer-lived
            // object) and the registering task then exited. There is nobody to
            // unblock, so discard the waiter without consuming permits; if the
            // `Acquire` is still alive and some other task polls it, it will
            // re-acquire from the (still available) permits.
            //
            // `remove_waiter` can reach this during execution cleanup, when the
            // task list is gone, so probe defensively and treat "can't tell" as
            // not stale (i.e. preserve the old behaviour).
            let front_is_stale = ExecutionState::try_with(|s| {
                !s.in_cleanup() && s.try_get(front.task_id()).is_some_and(|t| t.finished())
            })
            .unwrap_or(false);
            if front_is_stale {
                let waiter = self.waiters.pop_front().unwrap();
                waiter.is_queued.store(false, Ordering::SeqCst);
                // Preserve the "queued <=> waker registered" invariant asserted
                // in `Acquire::poll`; waking a finished task's waker is a no-op.
                waiter.waker.lock().unwrap().take();
                trace!("dropping stale waiter {:?} for finished task", waiter);
                continue;
            }
            if front.num_permits <= self.permits_available.available() {
                let waiter = self.waiters.pop_front().unwrap();

                crate::annotations::record_semaphore_acquire_unblocked(
                    self.id.unwrap(),
                    waiter.task_id(),
                    waiter.num_permits,
                );

                // The clock we pass into the semaphore is the clock of the
                // waiter, corresponding to the point at which the waiter was
                // enqueued. The clock we get in return corresponds to the
                // join of the clocks of the acquired permits, used to update
                // the waiter's clock to causally depend on the release events.
                let clock = self
                    .permits_available
                    .acquire(waiter.num_permits, waiter.clock.clone())
                    .unwrap();
                trace!("granted {:?} permits to waiter {:?}", waiter.num_permits, waiter);

                // Update waiter state as it is no longer in the queue
                assert!(waiter.is_queued.swap(false, Ordering::SeqCst));
                assert!(!waiter.has_permits.swap(true, Ordering::SeqCst));
                ExecutionState::with(|s| {
                    let task = s.get_mut(waiter.task_id());
                    assert!(!task.finished());
                    // The acquiry is causally dependent on the event
                    // which released the acquired permits.
                    task.clock.update(&clock);
                    task.unblock();
                });
                let mut maybe_waker = waiter.waker.lock().unwrap();
                if let Some(waker) = maybe_waker.take() {
                    waker.wake();
                }
            } else {
                return;
            }
        }
    }
}

/// Counting semaphore
#[derive(Debug)]
pub struct BatchSemaphore {
    state: RefCell<BatchSemaphoreState>,
    fairness: Fairness,
    #[allow(unused)]
    signature: ResourceSignature,
}

/// Error returned from the [`BatchSemaphore::try_acquire`] function.
#[derive(Debug, PartialEq, Eq)]
pub enum TryAcquireError {
    /// The semaphore has been closed and cannot issue new permits.
    Closed,

    /// The semaphore has no available permits.
    NoPermits,
}

/// Error returned from the [`BatchSemaphore::acquire`] function.
///
/// An `acquire*` operation can only fail if the semaphore has been
/// closed.
#[derive(Debug)]
pub struct AcquireError(());

impl AcquireError {
    fn closed() -> AcquireError {
        AcquireError(())
    }
}

impl fmt::Display for AcquireError {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(fmt, "semaphore closed")
    }
}

impl std::error::Error for AcquireError {}

impl BatchSemaphore {
    /// Creates a new semaphore with the initial number of permits.
    #[track_caller]
    pub fn new(num_permits: usize, fairness: Fairness) -> Self {
        Self::new_with_signature(
            num_permits,
            fairness,
            ExecutionState::new_resource_signature(ResourceType::BatchSemaphore),
        )
    }

    pub fn new_with_signature(num_permits: usize, fairness: Fairness, signature: ResourceSignature) -> Self {
        let state = RefCell::new(BatchSemaphoreState {
            id: Some(crate::annotations::record_semaphore_created()),
            waiters: VecDeque::new(),
            permits_available: PermitsAvailable::new(num_permits),
            closed: false,
        });
        Self {
            state,
            fairness,
            signature,
        }
    }

    /// Creates a new semaphore with the initial number of permits.
    #[track_caller]
    pub const fn const_new(num_permits: usize, fairness: Fairness) -> Self {
        Self::const_new_with_signature(
            num_permits,
            fairness,
            ResourceSignature::new_const(ResourceType::BatchSemaphore),
        )
    }

    pub const fn const_new_with_signature(
        num_permits: usize,
        fairness: Fairness,
        signature: ResourceSignature,
    ) -> Self {
        let state = RefCell::new(BatchSemaphoreState {
            id: None,
            waiters: VecDeque::new(),
            permits_available: PermitsAvailable::const_new(num_permits),
            closed: false,
        });
        Self {
            state,
            fairness,
            signature,
        }
    }

    /// Returns the current number of available permits.
    pub fn available_permits(&self) -> usize {
        let state = self.state.borrow();
        state.permits_available.available()
    }

    fn init_object_id(&self) {
        let mut state = self.state.borrow_mut();
        if state.id.is_none() {
            state.id = Some(crate::annotations::record_semaphore_created());
        }
    }

    /// Closes the semaphore. This prevents the semaphore from issuing new
    /// permits and notifies all pending waiters.
    pub fn close(&self) {
        thread::switch();
        self.close_no_scheduling_point();
    }

    /// Closes the semaphore without invoking `thread::switch`
    pub fn close_no_scheduling_point(&self) {
        self.init_object_id();
        let mut state = self.state.borrow_mut();
        if state.closed {
            return;
        }
        crate::annotations::record_semaphore_closed(state.id.unwrap());
        state.closed = true;

        // Wake up all the waiters.  Since we've marked the state as closed, they
        // will all return `AcquireError::closed` from their acquire calls.
        let ptr = &*state as *const BatchSemaphoreState;
        for waiter in state.waiters.drain(..) {
            trace!(
                "semaphore {:p} removing and waking up waiter {:?} on close",
                ptr,
                waiter,
            );
            assert!(waiter.is_queued.swap(false, Ordering::SeqCst));
            assert!(!waiter.has_permits.load(Ordering::SeqCst)); // sanity check
            ExecutionState::with(|exec_state| {
                // A waiter whose task has finished is stale (its `Acquire` was
                // cancelled and the task exited); there is nothing to unblock.
                if !exec_state.in_cleanup() && !exec_state.get(waiter.task_id()).finished() {
                    exec_state.get_mut(waiter.task_id()).unblock();
                }
            });
            let mut maybe_waker = waiter.waker.lock().unwrap();
            if let Some(waker) = maybe_waker.take() {
                waker.wake();
            }
        }
    }

    /// Returns true iff the semaphore is closed.
    pub fn is_closed(&self) -> bool {
        let state = self.state.borrow();
        state.closed
    }

    /// Try to acquire the specified number of permits from the Semaphore.
    /// If the permits are available, returns Ok(())
    /// If the semaphore is closed, returns `Err(TryAcquireError::Closed)`
    /// If there aren't enough permits, returns `Err(TryAcquireError::NoPermits)`
    pub fn try_acquire(&self, num_permits: usize) -> Result<(), TryAcquireError> {
        thread::switch();

        self.init_object_id();
        let mut state = self.state.borrow_mut();
        let id = state.id.unwrap();
        let res = state
            .acquire_permits(num_permits, self.fairness, Priority::Back)
            .inspect_err(|_err| {
                // Conservatively, the requester causally depends on the
                // last successful acquire.
                // TODO: This is not precise, but `try_acquire` causal dependency
                // TODO: is both hard to define, and is most likely not worth the
                // TODO: effort. The cases where causality would be tracked
                // TODO: "imprecisely" do not correspond to commonly used sync.
                // TODO: primitives, such as mutexes, mutexes, or condvars.
                // TODO: An example would be a counting semaphore used to guard
                // TODO: access to N homogenous resources (as opposed to FIFO,
                // TODO: heterogenous resources).
                // TODO: More precision could be gained by tracking clocks for all
                // TODO: current permit holders, with a data structure similar to
                // TODO: `permits_available`.
                ExecutionState::with(|s| {
                    s.update_clock(&state.permits_available.last_acquire);
                });
            });
        drop(state);

        // If we won the race for permits of an unfair semaphore, re-block
        // other waiting threads that can no longer succeed.
        if res.is_ok() {
            self.reblock_if_unfair();
        }

        crate::annotations::record_semaphore_try_acquire(id, num_permits, res.is_ok());

        res
    }

    /// Clean-up method used when a thread succeeds in acquiring permits. If
    /// the semaphore is unfair, a preceding `release` may have unblocked a
    /// number of threads, some of which may no longer be able to succeed with
    /// the permits remaining in the semaphore.
    fn reblock_if_unfair(&self) {
        if self.fairness == Fairness::Unfair {
            let state = self.state.borrow_mut();
            ExecutionState::with(|s| {
                for waiter in &state.waiters {
                    let available = state.permits_available.available();
                    // Skip stale waiters: the task that registered the waiter
                    // has finished, so there is nothing to block.
                    if available < waiter.num_permits && s.try_get(waiter.task_id()).is_some_and(|t| !t.finished()) {
                        // Block this waiter: it cannot succeed (there are not
                        // enough permits available); its `poll` would return
                        // without resolving.
                        s.get_mut(waiter.task_id()).block(false);
                    }
                }
            });
        }
    }

    fn enqueue_waiter(&self, waiter: &Arc<Waiter>, priority: Priority) {
        let mut state = self.state.borrow_mut();

        trace!(
            "enqueuing waiter {:?} ({priority:?}) for semaphore {:p}",
            waiter,
            &self.state
        );
        match priority {
            Priority::Back => state.waiters.push_back(waiter.clone()),
            // Overtakes the queue rather than joining its tail. Key invariant (1)
            // still holds: we only get here because the acquire failed, and a
            // `Priority::Front` acquire only fails when there really aren't
            // enough permits available, so the new head cannot be grantable.
            Priority::Front => state.waiters.push_front(waiter.clone()),
        }

        assert!(!waiter.has_permits.load(Ordering::SeqCst));
        assert!(!waiter.is_queued.swap(true, Ordering::SeqCst));
    }

    fn remove_waiter(&self, waiter: &Arc<Waiter>) {
        let mut state = self.state.borrow_mut();

        trace!(waiters = ?state.waiters, "removing waiter {:?} from semaphore {:p}", waiter, &self.state);

        // sanity checks
        assert!(!state.closed);
        assert!(!waiter.has_permits.load(Ordering::SeqCst));

        let index = state
            .waiters
            .iter()
            .position(|x| Arc::ptr_eq(x, waiter))
            .expect("did not find waiter");

        state.waiters.remove(index).unwrap();
        assert!(waiter.is_queued.swap(false, Ordering::SeqCst));

        match self.fairness {
            Fairness::StrictlyFair => {
                if index == 0 {
                    // If the semaphore is strictly fair, and we removed the first waiter, check if its
                    // removal unblocks remaining waiters.  This can happen in the following situation:
                    // - the semahore has 1 permit available
                    // - there are 2 waiters W1 and W2 where W1 wants 2 permits, and W2 wants 1 permit
                    // - if W1 gives up and drops out, we want to ensure W2 is granted the semaphore
                    state.unblock_waiters_from_front();
                }
            }
            Fairness::Unfair => {}
        }
    }

    /// Acquire the specified number of permits (async API)
    pub fn acquire(&self, num_permits: usize) -> Acquire<'_> {
        // No switch here; switch should be triggered on polling future
        self.init_object_id();
        Acquire::new(self, num_permits, Priority::Back)
    }

    /// Acquire the specified number of permits (blocking API)
    pub fn acquire_blocking(&self, num_permits: usize) -> Result<(), AcquireError> {
        crate::future::block_on(self.acquire(num_permits))
    }

    /// Release `num_permits` back to the Semaphore
    pub fn release(&self, num_permits: usize) {
        thread::switch();

        self.init_object_id();
        if num_permits == 0 {
            return;
        }

        let mut state = self.state.borrow_mut();

        crate::annotations::record_semaphore_release(state.id.unwrap(), num_permits);

        if ExecutionState::should_stop() {
            // In case we are panicking, we release permits, but also clear
            // the waiters queue: we should not unblock the threads at this
            // point. However, the permits are released such that future
            // acquires may succeed, as long as the requesters were not
            // blocking on the semaphore at the time of the panic. This is
            // used to correctly model lock poisoning.
            state.permits_available.release(num_permits, VectorClock::new());
            for waiter in &state.waiters {
                waiter.is_queued.swap(false, Ordering::SeqCst);
            }
            state.waiters.clear();
            state.closed = true;
            return;
        }

        // Permits released into the semaphore reflect the releasing thread's
        // clock; future acquires of those permits are causally dependent on
        // this event.
        ExecutionState::with(|s| {
            let clock = s.increment_clock();
            state.permits_available.release(num_permits, clock.clone());
        });

        // `ExecutionState::me()` is only wanted for this trace, so let the macro's
        // level check decide whether to pay for it. Computing it eagerly cost an
        // `ExecutionState::with` on every release even with tracing disabled.
        trace!(task = ?ExecutionState::me(), avail = ?state.permits_available, waiters = ?state.waiters, "released {} permits for semaphore {:p}", num_permits, &self.state);

        match self.fairness {
            Fairness::StrictlyFair => {
                // in a strictly fair mode we will grant permits to waiters from the front
                // of the queue, as long as there are enough permits available
                state.unblock_waiters_from_front();
            }
            Fairness::Unfair => {
                // in an unfair mode, we will unblock all the waiters for which
                // there are enough permits available, then let them race
                let num_available = state.permits_available.available();
                for waiter in &mut state.waiters {
                    if waiter.num_permits <= num_available {
                        // A waiter whose task has already finished is stale: its
                        // `Acquire` was cancelled (and possibly cached in a
                        // longer-lived object) and the task then exited. Unlike
                        // the strictly fair case there is nothing to clean up —
                        // an unfair waiter holds no permits, so it blocks
                        // nobody — but there is also nobody to unblock.
                        let stale = ExecutionState::with(|s| {
                            let task = s.get_mut(waiter.task_id());
                            if task.finished() {
                                true
                            } else {
                                task.unblock();
                                false
                            }
                        });
                        if stale {
                            continue;
                        }
                        let maybe_waker = waiter.waker.lock().unwrap();
                        if let Some(waker) = maybe_waker.as_ref() {
                            waker.wake_by_ref();
                        }
                    }
                }
            }
        }
        drop(state);
    }

    /// Atomically `upgrade` from holding `permits_currently_held` permits to holding
    /// `permits_to_be_held`, without ever dropping below `permits_currently_held` in between.
    /// The motivating use case is `parking_lot`'s `RwLockUpgradableReadGuard::upgrade`, which must
    /// take a read guard to a write guard without letting any writer in along the way.
    ///
    /// This is implemented by acquiring only the *missing* permits
    /// (`permits_to_be_held - permits_currently_held`), with priority over any waiter already
    /// queued, so that the request overtakes the queue. Both halves of that matter:
    ///
    /// * Keeping the held permits means no other task can claim the resource mid-upgrade. Releasing
    ///   them first (even for an instant) would hand the resource to a queued waiter, which for an
    ///   `RwLock` means a writer mutating the data an upgradable reader had already observed.
    /// * Overtaking the queue is what makes that safe rather than deadlock-prone. Since we hold
    ///   permits we will not release, a queued waiter ahead of us may be unsatisfiable (an `RwLock`
    ///   writer wants *all* permits), so waiting our turn behind it could deadlock. Queued waiters
    ///   hold no permits, so overtaking them costs nothing but their place in line -- which is
    ///   exactly the priority a real upgradable read lock gives an upgrade.
    ///
    /// The upgrade therefore blocks only on tasks that *currently hold* permits, and is granted as
    /// soon as they release. The returned future must be driven to completion; if it is dropped
    /// first, the caller still holds `permits_currently_held`.
    ///
    /// At most one `upgrade` may be in flight on a semaphore at a time. Two concurrent upgraders
    /// could each be waiting for permits the other holds, which no queue discipline can resolve.
    /// Callers are expected to enforce this (an `RwLock` does: there is only ever one upgradable
    /// reader).
    pub fn upgrade(&self, permits_currently_held: usize, permits_to_be_held: usize) -> Acquire<'_> {
        assert!(permits_currently_held > 0);
        assert!(permits_to_be_held > permits_currently_held);

        self.init_object_id();
        Acquire::new(self, permits_to_be_held - permits_currently_held, Priority::Front)
    }

    /// The non-blocking analogue of [`BatchSemaphore::upgrade`]: succeeds only if the missing
    /// permits are available right now, and never blocks or queues.
    ///
    /// Like `upgrade`, this ignores queued waiters (they hold no permits, so they cannot be the
    /// reason the upgrade is short of permits). A `try_upgrade` therefore fails only when some
    /// other task actually *holds* permits the upgrade needs.
    pub fn try_upgrade(&self, permits_currently_held: usize, permits_to_be_held: usize) -> Result<(), TryAcquireError> {
        assert!(permits_currently_held > 0);
        assert!(permits_to_be_held > permits_currently_held);

        thread::switch();

        self.init_object_id();
        let num_permits = permits_to_be_held - permits_currently_held;
        let mut state = self.state.borrow_mut();
        let id = state.id.unwrap();
        let res = state
            .acquire_permits(num_permits, self.fairness, Priority::Front)
            .inspect_err(|_err| {
                // Conservatively, the requester causally depends on the last successful acquire;
                // see the equivalent reasoning in `try_acquire`.
                ExecutionState::with(|s| {
                    s.update_clock(&state.permits_available.last_acquire);
                });
            });
        drop(state);

        // If we took permits from an unfair semaphore, re-block waiting threads that can no longer
        // succeed.
        if res.is_ok() {
            self.reblock_if_unfair();
        }

        crate::annotations::record_semaphore_try_acquire(id, num_permits, res.is_ok());

        res
    }
}

// Safety: Semaphore is never actually passed across true threads, only across continuations. The
// RefCell<_> type therefore can't be preempted mid-bookkeeping-operation.
// TODO we shouldn't need to do this, but RefCell is not Send, and anything we put within a Semaphore
// TODO needs to be Send.
unsafe impl Send for BatchSemaphore {}
unsafe impl Sync for BatchSemaphore {}

impl Default for BatchSemaphore {
    #[track_caller]
    fn default() -> Self {
        Self::new(Default::default(), Fairness::StrictlyFair)
    }
}

/// The future that results from async calls to `acquire*`.
/// Callers must `await` on this future to obtain the necessary permits.
pub struct Acquire<'a> {
    semaphore: &'a BatchSemaphore,
    num_permits: usize,

    /// Where this acquire sits relative to waiters already queued on a fair
    /// semaphore. Only [`BatchSemaphore::upgrade`] uses [`Priority::Front`]; see
    /// there for why an upgrade must overtake the queue.
    priority: Priority,

    /// Snapshotted when this `Acquire` is created, and moved into the `Waiter` if
    /// this acquire ends up blocking. See `Waiter::new` for why the snapshot must
    /// happen here rather than at enqueue time.
    clock: VectorClock,

    /// The shared part of this acquire, allocated only once the acquire has to
    /// block. An acquire that gets its permits immediately is never visible to
    /// any other task, so it needs no shared state and no allocation. While this
    /// is `None`, `has_permits` below is authoritative.
    waiter: Option<Arc<Waiter>>,

    /// Whether permits have been granted, for the case where no `Waiter` exists.
    /// Once one does, the releasing task writes `Waiter::has_permits` instead and
    /// this field is unused; read through `Acquire::has_permits`.
    has_permits: bool,

    completed: bool, // Has the future completed yet?
    never_polled: bool,
}

// Implement Debug in order to not output the `VectorClock`, matching `Waiter`.
impl fmt::Debug for Acquire<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Acquire")
            .field("num_permits", &self.num_permits)
            .field("priority", &self.priority)
            .field("waiter", &self.waiter)
            .field("has_permits", &self.has_permits())
            .field("completed", &self.completed)
            .finish()
    }
}

impl<'a> Acquire<'a> {
    fn new(semaphore: &'a BatchSemaphore, num_permits: usize, priority: Priority) -> Self {
        Self {
            semaphore,
            num_permits,
            priority,
            clock: current::clock(),
            waiter: None,
            has_permits: false,
            completed: false,
            never_polled: true,
        }
    }

    /// Have permits been granted to this acquire? Once a `Waiter` exists the
    /// releasing task owns that flag, so the shared copy is authoritative.
    fn has_permits(&self) -> bool {
        match &self.waiter {
            Some(waiter) => waiter.has_permits.load(Ordering::SeqCst),
            None => self.has_permits,
        }
    }

    /// Is this acquire in the semaphore's waiter queue? Only possible once a
    /// `Waiter` has been allocated, since the queue holds `Arc<Waiter>`.
    fn is_queued(&self) -> bool {
        match &self.waiter {
            Some(waiter) => waiter.is_queued.load(Ordering::SeqCst),
            None => false,
        }
    }

    fn grant_permits(&mut self) {
        match &self.waiter {
            Some(waiter) => waiter.has_permits.store(true, Ordering::SeqCst),
            None => self.has_permits = true,
        }
    }

    /// The shared `Waiter` for this acquire, allocating it if this is the first
    /// time the acquire has had to block. Returns an owned handle so callers can
    /// still use `self.semaphore` without holding a borrow of `self`.
    fn waiter_for_blocking(&mut self) -> Arc<Waiter> {
        if let Some(waiter) = &self.waiter {
            return Arc::clone(waiter);
        }
        let waiter = Arc::new(Waiter::new(self.num_permits, self.clock.clone()));
        self.waiter = Some(Arc::clone(&waiter));
        waiter
    }
}

impl Future for Acquire<'_> {
    type Output = Result<(), AcquireError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        assert!(!self.completed);

        // One borrow of the semaphore state rather than two (`is_closed` and
        // `available_permits` each took their own). Both reads describe the same
        // instant, before the scheduling point below, so merging them is sound.
        // Reads *after* the switch must stay separate and fresh, because other
        // tasks may have run in between.
        let will_succeed = self.has_permits() || {
            let state = self.semaphore.state.borrow();
            state.closed || state.permits_available.available() >= self.num_permits
        };

        // If the acquire will succeed on the first try, we need to context switch once to allow the previous
        // event to become visible. If we won't succeed, then we still need to context switch if the act of
        // blocking does not commute with other operations on `batch_semaphore` (double-yield optimization,
        // reasoning below).
        //
        // Fair Semaphores: blocking adds the current task to an *ordered* waiter queue. Two blocking acquires
        // *do not commute* because in one ordering the queue will be [T1 T2] and in the other ordering [T2 T1].
        // Thus we cannot apply the double-yield optimization for fair semaphores.
        //
        // Unfair Semaphores: blocking adds the current task to an *unordered set* of waiters. To check if the
        // double-yield is valid we check if each operation (Z) on the semaphore commutes with a blocking acquire (Y1):
        //
        //     - Blocking Acquire: in both orderings `Z Y1` and `Y1 Z`, the waiter set has the same members, thus
        //       the operations commute.
        //     - Try Acquire: the try-acquire will fail in both orderings without changing the state of the semaphore
        //     - Release: if the release unblocks Y1, then the optimization is not applicable. Otherwise, it must
        //       unblock another task in the waiter set. As waiter-set insertion and removal for disjoint elements
        //       commutes, release operations also commute in this case.
        //
        // Thus we apply the double-yield optimization for *unfair* semaphores only
        let blocking_is_not_commutative = self.semaphore.fairness == Fairness::StrictlyFair;

        if self.never_polled && (will_succeed || blocking_is_not_commutative) {
            thread::switch();
        }
        self.never_polled = false;

        let out = if self.has_permits() {
            assert!(!self.is_queued());
            self.completed = true;
            trace!("Acquire::poll for {:?} with permits", self);
            Poll::Ready(Ok(()))
        } else if self.semaphore.is_closed() {
            assert!(!self.is_queued());
            self.completed = true;
            trace!("Acquire::poll for {:?} with closed", self);
            Poll::Ready(Err(AcquireError::closed()))
        } else {
            let is_queued = self.is_queued();
            trace!("Acquire::poll for {:?}; is queued: {is_queued:?}", self);

            // Sanity check: there should be a waker if the waiter is in
            // the queue. Also true for unfair semaphores, which wake by ref.
            //
            // `debug_assert` rather than `assert`: this takes a `std::sync::Mutex`
            // on every poll, including the uncontended fast path, purely to check
            // an internal invariant.
            debug_assert_eq!(
                is_queued,
                self.waiter
                    .as_ref()
                    .is_some_and(|waiter| waiter.waker.lock().unwrap().is_some())
            );

            // Should the waiter try to acquire permits here? Four cases:
            // 1. unfair semaphore, waiter not yet enqueued;
            // 2. fair semaphore, waiter not yet enqueued;
            // 3. unfair semaphore, waiter already enqueued.
            // 4. fair semaphore, waiter already enqueued;
            //
            // 1. and 2. are similar: the future was polled for the first time,
            // so the waiter will try to acquire some permits. If successful,
            // the waiter need not be enqueued, and the future is resolved.
            // Otherwise, the waiter is added to the queue.
            //
            // 3. is slightly different: the future was polled, even though the
            // waiter was already in the queue. This can happen either because
            // the semaphore just received some permits and woke the waiter up,
            // or because the future itself was polled manually. Either way,
            // the semaphore is queried.
            //
            // 4. is a case where we do not try to acquire permits. The request
            // would always fail, and the waiter should remain suspended until
            // the semaphore has explicitly unblocked it and given it permits
            // during a `release` call.
            let try_to_acquire = match (self.semaphore.fairness, is_queued) {
                // written this way to mirror the cases described above
                (Fairness::Unfair, false) | (Fairness::StrictlyFair, false) | (Fairness::Unfair, true) => true,
                (Fairness::StrictlyFair, true) => false,
            };

            if try_to_acquire {
                // Access the semaphore state directly instead of `try_acquire`,
                // because in case of `NoPermits`, we do not want to update the
                // clock, as this thread will be blocked below.
                let mut state = self.semaphore.state.borrow_mut();
                let id = state.id.unwrap();
                let acquire_result = state.acquire_permits(self.num_permits, self.semaphore.fairness, self.priority);
                drop(state);

                match acquire_result {
                    Ok(()) => {
                        if is_queued {
                            let waiter = self
                                .waiter
                                .clone()
                                .expect("a queued acquire must have an allocated waiter");
                            crate::annotations::record_semaphore_acquire_unblocked(
                                id,
                                waiter.task_id(),
                                waiter.num_permits,
                            );
                            self.semaphore.remove_waiter(&waiter);
                        } else {
                            crate::annotations::record_semaphore_acquire_fast(id, self.num_permits);
                        }
                        self.grant_permits();
                        self.completed = true;
                        trace!("Acquire::poll for {:?} that got permits", self);

                        // If the semaphore is unfair, re-block other waiting
                        // threads that can no longer succeed.
                        self.semaphore.reblock_if_unfair();

                        Poll::Ready(Ok(()))
                    }
                    Err(TryAcquireError::NoPermits) => {
                        // This acquire has to block, so it now becomes visible to
                        // whichever task releases permits. That is the first point
                        // at which shared state is needed, so it is where the
                        // `Waiter` gets allocated.
                        let waiter = self.waiter_for_blocking();

                        let mut maybe_waker = waiter.waker.lock().unwrap();
                        *maybe_waker = Some(cx.waker().clone());
                        drop(maybe_waker);

                        // Point the waiter at whoever is polling now: this future
                        // may have been created by a different task.
                        waiter.set_task_id(ExecutionState::me());

                        if !is_queued {
                            crate::annotations::record_semaphore_acquire_blocked(id, self.num_permits);
                            // `enqueue_waiter` sets `is_queued` itself.
                            self.semaphore.enqueue_waiter(&waiter, self.priority);
                        }
                        trace!("Acquire::poll for {:?} that is enqueued", self);
                        Poll::Pending
                    }
                    Err(TryAcquireError::Closed) => unreachable!(),
                }
            } else {
                // No progress made, future is still pending. The waiter stays in
                // the queue, but re-point it at the current poller and refresh
                // its waker: this future may have been created by (or last
                // polled by) another task, and `release` must wake whoever is
                // waiting now. Without this, a permit granted to this waiter
                // would unblock a task that is no longer interested, and the
                // actual poller would never be woken.
                let waiter = self
                    .waiter
                    .as_ref()
                    .expect("a queued acquire must have an allocated waiter");
                *waiter.waker.lock().unwrap() = Some(cx.waker().clone());
                waiter.set_task_id(ExecutionState::me());
                Poll::Pending
            }
        };
        if matches!(out, Poll::Pending) {
            // `Backtrace::capture()` is a noop (it returns the constant `disabled()`) if `RUST_BACKTRACE`/`RUST_LIB_BACKTRACE` is not set.
            ExecutionState::with(|state| {
                state.current_mut().backtrace = if backtrace_enabled() {
                    Some(std::backtrace::Backtrace::force_capture())
                } else {
                    None
                }
            })
        }
        out
    }
}

impl Drop for Acquire<'_> {
    fn drop(&mut self) {
        trace!("Acquire::drop for {:?}", self);
        if self.is_queued() {
            // If the associated waiter is in the wait list, remove it
            let waiter = self
                .waiter
                .clone()
                .expect("a queued acquire must have an allocated waiter");
            self.semaphore.remove_waiter(&waiter);
        } else if self.has_permits() && !self.completed {
            // If the waiter was granted permits, release them. Note this must also
            // fire for an acquire that got its permits without ever allocating a
            // waiter, otherwise the semaphore leaks permits.
            self.semaphore.release(self.num_permits);
        }
    }
}

impl crate::annotations::WithName for &BatchSemaphore {
    fn with_name_and_kind(self, name: Option<&str>, kind: Option<&str>) -> Self {
        self.init_object_id();
        crate::annotations::record_name_for_object(self.state.borrow().id.unwrap(), name, kind);
        self
    }
}

impl crate::annotations::WithName for BatchSemaphore {
    fn with_name_and_kind(self, name: Option<&str>, kind: Option<&str>) -> Self {
        (&self).with_name_and_kind(name, kind);
        self
    }
}

impl BatchSemaphore {
    /// Returns a reference to this semaphore's resource signature.
    pub fn signature(&self) -> &ResourceSignature {
        &self.signature
    }
}