Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat:zmax && zmin #2966

Merged
merged 2 commits into from
Dec 11, 2024
Merged

feat:zmax && zmin #2966

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8397ae9
zmax && zmin
brother-jin Dec 5, 2024
1e7ce6e
zopmax && zpopmin
brother-jin Dec 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions include/pika_cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,12 +166,17 @@ class PikaCache : public pstd::noncopyable, public std::enable_shared_from_this<
const std::shared_ptr<DB>& db);
rocksdb::Status ZRevrangebylex(std::string& key, std::string& min, std::string& max, std::vector<std::string>* members,
const std::shared_ptr<DB>& db);
rocksdb::Status ZRevrank(std::string& key, std::string& member, int64_t *rank, const std::shared_ptr<DB>& db);
rocksdb::Status ZRevrank(std::string& key, std::string& member, int64_t* rank, const std::shared_ptr<DB>& db);
rocksdb::Status ZScore(std::string& key, std::string& member, double* score, const std::shared_ptr<DB>& db);
rocksdb::Status ZRangebylex(std::string& key, std::string& min, std::string& max, std::vector<std::string>* members, const std::shared_ptr<DB>& db);
rocksdb::Status ZRangebylex(std::string& key, std::string& min, std::string& max, std::vector<std::string>* members,
const std::shared_ptr<DB>& db);
rocksdb::Status ZLexcount(std::string& key, std::string& min, std::string& max, uint64_t* len,
const std::shared_ptr<DB>& db);
rocksdb::Status ZRemrangebylex(std::string& key, std::string& min, std::string& max, const std::shared_ptr<DB>& db);
rocksdb::Status ZPopMin(std::string& key, int64_t count, std::vector<storage::ScoreMember>* score_members,
const std::shared_ptr<DB>& db);
rocksdb::Status ZPopMax(std::string& key, int64_t count, std::vector<storage::ScoreMember>* score_members,
const std::shared_ptr<DB>& db);

