virt_mshv_vtl: Cleanup some vpindexes (#980)

smalis-msft · web-flow · commit fb4fc6b898a7 · 2025-03-11T09:37:03.000-07:00
Try to use VpIndex everywhere instead of usize and u32. The main place
where we aren't today is in the hypercall traits, which are defined
elsewhere.

This also removes a TODO about cpu_index potentially being different
than vp_index, and the cpu_index field entirely, under the assumption
that if it hasn't been an issue up till now it's not going to be. If we
do want to keep this field we should probably have some docs and a more
careful audit of when we should be using vp_index vs cpu_index.
diff --git a/openhcl/virt_mshv_vtl/src/lib.rs b/openhcl/virt_mshv_vtl/src/lib.rs
@@ -654,12 +654,20 @@ struct UhVpInner {
     message_queues: VtlArray<MessageQueues, 2>,
     #[inspect(skip)]
     vp_info: TargetVpInfo,
+    /// The Linux kernel's CPU index for this VP. This should be used instead of VpIndex
+    /// when interacting with non-MSHV kernel interfaces.
     cpu_index: u32,
     #[inspect(with = "|arr| inspect::iter_by_index(arr.iter().map(|v| v.lock().is_some()))")]
     hv_start_enable_vtl_vp: VtlArray<Mutex<Option<Box<VpStartEnableVtl>>>, 2>,
     sidecar_exit_reason: Mutex<Option<SidecarExitReason>>,
 }
 
+impl UhVpInner {
+    pub fn vp_index(&self) -> VpIndex {
+        self.vp_info.base.vp_index
+    }
+}
+
 #[cfg_attr(not(guest_arch = "x86_64"), allow(dead_code))]
 #[derive(Debug, Inspect)]
 /// Which operation is setting the initial vp context
@@ -855,11 +863,16 @@ impl UhPartitionInner {
 
     /// For requester VP to issue `proxy_irr_blocked` update to other VPs
     #[cfg(guest_arch = "x86_64")]
-    fn request_proxy_irr_filter_update(&self, vtl: GuestVtl, device_vector: u8, req_vp_index: u32) {
+    fn request_proxy_irr_filter_update(
+        &self,
+        vtl: GuestVtl,
+        device_vector: u8,
+        req_vp_index: VpIndex,
+    ) {
         tracing::debug!(
             ?vtl,
             device_vector,
-            req_vp_index,
+            req_vp_index = req_vp_index.index(),
             "request_proxy_irr_filter_update"
         );
 
@@ -871,7 +884,7 @@ impl UhPartitionInner {
 
         // Wake all other VPs for their `proxy_irr_blocked` filter update
         for vp in self.vps.iter() {
-            if vp.cpu_index != req_vp_index {
+            if vp.vp_index() != req_vp_index {
                 vp.wake(vtl, WakeReason::UPDATE_PROXY_IRR_FILTER);
             }
         }
@@ -1593,7 +1606,7 @@ impl<'a> UhProtoPartition<'a> {
             .vps_arch()
             .map(|vp_info| {
                 // TODO: determine CPU index, which in theory could be different
-                // from the VP index.
+                // from the VP index, though this hasn't happened yet.
                 let cpu_index = vp_info.base.vp_index.index();
                 UhVpInner::new(cpu_index, vp_info)
             })
diff --git a/openhcl/virt_mshv_vtl/src/processor/hardware_cvm/mod.rs b/openhcl/virt_mshv_vtl/src/processor/hardware_cvm/mod.rs
@@ -603,7 +603,7 @@ impl<T: CpuIo, B: HardwareIsolatedBacking> UhHypercallHandler<'_, '_, T, B> {
         self.vp.partition.request_proxy_irr_filter_update(
             self.intercepted_vtl,
             vector as u8,
-            self.vp.vp_index().index(),
+            self.vp.vp_index(),
         );
 
         // Update `proxy_irr_blocked` for this VP itself
@@ -1414,8 +1414,7 @@ impl<B: HardwareIsolatedBacking> UhProcessor<'_, B> {
 
     /// Returns the per-vp cvm inner state for this vp
     pub fn cvm_vp_inner(&self) -> &'_ crate::UhCvmVpInner {
-        self.cvm_partition()
-            .vp_inner(self.inner.vp_info.base.vp_index.index())
+        self.cvm_partition().vp_inner(self.vp_index().index())
     }
 
     /// Returns the appropriately backed TLB flush and lock access
@@ -1481,7 +1480,7 @@ impl<B: HardwareIsolatedBacking> UhProcessor<'_, B> {
             }
 
             tracing::debug!(
-                vp_index = self.inner.cpu_index,
+                vp_index = self.vp_index().index(),
                 ?vtl,
                 ?start_enable_vtl_state.operation,
                 "setting up vp with initial registers"
diff --git a/openhcl/virt_mshv_vtl/src/processor/hardware_cvm/tlb_lock.rs b/openhcl/virt_mshv_vtl/src/processor/hardware_cvm/tlb_lock.rs
@@ -9,9 +9,10 @@ use crate::UhProcessor;
 use hcl::GuestVtl;
 use hvdef::Vtl;
 use std::sync::atomic::Ordering;
+use virt::VpIndex;
 
 pub struct TlbLockAccess<'a> {
-    pub vp_index: usize,
+    pub vp_index: VpIndex,
     pub cvm_partition: &'a UhCvmPartitionState,
 }
 
@@ -21,17 +22,15 @@ impl TlbLockAccess<'_> {
         // those VPs that hold the lock at this point can block progress, because
         // any VP that acquires the lock after this point is guaranteed to see
         // state that this VP has already flushed.
-        let self_lock = &self
-            .cvm_partition
-            .vp_inner(self.vp_index as u32)
-            .tlb_lock_info[target_vtl];
+        let self_index = self.vp_index.index() as usize;
+        let self_lock = &self.cvm_partition.vp_inner(self_index as u32).tlb_lock_info[target_vtl];
         for vp in self.cvm_partition.tlb_locked_vps[target_vtl]
             .clone()
             .iter_ones()
         {
             // Never wait on the current VP, since the current VP will always
             // release its locks correctly when returning to the target VTL.
-            if vp == self.vp_index {
+            if vp == self_index {
                 continue;
             }
 
@@ -49,15 +48,15 @@ impl TlbLockAccess<'_> {
             // be new, this bit should not already be set.
             let other_lock_blocked =
                 &self.cvm_partition.vp_inner(vp as u32).tlb_lock_info[target_vtl].blocked_vps;
-            let _was_other_lock_blocked = other_lock_blocked.set_aliased(self.vp_index, true);
+            let _was_other_lock_blocked = other_lock_blocked.set_aliased(self_index, true);
             debug_assert!(!_was_other_lock_blocked);
 
             // It is possible that the target VP released the TLB lock before
             // the current VP was added to its blocked set. Check again to
             // see whether the TLB lock is still held, and if not, remove the
             // block.
             if !self.cvm_partition.tlb_locked_vps[target_vtl][vp] {
-                other_lock_blocked.set_aliased(self.vp_index, false);
+                other_lock_blocked.set_aliased(self_index, false);
                 if self_lock.blocking_vps.set_aliased(vp, false) {
                     self_lock.blocking_vp_count.fetch_sub(1, Ordering::Relaxed);
                 }
@@ -72,7 +71,7 @@ impl<B: HardwareIsolatedBacking> UhProcessor<'_, B> {
     /// concurrent instruction emulation.
     pub fn set_wait_for_tlb_locks(&mut self, target_vtl: GuestVtl) {
         TlbLockAccess {
-            vp_index: self.vp_index().index() as usize,
+            vp_index: self.vp_index(),
             cvm_partition: B::cvm_partition_state(self.shared),
         }
         .set_wait_for_tlb_locks(target_vtl);
diff --git a/openhcl/virt_mshv_vtl/src/processor/mod.rs b/openhcl/virt_mshv_vtl/src/processor/mod.rs
@@ -649,11 +649,9 @@ impl<'p, T: Backing> Processor for UhProcessor<'p, T> {
                 return Err(VpHaltReason::Cancel);
             }
         } else {
-            {
-                let mut current = Default::default();
-                affinity::get_current_thread_affinity(&mut current).unwrap();
-                assert_eq!(&current, CpuSet::new().set(self.inner.cpu_index));
-            }
+            let mut current = Default::default();
+            affinity::get_current_thread_affinity(&mut current).unwrap();
+            assert_eq!(&current, CpuSet::new().set(self.inner.cpu_index));
 
             // Lower the priority of this VP thread so that the VM does not return
             // to VTL0 while there is still outstanding VTL2 work to do.
@@ -791,7 +789,7 @@ impl<'a, T: Backing> UhProcessor<'a, T> {
         let inner = partition.vp(vp_info.base.vp_index).unwrap();
         let mut runner = partition
             .hcl
-            .runner(inner.cpu_index, idle_control.is_none())
+            .runner(inner.vp_index().index(), idle_control.is_none())
             .unwrap();
 
         let backing_shared = T::shared(&partition.backing_shared);
@@ -918,7 +916,7 @@ impl<'a, T: Backing> UhProcessor<'a, T> {
     }
 
     fn vp_index(&self) -> VpIndex {
-        self.inner.vp_info.base.vp_index
+        self.inner.vp_index()
     }
 
     #[cfg(guest_arch = "x86_64")]
diff --git a/openhcl/virt_mshv_vtl/src/processor/snp/mod.rs b/openhcl/virt_mshv_vtl/src/processor/snp/mod.rs
@@ -227,7 +227,7 @@ impl HardwareIsolatedBacking for SnpBacked {
         shared: &'a Self::Shared,
     ) -> impl TlbFlushLockAccess + 'a {
         SnpTlbLockFlushAccess {
-            vp_index: vp_index.index() as usize,
+            vp_index,
             partition,
             shared,
         }
@@ -2352,7 +2352,7 @@ impl<T: CpuIo> UhHypercallHandler<'_, '_, T, SnpBacked> {
 }
 
 struct SnpTlbLockFlushAccess<'a> {
-    vp_index: usize,
+    vp_index: VpIndex,
     partition: &'a UhPartitionInner,
     shared: &'a SnpBackedShared,
 }
diff --git a/openhcl/virt_mshv_vtl/src/processor/tdx/mod.rs b/openhcl/virt_mshv_vtl/src/processor/tdx/mod.rs
@@ -537,7 +537,7 @@ impl HardwareIsolatedBacking for TdxBacked {
         shared: &'a Self::Shared,
     ) -> impl TlbFlushLockAccess + 'a {
         TdxTlbLockFlushAccess {
-            vp_index: vp_index.index() as usize,
+            vp_index,
             partition,
             shared,
         }
@@ -3425,7 +3425,7 @@ impl<T: CpuIo> hv1_hypercall::FlushVirtualAddressListEx
 
         // Send flush IPIs to the specified VPs.
         TdxTlbLockFlushAccess {
-            vp_index: self.vp.vp_index().index() as usize,
+            vp_index: self.vp.vp_index(),
             partition: self.vp.partition,
             shared: self.vp.shared,
         }
@@ -3478,7 +3478,7 @@ impl<T: CpuIo> hv1_hypercall::FlushVirtualAddressSpaceEx
 
         // Send flush IPIs to the specified VPs.
         TdxTlbLockFlushAccess {
-            vp_index: self.vp.vp_index().index() as usize,
+            vp_index: self.vp.vp_index(),
             partition: self.vp.partition,
             shared: self.vp.shared,
         }
@@ -3552,7 +3552,7 @@ impl TdxTlbLockFlushAccess<'_> {
         // for each VP.
         std::sync::atomic::fence(Ordering::SeqCst);
         for target_vp in processors {
-            if self.vp_index != target_vp
+            if self.vp_index.index() as usize != target_vp
                 && self.shared.active_vtl[target_vp].load(Ordering::Relaxed) == target_vtl as u8
             {
                 self.partition.vps[target_vp].wake_vtl2();
@@ -3564,7 +3564,7 @@ impl TdxTlbLockFlushAccess<'_> {
 }
 
 struct TdxTlbLockFlushAccess<'a> {
-    vp_index: usize,
+    vp_index: VpIndex,
     partition: &'a UhPartitionInner,
     shared: &'a TdxBackedShared,
 }

Original file line number	Diff line number	Diff line change
`@@ -227,7 +227,7 @@ impl HardwareIsolatedBacking for SnpBacked {`
`227`	`227`	`shared: &'a Self::Shared,`
`228`	`228`	`) -> impl TlbFlushLockAccess + 'a {`
`229`	`229`	`SnpTlbLockFlushAccess {`
`230`		`- vp_index: vp_index.index() as usize,`
	`230`	`+ vp_index,`
`231`	`231`	`partition,`
`232`	`232`	`shared,`
`233`	`233`	`}`
`@@ -2352,7 +2352,7 @@ impl<T: CpuIo> UhHypercallHandler<'_, '_, T, SnpBacked> {`
`2352`	`2352`	`}`
`2353`	`2353`
`2354`	`2354`	`struct SnpTlbLockFlushAccess<'a> {`
`2355`		`- vp_index: usize,`
	`2355`	`+ vp_index: VpIndex,`
`2356`	`2356`	`partition: &'a UhPartitionInner,`
`2357`	`2357`	`shared: &'a SnpBackedShared,`
`2358`	`2358`	`}`