From bb03b896dfed6813a6bd18c86c5097f0bd308bef Mon Sep 17 00:00:00 2001 From: Akshit Gaur Date: Tue, 28 Jul 2026 15:20:30 +0000 Subject: [PATCH] Remove linear scan --- src/arch/x86_shared/interrupt/ipi.rs | 7 +- src/context/mod.rs | 67 ++++--- src/context/switch.rs | 260 ++++++++++++--------------- src/scheme/proc.rs | 24 ++- src/scheme/user.rs | 6 +- src/startup/mod.rs | 5 +- 6 files changed, 175 insertions(+), 194 deletions(-) diff --git a/src/arch/x86_shared/interrupt/ipi.rs b/src/arch/x86_shared/interrupt/ipi.rs index 20e9d436..3c6f9d24 100644 --- a/src/arch/x86_shared/interrupt/ipi.rs +++ b/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() }; }); diff --git a/src/context/mod.rs b/src/context/mod.rs index 63f9fbde..734688d3 100644 --- a/src/context/mod.rs +++ b/src/context/mod.rs @@ -84,11 +84,7 @@ static CONTEXTS: RwLock> = RwLock::new(BTreeSet::new()) // Actual context store for the scheduler static RUN_CONTEXTS: Mutex = Mutex::new(RunContextData::new()); -// Context that has been pushed out from RUN_CONTEXTS after being idle -static IDLE_CONTEXTS: Mutex> = Mutex::new(VecDeque::new()); - pub struct RunContextData { - // queue: VecDeque, queue: BTreeMap<(u64, Reverse, 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> { - IDLE_CONTEXTS.lock(token) -} - -pub fn idle_contexts_try( - token: LockToken<'_, L1>, -) -> Option>> { - 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, 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, cpu_id: Option) { +pub fn wakeup_context( + context_lock: &Arc, + cpu_id: Option, + 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) } diff --git a/src/context/switch.rs b/src/context/switch.rs index a2b07e80..f54211d5 100644 --- a/src/context/switch.rs +++ b/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, 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, ¤t_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, // wakeups - pub(crate) wakeup_list: SpinMutex>, + pub(crate) ipi_context_wakeup_list: Mutex>, + pub(crate) local_wakeup_list: RefCell>, + tmp_wakeups: RefCell, WeakContextRef)>>, + tmp_timers: RefCell>, } 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), diff --git a/src/scheme/proc.rs b/src/scheme/proc.rs index 20253dff..9c62e63c 100644 --- a/src/scheme/proc.rs +++ b/src/scheme/proc.rs @@ -84,17 +84,25 @@ fn try_stop_context( 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::()) } @@ -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::()) } } diff --git a/src/scheme/user.rs b/src/scheme/user.rs index c2c0e6f5..c5fdb4a8 100644 --- a/src/scheme/user.rs +++ b/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) diff --git a/src/startup/mod.rs b/src/startup/mod.rs index a4251c56..f571852a 100644 --- a/src/startup/mod.rs +++ b/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);