[GCS]Add task lease methods to TaskInfoAccessor (#6645)

src/ray/gcs/accessor.h

@@ -197,13 +197,41 @@ class TaskInfoAccessor {
       const StatusCallback &done) = 0;
 
   /// Cancel subscription to a task asynchronously.
-  /// This method is for node only (core worker shouldn't use this method).
   ///
   /// \param task_id The ID of the task to be unsubscribed to.
   /// \param done Callback that will be called when unsubscribe is complete.
   /// \return Status
   virtual Status AsyncUnsubscribe(const TaskID &task_id, const StatusCallback &done) = 0;
 
+  /// Add a task lease to GCS asynchronously.
+  ///
+  /// \param data_ptr The task lease that will be added to GCS.
+  /// \param callback Callback that will be called after task lease has been added
+  /// to GCS.
+  /// \return Status
+  virtual Status AsyncAddTaskLease(const std::shared_ptr<rpc::TaskLeaseData> &data_ptr,
+                                   const StatusCallback &callback) = 0;
+
+  /// Subscribe asynchronously to the event that the given task lease is added in GCS.
+  ///
+  /// \param task_id The ID of the task to be subscribed to.
+  /// \param subscribe Callback that will be called each time when the task lease is
+  /// updated or the task lease is empty currently.
+  /// \param done Callback that will be called when subscription is complete.
+  /// \return Status
+  virtual Status AsyncSubscribeTaskLease(
+      const TaskID &task_id,
+      const SubscribeCallback<TaskID, boost::optional<rpc::TaskLeaseData>> &subscribe,
+      const StatusCallback &done) = 0;
+
+  /// Cancel subscription to a task lease asynchronously.
+  ///
+  /// \param task_id The ID of the task to be unsubscribed to.
+  /// \param done Callback that will be called when unsubscribe is complete.
+  /// \return Status
+  virtual Status AsyncUnsubscribeTaskLease(const TaskID &task_id,
+                                           const StatusCallback &done) = 0;
+
  protected:
   TaskInfoAccessor() = default;
 };

src/ray/gcs/redis_accessor.cc

@@ -247,7 +247,9 @@ Status RedisJobInfoAccessor::AsyncSubscribeToFinishedJobs(
 }
 
 RedisTaskInfoAccessor::RedisTaskInfoAccessor(RedisGcsClient *client_impl)
-    : client_impl_(client_impl), task_sub_executor_(client_impl->raylet_task_table()) {}
+    : client_impl_(client_impl),
+      task_sub_executor_(client_impl->raylet_task_table()),
+      task_lease_sub_executor_(client_impl->task_lease_table()) {}
 
 Status RedisTaskInfoAccessor::AsyncAdd(const std::shared_ptr<TaskTableData> &data_ptr,
                                        const StatusCallback &callback) {
@@ -301,6 +303,31 @@ Status RedisTaskInfoAccessor::AsyncUnsubscribe(const TaskID &task_id,
   return task_sub_executor_.AsyncUnsubscribe(subscribe_id_, task_id, done);
 }
 
+Status RedisTaskInfoAccessor::AsyncAddTaskLease(
+    const std::shared_ptr<TaskLeaseData> &data_ptr, const StatusCallback &callback) {
+  TaskLeaseTable::WriteCallback on_done = nullptr;
+  if (callback != nullptr) {
+    on_done = [callback](RedisGcsClient *client, const TaskID &id,
+                         const TaskLeaseData &data) { callback(Status::OK()); };
+  }
+  TaskID task_id = TaskID::FromBinary(data_ptr->task_id());
+  TaskLeaseTable &task_lease_table = client_impl_->task_lease_table();
+  return task_lease_table.Add(JobID::Nil(), task_id, data_ptr, on_done);
+}
+
+Status RedisTaskInfoAccessor::AsyncSubscribeTaskLease(
+    const TaskID &task_id,
+    const SubscribeCallback<TaskID, boost::optional<TaskLeaseData>> &subscribe,
+    const StatusCallback &done) {
+  RAY_CHECK(subscribe != nullptr);
+  return task_lease_sub_executor_.AsyncSubscribe(subscribe_id_, task_id, subscribe, done);
+}
+
+Status RedisTaskInfoAccessor::AsyncUnsubscribeTaskLease(const TaskID &task_id,
+                                                        const StatusCallback &done) {
+  return task_lease_sub_executor_.AsyncUnsubscribe(subscribe_id_, task_id, done);
+}
+
 RedisObjectInfoAccessor::RedisObjectInfoAccessor(RedisGcsClient *client_impl)
     : client_impl_(client_impl), object_sub_executor_(client_impl->object_table()) {}
 

src/ray/gcs/redis_accessor.h

@@ -122,21 +122,33 @@ class RedisTaskInfoAccessor : public TaskInfoAccessor {
  public:
   explicit RedisTaskInfoAccessor(RedisGcsClient *client_impl);
 
-  ~RedisTaskInfoAccessor() {}
+  virtual ~RedisTaskInfoAccessor() {}
 
   Status AsyncAdd(const std::shared_ptr<TaskTableData> &data_ptr,
-                  const StatusCallback &callback);
+                  const StatusCallback &callback) override;
 
   Status AsyncGet(const TaskID &task_id,
-                  const OptionalItemCallback<TaskTableData> &callback);
+                  const OptionalItemCallback<TaskTableData> &callback) override;
 
