more refactoring. getting better.

Author: santisq

src/PSParallelPipeline/Commands/InvokeParallelCommand.cs

@@ -66,13 +66,13 @@ protected override void BeginProcessing()
             InitialSessionState = iss
         };
 
-        WorkerSettings workerSettings = new()
+        TaskSettings workerSettings = new()
         {
-            Token = _cts.Token,
+            Script = ScriptBlock.ToString(),
             UsingStatements = ScriptBlock.GetUsingParameters(this)
         };
 
-        _worker = new Worker(poolSettings, workerSettings);
+        _worker = new Worker(poolSettings, workerSettings, _cts.Token);
         _worker.Run();
     }
 
@@ -83,7 +83,7 @@ protected override void ProcessRecord()
 
         try
         {
-            _worker.Enqueue(InputObject, ScriptBlock);
+            _worker.Enqueue(InputObject);
             while (_worker.TryTake(out PSOutputData data))
             {
                 ProcessOutput(data);

src/PSParallelPipeline/PSTask.cs

@@ -30,10 +30,17 @@ private PSTask(RunspacePool pool)
         _pool = pool;
     }
 
-    static internal PSTask Create(RunspacePool runspacePool)
+    static internal async Task<PSTask> CreateAsync(
+        object? input,
+        RunspacePool runspacePool,
+        TaskSettings settings)
     {
         PSTask ps = new(runspacePool);
         HookStreams(ps._internalStreams, runspacePool.Streams);
+        ps.Runspace = await runspacePool.GetRunspaceAsync();
+        ps.AddInput(input);
+        ps.AddScript(settings.Script);
+        ps.AddUsingStatements(settings.UsingStatements);
         return ps;
     }
 
@@ -56,7 +63,7 @@ private static Task InvokePowerShellAsync(
             powerShell.BeginInvoke<PSObject, PSObject>(null, output),
             powerShell.EndInvoke);
 
-    internal PSTask AddInput(object? inputObject)
+    private void AddInput(object? inputObject)
     {
         if (inputObject is not null)
         {
@@ -65,17 +72,12 @@ internal PSTask AddInput(object? inputObject)
                 .AddArgument("_")
                 .AddArgument(inputObject);
         }
-
-        return this;
     }
 
-    internal PSTask AddScript(ScriptBlock script)
-    {
-        _powershell.AddScript(script.ToString(), useLocalScope: true);
-        return this;
-    }
+    private void AddScript(string script) =>
+        _powershell.AddScript(script, useLocalScope: true);
 
-    internal PSTask AddUsingStatements(Dictionary<string, object?> usingParams)
+    private void AddUsingStatements(Dictionary<string, object?> usingParams)
     {
         if (usingParams.Count > 0)
         {
@@ -84,18 +86,17 @@ internal PSTask AddUsingStatements(Dictionary<string, object?> usingParams)
                 ["--%"] = usingParams
             });
         }
