feat: thread-local storage (TLS)

Thread-local storage is commonly needed in all sorts of places, from storing a harts `TrapFrame`, ID or other state that should be statically available but still per-hart.

Instead of coming up with something homegrown here we use Rusts/LLVMs [`thread_local` attribute](rust-lang/rust#29594).

This requires us to do a couple things:
1. allocate one or more frames for TLS
2. (identity) map the frames
3. *on each hart* set the `tp` register to point to the TLS base

.cargo/config.toml

@@ -3,4 +3,4 @@ target = "riscv64gc-unknown-none-elf"
 rustflags = ["-C", "force-unwind-tables=true"]
 
 [target.riscv64gc-unknown-none-elf]
-runner = "qemu-system-riscv64 -machine virt -cpu rv64 -d guest_errors,unimp -smp 1 -m 128M -nographic -serial mon:stdio -bios default -device virtio-rng-device -device virtio-gpu-device -device virtio-net-device -device virtio-tablet-device -device virtio-keyboard-device -kernel "
+runner = "qemu-system-riscv64 -machine virt -cpu rv64 -d guest_errors,unimp -smp 4 -m 128M -nographic -serial mon:stdio -bios default -device virtio-rng-device -device virtio-gpu-device -device virtio-net-device -device virtio-tablet-device -device virtio-keyboard-device -kernel "

crates/kernel/riscv64-virt.ld

@@ -39,6 +39,16 @@ SECTIONS {
         __bss_end = .;
     }
 
+    .tls : {
+        __tdata_start = .;
+        *(.tdata .tdata.*)
+        __tdata_end = .;
+
+        __tbss_start = .;
+        *(.tbss .tbss.*)
+        __tbss_end = .;
+    }
+
     .eh_frame : {
         . = ALIGN(8);
         __eh_frame_start = .;

crates/kernel/src/arch/mod.rs

@@ -4,6 +4,7 @@ pub mod backtrace;
 pub mod interrupt;
 pub mod paging;
 mod start;
+pub mod tls;
 pub mod trap;
 
 pub const PAGE_SIZE: usize = 4096;

crates/kernel/src/arch/paging/mapper.rs

@@ -56,6 +56,10 @@ impl Mapper {
         &self.allocator
     }
 
+    pub fn allocator_mut(&mut self) -> &mut FrameAllocator {
+        &mut self.allocator
+    }
+
     pub fn root_table(&self) -> Table {
         Table::from_address(self.root_table, MAX_LEVEL)
     }

crates/kernel/src/arch/paging/mod.rs

@@ -4,12 +4,11 @@ use crate::arch::paging::mapper::Mapper;
 use crate::board_info::BoardInfo;
 use crate::paging::frame_alloc::FrameAllocator;
 use crate::paging::PhysicalAddress;
+use crate::sync::Mutex;
 use core::ops::Range;
 use core::ptr::addr_of;
-use riscv::register::satp;
-use riscv::register::satp::Mode;
 
-mod entry;
+pub mod entry;
 mod flush;
 mod mapper;
 mod table;
@@ -24,12 +23,27 @@ extern "C" {
     static __rodata_end: u8;
 }
 
+pub static mut MAPPER: Mutex<Option<Mapper>> = Mutex::new(None);
+
+pub fn activate() -> crate::Result<()> {
+    let mut mapper = unsafe { MAPPER.lock() };
+    let mapper = mapper.as_mut().unwrap();
+
+    // mapper.root_table().print_table();
+
+    log::debug!("activating paging...");
+    mapper.activate()?;
+    log::debug!("paging activated");
+
+    Ok(())
+}
+
 /// Initialize virtual memory management
 ///
 /// This will set up the page table, identity map the kernel and stack, and enable paging.
 ///
 /// TODO lift this out of the arch module
-pub fn init(board_info: &BoardInfo) -> crate::Result<()> {
+pub fn setup(board_info: &BoardInfo) -> crate::Result<()> {
     let stack_start = unsafe { addr_of!(__stack_start) as usize };
     let text_start = unsafe { addr_of!(__text_start) as usize };
     let text_end = unsafe { addr_of!(__text_end) as usize };
@@ -55,7 +69,7 @@ pub fn init(board_info: &BoardInfo) -> crate::Result<()> {
     // Step 2: initialize allocator
     let allocator = unsafe { FrameAllocator::new(&regions) };
 
-    // Step 4: create mapper
+    // Step 3: create mapper
     let mut mapper = Mapper::new(allocator)?;
 
     // helper function to identity map a region
@@ -83,24 +97,21 @@ pub fn init(board_info: &BoardInfo) -> crate::Result<()> {
         Ok(())
     };
 
-    // Step 4: map kernel
+    // Step 5: map kernel
     log::debug!("mapping kernel region: {:?}", kernel_region);
     identity_map_range(kernel_region.clone())?;
 
-    // Step 5: map stack
+    // Step 6: map stack
     log::debug!("mapping stack region: {:?}", stack_region);
     identity_map_range(stack_region.clone())?;
 
-    // Step 6: map MMIO regions (UART)
+    // Step 7: map MMIO regions (UART)
     log::debug!("mapping mmio region: {:?}", board_info.serial.mmio_regs);
     identity_map_range(align_range(board_info.serial.mmio_regs.clone()))?;
 
-    mapper.root_table().print_table();
-
-    // Step 7: enable paging
-    log::debug!("enabling paging... {:?}", mapper.root_table().address());
-    mapper.activate()?;
-    log::debug!("paging enabled");
+    unsafe {
+        MAPPER.lock().replace(mapper);
+    }
 
     Ok(())
 }

crates/kernel/src/arch/start.rs

@@ -1,6 +1,7 @@
 use crate::arch;
 use crate::board_info::BoardInfo;
 use core::arch::asm;
+use core::ops::Range;
 use core::ptr::addr_of_mut;
 
 /// Sets the harts stack pointer to the top of the stack.
@@ -102,21 +103,24 @@ extern "C" fn start(hartid: usize, opaque: *const u8) -> ! {
     extern "C" {
         static mut __bss_start: u64;
         static mut __bss_end: u64;
+        static mut __tbss_start: u64;
+        static mut __tbss_end: u64;
     }
+
     unsafe {
-        let mut ptr = addr_of_mut!(__bss_start);
-        let end = addr_of_mut!(__bss_end);
-        while ptr < end {
-            ptr.write_volatile(0);
-            ptr = ptr.offset(1);
-        }
+        // zero out the BSS section
+        zero_range(addr_of_mut!(__bss_start)..addr_of_mut!(__bss_end));
+        // zero out the TBSS section (thread local BSS)
+        zero_range(addr_of_mut!(__tbss_start)..addr_of_mut!(__tbss_end));
     }
 
     let board_info = BoardInfo::from_raw(opaque).unwrap();
 
     crate::logger::init(&board_info, 38400).unwrap();
 
-    arch::paging::init(&board_info).unwrap();
+    arch::paging::setup(&board_info).unwrap();
+
+    arch::tls::setup(&board_info).unwrap();
 
     for hart in 0..board_info.cpus {
         if hart != hartid {
@@ -126,3 +130,18 @@ extern "C" fn start(hartid: usize, opaque: *const u8) -> ! {
 
     crate::kmain(hartid)
 }
+
+/// Utility function to zero out a range of memory.
+/// This is used to zero out the BSS and TBSS sections.
+///
+/// We explicitly use `*mut u64` here so that we generate 64-bit stores (`sd` instructions) even
+/// on 32-bit systems.
+fn zero_range(range: Range<*mut u64>) {
+    let mut ptr = range.start;
+    while ptr < range.end {
+        unsafe {
+            ptr.write_volatile(0);
+        }
+        ptr = unsafe { ptr.offset(1) };
+    }
+}

crates/kernel/src/arch/tls.rs

@@ -0,0 +1,60 @@
+use crate::arch;
+use crate::arch::paging::{entry::PageFlags, MAPPER};
+use crate::board_info::BoardInfo;
+use crate::paging::PhysicalAddress;
+use core::arch::asm;
+use core::ptr::addr_of;
+
+// Example of TLS usage:
+// #[thread_local]
+// static mut HART_ID: usize = 1;
+
+extern "C" {
+    static __tdata_start: u8;
+    static __tbss_end: u8;
+}
+
+static mut TLS_BASE: Option<PhysicalAddress> = None;
+
+pub fn setup(board_info: &BoardInfo) -> crate::Result<()> {
+    let tls_start = unsafe { addr_of!(__tdata_start) as usize };
+    let tls_end = unsafe { addr_of!(__tbss_end) as usize };
+
+    let tls_region = unsafe { PhysicalAddress::new(tls_start)..PhysicalAddress::new(tls_end) };
+
+    if !tls_region.is_empty() {
+        let mut mapper = unsafe { MAPPER.lock() };
+        let mapper = mapper.as_mut().unwrap();
+
+        let num_pages = (tls_region.end.as_raw() - tls_region.start.as_raw())
+            .div_ceil(arch::PAGE_SIZE)
+            * board_info.cpus;
+        let start = mapper.allocator_mut().allocate_frames(num_pages)?;
+
+        log::debug!(
+            "mapping TLS region: {:?}",
+            start..start.add(num_pages * arch::PAGE_SIZE)
+        );
+
+        for i in 0..num_pages {
+            let phys = start.add(i * arch::PAGE_SIZE);
+            let flush = mapper.map_identity(phys, PageFlags::READ | PageFlags::WRITE)?;
+            unsafe {
+                flush.ignore();
+            }
+        }
+
+        unsafe {
+            TLS_BASE = Some(start);
+        }
+    }
+
+    Ok(())
+}
+
+pub fn activate() {
+    unsafe {
+        let base = TLS_BASE.unwrap();
+        asm!("mv tp, {0}", in(reg) base.as_raw());
+    }
+}

crates/kernel/src/main.rs

@@ -1,7 +1,8 @@
 #![no_std]
 #![no_main]
-#![feature(naked_functions, asm_const, error_in_core)]
+#![feature(naked_functions, asm_const, error_in_core, thread_local)]
 
+// mod allocator;
 mod arch;
 mod backtrace;
 mod board_info;
@@ -17,14 +18,14 @@ pub type Result<T> = core::result::Result<T, Error>;
 /// This is the main function of the kernel.
 ///
 /// After performing arch & board specific initialization, all harts will end up in this function.
-/// This function should set up hart-local state, and then ?. It should never return.
+/// This function should set up hart-local state, perform arch-agnostic setup and then ?. It should never return.
 pub fn kmain(hartid: usize) -> ! {
-    // arch-agnostic initialization
-    // per-hart initialization
+    arch::tls::activate();
+    arch::paging::activate().unwrap();
 
     log::info!("Hello world from hart {hartid}!");
 
     arch::trap::init().unwrap();
 
-    todo!()
+    arch::halt();
 }

crates/kernel/src/paging/frame_alloc.rs

@@ -1,5 +1,6 @@
 use crate::arch::PAGE_SIZE;
 use crate::paging::PhysicalAddress;
+use crate::Error;
 use core::ops::Range;
 
 pub struct FrameAllocator {
@@ -20,6 +21,32 @@ impl FrameAllocator {
             .ok_or(crate::Error::OutOfMemory)
     }
 
+    pub fn allocate_frames(&mut self, n: usize) -> crate::Result<PhysicalAddress> {
+        let mut current = self.free_frame_list.head.as_ref();
+        let mut first_frame = current;
+        let mut remaining = n;
+
+        while let Some(frame) = current {
+            if remaining == 0 {
+                return Ok(first_frame.unwrap().as_ptr());
+            } else if let Some(next) = frame.next.as_ref() {
+                // check if the next frame is consecutive
+                if next.as_ptr().as_raw() == frame.as_ptr().as_raw() + PAGE_SIZE {
+                    remaining -= 1;
+                    current = Some(next);
+                } else {
+                    // we found a free frame but its in a different region
+                    // so we need to start over
+                    remaining = n;
+                    first_frame = current;
+                }
+            }
+        }
+
+        // we didn't find enough free frames
+        Err(Error::OutOfMemory)
+    }
+
     pub fn deallocate_frame(&mut self, phys: PhysicalAddress) {
         let ptr = phys.0 as *mut FreeFrame;
         unsafe {