Remove monitor locks in TestScheduler.

* That takes care of multi-threaded tests where we hold the monitor lock
for longer than necessary.

Test: Added multi-threaded @Large tests.
Fixes: b/139385052
Change-Id: I76c8ceb5ab683994383820377ff25fcc560bf70d
diff --git a/work/workmanager-testing/src/androidTest/java/androidx/work/testing/TestSchedulerTest.java b/work/workmanager-testing/src/androidTest/java/androidx/work/testing/TestSchedulerTest.java
index 868ce63..83f2d59 100644
--- a/work/workmanager-testing/src/androidTest/java/androidx/work/testing/TestSchedulerTest.java
+++ b/work/workmanager-testing/src/androidTest/java/androidx/work/testing/TestSchedulerTest.java
@@ -20,10 +20,15 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 
 import android.content.Context;
+import android.os.Handler;
+import android.os.Looper;
+import android.util.Log;
 
+import androidx.lifecycle.Observer;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
+import androidx.work.Configuration;
 import androidx.work.Constraints;
 import androidx.work.NetworkType;
 import androidx.work.OneTimeWorkRequest;
@@ -41,7 +46,11 @@
 import org.junit.runner.RunWith;
 
 import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 
 @RunWith(AndroidJUnit4.class)
@@ -50,11 +59,18 @@
 
     private Context mContext;
     private TestDriver mTestDriver;
+    private Handler mHandler;
 
     @Before
     public void setUp() {
         mContext = ApplicationProvider.getApplicationContext();
-        WorkManagerTestInitHelper.initializeTestWorkManager(mContext);
+        mHandler = new Handler(Looper.getMainLooper());
+        // Don't set the task executor
+        Configuration configuration = new Configuration.Builder()
+                .setExecutor(new SynchronousExecutor())
+                .setMinimumLoggingLevel(Log.DEBUG)
+                .build();
+        WorkManagerTestInitHelper.initializeTestWorkManager(mContext, configuration);
         mTestDriver = WorkManagerTestInitHelper.getTestDriver(mContext);
         CountingTestWorker.COUNT.set(0);
     }
@@ -186,6 +202,106 @@
         mTestDriver.setPeriodDelayMet(request.getId());
     }
 
+    @Test
+    @LargeTest
+    public void testWorker_multipleSetInitialDelayMet_noDeadLock()
+            throws InterruptedException, ExecutionException {
+
+        Configuration configuration = new Configuration.Builder()
+                .setMinimumLoggingLevel(Log.DEBUG)
+                .build();
+        WorkManagerTestInitHelper.initializeTestWorkManager(mContext, configuration);
+        mTestDriver = WorkManagerTestInitHelper.getTestDriver(mContext);
+
+        // This should not deadlock
+        final OneTimeWorkRequest request = createWorkRequest();
+        final WorkManager workManager = WorkManager.getInstance(mContext);
+        workManager.enqueue(request).getResult().get();
+        mTestDriver.setInitialDelayMet(request.getId());
+        mTestDriver.setInitialDelayMet(request.getId());
+
+        final CountDownLatch latch = new CountDownLatch(1);
+        // Use the main looper to observe LiveData because we are using a SerialExecutor which is
+        // wrapping a SynchronousExecutor.
+        mHandler.post(new Runnable() {
+            @Override
+            public void run() {
+                workManager.getWorkInfoByIdLiveData(request.getId()).observeForever(
+                        new Observer<WorkInfo>() {
+                            @Override
+                            public void onChanged(WorkInfo workInfo) {
+                                if (workInfo != null && workInfo.getState().isFinished()) {
+                                    latch.countDown();
+                                }
+                            }
+                        });
+            }
+        });
+
+        latch.await(5, TimeUnit.SECONDS);
+        assertThat(latch.getCount(), is(0L));
+    }
+
+    @Test
+    @LargeTest
+    public void testWorker_multipleSetInitialDelayMetMultiThreaded_noDeadLock()
+            throws InterruptedException {
+
+        Configuration configuration = new Configuration.Builder()
+                .setMinimumLoggingLevel(Log.DEBUG)
+                .build();
+        WorkManagerTestInitHelper.initializeTestWorkManager(mContext, configuration);
+        mTestDriver = WorkManagerTestInitHelper.getTestDriver(mContext);
+
+        // This should not deadlock
+        final WorkManager workManager = WorkManager.getInstance(mContext);
+        int numberOfWorkers = 10;
+        final ExecutorService service = Executors.newFixedThreadPool(numberOfWorkers);
+        for (int i = 0; i < numberOfWorkers; i++) {
+            service.submit(new Runnable() {
+                @Override
+                public void run() {
+                    final OneTimeWorkRequest request = createWorkRequest();
+                    workManager.enqueue(request);
+                    mTestDriver.setInitialDelayMet(request.getId());
+                    mTestDriver.setInitialDelayMet(request.getId());
+                }
+            });
+        }
+
+        final CountDownLatch latch = new CountDownLatch(1);
+        // Use the main looper to observe LiveData because we are using a SerialExecutor which is
+        // wrapping a SynchronousExecutor.
+        mHandler.post(new Runnable() {
+            @Override
+            public void run() {
+                // Using the implicit tag name.
+                workManager.getWorkInfosByTagLiveData(TestWorker.class.getName()).observeForever(
+                        new Observer<List<WorkInfo>>() {
+                            @Override
+                            public void onChanged(List<WorkInfo> workInfos) {
+                                boolean completed = true;
+                                if (workInfos != null && !workInfos.isEmpty()) {
+                                    for (WorkInfo workInfo : workInfos) {
+                                        if (!workInfo.getState().isFinished()) {
+                                            completed = false;
+                                            break;
+                                        }
+                                    }
+                                }
+                                if (completed) {
+                                    latch.countDown();
+                                }
+                            }
+                        });
+            }
+        });
+
+        latch.await(10, TimeUnit.SECONDS);
+        service.shutdownNow();
+        assertThat(latch.getCount(), is(0L));
+    }
+
     private static OneTimeWorkRequest createWorkRequest() {
         return new OneTimeWorkRequest.Builder(TestWorker.class).build();
     }
