Why does my async method builder have to be a class or run in Debug mode?
Asked Answered
G

2

11

I'm trying to implement my own async method builder for a custom awaitable type. My awaitable type is just a struct containing a ValueTask<T>.

The problem is my asynchronous method builder only works when it's a class or compiled in Debug mode, not a struct and in Release mode.

Here's a minimal, reproducible example. You have to copy this code into a new console project on your local PC and run it in Release mode; .NET Fiddle apparently runs snippets in Debug mode. And of course this requires .Net 5+: https://dotnetfiddle.net/S6F9Hd

This code completes successfully when CustomAwaitableAsyncMethodBuilder<T> is a class or it is compiled in Debug mode. But it hangs and fails to complete otherwise:

class Program
{
    static async Task Main()
    {
        var expected = Guid.NewGuid().ToString();
        async CustomAwaitable<string> GetValueAsync()
        {
            await Task.Yield();
            return expected;
        }

        var actual = await GetValueAsync();

        if (!ReferenceEquals(expected, actual))
            throw new Exception();

        Console.WriteLine("Done!");
    }
}

Here is my custom awaitable type:

[AsyncMethodBuilder(typeof(CustomAwaitableAsyncMethodBuilder<>))]
public readonly struct CustomAwaitable<T>
{
    readonly ValueTask<T> _valueTask;

    public CustomAwaitable(ValueTask<T> valueTask)
    {
        _valueTask = valueTask;
    }

    public ValueTaskAwaiter<T> GetAwaiter() => _valueTask.GetAwaiter();
}

And here is my custom async method builder. Again, all I have to do to make the code run is change this from a struct to a class:

public struct CustomAwaitableAsyncMethodBuilder<T>
{
    Exception? _exception;
    bool _hasResult;
    SpinLock _lock;
    T? _result;
    TaskCompletionSource<T>? _source;

    public CustomAwaitable<T> Task
    {
        get
        {
            var lockTaken = false;
            try
            {
                _lock.Enter(ref lockTaken);
                if (_exception is not null)
                    return new CustomAwaitable<T>(ValueTask.FromException<T>(_exception));
                if (_hasResult)
                    return new CustomAwaitable<T>(ValueTask.FromResult(_result!));
                return new CustomAwaitable<T>(
                    new ValueTask<T>(
                        (_source ??= new TaskCompletionSource<T>(TaskCreationOptions.RunContinuationsAsynchronously))
                        .Task
                    )
                );
            }
            finally
            {
                if (lockTaken)
                    _lock.Exit();
            }
        }
    }

    public void AwaitOnCompleted<TAwaiter, TStateMachine>(
        ref TAwaiter awaiter,
        ref TStateMachine stateMachine)
        where TAwaiter : INotifyCompletion
        where TStateMachine : IAsyncStateMachine =>
        awaiter.OnCompleted(stateMachine.MoveNext);

    public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(
        ref TAwaiter awaiter,
        ref TStateMachine stateMachine)
        where TAwaiter : ICriticalNotifyCompletion
        where TStateMachine : IAsyncStateMachine =>
        awaiter.UnsafeOnCompleted(stateMachine.MoveNext);

    public static CustomAwaitableAsyncMethodBuilder<T> Create() => new()
    {
        _lock = new SpinLock(Debugger.IsAttached)
    };

    public void SetException(Exception exception)
    {
        var lockTaken = false;
        try
        {
            _lock.Enter(ref lockTaken);
            if (Volatile.Read(ref _source) is {} source)
            {
                source.TrySetException(exception);
            }
            else
            {
                _exception = exception;
            }
        }
        finally
        {
            if (lockTaken)
                _lock.Exit();
        }
    }

    public void SetResult(T result)
    {
        var lockTaken = false;
        try
        {
            _lock.Enter(ref lockTaken);
            if (Volatile.Read(ref _source) is {} source)
            {
                source.TrySetResult(result);
            }
            else
            {
                _result = result;
                _hasResult = true;
            }
        }
        finally
        {
            if (lockTaken)
                _lock.Exit();
        }
    }

    public void SetStateMachine(IAsyncStateMachine stateMachine) {}

    public void Start<TStateMachine>(ref TStateMachine stateMachine)
        where TStateMachine : IAsyncStateMachine => stateMachine.MoveNext();
}

