Split query execution into compile and execute calls

src/Features/Core/Portable/SemanticSearch/AbstractSemanticSearchService.cs

@@ -54,6 +54,19 @@ protected override IntPtr LoadUnmanagedDll(string unmanagedDllName)
             => IntPtr.Zero;
     }
 
+    private readonly struct CompiledQuery(MemoryStream peStream, MemoryStream pdbStream, SourceText text) : IDisposable
+    {
+        public MemoryStream PEStream { get; } = peStream;
+        public MemoryStream PdbStream { get; } = pdbStream;
+        public SourceText Text { get; } = text;
+
+        public void Dispose()
+        {
+            PEStream.Dispose();
+            PdbStream.Dispose();
+        }
+    }
+
     /// <summary>
     /// Mapping from the parameter type of the <c>Find</c> method to the <see cref="QueryKind"/> value.
     /// </summary>
@@ -66,76 +79,93 @@ protected override IntPtr LoadUnmanagedDll(string unmanagedDllName)
         .Add(typeof(IPropertySymbol), QueryKind.Property)
         .Add(typeof(IEventSymbol), QueryKind.Event);
 
+    private ImmutableDictionary<CompiledQueryId, CompiledQuery> _compiledQueries = ImmutableDictionary<CompiledQueryId, CompiledQuery>.Empty;
+
     protected abstract Compilation CreateCompilation(SourceText query, IEnumerable<MetadataReference> references, SolutionServices services, out SyntaxTree queryTree, CancellationToken cancellationToken);
 
