Proposal: ManualResetValueTaskSource{Logic} types

[stephentoub]

Related to https://github.com/dotnet/corefx/issues/32640.

Background

In .NET Core 2.1, we added the ability for a ValueTask<T> (and the non-generic ValueTask, which we also added) to be constructed not only from a T and Task<T>, but also from an IValueTaskSource<T>. The primary benefit of this is that it allows the ValueTask<T> to be backed by an object that can be reused in order to amortize the cost of the allocation. For example, also in .NET Core 2.1 we optimized the new ValueTask<int>-returning Socket.ReceiveAsync method, so that as long as it’s not used concurrently with itself on the same Socket instance, it won’t need to allocate a new promise instance to hand back, whether the operation completes synchronously (in which case it can just hand back a ValueTask<T> constructed from a T) or asynchronously (in which case it can hand back a ValueTask<T> constructed from a reusable object that implements IValueTaskSource<T>). However, IValueTaskSource<T> is an advanced interface that is not easy to implement, and the implementations we provided in .NET Core 2.1 were all internal as internal implementation details, e.g. in Socket, in System.Threading.Channels, and in System.IO.Pipelines.

In .NET Core 3.0, we’re adding support for IAsyncEnumerable<T>, and the C# compiler will be able to generate implementations for async iterators, just as it can for synchronous iterators today. We’ve crafted the async enumerator interface such that the various asynchronous methods return ValueTasks, which allows the compiler to allocate a single object that can then be reused for all asynchronous operations on that enumerable. In fact, that object can be the enumerable itself, which is also reused for the enumerator. In other words, by providing an object that appropriately implements IAsyncEnumerable<T>, IAsyncEnumerator<T>, and IValueTaskSource<bool>, the compiler can generate an async iterator with only a very minimal allocation overhead.

To do that, though, the compiler needs to implement IValueTaskSource<bool>, and as noted earlier, doing so is challenging, requiring a non-trivial amount of code. To make it easier on the compiler and on anyone else that wants to do similar things, we should provide an implementation of the core logic needed for this interface, such that implementing the interface is then as simple as delegating to the helpers provided.

Proposal

namespace System.Threading.Tasks.Sources // same namespace as IValueTaskSource
{
    public struct ManualResetValueTaskSourceLogic<TResult>
    {
        // Named the same as TaskCreationOptions.RunContinuationsAsynchronously.  Controls whether calls to Set* will invoke
        // continuations synchronously if possible, or whether continuations will always be forced to be queued to the thread pool.
        public bool RunContinuationsAsynchronously { get; set; }

        // Completes the instance.
        public void SetResult(TResult result);
        public void SetException(Exception error);

        // Resets the instance so that it can be used again.
        public void Reset();

        // Directly correspond to the IValueTaskSource interface, and as such an implementation of that interface can just delegate to these members.
        public short Version { get; }
        public ValueTaskSourceStatus GetStatus(short token);
        public TResult GetResult(short token);
        public void OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags);
    }
}

This mutable struct can then be used as a field on another object to help that object implement the IValueTaskSource interface. The C# compiler, for example, will implement the interface on its single enumerable/enumerator object, and that implementation will just delegate to the corresponding members of this type on a field on the object.

For example, here’s the full implementation of what a class-based ManualResetValueTaskSource<TResult> could look like when using this Logic type to implement it (we could optionally add this as well, now or later, if desired... developers that just want to be able to reuse an instance could use the friendlier class-based version, and developers that want to take it a step further and reuse an existing object could use the mutable struct-based version.)

namespace System.Threading.Tasks.Sources
{
    public class ManualResetValueTaskSource<T> : IValueTaskSource<T>, IValueTaskSource
    {
        private ManualResetValueTaskSourceLogic<T> _logic; // mutable struct; do not make this readonly

        public bool RunContinuationsAsynchronously { get => _logic.RunContinuationsAsynchronously; set => _logic.RunContinuationsAsynchronously = value; }
        public void Reset() => _logic.Reset();
        public void SetResult(T result) => _logic.SetResult(result);
        public void SetException(Exception error) => _logic.SetException(error);

        short IValueTaskSource.Version => _logic.Version;
        short IValueTaskSource<T>.Version => _logic.Version;

        void IValueTaskSource.GetResult(short token) => _logic.GetResult(token);
        T IValueTaskSource<T>.GetResult(short token) => _logic.GetResult(token);

        ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _logic.GetStatus(token);
        ValueTaskSourceStatus IValueTaskSource<T>.GetStatus(short token) => _logic.GetStatus(token);

        void IValueTaskSource.OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags) => _logic.OnCompleted(continuation, state, token, flags);
        void IValueTaskSource<T>.OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags) => _logic.OnCompleted(continuation, state, token, flags);
    }
}

To make the ManualResetValueTaskSourceLogic<TResult> work fully, we need some additional support, and for that I see three main options:

Option 1

Put ManualResetValueTaskSourceLogic<TResult> in System.Private.CoreLib. This will give it access to internals we can use on ExecutionContext.

Option 2

Add a public overload of ExecutionContext.Run. The type already provides:

