Add the PutEntity API to the stress/crash tests (facebook#10760)

Summary:
The patch adds the `PutEntity` API to the non-batched, batched, and
CF consistency stress tests. Namely, when the new `db_stress` command
line parameter `use_put_entity_one_in` is greater than zero, one in
N writes on average is performed using `PutEntity` rather than `Put`.
The wide-column entity written has the generated value in its default
column; in addition, it contains up to three additional columns where
the original generated value is divided up between the column name and the
column value (with the column name containing the first k characters of
the generated value, and the column value containing the rest). Whether
`PutEntity` is used (and if so, how many columns the entity has) is completely
determined by the "value base" used to generate the value (that is, there is
no randomness involved). Assuming the same `use_put_entity_one_in` setting
is used across `db_stress` invocations, this enables us to reconstruct and
validate the entity during subsequent `db_stress` runs.

Note that `PutEntity` is currently incompatible with `Merge`, transactions, and
user-defined timestamps; these combinations are currently disabled/disallowed.

Pull Request resolved: facebook#10760

Test Plan: Ran some batched, non-batched, and CF consistency stress tests using the script.

Reviewed By: riversand963

Differential Revision: D39939032

Pulled By: ltamasi

fbshipit-source-id: eafdf124e95993fb7d73158e3b006d11819f7fa9

db_stress_tool/batched_ops_stress.cc

@@ -26,31 +26,41 @@ class BatchedOpsStressTest : public StressTest {
                  const std::vector<int>& rand_column_families,
                  const std::vector<int64_t>& rand_keys,
                  char (&value)[100]) override {
-    uint32_t value_base =
+    assert(!rand_column_families.empty());
+    assert(!rand_keys.empty());
+
+    const std::string key_suffix = Key(rand_keys[0]);
+
+    const uint32_t value_base =
         thread->rand.Next() % thread->shared->UNKNOWN_SENTINEL;
-    size_t sz = GenerateValue(value_base, value, sizeof(value));
-    Slice v(value, sz);
-    std::string keys[10] = {"9", "8", "7", "6", "5", "4", "3", "2", "1", "0"};
-    std::string values[10] = {"9", "8", "7", "6", "5", "4", "3", "2", "1", "0"};
-    Slice value_slices[10];
+    const size_t sz = GenerateValue(value_base, value, sizeof(value));
+    const std::string value_suffix = Slice(value, sz).ToString();
+
     WriteBatch batch(0 /* reserved_bytes */, 0 /* max_bytes */,
                      FLAGS_batch_protection_bytes_per_key,
                      FLAGS_user_timestamp_size);
-    Status s;
-    auto cfh = column_families_[rand_column_families[0]];
-    std::string key_str = Key(rand_keys[0]);
-    for (int i = 0; i < 10; i++) {
-      keys[i] += key_str;
-      values[i] += v.ToString();
-      value_slices[i] = values[i];
+
+    ColumnFamilyHandle* const cfh = column_families_[rand_column_families[0]];
+    assert(cfh);
+
+    for (int i = 9; i >= 0; --i) {
+      const std::string prefix = std::to_string(i);
+
+      const std::string k = prefix + key_suffix;
+      const std::string v = prefix + value_suffix;
+
       if (FLAGS_use_merge) {
-        batch.Merge(cfh, keys[i], value_slices[i]);
+        batch.Merge(cfh, k, v);
+      } else if (FLAGS_use_put_entity_one_in > 0 &&
+                 (value_base % FLAGS_use_put_entity_one_in) == 0) {
+        batch.PutEntity(cfh, k, GenerateWideColumns(value_base, v));
       } else {
-        batch.Put(cfh, keys[i], value_slices[i]);
+        batch.Put(cfh, k, v);
       }
     }
 
-    s = db_->Write(write_opts, &batch);
+    const Status s = db_->Write(write_opts, &batch);
+
     if (!s.ok()) {
       fprintf(stderr, "multiput error: %s\n", s.ToString().c_str());
       thread->stats.AddErrors(1);

db_stress_tool/cf_consistency_stress.cc

@@ -25,22 +25,36 @@ class CfConsistencyStressTest : public StressTest {
                  const std::vector<int>& rand_column_families,
                  const std::vector<int64_t>& rand_keys,
                  char (&value)[100]) override {
-    std::string key_str = Key(rand_keys[0]);
-    Slice key = key_str;
-    uint64_t value_base = batch_id_.fetch_add(1);
-    size_t sz =
-        GenerateValue(static_cast<uint32_t>(value_base), value, sizeof(value));
-    Slice v(value, sz);
+    assert(!rand_column_families.empty());
+    assert(!rand_keys.empty());
+
+    const std::string k = Key(rand_keys[0]);
+
+    const uint32_t value_base = batch_id_.fetch_add(1);
+    const size_t sz = GenerateValue(value_base, value, sizeof(value));
+    const Slice v(value, sz);
+
     WriteBatch batch;
+
+    const bool use_put_entity = !FLAGS_use_merge &&
+                                FLAGS_use_put_entity_one_in > 0 &&
+                                (value_base % FLAGS_use_put_entity_one_in) == 0;
+
     for (auto cf : rand_column_families) {
-      ColumnFamilyHandle* cfh = column_families_[cf];
+      ColumnFamilyHandle* const cfh = column_families_[cf];
+      assert(cfh);
+
       if (FLAGS_use_merge) {
-        batch.Merge(cfh, key, v);
-      } else { /* !FLAGS_use_merge */
-        batch.Put(cfh, key, v);
+        batch.Merge(cfh, k, v);
+      } else if (use_put_entity) {
+        batch.PutEntity(cfh, k, GenerateWideColumns(value_base, v));
+      } else {
+        batch.Put(cfh, k, v);
       }
     }
+
     Status s = db_->Write(write_opts, &batch);
+
     if (!s.ok()) {
       fprintf(stderr, "multi put or merge error: %s\n", s.ToString().c_str());
       thread->stats.AddErrors(1);
@@ -538,7 +552,7 @@ class CfConsistencyStressTest : public StressTest {
   }
 
  private:
-  std::atomic<int64_t> batch_id_;
+  std::atomic<uint32_t> batch_id_;
 };
 
 StressTest* CreateCfConsistencyStressTest() {

db_stress_tool/db_stress_common.cc

@@ -239,6 +239,40 @@ uint32_t GetValueBase(Slice s) {
   return res;
 }
 
+WideColumns GenerateWideColumns(uint32_t value_base, const Slice& slice) {
+  WideColumns columns;
+
+  constexpr size_t max_columns = 4;
+  const size_t num_columns = (value_base % max_columns) + 1;
+
+  columns.reserve(num_columns);
+
+  assert(slice.size() >= num_columns);
+
+  columns.emplace_back(kDefaultWideColumnName, slice);
+
+  for (size_t i = 1; i < num_columns; ++i) {
+    const Slice name(slice.data(), i);
+    const Slice value(slice.data() + i, slice.size() - i);
+
+    columns.emplace_back(name, value);
+  }
+
+  return columns;
+}
+
+WideColumns GenerateExpectedWideColumns(uint32_t value_base,
+                                        const Slice& slice) {
+  WideColumns columns = GenerateWideColumns(value_base, slice);
+
+  std::sort(columns.begin(), columns.end(),
+            [](const WideColumn& lhs, const WideColumn& rhs) {
+              return lhs.name().compare(rhs.name()) < 0;
+            });
+
+  return columns;
+}
+
 std::string GetNowNanos() {
   uint64_t t = db_stress_env->NowNanos();
   std::string ret;

db_stress_tool/db_stress_common.h

@@ -235,6 +235,7 @@ DECLARE_bool(in_place_update);
 DECLARE_string(memtablerep);
 DECLARE_int32(prefix_size);
 DECLARE_bool(use_merge);
+DECLARE_uint32(use_put_entity_one_in);
 DECLARE_bool(use_full_merge_v1);
 DECLARE_int32(sync_wal_one_in);
 DECLARE_bool(avoid_unnecessary_blocking_io);
@@ -620,6 +621,10 @@ extern std::vector<int64_t> GenerateNKeys(ThreadState* thread, int num_keys,
 extern size_t GenerateValue(uint32_t rand, char* v, size_t max_sz);
 extern uint32_t GetValueBase(Slice s);
 
+extern WideColumns GenerateWideColumns(uint32_t value_base, const Slice& slice);
+extern WideColumns GenerateExpectedWideColumns(uint32_t value_base,
+                                               const Slice& slice);
+
 extern StressTest* CreateCfConsistencyStressTest();
 extern StressTest* CreateBatchedOpsStressTest();
 extern StressTest* CreateNonBatchedOpsStressTest();

db_stress_tool/db_stress_gflags.cc

@@ -858,6 +858,10 @@ DEFINE_bool(use_merge, false,
             "On true, replaces all writes with a Merge "
             "that behaves like a Put");
 
+DEFINE_uint32(use_put_entity_one_in, 0,
+              "If greater than zero, PutEntity will be used once per every N "
+              "write ops on average.");
+
 DEFINE_bool(use_full_merge_v1, false,
             "On true, use a merge operator that implement the deprecated "
             "version of FullMerge");

db_stress_tool/db_stress_test_base.cc

@@ -457,11 +457,14 @@ void StressTest::PreloadDbAndReopenAsReadOnly(int64_t number_of_keys,
   Status s;
   for (auto cfh : column_families_) {
     for (int64_t k = 0; k != number_of_keys; ++k) {
-      std::string key_str = Key(k);
-      Slice key = key_str;
-      size_t sz = GenerateValue(0 /*value_base*/, value, sizeof(value));
-      Slice v(value, sz);
-      shared->Put(cf_idx, k, 0, true /* pending */);
+      const std::string key = Key(k);
+
+      constexpr uint32_t value_base = 0;
+      const size_t sz = GenerateValue(value_base, value, sizeof(value));
+
+      const Slice v(value, sz);
+
+      shared->Put(cf_idx, k, value_base, true /* pending */);
 
       if (FLAGS_use_merge) {
         if (!FLAGS_use_txn) {
@@ -478,13 +481,13 @@ void StressTest::PreloadDbAndReopenAsReadOnly(int64_t number_of_keys,
           }
 #endif
         }
+      } else if (FLAGS_use_put_entity_one_in > 0) {
+        s = db_->PutEntity(write_opts, cfh, key,
+                           GenerateWideColumns(value_base, v));
       } else {
         if (!FLAGS_use_txn) {
-          std::string ts_str;
-          Slice ts;
           if (FLAGS_user_timestamp_size > 0) {
-            ts_str = GetNowNanos();
-            ts = ts_str;
+            const std::string ts = GetNowNanos();
             s = db_->Put(write_opts, cfh, key, ts, v);
           } else {
             s = db_->Put(write_opts, cfh, key, v);
@@ -503,7 +506,7 @@ void StressTest::PreloadDbAndReopenAsReadOnly(int64_t number_of_keys,
         }
       }
 
-      shared->Put(cf_idx, k, 0, false /* pending */);
+      shared->Put(cf_idx, k, value_base, false /* pending */);
       if (!s.ok()) {
         break;
       }

db_stress_tool/db_stress_tool.cc

@@ -294,6 +294,15 @@ int db_stress_tool(int argc, char** argv) {
     exit(1);
   }
 
+  if (FLAGS_use_put_entity_one_in > 0 &&
+      (FLAGS_use_merge || FLAGS_use_full_merge_v1 || FLAGS_use_txn ||
+       FLAGS_test_multi_ops_txns || FLAGS_user_timestamp_size > 0)) {
+    fprintf(stderr,
+            "PutEntity is currently incompatible with Merge, transactions, and "
+            "user-defined timestamps\n");
+    exit(1);
+  }
+
 #ifndef NDEBUG
   KillPoint* kp = KillPoint::GetInstance();
   kp->rocksdb_kill_odds = FLAGS_kill_random_test;

db_stress_tool/expected_state.cc

@@ -7,6 +7,7 @@
 
 #include "db_stress_tool/expected_state.h"
 
+#include "db/wide/wide_column_serialization.h"
 #include "db_stress_tool/db_stress_common.h"
 #include "db_stress_tool/db_stress_shared_state.h"
 #include "rocksdb/trace_reader_writer.h"
@@ -405,6 +406,50 @@ class ExpectedStateTraceRecordHandler : public TraceRecord::Handler,
     return Status::OK();
   }
 
+  Status PutEntityCF(uint32_t column_family_id, const Slice& key_with_ts,
+                     const Slice& entity) override {
+    Slice key =
+        StripTimestampFromUserKey(key_with_ts, FLAGS_user_timestamp_size);
+
+    uint64_t key_id = 0;
+    if (!GetIntVal(key.ToString(), &key_id)) {
+      return Status::Corruption("Unable to parse key", key.ToString());
+    }
+
+    Slice entity_copy = entity;
+    WideColumns columns;
+    if (!WideColumnSerialization::Deserialize(entity_copy, columns).ok()) {
+      return Status::Corruption("Unable to deserialize entity",
+                                entity.ToString(/* hex */ true));
+    }
+
+    if (columns.empty() || columns[0].name() != kDefaultWideColumnName) {
+      return Status::Corruption("Cannot find default column in entity",
+                                entity.ToString(/* hex */ true));
+    }
+
+    const Slice& value_of_default = columns[0].value();
+
+    const uint32_t value_base = GetValueBase(value_of_default);
+
+    if (columns != GenerateExpectedWideColumns(value_base, value_of_default)) {
+      return Status::Corruption("Wide columns in entity inconsistent",
+                                entity.ToString(/* hex */ true));
+    }
+
+    if (buffered_writes_) {
+      return WriteBatchInternal::PutEntity(buffered_writes_.get(),
+                                           column_family_id, key, columns);
+    }
+
+    state_->Put(column_family_id, static_cast<int64_t>(key_id), value_base,
+                false /* pending */);
+
+    ++num_write_ops_;
+
+    return Status::OK();
+  }
+
   Status DeleteCF(uint32_t column_family_id,
                   const Slice& key_with_ts) override {
     Slice key =