From 5bff707d20d8058285956da38d938cff597afa05 Mon Sep 17 00:00:00 2001 From: fangfengbin <869218239a@zju.edu.cn> Date: Sun, 26 Apr 2020 19:09:26 +0800 Subject: [PATCH] [GCS]Add in-memory store client (#8144) --- BUILD.bazel | 28 ++++ .../test/redis_gcs_table_storage_test.cc | 4 + .../store_client/in_memory_store_client.cc | 113 +++++++++++++ .../gcs/store_client/in_memory_store_client.h | 83 ++++++++++ .../test/in_memory_store_client_test.cc | 48 ++++++ .../test/redis_store_client_test.cc | 112 +------------ .../test/store_client_test_base.h | 152 +++++++++++++++++- 7 files changed, 434 insertions(+), 106 deletions(-) create mode 100644 src/ray/gcs/store_client/in_memory_store_client.cc create mode 100644 src/ray/gcs/store_client/in_memory_store_client.h create mode 100644 src/ray/gcs/store_client/test/in_memory_store_client_test.cc diff --git a/BUILD.bazel b/BUILD.bazel index 49c4449de..1eb15c48c 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1086,6 +1086,23 @@ cc_library( ], ) +cc_library( + name = "gcs_in_memory_store_client", + srcs = [ + "src/ray/gcs/store_client/in_memory_store_client.cc", + ], + hdrs = [ + "src/ray/gcs/callback.h", + "src/ray/gcs/store_client/in_memory_store_client.h", + "src/ray/gcs/store_client/store_client.h", + ], + copts = COPTS, + deps = [ + ":ray_common", + ":ray_util", + ], +) + cc_library( name = "store_client_test_lib", hdrs = [ @@ -1113,6 +1130,17 @@ cc_test( ], ) +cc_test( + name = "in_memory_store_client_test", + srcs = ["src/ray/gcs/store_client/test/in_memory_store_client_test.cc"], + copts = COPTS, + deps = [ + ":gcs_in_memory_store_client", + ":store_client_test_lib", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "gcs", srcs = glob( 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 index b9a71372d..588550bf3 100644 --- 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 @@ -27,6 +27,10 @@ class GcsTableStorageTest : public gcs::StoreClientTestBase { virtual ~GcsTableStorageTest() {} + static void SetUpTestCase() { RedisServiceManagerForTest::SetUpTestCase(); } + + static void TearDownTestCase() { RedisServiceManagerForTest::TearDownTestCase(); } + void InitStoreClient() override { gcs::RedisClientOptions options("127.0.0.1", REDIS_SERVER_PORT, "", true); redis_client_ = std::make_shared(options); diff --git a/src/ray/gcs/store_client/in_memory_store_client.cc b/src/ray/gcs/store_client/in_memory_store_client.cc new file mode 100644 index 000000000..2a001141d --- /dev/null +++ b/src/ray/gcs/store_client/in_memory_store_client.cc @@ -0,0 +1,113 @@ +// 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 "ray/gcs/store_client/in_memory_store_client.h" + +namespace ray { + +namespace gcs { + +Status InMemoryStoreClient::AsyncPut(const std::string &table_name, + const std::string &key, const std::string &data, + const StatusCallback &callback) { + auto table = GetOrCreateTable(table_name); + absl::MutexLock lock(&(table->mutex_)); + table->records_[key] = data; + main_io_service_.post([callback]() { callback(Status::OK()); }); + return Status::OK(); +} + +Status InMemoryStoreClient::AsyncPutWithIndex(const std::string &table_name, + const std::string &key, + const std::string &index_key, + const std::string &data, + const StatusCallback &callback) { + auto table = GetOrCreateTable(table_name); + absl::MutexLock lock(&(table->mutex_)); + table->records_[key] = data; + table->index_keys_[index_key].emplace_back(key); + main_io_service_.post([callback]() { callback(Status::OK()); }); + return Status::OK(); +} + +Status InMemoryStoreClient::AsyncGet(const std::string &table_name, + const std::string &key, + const OptionalItemCallback &callback) { + auto table = GetOrCreateTable(table_name); + absl::MutexLock lock(&(table->mutex_)); + auto iter = table->records_.find(key); + if (iter != table->records_.end()) { + auto data = iter->second; + main_io_service_.post([callback, data]() { callback(Status::OK(), data); }); + } else { + main_io_service_.post([callback]() { callback(Status::OK(), boost::none); }); + } + return Status::OK(); +} + +Status InMemoryStoreClient::AsyncGetAll( + const std::string &table_name, + const SegmentedCallback> &callback) { + auto table = GetOrCreateTable(table_name); + absl::MutexLock lock(&(table->mutex_)); + std::vector> result; + for (auto &record : table->records_) { + result.emplace_back(std::make_pair(record.first, record.second)); + } + main_io_service_.post([result, callback]() { callback(Status::OK(), false, result); }); + return Status::OK(); +} + +Status InMemoryStoreClient::AsyncDelete(const std::string &table_name, + const std::string &key, + const StatusCallback &callback) { + auto table = GetOrCreateTable(table_name); + absl::MutexLock lock(&(table->mutex_)); + table->records_.erase(key); + main_io_service_.post([callback]() { callback(Status::OK()); }); + return Status::OK(); +} + +Status InMemoryStoreClient::AsyncDeleteByIndex(const std::string &table_name, + const std::string &index_key, + const StatusCallback &callback) { + auto table = GetOrCreateTable(table_name); + absl::MutexLock lock(&(table->mutex_)); + auto iter = table->index_keys_.find(index_key); + if (iter != table->index_keys_.end()) { + for (auto &key : iter->second) { + table->records_.erase(key); + } + table->index_keys_.erase(iter); + } + main_io_service_.post([callback]() { callback(Status::OK()); }); + return Status::OK(); +} + +std::shared_ptr InMemoryStoreClient::GetOrCreateTable( + const std::string &table_name) { + absl::MutexLock lock(&mutex_); + auto iter = tables_.find(table_name); + if (iter != tables_.end()) { + return iter->second; + } else { + auto table = std::make_shared(); + tables_[table_name] = table; + return table; + } +} + +} // namespace gcs + +} // namespace ray diff --git a/src/ray/gcs/store_client/in_memory_store_client.h b/src/ray/gcs/store_client/in_memory_store_client.h new file mode 100644 index 000000000..b07676af0 --- /dev/null +++ b/src/ray/gcs/store_client/in_memory_store_client.h @@ -0,0 +1,83 @@ +// 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_STORE_CLIENT_IN_MEMORY_STORE_CLIENT_H +#define RAY_GCS_STORE_CLIENT_IN_MEMORY_STORE_CLIENT_H + +#include "absl/container/flat_hash_map.h" +#include "absl/synchronization/mutex.h" +#include "ray/gcs/store_client/store_client.h" +#include "ray/protobuf/gcs.pb.h" + +namespace ray { + +namespace gcs { + +/// \class InMemoryStoreClient +/// +/// This class is thread safe. +class InMemoryStoreClient : public StoreClient { + public: + explicit InMemoryStoreClient(boost::asio::io_service &main_io_service) + : main_io_service_(main_io_service) {} + + Status AsyncPut(const std::string &table_name, const std::string &key, + const std::string &data, const StatusCallback &callback) override; + + 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 std::string &key, + const OptionalItemCallback &callback) override; + + Status AsyncGetAll( + const std::string &table_name, + const SegmentedCallback> &callback) override; + + Status AsyncDelete(const std::string &table_name, const std::string &key, + const StatusCallback &callback) override; + + Status AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key, + const StatusCallback &callback) override; + + private: + struct InMemoryTable { + /// Mutex to protect the records_ field and the index_keys_ field. + absl::Mutex mutex_; + // Mapping from key to data. + absl::flat_hash_map records_ GUARDED_BY(mutex_); + // Mapping from index key to keys. + absl::flat_hash_map> index_keys_ + GUARDED_BY(mutex_); + }; + + std::shared_ptr GetOrCreateTable( + const std::string &table_name); + + /// Mutex to protect the tables_ field. + absl::Mutex mutex_; + absl::flat_hash_map> tables_ + GUARDED_BY(mutex_); + + /// Async API Callback needs to post to main_io_service_ to ensure the orderly execution + /// of the callback. + boost::asio::io_service &main_io_service_; +}; + +} // namespace gcs + +} // namespace ray + +#endif // RAY_GCS_STORE_CLIENT_IN_MEMORY_STORE_CLIENT_H diff --git a/src/ray/gcs/store_client/test/in_memory_store_client_test.cc b/src/ray/gcs/store_client/test/in_memory_store_client_test.cc new file mode 100644 index 000000000..f302d3d80 --- /dev/null +++ b/src/ray/gcs/store_client/test/in_memory_store_client_test.cc @@ -0,0 +1,48 @@ +// 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 "ray/gcs/store_client/in_memory_store_client.h" +#include "ray/gcs/store_client/test/store_client_test_base.h" + +namespace ray { + +namespace gcs { + +class InMemoryStoreClientTest : public StoreClientTestBase { + public: + void InitStoreClient() override { + store_client_ = std::make_shared(*(io_service_pool_->Get())); + } + + void DisconnectStoreClient() override {} +}; + +TEST_F(InMemoryStoreClientTest, AsyncPutAndAsyncGetTest) { TestAsyncPutAndAsyncGet(); } + +TEST_F(InMemoryStoreClientTest, AsyncDeleteTest) { TestAsyncDelete(); } + +TEST_F(InMemoryStoreClientTest, AsyncGetAllTest) { TestAsyncGetAll(); } + +TEST_F(InMemoryStoreClientTest, AsyncPutAndDeleteWithIndexTest) { + TestAsyncPutAndDeleteWithIndex(); +} + +} // namespace gcs + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} 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 24f705d74..a4f7408de 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 @@ -26,6 +26,10 @@ class RedisStoreClientTest : public StoreClientTestBase { virtual ~RedisStoreClientTest() {} + static void SetUpTestCase() { RedisServiceManagerForTest::SetUpTestCase(); } + + static void TearDownTestCase() { RedisServiceManagerForTest::TearDownTestCase(); } + void InitStoreClient() override { RedisClientOptions options("127.0.0.1", REDIS_SERVER_PORT, "", true); redis_client_ = std::make_shared(options); @@ -40,113 +44,11 @@ class RedisStoreClientTest : public StoreClientTestBase { std::shared_ptr redis_client_; }; -TEST_F(RedisStoreClientTest, AsyncPutAndAsyncGetTest) { - // AsyncPut without index. - auto put_calllback = [this](const Status &status) { - RAY_CHECK_OK(status); - --pending_count_; - }; - for (const auto &elem : key_to_value_) { - ++pending_count_; - RAY_CHECK_OK(store_client_->AsyncPut(table_name_, elem.first.Binary(), - elem.second.SerializeAsString(), put_calllback)); - } - WaitPendingDone(); +TEST_F(RedisStoreClientTest, AsyncPutAndAsyncGetTest) { TestAsyncPutAndAsyncGet(); } - // AsyncGet - auto get_callback = [this](const Status &status, - const boost::optional &result) { - RAY_CHECK_OK(status); - RAY_CHECK(result); - 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_; - RAY_CHECK_OK(store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback)); - } - WaitPendingDone(); -} +TEST_F(RedisStoreClientTest, AsyncDeleteTest) { TestAsyncDelete(); } -TEST_F(RedisStoreClientTest, AsyncDeleteTest) { - // AsyncPut - auto put_calllback = [this](const Status &status) { --pending_count_; }; - for (const auto &elem : key_to_value_) { - ++pending_count_; - RAY_CHECK_OK(store_client_->AsyncPut(table_name_, elem.first.Binary(), - elem.second.SerializeAsString(), put_calllback)); - } - WaitPendingDone(); - - // AsyncDelete - auto delete_calllback = [this](const Status &status) { - RAY_CHECK_OK(status); - --pending_count_; - }; - for (const auto &elem : key_to_value_) { - ++pending_count_; - RAY_CHECK_OK( - store_client_->AsyncDelete(table_name_, elem.first.Binary(), delete_calllback)); - } - WaitPendingDone(); - - // AsyncGet - 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_; - RAY_CHECK_OK(store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback)); - } - WaitPendingDone(); -} - -TEST_F(RedisStoreClientTest, DISABLED_AsyncGetAllTest) { - // AsyncPut - 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; - 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](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 = ActorID::FromBinary(item.first); - auto it = received_keys.find(actor_id); - RAY_CHECK(it == received_keys.end()); - received_keys.emplace(actor_id); - - auto map_it = key_to_value_.find(actor_id); - RAY_CHECK(map_it != key_to_value_.end()); - } - if (!has_more) { - RAY_CHECK(received_keys.size() == key_to_value_.size()); - } - pending_count_ -= result.size(); - }; - - pending_count_ += key_to_value_.size(); - RAY_CHECK_OK(store_client_->AsyncGetAll(table_name_, get_all_callback)); - WaitPendingDone(); -} +TEST_F(RedisStoreClientTest, DISABLED_AsyncGetAllTest) { TestAsyncGetAll(); } } // namespace gcs 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 0d23734ff..0a6739ce7 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 @@ -29,7 +29,7 @@ namespace ray { namespace gcs { -class StoreClientTestBase : public RedisServiceManagerForTest { +class StoreClientTestBase : public ::testing::Test { public: StoreClientTestBase() {} @@ -57,6 +57,156 @@ class StoreClientTestBase : public RedisServiceManagerForTest { virtual void DisconnectStoreClient() = 0; protected: + void TestAsyncPutAndAsyncGet() { + // AsyncPut without index. + auto put_calllback = [this](const Status &status) { + RAY_CHECK_OK(status); + --pending_count_; + }; + for (const auto &elem : key_to_value_) { + ++pending_count_; + RAY_CHECK_OK(store_client_->AsyncPut(table_name_, elem.first.Binary(), + elem.second.SerializeAsString(), + put_calllback)); + } + WaitPendingDone(); + + // AsyncGet + auto get_callback = [this](const Status &status, + const boost::optional &result) { + RAY_CHECK_OK(status); + RAY_CHECK(result); + 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_; + RAY_CHECK_OK( + store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback)); + } + WaitPendingDone(); + } + + void TestAsyncDelete() { + // AsyncPut + auto put_calllback = [this](const Status &status) { --pending_count_; }; + for (const auto &elem : key_to_value_) { + ++pending_count_; + RAY_CHECK_OK(store_client_->AsyncPut(table_name_, elem.first.Binary(), + elem.second.SerializeAsString(), + put_calllback)); + } + WaitPendingDone(); + + // AsyncDelete + auto delete_calllback = [this](const Status &status) { + RAY_CHECK_OK(status); + --pending_count_; + }; + for (const auto &elem : key_to_value_) { + ++pending_count_; + RAY_CHECK_OK( + store_client_->AsyncDelete(table_name_, elem.first.Binary(), delete_calllback)); + } + WaitPendingDone(); + + // AsyncGet + 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_; + RAY_CHECK_OK( + store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback)); + } + WaitPendingDone(); + } + + void TestAsyncPutAndDeleteWithIndex() { + // AsyncPut with index + auto put_calllback = [this](const Status &status) { --pending_count_; }; + for (const auto &elem : key_to_value_) { + ++pending_count_; + RAY_CHECK_OK(store_client_->AsyncPutWithIndex( + table_name_, elem.first.Binary(), key_to_index_[elem.first].Binary(), + elem.second.SerializeAsString(), put_calllback)); + } + WaitPendingDone(); + + // AsyncDelete by index + auto delete_calllback = [this](const Status &status) { + RAY_CHECK_OK(status); + --pending_count_; + }; + for (const auto &elem : index_to_keys_) { + ++pending_count_; + RAY_CHECK_OK(store_client_->AsyncDeleteByIndex(table_name_, elem.first.Binary(), + delete_calllback)); + } + WaitPendingDone(); + + // AsyncGet + 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_; + RAY_CHECK_OK( + store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback)); + } + WaitPendingDone(); + } + + void TestAsyncGetAll() { + // AsyncPut + 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; + 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](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 = ActorID::FromBinary(item.first); + auto it = received_keys.find(actor_id); + RAY_CHECK(it == received_keys.end()); + received_keys.emplace(actor_id); + + auto map_it = key_to_value_.find(actor_id); + RAY_CHECK(map_it != key_to_value_.end()); + } + if (!has_more) { + RAY_CHECK(received_keys.size() == key_to_value_.size()); + } + pending_count_ -= result.size(); + }; + + pending_count_ += key_to_value_.size(); + RAY_CHECK_OK(store_client_->AsyncGetAll(table_name_, get_all_callback)); + WaitPendingDone(); + } + void GenTestData() { for (size_t i = 0; i < key_count_; i++) { rpc::ActorTableData actor;