src/ex/thread_pool.cpp

99.3% Lines (150 / 151) 100.0% Functions (28 / 28)

[ ] prev/next file n p prev/next uncovered line

thread_pool.cpp

f(x) Functions (28)

Function Calls Lines Blocks

Line	Hits	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	20991x	void push(continuation* c) noexcept
73		{
74	20991x	c->reserved = nullptr;
75	20991x	if(tail_)
76	2400x	tail_->reserved = c;
77		else
78	18591x	head_ = c;
79	20991x	tail_ = c;
80	20991x	}
81
82	23871x	continuation* pop() noexcept
83		{
84	23871x	if(!head_)
85	2880x	return nullptr;
86	20991x	continuation* c = head_;
87	20991x	head_ = static_cast<continuation*>(head_->reserved);
88	20991x	if(!head_)
89	18591x	tail_ = nullptr;
90	20991x	return c;
91		}
92
93	40012x	bool empty() const noexcept
94		{
95	40012x	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	1440x	~impl() = default;
111
112		bool
113	2721x	running_in_this_thread() const noexcept
114		{
115	2721x	return current_.get() == this;
116		}
117
118		// Abandon queued coroutines. Each is taken under the lock, as other
119		// threads may still post, and abandoned outside it, as abandoning
120		// can post.
121		void
122	2880x	drain_abandoned() noexcept
123		{
124		for(;;)
125		{
126		continuation* c;
127		{
128	3437x	std::lock_guard<std::mutex> lock(mutex_);
129	3437x	c = pop();
130	3437x	}
131	3437x	if(!c)
132	2880x	return;
133	557x	if(c->h && c->h != std::noop_coroutine())
134	300x	capy::abandon(*c);
135	557x	}
136		}
137
138	1440x	impl(std::size_t num_threads, std::string_view thread_name_prefix)
139	1440x	: num_threads_(num_threads)
140		{
141	1440x	if(num_threads_ == 0)
142	8x	num_threads_ = std::max(
143	4x	std::thread::hardware_concurrency(), 1u);
144
145		// Truncate prefix to 12 chars, leaving room for up to 3-digit index.
146	1440x	auto n = thread_name_prefix.copy(thread_name_prefix_, 12);
147	1440x	thread_name_prefix_[n] = '\0';
148	1440x	}
149
150		void
151	20991x	post(continuation& c)
152		{
153	20991x	if(!started_.load(std::memory_order_acquire))
154		{
155		{
156	603x	std::lock_guard<std::mutex> lock(mutex_);
157	603x	if(stop_)
158		{
159		// Stopped before any thread started: nothing will
160		// run it, and the destructor's drain abandons it.
161	220x	push(&c);
162	220x	return;
163		}
164	603x	}
165	383x	ensure_started();
166		}
167	20771x	std::lock_guard<std::mutex> lock(mutex_);
168	20771x	push(&c);
169		// Under the lock so the pool cannot drain, join, and
170		// destroy the condition variable mid-signal.
171	20771x	work_cv_.notify_one();
172	20771x	}
173
174		void
175	3598x	on_work_started() noexcept
176		{
177	3598x	outstanding_work_.fetch_add(1, std::memory_order_acq_rel);
178	3598x	}
179
180		void
181	3298x	on_work_finished() noexcept
182		{
183	3298x	auto n = outstanding_work_.load(std::memory_order_relaxed);
184	3307x	while(n > 1)
185		{
186	1402x	if(outstanding_work_.compare_exchange_weak(
187		n, n - 1, std::memory_order_acq_rel))
188	692x	return;
189		}
190	2606x	std::lock_guard<std::mutex> lock(mutex_);
191	2606x	outstanding_work_.fetch_sub(1, std::memory_order_acq_rel);
192		// Re-read: a concurrent on_work_started() may have raised it.
193	2606x	if(outstanding_work_.load(
194	2606x	std::memory_order_acquire) == 0 && joined_ && !stop_)
195		{
196	701x	stop_ = true;
197	701x	done_cv_.notify_all();
198	701x	work_cv_.notify_all();
199		}
200	2606x	}
201
202		void
203	2737x	join() noexcept
204		{
205		{
206	2737x	std::unique_lock<std::mutex> lock(mutex_);
207	2737x	if(joined_)
208	1297x	return;
209	1440x	joined_ = true;
210
211	1440x	if(outstanding_work_.load(
212	1440x	std::memory_order_acquire) == 0)
213		{
214	381x	stop_ = true;
215	381x	work_cv_.notify_all();
216		}
217		else
218		{
219	1059x	done_cv_.wait(lock, [this]{
220	2061x	return stop_;
221		});
222		}
223	2737x	}
224
225	1947x	for(auto& t : threads_)
226	507x	if(t.joinable())
227	507x	t.join();
228		}
229
230		void
231	1743x	stop() noexcept
232		{
233		// Under the lock: the join() this releases may return and let
234		// the pool be destroyed as soon as the lock is free.
235	1743x	std::lock_guard<std::mutex> lock(mutex_);
236	1743x	stop_ = true;
237	1743x	work_cv_.notify_all();
238	1743x	done_cv_.notify_all();
239	1743x	}
240
241		private:
242		void
243	383x	ensure_started()
244		{
245		// Under the lock, so join() either sees every thread started
246		// here or makes this start none. A joined pool still starts them:
247		// join() may be waiting on work that needs them.
248	383x	std::lock_guard<std::mutex> lock(mutex_);
249	383x	if(started_.load(std::memory_order_relaxed) \|\| stop_)
250	✗	return;
251	383x	threads_.reserve(num_threads_);
252	890x	for(std::size_t i = 0; i < num_threads_; ++i)
253	1014x	threads_.emplace_back([this, i]{ run(i); });
254	383x	started_.store(true, std::memory_order_release);
255	383x	}
256
257		void
258	507x	run(std::size_t index)
259		{
260		// Build name; set_current_thread_name truncates to platform limits.
261		char name[16];
262	507x	std::snprintf(name, sizeof(name), "%s%zu", thread_name_prefix_, index);
263	507x	set_current_thread_name(name);
264
265		// Mark this thread as a worker of this pool so dispatch()
266		// can symmetric-transfer when called from within pool work.
267		struct scoped_pool
268		{
269	507x	scoped_pool(impl const* p) noexcept { current_.set(p); }
270	507x	~scoped_pool() noexcept { current_.set(nullptr); }
271	507x	} guard(this);
272
273		for(;;)
274		{
275	20941x	continuation* c = nullptr;
276		{
277	20941x	std::unique_lock<std::mutex> lock(mutex_);
278	20941x	work_cv_.wait(lock, [this]{
279	59487x	return !empty() \|\|
280	59487x	stop_;
281		});
282	20941x	if(stop_)
283	1014x	return;
284	20434x	c = pop();
285	20941x	}
286	20434x	if(c)
287	20434x	safe_resume(c->h);
288	20434x	}
289	507x	}
290		};
291
292		//------------------------------------------------------------------------------
293
294	1440x	thread_pool::
295		~thread_pool()
296		{
297	1440x	impl_->stop();
298	1440x	impl_->join();
299	1440x	shutdown();
300		// Takes what the unwinds post too, and the work of any service they
301		// create, which is shut down as it is created.
302	1440x	impl_->drain_abandoned();
303	1440x	destroy();
304		// Services give up their work in shutdown(), but anything a
305		// destructor posts anyway is still destroyed.
306	1440x	impl_->drain_abandoned();
307	1440x	delete impl_;
308	1440x	}
309
310	1440x	thread_pool::
311	1440x	thread_pool(std::size_t num_threads, std::string_view thread_name_prefix)
312	1440x	: impl_(new impl(num_threads, thread_name_prefix))
313		{
314	1440x	this->set_frame_allocator(std::allocator<void>{});
315	1440x	}
316
317		void
318	1297x	thread_pool::
319		join() noexcept
320		{
321	1297x	impl_->join();
322	1297x	}
323
324		void
325	303x	thread_pool::
326		stop() noexcept
327		{
328	303x	impl_->stop();
329	303x	}
330
331		//------------------------------------------------------------------------------
332
333		thread_pool::executor_type
334	14857x	thread_pool::
335		get_executor() const noexcept
336		{
337	14857x	return executor_type(
338	14857x	const_cast<thread_pool&>(*this));
339		}
340
341		void
342	3598x	thread_pool::executor_type::
343		on_work_started() const noexcept
344		{
345	3598x	pool_->impl_->on_work_started();
346	3598x	}
347
348		void
349	3298x	thread_pool::executor_type::
350		on_work_finished() const noexcept
351		{
352	3298x	pool_->impl_->on_work_finished();
353	3298x	}
354
355		void
356	18277x	thread_pool::executor_type::
357		post(continuation& c) const
358		{
359	18277x	pool_->impl_->post(c);
360	18277x	}
361
362		std::coroutine_handle<>
363	2721x	thread_pool::executor_type::
364		dispatch(continuation& c) const
365		{
366	2721x	if(pool_->impl_->running_in_this_thread())
367	7x	return c.h;
368	2714x	pool_->impl_->post(c);
369	2714x	return std::noop_coroutine();
370		}
371
372		} // capy
373		} // boost
374