wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-17 11:32:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

[GCS] impl RedisStoreClient for GCS Service (#7675)

Browse Source
This commit is contained in:
micafan
2020-04-01 21:18:19 +08:00
committed by GitHub
parent e153e3179f
commit 780c1c3b08
16 changed files with 1227 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
BUILD.bazel
+67
View File
@@ -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(
src/ray/common/test_util.cc
+9 -3
View File
@@ -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);
src/ray/gcs/callback.h
+10
View File
@@ -53,6 +53,16 @@ using SubscribeCallback = std::function<void(const ID &id, const Data &result)>;
template <typename Data>
using ItemCallback = std::function<void(const Data &result)>;
/// 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 <typename Data>
using SegmentedCallback =
std::function<void(Status status, bool has_more, const std::vector<Data> &result)>;
} // namespace gcs
} // namespace ray
src/ray/gcs/redis_client.cc
+178
View File
@@ -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 <unistd.h>
#include "ray/common/ray_config.h"
#include "ray/gcs/redis_context.h"
namespace ray {
namespace gcs {
static void GetRedisShards(redisContext *context, std::vector<std::string> *addresses,
std::vector<int> *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<redisReply *>(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<redisReply *>(redisCommand(context, "LRANGE RedisShards 0 -1"));
if (static_cast<int>(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<boost::asio::io_service *> io_services;
io_services.emplace_back(&io_service);
return Connect(io_services);
}
Status RedisClient::Connect(std::vector<boost::asio::io_service *> 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<RedisContext>(*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<std::string> addresses;
std::vector<int> 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<RedisContext>(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<RedisContext>(*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<RedisContext> 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<RedisContext> RedisClient::GetShardContext(const std::string &shard_key) {
static std::hash<std::string> hash;
size_t index = hash(shard_key) % shard_contexts_.size();
return shard_contexts_[index];
}
} // namespace gcs
} // namespace ray
src/ray/gcs/redis_client.h
+111
View File
@@ -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 <map>
#include <string>
#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<boost::asio::io_service *> io_services);
/// Disconnect with Redis. Non-thread safe.
void Disconnect();
std::vector<std::shared_ptr<RedisContext>> GetShardContexts() {
return shard_contexts_;
}
std::shared_ptr<RedisContext> GetShardContext(const std::string &shard_key);
std::shared_ptr<RedisContext> 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<std::shared_ptr<RedisContext>> shard_contexts_;
std::vector<std::unique_ptr<RedisAsioClient>> shard_asio_async_clients_;
std::vector<std::unique_ptr<RedisAsioClient>> shard_asio_subscribe_clients_;
// The following context writes everything to the primary shard
std::shared_ptr<RedisContext> primary_context_;
std::unique_ptr<RedisAsioClient> asio_async_auxiliary_client_;
std::unique_ptr<RedisAsioClient> asio_subscribe_auxiliary_client_;
};
} // namespace gcs
} // namespace ray
#endif // RAY_GCS_REDIS_CLIENT_H
src/ray/gcs/redis_context.cc
+52 -9
View File
@@ -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<size_t>(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<size_t>(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<std::string> *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<std::string> &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<CallbackReply> RedisContext::RunArgvSync(
return callback_reply;
}
Status RedisContext::RunArgvAsync(const std::vector<std::string> &args) {
Status RedisContext::RunArgvAsync(const std::vector<std::string> &args,
const RedisCallback &redis_callback) {
RAY_CHECK(redis_async_context_);
// Build the arguments.
std::vector<const char *> argv;
@@ -311,9 +351,12 @@ Status RedisContext::RunArgvAsync(const std::vector<std::string> &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<redisCallbackFn *>(&GlobalRedisCallback),
reinterpret_cast<void *>(callback_index), args.size(), argv.data(), argc.data());
return status;
}
src/ray/gcs/redis_context.h
+21 -2
View File
@@ -70,7 +70,19 @@ class CallbackReply {
/// Read this reply data as a string array.
const std::vector<std::string> &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<std::string> *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<std::string> 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<std::string> &args);
Status RunArgvAsync(const std::vector<std::string> &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_;
src/ray/gcs/redis_gcs_client.cc
+30 -132
View File
@@ -19,143 +19,56 @@
#include "ray/gcs/redis_accessor.h"
#include "ray/gcs/redis_context.h"
static void GetRedisShards(redisContext *context, std::vector<std::string> &addresses,
std::vector<int> &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<redisReply *>(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<redisReply *>(redisCommand(context, "LRANGE RedisShards 0 -1"));
if (static_cast<int>(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<RedisContext>(io_service);
std::shared_ptr<RedisContext> primary_context = redis_client_->GetPrimaryContext();
std::vector<std::shared_ptr<RedisContext>> 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<std::string> addresses;
std::vector<int> 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<RedisContext>(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<RedisContext>(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<RedisContext> 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 {
src/ray/gcs/redis_gcs_client.h
+10 -15
View File
@@ -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<void(const std::string &data)>;
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<std::shared_ptr<RedisContext>> shard_contexts() { return shard_contexts_; }
std::shared_ptr<RedisContext> primary_context() { return primary_context_; }
std::vector<std::shared_ptr<RedisContext>> shard_contexts() {
return redis_client_->GetShardContexts();
}
std::shared_ptr<RedisContext> 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<RedisClient> redis_client_;
std::unique_ptr<ObjectTable> object_table_;
std::unique_ptr<raylet::TaskTable> raylet_task_table_;
std::unique_ptr<LogBasedActorTable> log_based_actor_table_;
@@ -127,15 +130,7 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
std::unique_ptr<ActorCheckpointIdTable> actor_checkpoint_id_table_;
std::unique_ptr<DynamicResourceTable> resource_table_;
std::unique_ptr<WorkerFailureTable> worker_failure_table_;
// The following contexts write to the data shard
std::vector<std::shared_ptr<RedisContext>> shard_contexts_;
std::vector<std::unique_ptr<RedisAsioClient>> shard_asio_async_clients_;
std::vector<std::unique_ptr<RedisAsioClient>> shard_asio_subscribe_clients_;
// The following context writes everything to the primary shard
std::shared_ptr<RedisContext> primary_context_;
std::unique_ptr<JobTable> job_table_;
std::unique_ptr<RedisAsioClient> asio_async_auxiliary_client_;
std::unique_ptr<RedisAsioClient> asio_subscribe_auxiliary_client_;
};
} // namespace gcs
src/ray/gcs/store_client/redis_store_client.cc
+151
View File
@@ -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 <functional>
#include "ray/common/ray_config.h"
#include "ray/gcs/redis_context.h"
#include "ray/util/logging.h"
namespace ray {
namespace gcs {
template <typename Key, typename Data, typename IndexKey>
Status RedisStoreClient<Key, Data, IndexKey>::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 <typename Key, typename Data, typename IndexKey>
Status RedisStoreClient<Key, Data, IndexKey>::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 <typename Key, typename Data, typename IndexKey>
Status RedisStoreClient<Key, Data, IndexKey>::DoPut(const std::string &key,
const std::string &data,
const StatusCallback &callback) {
std::vector<std::string> args = {"SET", key, data};
RedisCallback write_callback = nullptr;
if (callback) {
write_callback = [callback](std::shared_ptr<CallbackReply> reply) {
auto status = reply->ReadAsStatus();
callback(status);
};
}
auto shard_context = redis_client_->GetShardContext(key);
return shard_context->RunArgvAsync(args, write_callback);
}
template <typename Key, typename Data, typename IndexKey>
Status RedisStoreClient<Key, Data, IndexKey>::AsyncGet(
const std::string &table_name, const Key &key,
const OptionalItemCallback<Data> &callback) {
RAY_CHECK(callback != nullptr);
auto redis_callback = [callback](std::shared_ptr<CallbackReply> reply) {
boost::optional<Data> 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<std::string> args = {"GET", full_key};
auto shard_context = redis_client_->GetShardContext(full_key);
return shard_context->RunArgvAsync(args, redis_callback);
}
template <typename Key, typename Data, typename IndexKey>
Status RedisStoreClient<Key, Data, IndexKey>::AsyncGetAll(
const std::string &table_name,
const SegmentedCallback<std::pair<Key, Data>> &callback) {
RAY_CHECK(0) << "Not implemented! Will implement this function in next PR.";
return Status::OK();
}
template <typename Key, typename Data, typename IndexKey>
Status RedisStoreClient<Key, Data, IndexKey>::AsyncDelete(
const std::string &table_name, const Key &key, const StatusCallback &callback) {
RedisCallback delete_callback = nullptr;
if (callback) {
delete_callback = [callback](std::shared_ptr<CallbackReply> 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<std::string> args = {"DEL", full_key};
auto shard_context = redis_client_->GetShardContext(full_key);
return shard_context->RunArgvAsync(args, delete_callback);
}
template <typename Key, typename Data, typename IndexKey>
Status RedisStoreClient<Key, Data, IndexKey>::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<ActorID, rpc::ActorTableData, JobID>;
} // namespace gcs
} // namespace ray
src/ray/gcs/store_client/redis_store_client.h
+69
View File
@@ -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 <memory>
#include <unordered_set>
#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 <typename Key, typename Data, typename IndexKey>
class RedisStoreClient : public StoreClient<Key, Data, IndexKey> {
public:
RedisStoreClient(std::shared_ptr<RedisClient> 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<Data> &callback) override;
Status AsyncGetAll(const std::string &table_name,
const SegmentedCallback<std::pair<Key, Data>> &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<RedisClient> redis_client_;
};
typedef RedisStoreClient<ActorID, rpc::ActorTableData, JobID> RedisActorStoreTable;
} // namespace gcs
} // namespace ray
#endif // RAY_GCS_STORE_CLIENT_REDIS_STORE_CLIENT_H
src/ray/gcs/store_client/store_client.h
+108
View File
@@ -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 <memory>
#include <string>
#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 <typename Key, typename Data, typename IndexKey>
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<Data> &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<std::pair<Key, Data>> &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<ActorID, rpc::ActorTableData, JobID> ActorStoreTable;
} // namespace gcs
} // namespace ray
#endif // RAY_GCS_STORE_CLIENT_STORE_CLIENT_H
src/ray/gcs/store_client/test/redis_store_client_test.cc
+166
View File
@@ -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<RedisClient>(options);
RAY_CHECK_OK(redis_client_->Connect(io_service_pool_->GetAll()));
store_client_ =
std::make_shared<RedisStoreClient<ActorID, rpc::ActorTableData, JobID>>(
redis_client_);
}
void DisconnectStoreClient() override { redis_client_->Disconnect(); }
protected:
std::shared_ptr<RedisClient> 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<rpc::ActorTableData> &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<rpc::ActorTableData> &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<std::pair<ActorID, rpc::ActorTableData>> &result) {
RAY_CHECK_OK(status);
static std::unordered_set<ActorID> 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();
}
src/ray/gcs/store_client/test/store_client_test_base.h
+112
View File
@@ -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 <atomic>
#include <chrono>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#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<IOServicePool>(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<ActorID> 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<int> &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<IOServicePool> io_service_pool_;
std::shared_ptr<StoreClient<ActorID, rpc::ActorTableData, JobID>> store_client_;
std::string table_name_{"test_table"};
size_t key_count_{5000};
size_t index_count_{100};
std::unordered_map<ActorID, rpc::ActorTableData> key_to_value_;
std::unordered_map<ActorID, JobID> key_to_index_;
std::unordered_map<JobID, std::unordered_set<ActorID>> index_to_keys_;
std::atomic<int> pending_count_{0};
std::chrono::milliseconds wait_pending_timeout_{5000};
};
} // namespace gcs
} // namespace ray
src/ray/util/io_service_pool.cc
+50
View File
@@ -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<boost::asio::io_service>(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
src/ray/util/io_service_pool.h
+83
View File
@@ -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 <atomic>
#include <boost/asio.hpp>
#include <thread>
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<boost::asio::io_service *> GetAll();
private:
size_t io_service_num_{0};
std::vector<std::thread> threads_;
std::vector<std::unique_ptr<boost::asio::io_service>> io_services_;
std::atomic<size_t> 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<boost::asio::io_service *> IOServicePool::GetAll() {
std::vector<boost::asio::io_service *> 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