Fix rate limiting for repeated small writes (#8)

Author: halter73

src/Microsoft.AspNetCore.Server.Kestrel.Core/Internal/FrameConnection.cs

@@ -435,18 +435,27 @@ public void StartTimingWrite(long size)
 
                 if (minResponseDataRate != null)
                 {
-                    var timeoutTicks = Math.Max(
-                        minResponseDataRate.GracePeriod.Ticks,
-                        TimeSpan.FromSeconds(size / minResponseDataRate.BytesPerSecond).Ticks);
-
                     // Add Heartbeat.Interval since this can be called right before the next heartbeat.
+                    var currentTimeUpperBound = _lastTimestamp + Heartbeat.Interval.Ticks;
+                    var ticksToCompleteWriteAtMinRate = TimeSpan.FromSeconds(size / minResponseDataRate.BytesPerSecond).Ticks;
 
-                    // Don't penalize a connection for completing previous flushes more quickly than required.
+                    // If ticksToCompleteWriteAtMinRate is less than the configured grace period,
+                    // allow that write to take up to the grace period to complete. Only add the grace period
+                    // to the current time and not to any accumulated timeout.
+                    var singleWriteTimeoutTimestamp = currentTimeUpperBound + Math.Max(
+                        minResponseDataRate.GracePeriod.Ticks,
+                        ticksToCompleteWriteAtMinRate);
+
+                    // Don't penalize a connection for completing previous writes more quickly than required.
                     // We don't want to kill a connection when flushing the chunk terminator just because the previous
                     // chunk was large if the previous chunk was flushed quickly.
-                    _writeTimingTimeoutTimestamp = Math.Max(_writeTimingTimeoutTimestamp, _lastTimestamp + Heartbeat.Interval.Ticks);
 
-                    _writeTimingTimeoutTimestamp += timeoutTicks;
+                    // Don't add any grace period to this accumulated timeout because the grace period could
+                    // get accumulated repeatedly making the timeout for a bunch of consecutive small writes
+                    // far too conservative.
+                    var accumulatedWriteTimeoutTimestamp = _writeTimingTimeoutTimestamp + ticksToCompleteWriteAtMinRate;
+
+                    _writeTimingTimeoutTimestamp = Math.Max(singleWriteTimeoutTimestamp, accumulatedWriteTimeoutTimestamp);
                     _writeTimingWrites++;
                 }
             }

test/Microsoft.AspNetCore.Server.Kestrel.Core.Tests/FrameConnectionTests.cs

@@ -529,5 +529,56 @@ public void WriteTimingTimeoutPushedOnConcurrentWrite()
             Assert.True(_frameConnection.TimedOut);
             Assert.True(aborted.Wait(TimeSpan.FromSeconds(10)));
         }
+
+        [Fact]
+        public async Task WriteTimingAbortsConnectionWhenRepeadtedSmallWritesDoNotCompleteWithMinimumDataRate()
+        {
+            var systemClock = new MockSystemClock();
+            var minResponseDataRate = new MinDataRate(bytesPerSecond: 100, gracePeriod: TimeSpan.FromSeconds(5));
+            var numWrites = 5;
+            var writeSize = 100;
+            var aborted = new TaskCompletionSource<object>();
+
+            _frameConnectionContext.ServiceContext.ServerOptions.Limits.MinResponseDataRate = minResponseDataRate;
+            _frameConnectionContext.ServiceContext.SystemClock = systemClock;
+
+            var mockLogger = new Mock<IKestrelTrace>();
+            _frameConnectionContext.ServiceContext.Log = mockLogger.Object;
+
+            _frameConnection.CreateFrame(new DummyApplication(), _frameConnectionContext.Input.Reader, _frameConnectionContext.Output);
+            _frameConnection.Frame.Reset();
+            _frameConnection.Frame.RequestAborted.Register(() =>
+            {
+                aborted.SetResult(null);
+            });
+
+            // Initialize timestamp
+            var startTime = systemClock.UtcNow;
+            _frameConnection.Tick(startTime);
+
+            // 5 consecutive 100 byte writes.
+            for (var i = 0; i < numWrites - 1; i++)
+            {
+                _frameConnection.StartTimingWrite(writeSize);
+                _frameConnection.StopTimingWrite();
+            }
+
+            // Stall the last write.
+            _frameConnection.StartTimingWrite(writeSize);
+
+            // Move the clock forward Heartbeat.Interval + MinDataRate.GracePeriod + 4 seconds.
+            // The grace period should only be added for the first write. The subsequent 4 100 byte writes should add 1 second each to the timeout given the 100 byte/s min rate.
+            systemClock.UtcNow += Heartbeat.Interval + minResponseDataRate.GracePeriod + TimeSpan.FromSeconds((numWrites - 1) * writeSize / minResponseDataRate.BytesPerSecond);
+            _frameConnection.Tick(systemClock.UtcNow);
+
+            Assert.False(_frameConnection.TimedOut);
+
+            // On more tick forward triggers the timeout.
+            systemClock.UtcNow += TimeSpan.FromTicks(1);
+            _frameConnection.Tick(systemClock.UtcNow);
+
+            Assert.True(_frameConnection.TimedOut);
+            await aborted.Task.TimeoutAfter(TimeSpan.FromSeconds(10));
+        }
     }
 }

test/Microsoft.AspNetCore.Server.Kestrel.FunctionalTests/ResponseTests.cs

@@ -2446,8 +2446,8 @@ await connection.Receive(
         public async Task ConnectionClosedWhenResponseDoesNotSatisfyMinimumDataRate()
         {
             // Smaller chunks allow for quicker response rate timeouts.
-            var chunkSize = 1024;
-            var chunks = 8192;
+            const int chunkSize = 1024;
+            const int chunks = 256 * 1024;;
 
             var responseRateTimeoutMessageLogged = new TaskCompletionSource<object>();
             var connectionStopMessageLogged = new TaskCompletionSource<object>();
@@ -2518,7 +2518,7 @@ await connection.Receive(
         public async Task HttpsConnectionClosedWhenResponseDoesNotSatisfyMinimumDataRate()
         {
             const int chunkSize = 1024;
-            const int chunks = 8192;
+            const int chunks = 256 * 1024;
 
             var certificate = new X509Certificate2(TestResources.TestCertificatePath, "testPassword");