Skip to content

TDX: Pull TLB flush ringbuf struct into support/, add tests #1387

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
May 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .config/nextest.toml
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,11 @@ fail-fast = false
path = "junit.xml"
store-success-output = "true"

[[profile.ci.overrides]]
# allow loom based tests more time, as they take a while
filter = 'test(loom)'
slow-timeout = { period = "30s", terminate-after = 2 }

[[profile.ci.overrides]]
# use fuzzy-matching for the package() to allow out-of-tree tests to use the
# same profile
Expand Down
11 changes: 11 additions & 0 deletions Cargo.lock
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,16 @@ version = "1.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"

[[package]]
name = "atomic_ringbuf"
version = "0.0.0"
dependencies = [
"cfg-if",
"inspect",
"loom",
"parking_lot",
]

[[package]]
name = "autocfg"
version = "1.4.0"
Expand Down Expand Up @@ -7949,6 +7959,7 @@ dependencies = [
"aarch64defs",
"aarch64emu",
"anyhow",
"atomic_ringbuf",
"bitfield-struct 0.10.1",
"bitvec",
"build_rs_guest_arch",
Expand Down
6 changes: 6 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ pipette_protocol = { path = "petri/pipette_protocol" }
# support crates
address_filter = { path = "support/address_filter" }
arc_cyclic_builder = { path = "support/arc_cyclic_builder" }
atomic_ringbuf = { path = "support/atomic_ringbuf" }
cache_topology = { path = "support/cache_topology" }
closeable_mutex = { path = "support/closeable_mutex" }
ci_logger = { path = "support/ci_logger" }
Expand Down Expand Up @@ -462,6 +463,7 @@ libfuzzer-sys = "0.4"
libtest-mimic = "0.8"
linkme = "0.3.9"
log = "0.4"
loom = "0.7.2"
macaddr = "1.0"
mbrman = "0.5"
mimalloc = { version = "0.1.39", default-features = false }
Expand Down Expand Up @@ -595,6 +597,10 @@ lto = 'fat'
[profile.dev.package.minimal_rt]
opt-level = 1

# Optimize loom even in the dev profile, as otherwise tests just take too long.
[profile.dev.package.loom]
opt-level = 3

[patch.crates-io]
# Pending <https://github.com/ferrilab/bitvec/pull/273>
bitvec = { git = "https://github.com/smalis-msft/bitvec", branch = "set-aliased-previous-val" }
Expand Down
1 change: 1 addition & 0 deletions openhcl/virt_mshv_vtl/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ vmcore.workspace = true
x86defs.workspace = true
x86emu.workspace = true

atomic_ringbuf.workspace = true
inspect_counters.workspace = true
inspect = { workspace = true, features = ["std"] }
mesh.workspace = true
Expand Down
107 changes: 8 additions & 99 deletions openhcl/virt_mshv_vtl/src/processor/tdx/tlb_flush.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,18 +5,15 @@

use crate::TdxBacked;
use crate::UhProcessor;
use atomic_ringbuf::AtomicRingBuffer;
use hcl::GuestVtl;
use hcl::ioctl::ProcessorRunner;
use hcl::ioctl::tdx::Tdx;
use hvdef::hypercall::HvGvaRange;
use inspect::Inspect;
use parking_lot::Mutex;
use parking_lot::MutexGuard;
use safeatomic::AtomicSliceOps;
use std::num::Wrapping;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use x86defs::tdx::TdGlaVmAndFlags;
use x86defs::tdx::TdxGlaListInfo;
Expand All @@ -29,7 +26,7 @@ pub(super) const FLUSH_GVA_LIST_SIZE: usize = 32;
#[derive(Debug, Inspect)]
pub(super) struct TdxPartitionFlushState {
/// A fixed-size ring buffer of GVAs that need to be flushed.
pub(super) gva_list: AtomicTlbRingBuffer,
pub(super) gva_list: AtomicRingBuffer<FLUSH_GVA_LIST_SIZE, HvGvaRange>,
/// The number of times an entire TLB flush has been requested.
pub(super) flush_entire_counter: AtomicU32,
/// The number of times a non-global TLB flush has been requested.
Expand All @@ -40,7 +37,7 @@ pub(super) struct TdxPartitionFlushState {
impl TdxPartitionFlushState {
pub(super) fn new() -> Self {
Self {
gva_list: AtomicTlbRingBuffer::new(),
gva_list: AtomicRingBuffer::new(),
flush_entire_counter: AtomicU32::new(0),
flush_entire_non_global_counter: AtomicU32::new(0),
}
Expand Down Expand Up @@ -112,7 +109,7 @@ impl UhProcessor<'_, TdxBacked> {
if flush_entire_required {
self_flush_state.flush_entire_counter = Wrapping(partition_flush_entire);
self_flush_state.flush_entire_non_global_counter = Wrapping(partition_flush_non_global);
self_flush_state.gva_list_count = Wrapping(partition_flush_state.gva_list.count());
self_flush_state.gva_list_count = partition_flush_state.gva_list.count();
Self::set_flush_entire(
true,
&mut self.backing.vtls[target_vtl].private_regs.vp_entry_flags,
Expand All @@ -139,7 +136,7 @@ impl UhProcessor<'_, TdxBacked> {
flush_page: &user_driver::memory::MemoryBlock,
) -> bool {
// Check quickly to see whether any new addresses are in the list.
let partition_list_count = Wrapping(partition_flush_state.gva_list.count());
let partition_list_count = partition_flush_state.gva_list.count();
if partition_list_count == *gva_list_count {
return true;
}
Expand All @@ -151,10 +148,10 @@ impl UhProcessor<'_, TdxBacked> {
}

// The last `count_diff` addresses are the new ones, copy them locally.
let flush_addrs = &mut [0; FLUSH_GVA_LIST_SIZE][..count_diff];
let flush_addrs = &mut [HvGvaRange(0); FLUSH_GVA_LIST_SIZE][..count_diff];
if !partition_flush_state
.gva_list
.try_copy(*gva_list_count, flush_addrs)
.try_copy(gva_list_count.0, flush_addrs)
{
return false;
}
Expand All @@ -165,7 +162,7 @@ impl UhProcessor<'_, TdxBacked> {

if count_diff == 1 {
runner
.invgla(gla_flags, TdxGlaListInfo::from(flush_addrs[0]))
.invgla(gla_flags, TdxGlaListInfo::from(flush_addrs[0].0))
.unwrap();
} else {
gla_flags.set_list(true);
Expand Down Expand Up @@ -198,91 +195,3 @@ impl UhProcessor<'_, TdxBacked> {
}
}
}

#[derive(Debug, Inspect)]
pub(super) struct AtomicTlbRingBuffer {
/// The contents of the buffer.
#[inspect(hex, with = "|x| inspect::iter_by_index(x.iter())")]
buffer: Box<[AtomicU64; FLUSH_GVA_LIST_SIZE]>,
/// The number of GVAs that have been added over the lifetime of the VM.
gva_list_count: AtomicUsize,
/// The number of GVAs that have started being added to the list over the
/// lifetime of the VM.
in_progress_count: AtomicUsize,
/// A guard to ensure that only one thread is writing to the list at a time.
write_lock: Mutex<()>,
}

pub(super) struct AtomicTlbRingBufferWriteGuard<'a> {
buf: &'a AtomicTlbRingBuffer,
_write_lock: MutexGuard<'a, ()>,
}

impl AtomicTlbRingBuffer {
fn new() -> Self {
Self {
buffer: Box::new(std::array::from_fn(|_| AtomicU64::new(0))),
gva_list_count: AtomicUsize::new(0),
in_progress_count: AtomicUsize::new(0),
write_lock: Mutex::new(()),
}
}

fn count(&self) -> usize {
self.gva_list_count.load(Ordering::Acquire)
}

pub fn write(&self) -> AtomicTlbRingBufferWriteGuard<'_> {
let write_lock = self.write_lock.lock();
AtomicTlbRingBufferWriteGuard {
buf: self,
_write_lock: write_lock,
}
}

fn try_copy(&self, start_count: Wrapping<usize>, flush_addrs: &mut [u64]) -> bool {
let mut index = start_count;
for flush_addr in flush_addrs.iter_mut() {
*flush_addr = self.buffer[index.0 % FLUSH_GVA_LIST_SIZE].load(Ordering::Relaxed);
index += 1;
}
std::sync::atomic::fence(Ordering::Acquire);

// Check to see whether any additional entries have been added
// that would have caused a wraparound. If so, the local list is
// incomplete and the copy has failed.
if (Wrapping(self.in_progress_count.load(Ordering::Acquire)) - start_count).0
> FLUSH_GVA_LIST_SIZE
{
return false;
}
true
}
}

impl AtomicTlbRingBufferWriteGuard<'_> {
pub fn extend(&self, items: impl ExactSizeIterator<Item = HvGvaRange>) {
debug_assert_eq!(
self.buf.in_progress_count.load(Ordering::Relaxed),
self.buf.gva_list_count.load(Ordering::Relaxed)
);
// Adding a new item to the buffer must be done in three steps:
// 1. Indicate that an entry is about to be added so that any flush
// code executing simultaneously will know that it might lose an
// entry that it is expecting to see.
// 2. Add the entry.
// 3. Increment the valid entry count so that any flush code executing
// simultaneously will know it is valid.
let len = items.len();
let start_count = self.buf.in_progress_count.load(Ordering::Relaxed);
let end_count = start_count.wrapping_add(len);
self.buf
.in_progress_count
.store(end_count, Ordering::Relaxed);
for (i, v) in items.enumerate() {
self.buf.buffer[(start_count.wrapping_add(i)) % FLUSH_GVA_LIST_SIZE]
.store(v.0, Ordering::Release);
}
self.buf.gva_list_count.store(end_count, Ordering::Release);
}
}
18 changes: 18 additions & 0 deletions support/atomic_ringbuf/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

[package]
name = "atomic_ringbuf"
rust-version.workspace = true
edition.workspace = true

[dependencies]
cfg-if.workspace = true
inspect.workspace = true
parking_lot.workspace = true

[target.'cfg(target_arch = "x86_64")'.dev-dependencies]
loom.workspace = true

[lints]
workspace = true
Loading