#[track_caller]
pub fn ax_sleep_until(deadline: crate::time::AxTimeValue) {
ax_runtime::task::thread::current::sleep_until(
u64::try_from(deadline.as_nanos())
.ok()
.and_then(ax_runtime::task::time::MonotonicDeadline::from_nanos)
.expect("absolute sleep deadline exceeds the kernel monotonic time domain"),
);
}
#[track_caller]
pub fn ax_yield_now() {
if let Err(error) = ax_runtime::task::thread::current::yield_current_cpu() {
panic!("ax_yield_now failed at a scheduler safe point: {error}");
}
}
#[track_caller]
pub fn ax_exit(exit_code: i32) -> ! {
ax_runtime::thread::exit_current(exit_code);
}
cfg_task! {
use core::time::Duration;
use {ax_runtime::task::sched::CpuId, ax_runtime::task::sched::CpuSet};
pub struct AxTaskHandle {
inner: ax_runtime::task::thread::ThreadHandle,
id: u64,
}
impl AxTaskHandle {
pub fn id(&self) -> u64 {
self.id
}
}
pub type AxCpuMask = ax_cpumask::CpuMask<{ ax_runtime::CPU_CAPACITY }>;
pub use {ax_runtime::task::sync::RawMutex as AxRawMutex};
pub struct AxWaitQueueHandle(ax_runtime::task::sync::WaitQueue);
impl AxWaitQueueHandle {
pub const fn new() -> Self {
Self(ax_runtime::task::sync::WaitQueue::new())
}
}
impl Default for AxWaitQueueHandle {
fn default() -> Self {
Self::new()
}
}
pub fn ax_current_task_id() -> u64 {
ax_runtime::task::thread::current::current_thread_id()
.unwrap_or_else(|error| panic!("current task is unavailable: {error}"))
.as_u64()
}
pub fn ax_spawn<F>(f: F, name: alloc::string::String, stack_size: usize) -> AxTaskHandle
where
F: FnOnce() + Send + 'static,
{
let inner = ax_runtime::thread::builder(name).stack_size(stack_size).spawn(f)
.unwrap_or_else(|error| panic!("failed to spawn task: {error}"));
AxTaskHandle {
id: inner.id().as_u64(),
inner,
}
}
#[track_caller]
pub fn ax_wait_for_exit(task: AxTaskHandle) -> i32 {
(task.inner).join()
.unwrap_or_else(|error| panic!("failed to join task: {error}"))
}
pub fn ax_set_current_priority(prio: isize) -> crate::ApiResult {
use {ax_runtime::task::sched::Nice, ax_runtime::task::sched::SchedulePolicy};
let nice = i8::try_from(prio)
.ok()
.and_then(|value| Nice::new(value).ok())
.ok_or(crate::ApiError::InvalidInput)?;
let thread = task_result(
ax_runtime::task::thread::current::current_thread_handle(),
"read current task handle",
)?;
let policy = thread.base_policy();
let SchedulePolicy::Fair { mode, .. } = policy else {
return Err(crate::ApiError::OperationNotSupported);
};
task_result(
thread.set_policy(SchedulePolicy::fair(nice, mode)),
"set current task priority",
)
}
#[track_caller]
pub fn ax_set_current_affinity(cpumask: AxCpuMask) -> crate::ApiResult {
let topology_len = task_result(
ax_runtime::task::sched::cpu_topology_len(),
"read task CPU topology",
)?;
let affinity = cpu_set_from_mask(cpumask, topology_len)?;
task_result(
ax_runtime::task::thread::current::set_current_thread_affinity(affinity),
"set current task affinity",
)
}
#[track_caller]
pub fn ax_wait_queue_wait(wq: &AxWaitQueueHandle, timeout: Option<Duration>) -> bool {
if let Some(dur) = timeout {
return wq.0.wait_timeout(dur);
}
wq.0.wait();
false
}
#[track_caller]
pub fn ax_wait_queue_wait_until(
wq: &AxWaitQueueHandle,
until_condition: impl Fn() -> bool,
timeout: Option<Duration>,
) -> bool {
if let Some(dur) = timeout {
return wq.0.wait_timeout_until(dur, until_condition);
}
wq.0.wait_until(until_condition);
false
}
#[track_caller]
pub fn ax_wait_queue_wait_until_deadline(
wq: &AxWaitQueueHandle,
deadline: Duration,
until_condition: impl Fn() -> bool,
) -> bool {
wq.0.wait_until_deadline(
u64::try_from(deadline.as_nanos())
.ok()
.and_then(ax_runtime::task::time::MonotonicDeadline::from_nanos)
.expect("wait deadline exceeds the kernel monotonic time domain"),
until_condition,
)
}
pub fn ax_wait_queue_wake(wq: &AxWaitQueueHandle, count: u32) -> usize {
let mut woken = 0;
if count == u32::MAX {
while wq.0.notify_one() {
woken += 1;
}
} else {
for _ in 0..count {
if !wq.0.notify_one() {
break;
}
woken += 1;
}
}
woken
}
fn task_result<T>(
result: Result<T, ax_runtime::task::thread::TaskError>,
operation: &'static str,
) -> crate::ApiResult<T> {
result.map_err(|error| {
ax_log::warn!("{operation} failed: {error}");
error.into()
})
}
fn cpu_set_from_mask(cpumask: AxCpuMask, topology_len: usize) -> crate::ApiResult<CpuSet> {
if cpumask.is_empty() {
return Err(crate::ApiError::InvalidInput);
}
let mut affinity = CpuSet::empty(topology_len);
for cpu_index in &cpumask {
let cpu_index =
u32::try_from(cpu_index).map_err(|_| crate::ApiError::InvalidInput)?;
if !affinity.insert(CpuId::new(cpu_index)) {
return Err(crate::ApiError::InvalidInput);
}
}
Ok(affinity)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cpu_mask_conversion_preserves_allowed_cpu() {
let affinity = cpu_set_from_mask(AxCpuMask::one_shot(0), 1).unwrap();
assert!(affinity.contains(CpuId::new(0)));
}
#[test]
fn cpu_mask_conversion_rejects_empty_mask() {
assert_eq!(
cpu_set_from_mask(AxCpuMask::new(), 1),
Err(crate::ApiError::InvalidInput)
);
}
#[test]
fn cpu_mask_conversion_rejects_cpu_outside_topology() {
assert_eq!(
cpu_set_from_mask(AxCpuMask::one_shot(0), 0),
Err(crate::ApiError::InvalidInput)
);
}
#[test]
fn pi_chain_limit_maps_to_bad_state() {
assert_eq!(
ax_io::IoError::from(crate::ApiError::from(
ax_runtime::task::thread::TaskError::PiChainLimit { limit: 8 }
)),
ax_io::IoError::BadState
);
}
#[test]
fn thread_capacity_maps_to_linux_eagain() {
assert_eq!(
ax_io::IoError::from(crate::ApiError::from(
ax_runtime::task::thread::TaskError::ThreadCapacity
)),
ax_io::IoError::WouldBlock
);
}
}
}