Run PowerShell class static methods in the correct Runspace/SessionState (#4209)

Summary
----------
When the same script file gets executed in multiple Runspaces, the RuntimeType generated from the powershell class defined in the file will be shared among those Runspaces. For different Runspaces, different SessionState should be used to run the static methods.

Fix
---
We use the `SessionStateKeeper` to query for the correct `SessionState` to use for a static method call.
We already use `SessionStateKeeper` to track the `Runspace/SessionState` where a PowerShell class is defined, and thus we can leverage this for the static method invocation.

src/System.Management.Automation/engine/parser/PSType.cs

@@ -253,7 +253,7 @@ namespace System.Management.Automation.Language
             internal readonly TypeBuilder _staticHelpersTypeBuilder;
             private readonly Dictionary<string, PropertyMemberAst> _definedProperties;
             private readonly Dictionary<string, List<Tuple<FunctionMemberAst, Type[]>>> _definedMethods;
-            internal readonly List<Tuple<string, object>> _fieldsToInitForMemberFunctions;
+            internal readonly List<(string fieldName, IParameterMetadataProvider bodyAst, bool isStatic)> _fieldsToInitForMemberFunctions;
             private bool _baseClassHasDefaultCtor;
 
             /// <summary>
@@ -276,7 +276,7 @@ namespace System.Management.Automation.Language
                 DefineCustomAttributes(_typeBuilder, typeDefinitionAst.Attributes, _parser, AttributeTargets.Class);
                 _typeDefinitionAst.Type = _typeBuilder;
 
-                _fieldsToInitForMemberFunctions = new List<Tuple<string, object>>();
+                _fieldsToInitForMemberFunctions = new List<(string, IParameterMetadataProvider, bool)>();
                 _definedMethods = new Dictionary<string, List<Tuple<FunctionMemberAst, Type[]>>>(StringComparer.OrdinalIgnoreCase);
                 _definedProperties = new Dictionary<string, PropertyMemberAst>(StringComparer.OrdinalIgnoreCase);
 
@@ -884,8 +884,7 @@ namespace System.Management.Automation.Language
                 ilGenerator.EmitCall(OpCodes.Call, invokeHelper, null);
                 ilGenerator.Emit(OpCodes.Ret);
 
-                var methodWrapper = new ScriptBlockMemberMethodWrapper(ipmp);
-                _fieldsToInitForMemberFunctions.Add(Tuple.Create(wrapperFieldName, (object)methodWrapper));
+                _fieldsToInitForMemberFunctions.Add((wrapperFieldName, ipmp, isStatic));
             }
         }
 
@@ -1171,16 +1170,22 @@ namespace System.Management.Automation.Language
                         runtimeTypeAssigned = true;
                         var helperType = helper._staticHelpersTypeBuilder.CreateType();
 
+                        SessionStateKeeper sessionStateKeeper = new SessionStateKeeper();
+                        helperType.GetField(s_sessionStateKeeperFieldName, BindingFlags.NonPublic | BindingFlags.Static).SetValue(null, sessionStateKeeper);
+                        
                         if (helper._fieldsToInitForMemberFunctions != null)
                         {
                             foreach (var tuple in helper._fieldsToInitForMemberFunctions)
                             {
-                                helperType.GetField(tuple.Item1, BindingFlags.NonPublic | BindingFlags.Static)
-                                    .SetValue(null, tuple.Item2);
+                                // If the wrapper is for a static method, we need the sessionStateKeeper to determine the right SessionState to run.
+                                // If the wrapper is for an instance method, we use the SessionState that the instance is bound to, and thus don't need sessionStateKeeper.
+                                var methodWrapper = tuple.isStatic
+                                    ? new ScriptBlockMemberMethodWrapper(tuple.bodyAst, sessionStateKeeper)
+                                    : new ScriptBlockMemberMethodWrapper(tuple.bodyAst);
+                                helperType.GetField(tuple.fieldName, BindingFlags.NonPublic | BindingFlags.Static)
+                                    .SetValue(null, methodWrapper);
                             }
                         }
