underhill_core: specify which CPUs have outstanding IO, not just if they have interrupts

openhcl/openhcl_boot/src/host_params/dt/mod.rs

@@ -379,6 +379,7 @@ struct PartitionTopology {
 struct PersistedPartitionTopology {
     topology: PartitionTopology,
     cpus_with_mapped_interrupts: Vec<u32>,
+    cpus_with_outstanding_io: Vec<u32>,
 }
 
 // Calculate the default mmio size for VTL2 when not specified by the host.
@@ -615,6 +616,7 @@ fn topology_from_persisted_state(
         partition_memory,
         partition_mmio,
         cpus_with_mapped_interrupts,
+        cpus_with_outstanding_io,
     } = parsed_protobuf;
 
     // FUTURE: should memory allocation mode should persist in saved state and
@@ -764,6 +766,7 @@ fn topology_from_persisted_state(
             memory_allocation_mode,
         },
         cpus_with_mapped_interrupts,
+        cpus_with_outstanding_io,
     })
 }
 

openhcl/underhill_core/src/dispatch/mod.rs

@@ -627,17 +627,14 @@ impl LoadedVm {
             };
 
             // Save the persisted state used by the next openhcl_boot.
-            let cpus_with_mapped_interrupts = match state
+            let nvme_vp_interrupt_state = crate::nvme_manager::save_restore_helpers::cpus_with_interrupts(state
                 .init_state
-                .nvme_state
-                .as_ref() {
-                    Some(nvme_state) => crate::nvme_manager::save_restore::cpus_with_interrupts(&nvme_state.nvme_state),
-                    None => vec![],
-                };
+                .nvme_state.as_ref().map(|n| &n.nvme_state));
 
             crate::loader::vtl2_config::write_persisted_info(
                 self.runtime_params.parsed_openhcl_boot(),
-                cpus_with_mapped_interrupts,
+                nvme_vp_interrupt_state.vps_with_mapped_interrupts,
+                nvme_vp_interrupt_state.vps_with_outstanding_io
             )
             .context("failed to write persisted info")?;
 

openhcl/underhill_core/src/loader/vtl2_config/mod.rs

@@ -204,6 +204,7 @@ impl Drop for Vtl2ParamsMap<'_> {
 pub fn write_persisted_info(
     parsed: &ParsedBootDtInfo,
     cpus_with_mapped_interrupts: Vec<u32>,
+    cpus_with_outstanding_io: Vec<u32>,
 ) -> anyhow::Result<()> {
     use loader_defs::shim::PersistedStateHeader;
     use loader_defs::shim::save_restore::MemoryEntry;
@@ -249,6 +250,7 @@ pub fn write_persisted_info(
             })
             .collect(),
         cpus_with_mapped_interrupts,
+        cpus_with_outstanding_io,
     };
 
     let protobuf = mesh_protobuf::encode(state);

openhcl/underhill_core/src/nvme_manager/mod.rs

@@ -58,6 +58,7 @@ use vmcore::vm_task::VmTaskDriverSource;
 pub mod device;
 pub mod manager;
 pub mod save_restore;
+pub mod save_restore_helpers;
 
 #[derive(Debug, Error)]
 #[error("nvme device {pci_id} error")]

openhcl/underhill_core/src/nvme_manager/save_restore.rs

@@ -2,7 +2,6 @@
 // Licensed under the MIT License.
 
 use mesh::payload::Protobuf;
-use std::collections::BTreeSet;
 use vmcore::save_restore::SavedStateRoot;
 
 #[derive(Protobuf, SavedStateRoot)]
@@ -22,15 +21,3 @@ pub struct NvmeSavedDiskConfig {
     #[mesh(2)]
     pub driver_state: nvme_driver::NvmeDriverSavedState,
 }
-
-/// Returns a sorted list of CPU IDs that have mapped device interrupts in the saved NVMe state.
-///
-/// This information is used to make heuristic decisions during restore, such as whether to
-/// disable sidecar for VMs with active device interrupts.
-pub fn cpus_with_interrupts(state: &NvmeManagerSavedState) -> Vec<u32> {
-    let mut cpus_with_interrupts = BTreeSet::new();
-    for disk in &state.nvme_disks {
-        cpus_with_interrupts.extend(disk.driver_state.worker_data.io.iter().map(|q| q.cpu));
-    }
-    cpus_with_interrupts.into_iter().collect()
-}

openhcl/underhill_core/src/nvme_manager/save_restore_helpers.rs

@@ -0,0 +1,57 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use crate::nvme_manager::save_restore::NvmeManagerSavedState;
+use std::collections::BTreeMap;
+use std::collections::btree_map::Entry;
+
+/// Useful state about how the VM's vCPUs interacted with NVMe device interrupts at the time of save.
+///
+/// This information is used to make heuristic decisions during restore, such as whether to
+/// disable sidecar for VMs with active device interrupts.
+pub struct VPInterruptState {
+    /// List of vCPUs with any mapped device interrupts, sorted by CPU ID.
+    pub vps_with_mapped_interrupts: Vec<u32>,
+
+    /// List of vCPUs with outstanding I/O at the time of save, sorted by CPU ID.
+    /// It is expected that this is a subset of `vps_with_mapped_interrupts`, since
+    /// only some queues will have in-flight I/O.
+    pub vps_with_outstanding_io: Vec<u32>,
+}
+
+/// Analyzes the saved NVMe manager state to determine which vCPUs had mapped device interrupts
+/// and which had outstanding I/O at the time of save.
+///
+/// See [`VPInterruptState`] for more details.
+pub fn cpus_with_interrupts(state: Option<&NvmeManagerSavedState>) -> VPInterruptState {
+    let mut vp_state = BTreeMap::new();
+
+    if let Some(state) = state {
+        for disk in &state.nvme_disks {
+            for q in &disk.driver_state.worker_data.io {
+                match vp_state.entry(q.cpu) {
+                    Entry::Vacant(e) => {
+                        e.insert(!q.queue_data.handler_data.pending_cmds.commands.is_empty());
+                    }
+                    Entry::Occupied(mut e) => {
+                        e.insert(
+                            e.get() | !q.queue_data.handler_data.pending_cmds.commands.is_empty(),
+                        );
+                    }
+                }
+            }
+        }
+    }
+
+    VPInterruptState {
+        vps_with_mapped_interrupts: vp_state.keys().cloned().collect(),
+        vps_with_outstanding_io: vp_state
+            .iter()
+            .filter_map(
+                |(&vp, &has_outstanding_io)| {
+                    if has_outstanding_io { Some(vp) } else { None }
+                },
+            )
+            .collect(),
+    }
+}

vm/loader/loader_defs/src/shim.rs

@@ -346,5 +346,18 @@ pub mod save_restore {
         /// still be functionally correct.
         #[mesh(3)]
         pub cpus_with_mapped_interrupts: Vec<u32>,
+        /// The list of CPUs with mapped device interrupts present at save time,
+        /// and that have outstanding IO on that CPU.
+        ///
+        /// While this list is today used as a semaphore (either there are device
+        /// interrupts mapped or not), in the future it may be used to provide more
+        /// granular restore hints. E.g., only start the CPUs with active
+        /// interrupts right away and defer other CPU startup until later.
+        ///
+        /// DEFAULT: For save state from prior versions, this will be empty.
+        /// This is fine: the restore heuristics might be less optimal, but will
+        /// still be functionally correct.
+        #[mesh(4)]
+        pub cpus_with_outstanding_io: Vec<u32>,
     }
 }