diff --git a/src/commands/commander.h b/src/commands/commander.h index a2cda41be22..874957a4053 100644 --- a/src/commands/commander.h +++ b/src/commands/commander.h @@ -113,6 +113,8 @@ struct CommandKeyRange { using CommandKeyRangeGen = std::function &)>; +using CommandKeyRangeVecGen = std::function(const std::vector &)>; + struct CommandAttributes { std::string name; @@ -129,6 +131,9 @@ struct CommandAttributes { // 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; + CommanderFactory factory; bool IsWrite() const { return (flags & kCmdWrite) != 0; } @@ -186,6 +191,7 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript ParseCommandFlags(description, name), {first_key, last_key, key_step}, {}, + {}, []() -> std::unique_ptr { return std::unique_ptr(new T()); }}; if ((first_key > 0 && key_step <= 0) || (first_key > 0 && last_key >= 0 && last_key < first_key)) { @@ -198,13 +204,23 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript template 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 { return std::unique_ptr(new T()); }}; + CommandAttributes attr{ + name, arity, + description, ParseCommandFlags(description, name), + {-1, 0, 0}, gen, + {}, []() -> std::unique_ptr { return std::unique_ptr(new T()); }}; + + return attr; +} + +template +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 { return std::unique_ptr(new T()); }}; return attr; } diff --git a/src/server/server.cc b/src/server/server.cc index a5604496b8b..9115ffea33b 100644 --- a/src/server/server.cc +++ b/src/server/server.cc @@ -1683,12 +1683,21 @@ void Server::UpdateWatchedKeysFromArgs(const std::vector &args, con if (attr.IsWrite() && watched_key_size_ > 0) { if (attr.key_range.first_key > 0) { updateWatchedKeysFromRange(args, attr.key_range); - } else if (attr.key_range.first_key < 0) { + } 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); + + if (vec_range.front().first_key > 0) { + for (const auto &range : vec_range) { + updateWatchedKeysFromRange(args, range); + } + } + } else { // support commands like flushdb (write flag && key range {0,0,0}) updateAllWatchedKeys();