Refactor VirtualTimeController to have a more flexible interface

We now support tasks choosing to extend virtual time.

Bug: 777763
Change-Id: I34d0f07ecee78fafe75a6e4d2280c6512c6f7e2f
Reviewed-on: https://chromium-review.googlesource.com/889298
Commit-Queue: Alex Clarke <alexclarke@chromium.org>
Reviewed-by: Eric Seckler <eseckler@chromium.org>
Cr-Commit-Position: refs/heads/master@{#532210}

headless/public/util/compositor_controller.cc

@@ -20,16 +20,20 @@ namespace headless {
 // Sends BeginFrames to advance animations while virtual time advances in
 // intervals.
 class CompositorController::AnimationBeginFrameTask
-    : public VirtualTimeController::RepeatingTask {
+    : public VirtualTimeController::RepeatingTask,
+      public VirtualTimeController::Observer,
+      public VirtualTimeController::StartDeferrer {
  public:
   explicit AnimationBeginFrameTask(CompositorController* compositor_controller)
-      : compositor_controller_(compositor_controller),
+      : RepeatingTask(StartPolicy::START_IMMEDIATELY, -1),
+        compositor_controller_(compositor_controller),
         weak_ptr_factory_(this) {}
 
   // VirtualTimeController::RepeatingTask implementation:
-  void IntervalElapsed(base::TimeDelta virtual_time_offset,
-                       const base::Closure& continue_callback) override {
-    continue_callback_ = continue_callback;
+  void IntervalElapsed(
+      base::TimeDelta virtual_time_offset,
+      base::OnceCallback<void(ContinuePolicy)> continue_callback) override {
+    interval_continue_callback_ = std::move(continue_callback);
 
     // Post a cancellable task that will issue a BeginFrame. This way, we can
     // cancel sending an animation-only BeginFrame if another virtual time task
@@ -43,20 +47,22 @@ class CompositorController::AnimationBeginFrameTask
         FROM_HERE, begin_frame_task_.callback());
   }
 
-  void BudgetRequested(base::TimeDelta virtual_time_offset,
-                       base::TimeDelta requested_budget,
-                       const base::Closure& continue_callback) override {
+  // VirtualTimeController::StartDeferrer implementation:
+  void DeferStart(base::OnceClosure continue_callback) override {
     // Run a BeginFrame if we cancelled it because the budged expired previously
     // and no other BeginFrame was sent while virtual time was paused.
     if (needs_begin_frame_on_virtual_time_resume_) {
-      continue_callback_ = continue_callback;
+      start_continue_callback_ = std::move(continue_callback);
       IssueAnimationBeginFrame();
       return;
     }
-    continue_callback.Run();
+    std::move(continue_callback).Run();
   }
 
-  void BudgetExpired(base::TimeDelta virtual_time_offset) override {
+  // VirtualTimeController::Observer implementation:
+  void VirtualTimeStarted(base::TimeDelta virtual_time_offset) override {}
+
+  void VirtualTimeStopped(base::TimeDelta virtual_time_offset) override {
     // Wait until a new budget was requested before sending another animation
     // BeginFrame, as it's likely that we will send a screenshotting BeginFrame.
     if (!begin_frame_task_.IsCancelled()) {
@@ -96,16 +102,20 @@ class CompositorController::AnimationBeginFrameTask
     TRACE_EVENT0(
         "headless",
         "CompositorController::AnimationBeginFrameTask::BeginFrameComplete");
-    DCHECK(continue_callback_);
-    auto callback = continue_callback_;
-    continue_callback_.Reset();
-    callback.Run();
+    DCHECK(interval_continue_callback_ || start_continue_callback_);
+    if (interval_continue_callback_)
+      std::move(interval_continue_callback_).Run(ContinuePolicy::NOT_REQUIRED);
+
+    if (start_continue_callback_)
+      std::move(start_continue_callback_).Run();
   }
 
   CompositorController* compositor_controller_;  // NOT OWNED
   bool needs_begin_frame_on_virtual_time_resume_ = false;
   base::CancelableClosure begin_frame_task_;
-  base::Closure continue_callback_;
+
+  base::OnceCallback<void(ContinuePolicy)> interval_continue_callback_;
+  base::OnceClosure start_continue_callback_;
   base::WeakPtrFactory<AnimationBeginFrameTask> weak_ptr_factory_;
 };
 
@@ -134,10 +144,14 @@ CompositorController::CompositorController(
       headless_experimental::EnableParams::Builder().Build());
   virtual_time_controller_->ScheduleRepeatingTask(
       animation_task_.get(), animation_begin_frame_interval_);
+  virtual_time_controller_->AddObserver(animation_task_.get());
+  virtual_time_controller_->SetStartDeferrer(animation_task_.get());
 }
 
 CompositorController::~CompositorController() {
+  virtual_time_controller_->RemoveObserver(animation_task_.get());
   virtual_time_controller_->CancelRepeatingTask(animation_task_.get());
+  virtual_time_controller_->SetStartDeferrer(nullptr);
   devtools_client_->GetHeadlessExperimental()
       ->GetExperimental()
       ->RemoveObserver(this);

headless/public/util/compositor_controller_browsertest.cc

@@ -155,16 +155,53 @@ class CompositorControllerBrowserTest
                             base::Unretained(this)));
   }
 
+  class AdditionalVirtualTimeBudget
+      : public VirtualTimeController::RepeatingTask,
+        public VirtualTimeController::Observer {
+   public:
+    AdditionalVirtualTimeBudget(VirtualTimeController* virtual_time_controller,
+                                CompositorControllerBrowserTest* test,
+                                base::TimeDelta budget)
+        : RepeatingTask(StartPolicy::START_IMMEDIATELY, 0),
+          virtual_time_controller_(virtual_time_controller),
+          test_(test) {
+      virtual_time_controller_->ScheduleRepeatingTask(this, budget);
+      virtual_time_controller_->AddObserver(this);
+      virtual_time_controller_->StartVirtualTime();
+    }
+
+    ~AdditionalVirtualTimeBudget() override {
+      virtual_time_controller_->RemoveObserver(this);
+      virtual_time_controller_->CancelRepeatingTask(this);
+    }
+
+    // headless::VirtualTimeController::RepeatingTask implementation:
+    void IntervalElapsed(
+        base::TimeDelta virtual_time,
+        base::OnceCallback<void(ContinuePolicy)> continue_callback) override {
+      std::move(continue_callback).Run(ContinuePolicy::NOT_REQUIRED);
+    }
+
+    // headless::VirtualTimeController::Observer:
+    void VirtualTimeStarted(base::TimeDelta virtual_time_offset) override {}
+
+    void VirtualTimeStopped(base::TimeDelta virtual_time_offset) override {
+      test_->OnVirtualTimeBudgetExpired();
+      delete this;
+    }
+
+   private:
+    headless::VirtualTimeController* const virtual_time_controller_;
+    CompositorControllerBrowserTest* test_;
+  };
+
   void OnRafReady(std::unique_ptr<runtime::EvaluateResult> result) {
     EXPECT_NE(nullptr, result);
     EXPECT_FALSE(result->HasExceptionDetails());
 
-    virtual_time_controller_->GrantVirtualTimeBudget(
-        emulation::VirtualTimePolicy::ADVANCE,
-        kNumFrames * kAnimationFrameInterval, base::Bind([]() {}),
-        base::BindRepeating(
-            &CompositorControllerBrowserTest::OnVirtualTimeBudgetExpired,
-            base::Unretained(this)));
+    // AdditionalVirtualTimeBudget will self delete.
+    new AdditionalVirtualTimeBudget(virtual_time_controller_.get(), this,
+                                    kNumFrames * kAnimationFrameInterval);
   }
 
   void OnVirtualTimeBudgetExpired() {

headless/public/util/virtual_time_controller.h

@@ -23,43 +23,73 @@ class HEADLESS_EXPORT VirtualTimeController
                         int max_task_starvation_count = 0);
   ~VirtualTimeController() override;
 
-  // Grants a |budget| of virtual time by applying the provided |policy|.
-  //
-  // |set_up_complete_callback|, if set, is run after the (initial) policy was
-  // applied via DevTools. |budget_expired_callback| will be called when the
-  // budget has been exhausted.
-  //
-  // Should not be called again until the previous invocation's
-  // budget_expired_callback was executed.
-  virtual void GrantVirtualTimeBudget(
-      emulation::VirtualTimePolicy policy,
-      base::TimeDelta budget,
-      const base::Callback<void()>& set_up_complete_callback,
-      const base::Callback<void()>& budget_expired_callback);
+  // Signals that virtual time should start advancing. If virtual time is
+  // already running, this does nothing.  When virtual time is ready to start
+  // the observers will be notified.
+  virtual void StartVirtualTime();
 
   class RepeatingTask {
    public:
+    // This policy controls whether or not StartVirtualTime() should wait for a
+    // navigation first.
+    enum StartPolicy {
+      WAIT_FOR_NAVIGATION,
+      START_IMMEDIATELY,
+    };
+
+    explicit RepeatingTask(StartPolicy start_policy, int priority)
+        : start_policy_(start_policy), priority_(priority) {}
+
     virtual ~RepeatingTask() {}
 
+    enum class ContinuePolicy {
+      CONTINUE_MORE_TIME_NEEDED,
+      NOT_REQUIRED,
+    };
+
     // Called when the tasks's requested virtual time interval has elapsed.
     // |virtual_time_offset| is the virtual time duration that has advanced
-    // since the page started loading (millisecond granularity). The task should
-    // call |continue_callback| when it is ready for virtual time to continue
-    // advancing.
+    // since the page started loading (millisecond granularity). When the task
+    // has completed it's perioodic work it should call |continue_callback|
+    // with CONTINUE_MORE_TIME_NEEDED if it wants virtual time to continue
+    // advancing, or NOT_REQUIRED otherwise.  Virtual time will continue to
+    // advance until all RepeatingTasks want it to stop.
     virtual void IntervalElapsed(
         base::TimeDelta virtual_time_offset,
-        const base::Callback<void()>& continue_callback) = 0;
+        base::OnceCallback<void(ContinuePolicy policy)> continue_callback) = 0;
 
-    // Called when a new virtual time budget grant was requested. The task
-    // should call |continue_callback| when it is ready for virtual time to
-    // continue advancing.
-    virtual void BudgetRequested(
-        base::TimeDelta virtual_time_offset,
-        base::TimeDelta requested_budget,
-        const base::Callback<void()>& continue_callback) = 0;
+    StartPolicy start_policy() const { return start_policy_; }
+
+    int priority() const { return priority_; }
 
-    // Called when the latest virtual time budget has been used up.
-    virtual void BudgetExpired(base::TimeDelta virtual_time_offset) = 0;
+   private:
+    const StartPolicy start_policy_;
+
+    // If more than one RepeatingTask is scheduled to run at any instant they
+    // are run in order of ascending |priority_|.
+    const int priority_;
+  };
+
+  // An API used by the CompositorController to defer the start of virtual time
+  // until it's ready.
+  class StartDeferrer {
+   public:
+    virtual ~StartDeferrer() {}
+
+    virtual void DeferStart(base::OnceClosure ready_callback) = 0;
+  };
+
+  class Observer {
+   public:
+    virtual ~Observer() {}
+
+    // Called when StartVirtualTime was called. May be delayed by a
+    // StartDeferrer.
+    virtual void VirtualTimeStarted(base::TimeDelta virtual_time_offset) = 0;
+
+    // Called when all RepeatingTasks have either voted for virtual time to stop
+    // advancing, or all have been removed.
+    virtual void VirtualTimeStopped(base::TimeDelta virtual_time_offset) = 0;
   };
 
   // Interleaves execution of the provided |task| with progression of virtual
@@ -72,6 +102,10 @@ class HEADLESS_EXPORT VirtualTimeController
                                      base::TimeDelta interval);
   virtual void CancelRepeatingTask(RepeatingTask* task);
 
+  // Adds an observer which is notified when virtual time starts and stops.
+  virtual void AddObserver(Observer* observer);
+  virtual void RemoveObserver(Observer* observer);
+
   // Returns the time that virtual time offsets are relative to.
   virtual base::Time GetVirtualTimeBase() const;
 
@@ -85,49 +119,59 @@ class HEADLESS_EXPORT VirtualTimeController
     return GetVirtualTimeBase() + GetCurrentVirtualTimeOffset();
   }
 
