Slice by state in CountMetricProducer

- When a log event is matched to the "what" atom of a metric,
MetricProducer parses the event for state primary key(s) and queries
StateTracker for the correct state values
- CountMetricProducer writes state information to the dump report proto
- Added e2e CountMetric tests that push events to StatsLogProcessor and
check that the resulting dump report is correct

Test: bit statsd_test:*
Bug: 136566566
Bug: 142124705
Change-Id: I9ca5b097151c0c73a617f3b16c8c3584b068cd82
diff --git a/cmds/statsd/src/HashableDimensionKey.cpp b/cmds/statsd/src/HashableDimensionKey.cpp
index 5e156bb..f9f11b2 100644
--- a/cmds/statsd/src/HashableDimensionKey.cpp
+++ b/cmds/statsd/src/HashableDimensionKey.cpp
@@ -59,10 +59,11 @@
     return JenkinsHashWhiten(hash);
 }
 
-bool filterValues(const Matcher& matcherField, const vector<FieldValue>& values, Value* output) {
+bool filterValues(const Matcher& matcherField, const vector<FieldValue>& values,
+                  FieldValue* output) {
     for (const auto& value : values) {
         if (value.mField.matches(matcherField)) {
-            (*output) = value.mValue;
+            (*output) = value;
             return true;
         }
     }
@@ -106,15 +107,34 @@
 
     size_t count = conditionDimension->getValues().size();
     if (count != links.conditionFields.size()) {
-        // ALOGE("WTF condition link is bad");
         return;
     }
 
     for (size_t i = 0; i < count; i++) {
         conditionDimension->mutableValue(i)->mField.setField(
-            links.conditionFields[i].mMatcher.getField());
+                links.conditionFields[i].mMatcher.getField());
         conditionDimension->mutableValue(i)->mField.setTag(
-            links.conditionFields[i].mMatcher.getTag());
+                links.conditionFields[i].mMatcher.getTag());
+    }
+}
+
+void getDimensionForState(const std::vector<FieldValue>& eventValues, const Metric2State& link,
+                          HashableDimensionKey* statePrimaryKey) {
+    // First, get the dimension from the event using the "what" fields from the
+    // MetricStateLinks.
+    filterValues(link.metricFields, eventValues, statePrimaryKey);
+
+    // Then check that the statePrimaryKey size equals the number of state fields
+    size_t count = statePrimaryKey->getValues().size();
+    if (count != link.stateFields.size()) {
+        return;
+    }
+
+    // For each dimension Value in the statePrimaryKey, set the field and tag
+    // using the state atom fields from MetricStateLinks.
+    for (size_t i = 0; i < count; i++) {
+        statePrimaryKey->mutableValue(i)->mField.setField(link.stateFields[i].mMatcher.getField());
+        statePrimaryKey->mutableValue(i)->mField.setTag(link.stateFields[i].mMatcher.getTag());
     }
 }
 
@@ -185,11 +205,11 @@
 
 bool MetricDimensionKey::operator==(const MetricDimensionKey& that) const {
     return mDimensionKeyInWhat == that.getDimensionKeyInWhat() &&
-           mDimensionKeyInCondition == that.getDimensionKeyInCondition();
+           mStateValuesKey == that.getStateValuesKey();
 };
 
 string MetricDimensionKey::toString() const {
-    return mDimensionKeyInWhat.toString() + mDimensionKeyInCondition.toString();
+    return mDimensionKeyInWhat.toString() + mStateValuesKey.toString();
 }
 
 bool MetricDimensionKey::operator<(const MetricDimensionKey& that) const {
@@ -199,7 +219,7 @@
         return false;
     }
 
-    return mDimensionKeyInCondition < that.getDimensionKeyInCondition();
+    return mStateValuesKey < that.getStateValuesKey();
 }
 
 }  // namespace statsd
diff --git a/cmds/statsd/src/HashableDimensionKey.h b/cmds/statsd/src/HashableDimensionKey.h
index a123850..b9b86ce 100644
--- a/cmds/statsd/src/HashableDimensionKey.h
+++ b/cmds/statsd/src/HashableDimensionKey.h
@@ -34,6 +34,12 @@
     std::vector<Matcher> conditionFields;
 };
 
+struct Metric2State {
+    int32_t stateAtomId;
+    std::vector<Matcher> metricFields;
+    std::vector<Matcher> stateFields;
+};
+
 class HashableDimensionKey {
 public:
     explicit HashableDimensionKey(const std::vector<FieldValue>& values) {
@@ -76,17 +82,16 @@
 };
 
 class MetricDimensionKey {
- public:
+public:
     explicit MetricDimensionKey(const HashableDimensionKey& dimensionKeyInWhat,
-                                const HashableDimensionKey& dimensionKeyInCondition)
-        : mDimensionKeyInWhat(dimensionKeyInWhat),
-          mDimensionKeyInCondition(dimensionKeyInCondition) {};
+                                const HashableDimensionKey& stateValuesKey)
+        : mDimensionKeyInWhat(dimensionKeyInWhat), mStateValuesKey(stateValuesKey){};
 
     MetricDimensionKey(){};
 
     MetricDimensionKey(const MetricDimensionKey& that)
         : mDimensionKeyInWhat(that.getDimensionKeyInWhat()),
-          mDimensionKeyInCondition(that.getDimensionKeyInCondition()) {};
+          mStateValuesKey(that.getStateValuesKey()){};
 
     MetricDimensionKey& operator=(const MetricDimensionKey& from) = default;
 
@@ -96,25 +101,25 @@
         return mDimensionKeyInWhat;
     }
 
-    inline const HashableDimensionKey& getDimensionKeyInCondition() const {
-        return mDimensionKeyInCondition;
+    inline const HashableDimensionKey& getStateValuesKey() const {
+        return mStateValuesKey;
     }
 
-    inline void setDimensionKeyInCondition(const HashableDimensionKey& key) {
-        mDimensionKeyInCondition = key;
+    inline void setStateValuesKey(const HashableDimensionKey& key) {
+        mStateValuesKey = key;
     }
 
-    bool hasDimensionKeyInCondition() const {
-        return mDimensionKeyInCondition.getValues().size() > 0;
+    bool hasStateValuesKey() const {
+        return mStateValuesKey.getValues().size() > 0;
     }
 
     bool operator==(const MetricDimensionKey& that) const;
 
     bool operator<(const MetricDimensionKey& that) const;
 
-  private:
-      HashableDimensionKey mDimensionKeyInWhat;
-      HashableDimensionKey mDimensionKeyInCondition;
+private:
+    HashableDimensionKey mDimensionKeyInWhat;
+    HashableDimensionKey mStateValuesKey;
 };
 
 android::hash_t hashDimension(const HashableDimensionKey& key);
@@ -124,7 +129,7 @@
  * The value of the FieldValue is output.
  */
 bool filterValues(const Matcher& matcherField, const std::vector<FieldValue>& values,
-                  Value* output);
+                  FieldValue* output);
 
 /**
  * Creating HashableDimensionKeys from FieldValues using matcher.
@@ -152,6 +157,13 @@
                               const Metric2Condition& links,
                               HashableDimensionKey* conditionDimension);
 
+/**
+ * Get dimension values using metric's "what" fields and fill statePrimaryKey's
+ * mField information using "state" fields.
+ */
+void getDimensionForState(const std::vector<FieldValue>& eventValues, const Metric2State& link,
+                          HashableDimensionKey* statePrimaryKey);
+
 }  // namespace statsd
 }  // namespace os
 }  // namespace android
@@ -172,7 +184,7 @@
 struct hash<MetricDimensionKey> {
     std::size_t operator()(const MetricDimensionKey& key) const {
         android::hash_t hash = hashDimension(key.getDimensionKeyInWhat());
-        hash = android::JenkinsHashMix(hash, hashDimension(key.getDimensionKeyInCondition()));
+        hash = android::JenkinsHashMix(hash, hashDimension(key.getStateValuesKey()));
         return android::JenkinsHashWhiten(hash);
     }
 };
diff --git a/cmds/statsd/src/StatsLogProcessor.h b/cmds/statsd/src/StatsLogProcessor.h
index 17f2770..b41771d 100644
--- a/cmds/statsd/src/StatsLogProcessor.h
+++ b/cmds/statsd/src/StatsLogProcessor.h
@@ -263,9 +263,10 @@
     FRIEND_TEST(MetricActivationE2eTest, TestCountMetricWithSameDeactivation);
     FRIEND_TEST(MetricActivationE2eTest, TestCountMetricWithTwoMetricsTwoDeactivations);
 
-    FRIEND_TEST(CountMetricE2eTest, TestWithSimpleState);
-    FRIEND_TEST(CountMetricE2eTest, TestWithMappedState);
-    FRIEND_TEST(CountMetricE2eTest, TestWithMultipleStates);
+    FRIEND_TEST(CountMetricE2eTest, TestSlicedState);
+    FRIEND_TEST(CountMetricE2eTest, TestSlicedStateWithMap);
+    FRIEND_TEST(CountMetricE2eTest, TestMultipleSlicedStates);
+    FRIEND_TEST(CountMetricE2eTest, TestSlicedStateWithPrimaryFields);
 
     FRIEND_TEST(DurationMetricE2eTest, TestOneBucket);
     FRIEND_TEST(DurationMetricE2eTest, TestTwoBuckets);
diff --git a/cmds/statsd/src/metrics/CountMetricProducer.cpp b/cmds/statsd/src/metrics/CountMetricProducer.cpp
index 4a06387..c29b32c 100644
--- a/cmds/statsd/src/metrics/CountMetricProducer.cpp
+++ b/cmds/statsd/src/metrics/CountMetricProducer.cpp
@@ -57,10 +57,9 @@
 const int FIELD_ID_DATA = 1;
 // for CountMetricData
 const int FIELD_ID_DIMENSION_IN_WHAT = 1;
-const int FIELD_ID_DIMENSION_IN_CONDITION = 2;
+const int FIELD_ID_SLICE_BY_STATE = 6;
 const int FIELD_ID_BUCKET_INFO = 3;
 const int FIELD_ID_DIMENSION_LEAF_IN_WHAT = 4;
-const int FIELD_ID_DIMENSION_LEAF_IN_CONDITION = 5;
 // for CountBucketInfo
 const int FIELD_ID_COUNT = 3;
 const int FIELD_ID_BUCKET_NUM = 4;
@@ -102,7 +101,13 @@
         mConditionSliced = true;
     }
 
-    // TODO(tsaichristine): b/142124705 handle metric state links
+    for (const auto& stateLink : metric.state_link()) {
+        Metric2State ms;
+        ms.stateAtomId = stateLink.state_atom_id();
+        translateFieldMatcher(stateLink.fields_in_what(), &ms.metricFields);
+        translateFieldMatcher(stateLink.fields_in_state(), &ms.stateFields);
+        mMetric2StateLinks.push_back(ms);
+    }
 
     flushIfNeededLocked(startTimeNs);
     // Adjust start for partial bucket
@@ -132,10 +137,9 @@
             (unsigned long)mCurrentSlicedCounter->size());
     if (verbose) {
         for (const auto& it : *mCurrentSlicedCounter) {
-            fprintf(out, "\t(what)%s\t(condition)%s  %lld\n",
-                it.first.getDimensionKeyInWhat().toString().c_str(),
-                it.first.getDimensionKeyInCondition().toString().c_str(),
-                (unsigned long long)it.second);
+            fprintf(out, "\t(what)%s\t(state)%s  %lld\n",
+                    it.first.getDimensionKeyInWhat().toString().c_str(),
+                    it.first.getStateValuesKey().toString().c_str(), (unsigned long long)it.second);
         }
     }
 }
@@ -196,22 +200,16 @@
                     FIELD_TYPE_MESSAGE | FIELD_ID_DIMENSION_IN_WHAT);
             writeDimensionToProto(dimensionKey.getDimensionKeyInWhat(), str_set, protoOutput);
             protoOutput->end(dimensionToken);
