Skip to content

Commit 06b15a4

Browse files
chejingebrother-jin
chejinge
and
brother-jin
authored
feat:mget supports cache querying multiple keys (OpenAtomFoundation#2675)
* multi get --------- Co-authored-by: chejinge <chejinge@360.cn>
1 parent 9fd3c51 commit 06b15a4

6 files changed

Lines changed: 253 additions & 40 deletions

File tree

‎include/pika_kv.h‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -268,14 +268,19 @@ class MgetCmd : public Cmd {
268268
void Merge() override;
269269
Cmd* Clone() override { return new MgetCmd(*this); }
270270

271+
private:
272+
void DoInitial() override;
273+
void MergeCachedAndDbResults();
274+
void AssembleResponseFromCache();
275+
271276
private:
272277
std::vector<std::string> keys_;
278+
std::vector<std::string> cache_miss_keys_;
273279
std::string value_;
280+
std::unordered_map<std::string, std::string> cache_hit_values_;
274281
std::vector<storage::ValueStatus> split_res_;
275282
std::vector<storage::ValueStatus> db_value_status_array_;
276283
std::vector<storage::ValueStatus> cache_value_status_array_;
277-
int64_t ttl_ = -1;
278-
void DoInitial() override;
279284
rocksdb::Status s_;
280285
};
281286

‎src/pika_bit.cc‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -164,12 +164,11 @@ void BitCountCmd::ReadCache() {
164164
int64_t count = 0;
165165
int64_t start = static_cast<long>(start_offset_);
166166
int64_t end = static_cast<long>(end_offset_);
167-
rocksdb::Status s;
167+
bool flag = true;
168168
if (count_all_) {
169-
s = db_->cache()->BitCount(key_, start, end, &count, 0);
170-
} else {
171-
s = db_->cache()->BitCount(key_, start, end, &count, 1);
169+
flag = false;
172170
}
171+
rocksdb::Status s = db_->cache()->BitCount(key_, start, end, &count, flag);
173172

174173
if (s.ok()) {
175174
res_.AppendInteger(count);

‎src/pika_command.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -678,15 +678,15 @@ void InitCmdTable(CmdTable* cmd_table) {
678678
cmd_table->insert(std::pair<std::string, std::unique_ptr<Cmd>>(kCmdNameBitSet, std::move(bitsetptr)));
679679
////bitgetCmd
680680
std::unique_ptr<Cmd> bitgetptr =
681-
std::make_unique<BitGetCmd>(kCmdNameBitGet, 3, kCmdFlagsRead | kCmdFlagsBit | kCmdFlagsSlow | kCmdFlagsDoThroughDB | kCmdFlagsReadCache | kCmdFlagsUpdateCache);
681+
std::make_unique<BitGetCmd>(kCmdNameBitGet, 3, kCmdFlagsRead | kCmdFlagsBit | kCmdFlagsSlow);
682682
cmd_table->insert(std::pair<std::string, std::unique_ptr<Cmd>>(kCmdNameBitGet, std::move(bitgetptr)));
683683
////bitcountCmd
684684
std::unique_ptr<Cmd> bitcountptr =
685685
std::make_unique<BitCountCmd>(kCmdNameBitCount, -2, kCmdFlagsRead | kCmdFlagsBit | kCmdFlagsSlow | kCmdFlagsDoThroughDB | kCmdFlagsReadCache | kCmdFlagsUpdateCache);
686686
cmd_table->insert(std::pair<std::string, std::unique_ptr<Cmd>>(kCmdNameBitCount, std::move(bitcountptr)));
687687
////bitposCmd
688688
std::unique_ptr<Cmd> bitposptr =
689-
std::make_unique<BitPosCmd>(kCmdNameBitPos, -3, kCmdFlagsRead | kCmdFlagsBit | kCmdFlagsSlow | kCmdFlagsDoThroughDB | kCmdFlagsReadCache | kCmdFlagsUpdateCache);
689+
std::make_unique<BitPosCmd>(kCmdNameBitPos, -3, kCmdFlagsRead | kCmdFlagsBit | kCmdFlagsSlow);
690690
cmd_table->insert(std::pair<std::string, std::unique_ptr<Cmd>>(kCmdNameBitPos, std::move(bitposptr)));
691691
////bitopCmd
692692
std::unique_ptr<Cmd> bitopptr =

‎src/pika_kv.cc‎

Lines changed: 76 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -509,26 +509,41 @@ void MgetCmd::DoInitial() {
509509
keys_ = argv_;
510510
keys_.erase(keys_.begin());
511511
split_res_.resize(keys_.size());
512+
cache_miss_keys_.clear();
513+
}
514+
515+
void MgetCmd::AssembleResponseFromCache() {
516+
res_.AppendArrayLenUint64(keys_.size());
517+
for (const auto& key : keys_) {
518+
auto it = cache_hit_values_.find(key);
519+
if (it != cache_hit_values_.end()) {
520+
res_.AppendStringLen(it->second.size());
521+
res_.AppendContent(it->second);
522+
} else {
523+
res_.SetRes(CmdRes::kErrOther, "Internal error during cache assembly");
524+
return;
525+
}
526+
}
512527
}
513528

514529
void MgetCmd::Do() {
530+
// Without using the cache and querying only the DB, we need to use keys_.
531+
// This line will only be assigned when querying the DB directly.
532+
if(cache_miss_keys_.size() == 0) {
533+
cache_miss_keys_ = keys_;
534+
}
515535
db_value_status_array_.clear();
516-
s_ = db_->storage()->MGet(keys_, &db_value_status_array_);
517-
if (s_.ok()) {
518-
res_.AppendArrayLenUint64(db_value_status_array_.size());
519-
for (const auto& vs : db_value_status_array_) {
520-
if (vs.status.ok()) {
521-
res_.AppendStringLenUint64(vs.value.size());
522-
res_.AppendContent(vs.value);
523-
} else {
524-
res_.AppendContent("$-1");
525-
}
536+
s_ = db_->storage()->MGet(cache_miss_keys_, &db_value_status_array_);
537+
if (!s_.ok()) {
538+
if (s_.IsInvalidArgument()) {
539+
res_.SetRes(CmdRes::kMultiKey);
540+
} else {
541+
res_.SetRes(CmdRes::kErrOther, s_.ToString());
526542
}
527-
} else if (s_.IsInvalidArgument()) {
528-
res_.SetRes(CmdRes::kMultiKey);
529-
} else {
530-
res_.SetRes(CmdRes::kErrOther, s_.ToString());
543+
return;
531544
}
545+
546+
MergeCachedAndDbResults();
532547
}
533548

534549
void MgetCmd::Split(const HintKeys& hint_keys) {
@@ -560,34 +575,67 @@ void MgetCmd::Merge() {
560575
}
561576
}
562577

578+
void MgetCmd::DoThroughDB() {
579+
res_.clear();
580+
Do();
581+
}
582+
563583
void MgetCmd::ReadCache() {
564-
if (1 < keys_.size()) {
565-
res_.SetRes(CmdRes::kCacheMiss);
566-
return;
584+
for (const auto key : keys_) {
585+
std::string value;
586+
auto s = db_->cache()->Get(const_cast<std::string&>(key), &value);
587+
if (s.ok()) {
588+
cache_hit_values_[key] = value;
589+
} else {
590+
cache_miss_keys_.push_back(key);
591+
}
567592
}
568-
auto s = db_->cache()->Get(keys_[0], &value_);
569-
if (s.ok()) {
570-
res_.AppendArrayLen(1);
571-
res_.AppendStringLen(value_.size());
572-
res_.AppendContent(value_);
593+
if (cache_miss_keys_.empty()) {
594+
AssembleResponseFromCache();
573595
} else {
574596
res_.SetRes(CmdRes::kCacheMiss);
575597
}
576598
}
577599

578-
void MgetCmd::DoThroughDB() {
579-
res_.clear();
580-
Do();
600+
void MgetCmd::DoUpdateCache() {
601+
size_t db_index = 0;
602+
for (const auto key : cache_miss_keys_) {
603+
if (db_index < db_value_status_array_.size() && db_value_status_array_[db_index].status.ok()) {
604+
db_->cache()->WriteKVToCache(const_cast<std::string&>(key), db_value_status_array_[db_index].value, db_value_status_array_[db_index].ttl);
605+
}
606+
db_index++;
607+
}
581608
}
582609

583-
void MgetCmd::DoUpdateCache() {
584-
for (size_t i = 0; i < keys_.size(); i++) {
610+
void MgetCmd::MergeCachedAndDbResults() {
611+
res_.AppendArrayLenUint64(keys_.size());
612+
613+
std::unordered_map<std::string, std::string> db_results_map;
614+
for (size_t i = 0; i < cache_miss_keys_.size(); ++i) {
585615
if (db_value_status_array_[i].status.ok()) {
586-
db_->cache()->WriteKVToCache(keys_[i], db_value_status_array_[i].value, db_value_status_array_[i].ttl);
616+
db_results_map[cache_miss_keys_[i]] = db_value_status_array_[i].value;
617+
}
618+
}
619+
620+
for (const auto& key : keys_) {
621+
auto cache_it = cache_hit_values_.find(key);
622+
623+
if (cache_it != cache_hit_values_.end()) {
624+
res_.AppendStringLen(cache_it->second.size());
625+
res_.AppendContent(cache_it->second);
626+
} else {
627+
auto db_it = db_results_map.find(key);
628+
if (db_it != db_results_map.end()) {
629+
res_.AppendStringLen(db_it->second.size());
630+
res_.AppendContent(db_it->second);
631+
} else {
632+
res_.AppendContent("$-1");
633+
}
587634
}
588635
}
589636
}
590637

638+
591639
void KeysCmd::DoInitial() {
592640
if (!CheckArg(argv_.size())) {
593641
res_.SetRes(CmdRes::kWrongNum, kCmdNameKeys);

‎src/pika_server.cc‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1688,9 +1688,7 @@ void PikaServer::DoCacheBGTask(void* arg) {
16881688
}
16891689

16901690
db->cache()->SetCacheStatus(PIKA_CACHE_STATUS_OK);
1691-
if (pCacheTaskArg->reenable_cache) {
1692-
pCacheTaskArg->conf->UnsetCacheDisableFlag();
1693-
}
1691+
g_pika_conf->UnsetCacheDisableFlag();
16941692
}
16951693

16961694
void PikaServer::ResetCacheConfig(std::shared_ptr<DB> db) {
@@ -1710,7 +1708,7 @@ void PikaServer::ClearHitRatio(std::shared_ptr<DB> db) {
17101708

17111709
void PikaServer::OnCacheStartPosChanged(int zset_cache_start_direction, std::shared_ptr<DB> db) {
17121710
ResetCacheConfig(db);
1713-
ClearCacheDbAsync(db);
1711+
ClearCacheDbAsyncV2(db);
17141712
}
17151713

17161714
void PikaServer::ClearCacheDbAsyncV2(std::shared_ptr<DB> db) {

‎tests/integration/cache_test.go‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,4 +201,167 @@ var _ = Describe("Cache test", func() {
201201
Expect(mGet4.Err()).NotTo(HaveOccurred())
202202
Expect(mGet4.Val()).To(Equal([]interface{}{nil, nil, nil, nil}))
203203
})
204+
205+
It("should mget for multi key in cache and db", func() {
206+
multiset1 := client.Set(ctx, "key1", "a", 3000*time.Millisecond)
207+
Expect(multiset1.Err()).NotTo(HaveOccurred())
208+
Expect(multiset1.Val()).To(Equal("OK"))
209+
210+
multiset2 := client.Set(ctx, "key2", "b", 3000*time.Millisecond)
211+
Expect(multiset2.Err()).NotTo(HaveOccurred())
212+
Expect(multiset2.Val()).To(Equal("OK"))
213+
214+
multiset3 := client.Set(ctx, "key3", "c", 3000*time.Millisecond)
215+
Expect(multiset3.Err()).NotTo(HaveOccurred())
216+
Expect(multiset3.Val()).To(Equal("OK"))
217+
218+
multiset4 := client.Set(ctx, "key4", "d", 3000*time.Millisecond)
219+
Expect(multiset4.Err()).NotTo(HaveOccurred())
220+
Expect(multiset4.Val()).To(Equal("OK"))
221+
222+
multikey1 := client.MGet(ctx, "key1")
223+
Expect(multikey1.Err()).NotTo(HaveOccurred())
224+
Expect(multikey1.Val()).To(Equal([]interface{}{"a"}))
225+
226+
MultiKey2 := client.Get(ctx, "key1")
227+
Expect(MultiKey2.Err()).NotTo(HaveOccurred())
228+
Expect(MultiKey2.Val()).To(Equal("a"))
229+
230+
MultiMget := client.MGet(ctx, "key1", "key2", "key3", "key4")
231+
Expect(MultiMget.Err()).NotTo(HaveOccurred())
232+
Expect(MultiMget.Val()).To(Equal([]interface{}{"a", "b", "c", "d"}))
233+
})
234+
235+
It("should mget for multi key in cache", func() {
236+
multiset1 := client.Set(ctx, "key1", "a", 3000*time.Millisecond)
237+
Expect(multiset1.Err()).NotTo(HaveOccurred())
238+
Expect(multiset1.Val()).To(Equal("OK"))
239+
240+
multiset2 := client.Set(ctx, "key2", "b", 3000*time.Millisecond)
241+
Expect(multiset2.Err()).NotTo(HaveOccurred())
242+
Expect(multiset2.Val()).To(Equal("OK"))
243+
244+
multiset3 := client.Set(ctx, "key3", "c", 3000*time.Millisecond)
245+
Expect(multiset3.Err()).NotTo(HaveOccurred())
246+
Expect(multiset3.Val()).To(Equal("OK"))
247+
248+
multiset4 := client.Set(ctx, "key4", "d", 3000*time.Millisecond)
249+
Expect(multiset4.Err()).NotTo(HaveOccurred())
250+
Expect(multiset4.Val()).To(Equal("OK"))
251+
252+
multikey1 := client.MGet(ctx, "key1")
253+
Expect(multikey1.Err()).NotTo(HaveOccurred())
254+
Expect(multikey1.Val()).To(Equal([]interface{}{"a"}))
255+
256+
MultiKey2 := client.Get(ctx, "key1")
257+
Expect(MultiKey2.Err()).NotTo(HaveOccurred())
258+
Expect(MultiKey2.Val()).To(Equal("a"))
259+
260+
MultiMget := client.MGet(ctx, "key1", "key2", "key3", "key4")
261+
Expect(MultiMget.Err()).NotTo(HaveOccurred())
262+
Expect(MultiMget.Val()).To(Equal([]interface{}{"a", "b", "c", "d"}))
263+
})
264+
265+
It("should mget for multi key in db", func() {
266+
multiset1 := client.Set(ctx, "key1", "a", 3000*time.Millisecond)
267+
Expect(multiset1.Err()).NotTo(HaveOccurred())
268+
Expect(multiset1.Val()).To(Equal("OK"))
269+
270+
multiset2 := client.Set(ctx, "key2", "b", 3000*time.Millisecond)
271+
Expect(multiset2.Err()).NotTo(HaveOccurred())
272+
Expect(multiset2.Val()).To(Equal("OK"))
273+
274+
multiset3 := client.Set(ctx, "key3", "c", 3000*time.Millisecond)
275+
Expect(multiset3.Err()).NotTo(HaveOccurred())
276+
Expect(multiset3.Val()).To(Equal("OK"))
277+
278+
multiset4 := client.Set(ctx, "key4", "d", 3000*time.Millisecond)
279+
Expect(multiset4.Err()).NotTo(HaveOccurred())
280+
Expect(multiset4.Val()).To(Equal("OK"))
281+
282+
multikey1 := client.MGet(ctx, "key1")
283+
Expect(multikey1.Err()).NotTo(HaveOccurred())
284+
Expect(multikey1.Val()).To(Equal([]interface{}{"a"}))
285+
286+
MultiKey2 := client.Get(ctx, "key1")
287+
Expect(MultiKey2.Err()).NotTo(HaveOccurred())
288+
Expect(MultiKey2.Val()).To(Equal("a"))
289+
290+
multikey3 := client.MGet(ctx, "key2")
291+
Expect(multikey3.Err()).NotTo(HaveOccurred())
292+
Expect(multikey3.Val()).To(Equal([]interface{}{"b"}))
293+
294+
multikey4 := client.MGet(ctx, "key3")
295+
Expect(multikey4.Err()).NotTo(HaveOccurred())
296+
Expect(multikey4.Val()).To(Equal([]interface{}{"c"}))
297+
298+
multikey5 := client.MGet(ctx, "key4")
299+
Expect(multikey5.Err()).NotTo(HaveOccurred())
300+
Expect(multikey5.Val()).To(Equal([]interface{}{"d"}))
301+
302+
MultiMget := client.MGet(ctx, "key1", "key2", "key3", "key4")
303+
Expect(MultiMget.Err()).NotTo(HaveOccurred())
304+
Expect(MultiMget.Val()).To(Equal([]interface{}{"a", "b", "c", "d"}))
305+
})
306+
307+
It("should mget for multi key in db", func() {
308+
multiset1 := client.Set(ctx, "key1", "a", 3000*time.Millisecond)
309+
Expect(multiset1.Err()).NotTo(HaveOccurred())
310+
Expect(multiset1.Val()).To(Equal("OK"))
311+
312+
multiset2 := client.Set(ctx, "key2", "b", 3000*time.Millisecond)
313+
Expect(multiset2.Err()).NotTo(HaveOccurred())
314+
Expect(multiset2.Val()).To(Equal("OK"))
315+
316+
multiset3 := client.Set(ctx, "key3", "c", 3000*time.Millisecond)
317+
Expect(multiset3.Err()).NotTo(HaveOccurred())
318+
Expect(multiset3.Val()).To(Equal("OK"))
319+
320+
multiset4 := client.Set(ctx, "key4", "d", 3000*time.Millisecond)
321+
Expect(multiset4.Err()).NotTo(HaveOccurred())
322+
Expect(multiset4.Val()).To(Equal("OK"))
323+
324+
MultiMget := client.MGet(ctx, "key1", "key2", "key3", "key4")
325+
Expect(MultiMget.Err()).NotTo(HaveOccurred())
326+
Expect(MultiMget.Val()).To(Equal([]interface{}{"a", "b", "c", "d"}))
327+
})
328+
329+
It("MGET against non existing key", func() {
330+
multiset1 := client.Set(ctx, "key1", "a", 3000*time.Millisecond)
331+
Expect(multiset1.Err()).NotTo(HaveOccurred())
332+
Expect(multiset1.Val()).To(Equal("OK"))
333+
334+
multiset3 := client.Set(ctx, "key3", "c", 3000*time.Millisecond)
335+
Expect(multiset3.Err()).NotTo(HaveOccurred())
336+
Expect(multiset3.Val()).To(Equal("OK"))
337+
338+
multiset4 := client.Set(ctx, "key4", "d", 3000*time.Millisecond)
339+
Expect(multiset4.Err()).NotTo(HaveOccurred())
340+
Expect(multiset4.Val()).To(Equal("OK"))
341+
342+
MultiMget := client.MGet(ctx, "key1", "key2", "key3", "key4")
343+
Expect(MultiMget.Err()).NotTo(HaveOccurred())
344+
Expect(MultiMget.Val()).To(Equal([]interface{}{"a", nil, "c", "d"}))
345+
})
346+
It("MGET against non-string key", func() {
347+
SetMultiKey := client.Set(ctx, "foo{t}", "BAR", 3000*time.Millisecond)
348+
Expect(SetMultiKey.Err()).NotTo(HaveOccurred())
349+
Expect(SetMultiKey.Val()).To(Equal("OK"))
350+
351+
SetMultiKey1 := client.Set(ctx, "bar{t}", "FOO", 3000*time.Millisecond)
352+
Expect(SetMultiKey1.Err()).NotTo(HaveOccurred())
353+
Expect(SetMultiKey1.Val()).To(Equal("OK"))
354+
355+
SaddMultiKey := client.SAdd(ctx, "myset{t}", "ciao")
356+
Expect(SaddMultiKey.Err()).NotTo(HaveOccurred())
357+
Expect(SaddMultiKey.Val()).To(Equal(int64(1)))
358+
359+
SaddMultiKey1 := client.SAdd(ctx, "myset{t}", "bau")
360+
Expect(SaddMultiKey1.Err()).NotTo(HaveOccurred())
361+
Expect(SaddMultiKey1.Val()).To(Equal(int64(1)))
362+
363+
MultiMget := client.MGet(ctx, "foo{t}", "baazz{t}", "bar{t}", "myset{t}")
364+
Expect(MultiMget.Err()).NotTo(HaveOccurred())
365+
Expect(MultiMget.Val()).To(Equal([]interface{}{"BAR", nil, "FOO", nil}))
366+
})
204367
})

0 commit comments

Comments
 (0)