Update setBroadcastSubscriber

change #setBroadcastSubscriber and #unsetBroadcastSubscriber
to avoid using intentsender

Bug: 146074295
Test: Ran GTS Tests
Change-Id: I1510e44bcdf49b579fd49f51081c6a40618039fa
diff --git a/apex/statsd/aidl/android/os/IStatsCompanionService.aidl b/apex/statsd/aidl/android/os/IStatsCompanionService.aidl
index 32413e2..21b7767 100644
--- a/apex/statsd/aidl/android/os/IStatsCompanionService.aidl
+++ b/apex/statsd/aidl/android/os/IStatsCompanionService.aidl
@@ -17,7 +17,6 @@
 package android.os;
 
 import android.os.IPullAtomCallback;
-import android.os.StatsDimensionsValue;
 import android.os.StatsLogEventWrapper;
 
 /**
@@ -66,15 +65,6 @@
     /** Pull the specified data. Results will be sent to statsd when complete. */
     StatsLogEventWrapper[] pullData(int pullCode);
 
-    /**
-     * Requests StatsCompanionService to send a broadcast using the given intentSender
-     * (which should cast to an IIntentSender), along with the other information specified.
-     */
-    oneway void sendSubscriberBroadcast(in IBinder intentSender, long configUid, long configId,
-                                        long subscriptionId, long subscriptionRuleId,
-                                        in String[] cookies,
-                                        in StatsDimensionsValue dimensionsValue);
-
     /** Tells StatsCompaionService to grab the uid map snapshot and send it to statsd. */
     oneway void triggerUidSnapshot();
 
diff --git a/apex/statsd/aidl/android/os/IStatsManagerService.aidl b/apex/statsd/aidl/android/os/IStatsManagerService.aidl
index 291140c..05583677 100644
--- a/apex/statsd/aidl/android/os/IStatsManagerService.aidl
+++ b/apex/statsd/aidl/android/os/IStatsManagerService.aidl
@@ -72,8 +72,17 @@
      * This function can only be called by the owner (uid) of the config. It must be called each
      * time statsd starts. Later calls overwrite previous calls; only one PendingIntent is stored.
      *
-     * Requires Manifest.permission.DUMP.
+     * Requires Manifest.permission.DUMP and Manifest.permission.PACKAGE_USAGE_STATS.
      */
     void setBroadcastSubscriber(long configKey, long subscriberId, in PendingIntent pendingIntent,
                                 in String packageName);
+
+    /**
+     * Undoes setBroadcastSubscriber() for the (configKey, subscriberId) pair.
+     * Any broadcasts associated with subscriberId will henceforth not be sent.
+     * No-op if this (configKey, subscriberId) pair was not associated with an PendingIntent.
+     *
+     * Requires Manifest.permission.DUMP and Manifest.permission.PACKAGE_USAGE_STATS.
+     */
+    void unsetBroadcastSubscriber(long configKey, long subscriberId, in String packageName);
 }
\ No newline at end of file
diff --git a/apex/statsd/aidl/android/os/IStatsd.aidl b/apex/statsd/aidl/android/os/IStatsd.aidl
index 323d56f..9358439 100644
--- a/apex/statsd/aidl/android/os/IStatsd.aidl
+++ b/apex/statsd/aidl/android/os/IStatsd.aidl
@@ -151,34 +151,31 @@
     void removeConfiguration(in long configKey, in String packageName);
 
     /**
-     * Set the IIntentSender (i.e. PendingIntent) to be used when broadcasting subscriber
+     * Set the PendingIntentRef to be used when broadcasting subscriber
      * information to the given subscriberId within the given config.
      *
-     * Suppose that the calling uid has added a config with key configKey, and that in this config
+     * Suppose that the calling uid has added a config with key configId, and that in this config
      * it is specified that when a particular anomaly is detected, a broadcast should be sent to
-     * a BroadcastSubscriber with id subscriberId. This function links the given intentSender with
-     * that subscriberId (for that config), so that this intentSender is used to send the broadcast
+     * a BroadcastSubscriber with id subscriberId. This function links the given pendingIntent with
+     * that subscriberId (for that config), so that this pendingIntent is used to send the broadcast
      * when the anomaly is detected.
      *
      * This function can only be called by the owner (uid) of the config. It must be called each
-     * time statsd starts. Later calls overwrite previous calls; only one intentSender is stored.
-     *
-     * intentSender must be convertible into an IntentSender using IntentSender(IBinder)
-     * and cannot be null.
+     * time statsd starts. Later calls overwrite previous calls; only one pendingIntent is stored.
      *
      * Requires Manifest.permission.DUMP.
      */
