[GCS]access task reconstruction in TaskInfoAccessor (#6688)

* add task lease interface to TaskInfoAccessor

* impl of task lease

* support accessing task lease in TaskInfoAccessor

* update raylet usage of task lease

* add comment

* fix lint

* fix UT of TaskDependencyManager

* fix UT of ReconstructionPolicy

* rm useless code from UT

* add task reconstruction methods to gcs

* fix ut of RedisGcsClient

* update test

* update comments

src/ray/gcs/accessor.h

@@ -232,6 +232,16 @@ class TaskInfoAccessor {
   virtual Status AsyncUnsubscribeTaskLease(const TaskID &task_id,
                                            const StatusCallback &done) = 0;
 
+  /// Attempt task reconstruction to GCS asynchronously.
+  ///
+  /// \param data_ptr The task reconstruction that will be added to GCS.
+  /// \param callback Callback that will be called after task reconstruction
+  /// has been added to GCS.
+  /// \return Status
+  virtual Status AttemptTaskReconstruction(
+      const std::shared_ptr<rpc::TaskReconstructionData> &data_ptr,
+      const StatusCallback &callback) = 0;
+
  protected:
   TaskInfoAccessor() = default;
 };

src/ray/gcs/redis_accessor.cc

@@ -331,6 +331,30 @@ Status RedisTaskInfoAccessor::AsyncUnsubscribeTaskLease(const TaskID &task_id,
   return task_lease_sub_executor_.AsyncUnsubscribe(subscribe_id_, task_id, done);
 }
 
+Status RedisTaskInfoAccessor::AttemptTaskReconstruction(
+    const std::shared_ptr<TaskReconstructionData> &data_ptr,
+    const StatusCallback &callback) {
+  TaskReconstructionLog::WriteCallback on_success = nullptr;
+  TaskReconstructionLog::WriteCallback on_failure = nullptr;
+  if (callback != nullptr) {
+    on_success = [callback](RedisGcsClient *client, const TaskID &id,
+                            const TaskReconstructionData &data) {
+      callback(Status::OK());
+    };
+    on_failure = [callback](RedisGcsClient *client, const TaskID &id,
+                            const TaskReconstructionData &data) {
+      callback(Status::Invalid("Updating task reconstruction failed."));
+    };
+  }
+
+  TaskID task_id = TaskID::FromBinary(data_ptr->task_id());
+  int reconstruction_attempt = data_ptr->num_reconstructions();
+  TaskReconstructionLog &task_reconstruction_log =
+      client_impl_->task_reconstruction_log();
+  return task_reconstruction_log.AppendAt(JobID::Nil(), task_id, data_ptr, on_success,
+                                          on_failure, reconstruction_attempt);
+}
+
 RedisObjectInfoAccessor::RedisObjectInfoAccessor(RedisGcsClient *client_impl)
     : client_impl_(client_impl), object_sub_executor_(client_impl->object_table()) {}
 

src/ray/gcs/redis_accessor.h

@@ -150,6 +150,10 @@ class RedisTaskInfoAccessor : public TaskInfoAccessor {
   Status AsyncUnsubscribeTaskLease(const TaskID &task_id,
                                    const StatusCallback &done) override;
 
+  Status AttemptTaskReconstruction(
+      const std::shared_ptr<TaskReconstructionData> &data_ptr,
+      const StatusCallback &callback) override;
+
  private:
   RedisGcsClient *client_impl_{nullptr};
   // Use a random ClientID for task subscription. Because:

src/ray/gcs/redis_gcs_client.h

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

src/ray/gcs/test/redis_gcs_client_test.cc

@@ -364,109 +364,115 @@ class TaskTableTestHelper {
 
 TEST_TASK_TABLE_MACRO(TestGcsWithAsio, TestTableLookup);
 
-void TestLogLookup(const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> client) {
-  // Append some entries to the log at an object ID.
-  TaskID task_id = RandomTaskId();
-  std::vector<std::string> node_manager_ids = {"abc", "def", "ghi"};
-  for (auto &node_manager_id : node_manager_ids) {
-    auto data = std::make_shared<TaskReconstructionData>();
-    data->set_node_manager_id(node_manager_id);
-    // Check that we added the correct object entries.
-    auto add_callback = [task_id, data](gcs::RedisGcsClient *client, const TaskID &id,
-                                        const TaskReconstructionData &d) {
+class LogLookupTestHelper {
+ public:
+  static void TestLogLookup(const JobID &job_id,
+                            std::shared_ptr<gcs::RedisGcsClient> client) {
+    // Append some entries to the log at an object ID.
+    TaskID task_id = RandomTaskId();
+    std::vector<std::string> node_manager_ids = {"abc", "def", "ghi"};
+    for (auto &node_manager_id : node_manager_ids) {
+      auto data = std::make_shared<TaskReconstructionData>();
+      data->set_node_manager_id(node_manager_id);
+      // Check that we added the correct object entries.
+      auto add_callback = [task_id, data](gcs::RedisGcsClient *client, const TaskID &id,
+                                          const TaskReconstructionData &d) {
+        ASSERT_EQ(id, task_id);
+        ASSERT_EQ(data->node_manager_id(), d.node_manager_id());
+      };
+      RAY_CHECK_OK(
+          client->task_reconstruction_log().Append(job_id, task_id, data, add_callback));
+    }
+
+    // Check that lookup returns the added object entries.
+    auto lookup_callback = [task_id, node_manager_ids](
+                               gcs::RedisGcsClient *client, const TaskID &id,
+                               const std::vector<TaskReconstructionData> &data) {
       ASSERT_EQ(id, task_id);
-      ASSERT_EQ(data->node_manager_id(), d.node_manager_id());
+      for (const auto &entry : data) {
+        ASSERT_EQ(entry.node_manager_id(), node_manager_ids[test->NumCallbacks()]);
+        test->IncrementNumCallbacks();
+      }
+      if (test->NumCallbacks() == node_manager_ids.size()) {
+        test->Stop();
+      }
     };
+
+    // Do a lookup at the object ID.
     RAY_CHECK_OK(
-        client->task_reconstruction_log().Append(job_id, task_id, data, add_callback));
+        client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback));
+    // 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(), node_manager_ids.size());
   }
 
-  // Check that lookup returns the added object entries.
-  auto lookup_callback = [task_id, node_manager_ids](
-                             gcs::RedisGcsClient *client, const TaskID &id,
-                             const std::vector<TaskReconstructionData> &data) {
-    ASSERT_EQ(id, task_id);
-    for (const auto &entry : data) {
-      ASSERT_EQ(entry.node_manager_id(), node_manager_ids[test->NumCallbacks()]);
-      test->IncrementNumCallbacks();
+  static void TestLogAppendAt(const JobID &job_id,
+                              std::shared_ptr<gcs::RedisGcsClient> client) {
+    TaskID task_id = RandomTaskId();
+    std::vector<std::string> node_manager_ids = {"A", "B"};
+    std::vector<std::shared_ptr<TaskReconstructionData>> data_log;
+    for (const auto &node_manager_id : node_manager_ids) {
+      auto data = std::make_shared<TaskReconstructionData>();
+      data->set_node_manager_id(node_manager_id);
+      data_log.push_back(data);
     }
-    if (test->NumCallbacks() == node_manager_ids.size()) {
-      test->Stop();
-    }
-  };
 
-  // Do a lookup at the object ID.
-  RAY_CHECK_OK(
-      client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback));
-  // 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(), node_manager_ids.size());
-}
+    // Check that we added the correct task.
+    auto failure_callback = [task_id](gcs::RedisGcsClient *client, const TaskID &id,
+                                      const TaskReconstructionData &d) {
+      ASSERT_EQ(id, task_id);
+      test->IncrementNumCallbacks();
+    };
+
+    // Will succeed.
+    RAY_CHECK_OK(client->task_reconstruction_log().Append(job_id, task_id,
+                                                          data_log.front(),
+                                                          /*done callback=*/nullptr));
+    // Append at index 0 will fail.
+    RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(
+        job_id, task_id, data_log[1],
+        /*done callback=*/nullptr, failure_callback, /*log_length=*/0));
+
+    // Append at index 2 will fail.
+    RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(
+        job_id, task_id, data_log[1],
+        /*done callback=*/nullptr, failure_callback, /*log_length=*/2));
+
+    // Append at index 1 will succeed.
+    RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(
+        job_id, task_id, data_log[1],
+        /*done callback=*/nullptr, failure_callback, /*log_length=*/1));
+
+    auto lookup_callback = [node_manager_ids](
+                               gcs::RedisGcsClient *client, const TaskID &id,
+                               const std::vector<TaskReconstructionData> &data) {
+      std::vector<std::string> appended_managers;
+      for (const auto &entry : data) {
+        appended_managers.push_back(entry.node_manager_id());
+      }
+      ASSERT_EQ(appended_managers, node_manager_ids);
+      test->Stop();
+    };
+    RAY_CHECK_OK(
+        client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback));
+    // 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(), 2);
+  }
+};
 
 TEST_F(TestGcsWithAsio, TestLogLookup) {
   test = this;
-  TestLogLookup(job_id_, client_);
+  LogLookupTestHelper::TestLogLookup(job_id_, client_);
 }
 
 TEST_TASK_TABLE_MACRO(TestGcsWithAsio, TestTableLookupFailure);
 
