1 files changed, 23 insertions, 15 deletions

diff --git a/neuralnetworks/utils/common/src/ProtectCallback.cpp b/neuralnetworks/utils/common/src/ProtectCallback.cpp
index abe4cb675e..18e1f3bf0b 100644
--- a/neuralnetworks/utils/common/src/ProtectCallback.cpp
+++ b/neuralnetworks/utils/common/src/ProtectCallback.cpp
@@ -35,19 +35,25 @@ void DeathRecipient::serviceDied(uint64_t /*cookie*/, const wp<hidl::base::V1_0:
     std::lock_guard guard(mMutex);
     std::for_each(mObjects.begin(), mObjects.end(),
                   [](IProtectedCallback* killable) { killable->notifyAsDeadObject(); });
+    mObjects.clear();
+    mIsDeadObject = true;
 }
 
 void DeathRecipient::add(IProtectedCallback* killable) const {
     CHECK(killable != nullptr);
     std::lock_guard guard(mMutex);
-    mObjects.push_back(killable);
+    if (mIsDeadObject) {
+        killable->notifyAsDeadObject();
+    } else {
+        mObjects.push_back(killable);
+    }
 }
 
 void DeathRecipient::remove(IProtectedCallback* killable) const {
     CHECK(killable != nullptr);
     std::lock_guard guard(mMutex);
-    const auto removedIter = std::remove(mObjects.begin(), mObjects.end(), killable);
-    mObjects.erase(removedIter);
+    const auto newEnd = std::remove(mObjects.begin(), mObjects.end(), killable);
+    mObjects.erase(newEnd, mObjects.end());
 }
 
 nn::GeneralResult<DeathHandler> DeathHandler::create(sp<hidl::base::V1_0::IBase> object) {
@@ -67,19 +73,16 @@ nn::GeneralResult<DeathHandler> DeathHandler::create(sp<hidl::base::V1_0::IBase>
 }
 
 DeathHandler::DeathHandler(sp<hidl::base::V1_0::IBase> object, sp<DeathRecipient> deathRecipient)
-    : kObject(std::move(object)), kDeathRecipient(std::move(deathRecipient)) {
-    CHECK(kObject != nullptr);
-    CHECK(kDeathRecipient != nullptr);
+    : mObject(std::move(object)), mDeathRecipient(std::move(deathRecipient)) {
+    CHECK(mObject != nullptr);
+    CHECK(mDeathRecipient != nullptr);
 }
 
 DeathHandler::~DeathHandler() {
-    if (kObject != nullptr && kDeathRecipient != nullptr) {
-        const auto ret = kObject->unlinkToDeath(kDeathRecipient);
-        const auto maybeSuccess = handleTransportError(ret);
-        if (!maybeSuccess.has_value()) {
-            LOG(ERROR) << maybeSuccess.error().message;
-        } else if (!maybeSuccess.value()) {
-            LOG(ERROR) << "IBase::linkToDeath returned false";
+    if (mObject != nullptr && mDeathRecipient != nullptr) {
+        const auto successful = mObject->unlinkToDeath(mDeathRecipient).isOk();
+        if (!successful) {
+            LOG(ERROR) << "IBase::linkToDeath failed";
         }
     }
 }
@@ -87,9 +90,14 @@ DeathHandler::~DeathHandler() {
 [[nodiscard]] base::ScopeGuard<DeathHandler::Cleanup> DeathHandler::protectCallback(
         IProtectedCallback* killable) const {
     CHECK(killable != nullptr);
-    kDeathRecipient->add(killable);
+    mDeathRecipient->add(killable);
     return base::make_scope_guard(
-            [deathRecipient = kDeathRecipient, killable] { deathRecipient->remove(killable); });
+            [deathRecipient = mDeathRecipient, killable] { deathRecipient->remove(killable); });
+}
+
+void DeathHandler::protectCallbackForLifetimeOfDeathHandler(IProtectedCallback* killable) const {
+    CHECK(killable != nullptr);
+    mDeathRecipient->add(killable);
 }
 
 }  // namespace android::hardware::neuralnetworks::utils