Allow default state, reset state, and nesting for binary states

Every state atom can have a customized default state and reset state by
annotating atoms.proto.

Binary state atoms (such as WakelockStateChanged)  can turn on nested
counting by annotating atoms.proto as well.

Generated atoms_info.h before change: https://paste.googleplex.com/4626190063108096
Generated atoms_info.h after change: https://paste.googleplex.com/5410938863747072
Generated atoms_info.cpp before change: https://paste.googleplex.com/5726061016907776
Generated atoms_info.cpp after change: https://paste.googleplex.com/5540983737417728

Test: bit statsd_test:*
Change-Id: I845616f103e013a7927de869b8e8228cfb244090
diff --git a/cmds/statsd/src/atom_field_options.proto b/cmds/statsd/src/atom_field_options.proto
index 946c550..9c875ba 100644
--- a/cmds/statsd/src/atom_field_options.proto
+++ b/cmds/statsd/src/atom_field_options.proto
@@ -24,46 +24,75 @@
 import "google/protobuf/descriptor.proto";
 
 enum StateField {
-    // Default value for fields that are not primary or exclusive state.
+    // Default value for fields that are not a primary or exclusive state field.
     STATE_FIELD_UNSET = 0;
     // Fields that represent the key that the state belongs to.
-    PRIMARY = 1;
+    // Used on simple proto fields. Do not use on attribution chains.
+    PRIMARY_FIELD = 1;
     // The field that represents the state. It's an exclusive state.
-    EXCLUSIVE = 2;
-
+    EXCLUSIVE_STATE = 2;
+    // Used on an attribution chain field to indicate that the first uid is the
+    // primary field.
     PRIMARY_FIELD_FIRST_UID = 3;
 }
 
-// Used to annotate an atom that reprsents a state change. A state change atom must have exactly ONE
-// exclusive state field, and any number of primary key fields.
-// For example,
-// message UidProcessStateChanged {
-//    optional int32 uid = 1 [(state_field_option).option = PRIMARY];
-//    optional android.app.ProcessStateEnum state = 2 [(state_field_option).option = EXCLUSIVE];
+// Used to annotate an atom that represents a state change. A state change atom must have exactly
+// ONE exclusive state field, and any number of primary key fields. For example, message
+// UidProcessStateChanged {
+//    optional int32 uid = 1 [(state_field_option).option = PRIMARY_FIELD];
+//    optional android.app.ProcessStateEnum state = 2 [(state_field_option).option =
+//    EXCLUSIVE_STATE];
 //  }
-// Each of this UidProcessStateChanged atom represents a state change for a specific uid.
+// Each UidProcessStateChanged atom event represents a state change for a specific uid.
 // A new state automatically overrides the previous state.
 //
-// If the atom has 2 or more primary fields, it means the combination of the primary fields are
-// the primary key.
+// If the atom has 2 or more primary fields, it means the combination of the
+// primary fields are the primary key.
 // For example:
 // message ThreadStateChanged {
-//    optional int32 pid = 1  [(state_field_option).option = PRIMARY];
-//    optional int32 tid = 2  [(state_field_option).option = PRIMARY];
-//    optional int32 state = 3 [(state_field_option).option = EXCLUSIVE];
+//    optional int32 pid = 1  [(state_field_option).option = PRIMARY_FIELD];
+//    optional int32 tid = 2  [(state_field_option).option = PRIMARY_FIELD];
+//    optional int32 state = 3 [(state_field_option).option = EXCLUSIVE_STATE];
 // }
 //
 // Sometimes, there is no primary key field, when the state is GLOBAL.
 // For example,
-//
 // message ScreenStateChanged {
-//    optional android.view.DisplayStateEnum state = 1 [(state_field_option).option = EXCLUSIVE];
+//    optional android.view.DisplayStateEnum state = 1 [(state_field_option).option =
+//    EXCLUSIVE_STATE];
 // }
 //
