Skip to content

Avoid empty segments from PipeReader.ReadAsync #72536

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jul 21, 2022
Merged

Avoid empty segments from PipeReader.ReadAsync #72536

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f0be4e6
Avoid empty segments from PipeReader.ReadAsync
BrennanConroy Jul 20, 2022
6c0d207
length :(
BrennanConroy Jul 20, 2022
5e13f99
fb
BrennanConroy Jul 20, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,17 @@ public void SetOwnedMemory(byte[] arrayPoolBuffer)
AvailableMemory = arrayPoolBuffer;
}

// Resets memory and internal state, should be called when removing the segment from the linked list
public void Reset()
{
ResetMemory();

Next = null;
RunningIndex = 0;
_next = null;
}

// Resets memory only, should be called when keeping the BufferSegment in the linked list and only swapping out the memory
public void ResetMemory()
{
IMemoryOwner<byte>? memoryOwner = _memoryOwner;
Expand All @@ -74,10 +85,8 @@ public void ResetMemory()
_array = null;
}

Next = null;
RunningIndex = 0;

Memory = default;
_next = null;
_end = 0;
AvailableMemory = default;
}
Expand Down
41 changes: 30 additions & 11 deletions src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,20 +197,41 @@ private void AllocateWriteHeadSynchronized(int sizeHint)
_writingHeadBytesBuffered = 0;
}

BufferSegment newSegment = AllocateSegment(sizeHint);
if (_writingHead.Length == 0)
Copy link
Member

@davidfowl davidfowl Jul 20, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would use _writingHeadBytesBuffered instead (but the logic would need to be reworked a bit)

Copy link
Member Author

@BrennanConroy BrennanConroy Jul 20, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think I tried using that field initially, but it gets set to 0 in Advance with a non-zero value, so you need to check some other value anyways. What's wrong with _writingHead.Length?

Copy link
Member

@davidfowl davidfowl Jul 20, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nothing wrong with it, it just avoids the indirection but this is fine.

{
// If we got here that means Advance was called with 0 bytes or GetMemory was called again without any writes occurring
// And, the newly requested memory size is greater than our unused segments internal memory buffer
// So we should reuse the BufferSegment and replace the memory it's holding, this way ReadAsync will not receive a buffer with one segment being empty
_writingHead.ResetMemory();
RentMemory(_writingHead, sizeHint);
}
else
{
BufferSegment newSegment = AllocateSegment(sizeHint);

_writingHead.SetNext(newSegment);
_writingHead = newSegment;
_writingHead.SetNext(newSegment);
_writingHead = newSegment;
}
}
}
}
}

private BufferSegment AllocateSegment(int sizeHint)
{
Debug.Assert(sizeHint >= 0);
BufferSegment newSegment = CreateSegmentUnsynchronized();

RentMemory(newSegment, sizeHint);

return newSegment;
}

private void RentMemory(BufferSegment segment, int sizeHint)
{
// Segment should be new or reset, otherwise a memory leak could occur
Debug.Assert(segment.MemoryOwner is null);
Debug.Assert(sizeHint >= 0);

MemoryPool<byte>? pool = null;
int maxSize = -1;

Expand All @@ -223,18 +244,16 @@ private BufferSegment AllocateSegment(int sizeHint)
if (sizeHint <= maxSize)
{
// Use the specified pool as it fits. Specified pool is not null as maxSize == -1 if _pool is null.
newSegment.SetOwnedMemory(pool!.Rent(GetSegmentSize(sizeHint, maxSize)));
segment.SetOwnedMemory(pool!.Rent(GetSegmentSize(sizeHint, maxSize)));
}
else
{
// Use the array pool
int sizeToRequest = GetSegmentSize(sizeHint);
newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(sizeToRequest));
segment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(sizeToRequest));
}

_writingHeadMemory = newSegment.AvailableMemory;

return newSegment;
_writingHeadMemory = segment.AvailableMemory;
}

private int GetSegmentSize(int sizeHint, int maxBufferSize = int.MaxValue)
Expand Down Expand Up @@ -553,7 +572,7 @@ void MoveReturnEndToNextBlock()
while (returnStart != null && returnStart != returnEnd)
{
BufferSegment? next = returnStart.NextSegment;
returnStart.ResetMemory();
returnStart.Reset();
ReturnSegmentUnsynchronized(returnStart);
returnStart = next;
}
Expand Down Expand Up @@ -866,7 +885,7 @@ private void CompletePipe()
BufferSegment returnSegment = segment;
segment = segment.NextSegment;

returnSegment.ResetMemory();
returnSegment.Reset();
}

