tlbmc/scheduler/scheduler.h

#ifndef THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_SCHEDULER_SCHEDULER_H_
#define THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_SCHEDULER_SCHEDULER_H_

#include <stdbool.h>

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <thread>

#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/functional/any_invocable.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "boost/asio/steady_timer.hpp"
#include "boost/circular_buffer.hpp"  //NOLINT: boost is commonly used in BMC
#include "time/clock.h"
#include "nlohmann/json.hpp"

namespace milotic_tlbmc {

// Represents the time information of a task.
struct TimeInfo {
  absl::Time last_scheduled_time = absl::InfiniteFuture();
  absl::Time last_run_time = absl::InfiniteFuture();
  boost::circular_buffer<absl::Duration> wait_times =
      boost::circular_buffer<absl::Duration>(1000);

  nlohmann::json ToJson() const;
};

// Represents the mode of a task.
//
// kPeriodic: The task will be scheduled periodically.
// kOnce: The task will be cancelled after it is executed once.
enum class TaskMode : uint8_t {
  kPeriodic = 0,
  kOnce,
};

// Represents a single task to be scheduled and run.
class Task : public std::enable_shared_from_this<Task> {
 public:
  Task(int task_id,
       absl::AnyInvocable<void(absl::AnyInvocable<void()>)> task_fn,
       absl::Duration period, boost::asio::io_context* io,
       absl::Duration task_timeout, TaskMode task_mode = TaskMode::kPeriodic,
       ecclesia::Clock* clock = ecclesia::Clock::RealClock());

  void Run() ABSL_LOCKS_EXCLUDED(task_mutex_, time_info_mutex_);
  void ScheduleNextRun() ABSL_LOCKS_EXCLUDED(timer_mutex_);
  absl::Duration GetPeriod() const;
  void SetPeriod(absl::Duration period);
  void Cancel() ABSL_LOCKS_EXCLUDED(timer_mutex_);
  void OnDone() ABSL_LOCKS_EXCLUDED(timer_mutex_, time_info_mutex_);
  bool IsCancelled() const { return cancelled_.load(); }
  int GetExecutionCount() const { return execution_count_.load(); }
  absl::Time GetLastRunTime() const ABSL_LOCKS_EXCLUDED(time_info_mutex_);
  absl::Duration GetAverageWaitTime() const
      ABSL_LOCKS_EXCLUDED(time_info_mutex_);
  absl::Duration GetTaskTimeout() const { return task_timeout_; }
  nlohmann::json ToJson() const;

 private:
  const int task_id_;
  const TaskMode task_mode_;
  // Ecclesia clock is thread safe.
  ecclesia::Clock* clock_ = nullptr;
  std::atomic<absl::Duration> period_ = absl::InfiniteDuration();
  std::atomic<bool> cancelled_ = false;
  std::atomic<int> execution_count_ = 0;

  // Task function to be executed.
  // task_fn_with_ack_: This is used for async tasks i.e tasks with
  // TaskMode kPeriodic that have an ack callback.
  absl::AnyInvocable<void(absl::AnyInvocable<void()>)> task_fn_with_ack_
      ABSL_GUARDED_BY(task_mutex_) = nullptr;
  mutable absl::Mutex task_mutex_;

  // TimeInfo to capture the last, next and average wait time.
  TimeInfo time_info_ ABSL_GUARDED_BY(time_info_mutex_);
  mutable absl::Mutex time_info_mutex_;

  // Timer to schedule the task.
  // Mainly used for async tasks.
  std::unique_ptr<boost::asio::steady_timer> timer_
      ABSL_GUARDED_BY(timer_mutex_) = nullptr;
  absl::Mutex timer_mutex_;
  const absl::Duration task_timeout_;
};

// Schedules tasks with a given periodicity.
//
// The scheduler is responsible for creating a thread for each period and
// scheduling the tasks to run on the corresponding thread based on the period.
// The scheduler also provides a way to reschedule a task with a new period and
// cancel a task.
//
// The scheduler is thread safe.
class TaskScheduler {
 public:
  // Options to configure the scheduler.
  struct Options {
    // Default options for the scheduler.
    //
    // Cleanup tasks period is set to 10 seconds arbitrarily, not shorter to
    // prevent overloading scheduler for cleanup tasks.
    //
    // Default task timeout is set to 60 seconds to allow for long running
    // tasks to complete which today aligns with the default timeout for
    // Redfish requests on client side.
    Options()
        : cleanup_tasks_period(absl::Seconds(10)),
          default_task_timeout(absl::Seconds(60)) {}

    absl::Duration cleanup_tasks_period;
    absl::Duration default_task_timeout;
  };

  explicit TaskScheduler(ecclesia::Clock* clock = ecclesia::Clock::RealClock(),
                         const Options& options = Options());
  virtual ~TaskScheduler();
  // Schedules the `task` for async invocation with the given `period`.
  //
  // The `task` will be invoked with a `on_done` callback as a parameter.
  // The `on_done` callback will be invoked when the task is done.
  // The `timeout` parameter is the timeout for the task to complete. If not
  // specified, the infinite timeout is used.
  int RunAndScheduleAsync(
      absl::AnyInvocable<void(absl::AnyInvocable<void()>)> task,
      absl::Duration period,
      std::optional<absl::Duration> timeout = absl::InfiniteDuration());

  // Schedule a oneshot task that is async.
  int ScheduleOneShotAsync(
      absl::AnyInvocable<void(absl::AnyInvocable<void()>)> task,
      absl::Duration period,
      std::optional<absl::Duration> timeout = std::nullopt);

  void UpdateTaskPeriod(int task_id, absl::Duration period)
      ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);
  void Cancel(int task_id) ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);
  void CancelAll() ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);
  void Stop() ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);
  size_t GetSchedulingAccuracyPercentage() const;
  nlohmann::json ToJson() const;

  // This only counts the number of user scheduled async tasks.
  size_t GetAllUserTaskCount() const ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);
  // This counts the number of all tasks including the internal scheduler tasks.
  size_t GetAllTaskCount() const ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

 protected:
  virtual absl::flat_hash_map<int, std::shared_ptr<Task>> GetSnapshot() const
      ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

  void CleanupInactiveTasks() ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

 private:
  std::atomic<int> next_task_id_;
  mutable absl::Mutex async_tasks_mutex_;
  absl::flat_hash_map<int, std::shared_ptr<Task>> id_to_async_tasks_
      ABSL_GUARDED_BY(async_tasks_mutex_);
  ecclesia::Clock* clock_;
  std::unique_ptr<boost::asio::io_context> io_;
  boost::asio::executor_work_guard<boost::asio::io_context::executor_type>
      work_guard_;
  std::thread async_executor_thread_;
  std::atomic<bool> stop_ = false;
  std::atomic<int> cleanup_task_id_ = -1;
  absl::Duration default_task_timeout_;
};

}  // namespace milotic_tlbmc

#endif  // THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_SCHEDULER_SCHEDULER_H_