Limit scheduled jobs to 100 per app

Packages that are entitled to schedule jobs on behalf of other uids
are not subject to the limit.  Also break the JobStore's monolithic
set of jobs into per-uid slices for efficiency and orthogonality.

Bug 27150350

Change-Id: I8f5f718bf200d55f9977a6fc53b7f617e7652ad9
diff --git a/services/core/java/com/android/server/job/JobSchedulerService.java b/services/core/java/com/android/server/job/JobSchedulerService.java
index 811c947..42ecc06 100644
--- a/services/core/java/com/android/server/job/JobSchedulerService.java
+++ b/services/core/java/com/android/server/job/JobSchedulerService.java
@@ -49,14 +49,13 @@
 import android.os.ServiceManager;
 import android.os.SystemClock;
 import android.os.UserHandle;
-import android.os.Process;
-import android.util.ArraySet;
 import android.util.Slog;
 import android.util.SparseArray;
 
 import com.android.internal.app.IBatteryStats;
 import com.android.server.DeviceIdleController;
 import com.android.server.LocalServices;
+import com.android.server.job.JobStore.JobStatusFunctor;
 import com.android.server.job.controllers.AppIdleController;
 import com.android.server.job.controllers.BatteryController;
 import com.android.server.job.controllers.ConnectivityController;
@@ -80,11 +79,15 @@
  */
 public final class JobSchedulerService extends com.android.server.SystemService
         implements StateChangedListener, JobCompletedListener {
+    static final String TAG = "JobSchedulerService";
     public static final boolean DEBUG = false;
+
     /** The number of concurrent jobs we run at one time. */
     private static final int MAX_JOB_CONTEXTS_COUNT
             = ActivityManager.isLowRamDeviceStatic() ? 3 : 6;
-    static final String TAG = "JobSchedulerService";
+    /** The maximum number of jobs that we allow an unprivileged app to schedule */
+    private static final int MAX_JOBS_PER_APP = 100;
+
     /** Global local for all job scheduler state. */
     final Object mLock = new Object();
     /** Master list of jobs. */
@@ -250,6 +253,15 @@
         if (DEBUG) Slog.d(TAG, "SCHEDULE: " + jobStatus.toShortString());
         JobStatus toCancel;
         synchronized (mLock) {
+            // Jobs on behalf of others don't apply to the per-app job cap
+            if (packageName == null) {
+                if (mJobs.countJobsForUid(uId) > MAX_JOBS_PER_APP) {
+                    Slog.w(TAG, "Too many jobs for uid " + uId);
+                    throw new IllegalStateException("Apps may not schedule more than "
+                                + MAX_JOBS_PER_APP + " distinct jobs");
+                }
+            }
+
             toCancel = mJobs.getJobByUidAndJobId(uId, job.getId());
         }
         startTrackingJob(jobStatus, toCancel);
@@ -261,17 +273,15 @@
     }
 
     public List<JobInfo> getPendingJobs(int uid) {
-        ArrayList<JobInfo> outList = new ArrayList<JobInfo>();
         synchronized (mLock) {
-            ArraySet<JobStatus> jobs = mJobs.getJobs();
-            for (int i=0; i<jobs.size(); i++) {
-                JobStatus job = jobs.valueAt(i);
-                if (job.getUid() == uid) {
-                    outList.add(job.getJob());
-                }
+            List<JobStatus> jobs = mJobs.getJobsByUid(uid);
+            ArrayList<JobInfo> outList = new ArrayList<JobInfo>(jobs.size());
+            for (int i = jobs.size() - 1; i >= 0; i--) {
+                JobStatus job = jobs.get(i);
+                outList.add(job.getJob());
             }
+            return outList;
         }
-        return outList;
     }
 
     void cancelJobsForUser(int userHandle) {
@@ -471,14 +481,16 @@
                                     getContext().getMainLooper()));
                 }
                 // Attach jobs to their controllers.
