Files
moss-kernel/src/sched/uspc_ret.rs
2025-12-21 01:37:14 -08:00

280 lines
13 KiB
Rust

use super::{SCHED_STATE, current_task, schedule, waker::create_waker};
use crate::process::TASK_LIST;
use crate::{
arch::{Arch, ArchImpl},
process::{
TaskState,
ctx::UserCtx,
exit::kernel_exit_with_signal,
thread_group::{
signal::{SigId, ksigaction::KSignalAction},
wait::ChildState,
},
},
};
use alloc::boxed::Box;
use core::{ptr, task::Poll};
enum State {
PickNewTask,
ProcessKernelWork,
ReturnToUserspace,
}
/// Prepares the kernel for a safe return to userspace, guaranteeing a valid
/// context frame is prepared.
///
/// This function is the primary gateway for transitioning from a kernel context
/// (e.g., after a syscall, interrupt, or fault) back to userspace code. It
/// ensures a task is in a valid state to resume execution, handling any
/// necessary kernel business before the final context restore.
///
/// The core of this function is a loop that guarantees progress. It will always
/// find a runnable task and prepare its context for userspace return, even if
/// the initial task needs to be terminated, put to sleep, or handle a signal
/// first.
///
/// The logical flow inside the loop for a chosen task is as follows:
///
/// 1. Schedule: A task is selected to run next. This may trigger a context
/// switch if the scheduler picks a different task than the one that entered
/// the kernel.
/// 2. Poll In-Flight Work: It checks for and polls any outstanding asynchronous
/// work associated with the task, such as signal delivery (`signal_work`) or
/// another kernel future (`kern_work`).
/// - If the work is `Poll::Pending`, the task is put to sleep, and the loop
/// continues to find another task.
/// - If the work completes successfully (`Poll::Ready(Ok(_))`), the
/// function proceeds.
/// - If signal delivery fails (`Poll::Ready(Err(_))`), it's a fatal error,
/// and the task is terminated. The loop continues to find another task.
/// 3. Action New Signals: If there is no in-flight work, it checks for any new
/// pending signals. Depending on the signal's disposition, the task may be
/// terminated or new `signal_work` may be queued. If new work is queued, the
/// loop restarts to process it immediately.
/// 4. Final Return: If a task has no pending work and no new signals to action,
/// its userspace context is restored into the `ctx` pointer, the current
/// address space is guaranteed to be active, and the function returns.
///
/// # Parameters
/// - `ctx`: A mutable pointer to the stack-allocated `UserCtx` (the exception
/// frame). This function will write the register context of the task being
/// returned to into this location, making it ready for the
/// architecture-specific return-to-userspace instructions.
///
/// # Guarantees
/// - This function will always return; it does not diverge (unless it panics).
/// - Upon return, the address space of the task being resumed will be active.
/// - Upon return, the memory pointed to by `ctx` will be populated with the
/// valid userspace register state for that task.
///
/// # Panics
/// This function will panic if it detects an attempt to process signals or
/// kernel work for the special idle task, as this indicates a critical bug in
/// the scheduler or task management.
pub fn dispatch_userspace_task(ctx: *mut UserCtx) {
let mut state = State::PickNewTask;
loop {
match state {
State::PickNewTask => {
// Pick a new task, poteintally context switching to a new task.
schedule();
state = State::ProcessKernelWork;
}
State::ProcessKernelWork => {
let task = current_task();
// First, let's handle signals. If there is any scheduled signal
// work (this has to be async to handle faults, etc).
let signal_work = task.ctx.lock_save_irq().take_signal_work();
if let Some(mut signal_work) = signal_work {
if task.is_idle_task() {
panic!("Signal processing for idle task");
}
match signal_work
.as_mut()
.poll(&mut core::task::Context::from_waker(&create_waker(
task.descriptor(),
))) {
Poll::Ready(Ok(state)) => {
// Signal actioning is complete. Return to userspace.
unsafe { ptr::copy_nonoverlapping(&state as _, ctx, 1) };
return;
}
Poll::Ready(Err(_)) => {
// If we errored, then we *cannot* progress the task.
// Delivery of the signal failed. Force the process to
// terminate.
kernel_exit_with_signal(SigId::SIGSEGV, true);
// Look for another task, this one is now dead.
state = State::PickNewTask;
continue;
}
Poll::Pending => {
task.ctx.lock_save_irq().put_signal_work(signal_work);
*task.state.lock_save_irq() = TaskState::Sleeping;
state = State::PickNewTask;
continue;
}
}
}
// Now let's handle any kernel work that's been spawned for this task.
let kern_work = task.ctx.lock_save_irq().take_kernel_work();
if let Some(mut kern_work) = kern_work {
if task.is_idle_task() {
panic!("Idle process should never have kernel work");
}
match kern_work
.as_mut()
.poll(&mut core::task::Context::from_waker(&create_waker(
task.descriptor(),
))) {
Poll::Ready(()) => {
// If the task just exited (entered the finished state),
// don't return to it's userspace, instead, find another
// task to execute, removing this task from the
// runqueue, reaping it's resouces.
if task.state.lock_save_irq().is_finished() {
SCHED_STATE
.borrow_mut()
.run_queue
.remove(&task.descriptor());
let mut task_list = TASK_LIST.lock_save_irq();
task_list.remove(&task.descriptor());
state = State::PickNewTask;
continue;
}
// Kernel work finished. Ensure we have no other new
// work to process (i.e. a signal was rasied). We
// don't need to clear the kernel context here as we
// used the *take* function above.
state = State::ProcessKernelWork;
continue;
}
Poll::Pending => {
// Kernel work hasn't finished. A wake up should
// have been scheduled by the future. Replace the
// kernel work context back into the task, set it's
// state to sleeping so it's not scheduled again and
// search for another task to execute.
task.ctx.lock_save_irq().put_kernel_work(kern_work);
*task.state.lock_save_irq() = TaskState::Sleeping;
state = State::PickNewTask;
continue;
}
}
}
// No kernel work. Check for any pending signals.
// We never handle signals for the idle task.
if task.is_idle_task() {
state = State::ReturnToUserspace;
continue;
}
// See if there are any signals we need to action.
let task = current_task();
let mut pending_task_sigs = task.pending_signals.lock_save_irq();
let mask = task.sig_mask.lock_save_irq();
if let Some((id, action)) = task
.process
.signals
.lock_save_irq()
.action_signal(*mask, &mut pending_task_sigs)
{
match action {
KSignalAction::Term | KSignalAction::Core => {
// Terminate the process, and find a new task.
kernel_exit_with_signal(id, false);
state = State::PickNewTask;
continue;
}
KSignalAction::Stop => {
// Default action: stop (suspend) the entire process.
let process = &task.process;
// Notify the parent that we have stopped (SIGCHLD).
if let Some(parent) = process
.parent
.lock_save_irq()
.as_ref()
.and_then(|p| p.upgrade())
{
parent
.child_notifiers
.child_update(process.tgid, ChildState::Stop { signal: id });
parent.signals.lock_save_irq().set_pending(SigId::SIGCHLD);
}
for thr_weak in process.threads.lock_save_irq().values() {
if let Some(thr) = thr_weak.upgrade() {
*thr.state.lock_save_irq() = TaskState::Stopped;
}
}
state = State::PickNewTask;
continue;
}
KSignalAction::Continue => {
let process = &task.process;
// Wake up all sleeping threads in the process.
for thr_weak in process.threads.lock_save_irq().values() {
if let Some(thr) = thr_weak.upgrade() {
let mut st = thr.state.lock_save_irq();
if *st == TaskState::Sleeping {
*st = TaskState::Runnable;
}
}
}
// Notify the parent that we have continued (SIGCHLD).
if let Some(parent) = process
.parent
.lock_save_irq()
.as_ref()
.and_then(|p| p.upgrade())
{
parent
.child_notifiers
.child_update(process.tgid, ChildState::Continue);
parent.signals.lock_save_irq().set_pending(SigId::SIGCHLD);
}
// Re-process kernel work for this task (there may be more to do).
state = State::ProcessKernelWork;
continue;
}
KSignalAction::Userspace(id, action) => {
let fut = ArchImpl::do_signal(id, action);
task.ctx.lock_save_irq().put_signal_work(Box::pin(fut));
state = State::ProcessKernelWork;
continue;
}
}
}
state = State::ReturnToUserspace;
}
State::ReturnToUserspace => {
// Real user-space return now.
current_task().ctx.lock_save_irq().restore_user_ctx(ctx);
return;
}
}
}
}