Merge changes from topic "tej-uid" into rvc-dev am: 0ac1c345a4 am: d4b44af54d

Change-Id: I75321a8cf5ad8409dc98385a8e249924fd54babc

apex/statsd/tests/libstatspull/src/com/android/internal/os/statsd/libstats/LibStatsPullTests.java

@@ -29,6 +29,7 @@ import androidx.test.runner.AndroidJUnit4;
 import com.android.internal.os.StatsdConfigProto.AtomMatcher;
 import com.android.internal.os.StatsdConfigProto.FieldFilter;
 import com.android.internal.os.StatsdConfigProto.GaugeMetric;
+import com.android.internal.os.StatsdConfigProto.PullAtomPackages;
 import com.android.internal.os.StatsdConfigProto.SimpleAtomMatcher;
 import com.android.internal.os.StatsdConfigProto.StatsdConfig;
 import com.android.internal.os.StatsdConfigProto.TimeUnit;
@@ -271,6 +272,9 @@ public class LibStatsPullTests {
                         .setSamplingType(GaugeMetric.SamplingType.FIRST_N_SAMPLES)
                         .setMaxNumGaugeAtomsPerBucket(1000)
                 )
+                .addPullAtomPackages(PullAtomPackages.newBuilder()
+                        .setAtomId(PULL_ATOM_TAG)
+                        .addPackages(LibStatsPullTests.class.getPackage().getName()))
                 .build();
         statsManager.addConfig(sConfigId, config.toByteArray());
         assertThat(StatsConfigUtils.verifyValidConfigExists(statsManager, sConfigId)).isTrue();

cmds/statsd/Android.bp

