This is the ninth part of the Async Wandering series. For your convenience you can find other parts in the table of contents in Part 1 – Why creating Form from WinForms in unit tests breaks async?

In previous part I mentioned multiple ways to await a task with timeout. Let’s see them.

Sketch

Let’s start with this code:

static async Task Main(string[] args)
{
	await Do();
}

public static async Task Do()
{
	await Hang(); // How to timeout here?
}

public static async Task Hang()
{
	await Task.Delay(TimeSpan.FromDays(1));
	Console.WriteLine("Worked!");
}

We will say that this code runs forever (even though it would finish after a day). The question is how do we call the method in line 8 with some timeout?

Let’s also add the following helper method:

public static async Task<T> ThrowTimeoutException<T>()
{
	await Task.Delay(TimeSpan.FromSeconds(1));
	throw new Exception("Timeout!");
}

It waits for a second and throws an exception, nothing big here.

WARNING: In examples below I don’t handle timeout exception properly. Don’t do it this way! Always await all your tasks or exceptions will be propagated by GC and kill your application in an out-of-band manner.

Solution 1 — timeouting manually

General trick is to wait for two tasks and throw exception when timeout happens:

public static async Task Do()
{
	var completed = await Task.WhenAny(Program.Hang(), Program.ThrowTimeoutException<bool>()).ConfigureAwait(false);
	await completed.ConfigureAwait(false);
}

Since Task.WhenAny returns task which resulted first, we need to wait for it as well.

Writing that code everywhere may be tedious. Let’s look further.

Solution 2 — extension method

public static async Task Do()
{
	await Program.Hang().Timeout().ConfigureAwait(false);
}

static class TaskExtensions
{
	public static async Task Timeout(this Task t)
	{
		await (await Task.WhenAny(t, Program.ThrowTimeoutException<bool>()).ConfigureAwait(false)).ConfigureAwait(false);
	}
}

That’s basically the same as before. However, this time we wrap the logic in one method and that’s probably the solution we are looking for.

Let’s explore some more.

Soltuion 3 — custom SynchronizationContext

This time we are going to use custom context and timeout over there:

public static async Task Do()
{
	await MyContext.Run(() => Program.Hang()).ConfigureAwait(false);
}

class MyTaskScheduler : TaskScheduler
{
	private readonly MyContext context;
	public BlockingCollection<Task> tasks = new BlockingCollection<Task>();

	public MyTaskScheduler(MyContext context)
	{
		this.context = context;
	}

	protected override IEnumerable<Task> GetScheduledTasks()
	{
		return tasks;
	}

	protected override void QueueTask(Task task)
	{
		Queue(task);
	}

	protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
	{
		return TryExecuteTask(task);
	}

	public new bool TryExecuteTask(Task task)
	{
		return base.TryExecuteTask(task);
	}

	public void Queue(Task task)
	{
		tasks.Add(task);
	}
}

class MyContext : SynchronizationContext
{
	public MyTaskScheduler scheduler;
	public TaskFactory factory;
	private int operations;

	public MyContext()
	{
		scheduler = new MyTaskScheduler(this);
		factory = new TaskFactory(CancellationToken.None, TaskCreationOptions.HideScheduler, TaskContinuationOptions.HideScheduler, scheduler);
	}

	public override void Post(SendOrPostCallback d, object state)
	{
		var task = factory.StartNew(() => d(state));
		scheduler.Queue(task);
	}

	public override void Send(SendOrPostCallback d, object state)
	{
		d(state);
	}

	public override void OperationCompleted()
	{
		operations--;
		if (operations == 0)
		{
			scheduler.tasks.CompleteAdding();
		}
	}

	public override void OperationStarted()
	{
		operations++;
	}

	public static Task Run(Action action)
	{
		TaskCompletionSource<bool> taskCompletionSource = new TaskCompletionSource<bool>(TaskContinuationOptions.RunContinuationsAsynchronously);
		Program.ThrowTimeoutException<bool>().ContinueWith(t => taskCompletionSource.SetException(t.Exception));

		Task.Run(() =>
		{
			var oldContext = SynchronizationContext.Current;
			var newContext = new MyContext();
			try
			{
				SynchronizationContext.SetSynchronizationContext(newContext);
				var spanningTask = newContext.factory.StartNew(action);
				foreach (var task in newContext.scheduler.tasks.GetConsumingEnumerable())
				{
					newContext.scheduler.TryExecuteTask(task);
					task.GetAwaiter().GetResult();
				}
				spanningTask.GetAwaiter().GetResult();
			}
			finally
			{
				SynchronizationContext.SetSynchronizationContext(oldContext);
			}
		}).ContinueWith(t => taskCompletionSource.SetException(t.Exception), TaskContinuationOptions.OnlyOnFaulted).ContinueWith(_ => taskCompletionSource.SetResult(true));

		return taskCompletionSource.Task;
	}
}

