.NET 6.0 中的 await 原理淺析

前言

看過不少關于 await 的原理的文章,也知道背后是編譯器給轉(zhuǎn)成了狀態(tài)機實現(xiàn)的,但是具體是怎么完成的,回調(diào)又是如何銜接的,一直都沒有搞清楚,這次下定決心把源碼自己跑了下,終于豁然開朗了

本文的演示代碼基于 VS2022 + .NET 6

示例

public class Program
{
    static int Work()
    {
        Console.WriteLine("In Task.Run");
        return 1;
    }

    static async Task TestAsync()
    {
        Console.WriteLine("Before Task.Run");
        await Task.Run(Work);
        Console.WriteLine("After Task.Run");
    }

    static void Main()
    {
        _ = TestAsync();
        Console.WriteLine("End");
        Console.ReadKey();
    }
}

很簡單的異步代碼,我們來看下,編譯器把它變成了啥
編譯后的代碼經(jīng)過我的整理,命名簡化了,更容易理解

點擊查看代碼

class Program
{
    static int Work()
    {
        Console.WriteLine("In Task.Run");
        return 1;
    }

    static Task TestAsync()
    {
        var stateMachine = new StateMachine()
        {
            _builder = AsyncTaskMethodBuilder.Create(),
            _state = -1
        };
        stateMachine._builder.Start(ref stateMachine);
        return stateMachine._builder.Task;
    }

    static void Main()
    {
        _ = TestAsync();
        Console.WriteLine("End");
        Console.ReadKey();
    }

    class StateMachine : IAsyncStateMachine
    {
        public int _state;
        public AsyncTaskMethodBuilder _builder;
        private TaskAwaiter<int> _awaiter;

        void IAsyncStateMachine.MoveNext()
        {
            int num = _state;
            try
            {
                TaskAwaiter<int> awaiter;
                if (num != 0)
                {
                    Console.WriteLine("Before Task.Run");
                    awaiter = Task.Run(Work).GetAwaiter();
                    if (!awaiter.IsCompleted)
                    {
                        _state = 0;
                        _awaiter = awaiter;
                        StateMachine stateMachine = this;
                        _builder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine);
                        return;
                    }
                }
                else
                {
                    awaiter = _awaiter;
                    _awaiter = default;
                    _state = -1;
                }
                awaiter.GetResult();
                Console.WriteLine("After Task.Run");
            }
            catch (Exception exception)
            {
                _state = -2;
                _builder.SetException(exception);
                return;
            }
            _state = -2;
            _builder.SetResult();
        }

        void IAsyncStateMachine.SetStateMachine(IAsyncStateMachine stateMachine) { }
    }
}

狀態(tài)機實現(xiàn)

我們看到實際是生成了一個隱藏的狀態(tài)機類 StateMachine
把狀態(tài)機的初始狀態(tài) _state 設置 -1
stateMachine._builder.Start(ref stateMachine); 啟動狀態(tài)機,內(nèi)部實際調(diào)用的就是狀態(tài)機的 MoveNext 方法
Task.Run 創(chuàng)建一個任務, 把委托放在 Task.m_action 字段,丟到線程池,等待調(diào)度

任務在線程池內(nèi)被調(diào)度完成后,是怎么回到這個狀態(tài)機繼續(xù)執(zhí)行后續(xù)代碼的呢?
_builder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine); 就是關鍵了, 跟下去,到了如下的代碼:

if (!this.AddTaskContinuation(stateMachineBox, false))
{
    ThreadPool.UnsafeQueueUserWorkItemInternal(stateMachineBox, true);
}
bool AddTaskContinuation(object tc, bool addBeforeOthers)
{
    return !this.IsCompleted && ((this.m_continuationObject == null && Interlocked.CompareExchange(ref this.m_continuationObject, tc, null) == null) || this.AddTaskContinuationComplex(tc, addBeforeOthers));
}

這里很清楚的看到,嘗試把狀態(tài)機對象(實際是狀態(tài)機的包裝類),賦值到 Task.m_continuationObject, 如果操作失敗,則把狀態(tài)機對象丟進線程池等待調(diào)度,這里為什么這么實現(xiàn),看一下線程池是怎么執(zhí)行的就清楚了

線程池實現(xiàn)

.NET6 的線程池實現(xiàn),實際是放到了 PortableThreadPool, 具體調(diào)試步驟我就不放了,直接說結(jié)果就是, 線程池線程從任務隊列中拿到任務后都執(zhí)行了 DispatchWorkItem 方法

