.NET AsyncLocal 避坑指南

2023-03-02 06:00:35

AsyncLocal 用法簡介

通過 AsyncLocal 我們可以在一個邏輯上下中維護一份資料,並且在後續程式碼中都可以存取和修改這份資料。

無論是在新建立的 Task 中還是 await 關鍵詞之後,我們都能夠存取前面設定的 AsyncLocal 的資料。

class Program
{
    private static AsyncLocal<string> _asyncLocal = new AsyncLocal<string>();
    
    static async Task Main(string[] args)
    {
        _asyncLocal.Value = "Hello World!";
        Task.Run(() => Console.WriteLine($"AsyncLocal in task: {_asyncLocal.Value}"));

        await FooAsync();
        Console.WriteLine($"AsyncLocal after await FooAsync: {_asyncLocal.Value}");
    }

    private static async Task FooAsync()
    {
        await Task.Delay(100);
        Console.WriteLine($"AsyncLocal after await in FooAsync: {_asyncLocal.Value}");
    }
}

輸出結果:

AsyncLocal in task: Hello World!
AsyncLocal after await in FooAsync: Hello World!
AsyncLocal after await FooAsync: Hello World!

AsyncLocal 實現原理

在我之前的部落格 揭祕 .NET 中的 AsyncLocal 中深入介紹了 AsyncLocal 的實現原理,這裡只做簡單的回顧。

AsyncLocal 的實際資料儲存在 ExecutionContext 中,而 ExecutionContext 作為執行緒的私有欄位與執行緒繫結,線上程會發生切換的地方,runtime 會將切換前的 ExecutionContext 儲存起來,切換後再恢復到新執行緒上。

這個儲存和恢復的過程是由 runtime 自動完成的,例如會發生在以下幾個地方:

  • new Thread(ThreadStart start).Start()
  • Task.Run(Action action)
  • ThreadPool.QueueUserWorkItem(WaitCallback callBack)
  • await 之後

以 await 為例,當我們在一個方法中使用了 await 關鍵詞,編譯器會將這個方法編譯成一個狀態機,這個狀態機會在 await 之前和之後分別儲存和恢復 ExecutionContext。

class Program
{
    private static AsyncLocal<string> _asyncLocal = new AsyncLocal<string>();
    
    static async Task Main(string[] args)
    {
        _asyncLocal.Value = "Hello World!";
        await FooAsync();
        Console.WriteLine($"AsyncLocal after await FooAsync: {_asyncLocal.Value}");
    }

    private static async Task FooAsync()
    {
        await Task.Delay(100);
    }
}

輸出結果:

AsyncLocal after await FooAsync: Hello World!

AsyncLocal 的坑

有時候我們會在 FooAsync 方法中去修改 AsyncLocal 的值,並希望在 Main 方法在 await FooAsync 之後能夠獲取到修改後的值,但是實際上這是不可能的。

class Program
{
    private static AsyncLocal<string> _asyncLocal = new AsyncLocal<string>();
    
    static async Task Main(string[] args)
    {
        _asyncLocal.Value = "A";
        Console.WriteLine($"AsyncLocal before FooAsync: {_asyncLocal.Value}");
        await FooAsync();
        Console.WriteLine($"AsyncLocal after await FooAsync: {_asyncLocal.Value}");
    }

    private static async Task FooAsync()
    {
        _asyncLocal.Value = "B";
        Console.WriteLine($"AsyncLocal before await in FooAsync: {_asyncLocal.Value}");
        await Task.Delay(100);
        Console.WriteLine($"AsyncLocal after await in FooAsync: {_asyncLocal.Value}");
    }
}

輸出結果:

AsyncLocal before FooAsync: A
AsyncLocal before await in FooAsync: B
AsyncLocal after await in FooAsync: B
AsyncLocal after await FooAsync: A

為什麼我們在 FooAsync 方法中修改了 AsyncLocal 的值,但是在 await FooAsync 之後,AsyncLocal 的值卻沒有被修改呢?

原因是 ExecutionContext 被設計成了一個不可變的物件,當我們在 FooAsync 方法中修改了 AsyncLocal 的值,實際上是建立了一個新的 ExecutionContext,原來的 AsyncLocal 的值被值拷貝到了新的 ExecutionContext 中,而原來的 ExecutionContext 仍然保持不變。

這樣的設計是為了保證執行緒的安全性,因為在多執行緒環境下,如果 ExecutionContext 是可變的,那麼在切換執行緒的時候,可能會出現資料不一致的情況。

我們通常把這種設計稱為 Copy On Write(簡稱COW),即在修改資料的時候,會先拷貝一份資料,然後在拷貝的資料上進行修改,這樣就不會影響到原來的資料。

ExecutionContext 中可能不止一個 AsyncLocal 的資料,修改任意一個 AsyncLocal 都會導致 ExecutionContext 的 COW。

所以上面程式碼的執行過程如下:

AsyncLocal 的避坑指南

那麼我們如何在 FooAsync 方法中修改 AsyncLocal 的值,並且在 Main 方法中獲取到修改後的值呢?

我們需要藉助一箇中介者,讓中介者來儲存 AsyncLocal 的值,然後在 FooAsync 方法中修改中介者的屬性值,這樣就可以在 Main 方法中獲取到修改後的值了。

下面我們設計一個 ValueHolder 來儲存 AsyncLocal 的值,修改 Value 並不會修改 AsyncLocal 的值,而是修改 ValueHolder 的屬性值,這樣就不會觸發 ExecutionContext 的 COW。

