Add store key support in key range

For store key related commands, we cannot extract the store
key, see apache#2051. This PR adds store_key to CommandKeyRange to
represent its index.

This PR handle these store key commands:
- zunionstore
- zinterstore
- zdiffstore
- georadius
- georadiusbymember
- geosearchstore

before:
```
> command getkeys zinterstore dst 2 src1 src2
1) "src1"
2) "src2"
> command getkeys GEORADIUS src 1 1 1 km store dst
1) "src"
> command getkeys GEORADIUSBYMEMBER src member radius m store dst
1) "src"
> command getkeys GEOSEARCHSTORE dst src frommember member byradius 10 m
1) "dst"
```

after:
```
> command getkeys zinterstore dst 2 src1 src2
1) "dst"
2) "src1"
3) "src2"
> command getkeys GEORADIUS src 1 1 1 km store dst
1) "src"
2) "dst"
> command getkeys GEORADIUSBYMEMBER src member radius m store dst
1) "src"
2) "dst"
> command getkeys GEOSEARCHSTORE dst src frommember member byradius 10 m
1) "dst"
2) "src"
```

src/commands/cmd_geo.cc

@@ -346,7 +346,25 @@ class CommandGeoRadius : public CommandGeoBase {
     return redis::Array(list);
   }
 
+  static CommandKeyRange Range(const std::vector<std::string> &args) {
+    int store_key = 0;
+
+    // Check for the presence of the stored key in the command args.
+    for (size_t i = 6; i < args.size(); i++) {
+      // For the case when a user specifies both "store" and "storedist" options,
+      // the second key will override the first key. The behavior is kept the same
+      // as in ParseRadiusExtraOption method.
+      if ((util::ToLower(args[i]) == "store" || util::ToLower(args[i]) == "storedist") && i + 1 < args.size()) {
+        store_key = (int)i + 1;
+        i++;
+      }
+    }
+
+    return {1, 1, 1, store_key};
+  }
+
  protected:
+  int option_start_index_ = 6;
   double radius_ = 0;
   bool with_coord_ = false;
   bool with_dist_ = false;
@@ -649,6 +667,23 @@ class CommandGeoRadiusByMember : public CommandGeoRadius {
 
     return Status::OK();
   }
