@nxtedition/rocksdb 13.5.13 → 15.0.1

package/deps/rocksdb/rocksdb/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

package/deps/rocksdb/rocksdb/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();
+}