Ensure calls to OOP for source-generator purposes see the exact same Solution-snapshot-instances across all calls they make

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

@@ -5,44 +5,186 @@
 using System;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.ErrorReporting;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
+using Microsoft.CodeAnalysis.Threading;
 
 namespace Microsoft.CodeAnalysis.Remote;
 
 internal interface IRemoteKeepAliveService
 {
     /// <summary>
-    /// Keeps alive this solution in the OOP process until the cancellation token is triggered.  Used so that long
-    /// running features (like 'inline rename' or 'lightbulbs') we can call into oop several times, with the same
-    /// snapshot, knowing that things will stay hydrated and alive on the OOP side.  Importantly, by keeping the
-    /// same <see cref="Solution"/> snapshot alive on the OOP side, computed attached values (like <see
-    /// cref="Compilation"/>s) will stay alive as well.
+    /// Keeps the solution identified by <paramref name="solutionChecksum"/> alive in the OOP process until <paramref
+    /// name="cancellationToken"/> is triggered. This enables long-running features (like inline rename or lightbulbs)
+    /// to make multiple OOP calls against the same snapshot, ensuring that computed values (like <see
+    /// cref="Compilation"/>s) remain cached rather than being rebuilt on each call.
     /// </summary>
-    ValueTask KeepAliveAsync(Checksum solutionChecksum, CancellationToken cancellationToken);
+    /// <param name="solutionChecksum">Checksum identifying the solution to pin.</param>
+    /// <param name="sessionId">Unique identifier for this session. The host uses this with <see
+    /// cref="WaitForSessionIdAsync"/> to block until the solution is actually pinned before proceeding with
+    /// dependent work.</param>
+    /// <param name="cancellationToken">Cancellation of this token releases the pinned solution.</param>
+    ValueTask KeepAliveAsync(Checksum solutionChecksum, long sessionId, CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Blocks until the session identified by <paramref name="sessionId"/> has fully synced and pinned its solution
+    /// in the OOP process. This ensures the host doesn't proceed with OOP calls until the solution is guaranteed to
+    /// be available.
+    /// </summary>
+    ValueTask WaitForSessionIdAsync(long sessionId, CancellationToken cancellationToken);
 }
 
 internal sealed class RemoteKeepAliveSession : IDisposable
 {
-    private readonly CancellationTokenSource _cancellationTokenSource = new();
+    private static long s_sessionId = 1;
+
+    /// <summary>
+    /// Unique identifier for this session. Used to coordinate between <see cref="IRemoteKeepAliveService.KeepAliveAsync"/>
+    /// (which syncs and pins the solution) and <see cref="IRemoteKeepAliveService.WaitForSessionIdAsync"/> (which blocks
+    /// until pinning completes).
+    /// </summary>
+    private long SessionId { get; } = Interlocked.Increment(ref s_sessionId);
 
-    private RemoteKeepAliveSession(
+    /// <summary>
+    /// Controls the lifetime of the OOP-side pinning. The <see cref="IRemoteKeepAliveService.KeepAliveAsync"/> call
+    /// blocks on this token; canceling it allows that call to return, releasing the pinned solution.
+    /// </summary>
+    private CancellationTokenSource KeepAliveTokenSource { get; } = new();
+
+    private RemoteKeepAliveSession()
+    {
+    }
+
+    /// <summary>
+    /// Creates and fully establishes a keep-alive session. Returns only after the solution is confirmed to be
+    /// pinned on the OOP side.
+    /// </summary>
+    /// <remarks>
+    /// <para>This method coordinates two concurrent OOP calls:</para>
+    /// <list type="number">
+    /// <item><see cref="IRemoteKeepAliveService.KeepAliveAsync"/>: Syncs the solution to OOP, then blocks until
+    /// <see cref="KeepAliveTokenSource"/> is canceled. This call is fire-and-forget from the host's perspective.</item>
+    /// <item><see cref="IRemoteKeepAliveService.WaitForSessionIdAsync"/>: Blocks until KeepAliveAsync has completed
+    /// syncing. This call is awaited, ensuring the solution is pinned before returning to the caller.</item>
+    /// </list>
+    /// <para>The two calls share <see cref="SessionId"/> so the OOP side can correlate them.</para>
+    /// </remarks>
+    private static async Task<RemoteKeepAliveSession> StartSessionAsync(
         SolutionCompilationState compilationState,
-        RemoteHostClient? client)
+        ProjectId? projectId,
+        RemoteHostClient? client,
+        CancellationToken callerCancellationToken)
     {
+        var session = new RemoteKeepAliveSession();
+
+        // When running in-process (no OOP client), return immediately. The caller holds the solution snapshot
+        // directly, so no pinning is needed.
         if (client is null)
-            return;
+            return session;
+
+        // Fire-and-forget: Start syncing and pinning the solution on the OOP side. This call will block on the OOP
+        // side until KeepAliveTokenSource is canceled (i.e., when this session is Disposed).
+        //
+        // Important: We pass KeepAliveTokenSource.Token (not callerCancellationToken) because:
+        // - The keep-alive must persist beyond this method, for the lifetime of the session
+        // - Disposing the session is what should cancel this work, not the caller's token
+        _ = InvokeKeepAliveAsync(compilationState, projectId, client, session);
+
+        // Block until the OOP side confirms the solution is pinned. This uses callerCancellationToken so the caller
+        // can abandon the wait if they no longer need the session.
+        await WaitForSessionIdAsync(compilationState, projectId, client, session, callerCancellationToken).ConfigureAwait(false);
+
+        return session;
 
-        // Now kick off the keep-alive work.  We don't wait on this as this will stick on the OOP side until
-        // the cancellation token triggers.
-        _ = client.TryInvokeAsync<IRemoteKeepAliveService>(
-            compilationState,
-            (service, solutionInfo, cancellationToken) => service.KeepAliveAsync(solutionInfo, cancellationToken),
-            _cancellationTokenSource.Token).AsTask();
+        static async Task InvokeKeepAliveAsync(
+            SolutionCompilationState compilationState,
+            ProjectId? projectId,
+            RemoteHostClient client,
+            RemoteKeepAliveSession session)
+        {
+            try
+            {
+                // Yield to allow StartSessionAsync to proceed to WaitForSessionIdAsync concurrently.
+                await Task.Yield().ConfigureAwait(false);
+
+                var sessionId = session.SessionId;
+                await client.TryInvokeAsync<IRemoteKeepAliveService>(
+                   compilationState,
+                   projectId,
+                   (service, solutionInfo, cancellationToken) => service.KeepAliveAsync(solutionInfo, sessionId, cancellationToken),
+                   session.KeepAliveTokenSource.Token).ConfigureAwait(false);
+            }
+            catch (Exception ex) when (FatalError.ReportAndCatchUnlessCanceled(ex))
+            {
+                // Non-cancellation exceptions indicate a catastrophic failure (e.g., broken OOP connection).
+                // We must dispose the session to:
+                // 1. Cancel KeepAliveTokenSource, which is linked into WaitForSessionIdAsync's token
+                // 2. Unblock WaitForSessionIdAsync so it doesn't hang forever
+                //
+                // Cancellation exceptions are expected and normal - they occur when the session is properly
+                // disposed, which cancels KeepAliveTokenSource and allows KeepAliveAsync to return.
+                session.Dispose();
+
+                // Don't rethrow: this is fire-and-forget. Errors were already reported via FatalError above.
+            }
+        }
+
+        static async Task WaitForSessionIdAsync(
+            SolutionCompilationState compilationState,
+            ProjectId? projectId,
+            RemoteHostClient client,
+            RemoteKeepAliveSession session,
+            CancellationToken callerCancellationToken)
+        {
+            try
+            {
+                // Link both cancellation sources so this call aborts if either:
+                // - The caller cancels (they no longer need the session)
+                // - InvokeKeepAliveAsync fails and disposes the session (which cancels KeepAliveTokenSource)
+                //
+                // Without the link to KeepAliveTokenSource, a failure in InvokeKeepAliveAsync would leave this
+                // call hanging indefinitely since the OOP side would never signal completion.
+                using var linkedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(
+                    session.KeepAliveTokenSource.Token,
+                    callerCancellationToken);
+
+                await client.TryInvokeAsync<IRemoteKeepAliveService>(
+                    compilationState,
+                    projectId,
+                    (service, _, cancellationToken) => service.WaitForSessionIdAsync(session.SessionId, cancellationToken),
+                    linkedTokenSource.Token).ConfigureAwait(false);
+            }
+            catch
+            {
+                // Any failure means we can't establish the session. The three lines below handle all cases:
+                //
+                // 1. Dispose(): Always clean up. The caller won't receive the session, so we must release it.
+                //    This also cancels KeepAliveTokenSource, which unblocks InvokeKeepAliveAsync if it's still running.
+                //
+                // 2. ThrowIfCancellationRequested(): If the caller's token caused the cancellation, rethrow with
+                //    that token to maintain proper cancellation semantics (the exception's CancellationToken
+                //    property should match what the caller passed in). This is a no-op if the caller didn't cancel.
+                //
+                // 3. throw: For all other failures (e.g., KeepAliveTokenSource canceled due to InvokeKeepAliveAsync
+                //    failing, or a non-cancellation exception), propagate the original exception.
+                session.Dispose();
+                callerCancellationToken.ThrowIfCancellationRequested();
+                throw;
+            }
+        }
     }
 
+    /// <summary>
+    /// Constructor for synchronous, best-effort session creation. Does not wait for the session to be established.
+    /// </summary>
     private RemoteKeepAliveSession(SolutionCompilationState compilationState, IAsynchronousOperationListener listener)
     {
-        var cancellationToken = _cancellationTokenSource.Token;
+        // Unlike the async entry-point, this constructor returns immediately without waiting for the solution to
+        // be pinned on the OOP side. This is acceptable for scenarios where:
+        // - The caller cannot await (e.g., in a constructor)
+        // - Best-effort pinning is sufficient (subsequent OOP calls will still work, just potentially slower)
+        //
+        // Track the async work so test infrastructure can detect outstanding operations.
         var token = listener.BeginAsyncOperation(nameof(RemoteKeepAliveSession));
 
         var task = CreateClientAndKeepAliveAsync();
@@ -52,16 +194,25 @@ private RemoteKeepAliveSession(SolutionCompilationState compilationState, IAsync
 
         async Task CreateClientAndKeepAliveAsync()
         {
+            var cancellationToken = this.KeepAliveTokenSource.Token;
             var client = await RemoteHostClient.TryGetClientAsync(compilationState.Services, cancellationToken).ConfigureAwait(false);
             if (client is null)
                 return;
 
-            // Now kick off the keep-alive work.  We don't wait on this as this will stick on the OOP side until
-            // the cancellation token triggers.
-            _ = client.TryInvokeAsync<IRemoteKeepAliveService>(
-                compilationState,
-                (service, solutionInfo, cancellationToken) => service.KeepAliveAsync(solutionInfo, cancellationToken),
-                cancellationToken).AsTask();
+            // Fire-and-forget: Start the keep-alive without waiting for confirmation. Unlike StartSessionAsync,
+            // we don't call WaitForSessionIdAsync because this is a best-effort, non-blocking path.
+            try
+            {
+                var sessionId = this.SessionId;
+                await client.TryInvokeAsync<IRemoteKeepAliveService>(
+                    compilationState,
+                    projectId: null,
+                    (service, solutionInfo, cancellationToken) => service.KeepAliveAsync(solutionInfo, sessionId, cancellationToken),
+                    cancellationToken).ConfigureAwait(false);
+            }
+            catch (Exception ex) when (FatalError.ReportAndCatchUnlessCanceled(ex))
+            {
+            }
         }
     }
 
@@ -76,43 +227,54 @@ async Task CreateClientAndKeepAliveAsync()
     public void Dispose()
     {
         GC.SuppressFinalize(this);
-        _cancellationTokenSource.Cancel();
-        _cancellationTokenSource.Dispose();
+
+        // Cancel rather than dispose the token source. CancellationTokenSource.Dispose() is only necessary when
+        // not canceling (to clean up internal wait handles), but we always cancel. The finalizer's Contract.Fail
+        // will catch any cases where Dispose is forgotten.
+        this.KeepAliveTokenSource.Cancel();
     }
 
     /// <summary>
-    /// Creates a session between the host and OOP, effectively pinning this <paramref name="solution"/> until <see
-    /// cref="IDisposable.Dispose"/> is called on it.  By pinning the solution we ensure that all calls to OOP for
-    /// the same solution during the life of this session do not need to resync the solution.  Nor do they then need
-    /// to rebuild any compilations they've already built due to the solution going away and then coming back.
+    /// Creates a best-effort keep-alive session synchronously. Returns immediately without waiting for the session
+    /// to be established on the OOP side.
     /// </summary>
     /// <remarks>
-    /// The <paramref name="listener"/> is not strictly necessary for this type.  This class functions just as an
-    /// optimization to hold onto data so it isn't resync'ed or recomputed.  However, we still want to let the
-    /// system know when unobserved async work is kicked off in case we have any tooling that keep track of this for
-    /// any reason (for example for tracking down problems in testing scenarios).
-    /// </remarks>
-    /// <remarks>
-    /// This synchronous entrypoint should be used only in contexts where using the async <see
-    /// cref="CreateAsync(Solution, CancellationToken)"/> is not possible (for example, in a constructor).
+    /// <para>Use this overload only when async code is not possible (e.g., in constructors). For guaranteed session
+    /// establishment, use <see cref="CreateAsync(Solution, CancellationToken)"/> instead.</para>
+    /// <para>Because this method doesn't wait for establishment, subsequent OOP calls may not benefit from the
+    /// pinned solution if they race ahead of the keep-alive setup. In practice this is rare, but callers requiring
+    /// guaranteed consistency must use the async overloads.</para>
+    /// <para>The <paramref name="listener"/> is used to track the async keep-alive work for testing infrastructure.</para>
     /// </remarks>
     public static RemoteKeepAliveSession Create(Solution solution, IAsynchronousOperationListener listener)
         => new(solution.CompilationState, listener);
 
     /// <summary>
-    /// Creates a session between the host and OOP, effectively pinning this <paramref name="solution"/> until <see
-    /// cref="IDisposable.Dispose"/> is called on it.  By pinning the solution we ensure that all calls to OOP for
-    /// the same solution during the life of this session do not need to resync the solution.  Nor do they then need
-    /// to rebuild any compilations they've already built due to the solution going away and then coming back.
+    /// Creates a keep-alive session, returning only after the session is fully established on the OOP side.
     /// </summary>
+    /// <remarks>
+    /// All subsequent OOP calls made while this session is alive will see the same pinned solution instance,
+    /// provided they pass matching solution/project-cone data. Mismatched calls (e.g., session created for full
+    /// solution but call made for project-cone) will not benefit from the pinning.
+    /// </remarks>
     public static Task<RemoteKeepAliveSession> CreateAsync(Solution solution, CancellationToken cancellationToken)
-        => CreateAsync(solution.CompilationState, cancellationToken);
+        => CreateAsync(solution, projectId: null, cancellationToken);
+
+    /// <inheritdoc cref="CreateAsync(Solution, CancellationToken)"/>
+    public static Task<RemoteKeepAliveSession> CreateAsync(Solution solution, ProjectId? projectId, CancellationToken cancellationToken)
+        => CreateAsync(solution.CompilationState, projectId, cancellationToken);
+
+    /// <inheritdoc cref="CreateAsync(Solution, CancellationToken)"/>
+    public static Task<RemoteKeepAliveSession> CreateAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken)
+        => CreateAsync(compilationState, projectId: null, cancellationToken);
 
     /// <inheritdoc cref="CreateAsync(Solution, CancellationToken)"/>
     public static async Task<RemoteKeepAliveSession> CreateAsync(
-        SolutionCompilationState compilationState, CancellationToken cancellationToken)
+        SolutionCompilationState compilationState, ProjectId? projectId, CancellationToken cancellationToken)
     {
-        var client = await RemoteHostClient.TryGetClientAsync(compilationState.Services, cancellationToken).ConfigureAwait(false);
-        return new RemoteKeepAliveSession(compilationState, client);
+        var client = await RemoteHostClient.TryGetClientAsync(
+            compilationState.Services, cancellationToken).ConfigureAwait(false);
+
+        return await StartSessionAsync(compilationState, projectId, client, cancellationToken).ConfigureAwait(false);
     }
 }

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

