PoshRSJob.psm1

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
$ScriptPath = Split-Path $MyInvocation.MyCommand.Path
#region Custom Object
Write-Verbose "Creating custom RSJob object"
Add-Type -TypeDefinition @"
    using System;
 
    namespace PoshRS.PowerShell
    {
        public class RSJob
        {
            public string Name;
            public int ID;
            public System.Management.Automation.PSInvocationState State;
            public System.Guid InstanceID;
            public object Handle;
            public object Runspace;
            public System.Management.Automation.PowerShell InnerJob;
            public System.Threading.ManualResetEvent Finished;
            public string Command;
            public System.Management.Automation.PSDataCollection<System.Management.Automation.ErrorRecord> Error;
            public System.Management.Automation.PSDataCollection<System.Management.Automation.VerboseRecord> Verbose;
            public System.Management.Automation.PSDataCollection<System.Management.Automation.DebugRecord> Debug;
            public System.Management.Automation.PSDataCollection<System.Management.Automation.WarningRecord> Warning;
            public System.Management.Automation.PSDataCollection<System.Management.Automation.ProgressRecord> Progress;
            public bool HasMoreData;
            public bool HasErrors;
            public object Output;
            public System.Guid RunspacePoolID;
            public bool Completed = false;
        }
        public class RSRunspacePool
        {
            public System.Management.Automation.Runspaces.RunspacePool RunspacePool;
            public System.Management.Automation.Runspaces.RunspacePoolState State;
            public int AvailableJobs;
            public int MaxJobs;
            public DateTime LastActivity = DateTime.MinValue;
            public System.Guid RunspacePoolID;
            public bool CanDispose = false;
        }
        public class V2UsingVariable
        {
            public string Name;
            public string NewName;
            public object Value;
            public string NewVarName;
        }
    }
"@
 -Language CSharpVersion3
#endregion Custom Object

#region RSJob Collections
Write-Verbose "Creating RS collections"
New-Variable Jobs -Value ([System.Collections.ArrayList]::Synchronized([System.Collections.ArrayList]@())) -Option AllScope,ReadOnly -Scope Global
New-Variable JobCleanup -Value ([hashtable]::Synchronized(@{})) -Option AllScope,ReadOnly -Scope Global
New-Variable JobID -Value ([int64]0) -Option AllScope,ReadOnly -Scope Global
New-Variable RunspacePools -Value ([System.Collections.ArrayList]::Synchronized([System.Collections.ArrayList]@())) -Option AllScope,ReadOnly -Scope Global
New-Variable RunspacePoolCleanup -Value ([hashtable]::Synchronized(@{})) -Option AllScope,ReadOnly -Scope Global
#endregion RSJob Collections

