diff --git a/CHANGELOG.MD b/CHANGELOG.MD index e4bcfdc9f..c46c5b4d1 100644 --- a/CHANGELOG.MD +++ b/CHANGELOG.MD @@ -2,6 +2,7 @@ ## Unreleased ### What's Changed +- Added pipeline-thread marshalling for terminating errors and `ShouldContinue` prompts raised after an asynchronous cmdlet await, while preserving derived cmdlet resource cleanup through virtual disposal. - Fixed module merging so `using` lines inside PowerShell here-strings and block comments remain in place instead of being moved into the generated module header. - Hardened the canonical `AsyncPSCmdlet` lifecycle with request-specific prompt and `ShouldProcess` replies while preserving PowerShell-thread execution until the first await. - Converted Task-backed App Store Connect and managed catalog cmdlets from sync-over-async blocking to cancellable asynchronous lifecycle hooks. diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs new file mode 100644 index 000000000..9af5cdc8c --- /dev/null +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Execution.cs @@ -0,0 +1,563 @@ +using System; +using System.Collections.Concurrent; +using System.Management.Automation; +using System.Threading; +using System.Threading.Tasks; + +namespace PSPublishModule; + +public abstract partial class AsyncPSCmdlet +{ + private void RunBlockInAsyncCore(Func task) + { + // The transport must remain lossless and non-blocking. The pipeline thread can enumerate + // user objects or invoke a host that waits for the same background producer that is writing + // here; applying bounded backpressure would deadlock both sides. + var outPipe = new BlockingCollection(); + Task blockTask; + var deferPipeDisposal = 0; + var pipeDisposed = 0; + var hookGeneration = Interlocked.Increment(ref _nextHookGeneration); + var synchronizationContext = SynchronizationContext.Current; + var hookSynchronizationContext = + new AsyncHookSynchronizationContext(this, hookGeneration); + + 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 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: + 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.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) + return; + + // 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)) + PumpItemPreservingCausalFailureRecords(item); + } + + void PumpThroughDirectAccessBarrier() + { + PipelineItem? completionMarker = null; + var barrier = new PipelineItem( + value: null, + PipelineType.DirectAccessBarrier, + hookGeneration: hookGeneration, + dropOnStop: true); + if (!TryQueue(barrier)) + { + ThrowIfStopped(); + throw new InvalidOperationException( + "No active PowerShell pipeline is available for direct access."); + } + + while (outPipe.TryTake(out var item)) + { + if (item.Type == PipelineType.HookCompleted) + { + completionMarker = item; + continue; + } + + PumpItemPreservingCausalFailureRecords(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."); + } + + if (pumpBoundPredecessor.Type == + PipelineType.HookCompleted) + { + completionMarker = + pumpBoundPredecessor; + continue; + } + + PumpItemPreservingCausalFailureRecords(pumpBoundPredecessor); + } + + if (completionMarker is not null) + { + outPipe.Add( + completionMarker); + } + + return; + } + } + } + + bool HasPumpBoundItems() + { + foreach (var queuedItem in outPipe.ToArray()) + { + if (queuedItem.IsPumpBound) + return true; + } + + return false; + } + + 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; + ThrowIfStopped(); + 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 PipelineStoppedException) + { + CancelSource(); + throw; + } + + 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); + PumpItemPreservingCausalFailureRecords(item); + if (item.Type == PipelineType.HookCompleted) + { + while (outPipe.TryTake(out var pumpBoundItem)) + { + ThrowIfStopped(); + PumpItemPreservingCausalFailureRecords(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) + { + ClearPipes(); + DisposePipeOnce(); + } + DeactivateHook(); + } + + if (pipelineException is OperationCanceledException && stopRequested) + throw new PipelineStoppedException(); + + throw; + } + + try + { + GetBlockTaskResult(blockTask); + } + finally + { + DeactivateHook(); + DisposePipeOnce(); + } + } +} diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs new file mode 100644 index 000000000..9cd5c6bd4 --- /dev/null +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.Pipeline.cs @@ -0,0 +1,462 @@ +using System; +using System.Collections.Concurrent; +using System.Management.Automation; +using System.Reflection; +using System.Runtime.ExceptionServices; +using System.Threading; +using System.Threading.Tasks; + +namespace PSPublishModule; + +public abstract partial class AsyncPSCmdlet +{ + /// Thread-safe progress bridge for asynchronous cmdlet code. + public new void WriteProgress(ProgressRecord progressRecord) + { + if (ShouldDropClosedCanceledStreamWrite()) + return; + + ThrowIfStopped(); + var item = new PipelineItem(SnapshotProgressRecord(progressRecord), PipelineType.Progress); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + + if (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineAccess(); + base.WriteProgress(progressRecord); + return; + } + + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + _ = TryQueue(item); + } + + private static readonly PropertyInfo? ProgressTotalProperty = + typeof(ProgressRecord).GetProperty("Total", BindingFlags.Instance | BindingFlags.Public); + + private static ProgressRecord SnapshotProgressRecord(ProgressRecord progressRecord) + { + var snapshot = new ProgressRecord( + progressRecord.ActivityId, + progressRecord.Activity, + progressRecord.StatusDescription) + { + CurrentOperation = progressRecord.CurrentOperation, + ParentActivityId = progressRecord.ParentActivityId, + PercentComplete = progressRecord.PercentComplete, + RecordType = progressRecord.RecordType, + SecondsRemaining = progressRecord.SecondsRemaining + }; + + if (ProgressTotalProperty is { CanRead: true, CanWrite: true }) + { + ProgressTotalProperty.SetValue( + snapshot, + ProgressTotalProperty.GetValue(progressRecord)); + } + + 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; + } + + private static ErrorRecord SnapshotErrorRecord(ErrorRecord errorRecord) + => new(errorRecord, errorRecord.Exception); + + /// Throws when PowerShell has requested cancellation. + protected internal void ThrowIfStopped() + { + if (_cancelSource.IsCancellationRequested) + 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() + { + bool cancelActiveBlocks; + lock (_lifecycleLock) + { + if (_disposeRequested) + return; + + _disposeRequested = true; + cancelActiveBlocks = _activeBlocks != 0; + Volatile.Write(ref _asyncLifecycleCompleted, 1); + } + + try + { + if (cancelActiveBlocks) + CancelSource(); + } + finally + { + lock (_lifecycleLock) + { + DisposeCancelSourceIfInactive(); + } + + Volatile.Write(ref _pipelineThreadId, 0); + } + } + + private bool IsPipelineThread + { + get + { + var pipelineThreadId = Volatile.Read(ref _pipelineThreadId); + return pipelineThreadId != 0 && + Environment.CurrentManagedThreadId == pipelineThreadId; + } + } + + private bool IsPumpingPipelineItem + => IsPipelineThread && Volatile.Read(ref _pipelinePumpDepth) != 0; + + 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 IDisposable EnterDirectPipelineAccess() + { + 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); + } + + private IDisposable EnterDirectPipelineInteraction() + => EnterDirectPipelineAccess(); + + 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 || IsConstructionThreadOutsideAsyncHook)) + 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 + { + blockTask.GetAwaiter().GetResult(); + } + catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) + { + throw new PipelineStoppedException(); + } + catch (PipelineStoppedException) + { + CancelSource(); + throw; + } + } + + private object? RequestPipelineReply(object? value, PipelineType type) + { + ThrowIfStopped(); + ValidateInteractionGeneration(); + var hookGeneration = _hookGeneration.Value; + var replyPipe = new PipelineReplyChannel(); + if (!TryQueue(new PipelineItem(value, type, replyPipe, hookGeneration))) + { + replyPipe.Abandon(); + ThrowIfStopped(); + throw new InvalidOperationException("No active PowerShell pipeline is available for the asynchronous request."); + } + + try + { + PipelineReply reply; + try + { + reply = replyPipe.Take(CancelToken); + } + catch (OperationCanceledException) when (_cancelSource.IsCancellationRequested) + { + throw new PipelineStoppedException(); + } + + ThrowIfStopped(); + if (reply.Rejection is not null) + ExceptionDispatchInfo.Capture(reply.Rejection).Throw(); + + return reply.Value; + } + finally + { + replyPipe.ReleaseRequester(); + } + } + + /// + /// Captures an output writer for callbacks whose producer does not flow the hook execution context. + /// + /// + /// Capture the writer inside an asynchronous PowerShell hook. Calls made after that hook ends are + /// rejected rather than being rebound to a later record lifecycle. + /// + protected Action CapturePipelineWriter(bool enumerateCollection = false) + { + var hookGeneration = _hookGeneration.Value; + if (hookGeneration == 0) + { + throw new InvalidOperationException( + "A lifecycle-bound pipeline writer can only be captured from an asynchronous PowerShell hook."); + } + + var pipelineType = enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output; + return value => _ = TryQueue( + new PipelineItem( + value, + pipelineType, + hookGeneration: hookGeneration, + dropOnStop: true)); + } + + /// + /// 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); + var pumpLease = _pipelinePumpLease.Value; + if (pumpLease is null) + { + var sharedPumpLease = + Volatile.Read(ref _currentPipelinePumpLease); + if (sharedPumpLease is not null && + (item.HookGeneration == 0 || + item.HookGeneration == sharedPumpLease.Generation)) + { + item.BindToHook(sharedPumpLease.Generation); + pumpLease = sharedPumpLease; + } + } + + var isPumpBound = pumpLease?.TryClaim(item.HookGeneration) == true; + try + { + lock (_hookAdmissionLock) + { + 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) + { + EnterAsyncBlock(); + try + { + RunBlockInAsyncCore(task); + } + finally + { + Volatile.Write(ref _pipelineThreadId, 0); + ExitAsyncBlock(); + } + } + + + 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() + { + lock (_lifecycleLock) + { + if (_cancelSourceDisposed) + return; + + _cancelSourceCancellationInProgress++; + } + + try + { + _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. + } + finally + { + lock (_lifecycleLock) + { + _cancelSourceCancellationInProgress--; + DisposeCancelSourceIfInactive(); + } + } + } + + private void DisposeCancelSourceIfInactive() + { + if (!_disposeRequested || + _activeBlocks != 0 || + _cancelSourceCancellationInProgress != 0 || + _cancelSourceDisposed) + return; + + _cancelSource.Dispose(); + _cancelSourceDisposed = true; + } +} diff --git a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs index e983aa210..4ae532773 100644 --- a/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs +++ b/PSPublishModule/Cmdlets/AsyncPSCmdlet.cs @@ -18,12 +18,66 @@ namespace PSPublishModule; /// Keep hook implementations asynchronous all the way through and pass to /// cancellable engine operations. Do not block with Task.Wait, Task.Result, or Task.WaitAll. /// -public abstract class AsyncPSCmdlet : PSCmdlet, IDisposable +public abstract partial class AsyncPSCmdlet : PSCmdlet, IDisposable { private sealed class AsyncHookSynchronizationContext : SynchronizationContext { + private readonly AsyncPSCmdlet _owner; + private readonly long _hookGeneration; + + internal AsyncHookSynchronizationContext( + AsyncPSCmdlet owner, + long hookGeneration) + { + _owner = owner; + _hookGeneration = hookGeneration; + } + public override void Post(SendOrPostCallback callback, object? state) - => ThreadPool.QueueUserWorkItem(_ => callback(state)); + => ThreadPool.QueueUserWorkItem(_ => + { + var priorHookGeneration = _owner._hookGeneration.Value; + try + { + _owner._hookGeneration.Value = _hookGeneration; + 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. + } + finally + { + _owner._hookGeneration.Value = priorHookGeneration; + } + }); + } + + 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 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 @@ -31,34 +85,311 @@ private enum PipelineType Output, OutputEnumerate, Error, + TerminatingError, Warning, Verbose, Debug, Information, + InformationWithTags, Progress, + CommandDetail, + ShouldProcessTarget, ShouldProcess, - PromptForCredential + ShouldProcessVerbose, + ShouldProcessReason, + ShouldContinue, + ShouldContinueAll, + ShouldContinueSecurity, + PromptForCredential, + PromptForCredentialOptions, + DirectAccessBarrier, + HookCompleted + } + + private sealed class PipelineReply + { + public PipelineReply(object? value, Exception? rejection = null) + { + Value = value; + Rejection = rejection; + } + + public object? Value { get; } + + public Exception? Rejection { get; } + } + + private sealed class PipelineReplyChannel + { + private readonly BlockingCollection _pipe = new(boundedCapacity: 1); + private int _owners = 2; + private int _pipelineOwner = 1; + private int _requesterState = 1; + + public PipelineReply Take(CancellationToken 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())); + + 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 (Interlocked.CompareExchange(ref _requesterState, 2, 1) != 1) + return; + + PipelineReply reply; + try + { + reply = createReply(); + } + catch (Exception exception) + { + TryPublish(new PipelineReply(value: null, exception)); + return; + } + + TryPublish(reply); + } + finally + { + ReleasePipeline(); + } + } + + private void TryPublish(PipelineReply reply) + { + try + { + _pipe.Add(reply); + } + catch (InvalidOperationException) + { + // The requester and pipeline can finish concurrently during cancellation. + } + } + + public void Abandon() + { + ReleaseRequester(); + ReleasePipeline(); + } + + public void ReleaseRequester() + { + if (Interlocked.Exchange(ref _requesterState, 0) != 0) + Release(); + } + + public void ReleasePipeline() + { + if (Interlocked.Exchange(ref _pipelineOwner, 0) == 1) + 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, + long hookGeneration = 0, + bool dropOnStop = false) { Value = value; Type = type; ReplyPipe = replyPipe; + HookGeneration = hookGeneration; + DropOnStop = dropOnStop; } public object? Value { get; } public PipelineType Type { get; } - public BlockingCollection? ReplyPipe { get; } + public PipelineReplyChannel? ReplyPipe { get; } + + public long HookGeneration { get; private set; } + + 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 + { + private readonly object _sync = new(); + private bool _active = true; + private int _claims; + + public PipelinePumpLease(long generation) + => Generation = generation; + + public long Generation { get; } + + public bool TryClaim(long generation) + { + lock (_sync) + { + if (!_active || generation != Generation) + return false; + + _claims++; + return true; + } + } + + 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); + } + } + } + + /// + /// 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(SnapshotErrorRecord(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(SnapshotInformationRecord(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, + dropOnStop: true)); } private readonly CancellationTokenSource _cancelSource = new(); - private static readonly SynchronizationContext HookSynchronizationContext = new AsyncHookSynchronizationContext(); + 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 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; + 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. @@ -90,30 +421,164 @@ 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) + { + ThrowIfStopped(); + if (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineInteraction(); + 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 (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineInteraction(); 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) + { + ThrowIfStopped(); + if (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineInteraction(); + 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) + { + ThrowIfStopped(); + if (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineInteraction(); + 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 (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineInteraction(); + return base.ShouldContinue(query, caption); + } + + 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) + { + ThrowIfStopped(); + if (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineInteraction(); + 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) + { + ThrowIfStopped(); + if (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineInteraction(); + 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 (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineInteraction(); 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) + { + using var pipelineContext = EnterDirectPipelineInteraction(); + return Host.UI.PromptForCredential( + caption, + message, + userName, + targetName, + allowedCredentialTypes, + options); + } - 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, allowedCredentialTypes, options), + PipelineType.PromptForCredentialOptions); } /// Thread-safe output bridge for asynchronous cmdlet code. @@ -123,236 +588,245 @@ protected override void StopProcessing() /// Thread-safe output bridge for asynchronous cmdlet code. public new void WriteObject(object? sendToPipeline, bool enumerateCollection) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + var item = new PipelineItem( + sendToPipeline, + enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output); + if (IsPumpingPipelineItem) + { + _ = TryQueue(item); + return; + } + + if (CanAccessPipelineDirectly) { + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteObject(sendToPipeline, enumerateCollection); return; } - _currentOutPipe.Add(new PipelineItem(sendToPipeline, enumerateCollection ? PipelineType.OutputEnumerate : PipelineType.Output), CancelToken); + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + _ = TryQueue(item); } /// Thread-safe error bridge for asynchronous cmdlet code. public new void WriteError(ErrorRecord errorRecord) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + var item = new PipelineItem(SnapshotErrorRecord(errorRecord), PipelineType.Error); + if (IsPumpingPipelineItem) { + _ = TryQueue(item); + return; + } + + if (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineAccess(); base.WriteError(errorRecord); return; } - _currentOutPipe.Add(new PipelineItem(errorRecord, PipelineType.Error), CancelToken); + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + _ = TryQueue(item); } - /// Thread-safe warning bridge for asynchronous cmdlet code. - public new void WriteWarning(string text) + /// Thread-safe terminating-error bridge for asynchronous cmdlet code. + public new void ThrowTerminatingError(ErrorRecord errorRecord) { ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (CanAccessPipelineDirectly) { - base.WriteWarning(text); + using var pipelineContext = EnterDirectPipelineAccess(); + base.ThrowTerminatingError(errorRecord); return; } - _currentOutPipe.Add(new PipelineItem(text, PipelineType.Warning), CancelToken); + if (!TryQueue(new PipelineItem( + SnapshotErrorRecord(errorRecord), + PipelineType.TerminatingError))) + { + ThrowIfStopped(); + throw new InvalidOperationException( + "No active PowerShell pipeline is available for the terminating error."); + } + + throw new PipelineStoppedException(); } - /// Thread-safe verbose bridge for asynchronous cmdlet code. - public new void WriteVerbose(string text) + /// Thread-safe warning bridge for asynchronous cmdlet code. + public new void WriteWarning(string message) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + var item = new PipelineItem(message, PipelineType.Warning); + if (IsPumpingPipelineItem) { - base.WriteVerbose(text); + _ = TryQueue(item); return; } - _currentOutPipe.Add(new PipelineItem(text, PipelineType.Verbose), CancelToken); - } - - /// Thread-safe debug bridge for asynchronous cmdlet code. - public new void WriteDebug(string text) - { - ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (CanAccessPipelineDirectly) { - base.WriteDebug(text); + using var pipelineContext = EnterDirectPipelineAccess(); + base.WriteWarning(message); return; } - _currentOutPipe.Add(new PipelineItem(text, PipelineType.Debug), CancelToken); + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + _ = TryQueue(item); } - /// Thread-safe information bridge for asynchronous cmdlet code. - public new void WriteInformation(InformationRecord informationRecord) + /// Thread-safe verbose bridge for asynchronous cmdlet code. + public new void WriteVerbose(string message) { + if (ShouldDropClosedCanceledStreamWrite()) + return; + ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + var item = new PipelineItem(message, PipelineType.Verbose); + if (IsPumpingPipelineItem) { - base.WriteInformation(informationRecord); + _ = TryQueue(item); return; } - _currentOutPipe.Add(new PipelineItem(informationRecord, PipelineType.Information), CancelToken); - } - - /// Thread-safe progress bridge for asynchronous cmdlet code. - public new void WriteProgress(ProgressRecord progressRecord) - { - ThrowIfStopped(); - if (_currentOutPipe is null || IsPipelineThread) + if (CanAccessPipelineDirectly) { - base.WriteProgress(progressRecord); + using var pipelineContext = EnterDirectPipelineAccess(); + base.WriteVerbose(message); return; } - _currentOutPipe.Add(new PipelineItem(progressRecord, PipelineType.Progress), CancelToken); - } + if (Volatile.Read(ref _currentOutPipe) is null) + return; - /// Throws when PowerShell has requested cancellation. - internal void ThrowIfStopped() - { - if (_cancelSource.IsCancellationRequested) - throw new PipelineStoppedException(); + _ = TryQueue(item); } - /// - public void Dispose() - => _cancelSource.Dispose(); - - private bool IsPipelineThread - => _pipelineThreadId != 0 && Environment.CurrentManagedThreadId == _pipelineThreadId; - - private void RunBlockInAsync(Func task) + /// Thread-safe debug bridge for asynchronous cmdlet code. + public new void WriteDebug(string message) { - using var outPipe = new BlockingCollection(); - Task blockTask; + if (ShouldDropClosedCanceledStreamWrite()) + return; - void ClearPipes() + ThrowIfStopped(); + var item = new PipelineItem(message, PipelineType.Debug); + if (IsPumpingPipelineItem) { - _currentOutPipe = null; - _pipelineThreadId = 0; - CompleteAddingIfNeeded(outPipe); + _ = TryQueue(item); + return; } - static void CompleteAddingIfNeeded(BlockingCollection pipe) + if (CanAccessPipelineDirectly) { - if (!pipe.IsAddingCompleted) - pipe.CompleteAdding(); + using var pipelineContext = EnterDirectPipelineAccess(); + base.WriteDebug(message); + return; } - 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.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.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; - } - } + if (Volatile.Read(ref _currentOutPipe) is null) + return; - void PumpQueuedItems() - { - while (outPipe.TryTake(out var item)) - PumpItem(item); - } + _ = TryQueue(item); + } - _pipelineThreadId = Environment.CurrentManagedThreadId; - _currentOutPipe = outPipe; + /// Thread-safe command-detail bridge for asynchronous cmdlet code. + public new void WriteCommandDetail(string text) + { + if (ShouldDropClosedCanceledStreamWrite()) + return; - var synchronizationContext = SynchronizationContext.Current; - try + ThrowIfStopped(); + var item = new PipelineItem(text, PipelineType.CommandDetail); + if (IsPumpingPipelineItem) { - SynchronizationContext.SetSynchronizationContext(HookSynchronizationContext); - blockTask = task(); + _ = TryQueue(item); + return; } - catch + + if (CanAccessPipelineDirectly) { - ClearPipes(); - throw; + using var pipelineContext = EnterDirectPipelineAccess(); + base.WriteCommandDetail(text); + return; } - finally + + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + _ = TryQueue(item); + } + + /// Thread-safe information bridge for asynchronous cmdlet code. + public new void WriteInformation(InformationRecord informationRecord) + { + if (ShouldDropClosedCanceledStreamWrite()) + return; + + ThrowIfStopped(); + var item = new PipelineItem( + SnapshotInformationRecord(informationRecord), + PipelineType.Information); + if (IsPumpingPipelineItem) { - SynchronizationContext.SetSynchronizationContext(synchronizationContext); + _ = TryQueue(item); + return; } - if (blockTask.IsCompleted) + if (CanAccessPipelineDirectly) { - CompleteAddingIfNeeded(outPipe); - try - { - PumpQueuedItems(); - } - finally - { - ClearPipes(); - } - - blockTask.GetAwaiter().GetResult(); + using var pipelineContext = EnterDirectPipelineAccess(); + base.WriteInformation(informationRecord); return; } - _ = blockTask.ContinueWith( - completed => ClearPipes(), - CancellationToken.None, - TaskContinuationOptions.ExecuteSynchronously, - TaskScheduler.Default); + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + _ = TryQueue(item); + } - try + /// 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()), + PipelineType.InformationWithTags); + if (IsPumpingPipelineItem) { - foreach (var item in outPipe.GetConsumingEnumerable(CancelToken)) - { - PumpItem(item); - } + _ = TryQueue(item); + return; } - catch - { - _cancelSource.Cancel(); - CompleteAddingIfNeeded(outPipe); - try - { - blockTask.GetAwaiter().GetResult(); - } - catch (Exception ex) when (ex is OperationCanceledException or PipelineStoppedException) - { - } - throw; + if (CanAccessPipelineDirectly) + { + using var pipelineContext = EnterDirectPipelineAccess(); + base.WriteInformation(messageData, tags ?? Array.Empty()); + return; } - blockTask.GetAwaiter().GetResult(); + if (Volatile.Read(ref _currentOutPipe) is null) + return; + + _ = TryQueue(item); } } diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs new file mode 100644 index 000000000..e5997fee8 --- /dev/null +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTestCommands.cs @@ -0,0 +1,958 @@ +using System.Collections.ObjectModel; +using System.Globalization; +using System.Management.Automation; +using System.Management.Automation.Host; +using System.Management.Automation.Runspaces; +using System.Reflection; +using System.Security; +using System.Threading; +using System.Threading.Tasks; +using PSPublishModule; + +namespace PowerForge.Tests; + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedCallback")] +public sealed class TestAsyncCapturedCallbackCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + var writeOutput = CapturePipelineWriter(); + var streams = CapturePipelineStreams(); + var completed = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + ThreadPool.UnsafeQueueUserWorkItem( + _ => + { + writeOutput("callback-output"); + streams.WriteWarning("callback-warning"); + completed.TrySetResult(true); + }, + null); + await completed.Task.ConfigureAwait(false); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCommandDetail")] +public sealed class TestAsyncCommandDetailCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteCommandDetail("worker-detail"); + WriteObject("completed"); + } +} + +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncReentrantPump", + SupportsShouldProcess = true)] +public sealed class TestAsyncReentrantPumpCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(new WarningEnumerable(this), enumerateCollection: true); + WriteWarning("after-enumeration"); + } + + private sealed class WarningEnumerable : IEnumerable + { + private readonly TestAsyncReentrantPumpCommand _command; + + public WarningEnumerable(TestAsyncReentrantPumpCommand command) + { + _command = command; + } + + public IEnumerator GetEnumerator() + { + Assert.True( + _command.ShouldProcess( + "enumerated-target")); + Task.Run(() => _command.WriteWarning("during-enumeration")).GetAwaiter().GetResult(); + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncReentrantFifo")] +public sealed class TestAsyncReentrantFifoCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(new WarningEnumerable(this), enumerateCollection: true); + } + + private sealed class WarningEnumerable : IEnumerable + { + private readonly TestAsyncReentrantFifoCommand _command; + + public WarningEnumerable(TestAsyncReentrantFifoCommand command) + { + _command = command; + } + + public IEnumerator GetEnumerator() + { + Task.Run(() => _command.WriteWarning("queued-first")).GetAwaiter().GetResult(); + _command.WriteWarning("direct-second"); + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedEnumeration")] +public sealed class TestAsyncCapturedEnumerationCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var streams = CapturePipelineStreams(); + WriteObject( + new CapturedWarningEnumerable( + message => streams.WriteWarning(message)), + enumerateCollection: true); + } + + private sealed class CapturedWarningEnumerable : IEnumerable + { + private readonly Action _writeWarning; + + public CapturedWarningEnumerable(Action writeWarning) + { + _writeWarning = writeWarning; + } + + public IEnumerator GetEnumerator() + { + using var completed = new ManualResetEventSlim(); + Exception? callbackException = null; + ThreadPool.UnsafeQueueUserWorkItem( + _ => + { + try + { + _writeWarning("captured-during-enumeration"); + } + catch (Exception exception) + { + callbackException = exception; + } + finally + { + completed.Set(); + } + }, + null); + Assert.True( + completed.Wait(TimeSpan.FromSeconds(5)), + "The context-free callback did not complete in time."); + if (callbackException is not null) + throw callbackException; + + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDerivedEndProcessing")] +public sealed class TestAsyncDerivedEndProcessingCommand : AsyncPSCmdlet +{ + protected override void EndProcessing() + { + WriteObject("before-base-end"); + base.EndProcessing(); + WriteObject("after-base-end"); + WriteWarning("after-base-warning"); + } +} + +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncDirectBarrierTail", + SupportsShouldProcess = true)] +public sealed class TestAsyncDirectBarrierTailCommand : AsyncPSCmdlet +{ + private static int _tailEnumerated; + + public static bool TailEnumeratedBeforeInteraction { get; private set; } + + public static void Reset() + { + Volatile.Write(ref _tailEnumerated, 0); + TailEnumeratedBeforeInteraction = false; + } + + protected override Task ProcessRecordAsync() + { + var streams = CapturePipelineStreams(); + Task.Run( + () => streams.WriteObject( + new OuterEnumerable(streams), + enumerateCollection: true)) + .GetAwaiter() + .GetResult(); + + _ = ShouldProcess("barrier-target"); + TailEnumeratedBeforeInteraction = Volatile.Read(ref _tailEnumerated) != 0; + return Task.CompletedTask; + } + + private sealed class OuterEnumerable(CapturedPipelineStreams streams) : IEnumerable + { + public IEnumerator GetEnumerator() + { + streams.WriteObject(new TailEnumerable(), enumerateCollection: true); + yield return "outer"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } + + private sealed class TailEnumerable : IEnumerable + { + public IEnumerator GetEnumerator() + { + Volatile.Write(ref _tailEnumerated, 1); + yield return "tail"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncReentrantInteractionTail", + SupportsShouldProcess = true)] +public sealed class TestAsyncReentrantInteractionTailCommand : AsyncPSCmdlet +{ + private static int _tailEnumerated; + + public static bool TailEnumeratedBeforeInteraction { get; private set; } + + public static void Reset() + { + Volatile.Write(ref _tailEnumerated, 0); + TailEnumeratedBeforeInteraction = false; + } + + protected override Task ProcessRecordAsync() + { + var streams = + CapturePipelineStreams(); + Task.Run( + () => streams.WriteObject( + new OuterEnumerable(this), + enumerateCollection: true)) + .GetAwaiter() + .GetResult(); + return Task.CompletedTask; + } + + private sealed class OuterEnumerable( + TestAsyncReentrantInteractionTailCommand command) + : IEnumerable + { + public IEnumerator GetEnumerator() + { + command.WriteObject( + new TailEnumerable(), + enumerateCollection: true); + _ = command.ShouldProcess( + "reentrant-target"); + TailEnumeratedBeforeInteraction = + Volatile.Read(ref _tailEnumerated) != 0; + yield return "outer"; + } + + System.Collections.IEnumerator + System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } + + private sealed class TailEnumerable : IEnumerable + { + public IEnumerator GetEnumerator() + { + Volatile.Write(ref _tailEnumerated, 1); + yield return "tail"; + } + + System.Collections.IEnumerator + System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncFailureRecoveryWrite")] +public sealed class TestAsyncFailureRecoveryWriteCommand : AsyncPSCmdlet +{ + public static bool TransportClearedAfterFailure { get; private set; } + + public static void Reset() + => TransportClearedAfterFailure = false; + + protected override void ProcessRecord() + { + try + { + base.ProcessRecord(); + } + catch (Exception) + { + TransportClearedAfterFailure = + typeof(AsyncPSCmdlet) + .GetField( + "_currentOutPipe", + BindingFlags.Instance | + BindingFlags.NonPublic)! + .GetValue(this) is null; + } + } + + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject( + new FailingEnumerable(), + enumerateCollection: true); + } + + private sealed class FailingEnumerable : IEnumerable + { + public IEnumerator GetEnumerator() + { + throw new InvalidOperationException( + "enumeration failed"); +#pragma warning disable CS0162 + yield break; +#pragma warning restore CS0162 + } + + System.Collections.IEnumerator + System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet( + VerbsDiagnostic.Test, + "AsyncHookContextIdentity")] +public sealed class TestAsyncHookContextIdentityCommand : AsyncPSCmdlet +{ + private SynchronizationContext? _beginContext; + + protected override Task BeginProcessingAsync() + { + _beginContext = + SynchronizationContext.Current; + return Task.CompletedTask; + } + + protected override Task ProcessRecordAsync() + { + WriteObject( + ReferenceEquals( + _beginContext, + SynchronizationContext.Current)); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyShouldProcess", SupportsShouldProcess = true)] +public sealed class TestAsyncEarlyShouldProcessCommand : AsyncPSCmdlet +{ + private bool _approved; + + protected override void ProcessRecord() + { + WriteObject("early-output"); + _approved = ShouldProcess("target"); + base.ProcessRecord(); + } + + protected override Task ProcessRecordAsync() + { + WriteObject(_approved); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncConstructorWrite")] +public sealed class TestAsyncConstructorWriteCommand : AsyncPSCmdlet +{ + public TestAsyncConstructorWriteCommand() + => WriteWarning("constructor output must not reach PowerShell"); + + protected override Task ProcessRecordAsync() + { + WriteObject("completed"); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncProgressSnapshot")] +public sealed class TestAsyncProgressSnapshotCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var progress = new ProgressRecord(17, "snapshot", "queued") + { + PercentComplete = 10 + }; + var totalProperty = progress.GetType().GetProperty("Total"); + if (totalProperty is not null) + { + totalProperty.SetValue( + progress, + Convert.ChangeType(42, totalProperty.PropertyType, CultureInfo.InvariantCulture)); + } + WriteProgress(progress); + progress.PercentComplete = 90; + if (totalProperty is not null) + { + totalProperty.SetValue( + progress, + Convert.ChangeType(99, totalProperty.PropertyType, CultureInfo.InvariantCulture)); + } + WriteObject("completed"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedErrorSnapshot")] +public sealed class TestAsyncCapturedErrorSnapshotCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var streams = CapturePipelineStreams(); + WriteObject(new ErrorSnapshotEnumerable(streams), enumerateCollection: true); + } + + private sealed class ErrorSnapshotEnumerable(CapturedPipelineStreams streams) : IEnumerable + { + public IEnumerator GetEnumerator() + { + var error = new ErrorRecord( + new InvalidOperationException("snapshot failure"), + "CapturedErrorSnapshot", + ErrorCategory.InvalidOperation, + targetObject: null) + { + ErrorDetails = new ErrorDetails("original details") + }; + streams.WriteError(error); + error.ErrorDetails = new ErrorDetails("mutated details"); + yield return "completed"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncTerminatingErrorSnapshot")] +public sealed class TestAsyncTerminatingErrorSnapshotCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var error = new ErrorRecord( + new InvalidOperationException("terminating snapshot failure"), + "TerminatingErrorSnapshot", + ErrorCategory.InvalidOperation, + targetObject: null) + { + ErrorDetails = new ErrorDetails("original terminating details") + }; + try + { + ThrowTerminatingError(error); + } + catch (PipelineStoppedException) + { + error.ErrorDetails = new ErrorDetails("mutated terminating details"); + throw; + } + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncStaleProgress")] +public sealed class TestAsyncStaleProgressCommand : AsyncPSCmdlet +{ + private IProgress? _firstRecordProgress; + private int _record; + + [Parameter(ValueFromPipeline = true)] + public int InputObject { get; set; } + + protected override async Task ProcessRecordAsync() + { + _record++; + if (_record == 1) + { + _firstRecordProgress = new Progress( + _ => WriteWarning("stale-first-record-warning")); + return; + } + + using var completed = new ManualResetEventSlim(); + Exception? callbackException = null; + ThreadPool.UnsafeQueueUserWorkItem( + _ => + { + try + { + _firstRecordProgress!.Report(InputObject); + } + catch (Exception exception) + { + callbackException = exception; + } + finally + { + completed.Set(); + } + }, + null); + Assert.True( + completed.Wait(TimeSpan.FromSeconds(5)), + "The context-free progress producer did not complete in time."); + await Task.Delay(100); + if (callbackException is not null) + throw callbackException; + WriteObject(InputObject); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncEnumeratorFailure")] +public sealed class TestAsyncEnumeratorFailureCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + var streams = CapturePipelineStreams(); + WriteObject(new FailingEnumerable(streams), enumerateCollection: true); + } + + private sealed class FailingEnumerable(CapturedPipelineStreams streams) : IEnumerable + { + public IEnumerator GetEnumerator() + { + streams.WriteWarning("before enumeration failure"); + yield return "first"; + throw new InvalidOperationException("enumeration failed"); + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncPipelineFailureWithThrowingCancellation")] +public sealed class TestAsyncPipelineFailureWithThrowingCancellationCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + _ = CancelToken.Register( + static () => throw new InvalidOperationException("cancellation callback failed")); + await Task.Yield(); + ThrowTerminatingError(new ErrorRecord( + new InvalidOperationException("pipeline failure"), + "PipelineFailure", + ErrorCategory.InvalidOperation, + targetObject: null)); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringHook")] +public sealed class TestAsyncDisposeDuringHookCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _tokenRead = new(); + + public static ManualResetEventSlim TokenRead => _tokenRead; + + public static bool ReadTokenAfterDispose { get; private set; } + + public static Exception? TokenReadException { get; private set; } + + public static void Reset() + { + _tokenRead.Dispose(); + _tokenRead = new ManualResetEventSlim(); + ReadTokenAfterDispose = false; + TokenReadException = null; + } + + protected override async Task ProcessRecordAsync() + { + Dispose(); + await Task.Yield(); + try + { + var token = CancelToken; + ReadTokenAfterDispose = true; + _tokenRead.Set(); + await Task.Delay(Timeout.InfiniteTimeSpan, token); + } + catch (OperationCanceledException) + { + throw; + } + catch (Exception exception) + { + TokenReadException = exception; + _tokenRead.Set(); + throw; + } + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposeDuringCancellation")] +public sealed class TestAsyncDisposeDuringCancellationCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _callbackCompleted = new(); + + public static ManualResetEventSlim CallbackCompleted => _callbackCompleted; + + public static bool TokenReadDuringCancellation { get; private set; } + + public static Exception? TokenReadException { get; private set; } + + public static void Reset() + { + _callbackCompleted.Dispose(); + _callbackCompleted = new ManualResetEventSlim(); + TokenReadDuringCancellation = false; + TokenReadException = null; + } + + protected override Task ProcessRecordAsync() + { + using var cancellationEntered = new ManualResetEventSlim(); + using var releaseCancellation = new ManualResetEventSlim(); + _ = CancelToken.Register( + () => + { + cancellationEntered.Set(); + releaseCancellation.Wait(TimeSpan.FromSeconds(5)); + try + { + _ = CancelToken; + TokenReadDuringCancellation = true; + } + catch (Exception exception) + { + TokenReadException = exception; + } + finally + { + _callbackCompleted.Set(); + } + }); + + _ = Task.Run(Dispose); + Assert.True( + cancellationEntered.Wait(TimeSpan.FromSeconds(5)), + "Cancellation did not enter the registered callback."); + ThreadPool.QueueUserWorkItem( + _ => + { + Thread.Sleep(100); + releaseCancellation.Set(); + }); + return Task.CompletedTask; + } + + protected override void EndProcessing() + { + // ProcessRecord intentionally disposes the command to exercise the cancellation race. + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDirectPipelineStop")] +public sealed class TestAsyncDirectPipelineStopCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _cancellationObserved = new(); + + public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + + public static void Reset() + { + _cancellationObserved.Dispose(); + _cancellationObserved = new ManualResetEventSlim(); + } + + protected override async Task ProcessRecordAsync() + { + _ = CancelToken.Register(_cancellationObserved.Set); + await Task.CompletedTask; + WriteObject(1); + WriteObject(2); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSuccessfulCompletion")] +public sealed class TestAsyncSuccessfulCompletionCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _cancellationObserved = new(); + + public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + + public static void Reset() + { + _cancellationObserved.Dispose(); + _cancellationObserved = new ManualResetEventSlim(); + } + + protected override async Task ProcessRecordAsync() + { + _ = CancelToken.Register(_cancellationObserved.Set); + await Task.Yield(); + WriteObject("completed"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncLateProgress")] +public sealed class TestAsyncLateProgressCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _callbackCompleted = new(); + private static ManualResetEventSlim _initialized = new(); + private static TestAsyncLateProgressCommand? _instance; + private static IProgress? _progress; + + public static ManualResetEventSlim CallbackCompleted => _callbackCompleted; + + public static ManualResetEventSlim Initialized => _initialized; + + public static Exception? LateWriteException { get; private set; } + + public static void Reset() + { + _callbackCompleted.Dispose(); + _initialized.Dispose(); + _callbackCompleted = new ManualResetEventSlim(); + _initialized = new ManualResetEventSlim(); + _instance = null; + _progress = null; + LateWriteException = null; + } + + public static void ReportAfterStop() + => _progress!.Report(50); + + public static void WriteAfterStop() + { + try + { + _instance!.WriteWarning("after stop"); + } + catch (Exception exception) + { + LateWriteException = exception; + } + } + + protected override async Task ProcessRecordAsync() + { + _instance = this; + _progress = new Progress(percent => + { + try + { + WriteProgress(new ProgressRecord(1, "late", "after stop") + { + PercentComplete = percent + }); + } + finally + { + _callbackCompleted.Set(); + } + }); + _initialized.Set(); + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousStop")] +public sealed class TestAsyncSynchronousStopCommand : AsyncPSCmdlet +{ + public void InvokeProcessRecord() + => base.ProcessRecord(); + + protected override Task ProcessRecordAsync() + { + Dispose(); + throw new OperationCanceledException("stopped synchronously"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncStoppedBeforeHook")] +public sealed class TestAsyncStoppedBeforeHookCommand : AsyncPSCmdlet +{ + public bool HookInvoked { get; private set; } + + public void InvokeProcessRecord() + => base.ProcessRecord(); + + public void InvokeStopProcessing() + => base.StopProcessing(); + + protected override Task ProcessRecordAsync() + { + HookInvoked = true; + return Task.CompletedTask; + } +} + +public sealed class ChoiceHost( + bool approved, + Exception? promptFailure = null, + ManualResetEventSlim? promptEntered = null, + ManualResetEventSlim? promptRelease = null) : PSHost +{ + private readonly Guid _id = Guid.NewGuid(); + private readonly ChoiceHostUserInterface _ui = + new(approved, promptFailure, promptEntered, promptRelease); + + public override Guid InstanceId => _id; + public override string Name => nameof(ChoiceHost); + public override Version Version => new(1, 0); + public override PSHostUserInterface UI => _ui; + + public IReadOnlyList ProgressRecords => _ui.ProgressRecords; + public override CultureInfo CurrentCulture => CultureInfo.InvariantCulture; + public override CultureInfo CurrentUICulture => CultureInfo.InvariantCulture; + public override void EnterNestedPrompt() { } + public override void ExitNestedPrompt() { } + public override void NotifyBeginApplication() { } + public override void NotifyEndApplication() { } + public override void SetShouldExit(int exitCode) { } +} + +public sealed class ChoiceHostUserInterface( + bool approved, + Exception? promptFailure, + ManualResetEventSlim? promptEntered, + ManualResetEventSlim? promptRelease) : PSHostUserInterface +{ + private readonly List _progressRecords = new(); + + public IReadOnlyList ProgressRecords => _progressRecords; + + public override PSHostRawUserInterface RawUI => null!; + + public override int PromptForChoice( + string caption, + string message, + Collection choices, + int defaultChoice) + { + promptEntered?.Set(); + Assert.True( + promptRelease?.Wait(TimeSpan.FromSeconds(5)) ?? true, + "The test host prompt was not released in time."); + return promptFailure is null ? (approved ? 0 : 1) : throw promptFailure; + } + + public override string ReadLine() => string.Empty; + public override SecureString ReadLineAsSecureString() => new(); + public override void Write(string value) { } + public override void Write(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value) { } + public override void WriteLine(string value) { } + public override void WriteErrorLine(string value) { } + public override void WriteDebugLine(string message) { } + public override void WriteProgress(long sourceId, ProgressRecord record) + => _progressRecords.Add(record); + public override void WriteVerboseLine(string message) { } + public override void WriteWarningLine(string message) { } + + public override PSCredential PromptForCredential( + string caption, + string message, + string userName, + string targetName) + => new(userName, new SecureString()); + + public override PSCredential PromptForCredential( + string caption, + string message, + string userName, + string targetName, + PSCredentialTypes allowedCredentialTypes, + PSCredentialUIOptions options) + => PromptForCredential(caption, message, userName, targetName); + + public override Dictionary Prompt( + string caption, + string message, + Collection descriptions) + => []; +} + +public sealed class ForwardingTaskScheduler : TaskScheduler +{ + private int _executeNextTaskInline; + private int _queuedTaskCount; + + public int QueuedTaskCount => Volatile.Read(ref _queuedTaskCount); + + public void Run(Action action) + { + Volatile.Write(ref _executeNextTaskInline, 1); + var task = Task.Factory.StartNew( + action, + CancellationToken.None, + TaskCreationOptions.None, + this); + task.GetAwaiter().GetResult(); + } + + protected override IEnumerable GetScheduledTasks() + => Array.Empty(); + + protected override void QueueTask(Task task) + { + if (Interlocked.Exchange(ref _executeNextTaskInline, 0) == 1) + { + Assert.True(TryExecuteTask(task)); + return; + } + + Interlocked.Increment(ref _queuedTaskCount); + ThreadPool.QueueUserWorkItem(_ => TryExecuteTask(task)); + } + + protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) + => false; +} + +public sealed class ForwardingSynchronizationContext : SynchronizationContext +{ + private int _postCount; + + public int PostCount => Volatile.Read(ref _postCount); + + public override void Post(SendOrPostCallback callback, object? state) + { + Interlocked.Increment(ref _postCount); + ThreadPool.QueueUserWorkItem(_ => callback(state)); + } +} diff --git a/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs b/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs new file mode 100644 index 000000000..1ce6ce308 --- /dev/null +++ b/PowerForge.Tests/AsyncPSCmdletLifecycleTests.cs @@ -0,0 +1,712 @@ +using System.Collections.ObjectModel; +using System.Globalization; +using System.Management.Automation; +using System.Management.Automation.Host; +using System.Management.Automation.Runspaces; +using System.Security; +using System.Threading; +using System.Threading.Tasks; +using PSPublishModule; + +namespace PowerForge.Tests; + +public sealed partial class AsyncPSCmdletTests +{ +[Fact] + public void AsyncPSCmdlet_dispose_without_an_active_hook_does_not_signal_stop() + { + using var command = new TestAsyncDisposableCommand(); + var stoppingToken = command.StoppingToken; + + command.Dispose(); + + Assert.False(stoppingToken.IsCancellationRequested); + } + + [Fact] + public void AsyncPSCmdlet_drops_pre_lifecycle_writes_from_other_threads() + { + using var command = new TestAsyncDisposableCommand(); + Exception? exception = null; + var workerThreadId = 0; + using var completed = new ManualResetEventSlim(); + var thread = new Thread(() => + { + workerThreadId = Environment.CurrentManagedThreadId; + exception = Record.Exception(() => command.WriteWarning("too-early")); + completed.Set(); + }); + + thread.Start(); + + Assert.True(completed.Wait(TimeSpan.FromSeconds(5))); + thread.Join(); + Assert.NotEqual(Environment.CurrentManagedThreadId, workerThreadId); + Assert.Null(exception); + } + + [Fact] + public void AsyncPSCmdlet_finishes_disposal_when_cancellation_callbacks_throw() + { + var command = new TestAsyncDisposableCommand(); + var stoppingToken = command.StoppingToken; + using var registration = stoppingToken.Register( + static () => throw new InvalidOperationException("cancellation callback failed")); + + command.InvokeStopProcessing(); + command.Dispose(); + + Assert.True(stoppingToken.IsCancellationRequested); + var sourceField = typeof(AsyncPSCmdlet).GetField( + "_cancelSource", + System.Reflection.BindingFlags.Instance | System.Reflection.BindingFlags.NonPublic); + var source = Assert.IsType(sourceField!.GetValue(command)); + Assert.Throws(source.Cancel); + } + + [Fact] + public void AsyncPSCmdlet_does_not_deadlock_when_a_synchronous_hook_queues_more_than_the_transport_capacity() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncLargeQueuedOutput", + typeof(TestAsyncLargeQueuedOutputCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncLargeQueuedOutput"); + + powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal(2048, powerShell.Streams.Warning.Count); + } + + [Fact] + public void AsyncPSCmdlet_restores_the_host_context_around_direct_pipeline_writes() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDirectContext", + typeof(TestAsyncDirectContextCommand), + helpFileName: null)); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-ObserveContext", + typeof(TestObserveContextCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncDirectContext") + .AddCommand("Test-ObserveContext"); + + var result = powerShell.Invoke(); + + Assert.Equal("context-output", Assert.Single(result).BaseObject); + Assert.Same(TestAsyncDirectContextCommand.HostContext, TestObserveContextCommand.ObservedContext); + } + + [Fact] + public void AsyncPSCmdlet_accepts_context_free_callbacks_through_a_captured_writer() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCapturedCallback", + typeof(TestAsyncCapturedCallbackCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCapturedCallback"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal("callback-output", Assert.Single(result).BaseObject); + Assert.Equal( + "callback-warning", + Assert.Single(powerShell.Streams.Warning).Message); + } + + [Fact] + public void AsyncPSCmdlet_marshals_command_details_to_the_pipeline_thread() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCommandDetail", + typeof(TestAsyncCommandDetailCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCommandDetail"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal("completed", Assert.Single(result).BaseObject); + } + + [Fact] + public void AsyncPSCmdlet_does_not_recursively_drain_when_enumeration_writes_to_the_pipeline() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncReentrantPump", + typeof(TestAsyncReentrantPumpCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncReentrantPump") + .AddParameter("Confirm", false); + + var result = powerShell.Invoke(); + + Assert.Collection(result, item => Assert.Equal("value", item.BaseObject)); + Assert.Equal( + ["after-enumeration", "during-enumeration"], + powerShell.Streams.Warning.Select(static warning => warning.Message).Order()); + } + + [Fact] + public void AsyncPSCmdlet_preserves_fifo_order_during_reentrant_queue_drains() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncReentrantFifo", + typeof(TestAsyncReentrantFifoCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncReentrantFifo"); + + var result = powerShell.Invoke(); + + Assert.Equal("value", Assert.Single(result).BaseObject); + Assert.Equal( + ["queued-first", "direct-second"], + powerShell.Streams.Warning.Select(static warning => warning.Message)); + } + + [Fact] + public void AsyncPSCmdlet_preserves_context_free_callbacks_while_their_output_is_pumped() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCapturedEnumeration", + typeof(TestAsyncCapturedEnumerationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCapturedEnumeration"); + + var result = powerShell.Invoke(); + + Assert.Equal("value", Assert.Single(result).BaseObject); + Assert.Equal( + "captured-during-enumeration", + Assert.Single( + powerShell.Streams.Warning).Message); + } + + [Fact] + public void AsyncPSCmdlet_keeps_derived_EndProcessing_active_after_the_base_call() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDerivedEndProcessing", + typeof(TestAsyncDerivedEndProcessingCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncDerivedEndProcessing"); + + var result = powerShell.Invoke(); + + Assert.Collection( + result, + item => Assert.Equal("before-base-end", item.BaseObject), + item => Assert.Equal("after-base-end", item.BaseObject)); + Assert.Equal( + "after-base-warning", + Assert.Single(powerShell.Streams.Warning).Message); + } + + [Fact] + public void AsyncPSCmdlet_drains_causal_tail_records_before_direct_interactions() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDirectBarrierTail", + typeof(TestAsyncDirectBarrierTailCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncDirectBarrierTail") + .AddParameter("Confirm", false); + TestAsyncDirectBarrierTailCommand.Reset(); + + var result = powerShell.Invoke(); + + Assert.Equal( + ["outer", "tail"], + result.Select(static item => item.BaseObject)); + Assert.True(TestAsyncDirectBarrierTailCommand.TailEnumeratedBeforeInteraction); + } + + [Fact] + public void AsyncPSCmdlet_drains_reentrant_records_before_direct_interactions() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncReentrantInteractionTail", + typeof( + TestAsyncReentrantInteractionTailCommand), + helpFileName: null)); + + using var runspace = + RunspaceFactory.CreateRunspace( + sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand( + "Test-AsyncReentrantInteractionTail") + .AddParameter( + "Confirm", + false); + TestAsyncReentrantInteractionTailCommand + .Reset(); + + var result = powerShell.Invoke(); + + Assert.Equal( + ["tail", "outer"], + result.Select( + static item => item.BaseObject)); + Assert.True( + TestAsyncReentrantInteractionTailCommand + .TailEnumeratedBeforeInteraction); + } + + [Fact] + public void AsyncPSCmdlet_clears_a_completed_pipe_before_a_recovery_write() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncFailureRecoveryWrite", + typeof( + TestAsyncFailureRecoveryWriteCommand), + helpFileName: null)); + + using var runspace = + RunspaceFactory.CreateRunspace( + sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand( + "Test-AsyncFailureRecoveryWrite"); + TestAsyncFailureRecoveryWriteCommand + .Reset(); + + var result = powerShell.Invoke(); + + Assert.Empty(result); + Assert.True( + TestAsyncFailureRecoveryWriteCommand + .TransportClearedAfterFailure); + } + + [Fact] + public void AsyncPSCmdlet_uses_a_distinct_context_for_each_hook() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncHookContextIdentity", + typeof( + TestAsyncHookContextIdentityCommand), + helpFileName: null)); + + using var runspace = + RunspaceFactory.CreateRunspace( + sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand( + "Test-AsyncHookContextIdentity"); + + var result = powerShell.Invoke(); + + Assert.False( + (bool)Assert.Single( + result).BaseObject); + } + + [Fact] + public void AsyncPSCmdlet_allows_ShouldProcess_before_the_async_base_hook_starts() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncEarlyShouldProcess", + typeof(TestAsyncEarlyShouldProcessCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncEarlyShouldProcess") + .AddParameter("Confirm", false); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + Assert.Collection( + result, + item => Assert.Equal("early-output", item.BaseObject), + item => Assert.True((bool)item.BaseObject)); + } + + [Fact] + public void AsyncPSCmdlet_drops_constructor_writes_before_pipeline_startup() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncConstructorWrite", + typeof(TestAsyncConstructorWriteCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncConstructorWrite"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Empty(powerShell.Streams.Warning); + Assert.Equal("completed", Assert.Single(result).BaseObject); + } + + [Fact] + public void AsyncPSCmdlet_snapshots_mutable_progress_before_queueing() + { + var sessionState = InitialSessionState.Create(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncProgressSnapshot", + typeof(TestAsyncProgressSnapshotCommand), + helpFileName: null)); + + var host = new ChoiceHost(approved: true); + using var runspace = RunspaceFactory.CreateRunspace(host, sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncProgressSnapshot"); + + var result = powerShell.Invoke(); + + Assert.Equal("completed", Assert.Single(result).BaseObject); + var captured = Assert.Single(host.ProgressRecords); + Assert.Equal(10, captured.PercentComplete); + var totalProperty = captured.GetType().GetProperty("Total"); + if (totalProperty is not null) + { + Assert.Equal( + Convert.ChangeType(42, totalProperty.PropertyType, CultureInfo.InvariantCulture), + totalProperty.GetValue(captured)); + } + } + + [Fact] + public void AsyncPSCmdlet_snapshots_captured_errors_before_queueing() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCapturedErrorSnapshot", + typeof(TestAsyncCapturedErrorSnapshotCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCapturedErrorSnapshot"); + + var result = powerShell.Invoke(); + + Assert.Equal("completed", Assert.Single(result).BaseObject); + var captured = Assert.Single(powerShell.Streams.Error); + Assert.Equal("original details", captured.ErrorDetails?.Message); + } + + [Fact] + public void AsyncPSCmdlet_snapshots_terminating_errors_before_queueing() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncTerminatingErrorSnapshot", + typeof(TestAsyncTerminatingErrorSnapshotCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncTerminatingErrorSnapshot"); + + var exception = Assert.Throws( + () => powerShell.Invoke()); + + Assert.Equal( + "original terminating details", + exception.ErrorRecord.ErrorDetails?.Message); + } + + [Fact] + public void AsyncPSCmdlet_rejects_context_callbacks_from_an_earlier_hook() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncStaleProgress", + typeof(TestAsyncStaleProgressCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddScript("1, 2 | Test-AsyncStaleProgress"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors); + Assert.Equal(2, Assert.Single(result).BaseObject); + Assert.Empty(powerShell.Streams.Warning); + } + + [Fact] + public void AsyncPSCmdlet_preserves_causal_stream_records_before_enumerator_failure() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncEnumeratorFailure", + typeof(TestAsyncEnumeratorFailureCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncEnumeratorFailure"); + + var exception = Record.Exception(() => powerShell.Invoke()); + + Assert.NotNull(exception); + Assert.Equal( + "before enumeration failure", + Assert.Single(powerShell.Streams.Warning).Message); + } + + [Fact] + public void AsyncPSCmdlet_preserves_pipeline_failure_when_cancellation_callbacks_throw() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncPipelineFailureWithThrowingCancellation", + typeof(TestAsyncPipelineFailureWithThrowingCancellationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncPipelineFailureWithThrowingCancellation"); + + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.StartsWith("PipelineFailure,", exception.ErrorRecord.FullyQualifiedErrorId, StringComparison.Ordinal); + Assert.Equal("pipeline failure", exception.InnerException?.Message); + } + + [Fact] + public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_the_async_hook_exits() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDisposeDuringHook", + typeof(TestAsyncDisposeDuringHookCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncDisposeDuringHook"); + TestAsyncDisposeDuringHookCommand.Reset(); + + powerShell.Invoke(); + + Assert.True( + TestAsyncDisposeDuringHookCommand.TokenRead.Wait(TimeSpan.FromSeconds(5)), + "The asynchronous hook did not finish reading the token in time."); + Assert.True(TestAsyncDisposeDuringHookCommand.ReadTokenAfterDispose); + Assert.Null(TestAsyncDisposeDuringHookCommand.TokenReadException); + } + + [Fact] + public void AsyncPSCmdlet_keeps_the_cancellation_source_alive_until_Cancel_returns() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDisposeDuringCancellation", + typeof(TestAsyncDisposeDuringCancellationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncDisposeDuringCancellation"); + TestAsyncDisposeDuringCancellationCommand.Reset(); + + powerShell.Invoke(); + + Assert.True( + TestAsyncDisposeDuringCancellationCommand.CallbackCompleted.Wait(TimeSpan.FromSeconds(5)), + "The cancellation callback did not complete in time."); + Assert.True(TestAsyncDisposeDuringCancellationCommand.TokenReadDuringCancellation); + Assert.Null(TestAsyncDisposeDuringCancellationCommand.TokenReadException); + } + + [Fact] + public void AsyncPSCmdlet_cancels_started_work_when_a_direct_write_stops_the_pipeline() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncDirectPipelineStop", + typeof(TestAsyncDirectPipelineStopCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncDirectPipelineStop") + .AddCommand("Select-Object") + .AddParameter("First", 1); + TestAsyncDirectPipelineStopCommand.Reset(); + + var result = powerShell.Invoke(); + + Assert.Equal(1, Assert.Single(result).BaseObject); + Assert.True( + TestAsyncDirectPipelineStopCommand.CancellationObserved.Wait(TimeSpan.FromSeconds(5)), + "The started operation was not canceled after the pipeline stopped."); + } + + [Fact] + public void AsyncPSCmdlet_does_not_cancel_after_successful_async_completion() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSuccessfulCompletion", + typeof(TestAsyncSuccessfulCompletionCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncSuccessfulCompletion"); + TestAsyncSuccessfulCompletionCommand.Reset(); + + var result = powerShell.Invoke(); + + Assert.Equal("completed", Assert.Single(result).BaseObject); + Assert.False(TestAsyncSuccessfulCompletionCommand.CancellationObserved.IsSet); + } + + [Fact] + public void AsyncPSCmdlet_suppresses_pipeline_stop_from_posted_progress_callbacks() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncLateProgress", + typeof(TestAsyncLateProgressCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncLateProgress"); + TestAsyncLateProgressCommand.Reset(); + + var invocation = powerShell.BeginInvoke(); + Assert.True( + TestAsyncLateProgressCommand.Initialized.Wait(TimeSpan.FromSeconds(5)), + "The progress callback was not initialized in time."); + powerShell.Stop(); + Assert.Throws(() => powerShell.EndInvoke(invocation)); + + TestAsyncLateProgressCommand.ReportAfterStop(); + TestAsyncLateProgressCommand.WriteAfterStop(); + + Assert.True( + TestAsyncLateProgressCommand.CallbackCompleted.Wait(TimeSpan.FromSeconds(5)), + "The posted progress callback did not complete in time."); + Assert.Null(TestAsyncLateProgressCommand.LateWriteException); + } + + [Fact] + public void AsyncPSCmdlet_normalizes_synchronous_cancellation_after_stop() + { + using var command = new TestAsyncSynchronousStopCommand(); + + Assert.Throws(command.InvokeProcessRecord); + } + + [Fact] + public void AsyncPSCmdlet_does_not_invoke_a_hook_after_stop() + { + using var command = new TestAsyncStoppedBeforeHookCommand(); + command.InvokeStopProcessing(); + + Assert.Throws(command.InvokeProcessRecord); + Assert.False(command.HookInvoked); + } +} diff --git a/PowerForge.Tests/AsyncPSCmdletTestCommands.cs b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs new file mode 100644 index 000000000..300962551 --- /dev/null +++ b/PowerForge.Tests/AsyncPSCmdletTestCommands.cs @@ -0,0 +1,663 @@ +using System.Collections.ObjectModel; +using System.Globalization; +using System.Management.Automation; +using System.Management.Automation.Host; +using System.Management.Automation.Runspaces; +using System.Security; +using System.Threading; +using System.Threading.Tasks; +using PSPublishModule; + +namespace PowerForge.Tests; + +[Cmdlet(VerbsDiagnostic.Test, "AsyncThreadAffinity")] +public sealed class TestAsyncThreadAffinityCommand : AsyncPSCmdlet +{ + private int _pipelineThreadId; + + protected override void BeginProcessing() + { + _pipelineThreadId = Environment.CurrentManagedThreadId; + base.BeginProcessing(); + } + + protected override async Task ProcessRecordAsync() + { + Assert.Equal(_pipelineThreadId, Environment.CurrentManagedThreadId); + await Task.Yield(); + WriteObject("post-await-output"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncQueuedOutput")] +public sealed class TestAsyncQueuedOutputCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + using var ready = new ManualResetEventSlim(); + Exception? workerException = null; + ThreadPool.QueueUserWorkItem(_ => + { + try + { + WriteObject("queued-output"); + } + catch (Exception ex) + { + workerException = ex; + } + finally + { + ready.Set(); + } + }); + + Assert.True(ready.Wait(TimeSpan.FromSeconds(5)), "Worker thread did not write output in time."); + if (workerException is not null) + throw workerException; + + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronizationContext")] +public sealed class TestAsyncSynchronizationContextCommand : AsyncPSCmdlet +{ + private ForwardingSynchronizationContext? _context; + + protected override void ProcessRecord() + { + var previousContext = SynchronizationContext.Current; + _context = new ForwardingSynchronizationContext(); + SynchronizationContext.SetSynchronizationContext(_context); + try + { + base.ProcessRecord(); + } + finally + { + SynchronizationContext.SetSynchronizationContext(previousContext); + } + } + + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(_context!.PostCount); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncTaskScheduler")] +public sealed class TestAsyncTaskSchedulerCommand : AsyncPSCmdlet +{ + private readonly ForwardingTaskScheduler _scheduler = new(); + private int _pipelineThreadId; + + protected override void ProcessRecord() + { + _pipelineThreadId = Environment.CurrentManagedThreadId; + _scheduler.Run(base.ProcessRecord); + } + + protected override async Task ProcessRecordAsync() + { + Assert.Equal(_pipelineThreadId, Environment.CurrentManagedThreadId); + Assert.NotSame(_scheduler, TaskScheduler.Current); + await Task.Factory.StartNew(static () => { }); + await Task.Yield(); + WriteObject(_scheduler.QueuedTaskCount); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncWriteOrdering")] +public sealed class TestAsyncWriteOrderingCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + Task.Run(() => WriteObject("first")).GetAwaiter().GetResult(); + WriteObject("second"); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncInformationTags")] +public sealed class TestAsyncInformationTagsCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + var tags = new[] { "before" }; + Task.Run(() => WriteInformation("message", tags)).GetAwaiter().GetResult(); + tags[0] = "after"; + + var informationRecord = new InformationRecord("record-message", "record-source"); + informationRecord.Tags.Add("record-before"); + Task.Run(() => WriteInformation(informationRecord)).GetAwaiter().GetResult(); + informationRecord.Tags[0] = "record-after"; + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousError")] +public sealed class TestAsyncSynchronousErrorCommand : AsyncPSCmdlet +{ + private static int _reachedAfterError; + + public static bool ReachedAfterError => Volatile.Read(ref _reachedAfterError) != 0; + + public static void Reset() + => Volatile.Write(ref _reachedAfterError, 0); + + protected override Task ProcessRecordAsync() + { + WriteError(new ErrorRecord( + new InvalidOperationException("stopping error"), + "AsyncSynchronousError", + ErrorCategory.InvalidOperation, + targetObject: null)); + Volatile.Write(ref _reachedAfterError, 1); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousEnumeration")] +public sealed class TestAsyncSynchronousEnumerationCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + var values = new[] { 1, 2 }; + WriteObject(values, enumerateCollection: true); + values[0] = 9; + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousReentrantDrain")] +public sealed class TestAsyncSynchronousReentrantDrainCommand : AsyncPSCmdlet +{ + [Parameter] + public SwitchParameter Fail { get; set; } + + protected override Task ProcessRecordAsync() + { + var streams = CapturePipelineStreams(); + WriteObject( + new ReentrantWarningEnumerable(streams.WriteWarning), + enumerateCollection: true); + + if (Fail) + throw new InvalidOperationException("synchronous reentrant drain failure"); + + return Task.CompletedTask; + } + + private sealed class ReentrantWarningEnumerable(Action writeWarning) : IEnumerable + { + public IEnumerator GetEnumerator() + { + writeWarning("reentrant-during-drain"); + yield return "value"; + } + + System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() + => GetEnumerator(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousFailure")] +public sealed class TestAsyncSynchronousFailureCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + Task.Run(() => WriteWarning("before-failure")).GetAwaiter().GetResult(); + throw new InvalidOperationException("synchronous hook failure"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronousPipelineStop")] +public sealed class TestAsyncSynchronousPipelineStopCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _cancellationObserved = new(); + + public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + + public static void Reset() + { + _cancellationObserved.Dispose(); + _cancellationObserved = new ManualResetEventSlim(); + } + + public void InvokeProcessRecord() + => base.ProcessRecord(); + + protected override Task ProcessRecordAsync() + { + _ = CancelToken.Register(_cancellationObserved.Set); + throw new PipelineStoppedException(); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncNullInformationTags")] +public sealed class TestAsyncNullInformationTagsCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteInformation("untagged", tags: null); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncStaleInteraction", SupportsShouldProcess = true)] +public sealed class TestAsyncStaleInteractionCommand : AsyncPSCmdlet +{ + private static readonly ManualResetEventSlim SecondRecordStarted = new(); + private static readonly ManualResetEventSlim StaleInteractionFinished = new(); + private static int _recordNumber; + + [Parameter(Mandatory = true, ValueFromPipeline = true)] + public string InputObject { get; set; } = string.Empty; + + public static Exception? StaleInteractionException { get; private set; } + + public static void Reset() + { + SecondRecordStarted.Reset(); + StaleInteractionFinished.Reset(); + Volatile.Write(ref _recordNumber, 0); + StaleInteractionException = null; + } + + protected override async Task ProcessRecordAsync() + { + if (Interlocked.Increment(ref _recordNumber) == 1) + { + _ = Task.Run(() => + { + Assert.True( + SecondRecordStarted.Wait(TimeSpan.FromSeconds(5)), + "The second lifecycle did not start in time."); + try + { + WriteWarning("stale-warning"); + _ = ShouldProcess("stale-target"); + } + catch (Exception exception) + { + StaleInteractionException = exception; + } + finally + { + StaleInteractionFinished.Set(); + } + }); + return; + } + + SecondRecordStarted.Set(); + await Task.Run(() => + { + Assert.True( + StaleInteractionFinished.Wait(TimeSpan.FromSeconds(5)), + "The stale interaction did not finish in time."); + }); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCancellation")] +public sealed class TestAsyncCancellationCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _started = new(); + + public static ManualResetEventSlim Started => _started; + + public static void Reset() + { + _started.Dispose(); + _started = new ManualResetEventSlim(); + } + + protected override async Task ProcessRecordAsync() + { + _started.Set(); + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCancellationWrite")] +public sealed class TestAsyncCancellationWriteCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _started = new(); + private static ManualResetEventSlim _writeAttempted = new(); + + public static ManualResetEventSlim Started => _started; + public static ManualResetEventSlim WriteAttempted => _writeAttempted; + + public static Exception? BackgroundWriteException { get; private set; } + + public static void Reset() + { + _started.Dispose(); + _writeAttempted.Dispose(); + _started = new ManualResetEventSlim(); + _writeAttempted = new ManualResetEventSlim(); + BackgroundWriteException = null; + } + + protected override async Task ProcessRecordAsync() + { + _started.Set(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } + catch (OperationCanceledException) when (CancelToken.IsCancellationRequested) + { + try + { + WriteProgress(new ProgressRecord(1, "cancelled", "finishing")); + WriteWarning("cancelled"); + } + catch (Exception exception) + { + BackgroundWriteException = exception; + } + finally + { + _writeAttempted.Set(); + } + } + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncCapturedCancellationWrite")] +public sealed class TestAsyncCapturedCancellationWriteCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _started = new(); + private static ManualResetEventSlim _writeAttempted = new(); + + public static ManualResetEventSlim Started => _started; + public static ManualResetEventSlim WriteAttempted => _writeAttempted; + + public static Exception? BackgroundWriteException { get; private set; } + + public static void Reset() + { + _started.Dispose(); + _writeAttempted.Dispose(); + _started = new ManualResetEventSlim(); + _writeAttempted = new ManualResetEventSlim(); + BackgroundWriteException = null; + } + + protected override async Task ProcessRecordAsync() + { + var streams = CapturePipelineStreams(); + _started.Set(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, CancelToken); + } + catch (OperationCanceledException) when (CancelToken.IsCancellationRequested) + { + var completed = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + ThreadPool.UnsafeQueueUserWorkItem( + _ => + { + try + { + streams.WriteWarning("cancelled callback"); + } + catch (Exception exception) + { + BackgroundWriteException = exception; + } + finally + { + _writeAttempted.Set(); + completed.TrySetResult(true); + } + }, + null); + await completed.Task.ConfigureAwait(false); + } + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncTerminatingError")] +public sealed class TestAsyncTerminatingErrorCommand : AsyncPSCmdlet +{ + private static int _reachedAfterTermination; + + public static bool ReachedAfterTermination => Volatile.Read(ref _reachedAfterTermination) != 0; + + public static void Reset() + => Volatile.Write(ref _reachedAfterTermination, 0); + + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + ThrowTerminatingError(new ErrorRecord( + new InvalidOperationException("async terminating failure"), + "AsyncTerminatingError", + ErrorCategory.InvalidOperation, + targetObject: null)); + Volatile.Write(ref _reachedAfterTermination, 1); + WriteObject("unreachable"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncEarlyTerminatingError")] +public sealed class TestAsyncEarlyTerminatingErrorCommand : AsyncPSCmdlet +{ + protected override void BeginProcessing() + => ThrowTerminatingError(new ErrorRecord( + new InvalidOperationException("early terminating failure"), + "EarlyTerminatingError", + ErrorCategory.InvalidOperation, + targetObject: null)); +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncShouldContinue")] +public sealed class TestAsyncShouldContinueCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + WriteObject(ShouldContinue("Proceed?", "Question")); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncClaimedShouldContinue")] +public sealed class TestAsyncClaimedShouldContinueCommand : AsyncPSCmdlet +{ + private static ManualResetEventSlim _replyObserved = new(); + private static ManualResetEventSlim _cancellationObserved = new(); + + public static ManualResetEventSlim ReplyObserved => _replyObserved; + public static ManualResetEventSlim CancellationObserved => _cancellationObserved; + public static bool SideEffectStarted { get; private set; } + + public static void Reset() + { + _replyObserved.Dispose(); + _cancellationObserved.Dispose(); + _replyObserved = new ManualResetEventSlim(); + _cancellationObserved = new ManualResetEventSlim(); + SideEffectStarted = false; + } + + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + _ = CancelToken.Register(_cancellationObserved.Set); + try + { + _ = ShouldContinue("Proceed?", "Question"); + SideEffectStarted = true; + } + finally + { + _replyObserved.Set(); + } + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncOperationCancellation")] +public sealed class TestAsyncOperationCancellationCommand : AsyncPSCmdlet +{ + protected override async Task ProcessRecordAsync() + { + await Task.Yield(); + throw new OperationCanceledException("operation timeout"); + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncLateWrite")] +public sealed class TestAsyncLateWriteCommand : AsyncPSCmdlet +{ + private static TestAsyncLateWriteCommand? _instance; + + public static Exception? LateWriteException { get; private set; } + + public static void Reset() + { + _instance = null; + LateWriteException = null; + } + + public static void WriteAfterCompletion() + { + try + { + _instance!.WriteWarning("late"); + } + catch (Exception exception) + { + LateWriteException = exception; + } + } + + protected override Task ProcessRecordAsync() + { + _instance = this; + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncLateInteraction", SupportsShouldProcess = true)] +public sealed class TestAsyncLateInteractionCommand : AsyncPSCmdlet +{ + private static TestAsyncLateInteractionCommand? _instance; + + public static Exception? LateInteractionException { get; private set; } + + public static void Reset() + { + _instance = null; + LateInteractionException = null; + } + + public static void InteractAfterCompletion() + => Task.Run(() => + { + try + { + _instance!.ShouldProcess("late-target"); + } + catch (Exception exception) + { + LateInteractionException = exception; + } + }).GetAwaiter().GetResult(); + + protected override Task ProcessRecordAsync() + { + _instance = this; + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDisposable")] +public sealed class TestAsyncDisposableCommand : AsyncPSCmdlet +{ + public CancellationToken StoppingToken => CancelToken; + + public void InvokeStopProcessing() + => base.StopProcessing(); +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncLargeQueuedOutput")] +public sealed class TestAsyncLargeQueuedOutputCommand : AsyncPSCmdlet +{ + protected override Task ProcessRecordAsync() + { + using var completed = new ManualResetEventSlim(); + Exception? workerException = null; + var thread = new Thread(() => + { + try + { + for (var i = 0; i < 2048; i++) + WriteWarning($"queued-{i}"); + } + catch (Exception exception) + { + workerException = exception; + } + finally + { + completed.Set(); + } + }); + + thread.Start(); + Assert.True(completed.Wait(TimeSpan.FromSeconds(5)), "The producer blocked while filling the pipeline transport."); + thread.Join(); + if (workerException is not null) + throw workerException; + + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "AsyncDirectContext")] +public sealed class TestAsyncDirectContextCommand : AsyncPSCmdlet +{ + public static SynchronizationContext? HostContext { get; private set; } + + protected override void ProcessRecord() + { + var previousContext = SynchronizationContext.Current; + HostContext = new ForwardingSynchronizationContext(); + SynchronizationContext.SetSynchronizationContext(HostContext); + try + { + base.ProcessRecord(); + } + finally + { + SynchronizationContext.SetSynchronizationContext(previousContext); + } + } + + protected override Task ProcessRecordAsync() + { + WriteObject("context-output"); + return Task.CompletedTask; + } +} + +[Cmdlet(VerbsDiagnostic.Test, "ObserveContext")] +public sealed class TestObserveContextCommand : PSCmdlet +{ + [Parameter(ValueFromPipeline = true)] + public object? InputObject { get; set; } + + public static SynchronizationContext? ObservedContext { get; private set; } + + protected override void ProcessRecord() + { + ObservedContext = SynchronizationContext.Current; + WriteObject(InputObject); + } +} diff --git a/PowerForge.Tests/AsyncPSCmdletTests.cs b/PowerForge.Tests/AsyncPSCmdletTests.cs index 31169c56b..ca8c1835a 100644 --- a/PowerForge.Tests/AsyncPSCmdletTests.cs +++ b/PowerForge.Tests/AsyncPSCmdletTests.cs @@ -1,12 +1,16 @@ +using System.Collections.ObjectModel; +using System.Globalization; using System.Management.Automation; +using System.Management.Automation.Host; using System.Management.Automation.Runspaces; +using System.Security; using System.Threading; using System.Threading.Tasks; using PSPublishModule; namespace PowerForge.Tests; -public sealed class AsyncPSCmdletTests +public sealed partial class AsyncPSCmdletTests { [Fact] public void AsyncPSCmdlet_starts_hooks_on_the_pipeline_thread_and_pumps_after_await() @@ -95,147 +99,503 @@ public void AsyncPSCmdlet_does_not_capture_a_custom_task_scheduler() var item = Assert.Single(result); Assert.Equal(0, item.BaseObject); } -} -[Cmdlet(VerbsDiagnostic.Test, "AsyncThreadAffinity")] -public sealed class TestAsyncThreadAffinityCommand : AsyncPSCmdlet -{ - private int _pipelineThreadId; + [Fact] + public void AsyncPSCmdlet_preserves_fifo_order_for_pipeline_thread_and_worker_writes() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncWriteOrdering", + typeof(TestAsyncWriteOrderingCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncWriteOrdering"); - protected override void BeginProcessing() + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + Assert.Collection( + result, + item => Assert.Equal("first", item.BaseObject), + item => Assert.Equal("second", item.BaseObject)); + } + + [Fact] + public void AsyncPSCmdlet_snapshots_information_tags_before_queueing() { - _pipelineThreadId = Environment.CurrentManagedThreadId; - base.BeginProcessing(); + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncInformationTags", + typeof(TestAsyncInformationTagsCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncInformationTags"); + + powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + Assert.Collection( + powerShell.Streams.Information, + record => + { + Assert.Equal("message", record.MessageData); + Assert.Equal(["before"], record.Tags); + }, + record => + { + Assert.Equal("record-message", record.MessageData); + Assert.Equal("record-source", record.Source); + Assert.Equal(["record-before"], record.Tags); + }); } - protected override async Task ProcessRecordAsync() + [Fact] + public void AsyncPSCmdlet_applies_stopping_preferences_before_the_hook_continues() { - Assert.Equal(_pipelineThreadId, Environment.CurrentManagedThreadId); - await Task.Yield(); - WriteObject("post-await-output"); + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSynchronousError", + typeof(TestAsyncSynchronousErrorCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncSynchronousError") + .AddParameter("ErrorAction", ActionPreference.Stop); + TestAsyncSynchronousErrorCommand.Reset(); + + _ = Assert.ThrowsAny(() => powerShell.Invoke()); + + Assert.False(TestAsyncSynchronousErrorCommand.ReachedAfterError); } -} -[Cmdlet(VerbsDiagnostic.Test, "AsyncQueuedOutput")] -public sealed class TestAsyncQueuedOutputCommand : AsyncPSCmdlet -{ - protected override Task ProcessRecordAsync() + [Fact] + public void AsyncPSCmdlet_enumerates_pipeline_thread_collections_before_the_hook_mutates_them() { - using var ready = new ManualResetEventSlim(); - Exception? workerException = null; - ThreadPool.QueueUserWorkItem(_ => - { - try - { - WriteObject("queued-output"); - } - catch (Exception ex) - { - workerException = ex; - } - finally - { - ready.Set(); - } - }); + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSynchronousEnumeration", + typeof(TestAsyncSynchronousEnumerationCommand), + helpFileName: null)); - Assert.True(ready.Wait(TimeSpan.FromSeconds(5)), "Worker thread did not write output in time."); - if (workerException is not null) - throw workerException; + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncSynchronousEnumeration"); + + var result = powerShell.Invoke(); - return Task.CompletedTask; + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + Assert.Collection( + result, + item => Assert.Equal(1, item.BaseObject), + item => Assert.Equal(2, item.BaseObject)); } -} -[Cmdlet(VerbsDiagnostic.Test, "AsyncSynchronizationContext")] -public sealed class TestAsyncSynchronizationContextCommand : AsyncPSCmdlet -{ - private ForwardingSynchronizationContext? _context; + [Theory] + [InlineData(false)] + [InlineData(true)] + public void AsyncPSCmdlet_drains_reentrant_records_from_synchronous_hooks(bool fail) + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSynchronousReentrantDrain", + typeof(TestAsyncSynchronousReentrantDrainCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncSynchronousReentrantDrain") + .AddParameter("Fail", fail); + + if (fail) + _ = Assert.ThrowsAny(() => powerShell.Invoke()); + else + Assert.Equal("value", Assert.Single(powerShell.Invoke()).BaseObject); + + Assert.Equal( + "reentrant-during-drain", + Assert.Single(powerShell.Streams.Warning).Message); + } - protected override void ProcessRecord() + [Fact] + public void AsyncPSCmdlet_drains_worker_records_before_a_synchronous_hook_failure() { - var previousContext = SynchronizationContext.Current; - _context = new ForwardingSynchronizationContext(); - SynchronizationContext.SetSynchronizationContext(_context); - try - { - base.ProcessRecord(); - } - finally - { - SynchronizationContext.SetSynchronizationContext(previousContext); - } + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncSynchronousFailure", + typeof(TestAsyncSynchronousFailureCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncSynchronousFailure"); + + _ = Assert.ThrowsAny(() => powerShell.Invoke()); + + var warning = Assert.Single(powerShell.Streams.Warning); + Assert.Equal("before-failure", warning.Message); } - protected override async Task ProcessRecordAsync() + [Fact] + public void AsyncPSCmdlet_cancels_the_shared_token_after_a_synchronous_pipeline_stop() { - await Task.Yield(); - WriteObject(_context!.PostCount); + using var command = new TestAsyncSynchronousPipelineStopCommand(); + TestAsyncSynchronousPipelineStopCommand.Reset(); + + Assert.Throws(command.InvokeProcessRecord); + Assert.True( + TestAsyncSynchronousPipelineStopCommand.CancellationObserved.Wait(TimeSpan.FromSeconds(5)), + "The shared cancellation token did not observe the synchronous pipeline stop."); } -} -[Cmdlet(VerbsDiagnostic.Test, "AsyncTaskScheduler")] -public sealed class TestAsyncTaskSchedulerCommand : AsyncPSCmdlet -{ - private readonly ForwardingTaskScheduler _scheduler = new(); + [Fact] + public void AsyncPSCmdlet_preserves_null_information_tags_after_await() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncNullInformationTags", + typeof(TestAsyncNullInformationTagsCommand), + helpFileName: null)); - protected override void ProcessRecord() - => _scheduler.Run(base.ProcessRecord); + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncNullInformationTags"); + + powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + var record = Assert.Single(powerShell.Streams.Information); + Assert.Equal("untagged", record.MessageData); + Assert.Empty(record.Tags); + } - protected override async Task ProcessRecordAsync() + [Fact] + public void AsyncPSCmdlet_rejects_interactions_from_an_older_record_lifecycle() { - Assert.Same(_scheduler, TaskScheduler.Current); - await Task.Yield(); - WriteObject(_scheduler.QueuedTaskCount); + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncStaleInteraction", + typeof(TestAsyncStaleInteractionCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell + .AddCommand("Test-AsyncStaleInteraction") + .AddParameter("Confirm", false); + TestAsyncStaleInteractionCommand.Reset(); + + powerShell.Invoke(new[] { "first", "second" }); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + Assert.IsType(TestAsyncStaleInteractionCommand.StaleInteractionException); + Assert.DoesNotContain(powerShell.Streams.Warning, static warning => warning.Message == "stale-warning"); } -} -public sealed class ForwardingTaskScheduler : TaskScheduler -{ - private int _executeNextTaskInline; - private int _queuedTaskCount; + [Fact] + public void AsyncPSCmdlet_translates_cancellation_to_a_pipeline_stop() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCancellation", + typeof(TestAsyncCancellationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCancellation"); + TestAsyncCancellationCommand.Reset(); - public int QueuedTaskCount => Volatile.Read(ref _queuedTaskCount); + var invocation = powerShell.BeginInvoke(); + Assert.True( + TestAsyncCancellationCommand.Started.Wait(TimeSpan.FromSeconds(5)), + "The asynchronous cmdlet did not start in time."); - public void Run(Action action) + powerShell.Stop(); + + var exception = Assert.Throws(() => powerShell.EndInvoke(invocation)); + Assert.IsNotType(exception); + } + + [Fact] + public void AsyncPSCmdlet_propagates_pipeline_stop_to_background_stream_writes() { - Volatile.Write(ref _executeNextTaskInline, 1); - var task = Task.Factory.StartNew( - action, - CancellationToken.None, - TaskCreationOptions.None, - this); - task.GetAwaiter().GetResult(); + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCancellationWrite", + typeof(TestAsyncCancellationWriteCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCancellationWrite"); + TestAsyncCancellationWriteCommand.Reset(); + + var invocation = powerShell.BeginInvoke(); + Assert.True( + TestAsyncCancellationWriteCommand.Started.Wait(TimeSpan.FromSeconds(5)), + "The asynchronous cmdlet did not start in time."); + + powerShell.Stop(); + Assert.Throws(() => powerShell.EndInvoke(invocation)); + Assert.True( + TestAsyncCancellationWriteCommand.WriteAttempted.Wait(TimeSpan.FromSeconds(5)), + "The background hook did not observe the stop in time."); + + Assert.IsType( + TestAsyncCancellationWriteCommand.BackgroundWriteException); } - protected override IEnumerable GetScheduledTasks() - => Array.Empty(); + [Fact] + public void AsyncPSCmdlet_drops_captured_callback_writes_after_pipeline_stop() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncCapturedCancellationWrite", + typeof(TestAsyncCapturedCancellationWriteCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncCapturedCancellationWrite"); + TestAsyncCapturedCancellationWriteCommand.Reset(); + + var invocation = powerShell.BeginInvoke(); + Assert.True( + TestAsyncCapturedCancellationWriteCommand.Started.Wait(TimeSpan.FromSeconds(5)), + "The asynchronous cmdlet did not start in time."); + + powerShell.Stop(); + Assert.Throws(() => powerShell.EndInvoke(invocation)); + Assert.True( + TestAsyncCapturedCancellationWriteCommand.WriteAttempted.Wait(TimeSpan.FromSeconds(5)), + "The captured callback did not observe the stop in time."); + Assert.Null(TestAsyncCapturedCancellationWriteCommand.BackgroundWriteException); + } - protected override void QueueTask(Task task) + [Fact] + public void AsyncPSCmdlet_marshals_terminating_errors_to_the_pipeline_thread() { - if (Interlocked.Exchange(ref _executeNextTaskInline, 0) == 1) - { - Assert.True(TryExecuteTask(task)); - return; - } + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncTerminatingError", + typeof(TestAsyncTerminatingErrorCommand), + helpFileName: null)); - Interlocked.Increment(ref _queuedTaskCount); - ThreadPool.QueueUserWorkItem(_ => TryExecuteTask(task)); + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncTerminatingError"); + TestAsyncTerminatingErrorCommand.Reset(); + + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.StartsWith("AsyncTerminatingError,", exception.ErrorRecord.FullyQualifiedErrorId, StringComparison.Ordinal); + Assert.Equal("async terminating failure", exception.InnerException?.Message); + Assert.False(TestAsyncTerminatingErrorCommand.ReachedAfterTermination); } - protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) - => false; -} + [Fact] + public void AsyncPSCmdlet_preserves_terminating_errors_before_the_initial_base_hook() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncEarlyTerminatingError", + typeof(TestAsyncEarlyTerminatingErrorCommand), + helpFileName: null)); -public sealed class ForwardingSynchronizationContext : SynchronizationContext -{ - private int _postCount; + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncEarlyTerminatingError"); + + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.StartsWith("EarlyTerminatingError,", exception.ErrorRecord.FullyQualifiedErrorId, StringComparison.Ordinal); + Assert.Equal("early terminating failure", exception.InnerException?.Message); + } + + [Theory] + [InlineData(true)] + [InlineData(false)] + public void AsyncPSCmdlet_marshals_ShouldContinue_to_the_pipeline_thread(bool approved) + { + var sessionState = InitialSessionState.Create(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncShouldContinue", + typeof(TestAsyncShouldContinueCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(new ChoiceHost(approved), sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncShouldContinue"); + + var result = powerShell.Invoke(); + + Assert.False(powerShell.HadErrors, string.Join(Environment.NewLine, powerShell.Streams.Error.Select(static error => error.ToString()))); + var item = Assert.Single(result); + Assert.Equal(approved, item.BaseObject); + } - public int PostCount => Volatile.Read(ref _postCount); + [Fact] + public void AsyncPSCmdlet_returns_host_interaction_failures_to_the_requesting_worker() + { + var sessionState = InitialSessionState.Create(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncShouldContinue", + typeof(TestAsyncShouldContinueCommand), + helpFileName: null)); + + var promptFailure = new InvalidOperationException("host prompt failed"); + using var runspace = RunspaceFactory.CreateRunspace(new ChoiceHost(approved: false, promptFailure), sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncShouldContinue"); - public override void Post(SendOrPostCallback callback, object? state) + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.Same(promptFailure, exception.InnerException); + } + + [Fact] + public async Task AsyncPSCmdlet_keeps_a_claimed_host_reply_observed_during_cancellation() { - Interlocked.Increment(ref _postCount); - ThreadPool.QueueUserWorkItem(_ => callback(state)); + using var promptEntered = new ManualResetEventSlim(); + using var promptRelease = new ManualResetEventSlim(); + var sessionState = InitialSessionState.Create(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncClaimedShouldContinue", + typeof(TestAsyncClaimedShouldContinueCommand), + helpFileName: null)); + + var host = new ChoiceHost( + approved: true, + promptEntered: promptEntered, + promptRelease: promptRelease); + using var runspace = RunspaceFactory.CreateRunspace(host, sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncClaimedShouldContinue"); + TestAsyncClaimedShouldContinueCommand.Reset(); + + var invocation = powerShell.BeginInvoke(); + Assert.True( + promptEntered.Wait(TimeSpan.FromSeconds(5)), + "The host interaction did not start in time."); + var stopTask = Task.Run(powerShell.Stop); + Assert.True( + TestAsyncClaimedShouldContinueCommand.CancellationObserved.Wait(TimeSpan.FromSeconds(5)), + "The claimed requester did not observe pipeline cancellation."); + promptRelease.Set(); + + await stopTask.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.Throws(() => powerShell.EndInvoke(invocation)); + Assert.True( + TestAsyncClaimedShouldContinueCommand.ReplyObserved.Wait(TimeSpan.FromSeconds(5)), + "The claimed requester abandoned the host reply during cancellation."); + Assert.False(TestAsyncClaimedShouldContinueCommand.SideEffectStarted); } + + [Fact] + public void AsyncPSCmdlet_preserves_operation_cancellation_that_is_not_a_pipeline_stop() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncOperationCancellation", + typeof(TestAsyncOperationCancellationCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncOperationCancellation"); + + var exception = Assert.Throws(() => powerShell.Invoke()); + + Assert.IsAssignableFrom(exception.InnerException); + Assert.IsNotType(exception.InnerException); + } + + [Fact] + public void AsyncPSCmdlet_drops_late_worker_writes_after_the_pipeline_closes() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncLateWrite", + typeof(TestAsyncLateWriteCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncLateWrite"); + TestAsyncLateWriteCommand.Reset(); + + powerShell.Invoke(); + TestAsyncLateWriteCommand.WriteAfterCompletion(); + + Assert.Null(TestAsyncLateWriteCommand.LateWriteException); + } + + [Fact] + public void AsyncPSCmdlet_rejects_late_worker_interactions_after_the_pipeline_closes() + { + var sessionState = InitialSessionState.CreateDefault(); + sessionState.Commands.Add(new SessionStateCmdletEntry( + "Test-AsyncLateInteraction", + typeof(TestAsyncLateInteractionCommand), + helpFileName: null)); + + using var runspace = RunspaceFactory.CreateRunspace(sessionState); + runspace.Open(); + using var powerShell = PowerShell.Create(); + powerShell.Runspace = runspace; + powerShell.AddCommand("Test-AsyncLateInteraction"); + TestAsyncLateInteractionCommand.Reset(); + + powerShell.Invoke(); + TestAsyncLateInteractionCommand.InteractAfterCompletion(); + + Assert.True( + TestAsyncLateInteractionCommand.LateInteractionException is + InvalidOperationException or PipelineStoppedException, + $"Unexpected late-interaction result: {TestAsyncLateInteractionCommand.LateInteractionException}"); + } + }