diff --git a/src/Renci.SshNet/Common/PipeStream.cs b/src/Renci.SshNet/Common/PipeStream.cs index 2a8d739e4..1344355cf 100644 --- a/src/Renci.SshNet/Common/PipeStream.cs +++ b/src/Renci.SshNet/Common/PipeStream.cs @@ -65,13 +65,7 @@ private int Read(Span 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 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; + } + /// public override void Write(byte[] buffer, int offset, int count) { @@ -134,17 +142,80 @@ private void WriteCore(ReadOnlySpan 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). + + /// + public override Task 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 + /// + public override ValueTask ReadAsync(Memory buffer, CancellationToken cancellationToken = default) + { + if (!cancellationToken.IsCancellationRequested && + TryReadWithoutBlocking(buffer.Span, out var bytesRead)) + { + return new ValueTask(bytesRead); + } + + return base.ReadAsync(buffer, cancellationToken); + } +#endif + + private bool TryReadWithoutBlocking(Span 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); + } + } /// public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) diff --git a/test/Renci.SshNet.Tests/Classes/Common/PipeStreamTest.cs b/test/Renci.SshNet.Tests/Classes/Common/PipeStreamTest.cs index 9759b4c60..a3466c245 100644 --- a/test/Renci.SshNet.Tests/Classes/Common/PipeStreamTest.cs +++ b/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 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 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 readTask = pipeStream.ReadAsync(new byte[4], 0, 4, token); + + Assert.IsTrue(readTask.IsCanceled); + + var ex = await Assert.ThrowsAsync(() => 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 readTask = pipeStream.ReadAsync(readBuffer.AsMemory()); + + Assert.IsTrue(readTask.IsCompletedSuccessfully); + Assert.AreEqual(3, await readTask); + CollectionAssert.AreEqual(new byte[] { 1, 2, 3 }, readBuffer); + + ValueTask 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 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() {