Using proc state to determine foreground status.

uidActive is true for bound foreground services. We would like to put a
quota on jobs for those processes as well, so switching to proc state
allows finer-grained control.

Bug: 117846754
Bug: 111423978
Test: atest com.android.server.job.controllers.QuotaControllerTest
Change-Id: I69b474b3dc0dda7a4d2234d75fd9023c3f041b67

3 files changed, 502 insertions, 36 deletions

diff --git a/core/proto/android/server/jobscheduler.proto b/core/proto/android/server/jobscheduler.proto
index 0ec8c1ada47e..7f3ea7a249ba 100644
--- a/core/proto/android/server/jobscheduler.proto
+++ b/core/proto/android/server/jobscheduler.proto
@@ -403,18 +403,23 @@ message StateControllerProto {
         optional bool is_charging = 1;
         optional bool is_in_parole = 2;
 
+        // List of UIDs currently in the foreground.
+        repeated int32 foreground_uids = 3;
+
         message TrackedJob {
             option (.android.msg_privacy).dest = DEST_AUTOMATIC;
 
             optional JobStatusShortInfoProto info = 1;
             optional int32 source_uid = 2;
             optional JobStatusDumpProto.Bucket effective_standby_bucket = 3;
-            optional bool has_quota = 4;
+            // If the job started while the app was in the TOP state.
+            optional bool is_top_started_job = 4;
+            optional bool has_quota = 5;
             // The amount of time that this job has remaining in its quota. This
             // can be negative if the job is out of quota.
-            optional int64 remaining_quota_ms = 5;
+            optional int64 remaining_quota_ms = 6;
         }
-        repeated TrackedJob tracked_jobs = 3;
+        repeated TrackedJob tracked_jobs = 4;
 
         message Package {
             option (.android.msg_privacy).dest = DEST_AUTOMATIC;
@@ -456,7 +461,7 @@ message StateControllerProto {
 
             repeated TimingSession saved_sessions = 3;
         }
-        repeated PackageStats package_stats = 4;
+        repeated PackageStats package_stats = 5;
     }
     message StorageController {
         option (.android.msg_privacy).dest = DEST_AUTOMATIC;
diff --git a/services/core/java/com/android/server/job/controllers/QuotaController.java b/services/core/java/com/android/server/job/controllers/QuotaController.java
index ac2dbdf9450e..c16d1b4ecec5 100644
--- a/services/core/java/com/android/server/job/controllers/QuotaController.java
+++ b/services/core/java/com/android/server/job/controllers/QuotaController.java
@@ -26,7 +26,10 @@ import static com.android.server.job.JobSchedulerService.sElapsedRealtimeClock;
 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.IUidObserver;
 import android.app.usage.UsageStatsManagerInternal;
 import android.app.usage.UsageStatsManagerInternal.AppIdleStateChangeListener;
 import android.content.BroadcastReceiver;
@@ -38,12 +41,14 @@ import android.os.BatteryManagerInternal;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
+import android.os.RemoteException;
 import android.os.UserHandle;
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Log;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.util.SparseBooleanArray;
 import android.util.proto.ProtoOutputStream;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -69,6 +74,11 @@ import java.util.function.Predicate;
  * bucket, it will be eligible to run. When a job's bucket changes, its new quota is immediately
  * applied to it.
  *
+ * Jobs are throttled while an app is not in a foreground state. All jobs are allowed to run
+ * freely when an app enters the foreground state and are restricted when the app leaves the
+ * foreground state. However, jobs that are started while the app is in the TOP state are not
+ * restricted regardless of the app's state change.
+ *
  * Test: atest com.android.server.job.controllers.QuotaControllerTest
  */
 public final class QuotaController extends StateController {
@@ -97,6 +107,12 @@ public final class QuotaController extends StateController {
             data.put(packageName, obj);
         }
 
+        public void clear() {
+            for (int i = 0; i < mData.size(); ++i) {
+                mData.valueAt(i).clear();
+            }
+        }
+
         /** Removes all the data for the user, if there was any. */
         public void delete(int userId) {
             mData.delete(userId);
@@ -119,6 +135,11 @@ public final class QuotaController extends StateController {
             return null;
         }
 
+        /** @see SparseArray#indexOfKey */
+        public int indexOfKey(int userId) {
+            return mData.indexOfKey(userId);
+        }
+
         /** Returns the userId at the given index. */
         public int keyAt(int index) {
             return mData.keyAt(index);
@@ -294,6 +315,17 @@ public final class QuotaController extends StateController {
     /** Cached calculation results for each app, with the standby buckets as the array indices. */
     private final UserPackageMap<ExecutionStats[]> mExecutionStatsCache = new UserPackageMap<>();
 
+    /** List of UIDs currently in the foreground. */
+    private final SparseBooleanArray mForegroundUids = new SparseBooleanArray();
+
+    /**
+     * List of jobs that started while the UID was in the TOP state. There will be no more than
+     * 16 ({@link JobSchedulerService.MAX_JOB_CONTEXTS_COUNT}) running at once, so an ArraySet is
+     * fine.
+     */
+    private final ArraySet<JobStatus> mTopStartedJobs = new ArraySet<>();
+
+    private final ActivityManagerInternal mActivityManagerInternal;
     private final AlarmManager mAlarmManager;
     private final ChargingTracker mChargeTracker;
     private final Handler mHandler;
@@ -343,6 +375,29 @@ public final class QuotaController extends StateController {
                 }
             };
 
+    private final IUidObserver mUidObserver = new IUidObserver.Stub() {
+        @Override
+        public void onUidStateChanged(int uid, int procState, long procStateSeq) {
+            mHandler.obtainMessage(MSG_UID_PROCESS_STATE_CHANGED, uid, procState).sendToTarget();
+        }
+
+        @Override
+        public void onUidGone(int uid, boolean disabled) {
+        }
+
+        @Override
+        public void onUidActive(int uid) {
+        }
+
+        @Override
+        public void onUidIdle(int uid, boolean disabled) {
+        }
+
+        @Override
+        public void onUidCachedChanged(int uid, boolean cached) {
+        }
+    };
+
     /**
      * The rolling window size for each standby bucket. Within each window, an app will have 10
      * minutes to run its jobs.
@@ -363,12 +418,15 @@ public final class QuotaController extends StateController {
     private static final int MSG_CLEAN_UP_SESSIONS = 1;
     /** Check if a package is now within its quota. */
     private static final int MSG_CHECK_PACKAGE = 2;
+    /** Process state for a UID has changed. */
+    private static final int MSG_UID_PROCESS_STATE_CHANGED = 3;
 
     public QuotaController(JobSchedulerService service) {
         super(service);
         mHandler = new QcHandler(mContext.getMainLooper());
         mChargeTracker = new ChargingTracker();
         mChargeTracker.startTracking();
+        mActivityManagerInternal = LocalServices.getService(ActivityManagerInternal.class);
         mAlarmManager = (AlarmManager) mContext.getSystemService(Context.ALARM_SERVICE);
 
         // Set up the app standby bucketing tracker
@@ -376,6 +434,14 @@ public final class QuotaController extends StateController {
                 UsageStatsManagerInternal.class);
         usageStats.addAppIdleStateChangeListener(new StandbyTracker());
 
+        try {
+            ActivityManager.getService().registerUidObserver(mUidObserver,
+                    ActivityManager.UID_OBSERVER_PROCSTATE,
+                    ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE, null);
+        } catch (RemoteException e) {
+            // ignored; both services live in system_server
+        }
+
         onConstantsUpdatedLocked();
     }
 
@@ -399,11 +465,15 @@ public final class QuotaController extends StateController {
         if (DEBUG) Slog.d(TAG, "Prepping for " + jobStatus.toShortString());
         final int userId = jobStatus.getSourceUserId();
         final String packageName = jobStatus.getSourcePackageName();
+        final int uid = jobStatus.getSourceUid();
         Timer timer = mPkgTimers.get(userId, packageName);
         if (timer == null) {
-            timer = new Timer(userId, packageName);
+            timer = new Timer(uid, userId, packageName);
             mPkgTimers.add(userId, packageName, timer);
         }
+        if (mActivityManagerInternal.getUidProcessState(uid) == ActivityManager.PROCESS_STATE_TOP) {
+            mTopStartedJobs.add(jobStatus);
+        }
         timer.startTrackingJob(jobStatus);
     }
 
@@ -421,6 +491,7 @@ public final class QuotaController extends StateController {
             if (jobs != null) {
                 jobs.remove(jobStatus);
             }
+            mTopStartedJobs.remove(jobStatus);
         }
     }
 
@@ -511,6 +582,7 @@ public final class QuotaController extends StateController {
             mInQuotaAlarmListeners.delete(userId, packageName);
         }
         mExecutionStatsCache.delete(userId, packageName);
+        mForegroundUids.delete(uid);
     }
 
     @Override
@@ -522,6 +594,20 @@ public final class QuotaController extends StateController {
         mExecutionStatsCache.delete(userId);
     }
 
+    private boolean isUidInForeground(int uid) {
+        if (UserHandle.isCore(uid)) {
+            return true;
+        }
+        synchronized (mLock) {
+            return mForegroundUids.get(uid);
+        }
+    }
+
+    /** @return true if the job was started while the app was in the TOP state. */
+    private boolean isTopStartedJob(@NonNull final JobStatus jobStatus) {
+        return mTopStartedJobs.contains(jobStatus);
+    }
+
     /**
      * Returns an appropriate standby bucket for the job, taking into account any standby
      * exemptions.
@@ -537,9 +623,15 @@ public final class QuotaController extends StateController {
 
     private boolean isWithinQuotaLocked(@NonNull final JobStatus jobStatus) {
         final int standbyBucket = getEffectiveStandbyBucket(jobStatus);
-        // Jobs for the active app should always be able to run.
-        return jobStatus.uidActive || isWithinQuotaLocked(
-                jobStatus.getSourceUserId(), jobStatus.getSourcePackageName(), standbyBucket);
+        Timer timer = mPkgTimers.get(jobStatus.getSourceUserId(), jobStatus.getSourcePackageName());
+        // A job is within quota if one of the following is true:
+        //   1. it was started while the app was in the TOP state
+        //   2. the app is currently in the foreground
+        //   3. the app overall is within its quota
+        return isTopStartedJob(jobStatus)
+                || isUidInForeground(jobStatus.getSourceUid())
+                || isWithinQuotaLocked(
+                      jobStatus.getSourceUserId(), jobStatus.getSourcePackageName(), standbyBucket);
     }
 
     private boolean isWithinQuotaLocked(final int userId, @NonNull final String packageName,
@@ -800,7 +892,7 @@ public final class QuotaController extends StateController {
         final boolean isCharging = mChargeTracker.isCharging();
         if (DEBUG) Slog.d(TAG, "handleNewChargingStateLocked: " + isCharging);
         // Deal with Timers first.
-        mPkgTimers.forEach((t) -> t.onChargingChanged(nowElapsed, isCharging));
+        mPkgTimers.forEach((t) -> t.onStateChanged(nowElapsed, isCharging));
         // Now update jobs.
         maybeUpdateAllConstraintsLocked();
     }
@@ -837,10 +929,15 @@ public final class QuotaController extends StateController {
         boolean changed = false;
         for (int i = jobs.size() - 1; i >= 0; --i) {
             final JobStatus js = jobs.valueAt(i);
-            if (js.uidActive) {
-                // Jobs for the active app should always be able to run.
+            if (isTopStartedJob(js)) {
+                // Job was started while the app was in the TOP state so we should allow it to
+                // finish.
                 changed |= js.setQuotaConstraintSatisfied(true);
-            } else if (realStandbyBucket == getEffectiveStandbyBucket(js)) {
+            } else if (realStandbyBucket != ACTIVE_INDEX
+                    && realStandbyBucket == getEffectiveStandbyBucket(js)) {
+                // An app in the ACTIVE bucket may be out of quota while the job could be in quota
+                // for some reason. Therefore, avoid setting the real value here and check each job
+                // individually.
                 changed |= js.setQuotaConstraintSatisfied(realInQuota);
             } else {
                 // This job is somehow exempted. Need to determine its own quota status.
@@ -854,7 +951,7 @@ public final class QuotaController extends StateController {
             maybeScheduleStartAlarmLocked(userId, packageName, realStandbyBucket);
         } else {
             QcAlarmListener alarmListener = mInQuotaAlarmListeners.get(userId, packageName);
-            if (alarmListener != null) {
+            if (alarmListener != null && alarmListener.isWaiting()) {
                 mAlarmManager.cancel(alarmListener);
                 // Set the trigger time to 0 so that the alarm doesn't think it's still waiting.
                 alarmListener.setTriggerTime(0);
@@ -863,6 +960,56 @@ public final class QuotaController extends StateController {
         return changed;
     }
 
+    private class UidConstraintUpdater implements Consumer<JobStatus> {
+        private final UserPackageMap<Integer> mToScheduleStartAlarms = new UserPackageMap<>();
+        public boolean wasJobChanged;
+
+        @Override
+        public void accept(JobStatus jobStatus) {
+            wasJobChanged |= jobStatus.setQuotaConstraintSatisfied(isWithinQuotaLocked(jobStatus));
+            final int userId = jobStatus.getSourceUserId();
+            final String packageName = jobStatus.getSourcePackageName();
+            final int realStandbyBucket = jobStatus.getStandbyBucket();
+            if (isWithinQuotaLocked(userId, packageName, realStandbyBucket)) {
+                QcAlarmListener alarmListener = mInQuotaAlarmListeners.get(userId, packageName);
+                if (alarmListener != null && alarmListener.isWaiting()) {
+                    mAlarmManager.cancel(alarmListener);
+                    // Set the trigger time to 0 so that the alarm doesn't think it's still waiting.
+                    alarmListener.setTriggerTime(0);
+                }
+            } else {
+                mToScheduleStartAlarms.add(userId, packageName, realStandbyBucket);
+            }
+        }
+
+        void postProcess() {
+            for (int u = 0; u < mToScheduleStartAlarms.numUsers(); ++u) {
+                final int userId = mToScheduleStartAlarms.keyAt(u);
+                for (int p = 0; p < mToScheduleStartAlarms.numPackagesForUser(userId); ++p) {
+                    final String packageName = mToScheduleStartAlarms.keyAt(u, p);
+                    final int standbyBucket = mToScheduleStartAlarms.get(userId, packageName);
+                    maybeScheduleStartAlarmLocked(userId, packageName, standbyBucket);
+                }
+            }
+        }
+
+        void reset() {
+            wasJobChanged = false;
+            mToScheduleStartAlarms.clear();
+        }
+    }
+
+    private final UidConstraintUpdater mUpdateUidConstraints = new UidConstraintUpdater();
+
+    private boolean maybeUpdateConstraintForUidLocked(final int uid) {
+        mService.getJobStore().forEachJobForSourceUid(uid, mUpdateUidConstraints);
+
+        mUpdateUidConstraints.postProcess();
+        boolean changed = mUpdateUidConstraints.wasJobChanged;
+        mUpdateUidConstraints.reset();
+        return changed;
+    }
+
     /**
      * Maybe schedule a non-wakeup alarm for the next time this package will have quota to run
      * again. This should only be called if the package is already out of quota.
@@ -1052,6 +1199,7 @@ public final class QuotaController extends StateController {
 
     private final class Timer {
         private final Package mPkg;
+        private final int mUid;
 
         // List of jobs currently running for this app that started when the app wasn't in the
         // foreground.
@@ -1059,16 +1207,18 @@ public final class QuotaController extends StateController {
         private long mStartTimeElapsed;
         private int mBgJobCount;
 
-        Timer(int userId, String packageName) {
+        Timer(int uid, int userId, String packageName) {
             mPkg = new Package(userId, packageName);
+            mUid = uid;
         }
 
         void startTrackingJob(@NonNull JobStatus jobStatus) {
-            if (jobStatus.uidActive) {
-                // We intentionally don't pay attention to fg state changes after a job has started.
+            if (isTopStartedJob(jobStatus)) {
+                // We intentionally don't pay attention to fg state changes after a TOP job has
+                // started.
                 if (DEBUG) {
                     Slog.v(TAG,
-                            "Timer ignoring " + jobStatus.toShortString() + " because uidActive");
+                            "Timer ignoring " + jobStatus.toShortString() + " because isTop");
                 }
                 return;
             }
@@ -1076,7 +1226,7 @@ public final class QuotaController extends StateController {
             synchronized (mLock) {
                 // Always track jobs, even when charging.
                 mRunningBgJobs.add(jobStatus);
-                if (!mChargeTracker.isCharging()) {
+                if (shouldTrackLocked()) {
                     mBgJobCount++;
                     if (mRunningBgJobs.size() == 1) {
                         // Started tracking the first job.
@@ -1142,6 +1292,10 @@ public final class QuotaController extends StateController {
             }
         }
 
+        boolean isRunning(JobStatus jobStatus) {
+            return mRunningBgJobs.contains(jobStatus);
+        }
+
         long getCurrentDuration(long nowElapsed) {
             synchronized (mLock) {
                 return !isActive() ? 0 : nowElapsed - mStartTimeElapsed;
@@ -1154,17 +1308,21 @@ public final class QuotaController extends StateController {
             }
         }
 
-        void onChargingChanged(long nowElapsed, boolean isCharging) {
+        private boolean shouldTrackLocked() {
+            return !mChargeTracker.isCharging() && !mForegroundUids.get(mUid);
+        }
+
+        void onStateChanged(long nowElapsed, boolean isQuotaFree) {
             synchronized (mLock) {
-                if (isCharging) {
+                if (isQuotaFree) {
                     emitSessionLocked(nowElapsed);
-                } else {
+                } else if (shouldTrackLocked()) {
                     // Start timing from unplug.
                     if (mRunningBgJobs.size() > 0) {
                         mStartTimeElapsed = nowElapsed;
                         // NOTE: this does have the unfortunate consequence that if the device is
-                        // repeatedly plugged in and unplugged, the job count for a package may be
-                        // artificially high.
+                        // repeatedly plugged in and unplugged, or an app changes foreground state
+                        // very frequently, the job count for a package may be artificially high.
                         mBgJobCount = mRunningBgJobs.size();
                         // Starting the timer means that all cached execution stats are now
                         // incorrect.
@@ -1371,6 +1529,38 @@ public final class QuotaController extends StateController {
                         }
                         break;
                     }
+                    case MSG_UID_PROCESS_STATE_CHANGED: {
+                        final int uid = msg.arg1;
+                        final int procState = msg.arg2;
+                        final int userId = UserHandle.getUserId(uid);
+                        final long nowElapsed = sElapsedRealtimeClock.millis();
+
+                        synchronized (mLock) {
+                            boolean isQuotaFree;
+                            if (procState <= ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE) {
+                                mForegroundUids.put(uid, true);
+                                isQuotaFree = true;
+                            } else {
+                                mForegroundUids.delete(uid);
+                                isQuotaFree = false;
+                            }
+                            // Update Timers first.
+                            final int userIndex = mPkgTimers.indexOfKey(userId);
+                            if (userIndex != -1) {
+                                final int numPkgs = mPkgTimers.numPackagesForUser(userId);
+                                for (int p = 0; p < numPkgs; ++p) {
+                                    Timer t = mPkgTimers.valueAt(userIndex, p);
+                                    if (t != null) {
+                                        t.onStateChanged(nowElapsed, isQuotaFree);
+                                    }
+                                }
+                            }
+                            if (maybeUpdateConstraintForUidLocked(uid)) {
+                                mStateChangedListener.onControllerStateChanged();
+                            }
+                        }
+                        break;
+                    }
                 }
             }
         }
@@ -1420,6 +1610,12 @@ public final class QuotaController extends StateController {
 
     @VisibleForTesting
     @NonNull
+    SparseBooleanArray getForegroundUids() {
+        return mForegroundUids;
+    }
+
+    @VisibleForTesting
+    @NonNull
     Handler getHandler() {
         return mHandler;
     }
@@ -1450,6 +1646,10 @@ public final class QuotaController extends StateController {
         pw.println("In parole: " + mInParole);
         pw.println();
 
+        pw.print("Foreground UIDs: ");
+        pw.println(mForegroundUids.toString());
+        pw.println();
+
         mTrackedJobs.forEach((jobs) -> {
             for (int j = 0; j < jobs.size(); j++) {
                 final JobStatus js = jobs.valueAt(j);
@@ -1460,6 +1660,9 @@ public final class QuotaController extends StateController {
                 js.printUniqueId(pw);
                 pw.print(" from ");
                 UserHandle.formatUid(pw, js.getSourceUid());
+                if (mTopStartedJobs.contains(js)) {
+                    pw.print(" (TOP)");
+                }
                 pw.println();
 
                 pw.increaseIndent();
@@ -1511,6 +1714,11 @@ public final class QuotaController extends StateController {
         proto.write(StateControllerProto.QuotaController.IS_CHARGING, mChargeTracker.isCharging());
         proto.write(StateControllerProto.QuotaController.IS_IN_PAROLE, mInParole);
 
+        for (int i = 0; i < mForegroundUids.size(); ++i) {
+            proto.write(StateControllerProto.QuotaController.FOREGROUND_UIDS,
+                    mForegroundUids.keyAt(i));
+        }
+
         mTrackedJobs.forEach((jobs) -> {
             for (int j = 0; j < jobs.size(); j++) {
                 final JobStatus js = jobs.valueAt(j);
@@ -1526,6 +1734,8 @@ public final class QuotaController extends StateController {
                 proto.write(
                         StateControllerProto.QuotaController.TrackedJob.EFFECTIVE_STANDBY_BUCKET,
                         getEffectiveStandbyBucket(js));
+                proto.write(StateControllerProto.QuotaController.TrackedJob.IS_TOP_STARTED_JOB,
+                        mTopStartedJobs.contains(js));
                 proto.write(StateControllerProto.QuotaController.TrackedJob.HAS_QUOTA,
                         js.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
                 proto.write(StateControllerProto.QuotaController.TrackedJob.REMAINING_QUOTA_MS,
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 f1cd0cd6d30c..57ee6dcad9f2 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
@@ -33,6 +33,7 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
@@ -43,7 +44,12 @@ import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import android.app.ActivityManager;
+import android.app.ActivityManagerInternal;
 import android.app.AlarmManager;
+import android.app.AppGlobals;
+import android.app.IActivityManager;
+import android.app.IUidObserver;
 import android.app.job.JobInfo;
 import android.app.usage.UsageStatsManager;
 import android.app.usage.UsageStatsManagerInternal;
@@ -56,13 +62,16 @@ import android.os.BatteryManager;
 import android.os.BatteryManagerInternal;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.RemoteException;
 import android.os.SystemClock;
+import android.util.SparseBooleanArray;
 
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.server.LocalServices;
 import com.android.server.job.JobSchedulerService;
 import com.android.server.job.JobSchedulerService.Constants;
+import com.android.server.job.JobStore;
 import com.android.server.job.controllers.QuotaController.ExecutionStats;
 import com.android.server.job.controllers.QuotaController.TimingSession;
 
@@ -96,9 +105,13 @@ public class QuotaControllerTest {
     private BroadcastReceiver mChargingReceiver;
     private Constants mConstants;
     private QuotaController mQuotaController;
+    private int mSourceUid;
+    private IUidObserver mUidObserver;
 
     private MockitoSession mMockingSession;
     @Mock
+    private ActivityManagerInternal mActivityMangerInternal;
+    @Mock
     private AlarmManager mAlarmManager;
     @Mock
     private Context mContext;
@@ -107,6 +120,8 @@ public class QuotaControllerTest {
     @Mock
     private UsageStatsManagerInternal mUsageStatsManager;
 
+    private JobStore mJobStore;
+
     @Before
     public void setUp() {
         mMockingSession = mockitoSession()
@@ -123,8 +138,17 @@ public class QuotaControllerTest {
         when(mJobSchedulerService.getLock()).thenReturn(mJobSchedulerService);
         when(mJobSchedulerService.getConstants()).thenReturn(mConstants);
         // Called in QuotaController constructor.
+        IActivityManager activityManager = ActivityManager.getService();
+        spyOn(activityManager);
+        try {
+            doNothing().when(activityManager).registerUidObserver(any(), anyInt(), anyInt(), any());
+        } catch (RemoteException e) {
+            fail("registerUidObserver threw exception: " + e.getMessage());
+        }
         when(mContext.getMainLooper()).thenReturn(Looper.getMainLooper());
         when(mContext.getSystemService(Context.ALARM_SERVICE)).thenReturn(mAlarmManager);
+        doReturn(mActivityMangerInternal)
+                .when(() -> LocalServices.getService(ActivityManagerInternal.class));
         doReturn(mock(BatteryManagerInternal.class))
                 .when(() -> LocalServices.getService(BatteryManagerInternal.class));
         doReturn(mUsageStatsManager)
@@ -132,6 +156,9 @@ public class QuotaControllerTest {
         // Used in JobStatus.
         doReturn(mock(PackageManagerInternal.class))
                 .when(() -> LocalServices.getService(PackageManagerInternal.class));
+        // Used in QuotaController.Handler.
+        mJobStore = JobStore.initAndGetForTesting(mContext, mContext.getFilesDir());
+        when(mJobSchedulerService.getJobStore()).thenReturn(mJobStore);
 
         // Freeze the clocks at 24 hours after this moment in time. Several tests create sessions
         // in the past, and QuotaController sometimes floors values at 0, so if the test time
@@ -150,10 +177,23 @@ public class QuotaControllerTest {
         // Capture the listeners.
         ArgumentCaptor<BroadcastReceiver> receiverCaptor =
                 ArgumentCaptor.forClass(BroadcastReceiver.class);
+        ArgumentCaptor<IUidObserver> uidObserverCaptor =
+                ArgumentCaptor.forClass(IUidObserver.class);
         mQuotaController = new QuotaController(mJobSchedulerService);
 
         verify(mContext).registerReceiver(receiverCaptor.capture(), any());
         mChargingReceiver = receiverCaptor.getValue();
+        try {
+            verify(activityManager).registerUidObserver(
+                    uidObserverCaptor.capture(),
+                    eq(ActivityManager.UID_OBSERVER_PROCSTATE),
+                    eq(ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE),
+                    any());
+            mUidObserver = uidObserverCaptor.getValue();
+            mSourceUid = AppGlobals.getPackageManager().getPackageUid(SOURCE_PACKAGE, 0, 0);
+        } catch (RemoteException e) {
+            fail(e.getMessage());
+        }
     }
 
     @After
@@ -182,6 +222,25 @@ public class QuotaControllerTest {
         mChargingReceiver.onReceive(mContext, intent);
     }
 
+    private void setProcessState(int procState) {
+        try {
+            doReturn(procState).when(mActivityMangerInternal).getUidProcessState(mSourceUid);
+            SparseBooleanArray foregroundUids = mQuotaController.getForegroundUids();
+            spyOn(foregroundUids);
+            mUidObserver.onUidStateChanged(mSourceUid, procState, 0);
+            if (procState <= ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE) {
+                verify(foregroundUids, timeout(SECOND_IN_MILLIS).times(1))
+                        .put(eq(mSourceUid), eq(true));
+                assertTrue(foregroundUids.get(mSourceUid));
+            } else {
+                verify(foregroundUids, timeout(SECOND_IN_MILLIS).times(1)).delete(eq(mSourceUid));
+                assertFalse(foregroundUids.get(mSourceUid));
+            }
+        } catch (RemoteException e) {
+            fail("registerUidObserver threw exception: " + e.getMessage());
+        }
+    }
+
     private void setStandbyBucket(int bucketIndex) {
         int bucket;
         switch (bucketIndex) {
@@ -204,9 +263,18 @@ public class QuotaControllerTest {
                 anyLong())).thenReturn(bucket);
     }
 
-    private void setStandbyBucket(int bucketIndex, JobStatus job) {
+    private void setStandbyBucket(int bucketIndex, JobStatus... jobs) {
         setStandbyBucket(bucketIndex);
-        job.setStandbyBucket(bucketIndex);
+        for (JobStatus job : jobs) {
+            job.setStandbyBucket(bucketIndex);
+        }
+    }
+
+    private void trackJobs(JobStatus... jobs) {
+        for (JobStatus job : jobs) {
+            mJobStore.add(job);
+            mQuotaController.maybeStartTrackingJobLocked(job, null);
+        }
     }
 
     private JobStatus createJobStatus(String testTag, int jobId) {
@@ -214,8 +282,11 @@ public class QuotaControllerTest {
                 new ComponentName(mContext, "TestQuotaJobService"))
                 .setMinimumLatency(Math.abs(jobId) + 1)
                 .build();
-        return JobStatus.createFromJobInfo(
+        JobStatus js = JobStatus.createFromJobInfo(
                 jobInfo, CALLING_UID, SOURCE_PACKAGE, SOURCE_USER_ID, testTag);
+        // Make sure tests aren't passing just because the default bucket is likely ACTIVE.
+        js.setStandbyBucket(FREQUENT_INDEX);
+        return js;
     }
 
     private TimingSession createTimingSession(long start, long duration, int count) {
@@ -709,6 +780,7 @@ public class QuotaControllerTest {
         verify(mAlarmManager, never()).set(anyInt(), anyLong(), eq(TAG_QUOTA_CHECK), any(), any());
 
         JobStatus jobStatus = createJobStatus("testMaybeScheduleStartAlarmLocked_Active", 1);
+        setStandbyBucket(standbyBucket, jobStatus);
         mQuotaController.maybeStartTrackingJobLocked(jobStatus, null);
         mQuotaController.prepareForExecutionLocked(jobStatus);
         advanceElapsedClock(5 * MINUTE_IN_MILLIS);
@@ -1339,19 +1411,23 @@ public class QuotaControllerTest {
         setDischarging();
 
         JobStatus jobStatus = createJobStatus("testTimerTracking_AllForeground", 1);
-        jobStatus.uidActive = true;
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
         mQuotaController.maybeStartTrackingJobLocked(jobStatus, null);
 
         assertNull(mQuotaController.getTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
 
         mQuotaController.prepareForExecutionLocked(jobStatus);
         advanceElapsedClock(5 * SECOND_IN_MILLIS);
+        // Change to a state that should still be considered foreground.
+        setProcessState(ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE);
+        advanceElapsedClock(5 * SECOND_IN_MILLIS);
         mQuotaController.maybeStopTrackingJobLocked(jobStatus, null, false);
         assertNull(mQuotaController.getTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
     }
 
     /**
-     * Tests that Timers properly track overlapping foreground and background jobs.
+     * Tests that Timers properly track sessions when switching between foreground and background
+     * states.
      */
     @Test
     public void testTimerTracking_ForegroundAndBackground() {
@@ -1360,7 +1436,6 @@ public class QuotaControllerTest {
         JobStatus jobBg1 = createJobStatus("testTimerTracking_ForegroundAndBackground", 1);
         JobStatus jobBg2 = createJobStatus("testTimerTracking_ForegroundAndBackground", 2);
         JobStatus jobFg3 = createJobStatus("testTimerTracking_ForegroundAndBackground", 3);
-        jobFg3.uidActive = true;
         mQuotaController.maybeStartTrackingJobLocked(jobBg1, null);
         mQuotaController.maybeStartTrackingJobLocked(jobBg2, null);
         mQuotaController.maybeStartTrackingJobLocked(jobFg3, null);
@@ -1368,6 +1443,7 @@ public class QuotaControllerTest {
         List<TimingSession> expected = new ArrayList<>();
 
         // UID starts out inactive.
+        setProcessState(ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
         long start = JobSchedulerService.sElapsedRealtimeClock.millis();
         mQuotaController.prepareForExecutionLocked(jobBg1);
         advanceElapsedClock(10 * SECOND_IN_MILLIS);
@@ -1379,48 +1455,223 @@ public class QuotaControllerTest {
 
         // Bg job starts while inactive, spans an entire active session, and ends after the
         // active session.
-        // Fg job starts after the bg job and ends before the bg job.
-        // Entire bg job duration should be counted since it started before active session. However,
-        // count should only be 1 since Timer shouldn't count fg jobs.
+        // App switching to foreground state then fg job starts.
+        // App remains in foreground state after coming to foreground, so there should only be one
+        // session.
         start = JobSchedulerService.sElapsedRealtimeClock.millis();
         mQuotaController.maybeStartTrackingJobLocked(jobBg2, null);
         mQuotaController.prepareForExecutionLocked(jobBg2);
         advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        expected.add(createTimingSession(start, 10 * SECOND_IN_MILLIS, 1));
+        setProcessState(ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE);
         mQuotaController.prepareForExecutionLocked(jobFg3);
         advanceElapsedClock(10 * SECOND_IN_MILLIS);
         mQuotaController.maybeStopTrackingJobLocked(jobFg3, null, false);
         advanceElapsedClock(10 * SECOND_IN_MILLIS);
         mQuotaController.maybeStopTrackingJobLocked(jobBg2, null, false);
-        expected.add(createTimingSession(start, 30 * SECOND_IN_MILLIS, 1));
         assertEquals(expected, mQuotaController.getTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
 
         advanceElapsedClock(SECOND_IN_MILLIS);
 
         // Bg job 1 starts, then fg job starts. Bg job 1 job ends. Shortly after, uid goes
         // "inactive" and then bg job 2 starts. Then fg job ends.
-        // This should result in two TimingSessions with a count of one each.
+        // This should result in two TimingSessions:
+        //  * The first should have a count of 1
+        //  * The second should have a count of 2 since it will include both jobs
         start = JobSchedulerService.sElapsedRealtimeClock.millis();
         mQuotaController.maybeStartTrackingJobLocked(jobBg1, null);
         mQuotaController.maybeStartTrackingJobLocked(jobBg2, null);
         mQuotaController.maybeStartTrackingJobLocked(jobFg3, null);
+        setProcessState(ActivityManager.PROCESS_STATE_LAST_ACTIVITY);
         mQuotaController.prepareForExecutionLocked(jobBg1);
         advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        expected.add(createTimingSession(start, 10 * SECOND_IN_MILLIS, 1));
+        setProcessState(ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE);
         mQuotaController.prepareForExecutionLocked(jobFg3);
         advanceElapsedClock(10 * SECOND_IN_MILLIS);
         mQuotaController.maybeStopTrackingJobLocked(jobBg1, jobBg1, true);
-        expected.add(createTimingSession(start, 20 * SECOND_IN_MILLIS, 1));
         advanceElapsedClock(10 * SECOND_IN_MILLIS); // UID "inactive" now
         start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        setProcessState(ActivityManager.PROCESS_STATE_TOP_SLEEPING);
         mQuotaController.prepareForExecutionLocked(jobBg2);
         advanceElapsedClock(10 * SECOND_IN_MILLIS);
         mQuotaController.maybeStopTrackingJobLocked(jobFg3, null, false);
         advanceElapsedClock(10 * SECOND_IN_MILLIS);
         mQuotaController.maybeStopTrackingJobLocked(jobBg2, null, false);
-        expected.add(createTimingSession(start, 20 * SECOND_IN_MILLIS, 1));
+        expected.add(createTimingSession(start, 20 * SECOND_IN_MILLIS, 2));
         assertEquals(expected, mQuotaController.getTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
     }
 
     /**
+     * Tests that Timers properly track overlapping top and background jobs.
+     */
+    @Test
+    public void testTimerTracking_TopAndNonTop() {
+        setDischarging();
+
+        JobStatus jobBg1 = createJobStatus("testTimerTracking_TopAndNonTop", 1);
+        JobStatus jobBg2 = createJobStatus("testTimerTracking_TopAndNonTop", 2);
+        JobStatus jobFg1 = createJobStatus("testTimerTracking_TopAndNonTop", 3);
+        JobStatus jobTop = createJobStatus("testTimerTracking_TopAndNonTop", 4);
+        mQuotaController.maybeStartTrackingJobLocked(jobBg1, null);
+        mQuotaController.maybeStartTrackingJobLocked(jobBg2, null);
+        mQuotaController.maybeStartTrackingJobLocked(jobFg1, null);
+        mQuotaController.maybeStartTrackingJobLocked(jobTop, null);
+        assertNull(mQuotaController.getTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
+        List<TimingSession> expected = new ArrayList<>();
+
+        // UID starts out inactive.
+        setProcessState(ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
+        long start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        mQuotaController.prepareForExecutionLocked(jobBg1);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        mQuotaController.maybeStopTrackingJobLocked(jobBg1, jobBg1, true);
+        expected.add(createTimingSession(start, 10 * SECOND_IN_MILLIS, 1));
+        assertEquals(expected, mQuotaController.getTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
+
+        advanceElapsedClock(SECOND_IN_MILLIS);
+
+        // Bg job starts while inactive, spans an entire active session, and ends after the
+        // active session.
+        // App switching to top state then fg job starts.
+        // App remains in top state after coming to top, so there should only be one
+        // session.
+        start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        mQuotaController.maybeStartTrackingJobLocked(jobBg2, null);
+        mQuotaController.prepareForExecutionLocked(jobBg2);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        expected.add(createTimingSession(start, 10 * SECOND_IN_MILLIS, 1));
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        mQuotaController.prepareForExecutionLocked(jobTop);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        mQuotaController.maybeStopTrackingJobLocked(jobTop, null, false);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        mQuotaController.maybeStopTrackingJobLocked(jobBg2, null, false);
+        assertEquals(expected, mQuotaController.getTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
+
+        advanceElapsedClock(SECOND_IN_MILLIS);
+
+        // Bg job 1 starts, then top job starts. Bg job 1 job ends. Then app goes to
+        // foreground_service and a new job starts. Shortly after, uid goes
+        // "inactive" and then bg job 2 starts. Then top job ends, followed by bg and fg jobs.
+        // This should result in two TimingSessions:
+        //  * The first should have a count of 1
+        //  * The second should have a count of 2, which accounts for the bg2 and fg, but not top
+        //    jobs.
+        start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        mQuotaController.maybeStartTrackingJobLocked(jobBg1, null);
+        mQuotaController.maybeStartTrackingJobLocked(jobBg2, null);
+        mQuotaController.maybeStartTrackingJobLocked(jobTop, null);
+        setProcessState(ActivityManager.PROCESS_STATE_LAST_ACTIVITY);
+        mQuotaController.prepareForExecutionLocked(jobBg1);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        expected.add(createTimingSession(start, 10 * SECOND_IN_MILLIS, 1));
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        mQuotaController.prepareForExecutionLocked(jobTop);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        mQuotaController.maybeStopTrackingJobLocked(jobBg1, jobBg1, true);
+        advanceElapsedClock(5 * SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE);
+        mQuotaController.prepareForExecutionLocked(jobFg1);
+        advanceElapsedClock(5 * SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS); // UID "inactive" now
+        start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        setProcessState(ActivityManager.PROCESS_STATE_TOP_SLEEPING);
+        mQuotaController.prepareForExecutionLocked(jobBg2);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        mQuotaController.maybeStopTrackingJobLocked(jobTop, null, false);
+        advanceElapsedClock(10 * SECOND_IN_MILLIS);
+        mQuotaController.maybeStopTrackingJobLocked(jobBg2, null, false);
+        mQuotaController.maybeStopTrackingJobLocked(jobFg1, null, false);
+        expected.add(createTimingSession(start, 20 * SECOND_IN_MILLIS, 2));
+        assertEquals(expected, mQuotaController.getTimingSessions(SOURCE_USER_ID, SOURCE_PACKAGE));
+    }
+
+    /**
+     * Tests that TOP jobs aren't stopped when an app runs out of quota.
+     */
+    @Test
+    public void testTracking_OutOfQuota_ForegroundAndBackground() {
+        setDischarging();
+
+        JobStatus jobBg = createJobStatus("testTracking_OutOfQuota_ForegroundAndBackground", 1);
+        JobStatus jobTop = createJobStatus("testTracking_OutOfQuota_ForegroundAndBackground", 2);
+        trackJobs(jobBg, jobTop);
+        setStandbyBucket(WORKING_INDEX, jobTop, jobBg); // 2 hour window
+        // Now the package only has 20 seconds to run.
+        final long remainingTimeMs = 20 * SECOND_IN_MILLIS;
+        mQuotaController.saveTimingSession(SOURCE_USER_ID, SOURCE_PACKAGE,
+                createTimingSession(
+                        JobSchedulerService.sElapsedRealtimeClock.millis() - HOUR_IN_MILLIS,
+                        10 * MINUTE_IN_MILLIS - remainingTimeMs, 1));
+
+        InOrder inOrder = inOrder(mJobSchedulerService);
+
+        // UID starts out inactive.
+        setProcessState(ActivityManager.PROCESS_STATE_SERVICE);
+        // Start the job.
+        mQuotaController.prepareForExecutionLocked(jobBg);
+        advanceElapsedClock(remainingTimeMs / 2);
+        // New job starts after UID is in the foreground. Since the app is now in the foreground, it
+        // should continue to have remainingTimeMs / 2 time remaining.
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        mQuotaController.prepareForExecutionLocked(jobTop);
+        advanceElapsedClock(remainingTimeMs);
+
+        // Wait for some extra time to allow for job processing.
+        inOrder.verify(mJobSchedulerService,
+                timeout(remainingTimeMs + 2 * SECOND_IN_MILLIS).times(0))
+                .onControllerStateChanged();
+        assertEquals(remainingTimeMs / 2, mQuotaController.getRemainingExecutionTimeLocked(jobBg));
+        assertEquals(remainingTimeMs / 2, mQuotaController.getRemainingExecutionTimeLocked(jobTop));
+        // Go to a background state.
+        setProcessState(ActivityManager.PROCESS_STATE_TOP_SLEEPING);
+        advanceElapsedClock(remainingTimeMs / 2 + 1);
+        inOrder.verify(mJobSchedulerService,
+                timeout(remainingTimeMs / 2 + 2 * SECOND_IN_MILLIS).times(1))
+                .onControllerStateChanged();
+        // Top job should still be allowed to run.
+        assertFalse(jobBg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertTrue(jobTop.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+
+        // New jobs to run.
+        JobStatus jobBg2 = createJobStatus("testTracking_OutOfQuota_ForegroundAndBackground", 3);
+        JobStatus jobTop2 = createJobStatus("testTracking_OutOfQuota_ForegroundAndBackground", 4);
+        JobStatus jobFg = createJobStatus("testTracking_OutOfQuota_ForegroundAndBackground", 5);
+        setStandbyBucket(WORKING_INDEX, jobBg2, jobTop2, jobFg);
+
+        advanceElapsedClock(20 * SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        inOrder.verify(mJobSchedulerService, timeout(SECOND_IN_MILLIS).times(1))
+                .onControllerStateChanged();
+        trackJobs(jobFg, jobTop);
+        mQuotaController.prepareForExecutionLocked(jobTop);
+        assertTrue(jobTop.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertTrue(jobFg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertTrue(jobBg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+
+        // App still in foreground so everything should be in quota.
+        advanceElapsedClock(20 * SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE);
+        assertTrue(jobTop.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertTrue(jobFg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertTrue(jobBg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+
+        advanceElapsedClock(20 * SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_SERVICE);
+        inOrder.verify(mJobSchedulerService, timeout(SECOND_IN_MILLIS).times(1))
+                .onControllerStateChanged();
+        // App is now in background and out of quota. Fg should now change to out of quota since it
+        // wasn't started. Top should remain in quota since it started when the app was in TOP.
+        assertTrue(jobTop.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertFalse(jobFg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertFalse(jobBg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        trackJobs(jobBg2);
+        assertFalse(jobBg2.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+    }
+
+    /**
      * Tests that a job is properly updated and JobSchedulerService is notified when a job reaches
      * its quota.
      */