Implement Cleanup { } Function Block

src/Microsoft.PowerShell.Commands.Utility/commands/utility/ImplicitRemotingCommands.cs

@@ -2819,15 +2819,16 @@ private void GenerateHelperFunctions(TextWriter writer)
 
             $positionalArguments.AddRange($args)
 
-            $clientSideParameters = Get-PSImplicitRemotingClientSideParameters $PSBoundParameters ${8}
-
-            $scriptCmd = {{ & $script:InvokeCommand `
-                            @clientSideParameters `
-                            -HideComputerName `
-                            -Session (Get-PSImplicitRemotingSession -CommandName '{0}') `
-                            -Arg ('{0}', $PSBoundParameters, $positionalArguments) `
-                            -Script {{ param($name, $boundParams, $unboundParams) & $name @boundParams @unboundParams }} `
-                         }}
+            $clientSideParameters = Get-PSImplicitRemotingClientSideParameters $PSBoundParameters ${9}
+
+            $scriptCmd = {{
+                & $script:InvokeCommand `
+                    @clientSideParameters `
+                    -HideComputerName `
+                    -Session (Get-PSImplicitRemotingSession -CommandName '{0}') `
+                    -Arg ('{0}', $PSBoundParameters, $positionalArguments) `
+                    -Script {{ param($name, $boundParams, $unboundParams) & $name @boundParams @unboundParams }} `
+            }}
 
             $steppablePipeline = $scriptCmd.GetSteppablePipeline($myInvocation.CommandOrigin)
             $steppablePipeline.Begin($myInvocation.ExpectingInput, $ExecutionContext)
@@ -2840,6 +2841,8 @@ private void GenerateHelperFunctions(TextWriter writer)
 
     End {{ {7} }}
 
+    Cleanup {{ {8} }}
+
     # .ForwardHelpTargetName {1}
     # .ForwardHelpCategory {5}
     # .RemoteHelpRunspace PSSession
@@ -2865,7 +2868,8 @@ private static void GenerateCommandProxy(TextWriter writer, CommandMetadata comm
                 /* 5 */ commandMetadata.WrappedCommandType,
                 /* 6 */ ProxyCommand.GetProcess(commandMetadata),
                 /* 7 */ ProxyCommand.GetEnd(commandMetadata),
-                /* 8 */ commandMetadata.WrappedAnyCmdlet);
+                /* 8 */ ProxyCommand.GetCleanup(commandMetadata),
+                /* 9 */ commandMetadata.WrappedAnyCmdlet);
         }
 
         private static void GenerateCommandProxy(TextWriter writer, IEnumerable<CommandMetadata> listOfCommandMetadata)

src/System.Management.Automation/engine/CommandMetadata.cs

@@ -869,8 +869,11 @@ internal string GetProxyCommand(string helpComment, bool generateDynamicParamete
 
 end
 {{{5}}}
+
+cleanup
+{{{6}}}
 <#
-{6}
+{7}
 #>
 ",
                 GetDecl(),
@@ -879,6 +882,7 @@ internal string GetProxyCommand(string helpComment, bool generateDynamicParamete
                 GetBeginBlock(),
                 GetProcessBlock(),
                 GetEndBlock(),
+                GetCleanupBlock(),
                 CodeGeneration.EscapeBlockCommentContent(helpComment));
 
             return result;
@@ -1107,6 +1111,17 @@ internal string GetEndBlock()
 ";
         }
 
+        internal string GetCleanupBlock()
+        {
+            return @"
+    try {
+        $steppablePipeline.Dispose()
+    } catch {
+        throw
+    }
+";
+        }
+
         #endregion
 
         #region Helper methods for restricting commands needed by implicit and interactive remoting

src/System.Management.Automation/engine/CommandProcessorBase.cs