I've noticed (by debugging this code in Release mode and setting lots of breakpoints) that when it's a struct in Release mode then this sequence of events happens:

  1. CustomAwaitableAsyncMethodBuilder<T>.Create() is called exactly once
  2. CustomAwaitableAsyncMethodBuilder<T>.Task property is accessed
  3. CustomAwaitableAsyncMethodBuilder<T>.SetResult(T) method is invoked
  4. The Task property and the SetResult method both execute concurrently
  5. The SpinLock does not mutually exclude the respective blocks of code (they end up racing)
  6. The effects of the code in the Task property are not visible to the code in the SetResult method, nor vice versa

#5 and #6 tell me that the Task property is accessed on one instance of the struct while the SetResult method is called on a different instance.

Why is that? What am I doing wrong?


I don't see much information out there for how exactly to implement an async method builder. The only things I have found are:

I've followed Microsoft's documentation. It even states

The builder type is a class or struct

...and also Microsoft's AsyncTaskMethodBuilder<TResult> is a struct.

There is one place where their documentation is incorrect (it states that the builder's "AwaitUnsafeOnCompleted() [method] should call awaiter.OnCompleted(action)", but AsyncTaskMethodBuilder doesn't do that). But other than that I assume their documentation is correct.

The only difference I can spot between my implementation and the two other implementations linked above is that mine stores result/exception inside the builder type itself, where they always delegate to a reference type to do that job.


Edit:

ILSpy says this is how my Main method's state machine is implemented:

// BrokenAsyncMethodBuilder.Program.<Main>d__0
using System;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading.Tasks;

[StructLayout(LayoutKind.Auto)]
[CompilerGenerated]
private struct <Main>d__0 : IAsyncStateMachine
{
    public int <>1__state;

    public AsyncTaskMethodBuilder <>t__builder;

    private <>c__DisplayClass0_0 <>8__1;

    private ValueTaskAwaiter<string> <>u__1;

    private void MoveNext()
    {
        int num = <>1__state;
        try
        {
            ValueTaskAwaiter<string> awaiter;
            if (num != 0)
            {
                <>8__1 = new <>c__DisplayClass0_0();
                <>8__1.expected = Guid.NewGuid().ToString();
                awaiter = GetValueAsync().GetAwaiter();
                if (!awaiter.IsCompleted)
                {
                    num = (<>1__state = 0);
                    <>u__1 = awaiter;
                    <>t__builder.AwaitUnsafeOnCompleted(ref awaiter, ref this);
                    return;
                }
            }
            else
            {
                awaiter = <>u__1;
                <>u__1 = default(ValueTaskAwaiter<string>);
                num = (<>1__state = -1);
            }
            string actual = awaiter.GetResult();
            if ((object)<>8__1.expected != actual)
            {
                throw new Exception();
            }
            Console.WriteLine("Done!");
        }
        catch (Exception exception)
        {
            <>1__state = -2;
            <>8__1 = null;
            <>t__builder.SetException(exception);
            return;
        }
        <>1__state = -2;
        <>8__1 = null;
        <>t__builder.SetResult();
        async CustomAwaitable<string> GetValueAsync()
        {
            await Task.Yield();
            return ((<>c__DisplayClass0_0)(object)this).expected;
        }
    }

    void IAsyncStateMachine.MoveNext()
    {
        //ILSpy generated this explicit interface implementation from .override directive in MoveNext
        this.MoveNext();
    }

    [DebuggerHidden]
    private void SetStateMachine(IAsyncStateMachine stateMachine)
    {
        <>t__builder.SetStateMachine(stateMachine);
    }

    void IAsyncStateMachine.SetStateMachine(IAsyncStateMachine stateMachine)
    {
        //ILSpy generated this explicit interface implementation from .override directive in SetStateMachine
        this.SetStateMachine(stateMachine);
    }
}

And this is how my GetValueAsync method's state machine is implemented:

// BrokenAsyncMethodBuilder.Program.<>c__DisplayClass0_0.<<Main>g__GetValueAsync|0>d
using System;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading.Tasks;

[StructLayout(LayoutKind.Auto)]
private struct <<Main>g__GetValueAsync|0>d : IAsyncStateMachine
{
    public int <>1__state;

