From 0541a65a9c9ff73fe532b7149b56382caf3dbeea Mon Sep 17 00:00:00 2001
From: leipeng <peng@topling.cn>
Date: Mon, 24 Apr 2023 21:05:34 +0800
Subject: [PATCH] Add MemTableRep::NeedsUserKeyCompareInGet() and relavant
 changes

---
 db/memtable.cc                | 9 ++++++++-
 db/memtable.h                 | 1 +
 include/rocksdb/memtablerep.h | 2 ++
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/db/memtable.cc b/db/memtable.cc
index 4f327c591..c8e2f2b1f 100644
--- a/db/memtable.cc
+++ b/db/memtable.cc
@@ -113,6 +113,7 @@ MemTable::MemTable(const InternalKeyComparator& cmp,
       oldest_key_time_(std::numeric_limits<uint64_t>::max()),
       atomic_flush_seqno_(kMaxSequenceNumber),
       approximate_memory_usage_(0) {
+  needs_user_key_cmp_in_get_ = table_->NeedsUserKeyCompareInGet();
   UpdateFlushState();
   // something went wrong if we need to flush before inserting anything
   assert(!ShouldScheduleFlush());
@@ -792,6 +793,7 @@ struct Saver {
   bool do_merge;
   bool allow_data_in_errors;
   bool is_zero_copy;
+  bool needs_user_key_cmp_in_get;
   bool CheckCallback(SequenceNumber _seq) {
     if (callback_) {
       return callback_->IsVisible(_seq);
@@ -832,8 +834,11 @@ static bool SaveValue(void* arg, const MemTableRep::KeyValuePair& pair) {
 #else
   constexpr size_t ts_sz = 0; // let compiler optimize it out
 #endif
-  if (user_comparator->EqualWithoutTimestamp(user_key_slice,
+  if (!s->needs_user_key_cmp_in_get ||
+      user_comparator->EqualWithoutTimestamp(user_key_slice,
                                              s->key->user_key())) {
+    assert(user_comparator->EqualWithoutTimestamp(user_key_slice,
+                                             s->key->user_key()));
     // Correct user key
     const uint64_t tag = DecodeFixed64(key_ptr + key_length - 8);
     ValueType type;
@@ -1233,6 +1238,7 @@ bool MemTable::Get(const LookupKey& key, PinnableSlice* value,
     saver.do_merge = do_merge;
     saver.allow_data_in_errors = moptions_.allow_data_in_errors;
     saver.is_zero_copy = read_opts.pinning_tls != nullptr;
+    saver.needs_user_key_cmp_in_get = needs_user_key_cmp_in_get_;
     if (LIKELY(value != nullptr)) {
       value->Reset();
     }
@@ -1332,6 +1338,7 @@ void MemTable::MultiGet(const ReadOptions& read_options, MultiGetRange* range,
     saver.do_merge = true;
     saver.allow_data_in_errors = moptions_.allow_data_in_errors;
     saver.is_zero_copy = read_options.pinning_tls != nullptr;
+    saver.needs_user_key_cmp_in_get = needs_user_key_cmp_in_get_;
     table_->Get(read_options, *(iter->lkey), &saver, SaveValue);
 
     if (!saver.found_final_value && saver.merge_in_progress) {
diff --git a/db/memtable.h b/db/memtable.h
index 39acaecf6..af2a17a58 100644
--- a/db/memtable.h
+++ b/db/memtable.h
@@ -557,6 +557,7 @@ class MemTable {
   // These are used to manage memtable flushes to storage
   bool flush_in_progress_;  // started the flush
   bool flush_completed_;    // finished the flush
+  bool needs_user_key_cmp_in_get_;
   uint64_t file_number_;    // filled up after flush is complete
 
   // The updates to be applied to the transaction log when this
diff --git a/include/rocksdb/memtablerep.h b/include/rocksdb/memtablerep.h
index 6fd2bfa0b..799472c7b 100644
--- a/include/rocksdb/memtablerep.h
+++ b/include/rocksdb/memtablerep.h
@@ -343,6 +343,8 @@ class MemTableRep {
   // Default: true
   virtual bool IsSnapshotSupported() const { return true; }
 
+  virtual bool NeedsUserKeyCompareInGet() const { return true; }
+
  protected:
   // When *key is an internal key concatenated with the value, returns the
   // user key.