Add a synchronous fast-path to PipeStream.ReadAsync

src/Renci.SshNet/Common/PipeStream.cs

@@ -65,13 +65,7 @@ private int Read(Span<byte> buffer)
                     _ = Monitor.Wait(_sync);
                 }
 
-                var bytesRead = Math.Min(buffer.Length, _buffer.ActiveLength);
-
-                _buffer.ActiveReadOnlySpan.Slice(0, bytesRead).CopyTo(buffer);
-
-                _buffer.Discard(bytesRead);
-
-                return bytesRead;
+                return ReadCore(buffer);
             }
         }
 
@@ -85,6 +79,20 @@ public override int ReadByte()
         }
 #endif
 
+        private int ReadCore(Span<byte> buffer)
+        {
+            Debug.Assert(Monitor.IsEntered(_sync));
+            Debug.Assert(_buffer.ActiveLength > 0 || _disposed);
+
+            var bytesRead = Math.Min(buffer.Length, _buffer.ActiveLength);
+
+            _buffer.ActiveReadOnlySpan.Slice(0, bytesRead).CopyTo(buffer);
+
+            _buffer.Discard(bytesRead);
+
+            return bytesRead;
+        }
+
         /// <inheritdoc/>
         public override void Write(byte[] buffer, int offset, int count)
         {
@@ -134,17 +142,80 @@ private void WriteCore(ReadOnlySpan<byte> buffer)
             Monitor.PulseAll(_sync);
         }
 
-        // We provide overrides for async Write methods but not async Read.
-        // The default implementations from the base class effectively call the
-        // sync methods on a threadpool thread, but only allowing one async
-        // operation at a time (for protecting thread-unsafe implementations).
-        // This constraint is desirable for reads because if there were multiple
-        // readers and no data coming in, our current Monitor.Wait implementation
-        // would just block as many threadpool threads as there are readers.
-        // But since a write is just short-lived buffer copying and can unblock
+        // We provide overrides for async Write methods, but the async Read overrides
+        // have only partial implementations (for the happy path when data is available
+        // to be read), and otherwise defer to the base implementation.
+        // The default implementations from the base class effectively call the sync
+        // methods on a threadpool thread, but only allowing one such async-over-sync
+        // operation at a time.
+        // This constraint is desirable for blocking reads because if there were
+        // multiple readers and no data coming in, our current Monitor.Wait
+        // implementation would just block as many threadpool threads as there
+        // are readers. But when data is already available, a read is just
+        // short-lived buffer copying, so we can complete synchronously and avoid
+        // the overhead of the base.
+        // Likewise, since a write is just short-lived buffer copying and can unblock
         // readers, it is beneficial to circumvent the one-at-a-time constraint,
         // as otherwise a waiting async read will block the async write that could
         // unblock it.
+        // Note that we do not override BeginRead/EndRead because the base ReadAsync
+        // takes a faster path in this way (and it can introduce circular dependency).
+
+        /// <inheritdoc/>
+        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+#if !NET
+            ThrowHelper.
+#endif
+            ValidateBufferArguments(buffer, offset, count);
+
+            if (!cancellationToken.IsCancellationRequested &&
+                TryReadWithoutBlocking(buffer.AsSpan(offset, count), out var bytesRead))
+            {
+                return Task.FromResult(bytesRead);
+            }
+
+            return base.ReadAsync(buffer, offset, count, cancellationToken);
+        }
+
+#if NET
+        /// <inheritdoc/>
+        public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
+        {
+            if (!cancellationToken.IsCancellationRequested &&
+                TryReadWithoutBlocking(buffer.Span, out var bytesRead))
+            {
+                return new ValueTask<int>(bytesRead);
+            }
+
+            return base.ReadAsync(buffer, cancellationToken);
+        }
+#endif
+
+        private bool TryReadWithoutBlocking(Span<byte> buffer, out int bytesRead)
+        {
+            if (!Monitor.TryEnter(_sync))
+            {
+                bytesRead = 0;
+                return false;
+            }
+
+            try
+            {
+                if (_buffer.ActiveLength == 0 && !_disposed)
+                {
+                    bytesRead = 0;
+                    return false;
+                }
+
+                bytesRead = ReadCore(buffer);
+                return true;
+            }
+            finally
+            {
+                Monitor.Exit(_sync);
+            }
+        }
 
         /// <inheritdoc/>
         public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)

test/Renci.SshNet.Tests/Classes/Common/PipeStreamTest.cs

@@ -1,5 +1,6 @@
 using System;
 using System.IO;
