p4: add test cases for bonus tasks (#653)

Signed-off-by: Alex Chi <[email protected]>

Author: skyzh

src/concurrency/transaction_manager.cpp

@@ -42,21 +42,52 @@ auto TransactionManager::Begin(IsolationLevel isolation_level) -> Transaction *
   auto *txn_ref = txn.get();
   txn_map_.insert(std::make_pair(txn_id, std::move(txn)));
 
-  // TODO(fall2023): set the timestamps and compute watermark.
+  // TODO(fall2023): set the timestamps here. Watermark updated below.
 
+  running_txns_.AddTxn(txn_ref->read_ts_);
   return txn_ref;
 }
 
+auto TransactionManager::VerifyTxn(Transaction *txn) -> bool { return true; }
+
 auto TransactionManager::Commit(Transaction *txn) -> bool {
-  std::lock_guard<std::mutex> commit_lck(commit_mutex_);
-  // TODO(fall2023): Implement me!
+  std::unique_lock<std::mutex> commit_lck(commit_mutex_);
+
+  if (txn->state_ != TransactionState::RUNNING) {
+    throw Exception("txn not in running state");
+  }
+
+  if (txn->GetIsolationLevel() == IsolationLevel::SERIALIZABLE) {
+    if (!VerifyTxn(txn)) {
+      commit_lck.unlock();
+      Abort(txn);
+      return false;
+    }
+  }
+
+  // TODO(fall2023): Implement the commit logic!
+
+  std::unique_lock<std::shared_mutex> lck(txn_map_mutex_);
+
+  // TODO(fall2023): set commit timestamp + update last committed timestamp here.
+
   txn->state_ = TransactionState::COMMITTED;
+  running_txns_.UpdateCommitTs(txn->commit_ts_);
+  running_txns_.RemoveTxn(txn->read_ts_);
+
   return true;
 }
 
 void TransactionManager::Abort(Transaction *txn) {
-  // TODO(fall2023): Implement me!
+  if (txn->state_ != TransactionState::RUNNING && txn->state_ != TransactionState::TAINTED) {
+    throw Exception("txn not in running / tainted state");
+  }
+
+  // TODO(fall2023): Implement the abort logic!
+
+  std::unique_lock<std::shared_mutex> lck(txn_map_mutex_);
   txn->state_ = TransactionState::ABORTED;
+  running_txns_.RemoveTxn(txn->read_ts_);
 }
 
 void TransactionManager::GarbageCollection() { UNIMPLEMENTED("not implemented"); }

src/concurrency/watermark.cpp

@@ -4,8 +4,16 @@
 
 namespace bustub {
 
-auto Watermark::AddTxn(timestamp_t read_ts) -> void { throw NotImplementedException("unimplemented"); }
+auto Watermark::AddTxn(timestamp_t read_ts) -> void {
+  if (read_ts < commit_ts_) {
+    throw Exception("read ts < commit ts");
+  }
 
-auto Watermark::RemoveTxn(timestamp_t read_ts) -> void { throw NotImplementedException("unimplemented"); }
+  // TODO(fall2023): implement me!
+}
+
+auto Watermark::RemoveTxn(timestamp_t read_ts) -> void {
+  // TODO(fall2023): implement me!
+}
 
 }  // namespace bustub

src/include/common/bustub_instance.h

@@ -121,7 +121,7 @@ class StringVectorWriter : public ResultWriter {
   void EndHeader() override {}
   void BeginRow() override { values_.emplace_back(); }
   void EndRow() override {}
-  void BeginTable(bool simplified_output) override {}
+  void BeginTable(bool simplified_output) override { values_.clear(); }
   void EndTable() override {}
 
   std::vector<std::vector<std::string>> values_;

src/include/concurrency/transaction_manager.h

@@ -127,6 +127,9 @@ class TransactionManager {
   Catalog *catalog_;
 
   std::atomic<txn_id_t> next_txn_id_{TXN_START_ID};
+
+ private:
+  auto VerifyTxn(Transaction *txn) -> bool;
 };
 
 }  // namespace bustub

test/txn/txn_common.h

@@ -248,28 +248,67 @@ void ExecuteTxn(BustubInstance &instance, const std::string &txn_var_name, Trans
   }
 }
 
