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;
 
@@ -220,7 +221,7 @@ public override async ValueTask<ReadResult> ReadAsync(CancellationToken cancella
                     if (UseZeroByteReads && _bufferedBytes == 0)
                     {
                         // Wait for data by doing 0 byte read before
-                        await InnerStream.ReadAsync(Memory<byte>.Empty, cancellationToken).ConfigureAwait(false);
+                        await InnerStream.ReadAsync(Memory<byte>.Empty, tokenSource.Token).ConfigureAwait(false);
                     }
 
                     AllocateReadTail();
@@ -259,6 +260,83 @@ public override async ValueTask<ReadResult> ReadAsync(CancellationToken cancella
             }
         }
 
+        protected override async ValueTask<ReadResult> ReadAtLeastAsyncCore(int minimumSize, CancellationToken cancellationToken)
+        {
+            // TODO ReadyAsync needs to throw if there are overlapping reads.
+            ThrowIfCompleted();
+
+            // 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)
+                {
+                    return readResult;
+                }
+            }
+
+            if (_isStreamCompleted)
+            {
+                return new ReadResult(buffer: default, isCanceled: false, isCompleted: true);
+            }
+
+            CancellationTokenRegistration reg = default;
+            if (cancellationToken.CanBeCanceled)
+            {
+                reg = cancellationToken.UnsafeRegister(state => ((StreamPipeReader)state!).Cancel(), this);
+            }
+
+            using (reg)
+            {
+                var isCanceled = false;
+                try
+                {
+                    while (_bufferedBytes < minimumSize && !_isStreamCompleted)
+                    {
+                        // This optimization only makes sense if we don't have anything buffered
+                        if (UseZeroByteReads && _bufferedBytes == 0)
+                        {
+                            // Wait for data by doing 0 byte read before
+                            await InnerStream.ReadAsync(Memory<byte>.Empty, tokenSource.Token).ConfigureAwait(false);
+                        }
+
+                        AllocateReadTail(minimumSize);
+
+                        Memory<byte> buffer = _readTail!.AvailableMemory.Slice(_readTail.End);
+
+                        int length = await InnerStream.ReadAsync(buffer, tokenSource.Token).ConfigureAwait(false);
+
+                        Debug.Assert(length + _readTail.End <= _readTail.AvailableMemory.Length);
+
+                        _readTail.End += length;
+                        _bufferedBytes += length;
+
+                        if (length == 0)
+                        {
+                            _isStreamCompleted = true;
+                        }
+                    }
+                }
+                catch (OperationCanceledException)
+                {
+                    ClearCancellationToken();
+
+                    if (tokenSource.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
+                    {
+                        // Catch cancellation and translate it into setting isCanceled = true
+                        isCanceled = true;
+                    }
+                    else
+                    {
+                        throw;
+                    }
+
+                }
+
+                return new ReadResult(GetCurrentReadOnlySequence(), isCanceled, _isStreamCompleted);
+            }
+        }
+
         /// <inheritdoc />
         public override async Task CopyToAsync(PipeWriter destination, CancellationToken cancellationToken = default)
         {
@@ -434,42 +512,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;
+        private 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/StreamPipeReaderTests.cs

@@ -31,18 +31,28 @@ public async Task CanRead()
             reader.Complete();
         }
 
-        [Fact]
-        public async Task CanReadAtLeast()
+        [Theory]
+        [InlineData(-1, false)]
+        [InlineData(-1, true)]
+        [InlineData(5, false)]
+        [InlineData(5, true)]
+        public async Task CanReadAtLeast(int bufferSize, bool bufferedRead)
         {
-            var stream = new MemoryStream(Encoding.ASCII.GetBytes("Hello World"));
-            var reader = PipeReader.Create(stream);
+            var stream = new MemoryStream(Encoding.ASCII.GetBytes("Hello Pipelines World"));
+            var reader = PipeReader.Create(stream, new StreamPipeReaderOptions(bufferSize: bufferSize));
+
+            if (bufferedRead)
+            {
+                ReadResult bufferedReadResult = await reader.ReadAsync();
+                Assert.NotEqual(0, bufferedReadResult.Buffer.Length);
+            }
 
-            ReadResult readResult = await reader.ReadAtLeastAsync(10);
+            ReadResult readResult = await reader.ReadAtLeastAsync(20);
             ReadOnlySequence<byte> buffer = readResult.Buffer;
 
-            Assert.Equal(11, buffer.Length);
-            Assert.True(buffer.IsSingleSegment);
-            Assert.Equal("Hello World", Encoding.ASCII.GetString(buffer.ToArray()));
+            Assert.Equal(21, buffer.Length);
+            Assert.Equal(bufferSize == -1 || !bufferedRead, buffer.IsSingleSegment);
+            Assert.Equal("Hello Pipelines World", Encoding.ASCII.GetString(buffer.ToArray()));
 
             reader.AdvanceTo(buffer.End);
             reader.Complete();