[xray] Implementing Gcs sharding (#2409)

Basically a re-implementation of #2281, with modifications of #2298 (A fix of #2334, for rebasing issues.).
[+] Implement sharding for gcs tables.
[+] Keep ClientTable and ErrorTable managed by the primary_shard. TaskTable is managed by the primary_shard for now, until a good hashing for tasks is implemented.
[+] Move AsyncGcsClient's initialization into Connect function.
[-] Move GetRedisShard and bool sharding from RedisContext's connect into AsyncGcsClient. This may make the interface cleaner.

Author: heyucongtom

src/global_scheduler/global_scheduler.cc

@@ -132,11 +132,6 @@ GlobalSchedulerState *GlobalSchedulerState_init(event_loop *loop,
                          "global_scheduler", node_ip_address,
                          std::vector<std::string>());
   db_attach(state->db, loop, false);
-
-  RAY_CHECK_OK(state->gcs_client.Connect(
-      std::string(redis_primary_addr), redis_primary_port, /*sharding=*/true));
-  RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(loop));
-  RAY_CHECK_OK(state->gcs_client.primary_context()->AttachToEventLoop(loop));
   state->policy_state = GlobalSchedulerPolicyState_init();
   return state;
 }

src/global_scheduler/global_scheduler.h

@@ -51,8 +51,6 @@ typedef struct {
   event_loop *loop;
   /** The global state store database. */
   DBHandle *db;
-  /** The handle to the GCS (modern version of the above). */
-  ray::gcs::AsyncGcsClient gcs_client;
   /** A hash table mapping local scheduler ID to the local schedulers that are
    *  connected to Redis. */
   std::unordered_map<DBClientID, LocalScheduler> local_schedulers;

src/local_scheduler/local_scheduler.cc

@@ -351,11 +351,6 @@ LocalSchedulerState *LocalSchedulerState_init(
     state->db = db_connect(std::string(redis_primary_addr), redis_primary_port,
                            "local_scheduler", node_ip_address, db_connect_args);
     db_attach(state->db, loop, false);
-
-    RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr),
-                                           redis_primary_port, true));
-    RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(loop));
-    RAY_CHECK_OK(state->gcs_client.primary_context()->AttachToEventLoop(loop));
   } else {
     state->db = NULL;
   }

src/local_scheduler/local_scheduler_shared.h