-// Only fields of primary types can be annotated. AttributionNode cannot be primary keys (and they
-// usually are not).
+// For state atoms with attribution chain, sometimes the primary key is the first uid in the chain.
+// For example:
+// message AudioStateChanged {
+//   repeated AttributionNode attribution_node = 1
+//       [(stateFieldOption).option = PRIMARY_KEY_FIRST_UID];
+//
+//    enum State {
+//      OFF = 0;
+//      ON = 1;
+//      // RESET indicates all audio stopped. Used when it (re)starts (e.g. after it crashes).
+//      RESET = 2;
+//    }
+//    optional State state = 2 [(stateFieldOption).option = EXCLUSIVE_STATE];
+// }
 message StateAtomFieldOption {
     optional StateField option = 1 [default = STATE_FIELD_UNSET];
+
+    // Note: We cannot annotate directly on the enums because many enums are imported from other
+    // proto files in the platform. proto-lite cc library does not support annotations unfortunately
+
+    // Knowing the default state value allows state trackers to remove entries that become the
+    // default state. If there is no default value specified, the default value is unknown, and all
+    // states will be tracked in memory.
+    optional int32 default_state_value = 2;
+
+    // A reset state signals all states go to default value. For example, BLE reset means all active
+    // BLE scans are to be turned off.
+    optional int32 reset_state_value = 3;
+
+    // If the state change needs to count nesting.
+    optional bool nested = 4 [default = true];
 }
 
 // Used to generate StatsLog.write APIs.
diff --git a/cmds/statsd/src/atoms.proto b/cmds/statsd/src/atoms.proto
index cbece78..4388a08 100644
--- a/cmds/statsd/src/atoms.proto
+++ b/cmds/statsd/src/atoms.proto
@@ -508,7 +508,8 @@
  */
 message ScreenStateChanged {
     // New screen state, from frameworks/base/core/proto/android/view/enums.proto.
-    optional android.view.DisplayStateEnum state = 1 [(state_field_option).option = EXCLUSIVE];
+    optional android.view.DisplayStateEnum state = 1
+            [(state_field_option).option = EXCLUSIVE_STATE, (state_field_option).nested = false];
 }
 
 /**
@@ -519,10 +520,11 @@
  *   frameworks/base/services/core/java/com/android/server/am/BatteryStatsService.java
  */
 message UidProcessStateChanged {
-    optional int32 uid = 1 [(state_field_option).option = PRIMARY, (is_uid) = true];
+    optional int32 uid = 1 [(state_field_option).option = PRIMARY_FIELD, (is_uid) = true];
 
     // The state, from frameworks/base/core/proto/android/app/enums.proto.
-    optional android.app.ProcessStateEnum state = 2 [(state_field_option).option = EXCLUSIVE];
+    optional android.app.ProcessStateEnum state = 2
+            [(state_field_option).option = EXCLUSIVE_STATE, (state_field_option).nested = false];
 }
 
 /**
@@ -554,7 +556,7 @@
         ASLEEP = 1;
         AWAKE = 2;
     }
-    optional State state = 1 [(state_field_option).option = EXCLUSIVE];
+    optional State state = 1 [(state_field_option).option = EXCLUSIVE_STATE];
 }
 
 /**
@@ -573,7 +575,7 @@
         CRITICAL = 4;   // critical memory.
 
     }
-    optional State factor = 1 [(state_field_option).option = EXCLUSIVE];
+    optional State factor = 1 [(state_field_option).option = EXCLUSIVE_STATE];
 }
 
 /**
@@ -660,7 +662,8 @@
  *   packages/apps/Bluetooth/src/com/android/bluetooth/gatt/AppScanStats.java
  */
 message BleScanStateChanged {
-    repeated AttributionNode attribution_node = 1;
+    repeated AttributionNode attribution_node = 1
+            [(state_field_option).option = PRIMARY_FIELD_FIRST_UID];
 
     enum State {
         OFF = 0;
@@ -668,14 +671,19 @@
         // RESET indicates all ble stopped. Used when it (re)starts (e.g. after it crashes).
         RESET = 2;
     }
-    optional State state = 2;
+    optional State state = 2 [
+        (state_field_option).option = EXCLUSIVE_STATE,
+        (state_field_option).default_state_value = 0 /* State.OFF */,
+        (state_field_option).reset_state_value = 2 /* State.RESET */,
+        (state_field_option).nested = true
+    ];
 
     // Does the scan have a filter.
-    optional bool is_filtered = 3;
+    optional bool is_filtered = 3 [(state_field_option).option = PRIMARY_FIELD];
     // Whether the scan is a CALLBACK_TYPE_FIRST_MATCH scan. Called 'background' scan internally.
-    optional bool is_first_match = 4;
+    optional bool is_first_match = 4 [(state_field_option).option = PRIMARY_FIELD];
     // Whether the scan set to piggy-back off the results of other scans (SCAN_MODE_OPPORTUNISTIC).
-    optional bool is_opportunistic = 5;
+    optional bool is_opportunistic = 5 [(state_field_option).option = PRIMARY_FIELD];
 }
 
 /**
@@ -919,11 +927,11 @@
 
     // The type (level) of the wakelock; e.g. a partial wakelock or a full wakelock.
     // From frameworks/base/core/proto/android/os/enums.proto.
-    optional android.os.WakeLockLevelEnum type = 2 [(state_field_option).option = PRIMARY];
+    optional android.os.WakeLockLevelEnum type = 2 [(state_field_option).option = PRIMARY_FIELD];
     ;
 
     // The wakelock tag (Called tag in the Java API, sometimes name elsewhere).
-    optional string tag = 3 [(state_field_option).option = PRIMARY];
+    optional string tag = 3 [(state_field_option).option = PRIMARY_FIELD];
 
     enum State {
         RELEASE = 0;
@@ -931,7 +939,11 @@
         CHANGE_RELEASE = 2;
         CHANGE_ACQUIRE = 3;
     }
-    optional State state = 4 [(state_field_option).option = EXCLUSIVE];
+    optional State state = 4 [
+        (state_field_option).option = EXCLUSIVE_STATE,
+        (state_field_option).default_state_value = 0,
+        (state_field_option).nested = true
+    ];
 }
 
 /**
@@ -3114,9 +3126,9 @@
  *     services/core/java/com/android/server/wm/Session.java
  */
 message OverlayStateChanged {
-    optional int32 uid = 1 [(state_field_option).option = PRIMARY, (is_uid) = true];
+    optional int32 uid = 1 [(state_field_option).option = PRIMARY_FIELD, (is_uid) = true];
 
-    optional string package_name = 2 [(state_field_option).option = PRIMARY];
+    optional string package_name = 2 [(state_field_option).option = PRIMARY_FIELD];
 
     optional bool using_alert_window = 3;
 
@@ -3124,7 +3136,11 @@
         ENTERED = 1;
         EXITED = 2;
     }