@@ -105,17 +105,24 @@ public ValueTask<bool> TryInvokeAsync<TService>(
         CancellationToken cancellationToken)
         where TService : class
     {
-        return TryInvokeAsync(solution.CompilationState, invocation, cancellationToken);
+        return TryInvokeAsync(solution.CompilationState, projectId: null, invocation, cancellationToken);
     }
 
+    /// <param name="projectId">If <see langword="null"/> the entire solution snapshot represented by <paramref
+    /// name="compilationState"/> will be synchronized with the OOP side.  If not <see langword="null"/> only the
+    /// project-cone represented by that id will be synchronized over.</param>
+    /// <returns></returns>
     public async ValueTask<bool> TryInvokeAsync<TService>(
         SolutionCompilationState compilationState,
+        ProjectId? projectId,
         Func<TService, Checksum, CancellationToken, ValueTask> invocation,
         CancellationToken cancellationToken)
         where TService : class
     {
         using var connection = CreateConnection<TService>(callbackTarget: null);
-        return await connection.TryInvokeAsync(compilationState, invocation, cancellationToken).ConfigureAwait(false);
+        return projectId is null
+            ? await connection.TryInvokeAsync(compilationState, invocation, cancellationToken).ConfigureAwait(false)
+            : await connection.TryInvokeAsync(compilationState, projectId, invocation, cancellationToken).ConfigureAwait(false);
     }
 
     public async ValueTask<Optional<TResult>> TryInvokeAsync<TService, TResult>(