Move onWakelockChanging updates to the Notifier thread

As both acquire and release wakelock calls are being notified to the BatteryStats via the notifier thread, wakelock changes should also follow the same pattern.

Bug: 354303438
Flag: com.android.server.power.feature.flags.improve_wakelock_latency
Test: atest NotifierTest
Change-Id: Ie383e485482a4d09c4570469e5b74ef10ea08693

2 files changed, 61 insertions, 11 deletions

diff --git a/services/core/java/com/android/server/power/Notifier.java b/services/core/java/com/android/server/power/Notifier.java
index aa5f5a24f179..b28da55b5196 100644
--- a/services/core/java/com/android/server/power/Notifier.java
+++ b/services/core/java/com/android/server/power/Notifier.java
@@ -375,9 +375,9 @@ public class Notifier {
             final boolean unimportantForLogging = newOwnerUid == Process.SYSTEM_UID
                     && (newFlags & PowerManager.UNIMPORTANT_FOR_LOGGING) != 0;
             try {
-                mBatteryStats.noteChangeWakelockFromSource(workSource, ownerPid, tag, historyTag,
-                        monitorType, newWorkSource, newOwnerPid, newTag, newHistoryTag,
-                        newMonitorType, unimportantForLogging);
+                notifyWakelockChanging(workSource, ownerPid, tag,
+                            historyTag, monitorType, newWorkSource, newOwnerPid, newTag,
+                            newHistoryTag, newMonitorType, unimportantForLogging);
             } catch (RemoteException ex) {
                 // Ignore
             }
@@ -1127,6 +1127,29 @@ public class Notifier {
         mWakeLockLog.onWakeLockReleased(tag, ownerUid, currentTime);
     }
 
+    @SuppressLint("AndroidFrameworkRequiresPermission")
+    private void notifyWakelockChanging(WorkSource workSource, int ownerPid, String tag,
+            String historyTag, int monitorType, WorkSource newWorkSource, int newOwnerPid,
+            String newTag, String newHistoryTag, int newMonitorType, boolean unimportantForLogging)
+            throws RemoteException {
+        if (!mFlags.improveWakelockLatency()) {
+            mBatteryStats.noteChangeWakelockFromSource(workSource, ownerPid, tag,
+                    historyTag, monitorType, newWorkSource, newOwnerPid, newTag,
+                    newHistoryTag, newMonitorType, unimportantForLogging);
+        } else {
+            mHandler.post(() -> {
+                try {
+                    mBatteryStats.noteChangeWakelockFromSource(workSource, ownerPid, tag,
+                            historyTag, monitorType, newWorkSource, newOwnerPid, newTag,
+                            newHistoryTag, newMonitorType, unimportantForLogging);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Failed to notify the wakelock changing from source via "
+                            + "Notifier." + e.getLocalizedMessage());
+                }
+            });
+        }
+    }
+
     private final class NotifierHandler extends Handler {
 
         public NotifierHandler(Looper looper) {
diff --git a/services/tests/powerservicetests/src/com/android/server/power/NotifierTest.java b/services/tests/powerservicetests/src/com/android/server/power/NotifierTest.java
index ce2bb95b790a..d45e31248d0b 100644
--- a/services/tests/powerservicetests/src/com/android/server/power/NotifierTest.java
+++ b/services/tests/powerservicetests/src/com/android/server/power/NotifierTest.java
@@ -44,6 +44,7 @@ import android.os.PowerManager;
 import android.os.RemoteException;
 import android.os.VibrationAttributes;
 import android.os.Vibrator;
+import android.os.WorkSource;
 import android.os.test.TestLooper;
 import android.provider.Settings;
 import android.testing.TestableContext;
@@ -60,6 +61,7 @@ import com.android.server.statusbar.StatusBarManagerInternal;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 import java.util.concurrent.Executor;
@@ -231,7 +233,7 @@ public class NotifierTest {
     }
 
     @Test
-    public void testOnWakeLockListener_RemoteException_NoRethrow() {
+    public void testOnWakeLockListener_RemoteException_NoRethrow() throws RemoteException {
         when(mPowerManagerFlags.improveWakelockLatency()).thenReturn(true);
         createNotifier();
         clearInvocations(mWakeLockLog, mBatteryStats, mAppOpsManager);
@@ -249,33 +251,58 @@ public class NotifierTest {
         verifyZeroInteractions(mWakeLockLog);
         mTestLooper.dispatchAll();
         verify(mWakeLockLog).onWakeLockReleased("wakelockTag", uid, 1);
-
+        clearInvocations(mBatteryStats);
         mNotifier.onWakeLockAcquired(PowerManager.PARTIAL_WAKE_LOCK, "wakelockTag",
                 "my.package.name", uid, pid, /* workSource= */ null, /* historyTag= */ null,
                 exceptingCallback);
+
+        verifyNoMoreInteractions(mWakeLockLog, mBatteryStats);
+        mTestLooper.dispatchAll();
+        verify(mWakeLockLog).onWakeLockAcquired("wakelockTag", uid,
+                PowerManager.PARTIAL_WAKE_LOCK, 1);
+        verify(mBatteryStats).noteStartWakelock(uid, pid, "wakelockTag", /* historyTag= */ null,
+                BatteryStats.WAKE_TYPE_PARTIAL, false);
+
+        verifyNoMoreInteractions(mWakeLockLog, mBatteryStats);
+        WorkSource worksourceOld = Mockito.mock(WorkSource.class);
+        WorkSource worksourceNew = Mockito.mock(WorkSource.class);
+
         mNotifier.onWakeLockChanging(PowerManager.PARTIAL_WAKE_LOCK, "wakelockTag",
-                "my.package.name", uid, pid, /* workSource= */ null, /* historyTag= */ null,
+                "my.package.name", uid, pid, worksourceOld, /* historyTag= */ null,
                 exceptingCallback,
                 PowerManager.SCREEN_BRIGHT_WAKE_LOCK, "wakelockTag",
-                "my.package.name", uid, pid, /* newWorkSource= */ null, /* newHistoryTag= */ null,
+                "my.package.name", uid, pid, worksourceNew, /* newHistoryTag= */ null,
                 exceptingCallback);
-        verifyNoMoreInteractions(mWakeLockLog);
         mTestLooper.dispatchAll();
-        verify(mWakeLockLog).onWakeLockAcquired("wakelockTag", uid,
-                PowerManager.PARTIAL_WAKE_LOCK, 1);
+        verify(mBatteryStats).noteChangeWakelockFromSource(worksourceOld, pid, "wakelockTag",
+                null, BatteryStats.WAKE_TYPE_PARTIAL, worksourceNew, pid, "wakelockTag",
+                null, BatteryStats.WAKE_TYPE_FULL, false);
         // If we didn't throw, we're good!
 
         // Test with improveWakelockLatency flag false, hence the wakelock log will run on the same
         // thread
-        clearInvocations(mWakeLockLog);
+        clearInvocations(mWakeLockLog, mBatteryStats);
         when(mPowerManagerFlags.improveWakelockLatency()).thenReturn(false);
 
+        // Acquire the wakelock
         mNotifier.onWakeLockAcquired(PowerManager.PARTIAL_WAKE_LOCK, "wakelockTag",
                 "my.package.name", uid, pid, /* workSource= */ null, /* historyTag= */ null,
                 exceptingCallback);
         verify(mWakeLockLog).onWakeLockAcquired("wakelockTag", uid,
                 PowerManager.PARTIAL_WAKE_LOCK, -1);
 
+        // Update the wakelock
+        mNotifier.onWakeLockChanging(PowerManager.PARTIAL_WAKE_LOCK, "wakelockTag",
+                "my.package.name", uid, pid, worksourceOld, /* historyTag= */ null,
+                exceptingCallback,
+                PowerManager.SCREEN_BRIGHT_WAKE_LOCK, "wakelockTag",
+                "my.package.name", uid, pid, worksourceNew, /* newHistoryTag= */ null,
+                exceptingCallback);
+        verify(mBatteryStats).noteChangeWakelockFromSource(worksourceOld, pid, "wakelockTag",
+                null, BatteryStats.WAKE_TYPE_PARTIAL, worksourceNew, pid, "wakelockTag",
+                null, BatteryStats.WAKE_TYPE_FULL, false);
+
+        // Release the wakelock
         mNotifier.onWakeLockReleased(PowerManager.PARTIAL_WAKE_LOCK, "wakelockTag",
                 "my.package.name", uid, pid, /* workSource= */ null, /* historyTag= */ null,
                 exceptingCallback);