Add a pulled metric for Wakelock duration based on uptime

Flag: com.android.server.power.feature.flags.framework_wakelock_info
Bug: 352602149
Test: atest PowerStatsTests:BatteryStatsPulledMetricsTest
Change-Id: I65878d769f70abeceafd5c166747db7126842a6d

6 files changed, 803 insertions, 1 deletions

diff --git a/services/core/java/com/android/server/power/feature/PowerManagerFlags.java b/services/core/java/com/android/server/power/feature/PowerManagerFlags.java
index fd60e06c0762..db57d11a75c0 100644
--- a/services/core/java/com/android/server/power/feature/PowerManagerFlags.java
+++ b/services/core/java/com/android/server/power/feature/PowerManagerFlags.java
@@ -47,6 +47,9 @@ public class PowerManagerFlags {
             Flags::perDisplayWakeByTouch
     );
 
+    private final FlagState mFrameworkWakelockInfo =
+            new FlagState(Flags.FLAG_FRAMEWORK_WAKELOCK_INFO, Flags::frameworkWakelockInfo);
+
     /** Returns whether early-screen-timeout-detector is enabled on not. */
     public boolean isEarlyScreenTimeoutDetectorEnabled() {
         return mEarlyScreenTimeoutDetectorFlagState.isEnabled();
@@ -67,6 +70,13 @@ public class PowerManagerFlags {
     }
 
     /**
+     * @return Whether FrameworkWakelockInfo atom logging is enabled or not.
+     */
+    public boolean isFrameworkWakelockInfoEnabled() {
+        return mFrameworkWakelockInfo.isEnabled();
+    }
+
+    /**
      * dumps all flagstates
      * @param pw printWriter
      */
@@ -75,6 +85,7 @@ public class PowerManagerFlags {
         pw.println(" " + mEarlyScreenTimeoutDetectorFlagState);
         pw.println(" " + mImproveWakelockLatency);
         pw.println(" " + mPerDisplayWakeByTouch);
+        pw.println(" " + mFrameworkWakelockInfo);
     }
 
     private static class FlagState {
diff --git a/services/core/java/com/android/server/power/feature/power_flags.aconfig b/services/core/java/com/android/server/power/feature/power_flags.aconfig
index 9cf3bb6df3db..8bb69ba0c5de 100644
--- a/services/core/java/com/android/server/power/feature/power_flags.aconfig
+++ b/services/core/java/com/android/server/power/feature/power_flags.aconfig
@@ -26,3 +26,10 @@ flag {
     bug: "343295183"
     is_fixed_read_only: true
 }
+
+flag {
+    name: "framework_wakelock_info"
+    namespace: "power"
+    description: "Feature flag to enable statsd pulling of FrameworkWakelockInfo atoms"
+    bug: "352602149"
+}
diff --git a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
index 936fadf0b12a..391071ff9fe8 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -145,6 +145,7 @@ import com.android.internal.util.XmlUtils;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
 import com.android.server.LocalServices;
+import com.android.server.power.feature.PowerManagerFlags;
 import com.android.server.power.optimization.Flags;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
 import com.android.server.power.stats.format.MobileRadioPowerStatsLayout;
@@ -5142,6 +5143,10 @@ public class BatteryStatsImpl extends BatteryStats {
             mFrameworkStatsLogger.wakelockStateChanged(mapIsolatedUid(uid), wc, name,
                     uidStats.mProcessState, true /* acquired */,
                     getPowerManagerWakeLockLevel(type));
+            if (mPowerManagerFlags.isFrameworkWakelockInfoEnabled()) {
+                mFrameworkEvents.noteStartWakeLock(
+                        mapIsolatedUid(uid), name, getPowerManagerWakeLockLevel(type), uptimeMs);
+            }
         }
     }
 
@@ -5187,6 +5192,10 @@ public class BatteryStatsImpl extends BatteryStats {
             mFrameworkStatsLogger.wakelockStateChanged(mapIsolatedUid(uid), wc, name,
                     uidStats.mProcessState, false/* acquired */,
                     getPowerManagerWakeLockLevel(type));
+            if (mPowerManagerFlags.isFrameworkWakelockInfoEnabled()) {
+                mFrameworkEvents.noteStopWakeLock(
+                        mapIsolatedUid(uid), name, getPowerManagerWakeLockLevel(type), uptimeMs);
+            }
 
             if (mappedUid != uid) {
                 // Decrement the ref count for the isolated uid and delete the mapping if uneeded.
@@ -11343,6 +11352,9 @@ public class BatteryStatsImpl extends BatteryStats {
         return mTmpCpuTimeInFreq;
     }
 
+    WakelockStatsFrameworkEvents mFrameworkEvents = new WakelockStatsFrameworkEvents();
+    PowerManagerFlags mPowerManagerFlags = new PowerManagerFlags();
+
     public BatteryStatsImpl(@NonNull BatteryStatsConfig config, @NonNull Clock clock,
             @NonNull MonotonicClock monotonicClock, @Nullable File systemDir,
             @NonNull Handler handler, @Nullable PlatformIdleStateCallback platformIdleStateCallback,
@@ -15964,6 +15976,10 @@ public class BatteryStatsImpl extends BatteryStats {
                 // Already plugged in. Schedule the long plug in alarm.
                 scheduleNextResetWhilePluggedInCheck();
             }
+
+            if (mPowerManagerFlags.isFrameworkWakelockInfoEnabled()) {
+                mFrameworkEvents.initialize(context);
+            }
         }
     }
 
@@ -16791,7 +16807,8 @@ public class BatteryStatsImpl extends BatteryStats {
         }
         int NSORPMS = in.readInt();
         if (NSORPMS > 10000) {
-            throw new ParcelFormatException("File corrupt: too many screen-off rpm stats " + NSORPMS);
+            throw new ParcelFormatException(
+                    "File corrupt: too many screen-off rpm stats " + NSORPMS);
         }
         for (int irpm = 0; irpm < NSORPMS; irpm++) {
             if (in.readInt() != 0) {
diff --git a/services/core/java/com/android/server/power/stats/WakelockStatsFrameworkEvents.java b/services/core/java/com/android/server/power/stats/WakelockStatsFrameworkEvents.java
new file mode 100644
index 000000000000..c9693bd20a08
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/WakelockStatsFrameworkEvents.java
@@ -0,0 +1,328 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+package com.android.server.power.stats;
+
+import android.app.StatsManager;
+import android.content.Context;
+import android.os.SystemClock;
+import android.util.Log;
+import android.util.StatsEvent;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.util.ConcurrentUtils;
+import com.android.internal.util.FrameworkStatsLog;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+/** A class to initialise and log metrics pulled by statsd. */
+public class WakelockStatsFrameworkEvents {
+    // statsd has a dimensional limit on the number of different keys it can handle.
+    // Beyond that limit, statsd will drop data.
+    //
+    // When we have seem SUMMARY_THRESHOLD distinct (uid, tag, wakeLockLevel) keys,
+    // we start summarizing new keys as (uid, OVERFLOW_TAG, OVERFLOW_LEVEL) to
+    // reduce the number of keys we pass to statsd.
+    //
+    // When we reach MAX_WAKELOCK_DIMENSIONS distinct keys, we summarize all new keys
+    // as (OVERFLOW_UID, HARD_CAP_TAG, OVERFLOW_LEVEL) to hard cap the number of
+    // distinct keys we pass to statsd.
+    @VisibleForTesting public static final int SUMMARY_THRESHOLD = 500;
+    @VisibleForTesting public static final int MAX_WAKELOCK_DIMENSIONS = 1000;
+
+    @VisibleForTesting public static final int HARD_CAP_UID = -1;
+    @VisibleForTesting public static final String OVERFLOW_TAG = "*overflow*";
+    @VisibleForTesting public static final String HARD_CAP_TAG = "*overflow hard cap*";
+    @VisibleForTesting public static final int OVERFLOW_LEVEL = 1;
+
+    private static class WakeLockKey {
+        private int uid;
+        private String tag;
+        private int powerManagerWakeLockLevel;
+        private int hashCode;
+
+        WakeLockKey(int uid, String tag, int powerManagerWakeLockLevel) {
+            this.uid = uid;
+            this.tag = new String(tag);
+            this.powerManagerWakeLockLevel = powerManagerWakeLockLevel;
+
+            this.hashCode = Objects.hash(uid, tag, powerManagerWakeLockLevel);
+        }
+
+        int getUid() {
+            return uid;
+        }
+
+        String getTag() {
+            return tag;
+        }
+
+        int getPowerManagerWakeLockLevel() {
+            return powerManagerWakeLockLevel;
+        }
+
+        void setOverflow() {
+            tag = OVERFLOW_TAG;
+            powerManagerWakeLockLevel = OVERFLOW_LEVEL;
+            this.hashCode = Objects.hash(uid, tag, powerManagerWakeLockLevel);
+        }
+
+        void setHardCap() {
+            uid = HARD_CAP_UID;
+            tag = HARD_CAP_TAG;
+            powerManagerWakeLockLevel = OVERFLOW_LEVEL;
+            this.hashCode = Objects.hash(uid, tag, powerManagerWakeLockLevel);
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || !(o instanceof WakeLockKey)) return false;
+
+            WakeLockKey that = (WakeLockKey) o;
+            return uid == that.uid
+                    && tag.equals(that.tag)
+                    && powerManagerWakeLockLevel == that.powerManagerWakeLockLevel;
+        }
+
+        @Override
+        public int hashCode() {
+            return this.hashCode;
+        }
+    }
+
+    private static class WakeLockStats {
+        // accumulated uptime attributed to this WakeLock since boot, where overlap
+        // (including nesting) is ignored
+        public long uptimeMillis = 0;
+
+        // count of WakeLocks that have been acquired and then released
+        public long completedCount = 0;
+    }
+
+    private final Object mLock = new Object();
+
+    @GuardedBy("mLock")
+    private final Map<WakeLockKey, WakeLockStats> mWakeLockStats = new HashMap<>();
+
+    private static class WakeLockData {
+        // uptime millis when first acquired
+        public long acquireUptimeMillis = 0;
+        public int refCount = 0;
+
+        WakeLockData(long uptimeMillis) {
+            acquireUptimeMillis = uptimeMillis;
+        }
+    }
+
+    @GuardedBy("mLock")
+    private final Map<WakeLockKey, WakeLockData> mOpenWakeLocks = new HashMap<>();
+
+    public void noteStartWakeLock(
+            int uid, String tag, int powerManagerWakeLockLevel, long eventUptimeMillis) {
+        final WakeLockKey key = new WakeLockKey(uid, tag, powerManagerWakeLockLevel);
+
+        synchronized (mLock) {
+            WakeLockData data =
+                    mOpenWakeLocks.computeIfAbsent(key, k -> new WakeLockData(eventUptimeMillis));
+            data.refCount++;
+            mOpenWakeLocks.put(key, data);
+        }
+    }
+
+    @VisibleForTesting
+    public boolean inOverflow() {
+        synchronized (mLock) {
+            return inOverflowLocked();
+        }
+    }
+
+    @GuardedBy("mLock")
+    private boolean inOverflowLocked() {
+        return mWakeLockStats.size() >= SUMMARY_THRESHOLD;
+    }
+
+    @VisibleForTesting
+    public boolean inHardCap() {
+        synchronized (mLock) {
+            return inHardCapLocked();
+        }
+    }
+
+    @GuardedBy("mLock")
+    private boolean inHardCapLocked() {
+        return mWakeLockStats.size() >= MAX_WAKELOCK_DIMENSIONS;
+    }
+
+    public void noteStopWakeLock(
+            int uid, String tag, int powerManagerWakeLockLevel, long eventUptimeMillis) {
+        WakeLockKey key = new WakeLockKey(uid, tag, powerManagerWakeLockLevel);
+
+        synchronized (mLock) {
+            WakeLockData data = mOpenWakeLocks.get(key);
+            if (data == null) {
+                Log.e(TAG, "WakeLock not found when stopping: " + uid + " " + tag);
+                return;
+            }
+
+            if (data.refCount == 1) {
+                mOpenWakeLocks.remove(key);
+                long wakeLockDur = eventUptimeMillis - data.acquireUptimeMillis;
+
+                // Rewrite key if in an overflow state.
+                if (inOverflowLocked() && !mWakeLockStats.containsKey(key)) {
+                    key.setOverflow();
+                    if (inHardCapLocked() && !mWakeLockStats.containsKey(key)) {
+                        key.setHardCap();
+                    }
+                }
+
+                WakeLockStats stats = mWakeLockStats.computeIfAbsent(key, k -> new WakeLockStats());
+                stats.uptimeMillis += wakeLockDur;
+                stats.completedCount++;
+                mWakeLockStats.put(key, stats);
+            } else {
+                data.refCount--;
+                mOpenWakeLocks.put(key, data);
+            }
+        }
+    }
+
+    public List<StatsEvent> pullFrameworkWakelockInfoAtoms() {
+        return pullFrameworkWakelockInfoAtoms(SystemClock.uptimeMillis());
+    }
+
+    public List<StatsEvent> pullFrameworkWakelockInfoAtoms(long nowMillis) {
+        List<StatsEvent> result = new ArrayList<>();
+        HashSet<WakeLockKey> keys = new HashSet<>();
+
+        // Used to collect open WakeLocks when in an overflow state.
+        HashMap<WakeLockKey, WakeLockStats> openOverflowStats = new HashMap<>();
+
+        synchronized (mLock) {
+            keys.addAll(mWakeLockStats.keySet());
+
+            // If we are in an overflow state, an open wakelock may have a new key
+            // that needs to be summarized.
+            if (inOverflowLocked()) {
+                for (WakeLockKey key : mOpenWakeLocks.keySet()) {
+                    if (!mWakeLockStats.containsKey(key)) {
+                        WakeLockData data = mOpenWakeLocks.get(key);
+
+                        key.setOverflow();
+                        if (inHardCapLocked() && !mWakeLockStats.containsKey(key)) {
+                            key.setHardCap();
+                        }
+                        keys.add(key);
+
+                        WakeLockStats stats =
+                                openOverflowStats.computeIfAbsent(key, k -> new WakeLockStats());
+                        stats.uptimeMillis += nowMillis - data.acquireUptimeMillis;
+                        openOverflowStats.put(key, stats);
+                    }
+                }
+            } else {
+                keys.addAll(mOpenWakeLocks.keySet());
+            }
+
+            for (WakeLockKey key : keys) {
+                long openWakeLockUptime = 0;
+                WakeLockData data = mOpenWakeLocks.get(key);
+                if (data != null) {
+                    openWakeLockUptime = nowMillis - data.acquireUptimeMillis;
+                }
+
+                WakeLockStats stats = mWakeLockStats.computeIfAbsent(key, k -> new WakeLockStats());
+                WakeLockStats extraTime =
+                        openOverflowStats.computeIfAbsent(key, k -> new WakeLockStats());
+
+                stats.uptimeMillis += openWakeLockUptime + extraTime.uptimeMillis;
+
+                StatsEvent event =
+                        StatsEvent.newBuilder()
+                                .setAtomId(FrameworkStatsLog.FRAMEWORK_WAKELOCK_INFO)
+                                .writeInt(key.getUid())
+                                .writeString(key.getTag())
+                                .writeInt(key.getPowerManagerWakeLockLevel())
+                                .writeLong(stats.uptimeMillis)
+                                .writeLong(stats.completedCount)
+                                .build();
+                result.add(event);
+            }
+        }
+
+        return result;
+    }
+
+    private static final String TAG = "BatteryStatsPulledMetrics";
+
+    private final StatsPullCallbackHandler mStatsPullCallbackHandler =
+            new StatsPullCallbackHandler();
+
+    private boolean mIsInitialized = false;
+
+    public void initialize(Context context) {
+        if (mIsInitialized) {
+            return;
+        }
+
+        final StatsManager statsManager = context.getSystemService(StatsManager.class);
+        if (statsManager == null) {
+            Log.e(
+                    TAG,
+                    "Error retrieving StatsManager. Cannot initialize BatteryStatsPulledMetrics.");
+        } else {
+            Log.d(TAG, "Registering callback with StatsManager");
+
+            // DIRECT_EXECUTOR means that callback will run on binder thread.
+            statsManager.setPullAtomCallback(
+                    FrameworkStatsLog.FRAMEWORK_WAKELOCK_INFO,
+                    null /* metadata */,
+                    ConcurrentUtils.DIRECT_EXECUTOR,
+                    mStatsPullCallbackHandler);
+            mIsInitialized = true;
+        }
+    }
+
+    private class StatsPullCallbackHandler implements StatsManager.StatsPullAtomCallback {
+        @Override
+        public int onPullAtom(int atomTag, List<StatsEvent> data) {
+            // handle the tags appropriately.
+            List<StatsEvent> events = pullEvents(atomTag);
+            if (events == null) {
+                return StatsManager.PULL_SKIP;
+            }
+
+            data.addAll(events);
+            return StatsManager.PULL_SUCCESS;
+        }
+
+        private List<StatsEvent> pullEvents(int atomTag) {
+            switch (atomTag) {
+                case FrameworkStatsLog.FRAMEWORK_WAKELOCK_INFO:
+                    return pullFrameworkWakelockInfoAtoms();
+                default:
+                    return null;
+            }
+        }
+    }
+}
diff --git a/services/tests/powerstatstests/Android.bp b/services/tests/powerstatstests/Android.bp
index d6ca10a23fb9..fab20315913d 100644
--- a/services/tests/powerstatstests/Android.bp
+++ b/services/tests/powerstatstests/Android.bp
@@ -27,6 +27,9 @@ android_test {
         "servicestests-utils",
         "platform-test-annotations",
         "flag-junit",
+        "statsdprotolite",
+        "StatsdTestUtils",
+        "platformprotoslite",
     ],
 
     libs: [
@@ -68,6 +71,9 @@ android_ravenwood_test {
         "androidx.test.uiautomator_uiautomator",
         "modules-utils-binary-xml",
         "flag-junit",
+        "statsdprotolite",
+        "StatsdTestUtils",
+        "platformprotoslite",
     ],
     srcs: [
         "src/com/android/server/power/stats/*.java",
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java
new file mode 100644
index 000000000000..cb644dbe8c96
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java
@@ -0,0 +1,433 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+package com.android.server.power.stats;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import android.os.WakeLockLevelEnum;
+import android.util.StatsEvent;
+import android.util.StatsEventTestUtils;
+
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.os.AtomsProto;
+import com.android.os.framework.FrameworkExtensionAtoms;
+import com.android.os.framework.FrameworkExtensionAtoms.FrameworkWakelockInfo;
+
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import com.google.protobuf.ExtensionRegistryLite;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class WakelockStatsFrameworkEventsTest {
+    private WakelockStatsFrameworkEvents mEvents;
+    private ExtensionRegistryLite mRegistry;
+
+    @Before
+    public void setup() {
+        mEvents = new WakelockStatsFrameworkEvents();
+        mRegistry = ExtensionRegistryLite.newInstance();
+        FrameworkExtensionAtoms.registerAllExtensions(mRegistry);
+    }
+
+    private static final int UID_1 = 1;
+    private static final int UID_2 = 2;
+
+    private static final String TAG_1 = "TAG1";
+    private static final String TAG_2 = "TAG2";
+
+    private static final WakeLockLevelEnum WAKELOCK_TYPE_1 = WakeLockLevelEnum.PARTIAL_WAKE_LOCK;
+    private static final WakeLockLevelEnum WAKELOCK_TYPE_2 = WakeLockLevelEnum.DOZE_WAKE_LOCK;
+
+    private static final long TS_1 = 1000;
+    private static final long TS_2 = 2000;
+    private static final long TS_3 = 3000;
+    private static final long TS_4 = 4000;
+    private static final long TS_5 = 5000;
+
+    // Assumes that mEvents is empty.
+    private void makeMetricsAlmostOverflow() throws Exception {
+        for (int i = 0; i < mEvents.SUMMARY_THRESHOLD - 1; i++) {
+            String tag = "forceOverflow" + i;
+            mEvents.noteStartWakeLock(UID_1, tag, WAKELOCK_TYPE_1.getNumber(), TS_1);
+            mEvents.noteStopWakeLock(UID_1, tag, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        }
+
+        assertFalse("not overflow", mEvents.inOverflow());
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo notOverflowInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.OVERFLOW_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("not overflow", notOverflowInfo, null);
+
+        // Add one more to hit an overflow state.
+        String lastTag = "forceOverflowLast";
+        mEvents.noteStartWakeLock(UID_1, lastTag, WAKELOCK_TYPE_2.getNumber(), TS_1);
+        mEvents.noteStopWakeLock(UID_1, lastTag, WAKELOCK_TYPE_2.getNumber(), TS_2);
+
+        assertTrue("overflow", mEvents.inOverflow());
+        info = pullResults(TS_4);
+
+        FrameworkWakelockInfo tag1Info =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(lastTag))
+                        .findFirst()
+                        .orElse(null);
+
+        assertTrue("lastTag found", tag1Info != null);
+        assertEquals("uid", UID_1, tag1Info.getAttributionUid());
+        assertEquals("tag", lastTag, tag1Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_2, tag1Info.getType());
+        assertEquals("duration", TS_2 - TS_1, tag1Info.getUptimeMillis());
+        assertEquals("count", 1, tag1Info.getCompletedCount());
+    }
+
+    // Assumes that mEvents is empty.
+    private void makeMetricsAlmostHardCap() throws Exception {
+        for (int i = 0; i < mEvents.MAX_WAKELOCK_DIMENSIONS - 1; i++) {
+            mEvents.noteStartWakeLock(i /* uid */, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+            mEvents.noteStopWakeLock(i /* uid */, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        }
+
+        assertFalse("not hard capped", mEvents.inHardCap());
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo notOverflowInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.HARD_CAP_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("not overflow", notOverflowInfo, null);
+
+        // Add one more to hit an hardcap state.
+        int hardCapUid = mEvents.MAX_WAKELOCK_DIMENSIONS;
+        mEvents.noteStartWakeLock(hardCapUid, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+        mEvents.noteStopWakeLock(hardCapUid, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_2);
+
+        assertTrue("hard capped", mEvents.inHardCap());
+        info = pullResults(TS_4);
+
+        FrameworkWakelockInfo tag2Info =
+                info.stream()
+                        .filter(i -> i.getAttributionUid() == hardCapUid)
+                        .findFirst()
+                        .orElse(null);
+
+        assertTrue("hardCapUid found", tag2Info != null);
+        assertEquals("uid", hardCapUid, tag2Info.getAttributionUid());
+        assertEquals("tag", mEvents.OVERFLOW_TAG, tag2Info.getAttributionTag());
+        assertEquals(
+                "type", WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL), tag2Info.getType());
+        assertEquals("duration", TS_2 - TS_1, tag2Info.getUptimeMillis());
+        assertEquals("count", 1, tag2Info.getCompletedCount());
+    }
+
+    private ArrayList<FrameworkWakelockInfo> pullResults(long timestamp) throws Exception {
+        ArrayList<FrameworkWakelockInfo> result = new ArrayList<>();
+        List<StatsEvent> events = mEvents.pullFrameworkWakelockInfoAtoms(timestamp);
+
+        for (StatsEvent e : events) {
+            // The returned atom does not have external extensions registered.
+            // So we serialize and then deserialize with extensions registered.
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(e);
+
+            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+            CodedOutputStream codedos = CodedOutputStream.newInstance(outputStream);
+            atom.writeTo(codedos);
+            codedos.flush();
+
+            ByteArrayInputStream inputStream = new ByteArrayInputStream(outputStream.toByteArray());
+            CodedInputStream codedis = CodedInputStream.newInstance(inputStream);
+            AtomsProto.Atom atomWithExtensions = AtomsProto.Atom.parseFrom(codedis, mRegistry);
+
+            assertTrue(
+                    atomWithExtensions.hasExtension(FrameworkExtensionAtoms.frameworkWakelockInfo));
+            FrameworkWakelockInfo info =
+                    atomWithExtensions.getExtension(FrameworkExtensionAtoms.frameworkWakelockInfo);
+            result.add(info);
+        }
+
+        return result;
+    }
+
+    @Test
+    public void singleWakelock() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_3);
+
+        assertEquals("size", 1, info.size());
+        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
+        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
+        assertEquals("duration", TS_2 - TS_1, info.get(0).getUptimeMillis());
+        assertEquals("count", 1, info.get(0).getCompletedCount());
+    }
+
+    @Test
+    public void wakelockOpen() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_3);
+
+        assertEquals("size", 1, info.size());
+        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
+        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
+        assertEquals("duration", TS_3 - TS_1, info.get(0).getUptimeMillis());
+        assertEquals("count", 0, info.get(0).getCompletedCount());
+    }
+
+    @Test
+    public void wakelockOpenOverlap() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_3);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+
+        assertEquals("size", 1, info.size());
+        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
+        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
+        assertEquals("duration", TS_4 - TS_1, info.get(0).getUptimeMillis());
+        assertEquals("count", 0, info.get(0).getCompletedCount());
+    }
+
+    @Test
+    public void testOverflow() throws Exception {
+        makeMetricsAlmostOverflow();
+
+        // This one gets tagged as an overflow.
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_2);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo overflowInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.OVERFLOW_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("uid", UID_1, overflowInfo.getAttributionUid());
+        assertEquals(
+                "type",
+                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
+                overflowInfo.getType());
+        assertEquals("duration", TS_2 - TS_1, overflowInfo.getUptimeMillis());
+        assertEquals("count", 1, overflowInfo.getCompletedCount());
+    }
+
+    @Test
+    public void testOverflowOpen() throws Exception {
+        makeMetricsAlmostOverflow();
+
+        // This is the open wakelock that overflows.
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo overflowInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.OVERFLOW_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("uid", UID_1, overflowInfo.getAttributionUid());
+        assertEquals(
+                "type",
+                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
+                overflowInfo.getType());
+        assertEquals("duration", (TS_4 - TS_1), overflowInfo.getUptimeMillis());
+        assertEquals("count", 0, overflowInfo.getCompletedCount());
+    }
+
+    @Test
+    public void testHardCap() throws Exception {
+        makeMetricsAlmostHardCap();
+
+        // This one gets tagged as a hard cap.
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_2);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo hardCapInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.HARD_CAP_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("uid", mEvents.HARD_CAP_UID, hardCapInfo.getAttributionUid());
+        assertEquals(
+                "type",
+                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
+                hardCapInfo.getType());
+        assertEquals("duration", TS_2 - TS_1, hardCapInfo.getUptimeMillis());
+        assertEquals("count", 1, hardCapInfo.getCompletedCount());
+    }
+
+    @Test
+    public void testHardCapOpen() throws Exception {
+        makeMetricsAlmostHardCap();
+
+        // This is the open wakelock that overflows.
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo hardCapInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.HARD_CAP_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("uid", mEvents.HARD_CAP_UID, hardCapInfo.getAttributionUid());
+        assertEquals(
+                "type",
+                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
+                hardCapInfo.getType());
+        assertEquals("duration", (TS_4 - TS_1), hardCapInfo.getUptimeMillis());
+        assertEquals("count", 0, hardCapInfo.getCompletedCount());
+    }
+
+    @Test
+    public void overlappingWakelocks() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+
+        assertEquals("size", 1, info.size());
+        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
+        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
+        assertEquals("duration", TS_4 - TS_1, info.get(0).getUptimeMillis());
+        assertEquals("count", 1, info.get(0).getCompletedCount());
+    }
+
+    @Test
+    public void diffUid() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_2, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        mEvents.noteStopWakeLock(UID_2, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+        assertEquals("size", 2, info.size());
+
+        FrameworkWakelockInfo uid1Info =
+                info.stream().filter(i -> i.getAttributionUid() == UID_1).findFirst().orElse(null);
+
+        assertTrue("UID_1 found", uid1Info != null);
+        assertEquals("uid", UID_1, uid1Info.getAttributionUid());
+        assertEquals("tag", TAG_1, uid1Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.getType());
+        assertEquals("duration", TS_4 - TS_1, uid1Info.getUptimeMillis());
+        assertEquals("count", 1, uid1Info.getCompletedCount());
+
+        FrameworkWakelockInfo uid2Info =
+                info.stream().filter(i -> i.getAttributionUid() == UID_2).findFirst().orElse(null);
+        assertTrue("UID_2 found", uid2Info != null);
+        assertEquals("uid", UID_2, uid2Info.getAttributionUid());
+        assertEquals("tag", TAG_1, uid2Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, uid2Info.getType());
+        assertEquals("duration", TS_3 - TS_2, uid2Info.getUptimeMillis());
+        assertEquals("count", 1, uid2Info.getCompletedCount());
+    }
+
+    @Test
+    public void diffTag() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_1.getNumber(), TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+        assertEquals("size", 2, info.size());
+
+        FrameworkWakelockInfo uid1Info =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(TAG_1))
+                        .findFirst()
+                        .orElse(null);
+
+        assertTrue("TAG_1 found", uid1Info != null);
+        assertEquals("uid", UID_1, uid1Info.getAttributionUid());
+        assertEquals("tag", TAG_1, uid1Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.getType());
+        assertEquals("duration", TS_4 - TS_1, uid1Info.getUptimeMillis());
+        assertEquals("count", 1, uid1Info.getCompletedCount());
+
+        FrameworkWakelockInfo uid2Info =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(TAG_2))
+                        .findFirst()
+                        .orElse(null);
+        assertTrue("TAG_2 found", uid2Info != null);
+        assertEquals("uid", UID_1, uid2Info.getAttributionUid());
+        assertEquals("tag", TAG_2, uid2Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, uid2Info.getType());
+        assertEquals("duration", TS_3 - TS_2, uid2Info.getUptimeMillis());
+        assertEquals("count", 1, uid2Info.getCompletedCount());
+    }
+
+    @Test
+    public void diffType() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_2.getNumber(), TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_2.getNumber(), TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+        assertEquals("size", 2, info.size());
+
+        FrameworkWakelockInfo uid1Info =
+                info.stream().filter(i -> i.getType() == WAKELOCK_TYPE_1).findFirst().orElse(null);
+
+        assertTrue("WAKELOCK_TYPE_1 found", uid1Info != null);
+        assertEquals("uid", UID_1, uid1Info.getAttributionUid());
+        assertEquals("tag", TAG_1, uid1Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.getType());
+        assertEquals("duration", TS_4 - TS_1, uid1Info.getUptimeMillis());
+        assertEquals("count", 1, uid1Info.getCompletedCount());
+
+        FrameworkWakelockInfo uid2Info =
+                info.stream().filter(i -> i.getType() == WAKELOCK_TYPE_2).findFirst().orElse(null);
+        assertTrue("WAKELOCK_TYPE_2 found", uid2Info != null);
+        assertEquals("uid", UID_1, uid2Info.getAttributionUid());
+        assertEquals("tag", TAG_1, uid2Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_2, uid2Info.getType());
+        assertEquals("duration", TS_3 - TS_2, uid2Info.getUptimeMillis());
+        assertEquals("count", 1, uid2Info.getCompletedCount());
+    }
+}