Skip to main content

sparreal_kernel/os/async/
task.rs

1//! 异步任务相关结构体
2//!
3//! 定义了异步执行器使用的任务、句柄和Waker实现。
4
5use core::future::Future;
6use core::pin::Pin;
7use core::task::{Context, Poll, Waker};
8
9use alloc::boxed::Box;
10use alloc::sync::Arc;
11
12use crate::os::sync::spinlock::IrqSpinlock;
13use crate::os::time;
14
15/// 任务唯一标识符
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
17pub struct TaskId(pub u64);
18
19impl TaskId {
20    /// 生成新的任务ID
21    pub fn new() -> Self {
22        use core::sync::atomic::{AtomicU64, Ordering};
23        static COUNTER: AtomicU64 = AtomicU64::new(1);
24        TaskId(COUNTER.fetch_add(1, Ordering::Relaxed))
25    }
26}
27
28impl Default for TaskId {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34/// 任务状态
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub enum TaskState {
37    /// 等待执行
38    Pending,
39    /// 已被唤醒,准备执行
40    Woken,
41    /// 正在执行
42    Running,
43    /// 已完成
44    Completed,
45}
46
47/// 任务元数据
48#[derive(Debug)]
49pub struct TaskMetadata {
50    /// 任务ID
51    pub id: TaskId,
52    /// 任务状态
53    pub state: TaskState,
54    /// 创建时间
55    pub created_at: u64,
56    /// 最后唤醒时间
57    pub last_wake_at: u64,
58    /// 最后执行时间
59    pub last_run_at: Option<u64>,
60    /// 执行次数
61    pub run_count: u64,
62}
63
64impl TaskMetadata {
65    /// 创建新的任务元数据
66    pub fn new(id: TaskId) -> Self {
67        let now = time::since_boot().as_millis() as u64;
68        Self {
69            id,
70            state: TaskState::Woken, // 新任务默认为 Woken 状态,可以立即执行
71            created_at: now,
72            last_wake_at: now,
73            last_run_at: None,
74            run_count: 0,
75        }
76    }
77
78    /// 检查任务是否超过指定时间未执行
79    pub fn is_expired(&self, timeout_ms: u64) -> bool {
80        if let Some(last_run) = self.last_run_at {
81            time::since_boot().as_millis() as u64 - last_run > timeout_ms
82        } else {
83            // 从未执行过的任务也检查创建时间
84            time::since_boot().as_millis() as u64 - self.created_at > timeout_ms
85        }
86    }
87
88    /// 更新执行时间
89    pub fn update_execution(&mut self) {
90        self.last_run_at = Some(time::since_boot().as_millis() as u64);
91        self.run_count += 1;
92        self.state = TaskState::Running;
93    }
94
95    /// 标记任务为已唤醒
96    pub fn mark_woken(&mut self) {
97        self.last_wake_at = time::since_boot().as_millis() as u64;
98        self.state = TaskState::Woken;
99    }
100
101    /// 标记任务为已完成
102    pub fn mark_completed(&mut self) {
103        self.state = TaskState::Completed;
104    }
105}
106
107/// 任务的优先级
108#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
109pub struct TaskPriority {
110    /// 是否为唤醒任务(最高优先级)
111    pub is_woken: bool,
112    /// 时间戳(越小越优先)
113    pub timestamp: u64,
114    /// 任务ID(作为平局打破器)
115    pub task_id: TaskId,
116}
117
118impl TaskPriority {
119    /// 创建新的任务优先级
120    pub fn new(metadata: &TaskMetadata) -> Self {
121        Self {
122            is_woken: metadata.state == TaskState::Woken,
123            timestamp: metadata.last_wake_at,
124            task_id: metadata.id,
125        }
126    }
127}
128
129/// 任务引用
130pub struct TaskRef {
131    /// 任务元数据
132    pub metadata: IrqSpinlock<TaskMetadata>,
133    /// Future对象的Pin引用
134    pub future: IrqSpinlock<Pin<Box<dyn Future<Output = ()> + Send>>>,
135}
136
137impl core::fmt::Debug for TaskRef {
138    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
139        let id = self.metadata.lock().id;
140        write!(f, "TaskRef({id:?})")
141    }
142}
143
144impl TaskRef {
145    /// 创建新的任务引用
146    pub fn new<F>(future: F, metadata: TaskMetadata) -> Self
147    where
148        F: Future<Output = ()> + Send + 'static,
149    {
150        Self {
151            metadata: IrqSpinlock::new(metadata),
152            future: IrqSpinlock::new(Box::pin(future)),
153        }
154    }
155
156    /// 获取任务ID
157    pub fn id(&self) -> TaskId {
158        self.metadata.lock().id
159    }
160
161    /// 检查任务是否已完成
162    pub fn is_completed(&self) -> bool {
163        self.metadata.lock().state == TaskState::Completed
164    }
165
166    /// 获取任务优先级
167    pub fn priority(&self) -> TaskPriority {
168        TaskPriority::new(&self.metadata.lock())
169    }
170
171    /// 轮询任务Future
172    pub fn poll(&self, waker: &Waker) -> Poll<()> {
173        let metadata = &mut self.metadata.lock();
174        if metadata.state == TaskState::Completed {
175            return Poll::Ready(());
176        }
177
178        metadata.update_execution();
179
180        let future = &mut self.future.lock();
181        let mut cx = Context::from_waker(waker);
182
183        match future.as_mut().poll(&mut cx) {
184            Poll::Ready(()) => {
185                metadata.mark_completed();
186                Poll::Ready(())
187            }
188            Poll::Pending => {
189                metadata.state = TaskState::Pending;
190                Poll::Pending
191            }
192        }
193    }
194}
195
196/// 任务句柄,用于与任务交互
197#[derive(Debug, Clone)]
198pub struct TaskHandle {
199    /// 任务ID
200    pub id: TaskId,
201    /// 内部任务引用(可选,用于唤醒)
202    task_ref: Option<Arc<TaskRef>>,
203}
204
205impl TaskHandle {
206    /// 创建新的任务句柄
207    pub fn new(id: TaskId) -> Self {
208        Self { id, task_ref: None }
209    }
210
211    /// 创建带有任务引用的句柄
212    pub fn with_ref(id: TaskId, task_ref: Arc<TaskRef>) -> Self {
213        Self {
214            id,
215            task_ref: Some(task_ref),
216        }
217    }
218
219    /// 唤醒对应的任务
220    pub fn wake(&self) -> bool {
221        if let Some(task_ref) = &self.task_ref {
222            task_ref.metadata.lock().mark_woken();
223            true
224        } else {
225            false
226        }
227    }
228}