Add options for configuring number of vCPUs and CPU affinity

Bug: 197358423
Test: atest MicrodroidHostTestCases

Change-Id: I61a7e746ddd83a1816d18166fb74f4aa5a2565ce
diff --git a/vm/src/main.rs b/vm/src/main.rs
index d53305b..a466a4c 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -77,6 +77,14 @@
         #[structopt(short, long)]
         mem: Option<u32>,
 
+        /// Number of vCPUs in the VM. If unspecified, defaults to 1.
+        #[structopt(long)]
+        cpus: Option<u32>,
+
+        /// Host CPUs where vCPUs are run on. If unspecified, vCPU runs on any host CPU.
+        #[structopt(long)]
+        cpu_affinity: Option<String>,
+
         /// Paths to extra idsig files.
         #[structopt(long)]
         extra_idsigs: Vec<PathBuf>,
@@ -91,6 +99,18 @@
         #[structopt(short, long)]
         daemonize: bool,
 
+        /// Number of vCPUs in the VM. If unspecified, defaults to 1.
+        #[structopt(long)]
+        cpus: Option<u32>,
+
+        /// Host CPUs where vCPUs are run on. If unspecified, vCPU runs on any host CPU. The format
+        /// can be either a comma-separated list of CPUs or CPU ranges to run vCPUs on (e.g.
+        /// "0,1-3,5" to choose host CPUs 0, 1, 2, 3, and 5, or a colon-separated list of
+        /// assignments of vCPU-to-host-CPU assignments e.g. "0=0:1=1:2=2" to map vCPU 0 to host
+        /// CPU 0 and so on.
+        #[structopt(long)]
+        cpu_affinity: Option<String>,
+
         /// Path to file for VM console output.
         #[structopt(long)]
         console: Option<PathBuf>,
@@ -155,6 +175,8 @@
             log,
             debug,
             mem,
+            cpus,
+            cpu_affinity,
             extra_idsigs,
         } => command_run_app(
             service,
@@ -167,10 +189,20 @@
             log.as_deref(),
             debug,
             mem,
+            cpus,
+            cpu_affinity,
             &extra_idsigs,
         ),
-        Opt::Run { config, daemonize, console } => {
-            command_run(service, &config, daemonize, console.as_deref(), /* mem */ None)
+        Opt::Run { config, daemonize, cpus, cpu_affinity, console } => {
+            command_run(
+                service,
+                &config,
+                daemonize,
+                console.as_deref(),
+                /* mem */ None,
+                cpus,
+                cpu_affinity,
+            )
         }
         Opt::Stop { cid } => command_stop(service, cid),
         Opt::List => command_list(service),
diff --git a/vm/src/run.rs b/vm/src/run.rs
index 7f5f9fc..19982ea 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -50,6 +50,8 @@
     log_path: Option<&Path>,
     debug_level: DebugLevel,
     mem: Option<u32>,
+    cpus: Option<u32>,
+    cpu_affinity: Option<String>,
     extra_idsigs: &[PathBuf],
 ) -> Result<(), Error> {
     let extra_apks = parse_extra_apk_list(apk, config_path)?;
@@ -98,6 +100,8 @@
         configPath: config_path.to_owned(),
         debugLevel: debug_level,
         memoryMib: mem.unwrap_or(0) as i32, // 0 means use the VM default
+        numCpus: cpus.unwrap_or(1) as i32,
+        cpuAffinity: cpu_affinity,
     });
     run(
         service,
@@ -116,6 +120,8 @@
     daemonize: bool,
     console_path: Option<&Path>,
     mem: Option<u32>,
+    cpus: Option<u32>,
+    cpu_affinity: Option<String>,
 ) -> Result<(), Error> {
     let config_file = File::open(config_path).context("Failed to open config file")?;
     let mut config =
@@ -123,6 +129,10 @@
     if let Some(mem) = mem {
         config.memoryMib = mem as i32;
     }
+    if let Some(cpus) = cpus {
+        config.numCpus = cpus as i32;
+    }
+    config.cpuAffinity = cpu_affinity;
     run(
         service,
         &VirtualMachineConfig::RawConfig(config),