cbsch-pslib

1.0.14

module/src/AsyncSampleCmdlet.cs

using System;
using System.Collections.Concurrent;
using System.Management.Automation;
using System.Threading;
using System.Threading.Tasks;
 
namespace Cbsch.PSlib {
    /// <summary>
    /// Base class for async-enabled cmdlets
    /// </summary>
    public abstract class AsyncSampleCmdlet : PSCmdlet {
        protected int BoundedCapacity { get; set; }
 
        protected AsyncSampleCmdlet(int boundedCapacity = 50) {
            this.BoundedCapacity = Math.Max(1, boundedCapacity);
        }
 
        #region sealed overrides
        protected sealed override void BeginProcessing() {
            AsyncCmdletSynchronizationContext.Async(BeginProcessingAsync, BoundedCapacity);
        }
 
        protected sealed override void ProcessRecord() {
            AsyncCmdletSynchronizationContext.Async(ProcessRecordAsync, BoundedCapacity);
        }
 
        protected sealed override void EndProcessing() {
            AsyncCmdletSynchronizationContext.Async(EndProcessingAsync, BoundedCapacity);
        }
 
        protected sealed override void StopProcessing() {
            AsyncCmdletSynchronizationContext.Async(StopProcessingAsync, BoundedCapacity);
        }
 
        #endregion sealed overrides
 
        #region intercepted methods
        public new void WriteDebug(string text) {
            AsyncCmdletSynchronizationContext.PostItem(new MarshalItemAction<string>(base.WriteDebug, text));
        }
 
        public new void WriteError(ErrorRecord errorRecord) {
            AsyncCmdletSynchronizationContext.PostItem(new MarshalItemAction<ErrorRecord>(base.WriteError, errorRecord));
        }
 
        public new void WriteObject(object sendToPipeline) {
            AsyncCmdletSynchronizationContext.PostItem(new MarshalItemAction<object>(base.WriteObject, sendToPipeline));
        }
 
        public new void WriteObject(object sendToPipeline, bool enumerateCollection) {
            AsyncCmdletSynchronizationContext.PostItem(new MarshalItemAction<object, bool>(base.WriteObject, sendToPipeline, enumerateCollection));
        }
 
        public new void WriteProgress(ProgressRecord progressRecord) {
            AsyncCmdletSynchronizationContext.PostItem(new MarshalItemAction<ProgressRecord>(base.WriteProgress, progressRecord));
        }
 
        public new void WriteVerbose(string text) {
            var workItem = new MarshalItemAction<string>(base.WriteVerbose, text);
            AsyncCmdletSynchronizationContext.PostItem(workItem);
        }
 
        public new void WriteWarning(string text) {
            AsyncCmdletSynchronizationContext.PostItem(new MarshalItemAction<string>(base.WriteWarning, text));
        }
 
        public new void WriteCommandDetail(string text) {
            AsyncCmdletSynchronizationContext.PostItem(new MarshalItemAction<string>(base.WriteCommandDetail, text));
        }
 
        public new bool ShouldProcess(string target) {
            var workItem = new MarshalItemFunc<string, bool>(base.ShouldProcess, target);
            AsyncCmdletSynchronizationContext.PostItem(workItem);
            return workItem.WaitForResult();
        }
 
        public new bool ShouldProcess(string target, string action) {
            var workItem = new MarshalItemFunc<string, string, bool>(base.ShouldProcess, target, action);
            AsyncCmdletSynchronizationContext.PostItem(workItem);
            return workItem.WaitForResult();
        }
 
        public new bool ShouldProcess(string verboseDescription,
            string verboseWarning, string caption) {
            var workItem = new MarshalItemFunc<string, string, string, bool>(base.ShouldProcess, verboseDescription,
                verboseWarning, caption);
            AsyncCmdletSynchronizationContext.PostItem(workItem);
            return workItem.WaitForResult();
        }
 
        public new bool ShouldProcess(string verboseDescription, string verboseWarning,
            string caption, out ShouldProcessReason shouldProcessReason) {
            var workItem = new MarshalItemFuncOut<string, string, string, bool, ShouldProcessReason>(
                base.ShouldProcess, verboseDescription, verboseWarning, caption);
            AsyncCmdletSynchronizationContext.PostItem(workItem);
            return workItem.WaitForResult(out shouldProcessReason);
        }
 
        public new bool ShouldContinue(string query, string caption) {
            var workItem = new MarshalItemFunc<string, string, bool>(base.ShouldContinue, query, caption);
            AsyncCmdletSynchronizationContext.PostItem(workItem);
            return workItem.WaitForResult();
        }
 
