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.
+