-void TestLogAppendAt(const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> client) {
-  TaskID task_id = RandomTaskId();
-  std::vector<std::string> node_manager_ids = {"A", "B"};
-  std::vector<std::shared_ptr<TaskReconstructionData>> data_log;
-  for (const auto &node_manager_id : node_manager_ids) {
-    auto data = std::make_shared<TaskReconstructionData>();
-    data->set_node_manager_id(node_manager_id);
-    data_log.push_back(data);
-  }
-
-  // Check that we added the correct task.
-  auto failure_callback = [task_id](gcs::RedisGcsClient *client, const TaskID &id,
-                                    const TaskReconstructionData &d) {
-    ASSERT_EQ(id, task_id);
-    test->IncrementNumCallbacks();
-  };
-
-  // Will succeed.
-  RAY_CHECK_OK(client->task_reconstruction_log().Append(job_id, task_id, data_log.front(),
-                                                        /*done callback=*/nullptr));
-  // Append at index 0 will fail.
-  RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(
-      job_id, task_id, data_log[1],
-      /*done callback=*/nullptr, failure_callback, /*log_length=*/0));
-
-  // Append at index 2 will fail.
-  RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(
-      job_id, task_id, data_log[1],
-      /*done callback=*/nullptr, failure_callback, /*log_length=*/2));
-
-  // Append at index 1 will succeed.
-  RAY_CHECK_OK(client->task_reconstruction_log().AppendAt(
-      job_id, task_id, data_log[1],
-      /*done callback=*/nullptr, failure_callback, /*log_length=*/1));
-
-  auto lookup_callback = [node_manager_ids](
-                             gcs::RedisGcsClient *client, const TaskID &id,
-                             const std::vector<TaskReconstructionData> &data) {
-    std::vector<std::string> appended_managers;
-    for (const auto &entry : data) {
-      appended_managers.push_back(entry.node_manager_id());
-    }
-    ASSERT_EQ(appended_managers, node_manager_ids);
-    test->Stop();
-  };
-  RAY_CHECK_OK(
-      client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback));
-  // 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(), 2);
-}
-
 TEST_F(TestGcsWithAsio, TestLogAppendAt) {
   test = this;
-  TestLogAppendAt(job_id_, client_);
+  LogLookupTestHelper::TestLogAppendAt(job_id_, client_);
 }
 
 class SetTestHelper {
@@ -806,52 +812,55 @@ TEST_F(TestGcsWithAsio, TestSet) {
   SetTestHelper::TestSet(job_id_, client_);
 }
 
-void TestDeleteKeysFromLog(
-    const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> client,
-    std::vector<std::shared_ptr<TaskReconstructionData>> &data_vector) {
-  std::vector<TaskID> ids;
-  TaskID task_id;
-  for (auto &data : data_vector) {
-    task_id = RandomTaskId();
-    ids.push_back(task_id);
-    // Check that we added the correct object entries.
-    auto add_callback = [task_id, data](gcs::RedisGcsClient *client, const TaskID &id,
-                                        const TaskReconstructionData &d) {
-      ASSERT_EQ(id, task_id);
-      ASSERT_EQ(data->node_manager_id(), d.node_manager_id());
-      test->IncrementNumCallbacks();
-    };
-    RAY_CHECK_OK(
-        client->task_reconstruction_log().Append(job_id, task_id, data, add_callback));
+class LogDeleteTestHelper {
+ public:
+  static void TestDeleteKeysFromLog(
+      const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> client,
+      std::vector<std::shared_ptr<TaskReconstructionData>> &data_vector) {
+    std::vector<TaskID> ids;
+    TaskID task_id;
+    for (auto &data : data_vector) {
+      task_id = RandomTaskId();
+      ids.push_back(task_id);
+      // Check that we added the correct object entries.
+      auto add_callback = [task_id, data](gcs::RedisGcsClient *client, const TaskID &id,
+                                          const TaskReconstructionData &d) {
+        ASSERT_EQ(id, task_id);
+        ASSERT_EQ(data->node_manager_id(), d.node_manager_id());
+        test->IncrementNumCallbacks();
+      };
+      RAY_CHECK_OK(
+          client->task_reconstruction_log().Append(job_id, task_id, data, add_callback));
+    }
+    for (const auto &task_id : ids) {
+      // Check that lookup returns the added object entries.
+      auto lookup_callback = [task_id, data_vector](
+                                 gcs::RedisGcsClient *client, const TaskID &id,
+                                 const std::vector<TaskReconstructionData> &data) {
+        ASSERT_EQ(id, task_id);
+        ASSERT_EQ(data.size(), 1);
+        test->IncrementNumCallbacks();
+      };
+      RAY_CHECK_OK(
+          client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback));
+    }
+    if (ids.size() == 1) {
+      client->task_reconstruction_log().Delete(job_id, ids[0]);
+    } else {
+      client->task_reconstruction_log().Delete(job_id, ids);
+    }
+    for (const auto &task_id : ids) {
+      auto lookup_callback = [task_id](gcs::RedisGcsClient *client, const TaskID &id,
+                                       const std::vector<TaskReconstructionData> &data) {
+        ASSERT_EQ(id, task_id);
+        ASSERT_TRUE(data.size() == 0);
+        test->IncrementNumCallbacks();
+      };
+      RAY_CHECK_OK(
+          client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback));
+    }
   }
-  for (const auto &task_id : ids) {
-    // Check that lookup returns the added object entries.
-    auto lookup_callback = [task_id, data_vector](
-                               gcs::RedisGcsClient *client, const TaskID &id,
-                               const std::vector<TaskReconstructionData> &data) {
-      ASSERT_EQ(id, task_id);
-      ASSERT_EQ(data.size(), 1);
-      test->IncrementNumCallbacks();
-    };
-    RAY_CHECK_OK(
-        client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback));
-  }
-  if (ids.size() == 1) {
-    client->task_reconstruction_log().Delete(job_id, ids[0]);
-  } else {
-    client->task_reconstruction_log().Delete(job_id, ids);
-  }
-  for (const auto &task_id : ids) {
-    auto lookup_callback = [task_id](gcs::RedisGcsClient *client, const TaskID &id,
-                                     const std::vector<TaskReconstructionData> &data) {
-      ASSERT_EQ(id, task_id);
-      ASSERT_TRUE(data.size() == 0);
-      test->IncrementNumCallbacks();
-    };
-    RAY_CHECK_OK(
-        client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback));
-  }
-}
+};
 
 // Test delete function for keys of Log or Table.
 void TestDeleteKeys(const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> client) {
@@ -867,7 +876,7 @@ void TestDeleteKeys(const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> cl
   // Test one element case.
   AppendTaskReconstructionData(1);
   ASSERT_EQ(task_reconstruction_vector.size(), 1);
-  TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector);
+  LogDeleteTestHelper::TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector);
   // Test the case for more than one elements and less than
   // maximum_gcs_deletion_batch_size.
   AppendTaskReconstructionData(RayConfig::instance().maximum_gcs_deletion_batch_size() /
@@ -875,14 +884,14 @@ void TestDeleteKeys(const JobID &job_id, std::shared_ptr<gcs::RedisGcsClient> cl
   ASSERT_GT(task_reconstruction_vector.size(), 1);
   ASSERT_LT(task_reconstruction_vector.size(),
             RayConfig::instance().maximum_gcs_deletion_batch_size());
-  TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector);
+  LogDeleteTestHelper::TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector);
   // Test the case for more than maximum_gcs_deletion_batch_size.
   // The Delete function will split the data into two commands.
   AppendTaskReconstructionData(RayConfig::instance().maximum_gcs_deletion_batch_size() /
                                2);
   ASSERT_GT(task_reconstruction_vector.size(),
             RayConfig::instance().maximum_gcs_deletion_batch_size());
