Merge pull request #12145 from heejaechang/leakfix

fixed 2 memory leaks

src/Compilers/Core/Portable/DiagnosticAnalyzer/AnalyzerManager.cs

@@ -235,6 +235,8 @@ public async Task<GeneratedCodeAnalysisFlags> GetGeneratedCodeAnalysisFlagsAsync
                 });
 
                 // Subscribe for exceptions from lazily evaluated localizable strings in the descriptors.
+                // REVIEW: find out better way to handle these exception handlers. right now, it can leak
+                //         so easily unless ClearAnalyzerState is called from host properly
                 foreach (var descriptor in supportedDiagnostics)
                 {
                     descriptor.Title.OnException += handler;
@@ -257,6 +259,7 @@ public async Task<GeneratedCodeAnalysisFlags> GetGeneratedCodeAnalysisFlagsAsync
             Tuple<ImmutableArray<DiagnosticDescriptor>, EventHandler<Exception>> value;
             lock (_descriptorMap)
             {
+                // REVIEW: how one knows exception handler is same as the one cached?
                 if (_descriptorMap.TryGetValue(analyzer, out value))
                 {
                     return value.Item1;

src/EditorFeatures/Core/Shared/Preview/PreviewSolutionCrawlerRegistrationService.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.Immutable;
 using System.Composition;
-using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Diagnostics;
+using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.SolutionCrawler;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Editor.Shared.Preview
 {
-    [ExportWorkspaceService(typeof(ISolutionCrawlerRegistrationService), WorkspaceKind.Preview), Shared]
-    internal class PreviewSolutionCrawlerRegistrationService : ISolutionCrawlerRegistrationService
+    [ExportWorkspaceServiceFactory(typeof(ISolutionCrawlerRegistrationService), WorkspaceKind.Preview), Shared]
+    internal class PreviewSolutionCrawlerRegistrationServiceFactory : IWorkspaceServiceFactory
     {
-        private static readonly ConditionalWeakTable<Workspace, CancellationTokenSource> s_cancellationTokens =
-            new ConditionalWeakTable<Workspace, CancellationTokenSource>();
-
-        private readonly IIncrementalAnalyzerProvider _provider;
+        private readonly DiagnosticAnalyzerService _analyzerService;
 
         [ImportingConstructor]
-        public PreviewSolutionCrawlerRegistrationService(IDiagnosticAnalyzerService diagnosticService)
+        public PreviewSolutionCrawlerRegistrationServiceFactory(IDiagnosticAnalyzerService analyzerService)
+        {
+            // this service is directly tied to DiagnosticAnalyzerService and
+            // depends on its implementation.
+            _analyzerService = analyzerService as DiagnosticAnalyzerService;
+            Contract.ThrowIfNull(_analyzerService);
+        }
+
+        public IWorkspaceService CreateService(HostWorkspaceServices workspaceServices)
         {
-            _provider = diagnosticService as IIncrementalAnalyzerProvider;
-            Contract.ThrowIfNull(_provider);
+            // to make life time management easier, just create new sevice per new workspace
+            return new Service(this, workspaceServices.Workspace);
         }
 
-        public async void Register(Workspace workspace)
+        // internal for testing
+        internal class Service : ISolutionCrawlerRegistrationService
         {
-            try
+            private readonly PreviewSolutionCrawlerRegistrationServiceFactory _owner;
+            private readonly Workspace _workspace;
+            private readonly CancellationTokenSource _source;
+
+            // since we now have one service for each one specific instance of workspace,
+            // we can have states for this specific workspace.
+            private Task _analyzeTask;
+
+            [ImportingConstructor]
+            public Service(PreviewSolutionCrawlerRegistrationServiceFactory owner, Workspace workspace)
+            {
+                _owner = owner;
+                _workspace = workspace;
+                _source = new CancellationTokenSource();
+            }
+
+            public void Register(Workspace workspace)
+            {
+                // given workspace must be owner of this workspace service
+                Contract.ThrowIfFalse(workspace == _workspace);
+
+                // this can't be called twice
+                Contract.ThrowIfFalse(_analyzeTask == null);
+
+                _analyzeTask = AnalyzeAsync();
+            }
+
+            private async Task AnalyzeAsync()
             {
-                var workerBackOffTimeSpanInMS = workspace.Options.GetOption(InternalSolutionCrawlerOptions.PreviewBackOffTimeSpanInMS);
+                var workerBackOffTimeSpanInMS = _workspace.Options.GetOption(InternalSolutionCrawlerOptions.PreviewBackOffTimeSpanInMS);
+                var diagnosticAnalyzer = _owner._analyzerService.CreateIncrementalAnalyzer(_workspace);
 
-                var analyzer = _provider.CreateIncrementalAnalyzer(workspace);
-                var source = s_cancellationTokens.GetValue(workspace, _ => new CancellationTokenSource());
+                var solution = _workspace.CurrentSolution;
+                var documentIds = _workspace.GetOpenDocumentIds().ToImmutableArray();
 
-                var solution = workspace.CurrentSolution;
-                foreach (var documentId in workspace.GetOpenDocumentIds())
+                try
                 {
-                    var document = solution.GetDocument(documentId);
-                    if (document == null)
+                    foreach (var documentId in documentIds)
                     {
-                        continue;
-                    }
+                        var document = solution.GetDocument(documentId);
+                        if (document == null)
+                        {
+                            continue;
+                        }
 
-                    // delay analyzing
-                    await Task.Delay(workerBackOffTimeSpanInMS).ConfigureAwait(false);
+                        // delay analyzing
+                        await Task.Delay(workerBackOffTimeSpanInMS, _source.Token).ConfigureAwait(false);
 
-                    // do actual analysis
-                    await analyzer.AnalyzeSyntaxAsync(document, source.Token).ConfigureAwait(false);
-                    await analyzer.AnalyzeDocumentAsync(document, bodyOpt: null, cancellationToken: source.Token).ConfigureAwait(false);
+                        // do actual analysis
+                        await diagnosticAnalyzer.AnalyzeSyntaxAsync(document, _source.Token).ConfigureAwait(false);
+                        await diagnosticAnalyzer.AnalyzeDocumentAsync(document, bodyOpt: null, cancellationToken: _source.Token).ConfigureAwait(false);
 
-                    // don't call project one.
+                        // don't call project one.
+                    }
+                }
+                catch (OperationCanceledException)
+                {
+                    // do nothing
                 }
             }
-            catch (OperationCanceledException)
-            {
-                // do nothing
-            }
-        }
 
-        public void Unregister(Workspace workspace, bool blockingShutdown = false)
-        {
-            CancellationTokenSource source;
-            if (s_cancellationTokens.TryGetValue(workspace, out source))
+            public async void Unregister(Workspace workspace, bool blockingShutdown = false)
             {
-                source.Cancel();
+                Contract.ThrowIfFalse(workspace == _workspace);
+                _source.Cancel();
+
+                // wait for analyzer work to be finished
+                await _analyzeTask.ConfigureAwait(false);
+
+                // ask it to reset its stages for the given workspace
+                _owner._analyzerService.ShutdownAnalyzerFrom(_workspace);
             }
         }
     }

src/EditorFeatures/Test/Preview/PreviewWorkspaceTests.cs

@@ -135,7 +135,7 @@ public void TestPreviewServices()
             using (var previewWorkspace = new PreviewWorkspace(MefV1HostServices.Create(TestExportProvider.ExportProviderWithCSharpAndVisualBasic.AsExportProvider())))
             {
                 var service = previewWorkspace.Services.GetService<ISolutionCrawlerRegistrationService>();
-                Assert.True(service is PreviewSolutionCrawlerRegistrationService);
+                Assert.True(service is PreviewSolutionCrawlerRegistrationServiceFactory.Service);
 
                 var persistentService = previewWorkspace.Services.GetService<IPersistentStorageService>();
                 Assert.NotNull(persistentService);

src/Features/Core/Portable/Diagnostics/BaseDiagnosticIncrementalAnalyzer.cs

@@ -219,10 +219,13 @@ public virtual bool NeedsReanalysisOnOptionChanged(object sender, OptionChangedE
             return false;
         }
 
-        public virtual void LogAnalyzerCountSummary()
+        public virtual void Shutdown()
         {
+            // virtual for v1 that got deleted in master
         }
 
+        public abstract void LogAnalyzerCountSummary();
+
         // internal for testing purposes.
         internal Action<Exception, DiagnosticAnalyzer, Diagnostic> GetOnAnalyzerException(ProjectId projectId)
         {

src/Features/Core/Portable/Diagnostics/DiagnosticAnalyzerService_IncrementalAnalyzer.cs

@@ -46,6 +46,16 @@ private BaseDiagnosticIncrementalAnalyzer GetOrCreateIncrementalAnalyzer(Workspa
             return _map.GetValue(workspace, _createIncrementalAnalyzer);
         }
 
+        public void ShutdownAnalyzerFrom(Workspace workspace)
+        {
+            // this should be only called once analyzer associated with the workspace is done.
+            BaseDiagnosticIncrementalAnalyzer analyzer;
+            if (_map.TryGetValue(workspace, out analyzer))
+            {
+                analyzer.Shutdown();
+            }
+        }
+
         private BaseDiagnosticIncrementalAnalyzer CreateIncrementalAnalyzerCallback(Workspace workspace)
         {
             // subscribe to active context changed event for new workspace
@@ -188,6 +198,11 @@ public override Task SynchronizeWithBuildAsync(Workspace workspace, ImmutableDic
             }
             #endregion
 
+            public override void Shutdown()
+            {
+                Analyzer.Shutdown();
+            }
+
             public override void LogAnalyzerCountSummary()
             {
                 Analyzer.LogAnalyzerCountSummary();

src/Features/Core/Portable/Diagnostics/DiagnosticService.cs

@@ -24,7 +24,7 @@ internal partial class DiagnosticService : IDiagnosticService
         private readonly SimpleTaskQueue _eventQueue;
 
         private readonly object _gate;
-        private readonly Dictionary<IDiagnosticUpdateSource, Dictionary<object, Data>> _map;
+        private readonly Dictionary<IDiagnosticUpdateSource, Dictionary<Workspace, Dictionary<object, Data>>> _map;
 
         [ImportingConstructor]
         public DiagnosticService([ImportMany] IEnumerable<Lazy<IAsynchronousOperationListener, FeatureMetadata>> asyncListeners) : this()
@@ -39,7 +39,7 @@ public DiagnosticService([ImportMany] IEnumerable<Lazy<IAsynchronousOperationLis
             _listener = new AggregateAsynchronousOperationListener(asyncListeners, FeatureAttribute.DiagnosticService);
 
             _gate = new object();
-            _map = new Dictionary<IDiagnosticUpdateSource, Dictionary<object, Data>>();
+            _map = new Dictionary<IDiagnosticUpdateSource, Dictionary<Workspace, Dictionary<object, Data>>>();
         }
 
         public event EventHandler<DiagnosticsUpdatedArgs> DiagnosticsUpdated
@@ -114,12 +114,28 @@ private bool UpdateDataMap(IDiagnosticUpdateSource source, DiagnosticsUpdatedArg
                     return false;
                 }
 
-                var diagnosticDataMap = _map.GetOrAdd(source, _ => new Dictionary<object, Data>());
+                // 2 different workspaces (ex, PreviewWorkspaces) can return same Args.Id, we need to
+                // distinguish them. so we separate diagnostics per workspace map.
+                var workspaceMap = _map.GetOrAdd(source, _ => new Dictionary<Workspace, Dictionary<object, Data>>());
+
+                if (args.Diagnostics.Length == 0 && !workspaceMap.ContainsKey(args.Workspace))
+                {
+                    // no new diagnostic, and we don't have workspace for it.
+                    return false;
+                }
+
+                var diagnosticDataMap = workspaceMap.GetOrAdd(args.Workspace, _ => new Dictionary<object, Data>());
 
                 diagnosticDataMap.Remove(args.Id);
                 if (diagnosticDataMap.Count == 0 && args.Diagnostics.Length == 0)
                 {
-                    _map.Remove(source);
+                    workspaceMap.Remove(args.Workspace);
+
+                    if (workspaceMap.Count == 0)
+                    {
+                        _map.Remove(source);
+                    }
+
                     return true;
                 }
 
@@ -258,10 +274,14 @@ public IEnumerable<UpdatedEventArgs> GetDiagnosticsUpdatedEventArgs(Workspace wo
         private void AppendMatchingData(
             IDiagnosticUpdateSource source, Workspace workspace, ProjectId projectId, DocumentId documentId, object id, List<Data> list)
         {
+            Contract.ThrowIfNull(workspace);
+
             lock (_gate)
             {
                 Dictionary<object, Data> current;
-                if (!_map.TryGetValue(source, out current))
+                Dictionary<Workspace, Dictionary<object, Data>> workspaceMap;
+                if (!_map.TryGetValue(source, out workspaceMap) ||
+                    !workspaceMap.TryGetValue(workspace, out current))
                 {
                     return;
                 }
@@ -279,9 +299,9 @@ public IEnumerable<UpdatedEventArgs> GetDiagnosticsUpdatedEventArgs(Workspace wo
 
                 foreach (var data in current.Values)
                 {
-                    if (TryAddData(documentId, data, d => d.DocumentId, list) ||
-                        TryAddData(projectId, data, d => d.ProjectId, list) ||
-                        TryAddData(workspace, data, d => d.Workspace, list))
+                    if (TryAddData(workspace, documentId, data, d => d.DocumentId, list) ||
+                        TryAddData(workspace, projectId, data, d => d.ProjectId, list) ||
+                        TryAddData(workspace, workspace, data, d => d.Workspace, list))
                     {
                         continue;
                     }
@@ -289,13 +309,19 @@ public IEnumerable<UpdatedEventArgs> GetDiagnosticsUpdatedEventArgs(Workspace wo
             }
         }
 
-        private bool TryAddData<T>(T key, Data data, Func<Data, T> keyGetter, List<Data> result) where T : class
+        private bool TryAddData<T>(Workspace workspace, T key, Data data, Func<Data, T> keyGetter, List<Data> result) where T : class
         {
             if (key == null)
             {
                 return false;
             }
 
+            // make sure data is from same workspace. project/documentId can be shared between 2 different workspace
+            if (workspace != data.Workspace)
+            {
+                return false;
+            }
+
             if (key == keyGetter(data))
             {
                 result.Add(data);
@@ -314,15 +340,15 @@ private void AssertIfNull(ImmutableArray<DiagnosticData> diagnostics)
         }
 
         [Conditional("DEBUG")]
-        private void AssertIfNull(DiagnosticData diagnostic)
+        private void AssertIfNull<T>(T obj) where T : class
         {
-            if (diagnostic == null)
+            if (obj == null)
             {
                 Contract.Requires(false, "who returns invalid data?");
             }
         }
 
-        private struct Data : IEquatable<Data>
+        private struct Data
         {
             public readonly Workspace Workspace;
             public readonly ProjectId ProjectId;
@@ -343,27 +369,6 @@ public Data(UpdatedEventArgs args, ImmutableArray<DiagnosticData> diagnostics)
                 this.Id = args.Id;
                 this.Diagnostics = diagnostics;
             }
-
-            public bool Equals(Data other)
-            {
-                return this.Workspace == other.Workspace &&
-                       this.ProjectId == other.ProjectId &&
-                       this.DocumentId == other.DocumentId &&
-                       this.Id == other.Id;
-            }
-
-            public override bool Equals(object obj)
-            {
-                return (obj is Data) && Equals((Data)obj);
-            }
-
-            public override int GetHashCode()
-            {
-                return Hash.Combine(Workspace,
-                       Hash.Combine(ProjectId,
-                       Hash.Combine(DocumentId,
-                       Hash.Combine(Id, 1))));
-            }
         }
     }
 }

src/Features/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer.StateManager.ProjectStates.cs

@@ -28,6 +28,12 @@ public ProjectStates(StateManager owner)
                     _stateMap = new ConcurrentDictionary<ProjectId, Entry>(concurrencyLevel: 2, capacity: 10);
                 }
 
+                public IEnumerable<StateSet> GetStateSets()
+                {
+                    // return existing state sets
+                    return _stateMap.Values.SelectMany(e => e.AnalyzerMap.Values).ToImmutableArray();
+                }
+
                 public IEnumerable<StateSet> GetStateSets(ProjectId projectId)
                 {
                     var map = GetCachedAnalyzerMap(projectId);

src/Features/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer.StateManager.cs

@@ -47,6 +47,15 @@ public IEnumerable<DiagnosticAnalyzer> GetOrCreateAnalyzers(Project project)
                 return _hostStates.GetAnalyzers(project.Language).Concat(_projectStates.GetOrCreateAnalyzers(project));
             }
 
+            /// <summary>
+            /// Return all <see cref="StateSet"/>.
+            /// This will never create new <see cref="StateSet"/> but will return ones already created.
+            /// </summary>
+            public IEnumerable<StateSet> GetStateSets()
+            {
+                return _hostStates.GetStateSets().Concat(_projectStates.GetStateSets());
+            }
+
             /// <summary>
             /// Return <see cref="StateSet"/>s for the given <see cref="ProjectId"/>. 
             /// This will never create new <see cref="StateSet"/> but will return ones already created.

src/Features/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer.StateSet.cs

@@ -83,6 +83,34 @@ public bool ContainsAnyDocumentOrProjectDiagnostics(ProjectId projectId)
                 return !projectState.IsEmpty();
             }
 
+            public IEnumerable<ProjectId> GetProjectsWithDiagnostics()
+            {
+                // quick bail out
+                if (_activeFileStates.IsEmpty && _projectStates.IsEmpty)
+                {
+                    return SpecializedCollections.EmptyEnumerable<ProjectId>();
+                }
+
+                if (_activeFileStates.Count == 1 && _projectStates.IsEmpty)
+                {
+                    // see whether we actually have diagnostics
+                    var kv = _activeFileStates.First();
+                    if (kv.Value.IsEmpty)
+                    {
+                        return SpecializedCollections.EmptyEnumerable<ProjectId>();
+                    }
+
+                    // we do have diagnostics
+                    return SpecializedCollections.SingletonEnumerable(kv.Key.ProjectId);
+                }
+
+                return new HashSet<ProjectId>(
+                    _activeFileStates.Where(kv => !kv.Value.IsEmpty)
+                                     .Select(kv => kv.Key.ProjectId)
+                                     .Concat(_projectStates.Where(kv => !kv.Value.IsEmpty())
+                                                           .Select(kv => kv.Key)));
+            }
+
             public IEnumerable<DocumentId> GetDocumentsWithDiagnostics(ProjectId projectId)
             {
                 HashSet<DocumentId> set = null;

src/Features/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer.cs

@@ -7,6 +7,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Diagnostics.Log;
+using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Roslyn.Utilities;
 
@@ -95,6 +96,24 @@ private void OnProjectAnalyzerReferenceChanged(object sender, ProjectAnalyzerRef
             ClearAllDiagnostics(stateSets, project.Id);
         }
 
+        public override void Shutdown()
+        {
+            var stateSets = _stateManager.GetStateSets();
+
+            Owner.RaiseBulkDiagnosticsUpdated(raiseEvents =>
+            {
+                var handleActiveFile = true;
+                foreach (var stateSet in stateSets)
+                {
+                    var projectIds = stateSet.GetProjectsWithDiagnostics();
+                    foreach (var projectId in projectIds)
+                    {
+                        RaiseProjectDiagnosticsRemoved(stateSet, projectId, stateSet.GetDocumentsWithDiagnostics(projectId), handleActiveFile, raiseEvents);
+                    }
+                }
+            });
+        }
+
         private void ClearAllDiagnostics(ImmutableArray<StateSet> stateSets, ProjectId projectId)
         {
             Owner.RaiseBulkDiagnosticsUpdated(raiseEvents =>

src/Features/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer_IncrementalAnalyzer.cs

@@ -14,7 +14,6 @@
 
 namespace Microsoft.CodeAnalysis.Diagnostics.EngineV2
 {
-    // TODO: make it to use cache
     internal partial class DiagnosticIncrementalAnalyzer : BaseDiagnosticIncrementalAnalyzer
     {
         public override Task AnalyzeSyntaxAsync(Document document, CancellationToken cancellationToken)

src/Features/Core/Portable/Diagnostics/HostAnalyzerManager.cs

@@ -166,9 +166,12 @@ public bool IsAnalyzerSuppressed(DiagnosticAnalyzer analyzer, Project project)
                 return false;
             }
 
+            // don't capture project
+            var projectId = project.Id;
+
             // Skip telemetry logging for supported diagnostics, as that can cause an infinite loop.
             Action<Exception, DiagnosticAnalyzer, Diagnostic> onAnalyzerException = (ex, a, diagnostic) =>
-                    AnalyzerHelper.OnAnalyzerException_NoTelemetryLogging(ex, a, diagnostic, _hostDiagnosticUpdateSource, project.Id);
+                    AnalyzerHelper.OnAnalyzerException_NoTelemetryLogging(ex, a, diagnostic, _hostDiagnosticUpdateSource, projectId);
 
             return CompilationWithAnalyzers.IsDiagnosticAnalyzerSuppressed(analyzer, options, onAnalyzerException);
         }