yggdrasil/lib/runtime/src/process/thread.rs

//! Runtime utilities for thread handling

use core::{
    mem::MaybeUninit,
    ptr,
    sync::atomic::{AtomicU32, Ordering},
};

use abi::{
    error::Error,
    mem::{MappingFlags, MappingSource},
    process::{ThreadEvent, ThreadSpawnOptions},
};
use alloc::{boxed::Box, sync::Arc};

use crate::{process::thread_local, sync::rwlock::RwLock};

use super::{signal, thread_local::TlsImage};

pub mod options {
    //! Thread option definitions
    pub use abi::process::thread::*;
}

/// Describes a runtime thread.
///
/// `R` generic parameter denotes the thread's return type.
pub struct Thread<R> {
    id: AtomicU32,
    result: RwLock<Option<R>>,
}

/// Describes a "handle" to some runtime thread. This is what gets returned to the caller
/// when a thread is spawned and can be used to, e.g., join the created thread later.
pub struct ThreadHandle<R> {
    #[allow(unused)]
    stack: Option<OwnedStack>,
    thread: Arc<Thread<R>>,
}

struct OwnedStack {
    base: usize,
    size: usize,
}

/// Describes how a thread's signal stack should be set up.
pub enum ThreadSignalStack {
    /// Thread signal stack will be allocated with the requested size.
    Allocate(usize),
    /// Do not set up a signal stack.
    None,
}

/// Describes how a thread's main stack should be set up.
pub enum ThreadStack {
    /// Thread stack will be allocated with the requested size.
    Allocate(usize),
    /// A raw value will be used for the stack pointer (unsafe).
    Raw(usize),
}

/// Describes the thread's function that will get entered once the thread starts running.
///
/// `A` parameter denotes the argument type that will be passed to the thread.
pub enum ThreadFunction<A, R> {
    /// A boxed closure will be called.
    Closure(Box<dyn FnOnce(A) -> R + 'static>),
    /// A function pointer will be called.
    Pointer(fn(A) -> R),
}

/// Describes how a thread should be constructed
pub struct ThreadCreateInfo<A, R> {
    /// See [ThreadSignalStack].
    pub signal_stack: ThreadSignalStack,
    /// See [ThreadStack].
    pub stack: ThreadStack,
    /// See [ThreadFunction].
    pub entry: ThreadFunction<A, R>,
    /// Image used to initialize the thread's TLS
    pub tls_image: Option<&'static TlsImage>,
}

struct ThreadArgument<A, R> {
    entry: Option<(A, ThreadFunction<A, R>)>,
    signal_stack: ThreadSignalStack,
    thread_self: Arc<Thread<R>>,
    set_thread_self: bool,
    tls_image: Option<&'static TlsImage>,
}

// TODO maybe put this under a `thread-self` feature and avoid a TLS allocation?
#[thread_local]
static mut SELF: usize = 0;

impl<R> Thread<R> {
    /// Creates a new thread with the requested options and passes it an argument.
    pub fn spawn<A>(
        info: ThreadCreateInfo<A, R>,
        argument: A,
        runtime: bool,
    ) -> Result<ThreadHandle<R>, Error> {
        let (stack, sp) = match info.stack {
            ThreadStack::Allocate(size) => {
                let (stack, sp) = OwnedStack::allocate(size)?;
                (Some(stack), sp)
            }
            ThreadStack::Raw(_) => todo!(),
        };

        let thread = Arc::new(Thread::<R> {
            id: AtomicU32::new(0),
            result: RwLock::new(None),
        });

        let thread_argument = Box::into_raw(Box::new(ThreadArgument {
            entry: Some((argument, info.entry)),
            thread_self: thread.clone(),
            signal_stack: info.signal_stack,
            set_thread_self: runtime,
            tls_image: info.tls_image,
        }))
        .addr();

        let spawn_options = ThreadSpawnOptions {
            entry: Self::thread_entry::<A>,
            argument: thread_argument,
            stack_top: sp,
        };

        let id = unsafe { crate::sys::spawn_thread(&spawn_options) }?;
        thread.id.store(id, Ordering::Release);

        Ok(ThreadHandle { stack, thread })
    }

    /// Returns this thread's ID.
    pub fn id(&self) -> u32 {
        self.id.load(Ordering::Acquire)
    }

    /// Sets the current thread name.
    pub fn set_name(name: &str) {
        let mut buffer = [0; 256];
        crate::process::set_thread_option_with::<options::Name>(&mut buffer, &name).ok();
    }

    /// # Safety
    ///
    /// This function has to be called with the same `Self` type as the thread it was created with.
    pub unsafe fn current() -> Arc<Self> {
        let raw: *const Self = ptr::with_exposed_provenance(SELF);
        if raw.is_null() {
            panic!("Thread::SELF == NULL, was spawn() called with runtime=false?");
        }
        // This will "move" the pointer, so an extra strong count increment is required.
        Arc::increment_strong_count(raw);
        Arc::from_raw(raw)
    }

