Prevent SO on documentation fetch (#1106)

* Prevent SO on documentation fetch

* Remove disposed exception

* Add Pop

* PR feedback

* Add lock to push/pop

* Add test

* Remove lock

Author: Mikhail Arkhipov

src/Analysis/Ast/Impl/Modules/PythonModule.cs

@@ -135,7 +135,8 @@ public virtual string Documentation {
                     _documentation = m.TryGetConstant<string>(out var s) ? s : string.Empty;
                     if (string.IsNullOrEmpty(_documentation)) {
                         m = GetMember($"_{Name}");
-                        _documentation = m?.GetPythonType()?.Documentation;
+                        var t = m?.GetPythonType();
+                        _documentation = t != null && !t.Equals(this) ? t.Documentation : null;
                         if (string.IsNullOrEmpty(_documentation)) {
                             _documentation = TryGetDocFromModuleInitFile();
                         }
@@ -496,7 +497,7 @@ private string TryGetDocFromModuleInitFile() {
                         // Also, handle quadruple+ quotes.
                         line = line.Trim();
                         line = line.All(c => c == quote[0]) ? quote : line;
-                        if (line.EndsWithOrdinal(quote) && line.IndexOf(quote, StringComparison.Ordinal) < line.LastIndexOf(quote, StringComparison.Ordinal)) { 
+                        if (line.EndsWithOrdinal(quote) && line.IndexOf(quote, StringComparison.Ordinal) < line.LastIndexOf(quote, StringComparison.Ordinal)) {
                             return line.Substring(quote.Length, line.Length - 2 * quote.Length).Trim();
                         }
                         var sb = new StringBuilder();

src/Analysis/Ast/Impl/Types/PythonClassType.cs

@@ -17,7 +17,6 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
-using System.Threading;
 using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Analysis.Specializations.Typing;
 using Microsoft.Python.Analysis.Types.Collections;
@@ -33,11 +32,11 @@ namespace Microsoft.Python.Analysis.Types {
     [DebuggerDisplay("Class {Name}")]
     internal class PythonClassType : PythonType, IPythonClassType, IPythonTemplateType, IEquatable<IPythonClassType> {
         private static readonly string[] _classMethods = { "mro", "__dict__", @"__weakref__" };
-        private readonly object _lock = new object();
-        private readonly AsyncLocal<IPythonClassType> _processing = new AsyncLocal<IPythonClassType>();
+        private IPythonClassType _processing;
         private List<IPythonType> _bases;
         private IReadOnlyList<IPythonType> _mro;
         private Dictionary<string, IPythonType> _genericParameters;
+        private string _documentation;
 
         // For tests
         internal PythonClassType(string name, Location location)
@@ -59,34 +58,31 @@ public PythonClassType(
 
         public override IEnumerable<string> GetMemberNames() {
             var names = new HashSet<string>();
-            lock (_lock) {
-                names.UnionWith(Members.Keys);
-            }
+            names.UnionWith(Members.Keys);
             foreach (var m in Mro.Skip(1)) {
                 names.UnionWith(m.GetMemberNames());
             }
             return DeclaringModule.Interpreter.LanguageVersion.Is3x() ? names.Concat(_classMethods).Distinct() : names;
         }
 
         public override IMember GetMember(string name) {
-            IMember member;
-            lock (_lock) {
-                if (Members.TryGetValue(name, out member)) {
-                    return member;
-                }
+            // Push/Pop should be lock protected.
+            if (Members.TryGetValue(name, out var member)) {
+                return member;
+            }
 
-                // Special case names that we want to add to our own Members dict
-                var is3x = DeclaringModule.Interpreter.LanguageVersion.Is3x();
-                switch (name) {
-                    case "__mro__":
-                    case "mro":
-                        return is3x ? PythonCollectionType.CreateList(DeclaringModule.Interpreter, Mro) : UnknownType;
-                    case "__dict__":
-                        return is3x ? DeclaringModule.Interpreter.GetBuiltinType(BuiltinTypeId.Dict) : UnknownType;
-                    case @"__weakref__":
-                        return is3x ? DeclaringModule.Interpreter.GetBuiltinType(BuiltinTypeId.Object) : UnknownType;
-                }
+            // Special case names that we want to add to our own Members dict
+            var is3x = DeclaringModule.Interpreter.LanguageVersion.Is3x();
+            switch (name) {
+                case "__mro__":
+                case "mro":
+                    return is3x ? PythonCollectionType.CreateList(DeclaringModule.Interpreter, Mro) : UnknownType;
+                case "__dict__":
+                    return is3x ? DeclaringModule.Interpreter.GetBuiltinType(BuiltinTypeId.Dict) : UnknownType;
+                case @"__weakref__":
+                    return is3x ? DeclaringModule.Interpreter.GetBuiltinType(BuiltinTypeId.Object) : UnknownType;
             }
+
             if (Push(this)) {
                 try {
                     foreach (var m in Mro.Reverse()) {
@@ -104,19 +100,33 @@ public override IMember GetMember(string name) {
 
         public override string Documentation {
             get {
-                // Try doc from the type (class definition AST node).
-                var doc = base.Documentation;
-                // Try docs __init__.
-                if (string.IsNullOrEmpty(doc)) {
-                    var init = GetMember("__init__") as IPythonFunctionType;
-                    doc = init?.DeclaringType == this ? init.Documentation : null;
+                if (!string.IsNullOrEmpty(_documentation)) {
+                    return _documentation;
                 }
-                // Try bases.
-                if (string.IsNullOrEmpty(doc) && Bases != null) {
-                    var o = DeclaringModule.Interpreter.GetBuiltinType(BuiltinTypeId.Object);
-                    doc = Bases.Except(Enumerable.Repeat(o, 1)).FirstOrDefault(b => !string.IsNullOrEmpty(b?.Documentation))?.Documentation;
+                // Make sure we do not cycle through bases back here.
+                if (!Push(this)) {
+                    return null;
                 }
-                return doc;
+                try {
+                    // Try doc from the type first (class definition AST node).
+                    _documentation = base.Documentation;
+                    if (string.IsNullOrEmpty(_documentation)) {
+                        // If not present, try docs __init__. IPythonFunctionType handles
+                        // __init__ in a special way so there is no danger of call coming
+                        // back here and causing stack overflow.
+                        _documentation = (GetMember("__init__") as IPythonFunctionType)?.Documentation;
+                    }
+
+                    if (string.IsNullOrEmpty(_documentation) && Bases != null) {
+                        // If still not found, try bases. 
+                        var o = DeclaringModule.Interpreter.GetBuiltinType(BuiltinTypeId.Object);
+                        _documentation = Bases.FirstOrDefault(b => b != o && !string.IsNullOrEmpty(b?.Documentation))
+                            ?.Documentation;
+                    }
+                } finally {
+                    Pop();
+                }
+                return _documentation;
             }
         }
 
@@ -158,17 +168,15 @@ public override IMember Index(IPythonInstance instance, object index) {
 
         public IReadOnlyList<IPythonType> Mro {
             get {
-                lock (_lock) {
-                    if (_mro != null) {
-                        return _mro;
-                    }
-                    if (_bases == null) {
-                        return new IPythonType[] { this };
-                    }
-                    _mro = new IPythonType[] { this };
-                    _mro = CalculateMro(this);
+                if (_mro != null) {
                     return _mro;
                 }
+                if (_bases == null) {
+                    return new IPythonType[] { this };
+                }
+                _mro = new IPythonType[] { this };
+                _mro = CalculateMro(this);
+                return _mro;
             }
         }
 
@@ -177,36 +185,34 @@ public IReadOnlyDictionary<string, IPythonType> GenericParameters
         #endregion
 
         internal void SetBases(IEnumerable<IPythonType> bases) {
-            lock (_lock) {
-                if (_bases != null) {
-                    return; // Already set
-                }
-
-                bases = bases != null ? bases.Where(b => !b.GetPythonType().IsUnknown()).ToArray() : Array.Empty<IPythonType>();
-                // For Python 3+ attach object as a base class by default except for the object class itself.
-                if (DeclaringModule.Interpreter.LanguageVersion.Is3x() && DeclaringModule.ModuleType != ModuleType.Builtins) {
-                    var objectType = DeclaringModule.Interpreter.GetBuiltinType(BuiltinTypeId.Object);
-                    // During processing of builtins module some types may not be available yet.
-                    // Specialization will attach proper base at the end.
-                    Debug.Assert(!objectType.IsUnknown());
-                    if (!bases.Any(b => objectType.Equals(b))) {
-                        bases = bases.Concat(Enumerable.Repeat(objectType, 1));
-                    }
-                }
+            if (_bases != null) {
+                return; // Already set
+            }
 
-                _bases = bases.ToList();
-                if (_bases.Count > 0) {
-                    AddMember("__base__", _bases[0], true);
-                }
-                // Invalidate MRO
-                _mro = null;
-                if (DeclaringModule is BuiltinsPythonModule) {
-                    // TODO: If necessary, we can set __bases__ on builtins when the module is fully analyzed.
-                    return;
+            bases = bases != null ? bases.Where(b => !b.GetPythonType().IsUnknown()).ToArray() : Array.Empty<IPythonType>();
+            // For Python 3+ attach object as a base class by default except for the object class itself.
+            if (DeclaringModule.Interpreter.LanguageVersion.Is3x() && DeclaringModule.ModuleType != ModuleType.Builtins) {
+                var objectType = DeclaringModule.Interpreter.GetBuiltinType(BuiltinTypeId.Object);
+                // During processing of builtins module some types may not be available yet.
+                // Specialization will attach proper base at the end.
+                Debug.Assert(!objectType.IsUnknown());
+                if (!bases.Any(b => objectType.Equals(b))) {
+                    bases = bases.Concat(Enumerable.Repeat(objectType, 1));
                 }
+            }
 
-                AddMember("__bases__", PythonCollectionType.CreateList(DeclaringModule.Interpreter, _bases), true);
+            _bases = bases.ToList();
+            if (_bases.Count > 0) {
+                AddMember("__base__", _bases[0], true);
+            }
+            // Invalidate MRO
+            _mro = null;
+            if (DeclaringModule is BuiltinsPythonModule) {
+                // TODO: If necessary, we can set __bases__ on builtins when the module is fully analyzed.
+                return;
             }
+
+            AddMember("__bases__", PythonCollectionType.CreateList(DeclaringModule.Interpreter, _bases), true);
         }
 
         /// <summary>
@@ -263,15 +269,17 @@ internal static IReadOnlyList<IPythonType> CalculateMro(IPythonType type, HashSe
             }
         }
 
+        #region Reentrancy guards
         private bool Push(IPythonClassType cls) {
-            if (_processing.Value == null) {
-                _processing.Value = cls;
+            if (_processing == null) {
+                _processing = cls;
                 return true;
             }
             return false;
         }
+        private void Pop() => _processing = null;
+        #endregion
 
-        private void Pop() => _processing.Value = null;
         public bool Equals(IPythonClassType other)
             => Name == other?.Name && DeclaringModule.Equals(other?.DeclaringModule);
 
@@ -375,7 +383,7 @@ public IPythonType CreateSpecificType(IArgumentSet args) {
             // Prevent reentrancy when resolving generic class where
             // method may be returning instance of type of the same class.
             if (!Push(classType)) {
-                return _processing.Value;
+                return _processing;
             }
 
             try {
@@ -400,6 +408,7 @@ public IPythonType CreateSpecificType(IArgumentSet args) {
                         }
                     }
                 }
+
                 // Set specific class bases
                 classType.SetBases(bases.Concat(newBases));
                 // Transfer members from generic to specific type.

src/Analysis/Ast/Impl/Types/PythonFunctionType.cs

@@ -81,8 +81,12 @@ Location location
             FunctionDefinition = fd;
             DeclaringType = declaringType;
 
+            // For __init__ documentation may either come from the function node of the the declaring
+            // type. Note that if there is no documentation on the class node, the class will try and
+            // get documentation from its __init__ function, delegating down to this type. So we need
+            // to set documentation statically for __init__ here or we may end up/ with stack overflows.
             if (fd.Name == "__init__") {
-                _documentation = declaringType?.Documentation;
+                _documentation = declaringType?.Documentation ?? fd.Documentation;
             }
             ProcessDecorators(fd);
         }
@@ -139,7 +143,7 @@ internal void Specialize(string[] dependencies) {
 
         internal ImmutableArray<string> Dependencies { get; private set; } = ImmutableArray<string>.Empty;
 
-        internal void AddOverload(IPythonFunctionOverload overload) 
+        internal void AddOverload(IPythonFunctionOverload overload)
             => _overloads = _overloads.Count > 0 ? _overloads.Add(overload) : ImmutableArray<IPythonFunctionOverload>.Create(overload);
 
         internal IPythonFunctionType ToUnbound() => new PythonUnboundMethod(this);

src/Analysis/Ast/Test/ClassesTests.cs

@@ -553,5 +553,24 @@ def a(self):
             analysis.Should().HaveVariable("a").OfType(BuiltinTypeId.Int);
         }
 
+        [TestMethod, Priority(0)]
+        public async Task DocumentationWithCycleInBases() {
+            const string code = @"
+class A(C):
+    '''class A doc'''
+
+class B(A):
+    def __init__(self):
+        '''class B doc'''
+        pass
+
+class C(B):
+    '''class C doc'''
+";
+            var analysis = await GetAnalysisAsync(code);
+            analysis.Should().HaveClass("A").Which.Should().HaveDocumentation("class A doc");
+            analysis.Should().HaveClass("B").Which.Should().HaveDocumentation("class B doc");
+            analysis.Should().HaveClass("C").Which.Should().HaveDocumentation("class C doc");
+        }
     }
 }

src/Core/Impl/Extensions/IOExtensions.cs

@@ -114,25 +114,22 @@ public static string ReadTextWithRetry(this IFileSystem fs, string file) {
             for (var retries = 100; retries > 0; --retries) {
                 try {
                     return fs.ReadAllText(file);
-                } catch (UnauthorizedAccessException) {
-                    Thread.Sleep(10);
-                } catch (IOException) {
+                } catch (Exception ex) when (ex is UnauthorizedAccessException || ex is IOException || ex is ObjectDisposedException) {
                     Thread.Sleep(10);
                 }
             }
             return null;
         }
 
         public static void WriteTextWithRetry(this IFileSystem fs, string filePath, string text) {
-            try {
-                using (var stream = fs.OpenWithRetry(filePath, FileMode.Create, FileAccess.Write, FileShare.Read)) {
-                    if (stream != null) {
-                        var bytes = Encoding.UTF8.GetBytes(text);
-                        stream.Write(bytes, 0, bytes.Length);
-                        return;
-                    }
+            for (var retries = 100; retries > 0; --retries) {
+                try {
+                    fs.WriteAllText(filePath, text);
+                    return;
+                } catch (Exception ex) when (ex is IOException || ex is UnauthorizedAccessException) {
+                    Thread.Sleep(10);
                 }
-            } catch (IOException) { } catch (UnauthorizedAccessException) { }
+            }
 
             try {
                 fs.DeleteFile(filePath);