Add options for configuring number of vCPUs and CPU affinity

Bug: 197358423
Test: atest MicrodroidHostTestCases

Change-Id: I61a7e746ddd83a1816d18166fb74f4aa5a2565ce
diff --git a/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java b/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
index acaead0..819061b 100644
--- a/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
+++ b/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
@@ -43,6 +43,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -119,7 +120,9 @@
                         packageName,
                         configPath,
                         /* debug */ true,
-                        /* use default memoryMib */ 0);
+                        /* use default memoryMib */ 0,
+                        Optional.empty(),
+                        Optional.empty());
         adbConnectToMicrodroid(androidDevice, sCid);
 
         // Root because authfs (started from shell in this test) currently require root to open
diff --git a/compos/tests/java/android/compos/test/ComposKeyTestCase.java b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
index 961f34a..d7c0058 100644
--- a/compos/tests/java/android/compos/test/ComposKeyTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
@@ -32,6 +32,8 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.Optional;
+
 @RootPermissionTest
 @RunWith(DeviceJUnit4ClassRunner.class)
 public final class ComposKeyTestCase extends VirtualizationTestCaseBase {
@@ -132,7 +134,9 @@
                         packageName,
                         "assets/vm_test_config.json",
                         /* debug */ true,
-                        /* use default memoryMib */ 0);
+                        /* use default memoryMib */ 0,
+                        Optional.empty(),
+                        Optional.empty());
         adbConnectToMicrodroid(getDevice(), mCid);
     }
 
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index dba4c8f..2ced5b0 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -34,6 +34,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.regex.Pattern;
 
 /**
  * Represents a configuration of a virtual machine. A configuration consists of hardware
@@ -51,6 +52,8 @@
     private static final String KEY_PAYLOADCONFIGPATH = "payloadConfigPath";
     private static final String KEY_DEBUGLEVEL = "debugLevel";
     private static final String KEY_MEMORY_MIB = "memoryMib";
+    private static final String KEY_NUM_CPUS = "numCpu";
+    private static final String KEY_CPU_AFFINITY = "cpuAffinity";
 
     // Paths to the APK file of this application.
     private final @NonNull String mApkPath;
@@ -85,6 +88,18 @@
     private final int mMemoryMib;
 
     /**
+     * Number of vCPUs in the VM. Defaults to 1 when not specified.
+     */
+    private final int mNumCpus;
+
+    /**
+     * Comma-separated list of CPUs or CPU ranges to run vCPUs on (e.g. 0,1-3,5), or
+     * colon-separated list of assignments of vCPU to host CPU assignments (e.g. 0=0:1=1:2=2).
+     * Default is no mask which means a vCPU can run on any host CPU.
+     */
+    private final String mCpuAffinity;
+
+    /**
      * Path within the APK to the payload config file that defines software aspects of this config.
      */
     private final @NonNull String mPayloadConfigPath;
@@ -96,12 +111,16 @@
             @NonNull Signature[] certs,
             @NonNull String payloadConfigPath,
             DebugLevel debugLevel,
-            int memoryMib) {
+            int memoryMib,
+            int numCpus,
+            String cpuAffinity) {
         mApkPath = apkPath;
         mCerts = certs;
         mPayloadConfigPath = payloadConfigPath;
         mDebugLevel = debugLevel;
         mMemoryMib = memoryMib;
+        mNumCpus = numCpus;
+        mCpuAffinity = cpuAffinity;
     }
 
     /** Loads a config from a stream, for example a file. */
@@ -131,7 +150,10 @@
         }
         final DebugLevel debugLevel = DebugLevel.values()[b.getInt(KEY_DEBUGLEVEL)];
         final int memoryMib = b.getInt(KEY_MEMORY_MIB);
-        return new VirtualMachineConfig(apkPath, certs, payloadConfigPath, debugLevel, memoryMib);
+        final int numCpus = b.getInt(KEY_NUM_CPUS);
+        final String cpuAffinity = b.getString(KEY_CPU_AFFINITY);
+        return new VirtualMachineConfig(apkPath, certs, payloadConfigPath, debugLevel, memoryMib,
+                numCpus, cpuAffinity);
     }
 
     /** Persists this config to a stream, for example a file. */
@@ -198,6 +220,8 @@
                 break;
         }
         parcel.memoryMib = mMemoryMib;