    public CustomAwaitableAsyncMethodBuilder<string> <>t__builder;

    public <>c__DisplayClass0_0 <>4__this;

    private YieldAwaitable.YieldAwaiter <>u__1;

    private void MoveNext()
    {
        int num = <>1__state;
        <>c__DisplayClass0_0 <>c__DisplayClass0_ = <>4__this;
        string expected;
        try
        {
            YieldAwaitable.YieldAwaiter awaiter;
            if (num != 0)
            {
                awaiter = Task.Yield().GetAwaiter();
                if (!awaiter.IsCompleted)
                {
                    num = (<>1__state = 0);
                    <>u__1 = awaiter;
                    <>t__builder.AwaitUnsafeOnCompleted(ref awaiter, ref this);
                    return;
                }
            }
            else
            {
                awaiter = <>u__1;
                <>u__1 = default(YieldAwaitable.YieldAwaiter);
                num = (<>1__state = -1);
            }
            awaiter.GetResult();
            expected = <>c__DisplayClass0_.expected;
        }
        catch (Exception exception)
        {
            <>1__state = -2;
            <>t__builder.SetException(exception);
            return;
        }
        <>1__state = -2;
        <>t__builder.SetResult(expected);
    }

    void IAsyncStateMachine.MoveNext()
    {
        //ILSpy generated this explicit interface implementation from .override directive in MoveNext
        this.MoveNext();
    }

    [DebuggerHidden]
    private void SetStateMachine(IAsyncStateMachine stateMachine)
    {
        <>t__builder.SetStateMachine(stateMachine);
    }

    void IAsyncStateMachine.SetStateMachine(IAsyncStateMachine stateMachine)
    {
        //ILSpy generated this explicit interface implementation from .override directive in SetStateMachine
        this.SetStateMachine(stateMachine);
    }
}

I'm unable to find this call into CustomAwaitableAsyncMethodBuilder<T>.Create(), other than it rather unhelpfully says it is used by the Main method (which decompiles to be very similar to the code at the top of this question). Perhaps I'm not familiar enough with ILSpy, or perhaps it has a bug:

Screenshot of call stack showing call into the Create method

Gain answered 30/9, 2021 at 17:57 Comment(2)
"If the state machine is implemented as a struct, then builder.SetStateMachine(stateMachine) is called with a boxed instance of the state machine that the builder can cache if necessary" - does that happen?Merbromin
@Merbromin I can confirm with breakpoints that it does not. I get the feeling the SetStateMachine method is deprecated. See the method that gets called by Microsoft's implementation of that method--all paths either throw or Debug.Fail. Also the comments in the link mention something something legacyGain
G
10

Found it! If you use ILSpy to disassemble the .dll compiled from the question's code (use the .NET Fiddle link and follow the question's instructions), and then turn ILSpy's language version down to C# 4 (which was the version before async/await was introduced), then you'll see that this is how the GetValueAsync method is implemented:

// BrokenAsyncMethodBuilder.Program.<>c__DisplayClass0_0
using System.Runtime.CompilerServices;

[AsyncStateMachine(typeof(<<Main>g__GetValueAsync|0>d))]
[return: System.Runtime.CompilerServices.Nullable(new byte[] { 0, 1 })]
internal CustomAwaitable<string> <Main>g__GetValueAsync|0()
{
    <<Main>g__GetValueAsync|0>d stateMachine = default(<<Main>g__GetValueAsync|0>d);
    stateMachine.<>t__builder = CustomAwaitableAsyncMethodBuilder<string>.Create();
    stateMachine.<>4__this = this;
    stateMachine.<>1__state = -1;
    stateMachine.<>t__builder.Start(ref stateMachine);
    return stateMachine.<>t__builder.Task;
}

You can reference the disassembly at the end of the question to learn how the <<Main>g__GetValueAsync|0>d type (the state machine for the GetValueAsync method) is implemented. You'll notice that the stateMachine.<>t__builder field is the CustomAwaitableAsyncMethodBuilder type (the async method builder in the question).

