Reduce closure allocations associated with AsyncLazy usage (#72449)

 Reduce closure allocations associated with AsyncLazy usage

This was noticed while looking at speedometer profiles. This class is used in several high traffic areas, and the current AsyncLazy design doesn't allow usage in a way that prevents closure allocations. This PR adds the "arg" usage pattern to AsyncLazy such that funcs passed to AsyncLazy can take advantage of the arg infrastructure to avoid closure allocations.

This is done by changing AsyncLazy<T> to be an abstract class with a single derivation that caches the argument data. Users still declare their type to be AsyncLazy<T>, but the AsyncLazy.Create methods have been altered to have optional overloads that take in data.

src/EditorFeatures/Core/RenameTracking/RenameTrackingTaggerProvider.RenameTrackingCommitter.cs

@@ -46,7 +46,9 @@ public RenameTrackingCommitter(
             _refactorNotifyServices = refactorNotifyServices;
             _undoHistoryRegistry = undoHistoryRegistry;
             _displayText = displayText;
-            _renameSymbolResultGetter = AsyncLazy.Create(c => RenameSymbolWorkerAsync(c));
+            _renameSymbolResultGetter = AsyncLazy.Create(
+                static (self, c) => self.RenameSymbolWorkerAsync(c),
+                arg: this);
         }
 
         /// <summary>

src/Features/Core/Portable/AddImport/SearchScopes/SourceSymbolsProjectSearchScope.cs

@@ -48,11 +48,12 @@ protected override async Task<ImmutableArray<ISymbol>> FindDeclarationsAsync(
             return declarations;
 
             static AsyncLazy<IAssemblySymbol?> CreateLazyAssembly(Project project)
-                => new(async c =>
+                => AsyncLazy.Create(static async (project, c) =>
                        {
                            var compilation = await project.GetRequiredCompilationAsync(c).ConfigureAwait(false);
-                           return compilation.Assembly;
-                       });
+                           return (IAssemblySymbol?)compilation.Assembly;
+                       },
+                       arg: project);
         }
     }
 }

src/Features/Core/Portable/Completion/SharedSyntaxContextsWithSpeculativeModel.cs

@@ -43,8 +43,8 @@ public Task<SyntaxContext> GetSyntaxContextAsync(Document document, Cancellation
         // Extract a local function to avoid creating a closure for code path of cache hit.
         static AsyncLazy<SyntaxContext> GetLazySyntaxContextWithSpeculativeModel(Document document, SharedSyntaxContextsWithSpeculativeModel self)
         {
-            return self._cache.GetOrAdd(document, d => AsyncLazy.Create(cancellationToken
-                => Utilities.CreateSyntaxContextWithExistingSpeculativeModelAsync(d, self._position, cancellationToken)));
+            return self._cache.GetOrAdd(document, d => AsyncLazy.Create(static (arg, cancellationToken)
+                => Utilities.CreateSyntaxContextWithExistingSpeculativeModelAsync(arg.d, arg._position, cancellationToken), (d, self._position)));
         }
     }
 }

src/Features/Core/Portable/DesignerAttribute/DesignerAttributeDiscoveryService.cs

@@ -82,8 +82,9 @@ static async Task<bool> HasDesignerCategoryTypeAsync(
             }
 
             var asyncLazy = s_metadataIdToDesignerAttributeInfo.GetValue(
-                metadataId, _ => AsyncLazy.Create(cancellationToken =>
-                    ComputeHasDesignerCategoryTypeAsync(solutionServices, solutionKey, peReference, cancellationToken)));
+                metadataId, _ => AsyncLazy.Create(static (arg, cancellationToken) =>
+                    ComputeHasDesignerCategoryTypeAsync(arg.solutionServices, arg.solutionKey, arg.peReference, cancellationToken),
+                    arg: (solutionServices, solutionKey, peReference)));
             return await asyncLazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
         }
 
