Fix connection leaking on cancellation (#26178)

* add callstack on debug to make investigation easier

* if GetPinnedScopeAsync throws cancellation exception, connection can be leaked not disposed.

fixed the issue and added test

* make sure cancellation is thrown in test

* PR feedbacks

* more feedbacks

* address PR feedback

src/VisualStudio/Core/Test.Next/Services/VisualStudioDiagnosticAnalyzerExecutorTests.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.IO;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -20,6 +22,7 @@
 using Microsoft.CodeAnalysis.VisualBasic.UseNullPropagation;
 using Microsoft.CodeAnalysis.Workspaces.Diagnostics;
 using Microsoft.VisualStudio.LanguageServices.Remote;
+using Roslyn.Test.Utilities;
 using Roslyn.Utilities;
 using Roslyn.VisualStudio.Next.UnitTests.Mocks;
 using Xunit;
@@ -121,6 +124,29 @@ public async Task TestCancellation()
             }
         }
 
+        [Fact, Trait(Traits.Feature, Traits.Features.RemoteHost)]
+        [WorkItem(26178, "https://github.com/dotnet/roslyn/pull/26178")]
+        public async Task TestCancellationOnSessionWithSolution()
+        {
+            var code = @"class Test { void Method() { } }";
+
+            using (var workspace = CreateWorkspace(LanguageNames.CSharp, code))
+            {
+                var solution = workspace.CurrentSolution;
+                var solutionChecksum = await solution.State.GetChecksumAsync(CancellationToken.None);
+
+                var source = new CancellationTokenSource();
+                var connection = new InvokeThrowsCancellationConnection(source);
+                var exception = await Assert.ThrowsAnyAsync<OperationCanceledException>(() => SessionWithSolution.CreateAsync(connection, solution, source.Token));
+                Assert.Equal(exception.CancellationToken, source.Token);
+
+                // make sure things that should have been cleaned up are cleaned up
+                var service = (RemotableDataServiceFactory.Service)solution.Workspace.Services.GetService<IRemotableDataService>();
+                Assert.Null(service.GetRemotableData_TestOnly(solutionChecksum, CancellationToken.None));
+                Assert.True(connection.Disposed);
+            }
+        }
+
         [Fact, Trait(Traits.Feature, Traits.Features.RemoteHost)]
         public async Task TestHostAnalyzers()
         {
@@ -293,5 +319,45 @@ public MyUpdateSource(Workspace workspace)
 
             public override Workspace Workspace => _workspace;
         }
+
+        private class InvokeThrowsCancellationConnection : RemoteHostClient.Connection
+        {
+            private readonly CancellationTokenSource _source;
+
+            public bool Disposed = false;
+
+            public InvokeThrowsCancellationConnection(CancellationTokenSource source)
+            {
+                _source = source;
+            }
+
+            public override Task InvokeAsync(string targetName, IReadOnlyList<object> arguments, CancellationToken cancellationToken)
+            {
+                // cancel and throw cancellation exception
+                _source.Cancel();
+                _source.Token.ThrowIfCancellationRequested();
+
+                throw Utilities.ExceptionUtilities.Unreachable;
+            }
+
+            public override Task<T> InvokeAsync<T>(
+                string targetName, IReadOnlyList<object> arguments, CancellationToken cancellationToken)
+                => throw new NotImplementedException();
+
+            public override Task InvokeAsync(
+                string targetName, IReadOnlyList<object> arguments, Func<Stream, CancellationToken, Task> funcWithDirectStreamAsync, CancellationToken cancellationToken)
+                => throw new NotImplementedException();
+
+            public override Task<T> InvokeAsync<T>(
+                string targetName, IReadOnlyList<object> arguments, Func<Stream, CancellationToken, Task<T>> funcWithDirectStreamAsync, CancellationToken cancellationToken)
+                => throw new NotImplementedException();
+
+            protected override void Dispose(bool disposing)
+            {
+                base.Dispose(disposing);
+
+                Disposed = true;
+            }
+        }
     }
 }

