Repository: theseus-os/Theseus (language: rust)

Issue:
Upcoming `tlb_shootdown` improvements
Note to ourselves so we don't forget about these easy improvements after #913 lands

1. `let mut old_lock_val = TLB_SHOOTDOWN_IPI_LOCK.load(Ordering::Relaxed);` should just be `let mut old_lock_val = false;`
2. `TLB_SHOOTDOWN_IPI_COUNT.store(cpu_count - 1, Ordering::SeqCst);` doesn't need sequential consistency
3. Make the following non-`pub`, and make `handle_tlb_shootdown_ipi` accept no parameters and access the `TLB_SHOOTDOWN_IPI_PAGES` directly within it
    * `pub static TLB_SHOOTDOWN_IPI_COUNT`
    * `pub static TLB_SHOOTDOWN_IPI_LOCK`
    * `pub static TLB_SHOOTDOWN_IPI_PAGES`
4. [Optional] `TLB_SHOOTDOWN_IPI_PAGES` doesn't need to be a `RwLockIrqSafe`, because `RwLock`s are very slow. Instead, it can probably just be a `crossbeam_utils::AtomicCell<PageRange>`, assuming that `PageRange` satisfies the const `is_lock_free()` condition (if not, we can use a type wrapper like for `Option<u8>`)




Discussion / hints:

Fix patch:
diff --git a/kernel/exceptions_full/src/lib.rs b/kernel/exceptions_full/src/lib.rs
index a43898cb54..c456e6a3d2 100644
--- a/kernel/exceptions_full/src/lib.rs
+++ b/kernel/exceptions_full/src/lib.rs
@@ -263,7 +263,9 @@ extern "x86-interrupt" fn debug_handler(stack_frame: InterruptStackFrame) {
     // don't halt here, this isn't a fatal/permanent failure, just a brief pause.
 }
 
-/// exception 0x02, also used for TLB Shootdown IPIs and sampling interrupts.
+/// Exception 0x02 is a Non-Maskable Interrupt (NMI).
+///
+/// Theseus uses this for TLB Shootdown IPIs and sampling interrupts.
 ///
 /// # Important Note
 /// Acquiring ANY locks in this function, even irq-safe ones, could cause a deadlock
@@ -271,16 +273,11 @@ extern "x86-interrupt" fn debug_handler(stack_frame: InterruptStackFrame) {
 /// another regular interrupt. 
 /// This includes printing to the log (e.g., `debug!()`) or the screen.
 extern "x86-interrupt" fn nmi_handler(stack_frame: InterruptStackFrame) {
+    // trace!("nmi_handler (CPU {})", cpu::current_cpu());
     let mut expected_nmi = false;
 
-    // currently we're using NMIs to send TLB shootdown IPIs
-    {
-        let pages_to_invalidate = tlb_shootdown::TLB_SHOOTDOWN_IPI_PAGES.read().clone();
-        if let Some(pages) = pages_to_invalidate {
-            // trace!("nmi_handler (AP {})", cpu::current_cpu());
-            tlb_shootdown::handle_tlb_shootdown_ipi(pages);
-            expected_nmi = true;
-        }
+    if tlb_shootdown::handle_tlb_shootdown_ipi() {
+        return;
     }
 
     // Performance monitoring hardware uses NMIs to trigger a sampling interrupt.
diff --git a/kernel/memory/src/lib.rs b/kernel/memory/src/lib.rs
index 0be49a9ac3..cb4cddb030 100644
--- a/kernel/memory/src/lib.rs
+++ b/kernel/memory/src/lib.rs
@@ -22,7 +22,7 @@ pub use self::paging::{
     translate,
 };
 
-pub use memory_structs::{Frame, Page, FrameRange, PageRange, VirtualAddress, PhysicalAddress};
+pub use memory_structs::*;
 pub use page_allocator::{
     AllocatedPages, allocate_pages, allocate_pages_at,
     allocate_pages_by_bytes, allocate_pages_by_bytes_at,
diff --git a/kernel/memory_structs/src/lib.rs b/kernel/memory_structs/src/lib.rs
index f82855165e..7080dfb959 100644
--- a/kernel/memory_structs/src/lib.rs
+++ b/kernel/memory_structs/src/lib.rs
@@ -469,6 +469,24 @@ macro_rules! implement_page_frame_range {
                     self.0
                 }
             }
+
+            
+            #[doc = "A `" $TypeName "` that implements `Copy`"]
+            #[derive(Clone, Copy)]
+            pub struct [<Copyable $TypeName>] {
+                start: $chunk,
+                end: $chunk,
+            }
+            impl From<$TypeName> for [<Copyable $TypeName>] {
+                fn from(r: $TypeName) -> Self {
+                    Self { start: *r.start(), end: *r.end() }
+                }
+            }
+            impl From<[<Copyable $TypeName>]> for $TypeName {
+                fn from(cr: [<Copyable $TypeName>]) -> Self {
+                    Self::new(cr.start, cr.end)
+                }
+            }
         }
     };
 }
