Changes to Async classes (#16)

SeeminglyScience · web-flow · commit 617c4d2f4f71 · 2017-08-12T08:45:39.000-04:00
- Fix creating async delegates with different return types

- Add ability to access delegate argument from $args, $PSItem,
  or the parameter name from the delegates invoke method.

- New instances of PSTaskFactory now have their own runspace pool

- Await is no longer blocking
diff --git a/module/Classes/Async.ps1 b/module/Classes/Async.ps1
@@ -3,49 +3,84 @@ using namespace System.Collections.ObjectModel
 using namespace System.Linq.Expressions
 using namespace System.Management.Automation
 using namespace System.Management.Automation.Runspaces
+using namespace System.Threading
 using namespace System.Threading.Tasks
 
 # Static class that facilitates the use of traditional .NET async techniques in PowerShell.
-class AsyncOps
-{
+class AsyncOps {
     static [PSTaskFactory] $Factory = [PSTaskFactory]::new();
-    static [RunspacePool] $RunspacePool;
 
-    # Hides the result property on a Task object. This is done because the getter for Result waits
-    # for the task to finish, even if just output to the console.
-    static [Task] HideResult([Task] $target) {
+    static [Task] ContinueWithCodeMethod([psobject] $instance, [scriptblock] $continuationAction) {
+        $delegate = [AsyncOps]::CreateAsyncDelegate(
+            $continuationAction,
+            [Action[Task[Collection[psobject]]]])
+
+        return [AsyncOps]::PrepareTask($instance.psadapted.ContinueWith($delegate))
+    }
+
+    # - Hides the result property on a Task object. This is done because the getter for Result waits
+    #   for the task to finish, even if just output to the console.
+    # - Adds ContinueWith code method that wraps scriptblocks with CreateAsyncDelegate
+    static [Task] PrepareTask([Task] $target) {
         $propertyList    = $target.psobject.Properties.Name -notmatch 'Result' -as [string[]]
         $propertySet     = [PSPropertySet]::new('DefaultDisplayPropertySet', $propertyList) -as [PSMemberInfo[]]
         $standardMembers = [PSMemberSet]::new('PSStandardMembers', $propertySet)
 
         $target.psobject.Members.Add($standardMembers)
 
+        $target.psobject.Methods.Add(
+            [PSCodeMethod]::new(
+                'ContinueWith',
+                [AsyncOps].GetMethod('ContinueWithCodeMethod')))
+
         return $target
     }
 
-    # Create a delegate from a scriptblock that can be used in threads without runspaces, like those
-    # used in Tasks or AsyncCallbacks.
     static [MulticastDelegate] CreateAsyncDelegate([scriptblock] $function, [type] $delegateType) {
-        # Create a runspace pool the first time this method is invoked.
-        if (-not [AsyncOps]::RunspacePool) {
-            [AsyncOps]::RunspacePool = [runspacefactory]::CreateRunspacePool(1, 4)
-            [AsyncOps]::RunspacePool.Open()
-        }
+        return [AsyncOps]::CreateAsyncDelegate($function, $delegateType, [AsyncOps]::Factory)
+    }
 
+    # Create a delegate from a scriptblock that can be used in threads without runspaces, like those
+    # used in Tasks or AsyncCallbacks.
+    static [MulticastDelegate] CreateAsyncDelegate([scriptblock] $function, [type] $delegateType, [PSTaskFactory] $factory) {
+        $invokeMethod = $delegateType.GetMethod('Invoke')
+        $returnType = $invokeMethod.ReturnType
         # Create a parameter expression for each parameter the delegate takes.
-        $parameters = $delegateType.
-            GetMethod('Invoke').
+        $parameters = $invokeMethod.
             GetParameters().
             ForEach{ [Expression]::Parameter($PSItem.ParameterType, $PSItem.Name) }
 
-        # Set AsyncState variable that will hold delegate arguments and/or state.
-        $preparedScript = 'param($AsyncState) . {{ {0} }}' -f $function
+        $scriptParameters = [string]::Empty
+        if ($parameters) {
+            $scriptParameters = '$' + ($invokeMethod.GetParameters().Name -join ', $')
+        }
+
+        # Allow access to parameters in the following ways:
+        # - By the name given to them by the delegate's invoke method
+        # - $args
+        # - $PSItem/$_  (first parameter only)
+        $preparedScript =
+            'param({0}) process {{ return {{ {1} }}.InvokeReturnAsIs($PSBoundParameters.Values) }}' -f
+            $scriptParameters,
+            $function
 
         # Prepare variable and constant expressions.
-        $pool       = [Expression]::Property($null, [AsyncOps], 'RunspacePool')
         $scriptText = [Expression]::Constant($preparedScript, [string])
         $ps         = [Expression]::Variable([powershell], 'ps')
-        $result     = [Expression]::Variable([Collection[psobject]], 'result')
+        $collectionResultType = [Collection[psobject]]
+        if ($returnType -ne [void] -and $returnType -ne [Collection[psobject]]) {
+            $collectionResultType = [Collection`1].MakeGenericType($returnType)
+        }
+
+        $result     = [Expression]::Variable($collectionResultType, 'result')
+        $psInput    = [Expression]::Variable([Object[]], 'psInput')
+        $guid       = [Expression]::Constant($factory.InstanceId, [guid])
+        $pool       = [Expression]::Property(
+            [Expression]::Property(
+                [Expression]::Property($null, [PSTaskFactory], 'Instances'),
+                'Item',
+                $guid),
+            'RunspacePool')
 
         # Group the expressions for the body by creating them in a scriptblock.
         [Expression[]]$expressions = & {
@@ -57,12 +92,34 @@ class AsyncOps
                 [Expression]::Call($ps, 'AddArgument', @(), $parameter)
             }
 
-            [Expression]::Assign($result, [Expression]::Call($ps, 'Invoke', @(), @()))
+            $invokeArgs = @()
+            if ($parameters) {
+                [Expression]::Assign(
+                    $psInput,
+                    [Expression]::NewArrayInit([object], $parameters[0] -as [Expression[]]))
+
+                $invokeArgs = @($psInput)
+            }
+
+            $invokeTypeArgs = @()
+            if ($returnType -ne [void] -and $returnType -ne [Collection[psobject]]) {
+                $invokeTypeArgs = @($returnType)
+            }
+
+            [Expression]::Assign($result, [Expression]::Call($ps, 'Invoke', $invokeTypeArgs, $invokeArgs))
             [Expression]::Call($ps, 'Dispose', @(), @())
-            $result
+            if ($returnType -ne [void]) {
+                [Expression]::Call(
+                    [LanguagePrimitives],
+                    'ConvertTo',
+                    @($returnType),
+                    $result -as [Expression[]])
+            } else {
+                $result
+            }
         }
 
-        $block  = [Expression]::Block([ParameterExpression[]]($ps, $result), $expressions)
+        $block  = [Expression]::Block([ParameterExpression[]]($ps, $result, $psInput), $expressions)
         $lambda = [Expression]::Lambda(
             $delegateType,
             $block,
@@ -73,16 +130,75 @@ class AsyncOps
 
 # A TaskFactory implementation that creates tasks that run scriptblocks in a runspace pool.
 class PSTaskFactory : TaskFactory[Collection[psobject]] {
+    hidden static [Dictionary[guid, PSTaskFactory]] $Instances = [Dictionary[guid, PSTaskFactory]]::new()
+
+    hidden [RunspacePool] $RunspacePool;
+    hidden [guid] $InstanceId;
+    hidden [bool] $IsDisposed = $false;
+
+    PSTaskFactory() : base() {
+        $this.Initialize()
+    }
+
+    PSTaskFactory([CancellationToken] $cancellationToken) : base($cancellationToken) {
+        $this.Initialize()
+    }
+
+    PSTaskFactory([TaskScheduler] $scheduler) : base($scheduler) {
+        $this.Initialize()
+    }
+
+    PSTaskFactory(
+        [TaskCreationOptions] $creationOptions,
+        [TaskContinuationOptions] $continuationOptions)
+        : base($creationOptions, $continuationOptions) {
+        $this.Initialize()
+    }
+
+    PSTaskFactory(
+        [CancellationToken] $cancellationToken,
+        [TaskCreationOptions] $creationOptions,
+        [TaskContinuationOptions] $continuationOptions,
+        [TaskScheduler] $scheduler)
+        : base($cancellationToken, $creationOptions, $continuationAction, $scheduler) {
+        $this.Initialize()
+    }
+
+    hidden [void] Initialize() {
+        $this.RunspacePool = [runspacefactory]::CreateRunspacePool(1, 4)
+        $this.RunspacePool.Open()
+        $this.InstanceId = [guid]::NewGuid()
+        [PSTaskFactory]::Instances.Add($this.InstanceId, $this)
+    }
+
+    # Can't implement IDisposable while inheriting a generic class because of a parse error, need
+    # to create an issue.
+    [void] Dispose() {
+        $this.AssertNotDisposed()
+        [PSTaskFactory]::Instances.Remove($this.InstanceId)
+        $this.RunspacePool.Dispose()
+        $this.IsDisposed = $true
+    }
+
+    hidden [void] AssertNotDisposed() {
+        if ($this.IsDisposed) {
+            throw [InvalidOperationException]::new(
+                'Cannot perform operation because object "PSTaskFactory" has already been disposed.')
+        }
+    }
+
     # Shortcut to AsyncOps.CreateAsyncDelegate
     hidden [MulticastDelegate] Wrap([scriptblock] $function, [type] $delegateType) {
-        return [AsyncOps]::CreateAsyncDelegate($function, $delegateType)
+        $this.AssertNotDisposed()
+        return [AsyncOps]::CreateAsyncDelegate($function, $delegateType, $this)
     }
 
     # The remaining functions implement methods from TaskFactory. All of these methods call the base
     # method after wrapping the scriptblock to create a delegate that will work in tasks.
     [Task[Collection[psobject]]] ContinueWhenAll([Task[]] $tasks, [scriptblock] $continuationAction) {
+        $this.AssertNotDisposed()
         $delegateType = [Func`2].MakeGenericType([Task[]], [Collection[psobject]])
-        return [AsyncOps]::HideResult(
+        return [AsyncOps]::PrepareTask(
             ([TaskFactory[Collection[psobject]]]$this).ContinueWhenAll(
                 $tasks,
                 $this.Wrap($continuationAction, $delegateType),
@@ -92,8 +208,9 @@ class PSTaskFactory : TaskFactory[Collection[psobject]] {
     }
 
     [Task[Collection[psobject]]] ContinueWhenAny([Task[]] $tasks, [scriptblock] $continuationAction) {
-        $delegateType = [Func`2].MakeGenericType([Task[]], [Collection[psobject]])
-        return [AsyncOps]::HideResult(
+        $this.AssertNotDisposed()
+        $delegateType = [Func`2].MakeGenericType([Task], [Collection[psobject]])
+        return [AsyncOps]::PrepareTask(
             ([TaskFactory[Collection[psobject]]]$this).ContinueWhenAny(
                 $tasks,
                 $this.Wrap($continuationAction, $delegateType),
@@ -103,7 +220,8 @@ class PSTaskFactory : TaskFactory[Collection[psobject]] {
     }
 
     [Task[Collection[psobject]]] StartNew([scriptblock] $function) {
-        return [AsyncOps]::HideResult(
+        $this.AssertNotDisposed()
+        return [AsyncOps]::PrepareTask(
             ([TaskFactory[Collection[psobject]]]$this).StartNew(
                 $this.Wrap($function, [Func[Collection[psobject]]]),
                 $this.CancellationToken,
@@ -112,7 +230,8 @@ class PSTaskFactory : TaskFactory[Collection[psobject]] {
     }
 
     [Task[Collection[psobject]]] StartNew([scriptblock] $function, [object] $state) {
-        return [AsyncOps]::HideResult(
+        $this.AssertNotDisposed()
+        return [AsyncOps]::PrepareTask(
             ([TaskFactory[Collection[psobject]]]$this).StartNew(
                 $this.Wrap($function, [Func[object, Collection[psobject]]]),
                 $state,
@@ -133,6 +252,7 @@ function async {
         $ArgumentList
     )
     process {
+        if (-not $scriptblock) { return }
         [AsyncOps]::Factory.StartNew($ScriptBlock, $ArgumentList)
     }
 }
@@ -148,10 +268,35 @@ function await {
         $taskList = [List[Task]]::new()
     }
     process {
-        $taskList.Add($Task)
+        if ($Task) { $taskList.Add($Task) }
     }
     end {
-        return $taskList.Result
+        if (-not $taskList.Count) { return }
+
+        $finished = $false
+        while (-not $finished) {
+            $finished = $taskList.TrueForAll({
+                param([Task]$task)
+
+                $task.IsCompleted -or
+                $task.IsCanceled -or
+                $task.IsFaulted
+            })
+        }
+        $taskList.ToArray().ForEach{
+            if ($PSItem.IsFaulted -and $PSItem.Exception) {
+                if (-not ($exception = $PSItem.Exception.InnerException)) {
+                    $exception = $PSItem.Exception
+                }
+                $PSCmdlet.WriteError(
+                    [ErrorRecord]::new(
+                        $exception,
+                        $exception.GetType().Name,
+                        [ErrorCategory]::InvalidResult,
+                        $PSItem))
+            }
+            $PSItem.Result
+        }
     }
 }
 
@@ -162,7 +307,7 @@ function ContinueWith {
         $ContinuationAction,
 
         [switch]
-        $Any,
+        $WhenAny,
 
         [Parameter(ValueFromPipeline)]
         [Task]
@@ -172,10 +317,12 @@ function ContinueWith {
         $taskList = [List[Task]]::new()
     }
     process {
-        $taskList.Add($Task)
+        if ($Task) { $taskList.Add($Task) }
     }
     end {
-        if ($Any.IsPresent) {
+        if (-not $taskList.Count) { return }
+
+        if ($WhenAny.IsPresent) {
             return [AsyncOps]::Factory.ContinueWhenAny($taskList, $ContinuationAction)
         }
         return [AsyncOps]::Factory.ContinueWhenAll($taskList, $ContinuationAction)
