Fix a race condition in MessageWriter.WriteMessage

This change fixes a race condition in MessageWriter.WriteMessage where
multiple calls to the method would end up writing at the same time.  This
is due to the use of 'await' calls and happens even when both calls occur
on the same synchronization context.  The subsequent calls proceed while
the first call is awaiting the completion of I/O writes.

The solution is to use an AsyncLock to synchronize I/O operations while
allowing more messages to be queued up.

Author: daviwil

src/PowerShellEditorServices.Host/MessageLoop.cs

@@ -31,7 +31,8 @@ public class MessageLoop : IHost
         private EditorSession editorSession;
         private MessageReader messageReader;
         private MessageWriter messageWriter;
-        private SynchronizationContext syncContext;
+        private SynchronizationContext applicationSyncContext;
+        private SynchronizationContext messageLoopSyncContext;
         private AsyncContextThread messageLoopThread;
 
         #endregion
@@ -61,7 +62,7 @@ public void Start()
         private async Task StartMessageLoop()
         {
             // Hold on to the current synchronization context
-            this.syncContext = SynchronizationContext.Current;
+            this.applicationSyncContext = SynchronizationContext.Current;
 
             // Start the message listener on another thread
             this.messageLoopThread = new AsyncContextThread(true);
@@ -70,6 +71,8 @@ private async Task StartMessageLoop()
 
         async Task ListenForMessages()
         {
+            this.messageLoopSyncContext = SynchronizationContext.Current;
+
             // Ensure that the console is using UTF-8 encoding
             System.Console.InputEncoding = Encoding.UTF8;
             System.Console.OutputEncoding = Encoding.UTF8;
@@ -161,21 +164,26 @@ await this.messageWriter.WriteMessage(
 
         async void DebugService_DebuggerStopped(object sender, DebuggerStopEventArgs e)
         {
-            await this.messageWriter.WriteMessage(
-                new StoppedEvent
+            // Push the write operation to the correct thread
+            this.messageLoopSyncContext.Post(
+                async (obj) =>
                 {
-                    Body = new StoppedEventBody
-                    {
-                        Source = new Source
+                    await this.messageWriter.WriteMessage(
+                        new StoppedEvent
                         {
-                            Path = e.InvocationInfo.ScriptName,
-                        },
-                        Line = e.InvocationInfo.ScriptLineNumber,
-                        Column = e.InvocationInfo.OffsetInLine,
-                        ThreadId = 1, // TODO: Change this based on context
-                        Reason = "breakpoint" // TODO: Change this based on context
-                    }
-                });
+                            Body = new StoppedEventBody
+                            {
+                                Source = new Source
+                                {
+                                    Path = e.InvocationInfo.ScriptName,
+                                },
+                                Line = e.InvocationInfo.ScriptLineNumber,
+                                Column = e.InvocationInfo.OffsetInLine,
+                                ThreadId = 1, // TODO: Change this based on context
+                                Reason = "breakpoint" // TODO: Change this based on context
+                            }
+                        });
+                }, null);
         }
 
         void PowerShellSession_BreakpointUpdated(object sender, BreakpointUpdatedEventArgs e)

src/PowerShellEditorServices.Transport.Stdio/Message/MessageWriter.cs

@@ -6,6 +6,7 @@
 using Microsoft.PowerShell.EditorServices.Utility;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
+using Nito.AsyncEx;
 using System.IO;
 using System.Text;
 using System.Threading.Tasks;
@@ -18,6 +19,7 @@ public class MessageWriter
 
         private Stream outputStream;
         private MessageTypeResolver messageTypeResolver;
+        private AsyncLock writeLock = new AsyncLock();
 
         private JsonSerializer loggingSerializer = 
             JsonSerializer.Create(
@@ -80,10 +82,16 @@ public async Task WriteMessage(MessageBase messageToWrite)
                         Constants.ContentLengthFormatString,
                         messageBytes.Length));
 
-            // Send the message
-            await this.outputStream.WriteAsync(headerBytes, 0, headerBytes.Length);
-            await this.outputStream.WriteAsync(messageBytes, 0, messageBytes.Length);
-            await this.outputStream.FlushAsync();
+            // Make sure only one call is writing at a time.  You might be thinking
+            // "Why not use a normal lock?"  We use an AsyncLock here so that the
+            // message loop doesn't get blocked while waiting for I/O to complete.
+            using (await this.writeLock.LockAsync())
+            {
+                // Send the message
+                await this.outputStream.WriteAsync(headerBytes, 0, headerBytes.Length);
+                await this.outputStream.WriteAsync(messageBytes, 0, messageBytes.Length);
+                await this.outputStream.FlushAsync();
+            }
         }
 
         #endregion