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 #pragma once
7 
8 #include <algorithm>
9 #include <cinttypes>
10 #include <functional>
11 #include <string>
12 #include <thread>
13 
14 #include "db/db_impl/db_impl.h"
15 #include "rocksdb/db.h"
16 #include "rocksdb/options.h"
17 #include "rocksdb/utilities/transaction.h"
18 #include "rocksdb/utilities/transaction_db.h"
19 #include "table/mock_table.h"
20 #include "test_util/fault_injection_test_env.h"
21 #include "test_util/sync_point.h"
22 #include "test_util/testharness.h"
23 #include "test_util/testutil.h"
24 #include "test_util/transaction_test_util.h"
25 #include "util/random.h"
26 #include "util/string_util.h"
27 #include "utilities/merge_operators.h"
28 #include "utilities/merge_operators/string_append/stringappend.h"
29 #include "utilities/transactions/pessimistic_transaction_db.h"
30 #include "utilities/transactions/write_unprepared_txn_db.h"
31 
32 #include "port/port.h"
33 
34 namespace ROCKSDB_NAMESPACE {
35 
36 // Return true if the ith bit is set in combination represented by comb
IsInCombination(size_t i,size_t comb)37 bool IsInCombination(size_t i, size_t comb) { return comb & (size_t(1) << i); }
38 
39 enum WriteOrdering : bool { kOrderedWrite, kUnorderedWrite };
40 
41 class TransactionTestBase : public ::testing::Test {
42  public:
43   TransactionDB* db;
44   FaultInjectionTestEnv* env;
45   std::string dbname;
46   Options options;
47 
48   TransactionDBOptions txn_db_options;
49   bool use_stackable_db_;
50 
TransactionTestBase(bool use_stackable_db,bool two_write_queue,TxnDBWritePolicy write_policy,WriteOrdering write_ordering)51   TransactionTestBase(bool use_stackable_db, bool two_write_queue,
52                       TxnDBWritePolicy write_policy,
53                       WriteOrdering write_ordering)
54       : db(nullptr), env(nullptr), use_stackable_db_(use_stackable_db) {
55     options.create_if_missing = true;
56     options.max_write_buffer_number = 2;
57     options.write_buffer_size = 4 * 1024;
58     options.unordered_write = write_ordering == kUnorderedWrite;
59     options.level0_file_num_compaction_trigger = 2;
60     options.merge_operator = MergeOperators::CreateFromStringId("stringappend");
61     env = new FaultInjectionTestEnv(Env::Default());
62     options.env = env;
63     options.two_write_queues = two_write_queue;
64     dbname = test::PerThreadDBPath("transaction_testdb");
65 
66     DestroyDB(dbname, options);
67     txn_db_options.transaction_lock_timeout = 0;
68     txn_db_options.default_lock_timeout = 0;
69     txn_db_options.write_policy = write_policy;
70     txn_db_options.rollback_merge_operands = true;
71     // This will stress write unprepared, by forcing write batch flush on every
72     // write.
73     txn_db_options.default_write_batch_flush_threshold = 1;
74     // Write unprepared requires all transactions to be named. This setting
75     // autogenerates the name so that existing tests can pass.
76     txn_db_options.autogenerate_name = true;
77     Status s;
78     if (use_stackable_db == false) {
79       s = TransactionDB::Open(options, txn_db_options, dbname, &db);
80     } else {
81       s = OpenWithStackableDB();
82     }
83     assert(s.ok());
84   }
85 
~TransactionTestBase()86   ~TransactionTestBase() {
87     delete db;
88     db = nullptr;
89     // This is to skip the assert statement in FaultInjectionTestEnv. There
90     // seems to be a bug in btrfs that the makes readdir return recently
91     // unlink-ed files. By using the default fs we simply ignore errors resulted
92     // from attempting to delete such files in DestroyDB.
93     options.env = Env::Default();
94     DestroyDB(dbname, options);
95     delete env;
96   }
97 
ReOpenNoDelete()98   Status ReOpenNoDelete() {
99     delete db;
100     db = nullptr;
101     env->AssertNoOpenFile();
102     env->DropUnsyncedFileData();
103     env->ResetState();
104     Status s;
105     if (use_stackable_db_ == false) {
106       s = TransactionDB::Open(options, txn_db_options, dbname, &db);
107     } else {
108       s = OpenWithStackableDB();
109     }
110     assert(!s.ok() || db != nullptr);
111     return s;
112   }
113 
ReOpenNoDelete(std::vector<ColumnFamilyDescriptor> & cfs,std::vector<ColumnFamilyHandle * > * handles)114   Status ReOpenNoDelete(std::vector<ColumnFamilyDescriptor>& cfs,
115                         std::vector<ColumnFamilyHandle*>* handles) {
116     for (auto h : *handles) {
117       delete h;
118     }
119     handles->clear();
120     delete db;
121     db = nullptr;
122     env->AssertNoOpenFile();
123     env->DropUnsyncedFileData();
124     env->ResetState();
125     Status s;
126     if (use_stackable_db_ == false) {
127       s = TransactionDB::Open(options, txn_db_options, dbname, cfs, handles,
128                               &db);
129     } else {
130       s = OpenWithStackableDB(cfs, handles);
131     }
132     assert(!s.ok() || db != nullptr);
133     return s;
134   }
135 
ReOpen()136   Status ReOpen() {
137     delete db;
138     db = nullptr;
139     DestroyDB(dbname, options);
140     Status s;
141     if (use_stackable_db_ == false) {
142       s = TransactionDB::Open(options, txn_db_options, dbname, &db);
143     } else {
144       s = OpenWithStackableDB();
145     }
146     assert(db != nullptr);
147     return s;
148   }
149 
OpenWithStackableDB(std::vector<ColumnFamilyDescriptor> & cfs,std::vector<ColumnFamilyHandle * > * handles)150   Status OpenWithStackableDB(std::vector<ColumnFamilyDescriptor>& cfs,
151                              std::vector<ColumnFamilyHandle*>* handles) {
152     std::vector<size_t> compaction_enabled_cf_indices;
153     TransactionDB::PrepareWrap(&options, &cfs, &compaction_enabled_cf_indices);
154     DB* root_db = nullptr;
155     Options options_copy(options);
156     const bool use_seq_per_batch =
157         txn_db_options.write_policy == WRITE_PREPARED ||
158         txn_db_options.write_policy == WRITE_UNPREPARED;
159     const bool use_batch_per_txn =
160         txn_db_options.write_policy == WRITE_COMMITTED ||
161         txn_db_options.write_policy == WRITE_PREPARED;
162     Status s = DBImpl::Open(options_copy, dbname, cfs, handles, &root_db,
163                             use_seq_per_batch, use_batch_per_txn);
164     StackableDB* stackable_db = new StackableDB(root_db);
165     if (s.ok()) {
166       assert(root_db != nullptr);
167       s = TransactionDB::WrapStackableDB(stackable_db, txn_db_options,
168                                          compaction_enabled_cf_indices,
169                                          *handles, &db);
170     }
171     if (!s.ok()) {
172       delete stackable_db;
173     }
174     return s;
175   }
176 
OpenWithStackableDB()177   Status OpenWithStackableDB() {
178     std::vector<size_t> compaction_enabled_cf_indices;
179     std::vector<ColumnFamilyDescriptor> column_families{ColumnFamilyDescriptor(
180         kDefaultColumnFamilyName, ColumnFamilyOptions(options))};
181 
182     TransactionDB::PrepareWrap(&options, &column_families,
183                                &compaction_enabled_cf_indices);
184     std::vector<ColumnFamilyHandle*> handles;
185     DB* root_db = nullptr;
186     Options options_copy(options);
187     const bool use_seq_per_batch =
188         txn_db_options.write_policy == WRITE_PREPARED ||
189         txn_db_options.write_policy == WRITE_UNPREPARED;
190     const bool use_batch_per_txn =
191         txn_db_options.write_policy == WRITE_COMMITTED ||
192         txn_db_options.write_policy == WRITE_PREPARED;
193     Status s = DBImpl::Open(options_copy, dbname, column_families, &handles,
194                             &root_db, use_seq_per_batch, use_batch_per_txn);
195     if (!s.ok()) {
196       delete root_db;
197       return s;
198     }
199     StackableDB* stackable_db = new StackableDB(root_db);
200     assert(root_db != nullptr);
201     assert(handles.size() == 1);
202     s = TransactionDB::WrapStackableDB(stackable_db, txn_db_options,
203                                        compaction_enabled_cf_indices, handles,
204                                        &db);
205     delete handles[0];
206     if (!s.ok()) {
207       delete stackable_db;
208     }
209     return s;
210   }
211 
212   std::atomic<size_t> linked = {0};
213   std::atomic<size_t> exp_seq = {0};
214   std::atomic<size_t> commit_writes = {0};
215   std::atomic<size_t> expected_commits = {0};
216   // Without Prepare, the commit does not write to WAL
217   std::atomic<size_t> with_empty_commits = {0};
218   std::function<void(size_t, Status)> txn_t0_with_status = [&](size_t index,
219                                                                Status exp_s) {
220     // Test DB's internal txn. It involves no prepare phase nor a commit marker.
221     WriteOptions wopts;
222     auto s = db->Put(wopts, "key" + std::to_string(index), "value");
223     ASSERT_EQ(exp_s, s);
224     if (txn_db_options.write_policy == TxnDBWritePolicy::WRITE_COMMITTED) {
225       // Consume one seq per key
226       exp_seq++;
227     } else {
228       // Consume one seq per batch
229       exp_seq++;
230       if (options.two_write_queues) {
231         // Consume one seq for commit
232         exp_seq++;
233       }
234     }
235     with_empty_commits++;
236   };
237   std::function<void(size_t)> txn_t0 = [&](size_t index) {
238     return txn_t0_with_status(index, Status::OK());
239   };
240   std::function<void(size_t)> txn_t1 = [&](size_t index) {
241     // Testing directly writing a write batch. Functionality-wise it is
242     // equivalent to commit without prepare.
243     WriteBatch wb;
244     auto istr = std::to_string(index);
245     ASSERT_OK(wb.Put("k1" + istr, "v1"));
246     ASSERT_OK(wb.Put("k2" + istr, "v2"));
247     ASSERT_OK(wb.Put("k3" + istr, "v3"));
248     WriteOptions wopts;
249     auto s = db->Write(wopts, &wb);
250     if (txn_db_options.write_policy == TxnDBWritePolicy::WRITE_COMMITTED) {
251       // Consume one seq per key
252       exp_seq += 3;
253     } else {
254       // Consume one seq per batch
255       exp_seq++;
256       if (options.two_write_queues) {
257         // Consume one seq for commit
258         exp_seq++;
259       }
260     }
261     ASSERT_OK(s);
262     with_empty_commits++;
263   };
264   std::function<void(size_t)> txn_t2 = [&](size_t index) {
265     // Commit without prepare. It should write to DB without a commit marker.
266     TransactionOptions txn_options;
267     WriteOptions write_options;
268     Transaction* txn = db->BeginTransaction(write_options, txn_options);
269     auto istr = std::to_string(index);
270     ASSERT_OK(txn->SetName("xid" + istr));
271     ASSERT_OK(txn->Put(Slice("foo" + istr), Slice("bar")));
272     ASSERT_OK(txn->Put(Slice("foo2" + istr), Slice("bar2")));
273     ASSERT_OK(txn->Put(Slice("foo3" + istr), Slice("bar3")));
274     ASSERT_OK(txn->Put(Slice("foo4" + istr), Slice("bar4")));
275     ASSERT_OK(txn->Commit());
276     if (txn_db_options.write_policy == TxnDBWritePolicy::WRITE_COMMITTED) {
277       // Consume one seq per key
278       exp_seq += 4;
279     } else if (txn_db_options.write_policy ==
280                TxnDBWritePolicy::WRITE_PREPARED) {
281       // Consume one seq per batch
282       exp_seq++;
283       if (options.two_write_queues) {
284         // Consume one seq for commit
285         exp_seq++;
286       }
287     } else {
288       // Flushed after each key, consume one seq per flushed batch
289       exp_seq += 4;
290       // WriteUnprepared implements CommitWithoutPrepareInternal by simply
291       // calling Prepare then Commit. Consume one seq for the prepare.
292       exp_seq++;
293     }
294     delete txn;
295     with_empty_commits++;
296   };
297   std::function<void(size_t)> txn_t3 = [&](size_t index) {
298     // A full 2pc txn that also involves a commit marker.
299     TransactionOptions txn_options;
300     WriteOptions write_options;
301     Transaction* txn = db->BeginTransaction(write_options, txn_options);
302     auto istr = std::to_string(index);
303     ASSERT_OK(txn->SetName("xid" + istr));
304     ASSERT_OK(txn->Put(Slice("foo" + istr), Slice("bar")));
305     ASSERT_OK(txn->Put(Slice("foo2" + istr), Slice("bar2")));
306     ASSERT_OK(txn->Put(Slice("foo3" + istr), Slice("bar3")));
307     ASSERT_OK(txn->Put(Slice("foo4" + istr), Slice("bar4")));
308     ASSERT_OK(txn->Put(Slice("foo5" + istr), Slice("bar5")));
309     expected_commits++;
310     ASSERT_OK(txn->Prepare());
311     commit_writes++;
312     ASSERT_OK(txn->Commit());
313     if (txn_db_options.write_policy == TxnDBWritePolicy::WRITE_COMMITTED) {
314       // Consume one seq per key
315       exp_seq += 5;
316     } else if (txn_db_options.write_policy ==
317                TxnDBWritePolicy::WRITE_PREPARED) {
318       // Consume one seq per batch
319       exp_seq++;
320       // Consume one seq per commit marker
321       exp_seq++;
322     } else {
323       // Flushed after each key, consume one seq per flushed batch
324       exp_seq += 5;
325       // Consume one seq per commit marker
326       exp_seq++;
327     }
328     delete txn;
329   };
330   std::function<void(size_t)> txn_t4 = [&](size_t index) {
331     // A full 2pc txn that also involves a commit marker.
332     TransactionOptions txn_options;
333     WriteOptions write_options;
334     Transaction* txn = db->BeginTransaction(write_options, txn_options);
335     auto istr = std::to_string(index);
336     ASSERT_OK(txn->SetName("xid" + istr));
337     ASSERT_OK(txn->Put(Slice("foo" + istr), Slice("bar")));
338     ASSERT_OK(txn->Put(Slice("foo2" + istr), Slice("bar2")));
339     ASSERT_OK(txn->Put(Slice("foo3" + istr), Slice("bar3")));
340     ASSERT_OK(txn->Put(Slice("foo4" + istr), Slice("bar4")));
341     ASSERT_OK(txn->Put(Slice("foo5" + istr), Slice("bar5")));
342     expected_commits++;
343     ASSERT_OK(txn->Prepare());
344     commit_writes++;
345     ASSERT_OK(txn->Rollback());
346     if (txn_db_options.write_policy == TxnDBWritePolicy::WRITE_COMMITTED) {
347       // No seq is consumed for deleting the txn buffer
348       exp_seq += 0;
349     } else if (txn_db_options.write_policy ==
350                TxnDBWritePolicy::WRITE_PREPARED) {
351       // Consume one seq per batch
352       exp_seq++;
353       // Consume one seq per rollback batch
354       exp_seq++;
355       if (options.two_write_queues) {
356         // Consume one seq for rollback commit
357         exp_seq++;
358       }
359     } else {
360       // Flushed after each key, consume one seq per flushed batch
361       exp_seq += 5;
362       // Consume one seq per rollback batch
363       exp_seq++;
364       if (options.two_write_queues) {
365         // Consume one seq for rollback commit
366         exp_seq++;
367       }
368     }
369     delete txn;
370   };
371 
372   // Test that we can change write policy after a clean shutdown (which would
373   // empty the WAL)
CrossCompatibilityTest(TxnDBWritePolicy from_policy,TxnDBWritePolicy to_policy,bool empty_wal)374   void CrossCompatibilityTest(TxnDBWritePolicy from_policy,
375                               TxnDBWritePolicy to_policy, bool empty_wal) {
376     TransactionOptions txn_options;
377     ReadOptions read_options;
378     WriteOptions write_options;
379     uint32_t index = 0;
380     Random rnd(1103);
381     options.write_buffer_size = 1024;  // To create more sst files
382     std::unordered_map<std::string, std::string> committed_kvs;
383     Transaction* txn;
384 
385     txn_db_options.write_policy = from_policy;
386     if (txn_db_options.write_policy == WRITE_COMMITTED) {
387       options.unordered_write = false;
388     }
389     ReOpen();
390 
391     for (int i = 0; i < 1024; i++) {
392       auto istr = std::to_string(index);
393       auto k = Slice("foo-" + istr).ToString();
394       auto v = Slice("bar-" + istr).ToString();
395       // For test the duplicate keys
396       auto v2 = Slice("bar2-" + istr).ToString();
397       auto type = rnd.Uniform(4);
398       switch (type) {
399         case 0:
400           committed_kvs[k] = v;
401           ASSERT_OK(db->Put(write_options, k, v));
402           committed_kvs[k] = v2;
403           ASSERT_OK(db->Put(write_options, k, v2));
404           break;
405         case 1: {
406           WriteBatch wb;
407           committed_kvs[k] = v;
408           wb.Put(k, v);
409           committed_kvs[k] = v2;
410           wb.Put(k, v2);
411           ASSERT_OK(db->Write(write_options, &wb));
412 
413         } break;
414         case 2:
415         case 3:
416           txn = db->BeginTransaction(write_options, txn_options);
417           ASSERT_OK(txn->SetName("xid" + istr));
418           committed_kvs[k] = v;
419           ASSERT_OK(txn->Put(k, v));
420           committed_kvs[k] = v2;
421           ASSERT_OK(txn->Put(k, v2));
422 
423           if (type == 3) {
424             ASSERT_OK(txn->Prepare());
425           }
426           ASSERT_OK(txn->Commit());
427           delete txn;
428           break;
429         default:
430           assert(0);
431       }
432 
433       index++;
434     }  // for i
435 
436     txn_db_options.write_policy = to_policy;
437     if (txn_db_options.write_policy == WRITE_COMMITTED) {
438       options.unordered_write = false;
439     }
440     auto db_impl = reinterpret_cast<DBImpl*>(db->GetRootDB());
441     // Before upgrade/downgrade the WAL must be emptied
442     if (empty_wal) {
443       db_impl->TEST_FlushMemTable();
444     } else {
445       db_impl->FlushWAL(true);
446     }
447     auto s = ReOpenNoDelete();
448     if (empty_wal) {
449       ASSERT_OK(s);
450     } else {
451       // Test that we can detect the WAL that is produced by an incompatible
452       // WritePolicy and fail fast before mis-interpreting the WAL.
453       ASSERT_TRUE(s.IsNotSupported());
454       return;
455     }
456     db_impl = reinterpret_cast<DBImpl*>(db->GetRootDB());
457     // Check that WAL is empty
458     VectorLogPtr log_files;
459     db_impl->GetSortedWalFiles(log_files);
460     ASSERT_EQ(0, log_files.size());
461 
462     for (auto& kv : committed_kvs) {
463       std::string value;
464       s = db->Get(read_options, kv.first, &value);
465       if (s.IsNotFound()) {
466         printf("key = %s\n", kv.first.c_str());
467       }
468       ASSERT_OK(s);
469       if (kv.second != value) {
470         printf("key = %s\n", kv.first.c_str());
471       }
472       ASSERT_EQ(kv.second, value);
473     }
474   }
475 };
476 
477 class TransactionTest
478     : public TransactionTestBase,
479       virtual public ::testing::WithParamInterface<
480           std::tuple<bool, bool, TxnDBWritePolicy, WriteOrdering>> {
481  public:
TransactionTest()482   TransactionTest()
483       : TransactionTestBase(std::get<0>(GetParam()), std::get<1>(GetParam()),
484                             std::get<2>(GetParam()), std::get<3>(GetParam())){};
485 };
486 
487 class TransactionStressTest : public TransactionTest {};
488 
489 class MySQLStyleTransactionTest
490     : public TransactionTestBase,
491       virtual public ::testing::WithParamInterface<
492           std::tuple<bool, bool, TxnDBWritePolicy, WriteOrdering, bool>> {
493  public:
MySQLStyleTransactionTest()494   MySQLStyleTransactionTest()
495       : TransactionTestBase(std::get<0>(GetParam()), std::get<1>(GetParam()),
496                             std::get<2>(GetParam()), std::get<3>(GetParam())),
497         with_slow_threads_(std::get<4>(GetParam())) {
498     if (with_slow_threads_ &&
499         (txn_db_options.write_policy == WRITE_PREPARED ||
500          txn_db_options.write_policy == WRITE_UNPREPARED)) {
501       // The corner case with slow threads involves the caches filling
502       // over which would not happen even with artifial delays. To help
503       // such cases to show up we lower the size of the cache-related data
504       // structures.
505       txn_db_options.wp_snapshot_cache_bits = 1;
506       txn_db_options.wp_commit_cache_bits = 10;
507       options.write_buffer_size = 1024;
508       EXPECT_OK(ReOpen());
509     }
510   };
511 
512  protected:
513   // Also emulate slow threads by addin artiftial delays
514   const bool with_slow_threads_;
515 };
516 
517 }  // namespace ROCKSDB_NAMESPACE
518