我們還需要設計一個 ValueAccessor 來封裝 ValueHolder 對值的存取和修改,這樣可以保證 ValueHolder 的值只能在 ValueAccessor 中被修改。

class ValueAccessor<T> : IValueAccessor<T>
{
    private static AsyncLocal<ValueHolder<T>> _asyncLocal = new AsyncLocal<ValueHolder<T>>();

    public T Value
    {
        get => _asyncLocal.Value != null ? _asyncLocal.Value.Value : default;
        set
        {
            _asyncLocal.Value ??= new ValueHolder<T>();

            _asyncLocal.Value.Value = value;
        }
    }
}

class ValueHolder<T>
{
    public T Value { get; set; }
}

class Program
{
    private static IValueAccessor<string> _valueAccessor = new ValueAccessor<string>();

    static async Task Main(string[] args)
    {
        _valueAccessor.Value = "A";
        Console.WriteLine($"ValueAccessor before await FooAsync in Main: {_valueAccessor.Value}");
        await FooAsync();
        Console.WriteLine($"ValueAccessor after await FooAsync in Main: {_valueAccessor.Value}");
    }

    private static async Task FooAsync()
    {
        _valueAccessor.Value = "B";
        Console.WriteLine($"ValueAccessor before await in FooAsync: {_valueAccessor.Value}");
        await Task.Delay(100);
        Console.WriteLine($"ValueAccessor after await in FooAsync: {_valueAccessor.Value}");
    }
}

輸出結果:

ValueAccessor before await FooAsync in Main: A
ValueAccessor before await in FooAsync: B
ValueAccessor after await in FooAsync: B
ValueAccessor after await FooAsync in Main: B

HttpContextAccessor 的實現原理

我們常用的 HttpContextAccessor 通過HttpContextHolder 來間接地在 AsyncLocal 中儲存 HttpContext。

如果要更新 HttpContext,只需要在 HttpContextHolder 中更新即可。因為 AsyncLocal 的值不會被修改,更新 HttpContext 時 ExecutionContext 也不會出現 COW 的情況。

不過 HttpContextAccessor 中的邏輯有點特殊,它的 HttpContextHolder 是為保證清除 HttpContext 時,這個 HttpContext 能在所有參照它的 ExecutionContext 中被清除(可能因為修改 HttpContextHolder 之外的 AsyncLocal 資料導致 ExecutionContext 已經 COW 很多次了)。

下面是 HttpContextAccessor 的實現,英文註釋是原文,中文註釋是我自己的理解。

/// </summary>
public class HttpContextAccessor : IHttpContextAccessor
{
    private static readonly AsyncLocal<HttpContextHolder> _httpContextCurrent = new AsyncLocal<HttpContextHolder>();

    /// <inheritdoc/>
    public HttpContext? HttpContext
    {
        get
        {
            return _httpContextCurrent.Value?.Context;
        }
        set
        {
            var holder = _httpContextCurrent.Value;
            if (holder != null)
            {
                // Clear current HttpContext trapped in the AsyncLocals, as its done.
                // 這邊的邏輯是為了保證清除 HttpContext 時,這個 HttpContext 能在所有參照它的 ExecutionContext 中被清除
                holder.Context = null;
            }

            if (value != null)
            {
                // Use an object indirection to hold the HttpContext in the AsyncLocal,
                // so it can be cleared in all ExecutionContexts when its cleared.
                // 這邊直接修改了 AsyncLocal 的值,所以會導致 ExecutionContext 的 COW。新的 HttpContext 不會被傳遞到原先的 ExecutionContext 中。
                _httpContextCurrent.Value = new HttpContextHolder { Context = value };
            }
        }
    }

    private sealed class HttpContextHolder
    {
        public HttpContext? Context;
    }
}

但 HttpContextAccessor 的實現並不允許將新賦值的非 null 的 HttpContext 傳遞到外層的 ExecutionContext 中,可以參考上面的原始碼及註釋理解。

class Program
{
    private static IHttpContextAccessor _httpContextAccessor = new HttpContextAccessor();
    
    static async Task Main(string[] args)
    {
        var httpContext = new DefaultHttpContext
        {
            Items = new Dictionary<object, object>
            {
                { "Name", "A"}
            }
        };
        _httpContextAccessor.HttpContext = httpContext;
        Console.WriteLine($"HttpContext before await FooAsync in Main: {_httpContextAccessor.HttpContext.Items["Name"]}");
        await FooAsync();
        // HttpContext 被清空了,下面這行輸出 null
        Console.WriteLine($"HttpContext after await FooAsync in Main: {_httpContextAccessor.HttpContext?.Items["Name"]}");
    }

    private static async Task FooAsync()
    {
        _httpContextAccessor.HttpContext = new DefaultHttpContext
        {
            Items = new Dictionary<object, object>
            {
                { "Name", "B"}
            }
        };
        Console.WriteLine($"HttpContext before await in FooAsync: {_httpContextAccessor.HttpContext.Items["Name"]}");
        await Task.Delay(1000);
        Console.WriteLine($"HttpContext after await in FooAsync: {_httpContextAccessor.HttpContext.Items["Name"]}");
    }
}

輸出結果:

HttpContext before await FooAsync in Main: A
HttpContext before await in FooAsync: B
HttpContext after await in FooAsync: B
HttpContext after await FooAsync in Main: