Stop sending after a partial frame write

LengthPrefixCommunicationChannel writes each message as a 7-bit encoded
byte count followed by the payload, through a BufferedStream. A payload
larger than the free buffer space forces the buffered prefix out first,
so a write that fails part way through can leave the prefix on the wire
without its body.

Send logged that error and left the channel usable. Callers that swallow
CommunicationException, such as the timer-driven cache flush, keep
sending, and the receiver reads the next message inside the frame the
previous one never finished.

A captured trace shows the whole sequence. A 33,185 byte StatsChange
flushes its 3 byte prefix a1 83 02, then the payload write fails with
SocketException 10038. 126 ms later another thread sends a 38,382 byte
message with prefix ee ab 02. The receiver consumes 33,185 bytes, which
is the second message's prefix plus the first 33,182 bytes of its
payload, decodes ee ab as U+FFFD and 02 as U+0002, and aborts with
"Unexpected character encountered while parsing value" at line 0,
position 0.

Latch the first send failure, dispose the stream so the peer sees a
disconnect rather than a shifted frame, and reject later sends. Skip
BinaryWriter.Dispose after a failure because it flushes the
BufferedStream and can re-emit bytes of the incomplete frame. Run the
whole Send body under the write lock so a concurrent send cannot pass
the latch check and then write.

Dispose takes the write lock with Monitor.TryEnter rather than waiting
for it. Waiting would block shutdown behind a stuck socket write, which
is what happens today: with an in-flight write held open, the existing
Dispose blocks until that write completes because BinaryWriter.Dispose
flushes the BufferedStream underneath it. Skipping the flush when the
lock is busy is safe, since Send flushes every message.

Shutdown behavior is otherwise unchanged. Send after Dispose still
writes, and the NotSupportedException tolerance for an owner-disposed
stream stays.

Add regression tests for a prefix-only write, disposal after a buffered
failure, a concurrent send queued behind a failing one, and disposal
racing an in-flight send. All four fail on unmodified main. A fifth
test pins the existing post-Dispose behavior so this change cannot
alter it silently.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 9e03d8a7-a869-4ca9-ab36-45529635e316

Author: Damian Horna

src/Microsoft.TestPlatform.CommunicationUtilities/LengthPrefixCommunicationChannel.cs