+
+  static CommandKeyRange Range(const std::vector<std::string> &args) {
+    int store_key = 0;
+
+    // Check for the presence of the stored key in the command args.
+    for (size_t i = 5; i < args.size(); i++) {
+      // For the case when a user specifies both "store" and "storedist" options,
+      // the second key will override the first key. The behavior is kept the same
+      // as in ParseRadiusExtraOption method.
+      if ((util::ToLower(args[i]) == "store" || util::ToLower(args[i]) == "storedist") && i + 1 < args.size()) {
+        store_key = (int)i + 1;
+        i++;
+      }
+    }
+
+    return {1, 1, 1, store_key};
+  }
 };
 
 class CommandGeoRadiusReadonly : public CommandGeoRadius {
@@ -665,11 +700,12 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandGeoAdd>("geoadd", -5, "write", 1, 1,
                         MakeCmdAttr<CommandGeoDist>("geodist", -4, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandGeoHash>("geohash", -3, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandGeoPos>("geopos", -3, "read-only", 1, 1, 1),
-                        MakeCmdAttr<CommandGeoRadius>("georadius", -6, "write", 1, 1, 1),
-                        MakeCmdAttr<CommandGeoRadiusByMember>("georadiusbymember", -5, "write", 1, 1, 1),
+                        MakeCmdAttr<CommandGeoRadius>("georadius", -6, "write", CommandGeoRadius::Range),
+                        MakeCmdAttr<CommandGeoRadiusByMember>("georadiusbymember", -5, "write",
+                                                              CommandGeoRadiusByMember::Range),
                         MakeCmdAttr<CommandGeoRadiusReadonly>("georadius_ro", -6, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandGeoRadiusByMemberReadonly>("georadiusbymember_ro", -5, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandGeoSearch>("geosearch", -7, "read-only", 1, 1, 1),
-                        MakeCmdAttr<CommandGeoSearchStore>("geosearchstore", -8, "write", 1, 1, 1))
+                        MakeCmdAttr<CommandGeoSearchStore>("geosearchstore", -8, "write", 2, 2, 1, 1))
 
 }  // namespace redis

src/commands/cmd_list.cc

@@ -225,6 +225,7 @@ class CommandLMPop : public Commander {
     // This parsing would always succeed as this cmd has been parsed before.
     auto num_key = *ParseInt<int32_t>(args[1], 10);
     range.last_key = range.first_key + num_key - 1;
+    range.store_key = 0;
     return range;
   };
 
@@ -446,6 +447,7 @@ class CommandBLMPop : public BlockingCommander {
     // This parsing would always succeed as this cmd has been parsed before.
     auto num_key = *ParseInt<int32_t>(args[2], 10);
     range.last_key = range.first_key + num_key - 1;
+    range.store_key = 0;
     return range;
   };
 

src/commands/cmd_zset.cc

@@ -1224,7 +1224,7 @@ class CommandZUnionStore : public Commander {
 
   static CommandKeyRange Range(const std::vector<std::string> &args) {
     int num_key = *ParseInt<int>(args[2], 10);
-    return {3, 2 + num_key, 1};
+    return {3, 2 + num_key, 1, 1};
   }
 
  protected:
@@ -1251,7 +1251,7 @@ class CommandZInterStore : public CommandZUnionStore {
 
   static CommandKeyRange Range(const std::vector<std::string> &args) {
     int num_key = *ParseInt<int>(args[2], 10);
-    return {3, 2 + num_key, 1};
+    return {3, 2 + num_key, 1, 1};
   }
 };
 
@@ -1505,7 +1505,7 @@ class CommandZDiffStore : public Commander {
 
   static CommandKeyRange Range(const std::vector<std::string> &args) {
     int num_key = *ParseInt<int>(args[2], 10);
-    return {3, 2 + num_key, 1};
+    return {3, 2 + num_key, 1, 1};
   }
 
  protected:

src/commands/commander.cc

@@ -81,9 +81,19 @@ void DumpKeyRange(std::vector<int> &keys_index, int argc, const CommandKeyRange
   auto last = range.last_key;
   if (last < 0) last = argc + last;
 
+  // The index of store key is before the first key
+  if (range.store_key > 0 && range.store_key < range.first_key) {
+    keys_index.emplace_back(range.store_key);
+  }
+
   for (int i = range.first_key; i <= last; i += range.key_step) {
     keys_index.emplace_back(i);
   }
+
+  // The index of store key is after the last key
+  if (range.store_key > 0 && range.store_key > last) {
+    keys_index.emplace_back(range.store_key);
+  }
 }
 
 Status CommandTable::GetKeysFromCommand(const CommandAttributes *attributes, const std::vector<std::string> &cmd_tokens,

src/commands/commander.h

@@ -93,18 +93,22 @@ class CommanderWithParseMove : Commander {
 using CommanderFactory = std::function<std::unique_ptr<Commander>()>;
 
 struct CommandKeyRange {
-  // index of the first key in command tokens
+  // index of the first key (non-store key) in command tokens
   // 0 stands for no key, since the first index of command arguments is command name
   int first_key;
 
-  // index of the last key in command tokens
+  // index of the last key (non-store key) in command tokens
   // in normal one-key commands, first key and last key index are both 1
   // -n stands for the n-th last index of the sequence, i.e. args.size() - n
   int last_key;
 
-  // step length of key position
+  // step length of key (non-store key) position
   // e.g. key step 2 means "key other key other ..." sequence
   int key_step;
+
+  // index of the store key in command tokens
+  // 0 stands for no store key, since the first index of command arguments is command name
+  int store_key;
 };
 
 using CommandKeyRangeGen = std::function<CommandKeyRange(const std::vector<std::string> &)>;
@@ -197,7 +201,7 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript
                          description,
                          ParseCommandFlags(description, name),
                          flag_gen,
-                         {first_key, last_key, key_step},
+                         {first_key, last_key, key_step, 0},
                          {},
                          {},
                          []() -> std::unique_ptr<Commander> { return std::unique_ptr<Commander>(new T()); }};
@@ -210,6 +214,27 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript
   return attr;
 }
 
+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 store_key, const AdditionalFlagGen &flag_gen = {}) {
+  CommandAttributes attr{name,
+                         arity,
+                         description,
+                         ParseCommandFlags(description, name),
+                         flag_gen,
+                         {first_key, last_key, key_step, store_key},
+                         {},
+                         {},
+                         []() -> 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) || (store_key < 0)) {
+    std::cout << fmt::format("Encountered invalid key range in command {}", name) << std::endl;
+    std::abort();
+  }
+
+  return attr;
+}
+
 template <typename T>
 auto MakeCmdAttr(const std::string &name, int arity, const std::string &description, const CommandKeyRangeGen &gen,
                  const AdditionalFlagGen &flag_gen = {}) {
@@ -218,7 +243,7 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript
                          description,
                          ParseCommandFlags(description, name),
                          flag_gen,
-                         {-1, 0, 0},
+                         {-1, 0, 0, 0},
                          gen,
                          {},
                          []() -> std::unique_ptr<Commander> { return std::unique_ptr<Commander>(new T()); }};
@@ -234,7 +259,7 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript
                          description,
                          ParseCommandFlags(description, name),
                          flag_gen,
-                         {-2, 0, 0},
+                         {-2, 0, 0, 0},
                          {},
                          vec_gen,
                          []() -> std::unique_ptr<Commander> { return std::unique_ptr<Commander>(new T()); }};

tests/gocase/unit/command/command_test.go

@@ -88,9 +88,10 @@ func TestCommand(t *testing.T) {
 		r := rdb.Do(ctx, "COMMAND", "GETKEYS", "ZINTERSTORE", "dst", "2", "src1", "src2")
 		vs, err := r.Slice()
 		require.NoError(t, err)
-		require.Len(t, vs, 2)
-		require.Equal(t, "src1", vs[0])
-		require.Equal(t, "src2", vs[1])
+		require.Len(t, vs, 3)
+		require.Equal(t, "dst", vs[0])
+		require.Equal(t, "src1", vs[1])
+		require.Equal(t, "src2", vs[2])
 	})
 
 	t.Run("COMMAND GETKEYS ZINTERCARD", func(t *testing.T) {
@@ -115,9 +116,10 @@ func TestCommand(t *testing.T) {
 		r := rdb.Do(ctx, "COMMAND", "GETKEYS", "ZUNIONSTORE", "dst", "2", "src1", "src2")
 		vs, err := r.Slice()
 		require.NoError(t, err)
-		require.Len(t, vs, 2)
-		require.Equal(t, "src1", vs[0])
-		require.Equal(t, "src2", vs[1])
+		require.Len(t, vs, 3)
+		require.Equal(t, "dst", vs[0])
+		require.Equal(t, "src1", vs[1])
+		require.Equal(t, "src2", vs[2])
 	})
 
 	t.Run("COMMAND GETKEYS ZDIFF", func(t *testing.T) {
@@ -133,9 +135,10 @@ func TestCommand(t *testing.T) {
 		r := rdb.Do(ctx, "COMMAND", "GETKEYS", "ZDIFFSTORE", "dst", "2", "src1", "src2")
 		vs, err := r.Slice()
 		require.NoError(t, err)
-		require.Len(t, vs, 2)
-		require.Equal(t, "src1", vs[0])
-		require.Equal(t, "src2", vs[1])
+		require.Len(t, vs, 3)
+		require.Equal(t, "dst", vs[0])
+		require.Equal(t, "src1", vs[1])
+		require.Equal(t, "src2", vs[2])
 	})
 
 	t.Run("COMMAND GETKEYS ZMPOP", func(t *testing.T) {
@@ -173,4 +176,59 @@ func TestCommand(t *testing.T) {
 		require.Equal(t, "key1", vs[0])
 		require.Equal(t, "key2", vs[1])
 	})
+
+	t.Run("COMMAND GETKEYS GEORADIUS", func(t *testing.T) {
+		r := rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUS", "src", "1", "1", "1", "km", "store", "dst")
+		vs, err := r.Slice()
+		require.NoError(t, err)
+		require.Len(t, vs, 2)
+		require.Equal(t, "src", vs[0])
+		require.Equal(t, "dst", vs[1])
+
+		r = rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUS", "src", "1", "1", "1", "km", "storedist", "dst")
+		vs, err = r.Slice()
+		require.NoError(t, err)
+		require.Len(t, vs, 2)
+		require.Equal(t, "src", vs[0])
+		require.Equal(t, "dst", vs[1])
+
+		r = rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUS", "src", "1", "1", "1", "km", "store", "dst1", "storedist", "dst2")
+		vs, err = r.Slice()
+		require.NoError(t, err)
+		require.Len(t, vs, 2)
+		require.Equal(t, "src", vs[0])
+		require.Equal(t, "dst2", vs[1])
+	})
+
+	t.Run("COMMAND GETKEYS GEORADIUSBYMEMBER", func(t *testing.T) {
+		r := rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUSBYMEMBER", "src", "member", "radius", "m", "store", "dst")
+		vs, err := r.Slice()
+		require.NoError(t, err)
+		require.Len(t, vs, 2)
+		require.Equal(t, "src", vs[0])
+		require.Equal(t, "dst", vs[1])
+
+		r = rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUSBYMEMBER", "src", "member", "radius", "m", "storedist", "dst")
+		vs, err = r.Slice()
+		require.NoError(t, err)
+		require.Len(t, vs, 2)
+		require.Equal(t, "src", vs[0])
+		require.Equal(t, "dst", vs[1])
+
+		r = rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUSBYMEMBER", "src", "member", "radius", "m", "store", "dst1", "storedist", "dst2")
+		vs, err = r.Slice()
+		require.NoError(t, err)
+		require.Len(t, vs, 2)
+		require.Equal(t, "src", vs[0])
+		require.Equal(t, "dst2", vs[1])
+	})
+
+	t.Run("COMMAND GETKEYS GEOSEARCHSTORE", func(t *testing.T) {
+		r := rdb.Do(ctx, "COMMAND", "GETKEYS", "GEOSEARCHSTORE", "dst", "src", "frommember", "member", "byradius", "10", "m")
+		vs, err := r.Slice()
+		require.NoError(t, err)
+		require.Len(t, vs, 2)
+		require.Equal(t, "dst", vs[0])
+		require.Equal(t, "src", vs[1])
+	})
 }