TaskScheduler: Clear PriorityQueue during SchedulerWorkerDelegate:: and...

TaskScheduler: Clear PriorityQueue during SchedulerWorkerDelegate:: and SchedulerSingleThreadTaskRunnerManager::JoinForTesting

Bug: 889029
Change-Id: I543909faf8b66abf9a3866bc5aea8cc89277c8bb
Reviewed-on: https://chromium-review.googlesource.com/c/1338140
Commit-Queue: Jesse McKenna <jessemckenna@google.com>
Reviewed-by: François Doray <fdoray@chromium.org>
Cr-Commit-Position: refs/heads/master@{#608595}

base/task/task_scheduler/priority_queue.cc

@@ -147,11 +147,30 @@ size_t PriorityQueue::Transaction::Size() const {
 
 PriorityQueue::PriorityQueue() = default;
 
-PriorityQueue::~PriorityQueue() = default;
+PriorityQueue::~PriorityQueue() {
+  if (is_flush_sequences_on_destroy_enabled_) {
+    while (!container_.empty()) {
+      scoped_refptr<Sequence> sequence = BeginTransaction()->PopSequence();
+      {
+        std::unique_ptr<Sequence::Transaction> sequence_transaction =
+            sequence->BeginTransaction();
+        while (!sequence_transaction->IsEmpty()) {
+          sequence_transaction->TakeTask();
+          sequence_transaction->Pop();
+        }
+      }
+    }
+  }
+}
 
 std::unique_ptr<PriorityQueue::Transaction> PriorityQueue::BeginTransaction() {
   return WrapUnique(new Transaction(this));
 }
 
+void PriorityQueue::EnableFlushSequencesOnDestroyForTesting() {
+  DCHECK(!is_flush_sequences_on_destroy_enabled_);
+  is_flush_sequences_on_destroy_enabled_ = true;
+}
+
 }  // namespace internal
 }  // namespace base

base/task/task_scheduler/priority_queue.h

@@ -90,6 +90,11 @@ class BASE_EXPORT PriorityQueue {
   // PriorityQueue.
   std::unique_ptr<Transaction> BeginTransaction();
 
+  // Set the PriorityQueue to empty all its Sequences of Tasks when it is
+  // destroyed; needed to prevent memory leaks caused by a reference cycle
+  // (Sequence -> Task -> TaskRunner -> Sequence...) during test teardown.
+  void EnableFlushSequencesOnDestroyForTesting();
+
   const SchedulerLock* container_lock() const { return &container_lock_; }
 
  private:
@@ -104,6 +109,9 @@ class BASE_EXPORT PriorityQueue {
 
   ContainerType container_;
 
+  // Should only be enabled by EnableFlushSequencesOnDestroyForTesting().
+  bool is_flush_sequences_on_destroy_enabled_ = false;
+
   DISALLOW_COPY_AND_ASSIGN(PriorityQueue);
 };
 

base/task/task_scheduler/scheduler_single_thread_task_runner_manager.cc

@@ -129,6 +129,10 @@ class SchedulerWorkerDelegate : public SchedulerWorker::Delegate {
 
   void OnMainExit(SchedulerWorker* /* worker */) override {}
 
+  void EnableFlushPriorityQueueSequencesOnDestroyForTesting() {
+    priority_queue_.EnableFlushSequencesOnDestroyForTesting();
+  }
+
  private:
   const std::string thread_name_;
   const SchedulerWorker::ThreadLabel thread_label_;
@@ -486,8 +490,11 @@ void SchedulerSingleThreadTaskRunnerManager::JoinForTesting() {
     local_workers = std::move(workers_);
   }
 
-  for (const auto& worker : local_workers)
+  for (const auto& worker : local_workers) {
+    static_cast<SchedulerWorkerDelegate*>(worker->delegate())
+        ->EnableFlushPriorityQueueSequencesOnDestroyForTesting();
     worker->JoinForTesting();
+  }
 
   {
     AutoSchedulerLock auto_lock(lock_);

base/task/task_scheduler/scheduler_worker_pool_impl.cc

@@ -355,6 +355,8 @@ void SchedulerWorkerPoolImpl::JoinForTesting() {
   for (const auto& worker : workers_copy)
     worker->JoinForTesting();
 
+  shared_priority_queue_.EnableFlushSequencesOnDestroyForTesting();
+
   AutoSchedulerLock auto_lock(lock_);
   DCHECK(workers_ == workers_copy);
   // Release |workers_| to clear their TrackedRef against |this|.

base/task/task_scheduler/sequence.cc

@@ -43,21 +43,21 @@ bool Sequence::Transaction::PushTask(Task task) {
 }
 
 Optional<Task> Sequence::Transaction::TakeTask() {
-  DCHECK(!sequence_->queue_.empty());
+  DCHECK(!IsEmpty());
   DCHECK(sequence_->queue_.front().task);
 
   return std::move(sequence_->queue_.front());
 }
 
 bool Sequence::Transaction::Pop() {
-  DCHECK(!sequence_->queue_.empty());
+  DCHECK(!IsEmpty());
   DCHECK(!sequence_->queue_.front().task);
   sequence_->queue_.pop();
-  return sequence_->queue_.empty();
+  return IsEmpty();
 }
 
 SequenceSortKey Sequence::Transaction::GetSortKey() const {
-  DCHECK(!sequence_->queue_.empty());
+  DCHECK(!IsEmpty());
 
   // Save the sequenced time of the next task in the sequence.
   base::TimeTicks next_task_sequenced_time =
@@ -67,6 +67,10 @@ SequenceSortKey Sequence::Transaction::GetSortKey() const {
                          next_task_sequenced_time);
 }
 
+bool Sequence::Transaction::IsEmpty() const {
+  return sequence_->queue_.empty();
+}
+
 void Sequence::SetHeapHandle(const HeapHandle& handle) {
   heap_handle_ = handle;
 }

base/task/task_scheduler/sequence.h

@@ -76,6 +76,8 @@ class BASE_EXPORT Sequence : public RefCountedThreadSafe<Sequence> {
     // Cannot be called on an empty Sequence.
     SequenceSortKey GetSortKey() const;
 
+    bool IsEmpty() const;
+
     scoped_refptr<Sequence> sequence() { return sequence_; }
 
    private: