Merge "Ref count isolated uid usage in BatteryStats"
diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index 4043060..719dc53 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -661,6 +661,10 @@
      * Mapping isolated uids to the actual owning app uid.
      */
     final SparseIntArray mIsolatedUids = new SparseIntArray();
+    /**
+     * Internal reference count of isolated uids.
+     */
+    final SparseIntArray mIsolatedUidRefCounts = new SparseIntArray();
 
     /**
      * The statistics we have collected organized by uids.
@@ -3855,6 +3859,7 @@
     public void addIsolatedUidLocked(int isolatedUid, int appUid,
             long elapsedRealtimeMs, long uptimeMs) {
         mIsolatedUids.put(isolatedUid, appUid);
+        mIsolatedUidRefCounts.put(isolatedUid, 1);
         final Uid u = getUidStatsLocked(appUid, elapsedRealtimeMs, uptimeMs);
         u.addIsolatedUid(isolatedUid);
     }
@@ -3873,19 +3878,51 @@
     }
 
     /**
-     * This should only be called after the cpu times have been read.
+     * Isolated uid should only be removed after all wakelocks associated with the uid are stopped
+     * and the cpu time-in-state has been read one last time for the uid.
+     *
      * @see #scheduleRemoveIsolatedUidLocked(int, int)
+     *
+     * @return true if the isolated uid is actually removed.
      */
     @GuardedBy("this")
-    public void removeIsolatedUidLocked(int isolatedUid, long elapsedRealtimeMs, long uptimeMs) {
+    public boolean maybeRemoveIsolatedUidLocked(int isolatedUid, long elapsedRealtimeMs,
+            long uptimeMs) {
+        final int refCount = mIsolatedUidRefCounts.get(isolatedUid, 0) - 1;
+        if (refCount > 0) {
+            // Isolated uid is still being tracked
+            mIsolatedUidRefCounts.put(isolatedUid, refCount);
+            return false;
+        }
+
         final int idx = mIsolatedUids.indexOfKey(isolatedUid);
         if (idx >= 0) {
             final int ownerUid = mIsolatedUids.valueAt(idx);
             final Uid u = getUidStatsLocked(ownerUid, elapsedRealtimeMs, uptimeMs);
             u.removeIsolatedUid(isolatedUid);
             mIsolatedUids.removeAt(idx);
+            mIsolatedUidRefCounts.delete(isolatedUid);
+        } else {
+            Slog.w(TAG, "Attempted to remove untracked isolated uid (" + isolatedUid + ")");
         }
         mPendingRemovedUids.add(new UidToRemove(isolatedUid, elapsedRealtimeMs));
+
+        return true;
+    }
+
+    /**
+     * Increment the ref count for an isolated uid.
+     * call #maybeRemoveIsolatedUidLocked to decrement.
+     */
+    public void incrementIsolatedUidRefCount(int uid) {
+        final int refCount = mIsolatedUidRefCounts.get(uid, 0);
+        if (refCount <= 0) {
+            // Uid is not mapped or referenced
+            Slog.w(TAG,
+                    "Attempted to increment ref counted of untracked isolated uid (" + uid + ")");
+            return;
+        }
+        mIsolatedUidRefCounts.put(uid, refCount + 1);
     }
 
     public int mapUid(int uid) {
@@ -4245,7 +4282,7 @@
 
     public void noteStartWakeLocked(int uid, int pid, WorkChain wc, String name, String historyName,
             int type, boolean unimportantForLogging, long elapsedRealtimeMs, long uptimeMs) {
-        uid = mapUid(uid);
+        final int mappedUid = mapUid(uid);
         if (type == WAKE_TYPE_PARTIAL) {
             // Only care about partial wake locks, since full wake locks
             // will be canceled when the user puts the screen to sleep.
@@ -4255,9 +4292,9 @@
             }
             if (mRecordAllHistory) {
                 if (mActiveEvents.updateState(HistoryItem.EVENT_WAKE_LOCK_START, historyName,
-                        uid, 0)) {
+                        mappedUid, 0)) {
                     addHistoryEventLocked(elapsedRealtimeMs, uptimeMs,
-                            HistoryItem.EVENT_WAKE_LOCK_START, historyName, uid);
+                            HistoryItem.EVENT_WAKE_LOCK_START, historyName, mappedUid);
                 }
             }
             if (mWakeLockNesting == 0) {
@@ -4266,7 +4303,7 @@
                         + Integer.toHexString(mHistoryCur.states));
                 mHistoryCur.wakelockTag = mHistoryCur.localWakelockTag;
                 mHistoryCur.wakelockTag.string = mInitialAcquireWakeName = historyName;
-                mHistoryCur.wakelockTag.uid = mInitialAcquireWakeUid = uid;
+                mHistoryCur.wakelockTag.uid = mInitialAcquireWakeUid = mappedUid;
                 mWakeLockImportant = !unimportantForLogging;
                 addHistoryRecordLocked(elapsedRealtimeMs, uptimeMs);
             } else if (!mWakeLockImportant && !unimportantForLogging
@@ -4276,14 +4313,19 @@
                     mHistoryLastWritten.wakelockTag = null;
                     mHistoryCur.wakelockTag = mHistoryCur.localWakelockTag;
                     mHistoryCur.wakelockTag.string = mInitialAcquireWakeName = historyName;
-                    mHistoryCur.wakelockTag.uid = mInitialAcquireWakeUid = uid;
+                    mHistoryCur.wakelockTag.uid = mInitialAcquireWakeUid = mappedUid;
                     addHistoryRecordLocked(elapsedRealtimeMs, uptimeMs);
                 }
                 mWakeLockImportant = true;
             }
             mWakeLockNesting++;
         }
