From 0fd4112354d433105a8945f51ffc2b25db16d169 Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Mon, 26 Mar 2018 16:00:43 -0700 Subject: [PATCH] Introduce a log interface for the new GCS (#1771) * 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 --- src/common/redis_module/ray_redis_module.cc | 223 +++++++++-------- src/ray/gcs/client_test.cc | 3 +- src/ray/gcs/format/gcs.fbs | 4 +- src/ray/gcs/redis_context.cc | 11 +- src/ray/gcs/redis_context.h | 2 +- src/ray/gcs/tables.cc | 257 ++++++++++++-------- src/ray/gcs/tables.h | 144 ++++++++--- 7 files changed, 407 insertions(+), 237 deletions(-) diff --git a/src/common/redis_module/ray_redis_module.cc b/src/common/redis_module/ray_redis_module.cc index 3db89c433..fa952e653 100644 --- a/src/common/redis_module/ray_redis_module.cc +++ b/src/common/redis_module/ray_redis_module.cc @@ -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(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(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> 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 +/// +/// \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(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 :. /// \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(fbb.GetBufferPointer()), + TableEntryToFlatbuf(table_key, id, fbb); + RedisModule_Call(ctx, "PUBLISH", "sb", client_channel, reinterpret_cast(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 :. /// \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) { diff --git a/src/ray/gcs/client_test.cc b/src/ray/gcs/client_test.cc index 659817f60..a58d09bea 100644 --- a/src/ray/gcs/client_test.cc +++ b/src/ray/gcs/client_test.cc @@ -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); diff --git a/src/ray/gcs/format/gcs.fbs b/src/ray/gcs/format/gcs.fbs index e0ef9ade2..2025cc46f 100644 --- a/src/ray/gcs/format/gcs.fbs +++ b/src/ray/gcs/format/gcs.fbs @@ -23,9 +23,9 @@ enum TablePubsub:int { ACTOR } -table GcsNotification { +table GcsTableEntry { id: string; - data: string; + entries: [string]; } table FunctionTableData { diff --git a/src/ray/gcs/redis_context.cc b/src/ray/gcs/redis_context.cc index 3a5cf38f9..294270a36 100644 --- a/src/ray/gcs/redis_context.cc +++ b/src/ray/gcs/redis_context.cc @@ -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 &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(privdata); redisReply *reply = reinterpret_cast(r); - std::vector 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(privdata); redisReply *reply = reinterpret_cast(r); - std::vector 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; } diff --git a/src/ray/gcs/redis_context.h b/src/ray/gcs/redis_context.h index 5b90265b0..fadd2fd7c 100644 --- a/src/ray/gcs/redis_context.h +++ b/src/ray/gcs/redis_context.h @@ -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 &)>; + using RedisCallback = std::function; static RedisCallbackManager &instance() { static RedisCallbackManager instance; diff --git a/src/ray/gcs/tables.cc b/src/ray/gcs/tables.cc index 8233efefc..9b5916318 100644 --- a/src/ray/gcs/tables.cc +++ b/src/ray/gcs/tables.cc @@ -7,13 +7,124 @@ namespace ray { namespace gcs { +template +Status Log::Append(const JobID &job_id, const ID &id, + std::shared_ptr data, const Callback &done) { + auto d = std::shared_ptr( + 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 +Status Log::Lookup(const JobID &job_id, const ID &id, const Callback &lookup) { + auto d = std::shared_ptr( + 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 results; + if (!data.empty()) { + auto root = flatbuffers::GetRoot(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(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 nil; + return context_->RunAsync("RAY.TABLE_LOOKUP", id, nil.data(), nil.size(), prefix_, + pubsub_channel_, callback_index); +} + +template +Status Log::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( + 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(data.data()); + ID id = UniqueID::nil(); + if (root->id()->size() > 0) { + id = from_flatbuf(*root->id()); + } + std::vector results; + for (size_t i = 0; i < root->entries()->size(); i++) { + DataT result; + auto data_root = + flatbuffers::GetRoot(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 +Status Log::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 +Status Log::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 Status Table::Add(const JobID &job_id, const ID &id, std::shared_ptr data, const Callback &done) { auto d = std::shared_ptr( - 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 &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::Add(const JobID &job_id, const ID &id, template Status Table::Lookup(const JobID &job_id, const ID &id, const Callback &lookup, const FailureCallback &failure) { - auto d = std::shared_ptr( - new CallbackData({id, nullptr, lookup, failure, nullptr, this, client_})); - int64_t callback_index = - RedisCallbackManager::instance().add([d](const std::vector &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[0].data()); - root->UnPackTo(&result); - (d->callback)(d->client, d->id, result); - } - } - return true; - }); - std::vector nil; - return context_->RunAsync("RAY.TABLE_LOOKUP", id, nil.data(), nil.size(), prefix_, - pubsub_channel_, callback_index); + return Log::Lookup(job_id, id, + [lookup, failure](AsyncGcsClient *client, const ID &id, + const std::vector &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 Status Table::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( - new CallbackData({client_id, nullptr, subscribe, nullptr, done, this, client_})); - int64_t callback_index = RedisCallbackManager::instance().add( - [this, d](const std::vector &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(data[0].data()); - ID id = UniqueID::nil(); - if (notification->id()->size() > 0) { - id = from_flatbuf(*notification->id()); - } - DataT result; - auto root = flatbuffers::GetRoot(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 -Status Table::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 -Status Table::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::Subscribe( + job_id, client_id, + [subscribe](AsyncGcsClient *client, const ID &id, const std::vector &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(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 ¬ifications) { + RAY_CHECK(log_key == client_log_key_); + for (auto ¬ification : 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 &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(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 &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; template class Table; template class Table; template class Table; diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h index 54cf65627..90e650b90 100644 --- a/src/ray/gcs/tables.h +++ b/src/ray/gcs/tables.h @@ -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 -class Table { +class Log { public: using DataT = typename Data::NativeTableType; - using Callback = - std::function; - /// The callback to call when a lookup fails because there is no entry at the - /// key. - using FailureCallback = std::function; + using Callback = std::function &data)>; /// The callback to call when a SUBSCRIBE call completes and we are ready to /// request and receive notifications. using SubscriptionCallback = std::function; @@ -44,45 +49,44 @@ class Table { ID id; std::shared_ptr 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 *table; + Log *log; AsyncGcsClient *client; }; - Table(const std::shared_ptr &context, AsyncGcsClient *client) + Log(const std::shared_ptr &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 data, - const Callback &done); + Status Append(const JobID &job_id, const ID &id, std::shared_ptr 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 +class Table : private Log { + public: + using DataT = typename Log::DataT; + using Callback = + std::function; + /// The callback to call when a Lookup call returns an empty entry. + using FailureCallback = std::function; + /// The callback to call when a Subscribe call completes and we are ready to + /// request and receive notifications. + using SubscriptionCallback = typename Log::SubscriptionCallback; + + struct CallbackData { + ID id; + std::shared_ptr 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 *log; + AsyncGcsClient *client; + }; + + Table(const std::shared_ptr &context, AsyncGcsClient *client) + : Log(context, client) {} + + using Log::RequestNotifications; + using Log::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 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::context_; + using Log::client_; + using Log::pubsub_channel_; + using Log::prefix_; +}; + class ObjectTable : public Table { public: ObjectTable(const std::shared_ptr &context, AsyncGcsClient *client) @@ -210,10 +281,9 @@ class TaskTable : public Table { std::shared_ptr data, const TestAndUpdateCallback &callback) { int64_t callback_index = RedisCallbackManager::instance().add( - [this, callback, id](const std::vector &data) { - RAY_CHECK(data.size() == 1); + [this, callback, id](const std::string &data) { auto result = std::make_shared(); - auto root = flatbuffers::GetRoot(data[0].data()); + auto root = flatbuffers::GetRoot(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 { +/// \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 { public: using ClientTableCallback = std::function; ClientTable(const std::shared_ptr &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 { private: /// Handle a client table notification. - void HandleNotification(AsyncGcsClient *client, const ClientID &channel_id, - const ClientTableDataT ¬ifications); + void HandleNotification(AsyncGcsClient *client, const ClientTableDataT ¬ifications); /// Handle this client's successful connection to the GCS. - void HandleConnected(AsyncGcsClient *client, const ClientID &client_id, - const ClientTableDataT ¬ifications); + void HandleConnected(AsyncGcsClient *client, const ClientTableDataT ¬ifications); + /// 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.