Make BufferSegment pool lock-free

src/System.IO.Pipelines/src/System.IO.Pipelines.csproj

@@ -10,6 +10,9 @@
     <Compile Include="System\IO\Pipelines\FlushResult.cs" />
     <Compile Include="System\IO\Pipelines\InlineScheduler.cs" />
     <Compile Include="System\IO\Pipelines\IDuplexPipe.cs" />
+    <Compile Include="$(CommonPath)\Internal\Padding.cs">
+      <Link>Common\Internal\Padding.cs</Link>
+    </Compile>
     <Compile Include="System\IO\Pipelines\Pipe.DefaultPipeReader.cs" />
     <Compile Include="System\IO\Pipelines\Pipe.DefaultPipeWriter.cs" />
     <Compile Include="System\IO\Pipelines\Pipe.cs" />
@@ -24,6 +27,7 @@
     <Compile Include="System\IO\Pipelines\PipeWriter.cs" />
     <Compile Include="System\IO\Pipelines\ReadResult.cs" />
     <Compile Include="System\IO\Pipelines\ResultFlags.cs" />
+    <Compile Include="System\IO\Pipelines\SingleProducerSingleConsumerPool.cs" />
     <Compile Include="System\IO\Pipelines\StreamPipeExtensions.cs" />
     <Compile Include="System\IO\Pipelines\ThrowHelper.cs" />
   </ItemGroup>

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

@@ -42,11 +42,7 @@ public sealed partial class Pipe
         private readonly PipeScheduler _readerScheduler;
         private readonly PipeScheduler _writerScheduler;
 
-        // Segment Pool
-        private int _pooledSegmentCount;
-        private readonly BufferSegment[] _bufferSegmentPool;
-        // Temporary list to hold Segments return while being reset
-        private readonly BufferSegment[] _bufferSegmentsToReturn;
+        private readonly SingleProducerSingleConsumerPool<BufferSegment> _bufferSegmentPool;
 
         private readonly DefaultPipeReader _reader;
         private readonly DefaultPipeWriter _writer;
@@ -98,8 +94,7 @@ public Pipe(PipeOptions options)
                 ThrowHelper.ThrowArgumentNullException(ExceptionArgument.options);
             }
 
-            _bufferSegmentPool = new BufferSegment[SegmentPoolSize];
-            _bufferSegmentsToReturn = new BufferSegment[SegmentPoolSize];
+            _bufferSegmentPool = new SingleProducerSingleConsumerPool<BufferSegment>(SegmentPoolSize);
 
             _operationState = default;
             _readerCompletion = default;
@@ -239,64 +234,25 @@ private int GetSegmentSize(int sizeHint, int maxBufferSize = int.MaxValue)
             return adjustedToMaximumSize;
         }
 
-        private BufferSegment CreateSegment()
-        {
-            BufferSegment[] segmentPool = _bufferSegmentPool;
-            lock (segmentPool)
-            {
-                int index = _pooledSegmentCount - 1;
-                if ((uint)index < (uint)segmentPool.Length)
-                {
-                    _pooledSegmentCount = index;
-                    return segmentPool[index];
-                }
-            }
-
-            return new BufferSegment();
-        }
+        private BufferSegment CreateSegment() =>
+            _bufferSegmentPool.TryDequeue(out BufferSegment segment) ? segment : new BufferSegment();
 
         private void ReturnSegments(BufferSegment from, BufferSegment toExclusive)
         {
             Debug.Assert(from != null);
             Debug.Assert(from != toExclusive);
 
-            // Reset the Segments and return their data out of lock
-            BufferSegment[] segmentToReturn = _bufferSegmentsToReturn;
-            int count = 0;
             do
             {
                 BufferSegment next = from.NextSegment;
                 Debug.Assert(next != null || toExclusive == null);
 
                 from.ResetMemory();
 
-                if ((uint)count < (uint)segmentToReturn.Length)
-                {
-                    // Store in temporary list while preforming expensive resets
-                    segmentToReturn[count] = from;
-                    count++;
-                }
+                _bufferSegmentPool.TryEnqueue(from);
 
                 from = next;
             } while (from != toExclusive);
-
-            // Add the Segments back to pool from the temporary list under lock
-            BufferSegment[] segmentPool = _bufferSegmentPool;
-            lock (segmentPool)
-            {
-                int index = _pooledSegmentCount;
-                for (int i = 0; i < count; i++)
-                {
-                    if ((uint)index < (uint)segmentPool.Length)
-                    {
-                        segmentPool[index] = segmentToReturn[i];
-                        index++;
-                    }
-                    segmentToReturn[i] = null;
-                }
-
-                _pooledSegmentCount = index;
-            }
         }
 
         internal bool CommitUnsynchronized()

src/System.IO.Pipelines/src/System/IO/Pipelines/SingleProducerSingleConsumerPool.cs

