diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index 7e286f0423848..a817119a735f9 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -8618,7 +8618,7 @@ public class BatteryStatsImpl extends BatteryStats {
          * inactive so can be dropped.
          */
         @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
-        public boolean reset(long uptimeUs, long realtimeUs) {
+        public boolean reset(long uptimeUs, long realtimeUs, int resetReason) {
             boolean active = false;
 
             mOnBatteryBackgroundTimeBase.init(uptimeUs, realtimeUs);
@@ -8688,7 +8688,11 @@ public class BatteryStatsImpl extends BatteryStats {
             resetIfNotNull(mBluetoothControllerActivity, false, realtimeUs);
             resetIfNotNull(mModemControllerActivity, false, realtimeUs);
 
-            MeasuredEnergyStats.resetIfNotNull(mUidMeasuredEnergyStats);
+            if (resetReason == RESET_REASON_MEASURED_ENERGY_BUCKETS_CHANGE) {
+                mUidMeasuredEnergyStats = null;
+            } else {
+                MeasuredEnergyStats.resetIfNotNull(mUidMeasuredEnergyStats);
+            }
 
             resetIfNotNull(mUserCpuTime, false, realtimeUs);
             resetIfNotNull(mSystemCpuTime, false, realtimeUs);
@@ -11371,7 +11375,7 @@ public class BatteryStatsImpl extends BatteryStats {
         mNumConnectivityChange = 0;
 
         for (int i=0; i<mUidStats.size(); i++) {
-            if (mUidStats.valueAt(i).reset(uptimeUs, elapsedRealtimeUs)) {
+            if (mUidStats.valueAt(i).reset(uptimeUs, elapsedRealtimeUs, resetReason)) {
                 mUidStats.valueAt(i).detachFromTimeBase();
                 mUidStats.remove(mUidStats.keyAt(i));
                 i--;
diff --git a/core/java/com/android/internal/os/CustomMeasuredPowerCalculator.java b/core/java/com/android/internal/os/CustomMeasuredPowerCalculator.java
index 9b51a8ef64107..bb307a0d29d8e 100644
--- a/core/java/com/android/internal/os/CustomMeasuredPowerCalculator.java
+++ b/core/java/com/android/internal/os/CustomMeasuredPowerCalculator.java
@@ -21,13 +21,18 @@ import android.os.BatteryStats;
 import android.os.BatteryUsageStats;
 import android.os.BatteryUsageStatsQuery;
 import android.os.UidBatteryConsumer;
+import android.util.Slog;
 import android.util.SparseArray;
 
+import java.util.Arrays;
+
 /**
  * Calculates the amount of power consumed by custom energy consumers (i.e. consumers of type
  * {@link android.hardware.power.stats.EnergyConsumerType#OTHER}).
  */
 public class CustomMeasuredPowerCalculator extends PowerCalculator {
+    private static final String TAG = "CustomMeasuredPowerCalc";
+
     public CustomMeasuredPowerCalculator(PowerProfile powerProfile) {
     }
 
@@ -76,9 +81,9 @@ public class CustomMeasuredPowerCalculator extends PowerCalculator {
             if (totalPowerMah == null) {
                 newTotalPowerMah = new double[customMeasuredPowerMah.length];
             } else if (totalPowerMah.length != customMeasuredPowerMah.length) {
-                newTotalPowerMah = new double[customMeasuredPowerMah.length];
-                System.arraycopy(totalPowerMah, 0, newTotalPowerMah, 0,
-                        customMeasuredPowerMah.length);
+                Slog.wtf(TAG, "Number of custom energy components is not the same for all apps: "
+                        + totalPowerMah.length + ", " + customMeasuredPowerMah.length);
+                newTotalPowerMah = Arrays.copyOf(totalPowerMah, customMeasuredPowerMah.length);
             } else {
                 newTotalPowerMah = totalPowerMah;
             }
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsSensorTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsSensorTest.java
index b851f0ad34142..0135fe84909af 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsSensorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsSensorTest.java
@@ -403,7 +403,7 @@ public class BatteryStatsSensorTest extends TestCase {
         assertNotNull(sensor.getSensorBackgroundTime());
 
         // Reset the stats. Since the sensor is still running, we should still see the timer
-        bi.getUidStatsLocked(UID).reset(clocks.uptime * 1000, clocks.realtime * 1000);
+        bi.getUidStatsLocked(UID).reset(clocks.uptime * 1000, clocks.realtime * 1000, 0);
 
         sensor = uid.getSensorStats().get(SENSOR_ID);
         assertNotNull(sensor);
@@ -413,7 +413,7 @@ public class BatteryStatsSensorTest extends TestCase {
         bi.noteStopSensorLocked(UID, SENSOR_ID);
 
         // Now the sensor timer has stopped so this reset should also take out the sensor.
-        bi.getUidStatsLocked(UID).reset(clocks.uptime * 1000, clocks.realtime * 1000);
+        bi.getUidStatsLocked(UID).reset(clocks.uptime * 1000, clocks.realtime * 1000, 0);
 
         sensor = uid.getSensorStats().get(SENSOR_ID);
         assertNull(sensor);
@@ -465,7 +465,7 @@ public class BatteryStatsSensorTest extends TestCase {
 
         // Reset the stats. Since the sensor is still running, we should still see the timer
         // but still with 0 times.
-        bi.getUidStatsLocked(UID).reset(clocks.uptime * 1000, clocks.realtime * 1000);
+        bi.getUidStatsLocked(UID).reset(clocks.uptime * 1000, clocks.realtime * 1000, 0);
         assertEquals(0, timer.getTotalTimeLocked(1000*clocks.realtime, which));
         assertEquals(0, timer.getTotalDurationMsLocked(clocks.realtime));
         assertEquals(0, bgTimer.getTotalTimeLocked(1000*clocks.realtime, which));
@@ -504,7 +504,7 @@ public class BatteryStatsSensorTest extends TestCase {
 
         // Reset the stats. Since the sensor is still running, we should still see the timer
         // but with 0 times.
-        bi.getUidStatsLocked(UID).reset(clocks.uptime * 1000, clocks.realtime * 1000);
+        bi.getUidStatsLocked(UID).reset(clocks.uptime * 1000, clocks.realtime * 1000, 0);
         assertEquals(0, timer.getTotalTimeLocked(1000*clocks.realtime, which));
         assertEquals(0, timer.getTotalDurationMsLocked(clocks.realtime));
         assertEquals(0, bgTimer.getTotalTimeLocked(1000*clocks.realtime, which));