Acquire wakelock before alarm dispatch

This is required to avoid stalls in the app if a suspend operation is
already under-way for some reason. The expectation is that the userspace
wakelock acquisition will either block or succeed, in both cases
guaranteeing that the app executes continously while processing the
alarm.

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

Flag: com.android.server.alarm.acquire_wakelock_before_send

Bug: 391413964
Change-Id: I56a0739a60aa58139d40e8d64523328d5430b6ec

2 files changed, 29 insertions, 5 deletions

diff --git a/apex/jobscheduler/service/aconfig/alarm.aconfig b/apex/jobscheduler/service/aconfig/alarm.aconfig
index a6e980726a9a..9181ef0c532a 100644
--- a/apex/jobscheduler/service/aconfig/alarm.aconfig
+++ b/apex/jobscheduler/service/aconfig/alarm.aconfig
@@ -7,3 +7,13 @@ flag {
     description: "Persist list of users with alarms scheduled and wakeup stopped users before alarms are due"
     bug: "314907186"
 }
+
+flag {
+    name: "acquire_wakelock_before_send"
+    namespace: "backstage_power"
+    description: "Acquire the userspace alarm wakelock before sending the alarm"
+    bug: "391413964"
+    metadata {
+      purpose: PURPOSE_BUGFIX
+    }
+}
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 829442aed6ac..f89b13dce307 100644
--- a/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
+++ b/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
@@ -5334,6 +5334,18 @@ public class AlarmManagerService extends SystemService {
         public void deliverLocked(Alarm alarm, long nowELAPSED) {
             final long workSourceToken = ThreadLocalWorkSource.setUid(
                     getAlarmAttributionUid(alarm));
+
+            if (Flags.acquireWakelockBeforeSend()) {
+                // Acquire the wakelock before starting the app. This needs to be done to avoid
+                // random stalls in the receiving app in case a suspend attempt is already in
+                // progress. See b/391413964 for an incident where this was found to happen.
+                if (mBroadcastRefCount == 0) {
+                    setWakelockWorkSource(alarm.workSource, alarm.creatorUid, alarm.statsTag, true);
+                    mWakeLock.acquire();
+                    mHandler.obtainMessage(AlarmHandler.REPORT_ALARMS_ACTIVE, 1, 0).sendToTarget();
+                }
+            }
+
             try {
                 if (alarm.operation != null) {
                     // PendingIntent alarm
@@ -5399,14 +5411,16 @@ public class AlarmManagerService extends SystemService {
                 ThreadLocalWorkSource.restore(workSourceToken);
             }
 
-            // The alarm is now in flight; now arrange wakelock and stats tracking
             if (DEBUG_WAKELOCK) {
                 Slog.d(TAG, "mBroadcastRefCount -> " + (mBroadcastRefCount + 1));
             }
-            if (mBroadcastRefCount == 0) {
-                setWakelockWorkSource(alarm.workSource, alarm.creatorUid, alarm.statsTag, true);
-                mWakeLock.acquire();
-                mHandler.obtainMessage(AlarmHandler.REPORT_ALARMS_ACTIVE, 1, 0).sendToTarget();
+            if (!Flags.acquireWakelockBeforeSend()) {
+                // The alarm is now in flight; now arrange wakelock and stats tracking
+                if (mBroadcastRefCount == 0) {
+                    setWakelockWorkSource(alarm.workSource, alarm.creatorUid, alarm.statsTag, true);
+                    mWakeLock.acquire();
+                    mHandler.obtainMessage(AlarmHandler.REPORT_ALARMS_ACTIVE, 1, 0).sendToTarget();
+                }
             }
             final InFlight inflight = new InFlight(AlarmManagerService.this, alarm, nowELAPSED);
             mInFlight.add(inflight);