diff --git a/src/Servers/Kestrel/Core/test/BufferWriterTests.cs b/src/Servers/Kestrel/Core/test/BufferWriterWithPipeWriterTests.cs similarity index 90% rename from src/Servers/Kestrel/Core/test/BufferWriterTests.cs rename to src/Servers/Kestrel/Core/test/BufferWriterWithPipeWriterTests.cs index 255c10b3c75c..d9033000e914 100644 --- a/src/Servers/Kestrel/Core/test/BufferWriterTests.cs +++ b/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 writer = new BufferWriter(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 writer = new BufferWriter(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 writer = new BufferWriter(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 writer = new BufferWriter(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 writer = new BufferWriter(Pipe.Writer); - writer.Write(new byte[] { 1, 2, 3 }); + writer.Write([1, 2, 3]); Assert.Equal(0, writer.BytesCommitted); writer.Commit(); diff --git a/src/Shared/ServerInfrastructure/BufferWriter.cs b/src/Shared/ServerInfrastructure/BufferWriter.cs index 828194949306..21a026d00553 100644 --- a/src/Shared/ServerInfrastructure/BufferWriter.cs +++ b/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 source) { + if (_span.IsEmpty && !source.IsEmpty) + { + EnsureMore(); + } + if (_span.Length >= source.Length) { source.CopyTo(_span); diff --git a/src/Shared/test/Shared.Tests/ServerInfrastructure/BufferWriterTests.cs b/src/Shared/test/Shared.Tests/ServerInfrastructure/BufferWriterTests.cs new file mode 100644 index 000000000000..19e3af27c7a0 --- /dev/null +++ b/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(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 + { + private byte[]? _currentLease; + + public int BufferAcquisitions { get; private set; } + + public List 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 GetMemory(int sizeHint = 0) => AcquireBuffer(sizeHint); + + public Span GetSpan(int sizeHint = 0) => AcquireBuffer(sizeHint); + + private byte[] AcquireBuffer(int sizeHint) + { + ArgumentOutOfRangeException.ThrowIfNegative(sizeHint); + + _currentLease = new byte[Math.Max(sizeHint, 16)]; + BufferAcquisitions++; + return _currentLease; + } + } +}