[ScopedTaskEnvironment] Add MOCK_TIME support for ThreadPool tasks

Multi-threaded MOCK_TIME advances when either RunLoop::Run() or
ScopedTaskEnvironment::FastForward*() is active and the system reaches
an idle phase (both main thread and thread pool are out of immediate
tasks). It then advances by the minimum time delta necessary to make a
task able to run.

In a follow-up CL, I'll split MOCK_TIME off of MainThreadType (now that
it's no longer main thread specific) and migrate callers.

Bug: 946657
Change-Id: Id3da34ac066623291a8fca682fd13c9473e83e0b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1686776Reviewed-by: Alex Clarke <alexclarke@chromium.org>
Reviewed-by: François Doray <fdoray@chromium.org>
Commit-Queue: Gabriel Charette <gab@chromium.org>
Cr-Commit-Position: refs/heads/master@{#676272}

base/task/sequence_manager/time_domain.h

@@ -47,10 +47,10 @@ class BASE_EXPORT TimeDomain {
   // TODO(alexclarke): Make this main thread only.
   virtual TimeTicks Now() const = 0;
 
-  // Computes the delay until the time when TimeDomain needs to wake up
-  // some TaskQueue. Specific time domains (e.g. virtual or throttled) may
-  // return TimeDelata() if TaskQueues have any delayed tasks they deem
-  // eligible to run. It's also allowed to advance time domains's internal
+  // Computes the delay until the time when TimeDomain needs to wake up some
+  // TaskQueue on the main thread. Specific time domains (e.g. virtual or
+  // throttled) may return TimeDelta() if TaskQueues have any delayed tasks they
+  // deem eligible to run. It's also allowed to advance time domains's internal
   // clock when this method is called.
   // Can be called from main thread only.
   // NOTE: |lazy_now| and the return value are in the SequenceManager's time.

base/task/thread_pool/delayed_task_manager.cc

@@ -117,6 +117,13 @@ void DelayedTaskManager::ProcessRipeTasks() {
   }
 }
 
+Optional<TimeTicks> DelayedTaskManager::NextScheduledRunTime() const {
+  CheckedAutoLock auto_lock(queue_lock_);
+  if (delayed_task_queue_.empty())
+    return nullopt;
+  return delayed_task_queue_.Min().task.delayed_run_time;
+}
+
 TimeTicks DelayedTaskManager::GetTimeToScheduleProcessRipeTasksLockRequired() {
   queue_lock_.AssertAcquired();
   if (delayed_task_queue_.empty())

base/task/thread_pool/delayed_task_manager.h

@@ -13,6 +13,7 @@
 #include "base/macros.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/ref_counted.h"
+#include "base/optional.h"
 #include "base/synchronization/atomic_flag.h"
 #include "base/task/common/checked_lock.h"
 #include "base/task/common/intrusive_heap.h"
@@ -49,6 +50,12 @@ class BASE_EXPORT DelayedTaskManager {
                       PostTaskNowCallback post_task_now_callback,
                       scoped_refptr<TaskRunner> task_runner);
 
+  // Pop and post all the ripe tasks in the delayed task queue.
+  void ProcessRipeTasks();
+
+  // Returns the |delayed_run_time| of the next scheduled task, if any.
+  Optional<TimeTicks> NextScheduledRunTime() const;
+
  private:
   struct DelayedTask {
     DelayedTask();
@@ -86,9 +93,6 @@ class BASE_EXPORT DelayedTaskManager {
     DISALLOW_COPY_AND_ASSIGN(DelayedTask);
   };
 
-  // Pop and post all the ripe tasks in the delayed task queue.
-  void ProcessRipeTasks();
-
   // Get the time at which to schedule the next |ProcessRipeTasks()| execution,
   // or TimeTicks::Max() if none needs to be scheduled (i.e. no task, or next
   // task already scheduled).
@@ -108,7 +112,7 @@ class BASE_EXPORT DelayedTaskManager {
   // it is never modified. It is therefore safe to access
   // |service_thread_task_runner_| without synchronization once it is observed
   // that it is non-null.
-  CheckedLock queue_lock_;
+  mutable CheckedLock queue_lock_;
 
   scoped_refptr<TaskRunner> service_thread_task_runner_;
 

base/task/thread_pool/task_tracker.h

@@ -160,6 +160,11 @@ class BASE_EXPORT TaskTracker {
     return tracked_ref_factory_.GetTrackedRef();
   }
 
+  // Returns true if there are task sources that haven't completed their
+  // execution (still queued or in progress). If it returns false: the side-
+  // effects of all completed tasks are guaranteed to be visible to the caller.
+  bool HasIncompleteTaskSourcesForTesting() const;
+
  protected:
   // Runs and deletes |task| if |can_run_task| is true. Otherwise, just deletes
   // |task|. |task| is always deleted in the environment where it runs or would
@@ -172,11 +177,6 @@ class BASE_EXPORT TaskTracker {
                              const TaskTraits& traits,
                              bool can_run_task);
 
-  // Returns true if there are task sources that haven't completed their
-  // execution (still queued or in progress). If it returns false: the side-
-  // effects of all completed tasks are guaranteed to be visible to the caller.
-  bool HasIncompleteTaskSourcesForTesting() const;
-
  private:
   friend class RegisteredTaskSource;
   class State;

base/task/thread_pool/thread_pool_impl.cc

@@ -64,11 +64,13 @@ bool HasDisableBestEffortTasksSwitch() {
 
 ThreadPoolImpl::ThreadPoolImpl(StringPiece histogram_label)
     : ThreadPoolImpl(histogram_label,
-                     std::make_unique<TaskTrackerImpl>(histogram_label)) {}
+                     std::make_unique<TaskTrackerImpl>(histogram_label),
+                     DefaultTickClock::GetInstance()) {}
 
 ThreadPoolImpl::ThreadPoolImpl(StringPiece histogram_label,
-                               std::unique_ptr<TaskTrackerImpl> task_tracker)
-    : thread_pool_clock_(DefaultTickClock::GetInstance()),
+                               std::unique_ptr<TaskTrackerImpl> task_tracker,
+                               const TickClock* tick_clock)
+    : thread_pool_clock_(tick_clock),
       task_tracker_(std::move(task_tracker)),
       service_thread_(std::make_unique<ServiceThread>(
           task_tracker_.get(),
@@ -265,6 +267,16 @@ ThreadPoolImpl::CreateUpdateableSequencedTaskRunner(const TaskTraits& traits) {
   return MakeRefCounted<PooledSequencedTaskRunner>(new_traits, this);
 }
 
+Optional<TimeTicks> ThreadPoolImpl::NextScheduledRunTimeForTesting() const {
+  if (task_tracker_->HasIncompleteTaskSourcesForTesting())
+    return ThreadPoolClock::Now();
+  return delayed_task_manager_.NextScheduledRunTime();
+}
+
+void ThreadPoolImpl::ProcessRipeDelayedTasksForTesting() {
+  delayed_task_manager_.ProcessRipeTasks();
+}
+
 int ThreadPoolImpl::GetMaxConcurrentNonBlockedTasksWithTraitsDeprecated(
     const TaskTraits& traits) const {
   // This method does not support getting the maximum number of BEST_EFFORT

base/task/thread_pool/thread_pool_impl.h

@@ -14,6 +14,7 @@
 #include "base/macros.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/ref_counted.h"
+#include "base/optional.h"
 #include "base/sequence_checker.h"
 #include "base/strings/string_piece.h"
 #include "base/synchronization/atomic_flag.h"
@@ -64,9 +65,11 @@ class BASE_EXPORT ThreadPoolImpl : public ThreadPoolInstance,
   //|histogram_label| is used to label histograms, it must not be empty.
   explicit ThreadPoolImpl(StringPiece histogram_label);
 
-  // For testing only. Creates a ThreadPoolImpl with a custom TaskTracker.
+  // For testing only. Creates a ThreadPoolImpl with a custom TaskTracker and
+  // TickClock.
   ThreadPoolImpl(StringPiece histogram_label,
-                 std::unique_ptr<TaskTrackerImpl> task_tracker);
+                 std::unique_ptr<TaskTrackerImpl> task_tracker,
+                 const TickClock* tick_clock);
 
   ~ThreadPoolImpl() override;
 
@@ -101,6 +104,16 @@ class BASE_EXPORT ThreadPoolImpl : public ThreadPoolInstance,
   scoped_refptr<UpdateableSequencedTaskRunner>
   CreateUpdateableSequencedTaskRunner(const TaskTraits& traits);
 
+  // Returns the TimeTicks of the next task scheduled on ThreadPool (Now() if
+  // immediate, nullopt if none). This is thread-safe, i.e., it's safe if tasks
+  // are being posted in parallel with this call but such a situation obviously
+  // results in a race as to whether this call will see the new tasks in time.
+  Optional<TimeTicks> NextScheduledRunTimeForTesting() const;
+
+  // Forces ripe delayed tasks to be posted (e.g. when time is mocked and
+  // advances faster than the real-time delay on ServiceThread).
+  void ProcessRipeDelayedTasksForTesting();
+
  private:
   // Invoked after |has_fence_| or |has_best_effort_fence_| is updated. Sets the
   // CanRunPolicy in TaskTracker and wakes up workers as appropriate.

base/test/scoped_task_environment.cc

@@ -19,6 +19,7 @@
 #include "base/task/sequence_manager/sequence_manager_impl.h"
 #include "base/task/sequence_manager/time_domain.h"
 #include "base/task/thread_pool/thread_pool.h"
+#include "base/task/thread_pool/thread_pool_clock.h"
 #include "base/task/thread_pool/thread_pool_impl.h"
 #include "base/test/bind_test_util.h"
 #include "base/test/test_mock_time_task_runner.h"
@@ -29,6 +30,7 @@
 #include "base/threading/thread_restrictions.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/time/clock.h"
+#include "base/time/default_tick_clock.h"
 #include "base/time/tick_clock.h"
 #include "base/time/time.h"
 #include "base/time/time_override.h"
@@ -112,6 +114,9 @@ class ScopedTaskEnvironment::TestTaskTracker
   // thread, waiting for it to complete results in a deadlock...).
   bool DisallowRunTasks();
 
+  // Returns true if tasks are currently allowed to run.
+  bool TasksAllowedToRun() const;
+
  private:
   friend class ScopedTaskEnvironment;
 
@@ -122,7 +127,7 @@ class ScopedTaskEnvironment::TestTaskTracker
                      bool can_run_task) override;
 
   // Synchronizes accesses to members below.
-  Lock lock_;
+  mutable Lock lock_;
 
   // True if running tasks is allowed.
   bool can_run_tasks_ GUARDED_BY(lock_) = true;
@@ -192,6 +197,14 @@ class ScopedTaskEnvironment::MockTimeDomain
     return nullptr;
   }
 
+  void SetThreadPool(internal::ThreadPoolImpl* thread_pool,
+                     const TestTaskTracker* thread_pool_task_tracker) {
+    DCHECK(!thread_pool_);
+    DCHECK(!thread_pool_task_tracker_);
+    thread_pool_ = thread_pool;
+    thread_pool_task_tracker_ = thread_pool_task_tracker;
+  }
+
   // sequence_manager::TimeDomain:
 
   sequence_manager::LazyNow CreateLazyNow() const override {
@@ -239,7 +252,8 @@ class ScopedTaskEnvironment::MockTimeDomain
     if (!auto_advance_on_idle_)
       return false;
 
-    return FastForwardToNextTaskOrCap(TimeTicks::Max());
+    return FastForwardToNextTaskOrCap(TimeTicks::Max()) ==
+           NextTaskSource::kMainThread;
   }
 
   const char* GetName() const override { return "MockTimeDomain"; }
@@ -247,20 +261,73 @@ class ScopedTaskEnvironment::MockTimeDomain
   // TickClock implementation:
   TimeTicks NowTicks() const override { return Now(); }
 
-  // Advances time to the next task or to |fast_forward_cap| (if it's not
-  // Max()). Returns true if there's additional immediate work as a result
-  // of this call.
-  bool FastForwardToNextTaskOrCap(TimeTicks fast_forward_cap) {
+  // Used by FastForwardToNextTaskOrCap() to return which task source time was
+  // advanced to.
+  enum class NextTaskSource {
+    // Out of tasks under |fast_forward_cap|.
+    kNone,
+    // There's now >=1 immediate task on the main thread.
+    kMainThread,
+    // There's now >=1 immediate task in the thread pool.
+    kThreadPool,
+  };
+
+  // Advances time to the first of : next main thread task, next thread pool
+  // task, or |fast_forward_cap| (if it's not Max()).
+  NextTaskSource FastForwardToNextTaskOrCap(TimeTicks fast_forward_cap) {
     DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
 
     // We don't need to call ReclaimMemory here because
     // DelayTillNextTask will have dealt with cancelled delayed tasks for us.
-    Optional<TimeTicks> next_task_time = NextScheduledRunTime();
+    Optional<TimeTicks> next_main_thread_task_time = NextScheduledRunTime();
+
+    // Consider the next thread pool tasks iff they're running.
+    Optional<TimeTicks> next_thread_pool_task_time;
+    if (thread_pool_ && thread_pool_task_tracker_->TasksAllowedToRun()) {
+      next_thread_pool_task_time =
+          thread_pool_->NextScheduledRunTimeForTesting();
+    }
+
+    // Custom comparison logic to consider nullopt the largest rather than
+    // smallest value. Could consider using TimeTicks::Max() instead of nullopt
+    // to represent out-of-tasks?
+    Optional<TimeTicks> next_task_time;
+    if (!next_main_thread_task_time) {
+      next_task_time = next_thread_pool_task_time;
+    } else if (!next_thread_pool_task_time) {
+      next_task_time = next_main_thread_task_time;
+    } else {
+      next_task_time =
+          std::min(*next_main_thread_task_time, *next_thread_pool_task_time);
+    }
 
     if (next_task_time && *next_task_time <= fast_forward_cap) {
-      AutoLock lock(now_ticks_lock_);
-      now_ticks_ = *next_task_time;
-      return true;
+      {
+        AutoLock lock(now_ticks_lock_);
+        // It's possible for |next_task_time| to be in the past in the following
+        // scenario:
+        // Start with Now() == 100ms
+        // Thread A : Post 200ms delayed task T (construct and enqueue)
+        // Thread B : Construct 20ms delayed task U
+        //              => |delayed_run_time| == 120ms.
+        // Thread A : FastForwardToNextTaskOrCap() => fast-forwards to T @
+        //            300ms (task U is not yet in queue).
+        // Thread B : Complete enqueue of task U.
+        // Thread A : FastForwardToNextTaskOrCap() => must stay at 300ms and run
+        //            U, not go back to 120ms.
+        // Hence we need std::max() to protect again this because construction
+        // and enqueuing isn't atomic in time (LazyNow support in
+        // base/task/thread_pool could help).
+        now_ticks_ = std::max(now_ticks_, *next_task_time);
+      }
+
+      if (next_task_time == next_thread_pool_task_time) {
+        // Let the thread pool know that it should post its now ripe delayed
+        // tasks.
+        thread_pool_->ProcessRipeDelayedTasksForTesting();
+        return NextTaskSource::kThreadPool;
+      }
+      return NextTaskSource::kMainThread;
     }
 
     if (!fast_forward_cap.is_max()) {
@@ -270,11 +337,15 @@ class ScopedTaskEnvironment::MockTimeDomain
       now_ticks_ = std::max(now_ticks_, fast_forward_cap);
     }
 
-    return false;
+    return NextTaskSource::kNone;
   }
 
-  void set_auto_advance_on_idle(bool auto_advance_on_idle) {
+  // Sets |auto_advance_on_idle_| to |auto_advance_on_idle| and returns its
+  // previous value.
+  bool SetAutoAdvanceOnIdle(bool auto_advance_on_idle) {
+    const auto previous = auto_advance_on_idle_;
     auto_advance_on_idle_ = auto_advance_on_idle;
+    return previous;
   }
 
  private:
@@ -287,6 +358,9 @@ class ScopedTaskEnvironment::MockTimeDomain
 
   sequence_manager::SequenceManager* const sequence_manager_;
 
+  internal::ThreadPoolImpl* thread_pool_ = nullptr;
+  const TestTaskTracker* thread_pool_task_tracker_ = nullptr;
+
   std::unique_ptr<subtle::ScopedTimeClockOverrides> time_overrides_;
 
   // Protects |now_ticks_|
@@ -391,8 +465,14 @@ void ScopedTaskEnvironment::InitializeThreadPool() {
 
   auto task_tracker = std::make_unique<TestTaskTracker>();
   task_tracker_ = task_tracker.get();
-  ThreadPoolInstance::Set(std::make_unique<internal::ThreadPoolImpl>(
-      "ScopedTaskEnvironment", std::move(task_tracker)));
+  const TickClock* tick_clock =
+      mock_time_domain_ ? static_cast<TickClock*>(mock_time_domain_.get())
+                        : DefaultTickClock::GetInstance();
+  auto thread_pool = std::make_unique<internal::ThreadPoolImpl>(
+      "ScopedTaskEnvironment", std::move(task_tracker), tick_clock);
+  if (mock_time_domain_)
+    mock_time_domain_->SetThreadPool(thread_pool.get(), task_tracker_);
+  ThreadPoolInstance::Set(std::move(thread_pool));
   ThreadPoolInstance::Get()->Start(init_params);
 }
 
@@ -571,14 +651,21 @@ void ScopedTaskEnvironment::RunUntilIdle() {
 
 void ScopedTaskEnvironment::FastForwardBy(TimeDelta delta) {
   DCHECK(mock_time_domain_);
-  mock_time_domain_->set_auto_advance_on_idle(false);
+  DCHECK_GE(delta, TimeDelta());
+
+  const bool was_auto_advancing = mock_time_domain_->SetAutoAdvanceOnIdle(false);
+  const bool could_run_tasks = task_tracker_->AllowRunTasks();
 
   const TimeTicks fast_forward_until = mock_time_domain_->NowTicks() + delta;
   do {
     RunUntilIdle();
-  } while (mock_time_domain_->FastForwardToNextTaskOrCap(fast_forward_until));
+  } while (mock_time_domain_->FastForwardToNextTaskOrCap(fast_forward_until) !=
+           MockTimeDomain::NextTaskSource::kNone);
 
-  mock_time_domain_->set_auto_advance_on_idle(true);
+  if (was_auto_advancing)
+    mock_time_domain_->SetAutoAdvanceOnIdle(true);
+  if (!could_run_tasks)
+    task_tracker_->DisallowRunTasks();
 }
 
 void ScopedTaskEnvironment::FastForwardUntilNoTasksRemain() {
@@ -640,6 +727,11 @@ bool ScopedTaskEnvironment::TestTaskTracker::AllowRunTasks() {
   return could_run_tasks;
 }
 
+bool ScopedTaskEnvironment::TestTaskTracker::TasksAllowedToRun() const {
+  AutoLock auto_lock(lock_);
+  return can_run_tasks_;
+}
+
 bool ScopedTaskEnvironment::TestTaskTracker::DisallowRunTasks() {
   AutoLock auto_lock(lock_);
 

base/test/scoped_task_environment.h

@@ -78,16 +78,12 @@ class ScopedTaskEnvironment {
     DEFAULT,
     // The main thread doesn't pump system messages and uses a mock clock for
     // delayed tasks (controllable via FastForward*() methods).
-    // TODO(gab): Make this the default |main_thread_type|.
-    // TODO(gab): Also mock the ThreadPoolInstance's clock simultaneously (this
-    // currently only mocks the main thread's clock).
+    // TODO(gab): Make MOCK_TIME configurable independent of MainThreadType.
     MOCK_TIME,
     // The main thread pumps UI messages.
     UI,
     // The main thread pumps UI messages and uses a mock clock for delayed tasks
     // (controllable via FastForward*() methods).
-    // TODO(gab@): Enable mock time on all threads and make MOCK_TIME
-    // configurable independent of MainThreadType.
     UI_MOCK_TIME,
     // The main thread pumps asynchronous IO messages and supports the
     // FileDescriptorWatcher API on POSIX.
@@ -195,9 +191,9 @@ class ScopedTaskEnvironment {
   void RunUntilIdle();
 
   // Only valid for instances with a MOCK_TIME MainThreadType. Fast-forwards
-  // virtual time by |delta|, causing all tasks on the main thread with a
-  // remaining delay less than or equal to |delta| to be executed in their
-  // natural order before this returns. |delta| must be non-negative. Upon
+  // virtual time by |delta|, causing all tasks on the main thread and thread
+  // pool with a remaining delay less than or equal to |delta| to be executed in
+  // their natural order before this returns. |delta| must be non-negative. Upon
   // returning from this method, NowTicks() will be >= the initial |NowTicks() +
   // delta|. It is guaranteed to be == iff tasks executed in this
   // FastForwardBy() didn't result in nested calls to time-advancing-methods.
@@ -229,13 +225,13 @@ class ScopedTaskEnvironment {
   //                  FastForwardBy(2ms)
   //                  ================================
   //   Result: NowTicks() is 5ms further in the future.
-  //
-  // TODO(gab): Make this apply to ThreadPool delayed tasks as well
-  // (currently only main thread time is mocked).
   void FastForwardBy(TimeDelta delta);
 
   // Only valid for instances with a MOCK_TIME MainThreadType.
   // Short for FastForwardBy(TimeDelta::Max()).
+  //
+  // WARNING: This has the same caveat as RunUntilIdle() and is even more likely
+  // to spin forever (any RepeatingTimer will cause this).
   void FastForwardUntilNoTasksRemain();
 
   // Only valid for instances with a MOCK_TIME MainThreadType. Returns a