+        parcel.numCpus = mNumCpus;
+        parcel.cpuAffinity = mCpuAffinity;
         return parcel;
     }
 
@@ -207,6 +231,8 @@
         private String mPayloadConfigPath;
         private DebugLevel mDebugLevel;
         private int mMemoryMib;
+        private int mNumCpus;
+        private String mCpuAffinity;
         // TODO(jiyong): add more items like # of cpu, size of ram, debuggability, etc.
 
         /** Creates a builder for the given context (APK), and the payload config file in APK. */
@@ -214,6 +240,8 @@
             mContext = context;
             mPayloadConfigPath = payloadConfigPath;
             mDebugLevel = DebugLevel.NONE;
+            mNumCpus = 1;
+            mCpuAffinity = null;
         }
 
         /** Sets the debug level */
@@ -231,6 +259,25 @@
             return this;
         }
 
+        /**
+         * Sets the number of vCPUs in the VM. Defaults to 1.
+         */
+        public Builder numCpus(int num) {
+            mNumCpus = num;
+            return this;
+        }
+
+        /**
+         * Sets on which host CPUs the vCPUs can run. The format is 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 this can be 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.
+         */
+        public Builder cpuAffinity(String affinity) {
+            mCpuAffinity = affinity;
+            return this;
+        }
+
         /** Builds an immutable {@link VirtualMachineConfig} */
         public @NonNull VirtualMachineConfig build() {
             final String apkPath = mContext.getPackageCodePath();
@@ -248,8 +295,22 @@
                 throw new RuntimeException(e);
             }
 
+            final int availableCpus = Runtime.getRuntime().availableProcessors();
+            if (mNumCpus < 0 || mNumCpus > availableCpus) {
+                throw new IllegalArgumentException("Number of vCPUs (" + mNumCpus + ") is out of "
+                        + "range [1, " + availableCpus + "]");
+            }
+
+            if (mCpuAffinity != null
+                    && !Pattern.matches("[\\d]+(-[\\d]+)?(,[\\d]+(-[\\d]+)?)*", mCpuAffinity)
+                    && !Pattern.matches("[\\d]+=[\\d]+(:[\\d]+=[\\d]+)*", mCpuAffinity)) {
+                throw new IllegalArgumentException("CPU affinity [" + mCpuAffinity + "]"
+                        + " is invalid");
+            }
+
             return new VirtualMachineConfig(
-                    apkPath, certs, mPayloadConfigPath, mDebugLevel, mMemoryMib);
+                    apkPath, certs, mPayloadConfigPath, mDebugLevel, mMemoryMib, mNumCpus,
+                    mCpuAffinity);
         }
     }
 }
diff --git a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
index c71d6ac..528f7c2 100644
--- a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
+++ b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
@@ -34,6 +34,7 @@
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.util.Arrays;
+import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.regex.Matcher;
@@ -182,7 +183,9 @@
             String packageName,
             String configPath,
             boolean debug,
-            int memoryMib)
+            int memoryMib,
+            Optional<Integer> numCpus,
+            Optional<String> cpuAffinity)
             throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(androidDevice);
 
@@ -216,6 +219,8 @@
                         "--daemonize",
                         "--log " + logPath,
                         "--mem " + memoryMib,
+                        numCpus.isPresent() ? "--cpus " + numCpus.get() : "",
+                        cpuAffinity.isPresent() ? "--cpu-affinity " + cpuAffinity.get() : "",
                         debugFlag,
                         apkPath,
                         outApkIdsigPath,
diff --git a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
index c2ef566..7d8cf85 100644
--- a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
+++ b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
@@ -29,6 +29,8 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.Optional;
+
 @RunWith(DeviceJUnit4ClassRunner.class)
 public class MicrodroidTestCase extends VirtualizationTestCaseBase {
     private static final String APK_NAME = "MicrodroidTestApp.apk";
@@ -37,6 +39,10 @@
     private static final int MIN_MEM_ARM64 = 256;
     private static final int MIN_MEM_X86_64 = 400;
 
+    // Number of vCPUs and their affinity to host CPUs for testing purpose
+    private static final int NUM_VCPUS = 3;
+    private static final String CPU_AFFINITY = "0,1,2";
+
     private int minMemorySize() throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(getDevice());
         String abi = android.run("getprop", "ro.product.cpu.abi");
@@ -61,7 +67,9 @@
                         PACKAGE_NAME,
                         configPath,
                         /* debug */ true,
-                        minMemorySize());
+                        minMemorySize(),
+                        Optional.of(NUM_VCPUS),
+                        Optional.of(CPU_AFFINITY));
         adbConnectToMicrodroid(getDevice(), cid);
 
         // Wait until logd-init starts. The service is one of the last services that are started in
