Merge "Consistently close BatteryUsageStats objects" into main
diff --git a/core/java/android/os/BatteryUsageStats.java b/core/java/android/os/BatteryUsageStats.java
index ddcb5c6..6c3c285 100644
--- a/core/java/android/os/BatteryUsageStats.java
+++ b/core/java/android/os/BatteryUsageStats.java
@@ -24,6 +24,8 @@
 import android.util.SparseArray;
 import android.util.proto.ProtoOutputStream;
 
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.os.BatteryStatsHistory;
 import com.android.internal.os.BatteryStatsHistoryIterator;
 import com.android.internal.os.MonotonicClock;
@@ -43,7 +45,9 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Contains a snapshot of battery attribution data, on a per-subsystem and per-UID basis.
@@ -126,6 +130,12 @@
     // Max window size. CursorWindow uses only as much memory as needed.
     private static final long BATTERY_CONSUMER_CURSOR_WINDOW_SIZE = 20_000_000; // bytes
 
+    /**
+     * Used by tests to ensure all BatteryUsageStats instances are closed.
+     */
+    @VisibleForTesting
+    public static boolean DEBUG_INSTANCE_COUNT;
+
     private static final int STATSD_PULL_ATOM_MAX_BYTES = 45000;
 
     private static final int[] UID_USAGE_TIME_PROCESS_STATES = {
@@ -153,7 +163,7 @@
     private final List<UserBatteryConsumer> mUserBatteryConsumers;
     private final AggregateBatteryConsumer[] mAggregateBatteryConsumers;
     private final BatteryStatsHistory mBatteryStatsHistory;
-    private BatteryConsumer.BatteryConsumerDataLayout mBatteryConsumerDataLayout;
+    private final BatteryConsumer.BatteryConsumerDataLayout mBatteryConsumerDataLayout;
     private CursorWindow mBatteryConsumersCursorWindow;
 
     private BatteryUsageStats(@NonNull Builder builder) {
@@ -873,6 +883,7 @@
 
     @Override
     public void close() throws IOException {
+        onCursorWindowReleased(mBatteryConsumersCursorWindow);
         mBatteryConsumersCursorWindow.close();
         mBatteryConsumersCursorWindow = null;
     }
@@ -880,6 +891,7 @@
     @Override
     protected void finalize() throws Throwable {
         if (mBatteryConsumersCursorWindow != null) {
+            // Do not decrement sOpenCusorWindowCount. All instances should be closed explicitly
             mBatteryConsumersCursorWindow.close();
         }
         super.finalize();
@@ -934,6 +946,7 @@
                 boolean includesPowerStateData, double minConsumedPowerThreshold) {
             mBatteryConsumersCursorWindow =
                     new CursorWindow(null, BATTERY_CONSUMER_CURSOR_WINDOW_SIZE);
+            onCursorWindowAllocated(mBatteryConsumersCursorWindow);
             mBatteryConsumerDataLayout = BatteryConsumer.createBatteryConsumerDataLayout(
                     customPowerComponentNames, includePowerModels, includeProcessStateData,
                     includeScreenStateData, includesPowerStateData);
@@ -996,6 +1009,7 @@
          */
         public void discard() {
             mBatteryConsumersCursorWindow.close();
+            onCursorWindowReleased(mBatteryConsumersCursorWindow);
         }
 
         /**
@@ -1264,4 +1278,50 @@
             }
         }
     }
+
+    @GuardedBy("BatteryUsageStats.class")
+    private static Map<CursorWindow, Exception> sInstances;
+
+    private static void onCursorWindowAllocated(CursorWindow window) {
+        if (!DEBUG_INSTANCE_COUNT) {
+            return;
+        }
+
+        synchronized (BatteryUsageStats.class) {
+            if (sInstances == null) {
+                sInstances = new HashMap<>();
+            }
+            sInstances.put(window, new Exception());
+        }
+    }
+
+    private static void onCursorWindowReleased(CursorWindow window) {
+        if (!DEBUG_INSTANCE_COUNT) {
+            return;
+        }
+
+        synchronized (BatteryUsageStats.class) {
+            sInstances.remove(window);
+        }
+    }
+
+    /**
+     * Used by tests to ensure all BatteryUsageStats instances are closed.
+     */
+    @VisibleForTesting
+    public static void assertAllInstancesClosed() {
+        if (!DEBUG_INSTANCE_COUNT) {
+            throw new IllegalStateException("DEBUG_INSTANCE_COUNT is false");
+        }
+
+        synchronized (BatteryUsageStats.class) {
+            if (!sInstances.isEmpty()) {
+                Exception callSite = sInstances.entrySet().iterator().next().getValue();
+                int count = sInstances.size();
+                sInstances.clear();
+                throw new IllegalStateException(
+                        "Instances of BatteryUsageStats not closed: " + count, callSite);
+            }
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index c47cad9..28b606c 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -3178,12 +3178,15 @@
             mStats.collectPowerStatsSamples();
         }
 
-        BatteryUsageStats batteryUsageStats =
-                mBatteryUsageStatsProvider.getBatteryUsageStats(mStats, query);
-        if (proto) {
-            batteryUsageStats.dumpToProto(fd);
-        } else {
-            batteryUsageStats.dump(pw, "  ");
+        try (BatteryUsageStats batteryUsageStats =
+                     mBatteryUsageStatsProvider.getBatteryUsageStats(mStats, query)) {
+            if (proto) {
+                batteryUsageStats.dumpToProto(fd);
+            } else {
+                batteryUsageStats.dump(pw, "  ");
+            }
+        } catch (IOException e) {
+            Slog.e(TAG, "Cannot close BatteryUsageStats", e);
         }
     }
 
diff --git a/services/core/java/com/android/server/power/stats/AccumulatedBatteryUsageStatsSection.java b/services/core/java/com/android/server/power/stats/AccumulatedBatteryUsageStatsSection.java
index dd6d5db..c0a06fc 100644
--- a/services/core/java/com/android/server/power/stats/AccumulatedBatteryUsageStatsSection.java
+++ b/services/core/java/com/android/server/power/stats/AccumulatedBatteryUsageStatsSection.java
@@ -51,6 +51,11 @@
         mBatteryUsageStats.build().dump(ipw, "");
     }
 
+    @Override
+    public void close() {
+        mBatteryUsageStats.discard();
+    }
+
     static class Reader implements PowerStatsSpan.SectionReader {
         @Override
         public String getType() {
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
index 265f1dfc..b996c43 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -295,7 +295,8 @@
                     stats.builder = ((AccumulatedBatteryUsageStatsSection) section)
                             .getBatteryUsageStatsBuilder();
                     stats.startWallClockTime = powerStatsSpan.getMetadata().getStartTime();
-                    stats.startMonotonicTime = powerStatsSpan.getMetadata().getStartMonotonicTime();
+                    stats.startMonotonicTime =
+                            powerStatsSpan.getMetadata().getStartMonotonicTime();
                     stats.endMonotonicTime = powerStatsSpan.getMetadata().getEndMonotonicTime();
                     break;
                 }
@@ -484,29 +485,30 @@
                 continue;
             }
 
-            PowerStatsSpan powerStatsSpan = mPowerStatsStore.loadPowerStatsSpan(
-                    spanMetadata.getId(), BatteryUsageStatsSection.TYPE);
-            if (powerStatsSpan == null) {
-                continue;
-            }
-
-            for (PowerStatsSpan.Section section : powerStatsSpan.getSections()) {
-                BatteryUsageStats snapshot =
-                        ((BatteryUsageStatsSection) section).getBatteryUsageStats();
-                if (!Arrays.equals(snapshot.getCustomPowerComponentNames(),
-                        customEnergyConsumerNames)) {
-                    Log.w(TAG, "Ignoring older BatteryUsageStats snapshot, which has different "
-                            + "custom power components: "
-                            + Arrays.toString(snapshot.getCustomPowerComponentNames()));
+            try (PowerStatsSpan powerStatsSpan = mPowerStatsStore.loadPowerStatsSpan(
+                    spanMetadata.getId(), BatteryUsageStatsSection.TYPE)) {
+                if (powerStatsSpan == null) {
                     continue;
                 }
 
-                if (includeProcessStateData && !snapshot.isProcessStateDataIncluded()) {
-                    Log.w(TAG, "Ignoring older BatteryUsageStats snapshot, which "
-                            + " does not include process state data");
-                    continue;
+                for (PowerStatsSpan.Section section : powerStatsSpan.getSections()) {
+                    BatteryUsageStats snapshot =
+                            ((BatteryUsageStatsSection) section).getBatteryUsageStats();
+                    if (!Arrays.equals(snapshot.getCustomPowerComponentNames(),
+                            customEnergyConsumerNames)) {
+                        Log.w(TAG, "Ignoring older BatteryUsageStats snapshot, which has different "
+                                + "custom power components: "
+                                + Arrays.toString(snapshot.getCustomPowerComponentNames()));
+                        continue;
+                    }
+
+                    if (includeProcessStateData && !snapshot.isProcessStateDataIncluded()) {
+                        Log.w(TAG, "Ignoring older BatteryUsageStats snapshot, which "
+                                + " does not include process state data");
+                        continue;
+                    }
+                    builder.add(snapshot);
                 }
-                builder.add(snapshot);
             }
         }
         return builder.build();
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsSection.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsSection.java
index af36524..eb896e9 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsSection.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsSection.java
@@ -18,6 +18,7 @@
 
 import android.os.BatteryUsageStats;
 import android.util.IndentingPrintWriter;
+import android.util.Slog;
 
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
@@ -28,6 +29,7 @@
 
 class BatteryUsageStatsSection extends PowerStatsSpan.Section {
     public static final String TYPE = "battery-usage-stats";
+    private static final String TAG = "BatteryUsageStatsSection";
 
     private final BatteryUsageStats mBatteryUsageStats;
 
@@ -50,6 +52,15 @@
         mBatteryUsageStats.dump(ipw, "");
     }
 
+    @Override
+    public void close() {
+        try {
+            mBatteryUsageStats.close();
+        } catch (IOException e) {
+            Slog.e(TAG, "Closing BatteryUsageStats", e);
+        }
+    }
+
     static class Reader implements PowerStatsSpan.SectionReader {
         @Override
         public String getType() {
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsSpan.java b/services/core/java/com/android/server/power/stats/PowerStatsSpan.java
index 5105272..4f560cf 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsSpan.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsSpan.java
@@ -51,7 +51,7 @@
 /**
  * Contains power stats of various kinds, aggregated over a time span.
  */
-public class PowerStatsSpan {
+public class PowerStatsSpan implements AutoCloseable {
     private static final String TAG = "PowerStatsStore";
 
     /**
@@ -321,6 +321,13 @@
         public void dump(IndentingPrintWriter ipw) {
             ipw.println(mType);
         }
+
+        /**
+         * Closes the section, releasing any resources it held. Once closed, the Section
+         * should not be used.
+         */
+        public void close() {
+        }
     }
 
     /**
@@ -484,4 +491,10 @@
             ipw.decreaseIndent();
         }
     }
+    @Override
+    public void close() {
+        for (int i = 0; i < mSections.size(); i++) {
+            mSections.get(i).close();
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsStore.java b/services/core/java/com/android/server/power/stats/PowerStatsStore.java
index 3673617..d83d355 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsStore.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsStore.java
@@ -331,9 +331,10 @@
         ipw.increaseIndent();
         List<PowerStatsSpan.Metadata> contents = getTableOfContents();
         for (PowerStatsSpan.Metadata metadata : contents) {
-            PowerStatsSpan span = loadPowerStatsSpan(metadata.getId());
-            if (span != null) {
-                span.dump(ipw);
+            try (PowerStatsSpan span = loadPowerStatsSpan(metadata.getId())) {
+                if (span != null) {
+                    span.dump(ipw);
+                }
             }
         }
         ipw.decreaseIndent();
diff --git a/services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java b/services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java
index c8170a1..b2442c8 100644
--- a/services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java
+++ b/services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java
@@ -105,19 +105,19 @@
                     maxEndTime = spanMaxTime;
                 }
 
-                PowerStatsSpan span = mPowerStatsStore.loadPowerStatsSpan(metadata.getId(),
-                        AggregatedPowerStatsSection.TYPE);
-                if (span == null) {
-                    Slog.e(TAG, "Could not read PowerStatsStore section " + metadata);
-                    continue;
-                }
-                List<PowerStatsSpan.Section> sections = span.getSections();
-                for (int k = 0; k < sections.size(); k++) {
-                    hasStoredSpans = true;
-                    PowerStatsSpan.Section section = sections.get(k);
-                    populateBatteryUsageStatsBuilder(batteryUsageStatsBuilder,
-                            ((AggregatedPowerStatsSection) section).getAggregatedPowerStats());
-                    // TODO(b/371614748): close the builder
+                try (PowerStatsSpan span = mPowerStatsStore.loadPowerStatsSpan(metadata.getId(),
+                        AggregatedPowerStatsSection.TYPE)) {
+                    if (span == null) {
+                        Slog.e(TAG, "Could not read PowerStatsStore section " + metadata);
+                        continue;
+                    }
+                    List<PowerStatsSpan.Section> sections = span.getSections();
+                    for (int k = 0; k < sections.size(); k++) {
+                        hasStoredSpans = true;
+                        PowerStatsSpan.Section section = sections.get(k);
+                        populateBatteryUsageStatsBuilder(batteryUsageStatsBuilder,
+                                ((AggregatedPowerStatsSection) section).getAggregatedPowerStats());
+                    }
                 }
             }
 
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
index 177f30a..0a1fc31 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
@@ -987,5 +987,7 @@
                         BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
                 .getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
                 .isEqualTo(60000);
+
+        span.close();
     }
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsManagerTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsManagerTest.java
index 9a64ce1..94997b2 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsManagerTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsManagerTest.java
@@ -41,7 +41,7 @@
     public final RavenwoodRule mRavenwood = new RavenwoodRule();
 
     @Test
-    public void testBatteryUsageStatsDataConsistency() {
+    public void testBatteryUsageStatsDataConsistency() throws Exception {
         BatteryStatsManager bsm = getContext().getSystemService(BatteryStatsManager.class);
         BatteryUsageStats stats = bsm.getBatteryUsageStats(
                 new BatteryUsageStatsQuery.Builder().setMaxStatsAgeMs(
@@ -70,5 +70,7 @@
                 }
             }
         }
+
+        stats.close();
     }
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
index 7d2bc17..813dd84 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
@@ -64,7 +64,7 @@
     private static final int UID_3 = 4000;
 
     @Test
-    public void testAtom_BatteryUsageStatsPerUid() {
+    public void testAtom_BatteryUsageStatsPerUid() throws Exception {
         final BatteryUsageStats bus = buildBatteryUsageStats();
         BatteryStatsService.FrameworkStatsLogger statsLogger =
                 mock(BatteryStatsService.FrameworkStatsLogger.class);
@@ -72,6 +72,8 @@
         List<StatsEvent> actual = new ArrayList<>();
         new BatteryStatsService.StatsPerUidLogger(statsLogger).logStats(bus, actual);
 
+        bus.close();
+
         if (DEBUG) {
             System.out.println(mockingDetails(statsLogger).printInvocations());
         }
@@ -284,7 +286,7 @@
     }
 
     @Test
-    public void testAtom_BatteryUsageStatsAtomsProto() {
+    public void testAtom_BatteryUsageStatsAtomsProto() throws Exception {
         final BatteryUsageStats bus = buildBatteryUsageStats();
         final byte[] bytes = bus.getStatsProto();
         BatteryUsageStatsAtomsProto proto;
@@ -347,6 +349,7 @@
 
         // UID_3 - Should be none, since no interesting data (done last for debugging convenience).
         assertEquals(3, proto.uidBatteryConsumers.length);
+        bus.close();
     }
 
     private void assertSameBatteryConsumer(String message, BatteryConsumer consumer,
@@ -582,7 +585,7 @@
     }
 
     @Test
-    public void testLargeAtomTruncated() {
+    public void testLargeAtomTruncated() throws Exception {
         final BatteryUsageStats.Builder builder =
                 new BatteryUsageStats.Builder(new String[0], true, false, false, false, 0);
         // If not truncated, this BatteryUsageStats object would generate a proto buffer
@@ -618,6 +621,8 @@
         assertThat(bytes.length).isGreaterThan(20000);
         assertThat(bytes.length).isLessThan(50000);
 
+        batteryUsageStats.close();
+
         BatteryUsageStatsAtomsProto proto;
         try {
             proto = BatteryUsageStatsAtomsProto.parseFrom(bytes);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
index e9d95fc..87a26d0 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
@@ -94,7 +94,7 @@
     }
 
     @Test
-    public void test_getBatteryUsageStats() {
+    public void test_getBatteryUsageStats() throws IOException {
         final BatteryUsageStats batteryUsageStats = prepareBatteryUsageStats(false);
 
         final List<UidBatteryConsumer> uidBatteryConsumers =
@@ -121,24 +121,27 @@
 
         assertThat(batteryUsageStats.getStatsStartTimestamp()).isEqualTo(12345);
         assertThat(batteryUsageStats.getStatsEndTimestamp()).isEqualTo(180 * MINUTE_IN_MS);
+        batteryUsageStats.close();
     }
 
     @Test
-    public void batteryLevelInfo_charging() {
+    public void batteryLevelInfo_charging() throws IOException {
         final BatteryUsageStats batteryUsageStats = prepareBatteryUsageStats(true);
         assertThat(batteryUsageStats.getBatteryCapacity()).isEqualTo(4000.0);
         assertThat(batteryUsageStats.getChargeTimeRemainingMs()).isEqualTo(1_200_000);
+        batteryUsageStats.close();
     }
 
     @Test
-    public void batteryLevelInfo_onBattery() {
+    public void batteryLevelInfo_onBattery() throws IOException {
         final BatteryUsageStats batteryUsageStats = prepareBatteryUsageStats(false);
         assertThat(batteryUsageStats.getBatteryCapacity()).isEqualTo(4000.0);
         assertThat(batteryUsageStats.getBatteryTimeRemainingMs()).isEqualTo(600_000);
+        batteryUsageStats.close();
     }
 
     @Test
-    public void test_selectPowerComponents() {
+    public void test_selectPowerComponents() throws IOException {
         BatteryStatsImpl batteryStats = prepareBatteryStats(false);
 
         BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
@@ -163,6 +166,8 @@
         assertThat(
                 uidBatteryConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
                 .isEqualTo(0);
+
+        batteryUsageStats.close();
     }
 
     private BatteryStatsImpl prepareBatteryStats(boolean plugInAtTheEnd) {
@@ -274,7 +279,7 @@
     }
 
     @Test
-    public void testWriteAndReadHistory() {
+    public void testWriteAndReadHistory() throws IOException {
         MockBatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
         synchronized (batteryStats) {
             batteryStats.setRecordAllHistoryLocked(true);
@@ -306,6 +311,8 @@
 
         Parcel in = Parcel.obtain();
         batteryUsageStats.writeToParcel(in, 0);
+        batteryUsageStats.close();
+
         final byte[] bytes = in.marshall();
 
         Parcel out = Parcel.obtain();
@@ -349,10 +356,12 @@
 
         assertThat(iterator.hasNext()).isFalse();
         assertThat(iterator.next()).isNull();
+
+        unparceled.close();
     }
 
     @Test
-    public void testWriteAndReadHistoryTags() {
+    public void testWriteAndReadHistoryTags() throws IOException {
         MockBatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
         synchronized (batteryStats) {
             batteryStats.setRecordAllHistoryLocked(true);
@@ -400,6 +409,8 @@
             assertThat(parcel.dataSize()).isAtMost(128_000);
         }
 
+        batteryUsageStats.close();
+
         parcel.setDataPosition(0);
 
         BatteryUsageStats unparceled = parcel.readParcelable(getClass().getClassLoader(),
@@ -461,7 +472,7 @@
     }
 
     @Test
-    public void testAggregateBatteryStats() {
+    public void testAggregateBatteryStats() throws IOException {
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
 
         setTime(5 * MINUTE_IN_MS);
@@ -563,6 +574,8 @@
                 .getConsumedPower(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
                 .isWithin(0.1)
                 .of(180.0);
+
+        stats.close();
     }
 
     @Test
@@ -689,6 +702,8 @@
                 .isEqualTo(60 * MINUTE_IN_MS);
 
         assertThat(count[0]).isEqualTo(expectedNumberOfUpdates);
+
+        stats.close();
     }
 
     private void setTime(long timeMs) {
@@ -733,6 +748,7 @@
                     .isWithin(PRECISION).of(8.33333);
             assertThat(uid.getConsumedPower(componentId1))
                     .isWithin(PRECISION).of(8.33333);
+            stats.close();
             return null;
         }).when(powerStatsStore).storeBatteryUsageStatsAsync(anyLong(), any());
 
@@ -750,7 +766,7 @@
     }
 
     @Test
-    public void testAggregateBatteryStats_incompatibleSnapshot() {
+    public void testAggregateBatteryStats_incompatibleSnapshot() throws IOException {
         MockBatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
         batteryStats.initMeasuredEnergyStats(new String[]{"FOO", "BAR"});
 
@@ -776,6 +792,8 @@
         when(powerStatsStore.loadPowerStatsSpan(1, BatteryUsageStatsSection.TYPE))
                 .thenReturn(span1);
 
+        span1.close();
+
         BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
                 mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
                 mStatsRule.getCpuScalingPolicies(), powerStatsStore, 0, mMockClock);
@@ -787,5 +805,7 @@
         assertThat(stats.getCustomPowerComponentNames())
                 .isEqualTo(batteryStats.getCustomEnergyConsumerNames());
         assertThat(stats.getStatsDuration()).isEqualTo(1234);
+
+        stats.close();
     }
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
index 52675f6..383616e 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
@@ -323,6 +323,7 @@
     }
 
     private void before() {
+        BatteryUsageStats.DEBUG_INSTANCE_COUNT = true;
         HandlerThread bgThread = new HandlerThread("bg thread");
         bgThread.setUncaughtExceptionHandler((thread, throwable)-> {
             mThrowable = throwable;
@@ -338,6 +339,10 @@
 
     private void after() throws Throwable {
         waitForBackgroundThread();
+        if (mBatteryUsageStats != null) {
+            mBatteryUsageStats.close();
+        }
+        BatteryUsageStats.assertAllInstancesClosed();
     }
 
     public void waitForBackgroundThread() throws Throwable {
@@ -409,6 +414,14 @@
     }
 
     BatteryUsageStats apply(BatteryUsageStatsQuery query, PowerCalculator... calculators) {
+        if (mBatteryUsageStats != null) {
+            try {
+                mBatteryUsageStats.close();
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+            mBatteryUsageStats = null;
+        }
         final String[] customPowerComponentNames = mBatteryStats.getCustomEnergyConsumerNames();
         final boolean includePowerModels = (query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_POWER_MODELS) != 0;
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsTest.java
index 88624f2..1b6b8c4 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsTest.java
@@ -76,22 +76,25 @@
     private static final int APP_UID2 = 314;
 
     @Test
-    public void testBuilder() {
+    public void testBuilder() throws Exception {
         BatteryUsageStats batteryUsageStats = buildBatteryUsageStats1(true).build();
         assertBatteryUsageStats1(batteryUsageStats, true);
+        batteryUsageStats.close();
     }
 
     @Test
-    public void testBuilder_noProcessStateData() {
+    public void testBuilder_noProcessStateData() throws Exception {
         BatteryUsageStats batteryUsageStats = buildBatteryUsageStats1(false).build();
         assertBatteryUsageStats1(batteryUsageStats, false);
+        batteryUsageStats.close();
     }
 
     @Test
-    public void testParcelability_smallNumberOfUids() {
+    public void testParcelability_smallNumberOfUids() throws Exception {
         final BatteryUsageStats outBatteryUsageStats = buildBatteryUsageStats1(true).build();
         final Parcel parcel = Parcel.obtain();
         parcel.writeParcelable(outBatteryUsageStats, 0);
+        outBatteryUsageStats.close();
 
         assertThat(parcel.dataSize()).isLessThan(100000);
 
@@ -101,10 +104,11 @@
                 parcel.readParcelable(getClass().getClassLoader());
         assertThat(inBatteryUsageStats).isNotNull();
         assertBatteryUsageStats1(inBatteryUsageStats, true);
+        inBatteryUsageStats.close();
     }
 
     @Test
-    public void testParcelability_largeNumberOfUids() {
+    public void testParcelability_largeNumberOfUids() throws Exception {
         final BatteryUsageStats.Builder builder =
                 new BatteryUsageStats.Builder(new String[0]);
 
@@ -141,22 +145,27 @@
             assertThat(uidBatteryConsumer).isNotNull();
             assertThat(uidBatteryConsumer.getConsumedPower()).isEqualTo(i * 100);
         }
+        inBatteryUsageStats.close();
+        outBatteryUsageStats.close();
     }
 
     @Test
-    public void testDefaultSessionDuration() {
+    public void testDefaultSessionDuration() throws Exception {
         final BatteryUsageStats stats =
                 buildBatteryUsageStats1(true).setStatsDuration(10000).build();
         assertThat(stats.getStatsDuration()).isEqualTo(10000);
+        stats.close();
     }
 
     @Test
-    public void testDump() {
+    public void testDump() throws Exception {
         final BatteryUsageStats stats = buildBatteryUsageStats1(true).build();
         final StringWriter out = new StringWriter();
         try (PrintWriter pw = new PrintWriter(out)) {
             stats.dump(pw, "  ");
         }
+        stats.close();
+
         final String dump = out.toString();
 
         assertThat(dump).contains("Capacity: 4000");
@@ -187,12 +196,14 @@
     }
 
     @Test
-    public void testDumpNoScreenOrPowerState() {
+    public void testDumpNoScreenOrPowerState() throws Exception {
         final BatteryUsageStats stats = buildBatteryUsageStats1(true, false, false).build();
         final StringWriter out = new StringWriter();
         try (PrintWriter pw = new PrintWriter(out)) {
             stats.dump(pw, "  ");
         }
+        stats.close();
+
         final String dump = out.toString();
 
         assertThat(dump).contains("Capacity: 4000");
@@ -222,7 +233,7 @@
     }
 
     @Test
-    public void testAdd() {
+    public void testAdd() throws Exception {
         final BatteryUsageStats stats1 = buildBatteryUsageStats1(false).build();
         final BatteryUsageStats stats2 = buildBatteryUsageStats2(new String[]{"FOO"}, true).build();
         final BatteryUsageStats sum =
@@ -261,24 +272,31 @@
         assertAggregateBatteryConsumer(sum,
                 BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE,
                 40211, 40422, 40633, 40844);
+        stats1.close();
+        stats2.close();
+        sum.close();
     }
 
     @Test
-    public void testAdd_customComponentMismatch() {
+    public void testAdd_customComponentMismatch() throws Exception {
         final BatteryUsageStats.Builder builder =
                 new BatteryUsageStats.Builder(new String[]{"FOO"}, true, true, true, true, 0);
         final BatteryUsageStats stats = buildBatteryUsageStats2(new String[]{"BAR"}, false).build();
 
         assertThrows(IllegalArgumentException.class, () -> builder.add(stats));
+        stats.close();
+        builder.discard();
     }
 
     @Test
-    public void testAdd_processStateDataMismatch() {
+    public void testAdd_processStateDataMismatch() throws Exception {
         final BatteryUsageStats.Builder builder =
                 new BatteryUsageStats.Builder(new String[]{"FOO"}, true, true, true, true, 0);
         final BatteryUsageStats stats = buildBatteryUsageStats2(new String[]{"FOO"}, false).build();
 
         assertThrows(IllegalArgumentException.class, () -> builder.add(stats));
+        stats.close();
+        builder.discard();
     }
 
     @Test
@@ -290,12 +308,14 @@
         final BatteryUsageStats stats = buildBatteryUsageStats1(true).build();
         stats.writeXml(serializer);
         serializer.endDocument();
+        stats.close();
 
         ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
         TypedXmlPullParser parser = Xml.newBinaryPullParser();
         parser.setInput(in, StandardCharsets.UTF_8.name());
         final BatteryUsageStats fromXml = BatteryUsageStats.createFromXml(parser);
         assertBatteryUsageStats1(fromXml, true);
+        fromXml.close();
     }
 
     private BatteryUsageStats.Builder buildBatteryUsageStats1(boolean includeUserBatteryConsumer) {