x86_64: impl IPI and remote flush tlb

Signed-off-by: WangRunji <wangrunji0408@163.com>

kernel/src/arch/x86_64/gdt.rs

@@ -1,11 +1,14 @@
+use super::ipi::IPIEventItem;
 use alloc::boxed::Box;
-
+use alloc::vec::*;
+use core::sync::atomic::{AtomicBool, Ordering};
 use x86_64::registers::model_specific::Msr;
 use x86_64::structures::gdt::*;
 use x86_64::structures::tss::TaskStateSegment;
 use x86_64::{PrivilegeLevel, VirtAddr};
 
 use crate::consts::MAX_CPU_NUM;
+use crate::sync::{Semaphore, SpinLock as Mutex};
 
 /// Init TSS & GDT.
 pub fn init() {
@@ -20,22 +23,16 @@ static mut CPUS: [Option<Cpu>; MAX_CPU_NUM] = [
     None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None,
     None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None,
     None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None,
-    None, None, None, None, None, None, None, None, None, None, None, None, None, None, None,
+    None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None,
-    None,
-    //    None, None, None, None, None, None, None, None,
-    //    None, None, None, None, None, None, None, None,
-    //    None, None, None, None, None, None, None, None,
-    //    None, None, None, None, None, None, None, None,
-    //    None, None, None, None, None, None, None, None,
-    //    None, None, None, None, None, None, None, None,
-    //    None, None, None, None, None, None, None, None,
-    //    None, None, None, None, None, None, None, None,
 ];
 
 pub struct Cpu {
     gdt: GlobalDescriptorTable,
     tss: TaskStateSegment,
     double_fault_stack: [u8; 0x100],
+    preemption_disabled: AtomicBool, // TODO: check this on timer(). This is currently unavailable since related code is in rcore_thread.
+    ipi_handler_queue: Mutex<Vec<IPIEventItem>>,
+    id: usize,
 }
 
 impl Cpu {
@@ -44,9 +41,47 @@ impl Cpu {
             gdt: GlobalDescriptorTable::new(),
             tss: TaskStateSegment::new(),
             double_fault_stack: [0u8; 0x100],
+            preemption_disabled: AtomicBool::new(false),
+            ipi_handler_queue: Mutex::new(vec![]),
+            id: 0,
         }
     }
 
+    pub fn foreach(mut f: impl FnMut(&mut Cpu)) {
+        unsafe {
+            CPUS.iter_mut()
+                .filter_map(|x| x.as_mut())
+                .for_each(f);
+        }
+    }
+    pub fn get_id(&self) -> usize {
+        self.id
+    }
+    pub fn notify_event(&mut self, item: IPIEventItem) {
+        let mut queue = self.ipi_handler_queue.lock();
+        queue.push(item);
+    }
+    pub fn current() -> &'static mut Cpu {
+        unsafe { CPUS[super::cpu::id()].as_mut().unwrap() }
+    }
+    pub fn ipi_handler(&mut self) {
+        let mut queue = self.ipi_handler_queue.lock();
+        let mut current_events: Vec<IPIEventItem> = vec![];
+        ::core::mem::swap(&mut current_events, queue.as_mut());
+        drop(queue);
+        for ev in current_events.iter() {
+            ev.call();
+        }
+    }
+    pub fn disable_preemption(&self) -> bool {
+        self.preemption_disabled.swap(true, Ordering::Relaxed)
+    }
+    pub fn restore_preemption(&self, val: bool) {
+        self.preemption_disabled.store(val, Ordering::Relaxed);
+    }
+    pub fn can_preempt(&self) -> bool {
+        self.preemption_disabled.load(Ordering::Relaxed)
+    }
     unsafe fn init(&'static mut self) {
         use x86_64::instructions::segmentation::{load_fs, set_cs};
         use x86_64::instructions::tables::load_tss;
@@ -63,7 +98,7 @@ impl Cpu {
         self.gdt.add_entry(UCODE);
         self.gdt.add_entry(Descriptor::tss_segment(&self.tss));
         self.gdt.load();
-
+        self.id = super::cpu::id();
         // reload code segment register
         set_cs(KCODE_SELECTOR);
         // load TSS

kernel/src/arch/x86_64/interrupt/consts.rs

@@ -45,3 +45,6 @@ pub const PIRQE: u8 = 20;
 pub const PIRQF: u8 = 21;
 pub const PIRQG: u8 = 22;
 pub const PIRQH: u8 = 23;
+
+// IPI constants
+pub const IPIFuncCall: u8 = 0xfc;

kernel/src/arch/x86_64/interrupt/handler.rs

@@ -109,6 +109,11 @@ pub extern "C" fn rust_trap(tf: &mut TrapFrame) {
         Syscall32 => syscall32(tf),
         InvalidOpcode => invalid_opcode(tf),
         DivideError | GeneralProtectionFault => error(tf),
+        IPIFuncCall => {
+            let irq = tf.trap_num as u8 - IRQ0;
+            super::ack(irq); // must ack before switching
+            super::super::gdt::Cpu::current().ipi_handler();
+        }
         _ => panic!("Unhandled interrupt {:x}", tf.trap_num),
     }
 }

kernel/src/arch/x86_64/ipi.rs

@@ -0,0 +1,82 @@
+// Interface for inter-processor interrupt.
+// This module wraps inter-processor interrupt into a broadcast-calling style.
+
+use crate::consts::KERNEL_OFFSET;
+use crate::sync::{Semaphore, SpinLock as Mutex};
+use alloc::boxed::Box;
+use alloc::sync::Arc;
+use apic::{LocalApic, XApic, LAPIC_ADDR};
+use lazy_static::*;
+use rcore_memory::Page;
+use x86_64::instructions::tlb;
+use x86_64::VirtAddr;
+
+struct IPIInvoke<'a, A>(&'a (Fn(&A) -> ()), &'a A);
+
+lazy_static! {
+    static ref IPI_INVOKE_LOCK: Mutex<()> = Mutex::new(());
+}
+
+pub trait InvokeEventHandle {
+    fn call(&self);
+}
+
+struct InvokeEvent<A: 'static> {
+    function: fn(&A) -> (),
+    argument: Arc<A>,
+    done_semaphore: Arc<Semaphore>,
+}
+
+impl<A> InvokeEventHandle for InvokeEvent<A> {
+    fn call(&self) {
+        let arg_ref = self.argument.as_ref();
+        (self.function)(arg_ref);
+        self.done_semaphore.release();
+    }
+}
+
+pub type IPIEventItem = Box<InvokeEventHandle>;
+
+// TODO: something fishy is going on here...
+// In fact, the argument lives as long as the Arc.
+fn create_item<A: 'static>(f: fn(&A) -> (), arg: &Arc<A>, sem: &Arc<Semaphore>) -> IPIEventItem {
+    Box::new(InvokeEvent {
+        function: f,
+        argument: arg.clone(),
+        done_semaphore: sem.clone(),
+    })
+}
+unsafe fn get_apic() -> XApic {
+    let mut lapic = unsafe { XApic::new(KERNEL_OFFSET + LAPIC_ADDR) };
+    lapic
+}
+pub fn invoke_on_allcpu<A: 'static>(f: fn(&A) -> (), arg: A, wait: bool) {
+    // Step 1: initialize
+    use super::interrupt::consts::IPIFuncCall;
+    let mut apic = unsafe { get_apic() };
+    let sem = Arc::new(Semaphore::new(0));
+    let arcarg = Arc::new(arg);
+    let mut cpu_count = 0;
+    // Step 2: invoke
+    super::gdt::Cpu::foreach(|cpu| {
+        let id = cpu.get_id();
+        cpu_count += 1;
+        cpu.notify_event(create_item(f, &arcarg, &sem));
+        apic.send_ipi(id as u8, IPIFuncCall);
+    });
+    if wait {
+        for _ in 0..cpu_count {
+            sem.acquire();
+        }
+    }
+}
+
+// Examples of such cases.
+
+pub fn tlb_shootdown(tuple: &(usize, usize)) {
+//    debug!("CPU {}: remote tlb flush {:x?}", super::cpu::id(), tuple);
+    let (start_addr, end_addr) = *tuple;
+    for p in Page::range_of(start_addr, end_addr) {
+        tlb::flush(VirtAddr::new(p.start_address() as u64));
+    }
+}

