enhance TLB shootdowns on TDX

common/Cargo.toml

@@ -9,4 +9,4 @@ snp-types = { path = "snp-types" }
 tdx-types = { path = "tdx-types" }
 
 [patch.crates-io]
-x86_64 = { git = "https://github.com/rust-osdev/x86_64.git", rev = "c5bc9fc" }
+x86_64 = { git = "https://github.com/rust-osdev/x86_64.git", rev = "3fc9106" }

common/constants/Cargo.toml

@@ -3,7 +3,8 @@ name = "constants"
 version = "0.1.0"
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[features]
+nightly = []
 
 [dependencies]
 x86_64 = { version = "0.15.1", default-features = false }

common/constants/src/lib.rs

@@ -1,8 +1,13 @@
-//! This crate contains constants shared between the kernel, loader and host executable.
+//! This crate contains constants and related types shared between the kernel,
+//! loader and host executable.
 #![cfg_attr(not(test), no_std)]
-#![forbid(unsafe_code)]
 
-use core::{marker::PhantomData, ops::RangeInclusive};
+use core::{
+    fmt::{self, Debug, Display},
+    marker::PhantomData,
+    ops::{BitAnd, BitAndAssign, BitOrAssign, Index, RangeInclusive},
+    sync::atomic::{AtomicU32, Ordering},
+};
 
 use x86_64::{
     structures::paging::{
@@ -13,6 +18,234 @@ use x86_64::{
 
 pub const MAX_APS_COUNT: u8 = 32;
 
+/// `ApIndex` represents the index of one vCPU thread running the workload
+/// kernel. It's maximum value is capped at compile time.
+#[derive(Clone, Copy, PartialEq, Eq)]
+#[repr(transparent)]
+pub struct ApIndex(u8);
+
+impl ApIndex {
+    /// Create a new `ApIndex` from an integer.
+    ///
+    /// # Panics
+    ///
+    /// This function panics if `idx` exceeds [`MAX_APS_COUNT`].
+    #[must_use]
+    pub const fn new(idx: u8) -> Self {
+        assert!(idx < MAX_APS_COUNT);
+        Self(idx)
+    }
+
+    /// Create a new `ApIndex` from an integer or return `None` if `idx`
+    /// exceeds [`MAX_APS_COUNT`].
+    #[must_use]
+    pub const fn try_new(idx: u8) -> Option<Self> {
+        if idx < MAX_APS_COUNT {
+            Some(Self::new(idx))
+        } else {
+            None
+        }
+    }
+
+    /// Returns `true` for the first AP that starts running.
+    #[must_use]
+    pub const fn is_first(&self) -> bool {
+        self.0 == 0
+    }
+
+    #[must_use]
+    pub fn as_u8(&self) -> u8 {
+        self.0
+    }
+}
+
+impl Display for ApIndex {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        Display::fmt(&self.0, f)
+    }
+}
+
+impl Debug for ApIndex {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        Debug::fmt(&self.0, f)
+    }
+}
+
+impl<T> Index<ApIndex> for [T; MAX_APS_COUNT as usize] {
+    type Output = T;
+
+    fn index(&self, index: ApIndex) -> &Self::Output {
+        unsafe { self.get_unchecked(usize::from(index.as_u8())) }
+    }
+}
+
+type BitmapType = u32;
+type AtomicBitmapType = AtomicU32;
+
+// Make sure that both types are of the same size.
+const _: () = assert!(size_of::<BitmapType>() == size_of::<AtomicBitmapType>());
+
+// Make sure that the bitmap type can fit all bits.
+const _: () = assert!((MAX_APS_COUNT as usize).div_ceil(8) <= size_of::<BitmapType>());
+
+/// A bitmap containing one bit for every vCPU thread running the workload.
+/// Its size is capped by [`MAX_APS_COUNT`] at compile-time.
+#[derive(Clone, Copy)]
+#[repr(transparent)]
+pub struct ApBitmap(BitmapType);
+
+impl ApBitmap {
+    /// Create a new bitmap with all bits set to `false`.
+    #[must_use]
+    pub const fn empty() -> Self {
+        Self(0)
+    }
+
+    /// Create a new bitmap with all bits set to `true`.
+    #[must_use]
+    pub const fn all() -> Self {
+        let mut bits = 0;
+        let mut i = 0;
+        while i < MAX_APS_COUNT {
+            bits |= 1 << i;
+            i += 1;
+        }
+        Self(bits)
+    }
+
+    /// Returns the bit for the given AP.
+    #[must_use]
+    pub const fn get(&self, idx: ApIndex) -> bool {
+        self.0 & (1 << idx.0) != 0
+    }
+
+    /// Sets the bit for the given AP.
+    #[cfg(feature = "nightly")] // TODO: Remove this when Rust 1.83 is released.
+    pub const fn set(&mut self, idx: ApIndex, value: bool) {
+        if value {
+            self.0 |= 1 << idx.0;
+        } else {
+            self.0 &= !(1 << idx.0);
+        }
+    }
+
+    /// Sets the bit for the given AP.
+    #[cfg(not(feature = "nightly"))] // TODO: Remove this when Rust 1.83 is released.
+    pub fn set(&mut self, idx: ApIndex, value: bool) {
+        if value {
+            self.0 |= 1 << idx.0;
+        } else {
+            self.0 &= !(1 << idx.0);
+        }
+    }
+
+    /// Returns whether all bits are `false`.
+    #[must_use]
+    pub const fn is_empty(&self) -> bool {
+        self.0 == 0
+    }
+
+    /// Returns the index of the first AP whose bit is not set.
+    #[must_use]
+    pub fn first_unset(&self) -> Option<ApIndex> {
+        let idx = self.0.trailing_ones() as u8;
+        ApIndex::try_new(idx)
+    }
+}
+
+impl BitOrAssign for ApBitmap {
+    fn bitor_assign(&mut self, rhs: Self) {
+        self.0 |= rhs.0;
+    }
+}
+
+impl BitAnd for ApBitmap {
+    type Output = Self;
+
+    fn bitand(self, rhs: Self) -> Self::Output {
+        Self(self.0 & rhs.0)
+    }
+}
+
+impl BitAndAssign for ApBitmap {
+    fn bitand_assign(&mut self, rhs: Self) {
+        *self = *self & rhs;
+    }
+}
+
+impl IntoIterator for ApBitmap {
+    type Item = ApIndex;
+    type IntoIter = ApBitmapIter;
+
+    /// Returns the indicies of all APs whose bit is `true`.
+    fn into_iter(self) -> Self::IntoIter {
+        ApBitmapIter(self)
+    }
+}
+
+/// Returns the indicies of all APs whose bit is `true`.
+pub struct ApBitmapIter(ApBitmap);
+
+impl Iterator for ApBitmapIter {
+    type Item = ApIndex;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let idx = self.0 .0.trailing_zeros();
+        let idx = ApIndex::try_new(idx as u8)?;
+        self.0.set(idx, false);
+        Some(idx)
+    }
+}
+
+impl Debug for ApBitmap {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_set().entries(*self).finish()
+    }
+}
+
+/// The atomic equivalent of [`ApBitmap`].
+#[repr(transparent)]
+pub struct AtomicApBitmap(AtomicBitmapType);
+
+impl AtomicApBitmap {
+    /// Create a new bitmap with all bits set to `false`.
+    pub const fn empty() -> Self {
+        Self(AtomicBitmapType::new(0))
+    }
+
+    /// Returns the bit for the given AP.
+    pub fn get(&self, idx: ApIndex) -> bool {
+        self.0.load(Ordering::SeqCst) & (1 << idx.0) != 0
+    }
+
+    /// Returns a copy of all bits.
+    pub fn get_all(&self) -> ApBitmap {
+        ApBitmap(self.0.load(Ordering::SeqCst))
+    }
+
+    /// Sets the bit for the given AP to `true`.
+    pub fn set(&self, idx: ApIndex) -> bool {
+        let mask = 1 << idx.0;
+        self.0.fetch_or(mask, Ordering::SeqCst) & mask != 0
+    }
+
+    /// Sets the bits for the given APs to `true`.
+    pub fn set_all(&self, aps: ApBitmap) {
+        self.0.fetch_or(aps.0, Ordering::SeqCst);
+    }
+
+    /// Atomically clear the bit for the given AP and return its value.
+    pub fn take(&self, idx: ApIndex) -> bool {
+        let mask = 1 << idx.0;
+        self.0.fetch_and(!mask, Ordering::SeqCst) & mask != 0
+    }
+
+    /// Atomically clears the bits for all APs and return their values.
+    pub fn take_all(&self) -> ApBitmap {
+        ApBitmap(self.0.swap(0, Ordering::SeqCst))
+    }
+}
+
 pub const FIRST_AP: u8 = 0x80;
 
 pub const EXIT_PORT: u16 = 0xf4;
