Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
daxian-dbw merged 4 commits into from
Jul 15, 2017
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8ecbd1e
Run PowerShell class static methods in the correct Runspace/SessionState
daxian-dbw Jul 8, 2017
11606c9
Address review comments
daxian-dbw Jul 10, 2017
01fd93f
Address review comments - 2
daxian-dbw Jul 10, 2017
7cb615f
Fix typo in comment
lzybkr Jul 14, 2017
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 13 additions & 8 deletions src/System.Management.Automation/engine/parser/PSType.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,7 @@ private class DefineTypeHelper
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>
Expand All @@ -276,7 +276,7 @@ public DefineTypeHelper(Parser parser, ModuleBuilder module, TypeDefinitionAst t
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);

Expand Down Expand Up @@ -884,8 +884,7 @@ private void DefineMethodBody(
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));
}
}

Expand Down Expand Up @@ -1171,16 +1170,22 @@ internal static Assembly DefineTypes(Parser parser, Ast rootAst, TypeDefinitionA
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)
{
Expand Down
175 changes: 140 additions & 35 deletions src/System.Management.Automation/engine/runtime/Operations/ClassOps.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -66,21 +66,31 @@ internal void RegisterRunspace()
}

/// <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;
}
}
Expand All @@ -91,31 +101,124 @@ public class ScriptBlockMemberMethodWrapper
/// <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>
/// The SessionStateKeeper associated with the helper type generated from PowerShell class.
/// We query it for the SessionState to run static method in.
/// </summary>
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>
/// We use ThreadLocal boundScriptBlock to allow multi-thread execution of instance methods.
/// Constructor to be called when the wrapper is for a static member method.
/// </summary>
private ThreadLocal<ScriptBlock> _boundScriptBlock;
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>
Expand All @@ -125,18 +228,19 @@ internal void InitAtRuntime()
/// <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>
Expand All @@ -148,18 +252,19 @@ public void InvokeHelper(object instance, object sessionStateInternal, object[]
/// <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);
}
}

Expand Down
Loading