Reflect PR #38 Partial: Assembly Manager Improvements

pythonnet · C-SELLERS · Feb 3, 2021 · Feb 3, 2021 · Feb 4, 2021 · Feb 4, 2021
commit c362816f478a10e3e32003fcc4f37749abe6ca18
diff --git a/src/runtime/assemblymanager.cs b/src/runtime/assemblymanager.cs
@@ -7,6 +7,7 @@
 using System.Linq;
 using System.Reflection;
 using System.Threading;
+using System.Threading.Tasks;
 
 namespace Python.Runtime
 {
@@ -25,18 +26,17 @@ internal class AssemblyManager
         //    than it can end up referring to assemblies that are already unloaded (default behavior after unload appDomain -
         //     unless LoaderOptimization.MultiDomain is used);
         //    So for multidomain support it is better to have the dict. recreated for each app-domain initialization
-        private static ConcurrentDictionary<string, ConcurrentDictionary<Assembly, string>> namespaces =
-            new ConcurrentDictionary<string, ConcurrentDictionary<Assembly, string>>();
-        //private static Dictionary<string, Dictionary<string, string>> generics;
-        private static AssemblyLoadEventHandler lhandler;
-        private static ResolveEventHandler rhandler;
+        private static ConcurrentDictionary<string, ConcurrentDictionary<Assembly, byte>> namespaces =
+            new ConcurrentDictionary<string, ConcurrentDictionary<Assembly, byte>>();
+        private static ConcurrentDictionary<string, Assembly> assembliesNamesCache =
+            new ConcurrentDictionary<string, Assembly>();
+        private static ConcurrentQueue<Assembly> assemblies = new ConcurrentQueue<Assembly>();
+        private static int pendingAssemblies;
 
         // updated only under GIL?
         private static Dictionary<string, int> probed = new Dictionary<string, int>(32);
-
-        // modified from event handlers below, potentially triggered from different .NET threads
-        private static ConcurrentQueue<Assembly> assemblies;
-        internal static List<string> pypath;
+        private static List<string> pypath = new List<string>(16);
+        private static Dictionary<string, HashSet<string>> filesInPath = new Dictionary<string, HashSet<string>>();
 
         private AssemblyManager()
         {
@@ -49,30 +49,34 @@ private AssemblyManager()
         /// </summary>
         internal static void Initialize()
         {
-            assemblies = new ConcurrentQueue<Assembly>();
-            pypath = new List<string>(16);
-
             AppDomain domain = AppDomain.CurrentDomain;
 
-            lhandler = new AssemblyLoadEventHandler(AssemblyLoadHandler);
-            domain.AssemblyLoad += lhandler;
-
-            rhandler = new ResolveEventHandler(ResolveHandler);
-            domain.AssemblyResolve += rhandler;
+            domain.AssemblyLoad += AssemblyLoadHandler;
+            domain.AssemblyResolve += ResolveHandler;
 
-            Assembly[] items = domain.GetAssemblies();
-            foreach (Assembly a in items)
+            foreach (var assembly in domain.GetAssemblies())
             {
                 try
                 {
-                    ScanAssembly(a);
-                    assemblies.Enqueue(a);
+                    LaunchAssemblyLoader(assembly);
                 }
                 catch (Exception ex)
                 {
-                    Debug.WriteLine("Error scanning assembly {0}. {1}", a, ex);
+                    Debug.WriteLine("Error scanning assembly {0}. {1}", assembly, ex);
                 }
             }
+
+            var safeCount = 0;
+            // lets wait until all assemblies are loaded
+            do
+            {
+                if (safeCount++ > 200)
+                {
+                    throw new TimeoutException("Timeout while waiting for assemblies to load");
+                }
+
+                Thread.Sleep(50);
+            } while (pendingAssemblies > 0);
         }
 
 
@@ -82,8 +86,8 @@ internal static void Initialize()
         internal static void Shutdown()
         {
             AppDomain domain = AppDomain.CurrentDomain;
-            domain.AssemblyLoad -= lhandler;
-            domain.AssemblyResolve -= rhandler;
+            domain.AssemblyLoad -= AssemblyLoadHandler;
+            domain.AssemblyResolve -= ResolveHandler;
         }
 
 
@@ -97,8 +101,35 @@ internal static void Shutdown()
         private static void AssemblyLoadHandler(object ob, AssemblyLoadEventArgs args)
         {
             Assembly assembly = args.LoadedAssembly;
-            assemblies.Enqueue(assembly);
-            ScanAssembly(assembly);
+            LaunchAssemblyLoader(assembly);
+        }
+
+        /// <summary>
+        /// Launches a new task that will load the provided assembly
+        /// </summary>
+        private static void LaunchAssemblyLoader(Assembly assembly)
+        {
+            if (assembly != null)
+            {
+                Interlocked.Increment(ref pendingAssemblies);
+                Task.Factory.StartNew(() =>
+                {
+                    try
+                    {
+                        if (assembliesNamesCache.TryAdd(assembly.GetName().Name, assembly))
+                        {
+                            assemblies.Enqueue(assembly);
+                            ScanAssembly(assembly);
+                        }
+                    }
+                    catch
+                    {
+                        // pass
+                    }
+
+                    Interlocked.Decrement(ref pendingAssemblies);
+                });
+            }
         }
 
 
@@ -112,12 +143,12 @@ private static void AssemblyLoadHandler(object ob, AssemblyLoadEventArgs args)
         private static Assembly ResolveHandler(object ob, ResolveEventArgs args)
         {
             string name = args.Name.ToLower();
-            foreach (Assembly a in assemblies)
+            foreach (var assembly in assemblies)
             {
-                string full = a.FullName.ToLower();
+                var full = assembly.FullName.ToLower();
                 if (full.StartsWith(name))
                 {
-                    return a;
+                    return assembly;
                 }
             }
             return LoadAssemblyPath(args.Name);
@@ -139,19 +170,61 @@ internal static void UpdatePath()
         {
             BorrowedReference list = Runtime.PySys_GetObject("path");
             var count = Runtime.PyList_Size(list);
+            var sep = Path.DirectorySeparatorChar;
+
             if (count != pypath.Count)
             {
                 pypath.Clear();
                 probed.Clear();
+                // add first the current path
+                pypath.Add("");
                 for (var i = 0; i < count; i++)
                 {
                     BorrowedReference item = Runtime.PyList_GetItem(list, i);
                     string path = Runtime.GetManagedString(item);
                     if (path != null)
                     {
-                        pypath.Add(path);
+                        pypath.Add(path == string.Empty ? path : path + sep);
                     }
                 }
+
+                // for performance we will search for all files in each directory in the path once
+                Parallel.ForEach(pypath.Where(s =>
+                {
+                    try
+                    {
+                        lock (filesInPath)
+                        {
+                            // only search in directory if it exists and we haven't already analyzed it
+                            return Directory.Exists(s) && !filesInPath.ContainsKey(s);
+                        }
+                    }
+                    catch
+                    {
+                        // just in case, file operations can throw
+                    }
+                    return false;
+                }), path =>
+                {
+                    var container = new HashSet<string>();
+                    try
+                    {
+                        foreach (var file in Directory.EnumerateFiles(path)
+                            .Where(file => file.EndsWith(".dll") || file.EndsWith(".exe")))
+                        {
+                            container.Add(Path.GetFileName(file));
+                        }
+                    }
+                    catch
+                    {
+                        // just in case, file operations can throw
+                    }
+
+                    lock (filesInPath)
+                    {
+                        filesInPath[path] = container;
+                    }
+                });
             }
         }
 
@@ -163,30 +236,18 @@ internal static void UpdatePath()
         /// </summary>
         public static string FindAssembly(string name)
         {
-            char sep = Path.DirectorySeparatorChar;
-
-            foreach (string head in pypath)
+            foreach (var kvp in filesInPath)
             {
-                string path;
-                if (head == null || head.Length == 0)
-                {
-                    path = name;
-                }
-                else
-                {
-                    path = head + sep + name;
-                }
-
-                string temp = path + ".dll";
-                if (File.Exists(temp))
+                var dll = $"{name}.dll";
+                if (kvp.Value.Contains(dll))
                 {
-                    return temp;
+                    return kvp.Key + dll;
                 }
 
-                temp = path + ".exe";
-                if (File.Exists(temp))
+                var executable = $"{name}.exe";
+                if (kvp.Value.Contains(executable))
                 {
-                    return temp;
+                    return kvp.Key + executable;
                 }
             }
             return null;
@@ -242,14 +303,8 @@ public static Assembly LoadAssemblyFullPath(string name)
         /// </summary>
         public static Assembly FindLoadedAssembly(string name)
         {
-            foreach (Assembly a in assemblies)
-            {
-                if (a.GetName().Name == name)
-                {
-                    return a;
-                }
-            }
-            return null;
+            Assembly result;
+            return assembliesNamesCache.TryGetValue(name, out result) ? result : null;
         }
 
         /// <summary>
@@ -264,16 +319,10 @@ internal static void ScanAssembly(Assembly assembly)
             {
                 return;
             }
-
-            // skip this assembly, it causes 'GetTypes' call to hang
-            if (assembly.FullName.StartsWith("System.Windows.Forms"))
-            {
-                return;
-            }
-
             // A couple of things we want to do here: first, we want to
             // gather a list of all of the namespaces contributed to by
             // the assembly.
+
             foreach (Type t in GetTypes(assembly))
             {
                 string ns = t.Namespace ?? "";
@@ -284,13 +333,13 @@ internal static void ScanAssembly(Assembly assembly)
                     for (var n = 0; n < names.Length; n++)
                     {
                         s = n == 0 ? names[0] : s + "." + names[n];
-                        namespaces.TryAdd(s, new ConcurrentDictionary<Assembly, string>());
+                        namespaces.TryAdd(s, new ConcurrentDictionary<Assembly, byte>());
                     }
                 }
 
                 if (ns != null)
                 {
-                    namespaces[ns].TryAdd(assembly, string.Empty);
+                    namespaces[ns].TryAdd(assembly, 1);
                 }
 
                 if (ns != null && t.IsGenericTypeDefinition)