TaskScheduler.SwitchTo

_usage.md

public static async IAsyncEnumerable<string> EnumerateFilesAsync(
    string path,
    string searchPattern,
    SearchOption searchOption,
    [EnumeratorCancellation] CancellationToken cancellationToken)
{
    // Get off the UI thread before starting the file enumeration
    await TaskScheduler.Default.SwitchTo(cancellationToken);

    foreach (var file in Directory.EnumerateFiles(path, searchPattern, searchOption))
    {
        yield return file;

        // Get off the UI thread before continuing the file enumeration
        await TaskScheduler.Default.SwitchTo(cancellationToken);
    }
}

AsyncExtensions.cs

using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;

// Reference: https://github.com/microsoft/vs-threading/blob/v16.4.33/src/Microsoft.VisualStudio.Threading/AwaitExtensions.cs

internal static class AsyncExtensions
{
    public static TaskSchedulerAwaitable SwitchTo(this TaskScheduler scheduler, CancellationToken cancellationToken)
    {
        return new TaskSchedulerAwaitable(scheduler, cancellationToken);
    }

    /// <summary>
    /// Compiler machinery implementing <see cref="SwitchTo"/>.
    /// </summary>
    public readonly struct TaskSchedulerAwaitable
    {
        private readonly TaskScheduler scheduler;
        private readonly CancellationToken cancellationToken;

        /// <summary>
        /// Compiler machinery implementing <see cref="SwitchTo"/>.
        /// </summary>
        public TaskSchedulerAwaitable(TaskScheduler scheduler, CancellationToken cancellationToken)
        {
            this.scheduler = scheduler ?? throw new ArgumentNullException(nameof(scheduler));
            this.cancellationToken = cancellationToken;
        }

        /// <summary>
        /// Compiler machinery implementing <see cref="SwitchTo"/>.
        /// </summary>
        public TaskSchedulerAwaiter GetAwaiter() => new TaskSchedulerAwaiter(scheduler, cancellationToken);

        /// <summary>
        /// Compiler machinery implementing <see cref="SwitchTo"/>.
        /// </summary>
        public readonly struct TaskSchedulerAwaiter : ICriticalNotifyCompletion
        {
            private readonly TaskScheduler scheduler;
            private readonly CancellationToken cancellationToken;

            /// <summary>
            /// Compiler machinery implementing <see cref="SwitchTo"/>.
            /// </summary>
            public TaskSchedulerAwaiter(TaskScheduler scheduler, CancellationToken cancellationToken)
            {
                this.scheduler = scheduler;
                this.cancellationToken = cancellationToken;
            }

            /// <summary>
            /// Compiler machinery implementing <see cref="SwitchTo"/>.
            /// </summary>
            public bool IsCompleted
            {
                get
                {
                    if (cancellationToken.IsCancellationRequested)
                    {
                        // Skip scheduling and go straight to GetResult which will throw OperationCancelledException.
                        return true;
                    }

                    // We special-case 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.

                    return (scheduler == TaskScheduler.Default && Thread.CurrentThread.IsThreadPoolThread)
                        || (scheduler == TaskScheduler.Current && TaskScheduler.Current != TaskScheduler.Default);
                }
            }

            /// <summary>
            /// Compiler machinery implementing <see cref="SwitchTo"/>.
            /// </summary>
            public void OnCompleted(Action continuation)
            {
                if (scheduler == TaskScheduler.Default)
                {
                    ThreadPool.QueueUserWorkItem(state => ((Action)state!)(), continuation);
                }
                else
                {
                    Task.Factory.StartNew(continuation, cancellationToken, TaskCreationOptions.None, scheduler);
                }
            }

            /// <summary>
            /// Compiler machinery implementing <see cref="SwitchTo"/>.
            /// </summary>
            public void UnsafeOnCompleted(Action continuation)
            {
                if (scheduler == TaskScheduler.Default)
                {
                    ThreadPool.UnsafeQueueUserWorkItem(state => ((Action)state!)(), continuation);
                }
                else
                {
                    // There is no API for scheduling a Task without capturing the ExecutionContext.
                    Task.Factory.StartNew(continuation, cancellationToken, TaskCreationOptions.None, scheduler);
                }
            }

            /// <summary>
            /// Compiler machinery implementing <see cref="SwitchTo"/>.
            /// </summary>
            public void GetResult()
            {
                cancellationToken.ThrowIfCancellationRequested();
            }
        }
    }

    #endregion
}