diff --git a/src/ray/gcs/pubsub/gcs_pub_sub.cc b/src/ray/gcs/pubsub/gcs_pub_sub.cc index 6ea273c0d..ee6ac955b 100644 --- a/src/ray/gcs/pubsub/gcs_pub_sub.cc +++ b/src/ray/gcs/pubsub/gcs_pub_sub.cc @@ -36,12 +36,13 @@ Status GcsPubSub::Publish(const std::string &channel, const std::string &id, Status GcsPubSub::Subscribe(const std::string &channel, const std::string &id, const Callback &subscribe, const StatusCallback &done) { - return SubscribeInternal(channel, subscribe, done, boost::optional(id)); + return SubscribeInternal(channel, subscribe, done, false, + boost::optional(id)); } Status GcsPubSub::SubscribeAll(const std::string &channel, const Callback &subscribe, const StatusCallback &done) { - return SubscribeInternal(channel, subscribe, done); + return SubscribeInternal(channel, subscribe, done, true); } Status GcsPubSub::Unsubscribe(const std::string &channel_name, const std::string &id) { @@ -59,6 +60,7 @@ Status GcsPubSub::Unsubscribe(const std::string &channel_name, const std::string Status GcsPubSub::SubscribeInternal(const std::string &channel_name, const Callback &subscribe, const StatusCallback &done, + bool is_sub_or_unsub_all, const boost::optional &id) { std::string pattern = GenChannelPattern(channel_name, id); @@ -71,7 +73,7 @@ Status GcsPubSub::SubscribeInternal(const std::string &channel_name, } // Add the SUBSCRIBE command to the queue. - channel->second.command_queue.push_back(Command(subscribe, done)); + channel->second.command_queue.push_back(Command(subscribe, done, is_sub_or_unsub_all)); // Process the first command on the queue, if possible. return ExecuteCommandIfPossible(channel->first, channel->second); @@ -140,15 +142,25 @@ Status GcsPubSub::ExecuteCommandIfPossible(const std::string &channel_key, } } }; - status = redis_client_->GetPrimaryContext()->PSubscribeAsync(channel_key, callback, - callback_index); + + if (command.is_sub_or_unsub_all) { + status = redis_client_->GetPrimaryContext()->PSubscribeAsync(channel_key, callback, + callback_index); + } else { + status = redis_client_->GetPrimaryContext()->SubscribeAsync(channel_key, callback, + callback_index); + } channel.pending_reply = true; channel.command_queue.pop_front(); } else if (!command.is_subscribe && channel.callback_index != -1) { // The next command is UNSUBSCRIBE and we are currently subscribed, so we // can execute the command. The reply for will be received through the // SUBSCRIBE command's callback. - status = redis_client_->GetPrimaryContext()->PUnsubscribeAsync(channel_key); + if (command.is_sub_or_unsub_all) { + status = redis_client_->GetPrimaryContext()->PUnsubscribeAsync(channel_key); + } else { + status = redis_client_->GetPrimaryContext()->UnsubscribeAsync(channel_key); + } channel.pending_reply = true; channel.command_queue.pop_front(); } else if (!channel.pending_reply) { diff --git a/src/ray/gcs/pubsub/gcs_pub_sub.h b/src/ray/gcs/pubsub/gcs_pub_sub.h index 0c4bef83b..19a0923ac 100644 --- a/src/ray/gcs/pubsub/gcs_pub_sub.h +++ b/src/ray/gcs/pubsub/gcs_pub_sub.h @@ -99,12 +99,14 @@ class GcsPubSub { /// channel. struct Command { /// SUBSCRIBE constructor. - Command(const Callback &subscribe_callback, const StatusCallback &done_callback) + Command(const Callback &subscribe_callback, const StatusCallback &done_callback, + bool is_sub_or_unsub_all) : is_subscribe(true), subscribe_callback(subscribe_callback), - done_callback(done_callback) {} + done_callback(done_callback), + is_sub_or_unsub_all(is_sub_or_unsub_all) {} /// UNSUBSCRIBE constructor. - Command() : is_subscribe(false) {} + Command() : is_subscribe(false), is_sub_or_unsub_all(false) {} /// True if this is a SUBSCRIBE command and false if UNSUBSCRIBE. const bool is_subscribe; /// Callback that is called whenever a new pubsub message is received from @@ -113,6 +115,8 @@ class GcsPubSub { /// Callback that is called once we have successfully subscribed to a /// channel. This should only be set if is_subscribe is true. const StatusCallback done_callback; + /// True if this is a SUBSCRIBE all or UNSUBSCRIBE all command else false. + const bool is_sub_or_unsub_all; }; struct Channel { @@ -147,7 +151,7 @@ class GcsPubSub { EXCLUSIVE_LOCKS_REQUIRED(mutex_); Status SubscribeInternal(const std::string &channel_name, const Callback &subscribe, - const StatusCallback &done, + const StatusCallback &done, bool is_sub_or_unsub_all, const boost::optional &id = boost::none); std::string GenChannelPattern(const std::string &channel, diff --git a/src/ray/gcs/redis_context.cc b/src/ray/gcs/redis_context.cc index dc61b12f0..840752721 100644 --- a/src/ray/gcs/redis_context.cc +++ b/src/ray/gcs/redis_context.cc @@ -93,7 +93,8 @@ CallbackReply::CallbackReply(redisReply *redis_reply) : reply_type_(redis_reply- is_subscribe_callback_ = true; // If the message is for the initial subscription call, return the empty // string as a response to signify that subscription was successful. - } else if (strcmp(message_type->str, "punsubscribe") == 0) { + } else if (strcmp(message_type->str, "punsubscribe") == 0 || + strcmp(message_type->str, "unsubscribe") == 0) { is_unsubscribe_callback_ = true; } else if (strcmp(message_type->str, "message") == 0) { // If the message is from a PUBLISH, make sure the data is nonempty. @@ -412,6 +413,30 @@ Status RedisContext::SubscribeAsync(const NodeID &client_id, return status; } +Status RedisContext::SubscribeAsync(const std::string &channel, + const RedisCallback &redisCallback, + int64_t callback_index) { + RAY_CHECK(async_redis_subscribe_context_); + + RAY_UNUSED(RedisCallbackManager::instance().AddCallback(redisCallback, true, + io_service_, callback_index)); + std::string redis_command = "SUBSCRIBE %b"; + return async_redis_subscribe_context_->RedisAsyncCommand( + reinterpret_cast(&GlobalRedisCallback), + reinterpret_cast(callback_index), redis_command.c_str(), channel.c_str(), + channel.size()); +} + +Status RedisContext::UnsubscribeAsync(const std::string &channel) { + RAY_CHECK(async_redis_subscribe_context_); + + std::string redis_command = "UNSUBSCRIBE %b"; + return async_redis_subscribe_context_->RedisAsyncCommand( + reinterpret_cast(&GlobalRedisCallback), + reinterpret_cast(-1), redis_command.c_str(), channel.c_str(), + channel.size()); +} + Status RedisContext::PSubscribeAsync(const std::string &pattern, const RedisCallback &redisCallback, int64_t callback_index) { diff --git a/src/ray/gcs/redis_context.h b/src/ray/gcs/redis_context.h index 36ce4bb72..c8bcb04a1 100644 --- a/src/ray/gcs/redis_context.h +++ b/src/ray/gcs/redis_context.h @@ -239,6 +239,23 @@ class RedisContext { Status SubscribeAsync(const NodeID &client_id, const TablePubsub pubsub_channel, const RedisCallback &redisCallback, int64_t *out_callback_index); + /// Subscribes the client to the given channel. + /// + /// \param channel The subscription channel. + /// \param redisCallback The callback function that the notification calls. + /// \param callback_index The index at which to add the callback. This index + /// must already be allocated in the callback manager via + /// RedisCallbackManager::AllocateCallbackIndex. + /// \return Status. + Status SubscribeAsync(const std::string &channel, const RedisCallback &redisCallback, + int64_t callback_index); + + /// Unsubscribes the client from the given channel. + /// + /// \param channel The unsubscription channel. + /// \return Status. + Status UnsubscribeAsync(const std::string &channel); + /// Subscribes the client to the given pattern. /// /// \param pattern The pattern of subscription channel.