static void DispatchWorkItem(object workItem, Thread currentThread)
{
    Task task = workItem as Task;
    if (task != null)
    {
        task.ExecuteFromThreadPool(currentThread);
        return;
    }
    Unsafe.As<IThreadPoolWorkItem>(workItem).Execute();
}
virtual void ExecuteFromThreadPool(Thread threadPoolThread)
{
    this.ExecuteEntryUnsafe(threadPoolThread);
}

我們看到, 線程池隊列中的任務都是 object 類型的, 這里進行了類型判斷, 如果是 Task , 直接執(zhí)行 task.ExecuteFromThreadPool, 更有意思的這個方法是個虛方法,后面說明

ExecuteFromThreadPool 繼續(xù)追下去,我們來到了這里,代碼做了簡化

private void ExecuteWithThreadLocal(ref Task currentTaskSlot, Thread threadPoolThread = null)
{
    this.InnerInvoke();
    this.Finish(true);
}
virtual void InnerInvoke()
{
    Action action = this.m_action as Action;
    if (action != null)
    {
        action();
        return;
    }
}

很明顯 this.InnerInvoke 就是執(zhí)行了最開始 Task.Run(Work) 封裝的委托了, 在 m_action 字段
this.Finish(true); 跟下去會發(fā)現(xiàn)會調(diào)用 FinishStageTwo 設置任務的完成狀態(tài),異常等, 繼續(xù)調(diào)用 FinishStageThree 就來了重點: FinishContinuations 這個方法就是銜接后續(xù)回調(diào)的核心
```
internal void FinishContinuations()
{
    object obj = Interlocked.Exchange(ref this.m_continuationObject, Task.s_taskCompletionSentinel);
    if (obj != null)
    {
        this.RunContinuations(obj);
    }
}
```

還記得狀態(tài)機實現(xiàn)么, Task.m_continuationObject 字段實際存儲的就是狀態(tài)機的包裝類,這里線程池線程也會判斷這個字段有值的話,就直接使用它執(zhí)行后續(xù)代碼了

void RunContinuations(object continuationObject)
{
    var asyncStateMachineBox = continuationObject as IAsyncStateMachineBox;
    if (asyncStateMachineBox != null)
    {
        AwaitTaskContinuation.RunOrScheduleAction(asyncStateMachineBox, flag2);
        return;
    }
}

static void RunOrScheduleAction(IAsyncStateMachineBox box, bool allowInlining)
{
    if (allowInlining && AwaitTaskContinuation.IsValidLocationForInlining)
    {
        box.MoveNext();
        return;
    }
}

總結(jié)

Task.Run 創(chuàng)建 Task, 把委托放在 m_action 字段, 把 Task 壓入線程池隊列,等待調(diào)度
_builder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine); 嘗試把狀態(tài)機對象放在 Task.m_continuationObject 字段上,等待線程池線程調(diào)度完成任務后使用(用來執(zhí)行后續(xù)),若操作失敗,直接把狀態(tài)機對象壓入線程池隊列,等待調(diào)度
線程池線程調(diào)度任務完成后,會判斷 Task.m_continuationObject 有值,直接執(zhí)行它的 MoveNext

備注

狀態(tài)機實現(xiàn)中,嘗試修改 Task.m_continuationObject,可能會失敗,
就會直接把狀態(tài)機對象壓入線程池, 但是線程池調(diào)度,不都是判斷是不是 Task 類型么, 其實狀態(tài)機的包裝類是 Task 的子類,哈哈,是不是明白了

class AsyncStateMachineBox<TStateMachine> : Task<TResult>, IAsyncStateMachineBox where TStateMachine : IAsyncStateMachine

static void DispatchWorkItem(object workItem, Thread currentThread)
{
    Task task = workItem as Task;
    if (task != null)
    {
        task.ExecuteFromThreadPool(currentThread);
        return;
    }
    Unsafe.As<IThreadPoolWorkItem>(workItem).Execute();
}

還有就是狀態(tài)機包裝類,重寫了 Task.ExecuteFromThreadPool,所以線程池調(diào)用 task.ExecuteFromThreadPool 就是直接調(diào)用了狀態(tài)機的 MoveNext 了, Soga ^_^
```
override void ExecuteFromThreadPool(Thread threadPoolThread)
{
    this.MoveNext(threadPoolThread);
}
```

參考鏈接

關于線程池和異步的更深刻的原理,大家可以參考下面的文章

posted @ 2023-11-15 11:19 Broadm 閱讀(1650) 評論(4) 收藏舉報

刷新頁面返回頂部

Broadm

Better late than never

.NET 6.0 中的 await 原理淺析

前言

示例

狀態(tài)機實現(xiàn)

線程池實現(xiàn)

總結(jié)

備注

參考鏈接

公告