Add active services count callback to lazy services.

Additionally, expose methods to tryUnregister/reRegister
services.

Bug: 176239128
Test: test aidl_lazy_test

Change-Id: I9aa3c9b681bd340ca340fe7ed818ba3533678af2
Merged-In: I9aa3c9b681bd340ca340fe7ed818ba3533678af2
(cherry picked from commit 2055e9b48d273e0e8997ce4745a00ac6a24b09c6)

2 files changed, 156 insertions, 43 deletions

diff --git a/libs/binder/LazyServiceRegistrar.cpp b/libs/binder/LazyServiceRegistrar.cpp
index f2c5139b56..bfd6942c3e 100644
--- a/libs/binder/LazyServiceRegistrar.cpp
+++ b/libs/binder/LazyServiceRegistrar.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include "log/log_main.h"
 #define LOG_TAG "AidlLazyServiceRegistrar"
 
 #include <binder/LazyServiceRegistrar.h>
@@ -37,6 +38,13 @@ public:
                          bool allowIsolated, int dumpFlags);
     void forcePersist(bool persist);
 
+    void setActiveServicesCountCallback(const std::function<bool(int)>&
+                                        activeServicesCountCallback);
+
+    bool tryUnregister();
+
+    void reRegister();
+
 protected:
     Status onClients(const sp<IBinder>& service, bool clients) override;
 
@@ -48,6 +56,14 @@ private:
     void tryShutdown();
 
     /**
+     * Try to shutdown the process, unless:
+     * - 'forcePersist' is 'true', or
+     * - The active services count callback returns 'true', or
+     * - Some services have clients.
+     */
+    void maybeTryShutdown();
+
+    /*
      * Counter of the number of services that currently have at least one client.
      */
     size_t mNumConnectedServices;
@@ -56,6 +72,8 @@ private:
         sp<IBinder> service;
         bool allowIsolated;
         int dumpFlags;
+
+        bool registered = true;
     };
     /**
      * Map of registered names and services
@@ -63,6 +81,9 @@ private:
     std::map<std::string, Service> mRegisteredServices;
 
     bool mForcePersist;
+
+    // Callback used to report the number of services with clients
+    std::function<bool(int)> mActiveServicesCountCallback;
 };
 
 class ClientCounterCallback {
@@ -77,6 +98,13 @@ public:
      */
     void forcePersist(bool persist);
 
+    void setActiveServicesCountCallback(const std::function<bool(int)>&
+                                        activeServicesCountCallback);
+
+    bool tryUnregister();
+
+    void reRegister();
+
 private:
     sp<ClientCounterCallbackImpl> mImpl;
 };
