Fix StatsCompanionService sometimes can be null

Bug: 75970648
Test: will add cts test for puller alarms
Change-Id: I51b7d13f855d3c8ded8325d7cf0f614531eceea5
diff --git a/cmds/statsd/src/StatsService.cpp b/cmds/statsd/src/StatsService.cpp
index b86646a..b03b4b4 100644
--- a/cmds/statsd/src/StatsService.cpp
+++ b/cmds/statsd/src/StatsService.cpp
@@ -49,33 +49,6 @@
 constexpr const char* kPermissionDump = "android.permission.DUMP";
 #define STATS_SERVICE_DIR "/data/misc/stats-service"
 
-/**
- * Watches for the death of the stats companion (system process).
- */
-class CompanionDeathRecipient : public IBinder::DeathRecipient {
-public:
-    CompanionDeathRecipient(const sp<AlarmMonitor>& anomalyAlarmMonitor,
-                            const sp<AlarmMonitor>& periodicAlarmMonitor,
-                            const sp<StatsLogProcessor>& processor)
-        : mAnomalyAlarmMonitor(anomalyAlarmMonitor),
-          mPeriodicAlarmMonitor(periodicAlarmMonitor),
-          mProcessor(processor) {}
-    virtual void binderDied(const wp<IBinder>& who);
-
-private:
-    sp<AlarmMonitor> mAnomalyAlarmMonitor;
-    sp<AlarmMonitor> mPeriodicAlarmMonitor;
-    sp<StatsLogProcessor> mProcessor;
-};
-
-void CompanionDeathRecipient::binderDied(const wp<IBinder>& who) {
-    ALOGW("statscompanion service died");
-    mProcessor->WriteDataToDisk();
-    mAnomalyAlarmMonitor->setStatsCompanionService(nullptr);
-    mPeriodicAlarmMonitor->setStatsCompanionService(nullptr);
-    SubscriberReporter::getInstance().setStatsCompanionService(nullptr);
-}
-
 StatsService::StatsService(const sp<Looper>& handlerLooper)
     : mAnomalyAlarmMonitor(new AlarmMonitor(MIN_DIFF_TO_UPDATE_REGISTERED_ALARM_SECS,
        [](const sp<IStatsCompanionService>& sc, int64_t timeMillis) {
@@ -791,21 +764,6 @@
     }
 }
 
-sp<IStatsCompanionService> StatsService::getStatsCompanionService() {
-    sp<IStatsCompanionService> statsCompanion = nullptr;
-    // Get statscompanion service from service manager
-    const sp<IServiceManager> sm(defaultServiceManager());
-    if (sm != nullptr) {
-        const String16 name("statscompanion");
-        statsCompanion = interface_cast<IStatsCompanionService>(sm->checkService(name));
-        if (statsCompanion == nullptr) {
-            ALOGW("statscompanion service unavailable!");
-            return nullptr;
-        }
-    }
-    return statsCompanion;
-}
-
 Status StatsService::statsCompanionReady() {
     VLOG("StatsService::statsCompanionReady was called");
 
@@ -821,9 +779,8 @@
                 "statscompanion unavailable despite it contacting statsd!");
     }
     VLOG("StatsService::statsCompanionReady linking to statsCompanion.");
-    IInterface::asBinder(statsCompanion)
-        ->linkToDeath(new CompanionDeathRecipient(
-            mAnomalyAlarmMonitor, mPeriodicAlarmMonitor, mProcessor));
+    IInterface::asBinder(statsCompanion)->linkToDeath(this);
+    mStatsPullerManager.SetStatsCompanionService(statsCompanion);
     mAnomalyAlarmMonitor->setStatsCompanionService(statsCompanion);
     mPeriodicAlarmMonitor->setStatsCompanionService(statsCompanion);
     SubscriberReporter::getInstance().setStatsCompanionService(statsCompanion);
@@ -969,6 +926,12 @@
 
 
 void StatsService::binderDied(const wp <IBinder>& who) {
+    ALOGW("statscompanion service died");
+    mProcessor->WriteDataToDisk();
+    mAnomalyAlarmMonitor->setStatsCompanionService(nullptr);
+    mPeriodicAlarmMonitor->setStatsCompanionService(nullptr);
+    SubscriberReporter::getInstance().setStatsCompanionService(nullptr);
+    mStatsPullerManager.SetStatsCompanionService(nullptr);
 }
 
 }  // namespace statsd