// http://stackoverflow.com/a/22587012/1768303

​

using System;

using System.Threading;

using System.Threading.Tasks;

using System.Collections.Concurrent;

​

public class Program

{

    static void Identify()

    {

        var thread = Thread.CurrentThread;

        string name = thread.IsThreadPoolThread

            ? "Thread pool" : thread.Name;

        if (string.IsNullOrEmpty(name))

            name = "#" + thread.ManagedThreadId;

        Console.WriteLine("Continuation on: " + name);

    }

    public static void Main()

    {

        Thread.CurrentThread.Name = "Main thread";

        var source = new TaskCompletionSource<int>();

​

        // library code

        var task = source.Task.ImposeAsync();

​

        Func<Task, Task> testAwait = async (t) =>

        {

            await t;

            Identify();

        };

​

        // testing original source.Task

        source.Task.ContinueWith(delegate

        {

            Identify();

        });

​

        // client code

        testAwait(task);

        task.ContinueWith(delegate

        {

            Identify();

        }, TaskContinuationOptions.ExecuteSynchronously);

​

        // library code

        source.SetResultAsync(123);

​

        Thread.Sleep(1000);

        Console.WriteLine("Press [return]");

        Console.ReadLine();

    }

}

​

public static class TaskExt

{

    static readonly ConcurrentDictionary<Task, Thread> s_tcsTasks =

        new ConcurrentDictionary<Task, Thread>();

​

    // SetResultAsync

    static public void SetResultAsync<TResult>(

        this TaskCompletionSource<TResult> @this,

        TResult result)

    {

        s_tcsTasks.TryAdd(@this.Task, Thread.CurrentThread);

        try

        {

            @this.SetResult(result);

        }

        finally

        {

            Thread thread;

            s_tcsTasks.TryRemove(@this.Task, out thread);

        }

    }

​

    // ContinueWithAsync, TODO: more overrides

    static public Task ContinueWithAsync<TResult>(

        this Task<TResult> @this,

        Action<Task<TResult>> action,

        TaskContinuationOptions continuationOptions = TaskContinuationOptions.None)

    {

        return @this.ContinueWith((Func<Task<TResult>, Task>)(t =>

        {

            Thread thread = null;

            s_tcsTasks.TryGetValue(t, out thread);

            if (Thread.CurrentThread == thread)

            {

                // same thread which called SetResultAsync, avoid potential deadlocks

​

                // using thread pool

                return Task.Run(() => action(t));

​

                // not using thread pool (TaskCreationOptions.LongRunning creates a normal thread)

                // return Task.Factory.StartNew(() => action(t), TaskCreationOptions.LongRunning);

            }

            else

            {

                // continue on the same thread

                var task = new Task(() => action(t));

                task.RunSynchronously();

                return Task.FromResult(task);

            }

        }), continuationOptions).Unwrap();

    }

​

    // ImposeAsync, TODO: more overrides

    static public Task<TResult> ImposeAsync<TResult>(this Task<TResult> @this)

    {

        return @this.ContinueWith(new Func<Task<TResult>, Task<TResult>>(antecedent =>

        {

            Thread thread = null;

            s_tcsTasks.TryGetValue(antecedent, out thread);

            if (Thread.CurrentThread == thread)

            {

                return antecedent.ContinueWith(t => t, TaskContinuationOptions.None).Unwrap();

            }

            else

            {

                return antecedent;

            }

        }), TaskContinuationOptions.ExecuteSynchronously).Unwrap();

    }

​

}

​