Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Stop loading ISourceGenerators within VS entirely. #72835

Merged
merged 17 commits into from
Apr 2, 2024
11 changes: 11 additions & 0 deletions src/Workspaces/Core/Portable/Remote/RemoteHostClient.cs
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,17 @@ public async ValueTask<bool> TryInvokeAsync<TService>(
return await connection.TryInvokeAsync(project, invocation, cancellationToken).ConfigureAwait(false);
}

public async ValueTask<Optional<TResult>> TryInvokeAsync<TService, TResult>(
SolutionCompilationState compilationState,
ProjectId projectId,
Func<TService, Checksum, CancellationToken, ValueTask<TResult>> invocation,
CancellationToken cancellationToken)
where TService : class
{
using var connection = CreateConnection<TService>(callbackTarget: null);
return await connection.TryInvokeAsync(compilationState, projectId, invocation, cancellationToken).ConfigureAwait(false);
}
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

helper overload for the case where we're asking about a particular project in the context of a particular SolutionCompilationState (first time we've needed to do that).


/// <summary>
/// Equivalent to <see cref="TryInvokeAsync{TService}(Solution, Func{TService, Checksum, CancellationToken, ValueTask}, CancellationToken)"/>
/// except that only the project (and its dependent projects) will be sync'ed to the remote host before executing.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,12 @@ internal interface IRemoteSourceGenerationService
/// </summary>
ValueTask<ImmutableArray<string>> GetContentsAsync(
Checksum solutionChecksum, ProjectId projectId, ImmutableArray<DocumentId> documentIds, CancellationToken cancellationToken);

/// <summary>
/// Whether or not the specified <paramref name="projectId"/> has source generators or not.
/// </summary>
ValueTask<bool> HasGeneratorsAsync(
Checksum solutionChecksum, ProjectId projectId, CancellationToken cancellationToken);
Copy link
Member Author

@CyrusNajmabadi CyrusNajmabadi Apr 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pretty self-explanatory.

}