@@ -124,7 +125,9 @@ public async ValueTask ProcessPriorityDocumentAsync(
 
         using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
         {
-            var lazyProjectVersion = AsyncLazy.Create(frozenProject.GetSemanticVersionAsync);
+            var lazyProjectVersion = AsyncLazy.Create(static (frozenProject, c) =>
+                frozenProject.GetSemanticVersionAsync(c),
+                arg: frozenProject);
 
             await ScanForDesignerCategoryUsageAsync(
                 frozenProject, frozenDocument, callback, lazyProjectVersion, cancellationToken).ConfigureAwait(false);
@@ -170,7 +173,9 @@ private async Task ProcessProjectAsync(
         // The top level project version for this project.  We only care if anything top level changes here.
         // Downstream impact will already happen due to us keying off of the references a project has (which will
         // change if anything it depends on changes).
-        var lazyProjectVersion = AsyncLazy.Create(project.GetSemanticVersionAsync);
+        var lazyProjectVersion = AsyncLazy.Create(static (project, c) =>
+            project.GetSemanticVersionAsync(c),
+            arg: project);
 
         await ScanForDesignerCategoryUsageAsync(
             project, specificDocument: null, callback, lazyProjectVersion, cancellationToken).ConfigureAwait(false);

src/Features/Core/Portable/EditAndContinue/EditAndContinueDocumentAnalysesCache.cs

@@ -183,18 +183,19 @@ private AsyncLazy<DocumentAnalysisResults> GetDocumentAnalysisNoLock(Project bas
         }
 
         var lazyResults = AsyncLazy.Create(
-            asynchronousComputeFunction: async cancellationToken =>
+            static async (arg, cancellationToken) =>
             {
                 try
                 {
-                    var analyzer = document.Project.Services.GetRequiredService<IEditAndContinueAnalyzer>();
-                    return await analyzer.AnalyzeDocumentAsync(baseProject, _baseActiveStatements, document, activeStatementSpans, _capabilities, cancellationToken).ConfigureAwait(false);
+                    var analyzer = arg.document.Project.Services.GetRequiredService<IEditAndContinueAnalyzer>();
+                    return await analyzer.AnalyzeDocumentAsync(arg.baseProject, arg.self._baseActiveStatements, arg.document, arg.activeStatementSpans, arg.self._capabilities, cancellationToken).ConfigureAwait(false);
                 }
                 catch (Exception e) when (FatalError.ReportAndPropagateUnlessCanceled(e, cancellationToken))
                 {
                     throw ExceptionUtilities.Unreachable();
                 }
-            });
+            },
+            arg: (self: this, document, baseProject, activeStatementSpans));
 
         // Previous results for this document id are discarded as they are no longer relevant.
         // The only relevant analysis is for the latest base and document snapshots.

src/Features/Core/Portable/EditAndContinue/EditSession.cs

@@ -108,10 +108,14 @@ internal EditSession(
         telemetry.SetBreakState(inBreakState);
 
         BaseActiveStatements = lazyActiveStatementMap ?? (inBreakState
-            ? AsyncLazy.Create(GetBaseActiveStatementsAsync)
-            : new AsyncLazy<ActiveStatementsMap>(ActiveStatementsMap.Empty));
-
-        Capabilities = AsyncLazy.Create(GetCapabilitiesAsync);
+            ? AsyncLazy.Create(static (self, cancellationToken) =>
+                self.GetBaseActiveStatementsAsync(cancellationToken),
+                arg: this)
+            : AsyncLazy.Create(ActiveStatementsMap.Empty));
+
+        Capabilities = AsyncLazy.Create(static (self, cancellationToken) =>
+            self.GetCapabilitiesAsync(cancellationToken),
+            arg: this);
         Analyses = new EditAndContinueDocumentAnalysesCache(BaseActiveStatements, Capabilities);
     }
 

src/Features/Core/Portable/NavigateTo/AbstractNavigateToSearchService.CachedDocumentSearch.cs

@@ -211,8 +211,9 @@ await ProcessIndexAsync(
         // match on disk anymore.
         var asyncLazy = cachedIndexMap.GetOrAdd(
             (storageService, documentKey, stringTable),
-            static t => AsyncLazy.Create(
-                c => TopLevelSyntaxTreeIndex.LoadAsync(t.service, t.documentKey, checksum: null, t.stringTable, c)));
+            static t => AsyncLazy.Create(static (t, c) =>
+                TopLevelSyntaxTreeIndex.LoadAsync(t.service, t.documentKey, checksum: null, t.stringTable, c),
+                arg: t));
         return asyncLazy.GetValueAsync(cancellationToken);
     }
 }

src/Interactive/Host/Interactive/Core/InteractiveHost.LazyRemoteService.cs

