|
@@ -8,32 +8,107 @@
|
|
|
#include <functional>
|
|
#include <functional>
|
|
|
#include <mutex>
|
|
#include <mutex>
|
|
|
#include <string>
|
|
#include <string>
|
|
|
|
|
+#include <type_traits>
|
|
|
#include <vector>
|
|
#include <vector>
|
|
|
#include <queue>
|
|
#include <queue>
|
|
|
|
|
|
|
|
-#include "common/common_types.h"
|
|
|
|
|
|
|
+#include "common/thread.h"
|
|
|
#include "common/unique_function.h"
|
|
#include "common/unique_function.h"
|
|
|
|
|
|
|
|
namespace Common {
|
|
namespace Common {
|
|
|
|
|
|
|
|
-class ThreadWorker final {
|
|
|
|
|
|
|
+template <class StateType = void>
|
|
|
|
|
+class StatefulThreadWorker {
|
|
|
|
|
+ static constexpr bool with_state = !std::is_same_v<StateType, void>;
|
|
|
|
|
+
|
|
|
|
|
+ struct DummyCallable {
|
|
|
|
|
+ int operator()() const noexcept {
|
|
|
|
|
+ return 0;
|
|
|
|
|
+ }
|
|
|
|
|
+ };
|
|
|
|
|
+
|
|
|
|
|
+ using Task =
|
|
|
|
|
+ std::conditional_t<with_state, UniqueFunction<void, StateType*>, UniqueFunction<void>>;
|
|
|
|
|
+ using StateMaker = std::conditional_t<with_state, std::function<StateType()>, DummyCallable>;
|
|
|
|
|
+
|
|
|
public:
|
|
public:
|
|
|
- explicit ThreadWorker(std::size_t num_workers, const std::string& name);
|
|
|
|
|
- ~ThreadWorker();
|
|
|
|
|
- void QueueWork(UniqueFunction<void> work);
|
|
|
|
|
- void WaitForRequests();
|
|
|
|
|
|
|
+ explicit StatefulThreadWorker(size_t num_workers, std::string name, StateMaker func = {})
|
|
|
|
|
+ : workers_queued{num_workers}, thread_name{std::move(name)} {
|
|
|
|
|
+ const auto lambda = [this, func] {
|
|
|
|
|
+ Common::SetCurrentThreadName(thread_name.c_str());
|
|
|
|
|
+ {
|
|
|
|
|
+ std::conditional_t<with_state, StateType, int> state{func()};
|
|
|
|
|
+ while (!stop) {
|
|
|
|
|
+ Task task;
|
|
|
|
|
+ {
|
|
|
|
|
+ std::unique_lock lock{queue_mutex};
|
|
|
|
|
+ if (requests.empty()) {
|
|
|
|
|
+ wait_condition.notify_all();
|
|
|
|
|
+ }
|
|
|
|
|
+ condition.wait(lock, [this] { return stop || !requests.empty(); });
|
|
|
|
|
+ if (stop) {
|
|
|
|
|
+ break;
|
|
|
|
|
+ }
|
|
|
|
|
+ task = std::move(requests.front());
|
|
|
|
|
+ requests.pop();
|
|
|
|
|
+ }
|
|
|
|
|
+ if constexpr (with_state) {
|
|
|
|
|
+ task(&state);
|
|
|
|
|
+ } else {
|
|
|
|
|
+ task();
|
|
|
|
|
+ }
|
|
|
|
|
+ ++work_done;
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+ ++workers_stopped;
|
|
|
|
|
+ wait_condition.notify_all();
|
|
|
|
|
+ };
|
|
|
|
|
+ for (size_t i = 0; i < num_workers; ++i) {
|
|
|
|
|
+ threads.emplace_back(lambda);
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ ~StatefulThreadWorker() {
|
|
|
|
|
+ {
|
|
|
|
|
+ std::unique_lock lock{queue_mutex};
|
|
|
|
|
+ stop = true;
|
|
|
|
|
+ }
|
|
|
|
|
+ condition.notify_all();
|
|
|
|
|
+ for (std::thread& thread : threads) {
|
|
|
|
|
+ thread.join();
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ void QueueWork(Task work) {
|
|
|
|
|
+ {
|
|
|
|
|
+ std::unique_lock lock{queue_mutex};
|
|
|
|
|
+ requests.emplace(std::move(work));
|
|
|
|
|
+ ++work_scheduled;
|
|
|
|
|
+ }
|
|
|
|
|
+ condition.notify_one();
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ void WaitForRequests() {
|
|
|
|
|
+ std::unique_lock lock{queue_mutex};
|
|
|
|
|
+ wait_condition.wait(lock, [this] {
|
|
|
|
|
+ return workers_stopped >= workers_queued || work_done >= work_scheduled;
|
|
|
|
|
+ });
|
|
|
|
|
+ }
|
|
|
|
|
|
|
|
private:
|
|
private:
|
|
|
std::vector<std::thread> threads;
|
|
std::vector<std::thread> threads;
|
|
|
- std::queue<UniqueFunction<void>> requests;
|
|
|
|
|
|
|
+ std::queue<Task> requests;
|
|
|
std::mutex queue_mutex;
|
|
std::mutex queue_mutex;
|
|
|
std::condition_variable condition;
|
|
std::condition_variable condition;
|
|
|
std::condition_variable wait_condition;
|
|
std::condition_variable wait_condition;
|
|
|
std::atomic_bool stop{};
|
|
std::atomic_bool stop{};
|
|
|
- std::atomic<u64> work_scheduled{};
|
|
|
|
|
- std::atomic<u64> work_done{};
|
|
|
|
|
- std::atomic<u64> workers_stopped{};
|
|
|
|
|
- std::atomic<u64> workers_queued{};
|
|
|
|
|
|
|
+ std::atomic<size_t> work_scheduled{};
|
|
|
|
|
+ std::atomic<size_t> work_done{};
|
|
|
|
|
+ std::atomic<size_t> workers_stopped{};
|
|
|
|
|
+ std::atomic<size_t> workers_queued{};
|
|
|
|
|
+ std::string thread_name;
|
|
|
};
|
|
};
|
|
|
|
|
|
|
|
|
|
+using ThreadWorker = StatefulThreadWorker<>;
|
|
|
|
|
+
|
|
|
} // namespace Common
|
|
} // namespace Common
|