        public new bool ShouldContinue(string query, string caption, ref bool yesToAll,
        ref bool noToAll) {
            var workItem = new MarshalItemFuncRef<string, string, bool, bool, bool>(
                base.ShouldContinue, query, caption, yesToAll, noToAll);
            AsyncCmdletSynchronizationContext.PostItem(workItem);
            return workItem.WaitForResult(ref yesToAll, ref noToAll);
        }
 
        public new bool TransactionAvailable() {
            var workItem = new MarshalItemFunc<bool>(base.TransactionAvailable);
            AsyncCmdletSynchronizationContext.PostItem(workItem);
            return workItem.WaitForResult();
        }
 
        public new void ThrowTerminatingError(ErrorRecord errorRecord) {
            AsyncCmdletSynchronizationContext.PostItem(new MarshalItemAction<ErrorRecord>(base.ThrowTerminatingError, errorRecord));
        }
 
        #endregion
 
        #region async processing methods
 
        protected virtual Task BeginProcessingAsync() {
            return Task.FromResult(0);
        }
 
 
        protected virtual Task EndProcessingAsync() {
            return Task.FromResult(0);
        }
 
 
        protected virtual Task ProcessRecordAsync() {
            return Task.FromResult(0);
        }
 
 
        protected virtual Task StopProcessingAsync() {
            return Task.FromResult(0);
        }
 
        #endregion async processing methods
 
        private class AsyncCmdletSynchronizationContext : SynchronizationContext, IDisposable {
            private BlockingCollection<MarshalItem> workItems;
            private static AsyncCmdletSynchronizationContext currentAsyncCmdletContext;
 
            private AsyncCmdletSynchronizationContext(int boundedCapacity) {
                this.workItems = new BlockingCollection<MarshalItem>(boundedCapacity);
            }
 
            public static void Async(Func<Task> handler, int boundedCapacity) {
                var previousContext = SynchronizationContext.Current;
 
                try {
                    using (var synchronizationContext = new AsyncCmdletSynchronizationContext(boundedCapacity)) {
                        SetSynchronizationContext(synchronizationContext);
                        currentAsyncCmdletContext = synchronizationContext;
 
                        var task = handler();
                        if (task == null) {
                            return;
                        }
 
                        var waitable = task.ContinueWith(t => synchronizationContext.Complete(),
                            scheduler: TaskScheduler.Default);
 
                        synchronizationContext.ProcessQueue();
 
                        waitable.GetAwaiter().GetResult();
                    }
                } finally {
                    SetSynchronizationContext(previousContext);
                    currentAsyncCmdletContext = previousContext as AsyncCmdletSynchronizationContext;
                }
            }
 
            internal static void PostItem(MarshalItem item) {
                currentAsyncCmdletContext.Post(item);
            }
 
            public void Dispose() {
                if (this.workItems != null) {
                    this.workItems.Dispose();
                    this.workItems = null;
                }
            }
 
            private void EnsureNotDisposed() {
                if (this.workItems == null) {
                    throw new ObjectDisposedException(nameof(AsyncCmdletSynchronizationContext));
                }
            }
 
            private void Complete() {
                EnsureNotDisposed();
 
                this.workItems.CompleteAdding();
            }
 
            private void ProcessQueue() {
                MarshalItem workItem;
                while (this.workItems.TryTake(out workItem, Timeout.Infinite)) {
                    workItem.Invoke();
                }
            }
 
            public override void Post(SendOrPostCallback callback, object state) {
                if (callback == null) {
                    throw new ArgumentNullException(nameof(callback));
                }
 
                Post(new MarshalItemAction<object>(s => callback(s), state));
            }
 
            private void Post(MarshalItem item) {
                EnsureNotDisposed();
 
                this.workItems.Add(item);
            }
        }
 
        #region items
        internal abstract class MarshalItem {
            internal abstract void Invoke();
        }
 
        abstract class MarshalItemFuncBase<TRet> : MarshalItem {
            private TRet retVal;
            private readonly Task<TRet> retValTask;
 
            protected MarshalItemFuncBase() {
                this.retValTask = new Task<TRet>(() => this.retVal);
            }
 
            internal sealed override void Invoke() {
                this.retVal = this.InvokeFunc();
                this.retValTask.Start();
            }
 
            internal TRet WaitForResult() {
                this.retValTask.Wait();
                return this.retValTask.Result;
            }
 
            internal abstract TRet InvokeFunc();
        }
        class MarshalItemAction<T> : MarshalItem {
            private readonly Action<T> action;
            private readonly T arg1;
 
            internal MarshalItemAction(Action<T> action, T arg1) {
                this.action = action;
                this.arg1 = arg1;
            }
 
            internal override void Invoke() {
                this.action(this.arg1);
            }
        }
        class MarshalItemAction<T1, T2> : MarshalItem {
            private readonly Action<T1, T2> action;
            private readonly T1 arg1;
            private readonly T2 arg2;
 
