Ref count isolated uid use for long partial wakelocks

The isolated uids need to be kept in the map while it is holding a long
partial wakelock.

Bug: 218298973
Test: atest BatteryStatsNoteTest
Change-Id: I86456137e6d6ad47f9974da281defc04d20fb32f
(cherry picked from commit fd48aa6fa6a87e0e159aca7ca8a429c47905291b)
Merged-In: I86456137e6d6ad47f9974da281defc04d20fb32f

2 files changed, 129 insertions, 8 deletions

diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index 3b46f627265b..fe2f80cf72df 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -5429,7 +5429,6 @@ public class BatteryStatsImpl extends BatteryStats {
     @GuardedBy("this")
     public void noteLongPartialWakelockStart(String name, String historyName, int uid,
             long elapsedRealtimeMs, long uptimeMs) {
-        uid = mapUid(uid);
         noteLongPartialWakeLockStartInternal(name, historyName, uid, elapsedRealtimeMs, uptimeMs);
     }
 
@@ -5464,15 +5463,21 @@ public class BatteryStatsImpl extends BatteryStats {
     @GuardedBy("this")
     private void noteLongPartialWakeLockStartInternal(String name, String historyName, int uid,
             long elapsedRealtimeMs, long uptimeMs) {
+        final int mappedUid = mapUid(uid);
         if (historyName == null) {
             historyName = name;
         }
-        if (!mActiveEvents.updateState(HistoryItem.EVENT_LONG_WAKE_LOCK_START, historyName, uid,
-                0)) {
+        if (!mActiveEvents.updateState(HistoryItem.EVENT_LONG_WAKE_LOCK_START, historyName,
+                mappedUid, 0)) {
             return;
         }
         addHistoryEventLocked(elapsedRealtimeMs, uptimeMs, HistoryItem.EVENT_LONG_WAKE_LOCK_START,
-                historyName, uid);
+                historyName, mappedUid);
+        if (mappedUid != uid) {
+            // Prevent the isolated uid mapping from being removed while the wakelock is
+            // being held.
+            incrementIsolatedUidRefCount(uid);
+        }
     }
 
     @GuardedBy("this")
@@ -5484,7 +5489,6 @@ public class BatteryStatsImpl extends BatteryStats {
     @GuardedBy("this")
     public void noteLongPartialWakelockFinish(String name, String historyName, int uid,
             long elapsedRealtimeMs, long uptimeMs) {
-        uid = mapUid(uid);
         noteLongPartialWakeLockFinishInternal(name, historyName, uid, elapsedRealtimeMs, uptimeMs);
     }
 
@@ -5519,15 +5523,20 @@ public class BatteryStatsImpl extends BatteryStats {
     @GuardedBy("this")
     private void noteLongPartialWakeLockFinishInternal(String name, String historyName, int uid,
             long elapsedRealtimeMs, long uptimeMs) {
+        final int mappedUid = mapUid(uid);
         if (historyName == null) {
             historyName = name;
         }
-        if (!mActiveEvents.updateState(HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH, historyName, uid,
-                0)) {
+        if (!mActiveEvents.updateState(HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH, historyName,
+                mappedUid, 0)) {
             return;
         }
         addHistoryEventLocked(elapsedRealtimeMs, uptimeMs, HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH,
-                historyName, uid);
+                historyName, mappedUid);
+        if (mappedUid != uid) {
+            // Decrement the ref count for the isolated uid and delete the mapping if uneeded.
+            maybeRemoveIsolatedUidLocked(uid, elapsedRealtimeMs, uptimeMs);
+        }
     }
 
     @GuardedBy("this")
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 7ccb9d963ee6..dbe1e8132e1d 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
@@ -26,6 +26,8 @@ import static android.os.BatteryStats.WAKE_TYPE_PARTIAL;
 import static com.android.internal.os.BatteryStatsImpl.ExternalStatsSync.UPDATE_CPU;
 import static com.android.internal.os.BatteryStatsImpl.ExternalStatsSync.UPDATE_DISPLAY;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.mockito.Mockito.mock;
 
 import android.app.ActivityManager;
@@ -228,6 +230,116 @@ public class BatteryStatsNoteTest extends TestCase {
         assertEquals(120_000, bgTime);
     }
 
+    /**
+     * Test BatteryStatsImpl.Uid.noteLongPartialWakelockStart for an isolated uid.
+     */
+    @SmallTest
+    public void testNoteLongPartialWakelockStart_isolatedUid() throws Exception {
+        final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+
+
+        bi.setRecordAllHistoryLocked(true);
+        bi.forceRecordAllHistory();
+
+        int pid = 10;
+        String name = "name";
+        String historyName = "historyName";
+
+        WorkSource.WorkChain isolatedWorkChain = new WorkSource.WorkChain();
+        isolatedWorkChain.addNode(ISOLATED_UID, name);
+
+        // Map ISOLATED_UID to UID.
+        bi.addIsolatedUidLocked(ISOLATED_UID, UID);
+
+        bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_TOP);
+        bi.noteLongPartialWakelockStart(name, historyName, ISOLATED_UID);
+
+        clocks.realtime = clocks.uptime = 100;
+        bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
+
+        clocks.realtime = clocks.uptime = 220;
+        bi.noteLongPartialWakelockFinish(name, historyName, ISOLATED_UID);
+
+        final BatteryStatsHistoryIterator iterator =
+                bi.createBatteryStatsHistoryIterator();
+
+        BatteryStats.HistoryItem item = new BatteryStats.HistoryItem();
+
+        while (iterator.next(item)) {
+            if (item.eventCode == HistoryItem.EVENT_LONG_WAKE_LOCK_START) break;
+        }
+        assertThat(item.eventCode).isEqualTo(HistoryItem.EVENT_LONG_WAKE_LOCK_START);
+        assertThat(item.eventTag).isNotNull();
+        assertThat(item.eventTag.string).isEqualTo(historyName);
+        assertThat(item.eventTag.uid).isEqualTo(UID);
+
+        while (iterator.next(item)) {
+            if (item.eventCode == HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH) break;
+        }
+        assertThat(item.eventCode).isEqualTo(HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH);
+        assertThat(item.eventTag).isNotNull();
+        assertThat(item.eventTag.string).isEqualTo(historyName);
+        assertThat(item.eventTag.uid).isEqualTo(UID);
+    }
+
+    /**
+     * Test BatteryStatsImpl.Uid.noteLongPartialWakelockStart for an isolated uid.
+     */
+    @SmallTest
+    public void testNoteLongPartialWakelockStart_isolatedUidRace() throws Exception {
+        final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+
+
+        bi.setRecordAllHistoryLocked(true);
+        bi.forceRecordAllHistory();
+
+        int pid = 10;
+        String name = "name";
+        String historyName = "historyName";
+
+        WorkSource.WorkChain isolatedWorkChain = new WorkSource.WorkChain();
+        isolatedWorkChain.addNode(ISOLATED_UID, name);
+
+        // Map ISOLATED_UID to UID.
+        bi.addIsolatedUidLocked(ISOLATED_UID, UID);
+
+        bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_TOP);
+        bi.noteLongPartialWakelockStart(name, historyName, ISOLATED_UID);
+
+        clocks.realtime = clocks.uptime = 100;
+        bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
+
+        clocks.realtime = clocks.uptime = 150;
+        bi.maybeRemoveIsolatedUidLocked(ISOLATED_UID, clocks.realtime, clocks.uptime);
+
+        clocks.realtime = clocks.uptime = 220;
+        bi.noteLongPartialWakelockFinish(name, historyName, ISOLATED_UID);
+
+        final BatteryStatsHistoryIterator iterator =
+                bi.createBatteryStatsHistoryIterator();
+
+        BatteryStats.HistoryItem item = new BatteryStats.HistoryItem();
+
+        while (iterator.next(item)) {
+            if (item.eventCode == HistoryItem.EVENT_LONG_WAKE_LOCK_START) break;
+        }
+        assertThat(item.eventCode).isEqualTo(HistoryItem.EVENT_LONG_WAKE_LOCK_START);
+        assertThat(item.eventTag).isNotNull();
+        assertThat(item.eventTag.string).isEqualTo(historyName);
+        assertThat(item.eventTag.uid).isEqualTo(UID);
+
+        while (iterator.next(item)) {
+            if (item.eventCode == HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH) break;
+        }
+        assertThat(item.eventCode).isEqualTo(HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH);
+        assertThat(item.eventTag).isNotNull();
+        assertThat(item.eventTag.string).isEqualTo(historyName);
+        assertThat(item.eventTag.uid).isEqualTo(UID);
+    }
 
     /**
      * Test BatteryStatsImpl.noteUidProcessStateLocked.