Fix bug with TaskQueueSelector and blocked queues

components/scheduler/base/work_queue_sets.cc

Index: components/scheduler/base/work_queue_sets.cc
diff --git a/components/scheduler/base/work_queue_sets.cc b/components/scheduler/base/work_queue_sets.cc
index f952334e733c31c99b3209cd3f9e806764795634..96c5da4a01fc2877eea7e897d6f7c31f1d85f327 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