@@ -18,10 +18,14 @@ namespace Microsoft.VisualStudio.TestPlatform.CommunicationUtilities;
 /// </summary>
 public class LengthPrefixCommunicationChannel : ICommunicationChannel
 {
+    private readonly Stream _stream;
+
     private readonly BinaryReader _reader;
 
     private readonly BinaryWriter _writer;
 
+    private volatile Exception? _sendFailure;
+
     /// <summary>
     /// Sync object for sending messages
     /// Write for binarywriter is NOT thread-safe
@@ -30,6 +34,7 @@ public class LengthPrefixCommunicationChannel : ICommunicationChannel
 
     public LengthPrefixCommunicationChannel(Stream stream)
     {
+        _stream = stream;
         _reader = new BinaryReader(stream, Encoding.UTF8, true);
 
         // Using the Buffered stream while writing, improves the write performance. By reducing the number of writes.
@@ -42,29 +47,43 @@ public LengthPrefixCommunicationChannel(Stream stream)
     /// <inheritdoc />
     public Task Send(string data)
     {
-        try
+        // Writing Message on binarywriter is not Thread-Safe
+        // Need to sync one by one to avoid buffer corruption
+        lock (_writeSyncObject)
         {
-            // Writing Message on binarywriter is not Thread-Safe
-            // Need to sync one by one to avoid buffer corruption
-            lock (_writeSyncObject)
+            if (_sendFailure is not null)
+            {
+                throw new CommunicationException("Unable to send data over channel because a previous send failed.", _sendFailure);
+            }
+
+            try
             {
                 _writer.Write(data);
                 _writer.Flush();
             }
-        }
-        catch (NotSupportedException ex) when (!_writer.BaseStream.CanWrite)
-        {
-            // As we are simply creating streams around some stream passed as ctor argument, we
-            // end up in some unsynchronized behavior where it's possible that the outside stream
-            // was disposed and we are still trying to write something. In such case we would fail
-            // with "System.NotSupportedException: Stream does not support writing.".
-            // To avoid being too generic in that catch, I am checking if the stream is not writable.
-            EqtTrace.Verbose("LengthPrefixCommunicationChannel.Send: BaseStream is not writable (most likely it was dispose). {0}", ex);
-        }
-        catch (Exception ex)
-        {
-            EqtTrace.Error("LengthPrefixCommunicationChannel.Send: Error sending data: {0}.", ex);
-            throw new CommunicationException("Unable to send data over channel.", ex);
+            catch (NotSupportedException ex) when (!_stream.CanWrite)
+            {
+                // The owner can dispose the stream while the channel is shutting down.
+                EqtTrace.Verbose("LengthPrefixCommunicationChannel.Send: BaseStream is not writable (most likely it was disposed). {0}", ex);
+            }
+            catch (Exception ex)
+            {
+                _sendFailure = ex;
+
+                try
+                {
+                    // A failed write may have left a length prefix or a partial payload on the wire.
+                    // Close the stream before another send can reuse the broken frame boundary.
+                    _stream.Dispose();
+                }
+                catch (Exception disposeException)
+                {
+                    EqtTrace.Error("LengthPrefixCommunicationChannel.Send: Error closing stream after send failure: {0}.", disposeException);
+                }
+
+                EqtTrace.Error("LengthPrefixCommunicationChannel.Send: Error sending data: {0}.", ex);
+                throw new CommunicationException("Unable to send data over channel.", ex);
+            }
         }
 
         return Task.CompletedTask;
@@ -107,16 +126,36 @@ public Task NotifyDataAvailable(CancellationToken cancellationToken)
     /// <inheritdoc />
     public void Dispose()
     {
+        // Dispose can run while another thread is inside Send. Take the write lock when it is
+        // free, so disposal never flushes the writer mid-frame, but never wait for it: blocking
+        // here would turn a stuck socket write into a shutdown hang.
+        var lockTaken = false;
         try
         {
+            Monitor.TryEnter(_writeSyncObject, ref lockTaken);
+
             EqtTrace.Verbose("LengthPrefixCommunicationChannel.Dispose: Dispose reader and writer.");
             _reader.Dispose();
-            _writer.Dispose();
+
+            // BinaryWriter.Dispose flushes its BufferedStream. That would re-emit bytes of an
+            // incomplete frame after a failed send, and race the writer while a send is in
+            // flight. Send flushes every message, so nothing is pending in either case.
+            if (lockTaken && _sendFailure is null)
+            {
+                _writer.Dispose();
+            }
         }
         catch (ObjectDisposedException)
         {
             // We don't own the underlying stream lifecycle so it's possible that it's already disposed.
         }
+        finally
+        {
+            if (lockTaken)
+            {
+                Monitor.Exit(_writeSyncObject);
+            }
+        }
 
         GC.SuppressFinalize(this);
     }

test/Microsoft.TestPlatform.CommunicationUtilities.UnitTests/LengthPrefixCommunicationChannelTests.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
+using System.Diagnostics;
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
@@ -25,6 +26,8 @@ public class LengthPrefixCommunicationChannelTests : IDisposable
 
     private readonly BinaryWriter _writer;
 
+    public TestContext TestContext { get; set; } = null!;
+
     public LengthPrefixCommunicationChannelTests()
     {
         _stream = new MemoryStream();
@@ -120,6 +123,27 @@ public void DisposeShouldNotCloseTheStream()
         Assert.IsTrue(_stream.CanWrite);
     }
 
+    /// <summary>
+    /// Characterization test. Send after Dispose writes to the stream today, because the writer
+    /// is created with leaveOpen and its disposal only flushes. Shutdown races rely on that:
+    /// Dispose runs while timer driven sends are still in flight. This test passes both before
+    /// and after the send failure latch, so it fails if that behavior is ever changed by accident.
+    /// </summary>
+    [TestMethod]
+    public async Task SendAfterDisposeShouldStillWriteToTheStream()
+    {
+        using var stream = new MemoryStream();
+        var channel = new LengthPrefixCommunicationChannel(stream);
+        channel.Dispose();
+
+        await channel.Send(Dummydata);
+
+        Assert.IsTrue(stream.CanWrite);
+        SeekToBeginning(stream);
+        using var reader = new BinaryReader(stream);
+        Assert.AreEqual(Dummydata, reader.ReadString());
+    }
+
     [TestMethod]
     public async Task DoNotFailWhenWritingOnADisposedBaseStream()
     {
@@ -138,6 +162,90 @@ public async Task DoNotFailWhenReadingFromADisposedBaseStream()
         await _channel.NotifyDataAvailable(new CancellationToken());
     }
 
+    [TestMethod]
+    public async Task SendShouldCloseStreamAndRejectLaterMessagesAfterPartialFrame()
+    {
+        using var stream = new FailingWriteStream(bytesBeforeFailure: 3);
+        var channel = new LengthPrefixCommunicationChannel(stream);
+        var message = new string('x', SocketConstants.BufferSize + 1);
+
+        await Assert.ThrowsExactlyAsync<CommunicationException>(() => channel.Send(message));
+
+        Assert.HasCount(3, stream.WrittenBytes);
+        using (var prefixStream = new MemoryStream(stream.WrittenBytes))
+        using (var prefixReader = new BinaryReader(prefixStream))
+        {
+            Assert.AreEqual(message.Length, Read7BitEncodedInt(prefixReader));
+        }
+
+        var writeCallCount = stream.WriteCallCount;
+        await Assert.ThrowsExactlyAsync<CommunicationException>(() => channel.Send(Dummydata));
+
+        Assert.IsTrue(stream.IsDisposed);
+        Assert.AreEqual(writeCallCount, stream.WriteCallCount);
+        Assert.HasCount(3, stream.WrittenBytes);
+    }
+
+    [TestMethod]
+    public async Task DisposeShouldNotFlushBufferedDataAfterSendFailure()
+    {
+        using var stream = new FailingWriteStream(bytesBeforeFailure: 0);
+        var channel = new LengthPrefixCommunicationChannel(stream);
+
+        await Assert.ThrowsExactlyAsync<CommunicationException>(() => channel.Send(Dummydata));
+        var writeCallCount = stream.WriteCallCount;
+
+        channel.Dispose();
+
+        Assert.AreEqual(writeCallCount, stream.WriteCallCount);
+    }
+
+    [TestMethod]
+    public async Task ConcurrentSendShouldNotWriteAfterFirstSendFails()
+    {
+        using var stream = new FailingWriteStream(bytesBeforeFailure: 3, blockFirstWrite: true);
+        var channel = new LengthPrefixCommunicationChannel(stream);
+        var message = new string('x', SocketConstants.BufferSize + 1);
+
+        var firstSend = Task.Run(() => channel.Send(message), TestContext.CancellationToken);
+        Assert.IsTrue(stream.WriteStarted.Wait(TimeSpan.FromSeconds(5), TestContext.CancellationToken));
+
+        var secondSend = Task.Run(() => channel.Send(Dummydata), TestContext.CancellationToken);
+        await Task.Delay(50, TestContext.CancellationToken);
+        Assert.IsFalse(secondSend.IsCompleted);
+
+        stream.ReleaseWrite.Set();
+
+        await Assert.ThrowsExactlyAsync<CommunicationException>(() => firstSend);
+        await Assert.ThrowsExactlyAsync<CommunicationException>(() => secondSend);
+        Assert.IsTrue(stream.IsDisposed);
+        Assert.AreEqual(1, stream.WriteCallCount);
+    }
+
+    [TestMethod]
+    public async Task DisposeShouldNeitherWriteNorBlockWhileASendIsInFlight()
+    {
+        using var stream = new FailingWriteStream(bytesBeforeFailure: 3, blockFirstWrite: true);
+        var channel = new LengthPrefixCommunicationChannel(stream);
+        var message = new string('x', SocketConstants.BufferSize + 1);
+
+        var send = Task.Run(() => channel.Send(message), TestContext.CancellationToken);
+        Assert.IsTrue(stream.WriteStarted.Wait(TimeSpan.FromSeconds(5), TestContext.CancellationToken));
+
+        var writeCallCount = stream.WriteCallCount;
+        var stopwatch = Stopwatch.StartNew();
+        channel.Dispose();
+        stopwatch.Stop();
+
+        // The in-flight write is held for up to five seconds. Disposal must not wait for it,
+        // and must not flush the writer underneath it.
+        Assert.IsLessThan(TimeSpan.FromSeconds(2), stopwatch.Elapsed, $"Dispose blocked for {stopwatch.Elapsed}.");
+        Assert.AreEqual(writeCallCount, stream.WriteCallCount);
+
+        stream.ReleaseWrite.Set();
+        await Assert.ThrowsExactlyAsync<CommunicationException>(() => send);
+    }
+
     // TODO
     // WriteFromMultilpleThreadShouldBeInSequence
     private static void SeekToBeginning(Stream stream)
@@ -173,4 +281,77 @@ private static int Read7BitEncodedInt(BinaryReader reader)
 
         return count;
     }
+
+    private sealed class FailingWriteStream : MemoryStream
+    {
+        private readonly int _bytesBeforeFailure;
+        private readonly bool _blockFirstWrite;
+
+        private int _acceptedBytes;
+        private bool _failureInjected;
+
+        public FailingWriteStream(int bytesBeforeFailure, bool blockFirstWrite = false)
+        {
+            _bytesBeforeFailure = bytesBeforeFailure;
+            _blockFirstWrite = blockFirstWrite;
+        }
+
+        public bool IsDisposed { get; private set; }
+
+        public int WriteCallCount { get; private set; }
+
+        public byte[] WrittenBytes => ToArray();
+
+        public ManualResetEventSlim WriteStarted { get; } = new(false);
+
+        public ManualResetEventSlim ReleaseWrite { get; } = new(false);
+
+        public override void Write(byte[] buffer, int offset, int count)
+        {
+            if (IsDisposed)
+            {
+                throw new ObjectDisposedException(nameof(FailingWriteStream));
+            }
+
+            WriteCallCount++;
+            WriteStarted.Set();
+            if (!_failureInjected)
+            {
+                if (_blockFirstWrite && !ReleaseWrite.Wait(TimeSpan.FromSeconds(5)))
+                {
+                    throw new TimeoutException("Timed out waiting to release the injected write failure.");
+                }
+
+                _failureInjected = true;
+                var remainingBytes = _bytesBeforeFailure - _acceptedBytes;
+                if (remainingBytes > 0)
+                {
+                    var acceptedCount = Math.Min(remainingBytes, count);
+                    base.Write(buffer, offset, acceptedCount);
+                    _acceptedBytes += acceptedCount;
+                }
+
+                throw new IOException("Injected write failure.");
+            }
+
+            base.Write(buffer, offset, count);
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            if (IsDisposed)
+            {
+                return;
+            }
+
+            IsDisposed = true;
+            if (disposing)
+            {
+                WriteStarted.Dispose();
+                ReleaseWrite.Dispose();
+            }
+
+            base.Dispose(disposing);
+        }
+    }
 }