@@ -101,6 +109,9 @@
         // Check that no denials have happened so far
         assertThat(runOnMicrodroid("logcat -d -e 'avc:[[:space:]]{1,2}denied'"), is(""));
 
+        assertThat(runOnMicrodroid("cat /proc/cpuinfo | grep processor | wc -l"),
+                is(Integer.toString(NUM_VCPUS)));
+
         shutdownMicrodroid(getDevice(), cid);
     }
 
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
index 0cb187c..8265f96 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
@@ -52,4 +52,16 @@
      * the value in microdroid.json, if any, or the crosvm default.
      */
     int memoryMib;
+
+    /**
+     * Number of vCPUs in the VM. Defaults to 1.
+     */
+    int numCpus = 1;
+
+    /**
+     * Comma-separated list of CPUs or CPU ranges to run vCPUs on (e.g. 0,1-3,5), or
+     * colon-separated list of assignments of vCPU to host CPU assignments (e.g. 0=0:1=1:2=2).
+     * Default is no mask which means a vCPU can run on any host CPU.
+     */
+    @nullable String cpuAffinity;
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
index c62117e..bb4eef0 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
@@ -45,4 +45,16 @@
 
     /** The amount of RAM to give the VM, in MiB. 0 or negative to use the default. */
     int memoryMib;
+
+    /**
+     * Number of vCPUs in the VM. Defaults to 1.
+     */
+    int numCpus = 1;
+
+    /**
+     * Comma-separated list of CPUs or CPU ranges to run vCPUs on (e.g. 0,1-3,5), or
+     * colon-separated list of assignments of vCPU to host CPU assignments (e.g. 0=0:1=1:2=2).
+     * Default is no mask which means a vCPU can run on any host CPU.
+     */
+    @nullable String cpuAffinity;
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 1bd7ee0..c264270 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -266,6 +266,8 @@
             params: config.params.to_owned(),
             protected,
             memory_mib: config.memoryMib.try_into().ok().and_then(NonZeroU32::new),
+            cpus: config.numCpus.try_into().ok().and_then(NonZeroU32::new),
+            cpu_affinity: config.cpuAffinity.clone(),
             console_fd,
             log_fd,
             indirect_files,
@@ -563,6 +565,9 @@
         vm_config.memoryMib = config.memoryMib;
     }
 
+    vm_config.numCpus = config.numCpus;
+    vm_config.cpuAffinity = config.cpuAffinity.clone();
+
     // Microdroid requires an additional payload disk image and the bootconfig partition.
     if os_name == "microdroid" {
         add_microdroid_images(
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index bf1ff0c..0b1429c 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -45,6 +45,8 @@
     pub params: Option<String>,
     pub protected: bool,
     pub memory_mib: Option<NonZeroU32>,
+    pub cpus: Option<NonZeroU32>,
+    pub cpu_affinity: Option<String>,
     pub console_fd: Option<File>,
     pub log_fd: Option<File>,
     pub indirect_files: Vec<File>,
@@ -246,6 +248,14 @@
         command.arg("--mem").arg(memory_mib.to_string());
     }
 
+    if let Some(cpus) = config.cpus {
+        command.arg("--cpus").arg(cpus.to_string());
+    }
+
+    if let Some(cpu_affinity) = config.cpu_affinity {
+        command.arg("--cpu-affinity").arg(cpu_affinity);
+    }
+
     // Keep track of what file descriptors should be mapped to the crosvm process.
     let mut preserved_fds = config.indirect_files.iter().map(|file| file.as_raw_fd()).collect();
 
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),
diff --git a/vmconfig/src/lib.rs b/vmconfig/src/lib.rs
index 3828742..9c3e671 100644
--- a/vmconfig/src/lib.rs
+++ b/vmconfig/src/lib.rs
@@ -95,6 +95,7 @@
             disks: self.disks.iter().map(DiskImage::to_parcelable).collect::<Result<_, Error>>()?,
             protectedVm: self.protected,
             memoryMib: memory_mib,
+            ..Default::default()
         })
     }
 }