TaskScheduler: Add FlushForTesting().

base/task_scheduler/task_tracker.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/task_scheduler/task_tracker.h"
 
 #include <limits>
 
-#include "base/atomicops.h"
 #include "base/callback.h"
 #include "base/debug/task_annotator.h"
 #include "base/logging.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/sequence_token.h"
+#include "base/synchronization/condition_variable.h"
 #include "base/threading/sequenced_task_runner_handle.h"
 #include "base/threading/thread_restrictions.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/trace_event/trace_event.h"
 
 namespace base {
 namespace internal {
 
 namespace {
 
@@ skipping 96 matching lines @@
   // modify-write) semantics however. For example, if two threads are racing to
   // call IncrementNumTasksBlockingShutdown() and StartShutdown() respectively,
   // either the first thread will win and the StartShutdown() call will see the
   // blocking task or the second thread will win and
   // IncrementNumTasksBlockingShutdown() will know that shutdown has started.
   subtle::Atomic32 bits_ = 0;
 
   DISALLOW_COPY_AND_ASSIGN(State);
 };
 
-TaskTracker::TaskTracker() : state_(new State) {}
+TaskTracker::TaskTracker()
+    : state_(new State),
+      flush_for_testing_cv_(flush_for_testing_lock_.CreateConditionVariable()),
+      shutdown_lock_(&flush_for_testing_lock_) {}
 TaskTracker::~TaskTracker() = default;
 
 void TaskTracker::Shutdown() {
+  PerformShutdown();
+  DCHECK(IsShutdownComplete());
+
+  // Unblock FlushForTesting() when shutdown completes.
+  AutoSchedulerLock auto_lock(flush_for_testing_lock_);
+  flush_for_testing_cv_->Signal();
+}
+
+void TaskTracker::FlushForTesting() {
+  AutoSchedulerLock auto_lock(flush_for_testing_lock_);
+  while (subtle::NoBarrier_Load(&num_pending_undelayed_tasks_) != 0 &&
+         !IsShutdownComplete()) {
+    flush_for_testing_cv_->Wait();
+  }
+}
+
+bool TaskTracker::WillPostTask(const Task* task) {
+  DCHECK(task);
+
+  if (!BeforePostTask(task->traits.shutdown_behavior()))
+    return false;
+
+  if (task->delayed_run_time.is_null())
+    subtle::NoBarrier_AtomicIncrement(&num_pending_undelayed_tasks_, 1);
+
+  debug::TaskAnnotator task_annotator;
+  task_annotator.DidQueueTask(kQueueFunctionName, *task);
+
+  return true;
+}
+
+bool TaskTracker::RunTask(const Task* task,
+                          const SequenceToken& sequence_token) {
+  DCHECK(task);
+  DCHECK(sequence_token.IsValid());
+
+  const TaskShutdownBehavior shutdown_behavior =
+      task->traits.shutdown_behavior();
+  const bool can_run_task = BeforeRunTask(shutdown_behavior);
+
+  if (can_run_task) {
+    // All tasks run through here and the scheduler itself doesn't use
+    // singletons. Therefore, it isn't necessary to reset the singleton allowed
+    // bit after running the task.
+    ThreadRestrictions::SetSingletonAllowed(
+        task->traits.shutdown_behavior() !=
+        TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN);
+
+    {
+      // Set up SequenceToken as expected for the scope of the task.
+      ScopedSetSequenceTokenForCurrentThread
+          scoped_set_sequence_token_for_current_thread(sequence_token);
+
+      // Set up TaskRunnerHandle as expected for the scope of the task.
+      std::unique_ptr<SequencedTaskRunnerHandle> sequenced_task_runner_handle;
+      std::unique_ptr<ThreadTaskRunnerHandle> single_thread_task_runner_handle;
+      DCHECK(!task->sequenced_task_runner_ref ||
+             !task->single_thread_task_runner_ref);
+      if (task->sequenced_task_runner_ref) {
+        sequenced_task_runner_handle.reset(
+            new SequencedTaskRunnerHandle(task->sequenced_task_runner_ref));
+      } else if (task->single_thread_task_runner_ref) {
+        single_thread_task_runner_handle.reset(
+            new ThreadTaskRunnerHandle(task->single_thread_task_runner_ref));
+      }
+
+      TRACE_TASK_EXECUTION(kRunFunctionName, *task);
+
+      debug::TaskAnnotator task_annotator;
+      task_annotator.RunTask(kQueueFunctionName, *task);
+    }
+
+    AfterRunTask(shutdown_behavior);
+  }
+
+  if (task->delayed_run_time.is_null())
+    DecrementNumPendingUndelayedTasks();
+
+  return can_run_task;
+}
+
+bool TaskTracker::HasShutdownStarted() const {
+  return state_->HasShutdownStarted();
+}
+
+bool TaskTracker::IsShutdownComplete() const {
+  AutoSchedulerLock auto_lock(shutdown_lock_);
+  return shutdown_event_ && shutdown_event_->IsSignaled();
+}
+
+void TaskTracker::SetHasShutdownStartedForTesting() {
+  state_->StartShutdown();
+}
+
+void TaskTracker::PerformShutdown() {
   {
     AutoSchedulerLock auto_lock(shutdown_lock_);
 
     // This method can only be called once.
     DCHECK(!shutdown_event_);
     DCHECK(!num_block_shutdown_tasks_posted_during_shutdown_);
     DCHECK(!state_->HasShutdownStarted());
 
     shutdown_event_.reset(
         new WaitableEvent(WaitableEvent::ResetPolicy::MANUAL,
@@ skipping 30 matching lines @@
     // posted during shutdown. Otherwise, the histogram has already been
     // recorded in BeforePostTask().
     if (num_block_shutdown_tasks_posted_during_shutdown_ <
         kMaxBlockShutdownTasksPostedDuringShutdown) {
       RecordNumBlockShutdownTasksPostedDuringShutdown(
           num_block_shutdown_tasks_posted_during_shutdown_);
     }
   }
 }
 
-bool TaskTracker::WillPostTask(const Task* task) {
-  DCHECK(task);
-
-  if (!BeforePostTask(task->traits.shutdown_behavior()))
-    return false;
-
-  debug::TaskAnnotator task_annotator;
-  task_annotator.DidQueueTask(kQueueFunctionName, *task);
-
-  return true;
-}
-
-bool TaskTracker::RunTask(const Task* task,
-                          const SequenceToken& sequence_token) {
-  DCHECK(task);
-  DCHECK(sequence_token.IsValid());
-
-  const TaskShutdownBehavior shutdown_behavior =
-      task->traits.shutdown_behavior();
-  if (!BeforeRunTask(shutdown_behavior))
-    return false;
-
-  // All tasks run through here and the scheduler itself doesn't use singletons.
-  // Therefore, it isn't necessary to reset the singleton allowed bit after
-  // running the task.
-  ThreadRestrictions::SetSingletonAllowed(
-      task->traits.shutdown_behavior() !=
-      TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN);
-
-  {
-    // Set up SequenceToken as expected for the scope of the task.
-    ScopedSetSequenceTokenForCurrentThread
-        scoped_set_sequence_token_for_current_thread(sequence_token);
-
-    // Set up TaskRunnerHandle as expected for the scope of the task.
-    std::unique_ptr<SequencedTaskRunnerHandle> sequenced_task_runner_handle;
-    std::unique_ptr<ThreadTaskRunnerHandle> single_thread_task_runner_handle;
-    DCHECK(!task->sequenced_task_runner_ref ||
-           !task->single_thread_task_runner_ref);
-    if (task->sequenced_task_runner_ref) {
-      sequenced_task_runner_handle.reset(
-          new SequencedTaskRunnerHandle(task->sequenced_task_runner_ref));
-    } else if (task->single_thread_task_runner_ref) {
-      single_thread_task_runner_handle.reset(
-          new ThreadTaskRunnerHandle(task->single_thread_task_runner_ref));
-    }
-
-    TRACE_TASK_EXECUTION(kRunFunctionName, *task);
-
-    debug::TaskAnnotator task_annotator;
-    task_annotator.RunTask(kQueueFunctionName, *task);
-  }
-
-  AfterRunTask(shutdown_behavior);
-
-  return true;
-}
-
-bool TaskTracker::HasShutdownStarted() const {
-  return state_->HasShutdownStarted();
-}
-
-bool TaskTracker::IsShutdownComplete() const {
-  AutoSchedulerLock auto_lock(shutdown_lock_);
-  return shutdown_event_ && shutdown_event_->IsSignaled();
-}
-
-void TaskTracker::SetHasShutdownStartedForTesting() {
-  state_->StartShutdown();
-}
-
 bool TaskTracker::BeforePostTask(TaskShutdownBehavior shutdown_behavior) {
   if (shutdown_behavior == TaskShutdownBehavior::BLOCK_SHUTDOWN) {
     // BLOCK_SHUTDOWN tasks block shutdown between the moment they are posted
     // and the moment they complete their execution.
     const bool shutdown_started = state_->IncrementNumTasksBlockingShutdown();
 
     if (shutdown_started) {
       AutoSchedulerLock auto_lock(shutdown_lock_);
 
       // A BLOCK_SHUTDOWN task posted after shutdown has completed is an
@@ skipping 79 matching lines @@
 void TaskTracker::OnBlockingShutdownTasksComplete() {
   AutoSchedulerLock auto_lock(shutdown_lock_);
 
   // This method can only be called after shutdown has started.
   DCHECK(state_->HasShutdownStarted());
   DCHECK(shutdown_event_);
 
   shutdown_event_->Signal();
 }
 
+void TaskTracker::DecrementNumPendingUndelayedTasks() {
+  const auto new_num_pending_undelayed_tasks =
+      subtle::NoBarrier_AtomicIncrement(&num_pending_undelayed_tasks_, -1);
+  DCHECK_GE(new_num_pending_undelayed_tasks, 0);
+  if (new_num_pending_undelayed_tasks == 0) {
+    AutoSchedulerLock auto_lock(flush_for_testing_lock_);
+    flush_for_testing_cv_->Signal();
+  }
+}
+
 }  // namespace internal
 }  // namespace base