diff --git a/src/cluster/cluster.cc b/src/cluster/cluster.cc index 16e72b62fbc..c4dd9b8c697 100644 --- a/src/cluster/cluster.cc +++ b/src/cluster/cluster.cc @@ -832,16 +832,16 @@ bool Cluster::IsWriteForbiddenSlot(int slot) const { Status Cluster::CanExecByMySelf(const redis::CommandAttributes *attributes, const std::vector &cmd_tokens, redis::Connection *conn, lua::ScriptRunCtx *script_run_ctx) { - std::vector keys_indexes; + std::vector key_indexes; - // No keys - if (auto s = redis::CommandTable::GetKeysFromCommand(attributes, cmd_tokens, &keys_indexes); !s.IsOK()) - return Status::OK(); + auto s = redis::CommandTable::GetKeysFromCommand(attributes, cmd_tokens); + if (!s) return Status::OK(); + key_indexes = *s; - if (keys_indexes.empty()) return Status::OK(); + if (key_indexes.empty()) return Status::OK(); int slot = -1; - for (auto i : keys_indexes) { + for (auto i : key_indexes) { if (i >= static_cast(cmd_tokens.size())) break; int cur_slot = GetSlotIdFromKey(cmd_tokens[i]); diff --git a/src/commands/cmd_server.cc b/src/commands/cmd_server.cc index 260edca1ebc..725b28d0cdf 100644 --- a/src/commands/cmd_server.cc +++ b/src/commands/cmd_server.cc @@ -676,18 +676,16 @@ class CommandCommand : public Commander { return {Status::RedisUnknownCmd, "Invalid command specified"}; } - std::vector keys_indexes; - auto s = CommandTable::GetKeysFromCommand( - cmd_iter->second, std::vector(args_.begin() + 2, args_.end()), &keys_indexes); - if (!s.IsOK()) return s; + auto key_indexes = GET_OR_RET(CommandTable::GetKeysFromCommand( + cmd_iter->second, std::vector(args_.begin() + 2, args_.end()))); - if (keys_indexes.size() == 0) { + if (key_indexes.size() == 0) { return {Status::RedisExecErr, "Invalid arguments specified for command"}; } std::vector keys; - keys.reserve(keys_indexes.size()); - for (const auto &key_index : keys_indexes) { + keys.reserve(key_indexes.size()); + for (const auto &key_index : key_indexes) { keys.emplace_back(args_[key_index + 2]); } *output = conn->MultiBulkString(keys); diff --git a/src/commands/commander.cc b/src/commands/commander.cc index ecab960f35a..063b1c469d0 100644 --- a/src/commands/commander.cc +++ b/src/commands/commander.cc @@ -50,9 +50,10 @@ std::string CommandTable::GetCommandInfo(const CommandAttributes *command_attrib command_flags.append(redis::MultiLen(1)); command_flags.append(redis::BulkString(command_attributes->InitialFlags() & 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)); - command.append(redis::Integer(command_attributes->key_range.key_step)); + auto key_range = command_attributes->InitialKeyRange().ValueOr({0, 0, 0}); + command.append(redis::Integer(key_range.first_key)); + command.append(redis::Integer(key_range.last_key)); + command.append(redis::Integer(key_range.key_step)); return command; } @@ -79,40 +80,31 @@ void CommandTable::GetCommandsInfo(std::string *info, const std::vector &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); - } -} - -Status CommandTable::GetKeysFromCommand(const CommandAttributes *attributes, const std::vector &cmd_tokens, - std::vector *keys_index) { - if (attributes->key_range.first_key == 0) { - return {Status::NotOK, "The command has no key arguments"}; - } - +StatusOr> CommandTable::GetKeysFromCommand(const CommandAttributes *attributes, + const std::vector &cmd_tokens) { int argc = static_cast(cmd_tokens.size()); - if ((attributes->arity > 0 && attributes->arity != argc) || argc < -attributes->arity) { + if (!attributes->CheckArity(argc)) { return {Status::NotOK, "Invalid number of arguments specified for command"}; } - 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"}; + Status status; + std::vector key_indexes; + + attributes->ForEachKeyRange( + [&](const std::vector &, CommandKeyRange key_range) { + key_range.ForEachKeyIndex([&](int i) { key_indexes.push_back(i); }, cmd_tokens.size()); + }, + cmd_tokens, + [&](const auto &) { + status = {Status::NotOK, "The command has no key arguments"}; + }); + + if (!status) { + return status; } - return Status::OK(); + return key_indexes; } bool CommandTable::IsExists(const std::string &name) { diff --git a/src/commands/commander.h b/src/commands/commander.h index 3c09fe87957..87efdb0131e 100644 --- a/src/commands/commander.h +++ b/src/commands/commander.h @@ -141,9 +141,18 @@ struct CommandKeyRange { template void ForEachKey(F &&f, const std::vector &args) const { for (size_t i = first_key; last_key > 0 ? i <= size_t(last_key) : i <= args.size() + last_key; i += key_step) { + if (i >= args.size()) continue; std::forward(f)(args[i]); } } + + template + void ForEachKeyIndex(F &&f, size_t arg_size) const { + for (size_t i = first_key; last_key > 0 ? i <= size_t(last_key) : i <= arg_size + last_key; i += key_step) { + if (i >= arg_size) continue; + std::forward(f)(i); + } + } }; using CommandKeyRangeGen = std::function &)>; @@ -161,20 +170,20 @@ struct CommandAttributes { : name(std::move(name)), arity(arity), category(category), - key_range{0, 0, 0}, factory(std::move(factory)), flags_(flags), - flag_gen_(std::move(flag_gen)) {} + flag_gen_(std::move(flag_gen)), + key_range_{0, 0, 0} {} CommandAttributes(std::string name, int arity, CommandCategory category, uint64_t flags, AdditionalFlagGen flag_gen, CommandKeyRange key_range, CommanderFactory factory) : name(std::move(name)), arity(arity), category(category), - key_range(key_range), factory(std::move(factory)), flags_(flags), - flag_gen_(std::move(flag_gen)) { + flag_gen_(std::move(flag_gen)), + key_range_(key_range) { if (key_range.first_key <= 0 || key_range.key_step <= 0 || (key_range.last_key >= 0 && key_range.last_key < key_range.first_key)) { std::cout << fmt::format("Encountered invalid key range in command {}", this->name) << std::endl; @@ -187,22 +196,22 @@ struct CommandAttributes { : name(std::move(name)), arity(arity), category(category), - key_range{-1, 0, 0}, - key_range_gen(std::move(key_range)), factory(std::move(factory)), flags_(flags), - flag_gen_(std::move(flag_gen)) {} + flag_gen_(std::move(flag_gen)), + key_range_{-1, 0, 0}, + key_range_gen_(std::move(key_range)) {} CommandAttributes(std::string name, int arity, CommandCategory category, uint64_t flags, AdditionalFlagGen flag_gen, CommandKeyRangeVecGen key_range, CommanderFactory factory) : name(std::move(name)), arity(arity), category(category), - key_range{-2, 0, 0}, - key_range_vec_gen(std::move(key_range)), factory(std::move(factory)), flags_(flags), - flag_gen_(std::move(flag_gen)) {} + flag_gen_(std::move(flag_gen)), + key_range_{-2, 0, 0}, + key_range_vec_gen_(std::move(key_range)) {} // command name std::string name; @@ -215,16 +224,6 @@ struct CommandAttributes { // category of this command, e.g. key, string, hash CommandCategory category; - // static determined key range - // if key_range.first_key == 0, there's no key in this command args - CommandKeyRange key_range; - - // if key_range.first_key == -1, key_range_gen is used instead - CommandKeyRangeGen key_range_gen; - - // if key_range.first_key == -2, key_range_vec_gen is used instead - CommandKeyRangeVecGen key_range_vec_gen; - // commander object generator CommanderFactory factory; @@ -240,33 +239,57 @@ struct CommandAttributes { return !((arity > 0 && cmd_size != arity) || (arity < 0 && cmd_size < -arity)); } - template - void ForEachKeyRange(F &&f, const std::vector &args) const { - if (key_range.first_key > 0) { - std::forward(f)(args, key_range); - } else if (key_range.first_key == -1) { - redis::CommandKeyRange range = key_range_gen(args); + StatusOr InitialKeyRange() const { + if (key_range_.first_key >= 0) return key_range_; + return {Status::NotOK, "key range is unavailable without command arguments"}; + } + + template + void ForEachKeyRange(F &&f, const std::vector &args, G &&g) const { + if (key_range_.first_key > 0) { + std::forward(f)(args, key_range_); + } else if (key_range_.first_key == -1) { + redis::CommandKeyRange range = key_range_gen_(args); if (range.first_key > 0) { std::forward(f)(args, range); } - } else if (key_range.first_key == -2) { - std::vector vec_range = key_range_vec_gen(args); + } else if (key_range_.first_key == -2) { + std::vector vec_range = key_range_vec_gen_(args); for (const auto &range : vec_range) { if (range.first_key > 0) { std::forward(f)(args, range); } } + } else if (key_range_.first_key == 0) { + // otherwise, if there's no key inside the command arguments + // e.g. FLUSHALL, with "write" flag but no key specified + std::forward(g)(args); } } + template + void ForEachKeyRange(F &&f, const std::vector &args) const { + ForEachKeyRange(std::forward(f), args, [](const auto &) {}); + } + private: // bitmap of enum CommandFlags uint64_t flags_; // additional flags regarding to dynamic command arguments AdditionalFlagGen flag_gen_; + + // static determined key range + // if key_range.first_key == 0, there's no key in this command args + CommandKeyRange key_range_; + + // if key_range.first_key == -1, key_range_gen is used instead + CommandKeyRangeGen key_range_gen_; + + // if key_range.first_key == -2, key_range_vec_gen is used instead + CommandKeyRangeVecGen key_range_vec_gen_; }; using CommandMap = std::map; @@ -374,8 +397,8 @@ struct CommandTable { static void GetAllCommandsInfo(std::string *info); static void GetCommandsInfo(std::string *info, const std::vector &cmd_names); static std::string GetCommandInfo(const CommandAttributes *command_attributes); - static Status GetKeysFromCommand(const CommandAttributes *attributes, const std::vector &cmd_tokens, - std::vector *keys_indexes); + static StatusOr> GetKeysFromCommand(const CommandAttributes *attributes, + const std::vector &cmd_tokens); static size_t Size(); static bool IsExists(const std::string &name); diff --git a/src/common/string_util.h b/src/common/string_util.h index ac3b2904fae..2dcb10800f5 100644 --- a/src/common/string_util.h +++ b/src/common/string_util.h @@ -41,8 +41,7 @@ std::string EscapeString(std::string_view s); std::string StringNext(std::string s); template -std::string StringJoin( - const T &con, F &&f = [](const auto &v) -> decltype(auto) { return v; }, std::string_view sep = ", ") { +std::string StringJoin(const T &con, F &&f, std::string_view sep = ", ") { std::string res; bool is_first = true; for (const auto &v : con) { diff --git a/src/server/server.cc b/src/server/server.cc index f20b6824fc7..1de5534dc54 100644 --- a/src/server/server.cc +++ b/src/server/server.cc @@ -1984,27 +1984,9 @@ void Server::updateAllWatchedKeys() { void Server::UpdateWatchedKeysFromArgs(const std::vector &args, const redis::CommandAttributes &attr) { if ((attr.GenerateFlags(args) & redis::kCmdWrite) && watched_key_size_ > 0) { - if (attr.key_range.first_key > 0) { - updateWatchedKeysFromRange(args, attr.key_range); - } else if (attr.key_range.first_key == -1) { - redis::CommandKeyRange range = attr.key_range_gen(args); - - if (range.first_key > 0) { - updateWatchedKeysFromRange(args, range); - } - } else if (attr.key_range.first_key == -2) { - std::vector vec_range = attr.key_range_vec_gen(args); - - for (const auto &range : vec_range) { - if (range.first_key > 0) { - updateWatchedKeysFromRange(args, range); - } - } - - } else { - // support commands like flushdb (write flag && key range {0,0,0}) - updateAllWatchedKeys(); - } + attr.ForEachKeyRange([this](const std::vector &args, + redis::CommandKeyRange range) { updateWatchedKeysFromRange(args, range); }, + args, [this](const std::vector &) { updateAllWatchedKeys(); }); } }