Increase kernel heap size to fit more ELFs, improve scheduler by

removing O(n) operations (ready and sleep queues now), caching timer
ticks per irq on init_interrupts aarch64, only saving context when
switching, add kill syscall and use it for exit, only push input to
init, fix accidental REL_X events drop with virtio, fix the possibility
of not having a process to switch to by either staying, switching or
idling, improve scheduler by removing double locking and splitting code
to multiple functions, do not switch every IRQ but only 100 IRQs,
improve rust app building, move x86_64 kernel fb so it doesn't hit heap,
update submodules

build_rust_app.sh

@@ -1,13 +1,25 @@
 #!/bin/bash
-if [[ "$1" = "init" || "$1" = "libxunil" ]]; then
+
-    cd user/$1
+if [[ "$1" == "init" || "$1" == "libxunil" ]]; then
+    base="user/$1"
 else
-    cd user/apps/$1
+    base="user/apps/$1"
 fi
 
-cargo build --target $KARCH-unknown-none --release --config profile.release.debug=true
+cd "$base" || exit 1
-if [[ "$1" = "init" || "$1" = "libxunil" ]]; then
+
-    cp ./target/$KARCH-unknown-none/release/$1 ../../assets/$KARCH/$1
+cargo build --target "$KARCH-unknown-none" --release \
-else
+    --config profile.release.debug=true
-    cp ./target/$KARCH-unknown-none/release/$1 ../../../assets/$KARCH/$1
+
-fi
+case "$1" in
+    init)
+        cp "./target/$KARCH-unknown-none/release/$1" \
+           "../../assets/$KARCH/$1"
+        ;;
+    libxunil)
+        ;;
+    *)
+        cp "./target/$KARCH-unknown-none/release/$1" \
+           "../../../assets/$KARCH/$1"
+        ;;
+esac

config.mk

@@ -1,4 +1,3 @@
-export KARCH ?= x86_64
 export OUTPUT ?= kernel
-export MEMORY ?= 1G
+export MEMORY ?= 4G
 export TIMER_FREQUENCY_HZ ?= 1000

kernel/src/arch/aarch64/heap.rs

