[GCS]Optimize subscription perf (#11669)

* [GCS]Optimize subscription perf

* fix review comment

Co-authored-by: 灵洵 <fengbin.ffb@antfin.com>

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<std::string>(id));
+  return SubscribeInternal(channel, subscribe, done, false,
+                           boost::optional<std::string>(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<std::string> &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) {

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<std::string> &id = boost::none);
 
   std::string GenChannelPattern(const std::string &channel,

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<redisCallbackFn *>(&GlobalRedisCallback),
+      reinterpret_cast<void *>(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<redisCallbackFn *>(&GlobalRedisCallback),
+      reinterpret_cast<void *>(-1), redis_command.c_str(), channel.c_str(),
+      channel.size());
+}
+
 Status RedisContext::PSubscribeAsync(const std::string &pattern,
                                      const RedisCallback &redisCallback,
                                      int64_t callback_index) {

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.