//! Stage 11: several CPUs, and the scheduler that shares them out.
//! - all 4 CPUs come up and take timer interrupts
//! - threads and processes really run on more than one CPU at a time
//! - locks still work when two CPUs race for them
//! - sleeping threads wake on time, in order, and cost no CPU while asleep
//! - wait queues block and wake without busy-polling
//! - spawning and exiting threads/processes in a loop leaks no memory
#![no_std]
#![no_main]
#![feature(custom_test_frameworks)]
#![test_runner(rustbelt::test_runner)]
#![reexport_test_harness_main = "test_main"]

extern crate alloc;

use alloc::{string::String, sync::Arc, vec::Vec};
use bootloader::{BootInfo, entry_point};
use core::panic::PanicInfo;
use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use rustbelt::interrupts::ticks;
use rustbelt::process::{self, ExitStatus, SpawnArgs};
use rustbelt::thread::{self, WaitQueue};
use rustbelt::{fs, memory, smp, stack};
use spin::Mutex;
use x86_64::instructions::interrupts::without_interrupts;

entry_point!(main);

fn main(boot_info: &'static BootInfo) -> ! {
    rustbelt::init();
    rustbelt::init_memory(boot_info);
    rustbelt::init_smp();
    fs::init();
    rustbelt::init_storage();
    test_main();
    rustbelt::hlt_loop();
}

/// Start /mnt/disk0/bin/<name> with these arguments.
fn run(name: &str, extra_args: &[&str]) -> Arc<process::Process> {
    let mut spawn = SpawnArgs::kernel(name);
    spawn.args.extend(extra_args.iter().map(|a| String::from(*a)));
    process::spawn_path(&alloc::format!("/mnt/disk0/bin/{}", name), spawn).expect("spawn failed")
}

/// (ticks, busy ticks) of every CPU right now.
fn cpu_ticks() -> Vec<(u64, u64)> {
    (0..smp::cpu_count()).map(thread::cpu_ticks).collect()
}

#[test_case]
fn four_cpus_are_online_and_ticking() {
    assert_eq!(smp::cpu_count(), 4, "QEMU runs with -smp 4");
    let before = cpu_ticks();
    thread::sleep_ticks(4);
    let after = cpu_ticks();
    for cpu in 0..4 {
        assert!(after[cpu].0 > before[cpu].0, "CPU {} got no timer ticks", cpu);
    }
}

// --- work runs on several CPUs at once --------------------------------------

static STOP: AtomicBool = AtomicBool::new(false);

fn spin_until_stopped() {
    while !STOP.load(Ordering::Relaxed) {
        core::hint::spin_loop();
    }
}

#[test_case]
fn busy_threads_run_on_several_cpus_in_parallel() {
    STOP.store(false, Ordering::Relaxed);
    let ids: Vec<_> = (0..4).map(|_| thread::spawn("spinner", spin_until_stopped)).collect();
    thread::sleep_ticks(2); // let them spread out
    let before = cpu_ticks();
    thread::sleep_ticks(6); // ~330 ms; we sleep, the 4 spinners get 4 CPUs
    let after = cpu_ticks();
    let mut ran_on = 0;
    for &id in &ids {
        ran_on |= thread::ran_on(id).expect("spinner vanished");
    }
    STOP.store(true, Ordering::Relaxed);
    for id in ids {
        thread::join(id);
    }
    // At least two CPUs were busy (>= half their ticks) at the same time.
    let busy_cpus = (0..4)
        .filter(|&c| {
            let total = after[c].0 - before[c].0;
            let busy = after[c].1 - before[c].1;
            total > 0 && busy * 2 >= total
        })
        .count();
    assert!(ran_on.count_ones() >= 2, "spinners ran on CPUs {:#b} only", ran_on);
    assert!(busy_cpus >= 2, "only {} CPUs were busy: {:?} -> {:?}", busy_cpus, before, after);
}

#[test_case]
fn user_processes_run_on_several_cpus() {
    let ps: Vec<_> = (0..3).map(|_| run("spin", &[])).collect();
    let mut cpus = 0;
    for p in &ps {
        assert_eq!(process::wait(p), ExitStatus::Exited(0));
        cpus |= p.ran_on.load(Ordering::Relaxed);
    }
    assert!(cpus.count_ones() >= 2, "3 spin processes ran on CPUs {:#b} only", cpus);
}

// --- locks across CPUs --------------------------------------------------------

static COUNTER: Mutex<u64> = Mutex::new(0);

fn add_a_lot() {
    for _ in 0..50_000 {
        // Read-modify-write under the lock: without it, two CPUs would
        // overwrite each other's updates and the total would come up short.
        let mut c = COUNTER.lock();
        *c = core::hint::black_box(*c) + 1;
    }
}

