Merge "Use LongArrayMultiStateCounter for proc-in-state time counting"
diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index 7c4de82..94a1efe 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -77,7 +77,6 @@
 import android.util.ArraySet;
 import android.util.AtomicFile;
 import android.util.IndentingPrintWriter;
-import android.util.IntArray;
 import android.util.KeyValueListParser;
 import android.util.Log;
 import android.util.LongSparseArray;
@@ -161,7 +160,7 @@
     private static final int MAGIC = 0xBA757475; // 'BATSTATS'
 
     // Current on-disk Parcel version
-    static final int VERSION = 200;
+    static final int VERSION = 201;
 
     // The maximum number of names wakelocks we will keep track of
     // per uid; once the limit is reached, we batch the remaining wakelocks
@@ -464,10 +463,12 @@
             uidStates = mPendingUids.clone();
             mPendingUids.clear();
         }
+        final long timestampMs = mClock.elapsedRealtime();
         for (int i = uidStates.size() - 1; i >= 0; --i) {
             final int uid = uidStates.keyAt(i);
             final int procState = uidStates.valueAt(i);
             final int[] isolatedUids;
+            final LongArrayMultiStateCounter[] isolatedUidTimeInFreqCounters;
             final Uid u;
             synchronized (BatteryStatsImpl.this) {
                 // It's possible that uid no longer exists and any internal references have
@@ -479,26 +480,44 @@
                 }
                 if (u.mChildUids == null) {
                     isolatedUids = null;
+                    isolatedUidTimeInFreqCounters = null;
                 } else {
-                    isolatedUids = u.mChildUids.toArray();
-                    for (int j = isolatedUids.length - 1; j >= 0; --j) {
-                        isolatedUids[j] = u.mChildUids.get(j);
+                    int childUidCount = u.mChildUids.size();
+                    isolatedUids = new int[childUidCount];
+                    isolatedUidTimeInFreqCounters = new LongArrayMultiStateCounter[childUidCount];
+                    for (int j = childUidCount - 1; j >= 0; --j) {
+                        isolatedUids[j] = u.mChildUids.keyAt(j);
+                        isolatedUidTimeInFreqCounters[j] = u.mChildUids.valueAt(j);
                     }
                 }
             }
-            long[] cpuTimesMs = mKernelSingleUidTimeReader.readDeltaMs(uid);
+
+            LongArrayMultiStateCounter onBatteryCounter =
+                    u.getProcStateTimeCounter().getCounter();
+            LongArrayMultiStateCounter onBatteryScreenOffCounter =
+                    u.getProcStateScreenOffTimeCounter().getCounter();
+
+            onBatteryCounter.setState(procState, timestampMs);
+            mKernelSingleUidTimeReader.addDelta(uid, onBatteryCounter, timestampMs);
+
+            onBatteryScreenOffCounter.setState(procState, timestampMs);
+            mKernelSingleUidTimeReader.addDelta(uid, onBatteryScreenOffCounter, timestampMs);
+
             if (isolatedUids != null) {
+                LongArrayMultiStateCounter.LongArrayContainer deltaContainer =
+                        new LongArrayMultiStateCounter.LongArrayContainer(getCpuFreqCount());
                 for (int j = isolatedUids.length - 1; j >= 0; --j) {
-                    cpuTimesMs = addCpuTimes(cpuTimesMs,
-                            mKernelSingleUidTimeReader.readDeltaMs(isolatedUids[j]));
+                    if (isolatedUidTimeInFreqCounters[j] != null) {
+                        mKernelSingleUidTimeReader.addDelta(isolatedUids[j],
+                                isolatedUidTimeInFreqCounters[j], timestampMs, deltaContainer);
+                        onBatteryCounter.addCounts(deltaContainer);
+                        onBatteryScreenOffCounter.addCounts(deltaContainer);
+                    }
                 }
             }
-            if (onBattery && cpuTimesMs != null) {
-                synchronized (BatteryStatsImpl.this) {
-                    u.addProcStateTimesMs(procState, cpuTimesMs, onBattery);
-                    u.addProcStateScreenOffTimesMs(procState, cpuTimesMs, onBatteryScreenOff);
-                }
-            }
+
+            onBatteryCounter.setState(u.mProcessState, timestampMs);
+            onBatteryScreenOffCounter.setState(u.mProcessState, timestampMs);
         }
     }
 
@@ -537,6 +556,7 @@
                 return;
             }
 
+            // TODO(b/197162116): just get a list of UIDs
             final SparseArray<long[]> allUidCpuFreqTimesMs =
                     mCpuUidFreqTimeReader.getAllUidCpuFreqTimeMs();
             // If the KernelSingleUidTimeReader has stale cpu times, then we shouldn't try to
@@ -553,24 +573,40 @@
                 if (u == null) {
                     continue;
                 }
-                final long[] cpuTimesMs = allUidCpuFreqTimesMs.valueAt(i);
-                if (cpuTimesMs == null) {
-                    continue;
+
+                final int procState;
+                final int idx = mPendingUids.indexOfKey(uid);
+                if (idx >= 0) {
+                    procState = mPendingUids.valueAt(idx);
+                    mPendingUids.removeAt(idx);
+                } else {
+                    procState = u.mProcessState;
                 }
-                final long[] deltaTimesMs = mKernelSingleUidTimeReader.computeDelta(
-                        uid, cpuTimesMs.clone());
-                if (onBattery && deltaTimesMs != null) {
-                    final int procState;
-                    final int idx = mPendingUids.indexOfKey(uid);
-                    if (idx >= 0) {
-                        procState = mPendingUids.valueAt(idx);
-                        mPendingUids.removeAt(idx);
-                    } else {
-                        procState = u.mProcessState;
-                    }
-                    if (procState >= 0 && procState < Uid.NUM_PROCESS_STATE) {
-                        u.addProcStateTimesMs(procState, deltaTimesMs, onBattery);
-                        u.addProcStateScreenOffTimesMs(procState, deltaTimesMs, onBatteryScreenOff);
+
+                final long timestampMs = mClock.elapsedRealtime();
+                final LongArrayMultiStateCounter onBatteryCounter =
+                        u.getProcStateTimeCounter().getCounter();
+                final LongArrayMultiStateCounter onBatteryScreenOffCounter =
+                        u.getProcStateScreenOffTimeCounter().getCounter();
+
+                onBatteryCounter.setState(procState, timestampMs);
+                mKernelSingleUidTimeReader.addDelta(uid, onBatteryCounter, timestampMs);
+
+                onBatteryScreenOffCounter.setState(procState, timestampMs);
+                mKernelSingleUidTimeReader.addDelta(uid, onBatteryScreenOffCounter, timestampMs);
+
+                if (u.mChildUids != null) {
+                    final LongArrayMultiStateCounter.LongArrayContainer deltaContainer =
+                            new LongArrayMultiStateCounter.LongArrayContainer(getCpuFreqCount());
+                    for (int j = u.mChildUids.size() - 1; j >= 0; --j) {
+                        final LongArrayMultiStateCounter counter = u.mChildUids.valueAt(j);
+                        if (counter != null) {
+                            final int isolatedUid = u.mChildUids.keyAt(j);
+                            mKernelSingleUidTimeReader.addDelta(isolatedUid,
+                                    counter, timestampMs, deltaContainer);
+                            onBatteryCounter.addCounts(deltaContainer);
+                            onBatteryScreenOffCounter.addCounts(deltaContainer);
+                        }
                     }
                 }
             }
@@ -1653,6 +1689,97 @@
         }
     }
 
