[GCS] refactor the GCS Client Object Interface (#5695)

BUILD.bazel

@@ -843,6 +843,22 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "redis_object_info_accessor_test",
+    srcs = ["src/ray/gcs/test/redis_object_info_accessor_test.cc"],
+    args = ["$(location redis-server) $(location redis-cli) $(location libray_redis_module.so)"],
+    copts = COPTS,
+    data = [
+        "//:libray_redis_module.so",
+        "//:redis-cli",
+        "//:redis-server",
+    ],
+    deps = [
+        ":gcs",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_test(
     name = "subscription_executor_test",
     srcs = ["src/ray/gcs/test/subscription_executor_test.cc"],

src/ray/gcs/accessor.h

@@ -3,6 +3,7 @@
 
 #include "ray/common/id.h"
 #include "ray/gcs/callback.h"
+#include "ray/gcs/entry_change_notification.h"
 #include "ray/protobuf/gcs.pb.h"
 
 namespace ray {
@@ -176,6 +177,64 @@ class TaskInfoAccessor {
   TaskInfoAccessor() = default;
 };
 
+/// `ObjectInfoAccessor` is a sub-interface of `GcsClient`.
+/// This class includes all the methods that are related to accessing
+/// object information in the GCS.
+class ObjectInfoAccessor {
+ public:
+  virtual ~ObjectInfoAccessor() {}
+
+  /// Get object's locations from GCS asynchronously.
+  ///
+  /// \param object_id The ID of object to lookup in GCS.
+  /// \param callback Callback that will be called after lookup finishes.
+  /// \return Status
+  virtual Status AsyncGetLocations(
+      const ObjectID &object_id,
+      const MultiItemCallback<rpc::ObjectTableData> &callback) = 0;
+
+  /// Add location of object to GCS asynchronously.
+  ///
+  /// \param object_id The ID of object which location will be added to GCS.
+  /// \param node_id The location that will be added to GCS.
+  /// \param callback Callback that will be called after object has been added to GCS.
+  /// \return Status
+  virtual Status AsyncAddLocation(const ObjectID &object_id, const ClientID &node_id,
+                                  const StatusCallback &callback) = 0;
+
+  /// Remove location of object from GCS asynchronously.
+  ///
+  /// \param object_id The ID of object which location will be removed from GCS.
+  /// \param node_id The location that will be removed from GCS.
+  /// \param callback Callback that will be called after the delete finished.
+  /// \return Status
+  virtual Status AsyncRemoveLocation(const ObjectID &object_id, const ClientID &node_id,
+                                     const StatusCallback &callback) = 0;
+
+  /// Subscribe to any update of an object's location.
+  ///
+  /// \param object_id The ID of the object to be subscribed to.
+  /// \param subscribe Callback that will be called each time when the object's
+  /// location is updated.
+  /// \param done Callback that will be called when subscription is complete.
+  /// \return Status
+  virtual Status AsyncSubscribeToLocations(
+      const ObjectID &object_id,
+      const SubscribeCallback<ObjectID, ObjectChangeNotification> &subscribe,
+      const StatusCallback &done) = 0;
+
+  /// Cancel subscription to any update of an object's location.
+  ///
+  /// \param object_id The ID of the object to be unsubscribed to.
+  /// \param done Callback that will be called when unsubscription is complete.
+  /// \return Status
+  virtual Status AsyncUnsubscribeToLocations(const ObjectID &object_id,
+                                             const StatusCallback &done) = 0;
+
+ protected:
+  ObjectInfoAccessor() = default;
+};
+
 /// \class NodeInfoAccessor
 /// `NodeInfoAccessor` is a sub-interface of `GcsClient`.
 /// This class includes all the methods that are related to accessing

src/ray/gcs/entry_change_notification.h

@@ -0,0 +1,57 @@
+#ifndef RAY_GCS_ENTRY_CHANGE_NOTIFICATION_H
+#define RAY_GCS_ENTRY_CHANGE_NOTIFICATION_H
+
+#include <ray/protobuf/gcs.pb.h>
+#include <vector>
+
+namespace ray {
+
+namespace gcs {
+
+/// \class EntryChangeNotification
+/// EntryChangeNotification class is a template class which represent
+/// notification of entry change from GCS.
+template <typename Data>
+class EntryChangeNotification {
+ public:
+  EntryChangeNotification(rpc::GcsChangeMode change_mode, std::vector<Data> data)
+      : change_mode_(change_mode), data_(std::move(data)) {}
+
+  EntryChangeNotification(EntryChangeNotification &&other) {
+    change_mode_ = other.change_mode_;
+    data_ = std::move(other.data_);
+  }
+
+  EntryChangeNotification &operator=(EntryChangeNotification &&other) {
+    change_mode_ = other.change_mode_;
+    data_ = std::move(other.data_);
+  }
+
+  /// Whether the entry data is removed from GCS.
+  bool IsRemoved() const { return change_mode_ == rpc::GcsChangeMode::REMOVE; }
+
+  /// Whether the entry data is added to GCS.
+  bool IsAdded() const { return change_mode_ == rpc::GcsChangeMode::APPEND_OR_ADD; }
+
+  /// Get change mode of this notification. For test only.
+  ///
+  /// \return rpc::GcsChangeMode
+  rpc::GcsChangeMode GetGcsChangeMode() const { return change_mode_; }
+
+  /// Get data of this notification.
+  ///
+  /// \return Data
+  const std::vector<Data> &GetData() const { return data_; }
+
+ private:
+  rpc::GcsChangeMode change_mode_;
+  std::vector<Data> data_;
+};
+
+typedef EntryChangeNotification<rpc::ObjectTableData> ObjectChangeNotification;
+
+}  // namespace gcs
+
+}  // namespace ray
+
+#endif  // RAY_GCS_ENTRY_CHANGE_NOTIFICATION_H

src/ray/gcs/gcs_client.h

@@ -74,6 +74,13 @@ class GcsClient : public std::enable_shared_from_this<GcsClient> {
     return *job_accessor_;
   }
 
+  /// Get the sub-interface for accessing object information in GCS.
+  /// This function is thread safe.
+  ObjectInfoAccessor &Objects() {
+    RAY_CHECK(object_accessor_ != nullptr);
+    return *object_accessor_;
+  }
+
   /// Get the sub-interface for accessing node information in GCS.
   /// This function is thread safe.
   NodeInfoAccessor &Nodes() {
@@ -101,6 +108,7 @@ class GcsClient : public std::enable_shared_from_this<GcsClient> {
 
   std::unique_ptr<ActorInfoAccessor> actor_accessor_;
   std::unique_ptr<JobInfoAccessor> job_accessor_;
+  std::unique_ptr<ObjectInfoAccessor> object_accessor_;
   std::unique_ptr<NodeInfoAccessor> node_accessor_;
   std::unique_ptr<TaskInfoAccessor> task_accessor_;
 };

src/ray/gcs/redis_accessor.cc

@@ -126,12 +126,12 @@ Status RedisActorInfoAccessor::AsyncSubscribe(
     const ActorID &actor_id, const SubscribeCallback<ActorID, ActorTableData> &subscribe,
     const StatusCallback &done) {
   RAY_CHECK(subscribe != nullptr);
-  return actor_sub_executor_.AsyncSubscribe(node_id_, actor_id, subscribe, done);
+  return actor_sub_executor_.AsyncSubscribe(subscribe_id_, actor_id, subscribe, done);
 }
 
 Status RedisActorInfoAccessor::AsyncUnsubscribe(const ActorID &actor_id,
                                                 const StatusCallback &done) {
-  return actor_sub_executor_.AsyncUnsubscribe(node_id_, actor_id, done);
+  return actor_sub_executor_.AsyncUnsubscribe(subscribe_id_, actor_id, done);
 }
 
 RedisJobInfoAccessor::RedisJobInfoAccessor(RedisGcsClient *client_impl)
@@ -228,6 +228,70 @@ Status RedisTaskInfoAccessor::AsyncUnsubscribe(const TaskID &task_id,
   return task_sub_executor_.AsyncUnsubscribe(subscribe_id_, task_id, done);
 }
 
+RedisObjectInfoAccessor::RedisObjectInfoAccessor(RedisGcsClient *client_impl)
+    : client_impl_(client_impl), object_sub_executor_(client_impl->object_table()) {}
+
+Status RedisObjectInfoAccessor::AsyncGetLocations(
+    const ObjectID &object_id, const MultiItemCallback<ObjectTableData> &callback) {
+  RAY_CHECK(callback != nullptr);
+  auto on_done = [callback](RedisGcsClient *client, const ObjectID &object_id,
+                            const std::vector<ObjectTableData> &data) {
+    callback(Status::OK(), data);
+  };
+
+  ObjectTable &object_table = client_impl_->object_table();
+  return object_table.Lookup(JobID::Nil(), object_id, on_done);
+}
+
+Status RedisObjectInfoAccessor::AsyncAddLocation(const ObjectID &object_id,
+                                                 const ClientID &node_id,
+                                                 const StatusCallback &callback) {
+  std::function<void(RedisGcsClient * client, const ObjectID &id,
+                     const ObjectTableData &data)>
+      on_done = nullptr;
+  if (callback != nullptr) {
+    on_done = [callback](RedisGcsClient *client, const ObjectID &object_id,
+                         const ObjectTableData &data) { callback(Status::OK()); };
+  }
+
+  std::shared_ptr<ObjectTableData> data_ptr = std::make_shared<ObjectTableData>();
+  data_ptr->set_manager(node_id.Binary());
+
+  ObjectTable &object_table = client_impl_->object_table();
+  return object_table.Add(JobID::Nil(), object_id, data_ptr, on_done);
+}
+
+Status RedisObjectInfoAccessor::AsyncRemoveLocation(const ObjectID &object_id,
+                                                    const ClientID &node_id,
+                                                    const StatusCallback &callback) {
+  std::function<void(RedisGcsClient * client, const ObjectID &id,
+                     const ObjectTableData &data)>
+      on_done = nullptr;
+  if (callback != nullptr) {
+    on_done = [callback](RedisGcsClient *client, const ObjectID &object_id,
+                         const ObjectTableData &data) { callback(Status::OK()); };
+  }
+
+  std::shared_ptr<ObjectTableData> data_ptr = std::make_shared<ObjectTableData>();
+  data_ptr->set_manager(node_id.Binary());
+
+  ObjectTable &object_table = client_impl_->object_table();
+  return object_table.Remove(JobID::Nil(), object_id, data_ptr, on_done);
+}
+
+Status RedisObjectInfoAccessor::AsyncSubscribeToLocations(
+    const ObjectID &object_id,
+    const SubscribeCallback<ObjectID, ObjectChangeNotification> &subscribe,
+    const StatusCallback &done) {
+  RAY_CHECK(subscribe != nullptr);
+  return object_sub_executor_.AsyncSubscribe(subscribe_id_, object_id, subscribe, done);
+}
+
+Status RedisObjectInfoAccessor::AsyncUnsubscribeToLocations(const ObjectID &object_id,
+                                                            const StatusCallback &done) {
+  return object_sub_executor_.AsyncUnsubscribe(subscribe_id_, object_id, done);
+}
+
 RedisNodeInfoAccessor::RedisNodeInfoAccessor(RedisGcsClient *client_impl)
     : client_impl_(client_impl) {}
 

src/ray/gcs/redis_accessor.h

@@ -54,7 +54,7 @@ class RedisActorInfoAccessor : public ActorInfoAccessor {
   // And because the new GCS Client will no longer hold the local ClientID, so we use
   // random ClientID instead.
   // TODO(micafan): Remove this random id, once GCS becomes a service.
-  ClientID node_id_{ClientID::FromRandom()};
+  ClientID subscribe_id_{ClientID::FromRandom()};
 
   typedef SubscriptionExecutor<ActorID, ActorTableData, ActorTable>
       ActorSubscriptionExecutor;
@@ -132,6 +132,48 @@ class RedisTaskInfoAccessor : public TaskInfoAccessor {
   TaskSubscriptionExecutor task_sub_executor_;
 };
 
+/// \class RedisObjectInfoAccessor
+/// RedisObjectInfoAccessor is an implementation of `ObjectInfoAccessor`
+/// that uses Redis as the backend storage.
+class RedisObjectInfoAccessor : public ObjectInfoAccessor {
+ public:
+  explicit RedisObjectInfoAccessor(RedisGcsClient *client_impl);
+
+  virtual ~RedisObjectInfoAccessor() {}
+
+  Status AsyncGetLocations(const ObjectID &object_id,
+                           const MultiItemCallback<ObjectTableData> &callback) override;
+
+  Status AsyncAddLocation(const ObjectID &object_id, const ClientID &node_id,
+                          const StatusCallback &callback) override;
+
+  Status AsyncRemoveLocation(const ObjectID &object_id, const ClientID &node_id,
+                             const StatusCallback &callback) override;
+
+  Status AsyncSubscribeToLocations(
+      const ObjectID &object_id,
+      const SubscribeCallback<ObjectID, ObjectChangeNotification> &subscribe,
+      const StatusCallback &done) override;
+
+  Status AsyncUnsubscribeToLocations(const ObjectID &object_id,
+                                     const StatusCallback &done) override;
+
+ private:
+  RedisGcsClient *client_impl_{nullptr};
+
+  // Use a random ClientID for object subscription. Because:
+  // If we use ClientID::Nil, GCS will still send all objects' updates to this GCS Client.
+  // Even we can filter out irrelevant updates, but there will be extra overhead.
+  // And because the new GCS Client will no longer hold the local ClientID, so we use
+  // random ClientID instead.
+  // TODO(micafan): Remove this random id, once GCS becomes a service.
+  ClientID subscribe_id_{ClientID::FromRandom()};
+
+  typedef SubscriptionExecutor<ObjectID, ObjectChangeNotification, ObjectTable>
+      ObjectSubscriptionExecutor;
+  ObjectSubscriptionExecutor object_sub_executor_;
+};
+
 /// \class RedisNodeInfoAccessor
 /// RedisNodeInfoAccessor is an implementation of `NodeInfoAccessor`
 /// that uses Redis as the backend storage.

src/ray/gcs/redis_gcs_client.cc

@@ -144,6 +144,7 @@ Status RedisGcsClient::Connect(boost::asio::io_service &io_service) {
 
   actor_accessor_.reset(new RedisActorInfoAccessor(this));
   job_accessor_.reset(new RedisJobInfoAccessor(this));
+  object_accessor_.reset(new RedisObjectInfoAccessor(this));
   node_accessor_.reset(new RedisNodeInfoAccessor(this));
   task_accessor_.reset(new RedisTaskInfoAccessor(this));
 

src/ray/gcs/redis_gcs_client.h

@@ -24,10 +24,12 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
   friend class RedisJobInfoAccessor;
   friend class RedisTaskInfoAccessor;
   friend class RedisNodeInfoAccessor;
+  friend class RedisObjectInfoAccessor;
   friend class SubscriptionExecutorTest;
   friend class LogSubscribeTestHelper;
   friend class TaskTableTestHelper;
   friend class ClientTableTestHelper;
+  friend class SetTestHelper;
 
  public:
   /// Constructor of RedisGcsClient.
@@ -58,7 +60,6 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
   void Disconnect();
 
   // TODO: Some API for getting the error on the driver
-  ObjectTable &object_table();
   TaskReconstructionLog &task_reconstruction_log();
   TaskLeaseTable &task_lease_table();
   HeartbeatTable &heartbeat_table();
@@ -98,6 +99,8 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
   ActorTable &actor_table();
   /// This method will be deprecated, use method Jobs() instead.
   JobTable &job_table();
+  /// This method will be deprecated, use method Objects() instead
+  ObjectTable &object_table();
   /// This method will be deprecated, use method Nodes() instead.
   ClientTable &client_table();
   /// This method will be deprecated, use method Tasks() instead.

src/ray/gcs/subscription_executor.cc

@@ -52,8 +52,6 @@ Status SubscriptionExecutor<ID, Data, Table>::AsyncSubscribeAll(
       return;
     }
 
-    RAY_LOG(DEBUG) << "Subscribe received update of id " << id;
-
     SubscribeCallback<ID, Data> sub_one_callback = nullptr;
     SubscribeCallback<ID, Data> sub_all_callback = nullptr;
     {
@@ -190,6 +188,7 @@ template class SubscriptionExecutor<ActorID, ActorTableData, ActorTable>;
 template class SubscriptionExecutor<ActorID, ActorTableData, DirectActorTable>;
 template class SubscriptionExecutor<JobID, JobTableData, JobTable>;
 template class SubscriptionExecutor<TaskID, TaskTableData, raylet::TaskTable>;
+template class SubscriptionExecutor<ObjectID, ObjectChangeNotification, ObjectTable>;
 
 }  // namespace gcs
 

src/ray/gcs/tables.cc

@@ -361,6 +361,21 @@ Status Set<ID, Data>::Remove(const JobID &job_id, const ID &id,
                                        prefix_, pubsub_channel_, std::move(callback));
 }
 
+template <typename ID, typename Data>
+Status Set<ID, Data>::Subscribe(const JobID &job_id, const ClientID &client_id,
+                                const NotificationCallback &subscribe,
+                                const SubscriptionCallback &done) {
+  auto on_subscribe = [subscribe](RedisGcsClient *client, const ID &id,
+                                  const GcsChangeMode change_mode,
+                                  const std::vector<Data> &data) {
+    EntryChangeNotification<Data> change_notification(change_mode, data);
+    std::vector<EntryChangeNotification<Data>> notification_vec;
+    notification_vec.emplace_back(std::move(change_notification));
+    subscribe(client, id, notification_vec);
+  };
+  return Log<ID, Data>::Subscribe(job_id, client_id, on_subscribe, done);
+}
+
 template <typename ID, typename Data>
 std::string Set<ID, Data>::DebugString() const {
   std::stringstream result;

src/ray/gcs/tables.h

@@ -12,6 +12,7 @@
 #include "ray/util/logging.h"
 
 #include "ray/gcs/callback.h"
+#include "ray/gcs/entry_change_notification.h"
 #include "ray/gcs/redis_context.h"
 #include "ray/protobuf/gcs.pb.h"
 
@@ -93,9 +94,11 @@ class Log : public LogInterface<ID, Data>, virtual public PubsubInterface<ID> {
  public:
   using Callback = std::function<void(RedisGcsClient *client, const ID &id,
                                       const std::vector<Data> &data)>;
+
   using NotificationCallback =
       std::function<void(RedisGcsClient *client, const ID &id,
                          const GcsChangeMode change_mode, const std::vector<Data> &data)>;
+
   /// The callback to call when a write to a key succeeds.
   using WriteCallback = typename LogInterface<ID, Data>::WriteCallback;
   /// The callback to call when a SUBSCRIBE call completes and we are ready to
@@ -164,6 +167,7 @@ class Log : public LogInterface<ID, Data>, virtual public PubsubInterface<ID> {
   /// called with an empty vector, then there was no data at the key.
   /// \return Status
   Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup);
+
   /// Subscribe to any Append operations to this table. The caller may choose
   /// requests notifications for. This may only be called once per Log
   ///
@@ -210,6 +214,28 @@ class Log : public LogInterface<ID, Data>, virtual public PubsubInterface<ID> {
   Status CancelNotifications(const JobID &job_id, const ID &id, const ClientID &client_id,
                              const StatusCallback &done);
 
+  /// Subscribe to any modifications to the key. The caller may choose
+  /// to subscribe to all modifications, or to subscribe only to keys that it
+  /// requests notifications for. This may only be called once per Log
+  /// instance. This function is different from public version due to
+  /// an additional parameter change_mode in NotificationCallback. Therefore this
+  /// function supports notifications of remove operations.
+  ///
+  /// \param job_id The ID of the job.
+  /// \param client_id The type of update to listen to. If this is nil, then a
+  /// message for each Add to the table will be received. Else, only
+  /// messages for the given client will be received. In the latter
+  /// case, the client may request notifications on specific keys in the
+  /// table via `RequestNotifications`.
+  /// \param subscribe Callback that is called on each received message. If the
+  /// callback is called with an empty vector, then there was no data at the key.
+  /// \param done Callback that is called when subscription is complete and we
+  /// are ready to receive messages.
+  /// \return Status
+  Status Subscribe(const JobID &job_id, const ClientID &client_id,
+                   const NotificationCallback &subscribe,
+                   const SubscriptionCallback &done);
+
   /// Delete an entire key from redis.
   ///
   /// \param job_id The ID of the job.
@@ -235,28 +261,6 @@ class Log : public LogInterface<ID, Data>, virtual public PubsubInterface<ID> {
     return shard_contexts_[index(id) % shard_contexts_.size()];
   }
 
-  /// Subscribe to any modifications to the key. The caller may choose
-  /// to subscribe to all modifications, or to subscribe only to keys that it
-  /// requests notifications for. This may only be called once per Log
-  /// instance. This function is different from public version due to
-  /// an additional parameter change_mode in NotificationCallback. Therefore this
-  /// function supports notifications of remove operations.
-  ///
-  /// \param job_id The ID of the job.
-  /// \param client_id The type of update to listen to. If this is nil, then a
-  /// message for each Add to the table will be received. Else, only
-  /// messages for the given client will be received. In the latter
-  /// case, the client may request notifications on specific keys in the
-  /// table via `RequestNotifications`.
-  /// \param subscribe Callback that is called on each received message. If the
-  /// callback is called with an empty vector, then there was no data at the key.
-  /// \param done Callback that is called when subscription is complete and we
-  /// are ready to receive messages.
-  /// \return Status
-  Status Subscribe(const JobID &job_id, const ClientID &client_id,
-                   const NotificationCallback &subscribe,
-                   const SubscriptionCallback &done);
-
   /// The connection to the GCS.
   std::vector<std::shared_ptr<RedisContext>> shard_contexts_;
   /// The GCS client.
@@ -435,7 +439,6 @@ class Set : private Log<ID, Data>,
  public:
   using Callback = typename Log<ID, Data>::Callback;
   using WriteCallback = typename Log<ID, Data>::WriteCallback;
-  using NotificationCallback = typename Log<ID, Data>::NotificationCallback;
   using SubscriptionCallback = typename Log<ID, Data>::SubscriptionCallback;
 
   Set(const std::vector<std::shared_ptr<RedisContext>> &contexts, RedisGcsClient *client)
@@ -468,11 +471,24 @@ class Set : private Log<ID, Data>,
   Status Remove(const JobID &job_id, const ID &id, const std::shared_ptr<Data> &data,
                 const WriteCallback &done);
 
+  using NotificationCallback =
+      std::function<void(RedisGcsClient *client, const ID &id,
+                         const std::vector<ObjectChangeNotification> &data)>;
+  /// Subscribe to any add or remove operations to this table.
+  ///
+  /// \param job_id The ID of the job.
+  /// \param client_id The type of update to listen to. If this is nil, then a
+  /// message for each add or remove to the table will be received. Else, only
+  /// messages for the given client will be received. In the latter
+  /// case, the client may request notifications on specific keys in the
+  /// table via `RequestNotifications`.
+  /// \param subscribe Callback that is called on each received message.
+  /// \param done Callback that is called when subscription is complete and we
+  /// are ready to receive messages.
+  /// \return Status
   Status Subscribe(const JobID &job_id, const ClientID &client_id,
                    const NotificationCallback &subscribe,
-                   const SubscriptionCallback &done) {
-    return Log<ID, Data>::Subscribe(job_id, client_id, subscribe, done);
-  }
+                   const SubscriptionCallback &done);
 
   /// Returns debug string for class.
   ///

src/ray/gcs/test/redis_gcs_client_test.cc

@@ -470,71 +470,341 @@ TEST_F(TestGcsWithAsio, TestLogAppendAt) {
   TestLogAppendAt(job_id_, client_);
 }
 
-void TestSet(const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> client) {
-  // Add some entries to the set at an object ID.
-  ObjectID object_id = ObjectID::FromRandom();
-  std::vector<std::string> managers = {"abc", "def", "ghi"};
-  for (auto &manager : managers) {
-    auto data = std::make_shared<ObjectTableData>();
-    data->set_manager(manager);
-    // Check that we added the correct object entries.
-    auto add_callback = [object_id, data](gcs::RedisGcsClient *client, const ObjectID &id,
-                                          const ObjectTableData &d) {
+class SetTestHelper {
+ public:
+  static void TestSet(const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> client) {
+    // Add some entries to the set at an object ID.
+    ObjectID object_id = ObjectID::FromRandom();
+    std::vector<std::string> managers = {"abc", "def", "ghi"};
+    for (auto &manager : managers) {
+      auto data = std::make_shared<ObjectTableData>();
+      data->set_manager(manager);
+      // Check that we added the correct object entries.
+      auto add_callback = [object_id, data](gcs::RedisGcsClient *client,
+                                            const ObjectID &id,
+                                            const ObjectTableData &d) {
+        ASSERT_EQ(id, object_id);
+        ASSERT_EQ(data->manager(), d.manager());
+        test->IncrementNumCallbacks();
+      };
+      RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, add_callback));
+    }
+
+    // Check that lookup returns the added object entries.
+    auto lookup_callback = [object_id, managers](
+                               gcs::RedisGcsClient *client, const ObjectID &id,
+                               const std::vector<ObjectTableData> &data) {
       ASSERT_EQ(id, object_id);
-      ASSERT_EQ(data->manager(), d.manager());
+      ASSERT_EQ(data.size(), managers.size());
       test->IncrementNumCallbacks();
     };
-    RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, add_callback));
-  }
 
-  // Check that lookup returns the added object entries.
-  auto lookup_callback = [object_id, managers](gcs::RedisGcsClient *client,
-                                               const ObjectID &id,
-                                               const std::vector<ObjectTableData> &data) {
-    ASSERT_EQ(id, object_id);
-    ASSERT_EQ(data.size(), managers.size());
-    test->IncrementNumCallbacks();
-  };
+    // Do a lookup at the object ID.
+    RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback));
 
-  // Do a lookup at the object ID.
-  RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback));
+    for (auto &manager : managers) {
+      auto data = std::make_shared<ObjectTableData>();
+      data->set_manager(manager);
+      // Check that we added the correct object entries.
+      auto remove_entry_callback = [object_id, data](gcs::RedisGcsClient *client,
+                                                     const ObjectID &id,
+                                                     const ObjectTableData &d) {
+        ASSERT_EQ(id, object_id);
+        ASSERT_EQ(data->manager(), d.manager());
+        test->IncrementNumCallbacks();
+      };
+      RAY_CHECK_OK(
+          client->object_table().Remove(job_id, object_id, data, remove_entry_callback));
+    }
 
-  for (auto &manager : managers) {
-    auto data = std::make_shared<ObjectTableData>();
-    data->set_manager(manager);
-    // Check that we added the correct object entries.
-    auto remove_entry_callback = [object_id, data](gcs::RedisGcsClient *client,
-                                                   const ObjectID &id,
-                                                   const ObjectTableData &d) {
+    // Check that the entries are removed.
+    auto lookup_callback2 = [object_id, managers](
+                                gcs::RedisGcsClient *client, const ObjectID &id,
+                                const std::vector<ObjectTableData> &data) {
       ASSERT_EQ(id, object_id);
-      ASSERT_EQ(data->manager(), d.manager());
+      ASSERT_EQ(data.size(), 0);
       test->IncrementNumCallbacks();
+      test->Stop();
     };
-    RAY_CHECK_OK(
-        client->object_table().Remove(job_id, object_id, data, remove_entry_callback));
+
+    // Do a lookup at the object ID.
+    RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback2));
+    // Run the event loop. The loop will only stop if the Lookup callback is
+    // called (or an assertion failure).
+    test->Start();
+    ASSERT_EQ(test->NumCallbacks(), managers.size() * 2 + 2);
   }
 
-  // Check that the entries are removed.
-  auto lookup_callback2 = [object_id, managers](
-                              gcs::RedisGcsClient *client, const ObjectID &id,
-                              const std::vector<ObjectTableData> &data) {
-    ASSERT_EQ(id, object_id);
-    ASSERT_EQ(data.size(), 0);
-    test->IncrementNumCallbacks();
-    test->Stop();
-  };
+  static void TestDeleteKeysFromSet(
+      const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> client,
+      std::vector<std::shared_ptr<ObjectTableData>> &data_vector) {
+    std::vector<ObjectID> ids;
+    ObjectID object_id;
+    for (auto &data : data_vector) {
+      object_id = ObjectID::FromRandom();
+      ids.push_back(object_id);
+      // Check that we added the correct object entries.
+      auto add_callback = [object_id, data](gcs::RedisGcsClient *client,
+                                            const ObjectID &id,
+                                            const ObjectTableData &d) {
+        ASSERT_EQ(id, object_id);
+        ASSERT_EQ(data->manager(), d.manager());
+        test->IncrementNumCallbacks();
+      };
+      RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, add_callback));
+    }
+    for (const auto &object_id : ids) {
+      // Check that lookup returns the added object entries.
+      auto lookup_callback = [object_id, data_vector](
+                                 gcs::RedisGcsClient *client, const ObjectID &id,
+                                 const std::vector<ObjectTableData> &data) {
+        ASSERT_EQ(id, object_id);
+        ASSERT_EQ(data.size(), 1);
+        test->IncrementNumCallbacks();
+      };
+      RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback));
+    }
+    if (ids.size() == 1) {
+      client->object_table().Delete(job_id, ids[0]);
+    } else {
+      client->object_table().Delete(job_id, ids);
+    }
+    for (const auto &object_id : ids) {
+      auto lookup_callback = [object_id](gcs::RedisGcsClient *client, const ObjectID &id,
+                                         const std::vector<ObjectTableData> &data) {
+        ASSERT_EQ(id, object_id);
+        ASSERT_TRUE(data.size() == 0);
+        test->IncrementNumCallbacks();
+      };
+      RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback));
+    }
+  }
 
