Improve efficiency in PointLockManager by using separate Condvar (#13731)

Summary:
PointLockManager manages point lock per key. The old implementation partition the per key lock into 16 stripes. Each stripe handles the point lock for a subset of keys. Each stripe have only one conditional variable. This conditional variable is used by all the transactions that are waiting for its turn to acquire a lock of a key that belongs to this stripe.

In production, we notified that when there are multiple transactions trying to write to the same key, all of them will wait on the same conditional variables. When the previous lock holder released the key, all of the transactions are woken up, but only one of them could proceed, and the rest goes back to sleep. This wasted a lot of CPU cycles. In addition, when there are other keys being locked/unlocked on the same lock stripe, the problem becomes even worse.

In order to solve this issue, we implemented a new PerKeyPointLockManager that keeps a transaction waiter queue at per key level. When a transaction could not acquire a lock immediately, it joins the waiter queue of the key and waits on a dedicated conditional variable. When previous lock holder released the lock, it wakes up the next set of transactions that are eligible to acquire the lock from the waiting queue. The queue respect FIFO order, except it prioritizes lock upgrade/downgrade operation.

However, this waiter queue change increases the deadlock detection cost, because the transaction waiting in the queue also needs to be considered during deadlock detection. To resolve this issue, a new deadlock_timeout_us (microseconds) configuration is introduced in transaction option. Essentially, when a transaction is waiting on a lock, it will join the wait queue and wait for the duration configured by deadlock_timeout_us without perform deadlock detection. If the transaction didn't get the lock after the deadlock_timeout_us timeout is reached, it will then perform deadlock detection and wait until lock_timeout is reached. This optimization takes the heuristic where majority of the transaction would be able to get the lock without perform deadlock detection.

The deadlock_timeout_us configuration needs to be tuned for different workload, if the likelihood of deadlock is very low, the deadlock_timeout_us could be configured close to a big higher than the average transaction execution time, so that majority of the transaction would be able to acquire the lock without performing deadlock detection. If the likelihood of deadlock is high, deadlock_timeout_us could be configured with lower value, so that deadlock would get detected faster.

The new PerKeyPointLockManager is disabled by default. It can be enabled by TransactionDBOptions.use_per_key_point_lock_mgr. The deadlock_timeout_us is only effective when PerKeyPointLockManager is used. When deadlock_timeout_us is set to 0, transaction will perform deadlock detection immediately before wait.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/13731

Test Plan:
Unit test.
Stress unit test that validates deadlock detection and exclusive, shared lock guarantee.
A new point_lock_bench binary is created to help perform performance test.

Reviewed By: pdillinger

Differential Revision: D77353607

Pulled By: xingbowang

fbshipit-source-id: 21cf93354f9a367a78c8666596ed14013ac7240b

BUCK

@@ -419,6 +419,8 @@ cpp_library_wrapper(name="rocksdb_tools_lib", srcs=[
 
 cpp_library_wrapper(name="rocksdb_cache_bench_tools_lib", srcs=["cache/cache_bench_tool.cc"], deps=[":rocksdb_lib"], headers=[], link_whole=False, extra_test_libs=False)
 
+cpp_library_wrapper(name="rocksdb_point_lock_bench_tools_lib", srcs=["utilities/transactions/lock/point/point_lock_bench_tool.cc"], deps=[":rocksdb_lib"], headers=[], link_whole=False, extra_test_libs=False)
+
 rocks_cpp_library_wrapper(name="rocksdb_stress_lib", srcs=[
         "db_stress_tool/batched_ops_stress.cc",
         "db_stress_tool/cf_consistency_stress.cc",
@@ -450,6 +452,8 @@ cpp_binary_wrapper(name="db_bench", srcs=["tools/db_bench.cc"], deps=[":rocksdb_
 
 cpp_binary_wrapper(name="cache_bench", srcs=["cache/cache_bench.cc"], deps=[":rocksdb_cache_bench_tools_lib"], extra_preprocessor_flags=[], extra_bench_libs=False)
 
+cpp_binary_wrapper(name="point_lock_bench", srcs=["utilities/transactions/lock/point/point_lock_bench.cc"], deps=[":rocksdb_point_lock_bench_tools_lib"], extra_preprocessor_flags=[], extra_bench_libs=False)
+
 cpp_binary_wrapper(name="ribbon_bench", srcs=["microbench/ribbon_bench.cc"], deps=[], extra_preprocessor_flags=[], extra_bench_libs=True)
 
 cpp_binary_wrapper(name="db_basic_bench", srcs=["microbench/db_basic_bench.cc"], deps=[], extra_preprocessor_flags=[], extra_bench_libs=True)
@@ -5381,6 +5385,12 @@ cpp_unittest_wrapper(name="plain_table_db_test",
             extra_compiler_flags=[])
 
 
+cpp_unittest_wrapper(name="point_lock_manager_stress_test",
+            srcs=["utilities/transactions/lock/point/point_lock_manager_stress_test.cc"],
+            deps=[":rocksdb_test_lib"],
+            extra_compiler_flags=[])
+
+
 cpp_unittest_wrapper(name="point_lock_manager_test",
             srcs=["utilities/transactions/lock/point/point_lock_manager_test.cc"],
             deps=[":rocksdb_test_lib"],

CMakeLists.txt

@@ -1503,6 +1503,7 @@ if(WITH_TESTS)
         utilities/transactions/optimistic_transaction_test.cc
         utilities/transactions/transaction_test.cc
         utilities/transactions/lock/point/point_lock_manager_test.cc
+        utilities/transactions/lock/point/point_lock_manager_stress_test.cc
         utilities/transactions/write_committed_transaction_ts_test.cc
         utilities/transactions/write_prepared_transaction_test.cc
         utilities/transactions/write_unprepared_transaction_test.cc
@@ -1613,6 +1614,12 @@ if(WITH_BENCHMARK_TOOLS)
     utilities/persistent_cache/hash_table_bench.cc)
   target_link_libraries(hash_table_bench${ARTIFACT_SUFFIX}
     ${ROCKSDB_LIB} ${GFLAGS_LIB} ${FOLLY_LIBS})
+
+  add_executable(point_lock_bench${ARTIFACT_SUFFIX}
+    utilities/transactions/lock/point/point_lock_bench.cc
+    utilities/transactions/lock/point/point_lock_bench_tool.cc)
+  target_link_libraries(point_lock_bench${ARTIFACT_SUFFIX}
+    ${ROCKSDB_LIB} ${GFLAGS_LIB} ${FOLLY_LIBS})
 endif()
 
 option(WITH_TRACE_TOOLS "build with trace tools" ON)

Makefile

@@ -636,13 +636,14 @@ endif
 TEST_OBJECTS = $(patsubst %.cc, $(OBJ_DIR)/%.o, $(TEST_LIB_SOURCES) $(MOCK_LIB_SOURCES)) $(GTEST)
 BENCH_OBJECTS = $(patsubst %.cc, $(OBJ_DIR)/%.o, $(BENCH_LIB_SOURCES))
 CACHE_BENCH_OBJECTS = $(patsubst %.cc, $(OBJ_DIR)/%.o, $(CACHE_BENCH_LIB_SOURCES))
+POINT_LOCK_BENCH_OBJECTS = $(patsubst %.cc, $(OBJ_DIR)/%.o, $(POINT_LOCK_BENCH_LIB_SOURCES))
 TOOL_OBJECTS = $(patsubst %.cc, $(OBJ_DIR)/%.o, $(TOOL_LIB_SOURCES))
 ANALYZE_OBJECTS = $(patsubst %.cc, $(OBJ_DIR)/%.o, $(ANALYZER_LIB_SOURCES))
 STRESS_OBJECTS =  $(patsubst %.cc, $(OBJ_DIR)/%.o, $(STRESS_LIB_SOURCES))
 
 # Exclude build_version.cc -- a generated source file -- from all sources.  Not needed for dependencies
 ALL_SOURCES  = $(filter-out util/build_version.cc, $(LIB_SOURCES)) $(TEST_LIB_SOURCES) $(MOCK_LIB_SOURCES) $(GTEST_DIR)/gtest/gtest-all.cc
-ALL_SOURCES += $(TOOL_LIB_SOURCES) $(BENCH_LIB_SOURCES) $(CACHE_BENCH_LIB_SOURCES) $(ANALYZER_LIB_SOURCES) $(STRESS_LIB_SOURCES)
+ALL_SOURCES += $(TOOL_LIB_SOURCES) $(BENCH_LIB_SOURCES) $(CACHE_BENCH_LIB_SOURCES) $(POINT_LOCK_BENCH_LIB_SOURCES) $(ANALYZER_LIB_SOURCES) $(STRESS_LIB_SOURCES)
 ALL_SOURCES += $(TEST_MAIN_SOURCES) $(TOOL_MAIN_SOURCES) $(BENCH_MAIN_SOURCES)
 ALL_SOURCES += $(ROCKSDB_PLUGIN_SOURCES) $(ROCKSDB_PLUGIN_TESTS)
 
@@ -1343,6 +1344,9 @@ block_cache_trace_analyzer: $(OBJ_DIR)/tools/block_cache_analyzer/block_cache_tr
 cache_bench: $(OBJ_DIR)/cache/cache_bench.o $(CACHE_BENCH_OBJECTS) $(LIBRARY)
 	$(AM_LINK)
 
+point_lock_bench: $(OBJ_DIR)/utilities/transactions/lock/point/point_lock_bench.o $(POINT_LOCK_BENCH_OBJECTS) $(LIBRARY)
+	$(AM_LINK)
+
 persistent_cache_bench: $(OBJ_DIR)/utilities/persistent_cache/persistent_cache_bench.o $(LIBRARY)
 	$(AM_LINK)
 
@@ -1879,6 +1883,9 @@ heap_test: $(OBJ_DIR)/util/heap_test.o $(TEST_LIBRARY) $(LIBRARY)
 point_lock_manager_test: utilities/transactions/lock/point/point_lock_manager_test.o $(TEST_LIBRARY) $(LIBRARY)
 	$(AM_LINK)
 
+point_lock_manager_stress_test: utilities/transactions/lock/point/point_lock_manager_stress_test.o $(TEST_LIBRARY) $(LIBRARY)
+	$(AM_LINK)
+
 transaction_test: $(OBJ_DIR)/utilities/transactions/transaction_test.o $(TEST_LIBRARY) $(LIBRARY)
 	$(AM_LINK)
 

buckifier/buckify_rocksdb.py

@@ -206,6 +206,12 @@ def generate_buck(repo_path, deps_map):
         src_mk.get("CACHE_BENCH_LIB_SOURCES", []),
         [":rocksdb_lib"],
     )
+    # rocksdb_point_lock_bench_tools_lib
+    BUCK.add_library(
+        "rocksdb_point_lock_bench_tools_lib",
+        src_mk.get("POINT_LOCK_BENCH_LIB_SOURCES", []),
+        [":rocksdb_lib"],
+    )
     # rocksdb_stress_lib
     BUCK.add_rocksdb_library(
         "rocksdb_stress_lib",
@@ -229,6 +235,12 @@ def generate_buck(repo_path, deps_map):
     BUCK.add_binary(
         "cache_bench", ["cache/cache_bench.cc"], [":rocksdb_cache_bench_tools_lib"]
     )
+    # point_lock_bench binary
+    BUCK.add_binary(
+        "point_lock_bench",
+        ["utilities/transactions/lock/point/point_lock_bench.cc"],
+        [":rocksdb_point_lock_bench_tools_lib"]
+    )
     # bench binaries
     for src in src_mk.get("MICROBENCH_SOURCES", []):
         name = src.rsplit("/", 1)[1].split(".")[0] if "/" in src else src.split(".")[0]

db/c.cc

@@ -6433,6 +6433,11 @@ void rocksdb_transactiondb_options_set_default_lock_timeout(
   opt->rep.default_lock_timeout = default_lock_timeout;
 }
 
+void rocksdb_transactiondb_options_set_use_per_key_point_lock_mgr(
+    rocksdb_transactiondb_options_t* opt, int use_per_key_point_lock_mgr) {
+  opt->rep.use_per_key_point_lock_mgr = use_per_key_point_lock_mgr;
+}
+
 rocksdb_transaction_options_t* rocksdb_transaction_options_create() {
   return new rocksdb_transaction_options_t;
 }

db_stress_tool/db_stress_common.h

@@ -284,6 +284,7 @@ DECLARE_bool(use_txn);
 // Options for TransactionDB (a.k.a. Pessimistic Transaction DB)
 DECLARE_uint64(txn_write_policy);
 DECLARE_bool(unordered_write);
+DECLARE_bool(use_per_key_point_lock_mgr);
 
 // Options for OptimisticTransactionDB
 DECLARE_bool(use_optimistic_txn);

db_stress_tool/db_stress_gflags.cc

@@ -724,6 +724,10 @@ DEFINE_uint64(txn_write_policy, 0,
               "TxnDBWritePolicy::WRITE_COMMITTED. Note that this should not be "
               "changed across crashes.");
 
+DEFINE_bool(use_per_key_point_lock_mgr, true,
+            "Use PointLockManager(false) or PerKeyPointLockManager(true) in "
+            "TransactionDB.");
+
 DEFINE_bool(use_optimistic_txn, false, "Use OptimisticTransactionDB.");
 DEFINE_uint64(occ_validation_policy, 1,
               "Optimistic Concurrency Control Validation Policy for "

db_stress_tool/db_stress_test_base.cc

@@ -3944,6 +3944,8 @@ void StressTest::Open(SharedState* shared, bool reopen) {
             static_cast<size_t>(FLAGS_wp_snapshot_cache_bits);
         txn_db_options.wp_commit_cache_bits =
             static_cast<size_t>(FLAGS_wp_commit_cache_bits);
+        txn_db_options.use_per_key_point_lock_mgr =
+            FLAGS_use_per_key_point_lock_mgr;
         PrepareTxnDbOptions(shared, txn_db_options);
         s = TransactionDB::Open(options_, txn_db_options, FLAGS_db,
                                 cf_descriptors, &column_families_, &txn_db_);

include/rocksdb/point_lock_bench_tool.h

@@ -0,0 +1,14 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#pragma once
+
+#include "rocksdb/rocksdb_namespace.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+int point_lock_bench_tool(int argc, char** argv);
+}  // namespace ROCKSDB_NAMESPACE

include/rocksdb/status.h

@@ -116,6 +116,7 @@ class Status {
     kMergeOperatorFailed = 15,
     kMergeOperandThresholdExceeded = 16,
     kPrefetchLimitReached = 17,
+    kNotExpectedCodePath = 18,
     kMaxSubCode
   };
 
@@ -329,6 +330,9 @@ class Status {
     return code() == kOk;
   }
 
+  // Assert the status is OK in debug mode
+  void AssertOK() const { assert(ok()); }
+
   // Returns true iff the status indicates success *with* something
   // overwritten
   bool IsOkOverwritten() const {

include/rocksdb/utilities/transaction.h

@@ -653,7 +653,12 @@ class Transaction {
   // Change the value of TransactionOptions.lock_timeout (in milliseconds) for
   // this transaction.
   // Has no effect on OptimisticTransactions.
-  virtual void SetLockTimeout(int64_t timeout) = 0;
+  virtual void SetLockTimeout(int64_t timeout_ms) = 0;
+
+  // Change the value of deadlock_timeout (in milliseconds) for this
+  // transaction.
+  // Has no effect on OptimisticTransactions.
+  virtual void SetDeadlockTimeout(int64_t timeout_ms) = 0;
 
   // Return the WriteOptions that will be used during Commit()
   virtual WriteOptions* GetWriteOptions() = 0;

include/rocksdb/utilities/transaction_db.h

@@ -217,6 +217,11 @@ struct TransactionDBOptions {
   // Other value means the user provides a custom lock manager.
   std::shared_ptr<LockManagerHandle> lock_mgr_handle;
 
+  // EXPERIMENTAL
+  //
+  // Flag to enable/disable the per key point lock manager.
+  bool use_per_key_point_lock_mgr = false;
+
   // If true, the TransactionDB implementation might skip concurrency control
   // unless it is overridden by TransactionOptions or
   // TransactionDBWriteOptimizations. This can be used in conjunction with
@@ -319,6 +324,22 @@ struct TransactionOptions {
   // If negative, TransactionDBOptions::transaction_lock_timeout will be used.
   int64_t lock_timeout = -1;
 
+  // Timeout in microseconds before perform dead lock detection.
+  // If 0, deadlock detection will be performed immediately.
+  //
+  // To optimize performance, this parameter could be tuned.
+  //
+  // When deadlock happens very frequently, deadlock timeout should be set to 0,
+  // so deadlock will be detected immediately.
+  //
+  // When deadlock happen very rarely, this timeout could be turned to be
+  // slightly longer than the typical transaction execution time, so that
+  // transaction will be waked up to take the lock before this timeout, which
+  // will allow the transaction to save the CPU time on deadlock detection.
+  //
+  // Deadlock timeout is always smaller than lock_timeout.
+  int64_t deadlock_timeout_us = 500;
+
   // Expiration duration in milliseconds.  If non-negative, transactions that
   // last longer than this many milliseconds will fail to commit.  If not set,
   // a forgotten transaction that is never committed, rolled back, or deleted

src.mk

@@ -386,9 +386,12 @@ BENCH_LIB_SOURCES =                                             \
   tools/tool_hooks.cc                                           \
   tools/simulated_hybrid_file_system.cc                         \
 
-CACHE_BENCH_LIB_SOURCES =					                              \
+CACHE_BENCH_LIB_SOURCES =                                       \
   cache/cache_bench_tool.cc                                     \
 
+POINT_LOCK_BENCH_LIB_SOURCES =                                  \
+  utilities/transactions/lock/point/point_lock_bench_tool.cc    \
+
 STRESS_LIB_SOURCES =                                           \
   db_stress_tool/batched_ops_stress.cc                         \
   db_stress_tool/cf_consistency_stress.cc                      \
@@ -651,6 +654,7 @@ TEST_MAIN_SOURCES =                                                     \
   utilities/transactions/lock/range/range_locking_test.cc               \
   utilities/transactions/transaction_test.cc                            \
   utilities/transactions/lock/point/point_lock_manager_test.cc          \
+  utilities/transactions/lock/point/point_lock_manager_stress_test.cc   \
   utilities/transactions/write_prepared_transaction_test.cc             \
   utilities/transactions/write_unprepared_transaction_test.cc           \
   utilities/transactions/write_committed_transaction_ts_test.cc         \

tools/db_crashtest.py

@@ -567,6 +567,7 @@ txn_params = {
     # NOTE: often passed in from command line overriding this
     "txn_write_policy": random.randint(0, 2),
     "unordered_write": random.randint(0, 1),
+    "use_per_key_point_lock_mgr": lambda: random.choice([0, 1]),
     # TODO: there is such a thing as transactions with WAL disabled. We should
     # cover that case.
     "disable_wal": 0,

unreleased_history/performance_improvements/improve_point_lock_manager_performance.md

@@ -0,0 +1,2 @@
+Add a new experimental PerKeyPointLockManager to improve efficiency under high lock contention. PointLockManager was not efficient when there is high write contention on same key, as it uses a single conditional variable per lock stripe. PerKeyPointLockManager uses per thread conditional variable supporting fifo order. Although this is an experimental feature. By default, it is disabled. A new boolean flag TransactionDBOptions::use_per_key_point_lock_mgr is added to optionally enable it. Search the flag in code for more info.
+Together, a new configuration TransactionOptions::deadlock_timeout_us is added, which allows the transaction to wait for a short period before perform deadlock detection. When the workload has low lock contention, the deadlock_timeout_us can be configured to be slightly higher than average transaction execution time, so that transaction would likely be able to take the lock before deadlock detection is performed when it is waiting for a lock. This allows transaction to reduce CPU cost on performing deadlock detection, which could be expensive in CPU time. When the workload has high lock contention, the deadlock_timeout_us can be configured to 0, so that transaction would perform deadlock detection immediately. By default the value is 0 to keep the behavior same as before.

utilities/transactions/lock/lock_manager.cc

@@ -17,8 +17,11 @@ std::shared_ptr<LockManager> NewLockManager(PessimisticTransactionDB* db,
     auto mgr = opt.lock_mgr_handle->getLockManager();
     return std::shared_ptr<LockManager>(opt.lock_mgr_handle, mgr);
   } else {
-    // Use a point lock manager by default
-    return std::shared_ptr<LockManager>(new PointLockManager(db, opt));
+    if (opt.use_per_key_point_lock_mgr) {
+      return std::make_shared<PerKeyPointLockManager>(db, opt);
+    } else {
+      return std::make_shared<PointLockManager>(db, opt);
+    }
   }
 }
 

utilities/transactions/lock/point/any_lock_manager_test.h

@@ -0,0 +1,239 @@
+#pragma once
+
+#include "utilities/transactions/lock/point/point_lock_manager_test.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+using init_func_t = void (*)(PointLockManagerTest*);
+
+class AnyLockManagerTest : public PointLockManagerTest,
+                           public testing::WithParamInterface<init_func_t> {
+ public:
+  void SetUp() override {
+    // If a custom setup function was provided, use it. Otherwise, use what we
+    // have inherited.
+    auto init_func = GetParam();
+    if (init_func) {
+      (*init_func)(this);
+    } else {
+      PointLockManagerTest::SetUp();
+    }
+  }
+};
+
+TEST_P(AnyLockManagerTest, ReentrantExclusiveLock) {
+  // Tests that a txn can acquire exclusive lock on the same key repeatedly.
+  MockColumnFamilyHandle cf(1);
+  locker_->AddColumnFamily(&cf);
+  auto txn = NewTxn();
+  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, true));
+  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, true));
+
+  // Cleanup
+  locker_->UnLock(txn, 1, "k", env_);
+
+  delete txn;
+}
+
+TEST_P(AnyLockManagerTest, ReentrantSharedLock) {
+  // Tests that a txn can acquire shared lock on the same key repeatedly.
+  MockColumnFamilyHandle cf(1);
+  locker_->AddColumnFamily(&cf);
+  auto txn = NewTxn();
+  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, false));
+  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, false));
+
+  // Cleanup
+  if (dynamic_cast<PointLockManager*>(locker_.get()) != nullptr &&
+      dynamic_cast<PerKeyPointLockManager*>(locker_.get()) == nullptr) {
+    // PointLockManager would create 2 entries in the lock manager, so it needs
+    // to unlock it twice.
+    locker_->UnLock(txn, 1, "k", env_);
+  }
+  locker_->UnLock(txn, 1, "k", env_);
+
+  delete txn;
+}
+
+TEST_P(AnyLockManagerTest, LockUpgrade) {
+  // Tests that a txn can upgrade from a shared lock to an exclusive lock.
+  MockColumnFamilyHandle cf(1);
+  locker_->AddColumnFamily(&cf);
+  auto txn = NewTxn();
+  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, false));
+  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, true));
+
+  // Cleanup
+  locker_->UnLock(txn, 1, "k", env_);
+  delete txn;
+}
+
+TEST_P(AnyLockManagerTest, LockDowngrade) {
+  // Tests that a txn can acquire a shared lock after acquiring an exclusive
+  // lock on the same key.
+  MockColumnFamilyHandle cf(1);
+  locker_->AddColumnFamily(&cf);
+  auto txn = NewTxn();
+  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, true));
+  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, false));
+
+  // Cleanup
+  locker_->UnLock(txn, 1, "k", env_);
+  delete txn;
+}
+
+TEST_P(AnyLockManagerTest, LockConflict) {
+  // Tests that lock conflicts lead to lock timeout.
+  MockColumnFamilyHandle cf(1);
+  locker_->AddColumnFamily(&cf);
+  auto txn1 = NewTxn();
+  auto txn2 = NewTxn();
+
+  {
+    // exclusive-exclusive conflict.
+    ASSERT_OK(locker_->TryLock(txn1, 1, "k1", env_, true));
+    auto s = locker_->TryLock(txn2, 1, "k1", env_, true);
+    ASSERT_TRUE(s.IsTimedOut());
+  }
+
+  {
+    // exclusive-shared conflict.
+    ASSERT_OK(locker_->TryLock(txn1, 1, "k2", env_, true));
+    auto s = locker_->TryLock(txn2, 1, "k2", env_, false);
+    ASSERT_TRUE(s.IsTimedOut());
+  }
+
+  {
+    // shared-exclusive conflict.
+    ASSERT_OK(locker_->TryLock(txn1, 1, "k2", env_, false));
+    auto s = locker_->TryLock(txn2, 1, "k2", env_, true);
+    ASSERT_TRUE(s.IsTimedOut());
+  }
+
+  // Cleanup
+  locker_->UnLock(txn1, 1, "k1", env_);
+  locker_->UnLock(txn1, 1, "k2", env_);
+
+  delete txn1;
+  delete txn2;
+}
+
+TEST_P(AnyLockManagerTest, SharedLocks) {
+  // Tests that shared locks can be concurrently held by multiple transactions.
+  MockColumnFamilyHandle cf(1);
+  locker_->AddColumnFamily(&cf);
+  auto txn1 = NewTxn();
+  auto txn2 = NewTxn();
+  ASSERT_OK(locker_->TryLock(txn1, 1, "k", env_, false));
+  ASSERT_OK(locker_->TryLock(txn2, 1, "k", env_, false));
+
+  // Cleanup
+  locker_->UnLock(txn1, 1, "k", env_);
+  locker_->UnLock(txn2, 1, "k", env_);
+
+  delete txn1;
+  delete txn2;
+}
+
+TEST_P(AnyLockManagerTest, Deadlock) {
+  // Tests that deadlock can be detected.
+  // Deadlock scenario:
+  // txn1 exclusively locks k1, and wants to lock k2;
+  // txn2 exclusively locks k2, and wants to lock k1.
+  MockColumnFamilyHandle cf(1);
+  locker_->AddColumnFamily(&cf);
+  TransactionOptions txn_opt;
+  // disable dead lock timeout, so that the dead lock detection behavior is
+  // consistent. This prevents the test to be flaky
+  txn_opt.deadlock_timeout_us = 0;
+  txn_opt.deadlock_detect = true;
+  txn_opt.lock_timeout = 1000000;
+  auto txn1 = NewTxn(txn_opt);
+  auto txn2 = NewTxn(txn_opt);
+
+  ASSERT_OK(locker_->TryLock(txn1, 1, "k1", env_, true));
+  ASSERT_OK(locker_->TryLock(txn2, 1, "k2", env_, true));
+
+  // txn1 tries to lock k2, will be blocked.
+  port::Thread t;
+  BlockUntilWaitingTxn(wait_sync_point_name_, t, [&]() {
+    // block because txn2 is holding a lock on k2.
+    ASSERT_OK(locker_->TryLock(txn1, 1, "k2", env_, true));
+  });
+
+  auto s = locker_->TryLock(txn2, 1, "k1", env_, true);
+  ASSERT_TRUE(s.IsBusy());
+  ASSERT_EQ(s.subcode(), Status::SubCode::kDeadlock);
+
+  std::vector<DeadlockPath> deadlock_paths = locker_->GetDeadlockInfoBuffer();
+  ASSERT_EQ(deadlock_paths.size(), 1u);
+  ASSERT_FALSE(deadlock_paths[0].limit_exceeded);
+
+  std::vector<DeadlockInfo> deadlocks = deadlock_paths[0].path;
+  ASSERT_EQ(deadlocks.size(), 2u);
+
+  ASSERT_EQ(deadlocks[0].m_txn_id, txn1->GetID());
+  ASSERT_EQ(deadlocks[0].m_cf_id, 1u);
+  ASSERT_TRUE(deadlocks[0].m_exclusive);
+  ASSERT_EQ(deadlocks[0].m_waiting_key, "k2");
+
+  ASSERT_EQ(deadlocks[1].m_txn_id, txn2->GetID());
+  ASSERT_EQ(deadlocks[1].m_cf_id, 1u);
+  ASSERT_TRUE(deadlocks[1].m_exclusive);
+  ASSERT_EQ(deadlocks[1].m_waiting_key, "k1");
+
+  locker_->UnLock(txn2, 1, "k2", env_);
+  t.join();
+
+  // Cleanup
+  locker_->UnLock(txn1, 1, "k1", env_);
+  locker_->UnLock(txn1, 1, "k2", env_);
+  delete txn2;
+  delete txn1;
+}
+
+TEST_P(AnyLockManagerTest, GetWaitingTxns_MultipleTxns) {
+  MockColumnFamilyHandle cf(1);
+  locker_->AddColumnFamily(&cf);
+
+  auto txn1 = NewTxn();
+  ASSERT_OK(locker_->TryLock(txn1, 1, "k", env_, false));
+
+  auto txn2 = NewTxn();
+  ASSERT_OK(locker_->TryLock(txn2, 1, "k", env_, false));
+
+  auto txn3 = NewTxn();
+  txn3->SetLockTimeout(10000);
+  port::Thread t1;
+  BlockUntilWaitingTxn(wait_sync_point_name_, t1, [&]() {
+    ASSERT_OK(locker_->TryLock(txn3, 1, "k", env_, true));
+    locker_->UnLock(txn3, 1, "k", env_);
+  });
+
+  // Ok, now txn3 is waiting for lock on "k", which is owned by two
+  // transactions. Check that GetWaitingTxns reports this correctly
+  uint32_t wait_cf_id;
+  std::string wait_key;
+  auto waiters = txn3->GetWaitingTxns(&wait_cf_id, &wait_key);
+
+  ASSERT_EQ(wait_cf_id, 1u);
+  ASSERT_EQ(wait_key, "k");
+  ASSERT_EQ(waiters.size(), 2);
+  bool waits_correct =
+      (waiters[0] == txn1->GetID() && waiters[1] == txn2->GetID()) ||
+      (waiters[1] == txn1->GetID() && waiters[0] == txn2->GetID());
+  ASSERT_EQ(waits_correct, true);
+
+  // Release locks so txn3 can proceed with execution
+  locker_->UnLock(txn1, 1, "k", env_);
+  locker_->UnLock(txn2, 1, "k", env_);
+
+  // Wait until txn3 finishes
+  t1.join();
+
+  delete txn1;
+  delete txn2;
+  delete txn3;
+}
+
+}  // namespace ROCKSDB_NAMESPACE

