API Proposal: AsyncMethodBuilderOverrideAttribute and PoolingAsyncValueTaskMethodBuilders

[stephentoub]

Background

Async Method Builder Overrides

The language compiler needs to know how to “build” the return type of an async method, e.g. if you have the method:

public async ValueTask<int> MethodAsync() { … }

the compiler transforms the method into a state machine, and it needs to know how to produce a ValueTask<int> that represents that method’s operation, how to create the task, how to store a result into the task, how to store an exception into the task, and when to do all of that. This is a achieve via an “async method builder”. To know which builder to use, the type used in the return position needs to be annotated to say which builder to use (Task and Task are exempted and explicitly known to the compiler). For example, ValueTask<T> is declared like this:

[AsyncMethodBuilder(typeof(AsyncValueTaskMethodBuilder<>))]
public readonly struct ValueTask<TResult> : IEquatable<ValueTask<TResult>>

With that information (and with AsyncValueTaskMethodBuilder<> having the right shape), the compiler allows ValueTask<TResult> to be used as the return type of an async method, and it uses that builder type as part of its generated code, e.g.

    public ValueTask<int> MethodAsync()
    {
        <MethodAsync>d__0 stateMachine = default(<MethodAsync>d__0);
        stateMachine.<>t__builder = AsyncValueTaskMethodBuilder<int>.Create();
        stateMachine.<>1__state = -1;
        stateMachine.<>t__builder.Start(ref stateMachine);
        return stateMachine.<>t__builder.Task;
    }

This, however, means there’s no way to override the builder logic while still using the same return type: every method that returns a ValueTask<T> uses the same builder, the one built-in to the runtime, which means there’s no way for someone to customize the behavior. To customize it, a developer must use a different return type, which a) means an APIs public signature ends up being impacted by implementation detail, b) means we start to see additional task-like types being promoted (even though we want standardization on {Value}Task{<T>} and even though we have optimizations in place that cater to all awaits being for those types), and c) means we lack the ability for developers to mostly utilize our builder logic but tweak it based on their needs.

To address this, in C# 10 we’re adding the ability for (an advanced) developer to override the builder that’s used. This is achieved with an attribute:

[AsyncMethodBuilderOverride(typeof(MyCustomBuilder))]

The attribute can be applied not only at the method level, but at larger scopes as well, and applies to every async method defined in that scope. This extends all the way up to the module. If I as a developer wants every ValueTask and ValueTask<T> method in my assembly to use my custom builder, that’s achievable with two lines in a .cs file somewhere in the project:

[module: AsyncMethodBuilderOverride(typeof(MyCustomBuilder))]
[module: AsyncMethodBuilderOverride(typeof(MyCustomBuilder<>))]

We need to add that attribute to the core libraries, both so developers can use the feature and so that we can in the core libraries as well.

Pooling Builder for ValueTask{<T>}

.NET 5 included an experimental environment variable that changed the default builder used for async ValueTask{<T>} to be one that pools. With such a pooling builder, you can get amortized allocation-free async methods, something developers have been very interested in (to the point where some have defined their own return type just so that they can pool it).

For multiple reasons, however, it’s not tenable for us to make the pooling behavior the default, and for .NET 6 we’ll be removing the environment variable and option. There are a few main issues:

1. Compatibility. ValueTask has various constraints on how it’s meant to be consumed, but when constructed in certain ways, as an implementation detail it can support more than that. For example, if you construct a ValueTask around a Task, it’ll let you hook up multiple continuations even though ValueTask is only ever supposed to be awaited once. This means by switching every async ValueTask automatically, we could easily break code that doesn’t abide by the rules but that’s getting away with it. We do have an analyzer that flags some such misuse, but that only gets us so far.
2. Overhead. Pooling has overhead. It’s very unlikely that making every async method into one that pools is the correct tradeoff.
3. Scope. Related to both (1) and (2), the environment variable ends up being a process-wide setting, impacting not just your own code but any libraries you load into the process. If that library itself makes poor assumptions about how it consumes ValueTasks returned from its own async methods, it can be broken.
4. Size. To enable the environment variable, all of the code that supports both non-pooling and pooling is included in every app, which affects not only IL size, but more so AOT size, if every async method ends up needing twice as much assembly in support of the option.

