task_set.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
 
namespace dispenso {
enum class ParentCascadeCancel { kOff, kOn };
 
// Implementation note: kHeavy routes through the pool's per-group steal rings
// (schedulePlaced / scheduleBulkPlaced), which avoids consumer-side contention
// on the central moodycamel queue at high core counts. kLightweight routes
// through the central queue, which has lower per-task overhead but does not
// scale past ~100 contending consumers.
enum class TaskCost { kLightweight, kHeavy };
} // namespace dispenso
 
#include <dispenso/detail/task_set_impl.h>
 
namespace dispenso {
 
constexpr ssize_t kDefaultStealingMultiplier = 4;
 
class TaskSet : public TaskSetBase {
 public:
  TaskSet(
      ThreadPool& p,
      ParentCascadeCancel registerForParentCancel,
      ssize_t stealingLoadMultiplier = kDefaultStealingMultiplier)
      : TaskSetBase(p, registerForParentCancel, stealingLoadMultiplier),
        token_(makeToken(p.work_)) {}
 
  TaskSet(ThreadPool& p) : TaskSet(p, ParentCascadeCancel::kOff, kDefaultStealingMultiplier) {}
  TaskSet(ThreadPool& p, ssize_t stealingLoadMultiplier)
      : TaskSet(p, ParentCascadeCancel::kOff, stealingLoadMultiplier) {}
 
  TaskSet(TaskSet&& other) = delete;
  TaskSet& operator=(TaskSet&& other) = delete;
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f) {
    if (DISPENSO_EXPECT(canceled(), false)) {
      return;
    }
    if (outstandingTaskCount_.load(std::memory_order_relaxed) > taskSetLoadFactor_) {
      f();
    } else {
      pool_.schedule(token_, packageTask(std::forward<F>(f)));
    }
  }
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f, ForceQueuingTag fq) {
    pool_.schedule(token_, packageTask(std::forward<F>(f)), fq);
  }
 
  template <typename Generator>
  void scheduleBulk(size_t count, Generator&& gen) {
    scheduleBulkImpl(count, std::forward<Generator>(gen), &token_);
  }
 
  template <typename Generator>
  void scheduleBulk(size_t count, Generator&& gen, ForceQueuingTag) {
    scheduleBulkImplForceQueue(count, std::forward<Generator>(gen), &token_);
  }
 
  DISPENSO_DLL_ACCESS bool wait();
 
  DISPENSO_DLL_ACCESS bool tryWait(size_t maxToExecute);
 
  void cancel() {
    TaskSetBase::cancel();
  }
 
  bool canceled() const {
    return TaskSetBase::canceled();
  }
 
  ~TaskSet() {
    wait();
  }
 
 private:
  DISPENSO_DLL_ACCESS moodycamel::ProducerToken makeToken(
      moodycamel::ConcurrentQueue<OnceFunction>& pool);
 
  moodycamel::ProducerToken token_;
 
  template <typename Result>
  friend class detail::FutureBase;
  template <typename Result>
  friend class detail::FutureImplBase;
};
 
class ConcurrentTaskSet : public TaskSetBase {
 public:
  ConcurrentTaskSet(
      ThreadPool& pool,
      ParentCascadeCancel registerForParentCancel,
      ssize_t stealingLoadMultiplier = kDefaultStealingMultiplier,
      TaskCost cost = TaskCost::kHeavy)
      : TaskSetBase(pool, registerForParentCancel, stealingLoadMultiplier), cost_(cost) {}
 
  ConcurrentTaskSet(ThreadPool& p)
      : ConcurrentTaskSet(p, ParentCascadeCancel::kOff, kDefaultStealingMultiplier) {}
  ConcurrentTaskSet(ThreadPool& p, TaskCost cost)
      : ConcurrentTaskSet(p, ParentCascadeCancel::kOff, kDefaultStealingMultiplier, cost) {}
  ConcurrentTaskSet(ThreadPool& p, ssize_t stealingLoadMultiplier)
      : ConcurrentTaskSet(p, ParentCascadeCancel::kOff, stealingLoadMultiplier) {}
  ConcurrentTaskSet(ThreadPool& p, TaskCost cost, ssize_t stealingLoadMultiplier)
      : ConcurrentTaskSet(p, ParentCascadeCancel::kOff, stealingLoadMultiplier, cost) {}
 
