Add TaskScheduler::FlushAsyncForTesting()

This performs a task flush similar to TaskScheduler::FlushForTesting()
except it is async. It calls the callback if TaskScheduler::Flush()
would have returned.

BUG=804930

Change-Id: I533ab5a2fbaccd79a220aad3a535fa37ed7df823
Reviewed-on: https://chromium-review.googlesource.com/883607
Commit-Queue: Robert Liao <robliao@chromium.org>
Reviewed-by: Gabriel Charette <gab@chromium.org>
Cr-Commit-Position: refs/heads/master@{#532258}

base/task_scheduler/task_scheduler.h

@@ -148,6 +148,13 @@ class BASE_EXPORT TaskScheduler {
   // other threads during the call. Returns immediately when shutdown completes.
   virtual void FlushForTesting() = 0;
 
+  // Returns and calls |flush_callback| when there are no incomplete undelayed
+  // tasks. |flush_callback| may be called back on any thread and should not
+  // perform a lot of work. May be used when additional work on the current
+  // thread needs to be performed during a flush. Only one
+  // FlushAsyncForTesting() may be pending at any given time.
+  virtual void FlushAsyncForTesting(OnceClosure flush_callback) = 0;
+
   // Joins all threads. Tasks that are already running are allowed to complete
   // their execution. This can only be called once. Using this task scheduler
   // instance to create task runners or post tasks is not permitted during or

base/task_scheduler/task_scheduler_impl.cc

@@ -189,6 +189,10 @@ void TaskSchedulerImpl::FlushForTesting() {
   task_tracker_->FlushForTesting();
 }
 
+void TaskSchedulerImpl::FlushAsyncForTesting(OnceClosure flush_callback) {
+  task_tracker_->FlushAsyncForTesting(std::move(flush_callback));
+}
+
 void TaskSchedulerImpl::JoinForTesting() {
 #if DCHECK_IS_ON()
   DCHECK(!join_for_testing_returned_.IsSet());

base/task_scheduler/task_scheduler_impl.h

@@ -83,6 +83,7 @@ class BASE_EXPORT TaskSchedulerImpl : public TaskScheduler {
       const TaskTraits& traits) const override;
   void Shutdown() override;
   void FlushForTesting() override;
+  void FlushAsyncForTesting(OnceClosure flush_callback) override;
   void JoinForTesting() override;
 
  private:

base/task_scheduler/task_scheduler_impl_unittest.cc

@@ -118,7 +118,9 @@ void VerifyTimeAndTaskEnvironmentAndSignalEvent(const TaskTraits& traits,
 scoped_refptr<TaskRunner> CreateTaskRunnerWithTraitsAndExecutionMode(
     TaskScheduler* scheduler,
     const TaskTraits& traits,
-    test::ExecutionMode execution_mode) {
+    test::ExecutionMode execution_mode,
+    SingleThreadTaskRunnerThreadMode default_single_thread_task_runner_mode =
+        SingleThreadTaskRunnerThreadMode::SHARED) {
   switch (execution_mode) {
     case test::ExecutionMode::PARALLEL:
       return scheduler->CreateTaskRunnerWithTraits(traits);
@@ -126,7 +128,7 @@ scoped_refptr<TaskRunner> CreateTaskRunnerWithTraitsAndExecutionMode(
       return scheduler->CreateSequencedTaskRunnerWithTraits(traits);
     case test::ExecutionMode::SINGLE_THREADED: {
       return scheduler->CreateSingleThreadTaskRunnerWithTraits(
-          traits, SingleThreadTaskRunnerThreadMode::SHARED);
+          traits, default_single_thread_task_runner_mode);
     }
   }
   ADD_FAILURE() << "Unknown ExecutionMode";
@@ -393,6 +395,32 @@ TEST_P(TaskSchedulerImplTest, AllTasksAreUserBlocking) {
   task_running.Wait();
 }
 
+// Verifies that FlushAsyncForTesting() calls back correctly for all trait and
+// execution mode pairs.
+TEST_P(TaskSchedulerImplTest, FlushAsyncForTestingSimple) {
+  StartTaskScheduler();
+
+  WaitableEvent unblock_task(WaitableEvent::ResetPolicy::MANUAL,
+                             WaitableEvent::InitialState::NOT_SIGNALED);
+  CreateTaskRunnerWithTraitsAndExecutionMode(
+      &scheduler_,
+      TaskTraits::Override(GetParam().traits, {WithBaseSyncPrimitives()}),
+      GetParam().execution_mode, SingleThreadTaskRunnerThreadMode::DEDICATED)
+      ->PostTask(FROM_HERE,
+                 BindOnce(&WaitableEvent::Wait, Unretained(&unblock_task)));
+
+  WaitableEvent flush_event(WaitableEvent::ResetPolicy::MANUAL,
+                            WaitableEvent::InitialState::NOT_SIGNALED);
+  scheduler_.FlushAsyncForTesting(
+      BindOnce(&WaitableEvent::Signal, Unretained(&flush_event)));
+  PlatformThread::Sleep(TestTimeouts::tiny_timeout());
+  EXPECT_FALSE(flush_event.IsSignaled());
+
+  unblock_task.Signal();
+
+  flush_event.Wait();
+}
+
 INSTANTIATE_TEST_CASE_P(OneTraitsExecutionModePair,
                         TaskSchedulerImplTest,
                         ::testing::ValuesIn(GetTraitsExecutionModePairs()));
@@ -604,5 +632,14 @@ TEST_F(TaskSchedulerImplTest, SequenceLocalStorage) {
   scheduler_.FlushForTesting();
 }
 
+TEST_F(TaskSchedulerImplTest, FlushAsyncNoTasks) {
+  StartTaskScheduler();
+  bool called_back = false;
+  scheduler_.FlushAsyncForTesting(
+      BindOnce([](bool* called_back) { *called_back = true; },
+               Unretained(&called_back)));
+  EXPECT_TRUE(called_back);
+}
+
 }  // namespace internal
 }  // namespace base