async/await的多線程問題

今天嘗試把.net4.5新增的異步編程模型async/await加入自己的框架，因為從第一印象看，使用async/await的寫法實在太方便了，以同步代碼的方式寫異步流程，寫起來更順暢，不容易打斷思路，異常傳遞、資源控制(lock和using)也都完美支持，即使有少量的性能損失，也完全可以接受。

首先我寫了一個測試代碼，以熟悉async/await模型，代碼如下：

static class TestCase
{
    static async Task Test2()
    {
        Console.WriteLine("3 {0}", Thread.CurrentThread.ManagedThreadId);
        await Task.Run(() =>
        {
            Thread.Sleep(1000);
        });// 這里可以換成Task.Delay(1000);
        Console.WriteLine("4 {0}", Thread.CurrentThread.ManagedThreadId);
    }

    static async Task Test1()
    {
        Console.WriteLine("1 {0}", Thread.CurrentThread.ManagedThreadId);
        await Test2();
        Console.WriteLine("2 {0}", Thread.CurrentThread.ManagedThreadId);
    }

    static void Main(string[] args)
    {
        Test1().Wait();
    }
}

輸出如下：

代碼很簡單，結果也看似正確，直到我發現一個大問題：

await后面的代碼是在另一個線程執行的！！！

這不僅會影響到框架的運行，還會導致線程爭用資源的問題，也就是說看似同一處的代碼，會運行在不同線程，並且有可能並行。如果訪問公共資源(如靜態變量)還需要加鎖，嚴重影響編程體驗和性能。我推測，很可能Task類自帶的創建Task的靜態函數所產生的Task，設置Complete狀態時都是在其他線程，因此導致了await的回調也在該線程執行。於是我把代碼稍微改了一下：

static class TestCase
{
    static TaskCompletionSource<object> source = new TaskCompletionSource<object>();

    static async Task Test2()
    {
        Console.WriteLine("3 {0}", Thread.CurrentThread.ManagedThreadId);
        await source.Task;
        Console.WriteLine("4 {0}", Thread.CurrentThread.ManagedThreadId);
    }

    static async Task Test1()
    {
        Console.WriteLine("1 {0}", Thread.CurrentThread.ManagedThreadId);
        await Test2();
        Console.WriteLine("2 {0}", Thread.CurrentThread.ManagedThreadId);
    }

    static void Main(string[] args)
    {
        Task task = Test1();
        Thread.Sleep(1000);
        source.SetResult(null);
        task.Wait(); // 這一句其實是沒有必要的，並且如果放在SetResult前，會導致死鎖
    }
}

輸出如下：

果然全部是在主線程執行了。其中source.SetResult(null);這一句放到框架中，當條件合適時就在主線程Loop中調用，以便喚醒await繼續執行剩下的過程。此外，Task類的靜態函數所產生的Task，也可以通過一個包裝函數，來讓在其他線程執行的SetResult，Queue到主線程調用，類似這樣：

public Task<T> Wrap<T>(Task<T> task)
{
    Loop loop = Current;
    TaskCompletionSource<T> source = new TaskCompletionSource<T>();
    task.GetAwaiter().OnCompleted(() =>
    {
        loop.Execute(() =>
        {
            if (task.IsCompleted)
                source.TrySetResult(task.Result);
            else if (task.IsCanceled)
                source.TrySetCanceled();
            else
                source.TrySetException(task.Exception);
        });
    });
    return source.Task;
}