diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index c420e6d..9015cbb 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -124,7 +124,7 @@
     private static final int MAGIC = 0xBA757475; // 'BATSTATS'
 
     // Current on-disk Parcel version
-    private static final int VERSION = 172 + (USE_OLD_HISTORY ? 1000 : 0);
+    private static final int VERSION = 173 + (USE_OLD_HISTORY ? 1000 : 0);
 
     // Maximum number of items we will record in the history.
     private static final int MAX_HISTORY_ITEMS;
@@ -289,7 +289,7 @@
     /**
      * Update per-freq cpu times for all the uids in {@link #mPendingUids}.
      */
-    public void updateProcStateCpuTimes() {
+    public void updateProcStateCpuTimes(boolean onBattery, boolean onBatteryScreenOff) {
         final SparseIntArray uidStates;
         synchronized (BatteryStatsImpl.this) {
             if(!initKernelSingleUidTimeReaderLocked()) {
@@ -307,7 +307,6 @@
             final int procState = uidStates.valueAt(i);
             final int[] isolatedUids;
             final Uid u;
-            final boolean onBattery;
             synchronized (BatteryStatsImpl.this) {
                 // It's possible that uid no longer exists and any internal references have
                 // already been deleted, so using {@link #getAvailableUidStatsLocked} to avoid
@@ -324,7 +323,6 @@
                         isolatedUids[j] = u.mChildUids.get(j);
                     }
                 }
-                onBattery = mOnBatteryInternal;
             }
             long[] cpuTimesMs = mKernelSingleUidTimeReader.readDeltaMs(uid);
             if (isolatedUids != null) {
@@ -335,20 +333,27 @@
             }
             if (onBattery && cpuTimesMs != null) {
                 synchronized (BatteryStatsImpl.this) {
-                    u.addProcStateTimesMs(procState, cpuTimesMs);
-                    u.addProcStateScreenOffTimesMs(procState, cpuTimesMs);
+                    u.addProcStateTimesMs(procState, cpuTimesMs, onBattery);
+                    u.addProcStateScreenOffTimesMs(procState, cpuTimesMs, onBatteryScreenOff);
                 }
             }
         }
     }
 
+    public void copyFromAllUidsCpuTimes() {
+        synchronized (BatteryStatsImpl.this) {
+            copyFromAllUidsCpuTimes(
+                    mOnBatteryTimeBase.isRunning(), mOnBatteryScreenOffTimeBase.isRunning());
+        }
+    }
+
     /**
      * When the battery/screen state changes, we don't attribute the cpu times to any process
      * but we still need to snapshots of all uids to get correct deltas later on. Since we
      * already read this data for updating per-freq cpu times, we can use the same data for
      * per-procstate cpu times.
      */