-  // Do a lookup at the object ID.
-  RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback2));
-  // Run the event loop. The loop will only stop if the Lookup callback is
-  // called (or an assertion failure).
-  test->Start();
-  ASSERT_EQ(test->NumCallbacks(), managers.size() * 2 + 2);
-}
+  static void TestSetSubscribeAll(const JobID &job_id,
+                                  std::shared_ptr<gcs::RedisGcsClient> client) {
+    std::vector<ObjectID> object_ids;
+    for (int i = 0; i < 3; i++) {
+      object_ids.emplace_back(ObjectID::FromRandom());
+    }
+    std::vector<std::string> managers = {"abc", "def", "ghi"};
+
+    // Callback for a notification.
+    auto notification_callback =
+        [object_ids, managers](
+            gcs::RedisGcsClient *client, const ObjectID &id,
+            const std::vector<ObjectChangeNotification> &notifications) {
+          if (test->NumCallbacks() < 3 * 3) {
+            ASSERT_EQ(notifications[0].GetGcsChangeMode(), GcsChangeMode::APPEND_OR_ADD);
+          } else {
+            ASSERT_EQ(notifications[0].GetGcsChangeMode(), GcsChangeMode::REMOVE);
+          }
+          ASSERT_EQ(id, object_ids[test->NumCallbacks() / 3 % 3]);
+          // Check that we get notifications in the same order as the writes.
+          for (const auto &entry : notifications[0].GetData()) {
+            ASSERT_EQ(entry.manager(), managers[test->NumCallbacks() % 3]);
+            test->IncrementNumCallbacks();
+          }
+          if (test->NumCallbacks() == object_ids.size() * 3 * 2) {
+            test->Stop();
+          }
+        };
+
+    // Callback for subscription success. We are guaranteed to receive
+    // notifications after this is called.
+    auto subscribe_callback = [job_id, object_ids,
+                               managers](gcs::RedisGcsClient *client) {
+      // We have subscribed. Do the writes to the table.
+      for (size_t i = 0; i < object_ids.size(); i++) {
+        for (size_t j = 0; j < managers.size(); j++) {
+          auto data = std::make_shared<ObjectTableData>();
+          data->set_manager(managers[j]);
+          for (int k = 0; k < 3; k++) {
+            // Add the same entry several times.
+            // Expect no notification if the entry already exists.
+            RAY_CHECK_OK(
+                client->object_table().Add(job_id, object_ids[i], data, nullptr));
+          }
+        }
+      }
+      for (size_t i = 0; i < object_ids.size(); i++) {
+        for (size_t j = 0; j < managers.size(); j++) {
+          auto data = std::make_shared<ObjectTableData>();
+          data->set_manager(managers[j]);
+          for (int k = 0; k < 3; k++) {
+            // Remove the same entry several times.
+            // Expect no notification if the entry doesn't exist.
+            RAY_CHECK_OK(
+                client->object_table().Remove(job_id, object_ids[i], data, nullptr));
+          }
+        }
+      }
+    };
+
+    // Subscribe to all driver table notifications. Once we have successfully
+    // subscribed, we will append to the key several times and check that we get
+    // notified for each.
+    RAY_CHECK_OK(client->object_table().Subscribe(
+        job_id, ClientID::Nil(), notification_callback, subscribe_callback));
+
+    // Run the event loop. The loop will only stop if the registered subscription
+    // callback is called (or an assertion failure).
+    test->Start();
+    // Check that we received one notification callback for each write.
+    ASSERT_EQ(test->NumCallbacks(), object_ids.size() * 3 * 2);
+  }
+
+  static void TestSetSubscribeId(const JobID &job_id,
+                                 std::shared_ptr<gcs::RedisGcsClient> client) {
+    // Add a set entry.
+    ObjectID object_id1 = ObjectID::FromRandom();
+    std::vector<std::string> managers1 = {"abc", "def", "ghi"};
+    auto data1 = std::make_shared<ObjectTableData>();
+    data1->set_manager(managers1[0]);
+    RAY_CHECK_OK(client->object_table().Add(job_id, object_id1, data1, nullptr));
+
+    // Add a set entry at a second key.
+    ObjectID object_id2 = ObjectID::FromRandom();
+    std::vector<std::string> managers2 = {"jkl", "mno", "pqr"};
+    auto data2 = std::make_shared<ObjectTableData>();
+    data2->set_manager(managers2[0]);
+    RAY_CHECK_OK(client->object_table().Add(job_id, object_id2, data2, nullptr));
+
+    // The callback for a notification from the table. This should only be
+    // received for keys that we requested notifications for.
+    auto notification_callback =
+        [object_id2, managers2](
+            gcs::RedisGcsClient *client, const ObjectID &id,
+            const std::vector<ObjectChangeNotification> &notifications) {
+          ASSERT_EQ(notifications[0].GetGcsChangeMode(), GcsChangeMode::APPEND_OR_ADD);
+          // Check that we only get notifications for the requested key.
+          ASSERT_EQ(id, object_id2);
+          // Check that we get notifications in the same order as the writes.
+          for (const auto &entry : notifications[0].GetData()) {
+            ASSERT_EQ(entry.manager(), managers2[test->NumCallbacks()]);
+            test->IncrementNumCallbacks();
+          }
+          if (test->NumCallbacks() == managers2.size()) {
+            test->Stop();
+          }
+        };
+
+    // The callback for subscription success. Once we've subscribed, request
+    // notifications for only one of the keys, then write to both keys.
+    auto subscribe_callback = [job_id, object_id1, object_id2, managers1,
+                               managers2](gcs::RedisGcsClient *client) {
+      // Request notifications for one of the keys.
+      RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id2,
+                                                               local_client_id, nullptr));
+      // Write both keys. We should only receive notifications for the key that
+      // we requested them for.
+      auto remaining = std::vector<std::string>(++managers1.begin(), managers1.end());
+      for (const auto &manager : remaining) {
+        auto data = std::make_shared<ObjectTableData>();
+        data->set_manager(manager);
+        RAY_CHECK_OK(client->object_table().Add(job_id, object_id1, data, nullptr));
+      }
+      remaining = std::vector<std::string>(++managers2.begin(), managers2.end());
+      for (const auto &manager : remaining) {
+        auto data = std::make_shared<ObjectTableData>();
+        data->set_manager(manager);
+        RAY_CHECK_OK(client->object_table().Add(job_id, object_id2, data, nullptr));
+      }
+    };
+
+    // Subscribe to notifications for this client. This allows us to request and
+    // receive notifications for specific keys.
+    RAY_CHECK_OK(client->object_table().Subscribe(
+        job_id, local_client_id, notification_callback, subscribe_callback));
+    // Run the event loop. The loop will only stop if the registered subscription
+    // callback is called for the requested key.
+    test->Start();
+    // Check that we received one notification callback for each write to the
+    // requested key.
+    ASSERT_EQ(test->NumCallbacks(), managers2.size());
+  }
+
+  static void TestSetSubscribeCancel(const JobID &job_id,
+                                     std::shared_ptr<gcs::RedisGcsClient> client) {
+    // Add a set entry.
+    ObjectID object_id = ObjectID::FromRandom();
+    std::vector<std::string> managers = {"jkl", "mno", "pqr"};
+    auto data = std::make_shared<ObjectTableData>();
+    data->set_manager(managers[0]);
+    RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, nullptr));
+
+    // The callback for a notification from the object table. This should only be
+    // received for the object that we requested notifications for.
+    auto notification_callback =
+        [object_id, managers](
+            gcs::RedisGcsClient *client, const ObjectID &id,
+            const std::vector<ObjectChangeNotification> &notifications) {
+          ASSERT_EQ(notifications[0].GetGcsChangeMode(), GcsChangeMode::APPEND_OR_ADD);
+          ASSERT_EQ(id, object_id);
+          // Check that we get a duplicate notification for the first write. We get a
+          // duplicate notification because notifications
+          // are canceled after the first write, then requested again.
+          const std::vector<ObjectTableData> &data = notifications[0].GetData();
+          if (data.size() == 1) {
+            // first notification
+            ASSERT_EQ(data[0].manager(), managers[0]);
+            test->IncrementNumCallbacks();
+          } else {
+            // second notification
+            ASSERT_EQ(data.size(), managers.size());
+            std::unordered_set<std::string> managers_set(managers.begin(),
+                                                         managers.end());
+            std::unordered_set<std::string> data_managers_set;
+            for (const auto &entry : data) {
+              data_managers_set.insert(entry.manager());
+              test->IncrementNumCallbacks();
+            }
+            ASSERT_EQ(managers_set, data_managers_set);
+          }
+          if (test->NumCallbacks() == managers.size() + 1) {
+            test->Stop();
+          }
+        };
+
+    // The callback for a notification from the table. This should only be
+    // received for keys that we requested notifications for.
+    auto subscribe_callback = [job_id, object_id, managers](gcs::RedisGcsClient *client) {
+      // Request notifications, then cancel immediately. We should receive a
+      // notification for the current value at the key.
+      RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id,
+                                                               local_client_id, nullptr));
+      RAY_CHECK_OK(client->object_table().CancelNotifications(job_id, object_id,
+                                                              local_client_id, nullptr));
+      // Add to the key. Since we canceled notifications, we should not
+      // receive a notification for these writes.
+      auto remaining = std::vector<std::string>(++managers.begin(), managers.end());
+      for (const auto &manager : remaining) {
+        auto data = std::make_shared<ObjectTableData>();
+        data->set_manager(manager);
+        RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, nullptr));
+      }
+      // Request notifications again. We should receive a notification for the
+      // current values at the key.
+      RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id,
+                                                               local_client_id, nullptr));
+    };
+
+    // Subscribe to notifications for this client. This allows us to request and
+    // receive notifications for specific keys.
+    RAY_CHECK_OK(client->object_table().Subscribe(
+        job_id, local_client_id, notification_callback, subscribe_callback));
+    // Run the event loop. The loop will only stop if the registered subscription
+    // callback is called for the requested key.
+    test->Start();
+    // Check that we received a notification callback for the first append to the
+    // key, then a notification for all of the appends, because we cancel
+    // notifications in between.
+    ASSERT_EQ(test->NumCallbacks(), managers.size() + 1);
+  }
+};
 
 TEST_F(TestGcsWithAsio, TestSet) {
   test = this;
-  TestSet(job_id_, client_);
+  SetTestHelper::TestSet(job_id_, client_);
 }
 
 void TestDeleteKeysFromLog(
@@ -584,50 +854,6 @@ void TestDeleteKeysFromLog(
   }
 }
 
