keos/mm/
tlb.rs

1//! TLB Shootdown helper.
2use crate::sync::{
3    RwLock,
4    atomic::{AtomicBool, AtomicUsize},
5};
6use abyss::{
7    MAX_CPU,
8    addressing::Va,
9    boot::ONLINE_CPU,
10    dev::x86_64::apic::{IPIDest, Mode},
11    interrupt::Registers,
12    spinlock::SpinLock,
13    x86_64::Cr3,
14};
15use core::sync::atomic::Ordering;
16
17#[doc(hidden)]
18static IN_PROGRESS: SpinLock<()> = SpinLock::new(());
19
20#[doc(hidden)]
21static REQUEST: RwLock<Option<TlbIpi>> = RwLock::new(None);
22
23#[doc(hidden)]
24#[allow(clippy::declare_interior_mutable_const)]
25const PER_CORE_STATUS: AtomicBool = AtomicBool::new(false);
26
27#[doc(hidden)]
28static HAVE_REQUEST: [AtomicBool; MAX_CPU] = [PER_CORE_STATUS; MAX_CPU];
29
30/// Struct for TLB request
31pub struct TlbIpi {
32    /// Destination Cr3
33    cr3: Cr3,
34
35    /// If va is Some, invalidate only that page. Otherwise, shutdown the whole
36    /// TLB.
37    va: Option<Va>,
38
39    /// Count of CPUs that processed this request.
40    processed: AtomicUsize,
41}
42
43impl TlbIpi {
44    /// Send the request and wait until the request is done for all CPUs
45    pub fn send(cr3: Cr3, va: Option<Va>) {
46        let guard = loop {
47            if let Ok(guard) = IN_PROGRESS.try_lock() {
48                break guard;
49            }
50
51            TlbIpi::handle();
52        };
53
54        // Publish the requests.
55        {
56            let mut request = REQUEST.write();
57
58            assert!(
59                request.is_none(),
60                "Before sending TLB Shootdown request, the request queue must be empty."
61            );
62
63            *request = Some(Self {
64                cr3,
65                va,
66                processed: AtomicUsize::new(0),
67            });
68        }
69
70        // Waiting the requests
71        {
72            let request = REQUEST.read();
73            let self_id = abyss::x86_64::intrinsics::cpuid();
74
75            let online_cpu_cnt = ONLINE_CPU
76                .iter()
77                .enumerate()
78                .filter(|(i, cpu)| {
79                    if *i == self_id {
80                        true
81                    } else if cpu.load(Ordering::SeqCst) {
82                        HAVE_REQUEST[*i].store(true);
83                        true
84                    } else {
85                        false
86                    }
87                })
88                .count();
89
90            unsafe {
91                abyss::dev::x86_64::apic::send_ipi(IPIDest::AllExcludingSelf, Mode::Fixed(0x7e));
92            }
93
94            let request_ref = request.as_ref().unwrap();
95
96            while request_ref.processed.load() < online_cpu_cnt - 1 {
97                core::hint::spin_loop();
98            }
99        }
100
101        // Clean up the request
102        {
103            let mut request = REQUEST.write();
104            *request = None;
105        }
106
107        guard.unlock();
108    }
109
110    fn handle() {
111        let core_id = abyss::x86_64::intrinsics::cpuid();
112        if HAVE_REQUEST[core_id].load() {
113            let request = REQUEST.read();
114
115            if let Some(request) = &*request {
116                if request.cr3 == Cr3::current() {
117                    match request.va {
118                        Some(va) => unsafe {
119                            core::arch::asm!(
120                                "invlpg [{0}]",
121                                in(reg) va.into_usize(),
122                                options(nostack)
123                            )
124                        },
125                        _ => unsafe {
126                            core::arch::asm! {
127                                "mov rax, cr3",
128                                "mov cr3, rax",
129                                out("rax") _,
130                                options(nostack)
131                            }
132                        },
133                    }
134                }
135
136                request.processed.fetch_add(1);
137                HAVE_REQUEST[core_id].store(false);
138            }
139        }
140    }
141}
142
143/// Event handler for TLB Shootdown request
144pub fn handler(_regs: &mut Registers) {
145    TlbIpi::handle();
146}
keos/mm/tlb.rs

keos/mm/
tlb.rs
