Fix bug with TaskQueueSelector and blocked queues

The TaskQueueSelector is only supposed to touch the blocked_selector_
iff queue->should_report_when_execution_blocked() is true. Unfortunately
TaskQueueSelector::EnableQueue unconditionally added queues to the
blocked_selector_ leading to a potential UAF.

BUG=581973, 584544, 582712, 585744

Review URL: https://codereview.chromium.org/1685093002

Cr-Commit-Position: refs/heads/master@{#374692}

10 files changed, 191 insertions, 102 deletions

diff --git a/components/scheduler/base/task_queue_impl.cc b/components/scheduler/base/task_queue_impl.cc
index 5f0b4a1..5ce23f1 100644
--- a/components/scheduler/base/task_queue_impl.cc
+++ b/components/scheduler/base/task_queue_impl.cc
@@ -127,8 +127,8 @@ void TaskQueueImpl::UnregisterTaskQueue() {
   main_thread_only().task_queue_manager = nullptr;
   main_thread_only().delayed_incoming_queue = std::priority_queue<Task>();
   any_thread().immediate_incoming_queue = std::queue<Task>();
-  main_thread_only().immediate_work_queue->Clear();
-  main_thread_only().delayed_work_queue->Clear();
+  main_thread_only().immediate_work_queue.reset();
+  main_thread_only().delayed_work_queue.reset();
 }
 
 bool TaskQueueImpl::RunsTasksOnCurrentThread() const {
@@ -488,7 +488,7 @@ const char* TaskQueueImpl::GetName() const {
 }
 
 void TaskQueueImpl::SetQueuePriority(QueuePriority priority) {
-  if (!main_thread_only().task_queue_manager)
+  if (!main_thread_only().task_queue_manager || priority == GetQueuePriority())
     return;
   main_thread_only().task_queue_manager->selector_.SetQueuePriority(this,
                                                                     priority);
diff --git a/components/scheduler/base/task_queue_manager.cc b/components/scheduler/base/task_queue_manager.cc
index 4e31a86..009d4ff 100644
--- a/components/scheduler/base/task_queue_manager.cc
+++ b/components/scheduler/base/task_queue_manager.cc
@@ -195,6 +195,8 @@ void TaskQueueManager::DoWork(base::TimeTicks run_time, bool from_main_thread) {
       break;
     }
 
+    bool should_trigger_wakeup = work_queue->task_queue()->wakeup_policy() ==
+                                 TaskQueue::WakeupPolicy::CAN_WAKE_OTHER_QUEUES;
     switch (ProcessTaskFromWorkQueue(work_queue, &previous_task)) {
       case ProcessTaskResult::DEFERRED:
         // If a task was deferred, try again with another task. Note that this
@@ -206,8 +208,8 @@ void TaskQueueManager::DoWork(base::TimeTicks run_time, bool from_main_thread) {
       case ProcessTaskResult::TASK_QUEUE_MANAGER_DELETED:
         return;  // The TaskQueueManager got deleted, we must bail out.
     }
-    bool should_trigger_wakeup = work_queue->task_queue()->wakeup_policy() ==
-                                 TaskQueue::WakeupPolicy::CAN_WAKE_OTHER_QUEUES;
+    work_queue = nullptr; // The queue may have been unregistered.
+
     UpdateWorkQueues(should_trigger_wakeup, &previous_task);
 
     // Only run a single task per batch in nested run loops so that we can
diff --git a/components/scheduler/base/task_queue_selector.cc b/components/scheduler/base/task_queue_selector.cc
index 3120232..9db8ab9 100644
--- a/components/scheduler/base/task_queue_selector.cc
+++ b/components/scheduler/base/task_queue_selector.cc
@@ -13,8 +13,8 @@ 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),
@@ -25,17 +25,25 @@ 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
   }
 }
 
@@ -45,9 +53,9 @@ void TaskQueueSelector::EnableQueue(internal::TaskQueueImpl* queue) {
   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);
 }
@@ -56,9 +64,10 @@ 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,
@@ -66,9 +75,15 @@ void TaskQueueSelector::SetQueuePriority(internal::TaskQueueImpl* queue,
   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());
 }