utilities/transactions/lock/point/point_lock_bench.cc

@@ -0,0 +1,18 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#ifndef GFLAGS
+#include <cstdio>
+int main() {
+  fprintf(stderr, "Please install gflags to run rocksdb tools\n");
+  return 1;
+}
+#else
+#include "rocksdb/point_lock_bench_tool.h"
+int main(int argc, char** argv) {
+  return ROCKSDB_NAMESPACE::point_lock_bench_tool(argc, argv);
+}
+#endif  // GFLAGS

utilities/transactions/lock/point/point_lock_bench_tool.cc

@@ -0,0 +1,159 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#ifdef GFLAGS
+
+#include <cstdio>
+#include <iostream>
+#include <memory>
+
+#include "port/stack_trace.h"
+#include "rocksdb/convenience.h"
+#include "rocksdb/env.h"
+#include "rocksdb/utilities/transaction_db.h"
+#include "util/gflags_compat.h"
+#include "utilities/transactions/lock/point/point_lock_manager.h"
+#include "utilities/transactions/lock/point/point_lock_validation_test_runner.h"
+#include "utilities/transactions/pessimistic_transaction_db.h"
+
+using GFLAGS_NAMESPACE::ParseCommandLineFlags;
+
+namespace ROCKSDB_NAMESPACE {
+
+DEFINE_string(db_dir, "/tmp/point_lock_manager_test",
+              "DB path for running the benchmark");
+DEFINE_uint32(stripe_count, 16, "Number of stripes in point lock manager");
+DEFINE_bool(is_per_key_point_lock_manager, false,
+            "Use PerKeyPointLockManager or PointLockManager");
+DEFINE_uint32(thread_count, 64,
+              "Number of threads to acquire release locks concurrently");
+DEFINE_uint32(key_count, 16, "Number of keys to acquire release locks upon");
+DEFINE_uint32(max_num_keys_to_lock_per_txn, 8,
+              "Max Number of keys to lock in a transaction");
+DEFINE_uint32(execution_time_sec, 10,
+              "Number of seconds to execute the benchmark");
+DEFINE_uint32(lock_type, 2,
+              "Lock type to test, 0: exclusive lock only; 1: shared lock only; "
+              "2: both shared and exclusive locks");
+DEFINE_int64(lock_timeout_ms, 1000,
+             "Lock acquisition request timeout in milliseconds.");
+DEFINE_int64(deadlock_timeout_us, 500,
+             "DeadLock detection timeout in microseconds.");
+DEFINE_int64(lock_expiration_ms, 100,
+             "Acquired Lock expiration time in milliseconds.");
+DEFINE_bool(allow_non_deadlock_error, true,
+            "Allow returned error code other than deadlock, such as timeout.");
+DEFINE_uint32(
+    max_sleep_after_lock_acquisition_ms, 5,
+    "Max number of milliseconds to sleep after acquiring all the locks in the "
+    "transaction. The actuall sleep time will be randomized from 0 to max. It "
+    "is used to simulate some useful work performed.");
+DEFINE_bool(check_thread_stuck, false,
+            "Check thread periodically to see whether they are stuck or not. "
+            "This is useful for detecting stuck transaction quickly. But it "
+            "could have false-positive when running with ASAN or running with "
+            "high thread count on a small number of CPUs");
+
+namespace {  // anonymous namespace
+
+class PointLockManagerBenchmark {
+ public:
+  PointLockManagerBenchmark() {
+    env_ = Env::Default();
+    env_->CreateDir(FLAGS_db_dir);
+
+    Options opt;
+    opt.create_if_missing = true;
+    txndb_opt_.num_stripes = FLAGS_stripe_count;
+
+    db_ = nullptr;
+
+    auto s = TransactionDB::Open(opt, txndb_opt_, FLAGS_db_dir, &db_);
+    ASSERT_OK(s);
+
+    if (FLAGS_is_per_key_point_lock_manager) {
+      locker_ = std::make_shared<PerKeyPointLockManager>(
+          static_cast<PessimisticTransactionDB*>(db_), txndb_opt_);
+    } else {
+      locker_ = std::make_shared<PointLockManager>(
+          static_cast<PessimisticTransactionDB*>(db_), txndb_opt_);
+    }
+
+    txn_opt_.deadlock_detect = true;
+    txn_opt_.lock_timeout = FLAGS_lock_timeout_ms;
+    txn_opt_.deadlock_timeout_us = FLAGS_deadlock_timeout_us;
+    txn_opt_.expiration = FLAGS_lock_expiration_ms;
+  }
+
+  // Disable copy and assignment
+  PointLockManagerBenchmark(const PointLockManagerBenchmark&) = delete;
+  PointLockManagerBenchmark& operator=(const PointLockManagerBenchmark&) =
+      delete;
+  PointLockManagerBenchmark(PointLockManagerBenchmark&&) = delete;
+  PointLockManagerBenchmark& operator=(PointLockManagerBenchmark&&) = delete;
+
+  ~PointLockManagerBenchmark() {
+    delete db_;
+    auto s = DestroyDir(env_, FLAGS_db_dir);
+    ASSERT_OK(s);
+  }
+
+  void run() {
+    PointLockValidationTestRunner test_runner(
+        env_, txndb_opt_, locker_, db_, txn_opt_, FLAGS_thread_count,
+        FLAGS_key_count, FLAGS_max_num_keys_to_lock_per_txn,
+        FLAGS_execution_time_sec, static_cast<LockTypeToTest>(FLAGS_lock_type),
+        FLAGS_allow_non_deadlock_error,
+        FLAGS_max_sleep_after_lock_acquisition_ms, FLAGS_check_thread_stuck);
+    test_runner.run();
+  }
+
+ private:
+  Env* env_;
+  TransactionDBOptions txndb_opt_;
+  std::shared_ptr<LockManager> locker_;
+  TransactionDB* db_;
+  TransactionOptions txn_opt_;
+};
+
+}  // anonymous namespace
+
+int point_lock_bench_tool(int argc, char** argv) {
+  ROCKSDB_NAMESPACE::port::InstallStackTraceHandler();
+  ParseCommandLineFlags(&argc, &argv, true);
+
+  // Print test configuration
+  std::vector<gflags::CommandLineFlagInfo> all_flags;
+  gflags::GetAllFlags(&all_flags);
+
+  for (const auto& flag : all_flags) {
+    // only show the flags defined in this file
+    if (flag.filename.find("point_lock_bench_tool.cc") != std::string::npos) {
+      std::cout << "-" << flag.name << "=";
+      if (flag.type == "bool") {
+        std::cout << (gflags::GetCommandLineFlagInfoOrDie(flag.name.c_str())
+                                  .current_value == "true"
+                          ? "true"
+                          : "false");
+      } else {
+        std::cout << gflags::GetCommandLineFlagInfoOrDie(flag.name.c_str())
+                         .current_value;
+      }
+      std::cout << " ";
+    }
+  }
+  std::cout << std::endl;
+
+  // Run the benchmark
+  PointLockManagerBenchmark benchmark;
+  benchmark.run();
+
+  return 0;
+}
+
+}  // namespace ROCKSDB_NAMESPACE
+
+#endif  // GFLAGS