-                ArraySet<JobStatus> jobs = mJobs.getJobs();
-                for (int i=0; i<jobs.size(); i++) {
-                    JobStatus job = jobs.valueAt(i);
-                    for (int controller=0; controller<mControllers.size(); controller++) {
-                        mControllers.get(controller).deviceIdleModeChanged(mDeviceIdleMode);
-                        mControllers.get(controller).maybeStartTrackingJobLocked(job, null);
+                mJobs.forEachJob(new JobStatusFunctor() {
+                    @Override
+                    public void process(JobStatus job) {
+                        for (int controller = 0; controller < mControllers.size(); controller++) {
+                            final StateController sc = mControllers.get(controller);
+                            sc.deviceIdleModeChanged(mDeviceIdleMode);
+                            sc.maybeStartTrackingJobLocked(job, null);
+                        }
                     }
-                }
+                });
                 // GO GO GO!
                 mHandler.obtainMessage(MSG_CHECK_JOB).sendToTarget();
             }
@@ -741,23 +753,13 @@
          * as many as we can.
          */
         private void queueReadyJobsForExecutionLockedH() {
-            ArraySet<JobStatus> jobs = mJobs.getJobs();
-            mPendingJobs.clear();
             if (DEBUG) {
                 Slog.d(TAG, "queuing all ready jobs for execution:");
             }
-            for (int i=0; i<jobs.size(); i++) {
-                JobStatus job = jobs.valueAt(i);
-                if (isReadyToBeExecutedLocked(job)) {
-                    if (DEBUG) {
-                        Slog.d(TAG, "    queued " + job.toShortString());
-                    }
-                    mPendingJobs.add(job);
-                } else if (areJobConstraintsNotSatisfiedLocked(job)) {
-                    stopJobOnServiceContextLocked(job,
-                            JobParameters.REASON_CONSTRAINTS_NOT_SATISFIED);
-                }
-            }
+            mPendingJobs.clear();
+            mJobs.forEachJob(mReadyQueueFunctor);
+            mReadyQueueFunctor.postProcess();
+
             if (DEBUG) {
                 final int queuedJobs = mPendingJobs.size();
                 if (queuedJobs == 0) {
@@ -768,6 +770,34 @@
             }
         }
 
+        class ReadyJobQueueFunctor implements JobStatusFunctor {
+            ArrayList<JobStatus> newReadyJobs;
+
+            @Override
+            public void process(JobStatus job) {
+                if (isReadyToBeExecutedLocked(job)) {
+                    if (DEBUG) {
+                        Slog.d(TAG, "    queued " + job.toShortString());
+                    }
+                    if (newReadyJobs == null) {
+                        newReadyJobs = new ArrayList<JobStatus>();
+                    }
+                    newReadyJobs.add(job);
+                } else if (areJobConstraintsNotSatisfiedLocked(job)) {
+                    stopJobOnServiceContextLocked(job,
+                            JobParameters.REASON_CONSTRAINTS_NOT_SATISFIED);
+                }
+            }
+
+            public void postProcess() {
+                if (newReadyJobs != null) {
+                    mPendingJobs.addAll(newReadyJobs);
+                }
+                newReadyJobs = null;
+            }
+        }
+        private final ReadyJobQueueFunctor mReadyQueueFunctor = new ReadyJobQueueFunctor();
+
         /**
          * The state of at least one job has changed. Here is where we could enforce various
          * policies on when we want to execute jobs.
@@ -777,18 +807,21 @@
          * If more than 4 jobs total are ready we send them all off.
          * TODO: It would be nice to consolidate these sort of high-level policies somewhere.
          */
-        private void maybeQueueReadyJobsForExecutionLockedH() {
-            mPendingJobs.clear();
-            int chargingCount = 0;
-            int idleCount =  0;
-            int backoffCount = 0;
-            int connectivityCount = 0;
-            int contentCount = 0;
-            List<JobStatus> runnableJobs = null;
-            ArraySet<JobStatus> jobs = mJobs.getJobs();
-            if (DEBUG) Slog.d(TAG, "Maybe queuing ready jobs...");
-            for (int i=0; i<jobs.size(); i++) {
-                JobStatus job = jobs.valueAt(i);
+        class MaybeReadyJobQueueFunctor implements JobStatusFunctor {
+            int chargingCount;
+            int idleCount;
+            int backoffCount;
+            int connectivityCount;
+            int contentCount;
+            List<JobStatus> runnableJobs;
+
+            public MaybeReadyJobQueueFunctor() {
+                reset();
+            }
+
+            // Functor method invoked for each job via JobStore.forEachJob()
+            @Override
+            public void process(JobStatus job) {
                 if (isReadyToBeExecutedLocked(job)) {
                     try {
                         if (ActivityManagerNative.getDefault().getAppStartMode(job.getUid(),
@@ -797,7 +830,7 @@
                             Slog.w(TAG, "Aborting job " + job.getUid() + ":"
                                     + job.getJob().toString() + " -- package not allowed to start");
                             mHandler.obtainMessage(MSG_STOP_JOB, job).sendToTarget();
-                            continue;
+                            return;
                         }
                     } catch (RemoteException e) {
                     }
@@ -825,23 +858,45 @@
                             JobParameters.REASON_CONSTRAINTS_NOT_SATISFIED);
                 }
             }
-            if (backoffCount > 0 ||
-                    idleCount >= MIN_IDLE_COUNT ||
-                    connectivityCount >= MIN_CONNECTIVITY_COUNT ||
-                    chargingCount >= MIN_CHARGING_COUNT ||
-                    contentCount  >= MIN_CONTENT_COUNT ||
-                    (runnableJobs != null && runnableJobs.size() >= MIN_READY_JOBS_COUNT)) {
-                if (DEBUG) {
-                    Slog.d(TAG, "maybeQueueReadyJobsForExecutionLockedH: Running jobs.");
+
+            public void postProcess() {
+                if (backoffCount > 0 ||
+                        idleCount >= MIN_IDLE_COUNT ||
+                        connectivityCount >= MIN_CONNECTIVITY_COUNT ||
+                        chargingCount >= MIN_CHARGING_COUNT ||
+                        contentCount  >= MIN_CONTENT_COUNT ||
+                        (runnableJobs != null && runnableJobs.size() >= MIN_READY_JOBS_COUNT)) {
+                    if (DEBUG) {
+                        Slog.d(TAG, "maybeQueueReadyJobsForExecutionLockedH: Running jobs.");
+                    }
+                    mPendingJobs.addAll(runnableJobs);
+                } else {
+                    if (DEBUG) {
+                        Slog.d(TAG, "maybeQueueReadyJobsForExecutionLockedH: Not running anything.");
+                    }
                 }
-                for (int i=0; i<runnableJobs.size(); i++) {
-                    mPendingJobs.add(runnableJobs.get(i));
-                }
-            } else {
-                if (DEBUG) {
-                    Slog.d(TAG, "maybeQueueReadyJobsForExecutionLockedH: Not running anything.");
-                }
+
+                // Be ready for next time
+                reset();
             }
+
+            private void reset() {
+                chargingCount = 0;
+                idleCount =  0;
+                backoffCount = 0;
+                connectivityCount = 0;
+                contentCount = 0;
+                runnableJobs = null;
+            }
+        }
+        private final MaybeReadyJobQueueFunctor mMaybeQueueFunctor = new MaybeReadyJobQueueFunctor();
+
+        private void maybeQueueReadyJobsForExecutionLockedH() {
+            if (DEBUG) Slog.d(TAG, "Maybe queuing ready jobs...");
+
+            mPendingJobs.clear();
+            mJobs.forEachJob(mMaybeQueueFunctor);
+            mMaybeQueueFunctor.postProcess();
         }
 
         /**
@@ -1090,17 +1145,27 @@
         @Override
         public int scheduleAsPackage(JobInfo job, String packageName, int userId)
                 throws RemoteException {
+            final int callerUid = Binder.getCallingUid();
             if (DEBUG) {
-                Slog.d(TAG, "Scheduling job: " + job.toString() + " on behalf of " + packageName);
+                Slog.d(TAG, "Caller uid " + callerUid + " scheduling job: " + job.toString()
+                        + " on behalf of " + packageName);
             }
-            final int uid = Binder.getCallingUid();
-            if (uid != Process.SYSTEM_UID) {
-                throw new IllegalArgumentException("Only system process is allowed"
-                        + "to set packageName");
+
+            if (packageName == null) {
+                throw new NullPointerException("Must specify a package for scheduleAsPackage()");
             }
+
+            int mayScheduleForOthers = getContext().checkCallingOrSelfPermission(
+                    android.Manifest.permission.UPDATE_DEVICE_STATS);
+            if (mayScheduleForOthers != PackageManager.PERMISSION_GRANTED) {
+                throw new SecurityException("Caller uid " + callerUid
+                        + " not permitted to schedule jobs for other apps");
+            }
+
             long ident = Binder.clearCallingIdentity();
             try {
-                return JobSchedulerService.this.scheduleAsPackage(job, uid, packageName, userId);
+                return JobSchedulerService.this.scheduleAsPackage(job, callerUid,
+                        packageName, userId);
             } finally {
                 Binder.restoreCallingIdentity(ident);
             }
@@ -1183,7 +1248,7 @@
         return s.toString();
     }
 
-    void dumpInternal(PrintWriter pw) {
+    void dumpInternal(final PrintWriter pw) {
         final long now = SystemClock.elapsedRealtime();
         synchronized (mLock) {
             pw.print("Started users: ");
@@ -1193,24 +1258,27 @@
             pw.println();
             pw.println("Registered jobs:");
             if (mJobs.size() > 0) {
-                ArraySet<JobStatus> jobs = mJobs.getJobs();
-                for (int i=0; i<jobs.size(); i++) {
-                    JobStatus job = jobs.valueAt(i);
-                    pw.print("  Job #"); pw.print(i); pw.print(": ");
-                    pw.println(job.toShortString());
-                    job.dump(pw, "    ");
-                    pw.print("    Ready: ");
-                    pw.print(mHandler.isReadyToBeExecutedLocked(job));
-                    pw.print(" (job=");
-                    pw.print(job.isReady());
-                    pw.print(" pending=");
-                    pw.print(mPendingJobs.contains(job));
-                    pw.print(" active=");
-                    pw.print(isCurrentlyActiveLocked(job));
-                    pw.print(" user=");
-                    pw.print(mStartedUsers.contains(job.getUserId()));
-                    pw.println(")");
-                }
+                mJobs.forEachJob(new JobStatusFunctor() {
+                    private int index = 0;
+
+                    @Override
+                    public void process(JobStatus job) {
+                        pw.print("  Job #"); pw.print(index++); pw.print(": ");
+                        pw.println(job.toShortString());
+                        job.dump(pw, "    ");
+                        pw.print("    Ready: ");
+                        pw.print(mHandler.isReadyToBeExecutedLocked(job));
+                        pw.print(" (job=");
+                        pw.print(job.isReady());
+                        pw.print(" pending=");
+                        pw.print(mPendingJobs.contains(job));
+                        pw.print(" active=");
+                        pw.print(isCurrentlyActiveLocked(job));
+                        pw.print(" user=");
+                        pw.print(mStartedUsers.contains(job.getUserId()));
+                        pw.println(")");
+                    }
+                });
             } else {
                 pw.println("  None.");
             }
diff --git a/services/core/java/com/android/server/job/JobStore.java b/services/core/java/com/android/server/job/JobStore.java
index 6020247..f8753d6 100644
--- a/services/core/java/com/android/server/job/JobStore.java
+++ b/services/core/java/com/android/server/job/JobStore.java
@@ -29,6 +29,7 @@
 import android.util.ArraySet;
 import android.util.Pair;
 import android.util.Slog;
+import android.util.SparseArray;
 import android.util.Xml;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -68,8 +69,8 @@
 
     /** Threshold to adjust how often we want to write to the db. */
     private static final int MAX_OPS_BEFORE_WRITE = 1;
-    final ArraySet<JobStatus> mJobSet;
     final Object mLock;
+    final JobSet mJobSet; // per-caller-uid tracking
     final Context mContext;
 
     private int mDirtyOperations;
@@ -114,7 +115,7 @@
         jobDir.mkdirs();
         mJobsFile = new AtomicFile(new File(jobDir, "jobs.xml"));
 
-        mJobSet = new ArraySet<JobStatus>();
+        mJobSet = new JobSet();
 
         readJobMapFromDisk(mJobSet);
     }
@@ -137,19 +138,6 @@
         return replaced;
     }
 
-    /**
-     * Whether this jobStatus object already exists in the JobStore.
-     */
-    public boolean containsJobIdForUid(int jobId, int uId) {
-        for (int i=mJobSet.size()-1; i>=0; i--) {
-            JobStatus ts = mJobSet.valueAt(i);
-            if (ts.getUid() == uId && ts.getJobId() == jobId) {
-                return true;
-            }
-        }
-        return false;
-    }
-
     boolean containsJob(JobStatus jobStatus) {
         return mJobSet.contains(jobStatus);
     }
@@ -158,6 +146,10 @@
         return mJobSet.size();
     }
 
+    public int countJobsForUid(int uid) {
+        return mJobSet.countJobsForUid(uid);
+    }
+
     /**
      * Remove the provided job. Will also delete the job if it was persisted.
      * @param writeBack If true, the job will be deleted (if it was persisted) immediately.
@@ -188,14 +180,7 @@
      * @return A list of all the jobs scheduled by the provided user. Never null.
      */
     public List<JobStatus> getJobsByUser(int userHandle) {
-        List<JobStatus> matchingJobs = new ArrayList<JobStatus>();
-        for (int i=mJobSet.size()-1; i>=0; i--) {
-            JobStatus ts = mJobSet.valueAt(i);
-            if (UserHandle.getUserId(ts.getUid()) == userHandle) {
-                matchingJobs.add(ts);
-            }
-        }
-        return matchingJobs;
+        return mJobSet.getJobsByUser(userHandle);
     }
 
     /**
@@ -203,14 +188,7 @@
      * @return All JobStatus objects for a given uid from the master list. Never null.
      */
     public List<JobStatus> getJobsByUid(int uid) {
-        List<JobStatus> matchingJobs = new ArrayList<JobStatus>();
-        for (int i=mJobSet.size()-1; i>=0; i--) {
-            JobStatus ts = mJobSet.valueAt(i);
-            if (ts.getUid() == uid) {
-                matchingJobs.add(ts);
-            }
-        }
-        return matchingJobs;
+        return mJobSet.getJobsByUid(uid);
     }
 
     /**
@@ -219,20 +197,21 @@
      * @return the JobStatus that matches the provided uId and jobId, or null if none found.
      */
     public JobStatus getJobByUidAndJobId(int uid, int jobId) {
-        for (int i=mJobSet.size()-1; i>=0; i--) {
-            JobStatus ts = mJobSet.valueAt(i);
-            if (ts.getUid() == uid && ts.getJobId() == jobId) {
-                return ts;
-            }
-        }
-        return null;
+        return mJobSet.get(uid, jobId);
     }
 
     /**
-     * @return The live array of JobStatus objects.
+     * Iterate over the set of all jobs, invoking the supplied functor on each.  This is for
+     * customers who need to examine each job; we'd much rather not have to generate
+     * transient unified collections for them to iterate over and then discard, or creating
+     * iterators every time a client needs to perform a sweep.
      */
-    public ArraySet<JobStatus> getJobs() {
-        return mJobSet;
+    public void forEachJob(JobStatusFunctor functor) {
+        mJobSet.forEachJob(functor);
+    }
+
+    public interface JobStatusFunctor {
+        public void process(JobStatus jobStatus);
     }
 
     /** Version of the db schema. */
@@ -261,7 +240,7 @@
     }
 
     @VisibleForTesting
-    public void readJobMapFromDisk(ArraySet<JobStatus> jobSet) {
+    public void readJobMapFromDisk(JobSet jobSet) {
         new ReadJobMapFromDiskRunnable(jobSet).run();
     }
 
@@ -273,21 +252,19 @@
         @Override
         public void run() {
             final long startElapsed = SystemClock.elapsedRealtime();
-            List<JobStatus> mStoreCopy = new ArrayList<JobStatus>();
+            final List<JobStatus> storeCopy = new ArrayList<JobStatus>();
             synchronized (mLock) {
-                // Copy over the jobs so we can release the lock before writing.
-                for (int i=0; i<mJobSet.size(); i++) {
-                    JobStatus jobStatus = mJobSet.valueAt(i);
-
-                    if (!jobStatus.isPersisted()){
-                        continue;
+                // Clone the jobs so we can release the lock before writing.
+                mJobSet.forEachJob(new JobStatusFunctor() {
+                    @Override
+                    public void process(JobStatus job) {
+                        if (job.isPersisted()) {
+                            storeCopy.add(new JobStatus(job));
+                        }
                     }
-
-                    JobStatus copy = new JobStatus(jobStatus);
-                    mStoreCopy.add(copy);
-                }
+                });
             }
-            writeJobsMapImpl(mStoreCopy);
+            writeJobsMapImpl(storeCopy);
             if (JobSchedulerService.DEBUG) {
                 Slog.v(TAG, "Finished writing, took " + (SystemClock.elapsedRealtime()
                         - startElapsed) + "ms");
@@ -440,13 +417,13 @@
      * need to go through {@link JobStore#add(com.android.server.job.controllers.JobStatus)}.
      */
     private class ReadJobMapFromDiskRunnable implements Runnable {
-        private final ArraySet<JobStatus> jobSet;
+        private final JobSet jobSet;
 
         /**
          * @param jobSet Reference to the (empty) set of JobStatus objects that back the JobStore,
          *               so that after disk read we can populate it directly.
          */
-        ReadJobMapFromDiskRunnable(ArraySet<JobStatus> jobSet) {
+        ReadJobMapFromDiskRunnable(JobSet jobSet) {
             this.jobSet = jobSet;
         }
 
@@ -759,4 +736,122 @@
             return Pair.create(earliestRunTimeElapsed, latestRunTimeElapsed);
         }
     }
+
+    static class JobSet {
+        // Key is the getUid() originator of the jobs in each sheaf
+        private SparseArray<ArraySet<JobStatus>> mJobs;
+
+        public JobSet() {
+            mJobs = new SparseArray<ArraySet<JobStatus>>();
+        }
+
+        public List<JobStatus> getJobsByUid(int uid) {
+            ArrayList<JobStatus> matchingJobs = new ArrayList<JobStatus>();
+            ArraySet<JobStatus> jobs = mJobs.get(uid);
+            if (jobs != null) {
+                matchingJobs.addAll(jobs);
+            }
+            return matchingJobs;
+        }
+
+        // By user, not by uid, so we need to traverse by key and check
+        public List<JobStatus> getJobsByUser(int userId) {
+            ArrayList<JobStatus> result = new ArrayList<JobStatus>();
+            for (int i = mJobs.size() - 1; i >= 0; i--) {
+                if (UserHandle.getUserId(mJobs.keyAt(i)) == userId) {
+                    ArraySet<JobStatus> jobs = mJobs.get(i);
+                    if (jobs != null) {
+                        result.addAll(jobs);
+                    }
+                }
+            }
+            return result;
+        }
+
+        public boolean add(JobStatus job) {
+            final int uid = job.getUid();
+            ArraySet<JobStatus> jobs = mJobs.get(uid);
+            if (jobs == null) {
+                jobs = new ArraySet<JobStatus>();
+                mJobs.put(uid, jobs);
+            }
+            return jobs.add(job);
+        }
+
+        public boolean remove(JobStatus job) {
+            final int uid = job.getUid();
+            ArraySet<JobStatus> jobs = mJobs.get(uid);
+            boolean didRemove = (jobs != null) ? jobs.remove(job) : false;
+            if (didRemove && jobs.size() == 0) {
+                // no more jobs for this uid; let the now-empty set object be GC'd.
+                mJobs.remove(uid);
+            }
+            return didRemove;
+        }
+
+        public boolean contains(JobStatus job) {
+            final int uid = job.getUid();
+            ArraySet<JobStatus> jobs = mJobs.get(uid);
+            return jobs != null && jobs.contains(job);
+        }
+
+        public JobStatus get(int uid, int jobId) {
+            ArraySet<JobStatus> jobs = mJobs.get(uid);
+            if (jobs != null) {
+                for (int i = jobs.size() - 1; i >= 0; i--) {
+                    JobStatus job = jobs.valueAt(i);
+                    if (job.getJobId() == jobId) {
+                        return job;
+                    }
+                }
+            }
+            return null;
+        }
+
+        // Inefficient; use only for testing
+        public List<JobStatus> getAllJobs() {
+            ArrayList<JobStatus> allJobs = new ArrayList<JobStatus>(size());
+            for (int i = mJobs.size(); i >= 0; i--) {
+                allJobs.addAll(mJobs.valueAt(i));
+            }
+            return allJobs;
+        }
+
+        public void clear() {
+            mJobs.clear();
+        }
+
+        public int size() {
+            int total = 0;
+            for (int i = mJobs.size() - 1; i >= 0; i--) {
+                total += mJobs.valueAt(i).size();
+            }
+            return total;
+        }
+
+        // We only want to count the jobs that this uid has scheduled on its own
+        // behalf, not those that the app has scheduled on someone else's behalf.
+        public int countJobsForUid(int uid) {
+            int total = 0;
+            ArraySet<JobStatus> jobs = mJobs.get(uid);
+            if (jobs != null) {
+                for (int i = jobs.size() - 1; i >= 0; i--) {
+                    JobStatus job = jobs.valueAt(i);
+                    if (job.getUid() == job.getSourceUid()) {
+                        total++;
+                    }
+                }
+            }
+            return total;
+        }
+
+        public void forEachJob(JobStatusFunctor functor) {
+            for (int uidIndex = mJobs.size() - 1; uidIndex >= 0; uidIndex--) {
+                ArraySet<JobStatus> jobs = mJobs.valueAt(uidIndex);
+                for (int i = jobs.size() - 1; i >= 0; i--) {
+                    functor.process(jobs.valueAt(i));
+                }
+            }
+        }
+    }
 }
diff --git a/services/tests/servicestests/src/com/android/server/job/JobStoreTest.java b/services/tests/servicestests/src/com/android/server/job/JobStoreTest.java
index c786036..e64481e 100644
--- a/services/tests/servicestests/src/com/android/server/job/JobStoreTest.java
+++ b/services/tests/servicestests/src/com/android/server/job/JobStoreTest.java
@@ -12,6 +12,7 @@
 import android.util.Log;
 import android.util.ArraySet;
 
+import com.android.server.job.JobStore.JobSet;
 import com.android.server.job.controllers.JobStatus;
 
 import java.util.Iterator;
@@ -62,11 +63,11 @@
         mTaskStoreUnderTest.add(ts);
         Thread.sleep(IO_WAIT);
         // Manually load tasks from xml file.
-        final ArraySet<JobStatus> jobStatusSet = new ArraySet<JobStatus>();
+        final JobSet jobStatusSet = new JobSet();
         mTaskStoreUnderTest.readJobMapFromDisk(jobStatusSet);
 
         assertEquals("Didn't get expected number of persisted tasks.", 1, jobStatusSet.size());
-        final JobStatus loadedTaskStatus = jobStatusSet.iterator().next();
+        final JobStatus loadedTaskStatus = jobStatusSet.getAllJobs().get(0);
         assertTasksEqual(task, loadedTaskStatus.getJob());
         assertTrue("JobStore#contains invalid.", mTaskStoreUnderTest.containsJob(ts));
         assertEquals("Different uids.", SOME_UID, loadedTaskStatus.getUid());
@@ -97,10 +98,10 @@
         mTaskStoreUnderTest.add(taskStatus2);
         Thread.sleep(IO_WAIT);
 
-        final ArraySet<JobStatus> jobStatusSet = new ArraySet<JobStatus>();
+        final JobSet jobStatusSet = new JobSet();
         mTaskStoreUnderTest.readJobMapFromDisk(jobStatusSet);
         assertEquals("Incorrect # of persisted tasks.", 2, jobStatusSet.size());
-        Iterator<JobStatus> it = jobStatusSet.iterator();
+        Iterator<JobStatus> it = jobStatusSet.getAllJobs().iterator();
         JobStatus loaded1 = it.next();
         JobStatus loaded2 = it.next();
 
@@ -145,10 +146,10 @@
         mTaskStoreUnderTest.add(taskStatus);
         Thread.sleep(IO_WAIT);
 
-        final ArraySet<JobStatus> jobStatusSet = new ArraySet<JobStatus>();
+        final JobSet jobStatusSet = new JobSet();
         mTaskStoreUnderTest.readJobMapFromDisk(jobStatusSet);
         assertEquals("Incorrect # of persisted tasks.", 1, jobStatusSet.size());
-        JobStatus loaded = jobStatusSet.iterator().next();
+        JobStatus loaded = jobStatusSet.getAllJobs().iterator().next();
         assertTasksEqual(task, loaded.getJob());
     }
     public void testWritingTaskWithSourcePackage() throws Exception {
@@ -163,10 +164,10 @@
         mTaskStoreUnderTest.add(taskStatus);
         Thread.sleep(IO_WAIT);
 
-        final ArraySet<JobStatus> jobStatusSet = new ArraySet<JobStatus>();
+        final JobSet jobStatusSet = new JobSet();
         mTaskStoreUnderTest.readJobMapFromDisk(jobStatusSet);
         assertEquals("Incorrect # of persisted tasks.", 1, jobStatusSet.size());
-        JobStatus loaded = jobStatusSet.iterator().next();
+        JobStatus loaded = jobStatusSet.getAllJobs().iterator().next();
         assertEquals("Source package not equal.", loaded.getSourcePackageName(),
                 taskStatus.getSourcePackageName());
         assertEquals("Source user not equal.", loaded.getSourceUserId(),
@@ -184,10 +185,10 @@
         mTaskStoreUnderTest.add(taskStatus);
         Thread.sleep(IO_WAIT);
 
-        final ArraySet<JobStatus> jobStatusSet = new ArraySet<JobStatus>();
+        final JobSet jobStatusSet = new JobSet();
         mTaskStoreUnderTest.readJobMapFromDisk(jobStatusSet);
         assertEquals("Incorrect # of persisted tasks.", 1, jobStatusSet.size());
-        JobStatus loaded = jobStatusSet.iterator().next();
+        JobStatus loaded = jobStatusSet.getAllJobs().iterator().next();
         assertEquals("Period not equal.", loaded.getJob().getIntervalMillis(),
                 taskStatus.getJob().getIntervalMillis());
         assertEquals("Flex not equal.", loaded.getJob().getFlexMillis(),
@@ -211,10 +212,10 @@
         mTaskStoreUnderTest.add(js);
         Thread.sleep(IO_WAIT);
 
-        final ArraySet<JobStatus> jobStatusSet = new ArraySet<JobStatus>();
+        final JobSet jobStatusSet = new JobSet();
         mTaskStoreUnderTest.readJobMapFromDisk(jobStatusSet);
         assertEquals("Incorrect # of persisted tasks.", 1, jobStatusSet.size());
-        JobStatus loaded = jobStatusSet.iterator().next();
+        JobStatus loaded = jobStatusSet.getAllJobs().iterator().next();
 
         // Assert early runtime was clamped to be under now + period. We can do <= here b/c we'll
         // call SystemClock.elapsedRealtime after doing the disk i/o.
@@ -234,9 +235,9 @@
         final JobStatus js = JobStatus.createFromJobInfo(b.build(), SOME_UID, null, -1);
         mTaskStoreUnderTest.add(js);
         Thread.sleep(IO_WAIT);
-        final ArraySet<JobStatus> jobStatusSet = new ArraySet<JobStatus>();
+        final JobSet jobStatusSet = new JobSet();
         mTaskStoreUnderTest.readJobMapFromDisk(jobStatusSet);
-        JobStatus loaded = jobStatusSet.iterator().next();
+        JobStatus loaded = jobStatusSet.getAllJobs().iterator().next();
         assertEquals("Priority not correctly persisted.", 42, loaded.getPriority());
     }
 
@@ -255,10 +256,10 @@
         JobStatus jsPersisted = JobStatus.createFromJobInfo(b.build(), SOME_UID, null, -1);
         mTaskStoreUnderTest.add(jsPersisted);
         Thread.sleep(IO_WAIT);
-        final ArraySet<JobStatus> jobStatusSet = new ArraySet<JobStatus>();
+        final JobSet jobStatusSet = new JobSet();
         mTaskStoreUnderTest.readJobMapFromDisk(jobStatusSet);
         assertEquals("Job count is incorrect.", 1, jobStatusSet.size());
-        JobStatus jobStatus = jobStatusSet.iterator().next();
+        JobStatus jobStatus = jobStatusSet.getAllJobs().iterator().next();
         assertEquals("Wrong job persisted.", 43, jobStatus.getJobId());
     }