  ConcurrentTaskSet(ConcurrentTaskSet&& other) = delete;
  ConcurrentTaskSet& operator=(ConcurrentTaskSet&& other) = delete;
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(
      F&& f,
      bool skipRecheck = false,
      float poolRecursiveLoadFactor = kDefaultPoolRecursiveLoadFactor) {
    if (cost_ == TaskCost::kHeavy) {
      schedulePlaced(std::forward<F>(f), skipRecheck, poolRecursiveLoadFactor);
      return;
    }
    // Combined inline decision (mirrors scheduleBulkImpl):
    // 1. TaskSet-level overload
    // 2. Pool-recursive: pool threads inline when workRemaining_ exceeds
    //    numThreads * poolRecursiveLoadFactor
    // 3. Pool global: non-recursive callers inline at the loose poolLoadFactor_
    // After this check, use ForceQueuingTag to skip the redundant check
    // in ThreadPool::schedule.
    if (outstandingTaskCount_.load(std::memory_order_relaxed) > taskSetLoadFactor_ &&
        DISPENSO_EXPECT(!canceled(), true) && detail::PerPoolPerThreadInfo::canInlineSchedule()) {
      detail::InlineDepthGuard depthGuard;
      f();
      return;
    }
    if (!skipRecheck) {
      ssize_t curWork = pool_.workRemaining_.load(std::memory_order_relaxed);
      ssize_t quickFactor =
          static_cast<ssize_t>(static_cast<float>(pool_.numThreads()) * poolRecursiveLoadFactor);
      if ((detail::PerPoolPerThreadInfo::isPoolRecursive(&pool_) && curWork > quickFactor) ||
          curWork > pool_.poolLoadFactor_.load(std::memory_order_relaxed)) {
        if (!detail::PerPoolPerThreadInfo::canInlineSchedule()) {
          pool_.schedule(packageTask(std::forward<F>(f)), ForceQueuingTag());
          return;
        }
        detail::InlineDepthGuard depthGuard;
        f();
        return;
      }
    }
    pool_.schedule(packageTask(std::forward<F>(f)), ForceQueuingTag());
  }
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f, ForceQueuingTag fq) {
    if (cost_ == TaskCost::kHeavy) {
      pool_.schedulePlaced(packageTask(std::forward<F>(f)), fq);
      return;
    }
    pool_.schedule(packageTask(std::forward<F>(f)), fq);
  }
 
  template <typename Generator>
  void scheduleBulk(size_t count, Generator&& gen) {
    if (cost_ == TaskCost::kHeavy) {
      scheduleBulkImplPlaced(count, std::forward<Generator>(gen));
      return;
    }
    scheduleBulkImpl(count, std::forward<Generator>(gen), nullptr);
  }
 
  template <typename Generator>
  void scheduleBulk(size_t count, Generator&& gen, ForceQueuingTag) {
    scheduleBulkImplForceQueue(count, std::forward<Generator>(gen), nullptr);
  }
 
  DISPENSO_DLL_ACCESS bool wait();
 
  DISPENSO_DLL_ACCESS bool tryWait(size_t maxToExecute);
 
  void cancel() {
    TaskSetBase::cancel();
  }
 
  bool canceled() const {
    return TaskSetBase::canceled();
  }
 
  ~ConcurrentTaskSet() {
    wait();
  }
 
 private:
  bool tryExecuteNext() {
    return pool_.tryExecuteNext();
  }
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedulePlaced(
      F&& f,
      bool skipRecheck = false,
      float poolRecursiveLoadFactor = kDefaultPoolRecursiveLoadFactor) {
    ssize_t placedThreshold = std::max(pool_.numThreads() + 1, taskSetLoadFactor_ / 2);
    if (outstandingTaskCount_.load(std::memory_order_relaxed) > placedThreshold &&
        DISPENSO_EXPECT(!canceled(), true) && detail::PerPoolPerThreadInfo::canInlineSchedule()) {
      detail::InlineDepthGuard depthGuard;
      f();
      return;
    }
    if (!skipRecheck) {
      ssize_t curWork = pool_.workRemaining_.load(std::memory_order_relaxed);
      ssize_t quickFactor =
          static_cast<ssize_t>(static_cast<float>(pool_.numThreads()) * poolRecursiveLoadFactor);
      if ((detail::PerPoolPerThreadInfo::isPoolRecursive(&pool_) && curWork > quickFactor) ||
          curWork > pool_.poolLoadFactor_.load(std::memory_order_relaxed)) {
        if (!detail::PerPoolPerThreadInfo::canInlineSchedule()) {
          pool_.schedulePlaced(packageTask(std::forward<F>(f)), ForceQueuingTag());
          return;
        }
        detail::InlineDepthGuard depthGuard;
        f();
        return;
      }
    }
    pool_.schedulePlaced(packageTask(std::forward<F>(f)), ForceQueuingTag());
  }
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedulePlaced(F&& f, ForceQueuingTag) {
    pool_.schedulePlaced(packageTask(std::forward<F>(f)), ForceQueuingTag());
  }
 
  template <typename Result>
  friend class detail::FutureBase;
  template <typename Result>
  friend class detail::FutureImplBase;
 
  friend class detail::LimitGatedScheduler;
 
  TaskCost cost_{TaskCost::kHeavy};
};
 
DISPENSO_DLL_ACCESS TaskSetBase* parentTaskSet();
 
} // namespace dispenso