Make concurrency manager code more testable.

1. Split up the job context assignment code to make it more testable.
2. Add a few basic tests for the code. Additional tests will be added
   later.

Bug: 141645789
Test: atest FrameworksMockingServicesTests:JobConcurrencyManagerTest
Change-Id: I346fb1d0696dff76b1ece73d303a707075e3dadc

2 files changed, 252 insertions, 26 deletions

diff --git a/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java b/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java
index 586038427c8d..20bca3530b63 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java
@@ -34,6 +34,7 @@ import android.content.IntentFilter;
 import android.content.pm.UserInfo;
 import android.os.BatteryStats;
 import android.os.Handler;
+import android.os.Looper;
 import android.os.PowerManager;
 import android.os.RemoteException;
 import android.os.ServiceManager;
@@ -181,6 +182,7 @@ class JobConcurrencyManager {
     private final JobSchedulerService mService;
     private final Context mContext;
     private final Handler mHandler;
+    private final Injector mInjector;
 
     private PowerManager mPowerManager;
 
@@ -378,9 +380,15 @@ class JobConcurrencyManager {
     }
 
     JobConcurrencyManager(JobSchedulerService service) {
+        this(service, new Injector());
+    }
+
+    @VisibleForTesting
+    JobConcurrencyManager(JobSchedulerService service, Injector injector) {
         mService = service;
         mLock = mService.mLock;
         mContext = service.getTestableContext();
+        mInjector = injector;
 
         mHandler = JobSchedulerBackgroundThread.getHandler();
 
@@ -414,7 +422,7 @@ class JobConcurrencyManager {
                 ServiceManager.getService(BatteryStats.SERVICE_NAME));
         for (int i = 0; i < STANDARD_CONCURRENCY_LIMIT; i++) {
             mIdleContexts.add(
-                    new JobServiceContext(mService, this, batteryStats,
+                    mInjector.createJobServiceContext(mService, this, batteryStats,
                             mService.mJobPackageTracker, mContext.getMainLooper()));
         }
     }
@@ -657,15 +665,40 @@ class JobConcurrencyManager {
             return;
         }
 
+        prepareForAssignmentDeterminationLocked(
+                mRecycledIdle, mRecycledPreferredUidOnly, mRecycledStoppable);
+
+        if (DEBUG) {
+            Slog.d(TAG, printAssignments("running jobs initial",
+                    mRecycledStoppable, mRecycledPreferredUidOnly));
+        }
+
+        determineAssignmentsLocked(
+                mRecycledChanged, mRecycledIdle, mRecycledPreferredUidOnly, mRecycledStoppable);
+
+        if (DEBUG) {
+            Slog.d(TAG, printAssignments("running jobs final",
+                    mRecycledStoppable, mRecycledPreferredUidOnly, mRecycledChanged));
+
+            Slog.d(TAG, "work count results: " + mWorkCountTracker);
+        }
+
+        carryOutAssignmentChangesLocked(mRecycledChanged);
+
+        cleanUpAfterAssignmentChangesLocked(
+                mRecycledChanged, mRecycledIdle, mRecycledPreferredUidOnly, mRecycledStoppable);
+
+        noteConcurrency();
+    }
+
+    @VisibleForTesting
+    @GuardedBy("mLock")
+    void prepareForAssignmentDeterminationLocked(final ArraySet<ContextAssignment> idle,
+            final List<ContextAssignment> preferredUidOnly,
+            final List<ContextAssignment> stoppable) {
         final PendingJobQueue pendingJobQueue = mService.getPendingJobQueue();
         final List<JobServiceContext> activeServices = mActiveServices;
 
-        // To avoid GC churn, we recycle the arrays.
-        final ArraySet<ContextAssignment> changed = mRecycledChanged;
-        final ArraySet<ContextAssignment> idle = mRecycledIdle;
-        final ArrayList<ContextAssignment> preferredUidOnly = mRecycledPreferredUidOnly;
-        final ArrayList<ContextAssignment> stoppable = mRecycledStoppable;
-
         updateCounterConfigLocked();
         // Reset everything since we'll re-evaluate the current state.
         mWorkCountTracker.resetCounts();
@@ -719,15 +752,21 @@ class JobConcurrencyManager {
             assignment.context = jsc;
             idle.add(assignment);
         }
-        if (DEBUG) {
-            Slog.d(TAG, printAssignments("running jobs initial", stoppable, preferredUidOnly));
-        }
 
         mWorkCountTracker.onCountDone();
+    }
 
-        JobStatus nextPending;
+    @VisibleForTesting
+    @GuardedBy("mLock")
+    void determineAssignmentsLocked(final ArraySet<ContextAssignment> changed,
+            final ArraySet<ContextAssignment> idle,
+            final List<ContextAssignment> preferredUidOnly,
+            final List<ContextAssignment> stoppable) {
+        final PendingJobQueue pendingJobQueue = mService.getPendingJobQueue();
+        final List<JobServiceContext> activeServices = mActiveServices;
         pendingJobQueue.resetIterator();
-        int projectedRunningCount = numRunningJobs;
+        JobStatus nextPending;
+        int projectedRunningCount = activeServices.size();
         while ((nextPending = pendingJobQueue.next()) != null) {
             if (mRunningJobs.contains(nextPending)) {
                 // Should never happen.
@@ -895,13 +934,10 @@ class JobConcurrencyManager {
                         packageStats);
             }
         }
-        if (DEBUG) {
-            Slog.d(TAG, printAssignments("running jobs final",
-                    stoppable, preferredUidOnly, changed));
-
-            Slog.d(TAG, "assignJobsToContexts: " + mWorkCountTracker.toString());
-        }
+    }
 
+    @GuardedBy("mLock")
+    private void carryOutAssignmentChangesLocked(final ArraySet<ContextAssignment> changed) {
         for (int c = changed.size() - 1; c >= 0; --c) {
             final ContextAssignment assignment = changed.valueAt(c);
             final JobStatus js = assignment.context.getRunningJobLocked();
@@ -925,6 +961,13 @@ class JobConcurrencyManager {
             assignment.clear();
             mContextAssignmentPool.release(assignment);
         }
+    }
+
+    @GuardedBy("mLock")
+    private void cleanUpAfterAssignmentChangesLocked(final ArraySet<ContextAssignment> changed,
+            final ArraySet<ContextAssignment> idle,
+            final List<ContextAssignment> preferredUidOnly,
+            final List<ContextAssignment> stoppable) {
         for (int s = stoppable.size() - 1; s >= 0; --s) {
             final ContextAssignment assignment = stoppable.get(s);
             assignment.clear();
@@ -947,7 +990,6 @@ class JobConcurrencyManager {
         preferredUidOnly.clear();
         mWorkCountTracker.resetStagingCount();
         mActivePkgStats.forEach(mPackageStatsStagingCountClearer);
-        noteConcurrency();
     }
 
     @GuardedBy("mLock")
@@ -1496,7 +1538,7 @@ class JobConcurrencyManager {
 
     @NonNull
     private JobServiceContext createNewJobServiceContext() {
-        return new JobServiceContext(mService, this,
+        return mInjector.createJobServiceContext(mService, this,
                 IBatteryStats.Stub.asInterface(
                         ServiceManager.getService(BatteryStats.SERVICE_NAME)),
                 mService.mJobPackageTracker, mContext.getMainLooper());
@@ -1778,6 +1820,10 @@ class JobConcurrencyManager {
     @VisibleForTesting
     static class WorkTypeConfig {
         @VisibleForTesting
+        static final String KEY_PREFIX_MAX = CONFIG_KEY_PREFIX_CONCURRENCY + "max_";
+        @VisibleForTesting
+        static final String KEY_PREFIX_MIN = CONFIG_KEY_PREFIX_CONCURRENCY + "min_";
+        @VisibleForTesting
         static final String KEY_PREFIX_MAX_TOTAL = CONFIG_KEY_PREFIX_CONCURRENCY + "max_total_";
         private static final String KEY_PREFIX_MAX_TOP = CONFIG_KEY_PREFIX_CONCURRENCY + "max_top_";
         private static final String KEY_PREFIX_MAX_FGS = CONFIG_KEY_PREFIX_CONCURRENCY + "max_fgs_";
@@ -2329,7 +2375,8 @@ class JobConcurrencyManager {
         }
     }
 
-    private static final class ContextAssignment {
+    @VisibleForTesting
+    static final class ContextAssignment {
         public JobServiceContext context;
         public int preferredUid = JobServiceContext.NO_PREFERRED_UID;
         public int workType = WORK_TYPE_NONE;
@@ -2378,4 +2425,15 @@ class JobConcurrencyManager {
         mActivePkgStats.add(userId, packageName, packageStats);
         return packageStats;
     }
+
+    @VisibleForTesting
+    static class Injector {
+        @NonNull
+        JobServiceContext createJobServiceContext(JobSchedulerService service,
+                JobConcurrencyManager concurrencyManager, IBatteryStats batteryStats,
+                JobPackageTracker tracker, Looper looper) {
+            return new JobServiceContext(service, concurrencyManager, batteryStats,
+                    tracker, looper);
+        }
+    }
 }
diff --git a/services/tests/mockingservicestests/src/com/android/server/job/JobConcurrencyManagerTest.java b/services/tests/mockingservicestests/src/com/android/server/job/JobConcurrencyManagerTest.java
index b354c7b3d1f3..f46877e5a8c6 100644
--- a/services/tests/mockingservicestests/src/com/android/server/job/JobConcurrencyManagerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/job/JobConcurrencyManagerTest.java
@@ -16,30 +16,48 @@
 
 package com.android.server.job;
 
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static com.android.server.job.JobConcurrencyManager.KEY_PKG_CONCURRENCY_LIMIT_EJ;
 import static com.android.server.job.JobConcurrencyManager.KEY_PKG_CONCURRENCY_LIMIT_REGULAR;
 
+import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
 
-import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.anyLong;
+import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import android.annotation.Nullable;
 import android.app.ActivityManagerInternal;
+import android.app.AppGlobals;
 import android.app.job.JobInfo;
 import android.content.ComponentName;
 import android.content.Context;
+import android.content.pm.IPackageManager;
 import android.content.pm.UserInfo;
 import android.content.res.Resources;
+import android.os.Looper;
 import android.os.UserHandle;
 import android.provider.DeviceConfig;
+import android.util.ArraySet;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
+import static com.android.server.job.JobConcurrencyManager.WORK_TYPE_BG;
+import static com.android.server.job.JobConcurrencyManager.WORK_TYPE_BGUSER;
+import static com.android.server.job.JobConcurrencyManager.WORK_TYPE_BGUSER_IMPORTANT;
+import static com.android.server.job.JobConcurrencyManager.WORK_TYPE_EJ;
+import static com.android.server.job.JobConcurrencyManager.WORK_TYPE_FGS;
+import static com.android.server.job.JobConcurrencyManager.WORK_TYPE_NONE;
+import static com.android.server.job.JobConcurrencyManager.WORK_TYPE_TOP;
+
 import com.android.internal.R;
+import com.android.internal.app.IBatteryStats;
 import com.android.server.LocalServices;
 import com.android.server.job.JobConcurrencyManager.GracePeriodObserver;
 import com.android.server.job.JobConcurrencyManager.WorkTypeConfig;
@@ -52,6 +70,12 @@ import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoSession;
+import org.mockito.quality.Strictness;
+
+import java.util.ArrayList;
+import java.util.List;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
@@ -64,10 +88,24 @@ public final class JobConcurrencyManagerTest {
     private int mDefaultUserId;
     private GracePeriodObserver mGracePeriodObserver;
     private Context mContext;
+    private InjectorForTest mInjector;
+    private MockitoSession mMockingSession;
     private Resources mResources;
     private PendingJobQueue mPendingJobQueue;
     private DeviceConfig.Properties.Builder mConfigBuilder;
 
+    @Mock
+    private IPackageManager mIPackageManager;
+
+    static class InjectorForTest extends JobConcurrencyManager.Injector {
+        @Override
+        JobServiceContext createJobServiceContext(JobSchedulerService service,
+                JobConcurrencyManager concurrencyManager, IBatteryStats batteryStats,
+                JobPackageTracker tracker, Looper looper) {
+            return mock(JobServiceContext.class);
+        }
+    }
+
     @BeforeClass
     public static void setUpOnce() {
         LocalServices.addService(UserManagerInternal.class, mock(UserManagerInternal.class));
@@ -83,6 +121,11 @@ public final class JobConcurrencyManagerTest {
 
     @Before
     public void setUp() {
+        mMockingSession = mockitoSession()
+                .initMocks(this)
+                .mockStatic(AppGlobals.class)
+                .strictness(Strictness.LENIENT)
+                .startMocking();
         final JobSchedulerService jobSchedulerService = mock(JobSchedulerService.class);
         mContext = mock(Context.class);
         mResources = mock(Resources.class);
@@ -93,7 +136,9 @@ public final class JobConcurrencyManagerTest {
         mConfigBuilder = new DeviceConfig.Properties.Builder(DeviceConfig.NAMESPACE_JOB_SCHEDULER);
         mPendingJobQueue = new PendingJobQueue();
         doReturn(mPendingJobQueue).when(jobSchedulerService).getPendingJobQueue();
-        mJobConcurrencyManager = new JobConcurrencyManager(jobSchedulerService);
+        doReturn(mIPackageManager).when(AppGlobals::getPackageManager);
+        mInjector = new InjectorForTest();
+        mJobConcurrencyManager = new JobConcurrencyManager(jobSchedulerService, mInjector);
         mGracePeriodObserver = mock(GracePeriodObserver.class);
         mUserManagerInternal = LocalServices.getService(UserManagerInternal.class);
         mActivityManagerInternal = LocalServices.getService(ActivityManagerInternal.class);
@@ -106,6 +151,74 @@ public final class JobConcurrencyManagerTest {
     @After
     public void tearDown() throws Exception {
         resetConfig();
+        if (mMockingSession != null) {
+            mMockingSession.finishMocking();
+        }
+    }
+
+    @Test
+    public void testPrepareForAssignmentDetermination_noJobs() {
+        mPendingJobQueue.clear();
+
+        final ArraySet<JobConcurrencyManager.ContextAssignment> idle = new ArraySet<>();
+        final List<JobConcurrencyManager.ContextAssignment> preferredUidOnly = new ArrayList<>();
+        final List<JobConcurrencyManager.ContextAssignment> stoppable = new ArrayList<>();
+        mJobConcurrencyManager
+                .prepareForAssignmentDeterminationLocked(idle, preferredUidOnly, stoppable);
+
+        assertEquals(JobConcurrencyManager.STANDARD_CONCURRENCY_LIMIT, idle.size());
+        assertEquals(0, preferredUidOnly.size());
+        assertEquals(0, stoppable.size());
+    }
+
+    @Test
+    public void testPrepareForAssignmentDetermination_onlyPendingJobs() {
+        final ArraySet<JobStatus> jobs = new ArraySet<>();
+        for (int i = 0; i < JobConcurrencyManager.STANDARD_CONCURRENCY_LIMIT; ++i) {
+            JobStatus job = createJob(mDefaultUserId * UserHandle.PER_USER_RANGE + i);
+            mPendingJobQueue.add(job);
+            jobs.add(job);
+        }
+
+        final ArraySet<JobConcurrencyManager.ContextAssignment> idle = new ArraySet<>();
+        final List<JobConcurrencyManager.ContextAssignment> preferredUidOnly = new ArrayList<>();
+        final List<JobConcurrencyManager.ContextAssignment> stoppable = new ArrayList<>();
+        mJobConcurrencyManager
+                .prepareForAssignmentDeterminationLocked(idle, preferredUidOnly, stoppable);
+
+        assertEquals(JobConcurrencyManager.STANDARD_CONCURRENCY_LIMIT, idle.size());
+        assertEquals(0, preferredUidOnly.size());
+        assertEquals(0, stoppable.size());
+    }
+
+    @Test
+    public void testDetermineAssignments_allRegular() throws Exception {
+        setConcurrencyConfig(JobConcurrencyManager.STANDARD_CONCURRENCY_LIMIT,
+                new TypeConfig(WORK_TYPE_BG, 0, JobConcurrencyManager.STANDARD_CONCURRENCY_LIMIT));
+        final ArraySet<JobStatus> jobs = new ArraySet<>();
+        for (int i = 0; i < JobConcurrencyManager.STANDARD_CONCURRENCY_LIMIT; ++i) {
+            final int uid = mDefaultUserId * UserHandle.PER_USER_RANGE + i;
+            final String sourcePkgName = "com.source.package." + UserHandle.getAppId(uid);
+            setPackageUid(sourcePkgName, uid);
+            final JobStatus job = createJob(uid, sourcePkgName);
+            mPendingJobQueue.add(job);
+            jobs.add(job);
+        }
+
+        final ArraySet<JobConcurrencyManager.ContextAssignment> changed = new ArraySet<>();
+        final ArraySet<JobConcurrencyManager.ContextAssignment> idle = new ArraySet<>();
+        final List<JobConcurrencyManager.ContextAssignment> preferredUidOnly = new ArrayList<>();
+        final List<JobConcurrencyManager.ContextAssignment> stoppable = new ArrayList<>();
+        mJobConcurrencyManager
+                .prepareForAssignmentDeterminationLocked(idle, preferredUidOnly, stoppable);
+        mJobConcurrencyManager
+                .determineAssignmentsLocked(changed, idle, preferredUidOnly, stoppable);
+
+        assertEquals(JobConcurrencyManager.STANDARD_CONCURRENCY_LIMIT, changed.size());
+        for (int i = changed.size() - 1; i >= 0; --i) {
+            jobs.remove(changed.valueAt(i).newJob);
+        }
+        assertTrue("Some jobs weren't assigned", jobs.isEmpty());
     }
 
     @Test
@@ -403,16 +516,58 @@ public final class JobConcurrencyManagerTest {
     }
 
     private static JobStatus createJob(int uid) {
-        return createJob(uid, 1);
+        return createJob(uid, 1, null);
+    }
+
+    private static JobStatus createJob(int uid, String sourcePackageName) {
+        return createJob(uid, 1, sourcePackageName);
     }
 
     private static JobStatus createJob(int uid, int jobId) {
+        return createJob(uid, jobId, null);
+    }
+
+    private static JobStatus createJob(int uid, int jobId, @Nullable String sourcePackageName) {
         return JobStatus.createFromJobInfo(
                 new JobInfo.Builder(jobId, new ComponentName("foo", "bar")).build(), uid,
-                null, UserHandle.getUserId(uid), "JobConcurrencyManagerTest");
+                sourcePackageName, UserHandle.getUserId(uid), "JobConcurrencyManagerTest");
     }
 
-    private void setConcurrencyConfig(int total) throws Exception {
+    private static final class TypeConfig {
+        public final String workTypeString;
+        public final int min;
+        public final int max;
+
+        private TypeConfig(@JobConcurrencyManager.WorkType int workType, int min, int max) {
+            switch (workType) {
+                case WORK_TYPE_TOP:
+                    workTypeString = "top";
+                    break;
+                case WORK_TYPE_FGS:
+                    workTypeString = "fgs";
+                    break;
+                case WORK_TYPE_EJ:
+                    workTypeString = "ej";
+                    break;
+                case WORK_TYPE_BG:
+                    workTypeString = "bg";
+                    break;
+                case WORK_TYPE_BGUSER:
+                    workTypeString = "bguser";
+                    break;
+                case WORK_TYPE_BGUSER_IMPORTANT:
+                    workTypeString = "bguser_important";
+                    break;
+                case WORK_TYPE_NONE:
+                default:
+                    throw new IllegalArgumentException("invalid work type: " + workType);
+            }
+            this.min = min;
+            this.max = max;
+        }
+    }
+
+    private void setConcurrencyConfig(int total, TypeConfig... typeConfigs) throws Exception {
         // Set the values for all memory states so we don't have to worry about memory on the device
         // during testing.
         final String[] identifiers = {
@@ -422,10 +577,23 @@ public final class JobConcurrencyManagerTest {
         for (String identifier : identifiers) {
             mConfigBuilder
                     .setInt(WorkTypeConfig.KEY_PREFIX_MAX_TOTAL + identifier, total);
+            for (TypeConfig config : typeConfigs) {
+                mConfigBuilder.setInt(
+                        WorkTypeConfig.KEY_PREFIX_MAX + config.workTypeString + "_" + identifier,
+                        config.max);
+                mConfigBuilder.setInt(
+                        WorkTypeConfig.KEY_PREFIX_MIN + config.workTypeString + "_" + identifier,
+                        config.min);
+            }
         }
         updateDeviceConfig();
     }
 
+    private void setPackageUid(final String pkgName, final int uid) throws Exception {
+        doReturn(uid).when(mIPackageManager)
+                .getPackageUid(eq(pkgName), anyLong(), eq(UserHandle.getUserId(uid)));
+    }
+
     private void updateDeviceConfig() throws Exception {
         DeviceConfig.setProperties(mConfigBuilder.build());
         mJobConcurrencyManager.updateConfigLocked();