Add gcs server as well as the unit test (#6401)

BUILD.bazel

@@ -94,6 +94,17 @@ python_grpc_compile(
     deps = [":node_manager_proto"],
 )
 
+proto_library(
+    name = "gcs_service_proto",
+    srcs = ["src/ray/protobuf/gcs_service.proto"],
+    deps = [":gcs_proto"],
+)
+
+cc_proto_library(
+    name = "gcs_service_cc_proto",
+    deps = [":gcs_service_proto"],
+)
+
 proto_library(
     name = "object_manager_proto",
     srcs = ["src/ray/protobuf/object_manager.proto"],
@@ -171,6 +182,31 @@ cc_library(
     ],
 )
 
+# gcs_service gRPC lib.
+cc_grpc_library(
+    name = "gcs_service_cc_grpc",
+    srcs = [":gcs_service_proto"],
+    grpc_only = True,
+    deps = [":gcs_service_cc_proto"],
+)
+
+# gcs rpc server and client.
+cc_library(
+    name = "gcs_service_rpc",
+    hdrs = glob([
+        "src/ray/rpc/gcs_server/gcs_rpc_server.h",
+        "src/ray/rpc/gcs_server/gcs_rpc_client.h",
+    ]),
+    copts = COPTS,
+    deps = [
+        ":gcs_service_cc_grpc",
+        ":grpc_common_lib",
+        ":ray_common",
+        "@boost//:asio",
+        "@com_github_grpc_grpc//:grpc++",
+    ],
+)
+
 # Object manager gRPC lib.
 cc_grpc_library(
     name = "object_manager_cc_grpc",
@@ -282,6 +318,41 @@ cc_binary(
     ],
 )
 
+cc_library(
+    name = "gcs_server_lib",
+    srcs = glob(
+        [
+            "src/ray/gcs/gcs_server/*.cc",
+        ],
+        exclude = [
+            "src/ray/gcs/gcs_server/gcs_server_main.cc",
+            "src/ray/gcs/gcs_server/test/*.cc",
+        ],
+    ),
+    hdrs = glob(
+        [
+            "src/ray/gcs/gcs_server/*.h",
+        ],
+    ),
+    copts = COPTS,
+    deps = [
+        ":gcs",
+        ":gcs_service_rpc",
+    ],
+)
+
+cc_binary(
+    name = "gcs_server",
+    srcs = [
+        "src/ray/gcs/gcs_server/gcs_server_main.cc",
+    ],
+    copts = COPTS,
+    deps = [
+        ":gcs_server_lib",
+        "@com_github_gflags_gflags//:gflags",
+    ],
+)
+
 cc_library(
     name = "stats_lib",
     srcs = glob(
@@ -622,6 +693,16 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "gcs_server_rpc_test",
+    srcs = ["src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc"],
+    copts = COPTS,
+    deps = [
+        ":gcs_server_lib",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "object_manager",
     srcs = glob([

src/ray/gcs/gcs_server/gcs_server.cc

@@ -0,0 +1,54 @@
+#include "gcs_server.h"
+#include "job_info_handler_impl.h"
+
+namespace ray {
+namespace gcs {
+
+GcsServer::GcsServer(const ray::gcs::GcsServerConfig &config)
+    : config_(config),
+      rpc_server_(config.grpc_server_name, config.grpc_server_port,
+                  config.grpc_server_thread_num) {}
+
+GcsServer::~GcsServer() { Stop(); }
+
+void GcsServer::Start() {
+  // Init backend client.
+  InitBackendClient();
+
+  // Register rpc service.
+  job_info_handler_ = InitJobInfoHandler();
+  job_info_service_.reset(new rpc::JobInfoGrpcService(main_service_, *job_info_handler_));
+  rpc_server_.RegisterService(*job_info_service_);
+
+  // Run rpc server.
+  rpc_server_.Run();
+
+  // Run the event loop.
+  // Using boost::asio::io_context::work to avoid ending the event loop when
+  // there are no events to handle.
+  boost::asio::io_context::work worker(main_service_);
+  main_service_.run();
+}
+
+void GcsServer::Stop() {
+  // Shutdown the rpc server
+  rpc_server_.Shutdown();
+
+  // Stop the event loop.
+  main_service_.stop();
+}
+
+void GcsServer::InitBackendClient() {
+  GcsClientOptions options(config_.redis_address, config_.redis_port,
+                           config_.redis_password);
+  redis_gcs_client_ = std::make_shared<RedisGcsClient>(options);
+  auto status = redis_gcs_client_->Connect(main_service_);
+  RAY_CHECK(status.ok()) << "Failed to init redis gcs client as " << status;
+}
+
+std::unique_ptr<rpc::JobInfoHandler> GcsServer::InitJobInfoHandler() {
+  return std::unique_ptr<rpc::DefaultJobInfoHandler>(new rpc::DefaultJobInfoHandler());
+}
+
+}  // namespace gcs
+}  // namespace ray

src/ray/gcs/gcs_server/gcs_server.h

@@ -0,0 +1,68 @@
+#ifndef RAY_GCS_GCS_SERVER_H
+#define RAY_GCS_GCS_SERVER_H
+
+#include <ray/gcs/redis_gcs_client.h>
+#include <ray/rpc/gcs_server/gcs_rpc_server.h>
+#include <ray/rpc/grpc_server.h>
+#include <string>
+
+namespace ray {
+namespace gcs {
+
+struct GcsServerConfig {
+  std::string grpc_server_name = "GcsServer";
+  uint16_t grpc_server_port = 0;
+  uint16_t grpc_server_thread_num = 1;
+  std::string redis_password;
+  std::string redis_address;
+  uint16_t redis_port = 6379;
+  bool retry_redis = true;
+};
+
+/// The GcsServer will take over all requests from ServiceBasedGcsClient and transparent
+/// transmit the command to the backend reliable storage for the time being.
+/// In the future, GCS server's main responsibility is to manage meta data
+/// and the management of actor creation.
+/// For more details, please see the design document.
+/// https://docs.google.com/document/d/1d-9qBlsh2UQHo-AWMWR0GptI_Ajwu4SKx0Q0LHKPpeI/edit#heading=h.csi0gaglj2pv
+class GcsServer {
+ public:
+  explicit GcsServer(const GcsServerConfig &config);
+  virtual ~GcsServer();
+
+  /// Start gcs server.
+  void Start();
+
+  /// Stop gcs server.
+  void Stop();
+
+  /// Get the port of this gcs server.
+  int GetPort() const { return rpc_server_.GetPort(); }
+
+ protected:
+  /// Initialize the backend storage client
+  /// The gcs server is just the proxy between the gcs client and reliable storage
+  /// for the time being, so we need a backend client to connect to the storage.
+  virtual void InitBackendClient();
+
+  /// The job info handler
+  virtual std::unique_ptr<rpc::JobInfoHandler> InitJobInfoHandler();
+
+ private:
+  /// Gcs server configuration
+  GcsServerConfig config_;
+  /// The grpc server
+  rpc::GrpcServer rpc_server_;
+  /// The main io service to drive event posted from grpc threads.
+  boost::asio::io_context main_service_;
+  /// Job info handler and service
+  std::unique_ptr<rpc::JobInfoHandler> job_info_handler_;
+  std::unique_ptr<rpc::JobInfoGrpcService> job_info_service_;
+  /// Backend client
+  std::shared_ptr<RedisGcsClient> redis_gcs_client_;
+};
+
+}  // namespace gcs
+}  // namespace ray
+
+#endif

src/ray/gcs/gcs_server/gcs_server_main.cc

@@ -0,0 +1,53 @@
+#include <iostream>
+
+#include "ray/common/ray_config.h"
+#include "ray/gcs/gcs_server/gcs_server.h"
+#include "ray/util/util.h"
+
+#include "gflags/gflags.h"
+
+DEFINE_string(redis_address, "", "The ip address of redis.");
+DEFINE_int32(redis_port, -1, "The port of redis.");
+DEFINE_string(config_list, "", "The config list of raylet.");
+DEFINE_string(redis_password, "", "The password of redis.");
+DEFINE_bool(retry_redis, false, "Whether we retry to connect to the redis.");
+
+int main(int argc, char *argv[]) {
+  InitShutdownRAII ray_log_shutdown_raii(ray::RayLog::StartRayLog,
+                                         ray::RayLog::ShutDownRayLog, argv[0],
+                                         ray::RayLogLevel::INFO, /*log_dir=*/"");
+  ray::RayLog::InstallFailureSignalHandler();
+
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+  const std::string redis_address = FLAGS_redis_address;
+  const int redis_port = static_cast<int>(FLAGS_redis_port);
+  const std::string config_list = FLAGS_config_list;
+  const std::string redis_password = FLAGS_redis_password;
+  const bool retry_redis = FLAGS_retry_redis;
+  gflags::ShutDownCommandLineFlags();
+
+  std::unordered_map<std::string, std::string> config_map;
+
+  // Parse the configuration list.
+  std::istringstream config_string(config_list);
+  std::string config_name;
+  std::string config_value;
+
+  while (std::getline(config_string, config_name, ',')) {
+    RAY_CHECK(std::getline(config_string, config_value, ','));
+    config_map[config_name] = config_value;
+  }
+
+  RayConfig::instance().initialize(config_map);
+
+  ray::gcs::GcsServerConfig gcs_server_config;
+  gcs_server_config.grpc_server_name = "GcsServer";
+  gcs_server_config.grpc_server_port = 0;
+  gcs_server_config.grpc_server_thread_num = 1;
+  gcs_server_config.redis_address = redis_address;
+  gcs_server_config.redis_port = redis_port;
+  gcs_server_config.redis_password = redis_password;
+  gcs_server_config.retry_redis = retry_redis;
+  ray::gcs::GcsServer gcs_server(gcs_server_config);
+  gcs_server.Start();
+}

src/ray/gcs/gcs_server/job_info_handler_impl.cc

@@ -0,0 +1,19 @@
+#include "job_info_handler_impl.h"
+
+namespace ray {
+namespace rpc {
+void DefaultJobInfoHandler::HandleAddJob(const rpc::AddJobRequest &request,
+                                         rpc::AddJobReply *reply,
+                                         rpc::SendReplyCallback send_reply_callback) {
+  RAY_LOG(DEBUG) << "Received new job ...";
+  // TODO(zsl): The detailed implementation will be committed in next PR.
+}
+
+void DefaultJobInfoHandler::HandleMarkJobFinished(
+    const rpc::MarkJobFinishedRequest &request, rpc::MarkJobFinishedReply *reply,
+    rpc::SendReplyCallback send_reply_callback) {
+  RAY_LOG(DEBUG) << "Mark job as finished ...";
+  // TODO(zsl): The detailed implementation will be committed in next PR.
+}
+}  // namespace rpc
+}  // namespace ray

src/ray/gcs/gcs_server/job_info_handler_impl.h

@@ -0,0 +1,24 @@
+#ifndef RAY_GCS_JOB_INFO_HANDLER_IMPL_H
+#define RAY_GCS_JOB_INFO_HANDLER_IMPL_H
+
+#include "ray/rpc/gcs_server/gcs_rpc_server.h"
+
+namespace ray {
+namespace rpc {
+
+/// This class is used to implement `JobInfoHandler`, but only two logs have been printed.
+/// The detailed implementation is reflected in the following PR.
+class DefaultJobInfoHandler : public rpc::JobInfoHandler {
+ public:
+  void HandleAddJob(const AddJobRequest &request, AddJobReply *reply,
+                    SendReplyCallback send_reply_callback) override;
+
+  void HandleMarkJobFinished(const MarkJobFinishedRequest &request,
+                             MarkJobFinishedReply *reply,
+                             SendReplyCallback send_reply_callback) override;
+};
+
+}  // namespace rpc
+}  // namespace ray
+
+#endif

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

@@ -0,0 +1,93 @@
+#include <boost/asio/basic_socket.hpp>
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "ray/gcs/gcs_server/gcs_server.h"
+#include "ray/rpc/gcs_server/gcs_rpc_client.h"
+
+namespace ray {
+class GcsServerRpcTest : public ::testing::Test {
+ protected:
+  template <typename T>
+  class MockedGcsServer : public gcs::GcsServer {
+   public:
+    explicit MockedGcsServer(const gcs::GcsServerConfig &config)
+        : gcs::GcsServer(config) {}
+
+   protected:
+    void InitBackendClient() override {}
+
+    std::unique_ptr<rpc::JobInfoHandler> InitJobInfoHandler() override {
+      return std::unique_ptr<rpc::JobInfoHandler>(new T);
+    }
+  };
+};
+
+TEST_F(GcsServerRpcTest, JobInfo) {
+  class MockedJobInfoHandler : public rpc::JobInfoHandler {
+   public:
+    void HandleAddJob(const rpc::AddJobRequest &request, rpc::AddJobReply *reply,
+                      rpc::SendReplyCallback send_reply_callback) override {
+      send_reply_callback(Status::OK(), nullptr, nullptr);
+    }
+
+    void HandleMarkJobFinished(const rpc::MarkJobFinishedRequest &request,
+                               rpc::MarkJobFinishedReply *reply,
+                               rpc::SendReplyCallback send_reply_callback) override {
+      send_reply_callback(Status::OK(), nullptr, nullptr);
+    }
+  };
+
+  gcs::GcsServerConfig config;
+  config.grpc_server_port = 0;
+  config.grpc_server_name = "MockedGcsServer";
+  config.grpc_server_thread_num = 1;
+  MockedGcsServer<MockedJobInfoHandler> server(config);
+  std::thread([&server] { server.Start(); }).detach();
+
+  // Wait until server starts listening.
+  while (server.GetPort() == 0) {
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+  }
+
+  boost::asio::io_context client_io_context;
+  boost::asio::io_context::work worker(client_io_context);
+  rpc::ClientCallManager client_call_manager(client_io_context);
+  rpc::GcsRpcClient client("0.0.0.0", server.GetPort(), client_call_manager);
+  std::thread([&client_io_context] { client_io_context.run(); }).detach();
+
+  rpc::AddJobRequest add_job_request;
+  std::promise<rpc::AddJobReply> add_job_reply_promise;
+  auto add_job_reply_future = add_job_reply_promise.get_future();
+  client.AddJob(add_job_request, [&add_job_reply_promise](const Status &status,
+                                                          const rpc::AddJobReply &reply) {
+    if (status.ok()) {
+      add_job_reply_promise.set_value(reply);
+    }
+  });
+  auto future_status = add_job_reply_future.wait_for(std::chrono::milliseconds(200));
+  ASSERT_EQ(future_status, std::future_status::ready);
+
+  rpc::MarkJobFinishedRequest mark_job_finished_request;
+  std::promise<rpc::MarkJobFinishedReply> mark_job_finished_reply_promise;
+  auto mark_job_finished_reply_future = mark_job_finished_reply_promise.get_future();
+  client.MarkJobFinished(
+      mark_job_finished_request,
+      [&mark_job_finished_reply_promise](const Status &status,
+                                         const rpc::MarkJobFinishedReply &reply) {
+        if (status.ok()) {
+          mark_job_finished_reply_promise.set_value(reply);
+        }
+      });
+  future_status = mark_job_finished_reply_future.wait_for(std::chrono::milliseconds(200));
+  ASSERT_EQ(future_status, std::future_status::ready);
+
+  client_io_context.stop();
+  server.Stop();
+}
+
+}  // namespace ray
+
+int main(int argc, char **argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

src/ray/protobuf/gcs_service.proto

@@ -0,0 +1,29 @@
+syntax = "proto3";
+
+package ray.rpc;
+
+import "src/ray/protobuf/gcs.proto";
+
+message AddJobRequest {
+  JobTableData data = 1;
+}
+
+message AddJobReply {
+  bool success = 1;
+}
+
+message MarkJobFinishedRequest {
+  bytes job_id = 1;
+}
+
+message MarkJobFinishedReply {
+  bool success = 1;
+}
+
+// Service for job info access.
+service JobInfoGcsService {
+  // Add job to gcs server.
+  rpc AddJob(AddJobRequest) returns (AddJobReply);
+  // Mark job as finished to gcs server.
+  rpc MarkJobFinished(MarkJobFinishedRequest) returns (MarkJobFinishedReply);
+}

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -0,0 +1,62 @@
+#ifndef RAY_RPC_GCS_RPC_CLIENT_H
+#define RAY_RPC_GCS_RPC_CLIENT_H
+
+#include <thread>
+
+#include <grpcpp/grpcpp.h>
+
+#include "src/ray/protobuf/gcs_service.pb.h"
+#include "src/ray/rpc/client_call.h"
+
+namespace ray {
+namespace rpc {
+
+/// Client used for communicating with gcs server.
+class GcsRpcClient {
+ public:
+  /// Constructor.
+  ///
+  /// \param[in] address Address of gcs server.
+  /// \param[in] port Port of the gcs server.
+  /// \param[in] client_call_manager The `ClientCallManager` used for managing requests.
+  GcsRpcClient(const std::string &address, const int port,
+               ClientCallManager &client_call_manager)
+      : client_call_manager_(client_call_manager) {
+    std::shared_ptr<grpc::Channel> channel = grpc::CreateChannel(
+        address + ":" + std::to_string(port), grpc::InsecureChannelCredentials());
+    stub_ = JobInfoGcsService::NewStub(channel);
+  };
+
+  /// Add job info to gcs server.
+  ///
+  /// \param request The request message.
+  /// \param callback The callback function that handles reply from server.
+  void AddJob(const AddJobRequest &request, const ClientCallback<AddJobReply> &callback) {
+    client_call_manager_.CreateCall<JobInfoGcsService, AddJobRequest, AddJobReply>(
+        *stub_, &JobInfoGcsService::Stub::PrepareAsyncAddJob, request, callback);
+  }
+
+  /// Mark job as finished to gcs server.
+  ///
+  /// \param request The request message.
+  /// \param callback The callback function that handles reply from server.
+  void MarkJobFinished(const MarkJobFinishedRequest &request,
+                       const ClientCallback<MarkJobFinishedReply> &callback) {
+    client_call_manager_
+        .CreateCall<JobInfoGcsService, MarkJobFinishedRequest, MarkJobFinishedReply>(
+            *stub_, &JobInfoGcsService::Stub::PrepareAsyncMarkJobFinished, request,
+            callback);
+  }
+
+ private:
+  /// The gRPC-generated stub.
+  std::unique_ptr<JobInfoGcsService::Stub> stub_;
+
+  /// The `ClientCallManager` used for managing requests.
+  ClientCallManager &client_call_manager_;
+};
+
+}  // namespace rpc
+}  // namespace ray
+
+#endif

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -0,0 +1,68 @@
+#ifndef RAY_RPC_GCS_RPC_SERVER_H
+#define RAY_RPC_GCS_RPC_SERVER_H
+
+#include "src/ray/rpc/grpc_server.h"
+#include "src/ray/rpc/server_call.h"
+
+#include "src/ray/protobuf/gcs_service.grpc.pb.h"
+
+namespace ray {
+namespace rpc {
+
+class JobInfoHandler {
+ public:
+  virtual ~JobInfoHandler() = default;
+
+  virtual void HandleAddJob(const AddJobRequest &request, AddJobReply *reply,
+                            SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleMarkJobFinished(const MarkJobFinishedRequest &request,
+                                     MarkJobFinishedReply *reply,
+                                     SendReplyCallback send_reply_callback) = 0;
+};
+
+/// The `GrpcService` for `JobInfoGcsService`.
+class JobInfoGrpcService : public GrpcService {
+ public:
+  /// Constructor.
+  ///
+  /// \param[in] handler The service handler that actually handle the requests.
+  explicit JobInfoGrpcService(boost::asio::io_service &io_service,
+                              JobInfoHandler &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 {
+    std::unique_ptr<ServerCallFactory> add_job_call_factory(
+        new ServerCallFactoryImpl<JobInfoGcsService, JobInfoHandler, AddJobRequest,
+                                  AddJobReply>(
+            service_, &JobInfoGcsService::AsyncService::RequestAddJob, service_handler_,
+            &JobInfoHandler::HandleAddJob, cq, main_service_));
+    server_call_factories_and_concurrencies->emplace_back(std::move(add_job_call_factory),
+                                                          1);
+
+    std::unique_ptr<ServerCallFactory> mark_job_finished_call_factory(
+        new ServerCallFactoryImpl<JobInfoGcsService, JobInfoHandler,
+                                  MarkJobFinishedRequest, MarkJobFinishedReply>(
+            service_, &JobInfoGcsService::AsyncService::RequestMarkJobFinished,
+            service_handler_, &JobInfoHandler::HandleMarkJobFinished, cq, main_service_));
+    server_call_factories_and_concurrencies->emplace_back(
+        std::move(mark_job_finished_call_factory), 1);
+  }
+
+ private:
+  /// The grpc async service object.
+  JobInfoGcsService::AsyncService service_;
+  /// The service handler that actually handle the requests.
+  JobInfoHandler &service_handler_;
+};
+
+}  // namespace rpc
+}  // namespace ray
+
+#endif