Merge "Add StateTracker classes"

cmds/statsd/Android.bp

@@ -99,6 +99,8 @@ cc_defaults {
         "src/shell/shell_config.proto",
         "src/shell/ShellSubscriber.cpp",
         "src/socket/StatsSocketListener.cpp",
+        "src/state/StateManager.cpp",
+        "src/state/StateTracker.cpp",
         "src/stats_log_util.cpp",
         "src/statscompanion_util.cpp",
         "src/statsd_config.proto",
@@ -253,6 +255,7 @@ cc_test {
         "tests/metrics/ValueMetricProducer_test.cpp",
         "tests/MetricsManager_test.cpp",
         "tests/shell/ShellSubscriber_test.cpp",
+        "tests/state/StateTracker_test.cpp",
         "tests/statsd_test_util.cpp",
         "tests/StatsLogProcessor_test.cpp",
         "tests/StatsService_test.cpp",

cmds/statsd/src/HashableDimensionKey.cpp

@@ -59,6 +59,16 @@ android::hash_t hashDimension(const HashableDimensionKey& value) {
     return JenkinsHashWhiten(hash);
 }
 
+bool filterValues(const Matcher& matcherField, const vector<FieldValue>& values, Value* output) {
+    for (const auto& value : values) {
+        if (value.mField.matches(matcherField)) {
+            (*output) = value.mValue;
+            return true;
+        }
+    }
+    return false;
+}
+
 bool filterValues(const vector<Matcher>& matcherFields, const vector<FieldValue>& values,
                   HashableDimensionKey* output) {
     size_t num_matches = 0;

cmds/statsd/src/HashableDimensionKey.h

@@ -119,6 +119,13 @@ class MetricDimensionKey {
 
 android::hash_t hashDimension(const HashableDimensionKey& key);
 
+/**
+ * Returns true if a FieldValue field matches the matcher field.
+ * The value of the FieldValue is output.
+ */
+bool filterValues(const Matcher& matcherField, const std::vector<FieldValue>& values,
+                  Value* output);
+
 /**
  * Creating HashableDimensionKeys from FieldValues using matcher.
  *
@@ -169,4 +176,4 @@ struct hash<MetricDimensionKey> {
         return android::JenkinsHashWhiten(hash);
     }
 };
-}  // namespace std
+}  // namespace std

cmds/statsd/src/atoms.proto

@@ -3069,9 +3069,9 @@ message PictureInPictureStateChanged {
  *     services/core/java/com/android/server/wm/Session.java
  */
 message OverlayStateChanged {
-    optional int32 uid = 1 [(is_uid) = true];
+    optional int32 uid = 1 [(state_field_option).option = PRIMARY, (is_uid) = true];
 
-    optional string package_name = 2;
+    optional string package_name = 2 [(state_field_option).option = PRIMARY];
 
     optional bool using_alert_window = 3;
 
@@ -3079,7 +3079,7 @@ message OverlayStateChanged {
         ENTERED = 1;
         EXITED = 2;
     }
-    optional State state = 4;
+    optional State state = 4 [(state_field_option).option = EXCLUSIVE];
 }
 
 /*

cmds/statsd/src/state/StateListener.h

@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2019, The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <utils/RefBase.h>
+
+#include "HashableDimensionKey.h"
+
+namespace android {
+namespace os {
+namespace statsd {
+
+class StateListener : public virtual RefBase {
+public:
+    StateListener(){};
+
+    virtual ~StateListener(){};
+
+    /**
+     * Interface for handling a state change.
+     *
+     * The old and new state values map to the original state values.
+     * StateTrackers only track the original state values and are unaware
+     * of higher-level state groups. MetricProducers hold information on
+     * state groups and are responsible for mapping original state values to
+     * the correct state group.
+     *
+     * [atomId]: The id of the state atom
+     * [primaryKey]: The primary field values of the state atom
+     * [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;
+};
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android

cmds/statsd/src/state/StateManager.cpp

@@ -0,0 +1,92 @@
+/*
+ * Copyright (C) 2019, The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#define DEBUG false  // STOPSHIP if true
+#include "Log.h"
+
+#include "StateManager.h"
+
+namespace android {
+namespace os {
+namespace statsd {
+
+StateManager& StateManager::getInstance() {
+    static StateManager sStateManager;
+    return sStateManager;
+}
+
+void StateManager::onLogEvent(const LogEvent& event) {
+    std::lock_guard<std::mutex> lock(mMutex);
+    if (mStateTrackers.find(event.GetTagId()) != mStateTrackers.end()) {
+        mStateTrackers[event.GetTagId()]->onLogEvent(event);
+    }
+}
+
+bool StateManager::registerListener(int stateAtomId, wp<StateListener> listener) {
+    std::lock_guard<std::mutex> lock(mMutex);
+
+    // Check if state tracker already exists
+    if (mStateTrackers.find(stateAtomId) == mStateTrackers.end()) {
+        // Create a new state tracker iff atom is a state atom
+        auto it = android::util::AtomsInfo::kStateAtomsFieldOptions.find(stateAtomId);
+        if (it != android::util::AtomsInfo::kStateAtomsFieldOptions.end()) {
+            mStateTrackers[stateAtomId] = new StateTracker(stateAtomId, it->second);
+        } else {
+            ALOGE("StateManager cannot register listener, Atom %d is not a state atom",
+                  stateAtomId);
+            return false;
+        }
+    }
+    mStateTrackers[stateAtomId]->registerListener(listener);
+    return true;
+}
+
+void StateManager::unregisterListener(int stateAtomId, wp<StateListener> listener) {
+    std::unique_lock<std::mutex> lock(mMutex);
+
+    // Hold the sp<> until the lock is released so that ~StateTracker() is
+    // not called while the lock is held.
+    sp<StateTracker> toRemove;
+
+    // Unregister listener from correct StateTracker
+    auto it = mStateTrackers.find(stateAtomId);
+    if (it != mStateTrackers.end()) {
+        it->second->unregisterListener(listener);
+
+        // Remove the StateTracker if it has no listeners
+        if (it->second->getListenersCount() == 0) {
+            toRemove = it->second;
+            mStateTrackers.erase(it);
+        }
+    } else {
+        ALOGE("StateManager cannot unregister listener, StateTracker for atom %d does not exist",
+              stateAtomId);
+    }
+    lock.unlock();
+}
+
+int StateManager::getState(int stateAtomId, const HashableDimensionKey& key) {
+    std::lock_guard<std::mutex> lock(mMutex);
+    if (mStateTrackers.find(stateAtomId) != mStateTrackers.end()) {
+        return mStateTrackers[stateAtomId]->getState(key);
+    }
+
+    return StateTracker::kStateUnknown;
+}
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android

cmds/statsd/src/state/StateManager.h

@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2019, The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+//#include <utils/Log.h>
+#include <utils/RefBase.h>
+#include "HashableDimensionKey.h"
+
+#include "state/StateListener.h"
+#include "state/StateTracker.h"
+
+namespace android {
+namespace os {
+namespace statsd {
+
+class StateManager : public virtual RefBase {
+public:
+    StateManager(){};
+
+    ~StateManager(){};
+
+    // Returns a pointer to the single, shared StateManager object.
+    static StateManager& getInstance();
+
+    // Notifies the correct StateTracker of an event.
+    void onLogEvent(const LogEvent& event);
+
+    // Returns true if stateAtomId is the id of a state atom and notifies the
+    // correct StateTracker to register the listener. If the correct
+    // StateTracker does not exist, a new StateTracker is created.
+    bool registerListener(int stateAtomId, 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 stateAtomId, wp<StateListener> listener);
+
+    // Queries the correct StateTracker for the state that is mapped to the given
+    // query key.
+    // If the StateTracker doesn't exist, returns StateTracker::kStateUnknown.
+    int getState(int stateAtomId, const HashableDimensionKey& queryKey);
+
+    inline int getStateTrackersCount() {
+        std::lock_guard<std::mutex> lock(mMutex);
+        return mStateTrackers.size();
+    }
+
+    inline int getListenersCount(int stateAtomId) {
+        std::lock_guard<std::mutex> lock(mMutex);
+        if (mStateTrackers.find(stateAtomId) != mStateTrackers.end()) {
+            return mStateTrackers[stateAtomId]->getListenersCount();
+        }
+        return -1;
+    }
+
+private:
+  mutable std::mutex mMutex;
+
+  // Maps state atom ids to StateTrackers
+  std::unordered_map<int, sp<StateTracker>> mStateTrackers;
+};
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android

cmds/statsd/src/state/StateTracker.cpp

@@ -0,0 +1,159 @@
+/*
+ * Copyright (C) 2019, The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#define DEBUG true  // STOPSHIP if true
+#include "Log.h"
+
+#include "stats_util.h"
+
+#include "StateTracker.h"
+
+namespace android {
+namespace os {
+namespace statsd {
+
+StateTracker::StateTracker(const int atomId,
+                           const util::StateAtomFieldOptions& stateAtomInfo)
+  : mAtomId(atomId),
+    mStateField(getSimpleMatcher(atomId, stateAtomInfo.exclusiveField)) {
+    // create matcher for each primary field
+    // TODO(tsaichristine): handle when primary field is first uid in chain
+    for (const auto& primary : stateAtomInfo.primaryFields) {
+        Matcher matcher = getSimpleMatcher(atomId, primary);
+        mPrimaryFields.push_back(matcher);
+    }
+
+    // TODO(tsaichristine): set default state, reset state, and nesting
+}
+
+void StateTracker::onLogEvent(const LogEvent& event) {
+    // parse event for primary field values i.e. primary key
+    HashableDimensionKey primaryKey;
+    if (mPrimaryFields.size() > 0) {
+        if (!filterValues(mPrimaryFields, event.getValues(), &primaryKey) ||
+            primaryKey.getValues().size() != mPrimaryFields.size()) {
+            ALOGE("StateTracker error extracting primary key from log event.");
+            handleReset();
+            return;
+        }
+    } else {
+        // atom has no primary fields
+        primaryKey = DEFAULT_DIMENSION_KEY;
+    }
+
+    // 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());
+        handlePartialReset(primaryKey);
+        return;
+    }
+    stateValue = state.int_value;
+
+    if (stateValue == mResetState) {
+        VLOG("StateTracker Reset state: %s", state.toString().c_str());
+        handleReset();
+    }
+
+    // track and update state
+    int32_t oldState = 0;
+    int32_t newState = 0;
+    updateState(primaryKey, stateValue, &oldState, &newState);
+
+    // notify all listeners if state has changed
+    if (oldState != newState) {
+        VLOG("StateTracker updated state");
+        for (auto listener : mListeners) {
+            auto sListener = listener.promote();  // safe access to wp<>
+            if (sListener != nullptr) {
+                sListener->onStateChanged(mAtomId, primaryKey, oldState, newState);
+            }
+        }
+    } else {
+        VLOG("StateTracker NO updated state");
+    }
+}
+
+void StateTracker::registerListener(wp<StateListener> listener) {
+    mListeners.insert(listener);
+}
+
+void StateTracker::unregisterListener(wp<StateListener> listener) {
+    mListeners.erase(listener);
+}
+
+int StateTracker::getState(const HashableDimensionKey& queryKey) const {
+    if (queryKey.getValues().size() == mPrimaryFields.size()) {
+        auto it = mStateMap.find(queryKey);
+        if (it != mStateMap.end()) {
+            return it->second.state;
+        }
+    } 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;
+}
+
+void StateTracker::handleReset() {
+    VLOG("StateTracker handle reset");
+    for (const auto pair : mStateMap) {
+        for (auto l : mListeners) {
+            auto sl = l.promote();
+            if (sl != nullptr) {
+                sl->onStateChanged(mAtomId, pair.first, pair.second.state, mDefaultState);
+            }
+        }
+    }
+    mStateMap.clear();
+}
+
+void StateTracker::handlePartialReset(const HashableDimensionKey& primaryKey) {
+    VLOG("StateTracker handle partial reset");
+    if (mStateMap.find(primaryKey) != mStateMap.end()) {
+        mStateMap.erase(primaryKey);
+    }
+}
+
+void StateTracker::updateState(const HashableDimensionKey& primaryKey, const int32_t eventState,
+                               int32_t* oldState, int32_t* newState) {
+    // get old state (either current state in map or default state)
+    auto it = mStateMap.find(primaryKey);
+    if (it != mStateMap.end()) {
+        *oldState = it->second.state;
+    } else {
+        *oldState = mDefaultState;
+    }
+
+    // update state map
+    if (eventState == mDefaultState) {
+        // remove (key, state) pair if state returns to default state
+        VLOG("\t StateTracker changed to default state")
+        mStateMap.erase(primaryKey);
+    } else {
+        mStateMap[primaryKey].state = eventState;
+        mStateMap[primaryKey].count = 1;
+    }
+    *newState = eventState;
+
+    // TODO(tsaichristine): support atoms with nested counting
+}
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android

cmds/statsd/src/state/StateTracker.h

@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2019, The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <statslog.h>
+#include <utils/RefBase.h>
+#include "HashableDimensionKey.h"
+#include "logd/LogEvent.h"
+
+#include "state/StateListener.h"
+
+#include <unordered_map>
+
+namespace android {
+namespace os {
+namespace statsd {
+
+class StateTracker : public virtual RefBase {
+public:
+    StateTracker(const int atomId, const util::StateAtomFieldOptions& stateAtomInfo);
+
+    virtual ~StateTracker(){};
+
+    // Updates state map and notifies all listeners if a state change occurs.
+    // Checks if a state change has occurred by getting the state value from
+    // the log event and comparing the old and new states.
+    void onLogEvent(const LogEvent& event);
+
+    // Adds new listeners to set of StateListeners. If a listener is already
+    // registered, it is ignored.
+    void registerListener(wp<StateListener> listener);
+
+    void unregisterListener(wp<StateListener> listener);
+
+    // Returns the state value mapped to the given query key.
+    // 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 getState(const HashableDimensionKey& queryKey) const;
+
+    inline int getListenersCount() const {
+        return mListeners.size();
+    }
+
+    const static int kStateUnknown = -1;
+
+private:
+    struct StateValueInfo {
+        int32_t state;  // state value
+        int count;      // nested count (only used for binary states)
+    };
+
+    const int32_t mAtomId;  // id of the state atom being tracked
+
+    Matcher mStateField;  // matches the atom's exclusive state field
+
+    std::vector<Matcher> mPrimaryFields;  // matches the atom's primary fields
+
+    int32_t mDefaultState = kStateUnknown;
+
+    int32_t mResetState;
+
+    // Maps primary key to state value info
+    std::unordered_map<HashableDimensionKey, StateValueInfo> mStateMap;
+
+    // Set of all StateListeners (objects listening for state changes)
+    std::set<wp<StateListener>> mListeners;
+
+    // Reset all state values in map to default state
+    void handleReset();
+
+    // Reset only the state value mapped to primary key to default state
+    void handlePartialReset(const HashableDimensionKey& primaryKey);
+
+    // Update the StateMap based on the received state value.
+    // Store the old and new states.
+    void updateState(const HashableDimensionKey& primaryKey, const int32_t eventState,
+                     int32_t* oldState, int32_t* newState);
+};
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android

cmds/statsd/tests/state/StateTracker_test.cpp

@@ -0,0 +1,356 @@
+/*
+ * Copyright (C) 2019, The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <gtest/gtest.h>
+#include "state/StateManager.h"
+#include "state/StateTracker.h"
+#include "state/StateListener.h"
+
+#include "tests/statsd_test_util.h"
+
+#ifdef __ANDROID__
+
+namespace android {
+namespace os {
+namespace statsd {
+
+/**
+ * Mock StateListener class for testing.
+ * Stores primary key and state pairs.
+ */
+class TestStateListener : public virtual StateListener {
+public:
+    TestStateListener(){};
+
+    virtual ~TestStateListener(){};
+
+    struct Update {
+        Update(const HashableDimensionKey& key, int state) : mKey(key), mState(state){};
+        HashableDimensionKey mKey;
+        int mState;
+    };
+
+    std::vector<Update> updates;
+
+    void onStateChanged(int stateAtomId, const HashableDimensionKey& primaryKey, int oldState,
+                        int newState) {
+        updates.emplace_back(primaryKey, newState);
+    }
+};
+
+// START: build event functions.
+// State with no primary fields - ScreenStateChanged
+std::shared_ptr<LogEvent> buildScreenEvent(int state) {
+    std::shared_ptr<LogEvent> event =
+            std::make_shared<LogEvent>(android::util::SCREEN_STATE_CHANGED, 1000 /*timestamp*/);
+    event->write((int32_t)state);
+    event->init();
+    return event;
+}
+
+// State with one primary field - UidProcessStateChanged
+std::shared_ptr<LogEvent> buildUidProcessEvent(int uid, int state) {
+    std::shared_ptr<LogEvent> event =
+            std::make_shared<LogEvent>(android::util::UID_PROCESS_STATE_CHANGED, 1000 /*timestamp*/);
+    event->write((int32_t)uid);
+    event->write((int32_t)state);
+    event->init();
+    return event;
+}
+
+// State with multiple primary fields - OverlayStateChanged
+std::shared_ptr<LogEvent> buildOverlayEvent(int uid, const std::string& packageName, int state) {
+    std::shared_ptr<LogEvent> event =
+            std::make_shared<LogEvent>(android::util::OVERLAY_STATE_CHANGED, 1000 /*timestamp*/);
+    event->write((int32_t)uid);
+    event->write(packageName);
+    event->write(true);  // using_alert_window
+    event->write((int32_t)state);
+    event->init();
+    return event;
+}
+
+std::shared_ptr<LogEvent> buildIncorrectOverlayEvent(int uid, const std::string& packageName, int state) {
+    std::shared_ptr<LogEvent> event =
+            std::make_shared<LogEvent>(android::util::OVERLAY_STATE_CHANGED, 1000 /*timestamp*/);
+    event->write((int32_t)uid);
+    event->write(packageName);
+    event->write((int32_t)state);
+    event->init();
+    return event;
+}
+// END: build event functions.
+
+// START: get primary key functions
+void getUidProcessKey(int uid, HashableDimensionKey* key) {
+    int pos1[] = {1, 0, 0};
+    Field field1(27 /* atom id */, pos1, 0 /* depth */);
+    Value value1((int32_t)uid);
+
+    key->addValue(FieldValue(field1, value1));
+}
+
+void getOverlayKey(int uid, string packageName, HashableDimensionKey* key) {
+    int pos1[] = {1, 0, 0};
+    int pos2[] = {2, 0, 0};
+
+    Field field1(59 /* atom id */, pos1, 0 /* depth */);
+    Field field2(59 /* atom id */, pos2, 0 /* depth */);
+
+    Value value1((int32_t)uid);
+    Value value2(packageName);
+
+    key->addValue(FieldValue(field1, value1));
+    key->addValue(FieldValue(field2, value2));
+}
+// END: get primary key functions
+
+TEST(StateListenerTest, TestStateListenerWeakPointer) {
+    sp<TestStateListener> listener = new TestStateListener();
+    wp<TestStateListener> wListener = listener;
+    listener = nullptr;  // let go of listener
+    EXPECT_TRUE(wListener.promote() == nullptr);
+}
+
+TEST(StateManagerTest, TestStateManagerGetInstance) {
+    sp<TestStateListener> listener1 = new TestStateListener();
+    StateManager& mgr = StateManager::getInstance();
+
+    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener1);
+    EXPECT_EQ(1, mgr.getStateTrackersCount());
+    EXPECT_EQ(1, StateManager::getInstance().getStateTrackersCount());
+}
+
+/**
+ * Test registering listeners to StateTrackers
+ *
+ * - StateManager will create a new StateTracker if it doesn't already exist
+ * and then register the listener to the StateTracker.
+ * - If a listener is already registered to a StateTracker, it is not added again.
+ * - StateTrackers are only created for atoms that are state atoms.
+ */
+TEST(StateTrackerTest, TestRegisterListener) {
+    sp<TestStateListener> listener1 = new TestStateListener();
+    sp<TestStateListener> listener2 = new TestStateListener();
+    StateManager mgr;
+
+    // Register listener to non-existing StateTracker
+    EXPECT_EQ(0, mgr.getStateTrackersCount());
+    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener1);
+    EXPECT_EQ(1, mgr.getStateTrackersCount());
+    EXPECT_EQ(1, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
+
+    // Register listener to existing StateTracker
+    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener2);
+    EXPECT_EQ(1, mgr.getStateTrackersCount());
+    EXPECT_EQ(2, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
+
+    // Register already registered listener to existing StateTracker
+    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener2);
+    EXPECT_EQ(1, mgr.getStateTrackersCount());
+    EXPECT_EQ(2, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
+
+    // Register listener to non-state atom
+    mgr.registerListener(android::util::BATTERY_LEVEL_CHANGED, listener2);
+    EXPECT_EQ(1, mgr.getStateTrackersCount());
+}
+
+/**
+ * Test unregistering listeners from StateTrackers
+ *
+ * - StateManager will unregister listeners from a StateTracker only if the
+ * StateTracker exists and the listener is registered to the StateTracker.
+ * - Once all listeners are removed from a StateTracker, the StateTracker
+ * is also removed.
+ */
+TEST(StateTrackerTest, TestUnregisterListener) {
+    sp<TestStateListener> listener1 = new TestStateListener();
+    sp<TestStateListener> listener2 = new TestStateListener();
+    StateManager mgr;
+
+    // Unregister listener from non-existing StateTracker
+    EXPECT_EQ(0, mgr.getStateTrackersCount());
+    mgr.unregisterListener(android::util::SCREEN_STATE_CHANGED, listener1);
+    EXPECT_EQ(0, mgr.getStateTrackersCount());
+    EXPECT_EQ(-1, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
+
+    // Unregister non-registered listener from existing StateTracker
+    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener1);
+    EXPECT_EQ(1, mgr.getStateTrackersCount());
+    EXPECT_EQ(1, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
+    mgr.unregisterListener(android::util::SCREEN_STATE_CHANGED, listener2);
+    EXPECT_EQ(1, mgr.getStateTrackersCount());
+    EXPECT_EQ(1, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
+
+    // Unregister second-to-last listener from StateTracker
+    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener2);
+    mgr.unregisterListener(android::util::SCREEN_STATE_CHANGED, listener1);
+    EXPECT_EQ(1, mgr.getStateTrackersCount());
+    EXPECT_EQ(1, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
+
+    // Unregister last listener from StateTracker
+    mgr.unregisterListener(android::util::SCREEN_STATE_CHANGED, listener2);
+    EXPECT_EQ(0, mgr.getStateTrackersCount());
+    EXPECT_EQ(-1, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
+}
+
+/**
+ * Test StateManager's onLogEvent and StateListener's onStateChanged correctly
+ * updates listener for states without primary keys.
+ */
+TEST(StateTrackerTest, TestStateChangeNoPrimaryFields) {
+    sp<TestStateListener> listener1 = new TestStateListener();
+    StateManager mgr;
+    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener1);
+
+    // log event
+    std::shared_ptr<LogEvent> event =
+            buildScreenEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON);
+    mgr.onLogEvent(*event);
+
+    // check listener was updated
+    EXPECT_EQ(1, listener1->updates.size());
+    EXPECT_EQ(DEFAULT_DIMENSION_KEY, listener1->updates[0].mKey);
+    EXPECT_EQ(2, listener1->updates[0].mState);
+
+    // check StateTracker was updated by querying for state
+    HashableDimensionKey queryKey = DEFAULT_DIMENSION_KEY;
+    EXPECT_EQ(2, mgr.getState(android::util::SCREEN_STATE_CHANGED, queryKey));
+}
+
+/**
+ * Test StateManager's onLogEvent and StateListener's onStateChanged correctly
+ * updates listener for states with primary keys.
+ */
+TEST(StateTrackerTest, TestStateChangeOnePrimaryField) {
+    sp<TestStateListener> listener1 = new TestStateListener();
+    StateManager mgr;
+    mgr.registerListener(android::util::UID_PROCESS_STATE_CHANGED, listener1);
+
+    // log event
+    std::shared_ptr<LogEvent> event = buildUidProcessEvent(
+            1000,
+            android::app::ProcessStateEnum::PROCESS_STATE_TOP);  //  state value: 1002
+    mgr.onLogEvent(*event);
+
+    // check listener was updated
+    EXPECT_EQ(1, listener1->updates.size());
+    EXPECT_EQ(1000, listener1->updates[0].mKey.getValues()[0].mValue.int_value);
+    EXPECT_EQ(1002, listener1->updates[0].mState);
+
+    // check StateTracker was updated by querying for state
+    HashableDimensionKey queryKey;
+    getUidProcessKey(1000, &queryKey);
+    EXPECT_EQ(1002, mgr.getState(android::util::UID_PROCESS_STATE_CHANGED, queryKey));
+}
+
+TEST(StateTrackerTest, TestStateChangeMultiplePrimaryFields) {
+    sp<TestStateListener> listener1 = new TestStateListener();
+    StateManager mgr;
+    mgr.registerListener(android::util::OVERLAY_STATE_CHANGED, listener1);
+
+    // log event
+    std::shared_ptr<LogEvent> event = buildOverlayEvent(1000, "package1", 1);  // state: ENTERED
+    mgr.onLogEvent(*event);
+
+    // check listener update
+    EXPECT_EQ(1, listener1->updates.size());
+    EXPECT_EQ(1000, listener1->updates[0].mKey.getValues()[0].mValue.int_value);
+    EXPECT_EQ(1, listener1->updates[0].mState);
+}
+
+/**
+ * Test StateManager's onLogEvent and StateListener's onStateChanged
+ * when there is an error extracting state from log event. Listener is not
+ * updated of state change.
+ */
+TEST(StateTrackerTest, TestStateChangeEventError) {
+    sp<TestStateListener> listener1 = new TestStateListener();
+    StateManager mgr;
+    mgr.registerListener(android::util::OVERLAY_STATE_CHANGED, listener1);
+
+    // log event
+    std::shared_ptr<LogEvent> event =
+            buildIncorrectOverlayEvent(1000, "package1", 1);  // state: ENTERED
+    mgr.onLogEvent(*event);
+
+    // check listener update
+    EXPECT_EQ(0, listener1->updates.size());
+}
+
+TEST(StateTrackerTest, TestStateQuery) {
+    sp<TestStateListener> listener1 = new TestStateListener();
+    sp<TestStateListener> listener2 = new TestStateListener();
+    sp<TestStateListener> listener3 = new TestStateListener();
+    StateManager mgr;
+    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener1);
+    mgr.registerListener(android::util::UID_PROCESS_STATE_CHANGED, listener2);
+    mgr.registerListener(android::util::OVERLAY_STATE_CHANGED, listener3);
+
+    std::shared_ptr<LogEvent> event1 = buildUidProcessEvent(
+            1000,
+            android::app::ProcessStateEnum::PROCESS_STATE_TOP);  // state value: 1002
+    std::shared_ptr<LogEvent> event2 = buildUidProcessEvent(
+            1001,
+            android::app::ProcessStateEnum::PROCESS_STATE_FOREGROUND_SERVICE);  // state value: 1003
+    std::shared_ptr<LogEvent> event3 = buildUidProcessEvent(
+            1002,
+            android::app::ProcessStateEnum::PROCESS_STATE_PERSISTENT);  // state value: 1000
+    std::shared_ptr<LogEvent> event4 = buildUidProcessEvent(
+            1001,
+            android::app::ProcessStateEnum::PROCESS_STATE_TOP);  // state value: 1002
+    std::shared_ptr<LogEvent> event5 =
+            buildScreenEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON);  // state value:
+    std::shared_ptr<LogEvent> event6 = buildOverlayEvent(1000, "package1", 1);
+    std::shared_ptr<LogEvent> event7 = buildOverlayEvent(1000, "package2", 2);
+
+    mgr.onLogEvent(*event1);
+    mgr.onLogEvent(*event2);
+    mgr.onLogEvent(*event3);
+    mgr.onLogEvent(*event5);
+    mgr.onLogEvent(*event5);
+    mgr.onLogEvent(*event6);
+    mgr.onLogEvent(*event7);
+
+    // Query for UidProcessState of uid 1001
+    HashableDimensionKey queryKey1;
+    getUidProcessKey(1001, &queryKey1);
+    EXPECT_EQ(1003, mgr.getState(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.getState(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.getState(android::util::UID_PROCESS_STATE_CHANGED, queryKey1));
+
+    // Query for ScreenState
+    EXPECT_EQ(2, mgr.getState(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.getState(android::util::OVERLAY_STATE_CHANGED, queryKey3));
+}
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android
+#else
+GTEST_LOG_(INFO) << "This test does nothing.\n";
+#endif