Skip to main content

tokio/sync/
semaphore.rs

1use super::batch_semaphore as ll; // low level implementation
2use super::{AcquireError, TryAcquireError};
3#[cfg(all(tokio_unstable, feature = "tracing"))]
4use crate::util::trace;
5use std::sync::Arc;
6
7/// Counting semaphore performing asynchronous permit acquisition.
8///
9/// A semaphore maintains a set of permits. Permits are used to synchronize
10/// access to a shared resource. A semaphore differs from a mutex in that it
11/// can allow more than one concurrent caller to access the shared resource at a
12/// time.
13///
14/// When `acquire` is called and the semaphore has remaining permits, the
15/// function immediately returns a permit. However, if no remaining permits are
16/// available, `acquire` (asynchronously) waits until an outstanding permit is
17/// dropped. At this point, the freed permit is assigned to the caller.
18///
19/// This `Semaphore` is fair, which means that permits are given out in the order
20/// they were requested. This fairness is also applied when `acquire_many` gets
21/// involved, so if a call to `acquire_many` at the front of the queue requests
22/// more permits than currently available, this can prevent a call to `acquire`
23/// from completing, even if the semaphore has enough permits complete the call
24/// to `acquire`.
25///
26/// To use the `Semaphore` in a poll function, you can use the [`PollSemaphore`]
27/// utility.
28///
29/// # Memory ordering
30///
31/// If a task writes some data and then releases a permit, any task that later
32/// acquires a permit is guaranteed to see that data. This makes it safe to use
33/// a semaphore to hand data off between tasks through shared state.
34///
35/// Stated more precisely in terms of atomic memory orderings: acquiring a
36/// permit (via [`acquire`], [`acquire_many`], [`try_acquire`],
37/// [`try_acquire_many`], or their `_owned` variants), releasing permits (by
38/// dropping a [`SemaphorePermit`] or [`OwnedSemaphorePermit`], or by calling
39/// [`add_permits`] or [`forget_permits`]), and closing the semaphore (via
40/// [`close`]) are all `AcqRel` operations. They are totally ordered, and each
41/// one synchronizes-with all such operations that precede it, giving the same
42/// guarantees as `AcqRel` operations on a single atomic.
43///
44/// A failed acquisition attempt (including [`TryAcquireError::NoPermits`] and
45/// [`TryAcquireError::Closed`]), along with the [`available_permits`] and
46/// [`is_closed`] methods, behave like an `Acquire` load.
47///
48/// [`acquire`]: Semaphore::acquire
49/// [`acquire_many`]: Semaphore::acquire_many
50/// [`try_acquire`]: Semaphore::try_acquire
51/// [`try_acquire_many`]: Semaphore::try_acquire_many
52/// [`add_permits`]: Semaphore::add_permits
53/// [`forget_permits`]: Semaphore::forget_permits
54/// [`close`]: Semaphore::close
55/// [`available_permits`]: Semaphore::available_permits
56/// [`is_closed`]: Semaphore::is_closed
57///
58/// # Examples
59///
60/// Basic usage:
61///
62/// ```
63/// use tokio::sync::{Semaphore, TryAcquireError};
64///
65/// # #[tokio::main(flavor = "current_thread")]
66/// # async fn main() {
67/// let semaphore = Semaphore::new(3);
68///
69/// let a_permit = semaphore.acquire().await.unwrap();
70/// let two_permits = semaphore.acquire_many(2).await.unwrap();
71///
72/// assert_eq!(semaphore.available_permits(), 0);
73///
74/// let permit_attempt = semaphore.try_acquire();
75/// assert_eq!(permit_attempt.err(), Some(TryAcquireError::NoPermits));
76/// # }
77/// ```
78///
79/// ## Limit the number of simultaneously opened files in your program
80///
81/// Most operating systems have limits on the number of open file
82/// handles. Even in systems without explicit limits, resource constraints
83/// implicitly set an upper bound on the number of open files. If your
84/// program attempts to open a large number of files and exceeds this
85/// limit, it will result in an error.
86///
87/// This example uses a Semaphore with 100 permits. By acquiring a permit from
88/// the Semaphore before accessing a file, you ensure that your program opens
89/// no more than 100 files at a time. When trying to open the 101st
90/// file, the program will wait until a permit becomes available before
91/// proceeding to open another file.
92/// ```
93/// # #[cfg(not(target_family = "wasm"))]
94/// # {
95/// use std::io::Result;
96/// use tokio::fs::File;
97/// use tokio::sync::Semaphore;
98/// use tokio::io::AsyncWriteExt;
99///
100/// static PERMITS: Semaphore = Semaphore::const_new(100);
101///
102/// async fn write_to_file(message: &[u8]) -> Result<()> {
103///     let _permit = PERMITS.acquire().await.unwrap();
104///     let mut buffer = File::create("example.txt").await?;
105///     buffer.write_all(message).await?;
106///     Ok(()) // Permit goes out of scope here, and is available again for acquisition
107/// }
108/// # }
109/// ```
110///
111/// ## Limit the number of outgoing requests being sent at the same time
112///
113/// In some scenarios, it might be required to limit the number of outgoing
114/// requests being sent in parallel. This could be due to limits of a consumed
115/// API or the network resources of the system the application is running on.
116///
117/// This example uses an `Arc<Semaphore>` with 10 permits. Each task spawned is
118/// given a reference to the semaphore by cloning the `Arc<Semaphore>`. Before
119/// a task sends a request, it must acquire a permit from the semaphore by
120/// calling [`Semaphore::acquire`]. This ensures that at most 10 requests are
121/// sent in parallel at any given time. After a task has sent a request, it
122/// drops the permit to allow other tasks to send requests.
123///
124/// ```
125/// use std::sync::Arc;
126/// use tokio::sync::Semaphore;
127///
128/// # #[tokio::main(flavor = "current_thread")]
129/// # async fn main() {
130/// // Define maximum number of parallel requests.
131/// let semaphore = Arc::new(Semaphore::new(5));
132/// // Spawn many tasks that will send requests.
133/// let mut jhs = Vec::new();
134/// for task_id in 0..50 {
135///     let semaphore = semaphore.clone();
136///     let jh = tokio::spawn(async move {
137///         // Acquire permit before sending request.
138///         let _permit = semaphore.acquire().await.unwrap();
139///         // Send the request.
140///         let response = send_request(task_id).await;
141///         // Drop the permit after the request has been sent.
142///         drop(_permit);
143///         // Handle response.
144///         // ...
145///
146///         response
147///     });
148///     jhs.push(jh);
149/// }
150/// // Collect responses from tasks.
151/// let mut responses = Vec::new();
152/// for jh in jhs {
153///     let response = jh.await.unwrap();
154///     responses.push(response);
155/// }
156/// // Process responses.
157/// // ...
158/// # }
159/// # async fn send_request(task_id: usize) {
160/// #     // Send request.
161/// # }
162/// ```
163///
164/// ## Limit the number of incoming requests being handled at the same time
165///
166/// Similar to limiting the number of simultaneously opened files, network handles
167/// are a limited resource. Allowing an unbounded amount of requests to be processed
168/// could result in a denial-of-service, among many other issues.
169///
170/// This example uses an `Arc<Semaphore>` instead of a global variable.
171/// To limit the number of requests that can be processed at the time,
172/// we acquire a permit for each task before spawning it. Once acquired,
173/// a new task is spawned; and once finished, the permit is dropped inside
174/// of the task to allow others to spawn. Permits must be acquired via
175/// [`Semaphore::acquire_owned`] to be movable across the task boundary.
176/// (Since our semaphore is not a global variable — if it was, then `acquire` would be enough.)
177///
178/// ```no_run
179/// # #[cfg(not(target_family = "wasm"))]
180/// # {
181/// use std::sync::Arc;
182/// use tokio::sync::Semaphore;
183/// use tokio::net::TcpListener;
184///
185/// #[tokio::main]
186/// async fn main() -> std::io::Result<()> {
187///     let semaphore = Arc::new(Semaphore::new(3));
188///     let listener = TcpListener::bind("127.0.0.1:8080").await?;
189///
190///     loop {
191///         // Acquire permit before accepting the next socket.
192///         //
193///         // We use `acquire_owned` so that we can move `permit` into
194///         // other tasks.
195///         let permit = semaphore.clone().acquire_owned().await.unwrap();
196///         let (mut socket, _) = listener.accept().await?;
197///
198///         tokio::spawn(async move {
199///             // Do work using the socket.
200///             handle_connection(&mut socket).await;
201///             // Drop socket while the permit is still live.
202///             drop(socket);
203///             // Drop the permit, so more tasks can be created.
204///             drop(permit);
205///         });
206///     }
207/// }
208/// # async fn handle_connection(_socket: &mut tokio::net::TcpStream) {
209/// #   // Do work
210/// # }
211/// # }
212/// ```
213///
214/// ## Prevent tests from running in parallel
215///
216/// By default, Rust runs tests in the same file in parallel. However, in some
217/// cases, running two tests in parallel may lead to problems. For example, this
218/// can happen when tests use the same database.
219///
220/// Consider the following scenario:
221/// 1. `test_insert`: Inserts a key-value pair into the database, then retrieves
222///    the value using the same key to verify the insertion.
223/// 2. `test_update`: Inserts a key, then updates the key to a new value and
224///    verifies that the value has been accurately updated.
225/// 3. `test_others`: A third test that doesn't modify the database state. It
226///    can run in parallel with the other tests.
227///
228/// In this example, `test_insert` and `test_update` need to run in sequence to
229/// work, but it doesn't matter which test runs first. We can leverage a
230/// semaphore with a single permit to address this challenge.
231///
232/// ```
233/// # use tokio::sync::Mutex;
234/// # use std::collections::BTreeMap;
235/// # struct Database {
236/// #   map: Mutex<BTreeMap<String, i32>>,
237/// # }
238/// # impl Database {
239/// #    pub const fn setup() -> Database {
240/// #        Database {
241/// #            map: Mutex::const_new(BTreeMap::new()),
242/// #        }
243/// #    }
244/// #    pub async fn insert(&self, key: &str, value: i32) {
245/// #        self.map.lock().await.insert(key.to_string(), value);
246/// #    }
247/// #    pub async fn update(&self, key: &str, value: i32) {
248/// #        self.map.lock().await
249/// #            .entry(key.to_string())
250/// #            .and_modify(|origin| *origin = value);
251/// #    }
252/// #    pub async fn delete(&self, key: &str) {
253/// #        self.map.lock().await.remove(key);
254/// #    }
255/// #    pub async fn get(&self, key: &str) -> i32 {
256/// #        *self.map.lock().await.get(key).unwrap()
257/// #    }
258/// # }
259/// use tokio::sync::Semaphore;
260///
261/// // Initialize a static semaphore with only one permit, which is used to
262/// // prevent test_insert and test_update from running in parallel.
263/// static PERMIT: Semaphore = Semaphore::const_new(1);
264///
265/// // Initialize the database that will be used by the subsequent tests.
266/// static DB: Database = Database::setup();
267///
268/// #[tokio::test]
269/// # async fn fake_test_insert() {}
270/// async fn test_insert() {
271///     // Acquire permit before proceeding. Since the semaphore has only one permit,
272///     // the test will wait if the permit is already acquired by other tests.
273///     let permit = PERMIT.acquire().await.unwrap();
274///
275///     // Do the actual test stuff with database
276///
277///     // Insert a key-value pair to database
278///     let (key, value) = ("name", 0);
279///     DB.insert(key, value).await;
280///
281///     // Verify that the value has been inserted correctly.
282///     assert_eq!(DB.get(key).await, value);
283///
284///     // Undo the insertion, so the database is empty at the end of the test.
285///     DB.delete(key).await;
286///
287///     // Drop permit. This allows the other test to start running.
288///     drop(permit);
289/// }
290///
291/// #[tokio::test]
292/// # async fn fake_test_update() {}
293/// async fn test_update() {
294///     // Acquire permit before proceeding. Since the semaphore has only one permit,
295///     // the test will wait if the permit is already acquired by other tests.
296///     let permit = PERMIT.acquire().await.unwrap();
297///
298///     // Do the same insert.
299///     let (key, value) = ("name", 0);
300///     DB.insert(key, value).await;
301///
302///     // Update the existing value with a new one.
303///     let new_value = 1;
304///     DB.update(key, new_value).await;
305///
306///     // Verify that the value has been updated correctly.
307///     assert_eq!(DB.get(key).await, new_value);
308///
309///     // Undo any modificattion.
310///     DB.delete(key).await;
311///
312///     // Drop permit. This allows the other test to start running.
313///     drop(permit);
314/// }
315///
316/// #[tokio::test]
317/// # async fn fake_test_others() {}
318/// async fn test_others() {
319///     // This test can run in parallel with test_insert and test_update,
320///     // so it does not use PERMIT.
321/// }
322/// # #[tokio::main(flavor = "current_thread")]
323/// # async fn main() {
324/// #   test_insert().await;
325/// #   test_update().await;
326/// #   test_others().await;
327/// # }
328/// ```
329///
330/// ## Rate limiting using a token bucket
331///
332/// This example showcases the [`add_permits`] and [`SemaphorePermit::forget`] methods.
333///
334/// Many applications and systems have constraints on the rate at which certain
335/// operations should occur. Exceeding this rate can result in suboptimal
336/// performance or even errors.
337///
338/// This example implements rate limiting using a [token bucket]. A token bucket is a form of rate
339/// limiting that doesn't kick in immediately, to allow for short bursts of incoming requests that
340/// arrive at the same time.
341///
342/// With a token bucket, each incoming request consumes a token, and the tokens are refilled at a
343/// certain rate that defines the rate limit. When a burst of requests arrives, tokens are
344/// immediately given out until the bucket is empty. Once the bucket is empty, requests will have to
345/// wait for new tokens to be added.
346///
347/// Unlike the example that limits how many requests can be handled at the same time, we do not add
348/// tokens back when we finish handling a request. Instead, tokens are added only by a timer task.
349///
350/// Note that this implementation is suboptimal when the duration is small, because it consumes a
351/// lot of cpu constantly looping and sleeping.
352///
353/// [token bucket]: https://en.wikipedia.org/wiki/Token_bucket
354/// [`add_permits`]: crate::sync::Semaphore::add_permits
355/// [`SemaphorePermit::forget`]: crate::sync::SemaphorePermit::forget
356/// ```
357/// use std::sync::Arc;
358/// use tokio::sync::Semaphore;
359/// use tokio::time::{interval, Duration};
360///
361/// struct TokenBucket {
362///     sem: Arc<Semaphore>,
363///     jh: tokio::task::JoinHandle<()>,
364/// }
365///
366/// impl TokenBucket {
367///     fn new(duration: Duration, capacity: usize) -> Self {
368///         let sem = Arc::new(Semaphore::new(capacity));
369///
370///         // refills the tokens at the end of each interval
371///         let jh = tokio::spawn({
372///             let sem = sem.clone();
373///             let mut interval = interval(duration);
374///             interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
375///
376///             async move {
377///                 loop {
378///                     interval.tick().await;
379///
380///                     if sem.available_permits() < capacity {
381///                         sem.add_permits(1);
382///                     }
383///                 }
384///             }
385///         });
386///
387///         Self { jh, sem }
388///     }
389///
390///     async fn acquire(&self) {
391///         // This can return an error if the semaphore is closed, but we
392///         // never close it, so this error can never happen.
393///         let permit = self.sem.acquire().await.unwrap();
394///         // To avoid releasing the permit back to the semaphore, we use
395///         // the `SemaphorePermit::forget` method.
396///         permit.forget();
397///     }
398/// }
399///
400/// impl Drop for TokenBucket {
401///     fn drop(&mut self) {
402///         // Kill the background task so it stops taking up resources when we
403///         // don't need it anymore.
404///         self.jh.abort();
405///     }
406/// }
407///
408/// # #[tokio::main(flavor = "current_thread")]
409/// # async fn _hidden() {}
410/// # #[tokio::main(flavor = "current_thread", start_paused = true)]
411/// # async fn main() {
412/// let capacity = 5;
413/// let update_interval = Duration::from_secs_f32(1.0 / capacity as f32);
414/// let bucket = TokenBucket::new(update_interval, capacity);
415///
416/// for _ in 0..5 {
417///     bucket.acquire().await;
418///
419///     // do the operation
420/// }
421/// # }
422/// ```
423///
424/// [`PollSemaphore`]: https://docs.rs/tokio-util/latest/tokio_util/sync/struct.PollSemaphore.html
425/// [`Semaphore::acquire_owned`]: crate::sync::Semaphore::acquire_owned
426#[derive(Debug)]
427pub struct Semaphore {
428    /// The low level semaphore
429    ll_sem: ll::Semaphore,
430    #[cfg(all(tokio_unstable, feature = "tracing"))]
431    resource_span: tracing::Span,
432}
433
434/// A permit from the semaphore.
435///
436/// This type is created by the [`acquire`] method.
437///
438/// [`acquire`]: crate::sync::Semaphore::acquire()
439#[must_use]
440#[clippy::has_significant_drop]
441#[derive(Debug)]
442pub struct SemaphorePermit<'a> {
443    sem: &'a Semaphore,
444    permits: u32,
445}
446
447/// An owned permit from the semaphore.
448///
449/// This type is created by the [`acquire_owned`] method.
450///
451/// [`acquire_owned`]: crate::sync::Semaphore::acquire_owned()
452#[must_use]
453#[clippy::has_significant_drop]
454#[derive(Debug)]
455pub struct OwnedSemaphorePermit {
456    sem: Arc<Semaphore>,
457    permits: u32,
458}
459
460#[test]
461#[cfg(not(loom))]
462fn bounds() {
463    fn check_unpin<T: Unpin>() {}
464    // This has to take a value, since the async fn's return type is unnameable.
465    fn check_send_sync_val<T: Send + Sync>(_t: T) {}
466    fn check_send_sync<T: Send + Sync>() {}
467    check_unpin::<Semaphore>();
468    check_unpin::<SemaphorePermit<'_>>();
469    check_send_sync::<Semaphore>();
470
471    let semaphore = Semaphore::new(0);
472    check_send_sync_val(semaphore.acquire());
473}
474
475impl Semaphore {
476    /// The maximum number of permits which a semaphore can hold. It is `usize::MAX >> 3`.
477    ///
478    /// Exceeding this limit typically results in a panic.
479    pub const MAX_PERMITS: usize = super::batch_semaphore::Semaphore::MAX_PERMITS;
480
481    /// Creates a new semaphore with the initial number of permits.
482    ///
483    /// Panics if `permits` exceeds [`Semaphore::MAX_PERMITS`].
484    #[track_caller]
485    pub fn new(permits: usize) -> Self {
486        #[cfg(all(tokio_unstable, feature = "tracing"))]
487        let resource_span = {
488            let location = std::panic::Location::caller();
489
490            tracing::trace_span!(
491                parent: None,
492                "runtime.resource",
493                concrete_type = "Semaphore",
494                kind = "Sync",
495                loc.file = location.file(),
496                loc.line = location.line(),
497                loc.col = location.column(),
498                inherits_child_attrs = true,
499            )
500        };
501
502        #[cfg(all(tokio_unstable, feature = "tracing"))]
503        let ll_sem = resource_span.in_scope(|| ll::Semaphore::new(permits));
504
505        #[cfg(any(not(tokio_unstable), not(feature = "tracing")))]
506        let ll_sem = ll::Semaphore::new(permits);
507
508        Self {
509            ll_sem,
510            #[cfg(all(tokio_unstable, feature = "tracing"))]
511            resource_span,
512        }
513    }
514
515    /// Creates a new semaphore with the initial number of permits.
516    ///
517    /// When using the `tracing` [unstable feature], a `Semaphore` created with
518    /// `const_new` will not be instrumented. As such, it will not be visible
519    /// in [`tokio-console`]. Instead, [`Semaphore::new`] should be used to
520    /// create an instrumented object if that is needed.
521    ///
522    /// # Examples
523    ///
524    /// ```
525    /// use tokio::sync::Semaphore;
526    ///
527    /// static SEM: Semaphore = Semaphore::const_new(10);
528    /// ```
529    ///
530    /// [`tokio-console`]: https://github.com/tokio-rs/console
531    /// [unstable feature]: crate#unstable-features
532    #[cfg(not(all(loom, test)))]
533    pub const fn const_new(permits: usize) -> Self {
534        Self {
535            ll_sem: ll::Semaphore::const_new(permits),
536            #[cfg(all(tokio_unstable, feature = "tracing"))]
537            resource_span: tracing::Span::none(),
538        }
539    }
540
541    /// Creates a new closed semaphore with 0 permits.
542    pub(crate) fn new_closed() -> Self {
543        Self {
544            ll_sem: ll::Semaphore::new_closed(),
545            #[cfg(all(tokio_unstable, feature = "tracing"))]
546            resource_span: tracing::Span::none(),
547        }
548    }
549
550    /// Creates a new closed semaphore with 0 permits.
551    #[cfg(not(all(loom, test)))]
552    pub(crate) const fn const_new_closed() -> Self {
553        Self {
554            ll_sem: ll::Semaphore::const_new_closed(),
555            #[cfg(all(tokio_unstable, feature = "tracing"))]
556            resource_span: tracing::Span::none(),
557        }
558    }
559
560    /// Returns the current number of available permits.
561    pub fn available_permits(&self) -> usize {
562        self.ll_sem.available_permits()
563    }
564
565    /// Adds `n` new permits to the semaphore.
566    ///
567    /// The maximum number of permits is [`Semaphore::MAX_PERMITS`], and this function will panic if the limit is exceeded.
568    pub fn add_permits(&self, n: usize) {
569        self.ll_sem.release(n);
570    }
571
572    /// Decrease a semaphore's permits by a maximum of `n`.
573    ///
574    /// If there are insufficient permits and it's not possible to reduce by `n`,
575    /// return the number of permits that were actually reduced.
576    pub fn forget_permits(&self, n: usize) -> usize {
577        self.ll_sem.forget_permits(n)
578    }
579
580    /// Acquires a permit from the semaphore.
581    ///
582    /// If the semaphore has been closed, this returns an [`AcquireError`].
583    /// Otherwise, this returns a [`SemaphorePermit`] representing the
584    /// acquired permit.
585    ///
586    /// # Cancel safety
587    ///
588    /// This method uses a queue to fairly distribute permits in the order they
589    /// were requested. Cancelling a call to `acquire` makes you lose your place
590    /// in the queue.
591    ///
592    /// # Examples
593    ///
594    /// ```
595    /// use tokio::sync::Semaphore;
596    ///
597    /// # #[tokio::main(flavor = "current_thread")]
598    /// # async fn main() {
599    /// let semaphore = Semaphore::new(2);
600    ///
601    /// let permit_1 = semaphore.acquire().await.unwrap();
602    /// assert_eq!(semaphore.available_permits(), 1);
603    ///
604    /// let permit_2 = semaphore.acquire().await.unwrap();
605    /// assert_eq!(semaphore.available_permits(), 0);
606    ///
607    /// drop(permit_1);
608    /// assert_eq!(semaphore.available_permits(), 1);
609    /// # }
610    /// ```
611    ///
612    /// [`AcquireError`]: crate::sync::AcquireError
613    /// [`SemaphorePermit`]: crate::sync::SemaphorePermit
614    pub async fn acquire(&self) -> Result<SemaphorePermit<'_>, AcquireError> {
615        #[cfg(all(tokio_unstable, feature = "tracing"))]
616        let inner = trace::async_op(
617            || self.ll_sem.acquire(1),
618            self.resource_span.clone(),
619            "Semaphore::acquire",
620            "poll",
621            true,
622        );
623        #[cfg(not(all(tokio_unstable, feature = "tracing")))]
624        let inner = self.ll_sem.acquire(1);
625
626        inner.await?;
627        Ok(SemaphorePermit {
628            sem: self,
629            permits: 1,
630        })
631    }
632
633    /// Acquires `n` permits from the semaphore.
634    ///
635    /// If the semaphore has been closed, this returns an [`AcquireError`].
636    /// Otherwise, this returns a [`SemaphorePermit`] representing the
637    /// acquired permits.
638    ///
639    /// # Cancel safety
640    ///
641    /// This method uses a queue to fairly distribute permits in the order they
642    /// were requested. Cancelling a call to `acquire_many` makes you lose your
643    /// place in the queue.
644    ///
645    /// # Examples
646    ///
647    /// ```
648    /// use tokio::sync::Semaphore;
649    ///
650    /// # #[tokio::main(flavor = "current_thread")]
651    /// # async fn main() {
652    /// let semaphore = Semaphore::new(5);
653    ///
654    /// let permit = semaphore.acquire_many(3).await.unwrap();
655    /// assert_eq!(semaphore.available_permits(), 2);
656    /// # }
657    /// ```
658    ///
659    /// [`AcquireError`]: crate::sync::AcquireError
660    /// [`SemaphorePermit`]: crate::sync::SemaphorePermit
661    pub async fn acquire_many(&self, n: u32) -> Result<SemaphorePermit<'_>, AcquireError> {
662        #[cfg(all(tokio_unstable, feature = "tracing"))]
663        trace::async_op(
664            || self.ll_sem.acquire(n as usize),
665            self.resource_span.clone(),
666            "Semaphore::acquire_many",
667            "poll",
668            true,
669        )
670        .await?;
671
672        #[cfg(not(all(tokio_unstable, feature = "tracing")))]
673        self.ll_sem.acquire(n as usize).await?;
674
675        Ok(SemaphorePermit {
676            sem: self,
677            permits: n,
678        })
679    }
680
681    /// Tries to acquire a permit from the semaphore.
682    ///
683    /// If the semaphore has been closed, this returns a [`TryAcquireError::Closed`]
684    /// and a [`TryAcquireError::NoPermits`] if there are no permits left. Otherwise,
685    /// this returns a [`SemaphorePermit`] representing the acquired permits.
686    ///
687    /// # Examples
688    ///
689    /// ```
690    /// use tokio::sync::{Semaphore, TryAcquireError};
691    ///
692    /// # fn main() {
693    /// let semaphore = Semaphore::new(2);
694    ///
695    /// let permit_1 = semaphore.try_acquire().unwrap();
696    /// assert_eq!(semaphore.available_permits(), 1);
697    ///
698    /// let permit_2 = semaphore.try_acquire().unwrap();
699    /// assert_eq!(semaphore.available_permits(), 0);
700    ///
701    /// let permit_3 = semaphore.try_acquire();
702    /// assert_eq!(permit_3.err(), Some(TryAcquireError::NoPermits));
703    /// # }
704    /// ```
705    ///
706    /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed
707    /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits
708    /// [`SemaphorePermit`]: crate::sync::SemaphorePermit
709    pub fn try_acquire(&self) -> Result<SemaphorePermit<'_>, TryAcquireError> {
710        match self.ll_sem.try_acquire(1) {
711            Ok(()) => Ok(SemaphorePermit {
712                sem: self,
713                permits: 1,
714            }),
715            Err(e) => Err(e),
716        }
717    }
718
719    /// Tries to acquire `n` permits from the semaphore.
720    ///
721    /// If the semaphore has been closed, this returns a [`TryAcquireError::Closed`]
722    /// and a [`TryAcquireError::NoPermits`] if there are not enough permits left.
723    /// Otherwise, this returns a [`SemaphorePermit`] representing the acquired permits.
724    ///
725    /// # Examples
726    ///
727    /// ```
728    /// use tokio::sync::{Semaphore, TryAcquireError};
729    ///
730    /// # fn main() {
731    /// let semaphore = Semaphore::new(4);
732    ///
733    /// let permit_1 = semaphore.try_acquire_many(3).unwrap();
734    /// assert_eq!(semaphore.available_permits(), 1);
735    ///
736    /// let permit_2 = semaphore.try_acquire_many(2);
737    /// assert_eq!(permit_2.err(), Some(TryAcquireError::NoPermits));
738    /// # }
739    /// ```
740    ///
741    /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed
742    /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits
743    /// [`SemaphorePermit`]: crate::sync::SemaphorePermit
744    pub fn try_acquire_many(&self, n: u32) -> Result<SemaphorePermit<'_>, TryAcquireError> {
745        match self.ll_sem.try_acquire(n as usize) {
746            Ok(()) => Ok(SemaphorePermit {
747                sem: self,
748                permits: n,
749            }),
750            Err(e) => Err(e),
751        }
752    }
753
754    /// Acquires a permit from the semaphore.
755    ///
756    /// The semaphore must be wrapped in an [`Arc`] to call this method.
757    /// If the semaphore has been closed, this returns an [`AcquireError`].
758    /// Otherwise, this returns a [`OwnedSemaphorePermit`] representing the
759    /// acquired permit.
760    ///
761    /// # Cancel safety
762    ///
763    /// This method uses a queue to fairly distribute permits in the order they
764    /// were requested. Cancelling a call to `acquire_owned` makes you lose your
765    /// place in the queue.
766    ///
767    /// # Examples
768    ///
769    /// ```
770    /// use std::sync::Arc;
771    /// use tokio::sync::Semaphore;
772    ///
773    /// # #[tokio::main(flavor = "current_thread")]
774    /// # async fn main() {
775    /// let semaphore = Arc::new(Semaphore::new(3));
776    /// let mut join_handles = Vec::new();
777    ///
778    /// for _ in 0..5 {
779    ///     let permit = semaphore.clone().acquire_owned().await.unwrap();
780    ///     join_handles.push(tokio::spawn(async move {
781    ///         // perform task...
782    ///         // explicitly own `permit` in the task
783    ///         drop(permit);
784    ///     }));
785    /// }
786    ///
787    /// for handle in join_handles {
788    ///     handle.await.unwrap();
789    /// }
790    /// # }
791    /// ```
792    ///
793    /// [`Arc`]: std::sync::Arc
794    /// [`AcquireError`]: crate::sync::AcquireError
795    /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit
796    pub async fn acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, AcquireError> {
797        #[cfg(all(tokio_unstable, feature = "tracing"))]
798        let inner = trace::async_op(
799            || self.ll_sem.acquire(1),
800            self.resource_span.clone(),
801            "Semaphore::acquire_owned",
802            "poll",
803            true,
804        );
805        #[cfg(not(all(tokio_unstable, feature = "tracing")))]
806        let inner = self.ll_sem.acquire(1);
807
808        inner.await?;
809        Ok(OwnedSemaphorePermit {
810            sem: self,
811            permits: 1,
812        })
813    }
814
815    /// Acquires `n` permits from the semaphore.
816    ///
817    /// The semaphore must be wrapped in an [`Arc`] to call this method.
818    /// If the semaphore has been closed, this returns an [`AcquireError`].
819    /// Otherwise, this returns a [`OwnedSemaphorePermit`] representing the
820    /// acquired permit.
821    ///
822    /// # Cancel safety
823    ///
824    /// This method uses a queue to fairly distribute permits in the order they
825    /// were requested. Cancelling a call to `acquire_many_owned` makes you lose
826    /// your place in the queue.
827    ///
828    /// # Examples
829    ///
830    /// ```
831    /// use std::sync::Arc;
832    /// use tokio::sync::Semaphore;
833    ///
834    /// # #[tokio::main(flavor = "current_thread")]
835    /// # async fn main() {
836    /// let semaphore = Arc::new(Semaphore::new(10));
837    /// let mut join_handles = Vec::new();
838    ///
839    /// for _ in 0..5 {
840    ///     let permit = semaphore.clone().acquire_many_owned(2).await.unwrap();
841    ///     join_handles.push(tokio::spawn(async move {
842    ///         // perform task...
843    ///         // explicitly own `permit` in the task
844    ///         drop(permit);
845    ///     }));
846    /// }
847    ///
848    /// for handle in join_handles {
849    ///     handle.await.unwrap();
850    /// }
851    /// # }
852    /// ```
853    ///
854    /// [`Arc`]: std::sync::Arc
855    /// [`AcquireError`]: crate::sync::AcquireError
856    /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit
857    pub async fn acquire_many_owned(
858        self: Arc<Self>,
859        n: u32,
860    ) -> Result<OwnedSemaphorePermit, AcquireError> {
861        #[cfg(all(tokio_unstable, feature = "tracing"))]
862        let inner = trace::async_op(
863            || self.ll_sem.acquire(n as usize),
864            self.resource_span.clone(),
865            "Semaphore::acquire_many_owned",
866            "poll",
867            true,
868        );
869        #[cfg(not(all(tokio_unstable, feature = "tracing")))]
870        let inner = self.ll_sem.acquire(n as usize);
871
872        inner.await?;
873        Ok(OwnedSemaphorePermit {
874            sem: self,
875            permits: n,
876        })
877    }
878
879    /// Tries to acquire a permit from the semaphore.
880    ///
881    /// The semaphore must be wrapped in an [`Arc`] to call this method. If
882    /// the semaphore has been closed, this returns a [`TryAcquireError::Closed`]
883    /// and a [`TryAcquireError::NoPermits`] if there are no permits left.
884    /// Otherwise, this returns a [`OwnedSemaphorePermit`] representing the
885    /// acquired permit.
886    ///
887    /// # Examples
888    ///
889    /// ```
890    /// use std::sync::Arc;
891    /// use tokio::sync::{Semaphore, TryAcquireError};
892    ///
893    /// # fn main() {
894    /// let semaphore = Arc::new(Semaphore::new(2));
895    ///
896    /// let permit_1 = Arc::clone(&semaphore).try_acquire_owned().unwrap();
897    /// assert_eq!(semaphore.available_permits(), 1);
898    ///
899    /// let permit_2 = Arc::clone(&semaphore).try_acquire_owned().unwrap();
900    /// assert_eq!(semaphore.available_permits(), 0);
901    ///
902    /// let permit_3 = semaphore.try_acquire_owned();
903    /// assert_eq!(permit_3.err(), Some(TryAcquireError::NoPermits));
904    /// # }
905    /// ```
906    ///
907    /// [`Arc`]: std::sync::Arc
908    /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed
909    /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits
910    /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit
911    pub fn try_acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, TryAcquireError> {
912        match self.ll_sem.try_acquire(1) {
913            Ok(()) => Ok(OwnedSemaphorePermit {
914                sem: self,
915                permits: 1,
916            }),
917            Err(e) => Err(e),
918        }
919    }
920
921    /// Tries to acquire `n` permits from the semaphore.
922    ///
923    /// The semaphore must be wrapped in an [`Arc`] to call this method. If
924    /// the semaphore has been closed, this returns a [`TryAcquireError::Closed`]
925    /// and a [`TryAcquireError::NoPermits`] if there are no permits left.
926    /// Otherwise, this returns a [`OwnedSemaphorePermit`] representing the
927    /// acquired permit.
928    ///
929    /// # Examples
930    ///
931    /// ```
932    /// use std::sync::Arc;
933    /// use tokio::sync::{Semaphore, TryAcquireError};
934    ///
935    /// # fn main() {
936    /// let semaphore = Arc::new(Semaphore::new(4));
937    ///
938    /// let permit_1 = Arc::clone(&semaphore).try_acquire_many_owned(3).unwrap();
939    /// assert_eq!(semaphore.available_permits(), 1);
940    ///
941    /// let permit_2 = semaphore.try_acquire_many_owned(2);
942    /// assert_eq!(permit_2.err(), Some(TryAcquireError::NoPermits));
943    /// # }
944    /// ```
945    ///
946    /// [`Arc`]: std::sync::Arc
947    /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed
948    /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits
949    /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit
950    pub fn try_acquire_many_owned(
951        self: Arc<Self>,
952        n: u32,
953    ) -> Result<OwnedSemaphorePermit, TryAcquireError> {
954        match self.ll_sem.try_acquire(n as usize) {
955            Ok(()) => Ok(OwnedSemaphorePermit {
956                sem: self,
957                permits: n,
958            }),
959            Err(e) => Err(e),
960        }
961    }
962
963    /// Closes the semaphore.
964    ///
965    /// This prevents the semaphore from issuing new permits and notifies all pending waiters.
966    ///
967    /// # Examples
968    ///
969    /// ```
970    /// use tokio::sync::Semaphore;
971    /// use std::sync::Arc;
972    /// use tokio::sync::TryAcquireError;
973    ///
974    /// # #[tokio::main(flavor = "current_thread")]
975    /// # async fn main() {
976    /// let semaphore = Arc::new(Semaphore::new(1));
977    /// let semaphore2 = semaphore.clone();
978    ///
979    /// tokio::spawn(async move {
980    ///     let permit = semaphore.acquire_many(2).await;
981    ///     assert!(permit.is_err());
982    ///     println!("waiter received error");
983    /// });
984    ///
985    /// println!("closing semaphore");
986    /// semaphore2.close();
987    ///
988    /// // Cannot obtain more permits
989    /// assert_eq!(semaphore2.try_acquire().err(), Some(TryAcquireError::Closed))
990    /// # }
991    /// ```
992    pub fn close(&self) {
993        self.ll_sem.close();
994    }
995
996    /// Returns true if the semaphore is closed
997    pub fn is_closed(&self) -> bool {
998        self.ll_sem.is_closed()
999    }
1000}
1001
1002impl<'a> SemaphorePermit<'a> {
1003    /// Forgets the permit **without** releasing it back to the semaphore.
1004    /// This can be used to reduce the amount of permits available from a
1005    /// semaphore.
1006    ///
1007    /// # Examples
1008    ///
1009    /// ```
1010    /// use std::sync::Arc;
1011    /// use tokio::sync::Semaphore;
1012    ///
1013    /// let sem = Arc::new(Semaphore::new(10));
1014    /// {
1015    ///     let permit = sem.try_acquire_many(5).unwrap();
1016    ///     assert_eq!(sem.available_permits(), 5);
1017    ///     permit.forget();
1018    /// }
1019    ///
1020    /// // Since we forgot the permit, available permits won't go back to its initial value
1021    /// // even after the permit is dropped.
1022    /// assert_eq!(sem.available_permits(), 5);
1023    /// ```
1024    pub fn forget(mut self) {
1025        self.permits = 0;
1026    }
1027
1028    /// Merge two [`SemaphorePermit`] instances together, consuming `other`
1029    /// without releasing the permits it holds.
1030    ///
1031    /// Permits held by both `self` and `other` are released when `self` drops.
1032    ///
1033    /// # Panics
1034    ///
1035    /// This function panics if permits from different [`Semaphore`] instances
1036    /// are merged.
1037    ///
1038    /// # Examples
1039    ///
1040    /// ```
1041    /// use std::sync::Arc;
1042    /// use tokio::sync::Semaphore;
1043    ///
1044    /// let sem = Arc::new(Semaphore::new(10));
1045    /// let mut permit = sem.try_acquire().unwrap();
1046    ///
1047    /// for _ in 0..9 {
1048    ///     let _permit = sem.try_acquire().unwrap();
1049    ///     // Merge individual permits into a single one.
1050    ///     permit.merge(_permit)
1051    /// }
1052    ///
1053    /// assert_eq!(sem.available_permits(), 0);
1054    ///
1055    /// // Release all permits in a single batch.
1056    /// drop(permit);
1057    ///
1058    /// assert_eq!(sem.available_permits(), 10);
1059    /// ```
1060    #[track_caller]
1061    pub fn merge(&mut self, mut other: Self) {
1062        assert!(
1063            std::ptr::eq(self.sem, other.sem),
1064            "merging permits from different semaphore instances"
1065        );
1066        self.permits += other.permits;
1067        other.permits = 0;
1068    }
1069
1070    /// Splits `n` permits from `self` and returns a new [`SemaphorePermit`] instance that holds `n` permits.
1071    ///
1072    /// If there are insufficient permits and it's not possible to reduce by `n`, returns `None`.
1073    ///
1074    /// # Examples
1075    ///
1076    /// ```
1077    /// use std::sync::Arc;
1078    /// use tokio::sync::Semaphore;
1079    ///
1080    /// let sem = Arc::new(Semaphore::new(3));
1081    ///
1082    /// let mut p1 = sem.try_acquire_many(3).unwrap();
1083    /// let p2 = p1.split(1).unwrap();
1084    ///
1085    /// assert_eq!(p1.num_permits(), 2);
1086    /// assert_eq!(p2.num_permits(), 1);
1087    /// ```
1088    pub fn split(&mut self, n: usize) -> Option<Self> {
1089        let n = u32::try_from(n).ok()?;
1090
1091        if n > self.permits {
1092            return None;
1093        }
1094
1095        self.permits -= n;
1096
1097        Some(Self {
1098            sem: self.sem,
1099            permits: n,
1100        })
1101    }
1102
1103    /// Returns the number of permits held by `self`.
1104    pub fn num_permits(&self) -> usize {
1105        self.permits as usize
1106    }
1107}
1108
1109impl OwnedSemaphorePermit {
1110    /// Forgets the permit **without** releasing it back to the semaphore.
1111    /// This can be used to reduce the amount of permits available from a
1112    /// semaphore.
1113    ///
1114    /// # Examples
1115    ///
1116    /// ```
1117    /// use std::sync::Arc;
1118    /// use tokio::sync::Semaphore;
1119    ///
1120    /// let sem = Arc::new(Semaphore::new(10));
1121    /// {
1122    ///     let permit = sem.clone().try_acquire_many_owned(5).unwrap();
1123    ///     assert_eq!(sem.available_permits(), 5);
1124    ///     permit.forget();
1125    /// }
1126    ///
1127    /// // Since we forgot the permit, available permits won't go back to its initial value
1128    /// // even after the permit is dropped.
1129    /// assert_eq!(sem.available_permits(), 5);
1130    /// ```
1131    pub fn forget(mut self) {
1132        self.permits = 0;
1133    }
1134
1135    /// Merge two [`OwnedSemaphorePermit`] instances together, consuming `other`
1136    /// without releasing the permits it holds.
1137    ///
1138    /// Permits held by both `self` and `other` are released when `self` drops.
1139    ///
1140    /// # Panics
1141    ///
1142    /// This function panics if permits from different [`Semaphore`] instances
1143    /// are merged.
1144    ///
1145    /// # Examples
1146    ///
1147    /// ```
1148    /// use std::sync::Arc;
1149    /// use tokio::sync::Semaphore;
1150    ///
1151    /// let sem = Arc::new(Semaphore::new(10));
1152    /// let mut permit = sem.clone().try_acquire_owned().unwrap();
1153    ///
1154    /// for _ in 0..9 {
1155    ///     let _permit = sem.clone().try_acquire_owned().unwrap();
1156    ///     // Merge individual permits into a single one.
1157    ///     permit.merge(_permit)
1158    /// }
1159    ///
1160    /// assert_eq!(sem.available_permits(), 0);
1161    ///
1162    /// // Release all permits in a single batch.
1163    /// drop(permit);
1164    ///
1165    /// assert_eq!(sem.available_permits(), 10);
1166    /// ```
1167    #[track_caller]
1168    pub fn merge(&mut self, mut other: Self) {
1169        assert!(
1170            Arc::ptr_eq(&self.sem, &other.sem),
1171            "merging permits from different semaphore instances"
1172        );
1173        self.permits += other.permits;
1174        other.permits = 0;
1175    }
1176
1177    /// Splits `n` permits from `self` and returns a new [`OwnedSemaphorePermit`] instance that holds `n` permits.
1178    ///
1179    /// If there are insufficient permits and it's not possible to reduce by `n`, returns `None`.
1180    ///
1181    /// # Note
1182    ///
1183    /// It will clone the owned `Arc<Semaphore>` to construct the new instance.
1184    ///
1185    /// # Examples
1186    ///
1187    /// ```
1188    /// use std::sync::Arc;
1189    /// use tokio::sync::Semaphore;
1190    ///
1191    /// let sem = Arc::new(Semaphore::new(3));
1192    ///
1193    /// let mut p1 = sem.try_acquire_many_owned(3).unwrap();
1194    /// let p2 = p1.split(1).unwrap();
1195    ///
1196    /// assert_eq!(p1.num_permits(), 2);
1197    /// assert_eq!(p2.num_permits(), 1);
1198    /// ```
1199    pub fn split(&mut self, n: usize) -> Option<Self> {
1200        let n = u32::try_from(n).ok()?;
1201
1202        if n > self.permits {
1203            return None;
1204        }
1205
1206        self.permits -= n;
1207
1208        Some(Self {
1209            sem: self.sem.clone(),
1210            permits: n,
1211        })
1212    }
1213
1214    /// Returns the [`Semaphore`] from which this permit was acquired.
1215    pub fn semaphore(&self) -> &Arc<Semaphore> {
1216        &self.sem
1217    }
1218
1219    /// Returns the number of permits held by `self`.
1220    pub fn num_permits(&self) -> usize {
1221        self.permits as usize
1222    }
1223}
1224
1225impl Drop for SemaphorePermit<'_> {
1226    fn drop(&mut self) {
1227        self.sem.add_permits(self.permits as usize);
1228    }
1229}
1230
1231impl Drop for OwnedSemaphorePermit {
1232    fn drop(&mut self) {
1233        self.sem.add_permits(self.permits as usize);
1234    }
1235}