Decrement while removing from secondary lists

Many times alarms are moved to secondary lists like
mPendingWhileIdleAlarms. A cancel or a remove request will then remove
these alarms from here, in which case we need to decrement the alarm
counts appropriately.
Also, stubbing the handler in the test to fix flaky behaviour.

Test: atest CtsAlarmManagerTestCases
atest FrameworksMockingServicesTests:AlarmManagerServiceTest

Bug: 129995049
Bug: 130376760
Change-Id: I353f2b55cfd9566945b27b0807689a9d3e7c50fc

2 files changed, 67 insertions, 22 deletions

diff --git a/services/core/java/com/android/server/AlarmManagerService.java b/services/core/java/com/android/server/AlarmManagerService.java
index 1fa62cec6e30..1220e82485e7 100644
--- a/services/core/java/com/android/server/AlarmManagerService.java
+++ b/services/core/java/com/android/server/AlarmManagerService.java
@@ -760,7 +760,7 @@ class AlarmManagerService extends SystemService {
                 if (predicate.test(alarm)) {
                     alarms.remove(i);
                     if (!reOrdering) {
-                        decrementAlarmCount(alarm.uid);
+                        decrementAlarmCount(alarm.uid, 1);
                     }
                     didRemove = true;
                     if (alarm.alarmClock != null) {
@@ -1764,7 +1764,7 @@ class AlarmManagerService extends SystemService {
                                 + ", callingPackage: " + callingPackage;
                 // STOPSHIP (b/128866264): Just to catch breakages. Remove before final release.
                 Slog.wtf(TAG, errorMsg);
-                // TODO b/129995049: Resume throwing once issue is resolved.
+                // TODO b/129995049: Resume throwing after some soak time without errors
                 // throw new UnsupportedOperationException(errorMsg);
             }
             setImplLocked(type, triggerAtTime, triggerElapsed, windowLength, maxElapsed,
@@ -3113,17 +3113,21 @@ class AlarmManagerService extends SystemService {
             }
         }
         for (int i = mPendingWhileIdleAlarms.size() - 1; i >= 0; i--) {
-            if (mPendingWhileIdleAlarms.get(i).matches(operation, directReceiver)) {
+            final Alarm alarm = mPendingWhileIdleAlarms.get(i);
+            if (alarm.matches(operation, directReceiver)) {
                 // Don't set didRemove, since this doesn't impact the scheduled alarms.
                 mPendingWhileIdleAlarms.remove(i);
+                decrementAlarmCount(alarm.uid, 1);
             }
         }
         for (int i = mPendingBackgroundAlarms.size() - 1; i >= 0; i--) {
             final ArrayList<Alarm> alarmsForUid = mPendingBackgroundAlarms.valueAt(i);
             for (int j = alarmsForUid.size() - 1; j >= 0; j--) {
-                if (alarmsForUid.get(j).matches(operation, directReceiver)) {
+                final Alarm alarm = alarmsForUid.get(j);
+                if (alarm.matches(operation, directReceiver)) {
                     // Don't set didRemove, since this doesn't impact the scheduled alarms.
                     alarmsForUid.remove(j);
+                    decrementAlarmCount(alarm.uid, 1);
                 }
             }
             if (alarmsForUid.size() == 0) {
@@ -3169,6 +3173,7 @@ class AlarmManagerService extends SystemService {
             if (a.uid == uid) {
                 // Don't set didRemove, since this doesn't impact the scheduled alarms.
                 mPendingWhileIdleAlarms.remove(i);
+                decrementAlarmCount(uid, 1);
             }
         }
         for (int i = mPendingBackgroundAlarms.size() - 1; i >= 0; i --) {
@@ -3176,6 +3181,7 @@ class AlarmManagerService extends SystemService {
             for (int j = alarmsForUid.size() - 1; j >= 0; j--) {
                 if (alarmsForUid.get(j).uid == uid) {
                     alarmsForUid.remove(j);
+                    decrementAlarmCount(uid, 1);
                 }
             }
             if (alarmsForUid.size() == 0) {
@@ -3221,13 +3227,16 @@ class AlarmManagerService extends SystemService {
             if (a.matches(packageName)) {
                 // Don't set didRemove, since this doesn't impact the scheduled alarms.
                 mPendingWhileIdleAlarms.remove(i);
+                decrementAlarmCount(a.uid, 1);
             }
         }
         for (int i = mPendingBackgroundAlarms.size() - 1; i >= 0; i --) {
             final ArrayList<Alarm> alarmsForUid = mPendingBackgroundAlarms.valueAt(i);
             for (int j = alarmsForUid.size() - 1; j >= 0; j--) {
-                if (alarmsForUid.get(j).matches(packageName)) {
+                final Alarm alarm = alarmsForUid.get(j);
+                if (alarm.matches(packageName)) {
                     alarmsForUid.remove(j);
+                    decrementAlarmCount(alarm.uid, 1);
                 }
             }
             if (alarmsForUid.size() == 0) {
@@ -3272,10 +3281,15 @@ class AlarmManagerService extends SystemService {
             if (a.uid == uid) {
                 // Don't set didRemove, since this doesn't impact the scheduled alarms.
                 mPendingWhileIdleAlarms.remove(i);
+                decrementAlarmCount(uid, 1);
             }
         }
         for (int i = mPendingBackgroundAlarms.size() - 1; i >= 0; i--) {
             if (mPendingBackgroundAlarms.keyAt(i) == uid) {
+                final ArrayList<Alarm> toRemove = mPendingBackgroundAlarms.valueAt(i);
+                if (toRemove != null) {
+                    decrementAlarmCount(uid, toRemove.size());
+                }
                 mPendingBackgroundAlarms.removeAt(i);
             }
         }
@@ -3308,11 +3322,18 @@ class AlarmManagerService extends SystemService {
             if (UserHandle.getUserId(mPendingWhileIdleAlarms.get(i).creatorUid)
                     == userHandle) {
                 // Don't set didRemove, since this doesn't impact the scheduled alarms.
-                mPendingWhileIdleAlarms.remove(i);
+                final Alarm removed = mPendingWhileIdleAlarms.remove(i);
+                decrementAlarmCount(removed.uid, 1);
             }
         }
         for (int i = mPendingBackgroundAlarms.size() - 1; i >= 0; i--) {
             if (UserHandle.getUserId(mPendingBackgroundAlarms.keyAt(i)) == userHandle) {
+                final ArrayList<Alarm> toRemove = mPendingBackgroundAlarms.valueAt(i);
+                if (toRemove != null) {
+                    for (int j = 0; j < toRemove.size(); j++) {
+                        decrementAlarmCount(toRemove.get(j).uid, 1);
+                    }
+                }
                 mPendingBackgroundAlarms.removeAt(i);
             }
         }
@@ -3844,7 +3865,7 @@ class AlarmManagerService extends SystemService {
                 Slog.w(TAG, "Failure sending alarm.", e);
             }
             Trace.traceEnd(Trace.TRACE_TAG_POWER);
-            decrementAlarmCount(alarm.uid);
+            decrementAlarmCount(alarm.uid, 1);
         }
     }
 
@@ -4198,7 +4219,7 @@ class AlarmManagerService extends SystemService {
                                 removeImpl(alarm.operation, null);
                             }
                         }
-                        decrementAlarmCount(alarm.uid);
+                        decrementAlarmCount(alarm.uid, 1);
                     }
                     break;
                 }
@@ -4828,16 +4849,21 @@ class AlarmManagerService extends SystemService {
         }
     }
 
-    private void decrementAlarmCount(int uid) {
+    private void decrementAlarmCount(int uid, int decrement) {
+        int oldCount = 0;
         final int uidIndex = mAlarmsPerUid.indexOfKey(uid);
         if (uidIndex >= 0) {
-            final int newCount = mAlarmsPerUid.valueAt(uidIndex) - 1;
-            if (newCount > 0) {
-                mAlarmsPerUid.setValueAt(uidIndex, newCount);
+            oldCount = mAlarmsPerUid.valueAt(uidIndex);
+            if (oldCount > decrement) {
+                mAlarmsPerUid.setValueAt(uidIndex, oldCount - decrement);
             } else {
                 mAlarmsPerUid.removeAt(uidIndex);
             }
         }
+        if (oldCount < decrement) {
+            Slog.wtf(TAG, "Attempt to decrement existing alarm count " + oldCount + " by "
+                    + decrement + " for uid " + uid);
+        }
     }
 
     private class ShellCmd extends ShellCommand {
diff --git a/services/tests/mockingservicestests/src/com/android/server/AlarmManagerServiceTest.java b/services/tests/mockingservicestests/src/com/android/server/AlarmManagerServiceTest.java
index c0a11b27c65c..6517303842ed 100644
--- a/services/tests/mockingservicestests/src/com/android/server/AlarmManagerServiceTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/AlarmManagerServiceTest.java
@@ -263,6 +263,9 @@ public class AlarmManagerServiceTest {
 
         mService.onStart();
         spyOn(mService.mHandler);
+        // Stubbing the handler. Test should simulate any handling of messages synchronously.
+        doReturn(true).when(mService.mHandler).sendMessageAtTime(any(Message.class), anyLong());
+
         assertEquals(mService.mSystemUiUid, SYSTEM_UI_UID);
         assertEquals(mService.mClockReceiver, mClockReceiver);
         assertEquals(mService.mWakeLock, mWakeLock);
@@ -617,11 +620,9 @@ public class AlarmManagerServiceTest {
         testQuotasNoDeferral(STANDBY_BUCKET_RARE);
     }
 
-    private void sendAndHandleBucketChanged(int bucket) {
+    private void assertAndHandleBucketChanged(int bucket) {
         when(mUsageStatsManagerInternal.getAppStandbyBucket(eq(TEST_CALLING_PACKAGE), anyInt(),
                 anyLong())).thenReturn(bucket);
-        // Stubbing the handler call to simulate it synchronously here.
-        doReturn(true).when(mService.mHandler).sendMessage(any(Message.class));
         mAppStandbyListener.onAppIdleStateChanged(TEST_CALLING_PACKAGE,
                 UserHandle.getUserId(TEST_CALLING_UID), false, bucket, 0);
         final ArgumentCaptor<Message> messageCaptor = ArgumentCaptor.forClass(Message.class);
@@ -652,7 +653,7 @@ public class AlarmManagerServiceTest {
         // The next upcoming alarm in queue should also be set as expected.
         assertEquals(firstTrigger + workingQuota - 1, mTestTimer.getElapsed());
         // Downgrading the bucket now
-        sendAndHandleBucketChanged(STANDBY_BUCKET_RARE);
+        assertAndHandleBucketChanged(STANDBY_BUCKET_RARE);
         final int rareQuota = mService.getQuotaForBucketLocked(STANDBY_BUCKET_RARE);
         // The last alarm should now be deferred.
         final long expectedNextTrigger = (firstTrigger + workingQuota - 1 - rareQuota)
@@ -680,15 +681,13 @@ public class AlarmManagerServiceTest {
         assertEquals(deferredTrigger, mTestTimer.getElapsed());
 
         // Upgrading the bucket now
-        sendAndHandleBucketChanged(STANDBY_BUCKET_ACTIVE);
+        assertAndHandleBucketChanged(STANDBY_BUCKET_ACTIVE);
         // The last alarm should now be rescheduled to go as per original expectations
         final long originalTrigger = firstTrigger + frequentQuota;
         assertEquals("Incorrect next alarm trigger", originalTrigger, mTestTimer.getElapsed());
     }
 
-    private void sendAndHandleParoleChanged(boolean parole) {
-        // Stubbing the handler call to simulate it synchronously here.
-        doReturn(true).when(mService.mHandler).sendMessage(any(Message.class));
+    private void assertAndHandleParoleChanged(boolean parole) {
         mAppStandbyListener.onParoleStateChanged(parole);
         final ArgumentCaptor<Message> messageCaptor = ArgumentCaptor.forClass(Message.class);
         verify(mService.mHandler, atLeastOnce()).sendMessage(messageCaptor.capture());
@@ -719,7 +718,7 @@ public class AlarmManagerServiceTest {
         // Any subsequent alarms in queue should all be deferred
         assertEquals(firstTrigger + mAppStandbyWindow + 1, mTestTimer.getElapsed());
         // Paroling now
-        sendAndHandleParoleChanged(true);
+        assertAndHandleParoleChanged(true);
 
         // Subsequent alarms should now go off as per original expectations.
         for (int i = 0; i < 5; i++) {
@@ -728,7 +727,7 @@ public class AlarmManagerServiceTest {
             mTestTimer.expire();
         }
         // Come out of parole
-        sendAndHandleParoleChanged(false);
+        assertAndHandleParoleChanged(false);
 
         // Subsequent alarms should again get deferred
         final long expectedNextTrigger = (firstTrigger + 5) + 1 + mAppStandbyWindow;
@@ -938,6 +937,26 @@ public class AlarmManagerServiceTest {
     }
 
     @Test
+    public void alarmCountOnRemoveFromPendingWhileIdle() {
+        mService.mPendingIdleUntil = mock(AlarmManagerService.Alarm.class);
+        final int numAlarms = 15;
+        final PendingIntent[] pis = new PendingIntent[numAlarms];
+        for (int i = 0; i < numAlarms; i++) {
+            pis[i] = getNewMockPendingIntent();
+            setTestAlarm(ELAPSED_REALTIME, mNowElapsedTest + i + 5, pis[i]);
+        }
+        assertEquals(numAlarms, mService.mAlarmsPerUid.get(TEST_CALLING_UID));
+        assertEquals(numAlarms, mService.mPendingWhileIdleAlarms.size());
+        final int toRemove = 8;
+        for (int i = 0; i < toRemove; i++) {
+            mService.removeLocked(pis[i], null);
+            assertEquals(numAlarms - i - 1, mService.mAlarmsPerUid.get(TEST_CALLING_UID, 0));
+        }
+        mService.removeLocked(TEST_CALLING_UID);
+        assertEquals(0, mService.mAlarmsPerUid.get(TEST_CALLING_UID, 0));
+    }
+
+    @Test
     public void alarmCountOnAlarmRemoved() {
         final int numAlarms = 10;
         final PendingIntent[] pis = new PendingIntent[numAlarms];