-void TestDeleteKeysFromSet(const JobID &job_id,
-                           std::shared_ptr<gcs::RedisGcsClient> client,
-                           std::vector<std::shared_ptr<ObjectTableData>> &data_vector) {
-  std::vector<ObjectID> ids;
-  ObjectID object_id;
-  for (auto &data : data_vector) {
-    object_id = ObjectID::FromRandom();
-    ids.push_back(object_id);
-    // Check that we added the correct object entries.
-    auto add_callback = [object_id, data](gcs::RedisGcsClient *client, const ObjectID &id,
-                                          const ObjectTableData &d) {
-      ASSERT_EQ(id, object_id);
-      ASSERT_EQ(data->manager(), d.manager());
-      test->IncrementNumCallbacks();
-    };
-    RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, add_callback));
-  }
-  for (const auto &object_id : ids) {
-    // Check that lookup returns the added object entries.
-    auto lookup_callback = [object_id, data_vector](
-                               gcs::RedisGcsClient *client, const ObjectID &id,
-                               const std::vector<ObjectTableData> &data) {
-      ASSERT_EQ(id, object_id);
-      ASSERT_EQ(data.size(), 1);
-      test->IncrementNumCallbacks();
-    };
-    RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback));
-  }
-  if (ids.size() == 1) {
-    client->object_table().Delete(job_id, ids[0]);
-  } else {
-    client->object_table().Delete(job_id, ids);
-  }
-  for (const auto &object_id : ids) {
-    auto lookup_callback = [object_id](gcs::RedisGcsClient *client, const ObjectID &id,
-                                       const std::vector<ObjectTableData> &data) {
-      ASSERT_EQ(id, object_id);
-      ASSERT_TRUE(data.size() == 0);
-      test->IncrementNumCallbacks();
-    };
-    RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback));
-  }
-}
-
 // Test delete function for keys of Log or Table.
 void TestDeleteKeys(const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> client) {
   // Test delete function for keys of Log.
@@ -695,20 +921,20 @@ void TestDeleteKeys(const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> cl
   // Test one element case.
   AppendObjectData(1);
   ASSERT_EQ(object_vector.size(), 1);
-  TestDeleteKeysFromSet(job_id, client, object_vector);
+  SetTestHelper::TestDeleteKeysFromSet(job_id, client, object_vector);
   // Test the case for more than one elements and less than
   // maximum_gcs_deletion_batch_size.
   AppendObjectData(RayConfig::instance().maximum_gcs_deletion_batch_size() / 2);
   ASSERT_GT(object_vector.size(), 1);
   ASSERT_LT(object_vector.size(),
             RayConfig::instance().maximum_gcs_deletion_batch_size());
-  TestDeleteKeysFromSet(job_id, client, object_vector);
+  SetTestHelper::TestDeleteKeysFromSet(job_id, client, object_vector);
   // Test the case for more than maximum_gcs_deletion_batch_size.
   // The Delete function will split the data into two commands.
   AppendObjectData(RayConfig::instance().maximum_gcs_deletion_batch_size() / 2);
   ASSERT_GT(object_vector.size(),
             RayConfig::instance().maximum_gcs_deletion_batch_size());
-  TestDeleteKeysFromSet(job_id, client, object_vector);
+  SetTestHelper::TestDeleteKeysFromSet(job_id, client, object_vector);
 }
 
 TEST_F(TestGcsWithAsio, TestDeleteKey) {
@@ -903,80 +1129,9 @@ TEST_F(TestGcsWithAsio, TestLogSubscribeAll) {
   LogSubscribeTestHelper::TestLogSubscribeAll(job_id_, client_);
 }
 
-void TestSetSubscribeAll(const JobID &job_id,
-                         std::shared_ptr<gcs::RedisGcsClient> client) {
-  std::vector<ObjectID> object_ids;
-  for (int i = 0; i < 3; i++) {
-    object_ids.emplace_back(ObjectID::FromRandom());
-  }
-  std::vector<std::string> managers = {"abc", "def", "ghi"};
-
-  // Callback for a notification.
-  auto notification_callback = [object_ids, managers](
-                                   gcs::RedisGcsClient *client, const ObjectID &id,
-                                   const GcsChangeMode change_mode,
-                                   const std::vector<ObjectTableData> data) {
-    if (test->NumCallbacks() < 3 * 3) {
-      ASSERT_EQ(change_mode, GcsChangeMode::APPEND_OR_ADD);
-    } else {
-      ASSERT_EQ(change_mode, GcsChangeMode::REMOVE);
-    }
-    ASSERT_EQ(id, object_ids[test->NumCallbacks() / 3 % 3]);
-    // Check that we get notifications in the same order as the writes.
-    for (const auto &entry : data) {
-      ASSERT_EQ(entry.manager(), managers[test->NumCallbacks() % 3]);
-      test->IncrementNumCallbacks();
-    }
-    if (test->NumCallbacks() == object_ids.size() * 3 * 2) {
-      test->Stop();
-    }
-  };
-
-  // Callback for subscription success. We are guaranteed to receive
-  // notifications after this is called.
-  auto subscribe_callback = [job_id, object_ids, managers](gcs::RedisGcsClient *client) {
-    // We have subscribed. Do the writes to the table.
-    for (size_t i = 0; i < object_ids.size(); i++) {
-      for (size_t j = 0; j < managers.size(); j++) {
-        auto data = std::make_shared<ObjectTableData>();
-        data->set_manager(managers[j]);
-        for (int k = 0; k < 3; k++) {
-          // Add the same entry several times.
-          // Expect no notification if the entry already exists.
-          RAY_CHECK_OK(client->object_table().Add(job_id, object_ids[i], data, nullptr));
-        }
-      }
-    }
-    for (size_t i = 0; i < object_ids.size(); i++) {
-      for (size_t j = 0; j < managers.size(); j++) {
-        auto data = std::make_shared<ObjectTableData>();
-        data->set_manager(managers[j]);
-        for (int k = 0; k < 3; k++) {
-          // Remove the same entry several times.
-          // Expect no notification if the entry doesn't exist.
-          RAY_CHECK_OK(
-              client->object_table().Remove(job_id, object_ids[i], data, nullptr));
-        }
-      }
-    }
-  };
-
-  // Subscribe to all driver table notifications. Once we have successfully
-  // subscribed, we will append to the key several times and check that we get
-  // notified for each.
-  RAY_CHECK_OK(client->object_table().Subscribe(
-      job_id, ClientID::Nil(), notification_callback, subscribe_callback));
-
-  // Run the event loop. The loop will only stop if the registered subscription
-  // callback is called (or an assertion failure).
-  test->Start();
-  // Check that we received one notification callback for each write.
-  ASSERT_EQ(test->NumCallbacks(), object_ids.size() * 3 * 2);
-}
-
 TEST_F(TestGcsWithAsio, TestSetSubscribeAll) {
   test = this;
-  TestSetSubscribeAll(job_id_, client_);
+  SetTestHelper::TestSetSubscribeAll(job_id_, client_);
 }
 
 TEST_TASK_TABLE_MACRO(TestGcsWithAsio, TestTableSubscribeId);
@@ -986,79 +1141,9 @@ TEST_F(TestGcsWithAsio, TestLogSubscribeId) {
   LogSubscribeTestHelper::TestLogSubscribeId(job_id_, client_);
 }
 
-void TestSetSubscribeId(const JobID &job_id,
-                        std::shared_ptr<gcs::RedisGcsClient> client) {
-  // Add a set entry.
-  ObjectID object_id1 = ObjectID::FromRandom();
-  std::vector<std::string> managers1 = {"abc", "def", "ghi"};
-  auto data1 = std::make_shared<ObjectTableData>();
-  data1->set_manager(managers1[0]);
-  RAY_CHECK_OK(client->object_table().Add(job_id, object_id1, data1, nullptr));
-
-  // Add a set entry at a second key.
-  ObjectID object_id2 = ObjectID::FromRandom();
-  std::vector<std::string> managers2 = {"jkl", "mno", "pqr"};
-  auto data2 = std::make_shared<ObjectTableData>();
-  data2->set_manager(managers2[0]);
-  RAY_CHECK_OK(client->object_table().Add(job_id, object_id2, data2, nullptr));
-
-  // The callback for a notification from the table. This should only be
-  // received for keys that we requested notifications for.
-  auto notification_callback = [object_id2, managers2](
-                                   gcs::RedisGcsClient *client, const ObjectID &id,
-                                   const GcsChangeMode change_mode,
-                                   const std::vector<ObjectTableData> &data) {
-    ASSERT_EQ(change_mode, GcsChangeMode::APPEND_OR_ADD);
-    // Check that we only get notifications for the requested key.
-    ASSERT_EQ(id, object_id2);
-    // Check that we get notifications in the same order as the writes.
-    for (const auto &entry : data) {
-      ASSERT_EQ(entry.manager(), managers2[test->NumCallbacks()]);
-      test->IncrementNumCallbacks();
-    }
-    if (test->NumCallbacks() == managers2.size()) {
-      test->Stop();
-    }
-  };
-
-  // The callback for subscription success. Once we've subscribed, request
-  // notifications for only one of the keys, then write to both keys.
-  auto subscribe_callback = [job_id, object_id1, object_id2, managers1,
-                             managers2](gcs::RedisGcsClient *client) {
-    // Request notifications for one of the keys.
-    RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id2,
-                                                             local_client_id, nullptr));
-    // Write both keys. We should only receive notifications for the key that
-    // we requested them for.
-    auto remaining = std::vector<std::string>(++managers1.begin(), managers1.end());
-    for (const auto &manager : remaining) {
-      auto data = std::make_shared<ObjectTableData>();
-      data->set_manager(manager);
-      RAY_CHECK_OK(client->object_table().Add(job_id, object_id1, data, nullptr));
-    }
-    remaining = std::vector<std::string>(++managers2.begin(), managers2.end());
-    for (const auto &manager : remaining) {
-      auto data = std::make_shared<ObjectTableData>();
-      data->set_manager(manager);
-      RAY_CHECK_OK(client->object_table().Add(job_id, object_id2, data, nullptr));
-    }
-  };
-
-  // Subscribe to notifications for this client. This allows us to request and
-  // receive notifications for specific keys.
-  RAY_CHECK_OK(client->object_table().Subscribe(
-      job_id, local_client_id, notification_callback, subscribe_callback));
-  // Run the event loop. The loop will only stop if the registered subscription
-  // callback is called for the requested key.
-  test->Start();
-  // Check that we received one notification callback for each write to the
-  // requested key.
-  ASSERT_EQ(test->NumCallbacks(), managers2.size());
-}
-
 TEST_F(TestGcsWithAsio, TestSetSubscribeId) {
   test = this;
-  TestSetSubscribeId(job_id_, client_);
+  SetTestHelper::TestSetSubscribeId(job_id_, client_);
 }
 
 TEST_TASK_TABLE_MACRO(TestGcsWithAsio, TestTableSubscribeCancel);
