extern crate alloc;
use alloc::boxed::Box;
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
use core::sync::atomic::{AtomicUsize, Ordering};
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
use core::time::Duration;
use ax_kspin::SpinNoIrq;
use ax_lazyinit::LazyInit;
use ax_timer_list::{TimeValue, TimerEvent, TimerList};
use crate::host::{HostTime, default_host};
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
static TOKEN: AtomicUsize = AtomicUsize::new(0);
struct VmTimerEvent {
#[cfg(target_arch = "x86_64")]
token: usize,
callback: Box<dyn FnOnce(TimeValue) + Send + 'static>,
}
impl VmTimerEvent {
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
fn new<F>(token: usize, callback: F) -> Self
where
F: FnOnce(TimeValue) + Send + 'static,
{
#[cfg(not(target_arch = "x86_64"))]
let _ = token;
Self {
#[cfg(target_arch = "x86_64")]
token,
callback: Box::new(callback),
}
}
}
impl TimerEvent for VmTimerEvent {
fn callback(self, now: TimeValue) {
(self.callback)(now);
}
}
#[ax_percpu::def_percpu]
static TIMER_LIST: LazyInit<SpinNoIrq<TimerList<VmTimerEvent>>> = LazyInit::new();
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
pub(crate) fn register_timer(
deadline_ns: u64,
callback: Box<dyn FnOnce(Duration) + Send + 'static>,
) -> usize {
let token = TOKEN.fetch_add(1, Ordering::Relaxed);
let next_deadline = {
let timer_list = unsafe { TIMER_LIST.current_ref_mut_raw() };
let mut timers = timer_list.lock();
timers.set(
TimeValue::from_nanos(deadline_ns),
VmTimerEvent::new(token, callback),
);
timers.next_deadline()
};
rearm_host_timer(next_deadline);
token
}
#[cfg(target_arch = "x86_64")]
pub(crate) fn cancel_timer(token: usize) {
let next_deadline = {
let timer_list = unsafe { TIMER_LIST.current_ref_mut_raw() };
let mut timers = timer_list.lock();
timers.cancel(|event| event.token == token);
timers.next_deadline()
};
rearm_host_timer(next_deadline);
}
pub(crate) fn check_events() {
let timer_list = unsafe { TIMER_LIST.current_ref_mut_raw() };
loop {
let now = default_host().monotonic_time();
let expired = timer_list.lock().expire_one(now);
if let Some((deadline, event)) = expired {
trace!("handle VM timer event scheduled at {deadline:#?}");
event.callback(now);
} else {
rearm_host_timer(timer_list.lock().next_deadline());
break;
}
}
}
fn rearm_host_timer(next_deadline: Option<TimeValue>) {
if let Some(deadline) = next_deadline {
default_host().set_oneshot_timer(deadline.as_nanos() as u64);
}
}
pub(crate) fn init_percpu() {
info!("Initializing AxVM timer wheel...");
let timer_list = unsafe { TIMER_LIST.current_ref_mut_raw() };
timer_list.init_once(SpinNoIrq::new(TimerList::new()));
}