Improve performance of ConsolidatorScanPriority comparison (#8452)

* Improve the performance of ConsolidatorScanPriority comparison in SubscriptionManager * address review comments
QuantConnect · Dec 11, 2024 · 7445f7e · 7445f7e
1 parent 8febb9a
commit 7445f7e
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 16 deletions.
diff --git a/Common/Data/ConsolidatorWrapper.cs b/Common/Data/ConsolidatorWrapper.cs
@@ -14,6 +14,7 @@
 */
 
 using System;
+using System.Collections.Generic;
 using System.Threading;
 using QuantConnect.Interfaces;
 using QuantConnect.Data.Consolidators;
@@ -138,8 +139,24 @@ private void AdvanceScanTime(object _ = null, IBaseData consolidated = null)
         }
     }
 
-    internal class ConsolidatorScanPriority : IComparable
+    internal class ConsolidatorScanPriority
     {
+        private sealed class UtcScanTimeIdRelationalComparer : IComparer<ConsolidatorScanPriority>
+        {
+            public int Compare(ConsolidatorScanPriority? x, ConsolidatorScanPriority? y)
+            {
+                if (ReferenceEquals(x, y)) return 0;
+                if (y is null) return 1;
+                if (x is null) return -1;
+                var utcScanTimeComparison = x.UtcScanTime.CompareTo(y.UtcScanTime);
+                if (utcScanTimeComparison != 0) return utcScanTimeComparison;
+                return x.Id.CompareTo(y.Id);
+            }
+        }
+
+        public static IComparer<ConsolidatorScanPriority> Comparer { get; } =
+            new UtcScanTimeIdRelationalComparer();
+
         /// <summary>
         /// The next utc scan time
         /// </summary>
@@ -155,19 +172,5 @@ public ConsolidatorScanPriority(DateTime utcScanTime, long id)
             Id = id;
             UtcScanTime = utcScanTime;
         }
-
-        public int CompareTo(object obj)
-        {
-            if (obj == null) return 1;
-
-            var other = (ConsolidatorScanPriority)obj;
-            var result = UtcScanTime.CompareTo(other.UtcScanTime);
-            if (result == 0)
-            {
-                // if they are the same let's compare Ids too
-                return Id.CompareTo(other.Id);
-            }
-            return result;
-        }
     }
 }
diff --git a/Common/Data/SubscriptionManager.cs b/Common/Data/SubscriptionManager.cs
@@ -66,7 +66,7 @@ public SubscriptionManager(ITimeKeeper timeKeeper)
         {
             _consolidators = new();
             _timeKeeper = timeKeeper;
-            _consolidatorsSortedByScanTime = new(1000);
+            _consolidatorsSortedByScanTime = new(1000, ConsolidatorScanPriority.Comparer);
             _threadSafeCollectionLock = new object();
         }
 

diff --git a/Tests/Common/Data/ConsolidatorWrapperTests.cs b/Tests/Common/Data/ConsolidatorWrapperTests.cs
@@ -163,6 +163,47 @@ public void ScanTimeOnWorkingBarDayLightSavings(int hoursShift, bool savingsStar
             Assert.AreEqual(consolidator.WorkingData.EndTime.ConvertToUtc(tz), wrapper.UtcScanTime);
         }
 
+        [Test]
+        public void ConsolidatorScanPriorityComparerComparesByUtcScanDate()
+        {
+            const int id = 1;
+            var utcScanTime = new DateTime(2024, 12, 10, 0, 0, 0, DateTimeKind.Utc);
+
+            var priority1 = new ConsolidatorScanPriority(utcScanTime, id);
+            var priority2 = new ConsolidatorScanPriority(utcScanTime.AddSeconds(1), id + 1);
+            var priority3 = new ConsolidatorScanPriority(utcScanTime, id + 1);
+
+            Assert.AreEqual(-1, ConsolidatorScanPriority.Comparer.Compare(priority1, priority2));
+            Assert.AreEqual(1, ConsolidatorScanPriority.Comparer.Compare(priority2, priority1));
+            Assert.AreEqual(1, ConsolidatorScanPriority.Comparer.Compare(priority3, priority1));
+            Assert.AreEqual(-1, ConsolidatorScanPriority.Comparer.Compare(priority3, priority2));
+            Assert.AreEqual(0, ConsolidatorScanPriority.Comparer.Compare(priority1, priority1));
+        }
+
+        [Test]
+        public void ConsolidatorScanPriorityComparerComparesByIdIfUtcScanTimesAreEqual()
+        {
+            const int id = 1;
+            var utcScanTime = new DateTime(2024, 12, 10, 0, 0, 0, DateTimeKind.Utc);
+
+            var priority1 = new ConsolidatorScanPriority(utcScanTime, id);
+            var priority2 = new ConsolidatorScanPriority(utcScanTime, id + 1);
+
+            Assert.AreEqual(1, ConsolidatorScanPriority.Comparer.Compare(priority2, priority1));
+        }
+
+        [Test]
+        public void ConsolidatorScanPriorityComparerTreatsNullsRight()
+        {
+            const int id = 1;
+            var utcScanTime = new DateTime(2024, 12, 10, 0, 0, 0, DateTimeKind.Utc);
+            var priority1 = new ConsolidatorScanPriority(utcScanTime, id);
+
+            Assert.AreEqual(1, ConsolidatorScanPriority.Comparer.Compare(priority1, null));
+            Assert.AreEqual(-1, ConsolidatorScanPriority.Comparer.Compare(null, priority1));
+            Assert.AreEqual(0, ConsolidatorScanPriority.Comparer.Compare(null, null));
+        }
+
         private class TestConsolidator : IDataConsolidator
         {
             public IBaseData Consolidated { get; set; }