utilities/transactions/lock/point/point_lock_manager.h

@@ -132,8 +132,12 @@ class PointLockManager : public LockManager {
   // this column family is no longer in use.
   void RemoveColumnFamily(const ColumnFamilyHandle* cf) override;
 
+  // Caller makes sure that a lock on the key is not requested again, unless it
+  // is an upgrade or downgrade.
   Status TryLock(PessimisticTransaction* txn, ColumnFamilyId column_family_id,
                  const std::string& key, Env* env, bool exclusive) override;
+  // Caller makes sure that a lock on the key is not requested again, unless it
+  // is an upgrade or downgrade.
   Status TryLock(PessimisticTransaction* txn, ColumnFamilyId column_family_id,
                  const Endpoint& start, const Endpoint& end, Env* env,
                  bool exclusive) override;
@@ -153,7 +157,7 @@ class PointLockManager : public LockManager {
 
   void Resize(uint32_t new_size) override;
 
- private:
+ protected:
   PessimisticTransactionDB* txn_db_impl_;
 
   // Default number of lock map stripes per column family
@@ -179,6 +183,11 @@ class PointLockManager : public LockManager {
   // to avoid acquiring a mutex in order to look up a LockMap
   std::unique_ptr<ThreadLocalPtr> lock_maps_cache_;
 
+  // Thread local variable for KeyLockWaiter. As one thread could only need one
+  // KeyLockWaiter.
+  // Lazy init on first time usage
+  ThreadLocalPtr key_lock_waiter_;
+
   // Must be held when modifying wait_txn_map_ and rev_wait_txn_map_.
   std::mutex wait_txn_map_mutex_;
 
@@ -196,18 +205,13 @@ class PointLockManager : public LockManager {
 
   std::shared_ptr<LockMap> GetLockMap(uint32_t column_family_id);
 
-  Status AcquireWithTimeout(PessimisticTransaction* txn, LockMap* lock_map,
-                            LockMapStripe* stripe, uint32_t column_family_id,
-                            const std::string& key, Env* env, int64_t timeout,
-                            const LockInfo& lock_info);
+  virtual Status AcquireWithTimeout(
+      PessimisticTransaction* txn, LockMap* lock_map, LockMapStripe* stripe,
+      uint32_t column_family_id, const std::string& key, Env* env,
+      int64_t timeout, int64_t deadlock_timeout_us, const LockInfo& lock_info);
 
-  Status AcquireLocked(LockMap* lock_map, LockMapStripe* stripe,
-                       const std::string& key, Env* env,
-                       const LockInfo& lock_info, uint64_t* wait_time,
-                       autovector<TransactionID>* txn_ids);
-
-  void UnLockKey(PessimisticTransaction* txn, const std::string& key,
-                 LockMapStripe* stripe, LockMap* lock_map, Env* env);
+  virtual void UnLockKey(PessimisticTransaction* txn, const std::string& key,
+                         LockMapStripe* stripe, LockMap* lock_map, Env* env);
 
   // Returns true if a deadlock is detected.
   // Will DecrementWaiters() if a deadlock is detected.
@@ -219,6 +223,56 @@ class PointLockManager : public LockManager {
                         const autovector<TransactionID>& wait_ids);
   void DecrementWaitersImpl(const PessimisticTransaction* txn,
                             const autovector<TransactionID>& wait_ids);
+
+ private:
+  Status AcquireLocked(LockMap* lock_map, LockMapStripe* stripe,
+                       const std::string& key, Env* env,
+                       const LockInfo& lock_info, uint64_t* wait_time,
+                       autovector<TransactionID>* txn_ids);
+};
+
+class PerKeyPointLockManager : public PointLockManager {
+ public:
+  PerKeyPointLockManager(PessimisticTransactionDB* db,
+                         const TransactionDBOptions& opt);
+  // No copying allowed
+  PerKeyPointLockManager(const PerKeyPointLockManager&) = delete;
+  PerKeyPointLockManager& operator=(const PerKeyPointLockManager&) = delete;
+  // No move allowed
+  PerKeyPointLockManager(PerKeyPointLockManager&&) = delete;
+  PerKeyPointLockManager& operator=(PerKeyPointLockManager&&) = delete;
+
+  ~PerKeyPointLockManager() override {}
+
+  void UnLock(PessimisticTransaction* txn, const LockTracker& tracker,
+              Env* env) override;
+  void UnLock(PessimisticTransaction* txn, ColumnFamilyId column_family_id,
+              const std::string& key, Env* env) override;
+  void UnLock(PessimisticTransaction* txn, ColumnFamilyId column_family_id,
+              const Endpoint& start, const Endpoint& end, Env* env) override;
+
+  void UnLockKey(PessimisticTransaction* txn, const std::string& key,
+                 LockMapStripe* stripe, LockMap* lock_map, Env* env) override;
+
+ protected:
+  Status AcquireWithTimeout(PessimisticTransaction* txn, LockMap* lock_map,
+                            LockMapStripe* stripe, uint32_t column_family_id,
+                            const std::string& key, Env* env, int64_t timeout,
+                            int64_t deadlock_timeout_us,
+                            const LockInfo& lock_info) override;
+
+ private:
+  Status AcquireLocked(LockMap* lock_map, LockMapStripe* stripe,
+                       const std::string& key, Env* env,
+                       const LockInfo& txn_lock_info, uint64_t* wait_time,
+                       autovector<TransactionID>* txn_ids,
+                       LockInfo** lock_info_ptr, bool* isUpgrade, bool fifo);
+
+  int64_t CalculateWaitEndTime(int64_t expire_time_hint, int64_t end_time);
+
+  Status FillWaitIds(LockInfo& lock_info, const LockInfo& txn_lock_info,
+                     autovector<TransactionID>* wait_ids, bool& isUpgrade,
+                     TransactionID& my_txn_id, const std::string& key);
 };
 
 }  // namespace ROCKSDB_NAMESPACE

utilities/transactions/lock/point/point_lock_manager_stress_test.cc

@@ -0,0 +1,103 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#include "utilities/transactions/lock/point/point_lock_manager_test.h"
+#include "utilities/transactions/lock/point/point_lock_validation_test_runner.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+struct PointLockCorrectnessCheckTestParam {
+  bool is_per_key_point_lock_manager;
+  uint32_t thread_count;
+  uint32_t key_count;
+  uint32_t max_num_keys_to_lock_per_txn;
+  uint32_t execution_time_sec;
+  LockTypeToTest lock_type;
+  int64_t lock_timeout_us;
+  int64_t lock_expiration_us;
+  bool allow_non_deadlock_error;
+  // to simulate some useful work
+  uint32_t max_sleep_after_lock_acquisition_ms;
+};
+
+class PointLockCorrectnessCheckTest
+    : public PointLockManagerTest,
+      public testing::WithParamInterface<PointLockCorrectnessCheckTestParam> {
+ public:
+  void SetUp() override {
+    init();
+    auto const& param = GetParam();
+    auto per_key_lock_manager = param.is_per_key_point_lock_manager;
+    if (per_key_lock_manager) {
+      locker_ = std::make_shared<PerKeyPointLockManager>(
+          static_cast<PessimisticTransactionDB*>(db_), txndb_opt_);
+    } else {
+      locker_ = std::make_shared<PointLockManager>(
+          static_cast<PessimisticTransactionDB*>(db_), txndb_opt_);
+    }
+
+    txn_opt_.deadlock_detect = true;
+    txn_opt_.lock_timeout = param.lock_timeout_us;
+    txn_opt_.expiration = param.lock_expiration_us;
+  }
+
+ protected:
+  TransactionOptions txn_opt_;
+};
+
+TEST_P(PointLockCorrectnessCheckTest, LockCorrectnessValidation) {
+  auto const& param = GetParam();
+  PointLockValidationTestRunner test_runner(
+      env_, txndb_opt_, locker_, db_, txn_opt_, param.thread_count,
+      param.key_count, param.max_num_keys_to_lock_per_txn,
+      param.execution_time_sec, static_cast<LockTypeToTest>(param.lock_type),
+      param.allow_non_deadlock_error,
+      param.max_sleep_after_lock_acquisition_ms);
+  test_runner.run();
+}
+
+constexpr auto X_S_LOCK = LockTypeToTest::EXCLUSIVE_AND_SHARED;
+constexpr auto X_LOCK = LockTypeToTest::EXCLUSIVE_ONLY;
+constexpr auto S_LOCK = LockTypeToTest::SHARED_ONLY;
+
+INSTANTIATE_TEST_CASE_P(
+    PointLockCorrectnessCheckTestSuite, PointLockCorrectnessCheckTest,
+    ::testing::ValuesIn(std::vector<PointLockCorrectnessCheckTestParam>{
+        // 2 second timeout and no expiration simulates myrocks default
+        // configuration
+        {true, 16, 16, 8, 10, X_S_LOCK, 2000, -1, true, 0},
+        {false, 16, 16, 8, 10, X_S_LOCK, 2000, -1, true, 0},
+        {true, 16, 16, 8, 10, X_LOCK, 2000, -1, true, 0},
+        {false, 16, 16, 8, 10, X_LOCK, 2000, -1, true, 0},
+        {true, 16, 16, 8, 10, S_LOCK, 2000, -1, true, 0},
+        {false, 16, 16, 8, 10, S_LOCK, 2000, -1, true, 0},
+        // short timeout and expiration to test lock stealing
+        {true, 16, 16, 8, 10, X_S_LOCK, 10, 10, true, 10},
+        {false, 16, 16, 8, 10, X_S_LOCK, 10, 10, true, 10},
+        {true, 16, 16, 8, 10, X_LOCK, 10, 10, true, 10},
+        {false, 16, 16, 8, 10, X_LOCK, 10, 10, true, 10},
+        {true, 16, 16, 8, 10, S_LOCK, 10, 10, true, 10},
+        {false, 16, 16, 8, 10, S_LOCK, 10, 10, true, 10},
+        // long timeout and expiration to test deadlock detection without
+        // timeout
+        {true, 16, 16, 8, 10, X_S_LOCK, 100000, 100000, false, 0},
+        {false, 16, 16, 8, 10, X_S_LOCK, 100000, 100000, false, 0},
+        {true, 16, 16, 8, 10, X_LOCK, 100000, 100000, false, 0},
+        {false, 16, 16, 8, 10, X_LOCK, 100000, 100000, false, 0},
+        {true, 16, 16, 8, 10, S_LOCK, 100000, 100000, false, 0},
+        {false, 16, 16, 8, 10, S_LOCK, 100000, 100000, false, 0},
+        // Low lock contention
+        {true, 4, 1024 * 1024, 2, 10, S_LOCK, 100000, 100000, false, 0},
+        {false, 4, 1024 * 1024, 2, 10, S_LOCK, 100000, 100000, false, 0},
+    }));
+
+}  // namespace ROCKSDB_NAMESPACE
+
+int main(int argc, char** argv) {
+  ROCKSDB_NAMESPACE::port::InstallStackTraceHandler();
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

utilities/transactions/lock/point/point_lock_manager_test.h

@@ -4,321 +4,99 @@
 //  COPYING file in the root directory) and Apache 2.0 License
 //  (found in the LICENSE.Apache file in the root directory).
 
+#pragma once
+
 #include "file/file_util.h"
-#include "port/port.h"
-#include "port/stack_trace.h"
 #include "rocksdb/utilities/transaction_db.h"
 #include "test_util/testharness.h"
-#include "test_util/testutil.h"
 #include "utilities/transactions/lock/point/point_lock_manager.h"
+#include "utilities/transactions/lock/point/point_lock_manager_test_common.h"
 #include "utilities/transactions/pessimistic_transaction_db.h"
-#include "utilities/transactions/transaction_db_mutex_impl.h"
 
 namespace ROCKSDB_NAMESPACE {
 
-class MockColumnFamilyHandle : public ColumnFamilyHandle {
- public:
-  explicit MockColumnFamilyHandle(ColumnFamilyId cf_id) : cf_id_(cf_id) {}
-
-  ~MockColumnFamilyHandle() override {}
-
-  const std::string& GetName() const override { return name_; }
-
-  ColumnFamilyId GetID() const override { return cf_id_; }
-
-  Status GetDescriptor(ColumnFamilyDescriptor*) override {
-    return Status::OK();
-  }
-
-  const Comparator* GetComparator() const override {
-    return BytewiseComparator();
-  }
-
- private:
-  ColumnFamilyId cf_id_;
-  std::string name_ = "MockCF";
-};
-
 class PointLockManagerTest : public testing::Test {
  public:
-  void SetUp() override {
+  void init() {
     env_ = Env::Default();
     db_dir_ = test::PerThreadDBPath("point_lock_manager_test");
     ASSERT_OK(env_->CreateDir(db_dir_));
 
     Options opt;
     opt.create_if_missing = true;
-    TransactionDBOptions txn_opt;
-    txn_opt.transaction_lock_timeout = 0;
+    // Reduce the number of stripes to 4 to increase contention in test
+    txndb_opt_.num_stripes = 4;
+    txndb_opt_.transaction_lock_timeout = 0;
 
-    ASSERT_OK(TransactionDB::Open(opt, txn_opt, db_dir_, &db_));
-
-    // CAUTION: This test creates a separate lock manager object (right, NOT
-    // the one that the TransactionDB is using!), and runs tests on it.
-    locker_.reset(new PointLockManager(
-        static_cast<PessimisticTransactionDB*>(db_), txn_opt));
+    ASSERT_OK(TransactionDB::Open(opt, txndb_opt_, db_dir_, &db_));
 
     wait_sync_point_name_ = "PointLockManager::AcquireWithTimeout:WaitingTxn";
   }
 
+  void SetUp() override {
+    init();
+    // CAUTION: This test creates a separate lock manager object (right, NOT
+    // the one that the TransactionDB is using!), and runs tests on it.
+    locker_.reset(new PointLockManager(
+        static_cast<PessimisticTransactionDB*>(db_), txndb_opt_));
+  }
+
   void TearDown() override {
+    std::string errmsg;
+    auto no_lock_held = verifyNoLocksHeld(locker_, errmsg);
+    ASSERT_TRUE(no_lock_held) << errmsg;
     delete db_;
     EXPECT_OK(DestroyDir(env_, db_dir_));
   }
 
   PessimisticTransaction* NewTxn(
       TransactionOptions txn_opt = TransactionOptions()) {
+    // override deadlock_timeout_us;
+    txn_opt.deadlock_timeout_us = deadlock_timeout_us;
     Transaction* txn = db_->BeginTransaction(WriteOptions(), txn_opt);
     return static_cast<PessimisticTransaction*>(txn);
   }
 
+  int64_t deadlock_timeout_us = 0;
+
+  void UsePerKeyPointLockManager() {
+    locker_.reset(new PerKeyPointLockManager(
+        static_cast<PessimisticTransactionDB*>(db_), txndb_opt_));
+  }
+
  protected:
   Env* env_;
+  TransactionDBOptions txndb_opt_;
   std::shared_ptr<LockManager> locker_;
   const char* wait_sync_point_name_;
   friend void PointLockManagerTestExternalSetup(PointLockManagerTest*);
 
- private:
   std::string db_dir_;
   TransactionDB* db_;
 };
 
-using init_func_t = void (*)(PointLockManagerTest*);
-
-class AnyLockManagerTest : public PointLockManagerTest,
-                           public testing::WithParamInterface<init_func_t> {
- public:
-  void SetUp() override {
-    // If a custom setup function was provided, use it. Otherwise, use what we
-    // have inherited.
-    auto init_func = GetParam();
-    if (init_func)
-      (*init_func)(this);
-    else
-      PointLockManagerTest::SetUp();
-  }
-};
-
-TEST_P(AnyLockManagerTest, ReentrantExclusiveLock) {
-  // Tests that a txn can acquire exclusive lock on the same key repeatedly.
-  MockColumnFamilyHandle cf(1);
-  locker_->AddColumnFamily(&cf);
-  auto txn = NewTxn();
-  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, true));
-  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, true));
-
-  // Cleanup
-  locker_->UnLock(txn, 1, "k", env_);
-
-  delete txn;
-}
-
-TEST_P(AnyLockManagerTest, ReentrantSharedLock) {
-  // Tests that a txn can acquire shared lock on the same key repeatedly.
-  MockColumnFamilyHandle cf(1);
-  locker_->AddColumnFamily(&cf);
-  auto txn = NewTxn();
-  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, false));
-  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, false));
-
-  // Cleanup
-  locker_->UnLock(txn, 1, "k", env_);
-
-  delete txn;
-}
-
-TEST_P(AnyLockManagerTest, LockUpgrade) {
-  // Tests that a txn can upgrade from a shared lock to an exclusive lock.
-  MockColumnFamilyHandle cf(1);
-  locker_->AddColumnFamily(&cf);
-  auto txn = NewTxn();
-  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, false));
-  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, true));
-
-  // Cleanup
-  locker_->UnLock(txn, 1, "k", env_);
-  delete txn;
-}
-
-TEST_P(AnyLockManagerTest, LockDowngrade) {
-  // Tests that a txn can acquire a shared lock after acquiring an exclusive
-  // lock on the same key.
-  MockColumnFamilyHandle cf(1);
-  locker_->AddColumnFamily(&cf);
-  auto txn = NewTxn();
-  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, true));
-  ASSERT_OK(locker_->TryLock(txn, 1, "k", env_, false));
-
-  // Cleanup
-  locker_->UnLock(txn, 1, "k", env_);
-  delete txn;
-}
-
-TEST_P(AnyLockManagerTest, LockConflict) {
-  // Tests that lock conflicts lead to lock timeout.
-  MockColumnFamilyHandle cf(1);
-  locker_->AddColumnFamily(&cf);
-  auto txn1 = NewTxn();
-  auto txn2 = NewTxn();
-
-  {
-    // exclusive-exclusive conflict.
-    ASSERT_OK(locker_->TryLock(txn1, 1, "k1", env_, true));
-    auto s = locker_->TryLock(txn2, 1, "k1", env_, true);
-    ASSERT_TRUE(s.IsTimedOut());
-  }
-
-  {
-    // exclusive-shared conflict.
-    ASSERT_OK(locker_->TryLock(txn1, 1, "k2", env_, true));
-    auto s = locker_->TryLock(txn2, 1, "k2", env_, false);
-    ASSERT_TRUE(s.IsTimedOut());
-  }
-
-  {
-    // shared-exclusive conflict.
-    ASSERT_OK(locker_->TryLock(txn1, 1, "k2", env_, false));
-    auto s = locker_->TryLock(txn2, 1, "k2", env_, true);
-    ASSERT_TRUE(s.IsTimedOut());
-  }
-
-  // Cleanup
-  locker_->UnLock(txn1, 1, "k1", env_);
-  locker_->UnLock(txn1, 1, "k2", env_);
-
-  delete txn1;
-  delete txn2;
-}
-
-port::Thread BlockUntilWaitingTxn(const char* sync_point_name,
-                                  std::function<void()> f) {
+void BlockUntilWaitingTxn(const char* sync_point_name, port::Thread& t,
+                          std::function<void()> f) {
   std::atomic<bool> reached(false);
   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->SetCallBack(
       sync_point_name, [&](void* /*arg*/) { reached.store(true); });
   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->EnableProcessing();
 
-  port::Thread t(f);
+  t = port::Thread(f);
 
-  while (!reached.load()) {
-    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+  // timeout after 30 seconds, so test does not hang forever
+  // 30 seconds should be enough for the test to reach the expected state
+  // without causing too much flakiness
+  for (int i = 0; i < 3000; i++) {
+    if (reached.load()) {
+      break;
+    }
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
   }
+
+  ASSERT_TRUE(reached.load());
   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->DisableProcessing();
   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->ClearAllCallBacks();
-
-  return t;
-}
-
-TEST_P(AnyLockManagerTest, SharedLocks) {
-  // Tests that shared locks can be concurrently held by multiple transactions.
-  MockColumnFamilyHandle cf(1);
-  locker_->AddColumnFamily(&cf);
-  auto txn1 = NewTxn();
-  auto txn2 = NewTxn();
-  ASSERT_OK(locker_->TryLock(txn1, 1, "k", env_, false));
-  ASSERT_OK(locker_->TryLock(txn2, 1, "k", env_, false));
-
-  // Cleanup
-  locker_->UnLock(txn1, 1, "k", env_);
-  locker_->UnLock(txn2, 1, "k", env_);
-
-  delete txn1;
-  delete txn2;
-}
-
-TEST_P(AnyLockManagerTest, Deadlock) {
-  // Tests that deadlock can be detected.
-  // Deadlock scenario:
-  // txn1 exclusively locks k1, and wants to lock k2;
-  // txn2 exclusively locks k2, and wants to lock k1.
-  MockColumnFamilyHandle cf(1);
-  locker_->AddColumnFamily(&cf);
-  TransactionOptions txn_opt;
-  txn_opt.deadlock_detect = true;
-  txn_opt.lock_timeout = 1000000;
-  auto txn1 = NewTxn(txn_opt);
-  auto txn2 = NewTxn(txn_opt);
-
-  ASSERT_OK(locker_->TryLock(txn1, 1, "k1", env_, true));
-  ASSERT_OK(locker_->TryLock(txn2, 1, "k2", env_, true));
-
-  // txn1 tries to lock k2, will block forever.
-  port::Thread t = BlockUntilWaitingTxn(wait_sync_point_name_, [&]() {
-    // block because txn2 is holding a lock on k2.
-    ASSERT_OK(locker_->TryLock(txn1, 1, "k2", env_, true));
-  });
-
-  auto s = locker_->TryLock(txn2, 1, "k1", env_, true);
-  ASSERT_TRUE(s.IsBusy());
-  ASSERT_EQ(s.subcode(), Status::SubCode::kDeadlock);
-
-  std::vector<DeadlockPath> deadlock_paths = locker_->GetDeadlockInfoBuffer();
-  ASSERT_EQ(deadlock_paths.size(), 1u);
-  ASSERT_FALSE(deadlock_paths[0].limit_exceeded);
-
-  std::vector<DeadlockInfo> deadlocks = deadlock_paths[0].path;
-  ASSERT_EQ(deadlocks.size(), 2u);
-
-  ASSERT_EQ(deadlocks[0].m_txn_id, txn1->GetID());
-  ASSERT_EQ(deadlocks[0].m_cf_id, 1u);
-  ASSERT_TRUE(deadlocks[0].m_exclusive);
-  ASSERT_EQ(deadlocks[0].m_waiting_key, "k2");
-
-  ASSERT_EQ(deadlocks[1].m_txn_id, txn2->GetID());
-  ASSERT_EQ(deadlocks[1].m_cf_id, 1u);
-  ASSERT_TRUE(deadlocks[1].m_exclusive);
-  ASSERT_EQ(deadlocks[1].m_waiting_key, "k1");
-
-  locker_->UnLock(txn2, 1, "k2", env_);
-  t.join();
-
-  // Cleanup
-  locker_->UnLock(txn1, 1, "k1", env_);
-  locker_->UnLock(txn1, 1, "k2", env_);
-  delete txn2;
-  delete txn1;
-}
-
-TEST_P(AnyLockManagerTest, GetWaitingTxns_MultipleTxns) {
-  MockColumnFamilyHandle cf(1);
-  locker_->AddColumnFamily(&cf);
-
-  auto txn1 = NewTxn();
-  ASSERT_OK(locker_->TryLock(txn1, 1, "k", env_, false));
-
-  auto txn2 = NewTxn();
-  ASSERT_OK(locker_->TryLock(txn2, 1, "k", env_, false));
-
-  auto txn3 = NewTxn();
-  txn3->SetLockTimeout(10000);
-  port::Thread t1 = BlockUntilWaitingTxn(wait_sync_point_name_, [&]() {
-    ASSERT_OK(locker_->TryLock(txn3, 1, "k", env_, true));
-    locker_->UnLock(txn3, 1, "k", env_);
-  });
-
-  // Ok, now txn3 is waiting for lock on "k", which is owned by two
-  // transactions. Check that GetWaitingTxns reports this correctly
-  uint32_t wait_cf_id;
-  std::string wait_key;
-  auto waiters = txn3->GetWaitingTxns(&wait_cf_id, &wait_key);
-
-  ASSERT_EQ(wait_cf_id, 1u);
-  ASSERT_EQ(wait_key, "k");
-  ASSERT_EQ(waiters.size(), 2);
-  bool waits_correct =
-      (waiters[0] == txn1->GetID() && waiters[1] == txn2->GetID()) ||
-      (waiters[1] == txn1->GetID() && waiters[0] == txn2->GetID());
-  ASSERT_EQ(waits_correct, true);
-
-  // Release locks so txn3 can proceed with execution
-  locker_->UnLock(txn1, 1, "k", env_);
-  locker_->UnLock(txn2, 1, "k", env_);
-
-  // Wait until txn3 finishes
-  t1.join();
-
-  delete txn1;
-  delete txn2;
-  delete txn3;
 }
 
 }  // namespace ROCKSDB_NAMESPACE

utilities/transactions/lock/point/point_lock_manager_test_common.h

@@ -0,0 +1,78 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#pragma once
+
+#include <sstream>
+
+#include "rocksdb/db.h"
+#include "utilities/transactions/lock/lock_manager.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+constexpr auto kLongTxnTimeoutMs = 100000;
+constexpr auto kShortTxnTimeoutMs = 100;
+
+class MockColumnFamilyHandle : public ColumnFamilyHandle {
+ public:
+  explicit MockColumnFamilyHandle(ColumnFamilyId cf_id) : cf_id_(cf_id) {}
+
+  // disable copy and assignment
+  MockColumnFamilyHandle(const MockColumnFamilyHandle&) = delete;
+  MockColumnFamilyHandle& operator=(const MockColumnFamilyHandle&) = delete;
+  // disable move
+  MockColumnFamilyHandle(MockColumnFamilyHandle&&) = delete;
+  MockColumnFamilyHandle& operator=(MockColumnFamilyHandle&&) = delete;
+
+  ~MockColumnFamilyHandle() override {}
+
+  const std::string& GetName() const override { return name_; }
+
+  ColumnFamilyId GetID() const override { return cf_id_; }
+
+  Status GetDescriptor(ColumnFamilyDescriptor*) override {
+    return Status::OK();
+  }
+
+  const Comparator* GetComparator() const override {
+    return BytewiseComparator();
+  }
+
+ private:
+  ColumnFamilyId cf_id_;
+  std::string name_ = "MockCF";
+};
+
+// Verify no lock was held. Return true, if success. False, if there is. Set
+// error message on False.
+bool verifyNoLocksHeld(std::shared_ptr<LockManager>& locker,
+                       std::string& errmsg) {
+  // Validate no lock was held at the end of the test
+  auto lock_status = locker->GetPointLockStatus();
+  // print the lock status for debugging
+  std::stringstream ss;
+  for (auto& s : lock_status) {
+    ss << "id " << s.first;
+    ss << " key " << s.second.key;
+    ss << " type " << (s.second.exclusive ? "exclusive" : "shared");
+    ss << " txn ids [";
+    for (auto& t : s.second.ids) {
+      ss << t << ",";
+    }
+    ss << "]";
+    ss << std::endl;
+  }
+
+  if (!lock_status.empty()) {
+    errmsg = std::to_string(lock_status.size()) +
+             " locks were held at the end. " + ss.str();
+    return false;
+  }
+
+  return true;
+}
+
+}  // namespace ROCKSDB_NAMESPACE

