Merge "Perform partial update for count metric"

7 files changed, 589 insertions, 58 deletions

diff --git a/cmds/statsd/src/metrics/CountMetricProducer.cpp b/cmds/statsd/src/metrics/CountMetricProducer.cpp
index 3dbb6ed47ff8..a8ef54a335c4 100644
--- a/cmds/statsd/src/metrics/CountMetricProducer.cpp
+++ b/cmds/statsd/src/metrics/CountMetricProducer.cpp
@@ -24,6 +24,7 @@
 #include <stdlib.h>
 
 #include "guardrail/StatsdStats.h"
+#include "metrics/parsing_utils/metrics_manager_util.h"
 #include "stats_log_util.h"
 #include "stats_util.h"
 
@@ -122,6 +123,47 @@ CountMetricProducer::~CountMetricProducer() {
     VLOG("~CountMetricProducer() called");
 }
 
+bool CountMetricProducer::onConfigUpdatedLocked(
+        const StatsdConfig& config, const int configIndex, const int metricIndex,
+        const vector<sp<AtomMatchingTracker>>& allAtomMatchingTrackers,
+        const unordered_map<int64_t, int>& oldAtomMatchingTrackerMap,
+        const unordered_map<int64_t, int>& newAtomMatchingTrackerMap,
+        const sp<EventMatcherWizard>& matcherWizard,
+        const vector<sp<ConditionTracker>>& allConditionTrackers,
+        const unordered_map<int64_t, int>& conditionTrackerMap, const sp<ConditionWizard>& wizard,
+        const unordered_map<int64_t, int>& metricToActivationMap,
+        unordered_map<int, vector<int>>& trackerToMetricMap,
+        unordered_map<int, vector<int>>& conditionToMetricMap,
+        unordered_map<int, vector<int>>& activationAtomTrackerToMetricMap,
+        unordered_map<int, vector<int>>& deactivationAtomTrackerToMetricMap,
+        vector<int>& metricsWithActivation) {
+    if (!MetricProducer::onConfigUpdatedLocked(
+                config, configIndex, metricIndex, allAtomMatchingTrackers,
+                oldAtomMatchingTrackerMap, newAtomMatchingTrackerMap, matcherWizard,
+                allConditionTrackers, conditionTrackerMap, wizard, metricToActivationMap,
+                trackerToMetricMap, conditionToMetricMap, activationAtomTrackerToMetricMap,
+                deactivationAtomTrackerToMetricMap, metricsWithActivation)) {
+        return false;
+    }
+
+    const CountMetric& metric = config.count_metric(configIndex);
+    int trackerIndex;
+    // Update appropriate indices, specifically mConditionIndex and MetricsManager maps.
+    if (!handleMetricWithAtomMatchingTrackers(metric.what(), metricIndex, false,
+                                              allAtomMatchingTrackers, newAtomMatchingTrackerMap,
+                                              trackerToMetricMap, trackerIndex)) {
+        return false;
+    }
+
+    if (metric.has_condition() &&
+        !handleMetricWithConditions(metric.condition(), metricIndex, conditionTrackerMap,
+                                    metric.links(), allConditionTrackers, mConditionTrackerIndex,
+                                    conditionToMetricMap)) {
+        return false;
+    }
+    return true;
+}
+
 void CountMetricProducer::onStateChanged(const int64_t eventTimeNs, const int32_t atomId,
                                          const HashableDimensionKey& primaryKey,
                                          const FieldValue& oldState, const FieldValue& newState) {
diff --git a/cmds/statsd/src/metrics/CountMetricProducer.h b/cmds/statsd/src/metrics/CountMetricProducer.h
index 6b2f2ca61ecc..0769ac459b8c 100644
--- a/cmds/statsd/src/metrics/CountMetricProducer.h
+++ b/cmds/statsd/src/metrics/CountMetricProducer.h
@@ -97,6 +97,22 @@ private:
     void flushCurrentBucketLocked(const int64_t& eventTimeNs,
                                   const int64_t& nextBucketStartTimeNs) override;
 
+    bool onConfigUpdatedLocked(
+            const StatsdConfig& config, const int configIndex, const int metricIndex,
+            const std::vector<sp<AtomMatchingTracker>>& allAtomMatchingTrackers,
+            const std::unordered_map<int64_t, int>& oldAtomMatchingTrackerMap,
+            const std::unordered_map<int64_t, int>& newAtomMatchingTrackerMap,
+            const sp<EventMatcherWizard>& matcherWizard,
+            const std::vector<sp<ConditionTracker>>& allConditionTrackers,
+            const std::unordered_map<int64_t, int>& conditionTrackerMap,
+            const sp<ConditionWizard>& wizard,
+            const std::unordered_map<int64_t, int>& metricToActivationMap,
+            std::unordered_map<int, std::vector<int>>& trackerToMetricMap,
+            std::unordered_map<int, std::vector<int>>& conditionToMetricMap,
+            std::unordered_map<int, std::vector<int>>& activationAtomTrackerToMetricMap,
+            std::unordered_map<int, std::vector<int>>& deactivationAtomTrackerToMetricMap,
+            std::vector<int>& metricsWithActivation) override;
+
     std::unordered_map<MetricDimensionKey, std::vector<CountBucket>> mPastBuckets;
 
     // The current bucket (may be a partial bucket).
diff --git a/cmds/statsd/src/metrics/MetricProducer.h b/cmds/statsd/src/metrics/MetricProducer.h
index 320b28581a38..4360010746ae 100644
--- a/cmds/statsd/src/metrics/MetricProducer.h
+++ b/cmds/statsd/src/metrics/MetricProducer.h
@@ -566,7 +566,9 @@ protected:
     FRIEND_TEST(MetricsManagerTest, TestInitialConditions);
 
     FRIEND_TEST(ConfigUpdateTest, TestUpdateMetricActivations);
+    FRIEND_TEST(ConfigUpdateTest, TestUpdateCountMetrics);
     FRIEND_TEST(ConfigUpdateTest, TestUpdateEventMetrics);
+    FRIEND_TEST(ConfigUpdateTest, TestUpdateMetricsMultipleTypes);
 };
 
 }  // namespace statsd