-        if (uid >= 0) {
+        if (mappedUid >= 0) {
+            if (mappedUid != uid) {
+                // Prevent the isolated uid mapping from being removed while the wakelock is
+                // being held.
+                incrementIsolatedUidRefCount(uid);
+            }
             if (mOnBatteryScreenOffTimeBase.isRunning()) {
                 // We only update the cpu time when a wake lock is acquired if the screen is off.
                 // If the screen is on, we don't distribute the power amongst partial wakelocks.
@@ -4293,7 +4335,7 @@
                 requestWakelockCpuUpdate();
             }
 
-            getUidStatsLocked(uid, elapsedRealtimeMs, uptimeMs)
+            getUidStatsLocked(mappedUid, elapsedRealtimeMs, uptimeMs)
                     .noteStartWakeLocked(pid, name, type, elapsedRealtimeMs);
 
             if (wc != null) {
@@ -4301,8 +4343,8 @@
                         wc.getTags(), getPowerManagerWakeLockLevel(type), name,
                         FrameworkStatsLog.WAKELOCK_STATE_CHANGED__STATE__ACQUIRE);
             } else {
-                FrameworkStatsLog.write_non_chained(FrameworkStatsLog.WAKELOCK_STATE_CHANGED, uid,
-                        null, getPowerManagerWakeLockLevel(type), name,
+                FrameworkStatsLog.write_non_chained(FrameworkStatsLog.WAKELOCK_STATE_CHANGED,
+                        mappedUid, null, getPowerManagerWakeLockLevel(type), name,
                         FrameworkStatsLog.WAKELOCK_STATE_CHANGED__STATE__ACQUIRE);
             }
         }
@@ -4316,7 +4358,7 @@
 
     public void noteStopWakeLocked(int uid, int pid, WorkChain wc, String name, String historyName,
             int type, long elapsedRealtimeMs, long uptimeMs) {
-        uid = mapUid(uid);
+        final int mappedUid = mapUid(uid);
         if (type == WAKE_TYPE_PARTIAL) {
             mWakeLockNesting--;
             if (mRecordAllHistory) {
@@ -4324,9 +4366,9 @@
                     historyName = name;
                 }
                 if (mActiveEvents.updateState(HistoryItem.EVENT_WAKE_LOCK_FINISH, historyName,
-                        uid, 0)) {
+                        mappedUid, 0)) {
                     addHistoryEventLocked(elapsedRealtimeMs, uptimeMs,
-                            HistoryItem.EVENT_WAKE_LOCK_FINISH, historyName, uid);
+                            HistoryItem.EVENT_WAKE_LOCK_FINISH, historyName, mappedUid);
                 }
             }
             if (mWakeLockNesting == 0) {
@@ -4338,7 +4380,7 @@
                 addHistoryRecordLocked(elapsedRealtimeMs, uptimeMs);
             }
         }
