1#![cfg_attr(not(feature = "sync"), allow(unreachable_pub, dead_code))]
2use 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
34pub(crate) struct Semaphore {
36 waiters: Mutex<Waitlist>,
37 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#[derive(Debug, PartialEq, Eq)]
52pub enum TryAcquireError {
53 Closed,
57
58 NoPermits,
60}
61#[derive(Debug)]
69pub struct AcquireError(());
70
71pub(crate) struct Acquire<'a> {
72 node: Waiter,
73 semaphore: &'a Semaphore,
74 num_permits: usize,
75 queued: bool,
83}
84
85struct Waiter {
87 state: AtomicUsize,
92
93 waker: UnsafeCell<Option<Waker>>,
99
100 pointers: linked_list::Pointers<Waiter>,
113
114 #[cfg(all(tokio_unstable, feature = "tracing"))]
115 ctx: trace::AsyncOpTracingCtx,
116
117 _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 pub(crate) const MAX_PERMITS: usize = usize::MAX >> 3;
138 const CLOSED: usize = 1;
139 const PERMIT_SHIFT: usize = 1;
143
144 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 #[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 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 #[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 pub(crate) fn available_permits(&self) -> usize {
232 self.permits.load(Acquire) >> Self::PERMIT_SHIFT
233 }
234
235 pub(crate) fn release(&self, added: usize) {
239 if added == 0 {
240 return;
241 }
242
243 self.add_permits_locked(added, self.waiters.lock());
245 }
246
247 pub(crate) fn close(&self) {
250 let mut waiters = self.waiters.lock();
251 self.permits.fetch_or(Self::CLOSED, Release);
259 waiters.closed = true;
260
261 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 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 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 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 if curr & Self::CLOSED == Self::CLOSED {
305 return Err(TryAcquireError::Closed);
306 }
307
308 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 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 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 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 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 #[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 #[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); wakers.wake_all();
402 }
403
404 assert_eq!(rem, 0);
405 }
406
407 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 let mut curr = self.permits.load(Acquire);
451 let mut waiters = loop {
452 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 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 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 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 #[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 node.waker.with_mut(|waker| {
555 let waker = unsafe { &mut *waker };
557 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 !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 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 #[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 #[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 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 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 !self.queued {
763 return;
764 }
765
766 let mut waiters = self.semaphore.waiters.lock();
770
771 let node = NonNull::from(&mut self.node);
773 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
783unsafe impl Sync for Acquire<'_> {}
789
790impl 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
806impl TryAcquireError {
809 #[allow(dead_code)] pub(crate) fn is_closed(&self) -> bool {
812 matches!(self, TryAcquireError::Closed)
813 }
814
815 #[allow(dead_code)] 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
834unsafe 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}