From 5785544f6a34763789b879627806589752da2eec Mon Sep 17 00:00:00 2001
From: Alex Peck <alexpeck@microsoft.com>
Date: Mon, 13 Nov 2023 12:41:42 -0800
Subject: [PATCH] trim volatile

---
 BitFaster.Caching/Lru/ConcurrentLruCore.cs | 58 ++++++++++++++++------
 1 file changed, 44 insertions(+), 14 deletions(-)

diff --git a/BitFaster.Caching/Lru/ConcurrentLruCore.cs b/BitFaster.Caching/Lru/ConcurrentLruCore.cs
index b3b75593..48c0977b 100644
--- a/BitFaster.Caching/Lru/ConcurrentLruCore.cs
+++ b/BitFaster.Caching/Lru/ConcurrentLruCore.cs
@@ -468,10 +468,9 @@ private void TrimExpired()
         // backcompat: make internal
         protected int TrimAllDiscardedItems()
         {
-            int itemsRemoved = 0;
-
-            void RemoveDiscardableItems(ConcurrentQueue<I> q, ref int queueCounter)
+            int RemoveDiscardableItems(ConcurrentQueue<I> q, ref int queueCounter)
             {
+                int itemsRemoved = 0;
                 int localCount = queueCounter;
 
                 for (int i = 0; i < localCount; i++)
@@ -490,13 +489,20 @@ void RemoveDiscardableItems(ConcurrentQueue<I> q, ref int queueCounter)
                         }
                     }
                 }
+
+                return itemsRemoved;
             }
 
-            RemoveDiscardableItems(coldQueue, ref this.counter.cold);
-            RemoveDiscardableItems(warmQueue, ref this.counter.warm);
-            RemoveDiscardableItems(hotQueue, ref this.counter.hot);
+            int coldRem = RemoveDiscardableItems(coldQueue, ref this.counter.cold);
+            int warmRem = RemoveDiscardableItems(warmQueue, ref this.counter.warm);
+            int hotRem = RemoveDiscardableItems(hotQueue, ref this.counter.hot);
+
+            if (warmRem > 0)
+            {
+                Volatile.Write(ref this.isWarm, false);
+            }
 
-            return itemsRemoved;
+            return coldRem + warmRem + hotRem;
         }
 
         private void TrimLiveItems(int itemsRemoved, int itemCount, ItemRemovedReason reason)
@@ -508,35 +514,59 @@ private void TrimLiveItems(int itemsRemoved, int itemCount, ItemRemovedReason re
             int trimWarmAttempts = 0;
             int maxWarmHotAttempts = (this.capacity.Warm * 2) + this.capacity.Hot;
 
+            int warmCount = Volatile.Read(ref this.counter.warm);
+            int coldCount = Volatile.Read(ref this.counter.cold);
+
             while (itemsRemoved < itemCount && trimWarmAttempts < maxWarmHotAttempts)
             {
-                if (Volatile.Read(ref this.counter.cold) > 0)
+                if (coldCount > 0)
                 {
                     if (TryRemoveCold(reason) == (ItemDestination.Remove, 0))
                     {
                         itemsRemoved++;
+                        coldCount--;
                         trimWarmAttempts = 0;
                     }
 
-                    TrimWarmOrHot(reason);
+                    TrimWarmOrHot(reason, ref warmCount, ref coldCount);
                 }
                 else
                 {
-                    TrimWarmOrHot(reason);
+                    TrimWarmOrHot(reason, ref warmCount, ref coldCount);
                     trimWarmAttempts++;
                 }
             }
+
+            if (warmCount < this.capacity.Warm)
+            {
+                Volatile.Write(ref this.isWarm, false);
+            }
         }
 
-        private void TrimWarmOrHot(ItemRemovedReason reason)
+        private void TrimWarmOrHot(ItemRemovedReason reason, ref int warmCount, ref int coldCount)
         {
-            if (Volatile.Read(ref this.counter.warm) > 0)
+            if (warmCount > 0)
             {
-                CycleWarmUnchecked(reason);
+                var (dest, count) = CycleWarmUnchecked(reason);
+                warmCount--;
+                UpdateCount(ref warmCount, ref coldCount, dest, count);
             }
             else if (Volatile.Read(ref this.counter.hot) > 0)
             {
-                CycleHotUnchecked(reason);
+                var (dest, count) = CycleHotUnchecked(reason);
+                UpdateCount(ref warmCount, ref coldCount, dest, count);
+            }
+
+            void UpdateCount(ref int warmCount, ref int coldCount, ItemDestination dest, int count)
+            {
+                if (dest == ItemDestination.Cold)
+                {
+                    coldCount = count;
+                }
+                else if (dest == ItemDestination.Warm)
+                {
+                    warmCount = count;
+                }
             }
         }