src/ex/thread_pool.cpp

100.0% Lines (153 / 153) 100.0% Functions (28 / 28)

[ ] prev/next file

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	23707x	void push(continuation* c) noexcept
73		{
74	23707x	c->reserved = nullptr;
75	23707x	if(tail_)
76	6627x	tail_->reserved = c;
77		else
78	17080x	head_ = c;
79	23707x	tail_ = c;
80	23707x	}
81
82	26567x	continuation* pop() noexcept
83		{
84	26567x	if(!head_)
85	2860x	return nullptr;
86	23707x	continuation* c = head_;
87	23707x	head_ = static_cast<continuation*>(head_->reserved);
88	23707x	if(!head_)
89	17080x	tail_ = nullptr;
90	23707x	return c;
91		}
92
93	38579x	bool empty() const noexcept
94		{
95	38579x	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	1428x	~impl() = default;
111
112		bool
113	2709x	running_in_this_thread() const noexcept
114		{
115	2709x	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	2860x	drain_abandoned() noexcept
126		{
127	2860x	bool any = false;
128		for(;;)
129		{
130		continuation* c;
131		{
132	5411x	std::lock_guard<std::mutex> lock(mutex_);
133	5411x	c = pop();
134	5411x	}
135	5411x	if(!c)
136	2860x	return any;
137	2551x	any = true;
138	2551x	if(c->h && c->h != std::noop_coroutine())
139	2300x	capy::abandon(*c);
140	2551x	}
141		}
142
143	1428x	impl(std::size_t num_threads, std::string_view thread_name_prefix)
144	1428x	: num_threads_(num_threads)
145		{
146	1428x	if(num_threads_ == 0)
147	8x	num_threads_ = std::max(
148	4x	std::thread::hardware_concurrency(), 1u);
149
150		// Truncate prefix to 12 chars, leaving room for up to 3-digit index.
151	1428x	auto n = thread_name_prefix.copy(thread_name_prefix_, 12);
152	1428x	thread_name_prefix_[n] = '\0';
153	1428x	}
154
155		void
156	23707x	post(continuation& c)
157		{
158	23707x	if(!started_.load(std::memory_order_acquire))
159		{
160		{
161	2576x	std::lock_guard<std::mutex> lock(mutex_);
162	2576x	if(stop_)
163		{
164		// Stopped before any thread started: nothing will
165		// run it, and the destructor's drain abandons it.
166	2204x	push(&c);
167	2204x	return;
168		}
169	2576x	}
170	372x	ensure_started();
171		}
172	21503x	std::lock_guard<std::mutex> lock(mutex_);
173	21503x	push(&c);
174		// Under the lock so the pool cannot drain, join, and
175		// destroy the condition variable mid-signal.
176	21503x	work_cv_.notify_one();
177	21503x	}
178
179		void
180	3586x	on_work_started() noexcept
181		{
182	3586x	outstanding_work_.fetch_add(1, std::memory_order_acq_rel);
183	3586x	}
184
185		void
186	3286x	on_work_finished() noexcept
187		{
188	3286x	auto n = outstanding_work_.load(std::memory_order_relaxed);
189	3292x	while(n > 1)
190		{
191	2956x	if(outstanding_work_.compare_exchange_weak(
192		n, n - 1, std::memory_order_acq_rel))
193	1472x	return;
194		}
195	1814x	std::lock_guard<std::mutex> lock(mutex_);
196	1814x	outstanding_work_.fetch_sub(1, std::memory_order_acq_rel);
197		// Re-read: a concurrent on_work_started() may have raised it.
198	1814x	if(outstanding_work_.load(
199	1814x	std::memory_order_acquire) == 0 && joined_ && !stop_)
200		{
201	694x	stop_ = true;
202	694x	done_cv_.notify_all();
203	694x	work_cv_.notify_all();
204		}
205	1814x	}
206
207		void
208	2713x	join() noexcept
209		{
210		{
211	2713x	std::unique_lock<std::mutex> lock(mutex_);
212	2713x	if(joined_)
213	1285x	return;
214	1428x	joined_ = true;
215
216	1428x	if(outstanding_work_.load(
217	1428x	std::memory_order_acquire) == 0)
218		{
219	378x	stop_ = true;
220	378x	work_cv_.notify_all();
221		}
222		else
223		{
224	1050x	done_cv_.wait(lock, [this]{
225	2045x	return stop_;
226		});
227		}
228	2713x	}
229
230	1923x	for(auto& t : threads_)
231	495x	if(t.joinable())
232	495x	t.join();
233		}
234
235		void
236	1731x	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	1731x	std::lock_guard<std::mutex> lock(mutex_);
241	1731x	stop_ = true;
242	1731x	work_cv_.notify_all();
243	1731x	done_cv_.notify_all();
244	1731x	}
245
246		private:
247		void
248	372x	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	372x	std::lock_guard<std::mutex> lock(mutex_);
254	372x	if(started_.load(std::memory_order_relaxed) \|\| stop_)
255	1x	return;
256	371x	threads_.reserve(num_threads_);
257	866x	for(std::size_t i = 0; i < num_threads_; ++i)
258	990x	threads_.emplace_back([this, i]{ run(i); });
259	371x	started_.store(true, std::memory_order_release);
260	372x	}
261
262		void
263	495x	run(std::size_t index)
264		{
265		// Build name; set_current_thread_name truncates to platform limits.
266		char name[16];
267	495x	std::snprintf(name, sizeof(name), "%s%zu", thread_name_prefix_, index);
268	495x	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	495x	scoped_pool(impl const* p) noexcept { current_.set(p); }
275	495x	~scoped_pool() noexcept { current_.set(nullptr); }
276	495x	} guard(this);
277
278		for(;;)
279		{
280	21651x	continuation* c = nullptr;
281		{
282	21651x	std::unique_lock<std::mutex> lock(mutex_);
283	21651x	work_cv_.wait(lock, [this]{
284	55903x	return !empty() \|\|
285	55903x	stop_;
286		});
287	21651x	if(stop_)
288	990x	return;
289	21156x	c = pop();
290	21651x	}
291	21156x	if(c)
292	21156x	safe_resume(c->h);
293	21156x	}
294	495x	}
295		};
296
297		//------------------------------------------------------------------------------
298
299	1428x	thread_pool::
300		~thread_pool()
301		{
302	1428x	impl_->stop();
303	1428x	impl_->join();
304	1428x	shutdown();
305	1428x	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	1432x	destroy();
311		}
312	1432x	while(impl_->drain_abandoned());
313	1428x	delete impl_;
314	1428x	}
315
316	1428x	thread_pool::
317	1428x	thread_pool(std::size_t num_threads, std::string_view thread_name_prefix)
318	1428x	: impl_(new impl(num_threads, thread_name_prefix))
319		{
320	1428x	this->set_frame_allocator(std::allocator<void>{});
321	1428x	}
322
323		void
324	1285x	thread_pool::
325		join() noexcept
326		{
327	1285x	impl_->join();
328	1285x	}
329
330		void
331	303x	thread_pool::
332		stop() noexcept
333		{
334	303x	impl_->stop();
335	303x	}
336
337		//------------------------------------------------------------------------------
338
339		thread_pool::executor_type
340	14845x	thread_pool::
341		get_executor() const noexcept
342		{
343	14845x	return executor_type(
344	14845x	const_cast<thread_pool&>(*this));
345		}
346
347		void
348	3586x	thread_pool::executor_type::
349		on_work_started() const noexcept
350		{
351	3586x	pool_->impl_->on_work_started();
352	3586x	}
353
354		void
355	3286x	thread_pool::executor_type::
356		on_work_finished() const noexcept
357		{
358	3286x	pool_->impl_->on_work_finished();
359	3286x	}
360
361		void
362	21005x	thread_pool::executor_type::
363		post(continuation& c) const
364		{
365	21005x	pool_->impl_->post(c);
366	21005x	}
367
368		std::coroutine_handle<>
369	2709x	thread_pool::executor_type::
370		dispatch(continuation& c) const
371		{
372	2709x	if(pool_->impl_->running_in_this_thread())
373	7x	return c.h;
374	2702x	pool_->impl_->post(c);
375	2702x	return std::noop_coroutine();
376		}
377
378		} // capy
379		} // boost
380