diff --git a/BUILD.bazel b/BUILD.bazel index 638eaa85b..d2413b354 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -893,6 +893,73 @@ alias( actual = "@com_github_antirez_redis//:hiredis", ) +cc_library( + name = "redis_client", + srcs = [ + "src/ray/gcs/asio.cc", + "src/ray/gcs/redis_async_context.cc", + "src/ray/gcs/redis_client.cc", + "src/ray/gcs/redis_context.cc", + ], + hdrs = [ + "src/ray/gcs/asio.h", + "src/ray/gcs/redis_async_context.h", + "src/ray/gcs/redis_client.h", + "src/ray/gcs/redis_context.h", + ], + copts = COPTS, + deps = [ + ":hiredis", + ":ray_common", + ":ray_util", + ":stats_lib", + "@boost//:asio", + ], +) + +cc_library( + name = "redis_store_client", + srcs = [ + "src/ray/gcs/store_client/redis_store_client.cc", + ], + hdrs = [ + "src/ray/gcs/callback.h", + "src/ray/gcs/store_client/redis_store_client.h", + "src/ray/gcs/store_client/store_client.h", + ], + copts = COPTS, + deps = [ + "redis_client", + ], +) + +cc_library( + name = "store_client_test_lib", + hdrs = [ + "src/ray/gcs/store_client/test/store_client_test_base.h", + ], + copts = COPTS, + deps = [ + "redis_store_client", + ], +) + +cc_test( + name = "redis_store_client_test", + srcs = ["src/ray/gcs/store_client/test/redis_store_client_test.cc"], + args = ["$(location redis-server) $(location redis-cli)"], + copts = COPTS, + data = [ + "//:redis-cli", + "//:redis-server", + ], + deps = [ + ":redis_store_client", + ":store_client_test_lib", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "gcs", srcs = glob( diff --git a/src/ray/common/test_util.cc b/src/ray/common/test_util.cc index 4edb590f8..c4d309522 100644 --- a/src/ray/common/test_util.cc +++ b/src/ray/common/test_util.cc @@ -29,9 +29,15 @@ void RedisServiceManagerForTest::SetUpTestCase() { // Use random port to avoid port conflicts between UTs. REDIS_SERVER_PORT = random_gen(gen); - std::string start_redis_command = - REDIS_SERVER_EXEC_PATH + " --loglevel warning --loadmodule " + - REDIS_MODULE_LIBRARY_PATH + " --port " + std::to_string(REDIS_SERVER_PORT) + " &"; + std::string load_module_command; + if (!REDIS_MODULE_LIBRARY_PATH.empty()) { + // Fill load module command. + load_module_command = "--loadmodule " + REDIS_MODULE_LIBRARY_PATH; + } + + std::string start_redis_command = REDIS_SERVER_EXEC_PATH + " --loglevel warning " + + load_module_command + " --port " + + std::to_string(REDIS_SERVER_PORT) + " &"; RAY_LOG(INFO) << "Start redis command is: " << start_redis_command; RAY_CHECK(system(start_redis_command.c_str()) == 0); usleep(200 * 1000); diff --git a/src/ray/gcs/callback.h b/src/ray/gcs/callback.h index 7545058f6..af805f239 100644 --- a/src/ray/gcs/callback.h +++ b/src/ray/gcs/callback.h @@ -53,6 +53,16 @@ using SubscribeCallback = std::function; template using ItemCallback = std::function; +/// This callback is used to receive a large amount of results. +/// \param status Status indicates whether the scan was successful. +/// \param has_more Whether more data will be called back. +/// If `has_more == true`, there are more data to be received. This callback will +/// be called again. +/// \param result The items returned by storage. +template +using SegmentedCallback = + std::function &result)>; + } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/redis_client.cc b/src/ray/gcs/redis_client.cc new file mode 100644 index 000000000..c4ef49b1f --- /dev/null +++ b/src/ray/gcs/redis_client.cc @@ -0,0 +1,178 @@ +// 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/redis_client.h" + +#include +#include "ray/common/ray_config.h" +#include "ray/gcs/redis_context.h" + +namespace ray { + +namespace gcs { + +static void GetRedisShards(redisContext *context, std::vector *addresses, + std::vector *ports) { + // Get the total number of Redis shards in the system. + int num_attempts = 0; + redisReply *reply = nullptr; + while (num_attempts < RayConfig::instance().redis_db_connect_retries()) { + // Try to read the number of Redis shards from the primary shard. If the + // entry is present, exit. + reply = reinterpret_cast(redisCommand(context, "GET NumRedisShards")); + if (reply->type != REDIS_REPLY_NIL) { + break; + } + + // Sleep for a little, and try again if the entry isn't there yet. + freeReplyObject(reply); + usleep(RayConfig::instance().redis_db_connect_wait_milliseconds() * 1000); + num_attempts++; + } + RAY_CHECK(num_attempts < RayConfig::instance().redis_db_connect_retries()) + << "No entry found for NumRedisShards"; + RAY_CHECK(reply->type == REDIS_REPLY_STRING) + << "Expected string, found Redis type " << reply->type << " for NumRedisShards"; + int num_redis_shards = atoi(reply->str); + RAY_CHECK(num_redis_shards >= 1) << "Expected at least one Redis shard, " + << "found " << num_redis_shards; + freeReplyObject(reply); + + // Get the addresses of all of the Redis shards. + num_attempts = 0; + while (num_attempts < RayConfig::instance().redis_db_connect_retries()) { + // Try to read the Redis shard locations from the primary shard. If we find + // that all of them are present, exit. + reply = + reinterpret_cast(redisCommand(context, "LRANGE RedisShards 0 -1")); + if (static_cast(reply->elements) == num_redis_shards) { + break; + } + + // Sleep for a little, and try again if not all Redis shard addresses have + // been added yet. + freeReplyObject(reply); + usleep(RayConfig::instance().redis_db_connect_wait_milliseconds() * 1000); + num_attempts++; + } + RAY_CHECK(num_attempts < RayConfig::instance().redis_db_connect_retries()) + << "Expected " << num_redis_shards << " Redis shard addresses, found " + << reply->elements; + + // Parse the Redis shard addresses. + for (size_t i = 0; i < reply->elements; ++i) { + // Parse the shard addresses and ports. + RAY_CHECK(reply->element[i]->type == REDIS_REPLY_STRING); + std::string addr; + std::stringstream ss(reply->element[i]->str); + getline(ss, addr, ':'); + addresses->emplace_back(std::move(addr)); + int port; + ss >> port; + ports->emplace_back(port); + } + freeReplyObject(reply); +} + +RedisClient::RedisClient(const RedisClientOptions &options) : options_(options) {} + +Status RedisClient::Connect(boost::asio::io_service &io_service) { + std::vector io_services; + io_services.emplace_back(&io_service); + return Connect(io_services); +} + +Status RedisClient::Connect(std::vector io_services) { + RAY_CHECK(!is_connected_); + RAY_CHECK(!io_services.empty()); + + if (options_.server_ip_.empty()) { + RAY_LOG(ERROR) << "Failed to connect, redis server address is empty."; + return Status::Invalid("Redis server address is invalid!"); + } + + primary_context_ = std::make_shared(*io_services[0]); + + RAY_CHECK_OK(primary_context_->Connect(options_.server_ip_, options_.server_port_, + /*sharding=*/true, + /*password=*/options_.password_)); + + if (!options_.is_test_client_) { + // Moving sharding into constructor defaultly means that sharding = true. + // This design decision may worth a look. + std::vector addresses; + std::vector ports; + GetRedisShards(primary_context_->sync_context(), &addresses, &ports); + if (addresses.empty()) { + RAY_CHECK(ports.empty()); + addresses.push_back(options_.server_ip_); + ports.push_back(options_.server_port_); + } + + for (size_t i = 0; i < addresses.size(); ++i) { + size_t io_service_index = (i + 1) % io_services.size(); + boost::asio::io_service &io_service = *io_services[io_service_index]; + // Populate shard_contexts. + shard_contexts_.push_back(std::make_shared(io_service)); + RAY_CHECK_OK(shard_contexts_[i]->Connect(addresses[i], ports[i], /*sharding=*/true, + /*password=*/options_.password_)); + } + } else { + shard_contexts_.push_back(std::make_shared(*io_services[0])); + RAY_CHECK_OK(shard_contexts_[0]->Connect(options_.server_ip_, options_.server_port_, + /*sharding=*/true, + /*password=*/options_.password_)); + } + + Attach(); + + is_connected_ = true; + RAY_LOG(INFO) << "RedisClient connected."; + + return Status::OK(); +} + +void RedisClient::Attach() { + // Take care of sharding contexts. + RAY_CHECK(shard_asio_async_clients_.empty()) << "Attach shall be called only once"; + for (std::shared_ptr context : shard_contexts_) { + boost::asio::io_service &io_service = context->io_service(); + shard_asio_async_clients_.emplace_back( + new RedisAsioClient(io_service, context->async_context())); + shard_asio_subscribe_clients_.emplace_back( + new RedisAsioClient(io_service, context->subscribe_context())); + } + + boost::asio::io_service &io_service = primary_context_->io_service(); + asio_async_auxiliary_client_.reset( + new RedisAsioClient(io_service, primary_context_->async_context())); + asio_subscribe_auxiliary_client_.reset( + new RedisAsioClient(io_service, primary_context_->subscribe_context())); +} + +void RedisClient::Disconnect() { + RAY_CHECK(is_connected_); + is_connected_ = false; + RAY_LOG(INFO) << "RedisClient disconnected."; +} + +std::shared_ptr RedisClient::GetShardContext(const std::string &shard_key) { + static std::hash hash; + size_t index = hash(shard_key) % shard_contexts_.size(); + return shard_contexts_[index]; +} + +} // namespace gcs + +} // namespace ray diff --git a/src/ray/gcs/redis_client.h b/src/ray/gcs/redis_client.h new file mode 100644 index 000000000..cc062f8e0 --- /dev/null +++ b/src/ray/gcs/redis_client.h @@ -0,0 +1,111 @@ +// 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_REDIS_CLIENT_H +#define RAY_GCS_REDIS_CLIENT_H + +#include +#include + +#include "ray/common/status.h" +#include "ray/gcs/asio.h" +#include "ray/util/logging.h" + +namespace ray { + +namespace gcs { + +class RedisContext; + +class RedisClientOptions { + public: + RedisClientOptions(const std::string &ip, int port, const std::string &password, + bool is_test_client = false) + : server_ip_(ip), + server_port_(port), + password_(password), + is_test_client_(is_test_client) {} + + // Redis server address + std::string server_ip_; + int server_port_; + + // Password of Redis. + std::string password_; + + // Whether this client is used for tests. + bool is_test_client_{false}; +}; + +/// \class RedisClient +/// This class is used to send commands to Redis. +class RedisClient { + public: + RedisClient(const RedisClientOptions &options); + + /// Connect to Redis. Non-thread safe. + /// Call this function before calling other functions. + /// + /// \param io_service The event loop for this client. + /// This io_service must be single-threaded. Because `RedisAsioClient` is + /// non-thread safe. + /// \return Status + Status Connect(boost::asio::io_service &io_service); + + // TODO(micafan) Maybe it's not necessary to use multi threads. + /// Connect to Redis. Non-thread safe. + /// Call this function before calling other functions. + /// + /// \param io_services The event loops for this client. Each RedisContext bind to + /// an event loop. Each io_service must be single-threaded. Because `RedisAsioClient` + /// is non-thread safe. + /// \return Status + Status Connect(std::vector io_services); + + /// Disconnect with Redis. Non-thread safe. + void Disconnect(); + + std::vector> GetShardContexts() { + return shard_contexts_; + } + + std::shared_ptr GetShardContext(const std::string &shard_key); + + std::shared_ptr GetPrimaryContext() { return primary_context_; } + + protected: + /// Attach this client to an asio event loop. Note that only + /// one event loop should be attached at a time. + void Attach(); + + RedisClientOptions options_; + + /// Whether this client is connected to redis. + bool is_connected_{false}; + + // The following contexts write to the data shard + std::vector> shard_contexts_; + std::vector> shard_asio_async_clients_; + std::vector> shard_asio_subscribe_clients_; + // The following context writes everything to the primary shard + std::shared_ptr primary_context_; + std::unique_ptr asio_async_auxiliary_client_; + std::unique_ptr asio_subscribe_auxiliary_client_; +}; + +} // namespace gcs + +} // namespace ray + +#endif // RAY_GCS_REDIS_CLIENT_H \ No newline at end of file diff --git a/src/ray/gcs/redis_context.cc b/src/ray/gcs/redis_context.cc index ea34b20ef..1cddcd781 100644 --- a/src/ray/gcs/redis_context.cc +++ b/src/ray/gcs/redis_context.cc @@ -99,13 +99,8 @@ CallbackReply::CallbackReply(redisReply *redis_reply) : reply_type_(redis_reply- string_reply_ = std::string(message->str, message->len); RAY_CHECK(!string_reply_.empty()) << "Empty message received on subscribe channel."; } else { - string_array_reply_.reserve(redis_reply->elements); - for (size_t i = 0; i < redis_reply->elements; ++i) { - auto *entry = redis_reply->element[i]; - RAY_CHECK(REDIS_REPLY_STRING == entry->type) - << "Unexcepted type: " << entry->type; - string_array_reply_.emplace_back(std::string(entry->str, entry->len)); - } + // Array replies are used for scan or get. + ParseAsStringArrayOrScanArray(redis_reply); } break; } @@ -115,6 +110,44 @@ CallbackReply::CallbackReply(redisReply *redis_reply) : reply_type_(redis_reply- } } +void CallbackReply::ParseAsStringArrayOrScanArray(redisReply *redis_reply) { + RAY_CHECK(REDIS_REPLY_ARRAY == redis_reply->type); + const auto array_size = static_cast(redis_reply->elements); + if (array_size == 2) { + auto *cursor_entry = redis_reply->element[0]; + auto *array_entry = redis_reply->element[1]; + if (REDIS_REPLY_ARRAY == array_entry->type) { + // Parse as a scan array + RAY_CHECK(REDIS_REPLY_STRING == cursor_entry->type); + std::string cursor_str(cursor_entry->str, cursor_entry->len); + next_scan_cursor_reply_ = std::stoi(cursor_str); + const auto scan_array_size = array_entry->elements; + string_array_reply_.reserve(scan_array_size); + for (size_t i = 0; i < scan_array_size; ++i) { + auto *entry = array_entry->element[i]; + RAY_CHECK(REDIS_REPLY_STRING == entry->type) + << "Unexcepted type: " << entry->type; + string_array_reply_.push_back(std::string(entry->str, entry->len)); + RAY_LOG(DEBUG) << "Element index is " << i << ", element length is " + << entry->len; + } + return; + } + } + ParseAsStringArray(redis_reply); +} + +void CallbackReply::ParseAsStringArray(redisReply *redis_reply) { + RAY_CHECK(REDIS_REPLY_ARRAY == redis_reply->type); + const auto array_size = static_cast(redis_reply->elements); + string_array_reply_.reserve(array_size); + for (size_t i = 0; i < array_size; ++i) { + auto *entry = redis_reply->element[i]; + RAY_CHECK(REDIS_REPLY_STRING == entry->type) << "Unexcepted type: " << entry->type; + string_array_reply_.push_back(std::string(entry->str, entry->len)); + } +} + bool CallbackReply::IsNil() const { return REDIS_REPLY_NIL == reply_type_; } int64_t CallbackReply::ReadAsInteger() const { @@ -137,6 +170,12 @@ const std::string &CallbackReply::ReadAsPubsubData() const { return string_reply_; } +size_t CallbackReply::ReadAsScanArray(std::vector *array) const { + RAY_CHECK(reply_type_ == REDIS_REPLY_ARRAY) << "Unexpected type: " << reply_type_; + *array = string_array_reply_; + return next_scan_cursor_reply_; +} + const std::vector &CallbackReply::ReadAsStringArray() const { RAY_CHECK(reply_type_ == REDIS_REPLY_ARRAY) << "Unexpected type: " << reply_type_; return string_array_reply_; @@ -302,7 +341,8 @@ std::unique_ptr RedisContext::RunArgvSync( return callback_reply; } -Status RedisContext::RunArgvAsync(const std::vector &args) { +Status RedisContext::RunArgvAsync(const std::vector &args, + const RedisCallback &redis_callback) { RAY_CHECK(redis_async_context_); // Build the arguments. std::vector argv; @@ -311,9 +351,12 @@ Status RedisContext::RunArgvAsync(const std::vector &args) { argv.push_back(args[i].data()); argc.push_back(args[i].size()); } + int64_t callback_index = + RedisCallbackManager::instance().add(redis_callback, false, io_service_); // Run the Redis command. Status status = redis_async_context_->RedisAsyncCommandArgv( - nullptr, nullptr, args.size(), argv.data(), argc.data()); + reinterpret_cast(&GlobalRedisCallback), + reinterpret_cast(callback_index), args.size(), argv.data(), argc.data()); return status; } diff --git a/src/ray/gcs/redis_context.h b/src/ray/gcs/redis_context.h index 5573d92f2..0af743efe 100644 --- a/src/ray/gcs/redis_context.h +++ b/src/ray/gcs/redis_context.h @@ -70,7 +70,19 @@ class CallbackReply { /// Read this reply data as a string array. const std::vector &ReadAsStringArray() const; + /// Read this reply data as a scan array. + /// + /// \param array The result array of scan. + /// \return size_t The next cursor for scan. + size_t ReadAsScanArray(std::vector *array) const; + private: + /// Parse redis reply as string array or scan array. + void ParseAsStringArrayOrScanArray(redisReply *redis_reply); + + /// Parse redis reply as string array. + void ParseAsStringArray(redisReply *redis_reply); + /// Flag indicating the type of reply this represents. int reply_type_; @@ -84,8 +96,11 @@ class CallbackReply { std::string string_reply_; /// Reply data if reply_type_ is REDIS_REPLY_ARRAY. - /// Note this is used when get actor table data. + /// Represent the reply of StringArray or ScanArray. std::vector string_array_reply_; + + /// Represent the reply of SCanArray, means the next scan cursor for scan request. + size_t next_scan_cursor_reply_{0}; }; /// Every callback should take in a vector of the results from the Redis @@ -204,8 +219,10 @@ class RedisContext { /// Run an arbitrary Redis command without a callback. /// /// \param args The vector of command args to pass to Redis. + /// \param redis_callback The Redis callback function. /// \return Status. - Status RunArgvAsync(const std::vector &args); + Status RunArgvAsync(const std::vector &args, + const RedisCallback &redis_callback = nullptr); /// Subscribe to a specific Pub-Sub channel. /// @@ -232,6 +249,8 @@ class RedisContext { return *async_redis_subscribe_context_; } + boost::asio::io_service &io_service() { return io_service_; } + private: boost::asio::io_service &io_service_; redisContext *context_; diff --git a/src/ray/gcs/redis_gcs_client.cc b/src/ray/gcs/redis_gcs_client.cc index 6d527e1f9..1500b6d62 100644 --- a/src/ray/gcs/redis_gcs_client.cc +++ b/src/ray/gcs/redis_gcs_client.cc @@ -19,143 +19,56 @@ #include "ray/gcs/redis_accessor.h" #include "ray/gcs/redis_context.h" -static void GetRedisShards(redisContext *context, std::vector &addresses, - std::vector &ports) { - // Get the total number of Redis shards in the system. - int num_attempts = 0; - redisReply *reply = nullptr; - while (num_attempts < RayConfig::instance().redis_db_connect_retries()) { - // Try to read the number of Redis shards from the primary shard. If the - // entry is present, exit. - reply = reinterpret_cast(redisCommand(context, "GET NumRedisShards")); - if (reply->type != REDIS_REPLY_NIL) { - break; - } - - // Sleep for a little, and try again if the entry isn't there yet. */ - freeReplyObject(reply); - usleep(RayConfig::instance().redis_db_connect_wait_milliseconds() * 1000); - num_attempts++; - } - RAY_CHECK(num_attempts < RayConfig::instance().redis_db_connect_retries()) - << "No entry found for NumRedisShards"; - RAY_CHECK(reply->type == REDIS_REPLY_STRING) - << "Expected string, found Redis type " << reply->type << " for NumRedisShards"; - int num_redis_shards = atoi(reply->str); - RAY_CHECK(num_redis_shards >= 1) << "Expected at least one Redis shard, " - << "found " << num_redis_shards; - freeReplyObject(reply); - - // Get the addresses of all of the Redis shards. - num_attempts = 0; - while (num_attempts < RayConfig::instance().redis_db_connect_retries()) { - // Try to read the Redis shard locations from the primary shard. If we find - // that all of them are present, exit. - reply = - reinterpret_cast(redisCommand(context, "LRANGE RedisShards 0 -1")); - if (static_cast(reply->elements) == num_redis_shards) { - break; - } - - // Sleep for a little, and try again if not all Redis shard addresses have - // been added yet. - freeReplyObject(reply); - usleep(RayConfig::instance().redis_db_connect_wait_milliseconds() * 1000); - num_attempts++; - } - RAY_CHECK(num_attempts < RayConfig::instance().redis_db_connect_retries()) - << "Expected " << num_redis_shards << " Redis shard addresses, found " - << reply->elements; - - // Parse the Redis shard addresses. - for (size_t i = 0; i < reply->elements; ++i) { - // Parse the shard addresses and ports. - RAY_CHECK(reply->element[i]->type == REDIS_REPLY_STRING); - std::string addr; - std::stringstream ss(reply->element[i]->str); - getline(ss, addr, ':'); - addresses.push_back(addr); - int port; - ss >> port; - ports.push_back(port); - } - freeReplyObject(reply); -} - namespace ray { namespace gcs { RedisGcsClient::RedisGcsClient(const GcsClientOptions &options) - : GcsClient(options), command_type_(CommandType::kRegular) {} + : RedisGcsClient(options, CommandType::kRegular) {} RedisGcsClient::RedisGcsClient(const GcsClientOptions &options, CommandType command_type) - : GcsClient(options), command_type_(command_type) {} + : GcsClient(options), command_type_(command_type) { + RedisClientOptions redis_client_options(options.server_ip_, options.server_port_, + options.password_, options.is_test_client_); + redis_client_.reset(new RedisClient(redis_client_options)); +} Status RedisGcsClient::Connect(boost::asio::io_service &io_service) { RAY_CHECK(!is_connected_); - if (options_.server_ip_.empty()) { - RAY_LOG(ERROR) << "Failed to connect, gcs service address is empty."; - return Status::Invalid("gcs service address is invalid!"); + Status status = redis_client_->Connect(io_service); + if (!status.ok()) { + RAY_LOG(INFO) << "RedisGcsClient::Connect failed, status " << status.ToString(); + return status; } - primary_context_ = std::make_shared(io_service); + std::shared_ptr primary_context = redis_client_->GetPrimaryContext(); + std::vector> shard_contexts = + redis_client_->GetShardContexts(); - RAY_CHECK_OK(primary_context_->Connect(options_.server_ip_, options_.server_port_, - /*sharding=*/true, - /*password=*/options_.password_)); - - if (!options_.is_test_client_) { - // Moving sharding into constructor defaultly means that sharding = true. - // This design decision may worth a look. - std::vector addresses; - std::vector ports; - GetRedisShards(primary_context_->sync_context(), addresses, ports); - if (addresses.empty()) { - RAY_CHECK(ports.empty()); - addresses.push_back(options_.server_ip_); - ports.push_back(options_.server_port_); - } - - for (size_t i = 0; i < addresses.size(); ++i) { - // Populate shard_contexts. - shard_contexts_.push_back(std::make_shared(io_service)); - RAY_CHECK_OK(shard_contexts_[i]->Connect(addresses[i], ports[i], /*sharding=*/true, - /*password=*/options_.password_)); - } - } else { - shard_contexts_.push_back(std::make_shared(io_service)); - RAY_CHECK_OK(shard_contexts_[0]->Connect(options_.server_ip_, options_.server_port_, - /*sharding=*/true, - /*password=*/options_.password_)); - } - - Attach(io_service); - - log_based_actor_table_.reset(new LogBasedActorTable({primary_context_}, this)); - actor_table_.reset(new ActorTable({primary_context_}, this)); + log_based_actor_table_.reset(new LogBasedActorTable({primary_context}, this)); + actor_table_.reset(new ActorTable({primary_context}, this)); // TODO(micafan) Modify ClientTable' Constructor(remove ClientID) in future. // We will use NodeID instead of ClientID. // For worker/driver, it might not have this field(NodeID). // For raylet, NodeID should be initialized in raylet layer(not here). - client_table_.reset(new ClientTable({primary_context_}, this)); + client_table_.reset(new ClientTable({primary_context}, this)); - error_table_.reset(new ErrorTable({primary_context_}, this)); - job_table_.reset(new JobTable({primary_context_}, this)); - heartbeat_batch_table_.reset(new HeartbeatBatchTable({primary_context_}, this)); + error_table_.reset(new ErrorTable({primary_context}, this)); + job_table_.reset(new JobTable({primary_context}, this)); + heartbeat_batch_table_.reset(new HeartbeatBatchTable({primary_context}, this)); // Tables below would be sharded. - object_table_.reset(new ObjectTable(shard_contexts_, this)); - raylet_task_table_.reset(new raylet::TaskTable(shard_contexts_, this, command_type_)); - task_reconstruction_log_.reset(new TaskReconstructionLog(shard_contexts_, this)); - task_lease_table_.reset(new TaskLeaseTable(shard_contexts_, this)); - heartbeat_table_.reset(new HeartbeatTable(shard_contexts_, this)); - profile_table_.reset(new ProfileTable(shard_contexts_, this)); - actor_checkpoint_table_.reset(new ActorCheckpointTable(shard_contexts_, this)); - actor_checkpoint_id_table_.reset(new ActorCheckpointIdTable(shard_contexts_, this)); - resource_table_.reset(new DynamicResourceTable({primary_context_}, this)); - worker_failure_table_.reset(new WorkerFailureTable(shard_contexts_, this)); + object_table_.reset(new ObjectTable(shard_contexts, this)); + raylet_task_table_.reset(new raylet::TaskTable(shard_contexts, this, command_type_)); + task_reconstruction_log_.reset(new TaskReconstructionLog(shard_contexts, this)); + task_lease_table_.reset(new TaskLeaseTable(shard_contexts, this)); + heartbeat_table_.reset(new HeartbeatTable(shard_contexts, this)); + profile_table_.reset(new ProfileTable(shard_contexts, this)); + actor_checkpoint_table_.reset(new ActorCheckpointTable(shard_contexts, this)); + actor_checkpoint_id_table_.reset(new ActorCheckpointIdTable(shard_contexts, this)); + resource_table_.reset(new DynamicResourceTable({primary_context}, this)); + worker_failure_table_.reset(new WorkerFailureTable(shard_contexts, this)); actor_accessor_.reset(new RedisLogBasedActorInfoAccessor(this)); job_accessor_.reset(new RedisJobInfoAccessor(this)); @@ -176,23 +89,8 @@ Status RedisGcsClient::Connect(boost::asio::io_service &io_service) { void RedisGcsClient::Disconnect() { RAY_CHECK(is_connected_); is_connected_ = false; + redis_client_->Disconnect(); RAY_LOG(INFO) << "RedisGcsClient Disconnected."; - // TODO(micafan): Synchronously unregister node if this client is Raylet. -} - -void RedisGcsClient::Attach(boost::asio::io_service &io_service) { - // Take care of sharding contexts. - RAY_CHECK(shard_asio_async_clients_.empty()) << "Attach shall be called only once"; - for (std::shared_ptr context : shard_contexts_) { - shard_asio_async_clients_.emplace_back( - new RedisAsioClient(io_service, context->async_context())); - shard_asio_subscribe_clients_.emplace_back( - new RedisAsioClient(io_service, context->subscribe_context())); - } - asio_async_auxiliary_client_.reset( - new RedisAsioClient(io_service, primary_context_->async_context())); - asio_subscribe_auxiliary_client_.reset( - new RedisAsioClient(io_service, primary_context_->subscribe_context())); } std::string RedisGcsClient::DebugString() const { diff --git a/src/ray/gcs/redis_gcs_client.h b/src/ray/gcs/redis_gcs_client.h index 55c51a90a..4bf8d6739 100644 --- a/src/ray/gcs/redis_gcs_client.h +++ b/src/ray/gcs/redis_gcs_client.h @@ -22,6 +22,7 @@ #include "ray/common/status.h" #include "ray/gcs/asio.h" #include "ray/gcs/gcs_client.h" +#include "ray/gcs/redis_client.h" #include "ray/gcs/tables.h" #include "ray/util/logging.h" @@ -67,14 +68,18 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { // We also need something to export generic code to run on workers from the // driver (to set the PYTHONPATH) - using GetExportCallback = std::function; Status AddExport(const std::string &job_id, std::string &export_data); Status GetExport(const std::string &job_id, int64_t export_index, const GetExportCallback &done_callback); - std::vector> shard_contexts() { return shard_contexts_; } - std::shared_ptr primary_context() { return primary_context_; } + std::vector> shard_contexts() { + return redis_client_->GetShardContexts(); + } + + std::shared_ptr primary_context() { + return redis_client_->GetPrimaryContext(); + } /// The following xxx_table methods implement the Accessor interfaces. /// Implements the Actors() interface. @@ -104,14 +109,12 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { WorkerFailureTable &worker_failure_table(); private: - /// Attach this client to an asio event loop. Note that only - /// one event loop should be attached at a time. - void Attach(boost::asio::io_service &io_service); - // GCS command type. If CommandType::kChain, chain-replicated versions of the tables // might be used, if available. CommandType command_type_{CommandType::kUnknown}; + std::unique_ptr redis_client_; + std::unique_ptr object_table_; std::unique_ptr raylet_task_table_; std::unique_ptr log_based_actor_table_; @@ -127,15 +130,7 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { std::unique_ptr actor_checkpoint_id_table_; std::unique_ptr resource_table_; std::unique_ptr worker_failure_table_; - // The following contexts write to the data shard - std::vector> shard_contexts_; - std::vector> shard_asio_async_clients_; - std::vector> shard_asio_subscribe_clients_; - // The following context writes everything to the primary shard - std::shared_ptr primary_context_; std::unique_ptr job_table_; - std::unique_ptr asio_async_auxiliary_client_; - std::unique_ptr asio_subscribe_auxiliary_client_; }; } // namespace gcs diff --git a/src/ray/gcs/store_client/redis_store_client.cc b/src/ray/gcs/store_client/redis_store_client.cc new file mode 100644 index 000000000..ffea3ff2f --- /dev/null +++ b/src/ray/gcs/store_client/redis_store_client.cc @@ -0,0 +1,151 @@ +// 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/redis_store_client.h" + +#include +#include "ray/common/ray_config.h" +#include "ray/gcs/redis_context.h" +#include "ray/util/logging.h" + +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); +} + +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) { + if (!status.ok()) { + // Run callback if failed. + if (callback != nullptr) { + callback(status); + } + return; + } + + // Write data to Redis. + std::string full_key = table_name + key.Binary(); + status = DoPut(full_key, data_str, callback); + + if (!status.ok()) { + // Run callback if failed. + if (callback != nullptr) { + callback(status); + } + } + }; + + // 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); +} + +template +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) { + auto status = reply->ReadAsStatus(); + callback(status); + }; + } + + auto shard_context = redis_client_->GetShardContext(key); + return shard_context->RunArgvAsync(args, write_callback); +} + +template +Status RedisStoreClient::AsyncGet( + const std::string &table_name, const Key &key, + const OptionalItemCallback &callback) { + RAY_CHECK(callback != nullptr); + + auto redis_callback = [callback](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)); + result = std::move(data); + } + } + callback(Status::OK(), result); + }; + + std::string full_key = table_name + key.Binary(); + 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( + const std::string &table_name, + 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) { + RedisCallback delete_callback = nullptr; + if (callback) { + delete_callback = [callback](std::shared_ptr reply) { + int64_t deleted_count = reply->ReadAsInteger(); + RAY_LOG(DEBUG) << "Delete done, total delete count " << deleted_count; + callback(Status::OK()); + }; + } + + std::string full_key = table_name + key.Binary(); + 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) { + RAY_CHECK(0) << "Not implemented! Will implement this function in next PR."; + return Status::OK(); +} + +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 new file mode 100644 index 000000000..25a2a9677 --- /dev/null +++ b/src/ray/gcs/store_client/redis_store_client.h @@ -0,0 +1,69 @@ +// 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_REDIS_STORE_CLIENT_H +#define RAY_GCS_STORE_CLIENT_REDIS_STORE_CLIENT_H + +#include +#include +#include "ray/gcs/redis_client.h" +#include "ray/gcs/redis_context.h" +#include "ray/gcs/store_client/store_client.h" +#include "ray/protobuf/gcs.pb.h" + +namespace ray { + +namespace gcs { + +template +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 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, + const StatusCallback &callback) override; + + Status AsyncGet(const std::string &table_name, const Key &key, + const OptionalItemCallback &callback) override; + + Status AsyncGetAll(const std::string &table_name, + const SegmentedCallback> &callback) override; + + Status AsyncDelete(const std::string &table_name, const Key &key, + const StatusCallback &callback) override; + + Status AsyncDeleteByIndex(const std::string &table_name, const IndexKey &index_key, + const StatusCallback &callback) override; + + private: + Status DoPut(const std::string &key, const std::string &data, + const StatusCallback &callback); + + std::shared_ptr redis_client_; +}; + +typedef RedisStoreClient RedisActorStoreTable; + +} // namespace gcs + +} // namespace ray + +#endif // RAY_GCS_STORE_CLIENT_REDIS_STORE_CLIENT_H diff --git a/src/ray/gcs/store_client/store_client.h b/src/ray/gcs/store_client/store_client.h new file mode 100644 index 000000000..941ac16b2 --- /dev/null +++ b/src/ray/gcs/store_client/store_client.h @@ -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. + +#ifndef RAY_GCS_STORE_CLIENT_STORE_CLIENT_H +#define RAY_GCS_STORE_CLIENT_STORE_CLIENT_H + +#include +#include +#include "ray/common/id.h" +#include "ray/common/status.h" +#include "ray/gcs/callback.h" +#include "ray/protobuf/gcs.pb.h" +#include "ray/util/io_service_pool.h" +#include "ray/util/logging.h" + +namespace ray { + +namespace gcs { + +/// \class StoreClient +/// Abstract interface of the storage client. +template +class StoreClient { + public: + virtual ~StoreClient() {} + + /// Write data to the given table asynchronously. + /// + /// \param table_name The name of the table to be written. + /// \param key The key that will be written to the table. + /// \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; + + /// Write data to the given table asynchronously. + /// + /// \param table_name The name of the table to be written. + /// \param key The key that will be written to the table. + /// \param index_key A secondary key that will be used for indexing the data. + /// \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, + const StatusCallback &callback) = 0; + + /// Get data from the given table asynchronously. + /// + /// \param table_name The name of the table to be read. + /// \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; + + /// Get all data from the given table asynchronously. + /// + /// \param table_name The name of the table to be read. + /// \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; + + /// Delete data from the given table asynchronously. + /// + /// \param table_name The name of the table from which data is to be deleted. + /// \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, + const StatusCallback &callback) = 0; + + /// Delete by index from the given table asynchronously. + /// + /// \param table_name The name of the table from which data is to be deleted. + /// \param index_key The secondary key that will be used to delete the indexed data. + /// from the table. + /// \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 StatusCallback &callback) = 0; + + protected: + StoreClient() = default; +}; + +typedef StoreClient ActorStoreTable; + +} // namespace gcs + +} // namespace ray + +#endif // RAY_GCS_STORE_CLIENT_STORE_CLIENT_H 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 new file mode 100644 index 000000000..9ccc641af --- /dev/null +++ b/src/ray/gcs/store_client/test/redis_store_client_test.cc @@ -0,0 +1,166 @@ +// 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/redis_store_client.h" +#include "ray/gcs/redis_client.h" +#include "ray/gcs/store_client/test/store_client_test_base.h" + +namespace ray { + +namespace gcs { + +class RedisStoreClientTest : public StoreClientTestBase { + public: + RedisStoreClientTest() {} + + virtual ~RedisStoreClientTest() {} + + void InitStoreClient() override { + 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())); + + store_client_ = + std::make_shared>( + redis_client_); + } + + void DisconnectStoreClient() override { redis_client_->Disconnect(); } + + protected: + std::shared_ptr redis_client_; +}; + +TEST_F(RedisStoreClientTest, AsyncPutAndAsyncGetTest) { + // AsyncPut without index. + auto put_calllback = [this](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); + } + WaitPendingDone(); + + // AsyncGet + auto get_callback = [this](Status status, + const boost::optional &result) { + RAY_CHECK_OK(status); + RAY_CHECK(result); + ActorID actor_id = ActorID::FromBinary(result->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); + } + WaitPendingDone(); +} + +TEST_F(RedisStoreClientTest, AsyncDeleteTest) { + // AsyncPut + auto put_calllback = [this](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); + } + WaitPendingDone(); + + // AsyncDelete + auto delete_calllback = [this](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); + } + WaitPendingDone(); + + // AsyncGet + auto get_callback = [this](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); + } + WaitPendingDone(); +} + +TEST_F(RedisStoreClientTest, DISABLED_AsyncGetAllTest) { + // AsyncPut + auto put_calllback = [this](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); + } + WaitPendingDone(); + + // AsyncGetAll + auto get_all_callback = + [this](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; + 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(); + Status status = store_client_->AsyncGetAll(table_name_, get_all_callback); + RAY_CHECK_OK(status); + WaitPendingDone(); +} + +} // namespace gcs + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + RAY_CHECK(argc == 3); + ray::REDIS_SERVER_EXEC_PATH = argv[1]; + ray::REDIS_CLIENT_EXEC_PATH = argv[2]; + return RUN_ALL_TESTS(); +} 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 new file mode 100644 index 000000000..5495b79d8 --- /dev/null +++ b/src/ray/gcs/store_client/test/store_client_test_base.h @@ -0,0 +1,112 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include "ray/common/id.h" +#include "ray/common/test_util.h" +#include "ray/gcs/store_client/store_client.h" +#include "ray/util/io_service_pool.h" +#include "ray/util/logging.h" + +namespace ray { + +namespace gcs { + +class StoreClientTestBase : public RedisServiceManagerForTest { + public: + StoreClientTestBase() {} + + virtual ~StoreClientTestBase() {} + + void SetUp() override { + io_service_pool_ = std::make_shared(io_service_num_); + io_service_pool_->Run(); + InitStoreClient(); + GenTestData(); + } + + void TearDown() override { + DisconnectStoreClient(); + + io_service_pool_->Stop(); + + key_to_value_.clear(); + key_to_index_.clear(); + index_to_keys_.clear(); + } + + virtual void InitStoreClient() = 0; + + virtual void DisconnectStoreClient() = 0; + + protected: + void GenTestData() { + for (size_t i = 0; i < key_count_; i++) { + rpc::ActorTableData actor; + actor.set_max_reconstructions(1); + actor.set_remaining_reconstructions(1); + JobID job_id = JobID::FromInt(i % index_count_); + actor.set_job_id(job_id.Binary()); + actor.set_state(rpc::ActorTableData::ALIVE); + ActorID actor_id = ActorID::Of(job_id, RandomTaskId(), /*parent_task_counter=*/i); + actor.set_actor_id(actor_id.Binary()); + + key_to_value_[actor_id] = actor; + + key_to_index_[actor_id] = job_id; + + auto it = index_to_keys_.find(job_id); + if (it != index_to_keys_.end()) { + it->second.emplace(actor_id); + } else { + std::unordered_set key_set; + key_set.emplace(actor_id); + index_to_keys_.emplace(job_id, std::move(key_set)); + } + } + } + + void WaitPendingDone() { WaitPendingDone(pending_count_); } + + void WaitPendingDone(std::atomic &pending_count) { + auto condition = [&pending_count]() { return pending_count == 0; }; + EXPECT_TRUE(WaitForCondition(condition, wait_pending_timeout_.count())); + } + + protected: + size_t io_service_num_{2}; + std::shared_ptr io_service_pool_; + + std::shared_ptr> store_client_; + + std::string table_name_{"test_table"}; + size_t key_count_{5000}; + size_t index_count_{100}; + std::unordered_map key_to_value_; + std::unordered_map key_to_index_; + std::unordered_map> index_to_keys_; + + std::atomic pending_count_{0}; + std::chrono::milliseconds wait_pending_timeout_{5000}; +}; + +} // namespace gcs + +} // namespace ray diff --git a/src/ray/util/io_service_pool.cc b/src/ray/util/io_service_pool.cc new file mode 100644 index 000000000..6b56062b1 --- /dev/null +++ b/src/ray/util/io_service_pool.cc @@ -0,0 +1,50 @@ +// 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/util/io_service_pool.h" +#include "ray/util/logging.h" + +namespace ray { + +IOServicePool::IOServicePool(size_t io_service_num) : io_service_num_(io_service_num) {} + +IOServicePool::~IOServicePool() {} + +void IOServicePool::Run() { + for (size_t i = 0; i < io_service_num_; ++i) { + io_services_.emplace_back( + std::unique_ptr(new boost::asio::io_service)); + boost::asio::io_service &io_service = *io_services_[i]; + threads_.emplace_back([&io_service] { + boost::asio::io_service::work work(io_service); + io_service.run(); + }); + } + + RAY_LOG(INFO) << "IOServicePool is running with " << io_service_num_ << " io_service."; +} + +void IOServicePool::Stop() { + for (auto &io_service : io_services_) { + io_service->stop(); + } + + for (auto &thread : threads_) { + thread.join(); + } + + RAY_LOG(INFO) << "IOServicePool is stopped."; +} + +} // namespace ray diff --git a/src/ray/util/io_service_pool.h b/src/ray/util/io_service_pool.h new file mode 100644 index 000000000..a923e7859 --- /dev/null +++ b/src/ray/util/io_service_pool.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_UTIL_IO_SERVICE_POOL_H +#define RAY_UTIL_IO_SERVICE_POOL_H + +#include +#include +#include + +namespace ray { + +/// \class IOServicePool +/// The io_service pool. Each io_service owns a thread. +/// To get io_service from this pool should call `Run()` first. +/// Before exit, `Stop()` must be called. +class IOServicePool { + public: + IOServicePool(size_t io_service_num); + + ~IOServicePool(); + + void Run(); + + void Stop(); + + /// Select io_service by round robin. + /// + /// \return io_service + boost::asio::io_service *Get(); + + /// Select io_service by hash. + /// + /// \param hash Use this hash to pick a io_service. + /// The same hash will alway get the same io_service. + /// \return io_service + boost::asio::io_service *Get(size_t hash); + + /// Get all io_service. + /// This is only use for RedisClient::Connect(). + std::vector GetAll(); + + private: + size_t io_service_num_{0}; + + std::vector threads_; + std::vector> io_services_; + + std::atomic current_index_; +}; + +inline boost::asio::io_service *IOServicePool::Get() { + size_t index = ++current_index_ % io_service_num_; + return io_services_[index].get(); +} + +inline boost::asio::io_service *IOServicePool::Get(size_t hash) { + size_t index = hash % io_service_num_; + return io_services_[index].get(); +} + +inline std::vector IOServicePool::GetAll() { + std::vector io_services; + for (auto &io_service : io_services_) { + io_services.emplace_back(io_service.get()); + } + return io_services; +} + +} // namespace ray + +#endif // RAY_UTIL_IO_SERVICE_POOL_H