Scheduling start alarm when job starts off out of quota.

Up until now, QuotaController was only scheduling the quota check alarm
when an app ran out of quota while a job was running, the UID proc state
or standby bucket changed, the device was unplugged, the parole state
changed, or quota controller constants changed. However, if a job was
scheduled and already out of quota (which could be the case due to
job count throttling), an alarm wasn't scheduled. This meant that alarms
throttled due to high job counts probably wouldn't run until the device
was plugged in or the app changed its standby bucket or proc state. Now,
we schedule an alarm if a newly scheduled job is already out of quota so
that it will come back into quota at the proper time.

Bug: 131267600
Test: atest com.android.server.job.controllers.QuotaControllerTest
Test: atest CtsJobSchedulerTestCases
Change-Id: I802b0aa076690451d901521327c4ddab111c42f6
diff --git a/core/proto/android/server/jobscheduler.proto b/core/proto/android/server/jobscheduler.proto
index dc2e6d5..0df2c83 100644
--- a/core/proto/android/server/jobscheduler.proto
+++ b/core/proto/android/server/jobscheduler.proto
@@ -459,6 +459,7 @@
 
         optional bool is_charging = 1;
         optional bool is_in_parole = 2;
+        optional int64 elapsed_realtime = 6;
 
         // List of UIDs currently in the foreground.
         repeated int32 foreground_uids = 3;
@@ -478,6 +479,16 @@
         }
         repeated TrackedJob tracked_jobs = 4;
 
+        message AlarmListener {
+            option (.android.msg_privacy).dest = DEST_AUTOMATIC;
+
+            // Whether the listener is waiting for an alarm or not.
+            optional bool is_waiting = 1;
+            // The time at which the alarm should go off, in the elapsed realtime timebase. Only
+            // valid if is_waiting is true.
+            optional int64 trigger_time_elapsed = 2;
+        }
+
         message ExecutionStats {
             option (.android.msg_privacy).dest = DEST_AUTOMATIC;
 
@@ -567,6 +578,8 @@
             repeated TimingSession saved_sessions = 3;
 
             repeated ExecutionStats execution_stats = 4;
+
+            optional AlarmListener in_quota_alarm_listener = 5;
         }
         repeated PackageStats package_stats = 5;
     }
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 ccd1db4..11f0939 100644
--- a/services/core/java/com/android/server/job/controllers/QuotaController.java
+++ b/services/core/java/com/android/server/job/controllers/QuotaController.java
@@ -511,17 +511,28 @@
 
     @Override
     public void maybeStartTrackingJobLocked(JobStatus jobStatus, JobStatus lastJob) {
+        final int userId = jobStatus.getSourceUserId();
+        final String pkgName = jobStatus.getSourcePackageName();
         // Still need to track jobs even if mShouldThrottle is false in case it's set to true at
         // some point.
-        ArraySet<JobStatus> jobs = mTrackedJobs.get(jobStatus.getSourceUserId(),
-                jobStatus.getSourcePackageName());
+        ArraySet<JobStatus> jobs = mTrackedJobs.get(userId, pkgName);
         if (jobs == null) {
             jobs = new ArraySet<>();
-            mTrackedJobs.add(jobStatus.getSourceUserId(), jobStatus.getSourcePackageName(), jobs);
+            mTrackedJobs.add(userId, pkgName, jobs);
         }
         jobs.add(jobStatus);
         jobStatus.setTrackingController(JobStatus.TRACKING_QUOTA);
-        jobStatus.setQuotaConstraintSatisfied(!mShouldThrottle || isWithinQuotaLocked(jobStatus));
+        if (mShouldThrottle) {
+            final boolean isWithinQuota = isWithinQuotaLocked(jobStatus);
+            jobStatus.setQuotaConstraintSatisfied(isWithinQuota);
+            if (!isWithinQuota) {
+                maybeScheduleStartAlarmLocked(userId, pkgName,
+                        getEffectiveStandbyBucket(jobStatus));
+            }
+        } else {
+            // QuotaController isn't throttling, so always set to true.
+            jobStatus.setQuotaConstraintSatisfied(true);
+        }
     }
 
     @Override
@@ -1628,6 +1639,9 @@
             if (isActive()) {
                 pw.print("started at ");
                 pw.print(mStartTimeElapsed);
+                pw.print(" (");
+                pw.print(sElapsedRealtimeClock.millis() - mStartTimeElapsed);
+                pw.print("ms ago)");
             } else {
                 pw.print("NOT active");
             }
@@ -1937,6 +1951,7 @@
         pw.println("Is throttling: " + mShouldThrottle);
         pw.println("Is charging: " + mChargeTracker.isCharging());
         pw.println("In parole: " + mInParole);
+        pw.println("Current elapsed time: " + sElapsedRealtimeClock.millis());
         pw.println();
 
         pw.print("Foreground UIDs: ");
@@ -2030,6 +2045,26 @@
             }
         }
         pw.decreaseIndent();
