src/commands/cmd_pubsub.cc - kvrocks - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  *
  */

 #include "commander.h"
 #include "error_constants.h"
 #include "server/server.h"
 #include "storage/redis_pubsub.h"

 namespace redis {

 class CommandPublish : public Commander {
  public:
   Status Execute([[maybe_unused]] engine::Context &ctx, Server *srv, [[maybe_unused]] Connection *conn,
                  std::string *output) override {
     if (!srv->IsSlave()) {
       // Compromise: can't replicate a message to sub-replicas in a cascading-like structure.
       // Replication relies on WAL seq; increasing the seq on a replica will break the replication process,
       // hence the compromise solution
       redis::PubSub pubsub_db(srv->storage);

       auto s = pubsub_db.Publish(ctx, args_[1], args_[2]);
       if (!s.ok()) {
         return {Status::RedisExecErr, s.ToString()};
       }
     }

     int receivers = srv->PublishMessage(args_[1], args_[2]);

     *output = redis::Integer(receivers);

     return Status::OK();
   }
 };

 class CommandMPublish : public Commander {
  public:
   Status Execute(engine::Context &ctx, Server *srv, [[maybe_unused]] Connection *conn, std::string *output) override {
     int total_receivers = 0;

     for (size_t i = 2; i < args_.size(); i++) {
       if (!srv->IsSlave()) {
         redis::PubSub pubsub_db(srv->storage);

         auto s = pubsub_db.Publish(ctx, args_[1], args_[i]);
         if (!s.ok()) {
           return {Status::RedisExecErr, s.ToString()};
         }
       }

       int receivers = srv->PublishMessage(args_[1], args_[i]);
       total_receivers += receivers;
     }

     *output = redis::Integer(total_receivers);

     return Status::OK();
   }
 };

 void SubscribeCommandReply(const Connection *conn, std::string *output, const std::string &name,
                            const std::string &sub_name, int num) {
   output->append(conn->HeaderOfPush(3));
   output->append(redis::BulkString(name));
   output->append(sub_name.empty() ? conn->NilString() : BulkString(sub_name));
   output->append(redis::Integer(num));
 }

 class CommandSubscribe : public Commander {
  public:
   Status Execute([[maybe_unused]] engine::Context &ctx, [[maybe_unused]] Server *srv, Connection *conn,
                  std::string *output) override {
     for (unsigned i = 1; i < args_.size(); i++) {
       conn->SubscribeChannel(args_[i]);
       SubscribeCommandReply(conn, output, "subscribe", args_[i],
                             conn->SubscriptionsCount() + conn->PSubscriptionsCount());
     }
     return Status::OK();
   }
 };

 class CommandUnSubscribe : public Commander {
  public:
   Status Execute([[maybe_unused]] engine::Context &ctx, [[maybe_unused]] Server *srv, Connection *conn,
                  std::string *output) override {
     if (args_.size() == 1) {
       conn->UnsubscribeAll([conn, output](const std::string &sub_name, int num) {
         SubscribeCommandReply(conn, output, "unsubscribe", sub_name, num);
       });
     } else {
       for (size_t i = 1; i < args_.size(); i++) {
         conn->UnsubscribeChannel(args_[i]);
         SubscribeCommandReply(conn, output, "unsubscribe", args_[i],
                               conn->SubscriptionsCount() + conn->PSubscriptionsCount());
       }
     }
     return Status::OK();
   }
 };

 class CommandPSubscribe : public Commander {
  public:
   Status Execute([[maybe_unused]] engine::Context &ctx, [[maybe_unused]] Server *srv, Connection *conn,
                  std::string *output) override {
     for (size_t i = 1; i < args_.size(); i++) {
       conn->PSubscribeChannel(args_[i]);
       SubscribeCommandReply(conn, output, "psubscribe", args_[i],
                             conn->SubscriptionsCount() + conn->PSubscriptionsCount());
     }
     return Status::OK();
   }
 };

 class CommandPUnSubscribe : public Commander {
  public:
   Status Execute([[maybe_unused]] engine::Context &ctx, [[maybe_unused]] Server *srv, Connection *conn,
                  std::string *output) override {
     if (args_.size() == 1) {
       conn->PUnsubscribeAll([conn, output](const std::string &sub_name, int num) {
         SubscribeCommandReply(conn, output, "punsubscribe", sub_name, num);
       });
     } else {
       for (size_t i = 1; i < args_.size(); i++) {
         conn->PUnsubscribeChannel(args_[i]);
         SubscribeCommandReply(conn, output, "punsubscribe", args_[i],
                               conn->SubscriptionsCount() + conn->PSubscriptionsCount());
       }
     }
     return Status::OK();
   }
 };

 class CommandSSubscribe : public Commander {
  public:
   Status Execute([[maybe_unused]] engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
     uint16_t slot = 0;
     if (srv->GetConfig()->cluster_enabled) {
       slot = GetSlotIdFromKey(args_[1]);
       for (unsigned int i = 2; i < args_.size(); i++) {
         if (GetSlotIdFromKey(args_[i]) != slot) {
           return {Status::RedisCrossSlot, "Keys in request don't hash to the same slot"};
         }
       }
     }

     for (unsigned int i = 1; i < args_.size(); i++) {
       conn->SSubscribeChannel(args_[i], slot);
       SubscribeCommandReply(conn, output, "ssubscribe", args_[i], conn->SSubscriptionsCount());
     }
     return Status::OK();
   }
 };

 class CommandSUnSubscribe : public Commander {
  public:
   Status Execute([[maybe_unused]] engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
     if (args_.size() == 1) {
       conn->SUnsubscribeAll([conn, output](const std::string &sub_name, int num) {
         SubscribeCommandReply(conn, output, "sunsubscribe", sub_name, num);
       });
     } else {
       for (size_t i = 1; i < args_.size(); i++) {
         conn->SUnsubscribeChannel(args_[i], srv->GetConfig()->cluster_enabled ? GetSlotIdFromKey(args_[i]) : 0);
         SubscribeCommandReply(conn, output, "sunsubscribe", args_[i], conn->SSubscriptionsCount());
       }
     }
     return Status::OK();
   }
 };

 class CommandPubSub : public Commander {
  public:
   Status Parse(const std::vector<std::string> &args) override {
     subcommand_ = util::ToLower(args[1]);
     if (subcommand_ == "numpat" && args.size() == 2) {
       return Status::OK();
     }

     if ((subcommand_ == "numsub" || subcommand_ == "shardnumsub") && args.size() >= 2) {
       if (args.size() > 2) {
         channels_ = std::vector<std::string>(args.begin() + 2, args.end());
       }
       return Status::OK();
     }

     if ((subcommand_ == "channels" || subcommand_ == "shardchannels") && args.size() <= 3) {
       if (args.size() == 3) {
         pattern_ = args[2];
       }
       return Status::OK();
     }

     return {Status::RedisInvalidCmd, errUnknownSubcommandOrWrongArguments};
   }

   Status Execute([[maybe_unused]] engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
     if (subcommand_ == "numpat") {
       *output = redis::Integer(srv->GetPubSubPatternSize());
       return Status::OK();
     }

     if (subcommand_ == "numsub" || subcommand_ == "shardnumsub") {
       std::vector<ChannelSubscribeNum> channel_subscribe_nums;
       if (subcommand_ == "numsub") {
         srv->ListChannelSubscribeNum(channels_, &channel_subscribe_nums);
       } else {
         srv->ListSChannelSubscribeNum(channels_, &channel_subscribe_nums);
       }

       output->append(redis::MultiLen(channel_subscribe_nums.size() * 2));
       for (const auto &chan_subscribe_num : channel_subscribe_nums) {
         output->append(redis::BulkString(chan_subscribe_num.channel));
         output->append(redis::Integer(chan_subscribe_num.subscribe_num));
       }

       return Status::OK();
     }

     if (subcommand_ == "channels" || subcommand_ == "shardchannels") {
       std::vector<std::string> channels;
       if (subcommand_ == "channels") {
         srv->GetChannelsByPattern(pattern_, &channels);
       } else {
         srv->GetSChannelsByPattern(pattern_, &channels);
       }
       *output = conn->MultiBulkString(channels);
       return Status::OK();
     }

     return {Status::RedisInvalidCmd, errUnknownSubcommandOrWrongArguments};
   }

  private:
   std::string pattern_;
   std::vector<std::string> channels_;
   std::string subcommand_;
 };

 REDIS_REGISTER_COMMANDS(Pubsub, MakeCmdAttr<CommandPublish>("publish", 3, "read-only", NO_KEY),
                         MakeCmdAttr<CommandMPublish>("mpublish", -3, "read-only", NO_KEY),
                         MakeCmdAttr<CommandSubscribe>("subscribe", -2, "read-only no-multi no-script", NO_KEY),
                         MakeCmdAttr<CommandUnSubscribe>("unsubscribe", -1, "read-only no-multi no-script", NO_KEY),
                         MakeCmdAttr<CommandPSubscribe>("psubscribe", -2, "read-only no-multi no-script", NO_KEY),
                         MakeCmdAttr<CommandPUnSubscribe>("punsubscribe", -1, "read-only no-multi no-script", NO_KEY),
                         MakeCmdAttr<CommandSSubscribe>("ssubscribe", -2, "read-only no-multi no-script", NO_KEY),
                         MakeCmdAttr<CommandSUnSubscribe>("sunsubscribe", -1, "read-only no-multi no-script", NO_KEY),
                         MakeCmdAttr<CommandPubSub>("pubsub", -2, "read-only no-script", NO_KEY), )

 }  // namespace redis
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*
	*/

	#include "commander.h"
	#include "error_constants.h"
	#include "server/server.h"
	#include "storage/redis_pubsub.h"

	namespace redis {

	class CommandPublish : public Commander {
	public:
	Status Execute([[maybe_unused]] engine::Context &ctx, Server srv, [[maybe_unused]] Connection conn,
	std::string *output) override {
	if (!srv->IsSlave()) {
	// Compromise: can't replicate a message to sub-replicas in a cascading-like structure.
	// Replication relies on WAL seq; increasing the seq on a replica will break the replication process,
	// hence the compromise solution
	redis::PubSub pubsub_db(srv->storage);

	auto s = pubsub_db.Publish(ctx, args_[1], args_[2]);
	if (!s.ok()) {
	return {Status::RedisExecErr, s.ToString()};
	}
	}

	int receivers = srv->PublishMessage(args_[1], args_[2]);

	*output = redis::Integer(receivers);

	return Status::OK();
	}
	};

	class CommandMPublish : public Commander {
	public:
	Status Execute(engine::Context &ctx, Server srv, [[maybe_unused]] Connection conn, std::string *output) override {
	int total_receivers = 0;

	for (size_t i = 2; i < args_.size(); i++) {
	if (!srv->IsSlave()) {
	redis::PubSub pubsub_db(srv->storage);

	auto s = pubsub_db.Publish(ctx, args_[1], args_[i]);
	if (!s.ok()) {
	return {Status::RedisExecErr, s.ToString()};
	}
	}

	int receivers = srv->PublishMessage(args_[1], args_[i]);
	total_receivers += receivers;
	}

	*output = redis::Integer(total_receivers);

	return Status::OK();
	}
	};

	void SubscribeCommandReply(const Connection conn, std::string output, const std::string &name,
	const std::string &sub_name, int num) {
	output->append(conn->HeaderOfPush(3));
	output->append(redis::BulkString(name));
	output->append(sub_name.empty() ? conn->NilString() : BulkString(sub_name));
	output->append(redis::Integer(num));
	}

	class CommandSubscribe : public Commander {
	public:
	Status Execute([[maybe_unused]] engine::Context &ctx, [[maybe_unused]] Server srv, Connection conn,
	std::string *output) override {
	for (unsigned i = 1; i < args_.size(); i++) {
	conn->SubscribeChannel(args_[i]);
	SubscribeCommandReply(conn, output, "subscribe", args_[i],
	conn->SubscriptionsCount() + conn->PSubscriptionsCount());
	}
	return Status::OK();
	}
	};

	class CommandUnSubscribe : public Commander {
	public:
	Status Execute([[maybe_unused]] engine::Context &ctx, [[maybe_unused]] Server srv, Connection conn,
	std::string *output) override {
	if (args_.size() == 1) {
	conn->UnsubscribeAll([conn, output](const std::string &sub_name, int num) {
	SubscribeCommandReply(conn, output, "unsubscribe", sub_name, num);
	});
	} else {
	for (size_t i = 1; i < args_.size(); i++) {
	conn->UnsubscribeChannel(args_[i]);
	SubscribeCommandReply(conn, output, "unsubscribe", args_[i],
	conn->SubscriptionsCount() + conn->PSubscriptionsCount());
	}
	}
	return Status::OK();
	}
	};

	class CommandPSubscribe : public Commander {
	public:
	Status Execute([[maybe_unused]] engine::Context &ctx, [[maybe_unused]] Server srv, Connection conn,
	std::string *output) override {
	for (size_t i = 1; i < args_.size(); i++) {
	conn->PSubscribeChannel(args_[i]);
	SubscribeCommandReply(conn, output, "psubscribe", args_[i],
	conn->SubscriptionsCount() + conn->PSubscriptionsCount());
	}
	return Status::OK();
	}
	};

	class CommandPUnSubscribe : public Commander {
	public:
	Status Execute([[maybe_unused]] engine::Context &ctx, [[maybe_unused]] Server srv, Connection conn,
	std::string *output) override {
	if (args_.size() == 1) {
	conn->PUnsubscribeAll([conn, output](const std::string &sub_name, int num) {
	SubscribeCommandReply(conn, output, "punsubscribe", sub_name, num);
	});
	} else {
	for (size_t i = 1; i < args_.size(); i++) {
	conn->PUnsubscribeChannel(args_[i]);
	SubscribeCommandReply(conn, output, "punsubscribe", args_[i],
	conn->SubscriptionsCount() + conn->PSubscriptionsCount());
	}
	}
	return Status::OK();
	}
	};

	class CommandSSubscribe : public Commander {
	public:
	Status Execute([[maybe_unused]] engine::Context &ctx, Server srv, Connection conn, std::string *output) override {
	uint16_t slot = 0;
	if (srv->GetConfig()->cluster_enabled) {
	slot = GetSlotIdFromKey(args_[1]);
	for (unsigned int i = 2; i < args_.size(); i++) {
	if (GetSlotIdFromKey(args_[i]) != slot) {
	return {Status::RedisCrossSlot, "Keys in request don't hash to the same slot"};
	}
	}
	}

	for (unsigned int i = 1; i < args_.size(); i++) {
	conn->SSubscribeChannel(args_[i], slot);
	SubscribeCommandReply(conn, output, "ssubscribe", args_[i], conn->SSubscriptionsCount());
	}
	return Status::OK();
	}
	};

	class CommandSUnSubscribe : public Commander {
	public:
	Status Execute([[maybe_unused]] engine::Context &ctx, Server srv, Connection conn, std::string *output) override {
	if (args_.size() == 1) {
	conn->SUnsubscribeAll([conn, output](const std::string &sub_name, int num) {
	SubscribeCommandReply(conn, output, "sunsubscribe", sub_name, num);
	});
	} else {
	for (size_t i = 1; i < args_.size(); i++) {
	conn->SUnsubscribeChannel(args_[i], srv->GetConfig()->cluster_enabled ? GetSlotIdFromKey(args_[i]) : 0);
	SubscribeCommandReply(conn, output, "sunsubscribe", args_[i], conn->SSubscriptionsCount());
	}
	}
	return Status::OK();
	}
	};

	class CommandPubSub : public Commander {
	public:
	Status Parse(const std::vector<std::string> &args) override {
	subcommand_ = util::ToLower(args[1]);
	if (subcommand_ == "numpat" && args.size() == 2) {
	return Status::OK();
	}

	if ((subcommand_ == "numsub" \|\| subcommand_ == "shardnumsub") && args.size() >= 2) {
	if (args.size() > 2) {
	channels_ = std::vector<std::string>(args.begin() + 2, args.end());
	}
	return Status::OK();
	}

	if ((subcommand_ == "channels" \|\| subcommand_ == "shardchannels") && args.size() <= 3) {
	if (args.size() == 3) {
	pattern_ = args[2];
	}
	return Status::OK();
	}

	return {Status::RedisInvalidCmd, errUnknownSubcommandOrWrongArguments};
	}

	Status Execute([[maybe_unused]] engine::Context &ctx, Server srv, Connection conn, std::string *output) override {
	if (subcommand_ == "numpat") {
	*output = redis::Integer(srv->GetPubSubPatternSize());
	return Status::OK();
	}

	if (subcommand_ == "numsub" \|\| subcommand_ == "shardnumsub") {
	std::vector<ChannelSubscribeNum> channel_subscribe_nums;
	if (subcommand_ == "numsub") {
	srv->ListChannelSubscribeNum(channels_, &channel_subscribe_nums);
	} else {
	srv->ListSChannelSubscribeNum(channels_, &channel_subscribe_nums);
	}

	output->append(redis::MultiLen(channel_subscribe_nums.size() * 2));
	for (const auto &chan_subscribe_num : channel_subscribe_nums) {
	output->append(redis::BulkString(chan_subscribe_num.channel));
	output->append(redis::Integer(chan_subscribe_num.subscribe_num));
	}

	return Status::OK();
	}

	if (subcommand_ == "channels" \|\| subcommand_ == "shardchannels") {
	std::vector<std::string> channels;
	if (subcommand_ == "channels") {
	srv->GetChannelsByPattern(pattern_, &channels);
	} else {
	srv->GetSChannelsByPattern(pattern_, &channels);
	}
	*output = conn->MultiBulkString(channels);
	return Status::OK();
	}

	return {Status::RedisInvalidCmd, errUnknownSubcommandOrWrongArguments};
	}

	private:
	std::string pattern_;
	std::vector<std::string> channels_;
	std::string subcommand_;
	};

	REDIS_REGISTER_COMMANDS(Pubsub, MakeCmdAttr<CommandPublish>("publish", 3, "read-only", NO_KEY),
	MakeCmdAttr<CommandMPublish>("mpublish", -3, "read-only", NO_KEY),
	MakeCmdAttr<CommandSubscribe>("subscribe", -2, "read-only no-multi no-script", NO_KEY),
	MakeCmdAttr<CommandUnSubscribe>("unsubscribe", -1, "read-only no-multi no-script", NO_KEY),
	MakeCmdAttr<CommandPSubscribe>("psubscribe", -2, "read-only no-multi no-script", NO_KEY),
	MakeCmdAttr<CommandPUnSubscribe>("punsubscribe", -1, "read-only no-multi no-script", NO_KEY),
	MakeCmdAttr<CommandSSubscribe>("ssubscribe", -2, "read-only no-multi no-script", NO_KEY),
	MakeCmdAttr<CommandSUnSubscribe>("sunsubscribe", -1, "read-only no-multi no-script", NO_KEY),
	MakeCmdAttr<CommandPubSub>("pubsub", -2, "read-only no-script", NO_KEY), )

	} // namespace redis