Repository: ray-project/ray
File: src/ray/gcs/gcs_resource_manager.cc (C++)

Code before:
// limitations under the License.

#include "ray/gcs/gcs_resource_manager.h"

#include <memory>
#include <string>
#include <utility>

#include "ray/common/ray_config.h"
#include "ray/gcs/state_util.h"
#include "ray/util/logging.h"

namespace ray {
namespace gcs {

GcsResourceManager::GcsResourceManager(instrumented_io_context &io_context,
                                       ClusterResourceManager &cluster_resource_manager,
                                       GcsNodeManager &gcs_node_manager,
                                       NodeID local_node_id,
                                       raylet::ClusterLeaseManager *cluster_lease_manager)
    : io_context_(io_context),
      cluster_resource_manager_(cluster_resource_manager),
      gcs_node_manager_(gcs_node_manager),
      local_node_id_(std::move(local_node_id)),
      cluster_lease_manager_(cluster_lease_manager) {
  gcs_node_manager_.AddNodeDrainingListener(
      [this](const NodeID &node_id, bool is_draining, int64_t deadline_timestamp_ms) {
        scheduling::NodeID scheduling_node_id(node_id.Binary());
        cluster_resource_manager_.SetNodeDraining(
            scheduling_node_id, is_draining, deadline_timestamp_ms);
      },
      io_context_);
}

void GcsResourceManager::ConsumeSyncMessage(
    std::shared_ptr<const rpc::syncer::RaySyncMessage> message) {
  // ConsumeSyncMessage is called by ray_syncer which might not run
  // in a dedicated thread for performance.
  // GcsResourceManager is a module always run in the main thread, so we just
  // delegate the work to the main thread for thread safety.
  // Ideally, all public api in GcsResourceManager need to be put into this
  // io context for thread safety.
  io_context_.dispatch(
      [this, message]() {
        if (message->message_type() == syncer::MessageType::COMMANDS) {
          // COMMANDS channel is currently unused.
        } else if (message->message_type() == syncer::MessageType::RESOURCE_VIEW) {
          syncer::ResourceViewSyncMessage resource_view_sync_message;
          resource_view_sync_message.ParseFromString(message->sync_message());
          UpdateFromResourceView(NodeID::FromBinary(message->node_id()),
                                 resource_view_sync_message);

Review comment (suggestion):
So some nit here.  Today the code base is arranged such that all listeners are wired up in GcsServer::InstallEventListeners().  I don't have a predeliction towards this pattern vs. the one you've used here BUT I am a stickler for pattern consistency.  Could we move the listener wiring up to there instead to keep things consistent?

Also....  Why is this even in this class?  I'm not sure I see why.

Code after:
// limitations under the License.

#include "ray/gcs/gcs_resource_manager.h"

#include <memory>
#include <string>
#include <utility>

#include "ray/common/ray_config.h"
#include "ray/gcs/state_util.h"
#include "ray/util/logging.h"

namespace ray {
namespace gcs {

GcsResourceManager::GcsResourceManager(instrumented_io_context &io_context,
                                       ClusterResourceManager &cluster_resource_manager,
                                       GcsNodeManager &gcs_node_manager,
                                       NodeID local_node_id,
                                       raylet::ClusterLeaseManager *cluster_lease_manager)
    : io_context_(io_context),
      cluster_resource_manager_(cluster_resource_manager),
      gcs_node_manager_(gcs_node_manager),
      local_node_id_(std::move(local_node_id)),
      cluster_lease_manager_(cluster_lease_manager) {}

void GcsResourceManager::ConsumeSyncMessage(
    std::shared_ptr<const rpc::syncer::RaySyncMessage> message) {
  // ConsumeSyncMessage is called by ray_syncer which might not run
  // in a dedicated thread for performance.
  // GcsResourceManager is a module always run in the main thread, so we just
  // delegate the work to the main thread for thread safety.
  // Ideally, all public api in GcsResourceManager need to be put into this
  // io context for thread safety.
  io_context_.dispatch(
      [this, message]() {
        if (message->message_type() == syncer::MessageType::COMMANDS) {
          // COMMANDS channel is currently unused.
        } else if (message->message_type() == syncer::MessageType::RESOURCE_VIEW) {
          syncer::ResourceViewSyncMessage resource_view_sync_message;
          resource_view_sync_message.ParseFromString(message->sync_message());
          UpdateFromResourceView(NodeID::FromBinary(message->node_id()),
                                 resource_view_sync_message);
        } else {
          RAY_LOG(FATAL) << "Unsupported message type: " << message->message_type();
        }
      },
      "GcsResourceManager::Update");
}

void GcsResourceManager::HandleGetDrainingNodes(