[xray] Put GCS data into the redis data shard (#2298)

src/global_scheduler/global_scheduler.cc

@@ -141,9 +141,10 @@ GlobalSchedulerState *GlobalSchedulerState_init(event_loop *loop,
                          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));
+  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/local_scheduler/local_scheduler.cc

@@ -358,8 +358,9 @@ LocalSchedulerState *LocalSchedulerState_init(
     db_attach(state->db, loop, false);
 
     RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr),
-                                           redis_primary_port));
+                                           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/plasma/plasma_manager.cc

@@ -487,8 +487,11 @@ PlasmaManagerState *PlasmaManagerState_init(const char *store_socket_name,
     db_attach(state->db, state->loop, false);
 
     RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr),
-                                           redis_primary_port));
+                                           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

@@ -8,14 +8,15 @@ namespace gcs {
 
 AsyncGcsClient::AsyncGcsClient(const ClientID &client_id, CommandType command_type) {
   context_ = std::make_shared<RedisContext>();
-  client_table_.reset(new ClientTable(context_, this, client_id));
+  primary_context_ = std::make_shared<RedisContext>();
+  client_table_.reset(new ClientTable(primary_context_, this, client_id));
   object_table_.reset(new ObjectTable(context_, this));
   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));
   heartbeat_table_.reset(new HeartbeatTable(context_, this));
-  error_table_.reset(new ErrorTable(context_, this));
+  error_table_.reset(new ErrorTable(primary_context_, this));
   command_type_ = command_type;
 }
 
@@ -34,8 +35,9 @@ AsyncGcsClient::AsyncGcsClient(CommandType command_type)
 
 AsyncGcsClient::AsyncGcsClient() : AsyncGcsClient(ClientID::from_random()) {}
 
-Status AsyncGcsClient::Connect(const std::string &address, int port) {
-  RAY_RETURN_NOT_OK(context_->Connect(address, port));
+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
@@ -53,6 +55,10 @@ 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()));
+  asio_async_auxiliary_client_.reset(
+      new RedisAsioClient(io_service, primary_context_->async_context()));
+  asio_subscribe_auxiliary_client_.reset(
+      new RedisAsioClient(io_service, primary_context_->subscribe_context()));
   return Status::OK();
 }
 

src/ray/gcs/client.h

@@ -36,8 +36,9 @@ class RAY_EXPORT AsyncGcsClient {
   ///
   /// \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);
+  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);
@@ -68,6 +69,7 @@ class RAY_EXPORT AsyncGcsClient {
                    const GetExportCallback &done_callback);
 
   std::shared_ptr<RedisContext> context() { return context_; }
+  std::shared_ptr<RedisContext> primary_context() { return primary_context_; }
 
  private:
   std::unique_ptr<FunctionTable> function_table_;
@@ -80,10 +82,14 @@ class RAY_EXPORT AsyncGcsClient {
   std::unique_ptr<HeartbeatTable> heartbeat_table_;
   std::unique_ptr<ErrorTable> error_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_;
-
+  // The following context writes everything to the primary shard
+  std::shared_ptr<RedisContext> primary_context_;
+  std::unique_ptr<RedisAsioClient> asio_async_auxiliary_client_;
+  std::unique_ptr<RedisAsioClient> asio_subscribe_auxiliary_client_;
   CommandType command_type_;
 };
 

src/ray/gcs/client_test.cc

@@ -29,7 +29,7 @@ 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));
+    RAY_CHECK_OK(client_->Connect("127.0.0.1", 6379, /*sharding=*/false));
 
     job_id_ = JobID::from_random();
   }

src/ray/gcs/redis_context.cc

@@ -2,6 +2,8 @@
 
 #include <unistd.h>
 
