Fix screen-off power estimation error caused by proportional distribution

Bug: 383037808
Test: atest PowerStatsTests
Flag: EXEMPT bugfix
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:ecd4cfb40fe79b977294894aabcf7622880f6866)
Merged-In: I77f7d1aa9841ccfaf959838de94e15e0de95c805
Change-Id: I77f7d1aa9841ccfaf959838de94e15e0de95c805
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 6f18107..c26eeed 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -5723,7 +5723,9 @@
                     displayStats.screenDozeTimer.stopRunningLocked(elapsedRealtimeMs);
                     shouldScheduleSync = true;
                     break;
-                case Display.STATE_OFF: // fallthrough
+                case Display.STATE_OFF:
+                    shouldScheduleSync = true;
+                    break;
                 case Display.STATE_UNKNOWN:
                     // Not tracked by timers.
                     break;
@@ -5756,7 +5758,9 @@
                     displayStats.screenDozeTimer.startRunningLocked(elapsedRealtimeMs);
                     shouldScheduleSync = true;
                     break;
-                case Display.STATE_OFF: // fallthrough
+                case Display.STATE_OFF:
+                    shouldScheduleSync = true;
+                    break;
                 case Display.STATE_UNKNOWN:
                     // Not tracked by timers.
                     break;
@@ -5873,7 +5877,7 @@
 
         if (shouldScheduleSync) {
             if (mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_SCREEN)) {
-                mScreenPowerStatsCollector.schedule();
+                mScreenPowerStatsCollector.onScreenStateChange();
             } else {
                 final int numDisplays = mPerDisplayBatteryStats.length;
                 final int[] displayStates = new int[numDisplays];
diff --git a/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
index c38904f..90039e8 100644
--- a/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
@@ -111,6 +111,22 @@
         return true;
     }
 
+    /**
+     * Must be called whenever the screen state (on/off/doze) changes.
+     */
+    public void onScreenStateChange() {
+        if (ensureInitialized() && mConsumedEnergyHelper.getEnergyConsumerCount() != 0) {
+            // Sync power monitor reading immediately, because the estimation algorithm
+            // distributes consumed power proportionally between screen states.
+            // Since screen power consumption differs dramatically between different states,
+            // this would lead an overestimation in the screen-off state.
+            forceSchedule();
+            return;
+        }
+        // Perhaps schedule a sync, allowing throttling
+        schedule();
+    }
+
     @Override
     public PowerStats collectStats() {
         if (!ensureInitialized()) {
@@ -126,7 +142,7 @@
             long screenOnTimeMs = mScreenUsageTimeRetriever.getScreenOnTimeMs(display);
             if (!mFirstSample) {
                 mLayout.setScreenOnDuration(mPowerStats.stats, display,
-                        screenOnTimeMs - mLastScreenOnTime[display]);
+                        Math.max(0, screenOnTimeMs - mLastScreenOnTime[display]));
             }
             mLastScreenOnTime[display] = screenOnTimeMs;
 
@@ -135,14 +151,15 @@
                         mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(display, level);
                 if (!mFirstSample) {
                     mLayout.setBrightnessLevelDuration(mPowerStats.stats, display, level,
-                            brightnessLevelTimeMs - mLastBrightnessLevelTime[display][level]);
+                            Math.max(0, brightnessLevelTimeMs
+                                    - mLastBrightnessLevelTime[display][level]));
                 }
                 mLastBrightnessLevelTime[display][level] = brightnessLevelTimeMs;
             }
             long screenDozeTimeMs = mScreenUsageTimeRetriever.getScreenDozeTimeMs(display);
             if (!mFirstSample) {
                 mLayout.setScreenDozeDuration(mPowerStats.stats, display,
-                        screenDozeTimeMs - mLastDozeTime[display]);
+                        Math.max(0, screenDozeTimeMs - mLastDozeTime[display]));
             }
             mLastDozeTime[display] = screenDozeTimeMs;
         }
@@ -162,7 +179,7 @@
             }
 
             mLayout.setUidTopActivityDuration(uidStats,
-                    mLayout.getUidTopActivityDuration(uidStats) + topActivityDuration);
+                    Math.max(0, mLayout.getUidTopActivityDuration(uidStats) + topActivityDuration));
         });
 
         long elapsedRealtime = mClock.elapsedRealtime();