Fix up delta apply acknowledgement and add logging (dotnet#17053) (do…

…tnet#17391)

* Fix up delta apply acknowledgement and add logging
* Apply suggestions from code review

Co-authored-by: Pranav K <prkrishn@hotmail.com>
Co-authored-by: Pranav K <prkrishn@hotmail.com>

Co-authored-by: Safia Abdalla <safia@microsoft.com>

src/BuiltInTools/BrowserRefresh/WebSocketScriptInjection.js

@@ -153,14 +153,10 @@ setTimeout(function () {
   }
 
   function sendDeltaApplied() {
-    const webAssemblyHotReloadSuccessPayload = new ArrayBuffer(1);
-    webAssemblyHotReloadSuccessPayload[0] = 1;
-    connection.send(webAssemblyHotReloadSuccessPayload);
+    connection.send(new Uint8Array([1]).buffer);
   }
 
   function sendDeltaNotApplied() {
-    const webAssemblyHotReloadFailurePayload = new ArrayBuffer(1);
-    webAssemblyHotReloadFailurePayload[0] = 0;
-    connection.send(webAssemblyHotReloadFailurePayload);
+    connection.send(new Uint8Array([0]).buffer);
   }
 }, 500);

src/BuiltInTools/dotnet-watch/Filters/BrowserRefreshServer.cs

@@ -32,8 +32,6 @@ public class BrowserRefreshServer : IAsyncDisposable
         private readonly TaskCompletionSource _taskCompletionSource;
         private IHost _refreshServer;
 
