Suspected CodeGen Bug in a User-Defined AsyncMethodBuilder

[LEI-Hongfaan]

Description

There appears to be a CodeGen bug within the SuspendTaskAsyncMethodBuilder.Start method. The assignment _Task = task; is erroneously emitted as a no-op. Note that changing SuspendTaskAsyncMethodBuilder from a value type (struct) to a reference type (class) results in the program running correctly.

Reproduction Steps

using MethodBuilders;
using System.Diagnostics;
using System.Runtime.CompilerServices;

namespace MethodBuilders {

    public struct SuspendTaskAsyncMethodBuilder<T> {

        public void SetStateMachine(IAsyncStateMachine stateMachine) {
            throw new NotImplementedException();
        }

        public static SuspendTaskAsyncMethodBuilder<T> Create()
            => new();

        public void SetResult(T reresult) {
        }

        public void SetException(Exception exception) {
            throw exception;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void Start<TStateMachine>(ref TStateMachine stateMachine)
            where TStateMachine : IAsyncStateMachine {
            var task = SuspendTask<T>.Create(ref stateMachine);
            _Task = task;
        }

        SuspendTask<T> _Task;

        public SuspendTask<T> Task => _Task;

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void AwaitOnCompleted<TAwaiter, TStateMachine>(
            ref TAwaiter awaiter, ref TStateMachine stateMachine)
            where TAwaiter : INotifyCompletion
            where TStateMachine : IAsyncStateMachine {
            throw new NotImplementedException();
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(
            ref TAwaiter awaiter, ref TStateMachine stateMachine)
            where TAwaiter : ICriticalNotifyCompletion
            where TStateMachine : IAsyncStateMachine {
            if (typeof(TAwaiter) == typeof(SuspendUninterceptedAwaiter<T>)) {
                ref SuspendUninterceptedAwaiter<T> suspendUninterceptedAwaiter = ref Unsafe.As<TAwaiter, SuspendUninterceptedAwaiter<T>>(ref awaiter);
                suspendUninterceptedAwaiter.continuation.result = _Task.continuation.result;
                return;
            }
        }
    }

    public struct SuspendTaskAwaiter<T> : INotifyCompletion {

        public bool IsCompleted {

            get => true;
        }

        public T GetResult() {
            return default;
        }

        public void OnCompleted(Action continuation) {
            throw new NotImplementedException();
        }
    }

    [AsyncMethodBuilder(typeof(SuspendTaskAsyncMethodBuilder<>))]
    public readonly struct SuspendTask<T> {

        internal readonly Continuation<T> continuation;

        public SuspendTask(Continuation<T> continuation) {
            this.continuation = continuation;
        }

        public SuspendTaskAwaiter<T> GetAwaiter() {
            return default;
        }

        public Continuation<T> createCoroutineUnintercepted() {
            return continuation;
        }

        internal static SuspendTask<T> Create<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine {
            return new SuspendTask<T>(new Continuation<T, TStateMachine>(ref stateMachine));
        }
    }

    public readonly struct SuspendUninterceptedAwaiter<T> : INotifyCompletion, ICriticalNotifyCompletion {

        internal readonly Continuation<T> continuation;

        public SuspendUninterceptedAwaiter(Continuation<T> continuation) : this() {
            this.continuation = continuation;
        }

        public bool IsCompleted {

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            get {
                return false;
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public T GetResult() {
            return continuation.result.ValueOrDefault;
        }

        public void OnCompleted(Action continuation) {
            throw new NotImplementedException();
        }

        public void UnsafeOnCompleted(Action continuation) {
            throw new NotImplementedException();
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public SuspendUninterceptedAwaiter<T> GetAwaiter() {
            return this;
        }
    }

    public struct Unit {
    }

    public readonly struct Result<T> {

        public readonly bool isFailure;

        public readonly Exception ExceptionOrDefault;

        public readonly T? ValueOrDefault;

        public bool isSuccess {

            get => !isFailure;
        }

        public T Value {

            get => isFailure ? throw new InvalidOperationException() : Value;
        }

        public T Exception {

            get => isFailure ? Exception : throw new InvalidOperationException();
        }

        public Exception? exceptionOrNull() {
            return ExceptionOrDefault;
        }

        public Result(T value) : this() {
            ValueOrDefault = value;
            isFailure = false;
        }

        public Result(Exception exception) : this() {
            ExceptionOrDefault = exception ?? throw new ArgumentNullException(nameof(exception));
            isFailure = true;
        }
    }

    static class Result {

        public static Result<T> success<T>(T value) {
            return new Result<T>(value);
        }
    }

    public abstract class Continuation<T> {

        internal Result<T> result;

        public void resumeWith(Result<T> result) {
            if (this.result.isFailure) {
                throw new InvalidOperationException();
            }
            if (result.isFailure) {
                throw result.ExceptionOrDefault;
            }
            this.result = result;
            MoveNextCore();
        }

        protected abstract void MoveNextCore();
    }

    public readonly struct COROUTINE_SUSPENDED_T {
    }

    public static class Suspend {

        public static readonly Unit Unit = default;

        public static readonly COROUTINE_SUSPENDED_T COROUTINE_SUSPENDED = default;

        public static SuspendUninterceptedAwaiter<T> suspendCoroutineUninterceptedOrReturn<T>(this PlaceholderContinuation<T> placeholder, COROUTINE_SUSPENDED_T COROUTINE_SUSPENDED) {
            return new SuspendUninterceptedAwaiter<T>(placeholder);
        }
    }

    public sealed class PlaceholderContinuation<T> : Continuation<T> {

        public new void resumeWith(Result<T> result) {
            throw new NotSupportedException();
        }

        protected override void MoveNextCore() {
            throw new NotSupportedException();
        }
    }

    public sealed class Continuation<T, TStateMachine> : Continuation<T>
        where TStateMachine : IAsyncStateMachine {

        internal TStateMachine stateMachine;

        public Continuation(ref TStateMachine stateMachine) {
            this.stateMachine = stateMachine;
        }

        protected override void MoveNextCore() {
            stateMachine.MoveNext();
        }
    }
}


namespace ThisAssembly {
    using static MethodBuilders.Suspend;

    class Program {

        const int N = 1_000_000;

        static int s_i = 0;

        static async SuspendTask<Unit> f_() {
            var __ = new PlaceholderContinuation<Unit>();
            for (var i = 0; i <= N - 1; ++i) {
                s_i = i;
                await __.suspendCoroutineUninterceptedOrReturn<Unit>(COROUTINE_SUSPENDED);
            }
            return default;
        }

        static SuspendTask<Unit> f = f_();

        static void Main(string[] args) {
            var t = Stopwatch.GetTimestamp();
            var c = f.createCoroutineUnintercepted();
            var s = Result.success(Unit);
            var r = 0L;
            for (var i = 0; i <= N - 1; ++i) {
                c.resumeWith(s);
                r += s_i;
            }
            Console.WriteLine($"{Stopwatch.GetElapsedTime(t)}, r = {r}");
        }
    }
}

Expected behavior

00:00:00.0235162, r = 499999500000

Actual behavior

Unhandled exception. System.NullReferenceException: Object reference not set to an instance of an object.

Regression?

No response

Known Workarounds

No response

Configuration

No response

Other information

No response

[stephentoub]

Why do you believe there's a codegen bug? You say "_Task = task;" is a nop, but is what you really mean that it's not actually being stored into your state machine? I would expect that's because you're storing it into the wrong one. By the time you get to the _Task = task line, you've already copied the state machine into your Continuation type... the _Task field you're storing into isn't the same one as is in the Continuation.

(In https://devblogs.microsoft.com/dotnet/how-async-await-really-works/, search for "Now comes a somewhat mind-bending step" for related discussion.)

[LEI-Hongfaan]

Why do you believe there's a codegen bug? You say "_Task = task;" is a nop, but is what you really mean that it's not actually being stored into your state machine? I would expect that's because you're storing it into the wrong one. By the time you get to the _Task = task line, you've already copied the state machine into your Continuation type... the _Task field you're storing into isn't the same one as is in the Continuation.

(In https://devblogs.microsoft.com/dotnet/how-async-await-really-works/, search for "Now comes a somewhat mind-bending step" for related discussion.)

I've observed that this issue generates new code in the newer versions of the .NET runtime. If this were a programming error, I would expect to see the corresponding code for _Task = task; in the disassembly when stepping into this method during debugging (with the older runtime), correct?

However, with the newer runtime, I've found that there are no issues under the debug configuration, but the problem persists under the release configuration. You might be right regarding the specific implementation issues of AsyncMethodBuilder that you mentioned. But the current phenomenon indeed reflects a bug.

[huoyaoyuan]

Tested with 8.0.6. The same piece of code works fine in Debug configuration, but throws NRE in Release. Probably a codegen bug.

[dotnet-policy-service]

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

[jakobbotsch]

The underlying state machine is emitted as a class by Roslyn under debug configuration. In release, it is emitted as a struct. Thus, Continuation<T, TStateMachine> has a reference to the state machine in debug, while it has a copy of the state machine in release. I expect that's the reason why it seems to work in debug, while the underlying issue is what @stephentoub described above.

I don't see any JIT issues here -- in debug no NRE is thrown. In release, the NRE is thrown whether or not the JIT is optimizing (DOTNET_JITMinOpts=1 disables optimizations). If this was a JIT bug I would expect JIT optimizations to affect the behavior when running the release version.

[dotnet-policy-service]

Tagging subscribers to this area: @dotnet/area-system-threading-tasks
See info in area-owners.md if you want to be subscribed.

[teo-tsirpanis]

FWIW I had encountered a very similar bug with a custom async method builder that fails only in Release mode, which I fixed in teo-tsirpanis/uom@98c951f.

From what @stephentoub said, it's likely a matter of wrong assignment order in SuspendTaskAsyncMethodBuilder<T>.Start(). You must in this order:

1. Create the SuspendTask<T>.
2. Assign it to the async method builder's _Task field.
3. Assign the state machine to the SuspendTask<T>.

Doing this might require some additional refactoring in your code.

[dotnet-policy-service]

This issue has been marked needs-author-action and may be missing some important information.