Make PowerShell class not affiliate with Runspace when declaring th…

…e `NoRunspaceAffinity` attribute (#18138)
PowerShell · Sep 22, 2022 · 79b8140 · 79b8140
1 parent c978d46
commit 79b8140
Show file tree

Hide file tree

Showing 6 changed files with 81 additions and 5 deletions.
diff --git a/src/System.Management.Automation/engine/parser/PSType.cs b/src/System.Management.Automation/engine/parser/PSType.cs
@@ -1472,4 +1472,18 @@ private static void EmitLdarg(ILGenerator emitter, int c)
             }
         }
     }
+
+    /// <summary>
+    /// The attribute for a PowerShell class to not affiliate with a particular Runspace\SessionState.
+    /// </summary>
+    [AttributeUsage(AttributeTargets.Class)]
+    public sealed class NoRunspaceAffinityAttribute : ParsingBaseAttribute
+    {
+        /// <summary>
+        /// Initializes a new instance of the attribute.
+        /// </summary>
+        public NoRunspaceAffinityAttribute()
+        {
+        }
+    }
 }
diff --git a/src/System.Management.Automation/engine/parser/TypeResolver.cs b/src/System.Management.Automation/engine/parser/TypeResolver.cs
@@ -750,6 +750,7 @@ internal static class CoreTypes
                     { typeof(CimConverter),                                new[] { "cimconverter" } },
                     { typeof(ModuleSpecification),                         null },
                     { typeof(IPEndPoint),                                  new[] { "IPEndpoint" } },
+                    { typeof(NoRunspaceAffinityAttribute),                 new[] { "NoRunspaceAffinity" } },
                     { typeof(NullString),                                  new[] { "NullString" } },
                     { typeof(OutputTypeAttribute),                         new[] { "OutputType" } },
                     { typeof(object[]),                                    null },

diff --git a/src/System.Management.Automation/engine/runtime/Operations/ClassOps.cs b/src/System.Management.Automation/engine/runtime/Operations/ClassOps.cs
@@ -22,7 +22,7 @@ namespace System.Management.Automation.Internal
     /// <summary>
     /// Every Runspace in one process contains SessionStateInternal per module (module SessionState).
     /// Every RuntimeType is associated to only one SessionState in the Runspace, which creates it:
-    /// it's ever global state or a module state.
+    /// it's either global state or a module state.
     /// In the former case, module can be imported from the different runspaces in the same process.
     /// And so runspaces will share RuntimeType. But in every runspace, Type is associated with just one SessionState.
     /// We want type methods to be able access $script: variables and module-specific methods.

diff --git a/src/System.Management.Automation/engine/runtime/Operations/MiscOps.cs b/src/System.Management.Automation/engine/runtime/Operations/MiscOps.cs
@@ -2300,8 +2300,13 @@ internal static void InitPowerShellTypesAtRuntime(TypeDefinitionAst[] types)
                 Diagnostics.Assert(t.Type != null, "TypeDefinitionAst.Type cannot be null");
                 if (t.IsClass)
                 {
-                    var helperType =
-                        t.Type.Assembly.GetType(t.Type.FullName + "_<staticHelpers>");
+                    if (t.Type.IsDefined(typeof(NoRunspaceAffinityAttribute), inherit: true))
+                    {
+                        // Skip the initialization for session state affinity.
+                        continue;
+                    }
+
+                    var helperType = t.Type.Assembly.GetType(t.Type.FullName + "_<staticHelpers>");
                     Diagnostics.Assert(helperType != null, "no corresponding " + t.Type.FullName + "_<staticHelpers> type found");
                     foreach (var p in helperType.GetFields(BindingFlags.Static | BindingFlags.NonPublic))
                     {

diff --git a/test/powershell/Language/Classes/Scripting.Classes.NoRunspaceAffinity.Tests.ps1 b/test/powershell/Language/Classes/Scripting.Classes.NoRunspaceAffinity.Tests.ps1
@@ -0,0 +1,52 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+Describe "Class can be defined without Runspace affinity" -Tags "CI" {
+
+    It "Applying the 'NoRunspaceAffinity' attribute make the class not affiliate with a particular Runspace/SessionState" {
+        [NoRunspaceAffinity()]
+        class NoAffinity {
+            [string] $Name;
+            [int] $RunspaceId;
+
+            NoAffinity() {
+               $this.RunspaceId = [runspace]::DefaultRunspace.Id
+            }
+
+            static [int] Echo() {
+                return [runspace]::DefaultRunspace.Id
+            }
+
+            [int] SetAndEcho([string] $value) {
+                $this.Name = $value
+                return [runspace]::DefaultRunspace.Id
+            }
+        }
+
+        $t = [NoAffinity]
+        $o = [NoAffinity]::new()
+
+        ## Running directly should use the current Runspace/SessionState.
+        $t::Echo() | Should -Be $Host.Runspace.Id
+        $o.RunspaceId | Should -Be $Host.Runspace.Id
+        $o.SetAndEcho('Blue') | Should -Be $Host.Runspace.Id
+        $o.Name | Should -Be 'Blue'
+
+        ## Running in a new Runspace should use that Runspace and its current SessionState.
+        try {
+            $ps = [powershell]::Create()
+            $ps.AddScript('function CallEcho($type) { $type::Echo() }').Invoke() > $null; $ps.Commands.Clear()
+            $ps.AddScript('function CallSetAndEcho($obj) { $obj.SetAndEcho(''Hello world'') }').Invoke() > $null; $ps.Commands.Clear()
+            $ps.AddScript('function GetName($obj) { $obj.Name }').Invoke() > $null; $ps.Commands.Clear()
+            $ps.AddScript('function NewObj($type) { $type::new().RunspaceId }').Invoke() > $null; $ps.Commands.Clear()
+
+            $ps.AddCommand('CallEcho').AddArgument($t).Invoke() | Should -Be $ps.Runspace.Id; $ps.Commands.Clear()
+            $ps.AddCommand('CallSetAndEcho').AddArgument($o).Invoke() | Should -Be $ps.Runspace.Id; $ps.Commands.Clear()
+            $ps.AddCommand('GetName').AddArgument($o).Invoke() | Should -Be 'Hello world'; $ps.Commands.Clear()
+            $ps.AddCommand('NewObj').AddArgument($t).Invoke() | Should -Be $ps.Runspace.Id; $ps.Commands.Clear()
+        }
+        finally {
+            $ps.Dispose()
+        }
+    }
+}
diff --git a/test/powershell/Language/Parser/TypeAccelerator.Tests.ps1 b/test/powershell/Language/Parser/TypeAccelerator.Tests.ps1
@@ -410,15 +410,19 @@ Describe "Type accelerators" -Tags "CI" {
                     Accelerator = 'ordered'
                     Type        = [System.Collections.Specialized.OrderedDictionary]
                 }
+                @{
+                    Accelerator = 'NoRunspaceAffinity'
+                    Type        = [System.Management.Automation.Language.NoRunspaceAffinityAttribute]
+                }
             )
 
             if ( !$IsWindows )
             {
-                $totalAccelerators = 101
+                $totalAccelerators = 102
             }
             else
             {
-                $totalAccelerators = 106
+                $totalAccelerators = 107
 
                 $extraFullPSAcceleratorTestCases = @(
                     @{