From c17404918c40b2824253ad544ea2c8e7b37d72f1 Mon Sep 17 00:00:00 2001 From: fangfengbin <869218239a@zju.edu.cn> Date: Wed, 15 Apr 2020 10:48:12 +0800 Subject: [PATCH] [GCS]Add gcs table storage interface (#7949) --- BUILD.bazel | 58 ++- .../test/service_based_gcs_client_test.cc | 107 +---- src/ray/gcs/gcs_server/gcs_table_storage.cc | 108 +++++ src/ray/gcs/gcs_server/gcs_table_storage.h | 444 ++++++++++++++++++ .../gcs_server/test/gcs_actor_manager_test.cc | 2 +- .../test/gcs_actor_scheduler_test.cc | 2 +- .../gcs_server/test/gcs_node_manager_test.cc | 2 +- .../gcs_server/test/gcs_server_rpc_test.cc | 108 ++--- .../test/redis_gcs_table_storage_test.cc | 129 +++++ .../gcs/store_client/redis_store_client.cc | 14 + .../gcs/{gcs_server => }/test/gcs_test_util.h | 75 ++- src/ray/protobuf/gcs.proto | 7 + 12 files changed, 874 insertions(+), 182 deletions(-) create mode 100644 src/ray/gcs/gcs_server/gcs_table_storage.cc create mode 100644 src/ray/gcs/gcs_server/gcs_table_storage.h create mode 100644 src/ray/gcs/gcs_server/test/redis_gcs_table_storage_test.cc rename src/ray/gcs/{gcs_server => }/test/gcs_test_util.h (82%) diff --git a/BUILD.bazel b/BUILD.bazel index eb8a26a10..adf15c46c 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -714,7 +714,10 @@ cc_test( cc_test( name = "gcs_server_rpc_test", - srcs = ["src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc"], + srcs = [ + "src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc", + "src/ray/gcs/test/gcs_test_util.h", + ], args = ["$(location redis-server) $(location redis-cli) $(location libray_redis_module.so)"], copts = COPTS, data = [ @@ -732,7 +735,7 @@ cc_test( name = "gcs_node_manager_test", srcs = [ "src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc", - "src/ray/gcs/gcs_server/test/gcs_test_util.h", + "src/ray/gcs/test/gcs_test_util.h", ], copts = COPTS, deps = [ @@ -745,7 +748,7 @@ cc_test( name = "gcs_actor_scheduler_test", srcs = [ "src/ray/gcs/gcs_server/test/gcs_actor_scheduler_test.cc", - "src/ray/gcs/gcs_server/test/gcs_test_util.h", + "src/ray/gcs/test/gcs_test_util.h", ], copts = COPTS, deps = [ @@ -758,7 +761,7 @@ cc_test( name = "gcs_actor_manager_test", srcs = [ "src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc", - "src/ray/gcs/gcs_server/test/gcs_test_util.h", + "src/ray/gcs/test/gcs_test_util.h", ], copts = COPTS, deps = [ @@ -767,6 +770,47 @@ cc_test( ], ) +cc_library( + name = "gcs_table_storage_lib", + srcs = glob( + [ + "src/ray/gcs/gcs_server/gcs_table_storage.cc", + ], + ), + hdrs = glob( + [ + "src/ray/gcs/gcs_server/gcs_table_storage.h", + ], + ), + copts = COPTS, + deps = [ + ":gcs", + ":ray_common", + ":redis_store_client", + ], +) + +cc_test( + name = "redis_gcs_table_storage_test", + srcs = [ + "src/ray/gcs/gcs_server/test/redis_gcs_table_storage_test.cc", + "src/ray/gcs/test/gcs_test_util.h", + ], + args = ["$(location redis-server) $(location redis-cli) $(location libray_redis_module.so)"], + copts = COPTS, + data = [ + "//:libray_redis_module.so", + "//:redis-cli", + "//:redis-server", + ], + deps = [ + ":gcs_server_lib", + ":gcs_table_storage_lib", + ":store_client_test_lib", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "service_based_gcs_client_lib", srcs = glob( @@ -783,12 +827,16 @@ cc_library( deps = [ ":gcs", ":gcs_service_rpc", + ":redis_store_client", ], ) cc_test( name = "gcs_server_test", - srcs = ["src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc"], + srcs = [ + "src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc", + "src/ray/gcs/test/gcs_test_util.h", + ], args = ["$(location redis-server) $(location redis-cli) $(location libray_redis_module.so)"], copts = COPTS, data = [ 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 62a44acb0..a9d08d2e8 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 @@ -17,6 +17,7 @@ #include "ray/common/test_util.h" #include "ray/gcs/gcs_client/service_based_accessor.h" #include "ray/gcs/gcs_server/gcs_server.h" +#include "ray/gcs/test/gcs_test_util.h" #include "ray/rpc/gcs_server/gcs_rpc_client.h" namespace ray { @@ -452,74 +453,6 @@ class ServiceBasedGcsClientTest : public RedisServiceManagerForTest { EXPECT_TRUE(WaitForCondition(condition, timeout_ms_.count())); } - std::shared_ptr GenJobTableData(JobID job_id) { - auto job_table_data = std::make_shared(); - job_table_data->set_job_id(job_id.Binary()); - job_table_data->set_is_dead(false); - job_table_data->set_timestamp(std::time(nullptr)); - job_table_data->set_node_manager_address("127.0.0.1"); - job_table_data->set_driver_pid(5667L); - return job_table_data; - } - - std::shared_ptr GenActorTableData(const JobID &job_id) { - auto actor_table_data = std::make_shared(); - ActorID actor_id = ActorID::Of(job_id, RandomTaskId(), 0); - actor_table_data->set_actor_id(actor_id.Binary()); - actor_table_data->set_job_id(job_id.Binary()); - actor_table_data->set_state( - rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE); - actor_table_data->set_max_reconstructions(1); - actor_table_data->set_remaining_reconstructions(1); - return actor_table_data; - } - - rpc::GcsNodeInfo GenGcsNodeInfo(const std::string &node_id) { - rpc::GcsNodeInfo gcs_node_info; - gcs_node_info.set_node_id(node_id); - gcs_node_info.set_state(rpc::GcsNodeInfo_GcsNodeState_ALIVE); - gcs_node_info.set_node_manager_address("127.0.0.1"); - return gcs_node_info; - } - - std::shared_ptr GenTaskTableData(const std::string &job_id, - const std::string &task_id) { - auto task_table_data = std::make_shared(); - rpc::Task task; - rpc::TaskSpec task_spec; - task_spec.set_job_id(job_id); - task_spec.set_task_id(task_id); - task.mutable_task_spec()->CopyFrom(task_spec); - task_table_data->mutable_task()->CopyFrom(task); - return task_table_data; - } - - std::shared_ptr GenTaskLeaseData(const std::string &task_id, - const std::string &node_id) { - auto task_lease_data = std::make_shared(); - task_lease_data->set_task_id(task_id); - task_lease_data->set_node_manager_id(node_id); - return task_lease_data; - } - - 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_; @@ -537,7 +470,7 @@ class ServiceBasedGcsClientTest : public RedisServiceManagerForTest { TEST_F(ServiceBasedGcsClientTest, TestJobInfo) { // Create job table data. JobID add_job_id = JobID::FromInt(1); - auto job_table_data = GenJobTableData(add_job_id); + auto job_table_data = Mocker::GenJobTableData(add_job_id); // Subscribe to finished jobs. std::atomic finished_job_count(0); @@ -555,7 +488,7 @@ TEST_F(ServiceBasedGcsClientTest, TestJobInfo) { TEST_F(ServiceBasedGcsClientTest, TestActorInfo) { // Create actor table data. JobID job_id = JobID::FromInt(1); - auto actor_table_data = GenActorTableData(job_id); + auto actor_table_data = Mocker::GenActorTableData(job_id); ActorID actor_id = ActorID::FromBinary(actor_table_data->actor_id()); // Subscribe to any update operations of an actor. @@ -586,7 +519,7 @@ TEST_F(ServiceBasedGcsClientTest, TestActorInfo) { TEST_F(ServiceBasedGcsClientTest, TestActorCheckpoint) { // Create actor checkpoint data. JobID job_id = JobID::FromInt(1); - auto actor_table_data = GenActorTableData(job_id); + auto actor_table_data = Mocker::GenActorTableData(job_id); ActorID actor_id = ActorID::FromBinary(actor_table_data->actor_id()); ActorCheckpointID checkpoint_id = ActorCheckpointID::FromRandom(); @@ -611,8 +544,8 @@ TEST_F(ServiceBasedGcsClientTest, TestActorCheckpoint) { 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); + auto actor_table_data1 = Mocker::GenActorTableData(job_id); + auto actor_table_data2 = Mocker::GenActorTableData(job_id); // Subscribe to any register or update operations of actors. std::atomic actor_update_count(0); @@ -630,8 +563,8 @@ TEST_F(ServiceBasedGcsClientTest, TestActorSubscribeAll) { TEST_F(ServiceBasedGcsClientTest, TestNodeInfo) { // Create gcs node info. - ClientID node1_id = ClientID::FromRandom(); - auto gcs_node1_info = GenGcsNodeInfo(node1_id.Binary()); + auto gcs_node1_info = Mocker::GenNodeInfo(); + ClientID node1_id = ClientID::FromBinary(gcs_node1_info->node_id()); // Subscribe to node addition and removal events from GCS. std::atomic register_count(0); @@ -647,16 +580,16 @@ TEST_F(ServiceBasedGcsClientTest, TestNodeInfo) { ASSERT_TRUE(SubscribeToNodeChange(on_subscribe)); // Register local node to GCS. - ASSERT_TRUE(RegisterSelf(gcs_node1_info)); + 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()); + 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 a node to GCS. - ClientID node2_id = ClientID::FromRandom(); - auto gcs_node2_info = GenGcsNodeInfo(node2_id.Binary()); - ASSERT_TRUE(RegisterNode(gcs_node2_info)); + auto gcs_node2_info = Mocker::GenNodeInfo(); + ClientID node2_id = ClientID::FromBinary(gcs_node2_info->node_id()); + ASSERT_TRUE(RegisterNode(*gcs_node2_info)); WaitPendingDone(register_count, 2); // Get information of all nodes from GCS. @@ -734,7 +667,7 @@ TEST_F(ServiceBasedGcsClientTest, TestNodeHeartbeat) { 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()); + auto task_table_data = Mocker::GenTaskTableData(job_id.Binary(), task_id.Binary()); // Subscribe to the event that the given task is added in GCS. std::atomic task_count(0); @@ -773,7 +706,7 @@ TEST_F(ServiceBasedGcsClientTest, TestTaskInfo) { // Add a task lease to GCS. ClientID node_id = ClientID::FromRandom(); - auto task_lease = GenTaskLeaseData(task_id.Binary(), node_id.Binary()); + auto task_lease = Mocker::GenTaskLeaseData(task_id.Binary(), node_id.Binary()); ASSERT_TRUE(AddTaskLease(task_lease)); WaitPendingDone(task_lease_count, 2); @@ -842,7 +775,7 @@ TEST_F(ServiceBasedGcsClientTest, TestObjectInfo) { TEST_F(ServiceBasedGcsClientTest, TestStats) { // Add profile data to GCS. ClientID node_id = ClientID::FromRandom(); - auto profile_table_data = GenProfileTableData(node_id); + auto profile_table_data = Mocker::GenProfileTableData(node_id); ASSERT_TRUE(AddProfileData(profile_table_data)); } @@ -856,7 +789,7 @@ TEST_F(ServiceBasedGcsClientTest, TestWorkerInfo) { ASSERT_TRUE(SubscribeToWorkerFailures(on_subscribe)); // Report a worker failure to GCS. - auto worker_failure_data = GenWorkerFailureData(); + auto worker_failure_data = Mocker::GenWorkerFailureData(); ASSERT_TRUE(ReportWorkerFailure(worker_failure_data)); WaitPendingDone(worker_failure_count, 1); } @@ -864,14 +797,14 @@ TEST_F(ServiceBasedGcsClientTest, TestWorkerInfo) { TEST_F(ServiceBasedGcsClientTest, TestErrorInfo) { // Report a job error to GCS. JobID job_id = JobID::FromInt(1); - auto error_table_data = GenErrorTableData(job_id); + auto error_table_data = Mocker::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); + auto job_table_data = Mocker::GenJobTableData(add_job_id); RAY_LOG(INFO) << "Initializing GCS service, port = " << gcs_server_->GetPort(); gcs_server_->Stop(); diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.cc b/src/ray/gcs/gcs_server/gcs_table_storage.cc new file mode 100644 index 000000000..2b41a4edc --- /dev/null +++ b/src/ray/gcs/gcs_server/gcs_table_storage.cc @@ -0,0 +1,108 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "gcs_table_storage.h" +#include "ray/common/id.h" +#include "ray/common/status.h" +#include "ray/gcs/callback.h" + +namespace ray { +namespace gcs { + +template +Status GcsTable::Put(const Key &key, const Data &value, + const StatusCallback &callback) { + return store_client_->AsyncPut(table_name_, key, value, callback); +} + +template +Status GcsTable::Get(const Key &key, + const OptionalItemCallback &callback) { + return store_client_->AsyncGet(table_name_, key, callback); +} + +template +Status GcsTable::GetAll( + const SegmentedCallback> &callback) { + return store_client_->AsyncGetAll(table_name_, callback); +} + +template +Status GcsTable::Delete(const Key &key, const StatusCallback &callback) { + return store_client_->AsyncDelete(table_name_, key, callback); +} + +template +Status GcsTable::BatchDelete(const std::vector &keys, + const StatusCallback &callback) { + // TODO(ffbin): We will use redis store client batch delete interface directly later. + auto finished_count = std::make_shared(0); + int size = keys.size(); + for (Key key : keys) { + auto done = [finished_count, size, callback](Status status) { + ++(*finished_count); + if (*finished_count == size) { + callback(Status::OK()); + } + }; + RAY_CHECK_OK(store_client_->AsyncDelete(table_name_, key, done)); + } + return Status::OK(); +} + +template +Status GcsTableWithJobId::Put(const Key &key, const Data &value, + const StatusCallback &callback) { + return this->store_client_->AsyncPutWithIndex(this->table_name_, key, + GetJobIdFromKey(key), value, callback); +} + +template +Status GcsTableWithJobId::GetByJobId( + const JobID &job_id, const SegmentedCallback> &callback) { + // TODO(ffbin): We will add this function after redis store client support + // AsyncGetByIndex interface. + return Status::NotImplemented("GetByJobId not implemented"); +} + +template +Status GcsTableWithJobId::DeleteByJobId(const JobID &job_id, + const StatusCallback &callback) { + return this->store_client_->AsyncDeleteByIndex(this->table_name_, job_id, callback); +} + +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTable; +template class GcsTableWithJobId; +template class GcsTableWithJobId; +template class GcsTableWithJobId; +template class GcsTableWithJobId; +template class GcsTableWithJobId; +template class GcsTableWithJobId; + +} // namespace gcs +} // namespace ray diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.h b/src/ray/gcs/gcs_server/gcs_table_storage.h new file mode 100644 index 000000000..1a69857ce --- /dev/null +++ b/src/ray/gcs/gcs_server/gcs_table_storage.h @@ -0,0 +1,444 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef RAY_GCS_GCS_TABLE_STORAGE_H_ +#define RAY_GCS_GCS_TABLE_STORAGE_H_ + +#include + +#include "ray/gcs/store_client/redis_store_client.h" +#include "ray/protobuf/gcs.pb.h" + +namespace ray { +namespace gcs { + +using rpc::ActorCheckpointData; +using rpc::ActorCheckpointIdData; +using rpc::ActorTableData; +using rpc::ErrorTableData; +using rpc::GcsNodeInfo; +using rpc::HeartbeatBatchTableData; +using rpc::HeartbeatTableData; +using rpc::JobTableData; +using rpc::ObjectTableData; +using rpc::ObjectTableDataList; +using rpc::ProfileTableData; +using rpc::ResourceMap; +using rpc::ResourceTableData; +using rpc::TaskLeaseData; +using rpc::TaskReconstructionData; +using rpc::TaskTableData; +using rpc::WorkerFailureData; + +/// \class GcsTable +/// +/// GcsTable is the storage interface for all GCS tables whose data do not belong to +/// specific jobs. This class is not meant to be used directly. All gcs table classes +/// without job id should derive from this class and override the table_name_ member with +/// a unique value for that table. +template +class GcsTable { + public: + explicit GcsTable(std::shared_ptr> store_client) + : store_client_(store_client) {} + + virtual ~GcsTable() = default; + + /// Write data to the table asynchronously. + /// + /// \param key The key that will be written to the table. + /// \param value The value of the key that will be written to the table. + /// \param callback Callback that will be called after write finishes. + /// \return Status + virtual Status Put(const Key &key, const Data &value, const StatusCallback &callback); + + /// Get data from the table asynchronously. + /// + /// \param key The key to lookup from the table. + /// \param callback Callback that will be called after read finishes. + /// \return Status + Status Get(const Key &key, const OptionalItemCallback &callback); + + /// Get all data from the table asynchronously. + /// + /// \param callback Callback that will be called after data has been received. + /// If the callback return `has_more == true` mean there's more data to be received. + /// \return Status + Status GetAll(const SegmentedCallback> &callback); + + /// Delete data from the table asynchronously. + /// + /// \param key The key that will be deleted from the table. + /// \param callback Callback that will be called after delete finishes. + /// \return Status + Status Delete(const Key &key, const StatusCallback &callback); + + /// Delete a batch of data from the table asynchronously. + /// + /// \param keys The batch key that will be deleted from the table. + /// \param callback Callback that will be called after delete finishes. + /// \return Status + Status BatchDelete(const std::vector &keys, const StatusCallback &callback); + + protected: + std::string table_name_; + std::shared_ptr> store_client_; +}; + +/// \class GcsTableWithJobId +/// +/// GcsTableWithJobId is the storage interface for all GCS tables whose data belongs to +/// specific jobs. This class is not meant to be used directly. All gcs table classes with +/// job id should derive from this class and override the table_name_ member with a unique +/// value for that table. +template +class GcsTableWithJobId : public GcsTable { + public: + explicit GcsTableWithJobId(std::shared_ptr> store_client) + : GcsTable(store_client) {} + + /// Write data to the table asynchronously. + /// + /// \param key The key that will be written to the table. The job id can be obtained + /// from the key. \param value The value of the key that will be written to the table. + /// \param callback Callback that will be called after write finishes. + /// \return Status + Status Put(const Key &key, const Data &value, const StatusCallback &callback) override; + + /// Get all the data of the specified job id from the table asynchronously. + /// + /// \param job_id The key to lookup from the table. + /// \param callback Callback that will be called after read finishes. + /// \return Status + Status GetByJobId(const JobID &job_id, + const SegmentedCallback> &callback); + + /// Delete all the data of the specified job id from the table asynchronously. + /// + /// \param job_id The key that will be deleted from the table. + /// \param callback Callback that will be called after delete finishes. + /// \return Status + Status DeleteByJobId(const JobID &job_id, const StatusCallback &callback); + + protected: + virtual JobID GetJobIdFromKey(const Key &key) = 0; +}; + +class GcsJobTable : public GcsTable { + public: + explicit GcsJobTable( + std::shared_ptr> store_client) + : GcsTable(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::JOB); + } +}; + +class GcsActorTable : public GcsTableWithJobId { + public: + explicit GcsActorTable( + std::shared_ptr> store_client) + : GcsTableWithJobId(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::ACTOR); + } + + private: + JobID GetJobIdFromKey(const ActorID &key) { return key.JobId(); } +}; + +class GcsActorCheckpointTable : public GcsTable { + public: + explicit GcsActorCheckpointTable( + std::shared_ptr> + store_client) + : GcsTable(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::ACTOR_CHECKPOINT); + } +}; + +class GcsActorCheckpointIdTable + : public GcsTableWithJobId { + public: + explicit GcsActorCheckpointIdTable( + std::shared_ptr> store_client) + : GcsTableWithJobId(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::ACTOR_CHECKPOINT_ID); + } + + private: + JobID GetJobIdFromKey(const ActorID &key) { return key.JobId(); } +}; + +class GcsTaskTable : public GcsTableWithJobId { + public: + explicit GcsTaskTable( + std::shared_ptr> store_client) + : GcsTableWithJobId(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::TASK); + } + + private: + JobID GetJobIdFromKey(const TaskID &key) { return key.ActorId().JobId(); } +}; + +class GcsTaskLeaseTable : public GcsTableWithJobId { + public: + explicit GcsTaskLeaseTable( + std::shared_ptr> store_client) + : GcsTableWithJobId(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::TASK_LEASE); + } + + private: + JobID GetJobIdFromKey(const TaskID &key) { return key.ActorId().JobId(); } +}; + +class GcsTaskReconstructionTable + : public GcsTableWithJobId { + public: + explicit GcsTaskReconstructionTable( + std::shared_ptr> store_client) + : GcsTableWithJobId(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::TASK_RECONSTRUCTION); + } + + private: + JobID GetJobIdFromKey(const TaskID &key) { return key.ActorId().JobId(); } +}; + +class GcsObjectTable : public GcsTableWithJobId { + public: + explicit GcsObjectTable( + std::shared_ptr> store_client) + : GcsTableWithJobId(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::OBJECT); + } + + private: + JobID GetJobIdFromKey(const ObjectID &key) { return key.TaskId().JobId(); } +}; + +class GcsNodeTable : public GcsTable { + public: + explicit GcsNodeTable( + std::shared_ptr> store_client) + : GcsTable(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::CLIENT); + } +}; + +class GcsNodeResourceTable : public GcsTable { + public: + explicit GcsNodeResourceTable( + std::shared_ptr> store_client) + : GcsTable(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::NODE_RESOURCE); + } +}; + +class GcsHeartbeatTable : public GcsTable { + public: + explicit GcsHeartbeatTable( + std::shared_ptr> store_client) + : GcsTable(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::HEARTBEAT); + } +}; + +class GcsHeartbeatBatchTable : public GcsTable { + public: + explicit GcsHeartbeatBatchTable( + std::shared_ptr> store_client) + : GcsTable(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::HEARTBEAT_BATCH); + } +}; + +class GcsErrorInfoTable : public GcsTable { + public: + explicit GcsErrorInfoTable( + std::shared_ptr> store_client) + : GcsTable(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::ERROR_INFO); + } +}; + +class GcsProfileTable : public GcsTable { + public: + explicit GcsProfileTable( + std::shared_ptr> store_client) + : GcsTable(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::PROFILE); + } +}; + +class GcsWorkerFailureTable : public GcsTable { + public: + explicit GcsWorkerFailureTable( + std::shared_ptr> store_client) + : GcsTable(std::move(store_client)) { + table_name_ = TablePrefix_Name(TablePrefix::WORKER_FAILURE); + } +}; + +/// \class GcsTableStorage +/// +/// This class is not meant to be used directly. All gcs table storage classes should +/// derive from this class and override class member variables. +class GcsTableStorage { + public: + GcsJobTable &JobTable() { + RAY_CHECK(job_table_ != nullptr); + return *job_table_; + } + + GcsActorTable &ActorTable() { + RAY_CHECK(actor_table_ != nullptr); + return *actor_table_; + } + + GcsActorCheckpointTable &ActorCheckpointTable() { + RAY_CHECK(actor_checkpoint_table_ != nullptr); + return *actor_checkpoint_table_; + } + + GcsActorCheckpointIdTable &ActorCheckpointIdTable() { + RAY_CHECK(actor_checkpoint_id_table_ != nullptr); + return *actor_checkpoint_id_table_; + } + + GcsTaskTable &TaskTable() { + RAY_CHECK(task_table_ != nullptr); + return *task_table_; + } + + GcsTaskLeaseTable &TaskLeaseTable() { + RAY_CHECK(task_lease_table_ != nullptr); + return *task_lease_table_; + } + + GcsTaskReconstructionTable &TaskReconstructionTable() { + RAY_CHECK(task_reconstruction_table_ != nullptr); + return *task_reconstruction_table_; + } + + GcsObjectTable &ObjectTable() { + RAY_CHECK(object_table_ != nullptr); + return *object_table_; + } + + GcsNodeTable &NodeTable() { + RAY_CHECK(node_table_ != nullptr); + return *node_table_; + } + + GcsNodeResourceTable &NodeResourceTable() { + RAY_CHECK(node_resource_table_ != nullptr); + return *node_resource_table_; + } + + GcsHeartbeatTable &HeartbeatTable() { + RAY_CHECK(heartbeat_table_ != nullptr); + return *heartbeat_table_; + } + + GcsHeartbeatBatchTable &HeartbeatBatchTable() { + RAY_CHECK(heartbeat_batch_table_ != nullptr); + return *heartbeat_batch_table_; + } + + GcsErrorInfoTable &ErrorInfoTable() { + RAY_CHECK(error_info_table_ != nullptr); + return *error_info_table_; + } + + GcsProfileTable &ProfileTable() { + RAY_CHECK(profile_table_ != nullptr); + return *profile_table_; + } + + GcsWorkerFailureTable &WorkerFailureTable() { + RAY_CHECK(worker_failure_table_ != nullptr); + return *worker_failure_table_; + } + + protected: + std::unique_ptr job_table_; + std::unique_ptr actor_table_; + std::unique_ptr actor_checkpoint_table_; + std::unique_ptr actor_checkpoint_id_table_; + std::unique_ptr task_table_; + std::unique_ptr task_lease_table_; + std::unique_ptr task_reconstruction_table_; + std::unique_ptr object_table_; + std::unique_ptr node_table_; + std::unique_ptr node_resource_table_; + std::unique_ptr heartbeat_table_; + std::unique_ptr heartbeat_batch_table_; + std::unique_ptr error_info_table_; + std::unique_ptr profile_table_; + std::unique_ptr worker_failure_table_; +}; + +/// \class RedisGcsTableStorage +/// RedisGcsTableStorage is an implementation of `GcsTableStorage` +/// that uses redis as storage. +class RedisGcsTableStorage : public GcsTableStorage { + public: + explicit RedisGcsTableStorage(std::shared_ptr redis_client) { + job_table_.reset(new GcsJobTable( + std::make_shared>(redis_client))); + actor_table_.reset(new GcsActorTable( + std::make_shared>( + redis_client))); + actor_checkpoint_table_.reset(new GcsActorCheckpointTable( + std::make_shared>( + redis_client))); + actor_checkpoint_id_table_.reset(new GcsActorCheckpointIdTable( + std::make_shared>( + redis_client))); + task_table_.reset(new GcsTaskTable( + std::make_shared>(redis_client))); + task_lease_table_.reset(new GcsTaskLeaseTable( + std::make_shared>(redis_client))); + task_reconstruction_table_.reset(new GcsTaskReconstructionTable( + std::make_shared>( + redis_client))); + object_table_.reset(new GcsObjectTable( + std::make_shared>( + redis_client))); + node_table_.reset(new GcsNodeTable( + std::make_shared>(redis_client))); + node_resource_table_.reset(new GcsNodeResourceTable( + std::make_shared>(redis_client))); + heartbeat_table_.reset(new GcsHeartbeatTable( + std::make_shared>( + redis_client))); + heartbeat_batch_table_.reset(new GcsHeartbeatBatchTable( + std::make_shared>( + redis_client))); + error_info_table_.reset(new GcsErrorInfoTable( + std::make_shared>(redis_client))); + profile_table_.reset(new GcsProfileTable( + std::make_shared>( + redis_client))); + worker_failure_table_.reset(new GcsWorkerFailureTable( + std::make_shared>( + redis_client))); + } +}; + +} // namespace gcs +} // namespace ray + +#endif // RAY_GCS_GCS_TABLE_STORAGE_H_ diff --git a/src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc b/src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc index 4dc5ccc88..3697f41ab 100644 --- a/src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc +++ b/src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include +#include #include #include "gtest/gtest.h" diff --git a/src/ray/gcs/gcs_server/test/gcs_actor_scheduler_test.cc b/src/ray/gcs/gcs_server/test/gcs_actor_scheduler_test.cc index 2c4b95c22..49865b4c9 100644 --- a/src/ray/gcs/gcs_server/test/gcs_actor_scheduler_test.cc +++ b/src/ray/gcs/gcs_server/test/gcs_actor_scheduler_test.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include +#include #include #include "gmock/gmock.h" diff --git a/src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc b/src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc index 6f2d96e2e..bb7f7dd39 100644 --- a/src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc +++ b/src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include +#include #include #include "gtest/gtest.h" 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 98aeae1e3..faf9b2386 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 @@ -15,6 +15,7 @@ #include "gtest/gtest.h" #include "ray/common/test_util.h" #include "ray/gcs/gcs_server/gcs_server.h" +#include "ray/gcs/test/gcs_test_util.h" #include "ray/rpc/gcs_server/gcs_rpc_client.h" namespace ray { @@ -382,56 +383,6 @@ class GcsServerTest : public RedisServiceManagerForTest { return status == std::future_status::ready; } - rpc::JobTableData GenJobTableData(JobID job_id) { - rpc::JobTableData job_table_data; - job_table_data.set_job_id(job_id.Binary()); - job_table_data.set_is_dead(false); - job_table_data.set_timestamp(std::time(nullptr)); - job_table_data.set_node_manager_address("127.0.0.1"); - job_table_data.set_driver_pid(5667L); - return job_table_data; - } - - rpc::ActorTableData GenActorTableData(const ActorID &actor_id) { - rpc::ActorTableData actor_table_data; - actor_table_data.set_actor_id(actor_id.Binary()); - actor_table_data.set_job_id(actor_id.JobId().Binary()); - actor_table_data.set_state( - rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE); - actor_table_data.set_max_reconstructions(1); - actor_table_data.set_remaining_reconstructions(1); - return actor_table_data; - } - - rpc::GcsNodeInfo GenGcsNodeInfo(const std::string &node_id) { - rpc::GcsNodeInfo gcs_node_info; - gcs_node_info.set_node_id(node_id); - gcs_node_info.set_node_manager_address("127.0.0.1"); - gcs_node_info.set_node_manager_port(6000); - gcs_node_info.set_state(rpc::GcsNodeInfo_GcsNodeState_ALIVE); - return gcs_node_info; - } - - rpc::TaskTableData GenTaskTableData(const std::string &job_id, - const std::string &task_id) { - rpc::TaskTableData task_table_data; - rpc::Task task; - rpc::TaskSpec task_spec; - task_spec.set_job_id(job_id); - task_spec.set_task_id(task_id); - task.mutable_task_spec()->CopyFrom(task_spec); - task_table_data.mutable_task()->CopyFrom(task); - return task_table_data; - } - - rpc::TaskLeaseData GenTaskLeaseData(const std::string &task_id, - const std::string &node_id) { - rpc::TaskLeaseData task_lease_data; - task_lease_data.set_task_id(task_id); - task_lease_data.set_node_manager_id(node_id); - return task_lease_data; - } - protected: // Gcs server std::unique_ptr gcs_server_; @@ -450,32 +401,31 @@ class GcsServerTest : public RedisServiceManagerForTest { TEST_F(GcsServerTest, TestActorInfo) { // Create actor_table_data JobID job_id = JobID::FromInt(1); - ActorID actor_id = ActorID::Of(job_id, RandomTaskId(), 0); - rpc::ActorTableData actor_table_data = GenActorTableData(actor_id); + auto actor_table_data = Mocker::GenActorTableData(job_id); // Register actor rpc::RegisterActorInfoRequest register_actor_info_request; - register_actor_info_request.mutable_actor_table_data()->CopyFrom(actor_table_data); + register_actor_info_request.mutable_actor_table_data()->CopyFrom(*actor_table_data); ASSERT_TRUE(RegisterActorInfo(register_actor_info_request)); - rpc::ActorTableData result = GetActorInfo(actor_table_data.actor_id()); + rpc::ActorTableData result = GetActorInfo(actor_table_data->actor_id()); ASSERT_TRUE(result.state() == rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE); // Update actor state rpc::UpdateActorInfoRequest update_actor_info_request; - actor_table_data.set_state( + actor_table_data->set_state( rpc::ActorTableData_ActorState::ActorTableData_ActorState_DEAD); - update_actor_info_request.set_actor_id(actor_table_data.actor_id()); - update_actor_info_request.mutable_actor_table_data()->CopyFrom(actor_table_data); + update_actor_info_request.set_actor_id(actor_table_data->actor_id()); + update_actor_info_request.mutable_actor_table_data()->CopyFrom(*actor_table_data); ASSERT_TRUE(UpdateActorInfo(update_actor_info_request)); - result = GetActorInfo(actor_table_data.actor_id()); + result = GetActorInfo(actor_table_data->actor_id()); ASSERT_TRUE(result.state() == rpc::ActorTableData_ActorState::ActorTableData_ActorState_DEAD); // Add actor checkpoint ActorCheckpointID checkpoint_id = ActorCheckpointID::FromRandom(); rpc::ActorCheckpointData checkpoint; - checkpoint.set_actor_id(actor_table_data.actor_id()); + checkpoint.set_actor_id(actor_table_data->actor_id()); checkpoint.set_checkpoint_id(checkpoint_id.Binary()); checkpoint.set_execution_dependency(checkpoint_id.Binary()); @@ -483,39 +433,38 @@ TEST_F(GcsServerTest, TestActorInfo) { add_actor_checkpoint_request.mutable_checkpoint_data()->CopyFrom(checkpoint); ASSERT_TRUE(AddActorCheckpoint(add_actor_checkpoint_request)); rpc::ActorCheckpointData checkpoint_result = - GetActorCheckpoint(actor_id.Binary(), checkpoint_id.Binary()); - ASSERT_TRUE(checkpoint_result.actor_id() == actor_table_data.actor_id()); + GetActorCheckpoint(actor_table_data->actor_id(), checkpoint_id.Binary()); + ASSERT_TRUE(checkpoint_result.actor_id() == actor_table_data->actor_id()); ASSERT_TRUE(checkpoint_result.checkpoint_id() == checkpoint_id.Binary()); rpc::ActorCheckpointIdData checkpoint_id_result = - GetActorCheckpointID(actor_table_data.actor_id()); - ASSERT_TRUE(checkpoint_id_result.actor_id() == actor_table_data.actor_id()); + GetActorCheckpointID(actor_table_data->actor_id()); + ASSERT_TRUE(checkpoint_id_result.actor_id() == actor_table_data->actor_id()); ASSERT_TRUE(checkpoint_id_result.checkpoint_ids_size() == 1); } TEST_F(GcsServerTest, TestJobInfo) { // Create job_table_data JobID job_id = JobID::FromInt(1); - rpc::JobTableData job_table_data = GenJobTableData(job_id); + auto job_table_data = Mocker::GenJobTableData(job_id); // Add job rpc::AddJobRequest add_job_request; - add_job_request.mutable_data()->CopyFrom(job_table_data); + add_job_request.mutable_data()->CopyFrom(*job_table_data); ASSERT_TRUE(AddJob(add_job_request)); // Mark job finished rpc::MarkJobFinishedRequest mark_job_finished_request; - mark_job_finished_request.set_job_id(job_table_data.job_id()); + mark_job_finished_request.set_job_id(job_table_data->job_id()); ASSERT_TRUE(MarkJobFinished(mark_job_finished_request)); } TEST_F(GcsServerTest, TestNodeInfo) { // Create gcs node info - ClientID node_id = ClientID::FromRandom(); - rpc::GcsNodeInfo gcs_node_info = GenGcsNodeInfo(node_id.Binary()); + auto gcs_node_info = Mocker::GenNodeInfo(); // Register node info rpc::RegisterNodeRequest register_node_info_request; - register_node_info_request.mutable_node_info()->CopyFrom(gcs_node_info); + register_node_info_request.mutable_node_info()->CopyFrom(*gcs_node_info); ASSERT_TRUE(RegisterNode(register_node_info_request)); std::vector node_info_list = GetAllNodeInfo(); ASSERT_TRUE(node_info_list.size() == 1); @@ -524,12 +473,12 @@ TEST_F(GcsServerTest, TestNodeInfo) { // Report heartbeat rpc::ReportHeartbeatRequest report_heartbeat_request; - report_heartbeat_request.mutable_heartbeat()->set_client_id(node_id.Binary()); + report_heartbeat_request.mutable_heartbeat()->set_client_id(gcs_node_info->node_id()); ASSERT_TRUE(ReportHeartbeat(report_heartbeat_request)); // Unregister node info rpc::UnregisterNodeRequest unregister_node_info_request; - unregister_node_info_request.set_node_id(node_id.Binary()); + unregister_node_info_request.set_node_id(gcs_node_info->node_id()); ASSERT_TRUE(UnregisterNode(unregister_node_info_request)); node_info_list = GetAllNodeInfo(); ASSERT_TRUE(node_info_list.size() == 1); @@ -538,21 +487,21 @@ TEST_F(GcsServerTest, TestNodeInfo) { // Update node resources rpc::UpdateResourcesRequest update_resources_request; - update_resources_request.set_node_id(node_id.Binary()); + update_resources_request.set_node_id(gcs_node_info->node_id()); rpc::ResourceTableData resource_table_data; resource_table_data.set_resource_capacity(1.0); std::string resource_name = "CPU"; (*update_resources_request.mutable_resources())[resource_name] = resource_table_data; ASSERT_TRUE(UpdateResources(update_resources_request)); - auto resources = GetResources(node_id.Binary()); + auto resources = GetResources(gcs_node_info->node_id()); ASSERT_TRUE(resources.size() == 1); // Delete node resources rpc::DeleteResourcesRequest delete_resources_request; - delete_resources_request.set_node_id(node_id.Binary()); + delete_resources_request.set_node_id(gcs_node_info->node_id()); delete_resources_request.add_resource_name_list(resource_name); ASSERT_TRUE(DeleteResources(delete_resources_request)); - resources = GetResources(node_id.Binary()); + resources = GetResources(gcs_node_info->node_id()); ASSERT_TRUE(resources.empty()); } @@ -591,11 +540,11 @@ TEST_F(GcsServerTest, TestTaskInfo) { // Create task_table_data JobID job_id = JobID::FromInt(1); TaskID task_id = TaskID::ForDriverTask(job_id); - rpc::TaskTableData job_table_data = GenTaskTableData(job_id.Binary(), task_id.Binary()); + auto job_table_data = Mocker::GenTaskTableData(job_id.Binary(), task_id.Binary()); // Add task rpc::AddTaskRequest add_task_request; - add_task_request.mutable_task_data()->CopyFrom(job_table_data); + add_task_request.mutable_task_data()->CopyFrom(*job_table_data); ASSERT_TRUE(AddTask(add_task_request)); rpc::TaskTableData result = GetTask(task_id.Binary()); ASSERT_TRUE(result.task().task_spec().job_id() == job_id.Binary()); @@ -609,10 +558,9 @@ TEST_F(GcsServerTest, TestTaskInfo) { // Add task lease ClientID node_id = ClientID::FromRandom(); - rpc::TaskLeaseData task_lease_data = - GenTaskLeaseData(task_id.Binary(), node_id.Binary()); + auto task_lease_data = Mocker::GenTaskLeaseData(task_id.Binary(), node_id.Binary()); rpc::AddTaskLeaseRequest add_task_lease_request; - add_task_lease_request.mutable_task_lease_data()->CopyFrom(task_lease_data); + add_task_lease_request.mutable_task_lease_data()->CopyFrom(*task_lease_data); ASSERT_TRUE(AddTaskLease(add_task_lease_request)); // Attempt task reconstruction diff --git a/src/ray/gcs/gcs_server/test/redis_gcs_table_storage_test.cc b/src/ray/gcs/gcs_server/test/redis_gcs_table_storage_test.cc new file mode 100644 index 000000000..b9a71372d --- /dev/null +++ b/src/ray/gcs/gcs_server/test/redis_gcs_table_storage_test.cc @@ -0,0 +1,129 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "gtest/gtest.h" +#include "ray/common/test_util.h" +#include "ray/gcs/gcs_server/gcs_table_storage.h" +#include "ray/gcs/store_client/redis_store_client.h" +#include "ray/gcs/store_client/test/store_client_test_base.h" +#include "ray/gcs/test/gcs_test_util.h" + +namespace ray { + +class GcsTableStorageTest : public gcs::StoreClientTestBase { + public: + GcsTableStorageTest() {} + + virtual ~GcsTableStorageTest() {} + + void InitStoreClient() override { + gcs::RedisClientOptions options("127.0.0.1", REDIS_SERVER_PORT, "", true); + redis_client_ = std::make_shared(options); + RAY_CHECK_OK(redis_client_->Connect(io_service_pool_->GetAll())); + + gcs_table_storage_ = std::make_shared(redis_client_); + } + + void DisconnectStoreClient() override { redis_client_->Disconnect(); } + + protected: + template + void Put(TABLE &table, const KEY &key, const VALUE &value) { + auto on_done = [this](Status status) { --pending_count_; }; + ++pending_count_; + RAY_CHECK_OK(table.Put(key, value, on_done)); + WaitPendingDone(); + } + + template + int Get(TABLE &table, const KEY &key, std::vector &values) { + auto on_done = [this, &values](Status status, const boost::optional &result) { + RAY_CHECK_OK(status); + --pending_count_; + values.clear(); + if (result) { + values.push_back(*result); + } + }; + ++pending_count_; + RAY_CHECK_OK(table.Get(key, on_done)); + WaitPendingDone(); + return values.size(); + } + + template + void Delete(TABLE &table, const KEY &key) { + auto on_done = [this](Status status) { + RAY_CHECK_OK(status); + --pending_count_; + }; + ++pending_count_; + RAY_CHECK_OK(table.Delete(key, on_done)); + WaitPendingDone(); + } + + std::shared_ptr redis_client_; + std::shared_ptr gcs_table_storage_; +}; + +TEST_F(GcsTableStorageTest, TestGcsTableApi) { + auto table = gcs_table_storage_->JobTable(); + JobID job1_id = JobID::FromInt(1); + JobID job2_id = JobID::FromInt(2); + auto job1_table_data = Mocker::GenJobTableData(job1_id); + auto job2_table_data = Mocker::GenJobTableData(job2_id); + + // Put. + Put(table, job1_id, *job1_table_data); + Put(table, job2_id, *job2_table_data); + + // Get. + std::vector values; + ASSERT_EQ(Get(table, job2_id, values), 1); + ASSERT_EQ(Get(table, job2_id, values), 1); + + // Delete. + Delete(table, job1_id); + ASSERT_EQ(Get(table, job1_id, values), 0); + ASSERT_EQ(Get(table, job2_id, values), 1); +} + +TEST_F(GcsTableStorageTest, TestGcsTableWithJobIdApi) { + auto table = gcs_table_storage_->ActorTable(); + JobID job_id = JobID::FromInt(3); + auto actor_table_data = Mocker::GenActorTableData(job_id); + ActorID actor_id = ActorID::FromBinary(actor_table_data->actor_id()); + + // Put. + Put(table, actor_id, *actor_table_data); + + // Get. + std::vector values; + ASSERT_EQ(Get(table, actor_id, values), 1); + + // Delete. + Delete(table, actor_id); + ASSERT_EQ(Get(table, actor_id, values), 0); +} + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + RAY_CHECK(argc == 4); + ray::REDIS_SERVER_EXEC_PATH = argv[1]; + ray::REDIS_CLIENT_EXEC_PATH = argv[2]; + ray::REDIS_MODULE_LIBRARY_PATH = argv[3]; + return RUN_ALL_TESTS(); +} diff --git a/src/ray/gcs/store_client/redis_store_client.cc b/src/ray/gcs/store_client/redis_store_client.cc index ffea3ff2f..b1959654f 100644 --- a/src/ray/gcs/store_client/redis_store_client.cc +++ b/src/ray/gcs/store_client/redis_store_client.cc @@ -145,6 +145,20 @@ Status RedisStoreClient::AsyncDeleteByIndex( } template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; +template class RedisStoreClient; } // namespace gcs diff --git a/src/ray/gcs/gcs_server/test/gcs_test_util.h b/src/ray/gcs/test/gcs_test_util.h similarity index 82% rename from src/ray/gcs/gcs_server/test/gcs_test_util.h rename to src/ray/gcs/test/gcs_test_util.h index 158cf9e13..9a6a837bf 100644 --- a/src/ray/gcs/gcs_server/test/gcs_test_util.h +++ b/src/ray/gcs/test/gcs_test_util.h @@ -15,13 +15,13 @@ #ifndef RAY_GCS_TEST_UTIL_H #define RAY_GCS_TEST_UTIL_H -#include -#include -#include -#include -#include -#include -#include +#include "src/ray/common/task/task.h" +#include "src/ray/common/task/task_util.h" +#include "src/ray/common/test_util.h" +#include "src/ray/gcs/gcs_server/gcs_actor_manager.h" +#include "src/ray/gcs/gcs_server/gcs_actor_scheduler.h" +#include "src/ray/gcs/gcs_server/gcs_node_manager.h" +#include "src/ray/util/asio_util.h" #include #include @@ -59,6 +59,67 @@ struct Mocker { return node; } + static std::shared_ptr GenJobTableData(JobID job_id) { + auto job_table_data = std::make_shared(); + job_table_data->set_job_id(job_id.Binary()); + job_table_data->set_is_dead(false); + job_table_data->set_timestamp(std::time(nullptr)); + job_table_data->set_node_manager_address("127.0.0.1"); + job_table_data->set_driver_pid(5667L); + return job_table_data; + } + + static std::shared_ptr GenActorTableData(const JobID &job_id) { + auto actor_table_data = std::make_shared(); + ActorID actor_id = ActorID::Of(job_id, RandomTaskId(), 0); + actor_table_data->set_actor_id(actor_id.Binary()); + actor_table_data->set_job_id(job_id.Binary()); + actor_table_data->set_state( + rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE); + actor_table_data->set_max_reconstructions(1); + actor_table_data->set_remaining_reconstructions(1); + return actor_table_data; + } + + static std::shared_ptr GenTaskTableData( + const std::string &job_id, const std::string &task_id) { + auto task_table_data = std::make_shared(); + rpc::Task task; + rpc::TaskSpec task_spec; + task_spec.set_job_id(job_id); + task_spec.set_task_id(task_id); + task.mutable_task_spec()->CopyFrom(task_spec); + task_table_data->mutable_task()->CopyFrom(task); + return task_table_data; + } + + static std::shared_ptr GenTaskLeaseData( + const std::string &task_id, const std::string &node_id) { + auto task_lease_data = std::make_shared(); + task_lease_data->set_task_id(task_id); + task_lease_data->set_node_manager_id(node_id); + return task_lease_data; + } + + static 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; + } + + static 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; + } + + static std::shared_ptr GenWorkerFailureData() { + auto worker_failure_data = std::make_shared(); + worker_failure_data->set_timestamp(std::time(nullptr)); + return worker_failure_data; + } + class MockWorkerClient : public rpc::CoreWorkerClientInterface { public: ray::Status PushNormalTask( diff --git a/src/ray/protobuf/gcs.proto b/src/ray/protobuf/gcs.proto index e0dcb0997..5b9b3e261 100644 --- a/src/ray/protobuf/gcs.proto +++ b/src/ray/protobuf/gcs.proto @@ -303,6 +303,13 @@ message WorkerFailureData { bool intentional_disconnect = 4; } +message ResourceMap { + map items = 1; +} + +message ObjectTableDataList { + repeated ObjectTableData items = 1; +} // This enum type is used as object's metadata to indicate the object's creating // task has failed because of a certain error. // TODO(hchen): We may want to make these errors more specific. E.g., we may want