ax_api/imp/
task.rs

1#[track_caller]
2pub fn ax_sleep_until(deadline: crate::time::AxTimeValue) {
3    ax_runtime::task::thread::current::sleep_until(
4        u64::try_from(deadline.as_nanos())
5            .ok()
6            .and_then(ax_runtime::task::time::MonotonicDeadline::from_nanos)
7            .expect("absolute sleep deadline exceeds the kernel monotonic time domain"),
8    );
9}
10
11#[track_caller]
12pub fn ax_yield_now() {
13    if let Err(error) = ax_runtime::task::thread::current::yield_current_cpu() {
14        panic!("ax_yield_now failed at a scheduler safe point: {error}");
15    }
16}
17
18#[track_caller]
19pub fn ax_exit(exit_code: i32) -> ! {
20    ax_runtime::thread::exit_current(exit_code);
21}
22
23cfg_task! {
24    use core::time::Duration;
25    use {ax_runtime::task::sched::CpuId, ax_runtime::task::sched::CpuSet};
26
27    /// A handle to a task.
28    pub struct AxTaskHandle {
29        inner: ax_runtime::task::thread::ThreadHandle,
30        id: u64,
31    }
32
33    impl AxTaskHandle {
34        /// Returns the task ID.
35        pub fn id(&self) -> u64 {
36            self.id
37        }
38    }
39
40    /// A mask to specify the CPU affinity.
41    pub type AxCpuMask = ax_cpumask::CpuMask<{ ax_runtime::CPU_CAPACITY }>;
42
43    pub use {ax_runtime::task::sync::RawMutex as AxRawMutex};
44
45    /// A handle to a wait queue.
46    ///
47    /// A wait queue is used to store sleeping tasks waiting for a certain event
48    /// to happen.
49    pub struct AxWaitQueueHandle(ax_runtime::task::sync::WaitQueue);
50
51    impl AxWaitQueueHandle {
52        /// Creates a new empty wait queue.
53        pub const fn new() -> Self {
54            Self(ax_runtime::task::sync::WaitQueue::new())
55        }
56    }
57
58    impl Default for AxWaitQueueHandle {
59        fn default() -> Self {
60            Self::new()
61        }
62    }
63
64    pub fn ax_current_task_id() -> u64 {
65        ax_runtime::task::thread::current::current_thread_id()
66            .unwrap_or_else(|error| panic!("current task is unavailable: {error}"))
67            .as_u64()
68    }
69
70    pub fn ax_spawn<F>(f: F, name: alloc::string::String, stack_size: usize) -> AxTaskHandle
71    where
72        F: FnOnce() + Send + 'static,
73    {
74        let inner = ax_runtime::thread::spawn_raw(f, name, stack_size)
75            .unwrap_or_else(|error| panic!("failed to spawn task: {error}"));
76        AxTaskHandle {
77            id: inner.id().as_u64(),
78            inner,
79        }
80    }
81
82    #[track_caller]
83    pub fn ax_wait_for_exit(task: AxTaskHandle) -> i32 {
84        ax_runtime::thread::join_thread(task.inner)
85            .unwrap_or_else(|error| panic!("failed to join task: {error}"))
86    }
87
88    pub fn ax_set_current_priority(prio: isize) -> crate::ApiResult {
89        use {ax_runtime::task::sched::Nice, ax_runtime::task::sched::SchedulePolicy};
90
91        let nice = i8::try_from(prio)
92            .ok()
93            .and_then(|value| Nice::new(value).ok())
94            .ok_or(crate::ApiError::InvalidInput)?;
95        let thread = task_result(
96            ax_runtime::task::thread::current::current_thread_handle(),
97            "read current task handle",
98        )?;
99        let policy = thread.base_policy();
100        let SchedulePolicy::Fair { mode, .. } = policy else {
101            return Err(crate::ApiError::OperationNotSupported);
102        };
103        task_result(
104            thread.set_policy(SchedulePolicy::fair(nice, mode)),
105            "set current task priority",
106        )
107    }
108
109    #[track_caller]
110    pub fn ax_set_current_affinity(cpumask: AxCpuMask) -> crate::ApiResult {
111        let topology_len = task_result(
112            ax_runtime::task::sched::cpu_topology_len(),
113            "read task CPU topology",
114        )?;
115        let affinity = cpu_set_from_mask(cpumask, topology_len)?;
116        task_result(
117            ax_runtime::task::thread::current::set_current_thread_affinity(affinity),
118            "set current task affinity",
119        )
120    }
121
122    #[track_caller]
123    pub fn ax_wait_queue_wait(wq: &AxWaitQueueHandle, timeout: Option<Duration>) -> bool {
124        if let Some(dur) = timeout {
125            return wq.0.wait_timeout(dur);
126        }
127
128        wq.0.wait();
129        false
130    }
131
132    #[track_caller]
133    pub fn ax_wait_queue_wait_until(
134        wq: &AxWaitQueueHandle,
135        until_condition: impl Fn() -> bool,
136        timeout: Option<Duration>,
137    ) -> bool {
138        if let Some(dur) = timeout {
139            return wq.0.wait_timeout_until(dur, until_condition);
140        }
141
142        wq.0.wait_until(until_condition);
143        false
144    }
145
146    /// Blocks until `until_condition` becomes true or the absolute monotonic
147    /// `deadline` elapses.
148    ///
149    /// Returns `true` only when the deadline wins.
150    #[track_caller]
151    pub fn ax_wait_queue_wait_until_deadline(
152        wq: &AxWaitQueueHandle,
153        deadline: Duration,
154        until_condition: impl Fn() -> bool,
155    ) -> bool {
156        wq.0.wait_until_deadline(
157            u64::try_from(deadline.as_nanos())
158                .ok()
159                .and_then(ax_runtime::task::time::MonotonicDeadline::from_nanos)
160                .expect("wait deadline exceeds the kernel monotonic time domain"),
161            until_condition,
162        )
163    }
164
165    pub fn ax_wait_queue_wake(wq: &AxWaitQueueHandle, count: u32) -> usize {
166        let mut woken = 0;
167        if count == u32::MAX {
168            while wq.0.notify_one() {
169                woken += 1;
170            }
171        } else {
172            for _ in 0..count {
173                if !wq.0.notify_one() {
174                    break;
175                }
176                woken += 1;
177            }
178        }
179        woken
180    }
181
182    fn task_result<T>(
183        result: Result<T, ax_runtime::task::thread::TaskError>,
184        operation: &'static str,
185    ) -> crate::ApiResult<T> {
186        result.map_err(|error| {
187            ax_log::warn!("{operation} failed: {error}");
188            error.into()
189        })
190    }
191
192    fn cpu_set_from_mask(cpumask: AxCpuMask, topology_len: usize) -> crate::ApiResult<CpuSet> {
193        if cpumask.is_empty() {
194            return Err(crate::ApiError::InvalidInput);
195        }
196        let mut affinity = CpuSet::empty(topology_len);
197        for cpu_index in &cpumask {
198            let cpu_index =
199                u32::try_from(cpu_index).map_err(|_| crate::ApiError::InvalidInput)?;
200            if !affinity.insert(CpuId::new(cpu_index)) {
201                return Err(crate::ApiError::InvalidInput);
202            }
203        }
204        Ok(affinity)
205    }
206
207    #[cfg(test)]
208    mod tests {
209        use super::*;
210
211        #[test]
212        fn cpu_mask_conversion_preserves_allowed_cpu() {
213            let affinity = cpu_set_from_mask(AxCpuMask::one_shot(0), 1).unwrap();
214
215            assert!(affinity.contains(CpuId::new(0)));
216        }
217
218        #[test]
219        fn cpu_mask_conversion_rejects_empty_mask() {
220            assert_eq!(
221                cpu_set_from_mask(AxCpuMask::new(), 1),
222                Err(crate::ApiError::InvalidInput)
223            );
224        }
225
226        #[test]
227        fn cpu_mask_conversion_rejects_cpu_outside_topology() {
228            assert_eq!(
229                cpu_set_from_mask(AxCpuMask::one_shot(0), 0),
230                Err(crate::ApiError::InvalidInput)
231            );
232        }
233
234        #[test]
235        fn pi_chain_limit_maps_to_bad_state() {
236            assert_eq!(
237                ax_io::IoError::from(crate::ApiError::from(
238                    ax_runtime::task::thread::TaskError::PiChainLimit { limit: 8 }
239                )),
240                ax_io::IoError::BadState
241            );
242        }
243
244        #[test]
245        fn thread_capacity_maps_to_linux_eagain() {
246            assert_eq!(
247                ax_io::IoError::from(crate::ApiError::from(
248                    ax_runtime::task::thread::TaskError::ThreadCapacity
249                )),
250                ax_io::IoError::WouldBlock
251            );
252        }
253    }
254
255}
ax_api/imp/task.rs

ax_api/imp/
task.rs