public delegate void ContextCallback(object state);
…
public static void Run(ExecutionContext executionContext, ContextCallback callback, object state);

and we add the following overload to complement that:

public delegate void ActionRef<T>(ref T arg);
...
public static void Run<TState>(ExecutionContext executionContext, ActionRef<ref T> callback, ref TState state);

This is needed so that a mutable struct (like ManualResetValueTaskSourceLogic<TResult>) can have a method invoked under an execution context and have any changes made in Run affect the original mutable struct.

Option 3

Options 1 and 2 both depend on being able to augment ExecutionContext, with Option 1 doing it internally and Option 2 doing it externally. There’s another alternative that doesn’t require this, but that instead introduces a new interface and requires the type wrapping the struct to implement it, the combination of which allows the existing ExecutionContext to be used.

.NET already provides an System.Runtime.CompilerServices.IStrongBox interface, but it’s non-generic, with an object Value { get; set; } property. To properly support this option, we first need a new generic variation of this interface, that’s not only generic but that has a ref getter:

namespace System.Runtime.CompilerServices
{
    public interface IStrongBox<T>
    {
        ref T Value { get; }
    }
}

This interface makes it possible for any object to have a field of type T, and via the interface, hand back a reference to that field. Then, we’d add this constructor to the struct:

public ManualResetValueTaskSourceLogic(IStrongBox<ManualResetValueTaskSourceLogic<TResult>> parent);

and when the class with the ManualResetValueTaskSourceLogic as a field constructs the struct, it passes itself in. That way, the MRVTSL can pass that parent object around, such as to the existing overload of ExecutionContext.Run, and then access itself by ref via the strong box, e.g.

ExecutionContext.Run(
    _executionContext,
    s => ((IStrongBox<ManualResetValueTaskSourceLogic<TResult>>)s).Value.InvokeContinuation(),
    _parent);

I prefer Options 1 and 2, in part because it’s simpler, and in part because it doesn’t involve the implementing object to expose part of its implementation via its interfaces. For example, if the C# compiler just implements its enumerables to implement IStrongBox<ManualResetValueTaskSourceLogic<bool>>, then it’s effectively exposing part of its implementation details out, such that code could cast to that interface, call get_Value, and now have direct ref access to the enumerable’s underlying MRVTSL. Code could defend against this by wrapping the MRVTSL in an externally unpronounceable type (e.g. a little wrapper struct internal to the assembly), but that’s yet another workaround and complication.

Option 3 does have a couple of benefits, though:

• The implementation can pass this object to other things that we can’t as easily control. In particular, there are two places where the prototype at https://github.com/dotnet/corefx/blob/master/src/Common/tests/System/Threading/Tasks/Sources/ManualResetValueTaskSource.cs passes the parent: to the ExecutionContext, and to a SynchronizationContext if one was captured by an awaiter and used to invoke the continuation. With the IStrongBox approach, that doesn’t involve additional allocations. With the proposed approach, we’ll need to incur an additional allocation to pass the continuation delegate and state objects in the single state object parameter, though at the expense of an additional field on the MRVTSL, we could likely cache and reuse that object across all subsequent Posts.
• Option 3 can be done OOB if we desire to do that at some point. However, we could still deliver ManualResetValueTaskSourceLogic<bool> in an OOB component, it would just have slightly impaired functionality in a corner case. Consider a ValueTask constructed from an MRVTSL. When you await one, the async method infrastructure ensures that ExecutionContext is properly flowed, and since it's flowing it, it tells the IValueTaskSource it need not flow it, in which case none of these options matter as the ManualResetValueTaskSourceLogic<bool> won't interact with ExecutionContext. However, if code used valueTask.GetAwaiter().OnCompleted(...) directly, then the IValueTaskSource implementation will be told to flow context, and without one of these three options, it won't be able to do so properly.

Manual vs Auto

This type is named ManualReset because it requires an explicit call to Reset to reset itself for additional use. An alternative we might want to look to provide in the future is AutoReset. That version would automatically reset itself when GetResult was called to retrieve the result.

The difficulty with this is that such an AutoResetValueTaskSourceLogic would often be used in situations where it’s possible for multiple concurrent operations to be happening. For example, Socket supports ReceiveAsync being used concurrently by any number of callers, but that’s rare, with the vast majority case being there only being a single ReceiveAsync at a time. Thus, Socket.ReceiveAsync maintains a single reusable IValueTaskSource implementation, which it hands out to that single consumer; if other calls end up occurring while that one is being used, the more expensive path is taken of handing out a new object. This means that ReceiveAsync needs to know whether the object is in use or not, and it’s not enough to just track whether the consumer of the object is done using it due to having called GetResult. The same issue applies to System.Threading.Channel’s use of IValueTaskSource implementations.

Thus, for an AutoResetValueTaskSourceLogic to be really useful, it has to interact as well with some kind of ticketing system, whereby the instance can be checked out for exclusive use and then put back for someone else to consume when GetResult is called. Getting that right requires additional design that we’ve not yet done.

[stephentoub]

Fixed by dotnet/corefx#33104