Enable batterystats for all display IDs

Batterystats used to be enabled only for the default display. This
change enables it for all displays (both for noting screen state and
brightness).

Bug: 366112793
Test: atest DisplayPowerControllerTest
Flag: com.android.server.display.feature.flags.enable_battery_stats_for_all_displays
Change-Id: I12e6ca4c9777fbdc3023e014dbd43fb2361eff06
diff --git a/core/java/com/android/internal/app/IBatteryStats.aidl b/core/java/com/android/internal/app/IBatteryStats.aidl
index fc198cb..a5e166b 100644
--- a/core/java/com/android/internal/app/IBatteryStats.aidl
+++ b/core/java/com/android/internal/app/IBatteryStats.aidl
@@ -146,7 +146,7 @@
     @EnforcePermission("UPDATE_DEVICE_STATS")
     void noteScreenState(int displayId, int state, int reason);
     @EnforcePermission("UPDATE_DEVICE_STATS")
-    void noteScreenBrightness(int brightness);
+    void noteScreenBrightness(int displayId, int brightness);
     @EnforcePermission("UPDATE_DEVICE_STATS")
     void noteUserActivity(int uid, int event);
     @EnforcePermission("UPDATE_DEVICE_STATS")
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index a7b08a5..5cb2a26 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -1855,7 +1855,7 @@
 
     @Override
     @EnforcePermission(UPDATE_DEVICE_STATS)
-    public void noteScreenBrightness(final int brightness) {
+    public void noteScreenBrightness(final int displayId, final int brightness) {
         super.noteScreenBrightness_enforcePermission();
 
         synchronized (mLock) {
@@ -1863,7 +1863,8 @@
             final long uptime = SystemClock.uptimeMillis();
             mHandler.post(() -> {
                 synchronized (mStats) {
-                    mStats.noteScreenBrightnessLocked(0, brightness, elapsedRealtime, uptime);
+                    mStats.noteScreenBrightnessLocked(
+                            displayId, brightness, elapsedRealtime, uptime);
                 }
             });
         }
diff --git a/services/core/java/com/android/server/display/DisplayPowerController.java b/services/core/java/com/android/server/display/DisplayPowerController.java
index 557c50d..04573f4 100644
--- a/services/core/java/com/android/server/display/DisplayPowerController.java
+++ b/services/core/java/com/android/server/display/DisplayPowerController.java
@@ -536,7 +536,9 @@
         mLastBrightnessEvent = new BrightnessEvent(mDisplayId);
         mTempBrightnessEvent = new BrightnessEvent(mDisplayId);
 
-        if (mDisplayId == Display.DEFAULT_DISPLAY) {
+        if (flags.isBatteryStatsEnabledForAllDisplays()) {
+            mBatteryStats = BatteryStatsService.getService();
+        } else if (mDisplayId == Display.DEFAULT_DISPLAY) {
             mBatteryStats = BatteryStatsService.getService();
         } else {
             mBatteryStats = null;
@@ -2791,8 +2793,7 @@
                 screenState, mDisplayStatsId, reason);
         if (mBatteryStats != null) {
             try {
-                // TODO(multi-display): make this multi-display
-                mBatteryStats.noteScreenState(/* displayId= */ 0, screenState, reason);
+                mBatteryStats.noteScreenState(mDisplayId, screenState, reason);
             } catch (RemoteException e) {
                 // same process
             }
@@ -2807,7 +2808,7 @@
                 int brightnessInt = mFlags.isBrightnessIntRangeUserPerceptionEnabled()
                         ? BrightnessSynchronizer.brightnessFloatToIntSetting(mContext, brightness)
                         : BrightnessSynchronizer.brightnessFloatToInt(brightness);
-                mBatteryStats.noteScreenBrightness(brightnessInt);
+                mBatteryStats.noteScreenBrightness(mDisplayId, brightnessInt);
             } catch (RemoteException e) {
                 // same process
             }
diff --git a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
index 69b67c8..f600e7f 100644
--- a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
+++ b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
@@ -200,6 +200,11 @@
             Flags::normalBrightnessForDozeParameter
     );
 
+    private final FlagState mEnableBatteryStatsForAllDisplays = new FlagState(
+            Flags.FLAG_ENABLE_BATTERY_STATS_FOR_ALL_DISPLAYS,
+            Flags::enableBatteryStatsForAllDisplays
+    );
+
     /**
      * @return {@code true} if 'port' is allowed in display layout configuration file.
      */
@@ -415,6 +420,14 @@
     }
 
     /**
+      * @return {@code true} if battery stats is enabled for all displays, not just the primary
+      * display.
+      */
+    public boolean isBatteryStatsEnabledForAllDisplays() {
+        return mEnableBatteryStatsForAllDisplays.isEnabled();
+    }
+
+    /**
      * dumps all flagstates
      * @param pw printWriter
      */
@@ -456,6 +469,7 @@
         pw.println(" " + mNewHdrBrightnessModifier);
         pw.println(" " + mNormalBrightnessForDozeParameter);
         pw.println(" " + mIdleScreenConfigInSubscribingLightSensor);
+        pw.println(" " + mEnableBatteryStatsForAllDisplays);
     }
 
     private static class FlagState {
diff --git a/services/core/java/com/android/server/display/feature/display_flags.aconfig b/services/core/java/com/android/server/display/feature/display_flags.aconfig
index da5063a..ce66e80 100644
--- a/services/core/java/com/android/server/display/feature/display_flags.aconfig
+++ b/services/core/java/com/android/server/display/feature/display_flags.aconfig
@@ -341,3 +341,11 @@
       purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "enable_battery_stats_for_all_displays"
+    namespace: "display_manager"
+    description: "Flag to enable battery stats for all displays."
+    bug: "366112793"
+    is_fixed_read_only: true
+}
\ No newline at end of file
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
index d0aec3b..bf5a692 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
@@ -75,6 +75,7 @@
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
 