#region Cleanup Routine
Write-Verbose "Creating routine to monitor RS jobs"
$jobCleanup.Flag=$True
$jobCleanup.Host = $Host
$jobcleanup.Runspace =[runspacefactory]::CreateRunspace()   
$jobcleanup.Runspace.Open()         
$jobcleanup.Runspace.SessionStateProxy.SetVariable("jobCleanup",$jobCleanup)     
$jobcleanup.Runspace.SessionStateProxy.SetVariable("jobs",$jobs) 
$jobCleanup.PowerShell = [PowerShell]::Create().AddScript({
    #Routine to handle completed runspaces
    Do {   
        [System.Threading.Monitor]::Enter($Jobs.syncroot) 
        Foreach($job in $jobs) {
            $job.state = $job.InnerJob.InvocationStateInfo.State
            If ($job.Handle.isCompleted -AND (-NOT $Job.Completed)) {   
                #$jobCleanup.Host.UI.WriteVerboseLine("$($Job.Id) completed")
                Try {           
                    $data = $job.InnerJob.EndInvoke($job.Handle)
                } Catch {
                    $CaughtErrors = $Error
                    #$jobCleanup.Host.UI.WriteVerboseLine("$($Job.Id) Caught terminating Error in job: $_")
                }
                #$jobCleanup.Host.UI.WriteVerboseLine("$($Job.Id) Checking for errors")
                If ($job.InnerJob.Streams.Error -OR $CaughtErrors) {
                    #$jobCleanup.Host.UI.WriteVerboseLine("$($Job.Id) Errors Found!")
                    $ErrorList = New-Object System.Management.Automation.PSDataCollection[System.Management.Automation.ErrorRecord]
                    If ($job.InnerJob.Streams.Error) {
                        ForEach ($Err in $job.InnerJob.Streams.Error) {
                            #$jobCleanup.Host.UI.WriteVerboseLine("`t$($Job.Id) Adding Error")
                            [void]$ErrorList.Add($Err)
                        }
                    }
                    If ($CaughtErrors) {
                        ForEach ($Err in $CaughtErrors) {
                            #$jobCleanup.Host.UI.WriteVerboseLine("`t$($Job.Id) Adding Error")
                            [void]$ErrorList.Add($Err)
                        }                    
                    }
                    $job.Error = $ErrorList
                }
                #$jobCleanup.Host.UI.WriteVerboseLine("$($Job.Id) Disposing job")
                $job.InnerJob.dispose() 
                $job.Completed = $True  
                If ((Get-Variable data -ErrorAction SilentlyContinue).Value) {
                    $job.output = $data
                    $job.HasMoreData = $True
                    Remove-Variable data -ErrorAction SilentlyContinue                    
                }                            
                $Error.Clear()
            } 
        }        
        [System.Threading.Monitor]::Exit($Jobs.syncroot)
        Start-Sleep -Milliseconds 100     
    } while ($jobCleanup.Flag)
})
$jobCleanup.PowerShell.Runspace = $jobcleanup.Runspace
$jobCleanup.Handle = $jobCleanup.PowerShell.BeginInvoke()  

Write-Verbose "Creating routine to monitor Runspace Pools"
$RunspacePoolCleanup.Flag=$True
$RunspacePoolCleanup.Host=$Host
#5 minute timeout for unused runspace pools
$RunspacePoolCleanup.Timeout = [timespan]::FromMinutes(1).Ticks
$RunspacePoolCleanup.Runspace =[runspacefactory]::CreateRunspace()
 
#Create Type Collection so the object will work properly
$Types = Get-ChildItem "$($PSScriptRoot)\TypeData" -Filter *Types* | Select -ExpandProperty Fullname 
$Types | ForEach { 
    $TypeEntry = New-Object System.Management.Automation.Runspaces.TypeConfigurationEntry -ArgumentList $_ 
    $RunspacePoolCleanup.Runspace.RunspaceConfiguration.types.Append($TypeEntry) 
} 
  
$RunspacePoolCleanup.Runspace.Open()         
$RunspacePoolCleanup.Runspace.SessionStateProxy.SetVariable("RunspacePoolCleanup",$RunspacePoolCleanup)     
$RunspacePoolCleanup.Runspace.SessionStateProxy.SetVariable("RunspacePools",$RunspacePools) 
$RunspacePoolCleanup.Runspace.SessionStateProxy.SetVariable("ParentHost",$Host) 
$RunspacePoolCleanup.PowerShell = [PowerShell]::Create().AddScript({
    #Routine to handle completed runspaces
    Do { 
        $DisposeRunspacePools=$False  
        If ($RunspacePools.Count -gt 0) { 
            #$ParentHost.ui.WriteVerboseLine("$($RunspacePools | Out-String)")
            [System.Threading.Monitor]::Enter($RunspacePools.syncroot) 
            Foreach($RunspacePool in $RunspacePools) {                
                #$ParentHost.ui.WriteVerboseLine("RunspacePool <$($RunspacePool.RunspaceID)> | MaxJobs: $($RunspacePool.MaxJobs) | AvailJobs: $($RunspacePool.AvailableJobs)")
                If (($RunspacePool.AvailableJobs -eq $RunspacePool.MaxJobs) -AND $RunspacePools.LastActivity.Ticks -ne 0) {
                    If ((Get-Date).Ticks - $RunspacePool.LastActivity.Ticks -gt $RunspacePoolCleanup.Timeout) {
                        #Dispose of runspace pool
                        $RunspacePool.RunspacePool.Close()
                        $RunspacePool.RunspacePool.Dispose()
                        $RunspacePool.CanDispose = $True
                        $DisposeRunspacePools=$True
                    }
                } Else {
                    $RunspacePool.LastActivity = (Get-Date)
                }               
            }       
            #Remove runspace pools
            If ($DisposeRunspacePools) {
                $TempCollection = $RunspacePools.Clone()
                $TempCollection | Where {
                    $_.CanDispose
                } | ForEach {
                    #$ParentHost.ui.WriteVerboseLine("Removing runspacepool <$($_.RunspaceID)>")
                    [void]$RunspacePools.Remove($_)
                }
            }
            Remove-Variable TempCollection
            [System.Threading.Monitor]::Exit($RunspacePools.syncroot)
        }
        Start-Sleep -Milliseconds 5000     
    } while ($RunspacePoolCleanup.Flag)
})
$RunspacePoolCleanup.PowerShell.Runspace = $RunspacePoolCleanup.Runspace
$RunspacePoolCleanup.Handle = $RunspacePoolCleanup.PowerShell.BeginInvoke() 
#endregion Cleanup Routine