-
-                        helperType.GetField(s_sessionStateKeeperFieldName, BindingFlags.NonPublic | BindingFlags.Static).SetValue(null, new SessionStateKeeper());
                     }
                     catch (TypeLoadException e)
                     {

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

@@ -32,8 +32,8 @@ namespace System.Management.Automation.Internal
     /// </summary>
     public class SessionStateKeeper
     {
-        // We use ConditionalWeakTable, because if GC already collect Runspace,
-        // then there is no way to call a ctor on the type in this Runspace.
+        // We use ConditionalWeakTable, because if GC already collect Runspace, then there
+        // is no way to call a ctor or a static method on the type in this Runspace.
         private readonly ConditionalWeakTable<Runspace, SessionStateInternal> _stateMap;
 
         internal SessionStateKeeper()
@@ -66,21 +66,31 @@ namespace System.Management.Automation.Internal
         }
 
         /// <summary>
-        /// This method should be called only from generated ctors for PowerShell classes.
+        /// This method should be called only from
+        ///  - generated ctors for PowerShell classes, AND
+        ///  - ScriptBlockMemberMethodWrapper when invoking static methods of PowerShell classes.
         /// It's not intended to be a public API, but because we generate type in a different assembly it has to be public.
         /// Return type should be SessionStateInternal, but it violates accessibility consistency, so we use object.
         /// </summary>
         /// <remarks>
         /// By default, PowerShell class instantiation usually happens in the same Runspace where the class is defined. In
         /// that case, the created instance will be bound to the session state used to define that class in the Runspace.
-        /// However, if the instantiation happens in a different Runspace where the class is not defined, then the created
-        /// instance won't be bound to any session state.
+        /// However, if the instantiation happens in a different Runspace where the class is not defined, or it happens on
+        /// a thread without a default Runspace, then the created instance won't be bound to any session state.
         /// </remarks>
         /// <returns>SessionStateInternal</returns>
         public object GetSessionState()
         {
             SessionStateInternal ss = null;
-            _stateMap.TryGetValue(Runspace.DefaultRunspace, out ss);
+
+            // DefaultRunspace could be null when we reach here. For example, create instance of
+            // a PowerShell class by using reflection on a thread without DefaultRunspace.
+            // Make sure we call 'TryGetValue' with a non-null key, otherwise ArgumentNullException will be thrown.
+            Runspace defaultRunspace = Runspace.DefaultRunspace;
+            if (defaultRunspace != null)
+            {
+                _stateMap.TryGetValue(defaultRunspace, out ss);
+            }
             return ss;
         }
     }
@@ -91,31 +101,124 @@ namespace System.Management.Automation.Internal
         /// <summary>Used in codegen</summary>
         public static readonly object[] _emptyArgumentArray = Utils.EmptyArray<object>(); // See TypeDefiner.DefineTypeHelper.DefineMethodBody
 
-        // we use this _scriptBlock instance for static methods.
-        private Lazy<ScriptBlock> _scriptBlock;
-        private IParameterMetadataProvider _ast;
+        /// <summary>
+        /// Indicate the wrapper is for a static member method.
+        /// </summary>
+        private readonly bool _isStatic;
 
         /// <summary>
-        /// We use ThreadLocal boundScriptBlock to allow multi-thread execution of instance methods.
+        /// The SessionStateKeeper associated with the helper type generated from PowerShell class.
+        /// We query it for the SessionState to run static method in.
         /// </summary>
-        private ThreadLocal<ScriptBlock> _boundScriptBlock;
+        private readonly SessionStateKeeper _sessionStateKeeper;
 
