Fix lazy service persistence for VS

As long as there is a VmInstance in existence, i.e. a client is
holding an IVirtualMachine reference, keep the virtualization service
running. Previously it would be killed by service manager when its
lazy IVirtualisationService interface was released, but that is
unhelpful to clients.

This reverts the changes in commit
7e54e2904f414369054114a0d5b9ed0e6d3a23cd that special-cased
debug_hold_vm in favor of the more general solution.

It also removes the workaround in compos_client to keep the service
from being killed.

Test: atest -p
Fixes: 200924402
Bug: 201628773
Change-Id: I8ba465b5970746b64205e08745d3ec4c931abae8
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)
     }
 }