XpandPwsh

1.221.0.22

Cmdlets/src/XpandPwsh.Cmdlets/AsyncCmdlet.cs

using System;
using System.Collections.Concurrent;
using System.Management.Automation;
using System.Threading;
using System.Threading.Tasks;
using System.Collections.Generic;
using System.Runtime.ExceptionServices;
 
namespace XpandPwsh.CmdLets{
 
    /// <summary>
    /// Base class for Cmdlets that run asynchronously.
    /// </summary>
    /// <remarks>
    /// Inherit from this class if your Cmdlet needs to use <c>async</c> / <c>await</c> functionality.
    /// </remarks>
    public abstract class AsyncCmdlet
        : PSCmdlet, IDisposable{
        /// <summary>
        /// The source for cancellation tokens that can be used to cancel the operation.
        /// </summary>
        private readonly CancellationTokenSource _cancellationSource = new CancellationTokenSource();
 
        /// <summary>
        /// Dispose of resources being used by the Cmdlet.
        /// </summary>
        public void Dispose(){
            Dispose(true);
            GC.SuppressFinalize(this);
        }
 
        /// <summary>
        /// Finaliser for <see cref="AsyncCmdlet" />.
        /// </summary>
        ~AsyncCmdlet(){
            Dispose(false);
        }
 
        /// <summary>
        /// Dispose of resources being used by the Cmdlet.
        /// </summary>
        /// <param name="disposing">
        /// Explicit disposal?
        /// </param>
        protected virtual void Dispose(bool disposing){
            if (disposing)
                _cancellationSource.Dispose();
        }
 
        /// <summary>
        /// Asynchronously perform Cmdlet pre-processing.
        /// </summary>
        /// <returns>
        /// A <see cref="Task" /> representing the asynchronous operation.
        /// </returns>
        protected virtual Task BeginProcessingAsync(){
            return BeginProcessingAsync(_cancellationSource.Token);
        }
 
        /// <summary>
        /// Asynchronously perform Cmdlet pre-processing.
        /// </summary>
        /// <param name="cancellationToken">
        /// A <see cref="CancellationToken" /> that can be used to cancel the asynchronous operation.
        /// </param>
        /// <returns>
        /// A <see cref="Task" /> representing the asynchronous operation.
        /// </returns>
        protected virtual Task BeginProcessingAsync(CancellationToken cancellationToken){
            return Task.CompletedTask;
        }
 
        /// <summary>
        /// Asynchronously perform Cmdlet processing.
        /// </summary>
        /// <returns>
        /// A <see cref="Task" /> representing the asynchronous operation.
        /// </returns>
        protected virtual Task ProcessRecordAsync(){
            return ProcessRecordAsync(_cancellationSource.Token);
        }
 
        /// <summary>
        /// Asynchronously perform Cmdlet processing.
        /// </summary>
        /// <param name="cancellationToken">
        /// A <see cref="CancellationToken" /> that can be used to cancel the asynchronous operation.
        /// </param>
        /// <returns>
        /// A <see cref="Task" /> representing the asynchronous operation.
        /// </returns>
        protected virtual Task ProcessRecordAsync(CancellationToken cancellationToken){
            return Task.CompletedTask;
        }
 
        /// <summary>
        /// Asynchronously perform Cmdlet post-processing.
        /// </summary>
        /// <returns>
        /// A <see cref="Task" /> representing the asynchronous operation.
        /// </returns>
        protected virtual Task EndProcessingAsync(){
            return EndProcessingAsync(_cancellationSource.Token);
        }
 
        /// <summary>
        /// Asynchronously perform Cmdlet post-processing.
        /// </summary>
        /// <param name="cancellationToken">
        /// A <see cref="CancellationToken" /> that can be used to cancel the asynchronous operation.
        /// </param>
        /// <returns>
        /// A <see cref="Task" /> representing the asynchronous operation.
        /// </returns>
        protected virtual Task EndProcessingAsync(CancellationToken cancellationToken){
            return Task.CompletedTask;
        }
 
        /// <summary>
        /// Perform Cmdlet pre-processing.
        /// </summary>
        protected sealed override void BeginProcessing(){
            ThreadAffinitiveSynchronizationContext.RunSynchronized(BeginProcessingAsync);
        }
 
        /// <summary>
        /// Perform Cmdlet processing.
        /// </summary>
        protected sealed override void ProcessRecord(){
            ThreadAffinitiveSynchronizationContext.RunSynchronized(ProcessRecordAsync);
        }
 
        /// <summary>
        /// Perform Cmdlet post-processing.
        /// </summary>
        protected sealed override void EndProcessing(){
            ThreadAffinitiveSynchronizationContext.RunSynchronized(EndProcessingAsync);
        }
 
        /// <summary>
        /// Interrupt Cmdlet processing (if possible).
        /// </summary>
        protected sealed override void StopProcessing(){
            _cancellationSource.Cancel();
 
            base.StopProcessing();
        }
 
        /// <summary>
        /// Write a progress record to the output stream, and as a verbose message.
        /// </summary>
        /// <param name="progressRecord">
        /// The progress record to write.
        /// </param>
        protected void WriteVerboseProgress(ProgressRecord progressRecord){
            if (progressRecord == null)
                throw new ArgumentNullException(nameof(progressRecord));
 
            WriteProgress(progressRecord);
            WriteVerbose(progressRecord.StatusDescription);
        }
 
        /// <summary>
        /// Write a progress record to the output stream, and as a verbose message.
        /// </summary>
        /// <param name="progressRecord">
        /// The progress record to write.
        /// </param>
        /// <param name="messageOrFormat">
        /// The message or message-format specifier.
        /// </param>
        /// <param name="formatArguments">
        /// Optional format arguments.
        /// </param>
        protected void WriteVerboseProgress(ProgressRecord progressRecord, string messageOrFormat,
            params object[] formatArguments){
            if (progressRecord == null)
                throw new ArgumentNullException(nameof(progressRecord));
 
            if (string.IsNullOrWhiteSpace(messageOrFormat))
                throw new ArgumentException(
                    "Argument cannot be null, empty, or composed entirely of whitespace: 'messageOrFormat'.",
                    nameof(messageOrFormat));
 
            if (formatArguments == null)
                throw new ArgumentNullException(nameof(formatArguments));
 
            progressRecord.StatusDescription = string.Format(messageOrFormat, formatArguments);
            WriteVerboseProgress(progressRecord);
        }
 
        /// <summary>
        /// Write a completed progress record to the output stream.
        /// </summary>
        /// <param name="progressRecord">
        /// The progress record to complete.
        /// </param>
        /// <param name="completionMessageOrFormat">
        /// The completion message or message-format specifier.
        /// </param>
        /// <param name="formatArguments">
        /// Optional format arguments.
        /// </param>
        protected void WriteProgressCompletion(ProgressRecord progressRecord, string completionMessageOrFormat,
            params object[] formatArguments){
            if (progressRecord == null)
                throw new ArgumentNullException(nameof(progressRecord));
 
            if (string.IsNullOrWhiteSpace(completionMessageOrFormat))
                throw new ArgumentException(
                    "Argument cannot be null, empty, or composed entirely of whitespace: 'completionMessageOrFormat'.",
                    nameof(completionMessageOrFormat));
 
            if (formatArguments == null)
                throw new ArgumentNullException(nameof(formatArguments));
 
            progressRecord.StatusDescription = string.Format(completionMessageOrFormat, formatArguments);
            progressRecord.PercentComplete = 100;
            progressRecord.RecordType = ProgressRecordType.Completed;
            WriteProgress(progressRecord);
            WriteVerbose(progressRecord.StatusDescription);
        }
    }
 