-  TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector);
+  LogDeleteTestHelper::TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector);
 
   // Test delete function for keys of Table.
   std::vector<std::shared_ptr<TaskTableData>> task_vector;

src/ray/protobuf/gcs.proto

@@ -68,10 +68,12 @@ message ObjectTableData {
 }
 
 message TaskReconstructionData {
+  // The ID of task.
+  bytes task_id = 1;
   // The number of times this task has been reconstructed so far.
-  uint64 num_reconstructions = 1;
+  uint64 num_reconstructions = 2;
   // The node manager that is trying to reconstruct the task.
-  bytes node_manager_id = 2;
+  bytes node_manager_id = 3;
 }
 
 message TaskTableData {

src/ray/raylet/node_manager.cc

@@ -98,8 +98,7 @@ 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_, object_directory_,
-          gcs_client_->task_reconstruction_log()),
+          self_node_id_, gcs_client_, object_directory_),
       task_dependency_manager_(
           object_manager, reconstruction_policy_, io_service, self_node_id_,
           RayConfig::instance().initial_reconstruction_timeout_milliseconds(),

src/ray/raylet/reconstruction_policy.cc

@@ -11,15 +11,13 @@ ReconstructionPolicy::ReconstructionPolicy(
     std::function<void(const TaskID &, const ObjectID &)> reconstruction_handler,
     int64_t initial_reconstruction_timeout_ms, const ClientID &client_id,
     std::shared_ptr<gcs::RedisGcsClient> gcs_client,
-    std::shared_ptr<ObjectDirectoryInterface> object_directory,
-    gcs::LogInterface<TaskID, TaskReconstructionData> &task_reconstruction_log)
+    std::shared_ptr<ObjectDirectoryInterface> object_directory)
     : io_service_(io_service),
       reconstruction_handler_(reconstruction_handler),
       initial_reconstruction_timeout_ms_(initial_reconstruction_timeout_ms),
       client_id_(client_id),
       gcs_client_(gcs_client),
