Allow an await to capture the current TaskScheduler instead of the current SynchronizationContext

[theodorzoulias]

Background and Motivation

The default behavior of the await operator is to capture the current SynchronizationContext, if one exists, and otherwise capture the current TaskScheduler. This behavior can be configured with the ConfigureAwait(false), which disables capturing both of these mechanisms. Currently it is not possible for a programmer to prioritize capturing the current TaskScheduler instead of the current SynchronizationContext, if they so wish for some reason. This results to the TaskScheduler mechanism being severely handicapped in the presence of an ambient SynchronizationContext. In this case it can schedule only the first part of an asynchronous method, until the first await. All continuations after the first await are presented with an ambient TaskScheduler.Current named ThreadPoolTaskScheduler. The programmatically configured TaskScheduler is ignored, and the programmer can't do really much about it.

Proposed API

A new overload of the ConfigureAwait method that accept a new ConfigureAwaitBehavior enum, having an OnlyCaptureTaskScheduler value:

namespace System.Threading.Tasks
{
    [Flags]
    public enum ConfigureAwaitBehavior
    {
        //...
        OnlyCaptureTaskScheduler = 0x8, // Ignore the synchronization context
    }

    public class Task
    {
        //...
        public ConfiguredTaskAwaitable ConfigureAwait(ConfigureAwaitBehavior behavior);
    }

    // The same for Task<TResult>
}

This is an extension of the API proposed on the issue #22144.

Usage Examples

It will be possible to build custom TaskSchedulers that are aware of the whole lifetime of an asynchronous method, whether or not an ambient SynchronizationContext is present. As an example one could build a TaskScheduler that measures the total duration that an asynchronous method is doing synchronous work or is blocked. In order for this to work, that method should have all its awaits configured with OnlyCaptureTaskScheduler.

var scheduler = new StopwatchScheduler();
await Task.Factory.StartNew(async () =>
{
    for (int i = 1; i <= 5; i++)
    {
        Thread.Sleep(500); // CPU-bound
        await Task.Delay(500).ConfigureAwait(
            ConfigureAwaitBehavior.OnlyCaptureTaskScheduler); // I/O-bound
    }
}, CancellationToken.None, TaskCreationOptions.None, scheduler).Unwrap();
Console.WriteLine("CPU-bound work: {scheduler.Elapsed.TotalMilliseconds:#,0} msec");

A complete example can be found here. The correct measurement is ~2,500 msec (five times Thread.Sleep(500)). The current await behavior gives an incorrect measurement of ~500 msec.

Another possibility would be to build asynchronous APIs that accept an async delegate together with a TaskScheduler, and ensure that all invocations of the delegate will be on that TaskScheduler. See for example the Retry method below:

public static Task<TResult> Retry<TResult>(Func<Task<TResult>> action,
    int retryCount, TaskScheduler taskScheduler)
{
    return Task.Factory.StartNew(async () =>
    {
        while (true)
        {
            try
            {
                return await action().ConfigureAwait(
                    ConfigureAwaitBehavior.OnlyCaptureTaskScheduler);
            }
            catch
            {
                if (--retryCount < 0) throw;
            }
        }
    }, default, TaskCreationOptions.DenyChildAttach, taskScheduler).Unwrap();
}

The current await behavior discourages async APIs with a TaskScheduler parameter, because the only reliable way to respect the parameter is to start a new configured task for each invocation.

Risks

In the above example, if the user neglects to configure all awaits inside their own action in the same way, the action will be invoked on the supplied taskScheduler only partially. This may be against their expectations.

[dotnet-issue-labeler]

I couldn't figure out the best area label to add to this issue. If you have write-permissions please help me learn by adding exactly one area label.

[ghost]

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

Issue Details

---

Background and Motivation

The default behavior of the await operator is to capture the current SynchronizationContext, if one exists, and otherwise capture the current TaskScheduler. This behavior can be configured with the ConfigureAwait(false), which disables capturing both of these mechanisms. Currently it is not possible for a programmer to prioritize capturing the current TaskScheduler instead of the current SynchronizationContext, if they so wish for some reason. This results to the TaskScheduler mechanism being severely handicapper in the presence of an ambient SynchronizationContext. In this case it can schedule only the first part of an asynchronous method, until the first await. All continuations after the first await are presented with an ambient TaskScheduler.Current named ThreadPoolTaskScheduler. The programmatically configured TaskScheduler is ignored, and the programmer can't do really much about it.

Proposed API

A new overload of the ConfigureAwait method that accept a new ConfigureAwaitBehavior enum, having an OnlyCaptureTaskScheduler value:

