Respect IBufferWriter leases after commit

src/Servers/Kestrel/Core/test/BufferWriterTests.cs → src/Servers/Kestrel/Core/test/BufferWriterWithPipeWriterTests.cs

@@ -1,17 +1,14 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Buffers;
-using System.Collections.Generic;
-using System.Linq;
-using Xunit;
 
 namespace System.IO.Pipelines.Tests;
 
-public class BufferWriterTests : IDisposable
+public class BufferWriterWithPipeWriterTests : IDisposable
 {
     protected Pipe Pipe;
-    public BufferWriterTests()
+    public BufferWriterWithPipeWriterTests()
     {
         Pipe = new Pipe(new PipeOptions(useSynchronizationContext: false, pauseWriterThreshold: 0, resumeWriterThreshold: 0));
     }
@@ -107,7 +104,7 @@ public void CanWriteIntoHeadlessBuffer()
     {
         BufferWriter<PipeWriter> writer = new BufferWriter<PipeWriter>(Pipe.Writer);
 
-        writer.Write(new byte[] { 1, 2, 3 });
+        writer.Write([1, 2, 3]);
         writer.Commit();
 
         Assert.Equal(3, writer.BytesCommitted);
@@ -119,9 +116,9 @@ public void CanWriteMultipleTimes()
     {
         BufferWriter<PipeWriter> writer = new BufferWriter<PipeWriter>(Pipe.Writer);
 
-        writer.Write(new byte[] { 1 });
-        writer.Write(new byte[] { 2 });
-        writer.Write(new byte[] { 3 });
+        writer.Write([1]);
+        writer.Write([2]);
+        writer.Write([3]);
         writer.Commit();
 
         Assert.Equal(3, writer.BytesCommitted);
@@ -151,7 +148,7 @@ public void EnsureAllocatesSpan()
         writer.Ensure(10);
         Assert.True(writer.Span.Length > 10);
         Assert.Equal(0, writer.BytesCommitted);
-        Assert.Equal(new byte[] { }, Read());
+        Assert.Equal([], Read());
     }
 
     [Fact]
@@ -160,7 +157,7 @@ public void ExposesSpan()
         int initialLength = Pipe.Writer.GetMemory().Length;
         BufferWriter<PipeWriter> writer = new BufferWriter<PipeWriter>(Pipe.Writer);
         Assert.Equal(initialLength, writer.Span.Length);
-        Assert.Equal(new byte[] { }, Read());
+        Assert.Equal([], Read());
     }
 
     [Fact]
@@ -170,12 +167,13 @@ public void SlicesSpanAndAdvancesAfterWrite()
 
         BufferWriter<PipeWriter> writer = new BufferWriter<PipeWriter>(Pipe.Writer);
 
-        writer.Write(new byte[] { 1, 2, 3 });
+        writer.Write([1, 2, 3]);
+        Assert.Equal(initialLength - 3, writer.Span.Length);
+
         writer.Commit();
 
         Assert.Equal(3, writer.BytesCommitted);
-        Assert.Equal(initialLength - 3, writer.Span.Length);
-        Assert.Equal(Pipe.Writer.GetMemory().Length, writer.Span.Length);
+        Assert.True(writer.Span.IsEmpty);
         Assert.Equal(new byte[] { 1, 2, 3 }, Read());
     }
 
@@ -184,7 +182,7 @@ public void BufferWriterCountsBytesCommitted()
     {
         BufferWriter<PipeWriter> writer = new BufferWriter<PipeWriter>(Pipe.Writer);
 
-        writer.Write(new byte[] { 1, 2, 3 });
+        writer.Write([1, 2, 3]);
         Assert.Equal(0, writer.BytesCommitted);
 
         writer.Commit();

src/Shared/ServerInfrastructure/BufferWriter.cs

@@ -69,6 +69,7 @@ public void Commit()
             _bytesCommitted += buffered;
             _buffered = 0;
             _output.Advance(buffered);
+            _span = default;
         }
     }
 
@@ -90,6 +91,11 @@ public void Advance(int count)
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public void Write(ReadOnlySpan<byte> source)
     {
+        if (_span.IsEmpty && !source.IsEmpty)
+        {
+            EnsureMore();
+        }
+
         if (_span.Length >= source.Length)
         {
             source.CopyTo(_span);

src/Shared/test/Shared.Tests/ServerInfrastructure/BufferWriterTests.cs

@@ -0,0 +1,65 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Buffers;
+
+namespace Microsoft.AspNetCore.Shared.Tests.ServerInfrastructure;
+
+public class BufferWriterTests
+{
+    [Fact]
+    public void WriteAfterCommitAcquiresNewBuffer()
+    {
+        var output = new StrictBufferWriter();
+        var writer = new BufferWriter<StrictBufferWriter>(output);
+
+        writer.Write([1, 2, 3]);
+        writer.Commit();
+        writer.Write([4, 5]);
+        writer.Commit();
+
+        Assert.Equal(2, output.BufferAcquisitions);
+        Assert.Collection(
+            output.CommittedBuffers,
+            buffer => Assert.Equal([ 1, 2, 3 ], buffer),
+            buffer => Assert.Equal([ 4, 5 ], buffer));
+    }
+
+    private sealed class StrictBufferWriter : IBufferWriter<byte>
+    {
+        private byte[]? _currentLease;
+
+        public int BufferAcquisitions { get; private set; }
+
+        public List<byte[]> CommittedBuffers { get; } = [];
+
+        public void Advance(int count)
+        {
+            if (_currentLease is null)
+            {
+                throw new InvalidOperationException("A new buffer must be acquired before advancing.");
+            }
+
+            if ((uint)count > (uint)_currentLease.Length)
+            {
+                throw new ArgumentOutOfRangeException(nameof(count));
+            }
+
+            CommittedBuffers.Add(_currentLease[..count]);
+            _currentLease = null;
+        }
+
+        public Memory<byte> GetMemory(int sizeHint = 0) => AcquireBuffer(sizeHint);
+
+        public Span<byte> GetSpan(int sizeHint = 0) => AcquireBuffer(sizeHint);
+
+        private byte[] AcquireBuffer(int sizeHint)
+        {
+            ArgumentOutOfRangeException.ThrowIfNegative(sizeHint);
+
+            _currentLease = new byte[Math.Max(sizeHint, 16)];
+            BufferAcquisitions++;
+            return _currentLease;
+        }
+    }
+}