// 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(); } } } } } } }