Add stats handler to gcs server (#6750)

src/ray/gcs/gcs_server/gcs_server.cc

@@ -3,6 +3,7 @@
 #include "job_info_handler_impl.h"
 #include "node_info_handler_impl.h"
 #include "object_info_handler_impl.h"
+#include "stats_handler_impl.h"
 #include "task_info_handler_impl.h"
 
 namespace ray {
@@ -44,6 +45,10 @@ void GcsServer::Start() {
       new rpc::TaskInfoGrpcService(main_service_, *task_info_handler_));
   rpc_server_.RegisterService(*task_info_service_);
 
+  stats_handler_ = InitStatsHandler();
+  stats_service_.reset(new rpc::StatsGrpcService(main_service_, *stats_handler_));
+  rpc_server_.RegisterService(*stats_service_);
+
   // Run rpc server.
   rpc_server_.Run();
 
@@ -95,5 +100,10 @@ std::unique_ptr<rpc::TaskInfoHandler> GcsServer::InitTaskInfoHandler() {
       new rpc::DefaultTaskInfoHandler(*redis_gcs_client_));
 }
 
+std::unique_ptr<rpc::StatsHandler> GcsServer::InitStatsHandler() {
+  return std::unique_ptr<rpc::DefaultStatsHandler>(
+      new rpc::DefaultStatsHandler(*redis_gcs_client_));
+}
+
 }  // namespace gcs
 }  // namespace ray

src/ray/gcs/gcs_server/gcs_server.h

