Reapply "openhcl_boot: set aside private pool memory based on heuristics (microsoft#2215)" (microsoft#2339)

This reverts commit c552f9a.

Author: mattkur

Cargo.lock

@@ -5021,7 +5021,9 @@ dependencies = [
  "sidecar_defs",
  "string_page_buf",
  "tdcall",
+ "test_with_tracing",
  "thiserror 2.0.16",
+ "tracing",
  "underhill_confidentiality",
  "x86defs",
  "zerocopy 0.8.25",

openhcl/openhcl_boot/Cargo.toml

@@ -37,6 +37,10 @@ safe_intrinsics.workspace = true
 tdcall.workspace = true
 x86defs.workspace = true
 
+[dev-dependencies]
+test_with_tracing.workspace = true
+tracing.workspace = true
+
 [build-dependencies]
 minimal_rt_build.workspace = true
 

openhcl/openhcl_boot/src/cmdline.rs

@@ -3,18 +3,23 @@
 
 //! Command line arguments and parsing for openhcl_boot.
 
+use crate::boot_logger::log;
 use underhill_confidentiality::OPENHCL_CONFIDENTIAL_DEBUG_ENV_VAR_NAME;
 
-/// Enable the private VTL2 GPA pool for page allocations. This is only enabled
-/// via the command line, because in order to support the VTL2 GPA pool
-/// generically, the boot shim must read serialized data from the previous
-/// OpenHCL instance on a servicing boot in order to guarantee the same memory
-/// layout is presented.
+/// Enable the private VTL2 GPA pool for page allocations.
 ///
-/// The value specified is the number of 4K pages to reserve for the pool.
+/// Possible values:
+/// * `release`: Use the release version of the lookup table (default), or device tree.
+/// * `debug`: Use the debug version of the lookup table, or device tree.
+/// * `off`: Disable the VTL2 GPA pool.
+/// * `<num_pages>`: Explicitly specify the size of the VTL2 GPA pool.
 ///
-/// TODO: Remove this commandline once support for reading saved state is
-/// supported in openhcl_boot.
+/// See `Vtl2GpaPoolConfig` for more details.
+const IGVM_VTL2_GPA_POOL_CONFIG: &str = "OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG=";
+
+/// Test-legacy/test-compat override for `OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG`.
+/// (otherwise, tests cannot modify the VTL2 GPA pool config different from what
+/// may be in the manifest).
 const ENABLE_VTL2_GPA_POOL: &str = "OPENHCL_ENABLE_VTL2_GPA_POOL=";
 
 /// Options controlling sidecar.
@@ -29,10 +34,59 @@ const SIDECAR: &str = "OPENHCL_SIDECAR=";
 /// Disable NVME keep alive regardless if the host supports it.
 const DISABLE_NVME_KEEP_ALIVE: &str = "OPENHCL_DISABLE_NVME_KEEP_ALIVE=";
 
+/// Lookup table to use for VTL2 GPA pool size heuristics.
+#[derive(Debug, PartialEq, Clone, Copy)]
+pub enum Vtl2GpaPoolLookupTable {
+    Release,
+    Debug,
+}
+
+#[derive(Debug, PartialEq, Clone, Copy)]
+pub enum Vtl2GpaPoolConfig {
+    /// Use heuristics to determine the VTL2 GPA pool size.
+    /// Reserve a default size based on the amount of VTL2 ram and
+    /// number of vCPUs. The point of this method is to account for cases where
+    /// we retrofit the private pool into existing deployments that do not
+    /// specify it explicitly.
+    ///
+    /// If the host specifies a size via the device tree, that size will be used
+    /// instead.
+    ///
+    /// The lookup table specifies whether to use the debug or release
+    /// heuristics (as the dev manifests provide different amounts of VTL2 RAM).
+    Heuristics(Vtl2GpaPoolLookupTable),
+
+    /// Explicitly disable the VTL2 private pool.
+    Off,
+
+    /// Explicitly specify the size of the VTL2 GPA pool in pages.
+    Pages(u64),
+}
+
+impl<S: AsRef<str>> From<S> for Vtl2GpaPoolConfig {
+    fn from(arg: S) -> Self {
+        match arg.as_ref() {
+            "debug" => Vtl2GpaPoolConfig::Heuristics(Vtl2GpaPoolLookupTable::Debug),
+            "release" => Vtl2GpaPoolConfig::Heuristics(Vtl2GpaPoolLookupTable::Release),
+            "off" => Vtl2GpaPoolConfig::Off,
+            _ => {
+                let num = arg.as_ref().parse::<u64>().unwrap_or(0);
+                // A size of 0 or failure to parse is treated as disabling
+                // the pool.
+                if num == 0 {
+                    Vtl2GpaPoolConfig::Off
+                } else {
+                    Vtl2GpaPoolConfig::Pages(num)
+                }
+            }
+        }
+    }
+}
+
 #[derive(Debug, PartialEq)]
 pub struct BootCommandLineOptions {
     pub confidential_debug: bool,
-    pub enable_vtl2_gpa_pool: Option<u64>,
+    pub enable_vtl2_gpa_pool: Vtl2GpaPoolConfig,
     pub sidecar: bool,
     pub sidecar_logging: bool,
     pub disable_nvme_keep_alive: bool,
@@ -42,7 +96,7 @@ impl BootCommandLineOptions {
     pub const fn new() -> Self {
         BootCommandLineOptions {
             confidential_debug: false,
-            enable_vtl2_gpa_pool: None,
+            enable_vtl2_gpa_pool: Vtl2GpaPoolConfig::Heuristics(Vtl2GpaPoolLookupTable::Release), // use the release config by default
             sidecar: true, // sidecar is enabled by default
             sidecar_logging: false,
             disable_nvme_keep_alive: false,
@@ -53,19 +107,25 @@ impl BootCommandLineOptions {
 impl BootCommandLineOptions {
     /// Parse arguments from a command line.
     pub fn parse(&mut self, cmdline: &str) {
+        let mut override_vtl2_gpa_pool: Option<Vtl2GpaPoolConfig> = None;
         for arg in cmdline.split_whitespace() {
             if arg.starts_with(OPENHCL_CONFIDENTIAL_DEBUG_ENV_VAR_NAME) {
                 let arg = arg.split_once('=').map(|(_, arg)| arg);
                 if arg.is_some_and(|a| a != "0") {
                     self.confidential_debug = true;
                 }
+            } else if arg.starts_with(IGVM_VTL2_GPA_POOL_CONFIG) {
+                if let Some((_, arg)) = arg.split_once('=') {
+                    self.enable_vtl2_gpa_pool = Vtl2GpaPoolConfig::from(arg);
+                } else {
+                    log!("WARNING: Missing value for IGVM_VTL2_GPA_POOL_CONFIG argument");
+                }
             } else if arg.starts_with(ENABLE_VTL2_GPA_POOL) {
-                self.enable_vtl2_gpa_pool = arg.split_once('=').and_then(|(_, arg)| {
-                    let num = arg.parse::<u64>().unwrap_or(0);
-                    // A size of 0 or failure to parse is treated as disabling
-                    // the pool.
-                    if num == 0 { None } else { Some(num) }
-                });
+                if let Some((_, arg)) = arg.split_once('=') {
+                    override_vtl2_gpa_pool = Some(Vtl2GpaPoolConfig::from(arg));
+                } else {
+                    log!("WARNING: Missing value for ENABLE_VTL2_GPA_POOL argument");
+                }
             } else if arg.starts_with(SIDECAR) {
                 if let Some((_, arg)) = arg.split_once('=') {
                     for arg in arg.split(',') {
@@ -84,6 +144,14 @@ impl BootCommandLineOptions {
                 }
             }
         }
+
+        if let Some(override_config) = override_vtl2_gpa_pool {
+            self.enable_vtl2_gpa_pool = override_config;
+            log!(
+                "INFO: Overriding VTL2 GPA pool config to {:?} from command line",
+                override_config
+            );
+        }
     }
 }
 
@@ -99,34 +167,53 @@ mod tests {
 
     #[test]
     fn test_vtl2_gpa_pool_parsing() {
-        assert_eq!(
-            parse_boot_command_line("OPENHCL_ENABLE_VTL2_GPA_POOL=1"),
-            BootCommandLineOptions {
-                enable_vtl2_gpa_pool: Some(1),
-                ..BootCommandLineOptions::new()
-            }
-        );
-        assert_eq!(
-            parse_boot_command_line("OPENHCL_ENABLE_VTL2_GPA_POOL=0"),
-            BootCommandLineOptions {
-                enable_vtl2_gpa_pool: None,
-                ..BootCommandLineOptions::new()
-            }
-        );
-        assert_eq!(
-            parse_boot_command_line("OPENHCL_ENABLE_VTL2_GPA_POOL=asdf"),
-            BootCommandLineOptions {
-                enable_vtl2_gpa_pool: None,
-                ..BootCommandLineOptions::new()
-            }
-        );
-        assert_eq!(
-            parse_boot_command_line("OPENHCL_ENABLE_VTL2_GPA_POOL=512"),
-            BootCommandLineOptions {
-                enable_vtl2_gpa_pool: Some(512),
-                ..BootCommandLineOptions::new()
-            }
-        );
+        for (cmdline, expected) in [
+            (
+                // default
+                "",
+                Vtl2GpaPoolConfig::Heuristics(Vtl2GpaPoolLookupTable::Release),
+            ),
+            (
+                "OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG=1",
+                Vtl2GpaPoolConfig::Pages(1),
+            ),
+            (
+                "OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG=0",
+                Vtl2GpaPoolConfig::Off,
+            ),
+            (
+                "OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG=asdf",
+                Vtl2GpaPoolConfig::Off,
+            ),
+            (
+                "OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG=512",
+                Vtl2GpaPoolConfig::Pages(512),
+            ),
+            (
+                "OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG=off",
+                Vtl2GpaPoolConfig::Off,
+            ),
+            (
+                "OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG=debug",
+                Vtl2GpaPoolConfig::Heuristics(Vtl2GpaPoolLookupTable::Debug),
+            ),
+            (
+                "OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG=release",
+                Vtl2GpaPoolConfig::Heuristics(Vtl2GpaPoolLookupTable::Release),
+            ),
+            (
+                // OPENHCL_ENABLE_VTL2_GPA_POOL= takes precedence over OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG=
+                "OPENHCL_IGVM_VTL2_GPA_POOL_CONFIG=release OPENHCL_ENABLE_VTL2_GPA_POOL=debug",
+                Vtl2GpaPoolConfig::Heuristics(Vtl2GpaPoolLookupTable::Debug),
+            ),
+        ] {
+            assert_eq!(
+                parse_boot_command_line(cmdline).enable_vtl2_gpa_pool,
+                expected,
+                "Failed parsing VTL2 GPA pool config from command line: {}",
+                cmdline
+            );
+        }
     }
 
     #[test]