Integrate credis with Ray & route task table entries into credis. (#1841)

src/ray/gcs/client.cc

@@ -6,20 +6,25 @@ namespace ray {
 
 namespace gcs {
 
-AsyncGcsClient::AsyncGcsClient(const ClientID &client_id) {
+AsyncGcsClient::AsyncGcsClient(const ClientID &client_id, CommandType command_type) {
   context_ = std::make_shared<RedisContext>();
   client_table_.reset(new ClientTable(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));
-  raylet_task_table_.reset(new raylet::TaskTable(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));
+  command_type_ = command_type;
 }
 
-AsyncGcsClient::AsyncGcsClient() : AsyncGcsClient(ClientID::from_random()) {}
+AsyncGcsClient::AsyncGcsClient(const ClientID &client_id)
+    : AsyncGcsClient(client_id, CommandType::kRegular) {}
 
-AsyncGcsClient::~AsyncGcsClient() {}
+AsyncGcsClient::AsyncGcsClient(CommandType command_type)
+    : AsyncGcsClient(ClientID::from_random(), 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));

src/ray/gcs/client.h

@@ -19,15 +19,18 @@ class RedisContext;
 
 class RAY_EXPORT AsyncGcsClient {
  public:
-  /// Start a GCS client with the given client ID. To read from the GCS tables,
-  /// Connect and then Attach must be called. To read and write from the GCS
-  /// tables requires a further call to Connect to the client table.
+  /// Start a GCS client with the given client ID and command type (regular or
+  /// chain-replicated). To read from the GCS tables, Connect() and then
+  /// Attach() must be called. To read and write from the GCS tables requires a
+  /// further call to Connect() to the client table.
   ///
   /// \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);
-  /// Start a GCS client with a random client ID.
+  AsyncGcsClient(CommandType command_type);
   AsyncGcsClient();
-  ~AsyncGcsClient();
 
   /// Connect to the GCS.
   ///
@@ -79,6 +82,8 @@ class RAY_EXPORT AsyncGcsClient {
   std::shared_ptr<RedisContext> context_;
   std::unique_ptr<RedisAsioClient> asio_async_client_;
   std::unique_ptr<RedisAsioClient> asio_subscribe_client_;
+
+  CommandType command_type_;
 };
 
 class SyncGcsClient {

src/ray/gcs/client_test.cc

@@ -12,6 +12,12 @@ extern "C" {
 
 namespace ray {
 
+namespace gcs {
+
+namespace {
+constexpr char kRandomId[] = "abcdefghijklmnopqrst";
+}  // namespace
+
 /* Flush redis. */
 static inline void flushall_redis(void) {
   redisContext *context = redisConnect("127.0.0.1", 6379);
@@ -21,8 +27,8 @@ static inline void flushall_redis(void) {
 
 class TestGcs : public ::testing::Test {
  public:
-  TestGcs() : num_callbacks_(0) {
-    client_ = std::make_shared<gcs::AsyncGcsClient>();
+  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));
 
     job_id_ = JobID::from_random();
@@ -43,6 +49,7 @@ class TestGcs : public ::testing::Test {
 
  protected:
   uint64_t num_callbacks_;
+  gcs::CommandType command_type_;
   std::shared_ptr<gcs::AsyncGcsClient> client_;
   JobID job_id_;
 };
@@ -51,10 +58,13 @@ TestGcs *test;
 
 class TestGcsWithAe : public TestGcs {
  public:
-  TestGcsWithAe() {
+  TestGcsWithAe(CommandType command_type) : TestGcs(command_type) {
     loop_ = aeCreateEventLoop(1024);
     RAY_CHECK_OK(client_->context()->AttachToEventLoop(loop_));
   }
+
+  TestGcsWithAe() : TestGcsWithAe(CommandType::kRegular) {}
+
   ~TestGcsWithAe() override {
     // Destroy the client first since it has a reference to the event loop.
     client_.reset();
@@ -67,11 +77,20 @@ class TestGcsWithAe : public TestGcs {
   aeEventLoop *loop_;
 };
 
+class TestGcsWithChainAe : public TestGcsWithAe {
+ public:
+  TestGcsWithChainAe() : TestGcsWithAe(gcs::CommandType::kChain){};
+};
+
 class TestGcsWithAsio : public TestGcs {
  public:
-  TestGcsWithAsio() : TestGcs(), io_service_(), work_(io_service_) {
+  TestGcsWithAsio(CommandType command_type)
+      : TestGcs(command_type), io_service_(), work_(io_service_) {
     RAY_CHECK_OK(client_->Attach(io_service_));
   }
+
+  TestGcsWithAsio() : TestGcsWithAsio(CommandType::kRegular) {}
+
   ~TestGcsWithAsio() {
     // Destroy the client first since it has a reference to the event loop.
     client_.reset();
@@ -86,6 +105,11 @@ class TestGcsWithAsio : public TestGcs {
   boost::asio::io_service::work work_;
 };
 
+class TestGcsWithChainAsio : public TestGcsWithAsio {
+ public:
+  TestGcsWithChainAsio() : TestGcsWithAsio(gcs::CommandType::kChain){};
+};
+
 void TestTableLookup(const JobID &job_id, std::shared_ptr<gcs::AsyncGcsClient> client) {
   TaskID task_id = TaskID::from_random();
   auto data = std::make_shared<protocol::TaskT>();
@@ -120,15 +144,20 @@ void TestTableLookup(const JobID &job_id, std::shared_ptr<gcs::AsyncGcsClient> c
   test->Start();
 }
 
-TEST_F(TestGcsWithAe, TestTableLookup) {
-  test = this;
-  TestTableLookup(job_id_, client_);
-}
+// Convenient macro to test across {ae, asio} x {regular, chain} x {the tests}.
+// Undefined at the end.
+#define TEST_MACRO(FIXTURE, TEST) \
+  TEST_F(FIXTURE, TEST) {         \
+    test = this;                  \
+    TEST(job_id_, client_);       \
+  }
 
-TEST_F(TestGcsWithAsio, TestTableLookup) {
-  test = this;
-  TestTableLookup(job_id_, client_);
-}
+TEST_MACRO(TestGcsWithAe, TestTableLookup);
+TEST_MACRO(TestGcsWithAsio, TestTableLookup);
+#if RAY_USE_NEW_GCS
+TEST_MACRO(TestGcsWithChainAe, TestTableLookup);
+TEST_MACRO(TestGcsWithChainAsio, TestTableLookup);
+#endif
 
 void TestLogLookup(const JobID &job_id, std::shared_ptr<gcs::AsyncGcsClient> client) {
   // Append some entries to the log at an object ID.
@@ -200,15 +229,12 @@ void TestTableLookupFailure(const JobID &job_id,
   test->Start();
 }
 
-TEST_F(TestGcsWithAe, TestTableLookupFailure) {
-  test = this;
-  TestTableLookupFailure(job_id_, client_);
-}
-
-TEST_F(TestGcsWithAsio, TestTableLookupFailure) {
-  test = this;
-  TestTableLookupFailure(job_id_, client_);
-}
+TEST_MACRO(TestGcsWithAe, TestTableLookupFailure);
+TEST_MACRO(TestGcsWithAsio, TestTableLookupFailure);
+#if RAY_USE_NEW_GCS
+TEST_MACRO(TestGcsWithChainAe, TestTableLookupFailure);
+TEST_MACRO(TestGcsWithChainAsio, TestTableLookupFailure);
+#endif
 
 void TestLogAppendAt(const JobID &job_id, std::shared_ptr<gcs::AsyncGcsClient> client) {
   TaskID task_id = TaskID::from_random();
@@ -227,14 +253,23 @@ void TestLogAppendAt(const JobID &job_id, std::shared_ptr<gcs::AsyncGcsClient> c
     test->IncrementNumCallbacks();
   };
 
+  // Will succeed.
   RAY_CHECK_OK(client->task_reconstruction_log().Append(job_id, task_id, data_log.front(),
-                                                        nullptr));
-  RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(job_id, task_id, data_log[1],
-                                                          nullptr, failure_callback, 0));
-  RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(job_id, task_id, data_log[1],
-                                                          nullptr, failure_callback, 2));
-  RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(job_id, task_id, data_log[1],
-                                                          nullptr, failure_callback, 1));
+                                                        /*done callback=*/nullptr));
+  // Append at index 0 will fail.
+  RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(
+      job_id, task_id, data_log[1],
+      /*done callback=*/nullptr, failure_callback, /*log_length=*/0));
+
+  // Append at index 2 will fail.
+  RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(
+      job_id, task_id, data_log[1],
+      /*done callback=*/nullptr, failure_callback, /*log_length=*/2));
+
+  // Append at index 1 will succeed.
+  RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(
+      job_id, task_id, data_log[1],
+      /*done callback=*/nullptr, failure_callback, /*log_length=*/1));
 
   auto lookup_callback = [managers](gcs::AsyncGcsClient *client, const UniqueID &id,
                                     const std::vector<TaskReconstructionDataT> &data) {
@@ -267,11 +302,24 @@ TEST_F(TestGcsWithAsio, TestLogAppendAt) {
 void TaskAdded(gcs::AsyncGcsClient *client, const TaskID &id,
                const TaskTableDataT &data) {
   ASSERT_EQ(data.scheduling_state, SchedulingState_SCHEDULED);
+  ASSERT_EQ(data.scheduler_id, kRandomId);
 }
 
+void TaskLookupHelper(gcs::AsyncGcsClient *client, const TaskID &id,
+                      const TaskTableDataT &data, bool do_stop) {
+  ASSERT_EQ(data.scheduling_state, SchedulingState_SCHEDULED);
+  ASSERT_EQ(data.scheduler_id, kRandomId);
+  if (do_stop) {
+    test->Stop();
+  }
+}
 void TaskLookup(gcs::AsyncGcsClient *client, const TaskID &id,
                 const TaskTableDataT &data) {
-  ASSERT_EQ(data.scheduling_state, SchedulingState_SCHEDULED);
+  TaskLookupHelper(client, id, data, /*do_stop=*/false);
+}
+void TaskLookupWithStop(gcs::AsyncGcsClient *client, const TaskID &id,
+                        const TaskTableDataT &data) {
+  TaskLookupHelper(client, id, data, /*do_stop=*/true);
 }
 
 void TaskLookupFailure(gcs::AsyncGcsClient *client, const TaskID &id) {
@@ -298,7 +346,7 @@ void TaskUpdateCallback(gcs::AsyncGcsClient *client, const TaskID &task_id,
 void TestTaskTable(const JobID &job_id, std::shared_ptr<gcs::AsyncGcsClient> client) {
   auto data = std::make_shared<TaskTableDataT>();
   data->scheduling_state = SchedulingState_SCHEDULED;
-  ClientID local_scheduler_id = ClientID::from_binary("abcdefghijklmnopqrst");
+  ClientID local_scheduler_id = ClientID::from_binary(kRandomId);
   data->scheduler_id = local_scheduler_id.binary();
   TaskID task_id = TaskID::from_random();
   RAY_CHECK_OK(client->task_table().Add(job_id, task_id, data, &TaskAdded));
@@ -317,15 +365,12 @@ void TestTaskTable(const JobID &job_id, std::shared_ptr<gcs::AsyncGcsClient> cli
   test->Start();
 }
 
-TEST_F(TestGcsWithAe, TestTaskTable) {
-  test = this;
-  TestTaskTable(job_id_, client_);
-}
-
-TEST_F(TestGcsWithAsio, TestTaskTable) {
-  test = this;
-  TestTaskTable(job_id_, client_);
-}
+TEST_MACRO(TestGcsWithAe, TestTaskTable);
+TEST_MACRO(TestGcsWithAsio, TestTaskTable);
+#if RAY_USE_NEW_GCS
+TEST_MACRO(TestGcsWithChainAe, TestTaskTable);
+TEST_MACRO(TestGcsWithChainAsio, TestTaskTable);
+#endif
 
 void TestTableSubscribeAll(const JobID &job_id,
                            std::shared_ptr<gcs::AsyncGcsClient> client) {
@@ -366,15 +411,12 @@ void TestTableSubscribeAll(const JobID &job_id,
   ASSERT_EQ(test->NumCallbacks(), task_specs.size());
 }
 
-TEST_F(TestGcsWithAe, TestTableSubscribeAll) {
-  test = this;
-  TestTableSubscribeAll(job_id_, client_);
-}
-
-TEST_F(TestGcsWithAsio, TestTableSubscribeAll) {
-  test = this;
-  TestTableSubscribeAll(job_id_, client_);
-}
+TEST_MACRO(TestGcsWithAe, TestTableSubscribeAll);
+TEST_MACRO(TestGcsWithAsio, TestTableSubscribeAll);
+#if RAY_USE_NEW_GCS
+TEST_MACRO(TestGcsWithChainAe, TestTableSubscribeAll);
+TEST_MACRO(TestGcsWithChainAsio, TestTableSubscribeAll);
+#endif
 
 void TestLogSubscribeAll(const JobID &job_id,
                          std::shared_ptr<gcs::AsyncGcsClient> client) {
@@ -498,15 +540,12 @@ void TestTableSubscribeId(const JobID &job_id,
   ASSERT_EQ(test->NumCallbacks(), task_specs2.size());
 }
 
-TEST_F(TestGcsWithAe, TestTableSubscribeId) {
-  test = this;
-  TestTableSubscribeId(job_id_, client_);
-}
-
-TEST_F(TestGcsWithAsio, TestTableSubscribeId) {
-  test = this;
-  TestTableSubscribeId(job_id_, client_);
-}
+TEST_MACRO(TestGcsWithAe, TestTableSubscribeId);
+TEST_MACRO(TestGcsWithAsio, TestTableSubscribeId);
+#if RAY_USE_NEW_GCS
+TEST_MACRO(TestGcsWithChainAe, TestTableSubscribeId);
+TEST_MACRO(TestGcsWithChainAsio, TestTableSubscribeId);
+#endif
 
 void TestLogSubscribeId(const JobID &job_id,
                         std::shared_ptr<gcs::AsyncGcsClient> client) {
@@ -650,15 +689,12 @@ void TestTableSubscribeCancel(const JobID &job_id,
   ASSERT_EQ(test->NumCallbacks(), 2);
 }
 
-TEST_F(TestGcsWithAe, TestTableSubscribeCancel) {
-  test = this;
-  TestTableSubscribeCancel(job_id_, client_);
-}
-
-TEST_F(TestGcsWithAsio, TestTableSubscribeCancel) {
-  test = this;
-  TestTableSubscribeCancel(job_id_, client_);
-}
+TEST_MACRO(TestGcsWithAe, TestTableSubscribeCancel);
+TEST_MACRO(TestGcsWithAsio, TestTableSubscribeCancel);
+#if RAY_USE_NEW_GCS
+TEST_MACRO(TestGcsWithChainAe, TestTableSubscribeCancel);
+TEST_MACRO(TestGcsWithChainAsio, TestTableSubscribeCancel);
+#endif
 
 void TestLogSubscribeCancel(const JobID &job_id,
                             std::shared_ptr<gcs::AsyncGcsClient> client) {
@@ -780,14 +816,14 @@ void TestClientTableDisconnect(const JobID &job_id,
   // event will stop the event loop.
   client->client_table().RegisterClientAddedCallback(
       [](gcs::AsyncGcsClient *client, const UniqueID &id, const ClientTableDataT &data) {
-        ClientTableNotification(client, id, data, true);
+        ClientTableNotification(client, id, data, /*is_insertion=*/true);
         // Disconnect from the client table. We should receive a notification
         // for the removal of our own entry.
         RAY_CHECK_OK(client->client_table().Disconnect());
       });
   client->client_table().RegisterClientRemovedCallback(
       [](gcs::AsyncGcsClient *client, const UniqueID &id, const ClientTableDataT &data) {
-        ClientTableNotification(client, id, data, false);
+        ClientTableNotification(client, id, data, /*is_insertion=*/false);
         test->Stop();
       });
   // Connect to the client table. We should receive notification for the
@@ -860,4 +896,7 @@ TEST_F(TestGcsWithAsio, TestClientTableMarkDisconnected) {
   TestClientTableMarkDisconnected(job_id_, client_);
 }
 
-}  // namespace
+#undef TEST_MACRO
+
+}  // namespace gcs
+}  // namespace ray

src/ray/gcs/redis_context.cc

@@ -3,9 +3,9 @@
 #include <unistd.h>
 
 extern "C" {
+#include "hiredis/adapters/ae.h"
 #include "hiredis/async.h"
 #include "hiredis/hiredis.h"
-#include "hiredis/adapters/ae.h"
 }
 
 // TODO(pcm): Integrate into the C++ tree.
@@ -55,9 +55,13 @@ void GlobalRedisCallback(void *c, void *r, void *privdata) {
   case (REDIS_REPLY_ERROR): {
     RAY_LOG(ERROR) << "Redis error " << reply->str;
   } break;
+  case (REDIS_REPLY_INTEGER): {
+    data = std::to_string(reply->integer);
+    break;
+  }
   default:
-    RAY_LOG(FATAL) << "Fatal redis error of type " << reply->type
-                   << " and with string " << reply->str;
+    RAY_LOG(FATAL) << "Fatal redis error of type " << reply->type << " and with string "
+                   << reply->str;
   }
   ProcessCallback(callback_index, data);
 }
@@ -99,9 +103,8 @@ void SubscribeRedisCallback(void *c, void *r, void *privdata) {
 }
 
 int64_t RedisCallbackManager::add(const RedisCallback &function) {
-  num_callbacks += 1;
-  callbacks_.emplace(num_callbacks, function);
-  return num_callbacks;
+  callbacks_.emplace(num_callbacks_, function);
+  return num_callbacks_++;
 }
 
 RedisCallback &RedisCallbackManager::get(int64_t callback_index) {
@@ -134,14 +137,13 @@ Status RedisContext::Connect(const std::string &address, int port) {
   int connection_attempts = 0;
   context_ = redisConnect(address.c_str(), port);
   while (context_ == nullptr || context_->err) {
-    if (connection_attempts >=
-        RayConfig::instance().redis_db_connect_retries()) {
+    if (connection_attempts >= RayConfig::instance().redis_db_connect_retries()) {
       if (context_ == nullptr) {
         RAY_LOG(FATAL) << "Could not allocate redis context.";
       }
       if (context_->err) {
-        RAY_LOG(FATAL) << "Could not establish connection to redis " << address
-                       << ":" << port;
+        RAY_LOG(FATAL) << "Could not establish connection to redis " << address << ":"
+                       << port;
       }
       break;
     }
@@ -159,8 +161,8 @@ Status RedisContext::Connect(const std::string &address, int port) {
   // Connect to async context
   async_context_ = redisAsyncConnect(address.c_str(), 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 " << address << ":"
+                   << port;
   }
   // Connect to subscribe context
   subscribe_context_ = redisAsyncConnect(address.c_str(), port);

src/ray/gcs/redis_context.h

@@ -25,7 +25,6 @@ using RedisCallback = std::function<bool(const std::string &)>;
 
 class RedisCallbackManager {
  public:
-
   static RedisCallbackManager &instance() {
     static RedisCallbackManager instance;
     return instance;
@@ -39,11 +38,11 @@ class RedisCallbackManager {
   void remove(int64_t callback_index);
 
  private:
-  RedisCallbackManager() : num_callbacks(0){};
+  RedisCallbackManager() : num_callbacks_(0){};
 
   ~RedisCallbackManager() { printf("shut down callback manager\n"); }
 
-  int64_t num_callbacks;
+  int64_t num_callbacks_ = 0;
   std::unordered_map<int64_t, RedisCallback> callbacks_;
 };
 

src/ray/gcs/tables.cc

@@ -11,7 +11,6 @@ template <typename ID, typename Data>
 Status Log<ID, Data>::Append(const JobID &job_id, const ID &id,
                              std::shared_ptr<DataT> &dataT, const WriteCallback &done) {
   auto callback = [this, id, dataT, done](const std::string &data) {
-    RAY_CHECK(data.empty());
     if (done != nullptr) {
       (done)(client_, id, *dataT);
     }
@@ -141,8 +140,15 @@ Status Table<ID, Data>::Add(const JobID &job_id, const ID &id,
   flatbuffers::FlatBufferBuilder fbb;
   fbb.ForceDefaults(true);
   fbb.Finish(Data::Pack(fbb, dataT.get()));
-  return context_->RunAsync("RAY.TABLE_ADD", id, fbb.GetBufferPointer(), fbb.GetSize(),
-                            prefix_, pubsub_channel_, std::move(callback));
+  if (command_type_ == CommandType::kRegular) {
+    return context_->RunAsync("RAY.TABLE_ADD", id, fbb.GetBufferPointer(), fbb.GetSize(),
+                              prefix_, pubsub_channel_, std::move(callback));
+  } else {
+    RAY_CHECK(command_type_ == CommandType::kChain);
+    return context_->RunAsync("RAY.CHAIN.TABLE_ADD", id, fbb.GetBufferPointer(),
+                              fbb.GetSize(), prefix_, pubsub_channel_,
+                              std::move(callback));
+  }
 }
 
 template <typename ID, typename Data>

src/ray/gcs/tables.h

@@ -27,6 +27,10 @@ class RedisContext;
 
 class AsyncGcsClient;
 
+/// Specifies whether commands issued to a table should be regular or chain-replicated
+/// (when available).
+enum class CommandType { kRegular, kChain };
+
 /// \class PubsubInterface
 ///
 /// The interface for a pubsub storage system. The client of a storage system
@@ -186,6 +190,9 @@ class Log : virtual public PubsubInterface<ID> {
   /// when we receive notifications. This is >= 0 iff we have subscribed to the
   /// table, otherwise -1.
   int64_t subscribe_callback_index_;
+
+  /// Commands to a GCS table can either be regular (default) or chain-replicated.
+  CommandType command_type_ = CommandType::kRegular;
 };
 
 template <typename ID, typename Data>
@@ -259,6 +266,7 @@ class Table : private Log<ID, Data>,
   using Log<ID, Data>::client_;
   using Log<ID, Data>::pubsub_channel_;
   using Log<ID, Data>::prefix_;
+  using Log<ID, Data>::command_type_;
 };
 
 class ObjectTable : public Log<ObjectID, ObjectTableData> {
@@ -320,6 +328,12 @@ class TaskTable : public Table<TaskID, ray::protocol::Task> {
     pubsub_channel_ = TablePubsub_RAYLET_TASK;
     prefix_ = TablePrefix_RAYLET_TASK;
   }
+
+  TaskTable(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client,
+            gcs::CommandType command_type)
+      : TaskTable(context, client) {
+    command_type_ = command_type;
+  };
 };
 
 }  // namespace raylet
@@ -331,7 +345,12 @@ class TaskTable : public Table<TaskID, TaskTableData> {
     pubsub_channel_ = TablePubsub_TASK;
     prefix_ = TablePrefix_TASK;
   };
-  ~TaskTable(){};
+
+  TaskTable(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client,
+            gcs::CommandType command_type)
+      : TaskTable(context, client) {
+    command_type_ = command_type;
+  }
 
   using TestAndUpdateCallback =
       std::function<void(AsyncGcsClient *client, const TaskID &id,

src/ray/test/run_gcs_tests.sh

@@ -7,7 +7,18 @@ set -e
 set -x
 
 # Start Redis.
-./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6379 &
+if [[ "${RAY_USE_NEW_GCS}" = "on" ]]; then
+    ./src/credis/redis/src/redis-server \
+        --loglevel warning \
+        --loadmodule ./src/credis/build/src/libmember.so \
+        --loadmodule ./src/common/redis_module/libray_redis_module.so \
+        --port 6379 &
+else
+    ./src/common/thirdparty/redis/src/redis-server \
+        --loglevel warning \
+        --loadmodule ./src/common/redis_module/libray_redis_module.so \
+        --port 6379 &
+fi
 sleep 1s
 
 ./src/ray/gcs/client_test

src/ray/test/run_object_manager_tests.sh

@@ -21,18 +21,25 @@ if [ ! -d "$RAY_ROOT/python" ]; then
 fi
 
 CORE_DIR="$RAY_ROOT/python/ray/core"
-REDIS_DIR="$CORE_DIR/src/common/thirdparty/redis/src"
 REDIS_MODULE="$CORE_DIR/src/common/redis_module/libray_redis_module.so"
-STORE_EXEC="$CORE_DIR/src/plasma/plasma_store"
+REDIS_DIR="$CORE_DIR/src/common/thirdparty/redis/src"
 
-echo "$STORE_EXEC"
-echo "$REDIS_DIR/redis-server --loglevel warning --loadmodule $REDIS_MODULE --port 6379"
-echo "$REDIS_DIR/redis-cli -p 6379 shutdown"
+if [[ "${RAY_USE_NEW_GCS}" = "on" ]]; then
+    REDIS_SERVER="$CORE_DIR/src/credis/redis/src/redis-server"
+
+    CREDIS_MODULE="$CORE_DIR/src/credis/build/src/libmember.so"
+    LOAD_MODULE_ARGS="--loadmodule ${CREDIS_MODULE} --loadmodule ${REDIS_MODULE}"
+else
+    REDIS_SERVER="${REDIS_DIR}/redis-server"
+    LOAD_MODULE_ARGS="--loadmodule ${REDIS_MODULE}"
+fi
+
+STORE_EXEC="$CORE_DIR/src/plasma/plasma_store"
 
 # Allow cleanup commands to fail.
 $REDIS_DIR/redis-cli -p 6379 shutdown || true
 sleep 1s
-$REDIS_DIR/redis-server --loglevel warning --loadmodule $REDIS_MODULE --port 6379 &
+${REDIS_SERVER} --loglevel warning ${LOAD_MODULE_ARGS} --port 6379 &
 sleep 1s
 # Run tests.
 $CORE_DIR/src/ray/object_manager/object_manager_stress_test $STORE_EXEC

src/ray/test/run_object_manager_valgrind.sh

@@ -26,15 +26,25 @@ REDIS_MODULE="$CORE_DIR/src/common/redis_module/libray_redis_module.so"
 STORE_EXEC="$CORE_DIR/src/plasma/plasma_store"
 VALGRIND_CMD="valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1"
 
+if [[ "${RAY_USE_NEW_GCS}" = "on" ]]; then
+    REDIS_SERVER="$CORE_DIR/src/credis/redis/src/redis-server"
+
+    CREDIS_MODULE="$CORE_DIR/src/credis/build/src/libmember.so"
+    LOAD_MODULE_ARGS="--loadmodule ${CREDIS_MODULE} --loadmodule ${REDIS_MODULE}"
+else
+    REDIS_SERVER="${REDIS_DIR}/redis-server"
+    LOAD_MODULE_ARGS="--loadmodule ${REDIS_MODULE}"
+fi
+
 echo "$STORE_EXEC"
-echo "$REDIS_DIR/redis-server --loglevel warning --loadmodule $REDIS_MODULE --port 6379"
+echo "${REDIS_SERVER} --loglevel warning ${LOAD_MODULE_ARGS} --port 6379"
 echo "$REDIS_DIR/redis-cli -p 6379 shutdown"
 
 # Allow cleanup commands to fail.
 killall plasma_store || true
 $REDIS_DIR/redis-cli -p 6379 shutdown || true
 sleep 1s
-$REDIS_DIR/redis-server --loglevel warning --loadmodule $REDIS_MODULE --port 6379 &
+${REDIS_SERVER} --loglevel warning ${LOAD_MODULE_ARGS} --port 6379 &
 sleep 1s
 
 # Run tests.

src/ray/test/start_raylets.sh

@@ -5,8 +5,24 @@
 # Cause the script to exit if a single command fails.
 set -e
 
+LaunchRedis() {
+    port=$1
+    if [[ "${RAY_USE_NEW_GCS}" = "on" ]]; then
+        ./src/credis/redis/src/redis-server \
+            --loglevel warning \
+            --loadmodule ./src/credis/build/src/libmember.so \
+            --loadmodule ./src/common/redis_module/libray_redis_module.so \
+            --port $port >/dev/null &
+    else
+        ./src/common/thirdparty/redis/src/redis-server \
+            --loglevel warning \
+            --loadmodule ./src/common/redis_module/libray_redis_module.so \
+            --port $port >/dev/null &
+    fi
+}
+
 # Start the GCS.
-./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6379 >/dev/null &
+LaunchRedis 6379
 sleep 1s
 
 if [[ $1 ]]; then