@@ -1068,85 +1153,9 @@ TEST_F(TestGcsWithAsio, TestLogSubscribeCancel) {
   LogSubscribeTestHelper::TestLogSubscribeCancel(job_id_, client_);
 }
 
-void TestSetSubscribeCancel(const JobID &job_id,
-                            std::shared_ptr<gcs::RedisGcsClient> client) {
-  // Add a set entry.
-  ObjectID object_id = ObjectID::FromRandom();
-  std::vector<std::string> managers = {"jkl", "mno", "pqr"};
-  auto data = std::make_shared<ObjectTableData>();
-  data->set_manager(managers[0]);
-  RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, nullptr));
-
-  // The callback for a notification from the object table. This should only be
-  // received for the object that we requested notifications for.
-  auto notification_callback = [object_id, managers](
-                                   gcs::RedisGcsClient *client, const ObjectID &id,
-                                   const GcsChangeMode change_mode,
-                                   const std::vector<ObjectTableData> &data) {
-    ASSERT_EQ(change_mode, GcsChangeMode::APPEND_OR_ADD);
-    ASSERT_EQ(id, object_id);
-    // Check that we get a duplicate notification for the first write. We get a
-    // duplicate notification because notifications
-    // are canceled after the first write, then requested again.
-    if (data.size() == 1) {
-      // first notification
-      ASSERT_EQ(data[0].manager(), managers[0]);
-      test->IncrementNumCallbacks();
-    } else {
-      // second notification
-      ASSERT_EQ(data.size(), managers.size());
-      std::unordered_set<std::string> managers_set(managers.begin(), managers.end());
-      std::unordered_set<std::string> data_managers_set;
-      for (const auto &entry : data) {
-        data_managers_set.insert(entry.manager());
-        test->IncrementNumCallbacks();
-      }
-      ASSERT_EQ(managers_set, data_managers_set);
-    }
-    if (test->NumCallbacks() == managers.size() + 1) {
-      test->Stop();
-    }
-  };
-
-  // The callback for a notification from the table. This should only be
-  // received for keys that we requested notifications for.
-  auto subscribe_callback = [job_id, object_id, managers](gcs::RedisGcsClient *client) {
-    // Request notifications, then cancel immediately. We should receive a
-    // notification for the current value at the key.
-    RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id,
-                                                             local_client_id, nullptr));
-    RAY_CHECK_OK(client->object_table().CancelNotifications(job_id, object_id,
-                                                            local_client_id, nullptr));
-    // Add to the key. Since we canceled notifications, we should not
-    // receive a notification for these writes.
-    auto remaining = std::vector<std::string>(++managers.begin(), managers.end());
-    for (const auto &manager : remaining) {
-      auto data = std::make_shared<ObjectTableData>();
-      data->set_manager(manager);
-      RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, nullptr));
-    }
-    // Request notifications again. We should receive a notification for the
-    // current values at the key.
-    RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id,
-                                                             local_client_id, nullptr));
-  };
-
-  // Subscribe to notifications for this client. This allows us to request and
-  // receive notifications for specific keys.
-  RAY_CHECK_OK(client->object_table().Subscribe(
-      job_id, local_client_id, notification_callback, subscribe_callback));
-  // Run the event loop. The loop will only stop if the registered subscription
-  // callback is called for the requested key.
-  test->Start();
-  // Check that we received a notification callback for the first append to the
-  // key, then a notification for all of the appends, because we cancel
-  // notifications in between.
-  ASSERT_EQ(test->NumCallbacks(), managers.size() + 1);
-}
-
 TEST_F(TestGcsWithAsio, TestSetSubscribeCancel) {
   test = this;
-  TestSetSubscribeCancel(job_id_, client_);
+  SetTestHelper::TestSetSubscribeCancel(job_id_, client_);
 }
 
 /// A helper class for ClientTable testing.

