Implement min read rate using new APIjohluo/min-rates

5 files changed, 141 insertions, 51 deletions

diff --git a/src/Servers/Kestrel/Core/src/Internal/Infrastructure/MinDataRateLimiter.cs b/src/Servers/Kestrel/Core/src/Internal/Infrastructure/MinDataRateLimiter.cs
new file mode 100644
index 0000000000..d3c06cfdf0
--- /dev/null
+++ b/src/Servers/Kestrel/Core/src/Internal/Infrastructure/MinDataRateLimiter.cs
@@ -0,0 +1,81 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Internal;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
+{
+    internal class MinDataRateLimiter : IRateLimiter
+    {
+        private long _bytesTokens;
+        private MinDataRate _minDataRate;
+        private long _initialTimestamp;
+        private long _lastTimestamp;
+        private double _rateRemainder;
+
+        public MinDataRateLimiter(MinDataRate minDataRate, long initialTimestamp)
+        {
+            _minDataRate = minDataRate;
+            _initialTimestamp = initialTimestamp;
+            _lastTimestamp = initialTimestamp;
+        }
+
+        public long EstimatedCount {
+            get {
+                if (Interlocked.Read(ref _lastTimestamp) - _initialTimestamp <= _minDataRate.GracePeriod.Ticks)
+                {
+                    return long.MaxValue;
+                }
+
+                return Interlocked.Read(ref _bytesTokens);
+            }
+        }
+
+        public ValueTask<bool> AcquireAsync(long requestedCount, CancellationToken cancellationToken = default)
+        {
+            Interlocked.Add(ref _bytesTokens, requestedCount);
+
+            return ValueTask.FromResult(true);
+        }
+
+        public bool TryAcquire(long requestedCount)
+        {
+            Interlocked.Add(ref _bytesTokens, requestedCount);
+
+            return true;
+        }
+
+        internal void Tick(long timestamp)
+        {
+            // Noop if less than 1 second has elapsed
+            var elapsedTicks = timestamp - Interlocked.Read(ref _lastTimestamp);
+            if (elapsedTicks < TimeSpan.TicksPerSecond)
+            {
+                return;
+            }
+
+            // No need to update byte counts count is non-positive
+            if (Interlocked.Read(ref _bytesTokens) <= 0)
+            {
+                _rateRemainder = 0.0;
+                Interlocked.Exchange(ref _lastTimestamp, timestamp);
+                return;
+            }
+
+            // Assume overly long tick intervals are the result of server resource starvation.
+            // Don't count extra time between ticks against the rate limit.
+            var bytes = _minDataRate.BytesPerSecond*(Math.Min(Heartbeat.Interval.Ticks, elapsedTicks)/TimeSpan.TicksPerSecond) + _rateRemainder;
+            var bytesToDecrement = (long)bytes;
+            _rateRemainder = bytes - bytesToDecrement;
+
+            var newTokens = Interlocked.Add(ref _bytesTokens, -bytesToDecrement);
+
+            if (newTokens < 0)
+            {
+                Interlocked.Add(ref _bytesTokens, newTokens);
+            }
+
+            Interlocked.Exchange(ref _lastTimestamp, timestamp);
+        }
+    }
+}
diff --git a/src/Servers/Kestrel/Core/src/Internal/Infrastructure/TimeoutControl.cs b/src/Servers/Kestrel/Core/src/Internal/Infrastructure/TimeoutControl.cs
index f4a460abce..d5e0eb5603 100644
--- a/src/Servers/Kestrel/Core/src/Internal/Infrastructure/TimeoutControl.cs
+++ b/src/Servers/Kestrel/Core/src/Internal/Infrastructure/TimeoutControl.cs
@@ -18,10 +18,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
 
         private readonly object _readTimingLock = new object();
         private MinDataRate? _minReadRate;
+        private MinDataRateLimiter? _minReadRateLimiter;
         private bool _readTimingEnabled;
         private bool _readTimingPauseRequested;
-        private long _readTimingElapsedTicks;
-        private long _readTimingBytesRead;
         private InputFlowControl? _connectionInputFlowControl;
         // The following are always 0 or 1 for HTTP/1.x
         private int _concurrentIncompleteRequestBodies;
@@ -98,19 +97,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
                     return;
                 }
 
-                // Assume overly long tick intervals are the result of server resource starvation.
-                // Don't count extra time between ticks against the rate limit.
-                _readTimingElapsedTicks += Math.Min(timestamp - _lastTimestamp, Heartbeat.Interval.Ticks);
-
                 Debug.Assert(_minReadRate != null);
+                Debug.Assert(_minReadRateLimiter != null);
 