@@ -7,7 +7,7 @@ use crate::{
 pub static ALLOCATOR: Locked<LinkedListAllocator> = Locked::new(LinkedListAllocator::new());
 
 pub const HEAP_START: usize = 0xffffffff90000000;
-pub const HEAP_SIZE: usize = 64 * 1024 * 1024; // 64 MiB
+pub const HEAP_SIZE: usize = 512 * 1024 * 1024; // 512 MiB
 
 pub fn init_heap(mapper: &mut AArchPageTable) {
     let pages = HEAP_SIZE / 4096;

kernel/src/arch/aarch64/interrupts.rs

@@ -6,7 +6,8 @@ use crate::{
     driver::io::virtio::{KEYBOARD_SLOT, MOUSE_SLOT, input::input_interrupt},
     task::{
         context::{UserContext, ctx_save},
-        scheduler::check_and_reschedule,
+        process::ProcessState,
+        scheduler::{SCHEDULER, check_and_reschedule, current_pid},
     },
 };
 
@@ -39,6 +40,8 @@ global_asm!(
   "#
 );
 
+static mut TIMER_TICKS_PER_IRQ: u64 = 0;
+
 fn timer_ticks_per_irq() -> u64 {
     let cntfrq: u64;
     unsafe { core::arch::asm!("mrs {}, cntfrq_el0", out(reg) cntfrq) };
@@ -161,6 +164,7 @@ unsafe fn gic_eoi(id: u32) {
 }
 
 pub fn init_interrupts() {
+    unsafe { TIMER_TICKS_PER_IRQ = timer_ticks_per_irq() };
     unsafe { use_exception_vectors() };
     unsafe { gic_init() };
     enable_interrupts();
@@ -203,17 +207,15 @@ unsafe extern "C" fn irq_handler(ctx: *mut UserContext) {
 
     match interrupt_id {
         30 => {
-            let ticks = timer_ticks_per_irq();
+            unsafe { core::arch::asm!("msr cntp_tval_el0, {}", in(reg) TIMER_TICKS_PER_IRQ) };
-            unsafe { core::arch::asm!("msr cntp_tval_el0, {}", in(reg) ticks) };
             unsafe { core::arch::asm!("msr cntp_ctl_el0, {}", in(reg) 1u64) };
 
             unsafe {
                 gic_eoi(interrupt_id);
             }
 
-            ctx_save(ctx);
             do_interrupt();
-            check_and_reschedule();
+            check_and_reschedule(&ctx);
         }
         interrupt_id if keyboard_irq == interrupt_id as u64 => {
             input_interrupt("kbd");
@@ -276,7 +278,18 @@ unsafe extern "C" fn sync_handler_user(ctx: *mut UserContext) {
                 )
             } as u64;
 
-            ctx_save(ctx as *const UserContext);
+            if let Some(pid) = current_pid() {
+                let guard = SCHEDULER.lock();
+                let needs_yield = guard
+                    .processes
+                    .get(&pid)
+                    .map_or(false, |p| !matches!(p.state, ProcessState::Running));
+                drop(guard);
+                if needs_yield {
+                    ctx_save(ctx);
+                    SCHEDULER.switch_next(true);
+                }
+            }
         }
         _ => handle_aborts(ec, ctx),
     };

kernel/src/arch/syscall.rs

@@ -1,4 +1,5 @@
 #![allow(dead_code, unused_imports)]
+use core::cmp::Reverse;
 use core::sync::atomic::Ordering;
 
 #[cfg(target_arch = "x86_64")]
@@ -85,7 +86,7 @@ type PageTable<'a> = OffsetPageTable<'a>;
 
 fn map_framebuffer() -> isize {
     let pid = current_pid().unwrap_or(0);
-    if pid == 0 {
+    if pid != 1 {
         return -1;
     }
 
@@ -208,6 +209,12 @@ fn close(fd: isize) -> isize {
 }
 
 fn input_read(user_ptr: *mut InputEvent, max_events: isize) -> isize {
+    let pid = current_pid().unwrap_or(0);
+
+    if pid != 1 {
+        return -1;
+    }
+
     #[cfg(target_arch = "x86_64")]
     process_scancodes();
 
@@ -217,12 +224,6 @@ fn input_read(user_ptr: *mut InputEvent, max_events: isize) -> isize {
         return -1;
     }
 
-    let pid = current_pid().unwrap_or(0);
-
-    if pid == 0 {
-        return -1;
-    }
-
     return SCHEDULER
         .with_process(pid as u64, |process| {
             let to_copy = (max_events as usize).min(process.input_buffer.len());
@@ -378,36 +379,18 @@ pub fn exec(arg0: isize) -> isize {
     0
 }
 
-pub fn exit() -> isize {
+pub fn kill(pid: isize, exit_code: isize) -> isize {
-    serial_print("Process Exited.");
+    let pid = pid as u64;
-    let pid = current_pid().unwrap_or(0);
+    println!("PID {} Exited.", pid);
-    if pid == 0 {
-        return 0;
-    }
 
-    let next_pid = {
+    SCHEDULER.terminate_process(pid, exit_code);
-        let mut sched = SCHEDULER.lock();
-
-        sched.processes.remove(&pid);
-
-        sched
-            .processes
-            .iter()
-            .find_map(|(other, proc)| {
-                if *other != pid && matches!(proc.state, ProcessState::Ready) {
-                    Some(*other)
-                } else {
-                    None
-                }
-            })
-            .unwrap_or(0)
-    };
-
-    if next_pid != 0 {
-        SCHEDULER.switch_to(next_pid, false);
-    }
 
+    if pid == current_pid().unwrap_or(0) {
+        SCHEDULER.switch_next(false);
         crate::arch::arch::infinite_idle();
+    } else {
+        0
+    }
 }
 
 fn sleep(ticks: isize) -> isize {
@@ -419,9 +402,15 @@ fn sleep(ticks: isize) -> isize {
         return 0;
     }
 
-    SCHEDULER.with_process(pid, |process| {
+    safe_lock(|| {
-        process.info.wake_tick = Some(TIMER.now().elapsed() + ticks as u64);
+        let mut guard = SCHEDULER.lock();
+        if let Some(process) = guard.processes.get_mut(&pid) {
+            let wake_at = TIMER.now().elapsed() + ticks as u64;
+            process.info.wake_tick = Some(wake_at);
             process.state = ProcessState::Blocked;
+            process.in_ready_queue = false;
+            guard.sleep_queue.push(Reverse((wake_at, pid)));
+        };
     });
 
     0
@@ -686,7 +675,7 @@ pub unsafe extern "C" fn syscall_dispatch(
         OPEN => open(arg0, arg1),
         CLOSE => close(arg0),
         LSEEK => vfs_lseek(arg0 as i64, arg1 as i64, arg2 as i32) as isize,
-        EXIT => exit(),
+        EXIT => kill(current_pid().unwrap() as isize, arg0),
         SLEEP => sleep(arg0),
         EXECVE => exec(arg0),
         CLOCK_GETTIME => ((TIMER.now().elapsed() as usize) * (TIMER_FREQUENCY_HZ / 1000)) as isize,
@@ -705,6 +694,7 @@ pub unsafe extern "C" fn syscall_dispatch(
         IPC_WRITE => ipc_write(arg0, arg1),
         IPC_MANAGE => ipc_manage(arg0, arg1, arg2),
         SHM_OPEN => shm_open(arg0, arg1),
+        KILL => kill(arg0, -9),
         FRAMEBUFFER_SWAP => {
             with_framebuffer(|fb| fb.present());
             0

kernel/src/arch/x86_64/heap.rs

@@ -9,7 +9,7 @@ use x86_64::structures::paging::{
 pub static ALLOCATOR: Locked<LinkedListAllocator> = Locked::new(LinkedListAllocator::new());
 
 pub const HEAP_START: usize = 0xffffffff90000000;
-pub const HEAP_SIZE: usize = 64 * 1024 * 1024; // 64 MiB
+pub const HEAP_SIZE: usize = 512 * 1024 * 1024; // 512 MiB
 
 pub fn init_heap(mapper: &mut OffsetPageTable) -> Result<(), MapToError<Size4KiB>> {
     let page_count = HEAP_SIZE / 4096;

kernel/src/arch/x86_64/interrupts.rs

@@ -1,9 +1,13 @@
 use core::arch::naked_asm;
 
 use crate::{
-    arch::x86_64::gdt,
+    arch::{arch::do_interrupt, x86_64::gdt},
     driver::io::ps2::{keyboard_interrupt, mouse_interrupt, push_scancode},
     println,
+    task::{
+        context::UserContext,
+        scheduler::{check_and_reschedule, current_pid},
+    },
 };
 use lazy_static::lazy_static;
 use pic8259::ChainedPics;
@@ -78,17 +82,21 @@ pub extern "x86-interrupt" fn page_fault_handler(
     error_code: PageFaultErrorCode,
 ) {
     panic!(
-        "EXCEPTION: PAGE FAULT\nAccessed Addresss: {:?}\nError Code: {:?}\n{:#?}",
+        "EXCEPTION: PAGE FAULT\nAccessed Address: {:?}\nError Code: {:?}\nCurrent PID: {}\nPER_CPU: {:#x}\n{:#?}",
         Cr2::read(),
         error_code,
+        current_pid().unwrap_or(0),
+        &raw const crate::arch::x86_64::syscall::PER_CPU as u64,
         stack_frame
     );
 }
 
 pub extern "x86-interrupt" fn gpf_handler(stack_frame: InterruptStackFrame, error_code: u64) {
     panic!(
-        "EXCEPTION: GENERAL PROTECTION FAULT\nError Code: {:?}\n{:#?}",
+        "EXCEPTION: GENERAL PROTECTION FAULT\nError Code: {:?}\nCurrent PID: {}\n{:#?}",
-        error_code, stack_frame
+        error_code,
+        current_pid().unwrap_or(0),
+        stack_frame
     );
 }
 
@@ -132,10 +140,7 @@ pub fn timer_interrupt_handler() {
 
         mov rdi, rsp
 
-        call ctx_save
+        call x86_interrupt
-        call do_interrupt
-        call eoi
-        call check_and_reschedule
         test rax, rax
         jnz .switched
 
@@ -202,6 +207,13 @@ pub fn timer_interrupt_handler() {
     )
 }
 
+#[unsafe(no_mangle)]
+extern "C" fn x86_interrupt(ctx: *mut UserContext) -> isize {
+    do_interrupt();
+    eoi();
+    check_and_reschedule(&ctx)
+}
+
 #[unsafe(no_mangle)]
 extern "C" fn eoi() {
     unsafe {

kernel/src/arch/x86_64/syscall.rs

@@ -3,8 +3,11 @@ use core::arch::asm;
 use x86_64::instructions::tlb::flush_all;
 
 use crate::{
+    arch::arch::safe_lock,
     arch::x86_64::gdt::{GDT, TSS_MUTEX},
     task::context::UserContext,
+    task::process::ProcessState,
+    task::scheduler::{SCHEDULER, current_pid},
 };
 
 const IA32_EFER: u32 = 0xC0000080;
@@ -129,6 +132,8 @@ unsafe extern "C" fn syscall_entry() {
             call ctx_save
             mov rsp, rbx
 
+            sti
+
             mov rdi, qword ptr [rbx + 112]
             mov rsi, qword ptr [rbx + 72]
             mov rdx, qword ptr [rbx + 64]
@@ -149,6 +154,9 @@ unsafe extern "C" fn syscall_entry() {
             call ctx_save
             mov rsp, rbx
 
+            mov rdi, rbx
+            call syscall_yield_check
+
             mov rax, qword ptr [rsp + 128]
             mov qword ptr [rsp + 96], rax
             mov rax, qword ptr [rsp + 136]
@@ -170,6 +178,7 @@ unsafe extern "C" fn syscall_entry() {
             mov r13, qword ptr [rsp + 16]
             mov r14, qword ptr [rsp + 8]
             mov r15, qword ptr [rsp + 0]
+            cli
             mov rsp, qword ptr gs:[0]
             swapgs
             sysretq
@@ -179,6 +188,22 @@ unsafe extern "C" fn syscall_entry() {
     );
 }
 
+#[unsafe(no_mangle)]
+unsafe extern "C" fn syscall_yield_check(_ctx: *mut UserContext) {
+    if let Some(pid) = current_pid() {
+        let needs_yield = {
+            let guard = safe_lock(|| SCHEDULER.lock());
+            guard
+                .processes
+                .get(&pid)
+                .map_or(false, |p| !matches!(p.state, ProcessState::Running))
+        };
+        if needs_yield {
+            SCHEDULER.switch_next(true);
+        }
+    }
+}
+
 #[unsafe(naked)]
 #[unsafe(no_mangle)]
 pub unsafe fn run_next(ctx: *const UserContext, user_rsp: u64, user_cs: u64, user_ss: u64) {

kernel/src/config.rs

@@ -1,2 +1,2 @@
 pub const TIMER_FREQUENCY_HZ: usize = 1000;
-pub const KARCH: &str = "x86_64";
+pub const KARCH: &str = "aarch64";

kernel/src/driver/elf/loader.rs

@@ -176,6 +176,8 @@ pub fn run_elf(file_bytes: &[u8], should_swapgs: bool, switch_to: bool) {
 
         #[cfg(target_arch = "aarch64")]
         SCHEDULER.with_process(process_pid, |process| {
+            use crate::task::context::UserContext;
+
             process.saved_ctx = Some(UserContext {
                 x0: 0,
                 x1: 0,

kernel/src/driver/framebuffer.rs

@@ -63,7 +63,7 @@ impl Framebuffer {
                 FrameAllocator, Mapper, Page, PageTableFlags, PhysFrame, Size4KiB,
             },
         };
-        const KERNEL_FB_BASE: u64 = 0xffffffffa0000000;
+        const KERNEL_FB_BASE: u64 = 0xffffffffc0000000;
         let buf_len = self.pitch * self.height;
         let byte_len = buf_len * core::mem::size_of::<u32>();
         let pixel_frames = (byte_len + 4095) / 4096;

kernel/src/driver/io/fs/assets.rs

@@ -1,6 +1,6 @@
-pub static INIT_ELF: &[u8] = include_bytes!("../../../../../assets/x86_64/init");
+pub static INIT_ELF: &[u8] = include_bytes!("../../../../../assets/aarch64/init");
 pub static DOOM_WAD: &[u8] = include_bytes!("../../../../../assets/doom1.wad");
-pub static DOOM_ELF: &[u8] = include_bytes!("../../../../../assets/x86_64/doomgeneric");
+pub static DOOM_ELF: &[u8] = include_bytes!("../../../../../assets/aarch64/doomgeneric");
-pub static HELLOWORLD_ELF: &[u8] = include_bytes!("../../../../../assets/x86_64/helloworld.elf");
+pub static HELLOWORLD_ELF: &[u8] = include_bytes!("../../../../../assets/aarch64/helloworld.elf");
-pub static BADAPPLE_ELF: &[u8] = include_bytes!("../../../../../assets/x86_64/badapple");
+pub static BADAPPLE_ELF: &[u8] = include_bytes!("../../../../../assets/aarch64/badapple");
-pub static SHELL_ELF: &[u8] = include_bytes!("../../../../../assets/x86_64/shell");
+pub static SHELL_ELF: &[u8] = include_bytes!("../../../../../assets/aarch64/shell");

kernel/src/driver/io/input.rs

@@ -139,7 +139,7 @@ pub fn process_input() {
         };
 
         let mut scheduler = SCHEDULER.lock();
-        for process in scheduler.processes.values_mut() {
+        if let Some(process) = scheduler.processes.get_mut(&1) {
             process.input_buffer.push(input_event);
         }
         drop(scheduler);

kernel/src/driver/io/virtio/input.rs

@@ -112,10 +112,6 @@ pub fn input_interrupt(device_type: &str) {
 
 #[allow(static_mut_refs)]
 pub fn handle_event(event: &VirtioInputEvent) {
-    if event.event_type == 0 || event.code == 0 {
-        return;
-    }
-
     enqueue_input_event(InputEvent {
         event_type: event.event_type,
         code: event.code,

kernel/src/main.rs

@@ -155,6 +155,7 @@ unsafe extern "C" fn kmain() -> ! {
                     kernel_main_x86_64()
                 }
 
+                #[cfg(target_arch = "aarch64")]
                 loop {}
             } else {
                 kernel_crash()

kernel/src/mm/shm.rs

@@ -3,7 +3,7 @@ use core::sync::atomic::AtomicU64;
 use alloc::{collections::btree_map::BTreeMap, string::String, vec::Vec};
 use spin::mutex::Mutex;
 
-pub const USER_SHM_BASE: u64 = 0x0000_7000_0000_0000;
+pub const USER_SHM_BASE: u64 = 0x0000_4000_0000_0000;
 pub static SHM_REGISTRY: Mutex<Option<BTreeMap<String, SharedMemory>>> = Mutex::new(None);
 pub static NEXT_SHM_ID: AtomicU64 = AtomicU64::new(1);
 pub const SHM_SLOT_SIZE: u64 = 64 * 1024 * 1024;

kernel/src/task/process.rs

@@ -15,7 +15,7 @@ pub enum ProcessState {
 }
 
 pub struct ProcessInfo {
-    pub exit_code: usize,
+    pub exit_code: isize,
     pub parent: usize,
     pub wake_tick: Option<u64>,
 }
@@ -33,6 +33,7 @@ pub struct Process {
     pub user_entry: u64,
     pub last_switch_tick: u64,
     pub info: ProcessInfo,
+    pub in_ready_queue: bool,
 }
 impl Process {
     pub fn new(
@@ -60,6 +61,7 @@ impl Process {
                 parent: 0,
                 wake_tick: None,
             },
+            in_ready_queue: false,
         }
     }
 }

kernel/src/task/scheduler.rs

@@ -1,10 +1,7 @@
 use core::sync::atomic::{AtomicU64, Ordering};
 
-use alloc::{collections::btree_map::BTreeMap, vec::Vec};
-
 use crate::{
     arch::arch::{GLOBAL_TICK_COUNT, safe_lock},
-    config::TIMER_FREQUENCY_HZ,
     driver::timer::TIMER,
     task::{
         context::UserContext,
@@ -12,6 +9,8 @@ use crate::{
     },
     util::Locked,
 };
+use alloc::collections::{binary_heap::BinaryHeap, btree_map::BTreeMap, vec_deque::VecDeque};
+use core::cmp::Reverse;
 
 #[cfg(target_arch = "aarch64")]
 use crate::arch::aarch64::interrupts::run_next;
@@ -21,6 +20,7 @@ use crate::arch::x86_64::{
     syscall::run_next,
     usermode::enter_usermode_x86_64,
 };
+use crate::task::context::ctx_save;
 #[cfg(target_arch = "x86_64")]
 use x86_64::structures::gdt::SegmentSelector;
 
@@ -44,8 +44,16 @@ pub fn set_current_pid(pid: Option<u64>) {
     CURRENT_PID.store(pid.unwrap_or(0), Ordering::Relaxed);
 }
 
+enum SwitchDecision {
+    Switch(Option<UserContext>, u64, u64),
+    Stay,
+    Idle,
+}
+
 pub struct Scheduler {
     pub processes: BTreeMap<u64, Process>,
+    ready_queue: VecDeque<u64>,
+    pub sleep_queue: BinaryHeap<Reverse<(u64, u64)>>,
     next_pid: u64,
 }
 
@@ -53,6 +61,8 @@ impl Scheduler {
     pub const fn new() -> Scheduler {
         Scheduler {
             processes: BTreeMap::new(),
+            ready_queue: VecDeque::new(),
+            sleep_queue: BinaryHeap::new(),
             next_pid: 1,
         }
     }
@@ -79,89 +89,87 @@ impl Locked<Scheduler> {
         );
         guard.processes.insert(pid, process);
 
+        if let Some(new_process) = guard.processes.get_mut(&pid) {
+            new_process.in_ready_queue = true;
+        }
+        guard.ready_queue.push_back(pid);
+
         Some(pid)
     }
 
-    pub fn next_task(&self) -> u64 {
+    pub fn switch_next(&self, should_swapgs: bool) -> bool {
-        if let Some(previous_pid) = current_pid() {
+        let previous_pid = current_pid().unwrap_or(0);
-            let mut guard = safe_lock(|| self.lock());
+        if previous_pid == 0 {
-
+            return false;
-            for process in guard.processes.values_mut() {
-                if process.info.wake_tick.is_some() {
-                    if TIMER.now().elapsed() >= process.info.wake_tick.unwrap() {
-                        process.state = ProcessState::Ready;
-                        process.info.wake_tick = None;
-                    }
-                }
         }
 
-            if let Some(process) = guard.processes.get_mut(&previous_pid) {
+        loop {
-                if matches!(process.state, ProcessState::Running) {
+            let decision = safe_lock(|| {
-                    process.state = ProcessState::Ready;
+                let mut guard = self.lock();
+                self.wake_processes(&mut guard);
+                self.mark_old_process_ready(&mut guard);
+
+                let mut chosen: Option<u64> = None;
+                while let Some(pid) = guard.ready_queue.pop_front() {
+                    if let Some(process) = guard.processes.get_mut(&pid) {
+                        process.in_ready_queue = false;
+                        if matches!(process.state, ProcessState::Ready) {
+                            chosen = Some(pid);
+                            break;
                         }
+                    };
                 }
 
-            let ready_pids: Vec<u64> = guard
+                let new_pid = match chosen {
+                    Some(pid) => pid,
+                    None => {
+                        let is_prev_ready = matches!(
+                            guard
                                 .processes
-                .iter()
+                                .get(&previous_pid)
-                .filter(|(_, process)| matches!(process.state, ProcessState::Ready))
+                                .map(|process| &process.state),
-                .map(|(&pid, _)| pid)
+                            Some(ProcessState::Ready)
-                .collect();
+                        );
 
-            if ready_pids.is_empty() {
+                        if is_prev_ready {
-                return previous_pid;
+                            previous_pid
-            }
-
-            let current_index = ready_pids.iter().position(|&pid| pid == previous_pid);
-
-            return match current_index {
-                Some(i) => {
-                    let next_index = (i + 1) % ready_pids.len();
-                    ready_pids[next_index]
-                }
-                None => ready_pids[0],
-            };
                         } else {
-            panic!("Could not get current PID when switching to next task")
+                            return SwitchDecision::Idle;
+                        }
+                    }
                 };
+
+                if current_pid() == Some(new_pid) {
+                    if let Some(p) = guard.processes.get_mut(&new_pid) {
+                        p.state = ProcessState::Running;
+                        p.last_switch_tick = GLOBAL_TICK_COUNT.load(Ordering::Relaxed);
                     }
 
-    #[allow(unused_variables)]
+                    return SwitchDecision::Stay;
-    pub fn switch_to(&self, pid: u64, should_swapgs: bool) {
-        let (ctx_opt, entry, stack_top, kernel_stack_top) = {
-            let mut guard = safe_lock(|| self.lock());
-
-            if let Some(previous_pid) = current_pid() {
-                if let Some(old_process) = guard.processes.get_mut(&previous_pid) {
-                    if matches!(old_process.state, ProcessState::Running) {
-                        old_process.state = ProcessState::Ready;
-                    }
-                } else {
-                    // no previous process
-                }
                 }
 
-            let new_process = guard.processes.get_mut(&pid).expect("Cant get new process");
+                let new_process = guard
+                    .processes
+                    .get_mut(&new_pid)
+                    .expect("Cant get new process");
+
                 new_process.state = ProcessState::Running;
+                new_process.last_switch_tick = GLOBAL_TICK_COUNT.load(Ordering::Relaxed);
+                new_process.in_ready_queue = false;
                 if let Some(address_space) = new_process.address_space.as_mut() {
                     address_space.use_address_space();
-            };
+                }
 
-            (
+                let ctx_opt = new_process.saved_ctx;
-                new_process.saved_ctx,
+                let entry = new_process.user_entry;
-                new_process.user_entry,
+                let stack_top = new_process.stack_top;
-                new_process.stack_top,
+                let kernel_stack_top = new_process.kernel_stack_top;
-                new_process.kernel_stack_top,
-            )
-        };
-
-        set_current_pid(Some(pid));
 
+                set_current_pid(Some(new_pid));
                 #[cfg(target_arch = "x86_64")]
                 unsafe {
-            use x86_64::VirtAddr;
-
                     use crate::arch::x86_64::{gdt::TSS_MUTEX, syscall::PER_CPU};
+                    use x86_64::VirtAddr;
 
                     PER_CPU.kernel_rsp = kernel_stack_top;
                     if let Some(tss) = TSS_MUTEX.lock().as_mut() {
@@ -169,6 +177,75 @@ impl Locked<Scheduler> {
                     }
                 }
 
+                #[cfg(target_arch = "aarch64")]
+                let _ = kernel_stack_top;
+
+                SwitchDecision::Switch(ctx_opt, entry, stack_top)
+            });
+
+            match decision {
+                SwitchDecision::Switch(ctx_opt, entry, stack_top) => {
+                    self.do_context_switch(ctx_opt, entry, stack_top, should_swapgs);
+                    return true;
+                }
+
+                SwitchDecision::Stay => return false,
+
+                SwitchDecision::Idle => {
+                    set_current_pid(None);
+
+                    #[cfg(target_arch = "x86_64")]
+                    {
+                        x86_64::instructions::interrupts::enable_and_hlt();
+                        x86_64::instructions::interrupts::disable();
+                    }
+
+                    #[cfg(target_arch = "aarch64")]
+                    unsafe {
+                        use aarch64_cpu::registers::{DAIF, Writeable};
+                        DAIF.write(
+                            DAIF::D::Masked + DAIF::A::Masked + DAIF::I::Unmasked + DAIF::F::Masked,
+                        );
+                        core::arch::asm!("wfi");
+                        DAIF.write(
+                            DAIF::D::Masked + DAIF::A::Masked + DAIF::I::Masked + DAIF::F::Masked,
+                        );
+                    }
+                }
+            }
+        }
+    }
+
+    pub fn wake_processes(&self, guard: &mut Scheduler) {
+        let now = TIMER.now().elapsed();
+        while let Some(&Reverse((wake_tick, pid))) = guard.sleep_queue.peek() {
+            if wake_tick > now {
+                break;
+            }
+
+            guard.sleep_queue.pop();
+
+            if let Some(process) = guard.processes.get_mut(&pid) {
+                if process.info.wake_tick == Some(wake_tick) {
+                    process.state = ProcessState::Ready;
+                    process.info.wake_tick = None;
+                    if !process.in_ready_queue {
+                        process.in_ready_queue = true;
+                        guard.ready_queue.push_back(process.pid);
+                    }
+                }
+            }
+        }
+    }
+
+    #[allow(unused_variables)]
+    fn do_context_switch(
+        &self,
+        ctx_opt: Option<UserContext>,
+        entry: u64,
+        stack_top: u64,
+        should_swapgs: bool,
+    ) {
         #[cfg(target_arch = "aarch64")]
         unsafe {
             let saved_ctx = ctx_opt.expect("Could not get user context");
@@ -193,6 +270,74 @@ impl Locked<Scheduler> {
         }
     }
 
+    pub fn switch_to(&self, pid: u64, should_swapgs: bool) {
+        let switch_info = safe_lock(|| {
+            let mut guard = self.lock();
+            self.wake_processes(&mut guard);
+            self.mark_old_process_ready(&mut guard);
+            let mut new_process = guard.processes.get_mut(&pid).expect("Cant get new process");
+            let (ctx_opt, entry, stack_top, kernel_stack_top) = self.use_process(&mut new_process);
+
+            set_current_pid(Some(pid));
+            #[cfg(target_arch = "x86_64")]
+            unsafe {
+                use crate::arch::x86_64::{gdt::TSS_MUTEX, syscall::PER_CPU};
+                use x86_64::VirtAddr;
+
+                PER_CPU.kernel_rsp = kernel_stack_top;
+                if let Some(tss) = TSS_MUTEX.lock().as_mut() {
+                    tss.privilege_stack_table[0] = VirtAddr::new(kernel_stack_top);
+                }
+            }
+
+            Some((ctx_opt, entry, stack_top))
+        });
+
+        if let Some((ctx_opt, entry, stack_top)) = switch_info {
+            self.do_context_switch(ctx_opt, entry, stack_top, should_swapgs);
+        }
+    }
+
+    #[inline]
+    pub fn use_process(&self, new_process: &mut Process) -> (Option<UserContext>, u64, u64, u64) {
+        new_process.last_switch_tick = GLOBAL_TICK_COUNT.load(Ordering::Relaxed);
+        new_process.state = ProcessState::Running;
+        new_process.in_ready_queue = false;
+        if let Some(address_space) = new_process.address_space.as_mut() {
+            address_space.use_address_space();
+        };
+
+        (
+            new_process.saved_ctx,
+            new_process.user_entry,
+            new_process.stack_top,
+            new_process.kernel_stack_top,
+        )
+    }
+
+    pub fn mark_old_process_ready(&self, guard: &mut Scheduler) {
+        if let Some(previous_pid) = current_pid() {
+            if let Some(old_process) = guard.processes.get_mut(&previous_pid) {
+                if matches!(old_process.state, ProcessState::Running) && !old_process.in_ready_queue
+                {
+                    old_process.state = ProcessState::Ready;
+                    old_process.in_ready_queue = true;
+                    guard.ready_queue.push_back(old_process.pid);
+                }
+            }
+        }
+    }
+
+    pub fn terminate_process(&self, pid: u64, exit_code: isize) {
+        let mut guard = self.lock();
+        if let Some(process) = guard.processes.get_mut(&pid) {
+            process.state = ProcessState::Zombie;
+            process.in_ready_queue = false;
+            process.input_buffer.clear();
+            process.info.exit_code = exit_code;
+        }
+    }
+
     pub fn with_process<F, R>(&self, index: u64, f: F) -> Option<R>
     where
         F: FnOnce(&mut Process) -> R,
@@ -206,35 +351,28 @@ impl Locked<Scheduler> {
 pub static SCHEDULER: Locked<Scheduler> = Locked::new(Scheduler::new());
 
 #[unsafe(no_mangle)]
-pub extern "C" fn check_and_reschedule() -> isize {
+pub extern "C" fn check_and_reschedule(ctx: &*mut UserContext) -> isize {
     let current_pid = CURRENT_PID.load(Ordering::Relaxed);
     let should = safe_lock(|| {
         let mut scheduler = SCHEDULER.lock();
 
+        SCHEDULER.wake_processes(&mut scheduler);
+
         if let Some(process) = scheduler.processes.get_mut(&current_pid) {
             let elapsed = GLOBAL_TICK_COUNT.load(Ordering::Relaxed) - process.last_switch_tick;
-            if elapsed >= (TIMER_FREQUENCY_HZ / 60) as u64 {
+            elapsed >= 100
-                process.last_switch_tick = GLOBAL_TICK_COUNT.load(Ordering::Relaxed);
-                true
-            } else {
-                false
-            }
         } else {
             false
         }
     });
 
     if should {
-        let next_task = SCHEDULER.next_task();
+        ctx_save(*ctx);
-
+        match SCHEDULER.switch_next(true) {
-        if next_task == current_pid {
+            true => return 1,
-            return 0;
+            false => return 0,
+        }
     }
 
-        SCHEDULER.switch_to(next_task, true);
-
-        1
-    } else {
     0
-    }
 }