-
-            if (dimensionKey.hasDimensionKeyInCondition()) {
-                uint64_t dimensionInConditionToken = protoOutput->start(
-                        FIELD_TYPE_MESSAGE | FIELD_ID_DIMENSION_IN_CONDITION);
-                writeDimensionToProto(dimensionKey.getDimensionKeyInCondition(),
-                                      str_set, protoOutput);
-                protoOutput->end(dimensionInConditionToken);
-            }
         } else {
             writeDimensionLeafNodesToProto(dimensionKey.getDimensionKeyInWhat(),
                                            FIELD_ID_DIMENSION_LEAF_IN_WHAT, str_set, protoOutput);
-            if (dimensionKey.hasDimensionKeyInCondition()) {
-                writeDimensionLeafNodesToProto(dimensionKey.getDimensionKeyInCondition(),
-                                               FIELD_ID_DIMENSION_LEAF_IN_CONDITION,
-                                               str_set, protoOutput);
-            }
+        }
+        // Then fill slice_by_state.
+        for (auto state : dimensionKey.getStateValuesKey().getValues()) {
+            uint64_t stateToken = protoOutput->start(FIELD_TYPE_MESSAGE | FIELD_COUNT_REPEATED |
+                                                     FIELD_ID_SLICE_BY_STATE);
+            writeStateToProto(state, protoOutput);
+            protoOutput->end(stateToken);
         }
         // Then fill bucket_info (CountBucketInfo).
         for (const auto& bucket : counter.second) {
@@ -282,7 +280,7 @@
     int64_t eventTimeNs = event.GetElapsedTimestampNs();
     flushIfNeededLocked(eventTimeNs);
 
-    if (condition == false) {
+    if (!condition) {
         return;
     }
 
diff --git a/cmds/statsd/src/metrics/CountMetricProducer.h b/cmds/statsd/src/metrics/CountMetricProducer.h
index 61e0892..8b17d88 100644
--- a/cmds/statsd/src/metrics/CountMetricProducer.h
+++ b/cmds/statsd/src/metrics/CountMetricProducer.h
@@ -52,7 +52,7 @@
 
     virtual ~CountMetricProducer();
 
-    void onStateChanged(int atomId, const HashableDimensionKey& primaryKey, int oldState,
+    void onStateChanged(int32_t atomId, const HashableDimensionKey& primaryKey, int oldState,
                         int newState) override;
 
 protected:
diff --git a/cmds/statsd/src/metrics/DurationMetricProducer.cpp b/cmds/statsd/src/metrics/DurationMetricProducer.cpp
index ab2a1c3..fee5e6e 100644
--- a/cmds/statsd/src/metrics/DurationMetricProducer.cpp
+++ b/cmds/statsd/src/metrics/DurationMetricProducer.cpp
@@ -53,10 +53,8 @@
 const int FIELD_ID_DATA = 1;
 // for DurationMetricData
 const int FIELD_ID_DIMENSION_IN_WHAT = 1;
-const int FIELD_ID_DIMENSION_IN_CONDITION = 2;
 const int FIELD_ID_BUCKET_INFO = 3;
 const int FIELD_ID_DIMENSION_LEAF_IN_WHAT = 4;
-const int FIELD_ID_DIMENSION_LEAF_IN_CONDITION = 5;
 // for DurationBucketInfo
 const int FIELD_ID_DURATION = 3;
 const int FIELD_ID_BUCKET_NUM = 4;
@@ -120,9 +118,8 @@
     mUseWhatDimensionAsInternalDimension = equalDimensions(mDimensionsInWhat, mInternalDimensions);
     if (mWizard != nullptr && mConditionTrackerIndex >= 0 &&
             mMetric2ConditionLinks.size() == 1) {
-        mHasLinksToAllConditionDimensionsInTracker =
-            mWizard->equalOutputDimensions(mConditionTrackerIndex,
-                                           mMetric2ConditionLinks.begin()->conditionFields);
+        mHasLinksToAllConditionDimensionsInTracker = mWizard->equalOutputDimensions(
+                mConditionTrackerIndex, mMetric2ConditionLinks.begin()->conditionFields);
     }
     flushIfNeededLocked(startTimeNs);
     // Adjust start for partial bucket
@@ -206,8 +203,7 @@
         mWizard->getTrueSlicedDimensions(mConditionTrackerIndex, &trueConditionDimensions);
         for (auto& whatIt : mCurrentSlicedDurationTrackerMap) {
             HashableDimensionKey linkedConditionDimensionKey;
-            getDimensionForCondition(whatIt.first.getValues(),
-                                     mMetric2ConditionLinks[0],
+            getDimensionForCondition(whatIt.first.getValues(), mMetric2ConditionLinks[0],
                                      &linkedConditionDimensionKey);
             if (trueConditionDimensions.find(linkedConditionDimensionKey) !=
                     trueConditionDimensions.end()) {
@@ -222,8 +218,7 @@
         if (currentUnSlicedPartCondition) {
             for (auto& whatIt : mCurrentSlicedDurationTrackerMap) {
                 HashableDimensionKey linkedConditionDimensionKey;
-                getDimensionForCondition(whatIt.first.getValues(),
-                                         mMetric2ConditionLinks[0],
+                getDimensionForCondition(whatIt.first.getValues(), mMetric2ConditionLinks[0],
                                          &linkedConditionDimensionKey);
                 if (dimensionsChangedToTrue->find(linkedConditionDimensionKey) !=
                         dimensionsChangedToTrue->end()) {
@@ -380,22 +375,9 @@
                     FIELD_TYPE_MESSAGE | FIELD_ID_DIMENSION_IN_WHAT);
             writeDimensionToProto(dimensionKey.getDimensionKeyInWhat(), str_set, protoOutput);
             protoOutput->end(dimensionToken);
-
-            if (dimensionKey.hasDimensionKeyInCondition()) {
-                uint64_t dimensionInConditionToken = protoOutput->start(
-                        FIELD_TYPE_MESSAGE | FIELD_ID_DIMENSION_IN_CONDITION);
-                writeDimensionToProto(dimensionKey.getDimensionKeyInCondition(),
-                                      str_set, protoOutput);
-                protoOutput->end(dimensionInConditionToken);
-            }
         } else {
             writeDimensionLeafNodesToProto(dimensionKey.getDimensionKeyInWhat(),
                                            FIELD_ID_DIMENSION_LEAF_IN_WHAT, str_set, protoOutput);
-            if (dimensionKey.hasDimensionKeyInCondition()) {
-                writeDimensionLeafNodesToProto(dimensionKey.getDimensionKeyInCondition(),
-                                               FIELD_ID_DIMENSION_LEAF_IN_CONDITION,
-                                               str_set, protoOutput);
-            }
         }
         // Then fill bucket_info (DurationBucketInfo).
         for (const auto& bucket : pair.second) {
@@ -472,7 +454,7 @@
     if (verbose) {
         for (const auto& whatIt : mCurrentSlicedDurationTrackerMap) {
             for (const auto& slice : whatIt.second) {
-                fprintf(out, "\t(what)%s\t(condition)%s\n", whatIt.first.toString().c_str(),
+                fprintf(out, "\t(what)%s\t(states)%s\n", whatIt.first.toString().c_str(),
                         slice.first.toString().c_str());
                 slice.second->dumpStates(out, verbose);
             }
@@ -483,8 +465,8 @@
 bool DurationMetricProducer::hitGuardRailLocked(const MetricDimensionKey& newKey) {
     auto whatIt = mCurrentSlicedDurationTrackerMap.find(newKey.getDimensionKeyInWhat());
     if (whatIt != mCurrentSlicedDurationTrackerMap.end()) {
-        auto condIt = whatIt->second.find(newKey.getDimensionKeyInCondition());
-        if (condIt != whatIt->second.end()) {
+        auto stateIt = whatIt->second.find(newKey.getStateValuesKey());
+        if (stateIt != whatIt->second.end()) {
             return false;
         }
         if (whatIt->second.size() > StatsdStats::kDimensionKeySizeSoftLimit - 1) {
@@ -493,8 +475,8 @@
                     mConfigKey, mMetricId, newTupleCount);
             // 2. Don't add more tuples, we are above the allowed threshold. Drop the data.
             if (newTupleCount > StatsdStats::kDimensionKeySizeHardLimit) {
-                ALOGE("DurationMetric %lld dropping data for condition dimension key %s",
-                    (long long)mMetricId, newKey.getDimensionKeyInCondition().toString().c_str());
+                ALOGE("DurationMetric %lld dropping data for state values key %s",
+                      (long long)mMetricId, newKey.getStateValuesKey().toString().c_str());
                 StatsdStats::getInstance().noteHardDimensionLimitReached(mMetricId);
                 return true;
             }
@@ -521,24 +503,24 @@
                                               const ConditionKey& conditionKeys,
                                               bool condition, const LogEvent& event) {
     const auto& whatKey = eventKey.getDimensionKeyInWhat();
-    const auto& condKey = eventKey.getDimensionKeyInCondition();
+    const auto& stateKey = eventKey.getStateValuesKey();
 
     auto whatIt = mCurrentSlicedDurationTrackerMap.find(whatKey);
     if (whatIt == mCurrentSlicedDurationTrackerMap.end()) {
         if (hitGuardRailLocked(eventKey)) {
             return;
         }
-        mCurrentSlicedDurationTrackerMap[whatKey][condKey] = createDurationTracker(eventKey);
+        mCurrentSlicedDurationTrackerMap[whatKey][stateKey] = createDurationTracker(eventKey);
     } else {
-        if (whatIt->second.find(condKey) == whatIt->second.end()) {
+        if (whatIt->second.find(stateKey) == whatIt->second.end()) {
             if (hitGuardRailLocked(eventKey)) {
                 return;
             }
-            mCurrentSlicedDurationTrackerMap[whatKey][condKey] = createDurationTracker(eventKey);
+            mCurrentSlicedDurationTrackerMap[whatKey][stateKey] = createDurationTracker(eventKey);
         }
     }
 
-    auto it = mCurrentSlicedDurationTrackerMap.find(whatKey)->second.find(condKey);
+    auto it = mCurrentSlicedDurationTrackerMap.find(whatKey)->second.find(stateKey);
     if (mUseWhatDimensionAsInternalDimension) {
         it->second->noteStart(whatKey, condition,
                               event.GetElapsedTimestampNs(), conditionKeys);
@@ -597,8 +579,8 @@
         if (mUseWhatDimensionAsInternalDimension) {
             auto whatIt = mCurrentSlicedDurationTrackerMap.find(dimensionInWhat);
             if (whatIt != mCurrentSlicedDurationTrackerMap.end()) {
-                for (const auto& condIt : whatIt->second) {
-                    condIt.second->noteStop(dimensionInWhat, event.GetElapsedTimestampNs(), false);
+                for (const auto& stateIt : whatIt->second) {
+                    stateIt.second->noteStop(dimensionInWhat, event.GetElapsedTimestampNs(), false);
                 }
             }
             return;
@@ -611,9 +593,9 @@
 
         auto whatIt = mCurrentSlicedDurationTrackerMap.find(dimensionInWhat);
         if (whatIt != mCurrentSlicedDurationTrackerMap.end()) {
-            for (const auto& condIt : whatIt->second) {
-                condIt.second->noteStop(
-                    internalDimensionKey, event.GetElapsedTimestampNs(), false);
+            for (const auto& stateIt : whatIt->second) {
+                stateIt.second->noteStop(internalDimensionKey, event.GetElapsedTimestampNs(),
+                                         false);
             }
         }
         return;
diff --git a/cmds/statsd/src/metrics/GaugeMetricProducer.cpp b/cmds/statsd/src/metrics/GaugeMetricProducer.cpp
index d0f88a8..64344e8 100644
--- a/cmds/statsd/src/metrics/GaugeMetricProducer.cpp
+++ b/cmds/statsd/src/metrics/GaugeMetricProducer.cpp
@@ -56,10 +56,8 @@
 const int FIELD_ID_SKIPPED_END_MILLIS = 4;
 // for GaugeMetricData
 const int FIELD_ID_DIMENSION_IN_WHAT = 1;
-const int FIELD_ID_DIMENSION_IN_CONDITION = 2;
 const int FIELD_ID_BUCKET_INFO = 3;
 const int FIELD_ID_DIMENSION_LEAF_IN_WHAT = 4;
-const int FIELD_ID_DIMENSION_LEAF_IN_CONDITION = 5;
 // for GaugeBucketInfo
 const int FIELD_ID_ATOM = 3;
 const int FIELD_ID_ELAPSED_ATOM_TIMESTAMP = 4;
@@ -166,10 +164,9 @@
             (unsigned long)mCurrentSlicedBucket->size());
     if (verbose) {
         for (const auto& it : *mCurrentSlicedBucket) {
-            fprintf(out, "\t(what)%s\t(condition)%s  %d atoms\n",
-                it.first.getDimensionKeyInWhat().toString().c_str(),
-                it.first.getDimensionKeyInCondition().toString().c_str(),
-                (int)it.second.size());
+            fprintf(out, "\t(what)%s\t(states)%s  %d atoms\n",
+                    it.first.getDimensionKeyInWhat().toString().c_str(),
+                    it.first.getStateValuesKey().toString().c_str(), (int)it.second.size());
         }
     }
 }
@@ -238,22 +235,9 @@
                     FIELD_TYPE_MESSAGE | FIELD_ID_DIMENSION_IN_WHAT);
             writeDimensionToProto(dimensionKey.getDimensionKeyInWhat(), str_set, protoOutput);
             protoOutput->end(dimensionToken);
-
-            if (dimensionKey.hasDimensionKeyInCondition()) {
-                uint64_t dimensionInConditionToken = protoOutput->start(
-                        FIELD_TYPE_MESSAGE | FIELD_ID_DIMENSION_IN_CONDITION);
-                writeDimensionToProto(dimensionKey.getDimensionKeyInCondition(),
-                                      str_set, protoOutput);
-                protoOutput->end(dimensionInConditionToken);
-            }
         } else {
             writeDimensionLeafNodesToProto(dimensionKey.getDimensionKeyInWhat(),
                                            FIELD_ID_DIMENSION_LEAF_IN_WHAT, str_set, protoOutput);
-            if (dimensionKey.hasDimensionKeyInCondition()) {
-                writeDimensionLeafNodesToProto(dimensionKey.getDimensionKeyInCondition(),
-                                               FIELD_ID_DIMENSION_LEAF_IN_CONDITION,
-                                               str_set, protoOutput);
-            }
         }
 
         // Then fill bucket_info (GaugeBucketInfo).
diff --git a/cmds/statsd/src/metrics/MetricProducer.cpp b/cmds/statsd/src/metrics/MetricProducer.cpp
index 2a700ef..2c8f0e3 100644
--- a/cmds/statsd/src/metrics/MetricProducer.cpp
+++ b/cmds/statsd/src/metrics/MetricProducer.cpp
@@ -16,8 +16,11 @@
 
 #define DEBUG false  // STOPSHIP if true
 #include "Log.h"
+
 #include "MetricProducer.h"
 
+#include "state/StateTracker.h"
+
 using android::util::FIELD_COUNT_REPEATED;
 using android::util::FIELD_TYPE_ENUM;
 using android::util::FIELD_TYPE_INT32;
@@ -92,9 +95,43 @@
         condition = mCondition == ConditionState::kTrue;
     }
 
+    // Stores atom id to primary key pairs for each state atom that the metric is
+    // sliced by.
+    std::map<int, HashableDimensionKey> statePrimaryKeys;
+
+    // For states with primary fields, use MetricStateLinks to get the primary
+    // field values from the log event. These values will form a primary key
+    // that will be used to query StateTracker for the correct state value.
+    for (const auto& stateLink : mMetric2StateLinks) {
+        getDimensionForState(event.getValues(), stateLink,
+                             &statePrimaryKeys[stateLink.stateAtomId]);
+    }
+
+    // For each sliced state, query StateTracker for the state value using
+    // either the primary key from the previous step or the DEFAULT_DIMENSION_KEY.
+    //
+    // Expected functionality: for any case where the MetricStateLinks are
+    // initialized incorrectly (ex. # of state links != # of primary fields, no
+    // links are provided for a state with primary fields, links are provided
+    // in the wrong order, etc.), StateTracker will simply return kStateUnknown
+    // when queried using an incorrect key.
+    HashableDimensionKey stateValuesKey;
+    for (auto atomId : mSlicedStateAtoms) {
+        FieldValue value;
+        if (statePrimaryKeys.find(atomId) != statePrimaryKeys.end()) {
+            // found a primary key for this state, query using the key
+            getMappedStateValue(atomId, statePrimaryKeys[atomId], &value);
+        } else {
+            // if no MetricStateLinks exist for this state atom,
+            // query using the default dimension key (empty HashableDimensionKey)
+            getMappedStateValue(atomId, DEFAULT_DIMENSION_KEY, &value);
+        }
+        stateValuesKey.addValue(value);
+    }
+
     HashableDimensionKey dimensionInWhat;
     filterValues(mDimensionsInWhat, event.getValues(), &dimensionInWhat);
-    MetricDimensionKey metricKey(dimensionInWhat, DEFAULT_DIMENSION_KEY);
+    MetricDimensionKey metricKey(dimensionInWhat, stateValuesKey);
     onMatchedLogEventInternalLocked(
             matcherIndex, metricKey, conditionKey, condition, event);
 }
@@ -227,6 +264,31 @@
     }
 }
 
+void MetricProducer::getMappedStateValue(const int32_t atomId, const HashableDimensionKey& queryKey,
+                                         FieldValue* value) {
+    if (!StateManager::getInstance().getStateValue(atomId, queryKey, value)) {
+        value->mValue = Value(StateTracker::kStateUnknown);
+        ALOGW("StateTracker not found for state atom %d", atomId);
+        return;
+    }
+
+    // check if there is a state map for this atom
+    auto atomIt = mStateGroupMap.find(atomId);
+    if (atomIt == mStateGroupMap.end()) {
+        return;
+    }
+    auto valueIt = atomIt->second.find(value->mValue.int_value);
+    if (valueIt == atomIt->second.end()) {
+        // state map exists, but value was not put in a state group
+        // so set mValue to kStateUnknown
+        // TODO(tsaichristine): handle incomplete state maps
+        value->mValue.setInt(StateTracker::kStateUnknown);
+    } else {
+        // set mValue to group_id
+        value->mValue.setLong(valueIt->second);
+    }
+}
+
 }  // namespace statsd
 }  // namespace os
 }  // namespace android
diff --git a/cmds/statsd/src/metrics/MetricProducer.h b/cmds/statsd/src/metrics/MetricProducer.h
index a72de22..d7cbcc8 100644
--- a/cmds/statsd/src/metrics/MetricProducer.h
+++ b/cmds/statsd/src/metrics/MetricProducer.h
@@ -30,6 +30,7 @@
 #include "matchers/matcher_util.h"
 #include "packages/PackageInfoListener.h"
 #include "state/StateListener.h"
+#include "state/StateManager.h"
 
 namespace android {
 namespace os {
@@ -340,6 +341,12 @@
         return (endNs - mTimeBaseNs) / mBucketSizeNs - 1;
     }
 
+    // Query StateManager for original state value.
+    // If no state map exists for this atom, return the original value.
+    // Otherwise, return the group_id mapped to the atom and original value.
+    void getMappedStateValue(const int32_t atomId, const HashableDimensionKey& queryKey,
+                             FieldValue* value);
+
     const int64_t mMetricId;
 
     const ConfigKey mConfigKey;
@@ -392,14 +399,19 @@
     bool mIsActive;
 
     // The slice_by_state atom ids defined in statsd_config.
-    std::vector<int> mSlicedStateAtoms;
+    std::vector<int32_t> mSlicedStateAtoms;
 
     // Maps atom ids and state values to group_ids (<atom_id, <value, group_id>>).
-    std::unordered_map<int, std::unordered_map<int, int64_t>> mStateGroupMap;
+    std::unordered_map<int32_t, std::unordered_map<int, int64_t>> mStateGroupMap;
 
-    FRIEND_TEST(CountMetricE2eTest, TestWithSimpleState);
-    FRIEND_TEST(CountMetricE2eTest, TestWithMappedState);
-    FRIEND_TEST(CountMetricE2eTest, TestWithMultipleStates);
+    // MetricStateLinks defined in statsd_config that link fields in the state
+    // atom to fields in the "what" atom.
+    std::vector<Metric2State> mMetric2StateLinks;
+
+    FRIEND_TEST(CountMetricE2eTest, TestSlicedState);
+    FRIEND_TEST(CountMetricE2eTest, TestSlicedStateWithMap);
+    FRIEND_TEST(CountMetricE2eTest, TestMultipleSlicedStates);
+    FRIEND_TEST(CountMetricE2eTest, TestSlicedStateWithPrimaryFields);
 
     FRIEND_TEST(DurationMetricE2eTest, TestOneBucket);
     FRIEND_TEST(DurationMetricE2eTest, TestTwoBuckets);
diff --git a/cmds/statsd/src/metrics/MetricsManager.h b/cmds/statsd/src/metrics/MetricsManager.h
index d184121..286610a 100644
--- a/cmds/statsd/src/metrics/MetricsManager.h
+++ b/cmds/statsd/src/metrics/MetricsManager.h
@@ -282,9 +282,10 @@
             TestActivationOnBootMultipleActivationsDifferentActivationTypes);
     FRIEND_TEST(StatsLogProcessorTest, TestActivationsPersistAcrossSystemServerRestart);
 
-    FRIEND_TEST(CountMetricE2eTest, TestWithSimpleState);
-    FRIEND_TEST(CountMetricE2eTest, TestWithMappedState);
-    FRIEND_TEST(CountMetricE2eTest, TestWithMultipleStates);
+    FRIEND_TEST(CountMetricE2eTest, TestSlicedState);
+    FRIEND_TEST(CountMetricE2eTest, TestSlicedStateWithMap);
+    FRIEND_TEST(CountMetricE2eTest, TestMultipleSlicedStates);
+    FRIEND_TEST(CountMetricE2eTest, TestSlicedStateWithPrimaryFields);
 
     FRIEND_TEST(DurationMetricE2eTest, TestOneBucket);
     FRIEND_TEST(DurationMetricE2eTest, TestTwoBuckets);
diff --git a/cmds/statsd/src/metrics/ValueMetricProducer.cpp b/cmds/statsd/src/metrics/ValueMetricProducer.cpp
index f53a745..eb78ebc5 100644
--- a/cmds/statsd/src/metrics/ValueMetricProducer.cpp
+++ b/cmds/statsd/src/metrics/ValueMetricProducer.cpp
@@ -59,10 +59,8 @@
 const int FIELD_ID_SKIPPED_END_MILLIS = 4;
 // for ValueMetricData
 const int FIELD_ID_DIMENSION_IN_WHAT = 1;
-const int FIELD_ID_DIMENSION_IN_CONDITION = 2;
 const int FIELD_ID_BUCKET_INFO = 3;
 const int FIELD_ID_DIMENSION_LEAF_IN_WHAT = 4;
-const int FIELD_ID_DIMENSION_LEAF_IN_CONDITION = 5;
 // for ValueBucketInfo
 const int FIELD_ID_VALUE_INDEX = 1;
 const int FIELD_ID_VALUE_LONG = 2;
@@ -266,21 +264,9 @@
                     protoOutput->start(FIELD_TYPE_MESSAGE | FIELD_ID_DIMENSION_IN_WHAT);
             writeDimensionToProto(dimensionKey.getDimensionKeyInWhat(), str_set, protoOutput);
             protoOutput->end(dimensionToken);
-            if (dimensionKey.hasDimensionKeyInCondition()) {
-                uint64_t dimensionInConditionToken =
-                        protoOutput->start(FIELD_TYPE_MESSAGE | FIELD_ID_DIMENSION_IN_CONDITION);
-                writeDimensionToProto(dimensionKey.getDimensionKeyInCondition(), str_set,
-                                      protoOutput);
-                protoOutput->end(dimensionInConditionToken);
-            }
         } else {
             writeDimensionLeafNodesToProto(dimensionKey.getDimensionKeyInWhat(),
                                            FIELD_ID_DIMENSION_LEAF_IN_WHAT, str_set, protoOutput);
-            if (dimensionKey.hasDimensionKeyInCondition()) {
-                writeDimensionLeafNodesToProto(dimensionKey.getDimensionKeyInCondition(),
-                                               FIELD_ID_DIMENSION_LEAF_IN_CONDITION, str_set,
-                                               protoOutput);
-            }
         }
 
         // Then fill bucket_info (ValueBucketInfo).
@@ -586,10 +572,10 @@
     if (verbose) {
         for (const auto& it : mCurrentSlicedBucket) {
           for (const auto& interval : it.second) {
-            fprintf(out, "\t(what)%s\t(condition)%s  (value)%s\n",
-                    it.first.getDimensionKeyInWhat().toString().c_str(),
-                    it.first.getDimensionKeyInCondition().toString().c_str(),
-                    interval.value.toString().c_str());
+              fprintf(out, "\t(what)%s\t(states)%s  (value)%s\n",
+                      it.first.getDimensionKeyInWhat().toString().c_str(),
+                      it.first.getStateValuesKey().toString().c_str(),
+                      interval.value.toString().c_str());
           }
         }
     }
diff --git a/cmds/statsd/src/metrics/metrics_manager_util.cpp b/cmds/statsd/src/metrics/metrics_manager_util.cpp
index 33e162e..6e76717 100644
--- a/cmds/statsd/src/metrics/metrics_manager_util.cpp
+++ b/cmds/statsd/src/metrics/metrics_manager_util.cpp
@@ -472,11 +472,13 @@
                                         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;
+            }
         }
 
-        // TODO(tsaichristine): add check for unequal number of MetricStateLinks
-        // and slice_by_states
-
         unordered_map<int, shared_ptr<Activation>> eventActivationMap;
         unordered_map<int, vector<shared_ptr<Activation>>> eventDeactivationMap;
         bool success = handleMetricActivation(config, metric.id(), metricIndex,
diff --git a/cmds/statsd/src/state/StateListener.h b/cmds/statsd/src/state/StateListener.h
index a31690a..f2b9a6b 100644
--- a/cmds/statsd/src/state/StateListener.h
+++ b/cmds/statsd/src/state/StateListener.h
@@ -43,8 +43,8 @@
      * [oldState]: Previous state value before state change
      * [newState]: Current state value after state change
      */
-    virtual void onStateChanged(int atomId, const HashableDimensionKey& primaryKey, int oldState,
-                                int newState) = 0;
+    virtual void onStateChanged(int32_t atomId, const HashableDimensionKey& primaryKey,
+                                int oldState, int newState) = 0;
 };
 
 }  // namespace statsd
diff --git a/cmds/statsd/src/state/StateManager.cpp b/cmds/statsd/src/state/StateManager.cpp
index 95b2c76..2fa28c9 100644
--- a/cmds/statsd/src/state/StateManager.cpp
+++ b/cmds/statsd/src/state/StateManager.cpp
@@ -35,7 +35,7 @@
     }
 }
 
-bool StateManager::registerListener(int atomId, wp<StateListener> listener) {
+bool StateManager::registerListener(int32_t atomId, wp<StateListener> listener) {
     std::lock_guard<std::mutex> lock(mMutex);
 
     // Check if state tracker already exists
@@ -53,7 +53,7 @@
     return true;
 }
 
-void StateManager::unregisterListener(int atomId, wp<StateListener> listener) {
+void StateManager::unregisterListener(int32_t atomId, wp<StateListener> listener) {
     std::unique_lock<std::mutex> lock(mMutex);
 
     // Hold the sp<> until the lock is released so that ~StateTracker() is
@@ -77,13 +77,15 @@
     lock.unlock();
 }
 
-int StateManager::getStateValue(int atomId, const HashableDimensionKey& key) {
+bool StateManager::getStateValue(int32_t atomId, const HashableDimensionKey& key,
+                                 FieldValue* output) const {
     std::lock_guard<std::mutex> lock(mMutex);
-    if (mStateTrackers.find(atomId) != mStateTrackers.end()) {
-        return mStateTrackers[atomId]->getStateValue(key);
-    }
 
-    return StateTracker::kStateUnknown;
+    auto it = mStateTrackers.find(atomId);
+    if (it != mStateTrackers.end()) {
+        return it->second->getStateValue(key, output);
+    }
+    return false;
 }
 
 }  // namespace statsd
diff --git a/cmds/statsd/src/state/StateManager.h b/cmds/statsd/src/state/StateManager.h
index 89ee6c0..272724c 100644
--- a/cmds/statsd/src/state/StateManager.h
+++ b/cmds/statsd/src/state/StateManager.h
@@ -42,26 +42,30 @@
     // Returns true if atomId is being tracked and is associated with a state
     // atom. StateManager notifies the correct StateTracker to register listener.
     // If the correct StateTracker does not exist, a new StateTracker is created.
-    bool registerListener(int atomId, wp<StateListener> listener);
+    bool registerListener(int32_t atomId, wp<StateListener> listener);
 
     // Notifies the correct StateTracker to unregister a listener
     // and removes the tracker if it no longer has any listeners.
-    void unregisterListener(int atomId, wp<StateListener> listener);
+    void unregisterListener(int32_t atomId, wp<StateListener> listener);
 
-    // Queries the correct StateTracker for the original/un-mapped state value
-    // that is mapped to the given query key.
-    // If the StateTracker doesn't exist, returns StateTracker::kStateUnknown.
-    int getStateValue(int atomId, const HashableDimensionKey& queryKey);
+    // Returns true if the StateTracker exists and queries for the
+    // original state value mapped to the given query key. The state value is
+    // stored and output in a FieldValue class.
+    // Returns false if the StateTracker doesn't exist.
+    bool getStateValue(int32_t atomId, const HashableDimensionKey& queryKey,
+                       FieldValue* output) const;
 
-    inline int getStateTrackersCount() {
+    inline int getStateTrackersCount() const {
         std::lock_guard<std::mutex> lock(mMutex);
         return mStateTrackers.size();
     }
 
-    inline int getListenersCount(int atomId) {
+    inline int getListenersCount(int32_t atomId) const {
         std::lock_guard<std::mutex> lock(mMutex);
-        if (mStateTrackers.find(atomId) != mStateTrackers.end()) {
-            return mStateTrackers[atomId]->getListenersCount();
+
+        auto it = mStateTrackers.find(atomId);
+        if (it != mStateTrackers.end()) {
+            return it->second->getListenersCount();
         }
         return -1;
     }
@@ -70,7 +74,7 @@
   mutable std::mutex mMutex;
 
   // Maps state atom ids to StateTrackers
-  std::unordered_map<int, sp<StateTracker>> mStateTrackers;
+  std::unordered_map<int32_t, sp<StateTracker>> mStateTrackers;
 };
 
 }  // namespace statsd
diff --git a/cmds/statsd/src/state/StateTracker.cpp b/cmds/statsd/src/state/StateTracker.cpp
index 323fc0e..e6f61226 100644
--- a/cmds/statsd/src/state/StateTracker.cpp
+++ b/cmds/statsd/src/state/StateTracker.cpp
@@ -25,10 +25,8 @@
 namespace os {
 namespace statsd {
 
-StateTracker::StateTracker(const int atomId,
-                           const util::StateAtomFieldOptions& stateAtomInfo)
-  : mAtomId(atomId),
-    mStateField(getSimpleMatcher(atomId, stateAtomInfo.exclusiveField)) {
+StateTracker::StateTracker(const int32_t atomId, const util::StateAtomFieldOptions& stateAtomInfo)
+    : mAtomId(atomId), mStateField(getSimpleMatcher(atomId, stateAtomInfo.exclusiveField)) {
     // create matcher for each primary field
     // TODO(tsaichristine): b/142108433 handle when primary field is first uid in chain
     for (const auto& primary : stateAtomInfo.primaryFields) {
@@ -55,24 +53,26 @@
     }
 
     // parse event for state value
-    Value state;
-    int32_t stateValue;
-    if (!filterValues(mStateField, event.getValues(), &state) || state.getType() != INT) {
-        ALOGE("StateTracker error extracting state from log event. Type: %d", state.getType());
+    FieldValue stateValue;
+    int32_t state;
+    if (!filterValues(mStateField, event.getValues(), &stateValue) ||
+        stateValue.mValue.getType() != INT) {
+        ALOGE("StateTracker error extracting state from log event. Type: %d",
+              stateValue.mValue.getType());
         handlePartialReset(primaryKey);
         return;
     }
-    stateValue = state.int_value;
+    state = stateValue.mValue.int_value;
 
-    if (stateValue == mResetState) {
-        VLOG("StateTracker Reset state: %s", state.toString().c_str());
+    if (state == mResetState) {
+        VLOG("StateTracker Reset state: %s", stateValue.mValue.toString().c_str());
         handleReset();
     }
 
     // track and update state
     int32_t oldState = 0;
     int32_t newState = 0;
-    updateState(primaryKey, stateValue, &oldState, &newState);
+    updateState(primaryKey, state, &oldState, &newState);
 
     // notify all listeners if state has changed
     if (oldState != newState) {
@@ -96,18 +96,27 @@
     mListeners.erase(listener);
 }
 
-int StateTracker::getStateValue(const HashableDimensionKey& queryKey) const {
+bool StateTracker::getStateValue(const HashableDimensionKey& queryKey, FieldValue* output) const {
+    output->mField = mStateField.mMatcher;
+
+    // Check that the query key has the correct number of primary fields.
     if (queryKey.getValues().size() == mPrimaryFields.size()) {
         auto it = mStateMap.find(queryKey);
         if (it != mStateMap.end()) {
-            return it->second.state;
+            output->mValue = it->second.state;
+            return true;
         }
     } else if (queryKey.getValues().size() > mPrimaryFields.size()) {
         ALOGE("StateTracker query key size > primary key size is illegal");
     } else {
         ALOGE("StateTracker query key size < primary key size is not supported");
     }
-    return mDefaultState;
+
+    // Set the state value to unknown if:
+    // - query key size is incorrect
+    // - query key is not found in state map
+    output->mValue = StateTracker::kStateUnknown;
+    return false;
 }
 
 void StateTracker::handleReset() {
diff --git a/cmds/statsd/src/state/StateTracker.h b/cmds/statsd/src/state/StateTracker.h
index cfa9fd8..450412d 100644
--- a/cmds/statsd/src/state/StateTracker.h
+++ b/cmds/statsd/src/state/StateTracker.h
@@ -30,7 +30,7 @@
 
 class StateTracker : public virtual RefBase {
 public:
-    StateTracker(const int atomId, const util::StateAtomFieldOptions& stateAtomInfo);
+    StateTracker(const int32_t atomId, const util::StateAtomFieldOptions& stateAtomInfo);
 
     virtual ~StateTracker(){};
 
@@ -45,10 +45,12 @@
 
     void unregisterListener(wp<StateListener> listener);
 
-    // Returns the state value mapped to the given query key.
+    // The output is a FieldValue object that has mStateField as the field and
+    // the original state value (found using the given query key) as the value.
+    //
     // If the key isn't mapped to a state or the key size doesn't match the
-    // primary key size, the default state is returned.
-    int getStateValue(const HashableDimensionKey& queryKey) const;
+    // number of primary fields, the output value is set to kStateUnknown.
+    bool getStateValue(const HashableDimensionKey& queryKey, FieldValue* output) const;
 
     inline int getListenersCount() const {
         return mListeners.size();
diff --git a/cmds/statsd/src/stats_log_util.cpp b/cmds/statsd/src/stats_log_util.cpp
index c22e3cc..76c1936 100644
--- a/cmds/statsd/src/stats_log_util.cpp
+++ b/cmds/statsd/src/stats_log_util.cpp
@@ -53,6 +53,11 @@
 
 const int DIMENSIONS_VALUE_TUPLE_VALUE = 1;
 
+// for StateValue Proto
+const int STATE_VALUE_ATOM_ID = 1;
+const int STATE_VALUE_CONTENTS_GROUP_ID = 2;
+const int STATE_VALUE_CONTENTS_VALUE = 3;
+
 // for PulledAtomStats proto
 const int FIELD_ID_PULLED_ATOM_STATS = 10;
 const int FIELD_ID_PULL_ATOM_ID = 1;
@@ -416,6 +421,23 @@
     protoOutput->end(atomToken);
 }
 
+void writeStateToProto(const FieldValue& state, util::ProtoOutputStream* protoOutput) {
+    protoOutput->write(FIELD_TYPE_INT32 | STATE_VALUE_ATOM_ID, state.mField.getTag());
+
+    switch (state.mValue.getType()) {
+        case INT:
+            protoOutput->write(FIELD_TYPE_INT32 | STATE_VALUE_CONTENTS_VALUE,
+                               state.mValue.int_value);
+            break;
+        case LONG:
+            protoOutput->write(FIELD_TYPE_INT64 | STATE_VALUE_CONTENTS_GROUP_ID,
+                               state.mValue.long_value);
+            break;
+        default:
+            break;
+    }
+}
+
 int64_t TimeUnitToBucketSizeInMillisGuardrailed(int uid, TimeUnit unit) {
     int64_t bucketSizeMillis = TimeUnitToBucketSizeInMillis(unit);
     if (bucketSizeMillis > 1000 && bucketSizeMillis < 5 * 60 * 1000LL && uid != AID_SHELL &&
diff --git a/cmds/statsd/src/stats_log_util.h b/cmds/statsd/src/stats_log_util.h
index bfb84cf..0a86363a 100644
--- a/cmds/statsd/src/stats_log_util.h
+++ b/cmds/statsd/src/stats_log_util.h
@@ -40,6 +40,8 @@
 void writeDimensionPathToProto(const std::vector<Matcher>& fieldMatchers,
                                util::ProtoOutputStream* protoOutput);
 
+void writeStateToProto(const FieldValue& state, util::ProtoOutputStream* protoOutput);
+
 // Convert the TimeUnit enum to the bucket size in millis with a guardrail on
 // bucket size.
 int64_t TimeUnitToBucketSizeInMillisGuardrailed(int uid, TimeUnit unit);
diff --git a/cmds/statsd/src/statsd_config.proto b/cmds/statsd/src/statsd_config.proto
index 0664867..a22805b 100644
--- a/cmds/statsd/src/statsd_config.proto
+++ b/cmds/statsd/src/statsd_config.proto
@@ -178,7 +178,7 @@
 }
 
 message MetricStateLink {
-  optional int64 state = 1;
+  optional int32 state_atom_id = 1;
 
   optional FieldMatcher fields_in_what = 2;
 
diff --git a/cmds/statsd/src/subscriber/IncidentdReporter.cpp b/cmds/statsd/src/subscriber/IncidentdReporter.cpp
index f2c6f1a..f1320c2 100644
--- a/cmds/statsd/src/subscriber/IncidentdReporter.cpp
+++ b/cmds/statsd/src/subscriber/IncidentdReporter.cpp
@@ -52,7 +52,6 @@
 const int FIELD_ID_TRIGGER_DETAILS_TRIGGER_METRIC = 1;
 const int FIELD_ID_METRIC_VALUE_METRIC_ID = 1;
 const int FIELD_ID_METRIC_VALUE_DIMENSION_IN_WHAT = 2;
-const int FIELD_ID_METRIC_VALUE_DIMENSION_IN_CONDITION = 3;
 const int FIELD_ID_METRIC_VALUE_VALUE = 4;
 
 const int FIELD_ID_PACKAGE_INFO = 3;
@@ -84,10 +83,8 @@
     writeDimensionToProto(dimensionKey.getDimensionKeyInWhat(), nullptr, &headerProto);
     headerProto.end(dimToken);
 
+    // deprecated field
     // optional DimensionsValue dimension_in_condition = 3;
-    dimToken = headerProto.start(FIELD_TYPE_MESSAGE | FIELD_ID_METRIC_VALUE_DIMENSION_IN_CONDITION);
-    writeDimensionToProto(dimensionKey.getDimensionKeyInCondition(), nullptr, &headerProto);
-    headerProto.end(dimToken);
 
     // optional int64 value = 4;
     headerProto.write(FIELD_TYPE_INT64 | FIELD_ID_METRIC_VALUE_VALUE, (long long)metricValue);
@@ -106,13 +103,6 @@
         }
     }
 
-    for (const auto& dim : dimensionKey.getDimensionKeyInCondition().getValues()) {
-        int uid = getUidIfExists(dim);
-        if (uid > 2000) {
-            uids.insert(uid);
-        }
-    }
-
     if (!uids.empty()) {
         uint64_t token = headerProto.start(FIELD_TYPE_MESSAGE | FIELD_ID_PACKAGE_INFO);
         UidMap::getInstance()->writeUidMapSnapshot(getElapsedRealtimeNs(), true, true, uids,
diff --git a/cmds/statsd/tests/e2e/CountMetric_e2e_test.cpp b/cmds/statsd/tests/e2e/CountMetric_e2e_test.cpp
index 6591d69..0f51c1b 100644
--- a/cmds/statsd/tests/e2e/CountMetric_e2e_test.cpp
+++ b/cmds/statsd/tests/e2e/CountMetric_e2e_test.cpp
@@ -18,6 +18,7 @@
 
 #include "src/StatsLogProcessor.h"
 #include "src/state/StateManager.h"
+#include "src/state/StateTracker.h"
 #include "tests/statsd_test_util.h"
 
 namespace android {
@@ -26,8 +27,20 @@
 
 #ifdef __ANDROID__
 
-TEST(CountMetricE2eTest, TestWithSimpleState) {
-    // Initialize config
+const int SCREEN_STATE_ATOM_ID = android::util::SCREEN_STATE_CHANGED;
+const int UID_PROCESS_STATE_ATOM_ID = android::util::UID_PROCESS_STATE_CHANGED;
+
+/**
+ * Test a count metric that has one slice_by_state with no primary fields.
+ *
+ * Once the CountMetricProducer is initialized, it has one atom id in
+ * mSlicedStateAtoms and no entries in mStateGroupMap.
+
+ * One StateTracker tracks the state atom, and it has one listener which is the
+ * CountMetricProducer that was initialized.
+ */
+TEST(CountMetricE2eTest, TestSlicedState) {
+    // Initialize config.
     StatsdConfig config;
     config.add_allowed_log_source("AID_ROOT");  // LogEvent defaults to UID of root.
 
@@ -37,30 +50,22 @@
     auto state = CreateScreenState();
     *config.add_state() = state;
 
-    // Create count metric that slices by screen state
+    // Create count metric that slices by screen state.
     int64_t metricId = 123456;
     auto countMetric = config.add_count_metric();
     countMetric->set_id(metricId);
     countMetric->set_what(syncStartMatcher.id());
-    countMetric->set_bucket(TimeUnit::ONE_MINUTE);
+    countMetric->set_bucket(TimeUnit::FIVE_MINUTES);
     countMetric->add_slice_by_state(state.id());
 
-    // Initialize StatsLogProcessor
-    const int64_t baseTimeNs = 0;                                   // 0:00
-    const int64_t configAddedTimeNs = baseTimeNs + 1 * NS_PER_SEC;  // 0:01
-    const int64_t bucketSizeNs =
-            TimeUnitToBucketSizeInMillis(config.count_metric(0).bucket()) * 1000LL * 1000LL;
-
+    // Initialize StatsLogProcessor.
+    const uint64_t bucketStartTimeNs = 10000000000;  // 0:10
+    const uint64_t bucketSizeNs =
+            TimeUnitToBucketSizeInMillis(config.count_metric(0).bucket()) * 1000000LL;
     int uid = 12345;
     int64_t cfgId = 98765;
     ConfigKey cfgKey(uid, cfgId);
-
-    auto processor = CreateStatsLogProcessor(baseTimeNs, configAddedTimeNs, config, cfgKey);
-
-    // Check that StateTrackers were properly initialized.
-    EXPECT_EQ(1, StateManager::getInstance().getStateTrackersCount());
-    EXPECT_EQ(1,
-              StateManager::getInstance().getListenersCount(android::util::SCREEN_STATE_CHANGED));
+    auto processor = CreateStatsLogProcessor(bucketStartTimeNs, bucketStartTimeNs, config, cfgKey);
 
     // Check that CountMetricProducer was initialized correctly.
     EXPECT_EQ(processor->mMetricsManagers.size(), 1u);
@@ -69,12 +74,118 @@
     EXPECT_EQ(metricsManager->mAllMetricProducers.size(), 1);
     sp<MetricProducer> metricProducer = metricsManager->mAllMetricProducers[0];
     EXPECT_EQ(metricProducer->mSlicedStateAtoms.size(), 1);
-    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(0), android::util::SCREEN_STATE_CHANGED);
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(0), SCREEN_STATE_ATOM_ID);
     EXPECT_EQ(metricProducer->mStateGroupMap.size(), 0);
+
+    // Check that StateTrackers were initialized correctly.
+    EXPECT_EQ(1, StateManager::getInstance().getStateTrackersCount());
+    EXPECT_EQ(1, StateManager::getInstance().getListenersCount(SCREEN_STATE_ATOM_ID));
+
+    /*
+               bucket #1                      bucket #2
+    |     1     2     3     4     5     6     7     8     9     10 (minutes)
+    |-----------------------------|-----------------------------|--
+            x                x         x    x        x      x       (syncStartEvents)
+          |                                       |                 (ScreenIsOnEvent)
+                   |     |                                          (ScreenIsOffEvent)
+                                                        |           (ScreenUnknownEvent)
+    */
+    // Initialize log events - first bucket.
+    int appUid = 123;
+    std::vector<AttributionNodeInternal> attributions1 = {CreateAttribution(appUid, "App1")};
+    std::vector<std::unique_ptr<LogEvent>> events;
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON,
+                                          bucketStartTimeNs + 50 * NS_PER_SEC));  // 1:00
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 75 * NS_PER_SEC));  // 1:25
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_OFF,
+                                          bucketStartTimeNs + 150 * NS_PER_SEC));  // 2:40
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_OFF,
+                                          bucketStartTimeNs + 200 * NS_PER_SEC));  // 3:30
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 250 * NS_PER_SEC));  // 4:20
+
+    // Initialize log events - second bucket.
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 350 * NS_PER_SEC));  // 6:00
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 400 * NS_PER_SEC));  // 6:50
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON,
+                                          bucketStartTimeNs + 450 * NS_PER_SEC));  // 7:40
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 475 * NS_PER_SEC));  // 8:05
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_UNKNOWN,
+                                          bucketStartTimeNs + 500 * NS_PER_SEC));  // 8:30
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 520 * NS_PER_SEC));  // 8:50
+
+    // Send log events to StatsLogProcessor.
+    for (auto& event : events) {
+        processor->OnLogEvent(event.get());
+    }
+
+    // Check dump report.
+    vector<uint8_t> buffer;
+    ConfigMetricsReportList reports;
+    processor->onDumpReport(cfgKey, bucketStartTimeNs + bucketSizeNs * 2 + 1, false, true, ADB_DUMP,
+                            FAST, &buffer);
+    EXPECT_GT(buffer.size(), 0);
+    EXPECT_TRUE(reports.ParseFromArray(&buffer[0], buffer.size()));
+    backfillDimensionPath(&reports);
+    backfillStringInReport(&reports);
+    backfillStartEndTimestamp(&reports);
+
+    EXPECT_EQ(1, reports.reports_size());
+    EXPECT_EQ(1, reports.reports(0).metrics_size());
+    EXPECT_TRUE(reports.reports(0).metrics(0).has_count_metrics());
+    EXPECT_EQ(3, reports.reports(0).metrics(0).count_metrics().data_size());
+
+    // For each CountMetricData, check StateValue info is correct and buckets
+    // have correct counts.
+    auto data = reports.reports(0).metrics(0).count_metrics().data(0);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_value());
+    EXPECT_EQ(android::view::DisplayStateEnum::DISPLAY_STATE_ON, data.slice_by_state(0).value());
+    EXPECT_EQ(2, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+    EXPECT_EQ(1, data.bucket_info(1).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(1);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_value());
+    EXPECT_EQ(android::view::DisplayStateEnum::DISPLAY_STATE_UNKNOWN, data.slice_by_state(0).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(2);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_value());
+    EXPECT_EQ(android::view::DisplayStateEnum::DISPLAY_STATE_OFF, data.slice_by_state(0).value());
+    EXPECT_EQ(2, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+    EXPECT_EQ(2, data.bucket_info(1).count());
 }
 
-TEST(CountMetricE2eTest, TestWithMappedState) {
-    // Initialize config
+/**
+ * Test a count metric that has one slice_by_state with a mapping and no
+ * primary fields.
+ *
+ * Once the CountMetricProducer is initialized, it has one atom id in
+ * mSlicedStateAtoms and has one entry per state value in mStateGroupMap.
+ *
+ * One StateTracker tracks the state atom, and it has one listener which is the
+ * CountMetricProducer that was initialized.
+ */
+TEST(CountMetricE2eTest, TestSlicedStateWithMap) {
+    // Initialize config.
     StatsdConfig config;
     config.add_allowed_log_source("AID_ROOT");  // LogEvent defaults to UID of root.
 
@@ -84,30 +195,26 @@
     auto state = CreateScreenStateWithOnOffMap();
     *config.add_state() = state;
 
-    // Create count metric that slices by screen state with on/off map
+    // Create count metric that slices by screen state with on/off map.
     int64_t metricId = 123456;
     auto countMetric = config.add_count_metric();
     countMetric->set_id(metricId);
     countMetric->set_what(syncStartMatcher.id());
-    countMetric->set_bucket(TimeUnit::ONE_MINUTE);
+    countMetric->set_bucket(TimeUnit::FIVE_MINUTES);
     countMetric->add_slice_by_state(state.id());
 
-    // Initialize StatsLogProcessor
-    const int64_t baseTimeNs = 0;                                   // 0:00
-    const int64_t configAddedTimeNs = baseTimeNs + 1 * NS_PER_SEC;  // 0:01
-    const int64_t bucketSizeNs =
-            TimeUnitToBucketSizeInMillis(config.count_metric(0).bucket()) * 1000LL * 1000LL;
-
+    // Initialize StatsLogProcessor.
+    const uint64_t bucketStartTimeNs = 10000000000;  // 0:10
+    const uint64_t bucketSizeNs =
+            TimeUnitToBucketSizeInMillis(config.count_metric(0).bucket()) * 1000000LL;
     int uid = 12345;
     int64_t cfgId = 98765;
     ConfigKey cfgKey(uid, cfgId);
+    auto processor = CreateStatsLogProcessor(bucketStartTimeNs, bucketStartTimeNs, config, cfgKey);
 
-    auto processor = CreateStatsLogProcessor(baseTimeNs, configAddedTimeNs, config, cfgKey);
-
-    // Check that StateTrackers were properly initialized.
+    // Check that StateTrackers were initialized correctly.
     EXPECT_EQ(1, StateManager::getInstance().getStateTrackersCount());
-    EXPECT_EQ(1,
-              StateManager::getInstance().getListenersCount(android::util::SCREEN_STATE_CHANGED));
+    EXPECT_EQ(1, StateManager::getInstance().getListenersCount(SCREEN_STATE_ATOM_ID));
 
     // Check that CountMetricProducer was initialized correctly.
     EXPECT_EQ(processor->mMetricsManagers.size(), 1u);
@@ -116,58 +223,371 @@
     EXPECT_EQ(metricsManager->mAllMetricProducers.size(), 1);
     sp<MetricProducer> metricProducer = metricsManager->mAllMetricProducers[0];
     EXPECT_EQ(metricProducer->mSlicedStateAtoms.size(), 1);
-    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(0), android::util::SCREEN_STATE_CHANGED);
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(0), SCREEN_STATE_ATOM_ID);
     EXPECT_EQ(metricProducer->mStateGroupMap.size(), 1);
 
     StateMap map = state.map();
     for (auto group : map.group()) {
         for (auto value : group.value()) {
-            EXPECT_EQ(metricProducer->mStateGroupMap[android::util::SCREEN_STATE_CHANGED][value],
+            EXPECT_EQ(metricProducer->mStateGroupMap[SCREEN_STATE_ATOM_ID][value],
                       group.group_id());
         }
     }
+
+    /*
+               bucket #1                      bucket #2
+    |     1     2     3     4     5     6     7     8     9     10 (minutes)
+    |-----------------------------|-----------------------------|--
+      x   x     x       x    x   x      x         x         x       (syncStartEvents)
+     -----------------------------------------------------------SCREEN_OFF events
+       |                                                            (ScreenStateUnknownEvent = 0)
+             |                  |                                   (ScreenStateOffEvent = 1)
+                          |                                         (ScreenStateDozeEvent = 3)
+                                                |                   (ScreenStateDozeSuspendEvent = 4)
+     -----------------------------------------------------------SCREEN_ON events
+                   |                                       |        (ScreenStateOnEvent = 2)
+                      |                                             (ScreenStateVrEvent = 5)
+                                            |                       (ScreenStateOnSuspendEvent = 6)
+    */
+    // Initialize log events - first bucket.
+    int appUid = 123;
+    std::vector<AttributionNodeInternal> attributions1 = {CreateAttribution(appUid, "App1")};
+
+    std::vector<std::unique_ptr<LogEvent>> events;
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 20 * NS_PER_SEC));  // 0:30
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_UNKNOWN,
+                                          bucketStartTimeNs + 30 * NS_PER_SEC));  // 0:40
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 60 * NS_PER_SEC));  // 1:10
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_OFF,
+                                          bucketStartTimeNs + 90 * NS_PER_SEC));  // 1:40
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 120 * NS_PER_SEC));  // 2:10
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON,
+                                          bucketStartTimeNs + 150 * NS_PER_SEC));  // 2:40
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_VR,
+                                          bucketStartTimeNs + 180 * NS_PER_SEC));  // 3:10
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 200 * NS_PER_SEC));  // 3:30
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_DOZE,
+                                          bucketStartTimeNs + 210 * NS_PER_SEC));  // 3:40
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 250 * NS_PER_SEC));  // 4:20
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_OFF,
+                                          bucketStartTimeNs + 280 * NS_PER_SEC));  // 4:50
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 285 * NS_PER_SEC));  // 4:55
+
+    // Initialize log events - second bucket.
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 360 * NS_PER_SEC));  // 6:10
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON_SUSPEND,
+                                          bucketStartTimeNs + 390 * NS_PER_SEC));  // 6:40
+    events.push_back(CreateScreenStateChangedEvent(
+            android::view::DisplayStateEnum::DISPLAY_STATE_DOZE_SUSPEND,
+            bucketStartTimeNs + 430 * NS_PER_SEC));  // 7:20
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 440 * NS_PER_SEC));  // 7:30
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON,
+                                          bucketStartTimeNs + 540 * NS_PER_SEC));  // 9:10
+    events.push_back(CreateSyncStartEvent(attributions1, "sync_name",
+                                          bucketStartTimeNs + 570 * NS_PER_SEC));  // 9:40
+
+    // Send log events to StatsLogProcessor.
+    for (auto& event : events) {
+        processor->OnLogEvent(event.get());
+    }
+
+    // Check dump report.
+    vector<uint8_t> buffer;
+    ConfigMetricsReportList reports;
+    processor->onDumpReport(cfgKey, bucketStartTimeNs + bucketSizeNs * 2 + 1, false, true, ADB_DUMP,
+                            FAST, &buffer);
+    EXPECT_GT(buffer.size(), 0);
+    EXPECT_TRUE(reports.ParseFromArray(&buffer[0], buffer.size()));
+    backfillDimensionPath(&reports);
+    backfillStringInReport(&reports);
+    backfillStartEndTimestamp(&reports);
+
+    EXPECT_EQ(1, reports.reports_size());
+    EXPECT_EQ(1, reports.reports(0).metrics_size());
+    EXPECT_TRUE(reports.reports(0).metrics(0).has_count_metrics());
+    EXPECT_EQ(3, reports.reports(0).metrics(0).count_metrics().data_size());
+
+    // For each CountMetricData, check StateValue info is correct and buckets
+    // have correct counts.
+    auto data = reports.reports(0).metrics(0).count_metrics().data(0);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_value());
+    EXPECT_EQ(-1 /* StateTracker::kStateUnknown */, data.slice_by_state(0).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(1);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_group_id());
+    EXPECT_EQ(StringToId("SCREEN_OFF"), data.slice_by_state(0).group_id());
+    EXPECT_EQ(2, data.bucket_info_size());
+    EXPECT_EQ(4, data.bucket_info(0).count());
+    EXPECT_EQ(2, data.bucket_info(1).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(2);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_group_id());
+    EXPECT_EQ(StringToId("SCREEN_ON"), data.slice_by_state(0).group_id());
+    EXPECT_EQ(2, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+    EXPECT_EQ(1, data.bucket_info(1).count());
 }
 