+
+        pw.println();
+        pw.println("In quota alarms:");
+        pw.increaseIndent();
+        for (int u = 0; u < mInQuotaAlarmListeners.numUsers(); ++u) {
+            final int userId = mInQuotaAlarmListeners.keyAt(u);
+            for (int p = 0; p < mInQuotaAlarmListeners.numPackagesForUser(userId); ++p) {
+                final String pkgName = mInQuotaAlarmListeners.keyAt(u, p);
+                QcAlarmListener alarmListener = mInQuotaAlarmListeners.valueAt(u, p);
+
+                pw.print(string(userId, pkgName));
+                pw.print(": ");
+                if (alarmListener.isWaiting()) {
+                    pw.println(alarmListener.getTriggerTimeElapsed());
+                } else {
+                    pw.println("NOT WAITING");
+                }
+            }
+        }
+        pw.decreaseIndent();
     }
 
     @Override
@@ -2040,6 +2075,8 @@
 
         proto.write(StateControllerProto.QuotaController.IS_CHARGING, mChargeTracker.isCharging());
         proto.write(StateControllerProto.QuotaController.IS_IN_PAROLE, mInParole);
+        proto.write(StateControllerProto.QuotaController.ELAPSED_REALTIME,
+                sElapsedRealtimeClock.millis());
 
         for (int i = 0; i < mForegroundUids.size(); ++i) {
             proto.write(StateControllerProto.QuotaController.FOREGROUND_UIDS,
@@ -2132,6 +2169,18 @@
                     }
                 }
 
+                QcAlarmListener alarmListener = mInQuotaAlarmListeners.get(userId, pkgName);
+                if (alarmListener != null) {
+                    final long alToken = proto.start(
+                            StateControllerProto.QuotaController.PackageStats.IN_QUOTA_ALARM_LISTENER);
+                    proto.write(StateControllerProto.QuotaController.AlarmListener.IS_WAITING,
+                            alarmListener.isWaiting());
+                    proto.write(
+                            StateControllerProto.QuotaController.AlarmListener.TRIGGER_TIME_ELAPSED,
+                            alarmListener.getTriggerTimeElapsed());
+                    proto.end(alToken);
+                }
+
                 proto.end(psToken);
             }
         }
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 f492d13..7c30f25 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
@@ -2197,4 +2197,51 @@
         assertFalse(jobStatus.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
         verify(handler, never()).sendMessageDelayed(any(), anyInt());
     }
+
+    /**
+     * Tests that the start alarm is properly scheduled when a job has been throttled due to the job
+     * count quota.
+     */
+    @Test
+    public void testStartAlarmScheduled_JobCount_AllowedTime() {
+        // saveTimingSession calls maybeScheduleCleanupAlarmLocked which interferes with these tests
+        // because it schedules an alarm too. Prevent it from doing so.
+        spyOn(mQuotaController);
+        doNothing().when(mQuotaController).maybeScheduleCleanupAlarmLocked();
+
+        final long start = JobSchedulerService.sElapsedRealtimeClock.millis();
+        final int standbyBucket = WORKING_INDEX;
+        setProcessState(ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
+
+        // No sessions saved yet.
+        mQuotaController.maybeScheduleStartAlarmLocked(SOURCE_USER_ID, SOURCE_PACKAGE,
+                standbyBucket);
+        verify(mAlarmManager, never()).set(anyInt(), anyLong(), eq(TAG_QUOTA_CHECK), any(), any());
+
+        // Ran jobs up to the job limit. All of them should be allowed to run.
+        for (int i = 0; i < mConstants.QUOTA_CONTROLLER_MAX_JOB_COUNT_PER_ALLOWED_TIME; ++i) {
+            JobStatus job = createJobStatus("testStartAlarmScheduled_JobCount_AllowedTime", i);
+            mQuotaController.maybeStartTrackingJobLocked(job, null);
+            assertTrue(job.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+            mQuotaController.prepareForExecutionLocked(job);
+            advanceElapsedClock(SECOND_IN_MILLIS);
+            mQuotaController.maybeStopTrackingJobLocked(job, null, false);
+            advanceElapsedClock(SECOND_IN_MILLIS);
+        }
+        // Start alarm shouldn't have been scheduled since the app was in quota up until this point.
+        verify(mAlarmManager, never()).set(anyInt(), anyLong(), eq(TAG_QUOTA_CHECK), any(), any());
+
+        // The app is now out of job count quota
+        JobStatus throttledJob = createJobStatus(
+                "testStartAlarmScheduled_JobCount_AllowedTime", 42);
+        mQuotaController.maybeStartTrackingJobLocked(throttledJob, null);
+        assertFalse(throttledJob.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+
+        ExecutionStats stats = mQuotaController.getExecutionStatsLocked(SOURCE_USER_ID,
+                SOURCE_PACKAGE, standbyBucket);
+        final long expectedWorkingAlarmTime =
+                stats.jobCountExpirationTimeElapsed + mConstants.QUOTA_CONTROLLER_ALLOWED_TIME_PER_PERIOD_MS;
+        verify(mAlarmManager, times(1))
+                .set(anyInt(), eq(expectedWorkingAlarmTime), eq(TAG_QUOTA_CHECK), any(), any());
+    }
 }