[Core] Add gRPC streaming support by moving to the async callback/reactor API.

src/ray/core_worker/core_worker.cc

@@ -104,7 +104,7 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
   // so that the worker (java/python .etc) can retrieve and handle the error
   // instead of crashing.
   auto grpc_client = rpc::NodeManagerWorkerClient::make(
-      options_.raylet_ip_address, options_.node_manager_port, *client_call_manager_);
+      options_.raylet_ip_address, options_.node_manager_port, io_service_);
   Status raylet_client_status;
   NodeID local_raylet_id;
   int assigned_port;
@@ -314,8 +314,7 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
   }
 
   auto raylet_client_factory = [this](const std::string ip_address, int port) {
-    auto grpc_client =
-        rpc::NodeManagerWorkerClient::make(ip_address, port, *client_call_manager_);
+    auto grpc_client = rpc::NodeManagerWorkerClient::make(ip_address, port, io_service_);
     return std::shared_ptr<raylet::RayletClient>(
         new raylet::RayletClient(std::move(grpc_client)));
   };
@@ -1319,8 +1318,7 @@ void CoreWorker::SpillOwnedObject(const ObjectID &object_id,
   RAY_LOG(DEBUG) << "Sending spill request to raylet for object " << object_id;
   auto raylet_client =
       std::make_shared<raylet::RayletClient>(rpc::NodeManagerWorkerClient::make(
-          node->node_manager_address(), node->node_manager_port(),
-          *client_call_manager_));
+          node->node_manager_address(), node->node_manager_port(), io_service_));
   raylet_client->RequestObjectSpillage(
       object_id, [object_id, callback](const Status &status,
                                        const rpc::RequestObjectSpillageReply &reply) {

src/ray/core_worker/core_worker_process.cc

@@ -183,9 +183,8 @@ void CoreWorkerProcessImpl::InitializeSystemConfig() {
   std::thread thread([&] {
     instrumented_io_context io_service;
     boost::asio::io_service::work work(io_service);
-    rpc::ClientCallManager client_call_manager(io_service);
     auto grpc_client = rpc::NodeManagerWorkerClient::make(
-        options_.raylet_ip_address, options_.node_manager_port, client_call_manager);
+        options_.raylet_ip_address, options_.node_manager_port, io_service);
     raylet::RayletClient raylet_client(grpc_client);
 
     std::function<void(int64_t)> get_once = [this, &get_once, &raylet_client, &promise,

src/ray/core_worker/gcs_server_address_updater.cc

@@ -20,9 +20,8 @@ namespace core {
 GcsServerAddressUpdater::GcsServerAddressUpdater(
     const std::string raylet_ip_address, const int port,
     std::function<void(std::string, int)> update_func)
-    : client_call_manager_(updater_io_service_),
-      raylet_client_(rpc::NodeManagerWorkerClient::make(raylet_ip_address, port,
-                                                        client_call_manager_)),
+    : raylet_client_(rpc::NodeManagerWorkerClient::make(raylet_ip_address, port,
+                                                        updater_io_service_)),
       update_func_(update_func),
       updater_runner_(updater_io_service_),
       updater_thread_([this] {

src/ray/core_worker/gcs_server_address_updater.h

@@ -37,7 +37,6 @@ class GcsServerAddressUpdater {
   /// Update gcs server address.
   void UpdateGcsServerAddress();
 
-  rpc::ClientCallManager client_call_manager_;
   /// A client connection to the raylet.
   raylet::RayletClient raylet_client_;
   std::function<void(std::string, int)> update_func_;

src/ray/gcs/gcs_server/gcs_server.cc

@@ -41,8 +41,7 @@ GcsServer::GcsServer(const ray::gcs::GcsServerConfig &config,
                   /*keepalive_time_ms=*/RayConfig::instance().grpc_keepalive_time_ms()),
       client_call_manager_(main_service,
                            RayConfig::instance().gcs_server_rpc_client_thread_num()),
-      raylet_client_pool_(
-          std::make_shared<rpc::NodeManagerClientPool>(client_call_manager_)),
+      raylet_client_pool_(std::make_shared<rpc::NodeManagerClientPool>(main_service)),
       pubsub_periodical_runner_(pubsub_io_service_),
       periodical_runner_(main_service),
       is_started_(false),

src/ray/raylet/node_manager.cc

@@ -2423,7 +2423,7 @@ void NodeManager::HandleFormatGlobalMemoryInfo(
   // Fetch from remote nodes.
   for (const auto &entry : remote_node_manager_addresses_) {
     auto client = std::make_unique<rpc::NodeManagerClient>(
-        entry.second.first, entry.second.second, client_call_manager_);
+        entry.second.first, entry.second.second, io_service_);
     client->GetNodeStats(
         stats_req, [replies, store_reply](const ray::Status &status,
                                           const rpc::GetNodeStatsReply &r) {

src/ray/raylet/node_manager.h

@@ -662,14 +662,14 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   std::shared_ptr<AgentManager> agent_manager_;
 
   /// The RPC server.
-  rpc::GrpcServer node_manager_server_;
+  rpc::GrpcCallbackServer node_manager_server_;
 
   /// The node manager RPC service.
-  rpc::NodeManagerGrpcService node_manager_service_;
+  rpc::NodeManagerServiceWithCallbacks node_manager_service_;
 
   /// The agent manager RPC service.
   std::unique_ptr<rpc::AgentManagerServiceHandler> agent_manager_service_handler_;
-  rpc::AgentManagerGrpcService agent_manager_service_;
+  rpc::AgentManagerServiceWithCallbacks agent_manager_service_;
 
   /// Manages all local objects that are pinned (primary
   /// copies), freed, and/or spilled.

src/ray/rpc/agent_manager/agent_manager_server.h

@@ -15,17 +15,13 @@
 #pragma once
 
 #include "ray/common/asio/instrumented_io_context.h"
-#include "ray/rpc/grpc_server.h"
-#include "ray/rpc/server_call.h"
+#include "ray/rpc/grpc_callback_server.h"
 #include "src/ray/protobuf/agent_manager.grpc.pb.h"
 #include "src/ray/protobuf/agent_manager.pb.h"
 
 namespace ray {
 namespace rpc {
 
-#define RAY_AGENT_MANAGER_RPC_HANDLERS \
-  RPC_SERVICE_HANDLER(AgentManagerService, RegisterAgent, -1)
-
 /// Implementations of the `AgentManagerGrpcService`, check interface in
 /// `src/ray/protobuf/agent_manager.proto`.
 class AgentManagerServiceHandler {
@@ -43,32 +39,10 @@ class AgentManagerServiceHandler {
                                    SendReplyCallback send_reply_callback) = 0;
 };
 
-/// The `GrpcService` for `AgentManagerGrpcService`.
-class AgentManagerGrpcService : public GrpcService {
- public:
-  /// Construct a `AgentManagerGrpcService`.
-  ///
-  /// \param[in] port See `GrpcService`.
-  /// \param[in] handler The service handler that actually handle the requests.
-  AgentManagerGrpcService(instrumented_io_context &io_service,
-                          AgentManagerServiceHandler &service_handler)
-      : GrpcService(io_service), service_handler_(service_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) override {
-    RAY_AGENT_MANAGER_RPC_HANDLERS
-  }
+#define RAY_AGENT_MANAGER_RPC_HANDLERS \
+  UNARY_CALLBACK_RPC_SERVICE_HANDLER(AgentManagerService, RegisterAgent)
 
- private:
-  /// The grpc async service object.
-  AgentManagerService::AsyncService service_;
-  /// The service handler that actually handle the requests.
-  AgentManagerServiceHandler &service_handler_;
-};
+CALLBACK_SERVICE(AgentManagerService, RAY_AGENT_MANAGER_RPC_HANDLERS)
 
 }  // namespace rpc
 }  // namespace ray