From c982f30d4e5d70e6d153ef47c9990a87f0ada5df Mon Sep 17 00:00:00 2001
From: Dongbo Wang <dongb.wang@gmail.com>
Date: Fri, 9 Dec 2016 11:06:07 -0800
Subject: [PATCH] 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.
---
 .../engine/runtime/Operations/ClassOps.cs     | 23 +++++++++--
 .../Scripting.Classes.Modules.Tests.ps1       | 40 +++++++++++++++++++
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/src/System.Management.Automation/engine/runtime/Operations/ClassOps.cs b/src/System.Management.Automation/engine/runtime/Operations/ClassOps.cs
index 6fb1d9d8ba..1e451eb865 100644
--- a/src/System.Management.Automation/engine/runtime/Operations/ClassOps.cs
+++ b/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>
diff --git a/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1 b/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1
index 00dfc40246..3098d2c381 100644
--- a/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1
+++ b/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
+    }
+}