Optimize Count(), fix debug assert, and fix SetExpiration cleanup

- Count(): fast-path when no expirations are due yet (peek min), iterate
  RawHeap instead of sortedSetDict for expired count
- Add RawHeap property to IndexedPriorityQueue for fast iteration
- Fix Debug.Assert to use correct minimum (SortedSetOverhead + DictionaryOverhead)
- Fix SetExpiration: remove from expirationQueue when expiration is in the past

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: hamdaankhalidmsft

libs/common/Collections/IndexedPriorityQueue.cs

@@ -22,6 +22,11 @@ public class IndexedPriorityQueue<TElement, TPriority>
         private (TElement element, TPriority priority)[] _heap = [];
         private int _count;
 
+        /// <summary>
+        /// Raw heap access to do iteration really fast
+        /// </summary>
+        public (TElement element, TPriority priority)[] RawHeap => _heap;
+
         /// <summary>
         /// Number of elements in the priority queue.
         /// </summary>

libs/server/Objects/SortedSet/SortedSetObject.cs

@@ -571,14 +571,18 @@ public bool TryGetScore(byte[] key, out double value)
         /// <returns>The count of elements in the sorted set.</returns>
         public int Count()
         {
-            if (!HasExpirableItems())
+            if (!HasExpirableItems() || (expirationQueue.TryPeek(out _, out var minExpiration) && minExpiration > DateTimeOffset.UtcNow.Ticks))
                 return sortedSetDict.Count;
 
             var expiredKeysCount = 0;
-            foreach (var item in sortedSetDict)
+            var rawHeap = expirationQueue.RawHeap;
+            var heapCount = expirationQueue.Count;
+            for (int i = 0; i < heapCount; i++)
             {
-                if (IsExpired(item.Key))
+                if (rawHeap[i].priority < DateTimeOffset.UtcNow.Ticks)
+                {
                     expiredKeysCount++;
+                }
             }
             return sortedSetDict.Count - expiredKeysCount;
         }
@@ -616,7 +620,7 @@ private void UpdateExpirationSize(bool add)
             else
             {
                 HeapMemorySize -= MemoryUtils.IndexedPriorityQueueEntryOverhead;
-                Debug.Assert(HeapMemorySize >= MemoryUtils.DictionaryOverhead);
+                Debug.Assert(HeapMemorySize >= MemoryUtils.DictionaryOverhead + MemoryUtils.SortedSetOverhead);
             }
         }
 
@@ -669,6 +673,7 @@ private int SetExpiration(byte[] key, long expiration, ExpireOption expireOption
             {
                 _ = sortedSetDict.Remove(key, out var value);
                 _ = sortedSet.Remove((value, key));
+                _ = expirationQueue.TryRemove(key);
                 UpdateSize(key, add: false);
                 return (int)SortedSetExpireResult.KeyAlreadyExpired;
             }