kernel/src/arch/x86_64/mod.rs

@@ -14,6 +14,7 @@ pub mod paging;
 pub mod rand;
 pub mod syscall;
 pub mod timer;
+pub mod ipi;
 
 static AP_CAN_INIT: AtomicBool = AtomicBool::new(false);
 

kernel/src/arch/x86_64/paging.rs

@@ -13,6 +13,7 @@ use x86_64::structures::paging::{
     FrameAllocator, FrameDeallocator,
 };
 use x86_64::PhysAddr;
+use core::sync::atomic::Ordering;
 
 pub trait PageExt {
     fn of_addr(address: usize) -> Self;
@@ -57,11 +58,13 @@ impl PageTable for ActivePageTable {
                 .unwrap()
                 .flush();
         }
+        flush_tlb_all(addr);
         unsafe { &mut *(get_entry_ptr(addr, 1)) }
     }
 
     fn unmap(&mut self, addr: usize) {
         self.0.unmap(Page::of_addr(addr)).unwrap().1.flush();
+        flush_tlb_all(addr);
     }
 
     fn get_entry(&mut self, addr: usize) -> Option<&mut Entry> {
@@ -93,6 +96,7 @@ impl Entry for PageEntry {
         use x86_64::{instructions::tlb::flush, VirtAddr};
         let addr = VirtAddr::new_unchecked((self as *const _ as u64) << 9);
         flush(addr);
+        flush_tlb_all(addr.as_u64() as usize);
     }
     fn accessed(&self) -> bool {
         self.0.flags().contains(EF::ACCESSED)
@@ -277,3 +281,11 @@ impl FrameDeallocator<Size4KiB> for FrameAllocatorForX86 {
         dealloc_frame(frame.start_address().as_u64() as usize);
     }
 }
+
+/// Flush TLB for `vaddr` on all CPU
+fn flush_tlb_all(vaddr: usize) {
+    if !super::AP_CAN_INIT.load(Ordering::Relaxed) {
+        return;
+    }
+    super::ipi::invoke_on_allcpu(super::ipi::tlb_shootdown, (vaddr, vaddr + 0x1000), false);
+}

user

@@ -1 +1 @@
-Subproject commit bb73d6ecce1ab0e6fae692c51e4335772b0335d4
+Subproject commit 9fb1d459b50bc14c7ac56d9fd94b4b8485620730