diff --git a/src/arch/x86_64/gdt.rs b/src/arch/x86_64/gdt.rs index 53a672665b..1463beff72 100644 --- a/src/arch/x86_64/gdt.rs +++ b/src/arch/x86_64/gdt.rs @@ -245,6 +245,8 @@ pub unsafe fn init_paging(stack_offset: usize, cpu_id: LogicalCpuId) { #[cfg(feature = "profiling")] profiling: None, }; + + crate::percpu::init_tlb_shootdown(cpu_id, &mut pcr.percpu); } #[derive(Copy, Clone, Debug)] #[repr(packed)] diff --git a/src/arch/x86_shared/device/local_apic.rs b/src/arch/x86_shared/device/local_apic.rs index b03c1ff5e8..5c7042ec87 100644 --- a/src/arch/x86_shared/device/local_apic.rs +++ b/src/arch/x86_shared/device/local_apic.rs @@ -149,7 +149,6 @@ impl LocalApic { } self.set_icr(icr); } - // Not used just yet, but allows triggering an NMI to another processor. pub fn ipi_nmi(&mut self, apic_id: u32) { let shift = if self.x2 { 32 } else { 56 }; self.set_icr((u64::from(apic_id) << shift) | (1 << 14) | (0b100 << 8)); diff --git a/src/context/arch/x86_64.rs b/src/context/arch/x86_64.rs index 5a999f7619..08f2fd05aa 100644 --- a/src/context/arch/x86_64.rs +++ b/src/context/arch/x86_64.rs @@ -231,31 +231,38 @@ pub unsafe fn switch_to(prev: &mut super::Context, next: &mut super::Context) { // to dyn or slice fat pointers), and NonNull optimization exists, map_or will hopefully be // optimized down to checking prev and next pointers, as next cannot be null. Some(ref next_space) => { - percpu_block.current_addrspace.set(Arc::as_ptr(next_space)); - if prev .addr_space .as_ref() .map_or(true, |prev_space| !Arc::ptr_eq(prev_space, next_space)) { - if let Some(ref prev_space) = prev.addr_space { - prev_space.inner.read().used_by.atomic_clear(this_cpu); - } + let prev_space_guard = prev.addr_space.as_ref().map(|asp| asp.inner.read()); // This lock needs to be held, because if the address space is write-locked by some // context that is e.g. unmapping memory, we either need to switch after its // changes have been made, or it needs to know this context is potentially using // this address space, at that time. - let next_space = next_space.inner.read(); + let next_space_guard = next_space.inner.read(); - next_space.used_by.atomic_set(this_cpu); - next_space.table.utable.make_current(); + if let Some(prev_space_guard) = prev_space_guard.as_deref() { + prev_space_guard.used_by.atomic_clear(this_cpu); + } + + percpu_block.current_addrspace.set(Arc::as_ptr(&next_space)); + next_space_guard.used_by.atomic_set(this_cpu); + next_space_guard.table.utable.make_current(); } } None => { // The next context is kernel-only, so switch to the page table without any user // mappings. if let Some(ref prev_space) = prev.addr_space { - prev_space.inner.read().used_by.atomic_clear(this_cpu); + let prev_space = prev_space.inner.read(); + prev_space.used_by.atomic_clear(this_cpu); + + // Do this while the lock is held, to prevent deadlocks. + percpu_block.current_addrspace.set(core::ptr::null()); + } else { + percpu_block.current_addrspace.set(core::ptr::null()); } RmmA::set_table(TableKind::User, empty_cr3()); } diff --git a/src/context/context.rs b/src/context/context.rs index 197dd72b86..e5c6290152 100644 --- a/src/context/context.rs +++ b/src/context/context.rs @@ -11,7 +11,7 @@ use crate::{ memory::{Frame, RaiiFrame}, paging::{RmmA, RmmArch}, scheme::{CallerCtx, FileHandle, SchemeNamespace}, - sync::WaitMap, + sync::WaitMap, percpu::PercpuBlock, }; use crate::syscall::{ @@ -399,9 +399,26 @@ impl Context { &mut self, addr_space: Arc, ) -> Option> { - if self.id == super::context_id() { + if let Some(ref addrsp) = self.addr_space && Arc::ptr_eq(addrsp, &addr_space) { + return Some(addr_space); + } + + let current_context_id = super::context_id(); + let current_cpu_id = crate::cpu_id(); + + if self.id == current_context_id { unsafe { - addr_space.inner.read().table.utable.make_current(); + let guard = addr_space.inner.read(); + + let prev_addrsp = self.addr_space.as_ref().map(|asp| asp.inner.read()); + + if let Some(prev_addrsp) = prev_addrsp.as_deref() { + prev_addrsp.used_by.atomic_clear(current_cpu_id); + } + + guard.used_by.atomic_set(current_cpu_id); + PercpuBlock::current().current_addrspace.set(Arc::as_ptr(&addr_space)); + guard.table.utable.make_current(); } } diff --git a/src/percpu.rs b/src/percpu.rs index 0ddc18af0e..3649b6d419 100644 --- a/src/percpu.rs +++ b/src/percpu.rs @@ -31,6 +31,10 @@ pub struct PercpuBlock { const NULL: AtomicPtr = AtomicPtr::new(core::ptr::null_mut()); static ALL_PERCPU_BLOCKS: [AtomicPtr; MAX_CPU_COUNT as usize] = [NULL; MAX_CPU_COUNT as usize]; +pub unsafe fn init_tlb_shootdown(id: LogicalCpuId, block: *mut PercpuBlock) { + ALL_PERCPU_BLOCKS[id.get() as usize].store(block, Ordering::Release) +} + // PercpuBlock::current() is implemented somewhere in the arch-specific modules bitflags::bitflags! { @@ -52,6 +56,7 @@ pub fn shootdown_tlb_ipi(_target: Option) {} pub fn shootdown_tlb_ipi(target: Option) { if let Some(target) = target { let Some(percpublock) = (unsafe { ALL_PERCPU_BLOCKS[target.get() as usize].load(Ordering::Acquire).as_ref() }) else { + log::warn!("Trying to TLB shootdown a CPU that doesn't exist or isn't initialized."); return; }; let bit = NmiReasons::TLB_SHOOTDOWN.bits();