os/kernel/scheduler/
wait_queue.rs

1//! 等待队列模块
2//!
3//! 定义了等待队列结构体及其相关操作
4use crate::kernel::task::SharedTask;
5use crate::kernel::{TaskQueue, sleep_task_with_block, wake_up_with_block, yield_task};
6use crate::sync::RawSpinLock;
7use alloc::vec::Vec;
8
9/// 等待队列结构体
10/// 用于管理等待某些事件的任务列表
11/// 提供将任务加入等待队列、从队列中唤醒任务等功能
12/// 内部使用任务队列和自旋锁来保证线程安全
13/// 使用示例：
14/// ```ignore
15/// let mut wait_queue = WaitQueue::new();
16/// wait_queue.sleep(task); // 将任务加入等待队列并阻塞
17/// wait_queue.wake_up(&task); // 唤醒指定任务
18/// wait_queue.wake_up_one(); // 唤醒队首任务
19/// wait_queue.wake_up_all(); // 唤醒所有任务
20/// ```
21#[derive(Debug)]
22pub struct WaitQueue {
23    tasks: TaskQueue,
24    lock: RawSpinLock,
25}
26
27impl WaitQueue {
28    /// 创建一个新的等待队列实例
29    pub fn new() -> Self {
30        WaitQueue {
31            tasks: TaskQueue::new(),
32            lock: RawSpinLock::new(),
33        }
34    }
35
36    /// 把任务加入等待队列，并调用 sleep_task（不会导致调度）
37    pub fn sleep(&mut self, task: SharedTask) {
38        let _g = self.lock.lock();
39        self.tasks.add_task(task.clone());
40        sleep_task_with_block(task, true);
41    }
42
43    /// 从等待队列中移除指定任务并在锁释放后唤醒
44    pub fn wake_up(&mut self, task: &SharedTask) {
45        let should_wake = {
46            let _g = self.lock.lock();
47            if self.tasks.contains(task) {
48                self.tasks.remove_task(task);
49                true
50            } else {
51                false
52            }
53        };
54        if should_wake {
55            wake_up_with_block(task.clone());
56        }
57    }
58
59    /// 唤醒队首一个任务：在临界区内 pop，然后在临界区外唤醒
60    pub fn wake_up_one(&mut self) {
61        let maybe_task = {
62            let _g = self.lock.lock();
63            self.tasks.pop_task()
64        };
65        if let Some(t) = maybe_task {
66            wake_up_with_block(t);
67        }
68    }
69
70    /// 唤醒队列中所有任务：一次性把要唤醒的任务收集出来，释放锁后逐个唤醒
71    pub fn wake_up_all(&mut self) {
72        let mut to_wake: Vec<SharedTask> = Vec::new();
73        {
74            let _g = self.lock.lock();
75            while let Some(t) = self.tasks.pop_task() {
76                to_wake.push(t);
77            }
78        }
79        for t in to_wake {
80            wake_up_with_block(t);
81        }
82    }
83
84    /// 将任务加入等待队列（不阻塞）
85    pub fn add_task(&mut self, task: SharedTask) {
86        let _g = self.lock.lock();
87        self.tasks.add_task(task);
88    }
89
90    /// 从等待队列中移除指定任务（不唤醒）
91    pub fn remove_task(&mut self, task: &SharedTask) {
92        let _g = self.lock.lock();
93        self.tasks.remove_task(task);
94    }
95
96    /// 检查任务是否在队列中
97    pub fn contains(&self, task: &SharedTask) -> bool {
98        let _g = self.lock.lock();
99        self.tasks.contains(task)
100    }
101
102    /// 检查等待队列是否为空
103    pub fn is_empty(&self) -> bool {
104        let _g = self.lock.lock();
105        self.tasks.is_empty()
106    }
107
108    /// 原子地检查条件并睡眠（用于防止 lost wakeup）
109    /// check_fn 在持有锁时被调用，如果返回 true 则不睡眠
110    pub fn sleep_if<F>(&mut self, task: SharedTask, check_fn: F) -> bool
111    where
112        F: FnOnce() -> bool,
113    {
114        let _g = self.lock.lock();
115        if check_fn() {
116            return false; // 条件满足，不睡眠
117        }
118        self.tasks.add_task(task.clone());
119        sleep_task_with_block(task, true);
120        true // 已睡眠
121    }
122}
123
124// SAFETY:
125// WaitQueue 内部使用 RawSpinLock 来保护任务队列的并发访问
126// 因此 WaitQueue 本身是线程安全的，可以在多线程环境中共享
127unsafe impl Send for WaitQueue {}
128unsafe impl Sync for WaitQueue {}
os/kernel/scheduler/wait_queue.rs

os/kernel/scheduler/
wait_queue.rs
