await-for-context-switching.cs

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

/*
実行結果の一例:
10: main thread (PostCount = 0)
6: task pool (PostCount = 0)
10: main thread (PostCount = 1)
6: task pool (PostCount = 1)
10: main thread (PostCount = 2)
10: main thread (PostCount = 2)
Press any key...
*/

namespace VainZero.Sandbox
{
    /// <summary>
    /// 処理をメインスレッドにディスパッチする同期コンテクストを表す。
    /// WPF の UI スレッドに関連づけられている同期コンテクストの簡易版。
    /// </summary>
    public sealed class MainThreadSynchronizationContext
        : SynchronizationContext
    {
        readonly Queue<Action> queue = new Queue<Action>();

        /// <summary>
        /// 処理がディスパッチされた回数を取得する。
        /// </summary>
        public int PostCount { get; private set; }

        public override void Send(SendOrPostCallback d, object state)
        {
            throw new NotImplementedException();
        }

        /// <summary>
        /// 処理をディスパッチする。
        /// このメソッドは、その処理の完了を待たずに制御を返す。
        /// </summary>
        /// <param name="d"></param>
        /// <param name="state"></param>
        public override void Post(SendOrPostCallback d, object state)
        {
            lock (queue)
            {
                queue.Enqueue(() => d(state));
                PostCount++;
            }
        }

        /// <summary>
        /// メッセージループを実行する。
        /// 指定されたタスクが完了したら自動的に終了する。
        /// </summary>
        public void Loop(Task task)
        {
            var wait = new SpinWait();

            while (true)
            {
                var action = default(Action);
                lock (queue)
                {
                    if (queue.Count == 0)
                    {
                        if (task.IsCompleted) goto OnCompleted;
                        goto Spin;
                    }
                    action = queue.Dequeue();
                }
                action();
                continue;

                Spin:
                wait.SpinOnce();
                continue;

                OnCompleted:
                task.Wait();
                break;
            }
        }
    }

    /// <summary>
    /// 継続をタスクプールで実行する awaitable を表す。
    /// 結果としては、await したスレッドの同期コンテクスト (または null) を返す。
    /// </summary>
    public struct SwitchToTaskPoolAwaitable
        : INotifyCompletion
    {
        readonly SynchronizationContext context;

        /// <summary>
        /// Invoked by await.
        /// </summary>
        /// <returns></returns>
        public SwitchToTaskPoolAwaitable GetAwaiter() => this;

        /// <summary>
        /// await したときに同期的に継続を実行できるかを取得する。
        /// 常に OnCompleted を起動してほしいので、完了していないことにする。
        /// </summary>
        public bool IsCompleted => false;

        /// <summary>
        /// await 元のスレッドの同期コンテクストを取得する。
        /// </summary>
        /// <returns></returns>
        public SynchronizationContext GetResult()
        {
            return context;
        }

        /// <summary>
        /// await の継続を実行する。
        /// </summary>
        /// <param name="continuation"></param>
        public void OnCompleted(Action continuation)
        {
            Task.Run(continuation);
        }

        internal SwitchToTaskPoolAwaitable(SynchronizationContext context)
        {
            this.context = context;
        }
    }

    /// <summary>
    /// 継続を指定された同期コンテクストで実行する awaiter を表す。
    /// </summary>
    public struct SynchronizationContextAwaiter
        : INotifyCompletion
    {
        readonly SynchronizationContext context;

        /// <summary>
        /// await したときに同期的に継続を実行できるかを取得する。
        /// context が null の場合は同期的に実行するしかないので、<c>true</c> を返す。
        /// 現在のスレッドの同期コンテクストが捕捉された同期コンテクストと同じ場合も、
        /// その実装は継続を同期的に実行するとみなしてよい (Rx.NET の実装がそうなっている) ので、
        /// <c>true</c> を返す。
        /// </summary>
        public bool IsCompleted =>
            context == null || SynchronizationContext.Current == context;

        /// <summary>
        /// Gets void.
        /// </summary>
        public void GetResult()
        {
        }

        /// <summary>
        /// 継続を同期コンテクストの中で実行する。
        /// </summary>
        /// <param name="continuation"></param>
        public void OnCompleted(Action continuation)
        {
            if (context == null)
            {
                // context == null の場合、await がこのメソッドを起動することはないが、
                // 手動で起動されることはありえる。
                continuation();
            }
            else
            {
                context.Post(_ => continuation(), default(object));
            }
        }

        internal SynchronizationContextAwaiter(SynchronizationContext context)
        {
            this.context = context;
        }
    }

    public static class TaskModule
    {
        public static SwitchToTaskPoolAwaitable SwitchToTaskPool()
        {
            return new SwitchToTaskPoolAwaitable(SynchronizationContext.Current);
        }
    }

    public static class SynchronizationContextExtension
    {
        public static SynchronizationContextAwaiter GetAwaiter(this SynchronizationContext @this)
        {
            return new SynchronizationContextAwaiter(@this);
        }
    }

    public sealed class Program
    {
        readonly MainThreadSynchronizationContext mainContext =
            new MainThreadSynchronizationContext();

        void Print(string threadName)
        {
            Console.WriteLine(
                "{0}: {1} (PostCount = {2})",
                // 現在いるスレッドの識別子
                Thread.CurrentThread.ManagedThreadId,
                // 現在いるスレッドの種類
                threadName,
                // メインスレッドへのディスパッチが起こった回数
                mainContext.PostCount
            );
        }

        async Task Start()
        {
            Print("main thread");

            // タスクプールに移動する。
            // ついでに、メインスレッドの同期コンテクストを捕捉する。
            var context = await TaskModule.SwitchToTaskPool();

            Print("task pool");

            // 同期コンテクスト (メインスレッド) にスイッチする。
            await context;

            Print("main thread");

            // タスクプールを忙しくさせながら、タスクプールにスイッチする。
            // おそらく新しいスレッドが開始されるので、先ほどとは異なるスレッドIDが表示される。
            for (var i = 0; i < 10; i++)
            {
                var _ = Task.Run(() => Thread.Sleep(500));
            }

            await TaskModule.SwitchToTaskPool();

            Print("task pool");

            // 同期コンテクストに戻る。
            await context;

            Print("main thread");

            // 既に同期コンテクストの内部 (メインスレッド上) にいる場合は Post されない。
            // すなわち、PostCount が増加しない。
            await context;

            Print("main thread");
        }

        public void Run()
        {
            // メインスレッドに同期コンテクストを設定する。
            SynchronizationContext.SetSynchronizationContext(mainContext);

            // タスクを開始する。
            var task = Start();

            // タスクが完了するまでメインループを行う。
            // (タスクから同期コンテクスト経由でディスパッチされた処理を実行していく。)
            mainContext.Loop(task);

            Console.WriteLine("Press any key...");
            Console.ReadKey();
        }

        public static void Main(string[] args)
        {
            new Program().Run();
        }
    }
}

license.txt

This is free and unencumbered software released into the public domain.

Anyone is free to copy, modify, publish, use, compile, sell, or
distribute this software, either in source code form or as a compiled
binary, for any purpose, commercial or non-commercial, and by any
means.

In jurisdictions that recognize copyright laws, the author or authors
of this software dedicate any and all copyright interest in the
software to the public domain. We make this dedication for the benefit
of the public at large and to the detriment of our heirs and
successors. We intend this dedication to be an overt act of
relinquishment in perpetuity of all present and future rights to this
software under copyright law.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.

For more information, please refer to <http://unlicense.org/>