-                if (_minReadRate.BytesPerSecond > 0 && _readTimingElapsedTicks > _minReadRate.GracePeriod.Ticks)
-                {
-                    var elapsedSeconds = (double)_readTimingElapsedTicks / TimeSpan.TicksPerSecond;
-                    var rate = _readTimingBytesRead / elapsedSeconds;
-
-                    timeout = rate < _minReadRate.BytesPerSecond && !Debugger.IsAttached;
-                }
+                timeout = _minReadRateLimiter.EstimatedCount < _minReadRate.BytesPerSecond && !Debugger.IsAttached;
+                _minReadRateLimiter.Tick(timestamp);
 
                 // PauseTimingReads() cannot just set _timingReads to false. It needs to go through at least one tick
                 // before pausing, otherwise _readTimingElapsed might never be updated if PauseTimingReads() is always
@@ -194,13 +185,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
                 Debug.Assert(_concurrentIncompleteRequestBodies == 0 || minRate == _minReadRate, "Multiple simultaneous read data rates are not supported.");
 
                 _minReadRate = minRate;
+                _minReadRateLimiter = new MinDataRateLimiter(minRate, _lastTimestamp);
                 _concurrentIncompleteRequestBodies++;
-
-                if (_concurrentIncompleteRequestBodies == 1)
-                {
-                    _readTimingElapsedTicks = 0;
-                    _readTimingBytesRead = 0;
-                }
             }
         }
 
@@ -247,9 +233,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
         {
             Debug.Assert(count >= 0, "BytesRead count must not be negative.");
 
+            // lock might not be needed?
             lock (_readTimingLock)
             {
-                _readTimingBytesRead += count;
+                _minReadRateLimiter?.TryAcquire(count);
             }
         }
 
diff --git a/src/Servers/Kestrel/Core/src/Microsoft.AspNetCore.Server.Kestrel.Core.csproj b/src/Servers/Kestrel/Core/src/Microsoft.AspNetCore.Server.Kestrel.Core.csproj
index ef7726f15f..22a6309c11 100644
--- a/src/Servers/Kestrel/Core/src/Microsoft.AspNetCore.Server.Kestrel.Core.csproj
+++ b/src/Servers/Kestrel/Core/src/Microsoft.AspNetCore.Server.Kestrel.Core.csproj
@@ -15,6 +15,7 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <Compile Include="$(SharedSourceRoot)RateLimit\*.cs" />
     <Compile Include="$(SharedSourceRoot)CertificateGeneration\**\*.cs" />
     <Compile Include="$(SharedSourceRoot)ValueTaskExtensions\**\*.cs" />
     <Compile Include="$(SharedSourceRoot)UrlDecoder\**\*.cs" />
diff --git a/src/Servers/Kestrel/Core/test/TimeoutControlTests.cs b/src/Servers/Kestrel/Core/test/TimeoutControlTests.cs
index 9477729e3a..40ef24ea75 100644
--- a/src/Servers/Kestrel/Core/test/TimeoutControlTests.cs
+++ b/src/Servers/Kestrel/Core/test/TimeoutControlTests.cs
@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
@@ -84,6 +84,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             now += TimeSpan.FromSeconds(1);
             _timeoutControl.Tick(now);
             now += TimeSpan.FromSeconds(1);
+            _timeoutControl.Tick(now);
             _timeoutControl.BytesRead(10);
             _timeoutControl.Tick(now);
 
@@ -92,7 +93,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
         }
 
         [Fact]
