3 files changed, 33 insertions, 23 deletions

diff --git a/cmds/statsd/src/config/ConfigManager.cpp b/cmds/statsd/src/config/ConfigManager.cpp
index 6d9c644bb40e..bbae3fef7934 100644
--- a/cmds/statsd/src/config/ConfigManager.cpp
+++ b/cmds/statsd/src/config/ConfigManager.cpp
@@ -43,20 +43,23 @@ using android::base::StringPrintf;
 using std::unique_ptr;
 
 struct ConfigReceiverDeathCookie {
-    ConfigReceiverDeathCookie(sp<ConfigManager> configManager, const ConfigKey& configKey,
-                              const shared_ptr<IPendingIntentRef>& pir):
-        mConfigManager(configManager),
-        mConfigKey(configKey),
-        mPir(pir) {}
+    ConfigReceiverDeathCookie(const wp<ConfigManager>& configManager, const ConfigKey& configKey,
+                              const shared_ptr<IPendingIntentRef>& pir) :
+            mConfigManager(configManager), mConfigKey(configKey), mPir(pir) {
+    }
 
-    sp<ConfigManager> mConfigManager;
+    wp<ConfigManager> mConfigManager;
     ConfigKey mConfigKey;
     shared_ptr<IPendingIntentRef> mPir;
 };
 
 void ConfigManager::configReceiverDied(void* cookie) {
     auto cookie_ = static_cast<ConfigReceiverDeathCookie*>(cookie);
-    sp<ConfigManager>& thiz = cookie_->mConfigManager;
+    sp<ConfigManager> thiz = cookie_->mConfigManager.promote();
+    if (!thiz) {
+        return;
+    }
+
     ConfigKey& configKey = cookie_->mConfigKey;
     shared_ptr<IPendingIntentRef>& pir = cookie_->mPir;
 
@@ -74,20 +77,23 @@ void ConfigManager::configReceiverDied(void* cookie) {
 }
 
 struct ActiveConfigChangedReceiverDeathCookie {
-    ActiveConfigChangedReceiverDeathCookie(sp<ConfigManager> configManager, const int uid,
-                                           const shared_ptr<IPendingIntentRef>& pir):
-        mConfigManager(configManager),
-        mUid(uid),
-        mPir(pir) {}
+    ActiveConfigChangedReceiverDeathCookie(const wp<ConfigManager>& configManager, const int uid,
+                                           const shared_ptr<IPendingIntentRef>& pir) :
+            mConfigManager(configManager), mUid(uid), mPir(pir) {
+    }
 
-    sp<ConfigManager> mConfigManager;
+    wp<ConfigManager> mConfigManager;
     int mUid;
     shared_ptr<IPendingIntentRef> mPir;
 };
 
 void ConfigManager::activeConfigChangedReceiverDied(void* cookie) {
     auto cookie_ = static_cast<ActiveConfigChangedReceiverDeathCookie*>(cookie);
-    sp<ConfigManager>& thiz = cookie_->mConfigManager;
+    sp<ConfigManager> thiz = cookie_->mConfigManager.promote();
+    if (!thiz) {
+        return;
+    }
+
     int uid = cookie_->mUid;
     shared_ptr<IPendingIntentRef>& pir = cookie_->mPir;
 
diff --git a/cmds/statsd/src/external/StatsPullerManager.cpp b/cmds/statsd/src/external/StatsPullerManager.cpp
index 79a7e8d318e2..ebe961014336 100644
--- a/cmds/statsd/src/external/StatsPullerManager.cpp
+++ b/cmds/statsd/src/external/StatsPullerManager.cpp
@@ -44,19 +44,23 @@ namespace statsd {
 // Stores the puller as a wp to avoid holding a reference in case it is unregistered and
 // pullAtomCallbackDied is never called.
 struct PullAtomCallbackDeathCookie {
-    PullAtomCallbackDeathCookie(sp<StatsPullerManager> pullerManager, const PullerKey& pullerKey,
-                                const wp<StatsPuller>& puller)
-        : mPullerManager(pullerManager), mPullerKey(pullerKey), mPuller(puller) {
+    PullAtomCallbackDeathCookie(const wp<StatsPullerManager>& pullerManager,
+                                const PullerKey& pullerKey, const wp<StatsPuller>& puller) :
+            mPullerManager(pullerManager), mPullerKey(pullerKey), mPuller(puller) {
     }
 
-    sp<StatsPullerManager> mPullerManager;
+    wp<StatsPullerManager> mPullerManager;
     PullerKey mPullerKey;
     wp<StatsPuller> mPuller;
 };
 
 void StatsPullerManager::pullAtomCallbackDied(void* cookie) {
     PullAtomCallbackDeathCookie* cookie_ = static_cast<PullAtomCallbackDeathCookie*>(cookie);
-    sp<StatsPullerManager>& thiz = cookie_->mPullerManager;
+    sp<StatsPullerManager> thiz = cookie_->mPullerManager.promote();
+    if (!thiz) {
+        return;
+    }
+
     const PullerKey& pullerKey = cookie_->mPullerKey;
     wp<StatsPuller> puller = cookie_->mPuller;
 
diff --git a/cmds/statsd/tests/external/StatsCallbackPuller_test.cpp b/cmds/statsd/tests/external/StatsCallbackPuller_test.cpp
index 6aff9ef80a71..4b9bac127dc8 100644
--- a/cmds/statsd/tests/external/StatsCallbackPuller_test.cpp
+++ b/cmds/statsd/tests/external/StatsCallbackPuller_test.cpp
@@ -190,13 +190,13 @@ TEST_F(StatsCallbackPullerTest, RegisterAndTimeout) {
     int32_t uid = 123;
     values.push_back(value);
 
-    StatsPullerManager pullerManager;
-    pullerManager.RegisterPullAtomCallback(uid, pullTagId, pullCoolDownNs, pullTimeoutNs,
-                                           vector<int32_t>(), cb);
+    sp<StatsPullerManager> pullerManager = new StatsPullerManager();
+    pullerManager->RegisterPullAtomCallback(uid, pullTagId, pullCoolDownNs, pullTimeoutNs,
+                                            vector<int32_t>(), cb);
     vector<shared_ptr<LogEvent>> dataHolder;
     int64_t startTimeNs = getElapsedRealtimeNs();
     // Returns false, since StatsPuller code will evaluate the timeout.
-    EXPECT_FALSE(pullerManager.Pull(pullTagId, {uid}, &dataHolder));
+    EXPECT_FALSE(pullerManager->Pull(pullTagId, {uid}, &dataHolder));
     int64_t endTimeNs = getElapsedRealtimeNs();
     int64_t actualPullDurationNs = endTimeNs - startTimeNs;