feat(server): Implement ZMSCORE command #357 (#367)

* feat(server): Implement ZMSCORE command (#357)
Signed-off-by: Logan Raarup <logan@logan.dk>

CONTRIBUTORS.md

@@ -4,6 +4,7 @@
 * **[Philipp Born](https://github.com/tamcore)**
   * Helm Chart
 * **[Braydn Moore](https://github.com/braydnm)**
+* **[Logan Raarup](https://github.com/logandk)**
 * **[Ryan Russell](https://github.com/ryanrussell)**
   * Docs & Code Readability
 * **[Ali-Akber Saifee](https://github.com/alisaifee)**

src/server/zset_family.cc

@@ -1274,6 +1274,35 @@ void ZSetFamily::ZScore(CmdArgList args, ConnectionContext* cntx) {
   }
 }
 
+void ZSetFamily::ZMScore(CmdArgList args, ConnectionContext* cntx) {
+  string_view key = ArgS(args, 1);
+
+  absl::InlinedVector<string_view, 8> members(args.size() - 2);
+  for (size_t i = 2; i < args.size(); ++i) {
+    members[i - 2] = ArgS(args, i);
+  }
+
+  auto cb = [&](Transaction* t, EngineShard* shard) {
+    return OpMScore(t->GetOpArgs(shard), key, members);
+  };
+
+  OpResult<MScoreResponse> result = cntx->transaction->ScheduleSingleHopT(std::move(cb));
+
+  if (result.status() == OpStatus::WRONG_TYPE) {
+    return (*cntx)->SendError(kWrongTypeErr);
+  }
+
+  (*cntx)->StartArray(result->size());
+  const MScoreResponse& array = result.value();
+  for (const auto& p : array) {
+    if (p) {
+      (*cntx)->SendDouble(*p);
+    } else {
+      (*cntx)->SendNull();
+    }
+  }
+}
+
 void ZSetFamily::ZScan(CmdArgList args, ConnectionContext* cntx) {
   string_view key = ArgS(args, 1);
   string_view token = ArgS(args, 2);
@@ -1603,6 +1632,32 @@ OpResult<double> ZSetFamily::OpScore(const OpArgs& op_args, string_view key, str
   return score;
 }
 
+OpResult<ZSetFamily::MScoreResponse> ZSetFamily::OpMScore(const OpArgs& op_args, string_view key, ArgSlice members) {
+  OpResult<PrimeIterator> res_it = op_args.shard->db_slice().Find(op_args.db_cntx, key, OBJ_ZSET);
+  if (!res_it)
+    return res_it.status();
+
+  MScoreResponse scores(members.size());
+
+  robj* zobj = res_it.value()->second.AsRObj();
+  sds& tmp_str = op_args.shard->tmp_str1;
+
+  for (size_t i = 0; i < members.size(); i++) {
+    const auto& m = members[i];
+
+    tmp_str = sdscpylen(tmp_str, m.data(), m.size());
+    double score;
+    int retval = zsetScore(zobj, tmp_str, &score);
+    if (retval == C_OK) {
+      scores[i] = score;
+    } else {
+      scores[i] = std::nullopt;
+    }
+  }
+
+  return scores;
+}
+
 auto ZSetFamily::OpRange(const ZRangeSpec& range_spec, const OpArgs& op_args, string_view key)
     -> OpResult<ScoredArray> {
   OpResult<PrimeIterator> res_it = op_args.shard->db_slice().Find(op_args.db_cntx, key, OBJ_ZSET);
@@ -1808,6 +1863,7 @@ void ZSetFamily::Register(CommandRegistry* registry) {
             << CI{"ZRANGEBYLEX", CO::READONLY, -4, 1, 1, 1}.HFUNC(ZRangeByLex)
             << CI{"ZRANGEBYSCORE", CO::READONLY, -4, 1, 1, 1}.HFUNC(ZRangeByScore)
             << CI{"ZSCORE", CO::READONLY | CO::FAST, 3, 1, 1, 1}.HFUNC(ZScore)
+            << CI{"ZMSCORE", CO::READONLY | CO::FAST, -3, 1, 1, 1}.HFUNC(ZMScore)
             << CI{"ZREMRANGEBYRANK", CO::WRITE, 4, 1, 1, 1}.HFUNC(ZRemRangeByRank)
             << CI{"ZREMRANGEBYSCORE", CO::WRITE, 4, 1, 1, 1}.HFUNC(ZRemRangeByScore)
             << CI{"ZREMRANGEBYLEX", CO::WRITE, 4, 1, 1, 1}.HFUNC(ZRemRangeByLex)

src/server/zset_family.h

@@ -62,6 +62,7 @@ class ZSetFamily {
   static void ZRank(CmdArgList args, ConnectionContext* cntx);
   static void ZRem(CmdArgList args, ConnectionContext* cntx);
   static void ZScore(CmdArgList args, ConnectionContext* cntx);
+  static void ZMScore(CmdArgList args, ConnectionContext* cntx);
   static void ZRangeByLex(CmdArgList args, ConnectionContext* cntx);
   static void ZRangeByScore(CmdArgList args, ConnectionContext* cntx);
   static void ZRemRangeByRank(CmdArgList args, ConnectionContext* cntx);
@@ -89,6 +90,9 @@ class ZSetFamily {
   static OpResult<unsigned> OpRem(const OpArgs& op_args, std::string_view key, ArgSlice members);
   static OpResult<double> OpScore(const OpArgs& op_args, std::string_view key,
                                   std::string_view member);
+  using MScoreResponse = std::vector<std::optional<double>>;
+  static OpResult<MScoreResponse> OpMScore(const OpArgs& op_args, std::string_view key,
+                                  ArgSlice members);
   static OpResult<ScoredArray> OpRange(const ZRangeSpec& range_spec, const OpArgs& op_args,
                                        std::string_view key);
   static OpResult<unsigned> OpRemRange(const OpArgs& op_args, std::string_view key,

src/server/zset_family_test.cc

@@ -64,6 +64,15 @@ TEST_F(ZSetFamilyTest, ZRem) {
   EXPECT_THAT(Run({"zrange", "x", "(-inf", "(+inf", "byscore"}), "a");
 }
 
+TEST_F(ZSetFamilyTest, ZMScore) {
+  Run({"zadd", "zms", "3.14", "a"});
+  Run({"zadd", "zms", "42", "another"});
+
+  auto resp = Run({"zmscore", "zms", "another", "a", "nofield"});
+  ASSERT_EQ(RespExpr::ARRAY, resp.type);
+  EXPECT_THAT(resp.GetVec(), ElementsAre("42", "3.14", ArgType(RespExpr::NIL)));
+}
+
 TEST_F(ZSetFamilyTest, ZRangeRank) {
   Run({"zadd", "x", "1.1", "a", "2.1", "b"});
   EXPECT_THAT(Run({"zrangebyscore", "x", "0", "(1.1"}), ArrLen(0));