From bb6d34c4a12f21aa00668f6983a005f1e303d675 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sat, 25 Jul 2026 22:50:20 +0200 Subject: [PATCH 01/27] Harden asynchronous cmdlet pipeline operations --- CHANGELOG.MD | 1 + PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 35 ++++++ PowerForge.Tests/AsyncPSCmdletTests.cs | 147 +++++++++++++++++++++++ 3 files changed, 183 insertions(+) diff --git a/CHANGELOG.MD b/CHANGELOG.MD index e4bcfdc9f..64e1f5a96 100644 --- a/CHANGELOG.MD +++ b/CHANGELOG.MD @@ -2,6 +2,7 @@ ## Unreleased ### What's Changed +- Added pipeline-thread marshalling for terminating errors and `ShouldContinue` prompts raised after an asynchronous cmdlet await. - Fixed module merging so `using` lines inside PowerShell here-strings and block comments remain in place instead of being moved into the generated module header. - Hardened the canonical `AsyncPSCmdlet` lifecycle with request-specific prompt and `ShouldProcess` replies while preserving PowerShell-thread execution until the first await. - Converted Task-backed App Store Connect and managed catalog cmdlets from sync-over-async blocking to cancellable asynchronous lifecycle hooks. diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index e983aa210..fc8ba53b9 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -31,12 +31,14 @@ private enum PipelineType Output, OutputEnumerate, Error, + TerminatingError, Warning, Verbose, Debug, Information, Progress, ShouldProcess, + ShouldContinue, PromptForCredential } @@ -104,6 +106,18 @@ protected override void StopProcessing() return (bool)replyPipe.Take(CancelToken)!; } + /// Thread-safe ShouldContinue bridge for asynchronous cmdlet code. + public new bool ShouldContinue(string query, string caption) + { + ThrowIfStopped(); + if (_currentOutPipe is null || IsPipelineThread) + return base.ShouldContinue(query, caption); + + using var replyPipe = new BlockingCollection(boundedCapacity: 1); + _currentOutPipe.Add(new PipelineItem((query, caption), PipelineType.ShouldContinue, replyPipe), CancelToken); + return (bool)replyPipe.Take(CancelToken)!; + } + /// Thread-safe credential prompt bridge for asynchronous cmdlet code. public PSCredential? PromptForCredential(string caption, string message, string userName, string targetName) { @@ -146,6 +160,20 @@ protected override void StopProcessing() _currentOutPipe.Add(new PipelineItem(errorRecord, PipelineType.Error), CancelToken); } + /// Thread-safe terminating-error bridge for asynchronous cmdlet code. + protected new void ThrowTerminatingError(ErrorRecord errorRecord) + { + ThrowIfStopped(); + if (_currentOutPipe is null || IsPipelineThread) + { + base.ThrowTerminatingError(errorRecord); + return; + } + + _currentOutPipe.Add(new PipelineItem(errorRecord, PipelineType.TerminatingError), CancelToken); + throw new PipelineStoppedException(); + } + /// Thread-safe warning bridge for asynchronous cmdlet code. public new void WriteWarning(string text) { @@ -256,6 +284,9 @@ void PumpItem(PipelineItem item) case PipelineType.Error: base.WriteError((ErrorRecord)item.Value!); break; + case PipelineType.TerminatingError: + base.ThrowTerminatingError((ErrorRecord)item.Value!); + break; case PipelineType.Warning: base.WriteWarning((string)item.Value!); break; @@ -275,6 +306,10 @@ void PumpItem(PipelineItem item) var should = ((string Target, string Action))item.Value!; item.ReplyPipe!.Add(base.ShouldProcess(should.Target, should.Action), CancelToken); break; + case PipelineType.ShouldContinue: + var shouldContinue = ((string Query, string Caption))item.Value!; + item.ReplyPipe!.Add(base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption), CancelToken); + break; case PipelineType.PromptForCredential: var prompt = ((string Caption, string Message, string UserName, string TargetName))item.Value!; item.ReplyPipe!.Add( diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 31169c56b..5edf884fd 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -1,5 +1,9 @@ +using System.Collections.ObjectModel; +using System.Globalization; using System.Management.Automation; +using System.Management.Automation.Host; using System.Management.Automation.Runspaces; +using System.Security; using System.Threading; using System.Threading.Tasks; using PSPublishModule; @@ -95,6 +99,53 @@ public void AsyncPSCmdlet_does_not_capture_a_custom_task_scheduler() var item = Assert.Single(result); Assert.Equal(0, item.BaseObject); } + + [Fact] + public void AsyncPSCmdlet_marshals_terminating_errors_to_the_pipeline_thread() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncTerminatingError", + typeof(TestAsyncTerminatingErrorCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncTerminatingError"); + TestAsyncTerminatingErrorCommand.Reset(); + + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.StartsWith("AsyncTerminatingError,", exception.ErrorRecord.FullyQualifiedErrorId, StringComparison.Ordinal); + Assert.Equal("async terminating failure", exception.InnerException?.Message); + Assert.False(TestAsyncTerminatingErrorCommand.ReachedAfterTermination); + } + + [Theory] + [InlineData(true)] + [InlineData(false)] + public void AsyncPSCmdlet_marshals_ShouldContinue_to_the_pipeline_thread(bool approved) + { + var sessionState = InitialSessionState.Create(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncShouldContinue", + typeof(TestAsyncShouldContinueCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(new ChoiceHost(approved), sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncShouldContinue"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + var item = Assert.Single(result); + Assert.Equal(approved, item.BaseObject); + } } [Cmdlet(VerbsDiagnostic.Test, "AsyncThreadAffinity")] @@ -190,6 +241,102 @@ protected override async Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncTerminatingError")] +public sealed class TestAsyncTerminatingErrorCommand : AsyncPSCmdlet +{ + private static int _reachedAfterTermination; + + public static bool ReachedAfterTermination => Volatile.Read(ref _reachedAfterTermination) != 0; + + public static void Reset() + => Volatile.Write(ref _reachedAfterTermination, 0); + + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + ThrowTerminatingError(new ErrorRecord( + new InvalidOperationException("async terminating failure"), + "AsyncTerminatingError", + ErrorCategory.InvalidOperation, + targetObject: null)); + Volatile.Write(ref _reachedAfterTermination, 1); + WriteObject("unreachable"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncShouldContinue")] +public sealed class TestAsyncShouldContinueCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(ShouldContinue("Proceed?", "Question")); + } +} + +public sealed class ChoiceHost(bool approved) : PSHost +{ + private readonly Guid _id = Guid.NewGuid(); + private readonly ChoiceHostUserInterface _ui = new(approved); + + public override Guid InstanceId => _id; + public override string Name => nameof(ChoiceHost); + public override Version Version => new(1, 0); + public override PSHostUserInterface UI => _ui; + public override CultureInfo CurrentCulture => CultureInfo.InvariantCulture; + public override CultureInfo CurrentUICulture => CultureInfo.InvariantCulture; + public override void EnterNestedPrompt() { } + public override void ExitNestedPrompt() { } + public override void NotifyBeginApplication() { } + public override void NotifyEndApplication() { } + public override void SetShouldExit(int exitCode) { } +} + +public sealed class ChoiceHostUserInterface(bool approved) : PSHostUserInterface +{ + public override PSHostRawUserInterface RawUI => null!; + + public override int PromptForChoice( + string caption, + string message, + Collection choices, + int defaultChoice) + => approved ? 0 : 1; + + public override string ReadLine() => string.Empty; + public override SecureString ReadLineAsSecureString() => new(); + public override void Write(string value) { } + public override void Write(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value) { } + public override void WriteLine(string value) { } + public override void WriteErrorLine(string value) { } + public override void WriteDebugLine(string message) { } + public override void WriteProgress(long sourceId, ProgressRecord record) { } + public override void WriteVerboseLine(string message) { } + public override void WriteWarningLine(string message) { } + + public override PSCredential PromptForCredential( + string caption, + string message, + string userName, + string targetName) + => new(userName, new SecureString()); + + public override PSCredential PromptForCredential( + string caption, + string message, + string userName, + string targetName, + PSCredentialTypes allowedCredentialTypes, + PSCredentialUIOptions options) + => PromptForCredential(caption, message, userName, targetName); + + public override Dictionary Prompt( + string caption, + string message, + Collection descriptions) + => []; +} + public sealed class ForwardingTaskScheduler : TaskScheduler { private int _executeNextTaskInline; From 2ddcac00580a7411543f61001842b16b5d266c26 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sat, 25 Jul 2026 23:06:15 +0200 Subject: [PATCH 02/27] Preserve derived async cmdlet cleanup --- CHANGELOG.MD | 2 +- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.MD b/CHANGELOG.MD index 64e1f5a96..c46c5b4d1 100644 --- a/CHANGELOG.MD +++ b/CHANGELOG.MD @@ -2,7 +2,7 @@ ## Unreleased ### What's Changed -- Added pipeline-thread marshalling for terminating errors and `ShouldContinue` prompts raised after an asynchronous cmdlet await. +- Added pipeline-thread marshalling for terminating errors and `ShouldContinue` prompts raised after an asynchronous cmdlet await, while preserving derived cmdlet resource cleanup through virtual disposal. - Fixed module merging so `using` lines inside PowerShell here-strings and block comments remain in place instead of being moved into the generated module header. - Hardened the canonical `AsyncPSCmdlet` lifecycle with request-specific prompt and `ShouldProcess` replies while preserving PowerShell-thread execution until the first await. - Converted Task-backed App Store Connect and managed catalog cmdlets from sync-over-async blocking to cancellable asynchronous lifecycle hooks. diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index fc8ba53b9..a14d87894 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -247,7 +247,7 @@ internal void ThrowIfStopped() } /// - public void Dispose() + public virtual void Dispose() => _cancelSource.Dispose(); private bool IsPipelineThread From 15b53be2529635adea80f077f8f432d962b364e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sat, 25 Jul 2026 23:12:28 +0200 Subject: [PATCH 03/27] Preserve derived async cancellation checks --- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index a14d87894..7f320c890 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -240,7 +240,7 @@ protected override void StopProcessing() } /// Throws when PowerShell has requested cancellation. - internal void ThrowIfStopped() + protected internal void ThrowIfStopped() { if (_cancelSource.IsCancellationRequested) throw new PipelineStoppedException(); From d41075fa3f4404ae0243daf8d8d891ff6edb3273 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sat, 25 Jul 2026 23:34:21 +0200 Subject: [PATCH 04/27] Normalize async cancellation as pipeline stops --- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 24 ++++++++++-- PowerForge.Tests/AsyncPSCmdletTests.cs | 47 ++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 4 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 7f320c890..4228e464e 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -253,6 +253,18 @@ public virtual void Dispose() private bool IsPipelineThread => _pipelineThreadId != 0 && Environment.CurrentManagedThreadId == _pipelineThreadId; + private static void GetBlockTaskResult(Task blockTask) + { + try + { + blockTask.GetAwaiter().GetResult(); + } + catch (OperationCanceledException) + { + throw new PipelineStoppedException(); + } + } + private void RunBlockInAsync(Func task) { using var outPipe = new BlockingCollection(); @@ -356,7 +368,7 @@ void PumpQueuedItems() ClearPipes(); } - blockTask.GetAwaiter().GetResult(); + GetBlockTaskResult(blockTask); return; } @@ -373,7 +385,7 @@ void PumpQueuedItems() PumpItem(item); } } - catch + catch (Exception pipelineException) { _cancelSource.Cancel(); CompleteAddingIfNeeded(outPipe); @@ -381,13 +393,17 @@ void PumpQueuedItems() { blockTask.GetAwaiter().GetResult(); } - catch (Exception ex) when (ex is OperationCanceledException or PipelineStoppedException) + catch (Exception completionException) when (completionException is OperationCanceledException or PipelineStoppedException) { + _ = completionException; } + if (pipelineException is OperationCanceledException) + throw new PipelineStoppedException(); + throw; } - blockTask.GetAwaiter().GetResult(); + GetBlockTaskResult(blockTask); } } diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 5edf884fd..eff71ddad 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -100,6 +100,33 @@ public void AsyncPSCmdlet_does_not_capture_a_custom_task_scheduler() Assert.Equal(0, item.BaseObject); } + [Fact] + public void AsyncPSCmdlet_translates_cancellation_to_a_pipeline_stop() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCancellation", + typeof(TestAsyncCancellationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCancellation"); + TestAsyncCancellationCommand.Reset(); + + var invocation = powerShell.BeginInvoke(); + Assert.True( + TestAsyncCancellationCommand.Started.Wait(TimeSpan.FromSeconds(5)), + "The asynchronous cmdlet did not start in time."); + + powerShell.Stop(); + + var exception = Assert.Throws(() => powerShell.EndInvoke(invocation)); + Assert.IsNotType(exception); + } + [Fact] public void AsyncPSCmdlet_marshals_terminating_errors_to_the_pipeline_thread() { @@ -241,6 +268,26 @@ protected override async Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncCancellation")] +public sealed class TestAsyncCancellationCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _started = new(); + + public static ManualResetEventSlim Started => _started; + + public static void Reset() + { + _started.Dispose(); + _started = new ManualResetEventSlim(); + } + + protected override async Task ProcessRecordAsync() + { + _started.Set(); + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncTerminatingError")] public sealed class TestAsyncTerminatingErrorCommand : AsyncPSCmdlet { From 0771013975f5f45588774cdc60878180a2eaec54 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 00:03:17 +0200 Subject: [PATCH 05/27] Harden async pipeline lifecycle contracts --- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 317 ++++++++++++++++++----- PowerForge.Tests/AsyncPSCmdletTests.cs | 102 ++++++++ 2 files changed, 361 insertions(+), 58 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 4228e464e..e2caaf7e1 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -37,14 +37,27 @@ private enum PipelineType Debug, Information, Progress, + ShouldProcessTarget, ShouldProcess, + ShouldProcessVerbose, + ShouldProcessReason, ShouldContinue, + ShouldContinueAll, + ShouldContinueSecurity, PromptForCredential } + private sealed class PipelineReply + { + public PipelineReply(object? value) + => Value = value; + + public object? Value { get; } + } + private sealed class PipelineItem { - public PipelineItem(object? value, PipelineType type, BlockingCollection? replyPipe = null) + public PipelineItem(object? value, PipelineType type, BlockingCollection? replyPipe = null) { Value = value; Type = type; @@ -55,13 +68,14 @@ public PipelineItem(object? value, PipelineType type, BlockingCollection? ReplyPipe { get; } + public BlockingCollection? ReplyPipe { get; } } private readonly CancellationTokenSource _cancelSource = new(); private static readonly SynchronizationContext HookSynchronizationContext = new AsyncHookSynchronizationContext(); private BlockingCollection? _currentOutPipe; private int _pipelineThreadId; + private int _disposed; /// Cancellation token triggered when PowerShell stops the cmdlet. protected internal CancellationToken CancelToken => _cancelSource.Token; @@ -94,40 +108,103 @@ protected virtual Task EndProcessingAsync() protected override void StopProcessing() => _cancelSource.Cancel(); + /// Thread-safe ShouldProcess bridge for asynchronous cmdlet code. + public new bool ShouldProcess(string? target) + { + if (IsPipelineThread) + return base.ShouldProcess(target ?? string.Empty); + + return (bool)RequestPipelineReply(target ?? string.Empty, PipelineType.ShouldProcessTarget)!; + } + /// Thread-safe ShouldProcess bridge for asynchronous cmdlet code. public new bool ShouldProcess(string? target, string action) { - ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) return base.ShouldProcess(target ?? string.Empty, action); - using var replyPipe = new BlockingCollection(boundedCapacity: 1); - _currentOutPipe.Add(new PipelineItem((target ?? string.Empty, action), PipelineType.ShouldProcess, replyPipe), CancelToken); - return (bool)replyPipe.Take(CancelToken)!; + return (bool)RequestPipelineReply((target ?? string.Empty, action), PipelineType.ShouldProcess)!; + } + + /// Thread-safe ShouldProcess bridge for asynchronous cmdlet code. + public new bool ShouldProcess(string verboseDescription, string verboseWarning, string caption) + { + if (IsPipelineThread) + return base.ShouldProcess(verboseDescription, verboseWarning, caption); + + return (bool)RequestPipelineReply( + (verboseDescription, verboseWarning, caption), + PipelineType.ShouldProcessVerbose)!; + } + + /// Thread-safe ShouldProcess bridge for asynchronous cmdlet code. + public new bool ShouldProcess( + string verboseDescription, + string verboseWarning, + string caption, + out ShouldProcessReason shouldProcessReason) + { + if (IsPipelineThread) + return base.ShouldProcess(verboseDescription, verboseWarning, caption, out shouldProcessReason); + + var reply = ((bool Result, ShouldProcessReason Reason))RequestPipelineReply( + (verboseDescription, verboseWarning, caption), + PipelineType.ShouldProcessReason)!; + shouldProcessReason = reply.Reason; + return reply.Result; } /// Thread-safe ShouldContinue bridge for asynchronous cmdlet code. public new bool ShouldContinue(string query, string caption) { - ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) return base.ShouldContinue(query, caption); - using var replyPipe = new BlockingCollection(boundedCapacity: 1); - _currentOutPipe.Add(new PipelineItem((query, caption), PipelineType.ShouldContinue, replyPipe), CancelToken); - return (bool)replyPipe.Take(CancelToken)!; + return (bool)RequestPipelineReply((query, caption), PipelineType.ShouldContinue)!; + } + + /// Thread-safe ShouldContinue bridge for asynchronous cmdlet code. + public new bool ShouldContinue(string query, string caption, ref bool yesToAll, ref bool noToAll) + { + if (IsPipelineThread) + return base.ShouldContinue(query, caption, ref yesToAll, ref noToAll); + + var reply = ((bool Result, bool YesToAll, bool NoToAll))RequestPipelineReply( + (query, caption, yesToAll, noToAll), + PipelineType.ShouldContinueAll)!; + yesToAll = reply.YesToAll; + noToAll = reply.NoToAll; + return reply.Result; + } + + /// Thread-safe ShouldContinue bridge for asynchronous cmdlet code. + public new bool ShouldContinue( + string query, + string caption, + bool hasSecurityImpact, + ref bool yesToAll, + ref bool noToAll) + { + if (IsPipelineThread) + return base.ShouldContinue(query, caption, hasSecurityImpact, ref yesToAll, ref noToAll); + + var reply = ((bool Result, bool YesToAll, bool NoToAll))RequestPipelineReply( + (query, caption, hasSecurityImpact, yesToAll, noToAll), + PipelineType.ShouldContinueSecurity)!; + yesToAll = reply.YesToAll; + noToAll = reply.NoToAll; + return reply.Result; } /// Thread-safe credential prompt bridge for asynchronous cmdlet code. public PSCredential? PromptForCredential(string caption, string message, string userName, string targetName) { - ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) return Host.UI.PromptForCredential(caption, message, userName, targetName); - using var replyPipe = new BlockingCollection(boundedCapacity: 1); - _currentOutPipe.Add(new PipelineItem((caption, message, userName, targetName), PipelineType.PromptForCredential, replyPipe), CancelToken); - return (PSCredential?)replyPipe.Take(CancelToken); + return (PSCredential?)RequestPipelineReply( + (caption, message, userName, targetName), + PipelineType.PromptForCredential); } /// Thread-safe output bridge for asynchronous cmdlet code. @@ -137,106 +214,129 @@ protected override void StopProcessing() /// Thread-safe output bridge for asynchronous cmdlet code. public new void WriteObject(object? sendToPipeline, bool enumerateCollection) { + if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) { base.WriteObject(sendToPipeline, enumerateCollection); return; } - _currentOutPipe.Add(new PipelineItem(sendToPipeline, enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output), CancelToken); + _ = TryQueue(new PipelineItem( + sendToPipeline, + enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output)); } /// Thread-safe error bridge for asynchronous cmdlet code. public new void WriteError(ErrorRecord errorRecord) { + if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) { base.WriteError(errorRecord); return; } - _currentOutPipe.Add(new PipelineItem(errorRecord, PipelineType.Error), CancelToken); + _ = TryQueue(new PipelineItem(errorRecord, PipelineType.Error)); } /// Thread-safe terminating-error bridge for asynchronous cmdlet code. - protected new void ThrowTerminatingError(ErrorRecord errorRecord) + public new void ThrowTerminatingError(ErrorRecord errorRecord) { ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) { base.ThrowTerminatingError(errorRecord); return; } - _currentOutPipe.Add(new PipelineItem(errorRecord, PipelineType.TerminatingError), CancelToken); + _ = TryQueue(new PipelineItem(errorRecord, PipelineType.TerminatingError)); throw new PipelineStoppedException(); } /// Thread-safe warning bridge for asynchronous cmdlet code. public new void WriteWarning(string text) { + if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) { base.WriteWarning(text); return; } - _currentOutPipe.Add(new PipelineItem(text, PipelineType.Warning), CancelToken); + _ = TryQueue(new PipelineItem(text, PipelineType.Warning)); } /// Thread-safe verbose bridge for asynchronous cmdlet code. public new void WriteVerbose(string text) { + if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) { base.WriteVerbose(text); return; } - _currentOutPipe.Add(new PipelineItem(text, PipelineType.Verbose), CancelToken); + _ = TryQueue(new PipelineItem(text, PipelineType.Verbose)); } /// Thread-safe debug bridge for asynchronous cmdlet code. public new void WriteDebug(string text) { + if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) { base.WriteDebug(text); return; } - _currentOutPipe.Add(new PipelineItem(text, PipelineType.Debug), CancelToken); + _ = TryQueue(new PipelineItem(text, PipelineType.Debug)); } /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(InformationRecord informationRecord) { + if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) { base.WriteInformation(informationRecord); return; } - _currentOutPipe.Add(new PipelineItem(informationRecord, PipelineType.Information), CancelToken); + _ = TryQueue(new PipelineItem(informationRecord, PipelineType.Information)); } /// Thread-safe progress bridge for asynchronous cmdlet code. public new void WriteProgress(ProgressRecord progressRecord) { + if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (IsPipelineThread) { base.WriteProgress(progressRecord); return; } - _currentOutPipe.Add(new PipelineItem(progressRecord, PipelineType.Progress), CancelToken); + _ = TryQueue(new PipelineItem(progressRecord, PipelineType.Progress)); } /// Throws when PowerShell has requested cancellation. @@ -248,39 +348,89 @@ protected internal void ThrowIfStopped() /// public virtual void Dispose() - => _cancelSource.Dispose(); + { + if (Interlocked.Exchange(ref _disposed, 1) != 0) + return; + + _cancelSource.Cancel(); + _cancelSource.Dispose(); + } private bool IsPipelineThread => _pipelineThreadId != 0 && Environment.CurrentManagedThreadId == _pipelineThreadId; - private static void GetBlockTaskResult(Task blockTask) + private void GetBlockTaskResult(Task blockTask) { try { blockTask.GetAwaiter().GetResult(); } - catch (OperationCanceledException) + catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) { throw new PipelineStoppedException(); } } + private object? RequestPipelineReply(object? value, PipelineType type) + { + ThrowIfStopped(); + var replyPipe = new BlockingCollection(boundedCapacity: 1); + if (!TryQueue(new PipelineItem(value, type, replyPipe))) + throw new InvalidOperationException("No active PowerShell pipeline is available for the asynchronous request."); + + var reply = replyPipe.Take(CancelToken); + replyPipe.Dispose(); + return reply.Value; + } + + private bool TryQueue(PipelineItem item) + { + var outPipe = Volatile.Read(ref _currentOutPipe); + if (outPipe is null) + return false; + + try + { + outPipe.Add(item, CancelToken); + return true; + } + catch (InvalidOperationException) + { + return false; + } + } + private void RunBlockInAsync(Func task) { - using var outPipe = new BlockingCollection(); + var outPipe = new BlockingCollection(); Task blockTask; + var deferPipeDisposal = 0; + var pipeDisposed = 0; void ClearPipes() { - _currentOutPipe = null; - _pipelineThreadId = 0; + if (ReferenceEquals(Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe), outPipe)) + _pipelineThreadId = 0; CompleteAddingIfNeeded(outPipe); } + void DisposePipeOnce() + { + if (Interlocked.Exchange(ref pipeDisposed, 1) == 0) + outPipe.Dispose(); + } + static void CompleteAddingIfNeeded(BlockingCollection pipe) { - if (!pipe.IsAddingCompleted) - pipe.CompleteAdding(); + try + { + if (!pipe.IsAddingCompleted) + pipe.CompleteAdding(); + } + catch (ObjectDisposedException) + { + // A deferred worker may race the one-time disposal after a pipeline failure. + } } void PumpItem(PipelineItem item) @@ -314,19 +464,61 @@ void PumpItem(PipelineItem item) case PipelineType.Progress: base.WriteProgress((ProgressRecord)item.Value!); break; + case PipelineType.ShouldProcessTarget: + item.ReplyPipe!.Add(new PipelineReply(base.ShouldProcess((string)item.Value!))); + break; case PipelineType.ShouldProcess: var should = ((string Target, string Action))item.Value!; - item.ReplyPipe!.Add(base.ShouldProcess(should.Target, should.Action), CancelToken); + item.ReplyPipe!.Add(new PipelineReply(base.ShouldProcess(should.Target, should.Action))); + break; + case PipelineType.ShouldProcessVerbose: + var verbose = ((string Description, string Warning, string Caption))item.Value!; + item.ReplyPipe!.Add(new PipelineReply( + base.ShouldProcess(verbose.Description, verbose.Warning, verbose.Caption))); + break; + case PipelineType.ShouldProcessReason: + var reasonRequest = ((string Description, string Warning, string Caption))item.Value!; + var result = base.ShouldProcess( + reasonRequest.Description, + reasonRequest.Warning, + reasonRequest.Caption, + out var reason); + item.ReplyPipe!.Add(new PipelineReply((result, reason))); break; case PipelineType.ShouldContinue: var shouldContinue = ((string Query, string Caption))item.Value!; - item.ReplyPipe!.Add(base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption), CancelToken); + item.ReplyPipe!.Add(new PipelineReply( + base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption))); + break; + case PipelineType.ShouldContinueAll: + var shouldContinueAll = + ((string Query, string Caption, bool YesToAll, bool NoToAll))item.Value!; + var yesToAll = shouldContinueAll.YesToAll; + var noToAll = shouldContinueAll.NoToAll; + var continueAll = base.ShouldContinue( + shouldContinueAll.Query, + shouldContinueAll.Caption, + ref yesToAll, + ref noToAll); + item.ReplyPipe!.Add(new PipelineReply((continueAll, yesToAll, noToAll))); + break; + case PipelineType.ShouldContinueSecurity: + var shouldContinueSecurity = + ((string Query, string Caption, bool HasSecurityImpact, bool YesToAll, bool NoToAll))item.Value!; + yesToAll = shouldContinueSecurity.YesToAll; + noToAll = shouldContinueSecurity.NoToAll; + var continueSecurity = base.ShouldContinue( + shouldContinueSecurity.Query, + shouldContinueSecurity.Caption, + shouldContinueSecurity.HasSecurityImpact, + ref yesToAll, + ref noToAll); + item.ReplyPipe!.Add(new PipelineReply((continueSecurity, yesToAll, noToAll))); break; case PipelineType.PromptForCredential: var prompt = ((string Caption, string Message, string UserName, string TargetName))item.Value!; - item.ReplyPipe!.Add( - Host.UI.PromptForCredential(prompt.Caption, prompt.Message, prompt.UserName, prompt.TargetName), - CancelToken); + item.ReplyPipe!.Add(new PipelineReply( + Host.UI.PromptForCredential(prompt.Caption, prompt.Message, prompt.UserName, prompt.TargetName))); break; } } @@ -366,6 +558,7 @@ void PumpQueuedItems() finally { ClearPipes(); + DisposePipeOnce(); } GetBlockTaskResult(blockTask); @@ -373,7 +566,12 @@ void PumpQueuedItems() } _ = blockTask.ContinueWith( - completed => ClearPipes(), + completed => + { + ClearPipes(); + if (Volatile.Read(ref deferPipeDisposal) != 0) + DisposePipeOnce(); + }, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default); @@ -387,23 +585,26 @@ void PumpQueuedItems() } catch (Exception pipelineException) { + var stopRequested = _cancelSource.IsCancellationRequested; + Volatile.Write(ref deferPipeDisposal, 1); _cancelSource.Cancel(); CompleteAddingIfNeeded(outPipe); - try - { - blockTask.GetAwaiter().GetResult(); - } - catch (Exception completionException) when (completionException is OperationCanceledException or PipelineStoppedException) - { - _ = completionException; - } + if (blockTask.IsCompleted) + DisposePipeOnce(); - if (pipelineException is OperationCanceledException) + if (pipelineException is OperationCanceledException && stopRequested) throw new PipelineStoppedException(); throw; } - GetBlockTaskResult(blockTask); + try + { + GetBlockTaskResult(blockTask); + } + finally + { + DisposePipeOnce(); + } } } diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index eff71ddad..f850f9281 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -173,6 +173,60 @@ public void AsyncPSCmdlet_marshals_ShouldContinue_to_the_pipeline_thread(bool ap var item = Assert.Single(result); Assert.Equal(approved, item.BaseObject); } + + [Fact] + public void AsyncPSCmdlet_preserves_operation_cancellation_that_is_not_a_pipeline_stop() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncOperationCancellation", + typeof(TestAsyncOperationCancellationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncOperationCancellation"); + + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.IsAssignableFrom(exception.InnerException); + Assert.IsNotType(exception.InnerException); + } + + [Fact] + public void AsyncPSCmdlet_drops_late_worker_writes_after_the_pipeline_closes() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncLateWrite", + typeof(TestAsyncLateWriteCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncLateWrite"); + TestAsyncLateWriteCommand.Reset(); + + powerShell.Invoke(); + TestAsyncLateWriteCommand.WriteAfterCompletion(); + + Assert.Null(TestAsyncLateWriteCommand.LateWriteException); + } + + [Fact] + public void AsyncPSCmdlet_dispose_requests_cancellation() + { + using var command = new TestAsyncDisposableCommand(); + var stoppingToken = command.StoppingToken; + + command.Dispose(); + + Assert.True(stoppingToken.IsCancellationRequested); + } } [Cmdlet(VerbsDiagnostic.Test, "AsyncThreadAffinity")] @@ -321,6 +375,54 @@ protected override async Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncOperationCancellation")] +public sealed class TestAsyncOperationCancellationCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + throw new OperationCanceledException("operation timeout"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncLateWrite")] +public sealed class TestAsyncLateWriteCommand : AsyncPSCmdlet +{ + private static TestAsyncLateWriteCommand? _instance; + + public static Exception? LateWriteException { get; private set; } + + public static void Reset() + { + _instance = null; + LateWriteException = null; + } + + public static void WriteAfterCompletion() + { + try + { + _instance!.WriteWarning("late"); + } + catch (Exception exception) + { + LateWriteException = exception; + } + } + + protected override Task ProcessRecordAsync() + { + _instance = this; + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposable")] +public sealed class TestAsyncDisposableCommand : AsyncPSCmdlet +{ + public CancellationToken StoppingToken => CancelToken; +} + public sealed class ChoiceHost(bool approved) : PSHost { private readonly Guid _id = Guid.NewGuid(); From b2f38522380b812edd4da928392932632bc34e58 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 00:19:28 +0200 Subject: [PATCH 06/27] Close async pipeline lifecycle races --- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 168 +++++++++++++++++++---- PowerForge.Tests/AsyncPSCmdletTests.cs | 119 ++++++++++++++++ 2 files changed, 259 insertions(+), 28 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index e2caaf7e1..356a72e26 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -72,10 +72,13 @@ public PipelineItem(object? value, PipelineType type, BlockingCollection? _currentOutPipe; + private bool _cancelSourceDisposed; + private bool _disposeRequested; + private int _activeBlocks; private int _pipelineThreadId; - private int _disposed; /// Cancellation token triggered when PowerShell stops the cmdlet. protected internal CancellationToken CancelToken => _cancelSource.Token; @@ -106,12 +109,13 @@ protected virtual Task EndProcessingAsync() /// protected override void StopProcessing() - => _cancelSource.Cancel(); + => CancelSource(); /// Thread-safe ShouldProcess bridge for asynchronous cmdlet code. public new bool ShouldProcess(string? target) { - if (IsPipelineThread) + ThrowIfStopped(); + if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) return base.ShouldProcess(target ?? string.Empty); return (bool)RequestPipelineReply(target ?? string.Empty, PipelineType.ShouldProcessTarget)!; @@ -120,7 +124,8 @@ protected override void StopProcessing() /// Thread-safe ShouldProcess bridge for asynchronous cmdlet code. public new bool ShouldProcess(string? target, string action) { - if (IsPipelineThread) + ThrowIfStopped(); + if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) return base.ShouldProcess(target ?? string.Empty, action); return (bool)RequestPipelineReply((target ?? string.Empty, action), PipelineType.ShouldProcess)!; @@ -129,7 +134,8 @@ protected override void StopProcessing() /// Thread-safe ShouldProcess bridge for asynchronous cmdlet code. public new bool ShouldProcess(string verboseDescription, string verboseWarning, string caption) { - if (IsPipelineThread) + ThrowIfStopped(); + if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) return base.ShouldProcess(verboseDescription, verboseWarning, caption); return (bool)RequestPipelineReply( @@ -144,7 +150,8 @@ protected override void StopProcessing() string caption, out ShouldProcessReason shouldProcessReason) { - if (IsPipelineThread) + ThrowIfStopped(); + if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) return base.ShouldProcess(verboseDescription, verboseWarning, caption, out shouldProcessReason); var reply = ((bool Result, ShouldProcessReason Reason))RequestPipelineReply( @@ -157,7 +164,8 @@ protected override void StopProcessing() /// Thread-safe ShouldContinue bridge for asynchronous cmdlet code. public new bool ShouldContinue(string query, string caption) { - if (IsPipelineThread) + ThrowIfStopped(); + if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) return base.ShouldContinue(query, caption); return (bool)RequestPipelineReply((query, caption), PipelineType.ShouldContinue)!; @@ -166,7 +174,8 @@ protected override void StopProcessing() /// Thread-safe ShouldContinue bridge for asynchronous cmdlet code. public new bool ShouldContinue(string query, string caption, ref bool yesToAll, ref bool noToAll) { - if (IsPipelineThread) + ThrowIfStopped(); + if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) return base.ShouldContinue(query, caption, ref yesToAll, ref noToAll); var reply = ((bool Result, bool YesToAll, bool NoToAll))RequestPipelineReply( @@ -185,7 +194,8 @@ protected override void StopProcessing() ref bool yesToAll, ref bool noToAll) { - if (IsPipelineThread) + ThrowIfStopped(); + if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) return base.ShouldContinue(query, caption, hasSecurityImpact, ref yesToAll, ref noToAll); var reply = ((bool Result, bool YesToAll, bool NoToAll))RequestPipelineReply( @@ -199,7 +209,8 @@ protected override void StopProcessing() /// Thread-safe credential prompt bridge for asynchronous cmdlet code. public PSCredential? PromptForCredential(string caption, string message, string userName, string targetName) { - if (IsPipelineThread) + ThrowIfStopped(); + if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) return Host.UI.PromptForCredential(caption, message, userName, targetName); return (PSCredential?)RequestPipelineReply( @@ -349,11 +360,20 @@ protected internal void ThrowIfStopped() /// public virtual void Dispose() { - if (Interlocked.Exchange(ref _disposed, 1) != 0) - return; + lock (_lifecycleLock) + { + if (_disposeRequested) + return; - _cancelSource.Cancel(); - _cancelSource.Dispose(); + _disposeRequested = true; + } + + CancelSource(); + + lock (_lifecycleLock) + { + DisposeCancelSourceIfInactive(); + } } private bool IsPipelineThread @@ -374,12 +394,14 @@ private void GetBlockTaskResult(Task blockTask) private object? RequestPipelineReply(object? value, PipelineType type) { ThrowIfStopped(); - var replyPipe = new BlockingCollection(boundedCapacity: 1); + using var replyPipe = new BlockingCollection(boundedCapacity: 1); if (!TryQueue(new PipelineItem(value, type, replyPipe))) + { + ThrowIfStopped(); throw new InvalidOperationException("No active PowerShell pipeline is available for the asynchronous request."); + } var reply = replyPipe.Take(CancelToken); - replyPipe.Dispose(); return reply.Value; } @@ -398,9 +420,26 @@ private bool TryQueue(PipelineItem item) { return false; } + catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) + { + return false; + } } private void RunBlockInAsync(Func task) + { + EnterAsyncBlock(); + try + { + RunBlockInAsyncCore(task); + } + finally + { + ExitAsyncBlock(); + } + } + + private void RunBlockInAsyncCore(Func task) { var outPipe = new BlockingCollection(); Task blockTask; @@ -538,9 +577,14 @@ void PumpQueuedItems() SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); blockTask = task(); } - catch + catch (Exception exception) { ClearPipes(); + DisposePipeOnce(); + + if (exception is OperationCanceledException && _cancelSource.IsCancellationRequested) + throw new PipelineStoppedException(); + throw; } finally @@ -565,16 +609,35 @@ void PumpQueuedItems() return; } - _ = blockTask.ContinueWith( - completed => - { - ClearPipes(); - if (Volatile.Read(ref deferPipeDisposal) != 0) - DisposePipeOnce(); - }, - CancellationToken.None, - TaskContinuationOptions.ExecuteSynchronously, - TaskScheduler.Default); + RetainAsyncBlock(); + try + { + _ = blockTask.ContinueWith( + completed => + { + try + { + if (completed.IsFaulted) + _ = completed.Exception; + + ClearPipes(); + if (Volatile.Read(ref deferPipeDisposal) != 0) + DisposePipeOnce(); + } + finally + { + ExitAsyncBlock(); + } + }, + CancellationToken.None, + TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + } + catch + { + ExitAsyncBlock(); + throw; + } try { @@ -587,7 +650,7 @@ void PumpQueuedItems() { var stopRequested = _cancelSource.IsCancellationRequested; Volatile.Write(ref deferPipeDisposal, 1); - _cancelSource.Cancel(); + CancelSource(); CompleteAddingIfNeeded(outPipe); if (blockTask.IsCompleted) DisposePipeOnce(); @@ -607,4 +670,53 @@ void PumpQueuedItems() DisposePipeOnce(); } } + + private void EnterAsyncBlock() + { + lock (_lifecycleLock) + { + if (_disposeRequested) + throw new ObjectDisposedException(GetType().FullName); + + _activeBlocks++; + } + } + + private void ExitAsyncBlock() + { + lock (_lifecycleLock) + { + _activeBlocks--; + DisposeCancelSourceIfInactive(); + } + } + + private void RetainAsyncBlock() + { + lock (_lifecycleLock) + { + _activeBlocks++; + } + } + + private void CancelSource() + { + try + { + _cancelSource.Cancel(); + } + catch (ObjectDisposedException) + { + // Disposal may race a late StopProcessing callback after all async hooks have exited. + } + } + + private void DisposeCancelSourceIfInactive() + { + if (!_disposeRequested || _activeBlocks != 0 || _cancelSourceDisposed) + return; + + _cancelSource.Dispose(); + _cancelSourceDisposed = true; + } } diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index f850f9281..82ca4c44b 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -227,6 +227,59 @@ public void AsyncPSCmdlet_dispose_requests_cancellation() Assert.True(stoppingToken.IsCancellationRequested); } + + [Fact] + public void AsyncPSCmdlet_allows_ShouldProcess_before_the_async_base_hook_starts() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncEarlyShouldProcess", + typeof(TestAsyncEarlyShouldProcessCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncEarlyShouldProcess") + .AddParameter("Confirm", false); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + Assert.True((bool)Assert.Single(result).BaseObject); + } + + [Fact] + public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_the_async_hook_exits() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDisposeDuringHook", + typeof(TestAsyncDisposeDuringHookCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncDisposeDuringHook"); + TestAsyncDisposeDuringHookCommand.Reset(); + + powerShell.Invoke(); + + Assert.True(TestAsyncDisposeDuringHookCommand.ReadTokenAfterDispose); + Assert.Null(TestAsyncDisposeDuringHookCommand.TokenReadException); + } + + [Fact] + public void AsyncPSCmdlet_normalizes_synchronous_cancellation_after_stop() + { + using var command = new TestAsyncSynchronousStopCommand(); + + Assert.Throws(command.InvokeProcessRecord); + } } [Cmdlet(VerbsDiagnostic.Test, "AsyncThreadAffinity")] @@ -423,6 +476,72 @@ public sealed class TestAsyncDisposableCommand : AsyncPSCmdlet public CancellationToken StoppingToken => CancelToken; } +[Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyShouldProcess", SupportsShouldProcess = true)] +public sealed class TestAsyncEarlyShouldProcessCommand : AsyncPSCmdlet +{ + private bool _approved; + + protected override void ProcessRecord() + { + _approved = ShouldProcess("target"); + base.ProcessRecord(); + } + + protected override Task ProcessRecordAsync() + { + WriteObject(_approved); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringHook")] +public sealed class TestAsyncDisposeDuringHookCommand : AsyncPSCmdlet +{ + public static bool ReadTokenAfterDispose { get; private set; } + + public static Exception? TokenReadException { get; private set; } + + public static void Reset() + { + ReadTokenAfterDispose = false; + TokenReadException = null; + } + + protected override async Task ProcessRecordAsync() + { + Dispose(); + await Task.Yield(); + try + { + var token = CancelToken; + ReadTokenAfterDispose = true; + await Task.Delay(Timeout.InfiniteTimeSpan, token); + } + catch (OperationCanceledException) + { + throw; + } + catch (Exception exception) + { + TokenReadException = exception; + throw; + } + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousStop")] +public sealed class TestAsyncSynchronousStopCommand : AsyncPSCmdlet +{ + public void InvokeProcessRecord() + => base.ProcessRecord(); + + protected override Task ProcessRecordAsync() + { + Dispose(); + throw new OperationCanceledException("stopped synchronously"); + } +} + public sealed class ChoiceHost(bool approved) : PSHost { private readonly Guid _id = Guid.NewGuid(); From e5487bb82e33db3fd3972dc4d62dc0b38c8121f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 00:37:33 +0200 Subject: [PATCH 07/27] Coordinate synchronous writes and prompt replies --- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 147 +++++++++++++++++------ PowerForge.Tests/AsyncPSCmdletTests.cs | 6 +- 2 files changed, 114 insertions(+), 39 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 356a72e26..b0d5eb8f6 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -55,9 +55,53 @@ public PipelineReply(object? value) public object? Value { get; } } + private sealed class PipelineReplyChannel + { + private readonly BlockingCollection _pipe = new(boundedCapacity: 1); + private int _owners = 2; + + public PipelineReply Take(CancellationToken cancellationToken) + => _pipe.Take(cancellationToken); + + public void Publish(Func createValue) + { + try + { + var reply = new PipelineReply(createValue()); + try + { + _pipe.Add(reply); + } + catch (InvalidOperationException) + { + // The requester and pipeline can finish concurrently during cancellation. + } + } + finally + { + Release(); + } + } + + public void Abandon() + { + Release(); + Release(); + } + + public void ReleaseRequester() + => Release(); + + private void Release() + { + if (Interlocked.Decrement(ref _owners) == 0) + _pipe.Dispose(); + } + } + private sealed class PipelineItem { - public PipelineItem(object? value, PipelineType type, BlockingCollection? replyPipe = null) + public PipelineItem(object? value, PipelineType type, PipelineReplyChannel? replyPipe = null) { Value = value; Type = type; @@ -68,7 +112,7 @@ public PipelineItem(object? value, PipelineType type, BlockingCollection? ReplyPipe { get; } + public PipelineReplyChannel? ReplyPipe { get; } } private readonly CancellationTokenSource _cancelSource = new(); @@ -374,6 +418,8 @@ public virtual void Dispose() { DisposeCancelSourceIfInactive(); } + + _pipelineThreadId = 0; } private bool IsPipelineThread @@ -394,15 +440,22 @@ private void GetBlockTaskResult(Task blockTask) private object? RequestPipelineReply(object? value, PipelineType type) { ThrowIfStopped(); - using var replyPipe = new BlockingCollection(boundedCapacity: 1); + var replyPipe = new PipelineReplyChannel(); if (!TryQueue(new PipelineItem(value, type, replyPipe))) { + replyPipe.Abandon(); ThrowIfStopped(); throw new InvalidOperationException("No active PowerShell pipeline is available for the asynchronous request."); } - var reply = replyPipe.Take(CancelToken); - return reply.Value; + try + { + return replyPipe.Take(CancelToken).Value; + } + finally + { + replyPipe.ReleaseRequester(); + } } private bool TryQueue(PipelineItem item) @@ -416,6 +469,10 @@ private bool TryQueue(PipelineItem item) outPipe.Add(item, CancelToken); return true; } + catch (ObjectDisposedException) + { + return false; + } catch (InvalidOperationException) { return false; @@ -448,8 +505,7 @@ private void RunBlockInAsyncCore(Func task) void ClearPipes() { - if (ReferenceEquals(Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe), outPipe)) - _pipelineThreadId = 0; + _ = Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe); CompleteAddingIfNeeded(outPipe); } @@ -504,60 +560,75 @@ void PumpItem(PipelineItem item) base.WriteProgress((ProgressRecord)item.Value!); break; case PipelineType.ShouldProcessTarget: - item.ReplyPipe!.Add(new PipelineReply(base.ShouldProcess((string)item.Value!))); + item.ReplyPipe!.Publish( + () => base.ShouldProcess((string)item.Value!)); break; case PipelineType.ShouldProcess: var should = ((string Target, string Action))item.Value!; - item.ReplyPipe!.Add(new PipelineReply(base.ShouldProcess(should.Target, should.Action))); + item.ReplyPipe!.Publish( + () => base.ShouldProcess(should.Target, should.Action)); break; case PipelineType.ShouldProcessVerbose: var verbose = ((string Description, string Warning, string Caption))item.Value!; - item.ReplyPipe!.Add(new PipelineReply( - base.ShouldProcess(verbose.Description, verbose.Warning, verbose.Caption))); + item.ReplyPipe!.Publish( + () => base.ShouldProcess(verbose.Description, verbose.Warning, verbose.Caption)); break; case PipelineType.ShouldProcessReason: var reasonRequest = ((string Description, string Warning, string Caption))item.Value!; - var result = base.ShouldProcess( - reasonRequest.Description, - reasonRequest.Warning, - reasonRequest.Caption, - out var reason); - item.ReplyPipe!.Add(new PipelineReply((result, reason))); + item.ReplyPipe!.Publish(() => + { + var result = base.ShouldProcess( + reasonRequest.Description, + reasonRequest.Warning, + reasonRequest.Caption, + out var reason); + return (result, reason); + }); break; case PipelineType.ShouldContinue: var shouldContinue = ((string Query, string Caption))item.Value!; - item.ReplyPipe!.Add(new PipelineReply( - base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption))); + item.ReplyPipe!.Publish( + () => base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption)); break; case PipelineType.ShouldContinueAll: var shouldContinueAll = ((string Query, string Caption, bool YesToAll, bool NoToAll))item.Value!; - var yesToAll = shouldContinueAll.YesToAll; - var noToAll = shouldContinueAll.NoToAll; - var continueAll = base.ShouldContinue( - shouldContinueAll.Query, - shouldContinueAll.Caption, - ref yesToAll, - ref noToAll); - item.ReplyPipe!.Add(new PipelineReply((continueAll, yesToAll, noToAll))); + item.ReplyPipe!.Publish(() => + { + var yesToAll = shouldContinueAll.YesToAll; + var noToAll = shouldContinueAll.NoToAll; + var continueAll = base.ShouldContinue( + shouldContinueAll.Query, + shouldContinueAll.Caption, + ref yesToAll, + ref noToAll); + return (continueAll, yesToAll, noToAll); + }); break; case PipelineType.ShouldContinueSecurity: var shouldContinueSecurity = ((string Query, string Caption, bool HasSecurityImpact, bool YesToAll, bool NoToAll))item.Value!; - yesToAll = shouldContinueSecurity.YesToAll; - noToAll = shouldContinueSecurity.NoToAll; - var continueSecurity = base.ShouldContinue( - shouldContinueSecurity.Query, - shouldContinueSecurity.Caption, - shouldContinueSecurity.HasSecurityImpact, - ref yesToAll, - ref noToAll); - item.ReplyPipe!.Add(new PipelineReply((continueSecurity, yesToAll, noToAll))); + item.ReplyPipe!.Publish(() => + { + var yesToAll = shouldContinueSecurity.YesToAll; + var noToAll = shouldContinueSecurity.NoToAll; + var continueSecurity = base.ShouldContinue( + shouldContinueSecurity.Query, + shouldContinueSecurity.Caption, + shouldContinueSecurity.HasSecurityImpact, + ref yesToAll, + ref noToAll); + return (continueSecurity, yesToAll, noToAll); + }); break; case PipelineType.PromptForCredential: var prompt = ((string Caption, string Message, string UserName, string TargetName))item.Value!; - item.ReplyPipe!.Add(new PipelineReply( - Host.UI.PromptForCredential(prompt.Caption, prompt.Message, prompt.UserName, prompt.TargetName))); + item.ReplyPipe!.Publish( + () => Host.UI.PromptForCredential( + prompt.Caption, + prompt.Message, + prompt.UserName, + prompt.TargetName)); break; } } diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 82ca4c44b..1d328eb5c 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -248,7 +248,10 @@ public void AsyncPSCmdlet_allows_ShouldProcess_before_the_async_base_hook_starts var result = powerShell.Invoke(); Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); - Assert.True((bool)Assert.Single(result).BaseObject); + Assert.Collection( + result, + item => Assert.Equal("early-output", item.BaseObject), + item => Assert.True((bool)item.BaseObject)); } [Fact] @@ -483,6 +486,7 @@ public sealed class TestAsyncEarlyShouldProcessCommand : AsyncPSCmdlet protected override void ProcessRecord() { + WriteObject("early-output"); _approved = ShouldProcess("target"); base.ProcessRecord(); } From 34d1d0a482ceec5fd310e19199b64219fc0f1083 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 01:11:19 +0200 Subject: [PATCH 08/27] Harden lifecycle boundaries and scheduler isolation --- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 137 +++++++++++++++-------- PowerForge.Tests/AsyncPSCmdletTests.cs | 93 ++++++++++++++- 2 files changed, 184 insertions(+), 46 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index b0d5eb8f6..8fcf56b38 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -36,6 +36,7 @@ private enum PipelineType Verbose, Debug, Information, + InformationWithTags, Progress, ShouldProcessTarget, ShouldProcess, @@ -122,6 +123,7 @@ public PipelineItem(object? value, PipelineType type, PipelineReplyChannel? repl private bool _cancelSourceDisposed; private bool _disposeRequested; private int _activeBlocks; + private int _asyncLifecycleStarted; private int _pipelineThreadId; /// Cancellation token triggered when PowerShell stops the cmdlet. @@ -159,7 +161,7 @@ protected override void StopProcessing() public new bool ShouldProcess(string? target) { ThrowIfStopped(); - if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) return base.ShouldProcess(target ?? string.Empty); return (bool)RequestPipelineReply(target ?? string.Empty, PipelineType.ShouldProcessTarget)!; @@ -169,7 +171,7 @@ protected override void StopProcessing() public new bool ShouldProcess(string? target, string action) { ThrowIfStopped(); - if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) return base.ShouldProcess(target ?? string.Empty, action); return (bool)RequestPipelineReply((target ?? string.Empty, action), PipelineType.ShouldProcess)!; @@ -179,7 +181,7 @@ protected override void StopProcessing() public new bool ShouldProcess(string verboseDescription, string verboseWarning, string caption) { ThrowIfStopped(); - if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) return base.ShouldProcess(verboseDescription, verboseWarning, caption); return (bool)RequestPipelineReply( @@ -195,7 +197,7 @@ protected override void StopProcessing() out ShouldProcessReason shouldProcessReason) { ThrowIfStopped(); - if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) return base.ShouldProcess(verboseDescription, verboseWarning, caption, out shouldProcessReason); var reply = ((bool Result, ShouldProcessReason Reason))RequestPipelineReply( @@ -209,7 +211,7 @@ protected override void StopProcessing() public new bool ShouldContinue(string query, string caption) { ThrowIfStopped(); - if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) return base.ShouldContinue(query, caption); return (bool)RequestPipelineReply((query, caption), PipelineType.ShouldContinue)!; @@ -219,7 +221,7 @@ protected override void StopProcessing() public new bool ShouldContinue(string query, string caption, ref bool yesToAll, ref bool noToAll) { ThrowIfStopped(); - if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) return base.ShouldContinue(query, caption, ref yesToAll, ref noToAll); var reply = ((bool Result, bool YesToAll, bool NoToAll))RequestPipelineReply( @@ -239,7 +241,7 @@ protected override void StopProcessing() ref bool noToAll) { ThrowIfStopped(); - if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) return base.ShouldContinue(query, caption, hasSecurityImpact, ref yesToAll, ref noToAll); var reply = ((bool Result, bool YesToAll, bool NoToAll))RequestPipelineReply( @@ -254,7 +256,7 @@ protected override void StopProcessing() public PSCredential? PromptForCredential(string caption, string message, string userName, string targetName) { ThrowIfStopped(); - if (IsPipelineThread || Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) return Host.UI.PromptForCredential(caption, message, userName, targetName); return (PSCredential?)RequestPipelineReply( @@ -269,16 +271,17 @@ protected override void StopProcessing() /// Thread-safe output bridge for asynchronous cmdlet code. public new void WriteObject(object? sendToPipeline, bool enumerateCollection) { - if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) - return; - - ThrowIfStopped(); - if (IsPipelineThread) + if (CanAccessPipelineDirectly) { + ThrowIfStopped(); base.WriteObject(sendToPipeline, enumerateCollection); return; } + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + ThrowIfStopped(); _ = TryQueue(new PipelineItem( sendToPipeline, enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output)); @@ -287,16 +290,17 @@ protected override void StopProcessing() /// Thread-safe error bridge for asynchronous cmdlet code. public new void WriteError(ErrorRecord errorRecord) { - if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) - return; - - ThrowIfStopped(); - if (IsPipelineThread) + if (CanAccessPipelineDirectly) { + ThrowIfStopped(); base.WriteError(errorRecord); return; } + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + ThrowIfStopped(); _ = TryQueue(new PipelineItem(errorRecord, PipelineType.Error)); } @@ -304,93 +308,121 @@ protected override void StopProcessing() public new void ThrowTerminatingError(ErrorRecord errorRecord) { ThrowIfStopped(); - if (IsPipelineThread) + if (CanAccessPipelineDirectly) { base.ThrowTerminatingError(errorRecord); return; } - _ = TryQueue(new PipelineItem(errorRecord, PipelineType.TerminatingError)); + if (!TryQueue(new PipelineItem(errorRecord, PipelineType.TerminatingError))) + { + ThrowIfStopped(); + throw new InvalidOperationException( + "No active PowerShell pipeline is available for the terminating error."); + } + throw new PipelineStoppedException(); } /// Thread-safe warning bridge for asynchronous cmdlet code. public new void WriteWarning(string text) { - if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) - return; - - ThrowIfStopped(); - if (IsPipelineThread) + if (CanAccessPipelineDirectly) { + ThrowIfStopped(); base.WriteWarning(text); return; } + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + ThrowIfStopped(); _ = TryQueue(new PipelineItem(text, PipelineType.Warning)); } /// Thread-safe verbose bridge for asynchronous cmdlet code. public new void WriteVerbose(string text) { - if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) - return; - - ThrowIfStopped(); - if (IsPipelineThread) + if (CanAccessPipelineDirectly) { + ThrowIfStopped(); base.WriteVerbose(text); return; } + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + ThrowIfStopped(); _ = TryQueue(new PipelineItem(text, PipelineType.Verbose)); } /// Thread-safe debug bridge for asynchronous cmdlet code. public new void WriteDebug(string text) { - if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) - return; - - ThrowIfStopped(); - if (IsPipelineThread) + if (CanAccessPipelineDirectly) { + ThrowIfStopped(); base.WriteDebug(text); return; } + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + ThrowIfStopped(); _ = TryQueue(new PipelineItem(text, PipelineType.Debug)); } /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(InformationRecord informationRecord) { - if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) - return; - - ThrowIfStopped(); - if (IsPipelineThread) + if (CanAccessPipelineDirectly) { + ThrowIfStopped(); base.WriteInformation(informationRecord); return; } + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + ThrowIfStopped(); _ = TryQueue(new PipelineItem(informationRecord, PipelineType.Information)); } - /// Thread-safe progress bridge for asynchronous cmdlet code. - public new void WriteProgress(ProgressRecord progressRecord) + /// Thread-safe information bridge for asynchronous cmdlet code. + public new void WriteInformation(object messageData, string[] tags) { - if (!IsPipelineThread && Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) + { + ThrowIfStopped(); + base.WriteInformation(messageData, tags); + return; + } + + if (Volatile.Read(ref _currentOutPipe) is null) return; ThrowIfStopped(); - if (IsPipelineThread) + _ = TryQueue(new PipelineItem((messageData, tags), PipelineType.InformationWithTags)); + } + + /// Thread-safe progress bridge for asynchronous cmdlet code. + public new void WriteProgress(ProgressRecord progressRecord) + { + if (CanAccessPipelineDirectly) { + ThrowIfStopped(); base.WriteProgress(progressRecord); return; } + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + ThrowIfStopped(); _ = TryQueue(new PipelineItem(progressRecord, PipelineType.Progress)); } @@ -425,6 +457,9 @@ public virtual void Dispose() private bool IsPipelineThread => _pipelineThreadId != 0 && Environment.CurrentManagedThreadId == _pipelineThreadId; + private bool CanAccessPipelineDirectly + => IsPipelineThread || Volatile.Read(ref _asyncLifecycleStarted) == 0; + private void GetBlockTaskResult(Task blockTask) { try @@ -556,6 +591,10 @@ void PumpItem(PipelineItem item) case PipelineType.Information: base.WriteInformation((InformationRecord)item.Value!); break; + case PipelineType.InformationWithTags: + var information = ((object MessageData, string[] Tags))item.Value!; + base.WriteInformation(information.MessageData, information.Tags); + break; case PipelineType.Progress: base.WriteProgress((ProgressRecord)item.Value!); break; @@ -639,6 +678,7 @@ void PumpQueuedItems() PumpItem(item); } + Volatile.Write(ref _asyncLifecycleStarted, 1); _pipelineThreadId = Environment.CurrentManagedThreadId; _currentOutPipe = outPipe; @@ -646,7 +686,14 @@ void PumpQueuedItems() try { SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); - blockTask = task(); + blockTask = TaskScheduler.Current == TaskScheduler.Default + ? task() + : Task.Factory.StartNew( + task, + CancellationToken.None, + TaskCreationOptions.DenyChildAttach, + TaskScheduler.Default) + .Unwrap(); } catch (Exception exception) { diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 1d328eb5c..65c6e97c0 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -150,6 +150,27 @@ public void AsyncPSCmdlet_marshals_terminating_errors_to_the_pipeline_thread() Assert.False(TestAsyncTerminatingErrorCommand.ReachedAfterTermination); } + [Fact] + public void AsyncPSCmdlet_preserves_terminating_errors_before_the_initial_base_hook() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncEarlyTerminatingError", + typeof(TestAsyncEarlyTerminatingErrorCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncEarlyTerminatingError"); + + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.StartsWith("EarlyTerminatingError,", exception.ErrorRecord.FullyQualifiedErrorId, StringComparison.Ordinal); + Assert.Equal("early terminating failure", exception.InnerException?.Message); + } + [Theory] [InlineData(true)] [InlineData(false)] @@ -217,6 +238,31 @@ public void AsyncPSCmdlet_drops_late_worker_writes_after_the_pipeline_closes() Assert.Null(TestAsyncLateWriteCommand.LateWriteException); } + [Fact] + public void AsyncPSCmdlet_rejects_late_worker_interactions_after_the_pipeline_closes() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncLateInteraction", + typeof(TestAsyncLateInteractionCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncLateInteraction"); + TestAsyncLateInteractionCommand.Reset(); + + powerShell.Invoke(); + TestAsyncLateInteractionCommand.InteractAfterCompletion(); + + Assert.True( + TestAsyncLateInteractionCommand.LateInteractionException is + InvalidOperationException or PipelineStoppedException, + $"Unexpected late-interaction result: {TestAsyncLateInteractionCommand.LateInteractionException}"); + } + [Fact] public void AsyncPSCmdlet_dispose_requests_cancellation() { @@ -372,7 +418,8 @@ protected override void ProcessRecord() protected override async Task ProcessRecordAsync() { - Assert.Same(_scheduler, TaskScheduler.Current); + Assert.Same(TaskScheduler.Default, TaskScheduler.Current); + await Task.Factory.StartNew(static () => { }); await Task.Yield(); WriteObject(_scheduler.QueuedTaskCount); } @@ -421,6 +468,17 @@ protected override async Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyTerminatingError")] +public sealed class TestAsyncEarlyTerminatingErrorCommand : AsyncPSCmdlet +{ + protected override void BeginProcessing() + => ThrowTerminatingError(new ErrorRecord( + new InvalidOperationException("early terminating failure"), + "EarlyTerminatingError", + ErrorCategory.InvalidOperation, + targetObject: null)); +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncShouldContinue")] public sealed class TestAsyncShouldContinueCommand : AsyncPSCmdlet { @@ -473,6 +531,39 @@ protected override Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncLateInteraction", SupportsShouldProcess = true)] +public sealed class TestAsyncLateInteractionCommand : AsyncPSCmdlet +{ + private static TestAsyncLateInteractionCommand? _instance; + + public static Exception? LateInteractionException { get; private set; } + + public static void Reset() + { + _instance = null; + LateInteractionException = null; + } + + public static void InteractAfterCompletion() + => Task.Run(() => + { + try + { + _instance!.ShouldProcess("late-target"); + } + catch (Exception exception) + { + LateInteractionException = exception; + } + }).GetAwaiter().GetResult(); + + protected override Task ProcessRecordAsync() + { + _instance = this; + return Task.CompletedTask; + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncDisposable")] public sealed class TestAsyncDisposableCommand : AsyncPSCmdlet { From 1b7882019f206ffdb0a506dddfe899cc5038178d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 01:25:29 +0200 Subject: [PATCH 09/27] Make async stream writes cancellation-safe --- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 10 +--- PowerForge.Tests/AsyncPSCmdletTests.cs | 76 ++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 9 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 8fcf56b38..91754a1df 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -281,7 +281,6 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - ThrowIfStopped(); _ = TryQueue(new PipelineItem( sendToPipeline, enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output)); @@ -300,7 +299,6 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - ThrowIfStopped(); _ = TryQueue(new PipelineItem(errorRecord, PipelineType.Error)); } @@ -337,7 +335,6 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - ThrowIfStopped(); _ = TryQueue(new PipelineItem(text, PipelineType.Warning)); } @@ -354,7 +351,6 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - ThrowIfStopped(); _ = TryQueue(new PipelineItem(text, PipelineType.Verbose)); } @@ -371,7 +367,6 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - ThrowIfStopped(); _ = TryQueue(new PipelineItem(text, PipelineType.Debug)); } @@ -388,7 +383,6 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - ThrowIfStopped(); _ = TryQueue(new PipelineItem(informationRecord, PipelineType.Information)); } @@ -405,7 +399,6 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - ThrowIfStopped(); _ = TryQueue(new PipelineItem((messageData, tags), PipelineType.InformationWithTags)); } @@ -422,7 +415,6 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - ThrowIfStopped(); _ = TryQueue(new PipelineItem(progressRecord, PipelineType.Progress)); } @@ -680,7 +672,7 @@ void PumpQueuedItems() Volatile.Write(ref _asyncLifecycleStarted, 1); _pipelineThreadId = Environment.CurrentManagedThreadId; - _currentOutPipe = outPipe; + Volatile.Write(ref _currentOutPipe, outPipe); var synchronizationContext = SynchronizationContext.Current; try diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 65c6e97c0..fd528b5e6 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -127,6 +127,33 @@ public void AsyncPSCmdlet_translates_cancellation_to_a_pipeline_stop() Assert.IsNotType(exception); } + [Fact] + public void AsyncPSCmdlet_drops_background_stream_writes_during_pipeline_stop() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCancellationWrite", + typeof(TestAsyncCancellationWriteCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCancellationWrite"); + TestAsyncCancellationWriteCommand.Reset(); + + var invocation = powerShell.BeginInvoke(); + Assert.True( + TestAsyncCancellationWriteCommand.Started.Wait(TimeSpan.FromSeconds(5)), + "The asynchronous cmdlet did not start in time."); + + powerShell.Stop(); + Assert.Throws(() => powerShell.EndInvoke(invocation)); + + Assert.Null(TestAsyncCancellationWriteCommand.BackgroundWriteException); + } + [Fact] public void AsyncPSCmdlet_marshals_terminating_errors_to_the_pipeline_thread() { @@ -318,6 +345,9 @@ public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_the_async_ho powerShell.Invoke(); + Assert.True( + TestAsyncDisposeDuringHookCommand.TokenRead.Wait(TimeSpan.FromSeconds(5)), + "The asynchronous hook did not finish reading the token in time."); Assert.True(TestAsyncDisposeDuringHookCommand.ReadTokenAfterDispose); Assert.Null(TestAsyncDisposeDuringHookCommand.TokenReadException); } @@ -445,6 +475,44 @@ protected override async Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncCancellationWrite")] +public sealed class TestAsyncCancellationWriteCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _started = new(); + + public static ManualResetEventSlim Started => _started; + + public static Exception? BackgroundWriteException { get; private set; } + + public static void Reset() + { + _started.Dispose(); + _started = new ManualResetEventSlim(); + BackgroundWriteException = null; + } + + protected override async Task ProcessRecordAsync() + { + _started.Set(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } + catch (OperationCanceledException) when (CancelToken.IsCancellationRequested) + { + try + { + WriteProgress(new ProgressRecord(1, "cancelled", "finishing")); + WriteWarning("cancelled"); + } + catch (Exception exception) + { + BackgroundWriteException = exception; + } + } + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncTerminatingError")] public sealed class TestAsyncTerminatingErrorCommand : AsyncPSCmdlet { @@ -592,12 +660,18 @@ protected override Task ProcessRecordAsync() [Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringHook")] public sealed class TestAsyncDisposeDuringHookCommand : AsyncPSCmdlet { + private static ManualResetEventSlim _tokenRead = new(); + + public static ManualResetEventSlim TokenRead => _tokenRead; + public static bool ReadTokenAfterDispose { get; private set; } public static Exception? TokenReadException { get; private set; } public static void Reset() { + _tokenRead.Dispose(); + _tokenRead = new ManualResetEventSlim(); ReadTokenAfterDispose = false; TokenReadException = null; } @@ -610,6 +684,7 @@ protected override async Task ProcessRecordAsync() { var token = CancelToken; ReadTokenAfterDispose = true; + _tokenRead.Set(); await Task.Delay(Timeout.InfiniteTimeSpan, token); } catch (OperationCanceledException) @@ -619,6 +694,7 @@ protected override async Task ProcessRecordAsync() catch (Exception exception) { TokenReadException = exception; + _tokenRead.Set(); throw; } } From d0b45e5c0d201dccfebb7718082a99dc96cfd1de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 01:36:15 +0200 Subject: [PATCH 10/27] Preserve async hook affinity and write ordering --- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 50 ++++++++++++++++-------- PowerForge.Tests/AsyncPSCmdletTests.cs | 44 ++++++++++++++++++++- 2 files changed, 76 insertions(+), 18 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 91754a1df..bba861e31 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -26,6 +26,18 @@ public override void Post(SendOrPostCallback callback, object? state) => ThreadPool.QueueUserWorkItem(_ => callback(state)); } + private sealed class AsyncHookTaskScheduler : TaskScheduler + { + protected override System.Collections.Generic.IEnumerable? GetScheduledTasks() + => null; + + protected override void QueueTask(Task task) + => ThreadPool.QueueUserWorkItem(_ => TryExecuteTask(task)); + + protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) + => TryExecuteTask(task); + } + private enum PipelineType { Output, @@ -119,6 +131,7 @@ public PipelineItem(object? value, PipelineType type, PipelineReplyChannel? repl private readonly CancellationTokenSource _cancelSource = new(); private readonly object _lifecycleLock = new(); private static readonly SynchronizationContext HookSynchronizationContext = new AsyncHookSynchronizationContext(); + private static readonly TaskScheduler HookTaskScheduler = new AsyncHookTaskScheduler(); private BlockingCollection? _currentOutPipe; private bool _cancelSourceDisposed; private bool _disposeRequested; @@ -271,7 +284,7 @@ protected override void StopProcessing() /// Thread-safe output bridge for asynchronous cmdlet code. public new void WriteObject(object? sendToPipeline, bool enumerateCollection) { - if (CanAccessPipelineDirectly) + if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) { ThrowIfStopped(); base.WriteObject(sendToPipeline, enumerateCollection); @@ -289,7 +302,7 @@ protected override void StopProcessing() /// Thread-safe error bridge for asynchronous cmdlet code. public new void WriteError(ErrorRecord errorRecord) { - if (CanAccessPipelineDirectly) + if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) { ThrowIfStopped(); base.WriteError(errorRecord); @@ -325,7 +338,7 @@ protected override void StopProcessing() /// Thread-safe warning bridge for asynchronous cmdlet code. public new void WriteWarning(string text) { - if (CanAccessPipelineDirectly) + if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) { ThrowIfStopped(); base.WriteWarning(text); @@ -341,7 +354,7 @@ protected override void StopProcessing() /// Thread-safe verbose bridge for asynchronous cmdlet code. public new void WriteVerbose(string text) { - if (CanAccessPipelineDirectly) + if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) { ThrowIfStopped(); base.WriteVerbose(text); @@ -357,7 +370,7 @@ protected override void StopProcessing() /// Thread-safe debug bridge for asynchronous cmdlet code. public new void WriteDebug(string text) { - if (CanAccessPipelineDirectly) + if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) { ThrowIfStopped(); base.WriteDebug(text); @@ -373,7 +386,7 @@ protected override void StopProcessing() /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(InformationRecord informationRecord) { - if (CanAccessPipelineDirectly) + if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) { ThrowIfStopped(); base.WriteInformation(informationRecord); @@ -389,7 +402,7 @@ protected override void StopProcessing() /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(object messageData, string[] tags) { - if (CanAccessPipelineDirectly) + if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) { ThrowIfStopped(); base.WriteInformation(messageData, tags); @@ -405,7 +418,7 @@ protected override void StopProcessing() /// Thread-safe progress bridge for asynchronous cmdlet code. public new void WriteProgress(ProgressRecord progressRecord) { - if (CanAccessPipelineDirectly) + if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) { ThrowIfStopped(); base.WriteProgress(progressRecord); @@ -678,14 +691,19 @@ void PumpQueuedItems() try { SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); - blockTask = TaskScheduler.Current == TaskScheduler.Default - ? task() - : Task.Factory.StartNew( - task, - CancellationToken.None, - TaskCreationOptions.DenyChildAttach, - TaskScheduler.Default) - .Unwrap(); + if (TaskScheduler.Current == TaskScheduler.Default) + { + blockTask = task(); + } + else + { + var invocationTask = new Task( + task, + CancellationToken.None, + TaskCreationOptions.DenyChildAttach); + invocationTask.RunSynchronously(HookTaskScheduler); + blockTask = invocationTask.GetAwaiter().GetResult(); + } } catch (Exception exception) { diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index fd528b5e6..5d718f044 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -100,6 +100,30 @@ public void AsyncPSCmdlet_does_not_capture_a_custom_task_scheduler() Assert.Equal(0, item.BaseObject); } + [Fact] + public void AsyncPSCmdlet_preserves_fifo_order_for_pipeline_thread_and_worker_writes() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncWriteOrdering", + typeof(TestAsyncWriteOrderingCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncWriteOrdering"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + Assert.Collection( + result, + item => Assert.Equal("first", item.BaseObject), + item => Assert.Equal("second", item.BaseObject)); + } + [Fact] public void AsyncPSCmdlet_translates_cancellation_to_a_pipeline_stop() { @@ -442,19 +466,35 @@ protected override async Task ProcessRecordAsync() public sealed class TestAsyncTaskSchedulerCommand : AsyncPSCmdlet { private readonly ForwardingTaskScheduler _scheduler = new(); + private int _pipelineThreadId; protected override void ProcessRecord() - => _scheduler.Run(base.ProcessRecord); + { + _pipelineThreadId = Environment.CurrentManagedThreadId; + _scheduler.Run(base.ProcessRecord); + } protected override async Task ProcessRecordAsync() { - Assert.Same(TaskScheduler.Default, TaskScheduler.Current); + Assert.Equal(_pipelineThreadId, Environment.CurrentManagedThreadId); + Assert.NotSame(_scheduler, TaskScheduler.Current); await Task.Factory.StartNew(static () => { }); await Task.Yield(); WriteObject(_scheduler.QueuedTaskCount); } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncWriteOrdering")] +public sealed class TestAsyncWriteOrderingCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + Task.Run(() => WriteObject("first")).GetAwaiter().GetResult(); + WriteObject("second"); + return Task.CompletedTask; + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncCancellation")] public sealed class TestAsyncCancellationCommand : AsyncPSCmdlet { From 4ccd7e0908eb14aef3b1ef6a9df3126ba0711176 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 01:45:00 +0200 Subject: [PATCH 11/27] Split async pipeline internals and snapshot tags --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 443 +++++++++++++++ PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 439 +-------------- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 504 ++++++++++++++++++ PowerForge.Tests/AsyncPSCmdletTests.cs | 504 +----------------- 4 files changed, 974 insertions(+), 916 deletions(-) create mode 100644 PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs create mode 100644 PowerForge.Tests/AsyncPSCmdletTestCommands.cs diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs new file mode 100644 index 000000000..0bf4a26ff --- /dev/null +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -0,0 +1,443 @@ +using System; +using System.Collections.Concurrent; +using System.Management.Automation; +using System.Threading; +using System.Threading.Tasks; + +namespace PSPublishModule; + +public abstract partial class AsyncPSCmdlet +{ + /// Thread-safe progress bridge for asynchronous cmdlet code. + public new void WriteProgress(ProgressRecord progressRecord) + { + if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) + { + ThrowIfStopped(); + base.WriteProgress(progressRecord); + return; + } + + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + _ = TryQueue(new PipelineItem(progressRecord, PipelineType.Progress)); + } + + /// Throws when PowerShell has requested cancellation. + protected internal void ThrowIfStopped() + { + if (_cancelSource.IsCancellationRequested) + throw new PipelineStoppedException(); + } + + /// + public virtual void Dispose() + { + lock (_lifecycleLock) + { + if (_disposeRequested) + return; + + _disposeRequested = true; + } + + CancelSource(); + + lock (_lifecycleLock) + { + DisposeCancelSourceIfInactive(); + } + + _pipelineThreadId = 0; + } + + private bool IsPipelineThread + => _pipelineThreadId != 0 && Environment.CurrentManagedThreadId == _pipelineThreadId; + + private bool CanAccessPipelineDirectly + => IsPipelineThread || Volatile.Read(ref _asyncLifecycleStarted) == 0; + + private void GetBlockTaskResult(Task blockTask) + { + try + { + blockTask.GetAwaiter().GetResult(); + } + catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) + { + throw new PipelineStoppedException(); + } + } + + private object? RequestPipelineReply(object? value, PipelineType type) + { + ThrowIfStopped(); + var replyPipe = new PipelineReplyChannel(); + if (!TryQueue(new PipelineItem(value, type, replyPipe))) + { + replyPipe.Abandon(); + ThrowIfStopped(); + throw new InvalidOperationException("No active PowerShell pipeline is available for the asynchronous request."); + } + + try + { + return replyPipe.Take(CancelToken).Value; + } + finally + { + replyPipe.ReleaseRequester(); + } + } + + private bool TryQueue(PipelineItem item) + { + var outPipe = Volatile.Read(ref _currentOutPipe); + if (outPipe is null) + return false; + + try + { + outPipe.Add(item, CancelToken); + return true; + } + catch (ObjectDisposedException) + { + return false; + } + catch (InvalidOperationException) + { + return false; + } + catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) + { + return false; + } + } + + private void RunBlockInAsync(Func task) + { + EnterAsyncBlock(); + try + { + RunBlockInAsyncCore(task); + } + finally + { + ExitAsyncBlock(); + } + } + + private void RunBlockInAsyncCore(Func task) + { + var outPipe = new BlockingCollection(); + Task blockTask; + var deferPipeDisposal = 0; + var pipeDisposed = 0; + + void ClearPipes() + { + _ = Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe); + CompleteAddingIfNeeded(outPipe); + } + + void DisposePipeOnce() + { + if (Interlocked.Exchange(ref pipeDisposed, 1) == 0) + outPipe.Dispose(); + } + + static void CompleteAddingIfNeeded(BlockingCollection pipe) + { + try + { + if (!pipe.IsAddingCompleted) + pipe.CompleteAdding(); + } + catch (ObjectDisposedException) + { + // A deferred worker may race the one-time disposal after a pipeline failure. + } + } + + void PumpItem(PipelineItem item) + { + switch (item.Type) + { + case PipelineType.Output: + base.WriteObject(item.Value); + break; + case PipelineType.OutputEnumerate: + base.WriteObject(item.Value, enumerateCollection: true); + break; + case PipelineType.Error: + base.WriteError((ErrorRecord)item.Value!); + break; + case PipelineType.TerminatingError: + base.ThrowTerminatingError((ErrorRecord)item.Value!); + break; + case PipelineType.Warning: + base.WriteWarning((string)item.Value!); + break; + case PipelineType.Verbose: + base.WriteVerbose((string)item.Value!); + break; + case PipelineType.Debug: + base.WriteDebug((string)item.Value!); + break; + case PipelineType.Information: + base.WriteInformation((InformationRecord)item.Value!); + break; + case PipelineType.InformationWithTags: + var information = ((object MessageData, string[] Tags))item.Value!; + base.WriteInformation(information.MessageData, information.Tags); + break; + case PipelineType.Progress: + base.WriteProgress((ProgressRecord)item.Value!); + break; + case PipelineType.ShouldProcessTarget: + item.ReplyPipe!.Publish( + () => base.ShouldProcess((string)item.Value!)); + break; + case PipelineType.ShouldProcess: + var should = ((string Target, string Action))item.Value!; + item.ReplyPipe!.Publish( + () => base.ShouldProcess(should.Target, should.Action)); + break; + case PipelineType.ShouldProcessVerbose: + var verbose = ((string Description, string Warning, string Caption))item.Value!; + item.ReplyPipe!.Publish( + () => base.ShouldProcess(verbose.Description, verbose.Warning, verbose.Caption)); + break; + case PipelineType.ShouldProcessReason: + var reasonRequest = ((string Description, string Warning, string Caption))item.Value!; + item.ReplyPipe!.Publish(() => + { + var result = base.ShouldProcess( + reasonRequest.Description, + reasonRequest.Warning, + reasonRequest.Caption, + out var reason); + return (result, reason); + }); + break; + case PipelineType.ShouldContinue: + var shouldContinue = ((string Query, string Caption))item.Value!; + item.ReplyPipe!.Publish( + () => base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption)); + break; + case PipelineType.ShouldContinueAll: + var shouldContinueAll = + ((string Query, string Caption, bool YesToAll, bool NoToAll))item.Value!; + item.ReplyPipe!.Publish(() => + { + var yesToAll = shouldContinueAll.YesToAll; + var noToAll = shouldContinueAll.NoToAll; + var continueAll = base.ShouldContinue( + shouldContinueAll.Query, + shouldContinueAll.Caption, + ref yesToAll, + ref noToAll); + return (continueAll, yesToAll, noToAll); + }); + break; + case PipelineType.ShouldContinueSecurity: + var shouldContinueSecurity = + ((string Query, string Caption, bool HasSecurityImpact, bool YesToAll, bool NoToAll))item.Value!; + item.ReplyPipe!.Publish(() => + { + var yesToAll = shouldContinueSecurity.YesToAll; + var noToAll = shouldContinueSecurity.NoToAll; + var continueSecurity = base.ShouldContinue( + shouldContinueSecurity.Query, + shouldContinueSecurity.Caption, + shouldContinueSecurity.HasSecurityImpact, + ref yesToAll, + ref noToAll); + return (continueSecurity, yesToAll, noToAll); + }); + break; + case PipelineType.PromptForCredential: + var prompt = ((string Caption, string Message, string UserName, string TargetName))item.Value!; + item.ReplyPipe!.Publish( + () => Host.UI.PromptForCredential( + prompt.Caption, + prompt.Message, + prompt.UserName, + prompt.TargetName)); + break; + } + } + + void PumpQueuedItems() + { + while (outPipe.TryTake(out var item)) + PumpItem(item); + } + + Volatile.Write(ref _asyncLifecycleStarted, 1); + _pipelineThreadId = Environment.CurrentManagedThreadId; + Volatile.Write(ref _currentOutPipe, outPipe); + + var synchronizationContext = SynchronizationContext.Current; + try + { + SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); + if (TaskScheduler.Current == TaskScheduler.Default) + { + blockTask = task(); + } + else + { + var invocationTask = new Task( + task, + CancellationToken.None, + TaskCreationOptions.DenyChildAttach); + invocationTask.RunSynchronously(HookTaskScheduler); + blockTask = invocationTask.GetAwaiter().GetResult(); + } + } + catch (Exception exception) + { + ClearPipes(); + DisposePipeOnce(); + + if (exception is OperationCanceledException && _cancelSource.IsCancellationRequested) + throw new PipelineStoppedException(); + + throw; + } + finally + { + SynchronizationContext.SetSynchronizationContext(synchronizationContext); + } + + if (blockTask.IsCompleted) + { + CompleteAddingIfNeeded(outPipe); + try + { + PumpQueuedItems(); + } + finally + { + ClearPipes(); + DisposePipeOnce(); + } + + GetBlockTaskResult(blockTask); + return; + } + + RetainAsyncBlock(); + try + { + _ = blockTask.ContinueWith( + completed => + { + try + { + if (completed.IsFaulted) + _ = completed.Exception; + + ClearPipes(); + if (Volatile.Read(ref deferPipeDisposal) != 0) + DisposePipeOnce(); + } + finally + { + ExitAsyncBlock(); + } + }, + CancellationToken.None, + TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + } + catch + { + ExitAsyncBlock(); + throw; + } + + try + { + foreach (var item in outPipe.GetConsumingEnumerable(CancelToken)) + { + PumpItem(item); + } + } + catch (Exception pipelineException) + { + var stopRequested = _cancelSource.IsCancellationRequested; + Volatile.Write(ref deferPipeDisposal, 1); + CancelSource(); + CompleteAddingIfNeeded(outPipe); + if (blockTask.IsCompleted) + DisposePipeOnce(); + + if (pipelineException is OperationCanceledException && stopRequested) + throw new PipelineStoppedException(); + + throw; + } + + try + { + GetBlockTaskResult(blockTask); + } + finally + { + DisposePipeOnce(); + } + } + + private void EnterAsyncBlock() + { + lock (_lifecycleLock) + { + if (_disposeRequested) + throw new ObjectDisposedException(GetType().FullName); + + _activeBlocks++; + } + } + + private void ExitAsyncBlock() + { + lock (_lifecycleLock) + { + _activeBlocks--; + DisposeCancelSourceIfInactive(); + } + } + + private void RetainAsyncBlock() + { + lock (_lifecycleLock) + { + _activeBlocks++; + } + } + + private void CancelSource() + { + try + { + _cancelSource.Cancel(); + } + catch (ObjectDisposedException) + { + // Disposal may race a late StopProcessing callback after all async hooks have exited. + } + } + + private void DisposeCancelSourceIfInactive() + { + if (!_disposeRequested || _activeBlocks != 0 || _cancelSourceDisposed) + return; + + _cancelSource.Dispose(); + _cancelSourceDisposed = true; + } +} diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index bba861e31..834551355 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -18,7 +18,7 @@ namespace PSPublishModule; /// Keep hook implementations asynchronous all the way through and pass to /// cancellable engine operations. Do not block with Task.Wait, Task.Result, or Task.WaitAll. /// -public abstract class AsyncPSCmdlet : PSCmdlet, IDisposable +public abstract partial class AsyncPSCmdlet : PSCmdlet, IDisposable { private sealed class AsyncHookSynchronizationContext : SynchronizationContext { @@ -412,439 +412,8 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem((messageData, tags), PipelineType.InformationWithTags)); - } - - /// Thread-safe progress bridge for asynchronous cmdlet code. - public new void WriteProgress(ProgressRecord progressRecord) - { - if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) - { - ThrowIfStopped(); - base.WriteProgress(progressRecord); - return; - } - - if (Volatile.Read(ref _currentOutPipe) is null) - return; - - _ = TryQueue(new PipelineItem(progressRecord, PipelineType.Progress)); - } - - /// Throws when PowerShell has requested cancellation. - protected internal void ThrowIfStopped() - { - if (_cancelSource.IsCancellationRequested) - throw new PipelineStoppedException(); - } - - /// - public virtual void Dispose() - { - lock (_lifecycleLock) - { - if (_disposeRequested) - return; - - _disposeRequested = true; - } - - CancelSource(); - - lock (_lifecycleLock) - { - DisposeCancelSourceIfInactive(); - } - - _pipelineThreadId = 0; - } - - private bool IsPipelineThread - => _pipelineThreadId != 0 && Environment.CurrentManagedThreadId == _pipelineThreadId; - - private bool CanAccessPipelineDirectly - => IsPipelineThread || Volatile.Read(ref _asyncLifecycleStarted) == 0; - - private void GetBlockTaskResult(Task blockTask) - { - try - { - blockTask.GetAwaiter().GetResult(); - } - catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) - { - throw new PipelineStoppedException(); - } - } - - private object? RequestPipelineReply(object? value, PipelineType type) - { - ThrowIfStopped(); - var replyPipe = new PipelineReplyChannel(); - if (!TryQueue(new PipelineItem(value, type, replyPipe))) - { - replyPipe.Abandon(); - ThrowIfStopped(); - throw new InvalidOperationException("No active PowerShell pipeline is available for the asynchronous request."); - } - - try - { - return replyPipe.Take(CancelToken).Value; - } - finally - { - replyPipe.ReleaseRequester(); - } - } - - private bool TryQueue(PipelineItem item) - { - var outPipe = Volatile.Read(ref _currentOutPipe); - if (outPipe is null) - return false; - - try - { - outPipe.Add(item, CancelToken); - return true; - } - catch (ObjectDisposedException) - { - return false; - } - catch (InvalidOperationException) - { - return false; - } - catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) - { - return false; - } - } - - private void RunBlockInAsync(Func task) - { - EnterAsyncBlock(); - try - { - RunBlockInAsyncCore(task); - } - finally - { - ExitAsyncBlock(); - } - } - - private void RunBlockInAsyncCore(Func task) - { - var outPipe = new BlockingCollection(); - Task blockTask; - var deferPipeDisposal = 0; - var pipeDisposed = 0; - - void ClearPipes() - { - _ = Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe); - CompleteAddingIfNeeded(outPipe); - } - - void DisposePipeOnce() - { - if (Interlocked.Exchange(ref pipeDisposed, 1) == 0) - outPipe.Dispose(); - } - - static void CompleteAddingIfNeeded(BlockingCollection pipe) - { - try - { - if (!pipe.IsAddingCompleted) - pipe.CompleteAdding(); - } - catch (ObjectDisposedException) - { - // A deferred worker may race the one-time disposal after a pipeline failure. - } - } - - void PumpItem(PipelineItem item) - { - switch (item.Type) - { - case PipelineType.Output: - base.WriteObject(item.Value); - break; - case PipelineType.OutputEnumerate: - base.WriteObject(item.Value, enumerateCollection: true); - break; - case PipelineType.Error: - base.WriteError((ErrorRecord)item.Value!); - break; - case PipelineType.TerminatingError: - base.ThrowTerminatingError((ErrorRecord)item.Value!); - break; - case PipelineType.Warning: - base.WriteWarning((string)item.Value!); - break; - case PipelineType.Verbose: - base.WriteVerbose((string)item.Value!); - break; - case PipelineType.Debug: - base.WriteDebug((string)item.Value!); - break; - case PipelineType.Information: - base.WriteInformation((InformationRecord)item.Value!); - break; - case PipelineType.InformationWithTags: - var information = ((object MessageData, string[] Tags))item.Value!; - base.WriteInformation(information.MessageData, information.Tags); - break; - case PipelineType.Progress: - base.WriteProgress((ProgressRecord)item.Value!); - break; - case PipelineType.ShouldProcessTarget: - item.ReplyPipe!.Publish( - () => base.ShouldProcess((string)item.Value!)); - break; - case PipelineType.ShouldProcess: - var should = ((string Target, string Action))item.Value!; - item.ReplyPipe!.Publish( - () => base.ShouldProcess(should.Target, should.Action)); - break; - case PipelineType.ShouldProcessVerbose: - var verbose = ((string Description, string Warning, string Caption))item.Value!; - item.ReplyPipe!.Publish( - () => base.ShouldProcess(verbose.Description, verbose.Warning, verbose.Caption)); - break; - case PipelineType.ShouldProcessReason: - var reasonRequest = ((string Description, string Warning, string Caption))item.Value!; - item.ReplyPipe!.Publish(() => - { - var result = base.ShouldProcess( - reasonRequest.Description, - reasonRequest.Warning, - reasonRequest.Caption, - out var reason); - return (result, reason); - }); - break; - case PipelineType.ShouldContinue: - var shouldContinue = ((string Query, string Caption))item.Value!; - item.ReplyPipe!.Publish( - () => base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption)); - break; - case PipelineType.ShouldContinueAll: - var shouldContinueAll = - ((string Query, string Caption, bool YesToAll, bool NoToAll))item.Value!; - item.ReplyPipe!.Publish(() => - { - var yesToAll = shouldContinueAll.YesToAll; - var noToAll = shouldContinueAll.NoToAll; - var continueAll = base.ShouldContinue( - shouldContinueAll.Query, - shouldContinueAll.Caption, - ref yesToAll, - ref noToAll); - return (continueAll, yesToAll, noToAll); - }); - break; - case PipelineType.ShouldContinueSecurity: - var shouldContinueSecurity = - ((string Query, string Caption, bool HasSecurityImpact, bool YesToAll, bool NoToAll))item.Value!; - item.ReplyPipe!.Publish(() => - { - var yesToAll = shouldContinueSecurity.YesToAll; - var noToAll = shouldContinueSecurity.NoToAll; - var continueSecurity = base.ShouldContinue( - shouldContinueSecurity.Query, - shouldContinueSecurity.Caption, - shouldContinueSecurity.HasSecurityImpact, - ref yesToAll, - ref noToAll); - return (continueSecurity, yesToAll, noToAll); - }); - break; - case PipelineType.PromptForCredential: - var prompt = ((string Caption, string Message, string UserName, string TargetName))item.Value!; - item.ReplyPipe!.Publish( - () => Host.UI.PromptForCredential( - prompt.Caption, - prompt.Message, - prompt.UserName, - prompt.TargetName)); - break; - } - } - - void PumpQueuedItems() - { - while (outPipe.TryTake(out var item)) - PumpItem(item); - } - - Volatile.Write(ref _asyncLifecycleStarted, 1); - _pipelineThreadId = Environment.CurrentManagedThreadId; - Volatile.Write(ref _currentOutPipe, outPipe); - - var synchronizationContext = SynchronizationContext.Current; - try - { - SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); - if (TaskScheduler.Current == TaskScheduler.Default) - { - blockTask = task(); - } - else - { - var invocationTask = new Task( - task, - CancellationToken.None, - TaskCreationOptions.DenyChildAttach); - invocationTask.RunSynchronously(HookTaskScheduler); - blockTask = invocationTask.GetAwaiter().GetResult(); - } - } - catch (Exception exception) - { - ClearPipes(); - DisposePipeOnce(); - - if (exception is OperationCanceledException && _cancelSource.IsCancellationRequested) - throw new PipelineStoppedException(); - - throw; - } - finally - { - SynchronizationContext.SetSynchronizationContext(synchronizationContext); - } - - if (blockTask.IsCompleted) - { - CompleteAddingIfNeeded(outPipe); - try - { - PumpQueuedItems(); - } - finally - { - ClearPipes(); - DisposePipeOnce(); - } - - GetBlockTaskResult(blockTask); - return; - } - - RetainAsyncBlock(); - try - { - _ = blockTask.ContinueWith( - completed => - { - try - { - if (completed.IsFaulted) - _ = completed.Exception; - - ClearPipes(); - if (Volatile.Read(ref deferPipeDisposal) != 0) - DisposePipeOnce(); - } - finally - { - ExitAsyncBlock(); - } - }, - CancellationToken.None, - TaskContinuationOptions.ExecuteSynchronously, - TaskScheduler.Default); - } - catch - { - ExitAsyncBlock(); - throw; - } - - try - { - foreach (var item in outPipe.GetConsumingEnumerable(CancelToken)) - { - PumpItem(item); - } - } - catch (Exception pipelineException) - { - var stopRequested = _cancelSource.IsCancellationRequested; - Volatile.Write(ref deferPipeDisposal, 1); - CancelSource(); - CompleteAddingIfNeeded(outPipe); - if (blockTask.IsCompleted) - DisposePipeOnce(); - - if (pipelineException is OperationCanceledException && stopRequested) - throw new PipelineStoppedException(); - - throw; - } - - try - { - GetBlockTaskResult(blockTask); - } - finally - { - DisposePipeOnce(); - } - } - - private void EnterAsyncBlock() - { - lock (_lifecycleLock) - { - if (_disposeRequested) - throw new ObjectDisposedException(GetType().FullName); - - _activeBlocks++; - } - } - - private void ExitAsyncBlock() - { - lock (_lifecycleLock) - { - _activeBlocks--; - DisposeCancelSourceIfInactive(); - } - } - - private void RetainAsyncBlock() - { - lock (_lifecycleLock) - { - _activeBlocks++; - } - } - - private void CancelSource() - { - try - { - _cancelSource.Cancel(); - } - catch (ObjectDisposedException) - { - // Disposal may race a late StopProcessing callback after all async hooks have exited. - } - } - - private void DisposeCancelSourceIfInactive() - { - if (!_disposeRequested || _activeBlocks != 0 || _cancelSourceDisposed) - return; - - _cancelSource.Dispose(); - _cancelSourceDisposed = true; + _ = TryQueue(new PipelineItem( + (messageData, (string[])tags.Clone()), + PipelineType.InformationWithTags)); } } diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs new file mode 100644 index 000000000..11131c41b --- /dev/null +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -0,0 +1,504 @@ +using System.Collections.ObjectModel; +using System.Globalization; +using System.Management.Automation; +using System.Management.Automation.Host; +using System.Management.Automation.Runspaces; +using System.Security; +using System.Threading; +using System.Threading.Tasks; +using PSPublishModule; + +namespace PowerForge.Tests; + +[Cmdlet(VerbsDiagnostic.Test, "AsyncThreadAffinity")] +public sealed class TestAsyncThreadAffinityCommand : AsyncPSCmdlet +{ + private int _pipelineThreadId; + + protected override void BeginProcessing() + { + _pipelineThreadId = Environment.CurrentManagedThreadId; + base.BeginProcessing(); + } + + protected override async Task ProcessRecordAsync() + { + Assert.Equal(_pipelineThreadId, Environment.CurrentManagedThreadId); + await Task.Yield(); + WriteObject("post-await-output"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncQueuedOutput")] +public sealed class TestAsyncQueuedOutputCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + using var ready = new ManualResetEventSlim(); + Exception? workerException = null; + ThreadPool.QueueUserWorkItem(_ => + { + try + { + WriteObject("queued-output"); + } + catch (Exception ex) + { + workerException = ex; + } + finally + { + ready.Set(); + } + }); + + Assert.True(ready.Wait(TimeSpan.FromSeconds(5)), "Worker thread did not write output in time."); + if (workerException is not null) + throw workerException; + + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronizationContext")] +public sealed class TestAsyncSynchronizationContextCommand : AsyncPSCmdlet +{ + private ForwardingSynchronizationContext? _context; + + protected override void ProcessRecord() + { + var previousContext = SynchronizationContext.Current; + _context = new ForwardingSynchronizationContext(); + SynchronizationContext.SetSynchronizationContext(_context); + try + { + base.ProcessRecord(); + } + finally + { + SynchronizationContext.SetSynchronizationContext(previousContext); + } + } + + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(_context!.PostCount); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncTaskScheduler")] +public sealed class TestAsyncTaskSchedulerCommand : AsyncPSCmdlet +{ + private readonly ForwardingTaskScheduler _scheduler = new(); + private int _pipelineThreadId; + + protected override void ProcessRecord() + { + _pipelineThreadId = Environment.CurrentManagedThreadId; + _scheduler.Run(base.ProcessRecord); + } + + protected override async Task ProcessRecordAsync() + { + Assert.Equal(_pipelineThreadId, Environment.CurrentManagedThreadId); + Assert.NotSame(_scheduler, TaskScheduler.Current); + await Task.Factory.StartNew(static () => { }); + await Task.Yield(); + WriteObject(_scheduler.QueuedTaskCount); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncWriteOrdering")] +public sealed class TestAsyncWriteOrderingCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + Task.Run(() => WriteObject("first")).GetAwaiter().GetResult(); + WriteObject("second"); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncInformationTags")] +public sealed class TestAsyncInformationTagsCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + var tags = new[] { "before" }; + Task.Run(() => WriteInformation("message", tags)).GetAwaiter().GetResult(); + tags[0] = "after"; + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCancellation")] +public sealed class TestAsyncCancellationCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _started = new(); + + public static ManualResetEventSlim Started => _started; + + public static void Reset() + { + _started.Dispose(); + _started = new ManualResetEventSlim(); + } + + protected override async Task ProcessRecordAsync() + { + _started.Set(); + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCancellationWrite")] +public sealed class TestAsyncCancellationWriteCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _started = new(); + + public static ManualResetEventSlim Started => _started; + + public static Exception? BackgroundWriteException { get; private set; } + + public static void Reset() + { + _started.Dispose(); + _started = new ManualResetEventSlim(); + BackgroundWriteException = null; + } + + protected override async Task ProcessRecordAsync() + { + _started.Set(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } + catch (OperationCanceledException) when (CancelToken.IsCancellationRequested) + { + try + { + WriteProgress(new ProgressRecord(1, "cancelled", "finishing")); + WriteWarning("cancelled"); + } + catch (Exception exception) + { + BackgroundWriteException = exception; + } + } + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncTerminatingError")] +public sealed class TestAsyncTerminatingErrorCommand : AsyncPSCmdlet +{ + private static int _reachedAfterTermination; + + public static bool ReachedAfterTermination => Volatile.Read(ref _reachedAfterTermination) != 0; + + public static void Reset() + => Volatile.Write(ref _reachedAfterTermination, 0); + + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + ThrowTerminatingError(new ErrorRecord( + new InvalidOperationException("async terminating failure"), + "AsyncTerminatingError", + ErrorCategory.InvalidOperation, + targetObject: null)); + Volatile.Write(ref _reachedAfterTermination, 1); + WriteObject("unreachable"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyTerminatingError")] +public sealed class TestAsyncEarlyTerminatingErrorCommand : AsyncPSCmdlet +{ + protected override void BeginProcessing() + => ThrowTerminatingError(new ErrorRecord( + new InvalidOperationException("early terminating failure"), + "EarlyTerminatingError", + ErrorCategory.InvalidOperation, + targetObject: null)); +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncShouldContinue")] +public sealed class TestAsyncShouldContinueCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(ShouldContinue("Proceed?", "Question")); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncOperationCancellation")] +public sealed class TestAsyncOperationCancellationCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + throw new OperationCanceledException("operation timeout"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncLateWrite")] +public sealed class TestAsyncLateWriteCommand : AsyncPSCmdlet +{ + private static TestAsyncLateWriteCommand? _instance; + + public static Exception? LateWriteException { get; private set; } + + public static void Reset() + { + _instance = null; + LateWriteException = null; + } + + public static void WriteAfterCompletion() + { + try + { + _instance!.WriteWarning("late"); + } + catch (Exception exception) + { + LateWriteException = exception; + } + } + + protected override Task ProcessRecordAsync() + { + _instance = this; + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncLateInteraction", SupportsShouldProcess = true)] +public sealed class TestAsyncLateInteractionCommand : AsyncPSCmdlet +{ + private static TestAsyncLateInteractionCommand? _instance; + + public static Exception? LateInteractionException { get; private set; } + + public static void Reset() + { + _instance = null; + LateInteractionException = null; + } + + public static void InteractAfterCompletion() + => Task.Run(() => + { + try + { + _instance!.ShouldProcess("late-target"); + } + catch (Exception exception) + { + LateInteractionException = exception; + } + }).GetAwaiter().GetResult(); + + protected override Task ProcessRecordAsync() + { + _instance = this; + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposable")] +public sealed class TestAsyncDisposableCommand : AsyncPSCmdlet +{ + public CancellationToken StoppingToken => CancelToken; +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyShouldProcess", SupportsShouldProcess = true)] +public sealed class TestAsyncEarlyShouldProcessCommand : AsyncPSCmdlet +{ + private bool _approved; + + protected override void ProcessRecord() + { + WriteObject("early-output"); + _approved = ShouldProcess("target"); + base.ProcessRecord(); + } + + protected override Task ProcessRecordAsync() + { + WriteObject(_approved); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringHook")] +public sealed class TestAsyncDisposeDuringHookCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _tokenRead = new(); + + public static ManualResetEventSlim TokenRead => _tokenRead; + + public static bool ReadTokenAfterDispose { get; private set; } + + public static Exception? TokenReadException { get; private set; } + + public static void Reset() + { + _tokenRead.Dispose(); + _tokenRead = new ManualResetEventSlim(); + ReadTokenAfterDispose = false; + TokenReadException = null; + } + + protected override async Task ProcessRecordAsync() + { + Dispose(); + await Task.Yield(); + try + { + var token = CancelToken; + ReadTokenAfterDispose = true; + _tokenRead.Set(); + await Task.Delay(Timeout.InfiniteTimeSpan, token); + } + catch (OperationCanceledException) + { + throw; + } + catch (Exception exception) + { + TokenReadException = exception; + _tokenRead.Set(); + throw; + } + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousStop")] +public sealed class TestAsyncSynchronousStopCommand : AsyncPSCmdlet +{ + public void InvokeProcessRecord() + => base.ProcessRecord(); + + protected override Task ProcessRecordAsync() + { + Dispose(); + throw new OperationCanceledException("stopped synchronously"); + } +} + +public sealed class ChoiceHost(bool approved) : PSHost +{ + private readonly Guid _id = Guid.NewGuid(); + private readonly ChoiceHostUserInterface _ui = new(approved); + + public override Guid InstanceId => _id; + public override string Name => nameof(ChoiceHost); + public override Version Version => new(1, 0); + public override PSHostUserInterface UI => _ui; + public override CultureInfo CurrentCulture => CultureInfo.InvariantCulture; + public override CultureInfo CurrentUICulture => CultureInfo.InvariantCulture; + public override void EnterNestedPrompt() { } + public override void ExitNestedPrompt() { } + public override void NotifyBeginApplication() { } + public override void NotifyEndApplication() { } + public override void SetShouldExit(int exitCode) { } +} + +public sealed class ChoiceHostUserInterface(bool approved) : PSHostUserInterface +{ + public override PSHostRawUserInterface RawUI => null!; + + public override int PromptForChoice( + string caption, + string message, + Collection choices, + int defaultChoice) + => approved ? 0 : 1; + + public override string ReadLine() => string.Empty; + public override SecureString ReadLineAsSecureString() => new(); + public override void Write(string value) { } + public override void Write(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value) { } + public override void WriteLine(string value) { } + public override void WriteErrorLine(string value) { } + public override void WriteDebugLine(string message) { } + public override void WriteProgress(long sourceId, ProgressRecord record) { } + public override void WriteVerboseLine(string message) { } + public override void WriteWarningLine(string message) { } + + public override PSCredential PromptForCredential( + string caption, + string message, + string userName, + string targetName) + => new(userName, new SecureString()); + + public override PSCredential PromptForCredential( + string caption, + string message, + string userName, + string targetName, + PSCredentialTypes allowedCredentialTypes, + PSCredentialUIOptions options) + => PromptForCredential(caption, message, userName, targetName); + + public override Dictionary Prompt( + string caption, + string message, + Collection descriptions) + => []; +} + +public sealed class ForwardingTaskScheduler : TaskScheduler +{ + private int _executeNextTaskInline; + private int _queuedTaskCount; + + public int QueuedTaskCount => Volatile.Read(ref _queuedTaskCount); + + public void Run(Action action) + { + Volatile.Write(ref _executeNextTaskInline, 1); + var task = Task.Factory.StartNew( + action, + CancellationToken.None, + TaskCreationOptions.None, + this); + task.GetAwaiter().GetResult(); + } + + protected override IEnumerable GetScheduledTasks() + => Array.Empty(); + + protected override void QueueTask(Task task) + { + if (Interlocked.Exchange(ref _executeNextTaskInline, 0) == 1) + { + Assert.True(TryExecuteTask(task)); + return; + } + + Interlocked.Increment(ref _queuedTaskCount); + ThreadPool.QueueUserWorkItem(_ => TryExecuteTask(task)); + } + + protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) + => false; +} + +public sealed class ForwardingSynchronizationContext : SynchronizationContext +{ + private int _postCount; + + public int PostCount => Volatile.Read(ref _postCount); + + public override void Post(SendOrPostCallback callback, object? state) + { + Interlocked.Increment(ref _postCount); + ThreadPool.QueueUserWorkItem(_ => callback(state)); + } +} diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 5d718f044..2b7c77ad0 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -124,6 +124,29 @@ public void AsyncPSCmdlet_preserves_fifo_order_for_pipeline_thread_and_worker_wr item => Assert.Equal("second", item.BaseObject)); } + [Fact] + public void AsyncPSCmdlet_snapshots_information_tags_before_queueing() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncInformationTags", + typeof(TestAsyncInformationTagsCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncInformationTags"); + + powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + var record = Assert.Single(powerShell.Streams.Information); + Assert.Equal("message", record.MessageData); + Assert.Equal(["before"], record.Tags); + } + [Fact] public void AsyncPSCmdlet_translates_cancellation_to_a_pipeline_stop() { @@ -384,484 +407,3 @@ public void AsyncPSCmdlet_normalizes_synchronous_cancellation_after_stop() Assert.Throws(command.InvokeProcessRecord); } } - -[Cmdlet(VerbsDiagnostic.Test, "AsyncThreadAffinity")] -public sealed class TestAsyncThreadAffinityCommand : AsyncPSCmdlet -{ - private int _pipelineThreadId; - - protected override void BeginProcessing() - { - _pipelineThreadId = Environment.CurrentManagedThreadId; - base.BeginProcessing(); - } - - protected override async Task ProcessRecordAsync() - { - Assert.Equal(_pipelineThreadId, Environment.CurrentManagedThreadId); - await Task.Yield(); - WriteObject("post-await-output"); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncQueuedOutput")] -public sealed class TestAsyncQueuedOutputCommand : AsyncPSCmdlet -{ - protected override Task ProcessRecordAsync() - { - using var ready = new ManualResetEventSlim(); - Exception? workerException = null; - ThreadPool.QueueUserWorkItem(_ => - { - try - { - WriteObject("queued-output"); - } - catch (Exception ex) - { - workerException = ex; - } - finally - { - ready.Set(); - } - }); - - Assert.True(ready.Wait(TimeSpan.FromSeconds(5)), "Worker thread did not write output in time."); - if (workerException is not null) - throw workerException; - - return Task.CompletedTask; - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronizationContext")] -public sealed class TestAsyncSynchronizationContextCommand : AsyncPSCmdlet -{ - private ForwardingSynchronizationContext? _context; - - protected override void ProcessRecord() - { - var previousContext = SynchronizationContext.Current; - _context = new ForwardingSynchronizationContext(); - SynchronizationContext.SetSynchronizationContext(_context); - try - { - base.ProcessRecord(); - } - finally - { - SynchronizationContext.SetSynchronizationContext(previousContext); - } - } - - protected override async Task ProcessRecordAsync() - { - await Task.Yield(); - WriteObject(_context!.PostCount); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncTaskScheduler")] -public sealed class TestAsyncTaskSchedulerCommand : AsyncPSCmdlet -{ - private readonly ForwardingTaskScheduler _scheduler = new(); - private int _pipelineThreadId; - - protected override void ProcessRecord() - { - _pipelineThreadId = Environment.CurrentManagedThreadId; - _scheduler.Run(base.ProcessRecord); - } - - protected override async Task ProcessRecordAsync() - { - Assert.Equal(_pipelineThreadId, Environment.CurrentManagedThreadId); - Assert.NotSame(_scheduler, TaskScheduler.Current); - await Task.Factory.StartNew(static () => { }); - await Task.Yield(); - WriteObject(_scheduler.QueuedTaskCount); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncWriteOrdering")] -public sealed class TestAsyncWriteOrderingCommand : AsyncPSCmdlet -{ - protected override Task ProcessRecordAsync() - { - Task.Run(() => WriteObject("first")).GetAwaiter().GetResult(); - WriteObject("second"); - return Task.CompletedTask; - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncCancellation")] -public sealed class TestAsyncCancellationCommand : AsyncPSCmdlet -{ - private static ManualResetEventSlim _started = new(); - - public static ManualResetEventSlim Started => _started; - - public static void Reset() - { - _started.Dispose(); - _started = new ManualResetEventSlim(); - } - - protected override async Task ProcessRecordAsync() - { - _started.Set(); - await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncCancellationWrite")] -public sealed class TestAsyncCancellationWriteCommand : AsyncPSCmdlet -{ - private static ManualResetEventSlim _started = new(); - - public static ManualResetEventSlim Started => _started; - - public static Exception? BackgroundWriteException { get; private set; } - - public static void Reset() - { - _started.Dispose(); - _started = new ManualResetEventSlim(); - BackgroundWriteException = null; - } - - protected override async Task ProcessRecordAsync() - { - _started.Set(); - try - { - await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); - } - catch (OperationCanceledException) when (CancelToken.IsCancellationRequested) - { - try - { - WriteProgress(new ProgressRecord(1, "cancelled", "finishing")); - WriteWarning("cancelled"); - } - catch (Exception exception) - { - BackgroundWriteException = exception; - } - } - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncTerminatingError")] -public sealed class TestAsyncTerminatingErrorCommand : AsyncPSCmdlet -{ - private static int _reachedAfterTermination; - - public static bool ReachedAfterTermination => Volatile.Read(ref _reachedAfterTermination) != 0; - - public static void Reset() - => Volatile.Write(ref _reachedAfterTermination, 0); - - protected override async Task ProcessRecordAsync() - { - await Task.Yield(); - ThrowTerminatingError(new ErrorRecord( - new InvalidOperationException("async terminating failure"), - "AsyncTerminatingError", - ErrorCategory.InvalidOperation, - targetObject: null)); - Volatile.Write(ref _reachedAfterTermination, 1); - WriteObject("unreachable"); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyTerminatingError")] -public sealed class TestAsyncEarlyTerminatingErrorCommand : AsyncPSCmdlet -{ - protected override void BeginProcessing() - => ThrowTerminatingError(new ErrorRecord( - new InvalidOperationException("early terminating failure"), - "EarlyTerminatingError", - ErrorCategory.InvalidOperation, - targetObject: null)); -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncShouldContinue")] -public sealed class TestAsyncShouldContinueCommand : AsyncPSCmdlet -{ - protected override async Task ProcessRecordAsync() - { - await Task.Yield(); - WriteObject(ShouldContinue("Proceed?", "Question")); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncOperationCancellation")] -public sealed class TestAsyncOperationCancellationCommand : AsyncPSCmdlet -{ - protected override async Task ProcessRecordAsync() - { - await Task.Yield(); - throw new OperationCanceledException("operation timeout"); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncLateWrite")] -public sealed class TestAsyncLateWriteCommand : AsyncPSCmdlet -{ - private static TestAsyncLateWriteCommand? _instance; - - public static Exception? LateWriteException { get; private set; } - - public static void Reset() - { - _instance = null; - LateWriteException = null; - } - - public static void WriteAfterCompletion() - { - try - { - _instance!.WriteWarning("late"); - } - catch (Exception exception) - { - LateWriteException = exception; - } - } - - protected override Task ProcessRecordAsync() - { - _instance = this; - return Task.CompletedTask; - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncLateInteraction", SupportsShouldProcess = true)] -public sealed class TestAsyncLateInteractionCommand : AsyncPSCmdlet -{ - private static TestAsyncLateInteractionCommand? _instance; - - public static Exception? LateInteractionException { get; private set; } - - public static void Reset() - { - _instance = null; - LateInteractionException = null; - } - - public static void InteractAfterCompletion() - => Task.Run(() => - { - try - { - _instance!.ShouldProcess("late-target"); - } - catch (Exception exception) - { - LateInteractionException = exception; - } - }).GetAwaiter().GetResult(); - - protected override Task ProcessRecordAsync() - { - _instance = this; - return Task.CompletedTask; - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposable")] -public sealed class TestAsyncDisposableCommand : AsyncPSCmdlet -{ - public CancellationToken StoppingToken => CancelToken; -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyShouldProcess", SupportsShouldProcess = true)] -public sealed class TestAsyncEarlyShouldProcessCommand : AsyncPSCmdlet -{ - private bool _approved; - - protected override void ProcessRecord() - { - WriteObject("early-output"); - _approved = ShouldProcess("target"); - base.ProcessRecord(); - } - - protected override Task ProcessRecordAsync() - { - WriteObject(_approved); - return Task.CompletedTask; - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringHook")] -public sealed class TestAsyncDisposeDuringHookCommand : AsyncPSCmdlet -{ - private static ManualResetEventSlim _tokenRead = new(); - - public static ManualResetEventSlim TokenRead => _tokenRead; - - public static bool ReadTokenAfterDispose { get; private set; } - - public static Exception? TokenReadException { get; private set; } - - public static void Reset() - { - _tokenRead.Dispose(); - _tokenRead = new ManualResetEventSlim(); - ReadTokenAfterDispose = false; - TokenReadException = null; - } - - protected override async Task ProcessRecordAsync() - { - Dispose(); - await Task.Yield(); - try - { - var token = CancelToken; - ReadTokenAfterDispose = true; - _tokenRead.Set(); - await Task.Delay(Timeout.InfiniteTimeSpan, token); - } - catch (OperationCanceledException) - { - throw; - } - catch (Exception exception) - { - TokenReadException = exception; - _tokenRead.Set(); - throw; - } - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousStop")] -public sealed class TestAsyncSynchronousStopCommand : AsyncPSCmdlet -{ - public void InvokeProcessRecord() - => base.ProcessRecord(); - - protected override Task ProcessRecordAsync() - { - Dispose(); - throw new OperationCanceledException("stopped synchronously"); - } -} - -public sealed class ChoiceHost(bool approved) : PSHost -{ - private readonly Guid _id = Guid.NewGuid(); - private readonly ChoiceHostUserInterface _ui = new(approved); - - public override Guid InstanceId => _id; - public override string Name => nameof(ChoiceHost); - public override Version Version => new(1, 0); - public override PSHostUserInterface UI => _ui; - public override CultureInfo CurrentCulture => CultureInfo.InvariantCulture; - public override CultureInfo CurrentUICulture => CultureInfo.InvariantCulture; - public override void EnterNestedPrompt() { } - public override void ExitNestedPrompt() { } - public override void NotifyBeginApplication() { } - public override void NotifyEndApplication() { } - public override void SetShouldExit(int exitCode) { } -} - -public sealed class ChoiceHostUserInterface(bool approved) : PSHostUserInterface -{ - public override PSHostRawUserInterface RawUI => null!; - - public override int PromptForChoice( - string caption, - string message, - Collection choices, - int defaultChoice) - => approved ? 0 : 1; - - public override string ReadLine() => string.Empty; - public override SecureString ReadLineAsSecureString() => new(); - public override void Write(string value) { } - public override void Write(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value) { } - public override void WriteLine(string value) { } - public override void WriteErrorLine(string value) { } - public override void WriteDebugLine(string message) { } - public override void WriteProgress(long sourceId, ProgressRecord record) { } - public override void WriteVerboseLine(string message) { } - public override void WriteWarningLine(string message) { } - - public override PSCredential PromptForCredential( - string caption, - string message, - string userName, - string targetName) - => new(userName, new SecureString()); - - public override PSCredential PromptForCredential( - string caption, - string message, - string userName, - string targetName, - PSCredentialTypes allowedCredentialTypes, - PSCredentialUIOptions options) - => PromptForCredential(caption, message, userName, targetName); - - public override Dictionary Prompt( - string caption, - string message, - Collection descriptions) - => []; -} - -public sealed class ForwardingTaskScheduler : TaskScheduler -{ - private int _executeNextTaskInline; - private int _queuedTaskCount; - - public int QueuedTaskCount => Volatile.Read(ref _queuedTaskCount); - - public void Run(Action action) - { - Volatile.Write(ref _executeNextTaskInline, 1); - var task = Task.Factory.StartNew( - action, - CancellationToken.None, - TaskCreationOptions.None, - this); - task.GetAwaiter().GetResult(); - } - - protected override IEnumerable GetScheduledTasks() - => Array.Empty(); - - protected override void QueueTask(Task task) - { - if (Interlocked.Exchange(ref _executeNextTaskInline, 0) == 1) - { - Assert.True(TryExecuteTask(task)); - return; - } - - Interlocked.Increment(ref _queuedTaskCount); - ThreadPool.QueueUserWorkItem(_ => TryExecuteTask(task)); - } - - protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) - => false; -} - -public sealed class ForwardingSynchronizationContext : SynchronizationContext -{ - private int _postCount; - - public int PostCount => Volatile.Read(ref _postCount); - - public override void Post(SendOrPostCallback callback, object? state) - { - Interlocked.Increment(ref _postCount); - ThreadPool.QueueUserWorkItem(_ => callback(state)); - } -} From 4488113064cbb1988347086ff4bb086d29add6e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 02:33:26 +0200 Subject: [PATCH 12/27] Preserve synchronous pipeline semantics --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 100 ++++++++++++-- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 124 ++++++++++++++---- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 109 +++++++++++++++ PowerForge.Tests/AsyncPSCmdletTests.cs | 115 ++++++++++++++++ 4 files changed, 417 insertions(+), 31 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index 0bf4a26ff..fc4e13cb8 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -11,9 +11,9 @@ public abstract partial class AsyncPSCmdlet /// Thread-safe progress bridge for asynchronous cmdlet code. public new void WriteProgress(ProgressRecord progressRecord) { - if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) { - ThrowIfStopped(); + PrepareDirectPipelineAccess(); base.WriteProgress(progressRecord); return; } @@ -58,6 +58,38 @@ private bool IsPipelineThread private bool CanAccessPipelineDirectly => IsPipelineThread || Volatile.Read(ref _asyncLifecycleStarted) == 0; + private void PrepareDirectPipelineAccess() + { + ThrowIfStopped(); + if (IsPipelineThread) + Volatile.Read(ref _pumpQueuedItems)?.Invoke(); + } + + private void PrepareDirectPipelineInteraction() + { + ThrowIfStopped(); + ValidateInteractionGeneration(); + if (IsPipelineThread) + Volatile.Read(ref _pumpQueuedItems)?.Invoke(); + } + + private void ValidateInteractionGeneration() + { + if (Volatile.Read(ref _asyncLifecycleStarted) == 0) + return; + + var activeGeneration = Volatile.Read(ref _activeHookGeneration); + var originatingGeneration = _hookGeneration.Value; + if (activeGeneration == 0 && originatingGeneration == 0 && IsPipelineThread) + return; + + if (originatingGeneration == 0 || originatingGeneration != activeGeneration) + { + throw new InvalidOperationException( + "The asynchronous PowerShell lifecycle that originated this request is no longer active."); + } + } + private void GetBlockTaskResult(Task blockTask) { try @@ -73,8 +105,10 @@ private void GetBlockTaskResult(Task blockTask) private object? RequestPipelineReply(object? value, PipelineType type) { ThrowIfStopped(); + ValidateInteractionGeneration(); + var hookGeneration = _hookGeneration.Value; var replyPipe = new PipelineReplyChannel(); - if (!TryQueue(new PipelineItem(value, type, replyPipe))) + if (!TryQueue(new PipelineItem(value, type, replyPipe, hookGeneration))) { replyPipe.Abandon(); ThrowIfStopped(); @@ -135,17 +169,24 @@ private void RunBlockInAsyncCore(Func task) Task blockTask; var deferPipeDisposal = 0; var pipeDisposed = 0; + var hookGeneration = Interlocked.Increment(ref _nextHookGeneration); void ClearPipes() { + Volatile.Write(ref _pumpQueuedItems, null); _ = Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe); + _ = Interlocked.CompareExchange(ref _activeHookGeneration, 0, hookGeneration); CompleteAddingIfNeeded(outPipe); } void DisposePipeOnce() { if (Interlocked.Exchange(ref pipeDisposed, 1) == 0) + { + while (outPipe.TryTake(out var abandonedItem)) + abandonedItem.ReplyPipe?.ReleasePipeline(); outPipe.Dispose(); + } } static void CompleteAddingIfNeeded(BlockingCollection pipe) @@ -163,6 +204,13 @@ static void CompleteAddingIfNeeded(BlockingCollection pipe) void PumpItem(PipelineItem item) { + if (item.ReplyPipe is not null && + item.HookGeneration != Volatile.Read(ref _activeHookGeneration)) + { + item.ReplyPipe.ReleasePipeline(); + return; + } + switch (item.Type) { case PipelineType.Output: @@ -190,8 +238,8 @@ void PumpItem(PipelineItem item) base.WriteInformation((InformationRecord)item.Value!); break; case PipelineType.InformationWithTags: - var information = ((object MessageData, string[] Tags))item.Value!; - base.WriteInformation(information.MessageData, information.Tags); + var information = ((object MessageData, string[]? Tags))item.Value!; + base.WriteInformation(information.MessageData, information.Tags!); break; case PipelineType.Progress: base.WriteProgress((ProgressRecord)item.Value!); @@ -267,6 +315,23 @@ void PumpItem(PipelineItem item) prompt.UserName, prompt.TargetName)); break; + case PipelineType.PromptForCredentialOptions: + var promptOptions = + ((string Caption, + string Message, + string UserName, + string TargetName, + PSCredentialTypes AllowedCredentialTypes, + PSCredentialUIOptions Options))item.Value!; + item.ReplyPipe!.Publish( + () => Host.UI.PromptForCredential( + promptOptions.Caption, + promptOptions.Message, + promptOptions.UserName, + promptOptions.TargetName, + promptOptions.AllowedCredentialTypes, + promptOptions.Options)); + break; } } @@ -278,19 +343,23 @@ void PumpQueuedItems() Volatile.Write(ref _asyncLifecycleStarted, 1); _pipelineThreadId = Environment.CurrentManagedThreadId; + Volatile.Write(ref _activeHookGeneration, hookGeneration); + Volatile.Write(ref _pumpQueuedItems, PumpQueuedItems); Volatile.Write(ref _currentOutPipe, outPipe); var synchronizationContext = SynchronizationContext.Current; + var priorHookGeneration = _hookGeneration.Value; try { SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); + _hookGeneration.Value = hookGeneration; if (TaskScheduler.Current == TaskScheduler.Default) { blockTask = task(); } else { - var invocationTask = new Task( + using var invocationTask = new Task( task, CancellationToken.None, TaskCreationOptions.DenyChildAttach); @@ -300,8 +369,19 @@ void PumpQueuedItems() } catch (Exception exception) { - ClearPipes(); - DisposePipeOnce(); + try + { + PumpQueuedItems(); + } + catch + { + // Preserve the hook failure after best-effort delivery of records written before it. + } + finally + { + ClearPipes(); + DisposePipeOnce(); + } if (exception is OperationCanceledException && _cancelSource.IsCancellationRequested) throw new PipelineStoppedException(); @@ -310,11 +390,15 @@ void PumpQueuedItems() } finally { + _hookGeneration.Value = priorHookGeneration; SynchronizationContext.SetSynchronizationContext(synchronizationContext); } if (blockTask.IsCompleted) { + if (blockTask.IsFaulted) + _ = blockTask.Exception; + CompleteAddingIfNeeded(outPipe); try { diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 834551355..c653caa3c 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -57,7 +57,8 @@ private enum PipelineType ShouldContinue, ShouldContinueAll, ShouldContinueSecurity, - PromptForCredential + PromptForCredential, + PromptForCredentialOptions } private sealed class PipelineReply @@ -72,6 +73,8 @@ private sealed class PipelineReplyChannel { private readonly BlockingCollection _pipe = new(boundedCapacity: 1); private int _owners = 2; + private int _pipelineOwner = 1; + private int _requesterOwner = 1; public PipelineReply Take(CancellationToken cancellationToken) => _pipe.Take(cancellationToken); @@ -80,6 +83,9 @@ public void Publish(Func createValue) { try { + if (Volatile.Read(ref _requesterOwner) == 0) + return; + var reply = new PipelineReply(createValue()); try { @@ -92,18 +98,27 @@ public void Publish(Func createValue) } finally { - Release(); + ReleasePipeline(); } } public void Abandon() { - Release(); - Release(); + ReleaseRequester(); + ReleasePipeline(); } public void ReleaseRequester() - => Release(); + { + if (Interlocked.Exchange(ref _requesterOwner, 0) == 1) + Release(); + } + + public void ReleasePipeline() + { + if (Interlocked.Exchange(ref _pipelineOwner, 0) == 1) + Release(); + } private void Release() { @@ -114,11 +129,16 @@ private void Release() private sealed class PipelineItem { - public PipelineItem(object? value, PipelineType type, PipelineReplyChannel? replyPipe = null) + public PipelineItem( + object? value, + PipelineType type, + PipelineReplyChannel? replyPipe = null, + long hookGeneration = 0) { Value = value; Type = type; ReplyPipe = replyPipe; + HookGeneration = hookGeneration; } public object? Value { get; } @@ -126,13 +146,19 @@ public PipelineItem(object? value, PipelineType type, PipelineReplyChannel? repl public PipelineType Type { get; } public PipelineReplyChannel? ReplyPipe { get; } + + public long HookGeneration { get; } } private readonly CancellationTokenSource _cancelSource = new(); + private readonly AsyncLocal _hookGeneration = new(); private readonly object _lifecycleLock = new(); private static readonly SynchronizationContext HookSynchronizationContext = new AsyncHookSynchronizationContext(); private static readonly TaskScheduler HookTaskScheduler = new AsyncHookTaskScheduler(); private BlockingCollection? _currentOutPipe; + private Action? _pumpQueuedItems; + private long _activeHookGeneration; + private long _nextHookGeneration; private bool _cancelSourceDisposed; private bool _disposeRequested; private int _activeBlocks; @@ -175,7 +201,10 @@ protected override void StopProcessing() { ThrowIfStopped(); if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); return base.ShouldProcess(target ?? string.Empty); + } return (bool)RequestPipelineReply(target ?? string.Empty, PipelineType.ShouldProcessTarget)!; } @@ -185,7 +214,10 @@ protected override void StopProcessing() { ThrowIfStopped(); if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); return base.ShouldProcess(target ?? string.Empty, action); + } return (bool)RequestPipelineReply((target ?? string.Empty, action), PipelineType.ShouldProcess)!; } @@ -195,7 +227,10 @@ protected override void StopProcessing() { ThrowIfStopped(); if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); return base.ShouldProcess(verboseDescription, verboseWarning, caption); + } return (bool)RequestPipelineReply( (verboseDescription, verboseWarning, caption), @@ -211,7 +246,10 @@ protected override void StopProcessing() { ThrowIfStopped(); if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); return base.ShouldProcess(verboseDescription, verboseWarning, caption, out shouldProcessReason); + } var reply = ((bool Result, ShouldProcessReason Reason))RequestPipelineReply( (verboseDescription, verboseWarning, caption), @@ -225,7 +263,10 @@ protected override void StopProcessing() { ThrowIfStopped(); if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); return base.ShouldContinue(query, caption); + } return (bool)RequestPipelineReply((query, caption), PipelineType.ShouldContinue)!; } @@ -235,7 +276,10 @@ protected override void StopProcessing() { ThrowIfStopped(); if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); return base.ShouldContinue(query, caption, ref yesToAll, ref noToAll); + } var reply = ((bool Result, bool YesToAll, bool NoToAll))RequestPipelineReply( (query, caption, yesToAll, noToAll), @@ -255,7 +299,10 @@ protected override void StopProcessing() { ThrowIfStopped(); if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); return base.ShouldContinue(query, caption, hasSecurityImpact, ref yesToAll, ref noToAll); + } var reply = ((bool Result, bool YesToAll, bool NoToAll))RequestPipelineReply( (query, caption, hasSecurityImpact, yesToAll, noToAll), @@ -270,13 +317,43 @@ protected override void StopProcessing() { ThrowIfStopped(); if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); return Host.UI.PromptForCredential(caption, message, userName, targetName); + } return (PSCredential?)RequestPipelineReply( (caption, message, userName, targetName), PipelineType.PromptForCredential); } + /// Thread-safe credential prompt bridge for asynchronous cmdlet code. + public PSCredential? PromptForCredential( + string caption, + string message, + string userName, + string targetName, + PSCredentialTypes allowedCredentialTypes, + PSCredentialUIOptions options) + { + ThrowIfStopped(); + if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); + return Host.UI.PromptForCredential( + caption, + message, + userName, + targetName, + allowedCredentialTypes, + options); + } + + return (PSCredential?)RequestPipelineReply( + (caption, message, userName, targetName, allowedCredentialTypes, options), + PipelineType.PromptForCredentialOptions); + } + /// Thread-safe output bridge for asynchronous cmdlet code. public new void WriteObject(object? sendToPipeline) => WriteObject(sendToPipeline, enumerateCollection: false); @@ -284,9 +361,9 @@ protected override void StopProcessing() /// Thread-safe output bridge for asynchronous cmdlet code. public new void WriteObject(object? sendToPipeline, bool enumerateCollection) { - if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) { - ThrowIfStopped(); + PrepareDirectPipelineAccess(); base.WriteObject(sendToPipeline, enumerateCollection); return; } @@ -302,9 +379,9 @@ protected override void StopProcessing() /// Thread-safe error bridge for asynchronous cmdlet code. public new void WriteError(ErrorRecord errorRecord) { - if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) { - ThrowIfStopped(); + PrepareDirectPipelineAccess(); base.WriteError(errorRecord); return; } @@ -321,6 +398,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { + PrepareDirectPipelineAccess(); base.ThrowTerminatingError(errorRecord); return; } @@ -338,9 +416,9 @@ protected override void StopProcessing() /// Thread-safe warning bridge for asynchronous cmdlet code. public new void WriteWarning(string text) { - if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) { - ThrowIfStopped(); + PrepareDirectPipelineAccess(); base.WriteWarning(text); return; } @@ -354,9 +432,9 @@ protected override void StopProcessing() /// Thread-safe verbose bridge for asynchronous cmdlet code. public new void WriteVerbose(string text) { - if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) { - ThrowIfStopped(); + PrepareDirectPipelineAccess(); base.WriteVerbose(text); return; } @@ -370,9 +448,9 @@ protected override void StopProcessing() /// Thread-safe debug bridge for asynchronous cmdlet code. public new void WriteDebug(string text) { - if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) { - ThrowIfStopped(); + PrepareDirectPipelineAccess(); base.WriteDebug(text); return; } @@ -386,9 +464,9 @@ protected override void StopProcessing() /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(InformationRecord informationRecord) { - if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) { - ThrowIfStopped(); + PrepareDirectPipelineAccess(); base.WriteInformation(informationRecord); return; } @@ -400,12 +478,12 @@ protected override void StopProcessing() } /// Thread-safe information bridge for asynchronous cmdlet code. - public new void WriteInformation(object messageData, string[] tags) + public new void WriteInformation(object messageData, string[]? tags) { - if (CanAccessPipelineDirectly && Volatile.Read(ref _currentOutPipe) is null) + if (CanAccessPipelineDirectly) { - ThrowIfStopped(); - base.WriteInformation(messageData, tags); + PrepareDirectPipelineAccess(); + base.WriteInformation(messageData, tags!); return; } @@ -413,7 +491,7 @@ protected override void StopProcessing() return; _ = TryQueue(new PipelineItem( - (messageData, (string[])tags.Clone()), + (messageData, tags is null ? null : (string[])tags.Clone()), PipelineType.InformationWithTags)); } } diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index 11131c41b..b4aba90a9 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -132,6 +132,115 @@ protected override Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousError")] +public sealed class TestAsyncSynchronousErrorCommand : AsyncPSCmdlet +{ + private static int _reachedAfterError; + + public static bool ReachedAfterError => Volatile.Read(ref _reachedAfterError) != 0; + + public static void Reset() + => Volatile.Write(ref _reachedAfterError, 0); + + protected override Task ProcessRecordAsync() + { + WriteError(new ErrorRecord( + new InvalidOperationException("stopping error"), + "AsyncSynchronousError", + ErrorCategory.InvalidOperation, + targetObject: null)); + Volatile.Write(ref _reachedAfterError, 1); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousEnumeration")] +public sealed class TestAsyncSynchronousEnumerationCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + var values = new[] { 1, 2 }; + WriteObject(values, enumerateCollection: true); + values[0] = 9; + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousFailure")] +public sealed class TestAsyncSynchronousFailureCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + Task.Run(() => WriteWarning("before-failure")).GetAwaiter().GetResult(); + throw new InvalidOperationException("synchronous hook failure"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncNullInformationTags")] +public sealed class TestAsyncNullInformationTagsCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteInformation("untagged", tags: null); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncStaleInteraction", SupportsShouldProcess = true)] +public sealed class TestAsyncStaleInteractionCommand : AsyncPSCmdlet +{ + private static readonly ManualResetEventSlim SecondRecordStarted = new(); + private static readonly ManualResetEventSlim StaleInteractionFinished = new(); + private static int _recordNumber; + + [Parameter(Mandatory = true, ValueFromPipeline = true)] + public string InputObject { get; set; } = string.Empty; + + public static Exception? StaleInteractionException { get; private set; } + + public static void Reset() + { + SecondRecordStarted.Reset(); + StaleInteractionFinished.Reset(); + Volatile.Write(ref _recordNumber, 0); + StaleInteractionException = null; + } + + protected override async Task ProcessRecordAsync() + { + if (Interlocked.Increment(ref _recordNumber) == 1) + { + _ = Task.Run(() => + { + Assert.True( + SecondRecordStarted.Wait(TimeSpan.FromSeconds(5)), + "The second lifecycle did not start in time."); + try + { + _ = ShouldProcess("stale-target"); + } + catch (Exception exception) + { + StaleInteractionException = exception; + } + finally + { + StaleInteractionFinished.Set(); + } + }); + return; + } + + SecondRecordStarted.Set(); + await Task.Run(() => + { + Assert.True( + StaleInteractionFinished.Wait(TimeSpan.FromSeconds(5)), + "The stale interaction did not finish in time."); + }); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncCancellation")] public sealed class TestAsyncCancellationCommand : AsyncPSCmdlet { diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 2b7c77ad0..cd62633f8 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -147,6 +147,121 @@ public void AsyncPSCmdlet_snapshots_information_tags_before_queueing() Assert.Equal(["before"], record.Tags); } + [Fact] + public void AsyncPSCmdlet_applies_stopping_preferences_before_the_hook_continues() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSynchronousError", + typeof(TestAsyncSynchronousErrorCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncSynchronousError") + .AddParameter("ErrorAction", ActionPreference.Stop); + TestAsyncSynchronousErrorCommand.Reset(); + + _ = Assert.ThrowsAny(() => powerShell.Invoke()); + + Assert.False(TestAsyncSynchronousErrorCommand.ReachedAfterError); + } + + [Fact] + public void AsyncPSCmdlet_enumerates_pipeline_thread_collections_before_the_hook_mutates_them() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSynchronousEnumeration", + typeof(TestAsyncSynchronousEnumerationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncSynchronousEnumeration"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + Assert.Collection( + result, + item => Assert.Equal(1, item.BaseObject), + item => Assert.Equal(2, item.BaseObject)); + } + + [Fact] + public void AsyncPSCmdlet_drains_worker_records_before_a_synchronous_hook_failure() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSynchronousFailure", + typeof(TestAsyncSynchronousFailureCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncSynchronousFailure"); + + _ = Assert.ThrowsAny(() => powerShell.Invoke()); + + var warning = Assert.Single(powerShell.Streams.Warning); + Assert.Equal("before-failure", warning.Message); + } + + [Fact] + public void AsyncPSCmdlet_preserves_null_information_tags_after_await() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncNullInformationTags", + typeof(TestAsyncNullInformationTagsCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncNullInformationTags"); + + powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + var record = Assert.Single(powerShell.Streams.Information); + Assert.Equal("untagged", record.MessageData); + Assert.Empty(record.Tags); + } + + [Fact] + public void AsyncPSCmdlet_rejects_interactions_from_an_older_record_lifecycle() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncStaleInteraction", + typeof(TestAsyncStaleInteractionCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncStaleInteraction") + .AddParameter("Confirm", false); + TestAsyncStaleInteractionCommand.Reset(); + + powerShell.Invoke(new[] { "first", "second" }); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + Assert.IsType(TestAsyncStaleInteractionCommand.StaleInteractionException); + } + [Fact] public void AsyncPSCmdlet_translates_cancellation_to_a_pipeline_stop() { From 5fd3a26faeb5459556fdf85e95a9e357574396e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 02:54:30 +0200 Subject: [PATCH 13/27] Bind async records to lifecycle generations --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 21 ++++++++++---- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 29 ++++++++++++++++--- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 1 + PowerForge.Tests/AsyncPSCmdletTests.cs | 1 + 4 files changed, 43 insertions(+), 9 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index fc4e13cb8..7fc175e9f 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -117,7 +117,11 @@ private void GetBlockTaskResult(Task blockTask) try { - return replyPipe.Take(CancelToken).Value; + var reply = replyPipe.Take(CancelToken); + if (reply.Rejection is not null) + throw reply.Rejection; + + return reply.Value; } finally { @@ -127,6 +131,7 @@ private void GetBlockTaskResult(Task blockTask) private bool TryQueue(PipelineItem item) { + item.BindToHook(_hookGeneration.Value); var outPipe = Volatile.Read(ref _currentOutPipe); if (outPipe is null) return false; @@ -175,16 +180,18 @@ void ClearPipes() { Volatile.Write(ref _pumpQueuedItems, null); _ = Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe); - _ = Interlocked.CompareExchange(ref _activeHookGeneration, 0, hookGeneration); CompleteAddingIfNeeded(outPipe); } + void DeactivateHook() + => _ = Interlocked.CompareExchange(ref _activeHookGeneration, 0, hookGeneration); + void DisposePipeOnce() { if (Interlocked.Exchange(ref pipeDisposed, 1) == 0) { while (outPipe.TryTake(out var abandonedItem)) - abandonedItem.ReplyPipe?.ReleasePipeline(); + abandonedItem.ReplyPipe?.Reject(); outPipe.Dispose(); } } @@ -204,10 +211,10 @@ static void CompleteAddingIfNeeded(BlockingCollection pipe) void PumpItem(PipelineItem item) { - if (item.ReplyPipe is not null && + if (Volatile.Read(ref _asyncLifecycleStarted) != 0 && item.HookGeneration != Volatile.Read(ref _activeHookGeneration)) { - item.ReplyPipe.ReleasePipeline(); + item.ReplyPipe?.Reject(); return; } @@ -380,6 +387,7 @@ void PumpQueuedItems() finally { ClearPipes(); + DeactivateHook(); DisposePipeOnce(); } @@ -407,6 +415,7 @@ void PumpQueuedItems() finally { ClearPipes(); + DeactivateHook(); DisposePipeOnce(); } @@ -459,6 +468,7 @@ void PumpQueuedItems() CompleteAddingIfNeeded(outPipe); if (blockTask.IsCompleted) DisposePipeOnce(); + DeactivateHook(); if (pipelineException is OperationCanceledException && stopRequested) throw new PipelineStoppedException(); @@ -472,6 +482,7 @@ void PumpQueuedItems() } finally { + DeactivateHook(); DisposePipeOnce(); } } diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index c653caa3c..201f3ebf6 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -63,10 +63,15 @@ private enum PipelineType private sealed class PipelineReply { - public PipelineReply(object? value) - => Value = value; + public PipelineReply(object? value, Exception? rejection = null) + { + Value = value; + Rejection = rejection; + } public object? Value { get; } + + public Exception? Rejection { get; } } private sealed class PipelineReplyChannel @@ -80,13 +85,23 @@ public PipelineReply Take(CancellationToken cancellationToken) => _pipe.Take(cancellationToken); public void Publish(Func createValue) + => PublishReply(() => new PipelineReply(createValue())); + + public void Reject() + => PublishReply( + () => new PipelineReply( + value: null, + new InvalidOperationException( + "The asynchronous PowerShell lifecycle that originated this request is no longer active."))); + + private void PublishReply(Func createReply) { try { if (Volatile.Read(ref _requesterOwner) == 0) return; - var reply = new PipelineReply(createValue()); + var reply = createReply(); try { _pipe.Add(reply); @@ -147,7 +162,13 @@ public PipelineItem( public PipelineReplyChannel? ReplyPipe { get; } - public long HookGeneration { get; } + public long HookGeneration { get; private set; } + + public void BindToHook(long hookGeneration) + { + if (HookGeneration == 0) + HookGeneration = hookGeneration; + } } private readonly CancellationTokenSource _cancelSource = new(); diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index b4aba90a9..c23648621 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -217,6 +217,7 @@ protected override async Task ProcessRecordAsync() "The second lifecycle did not start in time."); try { + WriteWarning("stale-warning"); _ = ShouldProcess("stale-target"); } catch (Exception exception) diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index cd62633f8..15129b0c2 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -260,6 +260,7 @@ public void AsyncPSCmdlet_rejects_interactions_from_an_older_record_lifecycle() Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); Assert.IsType(TestAsyncStaleInteractionCommand.StaleInteractionException); + Assert.DoesNotContain(powerShell.Streams.Warning, static warning => warning.Message == "stale-warning"); } [Fact] From 84dda05b5164baa88b9bf6adb602f9dfa6b51c05 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 03:16:20 +0200 Subject: [PATCH 14/27] Harden async pipeline edge cases --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 30 +++++++++++++--- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 1 + PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 30 ++++++++++++++++ PowerForge.Tests/AsyncPSCmdletTests.cs | 35 +++++++++++++++++++ 4 files changed, 92 insertions(+), 4 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index 7fc175e9f..41762a949 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -56,7 +56,9 @@ private bool IsPipelineThread => _pipelineThreadId != 0 && Environment.CurrentManagedThreadId == _pipelineThreadId; private bool CanAccessPipelineDirectly - => IsPipelineThread || Volatile.Read(ref _asyncLifecycleStarted) == 0; + => IsPipelineThread || + (Volatile.Read(ref _asyncLifecycleStarted) == 0 && + Environment.CurrentManagedThreadId == _constructionThreadId); private void PrepareDirectPipelineAccess() { @@ -117,7 +119,16 @@ private void GetBlockTaskResult(Task blockTask) try { - var reply = replyPipe.Take(CancelToken); + PipelineReply reply; + try + { + reply = replyPipe.Take(CancelToken); + } + catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) + { + throw new PipelineStoppedException(); + } + if (reply.Rejection is not null) throw reply.Rejection; @@ -174,6 +185,7 @@ private void RunBlockInAsyncCore(Func task) Task blockTask; var deferPipeDisposal = 0; var pipeDisposed = 0; + var pumpingQueuedItems = 0; var hookGeneration = Interlocked.Increment(ref _nextHookGeneration); void ClearPipes() @@ -344,8 +356,18 @@ void PumpItem(PipelineItem item) void PumpQueuedItems() { - while (outPipe.TryTake(out var item)) - PumpItem(item); + if (Interlocked.Exchange(ref pumpingQueuedItems, 1) != 0) + return; + + try + { + while (outPipe.TryTake(out var item)) + PumpItem(item); + } + finally + { + Volatile.Write(ref pumpingQueuedItems, 0); + } } Volatile.Write(ref _asyncLifecycleStarted, 1); diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 201f3ebf6..33d2ad2c6 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -173,6 +173,7 @@ public void BindToHook(long hookGeneration) private readonly CancellationTokenSource _cancelSource = new(); private readonly AsyncLocal _hookGeneration = new(); + private readonly int _constructionThreadId = Environment.CurrentManagedThreadId; private readonly object _lifecycleLock = new(); private static readonly SynchronizationContext HookSynchronizationContext = new AsyncHookSynchronizationContext(); private static readonly TaskScheduler HookTaskScheduler = new AsyncHookTaskScheduler(); diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index c23648621..e47b23d27 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -425,6 +425,36 @@ public sealed class TestAsyncDisposableCommand : AsyncPSCmdlet public CancellationToken StoppingToken => CancelToken; } +[Cmdlet(VerbsDiagnostic.Test, "AsyncReentrantPump")] +public sealed class TestAsyncReentrantPumpCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(new WarningEnumerable(this), enumerateCollection: true); + WriteWarning("after-enumeration"); + } + + private sealed class WarningEnumerable : IEnumerable + { + private readonly TestAsyncReentrantPumpCommand _command; + + public WarningEnumerable(TestAsyncReentrantPumpCommand command) + { + _command = command; + } + + public IEnumerator GetEnumerator() + { + _command.WriteWarning("during-enumeration"); + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyShouldProcess", SupportsShouldProcess = true)] public sealed class TestAsyncEarlyShouldProcessCommand : AsyncPSCmdlet { diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 15129b0c2..1cf8d18ef 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -464,6 +464,41 @@ public void AsyncPSCmdlet_dispose_requests_cancellation() Assert.True(stoppingToken.IsCancellationRequested); } + [Fact] + public async Task AsyncPSCmdlet_drops_pre_lifecycle_writes_from_other_threads() + { + using var command = new TestAsyncDisposableCommand(); + + var exception = await Record.ExceptionAsync( + () => Task.Run(() => command.WriteWarning("too-early"))); + + Assert.Null(exception); + } + + [Fact] + public void AsyncPSCmdlet_does_not_recursively_drain_when_enumeration_writes_to_the_pipeline() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncReentrantPump", + typeof(TestAsyncReentrantPumpCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncReentrantPump"); + + var result = powerShell.Invoke(); + + Assert.Collection(result, item => Assert.Equal("value", item.BaseObject)); + Assert.Collection( + powerShell.Streams.Warning, + warning => Assert.Equal("during-enumeration", warning.Message), + warning => Assert.Equal("after-enumeration", warning.Message)); + } + [Fact] public void AsyncPSCmdlet_allows_ShouldProcess_before_the_async_base_hook_starts() { From 8f478d028647dc505c60511d3d0a654e1ddde11f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 03:47:08 +0200 Subject: [PATCH 15/27] Propagate host interaction failures --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 4 ++- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 25 +++++++++++++++---- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 8 +++--- PowerForge.Tests/AsyncPSCmdletTests.cs | 21 ++++++++++++++++ 4 files changed, 48 insertions(+), 10 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index 41762a949..49c06d986 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -258,7 +258,9 @@ void PumpItem(PipelineItem item) break; case PipelineType.InformationWithTags: var information = ((object MessageData, string[]? Tags))item.Value!; - base.WriteInformation(information.MessageData, information.Tags!); + base.WriteInformation( + information.MessageData, + information.Tags ?? Array.Empty()); break; case PipelineType.Progress: base.WriteProgress((ProgressRecord)item.Value!); diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 33d2ad2c6..a428cf994 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -101,15 +101,18 @@ private void PublishReply(Func createReply) if (Volatile.Read(ref _requesterOwner) == 0) return; - var reply = createReply(); + PipelineReply reply; try { - _pipe.Add(reply); + reply = createReply(); } - catch (InvalidOperationException) + catch (Exception exception) { - // The requester and pipeline can finish concurrently during cancellation. + TryPublish(new PipelineReply(value: null, exception)); + throw; } + + TryPublish(reply); } finally { @@ -117,6 +120,18 @@ private void PublishReply(Func createReply) } } + private void TryPublish(PipelineReply reply) + { + try + { + _pipe.Add(reply); + } + catch (InvalidOperationException) + { + // The requester and pipeline can finish concurrently during cancellation. + } + } + public void Abandon() { ReleaseRequester(); @@ -505,7 +520,7 @@ protected override void StopProcessing() if (CanAccessPipelineDirectly) { PrepareDirectPipelineAccess(); - base.WriteInformation(messageData, tags!); + base.WriteInformation(messageData, tags ?? Array.Empty()); return; } diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index e47b23d27..e589d2b43 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -530,10 +530,10 @@ protected override Task ProcessRecordAsync() } } -public sealed class ChoiceHost(bool approved) : PSHost +public sealed class ChoiceHost(bool approved, Exception? promptFailure = null) : PSHost { private readonly Guid _id = Guid.NewGuid(); - private readonly ChoiceHostUserInterface _ui = new(approved); + private readonly ChoiceHostUserInterface _ui = new(approved, promptFailure); public override Guid InstanceId => _id; public override string Name => nameof(ChoiceHost); @@ -548,7 +548,7 @@ public override void NotifyEndApplication() { } public override void SetShouldExit(int exitCode) { } } -public sealed class ChoiceHostUserInterface(bool approved) : PSHostUserInterface +public sealed class ChoiceHostUserInterface(bool approved, Exception? promptFailure) : PSHostUserInterface { public override PSHostRawUserInterface RawUI => null!; @@ -557,7 +557,7 @@ public override int PromptForChoice( string message, Collection choices, int defaultChoice) - => approved ? 0 : 1; + => promptFailure is null ? (approved ? 0 : 1) : throw promptFailure; public override string ReadLine() => string.Empty; public override SecureString ReadLineAsSecureString() => new(); diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 1cf8d18ef..b0219fdba 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -385,6 +385,27 @@ public void AsyncPSCmdlet_marshals_ShouldContinue_to_the_pipeline_thread(bool ap Assert.Equal(approved, item.BaseObject); } + [Fact] + public void AsyncPSCmdlet_returns_host_interaction_failures_to_the_requesting_worker() + { + var sessionState = InitialSessionState.Create(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncShouldContinue", + typeof(TestAsyncShouldContinueCommand), + helpFileName: null)); + + var promptFailure = new InvalidOperationException("host prompt failed"); + using var runspace = RunspaceFactory.CreateRunspace(new ChoiceHost(approved: false, promptFailure), sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncShouldContinue"); + + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.Same(promptFailure, exception.InnerException); + } + [Fact] public void AsyncPSCmdlet_preserves_operation_cancellation_that_is_not_a_pipeline_stop() { From 4e66e9184e53c2349ac6c7bd4d6dd507646d4d79 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 04:12:38 +0200 Subject: [PATCH 16/27] Harden async pipeline callback ownership --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 38 +++++++-- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 35 ++++++--- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 29 +++++++ PowerForge.Tests/AsyncPSCmdletTests.cs | 78 ++++++++++++++++++- 4 files changed, 162 insertions(+), 18 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index 49c06d986..3903b805c 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -42,14 +42,19 @@ public virtual void Dispose() _disposeRequested = true; } - CancelSource(); - - lock (_lifecycleLock) + try { - DisposeCancelSourceIfInactive(); + CancelSource(); } + finally + { + lock (_lifecycleLock) + { + DisposeCancelSourceIfInactive(); + } - _pipelineThreadId = 0; + _pipelineThreadId = 0; + } } private bool IsPipelineThread @@ -140,6 +145,26 @@ private void GetBlockTaskResult(Task blockTask) } } + /// + /// Captures an output writer for callbacks whose producer does not flow the hook execution context. + /// + /// + /// Capture the writer inside an asynchronous PowerShell hook. Calls made after that hook ends are + /// rejected rather than being rebound to a later record lifecycle. + /// + protected Action CapturePipelineWriter(bool enumerateCollection = false) + { + var hookGeneration = _hookGeneration.Value; + if (hookGeneration == 0) + { + throw new InvalidOperationException( + "A lifecycle-bound pipeline writer can only be captured from an asynchronous PowerShell hook."); + } + + var pipelineType = enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output; + return value => _ = TryQueue(new PipelineItem(value, pipelineType, hookGeneration: hookGeneration)); + } + private bool TryQueue(PipelineItem item) { item.BindToHook(_hookGeneration.Value); @@ -265,6 +290,9 @@ void PumpItem(PipelineItem item) case PipelineType.Progress: base.WriteProgress((ProgressRecord)item.Value!); break; + case PipelineType.CommandDetail: + base.WriteCommandDetail((string)item.Value!); + break; case PipelineType.ShouldProcessTarget: item.ReplyPipe!.Publish( () => base.ShouldProcess((string)item.Value!)); diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index a428cf994..f6bd48ddc 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -50,6 +50,7 @@ private enum PipelineType Information, InformationWithTags, Progress, + CommandDetail, ShouldProcessTarget, ShouldProcess, ShouldProcessVerbose, @@ -451,51 +452,67 @@ protected override void StopProcessing() } /// Thread-safe warning bridge for asynchronous cmdlet code. - public new void WriteWarning(string text) + public new void WriteWarning(string message) { if (CanAccessPipelineDirectly) { PrepareDirectPipelineAccess(); - base.WriteWarning(text); + base.WriteWarning(message); return; } if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(text, PipelineType.Warning)); + _ = TryQueue(new PipelineItem(message, PipelineType.Warning)); } /// Thread-safe verbose bridge for asynchronous cmdlet code. - public new void WriteVerbose(string text) + public new void WriteVerbose(string message) { if (CanAccessPipelineDirectly) { PrepareDirectPipelineAccess(); - base.WriteVerbose(text); + base.WriteVerbose(message); return; } if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(text, PipelineType.Verbose)); + _ = TryQueue(new PipelineItem(message, PipelineType.Verbose)); } /// Thread-safe debug bridge for asynchronous cmdlet code. - public new void WriteDebug(string text) + public new void WriteDebug(string message) { if (CanAccessPipelineDirectly) { PrepareDirectPipelineAccess(); - base.WriteDebug(text); + base.WriteDebug(message); return; } if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(text, PipelineType.Debug)); + _ = TryQueue(new PipelineItem(message, PipelineType.Debug)); + } + + /// Thread-safe command-detail bridge for asynchronous cmdlet code. + public new void WriteCommandDetail(string text) + { + if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineAccess(); + base.WriteCommandDetail(text); + return; + } + + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + _ = TryQueue(new PipelineItem(text, PipelineType.CommandDetail)); } /// Thread-safe information bridge for asynchronous cmdlet code. diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index e589d2b43..9ccb29bb2 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -425,6 +425,35 @@ public sealed class TestAsyncDisposableCommand : AsyncPSCmdlet public CancellationToken StoppingToken => CancelToken; } +[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedCallback")] +public sealed class TestAsyncCapturedCallbackCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + var writeOutput = CapturePipelineWriter(); + var completed = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + ThreadPool.UnsafeQueueUserWorkItem( + _ => + { + writeOutput("callback-output"); + completed.TrySetResult(true); + }, + null); + await completed.Task.ConfigureAwait(false); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCommandDetail")] +public sealed class TestAsyncCommandDetailCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteCommandDetail("worker-detail"); + WriteObject("completed"); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncReentrantPump")] public sealed class TestAsyncReentrantPumpCommand : AsyncPSCmdlet { diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index b0219fdba..8850f4779 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -486,14 +486,84 @@ public void AsyncPSCmdlet_dispose_requests_cancellation() } [Fact] - public async Task AsyncPSCmdlet_drops_pre_lifecycle_writes_from_other_threads() + public void AsyncPSCmdlet_drops_pre_lifecycle_writes_from_other_threads() { using var command = new TestAsyncDisposableCommand(); + Exception? exception = null; + var workerThreadId = 0; + using var completed = new ManualResetEventSlim(); + var thread = new Thread(() => + { + workerThreadId = Environment.CurrentManagedThreadId; + exception = Record.Exception(() => command.WriteWarning("too-early")); + completed.Set(); + }); + + thread.Start(); + + Assert.True(completed.Wait(TimeSpan.FromSeconds(5))); + thread.Join(); + Assert.NotEqual(Environment.CurrentManagedThreadId, workerThreadId); + Assert.Null(exception); + } + + [Fact] + public void AsyncPSCmdlet_finishes_disposal_when_cancellation_callbacks_throw() + { + var command = new TestAsyncDisposableCommand(); + var stoppingToken = command.StoppingToken; + using var registration = stoppingToken.Register( + static () => throw new InvalidOperationException("cancellation callback failed")); - var exception = await Record.ExceptionAsync( - () => Task.Run(() => command.WriteWarning("too-early"))); + Assert.Throws(command.Dispose); - Assert.Null(exception); + var sourceField = typeof(AsyncPSCmdlet).GetField( + "_cancelSource", + System.Reflection.BindingFlags.Instance | System.Reflection.BindingFlags.NonPublic); + var source = Assert.IsType(sourceField!.GetValue(command)); + Assert.Throws(source.Cancel); + } + + [Fact] + public void AsyncPSCmdlet_accepts_context_free_callbacks_through_a_captured_writer() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCapturedCallback", + typeof(TestAsyncCapturedCallbackCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCapturedCallback"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal("callback-output", Assert.Single(result).BaseObject); + } + + [Fact] + public void AsyncPSCmdlet_marshals_command_details_to_the_pipeline_thread() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCommandDetail", + typeof(TestAsyncCommandDetailCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCommandDetail"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal("completed", Assert.Single(result).BaseObject); } [Fact] From a8f9438e9cd18f30a3fc89b9c4bf267c6e3751a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 04:38:31 +0200 Subject: [PATCH 17/27] Harden async lifecycle thread ownership --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 272 +++++++++--------- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 8 +- PowerForge.Tests/AsyncPSCmdletTests.cs | 4 +- 3 files changed, 153 insertions(+), 131 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index 3903b805c..685064733 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -60,10 +60,12 @@ public virtual void Dispose() private bool IsPipelineThread => _pipelineThreadId != 0 && Environment.CurrentManagedThreadId == _pipelineThreadId; + private bool IsConstructionThreadOutsideAsyncHook + => Volatile.Read(ref _currentOutPipe) is null && + Environment.CurrentManagedThreadId == _constructionThreadId; + private bool CanAccessPipelineDirectly - => IsPipelineThread || - (Volatile.Read(ref _asyncLifecycleStarted) == 0 && - Environment.CurrentManagedThreadId == _constructionThreadId); + => IsPipelineThread || IsConstructionThreadOutsideAsyncHook; private void PrepareDirectPipelineAccess() { @@ -87,7 +89,9 @@ private void ValidateInteractionGeneration() var activeGeneration = Volatile.Read(ref _activeHookGeneration); var originatingGeneration = _hookGeneration.Value; - if (activeGeneration == 0 && originatingGeneration == 0 && IsPipelineThread) + if (activeGeneration == 0 && + originatingGeneration == 0 && + (IsPipelineThread || IsConstructionThreadOutsideAsyncHook)) return; if (originatingGeneration == 0 || originatingGeneration != activeGeneration) @@ -200,6 +204,7 @@ private void RunBlockInAsync(Func task) } finally { + _pipelineThreadId = 0; ExitAsyncBlock(); } } @@ -255,132 +260,141 @@ void PumpItem(PipelineItem item) return; } - switch (item.Type) + var priorItemGeneration = _hookGeneration.Value; + try { - case PipelineType.Output: - base.WriteObject(item.Value); - break; - case PipelineType.OutputEnumerate: - base.WriteObject(item.Value, enumerateCollection: true); - break; - case PipelineType.Error: - base.WriteError((ErrorRecord)item.Value!); - break; - case PipelineType.TerminatingError: - base.ThrowTerminatingError((ErrorRecord)item.Value!); - break; - case PipelineType.Warning: - base.WriteWarning((string)item.Value!); - break; - case PipelineType.Verbose: - base.WriteVerbose((string)item.Value!); - break; - case PipelineType.Debug: - base.WriteDebug((string)item.Value!); - break; - case PipelineType.Information: - base.WriteInformation((InformationRecord)item.Value!); - break; - case PipelineType.InformationWithTags: - var information = ((object MessageData, string[]? Tags))item.Value!; - base.WriteInformation( - information.MessageData, - information.Tags ?? Array.Empty()); - break; - case PipelineType.Progress: - base.WriteProgress((ProgressRecord)item.Value!); - break; - case PipelineType.CommandDetail: - base.WriteCommandDetail((string)item.Value!); - break; - case PipelineType.ShouldProcessTarget: - item.ReplyPipe!.Publish( - () => base.ShouldProcess((string)item.Value!)); - break; - case PipelineType.ShouldProcess: - var should = ((string Target, string Action))item.Value!; - item.ReplyPipe!.Publish( - () => base.ShouldProcess(should.Target, should.Action)); - break; - case PipelineType.ShouldProcessVerbose: - var verbose = ((string Description, string Warning, string Caption))item.Value!; - item.ReplyPipe!.Publish( - () => base.ShouldProcess(verbose.Description, verbose.Warning, verbose.Caption)); - break; - case PipelineType.ShouldProcessReason: - var reasonRequest = ((string Description, string Warning, string Caption))item.Value!; - item.ReplyPipe!.Publish(() => - { - var result = base.ShouldProcess( - reasonRequest.Description, - reasonRequest.Warning, - reasonRequest.Caption, - out var reason); - return (result, reason); - }); - break; - case PipelineType.ShouldContinue: - var shouldContinue = ((string Query, string Caption))item.Value!; - item.ReplyPipe!.Publish( - () => base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption)); - break; - case PipelineType.ShouldContinueAll: - var shouldContinueAll = - ((string Query, string Caption, bool YesToAll, bool NoToAll))item.Value!; - item.ReplyPipe!.Publish(() => - { - var yesToAll = shouldContinueAll.YesToAll; - var noToAll = shouldContinueAll.NoToAll; - var continueAll = base.ShouldContinue( - shouldContinueAll.Query, - shouldContinueAll.Caption, - ref yesToAll, - ref noToAll); - return (continueAll, yesToAll, noToAll); - }); - break; - case PipelineType.ShouldContinueSecurity: - var shouldContinueSecurity = - ((string Query, string Caption, bool HasSecurityImpact, bool YesToAll, bool NoToAll))item.Value!; - item.ReplyPipe!.Publish(() => - { - var yesToAll = shouldContinueSecurity.YesToAll; - var noToAll = shouldContinueSecurity.NoToAll; - var continueSecurity = base.ShouldContinue( - shouldContinueSecurity.Query, - shouldContinueSecurity.Caption, - shouldContinueSecurity.HasSecurityImpact, - ref yesToAll, - ref noToAll); - return (continueSecurity, yesToAll, noToAll); - }); - break; - case PipelineType.PromptForCredential: - var prompt = ((string Caption, string Message, string UserName, string TargetName))item.Value!; - item.ReplyPipe!.Publish( - () => Host.UI.PromptForCredential( - prompt.Caption, - prompt.Message, - prompt.UserName, - prompt.TargetName)); - break; - case PipelineType.PromptForCredentialOptions: - var promptOptions = - ((string Caption, - string Message, - string UserName, - string TargetName, - PSCredentialTypes AllowedCredentialTypes, - PSCredentialUIOptions Options))item.Value!; - item.ReplyPipe!.Publish( - () => Host.UI.PromptForCredential( - promptOptions.Caption, - promptOptions.Message, - promptOptions.UserName, - promptOptions.TargetName, - promptOptions.AllowedCredentialTypes, - promptOptions.Options)); - break; + _hookGeneration.Value = item.HookGeneration; + switch (item.Type) + { + case PipelineType.Output: + base.WriteObject(item.Value); + break; + case PipelineType.OutputEnumerate: + base.WriteObject(item.Value, enumerateCollection: true); + break; + case PipelineType.Error: + base.WriteError((ErrorRecord)item.Value!); + break; + case PipelineType.TerminatingError: + base.ThrowTerminatingError((ErrorRecord)item.Value!); + break; + case PipelineType.Warning: + base.WriteWarning((string)item.Value!); + break; + case PipelineType.Verbose: + base.WriteVerbose((string)item.Value!); + break; + case PipelineType.Debug: + base.WriteDebug((string)item.Value!); + break; + case PipelineType.Information: + base.WriteInformation((InformationRecord)item.Value!); + break; + case PipelineType.InformationWithTags: + var information = ((object MessageData, string[]? Tags))item.Value!; + base.WriteInformation( + information.MessageData, + information.Tags ?? Array.Empty()); + break; + case PipelineType.Progress: + base.WriteProgress((ProgressRecord)item.Value!); + break; + case PipelineType.CommandDetail: + base.WriteCommandDetail((string)item.Value!); + break; + case PipelineType.ShouldProcessTarget: + item.ReplyPipe!.Publish( + () => base.ShouldProcess((string)item.Value!)); + break; + case PipelineType.ShouldProcess: + var should = ((string Target, string Action))item.Value!; + item.ReplyPipe!.Publish( + () => base.ShouldProcess(should.Target, should.Action)); + break; + case PipelineType.ShouldProcessVerbose: + var verbose = ((string Description, string Warning, string Caption))item.Value!; + item.ReplyPipe!.Publish( + () => base.ShouldProcess(verbose.Description, verbose.Warning, verbose.Caption)); + break; + case PipelineType.ShouldProcessReason: + var reasonRequest = ((string Description, string Warning, string Caption))item.Value!; + item.ReplyPipe!.Publish(() => + { + var result = base.ShouldProcess( + reasonRequest.Description, + reasonRequest.Warning, + reasonRequest.Caption, + out var reason); + return (result, reason); + }); + break; + case PipelineType.ShouldContinue: + var shouldContinue = ((string Query, string Caption))item.Value!; + item.ReplyPipe!.Publish( + () => base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption)); + break; + case PipelineType.ShouldContinueAll: + var shouldContinueAll = + ((string Query, string Caption, bool YesToAll, bool NoToAll))item.Value!; + item.ReplyPipe!.Publish(() => + { + var yesToAll = shouldContinueAll.YesToAll; + var noToAll = shouldContinueAll.NoToAll; + var continueAll = base.ShouldContinue( + shouldContinueAll.Query, + shouldContinueAll.Caption, + ref yesToAll, + ref noToAll); + return (continueAll, yesToAll, noToAll); + }); + break; + case PipelineType.ShouldContinueSecurity: + var shouldContinueSecurity = + ((string Query, string Caption, bool HasSecurityImpact, bool YesToAll, bool NoToAll))item.Value!; + item.ReplyPipe!.Publish(() => + { + var yesToAll = shouldContinueSecurity.YesToAll; + var noToAll = shouldContinueSecurity.NoToAll; + var continueSecurity = base.ShouldContinue( + shouldContinueSecurity.Query, + shouldContinueSecurity.Caption, + shouldContinueSecurity.HasSecurityImpact, + ref yesToAll, + ref noToAll); + return (continueSecurity, yesToAll, noToAll); + }); + break; + case PipelineType.PromptForCredential: + var prompt = ((string Caption, string Message, string UserName, string TargetName))item.Value!; + item.ReplyPipe!.Publish( + () => Host.UI.PromptForCredential( + prompt.Caption, + prompt.Message, + prompt.UserName, + prompt.TargetName)); + break; + case PipelineType.PromptForCredentialOptions: + var promptOptions = + ((string Caption, + string Message, + string UserName, + string TargetName, + PSCredentialTypes AllowedCredentialTypes, + PSCredentialUIOptions Options))item.Value!; + item.ReplyPipe!.Publish( + () => Host.UI.PromptForCredential( + promptOptions.Caption, + promptOptions.Message, + promptOptions.UserName, + promptOptions.TargetName, + promptOptions.AllowedCredentialTypes, + promptOptions.Options)); + break; + } + } + finally + { + _hookGeneration.Value = priorItemGeneration; } } diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index 9ccb29bb2..54d3f3bb5 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -454,7 +454,10 @@ protected override async Task ProcessRecordAsync() } } -[Cmdlet(VerbsDiagnostic.Test, "AsyncReentrantPump")] +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncReentrantPump", + SupportsShouldProcess = true)] public sealed class TestAsyncReentrantPumpCommand : AsyncPSCmdlet { protected override async Task ProcessRecordAsync() @@ -475,6 +478,9 @@ public WarningEnumerable(TestAsyncReentrantPumpCommand command) public IEnumerator GetEnumerator() { + Assert.True( + _command.ShouldProcess( + "enumerated-target")); _command.WriteWarning("during-enumeration"); yield return "value"; } diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 8850f4779..c774a201e 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -579,7 +579,9 @@ public void AsyncPSCmdlet_does_not_recursively_drain_when_enumeration_writes_to_ runspace.Open(); using var powerShell = PowerShell.Create(); powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncReentrantPump"); + powerShell + .AddCommand("Test-AsyncReentrantPump") + .AddParameter("Confirm", false); var result = powerShell.Invoke(); From 0f7b869510dd200f3ab89120020beb74ecc10834 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 04:58:30 +0200 Subject: [PATCH 18/27] Handle delayed async lifecycle edge cases --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 37 ++++++++--- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 2 +- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 58 ++++++++++++++++- PowerForge.Tests/AsyncPSCmdletTests.cs | 65 +++++++++++++++++++ 4 files changed, 152 insertions(+), 10 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index 685064733..a879fb377 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -21,9 +21,19 @@ public abstract partial class AsyncPSCmdlet if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(progressRecord, PipelineType.Progress)); + _ = TryQueue(new PipelineItem(SnapshotProgressRecord(progressRecord), PipelineType.Progress)); } + private static ProgressRecord SnapshotProgressRecord(ProgressRecord progressRecord) + => new(progressRecord.ActivityId, progressRecord.Activity, progressRecord.StatusDescription) + { + CurrentOperation = progressRecord.CurrentOperation, + ParentActivityId = progressRecord.ParentActivityId, + PercentComplete = progressRecord.PercentComplete, + RecordType = progressRecord.RecordType, + SecondsRemaining = progressRecord.SecondsRemaining + }; + /// Throws when PowerShell has requested cancellation. protected internal void ThrowIfStopped() { @@ -62,7 +72,8 @@ private bool IsPipelineThread private bool IsConstructionThreadOutsideAsyncHook => Volatile.Read(ref _currentOutPipe) is null && - Environment.CurrentManagedThreadId == _constructionThreadId; + Environment.CurrentManagedThreadId == _constructionThreadId && + CommandRuntime is not null; private bool CanAccessPipelineDirectly => IsPipelineThread || IsConstructionThreadOutsideAsyncHook; @@ -211,7 +222,7 @@ private void RunBlockInAsync(Func task) private void RunBlockInAsyncCore(Func task) { - var outPipe = new BlockingCollection(); + var outPipe = new BlockingCollection(boundedCapacity: 1024); Task blockTask; var deferPipeDisposal = 0; var pipeDisposed = 0; @@ -530,11 +541,21 @@ void PumpQueuedItems() { var stopRequested = _cancelSource.IsCancellationRequested; Volatile.Write(ref deferPipeDisposal, 1); - CancelSource(); - CompleteAddingIfNeeded(outPipe); - if (blockTask.IsCompleted) - DisposePipeOnce(); - DeactivateHook(); + try + { + CancelSource(); + } + catch (AggregateException) + { + // Preserve the pipeline failure while cancellation callbacks observe the same stop. + } + finally + { + CompleteAddingIfNeeded(outPipe); + if (blockTask.IsCompleted) + DisposePipeOnce(); + DeactivateHook(); + } if (pipelineException is OperationCanceledException && stopRequested) throw new PipelineStoppedException(); diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index f6bd48ddc..5c4f9dd8a 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -110,7 +110,7 @@ private void PublishReply(Func createReply) catch (Exception exception) { TryPublish(new PipelineReply(value: null, exception)); - throw; + return; } TryPublish(reply); diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index 54d3f3bb5..44eef46e3 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -509,6 +509,51 @@ protected override Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncConstructorWrite")] +public sealed class TestAsyncConstructorWriteCommand : AsyncPSCmdlet +{ + public TestAsyncConstructorWriteCommand() + => WriteWarning("constructor output must not reach PowerShell"); + + protected override Task ProcessRecordAsync() + { + WriteObject("completed"); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncProgressSnapshot")] +public sealed class TestAsyncProgressSnapshotCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var progress = new ProgressRecord(17, "snapshot", "queued") + { + PercentComplete = 10 + }; + WriteProgress(progress); + progress.PercentComplete = 90; + WriteObject("completed"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncPipelineFailureWithThrowingCancellation")] +public sealed class TestAsyncPipelineFailureWithThrowingCancellationCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + _ = CancelToken.Register( + static () => throw new InvalidOperationException("cancellation callback failed")); + await Task.Yield(); + ThrowTerminatingError(new ErrorRecord( + new InvalidOperationException("pipeline failure"), + "PipelineFailure", + ErrorCategory.InvalidOperation, + targetObject: null)); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringHook")] public sealed class TestAsyncDisposeDuringHookCommand : AsyncPSCmdlet { @@ -574,6 +619,8 @@ public sealed class ChoiceHost(bool approved, Exception? promptFailure = null) : public override string Name => nameof(ChoiceHost); public override Version Version => new(1, 0); public override PSHostUserInterface UI => _ui; + + public IReadOnlyList ProgressRecords => _ui.ProgressRecords; public override CultureInfo CurrentCulture => CultureInfo.InvariantCulture; public override CultureInfo CurrentUICulture => CultureInfo.InvariantCulture; public override void EnterNestedPrompt() { } @@ -585,6 +632,10 @@ public override void SetShouldExit(int exitCode) { } public sealed class ChoiceHostUserInterface(bool approved, Exception? promptFailure) : PSHostUserInterface { + private readonly List _progressRecords = new(); + + public IReadOnlyList ProgressRecords => _progressRecords; + public override PSHostRawUserInterface RawUI => null!; public override int PromptForChoice( @@ -601,7 +652,12 @@ public override void Write(ConsoleColor foregroundColor, ConsoleColor background public override void WriteLine(string value) { } public override void WriteErrorLine(string value) { } public override void WriteDebugLine(string message) { } - public override void WriteProgress(long sourceId, ProgressRecord record) { } + public override void WriteProgress(long sourceId, ProgressRecord record) + => _progressRecords.Add( + new ProgressRecord(record.ActivityId, record.Activity, record.StatusDescription) + { + PercentComplete = record.PercentComplete + }); public override void WriteVerboseLine(string message) { } public override void WriteWarningLine(string message) { } diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index c774a201e..5a4cf2516 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -618,6 +618,71 @@ public void AsyncPSCmdlet_allows_ShouldProcess_before_the_async_base_hook_starts item => Assert.True((bool)item.BaseObject)); } + [Fact] + public void AsyncPSCmdlet_drops_constructor_writes_before_pipeline_startup() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncConstructorWrite", + typeof(TestAsyncConstructorWriteCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncConstructorWrite"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Empty(powerShell.Streams.Warning); + Assert.Equal("completed", Assert.Single(result).BaseObject); + } + + [Fact] + public void AsyncPSCmdlet_snapshots_mutable_progress_before_queueing() + { + var sessionState = InitialSessionState.Create(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncProgressSnapshot", + typeof(TestAsyncProgressSnapshotCommand), + helpFileName: null)); + + var host = new ChoiceHost(approved: true); + using var runspace = RunspaceFactory.CreateRunspace(host, sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncProgressSnapshot"); + + var result = powerShell.Invoke(); + + Assert.Equal("completed", Assert.Single(result).BaseObject); + Assert.Equal(10, Assert.Single(host.ProgressRecords).PercentComplete); + } + + [Fact] + public void AsyncPSCmdlet_preserves_pipeline_failure_when_cancellation_callbacks_throw() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncPipelineFailureWithThrowingCancellation", + typeof(TestAsyncPipelineFailureWithThrowingCancellationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncPipelineFailureWithThrowingCancellation"); + + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.StartsWith("PipelineFailure,", exception.ErrorRecord.FullyQualifiedErrorId, StringComparison.Ordinal); + Assert.Equal("pipeline failure", exception.InnerException?.Message); + } + [Fact] public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_the_async_hook_exits() { From 44473c982df1f577b3dee450ea0a446dc4480342 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 05:42:56 +0200 Subject: [PATCH 19/27] Harden AsyncPSCmdlet lifecycle and callback transport --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 69 ++++++++-- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 129 +++++++++++++++--- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 91 +++++++++++- PowerForge.Tests/AsyncPSCmdletTests.cs | 74 ++++++++-- 4 files changed, 325 insertions(+), 38 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index a879fb377..07675ff18 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -11,9 +11,10 @@ public abstract partial class AsyncPSCmdlet /// Thread-safe progress bridge for asynchronous cmdlet code. public new void WriteProgress(ProgressRecord progressRecord) { + ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteProgress(progressRecord); return; } @@ -44,17 +45,21 @@ protected internal void ThrowIfStopped() /// public virtual void Dispose() { + bool cancelActiveBlocks; lock (_lifecycleLock) { if (_disposeRequested) return; _disposeRequested = true; + cancelActiveBlocks = _activeBlocks != 0; + Volatile.Write(ref _asyncLifecycleCompleted, 1); } try { - CancelSource(); + if (cancelActiveBlocks) + CancelSource(); } finally { @@ -72,25 +77,39 @@ private bool IsPipelineThread private bool IsConstructionThreadOutsideAsyncHook => Volatile.Read(ref _currentOutPipe) is null && + Volatile.Read(ref _asyncLifecycleCompleted) == 0 && Environment.CurrentManagedThreadId == _constructionThreadId && CommandRuntime is not null; private bool CanAccessPipelineDirectly => IsPipelineThread || IsConstructionThreadOutsideAsyncHook; - private void PrepareDirectPipelineAccess() + private IDisposable EnterDirectPipelineAccess() { ThrowIfStopped(); + ValidateInteractionGeneration(); if (IsPipelineThread) + { Volatile.Read(ref _pumpQueuedItems)?.Invoke(); + return new SynchronizationContextScope( + Volatile.Read(ref _pipelineSynchronizationContext)); + } + + return new SynchronizationContextScope(SynchronizationContext.Current); } - private void PrepareDirectPipelineInteraction() + private IDisposable EnterDirectPipelineInteraction() { ThrowIfStopped(); ValidateInteractionGeneration(); if (IsPipelineThread) + { Volatile.Read(ref _pumpQueuedItems)?.Invoke(); + return new SynchronizationContextScope( + Volatile.Read(ref _pipelineSynchronizationContext)); + } + + return new SynchronizationContextScope(SynchronizationContext.Current); } private void ValidateInteractionGeneration() @@ -180,6 +199,21 @@ private void GetBlockTaskResult(Task blockTask) return value => _ = TryQueue(new PipelineItem(value, pipelineType, hookGeneration: hookGeneration)); } + /// + /// Captures lifecycle-bound typed stream writers for callbacks that do not flow execution context. + /// + protected CapturedPipelineStreams CapturePipelineStreams() + { + var hookGeneration = _hookGeneration.Value; + if (hookGeneration == 0) + { + throw new InvalidOperationException( + "Lifecycle-bound pipeline streams can only be captured from an asynchronous PowerShell hook."); + } + + return new CapturedPipelineStreams(this, hookGeneration); + } + private bool TryQueue(PipelineItem item) { item.BindToHook(_hookGeneration.Value); @@ -202,6 +236,9 @@ private bool TryQueue(PipelineItem item) } catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) { + if (item.HookGeneration != 0) + throw new PipelineStoppedException(); + return false; } } @@ -222,7 +259,10 @@ private void RunBlockInAsync(Func task) private void RunBlockInAsyncCore(Func task) { - var outPipe = new BlockingCollection(boundedCapacity: 1024); + // The transport must remain lossless and non-blocking. The pipeline thread can enumerate + // user objects or invoke a host that waits for the same background producer that is writing + // here; applying bounded backpressure would deadlock both sides. + var outPipe = new BlockingCollection(); Task blockTask; var deferPipeDisposal = 0; var pipeDisposed = 0; @@ -435,6 +475,7 @@ void PumpQueuedItems() var priorHookGeneration = _hookGeneration.Value; try { + Volatile.Write(ref _pipelineSynchronizationContext, synchronizationContext); SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); _hookGeneration.Value = hookGeneration; if (TaskScheduler.Current == TaskScheduler.Default) @@ -477,6 +518,7 @@ void PumpQueuedItems() { _hookGeneration.Value = priorHookGeneration; SynchronizationContext.SetSynchronizationContext(synchronizationContext); + Volatile.Write(ref _pipelineSynchronizationContext, null); } if (blockTask.IsCompleted) @@ -484,7 +526,6 @@ void PumpQueuedItems() if (blockTask.IsFaulted) _ = blockTask.Exception; - CompleteAddingIfNeeded(outPipe); try { PumpQueuedItems(); @@ -511,9 +552,11 @@ void PumpQueuedItems() if (completed.IsFaulted) _ = completed.Exception; - ClearPipes(); if (Volatile.Read(ref deferPipeDisposal) != 0) + { + ClearPipes(); DisposePipeOnce(); + } } finally { @@ -532,10 +575,13 @@ void PumpQueuedItems() try { - foreach (var item in outPipe.GetConsumingEnumerable(CancelToken)) + while (!blockTask.IsCompleted || outPipe.Count != 0) { - PumpItem(item); + if (outPipe.TryTake(out var item, millisecondsTimeout: 50, CancelToken)) + PumpItem(item); } + + ClearPipes(); } catch (Exception pipelineException) { @@ -608,6 +654,11 @@ private void CancelSource() { _cancelSource.Cancel(); } + catch (AggregateException) + { + // Cancellation callbacks are third-party code. A failing callback must not escape + // StopProcessing or mask the pipeline failure that initiated cancellation. + } catch (ObjectDisposedException) { // Disposal may race a late StopProcessing callback after all async hooks have exited. diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 5c4f9dd8a..1025f5ec4 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -38,6 +38,20 @@ protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQu => TryExecuteTask(task); } + private sealed class SynchronizationContextScope : IDisposable + { + private readonly SynchronizationContext? _previous; + + public SynchronizationContextScope(SynchronizationContext? replacement) + { + _previous = SynchronizationContext.Current; + SynchronizationContext.SetSynchronizationContext(replacement); + } + + public void Dispose() + => SynchronizationContext.SetSynchronizationContext(_previous); + } + private enum PipelineType { Output, @@ -187,6 +201,64 @@ public void BindToHook(long hookGeneration) } } + /// + /// Lifecycle-bound stream writers for callbacks that do not flow the hook execution context. + /// + protected sealed class CapturedPipelineStreams + { + private readonly long _hookGeneration; + private readonly AsyncPSCmdlet _owner; + + internal CapturedPipelineStreams(AsyncPSCmdlet owner, long hookGeneration) + { + _owner = owner; + _hookGeneration = hookGeneration; + } + + /// Queues an output record for the originating hook. + public void WriteObject(object? value, bool enumerateCollection = false) + => Queue( + value, + enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output); + + /// Queues an error record for the originating hook. + public void WriteError(ErrorRecord errorRecord) + => Queue(errorRecord, PipelineType.Error); + + /// Queues a warning record for the originating hook. + public void WriteWarning(string message) + => Queue(message, PipelineType.Warning); + + /// Queues a verbose record for the originating hook. + public void WriteVerbose(string message) + => Queue(message, PipelineType.Verbose); + + /// Queues a debug record for the originating hook. + public void WriteDebug(string message) + => Queue(message, PipelineType.Debug); + + /// Queues an information record for the originating hook. + public void WriteInformation(InformationRecord informationRecord) + => Queue(informationRecord, PipelineType.Information); + + /// Queues tagged information for the originating hook. + public void WriteInformation(object messageData, string[]? tags) + => Queue( + (messageData, tags is null ? null : (string[])tags.Clone()), + PipelineType.InformationWithTags); + + /// Queues a progress record for the originating hook. + public void WriteProgress(ProgressRecord progressRecord) + => Queue(SnapshotProgressRecord(progressRecord), PipelineType.Progress); + + /// Queues command-detail text for the originating hook. + public void WriteCommandDetail(string text) + => Queue(text, PipelineType.CommandDetail); + + private void Queue(object? value, PipelineType type) + => _ = _owner.TryQueue(new PipelineItem(value, type, hookGeneration: _hookGeneration)); + } + private readonly CancellationTokenSource _cancelSource = new(); private readonly AsyncLocal _hookGeneration = new(); private readonly int _constructionThreadId = Environment.CurrentManagedThreadId; @@ -195,11 +267,13 @@ public void BindToHook(long hookGeneration) private static readonly TaskScheduler HookTaskScheduler = new AsyncHookTaskScheduler(); private BlockingCollection? _currentOutPipe; private Action? _pumpQueuedItems; + private SynchronizationContext? _pipelineSynchronizationContext; private long _activeHookGeneration; private long _nextHookGeneration; private bool _cancelSourceDisposed; private bool _disposeRequested; private int _activeBlocks; + private int _asyncLifecycleCompleted; private int _asyncLifecycleStarted; private int _pipelineThreadId; @@ -224,7 +298,16 @@ protected virtual Task ProcessRecordAsync() /// protected override void EndProcessing() - => RunBlockInAsync(EndProcessingAsync); + { + try + { + RunBlockInAsync(EndProcessingAsync); + } + finally + { + Volatile.Write(ref _asyncLifecycleCompleted, 1); + } + } /// Asynchronous end hook. protected virtual Task EndProcessingAsync() @@ -240,7 +323,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldProcess(target ?? string.Empty); } @@ -253,7 +336,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldProcess(target ?? string.Empty, action); } @@ -266,7 +349,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldProcess(verboseDescription, verboseWarning, caption); } @@ -285,7 +368,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldProcess(verboseDescription, verboseWarning, caption, out shouldProcessReason); } @@ -302,7 +385,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldContinue(query, caption); } @@ -315,7 +398,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldContinue(query, caption, ref yesToAll, ref noToAll); } @@ -338,7 +421,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldContinue(query, caption, hasSecurityImpact, ref yesToAll, ref noToAll); } @@ -356,7 +439,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return Host.UI.PromptForCredential(caption, message, userName, targetName); } @@ -377,7 +460,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return Host.UI.PromptForCredential( caption, message, @@ -399,9 +482,10 @@ protected override void StopProcessing() /// Thread-safe output bridge for asynchronous cmdlet code. public new void WriteObject(object? sendToPipeline, bool enumerateCollection) { + ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteObject(sendToPipeline, enumerateCollection); return; } @@ -417,9 +501,10 @@ protected override void StopProcessing() /// Thread-safe error bridge for asynchronous cmdlet code. public new void WriteError(ErrorRecord errorRecord) { + ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteError(errorRecord); return; } @@ -436,7 +521,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.ThrowTerminatingError(errorRecord); return; } @@ -454,9 +539,10 @@ protected override void StopProcessing() /// Thread-safe warning bridge for asynchronous cmdlet code. public new void WriteWarning(string message) { + ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteWarning(message); return; } @@ -470,9 +556,10 @@ protected override void StopProcessing() /// Thread-safe verbose bridge for asynchronous cmdlet code. public new void WriteVerbose(string message) { + ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteVerbose(message); return; } @@ -486,9 +573,10 @@ protected override void StopProcessing() /// Thread-safe debug bridge for asynchronous cmdlet code. public new void WriteDebug(string message) { + ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteDebug(message); return; } @@ -502,9 +590,10 @@ protected override void StopProcessing() /// Thread-safe command-detail bridge for asynchronous cmdlet code. public new void WriteCommandDetail(string text) { + ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteCommandDetail(text); return; } @@ -518,9 +607,10 @@ protected override void StopProcessing() /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(InformationRecord informationRecord) { + ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteInformation(informationRecord); return; } @@ -534,9 +624,10 @@ protected override void StopProcessing() /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(object messageData, string[]? tags) { + ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteInformation(messageData, tags ?? Array.Empty()); return; } diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index 44eef46e3..f50fbfb9d 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -266,15 +266,19 @@ protected override async Task ProcessRecordAsync() public sealed class TestAsyncCancellationWriteCommand : AsyncPSCmdlet { private static ManualResetEventSlim _started = new(); + private static ManualResetEventSlim _writeAttempted = new(); public static ManualResetEventSlim Started => _started; + public static ManualResetEventSlim WriteAttempted => _writeAttempted; public static Exception? BackgroundWriteException { get; private set; } public static void Reset() { _started.Dispose(); + _writeAttempted.Dispose(); _started = new ManualResetEventSlim(); + _writeAttempted = new ManualResetEventSlim(); BackgroundWriteException = null; } @@ -296,6 +300,10 @@ protected override async Task ProcessRecordAsync() { BackgroundWriteException = exception; } + finally + { + _writeAttempted.Set(); + } } } } @@ -423,6 +431,85 @@ protected override Task ProcessRecordAsync() public sealed class TestAsyncDisposableCommand : AsyncPSCmdlet { public CancellationToken StoppingToken => CancelToken; + + public void InvokeStopProcessing() + => base.StopProcessing(); +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncLargeQueuedOutput")] +public sealed class TestAsyncLargeQueuedOutputCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + using var completed = new ManualResetEventSlim(); + Exception? workerException = null; + var thread = new Thread(() => + { + try + { + for (var i = 0; i < 2048; i++) + WriteWarning($"queued-{i}"); + } + catch (Exception exception) + { + workerException = exception; + } + finally + { + completed.Set(); + } + }); + + thread.Start(); + Assert.True(completed.Wait(TimeSpan.FromSeconds(5)), "The producer blocked while filling the pipeline transport."); + thread.Join(); + if (workerException is not null) + throw workerException; + + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDirectContext")] +public sealed class TestAsyncDirectContextCommand : AsyncPSCmdlet +{ + public static SynchronizationContext? HostContext { get; private set; } + + protected override void ProcessRecord() + { + var previousContext = SynchronizationContext.Current; + HostContext = new ForwardingSynchronizationContext(); + SynchronizationContext.SetSynchronizationContext(HostContext); + try + { + base.ProcessRecord(); + } + finally + { + SynchronizationContext.SetSynchronizationContext(previousContext); + } + } + + protected override Task ProcessRecordAsync() + { + WriteObject("context-output"); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "ObserveContext")] +public sealed class TestObserveContextCommand : PSCmdlet +{ + [Parameter(ValueFromPipeline = true)] + public object? InputObject { get; set; } + + public static SynchronizationContext? ObservedContext { get; private set; } + + protected override void ProcessRecord() + { + ObservedContext = SynchronizationContext.Current; + WriteObject(InputObject); + } } [Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedCallback")] @@ -431,11 +518,13 @@ public sealed class TestAsyncCapturedCallbackCommand : AsyncPSCmdlet protected override async Task ProcessRecordAsync() { var writeOutput = CapturePipelineWriter(); + var streams = CapturePipelineStreams(); var completed = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); ThreadPool.UnsafeQueueUserWorkItem( _ => { writeOutput("callback-output"); + streams.WriteWarning("callback-warning"); completed.TrySetResult(true); }, null); @@ -481,7 +570,7 @@ public IEnumerator GetEnumerator() Assert.True( _command.ShouldProcess( "enumerated-target")); - _command.WriteWarning("during-enumeration"); + Task.Run(() => _command.WriteWarning("during-enumeration")).GetAwaiter().GetResult(); yield return "value"; } diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 5a4cf2516..88d841b9b 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -291,7 +291,7 @@ public void AsyncPSCmdlet_translates_cancellation_to_a_pipeline_stop() } [Fact] - public void AsyncPSCmdlet_drops_background_stream_writes_during_pipeline_stop() + public void AsyncPSCmdlet_propagates_pipeline_stop_to_background_stream_writes() { var sessionState = InitialSessionState.CreateDefault(); sessionState.Commands.Add(new SessionStateCmdletEntry( @@ -313,8 +313,12 @@ public void AsyncPSCmdlet_drops_background_stream_writes_during_pipeline_stop() powerShell.Stop(); Assert.Throws(() => powerShell.EndInvoke(invocation)); + Assert.True( + TestAsyncCancellationWriteCommand.WriteAttempted.Wait(TimeSpan.FromSeconds(5)), + "The background hook did not observe the stop in time."); - Assert.Null(TestAsyncCancellationWriteCommand.BackgroundWriteException); + Assert.IsType( + TestAsyncCancellationWriteCommand.BackgroundWriteException); } [Fact] @@ -475,14 +479,14 @@ TestAsyncLateInteractionCommand.LateInteractionException is } [Fact] - public void AsyncPSCmdlet_dispose_requests_cancellation() + public void AsyncPSCmdlet_dispose_without_an_active_hook_does_not_signal_stop() { using var command = new TestAsyncDisposableCommand(); var stoppingToken = command.StoppingToken; command.Dispose(); - Assert.True(stoppingToken.IsCancellationRequested); + Assert.False(stoppingToken.IsCancellationRequested); } [Fact] @@ -515,8 +519,10 @@ public void AsyncPSCmdlet_finishes_disposal_when_cancellation_callbacks_throw() using var registration = stoppingToken.Register( static () => throw new InvalidOperationException("cancellation callback failed")); - Assert.Throws(command.Dispose); + command.InvokeStopProcessing(); + command.Dispose(); + Assert.True(stoppingToken.IsCancellationRequested); var sourceField = typeof(AsyncPSCmdlet).GetField( "_cancelSource", System.Reflection.BindingFlags.Instance | System.Reflection.BindingFlags.NonPublic); @@ -524,6 +530,54 @@ public void AsyncPSCmdlet_finishes_disposal_when_cancellation_callbacks_throw() Assert.Throws(source.Cancel); } + [Fact] + public void AsyncPSCmdlet_does_not_deadlock_when_a_synchronous_hook_queues_more_than_the_transport_capacity() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncLargeQueuedOutput", + typeof(TestAsyncLargeQueuedOutputCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncLargeQueuedOutput"); + + powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal(2048, powerShell.Streams.Warning.Count); + } + + [Fact] + public void AsyncPSCmdlet_restores_the_host_context_around_direct_pipeline_writes() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDirectContext", + typeof(TestAsyncDirectContextCommand), + helpFileName: null)); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-ObserveContext", + typeof(TestObserveContextCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncDirectContext") + .AddCommand("Test-ObserveContext"); + + var result = powerShell.Invoke(); + + Assert.Equal("context-output", Assert.Single(result).BaseObject); + Assert.Same(TestAsyncDirectContextCommand.HostContext, TestObserveContextCommand.ObservedContext); + } + [Fact] public void AsyncPSCmdlet_accepts_context_free_callbacks_through_a_captured_writer() { @@ -543,6 +597,9 @@ public void AsyncPSCmdlet_accepts_context_free_callbacks_through_a_captured_writ Assert.False(powerShell.HadErrors); Assert.Equal("callback-output", Assert.Single(result).BaseObject); + Assert.Equal( + "callback-warning", + Assert.Single(powerShell.Streams.Warning).Message); } [Fact] @@ -586,10 +643,9 @@ public void AsyncPSCmdlet_does_not_recursively_drain_when_enumeration_writes_to_ var result = powerShell.Invoke(); Assert.Collection(result, item => Assert.Equal("value", item.BaseObject)); - Assert.Collection( - powerShell.Streams.Warning, - warning => Assert.Equal("during-enumeration", warning.Message), - warning => Assert.Equal("after-enumeration", warning.Message)); + Assert.Equal( + ["after-enumeration", "during-enumeration"], + powerShell.Streams.Warning.Select(static warning => warning.Message).Order()); } [Fact] From 1acf573b03a7a91725565fa4b1208ae2e49a252c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 06:25:28 +0200 Subject: [PATCH 20/27] Close AsyncPSCmdlet producer and host-context races --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 105 ++++++++++++++++-- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 33 +++++- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 74 +++++++++++- PowerForge.Tests/AsyncPSCmdletTests.cs | 39 ++++++- 4 files changed, 231 insertions(+), 20 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index 07675ff18..c2bfeb0c3 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -1,6 +1,7 @@ using System; using System.Collections.Concurrent; using System.Management.Automation; +using System.Reflection; using System.Threading; using System.Threading.Tasks; @@ -25,8 +26,15 @@ public abstract partial class AsyncPSCmdlet _ = TryQueue(new PipelineItem(SnapshotProgressRecord(progressRecord), PipelineType.Progress)); } + private static readonly PropertyInfo? ProgressTotalProperty = + typeof(ProgressRecord).GetProperty("Total", BindingFlags.Instance | BindingFlags.Public); + private static ProgressRecord SnapshotProgressRecord(ProgressRecord progressRecord) - => new(progressRecord.ActivityId, progressRecord.Activity, progressRecord.StatusDescription) + { + var snapshot = new ProgressRecord( + progressRecord.ActivityId, + progressRecord.Activity, + progressRecord.StatusDescription) { CurrentOperation = progressRecord.CurrentOperation, ParentActivityId = progressRecord.ParentActivityId, @@ -35,6 +43,16 @@ private static ProgressRecord SnapshotProgressRecord(ProgressRecord progressReco SecondsRemaining = progressRecord.SecondsRemaining }; + if (ProgressTotalProperty is { CanRead: true, CanWrite: true }) + { + ProgressTotalProperty.SetValue( + snapshot, + ProgressTotalProperty.GetValue(progressRecord)); + } + + return snapshot; + } + /// Throws when PowerShell has requested cancellation. protected internal void ThrowIfStopped() { @@ -90,9 +108,18 @@ private IDisposable EnterDirectPipelineAccess() ValidateInteractionGeneration(); if (IsPipelineThread) { - Volatile.Read(ref _pumpQueuedItems)?.Invoke(); - return new SynchronizationContextScope( + var pipelineContext = new SynchronizationContextScope( Volatile.Read(ref _pipelineSynchronizationContext)); + try + { + Volatile.Read(ref _pumpQueuedItems)?.Invoke(); + return pipelineContext; + } + catch + { + pipelineContext.Dispose(); + throw; + } } return new SynchronizationContextScope(SynchronizationContext.Current); @@ -104,9 +131,18 @@ private IDisposable EnterDirectPipelineInteraction() ValidateInteractionGeneration(); if (IsPipelineThread) { - Volatile.Read(ref _pumpQueuedItems)?.Invoke(); - return new SynchronizationContextScope( + var pipelineContext = new SynchronizationContextScope( Volatile.Read(ref _pipelineSynchronizationContext)); + try + { + Volatile.Read(ref _pumpQueuedItems)?.Invoke(); + return pipelineContext; + } + catch + { + pipelineContext.Dispose(); + throw; + } } return new SynchronizationContextScope(SynchronizationContext.Current); @@ -196,7 +232,12 @@ private void GetBlockTaskResult(Task blockTask) } var pipelineType = enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output; - return value => _ = TryQueue(new PipelineItem(value, pipelineType, hookGeneration: hookGeneration)); + return value => _ = TryQueue( + new PipelineItem( + value, + pipelineType, + hookGeneration: hookGeneration, + dropOnStop: true)); } /// @@ -217,6 +258,17 @@ protected CapturedPipelineStreams CapturePipelineStreams() private bool TryQueue(PipelineItem item) { item.BindToHook(_hookGeneration.Value); + var pumpLease = _pipelinePumpLease.Value; + if (item.HookGeneration != 0 && + item.HookGeneration != Volatile.Read(ref _acceptingHookWritesGeneration) && + (pumpLease is null || + !pumpLease.IsActive || + item.HookGeneration != pumpLease.Generation)) + { + item.ReplyPipe?.Reject(); + return false; + } + var outPipe = Volatile.Read(ref _currentOutPipe); if (outPipe is null) return false; @@ -236,7 +288,7 @@ private bool TryQueue(PipelineItem item) } catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) { - if (item.HookGeneration != 0) + if (item.HookGeneration != 0 && !item.DropOnStop) throw new PipelineStoppedException(); return false; @@ -268,11 +320,14 @@ private void RunBlockInAsyncCore(Func task) var pipeDisposed = 0; var pumpingQueuedItems = 0; var hookGeneration = Interlocked.Increment(ref _nextHookGeneration); + var synchronizationContext = SynchronizationContext.Current; void ClearPipes() { + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); Volatile.Write(ref _pumpQueuedItems, null); _ = Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe); + _ = Interlocked.CompareExchange(ref _pipelineSynchronizationContext, null, synchronizationContext); CompleteAddingIfNeeded(outPipe); } @@ -312,9 +367,12 @@ void PumpItem(PipelineItem item) } var priorItemGeneration = _hookGeneration.Value; + var priorPumpLease = _pipelinePumpLease.Value; + var pumpLease = new PipelinePumpLease(item.HookGeneration); try { _hookGeneration.Value = item.HookGeneration; + _pipelinePumpLease.Value = pumpLease; switch (item.Type) { case PipelineType.Output: @@ -441,10 +499,14 @@ void PumpItem(PipelineItem item) promptOptions.AllowedCredentialTypes, promptOptions.Options)); break; + case PipelineType.HookCompleted: + break; } } finally { + pumpLease.Close(); + _pipelinePumpLease.Value = priorPumpLease; _hookGeneration.Value = priorItemGeneration; } } @@ -468,10 +530,10 @@ void PumpQueuedItems() Volatile.Write(ref _asyncLifecycleStarted, 1); _pipelineThreadId = Environment.CurrentManagedThreadId; Volatile.Write(ref _activeHookGeneration, hookGeneration); + Volatile.Write(ref _acceptingHookWritesGeneration, hookGeneration); Volatile.Write(ref _pumpQueuedItems, PumpQueuedItems); Volatile.Write(ref _currentOutPipe, outPipe); - var synchronizationContext = SynchronizationContext.Current; var priorHookGeneration = _hookGeneration.Value; try { @@ -518,11 +580,11 @@ void PumpQueuedItems() { _hookGeneration.Value = priorHookGeneration; SynchronizationContext.SetSynchronizationContext(synchronizationContext); - Volatile.Write(ref _pipelineSynchronizationContext, null); } if (blockTask.IsCompleted) { + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); if (blockTask.IsFaulted) _ = blockTask.Exception; @@ -549,9 +611,31 @@ void PumpQueuedItems() { try { + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); if (completed.IsFaulted) _ = completed.Exception; + try + { + if (!outPipe.IsAddingCompleted) + { + outPipe.Add( + new PipelineItem( + value: null, + PipelineType.HookCompleted, + hookGeneration: hookGeneration, + dropOnStop: true)); + } + } + catch (ObjectDisposedException) + { + // A pipeline failure may dispose the transport before the hook completes. + } + catch (InvalidOperationException) + { + // The pipeline completed adding while the hook completion was published. + } + if (Volatile.Read(ref deferPipeDisposal) != 0) { ClearPipes(); @@ -577,8 +661,7 @@ void PumpQueuedItems() { while (!blockTask.IsCompleted || outPipe.Count != 0) { - if (outPipe.TryTake(out var item, millisecondsTimeout: 50, CancelToken)) - PumpItem(item); + PumpItem(outPipe.Take(CancelToken)); } ClearPipes(); diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 1025f5ec4..d5bd89480 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -73,7 +73,8 @@ private enum PipelineType ShouldContinueAll, ShouldContinueSecurity, PromptForCredential, - PromptForCredentialOptions + PromptForCredentialOptions, + HookCompleted } private sealed class PipelineReply @@ -178,12 +179,14 @@ public PipelineItem( object? value, PipelineType type, PipelineReplyChannel? replyPipe = null, - long hookGeneration = 0) + long hookGeneration = 0, + bool dropOnStop = false) { Value = value; Type = type; ReplyPipe = replyPipe; HookGeneration = hookGeneration; + DropOnStop = dropOnStop; } public object? Value { get; } @@ -194,6 +197,8 @@ public PipelineItem( public long HookGeneration { get; private set; } + public bool DropOnStop { get; } + public void BindToHook(long hookGeneration) { if (HookGeneration == 0) @@ -201,6 +206,21 @@ public void BindToHook(long hookGeneration) } } + private sealed class PipelinePumpLease + { + private int _active = 1; + + public PipelinePumpLease(long generation) + => Generation = generation; + + public long Generation { get; } + + public bool IsActive => Volatile.Read(ref _active) != 0; + + public void Close() + => Volatile.Write(ref _active, 0); + } + /// /// Lifecycle-bound stream writers for callbacks that do not flow the hook execution context. /// @@ -256,11 +276,17 @@ public void WriteCommandDetail(string text) => Queue(text, PipelineType.CommandDetail); private void Queue(object? value, PipelineType type) - => _ = _owner.TryQueue(new PipelineItem(value, type, hookGeneration: _hookGeneration)); + => _ = _owner.TryQueue( + new PipelineItem( + value, + type, + hookGeneration: _hookGeneration, + dropOnStop: true)); } private readonly CancellationTokenSource _cancelSource = new(); private readonly AsyncLocal _hookGeneration = new(); + private readonly AsyncLocal _pipelinePumpLease = new(); private readonly int _constructionThreadId = Environment.CurrentManagedThreadId; private readonly object _lifecycleLock = new(); private static readonly SynchronizationContext HookSynchronizationContext = new AsyncHookSynchronizationContext(); @@ -269,6 +295,7 @@ private void Queue(object? value, PipelineType type) private Action? _pumpQueuedItems; private SynchronizationContext? _pipelineSynchronizationContext; private long _activeHookGeneration; + private long _acceptingHookWritesGeneration; private long _nextHookGeneration; private bool _cancelSourceDisposed; private bool _disposeRequested; diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index f50fbfb9d..2eadbf7f5 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -308,6 +308,61 @@ protected override async Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedCancellationWrite")] +public sealed class TestAsyncCapturedCancellationWriteCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _started = new(); + private static ManualResetEventSlim _writeAttempted = new(); + + public static ManualResetEventSlim Started => _started; + public static ManualResetEventSlim WriteAttempted => _writeAttempted; + + public static Exception? BackgroundWriteException { get; private set; } + + public static void Reset() + { + _started.Dispose(); + _writeAttempted.Dispose(); + _started = new ManualResetEventSlim(); + _writeAttempted = new ManualResetEventSlim(); + BackgroundWriteException = null; + } + + protected override async Task ProcessRecordAsync() + { + var streams = CapturePipelineStreams(); + _started.Set(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } + catch (OperationCanceledException) when (CancelToken.IsCancellationRequested) + { + var completed = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + ThreadPool.UnsafeQueueUserWorkItem( + _ => + { + try + { + streams.WriteWarning("cancelled callback"); + } + catch (Exception exception) + { + BackgroundWriteException = exception; + } + finally + { + _writeAttempted.Set(); + completed.TrySetResult(true); + } + }, + null); + await completed.Task.ConfigureAwait(false); + } + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncTerminatingError")] public sealed class TestAsyncTerminatingErrorCommand : AsyncPSCmdlet { @@ -621,8 +676,21 @@ protected override async Task ProcessRecordAsync() { PercentComplete = 10 }; + var totalProperty = progress.GetType().GetProperty("Total"); + if (totalProperty is not null) + { + totalProperty.SetValue( + progress, + Convert.ChangeType(42, totalProperty.PropertyType, CultureInfo.InvariantCulture)); + } WriteProgress(progress); progress.PercentComplete = 90; + if (totalProperty is not null) + { + totalProperty.SetValue( + progress, + Convert.ChangeType(99, totalProperty.PropertyType, CultureInfo.InvariantCulture)); + } WriteObject("completed"); } } @@ -742,11 +810,7 @@ public override void WriteLine(string value) { } public override void WriteErrorLine(string value) { } public override void WriteDebugLine(string message) { } public override void WriteProgress(long sourceId, ProgressRecord record) - => _progressRecords.Add( - new ProgressRecord(record.ActivityId, record.Activity, record.StatusDescription) - { - PercentComplete = record.PercentComplete - }); + => _progressRecords.Add(record); public override void WriteVerboseLine(string message) { } public override void WriteWarningLine(string message) { } diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 88d841b9b..03e947f29 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -321,6 +321,35 @@ public void AsyncPSCmdlet_propagates_pipeline_stop_to_background_stream_writes() TestAsyncCancellationWriteCommand.BackgroundWriteException); } + [Fact] + public void AsyncPSCmdlet_drops_captured_callback_writes_after_pipeline_stop() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCapturedCancellationWrite", + typeof(TestAsyncCapturedCancellationWriteCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCapturedCancellationWrite"); + TestAsyncCapturedCancellationWriteCommand.Reset(); + + var invocation = powerShell.BeginInvoke(); + Assert.True( + TestAsyncCapturedCancellationWriteCommand.Started.Wait(TimeSpan.FromSeconds(5)), + "The asynchronous cmdlet did not start in time."); + + powerShell.Stop(); + Assert.Throws(() => powerShell.EndInvoke(invocation)); + Assert.True( + TestAsyncCapturedCancellationWriteCommand.WriteAttempted.Wait(TimeSpan.FromSeconds(5)), + "The captured callback did not observe the stop in time."); + Assert.Null(TestAsyncCapturedCancellationWriteCommand.BackgroundWriteException); + } + [Fact] public void AsyncPSCmdlet_marshals_terminating_errors_to_the_pipeline_thread() { @@ -715,7 +744,15 @@ public void AsyncPSCmdlet_snapshots_mutable_progress_before_queueing() var result = powerShell.Invoke(); Assert.Equal("completed", Assert.Single(result).BaseObject); - Assert.Equal(10, Assert.Single(host.ProgressRecords).PercentComplete); + var captured = Assert.Single(host.ProgressRecords); + Assert.Equal(10, captured.PercentComplete); + var totalProperty = captured.GetType().GetProperty("Total"); + if (totalProperty is not null) + { + Assert.Equal( + Convert.ChangeType(42, totalProperty.PropertyType, CultureInfo.InvariantCulture), + totalProperty.GetValue(captured)); + } } [Fact] From 4a29240b385eba48fb051839f7b0ec6b99ffc1b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 06:49:55 +0200 Subject: [PATCH 21/27] Harden async lifecycle completion and admission --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 210 ++++++++++-------- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 14 +- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 181 +++++++++++++++ PowerForge.Tests/AsyncPSCmdletTests.cs | 118 ++++++++++ 4 files changed, 418 insertions(+), 105 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index c2bfeb0c3..d1317b895 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -86,12 +86,19 @@ public virtual void Dispose() DisposeCancelSourceIfInactive(); } - _pipelineThreadId = 0; + Volatile.Write(ref _pipelineThreadId, 0); } } private bool IsPipelineThread - => _pipelineThreadId != 0 && Environment.CurrentManagedThreadId == _pipelineThreadId; + { + get + { + var pipelineThreadId = Volatile.Read(ref _pipelineThreadId); + return pipelineThreadId != 0 && + Environment.CurrentManagedThreadId == pipelineThreadId; + } + } private bool IsConstructionThreadOutsideAsyncHook => Volatile.Read(ref _currentOutPipe) is null && @@ -126,27 +133,7 @@ private IDisposable EnterDirectPipelineAccess() } private IDisposable EnterDirectPipelineInteraction() - { - ThrowIfStopped(); - ValidateInteractionGeneration(); - if (IsPipelineThread) - { - var pipelineContext = new SynchronizationContextScope( - Volatile.Read(ref _pipelineSynchronizationContext)); - try - { - Volatile.Read(ref _pumpQueuedItems)?.Invoke(); - return pipelineContext; - } - catch - { - pipelineContext.Dispose(); - throw; - } - } - - return new SynchronizationContextScope(SynchronizationContext.Current); - } + => EnterDirectPipelineAccess(); private void ValidateInteractionGeneration() { @@ -259,39 +246,44 @@ private bool TryQueue(PipelineItem item) { item.BindToHook(_hookGeneration.Value); var pumpLease = _pipelinePumpLease.Value; - if (item.HookGeneration != 0 && - item.HookGeneration != Volatile.Read(ref _acceptingHookWritesGeneration) && - (pumpLease is null || - !pumpLease.IsActive || - item.HookGeneration != pumpLease.Generation)) - { - item.ReplyPipe?.Reject(); - return false; - } + lock (_hookAdmissionLock) + { + if (item.HookGeneration != 0 && + item.HookGeneration != Volatile.Read(ref _acceptingHookWritesGeneration) && + (pumpLease is null || + !pumpLease.IsActive || + item.HookGeneration != pumpLease.Generation) && + (!item.DropOnStop || + item.HookGeneration != Volatile.Read(ref _pumpingHookGeneration))) + { + item.ReplyPipe?.Reject(); + return false; + } - var outPipe = Volatile.Read(ref _currentOutPipe); - if (outPipe is null) - return false; + var outPipe = Volatile.Read(ref _currentOutPipe); + if (outPipe is null) + return false; - try - { - outPipe.Add(item, CancelToken); - return true; - } - catch (ObjectDisposedException) - { - return false; - } - catch (InvalidOperationException) - { - return false; - } - catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) - { - if (item.HookGeneration != 0 && !item.DropOnStop) - throw new PipelineStoppedException(); + try + { + outPipe.Add(item, CancelToken); + return true; + } + catch (ObjectDisposedException) + { + return false; + } + catch (InvalidOperationException) + { + return false; + } + catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) + { + if (item.HookGeneration != 0 && !item.DropOnStop) + throw new PipelineStoppedException(); - return false; + return false; + } } } @@ -304,7 +296,7 @@ private void RunBlockInAsync(Func task) } finally { - _pipelineThreadId = 0; + Volatile.Write(ref _pipelineThreadId, 0); ExitAsyncBlock(); } } @@ -318,17 +310,19 @@ private void RunBlockInAsyncCore(Func task) Task blockTask; var deferPipeDisposal = 0; var pipeDisposed = 0; - var pumpingQueuedItems = 0; var hookGeneration = Interlocked.Increment(ref _nextHookGeneration); var synchronizationContext = SynchronizationContext.Current; void ClearPipes() { - _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); - Volatile.Write(ref _pumpQueuedItems, null); - _ = Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe); - _ = Interlocked.CompareExchange(ref _pipelineSynchronizationContext, null, synchronizationContext); - CompleteAddingIfNeeded(outPipe); + lock (_hookAdmissionLock) + { + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); + Volatile.Write(ref _pumpQueuedItems, null); + _ = Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe); + _ = Interlocked.CompareExchange(ref _pipelineSynchronizationContext, null, synchronizationContext); + CompleteAddingIfNeeded(outPipe); + } } void DeactivateHook() @@ -368,11 +362,13 @@ void PumpItem(PipelineItem item) var priorItemGeneration = _hookGeneration.Value; var priorPumpLease = _pipelinePumpLease.Value; + var priorPumpingGeneration = Volatile.Read(ref _pumpingHookGeneration); var pumpLease = new PipelinePumpLease(item.HookGeneration); try { _hookGeneration.Value = item.HookGeneration; _pipelinePumpLease.Value = pumpLease; + Volatile.Write(ref _pumpingHookGeneration, item.HookGeneration); switch (item.Type) { case PipelineType.Output: @@ -506,6 +502,7 @@ void PumpItem(PipelineItem item) finally { pumpLease.Close(); + Volatile.Write(ref _pumpingHookGeneration, priorPumpingGeneration); _pipelinePumpLease.Value = priorPumpLease; _hookGeneration.Value = priorItemGeneration; } @@ -513,26 +510,19 @@ void PumpItem(PipelineItem item) void PumpQueuedItems() { - if (Interlocked.Exchange(ref pumpingQueuedItems, 1) != 0) - return; - - try - { - while (outPipe.TryTake(out var item)) - PumpItem(item); - } - finally - { - Volatile.Write(ref pumpingQueuedItems, 0); - } + while (outPipe.TryTake(out var item)) + PumpItem(item); } Volatile.Write(ref _asyncLifecycleStarted, 1); - _pipelineThreadId = Environment.CurrentManagedThreadId; + Volatile.Write(ref _pipelineThreadId, Environment.CurrentManagedThreadId); Volatile.Write(ref _activeHookGeneration, hookGeneration); - Volatile.Write(ref _acceptingHookWritesGeneration, hookGeneration); Volatile.Write(ref _pumpQueuedItems, PumpQueuedItems); - Volatile.Write(ref _currentOutPipe, outPipe); + lock (_hookAdmissionLock) + { + Volatile.Write(ref _acceptingHookWritesGeneration, hookGeneration); + Volatile.Write(ref _currentOutPipe, outPipe); + } var priorHookGeneration = _hookGeneration.Value; try @@ -584,13 +574,22 @@ void PumpQueuedItems() if (blockTask.IsCompleted) { - _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); + lock (_hookAdmissionLock) + { + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); + } if (blockTask.IsFaulted) _ = blockTask.Exception; try { PumpQueuedItems(); + GetBlockTaskResult(blockTask); + } + catch (PipelineStoppedException) + { + CancelSource(); + throw; } finally { @@ -599,7 +598,6 @@ void PumpQueuedItems() DisposePipeOnce(); } - GetBlockTaskResult(blockTask); return; } @@ -611,29 +609,32 @@ void PumpQueuedItems() { try { - _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); if (completed.IsFaulted) _ = completed.Exception; - try + lock (_hookAdmissionLock) { - if (!outPipe.IsAddingCompleted) + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); + try { - outPipe.Add( - new PipelineItem( - value: null, - PipelineType.HookCompleted, - hookGeneration: hookGeneration, - dropOnStop: true)); + if (!outPipe.IsAddingCompleted) + { + outPipe.Add( + new PipelineItem( + value: null, + PipelineType.HookCompleted, + hookGeneration: hookGeneration, + dropOnStop: true)); + } + } + catch (ObjectDisposedException) + { + // A pipeline failure may dispose the transport before the hook completes. + } + catch (InvalidOperationException) + { + // The pipeline completed adding while the hook completion was published. } - } - catch (ObjectDisposedException) - { - // A pipeline failure may dispose the transport before the hook completes. - } - catch (InvalidOperationException) - { - // The pipeline completed adding while the hook completion was published. } if (Volatile.Read(ref deferPipeDisposal) != 0) @@ -733,6 +734,14 @@ private void RetainAsyncBlock() private void CancelSource() { + lock (_lifecycleLock) + { + if (_cancelSourceDisposed) + return; + + _cancelSourceCancellationInProgress++; + } + try { _cancelSource.Cancel(); @@ -746,11 +755,22 @@ private void CancelSource() { // Disposal may race a late StopProcessing callback after all async hooks have exited. } + finally + { + lock (_lifecycleLock) + { + _cancelSourceCancellationInProgress--; + DisposeCancelSourceIfInactive(); + } + } } private void DisposeCancelSourceIfInactive() { - if (!_disposeRequested || _activeBlocks != 0 || _cancelSourceDisposed) + if (!_disposeRequested || + _activeBlocks != 0 || + _cancelSourceCancellationInProgress != 0 || + _cancelSourceDisposed) return; _cancelSource.Dispose(); diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index d5bd89480..79428760d 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -288,6 +288,7 @@ private void Queue(object? value, PipelineType type) private readonly AsyncLocal _hookGeneration = new(); private readonly AsyncLocal _pipelinePumpLease = new(); private readonly int _constructionThreadId = Environment.CurrentManagedThreadId; + private readonly object _hookAdmissionLock = new(); private readonly object _lifecycleLock = new(); private static readonly SynchronizationContext HookSynchronizationContext = new AsyncHookSynchronizationContext(); private static readonly TaskScheduler HookTaskScheduler = new AsyncHookTaskScheduler(); @@ -297,6 +298,8 @@ private void Queue(object? value, PipelineType type) private long _activeHookGeneration; private long _acceptingHookWritesGeneration; private long _nextHookGeneration; + private long _pumpingHookGeneration; + private int _cancelSourceCancellationInProgress; private bool _cancelSourceDisposed; private bool _disposeRequested; private int _activeBlocks; @@ -325,16 +328,7 @@ protected virtual Task ProcessRecordAsync() /// protected override void EndProcessing() - { - try - { - RunBlockInAsync(EndProcessingAsync); - } - finally - { - Volatile.Write(ref _asyncLifecycleCompleted, 1); - } - } + => RunBlockInAsync(EndProcessingAsync); /// Asynchronous end hook. protected virtual Task EndProcessingAsync() diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index 2eadbf7f5..368af401c 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -634,6 +634,103 @@ System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncReentrantFifo")] +public sealed class TestAsyncReentrantFifoCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(new WarningEnumerable(this), enumerateCollection: true); + } + + private sealed class WarningEnumerable : IEnumerable + { + private readonly TestAsyncReentrantFifoCommand _command; + + public WarningEnumerable(TestAsyncReentrantFifoCommand command) + { + _command = command; + } + + public IEnumerator GetEnumerator() + { + Task.Run(() => _command.WriteWarning("queued-first")).GetAwaiter().GetResult(); + _command.WriteWarning("direct-second"); + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedEnumeration")] +public sealed class TestAsyncCapturedEnumerationCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var streams = CapturePipelineStreams(); + WriteObject( + new CapturedWarningEnumerable( + message => streams.WriteWarning(message)), + enumerateCollection: true); + } + + private sealed class CapturedWarningEnumerable : IEnumerable + { + private readonly Action _writeWarning; + + public CapturedWarningEnumerable(Action writeWarning) + { + _writeWarning = writeWarning; + } + + public IEnumerator GetEnumerator() + { + using var completed = new ManualResetEventSlim(); + Exception? callbackException = null; + ThreadPool.UnsafeQueueUserWorkItem( + _ => + { + try + { + _writeWarning("captured-during-enumeration"); + } + catch (Exception exception) + { + callbackException = exception; + } + finally + { + completed.Set(); + } + }, + null); + Assert.True( + completed.Wait(TimeSpan.FromSeconds(5)), + "The context-free callback did not complete in time."); + if (callbackException is not null) + throw callbackException; + + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDerivedEndProcessing")] +public sealed class TestAsyncDerivedEndProcessingCommand : AsyncPSCmdlet +{ + protected override void EndProcessing() + { + base.EndProcessing(); + WriteObject("derived-end"); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyShouldProcess", SupportsShouldProcess = true)] public sealed class TestAsyncEarlyShouldProcessCommand : AsyncPSCmdlet { @@ -754,6 +851,90 @@ protected override async Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringCancellation")] +public sealed class TestAsyncDisposeDuringCancellationCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _callbackCompleted = new(); + + public static ManualResetEventSlim CallbackCompleted => _callbackCompleted; + + public static bool TokenReadDuringCancellation { get; private set; } + + public static Exception? TokenReadException { get; private set; } + + public static void Reset() + { + _callbackCompleted.Dispose(); + _callbackCompleted = new ManualResetEventSlim(); + TokenReadDuringCancellation = false; + TokenReadException = null; + } + + protected override Task ProcessRecordAsync() + { + using var cancellationEntered = new ManualResetEventSlim(); + using var releaseCancellation = new ManualResetEventSlim(); + _ = CancelToken.Register( + () => + { + cancellationEntered.Set(); + releaseCancellation.Wait(TimeSpan.FromSeconds(5)); + try + { + _ = CancelToken; + TokenReadDuringCancellation = true; + } + catch (Exception exception) + { + TokenReadException = exception; + } + finally + { + _callbackCompleted.Set(); + } + }); + + _ = Task.Run(Dispose); + Assert.True( + cancellationEntered.Wait(TimeSpan.FromSeconds(5)), + "Cancellation did not enter the registered callback."); + ThreadPool.QueueUserWorkItem( + _ => + { + Thread.Sleep(100); + releaseCancellation.Set(); + }); + return Task.CompletedTask; + } + + protected override void EndProcessing() + { + // ProcessRecord intentionally disposes the command to exercise the cancellation race. + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDirectPipelineStop")] +public sealed class TestAsyncDirectPipelineStopCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _cancellationObserved = new(); + + public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + + public static void Reset() + { + _cancellationObserved.Dispose(); + _cancellationObserved = new ManualResetEventSlim(); + } + + protected override async Task ProcessRecordAsync() + { + _ = CancelToken.Register(_cancellationObserved.Set); + await Task.CompletedTask; + WriteObject(1); + WriteObject(2); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousStop")] public sealed class TestAsyncSynchronousStopCommand : AsyncPSCmdlet { diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 03e947f29..1e29b0854 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -677,6 +677,72 @@ public void AsyncPSCmdlet_does_not_recursively_drain_when_enumeration_writes_to_ powerShell.Streams.Warning.Select(static warning => warning.Message).Order()); } + [Fact] + public void AsyncPSCmdlet_preserves_fifo_order_during_reentrant_queue_drains() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncReentrantFifo", + typeof(TestAsyncReentrantFifoCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncReentrantFifo"); + + var result = powerShell.Invoke(); + + Assert.Equal("value", Assert.Single(result).BaseObject); + Assert.Equal( + ["queued-first", "direct-second"], + powerShell.Streams.Warning.Select(static warning => warning.Message)); + } + + [Fact] + public void AsyncPSCmdlet_accepts_context_free_callbacks_while_their_lazy_item_is_pumped() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCapturedEnumeration", + typeof(TestAsyncCapturedEnumerationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCapturedEnumeration"); + + var result = powerShell.Invoke(); + + Assert.Equal("value", Assert.Single(result).BaseObject); + Assert.Equal( + "captured-during-enumeration", + Assert.Single(powerShell.Streams.Warning).Message); + } + + [Fact] + public void AsyncPSCmdlet_keeps_derived_EndProcessing_active_after_the_base_call() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDerivedEndProcessing", + typeof(TestAsyncDerivedEndProcessingCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncDerivedEndProcessing"); + + var result = powerShell.Invoke(); + + Assert.Equal("derived-end", Assert.Single(result).BaseObject); + } + [Fact] public void AsyncPSCmdlet_allows_ShouldProcess_before_the_async_base_hook_starts() { @@ -801,6 +867,58 @@ public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_the_async_ho Assert.Null(TestAsyncDisposeDuringHookCommand.TokenReadException); } + [Fact] + public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_Cancel_returns() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDisposeDuringCancellation", + typeof(TestAsyncDisposeDuringCancellationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncDisposeDuringCancellation"); + TestAsyncDisposeDuringCancellationCommand.Reset(); + + powerShell.Invoke(); + + Assert.True( + TestAsyncDisposeDuringCancellationCommand.CallbackCompleted.Wait(TimeSpan.FromSeconds(5)), + "The cancellation callback did not complete in time."); + Assert.True(TestAsyncDisposeDuringCancellationCommand.TokenReadDuringCancellation); + Assert.Null(TestAsyncDisposeDuringCancellationCommand.TokenReadException); + } + + [Fact] + public void AsyncPSCmdlet_cancels_started_work_when_a_direct_write_stops_the_pipeline() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDirectPipelineStop", + typeof(TestAsyncDirectPipelineStopCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncDirectPipelineStop") + .AddCommand("Select-Object") + .AddParameter("First", 1); + TestAsyncDirectPipelineStopCommand.Reset(); + + var result = powerShell.Invoke(); + + Assert.Equal(1, Assert.Single(result).BaseObject); + Assert.True( + TestAsyncDirectPipelineStopCommand.CancellationObserved.Wait(TimeSpan.FromSeconds(5)), + "The started operation was not canceled after the pipeline stopped."); + } + [Fact] public void AsyncPSCmdlet_normalizes_synchronous_cancellation_after_stop() { From c22b6634811ecaba374b26231872537ecc9588f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 07:11:41 +0200 Subject: [PATCH 22/27] Close final async lifecycle races --- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 15 ++++- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 13 +++- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 65 +++++++++++++++++++ PowerForge.Tests/AsyncPSCmdletTests.cs | 52 +++++++++++++++ 4 files changed, 142 insertions(+), 3 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index d1317b895..bd17210a0 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -510,7 +510,8 @@ void PumpItem(PipelineItem item) void PumpQueuedItems() { - while (outPipe.TryTake(out var item)) + var queuedAtEntry = outPipe.Count; + while (queuedAtEntry-- > 0 && outPipe.TryTake(out var item)) PumpItem(item); } @@ -546,6 +547,11 @@ void PumpQueuedItems() } catch (Exception exception) { + lock (_hookAdmissionLock) + { + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); + } + SynchronizationContext.SetSynchronizationContext(synchronizationContext); try { PumpQueuedItems(); @@ -607,11 +613,15 @@ void PumpQueuedItems() _ = blockTask.ContinueWith( completed => { + var retainedBlockOwned = true; try { if (completed.IsFaulted) _ = completed.Exception; + ExitAsyncBlock(); + retainedBlockOwned = false; + lock (_hookAdmissionLock) { _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); @@ -645,7 +655,8 @@ void PumpQueuedItems() } finally { - ExitAsyncBlock(); + if (retainedBlockOwned) + ExitAsyncBlock(); } }, CancellationToken.None, diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 79428760d..b20d7caa7 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -23,7 +23,18 @@ public abstract partial class AsyncPSCmdlet : PSCmdlet, IDisposable private sealed class AsyncHookSynchronizationContext : SynchronizationContext { public override void Post(SendOrPostCallback callback, object? state) - => ThreadPool.QueueUserWorkItem(_ => callback(state)); + => ThreadPool.QueueUserWorkItem(_ => + { + try + { + callback(state); + } + catch (PipelineStoppedException) + { + // Fire-and-forget callbacks such as Progress can run after StopProcessing. + // Await continuations capture their own exceptions into the hook task. + } + }); } private sealed class AsyncHookTaskScheduler : TaskScheduler diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index 368af401c..c026f94a3 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -935,6 +935,71 @@ protected override async Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncSuccessfulCompletion")] +public sealed class TestAsyncSuccessfulCompletionCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _cancellationObserved = new(); + + public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + + public static void Reset() + { + _cancellationObserved.Dispose(); + _cancellationObserved = new ManualResetEventSlim(); + } + + protected override async Task ProcessRecordAsync() + { + _ = CancelToken.Register(_cancellationObserved.Set); + await Task.Yield(); + WriteObject("completed"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncLateProgress")] +public sealed class TestAsyncLateProgressCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _callbackCompleted = new(); + private static ManualResetEventSlim _initialized = new(); + private static IProgress? _progress; + + public static ManualResetEventSlim CallbackCompleted => _callbackCompleted; + + public static ManualResetEventSlim Initialized => _initialized; + + public static void Reset() + { + _callbackCompleted.Dispose(); + _initialized.Dispose(); + _callbackCompleted = new ManualResetEventSlim(); + _initialized = new ManualResetEventSlim(); + _progress = null; + } + + public static void ReportAfterStop() + => _progress!.Report(50); + + protected override async Task ProcessRecordAsync() + { + _progress = new Progress(percent => + { + try + { + WriteProgress(new ProgressRecord(1, "late", "after stop") + { + PercentComplete = percent + }); + } + finally + { + _callbackCompleted.Set(); + } + }); + _initialized.Set(); + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousStop")] public sealed class TestAsyncSynchronousStopCommand : AsyncPSCmdlet { diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 1e29b0854..bfc56d3cb 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -919,6 +919,58 @@ public void AsyncPSCmdlet_cancels_started_work_when_a_direct_write_stops_the_pip "The started operation was not canceled after the pipeline stopped."); } + [Fact] + public void AsyncPSCmdlet_does_not_cancel_after_successful_async_completion() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSuccessfulCompletion", + typeof(TestAsyncSuccessfulCompletionCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncSuccessfulCompletion"); + TestAsyncSuccessfulCompletionCommand.Reset(); + + var result = powerShell.Invoke(); + + Assert.Equal("completed", Assert.Single(result).BaseObject); + Assert.False(TestAsyncSuccessfulCompletionCommand.CancellationObserved.IsSet); + } + + [Fact] + public void AsyncPSCmdlet_suppresses_pipeline_stop_from_posted_progress_callbacks() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncLateProgress", + typeof(TestAsyncLateProgressCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncLateProgress"); + TestAsyncLateProgressCommand.Reset(); + + var invocation = powerShell.BeginInvoke(); + Assert.True( + TestAsyncLateProgressCommand.Initialized.Wait(TimeSpan.FromSeconds(5)), + "The progress callback was not initialized in time."); + powerShell.Stop(); + Assert.Throws(() => powerShell.EndInvoke(invocation)); + + TestAsyncLateProgressCommand.ReportAfterStop(); + + Assert.True( + TestAsyncLateProgressCommand.CallbackCompleted.Wait(TimeSpan.FromSeconds(5)), + "The posted progress callback did not complete in time."); + } + [Fact] public void AsyncPSCmdlet_normalizes_synchronous_cancellation_after_stop() { From e81b733ac0341f03c8be6c85a69007201bb35f2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 08:11:06 +0200 Subject: [PATCH 23/27] Harden async cmdlet pipeline lifecycle --- .../Cmdlets/AsyncPSCmdlet.Execution.cs | 462 ++++++++++++++ .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 452 ++------------ PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 98 ++- .../AsyncPSCmdletLifecycleTestCommands.cs | 578 ++++++++++++++++++ .../AsyncPSCmdletLifecycleTests.cs | 485 +++++++++++++++ PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 571 +---------------- PowerForge.Tests/AsyncPSCmdletTests.cs | 489 +-------------- 7 files changed, 1661 insertions(+), 1474 deletions(-) create mode 100644 PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs create mode 100644 PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs create mode 100644 PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs new file mode 100644 index 000000000..adcf094da --- /dev/null +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs @@ -0,0 +1,462 @@ +using System; +using System.Collections.Concurrent; +using System.Management.Automation; +using System.Threading; +using System.Threading.Tasks; + +namespace PSPublishModule; + +public abstract partial class AsyncPSCmdlet +{ + private void RunBlockInAsyncCore(Func task) + { + // The transport must remain lossless and non-blocking. The pipeline thread can enumerate + // user objects or invoke a host that waits for the same background producer that is writing + // here; applying bounded backpressure would deadlock both sides. + var outPipe = new BlockingCollection(); + Task blockTask; + var deferPipeDisposal = 0; + var pipeDisposed = 0; + var hookGeneration = Interlocked.Increment(ref _nextHookGeneration); + var synchronizationContext = SynchronizationContext.Current; + + void ClearPipes() + { + lock (_hookAdmissionLock) + { + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); + Volatile.Write(ref _pumpQueuedItems, null); + _ = Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe); + _ = Interlocked.CompareExchange(ref _pipelineSynchronizationContext, null, synchronizationContext); + CompleteAddingIfNeeded(outPipe); + } + } + + void DeactivateHook() + => _ = Interlocked.CompareExchange(ref _activeHookGeneration, 0, hookGeneration); + + void DisposePipeOnce() + { + if (Interlocked.Exchange(ref pipeDisposed, 1) == 0) + { + while (outPipe.TryTake(out var abandonedItem)) + abandonedItem.ReplyPipe?.Reject(); + outPipe.Dispose(); + } + } + + static void CompleteAddingIfNeeded(BlockingCollection pipe) + { + try + { + if (!pipe.IsAddingCompleted) + pipe.CompleteAdding(); + } + catch (ObjectDisposedException) + { + // A deferred worker may race the one-time disposal after a pipeline failure. + } + } + + void PumpItem(PipelineItem item) + { + if (Volatile.Read(ref _asyncLifecycleStarted) != 0 && + item.HookGeneration != Volatile.Read(ref _activeHookGeneration)) + { + item.ReplyPipe?.Reject(); + return; + } + + var priorItemGeneration = _hookGeneration.Value; + var priorPumpLease = _pipelinePumpLease.Value; + var pumpLease = new PipelinePumpLease(item.HookGeneration); + try + { + _ = Interlocked.Increment(ref _pipelinePumpDepth); + _hookGeneration.Value = item.HookGeneration; + _pipelinePumpLease.Value = pumpLease; + switch (item.Type) + { + case PipelineType.Output: + base.WriteObject(item.Value); + break; + case PipelineType.OutputEnumerate: + base.WriteObject(item.Value, enumerateCollection: true); + break; + case PipelineType.Error: + base.WriteError((ErrorRecord)item.Value!); + break; + case PipelineType.TerminatingError: + base.ThrowTerminatingError((ErrorRecord)item.Value!); + break; + case PipelineType.Warning: + base.WriteWarning((string)item.Value!); + break; + case PipelineType.Verbose: + base.WriteVerbose((string)item.Value!); + break; + case PipelineType.Debug: + base.WriteDebug((string)item.Value!); + break; + case PipelineType.Information: + base.WriteInformation((InformationRecord)item.Value!); + break; + case PipelineType.InformationWithTags: + var information = ((object MessageData, string[]? Tags))item.Value!; + base.WriteInformation( + information.MessageData, + information.Tags ?? Array.Empty()); + break; + case PipelineType.Progress: + base.WriteProgress((ProgressRecord)item.Value!); + break; + case PipelineType.CommandDetail: + base.WriteCommandDetail((string)item.Value!); + break; + case PipelineType.ShouldProcessTarget: + item.ReplyPipe!.Publish( + () => base.ShouldProcess((string)item.Value!)); + break; + case PipelineType.ShouldProcess: + var should = ((string Target, string Action))item.Value!; + item.ReplyPipe!.Publish( + () => base.ShouldProcess(should.Target, should.Action)); + break; + case PipelineType.ShouldProcessVerbose: + var verbose = ((string Description, string Warning, string Caption))item.Value!; + item.ReplyPipe!.Publish( + () => base.ShouldProcess(verbose.Description, verbose.Warning, verbose.Caption)); + break; + case PipelineType.ShouldProcessReason: + var reasonRequest = ((string Description, string Warning, string Caption))item.Value!; + item.ReplyPipe!.Publish(() => + { + var result = base.ShouldProcess( + reasonRequest.Description, + reasonRequest.Warning, + reasonRequest.Caption, + out var reason); + return (result, reason); + }); + break; + case PipelineType.ShouldContinue: + var shouldContinue = ((string Query, string Caption))item.Value!; + item.ReplyPipe!.Publish( + () => base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption)); + break; + case PipelineType.ShouldContinueAll: + var shouldContinueAll = + ((string Query, string Caption, bool YesToAll, bool NoToAll))item.Value!; + item.ReplyPipe!.Publish(() => + { + var yesToAll = shouldContinueAll.YesToAll; + var noToAll = shouldContinueAll.NoToAll; + var continueAll = base.ShouldContinue( + shouldContinueAll.Query, + shouldContinueAll.Caption, + ref yesToAll, + ref noToAll); + return (continueAll, yesToAll, noToAll); + }); + break; + case PipelineType.ShouldContinueSecurity: + var shouldContinueSecurity = + ((string Query, string Caption, bool HasSecurityImpact, bool YesToAll, bool NoToAll))item.Value!; + item.ReplyPipe!.Publish(() => + { + var yesToAll = shouldContinueSecurity.YesToAll; + var noToAll = shouldContinueSecurity.NoToAll; + var continueSecurity = base.ShouldContinue( + shouldContinueSecurity.Query, + shouldContinueSecurity.Caption, + shouldContinueSecurity.HasSecurityImpact, + ref yesToAll, + ref noToAll); + return (continueSecurity, yesToAll, noToAll); + }); + break; + case PipelineType.PromptForCredential: + var prompt = ((string Caption, string Message, string UserName, string TargetName))item.Value!; + item.ReplyPipe!.Publish( + () => Host.UI.PromptForCredential( + prompt.Caption, + prompt.Message, + prompt.UserName, + prompt.TargetName)); + break; + case PipelineType.PromptForCredentialOptions: + var promptOptions = + ((string Caption, + string Message, + string UserName, + string TargetName, + PSCredentialTypes AllowedCredentialTypes, + PSCredentialUIOptions Options))item.Value!; + item.ReplyPipe!.Publish( + () => Host.UI.PromptForCredential( + promptOptions.Caption, + promptOptions.Message, + promptOptions.UserName, + promptOptions.TargetName, + promptOptions.AllowedCredentialTypes, + promptOptions.Options)); + break; + case PipelineType.DirectAccessBarrier: + case PipelineType.HookCompleted: + break; + } + } + finally + { + pumpLease.Close(); + _pipelinePumpLease.Value = priorPumpLease; + _hookGeneration.Value = priorItemGeneration; + _ = Interlocked.Decrement(ref _pipelinePumpDepth); + } + } + + void PumpQueuedItems() + { + if (IsPumpingPipelineItem) + return; + + int queuedAtEntry; + lock (_hookAdmissionLock) + { + queuedAtEntry = outPipe.Count; + } + + while (queuedAtEntry-- > 0 && outPipe.TryTake(out var item)) + PumpItem(item); + } + + void PumpThroughDirectAccessBarrier() + { + if (IsPumpingPipelineItem) + return; + + var barrier = new PipelineItem( + value: null, + PipelineType.DirectAccessBarrier, + hookGeneration: hookGeneration, + dropOnStop: true); + if (!TryQueue(barrier)) + return; + + while (outPipe.TryTake(out var item)) + { + PumpItem(item); + if (ReferenceEquals(item, barrier)) + return; + } + } + + Volatile.Write(ref _asyncLifecycleStarted, 1); + Volatile.Write(ref _pipelineThreadId, Environment.CurrentManagedThreadId); + Volatile.Write(ref _activeHookGeneration, hookGeneration); + Volatile.Write(ref _pumpQueuedItems, PumpThroughDirectAccessBarrier); + lock (_hookAdmissionLock) + { + Volatile.Write(ref _acceptingHookWritesGeneration, hookGeneration); + Volatile.Write(ref _currentOutPipe, outPipe); + } + + var priorHookGeneration = _hookGeneration.Value; + try + { + Volatile.Write(ref _pipelineSynchronizationContext, synchronizationContext); + SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); + _hookGeneration.Value = hookGeneration; + if (TaskScheduler.Current == TaskScheduler.Default) + { + blockTask = task(); + } + else + { + using var invocationTask = new Task( + task, + CancellationToken.None, + TaskCreationOptions.DenyChildAttach); + invocationTask.RunSynchronously(HookTaskScheduler); + blockTask = invocationTask.GetAwaiter().GetResult(); + } + } + catch (Exception exception) + { + lock (_hookAdmissionLock) + { + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); + } + SynchronizationContext.SetSynchronizationContext(synchronizationContext); + try + { + PumpQueuedItems(); + } + catch + { + // Preserve the hook failure after best-effort delivery of records written before it. + } + finally + { + ClearPipes(); + DeactivateHook(); + DisposePipeOnce(); + } + + if (exception is OperationCanceledException && _cancelSource.IsCancellationRequested) + throw new PipelineStoppedException(); + + throw; + } + finally + { + _hookGeneration.Value = priorHookGeneration; + SynchronizationContext.SetSynchronizationContext(synchronizationContext); + } + + if (blockTask.IsCompleted) + { + lock (_hookAdmissionLock) + { + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); + } + if (blockTask.IsFaulted) + _ = blockTask.Exception; + + try + { + ThrowIfStopped(); + PumpQueuedItems(); + GetBlockTaskResult(blockTask); + } + catch (PipelineStoppedException) + { + CancelSource(); + throw; + } + finally + { + ClearPipes(); + DeactivateHook(); + DisposePipeOnce(); + } + + return; + } + + RetainAsyncBlock(); + try + { + _ = blockTask.ContinueWith( + completed => + { + var retainedBlockOwned = true; + try + { + if (completed.IsFaulted) + _ = completed.Exception; + + ExitAsyncBlock(); + retainedBlockOwned = false; + + lock (_hookAdmissionLock) + { + _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); + try + { + if (!outPipe.IsAddingCompleted) + { + outPipe.Add( + new PipelineItem( + value: null, + PipelineType.HookCompleted, + hookGeneration: hookGeneration, + dropOnStop: true)); + } + } + catch (ObjectDisposedException) + { + // A pipeline failure may dispose the transport before the hook completes. + } + catch (InvalidOperationException) + { + // The pipeline completed adding while the hook completion was published. + } + } + + if (Volatile.Read(ref deferPipeDisposal) != 0) + { + ClearPipes(); + DisposePipeOnce(); + } + } + finally + { + if (retainedBlockOwned) + ExitAsyncBlock(); + } + }, + CancellationToken.None, + TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + } + catch + { + ExitAsyncBlock(); + throw; + } + + try + { + while (true) + { + var item = outPipe.Take(CancelToken); + PumpItem(item); + if (item.Type == PipelineType.HookCompleted) + { + while (outPipe.TryTake(out var pumpBoundItem)) + PumpItem(pumpBoundItem); + break; + } + } + + ClearPipes(); + } + catch (Exception pipelineException) + { + var stopRequested = _cancelSource.IsCancellationRequested; + Volatile.Write(ref deferPipeDisposal, 1); + try + { + CancelSource(); + } + catch (AggregateException) + { + // Preserve the pipeline failure while cancellation callbacks observe the same stop. + } + finally + { + CompleteAddingIfNeeded(outPipe); + if (blockTask.IsCompleted) + DisposePipeOnce(); + DeactivateHook(); + } + + if (pipelineException is OperationCanceledException && stopRequested) + throw new PipelineStoppedException(); + + throw; + } + + try + { + GetBlockTaskResult(blockTask); + } + finally + { + DeactivateHook(); + DisposePipeOnce(); + } + } +} + diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index bd17210a0..ba861000a 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -13,6 +13,13 @@ public abstract partial class AsyncPSCmdlet public new void WriteProgress(ProgressRecord progressRecord) { ThrowIfStopped(); + var item = new PipelineItem(SnapshotProgressRecord(progressRecord), PipelineType.Progress); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + if (CanAccessPipelineDirectly) { using var pipelineContext = EnterDirectPipelineAccess(); @@ -23,7 +30,7 @@ public abstract partial class AsyncPSCmdlet if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(SnapshotProgressRecord(progressRecord), PipelineType.Progress)); + _ = TryQueue(item); } private static readonly PropertyInfo? ProgressTotalProperty = @@ -53,6 +60,21 @@ private static ProgressRecord SnapshotProgressRecord(ProgressRecord progressReco return snapshot; } + private static InformationRecord SnapshotInformationRecord(InformationRecord informationRecord) + { + var snapshot = new InformationRecord(informationRecord.MessageData, informationRecord.Source) + { + TimeGenerated = informationRecord.TimeGenerated, + User = informationRecord.User, + Computer = informationRecord.Computer, + ProcessId = informationRecord.ProcessId, + NativeThreadId = informationRecord.NativeThreadId, + ManagedThreadId = informationRecord.ManagedThreadId + }; + snapshot.Tags.AddRange(informationRecord.Tags); + return snapshot; + } + /// Throws when PowerShell has requested cancellation. protected internal void ThrowIfStopped() { @@ -100,6 +122,9 @@ private bool IsPipelineThread } } + private bool IsPumpingPipelineItem + => IsPipelineThread && Volatile.Read(ref _pipelinePumpDepth) != 0; + private bool IsConstructionThreadOutsideAsyncHook => Volatile.Read(ref _currentOutPipe) is null && Volatile.Read(ref _asyncLifecycleCompleted) == 0 && @@ -119,7 +144,8 @@ private IDisposable EnterDirectPipelineAccess() Volatile.Read(ref _pipelineSynchronizationContext)); try { - Volatile.Read(ref _pumpQueuedItems)?.Invoke(); + if (!IsPumpingPipelineItem) + Volatile.Read(ref _pumpQueuedItems)?.Invoke(); return pipelineContext; } catch @@ -164,6 +190,11 @@ private void GetBlockTaskResult(Task blockTask) { throw new PipelineStoppedException(); } + catch (PipelineStoppedException) + { + CancelSource(); + throw; + } } private object? RequestPipelineReply(object? value, PipelineType type) @@ -252,9 +283,7 @@ private bool TryQueue(PipelineItem item) item.HookGeneration != Volatile.Read(ref _acceptingHookWritesGeneration) && (pumpLease is null || !pumpLease.IsActive || - item.HookGeneration != pumpLease.Generation) && - (!item.DropOnStop || - item.HookGeneration != Volatile.Read(ref _pumpingHookGeneration))) + item.HookGeneration != pumpLease.Generation)) { item.ReplyPipe?.Reject(); return false; @@ -301,419 +330,6 @@ private void RunBlockInAsync(Func task) } } - private void RunBlockInAsyncCore(Func task) - { - // The transport must remain lossless and non-blocking. The pipeline thread can enumerate - // user objects or invoke a host that waits for the same background producer that is writing - // here; applying bounded backpressure would deadlock both sides. - var outPipe = new BlockingCollection(); - Task blockTask; - var deferPipeDisposal = 0; - var pipeDisposed = 0; - var hookGeneration = Interlocked.Increment(ref _nextHookGeneration); - var synchronizationContext = SynchronizationContext.Current; - - void ClearPipes() - { - lock (_hookAdmissionLock) - { - _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); - Volatile.Write(ref _pumpQueuedItems, null); - _ = Interlocked.CompareExchange(ref _currentOutPipe, null, outPipe); - _ = Interlocked.CompareExchange(ref _pipelineSynchronizationContext, null, synchronizationContext); - CompleteAddingIfNeeded(outPipe); - } - } - - void DeactivateHook() - => _ = Interlocked.CompareExchange(ref _activeHookGeneration, 0, hookGeneration); - - void DisposePipeOnce() - { - if (Interlocked.Exchange(ref pipeDisposed, 1) == 0) - { - while (outPipe.TryTake(out var abandonedItem)) - abandonedItem.ReplyPipe?.Reject(); - outPipe.Dispose(); - } - } - - static void CompleteAddingIfNeeded(BlockingCollection pipe) - { - try - { - if (!pipe.IsAddingCompleted) - pipe.CompleteAdding(); - } - catch (ObjectDisposedException) - { - // A deferred worker may race the one-time disposal after a pipeline failure. - } - } - - void PumpItem(PipelineItem item) - { - if (Volatile.Read(ref _asyncLifecycleStarted) != 0 && - item.HookGeneration != Volatile.Read(ref _activeHookGeneration)) - { - item.ReplyPipe?.Reject(); - return; - } - - var priorItemGeneration = _hookGeneration.Value; - var priorPumpLease = _pipelinePumpLease.Value; - var priorPumpingGeneration = Volatile.Read(ref _pumpingHookGeneration); - var pumpLease = new PipelinePumpLease(item.HookGeneration); - try - { - _hookGeneration.Value = item.HookGeneration; - _pipelinePumpLease.Value = pumpLease; - Volatile.Write(ref _pumpingHookGeneration, item.HookGeneration); - switch (item.Type) - { - case PipelineType.Output: - base.WriteObject(item.Value); - break; - case PipelineType.OutputEnumerate: - base.WriteObject(item.Value, enumerateCollection: true); - break; - case PipelineType.Error: - base.WriteError((ErrorRecord)item.Value!); - break; - case PipelineType.TerminatingError: - base.ThrowTerminatingError((ErrorRecord)item.Value!); - break; - case PipelineType.Warning: - base.WriteWarning((string)item.Value!); - break; - case PipelineType.Verbose: - base.WriteVerbose((string)item.Value!); - break; - case PipelineType.Debug: - base.WriteDebug((string)item.Value!); - break; - case PipelineType.Information: - base.WriteInformation((InformationRecord)item.Value!); - break; - case PipelineType.InformationWithTags: - var information = ((object MessageData, string[]? Tags))item.Value!; - base.WriteInformation( - information.MessageData, - information.Tags ?? Array.Empty()); - break; - case PipelineType.Progress: - base.WriteProgress((ProgressRecord)item.Value!); - break; - case PipelineType.CommandDetail: - base.WriteCommandDetail((string)item.Value!); - break; - case PipelineType.ShouldProcessTarget: - item.ReplyPipe!.Publish( - () => base.ShouldProcess((string)item.Value!)); - break; - case PipelineType.ShouldProcess: - var should = ((string Target, string Action))item.Value!; - item.ReplyPipe!.Publish( - () => base.ShouldProcess(should.Target, should.Action)); - break; - case PipelineType.ShouldProcessVerbose: - var verbose = ((string Description, string Warning, string Caption))item.Value!; - item.ReplyPipe!.Publish( - () => base.ShouldProcess(verbose.Description, verbose.Warning, verbose.Caption)); - break; - case PipelineType.ShouldProcessReason: - var reasonRequest = ((string Description, string Warning, string Caption))item.Value!; - item.ReplyPipe!.Publish(() => - { - var result = base.ShouldProcess( - reasonRequest.Description, - reasonRequest.Warning, - reasonRequest.Caption, - out var reason); - return (result, reason); - }); - break; - case PipelineType.ShouldContinue: - var shouldContinue = ((string Query, string Caption))item.Value!; - item.ReplyPipe!.Publish( - () => base.ShouldContinue(shouldContinue.Query, shouldContinue.Caption)); - break; - case PipelineType.ShouldContinueAll: - var shouldContinueAll = - ((string Query, string Caption, bool YesToAll, bool NoToAll))item.Value!; - item.ReplyPipe!.Publish(() => - { - var yesToAll = shouldContinueAll.YesToAll; - var noToAll = shouldContinueAll.NoToAll; - var continueAll = base.ShouldContinue( - shouldContinueAll.Query, - shouldContinueAll.Caption, - ref yesToAll, - ref noToAll); - return (continueAll, yesToAll, noToAll); - }); - break; - case PipelineType.ShouldContinueSecurity: - var shouldContinueSecurity = - ((string Query, string Caption, bool HasSecurityImpact, bool YesToAll, bool NoToAll))item.Value!; - item.ReplyPipe!.Publish(() => - { - var yesToAll = shouldContinueSecurity.YesToAll; - var noToAll = shouldContinueSecurity.NoToAll; - var continueSecurity = base.ShouldContinue( - shouldContinueSecurity.Query, - shouldContinueSecurity.Caption, - shouldContinueSecurity.HasSecurityImpact, - ref yesToAll, - ref noToAll); - return (continueSecurity, yesToAll, noToAll); - }); - break; - case PipelineType.PromptForCredential: - var prompt = ((string Caption, string Message, string UserName, string TargetName))item.Value!; - item.ReplyPipe!.Publish( - () => Host.UI.PromptForCredential( - prompt.Caption, - prompt.Message, - prompt.UserName, - prompt.TargetName)); - break; - case PipelineType.PromptForCredentialOptions: - var promptOptions = - ((string Caption, - string Message, - string UserName, - string TargetName, - PSCredentialTypes AllowedCredentialTypes, - PSCredentialUIOptions Options))item.Value!; - item.ReplyPipe!.Publish( - () => Host.UI.PromptForCredential( - promptOptions.Caption, - promptOptions.Message, - promptOptions.UserName, - promptOptions.TargetName, - promptOptions.AllowedCredentialTypes, - promptOptions.Options)); - break; - case PipelineType.HookCompleted: - break; - } - } - finally - { - pumpLease.Close(); - Volatile.Write(ref _pumpingHookGeneration, priorPumpingGeneration); - _pipelinePumpLease.Value = priorPumpLease; - _hookGeneration.Value = priorItemGeneration; - } - } - - void PumpQueuedItems() - { - var queuedAtEntry = outPipe.Count; - while (queuedAtEntry-- > 0 && outPipe.TryTake(out var item)) - PumpItem(item); - } - - Volatile.Write(ref _asyncLifecycleStarted, 1); - Volatile.Write(ref _pipelineThreadId, Environment.CurrentManagedThreadId); - Volatile.Write(ref _activeHookGeneration, hookGeneration); - Volatile.Write(ref _pumpQueuedItems, PumpQueuedItems); - lock (_hookAdmissionLock) - { - Volatile.Write(ref _acceptingHookWritesGeneration, hookGeneration); - Volatile.Write(ref _currentOutPipe, outPipe); - } - - var priorHookGeneration = _hookGeneration.Value; - try - { - Volatile.Write(ref _pipelineSynchronizationContext, synchronizationContext); - SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); - _hookGeneration.Value = hookGeneration; - if (TaskScheduler.Current == TaskScheduler.Default) - { - blockTask = task(); - } - else - { - using var invocationTask = new Task( - task, - CancellationToken.None, - TaskCreationOptions.DenyChildAttach); - invocationTask.RunSynchronously(HookTaskScheduler); - blockTask = invocationTask.GetAwaiter().GetResult(); - } - } - catch (Exception exception) - { - lock (_hookAdmissionLock) - { - _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); - } - SynchronizationContext.SetSynchronizationContext(synchronizationContext); - try - { - PumpQueuedItems(); - } - catch - { - // Preserve the hook failure after best-effort delivery of records written before it. - } - finally - { - ClearPipes(); - DeactivateHook(); - DisposePipeOnce(); - } - - if (exception is OperationCanceledException && _cancelSource.IsCancellationRequested) - throw new PipelineStoppedException(); - - throw; - } - finally - { - _hookGeneration.Value = priorHookGeneration; - SynchronizationContext.SetSynchronizationContext(synchronizationContext); - } - - if (blockTask.IsCompleted) - { - lock (_hookAdmissionLock) - { - _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); - } - if (blockTask.IsFaulted) - _ = blockTask.Exception; - - try - { - PumpQueuedItems(); - GetBlockTaskResult(blockTask); - } - catch (PipelineStoppedException) - { - CancelSource(); - throw; - } - finally - { - ClearPipes(); - DeactivateHook(); - DisposePipeOnce(); - } - - return; - } - - RetainAsyncBlock(); - try - { - _ = blockTask.ContinueWith( - completed => - { - var retainedBlockOwned = true; - try - { - if (completed.IsFaulted) - _ = completed.Exception; - - ExitAsyncBlock(); - retainedBlockOwned = false; - - lock (_hookAdmissionLock) - { - _ = Interlocked.CompareExchange(ref _acceptingHookWritesGeneration, 0, hookGeneration); - try - { - if (!outPipe.IsAddingCompleted) - { - outPipe.Add( - new PipelineItem( - value: null, - PipelineType.HookCompleted, - hookGeneration: hookGeneration, - dropOnStop: true)); - } - } - catch (ObjectDisposedException) - { - // A pipeline failure may dispose the transport before the hook completes. - } - catch (InvalidOperationException) - { - // The pipeline completed adding while the hook completion was published. - } - } - - if (Volatile.Read(ref deferPipeDisposal) != 0) - { - ClearPipes(); - DisposePipeOnce(); - } - } - finally - { - if (retainedBlockOwned) - ExitAsyncBlock(); - } - }, - CancellationToken.None, - TaskContinuationOptions.ExecuteSynchronously, - TaskScheduler.Default); - } - catch - { - ExitAsyncBlock(); - throw; - } - - try - { - while (!blockTask.IsCompleted || outPipe.Count != 0) - { - PumpItem(outPipe.Take(CancelToken)); - } - - ClearPipes(); - } - catch (Exception pipelineException) - { - var stopRequested = _cancelSource.IsCancellationRequested; - Volatile.Write(ref deferPipeDisposal, 1); - try - { - CancelSource(); - } - catch (AggregateException) - { - // Preserve the pipeline failure while cancellation callbacks observe the same stop. - } - finally - { - CompleteAddingIfNeeded(outPipe); - if (blockTask.IsCompleted) - DisposePipeOnce(); - DeactivateHook(); - } - - if (pipelineException is OperationCanceledException && stopRequested) - throw new PipelineStoppedException(); - - throw; - } - - try - { - GetBlockTaskResult(blockTask); - } - finally - { - DeactivateHook(); - DisposePipeOnce(); - } - } private void EnterAsyncBlock() { diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index b20d7caa7..a66794952 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -85,6 +85,7 @@ private enum PipelineType ShouldContinueSecurity, PromptForCredential, PromptForCredentialOptions, + DirectAccessBarrier, HookCompleted } @@ -270,7 +271,7 @@ public void WriteDebug(string message) /// Queues an information record for the originating hook. public void WriteInformation(InformationRecord informationRecord) - => Queue(informationRecord, PipelineType.Information); + => Queue(SnapshotInformationRecord(informationRecord), PipelineType.Information); /// Queues tagged information for the originating hook. public void WriteInformation(object messageData, string[]? tags) @@ -309,13 +310,13 @@ private void Queue(object? value, PipelineType type) private long _activeHookGeneration; private long _acceptingHookWritesGeneration; private long _nextHookGeneration; - private long _pumpingHookGeneration; private int _cancelSourceCancellationInProgress; private bool _cancelSourceDisposed; private bool _disposeRequested; private int _activeBlocks; private int _asyncLifecycleCompleted; private int _asyncLifecycleStarted; + private int _pipelinePumpDepth; private int _pipelineThreadId; /// Cancellation token triggered when PowerShell stops the cmdlet. @@ -339,7 +340,16 @@ protected virtual Task ProcessRecordAsync() /// protected override void EndProcessing() - => RunBlockInAsync(EndProcessingAsync); + { + try + { + RunBlockInAsync(EndProcessingAsync); + } + finally + { + Volatile.Write(ref _asyncLifecycleCompleted, 1); + } + } /// Asynchronous end hook. protected virtual Task EndProcessingAsync() @@ -515,6 +525,15 @@ protected override void StopProcessing() public new void WriteObject(object? sendToPipeline, bool enumerateCollection) { ThrowIfStopped(); + var item = new PipelineItem( + sendToPipeline, + enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + if (CanAccessPipelineDirectly) { using var pipelineContext = EnterDirectPipelineAccess(); @@ -525,15 +544,20 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem( - sendToPipeline, - enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output)); + _ = TryQueue(item); } /// Thread-safe error bridge for asynchronous cmdlet code. public new void WriteError(ErrorRecord errorRecord) { ThrowIfStopped(); + var item = new PipelineItem(errorRecord, PipelineType.Error); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + if (CanAccessPipelineDirectly) { using var pipelineContext = EnterDirectPipelineAccess(); @@ -544,7 +568,7 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(errorRecord, PipelineType.Error)); + _ = TryQueue(item); } /// Thread-safe terminating-error bridge for asynchronous cmdlet code. @@ -572,6 +596,13 @@ protected override void StopProcessing() public new void WriteWarning(string message) { ThrowIfStopped(); + var item = new PipelineItem(message, PipelineType.Warning); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + if (CanAccessPipelineDirectly) { using var pipelineContext = EnterDirectPipelineAccess(); @@ -582,13 +613,20 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(message, PipelineType.Warning)); + _ = TryQueue(item); } /// Thread-safe verbose bridge for asynchronous cmdlet code. public new void WriteVerbose(string message) { ThrowIfStopped(); + var item = new PipelineItem(message, PipelineType.Verbose); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + if (CanAccessPipelineDirectly) { using var pipelineContext = EnterDirectPipelineAccess(); @@ -599,13 +637,20 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(message, PipelineType.Verbose)); + _ = TryQueue(item); } /// Thread-safe debug bridge for asynchronous cmdlet code. public new void WriteDebug(string message) { ThrowIfStopped(); + var item = new PipelineItem(message, PipelineType.Debug); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + if (CanAccessPipelineDirectly) { using var pipelineContext = EnterDirectPipelineAccess(); @@ -616,13 +661,20 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(message, PipelineType.Debug)); + _ = TryQueue(item); } /// Thread-safe command-detail bridge for asynchronous cmdlet code. public new void WriteCommandDetail(string text) { ThrowIfStopped(); + var item = new PipelineItem(text, PipelineType.CommandDetail); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + if (CanAccessPipelineDirectly) { using var pipelineContext = EnterDirectPipelineAccess(); @@ -633,13 +685,22 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(text, PipelineType.CommandDetail)); + _ = TryQueue(item); } /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(InformationRecord informationRecord) { ThrowIfStopped(); + var item = new PipelineItem( + SnapshotInformationRecord(informationRecord), + PipelineType.Information); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + if (CanAccessPipelineDirectly) { using var pipelineContext = EnterDirectPipelineAccess(); @@ -650,13 +711,22 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem(informationRecord, PipelineType.Information)); + _ = TryQueue(item); } /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(object messageData, string[]? tags) { ThrowIfStopped(); + var item = new PipelineItem( + (messageData, tags is null ? null : (string[])tags.Clone()), + PipelineType.InformationWithTags); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + if (CanAccessPipelineDirectly) { using var pipelineContext = EnterDirectPipelineAccess(); @@ -667,8 +737,6 @@ protected override void StopProcessing() if (Volatile.Read(ref _currentOutPipe) is null) return; - _ = TryQueue(new PipelineItem( - (messageData, tags is null ? null : (string[])tags.Clone()), - PipelineType.InformationWithTags)); + _ = TryQueue(item); } } diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs new file mode 100644 index 000000000..1f4eed035 --- /dev/null +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs @@ -0,0 +1,578 @@ +using System.Collections.ObjectModel; +using System.Globalization; +using System.Management.Automation; +using System.Management.Automation.Host; +using System.Management.Automation.Runspaces; +using System.Security; +using System.Threading; +using System.Threading.Tasks; +using PSPublishModule; + +namespace PowerForge.Tests; + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedCallback")] +public sealed class TestAsyncCapturedCallbackCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + var writeOutput = CapturePipelineWriter(); + var streams = CapturePipelineStreams(); + var completed = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + ThreadPool.UnsafeQueueUserWorkItem( + _ => + { + writeOutput("callback-output"); + streams.WriteWarning("callback-warning"); + completed.TrySetResult(true); + }, + null); + await completed.Task.ConfigureAwait(false); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCommandDetail")] +public sealed class TestAsyncCommandDetailCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteCommandDetail("worker-detail"); + WriteObject("completed"); + } +} + +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncReentrantPump", + SupportsShouldProcess = true)] +public sealed class TestAsyncReentrantPumpCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(new WarningEnumerable(this), enumerateCollection: true); + WriteWarning("after-enumeration"); + } + + private sealed class WarningEnumerable : IEnumerable + { + private readonly TestAsyncReentrantPumpCommand _command; + + public WarningEnumerable(TestAsyncReentrantPumpCommand command) + { + _command = command; + } + + public IEnumerator GetEnumerator() + { + Assert.True( + _command.ShouldProcess( + "enumerated-target")); + Task.Run(() => _command.WriteWarning("during-enumeration")).GetAwaiter().GetResult(); + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncReentrantFifo")] +public sealed class TestAsyncReentrantFifoCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(new WarningEnumerable(this), enumerateCollection: true); + } + + private sealed class WarningEnumerable : IEnumerable + { + private readonly TestAsyncReentrantFifoCommand _command; + + public WarningEnumerable(TestAsyncReentrantFifoCommand command) + { + _command = command; + } + + public IEnumerator GetEnumerator() + { + Task.Run(() => _command.WriteWarning("queued-first")).GetAwaiter().GetResult(); + _command.WriteWarning("direct-second"); + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedEnumeration")] +public sealed class TestAsyncCapturedEnumerationCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var streams = CapturePipelineStreams(); + WriteObject( + new CapturedWarningEnumerable( + message => streams.WriteWarning(message)), + enumerateCollection: true); + } + + private sealed class CapturedWarningEnumerable : IEnumerable + { + private readonly Action _writeWarning; + + public CapturedWarningEnumerable(Action writeWarning) + { + _writeWarning = writeWarning; + } + + public IEnumerator GetEnumerator() + { + using var completed = new ManualResetEventSlim(); + Exception? callbackException = null; + ThreadPool.UnsafeQueueUserWorkItem( + _ => + { + try + { + _writeWarning("captured-during-enumeration"); + } + catch (Exception exception) + { + callbackException = exception; + } + finally + { + completed.Set(); + } + }, + null); + Assert.True( + completed.Wait(TimeSpan.FromSeconds(5)), + "The context-free callback did not complete in time."); + if (callbackException is not null) + throw callbackException; + + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDerivedEndProcessing")] +public sealed class TestAsyncDerivedEndProcessingCommand : AsyncPSCmdlet +{ + protected override void EndProcessing() + { + WriteObject("derived-end"); + base.EndProcessing(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyShouldProcess", SupportsShouldProcess = true)] +public sealed class TestAsyncEarlyShouldProcessCommand : AsyncPSCmdlet +{ + private bool _approved; + + protected override void ProcessRecord() + { + WriteObject("early-output"); + _approved = ShouldProcess("target"); + base.ProcessRecord(); + } + + protected override Task ProcessRecordAsync() + { + WriteObject(_approved); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncConstructorWrite")] +public sealed class TestAsyncConstructorWriteCommand : AsyncPSCmdlet +{ + public TestAsyncConstructorWriteCommand() + => WriteWarning("constructor output must not reach PowerShell"); + + protected override Task ProcessRecordAsync() + { + WriteObject("completed"); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncProgressSnapshot")] +public sealed class TestAsyncProgressSnapshotCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var progress = new ProgressRecord(17, "snapshot", "queued") + { + PercentComplete = 10 + }; + var totalProperty = progress.GetType().GetProperty("Total"); + if (totalProperty is not null) + { + totalProperty.SetValue( + progress, + Convert.ChangeType(42, totalProperty.PropertyType, CultureInfo.InvariantCulture)); + } + WriteProgress(progress); + progress.PercentComplete = 90; + if (totalProperty is not null) + { + totalProperty.SetValue( + progress, + Convert.ChangeType(99, totalProperty.PropertyType, CultureInfo.InvariantCulture)); + } + WriteObject("completed"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncPipelineFailureWithThrowingCancellation")] +public sealed class TestAsyncPipelineFailureWithThrowingCancellationCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + _ = CancelToken.Register( + static () => throw new InvalidOperationException("cancellation callback failed")); + await Task.Yield(); + ThrowTerminatingError(new ErrorRecord( + new InvalidOperationException("pipeline failure"), + "PipelineFailure", + ErrorCategory.InvalidOperation, + targetObject: null)); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringHook")] +public sealed class TestAsyncDisposeDuringHookCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _tokenRead = new(); + + public static ManualResetEventSlim TokenRead => _tokenRead; + + public static bool ReadTokenAfterDispose { get; private set; } + + public static Exception? TokenReadException { get; private set; } + + public static void Reset() + { + _tokenRead.Dispose(); + _tokenRead = new ManualResetEventSlim(); + ReadTokenAfterDispose = false; + TokenReadException = null; + } + + protected override async Task ProcessRecordAsync() + { + Dispose(); + await Task.Yield(); + try + { + var token = CancelToken; + ReadTokenAfterDispose = true; + _tokenRead.Set(); + await Task.Delay(Timeout.InfiniteTimeSpan, token); + } + catch (OperationCanceledException) + { + throw; + } + catch (Exception exception) + { + TokenReadException = exception; + _tokenRead.Set(); + throw; + } + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringCancellation")] +public sealed class TestAsyncDisposeDuringCancellationCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _callbackCompleted = new(); + + public static ManualResetEventSlim CallbackCompleted => _callbackCompleted; + + public static bool TokenReadDuringCancellation { get; private set; } + + public static Exception? TokenReadException { get; private set; } + + public static void Reset() + { + _callbackCompleted.Dispose(); + _callbackCompleted = new ManualResetEventSlim(); + TokenReadDuringCancellation = false; + TokenReadException = null; + } + + protected override Task ProcessRecordAsync() + { + using var cancellationEntered = new ManualResetEventSlim(); + using var releaseCancellation = new ManualResetEventSlim(); + _ = CancelToken.Register( + () => + { + cancellationEntered.Set(); + releaseCancellation.Wait(TimeSpan.FromSeconds(5)); + try + { + _ = CancelToken; + TokenReadDuringCancellation = true; + } + catch (Exception exception) + { + TokenReadException = exception; + } + finally + { + _callbackCompleted.Set(); + } + }); + + _ = Task.Run(Dispose); + Assert.True( + cancellationEntered.Wait(TimeSpan.FromSeconds(5)), + "Cancellation did not enter the registered callback."); + ThreadPool.QueueUserWorkItem( + _ => + { + Thread.Sleep(100); + releaseCancellation.Set(); + }); + return Task.CompletedTask; + } + + protected override void EndProcessing() + { + // ProcessRecord intentionally disposes the command to exercise the cancellation race. + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDirectPipelineStop")] +public sealed class TestAsyncDirectPipelineStopCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _cancellationObserved = new(); + + public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + + public static void Reset() + { + _cancellationObserved.Dispose(); + _cancellationObserved = new ManualResetEventSlim(); + } + + protected override async Task ProcessRecordAsync() + { + _ = CancelToken.Register(_cancellationObserved.Set); + await Task.CompletedTask; + WriteObject(1); + WriteObject(2); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSuccessfulCompletion")] +public sealed class TestAsyncSuccessfulCompletionCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _cancellationObserved = new(); + + public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + + public static void Reset() + { + _cancellationObserved.Dispose(); + _cancellationObserved = new ManualResetEventSlim(); + } + + protected override async Task ProcessRecordAsync() + { + _ = CancelToken.Register(_cancellationObserved.Set); + await Task.Yield(); + WriteObject("completed"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncLateProgress")] +public sealed class TestAsyncLateProgressCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _callbackCompleted = new(); + private static ManualResetEventSlim _initialized = new(); + private static IProgress? _progress; + + public static ManualResetEventSlim CallbackCompleted => _callbackCompleted; + + public static ManualResetEventSlim Initialized => _initialized; + + public static void Reset() + { + _callbackCompleted.Dispose(); + _initialized.Dispose(); + _callbackCompleted = new ManualResetEventSlim(); + _initialized = new ManualResetEventSlim(); + _progress = null; + } + + public static void ReportAfterStop() + => _progress!.Report(50); + + protected override async Task ProcessRecordAsync() + { + _progress = new Progress(percent => + { + try + { + WriteProgress(new ProgressRecord(1, "late", "after stop") + { + PercentComplete = percent + }); + } + finally + { + _callbackCompleted.Set(); + } + }); + _initialized.Set(); + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousStop")] +public sealed class TestAsyncSynchronousStopCommand : AsyncPSCmdlet +{ + public void InvokeProcessRecord() + => base.ProcessRecord(); + + protected override Task ProcessRecordAsync() + { + Dispose(); + throw new OperationCanceledException("stopped synchronously"); + } +} + +public sealed class ChoiceHost(bool approved, Exception? promptFailure = null) : PSHost +{ + private readonly Guid _id = Guid.NewGuid(); + private readonly ChoiceHostUserInterface _ui = new(approved, promptFailure); + + public override Guid InstanceId => _id; + public override string Name => nameof(ChoiceHost); + public override Version Version => new(1, 0); + public override PSHostUserInterface UI => _ui; + + public IReadOnlyList ProgressRecords => _ui.ProgressRecords; + public override CultureInfo CurrentCulture => CultureInfo.InvariantCulture; + public override CultureInfo CurrentUICulture => CultureInfo.InvariantCulture; + public override void EnterNestedPrompt() { } + public override void ExitNestedPrompt() { } + public override void NotifyBeginApplication() { } + public override void NotifyEndApplication() { } + public override void SetShouldExit(int exitCode) { } +} + +public sealed class ChoiceHostUserInterface(bool approved, Exception? promptFailure) : PSHostUserInterface +{ + private readonly List _progressRecords = new(); + + public IReadOnlyList ProgressRecords => _progressRecords; + + public override PSHostRawUserInterface RawUI => null!; + + public override int PromptForChoice( + string caption, + string message, + Collection choices, + int defaultChoice) + => promptFailure is null ? (approved ? 0 : 1) : throw promptFailure; + + public override string ReadLine() => string.Empty; + public override SecureString ReadLineAsSecureString() => new(); + public override void Write(string value) { } + public override void Write(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value) { } + public override void WriteLine(string value) { } + public override void WriteErrorLine(string value) { } + public override void WriteDebugLine(string message) { } + public override void WriteProgress(long sourceId, ProgressRecord record) + => _progressRecords.Add(record); + public override void WriteVerboseLine(string message) { } + public override void WriteWarningLine(string message) { } + + public override PSCredential PromptForCredential( + string caption, + string message, + string userName, + string targetName) + => new(userName, new SecureString()); + + public override PSCredential PromptForCredential( + string caption, + string message, + string userName, + string targetName, + PSCredentialTypes allowedCredentialTypes, + PSCredentialUIOptions options) + => PromptForCredential(caption, message, userName, targetName); + + public override Dictionary Prompt( + string caption, + string message, + Collection descriptions) + => []; +} + +public sealed class ForwardingTaskScheduler : TaskScheduler +{ + private int _executeNextTaskInline; + private int _queuedTaskCount; + + public int QueuedTaskCount => Volatile.Read(ref _queuedTaskCount); + + public void Run(Action action) + { + Volatile.Write(ref _executeNextTaskInline, 1); + var task = Task.Factory.StartNew( + action, + CancellationToken.None, + TaskCreationOptions.None, + this); + task.GetAwaiter().GetResult(); + } + + protected override IEnumerable GetScheduledTasks() + => Array.Empty(); + + protected override void QueueTask(Task task) + { + if (Interlocked.Exchange(ref _executeNextTaskInline, 0) == 1) + { + Assert.True(TryExecuteTask(task)); + return; + } + + Interlocked.Increment(ref _queuedTaskCount); + ThreadPool.QueueUserWorkItem(_ => TryExecuteTask(task)); + } + + protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) + => false; +} + +public sealed class ForwardingSynchronizationContext : SynchronizationContext +{ + private int _postCount; + + public int PostCount => Volatile.Read(ref _postCount); + + public override void Post(SendOrPostCallback callback, object? state) + { + Interlocked.Increment(ref _postCount); + ThreadPool.QueueUserWorkItem(_ => callback(state)); + } +} + diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs new file mode 100644 index 000000000..cd6d52326 --- /dev/null +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs @@ -0,0 +1,485 @@ +using System.Collections.ObjectModel; +using System.Globalization; +using System.Management.Automation; +using System.Management.Automation.Host; +using System.Management.Automation.Runspaces; +using System.Security; +using System.Threading; +using System.Threading.Tasks; +using PSPublishModule; + +namespace PowerForge.Tests; + +public sealed partial class AsyncPSCmdletTests +{ +[Fact] + public void AsyncPSCmdlet_dispose_without_an_active_hook_does_not_signal_stop() + { + using var command = new TestAsyncDisposableCommand(); + var stoppingToken = command.StoppingToken; + + command.Dispose(); + + Assert.False(stoppingToken.IsCancellationRequested); + } + + [Fact] + public void AsyncPSCmdlet_drops_pre_lifecycle_writes_from_other_threads() + { + using var command = new TestAsyncDisposableCommand(); + Exception? exception = null; + var workerThreadId = 0; + using var completed = new ManualResetEventSlim(); + var thread = new Thread(() => + { + workerThreadId = Environment.CurrentManagedThreadId; + exception = Record.Exception(() => command.WriteWarning("too-early")); + completed.Set(); + }); + + thread.Start(); + + Assert.True(completed.Wait(TimeSpan.FromSeconds(5))); + thread.Join(); + Assert.NotEqual(Environment.CurrentManagedThreadId, workerThreadId); + Assert.Null(exception); + } + + [Fact] + public void AsyncPSCmdlet_finishes_disposal_when_cancellation_callbacks_throw() + { + var command = new TestAsyncDisposableCommand(); + var stoppingToken = command.StoppingToken; + using var registration = stoppingToken.Register( + static () => throw new InvalidOperationException("cancellation callback failed")); + + command.InvokeStopProcessing(); + command.Dispose(); + + Assert.True(stoppingToken.IsCancellationRequested); + var sourceField = typeof(AsyncPSCmdlet).GetField( + "_cancelSource", + System.Reflection.BindingFlags.Instance | System.Reflection.BindingFlags.NonPublic); + var source = Assert.IsType(sourceField!.GetValue(command)); + Assert.Throws(source.Cancel); + } + + [Fact] + public void AsyncPSCmdlet_does_not_deadlock_when_a_synchronous_hook_queues_more_than_the_transport_capacity() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncLargeQueuedOutput", + typeof(TestAsyncLargeQueuedOutputCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncLargeQueuedOutput"); + + powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal(2048, powerShell.Streams.Warning.Count); + } + + [Fact] + public void AsyncPSCmdlet_restores_the_host_context_around_direct_pipeline_writes() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDirectContext", + typeof(TestAsyncDirectContextCommand), + helpFileName: null)); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-ObserveContext", + typeof(TestObserveContextCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncDirectContext") + .AddCommand("Test-ObserveContext"); + + var result = powerShell.Invoke(); + + Assert.Equal("context-output", Assert.Single(result).BaseObject); + Assert.Same(TestAsyncDirectContextCommand.HostContext, TestObserveContextCommand.ObservedContext); + } + + [Fact] + public void AsyncPSCmdlet_accepts_context_free_callbacks_through_a_captured_writer() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCapturedCallback", + typeof(TestAsyncCapturedCallbackCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCapturedCallback"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal("callback-output", Assert.Single(result).BaseObject); + Assert.Equal( + "callback-warning", + Assert.Single(powerShell.Streams.Warning).Message); + } + + [Fact] + public void AsyncPSCmdlet_marshals_command_details_to_the_pipeline_thread() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCommandDetail", + typeof(TestAsyncCommandDetailCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCommandDetail"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal("completed", Assert.Single(result).BaseObject); + } + + [Fact] + public void AsyncPSCmdlet_does_not_recursively_drain_when_enumeration_writes_to_the_pipeline() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncReentrantPump", + typeof(TestAsyncReentrantPumpCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncReentrantPump") + .AddParameter("Confirm", false); + + var result = powerShell.Invoke(); + + Assert.Collection(result, item => Assert.Equal("value", item.BaseObject)); + Assert.Equal( + ["after-enumeration", "during-enumeration"], + powerShell.Streams.Warning.Select(static warning => warning.Message).Order()); + } + + [Fact] + public void AsyncPSCmdlet_preserves_fifo_order_during_reentrant_queue_drains() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncReentrantFifo", + typeof(TestAsyncReentrantFifoCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncReentrantFifo"); + + var result = powerShell.Invoke(); + + Assert.Equal("value", Assert.Single(result).BaseObject); + Assert.Equal( + ["queued-first", "direct-second"], + powerShell.Streams.Warning.Select(static warning => warning.Message)); + } + + [Fact] + public void AsyncPSCmdlet_rejects_context_free_callbacks_after_their_hook_completes() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCapturedEnumeration", + typeof(TestAsyncCapturedEnumerationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCapturedEnumeration"); + + var result = powerShell.Invoke(); + + Assert.Equal("value", Assert.Single(result).BaseObject); + Assert.Empty(powerShell.Streams.Warning); + } + + [Fact] + public void AsyncPSCmdlet_keeps_derived_EndProcessing_active_until_the_base_call() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDerivedEndProcessing", + typeof(TestAsyncDerivedEndProcessingCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncDerivedEndProcessing"); + + var result = powerShell.Invoke(); + + Assert.Equal("derived-end", Assert.Single(result).BaseObject); + } + + [Fact] + public void AsyncPSCmdlet_allows_ShouldProcess_before_the_async_base_hook_starts() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncEarlyShouldProcess", + typeof(TestAsyncEarlyShouldProcessCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncEarlyShouldProcess") + .AddParameter("Confirm", false); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + Assert.Collection( + result, + item => Assert.Equal("early-output", item.BaseObject), + item => Assert.True((bool)item.BaseObject)); + } + + [Fact] + public void AsyncPSCmdlet_drops_constructor_writes_before_pipeline_startup() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncConstructorWrite", + typeof(TestAsyncConstructorWriteCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncConstructorWrite"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Empty(powerShell.Streams.Warning); + Assert.Equal("completed", Assert.Single(result).BaseObject); + } + + [Fact] + public void AsyncPSCmdlet_snapshots_mutable_progress_before_queueing() + { + var sessionState = InitialSessionState.Create(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncProgressSnapshot", + typeof(TestAsyncProgressSnapshotCommand), + helpFileName: null)); + + var host = new ChoiceHost(approved: true); + using var runspace = RunspaceFactory.CreateRunspace(host, sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncProgressSnapshot"); + + var result = powerShell.Invoke(); + + Assert.Equal("completed", Assert.Single(result).BaseObject); + var captured = Assert.Single(host.ProgressRecords); + Assert.Equal(10, captured.PercentComplete); + var totalProperty = captured.GetType().GetProperty("Total"); + if (totalProperty is not null) + { + Assert.Equal( + Convert.ChangeType(42, totalProperty.PropertyType, CultureInfo.InvariantCulture), + totalProperty.GetValue(captured)); + } + } + + [Fact] + public void AsyncPSCmdlet_preserves_pipeline_failure_when_cancellation_callbacks_throw() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncPipelineFailureWithThrowingCancellation", + typeof(TestAsyncPipelineFailureWithThrowingCancellationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncPipelineFailureWithThrowingCancellation"); + + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.StartsWith("PipelineFailure,", exception.ErrorRecord.FullyQualifiedErrorId, StringComparison.Ordinal); + Assert.Equal("pipeline failure", exception.InnerException?.Message); + } + + [Fact] + public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_the_async_hook_exits() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDisposeDuringHook", + typeof(TestAsyncDisposeDuringHookCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncDisposeDuringHook"); + TestAsyncDisposeDuringHookCommand.Reset(); + + powerShell.Invoke(); + + Assert.True( + TestAsyncDisposeDuringHookCommand.TokenRead.Wait(TimeSpan.FromSeconds(5)), + "The asynchronous hook did not finish reading the token in time."); + Assert.True(TestAsyncDisposeDuringHookCommand.ReadTokenAfterDispose); + Assert.Null(TestAsyncDisposeDuringHookCommand.TokenReadException); + } + + [Fact] + public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_Cancel_returns() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDisposeDuringCancellation", + typeof(TestAsyncDisposeDuringCancellationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncDisposeDuringCancellation"); + TestAsyncDisposeDuringCancellationCommand.Reset(); + + powerShell.Invoke(); + + Assert.True( + TestAsyncDisposeDuringCancellationCommand.CallbackCompleted.Wait(TimeSpan.FromSeconds(5)), + "The cancellation callback did not complete in time."); + Assert.True(TestAsyncDisposeDuringCancellationCommand.TokenReadDuringCancellation); + Assert.Null(TestAsyncDisposeDuringCancellationCommand.TokenReadException); + } + + [Fact] + public void AsyncPSCmdlet_cancels_started_work_when_a_direct_write_stops_the_pipeline() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDirectPipelineStop", + typeof(TestAsyncDirectPipelineStopCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncDirectPipelineStop") + .AddCommand("Select-Object") + .AddParameter("First", 1); + TestAsyncDirectPipelineStopCommand.Reset(); + + var result = powerShell.Invoke(); + + Assert.Equal(1, Assert.Single(result).BaseObject); + Assert.True( + TestAsyncDirectPipelineStopCommand.CancellationObserved.Wait(TimeSpan.FromSeconds(5)), + "The started operation was not canceled after the pipeline stopped."); + } + + [Fact] + public void AsyncPSCmdlet_does_not_cancel_after_successful_async_completion() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSuccessfulCompletion", + typeof(TestAsyncSuccessfulCompletionCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncSuccessfulCompletion"); + TestAsyncSuccessfulCompletionCommand.Reset(); + + var result = powerShell.Invoke(); + + Assert.Equal("completed", Assert.Single(result).BaseObject); + Assert.False(TestAsyncSuccessfulCompletionCommand.CancellationObserved.IsSet); + } + + [Fact] + public void AsyncPSCmdlet_suppresses_pipeline_stop_from_posted_progress_callbacks() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncLateProgress", + typeof(TestAsyncLateProgressCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncLateProgress"); + TestAsyncLateProgressCommand.Reset(); + + var invocation = powerShell.BeginInvoke(); + Assert.True( + TestAsyncLateProgressCommand.Initialized.Wait(TimeSpan.FromSeconds(5)), + "The progress callback was not initialized in time."); + powerShell.Stop(); + Assert.Throws(() => powerShell.EndInvoke(invocation)); + + TestAsyncLateProgressCommand.ReportAfterStop(); + + Assert.True( + TestAsyncLateProgressCommand.CallbackCompleted.Wait(TimeSpan.FromSeconds(5)), + "The posted progress callback did not complete in time."); + } + + [Fact] + public void AsyncPSCmdlet_normalizes_synchronous_cancellation_after_stop() + { + using var command = new TestAsyncSynchronousStopCommand(); + + Assert.Throws(command.InvokeProcessRecord); + } +} + diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index c026f94a3..294dee69b 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -128,6 +128,11 @@ protected override Task ProcessRecordAsync() var tags = new[] { "before" }; Task.Run(() => WriteInformation("message", tags)).GetAwaiter().GetResult(); tags[0] = "after"; + + var informationRecord = new InformationRecord("record-message", "record-source"); + informationRecord.Tags.Add("record-before"); + Task.Run(() => WriteInformation(informationRecord)).GetAwaiter().GetResult(); + informationRecord.Tags[0] = "record-after"; return Task.CompletedTask; } } @@ -566,569 +571,3 @@ protected override void ProcessRecord() WriteObject(InputObject); } } - -[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedCallback")] -public sealed class TestAsyncCapturedCallbackCommand : AsyncPSCmdlet -{ - protected override async Task ProcessRecordAsync() - { - var writeOutput = CapturePipelineWriter(); - var streams = CapturePipelineStreams(); - var completed = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); - ThreadPool.UnsafeQueueUserWorkItem( - _ => - { - writeOutput("callback-output"); - streams.WriteWarning("callback-warning"); - completed.TrySetResult(true); - }, - null); - await completed.Task.ConfigureAwait(false); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncCommandDetail")] -public sealed class TestAsyncCommandDetailCommand : AsyncPSCmdlet -{ - protected override async Task ProcessRecordAsync() - { - await Task.Yield(); - WriteCommandDetail("worker-detail"); - WriteObject("completed"); - } -} - -[Cmdlet( - VerbsDiagnostic.Test, - "AsyncReentrantPump", - SupportsShouldProcess = true)] -public sealed class TestAsyncReentrantPumpCommand : AsyncPSCmdlet -{ - protected override async Task ProcessRecordAsync() - { - await Task.Yield(); - WriteObject(new WarningEnumerable(this), enumerateCollection: true); - WriteWarning("after-enumeration"); - } - - private sealed class WarningEnumerable : IEnumerable - { - private readonly TestAsyncReentrantPumpCommand _command; - - public WarningEnumerable(TestAsyncReentrantPumpCommand command) - { - _command = command; - } - - public IEnumerator GetEnumerator() - { - Assert.True( - _command.ShouldProcess( - "enumerated-target")); - Task.Run(() => _command.WriteWarning("during-enumeration")).GetAwaiter().GetResult(); - yield return "value"; - } - - System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() - => GetEnumerator(); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncReentrantFifo")] -public sealed class TestAsyncReentrantFifoCommand : AsyncPSCmdlet -{ - protected override async Task ProcessRecordAsync() - { - await Task.Yield(); - WriteObject(new WarningEnumerable(this), enumerateCollection: true); - } - - private sealed class WarningEnumerable : IEnumerable - { - private readonly TestAsyncReentrantFifoCommand _command; - - public WarningEnumerable(TestAsyncReentrantFifoCommand command) - { - _command = command; - } - - public IEnumerator GetEnumerator() - { - Task.Run(() => _command.WriteWarning("queued-first")).GetAwaiter().GetResult(); - _command.WriteWarning("direct-second"); - yield return "value"; - } - - System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() - => GetEnumerator(); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedEnumeration")] -public sealed class TestAsyncCapturedEnumerationCommand : AsyncPSCmdlet -{ - protected override async Task ProcessRecordAsync() - { - await Task.Yield(); - var streams = CapturePipelineStreams(); - WriteObject( - new CapturedWarningEnumerable( - message => streams.WriteWarning(message)), - enumerateCollection: true); - } - - private sealed class CapturedWarningEnumerable : IEnumerable - { - private readonly Action _writeWarning; - - public CapturedWarningEnumerable(Action writeWarning) - { - _writeWarning = writeWarning; - } - - public IEnumerator GetEnumerator() - { - using var completed = new ManualResetEventSlim(); - Exception? callbackException = null; - ThreadPool.UnsafeQueueUserWorkItem( - _ => - { - try - { - _writeWarning("captured-during-enumeration"); - } - catch (Exception exception) - { - callbackException = exception; - } - finally - { - completed.Set(); - } - }, - null); - Assert.True( - completed.Wait(TimeSpan.FromSeconds(5)), - "The context-free callback did not complete in time."); - if (callbackException is not null) - throw callbackException; - - yield return "value"; - } - - System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() - => GetEnumerator(); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncDerivedEndProcessing")] -public sealed class TestAsyncDerivedEndProcessingCommand : AsyncPSCmdlet -{ - protected override void EndProcessing() - { - base.EndProcessing(); - WriteObject("derived-end"); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyShouldProcess", SupportsShouldProcess = true)] -public sealed class TestAsyncEarlyShouldProcessCommand : AsyncPSCmdlet -{ - private bool _approved; - - protected override void ProcessRecord() - { - WriteObject("early-output"); - _approved = ShouldProcess("target"); - base.ProcessRecord(); - } - - protected override Task ProcessRecordAsync() - { - WriteObject(_approved); - return Task.CompletedTask; - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncConstructorWrite")] -public sealed class TestAsyncConstructorWriteCommand : AsyncPSCmdlet -{ - public TestAsyncConstructorWriteCommand() - => WriteWarning("constructor output must not reach PowerShell"); - - protected override Task ProcessRecordAsync() - { - WriteObject("completed"); - return Task.CompletedTask; - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncProgressSnapshot")] -public sealed class TestAsyncProgressSnapshotCommand : AsyncPSCmdlet -{ - protected override async Task ProcessRecordAsync() - { - await Task.Yield(); - var progress = new ProgressRecord(17, "snapshot", "queued") - { - PercentComplete = 10 - }; - var totalProperty = progress.GetType().GetProperty("Total"); - if (totalProperty is not null) - { - totalProperty.SetValue( - progress, - Convert.ChangeType(42, totalProperty.PropertyType, CultureInfo.InvariantCulture)); - } - WriteProgress(progress); - progress.PercentComplete = 90; - if (totalProperty is not null) - { - totalProperty.SetValue( - progress, - Convert.ChangeType(99, totalProperty.PropertyType, CultureInfo.InvariantCulture)); - } - WriteObject("completed"); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncPipelineFailureWithThrowingCancellation")] -public sealed class TestAsyncPipelineFailureWithThrowingCancellationCommand : AsyncPSCmdlet -{ - protected override async Task ProcessRecordAsync() - { - _ = CancelToken.Register( - static () => throw new InvalidOperationException("cancellation callback failed")); - await Task.Yield(); - ThrowTerminatingError(new ErrorRecord( - new InvalidOperationException("pipeline failure"), - "PipelineFailure", - ErrorCategory.InvalidOperation, - targetObject: null)); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringHook")] -public sealed class TestAsyncDisposeDuringHookCommand : AsyncPSCmdlet -{ - private static ManualResetEventSlim _tokenRead = new(); - - public static ManualResetEventSlim TokenRead => _tokenRead; - - public static bool ReadTokenAfterDispose { get; private set; } - - public static Exception? TokenReadException { get; private set; } - - public static void Reset() - { - _tokenRead.Dispose(); - _tokenRead = new ManualResetEventSlim(); - ReadTokenAfterDispose = false; - TokenReadException = null; - } - - protected override async Task ProcessRecordAsync() - { - Dispose(); - await Task.Yield(); - try - { - var token = CancelToken; - ReadTokenAfterDispose = true; - _tokenRead.Set(); - await Task.Delay(Timeout.InfiniteTimeSpan, token); - } - catch (OperationCanceledException) - { - throw; - } - catch (Exception exception) - { - TokenReadException = exception; - _tokenRead.Set(); - throw; - } - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringCancellation")] -public sealed class TestAsyncDisposeDuringCancellationCommand : AsyncPSCmdlet -{ - private static ManualResetEventSlim _callbackCompleted = new(); - - public static ManualResetEventSlim CallbackCompleted => _callbackCompleted; - - public static bool TokenReadDuringCancellation { get; private set; } - - public static Exception? TokenReadException { get; private set; } - - public static void Reset() - { - _callbackCompleted.Dispose(); - _callbackCompleted = new ManualResetEventSlim(); - TokenReadDuringCancellation = false; - TokenReadException = null; - } - - protected override Task ProcessRecordAsync() - { - using var cancellationEntered = new ManualResetEventSlim(); - using var releaseCancellation = new ManualResetEventSlim(); - _ = CancelToken.Register( - () => - { - cancellationEntered.Set(); - releaseCancellation.Wait(TimeSpan.FromSeconds(5)); - try - { - _ = CancelToken; - TokenReadDuringCancellation = true; - } - catch (Exception exception) - { - TokenReadException = exception; - } - finally - { - _callbackCompleted.Set(); - } - }); - - _ = Task.Run(Dispose); - Assert.True( - cancellationEntered.Wait(TimeSpan.FromSeconds(5)), - "Cancellation did not enter the registered callback."); - ThreadPool.QueueUserWorkItem( - _ => - { - Thread.Sleep(100); - releaseCancellation.Set(); - }); - return Task.CompletedTask; - } - - protected override void EndProcessing() - { - // ProcessRecord intentionally disposes the command to exercise the cancellation race. - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncDirectPipelineStop")] -public sealed class TestAsyncDirectPipelineStopCommand : AsyncPSCmdlet -{ - private static ManualResetEventSlim _cancellationObserved = new(); - - public static ManualResetEventSlim CancellationObserved => _cancellationObserved; - - public static void Reset() - { - _cancellationObserved.Dispose(); - _cancellationObserved = new ManualResetEventSlim(); - } - - protected override async Task ProcessRecordAsync() - { - _ = CancelToken.Register(_cancellationObserved.Set); - await Task.CompletedTask; - WriteObject(1); - WriteObject(2); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncSuccessfulCompletion")] -public sealed class TestAsyncSuccessfulCompletionCommand : AsyncPSCmdlet -{ - private static ManualResetEventSlim _cancellationObserved = new(); - - public static ManualResetEventSlim CancellationObserved => _cancellationObserved; - - public static void Reset() - { - _cancellationObserved.Dispose(); - _cancellationObserved = new ManualResetEventSlim(); - } - - protected override async Task ProcessRecordAsync() - { - _ = CancelToken.Register(_cancellationObserved.Set); - await Task.Yield(); - WriteObject("completed"); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncLateProgress")] -public sealed class TestAsyncLateProgressCommand : AsyncPSCmdlet -{ - private static ManualResetEventSlim _callbackCompleted = new(); - private static ManualResetEventSlim _initialized = new(); - private static IProgress? _progress; - - public static ManualResetEventSlim CallbackCompleted => _callbackCompleted; - - public static ManualResetEventSlim Initialized => _initialized; - - public static void Reset() - { - _callbackCompleted.Dispose(); - _initialized.Dispose(); - _callbackCompleted = new ManualResetEventSlim(); - _initialized = new ManualResetEventSlim(); - _progress = null; - } - - public static void ReportAfterStop() - => _progress!.Report(50); - - protected override async Task ProcessRecordAsync() - { - _progress = new Progress(percent => - { - try - { - WriteProgress(new ProgressRecord(1, "late", "after stop") - { - PercentComplete = percent - }); - } - finally - { - _callbackCompleted.Set(); - } - }); - _initialized.Set(); - await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); - } -} - -[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousStop")] -public sealed class TestAsyncSynchronousStopCommand : AsyncPSCmdlet -{ - public void InvokeProcessRecord() - => base.ProcessRecord(); - - protected override Task ProcessRecordAsync() - { - Dispose(); - throw new OperationCanceledException("stopped synchronously"); - } -} - -public sealed class ChoiceHost(bool approved, Exception? promptFailure = null) : PSHost -{ - private readonly Guid _id = Guid.NewGuid(); - private readonly ChoiceHostUserInterface _ui = new(approved, promptFailure); - - public override Guid InstanceId => _id; - public override string Name => nameof(ChoiceHost); - public override Version Version => new(1, 0); - public override PSHostUserInterface UI => _ui; - - public IReadOnlyList ProgressRecords => _ui.ProgressRecords; - public override CultureInfo CurrentCulture => CultureInfo.InvariantCulture; - public override CultureInfo CurrentUICulture => CultureInfo.InvariantCulture; - public override void EnterNestedPrompt() { } - public override void ExitNestedPrompt() { } - public override void NotifyBeginApplication() { } - public override void NotifyEndApplication() { } - public override void SetShouldExit(int exitCode) { } -} - -public sealed class ChoiceHostUserInterface(bool approved, Exception? promptFailure) : PSHostUserInterface -{ - private readonly List _progressRecords = new(); - - public IReadOnlyList ProgressRecords => _progressRecords; - - public override PSHostRawUserInterface RawUI => null!; - - public override int PromptForChoice( - string caption, - string message, - Collection choices, - int defaultChoice) - => promptFailure is null ? (approved ? 0 : 1) : throw promptFailure; - - public override string ReadLine() => string.Empty; - public override SecureString ReadLineAsSecureString() => new(); - public override void Write(string value) { } - public override void Write(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value) { } - public override void WriteLine(string value) { } - public override void WriteErrorLine(string value) { } - public override void WriteDebugLine(string message) { } - public override void WriteProgress(long sourceId, ProgressRecord record) - => _progressRecords.Add(record); - public override void WriteVerboseLine(string message) { } - public override void WriteWarningLine(string message) { } - - public override PSCredential PromptForCredential( - string caption, - string message, - string userName, - string targetName) - => new(userName, new SecureString()); - - public override PSCredential PromptForCredential( - string caption, - string message, - string userName, - string targetName, - PSCredentialTypes allowedCredentialTypes, - PSCredentialUIOptions options) - => PromptForCredential(caption, message, userName, targetName); - - public override Dictionary Prompt( - string caption, - string message, - Collection descriptions) - => []; -} - -public sealed class ForwardingTaskScheduler : TaskScheduler -{ - private int _executeNextTaskInline; - private int _queuedTaskCount; - - public int QueuedTaskCount => Volatile.Read(ref _queuedTaskCount); - - public void Run(Action action) - { - Volatile.Write(ref _executeNextTaskInline, 1); - var task = Task.Factory.StartNew( - action, - CancellationToken.None, - TaskCreationOptions.None, - this); - task.GetAwaiter().GetResult(); - } - - protected override IEnumerable GetScheduledTasks() - => Array.Empty(); - - protected override void QueueTask(Task task) - { - if (Interlocked.Exchange(ref _executeNextTaskInline, 0) == 1) - { - Assert.True(TryExecuteTask(task)); - return; - } - - Interlocked.Increment(ref _queuedTaskCount); - ThreadPool.QueueUserWorkItem(_ => TryExecuteTask(task)); - } - - protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) - => false; -} - -public sealed class ForwardingSynchronizationContext : SynchronizationContext -{ - private int _postCount; - - public int PostCount => Volatile.Read(ref _postCount); - - public override void Post(SendOrPostCallback callback, object? state) - { - Interlocked.Increment(ref _postCount); - ThreadPool.QueueUserWorkItem(_ => callback(state)); - } -} diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index bfc56d3cb..c389aeb65 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -10,7 +10,7 @@ namespace PowerForge.Tests; -public sealed class AsyncPSCmdletTests +public sealed partial class AsyncPSCmdletTests { [Fact] public void AsyncPSCmdlet_starts_hooks_on_the_pipeline_thread_and_pumps_after_await() @@ -142,9 +142,19 @@ public void AsyncPSCmdlet_snapshots_information_tags_before_queueing() powerShell.Invoke(); Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); - var record = Assert.Single(powerShell.Streams.Information); - Assert.Equal("message", record.MessageData); - Assert.Equal(["before"], record.Tags); + Assert.Collection( + powerShell.Streams.Information, + record => + { + Assert.Equal("message", record.MessageData); + Assert.Equal(["before"], record.Tags); + }, + record => + { + Assert.Equal("record-message", record.MessageData); + Assert.Equal("record-source", record.Source); + Assert.Equal(["record-before"], record.Tags); + }); } [Fact] @@ -507,475 +517,4 @@ TestAsyncLateInteractionCommand.LateInteractionException is $"Unexpected late-interaction result: {TestAsyncLateInteractionCommand.LateInteractionException}"); } - [Fact] - public void AsyncPSCmdlet_dispose_without_an_active_hook_does_not_signal_stop() - { - using var command = new TestAsyncDisposableCommand(); - var stoppingToken = command.StoppingToken; - - command.Dispose(); - - Assert.False(stoppingToken.IsCancellationRequested); - } - - [Fact] - public void AsyncPSCmdlet_drops_pre_lifecycle_writes_from_other_threads() - { - using var command = new TestAsyncDisposableCommand(); - Exception? exception = null; - var workerThreadId = 0; - using var completed = new ManualResetEventSlim(); - var thread = new Thread(() => - { - workerThreadId = Environment.CurrentManagedThreadId; - exception = Record.Exception(() => command.WriteWarning("too-early")); - completed.Set(); - }); - - thread.Start(); - - Assert.True(completed.Wait(TimeSpan.FromSeconds(5))); - thread.Join(); - Assert.NotEqual(Environment.CurrentManagedThreadId, workerThreadId); - Assert.Null(exception); - } - - [Fact] - public void AsyncPSCmdlet_finishes_disposal_when_cancellation_callbacks_throw() - { - var command = new TestAsyncDisposableCommand(); - var stoppingToken = command.StoppingToken; - using var registration = stoppingToken.Register( - static () => throw new InvalidOperationException("cancellation callback failed")); - - command.InvokeStopProcessing(); - command.Dispose(); - - Assert.True(stoppingToken.IsCancellationRequested); - var sourceField = typeof(AsyncPSCmdlet).GetField( - "_cancelSource", - System.Reflection.BindingFlags.Instance | System.Reflection.BindingFlags.NonPublic); - var source = Assert.IsType(sourceField!.GetValue(command)); - Assert.Throws(source.Cancel); - } - - [Fact] - public void AsyncPSCmdlet_does_not_deadlock_when_a_synchronous_hook_queues_more_than_the_transport_capacity() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncLargeQueuedOutput", - typeof(TestAsyncLargeQueuedOutputCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncLargeQueuedOutput"); - - powerShell.Invoke(); - - Assert.False(powerShell.HadErrors); - Assert.Equal(2048, powerShell.Streams.Warning.Count); - } - - [Fact] - public void AsyncPSCmdlet_restores_the_host_context_around_direct_pipeline_writes() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncDirectContext", - typeof(TestAsyncDirectContextCommand), - helpFileName: null)); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-ObserveContext", - typeof(TestObserveContextCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell - .AddCommand("Test-AsyncDirectContext") - .AddCommand("Test-ObserveContext"); - - var result = powerShell.Invoke(); - - Assert.Equal("context-output", Assert.Single(result).BaseObject); - Assert.Same(TestAsyncDirectContextCommand.HostContext, TestObserveContextCommand.ObservedContext); - } - - [Fact] - public void AsyncPSCmdlet_accepts_context_free_callbacks_through_a_captured_writer() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncCapturedCallback", - typeof(TestAsyncCapturedCallbackCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncCapturedCallback"); - - var result = powerShell.Invoke(); - - Assert.False(powerShell.HadErrors); - Assert.Equal("callback-output", Assert.Single(result).BaseObject); - Assert.Equal( - "callback-warning", - Assert.Single(powerShell.Streams.Warning).Message); - } - - [Fact] - public void AsyncPSCmdlet_marshals_command_details_to_the_pipeline_thread() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncCommandDetail", - typeof(TestAsyncCommandDetailCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncCommandDetail"); - - var result = powerShell.Invoke(); - - Assert.False(powerShell.HadErrors); - Assert.Equal("completed", Assert.Single(result).BaseObject); - } - - [Fact] - public void AsyncPSCmdlet_does_not_recursively_drain_when_enumeration_writes_to_the_pipeline() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncReentrantPump", - typeof(TestAsyncReentrantPumpCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell - .AddCommand("Test-AsyncReentrantPump") - .AddParameter("Confirm", false); - - var result = powerShell.Invoke(); - - Assert.Collection(result, item => Assert.Equal("value", item.BaseObject)); - Assert.Equal( - ["after-enumeration", "during-enumeration"], - powerShell.Streams.Warning.Select(static warning => warning.Message).Order()); - } - - [Fact] - public void AsyncPSCmdlet_preserves_fifo_order_during_reentrant_queue_drains() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncReentrantFifo", - typeof(TestAsyncReentrantFifoCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncReentrantFifo"); - - var result = powerShell.Invoke(); - - Assert.Equal("value", Assert.Single(result).BaseObject); - Assert.Equal( - ["queued-first", "direct-second"], - powerShell.Streams.Warning.Select(static warning => warning.Message)); - } - - [Fact] - public void AsyncPSCmdlet_accepts_context_free_callbacks_while_their_lazy_item_is_pumped() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncCapturedEnumeration", - typeof(TestAsyncCapturedEnumerationCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncCapturedEnumeration"); - - var result = powerShell.Invoke(); - - Assert.Equal("value", Assert.Single(result).BaseObject); - Assert.Equal( - "captured-during-enumeration", - Assert.Single(powerShell.Streams.Warning).Message); - } - - [Fact] - public void AsyncPSCmdlet_keeps_derived_EndProcessing_active_after_the_base_call() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncDerivedEndProcessing", - typeof(TestAsyncDerivedEndProcessingCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncDerivedEndProcessing"); - - var result = powerShell.Invoke(); - - Assert.Equal("derived-end", Assert.Single(result).BaseObject); - } - - [Fact] - public void AsyncPSCmdlet_allows_ShouldProcess_before_the_async_base_hook_starts() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncEarlyShouldProcess", - typeof(TestAsyncEarlyShouldProcessCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell - .AddCommand("Test-AsyncEarlyShouldProcess") - .AddParameter("Confirm", false); - - var result = powerShell.Invoke(); - - Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); - Assert.Collection( - result, - item => Assert.Equal("early-output", item.BaseObject), - item => Assert.True((bool)item.BaseObject)); - } - - [Fact] - public void AsyncPSCmdlet_drops_constructor_writes_before_pipeline_startup() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncConstructorWrite", - typeof(TestAsyncConstructorWriteCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncConstructorWrite"); - - var result = powerShell.Invoke(); - - Assert.False(powerShell.HadErrors); - Assert.Empty(powerShell.Streams.Warning); - Assert.Equal("completed", Assert.Single(result).BaseObject); - } - - [Fact] - public void AsyncPSCmdlet_snapshots_mutable_progress_before_queueing() - { - var sessionState = InitialSessionState.Create(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncProgressSnapshot", - typeof(TestAsyncProgressSnapshotCommand), - helpFileName: null)); - - var host = new ChoiceHost(approved: true); - using var runspace = RunspaceFactory.CreateRunspace(host, sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncProgressSnapshot"); - - var result = powerShell.Invoke(); - - Assert.Equal("completed", Assert.Single(result).BaseObject); - var captured = Assert.Single(host.ProgressRecords); - Assert.Equal(10, captured.PercentComplete); - var totalProperty = captured.GetType().GetProperty("Total"); - if (totalProperty is not null) - { - Assert.Equal( - Convert.ChangeType(42, totalProperty.PropertyType, CultureInfo.InvariantCulture), - totalProperty.GetValue(captured)); - } - } - - [Fact] - public void AsyncPSCmdlet_preserves_pipeline_failure_when_cancellation_callbacks_throw() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncPipelineFailureWithThrowingCancellation", - typeof(TestAsyncPipelineFailureWithThrowingCancellationCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncPipelineFailureWithThrowingCancellation"); - - var exception = Assert.Throws(() => powerShell.Invoke()); - - Assert.StartsWith("PipelineFailure,", exception.ErrorRecord.FullyQualifiedErrorId, StringComparison.Ordinal); - Assert.Equal("pipeline failure", exception.InnerException?.Message); - } - - [Fact] - public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_the_async_hook_exits() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncDisposeDuringHook", - typeof(TestAsyncDisposeDuringHookCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncDisposeDuringHook"); - TestAsyncDisposeDuringHookCommand.Reset(); - - powerShell.Invoke(); - - Assert.True( - TestAsyncDisposeDuringHookCommand.TokenRead.Wait(TimeSpan.FromSeconds(5)), - "The asynchronous hook did not finish reading the token in time."); - Assert.True(TestAsyncDisposeDuringHookCommand.ReadTokenAfterDispose); - Assert.Null(TestAsyncDisposeDuringHookCommand.TokenReadException); - } - - [Fact] - public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_Cancel_returns() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncDisposeDuringCancellation", - typeof(TestAsyncDisposeDuringCancellationCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncDisposeDuringCancellation"); - TestAsyncDisposeDuringCancellationCommand.Reset(); - - powerShell.Invoke(); - - Assert.True( - TestAsyncDisposeDuringCancellationCommand.CallbackCompleted.Wait(TimeSpan.FromSeconds(5)), - "The cancellation callback did not complete in time."); - Assert.True(TestAsyncDisposeDuringCancellationCommand.TokenReadDuringCancellation); - Assert.Null(TestAsyncDisposeDuringCancellationCommand.TokenReadException); - } - - [Fact] - public void AsyncPSCmdlet_cancels_started_work_when_a_direct_write_stops_the_pipeline() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncDirectPipelineStop", - typeof(TestAsyncDirectPipelineStopCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell - .AddCommand("Test-AsyncDirectPipelineStop") - .AddCommand("Select-Object") - .AddParameter("First", 1); - TestAsyncDirectPipelineStopCommand.Reset(); - - var result = powerShell.Invoke(); - - Assert.Equal(1, Assert.Single(result).BaseObject); - Assert.True( - TestAsyncDirectPipelineStopCommand.CancellationObserved.Wait(TimeSpan.FromSeconds(5)), - "The started operation was not canceled after the pipeline stopped."); - } - - [Fact] - public void AsyncPSCmdlet_does_not_cancel_after_successful_async_completion() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncSuccessfulCompletion", - typeof(TestAsyncSuccessfulCompletionCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncSuccessfulCompletion"); - TestAsyncSuccessfulCompletionCommand.Reset(); - - var result = powerShell.Invoke(); - - Assert.Equal("completed", Assert.Single(result).BaseObject); - Assert.False(TestAsyncSuccessfulCompletionCommand.CancellationObserved.IsSet); - } - - [Fact] - public void AsyncPSCmdlet_suppresses_pipeline_stop_from_posted_progress_callbacks() - { - var sessionState = InitialSessionState.CreateDefault(); - sessionState.Commands.Add(new SessionStateCmdletEntry( - "Test-AsyncLateProgress", - typeof(TestAsyncLateProgressCommand), - helpFileName: null)); - - using var runspace = RunspaceFactory.CreateRunspace(sessionState); - runspace.Open(); - using var powerShell = PowerShell.Create(); - powerShell.Runspace = runspace; - powerShell.AddCommand("Test-AsyncLateProgress"); - TestAsyncLateProgressCommand.Reset(); - - var invocation = powerShell.BeginInvoke(); - Assert.True( - TestAsyncLateProgressCommand.Initialized.Wait(TimeSpan.FromSeconds(5)), - "The progress callback was not initialized in time."); - powerShell.Stop(); - Assert.Throws(() => powerShell.EndInvoke(invocation)); - - TestAsyncLateProgressCommand.ReportAfterStop(); - - Assert.True( - TestAsyncLateProgressCommand.CallbackCompleted.Wait(TimeSpan.FromSeconds(5)), - "The posted progress callback did not complete in time."); - } - - [Fact] - public void AsyncPSCmdlet_normalizes_synchronous_cancellation_after_stop() - { - using var command = new TestAsyncSynchronousStopCommand(); - - Assert.Throws(command.InvokeProcessRecord); - } } From 719a81b9fef61433686520bd7210c83c97b5a833 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 08:38:32 +0200 Subject: [PATCH 24/27] Harden synchronous async pipeline races --- .../Cmdlets/AsyncPSCmdlet.Execution.cs | 18 +++-- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 30 ++++++- .../AsyncPSCmdletLifecycleTestCommands.cs | 24 ++++-- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 81 +++++++++++++++++++ PowerForge.Tests/AsyncPSCmdletTests.cs | 80 ++++++++++++++++++ 5 files changed, 216 insertions(+), 17 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs index adcf094da..e95953337 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs @@ -220,13 +220,10 @@ void PumpQueuedItems() if (IsPumpingPipelineItem) return; - int queuedAtEntry; - lock (_hookAdmissionLock) - { - queuedAtEntry = outPipe.Count; - } - - while (queuedAtEntry-- > 0 && outPipe.TryTake(out var item)) + // Both callers close ordinary admission before entering this drain. Only a pipeline + // item that is currently being pumped can enqueue more work through its flow-local + // lease, so continue until that causal tail is empty. + while (outPipe.TryTake(out var item)) PumpItem(item); } @@ -303,6 +300,12 @@ void PumpThroughDirectAccessBarrier() DisposePipeOnce(); } + if (exception is PipelineStoppedException) + { + CancelSource(); + throw; + } + if (exception is OperationCanceledException && _cancelSource.IsCancellationRequested) throw new PipelineStoppedException(); @@ -459,4 +462,3 @@ void PumpThroughDirectAccessBarrier() } } } - diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index a66794952..725f315ca 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -107,10 +107,32 @@ private sealed class PipelineReplyChannel private readonly BlockingCollection _pipe = new(boundedCapacity: 1); private int _owners = 2; private int _pipelineOwner = 1; - private int _requesterOwner = 1; + private int _requesterState = 1; public PipelineReply Take(CancellationToken cancellationToken) - => _pipe.Take(cancellationToken); + { + try + { + return _pipe.Take(cancellationToken); + } + catch (OperationCanceledException) + { + if (Interlocked.CompareExchange(ref _requesterState, 0, 1) == 1) + { + Release(); + throw; + } + + if (Volatile.Read(ref _requesterState) == 2) + { + // Once the pipeline claims the request, the host interaction cannot be canceled. + // Keep observing its reply so cancellation cannot abandon an in-flight prompt. + return _pipe.Take(CancellationToken.None); + } + + throw; + } + } public void Publish(Func createValue) => PublishReply(() => new PipelineReply(createValue())); @@ -126,7 +148,7 @@ private void PublishReply(Func createReply) { try { - if (Volatile.Read(ref _requesterOwner) == 0) + if (Interlocked.CompareExchange(ref _requesterState, 2, 1) != 1) return; PipelineReply reply; @@ -168,7 +190,7 @@ public void Abandon() public void ReleaseRequester() { - if (Interlocked.Exchange(ref _requesterOwner, 0) == 1) + if (Interlocked.Exchange(ref _requesterState, 0) != 0) Release(); } diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs index 1f4eed035..8da5dcc13 100644 --- a/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs @@ -456,10 +456,15 @@ protected override Task ProcessRecordAsync() } } -public sealed class ChoiceHost(bool approved, Exception? promptFailure = null) : PSHost +public sealed class ChoiceHost( + bool approved, + Exception? promptFailure = null, + ManualResetEventSlim? promptEntered = null, + ManualResetEventSlim? promptRelease = null) : PSHost { private readonly Guid _id = Guid.NewGuid(); - private readonly ChoiceHostUserInterface _ui = new(approved, promptFailure); + private readonly ChoiceHostUserInterface _ui = + new(approved, promptFailure, promptEntered, promptRelease); public override Guid InstanceId => _id; public override string Name => nameof(ChoiceHost); @@ -476,7 +481,11 @@ public override void NotifyEndApplication() { } public override void SetShouldExit(int exitCode) { } } -public sealed class ChoiceHostUserInterface(bool approved, Exception? promptFailure) : PSHostUserInterface +public sealed class ChoiceHostUserInterface( + bool approved, + Exception? promptFailure, + ManualResetEventSlim? promptEntered, + ManualResetEventSlim? promptRelease) : PSHostUserInterface { private readonly List _progressRecords = new(); @@ -489,7 +498,13 @@ public override int PromptForChoice( string message, Collection choices, int defaultChoice) - => promptFailure is null ? (approved ? 0 : 1) : throw promptFailure; + { + promptEntered?.Set(); + Assert.True( + promptRelease?.Wait(TimeSpan.FromSeconds(5)) ?? true, + "The test host prompt was not released in time."); + return promptFailure is null ? (approved ? 0 : 1) : throw promptFailure; + } public override string ReadLine() => string.Empty; public override SecureString ReadLineAsSecureString() => new(); @@ -575,4 +590,3 @@ public override void Post(SendOrPostCallback callback, object? state) ThreadPool.QueueUserWorkItem(_ => callback(state)); } } - diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index 294dee69b..d39c02c89 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -171,6 +171,38 @@ protected override Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousReentrantDrain")] +public sealed class TestAsyncSynchronousReentrantDrainCommand : AsyncPSCmdlet +{ + [Parameter] + public SwitchParameter Fail { get; set; } + + protected override Task ProcessRecordAsync() + { + var streams = CapturePipelineStreams(); + WriteObject( + new ReentrantWarningEnumerable(streams.WriteWarning), + enumerateCollection: true); + + if (Fail) + throw new InvalidOperationException("synchronous reentrant drain failure"); + + return Task.CompletedTask; + } + + private sealed class ReentrantWarningEnumerable(Action writeWarning) : IEnumerable + { + public IEnumerator GetEnumerator() + { + writeWarning("reentrant-during-drain"); + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousFailure")] public sealed class TestAsyncSynchronousFailureCommand : AsyncPSCmdlet { @@ -181,6 +213,29 @@ protected override Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousPipelineStop")] +public sealed class TestAsyncSynchronousPipelineStopCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _cancellationObserved = new(); + + public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + + public static void Reset() + { + _cancellationObserved.Dispose(); + _cancellationObserved = new ManualResetEventSlim(); + } + + public void InvokeProcessRecord() + => base.ProcessRecord(); + + protected override Task ProcessRecordAsync() + { + _ = CancelToken.Register(_cancellationObserved.Set); + throw new PipelineStoppedException(); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncNullInformationTags")] public sealed class TestAsyncNullInformationTagsCommand : AsyncPSCmdlet { @@ -412,6 +467,32 @@ protected override async Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncClaimedShouldContinue")] +public sealed class TestAsyncClaimedShouldContinueCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _replyObserved = new(); + private static ManualResetEventSlim _cancellationObserved = new(); + + public static ManualResetEventSlim ReplyObserved => _replyObserved; + public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + + public static void Reset() + { + _replyObserved.Dispose(); + _cancellationObserved.Dispose(); + _replyObserved = new ManualResetEventSlim(); + _cancellationObserved = new ManualResetEventSlim(); + } + + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + _ = CancelToken.Register(_cancellationObserved.Set); + _ = ShouldContinue("Proceed?", "Question"); + _replyObserved.Set(); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncOperationCancellation")] public sealed class TestAsyncOperationCancellationCommand : AsyncPSCmdlet { diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index c389aeb65..0bb12756e 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -204,6 +204,35 @@ public void AsyncPSCmdlet_enumerates_pipeline_thread_collections_before_the_hook item => Assert.Equal(2, item.BaseObject)); } + [Theory] + [InlineData(false)] + [InlineData(true)] + public void AsyncPSCmdlet_drains_reentrant_records_from_synchronous_hooks(bool fail) + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSynchronousReentrantDrain", + typeof(TestAsyncSynchronousReentrantDrainCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncSynchronousReentrantDrain") + .AddParameter("Fail", fail); + + if (fail) + _ = Assert.ThrowsAny(() => powerShell.Invoke()); + else + Assert.Equal("value", Assert.Single(powerShell.Invoke()).BaseObject); + + Assert.Equal( + "reentrant-during-drain", + Assert.Single(powerShell.Streams.Warning).Message); + } + [Fact] public void AsyncPSCmdlet_drains_worker_records_before_a_synchronous_hook_failure() { @@ -225,6 +254,18 @@ public void AsyncPSCmdlet_drains_worker_records_before_a_synchronous_hook_failur Assert.Equal("before-failure", warning.Message); } + [Fact] + public void AsyncPSCmdlet_cancels_the_shared_token_after_a_synchronous_pipeline_stop() + { + using var command = new TestAsyncSynchronousPipelineStopCommand(); + TestAsyncSynchronousPipelineStopCommand.Reset(); + + Assert.Throws(command.InvokeProcessRecord); + Assert.True( + TestAsyncSynchronousPipelineStopCommand.CancellationObserved.Wait(TimeSpan.FromSeconds(5)), + "The shared cancellation token did not observe the synchronous pipeline stop."); + } + [Fact] public void AsyncPSCmdlet_preserves_null_information_tags_after_await() { @@ -449,6 +490,45 @@ public void AsyncPSCmdlet_returns_host_interaction_failures_to_the_requesting_wo Assert.Same(promptFailure, exception.InnerException); } + [Fact] + public async Task AsyncPSCmdlet_keeps_a_claimed_host_reply_observed_during_cancellation() + { + using var promptEntered = new ManualResetEventSlim(); + using var promptRelease = new ManualResetEventSlim(); + var sessionState = InitialSessionState.Create(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncClaimedShouldContinue", + typeof(TestAsyncClaimedShouldContinueCommand), + helpFileName: null)); + + var host = new ChoiceHost( + approved: true, + promptEntered: promptEntered, + promptRelease: promptRelease); + using var runspace = RunspaceFactory.CreateRunspace(host, sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncClaimedShouldContinue"); + TestAsyncClaimedShouldContinueCommand.Reset(); + + var invocation = powerShell.BeginInvoke(); + Assert.True( + promptEntered.Wait(TimeSpan.FromSeconds(5)), + "The host interaction did not start in time."); + var stopTask = Task.Run(powerShell.Stop); + Assert.True( + TestAsyncClaimedShouldContinueCommand.CancellationObserved.Wait(TimeSpan.FromSeconds(5)), + "The claimed requester did not observe pipeline cancellation."); + promptRelease.Set(); + + await stopTask.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.Throws(() => powerShell.EndInvoke(invocation)); + Assert.True( + TestAsyncClaimedShouldContinueCommand.ReplyObserved.Wait(TimeSpan.FromSeconds(5)), + "The claimed requester abandoned the host reply during cancellation."); + } + [Fact] public void AsyncPSCmdlet_preserves_operation_cancellation_that_is_not_a_pipeline_stop() { From cac73cc2e9f7e6e058f08852a60ab89dff1588d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 09:26:42 +0200 Subject: [PATCH 25/27] Harden async lifecycle tail handling --- .../Cmdlets/AsyncPSCmdlet.Execution.cs | 34 ++++++++++- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 10 +++- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 16 ++--- .../AsyncPSCmdletLifecycleTestCommands.cs | 60 ++++++++++++++++++- .../AsyncPSCmdletLifecycleTests.cs | 37 +++++++++++- 5 files changed, 139 insertions(+), 18 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs index e95953337..bf49eabd6 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs @@ -238,14 +238,43 @@ void PumpThroughDirectAccessBarrier() hookGeneration: hookGeneration, dropOnStop: true); if (!TryQueue(barrier)) - return; + { + ThrowIfStopped(); + throw new InvalidOperationException( + "No active PowerShell pipeline is available for direct access."); + } while (outPipe.TryTake(out var item)) { PumpItem(item); if (ReferenceEquals(item, barrier)) + { + while (HasPumpBoundItems()) + { + ThrowIfStopped(); + if (!outPipe.TryTake(out var pumpBoundPredecessor)) + { + throw new InvalidOperationException( + "The PowerShell pipeline closed while causal records were pending."); + } + + PumpItem(pumpBoundPredecessor); + } + return; + } + } + } + + bool HasPumpBoundItems() + { + foreach (var queuedItem in outPipe.ToArray()) + { + if (queuedItem.IsPumpBound) + return true; } + + return false; } Volatile.Write(ref _asyncLifecycleStarted, 1); @@ -418,7 +447,10 @@ void PumpThroughDirectAccessBarrier() if (item.Type == PipelineType.HookCompleted) { while (outPipe.TryTake(out var pumpBoundItem)) + { + ThrowIfStopped(); PumpItem(pumpBoundItem); + } break; } } diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index ba861000a..a8cec1f20 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -277,18 +277,22 @@ private bool TryQueue(PipelineItem item) { item.BindToHook(_hookGeneration.Value); var pumpLease = _pipelinePumpLease.Value; + var isPumpBound = + pumpLease is { IsActive: true } && + item.HookGeneration == pumpLease.Generation; lock (_hookAdmissionLock) { if (item.HookGeneration != 0 && item.HookGeneration != Volatile.Read(ref _acceptingHookWritesGeneration) && - (pumpLease is null || - !pumpLease.IsActive || - item.HookGeneration != pumpLease.Generation)) + !isPumpBound) { item.ReplyPipe?.Reject(); return false; } + if (isPumpBound) + item.BindToPump(); + var outPipe = Volatile.Read(ref _currentOutPipe); if (outPipe is null) return false; diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 725f315ca..d01b75eec 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -233,11 +233,16 @@ public PipelineItem( public bool DropOnStop { get; } + public bool IsPumpBound { get; private set; } + public void BindToHook(long hookGeneration) { if (HookGeneration == 0) HookGeneration = hookGeneration; } + + public void BindToPump() + => IsPumpBound = true; } private sealed class PipelinePumpLease @@ -362,16 +367,7 @@ protected virtual Task ProcessRecordAsync() /// protected override void EndProcessing() - { - try - { - RunBlockInAsync(EndProcessingAsync); - } - finally - { - Volatile.Write(ref _asyncLifecycleCompleted, 1); - } - } + => RunBlockInAsync(EndProcessingAsync); /// Asynchronous end hook. protected virtual Task EndProcessingAsync() diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs index 8da5dcc13..0784347ee 100644 --- a/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs @@ -169,8 +169,66 @@ public sealed class TestAsyncDerivedEndProcessingCommand : AsyncPSCmdlet { protected override void EndProcessing() { - WriteObject("derived-end"); + WriteObject("before-base-end"); base.EndProcessing(); + WriteObject("after-base-end"); + WriteWarning("after-base-warning"); + } +} + +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncDirectBarrierTail", + SupportsShouldProcess = true)] +public sealed class TestAsyncDirectBarrierTailCommand : AsyncPSCmdlet +{ + private static int _tailEnumerated; + + public static bool TailEnumeratedBeforeInteraction { get; private set; } + + public static void Reset() + { + Volatile.Write(ref _tailEnumerated, 0); + TailEnumeratedBeforeInteraction = false; + } + + protected override Task ProcessRecordAsync() + { + var streams = CapturePipelineStreams(); + Task.Run( + () => streams.WriteObject( + new OuterEnumerable(streams), + enumerateCollection: true)) + .GetAwaiter() + .GetResult(); + + _ = ShouldProcess("barrier-target"); + TailEnumeratedBeforeInteraction = Volatile.Read(ref _tailEnumerated) != 0; + return Task.CompletedTask; + } + + private sealed class OuterEnumerable(CapturedPipelineStreams streams) : IEnumerable + { + public IEnumerator GetEnumerator() + { + streams.WriteObject(new TailEnumerable(), enumerateCollection: true); + yield return "outer"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } + + private sealed class TailEnumerable : IEnumerable + { + public IEnumerator GetEnumerator() + { + Volatile.Write(ref _tailEnumerated, 1); + yield return "tail"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); } } diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs index cd6d52326..d5cc38ef9 100644 --- a/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs @@ -227,7 +227,7 @@ public void AsyncPSCmdlet_rejects_context_free_callbacks_after_their_hook_comple } [Fact] - public void AsyncPSCmdlet_keeps_derived_EndProcessing_active_until_the_base_call() + public void AsyncPSCmdlet_keeps_derived_EndProcessing_active_after_the_base_call() { var sessionState = InitialSessionState.CreateDefault(); sessionState.Commands.Add(new SessionStateCmdletEntry( @@ -243,7 +243,39 @@ public void AsyncPSCmdlet_keeps_derived_EndProcessing_active_until_the_base_call var result = powerShell.Invoke(); - Assert.Equal("derived-end", Assert.Single(result).BaseObject); + Assert.Collection( + result, + item => Assert.Equal("before-base-end", item.BaseObject), + item => Assert.Equal("after-base-end", item.BaseObject)); + Assert.Equal( + "after-base-warning", + Assert.Single(powerShell.Streams.Warning).Message); + } + + [Fact] + public void AsyncPSCmdlet_drains_causal_tail_records_before_direct_interactions() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDirectBarrierTail", + typeof(TestAsyncDirectBarrierTailCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncDirectBarrierTail") + .AddParameter("Confirm", false); + TestAsyncDirectBarrierTailCommand.Reset(); + + var result = powerShell.Invoke(); + + Assert.Equal( + ["outer", "tail"], + result.Select(static item => item.BaseObject)); + Assert.True(TestAsyncDirectBarrierTailCommand.TailEnumeratedBeforeInteraction); } [Fact] @@ -482,4 +514,3 @@ public void AsyncPSCmdlet_normalizes_synchronous_cancellation_after_stop() Assert.Throws(command.InvokeProcessRecord); } } - From b26b12153fa0cb5392e9a0894f541bef05ec54d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 10:58:33 +0200 Subject: [PATCH 26/27] Harden async pipeline causality --- .../Cmdlets/AsyncPSCmdlet.Execution.cs | 87 ++++++- .../Cmdlets/AsyncPSCmdlet.Pipeline.cs | 128 +++++++--- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 67 ++++- .../AsyncPSCmdletLifecycleTestCommands.cs | 234 ++++++++++++++++++ .../AsyncPSCmdletLifecycleTests.cs | 155 +++++++++++- PowerForge.Tests/AsyncPSCmdletTestCommands.cs | 13 +- PowerForge.Tests/AsyncPSCmdletTests.cs | 1 + 7 files changed, 626 insertions(+), 59 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs index bf49eabd6..18cdab2ad 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs @@ -19,6 +19,8 @@ private void RunBlockInAsyncCore(Func task) var pipeDisposed = 0; var hookGeneration = Interlocked.Increment(ref _nextHookGeneration); var synchronizationContext = SynchronizationContext.Current; + var hookSynchronizationContext = + new AsyncHookSynchronizationContext(); void ClearPipes() { @@ -69,12 +71,17 @@ void PumpItem(PipelineItem item) var priorItemGeneration = _hookGeneration.Value; var priorPumpLease = _pipelinePumpLease.Value; + var priorSharedPumpLease = + Volatile.Read(ref _currentPipelinePumpLease); var pumpLease = new PipelinePumpLease(item.HookGeneration); try { _ = Interlocked.Increment(ref _pipelinePumpDepth); _hookGeneration.Value = item.HookGeneration; _pipelinePumpLease.Value = pumpLease; + Volatile.Write( + ref _currentPipelinePumpLease, + pumpLease); switch (item.Type) { case PipelineType.Output: @@ -208,13 +215,50 @@ void PumpItem(PipelineItem item) } finally { - pumpLease.Close(); + pumpLease.CloseAndWait(); + Volatile.Write( + ref _currentPipelinePumpLease, + priorSharedPumpLease); _pipelinePumpLease.Value = priorPumpLease; _hookGeneration.Value = priorItemGeneration; _ = Interlocked.Decrement(ref _pipelinePumpDepth); } } + void DrainPumpBoundItemsAfterFailure() + { + while (outPipe.TryTake(out var causalItem)) + { + if (!causalItem.IsPumpBound) + { + causalItem.ReplyPipe?.Reject(); + continue; + } + + try + { + PumpItem(causalItem); + } + catch + { + causalItem.ReplyPipe?.Reject(); + } + } + } + + void PumpItemPreservingCausalFailureRecords(PipelineItem item) + { + try + { + PumpItem(item); + } + catch + { + DrainPumpBoundItemsAfterFailure(); + throw; + } + } + void PumpQueuedItems() { if (IsPumpingPipelineItem) @@ -224,14 +268,12 @@ void PumpQueuedItems() // item that is currently being pumped can enqueue more work through its flow-local // lease, so continue until that causal tail is empty. while (outPipe.TryTake(out var item)) - PumpItem(item); + PumpItemPreservingCausalFailureRecords(item); } void PumpThroughDirectAccessBarrier() { - if (IsPumpingPipelineItem) - return; - + PipelineItem? completionMarker = null; var barrier = new PipelineItem( value: null, PipelineType.DirectAccessBarrier, @@ -246,7 +288,13 @@ void PumpThroughDirectAccessBarrier() while (outPipe.TryTake(out var item)) { - PumpItem(item); + if (item.Type == PipelineType.HookCompleted) + { + completionMarker = item; + continue; + } + + PumpItemPreservingCausalFailureRecords(item); if (ReferenceEquals(item, barrier)) { while (HasPumpBoundItems()) @@ -258,7 +306,21 @@ void PumpThroughDirectAccessBarrier() "The PowerShell pipeline closed while causal records were pending."); } - PumpItem(pumpBoundPredecessor); + if (pumpBoundPredecessor.Type == + PipelineType.HookCompleted) + { + completionMarker = + pumpBoundPredecessor; + continue; + } + + PumpItemPreservingCausalFailureRecords(pumpBoundPredecessor); + } + + if (completionMarker is not null) + { + outPipe.Add( + completionMarker); } return; @@ -291,8 +353,10 @@ bool HasPumpBoundItems() try { Volatile.Write(ref _pipelineSynchronizationContext, synchronizationContext); - SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); + SynchronizationContext.SetSynchronizationContext( + hookSynchronizationContext); _hookGeneration.Value = hookGeneration; + ThrowIfStopped(); if (TaskScheduler.Current == TaskScheduler.Default) { blockTask = task(); @@ -443,13 +507,13 @@ bool HasPumpBoundItems() while (true) { var item = outPipe.Take(CancelToken); - PumpItem(item); + PumpItemPreservingCausalFailureRecords(item); if (item.Type == PipelineType.HookCompleted) { while (outPipe.TryTake(out var pumpBoundItem)) { ThrowIfStopped(); - PumpItem(pumpBoundItem); + PumpItemPreservingCausalFailureRecords(pumpBoundItem); } break; } @@ -473,7 +537,10 @@ bool HasPumpBoundItems() { CompleteAddingIfNeeded(outPipe); if (blockTask.IsCompleted) + { + ClearPipes(); DisposePipeOnce(); + } DeactivateHook(); } diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs index a8cec1f20..9cd5c6bd4 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -2,6 +2,7 @@ using System.Collections.Concurrent; using System.Management.Automation; using System.Reflection; +using System.Runtime.ExceptionServices; using System.Threading; using System.Threading.Tasks; @@ -12,6 +13,9 @@ public abstract partial class AsyncPSCmdlet /// Thread-safe progress bridge for asynchronous cmdlet code. public new void WriteProgress(ProgressRecord progressRecord) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); var item = new PipelineItem(SnapshotProgressRecord(progressRecord), PipelineType.Progress); if (IsPumpingPipelineItem) @@ -75,6 +79,9 @@ private static InformationRecord SnapshotInformationRecord(InformationRecord inf return snapshot; } + private static ErrorRecord SnapshotErrorRecord(ErrorRecord errorRecord) + => new(errorRecord, errorRecord.Exception); + /// Throws when PowerShell has requested cancellation. protected internal void ThrowIfStopped() { @@ -82,6 +89,19 @@ protected internal void ThrowIfStopped() throw new PipelineStoppedException(); } + private bool ShouldDropClosedCanceledStreamWrite() + { + if (!_cancelSource.IsCancellationRequested || + Volatile.Read(ref _currentOutPipe) is not null) + { + return false; + } + + var originatingGeneration = _hookGeneration.Value; + return originatingGeneration == 0 || + originatingGeneration != Volatile.Read(ref _activeHookGeneration); + } + /// public virtual void Dispose() { @@ -144,8 +164,7 @@ private IDisposable EnterDirectPipelineAccess() Volatile.Read(ref _pipelineSynchronizationContext)); try { - if (!IsPumpingPipelineItem) - Volatile.Read(ref _pumpQueuedItems)?.Invoke(); + Volatile.Read(ref _pumpQueuedItems)?.Invoke(); return pipelineContext; } catch @@ -222,8 +241,9 @@ private void GetBlockTaskResult(Task blockTask) throw new PipelineStoppedException(); } + ThrowIfStopped(); if (reply.Rejection is not null) - throw reply.Rejection; + ExceptionDispatchInfo.Capture(reply.Rejection).Throw(); return reply.Value; } @@ -277,47 +297,79 @@ private bool TryQueue(PipelineItem item) { item.BindToHook(_hookGeneration.Value); var pumpLease = _pipelinePumpLease.Value; - var isPumpBound = - pumpLease is { IsActive: true } && - item.HookGeneration == pumpLease.Generation; - lock (_hookAdmissionLock) - { - if (item.HookGeneration != 0 && - item.HookGeneration != Volatile.Read(ref _acceptingHookWritesGeneration) && - !isPumpBound) + if (pumpLease is null) + { + var sharedPumpLease = + Volatile.Read(ref _currentPipelinePumpLease); + if (sharedPumpLease is not null && + (item.HookGeneration == 0 || + item.HookGeneration == sharedPumpLease.Generation)) { - item.ReplyPipe?.Reject(); - return false; + item.BindToHook(sharedPumpLease.Generation); + pumpLease = sharedPumpLease; } + } - if (isPumpBound) - item.BindToPump(); - - var outPipe = Volatile.Read(ref _currentOutPipe); - if (outPipe is null) - return false; - - try - { - outPipe.Add(item, CancelToken); - return true; - } - catch (ObjectDisposedException) - { - return false; - } - catch (InvalidOperationException) - { - return false; - } - catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) + var isPumpBound = pumpLease?.TryClaim(item.HookGeneration) == true; + try + { + lock (_hookAdmissionLock) { - if (item.HookGeneration != 0 && !item.DropOnStop) - throw new PipelineStoppedException(); - - return false; + var acceptingGeneration = + Volatile.Read(ref _acceptingHookWritesGeneration); + if (item.HookGeneration == 0 && + !isPumpBound) + { + if (acceptingGeneration == 0) + { + item.ReplyPipe?.Reject(); + return false; + } + + item.BindToHook(acceptingGeneration); + } + + if (item.HookGeneration != acceptingGeneration && + !isPumpBound) + { + item.ReplyPipe?.Reject(); + return false; + } + + if (isPumpBound) + item.BindToPump(); + + var outPipe = Volatile.Read(ref _currentOutPipe); + if (outPipe is null) + return false; + + try + { + outPipe.Add(item, CancelToken); + return true; + } + catch (ObjectDisposedException) + { + return false; + } + catch (InvalidOperationException) + { + return false; + } + catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) + { + if (item.HookGeneration != 0 && !item.DropOnStop) + throw new PipelineStoppedException(); + + return false; + } } } + finally + { + if (isPumpBound) + pumpLease!.ReleaseClaim(); + } } private void RunBlockInAsync(Func task) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index d01b75eec..195bd19a8 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -247,17 +247,46 @@ public void BindToPump() private sealed class PipelinePumpLease { - private int _active = 1; + private readonly object _sync = new(); + private bool _active = true; + private int _claims; public PipelinePumpLease(long generation) => Generation = generation; public long Generation { get; } - public bool IsActive => Volatile.Read(ref _active) != 0; + public bool TryClaim(long generation) + { + lock (_sync) + { + if (!_active || generation != Generation) + return false; + + _claims++; + return true; + } + } - public void Close() - => Volatile.Write(ref _active, 0); + public void ReleaseClaim() + { + lock (_sync) + { + _claims--; + if (!_active && _claims == 0) + Monitor.PulseAll(_sync); + } + } + + public void CloseAndWait() + { + lock (_sync) + { + _active = false; + while (_claims != 0) + Monitor.Wait(_sync); + } + } } /// @@ -282,7 +311,7 @@ public void WriteObject(object? value, bool enumerateCollection = false) /// Queues an error record for the originating hook. public void WriteError(ErrorRecord errorRecord) - => Queue(errorRecord, PipelineType.Error); + => Queue(SnapshotErrorRecord(errorRecord), PipelineType.Error); /// Queues a warning record for the originating hook. public void WriteWarning(string message) @@ -329,11 +358,11 @@ private void Queue(object? value, PipelineType type) private readonly int _constructionThreadId = Environment.CurrentManagedThreadId; private readonly object _hookAdmissionLock = new(); private readonly object _lifecycleLock = new(); - private static readonly SynchronizationContext HookSynchronizationContext = new AsyncHookSynchronizationContext(); private static readonly TaskScheduler HookTaskScheduler = new AsyncHookTaskScheduler(); private BlockingCollection? _currentOutPipe; private Action? _pumpQueuedItems; private SynchronizationContext? _pipelineSynchronizationContext; + private PipelinePumpLease? _currentPipelinePumpLease; private long _activeHookGeneration; private long _acceptingHookWritesGeneration; private long _nextHookGeneration; @@ -542,6 +571,9 @@ protected override void StopProcessing() /// Thread-safe output bridge for asynchronous cmdlet code. public new void WriteObject(object? sendToPipeline, bool enumerateCollection) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); var item = new PipelineItem( sendToPipeline, @@ -568,8 +600,11 @@ protected override void StopProcessing() /// Thread-safe error bridge for asynchronous cmdlet code. public new void WriteError(ErrorRecord errorRecord) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); - var item = new PipelineItem(errorRecord, PipelineType.Error); + var item = new PipelineItem(SnapshotErrorRecord(errorRecord), PipelineType.Error); if (IsPumpingPipelineItem) { _ = TryQueue(item); @@ -613,6 +648,9 @@ protected override void StopProcessing() /// Thread-safe warning bridge for asynchronous cmdlet code. public new void WriteWarning(string message) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); var item = new PipelineItem(message, PipelineType.Warning); if (IsPumpingPipelineItem) @@ -637,6 +675,9 @@ protected override void StopProcessing() /// Thread-safe verbose bridge for asynchronous cmdlet code. public new void WriteVerbose(string message) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); var item = new PipelineItem(message, PipelineType.Verbose); if (IsPumpingPipelineItem) @@ -661,6 +702,9 @@ protected override void StopProcessing() /// Thread-safe debug bridge for asynchronous cmdlet code. public new void WriteDebug(string message) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); var item = new PipelineItem(message, PipelineType.Debug); if (IsPumpingPipelineItem) @@ -685,6 +729,9 @@ protected override void StopProcessing() /// Thread-safe command-detail bridge for asynchronous cmdlet code. public new void WriteCommandDetail(string text) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); var item = new PipelineItem(text, PipelineType.CommandDetail); if (IsPumpingPipelineItem) @@ -709,6 +756,9 @@ protected override void StopProcessing() /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(InformationRecord informationRecord) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); var item = new PipelineItem( SnapshotInformationRecord(informationRecord), @@ -735,6 +785,9 @@ protected override void StopProcessing() /// Thread-safe information bridge for asynchronous cmdlet code. public new void WriteInformation(object messageData, string[]? tags) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); var item = new PipelineItem( (messageData, tags is null ? null : (string[])tags.Clone()), diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs index 0784347ee..d128916cc 100644 --- a/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs @@ -3,6 +3,7 @@ using System.Management.Automation; using System.Management.Automation.Host; using System.Management.Automation.Runspaces; +using System.Reflection; using System.Security; using System.Threading; using System.Threading.Tasks; @@ -232,6 +233,147 @@ System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() } } +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncReentrantInteractionTail", + SupportsShouldProcess = true)] +public sealed class TestAsyncReentrantInteractionTailCommand : AsyncPSCmdlet +{ + private static int _tailEnumerated; + + public static bool TailEnumeratedBeforeInteraction { get; private set; } + + public static void Reset() + { + Volatile.Write(ref _tailEnumerated, 0); + TailEnumeratedBeforeInteraction = false; + } + + protected override Task ProcessRecordAsync() + { + var streams = + CapturePipelineStreams(); + Task.Run( + () => streams.WriteObject( + new OuterEnumerable(this), + enumerateCollection: true)) + .GetAwaiter() + .GetResult(); + return Task.CompletedTask; + } + + private sealed class OuterEnumerable( + TestAsyncReentrantInteractionTailCommand command) + : IEnumerable + { + public IEnumerator GetEnumerator() + { + command.WriteObject( + new TailEnumerable(), + enumerateCollection: true); + _ = command.ShouldProcess( + "reentrant-target"); + TailEnumeratedBeforeInteraction = + Volatile.Read(ref _tailEnumerated) != 0; + yield return "outer"; + } + + System.Collections.IEnumerator + System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } + + private sealed class TailEnumerable : IEnumerable + { + public IEnumerator GetEnumerator() + { + Volatile.Write(ref _tailEnumerated, 1); + yield return "tail"; + } + + System.Collections.IEnumerator + System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncFailureRecoveryWrite")] +public sealed class TestAsyncFailureRecoveryWriteCommand : AsyncPSCmdlet +{ + public static bool TransportClearedAfterFailure { get; private set; } + + public static void Reset() + => TransportClearedAfterFailure = false; + + protected override void ProcessRecord() + { + try + { + base.ProcessRecord(); + } + catch (Exception) + { + TransportClearedAfterFailure = + typeof(AsyncPSCmdlet) + .GetField( + "_currentOutPipe", + BindingFlags.Instance | + BindingFlags.NonPublic)! + .GetValue(this) is null; + } + } + + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject( + new FailingEnumerable(), + enumerateCollection: true); + } + + private sealed class FailingEnumerable : IEnumerable + { + public IEnumerator GetEnumerator() + { + throw new InvalidOperationException( + "enumeration failed"); +#pragma warning disable CS0162 + yield break; +#pragma warning restore CS0162 + } + + System.Collections.IEnumerator + System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncHookContextIdentity")] +public sealed class TestAsyncHookContextIdentityCommand : AsyncPSCmdlet +{ + private SynchronizationContext? _beginContext; + + protected override Task BeginProcessingAsync() + { + _beginContext = + SynchronizationContext.Current; + return Task.CompletedTask; + } + + protected override Task ProcessRecordAsync() + { + WriteObject( + ReferenceEquals( + _beginContext, + SynchronizationContext.Current)); + return Task.CompletedTask; + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyShouldProcess", SupportsShouldProcess = true)] public sealed class TestAsyncEarlyShouldProcessCommand : AsyncPSCmdlet { @@ -293,6 +435,62 @@ protected override async Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedErrorSnapshot")] +public sealed class TestAsyncCapturedErrorSnapshotCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var streams = CapturePipelineStreams(); + WriteObject(new ErrorSnapshotEnumerable(streams), enumerateCollection: true); + } + + private sealed class ErrorSnapshotEnumerable(CapturedPipelineStreams streams) : IEnumerable + { + public IEnumerator GetEnumerator() + { + var error = new ErrorRecord( + new InvalidOperationException("snapshot failure"), + "CapturedErrorSnapshot", + ErrorCategory.InvalidOperation, + targetObject: null) + { + ErrorDetails = new ErrorDetails("original details") + }; + streams.WriteError(error); + error.ErrorDetails = new ErrorDetails("mutated details"); + yield return "completed"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncEnumeratorFailure")] +public sealed class TestAsyncEnumeratorFailureCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var streams = CapturePipelineStreams(); + WriteObject(new FailingEnumerable(streams), enumerateCollection: true); + } + + private sealed class FailingEnumerable(CapturedPipelineStreams streams) : IEnumerable + { + public IEnumerator GetEnumerator() + { + streams.WriteWarning("before enumeration failure"); + yield return "first"; + throw new InvalidOperationException("enumeration failed"); + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncPipelineFailureWithThrowingCancellation")] public sealed class TestAsyncPipelineFailureWithThrowingCancellationCommand : AsyncPSCmdlet { @@ -462,26 +660,44 @@ public sealed class TestAsyncLateProgressCommand : AsyncPSCmdlet { private static ManualResetEventSlim _callbackCompleted = new(); private static ManualResetEventSlim _initialized = new(); + private static TestAsyncLateProgressCommand? _instance; private static IProgress? _progress; public static ManualResetEventSlim CallbackCompleted => _callbackCompleted; public static ManualResetEventSlim Initialized => _initialized; + public static Exception? LateWriteException { get; private set; } + public static void Reset() { _callbackCompleted.Dispose(); _initialized.Dispose(); _callbackCompleted = new ManualResetEventSlim(); _initialized = new ManualResetEventSlim(); + _instance = null; _progress = null; + LateWriteException = null; } public static void ReportAfterStop() => _progress!.Report(50); + public static void WriteAfterStop() + { + try + { + _instance!.WriteWarning("after stop"); + } + catch (Exception exception) + { + LateWriteException = exception; + } + } + protected override async Task ProcessRecordAsync() { + _instance = this; _progress = new Progress(percent => { try @@ -514,6 +730,24 @@ protected override Task ProcessRecordAsync() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncStoppedBeforeHook")] +public sealed class TestAsyncStoppedBeforeHookCommand : AsyncPSCmdlet +{ + public bool HookInvoked { get; private set; } + + public void InvokeProcessRecord() + => base.ProcessRecord(); + + public void InvokeStopProcessing() + => base.StopProcessing(); + + protected override Task ProcessRecordAsync() + { + HookInvoked = true; + return Task.CompletedTask; + } +} + public sealed class ChoiceHost( bool approved, Exception? promptFailure = null, diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs index d5cc38ef9..22d461bec 100644 --- a/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs @@ -206,7 +206,7 @@ public void AsyncPSCmdlet_preserves_fifo_order_during_reentrant_queue_drains() } [Fact] - public void AsyncPSCmdlet_rejects_context_free_callbacks_after_their_hook_completes() + public void AsyncPSCmdlet_preserves_context_free_callbacks_while_their_output_is_pumped() { var sessionState = InitialSessionState.CreateDefault(); sessionState.Commands.Add(new SessionStateCmdletEntry( @@ -223,7 +223,10 @@ public void AsyncPSCmdlet_rejects_context_free_callbacks_after_their_hook_comple var result = powerShell.Invoke(); Assert.Equal("value", Assert.Single(result).BaseObject); - Assert.Empty(powerShell.Streams.Warning); + Assert.Equal( + "captured-during-enumeration", + Assert.Single( + powerShell.Streams.Warning).Message); } [Fact] @@ -278,6 +281,97 @@ public void AsyncPSCmdlet_drains_causal_tail_records_before_direct_interactions( Assert.True(TestAsyncDirectBarrierTailCommand.TailEnumeratedBeforeInteraction); } + [Fact] + public void AsyncPSCmdlet_drains_reentrant_records_before_direct_interactions() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncReentrantInteractionTail", + typeof( + TestAsyncReentrantInteractionTailCommand), + helpFileName: null)); + + using var runspace = + RunspaceFactory.CreateRunspace( + sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand( + "Test-AsyncReentrantInteractionTail") + .AddParameter( + "Confirm", + false); + TestAsyncReentrantInteractionTailCommand + .Reset(); + + var result = powerShell.Invoke(); + + Assert.Equal( + ["tail", "outer"], + result.Select( + static item => item.BaseObject)); + Assert.True( + TestAsyncReentrantInteractionTailCommand + .TailEnumeratedBeforeInteraction); + } + + [Fact] + public void AsyncPSCmdlet_clears_a_completed_pipe_before_a_recovery_write() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncFailureRecoveryWrite", + typeof( + TestAsyncFailureRecoveryWriteCommand), + helpFileName: null)); + + using var runspace = + RunspaceFactory.CreateRunspace( + sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand( + "Test-AsyncFailureRecoveryWrite"); + TestAsyncFailureRecoveryWriteCommand + .Reset(); + + var result = powerShell.Invoke(); + + Assert.Empty(result); + Assert.True( + TestAsyncFailureRecoveryWriteCommand + .TransportClearedAfterFailure); + } + + [Fact] + public void AsyncPSCmdlet_uses_a_distinct_context_for_each_hook() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncHookContextIdentity", + typeof( + TestAsyncHookContextIdentityCommand), + helpFileName: null)); + + using var runspace = + RunspaceFactory.CreateRunspace( + sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand( + "Test-AsyncHookContextIdentity"); + + var result = powerShell.Invoke(); + + Assert.False( + (bool)Assert.Single( + result).BaseObject); + } + [Fact] public void AsyncPSCmdlet_allows_ShouldProcess_before_the_async_base_hook_starts() { @@ -356,6 +450,51 @@ public void AsyncPSCmdlet_snapshots_mutable_progress_before_queueing() } } + [Fact] + public void AsyncPSCmdlet_snapshots_captured_errors_before_queueing() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCapturedErrorSnapshot", + typeof(TestAsyncCapturedErrorSnapshotCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCapturedErrorSnapshot"); + + var result = powerShell.Invoke(); + + Assert.Equal("completed", Assert.Single(result).BaseObject); + var captured = Assert.Single(powerShell.Streams.Error); + Assert.Equal("original details", captured.ErrorDetails?.Message); + } + + [Fact] + public void AsyncPSCmdlet_preserves_causal_stream_records_before_enumerator_failure() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncEnumeratorFailure", + typeof(TestAsyncEnumeratorFailureCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncEnumeratorFailure"); + + var exception = Record.Exception(() => powerShell.Invoke()); + + Assert.NotNull(exception); + Assert.Equal( + "before enumeration failure", + Assert.Single(powerShell.Streams.Warning).Message); + } + [Fact] public void AsyncPSCmdlet_preserves_pipeline_failure_when_cancellation_callbacks_throw() { @@ -500,10 +639,12 @@ public void AsyncPSCmdlet_suppresses_pipeline_stop_from_posted_progress_callback Assert.Throws(() => powerShell.EndInvoke(invocation)); TestAsyncLateProgressCommand.ReportAfterStop(); + TestAsyncLateProgressCommand.WriteAfterStop(); Assert.True( TestAsyncLateProgressCommand.CallbackCompleted.Wait(TimeSpan.FromSeconds(5)), "The posted progress callback did not complete in time."); + Assert.Null(TestAsyncLateProgressCommand.LateWriteException); } [Fact] @@ -513,4 +654,14 @@ public void AsyncPSCmdlet_normalizes_synchronous_cancellation_after_stop() Assert.Throws(command.InvokeProcessRecord); } + + [Fact] + public void AsyncPSCmdlet_does_not_invoke_a_hook_after_stop() + { + using var command = new TestAsyncStoppedBeforeHookCommand(); + command.InvokeStopProcessing(); + + Assert.Throws(command.InvokeProcessRecord); + Assert.False(command.HookInvoked); + } } diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs index d39c02c89..300962551 100644 --- a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -475,6 +475,7 @@ public sealed class TestAsyncClaimedShouldContinueCommand : AsyncPSCmdlet public static ManualResetEventSlim ReplyObserved => _replyObserved; public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + public static bool SideEffectStarted { get; private set; } public static void Reset() { @@ -482,14 +483,22 @@ public static void Reset() _cancellationObserved.Dispose(); _replyObserved = new ManualResetEventSlim(); _cancellationObserved = new ManualResetEventSlim(); + SideEffectStarted = false; } protected override async Task ProcessRecordAsync() { await Task.Yield(); _ = CancelToken.Register(_cancellationObserved.Set); - _ = ShouldContinue("Proceed?", "Question"); - _replyObserved.Set(); + try + { + _ = ShouldContinue("Proceed?", "Question"); + SideEffectStarted = true; + } + finally + { + _replyObserved.Set(); + } } } diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 0bb12756e..ca8c1835a 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -527,6 +527,7 @@ public async Task AsyncPSCmdlet_keeps_a_claimed_host_reply_observed_during_cance Assert.True( TestAsyncClaimedShouldContinueCommand.ReplyObserved.Wait(TimeSpan.FromSeconds(5)), "The claimed requester abandoned the host reply during cancellation."); + Assert.False(TestAsyncClaimedShouldContinueCommand.SideEffectStarted); } [Fact] From b4016f79d2f27043bc523addc49f553d1d8a5b96 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 11:28:54 +0200 Subject: [PATCH 27/27] Bind queued lifecycle data to its originating hook --- .../Cmdlets/AsyncPSCmdlet.Execution.cs | 2 +- PSPublishModule/Cmdlets/AsyncPSCmdlet.cs | 21 +++++- .../AsyncPSCmdletLifecycleTestCommands.cs | 74 +++++++++++++++++++ .../AsyncPSCmdletLifecycleTests.cs | 45 +++++++++++ 4 files changed, 140 insertions(+), 2 deletions(-) diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs index 18cdab2ad..9af5cdc8c 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs @@ -20,7 +20,7 @@ private void RunBlockInAsyncCore(Func task) var hookGeneration = Interlocked.Increment(ref _nextHookGeneration); var synchronizationContext = SynchronizationContext.Current; var hookSynchronizationContext = - new AsyncHookSynchronizationContext(); + new AsyncHookSynchronizationContext(this, hookGeneration); void ClearPipes() { diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index 195bd19a8..4ae532773 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -22,11 +22,24 @@ public abstract partial class AsyncPSCmdlet : PSCmdlet, IDisposable { private sealed class AsyncHookSynchronizationContext : SynchronizationContext { + private readonly AsyncPSCmdlet _owner; + private readonly long _hookGeneration; + + internal AsyncHookSynchronizationContext( + AsyncPSCmdlet owner, + long hookGeneration) + { + _owner = owner; + _hookGeneration = hookGeneration; + } + public override void Post(SendOrPostCallback callback, object? state) => ThreadPool.QueueUserWorkItem(_ => { + var priorHookGeneration = _owner._hookGeneration.Value; try { + _owner._hookGeneration.Value = _hookGeneration; callback(state); } catch (PipelineStoppedException) @@ -34,6 +47,10 @@ public override void Post(SendOrPostCallback callback, object? state) // Fire-and-forget callbacks such as Progress can run after StopProcessing. // Await continuations capture their own exceptions into the hook task. } + finally + { + _owner._hookGeneration.Value = priorHookGeneration; + } }); } @@ -635,7 +652,9 @@ protected override void StopProcessing() return; } - if (!TryQueue(new PipelineItem(errorRecord, PipelineType.TerminatingError))) + if (!TryQueue(new PipelineItem( + SnapshotErrorRecord(errorRecord), + PipelineType.TerminatingError))) { ThrowIfStopped(); throw new InvalidOperationException( diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs index d128916cc..e5997fee8 100644 --- a/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs @@ -467,6 +467,80 @@ System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() } } +[Cmdlet(VerbsDiagnostic.Test, "AsyncTerminatingErrorSnapshot")] +public sealed class TestAsyncTerminatingErrorSnapshotCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var error = new ErrorRecord( + new InvalidOperationException("terminating snapshot failure"), + "TerminatingErrorSnapshot", + ErrorCategory.InvalidOperation, + targetObject: null) + { + ErrorDetails = new ErrorDetails("original terminating details") + }; + try + { + ThrowTerminatingError(error); + } + catch (PipelineStoppedException) + { + error.ErrorDetails = new ErrorDetails("mutated terminating details"); + throw; + } + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncStaleProgress")] +public sealed class TestAsyncStaleProgressCommand : AsyncPSCmdlet +{ + private IProgress? _firstRecordProgress; + private int _record; + + [Parameter(ValueFromPipeline = true)] + public int InputObject { get; set; } + + protected override async Task ProcessRecordAsync() + { + _record++; + if (_record == 1) + { + _firstRecordProgress = new Progress( + _ => WriteWarning("stale-first-record-warning")); + return; + } + + using var completed = new ManualResetEventSlim(); + Exception? callbackException = null; + ThreadPool.UnsafeQueueUserWorkItem( + _ => + { + try + { + _firstRecordProgress!.Report(InputObject); + } + catch (Exception exception) + { + callbackException = exception; + } + finally + { + completed.Set(); + } + }, + null); + Assert.True( + completed.Wait(TimeSpan.FromSeconds(5)), + "The context-free progress producer did not complete in time."); + await Task.Delay(100); + if (callbackException is not null) + throw callbackException; + WriteObject(InputObject); + } +} + [Cmdlet(VerbsDiagnostic.Test, "AsyncEnumeratorFailure")] public sealed class TestAsyncEnumeratorFailureCommand : AsyncPSCmdlet { diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs index 22d461bec..1ce6ce308 100644 --- a/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs @@ -472,6 +472,51 @@ public void AsyncPSCmdlet_snapshots_captured_errors_before_queueing() Assert.Equal("original details", captured.ErrorDetails?.Message); } + [Fact] + public void AsyncPSCmdlet_snapshots_terminating_errors_before_queueing() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncTerminatingErrorSnapshot", + typeof(TestAsyncTerminatingErrorSnapshotCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncTerminatingErrorSnapshot"); + + var exception = Assert.Throws( + () => powerShell.Invoke()); + + Assert.Equal( + "original terminating details", + exception.ErrorRecord.ErrorDetails?.Message); + } + + [Fact] + public void AsyncPSCmdlet_rejects_context_callbacks_from_an_earlier_hook() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncStaleProgress", + typeof(TestAsyncStaleProgressCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddScript("1, 2 | Test-AsyncStaleProgress"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal(2, Assert.Single(result).BaseObject); + Assert.Empty(powerShell.Streams.Warning); + } + [Fact] public void AsyncPSCmdlet_preserves_causal_stream_records_before_enumerator_failure() {