contextのキャンセル時にtesting.T.Cleanupのようにクリーンアップ処理を実行できるパッケージdonegroupを作った
contextのキャンセルとそれに対応したクリーンアップ処理
Goにおいてcontextはさまざまな値を伝搬させるために使用します。その1つがキャンセル信号です。
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
contextを通じてキャンセル信号を送ることによって、そのcontextが伝播した関係するある種のcontextグループでキャンセルに対応した処理を行うことができます。
このcontextのキャンセルに対応した処理を、このエントリでは「クリーンアップ処理」と呼ぶことにします。
クリーンアップ処理の具体例: Graceful shutdown
あるミドルウェア（プロキシサーバでもデータベースサーバでもよい）があったとして、そのミドルウェアにはメインの機能以外に「ログやメトリクスを定期的にまとめて外部サーバに転送する機能（ログ転送機能）」があるとします。
ミドルウェアのプロセスを停止する際に、ログ転送機能が転送中であれば転送完了までプロセスの停止を待って欲しいですし、次の起動時のためにどこまで転送完了したかの記録するまでプロセスの停止を待って欲しいと考えます。
このような、いわゆる「Graceful shutdown」の実現に必要な後始末処理は、本エントリにおける「クリーンアップ処理」の1つの具体例になります。
cancel() はキャンセル信号の伝搬しかしない
キャンセル処理の伝搬は大抵 context.WithCancel() で作成した cancel() を使って行います。
cancel() を実行することでcontextや、そのcontextから派生した子contextにキャンセル信号が伝搬され、その信号は context.Done() から補足できます（実体は <-chan struct{}）。
contextを引数で受け取っている各処理は、context.Done() から信号を受け取ったときにクリーンアップ処理を行うようにする（処理を書く）のですが、このクリーンアップ処理の実行の完遂は context.WithCancel() と cancel() で実現しているキャンセル信号の伝搬の仕組みでは保証されません。
次のGo Playgroundに簡単な例を置きました。クリーンアップ処理が完了する前に main() が終了してしまい、クリーンアップ処理が完了していないことがわかると思います。
https://go.dev/play/p/F-4qnF5L9bf
クリーンアップ処理の実行を保証したい
少し複雑なGoのプログラムを書くと context はさまざまな箇所に渡されます。「（終了処理実現のために）関数の第一引数にcontextを渡す」という慣習もありますし、その深さ（どれくらい関数をまたいでcontextが渡されてきているか）は意図せず深くなっていきます。
また、goroutineの存在により並行処理もよく活用します。開始した並行処理は終了処理を書く必要があり、必要であればクリーンアップ処理を書く必要があります。そもそも context を渡しているのはgoroutineによる並行処理の終了処理のためでもあります。
終了処理のトリガーにはcontextのキャンセルを検知できる context.Done() を使います。
では、クリーンアップ処理はどう実装したら良いでしょう？クリーンアップ処理の完了を待つための仕組みが必要です。
testing.T.Cleanup
ところで、testing.T.Cleanup() はテストのクリーンアップ処理を保証するための仕組みです。t.Cleanupでクリーンアップ処理を登録しておくことによって、テスト終了時にそれらのクリーンアップ処理を実行してくれます。
私はcontextの文脈でもこの testing.T.Cleanup() のような機能が欲しいと考えました。
donegroup
上記のようなモチベーションで作成したパッケージが donegroup です。
donegroupはcontextのキャンセル時にtesting.T.Cleanupのようにクリーンアップ処理を実行できる仕組みを提供するパッケージです。
準備
まず context.WithCancel() の代わりに donegroup.WithCancel() を使ってcontextとcancel()の対を作成します。
そして、cancel() 実行後にクリーンアップ処理を待つため、donegroup.Wait() を実行します。
これだけで準備は完了です。
具体的には
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
を
ctx, cancel := donegroup.WithCancel(context.Background()) defer func() { cancel() if err := donegroup.Wait(ctx); err != nil { log.Fatal(err) } }()
に書き換えるだけです。
クリーンアップ処理の登録（ donegroup.Cleanup ）
クリーンアップ処理の登録も簡単です。任意の場所で donegroup.Cleanup() を使ってクリーンアップ処理を登録するだけです。
err := donegroup.Cleanup(ctx, func(_ context.Context) error { // 何かしらのクリーンアップ処理 fmt.Println("cleanup") return nil })
これだけで、contextのキャンセル時にクリーンアップ処理の実行をしてくれます。クリーンアップ処理はいくつでも登録できます。testing.T.Cleanup() との違いはクリーアップ処理の実行順序の保証がないということだけです。
関数実行の保証（ donegroup.Awaiter / donegroup.Awaitable ）
donegroupは、次のように「実行に時間がかかる処理の実行を待つ（保証する）」こともできます。
この場合は donegroup.Awaiter() を使います。completed() が実行されるまで donegroup.Wait() は待ってくれます。
ctx, cancel := donegroup.WithCancel(context.Background()) go func() { completed, err := donegroup.Awaiter(ctx) if err != nil { log.Fatal(err) return } // 実行に時間がかかる処理 time.Sleep(100 * time.Second) fmt.Println("do something") completed() }() // メインの処理 fmt.Println("main") time.Sleep(10 * time.Millisecond) cancel() if err := donegroup.Wait(ctx); err != nil { log.Fatal(err) } fmt.Println("finish") // Output: // main // do something // finish
「この関数ブロックの実行を保証したい」というような場合は defer donegroup.Awaitable(ctx)() が使いやすいです。
go func() { defer donegroup.Awaiter(ctx)() // この関数ブロックの実行を保証する。 if err != nil { log.Fatal(err) return } // 実行に時間がかかる処理 time.Sleep(100 * time.Second) fmt.Println("do something") }()
クリーンアップ処理の実行待ちにタイムアウトを設定したい
donegroup.Wait() の代わりに donegroup.WaitWithTimeout() を使うことでクリーンアップ処理待ちにタイムアウトを設定することができます。
この場合、donegroup.Cleanup() や donegroup.Awaiter() donegroup.Awaitable() にタイムアウトが伝播します。
具体的には donegroup.Cleanup() に渡すクリーンアップ処理 func(ctx context.Context) error の ctx にキャンセル信号が届くので ctx.Done() を使って適切にキャンセル処理を書きましょう。
donegroupのメリット
donegroupのメリットは、contextさえ渡しておけば関数シグネチャなどを変えずにいつでもクリーンアップ処理を追加実装できる点です。
また、クリーンアップ処理のエラーもまとめて受け取れます。
実装の深い場所でクリーンアップ処理が必要な並行処理を実装する場合も安心です。
というわけで
散らばっている並行処理の後始末に是非使ってみてください。