utilities/transactions/lock/point/point_lock_validation_test_runner.h

@@ -0,0 +1,466 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#pragma once
+
+#include <cinttypes>
+#include <cstddef>
+#include <cstdio>
+#include <iostream>
+#include <memory>
+
+#include "rocksdb/convenience.h"
+#include "rocksdb/db.h"
+#include "rocksdb/env.h"
+#include "rocksdb/utilities/transaction_db.h"
+#include "utilities/transactions/lock/lock_manager.h"
+#include "utilities/transactions/lock/point/point_lock_manager_test_common.h"
+#include "utilities/transactions/pessimistic_transaction.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+constexpr bool kDebugLog = false;
+
+// Since this code is executed both with and without gtest, it supports assert
+// with different ways.
+#ifdef ASSERT_TRUE
+#define ASSERT_TRUE_WITH_MSG(expr, errmsg) ASSERT_TRUE(expr) << (errmsg)
+#else
+#define ASSERT_TRUE_WITH_MSG(expr, errmsg)                             \
+  if (!(expr)) {                                                       \
+    std::cerr << "Assert true failed with error message: " << (errmsg) \
+              << std::endl;                                            \
+    abort();                                                           \
+  }
+#endif
+
+#ifndef ASSERT_OK
+#define ASSERT_OK(s) \
+  ASSERT_TRUE_WITH_MSG(s.ok(), "Failed with " + s.ToString());
+#endif
+
+#define ASSERT_TRUE_WITH_INFO(X) \
+  ASSERT_TRUE_WITH_MSG(          \
+      (X), " Txn " + std::to_string(txn_id) + " key " + std::to_string(key))
+
+#define ASSERT_EQ_WITH_INFO(X, Y) ASSERT_TRUE_WITH_INFO((X) == (Y))
+
+#define DEBUG_LOG(...)            \
+  if (kDebugLog) {                \
+    fprintf(stderr, __VA_ARGS__); \
+    fflush(stderr);               \
+  }
+
+#define DEBUG_LOG_WITH_PREFIX(format, ...) \
+  DEBUG_LOG("Txn %" PRIu64 " " format, txn_id, ##__VA_ARGS__);
+
+enum class LockTypeToTest : int8_t {
+  EXCLUSIVE_ONLY = 0,
+  SHARED_ONLY = 1,
+  EXCLUSIVE_AND_SHARED = 2,
+};
+
+struct KeyStatus {
+  KeyStatus(uint32_t k, bool ex, int v) : key(k), exclusive(ex), value(v) {}
+  uint32_t key;
+  bool exclusive;
+  int value;
+};
+
+class PointLockValidationTestRunner {
+ public:
+  PointLockValidationTestRunner(
+      Env* env, TransactionDBOptions txndb_opt,
+      std::shared_ptr<LockManager> locker, TransactionDB* db,
+      TransactionOptions txn_opt, uint32_t thd_cnt, uint32_t key_cnt,
+      uint32_t max_num_keys_to_lock_per_txn, uint32_t execution_time_sec,
+      LockTypeToTest lock_type, bool allow_non_deadlock_error,
+      uint32_t max_sleep_after_lock_acquisition_ms,
+      bool enable_per_thread_lock_count_assertion = false)
+      : env_(env),
+        txndb_opt_(std::move(txndb_opt)),
+        locker_(std::move(locker)),
+        db_(db),
+        txn_opt_(std::move(txn_opt)),
+        thread_count_(thd_cnt),
+        key_count_(key_cnt),
+        max_num_keys_to_lock_per_txn_(max_num_keys_to_lock_per_txn),
+        execution_time_sec_(execution_time_sec),
+        lock_type_(lock_type),
+        allow_non_deadlock_error_(allow_non_deadlock_error),
+        max_sleep_after_lock_acquisition_ms_(
+            max_sleep_after_lock_acquisition_ms),
+        enable_per_thread_lock_count_assertion_(
+            enable_per_thread_lock_count_assertion),
+        shutdown_(false) {
+    // Only enable lock status validation when lock expiration/stealing isk
+    // disabled.
+    enable_lock_status_validation_ = txn_opt_.expiration == -1;
+    values_.resize(key_count_, 0);
+    exclusive_lock_status_.resize(key_count_, 0);
+
+    // init counters and values
+    for (size_t i = 0; i < key_count_; i++) {
+      counters_.emplace_back(std::make_unique<std::atomic_int>(0));
+      shared_lock_count_.emplace_back(std::make_unique<std::atomic_int>(0));
+    }
+
+    for (size_t i = 0; i < thread_count_; i++) {
+      num_of_locks_acquired_per_thread_.emplace_back(
+          std::make_unique<std::atomic_int64_t>(0));
+    }
+  }
+
+  // Decide which lock type to acquire
+  // If the key is already locked and only one type of locks to be tested,
+  // return false, so caller could try to lock a different key.
+  // Otherwise, return true.
+  bool DecideLockType(
+      bool& acquire_exclusive_lock, uint32_t key,
+      std::unordered_map<uint32_t, KeyStatus>& locked_key_status,
+      bool& isUpgrade, bool& isDowngrade) {
+    // Decide lock type
+    acquire_exclusive_lock = Random::GetTLSInstance()->OneIn(2);
+
+    // check whether a lock on the same key is already held
+    auto it = locked_key_status.find(key);
+    if (it != locked_key_status.end()) {
+      // a lock on the same key is already held.
+      if (lock_type_ == LockTypeToTest::EXCLUSIVE_AND_SHARED) {
+        // if test both shared and exclusive locks, switch their type
+        if (it->second.exclusive == false) {
+          // If it is a shared lock, upgrade to an exclusive lock
+          acquire_exclusive_lock = true;
+          isUpgrade = true;
+        } else {
+          // If it is an exclusive lock, downgrade to a shared lock
+          acquire_exclusive_lock = false;
+          isDowngrade = true;
+        }
+      } else {
+        // Only one type of lock to test, and the key is already locked,
+        return false;
+      }
+    }
+
+    // This is a new key to lock or the lock type is switched.
+    if (lock_type_ != LockTypeToTest::EXCLUSIVE_AND_SHARED) {
+      // if only one type of locks to be acquired, update its type
+      acquire_exclusive_lock = (lock_type_ == LockTypeToTest::EXCLUSIVE_ONLY);
+    }
+    return true;
+  }
+
+  void run() {
+    // Verify lock guarantee. Exclusive lock provide unique access guarantee.
+    // Shared lock provide shared access guarantee.
+    // Create multiple threads. Each try to grab a lock with random type on
+    // random key.
+
+    // To validate lock exclusive guarantee, each key has a value and a counter
+    // used for tracking the number of exclusive locks have been acquired on it
+    // in each test run across all threads.
+
+    // Every time an exclusive lock is acquired, both the counter and the value
+    // are bumped by 1. The difference between the counter and the value is that
+    // counter is atomic, so it is guaranteed that it would not lose update,
+    // while value is not atomic. Its correctness is only guaranteed by the
+    // exclusiveness provided by the lock manager which is being tested. If the
+    // lock manager does not guarantee exclusiveness, the value would lose
+    // update, and the counter would mismatch with the value, which fails the
+    // test.
+
+    // To validate lock shared guarantee, after a shared lock is acquired, the
+    // counter and value are read and stored in a local variable inside the
+    // thread. Before the lock is released, the local copy is compared against
+    // the counter and value. If they mismatch, it means the shared lock
+    // guaranteed is violated.
+
+    MockColumnFamilyHandle cf(1);
+    locker_->AddColumnFamily(&cf);
+
+    for (uint32_t thd_idx = 0; thd_idx < thread_count_; thd_idx++) {
+      threads_.emplace_back([this, thd_idx]() {
+        auto txn = static_cast<PessimisticTransaction*>(
+            db_->BeginTransaction(WriteOptions(), txn_opt_));
+        auto txn_id = txn->GetID();
+        DEBUG_LOG_WITH_PREFIX("Thd %" PRIu32 " new txn\n", thd_idx);
+        while (!shutdown_) {
+          std::unordered_map<uint32_t, KeyStatus> locked_key_status;
+          auto num_key_to_lock = max_num_keys_to_lock_per_txn_;
+          Status s;
+
+          for (uint32_t j = 0; j < num_key_to_lock; j++) {
+            uint32_t key = 0;
+            key = Random::GetTLSInstance()->Uniform(key_count_);
+            auto key_str = std::to_string(key);
+            bool isUpgrade = false;
+            bool isDowngrade = false;
+            bool exclusive_lock_type;
+
+            if (!DecideLockType(exclusive_lock_type, key, locked_key_status,
+                                isUpgrade, isDowngrade)) {
+              // try a different key
+              j--;
+              continue;
+            }
+
+            if (enable_lock_status_validation_) {
+              if (isDowngrade) {
+                // Before downgrade, validate the lock is in exlusive status
+                // This could not be done after downgrade, as another thread
+                // could take a shared lock and update lock status
+                ASSERT_TRUE_WITH_INFO(exclusive_lock_status_[key]);
+                ASSERT_EQ_WITH_INFO(*shared_lock_count_[key], 0);
+                // for downgrade, update the lock status before acquiring the
+                // lock, as afterwards, it will not have exclusive access to it
+                exclusive_lock_status_[key] = 0;
+              }
+            }
+
+            // try to acquire the lock
+            DEBUG_LOG_WITH_PREFIX("try to acquire lock %" PRIu32 " type %s\n",
+                                  key,
+                                  exclusive_lock_type ? "exclusive" : "shared");
+            s = locker_->TryLock(txn, 1, key_str, env_, exclusive_lock_type);
+
+            if (s.ok()) {
+              DEBUG_LOG_WITH_PREFIX(
+                  "acquired lock %" PRIu32 " type %s\n", key,
+                  exclusive_lock_type ? "exclusive" : "shared");
+
+              auto it = locked_key_status.find(key);
+              if (isUpgrade || isDowngrade) {
+                // If it is either upgrade or downgrade, the key should exist
+                // already.
+                ASSERT_TRUE_WITH_INFO(it != locked_key_status.end());
+              } else {
+                locked_key_status.emplace(
+                    std::piecewise_construct, std::forward_as_tuple(key),
+                    std::forward_as_tuple(key, exclusive_lock_type,
+                                          values_[key]));
+              }
+              // update local lock status
+              if (exclusive_lock_type) {
+                if (isUpgrade) {
+                  it->second.exclusive = true;
+                }
+                num_of_exclusive_locks_acquired_++;
+              } else {
+                if (isDowngrade) {
+                  it->second.exclusive = false;
+                }
+                num_of_shared_locks_acquired_++;
+              }
+              num_of_locks_acquired_++;
+              (*num_of_locks_acquired_per_thread_[thd_idx])++;
+
+              if (enable_lock_status_validation_) {
+                if (exclusive_lock_type) {
+                  // validate the lock is not in exclusive status
+                  ASSERT_TRUE_WITH_INFO(!exclusive_lock_status_[key]);
+                  if (isUpgrade) {
+                    // validate the lock is in shared status and only had one
+                    // shared lock
+                    ASSERT_EQ_WITH_INFO(*shared_lock_count_[key], 1);
+                    shared_lock_count_[key]->fetch_sub(1);
+                  } else {
+                    ASSERT_EQ_WITH_INFO(*shared_lock_count_[key], 0);
+                  }
+                  // update the lock status
+                  exclusive_lock_status_[key] = 1;
+                } else {
+                  shared_lock_count_[key]->fetch_add(1);
+                  ASSERT_TRUE_WITH_INFO(!exclusive_lock_status_[key]);
+                }
+              }
+            } else {
+              if (!allow_non_deadlock_error_) {
+                ASSERT_TRUE_WITH_INFO(s.IsDeadlock());
+              }
+              if (s.IsDeadlock()) {
+                DEBUG_LOG_WITH_PREFIX(
+                    "detected deadlock on key %" PRIu32 ", abort\n", key);
+                num_of_deadlock_detected_++;
+                // for deadlock, release all locks acquired
+                break;
+              } else {
+                // for other errors, try again
+                DEBUG_LOG_WITH_PREFIX("failed to acquire lock on key %" PRIu32
+                                      ", due to "
+                                      "%s, "
+                                      "abort\n",
+                                      key, s.ToString().c_str());
+              }
+            }
+          }
+
+          // After all of the locks are acquired, try to sleep a bit to simulate
+          // some useful work to be done
+          if (max_sleep_after_lock_acquisition_ms_ != 0 && s.ok()) {
+            auto sleep_time_us = Random::GetTLSInstance()->Uniform(
+                static_cast<uint32_t>(max_sleep_after_lock_acquisition_ms_));
+            std::this_thread::sleep_for(
+                std::chrono::milliseconds(sleep_time_us));
+          }
+
+          // release all locks
+          for (const auto& pair : locked_key_status) {
+            auto key_status = pair.second;
+            auto key = key_status.key;
+            ASSERT_TRUE_WITH_INFO(key < key_count_);
+            if (enable_lock_status_validation_) {
+              ASSERT_EQ_WITH_INFO(counters_[key]->load(), values_[key]);
+              auto exclusive = key_status.exclusive;
+              if (exclusive) {
+                // for exclusive lock, bump the value by 1
+                (*counters_[key])++;
+                values_[key]++;
+                DEBUG_LOG_WITH_PREFIX("bump key %" PRIu32 " by 1 to %d\n", key,
+                                      values_[key]);
+                ASSERT_EQ_WITH_INFO(counters_[key]->load(), values_[key]);
+              } else {
+                // shared lock, validate the value has not changed since it was
+                // read
+                ASSERT_EQ_WITH_INFO(counters_[key]->load(), key_status.value);
+                ASSERT_EQ_WITH_INFO(values_[key], key_status.value);
+              }
+              if (exclusive) {
+                ASSERT_TRUE_WITH_INFO(exclusive_lock_status_[key]);
+                ASSERT_EQ_WITH_INFO(*shared_lock_count_[key], 0);
+                exclusive_lock_status_[key] = 0;
+              } else {
+                ASSERT_TRUE_WITH_INFO(!exclusive_lock_status_[key]);
+                ASSERT_TRUE_WITH_INFO(shared_lock_count_[key]->fetch_sub(1) >=
+                                      1);
+              }
+            }
+            DEBUG_LOG_WITH_PREFIX("release lock %" PRIu32 "\n", key);
+            locker_->UnLock(txn, 1, std::to_string(key), env_);
+          }
+        }
+        delete txn;
+      });
+    }
+
+    // run test for a few seconds
+    // print progress
+    auto prev_num_of_locks_acquired = num_of_locks_acquired_.load();
+    std::vector<int64_t> prev_num_of_locks_acquired_per_thread(thread_count_,
+                                                               0);
+    int64_t measured_locks_acquired = 0;
+    for (uint32_t i = 0; i < execution_time_sec_; i++) {
+      std::this_thread::sleep_for(std::chrono::seconds(1));
+      auto num_of_locks_acquired = num_of_locks_acquired_.load();
+      DEBUG_LOG("num_of_locks_acquired: %" PRId64 "\n", num_of_locks_acquired);
+      DEBUG_LOG("num_of_exclusive_locks_acquired: %" PRId64 "\n",
+                num_of_exclusive_locks_acquired_.load());
+      DEBUG_LOG("num_of_shared_locks_acquired: %" PRId64 "\n",
+                num_of_shared_locks_acquired_.load());
+      DEBUG_LOG("num_of_deadlock_detected: %" PRId64 "\n",
+                num_of_deadlock_detected_.load());
+      ASSERT_TRUE_WITH_MSG(num_of_locks_acquired > prev_num_of_locks_acquired,
+                           "No locks were acquired in the last 1 second");
+      for (uint32_t thd_idx = 0; thd_idx < thread_count_; thd_idx++) {
+        auto num_of_locks_acquired_per_thread =
+            num_of_locks_acquired_per_thread_[thd_idx]->load();
+        DEBUG_LOG("thread: %" PRIu32 " acquired %" PRId64 " locks\n", thd_idx,
+                  num_of_locks_acquired_per_thread);
+        if (enable_per_thread_lock_count_assertion_) {
+          ASSERT_TRUE_WITH_MSG(
+              num_of_locks_acquired_per_thread >
+                  prev_num_of_locks_acquired_per_thread[thd_idx],
+              "No locks were acquired in the last 1 second on thread " +
+                  std::to_string(thd_idx));
+        }
+        prev_num_of_locks_acquired_per_thread[thd_idx] =
+            num_of_locks_acquired_per_thread;
+      }
+      prev_num_of_locks_acquired = num_of_locks_acquired;
+      if (i == 0) {
+        measured_locks_acquired = num_of_locks_acquired;
+      }
+      if (i == execution_time_sec_ - 1) {
+        measured_locks_acquired =
+            num_of_locks_acquired - measured_locks_acquired;
+        // Skip the first second, as threads are warming up
+        printf("measured_num_of_locks_acquired: %" PRId64 "\n",
+               measured_locks_acquired / (execution_time_sec_ - 1));
+      }
+    }
+
+    shutdown_ = true;
+    for (auto& t : threads_) {
+      t.join();
+    }
+
+    // validate values against counters
+    for (uint32_t i = 0; i < key_count_; i++) {
+      ASSERT_TRUE_WITH_MSG(counters_[i]->load() == values_[i],
+                           "Exclusive lock guarantee is violated.");
+    }
+
+    ASSERT_TRUE_WITH_MSG(num_of_locks_acquired_.load() >= 0,
+                         "No lock were acquired at all");
+    printf("num_of_locks_acquired: %" PRId64 "\n",
+           num_of_locks_acquired_.load());
+
+    std::string errmsg;
+    auto no_lock_held = verifyNoLocksHeld(locker_, errmsg);
+    ASSERT_TRUE_WITH_MSG(no_lock_held, errmsg);
+  }
+
+ private:
+  // test configuration
+  Env* env_;
+  TransactionDBOptions txndb_opt_;
+  std::shared_ptr<LockManager> locker_;
+
+  TransactionDB* db_;
+  TransactionOptions txn_opt_;
+
+  uint32_t thread_count_;
+  uint32_t key_count_;
+  uint32_t max_num_keys_to_lock_per_txn_;
+  uint32_t execution_time_sec_;
+  LockTypeToTest lock_type_;
+  bool allow_non_deadlock_error_;
+  uint32_t max_sleep_after_lock_acquisition_ms_;
+
+  // In some of the test run, due to debug or ASAN build and short lock timeout,
+  // a thread may not be able to acquire any lock within a second. So skip this
+  // assertion by default. However, this could be useful for quickly detecting
+  // stuck thread, when running locally with longer timeout.
+  bool enable_per_thread_lock_count_assertion_;
+
+  // Internal test variables
+
+  bool enable_lock_status_validation_;
+  std::vector<std::thread> threads_;
+  std::vector<std::unique_ptr<std::atomic_int>> counters_;
+  std::vector<int> values_;
+
+  // track whether the lock is in exclusive status or
+  // not. vector<bool> does something special underneath, causing consistency
+  // issue. Therefore int64_t is used.
+  std::vector<int64_t> exclusive_lock_status_;
+
+  // A counter to track number of shared locks for tracking shared lock status
+  std::vector<std::unique_ptr<std::atomic_int>> shared_lock_count_;
+
+  // shutdown flag to signal threads to exit
+  std::atomic_bool shutdown_ = false;
+
+  // test statistics
+  std::atomic_int64_t num_of_locks_acquired_ = 0;
+  std::atomic_int64_t num_of_shared_locks_acquired_ = 0;
+  std::atomic_int64_t num_of_exclusive_locks_acquired_ = 0;
+  std::atomic_int64_t num_of_deadlock_detected_ = 0;
+  std::vector<std::unique_ptr<std::atomic_int64_t>>
+      num_of_locks_acquired_per_thread_;
+};
+
+}  // namespace ROCKSDB_NAMESPACE