+  virtual void SetStartDeferrer(StartDeferrer* start_deferrer);
+
  private:
   struct TaskEntry {
-    RepeatingTask* task;
     base::TimeDelta interval;
     base::TimeDelta next_execution_time;
     bool ready_to_advance = true;
+    RepeatingTask::ContinuePolicy continue_policy =
+        RepeatingTask::ContinuePolicy::CONTINUE_MORE_TIME_NEEDED;
   };
 
+  void ObserverReadyToStart();
+
   // emulation::Observer implementation:
   void OnVirtualTimeBudgetExpired(
       const emulation::VirtualTimeBudgetExpiredParams& params) override;
 
   void NotifyTasksAndAdvance();
   void NotifyTaskIntervalElapsed(TaskEntry* entry);
-  void NotifyTaskBudgetRequested(TaskEntry* entry, base::TimeDelta budget);
-  void TaskReadyToAdvance(TaskEntry* entry);
-
-  void DeleteTasksIfRequested();
+  void NotifyTaskVirtualTimeStarted(TaskEntry* entry);
+  void TaskReadyToAdvance(TaskEntry* entry,
+                          RepeatingTask::ContinuePolicy continue_policy);
 
-  void SetVirtualTimePolicy(base::TimeDelta next_budget);
+  void SetVirtualTimePolicy(base::TimeDelta next_budget,
+                            bool wait_for_navigation);
   void SetVirtualTimePolicyDone(
       std::unique_ptr<emulation::SetVirtualTimePolicyResult>);
 
   HeadlessDevToolsClient* const devtools_client_;  // NOT OWNED