@@ -109,8 +137,60 @@ bool ClientCounterCallbackImpl::registerService(const sp<IBinder>& service, cons
 
 void ClientCounterCallbackImpl::forcePersist(bool persist) {
     mForcePersist = persist;
-    if(!mForcePersist) {
+    if (!mForcePersist && mNumConnectedServices == 0) {
         // Attempt a shutdown in case the number of clients hit 0 while the flag was on
+        maybeTryShutdown();
+    }
+}
+
+bool ClientCounterCallbackImpl::tryUnregister() {
+    auto manager = interface_cast<AidlServiceManager>(asBinder(defaultServiceManager()));
+
+    for (auto& [name, entry] : mRegisteredServices) {
+        bool success = manager->tryUnregisterService(name, entry.service).isOk();
+
+        if (!success) {
+            ALOGI("Failed to unregister service %s", name.c_str());
+            return false;
+        }
+        entry.registered = false;
+    }
+
+    return true;
+}
+
+void ClientCounterCallbackImpl::reRegister() {
+    for (auto& [name, entry] : mRegisteredServices) {
+        // re-register entry if not already registered
+        if (entry.registered) {
+            continue;
+        }
+
+        if (!registerService(entry.service, name, entry.allowIsolated,
+                             entry.dumpFlags)) {
+            // Must restart. Otherwise, clients will never be able to get a hold of this service.
+            LOG_ALWAYS_FATAL("Bad state: could not re-register services");
+        }
+
+        entry.registered = true;
+    }
+}
+
+void ClientCounterCallbackImpl::maybeTryShutdown() {
+    if (mForcePersist) {
+        ALOGI("Shutdown prevented by forcePersist override flag.");
+        return;
+    }
+
+    bool handledInCallback = false;
+    if (mActiveServicesCountCallback != nullptr) {
+        handledInCallback = mActiveServicesCountCallback(mNumConnectedServices);
+    }
+
+    // If there is no callback defined or the callback did not handle this
+    // client count change event, try to shutdown the process if its services
+    // have no clients.
+    if (!handledInCallback && mNumConnectedServices == 0) {
         tryShutdown();
     }
 }
@@ -130,53 +210,24 @@ Status ClientCounterCallbackImpl::onClients(const sp<IBinder>& service, bool cli
           mNumConnectedServices, mRegisteredServices.size(),
           String8(service->getInterfaceDescriptor()).string(), clients);
 
-    tryShutdown();
+    maybeTryShutdown();
     return Status::ok();
 }
 
-void ClientCounterCallbackImpl::tryShutdown() {
-    if(mNumConnectedServices > 0) {
-        // Should only shut down if there are no clients
-        return;
-    }
-
-    if(mForcePersist) {
-        ALOGI("Shutdown prevented by forcePersist override flag.");
-        return;
-    }
-
-    ALOGI("Trying to shut down the service. No clients in use for any service in process.");
-
-    auto manager = interface_cast<AidlServiceManager>(asBinder(defaultServiceManager()));
+ void ClientCounterCallbackImpl::tryShutdown() {
+     ALOGI("Trying to shut down the service. No clients in use for any service in process.");
 
-    auto unRegisterIt = mRegisteredServices.begin();
-    for (; unRegisterIt != mRegisteredServices.end(); ++unRegisterIt) {
-        auto& entry = (*unRegisterIt);
+    if (tryUnregister()) {
+         ALOGI("Unregistered all clients and exiting");
+         exit(EXIT_SUCCESS);
+     }
 
-        bool success = manager->tryUnregisterService(entry.first, entry.second.service).isOk();
-
-        if (!success) {
-            ALOGI("Failed to unregister service %s", entry.first.c_str());
-            break;
-        }
-    }
-
-    if (unRegisterIt == mRegisteredServices.end()) {
-        ALOGI("Unregistered all clients and exiting");
-        exit(EXIT_SUCCESS);
-    }
-
-    for (auto reRegisterIt = mRegisteredServices.begin(); reRegisterIt != unRegisterIt;
-         reRegisterIt++) {
-        auto& entry = (*reRegisterIt);
+    reRegister();
+}
 
-        // re-register entry
-        if (!registerService(entry.second.service, entry.first, entry.second.allowIsolated,
-                             entry.second.dumpFlags)) {
-            // Must restart. Otherwise, clients will never be able to get a hold of this service.
-            ALOGE("Bad state: could not re-register services");
-        }
-    }
+void ClientCounterCallbackImpl::setActiveServicesCountCallback(const std::function<bool(int)>&
+                                                               activeServicesCountCallback) {
+    mActiveServicesCountCallback = activeServicesCountCallback;
 }
 
 ClientCounterCallback::ClientCounterCallback() {
@@ -192,6 +243,19 @@ void ClientCounterCallback::forcePersist(bool persist) {
     mImpl->forcePersist(persist);
 }
 
+void ClientCounterCallback::setActiveServicesCountCallback(const std::function<bool(int)>&
+                                                           activeServicesCountCallback) {
+    mImpl->setActiveServicesCountCallback(activeServicesCountCallback);
+}
+
+bool ClientCounterCallback::tryUnregister() {
+    return mImpl->tryUnregister();
+}
+
+void ClientCounterCallback::reRegister() {
+    mImpl->reRegister();
+}
+
 }  // namespace internal
 
 LazyServiceRegistrar::LazyServiceRegistrar() {
@@ -215,5 +279,18 @@ void LazyServiceRegistrar::forcePersist(bool persist) {
     mClientCC->forcePersist(persist);
 }
 
+void LazyServiceRegistrar::setActiveServicesCountCallback(const std::function<bool(int)>&
+                                                          activeServicesCountCallback) {
+    mClientCC->setActiveServicesCountCallback(activeServicesCountCallback);
+}
+
+bool LazyServiceRegistrar::tryUnregister() {
+    return mClientCC->tryUnregister();
+}
+
+void LazyServiceRegistrar::reRegister() {
+    mClientCC->reRegister();
+}
+
 }  // namespace hardware
-}  // namespace android
\ No newline at end of file
+}  // namespace android
diff --git a/libs/binder/include/binder/LazyServiceRegistrar.h b/libs/binder/include/binder/LazyServiceRegistrar.h
index d18c88ec0e..73893c7dd0 100644
--- a/libs/binder/include/binder/LazyServiceRegistrar.h
+++ b/libs/binder/include/binder/LazyServiceRegistrar.h
@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include <functional>
+
 #include <binder/IServiceManager.h>
 #include <binder/Status.h>
 #include <utils/StrongPointer.h>
@@ -53,6 +55,40 @@ class LazyServiceRegistrar {
       */
      void forcePersist(bool persist);
 
+     /**
+      * Set a callback that is executed when the total number of services with
+      * clients changes.
+      * The callback takes an argument, which is the number of registered
+      * lazy services for this process which have clients.
+      *
+      * Callback return value:
+      * - false: Default behavior for lazy services (shut down the process if there
+      *          are no clients).
+      * - true:  Don't shut down the process even if there are no clients.
+      *
+      * This callback gives a chance to:
+      * 1 - Perform some additional operations before exiting;
+      * 2 - Prevent the process from exiting by returning "true" from the
+      *     callback.
+      *
+      * This method should be called before 'registerService' to avoid races.
+      */
+     void setActiveServicesCountCallback(const std::function<bool(int)>&
+                                         activeServicesCountCallback);
+
+    /**
+      * Try to unregister all services previously registered with 'registerService'.
+      * Returns 'true' if successful.
+      */
+     bool tryUnregister();
+
+     /**
+      * Re-register services that were unregistered by 'tryUnregister'.
+      * This method should be called in the case 'tryUnregister' fails
+      * (and should be called on the same thread).
+      */
+     void reRegister();
+
    private:
      std::shared_ptr<internal::ClientCounterCallback> mClientCC;
      LazyServiceRegistrar();