timed_task.h

/*
 * 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 <chrono>
#include <functional>
#include <memory>
#include <queue>
 
#include <dispenso/detail/timed_task_impl.h>
#include <dispenso/priority.h>
#include <dispenso/timing.h>
 
namespace dispenso {
 
enum class TimedTaskType {
  kNormal 
  ,
  kSteady 
};
 
class TimedTask {
 public:
  TimedTask(const TimedTask&) = delete;
  TimedTask(TimedTask&& other) : impl_(std::move(other.impl_)) {}
 
  TimedTask& operator=(const TimedTask&) = delete;
  TimedTask& operator=(TimedTask&& other) {
    impl_ = std::move(other.impl_);
    return *this;
  }
 
  void cancel() {
    impl_->timesToRun.store(0, std::memory_order_release);
    impl_->flags.fetch_or(detail::kFFlagsCancelled, std::memory_order_release);
  }
 
  void detach() {
    impl_->flags.fetch_or(detail::kFFlagsDetached, std::memory_order_release);
  }
 
  size_t calls() const {
    return impl_->count.load(std::memory_order_acquire);
  }
 
  ~TimedTask() {
    if (!impl_ || impl_->flags.load(std::memory_order_acquire) & detail::kFFlagsDetached) {
      return;
    }
    cancel();
    while (impl_->inProgress.load(std::memory_order_acquire)) {
    }
    // Now we can safely destroy the underlying function.  We do this here because we can't risk
    // that func may call code in it's destructor that may no longer be relevant after this
    // TimedTask destructor completes.
    impl_->func = {};
  }
 
 private:
  template <typename Schedulable, typename F>
  TimedTask(
      Schedulable& sched,
      F&& f,
      double nextRunAbs,
      double period = 0.0,
      size_t timesToRun = 1,
      TimedTaskType type = TimedTaskType::kNormal)
      : impl_(
            detail::make_shared<detail::TimedTaskImpl>(
                timesToRun,
                nextRunAbs,
                period,
                std::forward<F>(f),
                sched,
                type == TimedTaskType::kSteady)) {}
 
  std::shared_ptr<detail::TimedTaskImpl> impl_;
 
  friend class TimedTaskScheduler;
}; // namespace dispenso
 
class TimedTaskScheduler {
 public:
  DISPENSO_DLL_ACCESS explicit TimedTaskScheduler(
      ThreadPriority priority = ThreadPriority::kNormal);
  DISPENSO_DLL_ACCESS ~TimedTaskScheduler();
 
  void setPriority(ThreadPriority priority) {
    std::lock_guard<std::mutex> lk(queueMutex_);
    priority_ = priority;
  }
 
  template <typename Schedulable, typename F>
  TimedTask schedule(
      Schedulable& sched,
      F&& func,
      double nextRunAbs,
      double period = 0.0,
      size_t timesToRun = 1,
      TimedTaskType type = TimedTaskType::kNormal) {
    TimedTask task(sched, std::forward<F>(func), nextRunAbs, period, timesToRun, type);
    addTimedTask(task.impl_);
    return task;
  }
 
  template <typename Schedulable, typename Rep, typename Period, typename F>
  TimedTask
  schedule(Schedulable& sched, F&& func, const std::chrono::duration<Rep, Period>& timeInFuture) {
    return schedule(sched, std::forward<F>(func), toNextRun(timeInFuture));
  }
 
  template <typename Schedulable, typename Clock, typename Duration, typename F>
  TimedTask schedule(
      Schedulable& sched,
      F&& func,
      const std::chrono::time_point<Clock, Duration>& nextRunTime) {
    return schedule(sched, std::forward<F>(func), toNextRun(nextRunTime));
  }
 
  template <typename Schedulable, typename Rep, typename Period, typename F>
  TimedTask schedule(
      Schedulable& sched,
      F&& func,
      const std::chrono::duration<Rep, Period>& timeInFuture,
      const std::chrono::duration<Rep, Period>& period,
      size_t timesToRun = std::numeric_limits<size_t>::max(),
      TimedTaskType type = TimedTaskType::kNormal) {
    return schedule(
        sched, std::forward<F>(func), toNextRun(timeInFuture), toPeriod(period), timesToRun, type);
  }
 
  template <
      typename Schedulable,
      typename Rep,
      typename Period,
      typename Clock,
      typename Duration,
      typename F>
  TimedTask schedule(
      Schedulable& sched,
      F&& func,
      const std::chrono::time_point<Clock, Duration>& nextRunTime,
      const std::chrono::duration<Rep, Period>& period,
      size_t timesToRun = std::numeric_limits<size_t>::max(),
      TimedTaskType type = TimedTaskType::kNormal) {
    return schedule(
        sched, std::forward<F>(func), toNextRun(nextRunTime), toPeriod(period), timesToRun, type);
  }
 
 private:
  template <class Rep, class Period>
  static double toNextRun(const std::chrono::duration<Rep, Period>& timeInFuture) {
    return getTime() + std::chrono::duration<double>(timeInFuture).count();
  }
 
  template <typename Clock, typename Duration>
  static double toNextRun(const std::chrono::time_point<Clock, Duration>& nextRunTime) {
    auto curTime = Clock::now();
    return toNextRun(nextRunTime - curTime);
  }
 
  template <class Rep, class Period>
  static double toPeriod(const std::chrono::duration<Rep, Period>& period) {
    return std::chrono::duration<double>(period).count();
  }
  DISPENSO_DLL_ACCESS void addTimedTask(std::shared_ptr<detail::TimedTaskImpl> task);
  void timeQueueRunLoop();
 
  void kickOffTask(std::shared_ptr<detail::TimedTaskImpl> next, double curTime);
 
  struct Compare {
    bool operator()(
        const std::shared_ptr<detail::TimedTaskImpl>& a,
        const std::shared_ptr<detail::TimedTaskImpl>& b) const {
      return a->nextAbsTime > b->nextAbsTime;
    }
  };
 
  // TODO(bbudge): Consider lock-free priority queue implementation.  I'd expect it to be minimally
  // beneficial for this use case though... timed tasks should rarely be super-high contention.
  std::mutex queueMutex_;
  std::priority_queue<
      std::shared_ptr<detail::TimedTaskImpl>,
      std::vector<std::shared_ptr<detail::TimedTaskImpl>>,
      Compare>
      tasks_;
  bool running_{true};
  detail::EpochWaiter epoch_;
  std::thread thread_;
  ThreadPriority priority_;
};
 
DISPENSO_DLL_ACCESS TimedTaskScheduler& globalTimedTaskScheduler();
 
} // namespace dispenso