Statsd implementation of puller API

Internal implementation of the puller API. Registers pullers by putting
them in the kAllPullAtomInfo map. Implements the actual pull,
with condition variables to timeout.

Lastly, keys the kAllPullAtom info by a PullerKey, which is a uid and
atom id. However, the uid is just set to a default of -1 for now. I will
work the security implementation in a follow up CL.
Test: builds, boots
Test: I will write unit tests in the future. It's very difficult to
write any without StatsEvent being completed.

Change-Id: Id602dd297b6ba7df811e2d5ab2e77efc0684e418
diff --git a/cmds/statsd/Android.bp b/cmds/statsd/Android.bp
index cf286e6..360ddd4 100644
--- a/cmds/statsd/Android.bp
+++ b/cmds/statsd/Android.bp
@@ -69,9 +69,11 @@
         "src/external/GpuStatsPuller.cpp",
         "src/external/Perfetto.cpp",
         "src/external/PowerStatsPuller.cpp",
+        "src/external/PullResultReceiver.cpp",
         "src/external/puller_util.cpp",
         "src/external/ResourceHealthManagerPuller.cpp",
         "src/external/StatsCallbackPuller.cpp",
+        "src/external/StatsCallbackPullerDeprecated.cpp",
         "src/external/StatsCompanionServicePuller.cpp",
         "src/external/StatsPuller.cpp",
         "src/external/StatsPullerManager.cpp",
diff --git a/cmds/statsd/src/StatsService.cpp b/cmds/statsd/src/StatsService.cpp
index f072c9c..c376677 100644
--- a/cmds/statsd/src/StatsService.cpp
+++ b/cmds/statsd/src/StatsService.cpp
@@ -1292,7 +1292,11 @@
 Status StatsService::registerPullAtomCallback(int32_t uid, int32_t atomTag, int64_t coolDownNs,
                                     int64_t timeoutNs, const std::vector<int32_t>& additiveFields,
                                     const sp<android::os::IPullAtomCallback>& pullerCallback) {
+    ENFORCE_UID(AID_SYSTEM);
+
     VLOG("StatsService::registerPuller called.");
+    mPullerManager->RegisterPullAtomCallback(uid, atomTag, coolDownNs, timeoutNs, additiveFields,
+                                             pullerCallback);
     return Status::ok();
 }
 
