os/sync/
semaphore.rs

1use crate::sync::UPSafeCell;
2use crate::task::{TaskControlBlock, block_current_and_run_next, current_task, wakeup_task};
3use alloc::{collections::VecDeque, sync::Arc};
4
5pub struct Semaphore {
6    pub inner: UPSafeCell<SemaphoreInner>,
7}
8
9pub struct SemaphoreInner {
10    pub count: isize,
11    pub wait_queue: VecDeque<Arc<TaskControlBlock>>,
12}
13
14impl Semaphore {
15    pub fn new(res_count: usize) -> Self {
16        Self {
17            inner: unsafe {
18                UPSafeCell::new(SemaphoreInner {
19                    count: res_count as isize,
20                    wait_queue: VecDeque::new(),
21                })
22            },
23        }
24    }
25
26    pub fn up(&self) {
27        let mut inner = self.inner.exclusive_access();
28        inner.count += 1;
29        if inner.count <= 0 {
30            if let Some(task) = inner.wait_queue.pop_front() {
31                wakeup_task(task);
32            }
33        }
34    }
35
36    pub fn down(&self) {
37        let mut inner = self.inner.exclusive_access();
38        inner.count -= 1;
39        if inner.count < 0 {
40            inner.wait_queue.push_back(current_task().unwrap());
41            drop(inner);
42            block_current_and_run_next();
43        }
44    }
45}