Merge changes from topics "b/200924402", "composd_race" am: 16a2b5b7ce

Original change: https://android-review.googlesource.com/c/platform/packages/modules/Virtualization/+/1836356

Change-Id: Ibd44735b3dae8bdc7119c256162c02a5a2f8f99c
diff --git a/binder_common/Android.bp b/binder_common/Android.bp
index 789a891..9e6d590 100644
--- a/binder_common/Android.bp
+++ b/binder_common/Android.bp
@@ -9,6 +9,7 @@
     edition: "2018",
     rustlibs: [
         "libbinder_rs",
+        "liblazy_static",
     ],
     apex_available: [
         "com.android.compos",
diff --git a/binder_common/lazy_service.rs b/binder_common/lazy_service.rs
new file mode 100644
index 0000000..a2b85db
--- /dev/null
+++ b/binder_common/lazy_service.rs
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//! Rust API for lazy (aka dynamic) AIDL services.
+//! See https://source.android.com/devices/architecture/aidl/dynamic-aidl.
+
+use binder::public_api::force_lazy_services_persist;
+use lazy_static::lazy_static;
+use std::sync::Mutex;
+
+// TODO(b/200924402): Move this class to libbinder_rs once the infrastructure needed exists.
+
+/// An RAII object to ensure a server of lazy services is not killed. During the lifetime of any of
+/// these objects the service manager will not not kill the current process even if none of its
+/// lazy services are in use.
+#[must_use]
+#[derive(Debug)]
+pub struct LazyServiceGuard {
+    // Prevent construction outside this module.
+    _private: (),
+}
+
+lazy_static! {
+    // Count of how many LazyServiceGuard objects are in existence.
+    static ref GUARD_COUNT: Mutex<u64> = Mutex::new(0);
+}
+
+impl LazyServiceGuard {
+    /// Create a new LazyServiceGuard to prevent the service manager prematurely killing this
+    /// process.
+    pub fn new() -> Self {
+        let mut count = GUARD_COUNT.lock().unwrap();
+        *count += 1;
+        if *count == 1 {
+            // It's important that we make this call with the mutex held, to make sure
+            // that multiple calls (e.g. if the count goes 1 -> 0 -> 1) are correctly
+            // sequenced. (That also means we can't just use an AtomicU64.)
+            force_lazy_services_persist(true);
+        }
+        Self { _private: () }
+    }
+}
+
+impl Drop for LazyServiceGuard {
+    fn drop(&mut self) {
+        let mut count = GUARD_COUNT.lock().unwrap();
+        *count -= 1;
+        if *count == 0 {
+            force_lazy_services_persist(false);
+        }
+    }
+}
+
+impl Clone for LazyServiceGuard {
+    fn clone(&self) -> Self {
+        Self::new()
+    }
+}
+
+impl Default for LazyServiceGuard {
+    fn default() -> Self {
+        Self::new()
+    }
+}
diff --git a/binder_common/lib.rs b/binder_common/lib.rs
index 54cb80e..055688a 100644
--- a/binder_common/lib.rs
+++ b/binder_common/lib.rs
@@ -16,6 +16,8 @@
 
 //! Common items useful for binder clients and/or servers.
 
+pub mod lazy_service;
+
 use binder::public_api::{ExceptionCode, Status};
 use std::ffi::CString;
 
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index e68deb8..79d8354 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -46,11 +46,8 @@
 
 /// This owns an instance of the CompOS VM.
 pub struct VmInstance {
-    #[allow(dead_code)] // Prevent service manager from killing the dynamic service
-    service: Strong<dyn IVirtualizationService>,
     #[allow(dead_code)] // Keeps the VM alive even if we don`t touch it
     vm: Strong<dyn IVirtualMachine>,
-    #[allow(dead_code)] // TODO: Do we need this?
     cid: i32,
 }
 
@@ -113,7 +110,7 @@
 
         let cid = vm_state.wait_until_ready()?;
 
-        Ok(VmInstance { service, vm, cid })
+        Ok(VmInstance { vm, cid })
     }
 
     /// Create and return an RPC Binder connection to the Comp OS service in the VM.
