Implement the client table for the new GCS (#1674)

* Add subscription callback to CallbackData

* Implement ClientTable

* Hook up ClientTable to AsyncGCSClient

* Add client_info to GCSClient Connect interface

* client table callbacks

* Unit test for client table

* Doc

* Fix idempotency check

* Fix mac build

* Fix memory issues in gcs client test

* Fix disconnection bug

* lint

src/common/redis_module/ray_redis_module.cc

@@ -420,9 +420,7 @@ int TableAdd_RedisCommand(RedisModuleCtx *ctx,
 
   // Publish a message on the requested pubsub channel if necessary.
   if (pubsub_channel == TablePubsub_TASK) {
-    size_t len = 0;
-    const char *buf = RedisModule_StringPtrLen(data, &len);
-
+    const char *buf = RedisModule_StringPtrLen(data, NULL);
     auto message = flatbuffers::GetRoot<TaskTableData>(buf);
 
     if (message->scheduling_state() == SchedulingState_WAITING ||
@@ -462,6 +460,46 @@ int TableAdd_RedisCommand(RedisModuleCtx *ctx,
       RedisModule_FreeString(ctx, publish_message);
       RedisModule_FreeString(ctx, publish_topic);
     }
+  } else if (pubsub_channel == TablePubsub_CLIENT) {
+    const char *buf = RedisModule_StringPtrLen(data, NULL);
+    auto client_data = flatbuffers::GetRoot<ClientTableData>(buf);
+
+    RedisModuleKey *clients_key = (RedisModuleKey *) RedisModule_OpenKey(
+        ctx, pubsub_channel_str, REDISMODULE_READ | REDISMODULE_WRITE);
+    // If this is a client addition, send all previous notifications, in order.
+    // NOTE(swang): This will go to all clients, so some clients will get
+    // duplicate notifications.
+    if (client_data->is_insertion() &&
+        RedisModule_KeyType(clients_key) != REDISMODULE_KEYTYPE_EMPTY) {
+      // NOTE(swang): Sets are not implemented yet, so we use ZSETs instead.
+      CHECK_ERROR(RedisModule_ZsetFirstInScoreRange(
+                      clients_key, REDISMODULE_NEGATIVE_INFINITE,
+                      REDISMODULE_POSITIVE_INFINITE, 1, 1),
+                  "Unable to initialize zset iterator");
+      do {
+        RedisModuleString *message =
+            RedisModule_ZsetRangeCurrentElement(clients_key, NULL);
+        RedisModuleCallReply *reply =
+            RedisModule_Call(ctx, "PUBLISH", "ss", pubsub_channel_str, message);
+        if (reply == NULL) {
+          RedisModule_CloseKey(clients_key);
+          RedisModule_ReplyWithError(ctx, "error during PUBLISH");
+        }
+      } while (RedisModule_ZsetRangeNext(clients_key));
+    }
+
+    // Append this notification to the past notifications so that it will get
+    // sent to new clients in the future.
+    size_t index = RedisModule_ValueLength(key);
+    RedisModule_ZsetAdd(clients_key, index, data, NULL);
+    // Publish the notification about this client.
+    RedisModuleCallReply *reply =
+        RedisModule_Call(ctx, "PUBLISH", "ss", pubsub_channel_str, data);
+    if (reply == NULL) {
+      RedisModule_ReplyWithError(ctx, "error during PUBLISH");
+    }
+
+    RedisModule_CloseKey(clients_key);
   } else if (pubsub_channel != TablePubsub_NO_PUBLISH) {
     // All other pubsub channels write the data back directly onto the channel.
     RedisModuleCallReply *reply =

src/global_scheduler/global_scheduler.cc

@@ -140,8 +140,14 @@ GlobalSchedulerState *GlobalSchedulerState_init(event_loop *loop,
                          "global_scheduler", node_ip_address,
                          std::vector<std::string>());
   db_attach(state->db, loop, false);
+
+  ClientTableDataT client_info;
+  client_info.client_id = get_db_client_id(state->db).binary();
+  client_info.node_manager_address = std::string(node_ip_address);
+  client_info.local_scheduler_port = 0;
+  client_info.object_manager_port = 0;
   RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr),
-                                         redis_primary_port));
+                                         redis_primary_port, client_info));
   RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(loop));
   state->policy_state = GlobalSchedulerPolicyState_init();
   return state;

src/local_scheduler/local_scheduler.cc

@@ -374,8 +374,14 @@ 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);
+
+    ClientTableDataT client_info;
+    client_info.client_id = get_db_client_id(state->db).binary();
+    client_info.node_manager_address = std::string(node_ip_address);
+    client_info.local_scheduler_port = 0;
+    client_info.object_manager_port = 0;
     RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr),
-                                           redis_primary_port));
+                                           redis_primary_port, client_info));
     RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(loop));
   } else {
     state->db = NULL;

src/plasma/plasma_manager.cc

@@ -486,8 +486,14 @@ 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);
+
+    ClientTableDataT client_info;
+    client_info.client_id = get_db_client_id(state->db).binary();
+    client_info.node_manager_address = std::string(manager_addr);
+    client_info.local_scheduler_port = 0;
+    client_info.object_manager_port = manager_port;
     RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr),