+    private static class TimeInFreqMultiStateCounter implements TimeBaseObs {
+        private final TimeBase mTimeBase;
+        private final LongArrayMultiStateCounter mCounter;
+
+        private TimeInFreqMultiStateCounter(TimeBase timeBase, Parcel in, long timestampMs) {
+            mTimeBase = timeBase;
+            mCounter = LongArrayMultiStateCounter.CREATOR.createFromParcel(in);
+            mCounter.setEnabled(mTimeBase.isRunning(), timestampMs);
+            timeBase.add(this);
+        }
+
+        private TimeInFreqMultiStateCounter(TimeBase timeBase, int stateCount, int cpuFreqCount,
+                long timestampMs) {
+            mTimeBase = timeBase;
+            mCounter = new LongArrayMultiStateCounter(stateCount, cpuFreqCount);
+            mCounter.setEnabled(mTimeBase.isRunning(), timestampMs);
+            timeBase.add(this);
+        }
+
+        private void writeToParcel(Parcel out) {
+            mCounter.writeToParcel(out, 0);
+        }
+
+        @Override
+        public void onTimeStarted(long elapsedRealtimeUs, long baseUptimeUs, long baseRealtimeUs) {
+            mCounter.setEnabled(true, elapsedRealtimeUs / 1000);
+        }
+
+        @Override
+        public void onTimeStopped(long elapsedRealtimeUs, long baseUptimeUs, long baseRealtimeUs) {
+            mCounter.setEnabled(false, elapsedRealtimeUs / 1000);
+        }
+
+        public LongArrayMultiStateCounter getCounter() {
+            return mCounter;
+        }
+
+        public int getStateCount() {
+            return mCounter.getStateCount();
+        }
+
+        public void setTrackingEnabled(boolean enabled, long timestampMs) {
+            mCounter.setEnabled(enabled && mTimeBase.isRunning(), timestampMs);
+        }
+
+        private void setState(int uidRunningState, long elapsedRealtimeMs) {
+            mCounter.setState(uidRunningState, elapsedRealtimeMs);
+        }
+
+        /**
+         * Returns accumulated counts for the specified state, or null if all counts are zero.
+         */
+        @Nullable
+        public long[] getCountsLocked(int which, int procState) {
+            LongArrayMultiStateCounter.LongArrayContainer longArrayContainer =
+                    new LongArrayMultiStateCounter.LongArrayContainer(mCounter.getArrayLength());
+            mCounter.getCounts(longArrayContainer, procState);
+            final long[] counts = new long[mCounter.getArrayLength()];
+            longArrayContainer.getValues(counts);
+
+            // Return counts only if at least one of the elements is non-zero.
+            for (int i = counts.length - 1; i >= 0; --i) {
+                if (counts[i] != 0) {
+                    return counts;
+                }
+            }
+            return null;
+        }
+
+        public void logState(Printer pw, String prefix) {
+            pw.println(prefix + "mCounter=" + mCounter);
+        }
+
+        /**
+         * Clears state of this counter.
+         */
+        @Override
+        public boolean reset(boolean detachIfReset, long elapsedRealtimeUs /* unused */) {
+            mCounter.reset();
+            if (detachIfReset) {
+                detach();
+            }
+            return true;
+        }
+
+        @Override
+        public void detach() {
+            mTimeBase.remove(this);
+        }
+    }
+
     @VisibleForTesting
     public static class LongSamplingCounter extends LongCounter implements TimeBaseObs {
         final TimeBase mTimeBase;
@@ -7300,10 +7427,10 @@
         LongSamplingCounterArray mScreenOffCpuFreqTimeMs;
         LongSamplingCounterArray mCpuClusterTimesMs;
 
-        LongSamplingCounterArray[] mProcStateTimeMs;
-        LongSamplingCounterArray[] mProcStateScreenOffTimeMs;
+        TimeInFreqMultiStateCounter mProcStateTimeMs;
+        TimeInFreqMultiStateCounter mProcStateScreenOffTimeMs;
 
-        IntArray mChildUids;
+        SparseArray<LongArrayMultiStateCounter> mChildUids;
 
         /**
          * The statistics we have collected for this uid's wake locks.
@@ -7455,8 +7582,10 @@
         }
 
         @VisibleForTesting
-        public void setProcessStateForTest(int procState) {
+        public void setProcessStateForTest(int procState, long elapsedTimeMs) {
             mProcessState = procState;
+            getProcStateTimeCounter().setState(procState, elapsedTimeMs);
+            getProcStateScreenOffTimeCounter().setState(procState, elapsedTimeMs);
         }
 
         @Override
@@ -7481,7 +7610,7 @@
 
         @Override
         public long[] getCpuFreqTimes(int which, int procState) {
-            if (which < 0 || which >= NUM_PROCESS_STATE) {
+            if (procState < 0 || procState >= NUM_PROCESS_STATE) {
                 return null;
             }
             if (mProcStateTimeMs == null) {
@@ -7491,12 +7620,13 @@
                 mProcStateTimeMs = null;
                 return null;
             }
-            return nullIfAllZeros(mProcStateTimeMs[procState], which);
+
+            return mProcStateTimeMs.getCountsLocked(which, procState);
         }
 
         @Override
         public long[] getScreenOffCpuFreqTimes(int which, int procState) {
-            if (which < 0 || which >= NUM_PROCESS_STATE) {
+            if (procState < 0 || procState >= NUM_PROCESS_STATE) {
                 return null;
             }
             if (mProcStateScreenOffTimeMs == null) {
@@ -7506,7 +7636,7 @@
                 mProcStateScreenOffTimeMs = null;
                 return null;
             }
-            return nullIfAllZeros(mProcStateScreenOffTimeMs[procState], which);
+            return mProcStateScreenOffTimeMs.getCountsLocked(which, procState);
         }
 
         public long getBinderCallCount() {
@@ -7525,15 +7655,26 @@
 
         public void addIsolatedUid(int isolatedUid) {
             if (mChildUids == null) {
-                mChildUids = new IntArray();
-            } else if (mChildUids.indexOf(isolatedUid) >= 0) {
+                mChildUids = new SparseArray<>();
+            } else if (mChildUids.indexOfKey(isolatedUid) >= 0) {
                 return;
             }
-            mChildUids.add(isolatedUid);
+            if (mBsi.trackPerProcStateCpuTimes()) {
+                LongArrayMultiStateCounter counter =
+                        new LongArrayMultiStateCounter(1, mBsi.getCpuFreqCount());
+                // Set initial values to all 0. This is a child UID and we want to include
+                // the entirety of its CPU time-in-freq stats into the parent's stats.
+                counter.updateValues(
+                        new LongArrayMultiStateCounter.LongArrayContainer(mBsi.getCpuFreqCount()),
+                        mBsi.mClock.elapsedRealtime());
+                mChildUids.put(isolatedUid, counter);
+            } else {
+                mChildUids.put(isolatedUid, null);
+            }
         }
 
         public void removeIsolatedUid(int isolatedUid) {
-            final int idx = mChildUids == null ? -1 : mChildUids.indexOf(isolatedUid);
+            final int idx = mChildUids == null ? -1 : mChildUids.indexOfKey(isolatedUid);
             if (idx < 0) {
                 return;
             }
@@ -7557,31 +7698,37 @@
             return null;
         }
 
-        private void addProcStateTimesMs(int procState, long[] cpuTimesMs, boolean onBattery) {
-            if (mProcStateTimeMs == null) {
-                mProcStateTimeMs = new LongSamplingCounterArray[NUM_PROCESS_STATE];
+        private void ensureMultiStateCounters() {
+            if (mProcStateTimeMs != null) {
+                return;
             }
-            if (mProcStateTimeMs[procState] == null
-                    || mProcStateTimeMs[procState].getSize() != cpuTimesMs.length) {
-                detachIfNotNull(mProcStateTimeMs[procState]);
-                mProcStateTimeMs[procState] = new LongSamplingCounterArray(
-                        mBsi.mOnBatteryTimeBase);
-            }
-            mProcStateTimeMs[procState].addCountLocked(cpuTimesMs, onBattery);
+
+            final long timestampMs = mBsi.mClock.elapsedRealtime();
+            mProcStateTimeMs =
+                    new TimeInFreqMultiStateCounter(mBsi.mOnBatteryTimeBase,
+                            NUM_PROCESS_STATE, mBsi.getCpuFreqCount(), timestampMs);
+            mProcStateScreenOffTimeMs =
+                    new TimeInFreqMultiStateCounter(mBsi.mOnBatteryScreenOffTimeBase,
+                            NUM_PROCESS_STATE, mBsi.getCpuFreqCount(), timestampMs);
         }
 
-        private void addProcStateScreenOffTimesMs(int procState, long[] cpuTimesMs,
-                boolean onBatteryScreenOff) {
-            if (mProcStateScreenOffTimeMs == null) {
-                mProcStateScreenOffTimeMs = new LongSamplingCounterArray[NUM_PROCESS_STATE];
+        private TimeInFreqMultiStateCounter getProcStateTimeCounter() {
+            ensureMultiStateCounters();
+            return mProcStateTimeMs;
+        }
+
+        private TimeInFreqMultiStateCounter getProcStateScreenOffTimeCounter() {
+            ensureMultiStateCounters();
+            return mProcStateScreenOffTimeMs;
+        }
+
+        private void setProcStateTimesTrackingEnabled(boolean enabled, long timestampMs) {
+            if (mProcStateTimeMs != null) {
+                mProcStateTimeMs.setTrackingEnabled(enabled, timestampMs);
             }
-            if (mProcStateScreenOffTimeMs[procState] == null
-                    || mProcStateScreenOffTimeMs[procState].getSize() != cpuTimesMs.length) {
-                detachIfNotNull(mProcStateScreenOffTimeMs[procState]);
-                mProcStateScreenOffTimeMs[procState] = new LongSamplingCounterArray(
-                        mBsi.mOnBatteryScreenOffTimeBase);
+            if (mProcStateScreenOffTimeMs != null) {
+                mProcStateScreenOffTimeMs.setTrackingEnabled(enabled, timestampMs);
             }
-            mProcStateScreenOffTimeMs[procState].addCountLocked(cpuTimesMs, onBatteryScreenOff);
         }
 
         @Override
@@ -9112,18 +9259,14 @@
             mCpuClusterTimesMs.writeToParcel(out);
 
             if (mProcStateTimeMs != null) {
-                out.writeInt(mProcStateTimeMs.length);
-                for (LongSamplingCounterArray counters : mProcStateTimeMs) {
-                    LongSamplingCounterArray.writeToParcel(out, counters);
-                }
+                out.writeInt(mProcStateTimeMs.getStateCount());
+                mProcStateTimeMs.writeToParcel(out);
             } else {
                 out.writeInt(0);
             }
             if (mProcStateScreenOffTimeMs != null) {
-                out.writeInt(mProcStateScreenOffTimeMs.length);
-                for (LongSamplingCounterArray counters : mProcStateScreenOffTimeMs) {
-                    LongSamplingCounterArray.writeToParcel(out, counters);
-                }
+                out.writeInt(mProcStateScreenOffTimeMs.getStateCount());
+                mProcStateScreenOffTimeMs.writeToParcel(out);
             } else {
                 out.writeInt(0);
             }
@@ -9416,22 +9559,27 @@
             mCpuActiveTimeMs = new LongSamplingCounter(mBsi.mOnBatteryTimeBase, in);
             mCpuClusterTimesMs = new LongSamplingCounterArray(mBsi.mOnBatteryTimeBase, in);
 
-            int length = in.readInt();
-            if (length == NUM_PROCESS_STATE) {
-                mProcStateTimeMs = new LongSamplingCounterArray[length];
-                for (int procState = 0; procState < length; ++procState) {
-                    mProcStateTimeMs[procState] = LongSamplingCounterArray.readFromParcel(
-                            in, mBsi.mOnBatteryTimeBase);
+            final long timestampMs = mBsi.mClock.elapsedRealtime();
+            int stateCount = in.readInt();
+            if (stateCount != 0) {
+                // Read the object from the Parcel, whether it's usable or not
+                TimeInFreqMultiStateCounter counter = new TimeInFreqMultiStateCounter(
+                        mBsi.mOnBatteryTimeBase, in, timestampMs);
+                if (stateCount == NUM_PROCESS_STATE) {
+                    mProcStateTimeMs = counter;
                 }
             } else {
                 mProcStateTimeMs = null;
             }
-            length = in.readInt();
-            if (length == NUM_PROCESS_STATE) {
-                mProcStateScreenOffTimeMs = new LongSamplingCounterArray[length];
-                for (int procState = 0; procState < length; ++procState) {
-                    mProcStateScreenOffTimeMs[procState] = LongSamplingCounterArray.readFromParcel(
-                            in, mBsi.mOnBatteryScreenOffTimeBase);
+
+            stateCount = in.readInt();
+            if (stateCount != 0) {
+                // Read the object from the Parcel, whether it's usable or not
+                TimeInFreqMultiStateCounter counter =
+                        new TimeInFreqMultiStateCounter(
+                                mBsi.mOnBatteryScreenOffTimeBase, in, timestampMs);
+                if (stateCount == NUM_PROCESS_STATE) {
+                    mProcStateScreenOffTimeMs = counter;
                 }
             } else {
                 mProcStateScreenOffTimeMs = null;
@@ -10607,6 +10755,16 @@
         return mCpuFreqs;
     }
 
+    /**
+     * Returns the total number of frequencies across all CPU clusters.
+     */
+    private int getCpuFreqCount() {
+        if (mCpuFreqs == null) {
+            mCpuFreqs = mCpuUidFreqTimeReader.readFreqs(mPowerProfile);
+        }
+        return mCpuFreqs != null ? mCpuFreqs.length : 0;
+    }
+
     public BatteryStatsImpl(File systemDir, Handler handler, PlatformIdleStateCallback cb,
             MeasuredEnergyRetriever energyStatsCb, UserInfoProvider userInfoProvider) {
         this(Clock.SYSTEM_CLOCK, systemDir, handler, cb, energyStatsCb, userInfoProvider);
@@ -14751,6 +14909,7 @@
                     DEFAULT_BATTERY_CHARGED_DELAY_MS);
         }
 
