Avoid empty segments from PipeReader.ReadAsync

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

@@ -59,7 +59,7 @@ public void SetOwnedMemory(byte[] arrayPoolBuffer)
             AvailableMemory = arrayPoolBuffer;
         }
 
-        public void ResetMemory()
+        public void ResetMemory(bool preserveIndex = false)
         {
             IMemoryOwner<byte>? memoryOwner = _memoryOwner;
             if (memoryOwner != null)
@@ -75,7 +75,8 @@ public void ResetMemory()
             }
 
             Next = null;
-            RunningIndex = 0;
+            // RunningIndex is used to determine the total length of the linked segments and contains the length of all the previous segments (not including this one)
+            RunningIndex = preserveIndex ? RunningIndex : 0;
             Memory = default;
             _next = null;
             _end = 0;

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

@@ -197,20 +197,41 @@ private void AllocateWriteHeadSynchronized(int sizeHint)
                             _writingHeadBytesBuffered = 0;
                         }
 
-                        BufferSegment newSegment = AllocateSegment(sizeHint);
+                        if (_writingHead.Length == 0)
+                        {
+                            // 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(preserveIndex: true);
+                            SetupSegment(_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();
 
+            SetupSegment(newSegment, sizeHint);
+
+            return newSegment;
+        }
+
+        private void SetupSegment(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;
 
@@ -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)

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

@@ -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()
         {