✋

C# の Event を await する

2022/03/19に公開

はじめに

ふと WPF の Animation / Storyboard で終了を await で待機したいなあ (UniTask の DOTween 対応みたいな) と思い立ちました。

WPF の Animation / Storyboard の完了は Timeline.Completed イベントで通知されるので、これを await で待機できればよいわけです。具体的には

Storyboard sb;

// 前処理
EventHandler? h;
h = (_, _) =>
{
    sb.Completed -= h;
    // 後処理
};
sb.Completed += h;
sb.Begin();

をこう書きたい。

Storyboard sb;

// 前処理
await sb.BeginAsync();
// 後処理

ということでいろいろ調べて 4 パターンほど実装してみました。きっかけが WPF の Animation なので、それを前提とした実装になっていますが、 C# の event を初めとした "Task / ValueTask 以外で await する" 際の基本的な内容になっているかなと思います。

private async void ButtonBase_OnClick(object sender, RoutedEventArgs e)
{
    if (TryFindResource("Storyboard") is Storyboard sb)
    {
        Button.Content = "Running";
        await sb.BeginTypeCAsync();
        Button.Content = null;
    }
}

結論

ほとんどの場合 ManualResetValueTaskSourceCore で用件が足りるはず
ManualResetValueTaskSourceCore が使えない (.NET Standard 2.1 より前) 場合は TaskCompletionSource を検討する
Awaitable/Awaiter パターンの独自実装は最後に考える

TaskCompletionSource

TaskCompletionSource は Task を利用して完了通知をするためのクラスで C# 5.0 の頃くらいから存在しているクラスです。典型的な Promise/Future パターンの実装になります。

public static Task BeginAsync(this Storyboard self, CancellationToken ct = default)
{
    var tcs = new TaskCompletionSource();
    EventHandler? h = null;
    h = (_, _) =>
    {
        self.Completed -= h;
        tcs.SetResult();
    };
    self.Completed += h;
    ct.Register(() =>
    {
        self.Completed -= h;
        tcs.SetCanceled();
    });

    try
    {
        self.Begin();
        return tcs.Task;
    }
    catch
    {
        self.Completed -= h;
        throw;
    }
}

TaskCompletionSource.SetResult をすると Task が完了状態になるので、 Completed イベントで SetResult をします。 Cancel 時の配慮も一応します。

ManualResetValueTaskSourceCore

TaskCompletionSource は .NET Framework でも使えますし、簡単便利に使えてよいのですが、中の実装が Task をベースになっているので高頻度に使われるとアロケーションが気になってきます。

.NET Core 3 (.NET Standard 2.1) 以降で "ManualResetValueTaskSourceCore" という IValueTaskSource を簡単に実装するための構造体が用意されています。

これを使うことで、より軽量な TaskCompletionSource (のようなもの) を実装することができます。下記の記事が大変参考になりました。ありがとうございます。

ポイントですが

IValueTaskSource を継承したクラスを定義する
IValueTaskSource の実装は基本的には ManualResetValueTaskSourceCore に丸投げで OK
ManualResetValueTaskSourceCore.Version が ValueTask の Token
使用する前に Reset をする
SetCanceled はないので (必要なら) "SetException(OperationCanceledException)" で代用する
アロケーションをなるべく避けるために ObjectPool (Microsoft.Extensions.ObjectPool) を活用する
- ObjectPool.Get で取得したらまず Reset
- GetResult のタイミングで ObjectPool.Return で戻す

ラッパー実装 (ManualResetValueTaskSource)

ManualResetValueTaskSourceCore を単純にほぼそのままラップした IValueTaskSource 実装をすると ValueTask ベースの TaskCompletionSource 的な実装になります。

public static ValueTask BeginAsync(this Storyboard self, CancellationToken ct = default)
{
    var vts = ManualResetValueTaskSource.Create();
    EventHandler? h = null;
    h = (_, _) =>
    {
        self.Completed -= h;
        vts.SetResult();
    };
    self.Completed += h;
    ct.Register(() =>
    {
        self.Completed -= h;
        vts.SetCanceled();
    });

    try
    {
        self.Begin();
        return vts.AsValueTask();
    }
    catch
    {
        self.Completed -= h;
        throw;
    }
}

使う側はほぼ TaskCompletionSource と同じになります。が、 ObjectPool を使っているのにキャプチャを伴うラムダ式を使っている段階で割と片手落ち感はあります。

イベント特化版 (EventValueTaskSource)

というわけでイベントにからむリソースを ValueTaskSource 内に保持する EventValueTaskSource を実装してみました。

public static ValueTask<EventArgs> BeginAsync(this Storyboard self, CancellationToken ct = default)
{
    var r = EventValueTaskSource<Timeline, EventHandler, EventArgs>.Create(self,
        (t, h) => t.Completed += h,
        (t, h) => t.Completed -= h,
        ct);
    try
    {
        self.Begin();
        return r.AsValueTask();
    }
    catch
    {
        r.Release();
        throw;
    }
}

イベント登録、解除に関する処理は EventValueTaskSource の中に押し込まれていますがやっている事の本質は ManualResetValueTaskSource と同じで Create のタイミングでイベントを登録 (addHandler を実行) し、 GetResult で解除 (releaseHandler を実行) します。キャプチャの伴うラムダ式は排除されているので、対イベント用としては ManualResetValueTaskSource より改善した実装になった・・・はず。

イベント特化なので使う時の手数も相当減ってよいのではないでしょうか。

Awaitable/Awaiter パターン

Awaitable/Awaiter パターンは所定の条件を満たした実装を行うと実装したクラスに対して await が可能になるというものです。

