Fix bug with TaskQueueSelector and blocked queues

components/scheduler/base/task_queue_selector.cc

 // Copyright 2014 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 "components/scheduler/base/task_queue_selector.h"
 
 #include "base/logging.h"
 #include "base/trace_event/trace_event_argument.h"
 #include "components/scheduler/base/task_queue_impl.h"
 #include "components/scheduler/base/work_queue.h"
 
 namespace scheduler {
 namespace internal {
 
 TaskQueueSelector::TaskQueueSelector()
-    : enabled_selector_(this),
-      blocked_selector_(this),
+    : enabled_selector_(this, "enabled"),
+      blocked_selector_(this, "blocked"),
       immediate_starvation_count_(0),
       high_priority_starvation_count_(0),
       num_blocked_queues_to_report_(0),
       task_queue_selector_observer_(nullptr) {}
 
 TaskQueueSelector::~TaskQueueSelector() {}
 
 void TaskQueueSelector::AddQueue(internal::TaskQueueImpl* queue) {
   DCHECK(main_thread_checker_.CalledOnValidThread());
   DCHECK(queue->IsQueueEnabled());
-  SetQueuePriority(queue, TaskQueue::NORMAL_PRIORITY);
+  enabled_selector_.AddQueue(queue, TaskQueue::NORMAL_PRIORITY);
 }
 
 void TaskQueueSelector::RemoveQueue(internal::TaskQueueImpl* queue) {
   DCHECK(main_thread_checker_.CalledOnValidThread());
   if (queue->IsQueueEnabled()) {
     enabled_selector_.RemoveQueue(queue);
+// The #if DCHECK_IS_ON() shouldn't be necessary but this doesn't compile on
+// chromeos bots without it :(
+#if DCHECK_IS_ON()
+    DCHECK(!blocked_selector_.CheckContainsQueueForTest(queue));
+#endif
   } else if (queue->should_report_when_execution_blocked()) {
     DCHECK_GT(num_blocked_queues_to_report_, 0u);
     num_blocked_queues_to_report_--;
     blocked_selector_.RemoveQueue(queue);
+#if DCHECK_IS_ON()
+    DCHECK(!enabled_selector_.CheckContainsQueueForTest(queue));
+#endif
   }
 }
 
 void TaskQueueSelector::EnableQueue(internal::TaskQueueImpl* queue) {
   DCHECK(main_thread_checker_.CalledOnValidThread());
   DCHECK(queue->IsQueueEnabled());
   if (queue->should_report_when_execution_blocked()) {
     DCHECK_GT(num_blocked_queues_to_report_, 0u);
     num_blocked_queues_to_report_--;
+    blocked_selector_.RemoveQueue(queue);
   }
-  blocked_selector_.RemoveQueue(queue);
-  enabled_selector_.AssignQueueToSet(queue, queue->GetQueuePriority());
+  enabled_selector_.AddQueue(queue, queue->GetQueuePriority());
   if (task_queue_selector_observer_)
     task_queue_selector_observer_->OnTaskQueueEnabled(queue);
 }
 
 void TaskQueueSelector::DisableQueue(internal::TaskQueueImpl* queue) {
   DCHECK(main_thread_checker_.CalledOnValidThread());
   DCHECK(!queue->IsQueueEnabled());
   enabled_selector_.RemoveQueue(queue);
-  blocked_selector_.AssignQueueToSet(queue, queue->GetQueuePriority());
-  if (queue->should_report_when_execution_blocked())
+  if (queue->should_report_when_execution_blocked()) {
+    blocked_selector_.AddQueue(queue, queue->GetQueuePriority());
     num_blocked_queues_to_report_++;
+  }
 }
 
 void TaskQueueSelector::SetQueuePriority(internal::TaskQueueImpl* queue,
                                          TaskQueue::QueuePriority priority) {
   DCHECK_LT(priority, TaskQueue::QUEUE_PRIORITY_COUNT);
   DCHECK(main_thread_checker_.CalledOnValidThread());
   if (queue->IsQueueEnabled()) {
-    enabled_selector_.AssignQueueToSet(queue, priority);
+    enabled_selector_.ChangeSetIndex(queue, priority);
+  } else if (queue->should_report_when_execution_blocked()) {
+    blocked_selector_.ChangeSetIndex(queue, priority);
   } else {
-    blocked_selector_.AssignQueueToSet(queue, priority);
+    // Normally blocked_selector_.ChangeSetIndex would assign the queue's
+    // priority, however if |queue->should_report_when_execution_blocked()| is
+    // false then the disabled queue is not in any set so we need to do it here.
+    queue->delayed_work_queue()->AssignSetIndex(priority);
+    queue->immediate_work_queue()->AssignSetIndex(priority);
   }
   DCHECK_EQ(priority, queue->GetQueuePriority());
 }
 
 TaskQueue::QueuePriority TaskQueueSelector::NextPriority(
     TaskQueue::QueuePriority priority) {
   DCHECK(priority < TaskQueue::QUEUE_PRIORITY_COUNT);
   return static_cast<TaskQueue::QueuePriority>(static_cast<int>(priority) + 1);
 }
 
 TaskQueueSelector::PrioritizingSelector::PrioritizingSelector(
-    TaskQueueSelector* task_queue_selector)
+    TaskQueueSelector* task_queue_selector,
+    const char* name)
     : task_queue_selector_(task_queue_selector),
-      delayed_work_queue_sets_(TaskQueue::QUEUE_PRIORITY_COUNT),
-      immediate_work_queue_sets_(TaskQueue::QUEUE_PRIORITY_COUNT) {}
+      delayed_work_queue_sets_(TaskQueue::QUEUE_PRIORITY_COUNT, name),
+      immediate_work_queue_sets_(TaskQueue::QUEUE_PRIORITY_COUNT, name) {}
 
-void TaskQueueSelector::PrioritizingSelector::AssignQueueToSet(
+void TaskQueueSelector::PrioritizingSelector::AddQueue(
     internal::TaskQueueImpl* queue,
     TaskQueue::QueuePriority priority) {
-  delayed_work_queue_sets_.AssignQueueToSet(queue->delayed_work_queue(),
+#if DCHECK_IS_ON()
+  DCHECK(!CheckContainsQueueForTest(queue));
+#endif
+  delayed_work_queue_sets_.AddQueue(queue->delayed_work_queue(), priority);
+  immediate_work_queue_sets_.AddQueue(queue->immediate_work_queue(), priority);
+#if DCHECK_IS_ON()
+  DCHECK(CheckContainsQueueForTest(queue));
+#endif
+}
+
+void TaskQueueSelector::PrioritizingSelector::ChangeSetIndex(
+    internal::TaskQueueImpl* queue,
+    TaskQueue::QueuePriority priority) {
+#if DCHECK_IS_ON()
+  DCHECK(CheckContainsQueueForTest(queue));
+#endif
+  delayed_work_queue_sets_.ChangeSetIndex(queue->delayed_work_queue(),
+                                          priority);
+  immediate_work_queue_sets_.ChangeSetIndex(queue->immediate_work_queue(),
                                             priority);
-  immediate_work_queue_sets_.AssignQueueToSet(queue->immediate_work_queue(),
-                                              priority);
-  // The #if DCHECK_IS_ON() shouldn't be necessary but this doesn't compile on
-  // chromeos bots without it :(
 #if DCHECK_IS_ON()
-  DCHECK_EQ(queue->delayed_work_queue()->Empty(),
-            !delayed_work_queue_sets_.ContainsWorkQueueForTest(
-                queue->delayed_work_queue()));
-  DCHECK_EQ(queue->immediate_work_queue()->Empty(),
-            !immediate_work_queue_sets_.ContainsWorkQueueForTest(
-                queue->immediate_work_queue()));
+  DCHECK(CheckContainsQueueForTest(queue));
 #endif
 }
 
 void TaskQueueSelector::PrioritizingSelector::RemoveQueue(
     internal::TaskQueueImpl* queue) {
 #if DCHECK_IS_ON()
-  DCHECK_EQ(queue->delayed_work_queue()->Empty(),
-            !delayed_work_queue_sets_.ContainsWorkQueueForTest(
-                queue->delayed_work_queue()))
-      << " Did you try to RemoveQueue twice? Thats not supported!";
-  DCHECK_EQ(queue->immediate_work_queue()->Empty(),
-            !immediate_work_queue_sets_.ContainsWorkQueueForTest(
-                queue->immediate_work_queue()))
-      << " Did you try to RemoveQueue twice? Thats not supported!";
+  DCHECK(CheckContainsQueueForTest(queue));
 #endif
   delayed_work_queue_sets_.RemoveQueue(queue->delayed_work_queue());
   immediate_work_queue_sets_.RemoveQueue(queue->immediate_work_queue());
 
 #if DCHECK_IS_ON()
-  DCHECK(!delayed_work_queue_sets_.ContainsWorkQueueForTest(
-      queue->delayed_work_queue()));
-  DCHECK(!immediate_work_queue_sets_.ContainsWorkQueueForTest(
-      queue->delayed_work_queue()));
+  DCHECK(!CheckContainsQueueForTest(queue));
 #endif
 }
 
 bool TaskQueueSelector::PrioritizingSelector::
     ChooseOldestImmediateTaskWithPriority(TaskQueue::QueuePriority priority,
                                           WorkQueue** out_work_queue) const {
   return immediate_work_queue_sets_.GetOldestQueueInSet(priority,
                                                         out_work_queue);
 }
 
@@ skipping 76 matching lines @@
   for (TaskQueue::QueuePriority priority = TaskQueue::HIGH_PRIORITY;
        priority < max_priority; priority = NextPriority(priority)) {
     if (ChooseOldestWithPriority(priority, out_chose_delayed_over_immediate,
                                  out_work_queue)) {
       return true;
     }
   }
   return false;
 }
 
+#if DCHECK_IS_ON() || !defined(NDEBUG)
+bool
+TaskQueueSelector::PrioritizingSelector::CheckContainsQueueForTest(
+    const internal::TaskQueueImpl* queue) const {
+  bool contains_delayed_work_queue =
+      delayed_work_queue_sets_.ContainsWorkQueueForTest(
+          queue->delayed_work_queue());
+
+  bool contains_immediate_work_queue =
+      immediate_work_queue_sets_.ContainsWorkQueueForTest(
+          queue->immediate_work_queue());
+
+  DCHECK_EQ(contains_delayed_work_queue, contains_immediate_work_queue);
+  return contains_delayed_work_queue;
+}
+#endif
+
 bool TaskQueueSelector::SelectWorkQueueToService(WorkQueue** out_work_queue) {
   DCHECK(main_thread_checker_.CalledOnValidThread());
   bool chose_delayed_over_immediate = false;
   bool found_queue = enabled_selector_.SelectWorkQueueToService(
       TaskQueue::QUEUE_PRIORITY_COUNT, out_work_queue,
       &chose_delayed_over_immediate);
   if (!found_queue) {
     TrySelectingBlockedQueue();
     return false;
   }
 
   TrySelectingBlockedQueueOverEnabledQueue(**out_work_queue);
   DidSelectQueueWithPriority(
       (*out_work_queue)->task_queue()->GetQueuePriority(),
       chose_delayed_over_immediate);
   return true;
 }
 
 void TaskQueueSelector::TrySelectingBlockedQueue() {
+  DCHECK(main_thread_checker_.CalledOnValidThread());
   if (!num_blocked_queues_to_report_ || !task_queue_selector_observer_)
     return;
   WorkQueue* chosen_blocked_queue;
   bool chose_delayed_over_immediate = false;
   // There was nothing unblocked to run, see if we could have run a blocked
   // task.
   if (blocked_selector_.SelectWorkQueueToService(
           TaskQueue::QUEUE_PRIORITY_COUNT, &chosen_blocked_queue,
           &chose_delayed_over_immediate)) {
     task_queue_selector_observer_->OnTriedToSelectBlockedWorkQueue(
         chosen_blocked_queue);
   }
 }
 
 void TaskQueueSelector::TrySelectingBlockedQueueOverEnabledQueue(
     const WorkQueue& chosen_enabled_queue) {
+  DCHECK(main_thread_checker_.CalledOnValidThread());
   if (!num_blocked_queues_to_report_ || !task_queue_selector_observer_)
     return;
 
   TaskQueue::QueuePriority max_priority =
       NextPriority(chosen_enabled_queue.task_queue()->GetQueuePriority());
 
   WorkQueue* chosen_blocked_queue;
   bool chose_delayed_over_immediate = false;
   bool found_queue = blocked_selector_.SelectWorkQueueToService(
       max_priority, &chosen_blocked_queue, &chose_delayed_over_immediate);
@@ skipping 50 matching lines @@
   state->SetInteger("immediate_starvation_count", immediate_starvation_count_);
   state->SetInteger("num_blocked_queues_to_report",
                     num_blocked_queues_to_report_);
 }
 
 void TaskQueueSelector::SetTaskQueueSelectorObserver(Observer* observer) {
   task_queue_selector_observer_ = observer;
 }
 
 bool TaskQueueSelector::EnabledWorkQueuesEmpty() const {
+  DCHECK(main_thread_checker_.CalledOnValidThread());
   for (TaskQueue::QueuePriority priority = TaskQueue::CONTROL_PRIORITY;
        priority < TaskQueue::QUEUE_PRIORITY_COUNT;
        priority = NextPriority(priority)) {
     if (!enabled_selector_.delayed_work_queue_sets()->IsSetEmpty(priority) ||
         !enabled_selector_.immediate_work_queue_sets()->IsSetEmpty(priority)) {
       return false;
     }
   }
   return true;
 }
 
 void TaskQueueSelector::SetImmediateStarvationCountForTest(
     size_t immediate_starvation_count) {
   immediate_starvation_count_ = immediate_starvation_count;
 }
 
 }  // namespace internal
 }  // namespace scheduler