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	46808x	void push(continuation* c) noexcept
73		{
74	46808x	c->reserved = nullptr;
75	46808x	if(tail_)
76	3252x	tail_->reserved = c;
77		else
78	43556x	head_ = c;
79	46808x	tail_ = c;
80	46808x	}
81
82	49680x	continuation* pop() noexcept
83		{
84	49680x	if(!head_)
85	2872x	return nullptr;
86	46808x	continuation* c = head_;
87	46808x	head_ = static_cast<continuation*>(head_->reserved);
88	46808x	if(!head_)
89	43556x	tail_ = nullptr;
90	46808x	return c;
91		}
92
93	90343x	bool empty() const noexcept
94		{
95	90343x	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	1436x	~impl() = default;
111
112		bool
113	13213x	running_in_this_thread() const noexcept
114		{
115	13213x	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	2872x	drain_abandoned() noexcept
123		{
124		for(;;)
125		{
126		continuation* c;
127		{
128	3409x	std::lock_guard<std::mutex> lock(mutex_);
129	3409x	c = pop();
130	3409x	}
131	3409x	if(!c)
132	2872x	return;
133	537x	if(c->h && c->h != std::noop_coroutine())
134	284x	capy::abandon(*c);
135	537x	}
136		}
137
138	1436x	impl(std::size_t num_threads, std::string_view thread_name_prefix)
139	1436x	: num_threads_(num_threads)
140		{
141	1436x	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	1436x	auto n = thread_name_prefix.copy(thread_name_prefix_, 12);
147	1436x	thread_name_prefix_[n] = '\0';
148	1436x	}
149
150		void
151	46810x	post(continuation& c)
152		{
153	46810x	if(!started_.load(std::memory_order_acquire))
154		{
155		{
156	601x	std::lock_guard<std::mutex> lock(mutex_);
157	601x	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	601x	}
165	381x	ensure_started();
166		}
167	46588x	std::lock_guard<std::mutex> lock(mutex_);
168	46588x	push(&c);
169		// Under the lock so the pool cannot drain, join, and
170		// destroy the condition variable mid-signal.
171	46588x	work_cv_.notify_one();
172	46588x	}
173
174		void
175	14096x	on_work_started() noexcept
176		{
177	14096x	outstanding_work_.fetch_add(1, std::memory_order_acq_rel);
178	14096x	}
179
180		void
181	13796x	on_work_finished() noexcept
182		{
183	13796x	auto n = outstanding_work_.load(std::memory_order_relaxed);
184	13806x	while(n > 1)
185		{
186	2014x	if(outstanding_work_.compare_exchange_weak(
187		n, n - 1, std::memory_order_acq_rel))
188	997x	return;
189		}
190	12799x	std::lock_guard<std::mutex> lock(mutex_);
191	12799x	outstanding_work_.fetch_sub(1, std::memory_order_acq_rel);
192		// Re-read: a concurrent on_work_started() may have raised it.
193	12799x	if(outstanding_work_.load(
194	12799x	std::memory_order_acquire) == 0 && joined_ && !stop_)
195		{
196	692x	stop_ = true;
197	692x	done_cv_.notify_all();
198	692x	work_cv_.notify_all();
199		}
200	12799x	}
201
202		void
203	2722x	join() noexcept
204		{
205		{
206	2722x	std::unique_lock<std::mutex> lock(mutex_);
207	2722x	if(joined_)
208	1286x	return;
209	1436x	joined_ = true;
210
211	1436x	if(outstanding_work_.load(
212	1436x	std::memory_order_acquire) == 0)
213		{
214	387x	stop_ = true;
215	387x	work_cv_.notify_all();
216		}
217		else
218		{
219	1049x	done_cv_.wait(lock, [this]{
220	2042x	return stop_;
221		});
222		}
223	2722x	}
224
225	1946x	for(auto& t : threads_)
226	510x	if(t.joinable())
227	510x	t.join();
228		}
229
230		void
231	1739x	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	1739x	std::lock_guard<std::mutex> lock(mutex_);
236	1739x	stop_ = true;
237	1739x	work_cv_.notify_all();
238	1739x	done_cv_.notify_all();
239	1739x	}
240
241		private:
242		void
243	381x	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	381x	std::lock_guard<std::mutex> lock(mutex_);
249	381x	if(started_.load(std::memory_order_relaxed) \|\| stop_)
250	✗	return;
251		// Start at the threads already running: if an earlier call
252		// threw partway, this one adds only the missing ones.
253	381x	threads_.reserve(num_threads_);
254	891x	for(std::size_t i = threads_.size(); i < num_threads_; ++i)
255	1022x	threads_.emplace_back([this, i]{ run(i); });
256	379x	started_.store(true, std::memory_order_release);
257	381x	}
258
259		void
260	510x	run(std::size_t index)
261		{
262		// Build name; set_current_thread_name truncates to platform limits.
263		char name[16];
264	510x	std::snprintf(name, sizeof(name), "%s%zu", thread_name_prefix_, index);
265	510x	set_current_thread_name(name);
266
267		// Mark this thread as a worker of this pool so dispatch()
268		// can symmetric-transfer when called from within pool work.
269		struct scoped_pool
270		{
271	510x	scoped_pool(impl const* p) noexcept { current_.set(p); }
272	510x	~scoped_pool() noexcept { current_.set(nullptr); }
273	510x	} guard(this);
274
275		for(;;)
276		{
277	46781x	continuation* c = nullptr;
278		{
279	46781x	std::unique_lock<std::mutex> lock(mutex_);
280	46781x	work_cv_.wait(lock, [this]{
281	134320x	return !empty() \|\|
282	134320x	stop_;
283		});
284	46781x	if(stop_)
285	1020x	return;
286	46271x	c = pop();
287	46781x	}
288	46271x	if(c)
289	46271x	safe_resume(c->h);
290	46271x	}
291	510x	}
292		};
293
294		//------------------------------------------------------------------------------
295
296	1436x	thread_pool::
297		~thread_pool()
298		{
299	1436x	impl_->stop();
300	1436x	impl_->join();
301	1436x	shutdown();
302		// Takes what the unwinds post too, and the work of any service they
303		// create, which is shut down as it is created.
304	1436x	impl_->drain_abandoned();
305	1436x	destroy();
306		// Services give up their work in shutdown(), but anything a
307		// destructor posts anyway is still destroyed.
308	1436x	impl_->drain_abandoned();
309	1436x	delete impl_;
310	1436x	}
311
312	1436x	thread_pool::
313	1436x	thread_pool(std::size_t num_threads, std::string_view thread_name_prefix)
314	1436x	: impl_(new impl(num_threads, thread_name_prefix))
315		{
316	1436x	this->set_frame_allocator(std::allocator<void>{});
317	1436x	}
318
319		void
320	1286x	thread_pool::
321		join() noexcept
322		{
323	1286x	impl_->join();
324	1286x	}
325
326		void
327	303x	thread_pool::
328		stop() noexcept
329		{
330	303x	impl_->stop();
331	303x	}
332
333		//------------------------------------------------------------------------------
334
335		thread_pool::executor_type
336	25360x	thread_pool::
337		get_executor() const noexcept
338		{
339	25360x	return executor_type(
340	25360x	const_cast<thread_pool&>(*this));
341		}
342
343		void
344	14096x	thread_pool::executor_type::
345		on_work_started() const noexcept
346		{
347	14096x	pool_->impl_->on_work_started();
348	14096x	}
349
350		void
351	13796x	thread_pool::executor_type::
352		on_work_finished() const noexcept
353		{
354	13796x	pool_->impl_->on_work_finished();
355	13796x	}
356
357		void
358	33604x	thread_pool::executor_type::
359		post(continuation& c) const
360		{
361	33604x	pool_->impl_->post(c);
362	33602x	}
363
364		std::coroutine_handle<>
365	13213x	thread_pool::executor_type::
366		dispatch(continuation& c) const
367		{
368	13213x	if(pool_->impl_->running_in_this_thread())
369	7x	return c.h;
370	13206x	pool_->impl_->post(c);
371	13206x	return std::noop_coroutine();
372		}
373
374		} // capy
375		} // boost
376