Remove linear scan

src/arch/x86_shared/interrupt/ipi.rs

@@ -1,8 +1,13 @@
 use crate::{
-    arch::device::local_apic::the_local_apic, context, percpu::PercpuBlock, sync::CleanLockToken,
+    arch::device::local_apic::the_local_apic,
+    context::{self, switch::drain_ipi_context_wakeups},
+    percpu::PercpuBlock,
+    sync::CleanLockToken,
 };
 
 interrupt!(wakeup, || {
+    let mut token = unsafe { CleanLockToken::new() };
+    drain_ipi_context_wakeups(&mut token);
     unsafe { the_local_apic().eoi() };
 });
 

src/context/mod.rs

@@ -84,11 +84,7 @@ static CONTEXTS: RwLock<L2, BTreeSet<ContextRef>> = RwLock::new(BTreeSet::new())
 // Actual context store for the scheduler
 static RUN_CONTEXTS: Mutex<L1, RunContextData> = Mutex::new(RunContextData::new());
 
-// Context that has been pushed out from RUN_CONTEXTS after being idle
-static IDLE_CONTEXTS: Mutex<L2, VecDeque<WeakContextRef>> = Mutex::new(VecDeque::new());
-
 pub struct RunContextData {
-    // queue: VecDeque<WeakContextRef>,
     queue: BTreeMap<(u64, Reverse<u64>, u32), (u64, u64, WeakContextRef)>, // ((vd, rem_slice, ctxt_id), (vtime, weight, context))
     timers: BTreeSet<(u128, WeakContextRef)>,                              // (wake, context)
     count: usize,
@@ -108,10 +104,6 @@ impl RunContextData {
             min_vtime: 0,
         }
     }
-    pub fn update_count(&mut self) -> usize {
-        self.count = self.queue.len();
-        self.count
-    }
 }
 
 /// Get the global schemes list, const
@@ -124,16 +116,6 @@ pub fn contexts_mut(token: LockToken<'_, L1>) -> RwLockWriteGuard<'_, L2, BTreeS
     CONTEXTS.write(token)
 }
 
