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 the BSD-style license found in the
   3 //  LICENSE file in the root directory of this source tree. An additional grant
   4 //  of patent rights can be found in the PATENTS file in the same 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/logging.h"
  22 #include "util/string_util.h"
  23 #include "util/testharness.h"
  24
  25 namespace rocksdb {
  26
  27 static std::string PrintContents(WriteBatch* b) {
  28   InternalKeyComparator cmp(BytewiseComparator());
  29   auto factory = std::make_shared<SkipListFactory>();
  30   Options options;
  31   options.memtable_factory = factory;
  32   ImmutableCFOptions ioptions(options);
  33   WriteBufferManager wb(options.db_write_buffer_size);
  34   MemTable* mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb,
  35                                kMaxSequenceNumber);
  36   mem->Ref();
  37   std::string state;
  38   ColumnFamilyMemTablesDefault cf_mems_default(mem);
  39   Status s = WriteBatchInternal::InsertInto(b, &cf_mems_default, nullptr);
  40   int count = 0;
  41   int put_count = 0;
  42   int delete_count = 0;
  43   int single_delete_count = 0;
  44   int delete_range_count = 0;
  45   int merge_count = 0;
  46   for (int i = 0; i < 2; ++i) {
  47     Arena arena;
  48     ScopedArenaIterator arena_iter_guard;
  49     std::unique_ptr<InternalIterator> iter_guard;
  50     InternalIterator* iter;
  51     if (i == 0) {
  52       iter = mem->NewIterator(ReadOptions(), &arena);
  53       arena_iter_guard.set(iter);
  54     } else {
  55       iter = mem->NewRangeTombstoneIterator(ReadOptions());
  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       memset((void*)&ikey, 0, sizeof(ikey));
  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     virtual 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     virtual Status DeleteCF(uint32_t column_family_id,
 250                             const Slice& key) override {
 251       if (column_family_id == 0) {
 252         seen += "Delete(" + key.ToString() + ")";
 253       } else {
 254         seen += "DeleteCF(" + ToString(column_family_id) + ", " +
 255                 key.ToString() + ")";
 256       }
 257       return Status::OK();
 258     }
 259     virtual Status SingleDeleteCF(uint32_t column_family_id,
 260                                   const Slice& key) override {
 261       if (column_family_id == 0) {
 262         seen += "SingleDelete(" + key.ToString() + ")";
 263       } else {
 264         seen += "SingleDeleteCF(" + ToString(column_family_id) + ", " +
 265                 key.ToString() + ")";
 266       }
 267       return Status::OK();
 268     }
 269     virtual Status DeleteRangeCF(uint32_t column_family_id,
 270                                  const Slice& begin_key,
 271                                  const Slice& end_key) override {
 272       if (column_family_id == 0) {
 273         seen += "DeleteRange(" + begin_key.ToString() + ", " +
 274                 end_key.ToString() + ")";
 275       } else {
 276         seen += "DeleteRangeCF(" + ToString(column_family_id) + ", " +
 277                 begin_key.ToString() + ", " + end_key.ToString() + ")";
 278       }
 279       return Status::OK();
 280     }
 281     virtual Status MergeCF(uint32_t column_family_id, const Slice& key,
 282                            const Slice& value) override {
 283       if (column_family_id == 0) {
 284         seen += "Merge(" + key.ToString() + ", " + value.ToString() + ")";
 285       } else {
 286         seen += "MergeCF(" + ToString(column_family_id) + ", " +
 287                 key.ToString() + ", " + value.ToString() + ")";
 288       }
 289       return Status::OK();
 290     }
 291     virtual void LogData(const Slice& blob) override {
 292       seen += "LogData(" + blob.ToString() + ")";
 293     }
 294     virtual Status MarkBeginPrepare() override {
 295       seen += "MarkBeginPrepare()";
 296       return Status::OK();
 297     }
 298     virtual Status MarkEndPrepare(const Slice& xid) override {
 299       seen += "MarkEndPrepare(" + xid.ToString() + ")";
 300       return Status::OK();
 301     }
 302     virtual Status MarkCommit(const Slice& xid) override {
 303       seen += "MarkCommit(" + xid.ToString() + ")";
 304       return Status::OK();
 305     }
 306     virtual Status MarkRollback(const Slice& xid) override {
 307       seen += "MarkRollback(" + xid.ToString() + ")";
 308       return Status::OK();
 309     }
 310   };
 311 }
 312
 313 TEST_F(WriteBatchTest, PutNotImplemented) {
 314   WriteBatch batch;
 315   batch.Put(Slice("k1"), Slice("v1"));
 316   ASSERT_EQ(1, batch.Count());
 317   ASSERT_EQ("Put(k1, v1)@0", PrintContents(&batch));
 318
 319   WriteBatch::Handler handler;
 320   ASSERT_OK(batch.Iterate(&handler));
 321 }
 322
 323 TEST_F(WriteBatchTest, DeleteNotImplemented) {
 324   WriteBatch batch;
 325   batch.Delete(Slice("k2"));
 326   ASSERT_EQ(1, batch.Count());
 327   ASSERT_EQ("Delete(k2)@0", PrintContents(&batch));
 328
 329   WriteBatch::Handler handler;
 330   ASSERT_OK(batch.Iterate(&handler));
 331 }
 332
 333 TEST_F(WriteBatchTest, SingleDeleteNotImplemented) {
 334   WriteBatch batch;
 335   batch.SingleDelete(Slice("k2"));
 336   ASSERT_EQ(1, batch.Count());
 337   ASSERT_EQ("SingleDelete(k2)@0", PrintContents(&batch));
 338
 339   WriteBatch::Handler handler;
 340   ASSERT_OK(batch.Iterate(&handler));
 341 }
 342
 343 TEST_F(WriteBatchTest, MergeNotImplemented) {
 344   WriteBatch batch;
 345   batch.Merge(Slice("foo"), Slice("bar"));
 346   ASSERT_EQ(1, batch.Count());
 347   ASSERT_EQ("Merge(foo, bar)@0", PrintContents(&batch));
 348
 349   WriteBatch::Handler handler;
 350   ASSERT_OK(batch.Iterate(&handler));
 351 }
 352
 353 TEST_F(WriteBatchTest, Blob) {
 354   WriteBatch batch;
 355   batch.Put(Slice("k1"), Slice("v1"));
 356   batch.Put(Slice("k2"), Slice("v2"));
 357   batch.Put(Slice("k3"), Slice("v3"));
 358   batch.PutLogData(Slice("blob1"));
 359   batch.Delete(Slice("k2"));
 360   batch.SingleDelete(Slice("k3"));
 361   batch.PutLogData(Slice("blob2"));
 362   batch.Merge(Slice("foo"), Slice("bar"));
 363   ASSERT_EQ(6, batch.Count());
 364   ASSERT_EQ(
 365       "Merge(foo, bar)@5"
 366       "Put(k1, v1)@0"
 367       "Delete(k2)@3"
 368       "Put(k2, v2)@1"
 369       "SingleDelete(k3)@4"
 370       "Put(k3, v3)@2",
 371       PrintContents(&batch));
 372
 373   TestHandler handler;
 374   batch.Iterate(&handler);
 375   ASSERT_EQ(
 376       "Put(k1, v1)"
 377       "Put(k2, v2)"
 378       "Put(k3, v3)"
 379       "LogData(blob1)"
 380       "Delete(k2)"
 381       "SingleDelete(k3)"
 382       "LogData(blob2)"
 383       "Merge(foo, bar)",
 384       handler.seen);
 385 }
 386
 387 TEST_F(WriteBatchTest, PrepareCommit) {
 388   WriteBatch batch;
 389   WriteBatchInternal::InsertNoop(&batch);
 390   batch.Put(Slice("k1"), Slice("v1"));
 391   batch.Put(Slice("k2"), Slice("v2"));
 392   batch.SetSavePoint();
 393   WriteBatchInternal::MarkEndPrepare(&batch, Slice("xid1"));
 394   Status s = batch.RollbackToSavePoint();
 395   ASSERT_EQ(s, Status::NotFound());
 396   WriteBatchInternal::MarkCommit(&batch, Slice("xid1"));
 397   WriteBatchInternal::MarkRollback(&batch, Slice("xid1"));
 398   ASSERT_EQ(2, batch.Count());
 399
 400   TestHandler handler;
 401   batch.Iterate(&handler);
 402   ASSERT_EQ(
 403       "MarkBeginPrepare()"
 404       "Put(k1, v1)"
 405       "Put(k2, v2)"
 406       "MarkEndPrepare(xid1)"
 407       "MarkCommit(xid1)"
 408       "MarkRollback(xid1)",
 409       handler.seen);
 410 }
 411
 412 // It requires more than 30GB of memory to run the test. With single memory
 413 // allocation of more than 30GB.
 414 // Not all platform can run it. Also it runs a long time. So disable it.
 415 TEST_F(WriteBatchTest, DISABLED_ManyUpdates) {
 416   // Insert key and value of 3GB and push total batch size to 12GB.
 417   static const size_t kKeyValueSize = 4u;
 418   static const uint32_t kNumUpdates = 3 << 30;
 419   std::string raw(kKeyValueSize, 'A');
 420   WriteBatch batch(kNumUpdates * (4 + kKeyValueSize * 2) + 1024u);
 421   char c = 'A';
 422   for (uint32_t i = 0; i < kNumUpdates; i++) {
 423     if (c > 'Z') {
 424       c = 'A';
 425     }
 426     raw[0] = c;
 427     raw[raw.length() - 1] = c;
 428     c++;
 429     batch.Put(raw, raw);
 430   }
 431
 432   ASSERT_EQ(kNumUpdates, batch.Count());
 433
 434   struct NoopHandler : public WriteBatch::Handler {
 435     uint32_t num_seen = 0;
 436     char expected_char = 'A';
 437     virtual Status PutCF(uint32_t column_family_id, const Slice& key,
 438                          const Slice& value) override {
 439       EXPECT_EQ(kKeyValueSize, key.size());
 440       EXPECT_EQ(kKeyValueSize, value.size());
 441       EXPECT_EQ(expected_char, key[0]);
 442       EXPECT_EQ(expected_char, value[0]);
 443       EXPECT_EQ(expected_char, key[kKeyValueSize - 1]);
 444       EXPECT_EQ(expected_char, value[kKeyValueSize - 1]);
 445       expected_char++;
 446       if (expected_char > 'Z') {
 447         expected_char = 'A';
 448       }
 449       ++num_seen;
 450       return Status::OK();
 451     }
 452     virtual Status DeleteCF(uint32_t column_family_id,
 453                             const Slice& key) override {
 454       EXPECT_TRUE(false);
 455       return Status::OK();
 456     }
 457     virtual Status SingleDeleteCF(uint32_t column_family_id,
 458                                   const Slice& key) override {
 459       EXPECT_TRUE(false);
 460       return Status::OK();
 461     }
 462     virtual Status MergeCF(uint32_t column_family_id, const Slice& key,
 463                            const Slice& value) override {
 464       EXPECT_TRUE(false);
 465       return Status::OK();
 466     }
 467     virtual void LogData(const Slice& blob) override { EXPECT_TRUE(false); }
 468     virtual bool Continue() override { return num_seen < kNumUpdates; }
 469   } handler;
 470
 471   batch.Iterate(&handler);
 472   ASSERT_EQ(kNumUpdates, handler.num_seen);
 473 }
 474
 475 // The test requires more than 18GB memory to run it, with single memory
 476 // allocation of more than 12GB. Not all the platform can run it. So disable it.
 477 TEST_F(WriteBatchTest, DISABLED_LargeKeyValue) {
 478   // Insert key and value of 3GB and push total batch size to 12GB.
 479   static const size_t kKeyValueSize = 3221225472u;
 480   std::string raw(kKeyValueSize, 'A');
 481   WriteBatch batch(size_t(12884901888ull + 1024u));
 482   for (char i = 0; i < 2; i++) {
 483     raw[0] = 'A' + i;
 484     raw[raw.length() - 1] = 'A' - i;
 485     batch.Put(raw, raw);
 486   }
 487
 488   ASSERT_EQ(2, batch.Count());
 489
 490   struct NoopHandler : public WriteBatch::Handler {
 491     int num_seen = 0;
 492     virtual Status PutCF(uint32_t column_family_id, const Slice& key,
 493                          const Slice& value) override {
 494       EXPECT_EQ(kKeyValueSize, key.size());
 495       EXPECT_EQ(kKeyValueSize, value.size());
 496       EXPECT_EQ('A' + num_seen, key[0]);
 497       EXPECT_EQ('A' + num_seen, value[0]);
 498       EXPECT_EQ('A' - num_seen, key[kKeyValueSize - 1]);
 499       EXPECT_EQ('A' - num_seen, value[kKeyValueSize - 1]);
 500       ++num_seen;
 501       return Status::OK();
 502     }
 503     virtual Status DeleteCF(uint32_t column_family_id,
 504                             const Slice& key) override {
 505       EXPECT_TRUE(false);
 506       return Status::OK();
 507     }
 508     virtual Status SingleDeleteCF(uint32_t column_family_id,
 509                                   const Slice& key) override {
 510       EXPECT_TRUE(false);
 511       return Status::OK();
 512     }
 513     virtual Status MergeCF(uint32_t column_family_id, const Slice& key,
 514                            const Slice& value) override {
 515       EXPECT_TRUE(false);
 516       return Status::OK();
 517     }
 518     virtual void LogData(const Slice& blob) override { EXPECT_TRUE(false); }
 519     virtual bool Continue() override { return num_seen < 2; }
 520   } handler;
 521
 522   batch.Iterate(&handler);
 523   ASSERT_EQ(2, handler.num_seen);
 524 }
 525
 526 TEST_F(WriteBatchTest, Continue) {
 527   WriteBatch batch;
 528
 529   struct Handler : public TestHandler {
 530     int num_seen = 0;
 531     virtual Status PutCF(uint32_t column_family_id, const Slice& key,
 532                          const Slice& value) override {
 533       ++num_seen;
 534       return TestHandler::PutCF(column_family_id, key, value);
 535     }
 536     virtual Status DeleteCF(uint32_t column_family_id,
 537                             const Slice& key) override {
 538       ++num_seen;
 539       return TestHandler::DeleteCF(column_family_id, key);
 540     }
 541     virtual Status SingleDeleteCF(uint32_t column_family_id,
 542                                   const Slice& key) override {
 543       ++num_seen;
 544       return TestHandler::SingleDeleteCF(column_family_id, key);
 545     }
 546     virtual Status MergeCF(uint32_t column_family_id, const Slice& key,
 547                            const Slice& value) override {
 548       ++num_seen;
 549       return TestHandler::MergeCF(column_family_id, key, value);
 550     }
 551     virtual void LogData(const Slice& blob) override {
 552       ++num_seen;
 553       TestHandler::LogData(blob);
 554     }
 555     virtual bool Continue() override { return num_seen < 5; }
 556   } handler;
 557
 558   batch.Put(Slice("k1"), Slice("v1"));
 559   batch.Put(Slice("k2"), Slice("v2"));
 560   batch.PutLogData(Slice("blob1"));
 561   batch.Delete(Slice("k1"));
 562   batch.SingleDelete(Slice("k2"));
 563   batch.PutLogData(Slice("blob2"));
 564   batch.Merge(Slice("foo"), Slice("bar"));
 565   batch.Iterate(&handler);
 566   ASSERT_EQ(
 567       "Put(k1, v1)"
 568       "Put(k2, v2)"
 569       "LogData(blob1)"
 570       "Delete(k1)"
 571       "SingleDelete(k2)",
 572       handler.seen);
 573 }
 574
 575 TEST_F(WriteBatchTest, PutGatherSlices) {
 576   WriteBatch batch;
 577   batch.Put(Slice("foo"), Slice("bar"));
 578
 579   {
 580     // Try a write where the key is one slice but the value is two
 581     Slice key_slice("baz");
 582     Slice value_slices[2] = { Slice("header"), Slice("payload") };
 583     batch.Put(SliceParts(&key_slice, 1),
 584               SliceParts(value_slices, 2));
 585   }
 586
 587   {
 588     // One where the key is composite but the value is a single slice
 589     Slice key_slices[3] = { Slice("key"), Slice("part2"), Slice("part3") };
 590     Slice value_slice("value");
 591     batch.Put(SliceParts(key_slices, 3),
 592               SliceParts(&value_slice, 1));
 593   }
 594
 595   WriteBatchInternal::SetSequence(&batch, 100);
 596   ASSERT_EQ("Put(baz, headerpayload)@101"
 597             "Put(foo, bar)@100"
 598             "Put(keypart2part3, value)@102",
 599             PrintContents(&batch));
 600   ASSERT_EQ(3, batch.Count());
 601 }
 602
 603 namespace {
 604 class ColumnFamilyHandleImplDummy : public ColumnFamilyHandleImpl {
 605  public:
 606   explicit ColumnFamilyHandleImplDummy(int id)
 607       : ColumnFamilyHandleImpl(nullptr, nullptr, nullptr), id_(id) {}
 608   uint32_t GetID() const override { return id_; }
 609   const Comparator* GetComparator() const override {
 610     return BytewiseComparator();
 611   }
 612
 613  private:
 614   uint32_t id_;
 615 };
 616 }  // namespace anonymous
 617
 618 TEST_F(WriteBatchTest, ColumnFamiliesBatchTest) {
 619   WriteBatch batch;
 620   ColumnFamilyHandleImplDummy zero(0), two(2), three(3), eight(8);
 621   batch.Put(&zero, Slice("foo"), Slice("bar"));
 622   batch.Put(&two, Slice("twofoo"), Slice("bar2"));
 623   batch.Put(&eight, Slice("eightfoo"), Slice("bar8"));
 624   batch.Delete(&eight, Slice("eightfoo"));
 625   batch.SingleDelete(&two, Slice("twofoo"));
 626   batch.DeleteRange(&two, Slice("3foo"), Slice("4foo"));
 627   batch.Merge(&three, Slice("threethree"), Slice("3three"));
 628   batch.Put(&zero, Slice("foo"), Slice("bar"));
 629   batch.Merge(Slice("omom"), Slice("nom"));
 630
 631   TestHandler handler;
 632   batch.Iterate(&handler);
 633   ASSERT_EQ(
 634       "Put(foo, bar)"
 635       "PutCF(2, twofoo, bar2)"
 636       "PutCF(8, eightfoo, bar8)"
 637       "DeleteCF(8, eightfoo)"
 638       "SingleDeleteCF(2, twofoo)"
 639       "DeleteRangeCF(2, 3foo, 4foo)"
 640       "MergeCF(3, threethree, 3three)"
 641       "Put(foo, bar)"
 642       "Merge(omom, nom)",
 643       handler.seen);
 644 }
 645
 646 #ifndef ROCKSDB_LITE
 647 TEST_F(WriteBatchTest, ColumnFamiliesBatchWithIndexTest) {
 648   WriteBatchWithIndex batch;
 649   ColumnFamilyHandleImplDummy zero(0), two(2), three(3), eight(8);
 650   batch.Put(&zero, Slice("foo"), Slice("bar"));
 651   batch.Put(&two, Slice("twofoo"), Slice("bar2"));
 652   batch.Put(&eight, Slice("eightfoo"), Slice("bar8"));
 653   batch.Delete(&eight, Slice("eightfoo"));
 654   batch.SingleDelete(&two, Slice("twofoo"));
 655   batch.DeleteRange(&two, Slice("twofoo"), Slice("threefoo"));
 656   batch.Merge(&three, Slice("threethree"), Slice("3three"));
 657   batch.Put(&zero, Slice("foo"), Slice("bar"));
 658   batch.Merge(Slice("omom"), Slice("nom"));
 659
 660   std::unique_ptr<WBWIIterator> iter;
 661
 662   iter.reset(batch.NewIterator(&eight));
 663   iter->Seek("eightfoo");
 664   ASSERT_OK(iter->status());
 665   ASSERT_TRUE(iter->Valid());
 666   ASSERT_EQ(WriteType::kPutRecord, iter->Entry().type);
 667   ASSERT_EQ("eightfoo", iter->Entry().key.ToString());
 668   ASSERT_EQ("bar8", iter->Entry().value.ToString());
 669
 670   iter->Next();
 671   ASSERT_OK(iter->status());
 672   ASSERT_TRUE(iter->Valid());
 673   ASSERT_EQ(WriteType::kDeleteRecord, iter->Entry().type);
 674   ASSERT_EQ("eightfoo", iter->Entry().key.ToString());
 675
 676   iter->Next();
 677   ASSERT_OK(iter->status());
 678   ASSERT_TRUE(!iter->Valid());
 679
 680   iter.reset(batch.NewIterator(&two));
 681   iter->Seek("twofoo");
 682   ASSERT_OK(iter->status());
 683   ASSERT_TRUE(iter->Valid());
 684   ASSERT_EQ(WriteType::kPutRecord, iter->Entry().type);
 685   ASSERT_EQ("twofoo", iter->Entry().key.ToString());
 686   ASSERT_EQ("bar2", iter->Entry().value.ToString());
 687
 688   iter->Next();
 689   ASSERT_OK(iter->status());
 690   ASSERT_TRUE(iter->Valid());
 691   ASSERT_EQ(WriteType::kSingleDeleteRecord, iter->Entry().type);
 692   ASSERT_EQ("twofoo", iter->Entry().key.ToString());
 693
 694   iter->Next();
 695   ASSERT_OK(iter->status());
 696   ASSERT_TRUE(iter->Valid());
 697   ASSERT_EQ(WriteType::kDeleteRangeRecord, iter->Entry().type);
 698   ASSERT_EQ("twofoo", iter->Entry().key.ToString());
 699   ASSERT_EQ("threefoo", iter->Entry().value.ToString());
 700
 701   iter->Next();
 702   ASSERT_OK(iter->status());
 703   ASSERT_TRUE(!iter->Valid());
 704
 705   iter.reset(batch.NewIterator());
 706   iter->Seek("gggg");
 707   ASSERT_OK(iter->status());
 708   ASSERT_TRUE(iter->Valid());
 709   ASSERT_EQ(WriteType::kMergeRecord, iter->Entry().type);
 710   ASSERT_EQ("omom", iter->Entry().key.ToString());
 711   ASSERT_EQ("nom", iter->Entry().value.ToString());
 712
 713   iter->Next();
 714   ASSERT_OK(iter->status());
 715   ASSERT_TRUE(!iter->Valid());
 716
 717   iter.reset(batch.NewIterator(&zero));
 718   iter->Seek("foo");
 719   ASSERT_OK(iter->status());
 720   ASSERT_TRUE(iter->Valid());
 721   ASSERT_EQ(WriteType::kPutRecord, iter->Entry().type);
 722   ASSERT_EQ("foo", iter->Entry().key.ToString());
 723   ASSERT_EQ("bar", iter->Entry().value.ToString());
 724
 725   iter->Next();
 726   ASSERT_OK(iter->status());
 727   ASSERT_TRUE(iter->Valid());
 728   ASSERT_EQ(WriteType::kPutRecord, iter->Entry().type);
 729   ASSERT_EQ("foo", iter->Entry().key.ToString());
 730   ASSERT_EQ("bar", iter->Entry().value.ToString());
 731
 732   iter->Next();
 733   ASSERT_OK(iter->status());
 734   ASSERT_TRUE(iter->Valid());
 735   ASSERT_EQ(WriteType::kMergeRecord, iter->Entry().type);
 736   ASSERT_EQ("omom", iter->Entry().key.ToString());
 737   ASSERT_EQ("nom", iter->Entry().value.ToString());
 738
 739   iter->Next();
 740   ASSERT_OK(iter->status());
 741   ASSERT_TRUE(!iter->Valid());
 742
 743   TestHandler handler;
 744   batch.GetWriteBatch()->Iterate(&handler);
 745   ASSERT_EQ(
 746       "Put(foo, bar)"
 747       "PutCF(2, twofoo, bar2)"
 748       "PutCF(8, eightfoo, bar8)"
 749       "DeleteCF(8, eightfoo)"
 750       "SingleDeleteCF(2, twofoo)"
 751       "DeleteRangeCF(2, twofoo, threefoo)"
 752       "MergeCF(3, threethree, 3three)"
 753       "Put(foo, bar)"
 754       "Merge(omom, nom)",
 755       handler.seen);
 756 }
 757 #endif  // !ROCKSDB_LITE
 758
 759 TEST_F(WriteBatchTest, SavePointTest) {
 760   Status s;
 761   WriteBatch batch;
 762   batch.SetSavePoint();
 763
 764   batch.Put("A", "a");
 765   batch.Put("B", "b");
 766   batch.SetSavePoint();
 767
 768   batch.Put("C", "c");
 769   batch.Delete("A");
 770   batch.SetSavePoint();
 771   batch.SetSavePoint();
 772
 773   ASSERT_OK(batch.RollbackToSavePoint());
 774   ASSERT_EQ(
 775       "Delete(A)@3"
 776       "Put(A, a)@0"
 777       "Put(B, b)@1"
 778       "Put(C, c)@2",
 779       PrintContents(&batch));
 780
 781   ASSERT_OK(batch.RollbackToSavePoint());
 782   ASSERT_OK(batch.RollbackToSavePoint());
 783   ASSERT_EQ(
 784       "Put(A, a)@0"
 785       "Put(B, b)@1",
 786       PrintContents(&batch));
 787
 788   batch.Delete("A");
 789   batch.Put("B", "bb");
 790
 791   ASSERT_OK(batch.RollbackToSavePoint());
 792   ASSERT_EQ("", PrintContents(&batch));
 793
 794   s = batch.RollbackToSavePoint();
 795   ASSERT_TRUE(s.IsNotFound());
 796   ASSERT_EQ("", PrintContents(&batch));
 797
 798   batch.Put("D", "d");
 799   batch.Delete("A");
 800
 801   batch.SetSavePoint();
 802
 803   batch.Put("A", "aaa");
 804
 805   ASSERT_OK(batch.RollbackToSavePoint());
 806   ASSERT_EQ(
 807       "Delete(A)@1"
 808       "Put(D, d)@0",
 809       PrintContents(&batch));
 810
 811   batch.SetSavePoint();
 812
 813   batch.Put("D", "d");
 814   batch.Delete("A");
 815
 816   ASSERT_OK(batch.RollbackToSavePoint());
 817   ASSERT_EQ(
 818       "Delete(A)@1"
 819       "Put(D, d)@0",
 820       PrintContents(&batch));
 821
 822   s = batch.RollbackToSavePoint();
 823   ASSERT_TRUE(s.IsNotFound());
 824   ASSERT_EQ(
 825       "Delete(A)@1"
 826       "Put(D, d)@0",
 827       PrintContents(&batch));
 828
 829   WriteBatch batch2;
 830
 831   s = batch2.RollbackToSavePoint();
 832   ASSERT_TRUE(s.IsNotFound());
 833   ASSERT_EQ("", PrintContents(&batch2));
 834
 835   batch2.Delete("A");
 836   batch2.SetSavePoint();
 837
 838   s = batch2.RollbackToSavePoint();
 839   ASSERT_OK(s);
 840   ASSERT_EQ("Delete(A)@0", PrintContents(&batch2));
 841
 842   batch2.Clear();
 843   ASSERT_EQ("", PrintContents(&batch2));
 844
 845   batch2.SetSavePoint();
 846
 847   batch2.Delete("B");
 848   ASSERT_EQ("Delete(B)@0", PrintContents(&batch2));
 849
 850   batch2.SetSavePoint();
 851   s = batch2.RollbackToSavePoint();
 852   ASSERT_OK(s);
 853   ASSERT_EQ("Delete(B)@0", PrintContents(&batch2));
 854
 855   s = batch2.RollbackToSavePoint();
 856   ASSERT_OK(s);
 857   ASSERT_EQ("", PrintContents(&batch2));
 858
 859   s = batch2.RollbackToSavePoint();
 860   ASSERT_TRUE(s.IsNotFound());
 861   ASSERT_EQ("", PrintContents(&batch2));
 862 }
 863
 864 TEST_F(WriteBatchTest, MemoryLimitTest) {
 865   Status s;
 866   // The header size is 12 bytes. The two Puts take 8 bytes which gives total
 867   // of 12 + 8 * 2 = 28 bytes.
 868   WriteBatch batch(0, 28);
 869
 870   ASSERT_OK(batch.Put("a", "...."));
 871   ASSERT_OK(batch.Put("b", "...."));
 872   s = batch.Put("c", "....");
 873   ASSERT_TRUE(s.IsMemoryLimit());
 874 }
 875
 876 }  // namespace rocksdb
 877
 878 int main(int argc, char** argv) {
 879   ::testing::InitGoogleTest(&argc, argv);
 880   return RUN_ALL_TESTS();
 881 }