@@ -34,7 +34,7 @@ private sealed class LazyRemoteService
 
             public LazyRemoteService(InteractiveHost host, InteractiveHostOptions options, int instanceId, bool skipInitialization)
             {
-                _lazyInitializedService = AsyncLazy.Create(TryStartAndInitializeProcessAsync);
+                _lazyInitializedService = AsyncLazy.Create(static (self, cancellationToken) => self.TryStartAndInitializeProcessAsync(cancellationToken), this);
                 _cancellationSource = new CancellationTokenSource();
                 InstanceId = instanceId;
                 Options = options;

src/Interactive/Host/Microsoft.CodeAnalysis.InteractiveHost.csproj

@@ -37,10 +37,12 @@
   </ItemGroup>
   <ItemGroup>
     <Compile Include="..\..\Compilers\Core\Portable\InternalUtilities\NonCopyableAttribute.cs" Link="Utilities\NonCopyableAttribute.cs" />
+    <Compile Include="..\..\Compilers\Core\Portable\InternalUtilities\VoidResult.cs" Link="Utilities\VoidResult.cs" />
     <Compile Include="..\..\Workspaces\SharedUtilitiesAndExtensions\Compiler\Core\TestHooks\IExpeditableDelaySource.cs" Link="Utilities\IExpeditableDelaySource.cs" />
     <Compile Include="..\..\Workspaces\SharedUtilitiesAndExtensions\Compiler\Core\Utilities\Contract.cs" Link="Utilities\Contract.cs" />
     <Compile Include="..\..\Workspaces\SharedUtilitiesAndExtensions\Compiler\Core\Utilities\Contract.InterpolatedStringHandlers.cs" Link="Utilities\Contract.InterpolatedStringHandlers.cs" />
     <Compile Include="..\..\Workspaces\SharedUtilitiesAndExtensions\Compiler\Core\Utilities\TaskExtensions.cs" Link="Utilities\TaskExtensions.cs" />
+    <Compile Include="..\..\Workspaces\SharedUtilitiesAndExtensions\Compiler\Core\Utilities\AsyncLazy.cs" Link="Utilities\AsyncLazy.cs" />
     <Compile Include="..\..\Workspaces\SharedUtilitiesAndExtensions\Compiler\Core\Utilities\AsyncLazy`1.cs" Link="Utilities\AsyncLazy`1.cs" />
     <Compile Include="..\..\Workspaces\SharedUtilitiesAndExtensions\Compiler\Core\Utilities\NonReentrantLock.cs" Link="Utilities\NonReentrantLock.cs" />
     <Compile Include="..\..\Compilers\Core\Portable\InternalUtilities\InterpolatedStringHandlerArgumentAttribute.cs" Link="Utilities\InterpolatedStringHandlerArgumentAttribute.cs" />

src/VisualStudio/Core/Def/ProjectSystem/FileChangeTracker.cs

@@ -20,7 +20,7 @@ internal sealed class FileChangeTracker : IVsFreeThreadedFileChangeEvents2, IDis
 {
     internal const _VSFILECHANGEFLAGS DefaultFileChangeFlags = _VSFILECHANGEFLAGS.VSFILECHG_Time | _VSFILECHANGEFLAGS.VSFILECHG_Add | _VSFILECHANGEFLAGS.VSFILECHG_Del | _VSFILECHANGEFLAGS.VSFILECHG_Size;
 
-    private static readonly AsyncLazy<uint?> s_none = new(value: null);
+    private static readonly AsyncLazy<uint?> s_none = AsyncLazy.Create(value: (uint?)null);
 
     private readonly IVsFileChangeEx _fileChangeService;
     private readonly string _filePath;
@@ -107,30 +107,34 @@ public Task StartFileChangeListeningAsync()
 
         Contract.ThrowIfTrue(_fileChangeCookie != s_none);
 
-        _fileChangeCookie = new AsyncLazy<uint?>(async cancellationToken =>
-        {
-            try
+        _fileChangeCookie = AsyncLazy.Create(
+            static async (self, cancellationToken) =>
             {
-                // TODO: Should we pass in cancellationToken here insead of CancellationToken.None?
-                return await ((IVsAsyncFileChangeEx2)_fileChangeService).AdviseFileChangeAsync(_filePath, _fileChangeFlags, this, CancellationToken.None).ConfigureAwait(false);
-            }
-            catch (Exception e) when (ReportException(e))
+                try
+                {
+                    // TODO: Should we pass in cancellationToken here instead of CancellationToken.None?
+                    uint? result = await ((IVsAsyncFileChangeEx2)self._fileChangeService).AdviseFileChangeAsync(self._filePath, self._fileChangeFlags, self, CancellationToken.None).ConfigureAwait(false);
+                    return result;
+                }
+                catch (Exception e) when (ReportException(e))
+                {
+                    return null;
+                }
+            },
+            static (self, cancellationToken) =>
             {
-                return null;
-            }
-        }, cancellationToken =>
-        {
-            try
-            {
-                Marshal.ThrowExceptionForHR(
-                    _fileChangeService.AdviseFileChange(_filePath, (uint)_fileChangeFlags, this, out var newCookie));
-                return newCookie;
-            }
-            catch (Exception e) when (ReportException(e))
-            {
-                return null;
-            }
-        });
+                try
+                {
+                    Marshal.ThrowExceptionForHR(
+                        self._fileChangeService.AdviseFileChange(self._filePath, (uint)self._fileChangeFlags, self, out var newCookie));
+                    return newCookie;
+                }
+                catch (Exception e) when (ReportException(e))
+                {
+                    return null;
+                }
+            },
+            arg: this);
 
         lock (s_lastBackgroundTaskGate)
         {

src/Workspaces/Core/Portable/FindSymbols/Declarations/DeclarationFinder_AllDeclarations.cs

@@ -63,7 +63,9 @@ internal static async Task<ImmutableArray<ISymbol>> FindAllDeclarationsWithNorma
 
         // Lazily produce the compilation.  We don't want to incur any costs (especially source generators) if there
         // are no results for this query in this project.
-        var lazyCompilation = AsyncLazy.Create(project.GetRequiredCompilationAsync);
+        var lazyCompilation = AsyncLazy.Create(static (project, cancellationToken) =>
+            project.GetRequiredCompilationAsync(cancellationToken),
+            arg: project);
 
         if (project.SupportsCompilation)
         {
@@ -114,12 +116,13 @@ async Task SearchMetadataReferencesAsync()
             {
                 using var _ = ArrayBuilder<ISymbol>.GetInstance(out var buffer);
 
-                var lazyAssembly = AsyncLazy.Create(async cancellationToken =>
-                {
-                    var compilation = await lazyCompilation.GetValueAsync(cancellationToken).ConfigureAwait(false);
-                    var assemblySymbol = compilation.GetAssemblyOrModuleSymbol(peReference) as IAssemblySymbol;
-                    return assemblySymbol;
-                });
+                var lazyAssembly = AsyncLazy.Create(static async (arg, cancellationToken) =>
+                    {
+                        var compilation = await arg.lazyCompilation.GetValueAsync(cancellationToken).ConfigureAwait(false);
+                        var assemblySymbol = compilation.GetAssemblyOrModuleSymbol(arg.peReference) as IAssemblySymbol;
+                        return assemblySymbol;
+                    },
+                    arg: (lazyCompilation, peReference));
 
                 await AddMetadataDeclarationsWithNormalQueryAsync(
                     project, lazyAssembly, peReference,

src/Workspaces/Core/Portable/FindSymbols/FindReferences/DependentTypeFinder_ProjectIndex.cs

@@ -35,8 +35,9 @@ public static Task<ProjectIndex> GetIndexAsync(
             if (!s_projectToIndex.TryGetValue(project.State, out var lazyIndex))
             {
                 lazyIndex = s_projectToIndex.GetValue(
-                    project.State, p => new AsyncLazy<ProjectIndex>(
-                        c => CreateIndexAsync(project, c)));
+                    project.State, p => AsyncLazy.Create(
+                        static (project, c) => CreateIndexAsync(project, c),
+                        project));
             }
 
             return lazyIndex.GetValueAsync(cancellationToken);

src/Workspaces/Core/Portable/FindSymbols/SymbolTree/SymbolTreeInfo.cs

@@ -123,7 +123,7 @@ public Task<ImmutableArray<ISymbol>> FindAsync(
         Contract.ThrowIfTrue(query.Kind == SearchKind.Custom, "Custom queries are not supported in this API");
 
         return this.FindAsync(
-            query, new AsyncLazy<IAssemblySymbol?>(assembly), filter, cancellationToken);
+            query, AsyncLazy.Create((IAssemblySymbol?)assembly), filter, cancellationToken);
     }
 
     public async Task<ImmutableArray<ISymbol>> FindAsync(

src/Workspaces/Core/Portable/FindSymbols/SymbolTree/SymbolTreeInfo_Metadata.cs

@@ -158,8 +158,9 @@ static async Task<SymbolTreeInfo> GetInfoForMetadataReferenceSlowAsync(
             // CreateMetadataSymbolTreeInfoAsync
             var asyncLazy = s_peReferenceToInfo.GetValue(
                 reference,
-                id => AsyncLazy.Create(
-                    c => CreateMetadataSymbolTreeInfoAsync(services, solutionKey, reference, checksum, c)));
+                id => AsyncLazy.Create(static (arg, c) =>
+                    CreateMetadataSymbolTreeInfoAsync(arg.services, arg.solutionKey, arg.reference, arg.checksum, c),
+                    arg: (services, solutionKey, reference, checksum)));
 
             return await asyncLazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
         }
@@ -177,14 +178,15 @@ static async Task<SymbolTreeInfo> CreateMetadataSymbolTreeInfoAsync(
 
             var asyncLazy = s_metadataIdToSymbolTreeInfo.GetValue(
                 metadataId,
-                metadataId => AsyncLazy.Create(
-                    cancellationToken => LoadOrCreateAsync(
-                        services,
-                        solutionKey,
-                        checksum,
-                        createAsync: checksum => new ValueTask<SymbolTreeInfo>(new MetadataInfoCreator(checksum, GetMetadataNoThrow(reference)).Create()),
-                        keySuffix: GetMetadataKeySuffix(reference),
-                        cancellationToken)));
+                metadataId => AsyncLazy.Create(static (arg, cancellationToken) =>
+                        LoadOrCreateAsync(
+                        arg.services,
+                        arg.solutionKey,
+                        arg.checksum,
+                        createAsync: checksum => new ValueTask<SymbolTreeInfo>(new MetadataInfoCreator(checksum, GetMetadataNoThrow(arg.reference)).Create()),
+                        keySuffix: GetMetadataKeySuffix(arg.reference),
+                        cancellationToken),
+                    arg: (services, solutionKey, checksum, reference)));
 
             var metadataIdSymbolTreeInfo = await asyncLazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
 

src/Workspaces/Core/Portable/FindSymbols/SymbolTree/SymbolTreeInfo_Source.cs

@@ -72,7 +72,10 @@ public static Task<SymbolTreeInfo> GetInfoForSourceAssemblyAsync(
     public static Task<Checksum> GetSourceSymbolsChecksumAsync(Project project, CancellationToken cancellationToken)
     {
         var lazy = s_projectToSourceChecksum.GetValue(
-            project.State, static p => AsyncLazy.Create(c => ComputeSourceSymbolsChecksumAsync(p, c)));
+            project.State,
+            static p => AsyncLazy.Create(
+                static (p, c) => ComputeSourceSymbolsChecksumAsync(p, c),
+                arg: p));
 
         return lazy.GetValueAsync(cancellationToken);
     }

src/Workspaces/Core/Portable/Workspace/Solution/AnalyzerConfigDocumentState.cs

@@ -38,9 +38,10 @@ public AnalyzerConfigDocumentState(
 
     private AsyncLazy<AnalyzerConfig> CreateAnalyzerConfigValueSource()
     {
-        return new AsyncLazy<AnalyzerConfig>(
-            asynchronousComputeFunction: async cancellationToken => AnalyzerConfig.Parse(await GetTextAsync(cancellationToken).ConfigureAwait(false), FilePath),
-            synchronousComputeFunction: cancellationToken => AnalyzerConfig.Parse(GetTextSynchronously(cancellationToken), FilePath));
+        return AsyncLazy.Create(
+            asynchronousComputeFunction: static async (self, cancellationToken) => AnalyzerConfig.Parse(await self.GetTextAsync(cancellationToken).ConfigureAwait(false), self.FilePath),
+            synchronousComputeFunction: static (self, cancellationToken) => AnalyzerConfig.Parse(self.GetTextSynchronously(cancellationToken), self.FilePath),
+            arg: this);
     }
 
     public AnalyzerConfig GetAnalyzerConfig(CancellationToken cancellationToken) => _analyzerConfigValueSource.GetValue(cancellationToken);

src/Workspaces/Core/Portable/Workspace/Solution/Document.cs

@@ -556,11 +556,12 @@ public Task<DocumentOptionSet> GetOptionsAsync(CancellationToken cancellationTok
 
     private void InitializeCachedOptions(OptionSet solutionOptions)
     {
-        var newAsyncLazy = AsyncLazy.Create(async cancellationToken =>
+        var newAsyncLazy = AsyncLazy.Create(static async (arg, cancellationToken) =>
         {
-            var options = await GetAnalyzerConfigOptionsAsync(cancellationToken).ConfigureAwait(false);
-            return new DocumentOptionSet(options, solutionOptions, Project.Language);
-        });
+            var options = await arg.self.GetAnalyzerConfigOptionsAsync(cancellationToken).ConfigureAwait(false);
+            return new DocumentOptionSet(options, arg.solutionOptions, arg.self.Project.Language);
+        },
+        arg: (self: this, solutionOptions));
 
         Interlocked.CompareExchange(ref _cachedOptions, newAsyncLazy, comparand: null);
     }