Retain physical CPU identity across BKL mailbox polling

.github/workflows/kernel-host-contracts.yml

@@ -6,12 +6,18 @@ on:
       - ".github/workflows/kernel-host-contracts.yml"
       - "src/smp.rs"
       - "src/smp/bkl.rs"
+      - "src/smp/bkl_cpu_tests.rs"
+      - "src/smp/quiescence.rs"
+      - "tests/bkl_contracts.rs"
       - "src/arch/x86_64/syscall_entry.rs"
   pull_request:
     paths:
       - ".github/workflows/kernel-host-contracts.yml"
       - "src/smp.rs"
       - "src/smp/bkl.rs"
+      - "src/smp/bkl_cpu_tests.rs"
+      - "src/smp/quiescence.rs"
+      - "tests/bkl_contracts.rs"
       - "src/arch/x86_64/syscall_entry.rs"
 
 permissions:
@@ -25,5 +31,5 @@ jobs:
       - uses: dtolnay/rust-toolchain@stable
       - name: Test the native BKL ownership module
         run: |
-          rustc --edition 2021 --test src/smp/bkl.rs -o /tmp/kernel-bkl-tests
+          rustc --edition 2021 --test tests/bkl_contracts.rs -o /tmp/kernel-bkl-tests
           /tmp/kernel-bkl-tests

src/smp.rs