src/Workspaces/Core/Portable/Remote/RemoteHostClient.cs

@@ -101,6 +101,9 @@ public abstract class Connection : IDisposable
 
             protected Connection()
             {
+#if DEBUG
+                _creationCallStack = Environment.StackTrace;
+#endif
                 _disposed = false;
             }
 
@@ -128,16 +131,22 @@ public void Dispose()
             }
 
 #if DEBUG
+            private readonly string _creationCallStack;
+#endif
             ~Connection()
             {
                 // this can happen if someone kills OOP. 
                 // when that happen, we don't want to crash VS, so this is debug only check
                 if (!Environment.HasShutdownStarted)
                 {
-                    Contract.Requires(false, $@"Should have been disposed!");
+#if DEBUG
+                    Contract.Fail($"Should have been disposed!\r\n {_creationCallStack}");
+#else
+                    Contract.Fail($"Should have been disposed!");
+#endif
                 }
             }
-#endif
+
         }
     }
 }

src/Workspaces/Core/Portable/Remote/RemoteHostClientExtensions.cs

@@ -45,14 +45,13 @@ internal static class RemoteHostClientExtensions
         public static async Task<SessionWithSolution> TryCreateSessionAsync(
             this RemoteHostClient client, string serviceName, Solution solution, object callbackTarget, CancellationToken cancellationToken)
         {
-            var session = await client.TryCreateConnectionAsync(serviceName, callbackTarget, cancellationToken).ConfigureAwait(false);
-            if (session == null)
+            var connection = await client.TryCreateConnectionAsync(serviceName, callbackTarget, cancellationToken).ConfigureAwait(false);
+            if (connection == null)
             {
                 return null;
             }
 
-            var scope = await GetPinnedScopeAsync(solution, cancellationToken).ConfigureAwait(false);
-            return await SessionWithSolution.CreateAsync(session, scope, cancellationToken).ConfigureAwait(false);
+            return await SessionWithSolution.CreateAsync(connection, solution, cancellationToken).ConfigureAwait(false);
         }
 
         /// <summary>

src/Workspaces/Core/Portable/Remote/RemoteHostSessionHelpers.cs

@@ -19,12 +19,16 @@ internal sealed class SessionWithSolution : IDisposable
         private readonly RemoteHostClient.Connection _connection;
         private readonly PinnedRemotableDataScope _scope;
 
-        public static async Task<SessionWithSolution> CreateAsync(RemoteHostClient.Connection connection, PinnedRemotableDataScope scope, CancellationToken cancellationToken)
+        public static async Task<SessionWithSolution> CreateAsync(RemoteHostClient.Connection connection, Solution solution, CancellationToken cancellationToken)
         {
-            var sessionWithSolution = new SessionWithSolution(connection, scope);
+            Contract.ThrowIfNull(connection);
+            Contract.ThrowIfNull(solution);
 
+            PinnedRemotableDataScope scope = null;
             try
             {
+                scope = await solution.GetPinnedScopeAsync(cancellationToken).ConfigureAwait(false);
+
                 // set connection state for this session.
                 // we might remove this in future. see https://github.com/dotnet/roslyn/issues/24836
                 await connection.InvokeAsync(
@@ -32,11 +36,14 @@ public static async Task<SessionWithSolution> CreateAsync(RemoteHostClient.Conne
                     new object[] { scope.SolutionInfo },
                     cancellationToken).ConfigureAwait(false);
 
-                return sessionWithSolution;
+                return new SessionWithSolution(connection, scope);
             }
             catch
             {
-                sessionWithSolution.Dispose();
+                // make sure disposable objects are disposed when
+                // exceptions are thrown
+                connection.Dispose();
+                scope?.Dispose();
 
                 // we only expect this to happen on cancellation. otherwise, rethrow
                 cancellationToken.ThrowIfCancellationRequested();