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,84 @@ 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
+ {
+ // 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);
+ }
+
+ while (_bufferedBytes < minimumSize)
+ {
+ 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;
+ break;
+ }
+ }
+ }
+ 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 +513,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();