With support for AsyncMethodBuilderOverride, we can move the pooling support into its own builders. A developer can then opt-in to using the pooling at whatever scope is most applicable to their scenario:

[AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))]
public readonly struct ValueTask<TResult> : IEquatable<ValueTask<TResult>>

We will use this surgically in our core libraries on places where end-to-end benchmarks show it to be impactful. As with AggressiveInlining, we’ll allow use of this, but PRs that apply it will need justification for its use, as it comes with tradeoffs.

Proposed API

namespace System.Runtime.CompilerServices
{
    // The shape of this attribute is the same as AsyncMethodBuilderAttribute, just with "Override" in the name
    // and different attribute targets.

    [AttributeUsage(AttributeTargets.Class | AttributeTargets.Struct | AttributeTargets.Interface | AttributeTargets.Method | AttributeTargets.Constructor | AttributeTargets.Event | AttributeTargets.Property | AttributeTargets.Module, Inherited = false, AllowMultiple = true)]
    public sealed partial class AsyncMethodBuilderOverrideAttribute : System.Attribute
    {
        public AsyncMethodBuilderOverrideAttribute(Type builderType);
        public Type BuilderType { get; }
    }

    // The shape of the below two types is exactly the same as AsyncValueTaskMethodBuilder{<T>},
    // (and is dictated by the compiler's requirements for the task-like pattern), just with a "Pooling"
    // prefix on the type names

    public struct PoolingAsyncValueTaskMethodBuilder
    {
        public static PoolingAsyncValueTaskMethodBuilder Create();
        public ValueTask Task { get; }
        public void SetResult();
        public void SetException(Exception exception);
        public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void Start<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine;
        public void SetStateMachine(IAsyncStateMachine stateMachine);
    }

    public struct PoolingAsyncValueTaskMethodBuilder<TResult>
    {
        public static PoolingAsyncValueTaskMethodBuilder<TResult> Create();
        public ValueTask<TResult> Task { get; }
        public void SetResult(TResult result);
        public void SetException(Exception exception);
        public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : System.Runtime.CompilerServices.ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void SetStateMachine(IAsyncStateMachine stateMachine);
        public void Start<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine;
    }
}

Adding these APIs and deleting the environment variable will serve to close #13633.

(After some experience with the API, we may want to add a few more overloads of Create, that would enable some level of control over pooling behavior, which developers could leverage by writing their own lightweight builder that just provides a Create method and returns one of these pooling builders constructed with one of those new overloads.)

cc: @jcouv

[ghost]

Tagging subscribers to this area:
See info in area-owners.md if you want to be subscribed.

Issue Details

---

Background

Async Method Builder Overrides

The language compiler needs to know how to “build” the return type of an async method, e.g. if you have the method:

public async ValueTask<int> MethodAsync() { … }

the compiler transforms the method into a state machine, and it needs to know how to produce a ValueTask<int> that represents that method’s operation, how to create the task, how to store a result into the task, how to store an exception into the task, and when to do all of that. This is a achieve via an “async method builder”. To know which builder to use, the type used in the return position needs to be annotated to say which builder to use (Task and Task are exempted and explicitly known to the compiler). For example, ValueTask<T> is declared like this:

[AsyncMethodBuilder(typeof(AsyncValueTaskMethodBuilder<>))]
public readonly struct ValueTask<TResult> : IEquatable<ValueTask<TResult>>

With that information (and with AsyncValueTaskMethodBuilder<> having the right shape), the compiler allows ValueTask<TResult> to be used as the return type of an async method, and it uses that builder type as part of its generated code, e.g.

    public ValueTask<int> MethodAsync()
    {
        <MethodAsync>d__0 stateMachine = default(<MethodAsync>d__0);
        stateMachine.<>t__builder = AsyncValueTaskMethodBuilder<int>.Create();
        stateMachine.<>1__state = -1;
        stateMachine.<>t__builder.Start(ref stateMachine);
        return stateMachine.<>t__builder.Task;
    }

