Utilities/detail/WorkerPool.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #include "Acts/Utilities/detail/WorkerPool.hpp"
0010
0011 #include <algorithm>
0012 #include <atomic>
0013 #include <condition_variable>
0014 #include <exception>
0015 #include <mutex>
0016 #include <thread>
0017 #include <utility>
0018 #include <vector>
0019
0020 namespace Acts::detail {
0021
0022 namespace {
0023 /// The pool whose loop body the current thread is executing, if any
0024 thread_local const void* tlActivePool = nullptr;
0025 }  // namespace
0026
0027 struct WorkerPool::Impl {
0028   std::vector<std::thread> threads;
0029
0030   /// Serializes concurrent parallelFor callers
0031   std::mutex callMutex;
0032
0033   std::mutex mutex;
0034   std::condition_variable wake;
0035   std::condition_variable done;
0036   std::size_t generation = 0;
0037   std::size_t busy = 0;
0038   bool stop = false;
0039
0040   // The current loop, written under `mutex` before the generation bump
0041   const std::function<void(std::size_t)>* body = nullptr;
0042   std::size_t n = 0;
0043   std::size_t grainSize = 1;
0044   std::atomic<std::size_t> next{0};
0045   std::atomic<bool> failed{false};
0046   std::exception_ptr error;
0047
0048   // Also runs when constructing WorkerPool fails after starting workers.
0049   // Join before any of the synchronization members are destroyed.
0050   ~Impl() {
0051     {
0052       std::lock_guard lock(mutex);
0053       stop = true;
0054     }
0055     wake.notify_all();
0056     for (auto& thread : threads) {
0057       thread.join();
0058     }
0059   }
0060
0061   void work() {
0062     tlActivePool = this;
0063     while (!failed.load(std::memory_order_relaxed)) {
0064       std::size_t begin = next.fetch_add(grainSize, std::memory_order_relaxed);
0065       if (begin >= n) {
0066         break;
0067       }
0068       std::size_t end = std::min(n, begin + grainSize);
0069       try {
0070         for (std::size_t i = begin; i < end; ++i) {
0071           (*body)(i);
0072         }
0073       } catch (...) {
0074         std::lock_guard lock(mutex);
0075         if (!error) {
0076           error = std::current_exception();
0077         }
0078         failed.store(true, std::memory_order_relaxed);
0079       }
0080     }
0081     tlActivePool = nullptr;
0082   }
0083
0084   void workerLoop() {
0085     std::size_t seen = 0;
0086     while (true) {
0087       {
0088         std::unique_lock lock(mutex);
0089         wake.wait(lock, [&] { return stop || generation != seen; });
0090         if (stop) {
0091           return;
0092         }
0093         seen = generation;
0094       }
0095       work();
0096       std::lock_guard lock(mutex);
0097       if (--busy == 0) {
0098         done.notify_one();
0099       }
0100     }
0101   }
0102 };
0103
0104 WorkerPool::WorkerPool(std::size_t numThreads)
0105     : m_impl(std::make_unique<Impl>()) {
0106   for (std::size_t i = 1; i < numThreads; ++i) {
0107     m_impl->threads.emplace_back([impl = m_impl.get()] { impl->workerLoop(); });
0108   }
0109 }
0110
0111 WorkerPool::~WorkerPool() = default;
0112
0113 std::size_t WorkerPool::size() const {
0114   return m_impl->threads.size() + 1;
0115 }
0116
0117 void WorkerPool::parallelFor(std::size_t n,
0118                              const std::function<void(std::size_t)>& body,
0119                              std::size_t grainSize) {
0120   // Nested calls stay on the current thread, including calls to another
0121   // pool, to avoid acquiring pool locks in a cycle.
0122   if (m_impl->threads.empty() || n <= 1 || tlActivePool != nullptr) {
0123     for (std::size_t i = 0; i < n; ++i) {
0124       body(i);
0125     }
0126     return;
0127   }
0128
0129   std::lock_guard call(m_impl->callMutex);
0130   Impl& impl = *m_impl;
0131   {
0132     std::lock_guard lock(impl.mutex);
0133     impl.body = &body;
0134     impl.n = n;
0135     impl.grainSize =
0136         grainSize > 0 ? grainSize : std::max<std::size_t>(1, n / (8 * size()));
0137     impl.next.store(0, std::memory_order_relaxed);
0138     impl.failed.store(false, std::memory_order_relaxed);
0139     impl.error = nullptr;
0140     impl.busy = impl.threads.size();
0141     ++impl.generation;
0142   }
0143   impl.wake.notify_all();
0144
0145   impl.work();
0146
0147   std::exception_ptr error;
0148   {
0149     std::unique_lock lock(impl.mutex);
0150     impl.done.wait(lock, [&] { return impl.busy == 0; });
0151     impl.body = nullptr;
0152     error = std::exchange(impl.error, nullptr);
0153   }
0154   if (error) {
0155     std::rethrow_exception(error);
0156   }
0157 }
0158
0159 }  // namespace Acts::detail