Allow pullAtomCallbacks to be unregistered in Java

Adds an API to unregister pullAtomCallbacks.
Fixes bug where the StatsEventParcels are null.
Makes onPullAtom a oneway call.
OnPullAtom returns an int code instead of a boolean
Made the APIs return RuntimeExceptions

Test: make, boots
Test: atest GtsStatsdHostTestCases
Bug: 146385842
Bug: 146385173
Bug: 144373250
Change-Id: I107a705a9024240c5c9f9e276293de8410e2b6f3
diff --git a/apex/statsd/aidl/android/os/IPullAtomCallback.aidl b/apex/statsd/aidl/android/os/IPullAtomCallback.aidl
index 88d3c3e..ff0b97b 100644
--- a/apex/statsd/aidl/android/os/IPullAtomCallback.aidl
+++ b/apex/statsd/aidl/android/os/IPullAtomCallback.aidl
@@ -26,6 +26,6 @@
     /**
      * Initiate a request for a pull for an atom.
      */
-     void onPullAtom(int atomTag, IPullAtomResultReceiver resultReceiver);
+     oneway void onPullAtom(int atomTag, IPullAtomResultReceiver resultReceiver);
 
 }
diff --git a/apex/statsd/aidl/android/os/IStatsCompanionService.aidl b/apex/statsd/aidl/android/os/IStatsCompanionService.aidl
index 22a2537..5a6118e 100644
--- a/apex/statsd/aidl/android/os/IStatsCompanionService.aidl
+++ b/apex/statsd/aidl/android/os/IStatsCompanionService.aidl
@@ -90,4 +90,7 @@
     /** Tells StatsCompanionService to tell statsd to register a puller for the given atom id */
     oneway void registerPullAtomCallback(int atomTag, long coolDownNs, long timeoutNs,
             in int[] additiveFields, IPullAtomCallback pullerCallback);
+
+    /** Tells StatsCompanionService to tell statsd to unregister a puller for the given atom id */
+    oneway void unregisterPullAtomCallback(int atomTag);
 }
diff --git a/apex/statsd/aidl/android/os/IStatsd.aidl b/apex/statsd/aidl/android/os/IStatsd.aidl
index cffc6ce..cce79fa 100644
--- a/apex/statsd/aidl/android/os/IStatsd.aidl
+++ b/apex/statsd/aidl/android/os/IStatsd.aidl
@@ -215,6 +215,11 @@
     */
    oneway void unregisterPullerCallback(int atomTag, String packageName);
 
+  /**
+   * Unregisters any pullAtomCallback for the given uid/atom.
+   */
+   oneway void unregisterPullAtomCallback(int uid, int atomTag);
+
     /**
      * The install requires staging.
      */
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 157da63..ff04462 100644
--- a/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
+++ b/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
@@ -2738,6 +2738,30 @@
         }
     }
 