-TEST(CountMetricE2eTest, TestWithMultipleStates) {
-    // Initialize config
+/**
+ * Test a count metric that has one slice_by_state with a primary field.
+
+ * Once the CountMetricProducer is initialized, it should have one
+ * MetricStateLink stored. State querying using a non-empty primary key
+ * should also work as intended.
+ */
+TEST(CountMetricE2eTest, TestSlicedStateWithPrimaryFields) {
+    // Initialize config.
     StatsdConfig config;
     config.add_allowed_log_source("AID_ROOT");  // LogEvent defaults to UID of root.
 
-    auto syncStartMatcher = CreateSyncStartAtomMatcher();
-    *config.add_atom_matcher() = syncStartMatcher;
+    auto appCrashMatcher =
+            CreateSimpleAtomMatcher("APP_CRASH_OCCURRED", android::util::APP_CRASH_OCCURRED);
+    *config.add_atom_matcher() = appCrashMatcher;
+
+    auto state = CreateUidProcessState();
+    *config.add_state() = state;
+
+    // Create count metric that slices by uid process state.
+    int64_t metricId = 123456;
+    auto countMetric = config.add_count_metric();
+    countMetric->set_id(metricId);
+    countMetric->set_what(appCrashMatcher.id());
+    countMetric->set_bucket(TimeUnit::FIVE_MINUTES);
+    countMetric->add_slice_by_state(state.id());
+    MetricStateLink* stateLink = countMetric->add_state_link();
+    stateLink->set_state_atom_id(UID_PROCESS_STATE_ATOM_ID);
+    auto fieldsInWhat = stateLink->mutable_fields_in_what();
+    *fieldsInWhat = CreateDimensions(android::util::APP_CRASH_OCCURRED, {1 /* uid */});
+    auto fieldsInState = stateLink->mutable_fields_in_state();
+    *fieldsInState = CreateDimensions(UID_PROCESS_STATE_ATOM_ID, {1 /* uid */});
+
+    // Initialize StatsLogProcessor.
+    const uint64_t bucketStartTimeNs = 10000000000;  // 0:10
+    const uint64_t bucketSizeNs =
+            TimeUnitToBucketSizeInMillis(config.count_metric(0).bucket()) * 1000000LL;
+    int uid = 12345;
+    int64_t cfgId = 98765;
+    ConfigKey cfgKey(uid, cfgId);
+    auto processor = CreateStatsLogProcessor(bucketStartTimeNs, bucketStartTimeNs, config, cfgKey);
+
+    // Check that StateTrackers were initialized correctly.
+    EXPECT_EQ(1, StateManager::getInstance().getStateTrackersCount());
+    EXPECT_EQ(1, StateManager::getInstance().getListenersCount(UID_PROCESS_STATE_ATOM_ID));
+
+    // Check that CountMetricProducer was initialized correctly.
+    EXPECT_EQ(processor->mMetricsManagers.size(), 1u);
+    sp<MetricsManager> metricsManager = processor->mMetricsManagers.begin()->second;
+    EXPECT_TRUE(metricsManager->isConfigValid());
+    EXPECT_EQ(metricsManager->mAllMetricProducers.size(), 1);
+    sp<MetricProducer> metricProducer = metricsManager->mAllMetricProducers[0];
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.size(), 1);
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(0), UID_PROCESS_STATE_ATOM_ID);
+    EXPECT_EQ(metricProducer->mStateGroupMap.size(), 0);
+    EXPECT_EQ(metricProducer->mMetric2StateLinks.size(), 1);
+
+    /*
+    NOTE: "1" or "2" represents the uid associated with the state/app crash event
+               bucket #1                      bucket #2
+    |     1     2     3     4     5     6     7     8     9     10
+    |-----------------------------|-----------------------------|--
+      1   1     1       1    1   2      1         1          2      (AppCrashEvents)
+     -----------------------------------------------------------PROCESS STATE events
+             1                  2                                   (ProcessStateTopEvent = 1002)
+                          1                 1                       (ProcessStateForegroundServiceEvent = 1003)
+                                                2                   (ProcessStateImportantBackgroundEvent = 1006)
+        1          1                                       1        (ProcessStateImportantForegroundEvent = 1005)
+
+    Based on the diagram above, an AppCrashEvent querying for process state value would return:
+    - StateTracker::kStateUnknown
+    - Important foreground
+    - Top
+    - Important foreground
+    - Foreground service
+    - Top (both the app crash and state still have matching uid = 2)
+
+    - Foreground service
+    - Foreground service
+    - Important background
+    */
+    // Initialize log events - first bucket.
+    std::vector<std::unique_ptr<LogEvent>> events;
+    events.push_back(
+            CreateAppCrashOccurredEvent(1 /* uid */, bucketStartTimeNs + 20 * NS_PER_SEC));  // 0:30
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_IMPORTANT_FOREGROUND,
+            bucketStartTimeNs + 30 * NS_PER_SEC));  // 0:40
+    events.push_back(
+            CreateAppCrashOccurredEvent(1 /* uid */, bucketStartTimeNs + 60 * NS_PER_SEC));  // 1:10
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_TOP,
+            bucketStartTimeNs + 90 * NS_PER_SEC));  // 1:40
+    events.push_back(CreateAppCrashOccurredEvent(1 /* uid */,
+                                                 bucketStartTimeNs + 120 * NS_PER_SEC));  // 2:10
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_IMPORTANT_FOREGROUND,
+            bucketStartTimeNs + 150 * NS_PER_SEC));  // 2:40
+    events.push_back(CreateAppCrashOccurredEvent(1 /* uid */,
+                                                 bucketStartTimeNs + 200 * NS_PER_SEC));  // 3:30
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_FOREGROUND_SERVICE,
+            bucketStartTimeNs + 210 * NS_PER_SEC));  // 3:40
+    events.push_back(CreateAppCrashOccurredEvent(1 /* uid */,
+                                                 bucketStartTimeNs + 250 * NS_PER_SEC));  // 4:20
+    events.push_back(CreateUidProcessStateChangedEvent(
+            2 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_TOP,
+            bucketStartTimeNs + 280 * NS_PER_SEC));  // 4:50
+    events.push_back(CreateAppCrashOccurredEvent(2 /* uid */,
+                                                 bucketStartTimeNs + 285 * NS_PER_SEC));  // 4:55
+
+    // Initialize log events - second bucket.
+    events.push_back(CreateAppCrashOccurredEvent(1 /* uid */,
+                                                 bucketStartTimeNs + 360 * NS_PER_SEC));  // 6:10
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_FOREGROUND_SERVICE,
+            bucketStartTimeNs + 390 * NS_PER_SEC));  // 6:40
+    events.push_back(CreateUidProcessStateChangedEvent(
+            2 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_IMPORTANT_BACKGROUND,
+            bucketStartTimeNs + 430 * NS_PER_SEC));  // 7:20
+    events.push_back(CreateAppCrashOccurredEvent(1 /* uid */,
+                                                 bucketStartTimeNs + 440 * NS_PER_SEC));  // 7:30
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_IMPORTANT_FOREGROUND,
+            bucketStartTimeNs + 540 * NS_PER_SEC));  // 9:10
+    events.push_back(CreateAppCrashOccurredEvent(2 /* uid */,
+                                                 bucketStartTimeNs + 570 * NS_PER_SEC));  // 9:40
+
+    // Send log events to StatsLogProcessor.
+    for (auto& event : events) {
+        processor->OnLogEvent(event.get());
+    }
+
+    // Check dump report.
+    vector<uint8_t> buffer;
+    ConfigMetricsReportList reports;
+    processor->onDumpReport(cfgKey, bucketStartTimeNs + bucketSizeNs * 2 + 1, false, true, ADB_DUMP,
+                            FAST, &buffer);
+    EXPECT_GT(buffer.size(), 0);
+    EXPECT_TRUE(reports.ParseFromArray(&buffer[0], buffer.size()));
+    backfillDimensionPath(&reports);
+    backfillStringInReport(&reports);
+    backfillStartEndTimestamp(&reports);
+
+    EXPECT_EQ(1, reports.reports_size());
+    EXPECT_EQ(1, reports.reports(0).metrics_size());
+    EXPECT_TRUE(reports.reports(0).metrics(0).has_count_metrics());
+    EXPECT_EQ(5, reports.reports(0).metrics(0).count_metrics().data_size());
+
+    // For each CountMetricData, check StateValue info is correct and buckets
+    // have correct counts.
+    auto data = reports.reports(0).metrics(0).count_metrics().data(0);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_value());
+    EXPECT_EQ(android::app::PROCESS_STATE_IMPORTANT_BACKGROUND, data.slice_by_state(0).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(1);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_value());
+    EXPECT_EQ(-1 /* StateTracker::kStateUnknown */, data.slice_by_state(0).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(2);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_value());
+    EXPECT_EQ(android::app::PROCESS_STATE_IMPORTANT_FOREGROUND, data.slice_by_state(0).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(2, data.bucket_info(0).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(3);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_value());
+    EXPECT_EQ(android::app::PROCESS_STATE_TOP, data.slice_by_state(0).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(2, data.bucket_info(0).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(4);
+    EXPECT_EQ(1, data.slice_by_state_size());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_value());
+    EXPECT_EQ(android::app::PROCESS_STATE_FOREGROUND_SERVICE, data.slice_by_state(0).value());
+    EXPECT_EQ(2, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+    EXPECT_EQ(2, data.bucket_info(1).count());
+}
+
+TEST(CountMetricE2eTest, TestMultipleSlicedStates) {
+    // Initialize config.
+    StatsdConfig config;
+    config.add_allowed_log_source("AID_ROOT");  // LogEvent defaults to UID of root.
+
+    auto appCrashMatcher =
+            CreateSimpleAtomMatcher("APP_CRASH_OCCURRED", android::util::APP_CRASH_OCCURRED);
+    *config.add_atom_matcher() = appCrashMatcher;
 
     auto state1 = CreateScreenStateWithOnOffMap();
     *config.add_state() = state1;
     auto state2 = CreateUidProcessState();
     *config.add_state() = state2;
 
-    // Create count metric that slices by screen state with on/off map
+    // Create count metric that slices by screen state with on/off map and
+    // slices by uid process state.
     int64_t metricId = 123456;
     auto countMetric = config.add_count_metric();
     countMetric->set_id(metricId);
-    countMetric->set_what(syncStartMatcher.id());
-    countMetric->set_bucket(TimeUnit::ONE_MINUTE);
+    countMetric->set_what(appCrashMatcher.id());
+    countMetric->set_bucket(TimeUnit::FIVE_MINUTES);
     countMetric->add_slice_by_state(state1.id());
     countMetric->add_slice_by_state(state2.id());
+    MetricStateLink* stateLink = countMetric->add_state_link();
+    stateLink->set_state_atom_id(UID_PROCESS_STATE_ATOM_ID);
+    auto fieldsInWhat = stateLink->mutable_fields_in_what();
+    *fieldsInWhat = CreateDimensions(android::util::APP_CRASH_OCCURRED, {1 /* uid */});
+    auto fieldsInState = stateLink->mutable_fields_in_state();
+    *fieldsInState = CreateDimensions(UID_PROCESS_STATE_ATOM_ID, {1 /* uid */});
 
-    // Initialize StatsLogProcessor
-    const int64_t baseTimeNs = 0;                                   // 0:00
-    const int64_t configAddedTimeNs = baseTimeNs + 1 * NS_PER_SEC;  // 0:01
-    const int64_t bucketSizeNs =
-            TimeUnitToBucketSizeInMillis(config.count_metric(0).bucket()) * 1000LL * 1000LL;
-
+    // Initialize StatsLogProcessor.
+    const uint64_t bucketStartTimeNs = 10000000000;  // 0:10
+    const uint64_t bucketSizeNs =
+            TimeUnitToBucketSizeInMillis(config.count_metric(0).bucket()) * 1000000LL;
     int uid = 12345;
     int64_t cfgId = 98765;
     ConfigKey cfgKey(uid, cfgId);
-
-    auto processor = CreateStatsLogProcessor(baseTimeNs, configAddedTimeNs, config, cfgKey);
+    auto processor = CreateStatsLogProcessor(bucketStartTimeNs, bucketStartTimeNs, config, cfgKey);
 
     // Check that StateTrackers were properly initialized.
     EXPECT_EQ(2, StateManager::getInstance().getStateTrackersCount());
-    EXPECT_EQ(1,
-              StateManager::getInstance().getListenersCount(android::util::SCREEN_STATE_CHANGED));
-    EXPECT_EQ(1, StateManager::getInstance().getListenersCount(
-                         android::util::UID_PROCESS_STATE_CHANGED));
+    EXPECT_EQ(1, StateManager::getInstance().getListenersCount(SCREEN_STATE_ATOM_ID));
+    EXPECT_EQ(1, StateManager::getInstance().getListenersCount(UID_PROCESS_STATE_ATOM_ID));
 
     // Check that CountMetricProducer was initialized correctly.
     EXPECT_EQ(processor->mMetricsManagers.size(), 1u);
@@ -176,17 +596,205 @@
     EXPECT_EQ(metricsManager->mAllMetricProducers.size(), 1);
     sp<MetricProducer> metricProducer = metricsManager->mAllMetricProducers[0];
     EXPECT_EQ(metricProducer->mSlicedStateAtoms.size(), 2);
-    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(0), android::util::SCREEN_STATE_CHANGED);
-    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(1), android::util::UID_PROCESS_STATE_CHANGED);
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(0), SCREEN_STATE_ATOM_ID);
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(1), UID_PROCESS_STATE_ATOM_ID);
     EXPECT_EQ(metricProducer->mStateGroupMap.size(), 1);