@@ -80,50 +95,50 @@ TaskQueue::QueuePriority TaskQueueSelector::NextPriority(
 }
 
 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::AddQueue(
+    internal::TaskQueueImpl* queue,
+    TaskQueue::QueuePriority priority) {
+#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::AssignQueueToSet(
+void TaskQueueSelector::PrioritizingSelector::ChangeSetIndex(
     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_.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
 }
 
@@ -220,6 +235,23 @@ bool TaskQueueSelector::PrioritizingSelector::SelectWorkQueueToService(
   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;
@@ -239,6 +271,7 @@ bool TaskQueueSelector::SelectWorkQueueToService(WorkQueue** out_work_queue) {
 }
 
 void TaskQueueSelector::TrySelectingBlockedQueue() {
+  DCHECK(main_thread_checker_.CalledOnValidThread());
   if (!num_blocked_queues_to_report_ || !task_queue_selector_observer_)
     return;
   WorkQueue* chosen_blocked_queue;
@@ -255,6 +288,7 @@ void TaskQueueSelector::TrySelectingBlockedQueue() {
 
 void TaskQueueSelector::TrySelectingBlockedQueueOverEnabledQueue(
     const WorkQueue& chosen_enabled_queue) {
+  DCHECK(main_thread_checker_.CalledOnValidThread());
   if (!num_blocked_queues_to_report_ || !task_queue_selector_observer_)
     return;
 
@@ -325,6 +359,7 @@ void TaskQueueSelector::SetTaskQueueSelectorObserver(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)) {
diff --git a/components/scheduler/base/task_queue_selector.h b/components/scheduler/base/task_queue_selector.h
index 80ec7c0..4a82adc 100644
--- a/components/scheduler/base/task_queue_selector.h
+++ b/components/scheduler/base/task_queue_selector.h
@@ -85,10 +85,13 @@ class SCHEDULER_EXPORT TaskQueueSelector {
  protected:
   class SCHEDULER_EXPORT PrioritizingSelector {
    public:
-    PrioritizingSelector(TaskQueueSelector* task_queue_selector);
+    PrioritizingSelector(TaskQueueSelector* task_queue_selector,
+                         const char* name);
 
-    void AssignQueueToSet(internal::TaskQueueImpl* queue,
-                          TaskQueue::QueuePriority priority);
+    void ChangeSetIndex(internal::TaskQueueImpl* queue,
+                        TaskQueue::QueuePriority priority);
+    void AddQueue(internal::TaskQueueImpl* queue,
+                  TaskQueue::QueuePriority priority);
     void RemoveQueue(internal::TaskQueueImpl* queue);
 
     bool SelectWorkQueueToService(TaskQueue::QueuePriority max_priority,
@@ -113,6 +116,10 @@ class SCHEDULER_EXPORT TaskQueueSelector {
                                   bool* out_chose_delayed_over_immediate,
                                   WorkQueue** out_work_queue) const;
 
+#if DCHECK_IS_ON() || !defined(NDEBUG)
+    bool CheckContainsQueueForTest(const internal::TaskQueueImpl* queue) const;
+#endif
+
    private:
     bool ChooseOldestImmediateTaskWithPriority(
         TaskQueue::QueuePriority priority,
diff --git a/components/scheduler/base/task_queue_selector_unittest.cc b/components/scheduler/base/task_queue_selector_unittest.cc
index 755e97e0..3d48ba9 100644
--- a/components/scheduler/base/task_queue_selector_unittest.cc
+++ b/components/scheduler/base/task_queue_selector_unittest.cc
@@ -126,6 +126,10 @@ class TaskQueueSelectorTest : public testing::Test {
   void TearDown() final {
     for (scoped_refptr<TaskQueueImpl>& task_queue : task_queues_) {
       task_queue->UnregisterTaskQueue();
+      // Note since this test doesn't have a TaskQueueManager we need to
+      // manually remove |task_queue| from the |selector_|.  Normally
+      // UnregisterTaskQueue would do that.
+      selector_.RemoveQueue(task_queue.get());
     }
   }
 
@@ -187,7 +191,7 @@ TEST_F(TaskQueueSelectorTest, TestObserverWithEnabledQueue) {
 
 TEST_F(TaskQueueSelectorTest,
        TestObserverWithSetQueuePriorityAndQueueAlreadyEnabled) {
-  selector_.SetQueuePriority(task_queues_[1].get(), TaskQueue::NORMAL_PRIORITY);
+  selector_.SetQueuePriority(task_queues_[1].get(), TaskQueue::HIGH_PRIORITY);
   MockObserver mock_observer;
   selector_.SetTaskQueueSelectorObserver(&mock_observer);
   EXPECT_CALL(mock_observer, OnTaskQueueEnabled(_)).Times(0);
@@ -294,7 +298,7 @@ TEST_F(TaskQueueSelectorTest, TestBestEffortGetsStarved) {
   PushTasks(queue_order, 2);
   selector_.SetQueuePriority(task_queues_[0].get(),
                              TaskQueue::BEST_EFFORT_PRIORITY);
-  selector_.SetQueuePriority(task_queues_[1].get(), TaskQueue::NORMAL_PRIORITY);
+  EXPECT_EQ(TaskQueue::NORMAL_PRIORITY, task_queues_[1]->GetQueuePriority());
   WorkQueue* chosen_work_queue = nullptr;
   for (int i = 0; i < 100; i++) {
     EXPECT_TRUE(selector_.SelectWorkQueueToService(&chosen_work_queue));
@@ -390,6 +394,7 @@ TEST_F(TaskQueueSelectorTest, TestObserverWithOneBlockedQueue) {
   EXPECT_FALSE(selector.SelectWorkQueueToService(&chosen_work_queue));
 
   task_queue->UnregisterTaskQueue();
+  selector.RemoveQueue(task_queue.get());
 }
 
 TEST_F(TaskQueueSelectorTest, TestObserverWithTwoBlockedQueues) {
@@ -422,12 +427,13 @@ TEST_F(TaskQueueSelectorTest, TestObserverWithTwoBlockedQueues) {
 
   // Removing the second queue and selecting again should result in another
   // notification.
+  task_queue->UnregisterTaskQueue();
   selector.RemoveQueue(task_queue.get());
   EXPECT_CALL(mock_observer, OnTriedToSelectBlockedWorkQueue(_)).Times(1);
   EXPECT_FALSE(selector.SelectWorkQueueToService(&chosen_work_queue));
 
-  task_queue->UnregisterTaskQueue();
   task_queue2->UnregisterTaskQueue();
+  selector.RemoveQueue(task_queue2.get());
 }
 
 struct ChooseOldestWithPriorityTestParam {
diff --git a/components/scheduler/base/work_queue.cc b/components/scheduler/base/work_queue.cc
index 6051a49..338c0b1 100644
--- a/components/scheduler/base/work_queue.cc
+++ b/components/scheduler/base/work_queue.cc
@@ -23,10 +23,9 @@ void WorkQueue::AsValueInto(base::trace_event::TracedValue* state) const {
   }
 }
 
-WorkQueue::~WorkQueue() {}
-
-void WorkQueue::Clear() {
-  work_queue_ = std::queue<TaskQueueImpl::Task>();
+WorkQueue::~WorkQueue() {
+  DCHECK(!work_queue_sets_) << task_queue_ ->GetName() << " : "
+      << work_queue_sets_->name() << " : " << name_;
 }
 
 const TaskQueueImpl::Task* WorkQueue::GetFrontTask() const {
@@ -43,21 +42,19 @@ bool WorkQueue::GetFrontTaskEnqueueOrder(EnqueueOrder* enqueue_order) const {
 }
 
 void WorkQueue::Push(TaskQueueImpl::Task&& task) {
-  DCHECK(work_queue_sets_);
   bool was_empty = work_queue_.empty();
   work_queue_.push(task);
-  if (was_empty)
+  if (was_empty && work_queue_sets_)
     work_queue_sets_->OnPushQueue(this);
 }
 
 void WorkQueue::PushAndSetEnqueueOrder(TaskQueueImpl::Task&& task,
                                        EnqueueOrder enqueue_order) {
-  DCHECK(work_queue_sets_);
   bool was_empty = work_queue_.empty();
   work_queue_.push(task);
   work_queue_.back().set_enqueue_order(enqueue_order);
 
-  if (was_empty)
+  if (was_empty && work_queue_sets_)
     work_queue_sets_->OnPushQueue(this);
 }
 
@@ -66,10 +63,9 @@ void WorkQueue::PopTaskForTest() {
 }
 
 void WorkQueue::SwapLocked(std::queue<TaskQueueImpl::Task>& incoming_queue) {
-  DCHECK(work_queue_sets_);
   std::swap(work_queue_, incoming_queue);
 
-  if (!work_queue_.empty())
+  if (!work_queue_.empty() && work_queue_sets_)
     work_queue_sets_->OnPushQueue(this);
   task_queue_->TraceQueueSize(true);
 }
@@ -84,9 +80,11 @@ TaskQueueImpl::Task WorkQueue::TakeTaskFromWorkQueue() {
   return pending_task;
 }
 
-void WorkQueue::AssignToWorkQueueSets(WorkQueueSets* work_queue_sets,
-                                      size_t work_queue_set_index) {
+void WorkQueue::AssignToWorkQueueSets(WorkQueueSets* work_queue_sets) {
   work_queue_sets_ = work_queue_sets;
+}
+
+void WorkQueue::AssignSetIndex(size_t work_queue_set_index) {
   work_queue_set_index_ = work_queue_set_index;
 }
 
diff --git a/components/scheduler/base/work_queue.h b/components/scheduler/base/work_queue.h
index 17a7a51..2384502 100644
--- a/components/scheduler/base/work_queue.h
+++ b/components/scheduler/base/work_queue.h
@@ -26,8 +26,10 @@ class SCHEDULER_EXPORT WorkQueue {
 
   // Associates this work queue with the given work queue sets. This must be
   // called before any tasks can be inserted into this work queue.
-  void AssignToWorkQueueSets(WorkQueueSets* work_queue_sets,
-                             size_t work_queue_set_index);
+  void AssignToWorkQueueSets(WorkQueueSets* work_queue_sets);
+
+  // Assigns the current set index.
+  void AssignSetIndex(size_t work_queue_set_index);
 
   void AsValueInto(base::trace_event::TracedValue* state) const;
 
diff --git a/components/scheduler/base/work_queue_sets.cc b/components/scheduler/base/work_queue_sets.cc
index f952334..96c5da4 100644
--- a/components/scheduler/base/work_queue_sets.cc
+++ b/components/scheduler/base/work_queue_sets.cc
@@ -10,15 +10,29 @@
 namespace scheduler {
 namespace internal {
 
-WorkQueueSets::WorkQueueSets(size_t num_sets)
-    : enqueue_order_to_work_queue_maps_(num_sets) {}
+WorkQueueSets::WorkQueueSets(size_t num_sets, const char* name)
+    : enqueue_order_to_work_queue_maps_(num_sets), name_(name) {}
 
 WorkQueueSets::~WorkQueueSets() {}
 
+void WorkQueueSets::AddQueue(WorkQueue* work_queue, size_t set_index) {
+  DCHECK(!work_queue->work_queue_sets());
+  DCHECK_LT(set_index, enqueue_order_to_work_queue_maps_.size());
+  EnqueueOrder enqueue_order;
+  bool has_enqueue_order = work_queue->GetFrontTaskEnqueueOrder(&enqueue_order);
+  work_queue->AssignToWorkQueueSets(this);
+  work_queue->AssignSetIndex(set_index);
+  if (!has_enqueue_order)
+    return;
+  enqueue_order_to_work_queue_maps_[set_index].insert(
+      std::make_pair(enqueue_order, work_queue));
+}
+
 void WorkQueueSets::RemoveQueue(WorkQueue* work_queue) {
   DCHECK_EQ(this, work_queue->work_queue_sets());
   EnqueueOrder enqueue_order;
   bool has_enqueue_order = work_queue->GetFrontTaskEnqueueOrder(&enqueue_order);
+  work_queue->AssignToWorkQueueSets(nullptr);
   if (!has_enqueue_order)
     return;
   size_t set_index = work_queue->work_queue_set_index();
@@ -29,13 +43,15 @@ void WorkQueueSets::RemoveQueue(WorkQueue* work_queue) {
   enqueue_order_to_work_queue_maps_[set_index].erase(enqueue_order);
 }
 
-void WorkQueueSets::AssignQueueToSet(WorkQueue* work_queue, size_t set_index) {
+void WorkQueueSets::ChangeSetIndex(WorkQueue* work_queue, size_t set_index) {
+  DCHECK_EQ(this, work_queue->work_queue_sets());
   DCHECK_LT(set_index, enqueue_order_to_work_queue_maps_.size());
   EnqueueOrder enqueue_order;
   bool has_enqueue_order = work_queue->GetFrontTaskEnqueueOrder(&enqueue_order);
   size_t old_set = work_queue->work_queue_set_index();
   DCHECK_LT(old_set, enqueue_order_to_work_queue_maps_.size());
-  work_queue->AssignToWorkQueueSets(this, set_index);
+  DCHECK_NE(old_set, set_index);
+  work_queue->AssignSetIndex(set_index);
   if (!has_enqueue_order)
     return;
   enqueue_order_to_work_queue_maps_[old_set].erase(enqueue_order);
@@ -44,6 +60,8 @@ void WorkQueueSets::AssignQueueToSet(WorkQueue* work_queue, size_t set_index) {
 }
 
 void WorkQueueSets::OnPushQueue(WorkQueue* work_queue) {
+  // NOTE if this funciton changes, we need to keep |WorkQueueSets::AddQueue| in
+  // sync.
   DCHECK_EQ(this, work_queue->work_queue_sets());
   EnqueueOrder enqueue_order;
   bool has_enqueue_order = work_queue->GetFrontTaskEnqueueOrder(&enqueue_order);
@@ -98,7 +116,8 @@ bool WorkQueueSets::IsSetEmpty(size_t set_index) const {
 }
 
 #if DCHECK_IS_ON() || !defined(NDEBUG)
-bool WorkQueueSets::ContainsWorkQueueForTest(WorkQueue* work_queue) const {
+bool WorkQueueSets::ContainsWorkQueueForTest(
+    const WorkQueue* work_queue) const {
   EnqueueOrder enqueue_order;
   bool has_enqueue_order = work_queue->GetFrontTaskEnqueueOrder(&enqueue_order);
 
@@ -114,6 +133,11 @@ bool WorkQueueSets::ContainsWorkQueueForTest(WorkQueue* work_queue) const {
     }
   }
 
+  if (work_queue->work_queue_sets() == this) {
+    DCHECK(!has_enqueue_order);
+    return true;
+  }
+
   return false;
 }
 #endif
diff --git a/components/scheduler/base/work_queue_sets.h b/components/scheduler/base/work_queue_sets.h
index 7dd451a..0c2c2dc 100644
--- a/components/scheduler/base/work_queue_sets.h
+++ b/components/scheduler/base/work_queue_sets.h
@@ -22,14 +22,17 @@ class TaskQueueImpl;
 
 class SCHEDULER_EXPORT WorkQueueSets {
  public:
-  explicit WorkQueueSets(size_t num_sets);
+  WorkQueueSets(size_t num_sets, const char* name);
   ~WorkQueueSets();
 
   // O(log num queues)
+  void AddQueue(WorkQueue* queue, size_t set_index);
+
+  // O(log num queues)
   void RemoveQueue(WorkQueue* work_queue);
 
   // O(log num queues)
-  void AssignQueueToSet(WorkQueue* queue, size_t set_index);
+  void ChangeSetIndex(WorkQueue* queue, size_t set_index);
 
   // O(log num queues)
   void OnPushQueue(WorkQueue* work_queue);
@@ -46,12 +49,15 @@ class SCHEDULER_EXPORT WorkQueueSets {
 #if DCHECK_IS_ON() || !defined(NDEBUG)
   // Note this iterates over everything in |enqueue_order_to_work_queue_maps_|.
   // It's intended for use with DCHECKS and for testing
-  bool ContainsWorkQueueForTest(WorkQueue* queue) const;
+  bool ContainsWorkQueueForTest(const WorkQueue* queue) const;
 #endif
 
+  const char* name() const { return name_; }
+
  private:
   typedef std::map<EnqueueOrder, WorkQueue*> EnqueueOrderToWorkQueueMap;
   std::vector<EnqueueOrderToWorkQueueMap> enqueue_order_to_work_queue_maps_;
+  const char* name_;
 
   DISALLOW_COPY_AND_ASSIGN(WorkQueueSets);
 };
diff --git a/components/scheduler/base/work_queue_sets_unittest.cc b/components/scheduler/base/work_queue_sets_unittest.cc
index efbdef0..4ffefde 100644
--- a/components/scheduler/base/work_queue_sets_unittest.cc
+++ b/components/scheduler/base/work_queue_sets_unittest.cc
@@ -16,7 +16,16 @@ namespace internal {
 
 class WorkQueueSetsTest : public testing::Test {
  public:
-  void SetUp() override { work_queue_sets_.reset(new WorkQueueSets(kNumSets)); }
+  void SetUp() override {
+    work_queue_sets_.reset(new WorkQueueSets(kNumSets, "test"));
+  }
+
+  void TearDown() override {
+    for (scoped_ptr<WorkQueue>& work_queue : work_queues_) {
+      if (work_queue->work_queue_sets())
+        work_queue_sets_->RemoveQueue(work_queue.get());
+    }
+  }
 
  protected:
   enum {
@@ -26,7 +35,7 @@ class WorkQueueSetsTest : public testing::Test {
   WorkQueue* NewTaskQueue(const char* queue_name) {
     WorkQueue* queue = new WorkQueue(nullptr, "test");
     work_queues_.push_back(make_scoped_ptr(queue));
-    work_queue_sets_->AssignQueueToSet(queue, TaskQueue::CONTROL_PRIORITY);
+    work_queue_sets_->AddQueue(queue, TaskQueue::CONTROL_PRIORITY);
     return queue;
   }
 
@@ -41,10 +50,10 @@ class WorkQueueSetsTest : public testing::Test {
   scoped_ptr<WorkQueueSets> work_queue_sets_;
 };
 
-TEST_F(WorkQueueSetsTest, AssignQueueToSet) {
+TEST_F(WorkQueueSetsTest, ChangeSetIndex) {
   WorkQueue* work_queue = NewTaskQueue("queue");
   size_t set = TaskQueue::NORMAL_PRIORITY;
-  work_queue_sets_->AssignQueueToSet(work_queue, set);
+  work_queue_sets_->ChangeSetIndex(work_queue, set);
 
   EXPECT_EQ(set, work_queue->work_queue_set_index());
 }
@@ -52,7 +61,7 @@ TEST_F(WorkQueueSetsTest, AssignQueueToSet) {
 TEST_F(WorkQueueSetsTest, GetOldestQueueInSet_QueueEmpty) {
   WorkQueue* work_queue = NewTaskQueue("queue");
   size_t set = TaskQueue::NORMAL_PRIORITY;
-  work_queue_sets_->AssignQueueToSet(work_queue, set);
+  work_queue_sets_->ChangeSetIndex(work_queue, set);
 
   WorkQueue* selected_work_queue;
   EXPECT_FALSE(
@@ -62,7 +71,7 @@ TEST_F(WorkQueueSetsTest, GetOldestQueueInSet_QueueEmpty) {
 TEST_F(WorkQueueSetsTest, OnPushQueue) {
   WorkQueue* work_queue = NewTaskQueue("queue");
   size_t set = TaskQueue::NORMAL_PRIORITY;
-  work_queue_sets_->AssignQueueToSet(work_queue, set);
+  work_queue_sets_->ChangeSetIndex(work_queue, set);
 
   WorkQueue* selected_work_queue;
   EXPECT_FALSE(
@@ -79,7 +88,7 @@ TEST_F(WorkQueueSetsTest, GetOldestQueueInSet_SingleTaskInSet) {
   WorkQueue* work_queue = NewTaskQueue("queue");
   work_queue->Push(FakeTaskWithEnqueueOrder(10));
   size_t set = 1;
-  work_queue_sets_->AssignQueueToSet(work_queue, set);
+  work_queue_sets_->ChangeSetIndex(work_queue, set);
 
   WorkQueue* selected_work_queue;
   EXPECT_TRUE(work_queue_sets_->GetOldestQueueInSet(set, &selected_work_queue));
@@ -94,9 +103,9 @@ TEST_F(WorkQueueSetsTest, GetOldestQueueInSet_MultipleAgesInSet) {
   queue2->Push(FakeTaskWithEnqueueOrder(5));
   queue3->Push(FakeTaskWithEnqueueOrder(4));
   size_t set = 2;
-  work_queue_sets_->AssignQueueToSet(queue1, set);
-  work_queue_sets_->AssignQueueToSet(queue2, set);
-  work_queue_sets_->AssignQueueToSet(queue3, set);
+  work_queue_sets_->ChangeSetIndex(queue1, set);
+  work_queue_sets_->ChangeSetIndex(queue2, set);
+  work_queue_sets_->ChangeSetIndex(queue3, set);
 
   WorkQueue* selected_work_queue;
   EXPECT_TRUE(work_queue_sets_->GetOldestQueueInSet(set, &selected_work_queue));
@@ -112,9 +121,9 @@ TEST_F(WorkQueueSetsTest, OnPopQueue) {
   queue2->Push(FakeTaskWithEnqueueOrder(1));
   queue3->Push(FakeTaskWithEnqueueOrder(4));
   size_t set = 3;
-  work_queue_sets_->AssignQueueToSet(queue1, set);
-  work_queue_sets_->AssignQueueToSet(queue2, set);
-  work_queue_sets_->AssignQueueToSet(queue3, set);
+  work_queue_sets_->ChangeSetIndex(queue1, set);
+  work_queue_sets_->ChangeSetIndex(queue2, set);
+  work_queue_sets_->ChangeSetIndex(queue3, set);
 
   WorkQueue* selected_work_queue;
   EXPECT_TRUE(work_queue_sets_->GetOldestQueueInSet(set, &selected_work_queue));
@@ -135,9 +144,9 @@ TEST_F(WorkQueueSetsTest, OnPopQueue_QueueBecomesEmpty) {
   queue2->Push(FakeTaskWithEnqueueOrder(5));
   queue3->Push(FakeTaskWithEnqueueOrder(4));
   size_t set = 4;
-  work_queue_sets_->AssignQueueToSet(queue1, set);
-  work_queue_sets_->AssignQueueToSet(queue2, set);
-  work_queue_sets_->AssignQueueToSet(queue3, set);
+  work_queue_sets_->ChangeSetIndex(queue1, set);
+  work_queue_sets_->ChangeSetIndex(queue2, set);
+  work_queue_sets_->ChangeSetIndex(queue3, set);
 
   WorkQueue* selected_work_queue;
   EXPECT_TRUE(work_queue_sets_->GetOldestQueueInSet(set, &selected_work_queue));
@@ -158,10 +167,10 @@ TEST_F(WorkQueueSetsTest,
   queue1->Push(FakeTaskWithEnqueueOrder(0x7ffffff1));
   queue2->Push(FakeTaskWithEnqueueOrder(0x7ffffff0));
   queue3->Push(FakeTaskWithEnqueueOrder(-0x7ffffff1));
-  size_t set = 0;
-  work_queue_sets_->AssignQueueToSet(queue1, set);
-  work_queue_sets_->AssignQueueToSet(queue2, set);
-  work_queue_sets_->AssignQueueToSet(queue3, set);
+  size_t set = 1;
+  work_queue_sets_->ChangeSetIndex(queue1, set);
+  work_queue_sets_->ChangeSetIndex(queue2, set);
+  work_queue_sets_->ChangeSetIndex(queue3, set);
 
   WorkQueue* selected_work_queue;
   EXPECT_TRUE(work_queue_sets_->GetOldestQueueInSet(set, &selected_work_queue));
@@ -176,9 +185,9 @@ TEST_F(WorkQueueSetsTest, GetOldestQueueInSet_MultipleAgesInSet_RemoveQueue) {
   queue2->Push(FakeTaskWithEnqueueOrder(5));
   queue3->Push(FakeTaskWithEnqueueOrder(4));
   size_t set = 1;
-  work_queue_sets_->AssignQueueToSet(queue1, set);
-  work_queue_sets_->AssignQueueToSet(queue2, set);
-  work_queue_sets_->AssignQueueToSet(queue3, set);
+  work_queue_sets_->ChangeSetIndex(queue1, set);
+  work_queue_sets_->ChangeSetIndex(queue2, set);
+  work_queue_sets_->ChangeSetIndex(queue3, set);
   work_queue_sets_->RemoveQueue(queue3);
 
   WorkQueue* selected_work_queue;
@@ -186,7 +195,7 @@ TEST_F(WorkQueueSetsTest, GetOldestQueueInSet_MultipleAgesInSet_RemoveQueue) {
   EXPECT_EQ(queue2, selected_work_queue);
 }
 
-TEST_F(WorkQueueSetsTest, AssignQueueToSet_Complex) {
+TEST_F(WorkQueueSetsTest, ChangeSetIndex_Complex) {
   WorkQueue* queue1 = NewTaskQueue("queue1");
   WorkQueue* queue2 = NewTaskQueue("queue2");
   WorkQueue* queue3 = NewTaskQueue("queue3");
@@ -197,10 +206,10 @@ TEST_F(WorkQueueSetsTest, AssignQueueToSet_Complex) {
   queue4->Push(FakeTaskWithEnqueueOrder(3));
   size_t set1 = 1;
   size_t set2 = 2;
-  work_queue_sets_->AssignQueueToSet(queue1, set1);
-  work_queue_sets_->AssignQueueToSet(queue2, set1);
-  work_queue_sets_->AssignQueueToSet(queue3, set2);
-  work_queue_sets_->AssignQueueToSet(queue4, set2);
+  work_queue_sets_->ChangeSetIndex(queue1, set1);
+  work_queue_sets_->ChangeSetIndex(queue2, set1);
+  work_queue_sets_->ChangeSetIndex(queue3, set2);
+  work_queue_sets_->ChangeSetIndex(queue4, set2);
 
   WorkQueue* selected_work_queue;
   EXPECT_TRUE(
@@ -211,7 +220,7 @@ TEST_F(WorkQueueSetsTest, AssignQueueToSet_Complex) {
       work_queue_sets_->GetOldestQueueInSet(set2, &selected_work_queue));
   EXPECT_EQ(queue4, selected_work_queue);
 
-  work_queue_sets_->AssignQueueToSet(queue4, set1);
+  work_queue_sets_->ChangeSetIndex(queue4, set1);
 
   EXPECT_TRUE(
       work_queue_sets_->GetOldestQueueInSet(set1, &selected_work_queue));
@@ -223,21 +232,21 @@ TEST_F(WorkQueueSetsTest, AssignQueueToSet_Complex) {
 }
 
 TEST_F(WorkQueueSetsTest, IsSetEmpty_NoWork) {
-  size_t set = 0;
+  size_t set = 2;
   EXPECT_TRUE(work_queue_sets_->IsSetEmpty(set));
 
   WorkQueue* work_queue = NewTaskQueue("queue");
-  work_queue_sets_->AssignQueueToSet(work_queue, set);
+  work_queue_sets_->ChangeSetIndex(work_queue, set);
   EXPECT_TRUE(work_queue_sets_->IsSetEmpty(set));
 }
 
 TEST_F(WorkQueueSetsTest, IsSetEmpty_Work) {
-  size_t set = 0;
+  size_t set = 2;
   EXPECT_TRUE(work_queue_sets_->IsSetEmpty(set));
 
   WorkQueue* work_queue = NewTaskQueue("queue");
   work_queue->Push(FakeTaskWithEnqueueOrder(1));
-  work_queue_sets_->AssignQueueToSet(work_queue, set);
+  work_queue_sets_->ChangeSetIndex(work_queue, set);
   EXPECT_FALSE(work_queue_sets_->IsSetEmpty(set));
 
   work_queue->PopTaskForTest();