Fixes flaky benchmarks.

* Ran the same benchmark 100 times locally and it looked stable.
* Delete the database at the end of the benchmark for a more realistic benchmark.
* Use a CompletableDeferred<Unit> instead of a Job() for more control on parenting.

Test: Tested benchmark locally.
Fixes: b/138850285
Change-Id: Ie9a65c08daf36c77c101fe0fcb4596ceebfe7091
diff --git a/work/workmanager-benchmark/src/androidTest/java/androidx/work/benchmark/DispatchingExecutor.kt b/work/workmanager-benchmark/src/androidTest/java/androidx/work/benchmark/DispatchingExecutor.kt
index 5b6b635..9d0fb58 100644
--- a/work/workmanager-benchmark/src/androidTest/java/androidx/work/benchmark/DispatchingExecutor.kt
+++ b/work/workmanager-benchmark/src/androidTest/java/androidx/work/benchmark/DispatchingExecutor.kt
@@ -16,21 +16,29 @@
 
 package androidx.work.benchmark
 
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
 import java.util.concurrent.Executor
 
 /**
  * An [Executor] where we can await termination of all commands.
  */
 class DispatchingExecutor : Executor {
-    val job = Job()
+    private val job = CompletableDeferred<Unit>()
     private val scope = CoroutineScope(Dispatchers.Default + job)
     override fun execute(command: Runnable) {
         scope.launch {
             command.run()
         }
     }
+
+    fun waitForIdle() {
+        runBlocking {
+            job.complete(Unit)
+            job.join()
+        }
+    }
 }
diff --git a/work/workmanager-benchmark/src/androidTest/java/androidx/work/benchmark/InitializeBenchmark.kt b/work/workmanager-benchmark/src/androidTest/java/androidx/work/benchmark/InitializeBenchmark.kt
index 3d33749..0d348fe 100644
--- a/work/workmanager-benchmark/src/androidTest/java/androidx/work/benchmark/InitializeBenchmark.kt
+++ b/work/workmanager-benchmark/src/androidTest/java/androidx/work/benchmark/InitializeBenchmark.kt
@@ -30,7 +30,6 @@
 import androidx.work.impl.WorkManagerImpl
 import androidx.work.impl.utils.SerialExecutor
 import androidx.work.impl.utils.taskexecutor.TaskExecutor
-import kotlinx.coroutines.runBlocking
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
@@ -89,12 +88,9 @@
             val database = WorkDatabase.create(context, configuration.taskExecutor, false)
             WorkManagerImpl(context, configuration, taskExecutor, database)
             runWithTimingDisabled {
-                runBlocking {
-                    executor.job.children.forEach {
-                        it.join()
-                    }
-                }
+                executor.waitForIdle()
                 database.close()
+                context.deleteDatabase(WorkDatabase.DB_NAME)
             }
         }
     }
@@ -117,13 +113,12 @@
             WorkManagerImpl(context, configuration, taskExecutor, database)
             // Prune records for the next run.
             runWithTimingDisabled {
-                runBlocking {
-                    executor.job.children.forEach {
-                        it.join()
-                    }
+                executor.waitForIdle()
+                with(database) {
+                    workSpecDao().pruneFinishedWorkWithZeroDependentsIgnoringKeepForAtLeast()
+                    close()
                 }
-                database.workSpecDao().pruneFinishedWorkWithZeroDependentsIgnoringKeepForAtLeast()
-                database.close()
+                context.deleteDatabase(WorkDatabase.DB_NAME)
             }
         }
     }
diff --git a/work/workmanager/src/main/java/androidx/work/impl/WorkDatabase.java b/work/workmanager/src/main/java/androidx/work/impl/WorkDatabase.java
index f0f35f8..2b8cec2 100644
--- a/work/workmanager/src/main/java/androidx/work/impl/WorkDatabase.java
+++ b/work/workmanager/src/main/java/androidx/work/impl/WorkDatabase.java
@@ -29,6 +29,7 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.RestrictTo;
+import androidx.annotation.VisibleForTesting;
 import androidx.room.Database;
 import androidx.room.Room;
 import androidx.room.RoomDatabase;
@@ -69,7 +70,8 @@
 @TypeConverters(value = {Data.class, WorkTypeConverters.class})
 public abstract class WorkDatabase extends RoomDatabase {
 
-    private static final String DB_NAME = "androidx.work.workdb";
+    @VisibleForTesting
+    public static final String DB_NAME = "androidx.work.workdb";
     // Delete rows in the workspec table that...
     private static final String PRUNE_SQL_FORMAT_PREFIX = "DELETE FROM workspec WHERE "
             // are completed...