Fix powershell class to use the current EngineSessionState for execution. (#2837)

When a PS class is defined in a module and the module gets reloaded, the class would still use the SessionState from the old module for execution, and thus it doesn't reflect changes to the module state during the reload. This fix is to make sure we always use the current EngineSessionState for PS class execution.

src/System.Management.Automation/engine/runtime/Operations/ClassOps.cs

@@ -43,9 +43,26 @@ namespace System.Management.Automation.Internal
 
         internal void RegisterRunspace()
         {
-            // it's not get, but really 'Add' value.
-            // ConditionalWeakTable.Add throw exception, when you are trying to add a value with the same key.
-            _stateMap.GetValue(Runspace.DefaultRunspace, runspace => runspace.ExecutionContext.EngineSessionState);
+            SessionStateInternal ssInMap = null;
+            Runspace rsToUse = Runspace.DefaultRunspace;
+            SessionStateInternal ssToUse = rsToUse.ExecutionContext.EngineSessionState;
+            
+            // Different threads will operate on different key/value pairs (default-runspace/session-state pairs),
+            // and a ConditionalWeakTable itself is thread safe, so there won't be race condition here.
+            if (!_stateMap.TryGetValue(rsToUse, out ssInMap))
+            {
+                // If the key doesn't exist yet, add it
+                _stateMap.Add(rsToUse, ssToUse);
+            }
+            else if (!ssInMap.Equals(ssToUse))
+            {
+                // If the key exists but the corresponding value is not what we should use, then remove the key/value pair and add the new pair.
+                // This could happen when a powershell class is defined in a module and the module gets reloaded. In such case, the same TypeDefinitionAst
+                // instance will get reused, but should be associated with the SessionState from the new module, instead of the one from the old module.
+                _stateMap.Remove(rsToUse);
+                _stateMap.Add(rsToUse, ssToUse);
+            }
+            // If the key exists and the corresponding value is the one we should use, then do nothing.
         }
 
         /// <summary>

test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1

@@ -71,3 +71,43 @@ Import-Module Random
     }
 
 }
+
+Describe 'Module reloading with Class definition' -Tags "CI" {
+
+    BeforeAll {
+        Set-Content -Path TestDrive:\TestModule.psm1 -Value @'
+$passedArgs = $args
+class Root { $passedIn = $passedArgs }
+function Get-PassedArgsRoot { [Root]::new().passedIn }
+function Get-PassedArgsNoRoot { $passedArgs }
+'@
+        $Arg_Hello = 'Hello'
+        $Arg_World = 'World'
+    }
+
+    AfterEach {
+        Remove-Module TestModule -Force -ErrorAction SilentlyContinue
+    }
+
+    It "Class execution reflects changes in module reloading with '-Force'" {
+        Import-Module TestDrive:\TestModule.psm1 -ArgumentList $Arg_Hello
+        Get-PassedArgsRoot | Should Be $Arg_Hello
+        Get-PassedArgsNoRoot | Should Be $Arg_Hello
+
+        Import-Module TestDrive:\TestModule.psm1 -ArgumentList $Arg_World -Force
+        Get-PassedArgsRoot | Should Be $Arg_World
+        Get-PassedArgsNoRoot | Should Be $Arg_World
+    }
+
+    It "Class execution reflects changes in module reloading with 'Remove-Module' and 'Import-Module'" {
+        Import-Module TestDrive:\TestModule.psm1 -ArgumentList $Arg_Hello
+        Get-PassedArgsRoot | Should Be $Arg_Hello
+        Get-PassedArgsNoRoot | Should Be $Arg_Hello
+
+        Remove-Module TestModule
+
+        Import-Module TestDrive:\TestModule.psm1 -ArgumentList $Arg_World
+        Get-PassedArgsRoot | Should Be $Arg_World
+        Get-PassedArgsNoRoot | Should Be $Arg_World
+    }
+}