thread_pool.h

Go to the documentation of this file.

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */
 
#pragma once
 
#include <atomic>
#include <cassert>
#include <condition_variable>
#include <cstdlib>
#include <deque>
#include <iterator>
#include <mutex>
#include <thread>
 
#include <moodycamel/concurrentqueue.h>
 
#include <dispenso/detail/epoch_waiter.h>
#include <dispenso/detail/per_thread_info.h>
#include <dispenso/once_function.h>
#include <dispenso/platform.h>
#include <dispenso/tsan_annotations.h>
 
namespace dispenso {
 
#if !defined(DISPENSO_WAKEUP_ENABLE)
#if defined(_WIN32) || defined(__linux__) || defined(__MACH__)
#define DISPENSO_WAKEUP_ENABLE 1
#else
#define DISPENSO_WAKEUP_ENABLE 0
#endif // platform
#endif // DISPENSO_WAKEUP_ENABLE
 
#if !defined(DISPENSO_POLL_PERIOD_US)
#if defined(_WIN32)
#define DISPENSO_POLL_PERIOD_US 1000
#else
#if !(DISPENSO_WAKEUP_ENABLE)
#define DISPENSO_POLL_PERIOD_US 200
#else
#define DISPENSO_POLL_PERIOD_US (1 << 15) // Determined empirically good on dual Xeon Linux
#endif // DISPENSO_WAKEUP_ENABLE
#endif // PLATFORM
#endif // DISPENSO_POLL_PERIOD_US
 
constexpr uint32_t kDefaultSleepLenUs = DISPENSO_POLL_PERIOD_US;
 
constexpr bool kDefaultWakeupEnable = DISPENSO_WAKEUP_ENABLE;
 
struct ForceQueuingTag {};
 
class alignas(kCacheLineSize) ThreadPool {
 public:
  DISPENSO_DLL_ACCESS ThreadPool(size_t n, size_t poolLoadMultiplier = 32);
 
  template <class Rep, class Period>
  void setSignalingWake(
      bool enable,
      const std::chrono::duration<Rep, Period>& sleepDuration =
          std::chrono::microseconds(kDefaultSleepLenUs)) {
    setSignalingWake(
        enable,
        static_cast<uint32_t>(
            std::chrono::duration_cast<std::chrono::microseconds>(sleepDuration).count()));
  }
 
  DISPENSO_DLL_ACCESS void resize(ssize_t n) {
    std::lock_guard<std::mutex> lk(threadsMutex_);
    resizeLocked(n);
  }
 
  ssize_t numThreads() const {
    return numThreads_.load(std::memory_order_relaxed);
  }
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f);
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f, ForceQueuingTag);
 
  template <typename Generator>
  void scheduleBulk(size_t count, Generator&& gen);
 
  DISPENSO_DLL_ACCESS ~ThreadPool();
 
 private:
  class PerThreadData {
   public:
    void setThread(std::thread&& t);
 
    bool running();
 
    void stop();
 
    ~PerThreadData();
 
   public:
    alignas(kCacheLineSize) std::thread thread_;
    std::atomic<bool> running_{true};
  };
 
  DISPENSO_DLL_ACCESS uint32_t wait(uint32_t priorEpoch);
  DISPENSO_DLL_ACCESS void wake();
  DISPENSO_DLL_ACCESS void wakeN(ssize_t n);
 
  void setSignalingWake(bool enable, uint32_t sleepDurationUs) {
    std::lock_guard<std::mutex> lk(threadsMutex_);
    ssize_t currentPoolSize = numThreads();
    resizeLocked(0);
    enableEpochWaiter_.store(enable, std::memory_order_release);
    sleepLengthUs_.store(sleepDurationUs, std::memory_order_release);
    resizeLocked(currentPoolSize);
  }
 
  DISPENSO_DLL_ACCESS void resizeLocked(ssize_t n);
 
  void executeNext(OnceFunction work);
 
  DISPENSO_DLL_ACCESS void threadLoop(PerThreadData& threadData);
 
  bool tryExecuteNext();
  bool tryExecuteNextFromProducerToken(moodycamel::ProducerToken& token);
 
  template <typename F>
  void schedule(moodycamel::ProducerToken& token, F&& f);
 
  template <typename F>
  void schedule(moodycamel::ProducerToken& token, F&& f, ForceQueuingTag);
 
  // Core bulk enqueue: unconditionally stage, enqueue, and wake for a chunk of tasks.
  // Caller is responsible for load factor checks. Count should be small (e.g. <= 2*numThreads).
  // When a producer token is provided, uses token-based enqueue for better throughput.
  template <typename Generator>
  void
  scheduleBulkEnqueue(size_t count, Generator&& gen, moodycamel::ProducerToken* token = nullptr);
 
  // The non-atomic read of two separate atomics (numSleeping_ and workRemaining_) can
  // race, potentially causing a missed wake in rare cases. This is benign: the epoch
  // waiter's sleep timeout (sleepLengthUs_) provides a safety net, so a missed wake only
  // delays wakeup by up to that duration, never causes a hang. Subsequent schedule()
  // calls will trigger wakes.
  void conditionallyWake() {
    if (enableEpochWaiter_.load(std::memory_order_acquire)) {
      ssize_t sleeping = numSleeping_.load(std::memory_order_acquire);
      if (sleeping > 0) {
        // Only wake if there's more pending work than awake threads can handle.
        // Don't count the caller — it may be a pool thread that will continue
        // dequeuing its own work, not processing what it just submitted.
        ssize_t awake = numThreads_.load(std::memory_order_relaxed) - sleeping;
        ssize_t pending = workRemaining_.load(std::memory_order_relaxed);
        if (pending > awake) {
          wake();
        }
      }
    }
  }
 
 public:
  // If we are not yet C++17, we provide aligned new/delete to avoid false sharing.
