Address deadlock and race conditions in WebSocketConnection

Directory.Build.props

@@ -67,7 +67,9 @@
 
     <!-- TFMs -->
     <SdkTargetFramework>$(NetCurrent)</SdkTargetFramework>
-    <VisualStudioServiceTargetFramework>net9.0</VisualStudioServiceTargetFramework>
+
+    <!-- TFM for projects that are deployed to run on .NET shipping with Visual Studio -->
+    <VisualStudioServiceTargetFramework>net10.0</VisualStudioServiceTargetFramework>
 
     <!-- NU1507 Disable multi-feed check as .NET uses multiple internal feeds intentionally -->
     <!-- NU5039 Disable NuGet is unable to find the readme file in the package. -->

src/Dotnet.Watch/AspireService/AspireServerService.cs

@@ -395,14 +395,16 @@ private async ValueTask<bool> SendMessageAsync(string dcpId, byte[] messageBytes
         }
 
         var success = false;
+        var lockAcquired = false;
         try
         {
             using var cancelTokenSource = CancellationTokenSource.CreateLinkedTokenSource(
                 cancellationToken, _shutdownCancellationSource.Token, connection.HttpRequestAborted);
 
             await _webSocketAccess.WaitAsync(cancelTokenSource.Token);
-            await connection.Socket.SendAsync(new ArraySegment<byte>(messageBytes), WebSocketMessageType.Text, endOfMessage: true, cancelTokenSource.Token);
+            lockAcquired = true;
 
+            await connection.Socket.SendAsync(new ArraySegment<byte>(messageBytes), WebSocketMessageType.Text, endOfMessage: true, cancelTokenSource.Token);
             success = true;
         }
         finally
@@ -413,7 +415,10 @@ private async ValueTask<bool> SendMessageAsync(string dcpId, byte[] messageBytes
                 _socketConnectionManager.RemoveSocketConnection(connection);
             }
 
-            _webSocketAccess.Release();
+            if (lockAcquired)
+            {
+                _webSocketAccess.Release();
+            }
         }
 
         return success;

src/Dotnet.Watch/AspireService/Helpers/ImmutableInterlockedExtensions.cs

@@ -0,0 +1,38 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Text;
+using System.Threading;
+
+namespace Aspire.Tools.Service;
+
+internal static class ImmutableInterlockedExtensions
+{
+    extension(ImmutableInterlocked)
+    {
+        public static (T oldValue, T newValue) Transform<T>(ref T location, Func<T, T> transformer) where T : class?
+        {
+            T oldValue = Volatile.Read(ref location);
+            while (true)
+            {
+                T newValue = transformer(oldValue);
+                if (ReferenceEquals(oldValue, newValue))
+                {
+                    // No change was actually required.
+                    return (oldValue, newValue);
+                }
+
+                T interlockedResult = Interlocked.CompareExchange(ref location, newValue, oldValue);
+                if (ReferenceEquals(oldValue, interlockedResult))
+                {
+                    return (oldValue, newValue);
+                }
+
+                oldValue = interlockedResult; // we already have a volatile read that we can reuse for the next loop
+            }
+        }
+    }
+}

src/Dotnet.Watch/AspireService/Helpers/SocketConnectionManager.cs

@@ -5,6 +5,8 @@
 
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Net.WebSockets;
 using System.Threading;
 using System.Threading.Tasks;
