scheduler.hpp

#pragma once
 
#include <atomic>
#include <coroutine>
#include <cstddef>
#include <cstdint>
#include <deque>
#include <exception>
#include <functional>
#include <limits>
#include <queue>
#include <string>
#include <thread>
#include <vector>
 
#include "ravel/clock.hpp"
#include "ravel/rng.hpp"
#include "ravel/task.hpp"
#include "ravel/trace.hpp"
 
namespace ravel {
 
using TaskId = std::size_t;
 
using TaskFactory = std::function<Task()>;
 
enum class RunStatus {
  Completed,         
  TaskThrew,         
  StepLimitReached,  
  TimeLimitReached,  
};
 
struct RunReport {
  RunStatus status = RunStatus::Completed;
  std::uint64_t steps = 0;  
  std::string failure;      
};
 
class Scheduler {
 public:
  Scheduler(VirtualClock& clock, VirtualRng& rng, Trace& trace) noexcept
      : clock_(clock), rng_(rng), trace_(trace) {}
 
  Scheduler(const Scheduler&) = delete;
  Scheduler& operator=(const Scheduler&) = delete;
 
  void spawn(std::string name, TaskFactory factory);
 
  // Awaitables for use inside a task:
  //   co_await scheduler.yield();     // let the scheduler pick who runs next
  //   co_await scheduler.sleep(50);   // suspend for 50 virtual ticks
 
  class YieldAwaiter {
   public:
    bool await_ready() const noexcept { return false; }
    void await_suspend(std::coroutine_handle<>) const { scheduler_.requeue_current_task(); }
    void await_resume() const noexcept {}
 
   private:
    friend class Scheduler;
    explicit YieldAwaiter(Scheduler& scheduler) noexcept : scheduler_(scheduler) {}
    Scheduler& scheduler_;
  };
 
  class SleepAwaiter {
   public:
    bool await_ready() const noexcept { return false; }
    void await_suspend(std::coroutine_handle<>) const {
      scheduler_.sleep_current_task(duration_);
    }
    void await_resume() const noexcept {}
 
   private:
    friend class Scheduler;
    SleepAwaiter(Scheduler& scheduler, VirtualClock::Tick duration) noexcept
        : scheduler_(scheduler), duration_(duration) {}
    Scheduler& scheduler_;
    VirtualClock::Tick duration_;
  };
 
  [[nodiscard]] YieldAwaiter yield() noexcept { return YieldAwaiter(*this); }
  [[nodiscard]] SleepAwaiter sleep(VirtualClock::Tick duration) noexcept {
    return SleepAwaiter(*this, duration);
  }
 
  RunReport run_until_quiescent(
      std::uint64_t max_steps,
      VirtualClock::Tick time_limit = std::numeric_limits<VirtualClock::Tick>::max());
 
  const std::string& task_name(TaskId id) const { return slots_.at(id).name; }
 
  // Building blocks for blocking primitives such as Channel; tasks rarely
  // need them directly.
 
  VirtualClock::Tick now() const noexcept { return clock_.now(); }
 
  TaskId current_task() const noexcept { return current_task_; }
 
  void make_runnable(TaskId id) { runnable_.push_back(id); }
 
  void call_after(VirtualClock::Tick delay, std::function<void()> action);
 
 private:
  struct TaskSlot {
    TaskSlot(std::string task_name, TaskFactory task_factory)
        : name(std::move(task_name)), factory(std::move(task_factory)) {}
 
    std::string name;
    TaskFactory factory;  // Declared before `task`: the coroutine may use it,
    Task task;            // so it must be destroyed after the coroutine.
  };
 
  struct Timer {
    VirtualClock::Tick due;
    std::uint64_t sequence;  
    std::function<void()> action;
  };
 
  struct DueLater {
    bool operator()(const Timer& a, const Timer& b) const noexcept {
      return a.due != b.due ? a.due > b.due : a.sequence > b.sequence;
    }
  };
 
  void requeue_current_task() { make_runnable(current_task_); }
  void sleep_current_task(VirtualClock::Tick duration);
 
  bool has_pending_work() const noexcept { return !runnable_.empty() || !timers_.empty(); }
  void fire_earliest_timers();
  TaskId take_random_runnable_task();
 
  std::exception_ptr resume_task(TaskId id);
 
  void record(TaskId task, TraceEventKind kind) {
    trace_.record({clock_.now(), task, kind});
  }
 
  VirtualClock& clock_;
  VirtualRng& rng_;
  Trace& trace_;
 
  std::deque<TaskSlot> slots_;
  std::vector<TaskId> runnable_;
  std::priority_queue<Timer, std::vector<Timer>, DueLater> timers_;
  std::uint64_t next_timer_sequence_ = 0;
  TaskId current_task_ = 0;  
 
  void require_running_thread(const char* operation) const;
  std::atomic<std::thread::id> running_thread_{};  
};
 
}  // namespace ravel