.NET AsyncLocal 避坑指南

黑洞視界發表於2023-03-01

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 中,新的 AsyncLocal 的值只會寫入到新的 ExecutionContext 中,而原來的 ExecutionContext 及其關聯的 AsyncLocal 仍然保持不變。

這樣的設計是為了保證執行緒的安全性,因為在多執行緒環境下,如果 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: 

歡迎關注個人技術公眾號

相關文章