Merge "Give apps a grace period for TOP." into sc-dev

2 files changed, 353 insertions, 60 deletions

diff --git a/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java b/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java
index 2f3ac225a190..784c63a6cf1f 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java
@@ -32,7 +32,6 @@ import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.UserIdInt;
 import android.app.ActivityManager;
-import android.app.ActivityManagerInternal;
 import android.app.AlarmManager;
 import android.app.AppGlobals;
 import android.app.IUidObserver;
@@ -348,12 +347,20 @@ public final class QuotaController extends StateController {
     private final SparseBooleanArray mTempAllowlistCache = new SparseBooleanArray();
 
     /**
-     * Mapping of app IDs to the when their temp allowlist grace period ends (in the elapsed
+     * Mapping of UIDs to the when their temp allowlist grace period ends (in the elapsed
      * realtime timebase).
      */
     private final SparseLongArray mTempAllowlistGraceCache = new SparseLongArray();
 
-    private final ActivityManagerInternal mActivityManagerInternal;
+    /** Current set of UIDs in the {@link ActivityManager#PROCESS_STATE_TOP} state. */
+    private final SparseBooleanArray mTopAppCache = new SparseBooleanArray();
+
+    /**
+     * Mapping of UIDs to the when their top app grace period ends (in the elapsed realtime
+     * timebase).
+     */
+    private final SparseLongArray mTopAppGraceCache = new SparseLongArray();
+
     private final AlarmManager mAlarmManager;
     private final ChargingTracker mChargeTracker;
     private final QcHandler mHandler;
@@ -412,7 +419,7 @@ public final class QuotaController extends StateController {
                 }
             };
 
