@nxtedition/rocksdb 15.4.0 → 15.5.0

package/deps/rocksdb/rocksdb/db/listener_test.cc

@@ -105,7 +105,7 @@ class TestCompactionListener : public EventListener {
     ASSERT_EQ(ci.output_files.size(), ci.output_file_infos.size());
 
     ASSERT_TRUE(test_);
-    ASSERT_EQ(test_->db_, db);
+    ASSERT_EQ(test_->db_.get(), db);
 
     std::vector<std::vector<FileMetaData>> files_by_level;
     test_->dbfull()->TEST_GetFilesMetaData(test_->handles_[ci.cf_id],
@@ -197,7 +197,7 @@ TEST_F(EventListenerTest, OnSingleDBCompactionTest) {
 
   ASSERT_EQ(listener->compacted_dbs_.size(), cf_names.size());
   for (size_t i = 0; i < cf_names.size(); ++i) {
-    ASSERT_EQ(listener->compacted_dbs_[i], db_);
+    ASSERT_EQ(listener->compacted_dbs_[i], db_.get());
   }
 }
 
@@ -268,7 +268,7 @@ class TestFlushListener : public EventListener {
     // that assumption does not hold (see the test case MultiDBMultiListeners
     // below).
     ASSERT_TRUE(test_);
-    if (db == test_->db_) {
+    if (db == test_->db_.get()) {
       std::vector<std::vector<FileMetaData>> files_by_level;
       ASSERT_LT(info.cf_id, test_->handles_.size());
       ASSERT_GE(info.cf_id, 0u);
@@ -343,7 +343,7 @@ TEST_F(EventListenerTest, OnSingleDBFlushTest) {
 
   // make sure callback functions are called in the right order
   for (size_t i = 0; i < cf_names.size(); ++i) {
-    ASSERT_EQ(listener->flushed_dbs_[i], db_);
+    ASSERT_EQ(listener->flushed_dbs_[i], db_.get());
     ASSERT_EQ(listener->flushed_column_family_names_[i], cf_names[i]);
   }
 }
@@ -387,7 +387,7 @@ TEST_F(EventListenerTest, MultiCF) {
       // make sure callback functions are called in the right order
       if (i == 7) {
         for (size_t j = 0; j < cf_names.size(); j++) {
-          ASSERT_EQ(listener->flushed_dbs_[j], db_);
+          ASSERT_EQ(listener->flushed_dbs_[j], db_.get());
           ASSERT_EQ(listener->flushed_column_family_names_[j], cf_names[j]);
         }
       }
@@ -422,22 +422,21 @@ TEST_F(EventListenerTest, MultiDBMultiListeners) {
   DBOptions db_opts(options);
   ColumnFamilyOptions cf_opts(options);
 
-  std::vector<DB*> dbs;
+  std::vector<std::unique_ptr<DB>> dbs;
   std::vector<std::vector<ColumnFamilyHandle*>> vec_handles;
 
   for (int d = 0; d < kNumDBs; ++d) {
     ASSERT_OK(DestroyDB(dbname_ + std::to_string(d), options));
-    DB* db;
+    ASSERT_OK(
+        DB::Open(options, dbname_ + std::to_string(d), &dbs.emplace_back()));
     std::vector<ColumnFamilyHandle*> handles;
-    ASSERT_OK(DB::Open(options, dbname_ + std::to_string(d), &db));
     for (size_t c = 0; c < cf_names.size(); ++c) {
       ColumnFamilyHandle* handle;
-      ASSERT_OK(db->CreateColumnFamily(cf_opts, cf_names[c], &handle));
+      ASSERT_OK(dbs.back()->CreateColumnFamily(cf_opts, cf_names[c], &handle));
       handles.push_back(handle);
     }
 
     vec_handles.push_back(std::move(handles));
-    dbs.push_back(db);
   }
 
   for (int d = 0; d < kNumDBs; ++d) {
@@ -450,23 +449,23 @@ TEST_F(EventListenerTest, MultiDBMultiListeners) {
   for (size_t c = 0; c < cf_names.size(); ++c) {
     for (int d = 0; d < kNumDBs; ++d) {
       ASSERT_OK(dbs[d]->Flush(FlushOptions(), vec_handles[d][c]));
-      ASSERT_OK(
-          static_cast_with_check<DBImpl>(dbs[d])->TEST_WaitForFlushMemTable());
+      ASSERT_OK(static_cast_with_check<DBImpl>(dbs[d].get())
+                    ->TEST_WaitForFlushMemTable());
     }
   }
 
   for (int d = 0; d < kNumDBs; ++d) {
     // Ensure background work is fully finished including listener callbacks
     // before accessing listener state.
-    ASSERT_OK(
-        static_cast_with_check<DBImpl>(dbs[d])->TEST_WaitForBackgroundWork());
+    ASSERT_OK(static_cast_with_check<DBImpl>(dbs[d].get())
+                  ->TEST_WaitForBackgroundWork());
   }
 
   for (auto* listener : listeners) {
     int pos = 0;
     for (size_t c = 0; c < cf_names.size(); ++c) {
       for (int d = 0; d < kNumDBs; ++d) {
-        ASSERT_EQ(listener->flushed_dbs_[pos], dbs[d]);
+        ASSERT_EQ(listener->flushed_dbs_[pos], dbs[d].get());
         ASSERT_EQ(listener->flushed_column_family_names_[pos], cf_names[c]);
         pos++;
       }
@@ -481,8 +480,8 @@ TEST_F(EventListenerTest, MultiDBMultiListeners) {
   }
   vec_handles.clear();
 
-  for (auto db : dbs) {
-    delete db;
+  for (auto& db : dbs) {
+    db.reset();
   }
 }
 
@@ -1309,16 +1308,21 @@ class BlobDBJobLevelEventListenerTest : public EventListener {
   explicit BlobDBJobLevelEventListenerTest(EventListenerTest* test)
       : test_(test), call_count_(0) {}
 
-  const VersionStorageInfo* GetVersionStorageInfo() const {
-    VersionSet* const versions = test_->dbfull()->GetVersionSet();
+  // NOTE: it's not safe to rely on test_->db_ for these functions because
+  // the DB may be in the process of closing when these are called, and the
+  // unique_ptr is set to nullptr before invoking ~DB()
+
+  const VersionStorageInfo* GetVersionStorageInfo(DB* db) const {
+    DBImpl* db_impl = static_cast_with_check<DBImpl>(db);
+    VersionSet* const versions = db_impl->GetVersionSet();
     assert(versions);
 
     ColumnFamilyData* const cfd = versions->GetColumnFamilySet()->GetDefault();
     EXPECT_NE(cfd, nullptr);
 
-    test_->dbfull()->TEST_LockMutex();
+    db_impl->TEST_LockMutex();
     Version* const current = cfd->current();
-    test_->dbfull()->TEST_UnlockMutex();
+    db_impl->TEST_UnlockMutex();
     EXPECT_NE(current, nullptr);
 
     const VersionStorageInfo* const storage_info = current->storage_info();
@@ -1328,8 +1332,9 @@ class BlobDBJobLevelEventListenerTest : public EventListener {
   }
 
   void CheckBlobFileAdditions(
+      DB* db,
       const std::vector<BlobFileAdditionInfo>& blob_file_addition_infos) const {
-    const auto* vstorage = GetVersionStorageInfo();
+    const auto* vstorage = GetVersionStorageInfo(db);
 
     EXPECT_FALSE(blob_file_addition_infos.empty());
 
@@ -1357,7 +1362,7 @@ class BlobDBJobLevelEventListenerTest : public EventListener {
     return result;
   }
 
-  void OnFlushCompleted(DB* /*db*/, const FlushJobInfo& info) override {
+  void OnFlushCompleted(DB* db, const FlushJobInfo& info) override {
     {
       std::lock_guard<std::mutex> lock(mutex_);
       IncreaseCallCount(/*mutex_locked*/ true);
@@ -1366,16 +1371,15 @@ class BlobDBJobLevelEventListenerTest : public EventListener {
 
     EXPECT_EQ(info.blob_compression_type, kNoCompression);
 
-    CheckBlobFileAdditions(info.blob_file_addition_infos);
+    CheckBlobFileAdditions(db, info.blob_file_addition_infos);
   }
 
-  void OnCompactionCompleted(DB* /*db*/,
-                             const CompactionJobInfo& info) override {
+  void OnCompactionCompleted(DB* db, const CompactionJobInfo& info) override {
     IncreaseCallCount(/*mutex_locked*/ false);
 
     EXPECT_EQ(info.blob_compression_type, kNoCompression);
 
-    CheckBlobFileAdditions(info.blob_file_addition_infos);
+    CheckBlobFileAdditions(db, info.blob_file_addition_infos);
 
     EXPECT_FALSE(info.blob_file_garbage_infos.empty());
 
@@ -1639,6 +1643,129 @@ TEST_F(EventListenerTest, BlobDBFileTest) {
   blob_event_listener->CheckCounters();
 }
 
+class BackgroundJobPressureTestListener : public EventListener {
+ public:
+  void OnBackgroundJobPressureChanged(
+      DB* /*db*/, const BackgroundJobPressure& snapshot) override {
+    std::lock_guard<std::mutex> lock(mu_);
+    snapshots_.push_back(snapshot);
+  }
+
+  std::vector<BackgroundJobPressure> GetSnapshots() {
+    std::lock_guard<std::mutex> lock(mu_);
+    return snapshots_;
+  }
+
+  void Reset() {
+    std::lock_guard<std::mutex> lock(mu_);
+    snapshots_.clear();
+  }
+
+ private:
+  std::mutex mu_;
+  std::vector<BackgroundJobPressure> snapshots_;
+};
+
+TEST_F(EventListenerTest, BackgroundJobPressure) {
+  auto listener = std::make_shared<BackgroundJobPressureTestListener>();
+  Options options = CurrentOptions();
+  options.create_if_missing = true;
+  options.max_background_jobs = 8;
+  options.disable_auto_compactions = false;
+  options.level0_slowdown_writes_trigger = 4;
+  options.level0_stop_writes_trigger = 100;
+  options.level0_file_num_compaction_trigger = 4;
+  options.listeners.push_back(listener);
+  DestroyAndReopen(options);
+
+  // Structural invariant checked on every snapshot.
+  auto CheckInvariants = [](const std::vector<BackgroundJobPressure>& snaps) {
+    ASSERT_GT(snaps.size(), 0u);
+    for (const auto& s : snaps) {
+      ASSERT_EQ(s.compaction_scheduled,
+                s.compaction_low_scheduled + s.compaction_bottom_scheduled);
+      ASSERT_EQ(s.compaction_running,
+                s.compaction_low_running + s.compaction_bottom_running);
+    }
+  };
+
+  // Block compaction so L0 files accumulate.
+  env_->SetBackgroundThreads(1, Env::Priority::LOW);
+  test::SleepingBackgroundTask sleeping_task;
+  env_->Schedule(&test::SleepingBackgroundTask::DoSleepTask, &sleeping_task,
+                 Env::Priority::LOW);
+  sleeping_task.WaitUntilSleeping();
+
+  // Phase 1: No pressure — 3 SST files (below slowdown trigger=4).
+  for (int i = 0; i < 3; i++) {
+    ASSERT_OK(Put("k" + std::to_string(i), std::string(100, 'x')));
+    ASSERT_OK(Flush());
+  }
+
+  auto snapshots = listener->GetSnapshots();
+  CheckInvariants(snapshots);
+  for (const auto& s : snapshots) {
+    ASSERT_EQ(s.compaction_running, 0);
+    ASSERT_EQ(s.compaction_scheduled, 0);
+    ASSERT_FALSE(s.compaction_speedup_active);
+    ASSERT_LT(s.write_stall_proximity_pct, 100);
+  }
+  // Note: we don't assert flush_running/flush_scheduled == 0 here because
+  // MaybeScheduleFlushOrCompaction() runs before the pressure callback and
+  // may schedule new flush work, making flush counts non-deterministic.
+
+  // Phase 2: Build pressure — flush past slowdown trigger (4 L0 SST files).
+  // Compaction is blocked, so L0 SST files pile up.
+  listener->Reset();
+  {
+    FlushOptions fo;
+    fo.allow_write_stall = true;
+    for (int i = 3; i < 10; i++) {
+      ASSERT_OK(Put("k" + std::to_string(i), std::string(100, 'x')));
+      ASSERT_OK(db_->Flush(fo));
+    }
+  }
+
+  snapshots = listener->GetSnapshots();
+  CheckInvariants(snapshots);
+  for (const auto& s : snapshots) {
+    ASSERT_EQ(s.compaction_running, 0);
+  }
+  // Scan history: pressure indicators appear as L0 SST files accumulate
+  bool found_speedup = false;
+  bool found_compaction_scheduled = false;
+  bool found_high_proximity = false;
+  for (const auto& s : snapshots) {
+    if (s.compaction_speedup_active) {
+      found_speedup = true;
+    }
+    if (s.compaction_scheduled > 0) {
+      found_compaction_scheduled = true;
+    }
+    if (s.write_stall_proximity_pct >= 100) {
+      found_high_proximity = true;
+    }
+  }
+  ASSERT_TRUE(found_speedup);
+  ASSERT_TRUE(found_compaction_scheduled);
+  ASSERT_TRUE(found_high_proximity);
+
+  // Phase 3: Relieve pressure — unblock compaction, wait for completion.
+  listener->Reset();
+  sleeping_task.WakeUp();
+  sleeping_task.WaitUntilDone();
+  ASSERT_OK(dbfull()->TEST_WaitForCompact());
+
+  snapshots = listener->GetSnapshots();
+  CheckInvariants(snapshots);
+  // Latest: all compactions finished, healthy
+  const auto& latest3 = snapshots.back();
+  ASSERT_EQ(latest3.compaction_running, 0);
+  ASSERT_EQ(latest3.compaction_scheduled, 0);
+  ASSERT_FALSE(latest3.compaction_speedup_active);
+  ASSERT_LT(latest3.write_stall_proximity_pct, 100);
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {

package/deps/rocksdb/rocksdb/db/manual_compaction_test.cc

@@ -98,7 +98,7 @@ class LogCompactionFilter : public CompactionFilter {
 
 TEST_F(ManualCompactionTest, CompactTouchesAllKeys) {
   for (int iter = 0; iter < 2; ++iter) {
-    DB* db;
+    std::unique_ptr<DB> db;
     Options options;
     if (iter == 0) {  // level compaction
       options.num_levels = 3;
@@ -128,7 +128,7 @@ TEST_F(ManualCompactionTest, CompactTouchesAllKeys) {
     delete itr;
 
     delete options.compaction_filter;
-    delete db;
+    db.reset();
     ASSERT_OK(DestroyDB(dbname_, options));
   }
 }
@@ -137,7 +137,7 @@ TEST_F(ManualCompactionTest, Test) {
   // Open database.  Disable compression since it affects the creation
   // of layers and the code below is trying to test against a very
   // specific scenario.
-  DB* db;
+  std::unique_ptr<DB> db;
   Options db_options;
   db_options.write_buffer_size = 1024;
   db_options.create_if_missing = true;
@@ -185,12 +185,12 @@ TEST_F(ManualCompactionTest, Test) {
   ASSERT_EQ(kNumKeys, num_keys) << "Bad number of keys";
 
   // close database
-  delete db;
+  db.reset();
   ASSERT_OK(DestroyDB(dbname_, Options()));
 }
 
 TEST_F(ManualCompactionTest, SkipLevel) {
-  DB* db;
+  std::unique_ptr<DB> db;
   Options options;
   options.level_compaction_dynamic_level_bytes = false;
   options.num_levels = 3;
@@ -298,7 +298,7 @@ TEST_F(ManualCompactionTest, SkipLevel) {
   }
 
   delete filter;
-  delete db;
+  db.reset();
   ASSERT_OK(DestroyDB(dbname_, options));
 }
 

package/deps/rocksdb/rocksdb/db/memtable.cc

@@ -72,7 +72,9 @@ ImmutableMemTableOptions::ImmutableMemTableOptions(
       allow_data_in_errors(ioptions.allow_data_in_errors),
       paranoid_memory_checks(mutable_cf_options.paranoid_memory_checks),
       memtable_veirfy_per_key_checksum_on_seek(
-          mutable_cf_options.memtable_veirfy_per_key_checksum_on_seek) {}
+          mutable_cf_options.memtable_veirfy_per_key_checksum_on_seek),
+      memtable_batch_lookup_optimization(
+          ioptions.memtable_batch_lookup_optimization) {}
 
 MemTable::MemTable(const InternalKeyComparator& cmp,
                    const ImmutableOptions& ioptions,
@@ -1575,65 +1577,179 @@ void MemTable::MultiGet(const ReadOptions& read_options, MultiGetRange* range,
       }
     }
   }
-  for (auto iter = temp_range.begin(); iter != temp_range.end(); ++iter) {
-    bool found_final_value{false};
-    bool merge_in_progress = iter->s->IsMergeInProgress();
-    if (!no_range_del) {
-      std::unique_ptr<FragmentedRangeTombstoneIterator> range_del_iter(
-          NewRangeTombstoneIteratorInternal(
-              read_options, GetInternalKeySeqno(iter->lkey->internal_key()),
-              immutable_memtable));
-      SequenceNumber covering_seq =
-          range_del_iter->MaxCoveringTombstoneSeqnum(iter->lkey->user_key());
-      if (covering_seq > iter->max_covering_tombstone_seq) {
-        iter->max_covering_tombstone_seq = covering_seq;
-        if (iter->timestamp) {
-          // Will be overwritten in SaveValue() if there is a point key with
-          // a higher seqno.
-          iter->timestamp->assign(range_del_iter->timestamp().data(),
-                                  range_del_iter->timestamp().size());
+
+  // Use batch lookup optimization when enabled
+  bool use_batch_optimization = moptions_.memtable_batch_lookup_optimization;
+  bool validate = moptions_.paranoid_memory_checks ||
+                  moptions_.memtable_veirfy_per_key_checksum_on_seek;
+
+  if (use_batch_optimization) {
+    // Phase 1: Handle range tombstones and set up Savers for batched lookup
+    std::array<Saver, MultiGetContext::MAX_BATCH_SIZE> savers{};
+    std::array<const char*, MultiGetContext::MAX_BATCH_SIZE> memtable_keys{};
+    std::array<void*, MultiGetContext::MAX_BATCH_SIZE> callback_args{};
+    std::array<bool, MultiGetContext::MAX_BATCH_SIZE> found_final_values{};
+    std::array<bool, MultiGetContext::MAX_BATCH_SIZE> merge_in_progresses{};
+    size_t num_keys = 0;
+
+    for (auto iter = temp_range.begin(); iter != temp_range.end(); ++iter) {
+      if (!no_range_del) {
+        std::unique_ptr<FragmentedRangeTombstoneIterator> range_del_iter(
+            NewRangeTombstoneIteratorInternal(
+                read_options, GetInternalKeySeqno(iter->lkey->internal_key()),
+                immutable_memtable));
+        SequenceNumber covering_seq =
+            range_del_iter->MaxCoveringTombstoneSeqnum(iter->lkey->user_key());
+        if (covering_seq > iter->max_covering_tombstone_seq) {
+          iter->max_covering_tombstone_seq = covering_seq;
+          if (iter->timestamp) {
+            iter->timestamp->assign(range_del_iter->timestamp().data(),
+                                    range_del_iter->timestamp().size());
+          }
         }
       }
+
+      found_final_values[num_keys] = false;
+      merge_in_progresses[num_keys] = iter->s->IsMergeInProgress();
+      Saver& saver = savers[num_keys];
+      saver.status = iter->s;
+      saver.found_final_value = &found_final_values[num_keys];
+      saver.merge_in_progress = &merge_in_progresses[num_keys];
+      saver.key = iter->lkey;
+      saver.value = iter->value ? iter->value->GetSelf() : nullptr;
+      saver.columns = iter->columns;
+      saver.timestamp = iter->timestamp;
+      saver.seq = kMaxSequenceNumber;
+      saver.mem = this;
+      saver.merge_context = &(iter->merge_context);
+      saver.max_covering_tombstone_seq = iter->max_covering_tombstone_seq;
+      saver.merge_operator = moptions_.merge_operator;
+      saver.logger = moptions_.info_log;
+      saver.inplace_update_support = moptions_.inplace_update_support;
+      saver.statistics = moptions_.statistics;
+      saver.clock = clock_;
+      saver.callback_ = callback;
+      saver.is_blob_index = &iter->is_blob_index;
+      saver.do_merge = true;
+      saver.allow_data_in_errors = moptions_.allow_data_in_errors;
+      saver.protection_bytes_per_key = moptions_.protection_bytes_per_key;
+
+      memtable_keys[num_keys] = iter->lkey->memtable_key().data();
+      callback_args[num_keys] = &savers[num_keys];
+      num_keys++;
     }
-    SequenceNumber dummy_seq;
-    GetFromTable(*(iter->lkey), iter->max_covering_tombstone_seq, true,
-                 callback, &iter->is_blob_index,
-                 iter->value ? iter->value->GetSelf() : nullptr, iter->columns,
-                 iter->timestamp, iter->s, &(iter->merge_context), &dummy_seq,
-                 &found_final_value, &merge_in_progress);
-
-    if (!found_final_value && merge_in_progress) {
-      if (iter->s->ok()) {
-        *(iter->s) = Status::MergeInProgress();
-      } else {
-        assert(iter->s->IsMergeInProgress());
+
+    // Phase 2: Batched lookup
+    if (num_keys > 0) {
+      Status check_s =
+          table_->MultiGet(num_keys, memtable_keys.data(), callback_args.data(),
+                           SaveValue, moptions_.allow_data_in_errors,
+                           validate ? moptions_.paranoid_memory_checks : false,
+                           validate ? key_validation_callback_ : nullptr);
+      if (check_s.IsCorruption()) {
+        // Mark all remaining keys as corruption
+        for (auto iter = temp_range.begin(); iter != temp_range.end(); ++iter) {
+          *(iter->s) = check_s;
+          range->MarkKeyDone(iter);
+        }
+        return;
       }
     }
 
-    if (found_final_value ||
-        (!iter->s->ok() && !iter->s->IsMergeInProgress())) {
-      // `found_final_value` should be set if an error/corruption occurs.
-      // The check on iter->s is just there in case GetFromTable() did not
-      // set `found_final_value` properly.
-      assert(found_final_value);
-      if (iter->value) {
-        iter->value->PinSelf();
-        range->AddValueSize(iter->value->size());
-      } else {
-        assert(iter->columns);
-        range->AddValueSize(iter->columns->serialized_size());
+    // Phase 3: Process results
+    size_t result_idx = 0;
+    for (auto iter = temp_range.begin(); iter != temp_range.end();
+         ++iter, ++result_idx) {
+      bool found_final_value = found_final_values[result_idx];
+      bool merge_in_progress = merge_in_progresses[result_idx];
+
+      if (!found_final_value && merge_in_progress) {
+        if (iter->s->ok()) {
+          *(iter->s) = Status::MergeInProgress();
+        } else {
+          assert(iter->s->IsMergeInProgress());
+        }
+      }
+
+      if (found_final_value ||
+          (!iter->s->ok() && !iter->s->IsMergeInProgress())) {
+        assert(found_final_value);
+        if (iter->value) {
+          iter->value->PinSelf();
+          range->AddValueSize(iter->value->size());
+        } else {
+          assert(iter->columns);
+          range->AddValueSize(iter->columns->serialized_size());
+        }
+
+        range->MarkKeyDone(iter);
+        RecordTick(moptions_.statistics, MEMTABLE_HIT);
+        if (range->GetValueSize() > read_options.value_size_soft_limit) {
+          for (auto range_iter = range->begin(); range_iter != range->end();
+               ++range_iter) {
+            range->MarkKeyDone(range_iter);
+            *(range_iter->s) = Status::Aborted();
+          }
+          break;
+        }
+      }
+    }
+  } else {
+    // Per-key lookup path
+    for (auto iter = temp_range.begin(); iter != temp_range.end(); ++iter) {
+      bool found_final_value{false};
+      bool merge_in_progress = iter->s->IsMergeInProgress();
+      if (!no_range_del) {
+        std::unique_ptr<FragmentedRangeTombstoneIterator> range_del_iter(
+            NewRangeTombstoneIteratorInternal(
+                read_options, GetInternalKeySeqno(iter->lkey->internal_key()),
+                immutable_memtable));
+        SequenceNumber covering_seq =
+            range_del_iter->MaxCoveringTombstoneSeqnum(iter->lkey->user_key());
+        if (covering_seq > iter->max_covering_tombstone_seq) {
+          iter->max_covering_tombstone_seq = covering_seq;
+          if (iter->timestamp) {
+            iter->timestamp->assign(range_del_iter->timestamp().data(),
+                                    range_del_iter->timestamp().size());
+          }
+        }
+      }
+      SequenceNumber dummy_seq;
+      GetFromTable(
+          *(iter->lkey), iter->max_covering_tombstone_seq, true, callback,
+          &iter->is_blob_index, iter->value ? iter->value->GetSelf() : nullptr,
+          iter->columns, iter->timestamp, iter->s, &(iter->merge_context),
+          &dummy_seq, &found_final_value, &merge_in_progress);
+
+      if (!found_final_value && merge_in_progress) {
+        if (iter->s->ok()) {
+          *(iter->s) = Status::MergeInProgress();
+        } else {
+          assert(iter->s->IsMergeInProgress());
+        }
       }
 
-      range->MarkKeyDone(iter);
-      RecordTick(moptions_.statistics, MEMTABLE_HIT);
-      if (range->GetValueSize() > read_options.value_size_soft_limit) {
-        // Set all remaining keys in range to Abort
-        for (auto range_iter = range->begin(); range_iter != range->end();
-             ++range_iter) {
-          range->MarkKeyDone(range_iter);
-          *(range_iter->s) = Status::Aborted();
+      if (found_final_value ||
+          (!iter->s->ok() && !iter->s->IsMergeInProgress())) {
+        assert(found_final_value);
+        if (iter->value) {
+          iter->value->PinSelf();
+          range->AddValueSize(iter->value->size());
+        } else {
+          assert(iter->columns);
+          range->AddValueSize(iter->columns->serialized_size());
+        }
+
+        range->MarkKeyDone(iter);
+        RecordTick(moptions_.statistics, MEMTABLE_HIT);
+        if (range->GetValueSize() > read_options.value_size_soft_limit) {
+          for (auto range_iter = range->begin(); range_iter != range->end();
+               ++range_iter) {
+            range->MarkKeyDone(range_iter);
+            *(range_iter->s) = Status::Aborted();
+          }
+          break;
         }
-        break;
       }
     }
   }
@@ -1841,6 +1957,36 @@ void MemTableRep::Get(const LookupKey& k, void* callback_args,
   }
 }
 
+Status MemTableRep::MultiGet(
+    size_t num_keys, const char* const* keys, void** callback_args,
+    bool (*callback_func)(void* arg, const char* entry),
+    bool allow_data_in_errors, bool detect_key_out_of_order,
+    const std::function<Status(const char*, bool)>& key_validation_callback) {
+  bool validate = detect_key_out_of_order || key_validation_callback != nullptr;
+  std::unique_ptr<Iterator> iter(GetIterator());
+  for (size_t i = 0; i < num_keys; ++i) {
+    Slice dummy;
+    if (validate) {
+      Status s = iter->SeekAndValidate(dummy, keys[i], allow_data_in_errors,
+                                       detect_key_out_of_order,
+                                       key_validation_callback);
+      for (; iter->Valid() && s.ok() &&
+             callback_func(callback_args[i], iter->key());
+           s = iter->NextAndValidate(allow_data_in_errors)) {
+      }
+      if (!s.ok()) {
+        return s;
+      }
+    } else {
+      iter->Seek(dummy, keys[i]);
+      for (; iter->Valid() && callback_func(callback_args[i], iter->key());
+           iter->Next()) {
+      }
+    }
+  }
+  return Status::OK();
+}
+
 void MemTable::RefLogContainingPrepSection(uint64_t log) {
   assert(log > 0);
   auto cur = min_prep_log_referenced_.load();

package/deps/rocksdb/rocksdb/db/memtable.h

@@ -65,6 +65,7 @@ struct ImmutableMemTableOptions {
   bool allow_data_in_errors;
   bool paranoid_memory_checks;
   bool memtable_veirfy_per_key_checksum_on_seek;
+  bool memtable_batch_lookup_optimization;
 };
 
 // Batched counters to updated when inserting keys in one write batch.
@@ -599,6 +600,11 @@ class MemTable final : public ReadOnlyMemTable {
                                                 std::memory_order_relaxed);
   }
 
+  // Returns true if a flush has already been scheduled for this memtable
+  bool HasFlushScheduled() const {
+    return flush_state_.load(std::memory_order_relaxed) == FLUSH_SCHEDULED;
+  }
+
   InternalIterator* NewIterator(
       const ReadOptions& read_options,
       UnownedPtr<const SeqnoToTimeMapping> seqno_to_time_mapping, Arena* arena,

package/deps/rocksdb/rocksdb/db/memtable_list_test.cc

@@ -33,12 +33,12 @@ std::string ValueWithWriteTime(std::string value, uint64_t write_time) {
 class MemTableListTest : public testing::Test {
  public:
   std::string dbname;
-  DB* db;
+  std::unique_ptr<DB> db;
   Options options;
   std::vector<ColumnFamilyHandle*> handles;
   std::atomic<uint64_t> file_number;
 
-  MemTableListTest() : db(nullptr), file_number(1) {
+  MemTableListTest() : file_number(1) {
     dbname = test::PerThreadDBPath("memtable_list_test");
     options.create_if_missing = true;
     EXPECT_OK(DestroyDB(dbname, options));
@@ -88,8 +88,7 @@ class MemTableListTest : public testing::Test {
         }
       }
       handles.clear();
-      delete db;
-      db = nullptr;
+      db.reset();
       EXPECT_OK(DestroyDB(dbname, options, cf_descs));
     }
   }