lock ordering

src/context/context.rs

@@ -19,6 +19,7 @@ use crate::{
     paging::{RmmA, RmmArch},
     percpu::PercpuBlock,
     scheme::{CallerCtx, FileHandle, SchemeId, SchemeNamespace},
+    sync::CleanLockToken,
 };
 
 use crate::syscall::error::{Error, Result, EAGAIN, EBADF, EEXIST, EINVAL, EMFILE, ESRCH};
@@ -428,11 +429,11 @@ pub struct BorrowedHtBuf {
     head_and_not_tail: bool,
 }
 impl BorrowedHtBuf {
-    pub fn head() -> Result<Self> {
+    pub fn head(token: &mut CleanLockToken) -> Result<Self> {
         Ok(Self {
             inner: Some(
                 context::current()
-                    .write()
+                    .write(token.token())
                     .syscall_head
                     .take()
                     .ok_or(Error::new(EAGAIN))?,
@@ -440,11 +441,11 @@ impl BorrowedHtBuf {
             head_and_not_tail: true,
         })
     }
-    pub fn tail() -> Result<Self> {
+    pub fn tail(token: &mut CleanLockToken) -> Result<Self> {
         Ok(Self {
             inner: Some(
                 context::current()
-                    .write()
+                    .write(token.token())
                     .syscall_tail
                     .take()
                     .ok_or(Error::new(EAGAIN))?,
@@ -495,7 +496,9 @@ impl Drop for BorrowedHtBuf {
         let Some(inner) = self.inner.take() else {
             return;
         };
-        match context.write() {
+        //TODO: do not allow drop so lock token can be passed in
+        let mut token = unsafe { CleanLockToken::new() };
+        match context.write(token.token()) {
             mut context => {
                 (if self.head_and_not_tail {
                     &mut context.syscall_head
@@ -871,10 +874,10 @@ impl FdTbl {
         FileHandle::from(start | UPPER_FDTBL_TAG)
     }
 
-    pub fn force_close_all(&mut self) {
+    pub fn force_close_all(&mut self, token: &mut CleanLockToken) {
         for file_opt in self.iter_mut() {
             if let Some(file) = file_opt.take() {
-                let _ = file.close();
+                let _ = file.close(token);
             }
         }
         self.active_count = 0;

src/context/file.rs

@@ -3,6 +3,7 @@
 use crate::{
     event,
     scheme::{self, SchemeId},
+    sync::CleanLockToken,
     syscall::error::{Error, Result, EBADF},
 };
 use alloc::sync::Arc;
@@ -70,22 +71,22 @@ pub struct FileDescriptor {
 impl FileDescription {
     /// Try closing a file, although at this point the description will be destroyed anyway, if
     /// doing so fails.
-    pub fn try_close(self) -> Result<()> {
+    pub fn try_close(self, token: &mut CleanLockToken) -> Result<()> {
         event::unregister_file(self.scheme, self.number);
 
-        let scheme = scheme::schemes()
+        let scheme = scheme::schemes(token.token())
             .get(self.scheme)
             .ok_or(Error::new(EBADF))?
             .clone();
 
-        scheme.close(self.number)
+        scheme.close(self.number, token)
     }
 }
 
 impl FileDescriptor {
-    pub fn close(self) -> Result<()> {
+    pub fn close(self, token: &mut CleanLockToken) -> Result<()> {
         if let Ok(file) = Arc::try_unwrap(self.description).map(RwLock::into_inner) {
-            file.try_close()?;
+            file.try_close(token)?;
         }
         Ok(())
     }

src/context/memory.rs

@@ -21,6 +21,7 @@ use crate::{
     paging::{Page, PageFlags, PageMapper, RmmA, TableKind, VirtualAddress},
     percpu::PercpuBlock,
     scheme::{self, KernelSchemes},
+    sync::CleanLockToken,
 };
 
 use super::{context::HardBlockedReason, file::FileDescription};
@@ -51,7 +52,7 @@ pub struct UnmapResult {
     pub flags: MunmapFlags,
 }
 impl UnmapResult {
-    pub fn unmap(mut self) -> Result<()> {
+    pub fn unmap(mut self, token: &mut CleanLockToken) -> Result<()> {
         let Some(GrantFileRef {
             base_offset,
             description,
@@ -64,14 +65,13 @@ impl UnmapResult {
             ref desc => (desc.scheme, desc.number),
         };
 
-        let funmap_result = scheme::schemes()
-            .get(scheme_id)
-            .cloned()
+        let scheme_opt = scheme::schemes(token.token()).get(scheme_id).cloned();
+        let funmap_result = scheme_opt
             .ok_or(Error::new(ENODEV))
-            .and_then(|scheme| scheme.kfunmap(number, base_offset, self.size, self.flags));
+            .and_then(|scheme| scheme.kfunmap(number, base_offset, self.size, self.flags, token));
 
         if let Ok(fd) = Arc::try_unwrap(description) {
-            fd.into_inner().try_close()?;
+            fd.into_inner().try_close(token)?;
         }
         funmap_result?;
 
@@ -502,7 +502,11 @@ impl AddrSpaceWrapper {
     }
     /// Borrows a page from user memory, requiring that the frame be Allocated and read/write. This
     /// is intended to be used for user-kernel shared memory.
-    pub fn borrow_frame_enforce_rw_allocated(self: &Arc<Self>, page: Page) -> Result<RaiiFrame> {
+    pub fn borrow_frame_enforce_rw_allocated(
+        self: &Arc<Self>,
+        page: Page,
+        token: &mut CleanLockToken,
+    ) -> Result<RaiiFrame> {
         let mut guard = self.acquire_write();
 
         let (_start_page, info) = guard.grants.contains(page).ok_or(Error::new(EINVAL))?;
@@ -519,8 +523,9 @@ impl AddrSpaceWrapper {
         {
             Frame::containing(f)
         } else {
-            let (frame, flush, new_guard) = correct_inner(self, guard, page, AccessMode::Write, 0)
-                .map_err(|_| Error::new(ENOMEM))?;
+            let (frame, flush, new_guard) =
+                correct_inner(self, guard, page, AccessMode::Write, 0, token)
+                    .map_err(|_| Error::new(ENOMEM))?;
             flush.flush();
             guard = new_guard;
 
@@ -1333,6 +1338,7 @@ impl Grant {
         lock: &AddrSpaceWrapper,
         mapper: &mut PageMapper,
         flusher: &mut Flusher,
+        token: &mut CleanLockToken,
     ) -> Result<Self> {
         if let Some(src) = src {
             let mut guard = src.addr_space_guard;
@@ -1359,6 +1365,7 @@ impl Grant {
                                     src_page,
                                     AccessMode::Read,
                                     0,
+                                    token,
                                 )
                                 .map_err(|_| Error::new(EIO))?;
                                 guard = new_guard;
@@ -1383,6 +1390,7 @@ impl Grant {
                                     src_page,
                                     AccessMode::Read,
                                     0,
+                                    token,
                                 )
                                 .map_err(|_| Error::new(EIO))?;
                                 guard = new_guard;
@@ -2232,6 +2240,9 @@ pub struct Table {
 
 impl Drop for AddrSpace {
     fn drop(&mut self) {
+        //TODO: DANGER: unmap requires a CleanLockToken, this cheats!
+        let mut token = unsafe { CleanLockToken::new() };
+
         for mut grant in core::mem::take(&mut self.grants).into_iter() {
             // Unpinning the grant is allowed, because pinning only occurs in UserScheme calls to
             // prevent unmapping the mapped range twice (which would corrupt only the scheme
@@ -2245,7 +2256,7 @@ impl Drop for AddrSpace {
             // the underlying pages (and send some funmaps).
             let res = grant.unmap(&mut self.table.utable, &mut NopFlusher);
 
-            let _ = res.unmap();
+            let _ = res.unmap(&mut token);
         }
     }
 }
@@ -2372,14 +2383,18 @@ pub unsafe fn copy_frame_to_frame_directly(dst: Frame, src: Frame) {
     }
 }
 
-pub fn try_correcting_page_tables(faulting_page: Page, access: AccessMode) -> Result<(), PfError> {
+pub fn try_correcting_page_tables(
+    faulting_page: Page,
+    access: AccessMode,
+    token: &mut CleanLockToken,
+) -> Result<(), PfError> {
     let Ok(addr_space_lock) = AddrSpace::current() else {
         debug!("User page fault without address space being set.");
         return Err(PfError::Segv);
     };
 
     let lock = &addr_space_lock;
-    let (_, flush, _) = correct_inner(lock, lock.acquire_write(), faulting_page, access, 0)?;
+    let (_, flush, _) = correct_inner(lock, lock.acquire_write(), faulting_page, access, 0, token)?;
 
     flush.flush();
 
@@ -2391,6 +2406,7 @@ fn correct_inner<'l>(
     faulting_page: Page,
     access: AccessMode,
     recursion_level: u32,
+    token: &mut CleanLockToken,
 ) -> Result<(Frame, PageFlush<RmmA>, RwLockWriteGuard<'l, AddrSpace>), PfError> {
     let mut addr_space = &mut *addr_space_guard;
     let mut flusher = Flusher::with_cpu_set(&mut addr_space.used_by, &addr_space_lock.tlb_ack);
@@ -2533,6 +2549,7 @@ fn correct_inner<'l>(
                                     src_page,
                                     AccessMode::Read,
                                     new_recursion_level,
+                                    token,
                                 )?
                             };
 
@@ -2610,7 +2627,7 @@ fn correct_inner<'l>(
             let (scheme_id, scheme_number) = match file_ref.description.read() {
                 ref desc => (desc.scheme, desc.number),
             };
-            let user_inner = scheme::schemes()
+            let user_inner = scheme::schemes(token.token())
                 .get(scheme_id)
                 .and_then(|s| {
                     if let KernelSchemes::User(user) = s {
@@ -2623,18 +2640,18 @@ fn correct_inner<'l>(
 
             let offset = file_ref.base_offset as u64 + (pages_from_grant_start * PAGE_SIZE) as u64;
             user_inner
-                .request_fmap(scheme_number, offset, 1, flags)
+                .request_fmap(scheme_number, offset, 1, flags, token)
                 .unwrap();
 
             let context_lock = crate::context::current();
             context_lock
-                .write()
+                .write(token.token())
                 .hard_block(HardBlockedReason::AwaitingMmap { file_ref });
 
-            super::switch();
+            super::switch(token);
 
             let frame = context_lock
-                .write()
+                .write(token.token())
                 .fmap_ret
                 .take()
                 .ok_or(PfError::NonfatalInternalError)?;
@@ -2678,9 +2695,9 @@ pub struct BorrowedFmapSource<'a> {
     pub addr_space_guard: RwLockWriteGuard<'a, AddrSpace>,
 }
 
-pub fn handle_notify_files(notify_files: Vec<UnmapResult>) {
+pub fn handle_notify_files(notify_files: Vec<UnmapResult>, token: &mut CleanLockToken) {
     for file in notify_files {
-        let _ = file.unmap();
+        let _ = file.unmap(token);
     }
 }
 

src/context/mod.rs

@@ -2,12 +2,8 @@
 //!
 //! For resources on contexts, please consult [wikipedia](https://en.wikipedia.org/wiki/Context_switch) and  [osdev](https://wiki.osdev.org/Context_Switching)
 
-use core::num::NonZeroUsize;
-
 use alloc::{collections::BTreeSet, sync::Arc};
-
-use spin::{RwLock, RwLockReadGuard, RwLockWriteGuard};
-use spinning_top::RwSpinlock;
+use core::num::NonZeroUsize;
 use syscall::ENOMEM;
 
 use crate::{
@@ -15,6 +11,10 @@ use crate::{
     cpu_set::LogicalCpuSet,
     paging::{RmmA, RmmArch, TableKind},
     percpu::PercpuBlock,
+    sync::{
+        ArcRwLockWriteGuard, CleanLockToken, LockToken, RwLock, RwLockReadGuard, RwLockWriteGuard,
+        L0, L1, L2,
+    },
     syscall::error::{Error, Result},
 };
 
@@ -24,6 +24,9 @@ pub use self::{
     switch::switch,
 };
 
+pub type ContextLock = RwLock<L2, Context>;
+pub type ArcContextLockWriteGuard = ArcRwLockWriteGuard<L2, Context>;
+
 #[cfg(target_arch = "aarch64")]
 #[path = "arch/aarch64.rs"]
 mod arch;
@@ -67,9 +70,21 @@ pub use self::arch::empty_cr3;
 
 // Set of weak references to all contexts available for scheduling. The only strong references are
 // the context file descriptors.
-static CONTEXTS: RwLock<BTreeSet<ContextRef>> = RwLock::new(BTreeSet::new());
+static CONTEXTS: RwLock<L1, BTreeSet<ContextRef>> = RwLock::new(BTreeSet::new());
 
-pub fn init() {
+/// Get the global schemes list, const
+pub fn contexts<'a>(token: LockToken<'a, L0>) -> RwLockReadGuard<'a, L1, BTreeSet<ContextRef>> {
+    CONTEXTS.read(token)
+}
+
+/// Get the global schemes list, mutable
+pub fn contexts_mut<'a>(
+    token: LockToken<'a, L0>,
+) -> RwLockWriteGuard<'a, L1, BTreeSet<ContextRef>> {
+    CONTEXTS.write(token)
+}
+
+pub fn init(token: &mut CleanLockToken) {
     let owner = None; // kmain not owned by any fd
     let mut context = Context::new(owner).expect("failed to create kmain context");
     context.sched_affinity = LogicalCpuSet::empty();
@@ -84,11 +99,9 @@ pub fn init() {
     context.running = true;
     context.cpu_id = Some(crate::cpu_id());
 
-    let context_lock = Arc::new(RwSpinlock::new(context));
+    let context_lock = Arc::new(ContextLock::new(context));
 
-    CONTEXTS
-        .write()
-        .insert(ContextRef(Arc::clone(&context_lock)));
+    contexts_mut(token.token()).insert(ContextRef(Arc::clone(&context_lock)));
 
     unsafe {
         let percpu = PercpuBlock::current();
@@ -99,35 +112,25 @@ pub fn init() {
     }
 }
 
-/// Get the global schemes list, const
-pub fn contexts() -> RwLockReadGuard<'static, BTreeSet<ContextRef>> {
-    CONTEXTS.read()
-}
-
-/// Get the global schemes list, mutable
-pub fn contexts_mut() -> RwLockWriteGuard<'static, BTreeSet<ContextRef>> {
-    CONTEXTS.write()
-}
-
-pub fn current() -> Arc<RwSpinlock<Context>> {
+pub fn current() -> Arc<ContextLock> {
     PercpuBlock::current()
         .switch_internals
         .with_context(|context| Arc::clone(context))
 }
-pub fn try_current() -> Option<Arc<RwSpinlock<Context>>> {
+pub fn try_current() -> Option<Arc<ContextLock>> {
     PercpuBlock::current()
         .switch_internals
         .try_with_context(|context| context.map(Arc::clone))
 }
-pub fn is_current(context: &Arc<RwSpinlock<Context>>) -> bool {
+pub fn is_current(context: &Arc<ContextLock>) -> bool {
     PercpuBlock::current()
         .switch_internals
         .with_context(|current| Arc::ptr_eq(context, current))
 }
 
-pub struct ContextRef(pub Arc<RwSpinlock<Context>>);
+pub struct ContextRef(pub Arc<ContextLock>);
 impl ContextRef {
-    pub fn upgrade(&self) -> Option<Arc<RwSpinlock<Context>>> {
+    pub fn upgrade(&self) -> Option<Arc<ContextLock>> {
         Some(Arc::clone(&self.0))
     }
 }
@@ -154,18 +157,17 @@ pub fn spawn(
     userspace_allowed: bool,
     owner_proc_id: Option<NonZeroUsize>,
     func: extern "C" fn(),
-) -> Result<Arc<RwSpinlock<Context>>> {
+    token: &mut CleanLockToken,
+) -> Result<Arc<ContextLock>> {
     let stack = Kstack::new()?;
 
-    let context_lock = Arc::try_new(RwSpinlock::new(Context::new(owner_proc_id)?))
+    let context_lock = Arc::try_new(ContextLock::new(Context::new(owner_proc_id)?))
         .map_err(|_| Error::new(ENOMEM))?;
 
-    CONTEXTS
-        .write()
-        .insert(ContextRef(Arc::clone(&context_lock)));
+    contexts_mut(token.token()).insert(ContextRef(Arc::clone(&context_lock)));
 
     {
-        let mut context = context_lock.write();
+        let mut context = context_lock.write(token.token());
         let _ = context.set_addr_space(Some(AddrSpaceWrapper::new()?));
         context
             .arch

src/context/signal.rs

@@ -1,10 +1,10 @@
 use core::sync::atomic::Ordering;
 
-use crate::{context, syscall::flag::SigcontrolFlags};
+use crate::{context, sync::CleanLockToken, syscall::flag::SigcontrolFlags};
 
-pub fn signal_handler() {
+pub fn signal_handler(token: &mut CleanLockToken) {
     let context_lock = context::current();
-    let mut context_guard = context_lock.write();
+    let mut context_guard = context_lock.write(token.token());
     let context = &mut *context_guard;
 
     let being_sigkilled = context.being_sigkilled;
@@ -12,7 +12,7 @@ pub fn signal_handler() {
     if being_sigkilled {
         drop(context_guard);
         drop(context_lock);
-        crate::syscall::process::exit_this_context(None);
+        crate::syscall::process::exit_this_context(None, token);
     }
 
     /*let thumbs_down = ptrace::breakpoint_callback(
@@ -72,9 +72,11 @@ pub fn signal_handler() {
     );
 }
 pub fn excp_handler(excp: syscall::Exception) {
+    let mut token = unsafe { CleanLockToken::new() };
+
     let current = context::current();
 
-    let mut context = current.write();
+    let context = current.write(token.token());
 
     let Some(eh) = context.sig.as_ref().and_then(|s| s.excp_handler) else {
         // TODO: Let procmgr print this?
@@ -87,7 +89,7 @@ pub fn excp_handler(excp: syscall::Exception) {
         drop(context);
         // TODO: Allow exceptions to be caught by tracer etc, without necessarily exiting the
         // context (closing files, dropping AddrSpace, etc)
-        crate::syscall::process::exit_this_context(Some(excp));
+        crate::syscall::process::exit_this_context(Some(excp), &mut token);
     };
     // TODO
     /*

src/context/switch.rs

@@ -9,15 +9,16 @@ use core::{
 };
 
 use alloc::sync::Arc;
-use spinning_top::{guard::ArcRwSpinlockWriteGuard, RwSpinlock};
 use syscall::PtraceFlags;
 
 use crate::{
-    context::{arch, contexts, Context},
+    context::{arch, contexts, ArcContextLockWriteGuard, Context, ContextLock},
     cpu_set::LogicalCpuId,
     cpu_stats,
     percpu::PercpuBlock,
-    ptrace, time,
+    ptrace,
+    sync::CleanLockToken,
+    time,
 };
 
 use super::ContextRef;
@@ -71,8 +72,8 @@ unsafe fn update_runnable(context: &mut Context, cpu_id: LogicalCpuId) -> Update
 }
 
 struct SwitchResultInner {
-    _prev_guard: ArcRwSpinlockWriteGuard<Context>,
-    _next_guard: ArcRwSpinlockWriteGuard<Context>,
+    _prev_guard: ArcContextLockWriteGuard,
+    _next_guard: ArcContextLockWriteGuard,
 }
 
 /// Tick function to update PIT ticks and trigger a context switch if necessary.
@@ -81,7 +82,7 @@ struct SwitchResultInner {
 /// switch if the counter reaches a set threshold (e.g., every 3 ticks).
 ///
 /// The function also calls the signal handler after switching contexts.
-pub fn tick() {
+pub fn tick(token: &mut CleanLockToken) {
     let ticks_cell = &PercpuBlock::current().switch_internals.pit_ticks;
 
     let new_ticks = ticks_cell.get() + 1;
@@ -89,8 +90,8 @@ pub fn tick() {
 
     // Trigger a context switch after every 3 ticks (approx. 6.75 ms).
     if new_ticks >= 3 {
-        switch();
-        crate::context::signal::signal_handler();
+        switch(token);
+        crate::context::signal::signal_handler(token);
     }
 }
 
@@ -136,7 +137,7 @@ pub enum SwitchResult {
 /// - `SwitchResult::Switched`: Indicates a successful switch to a new context.
 /// - `SwitchResult::AllContextsIdle`: Indicates all contexts are idle, and the CPU will switch
 ///   to an idle context.
-pub fn switch() -> SwitchResult {
+pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
     let percpu = PercpuBlock::current();
     cpu_stats::add_context_switch();
     percpu
@@ -161,11 +162,12 @@ pub fn switch() -> SwitchResult {
 
     let mut switch_context_opt = None;
     {
-        let contexts = contexts();
+        let contexts = contexts(token.token());
 
         // Lock the previous context.
         let prev_context_lock = crate::context::current();
-        let prev_context_guard = prev_context_lock.write_arc();
+        // We are careful not to lock this context twice
+        let prev_context_guard = unsafe { prev_context_lock.write_arc() };
 
         let idle_context = percpu.switch_internals.idle_context();
 
@@ -201,7 +203,8 @@ pub fn switch() -> SwitchResult {
 
             {
                 // Lock next context
-                let mut next_context_guard = next_context_lock.write_arc();
+                // We are careful not to lock this context twice
+                let mut next_context_guard = unsafe { next_context_lock.write_arc() };
 
                 // Check if the context is runnable and can be switched to.
                 if let UpdateResult::CanSwitch =
@@ -234,7 +237,7 @@ pub fn switch() -> SwitchResult {
             let percpu = PercpuBlock::current();
             unsafe {
                 percpu.switch_internals.set_current_context(Arc::clone(
-                    ArcRwSpinlockWriteGuard::rwlock(&next_context_guard),
+                    ArcContextLockWriteGuard::rwlock(&next_context_guard),
                 ));
             }
 
@@ -319,10 +322,10 @@ pub struct ContextSwitchPercpu {
     switch_result: Cell<Option<SwitchResultInner>>,
     pit_ticks: Cell<usize>,
 
-    current_ctxt: RefCell<Option<Arc<RwSpinlock<Context>>>>,
+    current_ctxt: RefCell<Option<Arc<ContextLock>>>,
 
     /// The idle process.
-    idle_ctxt: RefCell<Option<Arc<RwSpinlock<Context>>>>,
+    idle_ctxt: RefCell<Option<Arc<ContextLock>>>,
 
     pub(crate) being_sigkilled: Cell<bool>,
 }
@@ -345,7 +348,7 @@ impl ContextSwitchPercpu {
     ///
     /// # Returns
     /// The result of applying `f` to the current context.
-    pub fn with_context<T>(&self, f: impl FnOnce(&Arc<RwSpinlock<Context>>) -> T) -> T {
+    pub fn with_context<T>(&self, f: impl FnOnce(&Arc<ContextLock>) -> T) -> T {
         f(self
             .current_ctxt
             .borrow()
@@ -360,7 +363,7 @@ impl ContextSwitchPercpu {
     ///
     /// # Returns
     /// The result of applying `f` to the current context if any.
-    pub fn try_with_context<T>(&self, f: impl FnOnce(Option<&Arc<RwSpinlock<Context>>>) -> T) -> T {
+    pub fn try_with_context<T>(&self, f: impl FnOnce(Option<&Arc<ContextLock>>) -> T) -> T {
         f(self.current_ctxt.borrow().as_ref())
     }
 
@@ -371,7 +374,7 @@ impl ContextSwitchPercpu {
     ///
     /// # Parameters
     /// - `new`: The new context to be set as the current context.
-    pub unsafe fn set_current_context(&self, new: Arc<RwSpinlock<Context>>) {
+    pub unsafe fn set_current_context(&self, new: Arc<ContextLock>) {
         *self.current_ctxt.borrow_mut() = Some(new);
     }
 
@@ -382,7 +385,7 @@ impl ContextSwitchPercpu {
     ///
     /// # Parameters
     /// - `new`: The new context to be set as the idle context.
-    pub unsafe fn set_idle_context(&self, new: Arc<RwSpinlock<Context>>) {
+    pub unsafe fn set_idle_context(&self, new: Arc<ContextLock>) {
         *self.idle_ctxt.borrow_mut() = Some(new);
     }
 
@@ -390,7 +393,7 @@ impl ContextSwitchPercpu {
     ///
     /// # Returns
     /// A reference to the idle context.
-    pub fn idle_context(&self) -> Arc<RwSpinlock<Context>> {
+    pub fn idle_context(&self) -> Arc<ContextLock> {
         Arc::clone(
             self.idle_ctxt
                 .borrow()

src/context/timeout.rs

@@ -1,9 +1,10 @@
 use alloc::collections::VecDeque;
-use spin::{Mutex, MutexGuard, Once};
+use spin::Once;
 
 use crate::{
     event,
     scheme::SchemeId,
+    sync::{CleanLockToken, LockToken, Mutex, MutexGuard, L0, L1},
     syscall::{
         data::TimeSpec,
         flag::{CLOCK_MONOTONIC, CLOCK_REALTIME, EVENT_READ},
@@ -21,20 +22,26 @@ struct Timeout {
 
 type Registry = VecDeque<Timeout>;
 
-static REGISTRY: Once<Mutex<Registry>> = Once::new();
+static REGISTRY: Once<Mutex<L1, Registry>> = Once::new();
 
 /// Initialize registry, called if needed
-fn init_registry() -> Mutex<Registry> {
+fn init_registry() -> Mutex<L1, Registry> {
     Mutex::new(Registry::new())
 }
 
 /// Get the global timeouts list
-fn registry() -> MutexGuard<'static, Registry> {
-    REGISTRY.call_once(init_registry).lock()
+fn registry<'a>(token: LockToken<'a, L0>) -> MutexGuard<'a, L1, Registry> {
+    REGISTRY.call_once(init_registry).lock(token)
 }
 
-pub fn register(scheme_id: SchemeId, event_id: usize, clock: usize, time: TimeSpec) {
-    let mut registry = registry();
+pub fn register(
+    scheme_id: SchemeId,
+    event_id: usize,
+    clock: usize,
+    time: TimeSpec,
+    token: &mut CleanLockToken,
+) {
+    let mut registry = registry(token.token());
     registry.push_back(Timeout {
         scheme_id,
         event_id,
@@ -43,34 +50,38 @@ pub fn register(scheme_id: SchemeId, event_id: usize, clock: usize, time: TimeSp
     });
 }
 
-pub fn trigger() {
-    let mut registry = registry();
-
+pub fn trigger(token: &mut CleanLockToken) {
     let mono = time::monotonic();
     let real = time::realtime();
 
     let mut i = 0;
-    while i < registry.len() {
-        let trigger = match registry[i].clock {
-            CLOCK_MONOTONIC => {
-                let time = registry[i].time;
-                mono >= time
-            }
-            CLOCK_REALTIME => {
-                let time = registry[i].time;
-                real >= time
-            }
-            clock => {
-                println!("timeout::trigger: unknown clock {}", clock);
-                true
-            }
-        };
+    loop {
+        let mut registry = registry(token.token());
+        let timeout = if i < registry.len() {
+            let trigger = match registry[i].clock {
+                CLOCK_MONOTONIC => {
+                    let time = registry[i].time;
+                    mono >= time
+                }
+                CLOCK_REALTIME => {
+                    let time = registry[i].time;
+                    real >= time
+                }
+                clock => {
+                    println!("timeout::trigger: unknown clock {}", clock);
+                    true
+                }
+            };
 
-        if trigger {
-            let timeout = registry.remove(i).unwrap();
-            event::trigger(timeout.scheme_id, timeout.event_id, EVENT_READ);
+            if trigger {
+                registry.remove(i).unwrap()
+            } else {
+                i += 1;
+                continue;
+            }
         } else {
-            i += 1;
-        }
+            break;
+        };
+        event::trigger(timeout.scheme_id, timeout.event_id, EVENT_READ);
     }
 }