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 };
}
 
#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_);
  }
 
  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;
};
 
class ConcurrentTaskSet : public TaskSetBase {
 public:
  ConcurrentTaskSet(
      ThreadPool& pool,
      ParentCascadeCancel registerForParentCancel,
      ssize_t stealingLoadMultiplier = kDefaultStealingMultiplier)
      : TaskSetBase(pool, registerForParentCancel, stealingLoadMultiplier) {}
 
  ConcurrentTaskSet(ThreadPool& p)
      : ConcurrentTaskSet(p, ParentCascadeCancel::kOff, kDefaultStealingMultiplier) {}
  ConcurrentTaskSet(ThreadPool& p, ssize_t stealingLoadMultiplier)
      : ConcurrentTaskSet(p, ParentCascadeCancel::kOff, stealingLoadMultiplier) {}
 
  ConcurrentTaskSet(ConcurrentTaskSet&& other) = delete;
  ConcurrentTaskSet& operator=(ConcurrentTaskSet&& other) = delete;
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f, bool skipRecheck = false) {
    if (outstandingTaskCount_.load(std::memory_order_relaxed) > taskSetLoadFactor_ &&
        DISPENSO_EXPECT(!canceled(), true)) {
      f();
    } else if (skipRecheck) {
      pool_.schedule(packageTask(std::forward<F>(f)), ForceQueuingTag());
    } else {
      pool_.schedule(packageTask(std::forward<F>(f)));
    }
  }
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f, ForceQueuingTag fq) {
    pool_.schedule(packageTask(std::forward<F>(f)), fq);
  }
 
  template <typename Generator>
  void scheduleBulk(size_t count, Generator&& gen) {
    scheduleBulkImpl(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 Result>
  friend class detail::FutureBase;
 
  friend class detail::LimitGatedScheduler;
};
 
DISPENSO_DLL_ACCESS TaskSetBase* parentTaskSet();
 
} // namespace dispenso