add initial key_range_vec

Signed-off-by: clundro <infdahai@outlook.com>
apache · Jul 8, 2023 · d99f251 · d99f251
1 parent bf49e4f
commit d99f251
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 8 deletions.
diff --git a/src/commands/commander.h b/src/commands/commander.h
@@ -113,6 +113,8 @@ struct CommandKeyRange {
 
 using CommandKeyRangeGen = std::function<CommandKeyRange(const std::vector<std::string> &)>;
 
+using CommandKeyRangeVecGen = std::function<std::vector<CommandKeyRange>(const std::vector<std::string> &)>;
+
 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<Commander> { return std::unique_ptr<Commander>(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 <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()); }};
+  CommandAttributes attr{
+      name,        arity,
+      description, ParseCommandFlags(description, name),
+      {-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()); }};
 
   return attr;
 }

diff --git a/src/server/server.cc b/src/server/server.cc
@@ -1683,12 +1683,21 @@ void Server::UpdateWatchedKeysFromArgs(const std::vector<std::string> &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<redis::CommandKeyRange> 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();