diff --git a/kernel/memory_x86_64/src/lib.rs b/kernel/memory_x86_64/src/lib.rs
index adf0c4626f..8ecf86cb6a 100644
--- a/kernel/memory_x86_64/src/lib.rs
+++ b/kernel/memory_x86_64/src/lib.rs
@@ -265,7 +265,7 @@ pub fn get_vga_mem_addr(
 
 /// Flushes the specific virtual address in TLB. 
 pub fn tlb_flush_virt_addr(vaddr: VirtualAddress) {
-    tlb::flush(x86_64::VirtAddr::new(vaddr.value() as u64));
+    tlb::flush(x86_64::VirtAddr::new_truncate(vaddr.value() as u64));
 }
 
 /// Flushes the whole TLB. 
diff --git a/kernel/tlb_shootdown/Cargo.toml b/kernel/tlb_shootdown/Cargo.toml
index 3728d1f539..ac0ea34dec 100644
--- a/kernel/tlb_shootdown/Cargo.toml
+++ b/kernel/tlb_shootdown/Cargo.toml
@@ -18,6 +18,3 @@ apic = { path = "../apic" }
 [target.'cfg(target_arch = "aarch64")'.dependencies]
 memory_aarch64 = { path = "../memory_aarch64" }
 interrupts = { path = "../interrupts" }
-
-[lib]
-crate-type = ["rlib"]
diff --git a/kernel/tlb_shootdown/src/lib.rs b/kernel/tlb_shootdown/src/lib.rs
index 87ae397dcf..c4b1efafe5 100644
--- a/kernel/tlb_shootdown/src/lib.rs
+++ b/kernel/tlb_shootdown/src/lib.rs
@@ -14,16 +14,15 @@ use memory_x86_64::tlb_flush_virt_addr;
 #[cfg(target_arch = "aarch64")]
 use memory_aarch64::tlb_flush_virt_addr;
 
-/// The number of remaining cores that still need to handle the current TLB shootdown IPI
-pub static TLB_SHOOTDOWN_IPI_COUNT: AtomicU32 = AtomicU32::new(0);
-/// The lock that makes sure only one set of TLB shootdown IPIs is concurrently happening
-pub static TLB_SHOOTDOWN_IPI_LOCK: AtomicBool = AtomicBool::new(false);
-/// The range of pages for a TLB shootdown IPI.
-pub static TLB_SHOOTDOWN_IPI_PAGES: RwLockIrqSafe<Option<PageRange>> = RwLockIrqSafe::new(None);
+/// The number of remaining CPUs that still need to handle the current TLB shootdown IPI.
+static TLB_SHOOTDOWN_IPI_COUNT: AtomicU32 = AtomicU32::new(0);
+/// This lock ensures only one round of TLB shootdown IPIs can occur concurrently.
+static TLB_SHOOTDOWN_IPI_LOCK: AtomicBool = AtomicBool::new(false);
+/// The range of virtual pages to be flushed for a TLB shootdown IPI.
+static TLB_SHOOTDOWN_IPI_PAGES: RwLockIrqSafe<Option<PageRange>> = RwLockIrqSafe::new(None);
 
 
 /// Initializes data, functions, and structures for the TLB shootdown. 
-/// TODO: redesign this, it's weird and silly just to set one callback.
 pub fn init() {
     memory::set_broadcast_tlb_shootdown_cb(broadcast_tlb_shootdown);
 
@@ -31,28 +30,31 @@ pub fn init() {
     interrupts::setup_ipi_handler(tlb_shootdown_ipi_handler, interrupts::TLB_SHOOTDOWN_IPI).unwrap();
 }
 
-/// Handles a TLB shootdown ipi by flushing the `VirtualAddress`es 
-/// covered by the given range of `pages_to_invalidate`.
-/// 
+/// Handles a TLB shootdown IPI requested by another CPU.
+///
 /// There is no need to invoke this directly, it will be called by an IPI interrupt handler.
-pub fn handle_tlb_shootdown_ipi(pages_to_invalidate: PageRange) {
-    // log::trace!("handle_tlb_shootdown_ipi(): AP {}, pages: {:?}", apic::current_cpu(), pages_to_invalidate);
-
-    for page in pages_to_invalidate {
-        tlb_flush_virt_addr(page.start_address());
+///
+/// ## Return
+/// Returns `true` if virtual addresses were actually flushed, `false` otherwise.
+pub fn handle_tlb_shootdown_ipi() -> bool {
+    let pages_to_invalidate = TLB_SHOOTDOWN_IPI_PAGES.read().clone();
+    if let Some(pages) = pages_to_invalidate {
+        // log::trace!("handle_tlb_shootdown_ipi(): CPU {}, pages: {:?}", apic::current_cpu(), pages);
+        for page in pages {
+            tlb_flush_virt_addr(page.start_address());
+        }
+        TLB_SHOOTDOWN_IPI_COUNT.fetch_sub(1, Ordering::Relaxed);
+        true
+    } else {
+        false
     }
-
-    TLB_SHOOTDOWN_IPI_COUNT.fetch_sub(1, Ordering::SeqCst);
 }
 
 
-/// Broadcasts TLB shootdown IPI to all other AP cores.
-///
-/// Do not invoke this directly, but rather pass it as a callback to the memory subsystem,
-/// which will invoke it as needed (on remap/unmap operations).
+/// Broadcasts a TLB shootdown IPI to all other CPUs, causing them to flush (invalidate)
+/// the given virtual pages in their TLBs.
 ///
-/// Sends an IPI to all other cores (except me) to trigger 
-/// a TLB flush of the given pages' virtual addresses.
+/// This is invoked by the memory subsystem as needed, e.g., on remap/unmap operations.
 fn broadcast_tlb_shootdown(pages_to_invalidate: PageRange) {        
     // skip sending IPIs if there are no other cores running
     let cpu_count = cpu_count();
@@ -70,17 +72,20 @@ fn broadcast_tlb_shootdown(pages_to_invalidate: PageRange) {
 
     // acquire lock
     // TODO: add timeout!!
-    let mut old_lock_val = TLB_SHOOTDOWN_IPI_LOCK.load(Ordering::Relaxed);
     loop {
-        match TLB_SHOOTDOWN_IPI_LOCK.compare_exchange_weak(old_lock_val, true, Ordering::AcqRel, Ordering::Relaxed) { 
-            Ok(_) => break,
-            Err(v) => old_lock_val = v,
+        if TLB_SHOOTDOWN_IPI_LOCK.compare_exchange_weak(
+            false,
+            true,
+            Ordering::AcqRel,
+            Ordering::Relaxed,
+        ).is_ok() {
+            break;
         }
         spin_loop();
     }
 
     *TLB_SHOOTDOWN_IPI_PAGES.write() = Some(pages_to_invalidate);
-    TLB_SHOOTDOWN_IPI_COUNT.store(cpu_count - 1, Ordering::SeqCst); // -1 to exclude this core 
+    TLB_SHOOTDOWN_IPI_COUNT.store(cpu_count - 1, Ordering::Relaxed); // -1 to exclude this core 
 
     #[cfg(target_arch = "x86_64")] {
         let my_lapic = apic::get_my_apic()
@@ -110,12 +115,7 @@ fn broadcast_tlb_shootdown(pages_to_invalidate: PageRange) {
 /// Interrupt Handler for TLB Shootdowns on aarch64
 #[cfg(target_arch = "aarch64")]
 extern "C" fn tlb_shootdown_ipi_handler(_exc: &interrupts::ExceptionContext) -> interrupts::EoiBehaviour {
-    if let Some(pages_to_invalidate) = TLB_SHOOTDOWN_IPI_PAGES.read().clone() {
-        // trace!("nmi_handler (AP {})", cpu::current_cpu());
-        handle_tlb_shootdown_ipi(pages_to_invalidate);
-    } else {
-        panic!("Unexpected TLB Shootdown IPI!");
-    }
-
+    let expected = handle_tlb_shootdown_ipi();
+    assert!(expected, "Unexpected TLB Shootdown IPI!");
     interrupts::EoiBehaviour::CallerMustSignalEoi
 }


Test patch: