add option for custom vCPU count

Only supported by VirtualMachineRawConfig.

It has been observed that >1 vCPU hurts performance in some cases when
not also setting the CPU affinities and performing other tweaks, as done
by MATCH_HOST, so we avoided adding this before. However, there are
probably use cases where it is the best option and we've been explicitly
asked by partners to add it.

Bug: 381281047
Test: m
Change-Id: Ic22c747d064e70505398eb438a210ef42b5d75e1
diff --git a/android/virtmgr/src/aidl.rs b/android/virtmgr/src/aidl.rs
index 79c7d81..15a80a6 100644
--- a/android/virtmgr/src/aidl.rs
+++ b/android/virtmgr/src/aidl.rs
@@ -653,6 +653,14 @@
         let (cpus, host_cpu_topology) = match config.cpuTopology {
             CpuTopology::MATCH_HOST => (None, true),
             CpuTopology::ONE_CPU => (NonZeroU32::new(1), false),
+            CpuTopology::CUSTOM => (
+                NonZeroU32::new(
+                    u32::try_from(config.customVcpuCount)
+                        .context("bad customVcpuCount")
+                        .or_binder_exception(ExceptionCode::ILLEGAL_ARGUMENT)?,
+                ),
+                false,
+            ),
             val => {
                 return Err(anyhow!("Failed to parse CPU topology value {:?}", val))
                     .with_log()
@@ -1228,6 +1236,9 @@
     vm_config.name.clone_from(&config.name);
     vm_config.protectedVm = config.protectedVm;
     vm_config.cpuTopology = config.cpuTopology;
+    if config.cpuTopology == CpuTopology::CUSTOM {
+        bail!("AppConfig doesn't support CpuTopology::CUSTOM");
+    }
     vm_config.hugePages = config.hugePages || vm_payload_config.hugepages;
     vm_config.boostUclamp = config.boostUclamp;
 
diff --git a/android/virtmgr/src/atom.rs b/android/virtmgr/src/atom.rs
index 45b020e..e0fed85 100644
--- a/android/virtmgr/src/atom.rs
+++ b/android/virtmgr/src/atom.rs
@@ -92,6 +92,26 @@
     }
 }
 
+fn get_num_vcpus(cpu_topology: CpuTopology, custom_vcpu_count: Option<i32>) -> i32 {
+    match cpu_topology {
+        CpuTopology::ONE_CPU => 1,
+        CpuTopology::MATCH_HOST => {
+            get_num_cpus().and_then(|v| v.try_into().ok()).unwrap_or_else(|| {
+                warn!("Failed to determine the number of CPUs in the host");
+                INVALID_NUM_CPUS
+            })
+        }
+        CpuTopology::CUSTOM => custom_vcpu_count.unwrap_or_else(|| {
+            warn!("AppConfig doesn't support CpuTopology::CUSTOM");
+            INVALID_NUM_CPUS
+        }),
+        _ => {
+            warn!("invalid CpuTopology: {cpu_topology:?}");
+            INVALID_NUM_CPUS
+        }
+    }
+}
+
 /// Write the stats of VMCreation to statsd
 /// The function creates a separate thread which waits for statsd to start to push atom
 pub fn write_vm_creation_stats(
@@ -115,33 +135,24 @@
             binder_exception_code = e.exception_code() as i32;
         }
     }
-    let (vm_identifier, config_type, cpu_topology, memory_mib, apexes) = match config {
+
+    let (vm_identifier, config_type, num_cpus, memory_mib, apexes) = match config {
         VirtualMachineConfig::AppConfig(config) => (
             config.name.clone(),
             vm_creation_requested::ConfigType::VirtualMachineAppConfig,
-            config.cpuTopology,
+            get_num_vcpus(config.cpuTopology, None),
             config.memoryMib,
             get_apex_list(config),
         ),
         VirtualMachineConfig::RawConfig(config) => (
             config.name.clone(),
             vm_creation_requested::ConfigType::VirtualMachineRawConfig,
-            config.cpuTopology,
+            get_num_vcpus(config.cpuTopology, Some(config.customVcpuCount)),
             config.memoryMib,
             String::new(),
         ),
     };
 
-    let num_cpus: i32 = match cpu_topology {
-        CpuTopology::MATCH_HOST => {
-            get_num_cpus().and_then(|v| v.try_into().ok()).unwrap_or_else(|| {
-                warn!("Failed to determine the number of CPUs in the host");
-                INVALID_NUM_CPUS
-            })
-        }
-        _ => 1,
-    };
-
     let atom = AtomVmCreationRequested {
         uid: get_calling_uid() as i32,
         vmIdentifier: vm_identifier,
diff --git a/android/virtualizationservice/aidl/android/system/virtualizationservice/CpuTopology.aidl b/android/virtualizationservice/aidl/android/system/virtualizationservice/CpuTopology.aidl
index 8a8e3d0..31a33f4 100644
--- a/android/virtualizationservice/aidl/android/system/virtualizationservice/CpuTopology.aidl
+++ b/android/virtualizationservice/aidl/android/system/virtualizationservice/CpuTopology.aidl
@@ -22,4 +22,6 @@
     ONE_CPU = 0,
     /** Match physical CPU topology of the host. */
     MATCH_HOST = 1,
+    /** Number of vCPUs specified in the config. */
+    CUSTOM = 2,
 }
diff --git a/android/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl b/android/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
index 3393546..d98fdcc 100644
--- a/android/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
+++ b/android/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
@@ -65,6 +65,9 @@
     /** The vCPU topology that will be generated for the VM. Default to 1 vCPU. */
     CpuTopology cpuTopology = CpuTopology.ONE_CPU;
 
+    /** The number of vCPUs. Ignored unless `cpuTopology == CUSTOM`. */
+    int customVcpuCount;
+
     /**
      * A version or range of versions of the virtual platform that this config is compatible with.
      * The format follows SemVer.