#if __cplusplus < 201703L
  static void* operator new(size_t sz) {
    return detail::alignedMalloc(sz);
  }
  static void operator delete(void* ptr) {
    return detail::alignedFree(ptr);
  }
#endif // __cplusplus
 
 private:
  // Number of tasks each thread accumulates before flushing workRemaining_.
  // This batching reduces atomic contention in threadLoop, but inflates
  // workRemaining_ by up to kWorkBatchSize * numThreads, so poolLoadFactor_
  // must be reduced accordingly for accurate load-shedding in schedule().
  static constexpr int kWorkBatchSize = 8;
 
  mutable std::mutex threadsMutex_;
  std::deque<PerThreadData> threads_;
  size_t poolLoadMultiplier_;
 
  // These atomics are read frequently in the hot schedule() path, so they need
  // cache-line alignment to avoid false sharing with the mutex/deque above.
  alignas(kCacheLineSize) std::atomic<ssize_t> poolLoadFactor_;
  std::atomic<ssize_t> numThreads_;
 
  moodycamel::ConcurrentQueue<OnceFunction> work_;
 
  alignas(kCacheLineSize) std::atomic<ssize_t> numSleeping_{0};
 
  alignas(kCacheLineSize) std::atomic<ssize_t> workRemaining_{0};
 
  alignas(kCacheLineSize) detail::EpochWaiter epochWaiter_;
  alignas(kCacheLineSize) std::atomic<bool> enableEpochWaiter_{kDefaultWakeupEnable};
  std::atomic<uint32_t> sleepLengthUs_{kDefaultSleepLenUs};
 
#if defined DISPENSO_DEBUG
  alignas(kCacheLineSize) std::atomic<ssize_t> outstandingTaskSets_{0};
#endif // NDEBUG
 
  friend class ConcurrentTaskSet;
  friend class TaskSet;
  friend class TaskSetBase;
};
 
DISPENSO_DLL_ACCESS ThreadPool& globalThreadPool();
 
DISPENSO_DLL_ACCESS void resizeGlobalThreadPool(size_t numThreads);
 
// ----------------------------- Implementation details -------------------------------------
 
template <typename F>
DISPENSO_REQUIRES(OnceCallableFunc<F>)
inline void ThreadPool::schedule(F&& f) {
  ssize_t curWork = workRemaining_.load(std::memory_order_relaxed);
  ssize_t quickLoadFactor = numThreads_.load(std::memory_order_relaxed);
  quickLoadFactor += quickLoadFactor / 2;
  if ((detail::PerPoolPerThreadInfo::isPoolRecursive(this) && curWork > quickLoadFactor) ||
      (curWork > poolLoadFactor_.load(std::memory_order_relaxed))) {
    f();
  } else {
    schedule(std::forward<F>(f), ForceQueuingTag());
  }
}
 
template <typename F>
DISPENSO_REQUIRES(OnceCallableFunc<F>)
inline void ThreadPool::schedule(F&& f, ForceQueuingTag) {
  if (auto* token =
          static_cast<moodycamel::ProducerToken*>(detail::PerPoolPerThreadInfo::producer(this))) {
    schedule(*token, std::forward<F>(f), ForceQueuingTag());
    return;
  }
 
  if (!numThreads_.load(std::memory_order_relaxed)) {
    f();
    return;
  }
  workRemaining_.fetch_add(1, std::memory_order_release);
  DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_BEGIN();
  bool enqueued = work_.enqueue({std::forward<F>(f)});
  DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_END();
  (void)(enqueued); // unused
  assert(enqueued);
 
  conditionallyWake();
}
 
template <typename F>
inline void ThreadPool::schedule(moodycamel::ProducerToken& token, F&& f) {
  ssize_t curWork = workRemaining_.load(std::memory_order_relaxed);
  ssize_t quickLoadFactor = numThreads_.load(std::memory_order_relaxed);
  quickLoadFactor += quickLoadFactor / 2;
  if ((detail::PerPoolPerThreadInfo::isPoolRecursive(this) && curWork > quickLoadFactor) ||
      (curWork > poolLoadFactor_.load(std::memory_order_relaxed))) {
    f();
  } else {
    schedule(token, std::forward<F>(f), ForceQueuingTag());
  }
}
 