-    void setBroadcastSubscriber(long configKey, long subscriberId, in IBinder intentSender,
-                                in String packageName);
+    void setBroadcastSubscriber(long configId, long subscriberId, in IPendingIntentRef pir,
+                                int callingUid);
 
     /**
-     * Undoes setBroadcastSubscriber() for the (configKey, subscriberId) pair.
+     * Undoes setBroadcastSubscriber() for the (configId, subscriberId) pair.
      * Any broadcasts associated with subscriberId will henceforth not be sent.
-     * No-op if this (configKey, subsriberId) pair was not associated with an IntentSender.
+     * No-op if this (configKey, subscriberId) pair was not associated with an PendingIntentRef.
      *
      * Requires Manifest.permission.DUMP.
      */
-    void unsetBroadcastSubscriber(long configKey, long subscriberId, in String packageName);
+    void unsetBroadcastSubscriber(long configId, long subscriberId, int callingUid);
 
     /**
      * Apps can send an atom via this application breadcrumb with the specified label and state for
diff --git a/apex/statsd/service/java/com/android/server/stats/StatsCompanion.java b/apex/statsd/service/java/com/android/server/stats/StatsCompanion.java
index 77ccfb6..4383b50 100644
--- a/apex/statsd/service/java/com/android/server/stats/StatsCompanion.java
+++ b/apex/statsd/service/java/com/android/server/stats/StatsCompanion.java
@@ -28,6 +28,7 @@
 
 import com.android.server.SystemService;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 
 /**
@@ -101,7 +102,7 @@
         private static final String EXTRA_LAST_REPORT_TIME = "android.app.extra.LAST_REPORT_TIME";
         private static final int CODE_DATA_BROADCAST = 1;
         private static final int CODE_ACTIVE_CONFIGS_BROADCAST = 1;
-
+        private static final int CODE_SUBSCRIBER_BROADCAST = 1;
 
         private final PendingIntent mPendingIntent;
         private final Context mContext;
@@ -141,7 +142,36 @@
         @Override
         public void sendSubscriberBroadcast(long configUid, long configId, long subscriptionId,
                 long subscriptionRuleId, String[] cookies, StatsDimensionsValue dimensionsValue) {
-            // no-op
+            enforceStatsCompanionPermission(mContext);
+            Intent intent =
+                    new Intent()
+                            .putExtra(StatsManager.EXTRA_STATS_CONFIG_UID, configUid)
+                            .putExtra(StatsManager.EXTRA_STATS_CONFIG_KEY, configId)
+                            .putExtra(StatsManager.EXTRA_STATS_SUBSCRIPTION_ID, subscriptionId)
+                            .putExtra(StatsManager.EXTRA_STATS_SUBSCRIPTION_RULE_ID,
+                                    subscriptionRuleId)
+                            .putExtra(StatsManager.EXTRA_STATS_DIMENSIONS_VALUE, dimensionsValue);
+
+            ArrayList<String> cookieList = new ArrayList<>(cookies.length);
+            cookieList.addAll(Arrays.asList(cookies));
+            intent.putStringArrayListExtra(
+                    StatsManager.EXTRA_STATS_BROADCAST_SUBSCRIBER_COOKIES, cookieList);
+
+            if (DEBUG) {
+                Slog.d(TAG,
+                        String.format(
+                                "Statsd sendSubscriberBroadcast with params {%d %d %d %d %s %s}",
+                                configUid, configId, subscriptionId, subscriptionRuleId,
+                                Arrays.toString(cookies),
+                                dimensionsValue));
+            }
+            try {
+                mPendingIntent.send(mContext, CODE_SUBSCRIBER_BROADCAST, intent, null, null);
+            } catch (PendingIntent.CanceledException e) {
+                Slog.w(TAG,
+                        "Unable to send using PendingIntent from uid " + configUid
+                                + "; presumably it had been cancelled.");
+            }
         }
     }
 }
diff --git a/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java b/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
index befce21..d57afee 100644
--- a/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
+++ b/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
@@ -50,7 +50,6 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
-import android.content.IntentSender;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
@@ -88,7 +87,6 @@
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.StatFs;
-import android.os.StatsDimensionsValue;
 import android.os.StatsLogEventWrapper;
 import android.os.SynchronousResultReceiver;
 import android.os.SystemClock;
@@ -202,7 +200,6 @@
     private static final int PACKAGE_NAME_FIELD_ID = 4;
     private static final int INSTALLER_FIELD_ID = 5;
 
-    public static final int CODE_SUBSCRIBER_BROADCAST = 1;
     public static final int DEATH_THRESHOLD = 10;
     /**
      * Which native processes to snapshot memory for.
@@ -442,43 +439,6 @@
                 KernelCpuThreadReaderSettingsObserver.getSettingsModifiedReader(mContext);
     }
 
-    @Override
-    public void sendSubscriberBroadcast(IBinder intentSenderBinder, long configUid, long configKey,
-            long subscriptionId, long subscriptionRuleId, String[] cookies,
-            StatsDimensionsValue dimensionsValue) {
-        StatsCompanion.enforceStatsCompanionPermission(mContext);
-        IntentSender intentSender = new IntentSender(intentSenderBinder);
-        Intent intent =
-                new Intent()
-                        .putExtra(StatsManager.EXTRA_STATS_CONFIG_UID, configUid)
-                        .putExtra(StatsManager.EXTRA_STATS_CONFIG_KEY, configKey)
-                        .putExtra(StatsManager.EXTRA_STATS_SUBSCRIPTION_ID, subscriptionId)
-                        .putExtra(StatsManager.EXTRA_STATS_SUBSCRIPTION_RULE_ID, subscriptionRuleId)
-                        .putExtra(StatsManager.EXTRA_STATS_DIMENSIONS_VALUE, dimensionsValue);
-
-        ArrayList<String> cookieList = new ArrayList<>(cookies.length);
-        for (String cookie : cookies) {
-            cookieList.add(cookie);
-        }
-        intent.putStringArrayListExtra(
-                StatsManager.EXTRA_STATS_BROADCAST_SUBSCRIBER_COOKIES, cookieList);
-
-        if (DEBUG) {
-            Slog.d(TAG,
-                    String.format("Statsd sendSubscriberBroadcast with params {%d %d %d %d %s %s}",
-                            configUid, configKey, subscriptionId, subscriptionRuleId,
-                            Arrays.toString(cookies),
-                            dimensionsValue));
-        }
-        try {
-            intentSender.sendIntent(mContext, CODE_SUBSCRIBER_BROADCAST, intent, null, null);
-        } catch (IntentSender.SendIntentException e) {
-            Slog.w(TAG,
-                    "Unable to send using IntentSender from uid " + configUid
-                            + "; presumably it had been cancelled.");
-        }
-    }
-
     private final static int[] toIntArray(List<Integer> list) {
         int[] ret = new int[list.size()];
         for (int i = 0; i < ret.length; i++) {
diff --git a/apex/statsd/service/java/com/android/server/stats/StatsManagerService.java b/apex/statsd/service/java/com/android/server/stats/StatsManagerService.java
index ef2b761..1d03e3b 100644
--- a/apex/statsd/service/java/com/android/server/stats/StatsManagerService.java
+++ b/apex/statsd/service/java/com/android/server/stats/StatsManagerService.java
@@ -61,6 +61,9 @@
     @GuardedBy("mLock")
     private ArrayMap<Integer, PendingIntentRef> mActiveConfigsPirMap =
             new ArrayMap<>();
+    @GuardedBy("mLock")
+    private ArrayMap<ConfigKey, ArrayMap<Long, PendingIntentRef>> mBroadcastSubscriberPirMap =
+            new ArrayMap<>();
 
     public StatsManagerService(Context context) {
         super();
@@ -189,11 +192,56 @@
     }
 
     @Override
-    public void setBroadcastSubscriber(long configKey, long subscriberId,
+    public void setBroadcastSubscriber(long configId, long subscriberId,
             PendingIntent pendingIntent, String packageName) {
-        //no-op
-        if (DEBUG) {
-            Slog.d(TAG, "setBroadcastSubscriber");
+        enforceDumpAndUsageStatsPermission(packageName);
+        int callingUid = Binder.getCallingUid();
+        final long token = Binder.clearCallingIdentity();
+        PendingIntentRef pir = new PendingIntentRef(pendingIntent, mContext);
+        ConfigKey key = new ConfigKey(callingUid, configId);
+        // We add the PIR to a map so we can reregister if statsd is unavailable.
+        synchronized (mLock) {
+            ArrayMap<Long, PendingIntentRef> innerMap = mBroadcastSubscriberPirMap
+                    .getOrDefault(key, new ArrayMap<>());
+            innerMap.put(subscriberId, pir);
+            mBroadcastSubscriberPirMap.put(key, innerMap);
+        }
+        try {
+            IStatsd statsd = getStatsdNonblocking();
+            if (statsd != null) {
+                statsd.setBroadcastSubscriber(
+                        configId, subscriberId, pir, callingUid);
+            }
+        } catch (RemoteException e) {
+            Slog.e(TAG, "Failed to setBroadcastSubscriber with statsd");
+        } finally {
+            Binder.restoreCallingIdentity(token);
+        }
+    }
+
+    @Override
+    public void unsetBroadcastSubscriber(long configId, long subscriberId, String packageName) {
+        enforceDumpAndUsageStatsPermission(packageName);
+        int callingUid = Binder.getCallingUid();
+        final long token = Binder.clearCallingIdentity();
+        ConfigKey key = new ConfigKey(callingUid, configId);
+        synchronized (mLock) {
+            ArrayMap<Long, PendingIntentRef> innerMap = mBroadcastSubscriberPirMap
+                    .getOrDefault(key, new ArrayMap<>());
+            innerMap.remove(subscriberId);
+            if (innerMap.isEmpty()) {
+                mBroadcastSubscriberPirMap.remove(key);
+            }
+        }
+        try {
+            IStatsd statsd = getStatsdNonblocking();
+            if (statsd != null) {
+                statsd.unsetBroadcastSubscriber(configId, subscriberId, callingUid);
+            }
+        } catch (RemoteException e) {
+            Slog.e(TAG, "Failed to unsetBroadcastSubscriber with statsd");
+        } finally {
+            Binder.restoreCallingIdentity(token);
         }
     }
 
@@ -285,11 +333,19 @@
         if (statsd == null) {
             return;
         }
+        // Since we do not want to make an IPC with the a lock held, we first create local deep
+        // copies of the data with the lock held before iterating through the maps.
         ArrayMap<ConfigKey, PendingIntentRef> dataFetchCopy;
         ArrayMap<Integer, PendingIntentRef> activeConfigsChangedCopy;
+        ArrayMap<ConfigKey, ArrayMap<Long, PendingIntentRef>> broadcastSubscriberCopy;
         synchronized (mLock) {
             dataFetchCopy = new ArrayMap<>(mDataFetchPirMap);
             activeConfigsChangedCopy = new ArrayMap<>(mActiveConfigsPirMap);
+            broadcastSubscriberCopy = new ArrayMap<>();
+            for (Map.Entry<ConfigKey, ArrayMap<Long, PendingIntentRef>> entry
+                    : mBroadcastSubscriberPirMap.entrySet()) {
+                broadcastSubscriberCopy.put(entry.getKey(), new ArrayMap<>(entry.getValue()));
+            }
         }
         for (Map.Entry<ConfigKey, PendingIntentRef> entry : dataFetchCopy.entrySet()) {
             ConfigKey key = entry.getKey();
@@ -307,5 +363,17 @@
                 Slog.e(TAG, "Failed to setActiveConfigsChangedOperation from pirMap");
             }
         }
+        for (Map.Entry<ConfigKey, ArrayMap<Long, PendingIntentRef>> entry
+                : broadcastSubscriberCopy.entrySet()) {
+            for (Map.Entry<Long, PendingIntentRef> subscriberEntry : entry.getValue().entrySet()) {
+                ConfigKey configKey = entry.getKey();
+                try {
+                    statsd.setBroadcastSubscriber(configKey.getConfigId(), subscriberEntry.getKey(),
+                            subscriberEntry.getValue(), configKey.getUid());
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Failed to setBroadcastSubscriber from pirMap");
+                }
+            }
+        }
     }
 }
diff --git a/cmds/statsd/src/StatsService.cpp b/cmds/statsd/src/StatsService.cpp
index 3a29b72..5ff0f97 100644
--- a/cmds/statsd/src/StatsService.cpp
+++ b/cmds/statsd/src/StatsService.cpp
@@ -1111,7 +1111,6 @@
     mPullerManager->SetStatsCompanionService(statsCompanion);
     mAnomalyAlarmMonitor->setStatsCompanionService(statsCompanion);
     mPeriodicAlarmMonitor->setStatsCompanionService(statsCompanion);
-    SubscriberReporter::getInstance().setStatsCompanionService(statsCompanion);
     return Status::ok();
 }
 
@@ -1241,26 +1240,24 @@
 
 Status StatsService::setBroadcastSubscriber(int64_t configId,
                                             int64_t subscriberId,
-                                            const sp<android::IBinder>& intentSender,
-                                            const String16& packageName) {
-    ENFORCE_DUMP_AND_USAGE_STATS(packageName);
+                                            const sp<IPendingIntentRef>& pir,
+                                            const int32_t callingUid) {
+    ENFORCE_UID(AID_SYSTEM);
 
     VLOG("StatsService::setBroadcastSubscriber called.");
-    IPCThreadState* ipc = IPCThreadState::self();
-    ConfigKey configKey(ipc->getCallingUid(), configId);
+    ConfigKey configKey(callingUid, configId);
     SubscriberReporter::getInstance()
-            .setBroadcastSubscriber(configKey, subscriberId, intentSender);
+            .setBroadcastSubscriber(configKey, subscriberId, pir);
     return Status::ok();
 }
 
 Status StatsService::unsetBroadcastSubscriber(int64_t configId,
                                               int64_t subscriberId,
-                                              const String16& packageName) {
-    ENFORCE_DUMP_AND_USAGE_STATS(packageName);
+                                              const int32_t callingUid) {
+    ENFORCE_UID(AID_SYSTEM);
 
     VLOG("StatsService::unsetBroadcastSubscriber called.");
-    IPCThreadState* ipc = IPCThreadState::self();
-    ConfigKey configKey(ipc->getCallingUid(), configId);
+    ConfigKey configKey(callingUid, configId);
     SubscriberReporter::getInstance()
             .unsetBroadcastSubscriber(configKey, subscriberId);
     return Status::ok();
@@ -1626,7 +1623,6 @@
     }
     mAnomalyAlarmMonitor->setStatsCompanionService(nullptr);
     mPeriodicAlarmMonitor->setStatsCompanionService(nullptr);
-    SubscriberReporter::getInstance().setStatsCompanionService(nullptr);
     mPullerManager->SetStatsCompanionService(nullptr);
 }
 
diff --git a/cmds/statsd/src/StatsService.h b/cmds/statsd/src/StatsService.h
index 0e89a2b..0565f3c 100644
--- a/cmds/statsd/src/StatsService.h
+++ b/cmds/statsd/src/StatsService.h
@@ -149,20 +149,19 @@
                                        const String16& packageName) override;
 
     /**
-     * Binder call to associate the given config's subscriberId with the given intentSender.
-     * intentSender must be convertible into an IntentSender (in Java) using IntentSender(IBinder).
+     * Binder call to associate the given config's subscriberId with the given pendingIntentRef.
      */
     virtual Status setBroadcastSubscriber(int64_t configId,
                                           int64_t subscriberId,
-                                          const sp<android::IBinder>& intentSender,
-                                          const String16& packageName) override;
+                                          const sp<IPendingIntentRef>& pir,
+                                          const int32_t callingUid) override;
 
     /**
-     * Binder call to unassociate the given config's subscriberId with any intentSender.
+     * Binder call to unassociate the given config's subscriberId with any pendingIntentRef.
      */
     virtual Status unsetBroadcastSubscriber(int64_t configId,
                                             int64_t subscriberId,
-                                            const String16& packageName) override;
+                                            const int32_t callingUid) override;
 
     /** Inform statsCompanion that statsd is ready. */
     virtual void sayHiToStatsCompanion();