This, however, means there’s no way to override the builder logic while still using the same return type: every method that returns a ValueTask<T> uses the same builder, the one built-in to the runtime, which means there’s no way for someone to customize the behavior. To customize it, a developer must use a different return type, which a) means an APIs public signature ends up being impacted by implementation detail, b) means we start to see additional task-like types being promoted (even though we want standardization on {Value}Task{<T>} and even though we have optimizations in place that cater to all awaits being for those types), and c) means we lack the ability for developers to mostly utilize our builder logic but tweak it based on their needs.

To address this, in C# 10 we’re adding the ability for (an advanced) developer to override the builder that’s used. This is achieved with an attribute:

[AsyncMethodBuilderOverride(typeof(MyCustomBuilder))]

The attribute can be applied not only at the method level, but at larger scopes as well, and applies to every async method defined in that scope. This extends all the way up to the module. If I as a developer wants every ValueTask and ValueTask method in my assembly to use my custom builder, that’s achievable with two lines in a .cs file somewhere in the project:

[module: AsyncMethodBuilderOverride(typeof(MyCustomBuilder))]
[module: AsyncMethodBuilderOverride(typeof(MyCustomBuilder<>))]

We need to add that attribute to the core libraries, both so developers can use the feature and so that we can in the core libraries as well.

Pooling Builder for ValueTask{<T>}

.NET 5 included an experimental environment variable that changed the default builder used for async ValueTask{<T>} to be one that pools. With such a pooling builder, you can get amortized allocation-free async methods, something developers have been very interested in (to the point where some have defined their own return type just so that they can pool it).

For multiple reasons, however, it’s not tenable for us to make the pooling behavior the default, and for .NET 6 we’ll be removing the environment variable and option. There are a few main issues:

1. Compatibility. ValueTask has various constraints on how it’s meant to be consumed, but when constructed in certain ways, as an implementation detail it can support more than that. For example, if you construct a ValueTask around a Task, it’ll let you hook up multiple continuations even though ValueTask is only ever supposed to be awaited once. This means by switching every async ValueTask automatically, we could easily break code that doesn’t abide by the rules but that’s getting away with it. We do have an analyzer that flags some such misuse, but that only gets us so far.
2. Overhead. Pooling has overhead. It’s very unlikely that making every async method into one that pools is the correct tradeoff.
3. Scope. Related to both (1) and (2), the environment variable ends up being a process-wide setting, impacting not just your own code but any libraries you load into the process. If that library itself makes poor assumptions about how it consumes ValueTasks returned from its own async methods, it can be broken.
4. Size. To enable the environment variable, all of the code that supports both non-pooling and pooling is included in every app, which affects not only IL size, but more so AOT size, if every async method ends up needing twice as much assembly in support of the option.

With support for AsyncMethodBuilderOverride, we can move the pooling support into its own builders. A developer can then opt-in to using the pooling at whatever scope is most applicable to their scenario:

[AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))]
public readonly struct ValueTask<TResult> : IEquatable<ValueTask<TResult>>

We will use this surgically in our core libraries on places where end-to-end benchmarks show it to be impactful. As with AggressiveInlining, we’ll allow use of this, but PRs that apply it will need justification for its use, as it comes with tradeoffs.

Proposed API

namespace System.Runtime.CompilerServices
{
    [AttributeUsage(AttributeTargets.Class | AttributeTargets.Struct | AttributeTargets.Interface | AttributeTargets.Method | AttributeTargets.Constructor | AttributeTargets.Event | AttributeTargets.Property | AttributeTargets.Module, Inherited = false, AllowMultiple = true)]
    public sealed partial class AsyncMethodBuilderOverrideAttribute : System.Attribute
    {
        public AsyncMethodBuilderOverrideAttribute(Type builderType);
        public Type BuilderType { get; }
    }

    // The shape of the below two types is exactly the same as AsyncValueTaskMethodBuilder{<T>},
    // (and is dictated by the compiler's requirements for the task-like pattern), just with a "Pooling"
    // prefix on the type names