@@ -333,6 +333,7 @@ cc_test {
         "tests/external/puller_util_test.cpp",
         "tests/external/StatsCallbackPuller_test.cpp",
         "tests/external/StatsPuller_test.cpp",
+        "tests/external/StatsPullerManager_test.cpp",
         "tests/FieldValue_test.cpp",
         "tests/guardrail/StatsdStats_test.cpp",
         "tests/indexed_priority_queue_test.cpp",

cmds/statsd/src/StatsLogProcessor.cpp

@@ -536,6 +536,7 @@ void StatsLogProcessor::OnConfigUpdatedLocked(
             new MetricsManager(key, config, mTimeBaseNs, timestampNs, mUidMap, mPullerManager,
                                mAnomalyAlarmMonitor, mPeriodicAlarmMonitor);
     if (newMetricsManager->isConfigValid()) {
+        newMetricsManager->init();
         mUidMap->OnConfigUpdated(key);
         newMetricsManager->refreshTtl(timestampNs);
         mMetricsManagers[key] = newMetricsManager;

cmds/statsd/src/StatsService.cpp

@@ -385,9 +385,11 @@ void StatsService::print_cmd_help(int out) {
     dprintf(out, "  PKG           Optional package name to print the uids of the package\n");
     dprintf(out, "\n");
     dprintf(out, "\n");
-    dprintf(out, "usage: adb shell cmd stats pull-source [int] \n");
+    dprintf(out, "usage: adb shell cmd stats pull-source ATOM_TAG [PACKAGE] \n");
     dprintf(out, "\n");
-    dprintf(out, "  Prints the output of a pulled metrics source (int indicates source)\n");
+    dprintf(out, "  Prints the output of a pulled atom\n");
+    dprintf(out, "  UID           The atom to pull\n");
+    dprintf(out, "  PACKAGE       The package to pull from. Default is AID_SYSTEM\n");
     dprintf(out, "\n");
     dprintf(out, "\n");
     dprintf(out, "usage: adb shell cmd stats write-to-disk \n");
@@ -806,8 +808,21 @@ status_t StatsService::cmd_log_binary_push(int out, const Vector<String8>& args)
 
 status_t StatsService::cmd_print_pulled_metrics(int out, const Vector<String8>& args) {
     int s = atoi(args[1].c_str());
-    vector<shared_ptr<LogEvent> > stats;
+    vector<int32_t> uids;
-    if (mPullerManager->Pull(s, &stats)) {
+    if (args.size() > 2) {
+        string package = string(args[2].c_str());
+        auto it = UidMap::sAidToUidMapping.find(package);
+        if (it != UidMap::sAidToUidMapping.end()) {
+            uids.push_back(it->second);
+        } else {
+            set<int32_t> uids_set = mUidMap->getAppUid(package);
+            uids.insert(uids.end(), uids_set.begin(), uids_set.end());
+        }
+    } else {
+        uids.push_back(AID_SYSTEM);
+    }
+    vector<shared_ptr<LogEvent>> stats;
+    if (mPullerManager->Pull(s, uids, &stats)) {
         for (const auto& it : stats) {
             dprintf(out, "Pull from %d: %s\n", s, it->ToString().c_str());
         }

cmds/statsd/src/external/PullUidProvider.h

@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2020 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 <utils/RefBase.h>
+
+#include "StatsPuller.h"
+#include "logd/LogEvent.h"
+
+namespace android {
+namespace os {
+namespace statsd {
+
+class PullUidProvider : virtual public RefBase {
+public:
+    virtual ~PullUidProvider() {}
+
+    /**
+     * @param atomId The atom for which to get the uids.
+     */
+    virtual vector<int32_t> getPullAtomUids(int32_t atomId) = 0;
+};
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android

cmds/statsd/src/external/StatsPullerManager.cpp

@@ -47,27 +47,73 @@ const int64_t NO_ALARM_UPDATE = INT64_MAX;
 StatsPullerManager::StatsPullerManager()
     : kAllPullAtomInfo({
               // TrainInfo.
-              {{.atomTag = util::TRAIN_INFO}, new TrainInfoPuller()},
+              {{.atomTag = util::TRAIN_INFO, .uid = -1}, new TrainInfoPuller()},
       }),
       mNextPullTimeNs(NO_ALARM_UPDATE) {
 }
 
-bool StatsPullerManager::Pull(int tagId, vector<shared_ptr<LogEvent>>* data) {
+bool StatsPullerManager::Pull(int tagId, const ConfigKey& configKey,
-    AutoMutex _l(mLock);
+                              vector<shared_ptr<LogEvent>>* data, bool useUids) {
-    return PullLocked(tagId, data);
+    std::lock_guard<std::mutex> _l(mLock);
+    return PullLocked(tagId, configKey, data, useUids);
 }
 
-bool StatsPullerManager::PullLocked(int tagId, vector<shared_ptr<LogEvent>>* data) {
+bool StatsPullerManager::Pull(int tagId, const vector<int32_t>& uids,
-    VLOG("Initiating pulling %d", tagId);
+                              vector<std::shared_ptr<LogEvent>>* data, bool useUids) {
+    std::lock_guard<std::mutex> _l(mLock);
+    return PullLocked(tagId, uids, data, useUids);
+}
 
-    if (kAllPullAtomInfo.find({.atomTag = tagId}) != kAllPullAtomInfo.end()) {
+bool StatsPullerManager::PullLocked(int tagId, const ConfigKey& configKey,
-        bool ret = kAllPullAtomInfo.find({.atomTag = tagId})->second->Pull(data);
+                                    vector<shared_ptr<LogEvent>>* data, bool useUids) {
-        VLOG("pulled %d items", (int)data->size());
+    vector<int32_t> uids;
-        if (!ret) {
+    if (useUids) {
-            StatsdStats::getInstance().notePullFailed(tagId);
+        auto uidProviderIt = mPullUidProviders.find(configKey);
+        if (uidProviderIt == mPullUidProviders.end()) {
+            ALOGE("Error pulling tag %d. No pull uid provider for config key %s", tagId,
+                  configKey.ToString().c_str());
+            return false;
         }
-        return ret;
+        sp<PullUidProvider> pullUidProvider = uidProviderIt->second.promote();
+        if (pullUidProvider == nullptr) {
+            ALOGE("Error pulling tag %d, pull uid provider for config %s is gone.", tagId,
+                  configKey.ToString().c_str());
+            return false;
+        }
+        uids = pullUidProvider->getPullAtomUids(tagId);
+    }
+    return PullLocked(tagId, uids, data, useUids);
+}
+
+bool StatsPullerManager::PullLocked(int tagId, const vector<int32_t>& uids,
+                                    vector<shared_ptr<LogEvent>>* data, bool useUids) {
+    VLOG("Initiating pulling %d", tagId);
+    if (useUids) {
+        for (int32_t uid : uids) {
+            PullerKey key = {.atomTag = tagId, .uid = uid};
+            auto pullerIt = kAllPullAtomInfo.find(key);
+            if (pullerIt != kAllPullAtomInfo.end()) {
+                bool ret = pullerIt->second->Pull(data);
+                VLOG("pulled %zu items", data->size());
+                if (!ret) {
+                    StatsdStats::getInstance().notePullFailed(tagId);
+                }
+                return ret;
+            }
+        }
+        ALOGW("StatsPullerManager: Unknown tagId %d", tagId);
+        return false;  // Return early since we don't know what to pull.
     } else {
+        PullerKey key = {.atomTag = tagId, .uid = -1};
+        auto pullerIt = kAllPullAtomInfo.find(key);
+        if (pullerIt != kAllPullAtomInfo.end()) {
+            bool ret = pullerIt->second->Pull(data);
+            VLOG("pulled %zu items", data->size());
+            if (!ret) {
+                StatsdStats::getInstance().notePullFailed(tagId);
+            }
+            return ret;
+        }
         ALOGW("StatsPullerManager: Unknown tagId %d", tagId);
         return false;  // Return early since we don't know what to pull.
     }
@@ -96,7 +142,7 @@ void StatsPullerManager::updateAlarmLocked() {
 
 void StatsPullerManager::SetStatsCompanionService(
         shared_ptr<IStatsCompanionService> statsCompanionService) {
-    AutoMutex _l(mLock);
+    std::lock_guard<std::mutex> _l(mLock);
     shared_ptr<IStatsCompanionService> tmpForLock = mStatsCompanionService;
     mStatsCompanionService = statsCompanionService;
     for (const auto& pulledAtom : kAllPullAtomInfo) {
@@ -107,10 +153,11 @@ void StatsPullerManager::SetStatsCompanionService(
     }
 }
 
-void StatsPullerManager::RegisterReceiver(int tagId, wp<PullDataReceiver> receiver,
+void StatsPullerManager::RegisterReceiver(int tagId, const ConfigKey& configKey,
-                                              int64_t nextPullTimeNs, int64_t intervalNs) {
+                                          wp<PullDataReceiver> receiver, int64_t nextPullTimeNs,
-    AutoMutex _l(mLock);
+                                          int64_t intervalNs) {
-    auto& receivers = mReceivers[tagId];
+    std::lock_guard<std::mutex> _l(mLock);
+    auto& receivers = mReceivers[{.atomTag = tagId, .configKey = configKey}];
     for (auto it = receivers.begin(); it != receivers.end(); it++) {
         if (it->receiver == receiver) {
             VLOG("Receiver already registered of %d", (int)receivers.size());
@@ -142,13 +189,15 @@ void StatsPullerManager::RegisterReceiver(int tagId, wp<PullDataReceiver> receiv
     VLOG("Puller for tagId %d registered of %d", tagId, (int)receivers.size());
 }
 
-void StatsPullerManager::UnRegisterReceiver(int tagId, wp<PullDataReceiver> receiver) {
+void StatsPullerManager::UnRegisterReceiver(int tagId, const ConfigKey& configKey,
-    AutoMutex _l(mLock);
+                                            wp<PullDataReceiver> receiver) {
-    if (mReceivers.find(tagId) == mReceivers.end()) {
+    std::lock_guard<std::mutex> _l(mLock);
+    auto receiversIt = mReceivers.find({.atomTag = tagId, .configKey = configKey});
+    if (receiversIt == mReceivers.end()) {
         VLOG("Unknown pull code or no receivers: %d", tagId);
         return;
     }
-    auto& receivers = mReceivers.find(tagId)->second;
+    std::list<ReceiverInfo>& receivers = receiversIt->second;
     for (auto it = receivers.begin(); it != receivers.end(); it++) {
         if (receiver == it->receiver) {
             receivers.erase(it);
@@ -158,16 +207,26 @@ void StatsPullerManager::UnRegisterReceiver(int tagId, wp<PullDataReceiver> rece
     }
 }
 
+void StatsPullerManager::RegisterPullUidProvider(const ConfigKey& configKey,
+                                                 wp<PullUidProvider> provider) {
+    std::lock_guard<std::mutex> _l(mLock);
+    mPullUidProviders[configKey] = provider;
+}
+
+void StatsPullerManager::UnregisterPullUidProvider(const ConfigKey& configKey) {
+    std::lock_guard<std::mutex> _l(mLock);
+    mPullUidProviders.erase(configKey);
+}
+
 void StatsPullerManager::OnAlarmFired(int64_t elapsedTimeNs) {
-    AutoMutex _l(mLock);
+    std::lock_guard<std::mutex> _l(mLock);
     int64_t wallClockNs = getWallClockNs();
 
     int64_t minNextPullTimeNs = NO_ALARM_UPDATE;
 
-    vector<pair<int, vector<ReceiverInfo*>>> needToPull =
+    vector<pair<const ReceiverKey*, vector<ReceiverInfo*>>> needToPull;
-            vector<pair<int, vector<ReceiverInfo*>>>();
     for (auto& pair : mReceivers) {
-        vector<ReceiverInfo*> receivers = vector<ReceiverInfo*>();
+        vector<ReceiverInfo*> receivers;
         if (pair.second.size() != 0) {
             for (ReceiverInfo& receiverInfo : pair.second) {
                 if (receiverInfo.nextPullTimeNs <= elapsedTimeNs) {
@@ -179,17 +238,16 @@ void StatsPullerManager::OnAlarmFired(int64_t elapsedTimeNs) {
                 }
             }
             if (receivers.size() > 0) {
-                needToPull.push_back(make_pair(pair.first, receivers));
+                needToPull.push_back(make_pair(&pair.first, receivers));
             }
         }
     }
-
     for (const auto& pullInfo : needToPull) {
         vector<shared_ptr<LogEvent>> data;
-        bool pullSuccess = PullLocked(pullInfo.first, &data);
+        bool pullSuccess = PullLocked(pullInfo.first->atomTag, pullInfo.first->configKey, &data);
         if (pullSuccess) {
-            StatsdStats::getInstance().notePullDelay(
+            StatsdStats::getInstance().notePullDelay(pullInfo.first->atomTag,
-                    pullInfo.first, getElapsedRealtimeNs() - elapsedTimeNs);
+                                                     getElapsedRealtimeNs() - elapsedTimeNs);
         } else {
             VLOG("pull failed at %lld, will try again later", (long long)elapsedTimeNs);
         }
@@ -248,18 +306,19 @@ int StatsPullerManager::ClearPullerCacheIfNecessary(int64_t timestampNs) {
 void StatsPullerManager::RegisterPullAtomCallback(const int uid, const int32_t atomTag,
                                                   const int64_t coolDownNs, const int64_t timeoutNs,
                                                   const vector<int32_t>& additiveFields,
-                                                  const shared_ptr<IPullAtomCallback>& callback) {
+                                                  const shared_ptr<IPullAtomCallback>& callback,
-    AutoMutex _l(mLock);
+                                                  bool useUid) {
+    std::lock_guard<std::mutex> _l(mLock);
     VLOG("RegisterPullerCallback: adding puller for tag %d", atomTag);
     // TODO(b/146439412): linkToDeath with the callback so that we can remove it
     // and delete the puller.
     StatsdStats::getInstance().notePullerCallbackRegistrationChanged(atomTag, /*registered=*/true);
-    kAllPullAtomInfo[{.atomTag = atomTag}] =
+    kAllPullAtomInfo[{.atomTag = atomTag, .uid = useUid ? uid : -1}] =
             new StatsCallbackPuller(atomTag, callback, coolDownNs, timeoutNs, additiveFields);
 }
 
 void StatsPullerManager::UnregisterPullAtomCallback(const int uid, const int32_t atomTag) {
-    AutoMutex _l(mLock);
+    std::lock_guard<std::mutex> _l(mLock);
     StatsdStats::getInstance().notePullerCallbackRegistrationChanged(atomTag, /*registered=*/false);
     kAllPullAtomInfo.erase({.atomTag = atomTag});
 }

cmds/statsd/src/external/StatsPullerManager.h

@@ -19,15 +19,16 @@
 #include <aidl/android/os/IPullAtomCallback.h>
 #include <aidl/android/os/IStatsCompanionService.h>
 #include <utils/RefBase.h>
-#include <utils/threads.h>
 
 #include <list>
 #include <vector>
 
 #include "PullDataReceiver.h"
+#include "PullUidProvider.h"
 #include "StatsPuller.h"
 #include "guardrail/StatsdStats.h"
 #include "logd/LogEvent.h"
+#include "packages/UidMap.h"
 
 using aidl::android::os::IPullAtomCallback;
 using aidl::android::os::IStatsCompanionService;
@@ -67,11 +68,20 @@ public:
 
     // Registers a receiver for tagId. It will be pulled on the nextPullTimeNs
     // and then every intervalNs thereafter.
-    virtual void RegisterReceiver(int tagId, wp<PullDataReceiver> receiver, int64_t nextPullTimeNs,
+    virtual void RegisterReceiver(int tagId, const ConfigKey& configKey,
+                                  wp<PullDataReceiver> receiver, int64_t nextPullTimeNs,
                                   int64_t intervalNs);
 
     // Stop listening on a tagId.
-    virtual void UnRegisterReceiver(int tagId, wp<PullDataReceiver> receiver);
+    virtual void UnRegisterReceiver(int tagId, const ConfigKey& configKey,
+                                    wp<PullDataReceiver> receiver);
+
+    // Registers a pull uid provider for the config key. When pulling atoms, it will be used to
+    // determine which atoms to pull from.
+    virtual void RegisterPullUidProvider(const ConfigKey& configKey, wp<PullUidProvider> provider);
+
+    // Unregister a pull uid provider.
+    virtual void UnregisterPullUidProvider(const ConfigKey& configKey);
 
     // Verify if we know how to pull for this matcher
     bool PullerForMatcherExists(int tagId) const;
@@ -85,9 +95,16 @@ public:
     // Returns false when
     //   1) the pull fails
     //   2) pull takes longer than mPullTimeoutNs (intrinsic to puller)
+    //   3) Either a PullUidProvider was not registered for the config, or the there was no puller
+    //      registered for any of the uids for this atom.
     // If the metric wants to make any change to the data, like timestamps, they
     // should make a copy as this data may be shared with multiple metrics.
-    virtual bool Pull(int tagId, vector<std::shared_ptr<LogEvent>>* data);
+    virtual bool Pull(int tagId, const ConfigKey& configKey,
+                      vector<std::shared_ptr<LogEvent>>* data, bool useUids = false);
+
+    // Same as above, but directly specify the allowed uids to pull from.
+    virtual bool Pull(int tagId, const vector<int32_t>& uids,
+                      vector<std::shared_ptr<LogEvent>>* data, bool useUids = false);
 
     // Clear pull data cache immediately.
     int ForceClearPullerCache();
@@ -99,7 +116,8 @@ public:
 
     void RegisterPullAtomCallback(const int uid, const int32_t atomTag, const int64_t coolDownNs,
                                   const int64_t timeoutNs, const vector<int32_t>& additiveFields,
-                                  const shared_ptr<IPullAtomCallback>& callback);
+                                  const shared_ptr<IPullAtomCallback>& callback,
+                                  bool useUid = false);
 
     void UnregisterPullAtomCallback(const int uid, const int32_t atomTag);
 
@@ -108,19 +126,36 @@ public:
 private:
     shared_ptr<IStatsCompanionService> mStatsCompanionService = nullptr;
 
+    // A struct containing an atom id and a Config Key
+    typedef struct ReceiverKey {
+        const int atomTag;
+        const ConfigKey configKey;
+
+        inline bool operator<(const ReceiverKey& that) const {
+            return atomTag == that.atomTag ? configKey < that.configKey : atomTag < that.atomTag;
+        }
+    } ReceiverKey;
+
     typedef struct {
         int64_t nextPullTimeNs;
         int64_t intervalNs;
         wp<PullDataReceiver> receiver;
     } ReceiverInfo;
 
-    // mapping from simple matcher tagId to receivers
+    // mapping from Receiver Key to receivers
-    std::map<int, std::list<ReceiverInfo>> mReceivers;
+    std::map<ReceiverKey, std::list<ReceiverInfo>> mReceivers;
 
-    bool PullLocked(int tagId, vector<std::shared_ptr<LogEvent>>* data);
+    // mapping from Config Key to the PullUidProvider for that config
+    std::map<ConfigKey, wp<PullUidProvider>> mPullUidProviders;
+
+    bool PullLocked(int tagId, const ConfigKey& configKey, vector<std::shared_ptr<LogEvent>>* data,
+                    bool useUids = false);
+
+    bool PullLocked(int tagId, const vector<int32_t>& uids, vector<std::shared_ptr<LogEvent>>* data,
+                    bool useUids);
 
     // locks for data receiver and StatsCompanionService changes
-    Mutex mLock;
+    std::mutex mLock;
 
     void updateAlarmLocked();
 

cmds/statsd/src/guardrail/StatsdStats.h

@@ -119,6 +119,8 @@ public:
 
     const static int kMaxLogSourceCount = 50;
 
+    const static int kMaxPullAtomPackages = 100;
+
     // Max memory allowed for storing metrics per configuration. If this limit is exceeded, statsd
     // drops the metrics data in memory.
     static const size_t kMaxMetricsBytesPerConfig = 2 * 1024 * 1024;

cmds/statsd/src/metrics/GaugeMetricProducer.cpp

@@ -75,11 +75,9 @@ GaugeMetricProducer::GaugeMetricProducer(
         const int atomId, const int64_t timeBaseNs, const int64_t startTimeNs,
         const sp<StatsPullerManager>& pullerManager,
         const unordered_map<int, shared_ptr<Activation>>& eventActivationMap,
-        const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap,
+        const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap)
-        const vector<int>& slicedStateAtoms,
-        const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap)
     : MetricProducer(metric.id(), key, timeBaseNs, conditionIndex, wizard, eventActivationMap,
-                     eventDeactivationMap, slicedStateAtoms, stateGroupMap),
+                     eventDeactivationMap, /*slicedStateAtoms=*/{}, /*stateGroupMap=*/{}),
       mWhatMatcherIndex(whatMatcherIndex),
       mEventMatcherWizard(matcherWizard),
       mPullerManager(pullerManager),
@@ -135,15 +133,12 @@ GaugeMetricProducer::GaugeMetricProducer(
     flushIfNeededLocked(startTimeNs);
     // Kicks off the puller immediately.
     if (mIsPulled && mSamplingType == GaugeMetric::RANDOM_ONE_SAMPLE) {
-        mPullerManager->RegisterReceiver(mPullTagId, this, getCurrentBucketEndTimeNs(),
+        mPullerManager->RegisterReceiver(mPullTagId, mConfigKey, this, getCurrentBucketEndTimeNs(),
                                          mBucketSizeNs);
     }
 
     // Adjust start for partial first bucket and then pull if needed
     mCurrentBucketStartTimeNs = startTimeNs;
-    if (mIsActive && mIsPulled && mSamplingType == GaugeMetric::RANDOM_ONE_SAMPLE) {
-        pullAndMatchEventsLocked(mCurrentBucketStartTimeNs);
-    }
 
     VLOG("Gauge metric %lld created. bucket size %lld start_time: %lld sliced %d",
          (long long)metric.id(), (long long)mBucketSizeNs, (long long)mTimeBaseNs,
@@ -153,7 +148,7 @@ GaugeMetricProducer::GaugeMetricProducer(
 GaugeMetricProducer::~GaugeMetricProducer() {
     VLOG("~GaugeMetricProducer() called");
     if (mIsPulled && mSamplingType == GaugeMetric::RANDOM_ONE_SAMPLE) {
-        mPullerManager->UnRegisterReceiver(mPullTagId, this);
+        mPullerManager->UnRegisterReceiver(mPullTagId, mConfigKey, this);
     }
 }
 
@@ -298,6 +293,11 @@ void GaugeMetricProducer::onDumpReportLocked(const int64_t dumpTimeNs,
     }
 }
 
+void GaugeMetricProducer::prepareFirstBucketLocked() {
+    if (mIsActive && mIsPulled && mSamplingType == GaugeMetric::RANDOM_ONE_SAMPLE) {
+        pullAndMatchEventsLocked(mCurrentBucketStartTimeNs);
+    }
+}
 
 void GaugeMetricProducer::pullAndMatchEventsLocked(const int64_t timestampNs) {
     bool triggerPuller = false;
@@ -323,7 +323,7 @@ void GaugeMetricProducer::pullAndMatchEventsLocked(const int64_t timestampNs) {
         return;
     }
     vector<std::shared_ptr<LogEvent>> allData;
-    if (!mPullerManager->Pull(mPullTagId, &allData)) {
+    if (!mPullerManager->Pull(mPullTagId, mConfigKey, &allData)) {
         ALOGE("Gauge Stats puller failed for tag: %d at %lld", mPullTagId, (long long)timestampNs);
         return;
     }

cmds/statsd/src/metrics/GaugeMetricProducer.h

@@ -64,9 +64,7 @@ public:
             const int64_t startTimeNs, const sp<StatsPullerManager>& pullerManager,
             const std::unordered_map<int, std::shared_ptr<Activation>>& eventActivationMap = {},
             const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
-                    eventDeactivationMap = {},
+                    eventDeactivationMap = {});
-            const vector<int>& slicedStateAtoms = {},
-            const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap = {});
 
     virtual ~GaugeMetricProducer();
 
@@ -129,6 +127,8 @@ private:
     void flushCurrentBucketLocked(const int64_t& eventTimeNs,
                                   const int64_t& nextBucketStartTimeNs) override;
 
+    void prepareFirstBucketLocked() override;
+
     void pullAndMatchEventsLocked(const int64_t timestampNs);
 
     const int mWhatMatcherIndex;

cmds/statsd/src/metrics/MetricProducer.h

@@ -207,6 +207,11 @@ public:
         return clearPastBucketsLocked(dumpTimeNs);
     }
 
+    void prepareFirstBucket() {
+        std::lock_guard<std::mutex> lock(mMutex);
+        prepareFirstBucketLocked();
+    }
+
     // Returns the memory in bytes currently used to store this metric's data. Does not change
     // state.
     size_t byteSize() const {
@@ -344,6 +349,7 @@ protected:
                                     std::set<string> *str_set,
                                     android::util::ProtoOutputStream* protoOutput) = 0;
     virtual void clearPastBucketsLocked(const int64_t dumpTimeNs) = 0;
+    virtual void prepareFirstBucketLocked(){};
     virtual size_t byteSizeLocked() const = 0;
     virtual void dumpStatesLocked(FILE* out, bool verbose) const = 0;
     virtual void dropDataLocked(const int64_t dropTimeNs) = 0;

cmds/statsd/src/metrics/MetricsManager.cpp

@@ -70,6 +70,7 @@ MetricsManager::MetricsManager(const ConfigKey& key, const StatsdConfig& config,
       mTtlEndNs(-1),
       mLastReportTimeNs(currentTimeNs),
       mLastReportWallClockNs(getWallClockNs()),
+      mPullerManager(pullerManager),
       mShouldPersistHistory(config.persist_locally()) {
     // Init the ttl end timestamp.
     refreshTtl(timeBaseNs);
@@ -86,6 +87,7 @@ MetricsManager::MetricsManager(const ConfigKey& key, const StatsdConfig& config,
     mVersionStringsInReport = config.version_strings_in_metric_report();
     mInstallerInReport = config.installer_in_metric_report();
 
+    // Init allowed pushed atom uids.
     if (config.allowed_log_source_size() == 0) {
         mConfigValid = false;
         ALOGE("Log source whitelist is empty! This config won't get any data. Suggest adding at "
@@ -108,6 +110,40 @@ MetricsManager::MetricsManager(const ConfigKey& key, const StatsdConfig& config,
         }
     }
 
+    // Init default allowed pull atom uids.
+    int numPullPackages = 0;
+    for (const string& pullSource : config.default_pull_packages()) {
+        auto it = UidMap::sAidToUidMapping.find(pullSource);
+        if (it != UidMap::sAidToUidMapping.end()) {
+            numPullPackages++;
+            mDefaultPullUids.insert(it->second);
+        } else {
+            ALOGE("Default pull atom packages must be in sAidToUidMapping");
+            mConfigValid = false;
+        }
+    }
+    // Init per-atom pull atom packages.
+    for (const PullAtomPackages& pullAtomPackages : config.pull_atom_packages()) {
+        int32_t atomId = pullAtomPackages.atom_id();
+        for (const string& pullPackage : pullAtomPackages.packages()) {
+            numPullPackages++;
+            auto it = UidMap::sAidToUidMapping.find(pullPackage);
+            if (it != UidMap::sAidToUidMapping.end()) {
+                mPullAtomUids[atomId].insert(it->second);
+            } else {
+                mPullAtomPackages[atomId].insert(pullPackage);
+            }
+        }
+    }
+    if (numPullPackages > StatsdStats::kMaxPullAtomPackages) {
+        ALOGE("Too many sources in default_pull_packages and pull_atom_packages. This is likely to "
+              "be an error in the config");
+        mConfigValid = false;
+    } else {
+        initPullAtomSources();
+    }
+    mPullerManager->RegisterPullUidProvider(mConfigKey, this);
+
     // Store the sub-configs used.
     for (const auto& annotation : config.annotation()) {
         mAnnotations.emplace_back(annotation.field_int64(), annotation.field_int32());
@@ -153,6 +189,7 @@ MetricsManager::~MetricsManager() {
             StateManager::getInstance().unregisterListener(atomId, it);
         }
     }
+    mPullerManager->UnregisterPullUidProvider(mConfigKey);
 
     VLOG("~MetricsManager()");
 }
@@ -173,6 +210,20 @@ void MetricsManager::initLogSourceWhiteList() {
     }
 }
 
+void MetricsManager::initPullAtomSources() {
+    std::lock_guard<std::mutex> lock(mAllowedLogSourcesMutex);
+    mCombinedPullAtomUids.clear();
+    for (const auto& [atomId, uids] : mPullAtomUids) {
+        mCombinedPullAtomUids[atomId].insert(uids.begin(), uids.end());
+    }
+    for (const auto& [atomId, packages] : mPullAtomPackages) {
+        for (const string& pkg : packages) {
+            set<int32_t> uids = mUidMap->getAppUid(pkg);
+            mCombinedPullAtomUids[atomId].insert(uids.begin(), uids.end());
+        }
+    }
+}
+
 bool MetricsManager::isConfigValid() const {
     return mConfigValid;
 }
@@ -184,12 +235,18 @@ void MetricsManager::notifyAppUpgrade(const int64_t& eventTimeNs, const string&
         it->notifyAppUpgrade(eventTimeNs, apk, uid, version);
     }
     // check if we care this package
-    if (std::find(mAllowedPkg.begin(), mAllowedPkg.end(), apk) == mAllowedPkg.end()) {
+    if (std::find(mAllowedPkg.begin(), mAllowedPkg.end(), apk) != mAllowedPkg.end()) {
-        return;
+        // We will re-initialize the whole list because we don't want to keep the multi mapping of
+        // UID<->pkg inside MetricsManager to reduce the memory usage.
+        initLogSourceWhiteList();
+    }
+
+    for (const auto& it : mPullAtomPackages) {
+        if (it.second.find(apk) != it.second.end()) {
+            initPullAtomSources();
+            return;
+        }
     }
-    // We will re-initialize the whole list because we don't want to keep the multi mapping of
-    // UID<->pkg inside MetricsManager to reduce the memory usage.
-    initLogSourceWhiteList();
 }
 
 void MetricsManager::notifyAppRemoved(const int64_t& eventTimeNs, const string& apk,
@@ -199,24 +256,49 @@ void MetricsManager::notifyAppRemoved(const int64_t& eventTimeNs, const string&
         it->notifyAppRemoved(eventTimeNs, apk, uid);
     }
     // check if we care this package
-    if (std::find(mAllowedPkg.begin(), mAllowedPkg.end(), apk) == mAllowedPkg.end()) {
+    if (std::find(mAllowedPkg.begin(), mAllowedPkg.end(), apk) != mAllowedPkg.end()) {
-        return;
+        // We will re-initialize the whole list because we don't want to keep the multi mapping of
+        // UID<->pkg inside MetricsManager to reduce the memory usage.
+        initLogSourceWhiteList();
+    }
+
+    for (const auto& it : mPullAtomPackages) {
+        if (it.second.find(apk) != it.second.end()) {
+            initPullAtomSources();
+            return;
+        }
     }
-    // We will re-initialize the whole list because we don't want to keep the multi mapping of
-    // UID<->pkg inside MetricsManager to reduce the memory usage.
-    initLogSourceWhiteList();
 }
 
 void MetricsManager::onUidMapReceived(const int64_t& eventTimeNs) {
     // Purposefully don't inform metric producers on a new snapshot
     // because we don't need to flush partial buckets.
     // This occurs if a new user is added/removed or statsd crashes.
+    initPullAtomSources();
+
     if (mAllowedPkg.size() == 0) {
         return;
     }
     initLogSourceWhiteList();
 }
 
+void MetricsManager::init() {
+    for (const auto& producer : mAllMetricProducers) {
+        producer->prepareFirstBucket();
+    }
+}
+
+vector<int32_t> MetricsManager::getPullAtomUids(int32_t atomId) {
+    std::lock_guard<std::mutex> lock(mAllowedLogSourcesMutex);
+    vector<int32_t> uids;
+    const auto& it = mCombinedPullAtomUids.find(atomId);
+    if (it != mCombinedPullAtomUids.end()) {
+        uids.insert(uids.end(), it->second.begin(), it->second.end());
+    }
+    uids.insert(uids.end(), mDefaultPullUids.begin(), mDefaultPullUids.end());
+    return uids;
+}
+
 void MetricsManager::dumpStates(FILE* out, bool verbose) {
     fprintf(out, "ConfigKey %s, allowed source:", mConfigKey.ToString().c_str());
     {

cmds/statsd/src/metrics/MetricsManager.h

@@ -35,7 +35,7 @@ namespace os {
 namespace statsd {
 
 // A MetricsManager is responsible for managing metrics from one single config source.
-class MetricsManager : public virtual android::RefBase {
+class MetricsManager : public virtual android::RefBase, public virtual PullUidProvider {
 public:
     MetricsManager(const ConfigKey& configKey, const StatsdConfig& config, const int64_t timeBaseNs,
                    const int64_t currentTimeNs, const sp<UidMap>& uidMap,
@@ -69,6 +69,10 @@ public:
 
     void onUidMapReceived(const int64_t& eventTimeNs);
 
+    void init();
+
+    vector<int32_t> getPullAtomUids(int32_t atomId) override;
+
     bool shouldWriteToDisk() const {
         return mNoReportMetricIds.size() != mAllMetricProducers.size();
     }
@@ -159,6 +163,8 @@ private:
     int64_t mLastReportTimeNs;
     int64_t mLastReportWallClockNs;
 
+    sp<StatsPullerManager> mPullerManager;
+
     // The uid log sources from StatsdConfig.
     std::vector<int32_t> mAllowedUid;
 
@@ -169,13 +175,27 @@ private:
     // Logs from uids that are not in the list will be ignored to avoid spamming.
     std::set<int32_t> mAllowedLogSources;
 
+    // To guard access to mAllowedLogSources
+    mutable std::mutex mAllowedLogSourcesMutex;
+
+    // We can pull any atom from these uids.
+    std::set<int32_t> mDefaultPullUids;
+
+    // Uids that specific atoms can pull from.
+    // This is a map<atom id, set<uids>>
+    std::map<int32_t, std::set<int32_t>> mPullAtomUids;
+
+    // Packages that specific atoms can be pulled from.
+    std::map<int32_t, std::set<std::string>> mPullAtomPackages;
+
+    // All uids to pull for this atom. NOTE: Does not include the default uids for memory.
+    std::map<int32_t, std::set<int32_t>> mCombinedPullAtomUids;
+
     // Contains the annotations passed in with StatsdConfig.
     std::list<std::pair<const int64_t, const int32_t>> mAnnotations;
 
     const bool mShouldPersistHistory;
 
-    // To guard access to mAllowedLogSources
-    mutable std::mutex mAllowedLogSourcesMutex;
 
     // All event tags that are interesting to my metrics.
     std::set<int> mTagIds;
@@ -238,6 +258,8 @@ private:
 
     void initLogSourceWhiteList();
 
+    void initPullAtomSources();
+
     // The metrics that don't need to be uploaded or even reported.
     std::set<int64_t> mNoReportMetricIds;
 
@@ -275,6 +297,8 @@ private:
     FRIEND_TEST(MetricActivationE2eTest, TestCountMetricWithSameDeactivation);
     FRIEND_TEST(MetricActivationE2eTest, TestCountMetricWithTwoMetricsTwoDeactivations);
 
+    FRIEND_TEST(MetricsManagerTest, TestLogSources);
+
     FRIEND_TEST(StatsLogProcessorTest, TestActiveConfigMetricDiskWriteRead);
     FRIEND_TEST(StatsLogProcessorTest, TestActivationOnBoot);
     FRIEND_TEST(StatsLogProcessorTest, TestActivationOnBootMultipleActivations);

cmds/statsd/src/metrics/ValueMetricProducer.cpp

@@ -156,7 +156,7 @@ ValueMetricProducer::ValueMetricProducer(
     flushIfNeededLocked(startTimeNs);
 
     if (mIsPulled) {
-        mPullerManager->RegisterReceiver(mPullTagId, this, getCurrentBucketEndTimeNs(),
+        mPullerManager->RegisterReceiver(mPullTagId, mConfigKey, this, getCurrentBucketEndTimeNs(),
                                          mBucketSizeNs);
     }
 
@@ -166,10 +166,6 @@ ValueMetricProducer::ValueMetricProducer(
     mCurrentBucketStartTimeNs = startTimeNs;
     mConditionTimer.newBucketStart(mCurrentBucketStartTimeNs);
 
-     // Kicks off the puller immediately if condition is true and diff based.
-    if (mIsActive && mIsPulled && mCondition == ConditionState::kTrue && mUseDiff) {
-        pullAndMatchEventsLocked(mCurrentBucketStartTimeNs);
-    }
     // Now that activations are processed, start the condition timer if needed.
     mConditionTimer.onConditionChanged(mIsActive && mCondition == ConditionState::kTrue,
                                        mCurrentBucketStartTimeNs);
@@ -181,7 +177,7 @@ ValueMetricProducer::ValueMetricProducer(
 ValueMetricProducer::~ValueMetricProducer() {
     VLOG("~ValueMetricProducer() called");
     if (mIsPulled) {
-        mPullerManager->UnRegisterReceiver(mPullTagId, this);
+        mPullerManager->UnRegisterReceiver(mPullTagId, mConfigKey, this);
     }
 }
 
@@ -503,9 +499,16 @@ void ValueMetricProducer::onConditionChangedLocked(const bool condition,
     mConditionTimer.onConditionChanged(mCondition, eventTimeNs);
 }
 
+void ValueMetricProducer::prepareFirstBucketLocked() {
+    // Kicks off the puller immediately if condition is true and diff based.
+    if (mIsActive && mIsPulled && mCondition == ConditionState::kTrue && mUseDiff) {
+        pullAndMatchEventsLocked(mCurrentBucketStartTimeNs);
+    }
+}
+
 void ValueMetricProducer::pullAndMatchEventsLocked(const int64_t timestampNs) {
     vector<std::shared_ptr<LogEvent>> allData;
-    if (!mPullerManager->Pull(mPullTagId, &allData)) {
+    if (!mPullerManager->Pull(mPullTagId, mConfigKey, &allData)) {
         ALOGE("Stats puller failed for tag: %d at %lld", mPullTagId, (long long)timestampNs);
         invalidateCurrentBucket(timestampNs, BucketDropReason::PULL_FAILED);
         return;

cmds/statsd/src/metrics/ValueMetricProducer.h

@@ -122,6 +122,8 @@ private:
     void flushCurrentBucketLocked(const int64_t& eventTimeNs,
                                   const int64_t& nextBucketStartTimeNs) override;
 
+    void prepareFirstBucketLocked() override;
+
     void dropDataLocked(const int64_t dropTimeNs) override;
 
     // Calculate previous bucket end time based on current time.

cmds/statsd/src/metrics/metrics_manager_util.cpp

@@ -803,9 +803,8 @@ bool initMetrics(const ConfigKey& key, const StatsdConfig& config, const int64_t
         if (!success) return false;
 
         sp<MetricProducer> gaugeProducer = new GaugeMetricProducer(
-                key, metric, conditionIndex, wizard,
+                key, metric, conditionIndex, wizard, trackerIndex, matcherWizard, pullTagId,
-                trackerIndex, matcherWizard, pullTagId, triggerAtomId, atomTagId,
+                triggerAtomId, atomTagId, timeBaseTimeNs, currentTimeNs, pullerManager,
-                timeBaseTimeNs, currentTimeNs, pullerManager,
                 eventActivationMap, eventDeactivationMap);
         allMetricProducers.push_back(gaugeProducer);
     }
@@ -964,6 +963,7 @@ bool initStatsdConfig(const ConfigKey& key, const StatsdConfig& config, UidMap&
         ALOGE("initConditionTrackers failed");
         return false;
     }
+
     if (!initStates(config, stateAtomIdMap, allStateGroupMaps)) {
         ALOGE("initStates failed");
         return false;

cmds/statsd/src/packages/UidMap.cpp

@@ -545,7 +545,15 @@ const std::map<string, uint32_t> UidMap::sAidToUidMapping = {{"AID_ROOT", 0},
                                                              {"AID_LMKD", 1069},
                                                              {"AID_LLKD", 1070},
                                                              {"AID_IORAPD", 1071},
+                                                             {"AID_GPU_SERVICE", 1072},
                                                              {"AID_NETWORK_STACK", 1073},
+                                                             {"AID_GSID", 1074},
+                                                             {"AID_FSVERITY_CERT", 1075},
+                                                             {"AID_CREDSTORE", 1076},
+                                                             {"AID_EXTERNAL_STORAGE", 1077},
+                                                             {"AID_EXT_DATA_RW", 1078},
+                                                             {"AID_EXT_OBB_RW", 1079},
+                                                             {"AID_CONTEXT_HUB", 1080},
                                                              {"AID_SHELL", 2000},
                                                              {"AID_CACHE", 2001},
                                                              {"AID_DIAG", 2002}};

cmds/statsd/src/packages/UidMap.h

@@ -149,7 +149,7 @@ public:
     // Get currently cached value of memory used by UID map.
     size_t getBytesUsed() const;
 
-    std::set<int32_t> getAppUid(const string& package) const;
+    virtual std::set<int32_t> getAppUid(const string& package) const;
 
     // Write current PackageInfoSnapshot to ProtoOutputStream.
     // interestingUids: If not empty, only write the package info for these uids. If empty, write

cmds/statsd/src/shell/ShellSubscriber.cpp

@@ -108,7 +108,21 @@ bool ShellSubscriber::readConfig(shared_ptr<SubscriptionInfo> subscriptionInfo)
         if (minInterval < 0 || pulled.freq_millis() < minInterval) {
             minInterval = pulled.freq_millis();
         }
-        subscriptionInfo->mPulledInfo.emplace_back(pulled.matcher(), pulled.freq_millis());
+
+        vector<string> packages;
+        vector<int32_t> uids;
+        for (const string& pkg : pulled.packages()) {
+            auto it = UidMap::sAidToUidMapping.find(pkg);
+            if (it != UidMap::sAidToUidMapping.end()) {
+                uids.push_back(it->second);
+            } else {
+                packages.push_back(pkg);
+            }
+        }
+
+        subscriptionInfo->mPulledInfo.emplace_back(pulled.matcher(), pulled.freq_millis(), packages,
+                                                   uids);
+        VLOG("adding matcher for pulled atom %d", pulled.matcher().atom_id());
     }
     subscriptionInfo->mPullIntervalMin = minInterval;
 
@@ -127,7 +141,15 @@ void ShellSubscriber::startPull(int64_t myToken) {
         for (auto& pullInfo : mSubscriptionInfo->mPulledInfo) {
             if (pullInfo.mPrevPullElapsedRealtimeMs + pullInfo.mInterval < nowMillis) {
                 vector<std::shared_ptr<LogEvent>> data;
-                mPullerMgr->Pull(pullInfo.mPullerMatcher.atom_id(), &data);
+                vector<int32_t> uids;
+                uids.insert(uids.end(), pullInfo.mPullUids.begin(), pullInfo.mPullUids.end());
+                // This is slow. Consider storing the uids per app and listening to uidmap updates.
+                for (const string& pkg : pullInfo.mPullPackages) {
+                    set<int32_t> uidsForPkg = mUidMap->getAppUid(pkg);
+                    uids.insert(uids.end(), uidsForPkg.begin(), uidsForPkg.end());
+                }
+                uids.push_back(DEFAULT_PULL_UID);
+                mPullerMgr->Pull(pullInfo.mPullerMatcher.atom_id(), uids, &data);
                 VLOG("pulled %zu atoms with id %d", data.size(), pullInfo.mPullerMatcher.atom_id());
 
                 // TODO(b/150969574): Don't write to a pipe while holding a lock.

cmds/statsd/src/shell/ShellSubscriber.h

@@ -16,15 +16,17 @@
 
 #pragma once
 
-#include "logd/LogEvent.h"
-
 #include <android/util/ProtoOutputStream.h>
+#include <private/android_filesystem_config.h>
+
 #include <condition_variable>
 #include <mutex>
 #include <thread>
+
 #include "external/StatsPullerManager.h"
 #include "frameworks/base/cmds/statsd/src/shell/shell_config.pb.h"
 #include "frameworks/base/cmds/statsd/src/statsd_config.pb.h"
+#include "logd/LogEvent.h"
 #include "packages/UidMap.h"
 
 namespace android {
@@ -66,12 +68,19 @@ public:
 
 private:
     struct PullInfo {
-        PullInfo(const SimpleAtomMatcher& matcher, int64_t interval)
+        PullInfo(const SimpleAtomMatcher& matcher, int64_t interval,
-            : mPullerMatcher(matcher), mInterval(interval), mPrevPullElapsedRealtimeMs(0) {
+                 const std::vector<std::string>& packages, const std::vector<int32_t>& uids)
+            : mPullerMatcher(matcher),
+              mInterval(interval),
+              mPrevPullElapsedRealtimeMs(0),
+              mPullPackages(packages),
+              mPullUids(uids) {
         }
         SimpleAtomMatcher mPullerMatcher;
         int64_t mInterval;
         int64_t mPrevPullElapsedRealtimeMs;
+        std::vector<std::string> mPullPackages;
+        std::vector<int32_t> mPullUids;
     };
 
     struct SubscriptionInfo {
@@ -109,6 +118,8 @@ private:
     std::shared_ptr<SubscriptionInfo> mSubscriptionInfo = nullptr;
 
     int mToken = 0;
+
+    const int32_t DEFAULT_PULL_UID = AID_SYSTEM;
 };
 
 }  // namespace statsd

cmds/statsd/src/shell/shell_config.proto

@@ -28,6 +28,9 @@ message PulledAtomSubscription {
 
     /* gap between two pulls in milliseconds */
     optional int32 freq_millis = 2;
+
+    /* Packages that the pull is requested from */
+    repeated string packages = 3;
 }
 
 message ShellSubscription {

cmds/statsd/src/statsd_config.proto

@@ -428,6 +428,12 @@ message MetricActivation {
   repeated EventActivation event_activation = 2;
 }
 
+message PullAtomPackages {
+    optional int32 atom_id = 1;
+
+    repeated string packages = 2;
+}
+
 message StatsdConfig {
   optional int64 id = 1;
 
@@ -475,6 +481,10 @@ message StatsdConfig {
 
   repeated State state = 21;
 
+  repeated string default_pull_packages = 22;
+
+  repeated PullAtomPackages pull_atom_packages = 23;
+
   // Field number 1000 is reserved for later use.
   reserved 1000;
 }

cmds/statsd/tests/MetricsManager_test.cpp

@@ -13,7 +13,15 @@
 // limitations under the License.
 
 #include <gtest/gtest.h>
+#include <private/android_filesystem_config.h>
+#include <stdio.h>
 
+#include <set>
+#include <unordered_map>
+#include <vector>
+
+#include "frameworks/base/cmds/statsd/src/statsd_config.pb.h"
+#include "metrics/metrics_test_helper.h"
 #include "src/condition/ConditionTracker.h"
 #include "src/matchers/LogMatchingTracker.h"
 #include "src/metrics/CountMetricProducer.h"
@@ -23,22 +31,21 @@
 #include "src/metrics/metrics_manager_util.h"
 #include "statsd_test_util.h"
 
-#include "frameworks/base/cmds/statsd/src/statsd_config.pb.h"
+using namespace testing;
-
-#include <stdio.h>
-#include <set>
-#include <unordered_map>
-#include <vector>
-
-using namespace android::os::statsd;
 using android::sp;
+using android::os::statsd::Predicate;
+using std::map;
 using std::set;
 using std::unordered_map;
 using std::vector;
-using android::os::statsd::Predicate;
 
 #ifdef __ANDROID__
 
+namespace android {
+namespace os {
+namespace statsd {
+
+namespace {
 const ConfigKey kConfigKey(0, 12345);
 const long kAlertId = 3;
 
@@ -268,6 +275,11 @@ StatsdConfig buildCirclePredicates() {
     return config;
 }
 
+bool isSubset(const set<int32_t>& set1, const set<int32_t>& set2) {
+    return std::includes(set2.begin(), set2.end(), set1.begin(), set1.end());
+}
+}  // anonymous namespace
+
 TEST(MetricsManagerTest, TestGoodConfig) {
     UidMap uidMap;
     sp<StatsPullerManager> pullerManager = new StatsPullerManager();
@@ -488,6 +500,101 @@ TEST(MetricsManagerTest, testAlertWithUnknownMetric) {
                                   noReportMetricIds));
 }
 
+TEST(MetricsManagerTest, TestLogSources) {
+    string app1 = "app1";
+    set<int32_t> app1Uids = {1111, 11111};
+    string app2 = "app2";
+    set<int32_t> app2Uids = {2222};
+    string app3 = "app3";
+    set<int32_t> app3Uids = {3333, 1111};
+
+    map<string, set<int32_t>> pkgToUids;
+    pkgToUids[app1] = app1Uids;
+    pkgToUids[app2] = app2Uids;
+    pkgToUids[app3] = app3Uids;
+
+    int32_t atom1 = 10;
+    int32_t atom2 = 20;
+    int32_t atom3 = 30;
+    sp<MockUidMap> uidMap = new StrictMock<MockUidMap>();
+    EXPECT_CALL(*uidMap, getAppUid(_))
+            .Times(4)
+            .WillRepeatedly(Invoke([&pkgToUids](const string& pkg) {
+                const auto& it = pkgToUids.find(pkg);
+                if (it != pkgToUids.end()) {
+                    return it->second;
+                }
+                return set<int32_t>();
+            }));
+    sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
+    EXPECT_CALL(*pullerManager, RegisterPullUidProvider(kConfigKey, _)).Times(1);
+    EXPECT_CALL(*pullerManager, UnregisterPullUidProvider(kConfigKey)).Times(1);
+
+    sp<AlarmMonitor> anomalyAlarmMonitor;
+    sp<AlarmMonitor> periodicAlarmMonitor;
+
+    StatsdConfig config = buildGoodConfig();
+    config.add_allowed_log_source("AID_SYSTEM");
+    config.add_allowed_log_source(app1);
+    config.add_default_pull_packages("AID_SYSTEM");
+    config.add_default_pull_packages("AID_ROOT");
+
+    const set<int32_t> defaultPullUids = {AID_SYSTEM, AID_ROOT};
+
+    PullAtomPackages* pullAtomPackages = config.add_pull_atom_packages();
+    pullAtomPackages->set_atom_id(atom1);
+    pullAtomPackages->add_packages(app1);
+    pullAtomPackages->add_packages(app3);
+
+    pullAtomPackages = config.add_pull_atom_packages();
+    pullAtomPackages->set_atom_id(atom2);
+    pullAtomPackages->add_packages(app2);
+    pullAtomPackages->add_packages("AID_STATSD");
+
+    MetricsManager metricsManager(kConfigKey, config, timeBaseSec, timeBaseSec, uidMap,
+                                  pullerManager, anomalyAlarmMonitor, periodicAlarmMonitor);
+
+    EXPECT_TRUE(metricsManager.isConfigValid());
+
+    ASSERT_EQ(metricsManager.mAllowedUid.size(), 1);
+    EXPECT_EQ(metricsManager.mAllowedUid[0], AID_SYSTEM);
+
+    ASSERT_EQ(metricsManager.mAllowedPkg.size(), 1);
+    EXPECT_EQ(metricsManager.mAllowedPkg[0], app1);
+
+    ASSERT_EQ(metricsManager.mAllowedLogSources.size(), 3);
+    EXPECT_TRUE(isSubset({AID_SYSTEM}, metricsManager.mAllowedLogSources));
+    EXPECT_TRUE(isSubset(app1Uids, metricsManager.mAllowedLogSources));
+
+    ASSERT_EQ(metricsManager.mDefaultPullUids.size(), 2);
+    EXPECT_TRUE(isSubset(defaultPullUids, metricsManager.mDefaultPullUids));
+    ;
+
+    vector<int32_t> atom1Uids = metricsManager.getPullAtomUids(atom1);
+    ASSERT_EQ(atom1Uids.size(), 5);
+    set<int32_t> expectedAtom1Uids;
+    expectedAtom1Uids.insert(defaultPullUids.begin(), defaultPullUids.end());
+    expectedAtom1Uids.insert(app1Uids.begin(), app1Uids.end());
+    expectedAtom1Uids.insert(app3Uids.begin(), app3Uids.end());
+    EXPECT_TRUE(isSubset(expectedAtom1Uids, set<int32_t>(atom1Uids.begin(), atom1Uids.end())));
+
+    vector<int32_t> atom2Uids = metricsManager.getPullAtomUids(atom2);
+    ASSERT_EQ(atom2Uids.size(), 4);
+    set<int32_t> expectedAtom2Uids;
+    expectedAtom1Uids.insert(defaultPullUids.begin(), defaultPullUids.end());
+    expectedAtom1Uids.insert(app2Uids.begin(), app2Uids.end());
+    expectedAtom1Uids.insert(AID_STATSD);
+    EXPECT_TRUE(isSubset(expectedAtom2Uids, set<int32_t>(atom2Uids.begin(), atom2Uids.end())));
+
+    vector<int32_t> atom3Uids = metricsManager.getPullAtomUids(atom3);
+    ASSERT_EQ(atom3Uids.size(), 2);
+    EXPECT_TRUE(isSubset(defaultPullUids, set<int32_t>(atom3Uids.begin(), atom3Uids.end())));
+}
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android
+
 #else
 GTEST_LOG_(INFO) << "This test does nothing.\n";
 #endif

cmds/statsd/tests/e2e/GaugeMetric_e2e_pull_test.cpp

@@ -38,6 +38,7 @@ StatsdConfig CreateStatsdConfig(const GaugeMetric::SamplingType sampling_type,
                                 bool useCondition = true) {
     StatsdConfig config;
     config.add_allowed_log_source("AID_ROOT"); // LogEvent defaults to UID of root.
+    config.add_default_pull_packages("AID_ROOT");  // Fake puller is registered with root.
     auto atomMatcher = CreateSimpleAtomMatcher("TestMatcher", ATOM_TAG);
     *config.add_atom_matcher() = atomMatcher;
     *config.add_atom_matcher() = CreateScreenTurnedOnAtomMatcher();

cmds/statsd/tests/e2e/PartialBucket_e2e_test.cpp

@@ -71,6 +71,7 @@ StatsdConfig MakeConfig() {
 StatsdConfig MakeValueMetricConfig(int64_t minTime) {
     StatsdConfig config;
     config.add_allowed_log_source("AID_ROOT");  // LogEvent defaults to UID of root.
+    config.add_default_pull_packages("AID_ROOT");  // Fake puller is registered with root.
 
     auto pulledAtomMatcher =
             CreateSimpleAtomMatcher("TestMatcher", util::SUBSYSTEM_SLEEP_STATE);
@@ -94,6 +95,7 @@ StatsdConfig MakeValueMetricConfig(int64_t minTime) {
 StatsdConfig MakeGaugeMetricConfig(int64_t minTime) {
     StatsdConfig config;
     config.add_allowed_log_source("AID_ROOT");  // LogEvent defaults to UID of root.
+    config.add_default_pull_packages("AID_ROOT");  // Fake puller is registered with root.
 
     auto pulledAtomMatcher =
                 CreateSimpleAtomMatcher("TestMatcher", util::SUBSYSTEM_SLEEP_STATE);

cmds/statsd/tests/e2e/ValueMetric_pull_e2e_test.cpp

@@ -36,6 +36,7 @@ const int64_t metricId = 123456;
 StatsdConfig CreateStatsdConfig(bool useCondition = true) {
     StatsdConfig config;
     config.add_allowed_log_source("AID_ROOT"); // LogEvent defaults to UID of root.
+    config.add_default_pull_packages("AID_ROOT");  // Fake puller is registered with root.
     auto pulledAtomMatcher =
             CreateSimpleAtomMatcher("TestMatcher", util::SUBSYSTEM_SLEEP_STATE);
     *config.add_atom_matcher() = pulledAtomMatcher;

cmds/statsd/tests/external/StatsCallbackPuller_test.cpp

@@ -187,15 +187,16 @@ TEST_F(StatsCallbackPullerTest, RegisterAndTimeout) {
     pullDelayNs = 500000000;  // 500 ms.
     pullTimeoutNs = 10000;    // 10 microsseconds.
     int64_t value = 4321;
+    int32_t uid = 123;
     values.push_back(value);
 
     StatsPullerManager pullerManager;
-    pullerManager.RegisterPullAtomCallback(/*uid=*/-1, pullTagId, pullCoolDownNs, pullTimeoutNs,
+    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, &dataHolder));
+    EXPECT_FALSE(pullerManager.Pull(pullTagId, {uid}, &dataHolder));
     int64_t endTimeNs = getElapsedRealtimeNs();
     int64_t actualPullDurationNs = endTimeNs - startTimeNs;
 

cmds/statsd/tests/external/StatsPullerManager_test.cpp

@@ -0,0 +1,150 @@
+// Copyright (C) 2020 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 "src/external/StatsPullerManager.h"
+
+#include <aidl/android/os/IPullAtomResultReceiver.h>
+#include <aidl/android/util/StatsEventParcel.h>
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "stats_event.h"
+#include "tests/statsd_test_util.h"
+
+using aidl::android::util::StatsEventParcel;
+using ::ndk::SharedRefBase;
+using std::make_shared;
+using std::shared_ptr;
+using std::vector;
+
+namespace android {
+namespace os {
+namespace statsd {
+
+namespace {
+
+int pullTagId1 = 10101;
+int pullTagId2 = 10102;
+int uid1 = 9999;
+int uid2 = 8888;
+ConfigKey configKey(50, 12345);
+ConfigKey badConfigKey(60, 54321);
+int unregisteredUid = 98765;
+int64_t coolDownNs = NS_PER_SEC;
+int64_t timeoutNs = NS_PER_SEC / 2;
+
+AStatsEvent* createSimpleEvent(int32_t atomId, int32_t value) {
+    AStatsEvent* event = AStatsEvent_obtain();
+    AStatsEvent_setAtomId(event, atomId);
+    AStatsEvent_writeInt32(event, value);
+    AStatsEvent_build(event);
+    return event;
+}
+
+class FakePullAtomCallback : public BnPullAtomCallback {
+public:
+    FakePullAtomCallback(int32_t uid) : mUid(uid){};
+    Status onPullAtom(int atomTag,
+                      const shared_ptr<IPullAtomResultReceiver>& resultReceiver) override {
+        vector<StatsEventParcel> parcels;
+        AStatsEvent* event = createSimpleEvent(atomTag, mUid);
+        size_t size;
+        uint8_t* buffer = AStatsEvent_getBuffer(event, &size);
+
+        StatsEventParcel p;
+        // vector.assign() creates a copy, but this is inevitable unless
+        // stats_event.h/c uses a vector as opposed to a buffer.
+        p.buffer.assign(buffer, buffer + size);
+        parcels.push_back(std::move(p));
+        AStatsEvent_release(event);
+        resultReceiver->pullFinished(atomTag, /*success*/ true, parcels);
+        return Status::ok();
+    }
+    int32_t mUid;
+};
+
+class FakePullUidProvider : public PullUidProvider {
+public:
+    vector<int32_t> getPullAtomUids(int atomId) override {
+        if (atomId == pullTagId1) {
+            return {uid2, uid1};
+        } else if (atomId == pullTagId2) {
+            return {uid2};
+        }
+        return {};
+    }
+};
+
+sp<StatsPullerManager> createPullerManagerAndRegister() {
+    sp<StatsPullerManager> pullerManager = new StatsPullerManager();
+    shared_ptr<FakePullAtomCallback> cb1 = SharedRefBase::make<FakePullAtomCallback>(uid1);
+    pullerManager->RegisterPullAtomCallback(uid1, pullTagId1, coolDownNs, timeoutNs, {}, cb1, true);
+    shared_ptr<FakePullAtomCallback> cb2 = SharedRefBase::make<FakePullAtomCallback>(uid2);
+    pullerManager->RegisterPullAtomCallback(uid2, pullTagId1, coolDownNs, timeoutNs, {}, cb2, true);
+    pullerManager->RegisterPullAtomCallback(uid1, pullTagId2, coolDownNs, timeoutNs, {}, cb1, true);
+    return pullerManager;
+}
+}  // anonymous namespace
+
+TEST(StatsPullerManagerTest, TestPullInvalidUid) {
+    sp<StatsPullerManager> pullerManager = createPullerManagerAndRegister();
+
+    vector<shared_ptr<LogEvent>> data;
+    EXPECT_FALSE(pullerManager->Pull(pullTagId1, {unregisteredUid}, &data, true));
+}
+
+TEST(StatsPullerManagerTest, TestPullChoosesCorrectUid) {
+    sp<StatsPullerManager> pullerManager = createPullerManagerAndRegister();
+
+    vector<shared_ptr<LogEvent>> data;
+    EXPECT_TRUE(pullerManager->Pull(pullTagId1, {uid1}, &data, true));
+    ASSERT_EQ(data.size(), 1);
+    EXPECT_EQ(data[0]->GetTagId(), pullTagId1);
+    ASSERT_EQ(data[0]->getValues().size(), 1);
+    EXPECT_EQ(data[0]->getValues()[0].mValue.int_value, uid1);
+}
+
+TEST(StatsPullerManagerTest, TestPullInvalidConfigKey) {
+    sp<StatsPullerManager> pullerManager = createPullerManagerAndRegister();
+    sp<FakePullUidProvider> uidProvider = new FakePullUidProvider();
+    pullerManager->RegisterPullUidProvider(configKey, uidProvider);
+
+    vector<shared_ptr<LogEvent>> data;
+    EXPECT_FALSE(pullerManager->Pull(pullTagId1, badConfigKey, &data, true));
+}
+
+TEST(StatsPullerManagerTest, TestPullConfigKeyGood) {
+    sp<StatsPullerManager> pullerManager = createPullerManagerAndRegister();
+    sp<FakePullUidProvider> uidProvider = new FakePullUidProvider();
+    pullerManager->RegisterPullUidProvider(configKey, uidProvider);
+
+    vector<shared_ptr<LogEvent>> data;
+    EXPECT_TRUE(pullerManager->Pull(pullTagId1, configKey, &data, true));
+    EXPECT_EQ(data[0]->GetTagId(), pullTagId1);
+    ASSERT_EQ(data[0]->getValues().size(), 1);
+    EXPECT_EQ(data[0]->getValues()[0].mValue.int_value, uid2);
+}
+
+TEST(StatsPullerManagerTest, TestPullConfigKeyNoPullerWithUid) {
+    sp<StatsPullerManager> pullerManager = createPullerManagerAndRegister();
+    sp<FakePullUidProvider> uidProvider = new FakePullUidProvider();
+    pullerManager->RegisterPullUidProvider(configKey, uidProvider);
+
+    vector<shared_ptr<LogEvent>> data;
+    EXPECT_FALSE(pullerManager->Pull(pullTagId2, configKey, &data, true));
+}
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android

cmds/statsd/tests/metrics/GaugeMetricProducer_test.cpp

@@ -105,6 +105,7 @@ TEST(GaugeMetricProducerTest, TestFirstBucket) {
                                       logEventMatcherIndex, eventMatcherWizard,
                                       -1, -1, tagId, 5, 600 * NS_PER_SEC + NS_PER_SEC / 2,
                                       pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     EXPECT_EQ(600500000000, gaugeProducer.mCurrentBucketStartTimeNs);
     EXPECT_EQ(10, gaugeProducer.mCurrentBucketNum);
@@ -132,10 +133,11 @@ TEST(GaugeMetricProducerTest, TestPulledEventsNoCondition) {
                     atomMatcherId, logEventMatcherIndex, atomMatcher, uidMap)});
 
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
-    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, _, _, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, kConfigKey, _, _, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, kConfigKey, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, Pull(tagId, _))
+    EXPECT_CALL(*pullerManager, Pull(tagId, kConfigKey, _, _))
-            .WillOnce(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+            .WillOnce(Invoke([](int tagId, const ConfigKey&,
+                                vector<std::shared_ptr<LogEvent>>* data, bool) {
                 data->clear();
                 data->push_back(makeLogEvent(tagId, bucketStartTimeNs + 10, 3, "some value", 11));
                 return true;
@@ -144,6 +146,7 @@ TEST(GaugeMetricProducerTest, TestPulledEventsNoCondition) {
     GaugeMetricProducer gaugeProducer(kConfigKey, metric, -1 /*-1 meaning no condition*/, wizard,
                                       logEventMatcherIndex, eventMatcherWizard, tagId, -1, tagId,
                                       bucketStartTimeNs, bucketStartTimeNs, pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     vector<shared_ptr<LogEvent>> allData;
     allData.clear();
@@ -222,6 +225,7 @@ TEST(GaugeMetricProducerTest, TestPushedEventsWithUpgrade) {
                                       logEventMatcherIndex, eventMatcherWizard,
                                       -1 /* -1 means no pulling */, -1, tagId, bucketStartTimeNs,
                                       bucketStartTimeNs, pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     sp<AnomalyTracker> anomalyTracker = gaugeProducer.addAnomalyTracker(alert, alarmMonitor);
     EXPECT_TRUE(anomalyTracker != nullptr);
@@ -287,19 +291,21 @@ TEST(GaugeMetricProducerTest, TestPulledWithUpgrade) {
                     atomMatcherId, logEventMatcherIndex, atomMatcher, uidMap)});
 
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
-    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, _, _, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, kConfigKey, _, _, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, kConfigKey, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, Pull(tagId, _))
+    EXPECT_CALL(*pullerManager, Pull(tagId, kConfigKey, _, _))
             .WillOnce(Return(false))
-            .WillOnce(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+            .WillOnce(Invoke(
-                data->clear();
+                    [](int tagId, const ConfigKey&, vector<std::shared_ptr<LogEvent>>* data, bool) {
-                data->push_back(CreateRepeatedValueLogEvent(tagId, eventUpgradeTimeNs, 2));
+                        data->clear();
-                return true;
+                        data->push_back(CreateRepeatedValueLogEvent(tagId, eventUpgradeTimeNs, 2));
-            }));
+                        return true;
+                    }));
 
     GaugeMetricProducer gaugeProducer(kConfigKey, metric, -1 /*-1 meaning no condition*/, wizard,
                                       logEventMatcherIndex, eventMatcherWizard, tagId, -1, tagId,
                                       bucketStartTimeNs, bucketStartTimeNs, pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     vector<shared_ptr<LogEvent>> allData;
     allData.push_back(CreateRepeatedValueLogEvent(tagId, bucketStartTimeNs + 1, 1));
@@ -351,13 +357,14 @@ TEST(GaugeMetricProducerTest, TestPulledWithAppUpgradeDisabled) {
                     atomMatcherId, logEventMatcherIndex, atomMatcher, uidMap)});
 
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
-    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, _, _, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, kConfigKey, _, _, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, kConfigKey, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, Pull(tagId, _)).WillOnce(Return(false));
+    EXPECT_CALL(*pullerManager, Pull(tagId, kConfigKey, _, _)).WillOnce(Return(false));
 
     GaugeMetricProducer gaugeProducer(kConfigKey, metric, -1 /*-1 meaning no condition*/, wizard,
                                       logEventMatcherIndex, eventMatcherWizard, tagId, -1, tagId,
                                       bucketStartTimeNs, bucketStartTimeNs, pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     vector<shared_ptr<LogEvent>> allData;
     allData.push_back(CreateRepeatedValueLogEvent(tagId, bucketStartTimeNs + 1, 1));
@@ -399,10 +406,11 @@ TEST(GaugeMetricProducerTest, TestPulledEventsWithCondition) {
                     atomMatcherId, logEventMatcherIndex, atomMatcher, uidMap)});
 
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
-    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, _, _, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, kConfigKey, _, _, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, kConfigKey, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, Pull(tagId, _))
+    EXPECT_CALL(*pullerManager, Pull(tagId, kConfigKey, _, _))
-            .WillOnce(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+            .WillOnce(Invoke([](int tagId, const ConfigKey&,
+                                vector<std::shared_ptr<LogEvent>>* data, bool) {
                 data->clear();
                 data->push_back(CreateRepeatedValueLogEvent(tagId, bucketStartTimeNs + 10, 100));
                 return true;
@@ -411,6 +419,7 @@ TEST(GaugeMetricProducerTest, TestPulledEventsWithCondition) {
     GaugeMetricProducer gaugeProducer(kConfigKey, metric, 1, wizard, logEventMatcherIndex,
                                       eventMatcherWizard, tagId, -1, tagId, bucketStartTimeNs,
                                       bucketStartTimeNs, pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     gaugeProducer.onConditionChanged(true, bucketStartTimeNs + 8);
     EXPECT_EQ(1UL, gaugeProducer.mCurrentSlicedBucket->size());
@@ -481,10 +490,11 @@ TEST(GaugeMetricProducerTest, TestPulledEventsWithSlicedCondition) {
                     }));
 
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
-    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, _, _, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, kConfigKey, _, _, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, kConfigKey, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, Pull(tagId, _))
+    EXPECT_CALL(*pullerManager, Pull(tagId, kConfigKey, _, _))
-            .WillOnce(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+            .WillOnce(Invoke([](int tagId, const ConfigKey&,
+                                vector<std::shared_ptr<LogEvent>>* data, bool) {
                 data->clear();
                 data->push_back(CreateTwoValueLogEvent(tagId, bucketStartTimeNs + 10, 1000, 100));
                 return true;
@@ -493,6 +503,7 @@ TEST(GaugeMetricProducerTest, TestPulledEventsWithSlicedCondition) {
     GaugeMetricProducer gaugeProducer(kConfigKey, metric, 1, wizard, logEventMatcherIndex,
                                       eventMatcherWizard, tagId, -1, tagId, bucketStartTimeNs,
                                       bucketStartTimeNs, pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     gaugeProducer.onSlicedConditionMayChange(true, bucketStartTimeNs + 8);
 
@@ -517,9 +528,9 @@ TEST(GaugeMetricProducerTest, TestPulledEventsAnomalyDetection) {
     sp<MockConditionWizard> wizard = new NaggyMock<MockConditionWizard>();
 
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
-    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, _, _, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, RegisterReceiver(tagId, kConfigKey, _, _, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, _)).WillOnce(Return());
+    EXPECT_CALL(*pullerManager, UnRegisterReceiver(tagId, kConfigKey, _)).WillOnce(Return());
-    EXPECT_CALL(*pullerManager, Pull(tagId, _)).WillOnce(Return(false));
+    EXPECT_CALL(*pullerManager, Pull(tagId, kConfigKey, _, _)).WillOnce(Return(false));
 
     GaugeMetric metric;
     metric.set_id(metricId);
@@ -539,6 +550,7 @@ TEST(GaugeMetricProducerTest, TestPulledEventsAnomalyDetection) {
     GaugeMetricProducer gaugeProducer(kConfigKey, metric, -1 /*-1 meaning no condition*/, wizard,
                                       logEventMatcherIndex, eventMatcherWizard, tagId, -1, tagId,
                                       bucketStartTimeNs, bucketStartTimeNs, pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     Alert alert;
     alert.set_id(101);
@@ -616,13 +628,15 @@ TEST(GaugeMetricProducerTest, TestPullOnTrigger) {
                     atomMatcherId, logEventMatcherIndex, atomMatcher, uidMap)});
 
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
-    EXPECT_CALL(*pullerManager, Pull(tagId, _))
+    EXPECT_CALL(*pullerManager, Pull(tagId, kConfigKey, _, _))
-            .WillOnce(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+            .WillOnce(Invoke([](int tagId, const ConfigKey&,
+                                vector<std::shared_ptr<LogEvent>>* data, bool) {
                 data->clear();
                 data->push_back(CreateRepeatedValueLogEvent(tagId, bucketStartTimeNs + 10, 4));
                 return true;
             }))
-            .WillOnce(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+            .WillOnce(Invoke([](int tagId, const ConfigKey&,
+                                vector<std::shared_ptr<LogEvent>>* data, bool) {
                 data->clear();
                 data->push_back(CreateRepeatedValueLogEvent(tagId, bucketStartTimeNs + 20, 5));
                 return true;
@@ -633,6 +647,7 @@ TEST(GaugeMetricProducerTest, TestPullOnTrigger) {
     GaugeMetricProducer gaugeProducer(kConfigKey, metric, -1 /*-1 meaning no condition*/, wizard,
                                       logEventMatcherIndex, eventMatcherWizard, tagId, triggerId,
                                       tagId, bucketStartTimeNs, bucketStartTimeNs, pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     EXPECT_EQ(0UL, gaugeProducer.mCurrentSlicedBucket->size());
 
@@ -682,18 +697,21 @@ TEST(GaugeMetricProducerTest, TestRemoveDimensionInOutput) {
                     atomMatcherId, logEventMatcherIndex, atomMatcher, uidMap)});
 
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
-    EXPECT_CALL(*pullerManager, Pull(tagId, _))
+    EXPECT_CALL(*pullerManager, Pull(tagId, kConfigKey, _, _))
-            .WillOnce(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+            .WillOnce(Invoke(
-                data->clear();
+                    [](int tagId, const ConfigKey&, vector<std::shared_ptr<LogEvent>>* data, bool) {
-                data->push_back(CreateTwoValueLogEvent(tagId, bucketStartTimeNs + 3, 3, 4));
+                        data->clear();
-                return true;
+                        data->push_back(CreateTwoValueLogEvent(tagId, bucketStartTimeNs + 3, 3, 4));
-            }))
+                        return true;
-            .WillOnce(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+                    }))
+            .WillOnce(Invoke([](int tagId, const ConfigKey&,
+                                vector<std::shared_ptr<LogEvent>>* data, bool) {
                 data->clear();
                 data->push_back(CreateTwoValueLogEvent(tagId, bucketStartTimeNs + 10, 4, 5));
                 return true;
             }))
-            .WillOnce(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+            .WillOnce(Invoke([](int tagId, const ConfigKey&,
+                                vector<std::shared_ptr<LogEvent>>* data, bool) {
                 data->clear();
                 data->push_back(CreateTwoValueLogEvent(tagId, bucketStartTimeNs + 20, 4, 6));
                 return true;
@@ -704,6 +722,7 @@ TEST(GaugeMetricProducerTest, TestRemoveDimensionInOutput) {
     GaugeMetricProducer gaugeProducer(kConfigKey, metric, -1 /*-1 meaning no condition*/, wizard,
                                       logEventMatcherIndex, eventMatcherWizard, tagId, triggerId,
                                       tagId, bucketStartTimeNs, bucketStartTimeNs, pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     LogEvent triggerEvent(/*uid=*/0, /*pid=*/0);
     CreateNoValuesLogEvent(&triggerEvent, triggerId, bucketStartTimeNs + 3);
@@ -752,18 +771,20 @@ TEST(GaugeMetricProducerTest_BucketDrop, TestBucketDropWhenBucketTooSmall) {
                     atomMatcherId, logEventMatcherIndex, atomMatcher, uidMap)});
 
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
-    EXPECT_CALL(*pullerManager, Pull(tagId, _))
+    EXPECT_CALL(*pullerManager, Pull(tagId, kConfigKey, _, _))
             // Bucket start.
-            .WillOnce(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+            .WillOnce(Invoke(
-                data->clear();
+                    [](int tagId, const ConfigKey&, vector<std::shared_ptr<LogEvent>>* data, bool) {
-                data->push_back(CreateRepeatedValueLogEvent(tagId, bucketStartTimeNs, 10));
+                        data->clear();
-                return true;
+                        data->push_back(CreateRepeatedValueLogEvent(tagId, bucketStartTimeNs, 10));
-            }));
+                        return true;
+                    }));
 
     int triggerId = 5;
     GaugeMetricProducer gaugeProducer(kConfigKey, metric, -1 /*-1 meaning no condition*/, wizard,
                                       logEventMatcherIndex, eventMatcherWizard, tagId, triggerId,
                                       tagId, bucketStartTimeNs, bucketStartTimeNs, pullerManager);
+    gaugeProducer.prepareFirstBucket();
 
     LogEvent triggerEvent(/*uid=*/0, /*pid=*/0);
     CreateNoValuesLogEvent(&triggerEvent, triggerId, bucketStartTimeNs + 3);

cmds/statsd/tests/metrics/metrics_test_helper.h

@@ -33,15 +33,24 @@ public:
 
 class MockStatsPullerManager : public StatsPullerManager {
 public:
-    MOCK_METHOD4(RegisterReceiver, void(int tagId, wp<PullDataReceiver> receiver,
+    MOCK_METHOD5(RegisterReceiver,
-                                        int64_t nextPulltimeNs, int64_t intervalNs));
+                 void(int tagId, const ConfigKey& key, wp<PullDataReceiver> receiver,
-    MOCK_METHOD2(UnRegisterReceiver, void(int tagId, wp<PullDataReceiver> receiver));
+                      int64_t nextPulltimeNs, int64_t intervalNs));
-    MOCK_METHOD2(Pull, bool(const int pullCode, vector<std::shared_ptr<LogEvent>>* data));
+    MOCK_METHOD3(UnRegisterReceiver,
+                 void(int tagId, const ConfigKey& key, wp<PullDataReceiver> receiver));
+    MOCK_METHOD4(Pull, bool(const int pullCode, const ConfigKey& key,
+                            vector<std::shared_ptr<LogEvent>>* data, bool useUids));
+    MOCK_METHOD4(Pull, bool(const int pullCode, const vector<int32_t>& uids,
+                            vector<std::shared_ptr<LogEvent>>* data, bool useUids));
+    MOCK_METHOD2(RegisterPullUidProvider,
+                 void(const ConfigKey& configKey, wp<PullUidProvider> provider));
+    MOCK_METHOD1(UnregisterPullUidProvider, void(const ConfigKey& configKey));
 };
 
 class MockUidMap : public UidMap {
  public:
   MOCK_CONST_METHOD1(getHostUidOrSelf, int(int uid));
+  MOCK_CONST_METHOD1(getAppUid, std::set<int32_t>(const string& package));
 };
 
 HashableDimensionKey getMockedDimensionKey(int tagId, int key, std::string value);

cmds/statsd/tests/shell/ShellSubscriber_test.cpp

@@ -187,8 +187,10 @@ TEST(ShellSubscriberTest, testPulledSubscription) {
     sp<MockUidMap> uidMap = new NaggyMock<MockUidMap>();
 
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
-    EXPECT_CALL(*pullerManager, Pull(10016, _))
+    const vector<int32_t> uids = {AID_SYSTEM};
-            .WillRepeatedly(Invoke([](int tagId, vector<std::shared_ptr<LogEvent>>* data) {
+    EXPECT_CALL(*pullerManager, Pull(10016, uids, _, _))
+            .WillRepeatedly(Invoke([](int tagId, const vector<int32_t>&,
+                                      vector<std::shared_ptr<LogEvent>>* data, bool) {
                 data->clear();
                 data->push_back(makeCpuActiveTimeAtom(/*uid=*/kUid1, /*timeMillis=*/kCpuTime1));
                 data->push_back(makeCpuActiveTimeAtom(/*uid=*/kUid2, /*timeMillis=*/kCpuTime2));

cmds/statsd/tools/localtools/src/com/android/statsd/shelltools/testdrive/TestDrive.java

@@ -19,6 +19,7 @@ import com.android.internal.os.StatsdConfigProto.AtomMatcher;
 import com.android.internal.os.StatsdConfigProto.EventMetric;
 import com.android.internal.os.StatsdConfigProto.FieldFilter;
 import com.android.internal.os.StatsdConfigProto.GaugeMetric;
+import com.android.internal.os.StatsdConfigProto.PullAtomPackages;
 import com.android.internal.os.StatsdConfigProto.SimpleAtomMatcher;
 import com.android.internal.os.StatsdConfigProto.StatsdConfig;
 import com.android.internal.os.StatsdConfigProto.TimeUnit;
@@ -33,6 +34,7 @@ import com.google.common.io.Files;
 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
@@ -49,19 +51,22 @@ public class TestDrive {
     private static final int VENDOR_PULLED_ATOM_START_TAG = 150000;
     private static final long CONFIG_ID = 54321;
     private static final String[] ALLOWED_LOG_SOURCES = {
-        "AID_GRAPHICS",
+            "AID_GRAPHICS",
-        "AID_INCIDENTD",
+            "AID_INCIDENTD",
-        "AID_STATSD",
+            "AID_STATSD",
-        "AID_RADIO",
+            "AID_RADIO",
-        "com.android.systemui",
+            "com.android.systemui",
-        "com.android.vending",
+            "com.android.vending",
-        "AID_SYSTEM",
+            "AID_SYSTEM",
-        "AID_ROOT",
+            "AID_ROOT",
-        "AID_BLUETOOTH",
+            "AID_BLUETOOTH",
-        "AID_LMKD",
+            "AID_LMKD",
-        "com.android.managedprovisioning",
+            "com.android.managedprovisioning",
-        "AID_MEDIA",
+            "AID_MEDIA",
-        "AID_NETWORK_STACK"
+            "AID_NETWORK_STACK"
+    };
+    private static final String[] DEFAULT_PULL_SOURCES = {
+            "AID_SYSTEM",
     };
     private static final Logger LOGGER = Logger.getLogger(TestDrive.class.getName());
 
@@ -158,6 +163,16 @@ public class TestDrive {
         StatsdConfig.Builder builder = StatsdConfig.newBuilder();
         builder
             .addAllAllowedLogSource(allowedSources)
+            .addAllDefaultPullPackages(Arrays.asList(DEFAULT_PULL_SOURCES))
+            .addPullAtomPackages(PullAtomPackages.newBuilder()
+                    .setAtomId(Atom.GPU_STATS_GLOBAL_INFO_FIELD_NUMBER)
+                    .addPackages("AID_GPU_SERVICE"))
+            .addPullAtomPackages(PullAtomPackages.newBuilder()
+                    .setAtomId(Atom.GPU_STATS_APP_INFO_FIELD_NUMBER)
+                    .addPackages("AID_GPU_SERVICE"))
+            .addPullAtomPackages(PullAtomPackages.newBuilder()
+                    .setAtomId(Atom.TRAIN_INFO_FIELD_NUMBER)
+                    .addPackages("AID_STATSD"))
             .setHashStringsInMetricReport(false);
 
         if (hasPulledAtom(atomIds)) {