Now pay attention to what happens:

  1. The async method builder is created
  2. Its Start method is invoked with a reference to the state machine
    1. That calls .MoveNext() on the state machine
    2. MoveNext calls AwaitUnsafeOnCompleted() on the async method builder
    3. The async method builder calls UnsafeOnCompleted() on the awaiter (the awaiter is a YieldAwaitable.YieldAwaiter), giving it the state machine's MoveNext method as a delegate
    4. YieldAwaitable.YieldAwaiter.UnsafeOnCompleted posts the MoveNext delegate to the current synchronization context to be executed in a little bit
  3. The async method builder's Task property is accessed and returned

Notice also how the next time that the state machine's MoveNext method is invoked, it will call SetResult on the async method builder.

So at this point there is a delegate for the state machine's MoveNext method floating around in the synchronization context. Which means a copy of the state machine has been boxed up and placed on the heap. And the state machine holds a CustomAwaitableAsyncMethodBuilder, which is itself a struct. So it gets copied too. If CustomAwaitableAsyncMethodBuilder were a class then it would live on the heap and its reference would be copied instead of its value.

So by the time the GetValueAsync method returns, there will be a boxed instance of CustomAwaitableAsyncMethodBuilder on the heap, and another instance of CustomAwaitableAsyncMethodBuilder on the stack.

I do not know why this does not happen in debug builds. Nor do I yet understand what to do about this.

But this explains why "the Task property is accessed on one instance of the struct while the SetResult method is called on a different instance."


Thinking about this a little more, I think this is just another disguise for the fact that C#'s async/await system requires heap allocation and dynamic dispatch by design.

I wrote about this in a little depth not long ago as I compared C#'s async/await with Rust's implementation. In that article I pointed out how C#'s awaitable expressions are composed with dynamic dispatch on heap objects (the compiler packages up a continuation/state machine step as an Action, and Action is a reference type). But in Rust, awaitable expressions are composed with static dispatch (which doesn't necessarily require the heap).

This is exactly the same thing, just looking at it from the async method builder's perspective instead of the awaitable expression's perspective. Continuations are packaged up as state machine steps (the MoveNext method). It doesn't matter if the state machine is itself a value type; it will always get boxed when you treat its MoveNext method as an Action delegate. And in this case, boxing the state machine struct also copied the async method builder struct.

Gain answered 30/9, 2021 at 19:24 Comment(2)
So it's either a bug or a case of a grossly outdated documentation? Worth bringing it up with the devs in any case, in my opinion.Merbromin
@Merbromin Good ideaGain
E
0

I’ve written one of these before, and it’s a pain to say the least. One thing that surprised me is that the Task that is originally created to do some work, and the Task that is used to return the result are not bound together in any way, hence why the FromResult/FromException/FromCancellation methods are used to complete another task upon builder execution end. Also, I always got an Unhandled Exception (only in Release mode, not Debug) when I called Awaiter.OnCompleted(stateMachine.MoveNext), so I ended up calling that MoveNext method directly from the builder/state machine itself.

The only reason I made my own builder class is to identify hotspots in my code, where ValueTask would be better used for “fast path” task executions (to avoid the allocations). That is, I would log CallerMemberName, CallerFilePath, CallerLine Number, and where that task resulted in a fast path or a slow path execution. At first, it was slower, of course, doing all the locking for the logging of this telemetry, but after telemetry was collected, I put a backdown threshold to reduce the telemetry entries to only once every 10 invocations, and then it got faster than the “value task or task guessing” strategy used by junior devs. Ideally, they would know not to always use Task, but many dont. It took me over a week to write it, and it doesnt have all the methods a .NET Task does, but I add them as needed.

Ultimately, I’ve resolved that I probably should send my codebase to Steven Toub so he can add it to the framework. He’s 100x better than me at the C++ that would be needed to truly make it feasible for a production (not just an optimization exercise) runtime environment.

I’d love to see the code you finally came up with for your solution. Some of it still confuses me, especially that unhandled exception part. Btw, I wrote this on dotnetfiddle, so I didnt have the benefit of stepping into the code.

Eldin answered 3/7, 2024 at 18:25 Comment(1)
"I’d love to see the code you finally came up with for your solution." IIRC this was an academic exercise for me because I wanted to learn more about async/await in C#. At my employer we haven't yet had the need for a custom async method builder. The most advanced thing we've needed in this area is a custom awaitable type, and even that is only for convenienceGain

© 2022 - 2025 — McMap. All rights reserved.