RustForge

rust-kernel

Rustbelt: a hobby x86_64 kernel written in Rust

rust-kernel / tests / threads.rs
2604 bytes ยท raw
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
//! Preemptive threads: busy-looping threads that never yield still all make
//! progress, because the timer interrupt switches between them.
#![no_std]
#![no_main]
#![feature(custom_test_frameworks)]
#![test_runner(rust_kernel::test_runner)]
#![reexport_test_harness_main = "test_main"]

use bootloader::{BootInfo, entry_point};
use core::panic::PanicInfo;
use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use rust_kernel::allocator;
use rust_kernel::interrupts::TICKS;
use rust_kernel::memory::{self, BootInfoFrameAllocator};
use rust_kernel::thread;
use x86_64::VirtAddr;

entry_point!(main);

fn main(boot_info: &'static BootInfo) -> ! {
    rust_kernel::init();
    let phys_mem_offset = VirtAddr::new(boot_info.physical_memory_offset);
    let mut mapper = unsafe { memory::init(phys_mem_offset) };
    let mut frame_allocator = unsafe { BootInfoFrameAllocator::init(&boot_info.memory_map) };
    allocator::init_heap(&mut mapper, &mut frame_allocator).expect("heap initialization failed");

    test_main();
    rust_kernel::hlt_loop();
}

static COUNT_1: AtomicU64 = AtomicU64::new(0);
static COUNT_2: AtomicU64 = AtomicU64::new(0);
static DONE_2: AtomicBool = AtomicBool::new(false);

/// Never yields, never returns.
fn spinner_1() {
    loop {
        COUNT_1.fetch_add(1, Ordering::Relaxed);
    }
}

/// Never yields, but returns after a while (tests thread exit).
fn spinner_2() {
    for _ in 0..1_000_000 {
        COUNT_2.fetch_add(1, Ordering::Relaxed);
    }
    DONE_2.store(true, Ordering::Relaxed);
}

#[test_case]
fn busy_threads_are_preempted() {
    thread::spawn("spin1", spinner_1);
    thread::spawn("spin2", spinner_2);

    // The main thread busy-waits too (no hlt): if the timer didn't preempt
    // us, the other threads would never run and we'd hit the timeout.
    let deadline = TICKS.load(Ordering::Relaxed) + 18 * 20; // ~20 s
    while !(DONE_2.load(Ordering::Relaxed) && COUNT_1.load(Ordering::Relaxed) > 1000) {
        assert!(TICKS.load(Ordering::Relaxed) < deadline, "threads made no progress");
        core::hint::spin_loop();
    }
    assert_eq!(COUNT_2.load(Ordering::Relaxed), 1_000_000);

    // spinner_2 has exited; spinner_1 must keep running alongside us.
    let before = COUNT_1.load(Ordering::Relaxed);
    let wait_until = TICKS.load(Ordering::Relaxed) + 5;
    while TICKS.load(Ordering::Relaxed) < wait_until {
        core::hint::spin_loop();
    }
    assert!(COUNT_1.load(Ordering::Relaxed) > before, "spinner_1 stopped running");
}

#[panic_handler]
fn panic(info: &PanicInfo) -> ! {
    rust_kernel::test_panic_handler(info)
}