+    EXPECT_EQ(metricProducer->mMetric2StateLinks.size(), 1);
 
     StateMap map = state1.map();
     for (auto group : map.group()) {
         for (auto value : group.value()) {
-            EXPECT_EQ(metricProducer->mStateGroupMap[android::util::SCREEN_STATE_CHANGED][value],
+            EXPECT_EQ(metricProducer->mStateGroupMap[SCREEN_STATE_ATOM_ID][value],
                       group.group_id());
         }
     }
+
+    /*
+               bucket #1                      bucket #2
+    |     1     2     3     4     5     6     7     8     9     10 (minutes)
+    |-----------------------------|-----------------------------|--
+      1   1     1       1    1   2      1         1           2     (AppCrashEvents)
+     -----------------------------------------------------------SCREEN_OFF events
+       |                                                            (ScreenStateUnknownEvent = 0)
+             |                                  |                   (ScreenStateOffEvent = 1)
+                          |                                         (ScreenStateDozeEvent = 3)
+     -----------------------------------------------------------SCREEN_ON events
+                    |                                    |          (ScreenStateOnEvent = 2)
+                                            |                       (ScreenStateOnSuspendEvent = 6)
+     -----------------------------------------------------------PROCESS STATE events
+             1                  2                                   (ProcessStateTopEvent = 1002)
+                                          1                         (ProcessStateForegroundServiceEvent = 1003)
+                                               2                    (ProcessStateImportantBackgroundEvent = 1006)
+     1             1                                       1        (ProcessStateImportantForegroundEvent = 1005)
+
+     Based on the diagram above, Screen State / Process State pairs for each
+     AppCrashEvent are:
+     - StateTracker::kStateUnknown / important foreground
+     - off / important foreground
+     - off / Top
+     - on / important foreground
+     - off / important foreground
+     - off / top
+
+     - off / important foreground
+     - off / foreground service
+     - on / important background
+
+    */
+    // Initialize log events - first bucket.
+    std::vector<std::unique_ptr<LogEvent>> events;
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_IMPORTANT_FOREGROUND,
+            bucketStartTimeNs + 5 * NS_PER_SEC));  // 0:15
+    events.push_back(
+            CreateAppCrashOccurredEvent(1 /* uid */, bucketStartTimeNs + 20 * NS_PER_SEC));  // 0:30
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_UNKNOWN,
+                                          bucketStartTimeNs + 30 * NS_PER_SEC));  // 0:40
+    events.push_back(
+            CreateAppCrashOccurredEvent(1 /* uid */, bucketStartTimeNs + 60 * NS_PER_SEC));  // 1:10
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_TOP,
+            bucketStartTimeNs + 90 * NS_PER_SEC));  // 1:40
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_OFF,
+                                          bucketStartTimeNs + 90 * NS_PER_SEC));  // 1:40
+    events.push_back(CreateAppCrashOccurredEvent(1 /* uid */,
+                                                 bucketStartTimeNs + 120 * NS_PER_SEC));  // 2:10
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_IMPORTANT_FOREGROUND,
+            bucketStartTimeNs + 150 * NS_PER_SEC));  // 2:40
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON,
+                                          bucketStartTimeNs + 160 * NS_PER_SEC));  // 2:50
+    events.push_back(CreateAppCrashOccurredEvent(1 /* uid */,
+                                                 bucketStartTimeNs + 200 * NS_PER_SEC));  // 3:30
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_DOZE,
+                                          bucketStartTimeNs + 210 * NS_PER_SEC));  // 3:40
+    events.push_back(CreateAppCrashOccurredEvent(1 /* uid */,
+                                                 bucketStartTimeNs + 250 * NS_PER_SEC));  // 4:20
+    events.push_back(CreateUidProcessStateChangedEvent(
+            2 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_TOP,
+            bucketStartTimeNs + 280 * NS_PER_SEC));  // 4:50
+    events.push_back(CreateAppCrashOccurredEvent(2 /* uid */,
+                                                 bucketStartTimeNs + 285 * NS_PER_SEC));  // 4:55
+
+    // Initialize log events - second bucket.
+    events.push_back(CreateAppCrashOccurredEvent(1 /* uid */,
+                                                 bucketStartTimeNs + 360 * NS_PER_SEC));  // 6:10
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_FOREGROUND_SERVICE,
+            bucketStartTimeNs + 380 * NS_PER_SEC));  // 6:30
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON_SUSPEND,
+                                          bucketStartTimeNs + 390 * NS_PER_SEC));  // 6:40
+    events.push_back(CreateUidProcessStateChangedEvent(
+            2 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_IMPORTANT_BACKGROUND,
+            bucketStartTimeNs + 420 * NS_PER_SEC));  // 7:10
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_OFF,
+                                          bucketStartTimeNs + 440 * NS_PER_SEC));  // 7:30
+    events.push_back(CreateAppCrashOccurredEvent(1 /* uid */,
+                                                 bucketStartTimeNs + 450 * NS_PER_SEC));  // 7:40
+    events.push_back(
+            CreateScreenStateChangedEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON,
+                                          bucketStartTimeNs + 520 * NS_PER_SEC));  // 8:50
+    events.push_back(CreateUidProcessStateChangedEvent(
+            1 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_IMPORTANT_FOREGROUND,
+            bucketStartTimeNs + 540 * NS_PER_SEC));  // 9:10
+    events.push_back(CreateAppCrashOccurredEvent(2 /* uid */,
+                                                 bucketStartTimeNs + 570 * NS_PER_SEC));  // 9:40
+
+    // Send log events to StatsLogProcessor.
+    for (auto& event : events) {
+        processor->OnLogEvent(event.get());
+    }
+
+    // Check dump report.
+    vector<uint8_t> buffer;
+    ConfigMetricsReportList reports;
+    processor->onDumpReport(cfgKey, bucketStartTimeNs + bucketSizeNs * 2 + 1, false, true, ADB_DUMP,
+                            FAST, &buffer);
+    EXPECT_GT(buffer.size(), 0);
+    EXPECT_TRUE(reports.ParseFromArray(&buffer[0], buffer.size()));
+    backfillDimensionPath(&reports);
+    backfillStringInReport(&reports);
+    backfillStartEndTimestamp(&reports);
+
+    EXPECT_EQ(1, reports.reports_size());
+    EXPECT_EQ(1, reports.reports(0).metrics_size());
+    EXPECT_TRUE(reports.reports(0).metrics(0).has_count_metrics());
+    EXPECT_EQ(6, reports.reports(0).metrics(0).count_metrics().data_size());
+
+    // For each CountMetricData, check StateValue info is correct and buckets
+    // have correct counts.
+    auto data = reports.reports(0).metrics(0).count_metrics().data(0);
+    EXPECT_EQ(2, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_group_id());
+    EXPECT_EQ(StringToId("SCREEN_OFF"), data.slice_by_state(0).group_id());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(1).atom_id());
+    EXPECT_TRUE(data.slice_by_state(1).has_value());
+    EXPECT_EQ(android::app::PROCESS_STATE_FOREGROUND_SERVICE, data.slice_by_state(1).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(1);
+    EXPECT_EQ(2, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_value());
+    EXPECT_EQ(-1, data.slice_by_state(0).value());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(1).atom_id());
+    EXPECT_TRUE(data.slice_by_state(1).has_value());
+    EXPECT_EQ(android::app::PROCESS_STATE_IMPORTANT_FOREGROUND, data.slice_by_state(1).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(2);
+    EXPECT_EQ(2, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_group_id());
+    EXPECT_EQ(StringToId("SCREEN_OFF"), data.slice_by_state(0).group_id());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(1).atom_id());
+    EXPECT_TRUE(data.slice_by_state(1).has_value());
+    EXPECT_EQ(android::app::PROCESS_STATE_IMPORTANT_FOREGROUND, data.slice_by_state(1).value());
+    EXPECT_EQ(2, data.bucket_info_size());
+    EXPECT_EQ(2, data.bucket_info(0).count());
+    EXPECT_EQ(1, data.bucket_info(1).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(3);
+    EXPECT_EQ(2, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_group_id());
+    EXPECT_EQ(StringToId("SCREEN_ON"), data.slice_by_state(0).group_id());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(1).atom_id());
+    EXPECT_TRUE(data.slice_by_state(1).has_value());
+    EXPECT_EQ(android::app::PROCESS_STATE_IMPORTANT_FOREGROUND, data.slice_by_state(1).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(4);
+    EXPECT_EQ(2, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_group_id());
+    EXPECT_EQ(StringToId("SCREEN_ON"), data.slice_by_state(0).group_id());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(1).atom_id());
+    EXPECT_TRUE(data.slice_by_state(1).has_value());
+    EXPECT_EQ(android::app::PROCESS_STATE_IMPORTANT_BACKGROUND, data.slice_by_state(1).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(1, data.bucket_info(0).count());
+
+    data = reports.reports(0).metrics(0).count_metrics().data(5);
+    EXPECT_EQ(2, data.slice_by_state_size());
+    EXPECT_EQ(SCREEN_STATE_ATOM_ID, data.slice_by_state(0).atom_id());
+    EXPECT_TRUE(data.slice_by_state(0).has_group_id());
+    EXPECT_EQ(StringToId("SCREEN_OFF"), data.slice_by_state(0).group_id());
+    EXPECT_EQ(UID_PROCESS_STATE_ATOM_ID, data.slice_by_state(1).atom_id());
+    EXPECT_TRUE(data.slice_by_state(1).has_value());
+    EXPECT_EQ(android::app::PROCESS_STATE_TOP, data.slice_by_state(1).value());
+    EXPECT_EQ(1, data.bucket_info_size());
+    EXPECT_EQ(2, data.bucket_info(0).count());
 }
 
 }  // namespace statsd