-        public void RequestBodyDataRateIsAveragedOverTimeSpentReadingRequestBody()
+        public void RequestBodyDataRateIsTickedDownOverTimeSpentReadingRequestBody()
         {
             var gracePeriod = TimeSpan.FromSeconds(2);
             var minRate = new MinDataRate(bytesPerSecond: 100, gracePeriod: gracePeriod);
@@ -104,14 +105,14 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             _timeoutControl.StartRequestBody(minRate);
             _timeoutControl.StartTimingRead();
 
-            // Set base data rate to 200 bytes/second
+            // Set base data rate to 300 bytes during grace period
             now += TimeSpan.FromSeconds(1);
             _timeoutControl.Tick(now);
             now += TimeSpan.FromSeconds(1);
-            _timeoutControl.BytesRead(400);
+            _timeoutControl.BytesRead(300);
             _timeoutControl.Tick(now);
 
-            // Data rate: 200 bytes/second
+            // Data rate: 300 bytes after tick
             now += TimeSpan.FromSeconds(1);
             _timeoutControl.BytesRead(200);
             _timeoutControl.Tick(now);
@@ -119,7 +120,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             // Not timed out
             _mockTimeoutHandler.Verify(h => h.OnTimeout(It.IsAny<TimeoutReason>()), Times.Never);
 
-            // Data rate: 150 bytes/second
+            // Data rate: 200 bytes after tick
             now += TimeSpan.FromSeconds(1);
             _timeoutControl.BytesRead(0);
             _timeoutControl.Tick(now);
@@ -127,7 +128,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             // Not timed out
             _mockTimeoutHandler.Verify(h => h.OnTimeout(It.IsAny<TimeoutReason>()), Times.Never);
 
-            // Data rate: 120 bytes/second
+            // Data rate: 100 bytes after tick
             now += TimeSpan.FromSeconds(1);
             _timeoutControl.BytesRead(0);
             _timeoutControl.Tick(now);
@@ -135,7 +136,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             // Not timed out
             _mockTimeoutHandler.Verify(h => h.OnTimeout(It.IsAny<TimeoutReason>()), Times.Never);
 
-            // Data rate: 100 bytes/second
+            // Data rate: 0 bytes after tick
             now += TimeSpan.FromSeconds(1);
             _timeoutControl.BytesRead(0);
             _timeoutControl.Tick(now);
@@ -143,12 +144,12 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             // Not timed out
             _mockTimeoutHandler.Verify(h => h.OnTimeout(It.IsAny<TimeoutReason>()), Times.Never);
 
-            // Data rate: ~85 bytes/second
+            // Data rate: Timed out
             now += TimeSpan.FromSeconds(1);
             _timeoutControl.BytesRead(0);
             _timeoutControl.Tick(now);
 
-            // Timed out
+            // Verify timed out
             _mockTimeoutHandler.Verify(h => h.OnTimeout(TimeoutReason.ReadDataRate), Times.Once);
         }
 
@@ -164,24 +165,24 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             _timeoutControl.StartRequestBody(minRate);
             _timeoutControl.StartTimingRead();
 
-            // Tick at 3s, expected counted time is 3s, expected data rate is 200 bytes/second
+            // Tick at 3s, calculated data rate is 300 bytes/second before rate tick 1
             _systemClock.UtcNow += TimeSpan.FromSeconds(1);
             _timeoutControl.Tick(_systemClock.UtcNow);
             _systemClock.UtcNow += TimeSpan.FromSeconds(1);
             _timeoutControl.Tick(_systemClock.UtcNow);
             _systemClock.UtcNow += TimeSpan.FromSeconds(1);
-            _timeoutControl.BytesRead(600);
+            _timeoutControl.BytesRead(400);
             _timeoutControl.Tick(_systemClock.UtcNow);
 
             // Pause at 3.5s
             _systemClock.UtcNow += TimeSpan.FromSeconds(0.5);
             _timeoutControl.StopTimingRead();
 
-            // Tick at 4s, expected counted time is 4s (first tick after pause goes through), expected data rate is 150 bytes/second
+            // Tick at 4s (first tick after pause goes through), expected data rate is 300 bytes/second before rate tick 2
             _systemClock.UtcNow += TimeSpan.FromSeconds(0.5);
             _timeoutControl.Tick(_systemClock.UtcNow);
 
-            // Tick at 6s, expected counted time is 4s, expected data rate is 150 bytes/second
+            // Tick at 6s, expected data rate is still 300 bytes/second
             _systemClock.UtcNow += TimeSpan.FromSeconds(2);
             _timeoutControl.Tick(_systemClock.UtcNow);
 
@@ -192,16 +193,18 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             _systemClock.UtcNow += TimeSpan.FromSeconds(0.5);
             _timeoutControl.StartTimingRead();
 
-            // Tick at 9s, expected counted time is 6s, expected data rate is 100 bytes/second
+            // Tick at 8s,  expected data rate is 200 bytes/second before rate tick 3
             _systemClock.UtcNow += TimeSpan.FromSeconds(1.0);
             _timeoutControl.Tick(_systemClock.UtcNow);
-            _systemClock.UtcNow += TimeSpan.FromSeconds(.5);
+
+            // Tick at 9s,  expected data rate is 100 bytes/second before rate tick 4
+            _systemClock.UtcNow += TimeSpan.FromSeconds(1.0);
             _timeoutControl.Tick(_systemClock.UtcNow);
 
             // Not timed out
             _mockTimeoutHandler.Verify(h => h.OnTimeout(It.IsAny<TimeoutReason>()), Times.Never);
 
-            // Tick at 10s, expected counted time is 7s, expected data rate drops below 100 bytes/second
+            // Tick at 10s, expected data rate drops below 100 bytes/second before rate tick 5
             _systemClock.UtcNow += TimeSpan.FromSeconds(1);
             _timeoutControl.Tick(_systemClock.UtcNow);
 