/// <summary>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -359,6 +359,9 @@ async Task<InProgressState> CollapseInProgressStateAsync(InProgressState initial
{
try
{
// Only bother keeping track of staleCompilationWithGeneratedDocuments for projects that
// actually have generators in them.
var hasSourceGenerators = await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false);
var currentState = initialState;

// Then, we serially process the chain while that parsing is happening concurrently.
Expand All @@ -368,7 +371,7 @@ async Task<InProgressState> CollapseInProgressStateAsync(InProgressState initial

// We have a list of transformations to get to our final compilation; take the first transformation and apply it.
var (compilationWithoutGeneratedDocuments, staleCompilationWithGeneratedDocuments, generatorInfo) =
await ApplyFirstTransformationAsync(currentState).ConfigureAwait(false);
await ApplyFirstTransformationAsync(currentState, hasSourceGenerators).ConfigureAwait(false);

// We have updated state, so store this new result; this allows us to drop the intermediate state we already processed
// even if we were to get cancelled at a later point.
Expand Down Expand Up @@ -397,7 +400,7 @@ async Task<InProgressState> CollapseInProgressStateAsync(InProgressState initial
}

async Task<(Compilation compilationWithoutGeneratedDocuments, Compilation? staleCompilationWithGeneratedDocuments, CompilationTrackerGeneratorInfo generatorInfo)>
ApplyFirstTransformationAsync(InProgressState inProgressState)
ApplyFirstTransformationAsync(InProgressState inProgressState, bool hasSourceGenerators)
{
Contract.ThrowIfTrue(inProgressState.PendingTranslationActions.IsEmpty);
var translationAction = inProgressState.PendingTranslationActions[0];
Expand All @@ -422,7 +425,7 @@ async Task<InProgressState> CollapseInProgressStateAsync(InProgressState initial
// Also transform the compilation that has generated files; we won't do that though if the transformation either would cause problems with
// the generated documents, or if don't have any source generators in the first place.
if (translationAction.CanUpdateCompilationWithStaleGeneratedTreesIfGeneratorsGiveSameOutput &&
GetSourceGenerators(translationAction.OldProjectState).Any())
hasSourceGenerators)
{
staleCompilationWithGeneratedDocuments = await translationAction.TransformCompilationAsync(staleCompilationWithGeneratedDocuments, cancellationToken).ConfigureAwait(false);
}
Expand Down Expand Up @@ -774,10 +777,8 @@ public async ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSour
SolutionCompilationState compilationState, CancellationToken cancellationToken)
{
// If we don't have any generators, then we know we have no generated files, so we can skip the computation entirely.
if (!GetSourceGenerators(this.ProjectState).Any())
{
if (!await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false))
return TextDocumentStates<SourceGeneratedDocumentState>.Empty;
}

var finalState = await GetOrBuildFinalStateAsync(
compilationState, cancellationToken: cancellationToken).ConfigureAwait(false);
Expand All @@ -787,10 +788,8 @@ public async ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSour
public async ValueTask<ImmutableArray<Diagnostic>> GetSourceGeneratorDiagnosticsAsync(
SolutionCompilationState compilationState, CancellationToken cancellationToken)
{
if (!GetSourceGenerators(this.ProjectState).Any())
{
if (!await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false))
return [];
}

var finalState = await GetOrBuildFinalStateAsync(
compilationState, cancellationToken: cancellationToken).ConfigureAwait(false);
Expand All @@ -816,10 +815,8 @@ public async ValueTask<ImmutableArray<Diagnostic>> GetSourceGeneratorDiagnostics

public async ValueTask<GeneratorDriverRunResult?> GetSourceGeneratorRunResultAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken)
{
if (!GetSourceGenerators(this.ProjectState).Any())
{
if (!await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false))
return null;
}

var finalState = await GetOrBuildFinalStateAsync(
compilationState, cancellationToken).ConfigureAwait(false);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -57,14 +57,13 @@ private partial class CompilationTracker : ICompilationTracker
if (result.HasValue)
{
// Since we ran the SG work out of process, we could not have created or modified the driver passed in.
// So just pass what we got in right back out.
return (result.Value.compilationWithGeneratedFiles, new(result.Value.generatedDocuments, generatorInfo.Driver));
// Just return `null` for the driver as there's nothing to track for it on the host side.
return (result.Value.compilationWithGeneratedFiles, new(result.Value.generatedDocuments, Driver: null));
}

// If that failed (OOP crash, or we are the OOP process ourselves), then generate the SG docs locally.
var telemetryCollector = compilationState.SolutionState.Services.GetService<ISourceGeneratorTelemetryCollectorWorkspaceService>();
var (compilationWithGeneratedFiles, nextGeneratedDocuments, nextGeneratorDriver) = await ComputeNewGeneratorInfoInCurrentProcessAsync(
telemetryCollector,
compilationState,
compilationWithoutGeneratedFiles,
generatorInfo.Documents,
generatorInfo.Driver,
Expand Down Expand Up @@ -222,18 +221,15 @@ await newGeneratedDocuments.States.Values.SelectAsArrayAsync(
}

private async Task<(Compilation compilationWithGeneratedFiles, TextDocumentStates<SourceGeneratedDocumentState> generatedDocuments, GeneratorDriver? generatorDriver)> ComputeNewGeneratorInfoInCurrentProcessAsync(
ISourceGeneratorTelemetryCollectorWorkspaceService? telemetryCollector,
SolutionCompilationState compilationState,
Compilation compilationWithoutGeneratedFiles,
TextDocumentStates<SourceGeneratedDocumentState> oldGeneratedDocuments,
GeneratorDriver? generatorDriver,
Compilation? compilationWithStaleGeneratedTrees,
CancellationToken cancellationToken)
{
// If we don't have any source generators. Trivially bail out. Note: this check is intentionally don't in
// the "InCurrentProcess" call so that it will normally run only in the OOP process, thus ensuring that we
// get accurate information about what SourceGenerators we actually have (say, in case they they are rebuilt
// by the user while VS is running).
if (!GetSourceGenerators(this.ProjectState).Any())
// If we don't have any source generators. Trivially bail out.
if (!await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false))
return (compilationWithoutGeneratedFiles, TextDocumentStates<SourceGeneratedDocumentState>.Empty, generatorDriver);

// If we don't already have an existing generator driver, create one from scratch
Expand Down Expand Up @@ -268,6 +264,7 @@ await newGeneratedDocuments.States.Values.SelectAsArrayAsync(

var runResult = generatorDriver.GetRunResult();

var telemetryCollector = compilationState.SolutionState.Services.GetService<ISourceGeneratorTelemetryCollectorWorkspaceService>();
telemetryCollector?.CollectRunResult(
runResult, generatorDriver.GetTimingInfo(),
g => GetAnalyzerReference(this.ProjectState, g));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,19 @@
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.Diagnostics;
using Microsoft.CodeAnalysis.Remote;
using Microsoft.CodeAnalysis.Shared.Collections;
using Microsoft.CodeAnalysis.SourceGeneration;
using Roslyn.Utilities;

namespace Microsoft.CodeAnalysis;

using AnalyzerReferencesToSourceGenerators = ConditionalWeakTable<IReadOnlyList<AnalyzerReference>, SolutionCompilationState.SourceGeneratorMap>;

internal partial class SolutionCompilationState
{
internal sealed record SourceGeneratorMap(
Expand All @@ -26,52 +28,92 @@ internal sealed record SourceGeneratorMap(
/// of things so that we don't cause source generators to be loaded (and fixed) within VS (which is .net framework
/// only).
/// </summary>
private static readonly ImmutableArray<(string language, AnalyzerReferencesToSourceGenerators referencesToGenerators, AnalyzerReferencesToSourceGenerators.CreateValueCallback callback)> s_languageToAnalyzerReferencesToSourceGeneratorsMap =
[
(LanguageNames.CSharp, new(), (static rs => ComputeSourceGenerators(rs, LanguageNames.CSharp))),
(LanguageNames.VisualBasic, new(), (static rs => ComputeSourceGenerators(rs, LanguageNames.VisualBasic))),
];

private static SourceGeneratorMap ComputeSourceGenerators(IReadOnlyList<AnalyzerReference> analyzerReferences, string language)
{
using var generators = TemporaryArray<ISourceGenerator>.Empty;
var generatorToAnalyzerReference = ImmutableDictionary.CreateBuilder<ISourceGenerator, AnalyzerReference>();
private static readonly ConditionalWeakTable<ProjectState, SourceGeneratorMap> s_projectStateToSourceGeneratorsMap = new();
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this was a simpler way to organize the existing information. ProjectState already has the language and analyzerreferences. It's also what callers already pass in. And it's a green node which rarely changes. Note: if it does change, we'll 'recompute' the SGs, but end up finding the same set, since the AnalyzerRef objects themselves cache and return the same instances.

CyrusNajmabadi marked this conversation as resolved.
Show resolved Hide resolved

foreach (var reference in analyzerReferences)
{
foreach (var generator in reference.GetGenerators(language).Distinct())
{
generators.Add(generator);
generatorToAnalyzerReference.Add(generator, reference);
}
}

return new(generators.ToImmutableAndClear(), generatorToAnalyzerReference.ToImmutable());
}
/// <summary>
/// Cached information about if a project has source generators or not. Note: this is distinct from <see
/// cref="s_projectStateToSourceGeneratorsMap"/> as we want to be able to compute it by calling over to our OOP
/// process (if present) and having it make the determination, without the host necessarily loading generators
/// itself.
/// </summary>
private static readonly ConditionalWeakTable<ProjectState, AsyncLazy<bool>> s_hasSourceGeneratorsMap = new();
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

new to this PR.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ConditionalWeakTable

Would it be possible to only have a single CWT and only access the SourceGeneratorMap part of the tuple when OOP?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

tried, but it started getting very ugly/allocy


/// <summary>
/// This method should only be called in a .net core host like our out of process server.
/// </summary>
private static ImmutableArray<ISourceGenerator> GetSourceGenerators(ProjectState projectState)
=> GetSourceGenerators(projectState.Language, projectState.AnalyzerReferences);

private static ImmutableArray<ISourceGenerator> GetSourceGenerators(string language, IReadOnlyList<AnalyzerReference> analyzerReferences)
{
var map = GetSourceGeneratorMap(language, analyzerReferences);
var map = GetSourceGeneratorMap(projectState);
return map is null ? [] : map.SourceGenerators;
}

/// <summary>
/// This method should only be called in a .net core host like our out of process server.
/// </summary>
private static AnalyzerReference GetAnalyzerReference(ProjectState projectState, ISourceGenerator sourceGenerator)
{
var map = GetSourceGeneratorMap(projectState.Language, projectState.AnalyzerReferences);
var map = GetSourceGeneratorMap(projectState);
Contract.ThrowIfNull(map);
CyrusNajmabadi marked this conversation as resolved.
Show resolved Hide resolved
return map.SourceGeneratorToAnalyzerReference[sourceGenerator];
}

private static SourceGeneratorMap? GetSourceGeneratorMap(string language, IReadOnlyList<AnalyzerReference> analyzerReferences)
private static SourceGeneratorMap? GetSourceGeneratorMap(ProjectState projectState)
{
var tupleOpt = s_languageToAnalyzerReferencesToSourceGeneratorsMap.FirstOrNull(static (t, language) => t.language == language, language);
if (tupleOpt is null)
if (!projectState.SupportsCompilation)
return null;

var tuple = tupleOpt.Value;
return tuple.referencesToGenerators.GetValue(analyzerReferences, tuple.callback);
return s_projectStateToSourceGeneratorsMap.GetValue(projectState, ComputeSourceGenerators);

static SourceGeneratorMap ComputeSourceGenerators(ProjectState projectState)
{
using var generators = TemporaryArray<ISourceGenerator>.Empty;
var generatorToAnalyzerReference = ImmutableDictionary.CreateBuilder<ISourceGenerator, AnalyzerReference>();

foreach (var reference in projectState.AnalyzerReferences)
{
foreach (var generator in reference.GetGenerators(projectState.Language).Distinct())
{
generators.Add(generator);
generatorToAnalyzerReference.Add(generator, reference);
}
}

return new(generators.ToImmutableAndClear(), generatorToAnalyzerReference.ToImmutable());
}
}

public async Task<bool> HasSourceGeneratorsAsync(ProjectId projectId, CancellationToken cancellationToken)
{
var projectState = this.SolutionState.GetRequiredProjectState(projectId);

if (!s_hasSourceGeneratorsMap.TryGetValue(projectState, out var lazy))
{
// Extracted into local function to prevent allocations in the case where we find a value in the cache.
lazy = GetLazy(projectState);
}

return await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false);

AsyncLazy<bool> GetLazy(ProjectState projectState)
=> s_hasSourceGeneratorsMap.GetValue(
projectState,
projectState => AsyncLazy.Create(cancellationToken => ComputeHasSourceGeneratorsAsync(projectState, cancellationToken)));
CyrusNajmabadi marked this conversation as resolved.
Show resolved Hide resolved

async Task<bool> ComputeHasSourceGeneratorsAsync(
ProjectState projectState, CancellationToken cancellationToken)
{
var client = await RemoteHostClient.TryGetClientAsync(this.Services, cancellationToken).ConfigureAwait(false);
// If in proc, just load the generators and see if we have any.
if (client is null)
return GetSourceGenerators(projectState).Any();

// Out of process, call to the remote to figure this out.
var result = await client.TryInvokeAsync<IRemoteSourceGenerationService, bool>(
this,
projectId,
(service, solution, cancellationToken) => service.HasGeneratorsAsync(solution, projectId, cancellationToken),
cancellationToken).ConfigureAwait(false);
return result.HasValue && result.Value;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -63,4 +63,12 @@ public ValueTask<ImmutableArray<string>> GetContentsAsync(
return result.ToImmutableAndClear();
}, cancellationToken);
}

public ValueTask<bool> HasGeneratorsAsync(Checksum solutionChecksum, ProjectId projectId, CancellationToken cancellationToken)
{
return RunServiceAsync(solutionChecksum, async solution =>
{
return await solution.CompilationState.HasSourceGeneratorsAsync(projectId, cancellationToken).ConfigureAwait(false);
}, cancellationToken);
}
}
Loading