1#include "raid/core/thread_pool.hpp"2#include <cstdlib>3#include <algorithm>4#include <memory>56namespace raid {78ThreadPool::ThreadPool(unsigned n) {9 if (n == 0) {10 if (const char* e = std::getenv("RAID_THREADS")) n = static_cast<unsigned>(std::max(1, std::atoi(e)));11 else n = std::max(1u, std::thread::hardware_concurrency());12 }13 for (unsigned i = 0; i < n; ++i) workers_.emplace_back([this] { worker(); });14}1516ThreadPool::~ThreadPool() {17 { std::lock_guard<std::mutex> lk(m_); stop_ = true; }18 cv_.notify_all();19 for (auto& t : workers_) t.join();20}2122void ThreadPool::worker() {23 for (;;) {24 std::function<void()> job;25 { std::unique_lock<std::mutex> lk(m_); cv_.wait(lk, [this] { return stop_ || !queue_.empty(); });26 if (stop_ && queue_.empty()) return;27 job = std::move(queue_.front()); queue_.pop_front(); }28 job();29 }30}3132void ThreadPool::parallel_for(usize n, const std::function<void(usize, usize)>& fn, usize grain) {33 if (n == 0) return;34 const unsigned nt = size();35 if (nt <= 1 || n == 1) { fn(0, n); return; }36 if (grain == 0) grain = std::max<usize>(1, (n + nt * 4 - 1) / (nt * 4));37 const usize chunks = (n + grain - 1) / grain;38 // Shared completion state: jobs may outlive this frame's locals by a few instructions, so the state lives on the heap39 // and is released by the last reference (job or caller).40 struct Done { std::mutex m; std::condition_variable cv; usize remaining; std::exception_ptr first_error; };41 auto done = std::make_shared<Done>(); done->remaining = chunks;42 for (usize c = 0; c < chunks; ++c) {43 const usize b = c * grain, e = std::min(n, b + grain);44 std::function<void()> job = [&fn, done, b, e] {45 std::exception_ptr err;46 try { fn(b, e); } catch (...) { err = std::current_exception(); }47 std::lock_guard<std::mutex> lk(done->m);48 if (err && !done->first_error) done->first_error = err;49 if (--done->remaining == 0) done->cv.notify_all();50 };51 { std::lock_guard<std::mutex> lk(m_); queue_.push_back(std::move(job)); }52 cv_.notify_one();53 }54 // Help from the calling thread as well, then wait.55 for (;;) {56 std::function<void()> job;57 { std::lock_guard<std::mutex> lk(m_); if (!queue_.empty()) { job = std::move(queue_.front()); queue_.pop_front(); } }58 if (!job) break;59 job();60 }61 std::exception_ptr first_error;62 { std::unique_lock<std::mutex> lk(done->m); done->cv.wait(lk, [&] { return done->remaining == 0; }); first_error = done->first_error; }63 if (first_error) std::rethrow_exception(first_error);64}6566ThreadPool& ThreadPool::global() { static ThreadPool pool; return pool; }6768} // namespace raid69