#region Load Functions
Try {
    Get-ChildItem "$ScriptPath\Scripts" -Filter *.ps1 | Select -Expand FullName | ForEach {
        $Function = Split-Path $_ -Leaf
        . $_
    }
} Catch {
    Write-Warning ("{0}: {1}" -f $Function,$_.Exception.Message)
    Continue
}
#endregion Load Functions

#region Format and Type Data
Update-FormatData "$ScriptPath\TypeData\PoshRSJob.Format.ps1xml"
Update-TypeData "$ScriptPath\TypeData\PoshRSJob.Types.ps1xml"
#endregion Format and Type Data

#region Private Functions
Function Increment {
    Set-Variable -Name JobId -Value ($JobId + 1) -Force -Scope Global
    Write-Output $JobId
}

Function GetUsingVariables {
    Param ([scriptblock]$ScriptBlock)
    $ScriptBlock.ast.FindAll({$args[0] -is [System.Management.Automation.Language.UsingExpressionAst]},$True)    
}

Function GetUsingVariableValues {
    Param ([System.Management.Automation.Language.UsingExpressionAst[]]$UsingVar)
    $UsingVar = $UsingVar | Group SubExpression | ForEach {$_.Group | Select -First 1}    
    ForEach ($Var in $UsingVar) {
        Write-Verbose "Using: $($Var)" -Verbose
        Try {
            $Value = Get-Variable -Name $Var.SubExpression.VariablePath.UserPath -ErrorAction Stop
            [pscustomobject]@{
                Name = $Var.SubExpression.Extent.Text
                Value = $Value.Value
                NewName = ('$__using_{0}' -f $Var.SubExpression.VariablePath.UserPath)
                NewVarName = ('__using_{0}' -f $Var.SubExpression.VariablePath.UserPath)
            }
        } Catch {
            Throw "$($Var.SubExpression.Extent.Text) is not a valid Using: variable!"
        }
    }
}

Function ConvertScript {
    Param (
        [scriptblock]$ScriptBlock
    )
    $UsingVariables = @(GetUsingVariables -ScriptBlock $ScriptBlock)
    $List = New-Object 'System.Collections.Generic.List`1[System.Management.Automation.Language.VariableExpressionAst]'
    $Params = New-Object System.Collections.ArrayList
    If ($Script:Add_) {
        [void]$Params.Add('$_')
    }
    If ($UsingVariables) {        
        ForEach ($Ast in $UsingVariables) {
            [void]$list.Add($Ast.SubExpression)
        }
        $UsingVariableData = @(GetUsingVariableValues $UsingVariables)
        [void]$Params.AddRange(@($UsingVariableData.NewName | Select -Unique))
    } 
    $NewParams = $Params -join ', '
    $Tuple=[Tuple]::Create($list,$NewParams)
    $bindingFlags = [Reflection.BindingFlags]"Default,NonPublic,Instance"

    $GetWithInputHandlingForInvokeCommandImpl = ($ScriptBlock.ast.gettype().GetMethod('GetWithInputHandlingForInvokeCommandImpl',$bindingFlags))
    $StringScriptBlock = $GetWithInputHandlingForInvokeCommandImpl.Invoke($ScriptBlock.ast,@($Tuple))
    If ([scriptblock]::Create($StringScriptBlock).ast.endblock[0].statements.extent.text.startswith('$input |')) {
        $StringScriptBlock = $StringScriptBlock -replace '\$Input \|'
    }
    If (-NOT $ScriptBlock.Ast.ParamBlock) {
        $StringScriptBlock = "Param($($NewParams))`n$($StringScriptBlock)"
        [scriptblock]::Create($StringScriptBlock)
    } Else {
        [scriptblock]::Create($StringScriptBlock)
    }
}