-        private const int MESSAGE_TIMEOUT = 5000;
-
         public BrowserRefreshServer(IReporter reporter)
         {
             _reporter = reporter;
@@ -143,9 +141,6 @@ public async ValueTask DisposeAsync()
 
         public async ValueTask<ValueWebSocketReceiveResult?> ReceiveAsync(Memory<byte> buffer, CancellationToken cancellationToken)
         {
-            var linkedCancellationToken = CancellationTokenSource.CreateLinkedTokenSource(
-                cancellationToken,
-                new CancellationTokenSource(MESSAGE_TIMEOUT).Token);
             for (int i = 0; i < _clientSockets.Count; i++)
             {
                 var clientSocket = _clientSockets[i];
@@ -157,7 +152,7 @@ public async ValueTask DisposeAsync()
 
                 try
                 {
-                    return await clientSocket.ReceiveAsync(buffer, linkedCancellationToken.Token);
+                    return await clientSocket.ReceiveAsync(buffer, cancellationToken);
                 }
                 catch (Exception ex)
                 {

src/BuiltInTools/dotnet-watch/HotReload/BlazorWebAssemblyDeltaApplier.cs

@@ -18,6 +18,8 @@ internal class BlazorWebAssemblyDeltaApplier : IDeltaApplier
         private readonly IReporter _reporter;
         private int _sequenceId;
 
+        private static readonly TimeSpan MESSAGE_TIMEOUT = TimeSpan.FromSeconds(5);
+
         public BlazorWebAssemblyDeltaApplier(IReporter reporter)
         {
             _reporter = reporter;
@@ -51,7 +53,7 @@ public async ValueTask<bool> Apply(DotNetWatchContext context, string changedFil
 
             await context.BrowserRefreshServer.SendJsonSerlialized(payload, cancellationToken);
 
-            return await VerifyDeltaApplied(context, cancellationToken);
+            return await VerifyDeltaApplied(context, cancellationToken).WaitAsync(MESSAGE_TIMEOUT);
         }
 
         public async ValueTask ReportDiagnosticsAsync(DotNetWatchContext context, IEnumerable<string> diagnostics, CancellationToken cancellationToken)
@@ -67,28 +69,47 @@ public async ValueTask ReportDiagnosticsAsync(DotNetWatchContext context, IEnume
             }
         }
 
-        private async ValueTask<bool> VerifyDeltaApplied(DotNetWatchContext context, CancellationToken cancellationToken)
+        private async Task<bool> VerifyDeltaApplied(DotNetWatchContext context, CancellationToken cancellationToken)
         {
             var _receiveBuffer = new byte[1];
             try 
             {
-                var result = await context.BrowserRefreshServer.ReceiveAsync(_receiveBuffer, cancellationToken);
+                ValueWebSocketReceiveResult? result;
+                // Sometimes we can received a Close message from the WebSocket before the
+                // browser has transmitted the acknowledgement. To address this, we wait until
+                // a result has been received and it has the expected message type (binary).
+                // There is a `WaitAsync` at the invocation site of `VerifyDeltaApplied` so this
+                // will time out if we have received an appropriate message within 5 seconds.
+                do
+                {
+                    result = await context.BrowserRefreshServer.ReceiveAsync(_receiveBuffer, cancellationToken);
+                } while (!result.HasValue || result?.MessageType is not WebSocketMessageType.Binary);
                 return IsDeltaApplied(result);
             }
-            catch when (!cancellationToken.IsCancellationRequested)
+            catch (TaskCanceledException)
             {
                 _reporter.Verbose("Timed out while waiting to verify delta was applied.");
                 return false;
             }
 
             bool IsDeltaApplied(ValueWebSocketReceiveResult? result)
             {
+                _reporter.Verbose($"Received {_receiveBuffer[0]} from browser in {Stringify(result)}.");
                 return result.HasValue
                     && result.Value.Count == 1 // Should have received 1 byte on the socket for the acknowledgement
                     && result.Value.MessageType is WebSocketMessageType.Binary 
                     && result.Value.EndOfMessage
                     && _receiveBuffer[0] == 1;
             }
+
+            static string Stringify(ValueWebSocketReceiveResult? result)
+            {
+                if (result is ValueWebSocketReceiveResult r)
+                {
+                    return $"Count: {r.Count}, MessageType: {r.MessageType}, EndOfMessage: {r.EndOfMessage}";
+                }
+                return "no result received.";
+            }
         }
 
         public void Dispose()

src/BuiltInTools/dotnet-watch/HotReload/CompilationHandler.cs

@@ -118,10 +118,12 @@ public async ValueTask<bool> TryHandleFileChange(DotNetWatchContext context, Fil
                 var diagnostics = GetDiagnostics(updatedSolution, cancellationToken);
                 if (diagnostics.IsDefaultOrEmpty)
                 {
+                    _reporter.Verbose("No deltas modified. Applying changes to clear diagnostics.");
                     await _deltaApplier.Apply(context, file.FilePath, updates, cancellationToken);
                 }
                 else
                 {
+                    _reporter.Verbose("Found compilation errors during hot reload. Reporting it in application UI.");
                     await _deltaApplier.ReportDiagnosticsAsync(context, diagnostics, cancellationToken);
                 }
 
@@ -146,6 +148,7 @@ public async ValueTask<bool> TryHandleFileChange(DotNetWatchContext context, Fil
             _currentSolution = updatedSolution;
 
             var applyState = await _deltaApplier.Apply(context, file.FilePath, updates, cancellationToken);
+            _reporter.Verbose($"Received {(applyState ? "successful" : "failed")} apply from delta applier.");
             HotReloadEventSource.Log.HotReloadEnd(HotReloadEventSource.StartType.CompilationHandler);
             return applyState;
         }

src/BuiltInTools/dotnet-watch/HotReloadDotNetWatcher.cs

@@ -170,8 +170,9 @@ public async Task WatchAsync(DotNetWatchContext context, CancellationToken cance
                         context.ChangedFile = changedFile;
                     }
                 }
-                catch
+                catch (Exception e)
                 {
+                    _reporter.Verbose($"Caught top-level exception from hot reload: {e}");
                     if (!currentRunCancellationSource.IsCancellationRequested)
                     {
                         currentRunCancellationSource.Cancel();