Add more ergonomic API for persisting lazy services.

This is being moved from binder_common.

Bug: 234019127
Test: atest compos_key_tests MicrodroidHostTestCases MicrodroidTestApp libbinder_rs-internal_test
Change-Id: I69c4fb78ff0d9886905d212253f1dfaf1ea51d57

3 files changed, 62 insertions, 2 deletions

diff --git a/libs/binder/rust/Android.bp b/libs/binder/rust/Android.bp
index 355b3b4b89..6c819e1d2b 100644
--- a/libs/binder/rust/Android.bp
+++ b/libs/binder/rust/Android.bp
@@ -15,9 +15,10 @@ rust_library {
         "libutils",
     ],
     rustlibs: [
-        "liblibc",
         "libbinder_ndk_sys",
         "libdowncast_rs",
+        "liblazy_static",
+        "liblibc",
     ],
     host_supported: true,
     vendor_available: true,
@@ -170,9 +171,10 @@ rust_test {
         "libbinder_ndk",
     ],
     rustlibs: [
-        "liblibc",
         "libbinder_ndk_sys",
         "libdowncast_rs",
+        "liblazy_static",
+        "liblibc",
     ],
 }
 
diff --git a/libs/binder/rust/src/lib.rs b/libs/binder/rust/src/lib.rs
index 67872a9f00..195d9ac160 100644
--- a/libs/binder/rust/src/lib.rs
+++ b/libs/binder/rust/src/lib.rs
@@ -111,6 +111,7 @@ pub use binder::{BinderFeatures, FromIBinder, IBinder, Interface, Strong, Weak};
 pub use error::{ExceptionCode, Status, StatusCode};
 pub use native::{
     add_service, force_lazy_services_persist, is_handling_transaction, register_lazy_service,
+    LazyServiceGuard,
 };
 pub use parcel::{ParcelFileDescriptor, Parcelable, ParcelableHolder};
 pub use proxy::{
diff --git a/libs/binder/rust/src/native.rs b/libs/binder/rust/src/native.rs
index 9e2cef1348..3a6daddffd 100644
--- a/libs/binder/rust/src/native.rs
+++ b/libs/binder/rust/src/native.rs
@@ -22,6 +22,7 @@ use crate::parcel::{BorrowedParcel, Serialize};
 use crate::proxy::SpIBinder;
 use crate::sys;
 
+use lazy_static::lazy_static;
 use std::convert::TryFrom;
 use std::ffi::{c_void, CStr, CString};
 use std::fs::File;
@@ -30,6 +31,7 @@ use std::ops::Deref;
 use std::os::raw::c_char;
 use std::os::unix::io::FromRawFd;
 use std::slice;
+use std::sync::Mutex;
 
 /// Rust wrapper around Binder remotable objects.
 ///
@@ -487,6 +489,8 @@ pub fn register_lazy_service(identifier: &str, mut binder: SpIBinder) -> Result<
 /// If persist is true then shut down will be blocked until this function is called again with
 /// persist false. If this is to be the initial state, call this function before calling
 /// register_lazy_service.
+///
+/// Consider using [`LazyServiceGuard`] rather than calling this directly.
 pub fn force_lazy_services_persist(persist: bool) {
     unsafe {
         // Safety: No borrowing or transfer of ownership occurs here.
@@ -494,6 +498,59 @@ pub fn force_lazy_services_persist(persist: bool) {
     }
 }
 
+/// An RAII object to ensure a process which registers lazy services is not killed. During the
+/// lifetime of any of these objects the service manager will not not kill the 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()
+    }
+}
+
 /// Tests often create a base BBinder instance; so allowing the unit
 /// type to be remotable translates nicely to Binder::new(()).
 impl Remotable for () {