Refactored unloading into single common function

dotnet · Jun 21, 2018 · 0683d09 · 0683d09
1 parent ba8db83
commit 0683d09
Showing 1 changed file with 35 additions and 85 deletions.
diff --git a/src/System.Private.CoreLib/src/System/Runtime/Loader/AssemblyLoadContext.cs b/src/System.Private.CoreLib/src/System/Runtime/Loader/AssemblyLoadContext.cs
@@ -29,7 +29,7 @@ public abstract class AssemblyLoadContext
         // synchronization primitive to protect against usage of this instance while unloading
         private readonly object unloadLock = new object();
 
-        // Indicates the state of this ALC (Alive or in Unloading/Unloaded state)
+        // Indicates the state of this ALC (Alive or in Unloading state)
         private InternalState state;
 
         [DllImport(JitHelpers.QCall, CharSet = CharSet.Unicode)]
@@ -95,27 +95,45 @@ internal AssemblyLoadContext(bool fRepresentsTPALoadContext, bool isCollectible)
 
         ~AssemblyLoadContext()
         {
-            // Only valid for a Collectible ALC
-            // We perform an implicit Unload if no explicit Unload has been done
-            if (IsCollectible)
+            // Only valid for a Collectible ALC. Non-collectible ALCs have the finalizer suppressed.
+            // We get here only in case the explicit Unload was not initiated.
+            Debug.Assert(state != InternalState.Unloading);
+            InitiateUnload();
+        }
+
+        private void InitiateUnload()
+        {
+            Debug.Assert(IsCollectible);
+
+            if (!_isProcessExiting)
             {
                 var unloading = Unloading;
                 Unloading = null;
-                // TODO: should we enclose this with a try catch?
                 unloading?.Invoke(this);
+            }
 
-                // When in Unloading state, we are not supposed to be called on the finalizer
-                // as the native side is holding a strong reference after calling Unload
-                lock (unloadLock)
+            // When in Unloading state, we are not supposed to be called on the finalizer
+            // as the native side is holding a strong reference after calling Unload
+            lock (unloadLock)
+            {
+                if (!_isProcessExiting)
                 {
-                    Debug.Assert(state != InternalState.Unloading);
-                    if (state == InternalState.Alive)
-                    {
-                        UnloadCollectible();
-                    }
+                    Debug.Assert(state == InternalState.Alive);
+
+                    // The underlying code will transform the original weak handle 
+                    // created by InitializeLoadContext to a strong handle
+                    PrepareForAssemblyLoadContextRelease(m_pNativeAssemblyLoadContext);
                 }
 
-                OnUnloading();
+                state = InternalState.Unloading;
+            }
+
+            if (!_isProcessExiting)
+            {
+                lock (ContextsToUnload)
+                {
+                    ContextsToUnload.Remove(id);
+                }
             }
         }
 
@@ -244,34 +262,7 @@ public void Unload()
             }
 
             GC.SuppressFinalize(this);
-
-            var unloading = Unloading;
-            Unloading = null;
-
-            unloading?.Invoke(this);
-
-            lock (unloadLock)
-            {
-                UnloadCollectible();
-            }
-
-            OnUnloading();
-        }
-
-        private void UnloadCollectible()
-        {
-            Debug.Assert(IsCollectible);
-            if (state == InternalState.Alive)
-            {
-                // The underlying code will transform the original weak handle 
-                // created by InitializeLoadContext to a strong handle
-                PrepareForAssemblyLoadContextRelease(m_pNativeAssemblyLoadContext);
-            }
-            else
-            {
-                throw new InvalidOperationException(SR.GetResourceString("AssemblyLoadContext_Unload_AlreadyUnloaded"));
-            }
-            state = InternalState.Unloading;
+            InitiateUnload();
         }
 
         private void VerifyIsAlive()
@@ -383,34 +374,6 @@ private Assembly ResolveUsingEvent(AssemblyName assemblyName)
             return assembly;
         }
 
-        /// <summary>
-        /// This method is called back by the native code after Unloading has been initiated
-        /// This method is called indirectly by the finalizer of the LoaderAllocator
-        /// </summary>
-        private void OnUnloading()
-        {
-            lock (unloadLock)
-            {
-                if (state == InternalState.Unloading)
-                {
-                    state = InternalState.Unloaded;
-                    lock (ContextsToUnload)
-                    {
-                        if (!_isProcessExiting)
-                        {
-                            ContextsToUnload.Remove(id);
-                        }
-                    }
-                }
-                else
-                {
-                    // Otherwise we didn't have time to be called here and we will be 
-                    // called via the finalizer
-                    return;
-                }
-            }
-        }
-
         public Assembly LoadFromAssemblyName(AssemblyName assemblyName)
         {
             // Attempt to load the assembly, using the same ordering as static load, in the current load context.
@@ -543,15 +506,7 @@ public void StartProfileOptimization(string profile)
 
         private void OnAppContextUnloading()
         {
-            lock (unloadLock)
-            {
-                if (state == InternalState.Alive)
-                {
-                    state = InternalState.Unloading;
-                }
-            }
-
-            OnUnloading();
+            InitiateUnload();
         }
 
         private static void OnAppContextUnloading(object sender, EventArgs e)
@@ -626,12 +581,7 @@ private enum InternalState
             /// The unload process has started, the Unloading event will be called
             /// once the underlying LoaderAllocator has been finalized
             /// </summary>
-            Unloading,
-
-            /// <summary>
-            /// The event Unloading has been called.
-            /// </summary>
-            Unloaded
+            Unloading
         }
     }