1+ // Licensed to the Apache Software Foundation (ASF) under one
2+ // or more contributor license agreements. See the NOTICE file
3+ // distributed with this work for additional information
4+ // regarding copyright ownership. The ASF licenses this file
5+ // to you under the Apache License, Version 2.0 (the
6+ // "License"); you may not use this file except in compliance
7+ // with the License. You may obtain a copy of the License at
8+ //
9+ // http://www.apache.org/licenses/LICENSE-2.0
10+ //
11+ // Unless required by applicable law or agreed to in writing,
12+ // software distributed under the License is distributed on an
13+ // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+ // KIND, either express or implied. See the License for the
15+ // specific language governing permissions and limitations
16+ // under the License.
17+
18+ #pragma once
19+
20+ #include < mutex>
21+ #include < thread>
22+
23+ #include " util/blocking_priority_queue.hpp"
24+ #include " util/thread_group.h"
25+
26+ namespace doris {
27+
28+ // Work-Stealing threadpool which processes items (of type T) in parallel which were placed on multi
29+ // blocking queues by Offer(). Each item is processed by a single user-supplied method.
30+ class PriorityWorkStealingThreadPool : public PriorityThreadPool {
31+ public:
32+
33+ // Creates a new thread pool and start num_threads threads.
34+ // -- num_threads: how many threads are part of this pool
35+ // -- num_queues: how many queues are part of this pool
36+ // -- queue_size: the maximum size of the queue on which work items are offered. If the
37+ // queue exceeds this size, subsequent calls to Offer will block until there is
38+ // capacity available.
39+ // -- work_function: the function to run every time an item is consumed from the queue
40+ PriorityWorkStealingThreadPool (uint32_t num_threads, uint32_t num_queues, uint32_t queue_size)
41+ : PriorityThreadPool(0 , 0 ) {
42+ DCHECK_GT (num_queues, 0 );
43+ DCHECK_GE (num_threads, num_queues);
44+ // init _work_queues first because the work thread needs it
45+ for (int i = 0 ; i < num_queues; ++i) {
46+ _work_queues.emplace_back (std::make_shared<BlockingPriorityQueue<Task>>(queue_size));
47+ }
48+ for (int i = 0 ; i < num_threads; ++i) {
49+ _threads.create_thread (
50+ std::bind<void >(std::mem_fn (&PriorityWorkStealingThreadPool::work_thread), this , i));
51+ }
52+ }
53+
54+ // Blocking operation that puts a work item on the queue. If the queue is full, blocks
55+ // until there is capacity available.
56+ //
57+ // 'work' is copied into the work queue, but may be referenced at any time in the
58+ // future. Therefore the caller needs to ensure that any data referenced by work (if T
59+ // is, e.g., a pointer type) remains valid until work has been processed, and it's up to
60+ // the caller to provide their own signalling mechanism to detect this (or to wait until
61+ // after DrainAndshutdown returns).
62+ //
63+ // Returns true if the work item was successfully added to the queue, false otherwise
64+ // (which typically means that the thread pool has already been shut down).
65+ bool offer (Task task) override {
66+ return _work_queues[task.queue_id ]->blocking_put (task);
67+ }
68+
69+ bool offer (WorkFunction func) override {
70+ PriorityThreadPool::Task task = {0 , func, 0 };
71+ return _work_queues[task.queue_id ]->blocking_put (task);
72+ }
73+
74+ // Shuts the thread pool down, causing the work queue to cease accepting offered work
75+ // and the worker threads to terminate once they have processed their current work item.
76+ // Returns once the shutdown flag has been set, does not wait for the threads to
77+ // terminate.
78+ void shutdown () override {
79+ PriorityThreadPool::shutdown ();
80+ for (auto work_queue : _work_queues) {
81+ work_queue->shutdown ();
82+ }
83+ }
84+
85+ // Blocks until all threads are finished. shutdown does not need to have been called,
86+ // since it may be called on a separate thread.
87+ void join () override { _threads.join_all (); }
88+
89+ uint32_t get_queue_size () const override {
90+ uint32_t size = 0 ;
91+ for (auto work_queue : _work_queues) {
92+ size += work_queue->get_size ();
93+ }
94+ return size;
95+ }
96+
97+ // Blocks until the work queue is empty, and then calls shutdown to stop the worker
98+ // threads and Join to wait until they are finished.
99+ // Any work Offer()'ed during DrainAndshutdown may or may not be processed.
100+ void drain_and_shutdown () override {
101+ {
102+ std::unique_lock<std::mutex> l (_lock);
103+ while (get_queue_size () != 0 ) {
104+ _empty_cv.wait (l);
105+ }
106+ }
107+ shutdown ();
108+ join ();
109+ }
110+
111+ private:
112+ // Driver method for each thread in the pool. Continues to read work from the queue
113+ // until the pool is shutdown.
114+ void work_thread (int thread_id) {
115+ auto queue_id = thread_id % _work_queues.size ();
116+ auto steal_queue_id = (queue_id + 1 ) % _work_queues.size ();
117+ while (!is_shutdown ()) {
118+ Task task;
119+ // avoid blocking get
120+ bool is_other_queues_empty = true ;
121+ // steal work in round-robin if nothing to do
122+ while (_work_queues[queue_id]->get_size () == 0 && queue_id != steal_queue_id && !is_shutdown ()) {
123+ if (_work_queues[steal_queue_id]->non_blocking_get (&task)) {
124+ is_other_queues_empty = false ;
125+ task.work_function ();
126+ }
127+ steal_queue_id = (steal_queue_id + 1 ) % _work_queues.size ();
128+ }
129+ if (queue_id == steal_queue_id) {
130+ steal_queue_id = (steal_queue_id + 1 ) % _work_queues.size ();
131+ }
132+ if (is_other_queues_empty && _work_queues[queue_id]->blocking_get (&task, config::doris_blocking_priority_queue_wait_timeout_ms)) {
133+ task.work_function ();
134+ }
135+ if (_work_queues[queue_id]->get_size () == 0 ) {
136+ _empty_cv.notify_all ();
137+ }
138+ }
139+ }
140+
141+ // Queue on which work items are held until a thread is available to process them in
142+ // FIFO order.
143+ std::vector<std::shared_ptr<BlockingPriorityQueue<Task>>> _work_queues;
144+
145+ // Collection of worker threads that process work from the queues.
146+ ThreadGroup _threads;
147+
148+ // Guards _empty_cv
149+ std::mutex _lock;
150+
151+ // Signalled when the queue becomes empty
152+ std::condition_variable _empty_cv;
153+ };
154+
155+ } // namespace doris
0 commit comments