@@ -60,8 +60,6 @@ struct LocalSchedulerState {
   std::unordered_map<ActorID, ActorMapEntry> actor_mapping;
   /** The handle to the database. */
   DBHandle *db;
-  /** The handle to the GCS (modern version of the above). */
-  ray::gcs::AsyncGcsClient gcs_client;
   /** The Plasma client. */
   plasma::PlasmaClient *plasma_conn;
   /** State for the scheduling algorithm. */

src/plasma/plasma_manager.cc

@@ -215,8 +215,6 @@ struct PlasmaManagerState {
    *  other plasma stores. */
   std::unordered_map<std::string, ClientConnection *> manager_connections;
   DBHandle *db;
-  /** The handle to the GCS (modern version of the above). */
-  ray::gcs::AsyncGcsClient gcs_client;
   /** Our address. */
   const char *addr;
   /** Our port. */
@@ -490,13 +488,6 @@ PlasmaManagerState *PlasmaManagerState_init(const char *store_socket_name,
     state->db = db_connect(std::string(redis_primary_addr), redis_primary_port,
                            "plasma_manager", manager_addr, db_connect_args);
     db_attach(state->db, state->loop, false);
-
-    RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr),
-                                           redis_primary_port,
-                                           /*sharding=*/true));
-    RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(state->loop));
-    RAY_CHECK_OK(
-        state->gcs_client.primary_context()->AttachToEventLoop(state->loop));
   } else {
     state->db = NULL;
     RAY_LOG(DEBUG) << "No db connection specified";

src/ray/gcs/client.cc

@@ -2,51 +2,152 @@
 
 #include "ray/gcs/redis_context.h"
 
+static void GetRedisShards(redisContext *context, std::vector<std::string> &addresses,
+                           std::vector<int> &ports) {
+  // Get the total number of Redis shards in the system.
+  int num_attempts = 0;
+  redisReply *reply = nullptr;
+  while (num_attempts < RayConfig::instance().redis_db_connect_retries()) {
+    // Try to read the number of Redis shards from the primary shard. If the
+    // entry is present, exit.
+    reply = reinterpret_cast<redisReply *>(redisCommand(context, "GET NumRedisShards"));
+    if (reply->type != REDIS_REPLY_NIL) {
+      break;
+    }
+
+    // Sleep for a little, and try again if the entry isn't there yet. */
+    freeReplyObject(reply);
+    usleep(RayConfig::instance().redis_db_connect_wait_milliseconds() * 1000);
+    num_attempts++;
+  }
+  RAY_CHECK(num_attempts < RayConfig::instance().redis_db_connect_retries())
+      << "No entry found for NumRedisShards";
+  RAY_CHECK(reply->type == REDIS_REPLY_STRING) << "Expected string, found Redis type "
+                                               << reply->type << " for NumRedisShards";
+  int num_redis_shards = atoi(reply->str);
+  RAY_CHECK(num_redis_shards >= 1) << "Expected at least one Redis shard, "
+                                   << "found " << num_redis_shards;
+  freeReplyObject(reply);
+
+  // Get the addresses of all of the Redis shards.
+  num_attempts = 0;
+  while (num_attempts < RayConfig::instance().redis_db_connect_retries()) {
+    // Try to read the Redis shard locations from the primary shard. If we find
+    // that all of them are present, exit.
+    reply =
+        reinterpret_cast<redisReply *>(redisCommand(context, "LRANGE RedisShards 0 -1"));
+    if (static_cast<int>(reply->elements) == num_redis_shards) {
+      break;
+    }
+
+    // Sleep for a little, and try again if not all Redis shard addresses have
+    // been added yet.
+    freeReplyObject(reply);
+    usleep(RayConfig::instance().redis_db_connect_wait_milliseconds() * 1000);
+    num_attempts++;
+  }
+  RAY_CHECK(num_attempts < RayConfig::instance().redis_db_connect_retries())
+      << "Expected " << num_redis_shards << " Redis shard addresses, found "
+      << reply->elements;
+
+  // Parse the Redis shard addresses.
+  for (size_t i = 0; i < reply->elements; ++i) {
+    // Parse the shard addresses and ports.
+    RAY_CHECK(reply->element[i]->type == REDIS_REPLY_STRING);
+    std::string addr;
+    std::stringstream ss(reply->element[i]->str);
+    getline(ss, addr, ':');
+    addresses.push_back(addr);
+    int port;
+    ss >> port;
+    ports.push_back(port);
+  }
+  freeReplyObject(reply);
+}
+
 namespace ray {
 
 namespace gcs {
 
-AsyncGcsClient::AsyncGcsClient(const ClientID &client_id, CommandType command_type) {
-  context_ = std::make_shared<RedisContext>();
+AsyncGcsClient::AsyncGcsClient(const std::string &address, int port,
+                               const ClientID &client_id, CommandType command_type,
+                               bool is_test_client = false) {
   primary_context_ = std::make_shared<RedisContext>();
-  client_table_.reset(new ClientTable(primary_context_, this, client_id));
-  object_table_.reset(new ObjectTable(context_, this, command_type));
-  actor_table_.reset(new ActorTable(context_, this));
-  task_table_.reset(new TaskTable(context_, this, command_type));
-  raylet_task_table_.reset(new raylet::TaskTable(context_, this, command_type));
-  task_reconstruction_log_.reset(new TaskReconstructionLog(context_, this));
-  task_lease_table_.reset(new TaskLeaseTable(context_, this));
-  heartbeat_table_.reset(new HeartbeatTable(context_, this));
-  driver_table_.reset(new DriverTable(primary_context_, this));
-  error_table_.reset(new ErrorTable(primary_context_, this));
-  profile_table_.reset(new ProfileTable(context_, this));
+
+  RAY_CHECK_OK(primary_context_->Connect(address, port, /*sharding=*/true));
+
+  if (!is_test_client) {
+    // Moving sharding into constructor defaultly means that sharding = true.
+    // This design decision may worth a look.
+    std::vector<std::string> addresses;
+    std::vector<int> ports;
+    GetRedisShards(primary_context_->sync_context(), addresses, ports);
+    if (addresses.size() == 0 || ports.size() == 0) {
+      addresses.push_back(address);
+      ports.push_back(port);
+    }
+
+    // Populate shard_contexts.
+    for (size_t i = 0; i < addresses.size(); ++i) {
+      shard_contexts_.push_back(std::make_shared<RedisContext>());
+    }
+
+    RAY_CHECK(shard_contexts_.size() == addresses.size());
+    for (size_t i = 0; i < addresses.size(); ++i) {
+      RAY_CHECK_OK(
+          shard_contexts_[i]->Connect(addresses[i], ports[i], /*sharding=*/true));
+    }
+  } else {
+    shard_contexts_.push_back(std::make_shared<RedisContext>());
+    RAY_CHECK_OK(shard_contexts_[0]->Connect(address, port, /*sharding=*/true));
+  }
+
+  client_table_.reset(new ClientTable({primary_context_}, this, client_id));
+  error_table_.reset(new ErrorTable({primary_context_}, this));
+  driver_table_.reset(new DriverTable({primary_context_}, this));
+  // Tables below would be sharded.
+  object_table_.reset(new ObjectTable(shard_contexts_, this, command_type));
+  actor_table_.reset(new ActorTable(shard_contexts_, this));
+  task_table_.reset(new TaskTable(shard_contexts_, this, command_type));
+  raylet_task_table_.reset(new raylet::TaskTable(shard_contexts_, this, command_type));
+  task_reconstruction_log_.reset(new TaskReconstructionLog(shard_contexts_, this));
+  task_lease_table_.reset(new TaskLeaseTable(shard_contexts_, this));
+  heartbeat_table_.reset(new HeartbeatTable(shard_contexts_, this));
+  profile_table_.reset(new ProfileTable(shard_contexts_, this));
   command_type_ = command_type;
+
+  // TODO(swang): Call the client table's Connect() method here. To do this,
+  // we need to make sure that we are attached to an event loop first. This
+  // currently isn't possible because the aeEventLoop, which we use for
+  // testing, requires us to connect to Redis first.
 }
 
 #if RAY_USE_NEW_GCS
 // Use of kChain currently only applies to Table::Add which affects only the
 // task table, and when RAY_USE_NEW_GCS is set at compile time.
-AsyncGcsClient::AsyncGcsClient(const ClientID &client_id)
-    : AsyncGcsClient(client_id, CommandType::kChain) {}
+AsyncGcsClient::AsyncGcsClient(const std::string &address, int port,
+                               const ClientID &client_id, bool is_test_client = false)
+    : AsyncGcsClient(address, port, client_id, CommandType::kChain, is_test_client) {}
 #else
-AsyncGcsClient::AsyncGcsClient(const ClientID &client_id)
-    : AsyncGcsClient(client_id, CommandType::kRegular) {}
+AsyncGcsClient::AsyncGcsClient(const std::string &address, int port,
+                               const ClientID &client_id, bool is_test_client = false)
+    : AsyncGcsClient(address, port, client_id, CommandType::kRegular, is_test_client) {}
 #endif  // RAY_USE_NEW_GCS
 
-AsyncGcsClient::AsyncGcsClient(CommandType command_type)
-    : AsyncGcsClient(ClientID::from_random(), command_type) {}
+AsyncGcsClient::AsyncGcsClient(const std::string &address, int port,
+                               CommandType command_type)
+    : AsyncGcsClient(address, port, ClientID::from_random(), command_type) {}
 
-AsyncGcsClient::AsyncGcsClient() : AsyncGcsClient(ClientID::from_random()) {}
+AsyncGcsClient::AsyncGcsClient(const std::string &address, int port,
+                               CommandType command_type, bool is_test_client)
+    : AsyncGcsClient(address, port, ClientID::from_random(), command_type,
+                     is_test_client) {}
 
-Status AsyncGcsClient::Connect(const std::string &address, int port, bool sharding) {
-  RAY_RETURN_NOT_OK(context_->Connect(address, port, sharding));
-  RAY_RETURN_NOT_OK(primary_context_->Connect(address, port, /*sharding=*/false));
-  // TODO(swang): Call the client table's Connect() method here. To do this,
-  // we need to make sure that we are attached to an event loop first. This
-  // currently isn't possible because the aeEventLoop, which we use for
-  // testing, requires us to connect to Redis first.
-  return Status::OK();
-}
+AsyncGcsClient::AsyncGcsClient(const std::string &address, int port)
+    : AsyncGcsClient(address, port, ClientID::from_random()) {}
+
+AsyncGcsClient::AsyncGcsClient(const std::string &address, int port, bool is_test_client)
+    : AsyncGcsClient(address, port, ClientID::from_random(), is_test_client) {}
 
 Status Attach(plasma::EventLoop &event_loop) {
   // TODO(pcm): Implement this via
@@ -55,9 +156,14 @@ Status Attach(plasma::EventLoop &event_loop) {
 }
 
 Status AsyncGcsClient::Attach(boost::asio::io_service &io_service) {
-  asio_async_client_.reset(new RedisAsioClient(io_service, context_->async_context()));
-  asio_subscribe_client_.reset(
-      new RedisAsioClient(io_service, context_->subscribe_context()));
+  // Take care of sharding contexts.
+  RAY_CHECK(shard_asio_async_clients_.empty()) << "Attach shall be called only once";
+  for (std::shared_ptr<RedisContext> context : shard_contexts_) {
+    shard_asio_async_clients_.emplace_back(
+        new RedisAsioClient(io_service, context->async_context()));
+    shard_asio_subscribe_clients_.emplace_back(
+        new RedisAsioClient(io_service, context->subscribe_context()));
+  }
   asio_async_auxiliary_client_.reset(
       new RedisAsioClient(io_service, primary_context_->async_context()));
   asio_subscribe_auxiliary_client_.reset(

src/ray/gcs/client.h

@@ -24,21 +24,22 @@ class RAY_EXPORT AsyncGcsClient {
   /// Attach() must be called. To read and write from the GCS tables requires a
   /// further call to Connect() to the client table.
   ///
+  /// \param address The GCS IP address.
+  /// \param port The GCS port.
+  /// \param sharding If true, use sharded redis for the GCS.
   /// \param client_id The ID to assign to the client.
   /// \param command_type GCS command type.  If CommandType::kChain, chain-replicated
   /// versions of the tables might be used, if available.
-  AsyncGcsClient(const ClientID &client_id, CommandType command_type);
-  AsyncGcsClient(const ClientID &client_id);
-  AsyncGcsClient(CommandType command_type);
-  AsyncGcsClient();
+  AsyncGcsClient(const std::string &address, int port, const ClientID &client_id,
+                 CommandType command_type, bool is_test_client);
+  AsyncGcsClient(const std::string &address, int port, const ClientID &client_id,
+                 bool is_test_client);
+  AsyncGcsClient(const std::string &address, int port, CommandType command_type);
+  AsyncGcsClient(const std::string &address, int port, CommandType command_type,
+                 bool is_test_client);
+  AsyncGcsClient(const std::string &address, int port);
+  AsyncGcsClient(const std::string &address, int port, bool is_test_client);
 
-  /// Connect to the GCS.
-  ///
-  /// \param address The GCS IP address.
-  /// \param port The GCS port.
-  /// \param sharding If true, use sharded redis for the GCS.
-  /// \return Status.
-  Status Connect(const std::string &address, int port, bool sharding);
   /// Attach this client to a plasma event loop. Note that only
   /// one event loop should be attached at a time.
   Status Attach(plasma::EventLoop &event_loop);
@@ -71,7 +72,7 @@ class RAY_EXPORT AsyncGcsClient {
   Status GetExport(const std::string &driver_id, int64_t export_index,
                    const GetExportCallback &done_callback);
 
-  std::shared_ptr<RedisContext> context() { return context_; }
+  std::vector<std::shared_ptr<RedisContext>> shard_contexts() { return shard_contexts_; }
   std::shared_ptr<RedisContext> primary_context() { return primary_context_; }
 
  private:
@@ -88,9 +89,9 @@ class RAY_EXPORT AsyncGcsClient {
   std::unique_ptr<ProfileTable> profile_table_;
   std::unique_ptr<ClientTable> client_table_;
   // The following contexts write to the data shard
-  std::shared_ptr<RedisContext> context_;
-  std::unique_ptr<RedisAsioClient> asio_async_client_;
-  std::unique_ptr<RedisAsioClient> asio_subscribe_client_;
+  std::vector<std::shared_ptr<RedisContext>> shard_contexts_;
+  std::vector<std::unique_ptr<RedisAsioClient>> shard_asio_async_clients_;
+  std::vector<std::unique_ptr<RedisAsioClient>> shard_asio_subscribe_clients_;
   // The following context writes everything to the primary shard
   std::shared_ptr<RedisContext> primary_context_;
   std::unique_ptr<DriverTable> driver_table_;

src/ray/gcs/client_test.cc

@@ -28,9 +28,8 @@ static inline void flushall_redis(void) {
 class TestGcs : public ::testing::Test {
  public:
   TestGcs(CommandType command_type) : num_callbacks_(0), command_type_(command_type) {
-    client_ = std::make_shared<gcs::AsyncGcsClient>(command_type_);
-    RAY_CHECK_OK(client_->Connect("127.0.0.1", 6379, /*sharding=*/false));
-
+    client_ = std::make_shared<gcs::AsyncGcsClient>("127.0.0.1", 6379, command_type_,
+                                                    /*is_test_client=*/true);
     job_id_ = JobID::from_random();
   }
 
@@ -60,7 +59,10 @@ class TestGcsWithAe : public TestGcs {
  public:
   TestGcsWithAe(CommandType command_type) : TestGcs(command_type) {
     loop_ = aeCreateEventLoop(1024);
-    RAY_CHECK_OK(client_->context()->AttachToEventLoop(loop_));
+    RAY_CHECK_OK(client_->primary_context()->AttachToEventLoop(loop_));
+    for (auto &context : client_->shard_contexts()) {
+      RAY_CHECK_OK(context->AttachToEventLoop(loop_));
+    }
   }
 
   TestGcsWithAe() : TestGcsWithAe(CommandType::kRegular) {}