Clear out battery stats for removed user more quickly

Bug: 229159621
Bug: 244349060
Test: atest com.android.internal.os.BatteryStatsUserLifecycleTests#testNoCpuDataForRemovedUser
Change-Id: I69cb8f80e26e6103753e5b8da6f222f7cd87a9d9
diff --git a/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java b/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
index 2fbf3fb..751f535 100644
--- a/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
+++ b/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
@@ -79,7 +79,12 @@
     private static final long MAX_WIFI_STATS_SAMPLE_ERROR_MILLIS = 750;
 
     // Delay for clearing out battery stats for UIDs corresponding to a removed user
-    public static final int UID_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS = 10_000;
+    public static final int UID_QUICK_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS = 2_000;
+
+    // Delay for the _final_ clean-up of battery stats after a user removal - just in case
+    // some UIDs took longer than UID_QUICK_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS to
+    // stop running.
+    public static final int UID_FINAL_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS = 10_000;
 
     private final ScheduledExecutorService mExecutorService =
             Executors.newSingleThreadScheduledExecutor(
@@ -336,11 +341,20 @@
     @Override
     public Future<?> scheduleCleanupDueToRemovedUser(int userId) {
         synchronized (BatteryExternalStatsWorker.this) {
+            // Initial quick clean-up after a user removal
+            mExecutorService.schedule(() -> {
+                synchronized (mStats) {
+                    mStats.clearRemovedUserUidsLocked(userId);
+                }
+            }, UID_QUICK_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS, TimeUnit.MILLISECONDS);
+
+            // Final clean-up after a user removal, to take care of UIDs that were running longer
+            // than expected
             return mExecutorService.schedule(() -> {
                 synchronized (mStats) {
                     mStats.clearRemovedUserUidsLocked(userId);
                 }
-            }, UID_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS, TimeUnit.MILLISECONDS);
+            }, UID_FINAL_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS, TimeUnit.MILLISECONDS);
         }
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java b/services/tests/servicestests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java
index cff3bf8..b27ba88 100644
--- a/services/tests/servicestests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java
+++ b/services/tests/servicestests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java
@@ -38,7 +38,6 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.BeforeClass;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -51,8 +50,8 @@
 
     private static final long POLL_INTERVAL_MS = 500;
     private static final long USER_REMOVE_TIMEOUT_MS = 5_000;
-    private static final long STOP_USER_TIMEOUT_MS = 10_000;
-    private static final long USER_UIDS_REMOVE_TIMEOUT_MS = 15_000;
+    private static final long STOP_USER_TIMEOUT_MS = 20_000;
+    private static final long USER_UIDS_REMOVE_TIMEOUT_MS = 20_000;
     private static final long BATTERYSTATS_POLLING_TIMEOUT_MS = 5_000;
 
     private static final String CPU_DATA_TAG = "cpu";
@@ -79,26 +78,29 @@
         batteryOnScreenOff();
     }
 
-    @Ignore("b/244349060")
     @Test
     public void testNoCpuDataForRemovedUser() throws Exception {
         mIam.startUserInBackground(mTestUserId);
         waitUntilTrue("No uids for started user " + mTestUserId,
                 () -> getNumberOfUidsInBatteryStats() > 0, BATTERYSTATS_POLLING_TIMEOUT_MS);
 
+        final boolean[] userStopped = new boolean[1];
         CountDownLatch stopUserLatch = new CountDownLatch(1);
         mIam.stopUser(mTestUserId, true, new IStopUserCallback.Stub() {
             @Override
             public void userStopped(int userId) throws RemoteException {
+                userStopped[0] = true;
                 stopUserLatch.countDown();
             }
 
             @Override
             public void userStopAborted(int userId) throws RemoteException {
+                stopUserLatch.countDown();
             }
         });
-        assertTrue("User " + mTestUserId + " could not be stopped",
+        assertTrue("User " + mTestUserId + " could not be stopped in " + STOP_USER_TIMEOUT_MS,
                 stopUserLatch.await(STOP_USER_TIMEOUT_MS, TimeUnit.MILLISECONDS));
+        assertTrue("User " + mTestUserId + " could not be stopped", userStopped[0]);
 
         mUm.removeUser(mTestUserId);
         waitUntilTrue("Unable to remove user " + mTestUserId, () -> {