diff --git a/cmds/statsd/src/subscriber/SubscriberReporter.cpp b/cmds/statsd/src/subscriber/SubscriberReporter.cpp
index 25d2257..a9a105f 100644
--- a/cmds/statsd/src/subscriber/SubscriberReporter.cpp
+++ b/cmds/statsd/src/subscriber/SubscriberReporter.cpp
@@ -19,7 +19,6 @@
 
 #include "SubscriberReporter.h"
 
-using android::IBinder;
 using std::lock_guard;
 using std::unordered_map;
 
@@ -29,12 +28,32 @@
 
 using std::vector;
 
+class BroadcastSubscriberDeathRecipient : public android::IBinder::DeathRecipient {
+    public:
+        BroadcastSubscriberDeathRecipient(const ConfigKey& configKey, int64_t subscriberId):
+            mConfigKey(configKey),
+            mSubscriberId(subscriberId) {}
+        ~BroadcastSubscriberDeathRecipient() override = default;
+    private:
+        ConfigKey mConfigKey;
+        int64_t mSubscriberId;
+
+    void binderDied(const android::wp<android::IBinder>& who) override {
+        if (IInterface::asBinder(SubscriberReporter::getInstance().getBroadcastSubscriber(
+              mConfigKey, mSubscriberId)) == who.promote()) {
+            SubscriberReporter::getInstance().unsetBroadcastSubscriber(mConfigKey, mSubscriberId);
+        }
+    }
+};
+
 void SubscriberReporter::setBroadcastSubscriber(const ConfigKey& configKey,
                                                 int64_t subscriberId,
-                                                const sp<IBinder>& intentSender) {
+                                                const sp<IPendingIntentRef>& pir) {
     VLOG("SubscriberReporter::setBroadcastSubscriber called.");
     lock_guard<std::mutex> lock(mLock);
-    mIntentMap[configKey][subscriberId] = intentSender;
+    mIntentMap[configKey][subscriberId] = pir;
+    IInterface::asBinder(pir)->linkToDeath(
+        new BroadcastSubscriberDeathRecipient(configKey, subscriberId));
 }
 
 void SubscriberReporter::unsetBroadcastSubscriber(const ConfigKey& configKey,
@@ -97,18 +116,13 @@
     sendBroadcastLocked(it2->second, configKey, subscription, cookies, dimKey);
 }
 