+        @GuardedBy("BatteryStatsImpl.this")
         private void updateTrackCpuTimesByProcStateLocked(boolean wasEnabled, boolean isEnabled) {
             TRACK_CPU_TIMES_BY_PROC_STATE = isEnabled;
             if (isEnabled && !wasEnabled) {
@@ -14761,6 +14920,10 @@
                 mNumBatchedSingleUidCpuTimeReads = 0;
                 mCpuTimeReadsTrackingStartTimeMs = mClock.uptimeMillis();
             }
+            final long timestampMs = mClock.elapsedRealtime();
+            for (int i = mUidStats.size() - 1; i >= 0; i--) {
+                mUidStats.valueAt(i).setProcStateTimesTrackingEnabled(isEnabled, timestampMs);
+            }
         }
 
         private void updateProcStateCpuTimesReadDelayMs(long oldDelayMillis, long newDelayMillis) {
@@ -15535,31 +15698,32 @@
             u.mCpuActiveTimeMs.readSummaryFromParcelLocked(in);
             u.mCpuClusterTimesMs.readSummaryFromParcelLocked(in);
 
-            int length = in.readInt();
-            if (length == Uid.NUM_PROCESS_STATE) {
-                detachIfNotNull(u.mProcStateTimeMs);
-                u.mProcStateTimeMs = new LongSamplingCounterArray[length];
-                for (int procState = 0; procState < length; ++procState) {
-                    u.mProcStateTimeMs[procState]
-                            = LongSamplingCounterArray.readSummaryFromParcelLocked(
-                                    in, mOnBatteryTimeBase);
+            detachIfNotNull(u.mProcStateTimeMs);
+            u.mProcStateTimeMs = null;
+
+            int stateCount = in.readInt();
+            if (stateCount != 0) {
+                // Read the object from the Parcel, whether it's usable or not
+                TimeInFreqMultiStateCounter counter = new TimeInFreqMultiStateCounter(
+                        mOnBatteryTimeBase, in, mClock.elapsedRealtime());
+                if (stateCount == Uid.NUM_PROCESS_STATE) {
+                    u.mProcStateTimeMs = counter;
                 }
-            } else {
-                detachIfNotNull(u.mProcStateTimeMs);
-                u.mProcStateTimeMs = null;
             }
-            length = in.readInt();
-            if (length == Uid.NUM_PROCESS_STATE) {
+
+            detachIfNotNull(u.mProcStateScreenOffTimeMs);
+            u.mProcStateScreenOffTimeMs = null;
+
+            stateCount = in.readInt();
+            if (stateCount != 0) {
                 detachIfNotNull(u.mProcStateScreenOffTimeMs);
-                u.mProcStateScreenOffTimeMs = new LongSamplingCounterArray[length];
-                for (int procState = 0; procState < length; ++procState) {
-                    u.mProcStateScreenOffTimeMs[procState]
-                            = LongSamplingCounterArray.readSummaryFromParcelLocked(
-                                    in, mOnBatteryScreenOffTimeBase);
+                // Read the object from the Parcel, whether it's usable or not
+                TimeInFreqMultiStateCounter counter =
+                        new TimeInFreqMultiStateCounter(
+                                mOnBatteryScreenOffTimeBase, in, mClock.elapsedRealtime());
+                if (stateCount == Uid.NUM_PROCESS_STATE) {
+                    u.mProcStateScreenOffTimeMs = counter;
                 }
-            } else {
-                detachIfNotNull(u.mProcStateScreenOffTimeMs);
-                u.mProcStateScreenOffTimeMs = null;
             }
 
             if (in.readInt() != 0) {
@@ -16062,18 +16226,15 @@
             u.mCpuClusterTimesMs.writeSummaryToParcelLocked(out);
 
             if (u.mProcStateTimeMs != null) {
-                out.writeInt(u.mProcStateTimeMs.length);
-                for (LongSamplingCounterArray counters : u.mProcStateTimeMs) {
-                    LongSamplingCounterArray.writeSummaryToParcelLocked(out, counters);
-                }
+                out.writeInt(u.mProcStateTimeMs.getStateCount());
+                u.mProcStateTimeMs.writeToParcel(out);
             } else {
                 out.writeInt(0);
             }
+
             if (u.mProcStateScreenOffTimeMs != null) {
-                out.writeInt(u.mProcStateScreenOffTimeMs.length);
-                for (LongSamplingCounterArray counters : u.mProcStateScreenOffTimeMs) {
-                    LongSamplingCounterArray.writeSummaryToParcelLocked(out, counters);
-                }
+                out.writeInt(u.mProcStateScreenOffTimeMs.getStateCount());
+                u.mProcStateScreenOffTimeMs.writeToParcel(out);
             } else {
                 out.writeInt(0);
             }
diff --git a/core/java/com/android/internal/os/KernelSingleUidTimeReader.java b/core/java/com/android/internal/os/KernelSingleUidTimeReader.java
index 31952eb..9052644 100644
--- a/core/java/com/android/internal/os/KernelSingleUidTimeReader.java
+++ b/core/java/com/android/internal/os/KernelSingleUidTimeReader.java
@@ -243,6 +243,23 @@
         }
     }
 
+    /**
+     * Retrieves CPU time-in-state data for the specified UID and adds the accumulated
+     * delta to the supplied counter.
+     */
+    public void addDelta(int uid, LongArrayMultiStateCounter counter, long timestampMs) {
+        mInjector.addDelta(uid, counter, timestampMs, null);
+    }
+
+    /**
+     * Same as {@link #addDelta(int, LongArrayMultiStateCounter, long)}, also returning
+     * the delta in the supplied array container.
+     */
+    public void addDelta(int uid, LongArrayMultiStateCounter counter, long timestampMs,
+            LongArrayMultiStateCounter.LongArrayContainer deltaContainer) {
+        mInjector.addDelta(uid, counter, timestampMs, deltaContainer);
+    }
+
     @VisibleForTesting
     public static class Injector {
         public byte[] readData(String procFile) throws IOException {
@@ -254,14 +271,19 @@
         /**
          * Reads CPU time-in-state data for the specified UID and adds the delta since the
          * previous call to the current state stats in the LongArrayMultiStateCounter.
+         *
+         * The delta is also returned via the optional deltaOut parameter.
          */
-        public boolean addDelta(int uid, LongArrayMultiStateCounter counter, long timestampMs) {
-            return addDeltaFromBpf(uid, counter.mNativeObject, timestampMs);
+        public boolean addDelta(int uid, LongArrayMultiStateCounter counter, long timestampMs,
+                LongArrayMultiStateCounter.LongArrayContainer deltaOut) {
+            return addDeltaFromBpf(uid, counter.mNativeObject, timestampMs,
+                    deltaOut != null ? deltaOut.mNativeObject : 0);
         }
 
         @CriticalNative
         private static native boolean addDeltaFromBpf(int uid,
-                long longArrayMultiStateCounterNativePointer, long timestampMs);
+                long longArrayMultiStateCounterNativePointer, long timestampMs,
+                long longArrayContainerNativePointer);
 
         /**
          * Used for testing.
@@ -269,13 +291,15 @@
          * Takes mock cpu-time-in-frequency data and uses it the same way eBPF data would be used.
          */
         public boolean addDeltaForTest(int uid, LongArrayMultiStateCounter counter,
-                long timestampMs, long[][] timeInFreqDataNanos) {
-            return addDeltaForTest(uid, counter.mNativeObject, timestampMs, timeInFreqDataNanos);
+                long timestampMs, long[][] timeInFreqDataNanos,
+                LongArrayMultiStateCounter.LongArrayContainer deltaOut) {
+            return addDeltaForTest(uid, counter.mNativeObject, timestampMs, timeInFreqDataNanos,
+                    deltaOut != null ? deltaOut.mNativeObject : 0);
         }
 
         private static native boolean addDeltaForTest(int uid,
                 long longArrayMultiStateCounterNativePointer, long timestampMs,
-                long[][] timeInFreqDataNanos);
+                long[][] timeInFreqDataNanos, long longArrayContainerNativePointer);
     }
 
     @VisibleForTesting
diff --git a/core/java/com/android/internal/os/LongArrayMultiStateCounter.java b/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
index aecab3d..d98b73f 100644
--- a/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
+++ b/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
@@ -26,6 +26,8 @@
 
 import libcore.util.NativeAllocationRegistry;
 
+import java.util.Arrays;
+
 /**
  * Performs per-state counting of multi-element values over time. The class' behavior is illustrated
  * by this example:
@@ -62,7 +64,9 @@
                 NativeAllocationRegistry.createMalloced(
                         LongArrayContainer.class.getClassLoader(), native_getReleaseFunc());
 
-        private final long mNativeObject;
+        // Visible to other objects in this package so that it can be passed to @CriticalNative
+        // methods.
+        final long mNativeObject;
         private final int mLength;
 
         public LongArrayContainer(int length) {
@@ -93,6 +97,13 @@
             native_getValues(mNativeObject, array);
         }
 
+        @Override
+        public String toString() {
+            final long[] array = new long[mLength];
+            getValues(array);
+            return Arrays.toString(array);
+        }
+
         @CriticalNative
         private static native long native_init(int length);
 
@@ -133,6 +144,14 @@
         mLength = native_getArrayLength(mNativeObject);
     }
 
+    public int getStateCount() {
+        return mStateCount;
+    }
+
+    public int getArrayLength() {
+        return mLength;
+    }
+
     /**
      * Enables or disables the counter.  When the counter is disabled, it does not
      * accumulate counts supplied by the {@link #updateValues} method.
@@ -147,7 +166,7 @@
     public void setState(int state, long timestampMs) {
         if (state < 0 || state >= mStateCount) {
             throw new IllegalArgumentException(
-                    "State: " + state + ", outside the range: [0-" + mStateCount + "]");
+                    "State: " + state + ", outside the range: [0-" + (mStateCount - 1) + "]");
         }
         native_setState(mNativeObject, state, timestampMs);
     }
@@ -167,6 +186,17 @@
     }
 
     /**
+     * Adds the supplied values to the current accumulated values in the counter.
+     */
+    public void addCounts(LongArrayContainer counts) {
+        if (counts.mLength != mLength) {
+            throw new IllegalArgumentException(
+                    "Invalid array length: " + counts.mLength + ", expected: " + mLength);
+        }
+        native_addCounts(mNativeObject, counts.mNativeObject);
+    }
+
+    /**
      * Resets the accumulated counts to 0.
      */
     public void reset() {
@@ -231,6 +261,10 @@
             long longArrayContainerNativeObject, long timestampMs);
 
     @CriticalNative
+    private static native void native_addCounts(long nativeObject,
+            long longArrayContainerNativeObject);
+
+    @CriticalNative
     private static native void native_reset(long nativeObject);
 
     @CriticalNative
diff --git a/core/jni/com_android_internal_os_KernelSingleUidTimeReader.cpp b/core/jni/com_android_internal_os_KernelSingleUidTimeReader.cpp
index 689b259..bea5ffe 100644
--- a/core/jni/com_android_internal_os_KernelSingleUidTimeReader.cpp
+++ b/core/jni/com_android_internal_os_KernelSingleUidTimeReader.cpp
@@ -50,7 +50,8 @@
  */
 static jboolean addCpuTimeInFreqDelta(
         jint uid, jlong counterNativePtr, jlong timestampMs,
-        std::optional<std::vector<std::vector<uint64_t>>> timeInFreqDataNanos) {
+        std::optional<std::vector<std::vector<uint64_t>>> timeInFreqDataNanos,
+        jlong deltaOutContainerNativePtr) {
     if (!timeInFreqDataNanos) {
         return false;
     }
@@ -70,20 +71,35 @@
     for (size_t i = 0; i < s; ++i) {
         flattened[i] /= NSEC_PER_MSEC;
     }
-    counter->updateValue(flattened, timestampMs);
+    if (s != counter->getCount(0).size()) { // Counter has at least one state
+        ALOGE("Mismatch between eBPF data size (%d) and the counter size (%d)", (int)s,
+              (int)counter->getCount(0).size());
+        return false;
+    }
+
+    const std::vector<uint64_t> &delta = counter->updateValue(flattened, timestampMs);
+    if (deltaOutContainerNativePtr) {
+        std::vector<uint64_t> *vector =
+                reinterpret_cast<std::vector<uint64_t> *>(deltaOutContainerNativePtr);
+        *vector = delta;
+    }
+
     return true;
 }
 
-static jboolean addDeltaFromBpf(jint uid, jlong counterNativePtr, jlong timestampMs) {
+static jboolean addDeltaFromBpf(jint uid, jlong counterNativePtr, jlong timestampMs,
+                                jlong deltaOutContainerNativePtr) {
     return addCpuTimeInFreqDelta(uid, counterNativePtr, timestampMs,
-                                 android::bpf::getUidCpuFreqTimes(uid));
+                                 android::bpf::getUidCpuFreqTimes(uid), deltaOutContainerNativePtr);
 }
 
 static jboolean addDeltaForTest(JNIEnv *env, jclass, jint uid, jlong counterNativePtr,
-                                jlong timestampMs, jobjectArray timeInFreqDataNanos) {
+                                jlong timestampMs, jobjectArray timeInFreqDataNanos,
+                                jlong deltaOutContainerNativePtr) {
     if (!timeInFreqDataNanos) {
         return addCpuTimeInFreqDelta(uid, counterNativePtr, timestampMs,
-                                     std::optional<std::vector<std::vector<uint64_t>>>());
+                                     std::optional<std::vector<std::vector<uint64_t>>>(),
+                                     deltaOutContainerNativePtr);
     }
 
     std::vector<std::vector<uint64_t>> timeInFreqData;
@@ -97,17 +113,18 @@
         }
         timeInFreqData.push_back(cluster);
     }
-    return addCpuTimeInFreqDelta(uid, counterNativePtr, timestampMs, std::optional(timeInFreqData));
+    return addCpuTimeInFreqDelta(uid, counterNativePtr, timestampMs, std::optional(timeInFreqData),
+                                 deltaOutContainerNativePtr);
 }
 
 static const JNINativeMethod g_single_methods[] = {
         {"readBpfData", "(I)[J", (void *)getUidCpuFreqTimeMs},
 
         // @CriticalNative
-        {"addDeltaFromBpf", "(IJJ)Z", (void *)addDeltaFromBpf},
+        {"addDeltaFromBpf", "(IJJJ)Z", (void *)addDeltaFromBpf},
 
         // Used for testing
-        {"addDeltaForTest", "(IJJ[[J)Z", (void *)addDeltaForTest},
+        {"addDeltaForTest", "(IJJ[[JJ)Z", (void *)addDeltaForTest},
 };
 
 int register_com_android_internal_os_KernelSingleUidTimeReader(JNIEnv *env) {
diff --git a/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp b/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
index 8326de2..6e7ebda 100644
--- a/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
+++ b/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
@@ -65,6 +65,14 @@
     counter->updateValue(*vector, timestamp);
 }
 
+static void native_addCounts(jlong nativePtr, jlong longArrayContainerNativePtr) {
+    battery::LongArrayMultiStateCounter *counter =
+            reinterpret_cast<battery::LongArrayMultiStateCounter *>(nativePtr);
+    std::vector<uint64_t> *vector =
+            reinterpret_cast<std::vector<uint64_t> *>(longArrayContainerNativePtr);
+    counter->addValue(*vector);
+}
+
 static void native_reset(jlong nativePtr) {
     battery::LongArrayMultiStateCounter *counter =
             reinterpret_cast<battery::LongArrayMultiStateCounter *>(nativePtr);
@@ -184,6 +192,8 @@
         // @CriticalNative
         {"native_updateValues", "(JJJ)V", (void *)native_updateValues},
         // @CriticalNative
+        {"native_addCounts", "(JJ)V", (void *)native_addCounts},
+        // @CriticalNative
         {"native_reset", "(J)V", (void *)native_reset},
         // @CriticalNative
         {"native_getCounts", "(JJI)V", (void *)native_getCounts},
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsImplTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsImplTest.java
index cca6642..8b060ff 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsImplTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsImplTest.java
@@ -27,10 +27,13 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
-import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.when;
 
+import android.app.ActivityManager;
 import android.os.BatteryStats;
 import android.util.SparseArray;
 import android.util.SparseIntArray;
@@ -40,20 +43,19 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.os.KernelCpuUidTimeReader.KernelCpuUidFreqTimeReader;
-import com.android.internal.util.ArrayUtils;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentMatcher;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.Arrays;
-
 @LargeTest
 @RunWith(AndroidJUnit4.class)
 public class BatteryStatsImplTest {
+    private static final long[] CPU_FREQS = {1, 2, 3, 4, 5};
+    private static final int NUM_CPU_FREQS = CPU_FREQS.length;
+
     @Mock
     private KernelCpuUidFreqTimeReader mKernelUidCpuFreqTimeReader;
     @Mock
@@ -61,13 +63,16 @@
 
     private MockBatteryStatsImpl mBatteryStatsImpl;
 
+    private MockClock mMockClock = new MockClock();
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
+        when(mKernelUidCpuFreqTimeReader.readFreqs(any())).thenReturn(CPU_FREQS);
         when(mKernelUidCpuFreqTimeReader.allUidTimesAvailable()).thenReturn(true);
         when(mKernelSingleUidTimeReader.singleUidCpuTimesAvailable()).thenReturn(true);
-        mBatteryStatsImpl = new MockBatteryStatsImpl()
+        mBatteryStatsImpl = new MockBatteryStatsImpl(mMockClock)
                 .setKernelCpuUidFreqTimeReader(mKernelUidCpuFreqTimeReader)
                 .setKernelSingleUidTimeReader(mKernelSingleUidTimeReader)
                 .setTrackingCpuByProcStateEnabled(true);
@@ -76,9 +81,17 @@
     @Test
     public void testUpdateProcStateCpuTimes() {
         mBatteryStatsImpl.setOnBatteryInternal(true);
-        mBatteryStatsImpl.updateTimeBasesLocked(false, Display.STATE_ON, 0, 0);
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_ON, 0, 0);
+        }
 
         final int[] testUids = {10032, 10048, 10145, 10139};
+        final int[] activityManagerProcStates = {
+                ActivityManager.PROCESS_STATE_RECEIVER,
+                ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE,
+                ActivityManager.PROCESS_STATE_TOP,
+                ActivityManager.PROCESS_STATE_CACHED_EMPTY
+        };
         final int[] testProcStates = {
                 PROCESS_STATE_BACKGROUND,
                 PROCESS_STATE_FOREGROUND_SERVICE,
@@ -86,14 +99,24 @@
                 PROCESS_STATE_CACHED
         };
         addPendingUids(testUids, testProcStates);
+
+        // Initialize time-in-freq counters
+        mMockClock.realtime = 1000;
+        for (int i = 0; i < testUids.length; ++i) {
+            mBatteryStatsImpl.noteUidProcessStateLocked(testUids[i], activityManagerProcStates[i]);
+            mockKernelSingleUidTimeReader(testUids[i], new long[5]);
+        }
+        mBatteryStatsImpl.updateProcStateCpuTimes(true, false);
+
+        // Obtain initial CPU time-in-freq counts
         final long[][] cpuTimes = {
-                {349734983, 394982394832l, 909834, 348934, 9838},
+                {349734983, 394982394832L, 909834, 348934, 9838},
                 {7498, 1239890, 988, 13298, 98980},
                 {989834, 384098, 98483, 23809, 4984},
                 {4859048, 348903, 4578967, 5973894, 298549}
         };
         for (int i = 0; i < testUids.length; ++i) {
-            when(mKernelSingleUidTimeReader.readDeltaMs(testUids[i])).thenReturn(cpuTimes[i]);
+            mockKernelSingleUidTimeReader(testUids[i], cpuTimes[i]);
 
             // Verify there are no cpu times initially.
             final BatteryStats.Uid u = mBatteryStatsImpl.getUidStatsLocked(testUids[i]);
@@ -102,7 +125,9 @@
                 assertNull(u.getScreenOffCpuFreqTimes(STATS_SINCE_CHARGED, procState));
             }
         }
+        addPendingUids(testUids, testProcStates);
 
+        mMockClock.realtime += 1000;
         mBatteryStatsImpl.updateProcStateCpuTimes(true, false);
 
         verifyNoPendingUids();
@@ -119,6 +144,7 @@
             }
         }
 
