Fix PythonAnalyzerSession leak (microsoft#866)

Author: AlexanderSher

src/Analysis/Ast/Impl/Analyzer/PythonAnalyzer.cs

@@ -16,6 +16,7 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Python.Analysis.Dependencies;
@@ -39,17 +40,20 @@ public sealed class PythonAnalyzer : IPythonAnalyzer, IDisposable {
         private readonly DisposeToken _disposeToken = DisposeToken.Create<PythonAnalyzer>();
         private readonly object _syncObj = new object();
         private readonly AsyncManualResetEvent _analysisCompleteEvent = new AsyncManualResetEvent();
+        private readonly Action<Task> _startNextSession;
         private readonly ProgressReporter _progress;
         private readonly ILogger _log;
         private readonly int _maxTaskRunning = Environment.ProcessorCount;
         private int _version;
-        private PythonAnalyzerSession _session;
+        private PythonAnalyzerSession _currentSession;
+        private PythonAnalyzerSession _nextSession;
 
         public PythonAnalyzer(IServiceManager services) {
             _services = services;
             _log = services.GetService<ILogger>();
             _dependencyResolver = new DependencyResolver<AnalysisModuleKey, PythonAnalyzerEntry>();
             _analysisCompleteEvent.Set();
+            _startNextSession = StartNextSession;
 
             _progress = new ProgressReporter(services.GetService<IProgressService>());
         }
@@ -195,24 +199,56 @@ private void AnalyzeDocument(AnalysisModuleKey key, PythonAnalyzerEntry entry, I
                 LoadMissingDocuments(entry.Module.Interpreter, walker.MissingKeys);
             }
 
-            PythonAnalyzerSession currentSession, nextSession;
+            if (TryCreateSession(walker, entry, cancellationToken, out var session)) {
+                session.Start(true);
+            }
+        }
+
+        private bool TryCreateSession(IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> walker, PythonAnalyzerEntry entry, CancellationToken cancellationToken, out PythonAnalyzerSession session) {
             lock (_syncObj) {
-                currentSession = _session;
-                if (currentSession != null && currentSession.Version > walker.Version) {
-                    return;
+                if (_currentSession == null) {
+                    _currentSession = session = CreateSession(walker, null, cancellationToken);
+                    return true;
+                }
+
+                if (_currentSession.Version > walker.Version || _nextSession != null && _nextSession.Version > walker.Version) {
+                    session = null;
+                    return false;
+                }
+
+                if (_version > walker.Version && (!_currentSession.IsCompleted || walker.AffectedValues.GetCount(e => e.NotAnalyzed) < _maxTaskRunning)) {
+                    session = null;
+                    return false;
+                }
+                
+                if (_currentSession.IsCompleted) {
+                    _currentSession = session = CreateSession(walker, null, cancellationToken);
+                    return true;
                 }
 
-                if (_version > walker.Version && currentSession != null && (!currentSession.IsCompleted || walker.AffectedValues.GetCount(e => e.NotAnalyzed) < _maxTaskRunning)) {
+                _currentSession.Cancel();
+                _nextSession = session = CreateSession(walker, entry.IsUserModule ? entry : null, cancellationToken);
+                return entry.IsUserModule;
+            }
+        }
+
+        private void StartNextSession(Task task) {
+            PythonAnalyzerSession session;
+            lock (_syncObj) {
+                if (_nextSession == null) {
                     return;
                 }
 
-                nextSession = new PythonAnalyzerSession(_services, _progress, _analysisCompleteEvent, _disposeToken.CancellationToken, cancellationToken, walker, _version);
-                _session = nextSession;
+                _currentSession = session = _nextSession;
+                _nextSession = null;
             }
 
-            nextSession.Start(currentSession, entry);
+            session.Start(false);
         }
 
+        private PythonAnalyzerSession CreateSession(IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> walker, PythonAnalyzerEntry entry, CancellationToken cancellationToken) 
+            => new PythonAnalyzerSession(_services, _progress, _analysisCompleteEvent, _startNextSession, _disposeToken.CancellationToken, cancellationToken, walker, _version, entry);
+
         private void LoadMissingDocuments(IPythonInterpreter interpreter, ImmutableArray<AnalysisModuleKey> missingKeys) {
             foreach (var (moduleName, _, isTypeshed) in missingKeys) {
                 var moduleResolution = isTypeshed ? interpreter.TypeshedResolution : interpreter.ModuleResolution;

src/Analysis/Ast/Impl/Analyzer/PythonAnalyzerSession.cs

@@ -14,7 +14,6 @@
 // permissions and limitations under the License.
 
 using System;
-using System.ComponentModel;
 using System.Diagnostics;
 using System.Linq;
 using System.Threading;
@@ -25,7 +24,6 @@
 using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Analysis.Types;
 using Microsoft.Python.Core;
-using Microsoft.Python.Core.Diagnostics;
 using Microsoft.Python.Core.Logging;
 using Microsoft.Python.Core.Services;
 using Microsoft.Python.Parsing.Ast;
@@ -34,21 +32,21 @@ namespace Microsoft.Python.Analysis.Analyzer {
     internal sealed class PythonAnalyzerSession {
         private readonly int _maxTaskRunning = Environment.ProcessorCount;
         private readonly object _syncObj = new object();
-        private readonly Action<Task> _startNextSession;
 
+        private readonly IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> _walker;
+        private readonly PythonAnalyzerEntry _entry;
         private readonly CancellationTokenSource _cts;
+        private readonly Action<Task> _startNextSession;
         private readonly IServiceManager _services;
         private readonly AsyncManualResetEvent _analysisCompleteEvent;
         private readonly IDiagnosticsService _diagnosticsService;
         private readonly IProgressReporter _progress;
         private readonly IPythonAnalyzer _analyzer;
         private readonly ILogger _log;
 
-        private State _state = State.NotStarted;
-        private IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> _walker;
+        private State _state;
+        private bool _isCanceled;
         private int _runningTasks;
-        private int _version;
-        private PythonAnalyzerSession _nextSession;
 
         public bool IsCompleted {
             get {
@@ -58,21 +56,25 @@ public bool IsCompleted {
             }
         }
 
-        public int Version => _version;
+        public int Version { get; }
 
         public PythonAnalyzerSession(IServiceManager services,
             IProgressReporter progress,
             AsyncManualResetEvent analysisCompleteEvent,
+            Action<Task> startNextSession,
             CancellationToken analyzerToken,
             CancellationToken sessionToken, 
             IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> walker,
-            int version) {
+            int version,
+            PythonAnalyzerEntry entry) {
 
             _services = services;
             _analysisCompleteEvent = analysisCompleteEvent;
-            _startNextSession = StartNextSession;
-            _version = version;
+            _startNextSession = startNextSession;
+            Version = version;
             _walker = walker;
+            _entry = entry;
+            _state = State.NotStarted;
             _cts = CancellationTokenSource.CreateLinkedTokenSource(analyzerToken, sessionToken);
 
             _diagnosticsService = _services.GetService<IDiagnosticsService>();
@@ -81,119 +83,64 @@ public PythonAnalyzerSession(IServiceManager services,
             _log = _services.GetService<ILogger>();
         }
 
-        public void Start(PythonAnalyzerSession previousSession, PythonAnalyzerEntry entry) {
-            IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> walker;
-            State previousSessionState;
-
+        public void Start(bool analyzeEntry) {
             lock (_syncObj) {
-                walker = _walker;
                 if (_state != State.NotStarted) {
-                    _cts.Dispose();
+                    analyzeEntry = false;
+                } else if (_state == State.Completed) {
                     return;
-                }
-
-                previousSessionState = previousSession?.CancelOrSchedule(this, _version) ?? State.Completed;
-                if (previousSessionState == State.Completed) {
+                } else {
                     _state = State.Started;
-                    _walker = null;
                 }
             }
 
-            switch (previousSessionState) {
-                case State.Started when entry.IsUserModule:
-                    StartAnalysis(entry, walker.Version);
-                    break;
-                case State.Completed:
-                    Start(walker);
-                    break;
+            if (analyzeEntry && _entry != null) {
+                Task.Run(() => Analyze(_entry, _walker.Version, _cts.Token), _cts.Token).DoNotWait();
+            } else {
+                StartAsync(_walker).ContinueWith(_startNextSession).DoNotWait();
             }
         }
 
-        private void StartNextSession(Task task) {
-            PythonAnalyzerSession nextSession;
+        public void Cancel() {
             lock (_syncObj) {
-                _walker = null;
-                nextSession = _nextSession;
-                _nextSession = null;
+                _isCanceled = true;
             }
-
-            nextSession?.TryStart();
         }
 
-        private State CancelOrSchedule(PythonAnalyzerSession nextSession, int version) {
-            lock (_syncObj) {
-                Check.InvalidOperation(_nextSession == null);
-                Check.InvalidOperation(nextSession != null);
-                switch (_state) {
-                    case State.NotStarted:
-                        _state = State.Completed;
-                        break;
-                    case State.Started:
-                        _nextSession = nextSession;
-                        Interlocked.Exchange(ref _version, version);
-                        break;
-                    case State.Completed:
-                        break;
-                    default:
-                        throw new ArgumentOutOfRangeException();
-                }
-
-                return _state;
-            }
-        }
-
-        private void TryStart() {
-            IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> walker = default;
-            lock (_syncObj) {
-                if (_state == State.NotStarted) {
-                    _state = State.Started;
-                    walker = _walker;
-                }
-
-                _walker = null;
-            }
-
-            if (walker != default) {
-                Start(walker);
-            }
-        }
-
-        private void Start(IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> walker) => StartAsync(walker).ContinueWith(_startNextSession).DoNotWait();
-
         private async Task StartAsync(IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> walker) {
-            int version;
             lock (_syncObj) {
-                version = _version;
                 var notAnalyzed = walker.AffectedValues.Count(e => e.NotAnalyzed);
 
-                if (version > walker.Version && notAnalyzed < _maxTaskRunning) {
+                if (_isCanceled && notAnalyzed < _maxTaskRunning) {
                     _state = State.Completed;
                     _cts.Dispose();
                     return;
                 }
             }
-            
+
+            var cancellationToken = _cts.Token;
             var stopWatch = Stopwatch.StartNew();
             foreach (var affectedEntry in walker.AffectedValues) {
-                affectedEntry.Invalidate(version);
+                affectedEntry.Invalidate(Version);
             }
 
             var originalRemaining = walker.Remaining;
             var remaining = originalRemaining;
             try {
                 _log?.Log(TraceEventType.Verbose, $"Analysis version {walker.Version} of {originalRemaining} entries has started.");
-                remaining = await AnalyzeAffectedEntriesAsync(walker, stopWatch, _cts.Token);
+                remaining = await AnalyzeAffectedEntriesAsync(walker, stopWatch, cancellationToken);
             } finally {
                 _cts.Dispose();
                 stopWatch.Stop();
 
+                bool isCanceled;
                 lock (_syncObj) {
-                    if (_version == walker.Version) {
-                        _progress.ReportRemaining(walker.Remaining);
-                    }
-
+                    isCanceled = _isCanceled;
                     _state = State.Completed;
-                    _cts.Dispose();
+                }
+
+                if (!isCanceled) {
+                    _progress.ReportRemaining(walker.Remaining);
                 }
 
                 if (_log != null) {
@@ -212,7 +159,12 @@ private async Task<int> AnalyzeAffectedEntriesAsync(IDependencyChainWalker<Analy
             IDependencyChainNode<PythonAnalyzerEntry> node;
             var remaining = 0;
             while ((node = await walker.GetNextAsync(cancellationToken)) != null) {
-                if (_version > walker.Version && !node.Value.NotAnalyzed) {
+                bool isCanceled;
+                lock (_syncObj) {
+                    isCanceled = _isCanceled;
+                }
+
+                if (isCanceled && !node.Value.NotAnalyzed) {
                     remaining++;
                     node.Skip();
                     continue;
@@ -227,7 +179,12 @@ private async Task<int> AnalyzeAffectedEntriesAsync(IDependencyChainWalker<Analy
 
             if (walker.MissingKeys.All(k => k.IsTypeshed)) {
                 Interlocked.Exchange(ref _runningTasks, 0);
-                if (_version == walker.Version) {
+                bool isCanceled;
+                lock (_syncObj) {
+                    isCanceled = _isCanceled;
+                }
+
+                if (!isCanceled) {
                     _analysisCompleteEvent.Set();
                 }
             }
@@ -269,16 +226,18 @@ private void Analyze(IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEnt
                 node.Commit();
                 _log?.Log(TraceEventType.Verbose, $"Analysis of {module.Name}({module.ModuleType}) failed.");
             } finally {
-                if (_version == walker.Version) {
+                bool isCanceled;
+                lock (_syncObj) {
+                    isCanceled = _isCanceled;
+                }
+
+                if (!isCanceled) {
                     _progress.ReportRemaining(walker.Remaining);
                 }
                 Interlocked.Decrement(ref _runningTasks);
             }
         }
 
-        private void StartAnalysis(PythonAnalyzerEntry entry, int version)
-            => Task.Run(() => Analyze(entry, version, _cts.Token), _cts.Token).DoNotWait();
-
         private void Analyze(PythonAnalyzerEntry entry, int version, CancellationToken cancellationToken) {
             var stopWatch = Stopwatch.StartNew();
             try {