-    public async Task<ExecuteQueryResult> ExecuteQueryAsync(
-        Solution solution,
+    public CompileQueryResult CompileQuery(
+        SolutionServices services,
         string query,
         string referenceAssembliesDir,
-        ISemanticSearchResultsObserver observer,
-        OptionsProvider<ClassificationOptions> classificationOptions,
         TraceSource traceSource,
         CancellationToken cancellationToken)
     {
-        try
-        {
-            // add progress items - one for compilation, one for emit and one for each project:
-            var remainingProgressItemCount = 2 + solution.ProjectIds.Count;
-            await observer.AddItemsAsync(remainingProgressItemCount, cancellationToken).ConfigureAwait(false);
-
-            var metadataService = solution.Services.GetRequiredService<IMetadataService>();
-            var metadataReferences = SemanticSearchUtilities.GetMetadataReferences(metadataService, referenceAssembliesDir);
-            var queryText = SemanticSearchUtilities.CreateSourceText(query);
-            var queryCompilation = CreateCompilation(queryText, metadataReferences, solution.Services, out var queryTree, cancellationToken);
+        var metadataService = services.GetRequiredService<IMetadataService>();
+        var metadataReferences = SemanticSearchUtilities.GetMetadataReferences(metadataService, referenceAssembliesDir);
+        var queryText = SemanticSearchUtilities.CreateSourceText(query);
+        var queryCompilation = CreateCompilation(queryText, metadataReferences, services, out var queryTree, cancellationToken);
 
-            cancellationToken.ThrowIfCancellationRequested();
+        cancellationToken.ThrowIfCancellationRequested();
 
-            // complete compilation progress item:
-            remainingProgressItemCount--;
-            await observer.ItemsCompletedAsync(1, cancellationToken).ConfigureAwait(false);
+        var emitOptions = new EmitOptions(
+            debugInformationFormat: DebugInformationFormat.PortablePdb,
+            instrumentationKinds: [InstrumentationKind.StackOverflowProbing, InstrumentationKind.ModuleCancellation]);
 
-            var emitOptions = new EmitOptions(
-                debugInformationFormat: DebugInformationFormat.PortablePdb,
-                instrumentationKinds: [InstrumentationKind.StackOverflowProbing, InstrumentationKind.ModuleCancellation]);
+        var peStream = new MemoryStream();
+        var pdbStream = new MemoryStream();
 
-            using var peStream = new MemoryStream();
-            using var pdbStream = new MemoryStream();
-
-            var emitDifferenceTimer = SharedStopwatch.StartNew();
-            var emitResult = queryCompilation.Emit(peStream, pdbStream, options: emitOptions, cancellationToken: cancellationToken);
-            var emitTime = emitDifferenceTimer.Elapsed;
-
-            var executionTime = TimeSpan.Zero;
+        var emitDifferenceTimer = SharedStopwatch.StartNew();
+        var emitResult = queryCompilation.Emit(peStream, pdbStream, options: emitOptions, cancellationToken: cancellationToken);
+        var emitTime = emitDifferenceTimer.Elapsed;
 
-            cancellationToken.ThrowIfCancellationRequested();
+        CompiledQueryId queryId;
+        ImmutableArray<QueryCompilationError> errors;
+        if (emitResult.Success)
+        {
+            queryId = CompiledQueryId.Create(queryCompilation.Language);
+            Contract.ThrowIfFalse(ImmutableInterlocked.TryAdd(ref _compiledQueries, queryId, new CompiledQuery(peStream, pdbStream, queryText)));
 
-            // complete compilation progress item:
-            remainingProgressItemCount--;
-            await observer.ItemsCompletedAsync(1, cancellationToken).ConfigureAwait(false);
+            errors = [];
+        }
+        else
+        {
+            queryId = default;
 
-            if (!emitResult.Success)
+            foreach (var diagnostic in emitResult.Diagnostics)
             {
-                foreach (var diagnostic in emitResult.Diagnostics)
+                if (diagnostic.Severity == DiagnosticSeverity.Error)
                 {
-                    if (diagnostic.Severity == DiagnosticSeverity.Error)
-                    {
-                        traceSource.TraceInformation($"Semantic search query compilation failed: {diagnostic}");
-                    }
+                    traceSource.TraceInformation($"Semantic search query compilation failed: {diagnostic}");
                 }
+            }
 
-                var errors = emitResult.Diagnostics.SelectAsArray(
-                        d => d.Severity == DiagnosticSeverity.Error,
-                        d => new QueryCompilationError(d.Id, d.GetMessage(), (d.Location.SourceTree == queryTree) ? d.Location.SourceSpan : default));
+            errors = emitResult.Diagnostics.SelectAsArray(
+                d => d.Severity == DiagnosticSeverity.Error,
+                d => new QueryCompilationError(d.Id, d.GetMessage(), (d.Location.SourceTree == queryTree) ? d.Location.SourceSpan : default));
+        }
 
-                return CreateResult(errors, FeaturesResources.Semantic_search_query_failed_to_compile);
-            }
+        return new CompileQueryResult(queryId, errors, emitTime);
+    }
+
+    public void DiscardQuery(CompiledQueryId queryId)
+    {
+        Contract.ThrowIfFalse(ImmutableInterlocked.TryRemove(ref _compiledQueries, queryId, out var compiledQuery));
+        compiledQuery.Dispose();
+    }
 
-            peStream.Position = 0;
-            pdbStream.Position = 0;
+    public async Task<ExecuteQueryResult> ExecuteQueryAsync(
+        Solution solution,
+        CompiledQueryId queryId,
+        ISemanticSearchResultsObserver observer,
+        OptionsProvider<ClassificationOptions> classificationOptions,
+        TraceSource traceSource,
+        CancellationToken cancellationToken)
+    {
+        Contract.ThrowIfFalse(ImmutableInterlocked.TryRemove(ref _compiledQueries, queryId, out var query));
+
+        try
+        {
+            var executionTime = TimeSpan.Zero;
+
+            var remainingProgressItemCount = solution.ProjectIds.Count;
+            await observer.AddItemsAsync(remainingProgressItemCount, cancellationToken).ConfigureAwait(false);
+
+            query.PEStream.Position = 0;
+            query.PdbStream.Position = 0;
             var loadContext = new LoadContext();
             try
             {
-                var queryAssembly = loadContext.LoadFromStream(peStream, pdbStream);
+                var queryAssembly = loadContext.LoadFromStream(query.PEStream, query.PdbStream);
                 SetModuleCancellationToken(queryAssembly, cancellationToken);
 
                 SetToolImplementations(
@@ -146,17 +176,17 @@ public async Task<ExecuteQueryResult> ExecuteQueryAsync(
                 if (!TryGetFindMethod(queryAssembly, out var findMethod, out var queryKind, out var errorMessage, out var errorMessageArgs))
                 {
                     traceSource.TraceInformation($"Semantic search failed: {errorMessage}");
-                    return CreateResult(compilationErrors: [], errorMessage, errorMessageArgs);
+                    return CreateResult(errorMessage, errorMessageArgs);
                 }
 
-                var invocationContext = new QueryExecutionContext(queryText, findMethod, observer, classificationOptions, traceSource);
+                var invocationContext = new QueryExecutionContext(query.Text, findMethod, observer, classificationOptions, traceSource);
                 try
                 {
                     await invocationContext.InvokeAsync(solution, queryKind, cancellationToken).ConfigureAwait(false);
 
                     if (invocationContext.TerminatedWithException)
                     {
-                        return CreateResult(compilationErrors: [], FeaturesResources.Semantic_search_query_terminated_with_exception);
+                        return CreateResult(FeaturesResources.Semantic_search_query_terminated_with_exception);
                     }
                 }
                 finally
@@ -176,15 +206,19 @@ public async Task<ExecuteQueryResult> ExecuteQueryAsync(
                 }
             }
 
-            return CreateResult(compilationErrors: [], errorMessage: null);
+            return CreateResult(errorMessage: null);
 
-            ExecuteQueryResult CreateResult(ImmutableArray<QueryCompilationError> compilationErrors, string? errorMessage, params string[]? args)
-                => new(compilationErrors, errorMessage, args, emitTime, executionTime);
+            ExecuteQueryResult CreateResult(string? errorMessage, params string[]? args)
+                => new(errorMessage, args, executionTime);
         }
         catch (Exception e) when (FatalError.ReportAndPropagateUnlessCanceled(e, cancellationToken, ErrorSeverity.Critical))
         {
             throw ExceptionUtilities.Unreachable();
         }
+        finally
+        {
+            query.Dispose();
+        }
     }
 
     private static void SetModuleCancellationToken(Assembly queryAssembly, CancellationToken cancellationToken)

src/Features/Core/Portable/SemanticSearch/ExecuteQueryResult.cs

@@ -5,25 +5,59 @@
 using System;
 using System.Collections.Immutable;
 using System.Runtime.Serialization;
+using System.Threading;
 
 namespace Microsoft.CodeAnalysis.SemanticSearch;
 
 /// <summary>
 /// The result of Semantic Search query execution.
 /// </summary>
-/// <param name="compilationErrors">Compilation errors.</param>
 /// <param name="ErrorMessage">An error message if the execution failed.</param>
 /// <param name="ErrorMessageArgs">
 /// Arguments to be substituted to <paramref name="ErrorMessage"/>.
 /// Use when the values may contain PII that needs to be obscured in telemetry.
 /// Otherwise, <paramref name="ErrorMessage"/> should contain the formatted message.
 /// </param>
-/// <param name="EmitTime">Time it took to emit the query compilation.</param>
 /// <param name="ExecutionTime">Time it took to execute the query.</param>
 [DataContract]
 internal readonly record struct ExecuteQueryResult(
-    [property: DataMember(Order = 0)] ImmutableArray<QueryCompilationError> compilationErrors,
-    [property: DataMember(Order = 1)] string? ErrorMessage,
-    [property: DataMember(Order = 2)] string[]? ErrorMessageArgs = null,
-    [property: DataMember(Order = 3)] TimeSpan EmitTime = default,
-    [property: DataMember(Order = 4)] TimeSpan ExecutionTime = default);
+    [property: DataMember(Order = 0)] string? ErrorMessage,
+    [property: DataMember(Order = 1)] string[]? ErrorMessageArgs = null,
+    [property: DataMember(Order = 2)] TimeSpan ExecutionTime = default);
+
+/// <summary>
+/// The result of Semantic Search query compilation.
+/// </summary>
+/// <param name="QueryId">Id of the compiled query if the compilation was successful.</param>
+/// <param name="CompilationErrors">Compilation errors.</param>
+/// <param name="EmitTime">Time it took to emit the query compilation.</param>
+[DataContract]
+internal readonly record struct CompileQueryResult(
+    [property: DataMember(Order = 0)] CompiledQueryId QueryId,
+    [property: DataMember(Order = 1)] ImmutableArray<QueryCompilationError> CompilationErrors,
+    [property: DataMember(Order = 2)] TimeSpan EmitTime = default);
+
+[DataContract]
+internal readonly record struct CompiledQueryId
+{
+    private static int s_id;
+
+    [DataMember(Order = 0)]
+#pragma warning disable IDE0052 // Remove unread private members (https://github.com/dotnet/roslyn/issues/77907)
+    private readonly int _id;
+#pragma warning restore IDE0052
+
+    [DataMember(Order = 1)]
+#pragma warning disable IDE0052 // Remove unread private members (https://github.com/dotnet/roslyn/issues/77907)
+    public readonly string Language;
+#pragma warning restore IDE0052
+
+    private CompiledQueryId(int id, string language)
+    {
+        _id = id;
+        Language = language;
+    }
+
+    public static CompiledQueryId Create(string language)
+        => new(Interlocked.Increment(ref s_id), language);
+}

src/Features/Core/Portable/SemanticSearch/IRemoteSemanticSearchService.cs

@@ -10,6 +10,7 @@
 using Microsoft.CodeAnalysis.Classification;
 using Microsoft.CodeAnalysis.ErrorReporting;
 using Microsoft.CodeAnalysis.FindUsages;
+using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Remote;
 
@@ -26,7 +27,9 @@ internal interface ICallback
         ValueTask ItemsCompletedAsync(RemoteServiceCallbackId callbackId, int itemCount, CancellationToken cancellationToken);
     }
 
-    ValueTask<ExecuteQueryResult> ExecuteQueryAsync(Checksum solutionChecksum, RemoteServiceCallbackId callbackId, string language, string query, string referenceAssembliesDir, CancellationToken cancellationToken);
+    ValueTask<CompileQueryResult> CompileQueryAsync(string query, string language, string referenceAssembliesDir, CancellationToken cancellationToken);
+    ValueTask<ExecuteQueryResult> ExecuteQueryAsync(Checksum solutionChecksum, RemoteServiceCallbackId callbackId, CompiledQueryId queryId, CancellationToken cancellationToken);
+    ValueTask DiscardQueryAsync(CompiledQueryId queryId, CancellationToken cancellationToken);
 }
 
 internal static class RemoteSemanticSearchServiceProxy
