Add ReadAtLeastAsync implementation for StreamPipeReader

src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeReader.cs

@@ -52,6 +52,7 @@ public StreamPipeReader(Stream readingStream, StreamPipeReaderOptions options)
         private bool LeaveOpen => _options.LeaveOpen;
         private bool UseZeroByteReads => _options.UseZeroByteReads;
         private int BufferSize => _options.BufferSize;
+        private int MaxBufferSize => _options.MaxBufferSize;
         private int MinimumReadThreshold => _options.MinimumReadSize;
         private MemoryPool<byte> Pool => _options.Pool;
 
@@ -188,74 +189,168 @@ public override void Complete(Exception? exception = null)
         }
 
         /// <inheritdoc />
-        public override async ValueTask<ReadResult> ReadAsync(CancellationToken cancellationToken = default)
+        public override ValueTask<ReadResult> ReadAsync(CancellationToken cancellationToken = default)
         {
             // TODO ReadyAsync needs to throw if there are overlapping reads.
             ThrowIfCompleted();
 
+            cancellationToken.ThrowIfCancellationRequested();
+
             // PERF: store InternalTokenSource locally to avoid querying it twice (which acquires a lock)
             CancellationTokenSource tokenSource = InternalTokenSource;
             if (TryReadInternal(tokenSource, out ReadResult readResult))
             {
-                return readResult;
+                return new ValueTask<ReadResult>(readResult);
             }
 
             if (_isStreamCompleted)
             {
-                return new ReadResult(buffer: default, isCanceled: false, isCompleted: true);
+                ReadResult completedResult = new ReadResult(buffer: default, isCanceled: false, isCompleted: true);
+                return new ValueTask<ReadResult>(completedResult);
             }
 
-            CancellationTokenRegistration reg = default;
-            if (cancellationToken.CanBeCanceled)
-            {
-                reg = cancellationToken.UnsafeRegister(state => ((StreamPipeReader)state!).Cancel(), this);
-            }
+            return Core(this, tokenSource, cancellationToken);
 
-            using (reg)
+            static async ValueTask<ReadResult> Core(StreamPipeReader reader, CancellationTokenSource tokenSource, CancellationToken cancellationToken)
             {
-                var isCanceled = false;
-                try
+                CancellationTokenRegistration reg = default;
+                if (cancellationToken.CanBeCanceled)
                 {
-                    // This optimization only makes sense if we don't have anything buffered
-                    if (UseZeroByteReads && _bufferedBytes == 0)
+                    reg = cancellationToken.UnsafeRegister(state => ((StreamPipeReader)state!).Cancel(), reader);
+                }
+
+                using (reg)
+                {
+                    var isCanceled = false;
+                    try
                     {
-                        // Wait for data by doing 0 byte read before
-                        await InnerStream.ReadAsync(Memory<byte>.Empty, cancellationToken).ConfigureAwait(false);
-                    }
+                        // This optimization only makes sense if we don't have anything buffered
+                        if (reader.UseZeroByteReads && reader._bufferedBytes == 0)
+                        {
+                            // Wait for data by doing 0 byte read before
+                            await reader.InnerStream.ReadAsync(Memory<byte>.Empty, tokenSource.Token).ConfigureAwait(false);
+                        }
 
-                    AllocateReadTail();
+                        reader.AllocateReadTail();
 
-                    Memory<byte> buffer = _readTail!.AvailableMemory.Slice(_readTail.End);
+                        Memory<byte> buffer = reader._readTail!.AvailableMemory.Slice(reader._readTail.End);
 
-                    int length = await InnerStream.ReadAsync(buffer, tokenSource.Token).ConfigureAwait(false);
+                        int length = await reader.InnerStream.ReadAsync(buffer, tokenSource.Token).ConfigureAwait(false);
 
-                    Debug.Assert(length + _readTail.End <= _readTail.AvailableMemory.Length);
+                        Debug.Assert(length + reader._readTail.End <= reader._readTail.AvailableMemory.Length);
 
-                    _readTail.End += length;
-                    _bufferedBytes += length;
+                        reader._readTail.End += length;
+                        reader._bufferedBytes += length;
 
-                    if (length == 0)
+                        if (length == 0)
+                        {
+                            reader._isStreamCompleted = true;
+                        }
+                    }
+                    catch (OperationCanceledException)
                     {
-                        _isStreamCompleted = true;
+                        reader.ClearCancellationToken();
+
+                        if (tokenSource.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
+                        {
+                            // Catch cancellation and translate it into setting isCanceled = true
+                            isCanceled = true;
+                        }
+                        else
+                        {
+                            throw;
+                        }
+
                     }
+
+                    return new ReadResult(reader.GetCurrentReadOnlySequence(), isCanceled, reader._isStreamCompleted);
                 }
-                catch (OperationCanceledException)
+            }
+        }
+
+        protected override ValueTask<ReadResult> ReadAtLeastAsyncCore(int minimumSize, CancellationToken cancellationToken)
+        {
+            // TODO ReadyAsync needs to throw if there are overlapping reads.
+            ThrowIfCompleted();
+
+            cancellationToken.ThrowIfCancellationRequested();
+
+            // PERF: store InternalTokenSource locally to avoid querying it twice (which acquires a lock)
+            CancellationTokenSource tokenSource = InternalTokenSource;
+            if (TryReadInternal(tokenSource, out ReadResult readResult))
+            {
+                if (readResult.Buffer.Length >= minimumSize || readResult.IsCompleted || readResult.IsCanceled)
                 {
-                    ClearCancellationToken();
+                    return new ValueTask<ReadResult>(readResult);
+                }
+            }
+
+            if (_isStreamCompleted)
+            {
+                ReadResult completedResult = new ReadResult(buffer: default, isCanceled: false, isCompleted: true);
+                return new ValueTask<ReadResult>(completedResult);
+            }
 
-                    if (tokenSource.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
+            return Core(this, minimumSize, tokenSource, cancellationToken);
+
+            static async ValueTask<ReadResult> Core(StreamPipeReader reader, int minimumSize, CancellationTokenSource tokenSource, CancellationToken cancellationToken)
+            {
+                CancellationTokenRegistration reg = default;
+                if (cancellationToken.CanBeCanceled)
+                {
+                    reg = cancellationToken.UnsafeRegister(state => ((StreamPipeReader)state!).Cancel(), reader);
+                }
+
+                using (reg)
+                {
+                    var isCanceled = false;
+                    try
                     {
-                        // Catch cancellation and translate it into setting isCanceled = true
-                        isCanceled = true;
+                        // This optimization only makes sense if we don't have anything buffered
+                        if (reader.UseZeroByteReads && reader._bufferedBytes == 0)
+                        {
+                            // Wait for data by doing 0 byte read before
+                            await reader.InnerStream.ReadAsync(Memory<byte>.Empty, tokenSource.Token).ConfigureAwait(false);
+                        }
+
+                        do
+                        {
+                            reader.AllocateReadTail(minimumSize);
+
+                            Memory<byte> buffer = reader._readTail!.AvailableMemory.Slice(reader._readTail.End);
+
+                            int length = await reader.InnerStream.ReadAsync(buffer, tokenSource.Token).ConfigureAwait(false);
+
+                            Debug.Assert(length + reader._readTail.End <= reader._readTail.AvailableMemory.Length);
+
+                            reader._readTail.End += length;
+                            reader._bufferedBytes += length;
+
+                            if (length == 0)
+                            {
+                                reader._isStreamCompleted = true;
+                                break;
+                            }
+                        } while (reader._bufferedBytes < minimumSize);
                     }
-                    else
+                    catch (OperationCanceledException)
                     {
-                        throw;
+                        reader.ClearCancellationToken();
+
+                        if (tokenSource.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
+                        {
+                            // Catch cancellation and translate it into setting isCanceled = true
+                            isCanceled = true;
+                        }
+                        else
+                        {
+                            throw;
+                        }
+
                     }
 
+                    return new ReadResult(reader.GetCurrentReadOnlySequence(), isCanceled, reader._isStreamCompleted);
                 }
-
-                return new ReadResult(GetCurrentReadOnlySequence(), isCanceled, _isStreamCompleted);
             }
         }
 
@@ -434,42 +529,58 @@ private ReadOnlySequence<byte> GetCurrentReadOnlySequence()
             return new ReadOnlySequence<byte>(_readHead, _readIndex, _readTail, _readTail.End);
         }
 
-        private void AllocateReadTail()
+        private void AllocateReadTail(int? minimumSize = null)
         {
             if (_readHead == null)
             {
                 Debug.Assert(_readTail == null);
-                _readHead = AllocateSegment();
+                _readHead = AllocateSegment(minimumSize);
                 _readTail = _readHead;
             }
             else
             {
                 Debug.Assert(_readTail != null);
                 if (_readTail.WritableBytes < MinimumReadThreshold)
                 {
-                    BufferSegment nextSegment = AllocateSegment();
+                    BufferSegment nextSegment = AllocateSegment(minimumSize);
                     _readTail.SetNext(nextSegment);
                     _readTail = nextSegment;
                 }
             }
         }
 
-        private BufferSegment AllocateSegment()
+        private BufferSegment AllocateSegment(int? minimumSize = null)
         {
             BufferSegment nextSegment = CreateSegmentUnsynchronized();
 
-            if (_options.IsDefaultSharedMemoryPool)
+            var bufferSize = minimumSize ?? BufferSize;
+            int maxSize = !_options.IsDefaultSharedMemoryPool ? _options.Pool.MaxBufferSize : -1;
+
+            if (bufferSize <= maxSize)
             {
-                nextSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(BufferSize));
+                // Use the specified pool as it fits.
+                int sizeToRequest = GetSegmentSize(bufferSize, maxSize);
+                nextSegment.SetOwnedMemory(_options.Pool.Rent(sizeToRequest));
             }
             else
             {
-                nextSegment.SetOwnedMemory(Pool.Rent(BufferSize));
+                // Use the array pool
+                int sizeToRequest = GetSegmentSize(bufferSize, MaxBufferSize);
+                nextSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(sizeToRequest));
             }
 
             return nextSegment;
         }
 
