rust-kernel

//! threads [N] -- several threads in one process, pthread-style.
//!
//! Threads share *everything* a process owns: memory, open files, signal
//! handlers. Each has only its own registers, stack and TLS block (and a
//! thread id). ulib's `spawn_thread` = pthread_create: clone with
//! CLONE_VM | CLONE_FILES | CLONE_SIGHAND | CLONE_THREAD | ..., a fresh
//! stack, CLONE_SETTLS for the TLS block, and CLONE_CHILD_CLEARTID so
//! `join` can sleep on a futex until the kernel says "thread gone".
//!
//! The demo: a parallel sum. 1..=COUNT is split into N slices (default 4,
//! one per CPU); each thread adds up its slice, then adds its result to a
//! shared total under a futex mutex. Many small increments of a shared
//! counter (also under the mutex) check the mutex really excludes; a
//! counter bumped *without* it shows why that matters.
#![no_std]
#![no_main]

#[macro_use]
#[path = "ulib.rs"]
mod ulib;

extern crate alloc;

use alloc::sync::Arc;
use alloc::vec::Vec;
use core::cell::UnsafeCell;
use core::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use ulib::*;

const COUNT: u64 = 40_000_000;
const BUMPS: u64 = 20_000;

/// Data guarded by a FutexMutex (a tiny `Mutex<T>`).
struct Shared {
    lock: FutexMutex,
    total: UnsafeCell<u64>,
    bumps: UnsafeCell<u64>,
    /// Bumped with a plain read + write, no lock: races lose updates.
    racy: UnsafeCell<u64>,
}
unsafe impl Sync for Shared {}
unsafe impl Send for Shared {}

fn main() -> i64 {
    let n: u64 = arg(1).and_then(|a| a.parse().ok()).unwrap_or(4).clamp(1, 16);
    let mut c = Checks::new("threads");
    let shared = Arc::new(Shared {
        lock: FutexMutex::new(),
        total: UnsafeCell::new(0),
        bumps: UnsafeCell::new(0),
        racy: UnsafeCell::new(0),
    });
    let cpus_seen = Arc::new(AtomicU64::new(0)); // bit i = some thread ran on CPU i
    let tids = Arc::new(AtomicU64::new(0)); // sum of the threads' tids
    let tls_ok = Arc::new(AtomicU32::new(0));
    let start = now_us();
    let mut threads = Vec::new();
    for i in 0..n {
        let (shared, cpus_seen, tids, tls_ok) = (shared.clone(), cpus_seen.clone(), tids.clone(), tls_ok.clone());
        let from = COUNT * i / n + 1;
        let to = COUNT * (i + 1) / n;
        // Each thread gets `i` as its TLS value.
        let t = spawn_thread(i, move || {
            let mut sum = 0u64;
            for x in from..=to {
                sum = sum.wrapping_add(core::hint::black_box(x));
            }
            cpus_seen.fetch_or(1 << sysinfo().this_cpu, Ordering::Relaxed);
            tids.fetch_add(gettid() as u64, Ordering::Relaxed);
            if tls_value() == i {
                tls_ok.fetch_add(1, Ordering::Relaxed);
            }
            for _ in 0..BUMPS {
                shared.lock.lock();
                unsafe { *shared.bumps.get() += 1 };
                shared.lock.unlock();
                unsafe {
                    let v = core::ptr::read_volatile(shared.racy.get());
                    core::ptr::write_volatile(shared.racy.get(), v + 1);
                }
            }
            shared.lock.lock();
            unsafe { *shared.total.get() += sum };
            shared.lock.unlock();
        });
        match t {
            Ok(t) => threads.push(t),
            Err(e) => {
                println!("threads: spawn_thread: {}", strerror(e));
                return 1;
            }
        }
    }
    let spawned_tids: u64 = threads.iter().map(|t| t.tid as u64).sum();
    for t in threads {
        t.join();
    }
    let ms = (now_us() - start) / 1000;
    let total = unsafe { *shared.total.get() };
    let bumps = unsafe { *shared.bumps.get() };
    let racy = unsafe { *shared.racy.get() };
    let cpus = cpus_seen.load(Ordering::Relaxed).count_ones();
    println!("threads: {} threads summed 1..={} = {} in {} ms on {} CPU(s)", n, COUNT, total, ms, cpus);
    println!("threads: locked counter {} (want {}), unlocked counter {} (races lose updates)", bumps, n * BUMPS, racy);
    c.check(total == COUNT * (COUNT + 1) / 2, "parallel sum is right");
    c.check(bumps == n * BUMPS, "the futex mutex never let two threads in at once");
    c.check(tls_ok.load(Ordering::Relaxed) == n as u32, "every thread saw its own TLS value (fs:[8])");
    c.check(tids.load(Ordering::Relaxed) == spawned_tids, "gettid in each thread = the tid clone returned");
    c.check(gettid() == getpid(), "the main thread's tid is the pid");

    // A futex wait whose value doesn't match returns at once (EAGAIN); one
    // that matches sleeps until its timeout.
    let word = AtomicU32::new(7);
    c.check(futex_wait(&word, 8, None) == -EAGAIN, "futex_wait with a stale value: EAGAIN");
    let t0 = now_us();
    let r = futex_wait(&word, 7, Some(50_000));
    c.check(r == -ETIMEDOUT && now_us() - t0 >= 40_000, "futex_wait times out after 50 ms");
    c.check(futex_wake(&word, 1) == 0, "futex_wake with nobody waiting wakes 0");

    // exit vs exit_group: a thread that calls exit (60) ends only itself.
    let alive = Arc::new(AtomicU32::new(0));
    let a2 = alive.clone();
    let t = spawn_thread(0, move || {
        a2.store(1, Ordering::SeqCst);
        // falls off the end: ulib calls exit(0) = just this thread
    })
    .unwrap();
    t.join();
    c.check(alive.load(Ordering::SeqCst) == 1, "a thread's exit ends only that thread");

    // tgkill: a signal for one particular thread (here: ourselves).
    c.check(tgkill(getpid(), gettid(), 0) == 0, "tgkill(pid, tid, 0) finds the thread");
    c.check(tgkill(getpid(), 99999, 0) == -ESRCH, "tgkill to a missing thread: ESRCH");
    c.done()
}