みんな大好き TaskCompletionSource。 TaskCompletionSource は死ぬほど重要かつ便利な class なのですが、TaskCompletionSource.Task を await した時の continuation (= await 後に行われる継続処理) がどこのスレッドで実行されるか意識しないと危険だよというお話。

既定では continuation は同期的に呼ばれる。(非推奨)

既定では TaskCompletionSource.Task を await した後の continuation は TrySetResult から同期的に呼ばれます。

var tcs = new TaskCompletionSource();

Console.WriteLine($"1. start : {Environment.CurrentManagedThreadId}");

var _ = Task.Run(async () =>
{
    Console.WriteLine($"2. before await in Task.Run : {Environment.CurrentManagedThreadId}");

    await tcs.Task;

    Console.WriteLine($"3. after await in Task.Run : {Environment.CurrentManagedThreadId}");

    Thread.Sleep(TimeSpan.FromSeconds(5));

    Console.WriteLine($"4. after sleep in Task.Run : {Environment.CurrentManagedThreadId}");
});


Console.WriteLine($"5. before TrySetResult : {Environment.CurrentManagedThreadId}");

// Task.Run の中の await 時に既に tcs.Task がすでに完了状態になっている事を防ぐ
Thread.Sleep(TimeSpan.FromSeconds(1));

tcs.TrySetResult();

Console.WriteLine($"6. after TrySetResult : {Environment.CurrentManagedThreadId}");

Console.ReadLine();

結果はこうなります。

1. start : 1
5. before TrySetResult : 1
2. before await in Task.Run : 6
3. after await in Task.Run : 1
4. after sleep in Task.Run : 1
6. after TrySetResult : 1

tcs.TrySetResult() を呼んだ後、 await tcs.Task の continuation は tcs.TrySetResult() したスレッドと同一のスレッドで同期的に実行されます。そして、tcs.TrySetResult() の後の行は await tcs.Task の continuation が完了してから呼ばれます。これがなかなか問題で、軽率にデッドロックであったりスレッドプールの枯渇を引き起こすことになります。つらい。

というわけでどうしましょう？というと、continuation は非同期的に呼ぶようにしようね、が解となります。

continuation を非同期的に呼ぶようにする。(推奨)

どうするか？TaskCompletionSource を new する際に TaskCreationOptions.RunContinuationsAsynchronously を渡してあげます。

// 非推奨
var tcs = new TaskCompletionSource();

// 推奨
var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);

先ほどのコードで TaskCreationOptions.RunContinuationsAsynchronously を設定してみて実行してみると...

var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);

Console.WriteLine($"1. start : {Environment.CurrentManagedThreadId}");

var _ = Task.Run(async () =>
{
    Console.WriteLine($"2. before await in Task.Run : {Environment.CurrentManagedThreadId}");

    await tcs.Task;

    Console.WriteLine($"3. after await in Task.Run : {Environment.CurrentManagedThreadId}");

    Thread.Sleep(TimeSpan.FromSeconds(5));

    Console.WriteLine($"4. after sleep in Task.Run : {Environment.CurrentManagedThreadId}");
});


Console.WriteLine($"5. before TrySetResult : {Environment.CurrentManagedThreadId}");

// Task.Run の中の await 時に既に tcs.Task がすでに完了状態になっている事を防ぐ
Thread.Sleep(TimeSpan.FromSeconds(1));

tcs.TrySetResult();

Console.WriteLine($"6. after TrySetResult : {Environment.CurrentManagedThreadId}");

Console.ReadLine();

以下のような結果が得られます。

1. start : 1
5. before TrySetResult : 1
2. before await in Task.Run : 6
6. after TrySetResult : 1
3. after await in Task.Run : 4
4. after sleep in Task.Run : 4

tcs.TrySetResult() を呼んだ後、 await tcs.Task の continuation は tcs.TrySetResult() を呼び出したしたスレッドとは無関係の別のスレッドで非同期的に実行される事がみてとれます。そして tcs.TrySetResult() を呼んだスレッドは素直に直後に記述されている処理を実行するようになります。

まとめ

continuation がどこで実行されるのかは常に意識しておきましょう。そして相当強い理由がない限り常に TaskCompletionSource を new する際は TaskCreationOptions.RunContinuationsAsynchronously を設定しておくとよいでしょう。