-        if (uid >= 0) {
+        if (mappedUid >= 0) {
             if (mOnBatteryScreenOffTimeBase.isRunning()) {
                 if (DEBUG_ENERGY_CPU) {
                     Slog.d(TAG, "Updating cpu time because of -wake_lock");
@@ -4346,17 +4388,22 @@
                 requestWakelockCpuUpdate();
             }
 
-            getUidStatsLocked(uid, elapsedRealtimeMs, uptimeMs)
+            getUidStatsLocked(mappedUid, elapsedRealtimeMs, uptimeMs)
                     .noteStopWakeLocked(pid, name, type, elapsedRealtimeMs);
             if (wc != null) {
                 FrameworkStatsLog.write(FrameworkStatsLog.WAKELOCK_STATE_CHANGED, wc.getUids(),
                         wc.getTags(), getPowerManagerWakeLockLevel(type), name,
                         FrameworkStatsLog.WAKELOCK_STATE_CHANGED__STATE__RELEASE);
             } else {
-                FrameworkStatsLog.write_non_chained(FrameworkStatsLog.WAKELOCK_STATE_CHANGED, uid,
-                        null, getPowerManagerWakeLockLevel(type), name,
+                FrameworkStatsLog.write_non_chained(FrameworkStatsLog.WAKELOCK_STATE_CHANGED,
+                        mappedUid, null, getPowerManagerWakeLockLevel(type), name,
                         FrameworkStatsLog.WAKELOCK_STATE_CHANGED__STATE__RELEASE);
             }
+
+            if (mappedUid != uid) {
+                // Decrement the ref count for the isolated uid and delete the mapping if uneeded.
+                maybeRemoveIsolatedUidLocked(uid, elapsedRealtimeMs, uptimeMs);
+            }
         }
     }
 
@@ -16769,6 +16816,15 @@
         pw.print("UIDs removed since the later of device start or stats reset: ");
         pw.println(mNumUidsRemoved);
 
+        pw.println("Currently mapped isolated uids:");
+        final int numIsolatedUids = mIsolatedUids.size();
+        for (int i = 0; i < numIsolatedUids; i++) {
+            final int isolatedUid = mIsolatedUids.keyAt(i);
+            final int ownerUid = mIsolatedUids.valueAt(i);
+            final int refCount = mIsolatedUidRefCounts.get(isolatedUid);
+            pw.println("  " + isolatedUid + "->" + ownerUid + " (ref count = " + refCount + ")");
+        }
+
         pw.println();
         dumpConstantsLocked(pw);
 
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 a43d32d..e8e4330 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
@@ -23,6 +23,8 @@
 import android.os.BatteryStats;
 import android.os.BatteryStats.HistoryItem;
 import android.os.BatteryStats.Uid.Sensor;
+import android.os.Process;
+import android.os.UserHandle;
 import android.os.WorkSource;
 import android.util.SparseLongArray;
 import android.view.Display;
@@ -53,6 +55,8 @@
 public class BatteryStatsNoteTest extends TestCase {
 
     private static final int UID = 10500;
+    private static final int ISOLATED_APP_ID = Process.FIRST_ISOLATED_UID + 23;
+    private static final int ISOLATED_UID = UserHandle.getUid(0, ISOLATED_APP_ID);
     private static final WorkSource WS = new WorkSource(UID);
 
     /**
@@ -114,6 +118,88 @@
         assertEquals(120_000, bgTime);
     }
 
+    /**
+     * Test BatteryStatsImpl.Uid.noteStartWakeLocked for an isolated uid.
+     */
+    @SmallTest
+    public void testNoteStartWakeLocked_isolatedUid() throws Exception {
+        final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+
+        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.noteStartWakeLocked(ISOLATED_UID, pid, isolatedWorkChain, name, historyName,
+                WAKE_TYPE_PARTIAL, false);
+
+        clocks.realtime = clocks.uptime = 100;
+        bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
+
+        clocks.realtime = clocks.uptime = 220;
+        bi.noteStopWakeLocked(ISOLATED_UID, pid, isolatedWorkChain, name, historyName,
+                WAKE_TYPE_PARTIAL);
+
+        // ISOLATED_UID wakelock time should be attributed to UID.
+        BatteryStats.Timer aggregTimer = bi.getUidStats().get(UID)
+                .getAggregatedPartialWakelockTimer();
+        long actualTime = aggregTimer.getTotalTimeLocked(300_000, STATS_SINCE_CHARGED);
+        long bgTime = aggregTimer.getSubTimer().getTotalTimeLocked(300_000, STATS_SINCE_CHARGED);
+        assertEquals(220_000, actualTime);
+        assertEquals(120_000, bgTime);
+    }
+
+    /**
+     * Test BatteryStatsImpl.Uid.noteStartWakeLocked for an isolated uid, with a race where the
+     * isolated uid is removed from batterystats before the wakelock has been stopped.
+     */
+    @SmallTest
+    public void testNoteStartWakeLocked_isolatedUidRace() throws Exception {
+        final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+
+        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.noteStartWakeLocked(ISOLATED_UID, pid, isolatedWorkChain, name, historyName,
+                WAKE_TYPE_PARTIAL, false);
+
+        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.noteStopWakeLocked(ISOLATED_UID, pid, isolatedWorkChain, name, historyName,
+                WAKE_TYPE_PARTIAL);
+
+        // ISOLATED_UID wakelock time should be attributed to UID.
+        BatteryStats.Timer aggregTimer = bi.getUidStats().get(UID)
+                .getAggregatedPartialWakelockTimer();
+        long actualTime = aggregTimer.getTotalTimeLocked(300_000, STATS_SINCE_CHARGED);
+        long bgTime = aggregTimer.getSubTimer().getTotalTimeLocked(300_000, STATS_SINCE_CHARGED);
+        assertEquals(220_000, actualTime);
+        assertEquals(120_000, bgTime);
+    }
+
 
     /**
      * Test BatteryStatsImpl.noteUidProcessStateLocked.
diff --git a/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java b/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
index 449f02e..3c6b682 100644
--- a/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
+++ b/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
@@ -492,7 +492,7 @@
                     for (int uid : uidsToRemove) {
                         FrameworkStatsLog.write(FrameworkStatsLog.ISOLATED_UID_CHANGED, -1, uid,
                                 FrameworkStatsLog.ISOLATED_UID_CHANGED__EVENT__REMOVED);
-                        mStats.removeIsolatedUidLocked(uid, SystemClock.elapsedRealtime(),
+                        mStats.maybeRemoveIsolatedUidLocked(uid, SystemClock.elapsedRealtime(),
                                 SystemClock.uptimeMillis());
                     }
                     mStats.clearPendingRemovedUidsLocked();