Fix reentrant task bugs with VirtualTimeController

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

headless/public/util/virtual_time_controller.cc

@@ -39,8 +39,14 @@ void VirtualTimeController::GrantVirtualTimeBudget(
   virtual_time_policy_ = policy;
 
   // Notify tasks of new budget request.
+  for (TaskEntry& entry : tasks_)
+    entry.ready_to_advance = false;
+
+  iterating_over_tasks_ = true;
   for (TaskEntry& entry : tasks_)
     NotifyTaskBudgetRequested(&entry, budget_ms);
+  iterating_over_tasks_ = false;
+  DeleteTasksIfRequested();
 
   // If there tasks, NotifyTasksAndAdvance is called from TaskReadyToAdvance.
   if (tasks_.empty())
@@ -58,22 +64,39 @@ void VirtualTimeController::NotifyTasksAndAdvance() {
 
   // Give at most as much virtual time as available until the next callback.
   bool ready_to_advance = true;
+  iterating_over_tasks_ = true;
   TimeDelta next_budget = requested_budget_ - accumulated_time_;
+  // TODO(alexclarke): This is a little ugly, find a nicer way.
+  for (TaskEntry& entry : tasks_) {
+    if (entry.next_execution_time <= total_elapsed_time_)
+      entry.ready_to_advance = false;
+  }
+
   for (TaskEntry& entry : tasks_) {
     if (entry.next_execution_time <= total_elapsed_time_)
       NotifyTaskIntervalElapsed(&entry);
 
-    ready_to_advance &= entry.ready_to_advance;
-    next_budget =
-        std::min(next_budget, entry.next_execution_time - total_elapsed_time_);
+    if (tasks_to_delete_.find(entry.task) == tasks_to_delete_.end()) {
+      ready_to_advance &= entry.ready_to_advance;
+      next_budget = std::min(next_budget,
+                             entry.next_execution_time - total_elapsed_time_);
+    }
   }
+  iterating_over_tasks_ = false;
+  DeleteTasksIfRequested();
 
   if (!ready_to_advance)
     return;
 
   if (accumulated_time_ >= requested_budget_) {
+    for (TaskEntry& entry : tasks_)
+      entry.ready_to_advance = false;
+
+    iterating_over_tasks_ = true;
     for (const TaskEntry& entry : tasks_)
       entry.task->BudgetExpired(total_elapsed_time_);
+    iterating_over_tasks_ = false;
+    DeleteTasksIfRequested();
 
     auto callback = budget_expired_callback_;
     budget_expired_callback_.Reset();
@@ -84,8 +107,18 @@ void VirtualTimeController::NotifyTasksAndAdvance() {
   SetVirtualTimePolicy(next_budget);
 }
 
+void VirtualTimeController::DeleteTasksIfRequested() {
+  DCHECK(!iterating_over_tasks_);
+
+  // It's not safe to delete tasks while iterating over |tasks_| so do it now.
+  while (!tasks_to_delete_.empty()) {
+    CancelRepeatingTask(*tasks_to_delete_.begin());
+    tasks_to_delete_.erase(tasks_to_delete_.begin());
+  }
+}
+
 void VirtualTimeController::NotifyTaskIntervalElapsed(TaskEntry* entry) {
-  entry->ready_to_advance = false;
+  DCHECK(!entry->ready_to_advance);
   entry->next_execution_time = total_elapsed_time_ + entry->interval;
 
   entry->task->IntervalElapsed(
@@ -96,7 +129,7 @@ void VirtualTimeController::NotifyTaskIntervalElapsed(TaskEntry* entry) {
 
 void VirtualTimeController::NotifyTaskBudgetRequested(TaskEntry* entry,
                                                       int budget_ms) {
-  entry->ready_to_advance = false;
+  DCHECK(!entry->ready_to_advance);
   entry->task->BudgetRequested(
       total_elapsed_time_, budget_ms,
       base::Bind(&VirtualTimeController::TaskReadyToAdvance,
@@ -158,6 +191,10 @@ void VirtualTimeController::ScheduleRepeatingTask(RepeatingTask* task,
 }
 
 void VirtualTimeController::CancelRepeatingTask(RepeatingTask* task) {
+  if (iterating_over_tasks_) {
+    tasks_to_delete_.insert(task);
+    return;
+  }
   tasks_.remove_if(
       [task](const TaskEntry& entry) { return entry.task == task; });
 }

headless/public/util/virtual_time_controller.h

@@ -78,7 +78,7 @@ class HEADLESS_EXPORT VirtualTimeController
     RepeatingTask* task;
     base::TimeDelta interval;
     base::TimeDelta next_execution_time;
-    bool ready_to_advance = false;
+    bool ready_to_advance = true;
   };
 
   // emulation::Observer implementation:
@@ -90,6 +90,8 @@ class HEADLESS_EXPORT VirtualTimeController
   void NotifyTaskBudgetRequested(TaskEntry* entry, int budget_ms);
   void TaskReadyToAdvance(TaskEntry* entry);
 
+  void DeleteTasksIfRequested();
+
   void SetVirtualTimePolicy(const base::TimeDelta& next_budget);
   void SetVirtualTimePolicyDone(
       std::unique_ptr<emulation::SetVirtualTimePolicyResult>);
@@ -109,7 +111,9 @@ class HEADLESS_EXPORT VirtualTimeController
   base::TimeDelta accumulated_time_;
 
   std::list<TaskEntry> tasks_;
+  std::set<RepeatingTask*> tasks_to_delete_;
   bool in_notify_tasks_and_advance_ = false;
+  bool iterating_over_tasks_ = false;
 };
 
 }  // namespace headless

headless/public/util/virtual_time_controller_test.cc

@@ -16,6 +16,8 @@
 
 namespace headless {
 
+using testing::ElementsAre;
+using testing::Mock;
 using testing::Return;
 using testing::_;
 
@@ -314,4 +316,144 @@ TEST_F(VirtualTimeControllerTest, MultipleTasks) {
   EXPECT_FALSE(budget_expired_);
 }
 
+class VirtualTimeTask : public VirtualTimeController::RepeatingTask {
+ public:
+  using Task = base::Callback<void(const base::TimeDelta& virtual_time)>;
+
+  VirtualTimeTask(VirtualTimeController* controller,
+                  Task budget_requested_task,
+                  Task interval_elapsed_task,
+                  Task budget_expired_task)
+      : controller_(controller),
+        budget_requested_task_(budget_requested_task),
+        interval_elapsed_task_(interval_elapsed_task),
+        budget_expired_task_(budget_expired_task) {}
+
+  void IntervalElapsed(
+      const base::TimeDelta& virtual_time,
+      const base::Callback<void()>& continue_callback) override {
+    interval_elapsed_task_.Run(virtual_time);
+    continue_callback.Run();
+  }
+
+  void BudgetRequested(
+      const base::TimeDelta& virtual_time,
+      int requested_budget_ms,
+      const base::Callback<void()>& continue_callback) override {
+    budget_requested_task_.Run(virtual_time);
+    continue_callback.Run();
+  }
+
+  void BudgetExpired(const base::TimeDelta& virtual_time) override {
+    budget_expired_task_.Run(virtual_time);
+  };
+
+  VirtualTimeController* controller_;  // NOT OWNED
+  Task budget_requested_task_;
+  Task interval_elapsed_task_;
+  Task budget_expired_task_;
+};
+
+TEST_F(VirtualTimeControllerTest, ReentrantTask) {
+#if defined(__clang__)
+  std::vector<std::string> log;
+  VirtualTimeTask task_b(
+      controller_.get(),
+      base::Bind(
+          [](std::vector<std::string>* log,
+             const base::TimeDelta& virtual_time) {
+            log->push_back(base::StringPrintf(
+                "B: budget requested @ %d",
+                static_cast<int>(virtual_time.InMilliseconds())));
+          },
+          &log),
+      base::Bind(
+          [](std::vector<std::string>* log, VirtualTimeController* controller,
+             VirtualTimeTask* task_b, const base::TimeDelta& virtual_time) {
+            log->push_back(base::StringPrintf(
+                "B: interval elapsed @ %d",
+                static_cast<int>(virtual_time.InMilliseconds())));
+            controller->CancelRepeatingTask(task_b);
+          },
+          &log, controller_.get(), &task_b),
+      base::Bind(
+          [](std::vector<std::string>* log,
+             const base::TimeDelta& virtual_time) {
+            log->push_back(base::StringPrintf(
+                "B: budget expired @ %d",
+                static_cast<int>(virtual_time.InMilliseconds())));
+          },
+          &log));
+
+  VirtualTimeTask task_a(
+      controller_.get(),
+      base::Bind(
+          [](std::vector<std::string>* log,
+             const base::TimeDelta& virtual_time) {
+            log->push_back(base::StringPrintf(
+                "A: budget requested @ %d",
+                static_cast<int>(virtual_time.InMilliseconds())));
+          },
+          &log),
+      base::Bind(
+          [](std::vector<std::string>* log, VirtualTimeController* controller,
+             VirtualTimeTask* task_a, VirtualTimeTask* task_b,
+             const base::TimeDelta& virtual_time) {
+            log->push_back(base::StringPrintf(
+                "A: interval elapsed @ %d",
+                static_cast<int>(virtual_time.InMilliseconds())));
+            controller->CancelRepeatingTask(task_a);
+            controller->ScheduleRepeatingTask(task_b, 1500);
+          },
+          &log, controller_.get(), &task_a, &task_b),
+      base::Bind(
+          [](std::vector<std::string>* log,
+             const base::TimeDelta& virtual_time) {
+            log->push_back(base::StringPrintf(
+                "A: budget expired @ %d",
+                static_cast<int>(virtual_time.InMilliseconds())));
+          },
+          &log));
+
+  controller_->ScheduleRepeatingTask(&task_a, 1000);
+
+  EXPECT_CALL(*mock_host_,
+              DispatchProtocolMessage(
+                  &client_,
+                  "{\"id\":0,\"method\":\"Emulation.setVirtualTimePolicy\","
+                  "\"params\":{\"budget\":1000,"
+                  "\"maxVirtualTimeTaskStarvationCount\":0,"
+                  "\"policy\":\"advance\"}}"))
+      .WillOnce(Return(true));
+
+  GrantVirtualTimeBudget(6000);
+  Mock::VerifyAndClearExpectations(&mock_host_);
+
+  EXPECT_CALL(*mock_host_,
+              DispatchProtocolMessage(
+                  &client_,
+                  "{\"id\":2,\"method\":\"Emulation.setVirtualTimePolicy\","
+                  "\"params\":{\"budget\":1500,"
+                  "\"maxVirtualTimeTaskStarvationCount\":0,"
+                  "\"policy\":\"advance\"}}"))
+      .WillOnce(Return(true));
+  SendVirtualTimeBudgetExpiredEvent();
+  Mock::VerifyAndClearExpectations(&mock_host_);
+
+  EXPECT_CALL(*mock_host_,
+              DispatchProtocolMessage(
+                  &client_,
+                  "{\"id\":4,\"method\":\"Emulation.setVirtualTimePolicy\","
+                  "\"params\":{\"budget\":3500,"
+                  "\"maxVirtualTimeTaskStarvationCount\":0,"
+                  "\"policy\":\"advance\"}}"))
+      .WillOnce(Return(true));
+  SendVirtualTimeBudgetExpiredEvent();
+
+  EXPECT_THAT(
+      log, ElementsAre("A: budget requested @ 0", "A: interval elapsed @ 1000",
+                       "B: interval elapsed @ 2500"));
+#endif
+}
+
 }  // namespace headless