Fix #923: Runaway memory usage when working with msticpy (#972)

* Fix #923: Runaway memory usage when working with msticpy

* - Fix sorting
- Mark session canceled if ast isn't set yet

Author: AlexanderSher

src/Analysis/Ast/Impl/Analyzer/Evaluation/ExpressionEval.cs

@@ -16,7 +16,6 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
-using System.Reflection;
 using Microsoft.Python.Analysis.Analyzer.Symbols;
 using Microsoft.Python.Analysis.Diagnostics;
 using Microsoft.Python.Analysis.Modules;

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

@@ -197,43 +197,73 @@ private void AnalyzeDocument(AnalysisModuleKey key, PythonAnalyzerEntry entry, I
             _analysisCompleteEvent.Reset();
             _log?.Log(TraceEventType.Verbose, $"Analysis of {entry.Module.Name}({entry.Module.ModuleType}) queued");
 
-            var snapshot = _dependencyResolver.NotifyChanges(key, entry, dependencies);
-            
+            var graphVersion = _dependencyResolver.ChangeValue(key, entry, dependencies);
+
             lock (_syncObj) {
-                if (_version > snapshot.Version) {
+                if (_version > graphVersion) {
                     return;
                 }
 
-                _version = snapshot.Version;
+                _version = graphVersion;
+                _currentSession?.Cancel();
             }
 
-            if (snapshot.MissingKeys.Count > 0) {
-                LoadMissingDocuments(entry.Module.Interpreter, snapshot.MissingKeys);
-            }
+            UpdateDependentEntriesDepth(entry, dependencies, graphVersion);
 
-            if (TryCreateSession(snapshot, entry, cancellationToken, out var session)) {
+            if (TryCreateSession(graphVersion, entry, cancellationToken, out var session)) {
                 session.Start(true);
             }
         }
 
-        private bool TryCreateSession(DependencyGraphSnapshot<AnalysisModuleKey, PythonAnalyzerEntry> snapshot, PythonAnalyzerEntry entry, CancellationToken cancellationToken, out PythonAnalyzerSession session) {
+        private void UpdateDependentEntriesDepth(PythonAnalyzerEntry entry, ImmutableArray<AnalysisModuleKey> dependentKeys, int graphVersion) {
+            if (dependentKeys.Count == 0) {
+                return;
+            }
+
+            var dependentEntries = new List<PythonAnalyzerEntry>();
+            lock (_syncObj) {
+                if (_version > graphVersion) {
+                    return;
+                }
+
+                foreach (var key in dependentKeys) {
+                    if (_analysisEntries.TryGetValue(key, out var value)) {
+                        dependentEntries.Add(value);
+                    }
+                }
+            }
+
+            if (dependentEntries.Count == 0) {
+                return;
+            }
+
+            var depth = entry.Depth;
+            foreach (var dependentEntry in dependentEntries) {
+                dependentEntry.SetDepth(graphVersion, depth);
+            }
+        }
+
+        private bool TryCreateSession(int graphVersion, PythonAnalyzerEntry entry, CancellationToken cancellationToken, out PythonAnalyzerSession session) {
             var analyzeUserModuleOutOfOrder = false;
             lock (_syncObj) {
                 if (_currentSession != null) {
-                    if (_currentSession.Version > snapshot.Version || _nextSession != null && _nextSession.Version > snapshot.Version) {
+                    if (_currentSession.Version > graphVersion || _nextSession != null && _nextSession.Version > graphVersion) {
                         session = null;
                         return false;
                     }
 
                     analyzeUserModuleOutOfOrder = !_currentSession.IsCompleted && entry.IsUserModule && _currentSession.AffectedEntriesCount >= _maxTaskRunning;
-                    if (_version > snapshot.Version && analyzeUserModuleOutOfOrder) {
+                    if (_version > graphVersion && analyzeUserModuleOutOfOrder) {
                         session = CreateSession(null, entry, cancellationToken);
                         return true;
                     }
                 }
             }
 
-            if (!_dependencyResolver.TryCreateWalker(snapshot.Version, out var walker)) {
+            var snapshot = _dependencyResolver.CurrentGraphSnapshot;
+            LoadMissingDocuments(entry.Module.Interpreter, snapshot.MissingKeys);
+
+            if (!_dependencyResolver.TryCreateWalker(snapshot, out var walker)) {
                 session = null;
                 return false;
             }
@@ -254,7 +284,6 @@ private bool TryCreateSession(DependencyGraphSnapshot<AnalysisModuleKey, PythonA
                     return true;
                 }
 
-                _currentSession.Cancel();
                 _nextSession = session = CreateSession(walker, analyzeUserModuleOutOfOrder ? entry : null, cancellationToken);
                 return analyzeUserModuleOutOfOrder;
             }
@@ -278,9 +307,37 @@ private PythonAnalyzerSession CreateSession(IDependencyChainWalker<AnalysisModul
             => 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) {
+            if (missingKeys.Count == 0) {
+                return;
+            }
+            
+            var foundKeys = ImmutableArray<AnalysisModuleKey>.Empty;
+            foreach (var missingKey in missingKeys) {
+                lock (_syncObj) {
+                    if (_analysisEntries.TryGetValue(missingKey, out _)) {
+                        continue;
+                    }
+                }
+
+                var (moduleName, _, isTypeshed) = missingKey;
                 var moduleResolution = isTypeshed ? interpreter.TypeshedResolution : interpreter.ModuleResolution;
-                moduleResolution.GetOrLoadModule(moduleName);
+                var module = moduleResolution.GetOrLoadModule(moduleName);
+                if (module != null && module.ModuleType != ModuleType.Unresolved) {
+                    foundKeys = foundKeys.Add(missingKey);
+                }
+            }
+
+            if (foundKeys.Count > 0) {
+                foreach (var foundKey in foundKeys) {
+                    PythonAnalyzerEntry entry;
+                    lock (_syncObj) {
+                        if (!_analysisEntries.TryGetValue(foundKey, out entry)) {
+                            continue;
+                        }
+                    }
+
+                    _dependencyResolver.TryAddValue(foundKey, entry, ImmutableArray<AnalysisModuleKey>.Empty);
+                }
             }
         }
     }

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

@@ -39,6 +39,7 @@ internal sealed class PythonAnalyzerEntry {
         private HashSet<AnalysisModuleKey> _analysisDependencies;
         private int _bufferVersion;
         private int _analysisVersion;
+        private int _depth;
 
         public IPythonModule Module {
             get {
@@ -74,11 +75,21 @@ public int AnalysisVersion {
             }
         }
 
+        public int Depth {
+            get {
+                lock (_syncObj) {
+                    return _depth;
+                }
+            }
+        }
+
         public bool NotAnalyzed => PreviousAnalysis is EmptyAnalysis;
 
         public PythonAnalyzerEntry(EmptyAnalysis emptyAnalysis) {
             _previousAnalysis = emptyAnalysis;
+            _module = emptyAnalysis.Document;
             _isUserModule = emptyAnalysis.Document.ModuleType == ModuleType.User;
+            _depth = _isUserModule ? 0 : -1;
 
             _bufferVersion = -1;
             _analysisVersion = 0;
@@ -92,10 +103,24 @@ public bool IsValidVersion(int version, out IPythonModule module, out PythonAst
             lock (_syncObj) {
                 module = _module;
                 ast = _ast;
+                if (ast == null || module == null) {
+                    return false;
+                }
+
                 return _previousAnalysis is EmptyAnalysis || _isUserModule || _analysisVersion <= version;
             }
         }
 
+        public void SetDepth(int version, int depth) {
+            lock (_syncObj) {
+                if (_analysisVersion > version) {
+                    return;
+                }
+
+                _depth = _depth == -1 ? depth : Math.Min(_depth, depth);
+            }
+        }
+
         public void TrySetAnalysis(IDocumentAnalysis analysis, int version) {
             lock (_syncObj) {
                 if (_previousAnalysis is EmptyAnalysis) {
@@ -139,7 +164,7 @@ public void TryCancel(OperationCanceledException oce, int version) {
 
         public void Invalidate(int analysisVersion) {
             lock (_syncObj) {
-                if (_analysisVersion >= analysisVersion) {
+                if (_analysisVersion >= analysisVersion || !_analysisTcs.Task.IsCompleted) {
                     return;
                 }
 
@@ -176,7 +201,6 @@ public bool Invalidate(ImmutableArray<IPythonModule> analysisDependencies, int a
                     return false;
                 }
 
-                UpdateAnalysisTcs(analysisVersion);
                 if (_analysisDependencies == null) {
                     _analysisDependencies = dependenciesHashSet;
                 } else {
@@ -187,6 +211,7 @@ public bool Invalidate(ImmutableArray<IPythonModule> analysisDependencies, int a
                     }
                 }
 
+                UpdateAnalysisTcs(analysisVersion);
                 dependencies = _parserDependencies != null
                     ? ImmutableArray<AnalysisModuleKey>.Create(_parserDependencies.Union(_analysisDependencies).ToArray())
                     : ImmutableArray<AnalysisModuleKey>.Create(_analysisDependencies);

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

@@ -193,7 +193,7 @@ private void LogResults(double elapsed, int originalRemaining, int remaining, in
             if (_log == null) {
                 return;
             }
-
+            
             if (remaining == 0) {
                 _log.Log(TraceEventType.Verbose, $"Analysis version {version} of {originalRemaining} entries has been completed in {elapsed} ms.");
             } else if (remaining < originalRemaining) {
@@ -208,8 +208,8 @@ private async Task<int> AnalyzeAffectedEntriesAsync(Stopwatch stopWatch, Cancell
             var remaining = 0;
             var ace = new AsyncCountdownEvent(0);
 
+            bool isCanceled;
             while ((node = await _walker.GetNextAsync(cancellationToken)) != null) {
-                bool isCanceled;
                 lock (_syncObj) {
                     isCanceled = _isCanceled;
                 }
@@ -229,16 +229,18 @@ private async Task<int> AnalyzeAffectedEntriesAsync(Stopwatch stopWatch, Cancell
 
             await ace.WaitAsync(cancellationToken);
 
-            if (_walker.MissingKeys.All(k => k.IsTypeshed)) {
-                Interlocked.Exchange(ref _runningTasks, 0);
-                bool isCanceled;
-                lock (_syncObj) {
-                    isCanceled = _isCanceled;
-                }
+            lock (_syncObj) {
+                isCanceled = _isCanceled;
+            }
 
+            if (_walker.MissingKeys.Count == 0 || _walker.MissingKeys.All(k => k.IsTypeshed)) {
+                Interlocked.Exchange(ref _runningTasks, 0);
+                
                 if (!isCanceled) {
                     _analysisCompleteEvent.Set();
                 }
+            } else if (!isCanceled && _log != null && _log.LogLevel >= TraceEventType.Verbose) {
+                _log?.Log(TraceEventType.Verbose, $"Missing keys: {string.Join(", ", _walker.MissingKeys)}");
             }
 
             return remaining;
@@ -259,6 +261,11 @@ private void Analyze(IDependencyChainNode<PythonAnalyzerEntry> node, AsyncCountd
                 ace?.AddOne();
                 var entry = node.Value;
                 if (!entry.IsValidVersion(_walker.Version, out module, out var ast)) {
+                    if (ast == null) {
+                        // Entry doesn't have ast yet. There should be at least one more session.
+                        Cancel();
+                    }
+
                     _log?.Log(TraceEventType.Verbose, $"Analysis of {module.Name}({module.ModuleType}) canceled.");
                     node.Skip();
                     return;
@@ -296,6 +303,11 @@ private void Analyze(PythonAnalyzerEntry entry, int version, CancellationToken c
             var stopWatch = Stopwatch.StartNew();
             try {
                 if (!entry.IsValidVersion(version, out var module, out var ast)) {
+                    if (ast == null) {
+                        // Entry doesn't have ast yet. There should be at least one more session.
+                        Cancel();
+                    }
+
                     _log?.Log(TraceEventType.Verbose, $"Analysis of {module.Name}({module.ModuleType}) canceled.");
                     return;
                 }