Function IsExistingParamBlock {
    Param([scriptblock]$ScriptBlock)
    $errors = [System.Management.Automation.PSParseError[]] @()
    $Tokens = [Management.Automation.PsParser]::Tokenize($ScriptBlock.tostring(), [ref] $errors)       
    $Finding=$True
    For ($i=0;$i -lt $Tokens.count; $i++) {       
        If ($Tokens[$i].Content -eq 'Param' -AND $Tokens[$i].Type -eq 'Keyword') {
            $HasParam = $True
            BREAK
        }
    }
    If ($HasParam) {
        $True
    } Else {
        $False
    }
}

Function GetUsingVariablesV2 {
    Param ([scriptblock]$ScriptBlock)
    $errors = [System.Management.Automation.PSParseError[]] @()
    $Results = [Management.Automation.PsParser]::Tokenize($ScriptBlock.tostring(), [ref] $errors)
    $Results | Where {
        $_.Content -match '^Using:' -AND $_.Type -eq 'Variable'
    }
}

Function GetUsingVariableValuesV2 {
    Param ([System.Management.Automation.PSToken[]]$UsingVar)
    $UsingVar | ForEach {
        $Name = $_.Content -replace 'Using:'
        New-Object PoshRS.PowerShell.V2UsingVariable -Property @{
            Name = $Name
            NewName = '$__using_{0}' -f $Name
            Value = (Get-Variable -Name $Name).Value
            NewVarName = ('__using_{0}') -f $Name
        }
    }
}