    /// <summary>
    /// A synchronisation context that runs all calls scheduled on it (via <see cref="SynchronizationContext.Post" />) on a
    /// single thread.
    /// </summary>
    /// <remarks>
    /// With thanks to Stephen Toub.
    /// </remarks>
    public sealed class ThreadAffinitiveSynchronizationContext
        : SynchronizationContext, IDisposable{
        /// <summary>
        /// A blocking collection (effectively a queue) of work items to execute, consisting of callback delegates and their
        /// callback state (if any).
        /// </summary>
        private BlockingCollection<KeyValuePair<SendOrPostCallback, object>> _workItemQueue =
            new BlockingCollection<KeyValuePair<SendOrPostCallback, object>>();
 
        /// <summary>
        /// Create a new thread-affinitive synchronisation context.
        /// </summary>
        private ThreadAffinitiveSynchronizationContext(){
        }
 
        /// <summary>
        /// Dispose of resources being used by the synchronisation context.
        /// </summary>
        void IDisposable.Dispose(){
            if (_workItemQueue != null){
                _workItemQueue.Dispose();
                _workItemQueue = null;
            }
        }
 
        /// <summary>
        /// Check if the synchronisation context has been disposed.
        /// </summary>
        private void CheckDisposed(){
            if (_workItemQueue == null)
                throw new ObjectDisposedException(GetType().Name);
        }
 
        /// <summary>
        /// Run the message pump for the callback queue on the current thread.
        /// </summary>
        private void RunMessagePump(){
            CheckDisposed();
 
            while (_workItemQueue.TryTake(out var workItem, Timeout.InfiniteTimeSpan)){
                workItem.Key(workItem.Value);
 
                // Has the synchronisation context been disposed?
                if (_workItemQueue == null)
                    break;
            }
        }
 
        /// <summary>
        /// Terminate the message pump once all callbacks have completed.
        /// </summary>
        private void TerminateMessagePump(){
            CheckDisposed();
 
            _workItemQueue.CompleteAdding();
        }
 
        /// <summary>
        /// Dispatch an asynchronous message to the synchronization context.
        /// </summary>
        /// <param name="callback">
        /// The <see cref="SendOrPostCallback" /> delegate to call in the synchronisation context.
        /// </param>
        /// <param name="callbackState">
        /// Optional state data passed to the callback.
        /// </param>
        /// <exception cref="InvalidOperationException">
        /// The message pump has already been started, and then terminated by calling <see cref="TerminateMessagePump" />.
        /// </exception>
        public override void Post(SendOrPostCallback callback, object callbackState){
            if (callback == null)
                throw new ArgumentNullException(nameof(callback));
 
            CheckDisposed();
 
            try{
                _workItemQueue.Add(
                    new KeyValuePair<SendOrPostCallback, object>(
                        callback,
                        callbackState
                    )
                );
            }
            catch (InvalidOperationException eMessagePumpAlreadyTerminated){
                throw new InvalidOperationException(
                    "Cannot enqueue the specified callback because the synchronisation context's message pump has already been terminated.",
                    eMessagePumpAlreadyTerminated
                );
            }
        }
 
        /// <summary>
        /// Run an asynchronous operation using the current thread as its synchronisation context.
        /// </summary>
        /// <param name="asyncOperation">
        /// A <see cref="Func{TResult}" /> delegate representing the asynchronous operation to run.
        /// </param>
        public static void RunSynchronized(Func<Task> asyncOperation){
            if (asyncOperation == null)
                throw new ArgumentNullException(nameof(asyncOperation));
 
            var savedContext = Current;
            try{
                using (var synchronizationContext = new ThreadAffinitiveSynchronizationContext()){
                    SetSynchronizationContext(synchronizationContext);
 
                    var rootOperationTask = asyncOperation();
                    if (rootOperationTask == null)
                        throw new InvalidOperationException("The asynchronous operation delegate cannot return null.");
 
                    rootOperationTask.ContinueWith(
                        operationTask =>
                            // ReSharper disable once AccessToDisposedClosure
                            synchronizationContext.TerminateMessagePump(),
                        TaskScheduler.Default
                    );
 
                    synchronizationContext.RunMessagePump();
 
                    try{
                        rootOperationTask
                            .GetAwaiter()
                            .GetResult();
                    }
                    catch (AggregateException eWaitForTask
                    ) // The TPL will almost always wrap an AggregateException around any exception thrown by the async operation.
                    {
                        // Is this just a wrapped exception?
                        var flattenedAggregate = eWaitForTask.Flatten();
                        if (flattenedAggregate.InnerExceptions.Count != 1)
                            throw; // Nope, genuine aggregate.
 
                        // Yep, so rethrow (preserving original stack-trace).
                        ExceptionDispatchInfo.Capture(flattenedAggregate.InnerExceptions[0]).Throw();
                    }
                }
            }
            finally{
                SetSynchronizationContext(savedContext);
            }
        }
 
        /// <summary>
        /// Run an asynchronous operation using the current thread as its synchronisation context.
        /// </summary>
        /// <typeparam name="TResult">
        /// The operation result type.
        /// </typeparam>
        /// <param name="asyncOperation">
        /// A <see cref="Func{TResult}" /> delegate representing the asynchronous operation to run.
        /// </param>
        /// <returns>
        /// The operation result.
        /// </returns>
        public static TResult RunSynchronized<TResult>(Func<Task<TResult>> asyncOperation){
            if (asyncOperation == null)
                throw new ArgumentNullException(nameof(asyncOperation));
 
            var savedContext = Current;
            try{
                using (var synchronizationContext = new ThreadAffinitiveSynchronizationContext()){
                    SetSynchronizationContext(synchronizationContext);
 
                    var rootOperationTask = asyncOperation();
                    if (rootOperationTask == null)
                        throw new InvalidOperationException("The asynchronous operation delegate cannot return null.");
 
                    rootOperationTask.ContinueWith(
                        operationTask =>
                            // ReSharper disable once AccessToDisposedClosure
                            synchronizationContext.TerminateMessagePump(),
                        TaskScheduler.Default
                    );
 
                    synchronizationContext.RunMessagePump();
 
                    try{
                        return
                            rootOperationTask
                                .GetAwaiter()
                                .GetResult();
                    }
                    catch (AggregateException eWaitForTask
                    ) // The TPL will almost always wrap an AggregateException around any exception thrown by the async operation.
                    {
                        // Is this just a wrapped exception?
                        var flattenedAggregate = eWaitForTask.Flatten();
                        if (flattenedAggregate.InnerExceptions.Count != 1)
                            throw; // Nope, genuine aggregate.
 
                        // Yep, so rethrow (preserving original stack-trace).
                        ExceptionDispatchInfo
                            .Capture(
                                flattenedAggregate
                                    .InnerExceptions[0]
                            )
                            .Throw();
 
                        throw; // Never reached.
                    }
                }
            }
            finally{
                SetSynchronizationContext(savedContext);
            }
        }
    }
}