@@ -59,6 +59,9 @@ class GcsServer {
   /// The task info handler
   virtual std::unique_ptr<rpc::TaskInfoHandler> InitTaskInfoHandler();
 
+  /// The stats handler
+  virtual std::unique_ptr<rpc::StatsHandler> InitStatsHandler();
+
  private:
   /// Gcs server configuration
   GcsServerConfig config_;
@@ -81,6 +84,9 @@ class GcsServer {
   /// Task info handler and service
   std::unique_ptr<rpc::TaskInfoHandler> task_info_handler_;
   std::unique_ptr<rpc::TaskInfoGrpcService> task_info_service_;
+  /// Stats handler and service
+  std::unique_ptr<rpc::StatsHandler> stats_handler_;
+  std::unique_ptr<rpc::StatsGrpcService> stats_service_;
   /// Backend client
   std::shared_ptr<RedisGcsClient> redis_gcs_client_;
 };

src/ray/gcs/gcs_server/stats_handler_impl.cc

@@ -0,0 +1,32 @@
+#include "stats_handler_impl.h"
+
+namespace ray {
+namespace rpc {
+
+void DefaultStatsHandler::HandleAddProfileData(const AddProfileDataRequest &request,
+                                               AddProfileDataReply *reply,
+                                               SendReplyCallback send_reply_callback) {
+  ClientID node_id = ClientID::FromBinary(request.profile_data().component_id());
+  RAY_LOG(DEBUG) << "Adding profile data, component type = "
+                 << request.profile_data().component_type() << ", node id = " << node_id;
+  auto profile_table_data = std::make_shared<ProfileTableData>();
+  profile_table_data->CopyFrom(request.profile_data());
+  auto on_done = [node_id, request, send_reply_callback](Status status) {
+    if (!status.ok()) {
+      RAY_LOG(ERROR) << "Failed to add profile data, component type = "
+                     << request.profile_data().component_type()
+                     << ", node id = " << node_id;
+    }
+    send_reply_callback(status, nullptr, nullptr);
+  };
+
+  Status status = gcs_client_.Stats().AsyncAddProfileData(profile_table_data, on_done);
+  if (!status.ok()) {
+    on_done(status);
+  }
+  RAY_LOG(DEBUG) << "Finished adding profile data, component type = "
+                 << request.profile_data().component_type() << ", node id = " << node_id;
+}
+
+}  // namespace rpc
+}  // namespace ray

src/ray/gcs/gcs_server/stats_handler_impl.h

@@ -0,0 +1,27 @@
+#ifndef RAY_GCS_STATS_HANDLER_IMPL_H
+#define RAY_GCS_STATS_HANDLER_IMPL_H
+
+#include "ray/gcs/redis_gcs_client.h"
+#include "ray/rpc/gcs_server/gcs_rpc_server.h"
+
+namespace ray {
+namespace rpc {
+
+/// This implementation class of `StatsHandler`.
+class DefaultStatsHandler : public rpc::StatsHandler {
+ public:
+  explicit DefaultStatsHandler(gcs::RedisGcsClient &gcs_client)
+      : gcs_client_(gcs_client) {}
+
+  void HandleAddProfileData(const AddProfileDataRequest &request,
+                            AddProfileDataReply *reply,
+                            SendReplyCallback send_reply_callback) override;
+
+ private:
+  gcs::RedisGcsClient &gcs_client_;
+};
+
+}  // namespace rpc
+}  // namespace ray
+
+#endif  // RAY_GCS_STATS_HANDLER_IMPL_H

src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc

@@ -345,6 +345,16 @@ class GcsServerTest : public RedisServiceManagerForTest {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
+  bool AddProfileData(const rpc::AddProfileDataRequest &request) {
+    std::promise<bool> promise;
+    client_->AddProfileData(
+        request, [&promise](const Status &status, const rpc::AddProfileDataReply &reply) {
+          RAY_CHECK_OK(status);
+          promise.set_value(true);
+        });
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   bool WaitReady(const std::future<bool> &future, uint64_t timeout_ms) {
     auto status = future.wait_for(std::chrono::milliseconds(timeout_ms));
     return status == std::future_status::ready;
@@ -595,6 +605,14 @@ TEST_F(GcsServerTest, TestTaskInfo) {
   ASSERT_TRUE(AttemptTaskReconstruction(attempt_task_reconstruction_request));
 }
 
+TEST_F(GcsServerTest, TestStats) {
+  rpc::ProfileTableData profile_table_data;
+  profile_table_data.set_component_id(ClientID::FromRandom().Binary());
+  rpc::AddProfileDataRequest add_profile_data_request;
+  add_profile_data_request.mutable_profile_data()->CopyFrom(profile_table_data);
+  ASSERT_TRUE(AddProfileData(add_profile_data_request));
+}
+
 }  // namespace ray
 
 int main(int argc, char **argv) {

src/ray/protobuf/gcs_service.proto

@@ -22,9 +22,9 @@ message MarkJobFinishedReply {
 
 // Service for job info access.
 service JobInfoGcsService {
-  // Add job to gcs server.
+  // Add job to GCS Service.
   rpc AddJob(AddJobRequest) returns (AddJobReply);
-  // Mark job as finished to gcs server.
+  // Mark job as finished to GCS Service.
   rpc MarkJobFinished(MarkJobFinishedRequest) returns (MarkJobFinishedReply);
 }
 
@@ -269,3 +269,16 @@ service TaskInfoGcsService {
   rpc AttemptTaskReconstruction(AttemptTaskReconstructionRequest)
       returns (AttemptTaskReconstructionReply);
 }
+
+message AddProfileDataRequest {
+  ProfileTableData profile_data = 1;
+}
+
+message AddProfileDataReply {
+}
+
+// Service for stats access.
+service StatsGcsService {
+  // Add profile data to GCS Service.
+  rpc AddProfileData(AddProfileDataRequest) returns (AddProfileDataReply);
+}

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -26,6 +26,8 @@ class GcsRpcClient {
         new GrpcClient<ObjectInfoGcsService>(address, port, client_call_manager));
     task_info_grpc_client_ = std::unique_ptr<GrpcClient<TaskInfoGcsService>>(
         new GrpcClient<TaskInfoGcsService>(address, port, client_call_manager));
+    stats_grpc_client_ = std::unique_ptr<GrpcClient<StatsGcsService>>(
+        new GrpcClient<StatsGcsService>(address, port, client_call_manager));
   };
 
   /// Add job info to gcs server.
@@ -109,6 +111,9 @@ class GcsRpcClient {
   VOID_RPC_CLIENT_METHOD(TaskInfoGcsService, AttemptTaskReconstruction,
                          task_info_grpc_client_, )
 
+  /// Add profile data to GCS Service.
+  VOID_RPC_CLIENT_METHOD(StatsGcsService, AddProfileData, stats_grpc_client_, )
+
  private:
   /// The gRPC-generated stub.
   std::unique_ptr<GrpcClient<JobInfoGcsService>> job_info_grpc_client_;
@@ -116,6 +121,7 @@ class GcsRpcClient {
   std::unique_ptr<GrpcClient<NodeInfoGcsService>> node_info_grpc_client_;
   std::unique_ptr<GrpcClient<ObjectInfoGcsService>> object_info_grpc_client_;
   std::unique_ptr<GrpcClient<TaskInfoGcsService>> task_info_grpc_client_;
+  std::unique_ptr<GrpcClient<StatsGcsService>> stats_grpc_client_;
 };
 
 }  // namespace rpc

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -24,6 +24,9 @@ namespace rpc {
 #define TASK_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
   RPC_SERVICE_HANDLER(TaskInfoGcsService, HANDLER, CONCURRENCY)
 
+#define STATS_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
+  RPC_SERVICE_HANDLER(StatsGcsService, HANDLER, CONCURRENCY)
+
 class JobInfoGcsServiceHandler {
  public:
   virtual ~JobInfoGcsServiceHandler() = default;
@@ -296,11 +299,48 @@ class TaskInfoGrpcService : public GrpcService {
   TaskInfoGcsServiceHandler &service_handler_;
 };
 
+class StatsGcsServiceHandler {
+ public:
+  virtual ~StatsGcsServiceHandler() = default;
+
+  virtual void HandleAddProfileData(const AddProfileDataRequest &request,
+                                    AddProfileDataReply *reply,
+                                    SendReplyCallback send_reply_callback) = 0;
+};
+
+/// The `GrpcService` for `StatsGcsService`.
+class StatsGrpcService : public GrpcService {
+ public:
+  /// Constructor.
+  ///
+  /// \param[in] handler The service handler that actually handle the requests.
+  explicit StatsGrpcService(boost::asio::io_service &io_service,
+                            StatsGcsServiceHandler &handler)
+      : GrpcService(io_service), service_handler_(handler){};
+
+ protected:
+  grpc::Service &GetGrpcService() override { return service_; }
+
+  void InitServerCallFactories(
+      const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
+      std::vector<std::pair<std::unique_ptr<ServerCallFactory>, int>>
+          *server_call_factories_and_concurrencies) override {
+    STATS_SERVICE_RPC_HANDLER(AddProfileData, 1);
+  }
+
+ private:
+  /// The grpc async service object.
+  StatsGcsService::AsyncService service_;
+  /// The service handler that actually handle the requests.
+  StatsGcsServiceHandler &service_handler_;
+};
+
 using JobInfoHandler = JobInfoGcsServiceHandler;
 using ActorInfoHandler = ActorInfoGcsServiceHandler;
 using NodeInfoHandler = NodeInfoGcsServiceHandler;
 using ObjectInfoHandler = ObjectInfoGcsServiceHandler;
 using TaskInfoHandler = TaskInfoGcsServiceHandler;
+using StatsHandler = StatsGcsServiceHandler;
 
 }  // namespace rpc
 }  // namespace ray