diff --git a/work/workmanager-testing/src/main/java/androidx/work/testing/TestScheduler.java b/work/workmanager-testing/src/main/java/androidx/work/testing/TestScheduler.java
index afc4f41..1595350 100644
--- a/work/workmanager-testing/src/main/java/androidx/work/testing/TestScheduler.java
+++ b/work/workmanager-testing/src/main/java/androidx/work/testing/TestScheduler.java
@@ -50,8 +50,6 @@
     private final Map<String, InternalWorkState> mPendingWorkStates;
     private final Map<String, InternalWorkState> mTerminatedWorkStates;
 
-    private static final Object sLock = new Object();
-
     TestScheduler(@NonNull Context context) {
         mContext = context;
         mPendingWorkStates = new HashMap<>();
@@ -64,28 +62,24 @@
             return;
         }
 
-        synchronized (sLock) {
-            List<String> workSpecIdsToSchedule = new ArrayList<>(workSpecs.length);
-            for (WorkSpec workSpec : workSpecs) {
-                if (!mPendingWorkStates.containsKey(workSpec.id)) {
-                    mPendingWorkStates.put(workSpec.id, new InternalWorkState(mContext, workSpec));
-                }
-                workSpecIdsToSchedule.add(workSpec.id);
+        List<String> workSpecIdsToSchedule = new ArrayList<>(workSpecs.length);
+        for (WorkSpec workSpec : workSpecs) {
+            if (!mPendingWorkStates.containsKey(workSpec.id)) {
+                mPendingWorkStates.put(workSpec.id, new InternalWorkState(mContext, workSpec));
             }
-            scheduleInternal(workSpecIdsToSchedule);
+            workSpecIdsToSchedule.add(workSpec.id);
         }
+        scheduleInternal(workSpecIdsToSchedule);
     }
 
     @Override
     public void cancel(@NonNull String workSpecId) {
-        synchronized (sLock) {
-            WorkManagerImpl.getInstance(mContext).stopWork(workSpecId);
-            mPendingWorkStates.remove(workSpecId);
-            // We don't need to keep track of cancelled workSpecs. This is because subsequent calls
-            // to enqueue() will no-op because insertWorkSpec in WorkDatabase has a conflict
-            // policy of @Ignore. So TestScheduler will _never_ be asked to schedule those
-            // WorkSpecs.
-        }
+        // We don't need to keep track of cancelled workSpecs. This is because subsequent calls
+        // to enqueue() will no-op because insertWorkSpec in WorkDatabase has a conflict
+        // policy of @Ignore. So TestScheduler will _never_ be asked to schedule those
+        // WorkSpecs.
+        WorkManagerImpl.getInstance(mContext).stopWork(workSpecId);
+        mPendingWorkStates.remove(workSpecId);
     }
 
     /**
@@ -96,17 +90,15 @@
      * @throws IllegalArgumentException if {@code workSpecId} is not enqueued
      */
     void setAllConstraintsMet(@NonNull UUID workSpecId) {
-        synchronized (sLock) {
-            String id = workSpecId.toString();
-            if (!mTerminatedWorkStates.containsKey(id)) {
-                InternalWorkState internalWorkState = mPendingWorkStates.get(id);
-                if (internalWorkState == null) {
-                    throw new IllegalArgumentException(
-                            "Work with id " + workSpecId + " is not enqueued!");
-                }
-                internalWorkState.mConstraintsMet = true;
-                scheduleInternal(Collections.singletonList(workSpecId.toString()));
+        String id = workSpecId.toString();
+        if (!mTerminatedWorkStates.containsKey(id)) {
+            InternalWorkState internalWorkState = mPendingWorkStates.get(id);
+            if (internalWorkState == null) {
+                throw new IllegalArgumentException(
+                        "Work with id " + workSpecId + " is not enqueued!");
             }
+            internalWorkState.mConstraintsMet = true;
+            scheduleInternal(Collections.singletonList(workSpecId.toString()));
         }
     }
 
@@ -118,17 +110,15 @@
      * @throws IllegalArgumentException if {@code workSpecId} is not enqueued
      */
     void setInitialDelayMet(@NonNull UUID workSpecId) {
-        synchronized (sLock) {
-            String id = workSpecId.toString();
-            if (!mTerminatedWorkStates.containsKey(id)) {
-                InternalWorkState internalWorkState = mPendingWorkStates.get(id);
-                if (internalWorkState == null) {
-                    throw new IllegalArgumentException(
-                            "Work with id " + workSpecId + " is not enqueued!");
-                }
-                internalWorkState.mInitialDelayMet = true;
-                scheduleInternal(Collections.singletonList(workSpecId.toString()));
+        String id = workSpecId.toString();
+        if (!mTerminatedWorkStates.containsKey(id)) {
+            InternalWorkState internalWorkState = mPendingWorkStates.get(id);
+            if (internalWorkState == null) {
+                throw new IllegalArgumentException(
+                        "Work with id " + workSpecId + " is not enqueued!");
             }
+            internalWorkState.mInitialDelayMet = true;
+            scheduleInternal(Collections.singletonList(workSpecId.toString()));
         }
     }
 
@@ -140,29 +130,25 @@
      * @throws IllegalArgumentException if {@code workSpecId} is not enqueued
      */
     void setPeriodDelayMet(@NonNull UUID workSpecId) {
-        synchronized (sLock) {
-            String id = workSpecId.toString();
-            InternalWorkState internalWorkState = mPendingWorkStates.get(id);
-            if (internalWorkState == null) {
-                throw new IllegalArgumentException(
-                        "Work with id " + workSpecId + " is not enqueued!");
-            }
-            internalWorkState.mPeriodDelayMet = true;
-            scheduleInternal(Collections.singletonList(workSpecId.toString()));
+        String id = workSpecId.toString();
+        InternalWorkState internalWorkState = mPendingWorkStates.get(id);
+        if (internalWorkState == null) {
+            throw new IllegalArgumentException(
+                    "Work with id " + workSpecId + " is not enqueued!");
         }
+        internalWorkState.mPeriodDelayMet = true;
+        scheduleInternal(Collections.singletonList(workSpecId.toString()));
     }
 
     @Override
     public void onExecuted(@NonNull String workSpecId, boolean needsReschedule) {
-        synchronized (sLock) {
-            InternalWorkState internalWorkState = mPendingWorkStates.get(workSpecId);
-            if (internalWorkState != null) {
-                if (internalWorkState.mWorkSpec.isPeriodic()) {
-                    internalWorkState.reset();
-                } else {
-                    mTerminatedWorkStates.put(workSpecId, internalWorkState);
-                    mPendingWorkStates.remove(workSpecId);
-                }
+        InternalWorkState internalWorkState = mPendingWorkStates.get(workSpecId);
+        if (internalWorkState != null) {
+            if (internalWorkState.mWorkSpec.isPeriodic()) {
+                internalWorkState.reset();
+            } else {
+                mTerminatedWorkStates.put(workSpecId, internalWorkState);
+                mPendingWorkStates.remove(workSpecId);
             }
         }
     }
diff --git a/work/workmanager-testing/src/main/java/androidx/work/testing/TestWorkManagerImpl.java b/work/workmanager-testing/src/main/java/androidx/work/testing/TestWorkManagerImpl.java
index 4e7c916..bb84e45 100644
--- a/work/workmanager-testing/src/main/java/androidx/work/testing/TestWorkManagerImpl.java
+++ b/work/workmanager-testing/src/main/java/androidx/work/testing/TestWorkManagerImpl.java
@@ -49,6 +49,14 @@
 
         // Note: This implies that the call to ForceStopRunnable() actually does nothing.
         // This is okay when testing.
+
+        // IMPORTANT: Leave the main thread executor as a Direct executor. This is very important.
+        // Otherwise we subtly change the order of callbacks. onExecuted() will execute after
+        // a call to StopWorkRunnable(). StopWorkRunnable() removes the pending WorkSpec and
+        // therefore the call to onExecuted() does not add the workSpecId to the list of
+        // terminated WorkSpecs. This is because internalWorkState == null.
+        // Also for PeriodicWorkRequests, Schedulers.schedule() will run before the call to
+        // onExecuted() and therefore PeriodicWorkRequests will always run twice.
         super(
                 context,
                 configuration,