use generic wakeup instead of specialized

src/clock/mod.rs

@@ -1,5 +1,6 @@
 pub mod realtime;
 pub mod syscalls;
+pub mod timer;
 pub mod timespec;
 
 pub enum ClockId {

src/clock/syscalls/itimer.rs

@@ -1,13 +1,16 @@
+use crate::clock::timer::{TimerNamespace, make_timer_id, parse_timer_id};
 use crate::clock::timespec::TimeSpec;
 use crate::drivers::timer::{Instant, now, uptime};
 use crate::memory::uaccess::{UserCopyable, copy_from_user, copy_to_user};
 use crate::process::thread_group::signal::SigId;
 use crate::process::{ITimer, Task, Tid, find_task_by_tid};
 use crate::sched::syscall_ctx::ProcessCtx;
+use alloc::boxed::Box;
 use core::time::Duration;
 use libkernel::memory::address::TUA;
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[repr(u32)]
 pub enum ITimerType {
     Real = 0,
     Virtual = 1,
@@ -61,7 +64,15 @@ impl Default for ITimerVal {
 
 unsafe impl UserCopyable for ITimerVal {}
 
-pub fn itimer_irq_handler(tid: Tid, ty: ITimerType) -> Option<Instant> {
+pub fn itimer_irq_handler(tid: Tid, id: u64) -> Option<Instant> {
+    let (namespace, ty) = parse_timer_id(id);
+    if namespace != TimerNamespace::ITimer {
+        return None; // Not an itimer, should not happen
+    }
+    let ty = match ITimerType::try_from(ty as i32) {
+        Ok(ty) => ty,
+        Err(_) => return None, // Invalid timer type, should not happen
+    };
     let task = find_task_by_tid(tid)?;
     match ty {
         ITimerType::Real => {
@@ -161,14 +172,22 @@ pub async fn sys_setitimer(
                 crate::drivers::timer::SYS_TIMER
                     .get()
                     .unwrap()
-                    .remove_scheduled_itimer(current_task.tid(), ITimerType::Real);
+                    .remove_scheduled_timer(
+                        current_task.tid(),
+                        make_timer_id(TimerNamespace::ITimer, ITimerType::Real as u32),
+                    );
             }
             if let Some(next) = next {
                 timers.real = Some(ITimer { interval, next });
                 crate::drivers::timer::SYS_TIMER
                     .get()
                     .unwrap()
-                    .schedule_itimer(current_task.tid(), ITimerType::Real, next);
+                    .schedule_timer(
+                        current_task.tid(),
+                        make_timer_id(TimerNamespace::ITimer, ITimerType::Real as u32),
+                        Box::new(itimer_irq_handler),
+                        next,
+                    );
             }
         }
         _ => unimplemented!(),

src/clock/timer.rs

@@ -0,0 +1,60 @@
+//! Handling of timers and other asynchronous callbacks
+//!
+//! We represent timer data as an u64:
+//! u16 namespace, u16 unused, u32 id
+//!
+//! This is used by the callback to identify the relevant timer. It is also used to remove timers
+
+/// Not a timer, generic callback
+/// Catch all for all invalid namespaces.
+const NAMESPACE_NONE: u16 = 0;
+const NAMESPACE_ITIMER: u16 = 1;
+const NAMESPACE_TIMERFD: u16 = 2;
+const NAMESPACE_TIMER: u16 = 3;
+
+#[derive(Copy, Clone, Debug, PartialEq)]
+#[repr(u16)]
+pub enum TimerNamespace {
+    None = NAMESPACE_NONE,
+    ITimer = NAMESPACE_ITIMER,
+    TimerFd = NAMESPACE_TIMERFD,
+    Timer = NAMESPACE_TIMER,
+}
+
+impl From<u16> for TimerNamespace {
+    fn from(value: u16) -> Self {
+        match value {
+            NAMESPACE_NONE => Self::None,
+            NAMESPACE_ITIMER => Self::ITimer,
+            NAMESPACE_TIMERFD => Self::TimerFd,
+            NAMESPACE_TIMER => Self::Timer,
+            _ => Self::None, // Default to None for invalid namespaces
+        }
+    }
+}
+
+pub fn make_timer_id(namespace: TimerNamespace, id: u32) -> u64 {
+    (namespace as u64) << 48 | (id as u64)
+}
+
+pub fn parse_timer_id(timer_id: u64) -> (TimerNamespace, u32) {
+    let namespace = TimerNamespace::from((timer_id >> 48) as u16);
+    let id = (timer_id & 0xFFFF_FFFF) as u32;
+    (namespace, id)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::clock::timer::{TimerNamespace, make_timer_id, parse_timer_id};
+    use moss_macros::ktest;
+
+    #[ktest]
+    fn test_timer_id_encoding() {
+        let namespace = TimerNamespace::ITimer;
+        let id = 42;
+        let timer_id = make_timer_id(namespace, id);
+        let (parsed_namespace, parsed_id) = parse_timer_id(timer_id);
+        assert_eq!(namespace, parsed_namespace);
+        assert_eq!(id, parsed_id);
+    }
+}

src/drivers/timer/mod.rs

@@ -1,9 +1,9 @@
 use super::Driver;
-use crate::clock::syscalls::itimer::ITimerType;
 use crate::interrupts::{InterruptDescriptor, InterruptHandler};
 use crate::per_cpu_private;
 use crate::process::Tid;
 use crate::sync::OnceLock;
+use alloc::boxed::Box;
 use alloc::{collections::binary_heap::BinaryHeap, sync::Arc};
 use core::{
     future::poll_fn,
@@ -77,7 +77,11 @@ enum WakeupKind {
     /// This wake up is for the kernel's preemption mechanism.
     Preempt,
 
-    ITimer(Tid, ITimerType),
+    Timer(
+        Tid,
+        u64,
+        Box<dyn Fn(Tid, u64) -> Option<Instant> + Send + Sync>,
+    ),
 }
 
 unsafe impl Send for WakeupKind {}
@@ -181,8 +185,14 @@ impl InterruptHandler for SysTimer {
                         // Do nothing, the IRQ return-to-userspace code will
                         // call schedule() for us.
                     }
-                    WakeupKind::ITimer(tid, ty) => {
-                        crate::clock::syscalls::itimer::itimer_irq_handler(tid, ty);
+                    WakeupKind::Timer(tid, timer_id, callback) => {
+                        if let Some(next_instant) = callback(tid, timer_id) {
+                            // Re-schedule the timer for its next expiration.
+                            wake_q.push(WakeupEvent {
+                                when: next_instant,
+                                what: WakeupKind::Timer(tid, timer_id, callback),
+                            });
+                        }
                     }
                 }
             } else {
@@ -241,12 +251,18 @@ impl SysTimer {
         .await;
     }
 
-    pub fn schedule_itimer(&self, tid: Tid, ty: ITimerType, when: Instant) {
+    pub fn schedule_timer(
+        &self,
+        tid: Tid,
+        id: u64,
+        callback: Box<dyn Fn(Tid, u64) -> Option<Instant> + Send + Sync>,
+        when: Instant,
+    ) {
         let mut wakeup_q = WAKEUP_Q.borrow_mut();
 
         wakeup_q.push(WakeupEvent {
             when,
-            what: WakeupKind::ITimer(tid, ty),
+            what: WakeupKind::Timer(tid, id, callback),
         });
 
         // After pushing, we must update the hardware timer in case our
@@ -256,18 +272,17 @@ impl SysTimer {
         }
     }
 
-    pub fn remove_scheduled_itimer(&self, tid: Tid, ty: ITimerType) {
+    pub fn remove_scheduled_timer(&self, tid: Tid, id: u64) {
         let mut wakeup_q = WAKEUP_Q.borrow_mut();
 
-        // Remove any matching itimer events from the queue.
+        // Remove any timers matching the given tid and id.
         wakeup_q.retain(|event| {
-            if let WakeupKind::ITimer(event_tid, event_ty) = event.what {
-                !(event_tid == tid && event_ty == ty)
+            if let WakeupKind::Timer(event_tid, event_id, _) = &event.what {
+                !(event_tid == &tid && event_id == &id)
             } else {
                 true
             }
         });
-        // TODO: Handle cross-CPU cases where the itimer event might be on a different CPU's queue.
 
         // After removing, we must update the hardware timer in case we removed
         // the earliest event.