diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.cc b/src/ray/gcs/gcs_server/gcs_table_storage.cc index 2b41a4edc..ed0931306 100644 --- a/src/ray/gcs/gcs_server/gcs_table_storage.cc +++ b/src/ray/gcs/gcs_server/gcs_table_storage.cc @@ -23,24 +23,46 @@ namespace gcs { template Status GcsTable::Put(const Key &key, const Data &value, const StatusCallback &callback) { - return store_client_->AsyncPut(table_name_, key, value, callback); + return store_client_->AsyncPut(table_name_, key.Binary(), value.SerializeAsString(), + callback); } template Status GcsTable::Get(const Key &key, const OptionalItemCallback &callback) { - return store_client_->AsyncGet(table_name_, key, callback); + auto on_done = [callback](const Status &status, + const boost::optional &result) { + boost::optional value; + if (result) { + Data data; + data.ParseFromString(*result); + value = std::move(data); + } + callback(status, value); + }; + return store_client_->AsyncGet(table_name_, key.Binary(), on_done); } template Status GcsTable::GetAll( const SegmentedCallback> &callback) { - return store_client_->AsyncGetAll(table_name_, callback); + auto on_done = [callback]( + const Status &status, bool has_more, + const std::vector> &result) { + std::vector> values; + for (auto &item : result) { + Data data; + data.ParseFromString(item.second); + values.emplace_back(std::move(std::make_pair(Key::FromBinary(item.first), data))); + } + callback(status, has_more, values); + }; + return store_client_->AsyncGetAll(table_name_, on_done); } template Status GcsTable::Delete(const Key &key, const StatusCallback &callback) { - return store_client_->AsyncDelete(table_name_, key, callback); + return store_client_->AsyncDelete(table_name_, key.Binary(), callback); } template @@ -50,13 +72,13 @@ Status GcsTable::BatchDelete(const std::vector &keys, auto finished_count = std::make_shared(0); int size = keys.size(); for (Key key : keys) { - auto done = [finished_count, size, callback](Status status) { + auto done = [finished_count, size, callback](const Status &status) { ++(*finished_count); if (*finished_count == size) { callback(Status::OK()); } }; - RAY_CHECK_OK(store_client_->AsyncDelete(table_name_, key, done)); + RAY_CHECK_OK(store_client_->AsyncDelete(table_name_, key.Binary(), done)); } return Status::OK(); } @@ -64,8 +86,9 @@ Status GcsTable::BatchDelete(const std::vector &keys, 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); + return this->store_client_->AsyncPutWithIndex(this->table_name_, key.Binary(), + GetJobIdFromKey(key).Binary(), + value.SerializeAsString(), callback); } template @@ -79,7 +102,8 @@ Status GcsTableWithJobId::GetByJobId( template Status GcsTableWithJobId::DeleteByJobId(const JobID &job_id, const StatusCallback &callback) { - return this->store_client_->AsyncDeleteByIndex(this->table_name_, job_id, callback); + return this->store_client_->AsyncDeleteByIndex(this->table_name_, job_id.Binary(), + callback); } template class GcsTable; diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.h b/src/ray/gcs/gcs_server/gcs_table_storage.h index 1a69857ce..3ef728779 100644 --- a/src/ray/gcs/gcs_server/gcs_table_storage.h +++ b/src/ray/gcs/gcs_server/gcs_table_storage.h @@ -50,7 +50,7 @@ using rpc::WorkerFailureData; template class GcsTable { public: - explicit GcsTable(std::shared_ptr> store_client) + explicit GcsTable(std::shared_ptr &store_client) : store_client_(store_client) {} virtual ~GcsTable() = default; @@ -93,7 +93,7 @@ class GcsTable { protected: std::string table_name_; - std::shared_ptr> store_client_; + std::shared_ptr store_client_; }; /// \class GcsTableWithJobId @@ -105,7 +105,7 @@ class GcsTable { template class GcsTableWithJobId : public GcsTable { public: - explicit GcsTableWithJobId(std::shared_ptr> store_client) + explicit GcsTableWithJobId(std::shared_ptr &store_client) : GcsTable(store_client) {} /// Write data to the table asynchronously. @@ -137,31 +137,27 @@ class GcsTableWithJobId : public GcsTable { class GcsJobTable : public GcsTable { public: - explicit GcsJobTable( - std::shared_ptr> store_client) - : GcsTable(std::move(store_client)) { + explicit GcsJobTable(std::shared_ptr &store_client) + : GcsTable(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)) { + explicit GcsActorTable(std::shared_ptr &store_client) + : GcsTableWithJobId(store_client) { table_name_ = TablePrefix_Name(TablePrefix::ACTOR); } private: - JobID GetJobIdFromKey(const ActorID &key) { return key.JobId(); } + JobID GetJobIdFromKey(const ActorID &key) override { return key.JobId(); } }; class GcsActorCheckpointTable : public GcsTable { public: - explicit GcsActorCheckpointTable( - std::shared_ptr> - store_client) - : GcsTable(std::move(store_client)) { + explicit GcsActorCheckpointTable(std::shared_ptr &store_client) + : GcsTable(store_client) { table_name_ = TablePrefix_Name(TablePrefix::ACTOR_CHECKPOINT); } }; @@ -169,124 +165,112 @@ class GcsActorCheckpointTable : public GcsTable { public: - explicit GcsActorCheckpointIdTable( - std::shared_ptr> store_client) - : GcsTableWithJobId(std::move(store_client)) { + explicit GcsActorCheckpointIdTable(std::shared_ptr &store_client) + : GcsTableWithJobId(store_client) { table_name_ = TablePrefix_Name(TablePrefix::ACTOR_CHECKPOINT_ID); } private: - JobID GetJobIdFromKey(const ActorID &key) { return key.JobId(); } + JobID GetJobIdFromKey(const ActorID &key) override { return key.JobId(); } }; class GcsTaskTable : public GcsTableWithJobId { public: - explicit GcsTaskTable( - std::shared_ptr> store_client) - : GcsTableWithJobId(std::move(store_client)) { + explicit GcsTaskTable(std::shared_ptr &store_client) + : GcsTableWithJobId(store_client) { table_name_ = TablePrefix_Name(TablePrefix::TASK); } private: - JobID GetJobIdFromKey(const TaskID &key) { return key.ActorId().JobId(); } + JobID GetJobIdFromKey(const TaskID &key) override { return key.ActorId().JobId(); } }; class GcsTaskLeaseTable : public GcsTableWithJobId { public: - explicit GcsTaskLeaseTable( - std::shared_ptr> store_client) - : GcsTableWithJobId(std::move(store_client)) { + explicit GcsTaskLeaseTable(std::shared_ptr &store_client) + : GcsTableWithJobId(store_client) { table_name_ = TablePrefix_Name(TablePrefix::TASK_LEASE); } private: - JobID GetJobIdFromKey(const TaskID &key) { return key.ActorId().JobId(); } + JobID GetJobIdFromKey(const TaskID &key) override { return key.ActorId().JobId(); } }; class GcsTaskReconstructionTable : public GcsTableWithJobId { public: - explicit GcsTaskReconstructionTable( - std::shared_ptr> store_client) - : GcsTableWithJobId(std::move(store_client)) { + explicit GcsTaskReconstructionTable(std::shared_ptr &store_client) + : GcsTableWithJobId(store_client) { table_name_ = TablePrefix_Name(TablePrefix::TASK_RECONSTRUCTION); } private: - JobID GetJobIdFromKey(const TaskID &key) { return key.ActorId().JobId(); } + JobID GetJobIdFromKey(const TaskID &key) override { return key.ActorId().JobId(); } }; class GcsObjectTable : public GcsTableWithJobId { public: - explicit GcsObjectTable( - std::shared_ptr> store_client) - : GcsTableWithJobId(std::move(store_client)) { + explicit GcsObjectTable(std::shared_ptr &store_client) + : GcsTableWithJobId(store_client) { table_name_ = TablePrefix_Name(TablePrefix::OBJECT); } private: - JobID GetJobIdFromKey(const ObjectID &key) { return key.TaskId().JobId(); } + JobID GetJobIdFromKey(const ObjectID &key) override { return key.TaskId().JobId(); } }; class GcsNodeTable : public GcsTable { public: - explicit GcsNodeTable( - std::shared_ptr> store_client) - : GcsTable(std::move(store_client)) { + explicit GcsNodeTable(std::shared_ptr &store_client) + : GcsTable(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)) { + explicit GcsNodeResourceTable(std::shared_ptr &store_client) + : GcsTable(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)) { + explicit GcsHeartbeatTable(std::shared_ptr &store_client) + : GcsTable(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)) { + explicit GcsHeartbeatBatchTable(std::shared_ptr &store_client) + : GcsTable(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)) { + explicit GcsErrorInfoTable(std::shared_ptr &store_client) + : GcsTable(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)) { + explicit GcsProfileTable(std::shared_ptr &store_client) + : GcsTable(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)) { + explicit GcsWorkerFailureTable(std::shared_ptr &store_client) + : GcsTable(store_client) { table_name_ = TablePrefix_Name(TablePrefix::WORKER_FAILURE); } }; @@ -373,6 +357,7 @@ class GcsTableStorage { } protected: + std::shared_ptr store_client_; std::unique_ptr job_table_; std::unique_ptr actor_table_; std::unique_ptr actor_checkpoint_table_; @@ -396,45 +381,22 @@ class GcsTableStorage { 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))); + store_client_ = std::make_shared(redis_client); + job_table_.reset(new GcsJobTable(store_client_)); + actor_table_.reset(new GcsActorTable(store_client_)); + actor_checkpoint_table_.reset(new GcsActorCheckpointTable(store_client_)); + actor_checkpoint_id_table_.reset(new GcsActorCheckpointIdTable(store_client_)); + task_table_.reset(new GcsTaskTable(store_client_)); + task_lease_table_.reset(new GcsTaskLeaseTable(store_client_)); + task_reconstruction_table_.reset(new GcsTaskReconstructionTable(store_client_)); + object_table_.reset(new GcsObjectTable(store_client_)); + node_table_.reset(new GcsNodeTable(store_client_)); + node_resource_table_.reset(new GcsNodeResourceTable(store_client_)); + heartbeat_table_.reset(new GcsHeartbeatTable(store_client_)); + heartbeat_batch_table_.reset(new GcsHeartbeatBatchTable(store_client_)); + error_info_table_.reset(new GcsErrorInfoTable(store_client_)); + profile_table_.reset(new GcsProfileTable(store_client_)); + worker_failure_table_.reset(new GcsWorkerFailureTable(store_client_)); } }; diff --git a/src/ray/gcs/store_client/redis_store_client.cc b/src/ray/gcs/store_client/redis_store_client.cc index b1959654f..cb20dc59f 100644 --- a/src/ray/gcs/store_client/redis_store_client.cc +++ b/src/ray/gcs/store_client/redis_store_client.cc @@ -23,24 +23,19 @@ namespace ray { namespace gcs { -template -Status RedisStoreClient::AsyncPut(const std::string &table_name, - const Key &key, const Data &data, - const StatusCallback &callback) { - std::string full_key = table_name + key.Binary(); - std::string data_str = data.SerializeAsString(); - RAY_CHECK(!data_str.empty()); - return DoPut(full_key, data_str, callback); +Status RedisStoreClient::AsyncPut(const std::string &table_name, const std::string &key, + const std::string &data, + const StatusCallback &callback) { + std::string full_key = table_name + key; + return DoPut(full_key, data, callback); } -template -Status RedisStoreClient::AsyncPutWithIndex( - const std::string &table_name, const Key &key, const IndexKey &index_key, - const Data &data, const StatusCallback &callback) { - std::string data_str = data.SerializeAsString(); - RAY_CHECK(!data_str.empty()); - - auto write_callback = [this, table_name, key, data_str, callback](Status status) { +Status RedisStoreClient::AsyncPutWithIndex(const std::string &table_name, + const std::string &key, + const std::string &index_key, + const std::string &data, + const StatusCallback &callback) { + auto write_callback = [this, table_name, key, data, callback](Status status) { if (!status.ok()) { // Run callback if failed. if (callback != nullptr) { @@ -50,8 +45,8 @@ Status RedisStoreClient::AsyncPutWithIndex( } // Write data to Redis. - std::string full_key = table_name + key.Binary(); - status = DoPut(full_key, data_str, callback); + std::string full_key = table_name + key; + status = DoPut(full_key, data, callback); if (!status.ok()) { // Run callback if failed. @@ -62,18 +57,16 @@ Status RedisStoreClient::AsyncPutWithIndex( }; // Write index to Redis. - std::string index_table_key = index_key.Binary() + table_name + key.Binary(); - return DoPut(index_table_key, key.Binary(), write_callback); + std::string index_table_key = index_key + table_name + key; + return DoPut(index_table_key, key, write_callback); } -template -Status RedisStoreClient::DoPut(const std::string &key, - const std::string &data, - const StatusCallback &callback) { +Status RedisStoreClient::DoPut(const std::string &key, const std::string &data, + const StatusCallback &callback) { std::vector args = {"SET", key, data}; RedisCallback write_callback = nullptr; if (callback) { - write_callback = [callback](std::shared_ptr reply) { + write_callback = [callback](const std::shared_ptr &reply) { auto status = reply->ReadAsStatus(); callback(status); }; @@ -83,43 +76,38 @@ Status RedisStoreClient::DoPut(const std::string &key, return shard_context->RunArgvAsync(args, write_callback); } -template -Status RedisStoreClient::AsyncGet( - const std::string &table_name, const Key &key, - const OptionalItemCallback &callback) { +Status RedisStoreClient::AsyncGet(const std::string &table_name, const std::string &key, + const OptionalItemCallback &callback) { RAY_CHECK(callback != nullptr); - auto redis_callback = [callback](std::shared_ptr reply) { - boost::optional result; + auto redis_callback = [callback](const std::shared_ptr &reply) { + boost::optional result; if (!reply->IsNil()) { - std::string data_str = reply->ReadAsString(); - if (!data_str.empty()) { - Data data; - RAY_CHECK(data.ParseFromString(data_str)); + std::string data = reply->ReadAsString(); + if (!data.empty()) { result = std::move(data); } } callback(Status::OK(), result); }; - std::string full_key = table_name + key.Binary(); + std::string full_key = table_name + key; std::vector args = {"GET", full_key}; auto shard_context = redis_client_->GetShardContext(full_key); return shard_context->RunArgvAsync(args, redis_callback); } -template -Status RedisStoreClient::AsyncGetAll( +Status RedisStoreClient::AsyncGetAll( const std::string &table_name, - const SegmentedCallback> &callback) { + const SegmentedCallback> &callback) { RAY_CHECK(0) << "Not implemented! Will implement this function in next PR."; return Status::OK(); } -template -Status RedisStoreClient::AsyncDelete( - const std::string &table_name, const Key &key, const StatusCallback &callback) { +Status RedisStoreClient::AsyncDelete(const std::string &table_name, + const std::string &key, + const StatusCallback &callback) { RedisCallback delete_callback = nullptr; if (callback) { delete_callback = [callback](std::shared_ptr reply) { @@ -129,37 +117,20 @@ Status RedisStoreClient::AsyncDelete( }; } - std::string full_key = table_name + key.Binary(); + std::string full_key = table_name + key; std::vector args = {"DEL", full_key}; auto shard_context = redis_client_->GetShardContext(full_key); return shard_context->RunArgvAsync(args, delete_callback); } -template -Status RedisStoreClient::AsyncDeleteByIndex( - const std::string &table_name, const IndexKey &index_key, - const StatusCallback &callback) { +Status RedisStoreClient::AsyncDeleteByIndex(const std::string &table_name, + const std::string &index_key, + const StatusCallback &callback) { RAY_CHECK(0) << "Not implemented! Will implement this function in next PR."; return Status::OK(); } -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 } // namespace ray diff --git a/src/ray/gcs/store_client/redis_store_client.h b/src/ray/gcs/store_client/redis_store_client.h index 25a2a9677..200288995 100644 --- a/src/ray/gcs/store_client/redis_store_client.h +++ b/src/ray/gcs/store_client/redis_store_client.h @@ -26,31 +26,29 @@ namespace ray { namespace gcs { -template -class RedisStoreClient : public StoreClient { +class RedisStoreClient : public StoreClient { public: RedisStoreClient(std::shared_ptr redis_client) : redis_client_(std::move(redis_client)) {} - virtual ~RedisStoreClient() {} + Status AsyncPut(const std::string &table_name, const std::string &key, + const std::string &data, const StatusCallback &callback) override; - Status AsyncPut(const std::string &table_name, const Key &key, const Data &data, - const StatusCallback &callback) override; - - Status AsyncPutWithIndex(const std::string &table_name, const Key &key, - const IndexKey &index_key, const Data &data, + Status AsyncPutWithIndex(const std::string &table_name, const std::string &key, + const std::string &index_key, const std::string &data, const StatusCallback &callback) override; - Status AsyncGet(const std::string &table_name, const Key &key, - const OptionalItemCallback &callback) override; + Status AsyncGet(const std::string &table_name, const std::string &key, + const OptionalItemCallback &callback) override; - Status AsyncGetAll(const std::string &table_name, - const SegmentedCallback> &callback) override; + Status AsyncGetAll( + const std::string &table_name, + const SegmentedCallback> &callback) override; - Status AsyncDelete(const std::string &table_name, const Key &key, + Status AsyncDelete(const std::string &table_name, const std::string &key, const StatusCallback &callback) override; - Status AsyncDeleteByIndex(const std::string &table_name, const IndexKey &index_key, + Status AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key, const StatusCallback &callback) override; private: @@ -60,8 +58,6 @@ class RedisStoreClient : public StoreClient { std::shared_ptr redis_client_; }; -typedef RedisStoreClient RedisActorStoreTable; - } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/store_client/store_client.h b/src/ray/gcs/store_client/store_client.h index 941ac16b2..d7e22ddf0 100644 --- a/src/ray/gcs/store_client/store_client.h +++ b/src/ray/gcs/store_client/store_client.h @@ -30,10 +30,9 @@ namespace gcs { /// \class StoreClient /// Abstract interface of the storage client. -template class StoreClient { public: - virtual ~StoreClient() {} + virtual ~StoreClient() = default; /// Write data to the given table asynchronously. /// @@ -42,8 +41,8 @@ class StoreClient { /// \param data 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 AsyncPut(const std::string &table_name, const Key &key, const Data &data, - const StatusCallback &callback) = 0; + virtual Status AsyncPut(const std::string &table_name, const std::string &key, + const std::string &data, const StatusCallback &callback) = 0; /// Write data to the given table asynchronously. /// @@ -53,8 +52,8 @@ class StoreClient { /// \param data 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 AsyncPutWithIndex(const std::string &table_name, const Key &key, - const IndexKey &index_key, const Data &data, + virtual Status AsyncPutWithIndex(const std::string &table_name, const std::string &key, + const std::string &index_key, const std::string &data, const StatusCallback &callback) = 0; /// Get data from the given table asynchronously. @@ -63,8 +62,8 @@ class StoreClient { /// \param key The key to lookup from the table. /// \param callback Callback that will be called after read finishes. /// \return Status - virtual Status AsyncGet(const std::string &table_name, const Key &key, - const OptionalItemCallback &callback) = 0; + virtual Status AsyncGet(const std::string &table_name, const std::string &key, + const OptionalItemCallback &callback) = 0; /// Get all data from the given table asynchronously. /// @@ -72,8 +71,9 @@ class StoreClient { /// \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 - virtual Status AsyncGetAll(const std::string &table_name, - const SegmentedCallback> &callback) = 0; + virtual Status AsyncGetAll( + const std::string &table_name, + const SegmentedCallback> &callback) = 0; /// Delete data from the given table asynchronously. /// @@ -81,7 +81,7 @@ class StoreClient { /// \param key The key that will be deleted from the table. /// \param callback Callback that will be called after delete finishes. /// \return Status - virtual Status AsyncDelete(const std::string &table_name, const Key &key, + virtual Status AsyncDelete(const std::string &table_name, const std::string &key, const StatusCallback &callback) = 0; /// Delete by index from the given table asynchronously. @@ -92,15 +92,13 @@ class StoreClient { /// \param callback Callback that will be called after delete finishes. /// \return Status virtual Status AsyncDeleteByIndex(const std::string &table_name, - const IndexKey &index_key, + const std::string &index_key, const StatusCallback &callback) = 0; protected: StoreClient() = default; }; -typedef StoreClient ActorStoreTable; - } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/store_client/test/redis_store_client_test.cc b/src/ray/gcs/store_client/test/redis_store_client_test.cc index 9ccc641af..24f705d74 100644 --- a/src/ray/gcs/store_client/test/redis_store_client_test.cc +++ b/src/ray/gcs/store_client/test/redis_store_client_test.cc @@ -31,9 +31,7 @@ class RedisStoreClientTest : public StoreClientTestBase { redis_client_ = std::make_shared(options); RAY_CHECK_OK(redis_client_->Connect(io_service_pool_->GetAll())); - store_client_ = - std::make_shared>( - redis_client_); + store_client_ = std::make_shared(redis_client_); } void DisconnectStoreClient() override { redis_client_->Disconnect(); } @@ -44,96 +42,94 @@ class RedisStoreClientTest : public StoreClientTestBase { TEST_F(RedisStoreClientTest, AsyncPutAndAsyncGetTest) { // AsyncPut without index. - auto put_calllback = [this](Status status) { + auto put_calllback = [this](const Status &status) { RAY_CHECK_OK(status); --pending_count_; }; for (const auto &elem : key_to_value_) { ++pending_count_; - Status status = - store_client_->AsyncPut(table_name_, elem.first, elem.second, put_calllback); - RAY_CHECK_OK(status); + RAY_CHECK_OK(store_client_->AsyncPut(table_name_, elem.first.Binary(), + elem.second.SerializeAsString(), put_calllback)); } WaitPendingDone(); // AsyncGet - auto get_callback = [this](Status status, - const boost::optional &result) { + auto get_callback = [this](const Status &status, + const boost::optional &result) { RAY_CHECK_OK(status); RAY_CHECK(result); - ActorID actor_id = ActorID::FromBinary(result->actor_id()); + rpc::ActorTableData data; + RAY_CHECK(data.ParseFromString(*result)); + ActorID actor_id = ActorID::FromBinary(data.actor_id()); auto it = key_to_value_.find(actor_id); RAY_CHECK(it != key_to_value_.end()); --pending_count_; }; for (const auto &elem : key_to_value_) { ++pending_count_; - Status status = store_client_->AsyncGet(table_name_, elem.first, get_callback); - RAY_CHECK_OK(status); + RAY_CHECK_OK(store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback)); } WaitPendingDone(); } TEST_F(RedisStoreClientTest, AsyncDeleteTest) { // AsyncPut - auto put_calllback = [this](Status status) { --pending_count_; }; + auto put_calllback = [this](const Status &status) { --pending_count_; }; for (const auto &elem : key_to_value_) { ++pending_count_; - Status status = - store_client_->AsyncPut(table_name_, elem.first, elem.second, put_calllback); - RAY_CHECK_OK(status); + RAY_CHECK_OK(store_client_->AsyncPut(table_name_, elem.first.Binary(), + elem.second.SerializeAsString(), put_calllback)); } WaitPendingDone(); // AsyncDelete - auto delete_calllback = [this](Status status) { + auto delete_calllback = [this](const Status &status) { RAY_CHECK_OK(status); --pending_count_; }; for (const auto &elem : key_to_value_) { ++pending_count_; - Status status = store_client_->AsyncDelete(table_name_, elem.first, delete_calllback); - RAY_CHECK_OK(status); + RAY_CHECK_OK( + store_client_->AsyncDelete(table_name_, elem.first.Binary(), delete_calllback)); } WaitPendingDone(); // AsyncGet - auto get_callback = [this](Status status, - const boost::optional &result) { + auto get_callback = [this](const Status &status, + const boost::optional &result) { RAY_CHECK_OK(status); RAY_CHECK(!result); --pending_count_; }; for (const auto &elem : key_to_value_) { ++pending_count_; - Status status = store_client_->AsyncGet(table_name_, elem.first, get_callback); - RAY_CHECK_OK(status); + RAY_CHECK_OK(store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback)); } WaitPendingDone(); } TEST_F(RedisStoreClientTest, DISABLED_AsyncGetAllTest) { // AsyncPut - auto put_calllback = [this](Status status) { --pending_count_; }; + auto put_calllback = [this](const Status &status) { --pending_count_; }; for (const auto &elem : key_to_value_) { ++pending_count_; // Get index auto it = key_to_index_.find(elem.first); const JobID &index = it->second; - Status status = store_client_->AsyncPutWithIndex(table_name_, elem.first, index, - elem.second, put_calllback); - RAY_CHECK_OK(status); + RAY_CHECK_OK( + store_client_->AsyncPutWithIndex(table_name_, elem.first.Binary(), index.Binary(), + elem.second.SerializeAsString(), put_calllback)); } WaitPendingDone(); // AsyncGetAll auto get_all_callback = - [this](Status status, bool has_more, - const std::vector> &result) { + [this](const Status &status, bool has_more, + const std::vector> &result) { RAY_CHECK_OK(status); static std::unordered_set received_keys; for (const auto &item : result) { - const ActorID &actor_id = item.first; + const ActorID &actor_id = ActorID::FromBinary(item.first); auto it = received_keys.find(actor_id); RAY_CHECK(it == received_keys.end()); received_keys.emplace(actor_id); @@ -148,8 +144,7 @@ TEST_F(RedisStoreClientTest, DISABLED_AsyncGetAllTest) { }; pending_count_ += key_to_value_.size(); - Status status = store_client_->AsyncGetAll(table_name_, get_all_callback); - RAY_CHECK_OK(status); + RAY_CHECK_OK(store_client_->AsyncGetAll(table_name_, get_all_callback)); WaitPendingDone(); } diff --git a/src/ray/gcs/store_client/test/store_client_test_base.h b/src/ray/gcs/store_client/test/store_client_test_base.h index 5495b79d8..0d23734ff 100644 --- a/src/ray/gcs/store_client/test/store_client_test_base.h +++ b/src/ray/gcs/store_client/test/store_client_test_base.h @@ -94,7 +94,7 @@ class StoreClientTestBase : public RedisServiceManagerForTest { size_t io_service_num_{2}; std::shared_ptr io_service_pool_; - std::shared_ptr> store_client_; + std::shared_ptr store_client_; std::string table_name_{"test_table"}; size_t key_count_{5000};