utilities/transactions/lock/range/range_locking_test.cc

@@ -5,22 +5,18 @@
 
 #ifndef OS_WIN
 
-#include <algorithm>
 #include <functional>
 #include <iomanip>
 #include <string>
 #include <thread>
 
 #include "db/db_impl/db_impl.h"
-#include "port/port.h"
 #include "rocksdb/db.h"
 #include "rocksdb/options.h"
-#include "rocksdb/perf_context.h"
 #include "rocksdb/utilities/transaction.h"
 #include "rocksdb/utilities/transaction_db.h"
-#include "utilities/transactions/lock/point/point_lock_manager_test.h"
-#include "utilities/transactions/pessimistic_transaction_db.h"
-#include "utilities/transactions/transaction_test.h"
+#include "utilities/transactions/lock/point/any_lock_manager_test.h"
+#include "utilities/transactions/transaction_db_mutex_impl.h"
 
 using std::string;
 

utilities/transactions/pessimistic_transaction.cc

@@ -84,6 +84,10 @@ void PessimisticTransaction::Initialize(const TransactionOptions& txn_options) {
         txn_db_impl_->GetTxnDBOptions().transaction_lock_timeout * 1000;
   }
 
+  // deadlock timeout should be lower than lock timeout
+  deadlock_timeout_us_ =
+      std::min(txn_options.deadlock_timeout_us, lock_timeout_);
+
   if (txn_options.expiration >= 0) {
     expiration_time_ = start_time_ + txn_options.expiration * 1000;
   } else {

utilities/transactions/pessimistic_transaction.h

@@ -81,7 +81,7 @@ class PessimisticTransaction : public TransactionBaseImpl {
     return ids;
   }
 
-  void SetWaitingTxn(autovector<TransactionID> ids, uint32_t column_family_id,
+  void SetWaitingTxn(autovector<TransactionID>& ids, uint32_t column_family_id,
                      const std::string* key, bool is_timed_out = false) {
     std::lock_guard<std::mutex> lock(wait_mutex_);
     waiting_txn_ids_ = ids;
@@ -114,6 +114,10 @@ class PessimisticTransaction : public TransactionBaseImpl {
   void SetLockTimeout(int64_t timeout) override {
     lock_timeout_ = timeout * 1000;
   }
+  int64_t GetDeadlockTimeout() const { return deadlock_timeout_us_; }
+  void SetDeadlockTimeout(int64_t timeout_ms) override {
+    deadlock_timeout_us_ = timeout_ms * 1000;
+  }
 
   // Returns true if locks were stolen successfully, false otherwise.
   bool TryStealingLocks();
@@ -213,6 +217,10 @@ class PessimisticTransaction : public TransactionBaseImpl {
   // Timeout in microseconds when locking a key or -1 if there is no timeout.
   int64_t lock_timeout_;
 
+  // Timeout in microseconds before perform dead lock detection.
+  // If 0, deadlock detection will be performed immediately.
+  int64_t deadlock_timeout_us_;
+
   // Whether to perform deadlock detection or not.
   bool deadlock_detect_;
 

utilities/transactions/timestamped_snapshot_test.cc

@@ -9,17 +9,26 @@
 #include "utilities/transactions/transaction_test.h"
 
 namespace ROCKSDB_NAMESPACE {
+
+constexpr std::array TimestampedSnapshotWithTsSanityCheck_Params = {
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite),
+    std::make_tuple(false, false, WRITE_UNPREPARED, kOrderedWrite)};
+
 INSTANTIATE_TEST_CASE_P(
     Unsupported, TimestampedSnapshotWithTsSanityCheck,
-    ::testing::Values(
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite),
-        std::make_tuple(false, false, WRITE_UNPREPARED, kOrderedWrite)));
+    ::testing::ValuesIn(WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(
+        WRAP_PARAM(bool, bool, TxnDBWritePolicy, WriteOrdering),
+        TimestampedSnapshotWithTsSanityCheck_Params)));
 
-INSTANTIATE_TEST_CASE_P(WriteCommitted, TransactionTest,
-                        ::testing::Combine(::testing::Bool(), ::testing::Bool(),
-                                           ::testing::Values(WRITE_COMMITTED),
-                                           ::testing::Values(kOrderedWrite)));
+INSTANTIATE_TEST_CASE_P(
+    WriteCommitted, TransactionTest,
+    ::testing::Combine(/*use_stackable_db=*/::testing::Bool(),
+                       /*two_write_queue=*/::testing::Bool(),
+                       ::testing::Values(WRITE_COMMITTED),
+                       ::testing::Values(kOrderedWrite),
+                       /*use_per_key_point_lock_mgr=*/::testing::Bool(),
+                       /*deadlock_timeout_us=*/::testing::Values(0, 1000)));
 
 namespace {
 // Not thread-safe. Caller needs to provide external synchronization.

utilities/transactions/transaction_base.h

@@ -250,6 +250,8 @@ class TransactionBaseImpl : public Transaction {
 
   void SetLockTimeout(int64_t /*timeout*/) override { /* Do nothing */ }
 
+  void SetDeadlockTimeout(int64_t /*timeout*/) override { /* Do nothing */ }
+
   const Snapshot* GetSnapshot() const override {
     // will return nullptr when there is no snapshot
     return snapshot_.get();

utilities/transactions/transaction_db_mutex_impl.cc

@@ -7,8 +7,9 @@
 
 #include <chrono>
 #include <condition_variable>
-#include <functional>
 #include <mutex>
+#include <sstream>
+#include <thread>
 
 #include "rocksdb/utilities/transaction_db_mutex.h"
 

utilities/transactions/transaction_test.cc

@@ -35,51 +35,71 @@
 
 namespace ROCKSDB_NAMESPACE {
 
+constexpr std::array DBAsBaseDB_TransactionTest_Params = {
+    std::make_tuple(false, false, WRITE_COMMITTED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_COMMITTED, kOrderedWrite),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite),
+    std::make_tuple(false, false, WRITE_UNPREPARED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_UNPREPARED, kOrderedWrite)};
+
 INSTANTIATE_TEST_CASE_P(
     DBAsBaseDB, TransactionTest,
-    ::testing::Values(
-        std::make_tuple(false, false, WRITE_COMMITTED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_COMMITTED, kOrderedWrite),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite),
-        std::make_tuple(false, false, WRITE_UNPREPARED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_UNPREPARED, kOrderedWrite)));
+    ::testing::ValuesIn(WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(
+        WRAP_PARAM(bool, bool, TxnDBWritePolicy, WriteOrdering),
+        DBAsBaseDB_TransactionTest_Params)));
+
+constexpr std::array DBAsBaseDB_TransactionStressTest_Params = {
+    std::make_tuple(false, false, WRITE_COMMITTED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_COMMITTED, kOrderedWrite),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite),
+    std::make_tuple(false, false, WRITE_UNPREPARED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_UNPREPARED, kOrderedWrite)};
+
 INSTANTIATE_TEST_CASE_P(
     DBAsBaseDB, TransactionStressTest,
-    ::testing::Values(
-        std::make_tuple(false, false, WRITE_COMMITTED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_COMMITTED, kOrderedWrite),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite),
-        std::make_tuple(false, false, WRITE_UNPREPARED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_UNPREPARED, kOrderedWrite)));
+    ::testing::ValuesIn(WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(
+        WRAP_PARAM(bool, bool, TxnDBWritePolicy, WriteOrdering),
+        DBAsBaseDB_TransactionStressTest_Params)));
+
+constexpr std::array StackableDBAsBaseDB_TransactionTest_Params = {
+    std::make_tuple(true, true, WRITE_COMMITTED, kOrderedWrite),
+    std::make_tuple(true, true, WRITE_PREPARED, kOrderedWrite),
+    std::make_tuple(true, true, WRITE_UNPREPARED, kOrderedWrite)};
+
 INSTANTIATE_TEST_CASE_P(
     StackableDBAsBaseDB, TransactionTest,
-    ::testing::Values(
-        std::make_tuple(true, true, WRITE_COMMITTED, kOrderedWrite),
-        std::make_tuple(true, true, WRITE_PREPARED, kOrderedWrite),
-        std::make_tuple(true, true, WRITE_UNPREPARED, kOrderedWrite)));
+    ::testing::ValuesIn(WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(
+        WRAP_PARAM(bool, bool, TxnDBWritePolicy, WriteOrdering),
+        StackableDBAsBaseDB_TransactionTest_Params)));
 
 // MySQLStyleTransactionTest takes far too long for valgrind to run. Only do it
 // in full mode (`ROCKSDB_FULL_VALGRIND_RUN` compiler flag is set).
 #if !defined(ROCKSDB_VALGRIND_RUN) || defined(ROCKSDB_FULL_VALGRIND_RUN)
