Fix a bug in GetMergeOperands() with continue_cb set (#13383)

Summary:
Noticed this while I was working on memtable code. Wrong status (MergeInProgress()) and wrong number of merge operands can be returned if the `continue_cb` stop at an immutable memtable. This is due to

1. Get from memtable sets MergeInProress() status https://github.com/facebook/rocksdb/blob/302254d928402c80ecabb0cf2b76d18be35a87b9/db/memtable.cc#L1461
2. Get from immutable memtable does not update status but stops the get:  https://github.com/facebook/rocksdb/blob/302254d928402c80ecabb0cf2b76d18be35a87b9/db/memtable.cc#L1364
3. GetImpl() only returns merge_operands for OK status: https://github.com/facebook/rocksdb/blob/302254d928402c80ecabb0cf2b76d18be35a87b9/db/db_impl/db_impl.cc#L2552

Also updated some comments for GetMergeOperands().

Pull Request resolved: https://github.com/facebook/rocksdb/pull/13383

Test Plan: added a unit test that fails GetMergeOperands() with MergeInProgress() status before this fix.

Reviewed By: jowlyzhang

Differential Revision: D69322133

Pulled By: cbi42

fbshipit-source-id: aebfccd8446e9640cff02877915076e2d10f7a5b

db/db_impl/db_impl.cc

@@ -2566,6 +2566,8 @@ Status DBImpl::GetImpl(const ReadOptions& read_options, const Slice& key,
         // Return all merge operands for get_impl_options.key
         *get_impl_options.number_of_operands =
             static_cast<int>(merge_context.GetNumOperands());
+        // OK status is returned, some merge operand is found.
+        assert(*get_impl_options.number_of_operands > 0);
         if (*get_impl_options.number_of_operands >
             get_impl_options.get_merge_operands_options
                 ->expected_max_number_of_operands) {

db/db_merge_operand_test.cc

@@ -596,6 +596,42 @@ TEST_F(DBMergeOperandTest, GetMergeOperandsBaseDeletionInImmMem) {
   }
 }
 
+TEST_F(DBMergeOperandTest, GetMergeOperandCallbackStopAtImm) {
+  Options options = CurrentOptions();
+  options.max_write_buffer_number = 10;
+  options.merge_operator = MergeOperators::CreateStringAppendOperator();
+  DestroyAndReopen(options);
+
+  Random rnd(301);
+  ASSERT_OK(db_->PauseBackgroundWork());
+  ASSERT_OK(Merge("key", "v1"));
+  ASSERT_OK(dbfull()->TEST_SwitchMemtable());
+  // Keep this merge in an immutable memtable
+  uint64_t num_imm = 0;
+  ASSERT_TRUE(
+      db_->GetIntProperty(DB::Properties::kNumImmutableMemTable, &num_imm));
+  ASSERT_EQ(num_imm, 1);
+  ASSERT_OK(Merge("key", "v2"));
+
+  std::vector<PinnableSlice> merge_operands(2);
+  GetMergeOperandsOptions merge_operands_info;
+  merge_operands_info.expected_max_number_of_operands = 2;
+  int num_fetched = 0;
+  merge_operands_info.continue_cb = [&num_fetched](Slice /* value */) {
+    num_fetched++;
+    // Stop in the first immutable memtable.
+    return num_fetched < 2;
+  };
+  int num_merge_operands = 0;
+  ASSERT_OK(db_->GetMergeOperands(ReadOptions(), db_->DefaultColumnFamily(),
+                                  "key", merge_operands.data(),
+                                  &merge_operands_info, &num_merge_operands));
+  ASSERT_EQ(2, num_merge_operands);
+  ASSERT_EQ(2, num_fetched);
+
+  ASSERT_EQ("v1", merge_operands[0]);
+  ASSERT_EQ("v2", merge_operands[1]);
+}
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {

db/memtable.cc

@@ -1360,7 +1360,10 @@ static bool SaveValue(void* arg, const char* entry) {
         if (merge_context->get_merge_operands_options != nullptr &&
             merge_context->get_merge_operands_options->continue_cb != nullptr &&
             !merge_context->get_merge_operands_options->continue_cb(v)) {
-          // We were told not to continue.
+          // We were told not to continue. `status` may be MergeInProress(),
+          // overwrite to signal the end of successful get. This status
+          // will be checked at the end of GetImpl().
+          *(s->status) = Status::OK();
           *(s->found_final_value) = true;
           return false;
         }

include/rocksdb/db.h

@@ -764,8 +764,9 @@ class DB {
   // Populates the `merge_operands` array with all the merge operands in the DB
   // for `key`, or a customizable suffix of merge operands when
   // `GetMergeOperandsOptions::continue_cb` is set. The `merge_operands` array
-  // will be populated in the order of insertion. The number of entries
-  // populated in `merge_operands` will be assigned to `*number_of_operands`.
+  // will be populated in the order of insertion (older insertions first). The
+  // number of entries populated in `merge_operands` will be assigned to
+  // `*number_of_operands`.
   //
   // If the number of merge operands to return for `key` is greater than
   // `merge_operands_options.expected_max_number_of_operands`,
@@ -780,6 +781,9 @@ class DB {
   // The caller should delete or `Reset()` the `merge_operands` entries when
   // they are no longer needed. All `merge_operands` entries must be destroyed
   // or `Reset()` before this DB is closed or destroyed.
+  // OK status is returned if any merge operand is found.
+  // NotFound status is returned if no merge operand is found.
+  // Error status is returned if there is an error.
   virtual Status GetMergeOperands(
       const ReadOptions& options, ColumnFamilyHandle* column_family,
       const Slice& key, PinnableSlice* merge_operands,

unreleased_history/bug_fixes/get_merge_operand_imm.md

@@ -0,0 +1 @@
+* Fix a bug in `GetMergeOperands()` that can return incorrect status (MergeInProgress) and incorrect number of merge operands. This can happen when `GetMergeOperandsOptions::continue_cb` is set, both active and immutable memtables have merge operands and the callback stops the look up at the immutable memtable.