Skip to main content

shuttle_engine/future/
batch_semaphore.rs

1//! A counting semaphore supporting both async and sync operations.
2use crate::runtime::execution::ExecutionState;
3use crate::runtime::task::{clock::VectorClock, TaskId};
4use crate::runtime::thread;
5use crate::sync_types::{ResourceSignature, ResourceType};
6use crate::{backtrace_enabled, current};
7use std::cell::RefCell;
8use std::collections::VecDeque;
9use std::fmt;
10use std::future::Future;
11use std::pin::Pin;
12use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
13use std::sync::Arc;
14use std::sync::Mutex;
15use std::task::{Context, Poll, Waker};
16use tracing::trace;
17
18struct Waiter {
19    /// The task waiting on this waiter's `Acquire`.
20    ///
21    /// Refreshed on every poll (like `waker`) rather than frozen at creation
22    /// time. An `Acquire` future is not necessarily owned by the task that
23    /// created it: it can be cached inside a longer-lived object and later
24    /// polled by a different task (tokio's `poll_recv(&mut self, cx)` is the
25    /// motivating example — the in-flight acquire lives in the `Receiver`, and
26    /// a `Receiver` may be moved between tasks). The semaphore must unblock
27    /// whoever is actually waiting now, so this follows the poller. This
28    /// mirrors tokio's own `batch_semaphore`, which refreshes its waiter's
29    /// `Waker` under a `will_wake` check.
30    ///
31    /// Stored as an atomic rather than a `Cell` to keep `Waiter` (and hence
32    /// `Acquire`) `Sync`.
33    task_id: AtomicUsize,
34    num_permits: usize,
35    /// How many permits must be available for this waiter to make progress: `num_permits` for a
36    /// plain acquire, and the threshold at which it reserves the semaphore for a reserving acquire
37    /// (see [`BatchSemaphore::acquire_reserving`]). Only unfair semaphores look at this.
38    min_permits: usize,
39    priority: Priority,
40    is_queued: AtomicBool,
41    has_permits: AtomicBool,
42    /// Clock of the task that created this waiter. Note this is *not* refreshed
43    /// when `task_id` is: it is only used to seed the causality of the acquired
44    /// permits, and keeping the original enqueue clock is conservative (it can
45    /// only add happens-before edges, never remove them).
46    clock: VectorClock,
47    waker: Mutex<Option<Waker>>,
48}
49
50// Implement debug in order to not output the `VectorClock`
51impl fmt::Debug for Waiter {
52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53        f.debug_struct("Waiter")
54            .field("task_id", &self.task_id())
55            .field("num_permits", &self.num_permits)
56            .field("min_permits", &self.min_permits)
57            .field("priority", &self.priority)
58            .field("is_queued", &self.is_queued)
59            .field("has_permits", &self.has_permits)
60            .field("waker", &self.waker)
61            .finish()
62    }
63}
64
65impl Waiter {
66    /// A `Waiter` is the part of an acquire that a *releasing* task can see and
67    /// mutate, so it only needs to exist once an acquire actually blocks.
68    ///
69    /// `clock` is passed in rather than read from the ambient execution state,
70    /// because it must be snapshotted when the `Acquire` was created, not when it
71    /// later blocks: it feeds the happens-before edge recorded in
72    /// `unblock_waiters_from_front`, and a scheduling point sits between those two
73    /// moments. `task_id`, in contrast, tracks the current poller (see
74    /// [`Waiter::task_id`]), so it is read here and refreshed on later polls.
75    fn new(num_permits: usize, min_permits: usize, priority: Priority, clock: VectorClock) -> Self {
76        Self {
77            task_id: AtomicUsize::new(ExecutionState::me().into()),
78            num_permits,
79            min_permits,
80            priority,
81            is_queued: AtomicBool::new(false),
82            has_permits: AtomicBool::new(false),
83            clock,
84            waker: Mutex::new(None),
85        }
86    }
87
88    /// The task currently waiting on this waiter. See [`Waiter::task_id`].
89    fn task_id(&self) -> TaskId {
90        TaskId::from(self.task_id.load(Ordering::SeqCst))
91    }
92
93    /// Point this waiter at the task that is polling it now, so that a later
94    /// `release` unblocks the current poller rather than whoever polled first.
95    fn set_task_id(&self, task_id: TaskId) {
96        self.task_id.store(task_id.into(), Ordering::SeqCst);
97    }
98}
99
100/// Number of permits (`num_available`) available to be acquired. The permits
101/// are grouped into batches in the `permit_clocks` deque, such that batches
102/// farther back correspond to later `release` calls. Each batch is a tuple
103/// of the permits remaining in that batch and the clock of the event whence
104/// the permits originate.
105struct PermitsAvailable {
106    // Invariant: the number of permits available is equal to the sum of the
107    // batch sizes in the queue.
108    num_available: usize,
109
110    /// Batches of permits with associated clocks (corresponding to the
111    /// `release` events that created them). This is an `Option` because the
112    /// deque is lazily initialized; see `const_new`.
113    permit_clocks: Option<VecDeque<(usize, VectorClock)>>,
114
115    /// The clock of the last successful acquire event. Used for causal
116    /// dependence in `try_acquire` failures.
117    last_acquire: VectorClock,
118}
119
120// Implement debug in order to not output the `VectorClock`s
121impl fmt::Debug for PermitsAvailable {
122    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123        f.debug_struct("PermitsAvailable")
124            .field("num_available", &self.num_available)
125            .finish()
126    }
127}
128
129impl PermitsAvailable {
130    fn new(num_permits: usize) -> Self {
131        let mut permit_clocks = VecDeque::new();
132        if num_permits > 0 {
133            permit_clocks.push_back((num_permits, current::clock()));
134        }
135        Self {
136            num_available: num_permits,
137            permit_clocks: Some(permit_clocks),
138            last_acquire: VectorClock::new(),
139        }
140    }
141
142    const fn const_new(num_permits: usize) -> Self {
143        // A `VecDeque` cannot be populated in a const fn, due to allocation.
144        // Instead, we set `permit_clocks` to `None`, and initialize it lazily
145        // when it is needed for the first time, to contain one batch of size
146        // `num_permits`.
147        Self {
148            num_available: num_permits,
149            permit_clocks: None,
150            last_acquire: VectorClock::new(),
151        }
152    }
153
154    fn available(&self) -> usize {
155        self.num_available
156    }
157
158    fn init_permit_clocks(&mut self) {
159        if self.permit_clocks.is_none() {
160            let mut permit_clocks = VecDeque::new();
161            if self.num_available > 0 {
162                permit_clocks.push_back((self.num_available, VectorClock::new()));
163            }
164            self.permit_clocks = Some(permit_clocks);
165        }
166    }
167
168    fn acquire(&mut self, mut num_permits: usize, acquire_clock: VectorClock) -> Result<VectorClock, TryAcquireError> {
169        // Acquiring zero permits is always possible, and is not causally
170        // dependent on any event.
171        if num_permits == 0 {
172            return Ok(VectorClock::new());
173        }
174
175        if num_permits <= self.num_available {
176            self.init_permit_clocks();
177            self.last_acquire.update(&acquire_clock);
178            self.num_available -= num_permits;
179
180            // Acquire `num_permits` from the available batches. This may
181            // consume one or more batches from the queue. The resulting clock
182            // is the join of all the batches used (fully or partially), since
183            // the acquiry causally depends on the releases that created those
184            // batches.
185            let mut clock = VectorClock::new();
186            let permit_clocks = self.permit_clocks.as_mut().unwrap();
187            while let Some((batch_size, batch_clock)) = permit_clocks.front_mut() {
188                clock.update(batch_clock);
189
190                if num_permits < *batch_size {
191                    // The current batch is larger than the number of permits
192                    // requested: diminish batch, finish loop.
193                    *batch_size -= num_permits;
194                    num_permits = 0;
195                } else {
196                    // The current batch is fully consumed by the request.
197                    // Remove it from the queue.
198                    num_permits -= *batch_size;
199                    permit_clocks.pop_front();
200                }
201
202                // Break early to avoid causally depending on the next batch.
203                if num_permits == 0 {
204                    break;
205                }
206            }
207
208            assert_eq!(num_permits, 0);
209            Ok(clock)
210        } else {
211            // There are not enough permits to fulfill the request.
212            Err(TryAcquireError::NoPermits)
213        }
214    }
215
216    fn release(&mut self, num_permits: usize, clock: VectorClock) {
217        self.init_permit_clocks();
218        self.num_available += num_permits;
219        self.permit_clocks.as_mut().unwrap().push_back((num_permits, clock));
220    }
221}
222
223/// Fairness mode for the semaphore. Determines which threads are woken when
224/// permits are released.
225#[derive(Clone, Copy, Debug, PartialEq, Eq)]
226pub enum Fairness {
227    /// The semaphore is strictly fair, so earlier requesters always get
228    /// priority over later ones, unless they opt into [`BatchSemaphore::acquire_bypassable`].
229    StrictlyFair,
230
231    /// The semaphore makes no guarantees about fairness. In particular,
232    /// a waiter can be starved by other threads.
233    Unfair,
234}
235
236/// Where an acquire request sits relative to waiters that are already queued on
237/// a [`Fairness::StrictlyFair`] semaphore. Ignored by an unfair semaphore, which
238/// has no queue order to speak of.
239#[derive(Clone, Copy, Debug, PartialEq, Eq)]
240enum Priority {
241    /// The default: queue behind existing waiters, except bypassable waiters
242    /// whose requests cannot be satisfied yet.
243    Back,
244
245    /// Queue as with `Back`, but let later requests pass while waiting for enough permits.
246    Bypassable,
247
248    /// Overtake every queued waiter: take available permits even when others are
249    /// waiting, and if there still aren't enough, queue at the *front*.
250    ///
251    /// This is only correct for a requester that already holds permits of this
252    /// semaphore and is escalating its own claim (see [`BatchSemaphore::upgrade`]).
253    /// Such a request cannot be satisfied by making the queue wait its turn --
254    /// queued waiters hold no permits, so they can never release what the
255    /// requester is missing, and the requester will not release what it holds.
256    /// Deadlock is avoided precisely by letting it overtake them.
257    Front,
258}
259
260/// A counting semaphore which permits waiting on multiple permits at once,
261/// and supports both asychronous and synchronous blocking operations.
262#[derive(Debug)]
263struct BatchSemaphoreState {
264    id: Option<crate::annotations::ObjectId>,
265
266    // Key invariants:
267    //
268    // (1) On a strictly fair semaphore, every waiter W up to and including the
269    // first non-bypassable waiter (or every waiter if all are bypassable) satisfies
270    // W.num_permits > permits_available.available(). Releases and cancellations
271    // maintain this via `unblock_waiters_from_front`, which skips bypassable
272    // waiters that lack permits instead of stopping at them.
273    //
274    // (2) W is in waiters iff W.is_queued
275    //
276    // (3) W.is_queued ==> !W.has_permits
277    // Note: the converse is not true.  We can have !W.has_permits && !W.is_queued
278    // when the Acquire is created but not yet polled.
279    //
280    // (4) closed ==> waiters.is_empty()
281    //
282    // (5) if `reservation` is `Some(R)`, then the semaphore is unfair, and
283    // !R.is_queued && !R.has_permits
284    //
285    // (6) closed ==> reservation.is_none()
286    waiters: VecDeque<Arc<Waiter>>,
287    /// The waiter that holds the semaphore's reservation, if any (see
288    /// [`BatchSemaphore::acquire_reserving`]). While it is set, the available
289    /// permits are kept for this waiter: no other request can take one, and the
290    /// waiter takes its `num_permits` as soon as that many are available.
291    reservation: Option<Arc<Waiter>>,
292    permits_available: PermitsAvailable,
293    // TODO: should there be a clock for the close event?
294    closed: bool,
295}
296
297impl BatchSemaphoreState {
298    /// The permits that a request can take now. While a reservation holds the
299    /// semaphore, that is none, except for the holder itself.
300    fn available(&self) -> usize {
301        if self.reservation.is_some() {
302            0
303        } else {
304            self.permits_available.available()
305        }
306    }
307
308    /// Is `waiter` the holder of the semaphore's reservation?
309    fn is_reserved_by(&self, waiter: &Arc<Waiter>) -> bool {
310        self.reservation.as_ref().is_some_and(|r| Arc::ptr_eq(r, waiter))
311    }
312
313    fn acquire_permits(
314        &mut self,
315        num_permits: usize,
316        fairness: Fairness,
317        priority: Priority,
318    ) -> Result<(), TryAcquireError> {
319        assert!(num_permits > 0);
320        if self.closed {
321            Err(TryAcquireError::Closed)
322        } else if self.reservation.is_some() {
323            // The available permits are kept for the holder of the reservation,
324            // which takes them with `take_permits`.
325            Err(TryAcquireError::NoPermits)
326        } else if self
327            .waiters
328            .iter()
329            .all(|waiter| waiter.priority == Priority::Bypassable)
330            || matches!(fairness, Fairness::Unfair)
331            || priority == Priority::Front
332        {
333            // Permits here can be acquired in one of three scenarios:
334            // - The waiter queue is empty or only has bypassable requests,
335            //   so if there are enough permits available, immediately succeed.
336            // - The semaphore is operating in an unfair mode; the current
337            //   thread is either requesting permits for the first time, or it
338            //   was woken and selected by the scheduler. In either case, the
339            //   thread may succeed, as long as there are enough permits.
340            // - The request has `Priority::Front`, so it deliberately overtakes
341            //   the queue (see `BatchSemaphore::upgrade`). Queued waiters hold
342            //   no permits, so they cannot prevent this request from succeeding.
343            self.take_permits(num_permits)
344        } else {
345            Err(TryAcquireError::NoPermits)
346        }
347    }
348
349    /// Take `num_permits` of the available permits for the current task, if
350    /// there are that many, regardless of the waiters and the reservation.
351    fn take_permits(&mut self, num_permits: usize) -> Result<(), TryAcquireError> {
352        let clock = self.permits_available.acquire(num_permits, current::clock())?;
353
354        // If successful, the acquiry is causally dependent on the event
355        // which released the acquired permits.
356        ExecutionState::with(|s| {
357            s.update_clock(&clock);
358        });
359
360        Ok(())
361    }
362
363    /// Unblock the waiters of an unfair semaphore that can now make progress,
364    /// and let them race. While a reservation holds the semaphore, only its
365    /// holder can, once enough permits are available for it.
366    fn wake_unfair_waiters(&mut self) {
367        if let Some(holder) = &self.reservation {
368            // Like a waiter in the queue, a holder whose task has already
369            // finished is stale (see `is_stale`). Drop the reservation, so
370            // that it does not keep the permits from the waiters below. If the
371            // `Acquire` is still alive and another task polls it, it will
372            // reserve or acquire again.
373            if is_stale(holder) {
374                trace!("dropping stale reservation {:?} for finished task", holder);
375                self.reservation = None;
376            } else {
377                if holder.num_permits <= self.permits_available.available() {
378                    ExecutionState::with(|s| s.get_mut(holder.task_id()).unblock());
379                    if let Some(waker) = holder.waker.lock().unwrap().as_ref() {
380                        waker.wake_by_ref();
381                    }
382                }
383                return;
384            }
385        }
386
387        // Unblock all the waiters for which there are enough permits available,
388        // then let them race.
389        let num_available = self.permits_available.available();
390        for waiter in &mut self.waiters {
391            if waiter.min_permits <= num_available {
392                // Unlike the strictly fair case, there is nothing to clean
393                // up for a stale waiter (see `is_stale`): an unfair waiter
394                // holds no permits, so it blocks nobody. But there is also
395                // nobody to unblock.
396                if !unblock_unless_stale(waiter) {
397                    continue;
398                }
399                let maybe_waker = waiter.waker.lock().unwrap();
400                if let Some(waker) = maybe_waker.as_ref() {
401                    waker.wake_by_ref();
402                }
403            }
404        }
405    }
406
407    fn unblock_waiters_from_front(&mut self) {
408        let mut index = 0;
409        while let Some(front) = self.waiters.get(index) {
410            // There is nobody to unblock for a stale waiter (see `is_stale`),
411            // so discard it without consuming permits; if the `Acquire` is
412            // still alive and some other task polls it, it will re-acquire
413            // from the (still available) permits.
414            if is_stale(front) {
415                let waiter = self.waiters.remove(index).unwrap();
416                waiter.is_queued.store(false, Ordering::SeqCst);
417                // Preserve the "queued <=> waker registered" invariant asserted
418                // in `Acquire::poll`; waking a finished task's waker is a no-op.
419                waiter.waker.lock().unwrap().take();
420                trace!("dropping stale waiter {:?} for finished task", waiter);
421                continue;
422            }
423            if front.num_permits <= self.permits_available.available() {
424                let waiter = self.waiters.remove(index).unwrap();
425
426                crate::annotations::record_semaphore_acquire_unblocked(
427                    self.id.unwrap(),
428                    waiter.task_id(),
429                    waiter.num_permits,
430                );
431
432                // The clock we pass into the semaphore is the clock of the
433                // waiter, corresponding to the point at which the waiter was
434                // enqueued. The clock we get in return corresponds to the
435                // join of the clocks of the acquired permits, used to update
436                // the waiter's clock to causally depend on the release events.
437                let clock = self
438                    .permits_available
439                    .acquire(waiter.num_permits, waiter.clock.clone())
440                    .unwrap();
441                trace!("granted {:?} permits to waiter {:?}", waiter.num_permits, waiter);
442
443                // Update waiter state as it is no longer in the queue
444                assert!(waiter.is_queued.swap(false, Ordering::SeqCst));
445                assert!(!waiter.has_permits.swap(true, Ordering::SeqCst));
446                ExecutionState::with(|s| {
447                    let task = s.get_mut(waiter.task_id());
448                    assert!(!task.finished());
449                    // The acquiry is causally dependent on the event
450                    // which released the acquired permits.
451                    task.clock.update(&clock);
452                    task.unblock();
453                });
454                let mut maybe_waker = waiter.waker.lock().unwrap();
455                if let Some(waker) = maybe_waker.take() {
456                    waker.wake();
457                }
458            } else if front.priority == Priority::Bypassable {
459                index += 1;
460            } else {
461                return;
462            }
463        }
464    }
465}
466
467/// Whether `waiter` is stale: the task that registered it has finished, after its `Acquire` future
468/// was cancelled (e.g. a `select!` branch lost, or a `poll_recv`-style API cached the `Acquire`
469/// inside a longer-lived object). If the `Acquire` is still alive, another task can poll it again.
470/// Can't tell outside an execution, and then says no, which preserves the old behaviour.
471#[inline]
472fn is_stale(waiter: &Waiter) -> bool {
473    ExecutionState::try_with(|s| s.try_get(waiter.task_id()).is_some_and(|task| task.finished())).unwrap_or(false)
474}
475
476/// Unblock the task that registered `waiter`, unless the waiter is stale (see `is_stale`). Returns
477/// whether it unblocked the task.
478#[inline]
479fn unblock_unless_stale(waiter: &Waiter) -> bool {
480    ExecutionState::with(|s| {
481        let task = s.get_mut(waiter.task_id());
482        if task.finished() {
483            false
484        } else {
485            task.unblock();
486            true
487        }
488    })
489}
490
491/// Counting semaphore
492#[derive(Debug)]
493pub struct BatchSemaphore {
494    state: RefCell<BatchSemaphoreState>,
495    fairness: Fairness,
496    #[allow(unused)]
497    signature: ResourceSignature,
498}
499
500/// Error returned from the [`BatchSemaphore::try_acquire`] function.
501#[derive(Debug, PartialEq, Eq)]
502pub enum TryAcquireError {
503    /// The semaphore has been closed and cannot issue new permits.
504    Closed,
505
506    /// The semaphore has no available permits.
507    NoPermits,
508}
509
510/// Error returned from the [`BatchSemaphore::acquire`] function.
511///
512/// An `acquire*` operation can only fail if the semaphore has been
513/// closed.
514#[derive(Debug)]
515pub struct AcquireError(());
516
517impl AcquireError {
518    fn closed() -> AcquireError {
519        AcquireError(())
520    }
521}
522
523impl fmt::Display for AcquireError {
524    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
525        write!(fmt, "semaphore closed")
526    }
527}
528
529impl std::error::Error for AcquireError {}
530
531impl BatchSemaphore {
532    /// Creates a new semaphore with the initial number of permits.
533    #[track_caller]
534    pub fn new(num_permits: usize, fairness: Fairness) -> Self {
535        Self::new_with_signature(
536            num_permits,
537            fairness,
538            ExecutionState::new_resource_signature(ResourceType::BatchSemaphore),
539        )
540    }
541
542    pub fn new_with_signature(num_permits: usize, fairness: Fairness, signature: ResourceSignature) -> Self {
543        let state = RefCell::new(BatchSemaphoreState {
544            id: Some(crate::annotations::record_semaphore_created()),
545            waiters: VecDeque::new(),
546            reservation: None,
547            permits_available: PermitsAvailable::new(num_permits),
548            closed: false,
549        });
550        Self {
551            state,
552            fairness,
553            signature,
554        }
555    }
556
557    /// Creates a new semaphore with the initial number of permits.
558    #[track_caller]
559    pub const fn const_new(num_permits: usize, fairness: Fairness) -> Self {
560        Self::const_new_with_signature(
561            num_permits,
562            fairness,
563            ResourceSignature::new_const(ResourceType::BatchSemaphore),
564        )
565    }
566
567    pub const fn const_new_with_signature(
568        num_permits: usize,
569        fairness: Fairness,
570        signature: ResourceSignature,
571    ) -> Self {
572        let state = RefCell::new(BatchSemaphoreState {
573            id: None,
574            waiters: VecDeque::new(),
575            reservation: None,
576            permits_available: PermitsAvailable::const_new(num_permits),
577            closed: false,
578        });
579        Self {
580            state,
581            fairness,
582            signature,
583        }
584    }
585
586    /// Returns the current number of available permits. While a reservation
587    /// holds the semaphore (see [`BatchSemaphore::acquire_reserving`]), this is
588    /// zero: the available permits are kept for the holder.
589    pub fn available_permits(&self) -> usize {
590        let state = self.state.borrow();
591        state.available()
592    }
593
594    fn init_object_id(&self) {
595        let mut state = self.state.borrow_mut();
596        if state.id.is_none() {
597            state.id = Some(crate::annotations::record_semaphore_created());
598        }
599    }
600
601    /// Closes the semaphore. This prevents the semaphore from issuing new
602    /// permits and notifies all pending waiters.
603    pub fn close(&self) {
604        thread::switch();
605        self.close_no_scheduling_point();
606    }
607
608    /// Closes the semaphore without invoking `thread::switch`
609    pub fn close_no_scheduling_point(&self) {
610        self.init_object_id();
611        let mut state = self.state.borrow_mut();
612        if state.closed {
613            return;
614        }
615        crate::annotations::record_semaphore_closed(state.id.unwrap());
616        state.closed = true;
617
618        // Wake up all the waiters, and the holder of the reservation, which waits
619        // too.  Since we've marked the state as closed, they will all return
620        // `AcquireError::closed` from their acquire calls.
621        let ptr = &*state as *const BatchSemaphoreState;
622        let holder = state.reservation.take();
623        let queued = state
624            .waiters
625            .drain(..)
626            .inspect(|waiter| assert!(waiter.is_queued.swap(false, Ordering::SeqCst)));
627        for waiter in queued.chain(holder) {
628            trace!(
629                "semaphore {:p} removing and waking up waiter {:?} on close",
630                ptr,
631                waiter,
632            );
633            assert!(!waiter.has_permits.load(Ordering::SeqCst)); // sanity check
634                                                                 // There is nothing to unblock for a stale waiter (see `is_stale`).
635            unblock_unless_stale(&waiter);
636            let mut maybe_waker = waiter.waker.lock().unwrap();
637            if let Some(waker) = maybe_waker.take() {
638                waker.wake();
639            }
640        }
641    }
642
643    /// Returns true iff the semaphore is closed.
644    pub fn is_closed(&self) -> bool {
645        let state = self.state.borrow();
646        state.closed
647    }
648
649    /// Try to acquire the specified number of permits from the Semaphore.
650    /// If the permits are available, returns Ok(())
651    /// If the semaphore is closed, returns `Err(TryAcquireError::Closed)`
652    /// If there aren't enough permits, returns `Err(TryAcquireError::NoPermits)`
653    pub fn try_acquire(&self, num_permits: usize) -> Result<(), TryAcquireError> {
654        thread::switch();
655
656        self.init_object_id();
657        let mut state = self.state.borrow_mut();
658        let id = state.id.unwrap();
659        let res = state
660            .acquire_permits(num_permits, self.fairness, Priority::Back)
661            .inspect_err(|_err| {
662                // Conservatively, the requester causally depends on the
663                // last successful acquire.
664                // TODO: This is not precise, but `try_acquire` causal dependency
665                // TODO: is both hard to define, and is most likely not worth the
666                // TODO: effort. The cases where causality would be tracked
667                // TODO: "imprecisely" do not correspond to commonly used sync.
668                // TODO: primitives, such as mutexes, mutexes, or condvars.
669                // TODO: An example would be a counting semaphore used to guard
670                // TODO: access to N homogenous resources (as opposed to FIFO,
671                // TODO: heterogenous resources).
672                // TODO: More precision could be gained by tracking clocks for all
673                // TODO: current permit holders, with a data structure similar to
674                // TODO: `permits_available`.
675                ExecutionState::with(|s| {
676                    s.update_clock(&state.permits_available.last_acquire);
677                });
678            });
679        drop(state);
680
681        // If we won the race for permits of an unfair semaphore, re-block
682        // other waiting threads that can no longer succeed.
683        if res.is_ok() {
684            self.reblock_if_unfair();
685        }
686
687        crate::annotations::record_semaphore_try_acquire(id, num_permits, res.is_ok());
688
689        res
690    }
691
692    /// Clean-up method used when a thread succeeds in acquiring permits. If
693    /// the semaphore is unfair, a preceding `release` may have unblocked a
694    /// number of threads, some of which may no longer be able to succeed with
695    /// the permits remaining in the semaphore.
696    fn reblock_if_unfair(&self) {
697        if self.fairness == Fairness::Unfair {
698            let state = self.state.borrow_mut();
699            ExecutionState::with(|s| {
700                let me = s.try_current().map(|task| task.id());
701                for waiter in &state.waiters {
702                    let available = state.permits_available.available();
703                    // A queued waiter cannot make progress while a reservation
704                    // keeps the available permits.
705                    let can_progress = state.reservation.is_none() && waiter.min_permits <= available;
706                    // Skip stale waiters (see `is_stale`): there is nobody to
707                    // block. And skip the current task's own waiters: it is
708                    // running, which an `Acquire` of its that is still queued
709                    // doesn't change, and it would only block itself.
710                    let task = waiter.task_id();
711                    if !can_progress && Some(task) != me && s.try_get(task).is_some_and(|t| !t.finished()) {
712                        // Block this waiter: it cannot succeed (there are not
713                        // enough permits available); its `poll` would return
714                        // without resolving.
715                        s.get_mut(task).block(false);
716                    }
717                }
718            });
719        }
720    }
721
722    fn enqueue_waiter(&self, waiter: &Arc<Waiter>, priority: Priority) {
723        let mut state = self.state.borrow_mut();
724
725        trace!(
726            "enqueuing waiter {:?} ({priority:?}) for semaphore {:p}",
727            waiter,
728            &self.state
729        );
730        match priority {
731            Priority::Back | Priority::Bypassable => state.waiters.push_back(waiter.clone()),
732            // Overtakes the queue rather than joining its tail. Key invariant (1)
733            // still holds: we only get here because the acquire failed, and a
734            // `Priority::Front` acquire only fails when there really aren't
735            // enough permits available, so the new head cannot be grantable.
736            Priority::Front => state.waiters.push_front(waiter.clone()),
737        }
738
739        assert!(!waiter.has_permits.load(Ordering::SeqCst));
740        assert!(!waiter.is_queued.swap(true, Ordering::SeqCst));
741    }
742
743    fn remove_waiter(&self, waiter: &Arc<Waiter>) {
744        let mut state = self.state.borrow_mut();
745
746        trace!(waiters = ?state.waiters, "removing waiter {:?} from semaphore {:p}", waiter, &self.state);
747
748        // sanity checks
749        assert!(!state.closed);
750        assert!(!waiter.has_permits.load(Ordering::SeqCst));
751
752        let index = state
753            .waiters
754            .iter()
755            .position(|x| Arc::ptr_eq(x, waiter))
756            .expect("did not find waiter");
757
758        state.waiters.remove(index).unwrap();
759        assert!(waiter.is_queued.swap(false, Ordering::SeqCst));
760
761        match self.fairness {
762            Fairness::StrictlyFair => {
763                if index == 0
764                    || state
765                        .waiters
766                        .front()
767                        .is_some_and(|w| w.priority == Priority::Bypassable)
768                {
769                    // Removing a waiter can unblock later ones, including when bypassable waiters
770                    // precede it. This can happen in the following situation:
771                    // - the semahore has 1 permit available
772                    // - there are 2 waiters W1 and W2 where W1 wants 2 permits, and W2 wants 1 permit
773                    // - if W1 gives up and drops out, we want to ensure W2 is granted the semaphore
774                    state.unblock_waiters_from_front();
775                }
776            }
777            Fairness::Unfair => {}
778        }
779    }
780
781    /// End the reservation that `waiter` holds, because its `Acquire` was
782    /// dropped before it was granted. The permits that the reservation kept
783    /// were never taken, so they are available again at once.
784    fn cancel_reservation(&self, waiter: &Arc<Waiter>) {
785        let mut state = self.state.borrow_mut();
786
787        trace!("cancelling reservation {:?} of semaphore {:p}", waiter, &self.state);
788
789        assert!(state.is_reserved_by(waiter));
790        state.reservation = None;
791
792        // Wake the waiters that can now take the permits, unless `release`
793        // wouldn't either (see `ExecutionState::should_stop`).
794        let can_wake = ExecutionState::try_with(|s| !s.stops(std::thread::panicking())).unwrap_or(false);
795        if can_wake {
796            state.wake_unfair_waiters();
797        }
798    }
799
800    /// Acquire the specified number of permits (async API)
801    pub fn acquire(&self, num_permits: usize) -> Acquire<'_> {
802        // No switch here; switch should be triggered on polling future
803        self.init_object_id();
804        Acquire::new(self, num_permits, Priority::Back)
805    }
806
807    /// Acquire permits, allowing later requests to pass while this request cannot be satisfied.
808    ///
809    /// On a strictly fair semaphore, this request still waits behind earlier ordinary requests,
810    /// and receives permits in queue order once enough are released. It does not prevent later
811    /// requests from taking permits that it cannot use yet. This lets an upgradable `RwLock`
812    /// reader wait for another upgradable reader without blocking plain readers.
813    ///
814    /// On an unfair semaphore, this is equivalent to [`Self::acquire`].
815    pub fn acquire_bypassable(&self, num_permits: usize) -> Acquire<'_> {
816        self.init_object_id();
817        Acquire::new(self, num_permits, Priority::Bypassable)
818    }
819
820    /// Acquire the specified number of permits (blocking API)
821    pub fn acquire_blocking(&self, num_permits: usize) -> Result<(), AcquireError> {
822        crate::future::block_on(self.acquire(num_permits))
823    }
824
825    /// Acquire `num_permits` permits, and reserve the semaphore for this request
826    /// as soon as at least `min_permits` permits are available (async API). Only
827    /// an unfair semaphore supports this.
828    ///
829    /// Until `min_permits` permits are available, the request waits like one
830    /// from [`BatchSemaphore::acquire`]: it holds nothing and stops no other
831    /// request. As soon as they are, it reserves the semaphore in the same step.
832    /// From then on, no other request can take a permit, and this request takes
833    /// its `num_permits` as soon as that many are available. The reservation
834    /// ends when the request is granted, or when the returned future is dropped.
835    /// While it lasts, [`BatchSemaphore::available_permits`] is zero.
836    ///
837    /// The motivating use case is a `parking_lot` `RwLock` writer. `parking_lot`
838    /// sets `WRITER_BIT` only when no writer or upgradable reader holds the lock,
839    /// and from then on, the bit stops new readers while the writer waits for
840    /// the current ones to leave. With `min_permits` above the permits that are
841    /// left while an upgradable reader holds the lock, the reservation is that
842    /// bit.
843    ///
844    /// At most one request can hold the reservation. Another reserving request
845    /// waits like any other until the reservation ends.
846    ///
847    /// # Panics
848    ///
849    /// Panics if the semaphore is strictly fair (its queue already keeps the
850    /// permits for its first waiter), if `num_permits` is zero, or if
851    /// `min_permits > num_permits`.
852    pub fn acquire_reserving(&self, min_permits: usize, num_permits: usize) -> Acquire<'_> {
853        assert_eq!(
854            self.fairness,
855            Fairness::Unfair,
856            "only an unfair semaphore supports reservations"
857        );
858        assert!(num_permits > 0);
859        assert!(min_permits <= num_permits);
860
861        self.init_object_id();
862        Acquire::new_reserving(self, num_permits, min_permits)
863    }
864
865    /// Release `num_permits` back to the Semaphore
866    pub fn release(&self, num_permits: usize) {
867        // Execution teardown can unwind a task's stack from this scheduling point, which is often in
868        // a destructor that releases a lock (see `ExecutionState::tear_down`). The permits must not
869        // be lost then, as destructors that run later can need them.
870        struct ReleaseOnUnwind<'a>(&'a BatchSemaphore, usize);
871        impl Drop for ReleaseOnUnwind<'_> {
872            fn drop(&mut self) {
873                self.0.release_no_scheduling_point(self.1);
874            }
875        }
876        let release_on_unwind = ReleaseOnUnwind(self, num_permits);
877        thread::switch();
878        std::mem::forget(release_on_unwind);
879
880        self.release_no_scheduling_point(num_permits);
881    }
882
883    /// `release` without its scheduling point.
884    #[inline]
885    fn release_no_scheduling_point(&self, num_permits: usize) {
886        self.init_object_id();
887        if num_permits == 0 {
888            return;
889        }
890
891        let mut state = self.state.borrow_mut();
892
893        crate::annotations::record_semaphore_release(state.id.unwrap(), num_permits);
894
895        if ExecutionState::should_stop() {
896            // In case we are panicking, we release permits, but also clear
897            // the waiters queue: we should not unblock the threads at this
898            // point. However, the permits are released such that future
899            // acquires may succeed, as long as the requesters were not
900            // blocking on the semaphore at the time of the panic. This is
901            // used to correctly model lock poisoning.
902            state.permits_available.release(num_permits, VectorClock::new());
903            for waiter in &state.waiters {
904                waiter.is_queued.swap(false, Ordering::SeqCst);
905            }
906            state.waiters.clear();
907            state.reservation = None;
908            state.closed = true;
909            return;
910        }
911
912        // Permits released into the semaphore reflect the releasing thread's
913        // clock; future acquires of those permits are causally dependent on
914        // this event.
915        ExecutionState::with(|s| {
916            let clock = s.increment_clock();
917            state.permits_available.release(num_permits, clock.clone());
918        });
919
920        // `ExecutionState::me()` is only wanted for this trace, so let the macro's
921        // level check decide whether to pay for it. Computing it eagerly cost an
922        // `ExecutionState::with` on every release even with tracing disabled.
923        trace!(task = ?ExecutionState::me(), avail = ?state.permits_available, waiters = ?state.waiters, "released {} permits for semaphore {:p}", num_permits, &self.state);
924
925        match self.fairness {
926            Fairness::StrictlyFair => {
927                // in a strictly fair mode we will grant permits to waiters from the front
928                // of the queue, as long as there are enough permits available
929                state.unblock_waiters_from_front();
930            }
931            Fairness::Unfair => {
932                // in an unfair mode, we will unblock all the waiters for which
933                // there are enough permits available, then let them race
934                state.wake_unfair_waiters();
935            }
936        }
937        drop(state);
938    }
939
940    /// Atomically `upgrade` from holding `permits_currently_held` permits to holding
941    /// `permits_to_be_held`, without ever dropping below `permits_currently_held` in between.
942    /// The motivating use case is `parking_lot`'s `RwLockUpgradableReadGuard::upgrade`, which must
943    /// take a read guard to a write guard without letting any writer in along the way.
944    ///
945    /// This is implemented by acquiring only the *missing* permits
946    /// (`permits_to_be_held - permits_currently_held`), with priority over any waiter already
947    /// queued, so that the request overtakes the queue. Both halves of that matter:
948    ///
949    /// * Keeping the held permits means no other task can claim the resource mid-upgrade. Releasing
950    ///   them first (even for an instant) would hand the resource to a queued waiter, which for an
951    ///   `RwLock` means a writer mutating the data an upgradable reader had already observed.
952    /// * Overtaking the queue is what makes that safe rather than deadlock-prone. Since we hold
953    ///   permits we will not release, a queued waiter ahead of us may be unsatisfiable (an `RwLock`
954    ///   writer wants *all* permits), so waiting our turn behind it could deadlock. Queued waiters
955    ///   hold no permits, so overtaking them costs nothing but their place in line -- which is
956    ///   exactly the priority a real upgradable read lock gives an upgrade.
957    ///
958    /// The upgrade therefore blocks only on tasks that *currently hold* permits, and is granted as
959    /// soon as they release. The returned future must be driven to completion; if it is dropped
960    /// first, the caller still holds `permits_currently_held`.
961    ///
962    /// An unfair semaphore has no queue to overtake, so there the upgrade reserves the semaphore
963    /// instead (see [`BatchSemaphore::acquire_reserving`]), at once unless another request holds
964    /// the reservation. From then on, no other request can take a permit, so the upgrade again
965    /// waits only for the tasks that hold permits, and nothing can overtake it.
966    ///
967    /// At most one `upgrade` may be in flight on a semaphore at a time. Two concurrent upgraders
968    /// could each be waiting for permits the other holds, which no queue discipline can resolve.
969    /// Callers are expected to enforce this (an `RwLock` does: there is only ever one upgradable
970    /// reader).
971    pub fn upgrade(&self, permits_currently_held: usize, permits_to_be_held: usize) -> Acquire<'_> {
972        assert!(permits_currently_held > 0);
973        assert!(permits_to_be_held > permits_currently_held);
974
975        self.init_object_id();
976        let num_permits = permits_to_be_held - permits_currently_held;
977        match self.fairness {
978            Fairness::StrictlyFair => Acquire::new(self, num_permits, Priority::Front),
979            Fairness::Unfair => Acquire::new_reserving(self, num_permits, 0),
980        }
981    }
982
983    /// The non-blocking analogue of [`BatchSemaphore::upgrade`]: succeeds only if the missing
984    /// permits are available right now, and never blocks or queues.
985    ///
986    /// Like `upgrade`, this ignores queued waiters (they hold no permits, so they cannot be the
987    /// reason the upgrade is short of permits). A `try_upgrade` therefore fails only when some
988    /// other task actually *holds* permits the upgrade needs.
989    pub fn try_upgrade(&self, permits_currently_held: usize, permits_to_be_held: usize) -> Result<(), TryAcquireError> {
990        assert!(permits_currently_held > 0);
991        assert!(permits_to_be_held > permits_currently_held);
992
993        thread::switch();
994
995        self.init_object_id();
996        let num_permits = permits_to_be_held - permits_currently_held;
997        let mut state = self.state.borrow_mut();
998        let id = state.id.unwrap();
999        let res = state
1000            .acquire_permits(num_permits, self.fairness, Priority::Front)
1001            .inspect_err(|_err| {
1002                // Conservatively, the requester causally depends on the last successful acquire;
1003                // see the equivalent reasoning in `try_acquire`.
1004                ExecutionState::with(|s| {
1005                    s.update_clock(&state.permits_available.last_acquire);
1006                });
1007            });
1008        drop(state);
1009
1010        // If we took permits from an unfair semaphore, re-block waiting threads that can no longer
1011        // succeed.
1012        if res.is_ok() {
1013            self.reblock_if_unfair();
1014        }
1015
1016        crate::annotations::record_semaphore_try_acquire(id, num_permits, res.is_ok());
1017
1018        res
1019    }
1020}
1021
1022// Safety: Semaphore is never actually passed across true threads, only across continuations. The
1023// RefCell<_> type therefore can't be preempted mid-bookkeeping-operation.
1024// TODO we shouldn't need to do this, but RefCell is not Send, and anything we put within a Semaphore
1025// TODO needs to be Send.
1026unsafe impl Send for BatchSemaphore {}
1027unsafe impl Sync for BatchSemaphore {}
1028
1029impl Default for BatchSemaphore {
1030    #[track_caller]
1031    fn default() -> Self {
1032        Self::new(Default::default(), Fairness::StrictlyFair)
1033    }
1034}
1035
1036/// The future that results from async calls to `acquire*`.
1037/// Callers must `await` on this future to obtain the necessary permits.
1038pub struct Acquire<'a> {
1039    semaphore: &'a BatchSemaphore,
1040    num_permits: usize,
1041
1042    /// Where this acquire sits relative to waiters already queued on a fair
1043    /// semaphore. Only [`BatchSemaphore::upgrade`] uses [`Priority::Front`]; see
1044    /// there for why an upgrade must overtake the queue.
1045    priority: Priority,
1046
1047    /// For a reserving acquire, the number of available permits at which it
1048    /// reserves the semaphore (see [`BatchSemaphore::acquire_reserving`]).
1049    /// `None` for every other acquire.
1050    reserve_at: Option<usize>,
1051
1052    /// Snapshotted when this `Acquire` is created, and moved into the `Waiter` if
1053    /// this acquire ends up blocking. See `Waiter::new` for why the snapshot must
1054    /// happen here rather than at enqueue time.
1055    clock: VectorClock,
1056
1057    /// The shared part of this acquire, allocated only once the acquire has to
1058    /// block. An acquire that gets its permits immediately is never visible to
1059    /// any other task, so it needs no shared state and no allocation. While this
1060    /// is `None`, `has_permits` below is authoritative.
1061    waiter: Option<Arc<Waiter>>,
1062
1063    /// Whether permits have been granted, for the case where no `Waiter` exists.
1064    /// Once one does, the releasing task writes `Waiter::has_permits` instead and
1065    /// this field is unused; read through `Acquire::has_permits`.
1066    has_permits: bool,
1067
1068    completed: bool, // Has the future completed yet?
1069    never_polled: bool,
1070}
1071
1072// Implement Debug in order to not output the `VectorClock`, matching `Waiter`.
1073impl fmt::Debug for Acquire<'_> {
1074    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1075        f.debug_struct("Acquire")
1076            .field("num_permits", &self.num_permits)
1077            .field("priority", &self.priority)
1078            .field("reserve_at", &self.reserve_at)
1079            .field("waiter", &self.waiter)
1080            .field("has_permits", &self.has_permits())
1081            .field("completed", &self.completed)
1082            .finish()
1083    }
1084}
1085
1086impl<'a> Acquire<'a> {
1087    fn new(semaphore: &'a BatchSemaphore, num_permits: usize, priority: Priority) -> Self {
1088        Self {
1089            semaphore,
1090            num_permits,
1091            priority,
1092            reserve_at: None,
1093            clock: current::clock(),
1094            waiter: None,
1095            has_permits: false,
1096            completed: false,
1097            never_polled: true,
1098        }
1099    }
1100
1101    fn new_reserving(semaphore: &'a BatchSemaphore, num_permits: usize, min_permits: usize) -> Self {
1102        let mut acquire = Self::new(semaphore, num_permits, Priority::Back);
1103        acquire.reserve_at = Some(min_permits);
1104        acquire
1105    }
1106
1107    /// How many permits must be available for this acquire to make progress
1108    /// (see `Waiter::min_permits`).
1109    fn min_permits(&self) -> usize {
1110        self.reserve_at.unwrap_or(self.num_permits)
1111    }
1112
1113    /// Does this acquire hold the semaphore's reservation?
1114    fn is_reserving(&self) -> bool {
1115        self.waiter
1116            .as_ref()
1117            .is_some_and(|waiter| self.semaphore.state.borrow().is_reserved_by(waiter))
1118    }
1119
1120    /// Have permits been granted to this acquire? Once a `Waiter` exists the
1121    /// releasing task owns that flag, so the shared copy is authoritative.
1122    fn has_permits(&self) -> bool {
1123        match &self.waiter {
1124            Some(waiter) => waiter.has_permits.load(Ordering::SeqCst),
1125            None => self.has_permits,
1126        }
1127    }
1128
1129    /// Is this acquire in the semaphore's waiter queue? Only possible once a
1130    /// `Waiter` has been allocated, since the queue holds `Arc<Waiter>`.
1131    fn is_queued(&self) -> bool {
1132        match &self.waiter {
1133            Some(waiter) => waiter.is_queued.load(Ordering::SeqCst),
1134            None => false,
1135        }
1136    }
1137
1138    fn grant_permits(&mut self) {
1139        match &self.waiter {
1140            Some(waiter) => waiter.has_permits.store(true, Ordering::SeqCst),
1141            None => self.has_permits = true,
1142        }
1143    }
1144
1145    /// The shared `Waiter` for this acquire, allocating it if this is the first
1146    /// time the acquire has had to block. Returns an owned handle so callers can
1147    /// still use `self.semaphore` without holding a borrow of `self`.
1148    fn waiter_for_blocking(&mut self) -> Arc<Waiter> {
1149        if let Some(waiter) = &self.waiter {
1150            return Arc::clone(waiter);
1151        }
1152        let waiter = Arc::new(Waiter::new(
1153            self.num_permits,
1154            self.min_permits(),
1155            self.priority,
1156            self.clock.clone(),
1157        ));
1158        self.waiter = Some(Arc::clone(&waiter));
1159        waiter
1160    }
1161
1162    /// The part of `poll` for a reserving acquire (see
1163    /// [`BatchSemaphore::acquire_reserving`]), once `poll` knows that it has no
1164    /// permits yet and that the semaphore is open. Only an unfair semaphore has
1165    /// reserving acquires.
1166    fn poll_reserving(&mut self, min_permits: usize, cx: &mut Context<'_>) -> Poll<Result<(), AcquireError>> {
1167        let semaphore = self.semaphore;
1168        let is_queued = self.is_queued();
1169        let is_reserving = self.is_reserving();
1170        trace!(
1171            "Acquire::poll for reserving {:?}; is queued: {is_queued:?}, is reserving: {is_reserving:?}",
1172            self
1173        );
1174
1175        let mut state = semaphore.state.borrow_mut();
1176        let id = state.id.unwrap();
1177        let available = state.permits_available.available();
1178
1179        if is_reserving {
1180            if available < self.num_permits {
1181                // Still waiting for the tasks that hold the rest. Like a queued
1182                // waiter, follow the current poller.
1183                drop(state);
1184                let waiter = self.waiter_for_blocking();
1185                *waiter.waker.lock().unwrap() = Some(cx.waker().clone());
1186                waiter.set_task_id(ExecutionState::me());
1187                return Poll::Pending;
1188            }
1189
1190            // The reservation kept the permits for us, so take them.
1191            state.reservation = None;
1192            state.take_permits(self.num_permits).unwrap();
1193            // Let the waiters race for any permits that are left.
1194            state.wake_unfair_waiters();
1195            drop(state);
1196
1197            let waiter = self
1198                .waiter
1199                .clone()
1200                .expect("a reserving acquire must have an allocated waiter");
1201            crate::annotations::record_semaphore_acquire_unblocked(id, waiter.task_id(), self.num_permits);
1202            self.grant_permits();
1203            self.completed = true;
1204            trace!("Acquire::poll for {:?} that got permits", self);
1205            return Poll::Ready(Ok(()));
1206        }
1207
1208        if state.reservation.is_some() || available < min_permits {
1209            // Wait, holding nothing, like any other waiter of an unfair
1210            // semaphore.
1211            drop(state);
1212            let waiter = self.waiter_for_blocking();
1213            *waiter.waker.lock().unwrap() = Some(cx.waker().clone());
1214            waiter.set_task_id(ExecutionState::me());
1215            if !is_queued {
1216                crate::annotations::record_semaphore_acquire_blocked(id, self.num_permits);
1217                semaphore.enqueue_waiter(&waiter, Priority::Back);
1218            }
1219            trace!("Acquire::poll for {:?} that is enqueued", self);
1220            return Poll::Pending;
1221        }
1222
1223        if available >= self.num_permits {
1224            // There are enough permits, so there is nothing to reserve.
1225            state.take_permits(self.num_permits).unwrap();
1226            drop(state);
1227            if is_queued {
1228                let waiter = self
1229                    .waiter
1230                    .clone()
1231                    .expect("a queued acquire must have an allocated waiter");
1232                crate::annotations::record_semaphore_acquire_unblocked(id, waiter.task_id(), self.num_permits);
1233                semaphore.remove_waiter(&waiter);
1234            } else {
1235                crate::annotations::record_semaphore_acquire_fast(id, self.num_permits);
1236            }
1237            self.grant_permits();
1238            self.completed = true;
1239            trace!("Acquire::poll for {:?} that got permits", self);
1240            semaphore.reblock_if_unfair();
1241            return Poll::Ready(Ok(()));
1242        }
1243
1244        // Reserve the semaphore, and wait for the rest of the permits.
1245        drop(state);
1246        let waiter = self.waiter_for_blocking();
1247        *waiter.waker.lock().unwrap() = Some(cx.waker().clone());
1248        waiter.set_task_id(ExecutionState::me());
1249        if is_queued {
1250            semaphore.remove_waiter(&waiter);
1251        } else {
1252            crate::annotations::record_semaphore_acquire_blocked(id, self.num_permits);
1253        }
1254        semaphore.state.borrow_mut().reservation = Some(waiter);
1255        trace!("Acquire::poll for {:?} that reserved the semaphore", self);
1256        // No waiter can take a permit now.
1257        semaphore.reblock_if_unfair();
1258        Poll::Pending
1259    }
1260}
1261
1262impl Future for Acquire<'_> {
1263    type Output = Result<(), AcquireError>;
1264
1265    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1266        assert!(!self.completed);
1267
1268        // One borrow of the semaphore state rather than two (`is_closed` and
1269        // `available_permits` each took their own). Both reads describe the same
1270        // instant, before the scheduling point below, so merging them is sound.
1271        // Reads *after* the switch must stay separate and fresh, because other
1272        // tasks may have run in between. A reserving acquire that will reserve
1273        // the semaphore changes its state as much as one that will succeed, so
1274        // it compares against the permits at which it reserves.
1275        let will_succeed = self.has_permits() || {
1276            let state = self.semaphore.state.borrow();
1277            state.closed || state.available() >= self.min_permits()
1278        };
1279
1280        // If the acquire will succeed on the first try, we need to context switch once to allow the previous
1281        // event to become visible. If we won't succeed, then we still need to context switch if the act of
1282        // blocking does not commute with other operations on `batch_semaphore` (double-yield optimization,
1283        // reasoning below).
1284        //
1285        // Fair Semaphores: blocking adds the current task to an *ordered* waiter queue. Two blocking acquires
1286        // *do not commute* because in one ordering the queue will be [T1 T2] and in the other ordering [T2 T1].
1287        // Thus we cannot apply the double-yield optimization for fair semaphores.
1288        //
1289        // Unfair Semaphores: blocking adds the current task to an *unordered set* of waiters. To check if the
1290        // double-yield is valid we check if each operation (Z) on the semaphore commutes with a blocking acquire (Y1):
1291        //
1292        //     - Blocking Acquire: in both orderings `Z Y1` and `Y1 Z`, the waiter set has the same members, thus
1293        //       the operations commute.
1294        //     - Try Acquire: the try-acquire will fail in both orderings without changing the state of the semaphore
1295        //     - Release: if the release unblocks Y1, then the optimization is not applicable. Otherwise, it must
1296        //       unblock another task in the waiter set. As waiter-set insertion and removal for disjoint elements
1297        //       commutes, release operations also commute in this case.
1298        //
1299        // Thus we apply the double-yield optimization for *unfair* semaphores only
1300        let blocking_is_not_commutative = self.semaphore.fairness == Fairness::StrictlyFair;
1301
1302        if self.never_polled && (will_succeed || blocking_is_not_commutative) {
1303            thread::switch();
1304        }
1305        self.never_polled = false;
1306
1307        let out = if self.has_permits() {
1308            assert!(!self.is_queued());
1309            self.completed = true;
1310            trace!("Acquire::poll for {:?} with permits", self);
1311            Poll::Ready(Ok(()))
1312        } else if self.semaphore.is_closed() {
1313            assert!(!self.is_queued());
1314            self.completed = true;
1315            trace!("Acquire::poll for {:?} with closed", self);
1316            Poll::Ready(Err(AcquireError::closed()))
1317        } else if let Some(min_permits) = self.reserve_at {
1318            self.poll_reserving(min_permits, cx)
1319        } else {
1320            let is_queued = self.is_queued();
1321            trace!("Acquire::poll for {:?}; is queued: {is_queued:?}", self);
1322
1323            // Sanity check: there should be a waker if the waiter is in
1324            // the queue. Also true for unfair semaphores, which wake by ref.
1325            //
1326            // `debug_assert` rather than `assert`: this takes a `std::sync::Mutex`
1327            // on every poll, including the uncontended fast path, purely to check
1328            // an internal invariant.
1329            debug_assert_eq!(
1330                is_queued,
1331                self.waiter
1332                    .as_ref()
1333                    .is_some_and(|waiter| waiter.waker.lock().unwrap().is_some())
1334            );
1335
1336            // Should the waiter try to acquire permits here? Four cases:
1337            // 1. unfair semaphore, waiter not yet enqueued;
1338            // 2. fair semaphore, waiter not yet enqueued;
1339            // 3. unfair semaphore, waiter already enqueued.
1340            // 4. fair semaphore, waiter already enqueued;
1341            //
1342            // 1. and 2. are similar: the future was polled for the first time,
1343            // so the waiter will try to acquire some permits. If successful,
1344            // the waiter need not be enqueued, and the future is resolved.
1345            // Otherwise, the waiter is added to the queue.
1346            //
1347            // 3. is slightly different: the future was polled, even though the
1348            // waiter was already in the queue. This can happen either because
1349            // the semaphore just received some permits and woke the waiter up,
1350            // or because the future itself was polled manually. Either way,
1351            // the semaphore is queried.
1352            //
1353            // 4. is a case where we do not try to acquire permits. The request
1354            // would always fail, and the waiter should remain suspended until
1355            // the semaphore has explicitly unblocked it and given it permits
1356            // during a `release` call.
1357            let try_to_acquire = match (self.semaphore.fairness, is_queued) {
1358                // written this way to mirror the cases described above
1359                (Fairness::Unfair, false) | (Fairness::StrictlyFair, false) | (Fairness::Unfair, true) => true,
1360                (Fairness::StrictlyFair, true) => false,
1361            };
1362
1363            if try_to_acquire {
1364                // Access the semaphore state directly instead of `try_acquire`,
1365                // because in case of `NoPermits`, we do not want to update the
1366                // clock, as this thread will be blocked below.
1367                let mut state = self.semaphore.state.borrow_mut();
1368                let id = state.id.unwrap();
1369                let acquire_result = state.acquire_permits(self.num_permits, self.semaphore.fairness, self.priority);
1370                drop(state);
1371
1372                match acquire_result {
1373                    Ok(()) => {
1374                        if is_queued {
1375                            let waiter = self
1376                                .waiter
1377                                .clone()
1378                                .expect("a queued acquire must have an allocated waiter");
1379                            crate::annotations::record_semaphore_acquire_unblocked(
1380                                id,
1381                                waiter.task_id(),
1382                                waiter.num_permits,
1383                            );
1384                            self.semaphore.remove_waiter(&waiter);
1385                        } else {
1386                            crate::annotations::record_semaphore_acquire_fast(id, self.num_permits);
1387                        }
1388                        self.grant_permits();
1389                        self.completed = true;
1390                        trace!("Acquire::poll for {:?} that got permits", self);
1391
1392                        // If the semaphore is unfair, re-block other waiting
1393                        // threads that can no longer succeed.
1394                        self.semaphore.reblock_if_unfair();
1395
1396                        Poll::Ready(Ok(()))
1397                    }
1398                    Err(TryAcquireError::NoPermits) => {
1399                        // This acquire has to block, so it now becomes visible to
1400                        // whichever task releases permits. That is the first point
1401                        // at which shared state is needed, so it is where the
1402                        // `Waiter` gets allocated.
1403                        let waiter = self.waiter_for_blocking();
1404
1405                        let mut maybe_waker = waiter.waker.lock().unwrap();
1406                        *maybe_waker = Some(cx.waker().clone());
1407                        drop(maybe_waker);
1408
1409                        // Point the waiter at whoever is polling now: this future
1410                        // may have been created by a different task.
1411                        waiter.set_task_id(ExecutionState::me());
1412
1413                        if !is_queued {
1414                            crate::annotations::record_semaphore_acquire_blocked(id, self.num_permits);
1415                            // `enqueue_waiter` sets `is_queued` itself.
1416                            self.semaphore.enqueue_waiter(&waiter, self.priority);
1417                        }
1418                        trace!("Acquire::poll for {:?} that is enqueued", self);
1419                        Poll::Pending
1420                    }
1421                    Err(TryAcquireError::Closed) => unreachable!(),
1422                }
1423            } else {
1424                // No progress made, future is still pending. The waiter stays in
1425                // the queue, but re-point it at the current poller and refresh
1426                // its waker: this future may have been created by (or last
1427                // polled by) another task, and `release` must wake whoever is
1428                // waiting now. Without this, a permit granted to this waiter
1429                // would unblock a task that is no longer interested, and the
1430                // actual poller would never be woken.
1431                let waiter = self
1432                    .waiter
1433                    .as_ref()
1434                    .expect("a queued acquire must have an allocated waiter");
1435                *waiter.waker.lock().unwrap() = Some(cx.waker().clone());
1436                waiter.set_task_id(ExecutionState::me());
1437                Poll::Pending
1438            }
1439        };
1440        if matches!(out, Poll::Pending) {
1441            // `Backtrace::capture()` is a noop (it returns the constant `disabled()`) if `RUST_BACKTRACE`/`RUST_LIB_BACKTRACE` is not set.
1442            ExecutionState::with(|state| {
1443                state.current_mut().backtrace = if backtrace_enabled() {
1444                    Some(std::backtrace::Backtrace::force_capture())
1445                } else {
1446                    None
1447                }
1448            })
1449        }
1450        out
1451    }
1452}
1453
1454impl Drop for Acquire<'_> {
1455    fn drop(&mut self) {
1456        trace!("Acquire::drop for {:?}", self);
1457        if self.is_queued() {
1458            // If the associated waiter is in the wait list, remove it
1459            let waiter = self
1460                .waiter
1461                .clone()
1462                .expect("a queued acquire must have an allocated waiter");
1463            self.semaphore.remove_waiter(&waiter);
1464        } else if self.is_reserving() {
1465            // If the acquire holds the reservation, end it, so that it does not
1466            // keep the permits from other requests.
1467            let waiter = self
1468                .waiter
1469                .clone()
1470                .expect("a reserving acquire must have an allocated waiter");
1471            self.semaphore.cancel_reservation(&waiter);
1472        } else if self.has_permits() && !self.completed {
1473            // If the waiter was granted permits, release them. Note this must also
1474            // fire for an acquire that got its permits without ever allocating a
1475            // waiter, otherwise the semaphore leaks permits.
1476            self.semaphore.release(self.num_permits);
1477        }
1478    }
1479}
1480
1481impl crate::annotations::WithName for &BatchSemaphore {
1482    fn with_name_and_kind(self, name: Option<&str>, kind: Option<&str>) -> Self {
1483        self.init_object_id();
1484        crate::annotations::record_name_for_object(self.state.borrow().id.unwrap(), name, kind);
1485        self
1486    }
1487}
1488
1489impl crate::annotations::WithName for BatchSemaphore {
1490    fn with_name_and_kind(self, name: Option<&str>, kind: Option<&str>) -> Self {
1491        (&self).with_name_and_kind(name, kind);
1492        self
1493    }
1494}
1495
1496impl BatchSemaphore {
1497    /// Returns a reference to this semaphore's resource signature.
1498    pub fn signature(&self) -> &ResourceSignature {
1499        &self.signature
1500    }
1501}