@@ -90,6 +90,7 @@ pub fn aps_alive() -> u32 {
 // ---------------------------------------------------------------------------
 
 mod bkl;
+use bkl::AcquisitionCpu;
 static BKL: bkl::BigKernelLock = bkl::BigKernelLock::new();
 #[cfg(all(target_arch = "x86_64", feature = "spec"))]
 pub(crate) static SYSCALL_IDLE_BACKEDGES: core::sync::atomic::AtomicU64 =
@@ -125,34 +126,37 @@ pub fn bkl_acquire() {
 /// Acquire BKL with the exact architectural entry snapshot available to a remote quiescence
 /// request. A true result means the controller adopted this entry; do not execute it again.
 pub(crate) fn bkl_acquire_for_user_entry(mut entry: Option<UserEntrySnapshot>) -> bool {
+    // IF remains clear throughout admission and mailbox waits. Quiescence may adopt
+    // a user TCB, but never migrates this kernel continuation or returns to user mode.
+    let cpu = AcquisitionCpu::sample(crate::arch::get_cpu_id, MAX_CPUS);
     let mut adopted = false;
     #[cfg(all(target_arch = "x86_64", feature = "spec"))]
-    observe_waiting_entry(&entry);
+    observe_waiting_entry(cpu, &entry);
     #[cfg(target_arch = "x86_64")]
-    service_retirement_shootdown();
-    adopted |= service_quiescence(&mut entry);
-    let me = crate::arch::get_cpu_id() + 1;
+    service_retirement_shootdown_for_cpu(cpu);
+    adopted |= service_quiescence(cpu, &mut entry);
+    let me = cpu.owner();
     // A nested acquire cannot mint a second exclusive KernelState borrower.
     assert_ne!(BKL.holder(), me, "nested BKL ownership");
     loop {
         match BKL.try_acquire(me) {
             true => {
                 #[cfg(all(target_arch = "x86_64", feature = "spec"))]
                 {
-                    BKL_WAITERS.fetch_and(!(1 << (me - 1)), Ordering::Release);
-                    WAITING_ENTRY_KIND[(me - 1) as usize].store(0, Ordering::Release);
+                    BKL_WAITERS.fetch_and(!(1 << cpu.id()), Ordering::Release);
+                    WAITING_ENTRY_KIND[cpu.index()].store(0, Ordering::Release);
                 }
                 return adopted;
             }
             false => {
                 #[cfg(all(target_arch = "x86_64", feature = "spec"))]
-                BKL_WAITERS.fetch_or(1 << (me - 1), Ordering::Release);
+                BKL_WAITERS.fetch_or(1 << cpu.id(), Ordering::Release);
                 while BKL.holder() != 0 {
                     // IF is clear here, so a retirement IPI cannot run its ISR. Service its
                     // independent mailbox without borrowing scheduler state or acquiring BKL.
                     #[cfg(target_arch = "x86_64")]
-                    service_retirement_shootdown();
-                    adopted |= service_quiescence(&mut entry);
+                    service_retirement_shootdown_for_cpu(cpu);
+                    adopted |= service_quiescence(cpu, &mut entry);
                     core::hint::spin_loop();
                 }
             }
@@ -425,7 +429,9 @@ pub(crate) fn retirement_spec_snapshot() -> (u32, [u32; MAX_CPUS], bool) {
 }
 
 #[cfg(target_arch = "x86_64")]
-fn flush_retired_translations() {
+fn flush_retired_translations(cpu: AcquisitionCpu) {
+    #[cfg(not(feature = "spec"))]
+    let _ = cpu;
     // Changing CR4.PGE invalidates translations for all PCIDs, including global entries. Both
     // writes preserve every other control bit, and work whether PGE was initially set or clear.
     unsafe {
@@ -440,16 +446,22 @@ fn flush_retired_translations() {
         );
     }
     #[cfg(feature = "spec")]
-    RETIREMENT_FLUSHES[crate::arch::get_cpu_id() as usize].fetch_add(1, Ordering::Relaxed);
+    RETIREMENT_FLUSHES[cpu.index()].fetch_add(1, Ordering::Relaxed);
 }
 
 /// Lock-independent retirement service. Called before an IPI tries BKL and while a kernel
 /// entrant spins for BKL with IF clear. It never touches a TCB, scheduler, or shared IPI queue.
 #[cfg(target_arch = "x86_64")]
 pub(crate) fn service_retirement_shootdown() {
-    let mailbox = &RETIREMENT_MAILBOXES[crate::arch::get_cpu_id() as usize];
+    let cpu = AcquisitionCpu::sample(crate::arch::get_cpu_id, MAX_CPUS);
+    service_retirement_shootdown_for_cpu(cpu);
+}
+
+#[cfg(target_arch = "x86_64")]
+fn service_retirement_shootdown_for_cpu(cpu: AcquisitionCpu) {
+    let mailbox = &RETIREMENT_MAILBOXES[cpu.index()];
     if mailbox.claim() {
-        flush_retired_translations();
+        flush_retired_translations(cpu);
         mailbox.acknowledge();
     }
 }
@@ -469,8 +481,8 @@ pub fn retire_vspace_translations(pml4_paddr: u64) {
 pub fn retire_all_translations() {
     assert!(!RETIREMENT_ACTIVE.swap(true, Ordering::AcqRel),
         "reentrant or concurrent translation retirement shootdown");
-    let me = crate::arch::get_cpu_id();
-    let targets = ONLINE_CPUS.load(Ordering::Acquire) & !(1 << me);
+    let me = AcquisitionCpu::sample(crate::arch::get_cpu_id, MAX_CPUS);
+    let targets = ONLINE_CPUS.load(Ordering::Acquire) & !(1 << me.id());
     for cpu in 0..MAX_CPUS as u32 {
         if targets & (1 << cpu) != 0 {
             RETIREMENT_MAILBOXES[cpu as usize].publish();
@@ -479,7 +491,7 @@ pub fn retire_all_translations() {
             crate::arch::x86_64::lapic::send_ipi(cpu as u8, IPI_VECTOR);
         }
     }
-    flush_retired_translations();
+    flush_retired_translations(me);
     for cpu in 0..MAX_CPUS as u32 {
         if targets & (1 << cpu) != 0 {
             while !RETIREMENT_MAILBOXES[cpu as usize].completed() {

src/smp/bkl.rs

@@ -1,6 +1,36 @@
 //! Atomic CPU ownership for the big kernel lock; waiting and quiescence live in the caller.
 use core::sync::atomic::{AtomicU32, Ordering};
 
+/// One trusted physical-CPU sample for a non-migrating kernel acquisition.
+/// The caller supplies the hardware reader, never user TLS or thread affinity.
+/// Discard this value before returning or dispatching to user mode.
+#[derive(Copy, Clone)]
+pub(super) struct AcquisitionCpu(u32);
+
+impl AcquisitionCpu {
+    pub(super) fn sample(read_hardware: impl FnOnce() -> u32, cpu_limit: usize) -> Self {
+        let cpu = read_hardware();
+        assert!(
+            (cpu as usize) < cpu_limit,
+            "physical CPU outside mailbox range"
+        );
+        assert!(cpu < u32::MAX, "physical CPU cannot encode BKL owner");
+        Self(cpu)
+    }
+
+    pub(super) fn id(self) -> u32 {
+        self.0
+    }
+
+    pub(super) fn index(self) -> usize {
+        self.0 as usize
+    }
+
+    pub(super) fn owner(self) -> u32 {
+        self.0 + 1
+    }
+}
+
 pub struct BigKernelLock {
     owner: AtomicU32,
 }
@@ -41,6 +71,10 @@ impl BigKernelLock {
     }
 }
 
+#[cfg(test)]
+#[path = "bkl_cpu_tests.rs"]
+mod cpu_tests;
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/smp/bkl_cpu_tests.rs

@@ -0,0 +1,126 @@
+use super::{AcquisitionCpu, BigKernelLock, ReleaseError};
+use core::cell::Cell;
+
+#[test]
+fn one_physical_sample_survives_failed_admission_and_mailbox_polling() {
+    for physical in 0..4 {
+        let samples = Cell::new(0);
+        let cpu = AcquisitionCpu::sample(
+            || {
+                samples.set(samples.get() + 1);
+                physical
+            },
+            4,
+        );
+        let lock = BigKernelLock::new();
+        let foreign_owner = (physical + 1) % 4 + 1;
+        while !lock.try_acquire(foreign_owner) {}
+        let mut pending_retirement = true;
+        let mut entry = Some(physical);
+        let mut adopted = false;
+        let mut flushed = [false; 4];
+        let mut acknowledged = [false; 4];
+        for _ in 0..16 {
+            assert!(!lock.try_acquire(cpu.owner()));
+            assert_eq!(lock.release(cpu.owner()), Err(ReleaseError::NotOwner));
+            if pending_retirement {
+                flushed[cpu.index()] = true;
+                assert!(flushed[cpu.index()]);
+                acknowledged[cpu.index()] = true;
+                pending_retirement = false;
+            }
+            adopted |= entry.take().is_some();
+            assert_eq!(cpu.id(), physical);
+        }
+        assert!(adopted);
+        assert!(entry.is_none());
+        assert!(acknowledged[physical as usize]);
+        assert_eq!(acknowledged.iter().filter(|ack| **ack).count(), 1);
+        assert_eq!(samples.get(), 1);
+        assert_eq!(lock.release(foreign_owner), Ok(()));
+        while !lock.try_acquire(cpu.owner()) {}
+        assert_eq!(lock.holder(), physical + 1);
+        assert_eq!(lock.release(cpu.owner()), Ok(()));
+    }
+}
+
+#[test]
+#[should_panic(expected = "physical CPU outside mailbox range")]
+fn sampled_physical_identity_must_fit_mailboxes() {
+    AcquisitionCpu::sample(|| 4, 4);
+}
+
+#[test]
+fn native_acquisition_samples_hardware_once_and_threads_identity_through_polls() {
+    let source = include_str!("../smp.rs");
+    let acquisition = source
+        .split_once("pub(crate) fn bkl_acquire_for_user_entry(")
+        .unwrap()
+        .1;
+    let acquisition = acquisition.split_once("/// Release the BKL").unwrap().0;
+    assert_eq!(acquisition.matches("crate::arch::get_cpu_id").count(), 1);
+    let sample = acquisition
+        .find("AcquisitionCpu::sample(crate::arch::get_cpu_id")
+        .unwrap();
+    let retirement = acquisition
+        .find("service_retirement_shootdown_for_cpu(cpu)")
+        .unwrap();
+    assert!(sample < retirement);
+    assert_eq!(
+        acquisition
+            .matches("service_retirement_shootdown_for_cpu(cpu)")
+            .count(),
+        2
+    );
+    assert_eq!(
+        acquisition
+            .matches("service_quiescence(cpu, &mut entry)")
+            .count(),
+        2
+    );
+    assert!(!acquisition.contains("service_retirement_shootdown();"));
+    assert!(!acquisition.contains("gs:"));
+
+    let quiescence = include_str!("quiescence.rs");
+    let poll = quiescence.split_once("fn service_quiescence(").unwrap().1;
+    let poll = poll.split_once("/// Synchronously withdraw").unwrap().0;
+    assert!(!poll.contains("get_cpu_id"));
+    assert!(poll.contains("service_retirement_shootdown_for_cpu(cpu)"));
+    assert!(poll.contains("QUIESCENCE[cpu.index()]"));
+
+    let flush = source
+        .split_once("fn flush_retired_translations(")
+        .unwrap()
+        .1;
+    let flush = flush
+        .split_once("/// Lock-independent retirement service")
+        .unwrap()
+        .0;
+    assert!(!flush.contains("get_cpu_id"));
+    assert!(flush.contains("RETIREMENT_FLUSHES[cpu.index()]"));
+    assert_eq!(flush.matches("mov cr4,").count(), 2);
+
+    let service = source
+        .split_once("fn service_retirement_shootdown_for_cpu(")
+        .unwrap()
+        .1;
+    let service = service
+        .split_once("/// Synchronous translation retirement")
+        .unwrap()
+        .0;
+    assert!(!service.contains("get_cpu_id"));
+    assert!(service.contains("RETIREMENT_MAILBOXES[cpu.index()]"));
+    assert!(
+        service.find("flush_retired_translations(cpu)").unwrap()
+            < service.find("mailbox.acknowledge()").unwrap()
+    );
+    let wrapper = source
+        .split_once("pub(crate) fn service_retirement_shootdown()")
+        .unwrap()
+        .1;
+    let wrapper = wrapper
+        .split_once("fn service_retirement_shootdown_for_cpu(")
+        .unwrap()
+        .0;
+    assert!(wrapper.contains("AcquisitionCpu::sample(crate::arch::get_cpu_id"));
+}

src/smp/quiescence.rs

@@ -60,25 +60,25 @@ static QUIESCENCE: [QuiescenceMailbox; MAX_CPUS] =
 static WAITING_ENTRY_KIND: [AtomicU32; MAX_CPUS] = [const { AtomicU32::new(0) }; MAX_CPUS];
 
 #[cfg(all(target_arch = "x86_64", feature = "spec"))]
-fn observe_waiting_entry(entry: &Option<UserEntrySnapshot>) {
+fn observe_waiting_entry(cpu: AcquisitionCpu, entry: &Option<UserEntrySnapshot>) {
     let kind = match entry.as_ref().map(|entry| entry.kind) {
         None => 0,
         Some(UserEntryKind::Syscall) => 1,
         Some(_) => 2,
     };
-    WAITING_ENTRY_KIND[crate::arch::get_cpu_id() as usize].store(kind, Ordering::Release);
+    WAITING_ENTRY_KIND[cpu.index()].store(kind, Ordering::Release);
 }
 
 // Runs before acquiring BKL, including IF-clear lock contention. It accesses only
 // this CPU's architectural state and mailbox, never KERNEL or a TCB reference.
-fn service_quiescence(entry: &mut Option<UserEntrySnapshot>) -> bool {
-    let mailbox = &QUIESCENCE[crate::arch::get_cpu_id() as usize];
+fn service_quiescence(cpu: AcquisitionCpu, entry: &mut Option<UserEntrySnapshot>) -> bool {
+    let mailbox = &QUIESCENCE[cpu.index()];
     if mailbox.phase.load(Ordering::Acquire) != 1 {
         return false;
     }
     let captured = entry.take();
     #[cfg(all(target_arch = "x86_64", feature = "spec"))]
-    observe_waiting_entry(entry);
+    observe_waiting_entry(cpu, entry);
     let adopted = captured.is_some();
     let snapshot = QuiescenceSnapshot {
         entry: captured,
@@ -97,7 +97,7 @@ fn service_quiescence(entry: &mut Option<UserEntrySnapshot>) -> bool {
     mailbox.phase.store(2, Ordering::Release);
     while mailbox.phase.load(Ordering::Acquire) != 3 {
         #[cfg(target_arch = "x86_64")]
-        service_retirement_shootdown();
+        service_retirement_shootdown_for_cpu(cpu);
         core::hint::spin_loop();
     }
     mailbox.phase.store(0, Ordering::Release);

tests/bkl_contracts.rs

@@ -0,0 +1,3 @@
+// Preserve the native child-module visibility boundary in the host test lane.
+#[path = "../src/smp/bkl.rs"]
+mod bkl;