From 97fe7ee7c128bbc701a4230de59dcc730fbdc218 Mon Sep 17 00:00:00 2001 From: Dongbo Wang Date: Fri, 2 Dec 2016 13:09:11 -0800 Subject: [PATCH 1/2] Fix powershell class to use the current EngineSessionState for execution. 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 | 38 +++++++++++++++++++ 2 files changed, 58 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 6fb1d9d8ba2..1e451eb8653 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 @@ internal SessionStateKeeper() 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. } /// diff --git a/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1 b/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1 index 00dfc40246b..e1475fa30c6 100644 --- a/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1 +++ b/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1 @@ -71,3 +71,41 @@ 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 } +'@ + } + + AfterEach { + Remove-Module TestModule -Force -ErrorAction SilentlyContinue + } + + It "Class execution reflects changes in module reloading with '-Force'" { + Import-Module TestDrive:\TestModule.psm1 -ArgumentList 'value-1' + Get-PassedArgsRoot | Should Be 'value-1' + Get-PassedArgsNoRoot | Should Be 'value-1' + + Import-Module TestDrive:\TestModule.psm1 -ArgumentList 'value-2' -Force + Get-PassedArgsRoot | Should Be 'value-2' + Get-PassedArgsNoRoot | Should Be 'value-2' + } + + It "Class execution reflects changes in module reloading with 'Remove-Module' and 'Import-Module'" { + Import-Module TestDrive:\TestModule.psm1 -ArgumentList 'value-1' + Get-PassedArgsRoot | Should Be 'value-1' + Get-PassedArgsNoRoot | Should Be 'value-1' + + Remove-Module TestModule + + Import-Module TestDrive:\TestModule.psm1 -ArgumentList 'value-2' + Get-PassedArgsRoot | Should Be 'value-2' + Get-PassedArgsNoRoot | Should Be 'value-2' + } +} From 33de8924bef4e11fe468884f11ee234ec793a8bf Mon Sep 17 00:00:00 2001 From: Dongbo Wang Date: Tue, 6 Dec 2016 11:08:45 -0800 Subject: [PATCH 2/2] Use variable for repeated argument strings --- .../Scripting.Classes.Modules.Tests.ps1 | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1 b/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1 index e1475fa30c6..3098d2c381b 100644 --- a/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1 +++ b/test/powershell/Language/Classes/Scripting.Classes.Modules.Tests.ps1 @@ -81,6 +81,8 @@ class Root { $passedIn = $passedArgs } function Get-PassedArgsRoot { [Root]::new().passedIn } function Get-PassedArgsNoRoot { $passedArgs } '@ + $Arg_Hello = 'Hello' + $Arg_World = 'World' } AfterEach { @@ -88,24 +90,24 @@ function Get-PassedArgsNoRoot { $passedArgs } } It "Class execution reflects changes in module reloading with '-Force'" { - Import-Module TestDrive:\TestModule.psm1 -ArgumentList 'value-1' - Get-PassedArgsRoot | Should Be 'value-1' - Get-PassedArgsNoRoot | Should Be 'value-1' + 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 'value-2' -Force - Get-PassedArgsRoot | Should Be 'value-2' - Get-PassedArgsNoRoot | Should Be 'value-2' + 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 'value-1' - Get-PassedArgsRoot | Should Be 'value-1' - Get-PassedArgsNoRoot | Should Be 'value-1' + 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 'value-2' - Get-PassedArgsRoot | Should Be 'value-2' - Get-PassedArgsNoRoot | Should Be 'value-2' + Import-Module TestDrive:\TestModule.psm1 -ArgumentList $Arg_World + Get-PassedArgsRoot | Should Be $Arg_World + Get-PassedArgsNoRoot | Should Be $Arg_World } }