@@ -366,7 +366,10 @@ internal void RestorePreviousScope()
         {
             OnRestorePreviousScope();
 
-            Context.EngineSessionState = _previousCommandSessionState;
+            if (_previousCommandSessionState != null)
+            {
+                Context.EngineSessionState = _previousCommandSessionState;
+            }
 
             if (_previousScope != null)
             {
@@ -645,13 +648,7 @@ internal void DoComplete()
                 _context.ShellFunctionErrorOutputPipe = oldErrorOutputPipe;
                 _context.CurrentCommandProcessor = oldCurrentCommandProcessor;
 
-                // Destroy the local scope at this point if there is one...
-                if (_useLocalScope && CommandScope != null)
-                {
-                    CommandSessionState.RemoveScope(CommandScope);
-                }
-
-                // and the previous scope...
+                // Restore the previous scope...
                 if (_previousScope != null)
                 {
                     // Restore the scope but use the same session state instance we
@@ -668,6 +665,68 @@ internal void DoComplete()
             }
         }
 
+        protected virtual void HandleScopedAction<T>(Action<T> action, T state, string traceMessage)
+        {
+            Pipe oldErrorOutputPipe = Context.ShellFunctionErrorOutputPipe;
+            CommandProcessorBase oldCurrentCommandProcessor = Context.CurrentCommandProcessor;
+
+            try
+            {
+                // On V1 the output pipe was redirected to the command's output pipe only when it
+                // was already redirected. This is the original comment explaining this behaviour:
+                //
+                //      NTRAID#Windows Out of Band Releases-926183-2005-12-15
+                //      MonadTestHarness has a bad dependency on an artifact of the current implementation
+                //      The following code only redirects the output pipe if it's already redirected
+                //      to preserve the artifact. The test suites need to be fixed and then this
+                //      the check can be removed and the assignment always done.
+                //
+                // However, this makes the hosting APIs behave differently than commands executed
+                // from the command-line host (for example, see bugs Win7:415915 and Win7:108670).
+                // The RedirectShellErrorOutputPipe flag is used by the V2 hosting API to force the
+                // redirection.
+                if (this.RedirectShellErrorOutputPipe || Context.ShellFunctionErrorOutputPipe != null)
+                {
+                    Context.ShellFunctionErrorOutputPipe = CommandRuntime.ErrorOutputPipe;
+                }
+
+                Context.CurrentCommandProcessor = this;
+                using (CommandRuntime.AllowThisCommandToWrite(permittedToWriteToPipeline: true))
+                using (ParameterBinderBase.bindingTracer.TraceScope(traceMessage))
+                {
+                    SetCurrentScopeToExecutionScope();
+                    action(state);
+                }
+            }
+            catch (Exception e)
+            {
+                throw ManageInvocationException(e);
+            }
+            finally
+            {
+                Context.ShellFunctionErrorOutputPipe = oldErrorOutputPipe;
+                Context.CurrentCommandProcessor = oldCurrentCommandProcessor;
+
+                // Destroy the local scope at this point if there is one...
+                if (_useLocalScope && CommandScope != null)
+                {
+                    CommandSessionState.RemoveScope(CommandScope);
+                }
+
+                RestorePreviousScope();
+            }
+        }
+
+        /// <summary>
+        /// Virtual method to be overridden by a command processor type that handles script functions, such as
+        /// <see cref="DlrScriptCommandProcessor"/>.
+        /// The base implementation does nothing.
+        /// </summary>
+        protected virtual void CleanupScriptCommand()
+        {
+            // Do nothing -- subclasses may use.
+        }
+
         /// <summary>
         /// For diagnostic purposes.
         /// </summary>
@@ -940,21 +999,50 @@ public void Dispose()
         private void Dispose(bool disposing)
         {
             if (_disposed)
+            {
                 return;
+            }
 
-            if (disposing)
+            try
             {
-                // 2004/03/05-JonN Look into using metadata to check
-                // whether IDisposable is implemented, in order to avoid
-                // this expensive reflection cast.
-                IDisposable id = Command as IDisposable;
-                if (id != null)
+                if (disposing)
                 {
-                    id.Dispose();
+                    try
+                    {
+                        CommandRuntime.RemoveVariableListsInPipe();
+
+                        // This method handles script commands' `cleanup {}` blocks.
+                        // It is a no-op for compiled cmdlets.
+                        CleanupScriptCommand();
+                    }
+                    finally
+                    {
+                        if (Command is PSScriptCmdlet scriptCmdlet)
+                        {
+                            bool isStopping = ExceptionHandlingOps.SuspendStoppingPipeline(Context);
+                            try
+                            {
+                                HandleScopedAction(
+                                    (cmdlet) => cmdlet.Dispose(),
+                                    scriptCmdlet,
+                                    traceMessage: "CALLING Cleanup");
+                            }
+                            finally
+                            {
+                                ExceptionHandlingOps.RestoreStoppingPipeline(Context, isStopping);
+                            }
+                        }
+                        else if (Command is IDisposable disposableCommand)
+                        {
+                            disposableCommand.Dispose();
+                        }
+                    }
                 }
             }
-
-            _disposed = true;
+            finally
+            {
+                _disposed = true;
+            }
         }
 
         #endregion IDispose

src/System.Management.Automation/engine/EventManager.cs

@@ -1844,6 +1844,11 @@ public sealed class PSEngineEvent
         /// </summary>
         internal const string OnScriptBlockInvoke = "PowerShell.OnScriptBlockInvoke";
 
+        /// <summary>
+        /// Called during invocation of cleanup{} for script commands.
+        /// </summary>
+        internal const string OnScriptCommandCleanup = "PowerShell.OnScriptCommandCleanup";
+
         /// <summary>
         /// Called during scriptblock invocation.
         /// </summary>

src/System.Management.Automation/engine/ProxyCommand.cs

@@ -247,6 +247,30 @@ public static string GetEnd(CommandMetadata commandMetadata)
             return commandMetadata.GetEndBlock();
         }
 
+        /// <summary>
+        /// This method constructs a string representing the dispose block of the command
+        /// specified by <paramref name="commandMetadata"/>. The returned string only contains the
+        /// script, it is not enclosed in "cleanup { }".
+        /// </summary>
+        /// <param name="commandMetadata">
+        /// An instance of CommandMetadata representing a command.
+        /// </param>
+        /// <returns>
+        /// A string representing the end block of the command.
+        /// </returns>
+        /// <exception cref="ArgumentNullException">
+        /// If <paramref name="commandMetadata"/> is null.
+        /// </exception>
+        public static string GetCleanup(CommandMetadata commandMetadata)
+        {
+            if (commandMetadata == null)
+            {
+                throw PSTraceSource.NewArgumentNullException(nameof(commandMetadata));
+            }
+
+            return commandMetadata.GetCleanupBlock();
+        }
+
         private static T GetProperty<T>(PSObject obj, string property) where T : class
         {
             T result = null;