Add AOD metrics in batterystats

Fixes: 64899521
Test: manual - flash build, reset batteryStats, use device > 1 hr
with alternating pattern between screen on/off/AOD. Dump stats, check
all screen related stats look normal, esp. record matches actual time
spent in each screen state. In raw bugreport:
Search "amount discharged" for % discharge;
Search "Screen on/off/doze discharge" for mAh discharge;
Search "time on battery" for up/real time in each state.

Test: Added two unit tests for note AOD screen state

Merged-In: I7193a36751124dd380818b2b665303c0f0d8c984
Change-Id: I51cead7f92abd9e4c620f7dfde393993cdad494e

6 files changed, 404 insertions, 155 deletions

diff --git a/core/java/android/os/BatteryStats.java b/core/java/android/os/BatteryStats.java
index 66b6b47887b9..bf0a264b49ed 100644
--- a/core/java/android/os/BatteryStats.java
+++ b/core/java/android/os/BatteryStats.java
@@ -78,17 +78,17 @@ public abstract class BatteryStats implements Parcelable {
      * A constant indicating a sensor timer.
      */
     public static final int SENSOR = 3;
-    
+
     /**
      * A constant indicating a a wifi running timer
      */
     public static final int WIFI_RUNNING = 4;
-    
+
     /**
      * A constant indicating a full wifi lock timer
      */
     public static final int FULL_WIFI_LOCK = 5;
-    
+
     /**
      * A constant indicating a wifi scan
      */
@@ -217,8 +217,10 @@ public abstract class BatteryStats implements Parcelable {
      *   - Package wakeup alarms are now on screen-off timebase
      * New in version 26:
      *   - Resource power manager (rpm) states [but screenOffRpm is disabled from working properly]
+     * New in version 27:
+     *   - Always On Display (screen doze mode) time and power
      */
-    static final String CHECKIN_VERSION = "26";
+    static final String CHECKIN_VERSION = "27";
 
     /**
      * Old version, we hit 9 and ran out of room, need to remove.
@@ -1381,12 +1383,12 @@ public abstract class BatteryStats implements Parcelable {
         public static final int STATE_PHONE_SCANNING_FLAG = 1<<21;
         public static final int STATE_SCREEN_ON_FLAG = 1<<20;       // consider moving to states2
         public static final int STATE_BATTERY_PLUGGED_FLAG = 1<<19; // consider moving to states2
-        // empty slot
+        public static final int STATE_SCREEN_DOZE_FLAG = 1 << 18;
         // empty slot
         public static final int STATE_WIFI_MULTICAST_ON_FLAG = 1<<16;
 
         public static final int MOST_INTERESTING_STATES =
-            STATE_BATTERY_PLUGGED_FLAG | STATE_SCREEN_ON_FLAG;
+                STATE_BATTERY_PLUGGED_FLAG | STATE_SCREEN_ON_FLAG | STATE_SCREEN_DOZE_FLAG;
 
         public static final int SETTLE_TO_ZERO_STATES = 0xffff0000 & ~MOST_INTERESTING_STATES;
 
@@ -1414,8 +1416,8 @@ public abstract class BatteryStats implements Parcelable {
         public static final int STATE2_BLUETOOTH_SCAN_FLAG = 1 << 20;
 
         public static final int MOST_INTERESTING_STATES2 =
-            STATE2_POWER_SAVE_FLAG | STATE2_WIFI_ON_FLAG | STATE2_DEVICE_IDLE_MASK
-            | STATE2_CHARGING_FLAG | STATE2_PHONE_IN_CALL_FLAG | STATE2_BLUETOOTH_ON_FLAG;
+                STATE2_POWER_SAVE_FLAG | STATE2_WIFI_ON_FLAG | STATE2_DEVICE_IDLE_MASK
+                | STATE2_CHARGING_FLAG | STATE2_PHONE_IN_CALL_FLAG | STATE2_BLUETOOTH_ON_FLAG;
 
         public static final int SETTLE_TO_ZERO_STATES2 = 0xffff0000 & ~MOST_INTERESTING_STATES2;
 
@@ -1863,6 +1865,21 @@ public abstract class BatteryStats implements Parcelable {
      */
     public abstract int getScreenOnCount(int which);
 
+    /**
+     * Returns the time in microseconds that the screen has been dozing while the device was
+     * running on battery.
+     *
+     * {@hide}
+     */
+    public abstract long getScreenDozeTime(long elapsedRealtimeUs, int which);
+
+    /**
+     * Returns the number of times the screen was turned dozing.
+     *
+     * {@hide}
+     */
+    public abstract int getScreenDozeCount(int which);
+
     public abstract long getInteractiveTime(long elapsedRealtimeUs, int which);
 
     public static final int SCREEN_BRIGHTNESS_DARK = 0;
@@ -2116,8 +2133,7 @@ public abstract class BatteryStats implements Parcelable {
         "group", "compl", "dorm", "uninit"
     };
 
-    public static final BitDescription[] HISTORY_STATE_DESCRIPTIONS
-            = new BitDescription[] {
+    public static final BitDescription[] HISTORY_STATE_DESCRIPTIONS = new BitDescription[] {
         new BitDescription(HistoryItem.STATE_CPU_RUNNING_FLAG, "running", "r"),
         new BitDescription(HistoryItem.STATE_WAKE_LOCK_FLAG, "wake_lock", "w"),
         new BitDescription(HistoryItem.STATE_SENSOR_ON_FLAG, "sensor", "s"),
@@ -2131,6 +2147,7 @@ public abstract class BatteryStats implements Parcelable {
         new BitDescription(HistoryItem.STATE_AUDIO_ON_FLAG, "audio", "a"),
         new BitDescription(HistoryItem.STATE_SCREEN_ON_FLAG, "screen", "S"),
         new BitDescription(HistoryItem.STATE_BATTERY_PLUGGED_FLAG, "plugged", "BP"),
+        new BitDescription(HistoryItem.STATE_SCREEN_DOZE_FLAG, "screen_doze", "Sd"),
         new BitDescription(HistoryItem.STATE_DATA_CONNECTION_MASK,
                 HistoryItem.STATE_DATA_CONNECTION_SHIFT, "data_conn", "Pcn",
                 DATA_CONNECTION_NAMES, DATA_CONNECTION_NAMES),
@@ -2467,6 +2484,18 @@ public abstract class BatteryStats implements Parcelable {
     public abstract int getDischargeAmountScreenOffSinceCharge();
 
     /**
+     * Get the amount the battery has discharged while the screen was doze,
+     * since the last time power was unplugged.
+     */
+    public abstract int getDischargeAmountScreenDoze();
+
+    /**
+     * Get the amount the battery has discharged while the screen was doze,
+     * since the last time the device was charged.
+     */
+    public abstract int getDischargeAmountScreenDozeSinceCharge();
+
+    /**
      * Returns the total, last, or current battery uptime in microseconds.
      *
      * @param curTime the elapsed realtime in microseconds.
@@ -2483,7 +2512,7 @@ public abstract class BatteryStats implements Parcelable {
     public abstract long computeBatteryRealtime(long curTime, int which);
 
     /**
-     * Returns the total, last, or current battery screen off uptime in microseconds.
+     * Returns the total, last, or current battery screen off/doze uptime in microseconds.
      *
      * @param curTime the elapsed realtime in microseconds.
      * @param which one of STATS_SINCE_CHARGED, STATS_SINCE_UNPLUGGED, or STATS_CURRENT.
@@ -2491,7 +2520,7 @@ public abstract class BatteryStats implements Parcelable {
     public abstract long computeBatteryScreenOffUptime(long curTime, int which);
 
     /**
-     * Returns the total, last, or current battery screen off realtime in microseconds.
+     * Returns the total, last, or current battery screen off/doze realtime in microseconds.
      *
      * @param curTime the current elapsed realtime in microseconds.
      * @param which one of STATS_SINCE_CHARGED, STATS_SINCE_UNPLUGGED, or STATS_CURRENT.
@@ -2590,18 +2619,24 @@ public abstract class BatteryStats implements Parcelable {
     };
 
     /**
-     * Return the counter keeping track of the amount of battery discharge while the screen was off,
-     * measured in micro-Ampere-hours. This will be non-zero only if the device's battery has
+     * Return the amount of battery discharge while the screen was off, measured in
+     * micro-Ampere-hours. This will be non-zero only if the device's battery has
      * a coulomb counter.
      */
-    public abstract LongCounter getDischargeScreenOffCoulombCounter();
+    public abstract long getMahDischargeScreenOff(int which);
 
     /**
-     * Return the counter keeping track of the amount of battery discharge measured in
+     * Return the amount of battery discharge while the screen was in doze mode, measured in
      * micro-Ampere-hours. This will be non-zero only if the device's battery has
      * a coulomb counter.
      */
-    public abstract LongCounter getDischargeCoulombCounter();
+    public abstract long getMahDischargeScreenDoze(int which);
+
+    /**
+     * Return the amount of battery discharge  measured in micro-Ampere-hours. This will be
+     * non-zero only if the device's battery has a coulomb counter.
+     */
+    public abstract long getMahDischarge(int which);
 
     /**
      * Returns the estimated real battery capacity, which may be less than the capacity
@@ -3112,6 +3147,7 @@ public abstract class BatteryStats implements Parcelable {
         final long totalRealtime = computeRealtime(rawRealtime, which);
         final long totalUptime = computeUptime(rawUptime, which);
         final long screenOnTime = getScreenOnTime(rawRealtime, which);
+        final long screenDozeTime = getScreenDozeTime(rawRealtime, which);
         final long interactiveTime = getInteractiveTime(rawRealtime, which);
         final long powerSaveModeEnabledTime = getPowerSaveModeEnabledTime(rawRealtime, which);
         final long deviceIdleModeLightTime = getDeviceIdleModeTime(DEVICE_IDLE_MODE_LIGHT,
@@ -3124,9 +3160,9 @@ public abstract class BatteryStats implements Parcelable {
                 rawRealtime, which);
         final int connChanges = getNumConnectivityChange(which);
         final long phoneOnTime = getPhoneOnTime(rawRealtime, which);
-        final long dischargeCount = getDischargeCoulombCounter().getCountLocked(which);
-        final long dischargeScreenOffCount = getDischargeScreenOffCoulombCounter()
-                .getCountLocked(which);
+        final long dischargeCount = getMahDischarge(which);
+        final long dischargeScreenOffCount = getMahDischargeScreenOff(which);
+        final long dischargeScreenDozeCount = getMahDischargeScreenDoze(which);
 
         final StringBuilder sb = new StringBuilder(128);
 
@@ -3143,7 +3179,8 @@ public abstract class BatteryStats implements Parcelable {
                 getStartClockTime(),
                 whichBatteryScreenOffRealtime / 1000, whichBatteryScreenOffUptime / 1000,
                 getEstimatedBatteryCapacity(),
-                getMinLearnedBatteryCapacity(), getMaxLearnedBatteryCapacity());
+                getMinLearnedBatteryCapacity(), getMaxLearnedBatteryCapacity(),
+                screenDozeTime / 1000);
 
 
         // Calculate wakelock times across all uids.
@@ -3295,13 +3332,15 @@ public abstract class BatteryStats implements Parcelable {
                     getDischargeStartLevel()-getDischargeCurrentLevel(),
                     getDischargeStartLevel()-getDischargeCurrentLevel(),
                     getDischargeAmountScreenOn(), getDischargeAmountScreenOff(),
-                    dischargeCount / 1000, dischargeScreenOffCount / 1000);
+                    dischargeCount / 1000, dischargeScreenOffCount / 1000,
+                    getDischargeAmountScreenDoze(), dischargeScreenDozeCount / 1000);
         } else {
             dumpLine(pw, 0 /* uid */, category, BATTERY_DISCHARGE_DATA,
                     getLowDischargeAmountSinceCharge(), getHighDischargeAmountSinceCharge(),
                     getDischargeAmountScreenOnSinceCharge(),
                     getDischargeAmountScreenOffSinceCharge(),
-                    dischargeCount / 1000, dischargeScreenOffCount / 1000);
+                    dischargeCount / 1000, dischargeScreenOffCount / 1000,
+                    getDischargeAmountScreenDozeSinceCharge(), dischargeScreenDozeCount / 1000);
         }
 
         if (reqUid < 0) {
@@ -3831,6 +3870,7 @@ public abstract class BatteryStats implements Parcelable {
                 which);
         final long batteryTimeRemaining = computeBatteryTimeRemaining(rawRealtime);
         final long chargeTimeRemaining = computeChargeTimeRemaining(rawRealtime);
+        final long screenDozeTime = getScreenDozeTime(rawRealtime, which);
 
         final StringBuilder sb = new StringBuilder(128);
 
@@ -3868,25 +3908,35 @@ public abstract class BatteryStats implements Parcelable {
 
         sb.setLength(0);
         sb.append(prefix);
-                sb.append("  Time on battery: ");
-                formatTimeMs(sb, whichBatteryRealtime / 1000); sb.append("(");
-                sb.append(formatRatioLocked(whichBatteryRealtime, totalRealtime));
-                sb.append(") realtime, ");
-                formatTimeMs(sb, whichBatteryUptime / 1000);
-                sb.append("("); sb.append(formatRatioLocked(whichBatteryUptime, totalRealtime));
-                sb.append(") uptime");
+        sb.append("  Time on battery: ");
+        formatTimeMs(sb, whichBatteryRealtime / 1000); sb.append("(");
+        sb.append(formatRatioLocked(whichBatteryRealtime, totalRealtime));
+        sb.append(") realtime, ");
+        formatTimeMs(sb, whichBatteryUptime / 1000);
+        sb.append("("); sb.append(formatRatioLocked(whichBatteryUptime, whichBatteryRealtime));
+        sb.append(") uptime");
         pw.println(sb.toString());
+
         sb.setLength(0);
         sb.append(prefix);
-                sb.append("  Time on battery screen off: ");
-                formatTimeMs(sb, whichBatteryScreenOffRealtime / 1000); sb.append("(");
-                sb.append(formatRatioLocked(whichBatteryScreenOffRealtime, totalRealtime));
-                sb.append(") realtime, ");
-                formatTimeMs(sb, whichBatteryScreenOffUptime / 1000);
-                sb.append("(");
-                sb.append(formatRatioLocked(whichBatteryScreenOffUptime, totalRealtime));
-                sb.append(") uptime");
+        sb.append("  Time on battery screen off: ");
+        formatTimeMs(sb, whichBatteryScreenOffRealtime / 1000); sb.append("(");
+        sb.append(formatRatioLocked(whichBatteryScreenOffRealtime, whichBatteryRealtime));
+        sb.append(") realtime, ");
+        formatTimeMs(sb, whichBatteryScreenOffUptime / 1000);
+        sb.append("(");
+        sb.append(formatRatioLocked(whichBatteryScreenOffUptime, whichBatteryRealtime));
+        sb.append(") uptime");
+        pw.println(sb.toString());
+
+        sb.setLength(0);
+        sb.append(prefix);
+        sb.append("  Time on battery screen doze: ");
+        formatTimeMs(sb, screenDozeTime / 1000); sb.append("(");
+        sb.append(formatRatioLocked(screenDozeTime, whichBatteryRealtime));
+        sb.append(")");
         pw.println(sb.toString());
+
         sb.setLength(0);
         sb.append(prefix);
                 sb.append("  Total run time: ");
@@ -3910,8 +3960,7 @@ public abstract class BatteryStats implements Parcelable {
             pw.println(sb.toString());
         }
 
-        final LongCounter dischargeCounter = getDischargeCoulombCounter();
-        final long dischargeCount = dischargeCounter.getCountLocked(which);
+        final long dischargeCount = getMahDischarge(which);
         if (dischargeCount >= 0) {
             sb.setLength(0);
             sb.append(prefix);
@@ -3921,8 +3970,7 @@ public abstract class BatteryStats implements Parcelable {
             pw.println(sb.toString());
         }
 
-        final LongCounter dischargeScreenOffCounter = getDischargeScreenOffCoulombCounter();
-        final long dischargeScreenOffCount = dischargeScreenOffCounter.getCountLocked(which);
+        final long dischargeScreenOffCount = getMahDischargeScreenOff(which);
         if (dischargeScreenOffCount >= 0) {
             sb.setLength(0);
             sb.append(prefix);
@@ -3932,7 +3980,18 @@ public abstract class BatteryStats implements Parcelable {
             pw.println(sb.toString());
         }
 
-        final long dischargeScreenOnCount = dischargeCount - dischargeScreenOffCount;
+        final long dischargeScreenDozeCount = getMahDischargeScreenDoze(which);
+        if (dischargeScreenDozeCount >= 0) {
+            sb.setLength(0);
+            sb.append(prefix);
+            sb.append("  Screen doze discharge: ");
+            sb.append(BatteryStatsHelper.makemAh(dischargeScreenDozeCount / 1000.0));
+            sb.append(" mAh");
+            pw.println(sb.toString());
+        }
+
+        final long dischargeScreenOnCount =
+                dischargeCount - dischargeScreenOffCount - dischargeScreenDozeCount;
         if (dischargeScreenOnCount >= 0) {
             sb.setLength(0);
             sb.append(prefix);
@@ -4340,20 +4399,24 @@ public abstract class BatteryStats implements Parcelable {
                         pw.println(getDischargeCurrentLevel());
             }
             pw.print(prefix); pw.print("    Amount discharged while screen on: ");
-                    pw.println(getDischargeAmountScreenOn());
+            pw.println(getDischargeAmountScreenOn());
             pw.print(prefix); pw.print("    Amount discharged while screen off: ");
-                    pw.println(getDischargeAmountScreenOff());
+            pw.println(getDischargeAmountScreenOff());
+            pw.print(prefix); pw.print("    Amount discharged while screen doze: ");
+            pw.println(getDischargeAmountScreenDoze());
             pw.println(" ");
         } else {
             pw.print(prefix); pw.println("  Device battery use since last full charge");
             pw.print(prefix); pw.print("    Amount discharged (lower bound): ");
-                    pw.println(getLowDischargeAmountSinceCharge());
+            pw.println(getLowDischargeAmountSinceCharge());
             pw.print(prefix); pw.print("    Amount discharged (upper bound): ");
-                    pw.println(getHighDischargeAmountSinceCharge());
+            pw.println(getHighDischargeAmountSinceCharge());
             pw.print(prefix); pw.print("    Amount discharged while screen on: ");
-                    pw.println(getDischargeAmountScreenOnSinceCharge());
+            pw.println(getDischargeAmountScreenOnSinceCharge());
             pw.print(prefix); pw.print("    Amount discharged while screen off: ");
-                    pw.println(getDischargeAmountScreenOffSinceCharge());
+            pw.println(getDischargeAmountScreenOffSinceCharge());
+            pw.print(prefix); pw.print("    Amount discharged while screen doze: ");
+            pw.println(getDischargeAmountScreenDozeSinceCharge());
             pw.println();
         }
 
@@ -5426,7 +5489,7 @@ public abstract class BatteryStats implements Parcelable {
             }
         }
     }
-    
+
     public void prepareForDumpLocked() {
     }
 
@@ -6248,7 +6311,7 @@ public abstract class BatteryStats implements Parcelable {
             pw.println();
         }
     }
-    
+
     @SuppressWarnings("unused")
     public void dumpCheckinLocked(Context context, PrintWriter pw,
             List<ApplicationInfo> apps, int flags, long histStart) {
diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index 4300b1880721..f26c0cdfe958 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -119,7 +119,7 @@ public class BatteryStatsImpl extends BatteryStats {
     private static final int MAGIC = 0xBA757475; // 'BATSTATS'
 
     // Current on-disk Parcel version
-    private static final int VERSION = 166 + (USE_OLD_HISTORY ? 1000 : 0);
+    private static final int VERSION = 167 + (USE_OLD_HISTORY ? 1000 : 0);
 
     // Maximum number of items we will record in the history.
     private static final int MAX_HISTORY_ITEMS;
@@ -324,8 +324,8 @@ public class BatteryStatsImpl extends BatteryStats {
     protected final TimeBase mOnBatteryTimeBase = new TimeBase();
 
     // These are the objects that will want to do something when the device
-    // is unplugged from power *and* the screen is off.
-    final TimeBase mOnBatteryScreenOffTimeBase = new TimeBase();
+    // is unplugged from power *and* the screen is off or doze.
+    protected final TimeBase mOnBatteryScreenOffTimeBase = new TimeBase();
 
     // Set to true when we want to distribute CPU across wakelocks for the next
     // CPU update, even if we aren't currently running wake locks.
@@ -419,8 +419,12 @@ public class BatteryStatsImpl extends BatteryStats {
     public boolean mRecordAllHistory;
     boolean mNoAutoReset;
 
-    int mScreenState = Display.STATE_UNKNOWN;
-    StopwatchTimer mScreenOnTimer;
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
+    protected int mScreenState = Display.STATE_UNKNOWN;
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
+    protected StopwatchTimer mScreenOnTimer;
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
+    protected StopwatchTimer mScreenDozeTimer;
 
     int mScreenBrightnessBin = -1;
     final StopwatchTimer[] mScreenBrightnessTimer = new StopwatchTimer[NUM_SCREEN_BRIGHTNESS_BINS];
@@ -565,12 +569,16 @@ public class BatteryStatsImpl extends BatteryStats {
     int mHighDischargeAmountSinceCharge;
     int mDischargeScreenOnUnplugLevel;
     int mDischargeScreenOffUnplugLevel;
+    int mDischargeScreenDozeUnplugLevel;
     int mDischargeAmountScreenOn;
     int mDischargeAmountScreenOnSinceCharge;
     int mDischargeAmountScreenOff;
     int mDischargeAmountScreenOffSinceCharge;
+    int mDischargeAmountScreenDoze;
+    int mDischargeAmountScreenDozeSinceCharge;
 
     private LongSamplingCounter mDischargeScreenOffCounter;
+    private LongSamplingCounter mDischargeScreenDozeCounter;
     private LongSamplingCounter mDischargeCounter;
 
     static final int MAX_LEVEL_STEPS = 200;
@@ -654,13 +662,18 @@ public class BatteryStatsImpl extends BatteryStats {
     }
 
     @Override
-    public LongCounter getDischargeScreenOffCoulombCounter() {
-        return mDischargeScreenOffCounter;
+    public long getMahDischarge(int which) {
+        return mDischargeCounter.getCountLocked(which);
+    }
+
+    @Override
+    public long getMahDischargeScreenOff(int which) {
+        return mDischargeScreenOffCounter.getCountLocked(which);
     }
 
     @Override
-    public LongCounter getDischargeCoulombCounter() {
-        return mDischargeCounter;
+    public long getMahDischargeScreenDoze(int which) {
+        return mDischargeScreenDozeCounter.getCountLocked(which);
     }
 
     @Override
@@ -3553,8 +3566,9 @@ public class BatteryStatsImpl extends BatteryStats {
         mActiveHistoryStates2 = 0xffffffff;
     }
 
-    public void updateTimeBasesLocked(boolean unplugged, boolean screenOff, long uptime,
+    public void updateTimeBasesLocked(boolean unplugged, int screenState, long uptime,
             long realtime) {
+        final boolean screenOff = isScreenOff(screenState) || isScreenDoze(screenState);
         final boolean updateOnBatteryTimeBase = unplugged != mOnBatteryTimeBase.isRunning();
         final boolean updateOnBatteryScreenOffTimeBase =
                 (unplugged && screenOff) != mOnBatteryScreenOffTimeBase.isRunning();
@@ -3571,20 +3585,21 @@ public class BatteryStatsImpl extends BatteryStats {
                 updateRpmStatsLocked(); // if either OnBattery or OnBatteryScreenOff timebase changes.
             }
             if (DEBUG_ENERGY_CPU) {
-                Slog.d(TAG, "Updating cpu time because screen is now " + (screenOff ? "off" : "on")
+                Slog.d(TAG, "Updating cpu time because screen is now "
+                        + Display.stateToString(screenState)
                         + " and battery is " + (unplugged ? "on" : "off"));
             }
             updateCpuTimeLocked(true /* updateCpuFreqData */);
-
             mOnBatteryTimeBase.setRunning(unplugged, uptime, realtime);
-            mOnBatteryScreenOffTimeBase.setRunning(unplugged && screenOff, uptime, realtime);
-            for (int i = mUidStats.size() - 1; i >= 0; --i) {
-                final Uid u = mUidStats.valueAt(i);
-                if (updateOnBatteryTimeBase) {
-                    u.updateOnBatteryBgTimeBase(uptime, realtime);
+            if (updateOnBatteryTimeBase) {
+                for (int i = mUidStats.size() - 1; i >= 0; --i) {
+                    mUidStats.valueAt(i).updateOnBatteryBgTimeBase(uptime, realtime);
                 }
-                if (updateOnBatteryScreenOffTimeBase) {
-                    u.updateOnBatteryScreenOffBgTimeBase(uptime, realtime);
+            }
+            if (updateOnBatteryScreenOffTimeBase) {
+                mOnBatteryScreenOffTimeBase.setRunning(unplugged && screenOff, uptime, realtime);
+                for (int i = mUidStats.size() - 1; i >= 0; --i) {
+                    mUidStats.valueAt(i).updateOnBatteryScreenOffBgTimeBase(uptime, realtime);
                 }
             }
         }
@@ -3844,8 +3859,10 @@ public class BatteryStatsImpl extends BatteryStats {
     }
 
     public void setPretendScreenOff(boolean pretendScreenOff) {
-        mPretendScreenOff = pretendScreenOff;
-        noteScreenStateLocked(pretendScreenOff ? Display.STATE_OFF : Display.STATE_ON);
+        if (mPretendScreenOff != pretendScreenOff) {
+            mPretendScreenOff = pretendScreenOff;
+            noteScreenStateLocked(pretendScreenOff ? Display.STATE_OFF : Display.STATE_ON);
+        }
     }
 
     private String mInitialAcquireWakeName;
@@ -4175,54 +4192,58 @@ public class BatteryStatsImpl extends BatteryStats {
                 }
             }
 
-            if (state == Display.STATE_ON) {
-                // Screen turning on.
-                final long elapsedRealtime = mClocks.elapsedRealtime();
-                final long uptime = mClocks.uptimeMillis();
+            final long elapsedRealtime = mClocks.elapsedRealtime();
+            final long uptime = mClocks.uptimeMillis();
+
+            boolean updateHistory = false;
+            if (isScreenDoze(state)) {
+                mHistoryCur.states |= HistoryItem.STATE_SCREEN_DOZE_FLAG;
+                mScreenDozeTimer.startRunningLocked(elapsedRealtime);
+                updateHistory = true;
+            } else if (isScreenDoze(oldState)) {
+                mHistoryCur.states &= ~HistoryItem.STATE_SCREEN_DOZE_FLAG;
+                mScreenDozeTimer.stopRunningLocked(elapsedRealtime);
+                updateHistory = true;
+            }
+            if (isScreenOn(state)) {
                 mHistoryCur.states |= HistoryItem.STATE_SCREEN_ON_FLAG;
                 if (DEBUG_HISTORY) Slog.v(TAG, "Screen on to: "
                         + Integer.toHexString(mHistoryCur.states));
-                addHistoryRecordLocked(elapsedRealtime, uptime);
                 mScreenOnTimer.startRunningLocked(elapsedRealtime);
                 if (mScreenBrightnessBin >= 0) {
                     mScreenBrightnessTimer[mScreenBrightnessBin].startRunningLocked(elapsedRealtime);
                 }
-
-                updateTimeBasesLocked(mOnBatteryTimeBase.isRunning(), false,
-                        mClocks.uptimeMillis() * 1000, elapsedRealtime * 1000);
-
-                // Fake a wake lock, so we consider the device waked as long
-                // as the screen is on.
-                noteStartWakeLocked(-1, -1, "screen", null, WAKE_TYPE_PARTIAL, false,
-                        elapsedRealtime, uptime);
-
-                // Update discharge amounts.
-                if (mOnBatteryInternal) {
-                    updateDischargeScreenLevelsLocked(false, true);
-                }
-            } else if (oldState == Display.STATE_ON) {
-                // Screen turning off or dozing.
-                final long elapsedRealtime = mClocks.elapsedRealtime();
-                final long uptime = mClocks.uptimeMillis();
+                updateHistory = true;
+            } else if (isScreenOn(oldState)) {
                 mHistoryCur.states &= ~HistoryItem.STATE_SCREEN_ON_FLAG;
                 if (DEBUG_HISTORY) Slog.v(TAG, "Screen off to: "
                         + Integer.toHexString(mHistoryCur.states));
-                addHistoryRecordLocked(elapsedRealtime, uptime);
                 mScreenOnTimer.stopRunningLocked(elapsedRealtime);
                 if (mScreenBrightnessBin >= 0) {
                     mScreenBrightnessTimer[mScreenBrightnessBin].stopRunningLocked(elapsedRealtime);
                 }
-
+                updateHistory = true;
+            }
+            if (updateHistory) {
+                if (DEBUG_HISTORY) Slog.v(TAG, "Screen state to: "
+                        + Display.stateToString(state));
+                addHistoryRecordLocked(elapsedRealtime, uptime);
+            }
+            if (isScreenOn(state)) {
+                updateTimeBasesLocked(mOnBatteryTimeBase.isRunning(), state,
+                        mClocks.uptimeMillis() * 1000, elapsedRealtime * 1000);
+                // Fake a wake lock, so we consider the device waked as long as the screen is on.
+                noteStartWakeLocked(-1, -1, "screen", null, WAKE_TYPE_PARTIAL, false,
+                        elapsedRealtime, uptime);
+            } else if (isScreenOn(oldState)) {
                 noteStopWakeLocked(-1, -1, "screen", "screen", WAKE_TYPE_PARTIAL,
                         elapsedRealtime, uptime);
-
-                updateTimeBasesLocked(mOnBatteryTimeBase.isRunning(), true,
+                updateTimeBasesLocked(mOnBatteryTimeBase.isRunning(), state,
                         mClocks.uptimeMillis() * 1000, elapsedRealtime * 1000);
-
-                // Update discharge amounts.
-                if (mOnBatteryInternal) {
-                    updateDischargeScreenLevelsLocked(true, false);
-                }
+            }
+            // Update discharge amounts.
+            if (mOnBatteryInternal) {
+                updateDischargeScreenLevelsLocked(oldState, state);
             }
         }
     }
@@ -5371,6 +5392,14 @@ public class BatteryStatsImpl extends BatteryStats {
         return mScreenOnTimer.getCountLocked(which);
     }
 
+    @Override public long getScreenDozeTime(long elapsedRealtimeUs, int which) {
+        return mScreenDozeTimer.getTotalTimeLocked(elapsedRealtimeUs, which);
+    }
+
+    @Override public int getScreenDozeCount(int which) {
+        return mScreenDozeTimer.getCountLocked(which);
+    }
+
     @Override public long getScreenBrightnessTime(int brightnessBin,
             long elapsedRealtimeUs, int which) {
         return mScreenBrightnessTimer[brightnessBin].getTotalTimeLocked(
@@ -8809,6 +8838,7 @@ public class BatteryStatsImpl extends BatteryStats {
         mHandler = new MyHandler(handler.getLooper());
         mStartCount++;
         mScreenOnTimer = new StopwatchTimer(mClocks, null, -1, null, mOnBatteryTimeBase);
+        mScreenDozeTimer = new StopwatchTimer(mClocks, null, -1, null, mOnBatteryTimeBase);
         for (int i=0; i<NUM_SCREEN_BRIGHTNESS_BINS; i++) {
             mScreenBrightnessTimer[i] = new StopwatchTimer(mClocks, null, -100-i, null,
                     mOnBatteryTimeBase);
@@ -8867,6 +8897,7 @@ public class BatteryStatsImpl extends BatteryStats {
         mCameraOnTimer = new StopwatchTimer(mClocks, null, -13, null, mOnBatteryTimeBase);
         mBluetoothScanTimer = new StopwatchTimer(mClocks, null, -14, null, mOnBatteryTimeBase);
         mDischargeScreenOffCounter = new LongSamplingCounter(mOnBatteryScreenOffTimeBase);
+        mDischargeScreenDozeCounter = new LongSamplingCounter(mOnBatteryTimeBase);
         mDischargeCounter = new LongSamplingCounter(mOnBatteryTimeBase);
         mOnBattery = mOnBatteryInternal = false;
         long uptime = mClocks.uptimeMillis() * 1000;
@@ -9410,8 +9441,16 @@ public class BatteryStatsImpl extends BatteryStats {
         return mCharging;
     }
 
-    public boolean isScreenOn() {
-        return mScreenState == Display.STATE_ON;
+    public boolean isScreenOn(int state) {
+        return state == Display.STATE_ON;
+    }
+
+    public boolean isScreenOff(int state) {
+        return state == Display.STATE_OFF;
+    }
+
+    public boolean isScreenDoze(int state) {
+        return state == Display.STATE_DOZE || state == Display.STATE_DOZE_SUSPEND;
     }
 
     void initTimes(long uptime, long realtime) {
@@ -9431,9 +9470,12 @@ public class BatteryStatsImpl extends BatteryStats {
         mDischargeAmountScreenOnSinceCharge = 0;
         mDischargeAmountScreenOff = 0;
         mDischargeAmountScreenOffSinceCharge = 0;
+        mDischargeAmountScreenDoze = 0;
+        mDischargeAmountScreenDozeSinceCharge = 0;
         mDischargeStepTracker.init();
         mChargeStepTracker.init();
         mDischargeScreenOffCounter.reset(false);
+        mDischargeScreenDozeCounter.reset(false);
         mDischargeCounter.reset(false);
     }
 
@@ -9451,15 +9493,22 @@ public class BatteryStatsImpl extends BatteryStats {
         mOnBatteryTimeBase.reset(uptime, realtime);
         mOnBatteryScreenOffTimeBase.reset(uptime, realtime);
         if ((mHistoryCur.states&HistoryItem.STATE_BATTERY_PLUGGED_FLAG) == 0) {
-            if (mScreenState == Display.STATE_ON) {
+            if (isScreenOn(mScreenState)) {
                 mDischargeScreenOnUnplugLevel = mHistoryCur.batteryLevel;
+                mDischargeScreenDozeUnplugLevel = 0;
+                mDischargeScreenOffUnplugLevel = 0;
+            } else if (isScreenDoze(mScreenState)) {
+                mDischargeScreenOnUnplugLevel = 0;
+                mDischargeScreenDozeUnplugLevel = mHistoryCur.batteryLevel;
                 mDischargeScreenOffUnplugLevel = 0;
             } else {
                 mDischargeScreenOnUnplugLevel = 0;
+                mDischargeScreenDozeUnplugLevel = 0;
                 mDischargeScreenOffUnplugLevel = mHistoryCur.batteryLevel;
             }
             mDischargeAmountScreenOn = 0;
             mDischargeAmountScreenOff = 0;
+            mDischargeAmountScreenDoze = 0;
         }
         initActiveHistoryEventsLocked(mSecRealtime, mSecUptime);
     }
@@ -9470,6 +9519,7 @@ public class BatteryStatsImpl extends BatteryStats {
         mStartCount = 0;
         initTimes(uptimeMillis * 1000, elapsedRealtimeMillis * 1000);
         mScreenOnTimer.reset(false);
+        mScreenDozeTimer.reset(false);
         for (int i=0; i<NUM_SCREEN_BRIGHTNESS_BINS; i++) {
             mScreenBrightnessTimer[i].reset(false);
         }
@@ -9606,33 +9656,52 @@ public class BatteryStatsImpl extends BatteryStats {
         }
     }
 
-    void updateDischargeScreenLevelsLocked(boolean oldScreenOn, boolean newScreenOn) {
-        if (oldScreenOn) {
+    void updateDischargeScreenLevelsLocked(int oldState, int newState) {
+        updateOldDischargeScreenLevelLocked(oldState);
+        updateNewDischargeScreenLevelLocked(newState);
+    }
+
+    private void updateOldDischargeScreenLevelLocked(int state) {
+        if (isScreenOn(state)) {
             int diff = mDischargeScreenOnUnplugLevel - mDischargeCurrentLevel;
             if (diff > 0) {
                 mDischargeAmountScreenOn += diff;
                 mDischargeAmountScreenOnSinceCharge += diff;
             }
-        } else {
+        } else if (isScreenDoze(state)) {
+            int diff = mDischargeScreenDozeUnplugLevel - mDischargeCurrentLevel;
+            if (diff > 0) {
+                mDischargeAmountScreenDoze += diff;
+                mDischargeAmountScreenDozeSinceCharge += diff;
+            }
+        } else if (isScreenOff(state)){
             int diff = mDischargeScreenOffUnplugLevel - mDischargeCurrentLevel;
             if (diff > 0) {
                 mDischargeAmountScreenOff += diff;
                 mDischargeAmountScreenOffSinceCharge += diff;
             }
         }
-        if (newScreenOn) {
+    }
+
+    private void updateNewDischargeScreenLevelLocked(int state) {
+        if (isScreenOn(state)) {
             mDischargeScreenOnUnplugLevel = mDischargeCurrentLevel;
             mDischargeScreenOffUnplugLevel = 0;
-        } else {
+            mDischargeScreenDozeUnplugLevel = 0;
+        } else if (isScreenDoze(state)){
             mDischargeScreenOnUnplugLevel = 0;
+            mDischargeScreenDozeUnplugLevel = mDischargeCurrentLevel;
+            mDischargeScreenOffUnplugLevel = 0;
+        } else if (isScreenOff(state)) {
+            mDischargeScreenOnUnplugLevel = 0;
+            mDischargeScreenDozeUnplugLevel = 0;
             mDischargeScreenOffUnplugLevel = mDischargeCurrentLevel;
         }
     }
 
     public void pullPendingStateUpdatesLocked() {
         if (mOnBatteryInternal) {
-            final boolean screenOn = mScreenState == Display.STATE_ON;
-            updateDischargeScreenLevelsLocked(screenOn, screenOn);
+            updateDischargeScreenLevelsLocked(mScreenState, mScreenState);
         }
     }
 
@@ -10687,8 +10756,8 @@ public class BatteryStatsImpl extends BatteryStats {
         return false;
     }
 
-    void setOnBatteryLocked(final long mSecRealtime, final long mSecUptime, final boolean onBattery,
-            final int oldStatus, final int level, final int chargeUAh) {
+    protected void setOnBatteryLocked(final long mSecRealtime, final long mSecUptime,
+            final boolean onBattery, final int oldStatus, final int level, final int chargeUAh) {
         boolean doWrite = false;
         Message m = mHandler.obtainMessage(MSG_REPORT_POWER_CHANGE);
         m.arg1 = onBattery ? 1 : 0;
@@ -10696,7 +10765,7 @@ public class BatteryStatsImpl extends BatteryStats {
 
         final long uptime = mSecUptime * 1000;
         final long realtime = mSecRealtime * 1000;
-        final boolean screenOn = mScreenState == Display.STATE_ON;
+        final int screenState = mScreenState;
         if (onBattery) {
             // We will reset our status if we are unplugging after the
             // battery was last full, or the level is at 100, or
@@ -10772,16 +10841,23 @@ public class BatteryStatsImpl extends BatteryStats {
             }
             addHistoryRecordLocked(mSecRealtime, mSecUptime);
             mDischargeCurrentLevel = mDischargeUnplugLevel = level;
-            if (screenOn) {
+            if (isScreenOn(screenState)) {
                 mDischargeScreenOnUnplugLevel = level;
+                mDischargeScreenDozeUnplugLevel = 0;
+                mDischargeScreenOffUnplugLevel = 0;
+            } else if (isScreenDoze(screenState)) {
+                mDischargeScreenOnUnplugLevel = 0;
+                mDischargeScreenDozeUnplugLevel = level;
                 mDischargeScreenOffUnplugLevel = 0;
             } else {
                 mDischargeScreenOnUnplugLevel = 0;
+                mDischargeScreenDozeUnplugLevel = 0;
                 mDischargeScreenOffUnplugLevel = level;
             }
             mDischargeAmountScreenOn = 0;
+            mDischargeAmountScreenDoze = 0;
             mDischargeAmountScreenOff = 0;
-            updateTimeBasesLocked(true, !screenOn, uptime, realtime);
+            updateTimeBasesLocked(true, screenState, uptime, realtime);
         } else {
             mLastChargingStateLevel = level;
             mOnBattery = mOnBatteryInternal = false;
@@ -10796,8 +10872,8 @@ public class BatteryStatsImpl extends BatteryStats {
                 mLowDischargeAmountSinceCharge += mDischargeUnplugLevel-level-1;
                 mHighDischargeAmountSinceCharge += mDischargeUnplugLevel-level;
             }
-            updateDischargeScreenLevelsLocked(screenOn, screenOn);
-            updateTimeBasesLocked(false, !screenOn, uptime, realtime);
+            updateDischargeScreenLevelsLocked(screenState, screenState);
+            updateTimeBasesLocked(false, screenState, uptime, realtime);
             mChargeStepTracker.init();
             mLastChargeStepLevel = level;
             mMaxChargeStepLevel = level;
@@ -10914,6 +10990,9 @@ public class BatteryStatsImpl extends BatteryStats {
                 final long chargeDiff = mHistoryCur.batteryChargeUAh - chargeUAh;
                 mDischargeCounter.addCountLocked(chargeDiff);
                 mDischargeScreenOffCounter.addCountLocked(chargeDiff);
+                if (isScreenDoze(mScreenState)) {
+                    mDischargeScreenDozeCounter.addCountLocked(chargeDiff);
+                }
             }
             mHistoryCur.batteryChargeUAh = chargeUAh;
             setOnBatteryLocked(elapsedRealtime, uptime, onBattery, oldStatus, level, chargeUAh);
@@ -10956,6 +11035,9 @@ public class BatteryStatsImpl extends BatteryStats {
                     final long chargeDiff = mHistoryCur.batteryChargeUAh - chargeUAh;
                     mDischargeCounter.addCountLocked(chargeDiff);
                     mDischargeScreenOffCounter.addCountLocked(chargeDiff);
+                    if (isScreenDoze(mScreenState)) {
+                        mDischargeScreenDozeCounter.addCountLocked(chargeDiff);
+                    }
                 }
                 mHistoryCur.batteryChargeUAh = chargeUAh;
                 changed = true;
@@ -11264,10 +11346,11 @@ public class BatteryStatsImpl extends BatteryStats {
         return dischargeAmount;
     }
 
+    @Override
     public int getDischargeAmountScreenOn() {
         synchronized(this) {
             int val = mDischargeAmountScreenOn;
-            if (mOnBattery && mScreenState == Display.STATE_ON
+            if (mOnBattery && isScreenOn(mScreenState)
                     && mDischargeCurrentLevel < mDischargeScreenOnUnplugLevel) {
                 val += mDischargeScreenOnUnplugLevel-mDischargeCurrentLevel;
             }
@@ -11275,10 +11358,11 @@ public class BatteryStatsImpl extends BatteryStats {
         }
     }
 
+    @Override
     public int getDischargeAmountScreenOnSinceCharge() {
         synchronized(this) {
             int val = mDischargeAmountScreenOnSinceCharge;
-            if (mOnBattery && mScreenState == Display.STATE_ON
+            if (mOnBattery && isScreenOn(mScreenState)
                     && mDischargeCurrentLevel < mDischargeScreenOnUnplugLevel) {
                 val += mDischargeScreenOnUnplugLevel-mDischargeCurrentLevel;
             }
@@ -11286,24 +11370,52 @@ public class BatteryStatsImpl extends BatteryStats {
         }
     }
 
+    @Override
     public int getDischargeAmountScreenOff() {
         synchronized(this) {
             int val = mDischargeAmountScreenOff;
-            if (mOnBattery && mScreenState != Display.STATE_ON
+            if (mOnBattery && isScreenOff(mScreenState)
                     && mDischargeCurrentLevel < mDischargeScreenOffUnplugLevel) {
                 val += mDischargeScreenOffUnplugLevel-mDischargeCurrentLevel;
             }
-            return val;
+            // For backward compatibility, doze discharge is counted into screen off.
+            return val + getDischargeAmountScreenDoze();
         }
     }
 
+    @Override
     public int getDischargeAmountScreenOffSinceCharge() {
         synchronized(this) {
             int val = mDischargeAmountScreenOffSinceCharge;
-            if (mOnBattery && mScreenState != Display.STATE_ON
+            if (mOnBattery && isScreenOff(mScreenState)
                     && mDischargeCurrentLevel < mDischargeScreenOffUnplugLevel) {
                 val += mDischargeScreenOffUnplugLevel-mDischargeCurrentLevel;
             }
+            // For backward compatibility, doze discharge is counted into screen off.
+            return val + getDischargeAmountScreenDozeSinceCharge();
+        }
+    }
+
+    @Override
+    public int getDischargeAmountScreenDoze() {
+        synchronized(this) {
+            int val = mDischargeAmountScreenDoze;
+            if (mOnBattery && isScreenDoze(mScreenState)
+                    && mDischargeCurrentLevel < mDischargeScreenDozeUnplugLevel) {
+                val += mDischargeScreenDozeUnplugLevel-mDischargeCurrentLevel;
+            }
+            return val;
+        }
+    }
+
+    @Override
+    public int getDischargeAmountScreenDozeSinceCharge() {
+        synchronized(this) {
+            int val = mDischargeAmountScreenDozeSinceCharge;
+            if (mOnBattery && isScreenDoze(mScreenState)
+                    && mDischargeCurrentLevel < mDischargeScreenDozeUnplugLevel) {
+                val += mDischargeScreenDozeUnplugLevel-mDischargeCurrentLevel;
+            }
             return val;
         }
     }
@@ -11661,12 +11773,14 @@ public class BatteryStatsImpl extends BatteryStats {
         mHighDischargeAmountSinceCharge = in.readInt();
         mDischargeAmountScreenOnSinceCharge = in.readInt();
         mDischargeAmountScreenOffSinceCharge = in.readInt();
+        mDischargeAmountScreenDozeSinceCharge = in.readInt();
         mDischargeStepTracker.readFromParcel(in);
         mChargeStepTracker.readFromParcel(in);
         mDailyDischargeStepTracker.readFromParcel(in);
         mDailyChargeStepTracker.readFromParcel(in);
         mDischargeCounter.readSummaryFromParcelLocked(in);
         mDischargeScreenOffCounter.readSummaryFromParcelLocked(in);
+        mDischargeScreenDozeCounter.readSummaryFromParcelLocked(in);
         int NPKG = in.readInt();
         if (NPKG > 0) {
             mDailyPackageChanges = new ArrayList<>(NPKG);
@@ -11689,6 +11803,7 @@ public class BatteryStatsImpl extends BatteryStats {
 
         mScreenState = Display.STATE_UNKNOWN;
         mScreenOnTimer.readSummaryFromParcelLocked(in);
+        mScreenDozeTimer.readSummaryFromParcelLocked(in);
         for (int i=0; i<NUM_SCREEN_BRIGHTNESS_BINS; i++) {
             mScreenBrightnessTimer[i].readSummaryFromParcelLocked(in);
         }
@@ -12084,12 +12199,14 @@ public class BatteryStatsImpl extends BatteryStats {
         out.writeInt(getHighDischargeAmountSinceCharge());
         out.writeInt(getDischargeAmountScreenOnSinceCharge());
         out.writeInt(getDischargeAmountScreenOffSinceCharge());
+        out.writeInt(getDischargeAmountScreenDozeSinceCharge());
         mDischargeStepTracker.writeToParcel(out);
         mChargeStepTracker.writeToParcel(out);
         mDailyDischargeStepTracker.writeToParcel(out);
         mDailyChargeStepTracker.writeToParcel(out);
         mDischargeCounter.writeSummaryFromParcelLocked(out);
         mDischargeScreenOffCounter.writeSummaryFromParcelLocked(out);
+        mDischargeScreenDozeCounter.writeSummaryFromParcelLocked(out);
         if (mDailyPackageChanges != null) {
             final int NPKG = mDailyPackageChanges.size();
             out.writeInt(NPKG);
@@ -12107,6 +12224,7 @@ public class BatteryStatsImpl extends BatteryStats {
         out.writeLong(mNextMaxDailyDeadline);
 
         mScreenOnTimer.writeSummaryFromParcelLocked(out, NOWREAL_SYS);
+        mScreenDozeTimer.writeSummaryFromParcelLocked(out, NOWREAL_SYS);
         for (int i=0; i<NUM_SCREEN_BRIGHTNESS_BINS; i++) {
             mScreenBrightnessTimer[i].writeSummaryFromParcelLocked(out, NOWREAL_SYS);
         }
@@ -12541,6 +12659,7 @@ public class BatteryStatsImpl extends BatteryStats {
 
         mScreenState = Display.STATE_UNKNOWN;
         mScreenOnTimer = new StopwatchTimer(mClocks, null, -1, null, mOnBatteryTimeBase, in);
+        mScreenDozeTimer = new StopwatchTimer(mClocks, null, -1, null, mOnBatteryTimeBase, in);
         for (int i=0; i<NUM_SCREEN_BRIGHTNESS_BINS; i++) {
             mScreenBrightnessTimer[i] = new StopwatchTimer(mClocks, null, -100-i, null,
                     mOnBatteryTimeBase, in);
@@ -12634,10 +12753,13 @@ public class BatteryStatsImpl extends BatteryStats {
         mDischargeAmountScreenOnSinceCharge = in.readInt();
         mDischargeAmountScreenOff = in.readInt();
         mDischargeAmountScreenOffSinceCharge = in.readInt();
+        mDischargeAmountScreenDoze = in.readInt();
+        mDischargeAmountScreenDozeSinceCharge = in.readInt();
         mDischargeStepTracker.readFromParcel(in);
         mChargeStepTracker.readFromParcel(in);
         mDischargeCounter = new LongSamplingCounter(mOnBatteryTimeBase, in);
-        mDischargeScreenOffCounter = new LongSamplingCounter(mOnBatteryTimeBase, in);
+        mDischargeScreenOffCounter = new LongSamplingCounter(mOnBatteryScreenOffTimeBase, in);
+        mDischargeScreenDozeCounter = new LongSamplingCounter(mOnBatteryTimeBase, in);
         mLastWriteTime = in.readLong();
 
         mRpmStats.clear();
@@ -12756,6 +12878,7 @@ public class BatteryStatsImpl extends BatteryStats {
         mOnBatteryScreenOffTimeBase.writeToParcel(out, uSecUptime, uSecRealtime);
 
         mScreenOnTimer.writeToParcel(out, uSecRealtime);
+        mScreenDozeTimer.writeToParcel(out, uSecRealtime);
         for (int i=0; i<NUM_SCREEN_BRIGHTNESS_BINS; i++) {
             mScreenBrightnessTimer[i].writeToParcel(out, uSecRealtime);
         }
@@ -12818,10 +12941,13 @@ public class BatteryStatsImpl extends BatteryStats {
         out.writeInt(mDischargeAmountScreenOnSinceCharge);
         out.writeInt(mDischargeAmountScreenOff);
         out.writeInt(mDischargeAmountScreenOffSinceCharge);
+        out.writeInt(mDischargeAmountScreenDoze);
+        out.writeInt(mDischargeAmountScreenDozeSinceCharge);
         mDischargeStepTracker.writeToParcel(out);
         mChargeStepTracker.writeToParcel(out);
         mDischargeCounter.writeToParcel(out);
         mDischargeScreenOffCounter.writeToParcel(out);
+        mDischargeScreenDozeCounter.writeToParcel(out);
         out.writeLong(mLastWriteTime);
 
         out.writeInt(mRpmStats.size());
@@ -12930,8 +13056,10 @@ public class BatteryStatsImpl extends BatteryStats {
             pw.println("mOnBatteryScreenOffTimeBase:");
             mOnBatteryScreenOffTimeBase.dump(pw, "  ");
             Printer pr = new PrintWriterPrinter(pw);
-            pr.println("*** Screen timer:");
+            pr.println("*** Screen on timer:");
             mScreenOnTimer.logState(pr, "  ");
+            pr.println("*** Screen doze timer:");
+            mScreenDozeTimer.logState(pr, "  ");
             for (int i=0; i<NUM_SCREEN_BRIGHTNESS_BINS; i++) {
                 pr.println("*** Screen brightness #" + i + ":");
                 mScreenBrightnessTimer[i].logState(pr, "  ");
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsBackgroundStatsTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsBackgroundStatsTest.java
index eff6ad947744..c611a01c6e16 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsBackgroundStatsTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsBackgroundStatsTest.java
@@ -22,6 +22,7 @@ import android.os.BatteryStats;
 import android.os.WorkSource;
 import android.support.test.filters.SmallTest;
 import android.util.ArrayMap;
+import android.view.Display;
 
 import junit.framework.TestCase;
 
@@ -44,7 +45,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
         // Off-battery, non-existent
         clocks.realtime = clocks.uptime = 10;
         cur = clocks.realtime * 1000;
-        bi.updateTimeBasesLocked(false, false, cur, cur); // off battery
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, cur, cur); // off battery
         assertFalse(bgtb.isRunning());
         assertEquals(0, bgtb.computeRealtime(cur, STATS_SINCE_CHARGED));
 
@@ -65,7 +66,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
         // On-battery, background
         clocks.realtime = clocks.uptime = 303;
         cur = clocks.realtime * 1000;
-        bi.updateTimeBasesLocked(true, false, cur, cur); // on battery
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, cur, cur); // on battery
         // still in ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND
         assertTrue(bgtb.isRunning());
         assertEquals(0, bgtb.computeRealtime(cur, STATS_SINCE_CHARGED));
@@ -73,7 +74,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
         // On-battery, background - but change screen state
         clocks.realtime = clocks.uptime = 409;
         cur = clocks.realtime * 1000;
-        bi.updateTimeBasesLocked(true, true, cur, cur); // on battery (again)
+        bi.updateTimeBasesLocked(true, Display.STATE_OFF, cur, cur); // on battery (again)
         assertTrue(bgtb.isRunning());
         assertEquals(106_000, bgtb.computeRealtime(cur, STATS_SINCE_CHARGED));
 
@@ -87,7 +88,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
         // Off-battery, foreground
         clocks.realtime = clocks.uptime = 530;
         cur = clocks.realtime * 1000;
-        bi.updateTimeBasesLocked(false, false, cur, cur); // off battery
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, cur, cur); // off battery
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND);
         assertFalse(bgtb.isRunning());
         assertEquals(227_000, bgtb.computeRealtime(cur, STATS_SINCE_CHARGED));
@@ -112,17 +113,17 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
         // battery=off, screen=off, background=off
         cur = (clocks.realtime = clocks.uptime = 100) * 1000;
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND);
-        bi.updateTimeBasesLocked(false, false, cur, cur);
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, cur, cur);
         assertFalse(bgtb.isRunning());
 
         // battery=on, screen=off, background=off
         cur = (clocks.realtime = clocks.uptime = 200) * 1000;
-        bi.updateTimeBasesLocked(true, false, cur, cur);
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, cur, cur);
         assertFalse(bgtb.isRunning());
 
         // battery=on, screen=on, background=off
         cur = (clocks.realtime = clocks.uptime = 300) * 1000;
-        bi.updateTimeBasesLocked(true, true, cur, cur);
+        bi.updateTimeBasesLocked(true, Display.STATE_OFF, cur, cur);
         assertFalse(bgtb.isRunning());
 
         // battery=on, screen=on, background=on
@@ -133,12 +134,12 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
 
         // battery=on, screen=off, background=on
         cur = (clocks.realtime = clocks.uptime = 550) * 1000;
-        bi.updateTimeBasesLocked(true, false, cur, cur);
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, cur, cur);
         assertFalse(bgtb.isRunning());
 
         // battery=off, screen=off, background=on
         cur = (clocks.realtime = clocks.uptime = 660) * 1000;
-        bi.updateTimeBasesLocked(false, false, cur, cur);
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, cur, cur);
         assertFalse(bgtb.isRunning());
 
         // battery=off, screen=off, background=off
@@ -157,7 +158,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
 
         // On battery
         curr = 1000 * (clocks.realtime = clocks.uptime = 100);
-        bi.updateTimeBasesLocked(true, false, curr, curr); // on battery
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, curr, curr); // on battery
         // App in foreground
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND);
 
@@ -171,7 +172,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
 
         // Off battery
         curr = 1000 * (clocks.realtime = clocks.uptime = 305);
-        bi.updateTimeBasesLocked(false, false, curr, curr); // off battery
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, curr, curr); // off battery
 
         // Stop timer
         curr = 1000 * (clocks.realtime = clocks.uptime = 409);
@@ -200,7 +201,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
 
         // On battery
         curr = 1000 * (clocks.realtime = clocks.uptime = 100);
-        bi.updateTimeBasesLocked(true, false, curr, curr); // on battery
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, curr, curr); // on battery
 
         // App in foreground
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND);
@@ -215,7 +216,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
 
         // Off battery
         curr = 1000 * (clocks.realtime = clocks.uptime = 305);
-        bi.updateTimeBasesLocked(false, false, curr, curr); // off battery
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, curr, curr); // off battery
 
         // Start timer (unoptimized)
         curr = 1000 * (clocks.realtime = clocks.uptime = 1000);
@@ -223,7 +224,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
 
         // On battery
         curr = 1000 * (clocks.realtime = clocks.uptime = 2001);
-        bi.updateTimeBasesLocked(true, false, curr, curr); // on battery
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, curr, curr); // on battery
 
         // Move to foreground
         curr = 1000 * (clocks.realtime = clocks.uptime = 3004);
@@ -270,7 +271,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
 
         // On battery
         curr = 1000 * (clocks.realtime = clocks.uptime = 100);
-        bi.updateTimeBasesLocked(true, false, curr, curr); // on battery
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, curr, curr); // on battery
         // App in foreground
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND);
 
@@ -292,7 +293,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
 
         // Off battery
         curr = 1000 * (clocks.realtime = clocks.uptime = 305);
-        bi.updateTimeBasesLocked(false, false, curr, curr); // off battery
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, curr, curr); // off battery
 
         // Stop timer
         curr = 1000 * (clocks.realtime = clocks.uptime = 409);
@@ -331,7 +332,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
 
         // On battery
         curr = 1000 * (clocks.realtime = clocks.uptime = 100);
-        bi.updateTimeBasesLocked(true, false, curr, curr); // on battery
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, curr, curr); // on battery
         // App in foreground
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND);
 
@@ -353,7 +354,7 @@ public class BatteryStatsBackgroundStatsTest extends TestCase {
 
         // Off battery
         curr = 1000 * (clocks.realtime = clocks.uptime = 305);
-        bi.updateTimeBasesLocked(false, false, curr, curr); // off battery
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, curr, curr); // off battery
 
         // Stop timer
         curr = 1000 * (clocks.realtime = clocks.uptime = 409);
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
index 7769e69d317b..461d537e86f9 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
@@ -19,9 +19,11 @@ import static android.os.BatteryStats.STATS_SINCE_CHARGED;
 import static android.os.BatteryStats.WAKE_TYPE_PARTIAL;
 
 import android.app.ActivityManager;
+import android.os.BatteryManager;
 import android.os.BatteryStats;
 import android.os.WorkSource;
 import android.support.test.filters.SmallTest;
+import android.view.Display;
 
 import junit.framework.TestCase;
 
@@ -50,7 +52,7 @@ public class BatteryStatsNoteTest extends TestCase{
     @SmallTest
     public void testNoteBluetoothScanResultLocked() throws Exception {
         MockBatteryStatsImpl bi = new MockBatteryStatsImpl(new MockClocks());
-        bi.updateTimeBasesLocked(true, true, 0, 0);
+        bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_TOP);
 
         bi.noteBluetoothScanResultsFromSourceLocked(WS, 1);
@@ -82,7 +84,7 @@ public class BatteryStatsNoteTest extends TestCase{
         int pid = 10;
         String name = "name";
 
-        bi.updateTimeBasesLocked(true, true, 0, 0);
+        bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_TOP);
         bi.getUidStatsLocked(UID).noteStartWakeLocked(pid, name, WAKE_TYPE_PARTIAL, clocks.realtime);
 
@@ -124,7 +126,7 @@ public class BatteryStatsNoteTest extends TestCase{
         stateRuntimeMap.put(ActivityManager.PROCESS_STATE_CACHED_ACTIVITY_CLIENT, 5309);
         stateRuntimeMap.put(ActivityManager.PROCESS_STATE_CACHED_EMPTY, 42);
 
-        bi.updateTimeBasesLocked(true, false, 0, 0);
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, 0, 0);
 
         for (Map.Entry<Integer, Integer> entry : stateRuntimeMap.entrySet()) {
             bi.noteUidProcessStateLocked(UID, entry.getKey());
@@ -189,4 +191,46 @@ public class BatteryStatsNoteTest extends TestCase{
         expectedRunTimeMs = stateRuntimeMap.get(ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE);
         assertEquals(expectedRunTimeMs * 1000, actualRunTimeUs);
     }
+
+    /** Test BatteryStatsImpl.updateTimeBasesLocked. */
+    @SmallTest
+    public void testUpdateTimeBasesLocked() throws Exception {
+        final MockClocks clocks = new MockClocks(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+
+        bi.updateTimeBasesLocked(false, Display.STATE_OFF, 0, 0);
+        assertFalse(bi.getOnBatteryTimeBase().isRunning());
+        bi.updateTimeBasesLocked(false, Display.STATE_DOZE, 10, 10);
+        assertFalse(bi.getOnBatteryTimeBase().isRunning());
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, 20, 20);
+        assertFalse(bi.getOnBatteryTimeBase().isRunning());
+
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, 30, 30);
+        assertTrue(bi.getOnBatteryTimeBase().isRunning());
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        bi.updateTimeBasesLocked(true, Display.STATE_DOZE, 40, 40);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        bi.updateTimeBasesLocked(true, Display.STATE_OFF, 40, 40);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+    }
+
+    /** Test BatteryStatsImpl.noteScreenStateLocked. */
+    @SmallTest
+    public void testNoteScreenStateLocked() throws Exception {
+        final MockClocks clocks = new MockClocks(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, 0, 0);
+        bi.noteScreenStateLocked(Display.STATE_ON);
+        bi.noteScreenStateLocked(Display.STATE_DOZE);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(bi.getScreenState(), Display.STATE_DOZE);
+        bi.noteScreenStateLocked(Display.STATE_ON);
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(bi.getScreenState(), Display.STATE_ON);
+        bi.noteScreenStateLocked(Display.STATE_OFF);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(bi.getScreenState(), Display.STATE_OFF);
+    }
+
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsSensorTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsSensorTest.java
index 0294095cca9e..a751f9060588 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsSensorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsSensorTest.java
@@ -18,6 +18,7 @@ package com.android.internal.os;
 import android.app.ActivityManager;
 import android.os.BatteryStats;
 import android.support.test.filters.SmallTest;
+import android.view.Display;
 
 import junit.framework.TestCase;
 
@@ -74,7 +75,7 @@ public class BatteryStatsSensorTest extends TestCase {
 
         // Plugged-in (battery=off, sensor=off)
         curr = 1000 * (clocks.realtime = clocks.uptime = 100);
-        bi.updateTimeBasesLocked(false, false, curr, curr);
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, curr, curr);
 
 
         // Start sensor (battery=off, sensor=on)
@@ -110,7 +111,7 @@ public class BatteryStatsSensorTest extends TestCase {
 
         // Unplugged (battery=on, sensor=off)
         curr = 1000 * (clocks.realtime = clocks.uptime = 100);
-        bi.updateTimeBasesLocked(true, false, curr, curr);
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, curr, curr);
 
         // Start sensor (battery=on, sensor=on)
         curr = 1000 * (clocks.realtime = clocks.uptime = 200);
@@ -145,7 +146,7 @@ public class BatteryStatsSensorTest extends TestCase {
 
         // On battery (battery=on, sensor=off)
         curr = 1000 * (clocks.realtime = clocks.uptime = 100);
-        bi.updateTimeBasesLocked(true, false, curr, curr);
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, curr, curr);
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND);
 
         // Start sensor (battery=on, sensor=on)
@@ -154,7 +155,7 @@ public class BatteryStatsSensorTest extends TestCase {
 
         // Off battery (battery=off, sensor=on)
         curr = 1000 * (clocks.realtime = clocks.uptime = 305);
-        bi.updateTimeBasesLocked(false, false, curr, curr);
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, curr, curr);
 
         // Stop sensor while off battery (battery=off, sensor=off)
         curr = 1000 * (clocks.realtime = clocks.uptime = 409);
@@ -191,7 +192,7 @@ public class BatteryStatsSensorTest extends TestCase {
 
         // Plugged-in (battery=off, sensor=off)
         curr = 1000 * (clocks.realtime = clocks.uptime = 100);
-        bi.updateTimeBasesLocked(false, false, curr, curr);
+        bi.updateTimeBasesLocked(false, Display.STATE_ON, curr, curr);
 
         // Start sensor (battery=off, sensor=on)
         curr = 1000 * (clocks.realtime = clocks.uptime = 200);
@@ -209,7 +210,7 @@ public class BatteryStatsSensorTest extends TestCase {
 
         // Unplug (battery=on, sensor=on)
         curr = 1000 * (clocks.realtime = clocks.uptime = 305);
-        bi.updateTimeBasesLocked(true, false, curr, curr);
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, curr, curr);
 
         //Test situation
         curr = 1000 * (clocks.realtime = clocks.uptime = 410);
@@ -243,7 +244,7 @@ public class BatteryStatsSensorTest extends TestCase {
         long curr = 0; // realtime in us
         // Entire test is on-battery
         curr = 1000 * (clocks.realtime = clocks.uptime = 1000);
-        bi.updateTimeBasesLocked(true, false, curr, curr);
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, curr, curr);
 
         // See below for a diagram of events.
 
diff --git a/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java b/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java
index a123fce0748c..ae42f78da87d 100644
--- a/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java
+++ b/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java
@@ -26,6 +26,10 @@ public class MockBatteryStatsImpl extends BatteryStatsImpl {
     MockBatteryStatsImpl(Clocks clocks) {
         super(clocks);
         this.clocks = mClocks;
+        mScreenOnTimer = new BatteryStatsImpl.StopwatchTimer(clocks, null, -1, null,
+                mOnBatteryTimeBase);
+        mScreenDozeTimer = new BatteryStatsImpl.StopwatchTimer(clocks, null, -1, null,
+                mOnBatteryTimeBase);
         mBluetoothScanTimer = new StopwatchTimer(mClocks, null, -14, null, mOnBatteryTimeBase);
     }
 
@@ -37,6 +41,14 @@ public class MockBatteryStatsImpl extends BatteryStatsImpl {
         return mOnBatteryTimeBase;
     }
 
+    public TimeBase getOnBatteryScreenOffTimeBase() {
+        return mOnBatteryScreenOffTimeBase;
+    }
+
+    public int getScreenState() {
+        return mScreenState;
+    }
+
     public boolean isOnBattery() {
         return mForceOnBattery ? true : super.isOnBattery();
     }