+import com.android.internal.app.IBatteryStats;
 import com.android.modules.utils.testing.ExtendedMockitoRule;
 import com.android.server.LocalServices;
 import com.android.server.am.BatteryStatsService;
@@ -158,6 +159,8 @@
     private DisplayManagerFlags mDisplayManagerFlagsMock;
     @Mock
     private DisplayManagerInternal.DisplayOffloadSession mDisplayOffloadSession;
+    @Mock
+    private IBatteryStats mMockBatteryStats;
     @Captor
     private ArgumentCaptor<SensorEventListener> mSensorEventListenerCaptor;
 
@@ -204,7 +207,8 @@
 
         doAnswer((Answer<Void>) invocationOnMock -> null).when(() ->
                 SystemProperties.set(anyString(), any()));
-        doAnswer((Answer<Void>) invocationOnMock -> null).when(BatteryStatsService::getService);
+        doAnswer((Answer<IBatteryStats>) invocationOnMock -> mMockBatteryStats)
+                .when(BatteryStatsService::getService);
         doAnswer((Answer<Boolean>) invocationOnMock -> false)
                 .when(ActivityManager::isLowRamDeviceStatic);
 
@@ -2227,6 +2231,52 @@
         verify(mHolder.brightnessSetting).saveIfNeeded();
     }
 
+    @Test
+    public void testBatteryStatNotes_enabledOnDefaultDisplayWhenDisabledOnOthers()
+            throws Exception {
+        when(mDisplayManagerFlagsMock.isBatteryStatsEnabledForAllDisplays()).thenReturn(false);
+
+        verifyNoteScreenState(Display.DEFAULT_DISPLAY, /* expectNote= */ true);
+    }
+
+    @Test
+    public void testBatteryStatNotes_enabledOnDefaultDisplayWhenEnabledOnOthers() throws Exception {
+        when(mDisplayManagerFlagsMock.isBatteryStatsEnabledForAllDisplays()).thenReturn(true);
+
+        verifyNoteScreenState(Display.DEFAULT_DISPLAY, /* expectNote= */ true);
+    }
+
+    @Test
+    public void testBatteryStatNotes_flagGuardedOnNonDefaultDisplays() throws Exception {
+        when(mDisplayManagerFlagsMock.isBatteryStatsEnabledForAllDisplays()).thenReturn(false);
+
+        verifyNoteScreenState(/* displayId= */ 2, /* expectNote= */ false);
+
+        when(mDisplayManagerFlagsMock.isBatteryStatsEnabledForAllDisplays()).thenReturn(true);
+
+        verifyNoteScreenState(/* displayId= */ 2, /* expectNote= */ true);
+    }
+
+    private void verifyNoteScreenState(int displayId, boolean expectNote) throws Exception {
+        mHolder = createDisplayPowerController(displayId, UNIQUE_ID);
+        DisplayPowerRequest dpr = new DisplayPowerRequest();
+        dpr.policy = DisplayPowerRequest.POLICY_BRIGHT;
+        when(mHolder.displayPowerState.getScreenState()).thenReturn(Display.STATE_ON);
+
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        if (expectNote) {
+            verify(mMockBatteryStats)
+                    .noteScreenState(
+                            displayId, Display.STATE_ON, Display.STATE_REASON_DEFAULT_POLICY);
+            verify(mMockBatteryStats).noteScreenBrightness(eq(displayId), anyInt());
+        } else {
+            verify(mMockBatteryStats, never()).noteScreenState(anyInt(), anyInt(), anyInt());
+            verify(mMockBatteryStats, never()).noteScreenBrightness(anyInt(), anyInt());
+        }
+    }
+
     /**
      * Creates a mock and registers it to {@link LocalServices}.
      */