[GCS]Optimize gcs client testcases (#7895)

src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 #include "ray/gcs/gcs_client/service_based_gcs_client.h"
-
 #include "gtest/gtest.h"
 #include "ray/common/test_util.h"
 #include "ray/gcs/gcs_client/service_based_accessor.h"
@@ -22,12 +21,8 @@
 
 namespace ray {
 
-static std::string redis_server_executable;
-static std::string redis_client_executable;
-static std::string libray_redis_module_path;
-
-class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
- public:
+class ServiceBasedGcsClientTest : public RedisServiceManagerForTest {
+ protected:
   void SetUp() override {
     config.grpc_server_port = 0;
     config.grpc_server_name = "MockedGcsServer";
@@ -51,7 +46,7 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
       std::this_thread::sleep_for(std::chrono::milliseconds(10));
     }
 
-    // Create gcs client
+    // Create GCS client.
     gcs::GcsClientOptions options(config.redis_address, config.redis_port,
                                   config.redis_password, config.is_test);
     gcs_client_.reset(new gcs::ServiceBasedGcsClient(options));
@@ -67,6 +62,14 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
     FlushAll();
   }
 
+  bool SubscribeToFinishedJobs(
+      const gcs::SubscribeCallback<JobID, rpc::JobTableData> &subscribe) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Jobs().AsyncSubscribeToFinishedJobs(
+        subscribe, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   bool AddJob(const std::shared_ptr<rpc::JobTableData> &job_table_data) {
     std::promise<bool> promise;
     RAY_CHECK_OK(gcs_client_->Jobs().AsyncAdd(
@@ -81,6 +84,31 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
+  bool SubscribeActor(
+      const ActorID &actor_id,
+      const gcs::SubscribeCallback<ActorID, rpc::ActorTableData> &subscribe) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Actors().AsyncSubscribe(
+        actor_id, subscribe,
+        [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool UnsubscribeActor(const ActorID &actor_id) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Actors().AsyncUnsubscribe(
+        actor_id, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool SubscribeAllActors(
+      const gcs::SubscribeCallback<ActorID, rpc::ActorTableData> &subscribe) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Actors().AsyncSubscribeAll(
+        subscribe, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   bool RegisterActor(const std::shared_ptr<rpc::ActorTableData> &actor_table_data) {
     std::promise<bool> promise;
     RAY_CHECK_OK(gcs_client_->Actors().AsyncRegister(
@@ -151,6 +179,14 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
     return actor_checkpoint_id_data;
   }
 
+  bool SubscribeToNodeChange(
+      const gcs::SubscribeCallback<ClientID, rpc::GcsNodeInfo> &subscribe) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Nodes().AsyncSubscribeToNodeChange(
+        subscribe, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   bool RegisterSelf(const rpc::GcsNodeInfo &local_node_info) {
     Status status = gcs_client_->Nodes().RegisterSelf(local_node_info);
     return status.ok();
@@ -188,6 +224,15 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
+  bool SubscribeToResources(
+      const gcs::SubscribeCallback<ClientID, gcs::ResourceChangeNotification>
+          &subscribe) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Nodes().AsyncSubscribeToResources(
+        subscribe, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   gcs::NodeInfoAccessor::ResourceMap GetResources(const ClientID &node_id) {
     gcs::NodeInfoAccessor::ResourceMap resource_map;
     std::promise<bool> promise;
@@ -222,6 +267,14 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
+  bool SubscribeBatchHeartbeat(
+      const gcs::ItemCallback<rpc::HeartbeatBatchTableData> &subscribe) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Nodes().AsyncSubscribeBatchHeartbeat(
+        subscribe, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   bool ReportHeartbeat(const std::shared_ptr<rpc::HeartbeatTableData> heartbeat) {
     std::promise<bool> promise;
     RAY_CHECK_OK(gcs_client_->Nodes().AsyncReportHeartbeat(
@@ -229,6 +282,23 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
+  bool SubscribeTask(
+      const TaskID &task_id,
+      const gcs::SubscribeCallback<TaskID, rpc::TaskTableData> &subscribe) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Tasks().AsyncSubscribe(
+        task_id, subscribe,
+        [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool UnsubscribeTask(const TaskID &task_id) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Tasks().AsyncUnsubscribe(
+        task_id, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   bool AddTask(const std::shared_ptr<rpc::TaskTableData> task) {
     std::promise<bool> promise;
     RAY_CHECK_OK(gcs_client_->Tasks().AsyncAdd(
@@ -258,6 +328,24 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
+  bool SubscribeTaskLease(
+      const TaskID &task_id,
+      const gcs::SubscribeCallback<TaskID, boost::optional<rpc::TaskLeaseData>>
+          &subscribe) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Tasks().AsyncSubscribeTaskLease(
+        task_id, subscribe,
+        [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool UnsubscribeTaskLease(const TaskID &task_id) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Tasks().AsyncUnsubscribeTaskLease(
+        task_id, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   bool AddTaskLease(const std::shared_ptr<rpc::TaskLeaseData> task_lease) {
     std::promise<bool> promise;
     RAY_CHECK_OK(gcs_client_->Tasks().AsyncAddTaskLease(
@@ -274,14 +362,90 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
- protected:
-  bool WaitReady(const std::future<bool> &future,
-                 const std::chrono::milliseconds &timeout_ms) {
-    auto status = future.wait_for(timeout_ms);
-    return status == std::future_status::ready;
+  bool SubscribeToLocations(
+      const ObjectID &object_id,
+      const gcs::SubscribeCallback<ObjectID, gcs::ObjectChangeNotification> &subscribe) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Objects().AsyncSubscribeToLocations(
+        object_id, subscribe,
+        [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
   }
 
-  void WaitPendingDone(int &current_count, int expected_count) {
+  bool UnsubscribeToLocations(const ObjectID &object_id) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Objects().AsyncUnsubscribeToLocations(
+        object_id, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool AddLocation(const ObjectID &object_id, const ClientID &node_id) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Objects().AsyncAddLocation(
+        object_id, node_id,
+        [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool RemoveLocation(const ObjectID &object_id, const ClientID &node_id) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Objects().AsyncRemoveLocation(
+        object_id, node_id,
+        [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  std::vector<rpc::ObjectTableData> GetLocations(const ObjectID &object_id) {
+    std::promise<bool> promise;
+    std::vector<rpc::ObjectTableData> locations;
+    RAY_CHECK_OK(gcs_client_->Objects().AsyncGetLocations(
+        object_id, [&locations, &promise](
+                       Status status, const std::vector<rpc::ObjectTableData> &result) {
+          locations = result;
+          promise.set_value(status.ok());
+        }));
+    EXPECT_TRUE(WaitReady(promise.get_future(), timeout_ms_));
+    return locations;
+  }
+
+  bool AddProfileData(const std::shared_ptr<rpc::ProfileTableData> &profile_table_data) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Stats().AsyncAddProfileData(
+        profile_table_data,
+        [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool ReportJobError(const std::shared_ptr<rpc::ErrorTableData> &error_table_data) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Errors().AsyncReportJobError(
+        error_table_data, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool SubscribeToWorkerFailures(
+      const gcs::SubscribeCallback<WorkerID, rpc::WorkerFailureData> &subscribe) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Workers().AsyncSubscribeToWorkerFailures(
+        subscribe, [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool ReportWorkerFailure(
+      const std::shared_ptr<rpc::WorkerFailureData> &worker_failure_data) {
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Workers().AsyncReportWorkerFailure(
+        worker_failure_data,
+        [&promise](Status status) { promise.set_value(status.ok()); }));
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool WaitReady(std::future<bool> future, const std::chrono::milliseconds &timeout_ms) {
+    auto status = future.wait_for(timeout_ms);
+    return status == std::future_status::ready && future.get();
+  }
+
+  void WaitPendingDone(std::atomic<int> &current_count, int expected_count) {
     auto condition = [&current_count, expected_count]() {
       return current_count == expected_count;
     };
@@ -338,87 +502,89 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
     return task_lease_data;
   }
 
-  // Gcs server
+  std::shared_ptr<rpc::ProfileTableData> GenProfileTableData(const ClientID &node_id) {
+    auto profile_table_data = std::make_shared<rpc::ProfileTableData>();
+    profile_table_data->set_component_id(node_id.Binary());
+    return profile_table_data;
+  }
+
+  std::shared_ptr<rpc::ErrorTableData> GenErrorTableData(const JobID &job_id) {
+    auto error_table_data = std::make_shared<rpc::ErrorTableData>();
+    error_table_data->set_job_id(job_id.Binary());
+    return error_table_data;
+  }
+
+  std::shared_ptr<rpc::WorkerFailureData> GenWorkerFailureData() {
+    auto worker_failure_data = std::make_shared<rpc::WorkerFailureData>();
+    worker_failure_data->set_timestamp(std::time(nullptr));
+    return worker_failure_data;
+  }
+
+  // GCS server.
   gcs::GcsServerConfig config;
   std::unique_ptr<gcs::GcsServer> gcs_server_;
   std::unique_ptr<std::thread> thread_io_service_;
   std::unique_ptr<std::thread> thread_gcs_server_;
   std::unique_ptr<boost::asio::io_service> io_service_;
 
-  // Gcs client
+  // GCS client.
   std::unique_ptr<gcs::GcsClient> gcs_client_;
 
-  // Timeout waiting for gcs server reply, default is 2s
+  // Timeout waiting for GCS server reply, default is 2s.
   const std::chrono::milliseconds timeout_ms_{2000};
 };
 
-TEST_F(ServiceBasedGcsGcsClientTest, TestJobInfo) {
-  // Create job_table_data
+TEST_F(ServiceBasedGcsClientTest, TestJobInfo) {
+  // Create job table data.
   JobID add_job_id = JobID::FromInt(1);
   auto job_table_data = GenJobTableData(add_job_id);
 
-  std::promise<bool> promise;
-  auto on_subscribe = [&promise, add_job_id](const JobID &job_id,
-                                             const gcs::JobTableData &data) {
-    ASSERT_TRUE(add_job_id == job_id);
-    promise.set_value(true);
+  // Subscribe to finished jobs.
+  std::atomic<int> finished_job_count(0);
+  auto on_subscribe = [&finished_job_count](const JobID &job_id,
+                                            const gcs::JobTableData &data) {
+    finished_job_count++;
   };
-  RAY_CHECK_OK(gcs_client_->Jobs().AsyncSubscribeToFinishedJobs(
-      on_subscribe, [](Status status) { RAY_CHECK_OK(status); }));
+  ASSERT_TRUE(SubscribeToFinishedJobs(on_subscribe));
 
   ASSERT_TRUE(AddJob(job_table_data));
   ASSERT_TRUE(MarkJobFinished(add_job_id));
-  ASSERT_TRUE(WaitReady(promise.get_future(), timeout_ms_));
+  WaitPendingDone(finished_job_count, 1);
 }
 
-TEST_F(ServiceBasedGcsGcsClientTest, TestActorInfo) {
-  // Create actor_table_data
+TEST_F(ServiceBasedGcsClientTest, TestActorInfo) {
+  // Create actor table data.
   JobID job_id = JobID::FromInt(1);
   auto actor_table_data = GenActorTableData(job_id);
   ActorID actor_id = ActorID::FromBinary(actor_table_data->actor_id());
 
-  // Subscribe
-  std::promise<bool> promise_subscribe;
-  std::atomic<int> subscribe_callback_count(0);
-  auto on_subscribe = [&subscribe_callback_count](const ActorID &actor_id,
-                                                  const gcs::ActorTableData &data) {
-    ++subscribe_callback_count;
+  // Subscribe to any update operations of an actor.
+  std::atomic<int> actor_update_count(0);
+  auto on_subscribe = [&actor_update_count](const ActorID &actor_id,
+                                            const gcs::ActorTableData &data) {
+    ++actor_update_count;
   };
-  RAY_CHECK_OK(gcs_client_->Actors().AsyncSubscribe(actor_id, on_subscribe,
-                                                    [&promise_subscribe](Status status) {
-                                                      RAY_CHECK_OK(status);
-                                                      promise_subscribe.set_value(true);
-                                                    }));
+  ASSERT_TRUE(SubscribeActor(actor_id, on_subscribe));
 
-  // Register actor
+  // Register an actor to GCS.
   ASSERT_TRUE(RegisterActor(actor_table_data));
   ASSERT_TRUE(GetActor(actor_id).state() ==
               rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE);
 
-  // Unsubscribe
-  std::promise<bool> promise_unsubscribe;
-  RAY_CHECK_OK(gcs_client_->Actors().AsyncUnsubscribe(
-      actor_id, [&promise_unsubscribe](Status status) {
-        RAY_CHECK_OK(status);
-        promise_unsubscribe.set_value(true);
-      }));
-  ASSERT_TRUE(WaitReady(promise_unsubscribe.get_future(), timeout_ms_));
+  // Cancel subscription to an actor.
+  ASSERT_TRUE(UnsubscribeActor(actor_id));
 
-  // Update actor
+  // Update dynamic states of actor in GCS.
   actor_table_data->set_state(
       rpc::ActorTableData_ActorState::ActorTableData_ActorState_DEAD);
   ASSERT_TRUE(UpdateActor(actor_id, actor_table_data));
   ASSERT_TRUE(GetActor(actor_id).state() ==
               rpc::ActorTableData_ActorState::ActorTableData_ActorState_DEAD);
-  ASSERT_TRUE(WaitReady(promise_subscribe.get_future(), timeout_ms_));
-  auto condition = [&subscribe_callback_count]() {
-    return 1 == subscribe_callback_count;
-  };
-  EXPECT_TRUE(WaitForCondition(condition, timeout_ms_.count()));
+  WaitPendingDone(actor_update_count, 1);
 }
 
-TEST_F(ServiceBasedGcsGcsClientTest, TestActorCheckpoint) {
-  // Create actor checkpoint
+TEST_F(ServiceBasedGcsClientTest, TestActorCheckpoint) {
+  // Create actor checkpoint data.
   JobID job_id = JobID::FromInt(1);
   auto actor_table_data = GenActorTableData(job_id);
   ActorID actor_id = ActorID::FromBinary(actor_table_data->actor_id());
@@ -429,118 +595,109 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestActorCheckpoint) {
   checkpoint->set_checkpoint_id(checkpoint_id.Binary());
   checkpoint->set_execution_dependency(checkpoint_id.Binary());
 
-  // Add checkpoint
+  // Add actor checkpoint data to GCS.
   ASSERT_TRUE(AddCheckpoint(checkpoint));
 
-  // Get Checkpoint
+  // Get actor checkpoint data from GCS.
   auto get_checkpoint_result = GetCheckpoint(actor_id, checkpoint_id);
   ASSERT_TRUE(get_checkpoint_result.actor_id() == actor_id.Binary());
 
-  // Get CheckpointID
+  // Get actor checkpoint id data from GCS.
   auto get_checkpoint_id_result = GetCheckpointID(actor_id);
   ASSERT_TRUE(get_checkpoint_id_result.checkpoint_ids_size() == 1);
   ASSERT_TRUE(get_checkpoint_id_result.checkpoint_ids(0) == checkpoint_id.Binary());
 }
 
-TEST_F(ServiceBasedGcsGcsClientTest, TestActorSubscribeAll) {
-  // Create actor_table_data
+TEST_F(ServiceBasedGcsClientTest, TestActorSubscribeAll) {
+  // Create actor table data.
   JobID job_id = JobID::FromInt(1);
   auto actor_table_data1 = GenActorTableData(job_id);
   auto actor_table_data2 = GenActorTableData(job_id);
 
-  // Subscribe all
-  std::promise<bool> promise_subscribe_all;
-  std::atomic<int> subscribe_all_callback_count(0);
-  auto on_subscribe_all = [&subscribe_all_callback_count](
-                              const ActorID &actor_id, const gcs::ActorTableData &data) {
-    ++subscribe_all_callback_count;
+  // Subscribe to any register or update operations of actors.
+  std::atomic<int> actor_update_count(0);
+  auto on_subscribe = [&actor_update_count](const ActorID &actor_id,
+                                            const gcs::ActorTableData &data) {
+    ++actor_update_count;
   };
-  RAY_CHECK_OK(gcs_client_->Actors().AsyncSubscribeAll(
-      on_subscribe_all, [&promise_subscribe_all](Status status) {
-        RAY_CHECK_OK(status);
-        promise_subscribe_all.set_value(true);
-      }));
-  ASSERT_TRUE(WaitReady(promise_subscribe_all.get_future(), timeout_ms_));
+  ASSERT_TRUE(SubscribeAllActors(on_subscribe));
 
-  // Register actor
+  // Register an actor to GCS.
   ASSERT_TRUE(RegisterActor(actor_table_data1));
   ASSERT_TRUE(RegisterActor(actor_table_data2));
-  auto condition = [&subscribe_all_callback_count]() {
-    return 2 == subscribe_all_callback_count;
-  };
-  EXPECT_TRUE(WaitForCondition(condition, timeout_ms_.count()));
+  WaitPendingDone(actor_update_count, 2);
 }
 
-TEST_F(ServiceBasedGcsGcsClientTest, TestNodeInfo) {
-  // Create gcs node info
+TEST_F(ServiceBasedGcsClientTest, TestNodeInfo) {
+  // Create gcs node info.
   ClientID node1_id = ClientID::FromRandom();
   auto gcs_node1_info = GenGcsNodeInfo(node1_id.Binary());
 
-  int register_count = 0;
-  int unregister_count = 0;
-  RAY_CHECK_OK(gcs_client_->Nodes().AsyncSubscribeToNodeChange(
-      [&register_count, &unregister_count](const ClientID &node_id,
-                                           const rpc::GcsNodeInfo &data) {
-        if (data.state() == rpc::GcsNodeInfo::ALIVE) {
-          ++register_count;
-        } else if (data.state() == rpc::GcsNodeInfo::DEAD) {
-          ++unregister_count;
-        }
-      },
-      nullptr));
+  // Subscribe to node addition and removal events from GCS.
+  std::atomic<int> register_count(0);
+  std::atomic<int> unregister_count(0);
+  auto on_subscribe = [&register_count, &unregister_count](const ClientID &node_id,
+                                                           const rpc::GcsNodeInfo &data) {
+    if (data.state() == rpc::GcsNodeInfo::ALIVE) {
+      ++register_count;
+    } else if (data.state() == rpc::GcsNodeInfo::DEAD) {
+      ++unregister_count;
+    }
+  };
+  ASSERT_TRUE(SubscribeToNodeChange(on_subscribe));
 
-  // Register self
+  // Register local node to GCS.
   ASSERT_TRUE(RegisterSelf(gcs_node1_info));
   sleep(1);
   EXPECT_EQ(gcs_client_->Nodes().GetSelfId(), node1_id);
   EXPECT_EQ(gcs_client_->Nodes().GetSelfInfo().node_id(), gcs_node1_info.node_id());
   EXPECT_EQ(gcs_client_->Nodes().GetSelfInfo().state(), gcs_node1_info.state());
 
-  // Register node
+  // Register a node to GCS.
   ClientID node2_id = ClientID::FromRandom();
   auto gcs_node2_info = GenGcsNodeInfo(node2_id.Binary());
   ASSERT_TRUE(RegisterNode(gcs_node2_info));
   WaitPendingDone(register_count, 2);
 
-  // Get node list
+  // Get information of all nodes from GCS.
   std::vector<rpc::GcsNodeInfo> node_list = GetNodeInfoList();
   EXPECT_EQ(node_list.size(), 2);
-  EXPECT_EQ(register_count, 2);
   ASSERT_TRUE(gcs_client_->Nodes().Get(node1_id));
   EXPECT_EQ(gcs_client_->Nodes().GetAll().size(), 2);
 
-  // Unregister self
+  // Cancel registration of local node to GCS.
   ASSERT_TRUE(UnregisterSelf());
 
-  // Unregister node
+  // Cancel registration of a node to GCS.
   ASSERT_TRUE(UnregisterNode(node2_id));
   WaitPendingDone(unregister_count, 2);
 
+  // Get information of all nodes from GCS.
   node_list = GetNodeInfoList();
   EXPECT_EQ(node_list.size(), 2);
   EXPECT_EQ(node_list[0].state(),
             rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_DEAD);
   EXPECT_EQ(node_list[1].state(),
             rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_DEAD);
-  EXPECT_EQ(unregister_count, 2);
   ASSERT_TRUE(gcs_client_->Nodes().IsRemoved(node2_id));
 }
 
-TEST_F(ServiceBasedGcsGcsClientTest, TestNodeResources) {
-  int add_count = 0;
-  int remove_count = 0;
-  auto subscribe = [&add_count, &remove_count](
-                       const ClientID &id,
-                       const gcs::ResourceChangeNotification &notification) {
+TEST_F(ServiceBasedGcsClientTest, TestNodeResources) {
+  // Subscribe to node resource changes.
+  std::atomic<int> add_count(0);
+  std::atomic<int> remove_count(0);
+  auto on_subscribe = [&add_count, &remove_count](
+                          const ClientID &id,
+                          const gcs::ResourceChangeNotification &notification) {
     if (notification.IsAdded()) {
       ++add_count;
     } else if (notification.IsRemoved()) {
       ++remove_count;
     }
   };
-  RAY_CHECK_OK(gcs_client_->Nodes().AsyncSubscribeToResources(subscribe, nullptr));
+  ASSERT_TRUE(SubscribeToResources(on_subscribe));
 
-  // Update resources
+  // Update resources of node in GCS.
   ClientID node_id = ClientID::FromRandom();
   gcs::NodeInfoAccessor::ResourceMap resource_map;
   std::string key = "CPU";
@@ -549,26 +706,24 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestNodeResources) {
   resource_map[key] = resource;
   ASSERT_TRUE(UpdateResources(node_id, resource_map));
   WaitPendingDone(add_count, 1);
-  auto get_resources_result = GetResources(node_id);
-  ASSERT_TRUE(get_resources_result.count(key));
+  ASSERT_TRUE(GetResources(node_id).count(key));
 
-  // Delete resources
+  // Delete resources of a node from GCS.
   ASSERT_TRUE(DeleteResources(node_id, {key}));
   WaitPendingDone(remove_count, 1);
-  get_resources_result = GetResources(node_id);
-  ASSERT_TRUE(get_resources_result.empty());
+  ASSERT_TRUE(GetResources(node_id).empty());
 }
 
-TEST_F(ServiceBasedGcsGcsClientTest, TestNodeHeartbeat) {
-  int heartbeat_batch_count = 0;
-  auto heartbeat_batch_subscribe =
+TEST_F(ServiceBasedGcsClientTest, TestNodeHeartbeat) {
+  // Subscribe batched state of all nodes from GCS.
+  std::atomic<int> heartbeat_batch_count(0);
+  auto on_subscribe =
       [&heartbeat_batch_count](const gcs::HeartbeatBatchTableData &result) {
         ++heartbeat_batch_count;
       };
-  RAY_CHECK_OK(gcs_client_->Nodes().AsyncSubscribeBatchHeartbeat(
-      heartbeat_batch_subscribe, nullptr));
+  ASSERT_TRUE(SubscribeBatchHeartbeat(on_subscribe));
 
-  // Report heartbeat
+  // Report heartbeat of a node to GCS.
   ClientID node_id = ClientID::FromRandom();
   auto heartbeat = std::make_shared<rpc::HeartbeatTableData>();
   heartbeat->set_client_id(node_id.Binary());
@@ -576,61 +731,149 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestNodeHeartbeat) {
   WaitPendingDone(heartbeat_batch_count, 1);
 }
 
-TEST_F(ServiceBasedGcsGcsClientTest, TestTaskInfo) {
+TEST_F(ServiceBasedGcsClientTest, TestTaskInfo) {
   JobID job_id = JobID::FromInt(1);
   TaskID task_id = TaskID::ForDriverTask(job_id);
   auto task_table_data = GenTaskTableData(job_id.Binary(), task_id.Binary());
 
-  int task_count = 0;
+  // Subscribe to the event that the given task is added in GCS.
+  std::atomic<int> task_count(0);
   auto task_subscribe = [&task_count](const TaskID &id,
                                       const rpc::TaskTableData &result) { ++task_count; };
-  RAY_CHECK_OK(gcs_client_->Tasks().AsyncSubscribe(task_id, task_subscribe, nullptr));
+  ASSERT_TRUE(SubscribeTask(task_id, task_subscribe));
 
-  // Add task
+  // Add a task to GCS.
   ASSERT_TRUE(AddTask(task_table_data));
-
   auto get_task_result = GetTask(task_id);
   ASSERT_TRUE(get_task_result.task().task_spec().task_id() == task_id.Binary());
   ASSERT_TRUE(get_task_result.task().task_spec().job_id() == job_id.Binary());
-  RAY_CHECK_OK(gcs_client_->Tasks().AsyncUnsubscribe(task_id, nullptr));
+
+  // Cancel subscription to a task.
+  ASSERT_TRUE(UnsubscribeTask(task_id));
+
+  // Add a task to GCS again.
   ASSERT_TRUE(AddTask(task_table_data));
 
-  // Delete task
-  std::vector<TaskID> task_ids = {task_id};
-  ASSERT_TRUE(DeleteTask(task_ids));
+  // Assert unsubscribe succeeded.
+  usleep(100 * 1000);
   EXPECT_EQ(task_count, 1);
 
-  // Add task lease
-  int task_lease_count = 0;
+  // Delete tasks from GCS.
+  std::vector<TaskID> task_ids = {task_id};
+  ASSERT_TRUE(DeleteTask(task_ids));
+
+  // Subscribe to the event that the given task lease is added in GCS.
+  std::atomic<int> task_lease_count(0);
   auto task_lease_subscribe = [&task_lease_count](
                                   const TaskID &id,
                                   const boost::optional<rpc::TaskLeaseData> &result) {
     ++task_lease_count;
   };
-  RAY_CHECK_OK(gcs_client_->Tasks().AsyncSubscribeTaskLease(task_id, task_lease_subscribe,
-                                                            nullptr));
+  ASSERT_TRUE(SubscribeTaskLease(task_id, task_lease_subscribe));
+
+  // Add a task lease to GCS.
   ClientID node_id = ClientID::FromRandom();
   auto task_lease = GenTaskLeaseData(task_id.Binary(), node_id.Binary());
   ASSERT_TRUE(AddTaskLease(task_lease));
   WaitPendingDone(task_lease_count, 2);
 
-  RAY_CHECK_OK(gcs_client_->Tasks().AsyncUnsubscribeTaskLease(task_id, nullptr));
+  // Cancel subscription to a task lease.
+  ASSERT_TRUE(UnsubscribeTaskLease(task_id));
+
+  // Add a task lease to GCS again.
   ASSERT_TRUE(AddTaskLease(task_lease));
+
+  // Assert unsubscribe succeeded.
+  usleep(100 * 1000);
   EXPECT_EQ(task_lease_count, 2);
 
-  // Attempt task reconstruction
+  // Attempt task reconstruction to GCS.
   auto task_reconstruction_data = std::make_shared<rpc::TaskReconstructionData>();
   task_reconstruction_data->set_task_id(task_id.Binary());
   task_reconstruction_data->set_num_reconstructions(0);
   ASSERT_TRUE(AttemptTaskReconstruction(task_reconstruction_data));
 }
 
-TEST_F(ServiceBasedGcsGcsClientTest, TestDetectGcsAvailability) {
-  // Create job_table_data
+TEST_F(ServiceBasedGcsClientTest, TestObjectInfo) {
+  ObjectID object_id = ObjectID::FromRandom();
+  ClientID node_id = ClientID::FromRandom();
+
+  // Subscribe to any update of an object's location.
+  std::atomic<int> object_add_count(0);
+  std::atomic<int> object_remove_count(0);
+  auto on_subscribe = [&object_add_count, &object_remove_count](
+                          const ObjectID &object_id,
+                          const gcs::ObjectChangeNotification &result) {
+    if (!result.GetData().empty()) {
+      if (result.IsAdded()) {
+        ++object_add_count;
+      } else if (result.IsRemoved()) {
+        ++object_remove_count;
+      }
+    }
+  };
+  ASSERT_TRUE(SubscribeToLocations(object_id, on_subscribe));
+
+  // Add location of object to GCS.
+  ASSERT_TRUE(AddLocation(object_id, node_id));
+  WaitPendingDone(object_add_count, 1);
+
+  // Get object's locations from GCS.
+  auto locations = GetLocations(object_id);
+  ASSERT_EQ(locations.size(), 1);
+  ASSERT_EQ(locations.back().manager(), node_id.Binary());
+
+  // Remove location of object from GCS.
+  ASSERT_TRUE(RemoveLocation(object_id, node_id));
+  WaitPendingDone(object_remove_count, 1);
+  ASSERT_TRUE(GetLocations(object_id).empty());
+
+  // Cancel subscription to any update of an object's location.
+  ASSERT_TRUE(UnsubscribeToLocations(object_id));
+
+  // Add location of object to GCS again.
+  ASSERT_TRUE(AddLocation(object_id, node_id));
+
+  // Assert unsubscribe succeeded.
+  usleep(100 * 1000);
+  ASSERT_EQ(object_add_count, 1);
+}
+
+TEST_F(ServiceBasedGcsClientTest, TestStats) {
+  // Add profile data to GCS.
+  ClientID node_id = ClientID::FromRandom();
+  auto profile_table_data = GenProfileTableData(node_id);
+  ASSERT_TRUE(AddProfileData(profile_table_data));
+}
+
+TEST_F(ServiceBasedGcsClientTest, TestWorkerInfo) {
+  // Subscribe to all unexpected failure of workers from GCS.
+  std::atomic<int> worker_failure_count(0);
+  auto on_subscribe = [&worker_failure_count](const WorkerID &worker_id,
+                                              const rpc::WorkerFailureData &result) {
+    ++worker_failure_count;
+  };
+  ASSERT_TRUE(SubscribeToWorkerFailures(on_subscribe));
+
+  // Report a worker failure to GCS.
+  auto worker_failure_data = GenWorkerFailureData();
+  ASSERT_TRUE(ReportWorkerFailure(worker_failure_data));
+  WaitPendingDone(worker_failure_count, 1);
+}
+
+TEST_F(ServiceBasedGcsClientTest, TestErrorInfo) {
+  // Report a job error to GCS.
+  JobID job_id = JobID::FromInt(1);
+  auto error_table_data = GenErrorTableData(job_id);
+  ASSERT_TRUE(ReportJobError(error_table_data));
+}
+
+TEST_F(ServiceBasedGcsClientTest, TestDetectGcsAvailability) {
+  // Create job table data.
   JobID add_job_id = JobID::FromInt(1);
   auto job_table_data = GenJobTableData(add_job_id);
 
-  RAY_LOG(INFO) << "Gcs service init port = " << gcs_server_->GetPort();
+  RAY_LOG(INFO) << "Initializing GCS service, port = " << gcs_server_->GetPort();
   gcs_server_->Stop();
   thread_gcs_server_->join();
 
@@ -641,7 +884,7 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestDetectGcsAvailability) {
   while (gcs_server_->GetPort() == 0) {
     std::this_thread::sleep_for(std::chrono::milliseconds(10));
   }
-  RAY_LOG(INFO) << "Gcs service restart success, port = " << gcs_server_->GetPort();
+  RAY_LOG(INFO) << "GCS service restarted, port = " << gcs_server_->GetPort();
 
   std::promise<bool> promise;
   RAY_CHECK_OK(gcs_client_->Jobs().AsyncAdd(
@@ -649,7 +892,7 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestDetectGcsAvailability) {
   promise.get_future().get();
 }
 
-TEST_F(ServiceBasedGcsGcsClientTest, TestGcsRedisFailureDetector) {
+TEST_F(ServiceBasedGcsClientTest, TestGcsRedisFailureDetector) {
   // Stop redis.
   TearDownTestCase();
 
@@ -658,7 +901,7 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestGcsRedisFailureDetector) {
   auto interval_ms = RayConfig::instance().gcs_redis_heartbeat_interval_milliseconds();
   std::this_thread::sleep_for(std::chrono::milliseconds(3 * interval_ms));
 
-  // Check if gcs server has exited.
+  // Check if GCS server has exited.
   RAY_CHECK(gcs_server_->IsStopped());
 }
 

src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc

@@ -19,10 +19,6 @@
 
 namespace ray {
 
-static std::string redis_server_executable;
-static std::string redis_client_executable;
-static std::string libray_redis_module_path;
-
 class GcsServerTest : public RedisServiceManagerForTest {
  public:
   void SetUp() override {

streaming/src/test/mock_actor.cc

@@ -421,7 +421,6 @@ class StreamingWorker {
     std::shared_ptr<ActorHandle> actor_handle(new ActorHandle(actor_handle_serialized));
     STREAMING_CHECK(actor_handle != nullptr);
     STREAMING_LOG(INFO) << " actor id from handle: " << actor_handle->GetActorID();
-    ;
 
     // STREAMING_LOG(INFO) << "actor_handle_serialized: " << actor_handle_serialized;
     // peer_actor_handle_ =