-auto BeginTxn(BustubInstance &instance, const std::string &txn_var_name) -> Transaction * {
-  auto txn = instance.txn_manager_->Begin();
-  fmt::println(stderr, "- txn_begin var={} id={} status={} read_ts={}", txn_var_name,
-               txn->GetTransactionIdHumanReadable(), txn->GetTransactionState(), txn->GetReadTs());
+auto BeginTxn(BustubInstance &instance, const std::string &txn_var_name,
+              IsolationLevel iso_lvl = IsolationLevel::SNAPSHOT_ISOLATION) -> Transaction * {
+  auto txn = instance.txn_manager_->Begin(iso_lvl);
+  fmt::println(stderr, "- txn_begin var={} id={} status={} read_ts={} iso_lvl={}", txn_var_name,
+               txn->GetTransactionIdHumanReadable(), txn->GetTransactionState(), txn->GetReadTs(),
+               txn->GetIsolationLevel());
   return txn;
 }
 
-auto CommitTxn(BustubInstance &instance, const std::string &txn_var_name, Transaction *txn) {
+auto BeginTxnSerializable(BustubInstance &instance, const std::string &txn_var_name) -> Transaction * {
+  return BeginTxn(instance, txn_var_name, IsolationLevel::SERIALIZABLE);
+}
+
+const bool EXPECT_FAIL = true;
+
+auto CommitTxn(BustubInstance &instance, const std::string &txn_var_name, Transaction *txn, bool expect_fail = false) {
   if (txn->GetTransactionState() != TransactionState::RUNNING) {
     fmt::println(stderr, "txn not running");
     std::terminate();
   }
-  if (!instance.txn_manager_->Commit(txn)) {
-    fmt::println(stderr, "failed to commit txn: var={} id={}", txn_var_name, txn->GetTransactionId());
+  auto res = instance.txn_manager_->Commit(txn);
+  if (!expect_fail) {
+    if (!res) {
+      fmt::println(stderr, "failed to commit txn: var={} id={}", txn_var_name, txn->GetTransactionId());
+      std::terminate();
+    }
+    if (txn->GetTransactionState() != TransactionState::COMMITTED) {
+      fmt::println(stderr, "should set to committed state var={} id={}", txn_var_name, txn->GetTransactionId());
+      std::terminate();
+    }
+    fmt::println(stderr, "- txn_commit var={} id={} status={} read_ts={} commit_ts={}", txn_var_name,
+                 txn->GetTransactionIdHumanReadable(), txn->GetTransactionState(), txn->GetReadTs(),
+                 txn->GetCommitTs());
+    return;
+  }
+  if (res) {
+    fmt::println(stderr, "expect txn fail to commit, but committed: var={} id={}", txn_var_name,
+                 txn->GetTransactionId());
     std::terminate();
   }
-  if (txn->GetTransactionState() != TransactionState::COMMITTED) {
-    fmt::println(stderr, "should set to committed state var={} id={}", txn_var_name, txn->GetTransactionId());
+  if (txn->GetTransactionState() != TransactionState::ABORTED) {
+    fmt::println(stderr, "should set to aborted state var={} id={}", txn_var_name, txn->GetTransactionId());
     std::terminate();
   }
-  fmt::println(stderr, "- txn_commit var={} id={} status={} read_ts={} commit_ts={}", txn_var_name,
-               txn->GetTransactionIdHumanReadable(), txn->GetTransactionState(), txn->GetReadTs(), txn->GetCommitTs());
+  fmt::println(stderr, "- txn_commit_fail var={} id={} status={} read_ts={}", txn_var_name,
+               txn->GetTransactionIdHumanReadable(), txn->GetTransactionState(), txn->GetReadTs());
+}
+
+auto AbortTxn(BustubInstance &instance, const std::string &txn_var_name, Transaction *txn) {
+  if (txn->GetTransactionState() != TransactionState::RUNNING &&
+      txn->GetTransactionState() != TransactionState::TAINTED) {
+    fmt::println(stderr, "txn not running / tainted");
+    std::terminate();
+  }
+  instance.txn_manager_->Abort(txn);
+  if (txn->GetTransactionState() != TransactionState::ABORTED) {
+    fmt::println(stderr, "should set to aborted state var={} id={}", txn_var_name, txn->GetTransactionId());
+    std::terminate();
+  }
+  fmt::println(stderr, "- txn_abort var={} id={} status={} read_ts={}", txn_var_name,
+               txn->GetTransactionIdHumanReadable(), txn->GetTransactionState(), txn->GetReadTs());
 }
 
 auto CheckTainted(BustubInstance &instance, const std::string &txn_var_name, Transaction *txn) {

test/txn/txn_index_concurrent_test.cpp

@@ -2,6 +2,7 @@
 #include <exception>
 #include <memory>
 #include <mutex>  // NOLINT
+#include <random>
 #include <string>
 #include <thread>  // NOLINT
 #include "common/bustub_instance.h"
@@ -149,6 +150,7 @@ TEST(TxnIndexTest, DISABLED_IndexConcurrentUpdateTest) {  // NOLINT
           if (add_delete_insert) {
             StringVectorWriter data_writer;
             BUSTUB_ENSURE(bustub->ExecuteSqlTxn(generate_select_sql(i), data_writer, txn), "cannot retrieve data");
+            BUSTUB_ENSURE(data_writer.values_.size() == 1, "more than 1 row fetched??");
             const auto b_val = std::stoi(data_writer.values_[0][0]);
             BUSTUB_ENSURE(bustub->ExecuteSqlTxn(generate_delete_sql(i), data_writer, txn), "cannot delete data");
             BUSTUB_ENSURE(bustub->ExecuteSqlTxn(generate_txn_insert_sql(b_val, i), data_writer, txn),
@@ -190,6 +192,99 @@ TEST(TxnIndexTest, DISABLED_IndexConcurrentUpdateTest) {  // NOLINT
   }
 }
 
+TEST(TxnIndexTest, DISABLED_IndexConcurrentUpdateAbortTest) {  // NOLINT
+  const auto generate_sql = [](int n) -> std::string {
+    return fmt::format("UPDATE maintable SET b = b + {} WHERE a = {}", 1, n);
+  };
+  const int thread_cnt = 8;
+  const int number_cnt = 5;
+  const auto generate_insert_sql = [](int n) -> std::string {
+    std::vector<std::string> data;
+    data.reserve(n);
+    for (int i = 0; i < n; i++) {
+      data.push_back(fmt::format("({}, {})", i, 0));
+    }
+    return fmt::format("INSERT INTO maintable VALUES {}", fmt::join(data, ","));
+  };
+  const int trials = 10;
+  const int operation_cnt = 100;
+  for (int n = 0; n < trials; n++) {
+    auto bustub = std::make_unique<BustubInstance>();
+    EnsureIndexScan(*bustub);
+    Execute(*bustub, "CREATE TABLE maintable(a int primary key, b int)");
+    std::vector<std::thread> update_threads;
+    Execute(*bustub, generate_insert_sql(number_cnt), false);
+    TableHeapEntryNoMoreThan(*bustub, bustub->catalog_->GetTable("maintable"), number_cnt);
+    update_threads.reserve(thread_cnt);
+    std::map<int, std::vector<int>> operation_result;
+    std::mutex result_mutex;
+    fmt::println(stderr, "trial {}: running with {} threads with {} rows ", n + 1, thread_cnt, number_cnt);
+    global_disable_execution_exception_print.store(true);
+    for (int thread = 0; thread < thread_cnt; thread++) {
+      update_threads.emplace_back([&bustub, thread, generate_sql, &result_mutex, &operation_result]() {
+        NoopWriter writer;
+        std::vector<int> result(number_cnt, 0);
+        std::random_device dev;
+        std::mt19937 rng(dev());
+        std::uniform_int_distribution<std::mt19937::result_type> dist(0, number_cnt - 1);
+        for (int i = 0; i < operation_cnt; i++) {
+          int x = 0;
+          int y = 0;
+          do {
+            x = dist(rng);
+            y = dist(rng);
+          } while (x == y);
+          auto *txn = bustub->txn_manager_->Begin();
+          auto sql = generate_sql(x);
+          if (!bustub->ExecuteSqlTxn(sql, writer, txn)) {
+            bustub->txn_manager_->Abort(txn);
+            std::this_thread::sleep_for(std::chrono::milliseconds(1));
+            continue;
+          }
+          sql = generate_sql(y);
+          if (!bustub->ExecuteSqlTxn(sql, writer, txn)) {
+            bustub->txn_manager_->Abort(txn);
+            std::this_thread::sleep_for(std::chrono::milliseconds(1));
+            continue;
+          }
+          BUSTUB_ENSURE(bustub->txn_manager_->Commit(txn), "cannot commit??");
+          result[x] += 1;
+          result[y] += 1;
+        }
+        {
+          std::lock_guard<std::mutex> lck(result_mutex);
+          operation_result.emplace(thread, std::move(result));
+        }
+      });
+    }
+    for (auto &&thread : update_threads) {
+      thread.join();
+    }
+    global_disable_execution_exception_print.store(false);
+    std::vector<std::vector<int>> expected_rows;
+    for (int i = 0; i < number_cnt; i++) {
+      int total = 0;
+      for (int j = 0; j < thread_cnt; j++) {
+        total += operation_result[j][i];
+      }
+      expected_rows.push_back({i, total});
+      if (total < 10) {
+        fmt::println(stderr, "abort rate too high, {} txn succeeded", total);
+        std::terminate();
+      }
+    }
+    auto *table_info = bustub->catalog_->GetTable("maintable");
+    auto query_txn = BeginTxn(*bustub, "query_txn");
+    WithTxn(query_txn, QueryShowResult(*bustub, _var, _txn, "SELECT * FROM maintable", expected_rows));
+    TableHeapEntryNoMoreThan(*bustub, table_info, number_cnt);
+    if (n >= trials - 2) {
+      SimpleStreamWriter writer(std::cerr);
+      fmt::println(stderr, "--- the following data might be manually inspected by TAs ---");
+      bustub->ExecuteSqlTxn("SELECT * FROM maintable", writer, query_txn);
+    }
+  }
+}
+
 // NOLINTEND(bugprone-unchecked-optional-access))
 
 }  // namespace bustub