Fix concurrency in BATTERY_USAGE_STATS_* statsd atom generation

Bug: 229906525
Test: atest --rerun-until-failure 100 FrameworksServicesTests:BatteryStatsManagerTest
Change-Id: Ib2d8ab6df40776c7cd2b5852bd956a3f9285b486
(cherry picked from commit 2a510708a117c9efe9102a623988c695aecc6fcb)
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 7fe3c1f..606a09cb 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -798,6 +798,7 @@
                 case FrameworkStatsLog.BATTERY_USAGE_STATS_SINCE_RESET:
                     final BatteryUsageStatsQuery querySinceReset =
                             new BatteryUsageStatsQuery.Builder()
+                                    .setMaxStatsAgeMs(0)
                                     .includeProcessStateData()
                                     .includeVirtualUids()
                                     .build();
@@ -806,6 +807,7 @@
                 case FrameworkStatsLog.BATTERY_USAGE_STATS_SINCE_RESET_USING_POWER_PROFILE_MODEL:
                     final BatteryUsageStatsQuery queryPowerProfile =
                             new BatteryUsageStatsQuery.Builder()
+                                    .setMaxStatsAgeMs(0)
                                     .includeProcessStateData()
                                     .includeVirtualUids()
                                     .powerProfileModeledOnly()
@@ -822,6 +824,7 @@
                     final long sessionEnd = mStats.getStartClockTime();
                     final BatteryUsageStatsQuery queryBeforeReset =
                             new BatteryUsageStatsQuery.Builder()
+                                    .setMaxStatsAgeMs(0)
                                     .includeProcessStateData()
                                     .includeVirtualUids()
                                     .aggregateSnapshots(sessionStart, sessionEnd)
diff --git a/services/tests/servicestests/AndroidManifest.xml b/services/tests/servicestests/AndroidManifest.xml
index 0afb182..0483a60 100644
--- a/services/tests/servicestests/AndroidManifest.xml
+++ b/services/tests/servicestests/AndroidManifest.xml
@@ -103,6 +103,7 @@
     <uses-permission android:name="android.permission.MODIFY_AUDIO_ROUTING" />
     <uses-permission android:name="android.permission.PACKAGE_VERIFICATION_AGENT" />
     <uses-permission android:name="android.permission.OBSERVE_ROLE_HOLDERS" />
+    <uses-permission android:name="android.permission.BATTERY_STATS" />
 
     <queries>
         <package android:name="com.android.servicestests.apps.suspendtestapp" />
diff --git a/services/tests/servicestests/src/com/android/server/power/stats/BatteryStatsManagerTest.java b/services/tests/servicestests/src/com/android/server/power/stats/BatteryStatsManagerTest.java
new file mode 100644
index 0000000..7ae1117
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/power/stats/BatteryStatsManagerTest.java
@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import static androidx.test.InstrumentationRegistry.getContext;
+
+import static org.junit.Assert.fail;
+
+import android.os.BatteryConsumer;
+import android.os.BatteryStatsManager;
+import android.os.BatteryUsageStats;
+import android.os.BatteryUsageStatsQuery;
+import android.os.UidBatteryConsumer;
+
+import org.junit.Test;
+
+/**
+ * Test BatteryStatsManager and CellularBatteryStats to ensure that valid data is being reported
+ * and that invalid data is not reported.
+ */
+public class BatteryStatsManagerTest {
+
+    @Test
+    public void testBatteryUsageStatsDataConsistency() {
+        BatteryStatsManager bsm = getContext().getSystemService(BatteryStatsManager.class);
+        BatteryUsageStats stats = bsm.getBatteryUsageStats(
+                new BatteryUsageStatsQuery.Builder().setMaxStatsAgeMs(
+                        0).includeProcessStateData().build());
+        final int[] components =
+                {BatteryConsumer.POWER_COMPONENT_CPU,
+                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
+                        BatteryConsumer.POWER_COMPONENT_WIFI,
+                        BatteryConsumer.POWER_COMPONENT_BLUETOOTH};
+        final int[] states =
+                {BatteryConsumer.PROCESS_STATE_FOREGROUND,
+                        BatteryConsumer.PROCESS_STATE_BACKGROUND,
+                        BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE,
+                        BatteryConsumer.PROCESS_STATE_CACHED};
+        for (UidBatteryConsumer ubc : stats.getUidBatteryConsumers()) {
+            for (int component : components) {
+                double consumedPower = ubc.getConsumedPower(ubc.getKey(component));
+                double sumStates = 0;
+                for (int state : states) {
+                    sumStates += ubc.getConsumedPower(ubc.getKey(component, state));
+                }
+                if (sumStates > consumedPower + 0.1) {
+                    fail("Sum of states exceeds total. UID = " + ubc.getUid() + " "
+                            + BatteryConsumer.powerComponentIdToString(component)
+                            + " total = " + consumedPower + " states = " + sumStates);
+                }
+            }
+        }
+    }
+}