diff --git a/cmds/statsd/tests/metrics/metrics_test_helper.cpp b/cmds/statsd/tests/metrics/metrics_test_helper.cpp
index 7b9c0d6..108df04b 100644
--- a/cmds/statsd/tests/metrics/metrics_test_helper.cpp
+++ b/cmds/statsd/tests/metrics/metrics_test_helper.cpp
@@ -26,10 +26,23 @@
     return dimension;
 }
 
+HashableDimensionKey getMockedDimensionKeyLongValue(int tagId, int key, int64_t value) {
+    HashableDimensionKey dimension;
+    int pos[] = {key, 0, 0};
+    dimension.addValue(FieldValue(Field(tagId, pos, 0), Value(value)));
+
+    return dimension;
+}
+
 MetricDimensionKey getMockedMetricDimensionKey(int tagId, int key, string value) {
     return MetricDimensionKey(getMockedDimensionKey(tagId, key, value), DEFAULT_DIMENSION_KEY);
 }
 
+MetricDimensionKey getMockedStateDimensionKey(int tagId, int key, int64_t value) {
+    return MetricDimensionKey(DEFAULT_DIMENSION_KEY,
+                              getMockedDimensionKeyLongValue(tagId, key, value));
+}
+
 void buildSimpleAtomFieldMatcher(const int tagId, FieldMatcher* matcher) {
     matcher->set_field(tagId);
 }
