Add dynamic flag generator to replace adhoc logic in ExecuteCommands (#…

…1655)

We add dynamic flag generator regarding to command arguments to replace adhoc exclusive check logic in Connection::ExecuteCommands().

And we also fix and refactor GetKeysFromCommand() in this PR.

src/cluster/cluster.cc

@@ -742,7 +742,7 @@ bool Cluster::IsWriteForbiddenSlot(int slot) { return svr_->slot_migrator->GetFo
 Status Cluster::CanExecByMySelf(const redis::CommandAttributes *attributes, const std::vector<std::string> &cmd_tokens,
                                 redis::Connection *conn) {
   std::vector<int> keys_indexes;
-  auto s = redis::GetKeysFromCommand(attributes->name, static_cast<int>(cmd_tokens.size()), &keys_indexes);
+  auto s = redis::GetKeysFromCommand(attributes, cmd_tokens, &keys_indexes);
   // No keys
   if (!s.IsOK()) return Status::OK();
 
@@ -773,7 +773,7 @@ Status Cluster::CanExecByMySelf(const redis::CommandAttributes *attributes, cons
     }
     // To keep data consistency, slot will be forbidden write while sending the last incremental data.
     // During this phase, the requests of the migrating slot has to be rejected.
-    if (attributes->IsWrite() && IsWriteForbiddenSlot(slot)) {
+    if ((attributes->flags & redis::kCmdWrite) && IsWriteForbiddenSlot(slot)) {
       return {Status::RedisExecErr, "TRYAGAIN Can't write to slot being migrated which is in write forbidden phase"};
     }
 
@@ -795,7 +795,7 @@ Status Cluster::CanExecByMySelf(const redis::CommandAttributes *attributes, cons
     return Status::OK();  // I'm serving the imported slot
   }
 
-  if (myself_ && myself_->role == kClusterSlave && !attributes->IsWrite() &&
+  if (myself_ && myself_->role == kClusterSlave && !(attributes->flags & redis::kCmdWrite) &&
       nodes_.find(myself_->master_id) != nodes_.end() && nodes_[myself_->master_id] == slots_nodes_[slot]) {
     return Status::OK();  // My master is serving this slot
   }

src/commands/cmd_cluster.cc

@@ -284,7 +284,16 @@ class CommandClusterX : public Commander {
   std::unique_ptr<SyncMigrateContext> sync_migrate_ctx_ = nullptr;
 };
 
-REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandCluster>("cluster", -2, "cluster no-script", 0, 0, 0),
-                        MakeCmdAttr<CommandClusterX>("clusterx", -2, "cluster no-script", 0, 0, 0), )
+static uint64_t GenerateClusterFlag(const std::vector<std::string> &args) {
+  if (args.size() >= 2 && Cluster::SubCommandIsExecExclusive(args[1])) {
+    return kCmdExclusive;
+  }
+
+  return 0;
+}
+
+REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandCluster>("cluster", -2, "cluster no-script", 0, 0, 0, GenerateClusterFlag),
+                        MakeCmdAttr<CommandClusterX>("clusterx", -2, "cluster no-script", 0, 0, 0,
+                                                     GenerateClusterFlag), )
 
 }  // namespace redis

src/commands/cmd_geo.cc

@@ -275,7 +275,7 @@ class CommandGeoRadius : public CommandGeoBase {
 
         count_ = *parse_result;
         i += 2;
-      } else if (attributes_->IsWrite() &&
+      } else if ((attributes_->flags & kCmdWrite) &&
                  (util::ToLower(args_[i]) == "store" || util::ToLower(args_[i]) == "storedist") &&
                  i + 1 < args_.size()) {
         store_key_ = args_[i + 1];

src/commands/cmd_server.cc

@@ -26,6 +26,7 @@
 #include "server/redis_connection.h"
 #include "server/server.h"
 #include "stats/disk_stats.h"
+#include "string_util.h"
 #include "time_util.h"
 
 namespace redis {
@@ -612,8 +613,14 @@ class CommandCommand : public Commander {
       } else if (sub_command == "info") {
         GetCommandsInfo(output, std::vector<std::string>(args_.begin() + 2, args_.end()));
       } else if (sub_command == "getkeys") {
+        auto cmd_iter = command_details::original_commands.find(util::ToLower(args_[2]));
+        if (cmd_iter == command_details::original_commands.end()) {
+          return {Status::RedisUnknownCmd, "Invalid command specified"};
+        }
+
         std::vector<int> keys_indexes;
-        auto s = GetKeysFromCommand(args_[2], static_cast<int>(args_.size()) - 2, &keys_indexes);
+        auto s = GetKeysFromCommand(cmd_iter->second, std::vector<std::string>(args_.begin() + 2, args_.end()),
+                                    &keys_indexes);
         if (!s.IsOK()) return s;
 
         if (keys_indexes.size() == 0) {
@@ -964,12 +971,20 @@ class CommandStats : public Commander {
   }
 };
 
+static uint64_t GenerateConfigFlag(const std::vector<std::string> &args) {
+  if (args.size() >= 2 && util::EqualICase(args[1], "set")) {
+    return kCmdExclusive;
+  }
+
+  return 0;
+}
+
 REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandAuth>("auth", 2, "read-only ok-loading", 0, 0, 0),
                         MakeCmdAttr<CommandPing>("ping", -1, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandSelect>("select", 2, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandInfo>("info", -1, "read-only ok-loading", 0, 0, 0),
                         MakeCmdAttr<CommandRole>("role", 1, "read-only ok-loading", 0, 0, 0),
-                        MakeCmdAttr<CommandConfig>("config", -2, "read-only", 0, 0, 0),
+                        MakeCmdAttr<CommandConfig>("config", -2, "read-only", 0, 0, 0, GenerateConfigFlag),
                         MakeCmdAttr<CommandNamespace>("namespace", -3, "read-only exclusive", 0, 0, 0),
                         MakeCmdAttr<CommandKeys>("keys", 2, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandFlushDB>("flushdb", 1, "write", 0, 0, 0),

src/commands/commander.cc

@@ -46,7 +46,7 @@ std::string GetCommandInfo(const CommandAttributes *command_attributes) {
   command.append(redis::BulkString(command_attributes->name));
   command.append(redis::Integer(command_attributes->arity));
   command_flags.append(redis::MultiLen(1));
-  command_flags.append(redis::BulkString(command_attributes->IsWrite() ? "write" : "readonly"));
+  command_flags.append(redis::BulkString(command_attributes->flags & kCmdWrite ? "write" : "readonly"));
   command.append(command_flags);
   command.append(redis::Integer(command_attributes->key_range.first_key));
   command.append(redis::Integer(command_attributes->key_range.last_key));
@@ -77,30 +77,37 @@ void GetCommandsInfo(std::string *info, const std::vector<std::string> &cmd_name
   }
 }
 
-Status GetKeysFromCommand(const std::string &cmd_name, int argc, std::vector<int> *keys_indexes) {
-  auto cmd_iter = command_details::original_commands.find(util::ToLower(cmd_name));
-  if (cmd_iter == command_details::original_commands.end()) {
-    return {Status::RedisUnknownCmd, "Invalid command specified"};
+void DumpKeyRange(std::vector<int> &keys_index, int argc, const CommandKeyRange &range) {
+  auto last = range.last_key;
+  if (last < 0) last = argc + last;
+
+  for (int i = range.first_key; i <= last; i += range.key_step) {
+    keys_index.emplace_back(i);
   }
+}
 
-  auto command_attribute = cmd_iter->second;
-  if (command_attribute->key_range.first_key == 0) {
+Status GetKeysFromCommand(const CommandAttributes *attributes, const std::vector<std::string> &cmd_tokens,
+                          std::vector<int> *keys_index) {
+  if (attributes->key_range.first_key == 0) {
     return {Status::NotOK, "The command has no key arguments"};
   }
 
-  if (command_attribute->key_range.first_key < 0) {
-    return {Status::NotOK, "The command has dynamic positions of key arguments"};
-  }
+  int argc = static_cast<int>(cmd_tokens.size());
 
-  if ((command_attribute->arity > 0 && command_attribute->arity != argc) || argc < -command_attribute->arity) {
+  if ((attributes->arity > 0 && attributes->arity != argc) || argc < -attributes->arity) {
     return {Status::NotOK, "Invalid number of arguments specified for command"};
   }
 
-  auto last = command_attribute->key_range.last_key;
-  if (last < 0) last = argc + last;
-
-  for (int j = command_attribute->key_range.first_key; j <= last; j += command_attribute->key_range.key_step) {
-    keys_indexes->emplace_back(j);
+  if (attributes->key_range.first_key > 0) {
+    DumpKeyRange(*keys_index, argc, attributes->key_range);
+  } else if (attributes->key_range.first_key == -1) {
+    DumpKeyRange(*keys_index, argc, attributes->key_range_gen(cmd_tokens));
+  } else if (attributes->key_range.first_key == -2) {
+    for (const auto &range : attributes->key_range_vec_gen(cmd_tokens)) {
+      DumpKeyRange(*keys_index, argc, range);
+    }
+  } else {
+    return {Status::NotOK, "The key range specification is invalid"};
   }
 
   return Status::OK();

src/commands/commander.h

@@ -50,7 +50,7 @@ namespace redis {
 class Connection;
 struct CommandAttributes;
 
-enum CommandFlags {
+enum CommandFlags : uint64_t {
   kCmdWrite = 1ULL << 0,        // "write" flag
   kCmdReadOnly = 1ULL << 1,     // "read-only" flag
   kCmdReplication = 1ULL << 2,  // "replication" flag
@@ -115,17 +115,27 @@ using CommandKeyRangeGen = std::function<CommandKeyRange(const std::vector<std::
 
 using CommandKeyRangeVecGen = std::function<std::vector<CommandKeyRange>(const std::vector<std::string> &)>;
 
+using AdditionalFlagGen = std::function<uint64_t(const std::vector<std::string> &)>;
+
 struct CommandAttributes {
+  // command name
   std::string name;
 
   // number of command arguments
   // positive number n means number of arguments is equal to n
   // negative number -n means number of arguments is equal to or large than n
   int arity;
 
+  // space-splitted flag strings to initialize flags
   std::string description;
+
+  // bitmap of enum CommandFlags
   uint64_t flags;
 
+  // additional flags regarding to dynamic command arguments
+  AdditionalFlagGen flag_gen;
+
+  // static determined key range
   CommandKeyRange key_range;
 
   // if key_range.first_key == -1, key_range_gen is used instead
@@ -134,13 +144,14 @@ struct CommandAttributes {
   // if key_range.first_key == -2, key_range_vec_gen is used instead
   CommandKeyRangeVecGen key_range_vec_gen;
 
+  // commander object generator
   CommanderFactory factory;
 
-  bool IsWrite() const { return (flags & kCmdWrite) != 0; }
-  bool IsOkLoading() const { return (flags & kCmdLoading) != 0; }
-  bool IsExclusive() const { return (flags & kCmdExclusive) != 0; }
-  bool IsMulti() const { return (flags & kCmdMulti) != 0; }
-  bool IsNoMulti() const { return (flags & kCmdNoMulti) != 0; }
+  auto GenerateFlags(const std::vector<std::string> &args) const {
+    uint64_t res = flags;
+    if (flag_gen) res |= flag_gen(args);
+    return res;
+  }
 };
 
 using CommandMap = std::map<std::string, const CommandAttributes *>;
@@ -184,11 +195,12 @@ inline uint64_t ParseCommandFlags(const std::string &description, const std::str
 
 template <typename T>
 auto MakeCmdAttr(const std::string &name, int arity, const std::string &description, int first_key, int last_key,
-                 int key_step) {
+                 int key_step, const AdditionalFlagGen &flag_gen = {}) {
   CommandAttributes attr{name,
                          arity,
                          description,
                          ParseCommandFlags(description, name),
+                         flag_gen,
                          {first_key, last_key, key_step},
                          {},
                          {},
@@ -203,24 +215,33 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript
 }
 
 template <typename T>
-auto MakeCmdAttr(const std::string &name, int arity, const std::string &description, const CommandKeyRangeGen &gen) {
-  CommandAttributes attr{
-      name,        arity,
-      description, ParseCommandFlags(description, name),
-      {-1, 0, 0},  gen,
-      {},          []() -> std::unique_ptr<Commander> { return std::unique_ptr<Commander>(new T()); }};
+auto MakeCmdAttr(const std::string &name, int arity, const std::string &description, const CommandKeyRangeGen &gen,
+                 const AdditionalFlagGen &flag_gen = {}) {
+  CommandAttributes attr{name,
+                         arity,
+                         description,
+                         ParseCommandFlags(description, name),
+                         flag_gen,
+                         {-1, 0, 0},
+                         gen,
+                         {},
+                         []() -> std::unique_ptr<Commander> { return std::unique_ptr<Commander>(new T()); }};
 
   return attr;
 }
 
 template <typename T>
 auto MakeCmdAttr(const std::string &name, int arity, const std::string &description,
-                 const CommandKeyRangeVecGen &vec_gen) {
-  CommandAttributes attr{
-      name,        arity,
-      description, ParseCommandFlags(description, name),
-      {-2, 0, 0},  {},
-      vec_gen,     []() -> std::unique_ptr<Commander> { return std::unique_ptr<Commander>(new T()); }};
+                 const CommandKeyRangeVecGen &vec_gen, const AdditionalFlagGen &flag_gen = {}) {
+  CommandAttributes attr{name,
+                         arity,
+                         description,
+                         ParseCommandFlags(description, name),
+                         flag_gen,
+                         {-2, 0, 0},
+                         {},
+                         vec_gen,
+                         []() -> std::unique_ptr<Commander> { return std::unique_ptr<Commander>(new T()); }};
 
   return attr;
 }
@@ -254,7 +275,8 @@ const CommandMap *GetOriginalCommands();
 void GetAllCommandsInfo(std::string *info);
 void GetCommandsInfo(std::string *info, const std::vector<std::string> &cmd_names);
 std::string GetCommandInfo(const CommandAttributes *command_attributes);
-Status GetKeysFromCommand(const std::string &name, int argc, std::vector<int> *keys_indexes);
+Status GetKeysFromCommand(const CommandAttributes *attributes, const std::vector<std::string> &cmd_tokens,
+                          std::vector<int> *keys_indexes);
 bool IsCommandExists(const std::string &name);
 
 }  // namespace redis