nuclei/sync/
semaphore.rs

1//! A fair, weighted counting semaphore whose waits are futures.
2
3use core::error::Error;
4use core::fmt;
5use core::future::Future;
6use core::pin::Pin;
7use core::task::{Context, Poll};
8use std::sync::Arc;
9
10use super::contend::{Closed, Lock};
11use super::internal;
12use super::raw_semaphore::{MAX, PermitWait, RawSemaphore, observed_closed};
13
14/// Error returned from [`Semaphore::acquire`] and the other waits for
15/// permits.
16///
17/// A wait can only fail if the semaphore has been
18/// [closed](Semaphore::close).
19#[derive(Debug)]
20pub struct AcquireError(());
21
22impl AcquireError {
23    pub(crate) fn closed() -> Self {
24        Self(())
25    }
26}
27
28impl fmt::Display for AcquireError {
29    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30        f.write_str("semaphore closed")
31    }
32}
33
34impl Error for AcquireError {}
35
36/// Error returned from [`Semaphore::try_acquire`] and the other attempts
37/// that never wait.
38#[derive(Debug, PartialEq, Eq)]
39pub enum TryAcquireError {
40    /// The semaphore has been [closed](Semaphore::close) and cannot issue
41    /// new permits.
42    Closed,
43
44    /// The semaphore has no available permits.
45    NoPermits,
46}
47
48impl TryAcquireError {
49    /// What a failed attempt that saw `state` means.
50    fn of(state: usize) -> Self {
51        if observed_closed(state) {
52            Self::Closed
53        } else {
54            Self::NoPermits
55        }
56    }
57}
58
59impl fmt::Display for TryAcquireError {
60    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61        match self {
62            Self::Closed => f.write_str("semaphore closed"),
63            Self::NoPermits => f.write_str("no permits available"),
64        }
65    }
66}
67
68impl Error for TryAcquireError {}
69
70/// A counting semaphore whose waits are futures.
71///
72/// A request for `n` permits takes them whenever `n` are free, even while
73/// others wait; under contention [`acquire_many`](Self::acquire_many)
74/// queues. A release wakes the oldest waiters the freed permits could
75/// serve, stopping at the first that does not fit, so a large request at
76/// the front is not skipped for smaller ones behind it. A waiter that
77/// loses [`MAX_BYPASS`](super::MAX_BYPASS) times is owed a handoff: until
78/// it is served, [`try_acquire`](Self::try_acquire) fails and releases go
79/// to the waiters in queue order.
80///
81/// [`close`](Self::close) ends every wait with an [`AcquireError`] and makes
82/// later acquires fail at once.
83///
84/// ```
85/// use nuclei::sync::{Semaphore, TryAcquireError};
86///
87/// let budget = Semaphore::new(4);
88/// let first = budget.try_acquire_many(3).expect("three of four are free");
89/// assert_eq!(budget.try_acquire_many(2).unwrap_err(), TryAcquireError::NoPermits);
90/// drop(first);
91/// assert_eq!(budget.available_permits(), 4);
92/// ```
93pub struct Semaphore {
94    raw: RawSemaphore,
95}
96
97impl Semaphore {
98    /// The most permits a semaphore can hold: `usize::MAX >> 7`. The state
99    /// word keeps seven low bits for the wait queue and the closed flag, so
100    /// this is `2^57 - 1` on a 64-bit target and `2^25 - 1` on a 32-bit one.
101    pub const MAX_PERMITS: usize = MAX;
102
103    loom_const_fn! {
104        /// A semaphore holding `permits` permits.
105        ///
106        /// # Panics
107        ///
108        /// When `permits` exceeds [`MAX_PERMITS`](Self::MAX_PERMITS).
109        pub fn new(permits: usize) -> Self {
110            Self { raw: RawSemaphore::new(permits) }
111        }
112    }
113
114    loom_const_fn! {
115        /// A semaphore holding `permits` permits, usable in a `static`.
116        ///
117        /// # Panics
118        ///
119        /// When `permits` exceeds [`MAX_PERMITS`](Self::MAX_PERMITS).
120        pub fn const_new(permits: usize) -> Self {
121            Self { raw: RawSemaphore::new(permits) }
122        }
123    }
124
125    /// The permits free right now.
126    pub fn available_permits(&self) -> usize {
127        self.raw.available()
128    }
129
130    /// Adds `n` new permits, handing them to waiters first. It makes the
131    /// progress a returned permit does (see [`SemaphorePermit`]). Permits
132    /// added to a closed semaphore are counted, but nothing is handed out.
133    ///
134    /// # Panics
135    ///
136    /// When the semaphore would hold more than
137    /// [`MAX_PERMITS`](Self::MAX_PERMITS).
138    pub fn add_permits(&self, n: usize) {
139        self.raw.add(n);
140    }
141
142    /// Removes up to `n` free permits for good and returns how many it
143    /// removed, which is fewer than `n` when fewer are free. Permits held by
144    /// outstanding [`SemaphorePermit`]s are not touched.
145    pub fn forget_permits(&self, n: usize) -> usize {
146        self.raw.forget(n)
147    }
148
149    /// Waits for one permit. Uncontended, the first poll is one atomic
150    /// read-modify-write. Contended, it is lock-free, not wait-free: see
151    /// [Progress](super#progress).
152    ///
153    /// Completes with an [`AcquireError`] once the semaphore is
154    /// [closed](Self::close).
155    pub fn acquire(&self) -> Acquire<'_> {
156        self.acquire_many(1)
157    }
158
159    /// Waits for `n` permits, which it takes together.
160    ///
161    /// Completes with an [`AcquireError`] once the semaphore is
162    /// [closed](Self::close).
163    ///
164    /// # Panics
165    ///
166    /// When `n` exceeds [`MAX_PERMITS`](Self::MAX_PERMITS), which only a
167    /// target with a word narrower than 64 bits allows and which no
168    /// semaphore can ever satisfy.
169    pub fn acquire_many(&self, n: u32) -> Acquire<'_> {
170        Acquire {
171            sem: self,
172            n: checked_permits(n),
173            wait: PermitWait::new(),
174        }
175    }
176
177    /// Waits for one permit through an `Arc`, so the permit can outlive the
178    /// borrow.
179    ///
180    /// Completes with an [`AcquireError`] once the semaphore is
181    /// [closed](Self::close).
182    pub fn acquire_owned(self: Arc<Self>) -> AcquireOwned {
183        self.acquire_many_owned(1)
184    }
185
186    /// Waits for `n` permits through an `Arc`, so the permits can outlive
187    /// the borrow.
188    ///
189    /// # Panics
190    ///
191    /// When `n` exceeds [`MAX_PERMITS`](Self::MAX_PERMITS), as for
192    /// [`acquire_many`](Self::acquire_many).
193    pub fn acquire_many_owned(self: Arc<Self>, n: u32) -> AcquireOwned {
194        AcquireOwned {
195            sem: self,
196            n: checked_permits(n),
197            wait: PermitWait::new(),
198        }
199    }
200
201    /// Takes one permit if it is free and no handoff is owed.
202    pub fn try_acquire(&self) -> Result<SemaphorePermit<'_>, TryAcquireError> {
203        self.try_acquire_many(1)
204    }
205
206    /// Takes `n` permits together if they are free and no handoff is owed.
207    ///
208    /// # Panics
209    ///
210    /// When `n` exceeds [`MAX_PERMITS`](Self::MAX_PERMITS), as for
211    /// [`acquire_many`](Self::acquire_many).
212    pub fn try_acquire_many(&self, n: u32) -> Result<SemaphorePermit<'_>, TryAcquireError> {
213        let permits = checked_permits(n);
214        self.raw
215            .try_take(permits)
216            .map(|()| SemaphorePermit { sem: self, permits })
217            .map_err(TryAcquireError::of)
218    }
219
220    /// Takes one permit through an `Arc` if it is free and no handoff is
221    /// owed.
222    pub fn try_acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, TryAcquireError> {
223        self.try_acquire_many_owned(1)
224    }
225
226    /// Takes `n` permits together through an `Arc` if they are free and no
227    /// handoff is owed.
228    ///
229    /// # Panics
230    ///
231    /// When `n` exceeds [`MAX_PERMITS`](Self::MAX_PERMITS), as for
232    /// [`acquire_many`](Self::acquire_many).
233    pub fn try_acquire_many_owned(
234        self: Arc<Self>,
235        n: u32,
236    ) -> Result<OwnedSemaphorePermit, TryAcquireError> {
237        let permits = checked_permits(n);
238        match self.raw.try_take(permits) {
239            Ok(()) => Ok(OwnedSemaphorePermit { sem: self, permits }),
240            Err(seen) => Err(TryAcquireError::of(seen)),
241        }
242    }
243
244    /// [`acquire`](Self::acquire) for a thread that may block: it parks the
245    /// calling thread until the permit is free. On a nuclei shard (inside one of its tasks) it does not park the thread: it runs a nested `block_on`, so the shard keeps running its other tasks meanwhile.
246    /// Inside another runtime's task it blocks that task's executor thread,
247    /// so do not.
248    ///
249    /// # Panics
250    ///
251    /// On `wasm32`, where no other thread can return a permit, when the
252    /// permit is not free: the call could never end.
253    pub fn blocking_acquire(&self) -> Result<SemaphorePermit<'_>, AcquireError> {
254        internal::blocking(self.acquire())
255    }
256
257    /// [`acquire_many`](Self::acquire_many) for a thread that may block: see
258    /// [`blocking_acquire`](Self::blocking_acquire).
259    ///
260    /// # Panics
261    ///
262    /// As for [`blocking_acquire`](Self::blocking_acquire), and when `n`
263    /// exceeds [`MAX_PERMITS`](Self::MAX_PERMITS).
264    pub fn blocking_acquire_many(&self, n: u32) -> Result<SemaphorePermit<'_>, AcquireError> {
265        internal::blocking(self.acquire_many(n))
266    }
267
268    /// [`acquire_owned`](Self::acquire_owned) for a thread that may block:
269    /// see [`blocking_acquire`](Self::blocking_acquire).
270    ///
271    /// # Panics
272    ///
273    /// As for [`blocking_acquire`](Self::blocking_acquire).
274    pub fn blocking_acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, AcquireError> {
275        internal::blocking(self.acquire_owned())
276    }
277
278    /// [`acquire_many_owned`](Self::acquire_many_owned) for a thread that may
279    /// block: see [`blocking_acquire`](Self::blocking_acquire).
280    ///
281    /// # Panics
282    ///
283    /// As for [`blocking_acquire`](Self::blocking_acquire), and when `n`
284    /// exceeds [`MAX_PERMITS`](Self::MAX_PERMITS).
285    pub fn blocking_acquire_many_owned(
286        self: Arc<Self>,
287        n: u32,
288    ) -> Result<OwnedSemaphorePermit, AcquireError> {
289        internal::blocking(self.acquire_many_owned(n))
290    }
291
292    /// Closes the semaphore. Every wait still queued ends in an
293    /// [`AcquireError`], and every later acquire fails at once: `try_acquire`
294    /// and its siblings return [`TryAcquireError::Closed`], whatever permits
295    /// are free. A wait that was handed its permits before the close keeps
296    /// them. Permits returned or added afterwards are counted, and
297    /// [`available_permits`](Self::available_permits) shows them, but none is
298    /// ever handed out. Closing twice does nothing.
299    ///
300    /// It makes the progress a returned permit does (see
301    /// [`SemaphorePermit`]).
302    pub fn close(&self) {
303        self.raw.close();
304    }
305
306    /// Whether [`close`](Self::close) was called.
307    pub fn is_closed(&self) -> bool {
308        self.raw.is_closed()
309    }
310}
311
312/// `n` as a permit count, or a panic when no semaphore could ever satisfy it.
313#[expect(
314    clippy::panic,
315    reason = "documented under `# Panics` of the calls that take permits"
316)]
317fn checked_permits(n: u32) -> usize {
318    match usize::try_from(n) {
319        Ok(permits) if permits <= MAX => permits,
320        _ => panic!("a request exceeds Semaphore::MAX_PERMITS"),
321    }
322}
323
324impl fmt::Debug for Semaphore {
325    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
326        f.debug_struct("Semaphore")
327            .field("available_permits", &self.available_permits())
328            .field("closed", &self.is_closed())
329            .finish()
330    }
331}
332
333/// Permits taken from a [`Semaphore`], returned when dropped.
334///
335/// The release is one atomic read-modify-write when no waiter needs telling.
336/// Otherwise it is lock-free, not wait-free: it returns after one pass over
337/// the queue, plus at most one for each operation of another thread that
338/// lands on the queue meanwhile (see [Progress](super#progress)).
339#[must_use = "dropping a permit returns it at once"]
340pub struct SemaphorePermit<'a> {
341    sem: &'a Semaphore,
342    permits: usize,
343}
344
345impl<'a> SemaphorePermit<'a> {
346    /// Keeps the permits out of the semaphore for good.
347    pub fn forget(mut self) {
348        self.permits = 0;
349    }
350
351    /// Merges `other` into this permit, so one drop returns both.
352    ///
353    /// # Panics
354    ///
355    /// When the permits come from different semaphores, or their total
356    /// overflows `usize`.
357    #[track_caller]
358    pub fn merge(&mut self, mut other: Self) {
359        assert!(
360            core::ptr::eq(self.sem, other.sem),
361            "merging permits from different semaphore instances"
362        );
363        // Both are of one semaphore, whose permits never pass `MAX_PERMITS`.
364        self.permits += other.permits;
365        other.permits = 0;
366    }
367
368    /// Splits `n` permits off this one, or returns `None` when it holds
369    /// fewer than `n`.
370    pub fn split(&mut self, n: usize) -> Option<Self> {
371        if n > self.permits {
372            return None;
373        }
374        self.permits -= n;
375        Some(Self {
376            sem: self.sem,
377            permits: n,
378        })
379    }
380
381    /// The semaphore these permits came from.
382    pub fn semaphore(&self) -> &'a Semaphore {
383        self.sem
384    }
385
386    /// How many permits this holds.
387    pub fn num_permits(&self) -> usize {
388        self.permits
389    }
390
391    /// How many permits this holds; the same as
392    /// [`num_permits`](Self::num_permits).
393    pub fn count(&self) -> usize {
394        self.permits
395    }
396}
397
398impl Drop for SemaphorePermit<'_> {
399    fn drop(&mut self) {
400        self.sem.raw.release(self.permits);
401    }
402}
403
404impl fmt::Debug for SemaphorePermit<'_> {
405    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
406        f.debug_struct("SemaphorePermit")
407            .field("permits", &self.permits)
408            .finish()
409    }
410}
411
412/// Permits taken from a [`Semaphore`] held through an `Arc`. The release
413/// makes the progress [`SemaphorePermit`]'s does.
414#[must_use = "dropping a permit returns it at once"]
415pub struct OwnedSemaphorePermit {
416    sem: Arc<Semaphore>,
417    permits: usize,
418}
419
420impl OwnedSemaphorePermit {
421    /// Keeps the permits out of the semaphore for good.
422    pub fn forget(mut self) {
423        self.permits = 0;
424    }
425
426    /// Merges `other` into this permit, so one drop returns both.
427    ///
428    /// # Panics
429    ///
430    /// When the permits come from different semaphores, or their total
431    /// overflows `usize`.
432    #[track_caller]
433    pub fn merge(&mut self, mut other: Self) {
434        assert!(
435            Arc::ptr_eq(&self.sem, &other.sem),
436            "merging permits from different semaphore instances"
437        );
438        // Both are of one semaphore, whose permits never pass `MAX_PERMITS`.
439        self.permits += other.permits;
440        other.permits = 0;
441    }
442
443    /// Splits `n` permits off this one, or returns `None` when it holds
444    /// fewer than `n`.
445    pub fn split(&mut self, n: usize) -> Option<Self> {
446        if n > self.permits {
447            return None;
448        }
449        self.permits -= n;
450        Some(Self {
451            sem: Arc::clone(&self.sem),
452            permits: n,
453        })
454    }
455
456    /// The semaphore these permits came from.
457    pub fn semaphore(&self) -> &Arc<Semaphore> {
458        &self.sem
459    }
460
461    /// How many permits this holds.
462    pub fn num_permits(&self) -> usize {
463        self.permits
464    }
465
466    /// How many permits this holds; the same as
467    /// [`num_permits`](Self::num_permits).
468    pub fn count(&self) -> usize {
469        self.permits
470    }
471}
472
473impl Drop for OwnedSemaphorePermit {
474    fn drop(&mut self) {
475        self.sem.raw.release(self.permits);
476    }
477}
478
479impl fmt::Debug for OwnedSemaphorePermit {
480    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
481        f.debug_struct("OwnedSemaphorePermit")
482            .field("permits", &self.permits)
483            .finish()
484    }
485}
486
487/// The future [`Semaphore::acquire`] and [`Semaphore::acquire_many`] return.
488///
489/// Dropping it withdraws the wait, and passes on permits handed to it
490/// meanwhile.
491#[must_use = "futures do nothing unless polled"]
492pub struct Acquire<'a> {
493    sem: &'a Semaphore,
494    n: usize,
495    wait: PermitWait,
496}
497
498impl<'a> Future for Acquire<'a> {
499    type Output = Result<SemaphorePermit<'a>, AcquireError>;
500
501    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
502        let this = self.get_mut();
503        this.wait
504            .poll_closable(&this.sem.raw, this.n, cx.waker())
505            .map(|taken| {
506                taken
507                    .map(|()| SemaphorePermit {
508                        sem: this.sem,
509                        permits: this.n,
510                    })
511                    .map_err(|Closed| AcquireError::closed())
512            })
513    }
514}
515
516impl Drop for Acquire<'_> {
517    fn drop(&mut self) {
518        self.wait.cancel(&self.sem.raw, self.n);
519    }
520}
521
522impl fmt::Debug for Acquire<'_> {
523    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
524        f.debug_struct("Acquire").field("n", &self.n).finish()
525    }
526}
527
528/// The future [`Semaphore::acquire_owned`] and
529/// [`Semaphore::acquire_many_owned`] return.
530#[must_use = "futures do nothing unless polled"]
531pub struct AcquireOwned {
532    sem: Arc<Semaphore>,
533    n: usize,
534    wait: PermitWait,
535}
536
537impl Future for AcquireOwned {
538    type Output = Result<OwnedSemaphorePermit, AcquireError>;
539
540    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
541        let this = self.get_mut();
542        this.wait
543            .poll_closable(&this.sem.raw, this.n, cx.waker())
544            .map(|taken| {
545                taken
546                    .map(|()| OwnedSemaphorePermit {
547                        sem: Arc::clone(&this.sem),
548                        permits: this.n,
549                    })
550                    .map_err(|Closed| AcquireError::closed())
551            })
552    }
553}
554
555impl Drop for AcquireOwned {
556    fn drop(&mut self) {
557        self.wait.cancel(&self.sem.raw, self.n);
558    }
559}
560
561impl fmt::Debug for AcquireOwned {
562    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
563        f.debug_struct("AcquireOwned").field("n", &self.n).finish()
564    }
565}
566
567#[cfg(all(test, not(loom)))]
568mod tests {
569    use super::*;
570    use crate::sync::test_support::{Polled, block_on};
571    use std::sync::atomic::{AtomicUsize, Ordering as StdOrdering};
572
573    const ITERATIONS: usize = if cfg!(miri) { 20 } else { 2_000 };
574
575    #[test]
576    fn uncontended_permits_come_and_go() {
577        let sem = Semaphore::new(3);
578        let two = sem.try_acquire_many(2).expect("two of three are free");
579        assert_eq!(sem.available_permits(), 1);
580        assert_eq!(
581            sem.try_acquire_many(2).unwrap_err(),
582            TryAcquireError::NoPermits
583        );
584        drop(two);
585        assert_eq!(sem.available_permits(), 3);
586    }
587
588    /// A semaphore of one permit whose waiter has been passed over
589    /// `MAX_BYPASS` times, so it is owed the permit; returns the holder.
590    fn owe<'a>(
591        sem: &'a Semaphore,
592        waiter: &mut Polled<Acquire<'a>>,
593        mut holder: SemaphorePermit<'a>,
594    ) -> SemaphorePermit<'a> {
595        for round in 0..crate::sync::MAX_BYPASS {
596            drop(holder);
597            assert_eq!(waiter.wakes() as u32, round + 1, "a release nudges it");
598            holder = sem.try_acquire().expect("a free permit is taken at once");
599            waiter.pending();
600        }
601        holder
602    }
603
604    #[test]
605    fn a_release_wakes_the_front_request_and_none_it_cannot_serve() {
606        let sem = Semaphore::new(2);
607        let held = sem.try_acquire_many(2).expect("free");
608        let mut big = Polled::new(sem.acquire_many(2));
609        big.pending();
610        let mut small = Polled::new(sem.acquire());
611        small.pending();
612        drop(held);
613        assert_eq!(big.wakes(), 1, "the front request is woken");
614        assert_eq!(small.wakes(), 0, "the freed permits are all the front's");
615        let big = big.ready().expect("open");
616        small.pending();
617        drop(big);
618        assert_eq!(small.wakes(), 1);
619        assert_eq!(small.ready().expect("open").num_permits(), 1);
620    }
621
622    #[test]
623    fn a_free_permit_is_taken_past_a_waiter_until_it_is_owed() {
624        let sem = Semaphore::new(1);
625        let holder = sem.try_acquire().expect("free");
626        let mut waiter = Polled::new(sem.acquire());
627        waiter.pending();
628        let holder = owe(&sem, &mut waiter, holder);
629        assert_eq!(
630            sem.try_acquire().unwrap_err(),
631            TryAcquireError::NoPermits,
632            "nobody barges past an owed waiter"
633        );
634        drop(holder);
635        assert_eq!(sem.available_permits(), 0, "the permit went to the waiter");
636        drop(waiter.ready());
637        assert!(
638            sem.try_acquire().is_ok(),
639            "barging resumes once it is served"
640        );
641    }
642
643    #[test]
644    fn a_large_request_is_not_starved_by_small_callers() {
645        let sem = Semaphore::new(2);
646        let mut small = sem.try_acquire().expect("free");
647        let other = sem.try_acquire().expect("free");
648        let mut big = Polled::new(sem.acquire_many(2));
649        big.pending();
650        for round in 0..crate::sync::MAX_BYPASS {
651            drop(small);
652            assert_eq!(
653                big.wakes() as u32,
654                round + 1,
655                "a release wakes the front whatever it asks for"
656            );
657            small = sem
658                .try_acquire()
659                .expect("a small caller takes the freed permit");
660            big.pending();
661        }
662        drop(small);
663        assert!(
664            sem.try_acquire().is_err(),
665            "the freed permit is kept for the owed request"
666        );
667        drop(other);
668        assert_eq!(big.ready().expect("open").num_permits(), 2);
669    }
670
671    #[test]
672    fn a_withdrawn_front_waiter_unblocks_the_one_behind() {
673        let sem = Semaphore::new(2);
674        let first = sem.try_acquire().expect("free");
675        let second = sem.try_acquire().expect("free");
676        let mut big = Polled::new(sem.acquire_many(2));
677        big.pending();
678        let mut small = Polled::new(sem.acquire());
679        small.pending();
680        drop(first);
681        assert_eq!(small.wakes(), 0, "one permit is not enough for the front");
682        drop(big);
683        assert_eq!(small.wakes(), 1);
684        drop(small.ready());
685        drop(second);
686        assert_eq!(sem.available_permits(), 2);
687    }
688
689    #[test]
690    fn a_handoff_dropped_unseen_passes_its_permit_on() {
691        let sem = Semaphore::new(1);
692        let holder = sem.try_acquire().expect("free");
693        let mut owed = Polled::new(sem.acquire());
694        owed.pending();
695        let holder = owe(&sem, &mut owed, holder);
696        let mut next = Polled::new(sem.acquire());
697        next.pending();
698        drop(holder);
699        drop(owed);
700        assert_eq!(next.wakes(), 1, "the dropped handoff moved on");
701        drop(next.ready());
702        assert_eq!(sem.available_permits(), 1);
703    }
704
705    #[test]
706    fn a_request_for_nothing_never_waits() {
707        let sem = Semaphore::new(0);
708        let mut waiter = Polled::new(sem.acquire());
709        waiter.pending();
710        assert_eq!(
711            Polled::new(sem.acquire_many(0))
712                .ready()
713                .expect("open")
714                .num_permits(),
715            0
716        );
717        sem.add_permits(1);
718        waiter.ready().expect("open").forget();
719        assert_eq!(sem.available_permits(), 0);
720    }
721
722    #[test]
723    fn owned_permits_outlive_the_borrow() {
724        let sem = Arc::new(Semaphore::new(1));
725        let permit = block_on(Arc::clone(&sem).acquire_owned()).expect("open");
726        let mut next = Polled::new(Arc::clone(&sem).acquire_owned());
727        next.pending();
728        drop(permit);
729        assert_eq!(next.ready().expect("open").num_permits(), 1);
730    }
731
732    #[test]
733    fn withdrawn_waiters_do_not_pile_up_behind_a_live_one() {
734        let sem = Semaphore::new(0);
735        let mut front = Polled::new(sem.acquire());
736        front.pending();
737        for _ in 0..1_000 {
738            Polled::new(sem.acquire()).pending();
739        }
740        assert!(
741            sem.raw.linked() <= 4,
742            "{} nodes still linked",
743            sem.raw.linked()
744        );
745        sem.add_permits(1);
746        front.ready().expect("open").forget();
747    }
748
749    #[test]
750    #[should_panic(expected = "at most Semaphore::MAX_PERMITS")]
751    fn overflowing_the_permit_count_fails_loudly() {
752        Semaphore::new(Semaphore::MAX_PERMITS).add_permits(1);
753    }
754
755    #[test]
756    fn closing_ends_every_queued_wait_and_refuses_later_ones() {
757        let sem = Semaphore::new(1);
758        let held = sem.try_acquire().expect("free");
759        let mut first = Polled::new(sem.acquire());
760        first.pending();
761        let mut second = Polled::new(sem.acquire_many(1));
762        second.pending();
763        assert!(!sem.is_closed());
764        sem.close();
765        assert!(sem.is_closed());
766        assert_eq!((first.wakes(), second.wakes()), (1, 1));
767        assert_eq!(first.ready().unwrap_err().to_string(), "semaphore closed");
768        assert!(second.ready().is_err());
769        assert!(
770            Polled::new(sem.acquire()).ready().is_err(),
771            "a later wait fails at once"
772        );
773        assert_eq!(sem.try_acquire().unwrap_err(), TryAcquireError::Closed);
774        drop(held);
775        assert_eq!(sem.available_permits(), 1, "the returned permit counts");
776        assert_eq!(
777            sem.try_acquire().unwrap_err(),
778            TryAcquireError::Closed,
779            "but is not handed out"
780        );
781        sem.add_permits(2);
782        assert_eq!(sem.available_permits(), 3, "added permits count too");
783        sem.close();
784    }
785
786    #[test]
787    fn a_wait_handed_its_permit_before_the_close_keeps_it() {
788        let sem = Semaphore::new(1);
789        let held = sem.try_acquire().expect("free");
790        let mut waiter = Polled::new(sem.acquire());
791        waiter.pending();
792        // Owe it the permit, then hand it over by the release.
793        let held = owe(&sem, &mut waiter, held);
794        drop(held);
795        sem.close();
796        let permit = waiter.ready().expect("the permit was its before the close");
797        drop(permit);
798        assert_eq!(sem.available_permits(), 1);
799    }
800
801    #[test]
802    fn a_closed_wait_dropped_unseen_returns_nothing_twice() {
803        let sem = Semaphore::new(1);
804        let held = sem.try_acquire().expect("free");
805        let mut waiter = Polled::new(sem.acquire());
806        waiter.pending();
807        sem.close();
808        drop(waiter);
809        drop(held);
810        assert_eq!(sem.available_permits(), 1, "no permit made up");
811    }
812
813    #[test]
814    fn forgetting_permits_takes_only_free_ones() {
815        let sem = Semaphore::new(3);
816        let held = sem.try_acquire().expect("free");
817        assert_eq!(sem.forget_permits(5), 2);
818        assert_eq!(sem.available_permits(), 0);
819        assert_eq!(sem.forget_permits(0), 0);
820        drop(held);
821        assert_eq!(sem.available_permits(), 1);
822        sem.close();
823        assert_eq!(sem.forget_permits(1), 1, "forgetting keeps the closed flag");
824        assert!(sem.is_closed());
825    }
826
827    #[test]
828    fn permits_merge_and_split() {
829        let sem = Semaphore::new(5);
830        let mut all = sem.try_acquire_many(2).expect("free");
831        all.merge(sem.try_acquire_many(3).expect("free"));
832        assert_eq!(all.num_permits(), 5);
833        let part = all.split(2).expect("held");
834        assert_eq!((all.num_permits(), part.num_permits()), (3, 2));
835        assert!(all.split(4).is_none());
836        assert!(core::ptr::eq(all.semaphore(), &sem));
837        drop(part);
838        assert_eq!(sem.available_permits(), 2);
839        drop(all);
840        assert_eq!(sem.available_permits(), 5);
841    }
842
843    #[test]
844    #[should_panic(expected = "different semaphore instances")]
845    fn merging_permits_of_two_semaphores_fails_loudly() {
846        let (a, b) = (Semaphore::new(1), Semaphore::new(1));
847        let mut one = a.try_acquire().expect("free");
848        one.merge(b.try_acquire().expect("free"));
849    }
850
851    #[test]
852    fn threads_contending_never_exceed_the_permits() {
853        let sem = Arc::new(Semaphore::new(2));
854        let inside = Arc::new(AtomicUsize::new(0));
855        let threads: Vec<_> = (0..4)
856            .map(|_| {
857                let (sem, inside) = (Arc::clone(&sem), Arc::clone(&inside));
858                std::thread::spawn(move || {
859                    for i in 0..ITERATIONS {
860                        let permit = block_on(sem.acquire_many(1 + (i % 2) as u32)).expect("open");
861                        let now = inside.fetch_add(permit.num_permits(), StdOrdering::SeqCst);
862                        assert!(
863                            now + permit.num_permits() <= 2,
864                            "more permits out than exist"
865                        );
866                        inside.fetch_sub(permit.num_permits(), StdOrdering::SeqCst);
867                    }
868                })
869            })
870            .collect();
871        for thread in threads {
872            thread.join().expect("worker");
873        }
874        assert_eq!(sem.available_permits(), 2);
875    }
876}
nuclei/sync/semaphore.rs

nuclei/sync/
semaphore.rs
