From 424454f4943a6e555c6ca66714e461641f38e719 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sat, 25 Jul 2026 23:18:46 +0200 Subject: [PATCH 01/24] Harden asynchronous PowerShell cmdlets --- .../Communication/AsyncPSCmdlet.cs | 466 ++++++++++++------ 1 file changed, 309 insertions(+), 157 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index d52aa9d0..b52c2cbf 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs @@ -7,45 +7,70 @@ namespace DnsClientX.PowerShell; /// -/// An abstract base class for asynchronous PowerShell cmdlets. +/// Base class for cmdlets that await asynchronous engine work while routing PowerShell pipeline writes +/// back through the synchronous cmdlet pipeline thread. /// -public abstract class AsyncPSCmdlet : PSCmdlet, IDisposable { - /// - /// Defines the types of pipelines used in the cmdlet. - /// - private enum PipelineType { +/// +/// Invoke asynchronous hooks on the PowerShell pipeline thread until their first incomplete await. +/// The base temporarily replaces the host synchronization context with an internal thread-pool +/// context while invoking each hook. This prevents continuations from capturing either the host +/// context or a custom task scheduler that may be running the PowerShell pipeline thread. +/// 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 +{ + private sealed class AsyncHookSynchronizationContext : SynchronizationContext + { + public override void Post(SendOrPostCallback callback, object? state) + => ThreadPool.QueueUserWorkItem(_ => callback(state)); + } + + private enum PipelineType + { Output, OutputEnumerate, Error, + TerminatingError, Warning, Verbose, Debug, Information, Progress, ShouldProcess, + ShouldContinue, + PromptForCredential } - /// - /// Cancels the processing of the cmdlet. - /// - private CancellationTokenSource _cancelSource = new(); + private sealed class PipelineItem + { + public PipelineItem(object? value, PipelineType type, BlockingCollection? replyPipe = null) + { + Value = value; + Type = type; + ReplyPipe = replyPipe; + } + + public object? Value { get; } + + public PipelineType Type { get; } + + public BlockingCollection? ReplyPipe { get; } + } - private BlockingCollection<(object?, PipelineType)>? _currentOutPipe; - private BlockingCollection? _currentReplyPipe; + private readonly CancellationTokenSource _cancelSource = new(); + private static readonly SynchronizationContext HookSynchronizationContext = new AsyncHookSynchronizationContext(); + private BlockingCollection? _currentOutPipe; + private int _pipelineThreadId; - /// - /// Gets the cancellation token that is triggered when the cmdlet is stopped. - /// - protected internal CancellationToken CancelToken { get => _cancelSource.Token; } + /// Cancellation token triggered when PowerShell stops the cmdlet. + protected internal CancellationToken CancelToken => _cancelSource.Token; /// protected override void BeginProcessing() => RunBlockInAsync(BeginProcessingAsync); - /// - /// Override this method to implement asynchronous begin processing logic. - /// - /// A task representing the asynchronous operation. + /// Asynchronous begin hook. protected virtual Task BeginProcessingAsync() => Task.CompletedTask; @@ -53,10 +78,7 @@ protected virtual Task BeginProcessingAsync() protected override void ProcessRecord() => RunBlockInAsync(ProcessRecordAsync); - /// - /// Override this method to implement asynchronous record processing logic. - /// - /// A task representing the asynchronous operation. + /// Asynchronous process-record hook. protected virtual Task ProcessRecordAsync() => Task.CompletedTask; @@ -64,178 +86,308 @@ protected virtual Task ProcessRecordAsync() protected override void EndProcessing() => RunBlockInAsync(EndProcessingAsync); - /// - /// Override this method to implement asynchronous end processing logic. - /// - /// A task representing the asynchronous operation. + /// Asynchronous end hook. protected virtual Task EndProcessingAsync() => Task.CompletedTask; /// protected override void StopProcessing() - => _cancelSource?.Cancel(); - - /// - /// Runs the specified task asynchronously and handles the output and reply pipelines. - /// - /// The task to run asynchronously. - private void RunBlockInAsync(Func task) { - using BlockingCollection<(object?, PipelineType)> outPipe = new(); - using BlockingCollection replyPipe = new(); - Task blockTask = Task.Run(async () => { - try { - _currentOutPipe = outPipe; - _currentReplyPipe = replyPipe; - await task(); - } finally { - _currentOutPipe = null; - _currentReplyPipe = null; - outPipe.CompleteAdding(); - replyPipe.CompleteAdding(); - } - }); - - foreach ((object? data, PipelineType pipelineType) in outPipe.GetConsumingEnumerable()) { - switch (pipelineType) { - case PipelineType.Output: - base.WriteObject(data); - break; + => _cancelSource.Cancel(); - case PipelineType.OutputEnumerate: - base.WriteObject(data, true); - break; + /// Thread-safe ShouldProcess bridge for asynchronous cmdlet code. + public new bool ShouldProcess(string? target, string action) + { + ThrowIfStopped(); + if (_currentOutPipe is null || IsPipelineThread) + return base.ShouldProcess(target ?? string.Empty, action); - case PipelineType.Error: - base.WriteError((ErrorRecord)data!); - break; + using var replyPipe = new BlockingCollection(boundedCapacity: 1); + _currentOutPipe.Add(new PipelineItem((target ?? string.Empty, action), PipelineType.ShouldProcess, replyPipe), CancelToken); + return (bool)replyPipe.Take(CancelToken)!; + } - case PipelineType.Warning: - base.WriteWarning((string)data!); - break; + /// 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); - case PipelineType.Verbose: - base.WriteVerbose((string)data!); - break; + using var replyPipe = new BlockingCollection(boundedCapacity: 1); + _currentOutPipe.Add(new PipelineItem((query, caption), PipelineType.ShouldContinue, replyPipe), CancelToken); + return (bool)replyPipe.Take(CancelToken)!; + } - case PipelineType.Debug: - base.WriteDebug((string)data!); - break; + /// 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) + return Host.UI.PromptForCredential(caption, message, userName, targetName); - case PipelineType.Information: - base.WriteInformation((InformationRecord)data!); - break; + using var replyPipe = new BlockingCollection(boundedCapacity: 1); + _currentOutPipe.Add(new PipelineItem((caption, message, userName, targetName), PipelineType.PromptForCredential, replyPipe), CancelToken); + return (PSCredential?)replyPipe.Take(CancelToken); + } - case PipelineType.Progress: - base.WriteProgress((ProgressRecord)data!); - break; + /// Thread-safe output bridge for asynchronous cmdlet code. + public new void WriteObject(object? sendToPipeline) + => WriteObject(sendToPipeline, enumerateCollection: false); - case PipelineType.ShouldProcess: - (string target, string action) = (ValueTuple)data!; - bool res = base.ShouldProcess(target, action); - replyPipe.Add(res); - break; - } + /// Thread-safe output bridge for asynchronous cmdlet code. + public new void WriteObject(object? sendToPipeline, bool enumerateCollection) + { + ThrowIfStopped(); + if (_currentOutPipe is null || IsPipelineThread) + { + base.WriteObject(sendToPipeline, enumerateCollection); + return; } - blockTask.GetAwaiter().GetResult(); + _currentOutPipe.Add(new PipelineItem(sendToPipeline, enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output), CancelToken); } - /// - /// Determines whether the cmdlet should continue processing. - /// - /// The target of the operation. - /// The action to be performed. - /// True if the cmdlet should continue processing; otherwise, false. - public new bool ShouldProcess(string target, string action) { + /// Thread-safe error bridge for asynchronous cmdlet code. + public new void WriteError(ErrorRecord errorRecord) + { ThrowIfStopped(); - _currentOutPipe?.Add(((target, action), PipelineType.ShouldProcess)); - return (bool)_currentReplyPipe?.Take(CancelToken)!; - } - - /// - /// Writes an object to the output pipeline. - /// - /// The object to send to the pipeline. - public new void WriteObject(object? sendToPipeline) => WriteObject(sendToPipeline, false); - - /// - /// Writes an object to the output pipeline, optionally enumerating collections. - /// - /// The object to send to the pipeline. - /// If true, enumerates the collection. - public new void WriteObject(object? sendToPipeline, bool enumerateCollection) { - ThrowIfStopped(); - _currentOutPipe?.Add( - (sendToPipeline, enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output)); + if (_currentOutPipe is null || IsPipelineThread) + { + base.WriteError(errorRecord); + return; + } + + _currentOutPipe.Add(new PipelineItem(errorRecord, PipelineType.Error), CancelToken); } - /// - /// Writes an error record to the error pipeline. - /// - /// The error record to write. - public new void WriteError(ErrorRecord errorRecord) { + /// Thread-safe terminating-error bridge for asynchronous cmdlet code. + protected new void ThrowTerminatingError(ErrorRecord errorRecord) + { ThrowIfStopped(); - _currentOutPipe?.Add((errorRecord, PipelineType.Error)); + if (_currentOutPipe is null || IsPipelineThread) + { + base.ThrowTerminatingError(errorRecord); + return; + } + + _currentOutPipe.Add(new PipelineItem(errorRecord, PipelineType.TerminatingError), CancelToken); + throw new PipelineStoppedException(); } - /// - /// Writes a warning message to the warning pipeline. - /// - /// The warning message to write. - public new void WriteWarning(string message) { + /// Thread-safe warning bridge for asynchronous cmdlet code. + public new void WriteWarning(string text) + { ThrowIfStopped(); - _currentOutPipe?.Add((message, PipelineType.Warning)); + if (_currentOutPipe is null || IsPipelineThread) + { + base.WriteWarning(text); + return; + } + + _currentOutPipe.Add(new PipelineItem(text, PipelineType.Warning), CancelToken); } - /// - /// Writes a verbose message to the verbose pipeline. - /// - /// The verbose message to write. - public new void WriteVerbose(string message) { + /// Thread-safe verbose bridge for asynchronous cmdlet code. + public new void WriteVerbose(string text) + { ThrowIfStopped(); - _currentOutPipe?.Add((message, PipelineType.Verbose)); + if (_currentOutPipe is null || IsPipelineThread) + { + base.WriteVerbose(text); + return; + } + + _currentOutPipe.Add(new PipelineItem(text, PipelineType.Verbose), CancelToken); } - /// - /// Writes a debug message to the debug pipeline. - /// - /// The debug message to write. - public new void WriteDebug(string message) { + /// Thread-safe debug bridge for asynchronous cmdlet code. + public new void WriteDebug(string text) + { ThrowIfStopped(); - _currentOutPipe?.Add((message, PipelineType.Debug)); + if (_currentOutPipe is null || IsPipelineThread) + { + base.WriteDebug(text); + return; + } + + _currentOutPipe.Add(new PipelineItem(text, PipelineType.Debug), CancelToken); } - /// - /// Writes an information record to the information pipeline. - /// - /// The information record to write. - public new void WriteInformation(InformationRecord informationRecord) { + /// Thread-safe information bridge for asynchronous cmdlet code. + public new void WriteInformation(InformationRecord informationRecord) + { ThrowIfStopped(); - _currentOutPipe?.Add((informationRecord, PipelineType.Information)); + if (_currentOutPipe is null || IsPipelineThread) + { + base.WriteInformation(informationRecord); + return; + } + + _currentOutPipe.Add(new PipelineItem(informationRecord, PipelineType.Information), CancelToken); } - /// - /// Writes a progress record to the progress pipeline. - /// - /// The progress record to write. - public new void WriteProgress(ProgressRecord progressRecord) { + /// Thread-safe progress bridge for asynchronous cmdlet code. + public new void WriteProgress(ProgressRecord progressRecord) + { ThrowIfStopped(); - _currentOutPipe?.Add((progressRecord, PipelineType.Progress)); + if (_currentOutPipe is null || IsPipelineThread) + { + base.WriteProgress(progressRecord); + return; + } + + _currentOutPipe.Add(new PipelineItem(progressRecord, PipelineType.Progress), CancelToken); } - /// - /// Throws a if the cmdlet has been stopped. - /// - internal void ThrowIfStopped() { - if (_cancelSource.IsCancellationRequested) { + /// Throws when PowerShell has requested cancellation. + protected internal void ThrowIfStopped() + { + if (_cancelSource.IsCancellationRequested) throw new PipelineStoppedException(); - } } - /// - /// Disposes the resources used by the cmdlet. - /// - public void Dispose() { - _cancelSource?.Dispose(); + /// + public virtual void Dispose() + => _cancelSource.Dispose(); + + private bool IsPipelineThread + => _pipelineThreadId != 0 && Environment.CurrentManagedThreadId == _pipelineThreadId; + + private void RunBlockInAsync(Func task) + { + using var outPipe = new BlockingCollection(); + Task blockTask; + + void ClearPipes() + { + _currentOutPipe = null; + _pipelineThreadId = 0; + CompleteAddingIfNeeded(outPipe); + } + + static void CompleteAddingIfNeeded(BlockingCollection pipe) + { + if (!pipe.IsAddingCompleted) + pipe.CompleteAdding(); + } + + 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.Progress: + base.WriteProgress((ProgressRecord)item.Value!); + break; + case PipelineType.ShouldProcess: + 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( + Host.UI.PromptForCredential(prompt.Caption, prompt.Message, prompt.UserName, prompt.TargetName), + CancelToken); + break; + } + } + + void PumpQueuedItems() + { + while (outPipe.TryTake(out var item)) + PumpItem(item); + } + + _pipelineThreadId = Environment.CurrentManagedThreadId; + _currentOutPipe = outPipe; + + var synchronizationContext = SynchronizationContext.Current; + try + { + SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); + blockTask = task(); + } + catch + { + ClearPipes(); + throw; + } + finally + { + SynchronizationContext.SetSynchronizationContext(synchronizationContext); + } + + if (blockTask.IsCompleted) + { + CompleteAddingIfNeeded(outPipe); + try + { + PumpQueuedItems(); + } + finally + { + ClearPipes(); + } + + blockTask.GetAwaiter().GetResult(); + return; + } + + _ = blockTask.ContinueWith( + completed => ClearPipes(), + CancellationToken.None, + TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + + try + { + foreach (var item in outPipe.GetConsumingEnumerable(CancelToken)) + { + PumpItem(item); + } + } + catch + { + _cancelSource.Cancel(); + CompleteAddingIfNeeded(outPipe); + try + { + blockTask.GetAwaiter().GetResult(); + } + catch (Exception ex) when (ex is OperationCanceledException or PipelineStoppedException) + { + } + + throw; + } + + blockTask.GetAwaiter().GetResult(); } -} \ No newline at end of file +} From 9cd6e92017eb88ba3e617b2da8f832a1ce04c39d 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 02/24] Normalize async cancellation as pipeline stops --- .../Communication/AsyncPSCmdlet.cs | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index b52c2cbf..bb838805 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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); } } From da467aebe0bc36719f1b40d18bbfd0572e5ebf56 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 03/24] Harden async pipeline lifecycle contracts --- .../Communication/AsyncPSCmdlet.cs | 317 ++++++++++++++---- 1 file changed, 259 insertions(+), 58 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index bb838805..e3b5ad8d 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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(); + } } } From 20d4b267337a986d2c65aebdde325bc3217301af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 00:32:54 +0200 Subject: [PATCH 04/24] Close async cmdlet lifecycle races --- .../Communication/AsyncPSCmdlet.cs | 168 +++++++++++++++--- 1 file changed, 140 insertions(+), 28 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index e3b5ad8d..ff7862da 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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; + } } From 92e0a624566a67fafe2680afe61b9da17c5bc97e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 00:46:02 +0200 Subject: [PATCH 05/24] Coordinate synchronous writes and prompt replies --- .../Communication/AsyncPSCmdlet.cs | 147 +++++++++++++----- 1 file changed, 109 insertions(+), 38 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index ff7862da..93a90f3d 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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; } } From fb1068e6728384cc3ddaa5fc26c9226407e0c1a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 01:11:37 +0200 Subject: [PATCH 06/24] Harden lifecycle boundaries and scheduler isolation --- .../Communication/AsyncPSCmdlet.cs | 137 ++++++++++++------ 1 file changed, 92 insertions(+), 45 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index 93a90f3d..853f789a 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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) { From b1795b385663fd7d7760bcf1bfa8f2f544977d3e 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 07/24] Make async stream writes cancellation-safe --- DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index 853f789a..ed2c1586 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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 From 60066fb8feda403dd931fc18a55d3b7abc1e9165 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 01:36:16 +0200 Subject: [PATCH 08/24] Preserve async hook affinity and write ordering --- .../Communication/AsyncPSCmdlet.cs | 49 +++++++++++++------ 1 file changed, 33 insertions(+), 16 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index ed2c1586..f05595df 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs @@ -25,6 +25,17 @@ private sealed class AsyncHookSynchronizationContext : SynchronizationContext 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 { @@ -119,6 +130,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 +283,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 +301,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 +337,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 +353,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 +369,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 +385,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 +401,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 +417,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 +690,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) { From 4990e8b546dab74142c6bb80d91a4e6372dfa53d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 01:45:01 +0200 Subject: [PATCH 09/24] Split async pipeline internals and snapshot tags --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 443 ++++++++++++++++++ .../Communication/AsyncPSCmdlet.cs | 439 +---------------- 2 files changed, 447 insertions(+), 435 deletions(-) create mode 100644 DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs new file mode 100644 index 00000000..7040a1ec --- /dev/null +++ b/DnsClientX.PowerShell/Communication/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 DnsClientX.PowerShell; + +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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index f05595df..ea5b9b8c 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs @@ -18,7 +18,7 @@ namespace DnsClientX.PowerShell; /// 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 { @@ -411,439 +411,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)); } } From efb3bce3b4f8ae924555a42869ad28f0b150b4a8 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 10/24] Align async cmdlet lifecycle semantics --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 100 ++++++++++++-- .../Communication/AsyncPSCmdlet.cs | 125 ++++++++++++++---- 2 files changed, 194 insertions(+), 31 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 7040a1ec..428231b9 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index ea5b9b8c..8c8c56bf 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs @@ -25,6 +25,7 @@ private sealed class AsyncHookSynchronizationContext : SynchronizationContext public override void Post(SendOrPostCallback callback, object? state) => ThreadPool.QueueUserWorkItem(_ => callback(state)); } + private sealed class AsyncHookTaskScheduler : TaskScheduler { protected override System.Collections.Generic.IEnumerable? GetScheduledTasks() @@ -56,7 +57,8 @@ private enum PipelineType ShouldContinue, ShouldContinueAll, ShouldContinueSecurity, - PromptForCredential + PromptForCredential, + PromptForCredentialOptions } private sealed class PipelineReply @@ -71,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); @@ -79,6 +83,9 @@ public void Publish(Func createValue) { try { + if (Volatile.Read(ref _requesterOwner) == 0) + return; + var reply = new PipelineReply(createValue()); try { @@ -91,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() { @@ -113,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; } @@ -125,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; @@ -174,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)!; } @@ -184,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)!; } @@ -194,7 +227,10 @@ protected override void StopProcessing() { ThrowIfStopped(); if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); return base.ShouldProcess(verboseDescription, verboseWarning, caption); + } return (bool)RequestPipelineReply( (verboseDescription, verboseWarning, caption), @@ -210,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), @@ -224,7 +263,10 @@ protected override void StopProcessing() { ThrowIfStopped(); if (CanAccessPipelineDirectly) + { + PrepareDirectPipelineInteraction(); return base.ShouldContinue(query, caption); + } return (bool)RequestPipelineReply((query, caption), PipelineType.ShouldContinue)!; } @@ -234,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), @@ -254,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), @@ -269,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); @@ -283,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; } @@ -301,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; } @@ -320,6 +398,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { + PrepareDirectPipelineAccess(); base.ThrowTerminatingError(errorRecord); return; } @@ -337,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; } @@ -353,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; } @@ -369,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; } @@ -385,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; } @@ -399,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; } @@ -412,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)); } } From 74a4bdf00a075ce6837a8a00ea396d4c87eb1c36 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 11/24] Bind async records to lifecycle generations --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 21 ++++++++++---- .../Communication/AsyncPSCmdlet.cs | 29 ++++++++++++++++--- 2 files changed, 41 insertions(+), 9 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 428231b9..615c478f 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index 8c8c56bf..42ee765f 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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(); From a94ec7c3159e424c5bef03019df96bbf5be77331 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 12/24] Harden async pipeline edge cases --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 30 ++++++++++++++++--- .../Communication/AsyncPSCmdlet.cs | 1 + 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 615c478f..18e2f34d 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index 42ee765f..cc108a0a 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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(); From a533ab1d3e155470e2d82fbd81babeb56a29b0bf 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 13/24] Propagate host interaction failures --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 4 ++- .../Communication/AsyncPSCmdlet.cs | 25 +++++++++++++++---- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 18e2f34d..91c52c66 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index cc108a0a..bc998842 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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; } From f0236f14dd878008356235ce32408d102785e8ed 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 14/24] Harden async pipeline callback ownership --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 35 ++++++++++++++++--- .../Communication/AsyncPSCmdlet.cs | 35 ++++++++++++++----- 2 files changed, 56 insertions(+), 14 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 91c52c66..5806de46 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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,23 @@ private void GetBlockTaskResult(Task blockTask) } } + /// + /// Captures an output writer for callbacks whose producer does not flow the hook execution context. + /// Calls made after the originating hook ends are rejected. + /// + 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 +287,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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index bc998842..bf973902 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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. From d6bf7fb827f85b927b049a8e88d90d796956cc3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 04:38:32 +0200 Subject: [PATCH 15/24] Align async lifecycle thread ownership --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 272 +++++++++--------- 1 file changed, 143 insertions(+), 129 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 5806de46..7799cf9f 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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) @@ -197,6 +201,7 @@ private void RunBlockInAsync(Func task) } finally { + _pipelineThreadId = 0; ExitAsyncBlock(); } } @@ -252,132 +257,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; } } From c8c1a438c6c4444ff905f3d4460e1631019a7100 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 04:58:31 +0200 Subject: [PATCH 16/24] Align async lifecycle edge cases --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 37 +++++++++++++++---- .../Communication/AsyncPSCmdlet.cs | 2 +- 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 7799cf9f..47beeb76 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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; @@ -208,7 +219,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; @@ -527,11 +538,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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index bf973902..fb4ae70f 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs @@ -110,7 +110,7 @@ private void PublishReply(Func createReply) catch (Exception exception) { TryPublish(new PipelineReply(value: null, exception)); - throw; + return; } TryPublish(reply); From ab1b56205b374b0936103cfcf247973be02cd156 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 05:43:00 +0200 Subject: [PATCH 17/24] Sync final AsyncPSCmdlet lifecycle hardening --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 70 +++++++++-- .../Communication/AsyncPSCmdlet.cs | 109 +++++++++++++++--- 2 files changed, 150 insertions(+), 29 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 47beeb76..0d01031c 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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,27 +77,40 @@ 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() { if (Volatile.Read(ref _asyncLifecycleStarted) == 0) @@ -177,6 +195,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); @@ -199,6 +232,9 @@ private bool TryQueue(PipelineItem item) } catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) { + if (item.HookGeneration != 0) + throw new PipelineStoppedException(); + return false; } } @@ -219,7 +255,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; @@ -432,6 +471,7 @@ void PumpQueuedItems() var priorHookGeneration = _hookGeneration.Value; try { + Volatile.Write(ref _pipelineSynchronizationContext, synchronizationContext); SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); _hookGeneration.Value = hookGeneration; if (TaskScheduler.Current == TaskScheduler.Default) @@ -474,6 +514,7 @@ void PumpQueuedItems() { _hookGeneration.Value = priorHookGeneration; SynchronizationContext.SetSynchronizationContext(synchronizationContext); + Volatile.Write(ref _pipelineSynchronizationContext, null); } if (blockTask.IsCompleted) @@ -481,7 +522,6 @@ void PumpQueuedItems() if (blockTask.IsFaulted) _ = blockTask.Exception; - CompleteAddingIfNeeded(outPipe); try { PumpQueuedItems(); @@ -508,9 +548,11 @@ void PumpQueuedItems() if (completed.IsFaulted) _ = completed.Exception; - ClearPipes(); if (Volatile.Read(ref deferPipeDisposal) != 0) + { + ClearPipes(); DisposePipeOnce(); + } } finally { @@ -529,10 +571,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) { @@ -605,6 +650,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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index fb4ae70f..b91509c8 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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,44 @@ 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 +247,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 +278,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 +303,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldProcess(target ?? string.Empty); } @@ -253,7 +316,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldProcess(target ?? string.Empty, action); } @@ -266,7 +329,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldProcess(verboseDescription, verboseWarning, caption); } @@ -285,7 +348,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldProcess(verboseDescription, verboseWarning, caption, out shouldProcessReason); } @@ -302,7 +365,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldContinue(query, caption); } @@ -315,7 +378,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return base.ShouldContinue(query, caption, ref yesToAll, ref noToAll); } @@ -338,7 +401,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 +419,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return Host.UI.PromptForCredential(caption, message, userName, targetName); } @@ -377,7 +440,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineInteraction(); + using var pipelineContext = EnterDirectPipelineInteraction(); return Host.UI.PromptForCredential( caption, message, @@ -399,9 +462,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 +481,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 +501,7 @@ protected override void StopProcessing() ThrowIfStopped(); if (CanAccessPipelineDirectly) { - PrepareDirectPipelineAccess(); + using var pipelineContext = EnterDirectPipelineAccess(); base.ThrowTerminatingError(errorRecord); return; } @@ -454,9 +519,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 +536,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 +553,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 +570,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 +587,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 +604,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; } From d014e69db3ac3fecdcf3762d522490824480f1bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 06:25:32 +0200 Subject: [PATCH 18/24] Sync AsyncPSCmdlet completion hardening --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 111 ++++++++++++++++-- .../Communication/AsyncPSCmdlet.cs | 75 +++++++++--- 2 files changed, 160 insertions(+), 26 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 0d01031c..783260ad 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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,13 +131,23 @@ 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); } + private void ValidateInteractionGeneration() { if (Volatile.Read(ref _asyncLifecycleStarted) == 0) @@ -180,8 +217,11 @@ private void GetBlockTaskResult(Task blockTask) /// /// Captures an output writer for callbacks whose producer does not flow the hook execution context. - /// Calls made after the originating hook ends are rejected. /// + /// + /// 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; @@ -192,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)); } /// @@ -213,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; @@ -232,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; @@ -264,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); } @@ -308,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: @@ -437,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; } } @@ -464,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 { @@ -514,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; @@ -545,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(); @@ -573,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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index b91509c8..74c72e45 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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,46 +206,87 @@ 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. /// - protected sealed class CapturedPipelineStreams { + protected sealed class CapturedPipelineStreams + { private readonly long _hookGeneration; private readonly AsyncPSCmdlet _owner; - internal CapturedPipelineStreams(AsyncPSCmdlet owner, long hookGeneration) { + 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); + => Queue( + value, + enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output); /// Queues an error record for the originating hook. - public void WriteError(ErrorRecord errorRecord) => Queue(errorRecord, PipelineType.Error); + 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); + 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); + 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); + 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); + 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); + => 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); + 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); + 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(); @@ -249,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; From b6c2620bfb5266ab0cbe2005e855450443b834cb 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 19/24] Sync final async lifecycle hardening --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 210 ++++++++++-------- .../Communication/AsyncPSCmdlet.cs | 14 +- 2 files changed, 119 insertions(+), 105 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 783260ad..03e10db5 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index 74c72e45..f75fee03 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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() From 854f83b8a4486c3f0155d6219c8520ae10f940c1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 07:11:42 +0200 Subject: [PATCH 20/24] Sync final async lifecycle race fixes --- .../Communication/AsyncPSCmdlet.Pipeline.cs | 15 +++++++++++++-- .../Communication/AsyncPSCmdlet.cs | 13 ++++++++++++- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 03e10db5..521fd2f1 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index f75fee03..5a652591 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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 From e6ae3207d8f6b67d633ff523ca0d9ab4bcd3f675 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 21/24] Harden async PowerShell pipeline lifecycle --- .../Communication/AsyncPSCmdlet.Execution.cs | 461 ++++++++++++++++++ .../Communication/AsyncPSCmdlet.Pipeline.cs | 452 ++--------------- .../Communication/AsyncPSCmdlet.cs | 98 +++- 3 files changed, 578 insertions(+), 433 deletions(-) create mode 100644 DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs new file mode 100644 index 00000000..49b51991 --- /dev/null +++ b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs @@ -0,0 +1,461 @@ +using System; +using System.Collections.Concurrent; +using System.Management.Automation; +using System.Threading; +using System.Threading.Tasks; + +namespace DnsClientX.PowerShell; + +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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 521fd2f1..85443fe1 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index 5a652591..e05d135e 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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); } } From 180c2a1ec0bd57bd193a1bbae5318682d922ec92 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 22/24] Synchronize async pipeline cancellation --- .../Communication/AsyncPSCmdlet.Execution.cs | 17 ++++++----- .../Communication/AsyncPSCmdlet.cs | 30 ++++++++++++++++--- 2 files changed, 36 insertions(+), 11 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs index 49b51991..cda13ea2 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs +++ b/DnsClientX.PowerShell/Communication/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(); diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index e05d135e..56253594 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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(); } From 586b19364b889c43f03a552d8edb741ccd5a7c3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Sun, 26 Jul 2026 09:26:43 +0200 Subject: [PATCH 23/24] Harden async lifecycle tail handling --- .../Communication/AsyncPSCmdlet.Execution.cs | 34 ++++++++++++++++++- .../Communication/AsyncPSCmdlet.Pipeline.cs | 10 ++++-- .../Communication/AsyncPSCmdlet.cs | 16 ++++----- 3 files changed, 46 insertions(+), 14 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs index cda13ea2..8939c95e 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index 85443fe1..b02fcc0b 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index 56253594..19736805 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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() From 360df4cabfead730889bf82a7ceeda1ec80b1616 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 24/24] Sync final async pipeline hardening --- .../Communication/AsyncPSCmdlet.Execution.cs | 87 ++++++++++-- .../Communication/AsyncPSCmdlet.Pipeline.cs | 128 ++++++++++++------ .../Communication/AsyncPSCmdlet.cs | 67 ++++++++- 3 files changed, 227 insertions(+), 55 deletions(-) diff --git a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs index 8939c95e..b5d483ae 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Execution.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs index b02fcc0b..90d74957 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.Pipeline.cs +++ b/DnsClientX.PowerShell/Communication/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/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs b/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs index 19736805..f61767eb 100644 --- a/DnsClientX.PowerShell/Communication/AsyncPSCmdlet.cs +++ b/DnsClientX.PowerShell/Communication/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()),