PS Lite - 源码解读

PostOffice 类

/**
 * \brief 系统的中心。
 */
class Postoffice {
 public:
  /**
   * \brief 返回单例对象。
   */
  static Postoffice* Get() {
    static Postoffice e; return &e;
  }

  /** \brief 返回持有的 Van 实例。 */
  Van* van() { return van_; }

  /**
   * \brief 启动系统。
   *
   * 本函数将会阻塞，直到所有节点都启动。
   * \param argv0 用于日志的程序名。
   * \param do_barrier 是否在所有节点都启动完成前阻塞当前线程。
   */
  void Start(int customer_id, const char* argv0, const bool do_barrier);

  /**
   * \brief 停止系统。
   *
   * 所有节点在退出前都应该调用该方法。
   * \param do_barrier 是否在所有节点都停止前阻塞当前线程，默认为 true。
   */
  void Finalize(const int customer_id, const bool do_barrier = true);

  /**
   * \brief 向本系统添加一个 Customer。线程安全。
   */
  void AddCustomer(Customer* customer);

  /**
   * \brief 根据编号移除一个 Customer。线程安全。
   */
  void RemoveCustomer(Customer* customer);

  /**
   * \brief 获取给定编号的 Customer。线程安全。
   * \param app_id 应用编号。
   * \param customer_id Customer 编号。
   * \param timeout 以秒位单位的超时。
   * \return 如果不存在或者超时，返回 nullptr。
   */
  Customer* GetCustomer(int app_id, int customer_id, int timeout = 0) const;

  /**
   * \brief 获取给定编号的节点或节点组。线程安全。
   *
   * 如果是一个节点组的编号，那么返回所属的全部节点。否则，返回的列表只有一个节点元素。
   */
  const std::vector<int>& GetNodeIDs(int node_id) const {
    const auto it = node_ids_.find(node_id);
    CHECK(it != node_ids_.cend()) << "node " << node_id << " doesn't exist";
    return it->second;
  }

  /**
   * \brief 返回每个 Server 节点对应的键域。
   */
  const std::vector<Range>& GetServerKeyRanges();

  /**
   * \brief 回调函数的别名。
   */
  using Callback = std::function<void()>;

  /**
   * \brief 向系统注册一个回调函数，在 Finalize() 方法完成后被调用。
   *
   * 以下代码是等价的：
   * \code {cpp}
   * RegisterExitCallback(cb);
   * Finalize();
   * \endcode
   *
   * \code {cpp}
   * Finalize();
   * cb();
   * \endcode
   * \param cb 回调函数。
   */
  void RegisterExitCallback(const Callback& cb) {
    exit_callback_ = cb;
  }

  /**
   * \brief 从 Worker 编号映射到节点编号。
   * \param rank the worker rank
   */
  static inline int WorkerRankToID(int rank) {
    return rank * 2 + 9;
  }

  /**
   * \brief 从 Server 编号映射到节点编号。
   * \param rank the server rank
   */
  static inline int ServerRankToID(int rank) {
    return rank * 2 + 8;
  }

  /**
   * \brief 从节点编号映射到 Server/Worker 编号。
   * \param id 节点编号。
   */
  static inline int IDtoRank(int id) {
    return std::max((id - 8) / 2, 0);
  }

  /** \brief 返回 Worker 数量。 */
  int num_workers() const { return num_workers_; }

  /** \brief 返回 Server 数量。 */
  int num_servers() const { return num_servers_; }

  /** \brief 返回节点在所属分组的编号。
   *
   * 每个 Worker 都有唯一的编号，范围是 [0, NumWorkers())。Server 也是如此。
   * 该方法仅当 Start() 方法被调用后有效。
   */
  int my_rank() const { return IDtoRank(van_->my_node().id); }

  /** \brief 是否是 Worker 节点。 */
  int is_worker() const { return is_worker_; }

  /** \brief 是否是 Server 节点。 */
  int is_server() const { return is_server_; }

  /** \brief 是否是 Scheduler 节点。 */
  int is_scheduler() const { return is_scheduler_; }

  /** \brief 返回日志级别。 */
  int verbose() const { return verbose_; }

  /** \brief 是否是可恢复的节点。 */
  bool is_recovery() const { return van_->my_node().is_recovery; }

  /**
   * \brief 屏障。
   * \param node_id 要同步的节点组的编号。
   */
  void Barrier(int customer_id, int node_group);

  /**
   * \brief 处理控制信息，由持有的 Van 实例调用。
   * \param recv 收到的消息。
   */
  void Manage(const Message& recv);

  /**
   * \brief 更新心跳记录。
   * \param node_id 节点编号。
   * \param t 心跳时间戳。
   */
  void UpdateHeartbeat(int node_id, time_t t) {
    std::lock_guard<std::mutex> lk(heartbeat_mu_);
    heartbeats_[node_id] = t;
  }

  /**
   * \brief 获取在过去 t 秒内未发送心跳的所有节点。
   * \param t 以秒为单位的超时。
   */
  std::vector<int> GetDeadNodes(int t = 60);
}

Van 类

/**
 * \brief Van 类负责发送消息到远端节点。
 *
 * 如果环境变量 PS_RESEND 被设置位 1，那么在 PS_RESEND_TIMEOUT 毫秒后没有收到 ACK 消息的情况下，Van 实例会重发消息。
 */
class Van {
 public:
    /**
     * \brief 实例化 Van 类。
     * \param type zmq、socket等。
     */
    static Van *Create(const std::string &type);

    /** \brief 空的构造器。使用 Start() 方法来真的启动。
    Van() {}

    /**\brief 空的析构器。使用 Stop() 方法类真的停止。
    virtual ~Van() {}

    /**
     * \brief 启动 Van。
     *
     * 必须先调用该方法，才能调用 Send() 方法。
     *
     * 该方法初始化到所有节点的连接，启动接受消息的线程。
     * 如果收到了控制信息，交给 PostOffice::Manage() 方法处理。否则，交给相应的应用处理。
     */
    virtual void Start(int customer_id);

    /**
     * \brief 发送一个消息。线程安全。
     * \return 返回发送的字节数。如果发送失败，返回 -1。
     */
    int Send(const Message &msg);

    /**
     * \brief 返回所在的节点。
     */
    inline const Node &my_node() const {
      CHECK(ready_) << "call Start() first";
      return my_node_;
    }

    /**
     * \brief 停止 Van。
     * 停止接受消息的线程。
     */
    virtual void Stop();

    /**
     * \brief 获取下一个可用的时间戳。线程安全。
     */
    inline int GetTimestamp() { return timestamp_++; }

    /**
     * \brief 是否可以发送消息。线程安全。
     */
    inline bool IsReady() { return ready_; }
}