+        /// <summary>
+        /// We use WeakReference object to point to the default SessionState because if GC already collect the SessionState,
+        /// or the Runspace it chains to is closed and disposed, then we cannot run the static method there anyways.
+        /// </summary>
+        /// <remakr>
+        /// The default SessionState is used only if a static method is called from a Runspace where the PowerShell class is
+        /// never defined, or is called on a thread without a default Runspace. Usage like those should be rare.
+        /// </remakr>
+        private readonly WeakReference<SessionStateInternal> _defaultSessionStateToUse;
+
+        /// <summary>
+        /// The body AST of the member method.
+        /// </summary>
+        private readonly IParameterMetadataProvider _ast;
+
+        /// <summary>
+        /// We use _scriptBlock instance to provide the shared CompiledScriptBlockData.
+        /// </summary>
+        private readonly Lazy<ScriptBlock> _scriptBlock;
+
+        /// <summary>
+        /// We use ThreadLocal boundScriptBlock to allow multi-thread execution of member methods.
+        /// </summary>
+        private readonly ThreadLocal<ScriptBlock> _boundScriptBlock;
+
+        /// <summary>
+        /// Constructor to be called when the wrapper is for a static member method.
+        /// </summary>
+        internal ScriptBlockMemberMethodWrapper(IParameterMetadataProvider ast, SessionStateKeeper sessionStateKeeper)
+            : this(ast)
+        {
+            _isStatic = true;
+            _sessionStateKeeper = sessionStateKeeper;
+            _defaultSessionStateToUse = new WeakReference<SessionStateInternal>(null);
+        }
+
+        /// <summary>
+        /// Constructor to be called when the wrapper is for an instance member method.
+        /// </summary>
         internal ScriptBlockMemberMethodWrapper(IParameterMetadataProvider ast)
         {
             _ast = ast;
+            // This 'Lazy<T>' constructor ensures that only a single thread can initialize the instance in a thread-safe manner.
             _scriptBlock = new Lazy<ScriptBlock>(() => new ScriptBlock(_ast, isFilter: false));
-            _boundScriptBlock = new ThreadLocal<ScriptBlock>(
-                () =>
-                {
-                    var sb = _scriptBlock.Value.Clone();
-                    return sb;
-                });
+            _boundScriptBlock = new ThreadLocal<ScriptBlock>(() => _scriptBlock.Value.Clone());
         }
 
+        /// <summary>
+        /// Initialization happens when the script that defines PowerShell class is executed.
+        /// This initialization is required only if this wrapper is for a static method.
+        /// </summary>
+        /// <remakr>
+        /// When the same script file gets executed multiple times, the .NET type generated from the PowerShell class
+        /// defined in the file will be shared in those executions, and thus this method will be called multiple times
+        /// possibly in the contexts of different Runspace/SessionState.
+        ///
+        /// We always use the SessionState from the most recent execution as the default SessionState, so be noted that
+        /// the default SessionState may change over time.
+        ///
+        /// This should be OK because the common usage is to run the static method in the same Runspace where the class
+        /// is declared, and thus we can always get the correct SessionState to use by querying the 'SessionStateKeeper'.
+        /// The default SessionState is used only if a static method is called from a Runspace where the class is never
+        /// defined, or is called on a thread without a default Runspace.
+        /// </remakr>
         internal void InitAtRuntime()
         {
-            var context = Runspace.DefaultRunspace.ExecutionContext;
-            _scriptBlock.Value.SessionStateInternal = context.EngineSessionState;
+            if (_isStatic)
+            {
+                // WeakReference<T>'s instance methods are not thread-safe, so we need the lock to guarantee
+                // 'SetTarget' and 'TryGetTarget' are not called by multiple threads at the same time.
+                lock (_defaultSessionStateToUse)
+                {
+                    var context = Runspace.DefaultRunspace.ExecutionContext;
+                    _defaultSessionStateToUse.SetTarget(context.EngineSessionState);
+                }
+            }
+        }
+
+        /// <summary>
+        /// Set the SessionState of the script block appropriately.
+        /// </summary>
+        private void PrepareScriptBlockToInvoke(object instance, object sessionStateInternal)
+        {
+            SessionStateInternal sessionStateToUse = null;
+            if (instance != null)
+            {
+                // Use the SessionState passed in, which is the one associated with the instance.
+                sessionStateToUse = (SessionStateInternal)sessionStateInternal;
+            }
+            else
+            {
+                // For static method, it's a little complex.
+                // - Check if the current default runspace is registered with the SessionStateKeeper. If so, use the registered SessionState.
+                // - Otherwise, check if default SessionState is still alive. If so, use the default SessionState.
+                // - Otherwise, the 'SessionStateInternal' property will be set to null, and thus the default runspace of the current thread will be used.
+                //              If the current thread doesn't have a default Runspace, then an InvalidOperationException will be thrown when invoking the
+                //              script block, which is expected.
+                sessionStateToUse = (SessionStateInternal)_sessionStateKeeper.GetSessionState();
+                if (sessionStateToUse == null)
+                {
+                    lock (_defaultSessionStateToUse)
+                    {
+                        _defaultSessionStateToUse.TryGetTarget(out sessionStateToUse);
+                    }
+                }
+            }
+            _boundScriptBlock.Value.SessionStateInternal = sessionStateToUse;
         }
 
         /// <summary>
