Skip to main content

tokio/sync/
batch_semaphore.rs

1#![cfg_attr(not(feature = "sync"), allow(unreachable_pub, dead_code))]
2//! # Implementation Details.
3//!
4//! The semaphore is implemented using an intrusive linked list of waiters. An
5//! atomic counter tracks the number of available permits. If the semaphore does
6//! not contain the required number of permits, the task attempting to acquire
7//! permits places its waker at the end of a queue. When new permits are made
8//! available (such as by releasing an initial acquisition), they are assigned
9//! to the task at the front of the queue, waking that task if its requested
10//! number of permits is met.
11//!
12//! Because waiters are enqueued at the back of the linked list and dequeued
13//! from the front, the semaphore is fair. Tasks trying to acquire large numbers
14//! of permits at a time will always be woken eventually, even if many other
15//! tasks are acquiring smaller numbers of permits. This means that in a
16//! use-case like tokio's read-write lock, writers will not be starved by
17//! readers.
18use crate::loom::cell::UnsafeCell;
19use crate::loom::sync::atomic::AtomicUsize;
20use crate::loom::sync::{Mutex, MutexGuard};
21use crate::util::linked_list::{self, LinkedList};
22#[cfg(all(tokio_unstable, feature = "tracing"))]
23use crate::util::trace;
24use crate::util::WakeList;
25
26use std::future::Future;
27use std::marker::PhantomPinned;
28use std::pin::Pin;
29use std::ptr::NonNull;
30use std::sync::atomic::Ordering::*;
31use std::task::{ready, Context, Poll, Waker};
32use std::{cmp, fmt};
33
34/// An asynchronous counting semaphore which permits waiting on multiple permits at once.
35pub(crate) struct Semaphore {
36    waiters: Mutex<Waitlist>,
37    /// The current number of available permits in the semaphore.
38    permits: AtomicUsize,
39    #[cfg(all(tokio_unstable, feature = "tracing"))]
40    resource_span: tracing::Span,
41}
42
43struct Waitlist {
44    queue: LinkedList<Waiter>,
45    closed: bool,
46}
47
48/// Error returned from the [`Semaphore::try_acquire`] function.
49///
50/// [`Semaphore::try_acquire`]: crate::sync::Semaphore::try_acquire
51#[derive(Debug, PartialEq, Eq)]
52pub enum TryAcquireError {
53    /// The semaphore has been [closed] and cannot issue new permits.
54    ///
55    /// [closed]: crate::sync::Semaphore::close
56    Closed,
57
58    /// The semaphore has no available permits.
59    NoPermits,
60}
61/// Error returned from the [`Semaphore::acquire`] function.
62///
63/// An `acquire` operation can only fail if the semaphore has been
64/// [closed].
65///
66/// [closed]: crate::sync::Semaphore::close
67/// [`Semaphore::acquire`]: crate::sync::Semaphore::acquire
68#[derive(Debug)]
69pub struct AcquireError(());
70
71pub(crate) struct Acquire<'a> {
72    node: Waiter,
73    semaphore: &'a Semaphore,
74    num_permits: usize,
75    /// Whether `Acquire::drop` must clean up after the waiter, by removing it
76    /// from the wait queue if it is still linked and returning the permits
77    /// assigned to it to the semaphore.
78    ///
79    /// This is set as soon as `poll_acquire` assigns permits to the waiter or
80    /// links it into the wait queue, and it is only cleared once the future
81    /// completes and the permits are handed to the caller.
82    queued: bool,
83}
84
85/// An entry in the wait queue.
86struct Waiter {
87    /// The current state of the waiter.
88    ///
89    /// This is either the number of remaining permits required by
90    /// the waiter, or a flag indicating that the waiter is not yet queued.
91    state: AtomicUsize,
92
93    /// The waker to notify the task awaiting permits.
94    ///
95    /// # Safety
96    ///
97    /// This may only be accessed while the wait queue is locked.
98    waker: UnsafeCell<Option<Waker>>,
99
100    /// Intrusive linked-list pointers.
101    ///
102    /// # Safety
103    ///
104    /// This may only be accessed while the wait queue is locked.
105    ///
106    /// TODO: Ideally, we would be able to use loom to enforce that
107    /// this isn't accessed concurrently. However, it is difficult to
108    /// use a `UnsafeCell` here, since the `Link` trait requires _returning_
109    /// references to `Pointers`, and `UnsafeCell` requires that checked access
110    /// take place inside a closure. We should consider changing `Pointers` to
111    /// use `UnsafeCell` internally.
112    pointers: linked_list::Pointers<Waiter>,
113
114    #[cfg(all(tokio_unstable, feature = "tracing"))]
115    ctx: trace::AsyncOpTracingCtx,
116
117    /// Should not be `Unpin`.
118    _p: PhantomPinned,
119}
120
121generate_addr_of_methods! {
122    impl<> Waiter {
123        unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> {
124            &self.pointers
125        }
126    }
127}
128
129impl Semaphore {
130    /// The maximum number of permits which a semaphore can hold.
131    ///
132    /// Note that this reserves three bits of flags in the permit counter, but
133    /// we only actually use one of them. However, the previous semaphore
134    /// implementation used three bits, so we will continue to reserve them to
135    /// avoid a breaking change if additional flags need to be added in the
136    /// future.
137    pub(crate) const MAX_PERMITS: usize = usize::MAX >> 3;
138    const CLOSED: usize = 1;
139    // The least-significant bit in the number of permits is reserved to use
140    // as a flag indicating that the semaphore has been closed. Consequently
141    // PERMIT_SHIFT is used to leave that bit for that purpose.
142    const PERMIT_SHIFT: usize = 1;
143
144    /// Creates a new semaphore with the initial number of permits
145    ///
146    /// Maximum number of permits on 32-bit platforms is `1<<29`.
147    pub(crate) fn new(permits: usize) -> Self {
148        assert!(
149            permits <= Self::MAX_PERMITS,
150            "a semaphore may not have more than MAX_PERMITS permits ({})",
151            Self::MAX_PERMITS
152        );
153
154        #[cfg(all(tokio_unstable, feature = "tracing"))]
155        let resource_span = {
156            let resource_span = tracing::trace_span!(
157                parent: None,
158                "runtime.resource",
159                concrete_type = "Semaphore",
160                kind = "Sync",
161                is_internal = true
162            );
163
164            resource_span.in_scope(|| {
165                tracing::trace!(
166                    target: "runtime::resource::state_update",
167                    permits = permits,
168                    permits.op = "override",
169                )
170            });
171            resource_span
172        };
173
174        Self {
175            permits: AtomicUsize::new(permits << Self::PERMIT_SHIFT),
176            waiters: Mutex::new(Waitlist {
177                queue: LinkedList::new(),
178                closed: false,
179            }),
180            #[cfg(all(tokio_unstable, feature = "tracing"))]
181            resource_span,
182        }
183    }
184
185    /// Creates a new semaphore with the initial number of permits.
186    ///
187    /// Maximum number of permits on 32-bit platforms is `1<<29`.
188    #[cfg(not(all(loom, test)))]
189    pub(crate) const fn const_new(permits: usize) -> Self {
190        assert!(permits <= Self::MAX_PERMITS);
191
192        Self {
193            permits: AtomicUsize::new(permits << Self::PERMIT_SHIFT),
194            waiters: Mutex::const_new(Waitlist {
195                queue: LinkedList::new(),
196                closed: false,
197            }),
198            #[cfg(all(tokio_unstable, feature = "tracing"))]
199            resource_span: tracing::Span::none(),
200        }
201    }
202
203    /// Creates a new closed semaphore with 0 permits.
204    pub(crate) fn new_closed() -> Self {
205        Self {
206            permits: AtomicUsize::new(Self::CLOSED),
207            waiters: Mutex::new(Waitlist {
208                queue: LinkedList::new(),
209                closed: true,
210            }),
211            #[cfg(all(tokio_unstable, feature = "tracing"))]
212            resource_span: tracing::Span::none(),
213        }
214    }
215
216    /// Creates a new closed semaphore with 0 permits.
217    #[cfg(not(all(loom, test)))]
218    pub(crate) const fn const_new_closed() -> Self {
219        Self {
220            permits: AtomicUsize::new(Self::CLOSED),
221            waiters: Mutex::const_new(Waitlist {
222                queue: LinkedList::new(),
223                closed: true,
224            }),
225            #[cfg(all(tokio_unstable, feature = "tracing"))]
226            resource_span: tracing::Span::none(),
227        }
228    }
229
230    /// Returns the current number of available permits.
231    pub(crate) fn available_permits(&self) -> usize {
232        self.permits.load(Acquire) >> Self::PERMIT_SHIFT
233    }
234
235    /// Adds `added` new permits to the semaphore.
236    ///
237    /// The maximum number of permits is `usize::MAX >> 3`, and this function will panic if the limit is exceeded.
238    pub(crate) fn release(&self, added: usize) {
239        if added == 0 {
240            return;
241        }
242
243        // Assign permits to the wait queue
244        self.add_permits_locked(added, self.waiters.lock());
245    }
246
247    /// Closes the semaphore. This prevents the semaphore from issuing new
248    /// permits and notifies all pending waiters.
249    pub(crate) fn close(&self) {
250        let mut waiters = self.waiters.lock();
251        // If the semaphore's permits counter has enough permits for an
252        // unqueued waiter to acquire all the permits it needs immediately,
253        // it won't touch the wait list. Therefore, we have to set a bit on
254        // the permit counter as well. However, we must do this while
255        // holding the lock --- otherwise, if we set the bit and then wait
256        // to acquire the lock we'll enter an inconsistent state where the
257        // permit counter is closed, but the wait list is not.
258        self.permits.fetch_or(Self::CLOSED, Release);
259        waiters.closed = true;
260
261        // Wake the waiters in batches, with the lock released, because a
262        // waker may call back into the semaphore when it is woken or dropped.
263        // No new waiters are queued once `closed` is set.
264        let mut wakers = WakeList::new();
265        loop {
266            while wakers.can_push() {
267                let Some(mut waiter) = waiters.queue.pop_back() else {
268                    break;
269                };
270                // Safety: we hold the lock, so we can access the waker.
271                let waker = unsafe { waiter.as_mut().waker.with_mut(|waker| (*waker).take()) };
272                if let Some(waker) = waker {
273                    wakers.push(waker);
274                }
275            }
276
277            // If the batch has room left, the queue is empty.
278            let is_empty = wakers.can_push();
279            drop(waiters);
280            wakers.wake_all();
281
282            if is_empty {
283                return;
284            }
285            waiters = self.waiters.lock();
286        }
287    }
288
289    /// Returns true if the semaphore is closed.
290    pub(crate) fn is_closed(&self) -> bool {
291        self.permits.load(Acquire) & Self::CLOSED == Self::CLOSED
292    }
293
294    pub(crate) fn try_acquire(&self, num_permits: usize) -> Result<(), TryAcquireError> {
295        assert!(
296            num_permits <= Self::MAX_PERMITS,
297            "a semaphore may not have more than MAX_PERMITS permits ({})",
298            Self::MAX_PERMITS
299        );
300        let num_permits = num_permits << Self::PERMIT_SHIFT;
301        let mut curr = self.permits.load(Acquire);
302        loop {
303            // Has the semaphore closed?
304            if curr & Self::CLOSED == Self::CLOSED {
305                return Err(TryAcquireError::Closed);
306            }
307
308            // Are there enough permits remaining?
309            if curr < num_permits {
310                return Err(TryAcquireError::NoPermits);
311            }
312
313            let next = curr - num_permits;
314
315            match self.permits.compare_exchange(curr, next, AcqRel, Acquire) {
316                Ok(_) => {
317                    // TODO: Instrument once issue has been solved
318                    return Ok(());
319                }
320                Err(actual) => curr = actual,
321            }
322        }
323    }
324
325    pub(crate) fn acquire(&self, num_permits: usize) -> Acquire<'_> {
326        Acquire::new(self, num_permits)
327    }
328
329    /// Release `rem` permits to the semaphore's wait list, starting from the
330    /// end of the queue.
331    ///
332    /// If `rem` exceeds the number of permits needed by the wait list, the
333    /// remainder are assigned back to the semaphore.
334    fn add_permits_locked(&self, mut rem: usize, waiters: MutexGuard<'_, Waitlist>) {
335        let mut wakers = WakeList::new();
336        let mut lock = Some(waiters);
337        let mut is_empty = false;
338        while rem > 0 {
339            let mut waiters = lock.take().unwrap_or_else(|| self.waiters.lock());
340            'inner: while wakers.can_push() {
341                // Was the waiter assigned enough permits to wake it?
342                let _assigned = match waiters.queue.last() {
343                    Some(waiter) => {
344                        let (should_remove, assigned) = waiter.assign_permits(&mut rem);
345                        if !should_remove {
346                            #[cfg(all(tokio_unstable, feature = "tracing"))]
347                            waiter.trace_assigned(assigned);
348                            break 'inner;
349                        }
350                        assigned
351                    }
352                    None => {
353                        is_empty = true;
354                        // If we assigned permits to all the waiters in the queue, and there are
355                        // still permits left over, assign them back to the semaphore.
356                        break 'inner;
357                    }
358                };
359                let mut waiter = waiters.queue.pop_back().unwrap();
360                if let Some(waker) =
361                    unsafe { waiter.as_mut().waker.with_mut(|waker| (*waker).take()) }
362                {
363                    wakers.push(waker);
364                }
365                // Safety: we have locked the wait list.
366                #[cfg(all(tokio_unstable, feature = "tracing"))]
367                unsafe { waiter.as_ref() }.trace_assigned(_assigned);
368            }
369
370            if rem > 0 && is_empty {
371                let permits = rem;
372                assert!(
373                    permits <= Self::MAX_PERMITS,
374                    "cannot add more than MAX_PERMITS permits ({})",
375                    Self::MAX_PERMITS
376                );
377                let prev = self.permits.fetch_add(rem << Self::PERMIT_SHIFT, Release);
378                let prev = prev >> Self::PERMIT_SHIFT;
379                assert!(
380                    prev + permits <= Self::MAX_PERMITS,
381                    "number of added permits ({}) would overflow MAX_PERMITS ({})",
382                    rem,
383                    Self::MAX_PERMITS
384                );
385
386                // add remaining permits back
387                #[cfg(all(tokio_unstable, feature = "tracing"))]
388                self.resource_span.in_scope(|| {
389                    tracing::trace!(
390                    target: "runtime::resource::state_update",
391                    permits = rem,
392                    permits.op = "add",
393                    )
394                });
395
396                rem = 0;
397            }
398
399            drop(waiters); // release the lock
400
401            wakers.wake_all();
402        }
403
404        assert_eq!(rem, 0);
405    }
406
407    /// Decrease a semaphore's permits by a maximum of `n`.
408    ///
409    /// If there are insufficient permits and it's not possible to reduce by `n`,
410    /// return the number of permits that were actually reduced.
411    pub(crate) fn forget_permits(&self, n: usize) -> usize {
412        if n == 0 {
413            return 0;
414        }
415
416        let mut curr_bits = self.permits.load(Acquire);
417        loop {
418            let curr = curr_bits >> Self::PERMIT_SHIFT;
419            let new = curr.saturating_sub(n);
420            match self.permits.compare_exchange_weak(
421                curr_bits,
422                (new << Self::PERMIT_SHIFT) | (curr_bits & Self::CLOSED),
423                AcqRel,
424                Acquire,
425            ) {
426                Ok(_) => return std::cmp::min(curr, n),
427                Err(actual) => curr_bits = actual,
428            };
429        }
430    }
431
432    fn poll_acquire(
433        &self,
434        cx: &mut Context<'_>,
435        num_permits: usize,
436        node: Pin<&mut Waiter>,
437        queued: &mut bool,
438    ) -> Poll<Result<(), AcquireError>> {
439        let mut acquired = 0;
440
441        let needed = if *queued {
442            node.state.load(Acquire) << Self::PERMIT_SHIFT
443        } else {
444            num_permits << Self::PERMIT_SHIFT
445        };
446
447        let mut lock = None;
448        // First, try to take the requested number of permits from the
449        // semaphore.
450        let mut curr = self.permits.load(Acquire);
451        let mut waiters = loop {
452            // Has the semaphore closed?
453            if curr & Self::CLOSED > 0 {
454                return Poll::Ready(Err(AcquireError::closed()));
455            }
456
457            let mut remaining = 0;
458            let total = curr
459                .checked_add(acquired)
460                .expect("number of permits must not overflow");
461            let (next, acq) = if total >= needed {
462                let next = curr - (needed - acquired);
463                (next, needed >> Self::PERMIT_SHIFT)
464            } else {
465                remaining = (needed - acquired) - curr;
466                (0, curr >> Self::PERMIT_SHIFT)
467            };
468
469            if remaining > 0 && lock.is_none() {
470                // No permits were immediately available, so this permit will
471                // (probably) need to wait. We'll need to acquire a lock on the
472                // wait queue before continuing. We need to do this _before_ the
473                // CAS that sets the new value of the semaphore's `permits`
474                // counter. Otherwise, if we subtract the permits and then
475                // acquire the lock, we might miss additional permits being
476                // added while waiting for the lock.
477                lock = Some(self.waiters.lock());
478            }
479
480            match self.permits.compare_exchange(curr, next, AcqRel, Acquire) {
481                Ok(_) => {
482                    acquired += acq;
483                    if remaining == 0 {
484                        if !*queued {
485                            // The waiter now holds all of its permits. Record
486                            // this in its state and set `queued`, so that
487                            // `Acquire::drop` returns the permits if the
488                            // future is dropped before it completes, e.g.
489                            // because the tracing subscriber panics below.
490                            node.state.store(0, Release);
491                            *queued = true;
492
493                            #[cfg(all(tokio_unstable, feature = "tracing"))]
494                            self.resource_span.in_scope(|| {
495                                tracing::trace!(
496                                    target: "runtime::resource::state_update",
497                                    permits = acquired,
498                                    permits.op = "sub",
499                                );
500                                tracing::trace!(
501                                    target: "runtime::resource::async_op::state_update",
502                                    permits_obtained = acquired,
503                                    permits.op = "add",
504                                )
505                            });
506
507                            return Poll::Ready(Ok(()));
508                        } else if lock.is_none() {
509                            break self.waiters.lock();
510                        }
511                    }
512                    break lock.expect("lock must be acquired before waiting");
513                }
514                Err(actual) => curr = actual,
515            }
516        };
517
518        if waiters.closed {
519            return Poll::Ready(Err(AcquireError::closed()));
520        }
521
522        // The waiter is about to be assigned permits or linked into the wait
523        // queue, so `Acquire::drop` must clean up after it from now on.
524        let was_queued = *queued;
525        *queued = true;
526
527        #[cfg(all(tokio_unstable, feature = "tracing"))]
528        let sub_permits = acquired;
529        let (should_remove, _assigned) = node.assign_permits(&mut acquired);
530
531        // The tracing events are emitted only once the permits have been
532        // assigned to the waiter, so that `Acquire::drop` returns them if the
533        // tracing subscriber panics.
534        #[cfg(all(tokio_unstable, feature = "tracing"))]
535        self.resource_span.in_scope(|| {
536            tracing::trace!(
537                target: "runtime::resource::state_update",
538                permits = sub_permits,
539                permits.op = "sub",
540            )
541        });
542        #[cfg(all(tokio_unstable, feature = "tracing"))]
543        node.trace_assigned(_assigned);
544
545        if should_remove {
546            self.add_permits_locked(acquired, waiters);
547            return Poll::Ready(Ok(()));
548        }
549
550        assert_eq!(acquired, 0);
551        let mut old_waker = None;
552
553        // Otherwise, register the waker & enqueue the node.
554        node.waker.with_mut(|waker| {
555            // Safety: the wait list is locked, so we may modify the waker.
556            let waker = unsafe { &mut *waker };
557            // Do we need to register the new waker?
558            if waker
559                .as_ref()
560                .map_or(true, |waker| !waker.will_wake(cx.waker()))
561            {
562                old_waker = waker.replace(cx.waker().clone());
563            }
564        });
565
566        // If the waiter is not already in the wait queue, enqueue it.
567        if !was_queued {
568            let node = unsafe {
569                let node = Pin::into_inner_unchecked(node) as *mut _;
570                NonNull::new_unchecked(node)
571            };
572
573            waiters.queue.push_front(node);
574        }
575        drop(waiters);
576        drop(old_waker);
577
578        Poll::Pending
579    }
580}
581
582impl fmt::Debug for Semaphore {
583    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
584        fmt.debug_struct("Semaphore")
585            .field("permits", &self.available_permits())
586            .finish()
587    }
588}
589
590impl Waiter {
591    fn new(
592        num_permits: usize,
593        #[cfg(all(tokio_unstable, feature = "tracing"))] ctx: trace::AsyncOpTracingCtx,
594    ) -> Self {
595        Waiter {
596            waker: UnsafeCell::new(None),
597            state: AtomicUsize::new(num_permits),
598            pointers: linked_list::Pointers::new(),
599            #[cfg(all(tokio_unstable, feature = "tracing"))]
600            ctx,
601            _p: PhantomPinned,
602        }
603    }
604
605    /// Assign permits to the waiter.
606    ///
607    /// Returns whether the waiter should be removed from the queue, and the
608    /// number of permits that were assigned to it.
609    fn assign_permits(&self, n: &mut usize) -> (bool, usize) {
610        let mut curr = self.state.load(Acquire);
611        loop {
612            let assign = cmp::min(curr, *n);
613            let next = curr - assign;
614            match self.state.compare_exchange(curr, next, AcqRel, Acquire) {
615                Ok(_) => {
616                    *n -= assign;
617                    return (next == 0, assign);
618                }
619                Err(actual) => curr = actual,
620            }
621        }
622    }
623
624    /// Emit a tracing event for `assigned` permits having been assigned to
625    /// the waiter.
626    #[cfg(all(tokio_unstable, feature = "tracing"))]
627    fn trace_assigned(&self, assigned: usize) {
628        self.ctx.async_op_span.in_scope(|| {
629            tracing::trace!(
630                target: "runtime::resource::async_op::state_update",
631                permits_obtained = assigned,
632                permits.op = "add",
633            );
634        });
635    }
636}
637
638impl Future for Acquire<'_> {
639    type Output = Result<(), AcquireError>;
640
641    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
642        ready!(crate::trace::trace_leaf(cx));
643
644        #[cfg(all(tokio_unstable, feature = "tracing"))]
645        let _resource_span = self.node.ctx.resource_span.clone().entered();
646        #[cfg(all(tokio_unstable, feature = "tracing"))]
647        let _async_op_span = self.node.ctx.async_op_span.clone().entered();
648        #[cfg(all(tokio_unstable, feature = "tracing"))]
649        let _async_op_poll_span = self.node.ctx.async_op_poll_span.clone().entered();
650
651        let (node, semaphore, needed, queued) = self.project();
652
653        // First, ensure the current task has enough budget to proceed.
654        #[cfg(all(tokio_unstable, feature = "tracing"))]
655        let coop = ready!(trace_poll_op!(
656            "poll_acquire",
657            crate::task::coop::poll_proceed(cx),
658        ));
659
660        #[cfg(not(all(tokio_unstable, feature = "tracing")))]
661        let coop = ready!(crate::task::coop::poll_proceed(cx));
662
663        let result = match semaphore.poll_acquire(cx, needed, node, queued) {
664            Poll::Pending => Poll::Pending,
665            Poll::Ready(r) => {
666                coop.made_progress();
667                r?;
668                Poll::Ready(Ok(()))
669            }
670        };
671
672        #[cfg(all(tokio_unstable, feature = "tracing"))]
673        let result = trace_poll_op!("poll_acquire", result);
674
675        // The permits are handed to the caller once the future completes, so
676        // `Acquire::drop` must no longer return them. Clear `queued` only after
677        // the last tracing event, so that the permits are still returned if
678        // the subscriber panics.
679        if result.is_ready() {
680            *queued = false;
681        }
682
683        result
684    }
685}
686
687impl<'a> Acquire<'a> {
688    fn new(semaphore: &'a Semaphore, num_permits: usize) -> Self {
689        assert!(
690            num_permits <= Semaphore::MAX_PERMITS,
691            "a semaphore may not have more than MAX_PERMITS permits ({})",
692            Semaphore::MAX_PERMITS
693        );
694
695        #[cfg(any(not(tokio_unstable), not(feature = "tracing")))]
696        return Self {
697            node: Waiter::new(num_permits),
698            semaphore,
699            num_permits,
700            queued: false,
701        };
702
703        #[cfg(all(tokio_unstable, feature = "tracing"))]
704        return semaphore.resource_span.in_scope(|| {
705            let async_op_span =
706                tracing::trace_span!("runtime.resource.async_op", source = "Acquire::new");
707            let async_op_poll_span = async_op_span.in_scope(|| {
708                tracing::trace!(
709                    target: "runtime::resource::async_op::state_update",
710                    permits_requested = num_permits,
711                    permits.op = "override",
712                );
713
714                tracing::trace!(
715                    target: "runtime::resource::async_op::state_update",
716                    permits_obtained = 0usize,
717                    permits.op = "override",
718                );
719
720                tracing::trace_span!("runtime.resource.async_op.poll")
721            });
722
723            let ctx = trace::AsyncOpTracingCtx {
724                async_op_span,
725                async_op_poll_span,
726                resource_span: semaphore.resource_span.clone(),
727            };
728
729            Self {
730                node: Waiter::new(num_permits, ctx),
731                semaphore,
732                num_permits,
733                queued: false,
734            }
735        });
736    }
737
738    fn project(self: Pin<&mut Self>) -> (Pin<&mut Waiter>, &Semaphore, usize, &mut bool) {
739        fn is_unpin<T: Unpin>() {}
740        unsafe {
741            // Safety: all fields other than `node` are `Unpin`
742
743            is_unpin::<&Semaphore>();
744            is_unpin::<&mut bool>();
745            is_unpin::<usize>();
746
747            let this = self.get_unchecked_mut();
748            (
749                Pin::new_unchecked(&mut this.node),
750                this.semaphore,
751                this.num_permits,
752                &mut this.queued,
753            )
754        }
755    }
756}
757
758impl Drop for Acquire<'_> {
759    fn drop(&mut self) {
760        // If the future is completed, there is no node in the wait list, so we
761        // can skip acquiring the lock.
762        if !self.queued {
763            return;
764        }
765
766        // This is where we ensure safety. The future is being dropped,
767        // which means we must ensure that the waiter entry is no longer stored
768        // in the linked list.
769        let mut waiters = self.semaphore.waiters.lock();
770
771        // remove the entry from the list
772        let node = NonNull::from(&mut self.node);
773        // Safety: we have locked the wait list.
774        unsafe { waiters.queue.remove(node) };
775
776        let acquired_permits = self.num_permits - self.node.state.load(Acquire);
777        if acquired_permits > 0 {
778            self.semaphore.add_permits_locked(acquired_permits, waiters);
779        }
780    }
781}
782
783// Safety: the `Acquire` future is not `Sync` automatically because it contains
784// a `Waiter`, which, in turn, contains an `UnsafeCell`. However, the
785// `UnsafeCell` is only accessed when the future is borrowed mutably (either in
786// `poll` or in `drop`). Therefore, it is safe (although not particularly
787// _useful_) for the future to be borrowed immutably across threads.
788unsafe impl Sync for Acquire<'_> {}
789
790// ===== impl AcquireError ====
791
792impl AcquireError {
793    fn closed() -> AcquireError {
794        AcquireError(())
795    }
796}
797
798impl fmt::Display for AcquireError {
799    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
800        write!(fmt, "semaphore closed")
801    }
802}
803
804impl std::error::Error for AcquireError {}
805
806// ===== impl TryAcquireError =====
807
808impl TryAcquireError {
809    /// Returns `true` if the error was caused by a closed semaphore.
810    #[allow(dead_code)] // may be used later!
811    pub(crate) fn is_closed(&self) -> bool {
812        matches!(self, TryAcquireError::Closed)
813    }
814
815    /// Returns `true` if the error was caused by calling `try_acquire` on a
816    /// semaphore with no available permits.
817    #[allow(dead_code)] // may be used later!
818    pub(crate) fn is_no_permits(&self) -> bool {
819        matches!(self, TryAcquireError::NoPermits)
820    }
821}
822
823impl fmt::Display for TryAcquireError {
824    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
825        match self {
826            TryAcquireError::Closed => write!(fmt, "semaphore closed"),
827            TryAcquireError::NoPermits => write!(fmt, "no permits available"),
828        }
829    }
830}
831
832impl std::error::Error for TryAcquireError {}
833
834/// # Safety
835///
836/// `Waiter` is forced to be !Unpin.
837unsafe impl linked_list::Link for Waiter {
838    type Handle = NonNull<Waiter>;
839    type Target = Waiter;
840
841    fn as_raw(handle: &Self::Handle) -> NonNull<Waiter> {
842        *handle
843    }
844
845    unsafe fn from_raw(ptr: NonNull<Waiter>) -> NonNull<Waiter> {
846        ptr
847    }
848
849    unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> {
850        unsafe { Waiter::addr_of_pointers(target) }
851    }
852}