// Bit Commands
rocksdb::Status SetBit(std::string& key, size_t offset, int64_t value);
Expand Down
4 changes: 4 additions & 0 deletions include/pika_zset.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -603,6 +603,8 @@ class ZPopmaxCmd : public Cmd {
void Do() override;
void Split(const HintKeys& hint_keys) override {};
void Merge() override {};
void DoThroughDB() override;
void DoUpdateCache() override;
Cmd* Clone() override { return new ZPopmaxCmd(*this); }

private:
Expand All @@ -623,6 +625,8 @@ class ZPopminCmd : public Cmd {
void Do() override;
void Split(const HintKeys& hint_keys) override {};
void Merge() override {};
void DoThroughDB() override;
void DoUpdateCache() override;
Cmd* Clone() override { return new ZPopminCmd(*this); }

private:
Expand Down
78 changes: 78 additions & 0 deletions src/cache/src/zset.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -405,5 +405,83 @@ Status RedisCache::ZRemrangebylex(std::string& key, std::string &min, std::strin
return Status::OK();
}


Status RedisCache::ZPopMin(std::string& key, int64_t count, std::vector<storage::ScoreMember>* score_members) {
zitem* items = nullptr;
unsigned long items_size = 0;
robj* kobj = createObject(OBJ_STRING, sdsnewlen(key.data(), key.size()));
DEFER {
DecrObjectsRefCount(kobj);
};

int ret = RcZrange(cache_, kobj, 0, -1, &items, &items_size);
if (C_OK != ret) {
if (REDIS_KEY_NOT_EXIST == ret) {
return Status::NotFound("key not in cache");
}
return Status::Corruption("RcZrange failed");
}

unsigned long to_return = std::min(static_cast<unsigned long>(count), items_size);
for (unsigned long i = 0; i < to_return; ++i) {
storage::ScoreMember sm;
sm.score = items[i].score;
sm.member.assign(items[i].member, sdslen(items[i].member));
score_members->push_back(sm);
}

robj** members_obj = (robj**)zcallocate(sizeof(robj*) * items_size);
for (unsigned long i = 0; i < items_size; ++i) {
members_obj[i] = createObject(OBJ_STRING, sdsnewlen(items[i].member, sdslen(items[i].member)));
}
DEFER {
FreeObjectList(members_obj, items_size);
};

RcZRem(cache_, kobj, members_obj, to_return);

FreeZitemList(items, items_size);
return Status::OK();
}
Comment on lines +409 to +445
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Performance concern: Consider optimizing range retrieval

The implementation fetches the entire sorted set using RcZrange(0, -1) even when only a small count is needed. This could be inefficient for large sets.

Consider implementing a more efficient approach:

-  int ret = RcZrange(cache_, kobj, 0, -1, &items, &items_size);
+  int ret = RcZrange(cache_, kobj, 0, count - 1, &items, &items_size);

Committable suggestion skipped: line range outside the PR's diff.


Status RedisCache::ZPopMax(std::string& key, int64_t count, std::vector<storage::ScoreMember>* score_members) {
zitem* items = nullptr;
unsigned long items_size = 0;
robj* kobj = createObject(OBJ_STRING, sdsnewlen(key.data(), key.size()));
DEFER {
DecrObjectsRefCount(kobj);
};

int ret = RcZrange(cache_, kobj, 0, -1, &items, &items_size);
if (C_OK != ret) {
if (REDIS_KEY_NOT_EXIST == ret) {
return Status::NotFound("key not in cache");
}
return Status::Corruption("RcZrange failed");
}

unsigned long to_return = std::min(static_cast<unsigned long>(count), items_size);
for (unsigned long i = items_size - to_return; i < items_size; ++i) {
storage::ScoreMember sm;
sm.score = items[i].score;
sm.member.assign(items[i].member, sdslen(items[i].member));
score_members->push_back(sm);
}

robj** members_obj = (robj**)zcallocate(sizeof(robj*) * items_size);
for (unsigned long i = items_size - 1; i >= 0; --i) {
members_obj[items_size - 1 - i] = createObject(OBJ_STRING, sdsnewlen(items[i].member, sdslen(items[i].member)));
}

DEFER {
FreeObjectList(members_obj, items_size);
};

RcZRem(cache_, kobj, members_obj, to_return);

FreeZitemList(items, items_size);
return Status::OK();
}
Comment on lines +447 to +484
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Critical: Potential integer underflow in ZPopMax

The loop for creating member objects in ZPopMax has a potential integer underflow issue when decrementing an unsigned long.

Fix the loop to prevent integer underflow:

-  for (unsigned long i = items_size - 1; i >= 0; --i) {
+  for (unsigned long i = items_size; i > 0; --i) {
+    unsigned long idx = i - 1;
-    members_obj[items_size - 1 - i] = createObject(OBJ_STRING, sdsnewlen(items[i].member, sdslen(items[i].member)));
+    members_obj[items_size - i] = createObject(OBJ_STRING, sdsnewlen(items[idx].member, sdslen(items[idx].member)));
   }

Additionally, consider the same performance optimization as suggested for ZPopMin:

-  int ret = RcZrange(cache_, kobj, 0, -1, &items, &items_size);
+  int ret = RcZrange(cache_, kobj, -count, -1, &items, &items_size);
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
Status RedisCache::ZPopMax(std::string& key, int64_t count, std::vector<storage::ScoreMember>* score_members) {
zitem* items = nullptr;
unsigned long items_size = 0;
robj* kobj = createObject(OBJ_STRING, sdsnewlen(key.data(), key.size()));
DEFER {
DecrObjectsRefCount(kobj);
};
int ret = RcZrange(cache_, kobj, 0, -1, &items, &items_size);
if (C_OK != ret) {
if (REDIS_KEY_NOT_EXIST == ret) {
return Status::NotFound("key not in cache");
}
return Status::Corruption("RcZrange failed");
}
unsigned long to_return = std::min(static_cast<unsigned long>(count), items_size);
for (unsigned long i = items_size - to_return; i < items_size; ++i) {
storage::ScoreMember sm;
sm.score = items[i].score;
sm.member.assign(items[i].member, sdslen(items[i].member));
score_members->push_back(sm);
}
robj** members_obj = (robj**)zcallocate(sizeof(robj*) * items_size);
for (unsigned long i = items_size - 1; i >= 0; --i) {
members_obj[items_size - 1 - i] = createObject(OBJ_STRING, sdsnewlen(items[i].member, sdslen(items[i].member)));
}
DEFER {
FreeObjectList(members_obj, items_size);
};
RcZRem(cache_, kobj, members_obj, to_return);
FreeZitemList(items, items_size);
return Status::OK();
}
Status RedisCache::ZPopMax(std::string& key, int64_t count, std::vector<storage::ScoreMember>* score_members) {
zitem* items = nullptr;
unsigned long items_size = 0;
robj* kobj = createObject(OBJ_STRING, sdsnewlen(key.data(), key.size()));
DEFER {
DecrObjectsRefCount(kobj);
};
int ret = RcZrange(cache_, kobj, -count, -1, &items, &items_size);
if (C_OK != ret) {
if (REDIS_KEY_NOT_EXIST == ret) {
return Status::NotFound("key not in cache");
}
return Status::Corruption("RcZrange failed");
}
unsigned long to_return = std::min(static_cast<unsigned long>(count), items_size);
for (unsigned long i = items_size - to_return; i < items_size; ++i) {
storage::ScoreMember sm;
sm.score = items[i].score;
sm.member.assign(items[i].member, sdslen(items[i].member));
score_members->push_back(sm);
}
robj** members_obj = (robj**)zcallocate(sizeof(robj*) * items_size);
for (unsigned long i = items_size; i > 0; --i) {
unsigned long idx = i - 1;
members_obj[items_size - i] = createObject(OBJ_STRING, sdsnewlen(items[idx].member, sdslen(items[idx].member)));
}
DEFER {
FreeObjectList(members_obj, items_size);
};
RcZRem(cache_, kobj, members_obj, to_return);
FreeZitemList(items, items_size);
return Status::OK();
}


} // namespace cache
/* EOF */
29 changes: 29 additions & 0 deletions src/pika_cache.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1465,6 +1465,35 @@ Status PikaCache::ZRemrangebylex(std::string& key, std::string &min, std::string
}
}

Status PikaCache::ZPopMin(std::string& key, int64_t count, std::vector<storage::ScoreMember>* score_members, const std::shared_ptr<DB>& db) {
int cache_index = CacheIndex(key);
std::lock_guard lm(*cache_mutexs_[cache_index]);

auto cache_obj = caches_[cache_index];
Status s;

if (cache_obj->Exists(key)) {
return cache_obj->ZPopMin(key, count, score_members);
} else {
return Status::NotFound("key not in cache");
}
}

Status PikaCache::ZPopMax(std::string& key, int64_t count, std::vector<storage::ScoreMember>* score_members, const std::shared_ptr<DB>& db) {
int cache_index = CacheIndex(key);
std::lock_guard lm(*cache_mutexs_[cache_index]);

auto cache_obj = caches_[cache_index];
Status s;

if (cache_obj->Exists(key)) {
return cache_obj->ZPopMax(key, count, score_members);
} else {
return Status::NotFound("key not in cache");
}
}


/*-----------------------------------------------------------------------------
* Bit Commands
*----------------------------------------------------------------------------*/
Expand Down
4 changes: 2 additions & 2 deletions src/pika_command.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -600,11 +600,11 @@ void InitCmdTable(CmdTable* cmd_table) {
cmd_table->insert(std::pair<std::string, std::unique_ptr<Cmd>>(kCmdNameZRemrangebylex, std::move(zremrangebylexptr)));
////ZPopmax
std::unique_ptr<Cmd> zpopmaxptr = std::make_unique<ZPopmaxCmd>(
kCmdNameZPopmax, -2, kCmdFlagsWrite | kCmdFlagsZset | kCmdFlagsFast);
kCmdNameZPopmax, -2, kCmdFlagsWrite | kCmdFlagsZset | kCmdFlagsFast | kCmdFlagsDoThroughDB | kCmdFlagsUpdateCache);
cmd_table->insert(std::pair<std::string, std::unique_ptr<Cmd>>(kCmdNameZPopmax, std::move(zpopmaxptr)));
////ZPopmin
std::unique_ptr<Cmd> zpopminptr = std::make_unique<ZPopminCmd>(
kCmdNameZPopmin, -2, kCmdFlagsWrite | kCmdFlagsZset | kCmdFlagsFast);
kCmdNameZPopmin, -2, kCmdFlagsWrite | kCmdFlagsZset | kCmdFlagsFast | kCmdFlagsDoThroughDB | kCmdFlagsUpdateCache);
cmd_table->insert(std::pair<std::string, std::unique_ptr<Cmd>>(kCmdNameZPopmin, std::move(zpopminptr)));

