Merge "Respect dalvik.vm.boot-dex2oat-[threads|cpu-set]"
diff --git a/compos/Android.bp b/compos/Android.bp
index 18fd28f..ae0c4c3 100644
--- a/compos/Android.bp
+++ b/compos/Android.bp
@@ -25,10 +25,12 @@
         "liblog_rust",
         "libminijail_rust",
         "libnix",
+        "libnum_cpus",
         "libodsign_proto_rust",
         "libprotobuf",
         "libregex",
         "libring",
+        "librustutils",
         "libscopeguard",
     ],
     prefer_rlib: true,
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 46f3020..b54a921 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -39,6 +39,7 @@
 use log::{info, warn};
 use std::fs::File;
 use std::io::{BufRead, BufReader};
+use std::num::NonZeroU32;
 use std::os::raw;
 use std::os::unix::io::IntoRawFd;
 use std::path::Path;
@@ -57,6 +58,11 @@
 pub struct VmParameters {
     /// Whether the VM should be debuggable.
     pub debug_mode: bool,
+    /// Number of vCPUs to have in the VM. If None, defaults to 1.
+    pub cpus: Option<NonZeroU32>,
+    /// Comma separated list of host CPUs where vCPUs are assigned to. If None, any host CPU can be
+    /// used to run any vCPU.
+    pub cpu_set: Option<String>,
     /// If present, overrides the path to the VM config JSON file
     pub config_path: Option<String>,
 }
@@ -113,6 +119,8 @@
             configPath: config_path.to_owned(),
             debugLevel: debug_level,
             extraIdsigs: vec![idsig_manifest_apk_fd],
+            numCpus: parameters.cpus.map_or(1, NonZeroU32::get) as i32,
+            cpuAffinity: parameters.cpu_set.clone(),
             ..Default::default()
         });
 
diff --git a/compos/common/lib.rs b/compos/common/lib.rs
index 66ce8cb..7b0eb7a 100644
--- a/compos/common/lib.rs
+++ b/compos/common/lib.rs
@@ -68,3 +68,9 @@
 /// The path within our config APK of the VM configuration file we use when compiling staged
 /// APEXes before reboot.
 pub const PREFER_STAGED_VM_CONFIG_PATH: &str = "assets/vm_config_staged.json";
+
+/// Number of CPUs to run dex2oat (actually the entire compos VM) with
+pub const DEX2OAT_THREADS_PROP_NAME: &str = "dalvik.vm.boot-dex2oat-threads";
+
+/// Set of host-side CPUs to run dex2oat (actually the entire compos VM) on
+pub const DEX2OAT_CPU_SET_PROP_NAME: &str = "dalvik.vm.boot-dex2oat-cpu-set";
diff --git a/compos/composd/src/instance_manager.rs b/compos/composd/src/instance_manager.rs
index 24ae576..4fc4ad1 100644
--- a/compos/composd/src/instance_manager.rs
+++ b/compos/composd/src/instance_manager.rs
@@ -23,7 +23,13 @@
 use compos_aidl_interface::aidl::com::android::compos::ICompOsService::ICompOsService;
 use compos_aidl_interface::binder::Strong;
 use compos_common::compos_client::VmParameters;
-use compos_common::{PENDING_INSTANCE_DIR, PREFER_STAGED_VM_CONFIG_PATH, TEST_INSTANCE_DIR};
+use compos_common::{
+    DEX2OAT_CPU_SET_PROP_NAME, DEX2OAT_THREADS_PROP_NAME, PENDING_INSTANCE_DIR,
+    PREFER_STAGED_VM_CONFIG_PATH, TEST_INSTANCE_DIR,
+};
+use rustutils::system_properties;
+use std::num::NonZeroU32;
+use std::str::FromStr;
 use std::sync::{Arc, Mutex, Weak};
 use virtualizationservice::IVirtualizationService::IVirtualizationService;
 
@@ -45,7 +51,12 @@
 
     pub fn start_pending_instance(&self) -> Result<Arc<CompOsInstance>> {
         let config_path = Some(PREFER_STAGED_VM_CONFIG_PATH.to_owned());
-        let vm_parameters = VmParameters { config_path, ..Default::default() };
+        let mut vm_parameters = VmParameters { config_path, ..Default::default() };
+        vm_parameters.cpus = NonZeroU32::from_str(
+            &system_properties::read(DEX2OAT_THREADS_PROP_NAME).unwrap_or_default(),
+        )
+        .ok();
+        vm_parameters.cpu_set = system_properties::read(DEX2OAT_CPU_SET_PROP_NAME).ok();
         self.start_instance(PENDING_INSTANCE_DIR, vm_parameters)
     }
 
diff --git a/compos/src/compsvc.rs b/compos/src/compsvc.rs
index 28bf5d9..0150f3c 100644
--- a/compos/src/compsvc.rs
+++ b/compos/src/compsvc.rs
@@ -21,6 +21,7 @@
 use anyhow::{Context, Result};
 use binder_common::new_binder_exception;
 use compos_common::binder::to_binder_result;
+use compos_common::DEX2OAT_THREADS_PROP_NAME;
 use log::warn;
 use std::default::Default;
 use std::path::PathBuf;
@@ -41,6 +42,7 @@
     BinderFeatures, ExceptionCode, Interface, Result as BinderResult, Strong,
 };
 use compos_common::odrefresh::ODREFRESH_PATH;
+use rustutils::system_properties;
 
 const AUTHFS_SERVICE_NAME: &str = "authfs_service";
 const DEX2OAT_PATH: &str = "/apex/com.android.art/bin/dex2oat64";
@@ -132,6 +134,10 @@
         fd_annotation: &FdAnnotation,
     ) -> BinderResult<CompilationResult> {
         let authfs_service = get_authfs_service()?;
+        to_binder_result(
+            system_properties::write(DEX2OAT_THREADS_PROP_NAME, &num_cpus::get().to_string())
+                .context(format!("Can't write {}", DEX2OAT_THREADS_PROP_NAME)),
+        )?;
         let output = to_binder_result(
             compile_cmd(&self.dex2oat_path, args, authfs_service, fd_annotation)
                 .context("Compilation failed"),