-                                           redis_primary_port));
+                                           redis_primary_port, client_info));
     RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(state->loop));
   } else {
     state->db = NULL;

src/ray/CMakeLists.txt

@@ -33,6 +33,7 @@ set(RAY_SRCS
   status.cc
   gcs/client.cc
   gcs/tables.cc
+  gcs/task_table.cc
   gcs/redis_context.cc
   gcs/asio.cc
   common/client_connection.cc

src/ray/gcs/client.cc

@@ -10,11 +10,17 @@ AsyncGcsClient::AsyncGcsClient() {}
 
 AsyncGcsClient::~AsyncGcsClient() {}
 
-Status AsyncGcsClient::Connect(const std::string &address, int port) {
+Status AsyncGcsClient::Connect(const std::string &address, int port,
+                               const ClientTableDataT &client_info) {
   context_.reset(new RedisContext());
   RAY_RETURN_NOT_OK(context_->Connect(address, port));
   object_table_.reset(new ObjectTable(context_, this));
   task_table_.reset(new TaskTable(context_, this));
+  client_table_.reset(new ClientTable(context_, this, client_info));
+  // 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();
 }
 
@@ -25,25 +31,21 @@ Status Attach(plasma::EventLoop &event_loop) {
 }
 
 Status AsyncGcsClient::Attach(boost::asio::io_service &io_service) {
-  asio_client_.reset(new RedisAsioClient(io_service, context_->async_context()));
+  asio_async_client_.reset(new RedisAsioClient(io_service, context_->async_context()));
+  asio_subscribe_client_.reset(
+      new RedisAsioClient(io_service, context_->subscribe_context()));
   return Status::OK();
 }
 
-ObjectTable &AsyncGcsClient::object_table() {
-  return *object_table_;
-}
+ObjectTable &AsyncGcsClient::object_table() { return *object_table_; }
 
-TaskTable &AsyncGcsClient::task_table() {
-  return *task_table_;
-}
+TaskTable &AsyncGcsClient::task_table() { return *task_table_; }
 
-FunctionTable &AsyncGcsClient::function_table() {
-  return *function_table_;
-}
+ClientTable &AsyncGcsClient::client_table() { return *client_table_; }
 
-ClassTable &AsyncGcsClient::class_table() {
-  return *class_table_;
-}
+FunctionTable &AsyncGcsClient::function_table() { return *function_table_; }
+
+ClassTable &AsyncGcsClient::class_table() { return *class_table_; }
 
 }  // namespace gcs
 

src/ray/gcs/client.h

@@ -22,7 +22,14 @@ class RAY_EXPORT AsyncGcsClient {
   AsyncGcsClient();
   ~AsyncGcsClient();
 
-  Status Connect(const std::string &address, int port);
+  /// Connect to the GCS.
+  ///
+  /// \param address The GCS IP address.
+  /// \param port The GCS port.
+  /// \param client_info Information about the local client to connect.
+  /// \return Status.
+  Status Connect(const std::string &address, int port,
+                 const ClientTableDataT &client_info);
   /// 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);
@@ -38,6 +45,7 @@ class RAY_EXPORT AsyncGcsClient {
   inline ConfigTable &config_table();
   ObjectTable &object_table();
   TaskTable &task_table();
+  ClientTable &client_table();
   inline ErrorTable &error_table();
 
   // We also need something to export generic code to run on workers from the
@@ -45,8 +53,7 @@ class RAY_EXPORT AsyncGcsClient {
 
   using GetExportCallback = std::function<void(const std::string &data)>;
   Status AddExport(const std::string &driver_id, std::string &export_data);
-  Status GetExport(const std::string &driver_id,
-                   int64_t export_index,
+  Status GetExport(const std::string &driver_id, int64_t export_index,
                    const GetExportCallback &done_callback);
 
   std::shared_ptr<RedisContext> context() { return context_; }
@@ -56,29 +63,23 @@ class RAY_EXPORT AsyncGcsClient {
   std::unique_ptr<ClassTable> class_table_;
   std::unique_ptr<ObjectTable> object_table_;
   std::unique_ptr<TaskTable> task_table_;
+  std::unique_ptr<ClientTable> client_table_;
   std::shared_ptr<RedisContext> context_;
-  std::unique_ptr<RedisAsioClient> asio_client_;
+  std::unique_ptr<RedisAsioClient> asio_async_client_;
+  std::unique_ptr<RedisAsioClient> asio_subscribe_client_;
 };
 
 class SyncGcsClient {
-  Status LogEvent(const std::string &key,
-                  const std::string &value,
-                  double timestamp);
+  Status LogEvent(const std::string &key, const std::string &value, double timestamp);
   Status NotifyError(const std::map<std::string, std::string> &error_info);
-  Status RegisterFunction(const JobID &job_id,
-                          const FunctionID &function_id,
-                          const std::string &language,
-                          const std::string &name,
+  Status RegisterFunction(const JobID &job_id, const FunctionID &function_id,
+                          const std::string &language, const std::string &name,
                           const std::string &data);
-  Status RetrieveFunction(const JobID &job_id,
-                          const FunctionID &function_id,
-                          std::string *name,
-                          std::string *data);
+  Status RetrieveFunction(const JobID &job_id, const FunctionID &function_id,
+                          std::string *name, std::string *data);
 
   Status AddExport(const std::string &driver_id, std::string &export_data);
-  Status GetExport(const std::string &driver_id,
-                   int64_t export_index,
-                   std::string *data);
+  Status GetExport(const std::string &driver_id, int64_t export_index, std::string *data);
 };
 
 }  // namespace gcs

src/ray/gcs/client_test.cc

@@ -2,9 +2,9 @@
 
 // TODO(pcm): get rid of this and replace with the type safe plasma event loop
 extern "C" {
+#include "hiredis/adapters/ae.h"
 #include "hiredis/async.h"
 #include "hiredis/hiredis.h"
-#include "hiredis/adapters/ae.h"
 }
 
 #include "ray/gcs/client.h"
@@ -12,22 +12,39 @@ extern "C" {
 
 namespace ray {
 
+/* Flush redis. */
+static inline void flushall_redis(void) {
+  redisContext *context = redisConnect("127.0.0.1", 6379);
+  freeReplyObject(redisCommand(context, "FLUSHALL"));
+  redisFree(context);
+}
+
 class TestGcs : public ::testing::Test {
  public:
   TestGcs() {
-    RAY_CHECK_OK(client_.Connect("127.0.0.1", 6379));
-    job_id_ = UniqueID::from_random();
+    client_ = std::make_shared<gcs::AsyncGcsClient>();
+    ClientTableDataT client_info;
+    client_info.client_id = ClientID::from_random().binary();
+    client_info.node_manager_address = "127.0.0.1";
+    client_info.local_scheduler_port = 0;
+    client_info.object_manager_port = 0;
+    RAY_CHECK_OK(client_->Connect("127.0.0.1", 6379, client_info));
+
+    job_id_ = JobID::from_random();
   }
 
-  virtual ~TestGcs(){};
+  virtual ~TestGcs() {
+    // Clear all keys in the GCS.
+    flushall_redis();
+  };
 
   virtual void Start() = 0;
 
   virtual void Stop() = 0;
 
  protected:
-  gcs::AsyncGcsClient client_;
-  UniqueID job_id_;
+  std::shared_ptr<gcs::AsyncGcsClient> client_;
+  JobID job_id_;
 };
 
 TestGcs *test;
@@ -36,9 +53,13 @@ class TestGcsWithAe : public TestGcs {
  public:
   TestGcsWithAe() {
     loop_ = aeCreateEventLoop(1024);
-    RAY_CHECK_OK(client_.context()->AttachToEventLoop(loop_));
+    RAY_CHECK_OK(client_->context()->AttachToEventLoop(loop_));
+  }
+  ~TestGcsWithAe() override {
+    // Destroy the client first since it has a reference to the event loop.
+    client_.reset();
+    aeDeleteEventLoop(loop_);
   }
-  ~TestGcsWithAe() override { aeDeleteEventLoop(loop_); }
   void Start() override { aeMain(loop_); }
   void Stop() override { aeStop(loop_); }
 
@@ -48,35 +69,42 @@ class TestGcsWithAe : public TestGcs {
 
 class TestGcsWithAsio : public TestGcs {
  public:
-  TestGcsWithAsio() { RAY_CHECK_OK(client_.Attach(io_service_)); }
+  TestGcsWithAsio() : TestGcs(), io_service_(), work_(io_service_) {
+    RAY_CHECK_OK(client_->Attach(io_service_));
+  }
+  ~TestGcsWithAsio() {
+    // Destroy the client first since it has a reference to the event loop.
+    client_.reset();
+  }
   void Start() override { io_service_.run(); }
   void Stop() override { io_service_.stop(); }
 
  private:
   boost::asio::io_service io_service_;
+  // Give the event loop some work so that it's forced to run until Stop() is
+  // called.
+  boost::asio::io_service::work work_;
 };
 
-void ObjectAdded(gcs::AsyncGcsClient *client,
-                 const UniqueID &id,
+void ObjectAdded(gcs::AsyncGcsClient *client, const UniqueID &id,
                  std::shared_ptr<ObjectTableDataT> data) {
   ASSERT_EQ(data->managers, std::vector<std::string>({"A", "B"}));
 }
 
-void Lookup(gcs::AsyncGcsClient *client,
-            const UniqueID &id,
+void Lookup(gcs::AsyncGcsClient *client, const UniqueID &id,
             std::shared_ptr<ObjectTableDataT> data) {
   // Check that the object entry was added.
   ASSERT_EQ(data->managers, std::vector<std::string>({"A", "B"}));
   test->Stop();
 }
 
-void TestObjectTable(const UniqueID &job_id, gcs::AsyncGcsClient &client) {
+void TestObjectTable(const JobID &job_id, std::shared_ptr<gcs::AsyncGcsClient> client) {
   auto data = std::make_shared<ObjectTableDataT>();
   data->managers.push_back("A");
   data->managers.push_back("B");
   ObjectID object_id = ObjectID::from_random();
-  RAY_CHECK_OK(client.object_table().Add(job_id, object_id, data, &ObjectAdded));
-  RAY_CHECK_OK(client.object_table().Lookup(job_id, object_id, &Lookup));
+  RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, &ObjectAdded));
+  RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, &Lookup));
   // Run the event loop. The loop will only stop if the Lookup callback is
   // called (or an assertion failure).
   test->Start();
@@ -92,41 +120,36 @@ TEST_F(TestGcsWithAsio, TestObjectTable) {
   TestObjectTable(job_id_, client_);
 }
 
-void TaskAdded(gcs::AsyncGcsClient *client,
-               const TaskID &id,
+void TaskAdded(gcs::AsyncGcsClient *client, const TaskID &id,
                std::shared_ptr<TaskTableDataT> data) {
   ASSERT_EQ(data->scheduling_state, SchedulingState_SCHEDULED);
 }
 
-void TaskLookup(gcs::AsyncGcsClient *client,
-                const TaskID &id,
+void TaskLookup(gcs::AsyncGcsClient *client, const TaskID &id,
                 std::shared_ptr<TaskTableDataT> data) {
   ASSERT_EQ(data->scheduling_state, SchedulingState_SCHEDULED);
 }
 
-void TaskLookupAfterUpdate(gcs::AsyncGcsClient *client,
-                           const TaskID &id,
+void TaskLookupAfterUpdate(gcs::AsyncGcsClient *client, const TaskID &id,
                            std::shared_ptr<TaskTableDataT> data) {
   ASSERT_EQ(data->scheduling_state, SchedulingState_LOST);
   test->Stop();
 }
 
-void TaskUpdateCallback(gcs::AsyncGcsClient *client,
-                        const TaskID &task_id,
-                        const TaskTableDataT &task,
-                        bool updated) {
-  RAY_CHECK_OK(client->task_table().Lookup(DriverID::nil(), task_id,
-                                           &TaskLookupAfterUpdate));
+void TaskUpdateCallback(gcs::AsyncGcsClient *client, const TaskID &task_id,
+                        const TaskTableDataT &task, bool updated) {
+  RAY_CHECK_OK(
+      client->task_table().Lookup(DriverID::nil(), task_id, &TaskLookupAfterUpdate));
 }
 
-void TestTaskTable(const UniqueID &job_id, gcs::AsyncGcsClient &client) {
+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");
   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));
-  RAY_CHECK_OK(client.task_table().Lookup(job_id, task_id, &TaskLookup));
+  RAY_CHECK_OK(client->task_table().Add(job_id, task_id, data, &TaskAdded));
+  RAY_CHECK_OK(client->task_table().Lookup(job_id, task_id, &TaskLookup));
   auto update = std::make_shared<TaskTableTestAndUpdateT>();
   update->test_scheduler_id = local_scheduler_id.binary();
   update->test_state_bitmask = SchedulingState_SCHEDULED;
@@ -134,7 +157,7 @@ void TestTaskTable(const UniqueID &job_id, gcs::AsyncGcsClient &client) {
   // After test-and-setting, the callback will lookup the current state of the
   // task.
   RAY_CHECK_OK(
-      client.task_table().TestAndUpdate(job_id, task_id, update, &TaskUpdateCallback));
+      client->task_table().TestAndUpdate(job_id, task_id, update, &TaskUpdateCallback));
   // Run the event loop. The loop will only stop if the lookup after the
   // test-and-set succeeds (or an assertion failure).
   test->Start();
@@ -155,11 +178,11 @@ void ObjectTableSubscribed(gcs::AsyncGcsClient *client, const UniqueID &id,
   test->Stop();
 }
 
-void TestSubscribeAll(const UniqueID &job_id, gcs::AsyncGcsClient &client) {
+void TestSubscribeAll(const JobID &job_id, std::shared_ptr<gcs::AsyncGcsClient> client) {
   // Subscribe to all object table notifications. The registered callback for
   // notifications will check whether the object below is added.
-  RAY_CHECK_OK(client.object_table().Subscribe(job_id, ClientID::nil(), &Lookup,
-                                               &ObjectTableSubscribed));
+  RAY_CHECK_OK(client->object_table().Subscribe(job_id, ClientID::nil(), &Lookup,
+                                                &ObjectTableSubscribed));
   // Run the event loop. The loop will only stop if the subscription succeeds.
   test->Start();
 
@@ -168,7 +191,7 @@ void TestSubscribeAll(const UniqueID &job_id, gcs::AsyncGcsClient &client) {
   data->managers.push_back("A");
   data->managers.push_back("B");
   ObjectID object_id = ObjectID::from_random();
-  RAY_CHECK_OK(client.object_table().Add(job_id, object_id, data, &ObjectAdded));
+  RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, &ObjectAdded));
   // Run the event loop. The loop will only stop if the registered subscription
   // callback is called (or an assertion failure).
   test->Start();
@@ -184,4 +207,63 @@ TEST_F(TestGcsWithAsio, TestSubscribeAll) {
   TestSubscribeAll(job_id_, client_);
 }
 
+void ClientTableNotification(gcs::AsyncGcsClient *client, const UniqueID &id,
+                             std::shared_ptr<ClientTableDataT> data, bool is_insertion) {
+  ClientID added_id = client->client_table().GetLocalClientId();
+  ASSERT_EQ(ClientID::from_binary(data->client_id), added_id);
+  ASSERT_EQ(data->is_insertion, is_insertion);
+
+  auto cached_client = client->client_table().GetClient(added_id);
+  ASSERT_EQ(ClientID::from_binary(cached_client.client_id), added_id);
+  ASSERT_EQ(cached_client.is_insertion, is_insertion);
+}
+
+void TestClientTableConnect(const JobID &job_id,
+                            std::shared_ptr<gcs::AsyncGcsClient> client) {
+  // Register callbacks for when a client gets added and removed. The latter
+  // event will stop the event loop.
+  client->client_table().RegisterClientAddedCallback(
+      [](gcs::AsyncGcsClient *client, const UniqueID &id,
+         std::shared_ptr<ClientTableDataT> data) {
+        ClientTableNotification(client, id, data, true);
+        test->Stop();
+      });
+  // Connect and disconnect to client table. We should receive notifications
+  // for the addition and removal of our own entry.
+  RAY_CHECK_OK(client->client_table().Connect());
+  test->Start();
+}
+
+TEST_F(TestGcsWithAsio, TestClientTableConnect) {
+  test = this;
+  TestClientTableConnect(job_id_, client_);
+}
+
+void TestClientTableDisconnect(const JobID &job_id,
+                               std::shared_ptr<gcs::AsyncGcsClient> client) {
+  // Register callbacks for when a client gets added and removed. The latter
+  // event will stop the event loop.
+  client->client_table().RegisterClientAddedCallback(
+      [](gcs::AsyncGcsClient *client, const UniqueID &id,
+         std::shared_ptr<ClientTableDataT> data) {
+        ClientTableNotification(client, id, data, true);
+      });
+  client->client_table().RegisterClientRemovedCallback(
+      [](gcs::AsyncGcsClient *client, const UniqueID &id,
+         std::shared_ptr<ClientTableDataT> data) {
+        ClientTableNotification(client, id, data, false);
+        test->Stop();
+      });
+  // Connect and disconnect to client table. We should receive notifications
+  // for the addition and removal of our own entry.
+  RAY_CHECK_OK(client->client_table().Connect());
+  RAY_CHECK_OK(client->client_table().Disconnect());
+  test->Start();
+}
+
+TEST_F(TestGcsWithAsio, TestClientTableDisconnect) {
+  test = this;
+  TestClientTableDisconnect(job_id_, client_);
+}
+
 }  // namespace

src/ray/gcs/format/gcs.fbs

@@ -74,13 +74,6 @@ table CustomSerializerData {
 table ConfigTableData {
 }
 
-table Resource {
-  // The type of the resource.
-  resource_name: string;
-  // The total capacity of this resource type.
-  resource_capacity: double;
-}
-
 table ClientTableData {
   // The client ID of the client that the message is about.
   client_id: string;
@@ -88,17 +81,22 @@ table ClientTableData {
   node_manager_address: string;
   // The port at which the client's node manager is listening for TCP
   // connections from other node managers.
-  node_manager_port: int;
+  local_scheduler_port: int;
   // The port at which the client's object manager is listening for TCP
   // connections from other object managers.
   object_manager_port: int;
-  // The total resources of this client.
-  resources_total: [Resource];
   // True if the message is about the addition of a client and false if it is
   // about the deletion of a client.
   is_insertion: bool;
 }
 
+table Resource {
+  // The type of the resource.
+  resource_name: string;
+  // The total capacity of this resource type.
+  resource_capacity: double;
+}
+
 table NodeManagerHeartbeat {
   // The available resources on this node manager. This information may be
   // stale.

src/ray/gcs/redis_context.h

@@ -56,6 +56,7 @@ class RedisContext {
   Status SubscribeAsync(const ClientID &client_id, const TablePubsub pubsub_channel,
                         int64_t callback_index);
   redisAsyncContext *async_context() { return async_context_; }
+  redisAsyncContext *subscribe_context() { return subscribe_context_; };
 
  private:
   redisContext *context_;

src/ray/gcs/tables.cc

@@ -2,101 +2,139 @@
 
 #include "ray/gcs/client.h"
 
-#include "common_protocol.h"
-
-namespace {
-
-std::shared_ptr<TaskTableDataT> MakeTaskTableData(const TaskExecutionSpec &execution_spec,
-                                                  const ClientID &local_scheduler_id,
-                                                  SchedulingState scheduling_state) {
-  auto data = std::make_shared<TaskTableDataT>();
-  data->scheduling_state = scheduling_state;
-  data->task_info =
-      std::string(execution_spec.Spec(), execution_spec.SpecSize());
-  data->scheduler_id = local_scheduler_id.binary();
-
-  flatbuffers::FlatBufferBuilder fbb;
-  auto execution_dependencies = CreateTaskExecutionDependencies(
-      fbb, to_flatbuf(fbb, execution_spec.ExecutionDependencies()));
-  fbb.Finish(execution_dependencies);
-
-  data->execution_dependencies =
-      std::string((const char *) fbb.GetBufferPointer(), fbb.GetSize());
-  data->spillback_count = execution_spec.SpillbackCount();
-
-  return data;
-}
-
-}  // namespace
-
 namespace ray {
 
 namespace gcs {
 
-// TODO(pcm): This is a helper method that should go away once we get rid of
-// the Task* datastructure and replace it with TaskTableDataT.
-Status TaskTableAdd(AsyncGcsClient *gcs_client, Task *task) {
-  TaskExecutionSpec &execution_spec = *Task_task_execution_spec(task);
-  TaskSpec *spec = execution_spec.Spec();
-  auto data = MakeTaskTableData(execution_spec, Task_local_scheduler(task),
-                                static_cast<SchedulingState>(Task_state(task)));
-  return gcs_client->task_table().Add(
-      ray::JobID::nil(), TaskSpec_task_id(spec), data,
-      [](gcs::AsyncGcsClient *client, const TaskID &id,
-         std::shared_ptr<TaskTableDataT> data) {});
+void ClientTable::RegisterClientAddedCallback(const Callback &callback) {
+  client_added_callback_ = callback;
+  // Call the callback for any added clients that are cached.
+  for (const auto &entry : client_cache_) {
+    if (!entry.first.is_nil() && entry.second.is_insertion) {
+      auto data = std::make_shared<ClientTableDataT>(entry.second);
+      client_added_callback_(client_, entry.first, data);
+    }
+  }
 }
 
-// TODO(pcm): This is a helper method that should go away once we get rid of
-// the Task* datastructure and replace it with TaskTableDataT.
-Status TaskTableTestAndUpdate(AsyncGcsClient *gcs_client, const TaskID &task_id,
-                              const ClientID &local_scheduler_id, int test_state_bitmask,
-                              SchedulingState update_state,
-                              const TaskTable::TestAndUpdateCallback &callback) {
-  auto data = std::make_shared<TaskTableTestAndUpdateT>();
-  data->test_scheduler_id = local_scheduler_id.binary();
-  data->test_state_bitmask = test_state_bitmask;
-  data->update_state = update_state;
-  return gcs_client->task_table().TestAndUpdate(ray::JobID::nil(), task_id,
-                                                data, callback);
+void ClientTable::RegisterClientRemovedCallback(const Callback &callback) {
+  client_removed_callback_ = callback;
+  // Call the callback for any removed clients that are cached.
+  for (const auto &entry : client_cache_) {
+    if (!entry.first.is_nil() && !entry.second.is_insertion) {
+      auto data = std::make_shared<ClientTableDataT>(entry.second);
+      client_removed_callback_(client_, entry.first, data);
+    }
+  }
 }
 
-void ClientConnected(gcs::AsyncGcsClient *client, const ClientID &client_id,
-                     std::shared_ptr<ObjectTableDataT> data) {}
+void ClientTable::HandleNotification(AsyncGcsClient *client, const ClientID &channel_id,
+                                     std::shared_ptr<ClientTableDataT> data) {
+  ClientID client_id = ClientID::from_binary(data->client_id);
+  // It's possible to get duplicate notifications from the client table, so
+  // check whether this notification is new.
+  auto entry = client_cache_.find(client_id);
+  bool is_new;
+  if (entry == client_cache_.end()) {
+    // If the entry is not in the cache, then the notification is new.
+    is_new = true;
+  } else {
+    // If the entry is in the cache, then the notification is new if the client
+    // was alive and is now dead.
+    bool was_inserted = entry->second.is_insertion;
+    bool is_deleted = !data->is_insertion;
+    is_new = (was_inserted && is_deleted);
+    // Once a client with a given ID has been removed, it should never be added
+    // again. If the entry was in the cache and the client was deleted, check
+    // that this new notification is not an insertion.
+    if (!entry->second.is_insertion) {
+      RAY_CHECK(!data->is_insertion)
+          << "Notification for addition of a client that was already removed:"
+          << client_id.hex();
+    }
+  }
 
-ClientTable::ClientTable(const std::shared_ptr<RedisContext> &context,
-                         AsyncGcsClient *client)
-    : Table(context, client), client_id_(UniqueID::from_random()) {}
+  // Add the notification to our cache. Notifications are idempotent.
+  client_cache_[client_id] = *data;
 
-Status ClientTable::Connect(ClientID *client_id) {
-  auto data = std::make_shared<ClientTableDataT>();
-  data->client_id = client_id_.binary();
-  // TODO(swang): Get the address and port from somewhere.
-  data->node_manager_address = "";
-  data->node_manager_port = 0;
+  // If the notification is new, call any registered callbacks.
+  if (is_new) {
+    if (data->is_insertion) {
+      if (client_added_callback_ != nullptr) {
+        client_added_callback_(client, client_id, data);
+      }
+    } else {
+      if (client_removed_callback_ != nullptr) {
+        client_removed_callback_(client, client_id, data);
+      }
+    }
+  }
+}
+
+void ClientTable::HandleConnected(AsyncGcsClient *client, const ClientID &client_id,
+                                  std::shared_ptr<ClientTableDataT> data) {
+  RAY_CHECK(client_id == client_id_) << client_id.hex() << " " << client_id_.hex();
+}
+
+const ClientID &ClientTable::GetLocalClientId() { return client_id_; }
+
+const ClientTableDataT &ClientTable::GetLocalClient() { return local_client_; }
+
+Status ClientTable::Connect() {
+  RAY_CHECK(!disconnected_) << "Tried to reconnect a disconnected client.";
+
+  auto data = std::make_shared<ClientTableDataT>(local_client_);
   data->is_insertion = true;
-
-  // TODO(swang):
-  // - Add ourselves to the client table (easier if this is synchronous).
-  // - Subscribe to the client table.
-  // - Once subscription is complete, read all client table entries once.
-
-  *client_id = client_id_;
-  return Status::OK();
+  // Callback for a notification from the client table.
+  auto notification_callback = [this](AsyncGcsClient *client, const ClientID &channel_id,
+                                      std::shared_ptr<ClientTableDataT> data) {
+    return HandleNotification(client, channel_id, data);
+  };
+  // Callback to handle our own successful connection once we've added
+  // ourselves.
+  auto add_callback = [this](AsyncGcsClient *client, const ClientID &id,
+                             std::shared_ptr<ClientTableDataT> data) {
+    HandleConnected(client, id, data);
+  };
+  // Callback to add ourselves once we've successfully subscribed.
+  auto subscription_callback = [this, data, add_callback](
+      AsyncGcsClient *c, const ClientID &id, std::shared_ptr<ClientTableDataT> d) {
+    // Mark ourselves as deleted if we called Disconnect() since the last
+    // Connect() call.
+    if (disconnected_) {
+      data->is_insertion = false;
+    }
+    return Add(JobID::nil(), client_id_, data, add_callback);
+  };
+  return Subscribe(JobID::nil(), ClientID::nil(), notification_callback,
+                   subscription_callback);
 }
 
 Status ClientTable::Disconnect() {
-  auto data = std::make_shared<ClientTableDataT>();
-  data->client_id = client_id_.binary();
-  // TODO(swang): Get the address and port from somewhere.
-  data->node_manager_address = "";
-  data->node_manager_port = 0;
-  data->is_insertion = false;
-
-  // TODO(swang):
-  // - Add ourselves to the client table (easier if this is synchronous).
+  auto data = std::make_shared<ClientTableDataT>(local_client_);
+  data->is_insertion = true;
+  auto add_callback = [this](AsyncGcsClient *client, const ClientID &id,
+                             std::shared_ptr<ClientTableDataT> data) {
+    HandleConnected(client, id, data);
+  };
+  RAY_RETURN_NOT_OK(Add(JobID::nil(), client_id_, data, add_callback));
+  // We successfully added the deletion entry. Mark ourselves as disconnected.
+  disconnected_ = true;
   return Status::OK();
 }
 
+const ClientTableDataT &ClientTable::GetClient(const ClientID &client_id) {
+  RAY_CHECK(!client_id.is_nil());
+  auto entry = client_cache_.find(client_id);
+  if (entry != client_cache_.end()) {
+    return entry->second;
+  } else {
+    // If the requested client was not found, return a reference to the nil
+    // client entry.
+    return client_cache_[ClientID::nil()];
+  }
+}
+
 }  // namespace gcs
 
 }  // namespace ray

src/ray/gcs/tables.h

@@ -30,13 +30,16 @@ template <typename ID, typename Data>
 class Table {
  public:
   using DataT = typename Data::NativeTableType;
-  using Callback = std::function<
-      void(AsyncGcsClient *client, const ID &id, std::shared_ptr<DataT> data)>;
+  using Callback = std::function<void(AsyncGcsClient *client, const ID &id,
+                                      std::shared_ptr<DataT> data)>;
 
   struct CallbackData {
     ID id;
     std::shared_ptr<DataT> data;
     Callback callback;
+    // An optional callback to call for subscription operations, where the
+    // first message is a notification of subscription success.
+    Callback subscription_callback;
     Table<ID, Data> *table;
     AsyncGcsClient *client;
   };
@@ -49,16 +52,19 @@ class Table {
   /// \param job_id The ID of the job (= driver).
   /// \param id The ID of the data that is added to the GCS.
   /// \param data Data that is added to the GCS.
-  /// \param done Callback that is called once the data has been written to the GCS.
+  /// \param done Callback that is called once the data has been written to the
+  ///        GCS.
   /// \return Status
-  Status Add(const JobID &job_id,
-             const ID &id,
-             std::shared_ptr<DataT> data,
+  Status Add(const JobID &job_id, const ID &id, std::shared_ptr<DataT> data,
              const Callback &done) {
     auto d = std::shared_ptr<CallbackData>(
-        new CallbackData({id, data, done, this, client_}));
-    int64_t callback_index = RedisCallbackManager::instance().add([d](
-        const std::string &data) { (d->callback)(d->client, d->id, d->data); });
+        new CallbackData({id, data, done, nullptr, this, client_}));
+    int64_t callback_index =
+        RedisCallbackManager::instance().add([d](const std::string &data) {
+          if (d->callback != nullptr) {
+            (d->callback)(d->client, d->id, d->data);
+          }
+        });
     flatbuffers::FlatBufferBuilder fbb;
     fbb.ForceDefaults(true);
     fbb.Finish(Data::Pack(fbb, data.get()));
@@ -75,13 +81,15 @@ class Table {
   /// \return Status
   Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup) {
     auto d = std::shared_ptr<CallbackData>(
-        new CallbackData({id, nullptr, lookup, this}));
+        new CallbackData({id, nullptr, lookup, nullptr, this, client_}));
     int64_t callback_index =
         RedisCallbackManager::instance().add([d](const std::string &data) {
           auto result = std::make_shared<DataT>();
           auto root = flatbuffers::GetRoot<Data>(data.data());
           root->UnPackTo(result.get());
-          (d->callback)(d->client, d->id, result);
+          if (d->callback != nullptr) {
+            (d->callback)(d->client, d->id, result);
+          }
         });
     std::vector<uint8_t> nil;
     RAY_RETURN_NOT_OK(context_->RunAsync("RAY.TABLE_LOOKUP", id, nil.data(), nil.size(),
@@ -102,14 +110,17 @@ class Table {
   Status Subscribe(const JobID &job_id, const ClientID &client_id,
                    const Callback &subscribe, const Callback &done) {
     auto d = std::shared_ptr<CallbackData>(
-        new CallbackData({client_id, nullptr, subscribe, this}));
+        new CallbackData({client_id, nullptr, subscribe, done, this, client_}));
     int64_t callback_index =
-        RedisCallbackManager::instance().add([done, d](const std::string &data) {
+        RedisCallbackManager::instance().add([d](const std::string &data) {
           if (data.empty()) {
             // No data is provided. This is the callback for the initial
             // subscription request.
-            done(d->client, d->id, nullptr);
+            if (d->subscription_callback != nullptr) {
+              (d->subscription_callback)(d->client, d->id, nullptr);
+            }
           } else {
+            // Data is provided. This is the callback for a message.
             auto result = std::make_shared<DataT>();
             auto root = flatbuffers::GetRoot<Data>(data.data());
             root->UnPackTo(result.get());
@@ -124,8 +135,7 @@ class Table {
   Status Remove(const JobID &job_id, const ID &id, const Callback &done);
 
  protected:
-  std::unordered_map<ID, std::unique_ptr<CallbackData>, UniqueIDHasher>
-      callback_data_;
+  std::unordered_map<ID, std::unique_ptr<CallbackData>, UniqueIDHasher> callback_data_;
   std::shared_ptr<RedisContext> context_;
   AsyncGcsClient *client_;
   TablePubsub pubsub_channel_;
@@ -149,10 +159,8 @@ class ObjectTable : public Table<ObjectID, ObjectTableData> {
   /// \param done_callback Callback to be called when subscription is installed.
   ///        This is only used for the tests.
   /// \return Status
-  Status SubscribeToNotifications(const JobID &job_id,
-                                  bool subscribe_all,
-                                  const Callback &object_available,
-                                  const Callback &done);
+  Status SubscribeToNotifications(const JobID &job_id, bool subscribe_all,
+                                  const Callback &object_available, const Callback &done);
 
   /// Request notifications about the availability of some objects from the
   /// object
@@ -180,10 +188,9 @@ class TaskTable : public Table<TaskID, TaskTableData> {
     pubsub_channel_ = TablePubsub_TASK;
   };
 
-  using TestAndUpdateCallback = std::function<void(AsyncGcsClient *client,
-                                                   const TaskID &id,
-                                                   const TaskTableDataT &task,
-                                                   bool updated)>;
+  using TestAndUpdateCallback =
+      std::function<void(AsyncGcsClient *client, const TaskID &id,
+                         const TaskTableDataT &task, bool updated)>;
   using SubscribeToTaskCallback =
       std::function<void(std::shared_ptr<TaskTableDataT> task)>;
   /// Update a task's scheduling information in the task table, if the current
@@ -201,8 +208,7 @@ class TaskTable : public Table<TaskID, TaskTableData> {
   ///        with, if the current state matches test_state_bitmask.
   /// \param callback Function to be called when database returns result.
   /// \return Status
-  Status TestAndUpdate(const JobID &job_id,
-                       const TaskID &id,
+  Status TestAndUpdate(const JobID &job_id, const TaskID &id,
                        std::shared_ptr<TaskTableTestAndUpdateT> data,
                        const TestAndUpdateCallback &callback) {
     int64_t callback_index = RedisCallbackManager::instance().add(
@@ -213,7 +219,6 @@ class TaskTable : public Table<TaskID, TaskTableData> {
           callback(client_, id, *result, root->updated());
         });
     flatbuffers::FlatBufferBuilder fbb;
-    TaskTableTestAndUpdateBuilder builder(fbb);
     fbb.Finish(TaskTableTestAndUpdate::Pack(fbb, data.get()));
     RAY_RETURN_NOT_OK(context_->RunAsync("RAY.TABLE_TEST_AND_UPDATE", id,
                                          fbb.GetBufferPointer(), fbb.GetSize(),
@@ -231,7 +236,8 @@ class TaskTable : public Table<TaskID, TaskTableData> {
   ///        events we want to listen to. If you want to subscribe to updates
   ///        from
   ///        all local schedulers, pass in NIL_ID.
-  /// \param subscribe_callback Callback that will be called when the task table is
+  /// \param subscribe_callback Callback that will be called when the task table
+  /// is
   ///        updated.
   /// \param state_filter Events we want to listen to. Can have values from the
   ///        enum "scheduling_state" in task.h.
@@ -257,50 +263,28 @@ Status TaskTableTestAndUpdate(AsyncGcsClient *gcs_client, const TaskID &task_id,
                               SchedulingState update_state,
                               const TaskTable::TestAndUpdateCallback &callback);
 
-/// \class ClientInformation
-///
-/// Represents information in the client table about a particular client. Each
-/// client has an associated node manager.
-class ClientInformation {
- public:
-  /// Create a client information object.
-  ///
-  /// \param client_table_entry A serialized client table entry flatbuffer.
-  ClientInformation(const ClientTableData &client_table_entry);
-
-  /// Get the client ID.
-  ///
-  /// \return The ID of this client.
-  const ClientID &GetClientId() const;
-
-  /// Get the IP address of the client's node manager.
-  ///
-  /// \return The IP address of the client's node manager.
-  const std::string GetIpAddress() const;
-
-  /// Get the port at which the client's node manager is listening for
-  /// TCP connections.
-  ///
-  /// \return The client's TCP port.
-  int GetPort() const;
-
-  /// Get whether the client is alive.
-  ///
-  /// \return Whether the client is alive.
-  bool IsAlive() const;
-};
-
 class ClientTable : private Table<ClientID, ClientTableData> {
  public:
-  ClientTable(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client);
+  ClientTable(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client,
+              const ClientTableDataT &local_client)
+      : Table(context, client),
+        disconnected_(false),
+        client_id_(ClientID::from_binary(local_client.client_id)),
+        local_client_(local_client) {
+    pubsub_channel_ = TablePubsub_CLIENT;
 
-  /// Connect as a client to the GCS. This registers us in the client table and
-  /// begins subscription to client table notifications.
+    // Add a nil client to the cache so that we can serve requests for clients
+    // that we have not heard about.
+    ClientTableDataT nil_client;
+    nil_client.client_id = ClientID::nil().binary();
+    client_cache_[ClientID::nil()] = nil_client;
+  };
+
+  /// Connect as a client to the GCS. This registers us in the client table
+  /// and begins subscription to client table notifications.
   ///
-  /// \param[out] client_id The assigned client ID will be written to this pointer.
   /// \return Status
-  // TODO(swang): Call this from AsyncGcsClient::Connect?
-  ray::Status Connect(ClientID *client_id);
+  ray::Status Connect();
 
   /// Disconnect the client from the GCS. The client ID assigned during
   /// registration should never be reused after disconnecting.
@@ -308,16 +292,54 @@ class ClientTable : private Table<ClientID, ClientTableData> {
   /// \return Status
   ray::Status Disconnect();
 
-  /// Get a client's information from the cache.
+  /// Register a callback to call when a new client is added.
+  ///
+  /// \param callback The callback to register.
+  void RegisterClientAddedCallback(const Callback &callback);
+
+  /// Register a callback to call when a client is removed.
+  ///
+  /// \param callback The callback to register.
+  void RegisterClientRemovedCallback(const Callback &callback);
+
+  /// Get a client's information from the cache. The cache only contains
+  /// information for clients that we've heard a notification for.
   ///
   /// \param client The client to get information about.
-  const ClientInformation &GetClientInformation(const ClientID &client);
+  /// \return A reference to the requested client. If the client is not in the
+  ///         cache, then an entry with a nil ClientID will be returned.
+  const ClientTableDataT &GetClient(const ClientID &client);
+
+  /// Get the local client's ID.
+  ///
+  /// \return The local client's ID.
+  const ClientID &GetLocalClientId();
+
+  /// Get the local client's information.
+  ///
+  /// \return The local client's information.
+  const ClientTableDataT &GetLocalClient();
 
  private:
+  /// Handle a client table notification.
+  void HandleNotification(AsyncGcsClient *client, const ClientID &channel_id,
+                          std::shared_ptr<ClientTableDataT>);
+  /// Handle this client's successful connection to the GCS.
+  void HandleConnected(AsyncGcsClient *client, const ClientID &client_id,
+                       std::shared_ptr<ClientTableDataT>);
+
+  /// Whether this client has called Disconnect().
+  bool disconnected_;
   /// This client's ID.
-  ClientID client_id_;
+  const ClientID client_id_;
+  /// Information about this client.
+  ClientTableDataT local_client_;
+  /// The callback to call when a new client is added.
+  Callback client_added_callback_;
+  /// The callback to call when a client is removed.
+  Callback client_removed_callback_;
   /// A cache for information about all clients.
-  std::unordered_map<ClientID, ClientInformation, UniqueIDHasher> client_cache_;
+  std::unordered_map<ClientID, ClientTableDataT, UniqueIDHasher> client_cache_;
 };
 
 }  // namespace gcs

src/ray/gcs/task_table.cc

@@ -0,0 +1,68 @@
+#include "ray/gcs/tables.h"
+
+#include "ray/gcs/client.h"
+
+#include "common_protocol.h"
+#include "task.h"
+
+// TODO(swang): This file extends tables.cc so that we can separate out the
+// part that depends on the Task* datasturcture from the build. This should be
+// merged with tables.cc once we get rid of the Task* datastructure.
+
+namespace {
+
+std::shared_ptr<TaskTableDataT> MakeTaskTableData(const TaskExecutionSpec &execution_spec,
+                                                  const ClientID &local_scheduler_id,
+                                                  SchedulingState scheduling_state) {
+  auto data = std::make_shared<TaskTableDataT>();
+  data->scheduling_state = scheduling_state;
+  data->task_info = std::string(execution_spec.Spec(), execution_spec.SpecSize());
+  data->scheduler_id = local_scheduler_id.binary();
+
+  flatbuffers::FlatBufferBuilder fbb;
+  auto execution_dependencies = CreateTaskExecutionDependencies(
+      fbb, to_flatbuf(fbb, execution_spec.ExecutionDependencies()));
+  fbb.Finish(execution_dependencies);
+
+  data->execution_dependencies =
+      std::string((const char *)fbb.GetBufferPointer(), fbb.GetSize());
+  data->spillback_count = execution_spec.SpillbackCount();
+
+  return data;
+}
+
+}  // namespace
+
+namespace ray {
+
+namespace gcs {
+
+// TODO(pcm): This is a helper method that should go away once we get rid of
+// the Task* datastructure and replace it with TaskTableDataT.
+Status TaskTableAdd(AsyncGcsClient *gcs_client, Task *task) {
+  TaskExecutionSpec &execution_spec = *Task_task_execution_spec(task);
+  TaskSpec *spec = execution_spec.Spec();
+  auto data = MakeTaskTableData(execution_spec, Task_local_scheduler(task),
+                                static_cast<SchedulingState>(Task_state(task)));
+  return gcs_client->task_table().Add(ray::JobID::nil(), TaskSpec_task_id(spec), data,
+                                      [](gcs::AsyncGcsClient *client, const TaskID &id,
+                                         std::shared_ptr<TaskTableDataT> data) {});
+}
+
+// TODO(pcm): This is a helper method that should go away once we get rid of
+// the Task* datastructure and replace it with TaskTableDataT.
+Status TaskTableTestAndUpdate(AsyncGcsClient *gcs_client, const TaskID &task_id,
+                              const ClientID &local_scheduler_id, int test_state_bitmask,
+                              SchedulingState update_state,
+                              const TaskTable::TestAndUpdateCallback &callback) {
+  auto data = std::make_shared<TaskTableTestAndUpdateT>();
+  data->test_scheduler_id = local_scheduler_id.binary();
+  data->test_state_bitmask = test_state_bitmask;
+  data->update_state = update_state;
+  return gcs_client->task_table().TestAndUpdate(ray::JobID::nil(), task_id, data,
+                                                callback);
+}
+
+}  // namespace gcs
+
+}  // namespace ray