// Set
Expand Down
22 changes: 22 additions & 0 deletions src/pika_zset.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1507,6 +1507,17 @@ void ZPopmaxCmd::Do() {
}
}

void ZPopmaxCmd::DoThroughDB(){
Do();
}

void ZPopmaxCmd::DoUpdateCache(){
std::vector<storage::ScoreMember> score_members;
if(s_.ok() || s_.IsNotFound()){
db_->cache()->ZPopMax(key_, count_, &score_members, db_);
}
}

void ZPopminCmd::DoInitial() {
if (!CheckArg(argv_.size())) {
res_.SetRes(CmdRes::kWrongNum, kCmdNameZPopmin);
Expand All @@ -1523,6 +1534,17 @@ void ZPopminCmd::DoInitial() {
}
}

void ZPopminCmd::DoThroughDB(){
Do();
}

void ZPopminCmd::DoUpdateCache(){
std::vector<storage::ScoreMember> score_members;
if(s_.ok() || s_.IsNotFound()){
db_->cache()->ZPopMin(key_, count_, &score_members, db_);
}
}

void ZPopminCmd::Do() {
std::vector<storage::ScoreMember> score_members;
rocksdb::Status s = db_->storage()->ZPopMin(key_, count_, &score_members);
Expand Down
26 changes: 26 additions & 0 deletions tests/integration/zset_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1440,6 +1440,32 @@ var _ = Describe("Zset Commands", func() {
}}))
})

