How to make TaskCompletionSource.Task complete using specific TaskScheduler

Question

How to make the completion of TaskCompletionSource.Task happen on specific TaskScheduler, when I call TaskCompletionSource.SetResult?

Currently, I'm using the idea I borrowed from this post:

static public Task<TResult> ContinueOnTaskScheduler<TResult>(
    this Task<TResult> @this, TaskScheduler scheduler)
{
    return @this.ContinueWith(
        antecedent => antecedent,
        CancellationToken.None,
        TaskContinuationOptions.ExecuteSynchronously,
        scheduler).Unwrap();
}

So whenever I would return TaskCompletionSource.Task to the caller, I now return TaskCompletionSource.Task.ContinueOnTaskScheduler(scheduler) instead.

Is it possible to somehow avoid this another level of indirection of ContinueWith?

Answer 1

It would be interesting to know your goals behind this. Anyway, if you like to avoid the overhead of ContinueWith (which I think is quite low), you'd probably have to come up with your own version of a pattern similar to TaskCompletionSource.

It's not that complex. E.g., something like Promise below can be used in the same way you use TaskCompletionSource, but would allow to provide a custom TaskScheduler for completion (disclaimer: almost untested):

public class Promise
{
    readonly Task _task;
    readonly CancellationTokenSource _cts;
    readonly object _lock = new Object();
    Action _completionAction = null;

    // public API

    public Promise()
    {
        _cts = new CancellationTokenSource();
        _task = new Task(InvokeCompletionAction, _cts.Token); 
    }

    public Task Task { get { return _task; } }

    public void SetCompleted(TaskScheduler sheduler = null)
    {
        lock(_lock)
            Complete(sheduler);
    }

    public void SetException(Exception ex, TaskScheduler sheduler = null)
    {
        lock (_lock)
        {
            _completionAction = () => { throw ex; };
            Complete(sheduler);
        }
    }

    public void SetException(System.Runtime.ExceptionServices.ExceptionDispatchInfo edi, TaskScheduler sheduler = null)
    {
        lock (_lock)
        {
            _completionAction = () => { edi.Throw(); };
            Complete(sheduler);
        }
    }

    public void SetCancelled(TaskScheduler sheduler = null)
    {
        lock (_lock)
        {
            // don't call _cts.Cancel() outside _completionAction
            // otherwise the cancellation won't be done on the sheduler
            _completionAction = () =>
            {
                _cts.Cancel();
                _cts.Token.ThrowIfCancellationRequested();
            };
            Complete(sheduler);
        }
    }

    // implementation

    void InvokeCompletionAction()
    {
        if (_completionAction != null)
            _completionAction();
    }

    void Complete(TaskScheduler sheduler)
    {
        if (Task.Status != TaskStatus.Created)
            throw new InvalidOperationException("Invalid task state.");
        _task.RunSynchronously(sheduler?? TaskScheduler.Current);
    }
}

On a side note, this version has an override for SetException(ExceptionDispatchInfo edi), so you could propagate the active exception's state from inside catch:

catch(Exception ex)
{
    var edi = ExceptionDispatchInfo.Capture(ex);
    promise.SetException(edi);
}

It's easy to create a generic version of this, too.

There's a downside of this approach, though. A 3rd party can do promise.Task.Run or promise.Task.RunSynchronously, as the Task is exposed in the TaskStatus.Created state.

You could add a check for that into InvokeCompletionAction, or you could probably hide it using nested tasks / Task.Unwrap (although the latter would bring some overhead back).