TaskScheduler.Default.SwitchTo

SwitchTo.cs

// md wfapp
// dotnet new winforms
// then copy/paste this

using System.Diagnostics;

namespace wfapp;

public partial class Form1 : Form
{
    // library code
    private static async Task LibApiAsync()
    {
        // on UI thread
        Debug.WriteLine(Thread.CurrentThread.IsThreadPoolThread);

        await TaskScheduler.Default.SwitchTo();

        // on ThreadPool thread
        Debug.WriteLine(Thread.CurrentThread.IsThreadPoolThread);

        // not blocking the UI
        Thread.Sleep(1000);
    }

    // client code
    private async void Form1_Load(object sender, EventArgs e)
    {
        // on UI thread
        Debug.WriteLine(Thread.CurrentThread.IsThreadPoolThread);

        var scheduler = TaskScheduler.Current;

        await LibApiAsync();

        // true
        Debug.WriteLine(scheduler == TaskScheduler.Current);

        // still on UI thread
        this.Text = "Hello!";

        Debug.WriteLine(Thread.CurrentThread.IsThreadPoolThread);
    }

    public Form1()
    {
        InitializeComponent();
    }
}


public static class TaskSchedulerExtensions
{
    /// <summary>
    /// Continue on the specified task scheduler, which becomes the current one
    /// Inspired by <see cref="https://github.com/dotnet/runtime/issues/20025"/>this GitHub issue</see>.
    /// </summary>
    /// <param name="@this">A task scheduler instance, e.g., <c>TaskScheduler.Default</c></param>
    /// <param name="alwaysSchedule">Always use the task scheduler to queue the continuations,
    /// even if it can be executed synchronously.
    /// </param>
    /// <example>
    /// <code>
    /// await await TaskScheduler.Default.SwitchTo(alwaysSchedule: true);
    /// </code>
    /// </example> 
    /// <returns></returns>
    public static TaskSchedulerAwaitable SwitchTo(this TaskScheduler @this, bool alwaysSchedule = false)
    {
        return new TaskSchedulerAwaitable(@this, alwaysSchedule);
    }

    public struct TaskSchedulerAwaiter : System.Runtime.CompilerServices.ICriticalNotifyCompletion
    {
        private readonly TaskScheduler _scheduler;
        private readonly bool _alwaysSchedule;

        public TaskSchedulerAwaiter(TaskScheduler scheduler, bool alwaysSchedule = false)
        {
            _scheduler = scheduler;
            _alwaysSchedule = alwaysSchedule;
        }

        // optimize if already on the default task scheduler
        // and on a thread pool thread without sync context
        private bool CanInline() =>
            !_alwaysSchedule &&
            _scheduler == TaskScheduler.Default &&
            TaskScheduler.Current == TaskScheduler.Default &&
            Thread.CurrentThread.IsThreadPoolThread &&
            SynchronizationContext.Current == null;

        public bool IsCompleted => CanInline();

        public void GetResult() { }

        // a safe version that has to flow the execution context
        public void OnCompleted(Action continuation)
        {
            throw new NotImplementedException(nameof(OnCompleted));
        }

        // an unsafe version that doesn't have to flow the execution context
        public void UnsafeOnCompleted(Action continuation)
        {
            // use ThreadPool.UnsafeQueueUserWorkItem to optimize for TaskScheduler.Default 
            if (_scheduler == TaskScheduler.Default)
            {
                ThreadPool.UnsafeQueueUserWorkItem(
                    c => ((Action)c!).Invoke(),
                    continuation,
                    preferLocal: true);
                return;
            }

            // use Task.Factory.StartNew for all non-default task schedulers
            _ = Task.Factory.StartNew(
                continuation,
                CancellationToken.None,
                TaskCreationOptions.None,
                _scheduler);
        }
    }

    public struct TaskSchedulerAwaitable
    {
        private readonly TaskSchedulerAwaiter _awaiter;

        public TaskSchedulerAwaitable(TaskScheduler scheduler, bool alwaysSchedule = false)
        {
            _awaiter = new TaskSchedulerAwaiter(scheduler, alwaysSchedule);
        }

        public TaskSchedulerAwaiter GetAwaiter()
        {
            return _awaiter;
        }
    }
}