Release wakelock only when no alarms are in-flight

Fix the bug where the alarm wakelock could be released before all alarms
completed.

Test: atest FrameworksMockingServicesTests:com.android.server.alarm

Flag: com.android.server.alarm.acquire_wakelock_before_send

Fixes: 400058421
Change-Id: I86210548b9a8fc6050577ca1f7e973b127c47862

2 files changed, 41 insertions, 4 deletions

diff --git a/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java b/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
index 251776e907d8..44e4999ccf44 100644
--- a/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
+++ b/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
@@ -5369,7 +5369,9 @@ public class AlarmManagerService extends SystemService {
                         // to do any wakelock or stats tracking, so we have nothing
                         // left to do here but go on to the next thing.
                         mSendFinishCount++;
-                        if (Flags.acquireWakelockBeforeSend()) {
+                        if (Flags.acquireWakelockBeforeSend() && mBroadcastRefCount == 0) {
+                            // No other alarms are in-flight and this dispatch failed. We will
+                            // acquire the wakelock again before the next dispatch.
                             mWakeLock.release();
                         }
                         return;
@@ -5409,7 +5411,9 @@ public class AlarmManagerService extends SystemService {
                         // stats management to do.  It threw before we posted the delayed
                         // timeout message, so we're done here.
                         mListenerFinishCount++;
-                        if (Flags.acquireWakelockBeforeSend()) {
+                        if (Flags.acquireWakelockBeforeSend() && mBroadcastRefCount == 0) {
+                            // No other alarms are in-flight and this dispatch failed. We will
+                            // acquire the wakelock again before the next dispatch.
                             mWakeLock.release();
                         }
                         return;
diff --git a/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmManagerServiceTest.java b/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmManagerServiceTest.java
index 2a513ae3a8e8..d79d88400cf9 100644
--- a/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmManagerServiceTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmManagerServiceTest.java
@@ -953,11 +953,13 @@ public final class AlarmManagerServiceTest {
 
     @Test
     @EnableFlags(Flags.FLAG_ACQUIRE_WAKELOCK_BEFORE_SEND)
-    public void testWakelockOrdering() throws Exception {
+    public void testWakelockOrderingFirstAlarm() throws Exception {
         final long triggerTime = mNowElapsedTest + 5000;
         final PendingIntent alarmPi = getNewMockPendingIntent();
         setTestAlarm(ELAPSED_REALTIME_WAKEUP, triggerTime, alarmPi);
 
+        // Pretend that it is the first alarm in this batch, or no other alarms are still processing
+        mService.mBroadcastRefCount = 0;
         mNowElapsedTest = mTestTimer.getElapsed();
         mTestTimer.expire();
 
@@ -975,20 +977,51 @@ public final class AlarmManagerServiceTest {
 
     @Test
     @EnableFlags(Flags.FLAG_ACQUIRE_WAKELOCK_BEFORE_SEND)
-    public void testWakelockReleasedWhenSendFails() throws Exception {
+    public void testWakelockOrderingNonFirst() throws Exception {
         final long triggerTime = mNowElapsedTest + 5000;
         final PendingIntent alarmPi = getNewMockPendingIntent();
         setTestAlarm(ELAPSED_REALTIME_WAKEUP, triggerTime, alarmPi);
 
+        // Pretend that some previous alarms are still processing.
+        mService.mBroadcastRefCount = 3;
+        mNowElapsedTest = mTestTimer.getElapsed();
+        mTestTimer.expire();
+
+        final ArgumentCaptor<PendingIntent.OnFinished> onFinishedCaptor =
+                ArgumentCaptor.forClass(PendingIntent.OnFinished.class);
+        verify(alarmPi).send(eq(mMockContext), eq(0), any(Intent.class), onFinishedCaptor.capture(),
+                any(Handler.class), isNull(), any());
+        onFinishedCaptor.getValue().onSendFinished(alarmPi, null, 0, null, null);
+
+        verify(mWakeLock, never()).acquire();
+        verify(mWakeLock, never()).release();
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ACQUIRE_WAKELOCK_BEFORE_SEND)
+    public void testWakelockReleasedWhenSendFails() throws Exception {
+        final PendingIntent alarmPi = getNewMockPendingIntent();
         doThrow(new PendingIntent.CanceledException("test")).when(alarmPi).send(eq(mMockContext),
                 eq(0), any(Intent.class), any(), any(Handler.class), isNull(), any());
 
+        setTestAlarm(ELAPSED_REALTIME_WAKEUP, mNowElapsedTest + 5000, alarmPi);
+
+        // Pretend that it is the first alarm in this batch, or no other alarms are still processing
+        mService.mBroadcastRefCount = 0;
         mNowElapsedTest = mTestTimer.getElapsed();
         mTestTimer.expire();
 
         final InOrder inOrder = Mockito.inOrder(mWakeLock);
         inOrder.verify(mWakeLock).acquire();
         inOrder.verify(mWakeLock).release();
+
+        setTestAlarm(ELAPSED_REALTIME_WAKEUP, mNowElapsedTest + 5000, alarmPi);
+
+        // Pretend that some previous alarms are still processing.
+        mService.mBroadcastRefCount = 4;
+        mNowElapsedTest = mTestTimer.getElapsed();
+        mTestTimer.expire();
+        inOrder.verifyNoMoreInteractions();
     }
 
     @Test