Cancellation management in InitializeApplicationAsync

src/Microsoft.VisualStudio.ProjectSystem.Managed.VS/ProjectSystem/VS/HotReload/ProjectHotReloadSessionManager.cs

@@ -73,20 +73,7 @@ async ValueTask<bool> TryCreatePendingSessionInternalAsync()
 
                 if (await ProjectSupportsStartupHooksAsync())
                 {
-                    HotReloadSessionState hotReloadSessionState = new((HotReloadSessionState sessionState) =>
-                    {
-                        int count;
-                        lock (_activeSessionStates)
-                        {
-                            Assumes.True(_activeSessionStates.Remove(sessionState), "Cannot remove unknown hot reload session.");
-                            count = _activeSessionStates.Count;
-                        }
-
-                        if (count == 0)
-                        {
-                            _threadingService.ExecuteSynchronously(() => _projectHotReloadNotificationService.Value.SetHotReloadStateAsync(isInHotReload: false));
-                        }
-                    }, _threadingService);
+                    HotReloadSessionState hotReloadSessionState = new(RemoveSessionState, _threadingService);
 
                     IProjectHotReloadSession projectHotReloadSession = _hotReloadAgent.CreateHotReloadSession(
                         name: projectName,
@@ -97,7 +84,7 @@ async ValueTask<bool> TryCreatePendingSessionInternalAsync()
                         debugLaunchOptions: launchOptions);
 
                     hotReloadSessionState.Session = projectHotReloadSession;
-                    await projectHotReloadSession.ApplyLaunchVariablesAsync(environmentVariables, default);
+                    await projectHotReloadSession.ApplyLaunchVariablesAsync(environmentVariables, hotReloadSessionState.CancellationToken);
 
                     _pendingSessionState = hotReloadSessionState;
 
@@ -106,14 +93,13 @@ async ValueTask<bool> TryCreatePendingSessionInternalAsync()
                 else
                 {
                     // If startup hooks are not supported then tell the user why Hot Reload isn't available.
-
                     WriteOutputMessage(
                         new HotReloadLogMessage(
                             HotReloadVerbosity.Minimal,
                             Resources.ProjectHotReloadSessionManager_StartupHooksDisabled,
                             projectName,
                             null,
-                            HotReloadDiagnosticOutputService.GetProcessId(),
+                            0,
                             HotReloadDiagnosticErrorLevel.Warning
                         ));
                 }
@@ -124,6 +110,26 @@ async ValueTask<bool> TryCreatePendingSessionInternalAsync()
         }
     }
 
+    private void RemoveSessionState(HotReloadSessionState sessionState)
+    {
+        _threadingService.RunAndForget(async () =>
+        {
+            DebugTrace("Disposing Hot Reload session.");
+
+            int count;
+            lock (_activeSessionStates)
+            {
+                Assumes.True(_activeSessionStates.Remove(sessionState));
+                count = _activeSessionStates.Count;
+            }
+
+            if (count == 0)
+            {
+                await _projectHotReloadNotificationService.Value.SetHotReloadStateAsync(isInHotReload: false);
+            }
+        }, _unconfiguredProject);
+    }
+
     protected override Task InitializeCoreAsync(CancellationToken cancellationToken)
     {
         return Task.CompletedTask;
@@ -135,12 +141,15 @@ protected override Task DisposeCoreAsync(bool initialized)
 
         Task DisposeCoreInternalAsync()
         {
-            foreach (HotReloadSessionState sessionState in _activeSessionStates)
+            lock (_activeSessionStates)
             {
-                sessionState.Process?.Dispose();
-            }
+                foreach (HotReloadSessionState sessionState in _activeSessionStates)
+                {
+                    DisposeSessionStateAndStopSession(sessionState);
+                }
 
-            _activeSessionStates.Clear();
+                _activeSessionStates.Clear();
+            }
 
             return Task.CompletedTask;
         }