#[test_case]
fn a_lock_shared_by_four_cpus_loses_no_updates() {
    *COUNTER.lock() = 0;
    let ids: Vec<_> = (0..4).map(|_| thread::spawn("adder", add_a_lot)).collect();
    for id in ids {
        thread::join(id);
    }
    assert_eq!(*COUNTER.lock(), 200_000);
}

// --- sleep -------------------------------------------------------------------

#[test_case]
fn sleep_lasts_as_long_as_asked() {
    let start = ticks();
    thread::sleep_ticks(10);
    let slept = ticks() - start;
    assert!((10..=12).contains(&slept), "asked for 10 ticks, slept {}", slept);
}

static WOKE: Mutex<Vec<u64>> = Mutex::new(Vec::new());

fn sleeper(n: u64) {
    thread::sleep_ticks(n);
    without_interrupts(|| WOKE.lock().push(n));
}
fn sleep_9() {
    sleeper(9)
}
fn sleep_3() {
    sleeper(3)
}
fn sleep_6() {
    sleeper(6)
}

#[test_case]
fn sleepers_wake_in_deadline_order() {
    WOKE.lock().clear();
    let ids = [thread::spawn("s9", sleep_9), thread::spawn("s3", sleep_3), thread::spawn("s6", sleep_6)];
    for id in ids {
        thread::join(id);
    }
    assert_eq!(*WOKE.lock(), [3, 6, 9]);
}

#[test_case]
fn sleeping_threads_use_no_cpu() {
    // Four threads asleep (plus us): every CPU should be idle nearly all
    // the time, because sleepers are off the run queue, not polling.
    let ids: Vec<_> = (0..4).map(|_| thread::spawn("s9", sleep_9)).collect();
    thread::sleep_ticks(1);
    let before = cpu_ticks();
    thread::sleep_ticks(6);
    let after = cpu_ticks();
    for id in ids {
        thread::join(id);
    }
    let total: u64 = (0..4).map(|c| after[c].0 - before[c].0).sum();
    let busy: u64 = (0..4).map(|c| after[c].1 - before[c].1).sum();
    assert!(total > 0);
    assert!(busy * 10 <= total, "{} of {} ticks busy while everyone slept", busy, total);
}

// --- wait queues ---------------------------------------------------------------

static TURN: AtomicU64 = AtomicU64::new(0);
static TURN_CHANGED: WaitQueue = WaitQueue::new();
const ROUNDS: u64 = 500;

/// Take a turn whenever TURN has parity `mine`, ROUNDS times.
fn take_turns(mine: u64) {
    for _ in 0..ROUNDS {
        TURN_CHANGED.wait_until(|| (TURN.load(Ordering::SeqCst) % 2 == mine).then_some(()));
        TURN.fetch_add(1, Ordering::SeqCst);
        TURN_CHANGED.wake_all();
    }
}
fn ping() {
    take_turns(0)
}
fn pong() {
    take_turns(1)
}

#[test_case]
fn wait_queue_ping_pong_between_cpus() {
    TURN.store(0, Ordering::SeqCst);
    let a = thread::spawn("ping", ping);
    let b = thread::spawn("pong", pong);
    thread::join(a);
    thread::join(b);
    assert_eq!(TURN.load(Ordering::SeqCst), 2 * ROUNDS);
}

// --- no leaks --------------------------------------------------------------------

fn nothing() {}

/// Spawn 8 do-nothing threads at once and join them.
fn spawn_and_join_8() {
    let ids: Vec<_> = (0..8).map(|_| thread::spawn("nop", nothing)).collect();
    for id in ids {
        thread::join(id);
    }
}

#[test_case]
fn spawning_and_exiting_threads_leaks_no_frames() {
    spawn_and_join_8(); // warm up: the first time may add stack page tables
    let frames = memory::frames_in_use();
    let threads = thread::thread_count();
    let stacks = stack::live_stacks();
    for _ in 0..25 {
        spawn_and_join_8();
    }
    assert_eq!(memory::frames_in_use(), frames, "frames leaked by 200 threads");
    assert_eq!(thread::thread_count(), threads);
    assert_eq!(stack::live_stacks(), stacks);
}

/// Run 4 `child 0` processes at once and wait for all of them.
fn run_4_children() {
    let ps: Vec<_> = (0..4).map(|_| run("child", &["0"])).collect();
    for p in ps {
        assert_eq!(process::wait(&p), ExitStatus::Exited(0));
    } // each Arc dropped here: the last reference, so the memory is freed
}

#[test_case]
fn spawning_and_exiting_processes_leaks_no_frames() {
    run_4_children(); // warm up (file caches, stack slots)
    let frames = memory::frames_in_use();
    let processes = process::count();
    for _ in 0..10 {
        run_4_children();
    }
    assert_eq!(process::count(), processes);
    assert_eq!(memory::frames_in_use(), frames, "frames leaked by 40 processes");
}

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