_writingHead = null;
Expand Down
5 changes: 3 additions & 2 deletions src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeReader.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,6 @@ private void AdvanceTo(BufferSegment? consumedSegment, int consumedIndex, Buffer
while (returnStart != returnEnd)
{
BufferSegment next = returnStart.NextSegment!;
returnStart.ResetMemory();
ReturnSegmentUnsynchronized(returnStart);
returnStart = next;
}
Expand Down Expand Up @@ -192,7 +191,7 @@ private bool CompleteAndGetNeedsDispose()
BufferSegment returnSegment = segment;
segment = segment.NextSegment;

returnSegment.ResetMemory();
returnSegment.Reset();
}

return !LeaveOpen;
Expand Down Expand Up @@ -624,6 +623,8 @@ private void ReturnSegmentUnsynchronized(BufferSegment segment)
Debug.Assert(segment != _readHead, "Returning _readHead segment that's in use!");
Debug.Assert(segment != _readTail, "Returning _readTail segment that's in use!");

segment.Reset();

if (_bufferSegmentPool.Count < MaxSegmentPoolSize)
{
_bufferSegmentPool.Push(segment);
Expand Down
3 changes: 1 addition & 2 deletions src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeWriter.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,7 @@ private BufferSegment CreateSegmentUnsynchronized()

private void ReturnSegmentUnsynchronized(BufferSegment segment)
{
segment.Reset();
if (_bufferSegmentPool.Count < MaxSegmentPoolSize)
{
_bufferSegmentPool.Push(segment);
Expand Down Expand Up @@ -323,7 +324,6 @@ private async ValueTask<FlushResult> FlushAsyncInternal(bool writeToStream, Read
await InnerStream.WriteAsync(returnSegment.Memory, localToken).ConfigureAwait(false);
}

returnSegment.ResetMemory();
ReturnSegmentUnsynchronized(returnSegment);

// Update the head segment after we return the current segment
Expand Down Expand Up @@ -400,7 +400,6 @@ private void FlushInternal(bool writeToStream)
#endif
}

returnSegment.ResetMemory();
ReturnSegmentUnsynchronized(returnSegment);

// Update the head segment after we return the current segment
Expand Down
65 changes: 65 additions & 0 deletions src/libraries/System.IO.Pipelines/tests/PipeWriterTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,71 @@ public async Task CanWriteNothingToBuffer()
await buffer.FlushAsync();
}

[Fact]
public async Task WriteNothingThenWriteToNewSegment()
{
// Regression test: write nothing to force a segment to be created, then do a large write that's larger than the currently empty segment to force another new segment
// Verify that no 0 length segments are returned from the Reader.
PipeWriter buffer = Pipe.Writer;
Memory<byte> memory = buffer.GetMemory();
buffer.Advance(0); // doing nothing, the hard way
await buffer.FlushAsync();

memory = buffer.GetMemory(memory.Length + 1);
buffer.Advance(memory.Length);
await buffer.FlushAsync();

var res = await Pipe.Reader.ReadAsync();
Assert.True(res.Buffer.IsSingleSegment);
Assert.Equal(memory.Length, res.Buffer.Length);
}

[Fact]
public async Task WriteNothingBetweenTwoFullWrites()
{
int totalWrittenLength = 0;
PipeWriter buffer = Pipe.Writer;
Memory<byte> memory = buffer.GetMemory();
buffer.Advance(memory.Length); // doing nothing, the hard way
totalWrittenLength += memory.Length;
await buffer.FlushAsync();

memory = buffer.GetMemory();
buffer.Advance(0); // doing nothing, the hard way
await buffer.FlushAsync();

memory = buffer.GetMemory(memory.Length + 1);
buffer.Advance(memory.Length);
totalWrittenLength += memory.Length;
await buffer.FlushAsync();

var res = await Pipe.Reader.ReadAsync();
var segmentCount = 0;
foreach (ReadOnlyMemory<byte> _ in res.Buffer)
{
segmentCount++;
}
Assert.Equal(2, segmentCount);
Assert.Equal(totalWrittenLength, res.Buffer.Length);
}

[Fact]
public async Task WriteNothingThenWriteSomeBytes()
{
PipeWriter buffer = Pipe.Writer;
_ = buffer.GetMemory();
buffer.Advance(0); // doing nothing, the hard way
await buffer.FlushAsync();

var memory = buffer.GetMemory();
buffer.Advance(memory.Length);
await buffer.FlushAsync();

var res = await Pipe.Reader.ReadAsync();
Assert.True(res.Buffer.IsSingleSegment);
Assert.Equal(memory.Length, res.Buffer.Length);
}

[Fact]
public void EmptyWriteDoesNotThrow()
{
Expand Down