Support using MultiGetEntity as verification method in stress tests (#11228)

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

Reviewed By: akankshamahajan15

Differential Revision: D43332120

Pulled By: ltamasi

fbshipit-source-id: 15f32cf335aecb7e654da24ecafc6e010dc65194

db_stress_tool/no_batched_ops_stress.cc

@@ -53,6 +53,7 @@ class NonBatchedOpsStressTest : public StressTest {
         kGet,
         kGetEntity,
         kMultiGet,
+        kMultiGetEntity,
         kGetMergeOperands,
         // Add any new items above kNumberOfMethods
         kNumberOfMethods
@@ -163,11 +164,11 @@ class NonBatchedOpsStressTest : public StressTest {
           Status s =
               db_->GetEntity(options, column_families_[cf], key, &columns);
 
-          const WideColumns& columns_from_db = columns.columns();
-
           std::string from_db;
 
           if (s.ok()) {
+            const WideColumns& columns_from_db = columns.columns();
+
             if (!columns_from_db.empty() &&
                 columns_from_db[0].name() == kDefaultWideColumnName) {
               from_db = columns_from_db[0].value().ToString();
@@ -200,14 +201,14 @@ class NonBatchedOpsStressTest : public StressTest {
           size_t batch_size = thread->rand.Uniform(128) + 1;
           batch_size = std::min<size_t>(batch_size, end - i);
 
-          std::vector<std::string> keystrs(batch_size);
+          std::vector<std::string> key_strs(batch_size);
           std::vector<Slice> keys(batch_size);
           std::vector<PinnableSlice> values(batch_size);
           std::vector<Status> statuses(batch_size);
 
           for (size_t j = 0; j < batch_size; ++j) {
-            keystrs[j] = Key(i + j);
-            keys[j] = Slice(keystrs[j].data(), keystrs[j].size());
+            key_strs[j] = Key(i + j);
+            keys[j] = Slice(key_strs[j]);
           }
 
           db_->MultiGet(options, column_families_[cf], batch_size, keys.data(),
@@ -226,6 +227,61 @@ class NonBatchedOpsStressTest : public StressTest {
             }
           }
 
+          i += batch_size;
+        }
+      } else if (method == VerificationMethod::kMultiGetEntity) {
+        for (int64_t i = start; i < end;) {
+          if (thread->shared->HasVerificationFailedYet()) {
+            break;
+          }
+
+          // Keep the batch size to some reasonable value
+          size_t batch_size = thread->rand.Uniform(128) + 1;
+          batch_size = std::min<size_t>(batch_size, end - i);
+
+          std::vector<std::string> key_strs(batch_size);
+          std::vector<Slice> keys(batch_size);
+          std::vector<PinnableWideColumns> results(batch_size);
+          std::vector<Status> statuses(batch_size);
+
+          for (size_t j = 0; j < batch_size; ++j) {
+            key_strs[j] = Key(i + j);
+            keys[j] = Slice(key_strs[j]);
+          }
+
+          db_->MultiGetEntity(options, column_families_[cf], batch_size,
+                              keys.data(), results.data(), statuses.data());
+
+          for (size_t j = 0; j < batch_size; ++j) {
+            std::string from_db;
+
+            if (statuses[j].ok()) {
+              const WideColumns& columns_from_db = results[j].columns();
+
+              if (!columns_from_db.empty() &&
+                  columns_from_db[0].name() == kDefaultWideColumnName) {
+                from_db = columns_from_db[0].value().ToString();
+              }
+
+              const WideColumns expected_columns =
+                  GenerateExpectedWideColumns(GetValueBase(from_db), from_db);
+              if (columns_from_db != expected_columns) {
+                VerificationAbort(shared, static_cast<int>(cf), i, from_db,
+                                  columns_from_db, expected_columns);
+              }
+            }
+
+            VerifyOrSyncValue(static_cast<int>(cf), i + j, options, shared,
+                              from_db,
+                              /* msg_prefix */ "MultiGetEntity verification",
+                              statuses[j], /* strict */ true);
+
+            if (!from_db.empty()) {
+              PrintKeyValue(static_cast<int>(cf), static_cast<uint32_t>(i + j),
+                            from_db.data(), from_db.size());
+            }
+          }
+
           i += batch_size;
         }
       } else {