@@ -41,4 +54,4 @@
 
 }  // namespace statsd
 }  // namespace os
-}  // namespace android
\ No newline at end of file
+}  // namespace android
diff --git a/cmds/statsd/tests/metrics/metrics_test_helper.h b/cmds/statsd/tests/metrics/metrics_test_helper.h
index 329e39f..09c4d9e 100644
--- a/cmds/statsd/tests/metrics/metrics_test_helper.h
+++ b/cmds/statsd/tests/metrics/metrics_test_helper.h
@@ -47,6 +47,9 @@
 HashableDimensionKey getMockedDimensionKey(int tagId, int key, std::string value);
 MetricDimensionKey getMockedMetricDimensionKey(int tagId, int key, std::string value);
 
+HashableDimensionKey getMockedDimensionKeyLongValue(int tagId, int key, int64_t value);
+MetricDimensionKey getMockedStateDimensionKey(int tagId, int key, int64_t value);
+
 // Utils to build FieldMatcher proto for simple one-depth atoms.
 void buildSimpleAtomFieldMatcher(const int tagId, const int atomFieldNum, FieldMatcher* matcher);
 void buildSimpleAtomFieldMatcher(const int tagId, FieldMatcher* matcher);
diff --git a/cmds/statsd/tests/state/StateTracker_test.cpp b/cmds/statsd/tests/state/StateTracker_test.cpp
index 8d38000..4208fef 100644
--- a/cmds/statsd/tests/state/StateTracker_test.cpp
+++ b/cmds/statsd/tests/state/StateTracker_test.cpp
@@ -50,6 +50,12 @@
     }
 };
 