@@ -112,19 +115,41 @@ public async ValueTask<ClassificationOptions> GetClassificationOptionsAsync(stri
         }
     }
 
-    public static async ValueTask<ExecuteQueryResult> ExecuteQueryAsync(Solution solution, string language, string query, string referenceAssembliesDir, ISemanticSearchResultsObserver results, OptionsProvider<ClassificationOptions> classificationOptions, CancellationToken cancellationToken)
+    public static async ValueTask<CompileQueryResult?> CompileQueryAsync(SolutionServices services, string query, string language, string referenceAssembliesDir, CancellationToken cancellationToken)
     {
-        var client = await RemoteHostClient.TryGetClientAsync(solution.Services, cancellationToken).ConfigureAwait(false);
+        var client = await RemoteHostClient.TryGetClientAsync(services, cancellationToken).ConfigureAwait(false);
         if (client == null)
         {
-            return new ExecuteQueryResult(compilationErrors: [], FeaturesResources.Semantic_search_only_supported_on_net_core);
+            return null;
         }
 
+        var result = await client.TryInvokeAsync<IRemoteSemanticSearchService, CompileQueryResult>(
+            (service, cancellationToken) => service.CompileQueryAsync(query, language, referenceAssembliesDir, cancellationToken),
+            cancellationToken).ConfigureAwait(false);
+
+        return result.Value;
+    }
+
+    public static async ValueTask DiscardQueryAsync(SolutionServices services, CompiledQueryId queryId, CancellationToken cancellationToken)
+    {
+        var client = await RemoteHostClient.TryGetClientAsync(services, cancellationToken).ConfigureAwait(false);
+        Contract.ThrowIfNull(client);
+
+        await client.TryInvokeAsync<IRemoteSemanticSearchService>(
+            (service, cancellationToken) => service.DiscardQueryAsync(queryId, cancellationToken),
+            cancellationToken).ConfigureAwait(false);
+    }
+
+    public static async ValueTask<ExecuteQueryResult> ExecuteQueryAsync(Solution solution, CompiledQueryId queryId, ISemanticSearchResultsObserver results, OptionsProvider<ClassificationOptions> classificationOptions, CancellationToken cancellationToken)
+    {
+        var client = await RemoteHostClient.TryGetClientAsync(solution.Services, cancellationToken).ConfigureAwait(false);
+        Contract.ThrowIfNull(client);
+
         var serverCallback = new ServerCallback(solution, results, classificationOptions);
 
         var result = await client.TryInvokeAsync<IRemoteSemanticSearchService, ExecuteQueryResult>(
             solution,
-            (service, solutionInfo, callbackId, cancellationToken) => service.ExecuteQueryAsync(solutionInfo, callbackId, language, query, referenceAssembliesDir, cancellationToken),
+            (service, solutionInfo, callbackId, cancellationToken) => service.ExecuteQueryAsync(solutionInfo, callbackId, queryId, cancellationToken),
             callbackTarget: serverCallback,
             cancellationToken).ConfigureAwait(false);
 

src/Features/Core/Portable/SemanticSearch/ISemanticSearchService.cs

@@ -13,21 +13,36 @@ namespace Microsoft.CodeAnalysis.SemanticSearch;
 internal interface ISemanticSearchService : ILanguageService
 {
     /// <summary>
-    /// Executes given <paramref name="query"/> query against <paramref name="solution"/>.
+    /// Compiles a query. The query has to be executed or discarded.
     /// </summary>
-    /// <param name="solution">The solution snapshot.</param>
     /// <param name="query">Query (top-level code).</param>
     /// <param name="referenceAssembliesDir">Directory that contains refernece assemblies to be used for compilation of the query.</param>
+    CompileQueryResult CompileQuery(
+        SolutionServices services,
+        string query,
+        string referenceAssembliesDir,
+        TraceSource traceSource,
+        CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Executes given query against <paramref name="solution"/> and discards it.
+    /// </summary>
+    /// <param name="solution">The solution snapshot.</param>
+    /// <param name="queryId">Id of a compiled query.</param>
     /// <param name="observer">Observer of the found symbols.</param>
     /// <param name="classificationOptions">Options to use to classify the textual representation of the found symbols.</param>
     /// <param name="cancellationToken">Cancellation token.</param>
-    /// <returns>Error message on failure.</returns>
     Task<ExecuteQueryResult> ExecuteQueryAsync(
         Solution solution,
-        string query,
-        string referenceAssembliesDir,
+        CompiledQueryId queryId,
         ISemanticSearchResultsObserver observer,
         OptionsProvider<ClassificationOptions> classificationOptions,
         TraceSource traceSource,
         CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Discards resources associated with compiled query.
+    /// Only call if the query is not executed.
+    /// </summary>
+    void DiscardQuery(CompiledQueryId queryId);
 }