@@ -17,66 +19,52 @@ namespace Aspire.Tools.Service;
 /// </summary>
 internal class SocketConnectionManager : IDisposable
 {
-    // Track a single connection per Dcp ID
-    private readonly object _socketConnectionsLock = new();
-    private readonly Dictionary<string, WebSocketConnection> _webSocketConnections = new(StringComparer.Ordinal);
+    // Track a single connection per DCP ID
+    private ImmutableDictionary<string, WebSocketConnection> _webSocketConnections =
+        ImmutableDictionary<string, WebSocketConnection>.Empty;
 
     private void CleanupSocketConnections()
     {
-        lock (_socketConnectionsLock)
-        {
-            foreach (var connection in _webSocketConnections)
-            {
-                connection.Value.Tcs.SetResult();
-                connection.Value.CancelTokenRegistration.Dispose();
-            }
+        var connections = Interlocked.Exchange(ref _webSocketConnections, ImmutableDictionary<string, WebSocketConnection>.Empty);
 
-            _webSocketConnections.Clear();
+        foreach (var (_, connection) in connections)
+        {
+            connection.Dispose();
         }
     }
 
     public void AddSocketConnection(WebSocket socket, TaskCompletionSource tcs, string dcpId, CancellationToken httpRequestAborted)
     {
-        // We only support one connection per DCP Id, therefore if there is
+        // We only support one connection per DCP ID, therefore if there is
         // already a connection, drop that one before adding this one
-        lock (_socketConnectionsLock)
-        {
-            if (_webSocketConnections.TryGetValue(dcpId, out var existingConnection))
-            {
-                _webSocketConnections.Remove(dcpId);
-                existingConnection.Dispose();
-            }
 
-            // Register with the cancel token so that if the socket goes bad, we
-            // get notified and can remove it from our list. We need to track the registrations as well
-            // so we can dispose of it later
-            var newConnection = new WebSocketConnection(socket, tcs, dcpId, httpRequestAborted);
-            newConnection.CancelTokenRegistration = httpRequestAborted.Register(() =>
-            {
-                RemoveSocketConnection(newConnection);
-            });
+        var newConnection = new WebSocketConnection(socket, tcs, dcpId, httpRequestAborted);
+
+        var (oldConnections, _) = ImmutableInterlocked.Transform(ref _webSocketConnections, connections => connections.SetItem(dcpId, newConnection));
 
-            _webSocketConnections[dcpId] = newConnection;
+        if (oldConnections.TryGetValue(dcpId, out var oldConnection))
+        {
+            oldConnection.Dispose();
         }
+
+        // Hook up removal from tracked connections on abort after the connection has been added:
+        newConnection.RegisterCancellationCallback(RemoveSocketConnection);
     }
 
     public void RemoveSocketConnection(WebSocketConnection connection)
     {
-        lock (_socketConnectionsLock)
+        // If the connection is not in the dictionary, then it has already been removed and disposed or replaced with another connection.
+        if (ImmutableInterlocked.Update(ref _webSocketConnections,
+            connections => connections.TryGetValue(connection.DcpId, out var currentConnection) && currentConnection == connection
+                ? connections.Remove(connection.DcpId)
+                : connections))
         {
-            _webSocketConnections.Remove(connection.DcpId);
             connection.Dispose();
         }
     }
 
     public WebSocketConnection? GetSocketConnection(string dcpId)
-    {
-        lock (_socketConnectionsLock)
-        {
-            _webSocketConnections.TryGetValue(dcpId, out var connection);
-            return connection;
-        }
-    }
+        => _webSocketConnections.GetValueOrDefault(dcpId);
 
     public void Dispose()
     {

src/Dotnet.Watch/AspireService/Helpers/WebSocketConnection.cs

@@ -13,25 +13,56 @@ namespace Aspire.Tools.Service;
 /// <summary>
 /// Used by the SocketConnectionManager to track one socket connection. It needs to be disposed when done with it
 /// </summary>
-internal class WebSocketConnection : IDisposable
+internal sealed class WebSocketConnection(WebSocket socket, TaskCompletionSource tcs, string dcpId, CancellationToken httpRequestAborted) : IDisposable
 {
-    public WebSocketConnection(WebSocket socket, TaskCompletionSource tcs, string dcpId, CancellationToken httpRequestAborted)
-    {
-        Socket = socket;
-        Tcs = tcs;
-        DcpId = dcpId;
-        HttpRequestAborted = httpRequestAborted;
-    }
+    public WebSocket Socket { get; } = socket;
+    public TaskCompletionSource Tcs { get; } = tcs;
+    public string DcpId { get; } = dcpId;
+    public CancellationToken HttpRequestAborted { get; } = httpRequestAborted;
 
-    public WebSocket Socket { get; }
-    public TaskCompletionSource Tcs { get; }
-    public string DcpId { get; }
-    public CancellationToken HttpRequestAborted { get; }
-    public CancellationTokenRegistration CancelTokenRegistration { get; set; }
+    private readonly Lock _cancelTokenRegistrationLock = new();
+    private CancellationTokenRegistration? _cancelTokenRegistration;
+    private bool _isDisposed;
 
     public void Dispose()
     {
-       Tcs.SetResult();
-       CancelTokenRegistration.Dispose();
+        Tcs.TrySetResult();
+
+        CancellationTokenRegistration? registrationToDispose = null;
+        lock (_cancelTokenRegistrationLock)
+        {
+            if (!_isDisposed)
+            {
+                _isDisposed = true;
+                registrationToDispose = _cancelTokenRegistration;
+                _cancelTokenRegistration = null;
+            }
+        }
+
+        // The callback might be called during disposal, do so outside of the lock:
+        registrationToDispose?.Dispose();
+    }
+
+    public void RegisterCancellationCallback(Action<WebSocketConnection> callback)
+    {
+        // Note that the callback can be called synchronously before Register returns.
+        var cancelTokenRegistration = HttpRequestAborted.Register(() => callback(this));
+
+        bool disposeRegistration;
+        lock (_cancelTokenRegistrationLock)
+        {
+            disposeRegistration = _isDisposed;
+
+            if (!disposeRegistration)
+            {
+                _cancelTokenRegistration = cancelTokenRegistration;
+            }
+        }
+
+        if (disposeRegistration)
+        {
+            // The callback might be called during disposal, do so outside of the lock:
+            cancelTokenRegistration.Dispose();
+        }
     }
 }

src/Dotnet.Watch/AspireService/Microsoft.WebTools.AspireService.Package.csproj

@@ -1,6 +1,6 @@
 <Project Sdk="Microsoft.NET.Sdk">
   <PropertyGroup>
-    <!-- This source package is used by Visual Studio WebTools -->
+    <!-- This source package is used by Visual Studio Project System for an OOP service -->
     <TargetFramework>$(VisualStudioServiceTargetFramework)</TargetFramework>
     <GenerateDocumentationFile>false</GenerateDocumentationFile>
     <DebugType>none</DebugType>

src/Dotnet.Watch/AspireService/Microsoft.WebTools.AspireService.projitems

@@ -9,18 +9,6 @@
     <Import_RootNamespace>Microsoft.WebTools.AspireService</Import_RootNamespace>
   </PropertyGroup>
   <ItemGroup>
-    <Compile Include="$(MSBuildThisFileDirectory)AspireServerService.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Contracts\IAspireServerEvents.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Helpers\CertGenerator.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Helpers\ExceptionExtensions.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Helpers\HttpContextExtensions.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Helpers\LoggerProvider.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Helpers\SocketConnectionManager.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Helpers\SocketUtilities.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Helpers\WebSocketConnection.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Models\ErrorResponse.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Models\InfoResponse.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Models\RunSessionRequest.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Models\SessionChangeNotification.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)**\*.cs" />
   </ItemGroup>
 </Project>