Async thread's Activiy.Current is set incorrectly and cannot be rolled back

[CuteLeon]

Description

Hi Team:
I found that Activity.Current is returning a outdated Activity when start new activity in an async scenario, because Activity.Current's type is AsyncLocal and Activity's Parent is passed to mutilple thread but Activity only stopped on single thread.;

I started the ActivityParent in ThreadParent, then call and await an async task on ThreadChild in ActivityA's scope, create the ActivityChild on ThreadChild, ActivityChild is having ActivityParent as its parent, which is good.
Then I stopped ActivityChild on ThreadChild, ThreadChild's Activity.Current is rolled back to ActivityChild's parent (ActicityParent).
Then program go back to ThreadParent and ActivityParent stops on ThreadParent, ThreadParent's Activity.Current rolls back to ActivityParent's parent (null), which is good.
ThreadChild is designed to be alive and waiting for other jobs, but ActivityParent will never stop again on ThreadChild, so that ThreadChild's Activity.Current keeps returnning outdated ActivityParent as new Activity's parent on ThreadChild, and can not be rolled back to null any more.

Reproduction Steps

static async Task Main(string[] args)
{
    var listener = new ActivityListener
    {
        ShouldListenTo = (x) => true,
        Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllDataAndRecorded,
        ActivityStarted = (x) =>
        {
            Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [TID={Environment.CurrentManagedThreadId,3}] ActivityStarted: {x.DisplayName} (Parent={x.Parent?.DisplayName ?? "[null]"}) (Current={Activity.Current?.DisplayName ?? "[null]"})");
        },
        ActivityStopped = (x) =>
        {
            Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [TID={Environment.CurrentManagedThreadId,3}] ActivityStopped: {x.DisplayName} (Parent={x.Parent?.DisplayName ?? "[null]"}) (Current={Activity.Current?.DisplayName ?? "[null]"})");
        }
    };
    ActivitySource.AddActivityListener(listener);

    Activity.CurrentChanged += (s, e) =>
    {
        Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [TID={Environment.CurrentManagedThreadId,3}] \t\tCurrentChanged: {e.Previous?.DisplayName ?? "[null]"} => {e.Current?.DisplayName ?? "[null]"}");
    };

    var source = new ActivitySource("TestActivitySource");
    Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [TID={Environment.CurrentManagedThreadId,3}] Current activity before Main: {Activity.Current?.DisplayName ?? "[null]"}");
    using (var activity = source.StartActivity("ActivityParent")!)
    {
        Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [TID={Environment.CurrentManagedThreadId,3}] Current activity in Main: {Activity.Current?.DisplayName ?? "[null]"}");

        var task = AsyncTask();
        Thread.Sleep(30);
        Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [TID={Environment.CurrentManagedThreadId,3}] Current activity in Main after async task: {Activity.Current?.DisplayName ?? "[null]"}");
        await task;
    }
    Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [TID={Environment.CurrentManagedThreadId,3}] Current activity after Main: {Activity.Current?.DisplayName ?? "[null]"}");
    Console.Read();
}

static async Task AsyncTask()
{
    var source = new ActivitySource("TestActivitySource");
    Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [TID={Environment.CurrentManagedThreadId,3}] Current activity before async task: {Activity.Current?.DisplayName ?? "[null]"}");
    using (var activity = source.StartActivity("ActivityChild")!)
    {
        Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [TID={Environment.CurrentManagedThreadId,3}] Current activity in async task: {Activity.Current?.DisplayName ?? "[null]"}");
        await Task.Delay(10);
    }
    // Here, ThreadChild's Activity.Current is ActivityParent instead of null, which is bad.
    Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [TID={Environment.CurrentManagedThreadId,3}] Current activity after async task: {Activity.Current?.DisplayName ?? "[null]"}");
}

Expected behavior

Return an Activity object with correct parent from ActivitySource.StartActivity() in async scenario.

Actual behavior

Activity.Current can not roll back correctly once activitywas passed over multiple threads.

Regression?

No response

Known Workarounds

No response

Configuration

.Net 8.0
Windows Server 2022
x64

Other information

No response

[dotnet-policy-service]

Tagging subscribers to this area: @tarekgh, @tommcdon, @pjanotti
See info in area-owners.md if you want to be subscribed.

[CuteLeon]

14:05:10.021 [TID=  1] Current activity before Main: [null]
14:05:10.031 [TID=  1]          CurrentChanged: [null] => ActivityParent
14:05:10.031 [TID=  1] ActivityStarted: ActivityParent (Parent=[null]) (Current=ActivityParent)
14:05:10.031 [TID=  1] Current activity in Main: ActivityParent
14:05:10.032 [TID=  1] Current activity before async task: ActivityParent
14:05:10.032 [TID=  1]          CurrentChanged: ActivityParent => ActivityChild
14:05:10.032 [TID=  1] ActivityStarted: ActivityChild (Parent=ActivityParent) (Current=ActivityChild)
14:05:10.032 [TID=  1] Current activity in async task: ActivityChild
14:05:10.047 [TID= 10] ActivityStopped: ActivityChild (Parent=ActivityParent) (Current=ActivityChild)
14:05:10.047 [TID= 10]          CurrentChanged: ActivityChild => ActivityParent
14:05:10.047 [TID= 10] Current activity after async task: ActivityParent
14:05:10.078 [TID=  1] Current activity in Main after async task: ActivityParent
14:05:10.078 [TID=  1] ActivityStopped: ActivityParent (Parent=[null]) (Current=ActivityParent)
14:05:10.078 [TID=  1]          CurrentChanged: ActivityParent => [null]
14:05:10.078 [TID=  1] Current activity after Main: [null]

You can see 14:05:10.047 [TID= 10] Current activity after async task: ActivityParent

[davidfowl]

This is unfortunately by design because of how async locals work. Also this example is inherently racy. What are you trying to achieve?