Stop sending after a partial frame write

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 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,32 +47,74 @@ 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 (!_stream.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: Stream is not writable (most likely it was disposed). {0}", ex);
+            }
+            catch (Exception ex)
+            {
+                FaultChannel(ex);
+                throw new CommunicationException("Unable to send data over channel.", ex);
+            }
         }
-        catch (NotSupportedException ex) when (!_writer.BaseStream.CanWrite)
+
+        return Task.CompletedTask;
+    }
+
+    /// <summary>
+    /// Marks the channel unusable after a failed write, and closes the transport.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// <see cref="BinaryWriter.Write(string)"/> does not report how many bytes it wrote, so once it
+    /// throws we cannot know whether the length prefix reached the peer without its payload. If it
+    /// did, the peer is waiting for a body that will never arrive, and every later message would be
+    /// consumed as part of that frame.
+    /// </para>
+    /// <para>
+    /// There is no way to tell the peer in band, because the framing it is reading is the thing that
+    /// broke. Closing the transport is the only signal left: the peer reads end of stream, which it
+    /// already handles.
+    /// </para>
+    /// <para>
+    /// Both callers, SocketClient and SocketServer, build this channel over a TcpClient stream and
+    /// close that client in their own error paths, so this closes a stream that was already on its
+    /// way down. Dispose still leaves the stream open, which DisposeShouldNotCloseTheStream pins.
+    /// </para>
+    /// </remarks>
+    private void FaultChannel(Exception ex)
+    {
+        _sendFailure = ex;
+
+        try
         {
-            // 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);
+            _stream.Dispose();
         }
-        catch (Exception ex)
+        catch (Exception disposeException)
         {
-            EqtTrace.Error("LengthPrefixCommunicationChannel.Send: Error sending data: {0}.", ex);
-            throw new CommunicationException("Unable to send data over channel.", ex);
+            EqtTrace.Error("LengthPrefixCommunicationChannel.Send: Error closing stream after send failure: {0}.", disposeException);
         }
 
-        return Task.CompletedTask;
+        EqtTrace.Error("LengthPrefixCommunicationChannel.Send: Error sending data: {0}.", ex);
     }
 
     /// <inheritdoc />
@@ -107,16 +154,37 @@ 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. Skip that flush after a failed
+            // send, where the buffer can still hold part of an incomplete frame that must not
+            // reach the peer, and skip it while a send holds the lock, to avoid racing that
+            // writer. Neither skip loses data: Send flushes after every message it completes.
+            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);
     }