-  Status AsyncDelete(const std::vector<TaskID> &task_ids, const StatusCallback &callback);
+  Status AsyncDelete(const std::vector<TaskID> &task_ids,
+                     const StatusCallback &callback) override;
 
   Status AsyncSubscribe(const TaskID &task_id,
                         const SubscribeCallback<TaskID, TaskTableData> &subscribe,
-                        const StatusCallback &done);
+                        const StatusCallback &done) override;
 
-  Status AsyncUnsubscribe(const TaskID &task_id, const StatusCallback &done);
+  Status AsyncUnsubscribe(const TaskID &task_id, const StatusCallback &done) override;
+
+  Status AsyncAddTaskLease(const std::shared_ptr<TaskLeaseData> &data_ptr,
+                           const StatusCallback &callback) override;
+
+  Status AsyncSubscribeTaskLease(
+      const TaskID &task_id,
+      const SubscribeCallback<TaskID, boost::optional<TaskLeaseData>> &subscribe,
+      const StatusCallback &done) override;
+
+  Status AsyncUnsubscribeTaskLease(const TaskID &task_id,
+                                   const StatusCallback &done) override;
 
  private:
   RedisGcsClient *client_impl_{nullptr};
@@ -151,6 +163,10 @@ class RedisTaskInfoAccessor : public TaskInfoAccessor {
   typedef SubscriptionExecutor<TaskID, TaskTableData, raylet::TaskTable>
       TaskSubscriptionExecutor;
   TaskSubscriptionExecutor task_sub_executor_;
+
+  typedef SubscriptionExecutor<TaskID, boost::optional<TaskLeaseData>, TaskLeaseTable>
+      TaskLeaseSubscriptionExecutor;
+  TaskLeaseSubscriptionExecutor task_lease_sub_executor_;
 };
 
 /// \class RedisObjectInfoAccessor

src/ray/gcs/redis_gcs_client.h

@@ -63,7 +63,6 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
 
   // TODO: Some API for getting the error on the driver
   TaskReconstructionLog &task_reconstruction_log();
-  TaskLeaseTable &task_lease_table();
   ErrorTable &error_table();
   ProfileTable &profile_table();
 
@@ -101,8 +100,9 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
   HeartbeatTable &heartbeat_table();
   HeartbeatBatchTable &heartbeat_batch_table();
   DynamicResourceTable &resource_table();
-  /// This method will be deprecated, use method Tasks() instead.
+  /// The following two methods will be deprecated, use method Tasks() instead.
   raylet::TaskTable &raylet_task_table();
+  TaskLeaseTable &task_lease_table();
 
   // GCS command type. If CommandType::kChain, chain-replicated versions of the tables
   // might be used, if available.

src/ray/gcs/subscription_executor.cc

@@ -188,6 +188,8 @@ template class SubscriptionExecutor<ActorID, ActorTableData, ActorTable>;
 template class SubscriptionExecutor<JobID, JobTableData, JobTable>;
 template class SubscriptionExecutor<TaskID, TaskTableData, raylet::TaskTable>;
 template class SubscriptionExecutor<ObjectID, ObjectChangeNotification, ObjectTable>;
+template class SubscriptionExecutor<TaskID, boost::optional<TaskLeaseData>,
+                                    TaskLeaseTable>;
 template class SubscriptionExecutor<ClientID, ResourceChangeNotification,
                                     DynamicResourceTable>;
 template class SubscriptionExecutor<ClientID, HeartbeatTableData, HeartbeatTable>;

src/ray/gcs/subscription_executor.h

