Extract TAsyncExpiringCache::Erase helper to always clear probation cookie

commit_hash:84008f71594abd4c8cc1bf3cd0e2539ce2fafe4a

2 files changed, 19 insertions, 22 deletions

diff --git a/yt/yt/core/misc/async_expiring_cache-inl.h b/yt/yt/core/misc/async_expiring_cache-inl.h
index 2723bef9f5..34499d55de 100644
--- a/yt/yt/core/misc/async_expiring_cache-inl.h
+++ b/yt/yt/core/misc/async_expiring_cache-inl.h
@@ -85,9 +85,7 @@ typename TAsyncExpiringCache<TKey, TValue>::TExtendedGetResult TAsyncExpiringCac
         if (auto it = Map_.find(key); it != Map_.end()) {
             auto& entry = it->second;
             if (entry->Promise.IsSet() && entry->IsExpired(now)) {
-                NConcurrency::TDelayedExecutor::CancelAndClear(entry->ProbationCookie);
-                Map_.erase(it);
-                OnRemoved(key);
+                Erase(it);
             } else {
                 HitCounter_.Increment();
                 entry->AccessDeadline = now + NProfiling::DurationToCpuDuration(Config()->ExpireAfterAccessTime);
@@ -173,12 +171,10 @@ TFuture<std::vector<TErrorOr<TValue>>> TAsyncExpiringCache<TKey, TValue>::GetMan
 
         for (auto index : preliminaryIndexesToPopulate) {
             const auto& key = keys[index];
-            if (auto it = Map_.find(keys[index]); it != Map_.end()) {
+            if (auto it = Map_.find(key); it != Map_.end()) {
                 auto& entry = it->second;
                 if (entry->Promise.IsSet() && entry->IsExpired(now)) {
-                    NConcurrency::TDelayedExecutor::CancelAndClear(entry->ProbationCookie);
-                    Map_.erase(it);
-                    OnRemoved(key);
+                    Erase(it);
                 } else {
                     handleHit(index, entry);
                     continue;
@@ -282,9 +278,7 @@ void TAsyncExpiringCache<TKey, TValue>::InvalidateActive(const TKey& key)
 
     auto guard = WriterGuard(SpinLock_);
     if (auto it = Map_.find(key); it != Map_.end() && it->second->Promise.IsSet()) {
-        NConcurrency::TDelayedExecutor::CancelAndClear(it->second->ProbationCookie);
-        Map_.erase(it);
-        OnRemoved(key);
+        Erase(it);
         SizeCounter_.Update(Map_.size());
     }
 }
@@ -309,9 +303,7 @@ void TAsyncExpiringCache<TKey, TValue>::InvalidateValue(const TKey& key, const T
     if (auto it = Map_.find(key); it != Map_.end() && it->second->Promise.IsSet()) {
         auto valueOrError = it->second->Promise.Get();
         if (valueOrError.IsOK() && valueOrError.Value() == value) {
-            NConcurrency::TDelayedExecutor::CancelAndClear(it->second->ProbationCookie);
-            Map_.erase(it);
-            OnRemoved(key);
+            Erase(it);
             SizeCounter_.Update(Map_.size());
         }
     }
@@ -385,8 +377,7 @@ void TAsyncExpiringCache<TKey, TValue>::Set(const TKey& key, TErrorOr<TValue> va
             promise = entry->Promise;
         }
         if (expirationTime == TDuration::Zero()) {
-            Map_.erase(it);
-            OnRemoved(key);
+            Erase(it);
         } else {
             entry->AccessDeadline = accessDeadline;
             entry->UpdateDeadline = updateDeadline;
@@ -517,9 +508,7 @@ void TAsyncExpiringCache<TKey, TValue>::SetResult(
     }
 
     if (entry->IsExpired(now) || (entryUpdated && expirationTime == TDuration::Zero())) {
-        Map_.erase(it);
-        OnRemoved(key);
-        SizeCounter_.Update(Map_.size());
+        Erase(it);
         return;
     }
 
@@ -564,8 +553,7 @@ bool TAsyncExpiringCache<TKey, TValue>::TryEraseExpired(const TEntryPtr& entry,
         auto writerGuard = WriterGuard(SpinLock_);
 
         if (auto it = Map_.find(key); it != Map_.end() && entry == it->second && now > it->second->AccessDeadline) {
-            Map_.erase(it);
-            OnRemoved(key);
+            Erase(it);
             SizeCounter_.Update(Map_.size());
         }
         return true;
@@ -574,6 +562,14 @@ bool TAsyncExpiringCache<TKey, TValue>::TryEraseExpired(const TEntryPtr& entry,
 }
 
 template <class TKey, class TValue>
+void TAsyncExpiringCache<TKey, TValue>::Erase(THashMap<TKey, TEntryPtr>::iterator it)
+{
+    NConcurrency::TDelayedExecutor::CancelAndClear(it->second->ProbationCookie);
+    OnRemoved(it->first);
+    Map_.erase(it);
+}
+
+template <class TKey, class TValue>
 void TAsyncExpiringCache<TKey, TValue>::InvokeGetMany(
     const std::vector<TWeakPtr<TEntry>>& entries,
     const std::vector<TKey>& keys,
@@ -666,8 +662,7 @@ void TAsyncExpiringCache<TKey, TValue>::UpdateAll()
                 const auto& [_, entry] = *it;
                 if (entry->Promise.IsSet()) {
                     if (now > entry->AccessDeadline) {
-                        Map_.erase(it);
-                        OnRemoved(key);
+                        Erase(it);
                         SizeCounter_.Update(Map_.size());
                     } else if (entry->Future.Get().IsOK()) {
                         keys.push_back(key);
diff --git a/yt/yt/core/misc/async_expiring_cache.h b/yt/yt/core/misc/async_expiring_cache.h
index 2031c6e31f..4fc7759cb2 100644
--- a/yt/yt/core/misc/async_expiring_cache.h
+++ b/yt/yt/core/misc/async_expiring_cache.h
@@ -152,6 +152,8 @@ private:
         const TEntryPtr& Entry,
         const TKey& key);
 
+    void Erase(THashMap<TKey, TEntryPtr>::iterator it);
+
     void UpdateAll();
 
     void ScheduleEntryRefresh(