[core] Remove gcs_rpc_server.h dependency from GcsTaskManager (#56001)

Splitting gRPC service interface from implementation.

Stacked on: #55999

---------

Signed-off-by: Edward Oakes <[email protected]>
Signed-off-by: Douglas Strodtman <[email protected]>

Author: edoakes

src/ray/gcs/gcs_server/BUILD.bazel

@@ -199,14 +199,15 @@ ray_cc_library(
     hdrs = ["gcs_task_manager.h"],
     deps = [
         ":gcs_usage_stats_client",
+        ":grpc_service_interfaces",
         "//src/ray/common:asio",
         "//src/ray/common:id",
         "//src/ray/common:ray_config",
         "//src/ray/common:status",
         "//src/ray/gcs:gcs_pb_util",
         "//src/ray/protobuf:events_event_aggregator_service_cc_proto",
         "//src/ray/protobuf:gcs_cc_proto",
-        "//src/ray/rpc:gcs_server",
+        "//src/ray/stats:stats_metric",
         "//src/ray/util:counter_map",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",

src/ray/gcs/gcs_server/gcs_server.cc

@@ -735,10 +735,14 @@ void GcsServer::InitGcsTaskManager() {
   auto &io_context = io_context_provider_.GetIOContext<GcsTaskManager>();
   gcs_task_manager_ = std::make_unique<GcsTaskManager>(io_context);
   // Register service.
-  rpc_server_.RegisterService(
-      std::make_unique<rpc::TaskInfoGrpcService>(io_context, *gcs_task_manager_));
-  rpc_server_.RegisterService(
-      std::make_unique<rpc::RayEventExportGrpcService>(io_context, *gcs_task_manager_));
+  rpc_server_.RegisterService(std::make_unique<rpc::TaskInfoGrpcService>(
+      io_context,
+      *gcs_task_manager_,
+      RayConfig::instance().gcs_max_active_rpcs_per_handler()));
+  rpc_server_.RegisterService(std::make_unique<rpc::RayEventExportGrpcService>(
+      io_context,
+      *gcs_task_manager_,
+      RayConfig::instance().gcs_max_active_rpcs_per_handler()));
 }
 
 void GcsServer::InstallEventListeners() {

src/ray/gcs/gcs_server/gcs_task_manager.h

@@ -24,9 +24,10 @@
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/synchronization/mutex.h"
+#include "ray/gcs/gcs_server/grpc_service_interfaces.h"
 #include "ray/gcs/gcs_server/usage_stats_client.h"
 #include "ray/gcs/pb_util.h"
-#include "ray/rpc/gcs/gcs_rpc_server.h"
+#include "ray/stats/metric_defs.h"
 #include "ray/util/counter_map.h"
 #include "src/ray/protobuf/gcs.pb.h"
 
@@ -92,7 +93,8 @@ class FinishedTaskActorTaskGcPolicy : public TaskEventsGcPolicyInterface {
 ///
 /// This class has its own io_context and io_thread, that's separate from other GCS
 /// services. All handling of all rpc should be posted to the single thread it owns.
-class GcsTaskManager : public rpc::TaskInfoHandler, public rpc::RayEventExportHandler {
+class GcsTaskManager : public rpc::TaskInfoGcsServiceHandler,
+                       public rpc::RayEventExportGcsServiceHandler {
  public:
   /// Create a GcsTaskManager.
   explicit GcsTaskManager(instrumented_io_context &io_service);

src/ray/gcs/gcs_server/grpc_service_interfaces.h

@@ -206,5 +206,26 @@ class InternalKVGcsServiceHandler {
                                        SendReplyCallback send_reply_callback) = 0;
 };
 
+class TaskInfoGcsServiceHandler {
+ public:
+  virtual ~TaskInfoGcsServiceHandler() = default;
+
+  virtual void HandleAddTaskEventData(AddTaskEventDataRequest request,
+                                      AddTaskEventDataReply *reply,
+                                      SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleGetTaskEvents(GetTaskEventsRequest request,
+                                   GetTaskEventsReply *reply,
+                                   SendReplyCallback send_reply_callback) = 0;
+};
+
+class RayEventExportGcsServiceHandler {
+ public:
+  virtual ~RayEventExportGcsServiceHandler() = default;
+  virtual void HandleAddEvents(events::AddEventsRequest request,
+                               events::AddEventsReply *reply,
+                               SendReplyCallback send_reply_callback) = 0;
+};
+
 }  // namespace rpc
 }  // namespace ray

src/ray/gcs/gcs_server/grpc_services.cc

@@ -114,5 +114,23 @@ void InternalKVGrpcService::InitServerCallFactories(
       InternalKVGcsService, GetInternalConfig, max_active_rpcs_per_handler_)
 }
 
+void TaskInfoGrpcService::InitServerCallFactories(
+    const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
+    std::vector<std::unique_ptr<ServerCallFactory>> *server_call_factories,
+    const ClusterID &cluster_id) {
+  RPC_SERVICE_HANDLER(TaskInfoGcsService, AddTaskEventData, max_active_rpcs_per_handler_)
+  RPC_SERVICE_HANDLER(TaskInfoGcsService, GetTaskEvents, max_active_rpcs_per_handler_)
+}
+
+using events::AddEventsReply;
+using events::AddEventsRequest;
+
+void RayEventExportGrpcService::InitServerCallFactories(
+    const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
+    std::vector<std::unique_ptr<ServerCallFactory>> *server_call_factories,
+    const ClusterID &cluster_id) {
+  RPC_SERVICE_HANDLER(RayEventExportGcsService, AddEvents, max_active_rpcs_per_handler_)
+}
+
 }  // namespace rpc
 }  // namespace ray

src/ray/gcs/gcs_server/grpc_services.h

@@ -196,5 +196,51 @@ class InternalKVGrpcService : public GrpcService {
   int64_t max_active_rpcs_per_handler_;
 };
 
+class TaskInfoGrpcService : public GrpcService {
+ public:
+  explicit TaskInfoGrpcService(instrumented_io_context &io_service,
+                               TaskInfoGcsServiceHandler &handler,
+                               int64_t max_active_rpcs_per_handler)
+      : GrpcService(io_service),
+        service_handler_(handler),
+        max_active_rpcs_per_handler_(max_active_rpcs_per_handler){};
+
+ protected:
+  grpc::Service &GetGrpcService() override { return service_; }
+
+  void InitServerCallFactories(
+      const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
+      std::vector<std::unique_ptr<ServerCallFactory>> *server_call_factories,
+      const ClusterID &cluster_id) override;
+
+ private:
+  TaskInfoGcsService::AsyncService service_;
+  TaskInfoGcsServiceHandler &service_handler_;
+  int64_t max_active_rpcs_per_handler_;
+};
+
+class RayEventExportGrpcService : public GrpcService {
+ public:
+  explicit RayEventExportGrpcService(instrumented_io_context &io_service,
+                                     RayEventExportGcsServiceHandler &handler,
+                                     int64_t max_active_rpcs_per_handler)
+      : GrpcService(io_service),
+        service_handler_(handler),
+        max_active_rpcs_per_handler_(max_active_rpcs_per_handler){};
+
+ protected:
+  grpc::Service &GetGrpcService() override { return service_; }
+
+  void InitServerCallFactories(
+      const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
+      std::vector<std::unique_ptr<ServerCallFactory>> *server_call_factories,
+      const ClusterID &cluster_id) override;
+
+ private:
+  RayEventExportGcsService::AsyncService service_;
+  RayEventExportGcsServiceHandler &service_handler_;
+  int64_t max_active_rpcs_per_handler_;
+};
+
 }  // namespace rpc
 }  // namespace ray