+
+constexpr std::array MySQLStyleTransactionTest_Params = {
+    std::make_tuple(false, false, WRITE_COMMITTED, kOrderedWrite, false),
+    std::make_tuple(false, true, WRITE_COMMITTED, kOrderedWrite, false),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, false),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, true),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, false),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, true),
+    std::make_tuple(false, false, WRITE_UNPREPARED, kOrderedWrite, false),
+    std::make_tuple(false, false, WRITE_UNPREPARED, kOrderedWrite, true),
+    std::make_tuple(false, true, WRITE_UNPREPARED, kOrderedWrite, false),
+    std::make_tuple(false, true, WRITE_UNPREPARED, kOrderedWrite, true),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, false),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, true)};
+
 INSTANTIATE_TEST_CASE_P(
     MySQLStyleTransactionTest, MySQLStyleTransactionTest,
-    ::testing::Values(
-        std::make_tuple(false, false, WRITE_COMMITTED, kOrderedWrite, false),
-        std::make_tuple(false, true, WRITE_COMMITTED, kOrderedWrite, false),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, false),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, true),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, false),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, true),
-        std::make_tuple(false, false, WRITE_UNPREPARED, kOrderedWrite, false),
-        std::make_tuple(false, false, WRITE_UNPREPARED, kOrderedWrite, true),
-        std::make_tuple(false, true, WRITE_UNPREPARED, kOrderedWrite, false),
-        std::make_tuple(false, true, WRITE_UNPREPARED, kOrderedWrite, true),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, false),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, true)));
+    ::testing::ValuesIn(WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(
+        WRAP_PARAM(bool, bool, TxnDBWritePolicy, WriteOrdering, bool),
+        MySQLStyleTransactionTest_Params)));
+
 #endif  // !defined(ROCKSDB_VALGRIND_RUN) || defined(ROCKSDB_FULL_VALGRIND_RUN)
 
 TEST_P(TransactionTest, TestUpperBoundUponDeletion) {
@@ -5777,8 +5797,8 @@ Status TransactionStressTestInserter(
   TransactionOptions txn_options;
   txn_options.use_only_the_last_commit_time_batch_for_recovery = true;
 
-  // Inside the inserter we might also retake the snapshot. We do both since two
-  // separte functions are engaged for each.
+  // Inside the inserter we might also retake the snapshot. We do both since
+  // two separte functions are engaged for each.
   txn_options.set_snapshot = rand->OneIn(2);
 
   RandomTransactionInserter inserter(
@@ -8862,7 +8882,7 @@ TEST_P(TransactionTest, SecondaryIndexOnKey) {
   }
 }
 
-TEST_F(TransactionDBTest, CollapseKey) {
+TEST_P(TransactionDBTest, CollapseKey) {
   ASSERT_OK(ReOpen());
   ASSERT_OK(db->Put({}, "hello", "world"));
   ASSERT_OK(db->Flush({}));
@@ -8911,7 +8931,7 @@ TEST_F(TransactionDBTest, CollapseKey) {
   }
 }
 
-TEST_F(TransactionDBTest, FlushedLogWithPendingPrepareIsSynced) {
+TEST_P(TransactionDBTest, FlushedLogWithPendingPrepareIsSynced) {
   // Repro for a bug where we missed a necessary sync of the old WAL during
   // memtable flush. It happened due to applying an optimization to skip syncing
   // the old WAL in too many scenarios (all memtable flushes on single CF
@@ -8956,8 +8976,9 @@ TEST_F(TransactionDBTest, FlushedLogWithPendingPrepareIsSynced) {
   }
 }
 
-class CommitBypassMemtableTest : public DBTestBase,
-                                 public ::testing::WithParamInterface<bool> {
+class CommitBypassMemtableTest
+    : public DBTestBase,
+      public ::testing::WithParamInterface<std::tuple<bool, bool>> {
  public:
   CommitBypassMemtableTest() : DBTestBase("commit_bypass_memtable_test", true) {
     SetUpTransactionDB();
@@ -8968,12 +8989,11 @@ class CommitBypassMemtableTest : public DBTestBase,
   Options options;
   TransactionDBOptions txn_db_opts;
 
-  void SetUpTransactionDB(
-      bool atomic_flush = false) {
+  void SetUpTransactionDB(bool atomic_flush = false) {
     options = CurrentOptions();
     options.create_if_missing = true;
     options.allow_2pc = true;
-    options.two_write_queues = GetParam();
+    options.two_write_queues = std::get<0>(GetParam());
     // Avoid write stall
     options.max_write_buffer_number = 8;
     options.atomic_flush = atomic_flush;
@@ -8982,13 +9002,16 @@ class CommitBypassMemtableTest : public DBTestBase,
     Destroy(options, true);
 
     txn_db_opts.write_policy = TxnDBWritePolicy::WRITE_COMMITTED;
+    txn_db_opts.use_per_key_point_lock_mgr = std::get<1>(GetParam());
     ASSERT_OK(TransactionDB::Open(options, txn_db_opts, dbname_, &txn_db));
     ASSERT_NE(txn_db, nullptr);
     db_ = txn_db;
   }
 };
 
-INSTANTIATE_TEST_CASE_P(, CommitBypassMemtableTest, testing::Bool());
+INSTANTIATE_TEST_CASE_P(, CommitBypassMemtableTest,
+                        ::testing::Combine(::testing::Bool(),
+                                           ::testing::Bool()));
 
 // TODO: parameterize other tests in the file with commit_bypass_memtable
 TEST_P(CommitBypassMemtableTest, SingleCFUpdate) {
@@ -9776,7 +9799,7 @@ TEST_P(CommitBypassMemtableTest, MergeMiniStress) {
   }
 }
 
-TEST_F(TransactionDBTest, SelfDeadlockBug) {
+TEST_P(TransactionDBTest, SelfDeadlockBug) {
   ASSERT_OK(ReOpen());
 
   // Create two transactions
@@ -9820,6 +9843,11 @@ TEST_F(TransactionDBTest, SelfDeadlockBug) {
   delete txn2;
 }
 
+INSTANTIATE_TEST_CASE_P(
+    TransactionDBBasicTest, TransactionDBTest,
+    ::testing::Combine(/*user_per_key_point_lock_manager=*/::testing::Bool(),
+                       /*deadlock_timeout_us=*/::testing::Values(0, 1000)));
+
 TEST_P(CommitBypassMemtableTest,
        OptimizeLargeTxnCommitWriteBatchSizeThreshold) {
   // Tests TransactionOptions::large_txn_commit_optimize_byte_threshold

utilities/transactions/transaction_test.h

@@ -49,14 +49,18 @@ class TransactionTestBase : public ::testing::Test {
 
   TransactionDBOptions txn_db_options;
   bool use_stackable_db_;
+  int64_t deadlock_timeout_us_;
 
   TransactionTestBase(bool use_stackable_db, bool two_write_queue,
                       TxnDBWritePolicy write_policy,
-                      WriteOrdering write_ordering)
+                      WriteOrdering write_ordering,
+                      bool use_per_key_point_lock_mgr,
+                      int64_t deadlock_timeout_us)
       : db(nullptr),
         special_env(Env::Default()),
         env(nullptr),
-        use_stackable_db_(use_stackable_db) {
+        use_stackable_db_(use_stackable_db),
+        deadlock_timeout_us_(deadlock_timeout_us) {
     options.create_if_missing = true;
     options.max_write_buffer_number = 2;
     options.write_buffer_size = 4 * 1024;
@@ -77,6 +81,7 @@ class TransactionTestBase : public ::testing::Test {
     txn_db_options.default_lock_timeout = 0;
     txn_db_options.write_policy = write_policy;
     txn_db_options.rollback_merge_operands = true;
+    txn_db_options.use_per_key_point_lock_mgr = use_per_key_point_lock_mgr;
     // This will stress write unprepared, by forcing write batch flush on every
     // write.
     txn_db_options.default_write_batch_flush_threshold = 1;
@@ -481,30 +486,35 @@ class TransactionTestBase : public ::testing::Test {
 
 class TransactionTest
     : public TransactionTestBase,
-      virtual public ::testing::WithParamInterface<
-          std::tuple<bool, bool, TxnDBWritePolicy, WriteOrdering>> {
+      virtual public ::testing::WithParamInterface<std::tuple<
+          bool, bool, TxnDBWritePolicy, WriteOrdering, bool, int64_t>> {
  public:
   TransactionTest()
       : TransactionTestBase(std::get<0>(GetParam()), std::get<1>(GetParam()),
-                            std::get<2>(GetParam()), std::get<3>(GetParam())){};
+                            std::get<2>(GetParam()), std::get<3>(GetParam()),
+                            std::get<4>(GetParam()), std::get<5>(GetParam())) {}
 };
 
-class TransactionDBTest : public TransactionTestBase {
+class TransactionDBTest
+    : public TransactionTestBase,
+      virtual public ::testing::WithParamInterface<std::tuple<bool, int64_t>> {
  public:
   TransactionDBTest()
-      : TransactionTestBase(false, false, WRITE_COMMITTED, kOrderedWrite) {}
+      : TransactionTestBase(false, false, WRITE_COMMITTED, kOrderedWrite,
+                            std::get<0>(GetParam()), std::get<1>(GetParam())) {}
 };
 
 class TransactionStressTest : public TransactionTest {};
 
 class MySQLStyleTransactionTest
     : public TransactionTestBase,
-      virtual public ::testing::WithParamInterface<
-          std::tuple<bool, bool, TxnDBWritePolicy, WriteOrdering, bool>> {
+      virtual public ::testing::WithParamInterface<std::tuple<
+          bool, bool, TxnDBWritePolicy, WriteOrdering, bool, bool, int64_t>> {
  public:
   MySQLStyleTransactionTest()
       : TransactionTestBase(std::get<0>(GetParam()), std::get<1>(GetParam()),
-                            std::get<2>(GetParam()), std::get<3>(GetParam())),
+                            std::get<2>(GetParam()), std::get<3>(GetParam()),
+                            std::get<5>(GetParam()), std::get<6>(GetParam())),
         with_slow_threads_(std::get<4>(GetParam())) {
     if (with_slow_threads_ &&
         (txn_db_options.write_policy == WRITE_PREPARED ||
@@ -527,11 +537,13 @@ class MySQLStyleTransactionTest
 
 class WriteCommittedTxnWithTsTest
     : public TransactionTestBase,
-      public ::testing::WithParamInterface<std::tuple<bool, bool, bool>> {
+      public ::testing::WithParamInterface<
+          std::tuple<bool, bool, bool, bool, int64_t>> {
  public:
   WriteCommittedTxnWithTsTest()
       : TransactionTestBase(std::get<0>(GetParam()), std::get<1>(GetParam()),
-                            WRITE_COMMITTED, kOrderedWrite) {}
+                            WRITE_COMMITTED, kOrderedWrite,
+                            std::get<3>(GetParam()), std::get<4>(GetParam())) {}
   ~WriteCommittedTxnWithTsTest() override {
     for (auto* h : handles_) {
       delete h;
@@ -567,12 +579,13 @@ class WriteCommittedTxnWithTsTest
 
 class TimestampedSnapshotWithTsSanityCheck
     : public TransactionTestBase,
-      public ::testing::WithParamInterface<
-          std::tuple<bool, bool, TxnDBWritePolicy, WriteOrdering>> {
+      public ::testing::WithParamInterface<std::tuple<
+          bool, bool, TxnDBWritePolicy, WriteOrdering, bool, int64_t>> {
  public:
   explicit TimestampedSnapshotWithTsSanityCheck()
       : TransactionTestBase(std::get<0>(GetParam()), std::get<1>(GetParam()),
-                            std::get<2>(GetParam()), std::get<3>(GetParam())) {}
+                            std::get<2>(GetParam()), std::get<3>(GetParam()),
+                            std::get<4>(GetParam()), std::get<5>(GetParam())) {}
   ~TimestampedSnapshotWithTsSanityCheck() override {
     for (auto* h : handles_) {
       delete h;
@@ -583,4 +596,58 @@ class TimestampedSnapshotWithTsSanityCheck
   std::vector<ColumnFamilyHandle*> handles_{};
 };
 
+// Wrap existing params with per-key point lock manager parameters
+template <typename TargetParamType, typename SourceParamType, std::size_t... Is>
+std::vector<TargetParamType> WrapParamWithPerKeyPointLockManagerParamsImpl(
+    SourceParamType&& source_param, std::index_sequence<Is...>) {
+  std::vector<TargetParamType> wrapped_params;
+  // Use original PointLockManager
+  wrapped_params.push_back(TargetParamType(
+      std::get<Is>(std::forward<SourceParamType>(source_param))..., false,
+      INT64_C(0)));
+  // Use PerKeyPointLockManager with deadlock timeout 0
+  wrapped_params.push_back(TargetParamType(
+      std::get<Is>(std::forward<SourceParamType>(source_param))..., true,
+      INT64_C(0)));
+  // Use PerKeyPointLockManager with deadlock timeout 1000
+  wrapped_params.push_back(TargetParamType(
+      std::get<Is>(std::forward<SourceParamType>(source_param))..., true,
+      INT64_C(1000)));
+
+  return wrapped_params;
+}
+
+template <typename TargetParamType, typename SourceParamType>
+std::vector<TargetParamType> WrapParamWithPerKeyPointLockManagerParams(
+    SourceParamType&& source_param) {
+  // Get the size of the source param
+  constexpr std::size_t N = std::tuple_size_v<std::decay_t<SourceParamType>>;
+  // Create an index sequence from 0 to N-1
+  return WrapParamWithPerKeyPointLockManagerParamsImpl<TargetParamType>(
+      std::forward<SourceParamType>(source_param),
+      std::make_index_sequence<N>{});
+}
+
+template <typename TargetParamType, typename SourceParamType, size_t M>
+std::vector<TargetParamType> WrapParamsWithPerKeyPointLockManagerParams(
+    std::array<SourceParamType, M> source_param) {
+  std::vector<TargetParamType> wrapped_params;
+  for (auto& param : source_param) {
+    // Create an index sequence from 0 to N-1
+    auto new_params =
+        WrapParamWithPerKeyPointLockManagerParams<TargetParamType>(
+            std::forward<SourceParamType>(param));
+    wrapped_params.insert(wrapped_params.end(), new_params.begin(),
+                          new_params.end());
+  }
+  return wrapped_params;
+}
+
+#define WRAP_PARAM(...) __VA_ARGS__
+
+#define WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(SOURCE_PARAM_TYPES, \
+                                                          PARAMS)             \
+  WrapParamsWithPerKeyPointLockManagerParams<                                 \
+      std::tuple<SOURCE_PARAM_TYPES, bool, int64_t>>(PARAMS)
+
 }  // namespace ROCKSDB_NAMESPACE

utilities/transactions/write_committed_transaction_ts_test.cc

@@ -14,26 +14,12 @@
 namespace ROCKSDB_NAMESPACE {
 
 INSTANTIATE_TEST_CASE_P(
-    DBAsBaseDB, WriteCommittedTxnWithTsTest,
-    ::testing::Values(std::make_tuple(false, /*two_write_queue=*/false,
-                                      /*enable_indexing=*/false),
-                      std::make_tuple(false, /*two_write_queue=*/true,
-                                      /*enable_indexing=*/false),
-                      std::make_tuple(false, /*two_write_queue=*/false,
-                                      /*enable_indexing=*/true),
-                      std::make_tuple(false, /*two_write_queue=*/true,
-                                      /*enable_indexing=*/true)));
-
-INSTANTIATE_TEST_CASE_P(
-    DBAsStackableDB, WriteCommittedTxnWithTsTest,
-    ::testing::Values(std::make_tuple(true, /*two_write_queue=*/false,
-                                      /*enable_indexing=*/false),
-                      std::make_tuple(true, /*two_write_queue=*/true,
-                                      /*enable_indexing=*/false),
-                      std::make_tuple(true, /*two_write_queue=*/false,
-                                      /*enable_indexing=*/true),
-                      std::make_tuple(true, /*two_write_queue=*/true,
-                                      /*enable_indexing=*/true)));
+    DBAsBaseDBAndStackableDB, WriteCommittedTxnWithTsTest,
+    ::testing::Combine(/*use_stackable_db=*/::testing::Bool(),
+                       /*two_write_queue=*/::testing::Bool(),
+                       /*enable_indexing=*/::testing::Bool(),
+                       /*use_per_key_point_lock_mgr=*/::testing::Bool(),
+                       /*deadlock_timeout_us=*/::testing::Values(0, 1000)));
 
 TEST_P(WriteCommittedTxnWithTsTest, SanityChecks) {
   ASSERT_OK(ReOpenNoDelete());

utilities/transactions/write_prepared_transaction_test.cc

@@ -354,9 +354,12 @@ class WritePreparedTransactionTestBase : public TransactionTestBase {
  public:
   WritePreparedTransactionTestBase(bool use_stackable_db, bool two_write_queue,
                                    TxnDBWritePolicy write_policy,
-                                   WriteOrdering write_ordering)
+                                   WriteOrdering write_ordering,
+                                   bool user_per_key_point_lock_mgr,
+                                   int64_t deadlock_timeout_us)
       : TransactionTestBase(use_stackable_db, two_write_queue, write_policy,
-                            write_ordering){};
+                            write_ordering, user_per_key_point_lock_mgr,
+                            deadlock_timeout_us) {}
 
  protected:
   void UpdateTransactionDBOptions(size_t snapshot_cache_bits,
@@ -528,27 +531,30 @@ class WritePreparedTransactionTestBase : public TransactionTestBase {
 
 class WritePreparedTransactionTest
     : public WritePreparedTransactionTestBase,
-      virtual public ::testing::WithParamInterface<
-          std::tuple<bool, bool, TxnDBWritePolicy, WriteOrdering>> {
+      virtual public ::testing::WithParamInterface<std::tuple<
+          bool, bool, TxnDBWritePolicy, WriteOrdering, bool, int64_t>> {
  public:
   WritePreparedTransactionTest()
       : WritePreparedTransactionTestBase(
             std::get<0>(GetParam()), std::get<1>(GetParam()),
-            std::get<2>(GetParam()), std::get<3>(GetParam())){};
+            std::get<2>(GetParam()), std::get<3>(GetParam()),
+            std::get<4>(GetParam()), std::get<5>(GetParam())) {}
 };
 
 #if !defined(ROCKSDB_VALGRIND_RUN) || defined(ROCKSDB_FULL_VALGRIND_RUN)
 class SnapshotConcurrentAccessTest
     : public WritePreparedTransactionTestBase,
-      virtual public ::testing::WithParamInterface<std::tuple<
-          bool, bool, TxnDBWritePolicy, WriteOrdering, size_t, size_t>> {
+      virtual public ::testing::WithParamInterface<
+          std::tuple<bool, bool, TxnDBWritePolicy, WriteOrdering, size_t,
+                     size_t, bool, int64_t>> {
  public:
   SnapshotConcurrentAccessTest()
       : WritePreparedTransactionTestBase(
             std::get<0>(GetParam()), std::get<1>(GetParam()),
-            std::get<2>(GetParam()), std::get<3>(GetParam())),
+            std::get<2>(GetParam()), std::get<3>(GetParam()),
+            std::get<6>(GetParam()), std::get<7>(GetParam())),
         split_id_(std::get<4>(GetParam())),
-        split_cnt_(std::get<5>(GetParam())){};
+        split_cnt_(std::get<5>(GetParam())) {}
 
  protected:
   // A test is split into split_cnt_ tests, each identified with split_id_ where
@@ -560,13 +566,15 @@ class SnapshotConcurrentAccessTest
 
 class SeqAdvanceConcurrentTest
     : public WritePreparedTransactionTestBase,
-      virtual public ::testing::WithParamInterface<std::tuple<
-          bool, bool, TxnDBWritePolicy, WriteOrdering, size_t, size_t>> {
+      virtual public ::testing::WithParamInterface<
+          std::tuple<bool, bool, TxnDBWritePolicy, WriteOrdering, size_t,
+                     size_t, bool, int64_t>> {
  public:
   SeqAdvanceConcurrentTest()
       : WritePreparedTransactionTestBase(
             std::get<0>(GetParam()), std::get<1>(GetParam()),
-            std::get<2>(GetParam()), std::get<3>(GetParam())),
+            std::get<2>(GetParam()), std::get<3>(GetParam()),
+            std::get<6>(GetParam()), std::get<7>(GetParam())),
         split_id_(std::get<4>(GetParam())),
         split_cnt_(std::get<5>(GetParam())) {
     special_env.skip_fsync_ = true;
@@ -579,120 +587,143 @@ class SeqAdvanceConcurrentTest
   size_t split_cnt_;
 };
 
+constexpr std::array WritePreparedTransactionTest_Params = {
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite)};
+
 INSTANTIATE_TEST_CASE_P(
     WritePreparedTransaction, WritePreparedTransactionTest,
-    ::testing::Values(
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite)));
+    ::testing::ValuesIn(WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(
+        WRAP_PARAM(bool, bool, TxnDBWritePolicy, WriteOrdering),
+        WritePreparedTransactionTest_Params)));
 
 #if !defined(ROCKSDB_VALGRIND_RUN) || defined(ROCKSDB_FULL_VALGRIND_RUN)
-INSTANTIATE_TEST_CASE_P(
-    TwoWriteQueues, SnapshotConcurrentAccessTest,
-    ::testing::Values(
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 0, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 1, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 2, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 3, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 4, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 5, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 6, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 7, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 8, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 9, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 10, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 11, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 12, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 13, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 14, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 15, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 16, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 17, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 18, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 19, 20),
 
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 0, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 1, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 2, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 3, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 4, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 5, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 6, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 7, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 8, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 9, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 10, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 11, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 12, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 13, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 14, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 15, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 16, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 17, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 18, 20),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 19, 20)));
+constexpr std::array TwoWriteQueue_SnapshotConcurrentAccessTest_Params = {
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 0, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 1, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 2, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 3, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 4, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 5, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 6, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 7, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 8, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 9, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 10, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 11, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 12, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 13, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 14, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 15, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 16, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 17, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 18, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 19, 20),
+
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 0, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 1, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 2, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 3, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 4, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 5, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 6, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 7, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 8, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 9, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 10, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 11, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 12, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 13, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 14, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 15, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 16, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 17, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 18, 20),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 19, 20)};
+
+INSTANTIATE_TEST_CASE_P(
+    TwoWriteQueuesPointLockManager, SnapshotConcurrentAccessTest,
+    ::testing::ValuesIn(WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(
+        WRAP_PARAM(bool, bool, TxnDBWritePolicy, WriteOrdering, size_t, size_t),
+        TwoWriteQueue_SnapshotConcurrentAccessTest_Params)));
+
+constexpr std::array OneWriteQueue_SnapshotConcurrentAccessTest_Params = {
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 0, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 1, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 2, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 3, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 4, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 5, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 6, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 7, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 8, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 9, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 10, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 11, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 12, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 13, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 14, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 15, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 16, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 17, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 18, 20),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 19, 20),
+};
 
 INSTANTIATE_TEST_CASE_P(
     OneWriteQueue, SnapshotConcurrentAccessTest,
-    ::testing::Values(
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 0, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 1, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 2, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 3, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 4, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 5, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 6, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 7, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 8, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 9, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 10, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 11, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 12, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 13, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 14, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 15, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 16, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 17, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 18, 20),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 19, 20)));
+    ::testing::ValuesIn(WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(
+        WRAP_PARAM(bool, bool, TxnDBWritePolicy, WriteOrdering, size_t, size_t),
+        OneWriteQueue_SnapshotConcurrentAccessTest_Params)));
+
+constexpr std::array TwoWriteQueues_SeqAdvanceConcurrentTest_Params = {
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 0, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 1, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 2, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 3, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 4, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 5, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 6, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 7, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 8, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 9, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 0, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 1, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 2, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 3, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 4, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 5, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 6, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 7, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 8, 10),
+    std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 9, 10)};
 
 INSTANTIATE_TEST_CASE_P(
     TwoWriteQueues, SeqAdvanceConcurrentTest,
-    ::testing::Values(
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 0, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 1, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 2, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 3, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 4, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 5, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 6, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 7, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 8, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kOrderedWrite, 9, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 0, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 1, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 2, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 3, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 4, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 5, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 6, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 7, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 8, 10),
-        std::make_tuple(false, true, WRITE_PREPARED, kUnorderedWrite, 9, 10)));
+    ::testing::ValuesIn(WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(
+        WRAP_PARAM(bool, bool, TxnDBWritePolicy, WriteOrdering, size_t, size_t),
+        TwoWriteQueues_SeqAdvanceConcurrentTest_Params)));
+
+constexpr std::array OneWriteQueue_SeqAdvanceConcurrentTest_Params = {
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 0, 10),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 1, 10),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 2, 10),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 3, 10),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 4, 10),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 5, 10),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 6, 10),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 7, 10),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 8, 10),
+    std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 9, 10)};
 
 INSTANTIATE_TEST_CASE_P(
     OneWriteQueue, SeqAdvanceConcurrentTest,
-    ::testing::Values(
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 0, 10),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 1, 10),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 2, 10),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 3, 10),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 4, 10),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 5, 10),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 6, 10),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 7, 10),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 8, 10),
-        std::make_tuple(false, false, WRITE_PREPARED, kOrderedWrite, 9, 10)));
+    ::testing::ValuesIn(WRAP_PARAM_WITH_PER_KEY_POINT_LOCK_MANAGER_PARAMS(
+        WRAP_PARAM(bool, bool, TxnDBWritePolicy, WriteOrdering, size_t, size_t),
+        OneWriteQueue_SeqAdvanceConcurrentTest_Params)));
+
 #endif  // !defined(ROCKSDB_VALGRIND_RUN) || defined(ROCKSDB_FULL_VALGRIND_RUN)
 
 TEST_P(WritePreparedTransactionTest, CommitMap) {

utilities/transactions/write_unprepared_transaction_test.cc

@@ -13,37 +13,43 @@ class WriteUnpreparedTransactionTestBase : public TransactionTestBase {
  public:
   WriteUnpreparedTransactionTestBase(bool use_stackable_db,
                                      bool two_write_queue,
-                                     TxnDBWritePolicy write_policy)
+                                     TxnDBWritePolicy write_policy,
+                                     bool use_per_key_point_lock_mgr,
+                                     int64_t deadlock_timeout_us)
       : TransactionTestBase(use_stackable_db, two_write_queue, write_policy,
-                            kOrderedWrite) {}
+                            kOrderedWrite, use_per_key_point_lock_mgr,
+                            deadlock_timeout_us) {}
 };
 
 class WriteUnpreparedTransactionTest
     : public WriteUnpreparedTransactionTestBase,
       virtual public ::testing::WithParamInterface<