@@ -125,18 +228,19 @@ namespace System.Management.Automation.Internal
         /// <param name="args"></param>
         public void InvokeHelper(object instance, object sessionStateInternal, object[] args)
         {
-            ScriptBlock sb;
-            if (instance != null)
+            try
             {
-                _boundScriptBlock.Value.SessionStateInternal = (SessionStateInternal)sessionStateInternal;
-                sb = _boundScriptBlock.Value;
+                PrepareScriptBlockToInvoke(instance, sessionStateInternal);
+                _boundScriptBlock.Value.InvokeAsMemberFunction(instance, args);
             }
-            else
+            finally
             {
-                sb = _scriptBlock.Value;
+                // '_boundScriptBlock.Value' for a thread will live until
+                //  - the thread is gone, OR
+                //  - the dyanmic assembly holding this wrapper instance is GC collected.
+                // We don't hold on the SessionState object, so that GC can collect it as appropriate.
+                _boundScriptBlock.Value.SessionStateInternal = null;
             }
-
-            sb.InvokeAsMemberFunction(instance, args);
         }
 
         /// <summary>
@@ -148,18 +252,19 @@ namespace System.Management.Automation.Internal
         /// <returns></returns>
         public T InvokeHelperT<T>(object instance, object sessionStateInternal, object[] args)
         {
-            ScriptBlock sb;
-            if (instance != null)
+            try
             {
-                _boundScriptBlock.Value.SessionStateInternal = (SessionStateInternal)sessionStateInternal;
-                sb = _boundScriptBlock.Value;
+                PrepareScriptBlockToInvoke(instance, sessionStateInternal);
+                return _boundScriptBlock.Value.InvokeAsMemberFunctionT<T>(instance, args);
             }
-            else
+            finally
             {
-                sb = _scriptBlock.Value;
+                // '_boundScriptBlock.Value' for a thread will live until
+                //  - the thread is gone, OR
+                //  - the dyanmic assembly holding this wrapper instance is GC collected.
+                // We don't hold on the SessionState object, so that GC can collect it as appropriate.
+                _boundScriptBlock.Value.SessionStateInternal = null;
             }
-
-            return sb.InvokeAsMemberFunctionT<T>(instance, args);
         }
     }
 

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