@@ -18,7 +18,7 @@ namespace gcs {
 template <typename ID, typename Data, typename Table>
 class SubscriptionExecutor {
  public:
-  SubscriptionExecutor(Table &table) : table_(table) {}
+  explicit SubscriptionExecutor(Table &table) : table_(table) {}
 
   ~SubscriptionExecutor() {}
 

src/ray/gcs/tables.cc

@@ -730,6 +730,25 @@ std::string ClientTable::DebugString() const {
   return result.str();
 }
 
+Status TaskLeaseTable::Subscribe(const JobID &job_id, const ClientID &client_id,
+                                 const Callback &subscribe,
+                                 const SubscriptionCallback &done) {
+  auto on_subscribe = [subscribe](RedisGcsClient *client, const TaskID &task_id,
+                                  const std::vector<TaskLeaseData> &data) {
+    std::vector<boost::optional<TaskLeaseData>> result;
+    for (const auto &item : data) {
+      boost::optional<TaskLeaseData> optional_item(item);
+      result.emplace_back(std::move(optional_item));
+    }
+    if (result.empty()) {
+      boost::optional<TaskLeaseData> optional_item;
+      result.emplace_back(std::move(optional_item));
+    }
+    subscribe(client, task_id, result);
+  };
+  return Table<TaskID, TaskLeaseData>::Subscribe(job_id, client_id, on_subscribe, done);
+}
+
 Status ActorCheckpointIdTable::AddCheckpointId(const JobID &job_id,
                                                const ActorID &actor_id,
                                                const ActorCheckpointID &checkpoint_id,

src/ray/gcs/tables.h

@@ -727,6 +727,12 @@ class TaskReconstructionLog : public Log<TaskID, TaskReconstructionData> {
 
 class TaskLeaseTable : public Table<TaskID, TaskLeaseData> {
  public:
+  /// Use boost::optional to represent subscription results, so that we can
+  /// notify raylet whether the entry of task lease is empty.
+  using Callback =
+      std::function<void(RedisGcsClient *client, const TaskID &task_id,
+                         const std::vector<boost::optional<TaskLeaseData>> &data)>;
+
   TaskLeaseTable(const std::vector<std::shared_ptr<RedisContext>> &contexts,
                  RedisGcsClient *client)
       : Table(contexts, client) {
@@ -749,6 +755,11 @@ class TaskLeaseTable : public Table<TaskID, TaskLeaseData> {
 
     return GetRedisContext(id)->RunArgvAsync(args);
   }
+
+  /// Implement this method for the subscription tools class SubscriptionExecutor.
+  /// In this way TaskLeaseTable() can also reuse class SubscriptionExecutor.
+  Status Subscribe(const JobID &job_id, const ClientID &client_id,
+                   const Callback &subscribe, const SubscriptionCallback &done);
 };
 
 class ActorCheckpointTable : public Table<ActorCheckpointID, ActorCheckpointData> {

src/ray/protobuf/gcs.proto

@@ -213,14 +213,16 @@ message HeartbeatBatchTableData {
 
 // Data for a lease on task execution.
 message TaskLeaseData {
+  // The task ID.
+  bytes task_id = 1;
   // Node manager client ID.
-  bytes node_manager_id = 1;
+  bytes node_manager_id = 2;
   // The time that the lease was last acquired at. NOTE(swang): This is the
   // system clock time according to the node that added the entry and is not
   // synchronized with other nodes.
-  uint64 acquired_at = 2;
+  uint64 acquired_at = 3;
   // The period that the lease is active for.
-  uint64 timeout = 3;
+  uint64 timeout = 4;
 }
 
 message JobTableData {

src/ray/raylet/node_manager.cc

@@ -96,12 +96,12 @@ NodeManager::NodeManager(boost::asio::io_service &io_service,
             HandleTaskReconstruction(task_id, required_object_id);
           },
           RayConfig::instance().initial_reconstruction_timeout_milliseconds(),
-          self_node_id_, gcs_client_->task_lease_table(), object_directory_,
+          self_node_id_, gcs_client_, object_directory_,
           gcs_client_->task_reconstruction_log()),
       task_dependency_manager_(
           object_manager, reconstruction_policy_, io_service, self_node_id_,
           RayConfig::instance().initial_reconstruction_timeout_milliseconds(),
-          gcs_client_->task_lease_table()),
+          gcs_client_),
       lineage_cache_(self_node_id_, gcs_client_, config.max_lineage_size),
       actor_registry_(),
       node_manager_server_("NodeManager", config.node_manager_port),
@@ -136,30 +136,7 @@ NodeManager::NodeManager(boost::asio::io_service &io_service,
 }
 
 ray::Status NodeManager::RegisterGcs() {
-  const auto task_lease_notification_callback = [this](gcs::RedisGcsClient *client,
-                                                       const TaskID &task_id,
-                                                       const TaskLeaseData &task_lease) {
-    const ClientID node_manager_id = ClientID::FromBinary(task_lease.node_manager_id());
-    if (gcs_client_->Nodes().IsRemoved(node_manager_id)) {
-      // The node manager that added the task lease is already removed. The
-      // lease is considered inactive.
-      reconstruction_policy_.HandleTaskLeaseNotification(task_id, 0);
-    } else {
-      // NOTE(swang): The task_lease.timeout is an overestimate of the lease's
-      // expiration period since the entry may have been in the GCS for some
-      // time already. For a more accurate estimate, the age of the entry in
-      // the GCS should be subtracted from task_lease.timeout.
-      reconstruction_policy_.HandleTaskLeaseNotification(task_id, task_lease.timeout());
-    }
-  };
-  const auto task_lease_empty_callback = [this](gcs::RedisGcsClient *client,
-                                                const TaskID &task_id) {
-    reconstruction_policy_.HandleTaskLeaseNotification(task_id, 0);
-  };
-  RAY_RETURN_NOT_OK(gcs_client_->task_lease_table().Subscribe(
-      JobID::Nil(), self_node_id_, task_lease_notification_callback,
-      task_lease_empty_callback, nullptr));
-
+  // The TaskLease subscription is done on demand in reconstruction policy.
   // Register a callback to handle actor notifications.
   auto actor_notification_callback = [this](const ActorID &actor_id,
                                             const ActorTableData &data) {

src/ray/raylet/reconstruction_policy.cc

@@ -10,14 +10,14 @@ ReconstructionPolicy::ReconstructionPolicy(
     boost::asio::io_service &io_service,
     std::function<void(const TaskID &, const ObjectID &)> reconstruction_handler,
     int64_t initial_reconstruction_timeout_ms, const ClientID &client_id,
-    gcs::PubsubInterface<TaskID> &task_lease_pubsub,
+    std::shared_ptr<gcs::RedisGcsClient> gcs_client,
     std::shared_ptr<ObjectDirectoryInterface> object_directory,
     gcs::LogInterface<TaskID, TaskReconstructionData> &task_reconstruction_log)
     : io_service_(io_service),
       reconstruction_handler_(reconstruction_handler),
       initial_reconstruction_timeout_ms_(initial_reconstruction_timeout_ms),
       client_id_(client_id),
-      task_lease_pubsub_(task_lease_pubsub),
+      gcs_client_(gcs_client),
       object_directory_(std::move(object_directory)),
       task_reconstruction_log_(task_reconstruction_log) {}
 
@@ -43,6 +43,11 @@ void ReconstructionPolicy::SetTaskTimeout(
             // received. The current lease is now considered expired.
             HandleTaskLeaseExpired(task_id);
           } else {
+            const auto task_lease_notification_callback =
+                [this](const TaskID &task_id,
+                       const boost::optional<rpc::TaskLeaseData> &task_lease) {
+                  OnTaskLeaseNotification(task_id, task_lease);
+                };
             // This task is still required, so subscribe to task lease
             // notifications.  Reconstruction will be triggered if the current
             // task lease expires, or if no one has acquired the task lease.
@@ -52,9 +57,8 @@ void ReconstructionPolicy::SetTaskTimeout(
             // required by the task are no longer needed soon after.  If the
             // task is still required after this initial period, then we now
             // subscribe to task lease notifications.
-            RAY_CHECK_OK(task_lease_pubsub_.RequestNotifications(JobID::Nil(), task_id,
-                                                                 client_id_,
-                                                                 /*done*/ nullptr));
+            RAY_CHECK_OK(gcs_client_->Tasks().AsyncSubscribeTaskLease(
+                task_id, task_lease_notification_callback, /*done*/ nullptr));
             it->second.subscribed = true;
           }
         } else {
@@ -64,6 +68,28 @@ void ReconstructionPolicy::SetTaskTimeout(
       });
 }
 
+void ReconstructionPolicy::OnTaskLeaseNotification(
+    const TaskID &task_id, const boost::optional<rpc::TaskLeaseData> &task_lease) {
+  if (!task_lease) {
+    // Task lease not exist.
+    HandleTaskLeaseNotification(task_id, 0);
+    return;
+  }
+
+  const ClientID node_manager_id = ClientID::FromBinary(task_lease->node_manager_id());
+  if (gcs_client_->Nodes().IsRemoved(node_manager_id)) {
+    // The node manager that added the task lease is already removed. The
+    // lease is considered inactive.
+    HandleTaskLeaseNotification(task_id, 0);
+  } else {
+    // NOTE(swang): The task_lease.timeout is an overestimate of the
+    // lease's expiration period since the entry may have been in the GCS
+    // for some time already. For a more accurate estimate, the age of the
+    // entry in the GCS should be subtracted from task_lease.timeout.
+    HandleTaskLeaseNotification(task_id, task_lease->timeout());
+  }
+}
+
 void ReconstructionPolicy::HandleReconstructionLogAppend(
     const TaskID &task_id, const ObjectID &required_object_id, bool success) {
   auto it = listening_tasks_.find(task_id);
@@ -201,9 +227,8 @@ void ReconstructionPolicy::Cancel(const ObjectID &object_id) {
   if (it->second.created_objects.empty()) {
     // Cancel notifications for the task lease if we were subscribed to them.
     if (it->second.subscribed) {
-      RAY_CHECK_OK(task_lease_pubsub_.CancelNotifications(JobID::Nil(), task_id,
-                                                          client_id_,
-                                                          /*done*/ nullptr));
+      RAY_CHECK_OK(
+          gcs_client_->Tasks().AsyncUnsubscribeTaskLease(task_id, /*done*/ nullptr));
     }
     listening_tasks_.erase(it);
   }

src/ray/raylet/reconstruction_policy.h

@@ -37,13 +37,13 @@ class ReconstructionPolicy : public ReconstructionPolicyInterface {
   /// will be triggered.
   /// \param client_id The client ID to use when requesting notifications from
   /// the GCS.
-  /// \param task_lease_pubsub The GCS pub-sub storage system to request task
+  /// \param gcs_client The Client of GCS.
   /// lease notifications from.
   ReconstructionPolicy(
       boost::asio::io_service &io_service,
       std::function<void(const TaskID &, const ObjectID &)> reconstruction_handler,
       int64_t initial_reconstruction_timeout_ms, const ClientID &client_id,
-      gcs::PubsubInterface<TaskID> &task_lease_pubsub,
+      std::shared_ptr<gcs::RedisGcsClient> gcs_client,
       std::shared_ptr<ObjectDirectoryInterface> object_directory,
       gcs::LogInterface<TaskID, TaskReconstructionData> &task_reconstruction_log);
 
@@ -109,6 +109,10 @@ class ReconstructionPolicy : public ReconstructionPolicyInterface {
   void SetTaskTimeout(std::unordered_map<TaskID, ReconstructionTask>::iterator task_it,
                       int64_t timeout_ms);
 
+  /// Handle task lease notification from GCS.
+  void OnTaskLeaseNotification(const TaskID &task_id,
+                               const boost::optional<rpc::TaskLeaseData> &task_lease);
+
   /// Attempt to re-execute a task to reconstruct the required object.
   ///
   /// \param task_id The task to attempt to re-execute.
@@ -138,8 +142,8 @@ class ReconstructionPolicy : public ReconstructionPolicyInterface {
   const int64_t initial_reconstruction_timeout_ms_;
   /// The client ID to use when requesting notifications from the GCS.
   const ClientID client_id_;
-  /// The GCS pub-sub storage system to request task lease notifications from.
-  gcs::PubsubInterface<TaskID> &task_lease_pubsub_;
+  /// A client connection to the GCS.
+  std::shared_ptr<gcs::RedisGcsClient> gcs_client_;
   /// The object directory used to lookup object locations.
   std::shared_ptr<ObjectDirectoryInterface> object_directory_;
   gcs::LogInterface<TaskID, TaskReconstructionData> &task_reconstruction_log_;

src/ray/raylet/reconstruction_policy_test.cc

@@ -7,6 +7,7 @@
 #include <boost/asio.hpp>
 
 #include "ray/gcs/callback.h"
+#include "ray/gcs/redis_accessor.h"
 
 #include "ray/raylet/format/node_manager_generated.h"
 #include "ray/raylet/reconstruction_policy.h"
@@ -84,44 +85,68 @@ class MockObjectDirectory : public ObjectDirectoryInterface {
   std::unordered_map<ObjectID, std::unordered_set<ClientID>> locations_;
 };
 
-class MockGcs : public gcs::PubsubInterface<TaskID>,
-                public ray::gcs::LogInterface<TaskID, TaskReconstructionData> {
+class MockNodeInfoAccessor : public gcs::RedisNodeInfoAccessor {
  public:
-  MockGcs() : notification_callback_(nullptr), failure_callback_(nullptr){};
+  MockNodeInfoAccessor(gcs::RedisGcsClient *client)
+      : gcs::RedisNodeInfoAccessor(client) {}
 
-  void Subscribe(const gcs::TaskLeaseTable::WriteCallback &notification_callback,
-                 const gcs::TaskLeaseTable::FailureCallback &failure_callback) {
-    notification_callback_ = notification_callback;
-    failure_callback_ = failure_callback;
-  }
+  bool IsRemoved(const ClientID &node_id) const override { return false; }
+};
 
-  void Add(const JobID &job_id, const TaskID &task_id,
-           const std::shared_ptr<TaskLeaseData> &task_lease_data) {
-    task_lease_table_[task_id] = task_lease_data;
-    if (subscribed_tasks_.count(task_id) == 1) {
-      notification_callback_(nullptr, task_id, *task_lease_data);
-    }
-  }
+class MockTaskInfoAccessor : public gcs::RedisTaskInfoAccessor {
+ public:
+  MockTaskInfoAccessor(gcs::RedisGcsClient *client) : RedisTaskInfoAccessor(client) {}
 
-  Status RequestNotifications(const JobID &job_id, const TaskID &task_id,
-                              const ClientID &client_id,
-                              const gcs::StatusCallback &done) {
+  Status AsyncSubscribeTaskLease(
+      const TaskID &task_id,
+      const gcs::SubscribeCallback<TaskID, boost::optional<TaskLeaseData>> &subscribe,
+      const gcs::StatusCallback &done) override {
+    subscribe_callback_ = subscribe;
     subscribed_tasks_.insert(task_id);
     auto entry = task_lease_table_.find(task_id);
     if (entry == task_lease_table_.end()) {
-      failure_callback_(nullptr, task_id);
+      boost::optional<TaskLeaseData> result;
+      subscribe(task_id, result);
     } else {
-      notification_callback_(nullptr, task_id, *entry->second);
+      boost::optional<TaskLeaseData> result(*entry->second);
+      subscribe(task_id, result);
     }
     return ray::Status::OK();
   }
 
-  Status CancelNotifications(const JobID &job_id, const TaskID &task_id,
-                             const ClientID &client_id, const gcs::StatusCallback &done) {
+  Status AsyncUnsubscribeTaskLease(const TaskID &task_id,
+                                   const gcs::StatusCallback &done) override {
     subscribed_tasks_.erase(task_id);
     return ray::Status::OK();
   }
 
+  Status AsyncAddTaskLease(const std::shared_ptr<TaskLeaseData> &task_lease_data,
+                           const gcs::StatusCallback &done) override {
+    TaskID task_id = TaskID::FromBinary(task_lease_data->task_id());
+    task_lease_table_[task_id] = task_lease_data;
+    if (subscribed_tasks_.count(task_id) == 1) {
+      boost::optional<TaskLeaseData> result(*task_lease_data);
+      subscribe_callback_(task_id, result);
+    }
+    return Status::OK();
+  }
+
+ private:
+  gcs::SubscribeCallback<TaskID, boost::optional<TaskLeaseData>> subscribe_callback_;
+  std::unordered_map<TaskID, std::shared_ptr<TaskLeaseData>> task_lease_table_;
+  std::unordered_set<TaskID> subscribed_tasks_;
+};
+
+class MockGcs : public gcs::RedisGcsClient,
+                public ray::gcs::LogInterface<TaskID, TaskReconstructionData> {
+ public:
+  MockGcs() : gcs::RedisGcsClient(gcs::GcsClientOptions("", 0, "")){};
+
+  void Init(gcs::TaskInfoAccessor *task_accessor, gcs::NodeInfoAccessor *node_accessor) {
+    task_accessor_.reset(task_accessor);
+    node_accessor_.reset(node_accessor);
+  }
+
   Status AppendAt(
       const JobID &job_id, const TaskID &task_id,
       const std::shared_ptr<TaskReconstructionData> &task_data,
@@ -150,10 +175,6 @@ class MockGcs : public gcs::PubsubInterface<TaskID>,
           const ray::gcs::LogInterface<TaskID, TaskReconstructionData>::WriteCallback &));
 
  private:
-  gcs::TaskLeaseTable::WriteCallback notification_callback_;
-  gcs::TaskLeaseTable::FailureCallback failure_callback_;
-  std::unordered_map<TaskID, std::shared_ptr<TaskLeaseData>> task_lease_table_;
-  std::unordered_set<TaskID> subscribed_tasks_;
   std::unordered_map<TaskID, std::vector<TaskReconstructionData>>
       task_reconstruction_log_;
 };
@@ -162,7 +183,9 @@ class ReconstructionPolicyTest : public ::testing::Test {
  public:
   ReconstructionPolicyTest()
       : io_service_(),
-        mock_gcs_(),
+        mock_gcs_(new MockGcs()),
+        task_accessor_(new MockTaskInfoAccessor(mock_gcs_.get())),
+        node_accessor_(new MockNodeInfoAccessor(mock_gcs_.get())),
         mock_object_directory_(std::make_shared<MockObjectDirectory>()),
         reconstruction_timeout_ms_(50),
         reconstruction_policy_(std::make_shared<ReconstructionPolicy>(
@@ -171,17 +194,19 @@ class ReconstructionPolicyTest : public ::testing::Test {
               TriggerReconstruction(task_id);
             },
             reconstruction_timeout_ms_, ClientID::FromRandom(), mock_gcs_,
-            mock_object_directory_, mock_gcs_)),
+            mock_object_directory_, *mock_gcs_)),
         timer_canceled_(false) {
-    mock_gcs_.Subscribe(
-        [this](gcs::RedisGcsClient *client, const TaskID &task_id,
-               const TaskLeaseData &task_lease) {
-          reconstruction_policy_->HandleTaskLeaseNotification(task_id,
-                                                              task_lease.timeout());
-        },
-        [this](gcs::RedisGcsClient *client, const TaskID &task_id) {
-          reconstruction_policy_->HandleTaskLeaseNotification(task_id, 0);
-        });
+    subscribe_callback_ = [this](const TaskID &task_id,
+                                 const boost::optional<TaskLeaseData> &task_lease) {
+      if (task_lease) {
+        reconstruction_policy_->HandleTaskLeaseNotification(task_id,
+                                                            task_lease->timeout());
+      } else {
+        reconstruction_policy_->HandleTaskLeaseNotification(task_id, 0);
+      }
+    };
+
+    mock_gcs_->Init(task_accessor_, node_accessor_);
   }
 
   void TriggerReconstruction(const TaskID &task_id) { reconstructed_tasks_[task_id]++; }
@@ -230,7 +255,10 @@ class ReconstructionPolicyTest : public ::testing::Test {
 
  protected:
   boost::asio::io_service io_service_;
-  MockGcs mock_gcs_;
+  std::shared_ptr<MockGcs> mock_gcs_;
+  MockTaskInfoAccessor *task_accessor_;
+  MockNodeInfoAccessor *node_accessor_;
+  gcs::SubscribeCallback<TaskID, boost::optional<TaskLeaseData>> subscribe_callback_;
   std::shared_ptr<MockObjectDirectory> mock_object_directory_;
   uint64_t reconstruction_timeout_ms_;
   std::shared_ptr<ReconstructionPolicy> reconstruction_policy_;
@@ -340,7 +368,8 @@ TEST_F(ReconstructionPolicyTest, TestReconstructionSuppressed) {
   task_lease_data->set_node_manager_id(ClientID::FromRandom().Binary());
   task_lease_data->set_acquired_at(absl::GetCurrentTimeNanos() / 1000000);
   task_lease_data->set_timeout(2 * test_period);
-  mock_gcs_.Add(JobID::Nil(), task_id, task_lease_data);
+  task_lease_data->set_task_id(task_id.Binary());
+  RAY_CHECK_OK(mock_gcs_->Tasks().AsyncAddTaskLease(task_lease_data, nullptr));
 
   // Listen for an object.
   reconstruction_policy_->ListenAndMaybeReconstruct(object_id);
@@ -368,7 +397,8 @@ TEST_F(ReconstructionPolicyTest, TestReconstructionContinuallySuppressed) {
     task_lease_data->set_node_manager_id(ClientID::FromRandom().Binary());
     task_lease_data->set_acquired_at(absl::GetCurrentTimeNanos() / 1000000);
     task_lease_data->set_timeout(reconstruction_timeout_ms_);
-    mock_gcs_.Add(JobID::Nil(), task_id, task_lease_data);
+    task_lease_data->set_task_id(task_id.Binary());
+    RAY_CHECK_OK(mock_gcs_->Tasks().AsyncAddTaskLease(task_lease_data, nullptr));
   });
   // Run the test for much longer than the reconstruction timeout.
   Run(reconstruction_timeout_ms_ * 2);
@@ -422,11 +452,11 @@ TEST_F(ReconstructionPolicyTest, TestSimultaneousReconstructionSuppressed) {
   task_reconstruction_data->set_node_manager_id(ClientID::FromRandom().Binary());
   task_reconstruction_data->set_num_reconstructions(0);
   RAY_CHECK_OK(
-      mock_gcs_.AppendAt(JobID::Nil(), task_id, task_reconstruction_data, nullptr,
-                         /*failure_callback=*/
-                         [](ray::gcs::RedisGcsClient *client, const TaskID &task_id,
-                            const TaskReconstructionData &data) { ASSERT_TRUE(false); },
-                         /*log_index=*/0));
+      mock_gcs_->AppendAt(JobID::Nil(), task_id, task_reconstruction_data, nullptr,
+                          /*failure_callback=*/
+                          [](ray::gcs::RedisGcsClient *client, const TaskID &task_id,
+                             const TaskReconstructionData &data) { ASSERT_TRUE(false); },
+                          /*log_index=*/0));
 
   // Listen for an object.
   reconstruction_policy_->ListenAndMaybeReconstruct(object_id);

src/ray/raylet/task_dependency_manager.cc

@@ -12,14 +12,13 @@ TaskDependencyManager::TaskDependencyManager(
     ObjectManagerInterface &object_manager,
     ReconstructionPolicyInterface &reconstruction_policy,
     boost::asio::io_service &io_service, const ClientID &client_id,
-    int64_t initial_lease_period_ms,
-    gcs::TableInterface<TaskID, TaskLeaseData> &task_lease_table)
+    int64_t initial_lease_period_ms, std::shared_ptr<gcs::RedisGcsClient> gcs_client)
     : object_manager_(object_manager),
       reconstruction_policy_(reconstruction_policy),
       io_service_(io_service),
       client_id_(client_id),
       initial_lease_period_ms_(initial_lease_period_ms),
-      task_lease_table_(task_lease_table) {}
+      gcs_client_(gcs_client) {}
 
 bool TaskDependencyManager::CheckObjectLocal(const ObjectID &object_id) const {
   return local_objects_.count(object_id) == 1;
@@ -383,10 +382,11 @@ void TaskDependencyManager::AcquireTaskLease(const TaskID &task_id) {
   }
 
   auto task_lease_data = std::make_shared<TaskLeaseData>();
+  task_lease_data->set_task_id(task_id.Binary());
   task_lease_data->set_node_manager_id(client_id_.Hex());
   task_lease_data->set_acquired_at(absl::GetCurrentTimeNanos() / 1000000);
   task_lease_data->set_timeout(it->second.lease_period);
-  RAY_CHECK_OK(task_lease_table_.Add(JobID::Nil(), task_id, task_lease_data, nullptr));
+  RAY_CHECK_OK(gcs_client_->Tasks().AsyncAddTaskLease(task_lease_data, nullptr));
 
   auto period = boost::posix_time::milliseconds(it->second.lease_period / 2);
   it->second.lease_timer->expires_from_now(period);

src/ray/raylet/task_dependency_manager.h

@@ -4,6 +4,7 @@
 // clang-format off
 #include "ray/common/id.h"
 #include "ray/common/task/task.h"
+#include "ray/gcs/redis_gcs_client.h"
 #include "ray/object_manager/object_manager.h"
 #include "ray/raylet/reconstruction_policy.h"
 // clang-format on
@@ -33,7 +34,7 @@ class TaskDependencyManager {
                         ReconstructionPolicyInterface &reconstruction_policy,
                         boost::asio::io_service &io_service, const ClientID &client_id,
                         int64_t initial_lease_period_ms,
-                        gcs::TableInterface<TaskID, TaskLeaseData> &task_lease_table);
+                        std::shared_ptr<gcs::RedisGcsClient> gcs_client);
 
   /// Check whether an object is locally available.
   ///
@@ -226,8 +227,8 @@ class TaskDependencyManager {
   /// added to the GCS. The lease expiration period is doubled every time the
   /// lease is renewed.
   const int64_t initial_lease_period_ms_;
-  /// The storage system for the task lease table.
-  gcs::TableInterface<TaskID, TaskLeaseData> &task_lease_table_;
+  /// A client connection to the GCS.
+  std::shared_ptr<gcs::RedisGcsClient> gcs_client_;
   /// A mapping from task ID of each subscribed task to its list of object
   /// dependencies, either task arguments or objects passed into `ray.get`.
   std::unordered_map<ray::TaskID, TaskDependencies> task_dependencies_;

src/ray/raylet/task_dependency_manager_test.cc

@@ -6,6 +6,8 @@
 #include <boost/asio.hpp>
 
 #include "ray/common/task/task_util.h"
+#include "ray/gcs/redis_accessor.h"
+#include "ray/gcs/redis_gcs_client.h"
 #include "ray/raylet/task_dependency_manager.h"
 #include "ray/util/test_util.h"
 
@@ -31,13 +33,23 @@ class MockReconstructionPolicy : public ReconstructionPolicyInterface {
   MOCK_METHOD1(Cancel, void(const ObjectID &object_id));
 };
 
-class MockGcs : public gcs::TableInterface<TaskID, TaskLeaseData> {
+class MockTaskInfoAccessor : public gcs::RedisTaskInfoAccessor {
  public:
-  MOCK_METHOD4(
-      Add,
-      ray::Status(const JobID &job_id, const TaskID &task_id,
-                  const std::shared_ptr<TaskLeaseData> &task_data,
-                  const gcs::TableInterface<TaskID, TaskLeaseData>::WriteCallback &done));
+  MockTaskInfoAccessor(gcs::RedisGcsClient *client)
+      : gcs::RedisTaskInfoAccessor(client) {}
+
+  MOCK_METHOD2(AsyncAddTaskLease,
+               ray::Status(const std::shared_ptr<TaskLeaseData> &data_ptr,
+                           const gcs::StatusCallback &callback));
+};
+
+class MockGcsClient : public gcs::RedisGcsClient {
+ public:
+  MockGcsClient(const gcs::GcsClientOptions &options) : gcs::RedisGcsClient(options) {}
+
+  void Init(MockTaskInfoAccessor *task_accessor_mock) {
+    task_accessor_.reset(task_accessor_mock);
+  }
 };
 
 class TaskDependencyManagerTest : public ::testing::Test {
@@ -46,11 +58,15 @@ class TaskDependencyManagerTest : public ::testing::Test {
       : object_manager_mock_(),
         reconstruction_policy_mock_(),
         io_service_(),
-        gcs_mock_(),
+        options_("", 1, ""),
+        gcs_client_mock_(new MockGcsClient(options_)),
+        task_accessor_mock_(new MockTaskInfoAccessor(gcs_client_mock_.get())),
         initial_lease_period_ms_(100),
         task_dependency_manager_(object_manager_mock_, reconstruction_policy_mock_,
                                  io_service_, ClientID::Nil(), initial_lease_period_ms_,
-                                 gcs_mock_) {}
+                                 gcs_client_mock_) {
+    gcs_client_mock_->Init(task_accessor_mock_);
+  }
 
   void Run(uint64_t timeout_ms) {
     auto timer_period = boost::posix_time::milliseconds(timeout_ms);
@@ -67,7 +83,9 @@ class TaskDependencyManagerTest : public ::testing::Test {
   MockObjectManager object_manager_mock_;
   MockReconstructionPolicy reconstruction_policy_mock_;
   boost::asio::io_service io_service_;
-  MockGcs gcs_mock_;
+  gcs::GcsClientOptions options_;
+  std::shared_ptr<MockGcsClient> gcs_client_mock_;
+  MockTaskInfoAccessor *task_accessor_mock_;
   int64_t initial_lease_period_ms_;
   TaskDependencyManager task_dependency_manager_;
 };
@@ -235,7 +253,7 @@ TEST_F(TaskDependencyManagerTest, TestTaskChain) {
 
     // Mark each task as pending. A lease entry should be added to the GCS for
     // each task.
-    EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _));
+    EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _));
     task_dependency_manager_.TaskPending(task);
 
     i++;
@@ -287,7 +305,7 @@ TEST_F(TaskDependencyManagerTest, TestDependentPut) {
   // it is pending execution.
   EXPECT_CALL(object_manager_mock_, CancelPull(put_id));
   EXPECT_CALL(reconstruction_policy_mock_, Cancel(put_id));
-  EXPECT_CALL(gcs_mock_, Add(_, task1.GetTaskSpecification().TaskId(), _, _));
+  EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _));
   task_dependency_manager_.TaskPending(task1);
 }
 
@@ -300,7 +318,7 @@ TEST_F(TaskDependencyManagerTest, TestTaskForwarding) {
     const auto &arguments = task.GetDependencies();
     static_cast<void>(task_dependency_manager_.SubscribeGetDependencies(
         task.GetTaskSpecification().TaskId(), arguments));
-    EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _));
+    EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _));
     task_dependency_manager_.TaskPending(task);
   }
 
@@ -404,7 +422,8 @@ TEST_F(TaskDependencyManagerTest, TestTaskLeaseRenewal) {
   // Mark a task as pending.
   auto task = ExampleTask({}, 0);
   // We expect an initial call to acquire the lease.
-  EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _));
+  EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _));
+
   task_dependency_manager_.TaskPending(task);
 
   // Check that while the task is still pending, there is one call to renew the
@@ -415,8 +434,7 @@ TEST_F(TaskDependencyManagerTest, TestTaskLeaseRenewal) {
   for (int i = 1; i <= num_expected_calls; i++) {
     sleep_time += i * initial_lease_period_ms_;
   }
-  EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _))
-      .Times(num_expected_calls);
+  EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _)).Times(num_expected_calls);
   Run(sleep_time);
 }
 
@@ -438,7 +456,7 @@ TEST_F(TaskDependencyManagerTest, TestRemoveTasksAndRelatedObjects) {
         task.GetTaskSpecification().TaskId(), arguments);
     // Mark each task as pending. A lease entry should be added to the GCS for
     // each task.
-    EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _));
+    EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _));
     task_dependency_manager_.TaskPending(task);
   }