fix: replication for Z...STORE commands (#5728)

* fix: replication for Z...STORE commands

Author: BorysTheDev

src/server/set_family.cc

@@ -488,7 +488,7 @@ OpResult<uint32_t> OpAdd(const OpArgs& op_args, std::string_view key, const NewE
         RecordJournal(op_args, "DEL"sv, ArgSlice{key});
       }
     }
-    return 0;
+    return OpStatus::OK;
   }
 
   // We can use std::nullopt here because we check the type later.
@@ -544,6 +544,8 @@ OpResult<uint32_t> OpAdd(const OpArgs& op_args, std::string_view key, const NewE
     res = StringSetWrapper{co, op_args.db_cntx}.Add(vals, UINT32_MAX, false);
   }
 
+  // TODO: consider optimization to record real command if the replica is in stable_sync state
+  // and there is no slot migration process going on.
   if (journal_update && op_args.shard->journal()) {
     if (overwrite) {
       RecordJournal(op_args, "DEL"sv, ArgSlice{key});

src/server/zset_family.cc

@@ -152,9 +152,9 @@ void OutputScoredArrayResult(const OpResult<ScoredArray>& result, SinkReplyBuild
   rb->SendScoredArray(result.value(), true /* with scores */);
 }
 
-OpResult<DbSlice::ItAndUpdater> FindZEntry(const ZSetFamily::ZParams& zparams,
-                                           const OpArgs& op_args, string_view key,
-                                           size_t member_len) {
+OpResult<DbSlice::ItAndUpdater> PrepareZEntry(const ZSetFamily::ZParams& zparams,
+                                              const OpArgs& op_args, string_view key,
+                                              size_t member_len) {
   auto& db_slice = op_args.GetDbSlice();
   if (zparams.flags & ZADD_IN_XX) {
     return db_slice.FindMutable(op_args.db_cntx, key, OBJ_ZSET);
@@ -1551,7 +1551,8 @@ void ZBooleanOperation(CmdArgList args, string_view cmd, bool is_union, bool sto
     auto store_cb = [&, dest_shard = Shard(dest_key, maps.size())](Transaction* t,
                                                                    EngineShard* shard) {
       if (shard->shard_id() == dest_shard)
-        ZSetFamily::OpAdd(t->GetOpArgs(shard), ZSetFamily::ZParams{.override = true}, dest_key,
+        ZSetFamily::OpAdd(t->GetOpArgs(shard),
+                          ZSetFamily::ZParams{.override = true, .journal_update = true}, dest_key,
                           smvec);
       return OpStatus::OK;
     };
@@ -1696,7 +1697,8 @@ void ZRangeInternal(CmdArgList args, ZSetFamily::RangeParams range_params, Trans
       mvec[i++] = {score, str};
     }
 
-    add_result = ZSetFamily::OpAdd(t->GetOpArgs(shard), ZSetFamily::ZParams{.override = true},
+    add_result = ZSetFamily::OpAdd(t->GetOpArgs(shard),
+                                   ZSetFamily::ZParams{.override = true, .journal_update = true},
                                    *range_params.store_key, mvec);
 
     return OpStatus::OK;
@@ -1967,6 +1969,9 @@ OpResult<ZSetFamily::AddResult> ZSetFamily::OpAdd(const OpArgs& op_args,
     if (res_it && IsValid(res_it->it)) {
       res_it->post_updater.Run();
       db_slice.Del(op_args.db_cntx, res_it->it);
+      if (zparams.journal_update && op_args.shard->journal()) {
+        RecordJournal(op_args, "DEL"sv, ArgSlice{key});
+      }
     }
     return OpStatus::OK;
   }
@@ -1976,7 +1981,7 @@ OpResult<ZSetFamily::AddResult> ZSetFamily::OpAdd(const OpArgs& op_args,
   size_t field_len = members.size() > server.zset_max_listpack_entries
                          ? UINT32_MAX
                          : members.front().second.size();
-  auto res_it = FindZEntry(zparams, op_args, key, field_len);
+  auto res_it = PrepareZEntry(zparams, op_args, key, field_len);
 
   if (!res_it)
     return res_it.status();
@@ -2039,6 +2044,26 @@ OpResult<ZSetFamily::AddResult> ZSetFamily::OpAdd(const OpArgs& op_args,
 
   if (op_status != OpStatus::OK)
     return op_status;
+
+  // TODO: consider optimization to record real command if the replica is in stable_sync state
+  // and there is no slot migration process going on.
+  if (zparams.journal_update && op_args.shard->journal()) {
+    if (zparams.override) {
+      RecordJournal(op_args, "DEL"sv, ArgSlice{key});
+    }
+
+    vector<string> scores;
+    vector<string_view> mapped;
+    scores.reserve(members.size());
+    mapped.reserve(members.size() * 2 + 1);
+    mapped.push_back(key);
+    for (const auto& [score, member] : members) {
+      scores.push_back(absl::StrCat(score));
+      mapped.push_back(scores.back());
+      mapped.push_back(member);
+    }
+    RecordJournal(op_args, "ZADD"sv, mapped);
+  }
   return aresult;
 }
 
@@ -2336,7 +2361,8 @@ void ZSetFamily::ZDiffStore(CmdArgList args, const CommandContext& cmd_cntx) {
 
   auto store_cb = [&](Transaction* t, EngineShard* shard) {
     if (shard->shard_id() == dest_shard)
-      ZSetFamily::OpAdd(t->GetOpArgs(shard), ZSetFamily::ZParams{.override = true}, dest_key,
+      ZSetFamily::OpAdd(t->GetOpArgs(shard),
+                        ZSetFamily::ZParams{.override = true, .journal_update = true}, dest_key,
                         smvec);
     return OpStatus::OK;
   };
@@ -2836,7 +2862,7 @@ constexpr uint32_t kZUnionStore = WRITE | SORTEDSET | SLOW;
 }  // namespace acl
 
 void ZSetFamily::Register(CommandRegistry* registry) {
-  constexpr uint32_t kStoreMask = CO::WRITE | CO::VARIADIC_KEYS | CO::DENYOOM;
+  constexpr uint32_t kStoreMask = CO::WRITE | CO::VARIADIC_KEYS | CO::DENYOOM | CO::NO_AUTOJOURNAL;
   registry->StartFamily();
   // TODO: to add support for SCRIPT for BZPOPMIN, BZPOPMAX similarly to BLPOP.
   *registry

src/server/zset_family.h

@@ -64,6 +64,7 @@ class ZSetFamily {
     unsigned flags = 0;  // mask of ZADD_IN_ macros.
     bool ch = false;     // Corresponds to CH option.
     bool override = false;
+    bool journal_update = false;
   };
 
   using ScoredMember = std::pair<std::string, double>;

tests/dragonfly/replication_test.py

@@ -674,6 +674,22 @@ async def check_expire(key):
         # Check SUNIONSTORE turns into DEL and SADD
         await check_list_ooo("SUNIONSTORE k set1 set2", [r"DEL k", r"SADD k (.*?)"])
 
+        # Check ZDIFFSTORE turns into DEL and ZADD
+        await c_master.execute_command("zadd zet1 1 v1 2 v2 3 v3")
+        await c_master.execute_command("zadd zet2 1 v1 2 v2")
+        await skip_cmd()
+        await skip_cmd()
+        await check_list("ZDIFFSTORE k 2 zet1 zet2", [r"DEL k", r"ZADD k 3 v3"])
+
+        # Check ZINTERSTORE turns into DEL and ZADD
+        await check_list("ZINTERSTORE k 2 zet1 zet2", [r"DEL k", r"ZADD k (.*?)"])
+
+        # Check ZRANGESTORE turns into SREM and ZADD
+        await check_list_ooo("ZRANGESTORE k zet1 2 -1", [r"DEL k", r"ZADD k 3 v3"])
+
+        # Check ZUNIONSTORE turns into DEL and ZADD
+        await check_list_ooo("ZUNIONSTORE k 2 zet1 zet2", [r"DEL k", r"ZADD k (.*?)"])
+
         await c_master.set("k1", "1000")
         await c_master.set("k2", "1100")
         await skip_cmd()
@@ -751,6 +767,7 @@ async def check_expire(key):
             [r"XTRIM k-stream MINID 0", r"SREM k-one-element-set value[12]"],
         )
 
+        # TODO next Z-tests won't work with no-point-in-time replication
         # check BZMPOP turns into ZPOPMAX and ZPOPMIN command
         await c_master.zadd("key", {"a": 1, "b": 2, "c": 3})
         await skip_cmd()