namespace System.Threading.Tasks
{
    [Flags]
    public enum ConfigureAwaitBehavior
    {
        //...
        OnlyCaptureTaskScheduler = 0x8, // Ignore the synchronization context
    }

    public class Task
    {
        //...
        public ConfiguredTaskAwaitable ConfigureAwait(ConfigureAwaitBehavior behavior);
    }

    // The same for Task<TResult>
}

This is an extension of the API proposed on the issue #22144.

Usage Examples

It will be possible to build custom TaskSchedulers that are aware of the whole lifetime of an asynchronous method, whether or not an ambient SynchronizationContext is present. As an example one could build a TaskScheduler that measures the total duration that an asynchronous method is doing synchronous work or is blocked. In order for this to work, that method should have all its awaits configured with OnlyCaptureTaskScheduler.

var scheduler = new StopwatchScheduler();
await Task.Factory.StartNew(async () =>
{
    for (int i = 1; i <= 5; i++)
    {
        Thread.Sleep(500); // CPU-bound
        await Task.Delay(500).ConfigureAwait(
            ConfigureAwaitBehavior.OnlyCaptureTaskScheduler); // I/O-bound
    }
}, CancellationToken.None, TaskCreationOptions.None, scheduler).Unwrap();
Console.WriteLine("CPU-bound work: {scheduler.Elapsed.TotalMilliseconds:#,0} msec");

A complete example can be found here. The correct measurement is ~2,500 msec (five times Thread.Sleep(500)). The current await behavior gives an incorrect measurement of ~500 msec.

Another possibility would be to build asynchronous APIs that accept an async delegate together with a TaskScheduler, and ensure that all invocations of the delegate will be on that TaskScheduler. See for example the Retry method below:

public static Task<TResult> Retry<TResult>(Func<Task<TResult>> action,
    int retryCount, TaskScheduler taskScheduler)
{
    return Task.Factory.StartNew(async () =>
    {
        while (true)
        {
            try
            {
                return await action().ConfigureAwait(
                    ConfigureAwaitBehavior.OnlyCaptureTaskScheduler);
            }
            catch
            {
                if (--retryCount < 0) throw;
            }
        }
    }, default, TaskCreationOptions.DenyChildAttach, taskScheduler).Unwrap();
}

The current await behavior discourages async APIs with a TaskScheduler parameter, because the only reliable way to respect the parameter is to start a new configured task for each invocation.

Risks

In the above example, if the user neglects to configure all awaits inside their own action in the same way, the action will be invoked on the supplied taskScheduler only partially. This may be against their expectations.

Author:	theodorzoulias
Assignees:	-
Labels:	api-suggestion, area-System.Threading.Tasks, untriaged
Milestone:	-

[tarekgh]

CC @stephentoub

[theodorzoulias]

When I made this suggestion I hadn't realized (initially) that it is possible to configure the await manually, by creating a custom awaitable. This is pretty cool actually! Since the functionality requested here will probably be useful to only a small number of people, maybe it would be more practical to just publish an implementation of such an awaitable, so that anyone who needs it can copy-paste it into their project. I wonder if the implementation below is OK, or has any bugs that I can't see.

public struct TaskSchedulerAwaitable : INotifyCompletion
{
    private readonly ConfiguredTaskAwaitable.ConfiguredTaskAwaiter _taskAwaiter;
    private readonly TaskScheduler _capturedScheduler;

    public TaskSchedulerAwaitable(Task task)
    {
        _taskAwaiter = task.ConfigureAwait(false).GetAwaiter();
        _capturedScheduler = TaskScheduler.Current;
    }
    public TaskSchedulerAwaitable GetAwaiter() => this;
    public bool IsCompleted { get { return _taskAwaiter.IsCompleted; } }
    public void GetResult() { _taskAwaiter.GetResult(); }
    public void OnCompleted(Action continuation)
    {
        var self = this;
        _taskAwaiter.OnCompleted(() =>
        {
            if (TaskScheduler.Current == self._capturedScheduler)
            {
                continuation();
            }
            else
            {
                Task.Factory.StartNew(continuation, default,
                    TaskCreationOptions.DenyChildAttach, self._capturedScheduler);
            }
        });
    }
}

Usage example:

await new TaskSchedulerAwaitable(Task.Delay(1000)); // Continue on the current TaskScheduler

[stephentoub]

maybe it would be more practical to just publish an implementation of such an awaitable, so that anyone who needs it can copy-paste it into their project

Yup. Thanks.