Merge "Adding device-idle policy time to alarms"

5 files changed, 429 insertions, 287 deletions

diff --git a/apex/jobscheduler/service/java/com/android/server/alarm/Alarm.java b/apex/jobscheduler/service/java/com/android/server/alarm/Alarm.java
index 04feef485048..a8c0f0ec3e00 100644
--- a/apex/jobscheduler/service/java/com/android/server/alarm/Alarm.java
+++ b/apex/jobscheduler/service/java/com/android/server/alarm/Alarm.java
@@ -31,6 +31,8 @@ import android.util.IndentingPrintWriter;
 import android.util.TimeUtils;
 import android.util.proto.ProtoOutputStream;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 import java.text.SimpleDateFormat;
 import java.util.Date;
 
@@ -39,17 +41,23 @@ import java.util.Date;
  * expires. The timer will wake up the device if the alarm is a "wakeup" alarm.
  */
 class Alarm {
-    private static final int NUM_POLICIES = 2;
+    @VisibleForTesting
+    public static final int NUM_POLICIES = 3;
     /**
      * Index used to store the time the alarm was requested to expire. To be used with
-     * {@link #setPolicyElapsed(int, long)}
+     * {@link #setPolicyElapsed(int, long)}.
      */
     public static final int REQUESTER_POLICY_INDEX = 0;
     /**
      * Index used to store the earliest time the alarm can expire based on app-standby policy.
-     * To be used with {@link #setPolicyElapsed(int, long)}
+     * To be used with {@link #setPolicyElapsed(int, long)}.
      */
     public static final int APP_STANDBY_POLICY_INDEX = 1;
+    /**
+     * Index used to store the earliest time the alarm can expire based on the device's doze policy.
+     * To be used with {@link #setPolicyElapsed(int, long)}.
+     */
+    public static final int DEVICE_IDLE_POLICY_INDEX = 2;
 
     public final int type;
     /**
@@ -128,11 +136,19 @@ class Alarm {
      * @param policyIndex The index of the policy. One of [{@link #REQUESTER_POLICY_INDEX},
      *                    {@link #APP_STANDBY_POLICY_INDEX}].
      */
-    public long getPolicyElapsed(int policyIndex) {
+    @VisibleForTesting
+    long getPolicyElapsed(int policyIndex) {
         return mPolicyWhenElapsed[policyIndex];
     }
 
     /**
+     * @return the time this alarm was requested to go off in the elapsed time base.
+     */
+    public long getRequestedElapsed() {
+        return mPolicyWhenElapsed[REQUESTER_POLICY_INDEX];
+    }
+
+    /**
      * Get the earliest time that this alarm should be delivered to the requesting app.
      */
     public long getWhenElapsed() {
@@ -205,8 +221,10 @@ class Alarm {
                 return "requester";
             case APP_STANDBY_POLICY_INDEX:
                 return "app_standby";
+            case DEVICE_IDLE_POLICY_INDEX:
+                return "device_idle";
             default:
-                return "unknown";
+                return "--unknown--";
         }
     }
 
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 f196567bc391..c8a04d674739 100644
--- a/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
+++ b/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
@@ -21,12 +21,14 @@ import static android.app.AlarmManager.ELAPSED_REALTIME;
 import static android.app.AlarmManager.ELAPSED_REALTIME_WAKEUP;
 import static android.app.AlarmManager.FLAG_ALLOW_WHILE_IDLE;
 import static android.app.AlarmManager.FLAG_ALLOW_WHILE_IDLE_UNRESTRICTED;
+import static android.app.AlarmManager.FLAG_IDLE_UNTIL;
 import static android.app.AlarmManager.RTC;
 import static android.app.AlarmManager.RTC_WAKEUP;
 import static android.content.pm.PackageManager.MATCH_SYSTEM_ONLY;
 import static android.os.UserHandle.USER_SYSTEM;
 
 import static com.android.server.alarm.Alarm.APP_STANDBY_POLICY_INDEX;
+import static com.android.server.alarm.Alarm.DEVICE_IDLE_POLICY_INDEX;
 import static com.android.server.alarm.Alarm.REQUESTER_POLICY_INDEX;
 
 import android.annotation.NonNull;
@@ -81,7 +83,6 @@ import android.util.ArraySet;
 import android.util.IndentingPrintWriter;
 import android.util.Log;
 import android.util.LongArrayQueue;
-import android.util.MutableBoolean;
 import android.util.NtpTrustedTime;
 import android.util.Pair;
 import android.util.Slog;
@@ -699,36 +700,31 @@ public class AlarmManagerService extends SystemService {
     final HashMap<String, PriorityClass> mPriorities = new HashMap<>();
     int mCurrentSeq = 0;
 
-    static final class WakeupEvent {
-        public long when;
-        public int uid;
-        public String action;
-
-        public WakeupEvent(long theTime, int theUid, String theAction) {
-            when = theTime;
-            uid = theUid;
-            action = theAction;
-        }
-    }
-
     final Comparator<Alarm> mAlarmDispatchComparator = new Comparator<Alarm>() {
         @Override
         public int compare(Alarm lhs, Alarm rhs) {
-            // priority class trumps everything.  TICK < WAKEUP < NORMAL
+
+            // Alarm to exit device_idle should go out first.
+            final boolean lhsIdleUntil = (lhs.flags & FLAG_IDLE_UNTIL) != 0;
+            final boolean rhsIdleUntil = (rhs.flags & FLAG_IDLE_UNTIL) != 0;
+            if (lhsIdleUntil != rhsIdleUntil) {
+                return lhsIdleUntil ? -1 : 1;
+            }
+
+            // Then, priority class trumps everything.  TICK < WAKEUP < NORMAL
             if (lhs.priorityClass.priority < rhs.priorityClass.priority) {
                 return -1;
             } else if (lhs.priorityClass.priority > rhs.priorityClass.priority) {
                 return 1;
             }
 
-            // within each class, sort by nominal delivery time
-            if (lhs.getWhenElapsed() < rhs.getWhenElapsed()) {
+            // within each class, sort by requested delivery time
+            if (lhs.getRequestedElapsed() < rhs.getRequestedElapsed()) {
                 return -1;
-            } else if (lhs.getWhenElapsed() > rhs.getWhenElapsed()) {
+            } else if (lhs.getRequestedElapsed() > rhs.getRequestedElapsed()) {
                 return 1;
             }
 
-            // same priority class + same target delivery time
             return 0;
         }
     };
@@ -777,10 +773,9 @@ public class AlarmManagerService extends SystemService {
     final AlarmStore mAlarmStore;
 
     // set to non-null if in idle mode; while in this mode, any alarms we don't want
-    // to run during this time are placed in mPendingWhileIdleAlarms
+    // to run during this time are rescehduled to go off after this alarm.
     Alarm mPendingIdleUntil = null;
     Alarm mNextWakeFromIdle = null;
-    ArrayList<Alarm> mPendingWhileIdleAlarms = new ArrayList<>();
 
     @VisibleForTesting
     AlarmManagerService(Context context, Injector injector) {
@@ -830,15 +825,21 @@ public class AlarmManagerService extends SystemService {
                 return restoreRequestedTime(a);
             });
 
-            if (mNextWakeFromIdle != null && isRtc(mNextWakeFromIdle.type)) {
-                // The next wake from idle got updated due to the rtc time change, implying we need
-                // to update the time we have to come out of idle too.
-                changed |= mAlarmStore.updateAlarmDeliveries(a -> {
-                    if (a != mPendingIdleUntil) {
-                        return false;
+            if (changed && mPendingIdleUntil != null) {
+                if (mNextWakeFromIdle != null && isRtc(mNextWakeFromIdle.type)) {
+                    // The next wake from idle got updated due to the rtc time change, so we need
+                    // to update the time we have to come out of idle too.
+                    final boolean idleUntilUpdated = mAlarmStore.updateAlarmDeliveries(a -> {
+                        if (a != mPendingIdleUntil) {
+                            return false;
+                        }
+                        return adjustIdleUntilTime(a);
+                    });
+                    if (idleUntilUpdated) {
+                        mAlarmStore.updateAlarmDeliveries(
+                                alarm -> adjustDeliveryTimeBasedOnDeviceIdle(alarm));
                     }
-                    return adjustIdleUntilTime(a);
-                });
+                }
             }
 
             if (changed) {
@@ -972,11 +973,10 @@ public class AlarmManagerService extends SystemService {
             // Recurring alarms may have passed several alarm intervals while the
             // alarm was kept pending. Send the appropriate trigger count.
             if (alarm.repeatInterval > 0) {
-                alarm.count += (nowELAPSED - alarm.getPolicyElapsed(REQUESTER_POLICY_INDEX))
-                        / alarm.repeatInterval;
+                alarm.count += (nowELAPSED - alarm.getRequestedElapsed()) / alarm.repeatInterval;
                 // Also schedule its next recurrence
                 final long delta = alarm.count * alarm.repeatInterval;
-                final long nextElapsed = alarm.getPolicyElapsed(REQUESTER_POLICY_INDEX) + delta;
+                final long nextElapsed = alarm.getRequestedElapsed() + delta;
                 final long nextMaxElapsed = maxTriggerTime(nowELAPSED, nextElapsed,
                         alarm.repeatInterval);
                 setImplLocked(alarm.type, alarm.origWhen + delta, nextElapsed,
@@ -1015,25 +1015,6 @@ public class AlarmManagerService extends SystemService {
         }
     }
 
-    void restorePendingWhileIdleAlarmsLocked() {
-        if (RECORD_DEVICE_IDLE_ALARMS) {
-            IdleDispatchEntry ent = new IdleDispatchEntry();
-            ent.uid = 0;
-            ent.pkg = "FINISH IDLE";
-            ent.elapsedRealtime = mInjector.getElapsedRealtime();
-            mAllowWhileIdleDispatches.add(ent);
-        }
-
-        // Bring pending alarms back into the main list.
-        if (mPendingWhileIdleAlarms.size() > 0) {
-            ArrayList<Alarm> alarms = mPendingWhileIdleAlarms;
-            mPendingWhileIdleAlarms = new ArrayList<>();
-            for (int i = alarms.size() - 1; i >= 0; i--) {
-                setImplLocked(alarms.get(i));
-            }
-        }
-    }
-
     static final class InFlight {
         final PendingIntent mPendingIntent;
         final long mWhenElapsed;
@@ -1571,7 +1552,7 @@ public class AlarmManagerService extends SystemService {
             return false;
         }
         restoreRequestedTime(alarm);
-        long triggerBeforeFuzz = alarm.getPolicyElapsed(REQUESTER_POLICY_INDEX);
+        long triggerBeforeFuzz = alarm.getRequestedElapsed();
         if (mNextWakeFromIdle != null && triggerBeforeFuzz > mNextWakeFromIdle.getWhenElapsed()) {
             triggerBeforeFuzz = mNextWakeFromIdle.getWhenElapsed();
         }
@@ -1591,6 +1572,35 @@ public class AlarmManagerService extends SystemService {
     }
 
     /**
+     * Adjusts the delivery time of the alarm based on device_idle (doze) rules.
+     *
+     * @param alarm The alarm to adjust
+     * @return {@code true} if the alarm delivery time was updated.
+     */
+    private boolean adjustDeliveryTimeBasedOnDeviceIdle(Alarm alarm) {
+        final long nowElapsed = mInjector.getElapsedRealtime();
+        if (mPendingIdleUntil == null || mPendingIdleUntil == alarm) {
+            return alarm.setPolicyElapsed(DEVICE_IDLE_POLICY_INDEX, nowElapsed);
+        }
+
+        final long deviceIdlePolicyTime;
+        if ((alarm.flags & (AlarmManager.FLAG_ALLOW_WHILE_IDLE_UNRESTRICTED
+                | AlarmManager.FLAG_WAKE_FROM_IDLE)) != 0) {
+            // Unrestricted.
+            deviceIdlePolicyTime = nowElapsed;
+        } else if ((alarm.flags & AlarmManager.FLAG_ALLOW_WHILE_IDLE) != 0) {
+            // Allowed but limited.
+            final long lastDispatch = mLastAllowWhileIdleDispatch.get(alarm.creatorUid, 0);
+            deviceIdlePolicyTime = (lastDispatch == 0) ? nowElapsed
+                    : lastDispatch + mConstants.ALLOW_WHILE_IDLE_LONG_TIME;
+        } else {
+            // Not allowed.
+            deviceIdlePolicyTime = mPendingIdleUntil.getWhenElapsed();
+        }
+        return alarm.setPolicyElapsed(DEVICE_IDLE_POLICY_INDEX, deviceIdlePolicyTime);
+    }
+
+    /**
      * Adjusts the alarm's policy time for app_standby.
      *
      * @param alarm The alarm to update.
@@ -1643,12 +1653,6 @@ public class AlarmManagerService extends SystemService {
         return alarm.setPolicyElapsed(APP_STANDBY_POLICY_INDEX, nowElapsed);
     }
 
-    private static boolean isAllowedWhileIdle(Alarm a) {
-        return ((a.flags & (AlarmManager.FLAG_ALLOW_WHILE_IDLE
-                | AlarmManager.FLAG_ALLOW_WHILE_IDLE_UNRESTRICTED
-                | AlarmManager.FLAG_WAKE_FROM_IDLE)) != 0);
-    }
-
     private void setImplLocked(Alarm a) {
         if ((a.flags & AlarmManager.FLAG_IDLE_UNTIL) != 0) {
             adjustIdleUntilTime(a);
@@ -1676,16 +1680,9 @@ public class AlarmManagerService extends SystemService {
                 mAlarmStore.remove(mPendingIdleUntil::equals);
             }
             mPendingIdleUntil = a;
-            final ArrayList<Alarm> notAllowedWhileIdleAlarms = mAlarmStore.remove(
-                    alarm -> !isAllowedWhileIdle(alarm));
-            mPendingWhileIdleAlarms.addAll(notAllowedWhileIdleAlarms);
+            mAlarmStore.updateAlarmDeliveries(alarm -> adjustDeliveryTimeBasedOnDeviceIdle(alarm));
         } else if (mPendingIdleUntil != null) {
-            // We currently have an idle until alarm scheduled; if the new alarm has
-            // not explicitly stated it wants to run while idle, then put it on hold.
-            if (!isAllowedWhileIdle(a)) {
-                mPendingWhileIdleAlarms.add(a);
-                return;
-            }
+            adjustDeliveryTimeBasedOnDeviceIdle(a);
         }
         if ((a.flags & AlarmManager.FLAG_WAKE_FROM_IDLE) != 0) {
             if (mNextWakeFromIdle == null || mNextWakeFromIdle.getWhenElapsed()
@@ -1694,12 +1691,17 @@ public class AlarmManagerService extends SystemService {
                 // If this wake from idle is earlier than whatever was previously scheduled,
                 // and we are currently idling, then the idle-until time needs to be updated.
                 if (mPendingIdleUntil != null) {
-                    mAlarmStore.updateAlarmDeliveries(alarm -> {
+                    final boolean updated = mAlarmStore.updateAlarmDeliveries(alarm -> {
                         if (alarm != mPendingIdleUntil) {
                             return false;
                         }
                         return adjustIdleUntilTime(alarm);
                     });
+                    if (updated) {
+                        // idle-until got updated, so also update all alarms not allowed while idle.
+                        mAlarmStore.updateAlarmDeliveries(
+                                alarm -> adjustDeliveryTimeBasedOnDeviceIdle(alarm));
+                    }
                 }
             }
         }
@@ -2095,7 +2097,7 @@ public class AlarmManagerService extends SystemService {
 
             mAppWakeupHistory.dump(pw, nowELAPSED);
 
-            if (mPendingIdleUntil != null || mPendingWhileIdleAlarms.size() > 0) {
+            if (mPendingIdleUntil != null) {
                 pw.println();
                 pw.println("Idle mode state:");
 
@@ -2107,8 +2109,6 @@ public class AlarmManagerService extends SystemService {
                 } else {
                     pw.println("null");
                 }
-                pw.println("Pending alarms:");
-                dumpAlarmList(pw, mPendingWhileIdleAlarms, nowELAPSED, sdf);
                 pw.decreaseIndent();
             }
             if (mNextWakeFromIdle != null) {
@@ -2428,10 +2428,6 @@ public class AlarmManagerService extends SystemService {
                 mPendingIdleUntil.dumpDebug(
                         proto, AlarmManagerServiceDumpProto.PENDING_IDLE_UNTIL, nowElapsed);
             }
-            for (Alarm a : mPendingWhileIdleAlarms) {
-                a.dumpDebug(proto, AlarmManagerServiceDumpProto.PENDING_WHILE_IDLE_ALARMS,
-                        nowElapsed);
-            }
             if (mNextWakeFromIdle != null) {
                 mNextWakeFromIdle.dumpDebug(proto, AlarmManagerServiceDumpProto.NEXT_WAKE_FROM_IDLE,
                         nowElapsed);
@@ -2752,21 +2748,13 @@ public class AlarmManagerService extends SystemService {
             return;
         }
 
-        final Predicate<Alarm> whichAlarms = (Alarm a) -> a.matches(operation, directReceiver);
-        final ArrayList<Alarm> removedAlarms = mAlarmStore.remove(whichAlarms);
+        final ArrayList<Alarm> removedAlarms = mAlarmStore.remove(
+                a -> a.matches(operation, directReceiver));
         for (final Alarm removed : removedAlarms) {
             decrementAlarmCount(removed.uid, 1);
         }
         final boolean didRemove = !removedAlarms.isEmpty();
 
-        for (int i = mPendingWhileIdleAlarms.size() - 1; i >= 0; i--) {
-            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--) {
@@ -2793,22 +2781,25 @@ public class AlarmManagerService extends SystemService {
             if (DEBUG_BATCH) {
                 Slog.v(TAG, "remove(operation) changed bounds; rebatching");
             }
-            boolean restorePending = false;
+            boolean idleUntilUpdated = false;
             if (mPendingIdleUntil != null && mPendingIdleUntil.matches(operation, directReceiver)) {
                 mPendingIdleUntil = null;
-                restorePending = true;
+                idleUntilUpdated = true;
             }
             if (mNextWakeFromIdle != null && mNextWakeFromIdle.matches(operation, directReceiver)) {
                 mNextWakeFromIdle = mAlarmStore.getNextWakeFromIdleAlarm();
-                mAlarmStore.updateAlarmDeliveries(alarm -> {
-                    if (alarm != mPendingIdleUntil) {
-                        return false;
-                    }
-                    return adjustIdleUntilTime(alarm);
-                });
+                if (mPendingIdleUntil != null) {
+                    idleUntilUpdated |= mAlarmStore.updateAlarmDeliveries(alarm -> {
+                        if (alarm != mPendingIdleUntil) {
+                            return false;
+                        }
+                        return adjustIdleUntilTime(alarm);
+                    });
+                }
             }
-            if (restorePending) {
-                restorePendingWhileIdleAlarmsLocked();
+            if (idleUntilUpdated) {
+                mAlarmStore.updateAlarmDeliveries(
+                        alarm -> adjustDeliveryTimeBasedOnDeviceIdle(alarm));
             }
             rescheduleKernelAlarmsLocked();
             updateNextAlarmClockLocked();
@@ -2821,21 +2812,12 @@ public class AlarmManagerService extends SystemService {
             return;
         }
 
-        final Predicate<Alarm> whichAlarms = (Alarm a) -> a.uid == uid;
-        final ArrayList<Alarm> removed = mAlarmStore.remove(whichAlarms);
+        final ArrayList<Alarm> removed = mAlarmStore.remove(a -> a.uid == uid);
         final boolean didRemove = !removed.isEmpty();
         if (didRemove) {
             decrementAlarmCount(uid, removed.size());
         }
 
-        for (int i = mPendingWhileIdleAlarms.size() - 1; i >= 0; i--) {
-            final Alarm a = mPendingWhileIdleAlarms.get(i);
-            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--) {
             final ArrayList<Alarm> alarmsForUid = mPendingBackgroundAlarms.valueAt(i);
             for (int j = alarmsForUid.size() - 1; j >= 0; j--) {
@@ -2850,20 +2832,26 @@ public class AlarmManagerService extends SystemService {
         }
         // If we're currently using this app's alarms to come out of doze,
         // make sure to reset to any remaining WAKE_FROM_IDLE alarms.
+        boolean idleUntilUpdated = false;
         if (mNextWakeFromIdle != null && mNextWakeFromIdle.uid == uid) {
             mNextWakeFromIdle = mAlarmStore.getNextWakeFromIdleAlarm();
-            mAlarmStore.updateAlarmDeliveries(alarm -> {
-                if (alarm != mPendingIdleUntil) {
-                    return false;
-                }
-                return adjustIdleUntilTime(alarm);
-            });
+            if (mPendingIdleUntil != null) {
+                idleUntilUpdated |= mAlarmStore.updateAlarmDeliveries(alarm -> {
+                    if (alarm != mPendingIdleUntil) {
+                        return false;
+                    }
+                    return adjustIdleUntilTime(alarm);
+                });
+            }
         }
         if (mPendingIdleUntil != null && mPendingIdleUntil.uid == uid) {
             // Should never happen - only the system uid is allowed to set idle-until alarms
             Slog.wtf(TAG, "Removed app uid " + uid + " set idle-until alarm!");
             mPendingIdleUntil = null;
-            restorePendingWhileIdleAlarmsLocked();
+            idleUntilUpdated = true;
+        }
+        if (idleUntilUpdated) {
+            mAlarmStore.updateAlarmDeliveries(alarm -> adjustDeliveryTimeBasedOnDeviceIdle(alarm));
         }
         if (didRemove) {
             if (DEBUG_BATCH) {
@@ -2883,29 +2871,12 @@ public class AlarmManagerService extends SystemService {
             return;
         }
 
-        final MutableBoolean removedNextWakeFromIdle = new MutableBoolean(false);
-        final Predicate<Alarm> whichAlarms = (Alarm a) -> {
-            final boolean didMatch = a.matches(packageName);
-            if (didMatch && a == mNextWakeFromIdle) {
-                removedNextWakeFromIdle.value = true;
-            }
-            return didMatch;
-        };
-
-        final ArrayList<Alarm> removed = mAlarmStore.remove(whichAlarms);
+        final ArrayList<Alarm> removed = mAlarmStore.remove(a -> a.matches(packageName));
         final boolean didRemove = !removed.isEmpty();
         if (didRemove) {
             decrementAlarmCount(removed.get(0).uid, removed.size());
         }
 
-        for (int i = mPendingWhileIdleAlarms.size() - 1; i >= 0; i--) {
-            final Alarm a = mPendingWhileIdleAlarms.get(i);
-            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--) {
@@ -2921,14 +2892,20 @@ public class AlarmManagerService extends SystemService {
         }
         // If we're currently using this app's alarms to come out of doze,
         // make sure to reset to any remaining WAKE_FROM_IDLE alarms.
-        if (removedNextWakeFromIdle.value) {
+        if (mNextWakeFromIdle != null && mNextWakeFromIdle.matches(packageName)) {
             mNextWakeFromIdle = mAlarmStore.getNextWakeFromIdleAlarm();
-            mAlarmStore.updateAlarmDeliveries(alarm -> {
-                if (alarm != mPendingIdleUntil) {
-                    return false;
+            if (mPendingIdleUntil != null) {
+                final boolean updated = mAlarmStore.updateAlarmDeliveries(alarm -> {
+                    if (alarm != mPendingIdleUntil) {
+                        return false;
+                    }
+                    return adjustIdleUntilTime(alarm);
+                });
+                if (updated) {
+                    mAlarmStore.updateAlarmDeliveries(
+                            alarm -> adjustDeliveryTimeBasedOnDeviceIdle(alarm));
                 }
-                return adjustIdleUntilTime(alarm);
-            });
+            }
         }
         if (didRemove) {
             if (DEBUG_BATCH) {
@@ -2953,14 +2930,6 @@ public class AlarmManagerService extends SystemService {
             decrementAlarmCount(uid, removed.size());
         }
 
-        for (int i = mPendingWhileIdleAlarms.size() - 1; i >= 0; i--) {
-            final Alarm a = mPendingWhileIdleAlarms.get(i);
-            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);
@@ -2992,14 +2961,6 @@ public class AlarmManagerService extends SystemService {
         }
         final boolean didRemove = !removedAlarms.isEmpty();
 
-        for (int i = mPendingWhileIdleAlarms.size() - 1; i >= 0; i--) {
-            if (UserHandle.getUserId(mPendingWhileIdleAlarms.get(i).creatorUid)
-                    == userHandle) {
-                // Don't set didRemove, since this doesn't impact the scheduled alarms.
-                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);
@@ -3016,6 +2977,21 @@ public class AlarmManagerService extends SystemService {
                 mLastAllowWhileIdleDispatch.removeAt(i);
             }
         }
+        if (mNextWakeFromIdle != null && whichAlarms.test(mNextWakeFromIdle)) {
+            mNextWakeFromIdle = mAlarmStore.getNextWakeFromIdleAlarm();
+            if (mPendingIdleUntil != null) {
+                final boolean updated = mAlarmStore.updateAlarmDeliveries(alarm -> {
+                    if (alarm != mPendingIdleUntil) {
+                        return false;
+                    }
+                    return adjustIdleUntilTime(alarm);
+                });
+                if (updated) {
+                    mAlarmStore.updateAlarmDeliveries(
+                            alarm -> adjustDeliveryTimeBasedOnDeviceIdle(alarm));
+                }
+            }
+        }
 
         if (didRemove) {
             if (DEBUG_BATCH) {
@@ -3062,12 +3038,6 @@ public class AlarmManagerService extends SystemService {
                 return true;
             }
         }
-        for (int i = 0; i < mPendingWhileIdleAlarms.size(); i++) {
-            final Alarm a = mPendingWhileIdleAlarms.get(i);
-            if (a.matches(packageName)) {
-                return true;
-            }
-        }
         return false;
     }
 
@@ -3134,16 +3104,10 @@ public class AlarmManagerService extends SystemService {
     private static native long getNextAlarm(long nativeData, int type);
 
     private long getWhileIdleMinIntervalLocked(int uid) {
-        final boolean dozing = mPendingIdleUntil != null;
         final boolean ebs = (mAppStateTracker != null)
                 && mAppStateTracker.isForceAllAppsStandbyEnabled();
-        if (!dozing && !ebs) {
-            return mConstants.ALLOW_WHILE_IDLE_SHORT_TIME;
-        }
-        if (dozing) {
-            return mConstants.ALLOW_WHILE_IDLE_LONG_TIME;
-        }
-        if (mUseAllowWhileIdleShortTime.get(uid)) {
+
+        if (!ebs || mUseAllowWhileIdleShortTime.get(uid)) {
             // if the last allow-while-idle went off while uid was fg, or the uid
             // recently came into fg, don't block the alarm for long.
             return mConstants.ALLOW_WHILE_IDLE_SHORT_TIME;
@@ -3201,16 +3165,12 @@ public class AlarmManagerService extends SystemService {
             }
             if (mPendingIdleUntil == alarm) {
                 mPendingIdleUntil = null;
-                restorePendingWhileIdleAlarmsLocked();
+                mAlarmStore.updateAlarmDeliveries(a -> adjustDeliveryTimeBasedOnDeviceIdle(a));
             }
             if (mNextWakeFromIdle == alarm) {
                 mNextWakeFromIdle = mAlarmStore.getNextWakeFromIdleAlarm();
-                mAlarmStore.updateAlarmDeliveries(a -> {
-                    if (a != mPendingIdleUntil) {
-                        return false;
-                    }
-                    return adjustIdleUntilTime(a);
-                });
+                // Note that we don't need to update mPendingIdleUntil because it should already
+                // be removed from the alarm store.
             }
 
             // Recurring alarms may have passed several alarm intervals while the
@@ -3218,11 +3178,10 @@ public class AlarmManagerService extends SystemService {
             if (alarm.repeatInterval > 0) {
                 // this adjustment will be zero if we're late by
                 // less than one full repeat interval
-                alarm.count += (nowELAPSED - alarm.getPolicyElapsed(REQUESTER_POLICY_INDEX))
-                        / alarm.repeatInterval;
+                alarm.count += (nowELAPSED - alarm.getRequestedElapsed()) / alarm.repeatInterval;
                 // Also schedule its next recurrence
                 final long delta = alarm.count * alarm.repeatInterval;
-                final long nextElapsed = alarm.getPolicyElapsed(REQUESTER_POLICY_INDEX) + delta;
+                final long nextElapsed = alarm.getRequestedElapsed() + delta;
                 final long nextMaxElapsed = maxTriggerTime(nowELAPSED, nextElapsed,
                         alarm.repeatInterval);
                 setImplLocked(alarm.type, alarm.origWhen + delta, nextElapsed,
@@ -4222,6 +4181,12 @@ public class AlarmManagerService extends SystemService {
             if (allowWhileIdle) {
                 // Record the last time this uid handled an ALLOW_WHILE_IDLE alarm.
                 mLastAllowWhileIdleDispatch.put(alarm.creatorUid, nowELAPSED);
+                mAlarmStore.updateAlarmDeliveries(a -> {
+                    if (a.creatorUid != alarm.creatorUid) {
+                        return false;
+                    }
+                    return adjustDeliveryTimeBasedOnDeviceIdle(a);
+                });
                 if ((mAppStateTracker == null)
                         || mAppStateTracker.isUidInForeground(alarm.creatorUid)) {
                     mUseAllowWhileIdleShortTime.put(alarm.creatorUid, true);
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 a5f083477863..db4aba519a6a 100644
--- a/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmManagerServiceTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmManagerServiceTest.java
@@ -17,8 +17,14 @@ package com.android.server.alarm;
 
 import static android.app.AlarmManager.ELAPSED_REALTIME;
 import static android.app.AlarmManager.ELAPSED_REALTIME_WAKEUP;
+import static android.app.AlarmManager.FLAG_ALLOW_WHILE_IDLE;
+import static android.app.AlarmManager.FLAG_ALLOW_WHILE_IDLE_UNRESTRICTED;
+import static android.app.AlarmManager.FLAG_IDLE_UNTIL;
+import static android.app.AlarmManager.FLAG_STANDALONE;
+import static android.app.AlarmManager.FLAG_WAKE_FROM_IDLE;
 import static android.app.AlarmManager.RTC;
 import static android.app.AlarmManager.RTC_WAKEUP;
+import static android.app.AlarmManager.WINDOW_EXACT;
 import static android.app.usage.UsageStatsManager.STANDBY_BUCKET_ACTIVE;
 import static android.app.usage.UsageStatsManager.STANDBY_BUCKET_FREQUENT;
 import static android.app.usage.UsageStatsManager.STANDBY_BUCKET_RARE;
@@ -66,7 +72,6 @@ import static org.mockito.Mockito.atLeastOnce;
 
 import android.app.ActivityManager;
 import android.app.ActivityManagerInternal;
-import android.app.AlarmManager;
 import android.app.IActivityManager;
 import android.app.IAlarmCompleteListener;
 import android.app.IAlarmListener;
@@ -85,7 +90,6 @@ import android.os.UserHandle;
 import android.platform.test.annotations.Presubmit;
 import android.provider.DeviceConfig;
 import android.provider.Settings;
-import android.util.IndentingPrintWriter;
 import android.util.Log;
 import android.util.SparseArray;
 
@@ -112,8 +116,6 @@ import org.mockito.MockitoSession;
 import org.mockito.quality.Strictness;
 import org.mockito.stubbing.Answer;
 
-import java.io.PrintWriter;
-import java.text.SimpleDateFormat;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.concurrent.Executor;
@@ -354,48 +356,51 @@ public class AlarmManagerServiceTest {
     }
 
     private void setTestAlarm(int type, long triggerTime, PendingIntent operation) {
-        setTestAlarm(type, triggerTime, operation, 0, AlarmManager.FLAG_STANDALONE,
-                TEST_CALLING_UID);
+        setTestAlarm(type, triggerTime, operation, 0, FLAG_STANDALONE, TEST_CALLING_UID);
     }
 
     private void setRepeatingTestAlarm(int type, long firstTrigger, long interval,
             PendingIntent pi) {
-        setTestAlarm(type, firstTrigger, pi, interval, AlarmManager.FLAG_STANDALONE,
-                TEST_CALLING_UID);
+        setTestAlarm(type, firstTrigger, pi, interval, FLAG_STANDALONE, TEST_CALLING_UID);
     }
 
     private void setIdleUntilAlarm(int type, long triggerTime, PendingIntent pi) {
-        setTestAlarm(type, triggerTime, pi, 0, AlarmManager.FLAG_IDLE_UNTIL, TEST_CALLING_UID);
+        setTestAlarm(type, triggerTime, pi, 0, FLAG_IDLE_UNTIL, TEST_CALLING_UID);
     }
 
     private void setWakeFromIdle(int type, long triggerTime, PendingIntent pi) {
         // Note: Only alarm clock alarms are allowed to include this flag in the actual service.
         // But this is a unit test so we'll only test the flag for granularity and convenience.
-        setTestAlarm(type, triggerTime, pi, 0,
-                AlarmManager.FLAG_WAKE_FROM_IDLE | AlarmManager.FLAG_STANDALONE, TEST_CALLING_UID);
+        setTestAlarm(type, triggerTime, pi, 0, FLAG_WAKE_FROM_IDLE | FLAG_STANDALONE,
+                TEST_CALLING_UID);
+    }
+
+    private void setAllowWhileIdleAlarm(int type, long triggerTime, PendingIntent pi,
+            boolean unrestricted) {
+        final int flags = unrestricted ? FLAG_ALLOW_WHILE_IDLE_UNRESTRICTED : FLAG_ALLOW_WHILE_IDLE;
+        setTestAlarm(type, triggerTime, pi, 0, flags, TEST_CALLING_UID);
     }
 
     private void setTestAlarm(int type, long triggerTime, PendingIntent operation, long interval,
             int flags, int callingUid) {
-        mService.setImpl(type, triggerTime, AlarmManager.WINDOW_EXACT, interval, operation, null,
-                "test", flags, null, null, callingUid, TEST_CALLING_PACKAGE);
+        mService.setImpl(type, triggerTime, WINDOW_EXACT, interval, operation, null, "test", flags,
+                null, null, callingUid, TEST_CALLING_PACKAGE);
     }
 
     private void setTestAlarmWithListener(int type, long triggerTime, IAlarmListener listener) {
-        mService.setImpl(type, triggerTime, AlarmManager.WINDOW_EXACT, 0,
-                null, listener, "test", AlarmManager.FLAG_STANDALONE, null, null,
-                TEST_CALLING_UID, TEST_CALLING_PACKAGE);
+        mService.setImpl(type, triggerTime, WINDOW_EXACT, 0, null, listener, "test",
+                FLAG_STANDALONE, null, null, TEST_CALLING_UID, TEST_CALLING_PACKAGE);
     }
 
 
     private PendingIntent getNewMockPendingIntent() {
-        return getNewMockPendingIntent(TEST_CALLING_UID);
+        return getNewMockPendingIntent(TEST_CALLING_UID, TEST_CALLING_PACKAGE);
     }
 
-    private PendingIntent getNewMockPendingIntent(int mockUid) {
+    private PendingIntent getNewMockPendingIntent(int creatorUid, String creatorPackage) {
         final PendingIntent mockPi = mock(PendingIntent.class, Answers.RETURNS_DEEP_STUBS);
-        when(mockPi.getCreatorUid()).thenReturn(mockUid);
-        when(mockPi.getCreatorPackage()).thenReturn(TEST_CALLING_PACKAGE);
+        when(mockPi.getCreatorUid()).thenReturn(creatorUid);
+        when(mockPi.getCreatorPackage()).thenReturn(creatorPackage);
         return mockPi;
     }
 
@@ -865,7 +870,7 @@ public class AlarmManagerServiceTest {
     public void alarmCountKeyedOnCallingUid() {
         final int mockCreatorUid = 431412;
         setTestAlarm(ELAPSED_REALTIME, mNowElapsedTest + 5,
-                getNewMockPendingIntent(mockCreatorUid));
+                getNewMockPendingIntent(mockCreatorUid, TEST_CALLING_PACKAGE));
         assertEquals(1, mService.mAlarmsPerUid.get(TEST_CALLING_UID));
         assertEquals(-1, mService.mAlarmsPerUid.get(mockCreatorUid, -1));
     }
@@ -1018,7 +1023,8 @@ public class AlarmManagerServiceTest {
         for (int i = 0; i < numAlarms; i++) {
             int mockUid = UserHandle.getUid(mockUserId, 1234 + i);
             setTestAlarm(ELAPSED_REALTIME, mNowElapsedTest + i + 10,
-                    getNewMockPendingIntent(mockUid), 0, AlarmManager.FLAG_STANDALONE, mockUid);
+                    getNewMockPendingIntent(mockUid, TEST_CALLING_PACKAGE), 0, FLAG_STANDALONE,
+                    mockUid);
         }
         assertEquals(numAlarms, mService.mAlarmsPerUid.size());
         mService.removeUserLocked(mockUserId);
@@ -1026,26 +1032,6 @@ public class AlarmManagerServiceTest {
     }
 
     @Test
-    public void alarmCountOnRemoveFromPendingWhileIdle() {
-        mService.mPendingIdleUntil = mock(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];
@@ -1268,6 +1254,134 @@ public class AlarmManagerServiceTest {
         assertEquals(mNowElapsedTest + 12, mService.mPendingIdleUntil.getWhenElapsed());
     }
 
+    @Test
+    public void allowWhileIdleAlarmsWhileDeviceIdle() throws Exception {
+        doReturn(0).when(mService).fuzzForDuration(anyLong());
+
+        final long awiDelayForTest = 23;
+        setDeviceConfigLong(KEY_ALLOW_WHILE_IDLE_LONG_TIME, awiDelayForTest);
+        setDeviceConfigLong(KEY_ALLOW_WHILE_IDLE_SHORT_TIME, 0);
+
+        setIdleUntilAlarm(ELAPSED_REALTIME_WAKEUP, mNowElapsedTest + 1000,
+                getNewMockPendingIntent());
+        assertNotNull(mService.mPendingIdleUntil);
+
+        final long seedTrigger = mNowElapsedTest + 3;
+        final int numAlarms = 10;
+        final PendingIntent[] pis = new PendingIntent[numAlarms];
+        for (int i = 0; i < numAlarms; i++) {
+            pis[i] = getNewMockPendingIntent();
+            setAllowWhileIdleAlarm(ELAPSED_REALTIME_WAKEUP, seedTrigger + i * i, pis[i], false);
+        }
+
+        long lastAwiDispatch = -1;
+        int i = 0;
+        while (i < numAlarms) {
+            final long nextDispatch = (lastAwiDispatch >= 0) ? (lastAwiDispatch + awiDelayForTest)
+                    : (seedTrigger + i * i);
+            assertEquals("Wrong allow-while-idle dispatch", nextDispatch, mTestTimer.getElapsed());
+
+            mNowElapsedTest = nextDispatch;
+            mTestTimer.expire();
+
+            while (i < numAlarms && (seedTrigger + i * i) <= nextDispatch) {
+                verify(pis[i]).send(eq(mMockContext), eq(0), any(Intent.class), any(),
+                        any(Handler.class), isNull(), any());
+                i++;
+            }
+            Log.d(TAG, "Dispatched alarms upto " + i + " at " + nextDispatch);
+            lastAwiDispatch = nextDispatch;
+        }
+    }
+
+    @Test
+    public void allowWhileIdleUnrestricted() throws Exception {
+        doReturn(0).when(mService).fuzzForDuration(anyLong());
+
+        final long awiDelayForTest = 127;
+        setDeviceConfigLong(KEY_ALLOW_WHILE_IDLE_LONG_TIME, awiDelayForTest);
+        setDeviceConfigLong(KEY_ALLOW_WHILE_IDLE_SHORT_TIME, 0);
+
+        setIdleUntilAlarm(ELAPSED_REALTIME_WAKEUP, mNowElapsedTest + 1000,
+                getNewMockPendingIntent());
+        assertNotNull(mService.mPendingIdleUntil);
+
+        final long seedTrigger = mNowElapsedTest + 3;
+        for (int i = 1; i <= 5; i++) {
+            setAllowWhileIdleAlarm(ELAPSED_REALTIME_WAKEUP, seedTrigger + i * i,
+                    getNewMockPendingIntent(), true);
+        }
+        for (int i = 1; i <= 5; i++) {
+            final long nextTrigger = mTestTimer.getElapsed();
+            assertEquals("Wrong trigger for alarm " + i, seedTrigger + i * i, nextTrigger);
+            mNowElapsedTest = nextTrigger;
+            mTestTimer.expire();
+        }
+    }
+
+    @Test
+    public void deviceIdleThrottling() throws Exception {
+        doReturn(0).when(mService).fuzzForDuration(anyLong());
+
+        final long deviceIdleUntil = mNowElapsedTest + 1234;
+        setIdleUntilAlarm(ELAPSED_REALTIME_WAKEUP, deviceIdleUntil, getNewMockPendingIntent());
+
+        assertEquals(deviceIdleUntil, mTestTimer.getElapsed());
+
+        final int numAlarms = 10;
+        final PendingIntent[] pis = new PendingIntent[numAlarms];
+        for (int i = 0; i < numAlarms; i++) {
+            setTestAlarm(ELAPSED_REALTIME_WAKEUP, mNowElapsedTest + i + 1,
+                    pis[i] = getNewMockPendingIntent());
+            assertEquals(deviceIdleUntil, mTestTimer.getElapsed());
+        }
+
+        mNowElapsedTest = mTestTimer.getElapsed();
+        mTestTimer.expire();
+        for (int i = 0; i < numAlarms; i++) {
+            verify(pis[i]).send(eq(mMockContext), eq(0), any(Intent.class), any(),
+                    any(Handler.class), isNull(), any());
+        }
+    }
+
+    @Test
+    public void dispatchOrder() throws Exception {
+        doReturn(0).when(mService).fuzzForDuration(anyLong());
+
+        final long deviceIdleUntil = mNowElapsedTest + 1234;
+        final PendingIntent idleUntilPi = getNewMockPendingIntent();
+        setIdleUntilAlarm(ELAPSED_REALTIME_WAKEUP, deviceIdleUntil, idleUntilPi);
+
+        assertEquals(deviceIdleUntil, mTestTimer.getElapsed());
+
+        final PendingIntent pi5wakeup = getNewMockPendingIntent();
+        final PendingIntent pi4wakeupPackage = getNewMockPendingIntent();
+        final PendingIntent pi2nonWakeup = getNewMockPendingIntent(57, "test.different.package");
+
+        setTestAlarm(ELAPSED_REALTIME_WAKEUP, mNowElapsedTest + 5, pi5wakeup);
+        setTestAlarm(ELAPSED_REALTIME, mNowElapsedTest + 4, pi4wakeupPackage);
+        setTestAlarm(ELAPSED_REALTIME, mNowElapsedTest + 2, pi2nonWakeup);
+
+        mNowElapsedTest = deviceIdleUntil;
+        mTestTimer.expire();
+
+        // The order of the alarms in delivery list should be:
+        // IdleUntil, all alarms of a package with any wakeup alarms, then the rest.
+        // Within a package, alarms should be ordered by requested delivery time.
+        final PendingIntent[] expectedOrder = new PendingIntent[]{
+                idleUntilPi, pi4wakeupPackage, pi5wakeup, pi2nonWakeup};
+
+        ArgumentCaptor<ArrayList<Alarm>> listCaptor = ArgumentCaptor.forClass(ArrayList.class);
+        verify(mService).deliverAlarmsLocked(listCaptor.capture(), anyLong());
+        final ArrayList<Alarm> deliveryList = listCaptor.getValue();
+
+        assertEquals(expectedOrder.length, deliveryList.size());
+        for (int i = 0; i < expectedOrder.length; i++) {
+            assertTrue("Unexpected alarm: " + deliveryList.get(i) + " at pos: " + i,
+                    deliveryList.get(i).matches(expectedOrder[i], null));
+        }
+    }
+
     @After
     public void tearDown() {
         if (mMockingSession != null) {
@@ -1275,12 +1389,4 @@ public class AlarmManagerServiceTest {
         }
         LocalServices.removeServiceForTest(AlarmManagerInternal.class);
     }
-
-    private void dumpAllAlarms(String tag, ArrayList<Alarm> alarms) {
-        System.out.println(tag + ": ");
-        IndentingPrintWriter ipw = new IndentingPrintWriter(new PrintWriter(System.out));
-        AlarmManagerService.dumpAlarmList(ipw, alarms, mNowElapsedTest,
-                new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS"));
-        ipw.close();
-    }
 }
diff --git a/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmStoreTest.java b/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmStoreTest.java
index f0490ce469dd..c4fc61a5aa6e 100644
--- a/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmStoreTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmStoreTest.java
@@ -20,6 +20,7 @@ import static com.android.server.alarm.Constants.TEST_CALLING_PACKAGE;
 import static com.android.server.alarm.Constants.TEST_CALLING_UID;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 
@@ -45,22 +46,17 @@ public class AlarmStoreTest {
         mAlarmStore = new BatchingAlarmStore(null);
     }
 
-    private static Alarm createAlarm(long whenElapsed, long windowLength,
-            AlarmManager.AlarmClockInfo alarmClock) {
-        return createAlarm(AlarmManager.ELAPSED_REALTIME, whenElapsed, windowLength,
-                alarmClock);
+    private static Alarm createAlarm(long whenElapsed, long windowLength) {
+        return createAlarm(AlarmManager.ELAPSED_REALTIME, whenElapsed, windowLength, 0);
     }
 
-    private static Alarm createWakeupAlarm(long whenElapsed, long windowLength,
-            AlarmManager.AlarmClockInfo alarmClock) {
-        return createAlarm(AlarmManager.ELAPSED_REALTIME_WAKEUP, whenElapsed, windowLength,
-                alarmClock);
+    private static Alarm createWakeupAlarm(long whenElapsed, long windowLength, int flags) {
+        return createAlarm(AlarmManager.ELAPSED_REALTIME_WAKEUP, whenElapsed, windowLength, flags);
     }
 
-    private static Alarm createAlarm(int type, long whenElapsed, long windowLength,
-            AlarmManager.AlarmClockInfo alarmClock) {
+    private static Alarm createAlarm(int type, long whenElapsed, long windowLength, int flags) {
         return new Alarm(type, whenElapsed, whenElapsed, windowLength, 0, mock(PendingIntent.class),
-                null, null, null, 0, alarmClock, TEST_CALLING_UID, TEST_CALLING_PACKAGE);
+                null, null, null, flags, null, TEST_CALLING_UID, TEST_CALLING_PACKAGE);
     }
 
     private void addAlarmsToStore(Alarm... alarms) {
@@ -71,11 +67,11 @@ public class AlarmStoreTest {
 
     @Test
     public void add() {
-        final Alarm a1 = createAlarm(1, 0, null);
+        final Alarm a1 = createAlarm(1, 0);
         mAlarmStore.add(a1);
         assertEquals(1, mAlarmStore.size());
 
-        final Alarm a2 = createAlarm(2, 0, null);
+        final Alarm a2 = createAlarm(2, 0);
         mAlarmStore.add(a2);
         assertEquals(2, mAlarmStore.size());
 
@@ -86,9 +82,9 @@ public class AlarmStoreTest {
 
     @Test
     public void remove() {
-        final Alarm a1 = createAlarm(1, 0, null);
-        final Alarm a2 = createAlarm(2, 0, null);
-        final Alarm a5 = createAlarm(5, 0, null);
+        final Alarm a1 = createAlarm(1, 0);
+        final Alarm a2 = createAlarm(2, 0);
+        final Alarm a5 = createAlarm(5, 0);
         addAlarmsToStore(a1, a2, a5);
 
         ArrayList<Alarm> removed = mAlarmStore.remove(a -> (a.getWhenElapsed() < 4));
@@ -96,7 +92,7 @@ public class AlarmStoreTest {
         assertEquals(1, mAlarmStore.size());
         assertTrue(removed.contains(a1) && removed.contains(a2));
 
-        final Alarm a8 = createAlarm(8, 0, null);
+        final Alarm a8 = createAlarm(8, 0);
         addAlarmsToStore(a8, a2, a1);
 
         removed = mAlarmStore.remove(unused -> false);
@@ -110,9 +106,9 @@ public class AlarmStoreTest {
 
     @Test
     public void removePendingAlarms() {
-        final Alarm a1to11 = createAlarm(1, 10, null);
-        final Alarm a2to5 = createAlarm(2, 3, null);
-        final Alarm a6to9 = createAlarm(6, 3, null);
+        final Alarm a1to11 = createAlarm(1, 10);
+        final Alarm a2to5 = createAlarm(2, 3);
+        final Alarm a6to9 = createAlarm(6, 3);
         addAlarmsToStore(a2to5, a6to9, a1to11);
 
         final ArrayList<Alarm> pendingAt0 = mAlarmStore.removePendingAlarms(0);
@@ -134,10 +130,10 @@ public class AlarmStoreTest {
 
     @Test
     public void getNextWakeupDeliveryTime() {
-        final Alarm a1to10 = createAlarm(1, 9, null);
-        final Alarm a3to8wakeup = createWakeupAlarm(3, 5, null);
-        final Alarm a6wakeup = createWakeupAlarm(6, 0, null);
-        final Alarm a5 = createAlarm(5, 0, null);
+        final Alarm a1to10 = createAlarm(1, 9);
+        final Alarm a3to8wakeup = createWakeupAlarm(3, 5, 0);
+        final Alarm a6wakeup = createWakeupAlarm(6, 0, 0);
+        final Alarm a5 = createAlarm(5, 0);
         addAlarmsToStore(a5, a6wakeup, a3to8wakeup, a1to10);
 
         // The wakeup alarms are [6] and [3, 8], hence 6 is the latest time till when we can
@@ -155,10 +151,10 @@ public class AlarmStoreTest {
 
     @Test
     public void getNextDeliveryTime() {
-        final Alarm a1to10 = createAlarm(1, 9, null);
-        final Alarm a3to8wakeup = createWakeupAlarm(3, 5, null);
-        final Alarm a6wakeup = createWakeupAlarm(6, 0, null);
-        final Alarm a5 = createAlarm(5, 0, null);
+        final Alarm a1to10 = createAlarm(1, 9);
+        final Alarm a3to8wakeup = createWakeupAlarm(3, 5, 0);
+        final Alarm a6wakeup = createWakeupAlarm(6, 0, 0);
+        final Alarm a5 = createAlarm(5, 0);
         addAlarmsToStore(a5, a6wakeup, a3to8wakeup, a1to10);
 
         assertTrue(mAlarmStore.getNextDeliveryTime() <= 5);
@@ -168,10 +164,32 @@ public class AlarmStoreTest {
     }
 
     @Test
+    public void getNextWakeFromIdle() {
+        final Alarm a3 = createWakeupAlarm(3, 0, AlarmManager.FLAG_WAKE_FROM_IDLE);
+        final Alarm a5 = createWakeupAlarm(5, 0, AlarmManager.FLAG_WAKE_FROM_IDLE);
+        final Alarm a7 = createWakeupAlarm(7, 0, AlarmManager.FLAG_WAKE_FROM_IDLE);
+
+        mAlarmStore.add(a5);
+        assertEquals(a5, mAlarmStore.getNextWakeFromIdleAlarm());
+
+        mAlarmStore.add(a7);
+        assertEquals(a5, mAlarmStore.getNextWakeFromIdleAlarm());
+
+        mAlarmStore.add(a3);
+        assertEquals(a3, mAlarmStore.getNextWakeFromIdleAlarm());
+
+        mAlarmStore.remove(a -> (a == a3) || (a == a5));
+        assertEquals(a7, mAlarmStore.getNextWakeFromIdleAlarm());
+
+        mAlarmStore.remove(a -> (a == a7));
+        assertNull(mAlarmStore.getNextWakeFromIdleAlarm());
+    }
+
+    @Test
     public void updateAlarmDeliveries() {
-        final Alarm a5 = createAlarm(5, 0, null);
-        final Alarm a8 = createAlarm(8, 0, null);
-        final Alarm a10 = createAlarm(10, 0, null);
+        final Alarm a5 = createAlarm(5, 0);
+        final Alarm a8 = createAlarm(8, 0);
+        final Alarm a10 = createAlarm(10, 0);
         addAlarmsToStore(a8, a10, a5);
 
         assertEquals(5, mAlarmStore.getNextDeliveryTime());
diff --git a/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmTest.java b/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmTest.java
index efcfae38c3d3..e80f0655e641 100644
--- a/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmTest.java
@@ -19,6 +19,7 @@ package com.android.server.alarm;
 import static android.app.AlarmManager.ELAPSED_REALTIME;
 
 import static com.android.server.alarm.Alarm.APP_STANDBY_POLICY_INDEX;
+import static com.android.server.alarm.Alarm.NUM_POLICIES;
 import static com.android.server.alarm.Alarm.REQUESTER_POLICY_INDEX;
 import static com.android.server.alarm.Constants.TEST_CALLING_PACKAGE;
 import static com.android.server.alarm.Constants.TEST_CALLING_UID;
@@ -36,6 +37,8 @@ import androidx.test.runner.AndroidJUnit4;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.Random;
+
 @Presubmit
 @RunWith(AndroidJUnit4.class)
 public class AlarmTest {
@@ -49,29 +52,54 @@ public class AlarmTest {
     @Test
     public void initSetsOnlyRequesterPolicy() {
         final Alarm a = createDefaultAlarm(4567, 2);
-        assertEquals(4567, a.getPolicyElapsed(REQUESTER_POLICY_INDEX));
-        assertEquals(0, a.getPolicyElapsed(APP_STANDBY_POLICY_INDEX));
+
+        for (int i = 0; i < NUM_POLICIES; i++) {
+            if (i == REQUESTER_POLICY_INDEX) {
+                assertEquals(4567, a.getPolicyElapsed(i));
+            } else {
+                assertEquals(0, a.getPolicyElapsed(i));
+            }
+        }
+    }
+
+    /**
+     * Generates a long matrix {@code A} of size {@code NxN}, with the property that the {@code i}th
+     * row will have the {@code i}th element largest in that row.
+     *
+     * In other words, {@code A[i][i]} will be the maximum of {@code A[i][j]} over all {@code j},
+     * {@code 0<=j<N}.
+     */
+    private static long[][] generatePolicyTestMatrix(int n) {
+        final long[][] data = new long[n][n];
+        final Random random = new Random(971);
+        for (int i = 0; i < n; i++) {
+            data[i][i] = 1;
+            for (int j = 0; j < n; j++) {
+                if (i != j) {
+                    data[i][j] = random.nextInt(1 << 20);
+                    data[i][i] += data[i][j];
+                }
+            }
+        }
+        return data;
     }
 
     @Test
     public void whenElapsed() {
         final Alarm a = createDefaultAlarm(0, 0);
 
-        a.setPolicyElapsed(REQUESTER_POLICY_INDEX, 4);
-        a.setPolicyElapsed(APP_STANDBY_POLICY_INDEX, 10);
-        assertEquals(10, a.getWhenElapsed());
-
-        a.setPolicyElapsed(REQUESTER_POLICY_INDEX, 12);
-        assertEquals(12, a.getWhenElapsed());
-
-        a.setPolicyElapsed(REQUESTER_POLICY_INDEX, 7);
-        assertEquals(10, a.getWhenElapsed());
-
-        a.setPolicyElapsed(APP_STANDBY_POLICY_INDEX, 2);
-        assertEquals(7, a.getWhenElapsed());
-
-        a.setPolicyElapsed(APP_STANDBY_POLICY_INDEX, 7);
-        assertEquals(7, a.getWhenElapsed());
+        final long[][] uniqueData = generatePolicyTestMatrix(NUM_POLICIES);
+        for (int i = 0; i < NUM_POLICIES; i++) {
+            for (int j = 0; j < NUM_POLICIES; j++) {
+                a.setPolicyElapsed(j, uniqueData[i][j]);
+            }
+            assertEquals(uniqueData[i][i], a.getWhenElapsed());
+        }
+
+        for (int i = 0; i < NUM_POLICIES; i++) {
+            a.setPolicyElapsed(i, 3);
+        }
+        assertEquals(3, a.getWhenElapsed());
     }
 
     @Test
@@ -85,18 +113,21 @@ public class AlarmTest {
         a.setPolicyElapsed(REQUESTER_POLICY_INDEX, 2);
         assertEquals(14, a.getMaxWhenElapsed());
 
-        a.setPolicyElapsed(APP_STANDBY_POLICY_INDEX, 5);
-        assertEquals(14, a.getMaxWhenElapsed());
-
-        a.setPolicyElapsed(APP_STANDBY_POLICY_INDEX, 16);
-        assertEquals(16, a.getMaxWhenElapsed());
-
-        a.setPolicyElapsed(APP_STANDBY_POLICY_INDEX, 12);
-        assertEquals(14, a.getMaxWhenElapsed());
+        for (int i = 0; i < NUM_POLICIES; i++) {
+            if (i == REQUESTER_POLICY_INDEX) {
+                continue;
+            }
+            a.setPolicyElapsed(i, 17);
+            // getWhenElapsed is 17, so getMaxWhenElapsed will return 17 too.
+            assertEquals(17, a.getMaxWhenElapsed());
+
+            a.setPolicyElapsed(i, 5);
+            assertEquals(14, a.getMaxWhenElapsed());
+        }
     }
 
     @Test
-    public void setPolicyElapsed() {
+    public void setPolicyElapsedExact() {
         final Alarm exactAlarm = createDefaultAlarm(10, 0);
 
         assertTrue(exactAlarm.setPolicyElapsed(REQUESTER_POLICY_INDEX, 4));
@@ -108,6 +139,10 @@ public class AlarmTest {
 
         assertTrue(exactAlarm.setPolicyElapsed(REQUESTER_POLICY_INDEX, 7));
 
+    }
+
+    @Test
+    public void setPolicyElapsedInexact() {
         final Alarm inexactAlarm = createDefaultAlarm(10, 5);
 
         assertTrue(inexactAlarm.setPolicyElapsed(REQUESTER_POLICY_INDEX, 4));