-    private final IUidObserver mUidObserver = new IUidObserver.Stub() {
+    private class QcUidObserver extends IUidObserver.Stub {
         @Override
         public void onUidStateChanged(int uid, int procState, long procStateSeq, int capability) {
             mHandler.obtainMessage(MSG_UID_PROCESS_STATE_CHANGED, uid, procState).sendToTarget();
@@ -433,7 +440,7 @@ public final class QuotaController extends StateController {
         @Override
         public void onUidCachedChanged(int uid, boolean cached) {
         }
-    };
+    }
 
     private final BroadcastReceiver mPackageAddedReceiver = new BroadcastReceiver() {
         @Override
@@ -548,8 +555,10 @@ public final class QuotaController extends StateController {
      */
     private long mEJRewardNotificationSeenMs = QcConstants.DEFAULT_EJ_REWARD_NOTIFICATION_SEEN_MS;
 
-    private long mEJTempAllowlistGracePeriodMs =
-            QcConstants.DEFAULT_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS;
+    private long mEJGracePeriodTempAllowlistMs =
+            QcConstants.DEFAULT_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS;
+
+    private long mEJGracePeriodTopAppMs = QcConstants.DEFAULT_EJ_GRACE_PERIOD_TOP_APP_MS;
 
     /** The package verifier app. */
     @Nullable
@@ -586,7 +595,6 @@ public final class QuotaController extends StateController {
         mHandler = new QcHandler(mContext.getMainLooper());
         mChargeTracker = new ChargingTracker();
         mChargeTracker.startTracking();
-        mActivityManagerInternal = LocalServices.getService(ActivityManagerInternal.class);
         mAlarmManager = (AlarmManager) mContext.getSystemService(Context.ALARM_SERVICE);
         mQcConstants = new QcConstants();
         mBackgroundJobsController = backgroundJobsController;
@@ -606,9 +614,12 @@ public final class QuotaController extends StateController {
         pai.registerTempAllowlistChangeListener(new TempAllowlistTracker());
 
         try {
-            ActivityManager.getService().registerUidObserver(mUidObserver,
+            ActivityManager.getService().registerUidObserver(new QcUidObserver(),
                     ActivityManager.UID_OBSERVER_PROCSTATE,
                     ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE, null);
+            ActivityManager.getService().registerUidObserver(new QcUidObserver(),
+                    ActivityManager.UID_OBSERVER_PROCSTATE,
+                    ActivityManager.PROCESS_STATE_TOP, null);
         } catch (RemoteException e) {
             // ignored; both services live in system_server
         }
@@ -658,7 +669,7 @@ public final class QuotaController extends StateController {
         }
 
         final int uid = jobStatus.getSourceUid();
-        if (mActivityManagerInternal.getUidProcessState(uid) <= ActivityManager.PROCESS_STATE_TOP) {
+        if (mTopAppCache.get(uid)) {
             if (DEBUG) {
                 Slog.d(TAG, jobStatus.toShortString() + " is top started job");
             }
@@ -715,6 +726,8 @@ public final class QuotaController extends StateController {
         mUidToPackageCache.remove(uid);
         mTempAllowlistCache.delete(uid);
         mTempAllowlistGraceCache.delete(uid);
+        mTopAppCache.delete(uid);
+        mTopAppGraceCache.delete(uid);
     }
 
     @Override
@@ -770,14 +783,10 @@ public final class QuotaController extends StateController {
 
     /** Returns the maximum amount of time this job could run for. */
     public long getMaxJobExecutionTimeMsLocked(@NonNull final JobStatus jobStatus) {
-        // Need to look at current proc state as well in the case where the job hasn't started yet.
-        final boolean isTop = mActivityManagerInternal
-                .getUidProcessState(jobStatus.getSourceUid()) <= ActivityManager.PROCESS_STATE_TOP;
-
         if (!jobStatus.shouldTreatAsExpeditedJob()) {
             // If quota is currently "free", then the job can run for the full amount of time.
             if (mChargeTracker.isCharging()
-                    || isTop
+                    || mTopAppCache.get(jobStatus.getSourceUid())
                     || isTopStartedJobLocked(jobStatus)
                     || isUidInForeground(jobStatus.getSourceUid())) {
                 return mConstants.RUNTIME_FREE_QUOTA_MAX_LIMIT_MS;
@@ -790,7 +799,7 @@ public final class QuotaController extends StateController {
         if (mChargeTracker.isCharging()) {
             return mConstants.RUNTIME_FREE_QUOTA_MAX_LIMIT_MS;
         }
-        if (isTop || isTopStartedJobLocked(jobStatus)) {
+        if (mTopAppCache.get(jobStatus.getSourceUid()) || isTopStartedJobLocked(jobStatus)) {
             return Math.max(mEJLimitsMs[ACTIVE_INDEX] / 2,
                     getTimeUntilEJQuotaConsumedLocked(
                             jobStatus.getSourceUserId(), jobStatus.getSourcePackageName()));
@@ -817,14 +826,23 @@ public final class QuotaController extends StateController {
         if (isTopStartedJobLocked(jobStatus) || isUidInForeground(jobStatus.getSourceUid())) {
             return true;
         }
+
+        final long nowElapsed = sElapsedRealtimeClock.millis();
         final long tempAllowlistGracePeriodEndElapsed =
                 mTempAllowlistGraceCache.get(jobStatus.getSourceUid());
         final boolean hasTempAllowlistExemption = mTempAllowlistCache.get(jobStatus.getSourceUid())
-                || sElapsedRealtimeClock.millis() < tempAllowlistGracePeriodEndElapsed;
+                || nowElapsed < tempAllowlistGracePeriodEndElapsed;
         if (hasTempAllowlistExemption) {
             return true;
         }
 
+        final long topAppGracePeriodEndElapsed = mTopAppGraceCache.get(jobStatus.getSourceUid());
+        final boolean hasTopAppExemption = mTopAppCache.get(jobStatus.getSourceUid())
+                || nowElapsed < topAppGracePeriodEndElapsed;
+        if (hasTopAppExemption) {
+            return true;
+        }
+
         Timer ejTimer = mEJPkgTimers.get(jobStatus.getSourceUserId(),
                 jobStatus.getSourcePackageName());
         // Any already executing expedited jobs should be allowed to finish.
@@ -2101,8 +2119,12 @@ public final class QuotaController extends StateController {
             final boolean hasTempAllowlistExemption = !mRegularJobTimer
                     && (mTempAllowlistCache.get(mUid)
                     || nowElapsed < tempAllowlistGracePeriodEndElapsed);
+            final long topAppGracePeriodEndElapsed = mTopAppGraceCache.get(mUid);
+            final boolean hasTopAppExemption = !mRegularJobTimer
+                    && (mTopAppCache.get(mUid) || nowElapsed < topAppGracePeriodEndElapsed);
             return (standbyBucket == RESTRICTED_INDEX || !mChargeTracker.isCharging())
-                    && !mForegroundUids.get(mUid) && !hasTempAllowlistExemption;
+                    && !mForegroundUids.get(mUid) && !hasTempAllowlistExemption
+                    && !hasTopAppExemption;
         }
 
         void onStateChangedLocked(long nowElapsed, boolean isQuotaFree) {
@@ -2421,11 +2443,11 @@ public final class QuotaController extends StateController {
         public void onAppRemoved(int uid) {
             synchronized (mLock) {
                 final long nowElapsed = sElapsedRealtimeClock.millis();
-                final long endElapsed = nowElapsed + mEJTempAllowlistGracePeriodMs;
+                final long endElapsed = nowElapsed + mEJGracePeriodTempAllowlistMs;
                 mTempAllowlistCache.delete(uid);
                 mTempAllowlistGraceCache.put(uid, endElapsed);
                 Message msg = mHandler.obtainMessage(MSG_END_GRACE_PERIOD, uid, 0);
-                mHandler.sendMessageDelayed(msg, mEJTempAllowlistGracePeriodMs);
+                mHandler.sendMessageDelayed(msg, mEJGracePeriodTempAllowlistMs);
             }
         }
     }
@@ -2571,12 +2593,37 @@ public final class QuotaController extends StateController {
 
                         synchronized (mLock) {
                             boolean isQuotaFree;
-                            if (procState <= ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE) {
+                            if (procState <= ActivityManager.PROCESS_STATE_TOP) {
+                                mTopAppCache.put(uid, true);
+                                mTopAppGraceCache.delete(uid);
+                                if (mForegroundUids.get(uid)) {
+                                    // Went from FGS to TOP. We don't need to reprocess timers or
+                                    // jobs.
+                                    break;
+                                }
                                 mForegroundUids.put(uid, true);
                                 isQuotaFree = true;
                             } else {
-                                mForegroundUids.delete(uid);
-                                isQuotaFree = false;
+                                final boolean reprocess;
+                                if (procState <= ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE) {
+                                    reprocess = !mForegroundUids.get(uid);
+                                    mForegroundUids.put(uid, true);
+                                    isQuotaFree = true;
+                                } else {
+                                    reprocess = true;
+                                    mForegroundUids.delete(uid);
+                                    isQuotaFree = false;
+                                }
+                                if (mTopAppCache.get(uid)) {
+                                    final long endElapsed = nowElapsed + mEJGracePeriodTopAppMs;
+                                    mTopAppCache.delete(uid);
+                                    mTopAppGraceCache.put(uid, endElapsed);
+                                    sendMessageDelayed(obtainMessage(MSG_END_GRACE_PERIOD, uid, 0),
+                                            mEJGracePeriodTopAppMs);
+                                }
+                                if (!reprocess) {
+                                    break;
+                                }
                             }
                             // Update Timers first.
                             if (mPkgTimers.indexOfKey(userId) >= 0
@@ -2644,20 +2691,31 @@ public final class QuotaController extends StateController {
                     case MSG_END_GRACE_PERIOD: {
                         final int uid = msg.arg1;
                         synchronized (mLock) {
-                            if (mTempAllowlistCache.get(uid)) {
-                                // App added back to the temp allowlist during the grace period.
+                            if (mTempAllowlistCache.get(uid) || mTopAppCache.get(uid)) {
+                                // App added back to the temp allowlist or became top again
+                                // during the grace period.
                                 if (DEBUG) {
                                     Slog.d(TAG, uid + " is still allowed");
                                 }
                                 break;
                             }
+                            final long nowElapsed = sElapsedRealtimeClock.millis();
+                            if (nowElapsed < mTempAllowlistGraceCache.get(uid)
+                                    || nowElapsed < mTopAppGraceCache.get(uid)) {
+                                // One of the grace periods is still in effect.
+                                if (DEBUG) {
+                                    Slog.d(TAG, uid + " is still in grace period");
+                                }
+                                break;
+                            }
                             if (DEBUG) {
                                 Slog.d(TAG, uid + " is now out of grace period");
                             }
+                            mTempAllowlistGraceCache.delete(uid);
+                            mTopAppGraceCache.delete(uid);
                             final ArraySet<String> packages = getPackagesForUidLocked(uid);
                             if (packages != null) {
                                 final int userId = UserHandle.getUserId(uid);
-                                final long nowElapsed = sElapsedRealtimeClock.millis();
                                 for (int i = packages.size() - 1; i >= 0; --i) {
                                     Timer t = mEJPkgTimers.get(userId, packages.valueAt(i));
                                     if (t != null) {
@@ -2989,8 +3047,11 @@ public final class QuotaController extends StateController {
         static final String KEY_EJ_REWARD_NOTIFICATION_SEEN_MS =
                 QC_CONSTANT_PREFIX + "ej_reward_notification_seen_ms";
         @VisibleForTesting
-        static final String KEY_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS =
-                QC_CONSTANT_PREFIX + "ej_temp_allowlist_grace_period_ms";
+        static final String KEY_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS =
+                QC_CONSTANT_PREFIX + "ej_grace_period_temp_allowlist_ms";
+        @VisibleForTesting
+        static final String KEY_EJ_GRACE_PERIOD_TOP_APP_MS =
+                QC_CONSTANT_PREFIX + "ej_grace_period_top_app_ms";
 
         private static final long DEFAULT_ALLOWED_TIME_PER_PERIOD_MS =
                 10 * 60 * 1000L; // 10 minutes
@@ -3043,7 +3104,8 @@ public final class QuotaController extends StateController {
         private static final long DEFAULT_EJ_REWARD_TOP_APP_MS = 10 * SECOND_IN_MILLIS;
         private static final long DEFAULT_EJ_REWARD_INTERACTION_MS = 15 * SECOND_IN_MILLIS;
         private static final long DEFAULT_EJ_REWARD_NOTIFICATION_SEEN_MS = 0;
-        private static final long DEFAULT_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS = 3 * MINUTE_IN_MILLIS;
+        private static final long DEFAULT_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS = 3 * MINUTE_IN_MILLIS;
+        private static final long DEFAULT_EJ_GRACE_PERIOD_TOP_APP_MS = 1 * MINUTE_IN_MILLIS;
 
         /** How much time each app will have to run jobs within their standby bucket window. */
         public long ALLOWED_TIME_PER_PERIOD_MS = DEFAULT_ALLOWED_TIME_PER_PERIOD_MS;
@@ -3275,7 +3337,12 @@ public final class QuotaController extends StateController {
          * How much additional grace period to add to the end of an app's temp allowlist
          * duration.
          */
-        public long EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS = DEFAULT_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS;
+        public long EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS = DEFAULT_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS;
+
+        /**
+         * How much additional grace period to give an app when it leaves the TOP state.
+         */
+        public long EJ_GRACE_PERIOD_TOP_APP_MS = DEFAULT_EJ_GRACE_PERIOD_TOP_APP_MS;
 
         public void processConstantLocked(@NonNull DeviceConfig.Properties properties,
                 @NonNull String key) {
@@ -3470,13 +3537,21 @@ public final class QuotaController extends StateController {
                     mEJRewardNotificationSeenMs = Math.min(5 * MINUTE_IN_MILLIS,
                             Math.max(0, EJ_REWARD_NOTIFICATION_SEEN_MS));
                     break;
-                case KEY_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS:
+                case KEY_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS:
                     // We don't need to re-evaluate execution stats or constraint status for this.
-                    EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS =
-                            properties.getLong(key, DEFAULT_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS);
+                    EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS =
+                            properties.getLong(key, DEFAULT_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS);
                     // Limit grace period to be in the range [0 minutes, 1 hour].
-                    mEJTempAllowlistGracePeriodMs = Math.min(HOUR_IN_MILLIS,
-                            Math.max(0, EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS));
+                    mEJGracePeriodTempAllowlistMs = Math.min(HOUR_IN_MILLIS,
+                            Math.max(0, EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS));
+                    break;
+                case KEY_EJ_GRACE_PERIOD_TOP_APP_MS:
+                    // We don't need to re-evaluate execution stats or constraint status for this.
+                    EJ_GRACE_PERIOD_TOP_APP_MS =
+                            properties.getLong(key, DEFAULT_EJ_GRACE_PERIOD_TOP_APP_MS);
+                    // Limit grace period to be in the range [0 minutes, 1 hour].
+                    mEJGracePeriodTopAppMs = Math.min(HOUR_IN_MILLIS,
+                            Math.max(0, EJ_GRACE_PERIOD_TOP_APP_MS));
                     break;
             }
         }
@@ -3739,8 +3814,9 @@ public final class QuotaController extends StateController {
             pw.print(KEY_EJ_REWARD_TOP_APP_MS, EJ_REWARD_TOP_APP_MS).println();
             pw.print(KEY_EJ_REWARD_INTERACTION_MS, EJ_REWARD_INTERACTION_MS).println();
             pw.print(KEY_EJ_REWARD_NOTIFICATION_SEEN_MS, EJ_REWARD_NOTIFICATION_SEEN_MS).println();
-            pw.print(KEY_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS,
-                    EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS).println();
+            pw.print(KEY_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS,
+                    EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS).println();
+            pw.print(KEY_EJ_GRACE_PERIOD_TOP_APP_MS, EJ_GRACE_PERIOD_TOP_APP_MS).println();
 
             pw.decreaseIndent();
         }
@@ -3853,6 +3929,16 @@ public final class QuotaController extends StateController {
     }
 
     @VisibleForTesting
+    long getEJGracePeriodTempAllowlistMs() {
+        return mEJGracePeriodTempAllowlistMs;
+    }
+
+    @VisibleForTesting
+    long getEJGracePeriodTopAppMs() {
+        return mEJGracePeriodTopAppMs;
+    }
+
+    @VisibleForTesting
     @NonNull
     long[] getEJLimitsMs() {
         return mEJLimitsMs;
@@ -3888,11 +3974,6 @@ public final class QuotaController extends StateController {
     }
 
     @VisibleForTesting
-    long getEJTempAllowlistGracePeriodMs() {
-        return mEJTempAllowlistGracePeriodMs;
-    }
-
-    @VisibleForTesting
     @Nullable
     List<TimingSession> getEJTimingSessions(int userId, String packageName) {
         return mEJTimingSessions.get(userId, packageName);
@@ -3964,6 +4045,17 @@ public final class QuotaController extends StateController {
         pw.println(mForegroundUids.toString());
         pw.println();
 
+        pw.print("Cached top apps: ");
+        pw.println(mTopAppCache.toString());
+        pw.print("Cached top app grace period: ");
+        pw.println(mTopAppGraceCache.toString());
+
+        pw.print("Cached temp allowlist: ");
+        pw.println(mTempAllowlistCache.toString());
+        pw.print("Cached temp allowlist grace period: ");
+        pw.println(mTempAllowlistGraceCache.toString());
+        pw.println();
+
         pw.println("Cached UID->package map:");
         pw.increaseIndent();
         for (int i = 0; i < mUidToPackageCache.size(); ++i) {
@@ -3975,12 +4067,6 @@ public final class QuotaController extends StateController {
         pw.decreaseIndent();
         pw.println();
 
-        pw.print("Cached temp allowlist: ");
-        pw.println(mTempAllowlistCache.toString());
-        pw.print("Cached temp allowlist grace period: ");
-        pw.println(mTempAllowlistGraceCache.toString());
-
-        pw.println();
         mTrackedJobs.forEach((jobs) -> {
             for (int j = 0; j < jobs.size(); j++) {
                 final JobStatus js = jobs.valueAt(j);
diff --git a/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java b/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java
index 7e4bc1e371b2..281c1aafe049 100644
--- a/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java
@@ -75,6 +75,7 @@ import android.os.Handler;
 import android.os.Looper;
 import android.os.RemoteException;
 import android.os.SystemClock;
+import android.platform.test.annotations.LargeTest;
 import android.provider.DeviceConfig;
 import android.util.SparseBooleanArray;
 
@@ -286,14 +287,20 @@ public class QuotaControllerTest {
             doReturn(procState).when(mActivityMangerInternal).getUidProcessState(uid);
             SparseBooleanArray foregroundUids = mQuotaController.getForegroundUids();
             spyOn(foregroundUids);
+            final boolean contained = foregroundUids.get(uid);
             mUidObserver.onUidStateChanged(uid, procState, 0,
                     ActivityManager.PROCESS_CAPABILITY_NONE);
             if (procState <= ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE) {
-                verify(foregroundUids, timeout(2 * SECOND_IN_MILLIS).times(1))
-                        .put(eq(uid), eq(true));
+                if (!contained) {
+                    verify(foregroundUids, timeout(2 * SECOND_IN_MILLIS).times(1))
+                            .put(eq(uid), eq(true));
+                }
                 assertTrue(foregroundUids.get(uid));
             } else {
-                verify(foregroundUids, timeout(2 * SECOND_IN_MILLIS).times(1)).delete(eq(uid));
+                if (contained) {
+                    verify(foregroundUids, timeout(2 * SECOND_IN_MILLIS).times(1))
+                            .delete(eq(uid));
+                }
                 assertFalse(foregroundUids.get(uid));
             }
             waitForNonDelayedMessagesProcessed();
@@ -1993,7 +2000,7 @@ public class QuotaControllerTest {
 
         setProcessState(ActivityManager.PROCESS_STATE_RECEIVER);
         final long gracePeriodMs = 15 * SECOND_IN_MILLIS;
-        setDeviceConfigLong(QcConstants.KEY_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS, gracePeriodMs);
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS, gracePeriodMs);
         Handler handler = mQuotaController.getHandler();
         spyOn(handler);
 
@@ -2017,6 +2024,52 @@ public class QuotaControllerTest {
         }
     }
 
+    /**
+     * Tests that Timers properly track sessions when an app becomes top and is closed.
+     */
+    @Test
+    public void testIsWithinEJQuotaLocked_TopApp() {
+        setDischarging();
+        JobStatus js = createExpeditedJobStatus("testIsWithinEJQuotaLocked_TopApp", 1);
+        setStandbyBucket(FREQUENT_INDEX, js);
+        setDeviceConfigLong(QcConstants.KEY_EJ_LIMIT_FREQUENT_MS, 10 * MINUTE_IN_MILLIS);
+        final long now = JobSchedulerService.sElapsedRealtimeClock.millis();
+        mQuotaController.saveTimingSession(SOURCE_USER_ID, SOURCE_PACKAGE,
+                createTimingSession(now - (HOUR_IN_MILLIS), 3 * MINUTE_IN_MILLIS, 5), true);
+        mQuotaController.saveTimingSession(SOURCE_USER_ID, SOURCE_PACKAGE,
+                createTimingSession(now - (30 * MINUTE_IN_MILLIS), 3 * MINUTE_IN_MILLIS, 5), true);
+        mQuotaController.saveTimingSession(SOURCE_USER_ID, SOURCE_PACKAGE,
+                createTimingSession(now - (5 * MINUTE_IN_MILLIS), 4 * MINUTE_IN_MILLIS, 5), true);
+        setProcessState(ActivityManager.PROCESS_STATE_RECEIVER);
+        synchronized (mQuotaController.mLock) {
+            assertFalse(mQuotaController.isWithinEJQuotaLocked(js));
+        }
+
+        final long gracePeriodMs = 15 * SECOND_IN_MILLIS;
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TOP_APP_MS, gracePeriodMs);
+        Handler handler = mQuotaController.getHandler();
+        spyOn(handler);
+
+        // Apps on top should be able to schedule & start EJs, even if they're out
+        // of quota (as long as they are in the top grace period).
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        synchronized (mQuotaController.mLock) {
+            assertTrue(mQuotaController.isWithinEJQuotaLocked(js));
+        }
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        // Still in grace period
+        synchronized (mQuotaController.mLock) {
+            assertTrue(mQuotaController.isWithinEJQuotaLocked(js));
+        }
+        advanceElapsedClock(6 * SECOND_IN_MILLIS);
+        // Out of grace period.
+        synchronized (mQuotaController.mLock) {
+            assertFalse(mQuotaController.isWithinEJQuotaLocked(js));
+        }
+    }
+
     @Test
     public void testMaybeScheduleCleanupAlarmLocked() {
         // No sessions saved yet.
@@ -2657,8 +2710,9 @@ public class QuotaControllerTest {
         setDeviceConfigLong(QcConstants.KEY_EJ_REWARD_TOP_APP_MS, 87 * SECOND_IN_MILLIS);
         setDeviceConfigLong(QcConstants.KEY_EJ_REWARD_INTERACTION_MS, 86 * SECOND_IN_MILLIS);
         setDeviceConfigLong(QcConstants.KEY_EJ_REWARD_NOTIFICATION_SEEN_MS, 85 * SECOND_IN_MILLIS);
-        setDeviceConfigLong(QcConstants.KEY_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS,
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS,
                 84 * SECOND_IN_MILLIS);
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TOP_APP_MS, 83 * SECOND_IN_MILLIS);
 
         assertEquals(5 * MINUTE_IN_MILLIS, mQuotaController.getAllowedTimePerPeriodMs());
         assertEquals(2 * MINUTE_IN_MILLIS, mQuotaController.getInQuotaBufferMs());
@@ -2697,7 +2751,8 @@ public class QuotaControllerTest {
         assertEquals(87 * SECOND_IN_MILLIS, mQuotaController.getEJRewardTopAppMs());
         assertEquals(86 * SECOND_IN_MILLIS, mQuotaController.getEJRewardInteractionMs());
         assertEquals(85 * SECOND_IN_MILLIS, mQuotaController.getEJRewardNotificationSeenMs());
-        assertEquals(84 * SECOND_IN_MILLIS, mQuotaController.getEJTempAllowlistGracePeriodMs());
+        assertEquals(84 * SECOND_IN_MILLIS, mQuotaController.getEJGracePeriodTempAllowlistMs());
+        assertEquals(83 * SECOND_IN_MILLIS, mQuotaController.getEJGracePeriodTopAppMs());
     }
 
     @Test
@@ -2737,7 +2792,8 @@ public class QuotaControllerTest {
         setDeviceConfigLong(QcConstants.KEY_EJ_REWARD_TOP_APP_MS, -1);
         setDeviceConfigLong(QcConstants.KEY_EJ_REWARD_INTERACTION_MS, -1);
         setDeviceConfigLong(QcConstants.KEY_EJ_REWARD_NOTIFICATION_SEEN_MS, -1);
-        setDeviceConfigLong(QcConstants.KEY_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS, -1);
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS, -1);
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TOP_APP_MS, -1);
 
         assertEquals(MINUTE_IN_MILLIS, mQuotaController.getAllowedTimePerPeriodMs());
         assertEquals(0, mQuotaController.getInQuotaBufferMs());
@@ -2773,7 +2829,8 @@ public class QuotaControllerTest {
         assertEquals(10 * SECOND_IN_MILLIS, mQuotaController.getEJRewardTopAppMs());
         assertEquals(5 * SECOND_IN_MILLIS, mQuotaController.getEJRewardInteractionMs());
         assertEquals(0, mQuotaController.getEJRewardNotificationSeenMs());
-        assertEquals(0, mQuotaController.getEJTempAllowlistGracePeriodMs());
+        assertEquals(0, mQuotaController.getEJGracePeriodTempAllowlistMs());
+        assertEquals(0, mQuotaController.getEJGracePeriodTopAppMs());
 
         // Invalid configurations.
         // In_QUOTA_BUFFER should never be greater than ALLOWED_TIME_PER_PERIOD
@@ -2807,7 +2864,8 @@ public class QuotaControllerTest {
         setDeviceConfigLong(QcConstants.KEY_EJ_REWARD_TOP_APP_MS, 25 * HOUR_IN_MILLIS);
         setDeviceConfigLong(QcConstants.KEY_EJ_REWARD_INTERACTION_MS, 25 * HOUR_IN_MILLIS);
         setDeviceConfigLong(QcConstants.KEY_EJ_REWARD_NOTIFICATION_SEEN_MS, 25 * HOUR_IN_MILLIS);
-        setDeviceConfigLong(QcConstants.KEY_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS, 25 * HOUR_IN_MILLIS);
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS, 25 * HOUR_IN_MILLIS);
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TOP_APP_MS, 25 * HOUR_IN_MILLIS);
 
         assertEquals(24 * HOUR_IN_MILLIS, mQuotaController.getAllowedTimePerPeriodMs());
         assertEquals(5 * MINUTE_IN_MILLIS, mQuotaController.getInQuotaBufferMs());
@@ -2833,7 +2891,8 @@ public class QuotaControllerTest {
         assertEquals(15 * MINUTE_IN_MILLIS, mQuotaController.getEJRewardTopAppMs());
         assertEquals(15 * MINUTE_IN_MILLIS, mQuotaController.getEJRewardInteractionMs());
         assertEquals(5 * MINUTE_IN_MILLIS, mQuotaController.getEJRewardNotificationSeenMs());
-        assertEquals(HOUR_IN_MILLIS, mQuotaController.getEJTempAllowlistGracePeriodMs());
+        assertEquals(HOUR_IN_MILLIS, mQuotaController.getEJGracePeriodTempAllowlistMs());
+        assertEquals(HOUR_IN_MILLIS, mQuotaController.getEJGracePeriodTopAppMs());
     }
 
     /** Tests that TimingSessions aren't saved when the device is charging. */
@@ -3433,7 +3492,7 @@ public class QuotaControllerTest {
         setDischarging();
         setProcessState(ActivityManager.PROCESS_STATE_RECEIVER);
         final long gracePeriodMs = 15 * SECOND_IN_MILLIS;
-        setDeviceConfigLong(QcConstants.KEY_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS, gracePeriodMs);
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS, gracePeriodMs);
         Handler handler = mQuotaController.getHandler();
         spyOn(handler);
 
@@ -4910,6 +4969,7 @@ public class QuotaControllerTest {
     @Test
     public void testEJTimerTracking_TopAndNonTop() {
         setDischarging();
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TOP_APP_MS, 0);
 
         JobStatus jobBg1 = createExpeditedJobStatus("testEJTimerTracking_TopAndNonTop", 1);
         JobStatus jobBg2 = createExpeditedJobStatus("testEJTimerTracking_TopAndNonTop", 2);
@@ -5032,7 +5092,7 @@ public class QuotaControllerTest {
         setDischarging();
         setProcessState(ActivityManager.PROCESS_STATE_RECEIVER);
         final long gracePeriodMs = 15 * SECOND_IN_MILLIS;
-        setDeviceConfigLong(QcConstants.KEY_EJ_TEMP_ALLOWLIST_GRACE_PERIOD_MS, gracePeriodMs);
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS, gracePeriodMs);
         Handler handler = mQuotaController.getHandler();
         spyOn(handler);
 
@@ -5134,11 +5194,158 @@ public class QuotaControllerTest {
     }
 
     /**
+     * Tests that Timers properly track sessions when TOP state and temp allowlisting overlaps.
+     */
+    @Test
+    @LargeTest
+    public void testEJTimerTracking_TopAndTempAllowlisting() throws Exception {
+        setDischarging();
+        setProcessState(ActivityManager.PROCESS_STATE_RECEIVER);
+        final long gracePeriodMs = 5 * SECOND_IN_MILLIS;
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TEMP_ALLOWLIST_MS, gracePeriodMs);
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TOP_APP_MS, gracePeriodMs);
+        Handler handler = mQuotaController.getHandler();
+        spyOn(handler);
+
+        JobStatus job1 = createExpeditedJobStatus("testEJTimerTracking_TopAndTempAllowlisting", 1);
+        JobStatus job2 = createExpeditedJobStatus("testEJTimerTracking_TopAndTempAllowlisting", 2);
+        JobStatus job3 = createExpeditedJobStatus("testEJTimerTracking_TopAndTempAllowlisting", 3);
+        JobStatus job4 = createExpeditedJobStatus("testEJTimerTracking_TopAndTempAllowlisting", 4);
+        JobStatus job5 = createExpeditedJobStatus("testEJTimerTracking_TopAndTempAllowlisting", 5);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.maybeStartTrackingJobLocked(job1, null);
+        }
+        assertNull(mQuotaController.getEJTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
+        List<TimingSession> expected = new ArrayList<>();
+
+        // Case 1: job starts in TA grace period then app becomes TOP
+        long start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        mTempAllowlistListener.onAppAdded(mSourceUid);
+        mTempAllowlistListener.onAppRemoved(mSourceUid);
+        advanceElapsedClock(gracePeriodMs / 2);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.prepareForExecutionLocked(job1);
+        }
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        advanceElapsedClock(gracePeriodMs);
+        // Wait for the grace period to expire so the handler can process the message.
+        Thread.sleep(gracePeriodMs);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.maybeStopTrackingJobLocked(job1, job1, true);
+        }
+        assertNull(mQuotaController.getEJTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
+
+        advanceElapsedClock(gracePeriodMs);
+
+        // Case 2: job starts in TOP grace period then is TAed
+        start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        setProcessState(ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
+        advanceElapsedClock(gracePeriodMs / 2);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.maybeStartTrackingJobLocked(job2, null);
+            mQuotaController.prepareForExecutionLocked(job2);
+        }
+        mTempAllowlistListener.onAppAdded(mSourceUid);
+        advanceElapsedClock(gracePeriodMs);
+        // Wait for the grace period to expire so the handler can process the message.
+        Thread.sleep(gracePeriodMs);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.maybeStopTrackingJobLocked(job2, null, false);
+        }
+        assertNull(mQuotaController.getEJTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
+
+        advanceElapsedClock(gracePeriodMs);
+
+        // Case 3: job starts in TA grace period then app becomes TOP; job ends after TOP grace
+        mTempAllowlistListener.onAppAdded(mSourceUid);
+        mTempAllowlistListener.onAppRemoved(mSourceUid);
+        advanceElapsedClock(gracePeriodMs / 2);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.maybeStartTrackingJobLocked(job3, null);
+            mQuotaController.prepareForExecutionLocked(job3);
+        }
+        advanceElapsedClock(SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        advanceElapsedClock(gracePeriodMs);
+        // Wait for the grace period to expire so the handler can process the message.
+        Thread.sleep(gracePeriodMs);
+        setProcessState(ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
+        advanceElapsedClock(gracePeriodMs);
+        start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        // Wait for the grace period to expire so the handler can process the message.
+        Thread.sleep(2 * gracePeriodMs);
+        advanceElapsedClock(gracePeriodMs);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.maybeStopTrackingJobLocked(job3, job3, true);
+        }
+        expected.add(createTimingSession(start, gracePeriodMs, 1));
+        assertEquals(expected,
+                mQuotaController.getEJTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
+
+        advanceElapsedClock(gracePeriodMs);
+
+        // Case 4: job starts in TOP grace period then app becomes TAed; job ends after TA grace
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        setProcessState(ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
+        advanceElapsedClock(gracePeriodMs / 2);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.maybeStartTrackingJobLocked(job4, null);
+            mQuotaController.prepareForExecutionLocked(job4);
+        }
+        advanceElapsedClock(SECOND_IN_MILLIS);
+        mTempAllowlistListener.onAppAdded(mSourceUid);
+        advanceElapsedClock(gracePeriodMs);
+        // Wait for the grace period to expire so the handler can process the message.
+        Thread.sleep(gracePeriodMs);
+        mTempAllowlistListener.onAppRemoved(mSourceUid);
+        advanceElapsedClock(gracePeriodMs);
+        start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        // Wait for the grace period to expire so the handler can process the message.
+        Thread.sleep(2 * gracePeriodMs);
+        advanceElapsedClock(gracePeriodMs);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.maybeStopTrackingJobLocked(job4, job4, true);
+        }
+        expected.add(createTimingSession(start, gracePeriodMs, 1));
+        assertEquals(expected,
+                mQuotaController.getEJTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
+
+        advanceElapsedClock(gracePeriodMs);
+
+        // Case 5: job starts during overlapping grace period
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        advanceElapsedClock(SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
+        advanceElapsedClock(SECOND_IN_MILLIS);
+        mTempAllowlistListener.onAppAdded(mSourceUid);
+        advanceElapsedClock(SECOND_IN_MILLIS);
+        mTempAllowlistListener.onAppRemoved(mSourceUid);
+        advanceElapsedClock(gracePeriodMs - SECOND_IN_MILLIS);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.maybeStartTrackingJobLocked(job5, null);
+            mQuotaController.prepareForExecutionLocked(job5);
+        }
+        advanceElapsedClock(SECOND_IN_MILLIS);
+        start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        // Wait for the grace period to expire so the handler can process the message.
+        Thread.sleep(2 * gracePeriodMs);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.maybeStopTrackingJobLocked(job5, job5, true);
+        }
+        expected.add(createTimingSession(start, 10 * SECOND_IN_MILLIS, 1));
+        assertEquals(expected,
+                mQuotaController.getEJTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
+    }
+
+    /**
      * Tests that expedited jobs aren't stopped when an app runs out of quota.
      */
     @Test
     public void testEJTracking_OutOfQuota_ForegroundAndBackground() {
         setDischarging();
+        setDeviceConfigLong(QcConstants.KEY_EJ_GRACE_PERIOD_TOP_APP_MS, 0);
 
         JobStatus jobBg =
                 createExpeditedJobStatus("testEJTracking_OutOfQuota_ForegroundAndBackground", 1);