-          std::tuple<bool, bool, TxnDBWritePolicy>> {
+          std::tuple<bool, bool, TxnDBWritePolicy, bool, int64_t>> {
  public:
   WriteUnpreparedTransactionTest()
-      : WriteUnpreparedTransactionTestBase(std::get<0>(GetParam()),
-                                           std::get<1>(GetParam()),
-                                           std::get<2>(GetParam())) {}
+      : WriteUnpreparedTransactionTestBase(
+            std::get<0>(GetParam()), std::get<1>(GetParam()),
+            std::get<2>(GetParam()), std::get<3>(GetParam()),
+            std::get<4>(GetParam())) {}
 };
 
 INSTANTIATE_TEST_CASE_P(
     WriteUnpreparedTransactionTest, WriteUnpreparedTransactionTest,
-    ::testing::Values(std::make_tuple(false, false, WRITE_UNPREPARED),
-                      std::make_tuple(false, true, WRITE_UNPREPARED)));
+    ::testing::Combine(::testing::Values(false), ::testing::Bool(),
+                       ::testing::Values(WRITE_UNPREPARED), ::testing::Bool(),
+                       ::testing::Values(0, 1000)));
 
 enum SnapshotAction { NO_SNAPSHOT, RO_SNAPSHOT, REFRESH_SNAPSHOT };
 enum VerificationOperation { VERIFY_GET, VERIFY_NEXT, VERIFY_PREV };
 class WriteUnpreparedSnapshotTest
     : public WriteUnpreparedTransactionTestBase,
