ceph/src/rocksdb/db/write_batch_test.cc

   1 //  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
   2 //  This source code is licensed under both the GPLv2 (found in the
   3 //  COPYING file in the root directory) and Apache 2.0 License
   4 //  (found in the LICENSE.Apache file in the root directory).
   5 //
   6 // Copyright (c) 2011 The LevelDB Authors. All rights reserved.
   7 // Use of this source code is governed by a BSD-style license that can be
   8 // found in the LICENSE file. See the AUTHORS file for names of contributors.
   9
  10 #include "rocksdb/db.h"
  11
  12 #include <memory>
  13 #include "db/column_family.h"
  14 #include "db/memtable.h"
  15 #include "db/write_batch_internal.h"
  16 #include "rocksdb/env.h"
  17 #include "rocksdb/memtablerep.h"
  18 #include "rocksdb/utilities/write_batch_with_index.h"
  19 #include "rocksdb/write_buffer_manager.h"
  20 #include "table/scoped_arena_iterator.h"
  21 #include "util/string_util.h"
  22 #include "util/testharness.h"
  23
  24 namespace rocksdb {
  25
  26 static std::string PrintContents(WriteBatch* b) {
  27   InternalKeyComparator cmp(BytewiseComparator());
  28   auto factory = std::make_shared<SkipListFactory>();
  29   Options options;
  30   options.memtable_factory = factory;
  31   ImmutableCFOptions ioptions(options);
  32   WriteBufferManager wb(options.db_write_buffer_size);
  33   MemTable* mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb,
  34                                kMaxSequenceNumber, 0 /* column_family_id */);
  35   mem->Ref();
  36   std::string state;
  37   ColumnFamilyMemTablesDefault cf_mems_default(mem);
  38   Status s = WriteBatchInternal::InsertInto(b, &cf_mems_default, nullptr);
  39   int count = 0;
  40   int put_count = 0;
  41   int delete_count = 0;
  42   int single_delete_count = 0;
  43   int delete_range_count = 0;
  44   int merge_count = 0;
  45   for (int i = 0; i < 2; ++i) {
  46     Arena arena;
  47     ScopedArenaIterator arena_iter_guard;
  48     std::unique_ptr<InternalIterator> iter_guard;
  49     InternalIterator* iter;
  50     if (i == 0) {
  51       iter = mem->NewIterator(ReadOptions(), &arena);
  52       arena_iter_guard.set(iter);
  53     } else {
  54       iter = mem->NewRangeTombstoneIterator(ReadOptions(),
  55                                             kMaxSequenceNumber /* read_seq */);
  56       iter_guard.reset(iter);
  57     }
  58     if (iter == nullptr) {
  59       continue;
  60     }
  61     for (iter->SeekToFirst(); iter->Valid(); iter->Next()) {
  62       ParsedInternalKey ikey;
  63       ikey.clear();
  64       EXPECT_TRUE(ParseInternalKey(iter->key(), &ikey));
  65       switch (ikey.type) {
  66         case kTypeValue:
  67           state.append("Put(");
  68           state.append(ikey.user_key.ToString());
  69           state.append(", ");
  70           state.append(iter->value().ToString());
  71           state.append(")");
  72           count++;
  73           put_count++;
  74           break;
  75         case kTypeDeletion:
  76           state.append("Delete(");
  77           state.append(ikey.user_key.ToString());
  78           state.append(")");
  79           count++;
  80           delete_count++;
  81           break;
  82         case kTypeSingleDeletion:
  83           state.append("SingleDelete(");
  84           state.append(ikey.user_key.ToString());
  85           state.append(")");
  86           count++;
  87           single_delete_count++;
  88           break;
  89         case kTypeRangeDeletion:
  90           state.append("DeleteRange(");
  91           state.append(ikey.user_key.ToString());
  92           state.append(", ");
  93           state.append(iter->value().ToString());
  94           state.append(")");
  95           count++;
  96           delete_range_count++;
  97           break;
  98         case kTypeMerge:
  99           state.append("Merge(");
 100           state.append(ikey.user_key.ToString());
 101           state.append(", ");
 102           state.append(iter->value().ToString());
 103           state.append(")");
 104           count++;
 105           merge_count++;
 106           break;
 107         default:
 108           assert(false);
 109           break;
 110       }
 111       state.append("@");
 112       state.append(NumberToString(ikey.sequence));
 113     }
 114   }
 115   EXPECT_EQ(b->HasPut(), put_count > 0);
 116   EXPECT_EQ(b->HasDelete(), delete_count > 0);
 117   EXPECT_EQ(b->HasSingleDelete(), single_delete_count > 0);
 118   EXPECT_EQ(b->HasDeleteRange(), delete_range_count > 0);
 119   EXPECT_EQ(b->HasMerge(), merge_count > 0);
 120   if (!s.ok()) {
 121     state.append(s.ToString());
 122   } else if (count != WriteBatchInternal::Count(b)) {
 123     state.append("CountMismatch()");
 124   }
 125   delete mem->Unref();
 126   return state;
 127 }
 128
 129 class WriteBatchTest : public testing::Test {};
 130
 131 TEST_F(WriteBatchTest, Empty) {
 132   WriteBatch batch;
 133   ASSERT_EQ("", PrintContents(&batch));
 134   ASSERT_EQ(0, WriteBatchInternal::Count(&batch));
 135   ASSERT_EQ(0, batch.Count());
 136 }
 137
 138 TEST_F(WriteBatchTest, Multiple) {
 139   WriteBatch batch;
 140   batch.Put(Slice("foo"), Slice("bar"));
 141   batch.Delete(Slice("box"));
 142   batch.DeleteRange(Slice("bar"), Slice("foo"));
 143   batch.Put(Slice("baz"), Slice("boo"));
 144   WriteBatchInternal::SetSequence(&batch, 100);
 145   ASSERT_EQ(100U, WriteBatchInternal::Sequence(&batch));
 146   ASSERT_EQ(4, WriteBatchInternal::Count(&batch));
 147   ASSERT_EQ(
 148       "Put(baz, boo)@103"
 149       "Delete(box)@101"
 150       "Put(foo, bar)@100"
 151       "DeleteRange(bar, foo)@102",
 152       PrintContents(&batch));
 153   ASSERT_EQ(4, batch.Count());
 154 }
 155
 156 TEST_F(WriteBatchTest, Corruption) {
 157   WriteBatch batch;
 158   batch.Put(Slice("foo"), Slice("bar"));
 159   batch.Delete(Slice("box"));
 160   WriteBatchInternal::SetSequence(&batch, 200);
 161   Slice contents = WriteBatchInternal::Contents(&batch);
 162   WriteBatchInternal::SetContents(&batch,
 163                                   Slice(contents.data(),contents.size()-1));
 164   ASSERT_EQ("Put(foo, bar)@200"
 165             "Corruption: bad WriteBatch Delete",
 166             PrintContents(&batch));
 167 }
 168
 169 TEST_F(WriteBatchTest, Append) {
 170   WriteBatch b1, b2;
 171   WriteBatchInternal::SetSequence(&b1, 200);
 172   WriteBatchInternal::SetSequence(&b2, 300);
 173   WriteBatchInternal::Append(&b1, &b2);
 174   ASSERT_EQ("",
 175             PrintContents(&b1));
 176   ASSERT_EQ(0, b1.Count());
 177   b2.Put("a", "va");
 178   WriteBatchInternal::Append(&b1, &b2);
 179   ASSERT_EQ("Put(a, va)@200",
 180             PrintContents(&b1));
 181   ASSERT_EQ(1, b1.Count());
 182   b2.Clear();
 183   b2.Put("b", "vb");
 184   WriteBatchInternal::Append(&b1, &b2);
 185   ASSERT_EQ("Put(a, va)@200"
 186             "Put(b, vb)@201",
 187             PrintContents(&b1));
 188   ASSERT_EQ(2, b1.Count());
 189   b2.Delete("foo");
 190   WriteBatchInternal::Append(&b1, &b2);
 191   ASSERT_EQ("Put(a, va)@200"
 192             "Put(b, vb)@202"
 193             "Put(b, vb)@201"
 194             "Delete(foo)@203",
 195             PrintContents(&b1));
 196   ASSERT_EQ(4, b1.Count());
 197   b2.Clear();
 198   b2.Put("c", "cc");
 199   b2.Put("d", "dd");
 200   b2.MarkWalTerminationPoint();
 201   b2.Put("e", "ee");
 202   WriteBatchInternal::Append(&b1, &b2, /*wal only*/ true);
 203   ASSERT_EQ(
 204       "Put(a, va)@200"
 205       "Put(b, vb)@202"
 206       "Put(b, vb)@201"
 207       "Put(c, cc)@204"
 208       "Put(d, dd)@205"
 209       "Delete(foo)@203",
 210       PrintContents(&b1));
 211   ASSERT_EQ(6, b1.Count());
 212   ASSERT_EQ(
 213       "Put(c, cc)@0"
 214       "Put(d, dd)@1"
 215       "Put(e, ee)@2",
 216       PrintContents(&b2));
 217   ASSERT_EQ(3, b2.Count());
 218 }
 219
 220 TEST_F(WriteBatchTest, SingleDeletion) {
 221   WriteBatch batch;
 222   WriteBatchInternal::SetSequence(&batch, 100);
 223   ASSERT_EQ("", PrintContents(&batch));
 224   ASSERT_EQ(0, batch.Count());
 225   batch.Put("a", "va");
 226   ASSERT_EQ("Put(a, va)@100", PrintContents(&batch));
 227   ASSERT_EQ(1, batch.Count());
 228   batch.SingleDelete("a");
 229   ASSERT_EQ(
 230       "SingleDelete(a)@101"
 231       "Put(a, va)@100",
 232       PrintContents(&batch));
 233   ASSERT_EQ(2, batch.Count());
 234 }
 235
 236 namespace {
 237   struct TestHandler : public WriteBatch::Handler {
 238     std::string seen;
 239     Status PutCF(uint32_t column_family_id, const Slice& key,
 240                  const Slice& value) override {
 241       if (column_family_id == 0) {
 242         seen += "Put(" + key.ToString() + ", " + value.ToString() + ")";
 243       } else {
 244         seen += "PutCF(" + ToString(column_family_id) + ", " +
 245                 key.ToString() + ", " + value.ToString() + ")";
 246       }
 247       return Status::OK();
 248     }
 249     Status DeleteCF(uint32_t column_family_id, const Slice& key) override {
 250       if (column_family_id == 0) {
 251         seen += "Delete(" + key.ToString() + ")";
 252       } else {
 253         seen += "DeleteCF(" + ToString(column_family_id) + ", " +
 254                 key.ToString() + ")";
 255       }
 256       return Status::OK();
 257     }
 258     Status SingleDeleteCF(uint32_t column_family_id,
 259                           const Slice& key) override {
 260       if (column_family_id == 0) {
 261         seen += "SingleDelete(" + key.ToString() + ")";
 262       } else {
 263         seen += "SingleDeleteCF(" + ToString(column_family_id) + ", " +
 264                 key.ToString() + ")";
 265       }
 266       return Status::OK();
 267     }
 268     Status DeleteRangeCF(uint32_t column_family_id, const Slice& begin_key,
 269                          const Slice& end_key) override {
 270       if (column_family_id == 0) {
 271         seen += "DeleteRange(" + begin_key.ToString() + ", " +
 272                 end_key.ToString() + ")";
 273       } else {
 274         seen += "DeleteRangeCF(" + ToString(column_family_id) + ", " +
 275                 begin_key.ToString() + ", " + end_key.ToString() + ")";
 276       }
 277       return Status::OK();
 278     }
 279     Status MergeCF(uint32_t column_family_id, const Slice& key,
 280                    const Slice& value) override {
 281       if (column_family_id == 0) {
 282         seen += "Merge(" + key.ToString() + ", " + value.ToString() + ")";
 283       } else {
 284         seen += "MergeCF(" + ToString(column_family_id) + ", " +
 285                 key.ToString() + ", " + value.ToString() + ")";
 286       }
 287       return Status::OK();
 288     }
 289     void LogData(const Slice& blob) override {
 290       seen += "LogData(" + blob.ToString() + ")";
 291     }
 292     Status MarkBeginPrepare(bool unprepare) override {
 293       seen +=
 294           "MarkBeginPrepare(" + std::string(unprepare ? "true" : "false") + ")";
 295       return Status::OK();
 296     }
 297     Status MarkEndPrepare(const Slice& xid) override {
 298       seen += "MarkEndPrepare(" + xid.ToString() + ")";
 299       return Status::OK();
 300     }
 301     Status MarkNoop(bool empty_batch) override {
 302       seen += "MarkNoop(" + std::string(empty_batch ? "true" : "false") + ")";
 303       return Status::OK();
 304     }
 305     Status MarkCommit(const Slice& xid) override {
 306       seen += "MarkCommit(" + xid.ToString() + ")";
 307       return Status::OK();
 308     }
 309     Status MarkRollback(const Slice& xid) override {
 310       seen += "MarkRollback(" + xid.ToString() + ")";
 311       return Status::OK();
 312     }
 313   };
 314 }
 315
 316 TEST_F(WriteBatchTest, PutNotImplemented) {
 317   WriteBatch batch;
 318   batch.Put(Slice("k1"), Slice("v1"));
 319   ASSERT_EQ(1, batch.Count());
 320   ASSERT_EQ("Put(k1, v1)@0", PrintContents(&batch));
 321
 322   WriteBatch::Handler handler;
 323   ASSERT_OK(batch.Iterate(&handler));
 324 }
 325
 326 TEST_F(WriteBatchTest, DeleteNotImplemented) {
 327   WriteBatch batch;
 328   batch.Delete(Slice("k2"));
 329   ASSERT_EQ(1, batch.Count());
 330   ASSERT_EQ("Delete(k2)@0", PrintContents(&batch));
 331
 332   WriteBatch::Handler handler;
 333   ASSERT_OK(batch.Iterate(&handler));
 334 }
 335
 336 TEST_F(WriteBatchTest, SingleDeleteNotImplemented) {
 337   WriteBatch batch;
 338   batch.SingleDelete(Slice("k2"));
 339   ASSERT_EQ(1, batch.Count());
 340   ASSERT_EQ("SingleDelete(k2)@0", PrintContents(&batch));
 341
 342   WriteBatch::Handler handler;
 343   ASSERT_OK(batch.Iterate(&handler));
 344 }
 345
 346 TEST_F(WriteBatchTest, MergeNotImplemented) {
 347   WriteBatch batch;
 348   batch.Merge(Slice("foo"), Slice("bar"));
 349   ASSERT_EQ(1, batch.Count());
 350   ASSERT_EQ("Merge(foo, bar)@0", PrintContents(&batch));
 351
 352   WriteBatch::Handler handler;
 353   ASSERT_OK(batch.Iterate(&handler));
 354 }
 355
 356 TEST_F(WriteBatchTest, Blob) {
 357   WriteBatch batch;
 358   batch.Put(Slice("k1"), Slice("v1"));
 359   batch.Put(Slice("k2"), Slice("v2"));
 360   batch.Put(Slice("k3"), Slice("v3"));
 361   batch.PutLogData(Slice("blob1"));
 362   batch.Delete(Slice("k2"));
 363   batch.SingleDelete(Slice("k3"));
 364   batch.PutLogData(Slice("blob2"));
 365   batch.Merge(Slice("foo"), Slice("bar"));
 366   ASSERT_EQ(6, batch.Count());
 367   ASSERT_EQ(
 368       "Merge(foo, bar)@5"
 369       "Put(k1, v1)@0"
 370       "Delete(k2)@3"
 371       "Put(k2, v2)@1"
 372       "SingleDelete(k3)@4"
 373       "Put(k3, v3)@2",
 374       PrintContents(&batch));
 375
 376   TestHandler handler;
 377   batch.Iterate(&handler);
 378   ASSERT_EQ(
 379       "Put(k1, v1)"
 380       "Put(k2, v2)"
 381       "Put(k3, v3)"
 382       "LogData(blob1)"
 383       "Delete(k2)"
 384       "SingleDelete(k3)"
 385       "LogData(blob2)"
 386       "Merge(foo, bar)",
 387       handler.seen);
 388 }
 389
 390 TEST_F(WriteBatchTest, PrepareCommit) {
 391   WriteBatch batch;
 392   WriteBatchInternal::InsertNoop(&batch);
 393   batch.Put(Slice("k1"), Slice("v1"));
 394   batch.Put(Slice("k2"), Slice("v2"));
 395   batch.SetSavePoint();
 396   WriteBatchInternal::MarkEndPrepare(&batch, Slice("xid1"));
 397   Status s = batch.RollbackToSavePoint();
 398   ASSERT_EQ(s, Status::NotFound());
 399   WriteBatchInternal::MarkCommit(&batch, Slice("xid1"));
 400   WriteBatchInternal::MarkRollback(&batch, Slice("xid1"));
 401   ASSERT_EQ(2, batch.Count());
 402
 403   TestHandler handler;
 404   batch.Iterate(&handler);
 405   ASSERT_EQ(
 406       "MarkBeginPrepare(false)"
 407       "Put(k1, v1)"
 408       "Put(k2, v2)"
 409       "MarkEndPrepare(xid1)"
 410       "MarkCommit(xid1)"
 411       "MarkRollback(xid1)",
 412       handler.seen);
 413 }
 414
 415 // It requires more than 30GB of memory to run the test. With single memory
 416 // allocation of more than 30GB.
 417 // Not all platform can run it. Also it runs a long time. So disable it.
 418 TEST_F(WriteBatchTest, DISABLED_ManyUpdates) {
 419   // Insert key and value of 3GB and push total batch size to 12GB.
 420   static const size_t kKeyValueSize = 4u;
 421   static const uint32_t kNumUpdates = uint32_t(3 << 30);
 422   std::string raw(kKeyValueSize, 'A');
 423   WriteBatch batch(kNumUpdates * (4 + kKeyValueSize * 2) + 1024u);
 424   char c = 'A';
 425   for (uint32_t i = 0; i < kNumUpdates; i++) {
 426     if (c > 'Z') {
 427       c = 'A';
 428     }
 429     raw[0] = c;
 430     raw[raw.length() - 1] = c;
 431     c++;
 432     batch.Put(raw, raw);
 433   }
 434
 435   ASSERT_EQ(kNumUpdates, batch.Count());
 436
 437   struct NoopHandler : public WriteBatch::Handler {
 438     uint32_t num_seen = 0;
 439     char expected_char = 'A';
 440     Status PutCF(uint32_t /*column_family_id*/, const Slice& key,
 441                  const Slice& value) override {
 442       EXPECT_EQ(kKeyValueSize, key.size());
 443       EXPECT_EQ(kKeyValueSize, value.size());
 444       EXPECT_EQ(expected_char, key[0]);
 445       EXPECT_EQ(expected_char, value[0]);
 446       EXPECT_EQ(expected_char, key[kKeyValueSize - 1]);
 447       EXPECT_EQ(expected_char, value[kKeyValueSize - 1]);
 448       expected_char++;
 449       if (expected_char > 'Z') {
 450         expected_char = 'A';
 451       }
 452       ++num_seen;
 453       return Status::OK();
 454     }
 455     Status DeleteCF(uint32_t /*column_family_id*/,
 456                     const Slice& /*key*/) override {
 457       ADD_FAILURE();
 458       return Status::OK();
 459     }
 460     Status SingleDeleteCF(uint32_t /*column_family_id*/,
 461                           const Slice& /*key*/) override {
 462       ADD_FAILURE();
 463       return Status::OK();
 464     }
 465     Status MergeCF(uint32_t /*column_family_id*/, const Slice& /*key*/,
 466                    const Slice& /*value*/) override {
 467       ADD_FAILURE();
 468       return Status::OK();
 469     }
 470     void LogData(const Slice& /*blob*/) override { ADD_FAILURE(); }
 471     bool Continue() override { return num_seen < kNumUpdates; }
 472   } handler;
 473
 474   batch.Iterate(&handler);
 475   ASSERT_EQ(kNumUpdates, handler.num_seen);
 476 }
 477
 478 // The test requires more than 18GB memory to run it, with single memory
 479 // allocation of more than 12GB. Not all the platform can run it. So disable it.
 480 TEST_F(WriteBatchTest, DISABLED_LargeKeyValue) {
 481   // Insert key and value of 3GB and push total batch size to 12GB.
 482   static const size_t kKeyValueSize = 3221225472u;
 483   std::string raw(kKeyValueSize, 'A');
 484   WriteBatch batch(size_t(12884901888ull + 1024u));
 485   for (char i = 0; i < 2; i++) {
 486     raw[0] = 'A' + i;
 487     raw[raw.length() - 1] = 'A' - i;
 488     batch.Put(raw, raw);
 489   }
 490
 491   ASSERT_EQ(2, batch.Count());
 492
 493   struct NoopHandler : public WriteBatch::Handler {
 494     int num_seen = 0;
 495     Status PutCF(uint32_t /*column_family_id*/, const Slice& key,
 496                  const Slice& value) override {
 497       EXPECT_EQ(kKeyValueSize, key.size());
 498       EXPECT_EQ(kKeyValueSize, value.size());
 499       EXPECT_EQ('A' + num_seen, key[0]);
 500       EXPECT_EQ('A' + num_seen, value[0]);
 501       EXPECT_EQ('A' - num_seen, key[kKeyValueSize - 1]);
 502       EXPECT_EQ('A' - num_seen, value[kKeyValueSize - 1]);
 503       ++num_seen;
 504       return Status::OK();
 505     }
 506     Status DeleteCF(uint32_t /*column_family_id*/,
 507                     const Slice& /*key*/) override {
 508       ADD_FAILURE();
 509       return Status::OK();
 510     }
 511     Status SingleDeleteCF(uint32_t /*column_family_id*/,
 512                           const Slice& /*key*/) override {
 513       ADD_FAILURE();
 514       return Status::OK();
 515     }
 516     Status MergeCF(uint32_t /*column_family_id*/, const Slice& /*key*/,
 517                    const Slice& /*value*/) override {
 518       ADD_FAILURE();
 519       return Status::OK();
 520     }
 521     void LogData(const Slice& /*blob*/) override { ADD_FAILURE(); }
 522     bool Continue() override { return num_seen < 2; }
 523   } handler;
 524
 525   batch.Iterate(&handler);
 526   ASSERT_EQ(2, handler.num_seen);
 527 }
 528
 529 TEST_F(WriteBatchTest, Continue) {
 530   WriteBatch batch;
 531
 532   struct Handler : public TestHandler {
 533     int num_seen = 0;
 534     Status PutCF(uint32_t column_family_id, const Slice& key,
 535                  const Slice& value) override {
 536       ++num_seen;
 537       return TestHandler::PutCF(column_family_id, key, value);
 538     }
 539     Status DeleteCF(uint32_t column_family_id, const Slice& key) override {
 540       ++num_seen;
 541       return TestHandler::DeleteCF(column_family_id, key);
 542     }
 543     Status SingleDeleteCF(uint32_t column_family_id,
 544                           const Slice& key) override {
 545       ++num_seen;
 546       return TestHandler::SingleDeleteCF(column_family_id, key);
 547     }
 548     Status MergeCF(uint32_t column_family_id, const Slice& key,
 549                    const Slice& value) override {
 550       ++num_seen;
 551       return TestHandler::MergeCF(column_family_id, key, value);
 552     }
 553     void LogData(const Slice& blob) override {
 554       ++num_seen;
 555       TestHandler::LogData(blob);
 556     }
 557     bool Continue() override { return num_seen < 5; }
 558   } handler;
 559
 560   batch.Put(Slice("k1"), Slice("v1"));
 561   batch.Put(Slice("k2"), Slice("v2"));
 562   batch.PutLogData(Slice("blob1"));
 563   batch.Delete(Slice("k1"));
 564   batch.SingleDelete(Slice("k2"));
 565   batch.PutLogData(Slice("blob2"));
 566   batch.Merge(Slice("foo"), Slice("bar"));
 567   batch.Iterate(&handler);
 568   ASSERT_EQ(
 569       "Put(k1, v1)"
 570       "Put(k2, v2)"
 571       "LogData(blob1)"
 572       "Delete(k1)"
 573       "SingleDelete(k2)",
 574       handler.seen);
 575 }
 576
 577 TEST_F(WriteBatchTest, PutGatherSlices) {
 578   WriteBatch batch;
 579   batch.Put(Slice("foo"), Slice("bar"));
 580
 581   {
 582     // Try a write where the key is one slice but the value is two
 583     Slice key_slice("baz");
 584     Slice value_slices[2] = { Slice("header"), Slice("payload") };
 585     batch.Put(SliceParts(&key_slice, 1),
 586               SliceParts(value_slices, 2));
 587   }
 588
 589   {
 590     // One where the key is composite but the value is a single slice
 591     Slice key_slices[3] = { Slice("key"), Slice("part2"), Slice("part3") };
 592     Slice value_slice("value");
 593     batch.Put(SliceParts(key_slices, 3),
 594               SliceParts(&value_slice, 1));
 595   }
 596
 597   WriteBatchInternal::SetSequence(&batch, 100);
 598   ASSERT_EQ("Put(baz, headerpayload)@101"
 599             "Put(foo, bar)@100"
 600             "Put(keypart2part3, value)@102",
 601             PrintContents(&batch));
 602   ASSERT_EQ(3, batch.Count());
 603 }
 604
 605 namespace {
 606 class ColumnFamilyHandleImplDummy : public ColumnFamilyHandleImpl {
 607  public:
 608   explicit ColumnFamilyHandleImplDummy(int id)
 609       : ColumnFamilyHandleImpl(nullptr, nullptr, nullptr), id_(id) {}
 610   uint32_t GetID() const override { return id_; }
 611   const Comparator* GetComparator() const override {
 612     return BytewiseComparator();
 613   }
 614
 615  private:
 616   uint32_t id_;
 617 };
 618 }  // namespace anonymous
 619
 620 TEST_F(WriteBatchTest, ColumnFamiliesBatchTest) {
 621   WriteBatch batch;
 622   ColumnFamilyHandleImplDummy zero(0), two(2), three(3), eight(8);
 623   batch.Put(&zero, Slice("foo"), Slice("bar"));
 624   batch.Put(&two, Slice("twofoo"), Slice("bar2"));
 625   batch.Put(&eight, Slice("eightfoo"), Slice("bar8"));
 626   batch.Delete(&eight, Slice("eightfoo"));
 627   batch.SingleDelete(&two, Slice("twofoo"));
 628   batch.DeleteRange(&two, Slice("3foo"), Slice("4foo"));
 629   batch.Merge(&three, Slice("threethree"), Slice("3three"));
 630   batch.Put(&zero, Slice("foo"), Slice("bar"));
 631   batch.Merge(Slice("omom"), Slice("nom"));
 632
 633   TestHandler handler;
 634   batch.Iterate(&handler);
 635   ASSERT_EQ(
 636       "Put(foo, bar)"
 637       "PutCF(2, twofoo, bar2)"
 638       "PutCF(8, eightfoo, bar8)"
 639       "DeleteCF(8, eightfoo)"
 640       "SingleDeleteCF(2, twofoo)"
 641       "DeleteRangeCF(2, 3foo, 4foo)"
 642       "MergeCF(3, threethree, 3three)"
 643       "Put(foo, bar)"
 644       "Merge(omom, nom)",
 645       handler.seen);
 646 }
 647
 648 #ifndef ROCKSDB_LITE
 649 TEST_F(WriteBatchTest, ColumnFamiliesBatchWithIndexTest) {
 650   WriteBatchWithIndex batch;
 651   ColumnFamilyHandleImplDummy zero(0), two(2), three(3), eight(8);
 652   batch.Put(&zero, Slice("foo"), Slice("bar"));
 653   batch.Put(&two, Slice("twofoo"), Slice("bar2"));
 654   batch.Put(&eight, Slice("eightfoo"), Slice("bar8"));
 655   batch.Delete(&eight, Slice("eightfoo"));
 656   batch.SingleDelete(&two, Slice("twofoo"));
 657   batch.DeleteRange(&two, Slice("twofoo"), Slice("threefoo"));
 658   batch.Merge(&three, Slice("threethree"), Slice("3three"));
 659   batch.Put(&zero, Slice("foo"), Slice("bar"));
 660   batch.Merge(Slice("omom"), Slice("nom"));
 661
 662   std::unique_ptr<WBWIIterator> iter;
 663
 664   iter.reset(batch.NewIterator(&eight));
 665   iter->Seek("eightfoo");
 666   ASSERT_OK(iter->status());
 667   ASSERT_TRUE(iter->Valid());
 668   ASSERT_EQ(WriteType::kPutRecord, iter->Entry().type);
 669   ASSERT_EQ("eightfoo", iter->Entry().key.ToString());
 670   ASSERT_EQ("bar8", iter->Entry().value.ToString());
 671
 672   iter->Next();
 673   ASSERT_OK(iter->status());
 674   ASSERT_TRUE(iter->Valid());
 675   ASSERT_EQ(WriteType::kDeleteRecord, iter->Entry().type);
 676   ASSERT_EQ("eightfoo", iter->Entry().key.ToString());
 677
 678   iter->Next();
 679   ASSERT_OK(iter->status());
 680   ASSERT_TRUE(!iter->Valid());
 681
 682   iter.reset(batch.NewIterator(&two));
 683   iter->Seek("twofoo");
 684   ASSERT_OK(iter->status());
 685   ASSERT_TRUE(iter->Valid());
 686   ASSERT_EQ(WriteType::kPutRecord, iter->Entry().type);
 687   ASSERT_EQ("twofoo", iter->Entry().key.ToString());
 688   ASSERT_EQ("bar2", iter->Entry().value.ToString());
 689
 690   iter->Next();
 691   ASSERT_OK(iter->status());
 692   ASSERT_TRUE(iter->Valid());
 693   ASSERT_EQ(WriteType::kSingleDeleteRecord, iter->Entry().type);
 694   ASSERT_EQ("twofoo", iter->Entry().key.ToString());
 695
 696   iter->Next();
 697   ASSERT_OK(iter->status());
 698   ASSERT_TRUE(iter->Valid());
 699   ASSERT_EQ(WriteType::kDeleteRangeRecord, iter->Entry().type);
 700   ASSERT_EQ("twofoo", iter->Entry().key.ToString());
 701   ASSERT_EQ("threefoo", iter->Entry().value.ToString());
 702
 703   iter->Next();
 704   ASSERT_OK(iter->status());
 705   ASSERT_TRUE(!iter->Valid());
 706
 707   iter.reset(batch.NewIterator());
 708   iter->Seek("gggg");
 709   ASSERT_OK(iter->status());
 710   ASSERT_TRUE(iter->Valid());
 711   ASSERT_EQ(WriteType::kMergeRecord, iter->Entry().type);
 712   ASSERT_EQ("omom", iter->Entry().key.ToString());
 713   ASSERT_EQ("nom", iter->Entry().value.ToString());
 714
 715   iter->Next();
 716   ASSERT_OK(iter->status());
 717   ASSERT_TRUE(!iter->Valid());
 718
 719   iter.reset(batch.NewIterator(&zero));
 720   iter->Seek("foo");
 721   ASSERT_OK(iter->status());
 722   ASSERT_TRUE(iter->Valid());
 723   ASSERT_EQ(WriteType::kPutRecord, iter->Entry().type);
 724   ASSERT_EQ("foo", iter->Entry().key.ToString());
 725   ASSERT_EQ("bar", iter->Entry().value.ToString());
 726
 727   iter->Next();
 728   ASSERT_OK(iter->status());
 729   ASSERT_TRUE(iter->Valid());
 730   ASSERT_EQ(WriteType::kPutRecord, iter->Entry().type);
 731   ASSERT_EQ("foo", iter->Entry().key.ToString());
 732   ASSERT_EQ("bar", iter->Entry().value.ToString());
 733
 734   iter->Next();
 735   ASSERT_OK(iter->status());
 736   ASSERT_TRUE(iter->Valid());
 737   ASSERT_EQ(WriteType::kMergeRecord, iter->Entry().type);
 738   ASSERT_EQ("omom", iter->Entry().key.ToString());
 739   ASSERT_EQ("nom", iter->Entry().value.ToString());
 740
 741   iter->Next();
 742   ASSERT_OK(iter->status());
 743   ASSERT_TRUE(!iter->Valid());
 744
 745   TestHandler handler;
 746   batch.GetWriteBatch()->Iterate(&handler);
 747   ASSERT_EQ(
 748       "Put(foo, bar)"
 749       "PutCF(2, twofoo, bar2)"
 750       "PutCF(8, eightfoo, bar8)"
 751       "DeleteCF(8, eightfoo)"
 752       "SingleDeleteCF(2, twofoo)"
 753       "DeleteRangeCF(2, twofoo, threefoo)"
 754       "MergeCF(3, threethree, 3three)"
 755       "Put(foo, bar)"
 756       "Merge(omom, nom)",
 757       handler.seen);
 758 }
 759 #endif  // !ROCKSDB_LITE
 760
 761 TEST_F(WriteBatchTest, SavePointTest) {
 762   Status s;
 763   WriteBatch batch;
 764   batch.SetSavePoint();
 765
 766   batch.Put("A", "a");
 767   batch.Put("B", "b");
 768   batch.SetSavePoint();
 769
 770   batch.Put("C", "c");
 771   batch.Delete("A");
 772   batch.SetSavePoint();
 773   batch.SetSavePoint();
 774
 775   ASSERT_OK(batch.RollbackToSavePoint());
 776   ASSERT_EQ(
 777       "Delete(A)@3"
 778       "Put(A, a)@0"
 779       "Put(B, b)@1"
 780       "Put(C, c)@2",
 781       PrintContents(&batch));
 782
 783   ASSERT_OK(batch.RollbackToSavePoint());
 784   ASSERT_OK(batch.RollbackToSavePoint());
 785   ASSERT_EQ(
 786       "Put(A, a)@0"
 787       "Put(B, b)@1",
 788       PrintContents(&batch));
 789
 790   batch.Delete("A");
 791   batch.Put("B", "bb");
 792
 793   ASSERT_OK(batch.RollbackToSavePoint());
 794   ASSERT_EQ("", PrintContents(&batch));
 795
 796   s = batch.RollbackToSavePoint();
 797   ASSERT_TRUE(s.IsNotFound());
 798   ASSERT_EQ("", PrintContents(&batch));
 799
 800   batch.Put("D", "d");
 801   batch.Delete("A");
 802
 803   batch.SetSavePoint();
 804
 805   batch.Put("A", "aaa");
 806
 807   ASSERT_OK(batch.RollbackToSavePoint());
 808   ASSERT_EQ(
 809       "Delete(A)@1"
 810       "Put(D, d)@0",
 811       PrintContents(&batch));
 812
 813   batch.SetSavePoint();
 814
 815   batch.Put("D", "d");
 816   batch.Delete("A");
 817
 818   ASSERT_OK(batch.RollbackToSavePoint());
 819   ASSERT_EQ(
 820       "Delete(A)@1"
 821       "Put(D, d)@0",
 822       PrintContents(&batch));
 823
 824   s = batch.RollbackToSavePoint();
 825   ASSERT_TRUE(s.IsNotFound());
 826   ASSERT_EQ(
 827       "Delete(A)@1"
 828       "Put(D, d)@0",
 829       PrintContents(&batch));
 830
 831   WriteBatch batch2;
 832
 833   s = batch2.RollbackToSavePoint();
 834   ASSERT_TRUE(s.IsNotFound());
 835   ASSERT_EQ("", PrintContents(&batch2));
 836
 837   batch2.Delete("A");
 838   batch2.SetSavePoint();
 839
 840   s = batch2.RollbackToSavePoint();
 841   ASSERT_OK(s);
 842   ASSERT_EQ("Delete(A)@0", PrintContents(&batch2));
 843
 844   batch2.Clear();
 845   ASSERT_EQ("", PrintContents(&batch2));
 846
 847   batch2.SetSavePoint();
 848
 849   batch2.Delete("B");
 850   ASSERT_EQ("Delete(B)@0", PrintContents(&batch2));
 851
 852   batch2.SetSavePoint();
 853   s = batch2.RollbackToSavePoint();
 854   ASSERT_OK(s);
 855   ASSERT_EQ("Delete(B)@0", PrintContents(&batch2));
 856
 857   s = batch2.RollbackToSavePoint();
 858   ASSERT_OK(s);
 859   ASSERT_EQ("", PrintContents(&batch2));
 860
 861   s = batch2.RollbackToSavePoint();
 862   ASSERT_TRUE(s.IsNotFound());
 863   ASSERT_EQ("", PrintContents(&batch2));
 864
 865   WriteBatch batch3;
 866
 867   s = batch3.PopSavePoint();
 868   ASSERT_TRUE(s.IsNotFound());
 869   ASSERT_EQ("", PrintContents(&batch3));
 870
 871   batch3.SetSavePoint();
 872   batch3.Delete("A");
 873
 874   s = batch3.PopSavePoint();
 875   ASSERT_OK(s);
 876   ASSERT_EQ("Delete(A)@0", PrintContents(&batch3));
 877 }
 878
 879 TEST_F(WriteBatchTest, MemoryLimitTest) {
 880   Status s;
 881   // The header size is 12 bytes. The two Puts take 8 bytes which gives total
 882   // of 12 + 8 * 2 = 28 bytes.
 883   WriteBatch batch(0, 28);
 884
 885   ASSERT_OK(batch.Put("a", "...."));
 886   ASSERT_OK(batch.Put("b", "...."));
 887   s = batch.Put("c", "....");
 888   ASSERT_TRUE(s.IsMemoryLimit());
 889 }
 890
 891 }  // namespace rocksdb
 892
 893 int main(int argc, char** argv) {
 894   ::testing::InitGoogleTest(&argc, argv);
 895   return RUN_ALL_TESTS();
 896 }