Update workchain when wakelock uids updated

There are currently 2 ways to update the worksource with the caller uids
for appropriate attribution - Workchain and uids array in the WorkSource
object. The array approach has been deprecated, and is a hidden utility.
Workchain is the recommended way to do the attribution. A special hack
was created for Audio wakelocks to create wakelocks in async, and then
update the corresponding uids using updateWakelockUids API. This API
would then update the uids array inside the Worksource, and hence the
changing uids never get logged in the form of atoms.

Bug: 331304805
Flag: com.android.server.power.feature.flags.wakelock_attribution_via_workchain
Test: atest PowerManagerServiceTest

Change-Id: I4d1b8559c4293c35be5ba8d30957b5b48311c1ab

6 files changed, 95 insertions, 4 deletions

diff --git a/AconfigFlags.bp b/AconfigFlags.bp
index 834398e5c2c2..ac756ea1d624 100644
--- a/AconfigFlags.bp
+++ b/AconfigFlags.bp
@@ -1573,6 +1573,13 @@ java_aconfig_library {
     defaults: ["framework-minus-apex-aconfig-java-defaults"],
 }
 
+java_aconfig_library {
+    name: "power_flags_lib_host",
+    aconfig_declarations: "power_flags",
+    host_supported: true,
+    defaults: ["framework-minus-apex-aconfig-java-defaults"],
+}
+
 // Content
 aconfig_declarations {
     name: "android.content.flags-aconfig",
diff --git a/core/java/android/os/PowerManager.java b/core/java/android/os/PowerManager.java
index 2a5666cbe83c..764a029611f9 100644
--- a/core/java/android/os/PowerManager.java
+++ b/core/java/android/os/PowerManager.java
@@ -4205,6 +4205,17 @@ public final class PowerManager {
             else mFlags &= ~UNIMPORTANT_FOR_LOGGING;
         }
 
+        /** @hide */
+        public void updateUids(int[] uids) {
+            synchronized (mToken) {
+                try {
+                    mService.updateWakeLockUids(mToken, uids);
+                } catch (RemoteException e) {
+                    throw e.rethrowFromSystemServer();
+                }
+            }
+        }
+
         @Override
         public String toString() {
             synchronized (mToken) {
diff --git a/services/core/java/com/android/server/power/PowerManagerService.java b/services/core/java/com/android/server/power/PowerManagerService.java
index 090707db50a5..8fae875eb29b 100644
--- a/services/core/java/com/android/server/power/PowerManagerService.java
+++ b/services/core/java/com/android/server/power/PowerManagerService.java
@@ -5979,10 +5979,19 @@ public final class PowerManagerService extends SystemService
 
             if (uids != null) {
                 ws = new WorkSource();
-                // XXX should WorkSource have a way to set uids as an int[] instead of adding them
-                // one at a time?
-                for (int uid : uids) {
-                    ws.add(uid);
+                if (mFeatureFlags.isWakelockAttributionViaWorkchainEnabled()) {
+                    int callingUid = Binder.getCallingUid();
+                    for (int uid : uids) {
+                        WorkChain workChain = ws.createWorkChain();
+                        workChain.addNode(uid, null);
+                        workChain.addNode(callingUid, null);
+                    }
+                } else {
+                    // XXX should WorkSource have a way to set uids as an int[] instead of
+                    // adding them one at a time?
+                    for (int uid : uids) {
+                        ws.add(uid);
+                    }
                 }
             }
             updateWakeLockWorkSource(lock, ws, null);
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 42b44013bea2..ebc50fd85f24 100644
--- a/services/core/java/com/android/server/power/feature/PowerManagerFlags.java
+++ b/services/core/java/com/android/server/power/feature/PowerManagerFlags.java
@@ -63,6 +63,10 @@ public class PowerManagerFlags {
     private final FlagState mMoveWscLoggingToNotifier =
             new FlagState(Flags.FLAG_MOVE_WSC_LOGGING_TO_NOTIFIER, Flags::moveWscLoggingToNotifier);
 
+    private final FlagState mWakelockAttributionViaWorkchain =
+            new FlagState(Flags.FLAG_WAKELOCK_ATTRIBUTION_VIA_WORKCHAIN,
+                    Flags::wakelockAttributionViaWorkchain);
+
     /** Returns whether early-screen-timeout-detector is enabled on not. */
     public boolean isEarlyScreenTimeoutDetectorEnabled() {
         return mEarlyScreenTimeoutDetectorFlagState.isEnabled();
@@ -110,6 +114,13 @@ public class PowerManagerFlags {
     }
 
     /**
+     * @return Whether the wakelock attribution via workchain is enabled
+     */
+    public boolean isWakelockAttributionViaWorkchainEnabled() {
+        return mWakelockAttributionViaWorkchain.isEnabled();
+    }
+
+    /**
      * dumps all flagstates
      * @param pw printWriter
      */
@@ -120,6 +131,7 @@ public class PowerManagerFlags {
         pw.println(" " + mPerDisplayWakeByTouch);
         pw.println(" " + mFrameworkWakelockInfo);
         pw.println(" " + mMoveWscLoggingToNotifier);
+        pw.println(" " + mWakelockAttributionViaWorkchain);
     }
 
     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 613daf820e34..fefe195dc337 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
@@ -23,6 +23,17 @@ flag {
 }
 
 flag {
+    name: "wakelock_attribution_via_workchain"
+    namespace: "power"
+    description: "Enables the attribution of wakelocks via WorkChain for updateWakelockUids"
+    bug: "331304805"
+    is_fixed_read_only: true
+    metadata {
+      purpose: PURPOSE_BUGFIX
+    }
+}
+
+flag {
     name: "improve_wakelock_latency"
     namespace: "power"
     description: "Feature flag for tracking the optimizations to improve the latency of acquiring and releasing a wakelock."
diff --git a/services/tests/powerservicetests/src/com/android/server/power/PowerManagerServiceTest.java b/services/tests/powerservicetests/src/com/android/server/power/PowerManagerServiceTest.java
index 6b138b986fe7..29a17e1c85ab 100644
--- a/services/tests/powerservicetests/src/com/android/server/power/PowerManagerServiceTest.java
+++ b/services/tests/powerservicetests/src/com/android/server/power/PowerManagerServiceTest.java
@@ -89,9 +89,13 @@ import android.os.PowerManager;
 import android.os.PowerManagerInternal;
 import android.os.PowerSaveState;
 import android.os.UserHandle;
+import android.os.WorkSource;
 import android.os.test.TestLooper;
 import android.platform.test.annotations.DisableFlags;
 import android.platform.test.annotations.EnableFlags;
+import android.platform.test.annotations.RequiresFlagsEnabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.DeviceConfig;
 import android.provider.Settings;
@@ -200,6 +204,9 @@ public class PowerManagerServiceTest {
 
     @Rule public SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+
     private PowerManagerService mService;
     private ContextWrapper mContextSpy;
     private BatteryReceiver mBatteryReceiver;
@@ -3045,6 +3052,40 @@ public class PowerManagerServiceTest {
     }
 
     /**
+     * Test IPowerManager.updateWakeLockUids() updates the workchain with the new uids
+     */
+    @Test
+    @RequiresFlagsEnabled({Flags.FLAG_WAKELOCK_ATTRIBUTION_VIA_WORKCHAIN})
+    public void test_updateWakelockUids_updatesWorkchain() {
+        createService();
+        startSystem();
+        final String tag = "wakelock1";
+        final String packageName = "pkg.name";
+        final IBinder token = new Binder();
+        int flags = PowerManager.PARTIAL_WAKE_LOCK;
+        final IWakeLockCallback callback1 = Mockito.mock(IWakeLockCallback.class);
+        final IBinder callbackBinder1 = Mockito.mock(Binder.class);
+        when(callback1.asBinder()).thenReturn(callbackBinder1);
+        WorkSource oldWorksource = new WorkSource();
+        oldWorksource.createWorkChain().addNode(1000, null);
+        mService.getBinderServiceInstance().acquireWakeLock(token, flags, tag, packageName,
+                oldWorksource, null /* historyTag */, Display.INVALID_DISPLAY, callback1);
+        verify(mNotifierMock).onWakeLockAcquired(anyInt(), eq(tag), eq(packageName),
+                anyInt(), anyInt(), eq(oldWorksource), any(), same(callback1));
+
+        WorkSource newWorksource = new WorkSource();
+        newWorksource.createWorkChain().addNode(1011, null)
+                .addNode(Binder.getCallingUid(), null);
+        newWorksource.createWorkChain().addNode(1012, null)
+                .addNode(Binder.getCallingUid(), null);
+        mService.getBinderServiceInstance().updateWakeLockUids(token, new int[]{1011, 1012});
+        verify(mNotifierMock).onWakeLockChanging(anyInt(), eq(tag), eq(packageName),
+                anyInt(), anyInt(), eq(oldWorksource), any(), any(),
+                anyInt(), eq(tag), eq(packageName), anyInt(), anyInt(), eq(newWorksource), any(),
+                any());
+    }
+
+    /**
      * Test IPowerManager.updateWakeLockCallback() with a new IWakeLockCallback.
      */
     @Test