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

Introduce a log interface for the new GCS (#1771)

Browse Source 
* TABLE_APPEND call

* Convert callbacks back to taking in a string...

* GCS returns flatbuffers, define Log class

* Cleanups

* Modify client table to use the Log interface

* Fix bug where we replied twice from redis

* Fixes

* lint
This commit is contained in:
Stephanie Wang
2018-03-26 16:00:43 -07:00
committed by Philipp Moritz
parent 7c4afa4b04
commit 0fd4112354
7 changed files with 407 additions and 237 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/redis_module/ray_redis_module.cc
+121 -102
View File
@@ -502,63 +502,6 @@ int TaskTableAdd(RedisModuleCtx *ctx,
return RedisModule_ReplyWithSimpleString(ctx, "OK");
}
// TODO(swang): Implement the client table as an append-only log so that we
// don't need this special case for client table publication.
int ClientTableAdd(RedisModuleCtx *ctx,
RedisModuleString *pubsub_channel_str,
RedisModuleString *data) {
const char *buf = RedisModule_StringPtrLen(data, NULL);
auto client_data = flatbuffers::GetRoot<ClientTableData>(buf);
RedisModuleKey *clients_key = (RedisModuleKey *) RedisModule_OpenKey(
ctx, pubsub_channel_str, REDISMODULE_READ | REDISMODULE_WRITE);
// If this is a client addition, send all previous notifications, in order.
// NOTE(swang): This will go to all clients, so some clients will get
// duplicate notifications.
if (client_data->is_insertion() &&
RedisModule_KeyType(clients_key) != REDISMODULE_KEYTYPE_EMPTY) {
// NOTE(swang): Sets are not implemented yet, so we use ZSETs instead.
CHECK_ERROR(RedisModule_ZsetFirstInScoreRange(
clients_key, REDISMODULE_NEGATIVE_INFINITE,
REDISMODULE_POSITIVE_INFINITE, 1, 1),
"Unable to initialize zset iterator");
do {
RedisModuleString *message =
RedisModule_ZsetRangeCurrentElement(clients_key, NULL);
RedisModuleCallReply *reply =
RedisModule_Call(ctx, "PUBLISH", "ss", pubsub_channel_str, message);
if (reply == NULL) {
RedisModule_CloseKey(clients_key);
return RedisModule_ReplyWithError(ctx, "error during PUBLISH");
}
} while (RedisModule_ZsetRangeNext(clients_key));
}
// Append this notification to the past notifications so that it will get
// sent to new clients in the future.
size_t index = RedisModule_ValueLength(clients_key);
// Serialize the notification to send.
flatbuffers::FlatBufferBuilder fbb;
auto message = CreateGcsNotification(fbb, fbb.CreateString(""),
RedisStringToFlatbuf(fbb, data));
fbb.Finish(message);
auto notification = RedisModule_CreateString(
ctx, reinterpret_cast<const char *>(fbb.GetBufferPointer()),
fbb.GetSize());
RedisModule_ZsetAdd(clients_key, index, notification, NULL);
// Publish the notification about this client.
RedisModuleCallReply *reply =
RedisModule_Call(ctx, "PUBLISH", "ss", pubsub_channel_str, notification);
RedisModule_FreeString(ctx, notification);
if (reply == NULL) {
RedisModule_CloseKey(clients_key);
return RedisModule_ReplyWithError(ctx, "error during PUBLISH");
}
RedisModule_CloseKey(clients_key);
return RedisModule_ReplyWithSimpleString(ctx, "OK");
}
/// Publish a notification for a new entry at a key. This publishes a
/// notification to all subscribers of the table, as well as every client that
/// has requested notifications for this key.
@@ -575,8 +518,9 @@ int PublishTableAdd(RedisModuleCtx *ctx,
RedisModuleString *data) {
// Serialize the notification to send.
flatbuffers::FlatBufferBuilder fbb;
auto message = CreateGcsNotification(fbb, RedisStringToFlatbuf(fbb, id),
RedisStringToFlatbuf(fbb, data));
auto data_flatbuf = RedisStringToFlatbuf(fbb, data);
auto message = CreateGcsTableEntry(fbb, RedisStringToFlatbuf(fbb, id),
fbb.CreateVector(&data_flatbuf, 1));
fbb.Finish(message);
// Write the data back to any subscribers that are listening to all table
@@ -654,9 +598,6 @@ int TableAdd_RedisCommand(RedisModuleCtx *ctx,
// TODO(swang): This is only necessary for legacy Ray and should be removed
// once we switch to using the new GCS API for the task table.
return TaskTableAdd(ctx, id, data);
} else if (pubsub_channel == TablePubsub_CLIENT) {
// Publish all previous client table additions to the new client.
return ClientTableAdd(ctx, pubsub_channel_str, data);
} else if (pubsub_channel != TablePubsub_NO_PUBLISH) {
// All other pubsub channels write the data back directly onto the channel.
return PublishTableAdd(ctx, pubsub_channel_str, id, data);
@@ -665,28 +606,109 @@ int TableAdd_RedisCommand(RedisModuleCtx *ctx,
}
}
// This is a temporary redis command that will be removed once
// the GCS uses https://github.com/pcmoritz/credis.
int TableAppend_RedisCommand(RedisModuleCtx *ctx,
RedisModuleString **argv,
int argc) {
if (argc != 5) {
return RedisModule_WrongArity(ctx);
}
RedisModuleString *prefix_str = argv[1];
RedisModuleString *pubsub_channel_str = argv[2];
RedisModuleString *id = argv[3];
RedisModuleString *data = argv[4];
// Set the keys in the table.
RedisModuleKey *key = OpenPrefixedKey(ctx, prefix_str, id,
REDISMODULE_READ | REDISMODULE_WRITE);
size_t index = RedisModule_ValueLength(key);
RedisModule_ZsetAdd(key, index, data, NULL);
RedisModule_CloseKey(key);
// Publish a message on the requested pubsub channel if necessary.
TablePubsub pubsub_channel = ParseTablePubsub(pubsub_channel_str);
if (pubsub_channel != TablePubsub_NO_PUBLISH) {
// All other pubsub channels write the data back directly onto the channel.
return PublishTableAdd(ctx, pubsub_channel_str, id, data);
} else {
return RedisModule_ReplyWithSimpleString(ctx, "OK");
}
}
/// A helper function to create and finish a GcsTableEntry, based on the
/// current value or values at the given key.
void TableEntryToFlatbuf(RedisModuleKey *table_key,
RedisModuleString *entry_id,
flatbuffers::FlatBufferBuilder &fbb) {
auto key_type = RedisModule_KeyType(table_key);
switch (key_type) {
case REDISMODULE_KEYTYPE_STRING: {
// Build the flatbuffer from the string data.
size_t data_len = 0;
char *data_buf =
RedisModule_StringDMA(table_key, &data_len, REDISMODULE_READ);
auto data = fbb.CreateString(data_buf, data_len);
auto message = CreateGcsTableEntry(fbb, RedisStringToFlatbuf(fbb, entry_id),
fbb.CreateVector(&data, 1));
fbb.Finish(message);
} break;
case REDISMODULE_KEYTYPE_ZSET: {
// Build the flatbuffer from the set of log entries.
RAY_CHECK(RedisModule_ZsetFirstInScoreRange(
table_key, REDISMODULE_NEGATIVE_INFINITE,
REDISMODULE_POSITIVE_INFINITE, 1, 1) == REDISMODULE_OK);
std::vector<flatbuffers::Offset<flatbuffers::String>> data;
for (; !RedisModule_ZsetRangeEndReached(table_key);
RedisModule_ZsetRangeNext(table_key)) {
data.push_back(RedisStringToFlatbuf(
fbb, RedisModule_ZsetRangeCurrentElement(table_key, NULL)));
}
auto message = CreateGcsTableEntry(fbb, RedisStringToFlatbuf(fbb, entry_id),
fbb.CreateVector(data));
fbb.Finish(message);
} break;
default:
RAY_LOG(FATAL) << "Invalid Redis type during lookup: " << key_type;
}
}
/// Lookup the current value or values at a key. Returns the current value or
/// values at the key.
///
/// This is called from a client with the command:
//
/// RAY.TABLE_LOOKUP <table_prefix> <pubsub_channel> <id>
///
/// \param table_prefix The prefix string for keys in this table.
/// \param pubsub_channel The pubsub channel name that notifications for
/// this key should be published to. This field is unused for lookups.
/// \param id The ID of the key to lookup.
/// \return nil if the key is empty, the current value if the key type is a
/// string, or an array of the current values if the key type is a set.
int TableLookup_RedisCommand(RedisModuleCtx *ctx,
RedisModuleString **argv,
int argc) {
if (argc != 4) {
if (argc < 4) {
return RedisModule_WrongArity(ctx);
}
RedisModuleString *prefix_str = argv[1];
RedisModuleString *id = argv[3];
RedisModuleKey *key = OpenPrefixedKey(ctx, prefix_str, id, REDISMODULE_READ);
if (key == nullptr) {
return RedisModule_ReplyWithNull(ctx);
// Lookup the data at the key.
RedisModuleKey *table_key =
OpenPrefixedKey(ctx, prefix_str, id, REDISMODULE_READ);
if (table_key == nullptr) {
RedisModule_ReplyWithNull(ctx);
} else {
// Serialize the data to a flatbuffer to return to the client.
flatbuffers::FlatBufferBuilder fbb;
TableEntryToFlatbuf(table_key, id, fbb);
RedisModule_ReplyWithStringBuffer(
ctx, reinterpret_cast<const char *>(fbb.GetBufferPointer()),
fbb.GetSize());
}
size_t len = 0;
const char *buf = RedisModule_StringDMA(key, &len, REDISMODULE_READ);
RedisModule_ReplyWithStringBuffer(ctx, buf, len);
RedisModule_CloseKey(key);
RedisModule_CloseKey(table_key);
return REDISMODULE_OK;
}
@@ -706,7 +728,8 @@ int TableLookup_RedisCommand(RedisModuleCtx *ctx,
/// client, the channel name should be <pubsub_channel>:<client_id>.
/// \param id The ID of the key to publish notifications for.
/// \param client_id The ID of the client that is being notified.
/// \return The current value at the key, or OK if there is no value.
/// \return nil if the key is empty, the current value if the key type is a
/// string, or an array of the current values if the key type is a set.
int TableRequestNotifications_RedisCommand(RedisModuleCtx *ctx,
RedisModuleString **argv,
int argc) {
@@ -728,36 +751,27 @@ int TableRequestNotifications_RedisCommand(RedisModuleCtx *ctx,
CHECK_ERROR(RedisModule_ZsetAdd(notification_key, 0.0, client_channel, NULL),
"ZsetAdd failed.");
RedisModule_CloseKey(notification_key);
RedisModule_FreeString(ctx, client_channel);
// Return the current value at the key, if any, to the client that requested
// a notification.
// Lookup the current value at the key.
RedisModuleKey *table_key =
OpenPrefixedKey(ctx, prefix_str, id, REDISMODULE_READ);
if (table_key != nullptr) {
// Serialize the notification to send.
size_t data_len = 0;
char *data_buf =
RedisModule_StringDMA(table_key, &data_len, REDISMODULE_READ);
// Publish the current value at the key to the client that is requesting
// notifications.
flatbuffers::FlatBufferBuilder fbb;
auto message = CreateGcsNotification(fbb, RedisStringToFlatbuf(fbb, id),
fbb.CreateString(data_buf, data_len));
fbb.Finish(message);
int result = RedisModule_ReplyWithStringBuffer(
ctx, reinterpret_cast<const char *>(fbb.GetBufferPointer()),
TableEntryToFlatbuf(table_key, id, fbb);
RedisModule_Call(ctx, "PUBLISH", "sb", client_channel, reinterpret_cast<const char *>(fbb.GetBufferPointer()),
fbb.GetSize());
RedisModule_CloseKey(table_key);
return result;
} else {
RedisModule_CloseKey(table_key);
RedisModule_ReplyWithSimpleString(ctx, "OK");
return REDISMODULE_OK;
}
RedisModule_CloseKey(table_key);
RedisModule_FreeString(ctx, client_channel);
return RedisModule_ReplyWithNull(ctx);
}
/// Cancel notifications for changes to a key. The client will no longer
/// receive notifications for this key.
/// receive notifications for this key. This does not check if the client
/// first requested notifications before canceling them.
///
/// This is called from a client with the command:
//
@@ -769,9 +783,8 @@ int TableRequestNotifications_RedisCommand(RedisModuleCtx *ctx,
/// this key should be published to. If publishing to a specific client,
/// then the channel name should be <pubsub_channel>:<client_id>.
/// \param id The ID of the key to publish notifications for.
/// \param client_id The ID of the client that is being notified.
/// \return OK if the requesting client was removed, or an error if the client
/// was not found.
/// \param client_id The ID of the client to cancel notifications for.
/// \return OK.
int TableCancelNotifications_RedisCommand(RedisModuleCtx *ctx,
RedisModuleString **argv,
int argc) {
@@ -789,10 +802,10 @@ int TableCancelNotifications_RedisCommand(RedisModuleCtx *ctx,
// there are changes to the key.
RedisModuleKey *notification_key = OpenBroadcastKey(
ctx, pubsub_channel_str, id, REDISMODULE_READ | REDISMODULE_WRITE);
RAY_CHECK(RedisModule_KeyType(notification_key) != REDISMODULE_KEYTYPE_EMPTY);
int deleted;
RedisModule_ZsetRem(notification_key, client_channel, &deleted);
RAY_CHECK(deleted);
if (RedisModule_KeyType(notification_key) != REDISMODULE_KEYTYPE_EMPTY) {
RAY_CHECK(RedisModule_ZsetRem(notification_key, client_channel, NULL) ==
REDISMODULE_OK);
}
RedisModule_CloseKey(notification_key);
RedisModule_ReplyWithSimpleString(ctx, "OK");
@@ -1661,6 +1674,12 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx,
return REDISMODULE_ERR;
}
if (RedisModule_CreateCommand(ctx, "ray.table_append",
TableAppend_RedisCommand, "write", 0, 0,
0) == REDISMODULE_ERR) {
return REDISMODULE_ERR;
}
if (RedisModule_CreateCommand(ctx, "ray.table_lookup",
TableLookup_RedisCommand, "readonly", 0, 0,
0) == REDISMODULE_ERR) {
src/ray/gcs/client_test.cc
+2 -1
View File
@@ -420,9 +420,10 @@ TEST_F(TestGcsWithAsio, TestSubscribeCancel) {
TestSubscribeCancel(job_id_, client_);
}
void ClientTableNotification(gcs::AsyncGcsClient *client, const UniqueID &id,
void ClientTableNotification(gcs::AsyncGcsClient *client, const ClientID &client_id,
const ClientTableDataT &data, bool is_insertion) {
ClientID added_id = client->client_table().GetLocalClientId();
ASSERT_EQ(client_id, added_id);
ASSERT_EQ(ClientID::from_binary(data.client_id), added_id);
ASSERT_EQ(data.is_insertion, is_insertion);
src/ray/gcs/format/gcs.fbs
+2 -2
View File
@@ -23,9 +23,9 @@ enum TablePubsub:int {
ACTOR
}
table GcsNotification {
table GcsTableEntry {
id: string;
data: string;
entries: [string];
}
table FunctionTableData {
src/ray/gcs/redis_context.cc
+5 -6
View File
@@ -15,7 +15,7 @@ namespace {
/// A helper function to call the callback and delete it from the callback
/// manager if necessary.
void ProcessCallback(int64_t callback_index, const std::vector<std::string> &data) {
void ProcessCallback(int64_t callback_index, const std::string &data) {
if (callback_index >= 0) {
bool delete_callback =
ray::gcs::RedisCallbackManager::instance().get(callback_index)(data);
@@ -40,14 +40,14 @@ void GlobalRedisCallback(void *c, void *r, void *privdata) {
}
int64_t callback_index = reinterpret_cast<int64_t>(privdata);
redisReply *reply = reinterpret_cast<redisReply *>(r);
std::vector<std::string> data;
std::string data = "";
// Parse the response.
switch (reply->type) {
case (REDIS_REPLY_NIL): {
// Do not add any data for a nil response.
} break;
case (REDIS_REPLY_STRING): {
data.push_back(std::string(reply->str, reply->len));
data = std::string(reply->str, reply->len);
} break;
case (REDIS_REPLY_STATUS): {
} break;
@@ -67,7 +67,7 @@ void SubscribeRedisCallback(void *c, void *r, void *privdata) {
}
int64_t callback_index = reinterpret_cast<int64_t>(privdata);
redisReply *reply = reinterpret_cast<redisReply *>(r);
std::vector<std::string> data;
std::string data = "";
// Parse the response.
switch (reply->type) {
case (REDIS_REPLY_ARRAY): {
@@ -76,13 +76,12 @@ void SubscribeRedisCallback(void *c, void *r, void *privdata) {
if (strcmp(message_type->str, "subscribe") == 0) {
// If the message is for the initial subscription call, return the empty
// string as a response to signify that subscription was successful.
data.push_back("");
} else if (strcmp(message_type->str, "message") == 0) {
// If the message is from a PUBLISH, make sure the data is nonempty.
redisReply *message = reply->element[reply->elements - 1];
auto notification = std::string(message->str, message->len);
RAY_CHECK(!notification.empty()) << "Empty message received on subscribe channel";
data.push_back(notification);
data = notification;
} else {
RAY_LOG(FATAL) << "Fatal redis error during subscribe" << message_type->str;
}
src/ray/gcs/redis_context.h
+1 -1
View File
@@ -24,7 +24,7 @@ class RedisCallbackManager {
/// Every callback should take in a vector of the results from the Redis
/// operation and return a bool indicating whether the callback should be
/// deleted once called.
using RedisCallback = std::function<bool(const std::vector<std::string> &)>;
using RedisCallback = std::function<bool(const std::string &)>;
static RedisCallbackManager &instance() {
static RedisCallbackManager instance;
src/ray/gcs/tables.cc
+162 -95
View File
@@ -7,13 +7,124 @@ namespace ray {
namespace gcs {
template <typename ID, typename Data>
Status Log<ID, Data>::Append(const JobID &job_id, const ID &id,
std::shared_ptr<DataT> data, const Callback &done) {
auto d = std::shared_ptr<CallbackData>(
new CallbackData({id, data, done, nullptr, this, client_}));
int64_t callback_index =
RedisCallbackManager::instance().add([d](const std::string &data) {
if (d->callback != nullptr) {
(d->callback)(d->client, d->id, {*d->data});
}
return true;
});
flatbuffers::FlatBufferBuilder fbb;
fbb.ForceDefaults(true);
fbb.Finish(Data::Pack(fbb, data.get()));
return context_->RunAsync("RAY.TABLE_APPEND", id, fbb.GetBufferPointer(), fbb.GetSize(),
prefix_, pubsub_channel_, callback_index);
}
template <typename ID, typename Data>
Status Log<ID, Data>::Lookup(const JobID &job_id, const ID &id, const Callback &lookup) {
auto d = std::shared_ptr<CallbackData>(
new CallbackData({id, nullptr, lookup, nullptr, this, client_}));
int64_t callback_index =
RedisCallbackManager::instance().add([d](const std::string &data) {
if (d->callback != nullptr) {
std::vector<DataT> results;
if (!data.empty()) {
auto root = flatbuffers::GetRoot<GcsTableEntry>(data.data());
RAY_CHECK(from_flatbuf(*root->id()) == d->id);
for (size_t i = 0; i < root->entries()->size(); i++) {
DataT result;
auto data_root =
flatbuffers::GetRoot<Data>(root->entries()->Get(i)->data());
data_root->UnPackTo(&result);
results.emplace_back(std::move(result));
}
}
(d->callback)(d->client, d->id, results);
}
return true;
});
std::vector<uint8_t> nil;
return context_->RunAsync("RAY.TABLE_LOOKUP", id, nil.data(), nil.size(), prefix_,
pubsub_channel_, callback_index);
}
template <typename ID, typename Data>
Status Log<ID, Data>::Subscribe(const JobID &job_id, const ClientID &client_id,
const Callback &subscribe,
const SubscriptionCallback &done) {
RAY_CHECK(subscribe_callback_index_ == -1)
<< "Client called Subscribe twice on the same table";
auto d = std::shared_ptr<CallbackData>(
new CallbackData({client_id, nullptr, subscribe, done, this, client_}));
int64_t callback_index = RedisCallbackManager::instance().add(
[this, d](const std::string &data) {
if (data.empty()) {
// No notification data is provided. This is the callback for the
// initial subscription request.
if (d->subscription_callback != nullptr) {
(d->subscription_callback)(d->client);
}
} else {
// Data is provided. This is the callback for a message.
if (d->callback != nullptr) {
// Parse the notification.
auto root = flatbuffers::GetRoot<GcsTableEntry>(data.data());
ID id = UniqueID::nil();
if (root->id()->size() > 0) {
id = from_flatbuf(*root->id());
}
std::vector<DataT> results;
for (size_t i = 0; i < root->entries()->size(); i++) {
DataT result;
auto data_root =
flatbuffers::GetRoot<Data>(root->entries()->Get(i)->data());
data_root->UnPackTo(&result);
results.emplace_back(std::move(result));
}
(d->callback)(d->client, id, results);
}
}
// We do not delete the callback after calling it since there may be
// more subscription messages.
return false;
});
subscribe_callback_index_ = callback_index;
return context_->SubscribeAsync(client_id, pubsub_channel_, callback_index);
}
template <typename ID, typename Data>
Status Log<ID, Data>::RequestNotifications(const JobID &job_id, const ID &id,
const ClientID &client_id) {
RAY_CHECK(subscribe_callback_index_ >= 0)
<< "Client requested notifications on a key before Subscribe completed";
return context_->RunAsync("RAY.TABLE_REQUEST_NOTIFICATIONS", id, client_id.data(),
client_id.size(), prefix_, pubsub_channel_,
/*callback_index=*/-1);
}
template <typename ID, typename Data>
Status Log<ID, Data>::CancelNotifications(const JobID &job_id, const ID &id,
const ClientID &client_id) {
RAY_CHECK(subscribe_callback_index_ >= 0)
<< "Client canceled notifications on a key before Subscribe completed";
return context_->RunAsync("RAY.TABLE_CANCEL_NOTIFICATIONS", id, client_id.data(),
client_id.size(), prefix_, pubsub_channel_,
/*callback_index=*/-1);
}
template <typename ID, typename Data>
Status Table<ID, Data>::Add(const JobID &job_id, const ID &id,
std::shared_ptr<DataT> data, const Callback &done) {
auto d = std::shared_ptr<CallbackData>(
new CallbackData({id, data, done, nullptr, nullptr, this, client_}));
new CallbackData({id, data, done, nullptr, this, client_}));
int64_t callback_index =
RedisCallbackManager::instance().add([d](const std::vector<std::string> &data) {
RedisCallbackManager::instance().add([d](const std::string &data) {
if (d->callback != nullptr) {
(d->callback)(d->client, d->id, *d->data);
}
@@ -29,88 +140,33 @@ Status Table<ID, Data>::Add(const JobID &job_id, const ID &id,
template <typename ID, typename Data>
Status Table<ID, Data>::Lookup(const JobID &job_id, const ID &id, const Callback &lookup,
const FailureCallback &failure) {
auto d = std::shared_ptr<CallbackData>(
new CallbackData({id, nullptr, lookup, failure, nullptr, this, client_}));
int64_t callback_index =
RedisCallbackManager::instance().add([d](const std::vector<std::string> &data) {
if (data.empty()) {
if (d->failure != nullptr) {
(d->failure)(d->client, d->id);
}
} else {
RAY_CHECK(data.size() == 1);
if (d->callback != nullptr) {
DataT result;
auto root = flatbuffers::GetRoot<Data>(data[0].data());
root->UnPackTo(&result);
(d->callback)(d->client, d->id, result);
}
}
return true;
});
std::vector<uint8_t> nil;
return context_->RunAsync("RAY.TABLE_LOOKUP", id, nil.data(), nil.size(), prefix_,
pubsub_channel_, callback_index);
return Log<ID, Data>::Lookup(job_id, id,
[lookup, failure](AsyncGcsClient *client, const ID &id,
const std::vector<DataT> &data) {
if (data.empty()) {
if (failure != nullptr) {
(failure)(client, id);
}
} else {
RAY_CHECK(data.size() == 1);
if (lookup != nullptr) {
(lookup)(client, id, data[0]);
}
}
});
}
template <typename ID, typename Data>
Status Table<ID, Data>::Subscribe(const JobID &job_id, const ClientID &client_id,
const Callback &subscribe,
const SubscriptionCallback &done) {
RAY_CHECK(subscribe_callback_index_ == -1)
<< "Client called Subscribe twice on the same table";
auto d = std::shared_ptr<CallbackData>(
new CallbackData({client_id, nullptr, subscribe, nullptr, done, this, client_}));
int64_t callback_index = RedisCallbackManager::instance().add(
[this, d](const std::vector<std::string> &data) {
if (data.size() == 1 && data[0] == "") {
// No notification data is provided. This is the callback for the
// initial subscription request.
if (d->subscription_callback != nullptr) {
(d->subscription_callback)(d->client);
}
} else {
// Data is provided. This is the callback for a message.
RAY_CHECK(data.size() == 1);
if (d->callback != nullptr) {
// Parse the notification.
auto notification = flatbuffers::GetRoot<GcsNotification>(data[0].data());
ID id = UniqueID::nil();
if (notification->id()->size() > 0) {
id = from_flatbuf(*notification->id());
}
DataT result;
auto root = flatbuffers::GetRoot<Data>(notification->data()->data());
root->UnPackTo(&result);
(d->callback)(d->client, id, result);
}
}
// We do not delete the callback after calling it since there may be
// more subscription messages.
return false;
});
subscribe_callback_index_ = callback_index;
return context_->SubscribeAsync(client_id, pubsub_channel_, callback_index);
}
template <typename ID, typename Data>
Status Table<ID, Data>::RequestNotifications(const JobID &job_id, const ID &id,
const ClientID &client_id) {
RAY_CHECK(subscribe_callback_index_ >= 0)
<< "Client requested notifications on a key before Subscribe completed";
return context_->RunAsync("RAY.TABLE_REQUEST_NOTIFICATIONS", id, client_id.data(),
client_id.size(), prefix_, pubsub_channel_,
subscribe_callback_index_);
}
template <typename ID, typename Data>
Status Table<ID, Data>::CancelNotifications(const JobID &job_id, const ID &id,
const ClientID &client_id) {
RAY_CHECK(subscribe_callback_index_ >= 0)
<< "Client canceled notifications on a key before Subscribe completed";
return context_->RunAsync("RAY.TABLE_CANCEL_NOTIFICATIONS", id, client_id.data(),
client_id.size(), prefix_, pubsub_channel_,
/*callback_index=*/-1);
return Log<ID, Data>::Subscribe(
job_id, client_id,
[subscribe](AsyncGcsClient *client, const ID &id, const std::vector<DataT> &data) {
RAY_CHECK(data.size() == 1);
subscribe(client, id, data[0]);
},
done);
}
void ClientTable::RegisterClientAddedCallback(const ClientTableCallback &callback) {
@@ -118,7 +174,7 @@ void ClientTable::RegisterClientAddedCallback(const ClientTableCallback &callbac
// Call the callback for any added clients that are cached.
for (const auto &entry : client_cache_) {
if (!entry.first.is_nil() && entry.second.is_insertion) {
client_added_callback_(client_, ClientID::nil(), entry.second);
client_added_callback_(client_, entry.first, entry.second);
}
}
}
@@ -128,12 +184,12 @@ void ClientTable::RegisterClientRemovedCallback(const ClientTableCallback &callb
// Call the callback for any removed clients that are cached.
for (const auto &entry : client_cache_) {
if (!entry.first.is_nil() && !entry.second.is_insertion) {
client_removed_callback_(client_, ClientID::nil(), entry.second);
client_removed_callback_(client_, entry.first, entry.second);
}
}
}
void ClientTable::HandleNotification(AsyncGcsClient *client, const ClientID &channel_id,
void ClientTable::HandleNotification(AsyncGcsClient *client,
const ClientTableDataT &data) {
ClientID client_id = ClientID::from_binary(data.client_id);
// It's possible to get duplicate notifications from the client table, so
@@ -176,9 +232,10 @@ void ClientTable::HandleNotification(AsyncGcsClient *client, const ClientID &cha
}
}
void ClientTable::HandleConnected(AsyncGcsClient *client, const ClientID &client_id,
const ClientTableDataT &data) {
RAY_CHECK(client_id == client_id_) << client_id.hex() << " " << client_id_.hex();
void ClientTable::HandleConnected(AsyncGcsClient *client, const ClientTableDataT &data) {
auto connected_client_id = ClientID::from_binary(data.client_id);
RAY_CHECK(client_id_ == connected_client_id) << connected_client_id.hex() << " "
<< client_id_.hex();
}
const ClientID &ClientTable::GetLocalClientId() { return client_id_; }
@@ -191,15 +248,21 @@ Status ClientTable::Connect() {
auto data = std::make_shared<ClientTableDataT>(local_client_);
data->is_insertion = true;
// Callback for a notification from the client table.
auto notification_callback = [this](AsyncGcsClient *client, const ClientID &channel_id,
const ClientTableDataT &data) {
return HandleNotification(client, channel_id, data);
auto notification_callback = [this](
AsyncGcsClient *client, const UniqueID &log_key,
const std::vector<ClientTableDataT> &notifications) {
RAY_CHECK(log_key == client_log_key_);
for (auto &notification : notifications) {
HandleNotification(client, notification);
}
};
// Callback to handle our own successful connection once we've added
// ourselves.
auto add_callback = [this](AsyncGcsClient *client, const ClientID &id,
const ClientTableDataT &data) {
HandleConnected(client, id, data);
auto add_callback = [this](AsyncGcsClient *client, const UniqueID &log_key,
const std::vector<ClientTableDataT> &data) {
RAY_CHECK(log_key == client_log_key_);
RAY_CHECK(data.size() == 1);
HandleConnected(client, data[0]);
};
// Callback to add ourselves once we've successfully subscribed.
auto subscription_callback = [this, data, add_callback](AsyncGcsClient *c) {
@@ -208,9 +271,10 @@ Status ClientTable::Connect() {
if (disconnected_) {
data->is_insertion = false;
}
return Add(JobID::nil(), client_id_, data, add_callback);
RAY_CHECK_OK(RequestNotifications(JobID::nil(), client_log_key_, client_id_));
RAY_CHECK_OK(Append(JobID::nil(), client_log_key_, data, add_callback));
};
return Subscribe(JobID::nil(), ClientID::nil(), notification_callback,
return Subscribe(JobID::nil(), client_id_, notification_callback,
subscription_callback);
}
@@ -218,10 +282,12 @@ Status ClientTable::Disconnect() {
auto data = std::make_shared<ClientTableDataT>(local_client_);
data->is_insertion = true;
auto add_callback = [this](AsyncGcsClient *client, const ClientID &id,
const ClientTableDataT &data) {
HandleConnected(client, id, data);
const std::vector<ClientTableDataT> &data) {
RAY_CHECK(data.size() == 1);
HandleConnected(client, data[0]);
RAY_CHECK_OK(CancelNotifications(JobID::nil(), client_log_key_, id));
};
RAY_RETURN_NOT_OK(Add(JobID::nil(), client_id_, data, add_callback));
RAY_RETURN_NOT_OK(Append(JobID::nil(), client_log_key_, data, add_callback));
// We successfully added the deletion entry. Mark ourselves as disconnected.
disconnected_ = true;
return Status::OK();
@@ -239,6 +305,7 @@ const ClientTableDataT &ClientTable::GetClient(const ClientID &client_id) {
}
}
template class Log<ObjectID, ObjectTableData>;
template class Table<TaskID, ray::protocol::Task>;
template class Table<TaskID, TaskTableData>;
template class Table<ObjectID, ObjectTableData>;
src/ray/gcs/tables.h
+114 -30
View File
@@ -27,15 +27,20 @@ class RedisContext;
class AsyncGcsClient;
/// \class Log
///
/// A GCS table where every entry is an append-only log.
/// Example tables backed by Log:
/// ObjectTable: Stores a log of which clients have added or evicted an
/// object.
/// ClientTable: Stores a log of which GCS clients have been added or deleted
/// from the system.
template <typename ID, typename Data>
class Table {
class Log {
public:
using DataT = typename Data::NativeTableType;
using Callback =
std::function<void(AsyncGcsClient *client, const ID &id, const DataT &data)>;
/// The callback to call when a lookup fails because there is no entry at the
/// key.
using FailureCallback = std::function<void(AsyncGcsClient *client, const ID &id)>;
using Callback = std::function<void(AsyncGcsClient *client, const ID &id,
const std::vector<DataT> &data)>;
/// The callback to call when a SUBSCRIBE call completes and we are ready to
/// request and receive notifications.
using SubscriptionCallback = std::function<void(AsyncGcsClient *client)>;
@@ -44,45 +49,44 @@ class Table {
ID id;
std::shared_ptr<DataT> data;
Callback callback;
FailureCallback failure;
// An optional callback to call for subscription operations, where the
// first message is a notification of subscription success.
SubscriptionCallback subscription_callback;
Table<ID, Data> *table;
Log<ID, Data> *log;
AsyncGcsClient *client;
};
Table(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client)
Log(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client)
: context_(context),
client_(client),
pubsub_channel_(TablePubsub_NO_PUBLISH),
prefix_(TablePrefix_UNUSED),
subscribe_callback_index_(-1){};
/// Add an entry to the table.
/// Append a log entry to a key.
///
/// \param job_id The ID of the job (= driver).
/// \param id The ID of the data that is added to the GCS.
/// \param data Data that is added to the GCS.
/// \param data Data to append to the log.
/// \param done Callback that is called once the data has been written to the
/// GCS.
/// \return Status
Status Add(const JobID &job_id, const ID &id, std::shared_ptr<DataT> data,
const Callback &done);
Status Append(const JobID &job_id, const ID &id, std::shared_ptr<DataT> data,
const Callback &done);
/// Lookup an entry asynchronously.
/// Lookup the log values at a key asynchronously.
///
/// \param job_id The ID of the job (= driver).
/// \param id The ID of the data that is looked up in the GCS.
/// \param lookup Callback that is called after lookup. If the callback is
/// called with an empty vector, then there was no data at the key.
/// \return Status
Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup,
const FailureCallback &failure);
Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup);
/// Subscribe to any Add operations to this table. The caller may choose to
/// subscribe to all Adds, or to subscribe only to keys that it requests
/// notifications for. This may only be called once per Table instance.
/// Subscribe to any Append operations to this table. The caller may choose
/// to subscribe to all Appends, or to subscribe only to keys that it
/// requests notifications for. This may only be called once per Log
/// instance.
///
/// \param job_id The ID of the job (= driver).
/// \param client_id The type of update to listen to. If this is nil, then a
@@ -103,8 +107,8 @@ class Table {
///
/// The notifications will be returned via the subscribe callback that was
/// registered by `Subscribe`. An initial notification will be returned for
/// the current value(s) at the key, if any, and a subsequent notification
/// will be published for every following `Add` to the key. Before
/// the current values at the key, if any, and a subsequent notification will
/// be published for every following `Append` to the key. Before
/// notifications can be requested, the caller must first call `Subscribe`,
/// with the same `client_id`.
///
@@ -143,6 +147,73 @@ class Table {
int64_t subscribe_callback_index_;
};
/// \class Table
///
/// A GCS table where every entry is a single data item.
/// Example tables backed by Log:
/// TaskTable: Stores Task metadata needed for executing the task.
template <typename ID, typename Data>
class Table : private Log<ID, Data> {
public:
using DataT = typename Log<ID, Data>::DataT;
using Callback =
std::function<void(AsyncGcsClient *client, const ID &id, const DataT &data)>;
/// The callback to call when a Lookup call returns an empty entry.
using FailureCallback = std::function<void(AsyncGcsClient *client, const ID &id)>;
/// The callback to call when a Subscribe call completes and we are ready to
/// request and receive notifications.
using SubscriptionCallback = typename Log<ID, Data>::SubscriptionCallback;
struct CallbackData {
ID id;
std::shared_ptr<DataT> data;
Callback callback;
// An optional callback to call for subscription operations, where the
// first message is a notification of subscription success.
SubscriptionCallback subscription_callback;
Log<ID, Data> *log;
AsyncGcsClient *client;
};
Table(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client)
: Log<ID, Data>(context, client) {}
using Log<ID, Data>::RequestNotifications;
using Log<ID, Data>::CancelNotifications;
/// Add an entry to the table. This overwrites any existing data at the key.
///
/// \param job_id The ID of the job (= driver).
/// \param id The ID of the data that is added to the GCS.
/// \param data Data that is added to the GCS.
/// \param done Callback that is called once the data has been written to the
/// GCS.
/// \return Status
Status Add(const JobID &job_id, const ID &id, std::shared_ptr<DataT> data,
const Callback &done);
/// Lookup an entry asynchronously.
///
/// \param job_id The ID of the job (= driver).
/// \param id The ID of the data that is looked up in the GCS.
/// \param lookup Callback that is called after lookup if there was data the
/// key.
/// \param failure Callback that is called after lookup if there was no data
/// at the key.
/// \return Status
Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup,
const FailureCallback &failure);
Status Subscribe(const JobID &job_id, const ClientID &client_id,
const Callback &subscribe, const SubscriptionCallback &done);
protected:
using Log<ID, Data>::context_;
using Log<ID, Data>::client_;
using Log<ID, Data>::pubsub_channel_;
using Log<ID, Data>::prefix_;
};
class ObjectTable : public Table<ObjectID, ObjectTableData> {
public:
ObjectTable(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client)
@@ -210,10 +281,9 @@ class TaskTable : public Table<TaskID, TaskTableData> {
std::shared_ptr<TaskTableTestAndUpdateT> data,
const TestAndUpdateCallback &callback) {
int64_t callback_index = RedisCallbackManager::instance().add(
[this, callback, id](const std::vector<std::string> &data) {
RAY_CHECK(data.size() == 1);
[this, callback, id](const std::string &data) {
auto result = std::make_shared<TaskTableDataT>();
auto root = flatbuffers::GetRoot<TaskTableData>(data[0].data());
auto root = flatbuffers::GetRoot<TaskTableData>(data.data());
root->UnPackTo(result.get());
callback(client_, id, *result, root->updated());
return true;
@@ -263,13 +333,25 @@ Status TaskTableTestAndUpdate(AsyncGcsClient *gcs_client, const TaskID &task_id,
SchedulingState update_state,
const TaskTable::TestAndUpdateCallback &callback);
class ClientTable : private Table<ClientID, ClientTableData> {
/// \class ClientTable
///
/// The ClientTable stores information about active and inactive clients. It is
/// structured as a single log stored at a key known to all clients. When a
/// client connects, it appends an entry to the log indicating that it is
/// alive. When a client disconnects, or if another client detects its failure,
/// it should append an entry to the log indicating that it is dead. A client
/// that is marked as dead should never again be marked as alive; if it needs
/// to reconnect, it must connect with a different ClientID.
class ClientTable : private Log<UniqueID, ClientTableData> {
public:
using ClientTableCallback = std::function<void(
AsyncGcsClient *client, const ClientID &id, const ClientTableDataT &data)>;
ClientTable(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client,
const ClientTableDataT &local_client)
: Table(context, client),
: Log(context, client),
// We set the client log's key equal to nil so that all instances of
// ClientTable have the same key.
client_log_key_(UniqueID::nil()),
disconnected_(false),
client_id_(ClientID::from_binary(local_client.client_id)),
local_client_(local_client) {
@@ -325,12 +407,14 @@ class ClientTable : private Table<ClientID, ClientTableData> {
private:
/// Handle a client table notification.
void HandleNotification(AsyncGcsClient *client, const ClientID &channel_id,
const ClientTableDataT &notifications);
void HandleNotification(AsyncGcsClient *client, const ClientTableDataT &notifications);
/// Handle this client's successful connection to the GCS.
void HandleConnected(AsyncGcsClient *client, const ClientID &client_id,
const ClientTableDataT &notifications);
void HandleConnected(AsyncGcsClient *client, const ClientTableDataT &notifications);
/// The key at which the log of client information is stored. This key must
/// be kept the same across all instances of the ClientTable, so that all
/// clients append and read from the same key.
UniqueID client_log_key_;
/// Whether this client has called Disconnect().
bool disconnected_;
/// This client's ID.