[scheduler] Support adding one queue to multiple budget pools

third_party/WebKit/Source/platform/scheduler/renderer/task_queue_throttler.cc

 // Copyright 2015 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 "platform/scheduler/renderer/task_queue_throttler.h"
 
 #include <cstdint>
 
 #include "base/format_macros.h"
 #include "base/logging.h"
@@ skipping 152 matching lines @@
   if (!allow_throttling_)
     return false;
 
   auto find_it = queue_details_.find(task_queue);
   if (find_it == queue_details_.end())
     return false;
   return find_it->second.throttling_ref_count > 0;
 }
 
 void TaskQueueThrottler::UnregisterTaskQueue(TaskQueue* task_queue) {
-  LazyNow lazy_now(tick_clock_);
   auto find_it = queue_details_.find(task_queue);
-
   if (find_it == queue_details_.end())
     return;
 
-  if (find_it->second.budget_pool)
-    find_it->second.budget_pool->RemoveQueue(lazy_now.Now(), task_queue);
+  LazyNow lazy_now(tick_clock_);
+  std::unordered_set<BudgetPool*> budget_pools = find_it->second.budget_pools;
+  for (BudgetPool* budget_pool : budget_pools) {
+    budget_pool->RemoveQueue(lazy_now.Now(), task_queue);
+  }
 
-  queue_details_.erase(find_it);
+  // Iterator may have been deleted by BudgetPool::RemoveQueue, so don't
+  // use it here.
+  queue_details_.erase(task_queue);
 }
 
 void TaskQueueThrottler::OnTimeDomainHasImmediateWork(TaskQueue* queue) {
   // Forward to the main thread if called from another thread
   if (!task_runner_->RunsTasksOnCurrentThread()) {
     task_runner_->PostTask(FROM_HERE,
                            base::Bind(forward_immediate_work_callback_, queue));
     return;
   }
   TRACE_EVENT0(tracing_category_,
@@ skipping 125 matching lines @@
   budget_pools_[time_budget_pool] = base::WrapUnique(time_budget_pool);
   return time_budget_pool;
 }
 
 void TaskQueueThrottler::OnTaskRunTimeReported(TaskQueue* task_queue,
                                                base::TimeTicks start_time,
                                                base::TimeTicks end_time) {
   if (!IsThrottled(task_queue))
     return;
 
-  BudgetPool* budget_pool = GetBudgetPoolForQueue(task_queue);
-  if (!budget_pool)
+  auto find_it = queue_details_.find(task_queue);
+  if (find_it == queue_details_.end())
     return;
 
-  budget_pool->RecordTaskRunTime(start_time, end_time);
-  if (!budget_pool->HasEnoughBudgetToRun(end_time))
-    budget_pool->BlockThrottledQueues(end_time);
+  for (BudgetPool* budget_pool : find_it->second.budget_pools) {
+    budget_pool->RecordTaskRunTime(start_time, end_time);
+    if (!budget_pool->HasEnoughBudgetToRun(end_time))
+      budget_pool->BlockThrottledQueues(end_time);
+  }
 }
 
 void TaskQueueThrottler::BlockQueue(base::TimeTicks now, TaskQueue* queue) {
   if (!IsThrottled(queue))
     return;
 
   queue->InsertFence(TaskQueue::InsertFencePosition::BEGINNING_OF_TIME);
   SchedulePumpQueue(FROM_HERE, now, queue);
 }
 
@@ skipping 26 matching lines @@
   state->EndDictionary();
 }
 
 void TaskQueueThrottler::AddQueueToBudgetPool(TaskQueue* queue,
                                               BudgetPool* budget_pool) {
   std::pair<TaskQueueMap::iterator, bool> insert_result =
       queue_details_.insert(std::make_pair(queue, Metadata()));
 
   Metadata& metadata = insert_result.first->second;
 
-  DCHECK(!metadata.budget_pool);
-  metadata.budget_pool = budget_pool;
+  DCHECK(metadata.budget_pools.find(budget_pool) ==
+         metadata.budget_pools.end());
+
+  metadata.budget_pools.insert(budget_pool);
 }
 
 void TaskQueueThrottler::RemoveQueueFromBudgetPool(TaskQueue* queue,
                                                    BudgetPool* budget_pool) {
   auto find_it = queue_details_.find(queue);
   DCHECK(find_it != queue_details_.end() &&
-         find_it->second.budget_pool == budget_pool);
+         find_it->second.budget_pools.find(budget_pool) !=
+             find_it->second.budget_pools.end());
 
-  find_it->second.budget_pool = nullptr;
+  find_it->second.budget_pools.erase(budget_pool);
 
   MaybeDeleteQueueMetadata(find_it);
 }
 
 void TaskQueueThrottler::UnregisterBudgetPool(BudgetPool* budget_pool) {
   budget_pools_.erase(budget_pool);
 }
 
 void TaskQueueThrottler::UnblockQueue(base::TimeTicks now, TaskQueue* queue) {
   SchedulePumpQueue(FROM_HERE, now, queue);
@@ skipping 11 matching lines @@
       NextTaskRunTime(&lazy_now, queue);
   if (!next_desired_run_time)
     return;
 
   base::Optional<base::TimeTicks> next_run_time =
       Max(next_desired_run_time, GetNextAllowedRunTime(now, queue));
 
   MaybeSchedulePumpThrottledTasks(from_here, now, next_run_time.value());
 }
 
-BudgetPool* TaskQueueThrottler::GetBudgetPoolForQueue(TaskQueue* queue) {
-  auto find_it = queue_details_.find(queue);
-  if (find_it == queue_details_.end())
-    return nullptr;
-  return find_it->second.budget_pool;
-}
-
 base::TimeTicks TaskQueueThrottler::GetNextAllowedRunTime(base::TimeTicks now,
                                                           TaskQueue* queue) {
-  BudgetPool* budget_pool = GetBudgetPoolForQueue(queue);
-  if (!budget_pool)
-    return now;
-  return std::max(now, budget_pool->GetNextAllowedRunTime());
+  base::TimeTicks next_run_time = now;
+
+  auto find_it = queue_details_.find(queue);
+  if (find_it == queue_details_.end())
+    return next_run_time;
+
+  for (BudgetPool* budget_pool : find_it->second.budget_pools) {
+    next_run_time =
+        std::max(next_run_time, budget_pool->GetNextAllowedRunTime());
+  }
+
+  return next_run_time;
 }
 
 void TaskQueueThrottler::MaybeDeleteQueueMetadata(TaskQueueMap::iterator it) {
-  if (it->second.throttling_ref_count == 0 && !it->second.budget_pool)
+  if (it->second.throttling_ref_count == 0 && it->second.budget_pools.empty())
     queue_details_.erase(it);
 }
 
 void TaskQueueThrottler::DisableThrottling() {
   if (!allow_throttling_)
     return;
 
   allow_throttling_ = false;
 
   for (const auto& map_entry : queue_details_) {
@@ skipping 32 matching lines @@
     queue->InsertFence(TaskQueue::InsertFencePosition::BEGINNING_OF_TIME);
     queue->SetTimeDomain(time_domain_.get());
     SchedulePumpQueue(FROM_HERE, lazy_now.Now(), queue);
   }
 
   TRACE_EVENT0(tracing_category_, "TaskQueueThrottler_EnableThrottling");
 }
 
 }  // namespace scheduler
 }  // namespace blink