Re-add check for Invoke-Command job dispose. (#11388)

src/System.Management.Automation/engine/remoting/commands/InvokeCommandCommand.cs

@@ -1204,8 +1204,15 @@ namespace Microsoft.PowerShell.Commands
                         // be connected to later.
                         WriteJobResults(false);
 
-                        // finally dispose the job.
-                        _job.Dispose();
+                        // Dispose job object if it is not returned to the user.
+                        // The _asjob field can change dynamically and needs to be checked before the job 
+                        // object is disposed. For example, if remote sessions are disconnected abruptly
+                        // via WinRM, a disconnected job object is created to facilitate a reconnect.
+                        // If the job object is disposed here, then a session reconnect cannot happen.
+                        if (!_asjob)
+                        {
+                            _job.Dispose();
+                        }
 
                         // We no longer need to call ClearInvokeCommandOnRunspaces() here because
                         // this command might finish before the foreach block finishes. previously,
@@ -1251,8 +1258,15 @@ namespace Microsoft.PowerShell.Commands
                             // be connected to later.
                             WriteJobResults(false);
 
-                            // finally dispose the job.
-                            _job.Dispose();
+                            // Dispose job object if it is not returned to the user.
+                            // The _asjob field can change dynamically and needs to be checked before the job 
+                            // object is disposed. For example, if remote sessions are disconnected abruptly
+                            // via WinRM, a disconnected job object is created to facilitate a reconnect.
+                            // If the job object is disposed here, then a session reconnect cannot happen.
+                            if (!_asjob)
+                            {
+                                _job.Dispose();
+                            }
                         }
                     }
                 }

test/powershell/engine/Remoting/RemoteSession.Disconnect.Tests.ps1

@@ -0,0 +1,82 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+Import-Module HelpersCommon
+
+Describe "WinRM based remoting session abrupt disconnect" -Tags 'Feature','RequireAdminOnWindows' {
+
+    BeforeAll {
+        $originalDefaultParameterValues = $PSDefaultParameterValues.Clone()
+        if (! $IsWindows)
+        {
+            $PSDefaultParameterValues["it:skip"] = $true
+            return
+        }
+
+        $disconnectScript = @'
+            param (
+                [int] $RunspaceId
+            )
+
+            $rs = Get-Runspace -Id $RunspaceId
+
+            if (! $rs.RunspaceIsRemote)
+            {
+                throw "Runspace $RunspaceId is not a remote runspace."
+            }
+
+            # Wait up to one minute for Runspace to begin running script
+            $count = 0
+            while (($rs.RunspaceAvailability -ne "busy") -and (++$count -le 60)) {
+                Start-Sleep 1
+            }
+
+            if ($rs.RunspaceAvailability -ne "busy")
+            {
+                throw "Runspace $RunspaceId is not running any script after one minute."
+            }
+
+            # Disconnect running runspace
+            $rs.Disconnect()
+'@
+
+        $endPointName = "PowerShell.$(${PSVersionTable}.GitCommitId)"
+        $endPoint = (Get-PSSessionConfiguration -Name $endPointName -ErrorAction SilentlyContinue).Name
+        if ($endPoint -eq $null)
+        {
+            Enable-PSRemoting -SkipNetworkProfileCheck
+            $endPoint = (Get-PSSessionConfiguration -Name $endPointName).Name
+        }
+        $session = New-RemoteSession -ConfigurationName $endPoint
+
+        $ps = [powershell]::Create("NewRunspace")
+        $ps.AddScript($disconnectScript).AddParameter("RunspaceId", $session.Runspace.Id)
+    }
+
+    AfterAll {
+        $global:PSDefaultParameterValues = $originalDefaultParameterValues
+
+        if ($ps -ne $null) { $ps.Dispose() }
+        if ($session -ne $null) { Remove-PSSession -Session $session }
+        if ($script:job -ne $null) { Remove-Job -Job $script:job -Force }
+    }
+
+    It "Verifies that an abruptly disconnected Invoke-Command session produces a valid disconnected job needed for reconnect" {
+
+        # Start disconnect script running.
+        $ps.BeginInvoke()
+
+        # Run script synchronously on remote session, and let disconnect script disconnect the remote session.
+        $null = Invoke-Command -Session $session -ScriptBlock {
+            1..60 | ForEach-Object { Start-Sleep 1; "Output $_" }
+        } -ErrorAction SilentlyContinue
+
+        # Session should be disconnected.
+        $session.State | Should -BeExactly 'Disconnected'
+
+        # A disconnected job should have been created for reconnect.
+        $script:job = Get-Job | Where-Object { $_.ChildJobs[0].Runspace.Id -eq $session.Runspace.Id }
+        $script:job | Should -Not -BeNullOrEmpty
+        $script:job.State | Should -BeExactly 'Disconnected'
+    }
+}