template <typename F>
inline void ThreadPool::schedule(moodycamel::ProducerToken& token, F&& f, ForceQueuingTag) {
  if (!numThreads_.load(std::memory_order_relaxed)) {
    f();
    return;
  }
  workRemaining_.fetch_add(1, std::memory_order_release);
  DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_BEGIN();
  bool enqueued = work_.enqueue(token, {std::forward<F>(f)});
  DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_END();
  (void)(enqueued); // unused
  assert(enqueued);
 
  conditionallyWake();
}
 
inline bool ThreadPool::tryExecuteNext() {
  OnceFunction next;
  DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_BEGIN();
  bool dequeued = work_.try_dequeue(next);
  DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_END();
  if (dequeued) {
    executeNext(std::move(next));
    return true;
  }
  return false;
}
 
inline bool ThreadPool::tryExecuteNextFromProducerToken(moodycamel::ProducerToken& token) {
  OnceFunction next;
  if (work_.try_dequeue_from_producer(token, next)) {
    executeNext(std::move(next));
    return true;
  }
  return false;
}
 
inline void ThreadPool::executeNext(OnceFunction next) {
  next();
  workRemaining_.fetch_add(-1, std::memory_order_relaxed);
}
 
namespace detail {
// Generating iterator for scheduleBulkEnqueue. Produces OnceFunction objects
// on-the-fly during enqueue_bulk, avoiding the need for a staging buffer.
// moodycamel's enqueue_bulk uses single-pass input iterator semantics.
template <typename Generator>
struct BulkGenIter {
  using difference_type = std::ptrdiff_t;
  using value_type = OnceFunction;
  using pointer = OnceFunction*;
  using reference = OnceFunction&;
  using iterator_category = std::input_iterator_tag;
 
  Generator* gen;
  size_t index;
  OnceFunction operator*() {
    return (*gen)(index);
  }
  BulkGenIter& operator++() {
    ++index;
    return *this;
  }
  BulkGenIter operator++(int) {
    BulkGenIter tmp = *this;
    ++index;
    return tmp;
  }
};
} // namespace detail
 
template <typename Generator>
void ThreadPool::scheduleBulkEnqueue(
    size_t count,
    Generator&& gen,
    moodycamel::ProducerToken* token) {
  detail::BulkGenIter<typename std::remove_reference<Generator>::type> it{&gen, 0};
 
  // Single atomic update + bulk enqueue
  workRemaining_.fetch_add(static_cast<ssize_t>(count), std::memory_order_release);
 
  DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_BEGIN();
  bool enqueued;
  if (token) {
    enqueued = work_.enqueue_bulk(*token, it, count);
  } else {
    enqueued = work_.enqueue_bulk(it, count);
  }
  DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_END();
  (void)(enqueued);
  assert(enqueued);
 
  // Wake appropriate threads
  if (enableEpochWaiter_.load(std::memory_order_acquire)) {
    ssize_t sleeping = numSleeping_.load(std::memory_order_acquire);
    ssize_t toWake = std::min(sleeping, static_cast<ssize_t>(count));
    if (toWake > 0) {
      wakeN(toWake);
    }
  }
}
 
template <typename Generator>
void ThreadPool::scheduleBulk(size_t count, Generator&& gen) {
  if (count == 0) {
    return;
  }
 
  ssize_t numPool = numThreads_.load(std::memory_order_relaxed);
  if (!numPool) {
    // No threads in pool - execute all inline
    for (size_t i = 0; i < count; ++i) {
      gen(i)();
    }
    return;
  }
 
  // Process in chunks, interleaving enqueue and inline execution based on load.
  size_t chunkSize = static_cast<size_t>(numPool) + static_cast<size_t>(numPool) / 2;
  if (chunkSize < 1) {
    chunkSize = 1;
  }
  size_t i = 0;
  while (i < count) {
    ssize_t curWork = workRemaining_.load(std::memory_order_relaxed);
    ssize_t loadFactor = poolLoadFactor_.load(std::memory_order_relaxed);
    if (curWork > loadFactor) {
      // Over load factor - execute one task inline, then re-check
      gen(i)();
      ++i;
    } else {
      // Under load factor - enqueue a chunk
      ssize_t room = loadFactor - curWork;
      size_t toEnqueue = std::min({count - i, chunkSize, static_cast<size_t>(room)});
      if (toEnqueue == 0) {
        toEnqueue = 1;
      }
      size_t base = i;
      scheduleBulkEnqueue(toEnqueue, [&gen, base](size_t j) { return gen(base + j); });
      i += toEnqueue;
    }
  }
}
 
} // namespace dispenso