-      object_directory_(std::move(object_directory)),
-      task_reconstruction_log_(task_reconstruction_log) {}
+      object_directory_(std::move(object_directory)) {}
 
 void ReconstructionPolicy::SetTaskTimeout(
     std::unordered_map<TaskID, ReconstructionTask>::iterator task_it,
@@ -134,21 +132,19 @@ void ReconstructionPolicy::AttemptReconstruction(const TaskID &task_id,
   // reconstruction log. This will fail if another node has already inserted
   // an entry for this reconstruction.
   auto reconstruction_entry = std::make_shared<TaskReconstructionData>();
+  reconstruction_entry->set_task_id(task_id.Binary());
   reconstruction_entry->set_num_reconstructions(reconstruction_attempt);
   reconstruction_entry->set_node_manager_id(client_id_.Binary());
-  RAY_CHECK_OK(task_reconstruction_log_.AppendAt(
-      JobID::Nil(), task_id, reconstruction_entry,
-      /*success_callback=*/
-      [this, required_object_id](gcs::RedisGcsClient *client, const TaskID &task_id,
-                                 const TaskReconstructionData &data) {
-        HandleReconstructionLogAppend(task_id, required_object_id, /*success=*/true);
-      },
-      /*failure_callback=*/
-      [this, required_object_id](gcs::RedisGcsClient *client, const TaskID &task_id,
-                                 const TaskReconstructionData &data) {
-        HandleReconstructionLogAppend(task_id, required_object_id, /*success=*/false);
-      },
-      reconstruction_attempt));
+  RAY_CHECK_OK(gcs_client_->Tasks().AttemptTaskReconstruction(
+      reconstruction_entry,
+      /*done=*/
+      [this, task_id, required_object_id](Status status) {
+        if (status.ok()) {
+          HandleReconstructionLogAppend(task_id, required_object_id, /*success=*/true);
+        } else {
+          HandleReconstructionLogAppend(task_id, required_object_id, /*success=*/false);
+        }
+      }));
 
   // Increment the number of times reconstruction has been attempted. This is
   // used to suppress duplicate reconstructions of the same task. If

src/ray/raylet/reconstruction_policy.h

@@ -44,8 +44,7 @@ class ReconstructionPolicy : public ReconstructionPolicyInterface {
       std::function<void(const TaskID &, const ObjectID &)> reconstruction_handler,
       int64_t initial_reconstruction_timeout_ms, const ClientID &client_id,
       std::shared_ptr<gcs::RedisGcsClient> gcs_client,
-      std::shared_ptr<ObjectDirectoryInterface> object_directory,
-      gcs::LogInterface<TaskID, TaskReconstructionData> &task_reconstruction_log);
+      std::shared_ptr<ObjectDirectoryInterface> object_directory);
 
   /// Listen for task lease notifications about an object that may require
   /// reconstruction. If no notifications are received within the initial
@@ -146,7 +145,6 @@ class ReconstructionPolicy : public ReconstructionPolicyInterface {
   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_;
   /// The tasks that we are currently subscribed to in the GCS.
   std::unordered_map<TaskID, ReconstructionTask> listening_tasks_;
 };

src/ray/raylet/reconstruction_policy_test.cc

@@ -131,14 +131,33 @@ class MockTaskInfoAccessor : public gcs::RedisTaskInfoAccessor {
     return Status::OK();
   }
 
+  Status AttemptTaskReconstruction(
+      const std::shared_ptr<TaskReconstructionData> &task_data,
+      const gcs::StatusCallback &done) override {
+    int log_index = task_data->num_reconstructions();
+    TaskID task_id = TaskID::FromBinary(task_data->task_id());
+    if (task_reconstruction_log_[task_id].size() == static_cast<size_t>(log_index)) {
+      task_reconstruction_log_[task_id].push_back(*task_data);
+      if (done != nullptr) {
+        done(Status::OK());
+      }
+    } else {
+      if (done != nullptr) {
+        done(Status::Invalid("Updating task reconstruction failed."));
+      }
+    }
+    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_;
+  std::unordered_map<TaskID, std::vector<TaskReconstructionData>>
+      task_reconstruction_log_;
 };
 
-class MockGcs : public gcs::RedisGcsClient,
-                public ray::gcs::LogInterface<TaskID, TaskReconstructionData> {
+class MockGcs : public gcs::RedisGcsClient {
  public:
   MockGcs() : gcs::RedisGcsClient(gcs::GcsClientOptions("", 0, "")){};
 
@@ -146,37 +165,6 @@ class MockGcs : public gcs::RedisGcsClient,
     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,
-      const ray::gcs::LogInterface<TaskID, TaskReconstructionData>::WriteCallback
-          &success_callback,
-      const ray::gcs::LogInterface<TaskID, TaskReconstructionData>::WriteCallback
-          &failure_callback,
-      int log_index) {
-    if (task_reconstruction_log_[task_id].size() == static_cast<size_t>(log_index)) {
-      task_reconstruction_log_[task_id].push_back(*task_data);
-      if (success_callback != nullptr) {
-        success_callback(nullptr, task_id, *task_data);
-      }
-    } else {
-      if (failure_callback != nullptr) {
-        failure_callback(nullptr, task_id, *task_data);
-      }
-    }
-    return Status::OK();
-  }
-
-  MOCK_METHOD4(
-      Append,
-      ray::Status(
-          const JobID &, const TaskID &, const std::shared_ptr<TaskReconstructionData> &,
-          const ray::gcs::LogInterface<TaskID, TaskReconstructionData>::WriteCallback &));
-
- private:
-  std::unordered_map<TaskID, std::vector<TaskReconstructionData>>
-      task_reconstruction_log_;
 };
 
 class ReconstructionPolicyTest : public ::testing::Test {
@@ -194,7 +182,7 @@ class ReconstructionPolicyTest : public ::testing::Test {
               TriggerReconstruction(task_id);
             },
             reconstruction_timeout_ms_, ClientID::FromRandom(), mock_gcs_,
-            mock_object_directory_, *mock_gcs_)),
+            mock_object_directory_)),
         timer_canceled_(false) {
     subscribe_callback_ = [this](const TaskID &task_id,
                                  const boost::optional<TaskLeaseData> &task_lease) {
@@ -449,14 +437,13 @@ TEST_F(ReconstructionPolicyTest, TestSimultaneousReconstructionSuppressed) {
   // reconstruction first. This should suppress this node's first attempt at
   // reconstruction.
   auto task_reconstruction_data = std::make_shared<TaskReconstructionData>();
+  task_reconstruction_data->set_task_id(task_id.Binary());
   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));
+  RAY_CHECK_OK(mock_gcs_->Tasks().AttemptTaskReconstruction(
+      task_reconstruction_data,
+      /*done=*/
+      [](Status status) { ASSERT_TRUE(status.ok()); }));
 
   // Listen for an object.
   reconstruction_policy_->ListenAndMaybeReconstruct(object_id);