    extern "C" fn thread_entry<A>(raw: usize) -> ! {
        let raw: *mut ThreadArgument<A, R> = ptr::with_exposed_provenance_mut(raw);
        // This scope will ensure all the stuff is dropped before thread exit is called.
        {
            let mut argument = unsafe { Box::from_raw(raw) };

            // Setup TLS as soon as possible
            if let Err(err) = thread_local::init_tls(argument.tls_image, true) {
                crate::debug_trace!(Warn, "thread_entry failed: TLS init error: {err:?}");
                // TODO result is uninit
                unsafe { crate::sys::exit_thread() };
            }

            // TODO there will be a better way to do this, I promise
            while argument.thread_self.id.load(Ordering::Acquire) == 0 {
                core::hint::spin_loop();
            }

            // Insert compiler fence just in case to prevent it from trying to go into
            // (uninitialized) TLS
            core::sync::atomic::compiler_fence(Ordering::Release);

            // Setup SELF if needed
            if argument.set_thread_self {
                unsafe {
                    SELF = Arc::into_raw(argument.thread_self.clone()).addr();
                    debug_assert!(Arc::ptr_eq(&Thread::current(), &argument.thread_self));
                }
            }

            // Setup signal stack if needed
            let result = match argument.signal_stack {
                ThreadSignalStack::None => Ok(()),
                ThreadSignalStack::Allocate(size) => signal::setup_signal_stack(size),
            };
            if let Err(err) = result {
                panic!("Failed to set up thread's signal stack: {err:?}");
            }

            // Call the inner function
            let result = match argument.entry.take() {
                Some((arg, ThreadFunction::Closure(f))) => f(arg),
                Some((arg, ThreadFunction::Pointer(f))) => f(arg),
                None => unreachable!(),
            };

            argument.thread_self.result.write().replace(result);
        }

        unsafe { crate::sys::exit_thread() };
    }
}

impl<R> ThreadHandle<R> {
    /// Returns the thread ID this handle is related to.
    pub fn id(&self) -> u32 {
        self.thread.id.load(Ordering::Acquire)
    }

    /// Waits for the thread to finish and returns its result.
    pub fn join(self) -> Result<Option<R>, Error> {
        // TODO prevent threads from attempting to join themselves?
        loop {
            let mut event = MaybeUninit::uninit();
            unsafe { crate::sys::wait_thread(self.thread.id.load(Ordering::Acquire), &mut event) }?;
            let event = unsafe { event.assume_init() };
            if matches!(event, ThreadEvent::Exited) {
                break;
            }
        }
        Ok(self.into_result())
    }

    /// Waits for the thread to finish, ignoring interrupts.
    ///
    /// # Panics
    ///
    /// Will panic if the kernel returns any error besides [Error::Interrupted].
    pub fn join_uninterruptible(self) -> Option<R> {
        loop {
            let mut event = MaybeUninit::uninit();
            match unsafe {
                crate::sys::wait_thread(self.thread.id.load(Ordering::Acquire), &mut event)
            } {
                Ok(_) => {
                    let event = unsafe { event.assume_init() };
                    if !matches!(event, ThreadEvent::Exited) {
                        continue;
                    }
                }
                Err(Error::Interrupted) => continue,
                Err(Error::ProcessNotFound) => (),
                Err(error) => panic!("wait_thread syscall returned error: {error:?}"),
            }
            return self.into_result();
        }
    }

    fn into_result(self) -> Option<R> {
        self.thread.result.write().take()
    }
}

impl OwnedStack {
    fn allocate(mut size: usize) -> Result<(Self, usize), Error> {
        if size < 0x1000 {
            size = 0x1000;
        }
        size = (size + 0xFFF) & !0xFFF;
        let base = unsafe {
            crate::sys::map_memory(None, size, MappingFlags::WRITE, &MappingSource::Anonymous)
        }?;
        let top = base + size;

        Ok((Self { base, size }, top))
    }
}

impl Drop for OwnedStack {
    fn drop(&mut self) {
        crate::debug_trace!(
            Debug,
            "Drop OwnedStack {:#x?}",
            self.base..self.base + self.size
        );
        unsafe { crate::sys::unmap_memory(self.base, self.size) }.ok();
    }
}

unsafe impl<R> Send for Thread<R> {}
unsafe impl<R> Sync for Thread<R> {}

/// Sets the `SELF` TLS variable
///
/// # Safety
///
/// Must only be called once by the runtime library during setup. TLS must be initialized prior to
/// this call.
pub unsafe fn track_main<R>() {
    let main = Arc::new(Thread::<R> {
        id: AtomicU32::new(0),
        result: RwLock::new(None),
    });
    SELF = Arc::into_raw(main).addr();
}