+        // Accumulate CPU time-in-freq deltas
         final long[][] delta1 = {
                 {9589, 148934, 309894, 3098493, 98754},
                 {21983, 94983, 4983, 9878493, 84854},
@@ -126,10 +152,15 @@
                 {843895, 43948, 949582, 99, 384}
         };
         for (int i = 0; i < testUids.length; ++i) {
-            when(mKernelSingleUidTimeReader.readDeltaMs(testUids[i])).thenReturn(delta1[i]);
+            long[] newCpuTimes = new long[cpuTimes[i].length];
+            for (int j = 0; j < cpuTimes[i].length; j++) {
+                newCpuTimes[j] = cpuTimes[i][j] + delta1[i][j];
+            }
+            mockKernelSingleUidTimeReader(testUids[i], newCpuTimes);
         }
         addPendingUids(testUids, testProcStates);
 
+        mMockClock.realtime += 1000;
         mBatteryStatsImpl.updateProcStateCpuTimes(true, false);
 
         verifyNoPendingUids();
@@ -150,7 +181,12 @@
             }
         }
 
-        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        // Validate the on-battery-screen-off counter
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0,
+                    mMockClock.realtime * 1000);
+        }
+
         final long[][] delta2 = {
                 {95932, 2943, 49834, 89034, 139},
                 {349, 89605, 5896, 845, 98444},
@@ -158,10 +194,15 @@
                 {488, 998, 8498, 394, 574}
         };
         for (int i = 0; i < testUids.length; ++i) {
-            when(mKernelSingleUidTimeReader.readDeltaMs(testUids[i])).thenReturn(delta2[i]);
+            long[] newCpuTimes = new long[cpuTimes[i].length];
+            for (int j = 0; j < cpuTimes[i].length; j++) {
+                newCpuTimes[j] = cpuTimes[i][j] + delta1[i][j] + delta2[i][j];
+            }
+            mockKernelSingleUidTimeReader(testUids[i], newCpuTimes);
         }
         addPendingUids(testUids, testProcStates);
 