やりようによってはあらゆる型を await 可能にできる (await ～と書けるようにする意味があるかどうかは要検討) 方法です。

一応動作する形で実装してみたのですが完全理解には至らずもやっとした感じです。 Timeline.Completed イベントに特化する形で実装したので

public static TimelineCompletedAwaitable BeginAsync(this Storyboard self)
{
    var r = TimelineCompletedAwaitable.Create(self);
    self.Begin();
    return r;
}

のような感じになります。

Awaiter

Awaiter は await 中の状態管理する機構 (ステートマシン) に対する制御を実装するものです。クラス、構造体どちらで実装しても OK です。多くの場合は readonly struct で実装し、 IValueTaskSourceCore のような状態管理をしているクラスに全て移譲する (このクラスは再利用前提) 、としているようです。

Awaiter は INotifyCompletion か ICriticalNotifyCompletion の実装を必要とします。 ICriticalNotifyCompletion は INotifyCompletion を継承しており、調べた範囲で公開されている既存の実装はほぼ ICriticalNotifyCompletion を実装していました。

また、 INotifyCompletion / ICriticalNotifyCompletion 以外にも実装が必要なものがあります。

bool IsCompleted { get; }
T GetResult() or void GetResult()
void INotifyCompletion.OnCompleted(Action continuation);
void ICriticalNotifyCompletion.UnsafeOnCompleted(Action continuation);

IsCompleted / GetResult

IsCompleted は非同期処理が完了しているかどうかの状態を返します。

GetResult は結果を返すメソッドで、非同期処理が完了する前に呼ばれた場合は完了するまで待つ必要があります (が今回は省略しました) 。

OnCompleted / UnsafeOnCompleted

この二つのメソッドは基本的には同じで、 await している処理が完了した後に呼び出す必要のある処理をスケジューリングする処理を実装します。 continuation が完了後に呼び出す処理そのものです。実装上は IsCompleted == false の時に呼び出されます。

OnCompleted と UnsafeOnCompleted の違いですが

Unlike OnCompleted, UnsafeOnCompleted doesn't have to propagate ExecutionContext information.

とあります。

この差異について ManualResetValueTaskSourceCore の実装を見ると確かにこの違いの配慮がされた実装になっています。一方、今回の TimelineCompletedAwaitable ではシングルスレッドでの非同期実行でスレッドをまたがないため配慮不要と判断し、 continuation のみを退避してそのまま実行しています。

Awaitable

Awaitable は "Awaiter を返す GetAwaiter() というメソッドがある型" が条件です。拡張メソッドでもよいので既存の型に対して await 可能にすることもできます。

public TimelineCompletedAwaiter GetAwaiter() => new TimelineCompletedAwaiter(this);

ManualResetValueTaskSourceCore を使って Awaitable/Awaiter パターンを実装する

Awaitable/Awaiter パターンをちゃんと実装するのは難しい？ならばということで ManualResetValueTaskSourceCore を使って実装をしてみました。

問題の OnCompleted/UnsafeOnCompleted もこんな感じでいけます。

public void OnCompleted(Action continuation) =>
    _awaitable._core.OnCompleted(h => ((Action?)h)?.Invoke(), continuation, _token,
        ValueTaskSourceOnCompletedFlags.FlowExecutionContext |
        ValueTaskSourceOnCompletedFlags.UseSchedulingContext);

public void UnsafeOnCompleted(Action continuation) =>
    _awaitable._core.OnCompleted(h => ((Action?)h)?.Invoke(), continuation, _token,
        ValueTaskSourceOnCompletedFlags.UseSchedulingContext);

ですが目的が Awaitable を作りたいわけではなく await したい、なはずなので ManualResetValueTaskSourceCore を使うなら普通に ValueTask と組み合わせて使えばよいのではと思います。

ちなみに ValueTaskSourceOnCompletedFlags はちゃんと意味がありまして、例えば

public class SleepAwaitable
{
    private ManualResetValueTaskSourceCore<int> _core;

    public SleepAwaitable(int millisec)
    {
        ThreadPool.QueueUserWorkItem(
            _ =>
            {
                Thread.Sleep(millisec);
                _core.SetResult(0);
            });
    }

のように別スレッドから SetResult を呼んで完了をするようにした場合、

public void OnCompleted(Action continuation) =>
    _awaitable._core.OnCompleted(h => ((Action?)h)?.Invoke(), continuation, _token,
        ValueTaskSourceOnCompletedFlags.FlowExecutionContext);

public void UnsafeOnCompleted(Action continuation) =>
    _awaitable._core.OnCompleted(h => ((Action?)h)?.Invoke(), continuation, _token,
        ValueTaskSourceOnCompletedFlags.None);

と "ValueTaskSourceOnCompletedFlags.UseSchedulingContext" を外すとそのままワーカースレッド上で continuation 実行されます (ので WPF などの UI とからめると例外で落ちます) 。

おわりに

現状、 ManualResetValueTaskSourceCore の使い勝手がよいのでこれを使って自前で IValueTaskSource の実装をするのが一番バランスのよい方法ではないかと思います。単純なラッパー実装 (これは標準で用意してくれてもよいような気はしますが・・・) の他、今回のように用途に合わせた専用実装に使うのもよいと思います。

Awaitable/Awaiter パターンは ManualResetValueTaskSourceCore で要件を満たせないケースで実装すればいいかなと思います。そのようなケースは限られてくると思うのですが、シングルスレッド非同期実行を前提としてもっと処理を省いて軽くしたい場合などでしょうか。プロダクションコードで Awaitable/Awaiter パターンを実装するかどうかは別として、試しに書いてみて C# の async/await の挙動を追ってみるのは勉強になるかなと思いました。