            internal MarshalItemAction(Action<T1, T2> action, T1 arg1, T2 arg2) {
                this.action = action;
                this.arg1 = arg1;
                this.arg2 = arg2;
            }
 
            internal override void Invoke() {
                this.action(this.arg1, this.arg2);
            }
        }
        class MarshalItemFunc<TRet> : MarshalItemFuncBase<TRet> {
            private readonly Func<TRet> func;
 
            internal MarshalItemFunc(Func<TRet> func) {
                this.func = func;
            }
 
            internal override TRet InvokeFunc() {
                return this.func();
            }
        }
        class MarshalItemFunc<T1, TRet> : MarshalItemFuncBase<TRet> {
            private readonly Func<T1, TRet> func;
            private readonly T1 arg1;
 
            internal MarshalItemFunc(Func<T1, TRet> func, T1 arg1) {
                this.func = func;
                this.arg1 = arg1;
            }
 
            internal override TRet InvokeFunc() {
                return this.func(this.arg1);
            }
        }
        class MarshalItemFunc<T1, T2, TRet> : MarshalItemFuncBase<TRet> {
            private readonly Func<T1, T2, TRet> func;
            private readonly T1 arg1;
            private readonly T2 arg2;
 
            internal MarshalItemFunc(Func<T1, T2, TRet> func, T1 arg1, T2 arg2) {
                this.func = func;
                this.arg1 = arg1;
                this.arg2 = arg2;
            }
 
            internal override TRet InvokeFunc() {
                return this.func(this.arg1, this.arg2);
            }
        }
        class MarshalItemFunc<T1, T2, T3, TRet> : MarshalItemFuncBase<TRet> {
            private readonly Func<T1, T2, T3, TRet> func;
            private readonly T1 arg1;
            private readonly T2 arg2;
            private readonly T3 arg3;
 
            internal MarshalItemFunc(Func<T1, T2, T3, TRet> func, T1 arg1, T2 arg2, T3 arg3) {
                this.func = func;
                this.arg1 = arg1;
                this.arg2 = arg2;
                this.arg3 = arg3;
            }
 
            internal override TRet InvokeFunc() {
                return this.func(this.arg1, this.arg2, this.arg3);
            }
        }
        class MarshalItemFuncOut<T1, T2, T3, TRet, TOut> : MarshalItem {
            private readonly FuncOut func;
            private readonly T1 arg1;
            private readonly T2 arg2;
            private readonly T3 arg3;
 
            internal delegate TRet FuncOut(T1 t1, T2 t2, T3 t3, out TOut tout);
 
            private TRet retVal;
            private TOut outVal;
            private readonly Task<TRet> retValTask;
 
            internal MarshalItemFuncOut(FuncOut func, T1 arg1, T2 arg2, T3 arg3) {
                this.func = func;
                this.arg1 = arg1;
                this.arg2 = arg2;
                this.arg3 = arg3;
                this.retValTask = new Task<TRet>(() => this.retVal);
            }
 
            internal override void Invoke() {
                this.retVal = this.func(this.arg1, this.arg2, this.arg3, out this.outVal);
                this.retValTask.Start();
            }
 
            internal TRet WaitForResult(out TOut val) {
                this.retValTask.Wait();
                val = this.outVal;
                return this.retValTask.Result;
            }
        }
        class MarshalItemFuncRef<T1, T2, TRet, TRef1, TRef2> : MarshalItem {
            internal delegate TRet FuncRef(T1 t1, T2 t2, ref TRef1 tref1, ref TRef2 tref2);
 
            private readonly Task<TRet> retValTask;
            private readonly FuncRef func;
            private readonly T1 arg1;
            private readonly T2 arg2;
            private TRef1 arg3;
            private TRef2 arg4;
            private TRet retVal;
 
            internal MarshalItemFuncRef(FuncRef func, T1 arg1, T2 arg2, TRef1 arg3, TRef2 arg4) {
                this.func = func;
                this.arg1 = arg1;
                this.arg2 = arg2;
                this.arg3 = arg3;
                this.arg4 = arg4;
                this.retValTask = new Task<TRet>(() => this.retVal);
            }
 
            internal override void Invoke() {
                this.retVal = this.func(this.arg1, this.arg2, ref this.arg3, ref this.arg4);
                this.retValTask.Start();
            }
 
            // ReSharper disable RedundantAssignment
            internal TRet WaitForResult(ref TRef1 ref1, ref TRef2 ref2) {
                this.retValTask.Wait();
                ref1 = this.arg3;
                ref2 = this.arg4;
                return this.retValTask.Result;
            }
            // ReSharper restore RedundantAssignment
        }
        #endregion items
    }
}