Remove separate LockLimiter, avoid expiring persistent locks

3 files changed, 68 insertions, 118 deletions

diff --git a/ydb/core/client/locks_ut.cpp b/ydb/core/client/locks_ut.cpp
index 47b7c3dfdd5..a6e9a9bcb59 100644
--- a/ydb/core/client/locks_ut.cpp
+++ b/ydb/core/client/locks_ut.cpp
@@ -1810,7 +1810,7 @@ static void LocksLimit() {
 
     using TLock = TSysTables::TLocksTable::TLock;
 
-    ui32 limit = NDataShard::TLockLocker::TLockLimiter::LockLimit();
+    ui32 limit = NDataShard::TLockLocker::LockLimit();
     const ui32 factor = 100;
 
     const char * query = R"((
@@ -1917,7 +1917,7 @@ static void ShardLocks() {
     TClient::TFlatQueryOptions opts;
 
 
-    ui32 limit = NDataShard::TLockLocker::TLockLimiter::LockLimit();
+    ui32 limit = NDataShard::TLockLocker::LockLimit();
     //const ui32 factor = 100;
 
     const char * setLock = R"___((
diff --git a/ydb/core/tx/datashard/datashard_locks.cpp b/ydb/core/tx/datashard/datashard_locks.cpp
index 7385b514a7f..ef6b7ef4b4e 100644
--- a/ydb/core/tx/datashard/datashard_locks.cpp
+++ b/ydb/core/tx/datashard/datashard_locks.cpp
@@ -488,29 +488,50 @@ void TLockLocker::RemoveBrokenRanges() {
 TLockInfo::TPtr TLockLocker::GetOrAddLock(ui64 lockId, ui32 lockNodeId) {
     auto it = Locks.find(lockId);
     if (it != Locks.end()) {
-        Limiter.TouchLock(lockId);
+        if (it->second->IsInList<TLockInfoExpireListTag>()) {
+            ExpireQueue.PushBack(it->second.Get());
+        }
         if (lockNodeId && !it->second->LockNodeId) {
-            // This shouldn't ever happen, but better safe than sorry
+            // This should never happen, but better safe than sorry
             it->second->LockNodeId = lockNodeId;
             PendingSubscribeLocks.emplace_back(lockId, lockNodeId);
         }
         return it->second;
     }
 
-    TLockInfo::TPtr lock = Limiter.TryAddLock(lockId, lockNodeId);
-    if (lock) {
-        Locks[lockId] = lock;
-        if (lockNodeId) {
-            PendingSubscribeLocks.emplace_back(lockId, lockNodeId);
+    while (Locks.size() >= LockLimit()) {
+        if (!BrokenLocks.Empty()) {
+            // We remove broken locks first
+            TLockInfo* lock = BrokenLocks.Front();
+            RemoveOneLock(lock->GetLockId());
+            continue;
         }
+        if (!ExpireQueue.Empty()) {
+            TLockInfo* lock = ExpireQueue.Front();
+            if (TAppData::TimeProvider->Now() - lock->GetCreationTime() >= LockTimeLimit()) {
+                RemoveOneLock(lock->GetLockId());
+                continue;
+            }
+        }
+        // We cannot add any more locks
+        return nullptr;
+    }
+
+    TLockInfo::TPtr lock(new TLockInfo(this, lockId, lockNodeId));
+    Y_VERIFY(!lock->IsPersistent());
+    Locks[lockId] = lock;
+    if (lockNodeId) {
+        PendingSubscribeLocks.emplace_back(lockId, lockNodeId);
     }
+    ExpireQueue.PushBack(lock.Get());
     return lock;
 }
 
 TLockInfo::TPtr TLockLocker::AddLock(ui64 lockId, ui32 lockNodeId, ui32 generation, ui64 counter, TInstant createTs) {
     Y_VERIFY(Locks.find(lockId) == Locks.end());
 
-    TLockInfo::TPtr lock = Limiter.AddLock(lockId, lockNodeId, generation, counter, createTs);
+    TLockInfo::TPtr lock(new TLockInfo(this, lockId, lockNodeId, generation, counter, createTs));
+    Y_VERIFY(lock->IsPersistent());
     Locks[lockId] = lock;
     if (lockNodeId) {
         PendingSubscribeLocks.emplace_back(lockId, lockNodeId);
@@ -529,6 +550,7 @@ void TLockLocker::RemoveOneLock(ui64 lockTxId, ILocksDb* db) {
             Self->IncCounter(COUNTER_LOCKS_REMOVED);
         }
 
+        ExpireQueue.Remove(txLock.Get());
         if (txLock->InBrokenLocks) {
             BrokenLocks.Remove(txLock.Get());
             --BrokenLocksCount_;
@@ -541,27 +563,15 @@ void TLockLocker::RemoveOneLock(ui64 lockTxId, ILocksDb* db) {
             Tables.at(tableId)->RemoveWriteLock(txLock.Get());
         }
         txLock->CleanupConflicts();
-        Limiter.RemoveLock(lockTxId);
         Locks.erase(it);
+
         if (txLock->IsPersistent()) {
-            if (db) {
-                txLock->PersistRemoveLock(db);
-            } else {
-                Y_VERIFY(txLock->IsBroken(), "Scheduling persistent lock removal that is not broken");
-                RemovedPersistentLocks.push_back(txLock);
-            }
+            Y_VERIFY(db, "Cannot remove persistent locks without a database");
+            txLock->PersistRemoveLock(db);
         }
     }
 }
 
-void TLockLocker::RemoveBrokenLocks() {
-    RemoveBrokenRanges();
-    while (BrokenLocks) {
-        auto* lock = BrokenLocks.PopFront();
-        RemoveOneLock(lock->GetLockId());
-    }
-}
-
 void TLockLocker::ForceBreakLock(ui64 lockId) {
     if (auto lock = GetLock(lockId, TRowVersion::Min())) {
         lock->SetBroken(TRowVersion::Min());
@@ -589,14 +599,12 @@ void TLockLocker::RemoveSchema(const TPathId& tableId) {
     CleanupPending.clear();
     CleanupCandidates.clear();
     PendingSubscribeLocks.clear();
-    RemovedPersistentLocks.clear();
-    Limiter.Clear();
 }
 
 bool TLockLocker::ForceShardLock(const TPathId& tableId) const {
     auto it = Tables.find(tableId);
     if (it != Tables.end()) {
-        if (it->second->RangeCount() > TLockLimiter::LockLimit()) {
+        if (it->second->RangeCount() > LockLimit()) {
             return true;
         }
     }
@@ -605,7 +613,7 @@ bool TLockLocker::ForceShardLock(const TPathId& tableId) const {
 
 bool TLockLocker::ForceShardLock(const TIntrusiveList<TTableLocks, TTableLocksReadListTag>& readTables) const {
     for (auto& table : readTables) {
-        if (table.RangeCount() > TLockLimiter::LockLimit())
+        if (table.RangeCount() > LockLimit())
             return true;
     }
     return false;
@@ -647,50 +655,6 @@ void TLockLocker::SaveBrokenPersistentLocks(ILocksDb* db) {
     }
 }
 
-// TLockLocker.TLockLimiter
-
-TLockInfo::TPtr TLockLocker::TLockLimiter::TryAddLock(ui64 lockId, ui32 lockNodeId) {
-#if 1
-    if (LocksQueue.Size() >= LockLimit()) {
-        Parent->RemoveBrokenLocks();
-    }
-#endif
-    if (LocksQueue.Size() >= LockLimit()) {
-        TInstant forgetTime = TAppData::TimeProvider->Now() - TDuration::MilliSeconds(TimeLimitMSec());
-        auto oldest = LocksQueue.FindOldest();
-        if (oldest.Value() >= forgetTime)
-            return nullptr;
-
-        if (Parent->Self->TabletCounters) {
-            Parent->Self->IncCounter(COUNTER_LOCKS_EVICTED);
-        }
-
-        Parent->RemoveOneLock(oldest.Key()); // erase LocksQueue inside
-    }
-
-    LocksQueue.Insert(lockId, TAppData::TimeProvider->Now());
-    return TLockInfo::TPtr(new TLockInfo(Parent, lockId, lockNodeId));
-}
-
-void TLockLocker::TLockLimiter::RemoveLock(ui64 lockId) {
-    auto it = LocksQueue.FindWithoutPromote(lockId);
-    if (it != LocksQueue.End())
-        LocksQueue.Erase(it);
-}
-
-void TLockLocker::TLockLimiter::TouchLock(ui64 lockId) {
-    LocksQueue.Find(lockId);
-}
-
-TLockInfo::TPtr TLockLocker::TLockLimiter::AddLock(ui64 lockId, ui32 lockNodeId, ui32 generation, ui64 counter, TInstant createTs) {
-    LocksQueue.Insert(lockId, createTs);
-    return TLockInfo::TPtr(new TLockInfo(Parent, lockId, lockNodeId, generation, counter, createTs));
-}
-
-void TLockLocker::TLockLimiter::Clear() {
-    LocksQueue.Clear();
-}
-
 // TSysLocks
 
 TVector<TSysLocks::TLock> TSysLocks::ApplyLocks() {
@@ -797,6 +761,8 @@ TVector<TSysLocks::TLock> TSysLocks::ApplyLocks() {
             if (lock->GetWriteTables() && !lock->IsPersistent()) {
                 // We need to persist a new lock
                 lock->PersistLock(Db);
+                // Persistent locks cannot expire
+                Locker.ExpireQueue.Remove(lock.Get());
             }
         }
     }
diff --git a/ydb/core/tx/datashard/datashard_locks.h b/ydb/core/tx/datashard/datashard_locks.h
index 6cb67bfe8cc..0731925188f 100644
--- a/ydb/core/tx/datashard/datashard_locks.h
+++ b/ydb/core/tx/datashard/datashard_locks.h
@@ -235,6 +235,7 @@ struct TLockInfoReadConflictListTag {};
 struct TLockInfoWriteConflictListTag {};
 struct TLockInfoBrokenListTag {};
 struct TLockInfoBrokenPersistentListTag {};
+struct TLockInfoExpireListTag {};
 
 /// Aggregates shard, point and range locks
 class TLockInfo
@@ -245,6 +246,7 @@ class TLockInfo
     , public TIntrusiveListItem<TLockInfo, TLockInfoWriteConflictListTag>
     , public TIntrusiveListItem<TLockInfo, TLockInfoBrokenListTag>
     , public TIntrusiveListItem<TLockInfo, TLockInfoBrokenPersistentListTag>
+    , public TIntrusiveListItem<TLockInfo, TLockInfoExpireListTag>
 {
     friend class TTableLocks;
     friend class TLockLocker;
@@ -257,6 +259,18 @@ public:
     TLockInfo(TLockLocker * locker, ui64 lockId, ui32 lockNodeId, ui32 generation, ui64 counter, TInstant createTs);
     ~TLockInfo();
 
+    template<class TTag>
+    bool IsInList() const {
+        using TItem = TIntrusiveListItem<TLockInfo, TTag>;
+        return !static_cast<const TItem*>(this)->Empty();
+    }
+
+    template<class TTag>
+    void UnlinkFromList() {
+        using TItem = TIntrusiveListItem<TLockInfo, TTag>;
+        static_cast<TItem*>(this)->Unlink();
+    }
+
     ui32 GetGeneration() const { return Generation; }
     ui64 GetCounter(const TRowVersion& at = TRowVersion::Max()) const { return !BreakVersion || at < *BreakVersion ? Counter : Max<ui64>(); }
     bool IsBroken(const TRowVersion& at = TRowVersion::Max()) const { return GetCounter(at) == Max<ui64>(); }
@@ -416,45 +430,22 @@ class TLockLocker {
 
 public:
     /// Prevent unlimited lock's count growth
-    class TLockLimiter {
-    public:
-        static constexpr ui32 TimeLimitMSec() { return 5 * 60 * 1000; }
-        static constexpr ui64 LockLimit() {
-            // Valgrind and sanitizers are too slow
-            // Some tests cannot exhaust default limit in under 5 minutes
-            return NValgrind::PlainOrUnderValgrind(
-                NSan::PlainOrUnderSanitizer(
-                    16 * 1024,
-                    1024),
-                1024);
-        }
-
-        TLockLimiter(TLockLocker * parent)
-            : Parent(parent)
-            , LocksQueue(2 * LockLimit()) // it should be greater than LockLimit
-        {}
-
-        ui64 LocksCount() const { return Parent->LocksCount(); }
-
-        TLockInfo::TPtr TryAddLock(ui64 lockId, ui32 lockNodeId);
-        void RemoveLock(ui64 lockId);
-        void TouchLock(ui64 lockId);
-
-        TLockInfo::TPtr AddLock(ui64 lockId, ui32 lockNodeId, ui32 generation, ui64 counter, TInstant createTs);
-        void Clear();
-
-        // TODO: AddPoint, AddRange
+    static constexpr ui64 LockLimit() {
+        // Valgrind and sanitizers are too slow
+        // Some tests cannot exhaust default limit in under 5 minutes
+        return NValgrind::PlainOrUnderValgrind(
+            NSan::PlainOrUnderSanitizer(
+                16 * 1024,
+                1024),
+            1024);
+    }
 
-    private:
-        TLockLocker * Parent;
-        TLRUCache<ui64, TInstant> LocksQueue;
-    };
+    /// We don't expire locks until this time limit after they are created
+    static constexpr TDuration LockTimeLimit() { return TDuration::Minutes(5); }
 
     template <typename T>
     TLockLocker(const T * self)
         : Self(new TLocksDataShardAdapter<T>(self))
-        , Limiter(this)
-        , Counter(0)
     {}
 
     ~TLockLocker() {
@@ -544,10 +535,6 @@ public:
     ui32 Generation() const { return Self->Generation(); }
     ui64 IncCounter() { return Counter++; };
 
-    TVector<TLockInfo::TPtr>& GetRemovedPersistentLocks() {
-        return RemovedPersistentLocks;
-    }
-
     void Clear() {
         for (auto& pr : Tables) {
             pr.second->Clear();
@@ -558,15 +545,15 @@ public:
         CleanupPending.clear();
         CleanupCandidates.clear();
         PendingSubscribeLocks.clear();
-        RemovedPersistentLocks.clear();
-        Limiter.Clear();
     }
 
 private:
-    THolder<TLocksDataShard> Self;
+    const THolder<TLocksDataShard> Self;
     THashMap<ui64, TLockInfo::TPtr> Locks; // key is LockId
     THashMap<TPathId, TTableLocks::TPtr> Tables;
     THashSet<ui64> ShardLocks;
+    // A list of locks that may be removed when enough time passes
+    TIntrusiveList<TLockInfo, TLockInfoExpireListTag> ExpireQueue;
     // A list of broken, but not yet removed locks
     TIntrusiveList<TLockInfo, TLockInfoBrokenPersistentListTag> BrokenPersistentLocks;
     TIntrusiveList<TLockInfo, TLockInfoBrokenListTag> BrokenLocks;
@@ -575,9 +562,7 @@ private:
     TVector<ui64> CleanupPending;
     TPriorityQueue<TVersionedLockId> CleanupCandidates;
     TList<TPendingSubscribeLock> PendingSubscribeLocks;
-    TVector<TLockInfo::TPtr> RemovedPersistentLocks;
-    TLockLimiter Limiter;
-    ui64 Counter;
+    ui64 Counter = 0;
 
     TTableLocks::TPtr GetTableLocks(const TTableId& table) const {
         auto it = Tables.find(table.PathId);
@@ -590,7 +575,6 @@ private:
     TLockInfo::TPtr GetOrAddLock(ui64 lockId, ui32 lockNodeId);
     TLockInfo::TPtr AddLock(ui64 lockId, ui32 lockNodeId, ui32 generation, ui64 counter, TInstant createTs);
     void RemoveOneLock(ui64 lockId, ILocksDb* db = nullptr);
-    void RemoveBrokenLocks();
 
     void SaveBrokenPersistentLocks(ILocksDb* db);
 };