-    optional State state = 4 [(state_field_option).option = EXCLUSIVE];
+    optional State state = 4 [
+        (state_field_option).option = EXCLUSIVE_STATE,
+        (state_field_option).nested = false,
+        (state_field_option).default_state_value = 2
+    ];
 }
 
 /*
@@ -3290,7 +3306,7 @@
  */
 message AppDied {
     // timestamp(elapsedRealtime) of record creation
-    optional uint64 timestamp_millis = 1 [(state_field_option).option = EXCLUSIVE];
+    optional uint64 timestamp_millis = 1 [(state_field_option).option = EXCLUSIVE_STATE];
 }
 
 /**
@@ -3685,7 +3701,7 @@
         DIALOG_LINE_ITEM = 5;
     }
 
-    optional Type type = 1 [(state_field_option).option = EXCLUSIVE];
+    optional Type type = 1 [(state_field_option).option = EXCLUSIVE_STATE];
 
     // Used if the type is LINE_ITEM
     optional string package_name = 2;
diff --git a/cmds/statsd/src/state/StateTracker.cpp b/cmds/statsd/src/state/StateTracker.cpp
index 3ad21e0..ab86127 100644
--- a/cmds/statsd/src/state/StateTracker.cpp
+++ b/cmds/statsd/src/state/StateTracker.cpp
@@ -26,7 +26,9 @@
 namespace statsd {
 
 StateTracker::StateTracker(const int32_t atomId, const util::StateAtomFieldOptions& stateAtomInfo)
-    : mAtomId(atomId), mStateField(getSimpleMatcher(atomId, stateAtomInfo.exclusiveField)) {
+    : mAtomId(atomId),
+      mStateField(getSimpleMatcher(atomId, stateAtomInfo.exclusiveField)),
+      mNested(stateAtomInfo.nested) {
     // create matcher for each primary field
     for (const auto& primaryField : stateAtomInfo.primaryFields) {
         if (primaryField == util::FIRST_UID_IN_CHAIN) {
@@ -38,7 +40,13 @@
         }
     }
 
-    // TODO(tsaichristine): b/142108433 set default state, reset state, and nesting
+    if (stateAtomInfo.defaultState != util::UNSET_VALUE) {
+        mDefaultState = stateAtomInfo.defaultState;
+    }
+
+    if (stateAtomInfo.resetState != util::UNSET_VALUE) {
+        mResetState = stateAtomInfo.resetState;
+    }
 }
 
 void StateTracker::onLogEvent(const LogEvent& event) {
@@ -60,7 +68,6 @@
 
     // Parse event for state value.
     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",
@@ -68,11 +75,12 @@
         handlePartialReset(eventTimeNs, primaryKey);
         return;
     }
-    state = stateValue.mValue.int_value;
 
+    int32_t state = stateValue.mValue.int_value;
     if (state == mResetState) {
         VLOG("StateTracker Reset state: %s", stateValue.mValue.toString().c_str());
         handleReset(eventTimeNs);
+        return;
     }
 
     // Track and update state.
@@ -113,15 +121,17 @@
             return true;
         }
     } else if (queryKey.getValues().size() > mPrimaryFields.size()) {
-        ALOGE("StateTracker query key size > primary key size is illegal");
+        ALOGE("StateTracker query key size %zu > primary key size %zu is illegal",
+              queryKey.getValues().size(), mPrimaryFields.size());
     } else {
-        ALOGE("StateTracker query key size < primary key size is not supported");
+        ALOGE("StateTracker query key size %zu < primary key size %zu is not supported",
+              queryKey.getValues().size(), mPrimaryFields.size());
     }
 
-    // Set the state value to unknown if:
+    // Set the state value to default state if:
     // - query key size is incorrect
     // - query key is not found in state map
-    output->mValue = StateTracker::kStateUnknown;
+    output->mValue = mDefaultState;
     return false;
 }
 
@@ -164,18 +174,52 @@
         *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;
+    // Update state map for non-nested counting case.
+    // Every state event triggers a state overwrite.
+    if (!mNested) {
+        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;
+        return;
     }
-    *newState = eventState;
 
-    // TODO(tsaichristine): support atoms with nested counting
+    // Update state map for nested counting case.
+    //
+    // Nested counting is only allowed for binary state events such as ON/OFF or
+    // ACQUIRE/RELEASE. For example, WakelockStateChanged might have the state
+    // events: ON, ON, OFF. The state will still be ON until we see the same
+    // number of OFF events as ON events.
+    //
+    // In atoms.proto, a state atom with nested counting enabled
+    // must only have 2 states and one of the states must be the default state.
+    it = mStateMap.find(primaryKey);
+    if (it != mStateMap.end()) {
+        *newState = it->second.state;
+        if (eventState == it->second.state) {
+            it->second.count++;
+        } else if (eventState == mDefaultState) {
+            if ((--it->second.count) == 0) {
+                mStateMap.erase(primaryKey);
+                *newState = mDefaultState;
+            }
+        } else {
+            ALOGE("StateTracker Nest counting state has a third state instead of the binary state "
+                  "limit.");
+            return;
+        }
+    } else {
+        if (eventState != mDefaultState) {
+            mStateMap[primaryKey].state = eventState;
+            mStateMap[primaryKey].count = 1;
+        }
+        *newState = eventState;
+    }
 }
 
 }  // namespace statsd
diff --git a/cmds/statsd/src/state/StateTracker.h b/cmds/statsd/src/state/StateTracker.h
index 70f1627..aeca2a5 100644
--- a/cmds/statsd/src/state/StateTracker.h
+++ b/cmds/statsd/src/state/StateTracker.h
@@ -74,6 +74,8 @@
 
     int32_t mResetState = kStateUnknown;
 
+    const bool mNested;
+
     // Maps primary key to state value info
     std::unordered_map<HashableDimensionKey, StateValueInfo> mStateMap;