+  StartDeferrer* start_deferrer_ = nullptr;        // NOT OWNED
   const int max_task_starvation_count_;
 
-  emulation::VirtualTimePolicy virtual_time_policy_ =
-      emulation::VirtualTimePolicy::ADVANCE;
-  base::Callback<void()> set_up_complete_callback_;
-  base::Callback<void()> budget_expired_callback_;
-
-  bool virtual_time_active_ = false;
   base::TimeDelta total_elapsed_time_offset_;
-  base::TimeDelta requested_budget_;
-  base::TimeDelta last_used_budget_;
-  base::TimeDelta accumulated_budget_portion_;
+  base::TimeDelta last_budget_;
   // Initial virtual time that virtual time offsets are relative to.
   base::Time virtual_time_base_;
 
-  std::list<TaskEntry> tasks_;
-  std::set<RepeatingTask*> tasks_to_delete_;
+  struct RepeatingTaskOrdering {
+    bool operator()(RepeatingTask* a, RepeatingTask* b) const {
+      if (a->priority() == b->priority())
+        return a < b;
+      return a->priority() < b->priority();
+    };
+  };
+
+  std::map<RepeatingTask*, TaskEntry, RepeatingTaskOrdering> tasks_;
+  std::set<Observer*> observers_;
   bool in_notify_tasks_and_advance_ = false;
