aspnetcore/src/Shared/ObjectMethodExecutor/ObjectMethodExecutorAwaitable.cs

// Copyright (c) .NET Foundation. All rights reserved.
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.

using System;
using System.Runtime.CompilerServices;

namespace Microsoft.Extensions.Internal
{
    /// <summary>
    /// Provides a common awaitable structure that <see cref="ObjectMethodExecutor.ExecuteAsync"/> can
    /// return, regardless of whether the underlying value is a System.Task, an FSharpAsync, or an
    /// application-defined custom awaitable.
    /// </summary>
    internal struct ObjectMethodExecutorAwaitable
    {
        private readonly object _customAwaitable;
        private readonly Func<object, object> _getAwaiterMethod;
        private readonly Func<object, bool> _isCompletedMethod;
        private readonly Func<object, object> _getResultMethod;
        private readonly Action<object, Action> _onCompletedMethod;
        private readonly Action<object, Action> _unsafeOnCompletedMethod;

        // Perf note: since we're requiring the customAwaitable to be supplied here as an object,
        // this will trigger a further allocation if it was a value type (i.e., to box it). We can't
        // fix this by making the customAwaitable type generic, because the calling code typically
        // does not know the type of the awaitable/awaiter at compile-time anyway.
        //
        // However, we could fix it by not passing the customAwaitable here at all, and instead
        // passing a func that maps directly from the target object (e.g., controller instance),
        // target method (e.g., action method info), and params array to the custom awaiter in the
        // GetAwaiter() method below. In effect, by delaying the actual method call until the
        // upstream code calls GetAwaiter on this ObjectMethodExecutorAwaitable instance.
        // This optimization is not currently implemented because:
        // [1] It would make no difference when the awaitable was an object type, which is
        //     by far the most common scenario (e.g., System.Task<T>).
        // [2] It would be complex - we'd need some kind of object pool to track all the parameter
        //     arrays until we needed to use them in GetAwaiter().
        // We can reconsider this in the future if there's a need to optimize for ValueTask<T>
        // or other value-typed awaitables.

        public ObjectMethodExecutorAwaitable(
            object customAwaitable,
            Func<object, object> getAwaiterMethod,
            Func<object, bool> isCompletedMethod,
            Func<object, object> getResultMethod,
            Action<object, Action> onCompletedMethod,
            Action<object, Action> unsafeOnCompletedMethod)
        {
            _customAwaitable = customAwaitable;
            _getAwaiterMethod = getAwaiterMethod;
            _isCompletedMethod = isCompletedMethod;
            _getResultMethod = getResultMethod;
            _onCompletedMethod = onCompletedMethod;
            _unsafeOnCompletedMethod = unsafeOnCompletedMethod;
        }

        public Awaiter GetAwaiter()
        {
            var customAwaiter = _getAwaiterMethod(_customAwaitable);
            return new Awaiter(customAwaiter, _isCompletedMethod, _getResultMethod, _onCompletedMethod, _unsafeOnCompletedMethod);
        }

        public struct Awaiter : ICriticalNotifyCompletion
        {
            private readonly object _customAwaiter;
            private readonly Func<object, bool> _isCompletedMethod;
            private readonly Func<object, object> _getResultMethod;
            private readonly Action<object, Action> _onCompletedMethod;
            private readonly Action<object, Action> _unsafeOnCompletedMethod;

            public Awaiter(
                object customAwaiter,
                Func<object, bool> isCompletedMethod,
                Func<object, object> getResultMethod,
                Action<object, Action> onCompletedMethod,
                Action<object, Action> unsafeOnCompletedMethod)
            {
                _customAwaiter = customAwaiter;
                _isCompletedMethod = isCompletedMethod;
                _getResultMethod = getResultMethod;
                _onCompletedMethod = onCompletedMethod;
                _unsafeOnCompletedMethod = unsafeOnCompletedMethod;
            }

            public bool IsCompleted => _isCompletedMethod(_customAwaiter);

            public object GetResult() => _getResultMethod(_customAwaiter);

            public void OnCompleted(Action continuation)
            {
                _onCompletedMethod(_customAwaiter, continuation);
            }

            public void UnsafeOnCompleted(Action continuation)
            {
                // If the underlying awaitable implements ICriticalNotifyCompletion, use its UnsafeOnCompleted.
                // If not, fall back on using its OnCompleted.
                //
                // Why this is safe:
                // - Implementing ICriticalNotifyCompletion is a way of saying the caller can choose whether it
                //   needs the execution context to be preserved (which it signals by calling OnCompleted), or
                //   that it doesn't (which it signals by calling UnsafeOnCompleted). Obviously it's faster *not*
                //   to preserve and restore the context, so we prefer that where possible.
                // - If a caller doesn't need the execution context to be preserved and hence calls UnsafeOnCompleted,
                //   there's no harm in preserving it anyway - it's just a bit of wasted cost. That's what will happen
                //   if a caller sees that the proxy implements ICriticalNotifyCompletion but the proxy chooses to
                //   pass the call on to the underlying awaitable's OnCompleted method.

                var underlyingMethodToUse = _unsafeOnCompletedMethod ?? _onCompletedMethod;
                underlyingMethodToUse(_customAwaiter, continuation);
            }
        }
    }
}