src/ray/gcs/test/redis_object_info_accessor_test.cc

@@ -0,0 +1,138 @@
+#include <unordered_map>
+#include <vector>
+#include "gtest/gtest.h"
+#include "ray/gcs/redis_accessor.h"
+#include "ray/gcs/redis_gcs_client.h"
+#include "ray/gcs/test/accessor_test_base.h"
+#include "ray/util/test_util.h"
+
+namespace ray {
+
+namespace gcs {
+
+class RedisObjectInfoAccessorTest : public AccessorTestBase<ObjectID, ObjectTableData> {
+ protected:
+  void GenTestData() {
+    for (size_t i = 0; i < object_count_; ++i) {
+      ObjectVector object_vec;
+      for (size_t j = 0; j < copy_count_; ++j) {
+        auto object = std::make_shared<ObjectTableData>();
+        ClientID node_id = ClientID::FromRandom();
+        object->set_manager(node_id.Binary());
+        object_vec.emplace_back(std::move(object));
+      }
+      ObjectID id = ObjectID::FromRandom();
+      object_id_to_data_[id] = object_vec;
+    }
+  }
+
+  typedef std::vector<std::shared_ptr<ObjectTableData>> ObjectVector;
+  std::unordered_map<ObjectID, ObjectVector> object_id_to_data_;
+
+  size_t object_count_{100};
+  size_t copy_count_{5};
+};
+
+TEST_F(RedisObjectInfoAccessorTest, TestGetAddRemove) {
+  ObjectInfoAccessor &object_accessor = gcs_client_->Objects();
+  // add && get
+  // add
+  for (const auto &elem : object_id_to_data_) {
+    for (const auto &item : elem.second) {
+      ++pending_count_;
+      ClientID node_id = ClientID::FromBinary(item->manager());
+      RAY_CHECK_OK(
+          object_accessor.AsyncAddLocation(elem.first, node_id, [this](Status status) {
+            RAY_CHECK_OK(status);
+            --pending_count_;
+          }));
+    }
+  }
+  WaitPendingDone(wait_pending_timeout_);
+  // get
+  for (const auto &elem : object_id_to_data_) {
+    ++pending_count_;
+    size_t total_size = elem.second.size();
+    RAY_CHECK_OK(object_accessor.AsyncGetLocations(
+        elem.first,
+        [this, total_size](Status status, const std::vector<ObjectTableData> &result) {
+          RAY_CHECK_OK(status);
+          RAY_CHECK(total_size == result.size());
+          --pending_count_;
+        }));
+  }
+  WaitPendingDone(wait_pending_timeout_);
+
+  RAY_LOG(INFO) << "Case Add && Get done.";
+
+  // subscribe && delete
+  // subscribe
+  std::atomic<int> sub_pending_count(0);
+  auto subscribe = [this, &sub_pending_count](const ObjectID &object_id,
+                                              const ObjectChangeNotification &result) {
+    const auto it = object_id_to_data_.find(object_id);
+    ASSERT_TRUE(it != object_id_to_data_.end());
+    static size_t response_count = 1;
+    size_t cur_count = response_count <= object_count_ ? copy_count_ : 1;
+    ASSERT_EQ(result.GetData().size(), cur_count);
+    rpc::GcsChangeMode change_mode = response_count <= object_count_
+                                         ? rpc::GcsChangeMode::APPEND_OR_ADD
+                                         : rpc::GcsChangeMode::REMOVE;
+    ASSERT_EQ(change_mode, result.GetGcsChangeMode());
+    ++response_count;
+    --sub_pending_count;
+  };
+  for (const auto &elem : object_id_to_data_) {
+    ++pending_count_;
+    ++sub_pending_count;
+    RAY_CHECK_OK(object_accessor.AsyncSubscribeToLocations(elem.first, subscribe,
+                                                           [this](Status status) {
+                                                             RAY_CHECK_OK(status);
+                                                             --pending_count_;
+                                                           }));
+  }
+  WaitPendingDone(wait_pending_timeout_);
+  WaitPendingDone(sub_pending_count, wait_pending_timeout_);
+  // delete
+  for (const auto &elem : object_id_to_data_) {
+    ++pending_count_;
+    ++sub_pending_count;
+    const ObjectVector &object_vec = elem.second;
+    ClientID node_id = ClientID::FromBinary(object_vec[0]->manager());
+    RAY_CHECK_OK(
+        object_accessor.AsyncRemoveLocation(elem.first, node_id, [this](Status status) {
+          RAY_CHECK_OK(status);
+          --pending_count_;
+        }));
+  }
+  WaitPendingDone(wait_pending_timeout_);
+  WaitPendingDone(sub_pending_count, wait_pending_timeout_);
+  // get
+  for (const auto &elem : object_id_to_data_) {
+    ++pending_count_;
+    size_t total_size = elem.second.size();
+    RAY_CHECK_OK(object_accessor.AsyncGetLocations(
+        elem.first,
+        [this, total_size](Status status, const std::vector<ObjectTableData> &result) {
+          RAY_CHECK_OK(status);
+          ASSERT_EQ(total_size - 1, result.size());
+          --pending_count_;
+        }));
+  }
+  WaitPendingDone(wait_pending_timeout_);
+
+  RAY_LOG(INFO) << "Case Subscribe && Delete done.";
+}
+
+}  // namespace gcs
+
+}  // namespace ray
+
+int main(int argc, char **argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  RAY_CHECK(argc == 4);
+  ray::REDIS_SERVER_EXEC_PATH = argv[1];
+  ray::REDIS_CLIENT_EXEC_PATH = argv[2];
+  ray::REDIS_MODULE_LIBRARY_PATH = argv[3];
+  return RUN_ALL_TESTS();
+}