+    @Override
+    public void unregisterPullAtomCallback(int atomTag) {
+        synchronized (sStatsdLock) {
+            // Always remove the puller in SCS.
+            // If statsd is down, we will not register it when it comes back up.
+            int callingUid = Binder.getCallingUid();
+            final long token = Binder.clearCallingIdentity();
+            PullerKey key = new PullerKey(callingUid, atomTag);
+            mPullers.remove(key);
+
+            if (sStatsd == null) {
+                Slog.w(TAG, "Could not access statsd for registering puller for atom " + atomTag);
+                return;
+            }
+            try {
+                sStatsd.unregisterPullAtomCallback(callingUid, atomTag);
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Failed to access statsd to register puller for atom " + atomTag);
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
+        }
+    }
+
     // Statsd related code
 
     /**
diff --git a/cmds/statsd/src/StatsService.cpp b/cmds/statsd/src/StatsService.cpp
index 4d38ba0..bb3a094 100644
--- a/cmds/statsd/src/StatsService.cpp
+++ b/cmds/statsd/src/StatsService.cpp
@@ -1320,6 +1320,13 @@
     return Status::ok();
 }
 
+Status StatsService::unregisterPullAtomCallback(int32_t uid, int32_t atomTag) {
+    ENFORCE_UID(AID_SYSTEM);
+    VLOG("StatsService::unregisterPullAtomCallback called.");
+    mPullerManager->UnregisterPullAtomCallback(uid, atomTag);
+    return Status::ok();
+}
+
 Status StatsService::sendBinaryPushStateChangedAtom(const android::String16& trainNameIn,
                                                     const int64_t trainVersionCodeIn,
                                                     const int options,
diff --git a/cmds/statsd/src/StatsService.h b/cmds/statsd/src/StatsService.h
index 9abf415..de55ca9 100644
--- a/cmds/statsd/src/StatsService.h
+++ b/cmds/statsd/src/StatsService.h
@@ -199,6 +199,11 @@
     virtual Status unregisterPullerCallback(int32_t atomTag, const String16& packageName) override;
 
     /**
+     * Binder call to unregister any existing callback for the given uid and atom.
+     */
+    virtual Status unregisterPullAtomCallback(int32_t uid, int32_t atomTag) override;
+
+    /**
      * Binder call to log BinaryPushStateChanged atom.
      */
     virtual Status sendBinaryPushStateChangedAtom(
diff --git a/cmds/statsd/src/external/StatsPuller.cpp b/cmds/statsd/src/external/StatsPuller.cpp
index 3c6bc2d..883bd28 100644
--- a/cmds/statsd/src/external/StatsPuller.cpp
+++ b/cmds/statsd/src/external/StatsPuller.cpp
@@ -47,6 +47,7 @@
         if (mHasGoodData) {
             (*data) = mCachedData;
             StatsdStats::getInstance().notePullFromCache(mTagId);
+
         }
         return mHasGoodData;
     }
diff --git a/cmds/statsd/src/external/StatsPullerManager.cpp b/cmds/statsd/src/external/StatsPullerManager.cpp
index 9ee627e..dbe73b4 100644
--- a/cmds/statsd/src/external/StatsPullerManager.cpp
+++ b/cmds/statsd/src/external/StatsPullerManager.cpp
@@ -290,6 +290,11 @@
 }
 
 bool StatsPullerManager::Pull(int tagId, vector<shared_ptr<LogEvent>>* data) {
+    AutoMutex _l(mLock);
+    return PullLocked(tagId, data);
+}
+
+bool StatsPullerManager::PullLocked(int tagId, vector<shared_ptr<LogEvent>>* data) {
     VLOG("Initiating pulling %d", tagId);
 
     if (kAllPullAtomInfo.find({.atomTag = tagId}) != kAllPullAtomInfo.end()) {
@@ -418,7 +423,7 @@
 
     for (const auto& pullInfo : needToPull) {
         vector<shared_ptr<LogEvent>> data;
-        bool pullSuccess = Pull(pullInfo.first, &data);
+        bool pullSuccess = PullLocked(pullInfo.first, &data);
         if (pullSuccess) {
             StatsdStats::getInstance().notePullDelay(
                     pullInfo.first, getElapsedRealtimeNs() - elapsedTimeNs);
@@ -518,6 +523,12 @@
     kAllPullAtomInfo.erase({.atomTag = atomTag});
 }
 
+void StatsPullerManager::UnregisterPullAtomCallback(const int uid, const int32_t atomTag) {
+    AutoMutex _l(mLock);
+    StatsdStats::getInstance().notePullerCallbackRegistrationChanged(atomTag, /*registered=*/false);
+    kAllPullAtomInfo.erase({.atomTag = atomTag});
+}
+
 }  // namespace statsd
 }  // namespace os
 }  // namespace android
diff --git a/cmds/statsd/src/external/StatsPullerManager.h b/cmds/statsd/src/external/StatsPullerManager.h
index 1bd9f92..349fd47 100644
--- a/cmds/statsd/src/external/StatsPullerManager.h
+++ b/cmds/statsd/src/external/StatsPullerManager.h
@@ -125,6 +125,8 @@
 
     void UnregisterPullerCallback(int32_t atomTag);
 
+    void UnregisterPullAtomCallback(const int uid, const int32_t atomTag);
+
     static std::map<PullerKey, PullAtomInfo> kAllPullAtomInfo;
 
 private:
@@ -139,6 +141,8 @@
     // mapping from simple matcher tagId to receivers
     std::map<int, std::list<ReceiverInfo>> mReceivers;
 
+    bool PullLocked(int tagId, vector<std::shared_ptr<LogEvent>>* data);
+
     // locks for data receiver and StatsCompanionService changes
     Mutex mLock;
 
diff --git a/core/java/android/app/StatsManager.java b/core/java/android/app/StatsManager.java
index cd855cf..2d851e0 100644
--- a/core/java/android/app/StatsManager.java
+++ b/core/java/android/app/StatsManager.java
@@ -107,6 +107,20 @@
      */
     public static final String ACTION_STATSD_STARTED = "android.app.action.STATSD_STARTED";
 