-void SubscriberReporter::sendBroadcastLocked(const sp<IBinder>& intentSender,
+void SubscriberReporter::sendBroadcastLocked(const sp<IPendingIntentRef>& pir,
                                              const ConfigKey& configKey,
                                              const Subscription& subscription,
                                              const vector<String16>& cookies,
                                              const MetricDimensionKey& dimKey) const {
     VLOG("SubscriberReporter::sendBroadcastLocked called.");
-    if (mStatsCompanionService == nullptr) {
-        ALOGW("Failed to send subscriber broadcast: could not access StatsCompanionService.");
-        return;
-    }
-    mStatsCompanionService->sendSubscriberBroadcast(
-            intentSender,
+    pir->sendSubscriberBroadcast(
             configKey.GetUid(),
             configKey.GetId(),
             subscription.id(),
@@ -117,6 +131,20 @@
             getStatsDimensionsValue(dimKey.getDimensionKeyInWhat()));
 }
 
+sp<IPendingIntentRef> SubscriberReporter::getBroadcastSubscriber(const ConfigKey& configKey,
+                                                                 int64_t subscriberId) {
+    lock_guard<std::mutex> lock(mLock);
+    auto subscriberMapIt = mIntentMap.find(configKey);
+    if (subscriberMapIt == mIntentMap.end()) {
+        return nullptr;
+    }
+    auto pirMapIt = subscriberMapIt->second.find(subscriberId);
+    if (pirMapIt == subscriberMapIt->second.end()) {
+        return nullptr;
+    }
+    return pirMapIt->second;
+}
+
 void getStatsDimensionsValueHelper(const vector<FieldValue>& dims, size_t* index, int depth,
                                    int prefix, vector<StatsDimensionsValue>* output) {
     size_t count = dims.size();
diff --git a/cmds/statsd/src/subscriber/SubscriberReporter.h b/cmds/statsd/src/subscriber/SubscriberReporter.h
index 2a7f771..8ccc8ee 100644
--- a/cmds/statsd/src/subscriber/SubscriberReporter.h
+++ b/cmds/statsd/src/subscriber/SubscriberReporter.h
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <android/os/IPendingIntentRef.h>
 #include <android/os/IStatsCompanionService.h>
 #include <utils/RefBase.h>
 
@@ -47,23 +48,11 @@
     void operator=(SubscriberReporter const&) = delete;
 
     /**
-     * Tells SubscriberReporter what IStatsCompanionService to use.
-     * May be nullptr, but SubscriberReporter will not send broadcasts for any calls
-     * to alertBroadcastSubscriber that occur while nullptr.
-     */
-    void setStatsCompanionService(sp<IStatsCompanionService> statsCompanionService) {
-        std::lock_guard<std::mutex> lock(mLock);
-        sp<IStatsCompanionService> tmpForLock = mStatsCompanionService;
-        mStatsCompanionService = statsCompanionService;
-    }
-
-    /**
      * Stores the given intentSender, associating it with the given (configKey, subscriberId) pair.
-     * intentSender must be convertible into an IntentSender (in Java) using IntentSender(IBinder).
      */
     void setBroadcastSubscriber(const ConfigKey& configKey,
                                 int64_t subscriberId,
-                                const sp<android::IBinder>& intentSender);
+                                const sp<IPendingIntentRef>& pir);
 
     /**
      * Erases any intentSender information from the given (configKey, subscriberId) pair.
@@ -82,6 +71,8 @@
                                   const Subscription& subscription,
                                   const MetricDimensionKey& dimKey) const;
 
+    sp<IPendingIntentRef> getBroadcastSubscriber(const ConfigKey& configKey, int64_t subscriberId);
+
     static StatsDimensionsValue getStatsDimensionsValue(const HashableDimensionKey& dim);
 
 private:
@@ -92,15 +83,15 @@
     /** Binder interface for communicating with StatsCompanionService. */
     sp<IStatsCompanionService> mStatsCompanionService = nullptr;
 
-    /** Maps <ConfigKey, SubscriberId> -> IBinder (which represents an IIntentSender). */
+    /** Maps <ConfigKey, SubscriberId> -> IPendingIntentRef (which represents a PendingIntent). */
     std::unordered_map<ConfigKey,
-            std::unordered_map<int64_t, sp<android::IBinder>>> mIntentMap;
+            std::unordered_map<int64_t, sp<IPendingIntentRef>>> mIntentMap;
 
     /**
      * Sends a broadcast via the given intentSender (using mStatsCompanionService), along
      * with the information in the other parameters.
      */
-    void sendBroadcastLocked(const sp<android::IBinder>& intentSender,
+    void sendBroadcastLocked(const sp<IPendingIntentRef>& pir,
                              const ConfigKey& configKey,
                              const Subscription& subscription,
                              const std::vector<String16>& cookies,
diff --git a/core/java/android/app/StatsManager.java b/core/java/android/app/StatsManager.java
index 51b2d40..83d1de6 100644
--- a/core/java/android/app/StatsManager.java
+++ b/core/java/android/app/StatsManager.java
@@ -255,18 +255,17 @@
             throws StatsUnavailableException {
         synchronized (sLock) {
             try {
-                IStatsd service = getIStatsdLocked();
+                IStatsManagerService service = getIStatsManagerServiceLocked();
                 if (pendingIntent != null) {
-                    // Extracts IIntentSender from the PendingIntent and turns it into an IBinder.
-                    IBinder intentSender = pendingIntent.getTarget().asBinder();
-                    service.setBroadcastSubscriber(configKey, subscriberId, intentSender,
+                    service.setBroadcastSubscriber(configKey, subscriberId, pendingIntent,
                             mContext.getOpPackageName());
                 } else {
                     service.unsetBroadcastSubscriber(configKey, subscriberId,
                             mContext.getOpPackageName());
                 }
             } catch (RemoteException e) {
-                Slog.e(TAG, "Failed to connect to statsd when adding broadcast subscriber", e);
+                Slog.e(TAG, "Failed to connect to statsmanager when adding broadcast subscriber",
+                        e);
                 throw new StatsUnavailableException("could not connect", e);
             } catch (SecurityException e) {
                 throw new StatsUnavailableException(e.getMessage(), e);