Skip to content

Commit

Permalink
Add ReadAtLeastAsync implementation for StreamPipeReader (#52246)
Browse files Browse the repository at this point in the history 
-  Increase segment size based on minimumSize
- Use MaxBufferSize from custom memory pool
- Align segment allocation logic on Pipe one
- Group ReadAtLeastAsync tests
  • Loading branch information
@manandre
manandre committed May 10, 2021
1 parent 7065bba commit 2ab0058
Show file tree
Hide file tree
Showing 10 changed files with 390 additions and 185 deletions.
191 changes: 151 additions & 40 deletions src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeReader.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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);
}
}

Expand Down Expand Up @@ -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))
Expand Down
5 changes: 5 additions & 0 deletions src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeReaderOptions.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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; }
Expand Down
11 changes: 11 additions & 0 deletions src/libraries/System.IO.Pipelines/tests/BasePipeReaderReadAtLeastAsyncTests.cs
View file
Original file line number Diff line number Diff line change
@@ -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));
}
}
3 changes: 2 additions & 1 deletion src/libraries/System.IO.Pipelines/tests/Infrastructure/TestMemoryPool.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()
{
Expand Down
Loading

0 comments on commit 2ab0058

Please sign in to comment.