    public struct PoolingAsyncValueTaskMethodBuilder
    {
        public static PoolingAsyncValueTaskMethodBuilder Create();
        public ValueTask Task { get; }
        public void SetResult();
        public void SetException(Exception exception);
        public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void Start<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine;
        public void SetStateMachine(IAsyncStateMachine stateMachine);
    }

    public struct PoolingAsyncValueTaskMethodBuilder<TResult>
    {
        public static PoolingAsyncValueTaskMethodBuilder<TResult> Create();
        public ValueTask<TResult> Task { get; }
        public void SetResult(TResult result);
        public void SetException(Exception exception);
        public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : System.Runtime.CompilerServices.ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void SetStateMachine(IAsyncStateMachine stateMachine);
        public void Start<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine;
    }
}

(After some experience with the API, we may want to add a few more overloads of Create, that would enable some level of control over pooling behavior, which developers could leverage by writing their own lightweight builder that just provides a Create method and returns one of these pooling builders constructed with one of those new overloads.)

cc: @jcouv

Author:	stephentoub
Assignees:	-
Labels:	api-ready-for-review, area-System.Threading.Tasks, blocking
Milestone:	6.0.0

[terrajobst]

Video

• Looks good as proposed

namespace System.Runtime.CompilerServices
{
    [AttributeUsage(AttributeTargets.Class |
                    AttributeTargets.Struct |
                    AttributeTargets.Interface |
                    AttributeTargets.Method |
                    AttributeTargets.Constructor |
                    AttributeTargets.Event |
                    AttributeTargets.Property |
                    AttributeTargets.Module, Inherited = false, AllowMultiple = true)]
    public sealed partial class AsyncMethodBuilderOverrideAttribute : System.Attribute
    {
        public AsyncMethodBuilderOverrideAttribute(Type builderType);
        public Type BuilderType { get; }
    }
    public struct PoolingAsyncValueTaskMethodBuilder
    {
        public static PoolingAsyncValueTaskMethodBuilder Create();
        public ValueTask Task { get; }
        public void SetResult();
        public void SetException(Exception exception);
        public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void Start<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine;
        public void SetStateMachine(IAsyncStateMachine stateMachine);
    }
    public struct PoolingAsyncValueTaskMethodBuilder<TResult>
    {
        public static PoolingAsyncValueTaskMethodBuilder<TResult> Create();
        public ValueTask<TResult> Task { get; }
        public void SetResult(TResult result);
        public void SetException(Exception exception);
        public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : System.Runtime.CompilerServices.ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine;
        public void SetStateMachine(IAsyncStateMachine stateMachine);
        public void Start<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine;
    }
}

[stephentoub]

It turns out the scoping support leads to a whole bunch of unwanted behaviors at the language/compiler level. For C# 10, we're culling this back to just attributing the target method directly. So, we're not adding AsyncMethodBuilderOverride now, and instead just enabling the existing AsyncMethodBuilder attribute to be applied to methods. Should we want to bring back scoping in the future, we can bring back the override attribute or address ambiguities in other ways at that time.

[aelij]

For C# 10, we're culling this back to just attributing the target method directly

Too bad, I was hoping to (ab)use this feature to apply ConfigureAwait(false) at the module scope.

[stephentoub]

I expect that would be difficult to do well / efficiently, even if you could use it for that purpose. How would you propose to do it?

[aelij]

Something like:

    public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine
    {
        if (awaiter is TaskAwaiter<TResult>)
        {
            var configuredAwaiter = Unsafe.As<TAwaiter, TaskAwaiterStruct>(ref awaiter).Task.ConfigureAwait(false).GetAwaiter();
            _builder.AwaitOnCompleted(ref configuredAwaiter, ref stateMachine);
            return;
        }

        _builder.AwaitOnCompleted(ref awaiter, ref stateMachine);
    }

I've seen other AwaitOnCompleted use Unsafe.As so I expect it would be ok perf-wise. And JIT should optimize the branch.

[stephentoub]

if (awaiter is TaskAwaiter<TResult>)

TResult isn't the only possible generic type, though. That's the return type of this async method, but the body of the async method can await anything.

[aelij]

Ah, you're right of course :)