It("should ZRem multiple members", func() {

err := client.ZAdd(ctx, "zset", redis.Z{Score: 1, Member: "one"}).Err()
Expect(err).NotTo(HaveOccurred())
err = client.ZAdd(ctx, "zset", redis.Z{Score: 2, Member: "two"}).Err()
Expect(err).NotTo(HaveOccurred())
err = client.ZAdd(ctx, "zset", redis.Z{Score: 3, Member: "three"}).Err()
Expect(err).NotTo(HaveOccurred())
err = client.ZAdd(ctx, "zset", redis.Z{Score: 4, Member: "four"}).Err()
Expect(err).NotTo(HaveOccurred())

zRem := client.ZRem(ctx, "zset", "two", "four")
Expect(zRem.Err()).NotTo(HaveOccurred())
Expect(zRem.Val()).To(Equal(int64(2)))

vals, err := client.ZRangeWithScores(ctx, "zset", 0, -1).Result()
Expect(err).NotTo(HaveOccurred())
Expect(vals).To(Equal([]redis.Z{{
Score: 1,
Member: "one",
}, {
Score: 3,
Member: "three",
}}))
})

It("should ZRemRangeByRank", func() {
err := client.ZAdd(ctx, "zset", redis.Z{Score: 1, Member: "one"}).Err()
Expect(err).NotTo(HaveOccurred())
Expand Down
Loading