Add active services count callback to lazy HALs.

Additionally, add methods to tryUnregister/reRegister services.

Bug: 176240491
Test: atest hidl_lazy_test

Change-Id: I9e67315e1308d0ce7370bf7d293d74f9c081a3c4

2 files changed, 123 insertions, 17 deletions

diff --git a/transport/HidlLazyUtils.cpp b/transport/HidlLazyUtils.cpp
index d00b461..ce37dd0 100644
--- a/transport/HidlLazyUtils.cpp
+++ b/transport/HidlLazyUtils.cpp
@@ -36,6 +36,13 @@ class ClientCounterCallback : public ::android::hidl::manager::V1_2::IClientCall
 
     bool addRegisteredService(const sp<IBase>& service, const std::string& name);
 
+    bool tryUnregister();
+
+    void reRegister();
+
+    void setActiveServicesCountCallback(
+            const std::function<bool(int)>& activeServicesCountCallback);
+
   protected:
     Return<void> onClients(const sp<IBase>& service, bool clients) override;
 
@@ -44,6 +51,8 @@ class ClientCounterCallback : public ::android::hidl::manager::V1_2::IClientCall
         sp<IBase> service;
         std::string name;
         bool clients = false;
+        // Used to keep track of unregistered services to allow re-registry
+        bool registered = true;
     };
 
     /**
@@ -67,6 +76,11 @@ class ClientCounterCallback : public ::android::hidl::manager::V1_2::IClientCall
      * Number of services that have been registered.
      */
     std::vector<Service> mRegisteredServices;