+int getStateInt(StateManager& mgr, int atomId, const HashableDimensionKey& queryKey) {
+    FieldValue output;
+    mgr.getStateValue(atomId, queryKey, &output);
+    return output.mValue.int_value;
+}
+
 // START: build event functions.
 // State with no primary fields - ScreenStateChanged
 std::shared_ptr<LogEvent> buildScreenEvent(int state) {
@@ -240,7 +246,7 @@
 
     // check StateTracker was updated by querying for state
     HashableDimensionKey queryKey = DEFAULT_DIMENSION_KEY;
-    EXPECT_EQ(2, mgr.getStateValue(android::util::SCREEN_STATE_CHANGED, queryKey));
+    EXPECT_EQ(2, getStateInt(mgr, android::util::SCREEN_STATE_CHANGED, queryKey));
 }
 
 /**
@@ -265,7 +271,7 @@
     // check StateTracker was updated by querying for state
     HashableDimensionKey queryKey;
     getUidProcessKey(1000 /* uid */, &queryKey);
-    EXPECT_EQ(1002, mgr.getStateValue(android::util::UID_PROCESS_STATE_CHANGED, queryKey));
+    EXPECT_EQ(1002, getStateInt(mgr, android::util::UID_PROCESS_STATE_CHANGED, queryKey));
 }
 
 /**
@@ -290,7 +296,7 @@
     // check StateTracker was updated by querying for state
     HashableDimensionKey queryKey;
     getOverlayKey(1000 /* uid */, "package1", &queryKey);
