1 use super::batch_semaphore as ll; // low level implementation
2 use std::sync::Arc;
3
4 /// Counting semaphore performing asynchronous permit aquisition.
5 ///
6 /// A semaphore maintains a set of permits. Permits are used to synchronize
7 /// access to a shared resource. A semaphore differs from a mutex in that it
8 /// can allow more than one concurrent caller to access the shared resource at a
9 /// time.
10 ///
11 /// When `acquire` is called and the semaphore has remaining permits, the
12 /// function immediately returns a permit. However, if no remaining permits are
13 /// available, `acquire` (asynchronously) waits until an outstanding permit is
14 /// dropped. At this point, the freed permit is assigned to the caller.
15 #[derive(Debug)]
16 pub struct Semaphore {
17 /// The low level semaphore
18 ll_sem: ll::Semaphore,
19 }
20
21 /// A permit from the semaphore.
22 ///
23 /// This type is created by the [`acquire`] method.
24 ///
25 /// [`acquire`]: crate::sync::Semaphore::acquire()
26 #[must_use]
27 #[derive(Debug)]
28 pub struct SemaphorePermit<'a> {
29 sem: &'a Semaphore,
30 permits: u16,
31 }
32
33 /// An owned permit from the semaphore.
34 ///
35 /// This type is created by the [`acquire_owned`] method.
36 ///
37 /// [`acquire_owned`]: crate::sync::Semaphore::acquire_owned()
38 #[must_use]
39 #[derive(Debug)]
40 pub struct OwnedSemaphorePermit {
41 sem: Arc<Semaphore>,
42 permits: u16,
43 }
44
45 /// Error returned from the [`Semaphore::try_acquire`] function.
46 ///
47 /// A `try_acquire` operation can only fail if the semaphore has no available
48 /// permits.
49 ///
50 /// [`Semaphore::try_acquire`]: Semaphore::try_acquire
51 #[derive(Debug)]
52 pub struct TryAcquireError(());
53
54 #[test]
55 #[cfg(not(loom))]
bounds()56 fn bounds() {
57 fn check_unpin<T: Unpin>() {}
58 // This has to take a value, since the async fn's return type is unnameable.
59 fn check_send_sync_val<T: Send + Sync>(_t: T) {}
60 fn check_send_sync<T: Send + Sync>() {}
61 check_unpin::<Semaphore>();
62 check_unpin::<SemaphorePermit<'_>>();
63 check_send_sync::<Semaphore>();
64
65 let semaphore = Semaphore::new(0);
66 check_send_sync_val(semaphore.acquire());
67 }
68
69 impl Semaphore {
70 /// Creates a new semaphore with the initial number of permits.
new(permits: usize) -> Self71 pub fn new(permits: usize) -> Self {
72 Self {
73 ll_sem: ll::Semaphore::new(permits),
74 }
75 }
76
77 /// Returns the current number of available permits.
available_permits(&self) -> usize78 pub fn available_permits(&self) -> usize {
79 self.ll_sem.available_permits()
80 }
81
82 /// Adds `n` new permits to the semaphore.
83 ///
84 /// The maximum number of permits is `usize::MAX >> 3`, and this function will panic if the limit is exceeded.
add_permits(&self, n: usize)85 pub fn add_permits(&self, n: usize) {
86 self.ll_sem.release(n);
87 }
88
89 /// Acquires permit from the semaphore.
acquire(&self) -> SemaphorePermit<'_>90 pub async fn acquire(&self) -> SemaphorePermit<'_> {
91 self.ll_sem.acquire(1).await.unwrap();
92 SemaphorePermit {
93 sem: &self,
94 permits: 1,
95 }
96 }
97
98 /// Tries to acquire a permit from the semaphore.
try_acquire(&self) -> Result<SemaphorePermit<'_>, TryAcquireError>99 pub fn try_acquire(&self) -> Result<SemaphorePermit<'_>, TryAcquireError> {
100 match self.ll_sem.try_acquire(1) {
101 Ok(_) => Ok(SemaphorePermit {
102 sem: self,
103 permits: 1,
104 }),
105 Err(_) => Err(TryAcquireError(())),
106 }
107 }
108
109 /// Acquires permit from the semaphore.
110 ///
111 /// The semaphore must be wrapped in an [`Arc`] to call this method.
112 ///
113 /// [`Arc`]: std::sync::Arc
acquire_owned(self: Arc<Self>) -> OwnedSemaphorePermit114 pub async fn acquire_owned(self: Arc<Self>) -> OwnedSemaphorePermit {
115 self.ll_sem.acquire(1).await.unwrap();
116 OwnedSemaphorePermit {
117 sem: self.clone(),
118 permits: 1,
119 }
120 }
121
122 /// Tries to acquire a permit from the semaphore.
123 ///
124 /// The semaphore must be wrapped in an [`Arc`] to call this method.
125 ///
126 /// [`Arc`]: std::sync::Arc
try_acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, TryAcquireError>127 pub fn try_acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, TryAcquireError> {
128 match self.ll_sem.try_acquire(1) {
129 Ok(_) => Ok(OwnedSemaphorePermit {
130 sem: self.clone(),
131 permits: 1,
132 }),
133 Err(_) => Err(TryAcquireError(())),
134 }
135 }
136 }
137
138 impl<'a> SemaphorePermit<'a> {
139 /// Forgets the permit **without** releasing it back to the semaphore.
140 /// This can be used to reduce the amount of permits available from a
141 /// semaphore.
forget(mut self)142 pub fn forget(mut self) {
143 self.permits = 0;
144 }
145 }
146
147 impl OwnedSemaphorePermit {
148 /// Forgets the permit **without** releasing it back to the semaphore.
149 /// This can be used to reduce the amount of permits available from a
150 /// semaphore.
forget(mut self)151 pub fn forget(mut self) {
152 self.permits = 0;
153 }
154 }
155
156 impl<'a> Drop for SemaphorePermit<'_> {
drop(&mut self)157 fn drop(&mut self) {
158 self.sem.add_permits(self.permits as usize);
159 }
160 }
161
162 impl Drop for OwnedSemaphorePermit {
drop(&mut self)163 fn drop(&mut self) {
164 self.sem.add_permits(self.permits as usize);
165 }
166 }
167