+        mMockClock.realtime += 1000;
         mBatteryStatsImpl.updateProcStateCpuTimes(true, true);
 
         verifyNoPendingUids();
@@ -184,6 +225,10 @@
             }
         }
 
+        // Verify handling of isolated UIDs - their time-in-freq must be directly
+        // added to that of the parent UID's.  The proc state of the isolated UID is
+        // assumed to be the same as that of the parent UID, so there is no per-state
+        // data for isolated UIDs.
         final long[][] delta3 = {
                 {98545, 95768795, 76586, 548945, 57846},
                 {788876, 586, 578459, 8776984, 9578923},
@@ -191,16 +236,20 @@
                 {9493, 784, 99895, 8974893, 9879843}
         };
         for (int i = 0; i < testUids.length; ++i) {
-            when(mKernelSingleUidTimeReader.readDeltaMs(testUids[i])).thenReturn(
-                    delta3[i].clone());
+            long[] newCpuTimes = new long[cpuTimes[i].length];
+            for (int j = 0; j < cpuTimes[i].length; j++) {
+                newCpuTimes[j] = cpuTimes[i][j] + delta1[i][j] + delta2[i][j] + delta3[i][j];
+            }
+            mockKernelSingleUidTimeReader(testUids[i], newCpuTimes);
         }
         addPendingUids(testUids, testProcStates);
         final int parentUid = testUids[1];
         final int childUid = 99099;
         addIsolatedUid(parentUid, childUid);
         final long[] isolatedUidCpuTimes = {495784, 398473, 4895, 4905, 30984093};
