diff --git a/CMakeLists.txt b/CMakeLists.txt index 88835d050..5d7dae4fe 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -29,6 +29,12 @@ add_third_party( /luaconf.h ${THIRD_PARTY_LIB_DIR}/lua/include ) +add_third_party( + dconv + URL https://github.com/google/double-conversion/archive/refs/tags/v3.2.0.tar.gz + LIB libdouble-conversion.a +) + Message(STATUS "THIRD_PARTY_LIB_DIR ${THIRD_PARTY_LIB_DIR}") include_directories(src) diff --git a/README.md b/README.md index 2667bb0d0..7dab7563f 100644 --- a/README.md +++ b/README.md @@ -205,15 +205,15 @@ API 2.0 - [X] Set Family - [X] SSCAN - [X] Sorted Set Family - - [ ] ZCOUNT + - [X] ZCOUNT - [ ] ZINTERSTORE - [ ] ZLEXCOUNT - [ ] ZRANGEBYLEX - - [ ] ZRANK + - [X] ZRANK - [ ] ZREMRANGEBYLEX - [X] ZREMRANGEBYRANK - [ ] ZREVRANGEBYSCORE - - [ ] ZREVRANK + - [X] ZREVRANK - [ ] ZUNIONSTORE - [ ] ZSCAN - [ ] HYPERLOGLOG Family diff --git a/src/facade/CMakeLists.txt b/src/facade/CMakeLists.txt index 48a93b60c..1f0e86760 100644 --- a/src/facade/CMakeLists.txt +++ b/src/facade/CMakeLists.txt @@ -1,7 +1,7 @@ add_library(dfly_facade dragonfly_listener.cc dragonfly_connection.cc facade.cc memcache_parser.cc redis_parser.cc reply_builder.cc) cxx_link(dfly_facade base uring_fiber_lib fibers_ext strings_lib http_server_lib - tls_lib TRDP::mimalloc) + tls_lib TRDP::mimalloc TRDP::dconv) add_library(facade_test facade_test.cc) cxx_link(facade_test dfly_facade gtest_main_ext) diff --git a/src/facade/reply_builder.cc b/src/facade/reply_builder.cc index 2490d11f6..246909956 100644 --- a/src/facade/reply_builder.cc +++ b/src/facade/reply_builder.cc @@ -5,12 +5,14 @@ #include #include +#include #include "base/logging.h" #include "facade/error.h" using namespace std; using absl::StrAppend; +using namespace double_conversion; namespace facade { @@ -100,7 +102,6 @@ void MCReplyBuilder::SendSimpleString(std::string_view str) { Send(v, ABSL_ARRAYSIZE(v)); } - void MCReplyBuilder::SendStored() { SendSimpleString("STORED"); } @@ -116,8 +117,7 @@ void MCReplyBuilder::SendMGetResponse(const OptResp* resp, uint32_t count) { for (unsigned i = 0; i < count; ++i) { if (resp[i]) { const auto& src = *resp[i]; - absl::StrAppend(&header, "VALUE ", src.key, " ", src.mc_flag, " ", - src.value.size()); + absl::StrAppend(&header, "VALUE ", src.key, " ", src.mc_flag, " ", src.value.size()); if (src.mc_ver) { absl::StrAppend(&header, " ", src.mc_ver); } @@ -224,7 +224,11 @@ void RedisReplyBuilder::SendLong(long num) { } void RedisReplyBuilder::SendDouble(double val) { - SendBulkString(absl::StrCat(val)); + char buf[64]; + StringBuilder sb(buf, sizeof(buf)); + CHECK(DoubleToStringConverter::EcmaScriptConverter().ToShortest(val, &sb)); + + SendBulkString(sb.Finalize()); } void RedisReplyBuilder::SendMGetResponse(const OptResp* resp, uint32_t count) { @@ -286,4 +290,4 @@ void ReqSerializer::SendCommand(std::string_view str) { ec_ = sink_->Write(v, ABSL_ARRAYSIZE(v)); } -} // namespace dfly +} // namespace facade diff --git a/src/redis/zset.h b/src/redis/zset.h index 2a6ae5e1f..4c534364d 100644 --- a/src/redis/zset.h +++ b/src/redis/zset.h @@ -75,7 +75,7 @@ unsigned long zsetLength(const robj* zobj); void zsetConvert(robj* zobj, int encoding); void zsetConvertToZiplistIfNeeded(robj* zobj, size_t maxelelen); int zsetScore(robj* zobj, sds member, double* score); -// unsigned long zslGetRank(zskiplist *zsl, double score, sds o); +unsigned long zslGetRank(zskiplist *zsl, double score, sds ele); int zsetAdd(robj* zobj, double score, sds ele, int in_flags, int* out_flags, double* newscore); long zsetRank(robj* zobj, sds ele, int reverse); int zsetDel(robj* zobj, sds ele); diff --git a/src/server/zset_family.cc b/src/server/zset_family.cc index 5f09e3bde..8d31d1f15 100644 --- a/src/server/zset_family.cc +++ b/src/server/zset_family.cc @@ -7,6 +7,7 @@ extern "C" { #include "redis/listpack.h" #include "redis/object.h" +#include "redis/util.h" #include "redis/zset.h" } @@ -28,10 +29,22 @@ using CI = CommandId; static const char kNxXxErr[] = "XX and NX options at the same time are not compatible"; static const char kScoreNaN[] = "resulting score is not a number (NaN)"; +static const char kRangeErr[] = "min or max is not a float"; + constexpr unsigned kMaxListPackValue = 64; +inline zrangespec GetZrangeSpec(const ZSetFamily::ScoreInterval& si) { + zrangespec range; + range.min = si.first.val; + range.max = si.second.val; + range.minex = si.first.is_open; + range.maxex = si.second.is_open; + + return range; +} + OpResult FindZEntry(unsigned flags, const OpArgs& op_args, string_view key, - size_t member_len) { + size_t member_len) { auto& db_slice = op_args.shard->db_slice(); if (flags & ZADD_IN_XX) { return db_slice.Find(op_args.db_ind, key, OBJ_ZSET); @@ -138,11 +151,7 @@ void IntervalVisitor::operator()(const ZSetFamily::IndexInterval& ii) { } void IntervalVisitor::operator()(const ZSetFamily::ScoreInterval& si) { - zrangespec range; - range.min = si.first.val; - range.max = si.second.val; - range.minex = si.first.is_open; - range.maxex = si.second.is_open; + zrangespec range = GetZrangeSpec(si); switch (action_) { case Action::RANGE: @@ -158,9 +167,9 @@ void IntervalVisitor::ActionRange(unsigned start, unsigned end) { unsigned rangelen = (end - start) + 1; if (zobj_->encoding == OBJ_ENCODING_LISTPACK) { - unsigned char* zl = (uint8_t*)zobj_->ptr; - unsigned char *eptr, *sptr; - unsigned char* vstr; + uint8_t* zl = (uint8_t*)zobj_->ptr; + uint8_t *eptr, *sptr; + uint8_t* vstr; unsigned int vlen; long long vlong; double score = 0.0; @@ -349,12 +358,15 @@ void IntervalVisitor::AddResult(const uint8_t* vstr, unsigned vlen, long long vl } bool ParseScore(string_view src, double* score) { + if (src.empty()) + return false; + if (src == "-inf") { *score = -HUGE_VAL; } else if (src == "+inf") { *score = HUGE_VAL; } else { - return absl::SimpleAtod(src, score); + return string2d(src.data(), src.size(), score); } return true; }; @@ -373,27 +385,6 @@ bool ParseBound(string_view src, ZSetFamily::Bound* bound) { } // namespace -void ZSetFamily::ZCard(CmdArgList args, ConnectionContext* cntx) { - string_view key = ArgS(args, 1); - - auto cb = [&](Transaction* t, EngineShard* shard) -> OpResult { - OpResult find_res = shard->db_slice().Find(t->db_index(), key, OBJ_ZSET); - if (!find_res) { - return find_res.status(); - } - - return zsetLength(find_res.value()->second.AsRObj()); - }; - - OpResult result = cntx->transaction->ScheduleSingleHopT(std::move(cb)); - if (result.status() == OpStatus::WRONG_TYPE) { - (*cntx)->SendError(kWrongTypeErr); - return; - } - - (*cntx)->SendLong(result.value()); -} - void ZSetFamily::ZAdd(CmdArgList args, ConnectionContext* cntx) { string_view key = ArgS(args, 1); @@ -437,7 +428,9 @@ void ZSetFamily::ZAdd(CmdArgList args, ConnectionContext* cntx) { return; } - if ((zparams.flags & ZADD_IN_NX) && (zparams.flags & (ZADD_IN_GT | ZADD_IN_LT))) { + constexpr auto kRangeOpt = ZADD_IN_GT | ZADD_IN_LT; + if (((zparams.flags & ZADD_IN_NX) && (zparams.flags & kRangeOpt)) || + ((zparams.flags & kRangeOpt) == kRangeOpt)) { (*cntx)->SendError("GT, LT, and/or NX options at the same time are not compatible"); return; } @@ -445,7 +438,8 @@ void ZSetFamily::ZAdd(CmdArgList args, ConnectionContext* cntx) { absl::InlinedVector members; for (; i < args.size(); i += 2) { string_view cur_arg = ArgS(args, i); - double val; + double val = 0; + if (!ParseScore(cur_arg, &val)) { return (*cntx)->SendError(kInvalidFloatErr); } @@ -471,18 +465,66 @@ void ZSetFamily::ZAdd(CmdArgList args, ConnectionContext* cntx) { } // KEY_NOTFOUND may happen in case of XX flag. - if (status == OpStatus::SKIPPED || status == OpStatus::KEY_NOTFOUND) { - return (*cntx)->SendNull(); - } - - if (add_result.is_nan) { - return (*cntx)->SendError(kScoreNaN); - } - - if (zparams.flags & ZADD_IN_INCR) { - (*cntx)->SendDouble(add_result.new_score); + if (status == OpStatus::KEY_NOTFOUND) { + if (zparams.flags & ZADD_IN_INCR) + (*cntx)->SendNull(); + else + (*cntx)->SendLong(0); + } else if (status == OpStatus::SKIPPED) { + (*cntx)->SendNull(); + } else if (add_result.is_nan) { + (*cntx)->SendError(kScoreNaN); } else { - (*cntx)->SendLong(add_result.num_updated); + if (zparams.flags & ZADD_IN_INCR) { + (*cntx)->SendDouble(add_result.new_score); + } else { + (*cntx)->SendLong(add_result.num_updated); + } + } +} + +void ZSetFamily::ZCard(CmdArgList args, ConnectionContext* cntx) { + string_view key = ArgS(args, 1); + + auto cb = [&](Transaction* t, EngineShard* shard) -> OpResult { + OpResult find_res = shard->db_slice().Find(t->db_index(), key, OBJ_ZSET); + if (!find_res) { + return find_res.status(); + } + + return zsetLength(find_res.value()->second.AsRObj()); + }; + + OpResult result = cntx->transaction->ScheduleSingleHopT(std::move(cb)); + if (result.status() == OpStatus::WRONG_TYPE) { + (*cntx)->SendError(kWrongTypeErr); + return; + } + + (*cntx)->SendLong(result.value()); +} + +void ZSetFamily::ZCount(CmdArgList args, ConnectionContext* cntx) { + string_view key = ArgS(args, 1); + + string_view min_s = ArgS(args, 2); + string_view max_s = ArgS(args, 3); + + ScoreInterval si; + if (!ParseBound(min_s, &si.first) || !ParseBound(max_s, &si.second)) { + return (*cntx)->SendError(kRangeErr); + } + + auto cb = [&](Transaction* t, EngineShard* shard) { + OpArgs op_args{shard, t->db_index()}; + return OpCount(op_args, key, si); + }; + + OpResult result = cntx->transaction->ScheduleSingleHopT(std::move(cb)); + if (result.status() == OpStatus::WRONG_TYPE) { + (*cntx)->SendError(kWrongTypeErr); + } else { + (*cntx)->SendLong(*result); } } @@ -531,10 +573,18 @@ void ZSetFamily::ZRange(CmdArgList args, ConnectionContext* cntx) { ZRangeGeneric(std::move(args), false, cntx); } +void ZSetFamily::ZRank(CmdArgList args, ConnectionContext* cntx) { + ZRankGeneric(std::move(args), false, cntx); +} + void ZSetFamily::ZRevRange(CmdArgList args, ConnectionContext* cntx) { ZRangeGeneric(std::move(args), true, cntx); } +void ZSetFamily::ZRevRank(CmdArgList args, ConnectionContext* cntx) { + ZRankGeneric(std::move(args), true, cntx); +} + void ZSetFamily::ZRangeByScore(CmdArgList args, ConnectionContext* cntx) { string_view key = ArgS(args, 1); string_view min_s = ArgS(args, 2); @@ -578,7 +628,7 @@ void ZSetFamily::ZRemRangeByScore(CmdArgList args, ConnectionContext* cntx) { ScoreInterval si; if (!ParseBound(min_s, &si.first) || !ParseBound(max_s, &si.second)) { - return (*cntx)->SendError("min or max is not a float"); + return (*cntx)->SendError(kRangeErr); } ZRangeSpec range_spec; @@ -635,7 +685,7 @@ void ZSetFamily::ZRangeByScoreInternal(string_view key, string_view min_s, strin ScoreInterval si; if (!ParseBound(min_s, &si.first) || !ParseBound(max_s, &si.second)) { - return (*cntx)->SendError("min or max is not a float"); + return (*cntx)->SendError(kRangeErr); } range_spec.interval = si; @@ -730,6 +780,26 @@ void ZSetFamily::ZRangeGeneric(CmdArgList args, bool reverse, ConnectionContext* OutputScoredArrayResult(result, range_params, cntx); } +void ZSetFamily::ZRankGeneric(CmdArgList args, bool reverse, ConnectionContext* cntx) { + string_view key = ArgS(args, 1); + string_view member = ArgS(args, 2); + + auto cb = [&](Transaction* t, EngineShard* shard) { + OpArgs op_args{shard, t->db_index()}; + + return OpRank(op_args, key, member, reverse); + }; + + OpResult result = cntx->transaction->ScheduleSingleHopT(std::move(cb)); + if (result) { + (*cntx)->SendLong(*result); + } else if (result.status() == OpStatus::KEY_NOTFOUND) { + (*cntx)->SendNull(); + } else { + (*cntx)->SendError(result.status()); + } +} + OpStatus ZSetFamily::OpAdd(const ZParams& zparams, const OpArgs& op_args, string_view key, ScoredMemberSpan members, AddResult* add_result) { DCHECK(!members.empty()); @@ -859,19 +929,108 @@ OpResult ZSetFamily::OpRemRange(const OpArgs& op_args, string_view key return iv.removed(); } +OpResult ZSetFamily::OpRank(const OpArgs& op_args, string_view key, string_view member, + bool reverse) { + 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(); + op_args.shard->tmp_str1 = sdscpylen(op_args.shard->tmp_str1, member.data(), member.size()); + + long res = zsetRank(zobj, op_args.shard->tmp_str1, reverse); + if (res < 0) + return OpStatus::KEY_NOTFOUND; + return res; +} + +OpResult ZSetFamily::OpCount(const OpArgs& op_args, std::string_view key, + const ScoreInterval& interval) { + 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(); + zrangespec range = GetZrangeSpec(interval); + unsigned count = 0; + + if (zobj->encoding == OBJ_ENCODING_LISTPACK) { + uint8_t* zl = (uint8_t*)zobj->ptr; + uint8_t *eptr, *sptr; + double score; + + /* Use the first element in range as the starting point */ + eptr = zzlFirstInRange(zl, &range); + + /* No "first" element */ + if (eptr == NULL) { + return 0; + } + + /* First element is in range */ + sptr = lpNext(zl, eptr); + score = zzlGetScore(sptr); + + DCHECK(zslValueLteMax(score, &range)); + + /* Iterate over elements in range */ + while (eptr) { + score = zzlGetScore(sptr); + + /* Abort when the node is no longer in range. */ + if (!zslValueLteMax(score, &range)) { + break; + } else { + count++; + zzlNext(zl, &eptr, &sptr); + } + } + } else { + CHECK_EQ(unsigned(OBJ_ENCODING_SKIPLIST), zobj->encoding); + zset* zs = (zset*)zobj->ptr; + zskiplist* zsl = zs->zsl; + zskiplistNode* zn; + unsigned long rank; + + /* Find first element in range */ + zn = zslFirstInRange(zsl, &range); + + /* Use rank of first element, if any, to determine preliminary count */ + if (zn == NULL) + return 0; + + rank = zslGetRank(zsl, zn->score, zn->ele); + count = (zsl->length - (rank - 1)); + + /* Find last element in range */ + zn = zslLastInRange(zsl, &range); + + /* Use rank of last element, if any, to determine the actual count */ + if (zn != NULL) { + rank = zslGetRank(zsl, zn->score, zn->ele); + count -= (zsl->length - rank); + } + } + + return count; +} + #define HFUNC(x) SetHandler(&ZSetFamily::x) void ZSetFamily::Register(CommandRegistry* registry) { - *registry << CI{"ZCARD", CO::FAST | CO::READONLY, 2, 1, 1, 1}.HFUNC(ZCard) - << CI{"ZADD", CO::FAST | CO::WRITE | CO::DENYOOM, -4, 1, 1, 1}.HFUNC(ZAdd) + *registry << CI{"ZADD", CO::FAST | CO::WRITE | CO::DENYOOM, -4, 1, 1, 1}.HFUNC(ZAdd) + << CI{"ZCARD", CO::FAST | CO::READONLY, 2, 1, 1, 1}.HFUNC(ZCard) + << CI{"ZCOUNT", CO::FAST | CO::READONLY, 4, 1, 1, 1}.HFUNC(ZCount) << CI{"ZINCRBY", CO::FAST | CO::WRITE | CO::DENYOOM, 4, 1, 1, 1}.HFUNC(ZIncrBy) << CI{"ZREM", CO::FAST | CO::WRITE, -3, 1, 1, 1}.HFUNC(ZRem) << CI{"ZRANGE", CO::READONLY, -4, 1, 1, 1}.HFUNC(ZRange) + << CI{"ZRANK", CO::READONLY | CO::FAST, 3, 1, 1, 1}.HFUNC(ZRank) << CI{"ZRANGEBYSCORE", CO::READONLY, -4, 1, 1, 1}.HFUNC(ZRangeByScore) << CI{"ZSCORE", CO::READONLY | CO::FAST, 3, 1, 1, 1}.HFUNC(ZScore) << CI{"ZREMRANGEBYRANK", CO::WRITE, 4, 1, 1, 1}.HFUNC(ZRemRangeByRank) << CI{"ZREMRANGEBYSCORE", CO::WRITE, 4, 1, 1, 1}.HFUNC(ZRemRangeByScore) - << CI{"ZREVRANGE", CO::WRITE, 4, 1, 1, 1}.HFUNC(ZRevRange); + << CI{"ZREVRANGE", CO::WRITE, 4, 1, 1, 1}.HFUNC(ZRevRange) + << CI{"ZREVRANK", CO::READONLY | CO::FAST, 3, 1, 1, 1}.HFUNC(ZRevRank); } } // namespace dfly diff --git a/src/server/zset_family.h b/src/server/zset_family.h index ff114de59..c894dbee6 100644 --- a/src/server/zset_family.h +++ b/src/server/zset_family.h @@ -45,16 +45,19 @@ class ZSetFamily { private: template using OpResult = facade::OpResult; - static void ZCard(CmdArgList args, ConnectionContext* cntx); static void ZAdd(CmdArgList args, ConnectionContext* cntx); + static void ZCard(CmdArgList args, ConnectionContext* cntx); + static void ZCount(CmdArgList args, ConnectionContext* cntx); static void ZIncrBy(CmdArgList args, ConnectionContext* cntx); static void ZRange(CmdArgList args, ConnectionContext* cntx); + static void ZRank(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 ZRemRangeByRank(CmdArgList args, ConnectionContext* cntx); static void ZRemRangeByScore(CmdArgList args, ConnectionContext* cntx); static void ZRevRange(CmdArgList args, ConnectionContext* cntx); + static void ZRevRank(CmdArgList args, ConnectionContext* cntx); static void ZRangeByScoreInternal(std::string_view key, std::string_view min_s, std::string_view max_s, const RangeParams& params, @@ -64,6 +67,7 @@ class ZSetFamily { static void ZRemRangeGeneric(std::string_view key, const ZRangeSpec& range_spec, ConnectionContext* cntx); static void ZRangeGeneric(CmdArgList args, bool reverse, ConnectionContext* cntx); + static void ZRankGeneric(CmdArgList args, bool reverse, ConnectionContext* cntx); struct ZParams { unsigned flags = 0; // mask of ZADD_IN_ macros. @@ -89,6 +93,12 @@ class ZSetFamily { std::string_view key); static OpResult OpRemRange(const OpArgs& op_args, std::string_view key, const ZRangeSpec& spec); + + static OpResult OpRank(const OpArgs& op_args, std::string_view key, + std::string_view member, bool reverse); + + static OpResult OpCount(const OpArgs& op_args, std::string_view key, + const ScoreInterval& interval); }; } // namespace dfly diff --git a/src/server/zset_family_test.cc b/src/server/zset_family_test.cc index eb872a174..c23596481 100644 --- a/src/server/zset_family_test.cc +++ b/src/server/zset_family_test.cc @@ -40,6 +40,16 @@ TEST_F(ZSetFamilyTest, Add) { resp = Run({"zcard", "x"}); EXPECT_THAT(resp[0], IntArg(1)); + + EXPECT_THAT(Run({"zadd", "x", "", "a"}), ElementsAre(ErrArg("not a valid float"))); + + EXPECT_THAT(Run({"zadd", "ztmp", "xx", "10", "member"}), ElementsAre(IntArg(0))); + + const char kHighPrecision[] = "0.79028573343077946"; + + Run({"zadd", "zs", kHighPrecision, "a"}); + EXPECT_THAT(Run({"zscore", "zs", "a"}), ElementsAre("0.7902857334307795")); + EXPECT_EQ(0.79028573343077946, 0.7902857334307795); } TEST_F(ZSetFamilyTest, ZRem) { @@ -55,10 +65,21 @@ TEST_F(ZSetFamilyTest, ZRem) { EXPECT_THAT(Run({"zrange", "x", "(-inf", "(+inf", "byscore"}), ElementsAre("a")); } -TEST_F(ZSetFamilyTest, ZRange) { +TEST_F(ZSetFamilyTest, ZRangeRank) { Run({"zadd", "x", "1.1", "a", "2.1", "b"}); EXPECT_THAT(Run({"zrangebyscore", "x", "0", "(1.1"}), ElementsAre(ArrLen(0))); EXPECT_THAT(Run({"zrangebyscore", "x", "-inf", "1.1"}), ElementsAre("a")); + + EXPECT_EQ(2, CheckedInt({"zcount", "x", "1.1", "2.1"})); + EXPECT_EQ(1, CheckedInt({"zcount", "x", "(1.1", "2.1"})); + EXPECT_EQ(0, CheckedInt({"zcount", "y", "(1.1", "2.1"})); + + EXPECT_EQ(0, CheckedInt({"zrank", "x", "a"})); + EXPECT_EQ(1, CheckedInt({"zrank", "x", "b"})); + EXPECT_EQ(1, CheckedInt({"zrevrank", "x", "a"})); + EXPECT_EQ(0, CheckedInt({"zrevrank", "x", "b"})); + EXPECT_THAT(Run({"zrevrank", "x", "c"}), ElementsAre(ArgType(RespExpr::NIL))); + EXPECT_THAT(Run({"zrank", "y", "c"}), ElementsAre(ArgType(RespExpr::NIL))); } TEST_F(ZSetFamilyTest, ZRemRangeRank) {