xref: /aosp_15_r20/external/cronet/base/task/thread_pool/priority_queue.h (revision 6777b5387eb2ff775bb5750e3f5d96f37fb7352b)
1 // Copyright 2016 The Chromium Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 #ifndef BASE_TASK_THREAD_POOL_PRIORITY_QUEUE_H_
6 #define BASE_TASK_THREAD_POOL_PRIORITY_QUEUE_H_
7 
8 #include <functional>
9 #include <memory>
10 
11 #include "base/base_export.h"
12 #include "base/containers/intrusive_heap.h"
13 #include "base/task/common/checked_lock.h"
14 #include "base/task/thread_pool/task_source.h"
15 #include "base/task/thread_pool/task_source_sort_key.h"
16 #include "base/types/cxx23_to_underlying.h"
17 
18 namespace base {
19 namespace internal {
20 
21 // A PriorityQueue holds TaskSources of Tasks. This class is not thread-safe
22 // (requires external synchronization).
23 class BASE_EXPORT PriorityQueue {
24  public:
25   PriorityQueue();
26   PriorityQueue(const PriorityQueue&) = delete;
27   PriorityQueue& operator=(const PriorityQueue&) = delete;
28   ~PriorityQueue();
29 
30   PriorityQueue& operator=(PriorityQueue&& other);
31 
32   // Inserts |task_source| in the PriorityQueue with |task_source_sort_key|.
33   void Push(RegisteredTaskSource task_source,
34             TaskSourceSortKey task_source_sort_key);
35 
36   // Returns a reference to the TaskSourceSortKey representing the priority of
37   // the highest pending task in this PriorityQueue. The reference becomes
38   // invalid the next time that this PriorityQueue is modified.
39   // Cannot be called on an empty PriorityQueue.
40   const TaskSourceSortKey& PeekSortKey() const;
41 
42   // Returns a reference to the highest priority TaskSource in this
43   // PriorityQueue. Cannot be called on an empty PriorityQueue. The returned
44   // task source may be modified as long as its sort key isn't affected.
45   RegisteredTaskSource& PeekTaskSource() const;
46 
47   // Removes and returns the highest priority TaskSource in this PriorityQueue.
48   // Cannot be called on an empty PriorityQueue.
49   RegisteredTaskSource PopTaskSource();
50 
51   // Removes |task_source| from the PriorityQueue. Returns a
52   // RegisteredTaskSource which evaluates to true if successful, or false if
53   // |task_source| is not currently in the PriorityQueue or the PriorityQueue is
54   // empty.
55   RegisteredTaskSource RemoveTaskSource(const TaskSource& task_source);
56 
57   // Updates the sort key of |task_source| to |sort_key|, reordering
58   // |task_source| in the queue if necessary. No-ops if the TaskSource is not in
59   // the PriorityQueue or the PriorityQueue is empty.
60   void UpdateSortKey(const TaskSource& task_source, TaskSourceSortKey sort_key);
61 
62   // Returns true if the PriorityQueue is empty.
63   bool IsEmpty() const;
64 
65   // Returns the number of TaskSources in the PriorityQueue.
66   size_t Size() const;
67 
68   // Returns the number of TaskSources with |priority|.
GetNumTaskSourcesWithPriority(TaskPriority priority)69   size_t GetNumTaskSourcesWithPriority(TaskPriority priority) const {
70     return num_task_sources_per_priority_[base::to_underlying(priority)];
71   }
72 
73   // Set the PriorityQueue to empty all its TaskSources of Tasks when it is
74   // destroyed; needed to prevent memory leaks caused by a reference cycle
75   // (TaskSource -> Task -> TaskRunner -> TaskSource...) during test teardown.
76   void EnableFlushTaskSourcesOnDestroyForTesting();
77 
78   void swap(PriorityQueue& other);
79 
80  private:
81   // A class combining a TaskSource and the TaskSourceSortKey that determines
82   // its position in a PriorityQueue.
83   class TaskSourceAndSortKey;
84 
85   using ContainerType = IntrusiveHeap<TaskSourceAndSortKey>;
86 
87   void DecrementNumTaskSourcesForPriority(TaskPriority priority);
88   void IncrementNumTaskSourcesForPriority(TaskPriority priority);
89 
90   ContainerType container_;
91 
92   std::array<size_t, static_cast<int>(TaskPriority::HIGHEST) + 1>
93       num_task_sources_per_priority_ = {};
94 
95   // Should only be enabled by EnableFlushTaskSourcesOnDestroyForTesting().
96   bool is_flush_task_sources_on_destroy_enabled_ = false;
97 };
98 
99 }  // namespace internal
100 }  // namespace base
101 
102 #endif  // BASE_TASK_THREAD_POOL_PRIORITY_QUEUE_H_
103