-        when(mKernelSingleUidTimeReader.readDeltaMs(childUid)).thenReturn(isolatedUidCpuTimes);
+        mockKernelSingleUidTimeReader(childUid, isolatedUidCpuTimes, isolatedUidCpuTimes);
 
+        mMockClock.realtime += 1000;
         mBatteryStatsImpl.updateProcStateCpuTimes(true, true);
 
         verifyNoPendingUids();
@@ -233,7 +282,11 @@
     @Test
     public void testCopyFromAllUidsCpuTimes() {
         mBatteryStatsImpl.setOnBatteryInternal(false);
-        mBatteryStatsImpl.updateTimeBasesLocked(false, Display.STATE_ON, 0, 0);
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_ON, 0, 0);
+        }
+
+        mMockClock.realtime = 1000;
 
         final int[] testUids = {10032, 10048, 10145, 10139};
         final int[] testProcStates = {
@@ -242,8 +295,14 @@
                 PROCESS_STATE_TOP,
                 PROCESS_STATE_CACHED
         };
-        final int[] pendingUidIdx = {1, 2};
-        updateProcessStates(testUids, testProcStates, pendingUidIdx);
+        addPendingUids(testUids, testProcStates);
+
+        for (int i = 0; i < testUids.length; ++i) {
+            BatteryStatsImpl.Uid uid = mBatteryStatsImpl.getUidStatsLocked(testUids[i]);
+            uid.setProcessStateForTest(testProcStates[i], mMockClock.elapsedRealtime());
+            mockKernelSingleUidTimeReader(testUids[i], new long[NUM_CPU_FREQS]);
+        }
+        mBatteryStatsImpl.updateProcStateCpuTimes(true, false);
 
         final SparseArray<long[]> allUidCpuTimes = new SparseArray<>();
         long[][] allCpuTimes = {
@@ -257,18 +316,16 @@
         }
         when(mKernelUidCpuFreqTimeReader.getAllUidCpuFreqTimeMs()).thenReturn(allUidCpuTimes);
         long[][] expectedCpuTimes = {
-                {843598745, 397843, 32749, 99854},
-                {9834, 5885, 487589, 394},
-                {203984, 439, 9859, 30948},
-                {9389, 858, 239, 349}
+                {843598745, 397843, 32749, 99854, 23454},
+                {9834, 5885, 487589, 394, 93933},
+                {203984, 439, 9859, 30948, 49494},
+                {9389, 858, 239, 349, 50505}
         };
         for (int i = 0; i < testUids.length; ++i) {
-            final int idx = i;
-            final ArgumentMatcher<long[]> matcher = times -> Arrays.equals(times, allCpuTimes[idx]);
-            when(mKernelSingleUidTimeReader.computeDelta(eq(testUids[i]), argThat(matcher)))
-                    .thenReturn(expectedCpuTimes[i]);
+            mockKernelSingleUidTimeReader(testUids[i], expectedCpuTimes[i]);
         }
 
