Skip to content

Instantly share code, notes, and snippets.

@sunnyone
Created June 2, 2014 15:06
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save sunnyone/fd0b85c381e63861a198 to your computer and use it in GitHub Desktop.
Save sunnyone/fd0b85c381e63861a198 to your computer and use it in GitHub Desktop.
A base class that enables to use async/await (cf http://sunnyone41.blogspot.jp/2014/06/pscmdletasyncawait.html)
using System;
using System.Collections.Concurrent;
using System.Management.Automation;
using System.Threading;
using System.Threading.Tasks;
namespace AwaitablePSCmdlet
{
public class AwaitablePSCmdlet : PSCmdlet
{
private class PSSynchronizationContextItem
{
public SendOrPostCallback Callback { get; private set; }
public object State { get; private set; }
public PSSynchronizationContextItem(SendOrPostCallback callback, object state)
{
Callback = callback;
State = state;
}
}
private class PSSynchronizationContext : SynchronizationContext
{
private readonly BlockingCollection<PSSynchronizationContextItem> queue =
new BlockingCollection<PSSynchronizationContextItem>();
public override void Post(SendOrPostCallback d, object state)
{
queue.Add(new PSSynchronizationContextItem(d, state));
}
public void RunLoop()
{
PSSynchronizationContextItem item;
while (queue.TryTake(out item, Timeout.Infinite))
{
if (item == null)
{
break;
}
item.Callback(item.State);
}
}
public void Complete()
{
// Null means break
queue.Add(null);
}
}
private void doWithSyncCtx(Func<Task> func)
{
var prevCtx = SynchronizationContext.Current;
try
{
SynchronizationContext.SetSynchronizationContext(psSyncCtx);
var task = func();
task.ContinueWith(t => psSyncCtx.Complete(), TaskScheduler.Default);
psSyncCtx.RunLoop();
task.GetAwaiter().GetResult();
}
finally
{
SynchronizationContext.SetSynchronizationContext(prevCtx);
}
}
private PSSynchronizationContext psSyncCtx = new PSSynchronizationContext();
protected override void BeginProcessing()
{
doWithSyncCtx(BeginProcessingAsync);
}
protected override void ProcessRecord()
{
doWithSyncCtx(ProcessRecordAsync);
}
protected override void EndProcessing()
{
doWithSyncCtx(EndProcessingAsync);
}
protected virtual Task ProcessRecordAsync()
{
return Task.FromResult(true);
}
protected virtual Task BeginProcessingAsync()
{
return Task.FromResult(true);
}
protected virtual Task EndProcessingAsync()
{
return Task.FromResult(true);
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment