Store context to per cpu

This commit is contained in:
Wildan M
2026-03-31 15:29:47 +07:00
parent efa6b00189
commit 6e0143cf0d
5 changed files with 100 additions and 72 deletions
+23 -10
View File
@@ -11,8 +11,8 @@ use crate::{
paging::{RmmA, RmmArch, TableKind},
percpu::PercpuBlock,
sync::{
ArcRwLockWriteGuard, CleanLockToken, LockToken, RwLock, RwLockReadGuard, RwLockWriteGuard,
L0, L1, L2, L4,
ArcRwLockWriteGuard, CleanLockToken, LockToken, Mutex, MutexGuard, RwLock, RwLockReadGuard,
RwLockWriteGuard, L0, L1, L2, L4,
},
syscall::error::Result,
};
@@ -69,16 +69,30 @@ 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<L1, BTreeSet<ContextRef>> = RwLock::new(BTreeSet::new());
static CONTEXTS: Mutex<L2, BTreeSet<ContextRef>> = Mutex::new(BTreeSet::new());
/// Get the global schemes list, const
pub fn contexts(token: LockToken<'_, L0>) -> RwLockReadGuard<'_, L1, BTreeSet<ContextRef>> {
CONTEXTS.read(token)
/// Try to get the global free contexts
pub fn free_contexts_try(
token: LockToken<'_, L1>,
) -> Option<MutexGuard<'_, L2, BTreeSet<ContextRef>>> {
CONTEXTS.try_lock(token)
}
/// Get the global schemes list, mutable
/// Get the global free contexts
pub fn free_contexts(token: LockToken<'_, L1>) -> MutexGuard<'_, L2, BTreeSet<ContextRef>> {
CONTEXTS.lock(token)
}
/// Get per cpu contexts, const
pub fn contexts(token: LockToken<'_, L0>) -> RwLockReadGuard<'_, L1, BTreeSet<ContextRef>> {
let percpu = PercpuBlock::current();
percpu.contexts.read(token)
}
/// Get per cpu contexts, mutable
pub fn contexts_mut(token: LockToken<'_, L0>) -> RwLockWriteGuard<'_, L1, BTreeSet<ContextRef>> {
CONTEXTS.write(token)
let percpu = PercpuBlock::current();
percpu.contexts.write(token)
}
pub fn init(token: &mut CleanLockToken) {
@@ -106,7 +120,6 @@ pub fn init(token: &mut CleanLockToken) {
percpu
.switch_internals
.set_current_context(Arc::clone(&context_lock));
percpu.switch_internals.set_idle_context(context_lock);
}
}
@@ -172,7 +185,7 @@ pub fn spawn(
let context_lock = Arc::new(ContextLock::new(context));
let context_ref = ContextRef(Arc::clone(&context_lock));
contexts_mut(token.token()).insert(context_ref);
free_contexts(token.token().downgrade()).insert(context_ref);
Ok(context_lock)
}
+38 -60
View File
@@ -12,7 +12,10 @@ use alloc::sync::Arc;
use syscall::PtraceFlags;
use crate::{
context::{arch, contexts, ArcContextLockWriteGuard, Context, ContextLock},
context::{
arch, contexts, contexts_mut, free_contexts_try, ArcContextLockWriteGuard,
Context, ContextLock,
},
cpu_set::LogicalCpuId,
cpu_stats,
percpu::PercpuBlock,
@@ -159,30 +162,23 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
let cpu_id = crate::cpu_id();
// Lock the previous context.
let prev_context_lock = crate::context::current();
// We are careful not to lock this context twice
let mut prev_context_guard = unsafe { prev_context_lock.write_arc() };
if !prev_context_guard.is_preemptable() {
// Unset global lock
arch::CONTEXT_SWITCH_LOCK.store(false, Ordering::SeqCst);
// Pretend to have finished switching, so CPU is not idled
return SwitchResult::Switched;
}
let mut switch_context_opt = None;
{
let contexts = contexts(token.token());
// Lock the previous context.
let prev_context_lock = crate::context::current();
// We are careful not to lock this context twice
let prev_context_guard = unsafe { prev_context_lock.write_arc() };
if !prev_context_guard.is_preemptable() {
// Unset global lock
arch::CONTEXT_SWITCH_LOCK.store(false, Ordering::SeqCst);
// Pretend to have finished switching, so CPU is not idled
return SwitchResult::Switched;
}
let idle_context = percpu.switch_internals.idle_context();
// Stateful flag used to skip the idle process the first time it shows up.
// After that, this flag is set to `false` so the idle process can be
// picked up.
let mut skip_idle = true;
// Attempt to locate the next context to switch to.
for next_context_lock in contexts
// Include all contexts with IDs greater than the current...
@@ -196,18 +192,9 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
Bound::Excluded(ContextRef(Arc::clone(&prev_context_lock))),
)))
.filter_map(ContextRef::upgrade)
// ... and the idle context...
.chain(Some(Arc::clone(&idle_context)))
// ... but not the current context (note the `Bound::Excluded`),
// which is already locked.
{
if Arc::ptr_eq(&next_context_lock, &idle_context) && skip_idle {
// Skip idle process the first time it shows up, but allow it
// to be picked up again the next time.
skip_idle = false;
continue;
}
{
// Lock next context
// We are careful not to lock this context twice
@@ -219,13 +206,32 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
{
// Store locks for previous and next context and break out from loop
// for the switch
switch_context_opt = Some((prev_context_guard, next_context_guard));
switch_context_opt = Some(next_context_guard);
break;
}
}
}
};
if switch_context_opt.is_none() {
let mut this_contexts = contexts_mut(token.token());
let (this_contexts, mut token) = this_contexts.token_split();
if let Some(mut free_contexts) = free_contexts_try(token.token()) {
if let Some(context) = free_contexts.pop_last() {
// Check if we can run this free context immediately
if let Some(next_context) = context.upgrade() {
let mut next_context_guard = unsafe { next_context.write_arc() };
if let UpdateResult::CanSwitch =
unsafe { update_runnable(&mut next_context_guard, cpu_id, switch_time) }
{
switch_context_opt = Some(next_context_guard);
}
}
this_contexts.insert(context);
}
}
}
// Update per-cpu times
let percpu_nanos = switch_time.saturating_sub(percpu.switch_internals.switch_time.get()) as u64;
let percpu_ms = percpu_nanos / 1_000_000;
@@ -234,7 +240,7 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
// Switch process states, TSS stack pointer, and store new context ID
match switch_context_opt {
Some((mut prev_context_guard, mut next_context_guard)) => {
Some(mut next_context_guard) => {
// Update context states and prepare for the switch.
let prev_context = &mut *prev_context_guard;
let next_context = &mut *next_context_guard;
@@ -339,9 +345,6 @@ pub struct ContextSwitchPercpu {
current_ctxt: RefCell<Option<Arc<ContextLock>>>,
/// The idle process.
idle_ctxt: RefCell<Option<Arc<ContextLock>>>,
pub(crate) being_sigkilled: Cell<bool>,
}
@@ -352,7 +355,6 @@ impl ContextSwitchPercpu {
switch_time: Cell::new(0),
pit_ticks: Cell::new(0),
current_ctxt: RefCell::new(None),
idle_ctxt: RefCell::new(None),
being_sigkilled: Cell::new(false),
}
}
@@ -393,28 +395,4 @@ impl ContextSwitchPercpu {
pub unsafe fn set_current_context(&self, new: Arc<ContextLock>) {
*self.current_ctxt.borrow_mut() = Some(new);
}
/// Sets the idle context to a new value.
///
/// # Safety
/// This function is unsafe as it modifies the idle context state directly.
///
/// # Parameters
/// - `new`: The new context to be set as the idle context.
pub unsafe fn set_idle_context(&self, new: Arc<ContextLock>) {
*self.idle_ctxt.borrow_mut() = Some(new);
}
/// Retrieves the current idle context.
///
/// # Returns
/// A reference to the idle context.
pub fn idle_context(&self) -> Arc<ContextLock> {
Arc::clone(
self.idle_ctxt
.borrow()
.as_ref()
.expect("no idle context present"),
)
}
}
+5 -1
View File
@@ -1,4 +1,5 @@
use alloc::{
collections::BTreeSet,
sync::{Arc, Weak},
vec::Vec,
};
@@ -12,10 +13,11 @@ use syscall::PtraceFlags;
use crate::{
arch::device::ArchPercpuMisc,
context::{empty_cr3, memory::AddrSpaceWrapper, switch::ContextSwitchPercpu},
context::{empty_cr3, memory::AddrSpaceWrapper, switch::ContextSwitchPercpu, ContextRef},
cpu_set::{LogicalCpuId, MAX_CPU_COUNT},
cpu_stats::{CpuStats, CpuStatsData},
ptrace::Session,
sync::{RwLock, L1},
syscall::debug::SyscallDebugInfo,
};
@@ -44,6 +46,7 @@ pub struct PercpuBlock {
pub misc_arch_info: crate::device::ArchPercpuMisc,
pub stats: CpuStats,
pub contexts: RwLock<L1, BTreeSet<ContextRef>>,
}
static ALL_PERCPU_BLOCKS: [AtomicPtr<PercpuBlock>; MAX_CPU_COUNT as usize] =
@@ -193,6 +196,7 @@ impl PercpuBlock {
misc_arch_info: ArchPercpuMisc::default(),
stats: CpuStats::default(),
contexts: RwLock::new(BTreeSet::new()),
}
}
}
+28
View File
@@ -412,6 +412,34 @@ impl<L: Level, T> RwLock<L, T> {
}
}
pub fn try_read<'a, LP: Lower<L> + 'a>(
&'a self,
lock_token: LockToken<'a, LP>,
) -> Option<RwLockReadGuard<'a, L, T>> {
let inner = match self.inner.try_read() {
Some(inner) => inner,
None => return None,
};
Some(RwLockReadGuard {
inner,
lock_token: LockToken::downgraded(lock_token),
})
}
pub fn try_write<'a, LP: Lower<L> + 'a>(
&'a self,
lock_token: LockToken<'a, LP>,
) -> Option<RwLockWriteGuard<'a, L, T>> {
let inner = match self.inner.try_write() {
Some(inner) => inner,
None => return None,
};
Some(RwLockWriteGuard {
inner,
lock_token: LockToken::downgraded(lock_token),
})
}
// Unsafe due to not using token, currently required by context::switch
pub unsafe fn write_arc(self: &Arc<Self>) -> ArcRwLockWriteGuard<L, T> {
core::mem::forget(self.inner.write());
+6 -1
View File
@@ -65,7 +65,12 @@ pub fn exit_this_context(excp: Option<syscall::Exception>, token: &mut CleanLock
);
}
{
let _ = context::contexts_mut(token.token()).remove(&ContextRef(context_lock));
if !context::contexts_mut(token.token()).remove(&ContextRef(context_lock)) {
#[cfg(feature = "drop_panic")]
{
panic!("This context is not in the cpu")
}
}
}
context::switch(token);
unreachable!();