From b0cf7811520bd33194a81b97e485a93f3e0311f5 Mon Sep 17 00:00:00 2001
From: fangfengbin <869218239a@zju.edu.cn>
Date: Wed, 27 May 2020 11:55:17 +0800
Subject: [PATCH] fix resubscribe miss callback index bug (#8604)

---
 src/ray/gcs/pubsub/gcs_pub_sub.cc | 17 +++++++++++++----
 src/ray/gcs/pubsub/gcs_pub_sub.h  |  1 +
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/ray/gcs/pubsub/gcs_pub_sub.cc b/src/ray/gcs/pubsub/gcs_pub_sub.cc
index 4004d9f8b..52277112f 100644
--- a/src/ray/gcs/pubsub/gcs_pub_sub.cc
+++ b/src/ray/gcs/pubsub/gcs_pub_sub.cc
@@ -45,8 +45,15 @@ Status GcsPubSub::SubscribeAll(const std::string &channel, const Callback &subsc
 }
 
 Status GcsPubSub::Unsubscribe(const std::string &channel, const std::string &id) {
-  return redis_client_->GetPrimaryContext()->PUnsubscribeAsync(
-      GenChannelPattern(channel, boost::optional<std::string>(id)));
+  std::string pattern = GenChannelPattern(channel, id);
+  {
+    absl::MutexLock lock(&mutex_);
+    auto it = subscribe_callback_index_.find(pattern);
+    RAY_CHECK(it != subscribe_callback_index_.end());
+    unsubscribe_callback_index_[pattern] = it->second;
+    subscribe_callback_index_.erase(it);
+  }
+  return redis_client_->GetPrimaryContext()->PUnsubscribeAsync(pattern);
 }
 
 Status GcsPubSub::SubscribeInternal(const std::string &channel, const Callback &subscribe,
@@ -58,8 +65,8 @@ Status GcsPubSub::SubscribeInternal(const std::string &channel, const Callback &
       if (reply->IsUnsubscribeCallback()) {
         absl::MutexLock lock(&mutex_);
         ray::gcs::RedisCallbackManager::instance().remove(
-            subscribe_callback_index_[pattern]);
-        subscribe_callback_index_.erase(pattern);
+            unsubscribe_callback_index_[pattern]);
+        unsubscribe_callback_index_.erase(pattern);
       } else if (reply->IsSubscribeCallback()) {
         if (done) {
           done(Status::OK());
@@ -80,6 +87,8 @@ Status GcsPubSub::SubscribeInternal(const std::string &channel, const Callback &
                                                                     &out_callback_index);
   if (id) {
     absl::MutexLock lock(&mutex_);
+    // If the same pattern has been subscribed more than once, the last subscription takes
+    // effect.
     subscribe_callback_index_[pattern] = out_callback_index;
   }
   return status;
diff --git a/src/ray/gcs/pubsub/gcs_pub_sub.h b/src/ray/gcs/pubsub/gcs_pub_sub.h
index 1a2704c4b..b3417c812 100644
--- a/src/ray/gcs/pubsub/gcs_pub_sub.h
+++ b/src/ray/gcs/pubsub/gcs_pub_sub.h
@@ -101,6 +101,7 @@ class GcsPubSub {
   absl::Mutex mutex_;
 
   std::unordered_map<std::string, int64_t> subscribe_callback_index_ GUARDED_BY(mutex_);
+  std::unordered_map<std::string, int64_t> unsubscribe_callback_index_ GUARDED_BY(mutex_);
 };
 
 }  // namespace gcs