fix(server): Use correct messages for sharded pubsub (#5818)

Author: dranikpg

src/facade/command_id.h

@@ -88,21 +88,6 @@ class CommandId {
     opt_mask_ |= flag;
   }
 
-  // PUBLISH/SUBSCRIBE/UNSUBSCRIBE variant
-  bool IsPubSub() const {
-    return is_pub_sub_;
-  }
-
-  // PSUBSCRIBE/PUNSUBSCRIBE variant
-  bool IsPSub() const {
-    return is_p_pub_sub_;
-  }
-
-  // SSUBSCRIBE/SUNSUBSCRIBE variant
-  bool IsShardedPSub() const {
-    return is_sharded_pub_sub_;
-  }
-
  protected:
   std::string name_;
 
@@ -119,10 +104,6 @@ class CommandId {
 
   // Whether the command can only be used by admin connections.
   bool restricted_ = false;
-
-  bool is_pub_sub_ = false;
-  bool is_sharded_pub_sub_ = false;
-  bool is_p_pub_sub_ = false;
 };
 
 }  // namespace facade

src/facade/dragonfly_connection.cc

@@ -519,7 +519,7 @@ void Connection::AsyncOperations::operator()(const AclUpdateMessage& msg) {
 }
 
 void Connection::AsyncOperations::operator()(const PubMessage& pub_msg) {
-  RedisReplyBuilder* rbuilder = (RedisReplyBuilder*)builder;
+  RedisReplyBuilder* rb = static_cast<RedisReplyBuilder*>(builder);
 
   // Discard stale messages to not break the protocol after exiting "pubsub" mode.
   // Even after removing all subscriptions, we still can receive messages delayed
@@ -529,20 +529,19 @@ void Connection::AsyncOperations::operator()(const PubMessage& pub_msg) {
       !base::_in(pub_msg.channel, {"unsubscribe", "punsubscribe"}))
     return;
 
-  if (pub_msg.should_unsubscribe) {
-    rbuilder->StartCollection(3, RedisReplyBuilder::CollectionType::PUSH);
-    rbuilder->SendBulkString("unsubscribe");
-    rbuilder->SendBulkString(pub_msg.channel);
-    rbuilder->SendLong(0);
-    auto* cntx = self->cntx();
-    cntx->Unsubscribe(pub_msg.channel);
+  if (pub_msg.force_unsubscribe) {
+    rb->StartCollection(3, RedisReplyBuilder::CollectionType::PUSH);
+    rb->SendBulkString("sunsubscribe");
+    rb->SendBulkString(pub_msg.channel);
+    rb->SendLong(0);
+    self->cntx()->Unsubscribe(pub_msg.channel);
     return;
   }
 
   unsigned i = 0;
   array<string_view, 4> arr;
   if (pub_msg.pattern.empty()) {
-    arr[i++] = "message";
+    arr[i++] = pub_msg.is_sharded ? "smessage" : "message";
   } else {
     arr[i++] = "pmessage";
     arr[i++] = pub_msg.pattern;
@@ -551,8 +550,8 @@ void Connection::AsyncOperations::operator()(const PubMessage& pub_msg) {
   arr[i++] = pub_msg.channel;
   arr[i++] = pub_msg.message;
 
-  rbuilder->SendBulkStrArr(absl::Span<string_view>{arr.data(), i},
-                           RedisReplyBuilder::CollectionType::PUSH);
+  rb->SendBulkStrArr(absl::Span<string_view>{arr.data(), i},
+                     RedisReplyBuilder::CollectionType::PUSH);
 }
 
 void Connection::AsyncOperations::operator()(Connection::PipelineMessage& msg) {

src/facade/dragonfly_connection.h

@@ -73,7 +73,10 @@ class Connection : public util::Connection {
     std::string pattern;                // non-empty for pattern subscriber
     std::shared_ptr<char[]> buf;        // stores channel name and message
     std::string_view channel, message;  // channel and message parts from buf
-    bool should_unsubscribe = false;    // unsubscribe from channel after sending the message
+    bool is_sharded = false;
+
+    // Unsubscribe simultaneously when sending unsubscribe message. Used for cluster migrations
+    bool force_unsubscribe = false;
   };
 
   // Pipeline message, accumulated Redis command to be executed.

src/facade/facade.cc

@@ -138,13 +138,6 @@ CommandId::CommandId(const char* name, uint32_t mask, int8_t arity, int8_t first
       first_key_(first_key),
       last_key_(last_key),
       acl_categories_(acl_categories) {
-  if (name_ == "PUBLISH" || name_ == "SUBSCRIBE" || name_ == "UNSUBSCRIBE") {
-    is_pub_sub_ = true;
-  } else if (name_ == "PSUBSCRIBE" || name_ == "PUNSUBSCRIBE") {
-    is_p_pub_sub_ = true;
-  } else if (name_ == "SPUBLISH" || name_ == "SSUBSCRIBE" || name_ == "SUNSUBSCRIBE") {
-    is_sharded_pub_sub_ = true;
-  }
 }
 
 static bool ParseHumanReadableBytes(std::string_view str, int64_t* num_bytes) {

src/server/acl/validator.cc

@@ -78,10 +78,10 @@ bool ValidateCommand(const std::vector<uint64_t>& acl_commands, const CommandId&
 
   std::pair<bool, AclLog::Reason> auth_res;
 
-  if (id.IsPubSub() || id.IsShardedPSub()) {
-    auth_res = IsPubSubCommandAuthorized(false, cntx.acl_commands, cntx.pub_sub, tail_args, id);
-  } else if (id.IsPSub()) {
-    auth_res = IsPubSubCommandAuthorized(true, cntx.acl_commands, cntx.pub_sub, tail_args, id);
+  if (auto pkind = id.PubSubKind(); pkind) {
+    bool is_pattern = *pkind == CO::PubSubKind::PATTERN;
+    auth_res =
+        IsPubSubCommandAuthorized(is_pattern, cntx.acl_commands, cntx.pub_sub, tail_args, id);
   } else {
     auth_res = IsUserAllowedToInvokeCommandGeneric(cntx, id, tail_args);
   }

src/server/channel_store.cc

@@ -19,7 +19,8 @@ using namespace std;
 namespace {
 
 // Build functor for sending messages to connection
-auto BuildSender(string_view channel, facade::ArgRange messages, bool unsubscribe = false) {
+auto BuildSender(string_view channel, facade::ArgRange messages, bool sharded = false,
+                 bool unsubscribe = false) {
   absl::FixedArray<string_view, 1> views(messages.Size());
   size_t messages_size = accumulate(messages.begin(), messages.end(), 0,
                                     [](int sum, string_view str) { return sum + str.size(); });
@@ -36,11 +37,12 @@ auto BuildSender(string_view channel, facade::ArgRange messages, bool unsubscrib
     }
   }
 
-  return [channel, buf = std::move(buf), views = std::move(views), unsubscribe](
+  return [channel, buf = std::move(buf), views = std::move(views), sharded, unsubscribe](
              facade::Connection* conn, string pattern) {
     string_view channel_view{buf.get(), channel.size()};
     for (std::string_view message_view : views) {
-      conn->SendPubMessageAsync({std::move(pattern), buf, channel_view, message_view, unsubscribe});
+      conn->SendPubMessageAsync(
+          {std::move(pattern), buf, channel_view, message_view, sharded, unsubscribe});
     }
   };
 }
@@ -117,7 +119,8 @@ void ChannelStore::Destroy() {
 
 ChannelStore::ControlBlock ChannelStore::control_block;
 
-unsigned ChannelStore::SendMessages(std::string_view channel, facade::ArgRange messages) const {
+unsigned ChannelStore::SendMessages(std::string_view channel, facade::ArgRange messages,
+                                    bool sharded) const {
   vector<Subscriber> subscribers = FetchSubscribers(channel);
   if (subscribers.empty())
     return 0;
@@ -144,7 +147,7 @@ unsigned ChannelStore::SendMessages(std::string_view channel, facade::ArgRange m
   }
 
   auto subscribers_ptr = make_shared<decltype(subscribers)>(std::move(subscribers));
-  auto cb = [subscribers_ptr, send = BuildSender(channel, messages)](unsigned idx, auto*) {
+  auto cb = [subscribers_ptr, send = BuildSender(channel, messages, sharded)](unsigned idx, auto*) {
     auto it = lower_bound(subscribers_ptr->begin(), subscribers_ptr->end(), idx,
                           ChannelStore::Subscriber::ByThreadId);
     while (it != subscribers_ptr->end() && it->LastKnownThreadId() == idx) {
@@ -217,13 +220,14 @@ void ChannelStore::UnsubscribeAfterClusterSlotMigration(const cluster::SlotSet&
   csu.ApplyAndUnsubscribe();
 }
 
+// TODO: Reuse common code with Send function
+// TODO: Find proper solution to hacky `force_unsubscribe` flag or at least move logic out of io
 void ChannelStore::UnsubscribeConnectionsFromDeletedSlots(const ChannelsSubMap& sub_map,
                                                           uint32_t idx) {
-  const bool should_unsubscribe = true;
   for (const auto& [channel, subscribers] : sub_map) {
     // ignored by pub sub handler because should_unsubscribe is true
     std::string msg = "__ignore__";
-    auto send = BuildSender(channel, {facade::ArgSlice{msg}}, should_unsubscribe);
+    auto send = BuildSender(channel, {facade::ArgSlice{msg}}, false, true);
 
     auto it = lower_bound(subscribers.begin(), subscribers.end(), idx,
                           ChannelStore::Subscriber::ByThreadId);

src/server/channel_store.h

@@ -59,7 +59,7 @@ class ChannelStore {
   ChannelStore();
 
   // Send messages to channel, block on connection backpressure
-  unsigned SendMessages(std::string_view channel, facade::ArgRange messages) const;
+  unsigned SendMessages(std::string_view channel, facade::ArgRange messages, bool sharded) const;
 
   // Fetch all subscribers for channel, including matching patterns.
   std::vector<Subscriber> FetchSubscribers(std::string_view channel) const;

src/server/cluster/cluster_family_test.cc

@@ -670,6 +670,34 @@ TEST_F(ClusterFamilyTest, ClusterModePubSubNotAllowed) {
               ErrArg("PUNSUBSCRIBE is not supported in cluster mode yet"));
 }
 
+// SSUBSCRIBE and SPUBLISH work in cluster mode
+TEST_F(ClusterFamilyTest, ClusterModePubSub) {
+  single_response_ = false;
+  ConfigSingleNodeCluster(GetMyId());
+
+  // Ssubscribe works as expected
+  auto resp = pp_->at(1)->Await([&] { return Run({"SSUBSCRIBE", "cluster-channel"}); });
+  EXPECT_THAT(resp, RespElementsAre("ssubscribe", "cluster-channel", IntArg(1)));
+
+  // Send-receive a single message
+  resp = pp_->at(0)->Await([&] {
+    return Run({"SPUBLISH", "cluster-channel", "a simple message"});
+  });
+  EXPECT_THAT(resp, IntArg(1));
+
+  pp_->AwaitFiberOnAll([](util::ProactorBase* pb) {});
+
+  ASSERT_EQ(1, SubscriberMessagesLen("IO1"));
+  const auto& msg = GetPublishedMessage("IO1", 0);
+  EXPECT_TRUE(msg.is_sharded);
+  EXPECT_EQ("cluster-channel", msg.channel);
+  EXPECT_EQ("a simple message", msg.message);
+
+  // Sunsubscribe
+  resp = pp_->at(1)->Await([&] { return Run({"SUNSUBSCRIBE", "cluster-channel"}); });
+  EXPECT_THAT(resp, RespElementsAre("sunsubscribe", "cluster-channel", IntArg(0)));
+}
+
 TEST_F(ClusterFamilyTest, ClusterFirstConfigCallDropsEntriesNotOwnedByNode) {
   InitWithDbFilename();
 

src/server/command_registry.cc

@@ -4,15 +4,16 @@
 
 #include "server/command_registry.h"
 
+#include <absl/container/inlined_vector.h>
+#include <absl/strings/match.h>
+#include <absl/strings/str_cat.h>
 #include <absl/strings/str_split.h>
 #include <absl/time/clock.h>
 
-#include "absl/container/inlined_vector.h"
-#include "absl/strings/match.h"
-#include "absl/strings/str_cat.h"
 #include "base/bits.h"
 #include "base/flags.h"
 #include "base/logging.h"
+#include "base/stl_util.h"
 #include "facade/dragonfly_connection.h"
 #include "facade/error.h"
 #include "server/acl/acl_commands_def.h"
@@ -137,6 +138,17 @@ CommandId::CommandId(const char* name, uint32_t mask, int8_t arity, int8_t first
                                    kLatencyHistogramPrecision, &hist);
   CHECK_EQ(init_result, 0) << "failed to initialize histogram for command " << name;
   latency_histogram_ = hist;
+
+  if (name_.rfind("EVAL", 0) == 0)
+    kind_multi_ctr_ = CO::MultiControlKind::EVAL;
+  else if (base::_in(name_, {"EXEC", "MULTI", "DISCARD"}))
+    kind_multi_ctr_ = CO::MultiControlKind::EXEC;
+  else if (base::_in(name_, {"PUBLISH", "SUBSCRIBE", "UNSUBSCRIBE"}))
+    kind_pubsub_ = CO::PubSubKind::REGULAR;
+  else if (base::_in(name_, {"PSUBSCRIBE", "PUNSUBSCRIBE"}))
+    kind_pubsub_ = CO::PubSubKind::PATTERN;
+  else if (base::_in(name_, {"SPUBLISH", "SSUBSCRIBE", "SUNSUBSCRIBE"}))
+    kind_pubsub_ = CO::PubSubKind::SHARDED;
 }
 
 CommandId::~CommandId() {
@@ -174,7 +186,7 @@ bool CommandId::IsTransactional() const {
 }
 
 bool CommandId::IsMultiTransactional() const {
-  return CO::IsTransKind(name()) || CO::IsEvalKind(name());
+  return kind_multi_ctr_.has_value();
 }
 
 uint64_t CommandId::Invoke(CmdArgList args, const CommandContext& cmd_cntx) const {

src/server/command_registry.h

@@ -55,17 +55,14 @@ enum CommandOpt : uint32_t {
   IDEMPOTENT = 1U << 18,
 };
 
-constexpr inline bool IsEvalKind(std::string_view name) {
-  return name.compare(0, 4, "EVAL") == 0;
-}
+enum class PubSubKind : uint8_t { REGULAR = 0, PATTERN = 1, SHARDED = 2 };
 
-constexpr inline bool IsTransKind(std::string_view name) {
-  return (name == "EXEC") || (name == "MULTI") || (name == "DISCARD");
-}
-
-static_assert(IsEvalKind("EVAL") && IsEvalKind("EVAL_RO") && IsEvalKind("EVALSHA") &&
-              IsEvalKind("EVALSHA_RO"));
-static_assert(!IsEvalKind(""));
+// Commands controlling any multi command execution.
+// They often need to be handled separately from regular commands in many contexts
+enum class MultiControlKind : uint8_t {
+  EVAL,  // EVAL, EVAL_RO, EVALSHA, EVALSHA_RO
+  EXEC,  // EXEC, MULTI, DISCARD
+};
 
 };  // namespace CO
 
@@ -114,6 +111,7 @@ class CommandId : public facade::CommandId {
   // server_state.h)
   CommandId(const char* name, uint32_t mask, int8_t arity, int8_t first_key, int8_t last_key,
             std::optional<uint32_t> acl_categories = std::nullopt);
+
   CommandId(CommandId&& o) = default;
 
   ~CommandId();
@@ -182,7 +180,19 @@ class CommandId : public facade::CommandId {
 
   hdr_histogram* LatencyHist() const;
 
+  std::optional<CO::PubSubKind> PubSubKind() const {
+    return kind_pubsub_;
+  }
+
+  // Returns value if this command controls multi command execution (EVAL, EXEC & helpers)
+  std::optional<CO::MultiControlKind> MultiControlKind() const {
+    return kind_multi_ctr_;
+  }
+
  private:
+  std::optional<CO::PubSubKind> kind_pubsub_;
+  std::optional<CO::MultiControlKind> kind_multi_ctr_;
+
   // The following fields must copy manually in the move constructor.
   bool implicit_acl_;
   bool is_alias_{false};