src/ray/gcs/test/subscription_executor_test.cc

@@ -1,6 +1,7 @@
 #include "ray/gcs/subscription_executor.h"
 #include "gtest/gtest.h"
 #include "ray/gcs/callback.h"
+#include "ray/gcs/entry_change_notification.h"
 #include "ray/gcs/redis_gcs_client.h"
 #include "ray/gcs/test/accessor_test_base.h"
 
@@ -190,6 +191,7 @@ TEST_F(SubscriptionExecutorTest, UnsubscribeTest) {
   WaitPendingDone(do_sub_pending_count_, wait_pending_timeout_);
   sub_pending_count_ = id_to_data_.size();
   AsyncRegisterActorToGcs();
+  WaitPendingDone(pending_count_, wait_pending_timeout_);
   WaitPendingDone(sub_pending_count_, wait_pending_timeout_);
 }
 

src/ray/object_manager/object_directory.cc

@@ -15,7 +15,7 @@ using ray::rpc::ObjectTableData;
 /// Process a notification of the object table entries and store the result in
 /// node_ids. This assumes that node_ids already contains the result of the
 /// object table entries up to but not including this notification.
-void UpdateObjectLocations(const GcsChangeMode change_mode,
+void UpdateObjectLocations(bool is_added,
                            const std::vector<ObjectTableData> &location_updates,
                            std::shared_ptr<gcs::RedisGcsClient> gcs_client,
                            std::unordered_set<ClientID> *node_ids) {
@@ -24,7 +24,7 @@ void UpdateObjectLocations(const GcsChangeMode change_mode,
   // addition or deletion.
   for (const auto &object_table_data : location_updates) {
     ClientID node_id = ClientID::FromBinary(object_table_data.manager());
-    if (change_mode != GcsChangeMode::REMOVE) {
+    if (is_added) {
       node_ids->insert(node_id);
     } else {
       node_ids->erase(node_id);
@@ -42,52 +42,12 @@ void UpdateObjectLocations(const GcsChangeMode change_mode,
 
 }  // namespace
 
-void ObjectDirectory::RegisterBackend() {
-  auto object_notification_callback =
-      [this](gcs::RedisGcsClient *client, const ObjectID &object_id,
-             const GcsChangeMode change_mode,
-             const std::vector<ObjectTableData> &location_updates) {
-        // Objects are added to this map in SubscribeObjectLocations.
-        auto it = listeners_.find(object_id);
-        // Do nothing for objects we are not listening for.
-        if (it == listeners_.end()) {
-          return;
-        }
-
-        // Once this flag is set to true, it should never go back to false.
-        it->second.subscribed = true;
-
-        // Update entries for this object.
-        UpdateObjectLocations(change_mode, location_updates, gcs_client_,
-                              &it->second.current_object_locations);
-        // Copy the callbacks so that the callbacks can unsubscribe without interrupting
-        // looping over the callbacks.
-        auto callbacks = it->second.callbacks;
-        // Call all callbacks associated with the object id locations we have
-        // received.  This notifies the client even if the list of locations is
-        // empty, since this may indicate that the objects have been evicted from
-        // all nodes.
-        for (const auto &callback_pair : callbacks) {
-          // It is safe to call the callback directly since this is already running
-          // in the subscription callback stack.
-          callback_pair.second(object_id, it->second.current_object_locations);
-        }
-      };
-  RAY_CHECK_OK(gcs_client_->object_table().Subscribe(
-      JobID::Nil(), gcs_client_->Nodes().GetSelfId(), object_notification_callback,
-      nullptr));
-}
-
 ray::Status ObjectDirectory::ReportObjectAdded(
     const ObjectID &object_id, const ClientID &client_id,
     const object_manager::protocol::ObjectInfoT &object_info) {
   RAY_LOG(DEBUG) << "Reporting object added to GCS " << object_id;
-  // Append the addition entry to the object table.
-  auto data = std::make_shared<ObjectTableData>();
-  data->set_manager(client_id.Binary());
-  data->set_object_size(object_info.data_size);
   ray::Status status =
-      gcs_client_->object_table().Add(JobID::Nil(), object_id, data, nullptr);
+      gcs_client_->Objects().AsyncAddLocation(object_id, client_id, nullptr);
   return status;
 }
 
@@ -95,12 +55,8 @@ ray::Status ObjectDirectory::ReportObjectRemoved(
     const ObjectID &object_id, const ClientID &client_id,
     const object_manager::protocol::ObjectInfoT &object_info) {
   RAY_LOG(DEBUG) << "Reporting object removed to GCS " << object_id;
-  // Append the eviction entry to the object table.
-  auto data = std::make_shared<ObjectTableData>();
-  data->set_manager(client_id.Binary());
-  data->set_object_size(object_info.data_size);
   ray::Status status =
-      gcs_client_->object_table().Remove(JobID::Nil(), object_id, data, nullptr);
+      gcs_client_->Objects().AsyncRemoveLocation(object_id, client_id, nullptr);
   return status;
 };
 
@@ -136,7 +92,7 @@ void ObjectDirectory::HandleClientRemoved(const ClientID &client_id) {
     if (listener.second.current_object_locations.count(client_id) > 0) {
       // If the subscribed object has the removed client as a location, update
       // its locations with an empty update so that the location will be removed.
-      UpdateObjectLocations(GcsChangeMode::APPEND_OR_ADD, {}, gcs_client_,
+      UpdateObjectLocations(/*is_added*/ true, {}, gcs_client_,
                             &listener.second.current_object_locations);
       // Re-call all the subscribed callbacks for the object, since its
       // locations have changed.
@@ -156,10 +112,41 @@ ray::Status ObjectDirectory::SubscribeObjectLocations(const UniqueID &callback_i
   auto it = listeners_.find(object_id);
   if (it == listeners_.end()) {
     it = listeners_.emplace(object_id, LocationListenerState()).first;
-    status = gcs_client_->object_table().RequestNotifications(
-        JobID::Nil(), object_id, gcs_client_->Nodes().GetSelfId(),
-        /*done*/ nullptr);
+
+    auto object_notification_callback =
+        [this](const ObjectID &object_id,
+               const gcs::ObjectChangeNotification &object_notification) {
+          // Objects are added to this map in SubscribeObjectLocations.
+          auto it = listeners_.find(object_id);
+          // Do nothing for objects we are not listening for.
+          if (it == listeners_.end()) {
+            return;
+          }
+
+          // Once this flag is set to true, it should never go back to false.
+          it->second.subscribed = true;
+
+          // Update entries for this object.
+          UpdateObjectLocations(object_notification.IsAdded(),
+                                object_notification.GetData(), gcs_client_,
+                                &it->second.current_object_locations);
+          // Copy the callbacks so that the callbacks can unsubscribe without interrupting
+          // looping over the callbacks.
+          auto callbacks = it->second.callbacks;
+          // Call all callbacks associated with the object id locations we have
+          // received.  This notifies the client even if the list of locations is
+          // empty, since this may indicate that the objects have been evicted from
+          // all nodes.
+          for (const auto &callback_pair : callbacks) {
+            // It is safe to call the callback directly since this is already running
+            // in the subscription callback stack.
+            callback_pair.second(object_id, it->second.current_object_locations);
+          }
+        };
+    status = gcs_client_->Objects().AsyncSubscribeToLocations(
+        object_id, object_notification_callback, /*done*/ nullptr);
   }
+
   auto &listener_state = it->second;
   // TODO(hme): Make this fatal after implementing Pull suppression.
   if (listener_state.callbacks.count(callback_id) > 0) {
@@ -185,8 +172,8 @@ ray::Status ObjectDirectory::UnsubscribeObjectLocations(const UniqueID &callback
   }
   entry->second.callbacks.erase(callback_id);
   if (entry->second.callbacks.empty()) {
-    status = gcs_client_->object_table().CancelNotifications(
-        JobID::Nil(), object_id, gcs_client_->Nodes().GetSelfId(), /*done*/ nullptr);
+    status =
+        gcs_client_->Objects().AsyncUnsubscribeToLocations(object_id, /*done*/ nullptr);
     listeners_.erase(entry);
   }
   return status;
@@ -208,14 +195,16 @@ ray::Status ObjectDirectory::LookupLocations(const ObjectID &object_id,
     // We do not have any locations cached due to a concurrent
     // SubscribeObjectLocations call, so look up the object's locations
     // directly from the GCS.
-    status = gcs_client_->object_table().Lookup(
-        JobID::Nil(), object_id,
-        [this, callback](gcs::RedisGcsClient *client, const ObjectID &object_id,
-                         const std::vector<ObjectTableData> &location_updates) {
+    status = gcs_client_->Objects().AsyncGetLocations(
+        object_id,
+        [this, object_id, callback](
+            Status status, const std::vector<ObjectTableData> &location_updates) {
+          RAY_CHECK(status.ok())
+              << "Failed to get object location from GCS: " << status.message();
           // Build the set of current locations based on the entries in the log.
           std::unordered_set<ClientID> node_ids;
-          UpdateObjectLocations(GcsChangeMode::APPEND_OR_ADD, location_updates,
-                                gcs_client_, &node_ids);
+          UpdateObjectLocations(/*is_added*/ true, location_updates, gcs_client_,
+                                &node_ids);
           // It is safe to call the callback directly since this is already running
           // in the GCS client's lookup callback stack.
           callback(object_id, node_ids);

src/ray/object_manager/object_directory.h

@@ -33,8 +33,6 @@ class ObjectDirectoryInterface {
  public:
   virtual ~ObjectDirectoryInterface() {}
 
-  virtual void RegisterBackend() = 0;
-
   /// Lookup how to connect to a remote object manager.
   ///
   /// \param connection_info The connection information to fill out. This
@@ -135,8 +133,6 @@ class ObjectDirectory : public ObjectDirectoryInterface {
 
   virtual ~ObjectDirectory() {}
 
-  void RegisterBackend() override;
-
   void LookupRemoteConnectionInfo(RemoteConnectionInfo &connection_info) const override;
 
   std::vector<RemoteConnectionInfo> LookupAllRemoteConnections() const override;

src/ray/object_manager/object_manager.cc

@@ -38,8 +38,6 @@ ObjectManager::ObjectManager(asio::io_service &main_service, const ClientID &sel
 
 ObjectManager::~ObjectManager() { StopRpcService(); }
 
-void ObjectManager::RegisterGcs() { object_directory_->RegisterBackend(); }
-
 void ObjectManager::RunRpcService() { rpc_service_.run(); }
 
 void ObjectManager::StartRpcService() {

src/ray/object_manager/object_manager.h

@@ -157,9 +157,6 @@ class ObjectManager : public ObjectManagerInterface,
 
   ~ObjectManager();
 
-  /// Register GCS-related functionality.
-  void RegisterGcs();
-
   /// Subscribe to notifications of objects added to local store.
   /// Upon subscribing, the callback will be invoked for all objects that
   ///

src/ray/object_manager/test/object_manager_stress_test.cc

@@ -55,7 +55,6 @@ class MockServer {
     node_info.set_object_manager_port(object_manager_port);
 
     ray::Status status = gcs_client_->Nodes().RegisterSelf(node_info);
-    object_manager_.RegisterGcs();
     return status;
   }
 

src/ray/object_manager/test/object_manager_test.cc

@@ -49,7 +49,6 @@ class MockServer {
     node_info.set_object_manager_port(object_manager_port);
 
     ray::Status status = gcs_client_->Nodes().RegisterSelf(node_info);
-    object_manager_.RegisterGcs();
     return status;
   }
 

src/ray/raylet/mock_gcs_client.cc

@@ -114,8 +114,6 @@ ClientID GcsClient::Register(const std::string &ip, uint16_t port) {
   return client_id;
 }
 
-ObjectTable &GcsClient::object_table() { return *object_table_; }
-
 ClientTable &GcsClient::client_table() { return *client_table_; }
 
 }  // namespace ray

src/ray/raylet/mock_gcs_client.h

@@ -80,7 +80,6 @@ class GcsClient {
   }
   // Register the ip and port of the connecting client.
   ClientID Register(const std::string &ip, uint16_t port);
-  ObjectTable &object_table();
   ClientTable &client_table();
 
  private:

src/ray/raylet/node_manager.cc

@@ -136,8 +136,6 @@ NodeManager::NodeManager(boost::asio::io_service &io_service,
 }
 
 ray::Status NodeManager::RegisterGcs() {
-  object_manager_.RegisterGcs();
-
   const auto task_lease_notification_callback = [this](gcs::RedisGcsClient *client,
                                                        const TaskID &task_id,
                                                        const TaskLeaseData &task_lease) {

src/ray/raylet/reconstruction_policy_test.cc

@@ -64,7 +64,7 @@ class MockObjectDirectory : public ObjectDirectoryInterface {
 
   std::string DebugString() const override { return ""; }
 
-  MOCK_METHOD0(RegisterBackend, void(void));
+  MOCK_METHOD0(GetLocalClientID, ray::ClientID());
   MOCK_CONST_METHOD1(LookupRemoteConnectionInfo, void(RemoteConnectionInfo &));
   MOCK_CONST_METHOD0(LookupAllRemoteConnections, std::vector<RemoteConnectionInfo>());
   MOCK_METHOD3(SubscribeObjectLocations,