simple_semaphore/
lib.rs

1use std::sync::{
2    atomic::{AtomicUsize, Ordering},
3    Arc, Condvar, Mutex,
4};
5
6/// A Semaphore maintains the number of permits it is still allowed to give.
7///
8/// * When `acquire()` is called and the semaphore still has permits to give, it will return a `Permit`. If there are no permits that can be given, it will wait for one permit to be given back from a thread so that it can return a new `Permit`.
9/// * When `try_acquire()` is called and the semaphore still has permits to give, it will return `Some(Permit)`. If there are no permits that can be given, it will return `None`.
10#[derive(Debug)]
11pub struct Semaphore {
12    permits: Arc<AtomicUsize>,
13    condvar: Condvar,
14    mutex: Mutex<()>,
15}
16
17impl Semaphore {
18    /// Returns a new `Arc<Semaphore>` with the limit of permits chosen by you.
19    pub fn new(permits: usize) -> Arc<Self> {
20        Arc::new(Semaphore {
21            permits: Arc::new(AtomicUsize::new(permits)),
22            condvar: Condvar::new(),
23            mutex: Mutex::new(()),
24        })
25    }
26
27    /// Returns a new `Arc<Semaphore>` with the limit of permits set to the number of CPU cores the computer has.
28    pub fn default() -> Arc<Self> {
29        Arc::new(Semaphore {
30            permits: Arc::new(AtomicUsize::new(num_cpus::get())),
31            condvar: Condvar::new(),
32            mutex: Mutex::new(()),
33        })
34    }
35
36    /// Returns the number of available permits
37    pub fn available_permits(self: &Arc<Self>) -> usize {
38        self.permits.load(Ordering::Relaxed)
39    }
40
41    /// Tries to get a `Permit`. If no more permits can be given, it will wait for one permit to be given back from a thread so that it can return a new `Permit`.
42    #[allow(unused_must_use)]
43    pub fn acquire(self: &Arc<Self>) -> Permit {
44        loop {
45            if self.permits.load(Ordering::Acquire) != 0 {
46                self.permits.fetch_sub(1, Ordering::AcqRel);
47                return Permit {
48                    semaphore: Arc::clone(self),
49                };
50            }
51            let guard = self.mutex.lock().unwrap();
52            self.condvar.wait(guard).unwrap();
53        }
54    }
55
56    /// Tries to get a `Option<Permit>`. If no more permits can be given, it will return `None`.
57    pub fn try_acquire(self: &Arc<Self>) -> Option<Permit> {
58        if self.permits.load(Ordering::Acquire) != 0 {
59            self.permits.fetch_sub(1, Ordering::Release);
60            return Some(Permit {
61                semaphore: Arc::clone(self),
62            });
63        }
64        None
65    }
66
67    /// Releases a permit. This is what `drop()` on `Permit` calls, ideally use `drop(permit)`.
68    pub fn release(&self) {
69        self.permits.fetch_add(1, Ordering::Release);
70        self.condvar.notify_one();
71    }
72}
73
74/// A permit that holds the Semaphore, so that `drop(permit)` can be called.
75#[derive(Debug)]
76pub struct Permit {
77    semaphore: Arc<Semaphore>,
78}
79
80impl Drop for Permit {
81    /// Releases the permit.
82    fn drop(&mut self) {
83        self.semaphore.release();
84    }
85}
simple_semaphore/lib.rs

simple_semaphore/
lib.rs