+        mMockClock.realtime += 1000;
         mBatteryStatsImpl.copyFromAllUidsCpuTimes(true, false);
 
         verifyNoPendingUids();
@@ -286,6 +343,39 @@
         }
     }
 
+    private void mockKernelSingleUidTimeReader(int testUid, long[] cpuTimes) {
+        doAnswer(invocation -> {
+            LongArrayMultiStateCounter counter = invocation.getArgument(1);
+            long timestampMs = invocation.getArgument(2);
+            LongArrayMultiStateCounter.LongArrayContainer container =
+                    new LongArrayMultiStateCounter.LongArrayContainer(NUM_CPU_FREQS);
+            container.setValues(cpuTimes);
+            counter.updateValues(container, timestampMs);
+            return null;
+        }).when(mKernelSingleUidTimeReader).addDelta(eq(testUid),
+                any(LongArrayMultiStateCounter.class), anyLong());
+    }
+
+    private void mockKernelSingleUidTimeReader(int testUid, long[] cpuTimes, long[] delta) {
+        doAnswer(invocation -> {
+            LongArrayMultiStateCounter counter = invocation.getArgument(1);
+            long timestampMs = invocation.getArgument(2);
+            LongArrayMultiStateCounter.LongArrayContainer deltaContainer =
+                    invocation.getArgument(3);
+
+            LongArrayMultiStateCounter.LongArrayContainer container =
+                    new LongArrayMultiStateCounter.LongArrayContainer(NUM_CPU_FREQS);
+            container.setValues(cpuTimes);
+            counter.updateValues(container, timestampMs);
+            if (deltaContainer != null) {
+                deltaContainer.setValues(delta);
+            }
+            return null;
+        }).when(mKernelSingleUidTimeReader).addDelta(eq(testUid),
+                any(LongArrayMultiStateCounter.class), anyLong(),
+                any(LongArrayMultiStateCounter.LongArrayContainer.class));
+    }
+
     @Test
     public void testAddCpuTimes() {
         long[] timesA = null;
@@ -314,7 +404,9 @@
         final int releaseTimeMs = 1005;
         final int currentTimeMs = 1011;
 
-        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        }
 
         // Create a Uid Object
         final BatteryStats.Uid u = mBatteryStatsImpl.getUidStatsLocked(testUid);
@@ -325,7 +417,7 @@
         u.noteWifiMulticastDisabledLocked(releaseTimeMs);
 
         // Get the total acquisition time
-        long totalTime = u.getWifiMulticastTime(currentTimeMs*1000,
+        long totalTime = u.getWifiMulticastTime(currentTimeMs * 1000,
                 BatteryStats.STATS_SINCE_CHARGED);
         assertEquals("Miscalculations of Multicast wakelock acquisition time",
                 (releaseTimeMs - acquireTimeMs) * 1000, totalTime);
@@ -337,7 +429,9 @@
         final int acquireTimeMs = 1000;
         final int currentTimeMs = 1011;
 
-        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        }
 
         // Create a Uid Object
         final BatteryStats.Uid u = mBatteryStatsImpl.getUidStatsLocked(testUid);
@@ -347,7 +441,7 @@
         u.noteWifiMulticastEnabledLocked(acquireTimeMs);
 
         // Get the total acquisition time
-        long totalTime =  u.getWifiMulticastTime(currentTimeMs*1000,
+        long totalTime = u.getWifiMulticastTime(currentTimeMs * 1000,
                 BatteryStats.STATS_SINCE_CHARGED);
         assertEquals("Miscalculations of Multicast wakelock acquisition time",
                 (currentTimeMs - acquireTimeMs) * 1000, totalTime);
@@ -363,7 +457,9 @@
         final int releaseTimeMs_2 = 1009;
         final int currentTimeMs = 1011;
 
-        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        }
 
         // Create a Uid Object
         final BatteryStats.Uid u = mBatteryStatsImpl.getUidStatsLocked(testUid);