Function ConvertScriptBlockV2 {
    Param ([scriptblock]$ScriptBlock)
    $UsingVariables = GetUsingVariablesV2 -ScriptBlock $ScriptBlock
    $UsingVariable = GetUsingVariableValuesV2 -UsingVar $UsingVariables
    $errors = [System.Management.Automation.PSParseError[]] @()
    $Tokens = [Management.Automation.PsParser]::Tokenize($ScriptBlock.tostring(), [ref] $errors)
    $StringBuilder = New-Object System.Text.StringBuilder
    $UsingHash = @{}
    $UsingVariable | ForEach {
        $UsingHash["Using:$($_.Name)"] = $_.NewVarName
    }
    $HasParam = IsExistingParamBlock -ScriptBlock $ScriptBlock
    $Params = New-Object System.Collections.ArrayList
    If ($Script:Add_) {
        [void]$Params.Add('$_')
    }
    If ($UsingVariable) {        
        [void]$Params.AddRange(($UsingVariable | Select -expand NewName))
    } 
    $NewParams = $Params -join ', '  
    If (-Not $HasParam) {
        [void]$StringBuilder.Append("Param($($NewParams))")
    }
    For ($i=0;$i -lt $Tokens.count; $i++){
        #Write-Verbose "Type: $($Tokens[$i].Type)"
        #Write-Verbose "Previous Line: $($Previous.StartLine) -- Current Line: $($Tokens[$i].StartLine)"
        If ($Previous.StartLine -eq $Tokens[$i].StartLine) {
            $Space = " " * [int]($Tokens[$i].StartColumn - $Previous.EndColumn)
            [void]$StringBuilder.Append($Space)
        }
        Switch ($Tokens[$i].Type) {
            'NewLine' {[void]$StringBuilder.Append("`n")}
            'Variable' {
                If ($UsingHash[$Tokens[$i].Content]) {
                    [void]$StringBuilder.Append(("`${0}" -f $UsingHash[$Tokens[$i].Content]))
                } Else {
                    [void]$StringBuilder.Append(("`${0}" -f $Tokens[$i].Content))
                }
            }
            'String' {
                [void]$StringBuilder.Append(("`"{0}`"" -f $Tokens[$i].Content))
            }
            'GroupStart' {
                $Script:GroupStart++
                If ($Script:AddUsing -AND $Script:GroupStart -eq 1) {
                    $Script:AddUsing = $False
                    [void]$StringBuilder.Append($Tokens[$i].Content)                    
                    If ($HasParam) {
                        [void]$StringBuilder.Append("$($NewParams),")
                    }
                } Else {
                    [void]$StringBuilder.Append($Tokens[$i].Content)
                }
            }
            'GroupEnd' {
                $Script:GroupStart--
                If ($Script:GroupStart -eq 0) {
                    $Script:Param = $False
                    [void]$StringBuilder.Append($Tokens[$i].Content)
                } Else {
                    [void]$StringBuilder.Append($Tokens[$i].Content)
                }
            }
            'KeyWord' {
                If ($Tokens[$i].Content -eq 'Param') {
                    $Script:Param = $True
                    $Script:AddUsing = $True
                    $Script:GroupStart=0
                    [void]$StringBuilder.Append($Tokens[$i].Content)
                } Else {
                    [void]$StringBuilder.Append($Tokens[$i].Content)
                }                
            }
            Default {
                [void]$StringBuilder.Append($Tokens[$i].Content)         
            }
        } 
        $Previous = $Tokens[$i]   
    }
    #$StringBuilder.ToString()
    [scriptblock]::Create($StringBuilder.ToString())
}

Function GetParamVariable {
    [CmdletBinding()]
    param (
        [scriptblock]$ScriptBlock
    )     
    # Tokenize the script
    $tokens = [Management.Automation.PSParser]::Tokenize($ScriptBlock, [ref]$null) | Where {
        $_.Type -ne 'NewLine'
    }

    # First Pass - Grab all tokens between the first param block.
    $paramsearch = $false
    $groupstart = 0
    $groupend = 0
    for ($i = 0; $i -lt $tokens.Count; $i++) {
        if (!$paramsearch) {
            if ($tokens[$i].Content -eq "param" ) {
                $paramsearch = $true
            }
        }
        if ($paramsearch) {
            if (($tokens[$i].Type -eq "GroupStart") -and ($tokens[$i].Content -eq '(') ) {
                $groupstart++
            }
            if (($tokens[$i].Type -eq "GroupEnd") -and ($tokens[$i].Content -eq ')') ) {
                $groupend++
            }
            if (($groupstart -ge 1) -and ($groupstart -eq $groupend)) {
                $paramsearch = $false
            }
            if (($tokens[$i].Type -eq 'Variable') -and ($tokens[($i-1)].Content -ne '=')) {
                if ((($groupstart - $groupend) -eq 1)) {
                    "$($tokens[$i].Content)"
                }
            }
        }
    }
}
#endregion Private Functions

#region Aliases
New-Alias -Name ssj -Value Start-RSJob
New-Alias -Name gsj -Value Get-RSJob
New-Alias -Name rsj -Value Receive-RSJob
New-Alias -Name rmsj -Value Remove-RSJob
New-Alias -Name spsj -Value Stop-RSJob
New-Alias -Name wsj -Value Wait-RSJob
#endregion Aliases

#region Handle Module Removal
$ExecutionContext.SessionState.Module.OnRemove ={
    $jobCleanup.Flag=$False
    $RunspacePoolCleanup.Flag=$False
    #Let sit for a second to make sure it has had time to stop
    Start-Sleep -Seconds 1
    $jobCleanup.PowerShell.EndInvoke($jobCleanup.Handle)
    $jobCleanup.PowerShell.Dispose()    
    $RunspacePoolCleanup.PowerShell.EndInvoke($RunspacePoolCleanup.Handle)
    $RunspacePoolCleanup.PowerShell.Dispose()
    Remove-Variable JobId -Scope Global -Force
    Remove-Variable Jobs -Scope Global -Force
    Remove-Variable JobCleanup -Scope Global -Force
    Remove-Variable RunspacePoolCleanup -Scope Global -Force
    Remove-Variable RunspacePools -Scope Global -Force
}
#endregion Handle Module Removal

Export-ModuleMember -Alias * -Function '*-RSJob'