-    public void copyFromAllUidsCpuTimes() {
+    public void copyFromAllUidsCpuTimes(boolean onBattery, boolean onBatteryScreenOff) {
         synchronized (BatteryStatsImpl.this) {
             if(!initKernelSingleUidTimeReaderLocked()) {
                 return;
@@ -368,7 +373,7 @@
                 }
                 final long[] deltaTimesMs = mKernelSingleUidTimeReader.computeDelta(
                         uid, cpuTimesMs.clone());
-                if (mOnBatteryInternal && deltaTimesMs != null) {
+                if (onBattery && deltaTimesMs != null) {
                     final int procState;
                     final int idx = mPendingUids.indexOfKey(uid);
                     if (idx >= 0) {
@@ -378,8 +383,8 @@
                         procState = u.mProcessState;
                     }
                     if (procState >= 0 && procState < Uid.NUM_PROCESS_STATE) {
-                        u.addProcStateTimesMs(procState, deltaTimesMs);
-                        u.addProcStateScreenOffTimesMs(procState, deltaTimesMs);
+                        u.addProcStateTimesMs(procState, deltaTimesMs, onBattery);
+                        u.addProcStateScreenOffTimesMs(procState, deltaTimesMs, onBatteryScreenOff);
                     }
                 }
             }
@@ -443,8 +448,8 @@
 
         Future<?> scheduleSync(String reason, int flags);
         Future<?> scheduleCpuSyncDueToRemovedUid(int uid);
-        Future<?> scheduleReadProcStateCpuTimes();
-        Future<?> scheduleCopyFromAllUidsCpuTimes();
+        Future<?> scheduleReadProcStateCpuTimes(boolean onBattery, boolean onBatteryScreenOff);
+        Future<?> scheduleCopyFromAllUidsCpuTimes(boolean onBattery, boolean onBatteryScreenOff);
     }
 
     public Handler mHandler;
@@ -1231,12 +1236,10 @@
         public long[] mCounts;
         public long[] mLoadedCounts;
         public long[] mUnpluggedCounts;
-        public long[] mPluggedCounts;
 
         private LongSamplingCounterArray(TimeBase timeBase, Parcel in) {
             mTimeBase = timeBase;
-            mPluggedCounts = in.createLongArray();
-            mCounts = copyArray(mPluggedCounts, mCounts);
+            mCounts = in.createLongArray();
             mLoadedCounts = in.createLongArray();
             mUnpluggedCounts = in.createLongArray();
             timeBase.add(this);
@@ -1255,17 +1258,16 @@
 
         @Override
         public void onTimeStarted(long elapsedRealTime, long baseUptime, long baseRealtime) {
-            mUnpluggedCounts = copyArray(mPluggedCounts, mUnpluggedCounts);
+            mUnpluggedCounts = copyArray(mCounts, mUnpluggedCounts);
         }
 
         @Override
         public void onTimeStopped(long elapsedRealtime, long baseUptime, long baseRealtime) {
-            mPluggedCounts = copyArray(mCounts, mPluggedCounts);
         }
 
         @Override
         public long[] getCountsLocked(int which) {
-            long[] val = copyArray(mTimeBase.isRunning() ? mCounts : mPluggedCounts, null);
+            long[] val = copyArray(mCounts, null);
             if (which == STATS_SINCE_UNPLUGGED) {
                 subtract(val, mUnpluggedCounts);
             } else if (which != STATS_SINCE_CHARGED) {
@@ -1278,15 +1280,18 @@
         public void logState(Printer pw, String prefix) {
             pw.println(prefix + "mCounts=" + Arrays.toString(mCounts)
                     + " mLoadedCounts=" + Arrays.toString(mLoadedCounts)
-                    + " mUnpluggedCounts=" + Arrays.toString(mUnpluggedCounts)
-                    + " mPluggedCounts=" + Arrays.toString(mPluggedCounts));
+                    + " mUnpluggedCounts=" + Arrays.toString(mUnpluggedCounts));
         }
 
         public void addCountLocked(long[] counts) {
+            addCountLocked(counts, mTimeBase.isRunning());
+        }
+
+        public void addCountLocked(long[] counts, boolean isRunning) {
             if (counts == null) {
                 return;
             }
-            if (mTimeBase.isRunning()) {
+            if (isRunning) {
                 if (mCounts == null) {
                     mCounts = new long[counts.length];
                 }
@@ -1306,7 +1311,6 @@
         public void reset(boolean detachIfReset) {
             fillArray(mCounts, 0);
             fillArray(mLoadedCounts, 0);
-            fillArray(mPluggedCounts, 0);
             fillArray(mUnpluggedCounts, 0);
             if (detachIfReset) {
                 detach();
@@ -1325,7 +1329,6 @@
             mCounts = in.createLongArray();
             mLoadedCounts = copyArray(mCounts, mLoadedCounts);
             mUnpluggedCounts = copyArray(mCounts, mUnpluggedCounts);
-            mPluggedCounts = copyArray(mCounts, mPluggedCounts);
         }
 
         public static void writeToParcel(Parcel out, LongSamplingCounterArray counterArray) {
@@ -3783,7 +3786,8 @@
                         + " and battery is " + (unplugged ? "on" : "off"));
             }
             updateCpuTimeLocked();
-            mExternalSync.scheduleCopyFromAllUidsCpuTimes();
+            mExternalSync.scheduleCopyFromAllUidsCpuTimes(mOnBatteryTimeBase.isRunning(),
+                    mOnBatteryScreenOffTimeBase.isRunning());
 
             mOnBatteryTimeBase.setRunning(unplugged, uptime, realtime);
             if (updateOnBatteryTimeBase) {
@@ -6648,7 +6652,7 @@
             return null;
         }
 
-        private void addProcStateTimesMs(int procState, long[] cpuTimesMs) {
+        private void addProcStateTimesMs(int procState, long[] cpuTimesMs, boolean onBattery) {
             if (mProcStateTimeMs == null) {
                 mProcStateTimeMs = new LongSamplingCounterArray[NUM_PROCESS_STATE];
             }
@@ -6657,10 +6661,11 @@
                 mProcStateTimeMs[procState] = new LongSamplingCounterArray(
                         mBsi.mOnBatteryTimeBase);
             }
-            mProcStateTimeMs[procState].addCountLocked(cpuTimesMs);
+            mProcStateTimeMs[procState].addCountLocked(cpuTimesMs, onBattery);
         }
 
-        private void addProcStateScreenOffTimesMs(int procState, long[] cpuTimesMs) {
+        private void addProcStateScreenOffTimesMs(int procState, long[] cpuTimesMs,
+                boolean onBatteryScreenOff) {
             if (mProcStateScreenOffTimeMs == null) {
                 mProcStateScreenOffTimeMs = new LongSamplingCounterArray[NUM_PROCESS_STATE];
             }
@@ -6669,7 +6674,7 @@
                 mProcStateScreenOffTimeMs[procState] = new LongSamplingCounterArray(
                         mBsi.mOnBatteryScreenOffTimeBase);
             }
-            mProcStateScreenOffTimeMs[procState].addCountLocked(cpuTimesMs);
+            mProcStateScreenOffTimeMs[procState].addCountLocked(cpuTimesMs, onBatteryScreenOff);
         }
 
         @Override
@@ -9408,7 +9413,9 @@
 
                     if (mBsi.mPerProcStateCpuTimesAvailable) {
                         if (mBsi.mPendingUids.size() == 0) {
-                            mBsi.mExternalSync.scheduleReadProcStateCpuTimes();
+                            mBsi.mExternalSync.scheduleReadProcStateCpuTimes(
+                                    mBsi.mOnBatteryTimeBase.isRunning(),
+                                    mBsi.mOnBatteryScreenOffTimeBase.isRunning());
                         }
                         if (mBsi.mPendingUids.indexOfKey(mUid) < 0
                                 || ArrayUtils.contains(CRITICAL_PROC_STATES, mProcessState)) {
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 3794b5f..a8094ea 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsImplTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsImplTest.java
@@ -71,7 +71,7 @@
     @Test
     public void testUpdateProcStateCpuTimes() {
         mBatteryStatsImpl.setOnBatteryInternal(true);
-        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_ON, 0, 0);
+        mBatteryStatsImpl.updateTimeBasesLocked(false, Display.STATE_ON, 0, 0);
 
         final int[] testUids = {10032, 10048, 10145, 10139};
         final int[] testProcStates = {
@@ -98,7 +98,7 @@
             }
         }
 
-        mBatteryStatsImpl.updateProcStateCpuTimes();
+        mBatteryStatsImpl.updateProcStateCpuTimes(true, false);
 
         verifyNoPendingUids();
         for (int i = 0; i < testUids.length; ++i) {
@@ -125,7 +125,7 @@
         }
         addPendingUids(testUids, testProcStates);
 
-        mBatteryStatsImpl.updateProcStateCpuTimes();
+        mBatteryStatsImpl.updateProcStateCpuTimes(true, false);
 
         verifyNoPendingUids();
         for (int i = 0; i < testUids.length; ++i) {
@@ -157,7 +157,7 @@
         }
         addPendingUids(testUids, testProcStates);
 
-        mBatteryStatsImpl.updateProcStateCpuTimes();
+        mBatteryStatsImpl.updateProcStateCpuTimes(true, true);
 
         verifyNoPendingUids();
         for (int i = 0; i < testUids.length; ++i) {
@@ -196,7 +196,7 @@
         final long[] isolatedUidCpuTimes = {495784, 398473, 4895, 4905, 30984093};
         when(mKernelSingleUidTimeReader.readDeltaMs(childUid)).thenReturn(isolatedUidCpuTimes);
 
-        mBatteryStatsImpl.updateProcStateCpuTimes();
+        mBatteryStatsImpl.updateProcStateCpuTimes(true, true);
 
         verifyNoPendingUids();
         for (int i = 0; i < testUids.length; ++i) {
@@ -227,8 +227,8 @@
 
     @Test
     public void testCopyFromAllUidsCpuTimes() {
-        mBatteryStatsImpl.setOnBatteryInternal(true);
-        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_ON, 0, 0);
+        mBatteryStatsImpl.setOnBatteryInternal(false);
+        mBatteryStatsImpl.updateTimeBasesLocked(false, Display.STATE_ON, 0, 0);
 
         final int[] testUids = {10032, 10048, 10145, 10139};
         final int[] testProcStates = {
@@ -264,7 +264,7 @@
                     .thenReturn(expectedCpuTimes[i]);
         }
 
-        mBatteryStatsImpl.copyFromAllUidsCpuTimes();
+        mBatteryStatsImpl.copyFromAllUidsCpuTimes(true, false);
 
         verifyNoPendingUids();
         for (int i = 0; i < testUids.length; ++i) {
diff --git a/core/tests/coretests/src/com/android/internal/os/LongSamplingCounterArrayTest.java b/core/tests/coretests/src/com/android/internal/os/LongSamplingCounterArrayTest.java
index 27aec56..37b4e41a 100644
--- a/core/tests/coretests/src/com/android/internal/os/LongSamplingCounterArrayTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/LongSamplingCounterArrayTest.java
@@ -16,7 +16,9 @@
 
 package com.android.internal.os;
 
+import static android.os.BatteryStats.STATS_CURRENT;
 import static android.os.BatteryStats.STATS_SINCE_CHARGED;
+import static android.os.BatteryStats.STATS_SINCE_UNPLUGGED;
 
 import static com.android.internal.os.BatteryStatsImpl.LongSamplingCounterArray;
 import static com.android.internal.os.BatteryStatsImpl.TimeBase;
@@ -61,7 +63,6 @@
 
     private static final long[] COUNTS = {1111, 2222, 3333, 4444};
     private static final long[] LOADED_COUNTS = {5555, 6666, 7777, 8888};
-    private static final long[] PLUGGED_COUNTS = {9999, 11111, 22222, 33333};
     private static final long[] UNPLUGGED_COUNTS = {44444, 55555, 66666, 77777};
     private static final long[] ZEROES = {0, 0, 0, 0};
 
@@ -83,11 +84,10 @@
         parcel.setDataPosition(0);
 
         // Now clear counterArray and verify values are read from parcel correctly.
-        updateCounts(null, null, null, null);
+        updateCounts(null, null, null);
         mCounterArray = LongSamplingCounterArray.readFromParcel(parcel, mTimeBase);
         assertArrayEquals(COUNTS, mCounterArray.mCounts, "Unexpected counts");
         assertArrayEquals(LOADED_COUNTS, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
-        assertArrayEquals(COUNTS, mCounterArray.mPluggedCounts, "Unexpected pluggedCounts");
         assertArrayEquals(UNPLUGGED_COUNTS, mCounterArray.mUnpluggedCounts,
                 "Unexpected unpluggedCounts");
         parcel.recycle();
@@ -101,11 +101,10 @@
         parcel.setDataPosition(0);
 
         // Now clear counterArray and verify values are read from parcel correctly.
-        updateCounts(null, null, null, null);
+        updateCounts(null, null, null);
         mCounterArray = LongSamplingCounterArray.readSummaryFromParcelLocked(parcel, mTimeBase);
         assertArrayEquals(COUNTS, mCounterArray.mCounts, "Unexpected counts");
         assertArrayEquals(COUNTS, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
-        assertArrayEquals(COUNTS, mCounterArray.mPluggedCounts, "Unexpected pluggedCounts");
         assertArrayEquals(COUNTS, mCounterArray.mUnpluggedCounts, "Unexpected unpluggedCounts");
         parcel.recycle();
     }
@@ -116,8 +115,7 @@
         mCounterArray.onTimeStarted(0, 0, 0);
         assertArrayEquals(COUNTS, mCounterArray.mCounts, "Unexpected counts");
         assertArrayEquals(LOADED_COUNTS, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
-        assertArrayEquals(PLUGGED_COUNTS, mCounterArray.mPluggedCounts, "Unexpected pluggedCounts");
-        assertArrayEquals(PLUGGED_COUNTS, mCounterArray.mUnpluggedCounts,
+        assertArrayEquals(COUNTS, mCounterArray.mUnpluggedCounts,
                 "Unexpected unpluggedCounts");
     }
 
@@ -127,7 +125,6 @@
         mCounterArray.onTimeStopped(0, 0, 0);
         assertArrayEquals(COUNTS, mCounterArray.mCounts, "Unexpected counts");
         assertArrayEquals(LOADED_COUNTS, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
-        assertArrayEquals(COUNTS, mCounterArray.mPluggedCounts, "Unexpected pluggedCounts");
         assertArrayEquals(UNPLUGGED_COUNTS, mCounterArray.mUnpluggedCounts,
                 "Unexpected unpluggedCounts");
     }
@@ -137,24 +134,50 @@
         initializeCounterArrayWithDefaultValues();
 
         when(mTimeBase.isRunning()).thenReturn(false);
-        long[] actualVal = mCounterArray.getCountsLocked(STATS_SINCE_CHARGED);
-        long[] expectedVal = PLUGGED_COUNTS;
-        assertArrayEquals(expectedVal, actualVal, "Unexpected values");
+        assertArrayEquals(COUNTS, mCounterArray.getCountsLocked(STATS_SINCE_CHARGED),
+                "Unexpected values");
+        assertArrayEquals(subtract(COUNTS, LOADED_COUNTS),
+                mCounterArray.getCountsLocked(STATS_CURRENT), "Unexpected values");
+        assertArrayEquals(subtract(COUNTS, UNPLUGGED_COUNTS),
+                mCounterArray.getCountsLocked(STATS_SINCE_UNPLUGGED), "Unexpected values");
 
         when(mTimeBase.isRunning()).thenReturn(true);
-        actualVal = mCounterArray.getCountsLocked(STATS_SINCE_CHARGED);
-        expectedVal = COUNTS;
-        assertArrayEquals(expectedVal, actualVal, "Unexpected values");
+        assertArrayEquals(COUNTS, mCounterArray.getCountsLocked(STATS_SINCE_CHARGED),
+                "Unexpected values");
+        assertArrayEquals(subtract(COUNTS, LOADED_COUNTS),
+                mCounterArray.getCountsLocked(STATS_CURRENT), "Unexpected values");
+        assertArrayEquals(subtract(COUNTS, UNPLUGGED_COUNTS),
+                mCounterArray.getCountsLocked(STATS_SINCE_UNPLUGGED), "Unexpected values");
+    }
+
+    private long[] subtract(long[] val, long[] toSubtract) {
+        final long[] result = val.clone();
+        if (toSubtract != null) {
+            for (int i = val.length - 1; i >= 0; --i) {
+                result[i] -= toSubtract[i];
+            }
+        }
+        return result;
     }
 
     @Test
     public void testAddCountLocked() {
+        updateCounts(null, null, null);
         final long[] deltas = {123, 234, 345, 456};
         when(mTimeBase.isRunning()).thenReturn(true);
         mCounterArray.addCountLocked(deltas);
         assertArrayEquals(deltas, mCounterArray.mCounts, "Unexpected counts");
         assertArrayEquals(null, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
-        assertArrayEquals(null, mCounterArray.mPluggedCounts, "Unexpected pluggedCounts");
+        assertArrayEquals(null, mCounterArray.mUnpluggedCounts, "Unexpected unpluggedCounts");
+
+        updateCounts(null, null, null);
+        mCounterArray.addCountLocked(deltas, false);
+        assertArrayEquals(null, mCounterArray.mCounts, "Unexpected counts");
+        assertArrayEquals(null, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
+        assertArrayEquals(null, mCounterArray.mUnpluggedCounts, "Unexpected unpluggedCounts");
+        mCounterArray.addCountLocked(deltas, true);
+        assertArrayEquals(deltas, mCounterArray.mCounts, "Unexpected counts");
+        assertArrayEquals(null, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
         assertArrayEquals(null, mCounterArray.mUnpluggedCounts, "Unexpected unpluggedCounts");
 
         initializeCounterArrayWithDefaultValues();
@@ -165,7 +188,18 @@
         mCounterArray.addCountLocked(deltas);
         assertArrayEquals(newCounts, mCounterArray.mCounts, "Unexpected counts");
         assertArrayEquals(LOADED_COUNTS, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
-        assertArrayEquals(PLUGGED_COUNTS, mCounterArray.mPluggedCounts, "Unexpected pluggedCounts");
+        assertArrayEquals(UNPLUGGED_COUNTS, mCounterArray.mUnpluggedCounts,
+                "Unexpected unpluggedCounts");
+
+        initializeCounterArrayWithDefaultValues();
+        mCounterArray.addCountLocked(deltas, false);
+        assertArrayEquals(COUNTS, mCounterArray.mCounts, "Unexpected counts");
+        assertArrayEquals(LOADED_COUNTS, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
+        assertArrayEquals(UNPLUGGED_COUNTS, mCounterArray.mUnpluggedCounts,
+                "Unexpected unpluggedCounts");
+        mCounterArray.addCountLocked(deltas, true);
+        assertArrayEquals(newCounts, mCounterArray.mCounts, "Unexpected counts");
+        assertArrayEquals(LOADED_COUNTS, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
         assertArrayEquals(UNPLUGGED_COUNTS, mCounterArray.mUnpluggedCounts,
                 "Unexpected unpluggedCounts");
     }
@@ -177,7 +211,6 @@
         mCounterArray.reset(false /* detachIfReset */);
         assertArrayEquals(ZEROES, mCounterArray.mCounts, "Unexpected counts");
         assertArrayEquals(ZEROES, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
-        assertArrayEquals(ZEROES, mCounterArray.mPluggedCounts, "Unexpected pluggedCounts");
         assertArrayEquals(ZEROES, mCounterArray.mUnpluggedCounts, "Unexpected unpluggedCounts");
         verifyZeroInteractions(mTimeBase);
 
@@ -186,7 +219,6 @@
         mCounterArray.reset(true /* detachIfReset */);
         assertArrayEquals(ZEROES, mCounterArray.mCounts, "Unexpected counts");
         assertArrayEquals(ZEROES, mCounterArray.mLoadedCounts, "Unexpected loadedCounts");
-        assertArrayEquals(ZEROES, mCounterArray.mPluggedCounts, "Unexpected pluggedCounts");
         assertArrayEquals(ZEROES, mCounterArray.mUnpluggedCounts, "Unexpected unpluggedCounts");
         verify(mTimeBase).remove(mCounterArray);
         verifyNoMoreInteractions(mTimeBase);
@@ -200,7 +232,7 @@
     }
 
     private void initializeCounterArrayWithDefaultValues() {
-        updateCounts(COUNTS, LOADED_COUNTS, PLUGGED_COUNTS, UNPLUGGED_COUNTS);
+        updateCounts(COUNTS, LOADED_COUNTS, UNPLUGGED_COUNTS);
     }
 
     private void assertArrayEquals(long[] expected, long[] actual, String msg) {
@@ -208,11 +240,9 @@
                 + ", actual: " + Arrays.toString(actual), Arrays.equals(expected, actual));
     }
 
-    private void updateCounts(long[] counts, long[] loadedCounts,
-            long[] pluggedCounts, long[] unpluggedCounts) {
-        mCounterArray.mCounts = counts;
-        mCounterArray.mLoadedCounts = loadedCounts;
-        mCounterArray.mPluggedCounts = pluggedCounts;
-        mCounterArray.mUnpluggedCounts = unpluggedCounts;
+    private void updateCounts(long[] counts, long[] loadedCounts, long[] unpluggedCounts) {
+        mCounterArray.mCounts = counts == null ? null : counts.clone();
+        mCounterArray.mLoadedCounts = loadedCounts == null ? null : loadedCounts.clone();
+        mCounterArray.mUnpluggedCounts = unpluggedCounts == null ? null : unpluggedCounts.clone();
     }
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java b/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java
index f19ff67..43b41a0 100644
--- a/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java
+++ b/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java
@@ -139,12 +139,14 @@
         }
 
         @Override
-        public Future<?> scheduleReadProcStateCpuTimes() {
+        public Future<?> scheduleReadProcStateCpuTimes(
+                boolean onBattery, boolean onBatteryScreenOff) {
             return null;
         }
 
         @Override
-        public Future<?> scheduleCopyFromAllUidsCpuTimes() {
+        public Future<?> scheduleCopyFromAllUidsCpuTimes(
+                boolean onBattery, boolean onBatteryScreenOff) {
             return null;
         }
 
diff --git a/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java b/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
index 4582430..28b8edf 100644
--- a/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
+++ b/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
@@ -35,6 +35,7 @@
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.os.BatteryStatsImpl;
+import com.android.internal.util.function.pooled.PooledLambda;
 
 import libcore.util.EmptyArray;
 
@@ -117,7 +118,7 @@
     }
 
     @Override
-    public Future<?> scheduleReadProcStateCpuTimes() {
+    public Future<?> scheduleReadProcStateCpuTimes(boolean onBattery, boolean onBatteryScreenOff) {
         synchronized (mStats) {
             if (!mStats.mPerProcStateCpuTimesAvailable) {
                 return null;
@@ -125,14 +126,17 @@
         }
         synchronized (BatteryExternalStatsWorker.this) {
             if (!mExecutorService.isShutdown()) {
-                return mExecutorService.submit(mReadProcStateCpuTimesTask);
+                return mExecutorService.submit(PooledLambda.obtainRunnable(
+                        BatteryStatsImpl::updateProcStateCpuTimes,
+                        mStats, onBattery, onBatteryScreenOff).recycleOnUse());
             }
         }
         return null;
     }
 
     @Override
-    public Future<?> scheduleCopyFromAllUidsCpuTimes() {
+    public Future<?> scheduleCopyFromAllUidsCpuTimes(
+            boolean onBattery, boolean onBatteryScreenOff) {
         synchronized (mStats) {
             if (!mStats.mPerProcStateCpuTimesAvailable) {
                 return null;
@@ -140,26 +144,14 @@
         }
         synchronized (BatteryExternalStatsWorker.this) {
             if (!mExecutorService.isShutdown()) {
-                return mExecutorService.submit(mCopyFromAllUidsCpuTimesTask);
+                return mExecutorService.submit(PooledLambda.obtainRunnable(
+                        BatteryStatsImpl::copyFromAllUidsCpuTimes,
+                        mStats, onBattery, onBatteryScreenOff).recycleOnUse());
             }
         }
         return null;
     }
 
-    private final Runnable mReadProcStateCpuTimesTask = new Runnable() {
-        @Override
-        public void run() {
-            mStats.updateProcStateCpuTimes();
-        }
-    };
-
-    private final Runnable mCopyFromAllUidsCpuTimesTask = new Runnable() {
-        @Override
-        public void run() {
-            mStats.copyFromAllUidsCpuTimes();
-        }
-    };
-
     public synchronized Future<?> scheduleWrite() {
         if (mExecutorService.isShutdown()) {
             return CompletableFuture.failedFuture(new IllegalStateException("worker shutdown"));