Fix for PowerShell hang on exit (#5356)

CoreCLR doesn't call finalizer on process exit. PowerShell relies on the CLR finalizer to clean up state on exit. In this case, a Runspace pool was not closed or disposed and any pipeline worker threads created to run concurrent scripts won't end, causing the hang. The same thing can happen if any individual Runspace is created to run a concurrent script and is not closed.

We cannot use the `AppDomain.DomainUnload` event because it's not supported by the default load context. The `AppDomain.ProcessExit` event is also not helpful since it is only called during application exit which means threads already have to be cleaned up.

The fix introduces a static property called `PrimaryRunspace` to `Runspace`. When the PrimaryRunspace is closing it means that the PowerShell session is ending and on exit clean should be performed. The static property `PrimaryRunspace` can only be set once per process.

src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHost.cs

@@ -1658,6 +1658,7 @@ namespace Microsoft.PowerShell
 #endif
             runspace.ThreadOptions = PSThreadOptions.ReuseThread;
             runspace.EngineActivityId = EtwActivity.GetActivityId();
+            Runspace.PrimaryRunspace = runspace;
 
             s_runspaceInitTracer.WriteLine("Calling Runspace.Open");
 

src/System.Management.Automation/engine/hostifaces/Connection.cs

@@ -520,6 +520,29 @@ namespace System.Management.Automation.Runspaces
             }
         }
 
+        /// <summary>
+        /// A PrimaryRunspace is a runspace that persists for the entire lifetime of the PowerShell session. It is only
+        /// closed or disposed when the session is ending.  So when the PrimaryRunspace is closing it will trigger on-exit
+        /// cleanup that includes closing any other local runspaces left open, and will allow the process to exit.
+        /// </summary>
+        internal static Runspace PrimaryRunspace
+        {
+            get
+            {
+                return s_primaryRunspace;
+            }
+
+            set
+            {
+                var result = Interlocked.CompareExchange<Runspace>(ref s_primaryRunspace, value, null);
+                if (result != null)
+                {
+                    throw new PSInvalidOperationException(RunspaceStrings.PrimaryRunspaceAlreadySet);
+                }
+            }
+        }
+        private static Runspace s_primaryRunspace;
+
         /// <summary>
         /// Returns true if Runspace.DefaultRunspace can be used to
         /// create an instance of the PowerShell class with

src/System.Management.Automation/engine/hostifaces/LocalConnection.cs

@@ -811,40 +811,34 @@ namespace System.Management.Automation.Runspaces
         /// </remarks>
         private void DoCloseHelper()
         {
-            // Stop any transcription if we're the last runspace to exit
-            ExecutionContext executionContext = this.GetExecutionContext;
-            if (executionContext != null)
+            var isPrimaryRunspace = (Runspace.PrimaryRunspace == this);
+            var haveOpenRunspaces = false;
+            foreach (Runspace runspace in RunspaceList)
             {
-                Runspace hostRunspace = null;
-                try
+                if (runspace.RunspaceStateInfo.State == RunspaceState.Opened)
                 {
-                    hostRunspace = executionContext.EngineHostInterface.Runspace;
+                    haveOpenRunspaces = true;
+                    break;
                 }
-                catch (PSNotImplementedException)
+            }
+
+            // When closing the primary runspace, ensure all other local runspaces are closed.
+            var closeAllOpenRunspaces = isPrimaryRunspace && haveOpenRunspaces;
+
+            // Stop all transcriptions and unitialize AMSI if we're the last runspace to exit or we are exiting the primary runspace.
+            if (!haveOpenRunspaces)
+            {
+                ExecutionContext executionContext = this.GetExecutionContext;
+                if (executionContext != null)
                 {
-                    // EngineHostInterface.Runspace throws PSNotImplementedException if there
-                    // is no interactive host.
+                    PSHostUserInterface hostUI = executionContext.EngineHostInterface.UI;
+                    if (hostUI != null)
+                    {
+                        hostUI.StopAllTranscribing();
+                    }
                 }
 
-                if ((hostRunspace == null) || (this == hostRunspace))
-                {
-                    // We should close transcripting only if we are closing the last opened runspace.
-                    foreach (Runspace runspace in RunspaceList)
-                    {
-                        // At this stage, the last opened runspace should be at closing state.
-                        if (runspace.RunspaceStateInfo.State == RunspaceState.Opened)
-                        {
-                            return;
-                        }
-                    }
-
-                    PSHostUserInterface host = executionContext.EngineHostInterface.UI;
-                    if (host != null)
-                    {
-                        host.StopAllTranscribing();
-                    }
-                    AmsiUtils.Uninitialize();
-                }
+                AmsiUtils.Uninitialize();
             }
 
             // Generate the shutdown event
@@ -852,7 +846,6 @@ namespace System.Management.Automation.Runspaces
                 Events.GenerateEvent(PSEngineEvent.Exiting, null, new object[] { }, null, true, false);
 
             //Stop all running pipelines
-
             //Note:Do not perform the Cancel in lock. Reason is
             //Pipeline executes in separate thread, say threadP.
             //When pipeline is canceled/failed/completed in
@@ -897,8 +890,18 @@ namespace System.Management.Automation.Runspaces
             //Raise Event
             RaiseRunspaceStateEvents();
 
-            // Report telemetry if we have no more open runspaces.
+            if (closeAllOpenRunspaces)
+            {
+                foreach (Runspace runspace in RunspaceList)
+                {
+                    if (runspace.RunspaceStateInfo.State == RunspaceState.Opened)
+                    {
+                        runspace.Dispose();
+                    }
+                }
+            }
 
+            // Report telemetry if we have no more open runspaces.
 #if LEGACYTELEMETRY
             bool allRunspacesClosed = true;
             bool hostProvidesExitTelemetry = false;

src/System.Management.Automation/resources/RunspaceStrings.resx

@@ -243,4 +243,7 @@
   <data name="RunspaceNotLocal" xml:space="preserve">
     <value>DefaultRunspace must be a LocalRunspace</value>
   </data>
+  <data name="PrimaryRunspaceAlreadySet" xml:space="preserve">
+    <value>The static PrimaryRunspace property can only be set once, and has already been set.</value>
+  </data>
 </root>