+    // Pull atom callback return codes.
+    /**
+     * Value indicating that this pull was successful and that the result should be used.
+     *
+     * @hide
+     **/
+    public static final int PULL_SUCCESS = 0;
+
+    /**
+     * Value indicating that this pull was unsuccessful and that the result should not be used.
+     * @hide
+     **/
+    public static final int PULL_SKIP = 1;
+
     private static final long DEFAULT_COOL_DOWN_NS = 1_000_000_000L; // 1 second.
     private static final long DEFAULT_TIMEOUT_NS = 10_000_000_000L; // 10 seconds.
 
@@ -508,13 +522,11 @@
      *                          additive fields for mapping isolated to host uids.
      * @param callback          The callback to be invoked when the stats service pulls the atom.
      * @param executor          The executor in which to run the callback
-     * @throws RemoteException  if unsuccessful due to failing to connect to system server.
      *
      * @hide
      */
     public void registerPullAtomCallback(int atomTag, @Nullable PullAtomMetadata metadata,
-            @NonNull StatsPullAtomCallback callback, @NonNull Executor executor)
-            throws RemoteException, SecurityException {
+            @NonNull StatsPullAtomCallback callback, @NonNull Executor executor) {
         long coolDownNs = metadata == null ? DEFAULT_COOL_DOWN_NS : metadata.mCoolDownNs;
         long timeoutNs = metadata == null ? DEFAULT_TIMEOUT_NS : metadata.mTimeoutNs;
         int[] additiveFields = metadata == null ? new int[0] : metadata.mAdditiveFields;
@@ -522,10 +534,34 @@
             additiveFields = new int[0];
         }
         synchronized (sLock) {
-            IStatsCompanionService service = getIStatsCompanionServiceLocked();
-            PullAtomCallbackInternal rec =
+            try {
+                IStatsCompanionService service = getIStatsCompanionServiceLocked();
+                PullAtomCallbackInternal rec =
                     new PullAtomCallbackInternal(atomTag, callback, executor);
-            service.registerPullAtomCallback(atomTag, coolDownNs, timeoutNs, additiveFields, rec);
+                service.registerPullAtomCallback(atomTag, coolDownNs, timeoutNs, additiveFields,
+                        rec);
+            } catch (RemoteException e) {
+                throw new RuntimeException("Unable to register pull callback", e);
+            }
+        }
+    }
+
+    /**
+     * Unregisters a callback for an atom when that atom is to be pulled. Note that any ongoing
+     * pulls will still occur.
+     *
+     * @param atomTag           The tag of the atom of which to unregister
+     *
+     * @hide
+     */
+    public void unregisterPullAtomCallback(int atomTag) {
+        synchronized (sLock) {
+            try {
+                IStatsCompanionService service = getIStatsCompanionServiceLocked();
+                service.unregisterPullAtomCallback(atomTag);
+            } catch (RemoteException e) {
+                throw new RuntimeException("Unable to unregister pull atom callback");
+            }
         }
     }
 
@@ -544,9 +580,11 @@
         public void onPullAtom(int atomTag, IPullAtomResultReceiver resultReceiver) {
             mExecutor.execute(() -> {
                 List<StatsEvent> data = new ArrayList<>();
-                boolean success = mCallback.onPullAtom(atomTag, data);
+                int successInt = mCallback.onPullAtom(atomTag, data);
+                boolean success = successInt == PULL_SUCCESS;
                 StatsEventParcel[] parcels = new StatsEventParcel[data.size()];
                 for (int i = 0; i < data.size(); i++) {
+                    parcels[i] = new StatsEventParcel();
                     parcels[i].buffer = data.get(i).getBytes();
                 }
                 try {
@@ -649,9 +687,9 @@
     public interface StatsPullAtomCallback {
         /**
          * Pull data for the specified atom tag, filling in the provided list of StatsEvent data.
-         * @return if the pull was successful
+         * @return {@link #PULL_SUCCESS} if the pull was successful, or {@link #PULL_SKIP} if not.
          */
-        boolean onPullAtom(int atomTag, List<StatsEvent> data);
+        int onPullAtom(int atomTag, List<StatsEvent> data);
     }
 
     private class StatsdDeathRecipient implements IBinder.DeathRecipient {