diff --git a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc index 36b3578bf..62a44acb0 100644 --- a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc +++ b/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 &subscribe) { + std::promise 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 &job_table_data) { std::promise 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 &subscribe) { + std::promise 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 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 &subscribe) { + std::promise 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 &actor_table_data) { std::promise 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 &subscribe) { + std::promise 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 + &subscribe) { + std::promise 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 promise; @@ -222,6 +267,14 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest { return WaitReady(promise.get_future(), timeout_ms_); } + bool SubscribeBatchHeartbeat( + const gcs::ItemCallback &subscribe) { + std::promise 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 heartbeat) { std::promise 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 &subscribe) { + std::promise 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 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 task) { std::promise 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> + &subscribe) { + std::promise 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 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 task_lease) { std::promise 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 &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 &subscribe) { + std::promise 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 ¤t_count, int expected_count) { + bool UnsubscribeToLocations(const ObjectID &object_id) { + std::promise 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 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 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 GetLocations(const ObjectID &object_id) { + std::promise promise; + std::vector locations; + RAY_CHECK_OK(gcs_client_->Objects().AsyncGetLocations( + object_id, [&locations, &promise]( + Status status, const std::vector &result) { + locations = result; + promise.set_value(status.ok()); + })); + EXPECT_TRUE(WaitReady(promise.get_future(), timeout_ms_)); + return locations; + } + + bool AddProfileData(const std::shared_ptr &profile_table_data) { + std::promise 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 &error_table_data) { + std::promise 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 &subscribe) { + std::promise 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 &worker_failure_data) { + std::promise 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 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 ¤t_count, int expected_count) { auto condition = [¤t_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 GenProfileTableData(const ClientID &node_id) { + auto profile_table_data = std::make_shared(); + profile_table_data->set_component_id(node_id.Binary()); + return profile_table_data; + } + + std::shared_ptr GenErrorTableData(const JobID &job_id) { + auto error_table_data = std::make_shared(); + error_table_data->set_job_id(job_id.Binary()); + return error_table_data; + } + + std::shared_ptr GenWorkerFailureData() { + auto worker_failure_data = std::make_shared(); + worker_failure_data->set_timestamp(std::time(nullptr)); + return worker_failure_data; + } + + // GCS server. gcs::GcsServerConfig config; std::unique_ptr gcs_server_; std::unique_ptr thread_io_service_; std::unique_ptr thread_gcs_server_; std::unique_ptr io_service_; - // Gcs client + // GCS client. std::unique_ptr 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 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 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 promise_subscribe; - std::atomic 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 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 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 promise_subscribe_all; - std::atomic 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 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( - [®ister_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 register_count(0); + std::atomic unregister_count(0); + auto on_subscribe = [®ister_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 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 ¬ification) { +TEST_F(ServiceBasedGcsClientTest, TestNodeResources) { + // Subscribe to node resource changes. + std::atomic add_count(0); + std::atomic remove_count(0); + auto on_subscribe = [&add_count, &remove_count]( + const ClientID &id, + const gcs::ResourceChangeNotification ¬ification) { 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 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(); 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 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 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 task_ids = {task_id}; + ASSERT_TRUE(DeleteTask(task_ids)); + + // Subscribe to the event that the given task lease is added in GCS. + std::atomic task_lease_count(0); auto task_lease_subscribe = [&task_lease_count]( const TaskID &id, const boost::optional &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(); 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 object_add_count(0); + std::atomic 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 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 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()); } diff --git a/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc b/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc index 766ddf31e..98aeae1e3 100644 --- a/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc +++ b/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 { diff --git a/streaming/src/test/mock_actor.cc b/streaming/src/test/mock_actor.cc index 21212ac28..1ef0bc0d8 100644 --- a/streaming/src/test/mock_actor.cc +++ b/streaming/src/test/mock_actor.cc @@ -421,7 +421,6 @@ class StreamingWorker { std::shared_ptr 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_ =