Work around potential race in PipeWriter Redux (#11065)

- Make sure we always await the last flush task before calling FlushAsync
  again instead of preemptively calling FlushAsync and checking to see
  if the ValueTask is incomplete before bothering to acquire the _flushLock
- This now acquires the _flushLock fore every call to Response.Body.Write
  whereas this only happened for truly async writes before.
- I don't think this is a big concern since this should normally be uncontested,
  and DefaultPipeWriter.FlushAsync/GetResult already acquire a lock.

Author: halter73

src/Servers/Kestrel/Core/src/Internal/Infrastructure/TimingPipeFlusher.cs

@@ -22,7 +22,13 @@ internal class TimingPipeFlusher
         private readonly IKestrelTrace _log;
 
         private readonly object _flushLock = new object();
-        private Task<FlushResult> _lastFlushTask = null;
+
+        // This field should only be get or set under the _flushLock. This is a ValueTask that was either:
+        // 1. The default value where "IsCompleted" is true
+        // 2. Created by an async method
+        // 3. Constructed explicitely from a completed result
+        // This means it should be safe to await a single _lastFlushTask instance multiple times.
+        private ValueTask<FlushResult> _lastFlushTask;
 
         public TimingPipeFlusher(
             PipeWriter writer,
@@ -51,35 +57,41 @@ public ValueTask<FlushResult> FlushAsync(MinDataRate minRate, long count)
 
         public ValueTask<FlushResult> FlushAsync(MinDataRate minRate, long count, IHttpOutputAborter outputAborter, CancellationToken cancellationToken)
         {
-            var flushValueTask = _writer.FlushAsync(cancellationToken);
+            // https://github.com/dotnet/corefxlab/issues/1334
+            // Pipelines don't support multiple awaiters on flush.
+            lock (_flushLock)
+            {
+                if (_lastFlushTask.IsCompleted)
+                {
+                    _lastFlushTask = TimeFlushAsync(minRate, count, outputAborter, cancellationToken);
+                }
+                else
+                {
+                    _lastFlushTask = AwaitLastFlushAndTimeFlushAsync(_lastFlushTask, minRate, count, outputAborter, cancellationToken);
+                }
+
+                return _lastFlushTask;
+            }
+        }
+
+        private ValueTask<FlushResult> TimeFlushAsync(MinDataRate minRate, long count, IHttpOutputAborter outputAborter, CancellationToken cancellationToken)
+        {
+            var pipeFlushTask = _writer.FlushAsync(cancellationToken);
 
             if (minRate != null)
             {
                 _timeoutControl.BytesWrittenToBuffer(minRate, count);
             }
 
-            if (flushValueTask.IsCompletedSuccessfully)
+            if (pipeFlushTask.IsCompletedSuccessfully)
             {
-                return new ValueTask<FlushResult>(flushValueTask.Result);
+                return new ValueTask<FlushResult>(pipeFlushTask.Result);
             }
 
-            // https://github.com/dotnet/corefxlab/issues/1334
-            // Pipelines don't support multiple awaiters on flush.
-            // While it's acceptable to call PipeWriter.FlushAsync again before the last FlushAsync completes,
-            // it is not acceptable to attach a new continuation (via await, AsTask(), etc..). In this case,
-            // we find previous flush Task which still accounts for any newly committed bytes and await that.
-            lock (_flushLock)
-            {
-                if (_lastFlushTask == null || _lastFlushTask.IsCompleted)
-                {
-                    _lastFlushTask = flushValueTask.AsTask();
-                }
-
-                return TimeFlushAsync(minRate, count, outputAborter, cancellationToken);
-            }
+            return TimeFlushAsyncAwaited(pipeFlushTask, minRate, count, outputAborter, cancellationToken);
         }
 
-        private async ValueTask<FlushResult> TimeFlushAsync(MinDataRate minRate, long count, IHttpOutputAborter outputAborter, CancellationToken cancellationToken)
+        private async ValueTask<FlushResult> TimeFlushAsyncAwaited(ValueTask<FlushResult> pipeFlushTask, MinDataRate minRate, long count, IHttpOutputAborter outputAborter, CancellationToken cancellationToken)
         {
             if (minRate != null)
             {
@@ -88,7 +100,7 @@ private async ValueTask<FlushResult> TimeFlushAsync(MinDataRate minRate, long co
 
             try
             {
-                return await _lastFlushTask;
+                return await pipeFlushTask;
             }
             catch (OperationCanceledException ex) when (outputAborter != null)
             {
@@ -111,5 +123,11 @@ private async ValueTask<FlushResult> TimeFlushAsync(MinDataRate minRate, long co
 
             return default;
         }
+
+        private async ValueTask<FlushResult> AwaitLastFlushAndTimeFlushAsync(ValueTask<FlushResult> lastFlushTask, MinDataRate minRate, long count, IHttpOutputAborter outputAborter, CancellationToken cancellationToken)
+        {
+            await lastFlushTask;
+            return await TimeFlushAsync(minRate, count, outputAborter, cancellationToken);
+        }
     }
 }

src/Servers/Kestrel/Core/test/Http1ConnectionTests.cs

@@ -42,7 +42,7 @@ public class Http1ConnectionTests : IDisposable
 
         public Http1ConnectionTests()
         {
-            _pipelineFactory = MemoryPoolFactory.Create();
+            _pipelineFactory = SlabMemoryPoolFactory.Create();
             var options = new PipeOptions(_pipelineFactory, readerScheduler: PipeScheduler.Inline, writerScheduler: PipeScheduler.Inline, useSynchronizationContext: false);
             var pair = DuplexPipe.CreateConnectionPair(options, options);
 

src/Servers/Kestrel/Core/test/HttpResponseHeadersTests.cs

@@ -21,7 +21,7 @@ public class HttpResponseHeadersTests
         [Fact]
         public void InitialDictionaryIsEmpty()
         {
-            using (var memoryPool = MemoryPoolFactory.Create())
+            using (var memoryPool = SlabMemoryPoolFactory.Create())
             {
                 var options = new PipeOptions(memoryPool, readerScheduler: PipeScheduler.Inline, writerScheduler: PipeScheduler.Inline, useSynchronizationContext: false);
                 var pair = DuplexPipe.CreateConnectionPair(options, options);

src/Servers/Kestrel/Core/test/OutputProducerTests.cs

@@ -21,7 +21,7 @@ public class OutputProducerTests : IDisposable
 
         public OutputProducerTests()
         {
-            _memoryPool = MemoryPoolFactory.Create();
+            _memoryPool = SlabMemoryPoolFactory.Create();
         }
 
         public void Dispose()

src/Servers/Kestrel/Core/test/PipelineExtensionTests.cs

@@ -16,7 +16,7 @@ public class PipelineExtensionTests : IDisposable
         private const int _ulongMaxValueLength = 20;
 
         private readonly Pipe _pipe;
-        private readonly MemoryPool<byte> _memoryPool = MemoryPoolFactory.Create();
+        private readonly MemoryPool<byte> _memoryPool = SlabMemoryPoolFactory.Create();
 
         public PipelineExtensionTests()
         {

src/Servers/Kestrel/Core/test/StartLineTests.cs

@@ -499,7 +499,7 @@ public void AuthorityForms(string rawTarget, string path, string query)
 
         public StartLineTests()
         {
-            MemoryPool = MemoryPoolFactory.Create();
+            MemoryPool = SlabMemoryPoolFactory.Create();
             var options = new PipeOptions(MemoryPool, readerScheduler: PipeScheduler.Inline, writerScheduler: PipeScheduler.Inline, useSynchronizationContext: false);
             var pair = DuplexPipe.CreateConnectionPair(options, options);
             Transport = pair.Transport;

src/Servers/Kestrel/Core/test/TestHelpers/TestInput.cs

@@ -23,7 +23,7 @@ class TestInput : IDisposable
 
         public TestInput()
         {
-            _memoryPool = MemoryPoolFactory.Create();
+            _memoryPool = SlabMemoryPoolFactory.Create();
             var options = new PipeOptions(pool: _memoryPool, readerScheduler: PipeScheduler.Inline, writerScheduler: PipeScheduler.Inline, useSynchronizationContext: false);
             var pair = DuplexPipe.CreateConnectionPair(options, options);
             Transport = pair.Transport;

src/Servers/Kestrel/Core/test/TimingPipeFlusherTests.cs

@@ -0,0 +1,67 @@
+// 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.IO.Pipelines;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
+using Moq;
+using Xunit;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
+{
+    public class TimingPipeFlusherTests
+    {
+        [Fact]
+        public async Task IfFlushIsCalledAgainBeforeTheLastFlushCompletedItWaitsForTheLastCall()
+        {
+            var mockPipeWriter = new Mock<PipeWriter>();
+            var pipeWriterFlushTcsArray = new[] {
+                new TaskCompletionSource<FlushResult>(TaskCreationOptions.RunContinuationsAsynchronously),
+                new TaskCompletionSource<FlushResult>(TaskCreationOptions.RunContinuationsAsynchronously),
+                new TaskCompletionSource<FlushResult>(TaskCreationOptions.RunContinuationsAsynchronously),
+            };
+            var pipeWriterFlushCallCount = 0;
+
+            mockPipeWriter.Setup(p => p.FlushAsync(CancellationToken.None)).Returns(() =>
+            {
+                return new ValueTask<FlushResult>(pipeWriterFlushTcsArray[pipeWriterFlushCallCount++].Task);
+            });
+
+            var timingPipeFlusher = new TimingPipeFlusher(mockPipeWriter.Object, null, null);
+
+            var flushTask0 = timingPipeFlusher.FlushAsync();
+            var flushTask1 = timingPipeFlusher.FlushAsync();
+            var flushTask2 = timingPipeFlusher.FlushAsync();
+
+            Assert.False(flushTask0.IsCompleted);
+            Assert.False(flushTask1.IsCompleted);
+            Assert.False(flushTask2.IsCompleted);
+            Assert.Equal(1, pipeWriterFlushCallCount);
+
+            pipeWriterFlushTcsArray[0].SetResult(default);
+            await flushTask0.AsTask().DefaultTimeout();
+
+            Assert.True(flushTask0.IsCompleted);
+            Assert.False(flushTask1.IsCompleted);
+            Assert.False(flushTask2.IsCompleted);
+            Assert.True(pipeWriterFlushCallCount <= 2);
+
+            pipeWriterFlushTcsArray[1].SetResult(default);
+            await flushTask1.AsTask().DefaultTimeout();
+
+            Assert.True(flushTask0.IsCompleted);
+            Assert.True(flushTask1.IsCompleted);
+            Assert.False(flushTask2.IsCompleted);
+            Assert.True(pipeWriterFlushCallCount <= 3);
+
+            pipeWriterFlushTcsArray[2].SetResult(default);
+            await flushTask2.AsTask().DefaultTimeout();
+
+            Assert.True(flushTask0.IsCompleted);
+            Assert.True(flushTask1.IsCompleted);
+            Assert.True(flushTask2.IsCompleted);
+            Assert.Equal(3, pipeWriterFlushCallCount);
+        }
+    }
+}

src/Servers/Kestrel/Transport.Libuv/src/LibuvTransportOptions.cs

@@ -31,7 +31,7 @@ public class LibuvTransportOptions
 
         public long? MaxWriteBufferSize { get; set; } = 64 * 1024;
 
-        internal Func<MemoryPool<byte>> MemoryPoolFactory { get; set; } = System.Buffers.MemoryPoolFactory.Create;
+        internal Func<MemoryPool<byte>> MemoryPoolFactory { get; set; } = System.Buffers.SlabMemoryPoolFactory.Create;
 
         private static int ProcessorThreadCount
         {

src/Servers/Kestrel/Transport.Libuv/test/LibuvOutputConsumerTests.cs

@@ -41,7 +41,7 @@ public class LibuvOutputConsumerTests : IDisposable
 
         public LibuvOutputConsumerTests()
         {
-            _memoryPool = MemoryPoolFactory.Create();
+            _memoryPool = SlabMemoryPoolFactory.Create();
             _mockLibuv = new MockLibuv();
 
             var context = new TestLibuvTransportContext();
@@ -560,15 +560,15 @@ await Task.Run(async () =>
                     Assert.False(task1Waits.IsFaulted);
 
                     // following tasks should wait.
-                    var task3Canceled = outputProducer.WriteDataAsync(fullBuffer, cancellationToken: abortedSource.Token);
+                    var task2Canceled = outputProducer.WriteDataAsync(fullBuffer, cancellationToken: abortedSource.Token);
 
                     // Give time for tasks to percolate
                     await _mockLibuv.OnPostTask;
 
-                    // Third task is not completed
-                    Assert.False(task3Canceled.IsCompleted);
-                    Assert.False(task3Canceled.IsCanceled);
-                    Assert.False(task3Canceled.IsFaulted);
+                    // Second task is not completed
+                    Assert.False(task2Canceled.IsCompleted);
+                    Assert.False(task2Canceled.IsCanceled);
+                    Assert.False(task2Canceled.IsFaulted);
 
                     abortedSource.Cancel();
 
@@ -578,29 +578,29 @@ await Task.Run(async () =>
                         await _libuvThread.PostAsync(cb => cb(0), triggerNextCompleted);
                     }
 
-                    // First task is  completed
+                    // First task is completed
                     Assert.True(task1Waits.IsCompleted);
                     Assert.False(task1Waits.IsCanceled);
                     Assert.False(task1Waits.IsFaulted);
 
-                    // A final write guarantees that the error is observed by OutputProducer,
-                    // but doesn't return a canceled/faulted task.
-                    var task4Success = outputProducer.WriteDataAsync(fullBuffer);
-                    Assert.True(task4Success.IsCompleted);
-                    Assert.False(task4Success.IsCanceled);
-                    Assert.False(task4Success.IsFaulted);
+                    // Second task is now canceled
+                    await Assert.ThrowsAsync<OperationCanceledException>(() => task2Canceled);
+                    Assert.True(task2Canceled.IsCanceled);
 
-                    // Third task is now canceled
-                    await Assert.ThrowsAsync<OperationCanceledException>(() => task3Canceled);
-                    Assert.True(task3Canceled.IsCanceled);
+                    // A final write can still succeed.
+                    var task3Success = outputProducer.WriteDataAsync(fullBuffer);
 
                     await _mockLibuv.OnPostTask;
 
-                    // Complete the 4th write
+                    // Complete the 3rd write
                     while (completeQueue.TryDequeue(out var triggerNextCompleted))
                     {
                         await _libuvThread.PostAsync(cb => cb(0), triggerNextCompleted);
                     }
+
+                    Assert.True(task3Success.IsCompleted);
+                    Assert.False(task3Success.IsCanceled);
+                    Assert.False(task3Success.IsFaulted);;
                 }
             });
         }