@@ -0,0 +1,167 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Diagnostics;
+using System.Runtime.InteropServices;
+
+using Internal;
+
+namespace System.IO.Pipelines
+{
+    internal sealed class SingleProducerSingleConsumerPool<T> where T : class
+    {
+        // Adapted from SingleProducerSingleConsumerQueue
+        /// <summary>The maximum size to use for segments (in number of elements).</summary>
+        private const int MaxSegmentSize = 0x1000000; // this could be made as large as Int32.MaxValue / 2
+        /// <summary>The data stored in this segment.</summary>
+        private readonly RefAsValueType[] _array;
+        /// <summary>Details about the segment.</summary>
+        private SegmentState _state; // separated out to enable StructLayout attribute to take effect
+
+        /// <summary>Initializes the queue.</summary>
+        public SingleProducerSingleConsumerPool(int size)
+        {
+            // Validate constants in ctor rather than in an explicit cctor that would cause perf degradation
+            Debug.Assert(size > 0, "Initial segment size must be > 0.");
+            Debug.Assert((size & (size - 1)) == 0, "Initial segment size must be a power of 2");
+            Debug.Assert(size <= MaxSegmentSize, "Initial segment size should be <= maximum.");
+            Debug.Assert(MaxSegmentSize < int.MaxValue / 2, "Max segment size * 2 must be < Int32.MaxValue, or else overflow could occur.");
+
+            // Initialize the pool
+            _array = new RefAsValueType[size];
+        }
+
+        /// <summary>Enqueues an item into the queue.</summary>
+        /// <param name="item">The item to enqueue.</param>
+        public bool TryEnqueue(T item)
+        {
+            RefAsValueType[] array = _array;
+            int last = _state._last; // local copy to avoid multiple volatile reads
+
+            // Fast path: there's obviously room
+            int tail2 = (last + 1) & (array.Length - 1);
+            if (tail2 != _state._firstCopy)
+            {
+                array[last] = item;
+                _state._last = tail2;
+                return true;
+            }
+            // Slow path: there may not be room
+            else
+            {
+                return TryEnqueueSlow(item);
+            }
+        }
+
+        /// <summary>Enqueues an item into the queue.</summary>
+        /// <param name="item">The item to enqueue.</param>
+        private bool TryEnqueueSlow(T item)
+        {
+            if (_state._firstCopy != _state._first)
+            {
+                _state._firstCopy = _state._first;
+                return TryEnqueue(item); // will only recur once for this enqueue operation
+            }
+
+            return false;
+        }
+
+        /// <summary>Attempts to dequeue an item from the queue.</summary>
+        /// <param name="result">The dequeued item.</param>
+        /// <returns>true if an item could be dequeued; otherwise, false.</returns>
+        public bool TryDequeue(out T result)
+        {
+            RefAsValueType[] array = _array;
+            int first = _state._first; // local copy to avoid multiple volatile reads
+
+            // Fast path: there's obviously data available
+            if (first != _state._lastCopy)
+            {
+                result = array[first];
+                array[first] = default; // Clear the slot to release the element
+                _state._first = (first + 1) & (array.Length - 1);
+                return true;
+            }
+            // Slow path: there may not be data available
+            else
+            {
+                return TryDequeueSlow(out result);
+            }
+        }
+
+        /// <summary>Attempts to dequeue an item from the queue.</summary>
+        /// <param name="result">The dequeued item.</param>
+        /// <returns>true if an item could be dequeued; otherwise, false.</returns>
+        private bool TryDequeueSlow(out T result)
+        {
+            RefAsValueType[] array = _array;
+            if (_state._last != _state._lastCopy)
+            {
+                _state._lastCopy = _state._last;
+                return TryDequeue(out result); // will only recur once for this dequeue operation
+            }
+
+            int first = _state._first; // local copy to avoid extraneous volatile reads
+
+            if (first == _state._last)
+            {
+                result = default;
+                return false;
+            }
+
+            result = array[first];
+            array[first] = default; // Clear the slot to release the element
+            _state._first = (first + 1) & (_array.Length - 1);
+            _state._lastCopy = _state._last; // Refresh _lastCopy to ensure that _first has not passed _lastCopy
+
+            return true;
+        }
+
+        /// <summary>Stores information about a segment.</summary>
+        [StructLayout(LayoutKind.Sequential)] // enforce layout so that padding reduces false sharing
+        private struct SegmentState
+        {
+            /// <summary>Padding to reduce false sharing between the segment's array and _first.</summary>
+            internal PaddingFor32 _pad0;
+
+            /// <summary>The index of the current head in the segment.</summary>
+            internal volatile int _first;
+            /// <summary>A copy of the current tail index.</summary>
+            internal int _lastCopy; // not volatile as read and written by the producer, except for IsEmpty, and there _lastCopy is only read after reading the volatile _first
+
+            /// <summary>Padding to reduce false sharing between the first and last.</summary>
+            internal PaddingFor32 _pad1;
+
+            /// <summary>A copy of the current head index.</summary>
+            internal int _firstCopy; // not volatile as only read and written by the consumer thread
+            /// <summary>The index of the current tail in the segment.</summary>
+            internal volatile int _last;
+
+            /// <summary>Padding to reduce false sharing with the last and what's after the segment.</summary>
+            internal PaddingFor32 _pad2;
+        }
+
+        /// <summary>
+        /// A simple struct we wrap reference types inside when storing in arrays to
+        /// bypass the CLR's covariant checks when writing to arrays.
+        /// </summary>
+        /// <remarks>
+        /// We use <see cref="RefAsValueType"/> as a wrapper to avoid paying the cost of covariant checks whenever
+        /// the underlying array that the <see cref="SingleProducerSingleConsumerPool{T}"/> class uses is written to. 
+        /// We've recognized this as a perf win in ETL traces for these stack frames:
+        /// clr!JIT_Stelem_Ref
+        ///   clr!ArrayStoreCheck
+        ///     clr!ObjIsInstanceOf
+        /// </remarks>
+        internal struct RefAsValueType
+        {
+            private T _value;
+
+            private RefAsValueType(T value) => _value = value;
+
+            public static implicit operator T(RefAsValueType value) => value._value;
+            public static implicit operator RefAsValueType(T value) => new RefAsValueType(value);
+        }
+    }
+}