Worker-safety check: fixes from review of the default-on change

ParallelKit.Harness.ps1

@@ -664,9 +664,52 @@ $f = Test-PkWorkerSafety -WorkerScript { param($Id) $add = { param($k, $old) $ol
 Test-Assert ($f.Count -eq 1 -and $f[0].Rule -eq 'DelegateCast') "a variable holding a .GetNewClosure() block, passed to a .NET method, is flagged"
 $f = Test-PkWorkerSafety -WorkerScript { param($Id) $sb = { $_ * 2 }; $n = 'x'; $x = @(1, 2).ForEach($sb); $null = $n.Contains($n) }
 Test-Assert ($f.Count -eq 0) "a scriptblock variable given to .ForEach(), and other variables given to .NET methods, are not flagged"
-$f = Test-PkWorkerSafety -WorkerScript { param($Id) $list = [System.Collections.Generic.List[scriptblock]]::new(); $list.Add({ 1 })  # pk:ignore
+$f = Test-PkWorkerSafety -WorkerScript { param($Id) $null = [System.Threading.Tasks.Task]::Run({ 1 })  # pk:ignore
 }
-Test-Assert ($f.Count -eq 0) "'# pk:ignore' accepts a method that really does take a scriptblock"
+Test-Assert ($f.Count -eq 0) "'# pk:ignore' accepts a finding on its line"
+$f = Test-PkWorkerSafety -WorkerScript { param($Id) $h = { 1 }; $m = [System.Collections.Generic.Dictionary[string,scriptblock]]::new(); $m.Add('a', $h); $l = [System.Collections.Generic.List[scriptblock]]::new(); $l.Add({ 2 }); $q = [System.Collections.Generic.Queue[scriptblock]]::new(); $q.Enqueue($h) }
+Test-Assert ($f.Count -eq 0) ("storing a scriptblock in a collection (.Add/.Enqueue) is not flagged (got " + $f.Count + ")")
+$f = Test-PkWorkerSafety -WorkerScript { param($Id) $h = { 1 }; $d = [System.Collections.Concurrent.ConcurrentDictionary[string,object]]::new(); $null = $d.GetOrAdd('a', $h) }
+Test-Assert ($f.Count -eq 1 -and $f[0].Rule -eq 'DelegateCast') "a scriptblock variable passed as a GetOrAdd factory is still flagged"
+
+# MissingFunction through an alias, and for a module a worker can't load.
+Set-Alias -Name hdbl -Value Get-HarnessDouble -Scope Global
+$f = Test-PkWorkerSafety -WorkerScript { param($Id) hdbl -Value $Id }
+Test-Assert ($f.Count -eq 1 -and $f[0].Rule -eq 'MissingFunction') "an alias to an unshared function of yours is flagged"
+Remove-Item -LiteralPath Alias:\hdbl
+$modName = "pk-temp-" + [guid]::NewGuid().ToString("N")
+$modPath = Join-Path $env:TEMP ($modName + ".psm1")
+[System.IO.File]::WriteAllText($modPath, 'function Get-PkTempModuleThing { 1 }')
+Import-Module -Name $modPath -Global
+$f = Test-PkWorkerSafety -WorkerScript { param($Id) Get-PkTempModuleThing }
+Test-Assert ($f.Count -eq 1 -and $f[0].Rule -eq 'MissingFunction' -and $f[0].Text -like '*PSModulePath*') "a function from a module imported by path is flagged"
+$f = Test-PkWorkerSafety -WorkerScript { param($Id) Get-PkTempModuleThing } -SharedFunctionNames @('Get-PkTempModuleThing')
+Test-Assert ($f.Count -eq 0) "and listing it in -SharedFunctionNames clears it"
+Remove-Module -Name $modName
+[System.IO.File]::Delete($modPath)
+
+# Job-value scan: nested values, typed arrays, and live shared collections.
+$mainThreadBlock = { 1 }
+Test-Assert ((Get-PkScriptBlockJobKey -Job @{ Steps = @(@{ Name = 'a'; Run = $mainThreadBlock }) }) -eq 'Steps') "a scriptblock two levels down in a job value is found"
+Test-Assert ((Get-PkScriptBlockJobKey -Job @{ Opts = [pscustomobject]@{ OnDone = $mainThreadBlock } }) -eq 'Opts') "a scriptblock in a [pscustomobject] property is found"
+$bigBytes = [byte[]]::new(20MB)
+$sw = [System.Diagnostics.Stopwatch]::StartNew()
+$key = Get-PkScriptBlockJobKey -Job @{ Data = $bigBytes; Paths = [string[]]::new(200000) }
+$sw.Stop()
+Test-Assert ($null -eq $key -and $sw.Elapsed.TotalSeconds -lt 1) ("a 20 MB byte[] and a 200k string[] are skipped, not walked (" + [int]$sw.Elapsed.TotalMilliseconds + "ms)")
+$sink = [System.Collections.ArrayList]::Synchronized([System.Collections.ArrayList]::new())
+for ($i = 0; $i -lt 20000; $i++) { $null = $sink.Add([pscustomobject]@{ N = $i }) }
+$writer = [PowerShell]::Create()
+$null = $writer.AddScript('param($s) while ($true) { $null = $s.Add(1); $s.RemoveAt(0) }').AddArgument($sink)
+$writerHandle = $writer.BeginInvoke()
+$scanError = $null
+$sw = [System.Diagnostics.Stopwatch]::StartNew()
+while ($sw.Elapsed.TotalSeconds -lt 3 -and $null -eq $scanError) {
+    try { $null = Get-PkScriptBlockJobKey -Job @{ Sink = $sink } } catch { $scanError = $_ }
+}
+$null = $writer.BeginStop($null, $null)
+$writer.Dispose()
+Test-Assert ($null -eq $scanError) ("scanning a collection other threads are changing doesn't throw (" + $scanError + ")")
 
 # With no Initialize-ParallelKit at all, starting a process sets up the
 # thread pool and the kill-on-close job object. Needs a fresh process.

ParallelKit.ps1

@@ -1582,39 +1582,84 @@ function Invoke-PkRunspaceJobs {
 # ==============================================================================
 
 function Get-PkCommandKind {
+    # What a worker would need to run $Name: 'Function' (one of your session's
+    # functions - must be shared), 'UnloadableModuleFunction' (from a module a
+    # worker can't autoload), or 'Other' (cmdlet, program, autoloadable module
+    # function, or unknown). Reads the Alias: and Function: drives directly:
+    # Get-Command on a name that doesn't resolve searches every module on
+    # $env:PSModulePath, and the check runs on every Invoke-PkRunspaceJobs call.
     param([string]$Name)
-    $cmd = Get-Command -Name $Name -ErrorAction SilentlyContinue | Select-Object -First 1
     $hops = 0
-    while ($null -ne $cmd -and $cmd.CommandType -eq 'Alias' -and $hops -lt 5) {
-        $cmd = $cmd.ResolvedCommand
+    $alias = Get-Item -LiteralPath ("Alias:\" + $Name) -ErrorAction SilentlyContinue
+    while ($null -ne $alias -and $hops -lt 5) {
+        $Name = $alias.Definition
+        $alias = Get-Item -LiteralPath ("Alias:\" + $Name) -ErrorAction SilentlyContinue
         $hops++
     }
-    if ($null -eq $cmd) { return 'Unknown' }
-    # A function from a module (Get-FileHash, say) autoloads in a worker
-    # like a cmdlet does; only functions defined in your session need sharing.
-    if ($cmd.CommandType -eq 'Function' -and -not [string]::IsNullOrEmpty($cmd.ModuleName)) { return 'ModuleFunction' }
-    return [string]$cmd.CommandType
+    $fn = Get-Item -LiteralPath ("Function:\" + $Name) -ErrorAction SilentlyContinue
+    if ($null -eq $fn) { return 'Other' }
+    if ([string]::IsNullOrEmpty($fn.ModuleName)) { return 'Function' }
+    # A module function (Get-FileHash, say) autoloads in a worker like a
+    # cmdlet does - but only if the module can be found by name. One imported
+    # by path from elsewhere can't.
+    if ($null -ne (Get-Module -ListAvailable -Name $fn.ModuleName -ErrorAction SilentlyContinue | Select-Object -First 1)) { return 'Other' }
+    return 'UnloadableModuleFunction'
 }
 
 function Get-PkScriptBlockJobKey {
     # The first key of a job hashtable whose value is a scriptblock, or holds
-    # one (an array/list item or a hashtable value - one level down), else
-    # $null. A scriptblock belongs to the runspace that made it: run in a
-    # worker it executes against the main thread's session state, which can
-    # corrupt it (the kit's own functions vanish mid-run), crash, or hang.
+    # one up to three levels down (in arrays, lists, dictionaries and
+    # [pscustomobject]s), else $null. A scriptblock belongs to the runspace
+    # that made it: run in a worker it executes against the main thread's
+    # session state, which can corrupt it (the kit's own functions vanish
+    # mid-run), crash, or hang.
     param([Parameter(Mandatory)][hashtable]$Job)
     foreach ($key in $Job.Keys) {
-        $value = $Job[$key]
-        if ($value -is [scriptblock]) { return [string]$key }
-        if ($value -is [System.Collections.IDictionary]) {
-            foreach ($inner in $value.Values) { if ($inner -is [scriptblock]) { return [string]$key } }
-        } elseif ($value -is [System.Collections.IList]) {
-            foreach ($inner in $value) { if ($inner -is [scriptblock]) { return [string]$key } }
-        }
+        if (Test-PkHoldsScriptBlock -Value $Job[$key] -Depth 3) { return [string]$key }
     }
     return $null
 }
 
+function Test-PkHoldsScriptBlock {
+    # Get-PkScriptBlockJobKey's walk. Leaves (strings, numbers) are checked
+    # inline rather than through a call, so a big object[] stays cheap, and
+    # a collection whose element type can't hold anything (byte[], string[],
+    # List[int]) isn't walked at all.
+    param($Value, [int]$Depth)
+    if ($null -eq $Value) { return $false }
+    if ($Value -is [scriptblock]) { return $true }
+    if ($Depth -le 0) { return $false }
+    $type = $Value.GetType()
+    if ($Value -is [System.Collections.IDictionary]) {
+        $items = $Value.Values
+    } elseif ($Value -is [System.Collections.IList]) {
+        $items = $Value
+    } elseif ($type.FullName -eq 'System.Management.Automation.PSCustomObject') {
+        $items = foreach ($p in $Value.PSObject.Properties) { if ($p -is [System.Management.Automation.PSNoteProperty]) { $p.Value } }
+    } else {
+        return $false
+    }
+    $elementType = $null
+    if ($type.IsArray) { $elementType = $type.GetElementType() }
+    elseif ($type.IsGenericType) { $genericArgs = $type.GetGenericArguments(); $elementType = $genericArgs[$genericArgs.Length - 1] }
+    if ($null -ne $elementType -and ($elementType.IsValueType -or $elementType -eq [string])) { return $false }
+    try {
+        foreach ($inner in $items) {
+            if ($null -eq $inner -or $inner -is [string] -or $inner -is [System.ValueType]) { continue }
+            if ($inner -is [scriptblock]) { return $true }
+            if ($Depth -gt 1 -and (Test-PkHoldsScriptBlock -Value $inner -Depth ($Depth - 1))) { return $true }
+        }
+    } catch {
+        # A collection other threads are changing (a results sink shared by
+        # running workers) can't be enumerated safely; PowerShell may wrap
+        # the "collection was modified" error, so catch it whatever its type.
+        # Skip the value rather than fail the run: such sinks hold results,
+        # not code.
+        return $false
+    }
+    return $false
+}
+
 function Get-PkScriptBlockJobMessage {
     # Shared by the fixed-list and -NextJobProvider checks in Invoke-PkRunspaceJobs.
     param([string]$Job, [string]$Key)
@@ -1673,6 +1718,9 @@ function Test-PkWorkerSafety {
             if ($kind -eq 'Function' -and -not $listed.Contains($cmdName)) {
                 $findings.Add(@{ Rule = 'MissingFunction'; Where = $t.Where; Line = $c.Extent.StartLineNumber
                                  Text = ($t.Where + " (line " + $c.Extent.StartLineNumber + ") calls '" + $cmdName + "', which is not in -SharedFunctionNames") })
+            } elseif ($kind -eq 'UnloadableModuleFunction' -and -not $listed.Contains($cmdName)) {
+                $findings.Add(@{ Rule = 'MissingFunction'; Where = $t.Where; Line = $c.Extent.StartLineNumber
+                                 Text = ($t.Where + " (line " + $c.Extent.StartLineNumber + ") calls '" + $cmdName + "', from a module a worker can't load by name (imported by path?) - put the module on `$env:PSModulePath, or list the function in -SharedFunctionNames") })
             }
         }
 
@@ -1714,6 +1762,11 @@ function Test-PkWorkerSafety {
         # { }.GetNewClosure()) anywhere in this code, then passed as
         # $lazy = [System.Lazy[object]]::new($factory). By name, not by flow:
         # a variable reassigned to something else in between is still flagged.
+        # Methods that only store a value in a collection ($list.Add($sb),
+        # $map.Add($k, $sb), $queue.Enqueue($sb)) aren't flagged: storing a
+        # scriptblock makes no delegate. Not caught as a result: a collection
+        # typed to hold delegates (List[Func[object]]), where .Add converts.
+        $storageMethods = @('Add', 'AddRange', 'Insert', 'Enqueue', 'Push', 'TryAdd', 'Contains', 'Remove', 'IndexOf', 'set_Item')
         $scriptBlockVars = [System.Collections.Generic.HashSet[string]]::new([System.StringComparer]::OrdinalIgnoreCase)
         foreach ($assign in $ast.FindAll({ param($node) $node -is [System.Management.Automation.Language.AssignmentStatementAst] }, $true)) {
             $target = $assign.Left
@@ -1729,6 +1782,7 @@ function Test-PkWorkerSafety {
             if ($null -eq $call.Arguments) { continue }
             $memberName = [string]$call.Member.Extent.Text
             if ($memberName -in @('ForEach', 'Where')) { continue }
+            if ($memberName -in $storageMethods) { continue }
             $passed = $null
             foreach ($argument in $call.Arguments) {
                 if ($argument -is [System.Management.Automation.Language.ScriptBlockExpressionAst]) { $passed = 'a scriptblock' }

README.md

@@ -166,7 +166,7 @@ Each worker runs in its own runspace: a separate PowerShell session inside the s
 - Each job hashtable's keys become the worker's parameters. The worker's `param()` block must declare every key. A key it doesn't declare is an error; with a fixed `-JobParams` list it's caught before any worker starts.
 - If the worker declares a `$CancelToken` parameter, it gets the run's cancellation token.
 - `-JobParams` takes a list or array of job hashtables. A single hashtable works too, so a list that happens to hold one job is fine.
-- A worker runspace starts empty. Name every one of your functions it calls in `-SharedFunctionNames`. The kit's own worker functions are always there.
+- A worker runspace starts empty. Name every one of your functions it calls in `-SharedFunctionNames`. The kit's own worker functions are always there, and cmdlets and functions from installed modules load on their own. A module you imported by path isn't installed, so list its functions too (or put it on `$env:PSModulePath`).
 - Whatever the worker returns is handed to `-ResultHandler` on the main thread, once per object returned.
 - `-NextJobProvider { ... }` can replace `-JobParams` to produce jobs lazily: it's called on the main thread for the next job hashtable, and returns `$null` when there are no more.
 - `-PollCallback { ... }` runs on the main thread on every loop pass, for progress display and the like.
@@ -185,7 +185,7 @@ Worker code is the worker script, every function in `-SharedFunctionNames`, and
 
 Everything else is fine in a worker, cmdlets included. Callbacks that run on the main thread (`-ResultHandler`, `-PollCallback`, `-NextJobProvider`, `-OnComplete`, the message handler) are not worker code and can use anything.
 
-`Invoke-PkRunspaceJobs` checks both rules before it starts anything, and refuses to launch code that breaks them, naming the line or the job key. Pass `-CheckWorkerSafety:$false` to turn that off. The code check is also available on its own as `Test-PkWorkerSafety -WorkerScript $w -SharedFunctionNames $names`, which returns a list of findings (including functions the worker calls but you didn't share), for use in your own tests. For rule 2 it flags written casts, and scriptblocks passed into a .NET method or constructor, either written in the call or through a variable they were assigned to. `# pk:ignore` on a line accepts a finding. It can't see commands run through a variable (`& $cmd`), or a scriptblock buried more than one level deep in a job value.
+`Invoke-PkRunspaceJobs` checks both rules before it starts anything, and refuses to launch code that breaks them, naming the line or the job key. Pass `-CheckWorkerSafety:$false` to turn that off. The code check is also available on its own as `Test-PkWorkerSafety -WorkerScript $w -SharedFunctionNames $names`, which returns a list of findings (including functions the worker calls but you didn't share), for use in your own tests. For rule 2 it flags written casts, and scriptblocks passed into a .NET method or constructor, either written in the call or through a variable they were assigned to. `# pk:ignore` on a line accepts a finding. Storing a scriptblock in a collection (`$list.Add($sb)`, `$queue.Enqueue($sb)`) isn't flagged, since that makes no delegate. It can't see commands run through a variable (`& $cmd`), a collection typed to hold delegates (`List[Func[object]]`, whose `.Add` converts), or a scriptblock buried more than three levels deep in a job value.
 
 ### Why these rules exist