Fix excessive timer allocation in rate limiter (#6354)

* Fix excessive timer allocation in rate limiter * Looser match * Rename to currentNextSlot
NethermindEth · Dec 11, 2023 · 1040052 · 1040052
1 parent 600b7b6
commit 1040052
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 18 deletions.
diff --git a/src/Nethermind/Nethermind.Core.Test/RateLimiterTests.cs b/src/Nethermind/Nethermind.Core.Test/RateLimiterTests.cs
@@ -23,14 +23,13 @@ public async Task RateLimiter_should_delay_wait_to_rate_limit(int eventPerSec, i
     {
         RateLimiter rateLimiter = new(eventPerSec);
 
-        TimeSpan duration = TimeSpan.FromMilliseconds(durationMs);
-        DateTimeOffset startTime = DateTimeOffset.Now;
-        DateTimeOffset deadline = startTime + duration;
+        long startTime = Environment.TickCount64;
+        long deadline = startTime + durationMs;
         long counter = 0;
 
         Task[] tasks = Enumerable.Range(0, concurrency).Select(async (_) =>
         {
-            while (DateTimeOffset.Now < deadline)
+            while (Environment.TickCount64 < deadline)
             {
                 Interlocked.Increment(ref counter);
                 await rateLimiter.WaitAsync(CancellationToken.None);
@@ -39,7 +38,7 @@ public async Task RateLimiter_should_delay_wait_to_rate_limit(int eventPerSec, i
 
         await Task.WhenAll(tasks);
 
-        int effectivePerSec = (int)(counter / (DateTimeOffset.Now - startTime).TotalSeconds);
+        int effectivePerSec = (int)(counter / ((Environment.TickCount64 - startTime) / 1000.0));
         effectivePerSec.Should().BeInRange((int)(eventPerSec * 0.5), (int)(eventPerSec * 1.1));
     }
 

diff --git a/src/Nethermind/Nethermind.Core/RateLimiter.cs b/src/Nethermind/Nethermind.Core/RateLimiter.cs
@@ -50,23 +50,22 @@ public bool IsThrottled()
 
     public async ValueTask WaitAsync(CancellationToken ctx)
     {
+        long currentNextSlot = _nextSlot;
         while (true)
         {
-            long originalNextSlot = _nextSlot;
-
-            // Technically its possible that two `GetCurrentTick()` call at the same time can return same value,
-            // but its very unlikely.
-            long now = GetCurrentTick();
-            if (now >= originalNextSlot
-                && Interlocked.CompareExchange(ref _nextSlot, now + _delay, originalNextSlot) == originalNextSlot)
+            long nextSlot = Interlocked.CompareExchange(ref _nextSlot, currentNextSlot + _delay, currentNextSlot);
+            if (nextSlot == currentNextSlot)
             {
-                return;
+                break;
             }
+            currentNextSlot = nextSlot;
+        }
 
-            long toWait = originalNextSlot - now;
-            if (toWait < 0) continue;
+        long now = GetCurrentTick();
+        long toWait = currentNextSlot - now;
 
-            await Task.Delay(TimeSpan.FromMilliseconds(TickToMs(toWait)), ctx);
-        }
+        if (toWait <= 0) return;
+
+        await Task.Delay(TimeSpan.FromMilliseconds(TickToMs(toWait)), ctx);
     }
 }
diff --git a/src/Nethermind/Nethermind.Network.Discovery.Test/DiscoveryManagerTests.cs b/src/Nethermind/Nethermind.Network.Discovery.Test/DiscoveryManagerTests.cs
@@ -208,7 +208,7 @@ public async Task RateLimitOutgoingMessage()
             await _discoveryManager.SendMessageAsync(msg);
             await _discoveryManager.SendMessageAsync(msg);
             await _discoveryManager.SendMessageAsync(msg);
-            sw.Elapsed.Should().BeGreaterOrEqualTo(TimeSpan.FromSeconds(1));
+            sw.Elapsed.Should().BeGreaterOrEqualTo(TimeSpan.FromSeconds(0.9));
         }
     }
 }