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 if !ExecutionState::should_stop() {
795 state.wake_unfair_waiters();
796 }
797 }
798
799 /// Acquire the specified number of permits (async API)
800 pub fn acquire(&self, num_permits: usize) -> Acquire<'_> {
801 // No switch here; switch should be triggered on polling future
802 self.init_object_id();
803 Acquire::new(self, num_permits, Priority::Back)
804 }
805
806 /// Acquire permits, allowing later requests to pass while this request cannot be satisfied.
807 ///
808 /// On a strictly fair semaphore, this request still waits behind earlier ordinary requests,
809 /// and receives permits in queue order once enough are released. It does not prevent later
810 /// requests from taking permits that it cannot use yet. This lets an upgradable `RwLock`
811 /// reader wait for another upgradable reader without blocking plain readers.
812 ///
813 /// On an unfair semaphore, this is equivalent to [`Self::acquire`].
814 pub fn acquire_bypassable(&self, num_permits: usize) -> Acquire<'_> {
815 self.init_object_id();
816 Acquire::new(self, num_permits, Priority::Bypassable)
817 }
818
819 /// Acquire the specified number of permits (blocking API)
820 pub fn acquire_blocking(&self, num_permits: usize) -> Result<(), AcquireError> {
821 crate::future::block_on(self.acquire(num_permits))
822 }
823
824 /// Acquire `num_permits` permits, and reserve the semaphore for this request
825 /// as soon as at least `min_permits` permits are available (async API). Only
826 /// an unfair semaphore supports this.
827 ///
828 /// Until `min_permits` permits are available, the request waits like one
829 /// from [`BatchSemaphore::acquire`]: it holds nothing and stops no other
830 /// request. As soon as they are, it reserves the semaphore in the same step.
831 /// From then on, no other request can take a permit, and this request takes
832 /// its `num_permits` as soon as that many are available. The reservation
833 /// ends when the request is granted, or when the returned future is dropped.
834 /// While it lasts, [`BatchSemaphore::available_permits`] is zero.
835 ///
836 /// The motivating use case is a `parking_lot` `RwLock` writer. `parking_lot`
837 /// sets `WRITER_BIT` only when no writer or upgradable reader holds the lock,
838 /// and from then on, the bit stops new readers while the writer waits for
839 /// the current ones to leave. With `min_permits` above the permits that are
840 /// left while an upgradable reader holds the lock, the reservation is that
841 /// bit.
842 ///
843 /// At most one request can hold the reservation. Another reserving request
844 /// waits like any other until the reservation ends.
845 ///
846 /// # Panics
847 ///
848 /// Panics if the semaphore is strictly fair (its queue already keeps the
849 /// permits for its first waiter), if `num_permits` is zero, or if
850 /// `min_permits > num_permits`.
851 pub fn acquire_reserving(&self, min_permits: usize, num_permits: usize) -> Acquire<'_> {
852 assert_eq!(
853 self.fairness,
854 Fairness::Unfair,
855 "only an unfair semaphore supports reservations"
856 );
857 assert!(num_permits > 0);
858 assert!(min_permits <= num_permits);
859
860 self.init_object_id();
861 Acquire::new_reserving(self, num_permits, min_permits)
862 }
863
864 /// Release `num_permits` back to the Semaphore
865 pub fn release(&self, num_permits: usize) {
866 // Execution teardown can unwind a task's stack from this scheduling point, which is often in
867 // a destructor that releases a lock (see `ExecutionState::tear_down`). The permits must not
868 // be lost then, as destructors that run later can need them.
869 struct ReleaseOnUnwind<'a>(&'a BatchSemaphore, usize);
870 impl Drop for ReleaseOnUnwind<'_> {
871 fn drop(&mut self) {
872 self.0.release_no_scheduling_point(self.1);
873 }
874 }
875 let release_on_unwind = ReleaseOnUnwind(self, num_permits);
876 thread::switch();
877 std::mem::forget(release_on_unwind);
878
879 self.release_no_scheduling_point(num_permits);
880 }
881
882 /// `release` without its scheduling point.
883 #[inline]
884 fn release_no_scheduling_point(&self, num_permits: usize) {
885 self.init_object_id();
886 if num_permits == 0 {
887 return;
888 }
889
890 let mut state = self.state.borrow_mut();
891
892 crate::annotations::record_semaphore_release(state.id.unwrap(), num_permits);
893
894 // A release by a panicking task (e.g., a lock guard dropped during unwinding) is an
895 // ordinary release, because the panic may be caught and the execution may continue.
896 // Shuttle's `Mutex` and `RwLock` record poisoning in their inner `std::sync::Mutex` and
897 // `std::sync::RwLock`. If the execution has stopped, or its state is borrowed, the waiters
898 // cannot be unblocked (see `ExecutionState::should_stop`). In that case, the release adds
899 // the permits, clears the waiter queue, and closes the semaphore.
900 if ExecutionState::should_stop() {
901 state.permits_available.release(num_permits, VectorClock::new());
902 for waiter in &state.waiters {
903 waiter.is_queued.swap(false, Ordering::SeqCst);
904 }
905 state.waiters.clear();
906 state.reservation = None;
907 state.closed = true;
908 return;
909 }
910
911 // Permits released into the semaphore reflect the releasing thread's
912 // clock; future acquires of those permits are causally dependent on
913 // this event.
914 ExecutionState::with(|s| {
915 let clock = s.increment_clock();
916 state.permits_available.release(num_permits, clock.clone());
917 });
918
919 // `ExecutionState::me()` is only wanted for this trace, so let the macro's
920 // level check decide whether to pay for it. Computing it eagerly cost an
921 // `ExecutionState::with` on every release even with tracing disabled.
922 trace!(task = ?ExecutionState::me(), avail = ?state.permits_available, waiters = ?state.waiters, "released {} permits for semaphore {:p}", num_permits, &self.state);
923
924 match self.fairness {
925 Fairness::StrictlyFair => {
926 // in a strictly fair mode we will grant permits to waiters from the front
927 // of the queue, as long as there are enough permits available
928 state.unblock_waiters_from_front();
929 }
930 Fairness::Unfair => {
931 // in an unfair mode, we will unblock all the waiters for which
932 // there are enough permits available, then let them race
933 state.wake_unfair_waiters();
934 }
935 }
936 drop(state);
937 }
938
939 /// Atomically `upgrade` from holding `permits_currently_held` permits to holding
940 /// `permits_to_be_held`, without ever dropping below `permits_currently_held` in between.
941 /// The motivating use case is `parking_lot`'s `RwLockUpgradableReadGuard::upgrade`, which must
942 /// take a read guard to a write guard without letting any writer in along the way.
943 ///
944 /// This is implemented by acquiring only the *missing* permits
945 /// (`permits_to_be_held - permits_currently_held`), with priority over any waiter already
946 /// queued, so that the request overtakes the queue. Both halves of that matter:
947 ///
948 /// * Keeping the held permits means no other task can claim the resource mid-upgrade. Releasing
949 /// them first (even for an instant) would hand the resource to a queued waiter, which for an
950 /// `RwLock` means a writer mutating the data an upgradable reader had already observed.
951 /// * Overtaking the queue is what makes that safe rather than deadlock-prone. Since we hold
952 /// permits we will not release, a queued waiter ahead of us may be unsatisfiable (an `RwLock`
953 /// writer wants *all* permits), so waiting our turn behind it could deadlock. Queued waiters
954 /// hold no permits, so overtaking them costs nothing but their place in line -- which is
955 /// exactly the priority a real upgradable read lock gives an upgrade.
956 ///
957 /// The upgrade therefore blocks only on tasks that *currently hold* permits, and is granted as
958 /// soon as they release. The returned future must be driven to completion; if it is dropped
959 /// first, the caller still holds `permits_currently_held`.
960 ///
961 /// An unfair semaphore has no queue to overtake, so there the upgrade reserves the semaphore
962 /// instead (see [`BatchSemaphore::acquire_reserving`]), at once unless another request holds
963 /// the reservation. From then on, no other request can take a permit, so the upgrade again
964 /// waits only for the tasks that hold permits, and nothing can overtake it.
965 ///
966 /// At most one `upgrade` may be in flight on a semaphore at a time. Two concurrent upgraders
967 /// could each be waiting for permits the other holds, which no queue discipline can resolve.
968 /// Callers are expected to enforce this (an `RwLock` does: there is only ever one upgradable
969 /// reader).
970 pub fn upgrade(&self, permits_currently_held: usize, permits_to_be_held: usize) -> Acquire<'_> {
971 assert!(permits_currently_held > 0);
972 assert!(permits_to_be_held > permits_currently_held);
973
974 self.init_object_id();
975 let num_permits = permits_to_be_held - permits_currently_held;
976 match self.fairness {
977 Fairness::StrictlyFair => Acquire::new(self, num_permits, Priority::Front),
978 Fairness::Unfair => Acquire::new_reserving(self, num_permits, 0),
979 }
980 }
981
982 /// The non-blocking analogue of [`BatchSemaphore::upgrade`]: succeeds only if the missing
983 /// permits are available right now, and never blocks or queues.
984 ///
985 /// Like `upgrade`, this ignores queued waiters (they hold no permits, so they cannot be the
986 /// reason the upgrade is short of permits). A `try_upgrade` therefore fails only when some
987 /// other task actually *holds* permits the upgrade needs.
988 pub fn try_upgrade(&self, permits_currently_held: usize, permits_to_be_held: usize) -> Result<(), TryAcquireError> {
989 assert!(permits_currently_held > 0);
990 assert!(permits_to_be_held > permits_currently_held);
991
992 thread::switch();
993
994 self.init_object_id();
995 let num_permits = permits_to_be_held - permits_currently_held;
996 let mut state = self.state.borrow_mut();
997 let id = state.id.unwrap();
998 let res = state
999 .acquire_permits(num_permits, self.fairness, Priority::Front)
1000 .inspect_err(|_err| {
1001 // Conservatively, the requester causally depends on the last successful acquire;
1002 // see the equivalent reasoning in `try_acquire`.
1003 ExecutionState::with(|s| {
1004 s.update_clock(&state.permits_available.last_acquire);
1005 });
1006 });
1007 drop(state);
1008
1009 // If we took permits from an unfair semaphore, re-block waiting threads that can no longer
1010 // succeed.
1011 if res.is_ok() {
1012 self.reblock_if_unfair();
1013 }
1014
1015 crate::annotations::record_semaphore_try_acquire(id, num_permits, res.is_ok());
1016
1017 res
1018 }
1019}
1020
1021// Safety: Semaphore is never actually passed across true threads, only across continuations. The
1022// RefCell<_> type therefore can't be preempted mid-bookkeeping-operation.
1023// TODO we shouldn't need to do this, but RefCell is not Send, and anything we put within a Semaphore
1024// TODO needs to be Send.
1025unsafe impl Send for BatchSemaphore {}
1026unsafe impl Sync for BatchSemaphore {}
1027
1028impl Default for BatchSemaphore {
1029 #[track_caller]
1030 fn default() -> Self {
1031 Self::new(Default::default(), Fairness::StrictlyFair)
1032 }
1033}
1034
1035/// The future that results from async calls to `acquire*`.
1036/// Callers must `await` on this future to obtain the necessary permits.
1037pub struct Acquire<'a> {
1038 semaphore: &'a BatchSemaphore,
1039 num_permits: usize,
1040
1041 /// Where this acquire sits relative to waiters already queued on a fair
1042 /// semaphore. Only [`BatchSemaphore::upgrade`] uses [`Priority::Front`]; see
1043 /// there for why an upgrade must overtake the queue.
1044 priority: Priority,
1045
1046 /// For a reserving acquire, the number of available permits at which it
1047 /// reserves the semaphore (see [`BatchSemaphore::acquire_reserving`]).
1048 /// `None` for every other acquire.
1049 reserve_at: Option<usize>,
1050
1051 /// Snapshotted when this `Acquire` is created, and moved into the `Waiter` if
1052 /// this acquire ends up blocking. See `Waiter::new` for why the snapshot must
1053 /// happen here rather than at enqueue time.
1054 clock: VectorClock,
1055
1056 /// The shared part of this acquire, allocated only once the acquire has to
1057 /// block. An acquire that gets its permits immediately is never visible to
1058 /// any other task, so it needs no shared state and no allocation. While this
1059 /// is `None`, `has_permits` below is authoritative.
1060 waiter: Option<Arc<Waiter>>,
1061
1062 /// Whether permits have been granted, for the case where no `Waiter` exists.
1063 /// Once one does, the releasing task writes `Waiter::has_permits` instead and
1064 /// this field is unused; read through `Acquire::has_permits`.
1065 has_permits: bool,
1066
1067 completed: bool, // Has the future completed yet?
1068 never_polled: bool,
1069}
1070
1071// Implement Debug in order to not output the `VectorClock`, matching `Waiter`.
1072impl fmt::Debug for Acquire<'_> {
1073 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1074 f.debug_struct("Acquire")
1075 .field("num_permits", &self.num_permits)
1076 .field("priority", &self.priority)
1077 .field("reserve_at", &self.reserve_at)
1078 .field("waiter", &self.waiter)
1079 .field("has_permits", &self.has_permits())
1080 .field("completed", &self.completed)
1081 .finish()
1082 }
1083}
1084
1085impl<'a> Acquire<'a> {
1086 fn new(semaphore: &'a BatchSemaphore, num_permits: usize, priority: Priority) -> Self {
1087 Self {
1088 semaphore,
1089 num_permits,
1090 priority,
1091 reserve_at: None,
1092 clock: current::clock(),
1093 waiter: None,
1094 has_permits: false,
1095 completed: false,
1096 never_polled: true,
1097 }
1098 }
1099
1100 fn new_reserving(semaphore: &'a BatchSemaphore, num_permits: usize, min_permits: usize) -> Self {
1101 let mut acquire = Self::new(semaphore, num_permits, Priority::Back);
1102 acquire.reserve_at = Some(min_permits);
1103 acquire
1104 }
1105
1106 /// How many permits must be available for this acquire to make progress
1107 /// (see `Waiter::min_permits`).
1108 fn min_permits(&self) -> usize {
1109 self.reserve_at.unwrap_or(self.num_permits)
1110 }
1111
1112 /// Does this acquire hold the semaphore's reservation?
1113 fn is_reserving(&self) -> bool {
1114 self.waiter
1115 .as_ref()
1116 .is_some_and(|waiter| self.semaphore.state.borrow().is_reserved_by(waiter))
1117 }
1118
1119 /// Have permits been granted to this acquire? Once a `Waiter` exists the
1120 /// releasing task owns that flag, so the shared copy is authoritative.
1121 fn has_permits(&self) -> bool {
1122 match &self.waiter {
1123 Some(waiter) => waiter.has_permits.load(Ordering::SeqCst),
1124 None => self.has_permits,
1125 }
1126 }
1127
1128 /// Is this acquire in the semaphore's waiter queue? Only possible once a
1129 /// `Waiter` has been allocated, since the queue holds `Arc<Waiter>`.
1130 fn is_queued(&self) -> bool {
1131 match &self.waiter {
1132 Some(waiter) => waiter.is_queued.load(Ordering::SeqCst),
1133 None => false,
1134 }
1135 }
1136
1137 fn grant_permits(&mut self) {
1138 match &self.waiter {
1139 Some(waiter) => waiter.has_permits.store(true, Ordering::SeqCst),
1140 None => self.has_permits = true,
1141 }
1142 }
1143
1144 /// The shared `Waiter` for this acquire, allocating it if this is the first
1145 /// time the acquire has had to block. Returns an owned handle so callers can
1146 /// still use `self.semaphore` without holding a borrow of `self`.
1147 fn waiter_for_blocking(&mut self) -> Arc<Waiter> {
1148 if let Some(waiter) = &self.waiter {
1149 return Arc::clone(waiter);
1150 }
1151 let waiter = Arc::new(Waiter::new(
1152 self.num_permits,
1153 self.min_permits(),
1154 self.priority,
1155 self.clock.clone(),
1156 ));
1157 self.waiter = Some(Arc::clone(&waiter));
1158 waiter
1159 }
1160
1161 /// The part of `poll` for a reserving acquire (see
1162 /// [`BatchSemaphore::acquire_reserving`]), once `poll` knows that it has no
1163 /// permits yet and that the semaphore is open. Only an unfair semaphore has
1164 /// reserving acquires.
1165 fn poll_reserving(&mut self, min_permits: usize, cx: &mut Context<'_>) -> Poll<Result<(), AcquireError>> {
1166 let semaphore = self.semaphore;
1167 let is_queued = self.is_queued();
1168 let is_reserving = self.is_reserving();
1169 trace!(
1170 "Acquire::poll for reserving {:?}; is queued: {is_queued:?}, is reserving: {is_reserving:?}",
1171 self
1172 );
1173
1174 let mut state = semaphore.state.borrow_mut();
1175 let id = state.id.unwrap();
1176 let available = state.permits_available.available();
1177
1178 if is_reserving {
1179 if available < self.num_permits {
1180 // Still waiting for the tasks that hold the rest. Like a queued
1181 // waiter, follow the current poller.
1182 drop(state);
1183 let waiter = self.waiter_for_blocking();
1184 *waiter.waker.lock().unwrap() = Some(cx.waker().clone());
1185 waiter.set_task_id(ExecutionState::me());
1186 return Poll::Pending;
1187 }
1188
1189 // The reservation kept the permits for us, so take them.
1190 state.reservation = None;
1191 state.take_permits(self.num_permits).unwrap();
1192 // Let the waiters race for any permits that are left.
1193 state.wake_unfair_waiters();
1194 drop(state);
1195
1196 let waiter = self
1197 .waiter
1198 .clone()
1199 .expect("a reserving acquire must have an allocated waiter");
1200 crate::annotations::record_semaphore_acquire_unblocked(id, waiter.task_id(), self.num_permits);
1201 self.grant_permits();
1202 self.completed = true;
1203 trace!("Acquire::poll for {:?} that got permits", self);
1204 return Poll::Ready(Ok(()));
1205 }
1206
1207 if state.reservation.is_some() || available < min_permits {
1208 // Wait, holding nothing, like any other waiter of an unfair
1209 // semaphore.
1210 drop(state);
1211 let waiter = self.waiter_for_blocking();
1212 *waiter.waker.lock().unwrap() = Some(cx.waker().clone());
1213 waiter.set_task_id(ExecutionState::me());
1214 if !is_queued {
1215 crate::annotations::record_semaphore_acquire_blocked(id, self.num_permits);
1216 semaphore.enqueue_waiter(&waiter, Priority::Back);
1217 }
1218 trace!("Acquire::poll for {:?} that is enqueued", self);
1219 return Poll::Pending;
1220 }
1221
1222 if available >= self.num_permits {
1223 // There are enough permits, so there is nothing to reserve.
1224 state.take_permits(self.num_permits).unwrap();
1225 drop(state);
1226 if is_queued {
1227 let waiter = self
1228 .waiter
1229 .clone()
1230 .expect("a queued acquire must have an allocated waiter");
1231 crate::annotations::record_semaphore_acquire_unblocked(id, waiter.task_id(), self.num_permits);
1232 semaphore.remove_waiter(&waiter);
1233 } else {
1234 crate::annotations::record_semaphore_acquire_fast(id, self.num_permits);
1235 }
1236 self.grant_permits();
1237 self.completed = true;
1238 trace!("Acquire::poll for {:?} that got permits", self);
1239 semaphore.reblock_if_unfair();
1240 return Poll::Ready(Ok(()));
1241 }
1242
1243 // Reserve the semaphore, and wait for the rest of the permits.
1244 drop(state);
1245 let waiter = self.waiter_for_blocking();
1246 *waiter.waker.lock().unwrap() = Some(cx.waker().clone());
1247 waiter.set_task_id(ExecutionState::me());
1248 if is_queued {
1249 semaphore.remove_waiter(&waiter);
1250 } else {
1251 crate::annotations::record_semaphore_acquire_blocked(id, self.num_permits);
1252 }
1253 semaphore.state.borrow_mut().reservation = Some(waiter);
1254 trace!("Acquire::poll for {:?} that reserved the semaphore", self);
1255 // No waiter can take a permit now.
1256 semaphore.reblock_if_unfair();
1257 Poll::Pending
1258 }
1259}
1260
1261impl Future for Acquire<'_> {
1262 type Output = Result<(), AcquireError>;
1263
1264 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1265 assert!(!self.completed);
1266
1267 // One borrow of the semaphore state rather than two (`is_closed` and
1268 // `available_permits` each took their own). Both reads describe the same
1269 // instant, before the scheduling point below, so merging them is sound.
1270 // Reads *after* the switch must stay separate and fresh, because other
1271 // tasks may have run in between. A reserving acquire that will reserve
1272 // the semaphore changes its state as much as one that will succeed, so
1273 // it compares against the permits at which it reserves.
1274 let will_succeed = self.has_permits() || {
1275 let state = self.semaphore.state.borrow();
1276 state.closed || state.available() >= self.min_permits()
1277 };
1278
1279 // If the acquire will succeed on the first try, we need to context switch once to allow the previous
1280 // event to become visible. If we won't succeed, then we still need to context switch if the act of
1281 // blocking does not commute with other operations on `batch_semaphore` (double-yield optimization,
1282 // reasoning below).
1283 //
1284 // Fair Semaphores: blocking adds the current task to an *ordered* waiter queue. Two blocking acquires
1285 // *do not commute* because in one ordering the queue will be [T1 T2] and in the other ordering [T2 T1].
1286 // Thus we cannot apply the double-yield optimization for fair semaphores.
1287 //
1288 // Unfair Semaphores: blocking adds the current task to an *unordered set* of waiters. To check if the
1289 // double-yield is valid we check if each operation (Z) on the semaphore commutes with a blocking acquire (Y1):
1290 //
1291 // - Blocking Acquire: in both orderings `Z Y1` and `Y1 Z`, the waiter set has the same members, thus
1292 // the operations commute.
1293 // - Try Acquire: the try-acquire will fail in both orderings without changing the state of the semaphore
1294 // - Release: if the release unblocks Y1, then the optimization is not applicable. Otherwise, it must
1295 // unblock another task in the waiter set. As waiter-set insertion and removal for disjoint elements
1296 // commutes, release operations also commute in this case.
1297 //
1298 // Thus we apply the double-yield optimization for *unfair* semaphores only
1299 let blocking_is_not_commutative = self.semaphore.fairness == Fairness::StrictlyFair;
1300
1301 if self.never_polled && (will_succeed || blocking_is_not_commutative) {
1302 thread::switch();
1303 }
1304 self.never_polled = false;
1305
1306 let out = if self.has_permits() {
1307 assert!(!self.is_queued());
1308 self.completed = true;
1309 trace!("Acquire::poll for {:?} with permits", self);
1310 Poll::Ready(Ok(()))
1311 } else if self.semaphore.is_closed() {
1312 assert!(!self.is_queued());
1313 self.completed = true;
1314 trace!("Acquire::poll for {:?} with closed", self);
1315 Poll::Ready(Err(AcquireError::closed()))
1316 } else if let Some(min_permits) = self.reserve_at {
1317 self.poll_reserving(min_permits, cx)
1318 } else {
1319 let is_queued = self.is_queued();
1320 trace!("Acquire::poll for {:?}; is queued: {is_queued:?}", self);
1321
1322 // Sanity check: there should be a waker if the waiter is in
1323 // the queue. Also true for unfair semaphores, which wake by ref.
1324 //
1325 // `debug_assert` rather than `assert`: this takes a `std::sync::Mutex`
1326 // on every poll, including the uncontended fast path, purely to check
1327 // an internal invariant.
1328 debug_assert_eq!(
1329 is_queued,
1330 self.waiter
1331 .as_ref()
1332 .is_some_and(|waiter| waiter.waker.lock().unwrap().is_some())
1333 );
1334
1335 // Should the waiter try to acquire permits here? Four cases:
1336 // 1. unfair semaphore, waiter not yet enqueued;
1337 // 2. fair semaphore, waiter not yet enqueued;
1338 // 3. unfair semaphore, waiter already enqueued.
1339 // 4. fair semaphore, waiter already enqueued;
1340 //
1341 // 1. and 2. are similar: the future was polled for the first time,
1342 // so the waiter will try to acquire some permits. If successful,
1343 // the waiter need not be enqueued, and the future is resolved.
1344 // Otherwise, the waiter is added to the queue.
1345 //
1346 // 3. is slightly different: the future was polled, even though the
1347 // waiter was already in the queue. This can happen either because
1348 // the semaphore just received some permits and woke the waiter up,
1349 // or because the future itself was polled manually. Either way,
1350 // the semaphore is queried.
1351 //
1352 // 4. is a case where we do not try to acquire permits. The request
1353 // would always fail, and the waiter should remain suspended until
1354 // the semaphore has explicitly unblocked it and given it permits
1355 // during a `release` call.
1356 let try_to_acquire = match (self.semaphore.fairness, is_queued) {
1357 // written this way to mirror the cases described above
1358 (Fairness::Unfair, false) | (Fairness::StrictlyFair, false) | (Fairness::Unfair, true) => true,
1359 (Fairness::StrictlyFair, true) => false,
1360 };
1361
1362 if try_to_acquire {
1363 // Access the semaphore state directly instead of `try_acquire`,
1364 // because in case of `NoPermits`, we do not want to update the
1365 // clock, as this thread will be blocked below.
1366 let mut state = self.semaphore.state.borrow_mut();
1367 let id = state.id.unwrap();
1368 let acquire_result = state.acquire_permits(self.num_permits, self.semaphore.fairness, self.priority);
1369 drop(state);
1370
1371 match acquire_result {
1372 Ok(()) => {
1373 if is_queued {
1374 let waiter = self
1375 .waiter
1376 .clone()
1377 .expect("a queued acquire must have an allocated waiter");
1378 crate::annotations::record_semaphore_acquire_unblocked(
1379 id,
1380 waiter.task_id(),
1381 waiter.num_permits,
1382 );
1383 self.semaphore.remove_waiter(&waiter);
1384 } else {
1385 crate::annotations::record_semaphore_acquire_fast(id, self.num_permits);
1386 }
1387 self.grant_permits();
1388 self.completed = true;
1389 trace!("Acquire::poll for {:?} that got permits", self);
1390
1391 // If the semaphore is unfair, re-block other waiting
1392 // threads that can no longer succeed.
1393 self.semaphore.reblock_if_unfair();
1394
1395 Poll::Ready(Ok(()))
1396 }
1397 Err(TryAcquireError::NoPermits) => {
1398 // This acquire has to block, so it now becomes visible to
1399 // whichever task releases permits. That is the first point
1400 // at which shared state is needed, so it is where the
1401 // `Waiter` gets allocated.
1402 let waiter = self.waiter_for_blocking();
1403
1404 let mut maybe_waker = waiter.waker.lock().unwrap();
1405 *maybe_waker = Some(cx.waker().clone());
1406 drop(maybe_waker);
1407
1408 // Point the waiter at whoever is polling now: this future
1409 // may have been created by a different task.
1410 waiter.set_task_id(ExecutionState::me());
1411
1412 if !is_queued {
1413 crate::annotations::record_semaphore_acquire_blocked(id, self.num_permits);
1414 // `enqueue_waiter` sets `is_queued` itself.
1415 self.semaphore.enqueue_waiter(&waiter, self.priority);
1416 }
1417 trace!("Acquire::poll for {:?} that is enqueued", self);
1418 Poll::Pending
1419 }
1420 Err(TryAcquireError::Closed) => unreachable!(),
1421 }
1422 } else {
1423 // No progress made, future is still pending. The waiter stays in
1424 // the queue, but re-point it at the current poller and refresh
1425 // its waker: this future may have been created by (or last
1426 // polled by) another task, and `release` must wake whoever is
1427 // waiting now. Without this, a permit granted to this waiter
1428 // would unblock a task that is no longer interested, and the
1429 // actual poller would never be woken.
1430 let waiter = self
1431 .waiter
1432 .as_ref()
1433 .expect("a queued acquire must have an allocated waiter");
1434 *waiter.waker.lock().unwrap() = Some(cx.waker().clone());
1435 waiter.set_task_id(ExecutionState::me());
1436 Poll::Pending
1437 }
1438 };
1439 if matches!(out, Poll::Pending) {
1440 // `Backtrace::capture()` is a noop (it returns the constant `disabled()`) if `RUST_BACKTRACE`/`RUST_LIB_BACKTRACE` is not set.
1441 ExecutionState::with(|state| {
1442 state.current_mut().backtrace = if backtrace_enabled() {
1443 Some(std::backtrace::Backtrace::force_capture())
1444 } else {
1445 None
1446 }
1447 })
1448 }
1449 out
1450 }
1451}
1452
1453impl Drop for Acquire<'_> {
1454 fn drop(&mut self) {
1455 trace!("Acquire::drop for {:?}", self);
1456 if self.is_queued() {
1457 // If the associated waiter is in the wait list, remove it
1458 let waiter = self
1459 .waiter
1460 .clone()
1461 .expect("a queued acquire must have an allocated waiter");
1462 self.semaphore.remove_waiter(&waiter);
1463 } else if self.is_reserving() {
1464 // If the acquire holds the reservation, end it, so that it does not
1465 // keep the permits from other requests.
1466 let waiter = self
1467 .waiter
1468 .clone()
1469 .expect("a reserving acquire must have an allocated waiter");
1470 self.semaphore.cancel_reservation(&waiter);
1471 } else if self.has_permits() && !self.completed {
1472 // If the waiter was granted permits, release them. Note this must also
1473 // fire for an acquire that got its permits without ever allocating a
1474 // waiter, otherwise the semaphore leaks permits.
1475 self.semaphore.release(self.num_permits);
1476 }
1477 }
1478}
1479
1480impl crate::annotations::WithName for &BatchSemaphore {
1481 fn with_name_and_kind(self, name: Option<&str>, kind: Option<&str>) -> Self {
1482 self.init_object_id();
1483 crate::annotations::record_name_for_object(self.state.borrow().id.unwrap(), name, kind);
1484 self
1485 }
1486}
1487
1488impl crate::annotations::WithName for BatchSemaphore {
1489 fn with_name_and_kind(self, name: Option<&str>, kind: Option<&str>) -> Self {
1490 (&self).with_name_and_kind(name, kind);
1491 self
1492 }
1493}
1494
1495impl BatchSemaphore {
1496 /// Returns a reference to this semaphore's resource signature.
1497 pub fn signature(&self) -> &ResourceSignature {
1498 &self.signature
1499 }
1500}