TLA Line data Source code
1 : //
2 : // Copyright (c) 2025 Vinnie Falco (vinnie.falco@gmail.com)
3 : // Copyright (c) 2026 Michael Vandeberg
4 : //
5 : // Distributed under the Boost Software License, Version 1.0. (See accompanying
6 : // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
7 : //
8 : // Official repository: https://github.com/boostorg/capy
9 : //
10 :
11 : #include <boost/capy/ex/thread_pool.hpp>
12 : #include <boost/capy/continuation.hpp>
13 : #include <boost/capy/detail/thread_local_ptr.hpp>
14 : #include <boost/capy/ex/abandon.hpp>
15 : #include <boost/capy/ex/frame_allocator.hpp>
16 : #include <boost/capy/test/thread_name.hpp>
17 : #include <algorithm>
18 : #include <atomic>
19 : #include <condition_variable>
20 : #include <cstdio>
21 : #include <mutex>
22 : #include <thread>
23 : #include <vector>
24 :
25 : /*
26 : Thread pool implementation using a shared work queue.
27 :
28 : Work items are continuations linked via their intrusive next pointer,
29 : stored in a single queue protected by a mutex. No per-post heap
30 : allocation: the continuation is owned by the caller and linked
31 : directly. Worker threads wait on a condition_variable until work
32 : is available or stop is requested.
33 :
34 : Threads are started lazily on first post(), under mutex_, to avoid
35 : spawning threads for pools that are constructed but never used. A
36 : stopped pool starts none. Each thread is named
37 : with a configurable prefix plus index for debugger visibility.
38 :
39 : Work tracking: on_work_started/on_work_finished maintain the atomic
40 : outstanding_work_ counter. on_work_started is lock-free, and so is
41 : on_work_finished while the count stays above one. The release that
42 : may take it to zero is made under mutex_, so join() can't see zero,
43 : return and let the pool be destroyed while that call still uses it.
44 : join() blocks until this counter reaches zero, then signals workers
45 : to stop and joins threads.
46 :
47 : Two shutdown paths:
48 : - join(): waits for outstanding work to drain, then stops workers.
49 : - stop(): immediately signals workers to exit; queued work stays queued.
50 : - Destructor: stop() then join(), then queued work is abandoned.
51 : */
52 :
53 : namespace boost {
54 : namespace capy {
55 :
56 : //------------------------------------------------------------------------------
57 :
58 : class thread_pool::impl
59 : {
60 : // Identifies the pool owning the current worker thread, or
61 : // nullptr if the calling thread is not a pool worker. Checked
62 : // by dispatch() to decide between symmetric transfer (inline
63 : // resume) and post.
64 : static inline detail::thread_local_ptr<impl const> current_;
65 :
66 : // Intrusive queue of continuations: the next link is stored in
67 : // continuation::reserved (typed continuation* round-tripped through
68 : // void*). No per-post allocation: the continuation is owned by the caller.
69 : continuation* head_ = nullptr;
70 : continuation* tail_ = nullptr;
71 :
72 HIT 25439 : void push(continuation* c) noexcept
73 : {
74 25439 : c->reserved = nullptr;
75 25439 : if(tail_)
76 6244 : tail_->reserved = c;
77 : else
78 19195 : head_ = c;
79 25439 : tail_ = c;
80 25439 : }
81 :
82 28300 : continuation* pop() noexcept
83 : {
84 28300 : if(!head_)
85 2861 : return nullptr;
86 25439 : continuation* c = head_;
87 25439 : head_ = static_cast<continuation*>(head_->reserved);
88 25439 : if(!head_)
89 19195 : tail_ = nullptr;
90 25439 : return c;
91 : }
92 :
93 42255 : bool empty() const noexcept
94 : {
95 42255 : return head_ == nullptr;
96 : }
97 :
98 : std::mutex mutex_;
99 : std::condition_variable work_cv_;
100 : std::condition_variable done_cv_;
101 : std::vector<std::thread> threads_;
102 : std::atomic<std::size_t> outstanding_work_{0};
103 : bool stop_{false};
104 : bool joined_{false};
105 : std::size_t num_threads_;
106 : char thread_name_prefix_[13]{}; // 12 chars max + null terminator
107 : std::atomic<bool> started_{false};
108 :
109 : public:
110 1429 : ~impl() = default;
111 :
112 : bool
113 3210 : running_in_this_thread() const noexcept
114 : {
115 3210 : return current_.get() == this;
116 : }
117 :
118 : // Abandon queued coroutines. Called after the services shut down,
119 : // since their shutdown can post more, and again after they are
120 : // destroyed, for anything their destructors post. Each is taken
121 : // under the lock, as other threads may still post, and abandoned
122 : // outside it, as abandoning can post. Returns whether there was
123 : // anything to abandon.
124 : bool
125 2861 : drain_abandoned() noexcept
126 : {
127 2861 : bool any = false;
128 : for(;;)
129 : {
130 : continuation* c;
131 : {
132 5404 : std::lock_guard<std::mutex> lock(mutex_);
133 5404 : c = pop();
134 5404 : }
135 5404 : if(!c)
136 2861 : return any;
137 2543 : any = true;
138 2543 : if(c->h && c->h != std::noop_coroutine())
139 2292 : capy::abandon(*c);
140 2543 : }
141 : }
142 :
143 1429 : impl(std::size_t num_threads, std::string_view thread_name_prefix)
144 1429 : : num_threads_(num_threads)
145 : {
146 1429 : if(num_threads_ == 0)
147 8 : num_threads_ = std::max(
148 4 : std::thread::hardware_concurrency(), 1u);
149 :
150 : // Truncate prefix to 12 chars, leaving room for up to 3-digit index.
151 1429 : auto n = thread_name_prefix.copy(thread_name_prefix_, 12);
152 1429 : thread_name_prefix_[n] = '\0';
153 1429 : }
154 :
155 : void
156 25439 : post(continuation& c)
157 : {
158 25439 : if(!started_.load(std::memory_order_acquire))
159 : {
160 : {
161 2576 : std::lock_guard<std::mutex> lock(mutex_);
162 2576 : if(stop_)
163 : {
164 : // Stopped before any thread started: nothing will
165 : // run it, and the destructor's drain abandons it.
166 2204 : push(&c);
167 2204 : return;
168 : }
169 2576 : }
170 372 : ensure_started();
171 : }
172 23235 : std::lock_guard<std::mutex> lock(mutex_);
173 23235 : push(&c);
174 : // Under the lock so the pool cannot drain, join, and
175 : // destroy the condition variable mid-signal.
176 23235 : work_cv_.notify_one();
177 23235 : }
178 :
179 : void
180 4087 : on_work_started() noexcept
181 : {
182 4087 : outstanding_work_.fetch_add(1, std::memory_order_acq_rel);
183 4087 : }
184 :
185 : void
186 3787 : on_work_finished() noexcept
187 : {
188 3787 : auto n = outstanding_work_.load(std::memory_order_relaxed);
189 3793 : while(n > 1)
190 : {
191 2604 : if(outstanding_work_.compare_exchange_weak(
192 : n, n - 1, std::memory_order_acq_rel))
193 1296 : return;
194 : }
195 2491 : std::lock_guard<std::mutex> lock(mutex_);
196 2491 : outstanding_work_.fetch_sub(1, std::memory_order_acq_rel);
197 : // Re-read: a concurrent on_work_started() may have raised it.
198 2491 : if(outstanding_work_.load(
199 2491 : std::memory_order_acquire) == 0 && joined_ && !stop_)
200 : {
201 703 : stop_ = true;
202 703 : done_cv_.notify_all();
203 703 : work_cv_.notify_all();
204 : }
205 2491 : }
206 :
207 : void
208 2715 : join() noexcept
209 : {
210 : {
211 2715 : std::unique_lock<std::mutex> lock(mutex_);
212 2715 : if(joined_)
213 1286 : return;
214 1429 : joined_ = true;
215 :
216 1429 : if(outstanding_work_.load(
217 1429 : std::memory_order_acquire) == 0)
218 : {
219 369 : stop_ = true;
220 369 : work_cv_.notify_all();
221 : }
222 : else
223 : {
224 1060 : done_cv_.wait(lock, [this]{
225 2064 : return stop_;
226 : });
227 : }
228 2715 : }
229 :
230 1926 : for(auto& t : threads_)
231 497 : if(t.joinable())
232 497 : t.join();
233 : }
234 :
235 : void
236 1732 : stop() noexcept
237 : {
238 : // Under the lock: the join() this releases may return and let
239 : // the pool be destroyed as soon as the lock is free.
240 1732 : std::lock_guard<std::mutex> lock(mutex_);
241 1732 : stop_ = true;
242 1732 : work_cv_.notify_all();
243 1732 : done_cv_.notify_all();
244 1732 : }
245 :
246 : private:
247 : void
248 372 : ensure_started()
249 : {
250 : // Under the lock, so join() either sees every thread started
251 : // here or makes this start none. A joined pool still starts them:
252 : // join() may be waiting on work that needs them.
253 372 : std::lock_guard<std::mutex> lock(mutex_);
254 372 : if(started_.load(std::memory_order_relaxed) || stop_)
255 MIS 0 : return;
256 HIT 372 : threads_.reserve(num_threads_);
257 869 : for(std::size_t i = 0; i < num_threads_; ++i)
258 994 : threads_.emplace_back([this, i]{ run(i); });
259 372 : started_.store(true, std::memory_order_release);
260 372 : }
261 :
262 : void
263 497 : run(std::size_t index)
264 : {
265 : // Build name; set_current_thread_name truncates to platform limits.
266 : char name[16];
267 497 : std::snprintf(name, sizeof(name), "%s%zu", thread_name_prefix_, index);
268 497 : set_current_thread_name(name);
269 :
270 : // Mark this thread as a worker of this pool so dispatch()
271 : // can symmetric-transfer when called from within pool work.
272 : struct scoped_pool
273 : {
274 497 : scoped_pool(impl const* p) noexcept { current_.set(p); }
275 497 : ~scoped_pool() noexcept { current_.set(nullptr); }
276 497 : } guard(this);
277 :
278 : for(;;)
279 : {
280 23393 : continuation* c = nullptr;
281 : {
282 23393 : std::unique_lock<std::mutex> lock(mutex_);
283 23393 : work_cv_.wait(lock, [this]{
284 61513 : return !empty() ||
285 61513 : stop_;
286 : });
287 23393 : if(stop_)
288 994 : return;
289 22896 : c = pop();
290 23393 : }
291 22896 : if(c)
292 22896 : safe_resume(c->h);
293 22896 : }
294 497 : }
295 : };
296 :
297 : //------------------------------------------------------------------------------
298 :
299 1429 : thread_pool::
300 : ~thread_pool()
301 : {
302 1429 : impl_->stop();
303 1429 : impl_->join();
304 1429 : shutdown();
305 1429 : impl_->drain_abandoned();
306 : // A frame the drain unwinds can create a service, whose destructor
307 : // can post again: repeat until nothing is left of either.
308 : do
309 : {
310 1432 : destroy();
311 : }
312 1432 : while(impl_->drain_abandoned());
313 1429 : delete impl_;
314 1429 : }
315 :
316 1429 : thread_pool::
317 1429 : thread_pool(std::size_t num_threads, std::string_view thread_name_prefix)
318 1429 : : impl_(new impl(num_threads, thread_name_prefix))
319 : {
320 1429 : this->set_frame_allocator(std::allocator<void>{});
321 1429 : }
322 :
323 : void
324 1286 : thread_pool::
325 : join() noexcept
326 : {
327 1286 : impl_->join();
328 1286 : }
329 :
330 : void
331 303 : thread_pool::
332 : stop() noexcept
333 : {
334 303 : impl_->stop();
335 303 : }
336 :
337 : //------------------------------------------------------------------------------
338 :
339 : thread_pool::executor_type
340 15346 : thread_pool::
341 : get_executor() const noexcept
342 : {
343 15346 : return executor_type(
344 15346 : const_cast<thread_pool&>(*this));
345 : }
346 :
347 : void
348 4087 : thread_pool::executor_type::
349 : on_work_started() const noexcept
350 : {
351 4087 : pool_->impl_->on_work_started();
352 4087 : }
353 :
354 : void
355 3787 : thread_pool::executor_type::
356 : on_work_finished() const noexcept
357 : {
358 3787 : pool_->impl_->on_work_finished();
359 3787 : }
360 :
361 : void
362 22236 : thread_pool::executor_type::
363 : post(continuation& c) const
364 : {
365 22236 : pool_->impl_->post(c);
366 22236 : }
367 :
368 : std::coroutine_handle<>
369 3210 : thread_pool::executor_type::
370 : dispatch(continuation& c) const
371 : {
372 3210 : if(pool_->impl_->running_in_this_thread())
373 7 : return c.h;
374 3203 : pool_->impl_->post(c);
375 3203 : return std::noop_coroutine();
376 : }
377 :
378 : } // capy
379 : } // boost
|