+using System.Threading;
 using System.Threading.Tasks;
 
 using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -157,6 +158,115 @@ public async Task Read_EmptySpan_OnlyReturnsZeroWhenDataAvailable()
         }
 #endif
 
+        [TestMethod]
+        public async Task ReadAsync_DataAvailable_CompletesSynchronously()
+        {
+            var pipeStream = new PipeStream();
+
+            pipeStream.Write(new byte[] { 1, 2, 3, 4 }, 0, 4);
+
+            var readBuffer = new byte[3];
+
+            Task<int> readTask = pipeStream.ReadAsync(readBuffer, 0, 3);
+
+            Assert.IsTrue(readTask.IsCompleted);
+            Assert.AreEqual(3, await readTask);
+            CollectionAssert.AreEqual(new byte[] { 1, 2, 3 }, readBuffer);
+
+            readTask = pipeStream.ReadAsync(readBuffer, 0, 3);
+
+            Assert.IsTrue(readTask.IsCompleted);
+            Assert.AreEqual(1, await readTask);
+            Assert.AreEqual(4, readBuffer[0]);
+
+            // No data left, so the next read should block.
+            readTask = pipeStream.ReadAsync(readBuffer, 0, 3);
+
+            await Task.Delay(50);
+
+            Assert.IsFalse(readTask.IsCompleted);
+
+            await pipeStream.WriteAsync(new byte[] { 5 }, 0, 1);
+
+            Assert.AreEqual(1, await readTask);
+            Assert.AreEqual(5, readBuffer[0]);
+        }
+
+        [TestMethod]
+        public async Task ReadAsync_AfterDispose_CompletesSynchronously()
+        {
+            var pipeStream = new PipeStream();
+
+            pipeStream.Write(new byte[] { 1, 2, 3, 4 }, 0, 4);
+
+            pipeStream.Dispose();
+
+            Task<int> readTask = pipeStream.ReadAsync(new byte[5], 0, 5);
+
+            Assert.IsTrue(readTask.IsCompleted);
+            Assert.AreEqual(4, await readTask);
+
+            readTask = pipeStream.ReadAsync(new byte[5], 0, 5);
+
+            Assert.IsTrue(readTask.IsCompleted);
+            Assert.AreEqual(0, await readTask);
+        }
+
+        [TestMethod]
+        public async Task ReadAsync_DataAvailable_CancelledToken_TaskIsCancelled()
+        {
+            var pipeStream = new PipeStream();
+
+            pipeStream.Write(new byte[] { 1, 2, 3, 4 }, 0, 4);
+
+            var token = new CancellationToken(canceled: true);
+
+            Task<int> readTask = pipeStream.ReadAsync(new byte[4], 0, 4, token);
+
+            Assert.IsTrue(readTask.IsCanceled);
+
+            var ex = await Assert.ThrowsAsync<OperationCanceledException>(() => readTask);
+
+            Assert.AreEqual(token, ex.CancellationToken);
+            Assert.AreEqual(4, pipeStream.Length);
+        }
+
+#if NET
+        [TestMethod]
+        public async Task ReadAsync_Memory_DataAvailable_CompletesSynchronously()
+        {
+            var pipeStream = new PipeStream();
+
+            pipeStream.Write(new byte[] { 1, 2, 3, 4 }, 0, 4);
+
+            var readBuffer = new byte[3];
+
+            ValueTask<int> readTask = pipeStream.ReadAsync(readBuffer.AsMemory());
+
+            Assert.IsTrue(readTask.IsCompletedSuccessfully);
+            Assert.AreEqual(3, await readTask);
+            CollectionAssert.AreEqual(new byte[] { 1, 2, 3 }, readBuffer);
+
+            ValueTask<int> readTask2 = pipeStream.ReadAsync(readBuffer.AsMemory());
+
+            Assert.IsTrue(readTask2.IsCompletedSuccessfully);
+            Assert.AreEqual(1, await readTask2);
+            Assert.AreEqual(4, readBuffer[0]);
+
+            // No data left, so the next read should block.
+            ValueTask<int> readTask3 = pipeStream.ReadAsync(readBuffer.AsMemory());
+
+            await Task.Delay(50);
+
+            Assert.IsFalse(readTask3.IsCompleted);
+
+            await pipeStream.WriteAsync(new byte[] { 5 });
+
+            Assert.AreEqual(1, await readTask3);
+            Assert.AreEqual(5, readBuffer[0]);
+        }
+#endif
+
         [TestMethod]
         public void Read_AfterDispose_StillWorks()
         {