diff --git a/cmds/statsd/src/external/PullResultReceiver.cpp b/cmds/statsd/src/external/PullResultReceiver.cpp
new file mode 100644
index 0000000..6bd0545
--- /dev/null
+++ b/cmds/statsd/src/external/PullResultReceiver.cpp
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "PullResultReceiver.h"
+
+using namespace android::binder;
+using namespace android::util;
+using namespace std;
+
+namespace android {
+namespace os {
+namespace statsd {
+
+PullResultReceiver::PullResultReceiver(
+        std::function<void(int32_t, bool, const vector<android::util::StatsEvent>&)> pullFinishCb)
+    : pullFinishCallback(std::move(pullFinishCb)) {
+}
+
+Status PullResultReceiver::pullFinished(int32_t atomTag, bool success,
+                                        const vector<StatsEvent>& output) {
+    pullFinishCallback(atomTag, success, output);
+    return Status::ok();
+}
+
+PullResultReceiver::~PullResultReceiver() {
+}
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android
\ No newline at end of file
diff --git a/cmds/statsd/src/external/PullResultReceiver.h b/cmds/statsd/src/external/PullResultReceiver.h
new file mode 100644
index 0000000..f731f77
--- /dev/null
+++ b/cmds/statsd/src/external/PullResultReceiver.h
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <android/os/BnPullAtomResultReceiver.h>
+
+using namespace std;
+
+namespace android {
+namespace os {
+namespace statsd {
+
+class PullResultReceiver : public BnPullAtomResultReceiver {
+public:
+    PullResultReceiver(function<void(int32_t, bool, const vector<android::util::StatsEvent>&)>
+                               pullFinishCallback);
+    ~PullResultReceiver();
+
+    /**
+     * Binder call for finishing a pull.
+     */
+    binder::Status pullFinished(int32_t atomTag, bool success,
+                                        const vector<android::util::StatsEvent>& output) override;
+
+private:
+    function<void(int32_t, bool, const vector<android::util::StatsEvent>&)> pullFinishCallback;
+};
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android
diff --git a/cmds/statsd/src/external/StatsCallbackPuller.cpp b/cmds/statsd/src/external/StatsCallbackPuller.cpp
index d718273..92db684 100644
--- a/cmds/statsd/src/external/StatsCallbackPuller.cpp
+++ b/cmds/statsd/src/external/StatsCallbackPuller.cpp
@@ -17,21 +17,27 @@
 #define DEBUG false  // STOPSHIP if true
 #include "Log.h"
 
-#include <android/os/IStatsPullerCallback.h>
-
 #include "StatsCallbackPuller.h"
+
+#include <android/os/IPullAtomCallback.h>
+#include <android/util/StatsEvent.h>
+
+#include "PullResultReceiver.h"
+#include "StatsPullerManager.h"
 #include "logd/LogEvent.h"
 #include "stats_log_util.h"
 
 using namespace android::binder;
+using namespace android::util;
+using namespace std;
 
 namespace android {
 namespace os {
 namespace statsd {
 
-StatsCallbackPuller::StatsCallbackPuller(int tagId, const sp<IStatsPullerCallback>& callback) :
-        StatsPuller(tagId), mCallback(callback) {
-        VLOG("StatsCallbackPuller created for tag %d", tagId);
+StatsCallbackPuller::StatsCallbackPuller(int tagId, const sp<IPullAtomCallback>& callback)
+    : StatsPuller(tagId), mCallback(callback) {
+    VLOG("StatsCallbackPuller created for tag %d", tagId);
 }
 
 bool StatsCallbackPuller::PullInternal(vector<shared_ptr<LogEvent>>* data) {
@@ -40,20 +46,57 @@
         ALOGW("No callback registered");
         return false;
     }
-    int64_t wallClockTimeNs = getWallClockNs();
-    int64_t elapsedTimeNs = getElapsedRealtimeNs();
-    vector<StatsLogEventWrapper> returned_value;
-    Status status = mCallback->pullData(mTagId, elapsedTimeNs, wallClockTimeNs, &returned_value);
+
+    // Shared variables needed in the result receiver.
+    shared_ptr<mutex> cv_mutex = make_shared<mutex>();
+    shared_ptr<condition_variable> cv = make_shared<condition_variable>();
+    shared_ptr<bool> pullFinish = make_shared<bool>(false);
+    shared_ptr<bool> pullSuccess = make_shared<bool>(false);
+    shared_ptr<vector<shared_ptr<LogEvent>>> sharedData =
+            make_shared<vector<shared_ptr<LogEvent>>>();
+
+    sp<PullResultReceiver> resultReceiver = new PullResultReceiver(
+            [cv_mutex, cv, pullFinish, pullSuccess, sharedData](
+                    int32_t atomTag, bool success, const vector<StatsEvent>& output) {
+                // This is the result of the pull, executing in a statsd binder thread.
+                // The pull could have taken a long time, and we should only modify
+                // data (the output param) if the pointer is in scope and the pull did not time out.
+                {
+                    lock_guard<mutex> lk(*cv_mutex);
+                    // TODO: fill the shared vector of LogEvents once StatsEvent is complete.
+                    *pullSuccess = success;
+                    *pullFinish = true;
+                }
+                cv->notify_one();
+            });
+
+    // Initiate the pull.
+    Status status = mCallback->onPullAtom(mTagId, resultReceiver);
     if (!status.isOk()) {
-        ALOGW("StatsCallbackPuller::pull failed for %d", mTagId);
         return false;
     }
-    data->clear();
-    for (const StatsLogEventWrapper& it: returned_value) {
-        LogEvent::createLogEvents(it, *data);
+
+    {
+        unique_lock<mutex> unique_lk(*cv_mutex);
+        int64_t pullTimeoutNs =
+                StatsPullerManager::kAllPullAtomInfo.at({.atomTag = mTagId}).pullTimeoutNs;
+        // Wait until the pull finishes, or until the pull timeout.
+        cv->wait_for(unique_lk, chrono::nanoseconds(pullTimeoutNs),
+                     [pullFinish] { return *pullFinish; });
+        if (!*pullFinish) {
+            // Note: The parent stats puller will also note that there was a timeout and that the
+            // cache should be cleared. Once we migrate all pullers to this callback, we could
+            // consolidate the logic.
+            return true;
+        } else {
+            // Only copy the data if we did not timeout and the pull was successful.
+            if (pullSuccess) {
+                *data = std::move(*sharedData);
+            }
+            VLOG("StatsCallbackPuller::pull succeeded for %d", mTagId);
+            return *pullSuccess;
+        }
     }
-    VLOG("StatsCallbackPuller::pull succeeded for %d", mTagId);
-    return true;
 }
 
 }  // namespace statsd
diff --git a/cmds/statsd/src/external/StatsCallbackPuller.h b/cmds/statsd/src/external/StatsCallbackPuller.h
index c4bfa89..ce506c7 100644
--- a/cmds/statsd/src/external/StatsCallbackPuller.h
+++ b/cmds/statsd/src/external/StatsCallbackPuller.h
@@ -16,8 +16,9 @@
 
 #pragma once
 
-#include <android/os/IStatsPullerCallback.h>
+#include <android/os/IPullAtomCallback.h>
 #include <utils/String16.h>
+
 #include "StatsPuller.h"
 
 namespace android {
@@ -26,11 +27,11 @@
 
 class StatsCallbackPuller : public StatsPuller {
 public:
-    explicit StatsCallbackPuller(int tagId, const sp<IStatsPullerCallback>& callback);
+    explicit StatsCallbackPuller(int tagId, const sp<IPullAtomCallback>& callback);
 
 private:
     bool PullInternal(vector<std::shared_ptr<LogEvent> >* data) override;
-    const sp<IStatsPullerCallback> mCallback;
+    const sp<IPullAtomCallback> mCallback;
 };
 
 }  // namespace statsd
diff --git a/cmds/statsd/src/external/StatsCallbackPullerDeprecated.cpp b/cmds/statsd/src/external/StatsCallbackPullerDeprecated.cpp
new file mode 100644
index 0000000..4f88a91
--- /dev/null
+++ b/cmds/statsd/src/external/StatsCallbackPullerDeprecated.cpp
@@ -0,0 +1,63 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#define DEBUG false  // STOPSHIP if true
+#include "Log.h"
+
+#include "StatsCallbackPullerDeprecated.h"
+
+#include <android/os/IStatsPullerCallback.h>
+
+#include "logd/LogEvent.h"
+#include "stats_log_util.h"
+
+using namespace android::binder;
+
+namespace android {
+namespace os {
+namespace statsd {
+
+StatsCallbackPullerDeprecated::StatsCallbackPullerDeprecated(
+        int tagId, const sp<IStatsPullerCallback>& callback)
+    : StatsPuller(tagId), mCallback(callback) {
+    VLOG("StatsCallbackPuller created for tag %d", tagId);
+}
+
+bool StatsCallbackPullerDeprecated::PullInternal(vector<shared_ptr<LogEvent>>* data) {
+    VLOG("StatsCallbackPuller called for tag %d", mTagId)
+    if (mCallback == nullptr) {
+        ALOGW("No callback registered");
+        return false;
+    }
+    int64_t wallClockTimeNs = getWallClockNs();
+    int64_t elapsedTimeNs = getElapsedRealtimeNs();
+    vector<StatsLogEventWrapper> returned_value;
+    Status status = mCallback->pullData(mTagId, elapsedTimeNs, wallClockTimeNs, &returned_value);
+    if (!status.isOk()) {
+        ALOGW("StatsCallbackPuller::pull failed for %d", mTagId);
+        return false;
+    }
+    data->clear();
+    for (const StatsLogEventWrapper& it : returned_value) {
+        LogEvent::createLogEvents(it, *data);
+    }
+    VLOG("StatsCallbackPuller::pull succeeded for %d", mTagId);
+    return true;
+}
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android
diff --git a/cmds/statsd/src/external/StatsCallbackPullerDeprecated.h b/cmds/statsd/src/external/StatsCallbackPullerDeprecated.h
new file mode 100644
index 0000000..0289029
--- /dev/null
+++ b/cmds/statsd/src/external/StatsCallbackPullerDeprecated.h
@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <android/os/IStatsPullerCallback.h>
+#include <utils/String16.h>
+
+#include "StatsPuller.h"
+
+namespace android {
+namespace os {
+namespace statsd {
+
+class StatsCallbackPullerDeprecated : public StatsPuller {
+public:
+    explicit StatsCallbackPullerDeprecated(int tagId, const sp<IStatsPullerCallback>& callback);
+
+private:
+    bool PullInternal(vector<std::shared_ptr<LogEvent> >* data) override;
+    const sp<IStatsPullerCallback> mCallback;
+};
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android
diff --git a/cmds/statsd/src/external/StatsPuller.cpp b/cmds/statsd/src/external/StatsPuller.cpp
index 9552c0a..3c6bc2d 100644
--- a/cmds/statsd/src/external/StatsPuller.cpp
+++ b/cmds/statsd/src/external/StatsPuller.cpp
@@ -40,8 +40,9 @@
     lock_guard<std::mutex> lock(mLock);
     int64_t elapsedTimeNs = getElapsedRealtimeNs();
     StatsdStats::getInstance().notePull(mTagId);
-    const bool shouldUseCache = elapsedTimeNs - mLastPullTimeNs <
-                                StatsPullerManager::kAllPullAtomInfo.at(mTagId).coolDownNs;
+    const bool shouldUseCache =
+            elapsedTimeNs - mLastPullTimeNs <
+            StatsPullerManager::kAllPullAtomInfo.at({.atomTag = mTagId}).coolDownNs;
     if (shouldUseCache) {
         if (mHasGoodData) {
             (*data) = mCachedData;
@@ -63,7 +64,8 @@
     const int64_t pullDurationNs = getElapsedRealtimeNs() - elapsedTimeNs;
     StatsdStats::getInstance().notePullTime(mTagId, pullDurationNs);
     const bool pullTimeOut =
-            pullDurationNs > StatsPullerManager::kAllPullAtomInfo.at(mTagId).pullTimeoutNs;
+            pullDurationNs >
+            StatsPullerManager::kAllPullAtomInfo.at({.atomTag = mTagId}).pullTimeoutNs;
     if (pullTimeOut) {
         // Something went wrong. Discard the data.
         clearCacheLocked();
@@ -100,7 +102,7 @@
 
 int StatsPuller::ClearCacheIfNecessary(int64_t timestampNs) {
     if (timestampNs - mLastPullTimeNs >
-        StatsPullerManager::kAllPullAtomInfo.at(mTagId).coolDownNs) {
+        StatsPullerManager::kAllPullAtomInfo.at({.atomTag = mTagId}).coolDownNs) {
         return clearCache();
     } else {
         return 0;
diff --git a/cmds/statsd/src/external/StatsPullerManager.cpp b/cmds/statsd/src/external/StatsPullerManager.cpp
index 535fcfb..ce27ce6 100644
--- a/cmds/statsd/src/external/StatsPullerManager.cpp
+++ b/cmds/statsd/src/external/StatsPullerManager.cpp
@@ -19,6 +19,7 @@
 
 #include "StatsPullerManager.h"
 
+#include <android/os/IPullAtomCallback.h>
 #include <android/os/IStatsCompanionService.h>
 #include <android/os/IStatsPullerCallback.h>
 #include <cutils/log.h>
@@ -37,6 +38,7 @@
 #include "PowerStatsPuller.h"
 #include "ResourceHealthManagerPuller.h"
 #include "StatsCallbackPuller.h"
+#include "StatsCallbackPullerDeprecated.h"
 #include "StatsCompanionServicePuller.h"
 #include "SubsystemSleepStatePuller.h"
 #include "SurfaceflingerStatsPuller.h"
@@ -57,224 +59,226 @@
 // Values smaller than this may require to update the alarm.
 const int64_t NO_ALARM_UPDATE = INT64_MAX;
 
-std::map<int, PullAtomInfo> StatsPullerManager::kAllPullAtomInfo = {
+std::map<PullerKey, PullAtomInfo> StatsPullerManager::kAllPullAtomInfo = {
         // wifi_bytes_transfer
-        {android::util::WIFI_BYTES_TRANSFER,
+        {{.atomTag = android::util::WIFI_BYTES_TRANSFER},
          {.additiveFields = {2, 3, 4, 5},
           .puller = new StatsCompanionServicePuller(android::util::WIFI_BYTES_TRANSFER)}},
         // wifi_bytes_transfer_by_fg_bg
-        {android::util::WIFI_BYTES_TRANSFER_BY_FG_BG,
+        {{.atomTag = android::util::WIFI_BYTES_TRANSFER_BY_FG_BG},
          {.additiveFields = {3, 4, 5, 6},
           .puller = new StatsCompanionServicePuller(android::util::WIFI_BYTES_TRANSFER_BY_FG_BG)}},
         // mobile_bytes_transfer
-        {android::util::MOBILE_BYTES_TRANSFER,
+        {{.atomTag = android::util::MOBILE_BYTES_TRANSFER},
          {.additiveFields = {2, 3, 4, 5},
           .puller = new StatsCompanionServicePuller(android::util::MOBILE_BYTES_TRANSFER)}},
         // mobile_bytes_transfer_by_fg_bg
-        {android::util::MOBILE_BYTES_TRANSFER_BY_FG_BG,
+        {{.atomTag = android::util::MOBILE_BYTES_TRANSFER_BY_FG_BG},
          {.additiveFields = {3, 4, 5, 6},
           .puller =
                   new StatsCompanionServicePuller(android::util::MOBILE_BYTES_TRANSFER_BY_FG_BG)}},
         // bluetooth_bytes_transfer
-        {android::util::BLUETOOTH_BYTES_TRANSFER,
+        {{.atomTag = android::util::BLUETOOTH_BYTES_TRANSFER},
          {.additiveFields = {2, 3},
           .puller = new StatsCompanionServicePuller(android::util::BLUETOOTH_BYTES_TRANSFER)}},
         // kernel_wakelock
-        {android::util::KERNEL_WAKELOCK,
+        {{.atomTag = android::util::KERNEL_WAKELOCK},
          {.puller = new StatsCompanionServicePuller(android::util::KERNEL_WAKELOCK)}},
         // subsystem_sleep_state
-        {android::util::SUBSYSTEM_SLEEP_STATE, {.puller = new SubsystemSleepStatePuller()}},
+        {{.atomTag = android::util::SUBSYSTEM_SLEEP_STATE},
+         {.puller = new SubsystemSleepStatePuller()}},
         // on_device_power_measurement
-        {android::util::ON_DEVICE_POWER_MEASUREMENT, {.puller = new PowerStatsPuller()}},
+        {{.atomTag = android::util::ON_DEVICE_POWER_MEASUREMENT},
+         {.puller = new PowerStatsPuller()}},
         // cpu_time_per_freq
-        {android::util::CPU_TIME_PER_FREQ,
+        {{.atomTag = android::util::CPU_TIME_PER_FREQ},
          {.additiveFields = {3},
           .puller = new StatsCompanionServicePuller(android::util::CPU_TIME_PER_FREQ)}},
         // cpu_time_per_uid
-        {android::util::CPU_TIME_PER_UID,
+        {{.atomTag = android::util::CPU_TIME_PER_UID},
          {.additiveFields = {2, 3},
           .puller = new StatsCompanionServicePuller(android::util::CPU_TIME_PER_UID)}},
         // cpu_time_per_uid_freq
         // the throttling is 3sec, handled in
         // frameworks/base/core/java/com/android/internal/os/KernelCpuProcReader
-        {android::util::CPU_TIME_PER_UID_FREQ,
+        {{.atomTag = android::util::CPU_TIME_PER_UID_FREQ},
          {.additiveFields = {4},
           .puller = new StatsCompanionServicePuller(android::util::CPU_TIME_PER_UID_FREQ)}},
         // cpu_active_time
         // the throttling is 3sec, handled in
         // frameworks/base/core/java/com/android/internal/os/KernelCpuProcReader
-        {android::util::CPU_ACTIVE_TIME,
+        {{.atomTag = android::util::CPU_ACTIVE_TIME},
          {.additiveFields = {2},
           .puller = new StatsCompanionServicePuller(android::util::CPU_ACTIVE_TIME)}},
         // cpu_cluster_time
         // the throttling is 3sec, handled in
         // frameworks/base/core/java/com/android/internal/os/KernelCpuProcReader
-        {android::util::CPU_CLUSTER_TIME,
+        {{.atomTag = android::util::CPU_CLUSTER_TIME},
          {.additiveFields = {3},
           .puller = new StatsCompanionServicePuller(android::util::CPU_CLUSTER_TIME)}},
         // wifi_activity_energy_info
-        {android::util::WIFI_ACTIVITY_INFO,
+        {{.atomTag = android::util::WIFI_ACTIVITY_INFO},
          {.puller = new StatsCompanionServicePuller(android::util::WIFI_ACTIVITY_INFO)}},
         // modem_activity_info
-        {android::util::MODEM_ACTIVITY_INFO,
+        {{.atomTag = android::util::MODEM_ACTIVITY_INFO},
          {.puller = new StatsCompanionServicePuller(android::util::MODEM_ACTIVITY_INFO)}},
         // bluetooth_activity_info
-        {android::util::BLUETOOTH_ACTIVITY_INFO,
+        {{.atomTag = android::util::BLUETOOTH_ACTIVITY_INFO},
          {.puller = new StatsCompanionServicePuller(android::util::BLUETOOTH_ACTIVITY_INFO)}},
         // system_elapsed_realtime
-        {android::util::SYSTEM_ELAPSED_REALTIME,
+        {{.atomTag = android::util::SYSTEM_ELAPSED_REALTIME},
          {.coolDownNs = NS_PER_SEC,
           .puller = new StatsCompanionServicePuller(android::util::SYSTEM_ELAPSED_REALTIME),
           .pullTimeoutNs = NS_PER_SEC / 2,
          }},
         // system_uptime
-        {android::util::SYSTEM_UPTIME,
+        {{.atomTag = android::util::SYSTEM_UPTIME},
          {.puller = new StatsCompanionServicePuller(android::util::SYSTEM_UPTIME)}},
         // remaining_battery_capacity
-        {android::util::REMAINING_BATTERY_CAPACITY,
+        {{.atomTag = android::util::REMAINING_BATTERY_CAPACITY},
          {.puller = new ResourceHealthManagerPuller(android::util::REMAINING_BATTERY_CAPACITY)}},
         // full_battery_capacity
-        {android::util::FULL_BATTERY_CAPACITY,
+        {{.atomTag = android::util::FULL_BATTERY_CAPACITY},
          {.puller = new ResourceHealthManagerPuller(android::util::FULL_BATTERY_CAPACITY)}},
         // battery_voltage
-        {android::util::BATTERY_VOLTAGE,
+        {{.atomTag = android::util::BATTERY_VOLTAGE},
          {.puller = new ResourceHealthManagerPuller(android::util::BATTERY_VOLTAGE)}},
         // battery_level
-        {android::util::BATTERY_LEVEL,
+        {{.atomTag = android::util::BATTERY_LEVEL},
          {.puller = new ResourceHealthManagerPuller(android::util::BATTERY_LEVEL)}},
         // battery_cycle_count
-        {android::util::BATTERY_CYCLE_COUNT,
+        {{.atomTag = android::util::BATTERY_CYCLE_COUNT},
          {.puller = new ResourceHealthManagerPuller(android::util::BATTERY_CYCLE_COUNT)}},
         // process_memory_state
-        {android::util::PROCESS_MEMORY_STATE,
+        {{.atomTag = android::util::PROCESS_MEMORY_STATE},
          {.additiveFields = {4, 5, 6, 7, 8},
           .puller = new StatsCompanionServicePuller(android::util::PROCESS_MEMORY_STATE)}},
         // process_memory_high_water_mark
-        {android::util::PROCESS_MEMORY_HIGH_WATER_MARK,
+        {{.atomTag = android::util::PROCESS_MEMORY_HIGH_WATER_MARK},
          {.puller =
                   new StatsCompanionServicePuller(android::util::PROCESS_MEMORY_HIGH_WATER_MARK)}},
         // process_memory_snapshot
-        {android::util::PROCESS_MEMORY_SNAPSHOT,
+        {{.atomTag = android::util::PROCESS_MEMORY_SNAPSHOT},
          {.puller = new StatsCompanionServicePuller(android::util::PROCESS_MEMORY_SNAPSHOT)}},
         // system_ion_heap_size
-        {android::util::SYSTEM_ION_HEAP_SIZE,
+        {{.atomTag = android::util::SYSTEM_ION_HEAP_SIZE},
          {.puller = new StatsCompanionServicePuller(android::util::SYSTEM_ION_HEAP_SIZE)}},
         // process_system_ion_heap_size
-        {android::util::PROCESS_SYSTEM_ION_HEAP_SIZE,
+        {{.atomTag = android::util::PROCESS_SYSTEM_ION_HEAP_SIZE},
          {.puller = new StatsCompanionServicePuller(android::util::PROCESS_SYSTEM_ION_HEAP_SIZE)}},
         // temperature
-        {android::util::TEMPERATURE,
+        {{.atomTag = android::util::TEMPERATURE},
          {.puller = new StatsCompanionServicePuller(android::util::TEMPERATURE)}},
         // cooling_device
-        {android::util::COOLING_DEVICE,
+        {{.atomTag = android::util::COOLING_DEVICE},
          {.puller = new StatsCompanionServicePuller(android::util::COOLING_DEVICE)}},
         // binder_calls
-        {android::util::BINDER_CALLS,
+        {{.atomTag = android::util::BINDER_CALLS},
          {.additiveFields = {4, 5, 6, 8, 12},
           .puller = new StatsCompanionServicePuller(android::util::BINDER_CALLS)}},
         // binder_calls_exceptions
-        {android::util::BINDER_CALLS_EXCEPTIONS,
+        {{.atomTag = android::util::BINDER_CALLS_EXCEPTIONS},
          {.puller = new StatsCompanionServicePuller(android::util::BINDER_CALLS_EXCEPTIONS)}},
         // looper_stats
-        {android::util::LOOPER_STATS,
+        {{.atomTag = android::util::LOOPER_STATS},
          {.additiveFields = {5, 6, 7, 8, 9},
           .puller = new StatsCompanionServicePuller(android::util::LOOPER_STATS)}},
         // Disk Stats
-        {android::util::DISK_STATS,
+        {{.atomTag = android::util::DISK_STATS},
          {.puller = new StatsCompanionServicePuller(android::util::DISK_STATS)}},
         // Directory usage
-        {android::util::DIRECTORY_USAGE,
+        {{.atomTag = android::util::DIRECTORY_USAGE},
          {.puller = new StatsCompanionServicePuller(android::util::DIRECTORY_USAGE)}},
         // Size of app's code, data, and cache
-        {android::util::APP_SIZE,
+        {{.atomTag = android::util::APP_SIZE},
          {.puller = new StatsCompanionServicePuller(android::util::APP_SIZE)}},
         // Size of specific categories of files. Eg. Music.
-        {android::util::CATEGORY_SIZE,
+        {{.atomTag = android::util::CATEGORY_SIZE},
          {.puller = new StatsCompanionServicePuller(android::util::CATEGORY_SIZE)}},
         // Number of fingerprints enrolled for each user.
-        {android::util::NUM_FINGERPRINTS_ENROLLED,
+        {{.atomTag = android::util::NUM_FINGERPRINTS_ENROLLED},
          {.puller = new StatsCompanionServicePuller(android::util::NUM_FINGERPRINTS_ENROLLED)}},
         // Number of faces enrolled for each user.
-        {android::util::NUM_FACES_ENROLLED,
+        {{.atomTag = android::util::NUM_FACES_ENROLLED},
          {.puller = new StatsCompanionServicePuller(android::util::NUM_FACES_ENROLLED)}},
         // ProcStats.
-        {android::util::PROC_STATS,
+        {{.atomTag = android::util::PROC_STATS},
          {.puller = new StatsCompanionServicePuller(android::util::PROC_STATS)}},
         // ProcStatsPkgProc.
-        {android::util::PROC_STATS_PKG_PROC,
+        {{.atomTag = android::util::PROC_STATS_PKG_PROC},
          {.puller = new StatsCompanionServicePuller(android::util::PROC_STATS_PKG_PROC)}},
         // Disk I/O stats per uid.
-        {android::util::DISK_IO,
+        {{.atomTag = android::util::DISK_IO},
          {.additiveFields = {2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
           .coolDownNs = 3 * NS_PER_SEC,
           .puller = new StatsCompanionServicePuller(android::util::DISK_IO)}},
         // PowerProfile constants for power model calculations.
-        {android::util::POWER_PROFILE,
+        {{.atomTag = android::util::POWER_PROFILE},
          {.puller = new StatsCompanionServicePuller(android::util::POWER_PROFILE)}},
         // Process cpu stats. Min cool-down is 5 sec, inline with what AcitivityManagerService uses.
-        {android::util::PROCESS_CPU_TIME,
+        {{.atomTag = android::util::PROCESS_CPU_TIME},
          {.coolDownNs = 5 * NS_PER_SEC /* min cool-down in seconds*/,
           .puller = new StatsCompanionServicePuller(android::util::PROCESS_CPU_TIME)}},
-        {android::util::CPU_TIME_PER_THREAD_FREQ,
+        {{.atomTag = android::util::CPU_TIME_PER_THREAD_FREQ},
          {.additiveFields = {7, 9, 11, 13, 15, 17, 19, 21},
           .puller = new StatsCompanionServicePuller(android::util::CPU_TIME_PER_THREAD_FREQ)}},
         // DeviceCalculatedPowerUse.
-        {android::util::DEVICE_CALCULATED_POWER_USE,
+        {{.atomTag = android::util::DEVICE_CALCULATED_POWER_USE},
          {.puller = new StatsCompanionServicePuller(android::util::DEVICE_CALCULATED_POWER_USE)}},
         // DeviceCalculatedPowerBlameUid.
-        {android::util::DEVICE_CALCULATED_POWER_BLAME_UID,
+        {{.atomTag = android::util::DEVICE_CALCULATED_POWER_BLAME_UID},
          {.puller = new StatsCompanionServicePuller(
                   android::util::DEVICE_CALCULATED_POWER_BLAME_UID)}},
         // DeviceCalculatedPowerBlameOther.
-        {android::util::DEVICE_CALCULATED_POWER_BLAME_OTHER,
+        {{.atomTag = android::util::DEVICE_CALCULATED_POWER_BLAME_OTHER},
          {.puller = new StatsCompanionServicePuller(
                   android::util::DEVICE_CALCULATED_POWER_BLAME_OTHER)}},
         // DebugElapsedClock.
-        {android::util::DEBUG_ELAPSED_CLOCK,
+        {{.atomTag = android::util::DEBUG_ELAPSED_CLOCK},
          {.additiveFields = {1, 2, 3, 4},
           .puller = new StatsCompanionServicePuller(android::util::DEBUG_ELAPSED_CLOCK)}},
         // DebugFailingElapsedClock.
-        {android::util::DEBUG_FAILING_ELAPSED_CLOCK,
+        {{.atomTag = android::util::DEBUG_FAILING_ELAPSED_CLOCK},
          {.additiveFields = {1, 2, 3, 4},
           .puller = new StatsCompanionServicePuller(android::util::DEBUG_FAILING_ELAPSED_CLOCK)}},
         // BuildInformation.
-        {android::util::BUILD_INFORMATION,
+        {{.atomTag = android::util::BUILD_INFORMATION},
          {.puller = new StatsCompanionServicePuller(android::util::BUILD_INFORMATION)}},
         // RoleHolder.
-        {android::util::ROLE_HOLDER,
+        {{.atomTag = android::util::ROLE_HOLDER},
          {.puller = new StatsCompanionServicePuller(android::util::ROLE_HOLDER)}},
         // PermissionState.
-        {android::util::DANGEROUS_PERMISSION_STATE,
+        {{.atomTag = android::util::DANGEROUS_PERMISSION_STATE},
          {.puller = new StatsCompanionServicePuller(android::util::DANGEROUS_PERMISSION_STATE)}},
         // TrainInfo.
-        {android::util::TRAIN_INFO, {.puller = new TrainInfoPuller()}},
+        {{.atomTag = android::util::TRAIN_INFO}, {.puller = new TrainInfoPuller()}},
         // TimeZoneDataInfo.
-        {android::util::TIME_ZONE_DATA_INFO,
+        {{.atomTag = android::util::TIME_ZONE_DATA_INFO},
          {.puller = new StatsCompanionServicePuller(android::util::TIME_ZONE_DATA_INFO)}},
         // ExternalStorageInfo
-        {android::util::EXTERNAL_STORAGE_INFO,
+        {{.atomTag = android::util::EXTERNAL_STORAGE_INFO},
          {.puller = new StatsCompanionServicePuller(android::util::EXTERNAL_STORAGE_INFO)}},
         // GpuStatsGlobalInfo
-        {android::util::GPU_STATS_GLOBAL_INFO,
+        {{.atomTag = android::util::GPU_STATS_GLOBAL_INFO},
          {.puller = new GpuStatsPuller(android::util::GPU_STATS_GLOBAL_INFO)}},
         // GpuStatsAppInfo
-        {android::util::GPU_STATS_APP_INFO,
+        {{.atomTag = android::util::GPU_STATS_APP_INFO},
          {.puller = new GpuStatsPuller(android::util::GPU_STATS_APP_INFO)}},
         // AppsOnExternalStorageInfo
-        {android::util::APPS_ON_EXTERNAL_STORAGE_INFO,
+        {{.atomTag = android::util::APPS_ON_EXTERNAL_STORAGE_INFO},
          {.puller = new StatsCompanionServicePuller(android::util::APPS_ON_EXTERNAL_STORAGE_INFO)}},
         // Face Settings
-        {android::util::FACE_SETTINGS,
+        {{.atomTag = android::util::FACE_SETTINGS},
          {.puller = new StatsCompanionServicePuller(android::util::FACE_SETTINGS)}},
         // App ops
-        {android::util::APP_OPS,
+        {{.atomTag = android::util::APP_OPS},
          {.puller = new StatsCompanionServicePuller(android::util::APP_OPS)}},
         // SurfaceflingerStatsGlobalInfo
-        {android::util::SURFACEFLINGER_STATS_GLOBAL_INFO,
+        {{.atomTag = android::util::SURFACEFLINGER_STATS_GLOBAL_INFO},
          {.puller =
                   new SurfaceflingerStatsPuller(android::util::SURFACEFLINGER_STATS_GLOBAL_INFO)}},
         // VmsClientStats
-        {android::util::VMS_CLIENT_STATS,
+        {{.atomTag = android::util::VMS_CLIENT_STATS},
          {.additiveFields = {5, 6, 7, 8, 9, 10},
           .puller = new CarStatsPuller(android::util::VMS_CLIENT_STATS)}},
 };
@@ -285,8 +289,8 @@
 bool StatsPullerManager::Pull(int tagId, vector<shared_ptr<LogEvent>>* data) {
     VLOG("Initiating pulling %d", tagId);
 
-    if (kAllPullAtomInfo.find(tagId) != kAllPullAtomInfo.end()) {
-        bool ret = kAllPullAtomInfo.find(tagId)->second.puller->Pull(data);
+    if (kAllPullAtomInfo.find({.atomTag = tagId}) != kAllPullAtomInfo.end()) {
+        bool ret = kAllPullAtomInfo.find({.atomTag = tagId})->second.puller->Pull(data);
         VLOG("pulled %d items", (int)data->size());
         if (!ret) {
             StatsdStats::getInstance().notePullFailed(tagId);
@@ -300,7 +304,8 @@
 
 bool StatsPullerManager::PullerForMatcherExists(int tagId) const {
     // Vendor pulled atoms might be registered after we parse the config.
-    return isVendorPulledAtom(tagId) || kAllPullAtomInfo.find(tagId) != kAllPullAtomInfo.end();
+    return isVendorPulledAtom(tagId) ||
+           kAllPullAtomInfo.find({.atomTag = tagId}) != kAllPullAtomInfo.end();
 }
 
 void StatsPullerManager::updateAlarmLocked() {
@@ -469,6 +474,7 @@
     return totalCleared;
 }
 
+// Deprecated, remove after puller API is complete.
 void StatsPullerManager::RegisterPullerCallback(int32_t atomTag,
         const sp<IStatsPullerCallback>& callback) {
     AutoMutex _l(mLock);
@@ -479,7 +485,22 @@
     }
     VLOG("RegisterPullerCallback: adding puller for tag %d", atomTag);
     StatsdStats::getInstance().notePullerCallbackRegistrationChanged(atomTag, /*registered=*/true);
-    kAllPullAtomInfo[atomTag] = {.puller = new StatsCallbackPuller(atomTag, callback)};
+    kAllPullAtomInfo[{.atomTag = atomTag}] = {
+            .puller = new StatsCallbackPullerDeprecated(atomTag, callback)};
+}
+
+void StatsPullerManager::RegisterPullAtomCallback(const int uid, const int32_t atomTag,
+                                                  const int64_t coolDownNs, const int64_t timeoutNs,
+                                                  const vector<int32_t>& additiveFields,
+                                                  const sp<IPullAtomCallback>& callback) {
+    AutoMutex _l(mLock);
+    VLOG("RegisterPullerCallback: adding puller for tag %d", atomTag);
+    // TODO: linkToDeath with the callback so that we can remove it and delete the puller.
+    StatsdStats::getInstance().notePullerCallbackRegistrationChanged(atomTag, /*registered=*/true);
+    kAllPullAtomInfo[{.atomTag = atomTag}] = {.additiveFields = additiveFields,
+                                              .coolDownNs = coolDownNs,
+                                              .pullTimeoutNs = timeoutNs,
+                                              .puller = new StatsCallbackPuller(atomTag, callback)};
 }
 
 void StatsPullerManager::UnregisterPullerCallback(int32_t atomTag) {
@@ -489,7 +510,7 @@
         return;
     }
     StatsdStats::getInstance().notePullerCallbackRegistrationChanged(atomTag, /*registered=*/false);
-    kAllPullAtomInfo.erase(atomTag);
+    kAllPullAtomInfo.erase({.atomTag = atomTag});
 }
 
 }  // namespace statsd
diff --git a/cmds/statsd/src/external/StatsPullerManager.h b/cmds/statsd/src/external/StatsPullerManager.h
index 4ea1386..1bd9f92 100644
--- a/cmds/statsd/src/external/StatsPullerManager.h
+++ b/cmds/statsd/src/external/StatsPullerManager.h
@@ -16,15 +16,18 @@
 
 #pragma once
 
+#include <android/os/IPullAtomCallback.h>
 #include <android/os/IStatsCompanionService.h>
 #include <android/os/IStatsPullerCallback.h>
 #include <binder/IServiceManager.h>
 #include <utils/RefBase.h>
 #include <utils/threads.h>
+
 #include <list>
 #include <string>
 #include <unordered_map>
 #include <vector>
+
 #include "PullDataReceiver.h"
 #include "StatsPuller.h"
 #include "guardrail/StatsdStats.h"
@@ -53,6 +56,27 @@
     int64_t pullTimeoutNs = StatsdStats::kPullMaxDelayNs;
 } PullAtomInfo;
 
+typedef struct PullerKey {
+    // The uid of the process that registers this puller.
+    const int uid = -1;
+    // The atom that this puller is for.
+    const int atomTag;
+
+    bool operator<(const PullerKey& that) const {
+        if (uid < that.uid) {
+            return true;
+        }
+        if (uid > that.uid) {
+            return false;
+        }
+        return atomTag < that.atomTag;
+    };
+
+    bool operator==(const PullerKey& that) const {
+        return uid == that.uid && atomTag == that.atomTag;
+    };
+} PullerKey;
+
 class StatsPullerManager : public virtual RefBase {
 public:
     StatsPullerManager();
@@ -92,12 +116,16 @@
 
     void SetStatsCompanionService(sp<IStatsCompanionService> statsCompanionService);
 
-    void RegisterPullerCallback(int32_t atomTag,
-                                const sp<IStatsPullerCallback>& callback);
+    // Deprecated, remove after puller API is complete.
+    void RegisterPullerCallback(int32_t atomTag, const sp<IStatsPullerCallback>& callback);
+
+    void RegisterPullAtomCallback(const int uid, const int32_t atomTag, const int64_t coolDownNs,
+                                  const int64_t timeoutNs, const vector<int32_t>& additiveFields,
+                                  const sp<IPullAtomCallback>& callback);
 
     void UnregisterPullerCallback(int32_t atomTag);
 
-    static std::map<int, PullAtomInfo> kAllPullAtomInfo;
+    static std::map<PullerKey, PullAtomInfo> kAllPullAtomInfo;
 
 private:
     sp<IStatsCompanionService> mStatsCompanionService = nullptr;
diff --git a/cmds/statsd/src/external/puller_util.cpp b/cmds/statsd/src/external/puller_util.cpp
index 0b9b6ab..53fa630 100644
--- a/cmds/statsd/src/external/puller_util.cpp
+++ b/cmds/statsd/src/external/puller_util.cpp
@@ -55,7 +55,7 @@
  */
 void mapAndMergeIsolatedUidsToHostUid(vector<shared_ptr<LogEvent>>& data, const sp<UidMap>& uidMap,
                                       int tagId) {
-    if (StatsPullerManager::kAllPullAtomInfo.find(tagId) ==
+    if (StatsPullerManager::kAllPullAtomInfo.find({.atomTag = tagId}) ==
         StatsPullerManager::kAllPullAtomInfo.end()) {
         VLOG("Unknown pull atom id %d", tagId);
         return;
@@ -121,7 +121,7 @@
 
     vector<shared_ptr<LogEvent>> mergedData;
     const vector<int>& additiveFieldsVec =
-            StatsPullerManager::kAllPullAtomInfo.find(tagId)->second.additiveFields;
+            StatsPullerManager::kAllPullAtomInfo.find({.atomTag = tagId})->second.additiveFields;
     const set<int> additiveFields(additiveFieldsVec.begin(), additiveFieldsVec.end());
     bool needMerge = true;
 
diff --git a/core/java/android/app/StatsManager.java b/core/java/android/app/StatsManager.java
index 92bfee2..38ae5da 100644
--- a/core/java/android/app/StatsManager.java
+++ b/core/java/android/app/StatsManager.java
@@ -517,7 +517,7 @@
         }
     }
 
-    private static class  PullAtomCallbackInternal extends IPullAtomCallback.Stub {
+    private static class PullAtomCallbackInternal extends IPullAtomCallback.Stub {
         public final int mAtomId;
         public final StatsPullAtomCallback mCallback;
         public final Executor mExecutor;