diff --git a/src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHost.cs b/src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHost.cs index ae01316a89..060fa942b5 100644 --- a/src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHost.cs +++ b/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"); diff --git a/src/System.Management.Automation/engine/hostifaces/Connection.cs b/src/System.Management.Automation/engine/hostifaces/Connection.cs index 398a728831..c948a2f205 100644 --- a/src/System.Management.Automation/engine/hostifaces/Connection.cs +++ b/src/System.Management.Automation/engine/hostifaces/Connection.cs @@ -520,6 +520,29 @@ namespace System.Management.Automation.Runspaces } } + /// + /// 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. + /// + internal static Runspace PrimaryRunspace + { + get + { + return s_primaryRunspace; + } + + set + { + var result = Interlocked.CompareExchange(ref s_primaryRunspace, value, null); + if (result != null) + { + throw new PSInvalidOperationException(RunspaceStrings.PrimaryRunspaceAlreadySet); + } + } + } + private static Runspace s_primaryRunspace; + /// /// Returns true if Runspace.DefaultRunspace can be used to /// create an instance of the PowerShell class with diff --git a/src/System.Management.Automation/engine/hostifaces/LocalConnection.cs b/src/System.Management.Automation/engine/hostifaces/LocalConnection.cs index 21804bdab5..e9c66ea910 100644 --- a/src/System.Management.Automation/engine/hostifaces/LocalConnection.cs +++ b/src/System.Management.Automation/engine/hostifaces/LocalConnection.cs @@ -811,40 +811,34 @@ namespace System.Management.Automation.Runspaces /// 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; diff --git a/src/System.Management.Automation/resources/RunspaceStrings.resx b/src/System.Management.Automation/resources/RunspaceStrings.resx index 81a34adc19..2f2cec0dcf 100644 --- a/src/System.Management.Automation/resources/RunspaceStrings.resx +++ b/src/System.Management.Automation/resources/RunspaceStrings.resx @@ -243,4 +243,7 @@ DefaultRunspace must be a LocalRunspace + + The static PrimaryRunspace property can only be set once, and has already been set. +