From 6e19478a4b281ff1027fabb3da080788e1455eb3 Mon Sep 17 00:00:00 2001
From: Ben Adams <thundercat@illyriad.co.uk>
Date: Mon, 25 May 2020 21:09:28 +0100
Subject: [PATCH] Pipelines fast-path, return single block outside of lock

---
 .../src/System/IO/Pipelines/BufferSegment.cs  | 25 ++++++++++
 .../src/System/IO/Pipelines/Pipe.cs           | 50 +++++++++++++++----
 .../System/IO/Pipelines/PipeOperationState.cs | 13 +++++
 3 files changed, 79 insertions(+), 9 deletions(-)

diff --git a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
index 800b1414260838..0a6bd77a629ecc 100644
--- a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
+++ b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
@@ -60,6 +60,31 @@ public void SetOwnedMemory(byte[] arrayPoolBuffer)
             AvailableMemory = arrayPoolBuffer;
         }
 
+        public IMemoryOwner<byte>? GetMemoryBlockAndResetSegment()
+        {
+            IMemoryOwner<byte>? memoryOwner = _memoryOwner;
+            if (memoryOwner != null)
+            {
+                _memoryOwner = null;
+                // We return the owner rather than disposing
+            }
+            else
+            {
+                Debug.Assert(_array != null);
+                ArrayPool<byte>.Shared.Return(_array);
+                _array = null;
+            }
+
+            Next = null;
+            RunningIndex = 0;
+            Memory = default;
+            _next = null;
+            _end = 0;
+            AvailableMemory = default;
+
+            return memoryOwner;
+        }
+
         public void ResetMemory()
         {
             IMemoryOwner<byte>? memoryOwner = _memoryOwner;
diff --git a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs
index b92832c053dad5..4d719ba10bbc5e 100644
--- a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs
+++ b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs
@@ -438,11 +438,10 @@ private void AdvanceReader(BufferSegment? consumedSegment, int consumedIndex, Bu
                 ThrowHelper.ThrowInvalidOperationException_InvalidExaminedOrConsumedPosition();
             }
 
-            BufferSegment? returnStart = null;
-            BufferSegment? returnEnd = null;
-
             CompletionData completionData = default;
 
+            IMemoryOwner<byte>? blockToReturn = null;
+            bool success;
             lock (_sync)
             {
                 var examinedEverything = false;
@@ -475,6 +474,8 @@ private void AdvanceReader(BufferSegment? consumedSegment, int consumedIndex, Bu
                     }
                 }
 
+                BufferSegment? returnStart = null;
+                BufferSegment? returnEnd = null;
                 if (consumedSegment != null)
                 {
                     if (_readHead == null)
@@ -541,18 +542,49 @@ void MoveReturnEndToNextBlock()
                     _readerAwaitable.SetUncompleted();
                 }
 
-                while (returnStart != null && returnStart != returnEnd)
+                if (returnStart != null && returnStart != returnEnd)
                 {
                     BufferSegment? next = returnStart.NextSegment;
-                    returnStart.ResetMemory();
-                    ReturnSegmentUnsynchronized(returnStart);
-                    returnStart = next;
+                    if (next == null || next == returnEnd)
+                    {
+                        // Fast-path, single block to return; we will return the block outside of the lock.
+                        blockToReturn = returnStart.GetMemoryBlockAndResetSegment();
+                        ReturnSegmentUnsynchronized(returnStart);
+                    }
+                    else
+                    {
+                        // Multiple blocks to return; we will do it inside of lock.
+                        returnStart.ResetMemory();
+                        ReturnSegmentUnsynchronized(returnStart);
+                        returnStart = next;
+                        do
+                        {
+                            next = returnStart.NextSegment;
+                            returnStart.ResetMemory();
+                            ReturnSegmentUnsynchronized(returnStart);
+                            returnStart = next;
+                        } while (returnStart != null && returnStart != returnEnd) ;
+                    }
                 }
 
-                _operationState.EndRead();
+                success = _operationState.TryEndRead();
             }
 
-            TrySchedule(_writerScheduler, completionData);
+            if (success)
+            {
+                TrySchedule(_writerScheduler, completionData);
+            }
+
+            // Return the block before throwing the exception if there is one.
+            if (blockToReturn != null)
+            {
+                blockToReturn.Dispose();
+            }
+
+            if (!success)
+            {
+                ThrowHelper.ThrowInvalidOperationException_NoReadToComplete();
+            }
         }
 
         internal void CompleteReader(Exception? exception)
diff --git a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/PipeOperationState.cs b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/PipeOperationState.cs
index 1271df5a9a11c8..9eca42c3a30f0b 100644
--- a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/PipeOperationState.cs
+++ b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/PipeOperationState.cs
@@ -46,6 +46,19 @@ public void EndRead()
             _state &= ~(State.Reading | State.ReadingTentative);
         }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public bool TryEndRead()
+        {
+            if ((_state & State.Reading) != State.Reading &&
+                (_state & State.ReadingTentative) != State.ReadingTentative)
+            {
+                return false;
+            }
+
+            _state &= ~(State.Reading | State.ReadingTentative);
+            return true;
+        }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public void BeginWrite()
         {