diff --git a/cmds/statsd/src/metrics/parsing_utils/config_update_utils.cpp b/cmds/statsd/src/metrics/parsing_utils/config_update_utils.cpp
index 12d3b9487635..5ff69753a4a3 100644
--- a/cmds/statsd/src/metrics/parsing_utils/config_update_utils.cpp
+++ b/cmds/statsd/src/metrics/parsing_utils/config_update_utils.cpp
@@ -587,6 +587,51 @@ bool updateMetrics(const ConfigKey& key, const StatsdConfig& config, const int64
 
     // Now, perform the update. Must iterate the metric types in the same order
     int metricIndex = 0;
+    for (int i = 0; i < config.count_metric_size(); i++, metricIndex++) {
+        newMetricProducerMap[config.count_metric(i).id()] = metricIndex;
+        const CountMetric& metric = config.count_metric(i);
+        switch (metricsToUpdate[metricIndex]) {
+            case UPDATE_PRESERVE: {
+                const auto& oldMetricProducerIt = oldMetricProducerMap.find(metric.id());
+                if (oldMetricProducerIt == oldMetricProducerMap.end()) {
+                    ALOGE("Could not find Metric %lld in the previous config, but expected it "
+                          "to be there",
+                          (long long)metric.id());
+                    return false;
+                }
+                const int oldIndex = oldMetricProducerIt->second;
+                sp<MetricProducer> producer = oldMetricProducers[oldIndex];
+                producer->onConfigUpdated(
+                        config, i, metricIndex, allAtomMatchingTrackers, oldAtomMatchingTrackerMap,
+                        newAtomMatchingTrackerMap, matcherWizard, allConditionTrackers,
+                        conditionTrackerMap, wizard, metricToActivationMap, trackerToMetricMap,
+                        conditionToMetricMap, activationAtomTrackerToMetricMap,
+                        deactivationAtomTrackerToMetricMap, metricsWithActivation);
+                newMetricProducers.push_back(producer);
+                break;
+            }
+            case UPDATE_REPLACE:
+            case UPDATE_NEW: {
+                sp<MetricProducer> producer = createCountMetricProducerAndUpdateMetadata(
+                        key, config, timeBaseNs, currentTimeNs, metric, metricIndex,
+                        allAtomMatchingTrackers, newAtomMatchingTrackerMap, allConditionTrackers,
+                        conditionTrackerMap, initialConditionCache, wizard, stateAtomIdMap,
+                        allStateGroupMaps, metricToActivationMap, trackerToMetricMap,
+                        conditionToMetricMap, activationAtomTrackerToMetricMap,
+                        deactivationAtomTrackerToMetricMap, metricsWithActivation);
+                if (producer == nullptr) {
+                    return false;
+                }
+                newMetricProducers.push_back(producer);
+                break;
+            }
+            default: {
+                ALOGE("Metric \"%lld\" update state is unknown. This should never happen",
+                      (long long)metric.id());
+                return false;
+            }
+        }
+    }
     for (int i = 0; i < config.event_metric_size(); i++, metricIndex++) {
         newMetricProducerMap[config.event_metric(i).id()] = metricIndex;
         const EventMetric& metric = config.event_metric(i);
diff --git a/cmds/statsd/src/metrics/parsing_utils/metrics_manager_util.cpp b/cmds/statsd/src/metrics/parsing_utils/metrics_manager_util.cpp
index 51df7df5174e..98c6222c4326 100644
--- a/cmds/statsd/src/metrics/parsing_utils/metrics_manager_util.cpp
+++ b/cmds/statsd/src/metrics/parsing_utils/metrics_manager_util.cpp
@@ -348,6 +348,81 @@ bool handleMetricActivationOnConfigUpdate(
     return true;
 }
 
+sp<MetricProducer> createCountMetricProducerAndUpdateMetadata(
+        const ConfigKey& key, const StatsdConfig& config, const int64_t timeBaseNs,
+        const int64_t currentTimeNs, const CountMetric& metric, const int metricIndex,
+        const vector<sp<AtomMatchingTracker>>& allAtomMatchingTrackers,
+        const unordered_map<int64_t, int>& atomMatchingTrackerMap,
+        vector<sp<ConditionTracker>>& allConditionTrackers,
+        const unordered_map<int64_t, int>& conditionTrackerMap,
+        const vector<ConditionState>& initialConditionCache, const sp<ConditionWizard>& wizard,
+        const unordered_map<int64_t, int>& stateAtomIdMap,
+        const unordered_map<int64_t, unordered_map<int, int64_t>>& allStateGroupMaps,
+        const unordered_map<int64_t, int>& metricToActivationMap,
+        unordered_map<int, vector<int>>& trackerToMetricMap,
+        unordered_map<int, vector<int>>& conditionToMetricMap,
+        unordered_map<int, vector<int>>& activationAtomTrackerToMetricMap,
+        unordered_map<int, vector<int>>& deactivationAtomTrackerToMetricMap,
+        vector<int>& metricsWithActivation) {
+    if (!metric.has_id() || !metric.has_what()) {
+        ALOGW("cannot find metric id or \"what\" in CountMetric \"%lld\"", (long long)metric.id());
+        return nullptr;
+    }
+    int trackerIndex;
+    if (!handleMetricWithAtomMatchingTrackers(metric.what(), metricIndex,
+                                              metric.has_dimensions_in_what(),
+                                              allAtomMatchingTrackers, atomMatchingTrackerMap,
+                                              trackerToMetricMap, trackerIndex)) {
+        return nullptr;
+    }
+
+    int conditionIndex = -1;
+    if (metric.has_condition()) {
+        if (!handleMetricWithConditions(metric.condition(), metricIndex, conditionTrackerMap,
+                                        metric.links(), allConditionTrackers, conditionIndex,
+                                        conditionToMetricMap)) {
+            return nullptr;
+        }
+    } else {
+        if (metric.links_size() > 0) {
+            ALOGW("metrics has a MetricConditionLink but doesn't have a condition");
+            return nullptr;
+        }
+    }
+
+    std::vector<int> slicedStateAtoms;
+    unordered_map<int, unordered_map<int, int64_t>> stateGroupMap;
+    if (metric.slice_by_state_size() > 0) {
+        if (!handleMetricWithStates(config, metric.slice_by_state(), stateAtomIdMap,
+                                    allStateGroupMaps, slicedStateAtoms, stateGroupMap)) {
+            return nullptr;
+        }
+    } else {
+        if (metric.state_link_size() > 0) {
+            ALOGW("CountMetric has a MetricStateLink but doesn't have a slice_by_state");
+            return nullptr;
+        }
+    }
+
+    unordered_map<int, shared_ptr<Activation>> eventActivationMap;
+    unordered_map<int, vector<shared_ptr<Activation>>> eventDeactivationMap;
+    if (!handleMetricActivation(config, metric.id(), metricIndex, metricToActivationMap,
+                                atomMatchingTrackerMap, activationAtomTrackerToMetricMap,
+                                deactivationAtomTrackerToMetricMap, metricsWithActivation,
+                                eventActivationMap, eventDeactivationMap)) {
+        return nullptr;
+    }
+
+    uint64_t metricHash;
+    if (!getMetricProtoHash(config, metric, metric.id(), metricToActivationMap, metricHash)) {
+        return nullptr;
+    }
+
+    return new CountMetricProducer(key, metric, conditionIndex, initialConditionCache, wizard,
+                                   metricHash, timeBaseNs, currentTimeNs, eventActivationMap,
+                                   eventDeactivationMap, slicedStateAtoms, stateGroupMap);
+}
+
 sp<MetricProducer> createEventMetricProducerAndUpdateMetadata(
         const ConfigKey& key, const StatsdConfig& config, const int64_t timeBaseNs,
         const EventMetric& metric, const int metricIndex,
@@ -550,68 +625,20 @@ bool initMetrics(const ConfigKey& key, const StatsdConfig& config, const int64_t
     // Build MetricProducers for each metric defined in config.
     // build CountMetricProducer
     for (int i = 0; i < config.count_metric_size(); i++) {
-        const CountMetric& metric = config.count_metric(i);
-        if (!metric.has_what()) {
-            ALOGW("cannot find \"what\" in CountMetric \"%lld\"", (long long)metric.id());
-            return false;
-        }
-
         int metricIndex = allMetricProducers.size();
+        const CountMetric& metric = config.count_metric(i);
         metricMap.insert({metric.id(), metricIndex});
-        int trackerIndex;
-        if (!handleMetricWithAtomMatchingTrackers(metric.what(), metricIndex,
-                                                  metric.has_dimensions_in_what(),
-                                                  allAtomMatchingTrackers, atomMatchingTrackerMap,
-                                                  trackerToMetricMap, trackerIndex)) {
-            return false;
-        }
-
-        int conditionIndex = -1;
-        if (metric.has_condition()) {
-            if (!handleMetricWithConditions(metric.condition(), metricIndex, conditionTrackerMap,
-                                            metric.links(), allConditionTrackers, conditionIndex,
-                                            conditionToMetricMap)) {
-                return false;
-            }
-        } else {
-            if (metric.links_size() > 0) {
-                ALOGW("metrics has a MetricConditionLink but doesn't have a condition");
-                return false;
-            }
-        }
-
-        std::vector<int> slicedStateAtoms;
-        unordered_map<int, unordered_map<int, int64_t>> stateGroupMap;
-        if (metric.slice_by_state_size() > 0) {
-            if (!handleMetricWithStates(config, metric.slice_by_state(), stateAtomIdMap,
-                                        allStateGroupMaps, slicedStateAtoms, stateGroupMap)) {
-                return false;
-            }
-        } else {
-            if (metric.state_link_size() > 0) {
-                ALOGW("CountMetric has a MetricStateLink but doesn't have a slice_by_state");
-                return false;
-            }
-        }
-
-        unordered_map<int, shared_ptr<Activation>> eventActivationMap;
-        unordered_map<int, vector<shared_ptr<Activation>>> eventDeactivationMap;
-        bool success = handleMetricActivation(
-                config, metric.id(), metricIndex, metricToActivationMap, atomMatchingTrackerMap,
+        sp<MetricProducer> producer = createCountMetricProducerAndUpdateMetadata(
+                key, config, timeBaseTimeNs, currentTimeNs, metric, metricIndex,
+                allAtomMatchingTrackers, atomMatchingTrackerMap, allConditionTrackers,
+                conditionTrackerMap, initialConditionCache, wizard, stateAtomIdMap,
+                allStateGroupMaps, metricToActivationMap, trackerToMetricMap, conditionToMetricMap,
                 activationAtomTrackerToMetricMap, deactivationAtomTrackerToMetricMap,
-                metricsWithActivation, eventActivationMap, eventDeactivationMap);
-        if (!success) return false;
-
-        uint64_t metricHash;
-        if (!getMetricProtoHash(config, metric, metric.id(), metricToActivationMap, metricHash)) {
+                metricsWithActivation);
+        if (producer == nullptr) {
             return false;
         }
-
-        sp<MetricProducer> countProducer = new CountMetricProducer(
-                key, metric, conditionIndex, initialConditionCache, wizard, metricHash,
-                timeBaseTimeNs, currentTimeNs, eventActivationMap, eventDeactivationMap,
-                slicedStateAtoms, stateGroupMap);
-        allMetricProducers.push_back(countProducer);
+        allMetricProducers.push_back(producer);
     }
 
     // build DurationMetricProducer
diff --git a/cmds/statsd/src/metrics/parsing_utils/metrics_manager_util.h b/cmds/statsd/src/metrics/parsing_utils/metrics_manager_util.h
index fdfe5cfa6491..f6b2b7d1602c 100644
--- a/cmds/statsd/src/metrics/parsing_utils/metrics_manager_util.h
+++ b/cmds/statsd/src/metrics/parsing_utils/metrics_manager_util.h
@@ -93,6 +93,27 @@ bool handleMetricActivationOnConfigUpdate(
         std::unordered_map<int, shared_ptr<Activation>>& newEventActivationMap,
         std::unordered_map<int, std::vector<shared_ptr<Activation>>>& newEventDeactivationMap);
 
+// Creates a CountMetricProducer and updates the vectors/maps used by MetricsManager with
+// the appropriate indices. Returns an sp to the producer, or null if there was an error.
+sp<MetricProducer> createCountMetricProducerAndUpdateMetadata(
+        const ConfigKey& key, const StatsdConfig& config, const int64_t timeBaseNs,
+        const int64_t currentTimeNs, const CountMetric& metric, const int metricIndex,
+        const std::vector<sp<AtomMatchingTracker>>& allAtomMatchingTrackers,
+        const std::unordered_map<int64_t, int>& atomMatchingTrackerMap,
+        std::vector<sp<ConditionTracker>>& allConditionTrackers,
+        const std::unordered_map<int64_t, int>& conditionTrackerMap,
+        const std::vector<ConditionState>& initialConditionCache, const sp<ConditionWizard>& wizard,
+        const std::unordered_map<int64_t, int>& stateAtomIdMap,
+        const std::unordered_map<int64_t, std::unordered_map<int, int64_t>>& allStateGroupMaps,
+        const std::unordered_map<int64_t, int>& metricToActivationMap,
+        std::unordered_map<int, std::vector<int>>& trackerToMetricMap,
+        std::unordered_map<int, std::vector<int>>& conditionToMetricMap,
+        std::unordered_map<int, std::vector<int>>& activationAtomTrackerToMetricMap,
+        std::unordered_map<int, std::vector<int>>& deactivationAtomTrackerToMetricMap,
+        std::vector<int>& metricsWithActivation);
+
+// Creates an EventMetricProducer and updates the vectors/maps used by MetricsManager with
+// the appropriate indices. Returns an sp to the producer, or null if there was an error.
 sp<MetricProducer> createEventMetricProducerAndUpdateMetadata(
         const ConfigKey& key, const StatsdConfig& config, const int64_t timeBaseNs,
         const EventMetric& metric, const int metricIndex,
diff --git a/cmds/statsd/tests/metrics/parsing_utils/config_update_utils_test.cpp b/cmds/statsd/tests/metrics/parsing_utils/config_update_utils_test.cpp
index 3b346c11e831..9e6758785911 100644
--- a/cmds/statsd/tests/metrics/parsing_utils/config_update_utils_test.cpp
+++ b/cmds/statsd/tests/metrics/parsing_utils/config_update_utils_test.cpp
@@ -97,6 +97,7 @@ public:
         metricsWithActivation.clear();
         oldStateHashes.clear();
         noReportMetricIds.clear();
+        StateManager::getInstance().clear();
     }
 };
 
@@ -121,6 +122,21 @@ EventMetric createEventMetric(string name, int64_t what, optional<int64_t> condi
     return metric;
 }
 
+CountMetric createCountMetric(string name, int64_t what, optional<int64_t> condition,
+                              vector<int64_t> states) {
+    CountMetric metric;
+    metric.set_id(StringToId(name));
+    metric.set_what(what);
+    metric.set_bucket(TEN_MINUTES);
+    if (condition) {
+        metric.set_condition(condition.value());
+    }
+    for (const int64_t state : states) {
+        metric.add_slice_by_state(state);
+    }
+    return metric;
+}
+
 }  // anonymous namespace
 
 TEST_F(ConfigUpdateTest, TestSimpleMatcherPreserve) {
@@ -1419,7 +1435,7 @@ TEST_F(ConfigUpdateTest, TestUpdateEventMetrics) {
     EXPECT_EQ(oldMetricProducers[oldMetricProducerMap.at(event1Id)],
               newMetricProducers[newMetricProducerMap.at(event1Id)]);
 
-    // Make sure replaced conditions are different and included in replacedConditions.
+    // Make sure replaced metrics are different.
     EXPECT_NE(oldMetricProducers[oldMetricProducerMap.at(event2Id)],
               newMetricProducers[newMetricProducerMap.at(event2Id)]);
     EXPECT_NE(oldMetricProducers[oldMetricProducerMap.at(event3Id)],
@@ -1478,6 +1494,239 @@ TEST_F(ConfigUpdateTest, TestUpdateEventMetrics) {
     EXPECT_EQ(oldConditionWizard->getStrongCount(), 1);
 }
 
+TEST_F(ConfigUpdateTest, TestUpdateCountMetrics) {
+    StatsdConfig config;
+
+    // Add atom matchers/predicates/states. These are mostly needed for initStatsdConfig.
+    AtomMatcher matcher1 = CreateScreenTurnedOnAtomMatcher();
+    int64_t matcher1Id = matcher1.id();
+    *config.add_atom_matcher() = matcher1;
+
+    AtomMatcher matcher2 = CreateScreenTurnedOffAtomMatcher();
+    int64_t matcher2Id = matcher2.id();
+    *config.add_atom_matcher() = matcher2;
+
+    AtomMatcher matcher3 = CreateStartScheduledJobAtomMatcher();
+    int64_t matcher3Id = matcher3.id();
+    *config.add_atom_matcher() = matcher3;
+
+    AtomMatcher matcher4 = CreateFinishScheduledJobAtomMatcher();
+    int64_t matcher4Id = matcher4.id();
+    *config.add_atom_matcher() = matcher4;
+
+    AtomMatcher matcher5 = CreateBatterySaverModeStartAtomMatcher();
+    int64_t matcher5Id = matcher5.id();
+    *config.add_atom_matcher() = matcher5;
+
+    Predicate predicate1 = CreateScreenIsOnPredicate();
+    int64_t predicate1Id = predicate1.id();
+    *config.add_predicate() = predicate1;
+
+    State state1 = CreateScreenStateWithOnOffMap(0x123, 0x321);
+    int64_t state1Id = state1.id();
+    *config.add_state() = state1;
+
+    State state2 = CreateScreenState();
+    int64_t state2Id = state2.id();
+    *config.add_state() = state2;
+
+    // Add a few count metrics.
+    // Will be preserved.
+    CountMetric count1 = createCountMetric("COUNT1", matcher1Id, predicate1Id, {state1Id});
+    int64_t count1Id = count1.id();
+    *config.add_count_metric() = count1;
+
+    // Will be replaced.
+    CountMetric count2 = createCountMetric("COUNT2", matcher2Id, nullopt, {});
+    int64_t count2Id = count2.id();
+    *config.add_count_metric() = count2;
+
+    // Will be replaced.
+    CountMetric count3 = createCountMetric("COUNT3", matcher3Id, nullopt, {});
+    int64_t count3Id = count3.id();
+    *config.add_count_metric() = count3;
+
+    // Will be replaced.
+    CountMetric count4 = createCountMetric("COUNT4", matcher4Id, nullopt, {state2Id});
+    int64_t count4Id = count4.id();
+    *config.add_count_metric() = count4;
+
+    // Will be deleted.
+    CountMetric count5 = createCountMetric("COUNT5", matcher5Id, nullopt, {});
+    int64_t count5Id = count5.id();
+    *config.add_count_metric() = count5;
+
+    EXPECT_TRUE(initConfig(config));
+
+    // Change bucket size of count2, causing it to be replaced.
+    count2.set_bucket(ONE_HOUR);
+
+    // Mark matcher 3 as replaced. Causes count3 to be replaced.
+    set<int64_t> replacedMatchers;
+    replacedMatchers.insert(matcher3Id);
+
+    // Mark state 2 as replaced and change the state to be about a different atom.
+    // Causes count4 to be replaced.
+    set<int64_t> replacedStates;
+    replacedStates.insert(state2Id);
+    state2.set_atom_id(util::BATTERY_SAVER_MODE_STATE_CHANGED);
+
+    // Fake that predicate 1 is true for count metric 1.
+    ASSERT_EQ(oldMetricProducers[0]->getMetricId(), count1Id);
+    oldMetricProducers[0]->onConditionChanged(true, /*timestamp=*/0);
+    EXPECT_EQ(oldMetricProducers[0]->mCondition, ConditionState::kTrue);
+
+    EXPECT_EQ(StateManager::getInstance().getStateTrackersCount(), 1);
+    // Tell the StateManager that the screen is on.
+    unique_ptr<LogEvent> event =
+            CreateScreenStateChangedEvent(0, android::view::DisplayStateEnum::DISPLAY_STATE_ON);
+    StateManager::getInstance().onLogEvent(*event);
+
+    // New count metric. Should have an initial condition of true since it depends on predicate1.
+    CountMetric count6 = createCountMetric("EVENT6", matcher2Id, predicate1Id, {state1Id});
+    int64_t count6Id = count6.id();
+
+    // Map the matchers and predicates in reverse order to force the indices to change.
+    std::unordered_map<int64_t, int> newAtomMatchingTrackerMap;
+    const int matcher5Index = 0;
+    newAtomMatchingTrackerMap[matcher5Id] = 0;
+    const int matcher4Index = 1;
+    newAtomMatchingTrackerMap[matcher4Id] = 1;
+    const int matcher3Index = 2;
+    newAtomMatchingTrackerMap[matcher3Id] = 2;
+    const int matcher2Index = 3;
+    newAtomMatchingTrackerMap[matcher2Id] = 3;
+    const int matcher1Index = 4;
+    newAtomMatchingTrackerMap[matcher1Id] = 4;
+    // Use the existing matchers. A bit hacky, but saves code and we don't rely on them.
+    vector<sp<AtomMatchingTracker>> newAtomMatchingTrackers(5);
+    std::reverse_copy(oldAtomMatchingTrackers.begin(), oldAtomMatchingTrackers.end(),
+                      newAtomMatchingTrackers.begin());
+
+    std::unordered_map<int64_t, int> newConditionTrackerMap;
+    const int predicate1Index = 0;
+    newConditionTrackerMap[predicate1Id] = 0;
+    // Use the existing conditionTrackers. A bit hacky, but saves code and we don't rely on them.
+    vector<sp<ConditionTracker>> newConditionTrackers(1);
+    std::reverse_copy(oldConditionTrackers.begin(), oldConditionTrackers.end(),
+                      newConditionTrackers.begin());
+    // Fake that predicate1 is true for all new metrics.
+    vector<ConditionState> conditionCache = {ConditionState::kTrue};
+
+    StatsdConfig newConfig;
+    *newConfig.add_count_metric() = count6;
+    const int count6Index = 0;
+    *newConfig.add_count_metric() = count3;
+    const int count3Index = 1;
+    *newConfig.add_count_metric() = count1;
+    const int count1Index = 2;
+    *newConfig.add_count_metric() = count4;
+    const int count4Index = 3;
+    *newConfig.add_count_metric() = count2;
+    const int count2Index = 4;
+
+    *newConfig.add_state() = state1;
+    *newConfig.add_state() = state2;
+
+    unordered_map<int64_t, int> stateAtomIdMap;
+    unordered_map<int64_t, unordered_map<int, int64_t>> allStateGroupMaps;
+    map<int64_t, uint64_t> stateProtoHashes;
+    EXPECT_TRUE(initStates(newConfig, stateAtomIdMap, allStateGroupMaps, stateProtoHashes));
+    EXPECT_EQ(stateAtomIdMap[state2Id], util::BATTERY_SAVER_MODE_STATE_CHANGED);
+
+    // Output data structures to validate.
+    unordered_map<int64_t, int> newMetricProducerMap;
+    vector<sp<MetricProducer>> newMetricProducers;
+    unordered_map<int, vector<int>> conditionToMetricMap;
+    unordered_map<int, vector<int>> trackerToMetricMap;
+    set<int64_t> noReportMetricIds;
+    unordered_map<int, vector<int>> activationAtomTrackerToMetricMap;
+    unordered_map<int, vector<int>> deactivationAtomTrackerToMetricMap;
+    vector<int> metricsWithActivation;
+    EXPECT_TRUE(updateMetrics(
+            key, newConfig, /*timeBaseNs=*/123, /*currentTimeNs=*/12345, new StatsPullerManager(),
+            oldAtomMatchingTrackerMap, newAtomMatchingTrackerMap, replacedMatchers,
+            newAtomMatchingTrackers, newConditionTrackerMap, /*replacedConditions=*/{},
+            newConditionTrackers, conditionCache, stateAtomIdMap, allStateGroupMaps, replacedStates,
+            oldMetricProducerMap, oldMetricProducers, newMetricProducerMap, newMetricProducers,
+            conditionToMetricMap, trackerToMetricMap, noReportMetricIds,
+            activationAtomTrackerToMetricMap, deactivationAtomTrackerToMetricMap,
+            metricsWithActivation));
+
+    unordered_map<int64_t, int> expectedMetricProducerMap = {
+            {count1Id, count1Index}, {count2Id, count2Index}, {count3Id, count3Index},
+            {count4Id, count4Index}, {count6Id, count6Index},
+    };
+    EXPECT_THAT(newMetricProducerMap, ContainerEq(expectedMetricProducerMap));
+
+    // Make sure preserved metrics are the same.
+    ASSERT_EQ(newMetricProducers.size(), 5);
+    EXPECT_EQ(oldMetricProducers[oldMetricProducerMap.at(count1Id)],
+              newMetricProducers[newMetricProducerMap.at(count1Id)]);
+
+    // Make sure replaced metrics are different.
+    EXPECT_NE(oldMetricProducers[oldMetricProducerMap.at(count2Id)],
+              newMetricProducers[newMetricProducerMap.at(count2Id)]);
+    EXPECT_NE(oldMetricProducers[oldMetricProducerMap.at(count3Id)],
+              newMetricProducers[newMetricProducerMap.at(count3Id)]);
+    EXPECT_NE(oldMetricProducers[oldMetricProducerMap.at(count4Id)],
+              newMetricProducers[newMetricProducerMap.at(count4Id)]);
+
+    // Verify the conditionToMetricMap.
+    ASSERT_EQ(conditionToMetricMap.size(), 1);
+    const vector<int>& condition1Metrics = conditionToMetricMap[predicate1Index];
+    EXPECT_THAT(condition1Metrics, UnorderedElementsAre(count1Index, count6Index));
+
+    // Verify the trackerToMetricMap.
+    ASSERT_EQ(trackerToMetricMap.size(), 4);
+    const vector<int>& matcher1Metrics = trackerToMetricMap[matcher1Index];
+    EXPECT_THAT(matcher1Metrics, UnorderedElementsAre(count1Index));
+    const vector<int>& matcher2Metrics = trackerToMetricMap[matcher2Index];
+    EXPECT_THAT(matcher2Metrics, UnorderedElementsAre(count2Index, count6Index));
+    const vector<int>& matcher3Metrics = trackerToMetricMap[matcher3Index];
+    EXPECT_THAT(matcher3Metrics, UnorderedElementsAre(count3Index));
+    const vector<int>& matcher4Metrics = trackerToMetricMap[matcher4Index];
+    EXPECT_THAT(matcher4Metrics, UnorderedElementsAre(count4Index));
+
+    // Verify event activation/deactivation maps.
+    ASSERT_EQ(activationAtomTrackerToMetricMap.size(), 0);
+    ASSERT_EQ(deactivationAtomTrackerToMetricMap.size(), 0);
+    ASSERT_EQ(metricsWithActivation.size(), 0);
+
+    // Verify tracker indices/ids/conditions/states are correct.
+    EXPECT_EQ(newMetricProducers[count1Index]->getMetricId(), count1Id);
+    EXPECT_EQ(newMetricProducers[count1Index]->mConditionTrackerIndex, predicate1Index);
+    EXPECT_EQ(newMetricProducers[count1Index]->mCondition, ConditionState::kTrue);
+    EXPECT_THAT(newMetricProducers[count1Index]->getSlicedStateAtoms(),
+                UnorderedElementsAre(util::SCREEN_STATE_CHANGED));
+    EXPECT_EQ(newMetricProducers[count2Index]->getMetricId(), count2Id);
+    EXPECT_EQ(newMetricProducers[count2Index]->mConditionTrackerIndex, -1);
+    EXPECT_EQ(newMetricProducers[count2Index]->mCondition, ConditionState::kTrue);
+    EXPECT_TRUE(newMetricProducers[count2Index]->getSlicedStateAtoms().empty());
+    EXPECT_EQ(newMetricProducers[count3Index]->getMetricId(), count3Id);
+    EXPECT_EQ(newMetricProducers[count3Index]->mConditionTrackerIndex, -1);
+    EXPECT_EQ(newMetricProducers[count3Index]->mCondition, ConditionState::kTrue);
+    EXPECT_TRUE(newMetricProducers[count3Index]->getSlicedStateAtoms().empty());
+    EXPECT_EQ(newMetricProducers[count4Index]->getMetricId(), count4Id);
+    EXPECT_EQ(newMetricProducers[count4Index]->mConditionTrackerIndex, -1);
+    EXPECT_EQ(newMetricProducers[count4Index]->mCondition, ConditionState::kTrue);
+    EXPECT_THAT(newMetricProducers[count4Index]->getSlicedStateAtoms(),
+                UnorderedElementsAre(util::BATTERY_SAVER_MODE_STATE_CHANGED));
+    EXPECT_EQ(newMetricProducers[count6Index]->getMetricId(), count6Id);
+    EXPECT_EQ(newMetricProducers[count6Index]->mConditionTrackerIndex, predicate1Index);
+    EXPECT_EQ(newMetricProducers[count6Index]->mCondition, ConditionState::kTrue);
+    EXPECT_THAT(newMetricProducers[count6Index]->getSlicedStateAtoms(),
+                UnorderedElementsAre(util::SCREEN_STATE_CHANGED));
+
+    oldMetricProducers.clear();
+    // Ensure that the screen state StateTracker did not get deleted and replaced.
+    EXPECT_EQ(StateManager::getInstance().getStateTrackersCount(), 2);
+    FieldValue screenState;
+    StateManager::getInstance().getStateValue(util::SCREEN_STATE_CHANGED, DEFAULT_DIMENSION_KEY,
+                                              &screenState);
+    EXPECT_EQ(screenState.mValue.int_value, android::view::DisplayStateEnum::DISPLAY_STATE_ON);
+}
+
 TEST_F(ConfigUpdateTest, TestUpdateMetricActivations) {
     StatsdConfig config;
     // Add atom matchers
@@ -1605,6 +1854,135 @@ TEST_F(ConfigUpdateTest, TestUpdateMetricActivations) {
     EXPECT_THAT(producer->mEventDeactivationMap[matcher2Index],
                 UnorderedElementsAre(matcher4Activation));
 }
+
+TEST_F(ConfigUpdateTest, TestUpdateMetricsMultipleTypes) {
+    StatsdConfig config;
+    // Add atom matchers/predicates/states. These are mostly needed for initStatsdConfig
+    AtomMatcher matcher1 = CreateScreenTurnedOnAtomMatcher();
+    int64_t matcher1Id = matcher1.id();
+    *config.add_atom_matcher() = matcher1;
+
+    AtomMatcher matcher2 = CreateScreenTurnedOffAtomMatcher();
+    int64_t matcher2Id = matcher2.id();
+    *config.add_atom_matcher() = matcher2;
+
+    Predicate predicate1 = CreateScreenIsOnPredicate();
+    int64_t predicate1Id = predicate1.id();
+    *config.add_predicate() = predicate1;
+
+    // Add a few count metrics.
+    // Will be preserved.
+    CountMetric countMetric = createCountMetric("COUNT1", matcher1Id, predicate1Id, {});
+    int64_t countMetricId = countMetric.id();
+    *config.add_count_metric() = countMetric;
+
+    // Will be replaced since matcher2 is replaced.
+    EventMetric eventMetric = createEventMetric("EVENT1", matcher2Id, nullopt);
+    int64_t eventMetricId = eventMetric.id();
+    *config.add_event_metric() = eventMetric;
+
+    EXPECT_TRUE(initConfig(config));
+
+    // Used later to ensure the condition wizard is replaced. Get it before doing the update.
+    sp<ConditionWizard> oldConditionWizard = oldMetricProducers[0]->mWizard;
+    EXPECT_EQ(oldConditionWizard->getStrongCount(), 3);
+
+    // Mark matcher 2 as replaced. Causes eventMetric to be replaced.
+    set<int64_t> replacedMatchers;
+    replacedMatchers.insert(matcher2Id);
+
+    // Map the matchers and predicates in reverse order to force the indices to change.
+    std::unordered_map<int64_t, int> newAtomMatchingTrackerMap;
+    const int matcher2Index = 0;
+    newAtomMatchingTrackerMap[matcher2Id] = 0;
+    const int matcher1Index = 1;
+    newAtomMatchingTrackerMap[matcher1Id] = 1;
+    // Use the existing matchers. A bit hacky, but saves code and we don't rely on them.
+    vector<sp<AtomMatchingTracker>> newAtomMatchingTrackers(2);
+    std::reverse_copy(oldAtomMatchingTrackers.begin(), oldAtomMatchingTrackers.end(),
+                      newAtomMatchingTrackers.begin());
+
+    std::unordered_map<int64_t, int> newConditionTrackerMap;
+    const int predicate1Index = 0;
+    newConditionTrackerMap[predicate1Id] = 0;
+    // Use the existing conditionTrackers. A bit hacky, but saves code and we don't rely on them.
+    vector<sp<ConditionTracker>> newConditionTrackers(1);
+    std::reverse_copy(oldConditionTrackers.begin(), oldConditionTrackers.end(),
+                      newConditionTrackers.begin());
+    vector<ConditionState> conditionCache = {ConditionState::kUnknown};
+
+    StatsdConfig newConfig;
+    *newConfig.add_count_metric() = countMetric;
+    const int countMetricIndex = 0;
+    *newConfig.add_event_metric() = eventMetric;
+    const int eventMetricIndex = 1;
+
+    // Output data structures to validate.
+    unordered_map<int64_t, int> newMetricProducerMap;
+    vector<sp<MetricProducer>> newMetricProducers;
+    unordered_map<int, vector<int>> conditionToMetricMap;
+    unordered_map<int, vector<int>> trackerToMetricMap;
+    set<int64_t> noReportMetricIds;
+    unordered_map<int, vector<int>> activationAtomTrackerToMetricMap;
+    unordered_map<int, vector<int>> deactivationAtomTrackerToMetricMap;
+    vector<int> metricsWithActivation;
+    EXPECT_TRUE(updateMetrics(
+            key, newConfig, /*timeBaseNs=*/123, /*currentTimeNs=*/12345, new StatsPullerManager(),
+            oldAtomMatchingTrackerMap, newAtomMatchingTrackerMap, replacedMatchers,
+            newAtomMatchingTrackers, newConditionTrackerMap, /*replacedConditions=*/{},
+            newConditionTrackers, conditionCache, /*stateAtomIdMap*/ {}, /*allStateGroupMaps=*/{},
+            /*replacedStates=*/{}, oldMetricProducerMap, oldMetricProducers, newMetricProducerMap,
+            newMetricProducers, conditionToMetricMap, trackerToMetricMap, noReportMetricIds,
+            activationAtomTrackerToMetricMap, deactivationAtomTrackerToMetricMap,
+            metricsWithActivation));
+
+    unordered_map<int64_t, int> expectedMetricProducerMap = {
+            {countMetricId, countMetricIndex},
+            {eventMetricId, eventMetricIndex},
+    };
+    EXPECT_THAT(newMetricProducerMap, ContainerEq(expectedMetricProducerMap));
+
+    // Make sure preserved metrics are the same.
+    ASSERT_EQ(newMetricProducers.size(), 2);
+    EXPECT_EQ(oldMetricProducers[oldMetricProducerMap.at(countMetricId)],
+              newMetricProducers[newMetricProducerMap.at(countMetricId)]);
+
+    // Make sure replaced metrics are different.
+    EXPECT_NE(oldMetricProducers[oldMetricProducerMap.at(eventMetricId)],
+              newMetricProducers[newMetricProducerMap.at(eventMetricId)]);
+
+    // Verify the conditionToMetricMap.
+    ASSERT_EQ(conditionToMetricMap.size(), 1);
+    const vector<int>& condition1Metrics = conditionToMetricMap[predicate1Index];
+    EXPECT_THAT(condition1Metrics, UnorderedElementsAre(countMetricIndex));
+
+    // Verify the trackerToMetricMap.
+    ASSERT_EQ(trackerToMetricMap.size(), 2);
+    const vector<int>& matcher1Metrics = trackerToMetricMap[matcher1Index];
+    EXPECT_THAT(matcher1Metrics, UnorderedElementsAre(countMetricIndex));
+    const vector<int>& matcher2Metrics = trackerToMetricMap[matcher2Index];
+    EXPECT_THAT(matcher2Metrics, UnorderedElementsAre(eventMetricIndex));
+
+    // Verify event activation/deactivation maps.
+    ASSERT_EQ(activationAtomTrackerToMetricMap.size(), 0);
+    ASSERT_EQ(deactivationAtomTrackerToMetricMap.size(), 0);
+    ASSERT_EQ(metricsWithActivation.size(), 0);
+
+    // Verify tracker indices/ids/conditions are correct.
+    EXPECT_EQ(newMetricProducers[countMetricIndex]->getMetricId(), countMetricId);
+    EXPECT_EQ(newMetricProducers[countMetricIndex]->mConditionTrackerIndex, predicate1Index);
+    EXPECT_EQ(newMetricProducers[countMetricIndex]->mCondition, ConditionState::kUnknown);
+    EXPECT_EQ(newMetricProducers[eventMetricIndex]->getMetricId(), eventMetricId);
+    EXPECT_EQ(newMetricProducers[eventMetricIndex]->mConditionTrackerIndex, -1);
+    EXPECT_EQ(newMetricProducers[eventMetricIndex]->mCondition, ConditionState::kTrue);
+
+    sp<ConditionWizard> newConditionWizard = newMetricProducers[0]->mWizard;
+    EXPECT_NE(newConditionWizard, oldConditionWizard);
+    EXPECT_EQ(newConditionWizard->getStrongCount(), 3);
+    oldMetricProducers.clear();
+    // Only reference to the old wizard should be the one in the test.
+    EXPECT_EQ(oldConditionWizard->getStrongCount(), 1);
+}
 }  // namespace statsd
 }  // namespace os
 }  // namespace android