@@ -178,7 +187,7 @@ async Task ApplyHotReloadUpdateInternalAsync()
             {
                 cancellationToken.ThrowIfCancellationRequested();
 
-                if (sessionState.Session is IProjectHotReloadSessionInternal { DeltaApplier: {} deltaApplier })
+                if (sessionState.Session is IProjectHotReloadSessionInternal { DeltaApplier: { } deltaApplier })
                 {
                     updateTasks ??= [];
                     updateTasks.Add(applyFunction(deltaApplier, cancellationToken));
@@ -199,44 +208,99 @@ public Task ActivateSessionAsync(IVsLaunchedProcess? launchedProcess, VsDebugTar
 
         async Task ActivateSessionInternalAsync()
         {
-            if (_pendingSessionState is { Session: IProjectHotReloadSession session })
+            // _pendingSessionState can be null if project doesn't support Hot Reload. i.e doesn't have SupportsHotReload capability
+            HotReloadSessionState? sessionState = Interlocked.Exchange(ref _pendingSessionState, null);
+            if (sessionState is null)
             {
-                Process process = Process.GetProcessById((int)vsDebugTargetProcessInfo.dwProcessId);
-                _pendingSessionState.LaunchedProcess = launchedProcess;
-                _pendingSessionState.Process = process;
+                return;
+            }
 
-                if (!process.HasExited)
-                {
-                    process.Exited += (sender, e) =>
-                    {
-                        _threadingService.ExecuteSynchronously(() => session.StopSessionAsync(CancellationToken.None));
-                    };
-                    process.EnableRaisingEvents = true;
+            DebugTrace("Hot Reload session started.");
+            lock (_activeSessionStates)
+            {
+                _activeSessionStates.Add(sessionState);
+            }
 
-                    await _pendingSessionState.Session.StartSessionAsync(CancellationToken.None);
-                    lock (_activeSessionStates)
-                    {
-                        _activeSessionStates.Add(_pendingSessionState);
-                    }
+            Process? process = null;
+            try
+            {
+                sessionState.DebuggerProcess = launchedProcess;
+                process = Process.GetProcessById((int)vsDebugTargetProcessInfo.dwProcessId);
+                sessionState.Process = process;
+            }
+            catch (ArgumentException)
+            {
+                // process might have been exited in some cases.
+                // in that case, we early return without starting hotreload session
+                // one way to mimic this is to hit control + C as fast as you can once hit F5/Control + F5
+                DisposeSessionStateAndStopSession(sessionState);
+                return;
+            }
+
+            process.EnableRaisingEvents = true;
+            process.Exited += (sender, e) =>
+            {
+                DebugTrace("Process exited");
+                DisposeSessionStateAndStopSession(sessionState);
+            };
+
+            if (process.HasExited)
+            {
+                DebugTrace("Process exited");
+                DisposeSessionStateAndStopSession(sessionState);
+            }
+            else
+            {
+                try
+                {
+                    await sessionState.Session.StartSessionAsync(sessionState.CancellationToken);
                     await _projectHotReloadNotificationService.Value.SetHotReloadStateAsync(isInHotReload: true);
                 }
-
-                _pendingSessionState = null;
+                catch (OperationCanceledException)
+                {
+                    DisposeSessionStateAndStopSession(sessionState);
+                }
             }
         }
     }
 
-    private sealed class HotReloadSessionState(
-        Action<HotReloadSessionState> removeSessionState,
-        IProjectThreadingService threadingService) : IProjectHotReloadSessionCallback
+    private void DisposeSessionStateAndStopSession(HotReloadSessionState sessionState)
     {
+        sessionState.Dispose();
+
+        // In some occasions, StopSessionAsync might be invoked before StartSessionAsync
+        // For example, if the process exits quickly after launch
+        // So we call StopSessionAsync unconditionally to ensure the session is stopped properly
+        _threadingService.ExecuteSynchronously(() => sessionState.Session.StopSessionAsync(CancellationToken.None));
+    }
+
+    private sealed class HotReloadSessionState : IProjectHotReloadSessionCallback, IDisposable
+    {
+        private int _disposed = 0;
+
+        private readonly CancellationTokenSource _cancellationTokenSource = new();
+        private readonly Action<HotReloadSessionState> _removeSessionState;
+        private readonly IProjectThreadingService _threadingService;
+
+        public HotReloadSessionState(
+            Action<HotReloadSessionState> removeSessionState,
+            IProjectThreadingService threadingService)
+        {
+            _removeSessionState = removeSessionState;
+            _threadingService = threadingService;
+            CancellationToken = _cancellationTokenSource.Token;
+        }
+
+        public CancellationToken CancellationToken { get; }
+
+        [Obsolete]
         public bool SupportsRestart => Session is not null;
 
-        public IProjectHotReloadSession? Session { get; set; }
+        public IProjectHotReloadSession Session { get => field ?? throw Assumes.NotReachable(); set; }
 
         public Process? Process { get; set; }
 
-        public IVsLaunchedProcess? LaunchedProcess { get; set; }
+        public IVsLaunchedProcess? DebuggerProcess { get; set; }
 
         public IDeltaApplier? GetDeltaApplier()
         {
@@ -248,39 +312,57 @@ public Task OnAfterChangesAppliedAsync(CancellationToken cancellationToken)
             return Task.CompletedTask;
         }
 
+        [Obsolete]
         public Task<bool> RestartProjectAsync(CancellationToken cancellationToken)
         {
             return TaskResult.True;
         }
 
         public async Task<bool> StopProjectAsync(CancellationToken cancellationToken)
         {
-            if (Session is null || (LaunchedProcess is null && Process is null))
+            if (DebuggerProcess is not null && Process is not null)
             {
-                // if session is null, or there's no active process asscociated with this session, early return
-                return true;
+                // We have both DebuggerProcess and Process, they point to the same process. But DebuggerProcess provides a nicer way to terminate process
+                // without affecting the entire debug session.
+                // So we prefer to use DebuggerProcess to terminate the process first.
+
+                await TerminateProcessGracefullyAsync();
+
+                // When DebuggerProcess.Terminate(ignoreLaunchFlags: 1) return, the process might not be terminated
+                // So we first terminate the process nicely,
+                // Then wait for the process to exit. If the process doesn't exit within 500ms, kill it using traditional way.
+                await Process.WaitForExitAsync(default).WithTimeout(TimeSpan.FromMilliseconds(500));
             }
 
-            // prefer to terminate launched process first if we have it
-            if (LaunchedProcess is not null)
+            if (Process is not null)
             {
-                // need to call on UI thread
-                await threadingService.SwitchToUIThread(cancellationToken);
+                TerminateProcess(Process);
+            }
+
+            Dispose();
+
+            return true;
+
+            async Task TerminateProcessGracefullyAsync()
+            {
+                // Terminate DebuggerProcess need to call on UI thread
+                await _threadingService.SwitchToUIThread(CancellationToken.None);
 
                 // Ignore the debug option launching flags since we're just terminating the process, not the entire debug session
-                LaunchedProcess.Terminate(ignoreLaunchFlags: 1);
+                // TODO consider if we can use the return value of Terminate here to control whether we need to subsequently kill the process
+                DebuggerProcess.Terminate(ignoreLaunchFlags: 1);
             }
-            else
+
+            static void TerminateProcess(Process process)
             {
-                // stop the process by killing it
                 try
                 {
-                    if (Process is not null && !Process.HasExited)
+                    if (!process.HasExited)
                     {
                         // First try to close the process nicely and if that doesn't work kill it.
-                        if (!Process.CloseMainWindow())
+                        if (!process.CloseMainWindow())
                         {
-                            Process.Kill();
+                            process.Kill();
                         }
                     }
                 }
@@ -289,16 +371,35 @@ public async Task<bool> StopProjectAsync(CancellationToken cancellationToken)
                     // Process has already exited.
                 }
             }
+        }
 
-            if (Session is not null)
+        public void Dispose()
+        {
+            if (Interlocked.Exchange(ref _disposed, 1) == 1)
             {
-                await Session.StopSessionAsync(cancellationToken);
-                removeSessionState(this);
-
-                Session = null;
+                return;
             }
 
-            return true;
+            _cancellationTokenSource.Cancel();
+            _cancellationTokenSource.Dispose();
+
+            Process?.Dispose();
+
+            _removeSessionState(this);
         }
     }
+
+    private void DebugTrace(string message)
+    {
+#if DEBUG
+        var projectName = _unconfiguredProject.GetProjectName();
+        _hotReloadDiagnosticOutputService.Value.WriteLine(
+            new HotReloadLogMessage(
+                HotReloadVerbosity.Detailed,
+                message,
+                projectName,
+                errorLevel: HotReloadDiagnosticErrorLevel.Info),
+            CancellationToken.None);
+#endif
+    }
 }

src/Microsoft.VisualStudio.ProjectSystem.Managed/ProjectSystem/HotReload/DeltaApplier.cs

@@ -3,6 +3,7 @@
 using Microsoft.DotNet.HotReload;
 using Microsoft.VisualStudio.Debugger.Contracts.HotReload;
 using Microsoft.VisualStudio.HotReload.Components.DeltaApplier;
+using Microsoft.VisualStudio.Threading;
 
 namespace Microsoft.VisualStudio.ProjectSystem.HotReload;
 
@@ -33,7 +34,10 @@ public async ValueTask<ImmutableArray<string>> GetCapabilitiesAsync(Cancellation
 
     public async ValueTask InitializeApplicationAsync(CancellationToken cancellationToken)
     {
-        _ = await client.GetUpdateCapabilitiesAsync(cancellationToken);
+        // Not all clients respond correctly to cancellation tokens.
+        // For example, `DefaultHotreloadClient.GetUpdateCapabilitiesAsync(ct)`doesn't listen to the passed ct.
+        // Since DefaultHotreloadClient is defined in a source package, it can't be modified directly from project-system.
+        await client.GetUpdateCapabilitiesAsync(cancellationToken).WithCancellation(cancellationToken);
 
         // TODO: apply initial updates?
         // https://devdiv.visualstudio.com/DevDiv/_workitems/edit/2571676

src/Microsoft.VisualStudio.ProjectSystem.Managed/ProjectSystem/HotReload/HotReloadDiagnosticOutputService.cs

@@ -1,6 +1,5 @@
 // Licensed to the .NET Foundation under one or more agreements. The .NET Foundation licenses this file to you under the MIT license. See the LICENSE.md file in the project root for more information.
 
-using System.Diagnostics;
 using Microsoft.VisualStudio.Debugger.Contracts.HotReload;
 
 namespace Microsoft.VisualStudio.ProjectSystem.HotReload;
@@ -22,9 +21,4 @@ public void WriteLine(HotReloadLogMessage hotReloadLogMessage, CancellationToken
     {
         _threadingService.RunAndForget(() => _hotReloadLogger.LogAsync(hotReloadLogMessage, cancellationToken).AsTask(), unconfiguredProject: null);
     }
-
-    public static uint GetProcessId(Process? process = null)
-    {
-        return (uint)(process?.Id ?? 0);
-    }
 }

src/Microsoft.VisualStudio.ProjectSystem.Managed/ProjectSystem/HotReload/ProjectHotReloadAgent.cs

@@ -12,6 +12,7 @@ internal sealed class ProjectHotReloadAgent(
     Lazy<IHotReloadAgentManagerClient> hotReloadAgentManagerClient,
     Lazy<IHotReloadDiagnosticOutputService> hotReloadDiagnosticOutputService,
     IProjectSystemOptions projectSystemOptions,
+    IProjectThreadingService threadingService,
     [Import(AllowDefault = true)] IHotReloadDebugStateProvider? debugStateProvider) // allow default until VS Code is updated: https://devdiv.visualstudio.com/DevDiv/_workitems/edit/2571211
         : IProjectHotReloadAgent 
 {
@@ -35,6 +36,7 @@ public IProjectHotReloadSession CreateHotReloadSession(
             launchProfile: launchProfile,
             debugLaunchOptions: debugLaunchOptions,
             projectSystemOptions,
+            threadingService,
             debugStateProvider ?? DefaultDebugStateProvider.Instance);
     }