From f09f5166364bafa0006e1cd941621561f320fe94 Mon Sep 17 00:00:00 2001 From: Roman Gershman Date: Sun, 6 Mar 2022 08:46:48 +0200 Subject: [PATCH] Add zcard, zscore, zrem methods --- README.md | 10 ++--- src/server/zset_family.cc | 77 ++++++++++++++++++++++++++++++++-- src/server/zset_family.h | 20 +++++---- src/server/zset_family_test.cc | 21 ++++++++++ 4 files changed, 110 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index f317fbc..5096366 100644 --- a/README.md +++ b/README.md @@ -104,16 +104,16 @@ API 1.0 - [X] RPOP - [ ] RPOPLPUSH - [X] RPUSH -- [ ] SortedSet Family - - [ ] ZADD - - [ ] ZCARD +- [X] SortedSet Family + - [X] ZADD + - [X] ZCARD - [ ] ZINCRBY - [ ] ZRANGE - [ ] ZRANGEBYSCORE - - [ ] ZREM + - [X] ZREM - [ ] ZREMRANGEBYSCORE - [ ] ZREVRANGE - - [ ] ZSCORE + - [X] ZSCORE - [ ] Not sure whether these are required for the initial release. - [X] AUTH - [ ] BGREWRITEAOF diff --git a/src/server/zset_family.cc b/src/server/zset_family.cc index 0ece971..34b0516 100644 --- a/src/server/zset_family.cc +++ b/src/server/zset_family.cc @@ -28,7 +28,6 @@ using CI = CommandId; static const char kNxXxErr[] = "XX and NX options at the same time are not compatible"; constexpr unsigned kMaxZiplistValue = 64; - OpResult FindZEntry(unsigned flags, const OpArgs& op_args, string_view key, size_t member_len) { auto& db_slice = op_args.shard->db_slice(); @@ -168,15 +167,47 @@ void ZSetFamily::ZRangeByScore(CmdArgList args, ConnectionContext* cntx) { } void ZSetFamily::ZRem(CmdArgList args, ConnectionContext* cntx) { - (*cntx)->SendLong(0); + std::string_view key = ArgS(args, 1); + + absl::InlinedVector 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) { + OpArgs op_args{shard, t->db_index()}; + return OpRem(op_args, key, members); + }; + + OpResult result = cntx->transaction->ScheduleSingleHopT(std::move(cb)); + if (result.status() == OpStatus::WRONG_TYPE) { + (*cntx)->SendError(kWrongTypeErr); + } else { + (*cntx)->SendLong(result.value()); + } } void ZSetFamily::ZScore(CmdArgList args, ConnectionContext* cntx) { - (*cntx)->SendDouble(0); + std::string_view key = ArgS(args, 1); + std::string_view member = ArgS(args, 2); + + auto cb = [&](Transaction* t, EngineShard* shard) { + OpArgs op_args{shard, t->db_index()}; + return OpScore(op_args, key, member); + }; + + OpResult result = cntx->transaction->ScheduleSingleHopT(std::move(cb)); + if (result.status() == OpStatus::WRONG_TYPE) { + (*cntx)->SendError(kWrongTypeErr); + } else if (!result) { + (*cntx)->SendNull(); + } else { + (*cntx)->SendDouble(result.value()); + } } OpResult ZSetFamily::OpAdd(const ZParams& zparams, const OpArgs& op_args, string_view key, - const ScoredMemberSpan& members) { + ScoredMemberSpan members) { DCHECK(!members.empty()); OpResult res_it = FindZEntry(zparams.flags, op_args, key, members.front().second.size()); @@ -215,6 +246,44 @@ OpResult ZSetFamily::OpAdd(const ZParams& zparams, const OpArgs& op_ar return zparams.ch ? added + updated : added; } +OpResult ZSetFamily::OpRem(const OpArgs& op_args, string_view key, ArgSlice members) { + OpResult res_it = op_args.shard->db_slice().Find(op_args.db_ind, key, OBJ_ZSET); + if (!res_it) + return res_it.status(); + + robj* zobj = res_it.value()->second.AsRObj(); + sds& tmp_str = op_args.shard->tmp_str1; + unsigned deleted = 0; + for (string_view member : members) { + tmp_str = sdscpylen(tmp_str, member.data(), member.size()); + deleted += zsetDel(zobj, tmp_str); + } + auto zlen = zsetLength(zobj); + res_it.value()->second.SyncRObj(); + + if (zlen == 0) { + CHECK(op_args.shard->db_slice().Del(op_args.db_ind, res_it.value())); + } + + return deleted; +} + +OpResult ZSetFamily::OpScore(const OpArgs& op_args, string_view key, string_view member) { + OpResult res_it = op_args.shard->db_slice().Find(op_args.db_ind, key, OBJ_ZSET); + if (!res_it) + return res_it.status(); + + robj* zobj = res_it.value()->second.AsRObj(); + sds& tmp_str = op_args.shard->tmp_str1; + tmp_str = sdscpylen(tmp_str, member.data(), member.size()); + double score; + int retval = zsetScore(zobj, tmp_str, &score); + if (retval != C_OK) { + return OpStatus::KEY_NOTFOUND; + } + return score; +} + #define HFUNC(x) SetHandler(&ZSetFamily::x) void ZSetFamily::Register(CommandRegistry* registry) { diff --git a/src/server/zset_family.h b/src/server/zset_family.h index 90e22ac..cada498 100644 --- a/src/server/zset_family.h +++ b/src/server/zset_family.h @@ -19,13 +19,13 @@ class ZSetFamily { static void Register(CommandRegistry* registry); private: - static void ZCard(CmdArgList args, ConnectionContext* cntx); - static void ZAdd(CmdArgList args, ConnectionContext* cntx); - static void ZIncrBy(CmdArgList args, ConnectionContext* cntx); - static void ZRange(CmdArgList args, ConnectionContext* cntx); - static void ZRem(CmdArgList args, ConnectionContext* cntx); - static void ZScore(CmdArgList args, ConnectionContext* cntx); - static void ZRangeByScore(CmdArgList args, ConnectionContext* cntx); + static void ZCard(CmdArgList args, ConnectionContext* cntx); + static void ZAdd(CmdArgList args, ConnectionContext* cntx); + static void ZIncrBy(CmdArgList args, ConnectionContext* cntx); + static void ZRange(CmdArgList args, ConnectionContext* cntx); + static void ZRem(CmdArgList args, ConnectionContext* cntx); + static void ZScore(CmdArgList args, ConnectionContext* cntx); + static void ZRangeByScore(CmdArgList args, ConnectionContext* cntx); struct ZParams { unsigned flags = 0; // mask of ZADD_IN_ macros. @@ -37,8 +37,10 @@ class ZSetFamily { template using OpResult = facade::OpResult; static OpResult OpAdd(const ZParams& zparams, const OpArgs& op_args, - std::string_view key, const ScoredMemberSpan& members); - + std::string_view key, ScoredMemberSpan members); + static OpResult OpRem(const OpArgs& op_args, std::string_view key, ArgSlice members); + static OpResult OpScore(const OpArgs& op_args, std::string_view key, + std::string_view member); }; } // namespace dfly diff --git a/src/server/zset_family_test.cc b/src/server/zset_family_test.cc index 3b2a6fd..6a07384 100644 --- a/src/server/zset_family_test.cc +++ b/src/server/zset_family_test.cc @@ -24,11 +24,32 @@ class ZSetFamilyTest : public BaseFamilyTest { TEST_F(ZSetFamilyTest, Add) { auto resp = Run({"zadd", "x", "1.1", "a"}); EXPECT_THAT(resp[0], IntArg(1)); + + resp = Run({"zscore", "x", "a"}); + EXPECT_THAT(resp[0], StrArg("1.1")); + resp = Run({"zadd", "x", "2", "a"}); EXPECT_THAT(resp[0], IntArg(0)); + resp = Run({"zscore", "x", "a"}); + EXPECT_THAT(resp[0], StrArg("2")); + resp = Run({"zadd", "x", "ch", "3", "a"}); EXPECT_THAT(resp[0], IntArg(1)); + resp = Run({"zscore", "x", "a"}); + EXPECT_THAT(resp[0], StrArg("3")); + + resp = Run({"zcard", "x"}); + EXPECT_THAT(resp[0], IntArg(1)); } +TEST_F(ZSetFamilyTest, ZRem) { + auto resp = Run({"zadd", "x", "1.1", "b", "2.1", "a"}); + EXPECT_THAT(resp[0], IntArg(2)); + + resp = Run({"zrem", "x", "b", "c"}); + EXPECT_THAT(resp[0], IntArg(1)); + resp = Run({"zcard", "x"}); + EXPECT_THAT(resp[0], IntArg(1)); +} } // namespace dfly