-
-        return this;
     }
 
     internal async Task InvokeAsync()
     {
         try
         {
-            using CancellationTokenRegistration _ = _pool.RegisterCancellation(_powershell.Dispose);
-            Runspace = await _pool.GetRunspaceAsync();
+            using CancellationTokenRegistration _ = _pool.RegisterCancellation(_powershell.Stop);
             await InvokePowerShellAsync(_powershell, OutputStreams.Success);
         }
+        catch (PipelineStoppedException)
+        { }
         catch (Exception exception)
         {
             OutputStreams.AddOutput(exception.CreateProcessingTaskError(this));

src/PSParallelPipeline/RunspacePool.cs

@@ -1,6 +1,5 @@
 using System;
 using System.Collections.Concurrent;
-using System.Collections.Generic;
 using System.Management.Automation.Runspaces;
 using System.Threading;
 using System.Threading.Tasks;
@@ -15,7 +14,9 @@ internal sealed class RunspacePool : IDisposable
 
     private readonly ConcurrentQueue<Runspace> _pool = [];
 
-    private readonly List<Runspace> _created;
+    // private readonly List<Runspace> _created;
+
+    private readonly ConcurrentDictionary<Guid, Runspace> _created;
 
     private readonly bool _useNew;
 
@@ -33,21 +34,18 @@ internal RunspacePool(
         (MaxRunspaces, _useNew, _iss) = settings;
         Streams = streams;
         Token = token;
-        _created = new List<Runspace>(MaxRunspaces);
         _semaphore = new SemaphoreSlim(MaxRunspaces, MaxRunspaces);
+        _created = new ConcurrentDictionary<Guid, Runspace>(
+            Environment.ProcessorCount,
+            MaxRunspaces);
     }
     internal void PushRunspace(Runspace runspace)
     {
         if (_useNew)
         {
             runspace.Dispose();
-
-            lock (_created)
-            {
-                _created.Remove(runspace);
-                runspace = CreateRunspace();
-                _created.Add(runspace);
-            }
+            _created.TryRemove(runspace.InstanceId, out _);
+            runspace = CreateRunspace();
         }
 
         _pool.Enqueue(runspace);
@@ -60,6 +58,7 @@ internal CancellationTokenRegistration RegisterCancellation(Action callback) =>
     private Runspace CreateRunspace()
     {
         Runspace rs = RunspaceFactory.CreateRunspace(_iss);
+        _created[rs.InstanceId] = rs;
         rs.Open();
         return rs;
     }
@@ -72,18 +71,12 @@ internal async Task<Runspace> GetRunspaceAsync()
             return runspace;
         }
 
-        lock (_created)
-        {
-            runspace = CreateRunspace();
-            _created.Add(runspace);
-        }
-
-        return runspace;
+        return CreateRunspace();
     }
 
     public void Dispose()
     {
-        foreach (Runspace runspace in _created)
+        foreach (Runspace runspace in _created.Values)
         {
             runspace.Dispose();
         }

src/PSParallelPipeline/TaskSettings.cs

@@ -0,0 +1,7 @@
+using System.Collections.Generic;
+
+namespace PSParallelPipeline;
+
+internal record struct TaskSettings(
+    string Script,
+    Dictionary<string, object?> UsingStatements);

src/PSParallelPipeline/Worker.cs

@@ -3,47 +3,39 @@
 using System.Threading.Tasks;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Management.Automation;
 
 namespace PSParallelPipeline;
 
 internal sealed class Worker
 {
     private Task? _worker;
 
-    private readonly BlockingCollection<PSTask> _input = [];
+    private readonly TaskSettings _taskSettings;
+
+    private readonly BlockingCollection<object?> _input = [];
 
     private readonly BlockingCollection<PSOutputData> _output = [];
 
     private readonly RunspacePool _pool;
 
     private readonly CancellationToken _token;
 
-    private readonly Dictionary<string, object?> _usingParams;
-
     private readonly PSOutputStreams _streams;
 
     internal Worker(
         PoolSettings poolSettings,
-        WorkerSettings workerSettings)
+        TaskSettings taskSettings,
+        CancellationToken token)
     {
-        (_usingParams, _token) = workerSettings;
+        _token = token;
+        _taskSettings = taskSettings;
         _streams = new PSOutputStreams(_output);
         _pool = new RunspacePool(poolSettings, _streams, _token);
     }
 
     internal void Wait() => _worker?.GetAwaiter().GetResult();
 
-    internal void Enqueue(object? input, ScriptBlock script)
-    {
-        _input.Add(
-            item: PSTask
-                .Create(_pool)
-                .AddInput(input)
-                .AddScript(script)
-                .AddUsingStatements(_usingParams),
-            cancellationToken: _token);
-    }
+    internal void Enqueue(object? input) => _input.Add(input, _token);
 
     internal bool TryTake(out PSOutputData output) =>
         _output.TryTake(out output, 0, _token);
@@ -68,9 +60,14 @@ private async Task Start()
                     await ProcessAnyAsync(tasks);
                 }
 
-                if (_input.TryTake(out PSTask ps, 0, _token))
+                if (_input.TryTake(out object? input, 0, _token))
                 {
-                    tasks.Add(ps.InvokeAsync());
+                    PSTask task = await PSTask.CreateAsync(
+                        input: input,
+                        runspacePool: _pool,
+                        settings: _taskSettings);
+
+                    tasks.Add(task.InvokeAsync());
                 }
             }
 

src/PSParallelPipeline/WorkerSettings.cs

@@ -92,9 +92,9 @@ Describe PSParallelPipeline {
 
     Context 'UseNewRunspace Parameter' {
         It 'Should reuse runspaces by default' {
-            $runspaces = 0..10 | Invoke-Parallel { [runspace]::DefaultRunspace } |
-                Select-Object -ExpandProperty InstanceId -Unique
-            $runspaces.Count | Should -BeLessOrEqual 5
+            0..10 | Invoke-Parallel { [runspace]::DefaultRunspace } |
+                Select-Object -ExpandProperty InstanceId -Unique |
+                Should -HaveCount 5
         }
 
         It 'Should use a new runspace when the -UseNewRunspace is used' {
@@ -120,7 +120,9 @@ Describe PSParallelPipeline {
             0..10 | Invoke-Parallel { Test-Function $_ } -Functions Test-Function |
                 Sort-Object |
                 Should -BeExactly @(0..10 | ForEach-Object { Test-Function $_ })
+        }
 
+        It 'Should throw if a function could not be found' {
             { Invoke-Parallel -Functions Test-NotExist { } } |
                 Should -Throw -ExceptionType ([CommandNotFoundException])
         }
@@ -346,14 +348,14 @@ Describe PSParallelPipeline {
 
             Assert-RunspaceCount {
                 { 0..1000 | Invoke-Parallel @invokeParallelSplat } |
-                    Should -Throw
+                    Should -Throw -ExceptionType ([TimeoutException])
             } -TestCount 50 # -WaitSeconds 1
 
             Assert-RunspaceCount {
                 $invokeParallelSplat['UseNewRunspace'] = $true
                 $invokeParallelSplat['ThrottleLimit'] = 1001
                 { 0..1000 | Invoke-Parallel @invokeParallelSplat } |
-                    Should -Throw
+                    Should -Throw -ExceptionType ([TimeoutException])
             } -TestCount 50 # -WaitSeconds 1
         }
     }