// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
namespace Microsoft.VisualStudio.Threading
{
using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.CompilerServices;
using System.Runtime.ExceptionServices;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Win32;
using Microsoft.Win32.SafeHandles;
///
/// Extension methods and awaitables for .NET 4.5.
///
public static partial class AwaitExtensions
{
///
/// Gets an awaiter that schedules continuations on the specified scheduler.
///
/// The task scheduler used to execute continuations.
/// An awaitable.
public static TaskSchedulerAwaiter GetAwaiter(this TaskScheduler scheduler)
{
Requires.NotNull(scheduler, nameof(scheduler));
return new TaskSchedulerAwaiter(scheduler);
}
///
/// Gets an awaitable that schedules continuations on the specified scheduler.
///
/// The task scheduler used to execute continuations.
/// A value indicating whether the caller should yield even if
/// already executing on the desired task scheduler.
/// An awaitable.
public static TaskSchedulerAwaitable SwitchTo(this TaskScheduler scheduler, bool alwaysYield = false)
{
Requires.NotNull(scheduler, nameof(scheduler));
return new TaskSchedulerAwaitable(scheduler, alwaysYield);
}
///
/// Provides await functionality for ordinary s.
///
/// The handle to wait on.
/// The awaiter.
public static TaskAwaiter GetAwaiter(this WaitHandle handle)
{
Requires.NotNull(handle, nameof(handle));
Task task = handle.ToTask();
return task.GetAwaiter();
}
///
/// Returns a task that completes when the process exits and provides the exit code of that process.
///
/// The process to wait for exit.
///
/// A token whose cancellation will cause the returned Task to complete
/// before the process exits in a faulted state with an .
/// This token has no effect on the itself.
///
/// A task whose result is the of the .
public static async Task WaitForExitAsync(this Process process, CancellationToken cancellationToken = default(CancellationToken))
{
Requires.NotNull(process, nameof(process));
var tcs = new TaskCompletionSource();
EventHandler exitHandler = (s, e) =>
{
tcs.TrySetResult(process.ExitCode);
};
try
{
process.EnableRaisingEvents = true;
process.Exited += exitHandler;
if (process.HasExited)
{
// Allow for the race condition that the process has already exited.
tcs.TrySetResult(process.ExitCode);
}
using (cancellationToken.Register(() => tcs.TrySetCanceled(cancellationToken)))
{
return await tcs.Task.ConfigureAwait(false);
}
}
finally
{
process.Exited -= exitHandler;
}
}
///
/// Returns a Task that completes when the specified registry key changes.
///
/// The registry key to watch for changes.
/// true to watch the keys descendent keys as well; false to watch only this key without descendents.
/// Indicates the kinds of changes to watch for.
/// A token that may be canceled to release the resources from watching for changes and complete the returned Task as canceled.
///
/// A task that completes when the registry key changes, the handle is closed, or upon cancellation.
///
public static Task WaitForChangeAsync(this RegistryKey registryKey, bool watchSubtree = true, RegistryChangeNotificationFilters change = RegistryChangeNotificationFilters.Value | RegistryChangeNotificationFilters.Subkey, CancellationToken cancellationToken = default(CancellationToken))
{
Requires.NotNull(registryKey, nameof(registryKey));
return WaitForRegistryChangeAsync(registryKey.Handle, watchSubtree, change, cancellationToken);
}
///
/// Converts a to a .
///
/// The result of .
/// A value indicating whether the continuation should run on the captured , if any.
/// An awaitable.
public static ConfiguredTaskYieldAwaitable ConfigureAwait(this YieldAwaitable yieldAwaitable, bool continueOnCapturedContext)
{
return new ConfiguredTaskYieldAwaitable(continueOnCapturedContext);
}
///
/// Gets an awaitable that schedules the continuation with a preference to executing synchronously on the callstack that completed the ,
/// without regard to thread ID or any that may be applied when the continuation is scheduled or when the antecedent completes.
///
/// The task to await on.
/// An awaitable.
///
/// If there is not enough stack space remaining on the thread that is completing the ,
/// the continuation may be scheduled on the threadpool.
///
public static ExecuteContinuationSynchronouslyAwaitable ConfigureAwaitRunInline(this Task antecedent)
{
Requires.NotNull(antecedent, nameof(antecedent));
return new ExecuteContinuationSynchronouslyAwaitable(antecedent);
}
///
/// Gets an awaitable that schedules the continuation with a preference to executing synchronously on the callstack that completed the ,
/// without regard to thread ID or any that may be applied when the continuation is scheduled or when the antecedent completes.
///
/// The type of value returned by the awaited .
/// The task to await on.
/// An awaitable.
///
/// If there is not enough stack space remaining on the thread that is completing the ,
/// the continuation may be scheduled on the threadpool.
///
public static ExecuteContinuationSynchronouslyAwaitable ConfigureAwaitRunInline(this Task antecedent)
{
Requires.NotNull(antecedent, nameof(antecedent));
return new ExecuteContinuationSynchronouslyAwaitable(antecedent);
}
///
/// Returns an awaitable that will throw from the property of the task if it faults.
///
/// The task to track for completion.
///
/// An awaitable that may throw .
///
/// Awaiting a with its default only throws the first exception within .
/// When you do not want to lose the detail of other inner exceptions, use this extension method.
///
/// Thrown when faults.
public static AggregateExceptionAwaitable ConfigureAwaitForAggregateException(this Task task, bool continueOnCapturedContext = true) => new AggregateExceptionAwaitable(task, continueOnCapturedContext);
///
/// Returns a Task that completes when the specified registry key changes.
///
/// The handle to the open registry key to watch for changes.
/// true to watch the keys descendent keys as well; false to watch only this key without descendents.
/// Indicates the kinds of changes to watch for.
/// A token that may be canceled to release the resources from watching for changes and complete the returned Task as canceled.
///
/// A task that completes when the registry key changes, the handle is closed, or upon cancellation.
///
private static async Task WaitForRegistryChangeAsync(SafeRegistryHandle registryKeyHandle, bool watchSubtree, RegistryChangeNotificationFilters change, CancellationToken cancellationToken)
{
IDisposable? dedicatedThreadReleaser = null;
try
{
using (var evt = new ManualResetEvent(false))
{
static void DoNotify(SafeRegistryHandle registryKeyHandle, bool watchSubtree, RegistryChangeNotificationFilters change, WaitHandle evt)
{
int win32Error = NativeMethods.RegNotifyChangeKeyValue(
registryKeyHandle,
watchSubtree,
change,
evt.SafeWaitHandle,
true);
if (win32Error != 0)
{
throw new Win32Exception(win32Error);
}
}
if (LightUps.IsWindows8OrLater)
{
change |= NativeMethods.REG_NOTIFY_THREAD_AGNOSTIC;
DoNotify(registryKeyHandle, watchSubtree, change, evt);
}
else
{
// Engage our downlevel support by using a single, dedicated thread to guarantee
// that we request notification on a thread that will not be destroyed later.
// Although we *could* await this, we synchronously block because our caller expects
// subscription to have begun before we return: for the async part to simply be notification.
// This async method we're calling uses .ConfigureAwait(false) internally so this won't
// deadlock if we're called on a thread with a single-thread SynchronizationContext.
Action registerAction = () => DoNotify(registryKeyHandle, watchSubtree, change, evt);
dedicatedThreadReleaser = DownlevelRegistryWatcherSupport.ExecuteOnDedicatedThreadAsync(registerAction).GetAwaiter().GetResult();
}
await evt.ToTask(cancellationToken: cancellationToken).ConfigureAwait(false);
}
}
finally
{
dedicatedThreadReleaser?.Dispose();
}
}
///
/// The result of to prepare a to be awaited while throwing with all inner exceptions.
///
public readonly struct AggregateExceptionAwaitable
{
private readonly Task task;
private readonly bool continueOnCapturedContext;
///
/// Initializes a new instance of the struct.
///
public AggregateExceptionAwaitable(Task task, bool continueOnCapturedContext)
{
this.task = task;
this.continueOnCapturedContext = continueOnCapturedContext;
}
///
/// Gets an awaitable that schedules continuations on the specified scheduler.
///
public AggregateExceptionAwaiter GetAwaiter()
{
return new AggregateExceptionAwaiter(this.task, this.continueOnCapturedContext);
}
}
///
/// The result of to prepare a to be awaited while throwing with all inner exceptions.
///
public readonly struct AggregateExceptionAwaiter : ICriticalNotifyCompletion
{
private readonly Task task;
private readonly bool continueOnCapturedContext;
///
/// Initializes a new instance of the struct.
///
public AggregateExceptionAwaiter(Task task, bool continueOnCapturedContext)
{
this.task = task;
this.continueOnCapturedContext = continueOnCapturedContext;
}
///
public bool IsCompleted => this.Awaiter.IsCompleted;
private ConfiguredTaskAwaitable.ConfiguredTaskAwaiter Awaiter => this.task.ConfigureAwait(this.continueOnCapturedContext).GetAwaiter();
///
public void OnCompleted(Action continuation) => this.Awaiter.OnCompleted(continuation);
///
public void UnsafeOnCompleted(Action continuation) => this.Awaiter.UnsafeOnCompleted(continuation);
///
/// Thrown if the task was canceled.
/// Thrown if the task faulted.
public void GetResult()
{
if (this.task.Status == TaskStatus.Faulted && this.task.Exception is object)
{
ExceptionDispatchInfo.Capture(this.task.Exception).Throw();
}
this.Awaiter.GetResult();
}
}
///
/// An awaitable that executes continuations on the specified task scheduler.
///
public readonly struct TaskSchedulerAwaitable
{
///
/// The scheduler for continuations.
///
private readonly TaskScheduler taskScheduler;
///
/// A value indicating whether the awaitable will always call the caller to yield.
///
private readonly bool alwaysYield;
///
/// Initializes a new instance of the struct.
///
/// The task scheduler used to execute continuations.
/// A value indicating whether the caller should yield even if
/// already executing on the desired task scheduler.
public TaskSchedulerAwaitable(TaskScheduler taskScheduler, bool alwaysYield = false)
{
Requires.NotNull(taskScheduler, nameof(taskScheduler));
this.taskScheduler = taskScheduler;
this.alwaysYield = alwaysYield;
}
///
/// Gets an awaitable that schedules continuations on the specified scheduler.
///
public TaskSchedulerAwaiter GetAwaiter()
{
return new TaskSchedulerAwaiter(this.taskScheduler, this.alwaysYield);
}
}
///
/// An awaiter returned from .
///
public readonly struct TaskSchedulerAwaiter : ICriticalNotifyCompletion
{
///
/// The scheduler for continuations.
///
private readonly TaskScheduler scheduler;
///
/// A value indicating whether
/// should always return false.
///
private readonly bool alwaysYield;
///
/// Initializes a new instance of the struct.
///
/// The scheduler for continuations.
/// A value indicating whether the caller should yield even if
/// already executing on the desired task scheduler.
public TaskSchedulerAwaiter(TaskScheduler scheduler, bool alwaysYield = false)
{
this.scheduler = scheduler;
this.alwaysYield = alwaysYield;
}
///
/// Gets a value indicating whether no yield is necessary.
///
/// true if the caller is already running on that TaskScheduler.
public bool IsCompleted
{
get
{
if (this.alwaysYield)
{
return false;
}
// We special case the TaskScheduler.Default since that is semantically equivalent to being
// on a ThreadPool thread, and there are various ways to get on those threads.
// TaskScheduler.Current is never null. Even if no scheduler is really active and the current
// thread is not a threadpool thread, TaskScheduler.Current == TaskScheduler.Default, so we have
// to protect against that case too.
bool isThreadPoolThread = Thread.CurrentThread.IsThreadPoolThread;
return (this.scheduler == TaskScheduler.Default && isThreadPoolThread)
|| (this.scheduler == TaskScheduler.Current && TaskScheduler.Current != TaskScheduler.Default);
}
}
///
/// Schedules a continuation to execute using the specified task scheduler.
///
/// The delegate to invoke.
public void OnCompleted(Action continuation)
{
if (this.scheduler == TaskScheduler.Default)
{
ThreadPool.QueueUserWorkItem(state => ((Action)state!)(), continuation);
}
else
{
Task.Factory.StartNew(continuation, CancellationToken.None, TaskCreationOptions.None, this.scheduler);
}
}
///
/// Schedules a continuation to execute using the specified task scheduler
/// without capturing the ExecutionContext.
///
/// The action.
public void UnsafeOnCompleted(Action continuation)
{
if (this.scheduler == TaskScheduler.Default)
{
ThreadPool.UnsafeQueueUserWorkItem(state => ((Action)state!)(), continuation);
}
else
{
#if NETFRAMEWORK // Only bother suppressing flow on .NET Framework where the perf would improve from doing so.
if (ExecutionContext.IsFlowSuppressed())
{
Task.Factory.StartNew(continuation, CancellationToken.None, TaskCreationOptions.None, this.scheduler);
}
else
{
using (ExecutionContext.SuppressFlow())
{
Task.Factory.StartNew(continuation, CancellationToken.None, TaskCreationOptions.None, this.scheduler);
}
}
#else
Task.Factory.StartNew(continuation, CancellationToken.None, TaskCreationOptions.None, this.scheduler);
#endif
}
}
///
/// Does nothing.
///
public void GetResult()
{
}
}
///
/// An awaitable that will always lead the calling async method to yield,
/// then immediately resume, possibly on the original .
///
public readonly struct ConfiguredTaskYieldAwaitable
{
///
/// A value indicating whether the continuation should run on the captured , if any.
///
private readonly bool continueOnCapturedContext;
///
/// Initializes a new instance of the struct.
///
/// A value indicating whether the continuation should run on the captured , if any.
public ConfiguredTaskYieldAwaitable(bool continueOnCapturedContext)
{
this.continueOnCapturedContext = continueOnCapturedContext;
}
///
/// Gets the awaiter.
///
/// The awaiter.
public ConfiguredTaskYieldAwaiter GetAwaiter() => new ConfiguredTaskYieldAwaiter(this.continueOnCapturedContext);
}
///
/// An awaiter that will always lead the calling async method to yield,
/// then immediately resume, possibly on the original .
///
public readonly struct ConfiguredTaskYieldAwaiter : ICriticalNotifyCompletion
{
///
/// A value indicating whether the continuation should run on the captured , if any.
///
private readonly bool continueOnCapturedContext;
///
/// Initializes a new instance of the struct.
///
/// A value indicating whether the continuation should run on the captured , if any.
public ConfiguredTaskYieldAwaiter(bool continueOnCapturedContext)
{
this.continueOnCapturedContext = continueOnCapturedContext;
}
///
/// Gets a value indicating whether the caller should yield.
///
/// Always false.
public bool IsCompleted => false;
///
/// Schedules a continuation to execute immediately (but not synchronously).
///
/// The delegate to invoke.
public void OnCompleted(Action continuation)
{
if (this.continueOnCapturedContext)
{
Task.Yield().GetAwaiter().OnCompleted(continuation);
}
else
{
ThreadPool.QueueUserWorkItem(state => ((Action)state!)(), continuation);
}
}
///
/// Schedules a delegate for execution at the conclusion of a task's execution
/// without capturing the ExecutionContext.
///
/// The action.
public void UnsafeOnCompleted(Action continuation)
{
if (this.continueOnCapturedContext)
{
Task.Yield().GetAwaiter().UnsafeOnCompleted(continuation);
}
else
{
ThreadPool.UnsafeQueueUserWorkItem(state => ((Action)state!)(), continuation);
}
}
///
/// Does nothing.
///
public void GetResult()
{
}
}
///
/// A Task awaitable that has affinity to executing callbacks synchronously on the completing callstack.
///
public readonly struct ExecuteContinuationSynchronouslyAwaitable
{
///
/// The task whose completion will execute the continuation.
///
private readonly Task antecedent;
///
/// Initializes a new instance of the struct.
///
/// The task whose completion will execute the continuation.
public ExecuteContinuationSynchronouslyAwaitable(Task antecedent)
{
Requires.NotNull(antecedent, nameof(antecedent));
this.antecedent = antecedent;
}
///
/// Gets the awaiter.
///
/// The awaiter.
public ExecuteContinuationSynchronouslyAwaiter GetAwaiter() => new ExecuteContinuationSynchronouslyAwaiter(this.antecedent);
}
///
/// A Task awaiter that has affinity to executing callbacks synchronously on the completing callstack.
///
public readonly struct ExecuteContinuationSynchronouslyAwaiter : INotifyCompletion
{
///
/// The task whose completion will execute the continuation.
///
private readonly Task antecedent;
///
/// Initializes a new instance of the struct.
///
/// The task whose completion will execute the continuation.
public ExecuteContinuationSynchronouslyAwaiter(Task antecedent)
{
Requires.NotNull(antecedent, nameof(antecedent));
this.antecedent = antecedent;
}
///
/// Gets a value indicating whether the antedent has already completed.
///
public bool IsCompleted => this.antecedent.IsCompleted;
///
/// Rethrows any exception thrown by the antecedent.
///
public void GetResult() => this.antecedent.GetAwaiter().GetResult();
///
/// Schedules a callback to run when the antecedent task completes.
///
/// The callback to invoke.
public void OnCompleted(Action continuation)
{
Requires.NotNull(continuation, nameof(continuation));
this.antecedent.ContinueWith(
(_, s) => ((Action)s!)(),
continuation,
CancellationToken.None,
TaskContinuationOptions.ExecuteSynchronously,
TaskScheduler.Default);
}
}
///
/// A Task awaitable that has affinity to executing callbacks synchronously on the completing callstack.
///
/// The type of value returned by the awaited .
public readonly struct ExecuteContinuationSynchronouslyAwaitable
{
///
/// The task whose completion will execute the continuation.
///
private readonly Task antecedent;
///
/// Initializes a new instance of the struct.
///
/// The task whose completion will execute the continuation.
public ExecuteContinuationSynchronouslyAwaitable(Task antecedent)
{
Requires.NotNull(antecedent, nameof(antecedent));
this.antecedent = antecedent;
}
///
/// Gets the awaiter.
///
/// The awaiter.
public ExecuteContinuationSynchronouslyAwaiter GetAwaiter() => new ExecuteContinuationSynchronouslyAwaiter(this.antecedent);
}
///
/// A Task awaiter that has affinity to executing callbacks synchronously on the completing callstack.
///
/// The type of value returned by the awaited .
public readonly struct ExecuteContinuationSynchronouslyAwaiter : INotifyCompletion
{
///
/// The task whose completion will execute the continuation.
///
private readonly Task antecedent;
///
/// Initializes a new instance of the struct.
///
/// The task whose completion will execute the continuation.
public ExecuteContinuationSynchronouslyAwaiter(Task antecedent)
{
Requires.NotNull(antecedent, nameof(antecedent));
this.antecedent = antecedent;
}
///
/// Gets a value indicating whether the antedent has already completed.
///
public bool IsCompleted => this.antecedent.IsCompleted;
///
/// Rethrows any exception thrown by the antecedent.
///
public T GetResult() => this.antecedent.GetAwaiter().GetResult();
///
/// Schedules a callback to run when the antecedent task completes.
///
/// The callback to invoke.
public void OnCompleted(Action continuation)
{
Requires.NotNull(continuation, nameof(continuation));
this.antecedent.ContinueWith(
(_, s) => ((Action)s!)(),
continuation,
CancellationToken.None,
TaskContinuationOptions.ExecuteSynchronously,
TaskScheduler.Default);
}
}
///
/// Provides a dedicated thread for requesting registry change notifications.
///
///
/// For versions of Windows prior to Windows 8, requesting registry change notifications
/// required that the thread that made the request remain alive or else the watcher would
/// simply signal the event and stop watching for changes.
/// This class provides a single, dedicated thread for requesting such notifications
/// so that they don't get canceled when a thread happens to exit.
/// The dedicated thread is released when no one is watching the registry any more.
///
private static class DownlevelRegistryWatcherSupport
{
///
/// The size of the stack allocated for a thread that expects to stay within just a few methods in depth.
///
///
/// The default stack size for a thread is 1MB.
///
private const int SmallThreadStackSize = 100 * 1024;
///
/// The object to lock when accessing any fields.
/// This is also the object that is waited on by the dedicated thread,
/// and may be pulsed by others to wake the dedicated thread to do some work.
///
private static readonly object SyncObject = new object();
///
/// A queue of actions the dedicated thread should take.
///
private static readonly Queue>> PendingWork = new Queue>>();
///
/// The number of callers that still have an interest in the survival of the dedicated thread.
/// The dedicated thread will exit when this value reaches 0.
///
private static int keepAliveCount;
///
/// The thread that should stay alive and be dequeuing .
///
private static Thread? liveThread;
///
/// Executes some action on a long-lived thread.
///
/// The delegate to execute.
///
/// A task that either faults with the exception thrown by
/// or completes after successfully executing the delegate
/// with a result that should be disposed when it is safe to terminate the long-lived thread.
///
///
/// This thread never posts to , so it is safe
/// to call this method and synchronously block on its result.
///
internal static async Task ExecuteOnDedicatedThreadAsync(Action action)
{
Requires.NotNull(action, nameof(action));
var tcs = new TaskCompletionSource();
bool keepAliveCountIncremented = false;
try
{
lock (SyncObject)
{
PendingWork.Enqueue(Tuple.Create(action, tcs));
try
{
// This block intentionally left blank.
}
finally
{
// We make these two assignments within a finally block
// to guard against an untimely ThreadAbortException causing
// us to execute just one of them.
keepAliveCountIncremented = true;
++keepAliveCount;
}
if (keepAliveCount == 1)
{
Assumes.Null(liveThread);
liveThread = new Thread(Worker, SmallThreadStackSize)
{
IsBackground = true,
Name = "Registry watcher",
};
liveThread.Start();
}
else
{
// There *could* temporarily be multiple threads in some race conditions.
// Pulse all of them so that the live one is sure to get the message.
Monitor.PulseAll(SyncObject);
}
}
await tcs.Task.ConfigureAwait(false);
return new ThreadHandleRelease();
}
catch
{
if (keepAliveCountIncremented)
{
// Our caller will never have a chance to release their claim on the dedicated thread,
// so do it for them.
ReleaseRefOnDedicatedThread();
}
throw;
}
}
///
/// Decrements the count of interested parties in the live thread,
/// and helps it to terminate if necessary.
///
private static void ReleaseRefOnDedicatedThread()
{
lock (SyncObject)
{
if (--keepAliveCount == 0)
{
liveThread = null;
// Wake up any obsolete thread(s) so they can go to exit.
Monitor.PulseAll(SyncObject);
}
}
}
///
/// Executes thread-affinitized work from a queue until both the queue is empty
/// and any lingering interest in the survival of the dedicated thread has been released.
///
///
/// This method serves as the for our dedicated thread.
///
private static void Worker()
{
while (true)
{
Tuple>? work = null;
lock (SyncObject)
{
if (Thread.CurrentThread != liveThread)
{
// Regardless of our PendingWork and keepAliveCount,
// it isn't meant for this thread any more.
// This happens when keepAliveCount (at least temporarily)
// hits 0, so this thread must be assumed to be on its exit path,
// and another thread will be spawned to process new requests.
Assumes.True(liveThread is object || (keepAliveCount == 0 && PendingWork.Count == 0));
return;
}
if (PendingWork.Count > 0)
{
work = PendingWork.Dequeue();
}
else if (keepAliveCount == 0)
{
// No work, and no reason to stay alive. Exit the thread.
return;
}
else
{
// Sleep until another thread wants to wake us up with a Pulse.
Monitor.Wait(SyncObject);
}
}
if (work is object)
{
try
{
work.Item1();
work.Item2.SetResult(EmptyStruct.Instance);
}
catch (Exception ex)
{
work.Item2.SetException(ex);
}
}
}
}
///
/// Decrements the dedicated thread use counter by at most one upon disposal.
///
private class ThreadHandleRelease : IDisposable
{
///
/// A value indicating whether this instance has already been disposed.
///
private bool disposed;
///
/// Release the keep alive count reserved by this instance.
///
public void Dispose()
{
lock (SyncObject)
{
if (!this.disposed)
{
this.disposed = true;
ReleaseRefOnDedicatedThread();
}
}
}
}
}
}
}