-      virtual public ::testing::WithParamInterface<
-          std::tuple<bool, SnapshotAction, VerificationOperation>> {
+      virtual public ::testing::WithParamInterface<std::tuple<
+          bool, SnapshotAction, VerificationOperation, bool, int64_t>> {
  public:
   WriteUnpreparedSnapshotTest()
-      : WriteUnpreparedTransactionTestBase(false, std::get<0>(GetParam()),
-                                           WRITE_UNPREPARED),
+      : WriteUnpreparedTransactionTestBase(
+            false, std::get<0>(GetParam()), WRITE_UNPREPARED,
+            std::get<3>(GetParam()), std::get<4>(GetParam())),
         action_(std::get<1>(GetParam())),
         verify_op_(std::get<2>(GetParam())) {}
   SnapshotAction action_;
@@ -56,10 +62,11 @@ class WriteUnpreparedSnapshotTest
 // verification operation
 INSTANTIATE_TEST_CASE_P(
     WriteUnpreparedSnapshotTest, WriteUnpreparedSnapshotTest,
-    ::testing::Combine(
-        ::testing::Bool(),
-        ::testing::Values(NO_SNAPSHOT, RO_SNAPSHOT, REFRESH_SNAPSHOT),
-        ::testing::Values(VERIFY_GET, VERIFY_NEXT, VERIFY_PREV)));
+    ::testing::Combine(::testing::Bool(),
+                       ::testing::Values(NO_SNAPSHOT, RO_SNAPSHOT,
+                                         REFRESH_SNAPSHOT),
+                       ::testing::Values(VERIFY_GET, VERIFY_NEXT, VERIFY_PREV),
+                       ::testing::Bool(), ::testing::Values(0, 1000)));
 
 TEST_P(WriteUnpreparedTransactionTest, ReadYourOwnWrite) {
   // The following tests checks whether reading your own write for