+
+    /**
+     * Callback for reporting the number of services with clients.
+     */
+    std::function<bool(int)> mActiveServicesCountCallback;
 };
 
 class LazyServiceRegistrarImpl {
@@ -75,6 +89,10 @@ class LazyServiceRegistrarImpl {
 
     status_t registerService(const sp<::android::hidl::base::V1_0::IBase>& service,
                              const std::string& name);
+    bool tryUnregister();
+    void reRegister();
+    void setActiveServicesCountCallback(
+            const std::function<bool(int)>& activeServicesCountCallback);
 
   private:
     sp<ClientCounterCallback> mClientCallback;
@@ -147,48 +165,75 @@ Return<void> ClientCounterCallback::onClients(const sp<::android::hidl::base::V1
               << " available) client(s) in use after notification " << getDescriptor(service.get())
               << "/" << registered.name << " has clients: " << clients;
 
-    if (numWithClients == 0) {
+    bool handledInCallback = false;
+    if (mActiveServicesCountCallback != nullptr) {
+        handledInCallback = mActiveServicesCountCallback(numWithClients);
+    }
+
+    // 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 && numWithClients == 0) {
         tryShutdown();
     }
 
     return Status::ok();
 }
 
-void ClientCounterCallback::tryShutdown() {
-    LOG(INFO) << "Trying to exit HAL. No clients in use for any service in process.";
-
+bool ClientCounterCallback::tryUnregister() {
     auto manager = hardware::defaultServiceManager1_2();
 
-    auto unRegisterIt = mRegisteredServices.begin();
-    for (; unRegisterIt != mRegisteredServices.end(); ++unRegisterIt) {
-        auto& entry = (*unRegisterIt);
-
+    for (Service& entry : mRegisteredServices) {
         const std::string descriptor = getDescriptor(entry.service.get());
         bool success = manager->tryUnregister(descriptor, entry.name, entry.service);
 
         if (!success) {
             LOG(INFO) << "Failed to unregister HAL " << descriptor << "/" << entry.name;
-            break;
+            return false;
         }
-    }
 
-    if (unRegisterIt == mRegisteredServices.end()) {
-        LOG(INFO) << "Unregistered all clients and exiting";
-        exit(EXIT_SUCCESS);
+        // Mark the entry unregistered, but do not remove it (may still be re-registered)
+        entry.registered = false;
     }
 
-    for (auto reRegisterIt = mRegisteredServices.begin(); reRegisterIt != unRegisterIt;
-         reRegisterIt++) {
-        auto& entry = (*reRegisterIt);
+    return true;
+}
+
+void ClientCounterCallback::reRegister() {
+    for (Service& entry : mRegisteredServices) {
+        // re-register entry if not already registered
+        if (entry.registered) {
+            continue;
+        }
 
-        // re-register entry
         if (!registerService(entry.service, entry.name)) {
             // Must restart. Otherwise, clients will never be able to get ahold of this service.
             LOG(FATAL) << "Bad state: could not re-register " << getDescriptor(entry.service.get());
         }
+
+        entry.registered = true;
     }
 }
 
+void ClientCounterCallback::tryShutdown() {
+    LOG(INFO) << "Trying to exit HAL. No clients in use for any service in process.";
+
+    if (tryUnregister()) {
+        LOG(INFO) << "Unregistered all clients and exiting";
+        exit(EXIT_SUCCESS);
+    }
+
+    // At this point, we failed to unregister some of the services, leaving the
+    // server in an inconsistent state. Re-register all services that were
+    // unregistered by tryUnregister().
+    reRegister();
+}
+
+void ClientCounterCallback::setActiveServicesCountCallback(
+        const std::function<bool(int)>& activeServicesCountCallback) {
+    mActiveServicesCountCallback = activeServicesCountCallback;
+}
+
 status_t LazyServiceRegistrarImpl::registerService(
     const sp<::android::hidl::base::V1_0::IBase>& service, const std::string& name) {
     if (!mClientCallback->addRegisteredService(service, name)) {
@@ -198,6 +243,19 @@ status_t LazyServiceRegistrarImpl::registerService(
     return ::android::OK;
 }
 
+bool LazyServiceRegistrarImpl::tryUnregister() {
+    return mClientCallback->tryUnregister();
+}
+
+void LazyServiceRegistrarImpl::reRegister() {
+    mClientCallback->reRegister();
+}
+
+void LazyServiceRegistrarImpl::setActiveServicesCountCallback(
+        const std::function<bool(int)>& activeServicesCountCallback) {
+    mClientCallback->setActiveServicesCountCallback(activeServicesCountCallback);
+}
+
 }  // namespace details
 
 LazyServiceRegistrar::LazyServiceRegistrar() {
@@ -214,5 +272,18 @@ status_t LazyServiceRegistrar::registerService(
     return mImpl->registerService(service, name);
 }
 
+bool LazyServiceRegistrar::tryUnregister() {
+    return mImpl->tryUnregister();
+}
+
+void LazyServiceRegistrar::reRegister() {
+    mImpl->reRegister();
+}
+
+void LazyServiceRegistrar::setActiveServicesCountCallback(
+        const std::function<bool(int)>& activeServicesCountCallback) {
+    mImpl->setActiveServicesCountCallback(activeServicesCountCallback);
+}
+
 }  // namespace hardware
 }  // namespace android
diff --git a/transport/include/hidl/HidlLazyUtils.h b/transport/include/hidl/HidlLazyUtils.h
index 97fe20e..44fbcb2 100644
--- a/transport/include/hidl/HidlLazyUtils.h
+++ b/transport/include/hidl/HidlLazyUtils.h
@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include <functional>
+
 #include <android/hidl/base/1.0/IBase.h>
 #include <utils/RefBase.h>
 #include <utils/StrongPointer.h>
@@ -41,6 +43,39 @@ class LazyServiceRegistrar {
      static LazyServiceRegistrar& getInstance();
      status_t registerService(const sp<::android::hidl::base::V1_0::IBase>& service,
                               const std::string& name = "default");
+     /**
+      * 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 HALs for this process which have clients.
+      *
+      * Callback return value:
+      * - false: Default behavior for lazy HALs (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<details::LazyServiceRegistrarImpl> mImpl;