-    EXPECT_EQ(1, mgr.getStateValue(android::util::OVERLAY_STATE_CHANGED, queryKey));
+    EXPECT_EQ(1, getStateInt(mgr, android::util::OVERLAY_STATE_CHANGED, queryKey));
 }
 
 /**
@@ -353,25 +359,25 @@
     // Query for UidProcessState of uid 1001
     HashableDimensionKey queryKey1;
     getUidProcessKey(1001, &queryKey1);
-    EXPECT_EQ(1003, mgr.getStateValue(android::util::UID_PROCESS_STATE_CHANGED, queryKey1));
+    EXPECT_EQ(1003, getStateInt(mgr, android::util::UID_PROCESS_STATE_CHANGED, queryKey1));
 
     // Query for UidProcessState of uid 1004 - not in state map
     HashableDimensionKey queryKey2;
     getUidProcessKey(1004, &queryKey2);
-    EXPECT_EQ(-1, mgr.getStateValue(android::util::UID_PROCESS_STATE_CHANGED,
-                                    queryKey2));  // default state
+    EXPECT_EQ(-1, getStateInt(mgr, android::util::UID_PROCESS_STATE_CHANGED,
+                              queryKey2));  // default state
 
     // Query for UidProcessState of uid 1001 - after change in state
     mgr.onLogEvent(*event4);
-    EXPECT_EQ(1002, mgr.getStateValue(android::util::UID_PROCESS_STATE_CHANGED, queryKey1));
+    EXPECT_EQ(1002, getStateInt(mgr, android::util::UID_PROCESS_STATE_CHANGED, queryKey1));
 
     // Query for ScreenState
-    EXPECT_EQ(2, mgr.getStateValue(android::util::SCREEN_STATE_CHANGED, DEFAULT_DIMENSION_KEY));
+    EXPECT_EQ(2, getStateInt(mgr, android::util::SCREEN_STATE_CHANGED, DEFAULT_DIMENSION_KEY));
 
     // Query for OverlayState of uid 1000, package name "package2"
     HashableDimensionKey queryKey3;
     getOverlayKey(1000, "package2", &queryKey3);
-    EXPECT_EQ(2, mgr.getStateValue(android::util::OVERLAY_STATE_CHANGED, queryKey3));
+    EXPECT_EQ(2, getStateInt(mgr, android::util::OVERLAY_STATE_CHANGED, queryKey3));
 }
 
 }  // namespace statsd
diff --git a/cmds/statsd/tests/statsd_test_util.cpp b/cmds/statsd/tests/statsd_test_util.cpp
index 38c22ab..d154b1b 100644
--- a/cmds/statsd/tests/statsd_test_util.cpp
+++ b/cmds/statsd/tests/statsd_test_util.cpp
@@ -254,28 +254,28 @@
 State CreateScreenState() {
     State state;
     state.set_id(StringToId("ScreenState"));
-    state.set_atom_id(29);
+    state.set_atom_id(android::util::SCREEN_STATE_CHANGED);
     return state;
 }
 
 State CreateUidProcessState() {
     State state;
     state.set_id(StringToId("UidProcessState"));
-    state.set_atom_id(27);
+    state.set_atom_id(android::util::UID_PROCESS_STATE_CHANGED);
     return state;
 }
 
 State CreateOverlayState() {
     State state;
     state.set_id(StringToId("OverlayState"));
-    state.set_atom_id(59);
+    state.set_atom_id(android::util::OVERLAY_STATE_CHANGED);
     return state;
 }
 
 State CreateScreenStateWithOnOffMap() {
     State state;
     state.set_id(StringToId("ScreenStateOnOff"));
-    state.set_atom_id(29);
+    state.set_atom_id(android::util::SCREEN_STATE_CHANGED);
 
     auto map = CreateScreenStateOnOffMap();
     *state.mutable_map() = map;
@@ -286,7 +286,7 @@
 State CreateScreenStateWithInDozeMap() {
     State state;
     state.set_id(StringToId("ScreenStateInDoze"));
-    state.set_atom_id(29);
+    state.set_atom_id(android::util::SCREEN_STATE_CHANGED);
 
     auto map = CreateScreenStateInDozeMap();
     *state.mutable_map() = map;
@@ -533,6 +533,15 @@
         uid, ProcessLifeCycleStateChanged::CRASHED, timestampNs);
 }
 
+std::unique_ptr<LogEvent> CreateAppCrashOccurredEvent(const int uid, uint64_t timestampNs) {
+    auto event = std::make_unique<LogEvent>(android::util::APP_CRASH_OCCURRED, timestampNs);
+    event->write(uid);
+    event->write("eventType");
+    event->write("processName");
+    event->init();
+    return event;
+}
+
 std::unique_ptr<LogEvent> CreateIsolatedUidChangedEvent(
     int isolatedUid, int hostUid, bool is_create, uint64_t timestampNs) {
     auto logEvent = std::make_unique<LogEvent>(
@@ -544,6 +553,15 @@
     return logEvent;
 }
 
+std::unique_ptr<LogEvent> CreateUidProcessStateChangedEvent(
+        int uid, const android::app::ProcessStateEnum state, uint64_t timestampNs) {
+    auto event = std::make_unique<LogEvent>(android::util::UID_PROCESS_STATE_CHANGED, timestampNs);
+    event->write(uid);
+    event->write(state);
+    event->init();
+    return event;
+}
+
 sp<StatsLogProcessor> CreateStatsLogProcessor(const int64_t timeBaseNs, const int64_t currentTimeNs,
                                               const StatsdConfig& config, const ConfigKey& key) {
     sp<UidMap> uidMap = new UidMap();
diff --git a/cmds/statsd/tests/statsd_test_util.h b/cmds/statsd/tests/statsd_test_util.h
index c026105..e1e134b 100644
--- a/cmds/statsd/tests/statsd_test_util.h
+++ b/cmds/statsd/tests/statsd_test_util.h
@@ -192,6 +192,9 @@
 std::unique_ptr<LogEvent> CreateAppCrashEvent(
     const int uid, uint64_t timestampNs);
 
+// Create log event for an app crash.
+std::unique_ptr<LogEvent> CreateAppCrashOccurredEvent(const int uid, uint64_t timestampNs);
+
 // Create log event for acquiring wakelock.
 std::unique_ptr<LogEvent> CreateAcquireWakelockEvent(
         const std::vector<AttributionNodeInternal>& attributions, const string& wakelockName,
@@ -206,6 +209,10 @@
 std::unique_ptr<LogEvent> CreateIsolatedUidChangedEvent(
     int isolatedUid, int hostUid, bool is_create, uint64_t timestampNs);
 
+// Create log event for uid process state change.
+std::unique_ptr<LogEvent> CreateUidProcessStateChangedEvent(
+        int uid, const android::app::ProcessStateEnum state, uint64_t timestampNs);
+
 // Helper function to create an AttributionNodeInternal proto.
 AttributionNodeInternal CreateAttribution(const int& uid, const string& tag);