@@ -129,6 +126,7 @@
 
     /// Return the CID of the VM.
     pub fn cid(&self) -> i32 {
+        // TODO: Do we actually need/use this?
         self.cid
     }
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 76c3a16..a072060 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -19,13 +19,11 @@
 use crate::payload::add_microdroid_images;
 use crate::{Cid, FIRST_GUEST_CID, SYSPROP_LAST_CID};
 
-use binder_common::new_binder_exception;
+use ::binder::unstable_api::AsNative;
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualMachine::{
-    BnVirtualMachine, IVirtualMachine,
-};
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
     DiskImage::DiskImage,
+    IVirtualMachine::{BnVirtualMachine, IVirtualMachine},
     IVirtualMachineCallback::IVirtualMachineCallback,
     IVirtualizationService::IVirtualizationService,
     PartitionType::PartitionType,
@@ -36,20 +34,24 @@
     VirtualMachineState::VirtualMachineState,
 };
 use android_system_virtualizationservice::binder::{
-    self, force_lazy_services_persist, BinderFeatures, ExceptionCode, Interface, ParcelFileDescriptor, Status, Strong, ThreadState,
+    self, BinderFeatures, ExceptionCode, Interface, ParcelFileDescriptor, Status, Strong,
+    ThreadState,
 };
-use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
-    VM_BINDER_SERVICE_PORT, VM_STREAM_SERVICE_PORT, BnVirtualMachineService, IVirtualMachineService,
+use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
+    IVirtualMachineService::{
+        BnVirtualMachineService, IVirtualMachineService, VM_BINDER_SERVICE_PORT,
+        VM_STREAM_SERVICE_PORT,
+    },
 };
 use anyhow::{anyhow, bail, Context, Result};
-use ::binder::unstable_api::AsNative;
+use binder_common::{lazy_service::LazyServiceGuard, new_binder_exception};
 use disk::QcowFile;
-use idsig::{V4Signature, HashAlgorithm};
-use log::{debug, error, warn, info};
+use idsig::{HashAlgorithm, V4Signature};
+use log::{debug, error, info, warn};
 use microdroid_payload_config::VmPayloadConfig;
 use rustutils::system_properties;
 use std::convert::TryInto;
-use std::fs::{File, OpenOptions, create_dir};
+use std::fs::{create_dir, File, OpenOptions};
 use std::io::{Error, ErrorKind, Write};
 use std::num::NonZeroU32;
 use std::os::unix::io::{FromRawFd, IntoRawFd};
@@ -557,11 +559,13 @@
 #[derive(Debug)]
 struct VirtualMachine {
     instance: Arc<VmInstance>,
+    /// Keeps our service process running as long as this VM instance exists.
+    lazy_service_guard: LazyServiceGuard,
 }
 
 impl VirtualMachine {
     fn create(instance: Arc<VmInstance>) -> Strong<dyn IVirtualMachine> {
-        let binder = VirtualMachine { instance };
+        let binder = VirtualMachine { instance, lazy_service_guard: Default::default() };
         BnVirtualMachine::new_binder(binder, BinderFeatures::default())
     }
 }
@@ -714,19 +718,12 @@
     /// Store a strong VM reference.
     fn debug_hold_vm(&mut self, vm: Strong<dyn IVirtualMachine>) {
         self.debug_held_vms.push(vm);
-        // Make sure our process is not shut down while we hold the VM reference
-        // on behalf of the caller.
-        force_lazy_services_persist(true);
     }
 
     /// Retrieve and remove a strong VM reference.
     fn debug_drop_vm(&mut self, cid: i32) -> Option<Strong<dyn IVirtualMachine>> {
         let pos = self.debug_held_vms.iter().position(|vm| vm.getCid() == Ok(cid))?;
         let vm = self.debug_held_vms.swap_remove(pos);
-        if self.debug_held_vms.is_empty() {
-            // Once we no longer hold any VM references it is ok for our process to be shut down.
-            force_lazy_services_persist(false);
-        }
         Some(vm)
     }
 }