-  bool iterating_over_tasks_ = false;
+  bool virtual_time_started_ = false;
+  bool virtual_time_paused_ = true;
+  bool should_send_start_notification_ = false;
+
+  base::WeakPtrFactory<VirtualTimeController> weak_ptr_factory_;
 };
 
 }  // namespace headless

headless/test/headless_render_test.cc

@@ -30,9 +30,56 @@ void HeadlessRenderTest::PostRunAsynchronousTest() {
   EXPECT_EQ(FINISHED, state_) << "The test did not finish.";
 }
 
+class HeadlessRenderTest::AdditionalVirtualTimeBudget
+    : public VirtualTimeController::RepeatingTask,
+      public VirtualTimeController::Observer {
+ public:
+  AdditionalVirtualTimeBudget(VirtualTimeController* virtual_time_controller,
+                              HeadlessRenderTest* test,
+                              base::RunLoop* run_loop,
+                              int budget_ms)
+      : RepeatingTask(StartPolicy::WAIT_FOR_NAVIGATION, 0),
+        virtual_time_controller_(virtual_time_controller),
+        test_(test),
+        run_loop_(run_loop) {
+    virtual_time_controller_->ScheduleRepeatingTask(
+        this, base::TimeDelta::FromMilliseconds(budget_ms));
+    virtual_time_controller_->AddObserver(this);
+    virtual_time_controller_->StartVirtualTime();
+  }
+
+  ~AdditionalVirtualTimeBudget() override {
+    virtual_time_controller_->RemoveObserver(this);
+    virtual_time_controller_->CancelRepeatingTask(this);
+  }
+
+  // headless::VirtualTimeController::RepeatingTask implementation:
+  void IntervalElapsed(
+      base::TimeDelta virtual_time,
+      base::OnceCallback<void(ContinuePolicy)> continue_callback) override {
+    std::move(continue_callback).Run(ContinuePolicy::NOT_REQUIRED);
+  }
+
+  // headless::VirtualTimeController::Observer:
+  void VirtualTimeStarted(base::TimeDelta virtual_time_offset) override {
+    run_loop_->Quit();
+  }
+
+  void VirtualTimeStopped(base::TimeDelta virtual_time_offset) override {
+    test_->HandleVirtualTimeExhausted();
+    delete this;
+  }
+
+ private:
+  headless::VirtualTimeController* const virtual_time_controller_;
+  HeadlessRenderTest* test_;
+  base::RunLoop* run_loop_;
+};
+
 void HeadlessRenderTest::RunDevTooledTest() {
   http_handler_->SetHeadlessBrowserContext(browser_context_);
 
+  // TODO(alexclarke): Use the compositor controller here too.
   virtual_time_controller_ =
       std::make_unique<VirtualTimeController>(devtools_client_.get());
 
@@ -64,6 +111,16 @@ void HeadlessRenderTest::RunDevTooledTest() {
     run_loop.Run();
   }
 
+  {
+    base::RunLoop run_loop;
+    // Note AdditionalVirtualTimeBudget will self delete.
+    new AdditionalVirtualTimeBudget(virtual_time_controller_.get(), this,
+                                    &run_loop, 5000);
+    base::MessageLoop::ScopedNestableTaskAllower nest_loop(
+        base::MessageLoop::current());
+    run_loop.Run();
+  }
+
   state_ = STARTING;
   devtools_client_->GetPage()->Navigate(url.spec());
   browser()->BrowserMainThread()->PostDelayedTask(
@@ -124,11 +181,6 @@ void HeadlessRenderTest::OnFrameStartedLoading(
   if (state_ == STARTING) {
     state_ = LOADING;
     main_frame_ = params.GetFrameId();
-    virtual_time_controller_->GrantVirtualTimeBudget(
-        emulation::VirtualTimePolicy::PAUSE_IF_NETWORK_FETCHES_PENDING,
-        base::TimeDelta::FromMilliseconds(5000), base::Closure(),
-        base::Bind(&HeadlessRenderTest::HandleVirtualTimeExhausted,
-                   weak_ptr_factory_.GetWeakPtr()));
   }
 
   auto it = unconfirmed_frame_redirects_.find(params.GetFrameId());

headless/test/headless_render_test.h

@@ -118,6 +118,8 @@ class HeadlessRenderTest : public HeadlessAsyncDevTooledBrowserTest,
   std::vector<std::string> js_exceptions_;
 
  private:
+  class AdditionalVirtualTimeBudget;
+
   void HandleVirtualTimeExhausted();
   void OnGetDomSnapshotDone(
       std::unique_ptr<dom_snapshot::GetSnapshotResult> result);