@@ -120,7 +353,7 @@ where
 mod tests {
     use x86_64::{structures::paging::Page, VirtAddr};
 
-    use crate::{check_ranges, PageRange};
+    use crate::{check_ranges, ApBitmap, PageRange, MAX_APS_COUNT};
 
     #[test]
     fn test_address_range() {
@@ -171,4 +404,14 @@ mod tests {
             PageRange::<Page>::new(0x2000..=0x2fff),
         ])
     }
+
+    #[test]
+    fn test_bitmap_range() {
+        let bitmap = ApBitmap::all();
+        let iter = bitmap.into_iter();
+        for (i, idx) in iter.enumerate() {
+            assert_eq!(i, usize::from(idx.as_u8()));
+        }
+        assert_eq!(bitmap.into_iter().count(), usize::from(MAX_APS_COUNT));
+    }
 }

common/profiler-types/src/lib.rs

@@ -1,20 +1,17 @@
 #![no_std]
 
-use core::{mem::size_of, sync::atomic::AtomicU8};
+use core::mem::size_of;
 
 use bytemuck::{NoUninit, Zeroable};
-use constants::MAX_APS_COUNT;
-
-const NOTIFY_BITS: usize = MAX_APS_COUNT as usize;
-pub const NOTIFY_BYTES: usize = NOTIFY_BITS.div_ceil(8);
+use constants::{AtomicApBitmap, MAX_APS_COUNT};
 
 #[repr(C)]
 pub struct ProfilerControl {
     /// This is a bit field containing a bit for each AP. Set to `true` by the
     /// kernel after it writes to a header. Set to `false` by the host after
     /// reading and processing the header. This mechanism aims to reduce
     /// contention.
-    pub notify_flags: [AtomicU8; NOTIFY_BYTES],
+    pub notify_flags: AtomicApBitmap,
     pub headers: [PerCpuHeader; MAX_APS_COUNT as usize],
     /// The effective frequency in MHz of the guest view of TSC.
     pub tsc_mhz: u64,

common/supervisor-services/Cargo.toml

@@ -11,3 +11,4 @@ kernel = []
 [dependencies]
 bit_field = "0.10.2"
 bytemuck = { version = "1.15.0", features = ["derive"] }
+constants = { workspace = true }

common/supervisor-services/src/allocation_buffer.rs

@@ -6,7 +6,9 @@
 //! the buffer by issuing the [`AllocateMemory`](crate::command_buffer::AllocateMemory)
 //! command.
 
-use core::sync::atomic::{AtomicU16, Ordering};
+use core::sync::atomic::AtomicU16;
+#[cfg(any(feature = "kernel", feature = "supervisor"))]
+use core::sync::atomic::Ordering;
 
 use bytemuck::{Pod, Zeroable};
 

common/supervisor-services/src/command_buffer.rs

@@ -4,7 +4,9 @@
 
 #[cfg(feature = "kernel")]
 use core::iter::once;
-use core::sync::atomic::{AtomicU16, AtomicU8, Ordering};
+#[cfg(any(feature = "kernel", feature = "supervisor"))]
+use core::sync::atomic::Ordering;
+use core::sync::atomic::{AtomicU16, AtomicU8};
 
 #[cfg(feature = "kernel")]
 use bytemuck::bytes_of;