@@ -377,7 +473,7 @@
         u.noteWifiMulticastDisabledLocked(releaseTimeMs_2);
 
         // Get the total acquisition time
-        long totalTime =  u.getWifiMulticastTime(currentTimeMs*1000,
+        long totalTime = u.getWifiMulticastTime(currentTimeMs * 1000,
                 BatteryStats.STATS_SINCE_CHARGED);
         assertEquals("Miscalculations of Multicast wakelock acquisition time",
                 (releaseTimeMs_2 - acquireTimeMs_1) * 1000, totalTime);
@@ -393,7 +489,9 @@
         final int releaseTimeMs_2 = 1009;
         final int currentTimeMs = 1011;
 
-        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        }
 
         // Create a Uid Object
         final BatteryStats.Uid u = mBatteryStatsImpl.getUidStatsLocked(testUid);
@@ -407,11 +505,11 @@
         u.noteWifiMulticastDisabledLocked(releaseTimeMs_2);
 
         // Get the total acquisition time
-        long totalTime =  u.getWifiMulticastTime(currentTimeMs*1000,
+        long totalTime = u.getWifiMulticastTime(currentTimeMs * 1000,
                 BatteryStats.STATS_SINCE_CHARGED);
         assertEquals("Miscalculations of Multicast wakelock acquisition time",
                 ((releaseTimeMs_1 - acquireTimeMs_1) + (releaseTimeMs_2 - acquireTimeMs_2))
-                * 1000, totalTime);
+                        * 1000, totalTime);
     }
 
     private void addIsolatedUid(int parentUid, int childUid) {
@@ -426,20 +524,6 @@
         }
     }
 
-    private void updateProcessStates(int[] uids, int[] procStates,
-            int[] pendingUidsIdx) {
-        final SparseIntArray pendingUids = mBatteryStatsImpl.getPendingUids();
-        for (int i = 0; i < uids.length; ++i) {
-            final BatteryStatsImpl.Uid u = mBatteryStatsImpl.getUidStatsLocked(uids[i]);
-            if (ArrayUtils.contains(pendingUidsIdx, i)) {
-                u.setProcessStateForTest(PROCESS_STATE_TOP);
-                pendingUids.put(uids[i], procStates[i]);
-            } else {
-                u.setProcessStateForTest(procStates[i]);
-            }
-        }
-    }
-
     private void verifyNoPendingUids() {
         final SparseIntArray pendingUids = mBatteryStatsImpl.getPendingUids();
         assertEquals("There shouldn't be any pending uids left: " + pendingUids,
diff --git a/core/tests/coretests/src/com/android/internal/os/BstatsCpuTimesValidationTest.java b/core/tests/coretests/src/com/android/internal/os/BstatsCpuTimesValidationTest.java
index 63e13fd..441e85d 100644
--- a/core/tests/coretests/src/com/android/internal/os/BstatsCpuTimesValidationTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BstatsCpuTimesValidationTest.java
@@ -183,7 +183,9 @@
         sCpuFreqTimesAvailable = totalCpuTimes != null;
         final long[] fgCpuTimes = getAllCpuFreqTimes(Process.SYSTEM_UID,
                 PROCESS_STATE_FOREGROUND);
-        sPerProcStateTimesAvailable = fgCpuTimes != null;
+        final long[] topCpuTimes = getAllCpuFreqTimes(Process.SYSTEM_UID,
+                PROCESS_STATE_TOP);
+        sPerProcStateTimesAvailable = fgCpuTimes != null || topCpuTimes != null;
     }
 
     @Test
@@ -563,7 +565,7 @@
 
             final long[] cpuTimesMs3 = getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_TOP);
             assertCpuTimesValid(cpuTimesMs3);
-            assertCpuTimesEqual(cpuTimesMs3, cpuTimesMs, 20,
+            assertCpuTimesEqual(cpuTimesMs3, cpuTimesMs, 500,
                     "Unexpected cpu times after turning on tracking");
 
             doSomeWork(PROCESS_STATE_TOP);
@@ -588,7 +590,8 @@
     private void assertCpuTimesEqual(long[] actual, long[] expected, long delta, String errMsg) {
         for (int i = actual.length - 1; i >= 0; --i) {
             if (actual[i] > expected[i] + delta || actual[i] < expected[i]) {
-                fail(errMsg + ", actual=" + actual + ", expected=" + expected + ", delta=" + delta);
+                fail(errMsg + ", actual=" + Arrays.toString(actual)
+                        + ", expected=" + Arrays.toString(expected) + ", delta=" + delta);
             }
         }
     }
diff --git a/core/tests/coretests/src/com/android/internal/os/KernelSingleUidTimeReaderTest.java b/core/tests/coretests/src/com/android/internal/os/KernelSingleUidTimeReaderTest.java
index 4182574..74ab644 100644
--- a/core/tests/coretests/src/com/android/internal/os/KernelSingleUidTimeReaderTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/KernelSingleUidTimeReaderTest.java
@@ -285,7 +285,7 @@
         LongArrayMultiStateCounter counter = new LongArrayMultiStateCounter(2, 5);
         counter.setState(0, 0);
         mInjector.setCpuTimeInStatePerClusterNs(new long[][]{{0, 0, 0}, {0, 0}});
-        boolean success = mInjector.addDelta(TEST_UID, counter, 0);
+        boolean success = mInjector.addDelta(TEST_UID, counter, 0, null);
         assertThat(success).isTrue();
 
         // Nanoseconds
@@ -294,13 +294,16 @@
                         {1_000_000, 2_000_000, 3_000_000},
                         {4_000_000, 5_000_000}});
 
-        success = mInjector.addDelta(TEST_UID, counter, 2000);
-        assertThat(success).isTrue();
-
         LongArrayMultiStateCounter.LongArrayContainer array =
                 new LongArrayMultiStateCounter.LongArrayContainer(5);
         long[] out = new long[5];
 
+        success = mInjector.addDelta(TEST_UID, counter, 2000, array);
+        assertThat(success).isTrue();
+
+        array.getValues(out);
+        assertThat(out).isEqualTo(new long[]{1, 2, 3, 4, 5});
+
         counter.getCounts(array, 0);
         array.getValues(out);
         assertThat(out).isEqualTo(new long[]{1, 2, 3, 4, 5});
@@ -312,9 +315,12 @@
                         {11_000_000, 22_000_000, 33_000_000},
                         {44_000_000, 55_000_000}});
 
-        success = mInjector.addDelta(TEST_UID, counter, 4000);
+        success = mInjector.addDelta(TEST_UID, counter, 4000, array);
         assertThat(success).isTrue();
 
+        array.getValues(out);
+        assertThat(out).isEqualTo(new long[]{10, 20, 30, 40, 50});
+
         counter.getCounts(array, 0);
         array.getValues(out);
         assertThat(out).isEqualTo(new long[]{1 + 5, 2 + 10, 3 + 15, 4 + 20, 5 + 25});
@@ -371,8 +377,10 @@
         }
 
         @Override
-        public boolean addDelta(int uid, LongArrayMultiStateCounter counter, long timestampMs) {
-            return addDeltaForTest(uid, counter, timestampMs, mCpuTimeInStatePerClusterNs);
+        public boolean addDelta(int uid, LongArrayMultiStateCounter counter, long timestampMs,
+                LongArrayMultiStateCounter.LongArrayContainer deltaOut) {
+            return addDeltaForTest(uid, counter, timestampMs, mCpuTimeInStatePerClusterNs,
+                    deltaOut);
         }
     }
 }