src/ray/gcs/gcs_server/tests/BUILD.bazel

@@ -149,7 +149,6 @@ ray_cc_test(
     ],
     tags = ["team:core"],
     deps = [
-        ":gcs_server_test_util",
         "//:ray_mock",
         "//src/ray/gcs/gcs_server:gcs_task_manager",
         "//src/ray/gcs/tests:gcs_test_util_lib",

src/ray/rpc/gcs/gcs_rpc_server.h

@@ -117,16 +117,6 @@ namespace rpc {
                       HANDLER,               \
                       RayConfig::instance().gcs_max_active_rpcs_per_handler())
 
-#define TASK_INFO_SERVICE_RPC_HANDLER(HANDLER) \
-  RPC_SERVICE_HANDLER(TaskInfoGcsService,      \
-                      HANDLER,                 \
-                      RayConfig::instance().gcs_max_active_rpcs_per_handler())
-
-#define RAY_EVENT_EXPORT_SERVICE_RPC_HANDLER(HANDLER) \
-  RPC_SERVICE_HANDLER(RayEventExportGcsService,       \
-                      HANDLER,                        \
-                      RayConfig::instance().gcs_max_active_rpcs_per_handler())
-
 #define OBJECT_INFO_SERVICE_RPC_HANDLER(HANDLER) \
   RPC_SERVICE_HANDLER(ObjectInfoGcsService,      \
                       HANDLER,                   \
@@ -291,83 +281,8 @@ class PlacementGroupInfoGrpcService : public GrpcService {
   PlacementGroupInfoGcsServiceHandler &service_handler_;
 };
 
-class TaskInfoGcsServiceHandler {
- public:
-  virtual ~TaskInfoGcsServiceHandler() = default;
-
-  virtual void HandleAddTaskEventData(AddTaskEventDataRequest request,
-                                      AddTaskEventDataReply *reply,
-                                      SendReplyCallback send_reply_callback) = 0;
-
-  virtual void HandleGetTaskEvents(GetTaskEventsRequest request,
-                                   GetTaskEventsReply *reply,
-                                   SendReplyCallback send_reply_callback) = 0;
-};
-
-/// The `GrpcService` for `TaskInfoGcsService`.
-class TaskInfoGrpcService : public GrpcService {
- public:
-  /// Constructor.
-  ///
-  /// \param[in] io_service IO service to run the handler.
-  /// \param[in] handler The service handler that actually handle the requests.
-  explicit TaskInfoGrpcService(instrumented_io_context &io_service,
-                               TaskInfoGcsServiceHandler &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::unique_ptr<ServerCallFactory>> *server_call_factories,
-      const ClusterID &cluster_id) override {
-    TASK_INFO_SERVICE_RPC_HANDLER(AddTaskEventData);
-    TASK_INFO_SERVICE_RPC_HANDLER(GetTaskEvents);
-  }
-
- private:
-  /// The grpc async service object.
-  TaskInfoGcsService::AsyncService service_;
-  /// The service handler that actually handle the requests.
-  TaskInfoGcsServiceHandler &service_handler_;
-};
-
-class RayEventExportGcsServiceHandler {
- public:
-  virtual ~RayEventExportGcsServiceHandler() = default;
-  virtual void HandleAddEvents(AddEventsRequest request,
-                               AddEventsReply *reply,
-                               SendReplyCallback send_reply_callback) = 0;
-};
-
-/// The `GrpcService` for `RayEventExportGcsService`.
-class RayEventExportGrpcService : public GrpcService {
- public:
-  explicit RayEventExportGrpcService(instrumented_io_context &io_service,
-                                     RayEventExportGcsServiceHandler &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::unique_ptr<ServerCallFactory>> *server_call_factories,
-      const ClusterID &cluster_id) override {
-    RAY_EVENT_EXPORT_SERVICE_RPC_HANDLER(AddEvents);
-  }
-
- private:
-  /// The grpc async service object.
-  RayEventExportGcsService::AsyncService service_;
-  /// The service handler that actually handle the requests.
-  RayEventExportGcsServiceHandler &service_handler_;
-};
-
 using ActorInfoHandler = ActorInfoGcsServiceHandler;
 using PlacementGroupInfoHandler = PlacementGroupInfoGcsServiceHandler;
-using TaskInfoHandler = TaskInfoGcsServiceHandler;
-using RayEventExportHandler = RayEventExportGcsServiceHandler;
 
 }  // namespace rpc
 }  // namespace ray