+#include <sstream>
+
 extern "C" {
 #include "hiredis/adapters/ae.h"
 #include "hiredis/async.h"
@@ -36,7 +38,7 @@ namespace gcs {
 // asynchronous redis call. It dispatches the appropriate callback
 // that was registered with the RedisCallbackManager.
 void GlobalRedisCallback(void *c, void *r, void *privdata) {
-  if (r == NULL) {
+  if (r == nullptr) {
     return;
   }
   int64_t callback_index = reinterpret_cast<int64_t>(privdata);
@@ -67,7 +69,7 @@ void GlobalRedisCallback(void *c, void *r, void *privdata) {
 }
 
 void SubscribeRedisCallback(void *c, void *r, void *privdata) {
-  if (r == NULL) {
+  if (r == nullptr) {
     return;
   }
   int64_t callback_index = reinterpret_cast<int64_t>(privdata);
@@ -133,7 +135,70 @@ RedisContext::~RedisContext() {
   }
 }
 
-Status RedisContext::Connect(const std::string &address, int port) {
+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);
+}
+
+Status RedisContext::Connect(const std::string &address, int port, bool sharding) {
   int connection_attempts = 0;
   context_ = redisConnect(address.c_str(), port);
   while (context_ == nullptr || context_->err) {
@@ -158,17 +223,31 @@ Status RedisContext::Connect(const std::string &address, int port) {
   REDIS_CHECK_ERROR(context_, reply);
   freeReplyObject(reply);
 
+  std::string redis_address;
+  int redis_port;
+  if (sharding) {
+    // Get the redis data shard
+    std::vector<std::string> addresses;
+    std::vector<int> ports;
+    GetRedisShards(context_, &addresses, &ports);
+    redis_address = addresses[0];
+    redis_port = ports[0];
+  } else {
+    redis_address = address;
+    redis_port = port;
+  }
+
   // Connect to async context
-  async_context_ = redisAsyncConnect(address.c_str(), port);
+  async_context_ = redisAsyncConnect(redis_address.c_str(), redis_port);
   if (async_context_ == nullptr || async_context_->err) {
-    RAY_LOG(FATAL) << "Could not establish connection to redis " << address << ":"
-                   << port;
+    RAY_LOG(FATAL) << "Could not establish connection to redis " << redis_address << ":"
+                   << redis_port;
   }
   // Connect to subscribe context
-  subscribe_context_ = redisAsyncConnect(address.c_str(), port);
+  subscribe_context_ = redisAsyncConnect(redis_address.c_str(), redis_port);
   if (subscribe_context_ == nullptr || subscribe_context_->err) {
-    RAY_LOG(FATAL) << "Could not establish subscribe connection to redis " << address
-                   << ":" << port;
+    RAY_LOG(FATAL) << "Could not establish subscribe connection to redis "
+                   << redis_address << ":" << redis_port;
   }
   return Status::OK();
 }

src/ray/gcs/redis_context.h

@@ -51,7 +51,7 @@ class RedisContext {
   RedisContext()
       : context_(nullptr), async_context_(nullptr), subscribe_context_(nullptr) {}
   ~RedisContext();
-  Status Connect(const std::string &address, int port);
+  Status Connect(const std::string &address, int port, bool sharding);
   Status AttachToEventLoop(aeEventLoop *loop);
 
   /// Run an operation on some table key.

src/ray/object_manager/test/object_manager_stress_test.cc

@@ -43,7 +43,7 @@ class MockServer {
 
  private:
   ray::Status RegisterGcs(boost::asio::io_service &io_service) {
-    RAY_RETURN_NOT_OK(gcs_client_->Connect("127.0.0.1", 6379));
+    RAY_RETURN_NOT_OK(gcs_client_->Connect("127.0.0.1", 6379, /*sharding=*/false));
     RAY_RETURN_NOT_OK(gcs_client_->Attach(io_service));
 
     boost::asio::ip::tcp::endpoint endpoint = object_manager_acceptor_.local_endpoint();

src/ray/object_manager/test/object_manager_test.cc

@@ -34,7 +34,7 @@ class MockServer {
 
  private:
   ray::Status RegisterGcs(boost::asio::io_service &io_service) {
-    RAY_RETURN_NOT_OK(gcs_client_->Connect("127.0.0.1", 6379));
+    RAY_RETURN_NOT_OK(gcs_client_->Connect("127.0.0.1", 6379, /*sharding=*/false));
     RAY_RETURN_NOT_OK(gcs_client_->Attach(io_service));
 
     boost::asio::ip::tcp::endpoint endpoint = object_manager_acceptor_.local_endpoint();

src/ray/raylet/monitor.cc

@@ -18,7 +18,7 @@ Monitor::Monitor(boost::asio::io_service &io_service, const std::string &redis_a
     : gcs_client_(),
       heartbeat_timeout_ms_(RayConfig::instance().num_heartbeats_timeout()),
       heartbeat_timer_(io_service) {
-  RAY_CHECK_OK(gcs_client_.Connect(redis_address, redis_port));
+  RAY_CHECK_OK(gcs_client_.Connect(redis_address, redis_port, /*sharding=*/true));
   RAY_CHECK_OK(gcs_client_.Attach(io_service));
 }
 

src/ray/raylet/node_manager.cc

@@ -581,8 +581,8 @@ void NodeManager::ProcessNodeManagerMessage(TcpClientConnection &node_manager_cl
   } break;
   case protocol::MessageType::DisconnectClient: {
     // TODO(rkn): We need to do some cleanup here.
-    RAY_LOG(INFO) << "Received disconnect message from remote node manager. "
-                  << "We need to do some cleanup here.";
+    RAY_LOG(DEBUG) << "Received disconnect message from remote node manager. "
+                   << "We need to do some cleanup here.";
   } break;
   default:
     RAY_LOG(FATAL) << "Received unexpected message type " << message_type;

src/ray/raylet/raylet.cc

@@ -54,7 +54,7 @@ ray::Status Raylet::RegisterGcs(const std::string &node_ip_address,
                                 const std::string &redis_address, int redis_port,
                                 boost::asio::io_service &io_service,
                                 const NodeManagerConfig &node_manager_config) {
-  RAY_RETURN_NOT_OK(gcs_client_->Connect(redis_address, redis_port));
+  RAY_RETURN_NOT_OK(gcs_client_->Connect(redis_address, redis_port, /*sharding=*/true));
   RAY_RETURN_NOT_OK(gcs_client_->Attach(io_service));
 
   ClientTableDataT client_info = gcs_client_->client_table().GetLocalClient();