fix(hll_family): Don't use destination key in PFMERGE operation

Skip destination key to be used with PFMERGE. Fixed bug in
hllMergeDense that writes to wrong register.

Fixes #4750

Signed-off-by: mkaruza <[email protected]>

Author: mkaruza

src/redis/hyperloglog.c

@@ -1370,8 +1370,7 @@ int64_t pfcountSingle(struct HllBufferPtr hll_ptr) {
 }
 
 /* Merge dense-encoded HLL */
-static void hllMergeDense(uint8_t* max, struct HllBufferPtr to) {
-  uint8_t* registers = max + HLL_HDR_SIZE;
+static void hllMergeDense(uint8_t* registers, struct HllBufferPtr to) {
   uint8_t val;
   struct hllhdr* hll_hdr = (struct hllhdr*)to.hll;
 

src/server/hll_family.cc

@@ -169,11 +169,18 @@ OpResult<int64_t> CountHllsSingle(const OpArgs& op_args, string_view key) {
   }
 }
 
-OpResult<vector<string>> ReadValues(const OpArgs& op_args, const ShardArgs& keys) {
+OpResult<vector<string>> ReadValues(const OpArgs& op_args, const ShardArgs& keys,
+                                    bool skip_dest_key) {
   try {
     vector<string> values;
-    for (string_view key : keys) {
-      auto it = op_args.GetDbSlice().FindReadOnly(op_args.db_cntx, key, OBJ_STRING);
+    ShardArgs::Iterator start_key = keys.begin(), end_key = keys.end();
+    if (skip_dest_key) {
+      // If destkey is only found on this shard we can return
+      if (++start_key == end_key)
+        return OpStatus::OK;
+    }
+    for (; start_key != end_key; ++start_key) {
+      auto it = op_args.GetDbSlice().FindReadOnly(op_args.db_cntx, *start_key, OBJ_STRING);
       if (it.ok()) {
         string hll;
         it.value()->second.GetString(&hll);
@@ -211,7 +218,7 @@ OpResult<int64_t> PFCountMulti(CmdArgList args, const CommandContext& cmd_cntx)
   auto cb = [&](Transaction* t, EngineShard* shard) {
     ShardId sid = shard->shard_id();
     ShardArgs shard_args = t->GetShardArgs(shard->shard_id());
-    auto result = ReadValues(t->GetOpArgs(shard), shard_args);
+    auto result = ReadValues(t->GetOpArgs(shard), shard_args, false /* no dest key */);
     if (result.ok()) {
       hlls[sid] = std::move(result.value());
     }
@@ -246,12 +253,15 @@ void PFCount(CmdArgList args, const CommandContext& cmd_cntx) {
 OpResult<int> PFMergeInternal(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder) {
   vector<vector<string>> hlls;
   hlls.resize(shard_set->size());
+  const string_view dest_key = ArgS(args, 0);
+  const ShardId dest_shard = Shard(dest_key, shard_set->size());
 
   atomic_bool success = true;
   auto cb = [&](Transaction* t, EngineShard* shard) {
     ShardId sid = shard->shard_id();
+    const bool skip_dest_key = sid == dest_shard;
     ShardArgs shard_args = t->GetShardArgs(shard->shard_id());
-    auto result = ReadValues(t->GetOpArgs(shard), shard_args);
+    auto result = ReadValues(t->GetOpArgs(shard), shard_args, skip_dest_key);
     if (result.ok()) {
       hlls[sid] = std::move(result.value());
     } else {
@@ -275,16 +285,15 @@ OpResult<int> PFMergeInternal(CmdArgList args, Transaction* tx, SinkReplyBuilder
   int result = pfmerge(ptrs.data(), ptrs.size(), StringToHllPtr(hll));
 
   auto set_cb = [&](Transaction* t, EngineShard* shard) {
-    string_view key = ArgS(args, 0);
     const OpArgs& op_args = t->GetOpArgs(shard);
     auto& db_slice = op_args.GetDbSlice();
-    auto op_res = db_slice.AddOrFind(t->GetDbContext(), key);
+    auto op_res = db_slice.AddOrFind(t->GetDbContext(), dest_key);
     RETURN_ON_BAD_STATUS(op_res);
     auto& res = *op_res;
     res.it->second.SetString(hll);
 
     if (op_args.shard->journal()) {
-      RecordJournal(op_args, "SET", ArgSlice{key, hll});
+      RecordJournal(op_args, "SET", ArgSlice{dest_key, hll});
     }
 
     return OpStatus::OK;

src/server/hll_family_test.cc

@@ -166,8 +166,13 @@ TEST_F(HllFamilyTest, MergeToNew) {
 TEST_F(HllFamilyTest, MergeToExisting) {
   EXPECT_EQ(CheckedInt({"pfadd", "key1", "1", "2", "3"}), 1);
   EXPECT_EQ(CheckedInt({"pfadd", "key2", "4", "5"}), 1);
-  EXPECT_EQ(Run({"pfmerge", "key2", "key1"}), "OK");
-  EXPECT_EQ(CheckedInt({"pfcount", "key2"}), 5);
+  EXPECT_EQ(Run({"pfmerge", "key3", "key2", "key1"}), "OK");
+  EXPECT_EQ(CheckedInt({"pfcount", "key3"}), 5);
+  EXPECT_EQ(Run({"pfmerge", "key3", "key3"}), "OK");
+  EXPECT_EQ(CheckedInt({"pfcount", "key3"}), 5);
+  EXPECT_EQ(CheckedInt({"pfadd", "key4", "4", "5", "6"}), 1);
+  EXPECT_EQ(Run({"pfmerge", "key3", "key3", "key4"}), "OK");
+  EXPECT_EQ(CheckedInt({"pfcount", "key3"}), 6);
 }
 
 TEST_F(HllFamilyTest, MergeNonExisting) {