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/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,