nuclei/sync/local/
semaphore.rs

1//! A strictly first-come, first-served counting semaphore for one shard.
2//!
3//! The permits live in a `Cell` and the waiters in the shared
4//! [`Waiters`](super::waiters) list, so no operation touches an atomic. A
5//! request is served in arrival order and nobody barges: permits that come
6//! free are kept for the front waiter until its whole request fits, and
7//! neither a later waiter nor [`Semaphore::try_acquire`] takes them meanwhile.
8//! Dropping a waiting [`Acquire`] gives up its place; dropping one that was
9//! already granted passes its permits on.
10
11use core::cell::Cell;
12use core::fmt;
13use core::future::Future;
14use core::pin::Pin;
15use core::task::{Context, Poll};
16
17use super::waiters::{Key, State, Waiters};
18use crate::sync::{AcquireError, TryAcquireError};
19
20/// A counting semaphore whose waits are futures, for the tasks of one shard.
21///
22/// A request for `n` permits waits in line behind the requests that arrived
23/// before it, and is granted when it reaches the front and `n` permits are
24/// free. A large request at the front is therefore never bypassed by smaller
25/// ones behind it, and [`try_acquire`](Self::try_acquire) fails while anyone
26/// waits. [`close`](Self::close) ends every wait with an [`AcquireError`].
27///
28/// It is `!Send` and `!Sync`: share it between the tasks of a shard with an
29/// `Rc`.
30///
31/// ```
32/// use nuclei::sync::TryAcquireError;
33/// use nuclei::sync::local::Semaphore;
34///
35/// let budget = Semaphore::new(4);
36/// let first = budget.try_acquire_many(3).expect("three of four are free");
37/// assert_eq!(budget.try_acquire_many(2).unwrap_err(), TryAcquireError::NoPermits);
38/// drop(first);
39/// assert_eq!(budget.available_permits(), 4);
40/// ```
41pub struct Semaphore {
42    permits: Cell<usize>,
43    closed: Cell<bool>,
44    waiters: Waiters,
45}
46
47impl Semaphore {
48    /// The most permits a semaphore can hold: `usize::MAX >> 3`.
49    pub const MAX_PERMITS: usize = usize::MAX >> 3;
50
51    /// A semaphore holding `permits` permits.
52    ///
53    /// # Panics
54    ///
55    /// When `permits` exceeds [`MAX_PERMITS`](Self::MAX_PERMITS).
56    pub const fn new(permits: usize) -> Self {
57        assert!(
58            permits <= Self::MAX_PERMITS,
59            "a semaphore cannot hold more than Semaphore::MAX_PERMITS"
60        );
61        Self {
62            permits: Cell::new(permits),
63            closed: Cell::new(false),
64            waiters: Waiters::new(),
65        }
66    }
67
68    /// The permits free right now, including the ones kept for the front
69    /// waiter until its request fits.
70    pub fn available_permits(&self) -> usize {
71        self.permits.get()
72    }
73
74    /// Adds `n` new permits, handing them to waiters first. Permits added to
75    /// a closed semaphore are counted, but nothing is handed out.
76    ///
77    /// # Panics
78    ///
79    /// When the semaphore would hold more than
80    /// [`MAX_PERMITS`](Self::MAX_PERMITS).
81    #[track_caller]
82    pub fn add_permits(&self, n: usize) {
83        let total = self.permits.get().saturating_add(n);
84        assert!(
85            total <= Self::MAX_PERMITS,
86            "a semaphore cannot hold more than Semaphore::MAX_PERMITS"
87        );
88        self.permits.set(total);
89        self.serve();
90    }
91
92    /// Removes up to `n` free permits for good and returns how many it
93    /// removed, which is fewer than `n` when fewer are free. Permits held by
94    /// outstanding [`SemaphorePermit`]s are not touched.
95    pub fn forget_permits(&self, n: usize) -> usize {
96        let removed = n.min(self.permits.get());
97        self.permits.set(self.permits.get() - removed);
98        removed
99    }
100
101    /// Waits for one permit.
102    ///
103    /// Completes with an [`AcquireError`] once the semaphore is
104    /// [closed](Self::close).
105    pub fn acquire(&self) -> Acquire<'_> {
106        self.acquire_many(1)
107    }
108
109    /// Waits for `n` permits, which it takes together.
110    ///
111    /// Completes with an [`AcquireError`] once the semaphore is
112    /// [closed](Self::close).
113    ///
114    /// # Panics
115    ///
116    /// When `n` exceeds [`MAX_PERMITS`](Self::MAX_PERMITS), which only a
117    /// target with a word narrower than 64 bits allows and which no
118    /// semaphore can ever satisfy.
119    #[track_caller]
120    pub fn acquire_many(&self, n: u32) -> Acquire<'_> {
121        Acquire {
122            sem: self,
123            n: checked_permits(n),
124            key: None,
125        }
126    }
127
128    /// Takes one permit if it is free and nobody waits.
129    pub fn try_acquire(&self) -> Result<SemaphorePermit<'_>, TryAcquireError> {
130        self.try_acquire_many(1)
131    }
132
133    /// Takes `n` permits together if they are free and nobody waits.
134    ///
135    /// # Panics
136    ///
137    /// When `n` exceeds [`MAX_PERMITS`](Self::MAX_PERMITS), as for
138    /// [`acquire_many`](Self::acquire_many).
139    #[track_caller]
140    pub fn try_acquire_many(&self, n: u32) -> Result<SemaphorePermit<'_>, TryAcquireError> {
141        let permits = checked_permits(n);
142        if self.closed.get() {
143            return Err(TryAcquireError::Closed);
144        }
145        if !self.waiters.is_empty() || self.permits.get() < permits {
146            return Err(TryAcquireError::NoPermits);
147        }
148        self.permits.set(self.permits.get() - permits);
149        Ok(SemaphorePermit { sem: self, permits })
150    }
151
152    /// Closes the semaphore. Every wait still queued ends in an
153    /// [`AcquireError`], and every later acquire fails at once:
154    /// `try_acquire` and its sibling return [`TryAcquireError::Closed`],
155    /// whatever permits are free. A wait that was granted its permits before
156    /// the close keeps them. Permits returned or added afterwards are
157    /// counted, and [`available_permits`](Self::available_permits) shows
158    /// them, but none is ever handed out. Closing twice does nothing.
159    pub fn close(&self) {
160        if self.closed.replace(true) {
161            return;
162        }
163        while let Some(waker) = self.waiters.close_front() {
164            if let Some(waker) = waker {
165                waker.wake();
166            }
167        }
168    }
169
170    /// Whether [`close`](Self::close) was called.
171    pub fn is_closed(&self) -> bool {
172        self.closed.get()
173    }
174
175    /// Takes `permits` back and serves the line.
176    pub(crate) fn release(&self, permits: usize) {
177        self.permits.set(self.permits.get().saturating_add(permits));
178        self.serve();
179    }
180
181    /// Grants the front waiters while their requests fit. Each waker is woken
182    /// outside any borrow, and the state is read again after the wake.
183    fn serve(&self) {
184        if self.closed.get() {
185            return;
186        }
187        while let Some((need, waker)) = self.waiters.grant_front_if(self.permits.get()) {
188            self.permits.set(self.permits.get() - need);
189            if let Some(waker) = waker {
190                waker.wake();
191            }
192        }
193    }
194}
195
196/// `n` as a permit count, or a panic when no semaphore could ever satisfy it.
197#[expect(
198    clippy::panic,
199    reason = "documented under `# Panics` of the calls that take permits"
200)]
201#[track_caller]
202fn checked_permits(n: u32) -> usize {
203    match usize::try_from(n) {
204        Ok(permits) if permits <= Semaphore::MAX_PERMITS => permits,
205        _ => panic!("a request exceeds Semaphore::MAX_PERMITS"),
206    }
207}
208
209impl fmt::Debug for Semaphore {
210    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
211        f.debug_struct("Semaphore")
212            .field("available_permits", &self.available_permits())
213            .field("closed", &self.is_closed())
214            .finish()
215    }
216}
217
218/// Permits taken from a [`Semaphore`], returned when dropped.
219#[must_use = "dropping a permit returns it at once"]
220pub struct SemaphorePermit<'a> {
221    sem: &'a Semaphore,
222    permits: usize,
223}
224
225impl<'a> SemaphorePermit<'a> {
226    pub(crate) fn new(sem: &'a Semaphore, permits: usize) -> Self {
227        Self { sem, permits }
228    }
229
230    /// Keeps the permits out of the semaphore for good.
231    pub fn forget(mut self) {
232        self.permits = 0;
233    }
234
235    /// How many permits this holds.
236    pub fn num_permits(&self) -> usize {
237        self.permits
238    }
239
240    /// Merges `other` into this permit, so one drop returns both.
241    ///
242    /// # Panics
243    ///
244    /// When the permits come from different semaphores.
245    #[track_caller]
246    pub fn merge(&mut self, mut other: Self) {
247        assert!(
248            core::ptr::eq(self.sem, other.sem),
249            "merging permits from different semaphore instances"
250        );
251        // Both are of one semaphore, whose permits never pass `MAX_PERMITS`.
252        self.permits += other.permits;
253        other.permits = 0;
254    }
255}
256
257impl Drop for SemaphorePermit<'_> {
258    fn drop(&mut self) {
259        if self.permits > 0 {
260            self.sem.release(self.permits);
261        }
262    }
263}
264
265impl fmt::Debug for SemaphorePermit<'_> {
266    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
267        f.debug_struct("SemaphorePermit")
268            .field("permits", &self.permits)
269            .finish()
270    }
271}
272
273/// The wait for permits of [`Semaphore::acquire`] and
274/// [`Semaphore::acquire_many`].
275///
276/// Cancel safe: dropping it gives up its place in line; if its permits were
277/// already granted, they go back and the next waiter is served.
278#[must_use = "futures do nothing unless polled"]
279pub struct Acquire<'a> {
280    sem: &'a Semaphore,
281    n: usize,
282    key: Option<Key>,
283}
284
285impl Acquire<'_> {
286    /// Gives up the place in line, or passes on the permits granted; idempotent.
287    /// One that never queued (the common acquire, served at once) has nothing
288    /// to give up: one inline look.
289    #[inline]
290    pub(crate) fn cancel(&mut self) {
291        if self.key.is_some() {
292            self.leave_line();
293        }
294    }
295
296    /// [`cancel`](Self::cancel) of an acquire in line, out of line.
297    #[inline(never)]
298    fn leave_line(&mut self) {
299        let Some(key) = self.key.take() else {
300            return;
301        };
302        let (state, was_front, waker) = self.sem.waiters.cancel(key);
303        drop(waker);
304        match state {
305            State::Waiting if was_front => self.sem.serve(),
306            State::Granted => self.sem.release(self.n),
307            _ => {}
308        }
309    }
310}
311
312impl<'a> Future for Acquire<'a> {
313    type Output = Result<SemaphorePermit<'a>, AcquireError>;
314
315    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
316        let this = self.get_mut();
317        let sem = this.sem;
318        if let Some(key) = this.key {
319            match sem.waiters.refresh(key, cx.waker()) {
320                State::Waiting => return Poll::Pending,
321                State::Granted => {
322                    sem.waiters.take(key);
323                    this.key = None;
324                    return Poll::Ready(Ok(SemaphorePermit::new(sem, this.n)));
325                }
326                State::Closed => {
327                    sem.waiters.take(key);
328                    this.key = None;
329                    return Poll::Ready(Err(AcquireError::closed()));
330                }
331                State::Free => this.key = None,
332            }
333        }
334        if crate::coop::exhausted(cx.waker()) {
335            return Poll::Pending;
336        }
337        if sem.closed.get() {
338            crate::coop::spent_one();
339            return Poll::Ready(Err(AcquireError::closed()));
340        }
341        if sem.waiters.is_empty() && sem.permits.get() >= this.n {
342            sem.permits.set(sem.permits.get() - this.n);
343            crate::coop::spent_one();
344            return Poll::Ready(Ok(SemaphorePermit::new(sem, this.n)));
345        }
346        this.key = Some(sem.waiters.push(this.n, cx.waker().clone()));
347        Poll::Pending
348    }
349}
350
351impl Drop for Acquire<'_> {
352    fn drop(&mut self) {
353        self.cancel();
354    }
355}
356
357impl fmt::Debug for Acquire<'_> {
358    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
359        f.debug_struct("Acquire")
360            .field("permits", &self.n)
361            .field("waiting", &self.key.is_some())
362            .finish()
363    }
364}
nuclei/sync/local/semaphore.rs

nuclei/sync/local/
semaphore.rs