-pub fn idle_contexts(token: LockToken<'_, L1>) -> MutexGuard<'_, L2, VecDeque<WeakContextRef>> {
-    IDLE_CONTEXTS.lock(token)
-}
-
-pub fn idle_contexts_try(
-    token: LockToken<'_, L1>,
-) -> Option<MutexGuard<'_, L2, VecDeque<WeakContextRef>>> {
-    IDLE_CONTEXTS.try_lock(token)
-}
-
 pub fn run_contexts(token: LockToken<'_, L0>) -> MutexGuard<'_, L1, RunContextData> {
     RUN_CONTEXTS.lock(token)
 }
@@ -146,42 +128,55 @@ pub fn unblock_context(context_lock: &Arc<ContextLock>, token: &mut LockToken<'_
     let cpu_id = {
         let mut guard = context_lock.write(token.token());
         if !guard.unblock_no_ipi() {
-            if guard.status.is_runnable() {
+            let guard_runnable = guard.status.is_runnable();
+            let guard_cpu_id = guard.cpu_id;
+            if guard_runnable {
                 // already set to runnable externally
-                wakeup_context(context_lock, guard.cpu_id);
+                drop(guard);
+                wakeup_context(context_lock, guard_cpu_id, token);
             }
             return false;
         }
         guard.cpu_id
     };
 
-    wakeup_context(context_lock, cpu_id);
+    wakeup_context(context_lock, cpu_id, token);
 
     true
 }
 
-pub fn wakeup_context(context_lock: &Arc<ContextLock>, cpu_id: Option<LogicalCpuId>) {
+pub fn wakeup_context(
+    context_lock: &Arc<ContextLock>,
+    cpu_id: Option<LogicalCpuId>,
+    token: &mut LockToken<'_, L3>,
+) {
     let weak = WeakContextRef(Arc::downgrade(context_lock));
     let curr_cpu = crate::cpu_id();
 
-    if let Some(target) = cpu_id
-        && target != curr_cpu
-    {
-        if let Some(percpu) = unsafe {
-            ALL_PERCPU_BLOCKS[target.get() as usize]
-                .load(Ordering::Acquire)
-                .as_ref()
-        } {
-            percpu.switch_internals.wakeup_list.lock().push(weak);
-            ipi(IpiKind::Wakeup, IpiTarget::Other);
-            return;
+    if let Some(target) = cpu_id {
+        if target != curr_cpu {
+            if let Some(percpu) = unsafe {
+                ALL_PERCPU_BLOCKS[target.get() as usize]
+                    .load(Ordering::Acquire)
+                    .as_ref()
+            } {
+                // cross core wakeup
+                percpu
+                    .switch_internals
+                    .ipi_context_wakeup_list
+                    .lock(token.token())
+                    .push(weak);
+                ipi(IpiKind::Wakeup, IpiTarget::Other);
+                return;
+            }
         }
     }
 
+    // local wakeup
     PercpuBlock::current()
         .switch_internals
-        .wakeup_list
-        .lock()
+        .local_wakeup_list
+        .borrow_mut()
         .push(weak);
 }
 
@@ -401,7 +396,7 @@ impl Drop for PreemptGuardL2<'_> {
 pub fn get_contexts_stats(token: &mut CleanLockToken) -> (usize, usize, usize) {
     let alive = contexts(token.downgrade()).len();
     let running = run_contexts(token.token()).count;
-    let blocked = idle_contexts(token.downgrade()).len();
+    let blocked = alive.saturating_sub(running);
 
     (alive, running, blocked)
 }

src/context/switch.rs

@@ -4,16 +4,18 @@
 
 use crate::{
     context::{
-        self, arch, idle_contexts, idle_contexts_try, memory::AddrSpaceSwitchReadGuard,
-        run_contexts, run_contexts_try, wakeup_context, ArcContextLockWriteGuard, Context,
-        ContextLock, WeakContextRef,
+        self, arch, memory::AddrSpaceSwitchReadGuard, run_contexts, run_contexts_try,
+        wakeup_context, ArcContextLockWriteGuard, Context, ContextLock, WeakContextRef,
     },
     cpu_set::LogicalCpuId,
     cpu_stats::{self, CpuState},
-    percpu::PercpuBlock,
-    sync::{ArcRwLockWriteGuard, CleanLockToken, L4},
+    percpu::{self, PercpuBlock},
+    sync::{ArcRwLockWriteGuard, CleanLockToken, Mutex, L4},
+};
+use alloc::{
+    sync::{Arc, Weak},
+    vec::Vec,
 };
-use alloc::{sync::Arc, vec::Vec};
 use core::{
     cell::{Cell, RefCell},
     cmp::Reverse,
@@ -138,6 +140,22 @@ pub unsafe extern "C" fn switch_finish_hook() {
     }
 }
 
+/// Drains the cross_cpu_wakeup_list into local_wakeup_list.
+/// This is called from the ipi handler.
+pub fn drain_ipi_context_wakeups(token: &mut CleanLockToken) {
+    let percpu = PercpuBlock::current();
+    let mut cross_cpu_wake = percpu
+        .switch_internals
+        .ipi_context_wakeup_list
+        .lock(token.token());
+    if cross_cpu_wake.is_empty() {
+        return;
+    }
+
+    let mut local_wake = percpu.switch_internals.local_wakeup_list.borrow_mut();
+    local_wake.extend(cross_cpu_wake.drain(..));
+}
+
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub enum SwitchResult {
     Switched,
@@ -187,104 +205,96 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
         return SwitchResult::Switched;
     }
 
-    // Alarm (previously in update_runnable)
-    let mut wakeups: SmallVec<[(Option<u128>, WeakContextRef); 16]> = wakeup_contexts(token)
-        .into_iter()
-        .map(|ctxt| (None, ctxt))
-        .collect();
-    let mut push_idle: SmallVec<[WeakContextRef; 16]> = SmallVec::new();
-
-    // These timers coukd have expired
-    let mut timers: SmallVec<[(u128, WeakContextRef); 16]> = SmallVec::new();
-    if let Some(mut run_contexts) = run_contexts_try(token.token()) {
-        // Pop Timers
-        while let Some((wake, _)) = run_contexts.timers.first() {
-            if *wake > switch_time {
-                break;
-            }
-
-            if let Some(entry) = run_contexts.timers.pop_first() {
-                timers.push(entry);
-            }
-        }
-    }
-
-    for (wake, context_ref) in timers {
-        let Some(context_lock) = context_ref.upgrade() else {
-            continue;
-        };
-
-        let guard = context_lock.read(token.token());
-        if guard.status.is_soft_blocked() && guard.wake == Some(wake) {
-            wakeups.push((Some(wake), context_ref));
-        }
-    }
-
-    // Drain from percpu
     {
-        if let Some(mut percpu_wake) = percpu.switch_internals.wakeup_list.try_lock() {
-            wakeups.extend(percpu_wake.drain(..).map(|ctx| (None, ctx)));
-        }
-    }
+        // Alarm (previously in update_runnable)
+        let mut wakeups = percpu.switch_internals.tmp_wakeups.borrow_mut();
 
-    if wakeups.len() > 0 {
-        let mut run_contexts = run_contexts(token.token());
-        for (wake_opt, context_ref) in wakeups {
+        // These timers coukd have expired
+        let mut timers = percpu.switch_internals.tmp_timers.borrow_mut();
+        timers.clear();
+        if let Some(mut run_contexts) = run_contexts_try(token.token()) {
+            let split_key = (switch_time.saturating_add(1), WeakContextRef(Weak::new()));
+
+            timers.extend(run_contexts.timers.extract_if(..split_key, |_| true));
+        }
+
+        timers.retain(|(wake, context_ref)| {
             let Some(context_lock) = context_ref.upgrade() else {
-                continue;
+                return false;
             };
 
-            let Some(mut guard) = (unsafe { context_lock.try_write_arc() }) else {
-                if let Some(wake) = wake_opt {
-                    run_contexts.timers.insert((wake, context_ref));
-                } else {
-                    push_idle.push(context_ref);
-                }
-                continue;
-            };
-
-            if let Some(wake) = wake_opt {
-                if guard.status.is_soft_blocked() && guard.wake == Some(wake) {
-                    guard.wake = None;
-                    guard.unblock_no_ipi();
+            if let Some(guard) = context_lock.try_read(token.token()) {
+                if guard.status.is_soft_blocked() && guard.wake == Some(*wake) {
+                    wakeups.push((Some(*wake), context_ref.clone()));
                 }
+                false
+            } else {
+                true
             }
+        });
 
-            if guard.running || !guard.status.is_runnable() {
-                continue;
-            }
-
-            let new_vtime = guard.vtime.max(run_contexts.v);
-            guard.vtime = new_vtime;
-
-            let weight = SCHED_PRIO_TO_WEIGHT[guard.prio] as u64;
-            let scaled_slice = (BASE_SLICE_TICKS as u128 * SCALE) / weight as u128;
-
-            if !guard.is_active {
-                guard.is_active = true;
-                run_contexts.total_weight += weight;
-            }
-
-            if let Some(old_key) = guard.queue_key.take() {
-                run_contexts.queue.remove(&old_key);
-            }
-
-            guard.vd = new_vtime + scaled_slice as u64;
-            guard.rem_slice = BASE_SLICE_TICKS * SCALE as u64;
-            let key = (guard.vd, Reverse(guard.rem_slice), guard.debug_id);
-            guard.queue_key = Some(key);
-            drop(guard);
-
-            run_contexts
-                .queue
-                .insert(key, (new_vtime, weight, context_ref));
+        // Drain from percpu
+        {
+            let mut local_wake = percpu.switch_internals.local_wakeup_list.borrow_mut();
+            wakeups.extend(local_wake.drain(..).map(|ctx| (None, ctx)));
         }
-    }
 
-    {
-        let mut idle_list = idle_contexts(token.downgrade());
-        for context_ref in push_idle {
-            idle_list.push_back(context_ref);
+        if wakeups.len() > 0 {
+            let mut run_contexts = run_contexts(token.token());
+            for (wake_opt, context_ref) in wakeups.drain(..) {
+                let Some(context_lock) = context_ref.upgrade() else {
+                    continue;
+                };
+
+                let Some(mut guard) = (unsafe { context_lock.try_write_arc() }) else {
+                    if let Some(wake) = wake_opt {
+                        run_contexts.timers.insert((wake, context_ref));
+                    } else {
+                        percpu
+                            .switch_internals
+                            .local_wakeup_list
+                            .borrow_mut()
+                            .push(context_ref);
+                    }
+                    continue;
+                };
+
+                if let Some(wake) = wake_opt {
+                    if guard.status.is_soft_blocked() && guard.wake == Some(wake) {
+                        guard.wake = None;
+                        guard.unblock_no_ipi();
+                    }
+                }
+
+                if guard.running || !guard.status.is_runnable() {
+                    continue;
+                }
+
+                let new_vtime = guard.vtime.max(run_contexts.v);
+                guard.vtime = new_vtime;
+
+                let weight = SCHED_PRIO_TO_WEIGHT[guard.prio] as u64;
+                let scaled_slice = (BASE_SLICE_TICKS as u128 * SCALE) / weight as u128;
+
+                if !guard.is_active {
+                    guard.is_active = true;
+                    run_contexts.total_weight += weight;
+                }
+
+                if let Some(old_key) = guard.queue_key.take() {
+                    run_contexts.queue.remove(&old_key);
+                }
+
+                guard.vd = new_vtime + scaled_slice as u64;
+                guard.rem_slice = BASE_SLICE_TICKS * SCALE as u64;
+                let key = (guard.vd, Reverse(guard.rem_slice), guard.debug_id);
+                guard.queue_key = Some(key);
+                drop(guard);
+
+                run_contexts
+                    .queue
+                    .insert(key, (new_vtime, weight, context_ref));
+            }
         }
     }
 
@@ -432,48 +442,6 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
     }
 }
 
-fn wakeup_contexts(token: &mut CleanLockToken) -> SmallVec<[WeakContextRef; 16]> {
-    // TODO: Optimise this somehow
-    let mut wakeups = SmallVec::new();
-    let current_context = context::current();
-    let Some(idle_contexts) = idle_contexts_try(token.downgrade()) else {
-        // other cpus may spawning or killing contexts so let's skip wakeups to avoid contention
-        return wakeups;
-    };
-    let (mut idle_contexts, mut token) = idle_contexts.into_split();
-    let len = idle_contexts.len();
-    for _ in 0..len {
-        let Some(context_ref) = idle_contexts.pop_front() else {
-            break;
-        };
-        let Some(context) = context_ref.upgrade() else {
-            continue;
-        };
-        if Arc::ptr_eq(&context, &current_context) {
-            idle_contexts.push_back(context_ref);
-            continue;
-        }
-        let Some(guard) = context.try_read(token.token()) else {
-            idle_contexts.push_back(context_ref);
-            continue;
-        };
-        if guard.status.is_dead() {
-            // TODO: who hold this dead context?
-            continue;
-        }
-
-        if guard.status.is_runnable() && !guard.running {
-            drop(guard);
-            wakeups.push(context_ref);
-            continue;
-        }
-
-        drop(guard);
-        idle_contexts.push_back(context_ref);
-    }
-    wakeups
-}
-
 /// This is the scheduler function which currently utilises EEVDF Scheduler
 fn select_next_context(
     token: &mut CleanLockToken,
@@ -593,8 +561,6 @@ fn select_next_context(
 
             contexts_to_remove.push((*vd, *rem_slice, *ctxt_id));
             drop(guard);
-            // TODO: Reenqueue should be handled by unblock
-            idle_contexts(token.token()).push_back(context_ref.clone());
             continue;
         }
 
@@ -703,9 +669,6 @@ fn select_next_context(
                         WeakContextRef(Arc::downgrade(&prev_context_lock)),
                     ),
                 );
-            } else if !is_idle && !is_timer {
-                idle_contexts(token.token())
-                    .push_back(WeakContextRef(Arc::downgrade(&prev_context_lock)));
             }
 
             return Some((chosen_guard, addr_space));
@@ -713,11 +676,6 @@ fn select_next_context(
             return None;
         }
     } else {
-        if !is_idle && !is_timer {
-            idle_contexts(token.token())
-                .push_back(WeakContextRef(Arc::downgrade(&prev_context_lock)));
-        }
-
         let prev_is_dead = !is_idle && !prev_context_guard.status.is_runnable();
         if (!was_idle || prev_is_dead) && !is_idle {
             return Some(unsafe { (idle_context.write_arc(), None) });
@@ -747,7 +705,10 @@ pub struct ContextSwitchPercpu {
     pub(crate) being_sigkilled: Cell<bool>,
 
     // wakeups
-    pub(crate) wakeup_list: SpinMutex<Vec<WeakContextRef>>,
+    pub(crate) ipi_context_wakeup_list: Mutex<L4, Vec<WeakContextRef>>,
+    pub(crate) local_wakeup_list: RefCell<Vec<WeakContextRef>>,
+    tmp_wakeups: RefCell<Vec<(Option<u128>, WeakContextRef)>>,
+    tmp_timers: RefCell<Vec<(u128, WeakContextRef)>>,
 }
 
 impl ContextSwitchPercpu {
@@ -759,7 +720,10 @@ impl ContextSwitchPercpu {
             current_ctxt: RefCell::new(None),
             idle_ctxt: RefCell::new(None),
             being_sigkilled: Cell::new(false),
-            wakeup_list: SpinMutex::new(Vec::new()),
+            ipi_context_wakeup_list: Mutex::new(Vec::new()),
+            local_wakeup_list: RefCell::new(Vec::new()),
+            tmp_wakeups: RefCell::new(Vec::new()),
+            tmp_timers: RefCell::new(Vec::new()),
 
             #[cfg(feature = "profiling")]
             current_dbg_id: core::sync::atomic::AtomicU32::new(!0),

src/scheme/proc.rs

@@ -84,17 +84,25 @@ fn try_stop_context<T>(
         running = context_ref.read(token.token()).running;
     }
 
-    let mut context = context_ref.write(token.token());
+    let mut guard = context_ref.write(token.token());
     assert!(
-        !context.running,
+        !guard.running,
         "process can't have been restarted, we stopped it!"
     );
 
-    let (context, token) = context.token_split();
-    let ret = callback(context, token);
+    let (context, l4_token) = guard.token_split();
+    let ret = callback(context, l4_token);
 
     context.status = prev_status;
 
+    let cpu_id = context.cpu_id;
+
+    let wake = context.status.is_runnable();
+    drop(guard);
+    if wake {
+        wakeup_context(&context_ref, cpu_id, &mut token.downgrade());
+    }
+
     ret
 }
 
@@ -1147,7 +1155,7 @@ impl ContextHandle {
                     }
                 };
 
-                wakeup_context(&context, cpu_id);
+                wakeup_context(&context, cpu_id, &mut token.downgrade());
                 Ok(buf.len())
             }
             ContextHandle::Filetable { .. } | ContextHandle::NewFiletable { .. } => {
@@ -1292,7 +1300,9 @@ impl ContextHandle {
                         } = guard.status
                         {
                             guard.status = Status::Runnable;
-                            wakeup_context(&context, guard.cpu_id);
+                            let cpu_id = guard.cpu_id;
+                            drop(guard);
+                            wakeup_context(&context, cpu_id, &mut token.downgrade());
                         }
                         Ok(size_of::<usize>())
                     }
@@ -1336,7 +1346,7 @@ impl ContextHandle {
                                 ctxt.being_sigkilled = true;
                                 ctxt.cpu_id
                             };
-                            wakeup_context(&context, cpu_id);
+                            wakeup_context(&context, cpu_id, &mut token.downgrade());
                             Ok(size_of::<usize>())
                         }
                     }

src/scheme/user.rs

@@ -878,6 +878,7 @@ impl UserInner {
                 let context_lock = context.upgrade().ok_or(Error::new(ESRCH))?;
 
                 let mut lock_token = token.token();
+                let mut wake = (false, None);
                 let (frame, _) = AddrSpace::current()?
                     .acquire_read(lock_token.downgrade())
                     .table
@@ -892,10 +893,13 @@ impl UserInner {
                     } = context.status
                     {
                         context.status = Status::Runnable;
-                        wakeup_context(&context_lock, context.cpu_id);
+                        wake = (true, context.cpu_id);
                     }
                     context.fmap_ret = Some(Frame::containing(frame));
                 }
+                if wake.0 {
+                    wakeup_context(&context_lock, wake.1, &mut token.downgrade());
+                }
             }
             ParsedCqe::TriggerFevent { number, flags } => {
                 event::trigger(self.scheme_id, number, flags, token)

src/startup/mod.rs

@@ -183,7 +183,10 @@ pub(crate) fn kmain(bootstrap: Bootstrap) -> ! {
             context.euid = 0;
             context.egid = 0;
 
-            wakeup_context(&context_lock, context.cpu_id);
+            let cpu_id = context.cpu_id;
+            drop(context);
+
+            wakeup_context(&context_lock, cpu_id, &mut token.downgrade());
         }
         Err(err) => {
             panic!("failed to spawn userspace_init: {:?}", err);