diff --git a/cmds/statsd/src/metrics/parsing_utils/config_update_utils.cpp b/cmds/statsd/src/metrics/parsing_utils/config_update_utils.cpp
index 0699ec01d66e6..39789cd86bb17 100644
--- a/cmds/statsd/src/metrics/parsing_utils/config_update_utils.cpp
+++ b/cmds/statsd/src/metrics/parsing_utils/config_update_utils.cpp
@@ -396,6 +396,23 @@ bool updateConditions(const ConfigKey& key, const StatsdConfig& config,
     return true;
 }
 
+bool updateStates(const StatsdConfig& config, const map<int64_t, uint64_t>& oldStateProtoHashes,
+                  unordered_map<int64_t, int>& stateAtomIdMap,
+                  unordered_map<int64_t, unordered_map<int, int64_t>>& allStateGroupMaps,
+                  map<int64_t, uint64_t>& newStateProtoHashes, set<int64_t>& replacedStates) {
+    // Share with metrics_manager_util.
+    if (!initStates(config, stateAtomIdMap, allStateGroupMaps, newStateProtoHashes)) {
+        return false;
+    }
+
+    for (const auto& [stateId, stateHash] : oldStateProtoHashes) {
+        const auto& it = newStateProtoHashes.find(stateId);
+        if (it != newStateProtoHashes.end() && it->second != stateHash) {
+            replacedStates.insert(stateId);
+        }
+    }
+    return true;
+}
 // Returns true if any matchers in the metric activation were replaced.
 bool metricActivationDepsChange(const StatsdConfig& config,
                                 const unordered_map<int64_t, int>& metricToActivationMap,
@@ -1042,6 +1059,7 @@ bool updateStatsdConfig(const ConfigKey& key, const StatsdConfig& config, const
                         set<int64_t>& noReportMetricIds) {
     set<int64_t> replacedMatchers;
     set<int64_t> replacedConditions;
+    set<int64_t> replacedStates;
     set<int64_t> replacedMetrics;
     vector<ConditionState> conditionCache;
     unordered_map<int64_t, int> stateAtomIdMap;
@@ -1053,7 +1071,6 @@ bool updateStatsdConfig(const ConfigKey& key, const StatsdConfig& config, const
         ALOGE("updateAtomMatchingTrackers failed");
         return false;
     }
-    VLOG("updateAtomMatchingTrackers succeeded");
 
     if (!updateConditions(key, config, newAtomMatchingTrackerMap, replacedMatchers,
                           oldConditionTrackerMap, oldConditionTrackers, newConditionTrackerMap,
@@ -1062,21 +1079,12 @@ bool updateStatsdConfig(const ConfigKey& key, const StatsdConfig& config, const
         ALOGE("updateConditions failed");
         return false;
     }
-    VLOG("updateConditions succeeded");
 
-    // Share with metrics_manager_util,
-    if (!initStates(config, stateAtomIdMap, allStateGroupMaps, newStateProtoHashes)) {
-        ALOGE("initStates failed");
+    if (!updateStates(config, oldStateProtoHashes, stateAtomIdMap, allStateGroupMaps,
+                      newStateProtoHashes, replacedStates)) {
+        ALOGE("updateStates failed");
         return false;
     }
-
-    set<int64_t> replacedStates;
-    for (const auto& [stateId, stateHash] : oldStateProtoHashes) {
-        const auto& it = newStateProtoHashes.find(stateId);
-        if (it != newStateProtoHashes.end() && it->second != stateHash) {
-            replacedStates.insert(stateId);
-        }
-    }
     if (!updateMetrics(key, config, timeBaseNs, currentTimeNs, pullerManager,
                        oldAtomMatchingTrackerMap, newAtomMatchingTrackerMap, replacedMatchers,
                        newAtomMatchingTrackers, newConditionTrackerMap, replacedConditions,
diff --git a/cmds/statsd/src/metrics/parsing_utils/config_update_utils.h b/cmds/statsd/src/metrics/parsing_utils/config_update_utils.h
index 178a9d220b4dd..8e2be68996996 100644
--- a/cmds/statsd/src/metrics/parsing_utils/config_update_utils.h
+++ b/cmds/statsd/src/metrics/parsing_utils/config_update_utils.h
@@ -126,6 +126,13 @@ bool updateConditions(const ConfigKey& key, const StatsdConfig& config,
                       std::vector<ConditionState>& conditionCache,
                       std::set<int64_t>& replacedConditions);
 
+bool updateStates(const StatsdConfig& config,
+                  const std::map<int64_t, uint64_t>& oldStateProtoHashes,
+                  std::unordered_map<int64_t, int>& stateAtomIdMap,
+                  std::unordered_map<int64_t, std::unordered_map<int, int64_t>>& allStateGroupMaps,
+                  std::map<int64_t, uint64_t>& newStateProtoHashes,
+                  std::set<int64_t>& replacedStates);
+
 // Function to determine the update status (preserve/replace/new) of all metrics in the config.
 // [config]: the input StatsdConfig
 // [oldMetricProducerMap]: metric id to index mapping in the existing MetricsManager
diff --git a/cmds/statsd/tests/metrics/parsing_utils/config_update_utils_test.cpp b/cmds/statsd/tests/metrics/parsing_utils/config_update_utils_test.cpp
index 66bab4ea70daf..d78c14c6ed82e 100644
--- a/cmds/statsd/tests/metrics/parsing_utils/config_update_utils_test.cpp
+++ b/cmds/statsd/tests/metrics/parsing_utils/config_update_utils_test.cpp
@@ -997,6 +997,57 @@ TEST_F(ConfigUpdateTest, TestUpdateConditions) {
     EXPECT_THAT(combinationTracker1->mSlicedChildren, IsEmpty());
 }
 
+TEST_F(ConfigUpdateTest, TestUpdateStates) {
+    StatsdConfig config;
+    // Add states.
+    // Will be replaced because we add a state map.
+    State state1 = CreateScreenState();
+    int64_t state1Id = state1.id();
+    *config.add_state() = state1;
+
+    // Will be preserved.
+    State state2 = CreateUidProcessState();
+    int64_t state2Id = state2.id();
+    *config.add_state() = state2;
+
+    // Will be replaced since the atom changes from overlay to screen.
+    State state3 = CreateOverlayState();
+    int64_t state3Id = state3.id();
+    *config.add_state() = state3;
+
+    EXPECT_TRUE(initConfig(config));
+
+    // Change definitions of state1 and state3.
+    int64_t screenOnId = 0x4321, screenOffId = 0x1234;
+    *state1.mutable_map() = CreateScreenStateSimpleOnOffMap(screenOnId, screenOffId);
+    state3.set_atom_id(util::SCREEN_STATE_CHANGED);
+
+    StatsdConfig newConfig;
+    *newConfig.add_state() = state3;
+    *newConfig.add_state() = state1;
+    *newConfig.add_state() = state2;
+
+    unordered_map<int64_t, int> stateAtomIdMap;
+    unordered_map<int64_t, unordered_map<int, int64_t>> allStateGroupMaps;
+    map<int64_t, uint64_t> newStateProtoHashes;
+    set<int64_t> replacedStates;
+    EXPECT_TRUE(updateStates(newConfig, oldStateHashes, stateAtomIdMap, allStateGroupMaps,
+                             newStateProtoHashes, replacedStates));
+    EXPECT_THAT(replacedStates, ContainerEq(set({state1Id, state3Id})));
+
+    unordered_map<int64_t, int> expectedStateAtomIdMap = {
+            {state1Id, util::SCREEN_STATE_CHANGED},
+            {state2Id, util::UID_PROCESS_STATE_CHANGED},
+            {state3Id, util::SCREEN_STATE_CHANGED}};
+    EXPECT_THAT(stateAtomIdMap, ContainerEq(expectedStateAtomIdMap));
+
+    unordered_map<int64_t, unordered_map<int, int64_t>> expectedStateGroupMaps = {
+            {state1Id,
+             {{android::view::DisplayStateEnum::DISPLAY_STATE_OFF, screenOffId},
+              {android::view::DisplayStateEnum::DISPLAY_STATE_ON, screenOnId}}}};
+    EXPECT_THAT(allStateGroupMaps, ContainerEq(expectedStateGroupMaps));
+}
+
 TEST_F(ConfigUpdateTest, TestEventMetricPreserve) {
     StatsdConfig config;
     AtomMatcher startMatcher = CreateScreenTurnedOnAtomMatcher();