+        private int GetSegmentSize(int sizeHint, int maxBufferSize)
+        {
+            // First we need to handle case where hint is smaller than minimum segment size
+            sizeHint = Math.Max(BufferSize, sizeHint);
+            // After that adjust it to fit into pools max buffer size
+            int adjustedToMaximumSize = Math.Min(maxBufferSize, sizeHint);
+            return adjustedToMaximumSize;
+        }
+
         private BufferSegment CreateSegmentUnsynchronized()
         {
             if (_bufferSegmentPool.TryPop(out BufferSegment? segment))

src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeReaderOptions.cs

@@ -9,6 +9,7 @@ namespace System.IO.Pipelines
     public class StreamPipeReaderOptions
     {
         private const int DefaultBufferSize = 4096;
+        internal const int DefaultMaxBufferSize = 2048 * 1024;
         private const int DefaultMinimumReadSize = 1024;
 
         internal static readonly StreamPipeReaderOptions s_default = new StreamPipeReaderOptions();
@@ -55,6 +56,10 @@ public StreamPipeReaderOptions(MemoryPool<byte>? pool = null, int bufferSize = -
         /// <value>The buffer size.</value>
         public int BufferSize { get; }
 
+        /// <summary>Gets the maximum buffer size to use when renting memory from the <see cref="System.IO.Pipelines.StreamPipeReaderOptions.Pool" />.</summary>
+        /// <value>The maximum buffer size.</value>
+        internal int MaxBufferSize { get; } = DefaultMaxBufferSize;
+
         /// <summary>Gets the threshold of remaining bytes in the buffer before a new buffer is allocated.</summary>
         /// <value>The minimum read size.</value>
         public int MinimumReadSize { get; }

src/libraries/System.IO.Pipelines/tests/BasePipeReaderReadAtLeastAsyncTests.cs

@@ -0,0 +1,11 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace System.IO.Pipelines.Tests
+{
+    public class BasePipeReaderReadAtLeastAsyncTests : ReadAtLeastAsyncTests
+    {
+        private PipeReader? _pipeReader;
+        protected override PipeReader PipeReader => _pipeReader ?? (_pipeReader = new BasePipeReader(Pipe.Reader));
+    }
+}

src/libraries/System.IO.Pipelines/tests/Infrastructure/TestMemoryPool.cs

@@ -26,7 +26,8 @@ protected override void Dispose(bool disposing)
             _disposed = true;
         }
 
-        public override int MaxBufferSize => 4096;
+        internal const int DefaultMaxBufferSize = 4096;
+        public override int MaxBufferSize => DefaultMaxBufferSize;
 
         internal void CheckDisposed()
         {