@@ -220,7 +223,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             _timeoutControl.StartRequestBody(minRate);
             _timeoutControl.StartTimingRead();
 
-            // Tick at 2s, expected counted time is 2s, expected data rate is 100 bytes/second
+            // Tick at 2s, expected counted time is 2s, expected data rate is 200 bytes/second
             _systemClock.UtcNow += TimeSpan.FromSeconds(1);
             _timeoutControl.Tick(_systemClock.UtcNow);
             _systemClock.UtcNow += TimeSpan.FromSeconds(1);
@@ -238,18 +241,20 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             _systemClock.UtcNow += TimeSpan.FromSeconds(0.25);
             _timeoutControl.StartTimingRead();
 
-            // Tick at 3s, expected counted time is 3s, expected data rate is 100 bytes/second
+            // Tick at 3s, expected counted time is 3s, expected data rate is 100 bytes/second after tick
             _systemClock.UtcNow += TimeSpan.FromSeconds(0.5);
-            _timeoutControl.BytesRead(100);
             _timeoutControl.Tick(_systemClock.UtcNow);
 
             // Not timed out
             _mockTimeoutHandler.Verify(h => h.OnTimeout(It.IsAny<TimeoutReason>()), Times.Never);
 
-            // Tick at 4s, expected counted time is 4s, expected data rate drops below 100 bytes/second
+            // Tick at 4s, expected counted time is 4s, expected data rate drops below 100 bytes/second after tick
             _systemClock.UtcNow += TimeSpan.FromSeconds(1);
             _timeoutControl.Tick(_systemClock.UtcNow);
 
+            // Tick at 4s, this time it will timeout
+            _timeoutControl.Tick(_systemClock.UtcNow);
+
             // Timed out
             _mockTimeoutHandler.Verify(h => h.OnTimeout(TimeoutReason.ReadDataRate), Times.Once);
         }
@@ -301,18 +306,20 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             _timeoutControl.StartRequestBody(minRate);
             _timeoutControl.StartTimingRead();
 
-            // Tick past grace period
+            // Tick past grace period, 99 bytes after tick 3
             now += TimeSpan.FromSeconds(1);
             _timeoutControl.BytesRead(100);
             _timeoutControl.Tick(now);
             now += TimeSpan.FromSeconds(1);
-            _timeoutControl.BytesRead(100);
+            _timeoutControl.BytesRead(99);
+            _timeoutControl.Tick(now);
+            now += TimeSpan.FromSeconds(1);
             _timeoutControl.Tick(now);
 
             // Induce low flow control availability
             flowControl.TryAdvance(2);
 
-            // Read 0 bytes in 1 second
+            // Read 0 bytes in 1 second, this tick should timeout since 99 bytes is below minimum, but low flow control means this is ignored
             now += TimeSpan.FromSeconds(1);
             _timeoutControl.Tick(now);
 
@@ -323,13 +330,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             flowControl.TryUpdateWindow(2, out _);
             _timeoutControl.Tick(now);
 
-            // Still not timed out
-            _mockTimeoutHandler.Verify(h => h.OnTimeout(It.IsAny<TimeoutReason>()), Times.Never);
-            
-            // Read 0 bytes in 1 second
-            now += TimeSpan.FromSeconds(1);
-            _timeoutControl.Tick(now);
-
             // Timed out
             _mockTimeoutHandler.Verify(h => h.OnTimeout(TimeoutReason.ReadDataRate), Times.Once);
         }
@@ -521,6 +521,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
 
             // Tick after grace period w/ low data rate
             _systemClock.UtcNow += TimeSpan.FromSeconds(1);
+            _timeoutControl.Tick(_systemClock.UtcNow);
             _timeoutControl.BytesRead(1);
             _timeoutControl.Tick(_systemClock.UtcNow);
         }
diff --git a/src/Shared/RateLimit/IRateLimiter.cs b/src/Shared/RateLimit/IRateLimiter.cs
new file mode 100644
index 0000000000..376de298b0
--- /dev/null
+++ b/src/Shared/RateLimit/IRateLimiter.cs
@@ -0,0 +1,20 @@
+
+// Single resource
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.Internal
+{
+    internal interface IRateLimiter
+    {
+        // an inaccurate view of resources
+        long EstimatedCount { get; }
+
+        // Fast synchronous attempt to acquire resources, it won't actually acquire the resource
+        bool TryAcquire(long requestedCount);
+
+        // Wait until the requested resources are available
+        ValueTask<bool> AcquireAsync(long requestedCount, CancellationToken cancellationToken = default);
+
+    }
+}
\ No newline at end of file