@@ -0,0 +1,120 @@
+Describe "Additional static method tests" -Tags "CI" {
+
+    Context "Basic static member methods" {
+        BeforeAll {
+            function Get-Name { "YES" }
+        }
+
+        It "test basic static constructor" {
+            class Foo {
+                static [string] $Name
+                static Foo() { [Foo]::Name = Get-Name }
+            }
+
+            [Foo]::Name | Should Be "Yes"
+        }
+
+        It "test basic static method" {
+            class Foo {
+                static [string] GetName() { return (Get-Name) }
+            }
+
+            [Foo]::GetName() | Should Be "Yes"
+        }
+    }
+
+    Context "Class defined in different Runspace" {
+        BeforeAll {
+@'
+class Foo
+{
+    static [string] $Name
+    static Foo() { [Foo]::Name = Get-Name }
+    
+    static [string] GetName()
+    {
+        return (Get-AnotherName)
+    }
+}
+'@ | Set-Content -Path $TestDrive\class.ps1 -Force
+
+            ## Define the functions that [Foo] depends on in the default Runspace.
+            function Get-Name { "Default Runspace - Name" }
+            function Get-AnotherName { "Default Runspace - AnotherName" }
+
+            ## Create another Runspace PS1
+            $ps1 = [powershell]::Create()
+            ## Create another Runspace PS2
+            $ps2 = [powershell]::Create()
+
+            function RunScriptInPS {
+                param(
+                    [powershell] $PowerShell,
+                    [string] $Script,
+                    [switch] $IgnoreResult
+                )
+                $result = $PowerShell.AddScript($Script).Invoke()
+                $PowerShell.Commands.Clear()
+
+                if (-not $IgnoreResult) {
+                    return $result
+                }
+            }
+
+            ## Define the functions that [Foo] depends on in PS1 Runspace.
+            RunScriptInPS -PowerShell $ps1 -Script "function Get-Name { 'PS1 Runspace - Name' }" -IgnoreResult
+            RunScriptInPS -PowerShell $ps1 -Script "function Get-AnotherName { 'PS1 Runspace - AnotherName' }" -IgnoreResult
+
+            # Dot source class.ps1 in the current Runspace
+            . $TestDrive\class.ps1
+            # And then dot source class.ps1 in the PS1 Runspace
+            RunScriptInPS -PowerShell $ps1 -Script ". $TestDrive\class.ps1" -IgnoreResult
+        }
+
+        AfterAll {
+            # Dispose both Runspaces
+            $ps1.Dispose()
+            $ps2.Dispose()
+        }
+
+        It "Static constructor should run in the triggering Runspace if the class has been defined in that Runspace" {
+            
+            ## The static constructor is triggered by accessing '[Foo]::Name' which happens in the current Runspace.
+            ## The class 'Foo' has been defined in the current Runspace, so it uses the current Runspace to run the
+            ## static constructor.
+            [Foo]::Name | Should Be "Default Runspace - Name"
+
+            ## Static constructor runs only once, so accessing the Name property in the PS1 Runspace will just return
+            ## the existing value.
+            RunScriptInPS -PowerShell $ps1 -Script "[Foo]::Name" | Should Be "Default Runspace - Name"
+        }
+
+        It "Static method use the Runspace where the call happens if the class has been defined in that Runspace" {
+
+            ## We call the static method in the current Runspace. The class 'Foo' has been defined
+            ## in the current Runspace, so it will use it to run the method.
+            [Foo]::GetName() | Should Be "Default Runspace - AnotherName"
+
+            ## We call the static method in PS1 Runspace. The class 'Foo' has been defined in the
+            ## PS1 Runspace, so it will use it to run the method.
+            RunScriptInPS -PowerShell $ps1 -Script "[Foo]::GetName()" | Should Be 'PS1 Runspace - AnotherName'
+        }
+
+        It "Static method use the default SessionState if it's called in a Runspace where the class is not defined" {
+
+            ## Define the functions that [Foo] depends on in PS2 Runspace.
+            RunScriptInPS -PowerShell $ps2 -Script "function Get-Name { 'PS2 Runspace - Name' }" -IgnoreResult
+            RunScriptInPS -PowerShell $ps2 -Script "function Get-AnotherName { 'PS2 Runspace - AnotherName' }" -IgnoreResult
+            
+            ## Define the function to call the static method 'GetName' on the passed-in type
+            RunScriptInPS -PowerShell $ps2 -Script 'function Call-GetName([type] $type) { $type::GetName() }' -IgnoreResult
+
+            ## We call the static method in PS2 Runspace. The class is not defined in this Runspace,
+            ## so the default SessionState will be used to run the method. The default SessionState
+            ## is always the one where the class was defined most recently. In this case, the class
+            ## was lastly defined in PS1 Runspace, os the method will be invoked in PS1 Runspace.
+            $result = $ps2.AddCommand("Call-GetName").AddParameter("type", [Foo]).Invoke()
+            $result | Should Be 'PS1 Runspace - AnotherName'
+        }
+    }
+}