os/kernel/task/
tid_allocator.rs

1//! 简单的任务ID分配器实现
2
3use core::sync::atomic::AtomicU32;
4
5/// 简单的任务ID分配器。
6/// 每次调用`allocate`方法时,都会返回一个唯一的任务ID。
7/// 任务ID从1开始递增。
8#[derive(Debug)]
9pub struct TidAllocator {
10    next_tid: AtomicU32,
11}
12
13impl TidAllocator {
14    /// 创建一个新的TidAllocator实例。
15    pub const fn new() -> Self {
16        TidAllocator {
17            next_tid: AtomicU32::new(1),
18        }
19    }
20
21    /// 分配一个新的任务ID。
22    pub fn allocate(&self) -> u32 {
23        self.next_tid
24            .fetch_add(1, core::sync::atomic::Ordering::SeqCst)
25    }
26}
27
28#[cfg(test)]
29mod tests {
30    use super::*;
31    use crate::{kassert, test_case};
32
33    // 顺序分配测试:检查分配值从1开始并依次递增
34    test_case!(test_tid_allocate_sequence, {
35        let alloc = TidAllocator::new();
36        kassert!(alloc.allocate() == 1);
37        kassert!(alloc.allocate() == 2);
38        kassert!(alloc.allocate() == 3);
39    });
40
41    // 多引用调用测试:通过多个引用连续分配,确保值唯一且递增
42    test_case!(test_tid_allocate_multiple_refs, {
43        let alloc = TidAllocator::new();
44        let a1 = alloc.allocate();
45        let a2 = alloc.allocate();
46        kassert!(a1 == 1);
47        kassert!(a2 == 2);
48
49        // 通过另一个不可变引用继续分配
50        let r = &alloc;
51        let a3 = r.allocate();
52        kassert!(a3 == 3);
53    });
54}