Ray 源码解析（一）：任务的状态转移和组织形式

enum class TaskState {
  // The task may be placed on a node.
  PLACEABLE,
  // The task has been placed on a node and is waiting for some object
  // dependencies to become local.
  WAITING,
  // The task has been placed on a node, all dependencies are satisfied, and is
  // waiting for resources to run.
  READY,
  // The task is running on a worker. The task may also be blocked in a ray.get
  // or ray.wait call, in which case it also has state BLOCKED.
  RUNNING,
  // The task has resources that cannot be satisfied by any node, as far as we
  // know.
  INFEASIBLE,
  // The task is an actor method and is waiting to learn where the actor was
  // created.
  WAITING_FOR_ACTOR_CREATION,
  // Swap queue for tasks that are in between states. This can happen when a
  // task is removed from one queue, and an async callback is responsible for
  // re-queuing the task. For example, a READY task that has just been assigned
  // to a worker will get moved to the SWAP queue while waiting for a response
  // from the worker. If the worker accepts the task, the task will be added to
  // the RUNNING queue, else it will be returned to READY.
  SWAP,
  // The number of task queues. All states that precede this enum must have an
  // associated TaskQueue in SchedulingQueue. All states that succeed
  // this enum do not have an associated TaskQueue, since the tasks
  // in those states may not have any associated task data.
  kNumTaskQueues,
  // The task is running but blocked in a ray.get or ray.wait call. Tasks that
  // were explicitly assigned by us may be both BLOCKED and RUNNING, while
  // tasks that were created out-of-band (e.g., the application created
  // multiple threads) are only BLOCKED.
  BLOCKED,
  // The task is a driver task.
  DRIVER,
};

class TaskQueue {
 public:
  virtual ~TaskQueue() {}

  // 任务的增删改查操作
  virtual bool AppendTask(const TaskID &task_id, const Task &task);
  virtual bool RemoveTask(const TaskID &task_id,
                          std::vector<Task> *removed_tasks = nullptr);
  bool HasTask(const TaskID &task_id) const;
  const std::list<Task> &GetTasks() const;
  const Task &GetTask(const TaskID &task_id) const;

  // 获取队列中所需资源总和
  const ResourceSet &GetCurrentResourceLoad() const;

 protected:
  // 链表+哈希组织，可以快速查找O(1)和线性遍历O(n)
  std::list<Task> task_list_;
  std::unordered_map<TaskID, std::list<Task>::iterator> task_map_;
  
  // 所有任务所需资源总和
  ResourceSet current_resource_load_;
};

class SchedulingQueue {
 public:
  /// 创建一个空的调度队列，初始化各个状态对应的任务队列，就绪队列被单独拿出来用 ReadyQueue 做初始化。
  SchedulingQueue() : ready_queue_(std::make_shared<ReadyQueue>()) {
    for (const auto &task_state : {
             TaskState::PLACEABLE,
             TaskState::WAITING,
             TaskState::READY,
             TaskState::RUNNING,
             TaskState::INFEASIBLE,
             TaskState::WAITING_FOR_ACTOR_CREATION,
             TaskState::SWAP,
         }) {
      if (task_state == TaskState::READY) {
        task_queues_[static_cast<int>(task_state)] = ready_queue_;
      } else {
        task_queues_[static_cast<int>(task_state)] = std::make_shared<TaskQueue>();
      }
    }
  }

  // 我觉得名字起得不好，他的实际操作是获取所有就绪任务资源需求之和
  ResourceSet GetResourceLoad() const;
  
  /// 单个任务的增删查
  bool HasTask(const TaskID &task_id) const;
  const Task &GetTaskOfState(const TaskID &task_id, TaskState task_state) const;
  bool RemoveTask(const TaskID &task_id, Task *removed_task, TaskState *removed_task_state = nullptr);
  
  // 按状态获取任务，对于就绪状态，还需要按资源进行聚集
  const std::unordered_map<ResourceSet, ordered_set<TaskID>> &GetReadyTasksWithResources() const;
  const std::list<Task> &GetTasks(TaskState task_state) const;
  
  // 一组任务的移来移去
  std::vector<Task> RemoveTasks(std::unordered_set<TaskID> &task_ids);
  void QueueTasks(const std::vector<Task> &tasks, TaskState task_state);
  void MoveTasks(std::unordered_set<TaskID> &tasks, TaskState src_state, TaskState dst_state);
  void FilterState(std::unordered_set<TaskID> &task_ids, TaskState filter_state) const;
  
  /// 这两个函数是按其他维度：Actor 和 Job 来获取一组任务
  std::unordered_set<TaskID> GetTaskIdsForJob(const JobID &job_id) const;
  std::unordered_set<TaskID> GetTaskIdsForActor(const ActorID &actor_id) const;
  
  /// 阻塞任务和用户进程增删改查
  const std::unordered_set<TaskID> &GetBlockedTaskIds() const;
  const std::unordered_set<TaskID> &GetDriverTaskIds() const;
  void AddBlockedTaskId(const TaskID &task_id);
  void RemoveBlockedTaskId(const TaskID &task_id);
  void AddDriverTaskId(const TaskID &task_id);
  void RemoveDriverTaskId(const TaskID &task_id);

  /// 用来调试和监控
  std::string DebugString() const;
  void RecordMetrics() const;
  
  /// 这个好像没啥用，都没实现
  ResourceSet GetReadyQueueResources() const;

 private:
  /// 一个辅助函数，由于调度队列算是有两层索引 task state -> (task id -> task)，
  /// 因此经常需要定位到某个状态对应的任务队列，进而获取其中的某个任务。
  const std::shared_ptr<TaskQueue> &GetTaskQueue(TaskState task_state) const;

  /// 两个辅助函数，用来在指定状态的任务队列中删除或者过滤任务的
  void RemoveTasksFromQueue(ray::raylet::TaskState task_state,
                            std::unordered_set<ray::TaskID> &task_ids,
                            std::vector<ray::Task> *removed_tasks);
  void FilterStateFromQueue(std::unordered_set<ray::TaskID> &task_ids,
                            TaskState task_state) const;

  
  // kNumTaskQueues 作用便在此，所有int值在其之前的状态都有对应的任务队列
  std::array<std::shared_ptr<TaskQueue>, static_cast<int>(TaskState::kNumTaskQueues)>
      task_queues_;
  // 调度时候，就绪队列用的比较多，就单独维护一个指针在此
  const std::shared_ptr<ReadyQueue> ready_queue_;

  // 这两个状态（blocked 和 driver）没有对应的任务队列，只是用集合来保存id
  std::unordered_set<TaskID> blocked_task_ids_;
  std::unordered_set<TaskID> driver_task_ids_;
};