From e3497bffca382c65b720c8193eb6bd3bc85e465f Mon Sep 17 00:00:00 2001
From: Jan Vorlicek <janvorli@microsoft.com>
Date: Fri, 29 Mar 2019 12:55:00 +0100
Subject: [PATCH] Fix accidental non-collectible context unloading

The OnProcessExit was iterating over the s_contextsToUnload and calling
Unload on each of the contexts. However the s_contextsToUnload is a
misnomer, it contains all the contexts, both collectible and
non-collectible ones. So we were calling Unload on non-collectible
contexts too there. That resulted in InvalidOperationException being
thrown from the OnProcessExit.

This change fixes that by calling unload only on ALCs whose
IsCollectible returns true. It also renames the s_contextsToUnload to
s_allContexts.

Close #22902
---
 .../System/Runtime/Loader/AssemblyLoadContext.cs   | 24 +++++++++++-----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/src/System.Private.CoreLib/shared/System/Runtime/Loader/AssemblyLoadContext.cs b/src/System.Private.CoreLib/shared/System/Runtime/Loader/AssemblyLoadContext.cs
index 835c5c6..d5d9b03 100644
--- a/src/System.Private.CoreLib/shared/System/Runtime/Loader/AssemblyLoadContext.cs
+++ b/src/System.Private.CoreLib/shared/System/Runtime/Loader/AssemblyLoadContext.cs
@@ -28,14 +28,14 @@ namespace System.Runtime.Loader
             Unloading
         }
 
-        private static readonly Dictionary<long, WeakReference<AssemblyLoadContext>> s_contextsToUnload = new Dictionary<long, WeakReference<AssemblyLoadContext>>();
+        private static readonly Dictionary<long, WeakReference<AssemblyLoadContext>> s_allContexts = new Dictionary<long, WeakReference<AssemblyLoadContext>>();
         private static long s_nextId;
         private static bool s_isProcessExiting;
 
         // Indicates the state of this ALC (Alive or in Unloading state)
         private InternalState _state;
 
-        // Id used by s_contextsToUnload
+        // Id used by s_allContexts
         private readonly long _id;
 
         // synchronization primitive to protect against usage of this instance while unloading
@@ -75,7 +75,7 @@ namespace System.Runtime.Loader
             }
 
             // Add this instance to the list of alive ALC
-            lock (s_contextsToUnload)
+            lock (s_allContexts)
             {
                 if (s_isProcessExiting)
                 {
@@ -88,7 +88,7 @@ namespace System.Runtime.Loader
                 _nativeAssemblyLoadContext = InitializeAssemblyLoadContext(thisHandlePtr, representsTPALoadContext, isCollectible);
 
                 _id = s_nextId++;
-                s_contextsToUnload.Add(_id, new WeakReference<AssemblyLoadContext>(this, true));
+                s_allContexts.Add(_id, new WeakReference<AssemblyLoadContext>(this, true));
             }
         }
 
@@ -132,9 +132,9 @@ namespace System.Runtime.Loader
 
             if (!s_isProcessExiting)
             {
-                lock (s_contextsToUnload)
+                lock (s_allContexts)
                 {
-                    s_contextsToUnload.Remove(_id);
+                    s_allContexts.Remove(_id);
                 }
             }
         }
@@ -201,10 +201,10 @@ namespace System.Runtime.Loader
                 AssemblyLoadContext d = AssemblyLoadContext.Default; // Ensure default is initialized
 
                 List<WeakReference<AssemblyLoadContext>> alcList = null;
-                lock (s_contextsToUnload)
+                lock (s_allContexts)
                 {
                     // To make this thread safe we need a quick snapshot while locked
-                    alcList = new List<WeakReference<AssemblyLoadContext>>(s_contextsToUnload.Values);
+                    alcList = new List<WeakReference<AssemblyLoadContext>>(s_allContexts.Values);
                 }
 
                 foreach (WeakReference<AssemblyLoadContext> weakAlc in alcList)
@@ -385,18 +385,18 @@ namespace System.Runtime.Loader
 
         internal static void OnProcessExit()
         {
-            lock (s_contextsToUnload)
+            lock (s_allContexts)
             {
                 s_isProcessExiting = true;
-                foreach (var alcAlive in s_contextsToUnload)
+                foreach (var alcAlive in s_allContexts)
                 {
-                    if (alcAlive.Value.TryGetTarget(out AssemblyLoadContext alc))
+                    if (alcAlive.Value.TryGetTarget(out AssemblyLoadContext alc) && alc.IsCollectible)
                     {
                         // Should we use a try/catch?
                         alc.InitiateUnload();
                     }
                 }
-                s_contextsToUnload.Clear();
+                s_allContexts.Clear();
             }
         }        
 
-- 
2.7.4