Schedule completion prior to pooling segments

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

@@ -204,7 +204,7 @@ private void AllocateWriteHeadUnsynchronized(int sizeHint)
 
         private BufferSegment AllocateSegment(int sizeHint)
         {
-            BufferSegment newSegment = CreateSegmentUnsynchronized();
+            BufferSegment newSegment = CreateSegment();
 
             if (_pool is null)
             {
@@ -236,23 +236,46 @@ private int GetSegmentSize(int sizeHint, int maxBufferSize = int.MaxValue)
             return adjustedToMaximumSize;
         }
 
-        private BufferSegment CreateSegmentUnsynchronized()
+        private BufferSegment CreateSegment()
         {
-            if (_pooledSegmentCount > 0)
+            BufferSegment[] segmentPool = _bufferSegmentPool;
+            lock (segmentPool)
             {
-                _pooledSegmentCount--;
-                return _bufferSegmentPool[_pooledSegmentCount];
+                int index = _pooledSegmentCount - 1;
+                if ((uint)index < (uint)segmentPool.Length)
+                {
+                    _pooledSegmentCount = index;
+                    return segmentPool[index];
+                }
             }
 
             return new BufferSegment();
         }
 
-        private void ReturnSegmentUnsynchronized(BufferSegment segment)
+        private void ReturnSegments(BufferSegment from, BufferSegment toExclusive)
         {
-            if (_pooledSegmentCount < _bufferSegmentPool.Length)
+            Debug.Assert(from != null);
+            Debug.Assert(from != toExclusive);
+
+            BufferSegment[] segmentPool = _bufferSegmentPool;
+            lock (segmentPool)
             {
-                _bufferSegmentPool[_pooledSegmentCount] = segment;
-                _pooledSegmentCount++;
+                int index = _pooledSegmentCount;
+                do
+                {
+                    BufferSegment next = from.NextSegment;
+                    from.ResetMemory();
+
+                    if ((uint)index < (uint)segmentPool.Length)
+                    {
+                        segmentPool[index] = from;
+                        index++;
+                    }
+
+                    from = next;
+                } while (from != toExclusive);
+
+                _pooledSegmentCount = index;
             }
         }
 
@@ -494,18 +517,15 @@ private void AdvanceReader(BufferSegment consumedSegment, int consumedIndex, Buf
                     _readerAwaitable.SetUncompleted();
                 }
 
-                while (returnStart != null && returnStart != returnEnd)
-                {
-                    BufferSegment next = returnStart.NextSegment;
-                    returnStart.ResetMemory();
-                    ReturnSegmentUnsynchronized(returnStart);
-                    returnStart = next;
-                }
-
                 _operationState.EndRead();
             }
 
             TrySchedule(_writerScheduler, completionData);
+
+            if (returnStart != null && returnStart != returnEnd)
+            {
+                ReturnSegments(returnStart, returnEnd);
+            }
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -735,6 +755,7 @@ private static void ExecuteWithExecutionContext(object state)
 
         private void CompletePipe()
         {
+            BufferSegment segment;
             lock (_sync)
             {
                 if (_disposed)
@@ -743,22 +764,20 @@ private void CompletePipe()
                 }
 
                 _disposed = true;
-                // Return all segments
+                // Get segments chain to return
                 // if _readHead is null we need to try return _commitHead
                 // because there might be a block allocated for writing
-                BufferSegment segment = _readHead ?? _readTail;
-                while (segment != null)
-                {
-                    BufferSegment returnSegment = segment;
-                    segment = segment.NextSegment;
-
-                    returnSegment.ResetMemory();
-                }
+                segment = _readHead ?? _readTail;
 
                 _writingHead = null;
                 _readHead = null;
                 _readTail = null;
             }
+
+            if (segment != null)
+            {
+                ReturnSegments(segment, toExclusive: null);
+            }
         }
 
         internal ValueTaskSourceStatus GetReadAsyncStatus()