That’s super similar to awaiting void which we have already seen in Part 5. Once we start executing method we fire a timeout as well. Bonus points, this works for async void methods as well. The advantage is we handle all continuations on the context so we can timeout them selectively, as needed. This is probably an overkill, though.

Solution 4 — custom task type

We can use custom task and timeout there:

public static async TimeoutableTask Do()
{
	await Program.Hang().ConfigureAwait(false);
}

[AsyncMethodBuilder(typeof(TimeoutableTaskMethodBuilder))]
public class TimeoutableTask
{
	public TaskCompletionSource<object> Promise { get; } = new TaskCompletionSource<object>();

	public Task AsTask() => Promise.Task;

	public TaskAwaiter<object> GetAwaiter()
	{
		return Promise.Task.GetAwaiter();
	}

	public static implicit operator Task(TimeoutableTask task) => task.AsTask();
}

public class TimeoutableTaskMethodBuilder
{
	public void Start<TStateMachine>(ref TStateMachine stateMachine)
		where TStateMachine : IAsyncStateMachine
	{
		Program.ThrowTimeoutException<bool>().ContinueWith(t => {
			if (!Task.GetAwaiter().IsCompleted)
			{
				Task.Promise.SetException(t.Exception);
			}
		});
		stateMachine.MoveNext();
	}

	public static TimeoutableTaskMethodBuilder Create()
	{
		return new TimeoutableTaskMethodBuilder();
	}

	public void SetStateMachine(IAsyncStateMachine stateMachine)
	{
	}

	public void SetResult()
	{
		Task.Promise.SetResult(null);
	}

	public void SetException(Exception exception)
	{
		Task.Promise.SetException(exception);
	}

	public TimeoutableTask Task { get; } = new TimeoutableTask();

	public void AwaitOnCompleted<TAwaiter, TStateMachine>(
		ref TAwaiter awaiter,
		ref TStateMachine stateMachine)
		where TAwaiter : INotifyCompletion
		where TStateMachine : IAsyncStateMachine
	{
		awaiter.OnCompleted(ResumeAfterAwait(stateMachine));
	}

	public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(
		ref TAwaiter awaiter,
		ref TStateMachine stateMachine)
		where TAwaiter : ICriticalNotifyCompletion
		where TStateMachine : IAsyncStateMachine
	{
		awaiter.UnsafeOnCompleted(ResumeAfterAwait(stateMachine));
	}

	private Action ResumeAfterAwait<TStateMachine>(TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine
	{
		return () =>
		{
			stateMachine.MoveNext();
		};
	}
}

Notice first line – we use async TimeoutableTask. We need to provide a builder for it which is trivial in our case. We just delegate to regular logic and also fire a timeout task in line 26.

This is the twentieth fourth part of the .NET Inside Out series. For your convenience you can find other parts in the table of contents in Part 1 – Virtual and non-virtual calls in C#

Let’s take this code:

using System;
using System.Threading;
using System.Threading.Tasks;
					
public class Program
{
	public static void Main()
	{
		Task.Factory.StartNew(Handle);
	}
	
	public static void Handle(){
		Console.WriteLine(Thread.CurrentThread.ManagedThreadId);
		Task.Run(() => {
			Console.WriteLine(Thread.CurrentThread.ManagedThreadId);
			Thread.Sleep(1000);
		}).Wait();
	}
}

and compare it with this one:

using System;
using System.Threading;
using System.Threading.Tasks;
					
public class Program
{
	public static void Main()
	{
		Task.Factory.StartNew(Handle);
	}
	
	public static void Handle(){
		Console.WriteLine(Thread.CurrentThread.ManagedThreadId);
		var tcs = new TaskCompletionSource<bool>();
		Task.Run(() => {
			Console.WriteLine(Thread.CurrentThread.ManagedThreadId);
			Thread.Sleep(1000);
			tcs.SetResult(true);
		});
		tcs.Task.Wait();
	}
}

They look very similar, however, they give different outputs. First code runs both tasks (Handle and one sleeping) on the same thread:

11
11

while the latter runs on different threads:

10
12

Why? That’s because scheduler tries to detect a situation when we are waiting on the task immediately (first code) and run it inline. In the second code this cannot be detected as we wait for “different” task.