Goのcontextパッケージの使い方がイマイチよく分からない。このように思っている方は、多いのではないでしょうか。

「なんかリクエスト情報を持っている」、今までなんとなくそう思いながら使っていました。そして、そのcontextの使い方を間違えて本番にリリースした結果・・・　見事に事故りました。

事故の経緯

今回使用したのは、graphqlというライブラリです。GraphQLクエリ実行に便利なクライアントですね。
https://github.com/machinebox/graphql

社内のAPIからコールしていたのは、このライブラリのgraphql.Run()メソッドです。今回の修正がこちら。
もともとcontext.Backgrount()で空のcontextを生成して使用していたのですが、同じリクエストスコープだからといって、パラメータのcontextを使用することにしたら、

・
・
・
事故りました−−−
client.Run()でエラーが連発していたんです（泣）

事故の原因

client.Run()の中身を覗いてみると、先頭でこのようなチェックをしています。

func (c *Client) Run(ctx context.Context, req *Request, resp interface{}) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}
	if len(req.files) > 0 && !c.useMultipartForm {

パラメータのcontextが終了しているか否かを、ctx.Done()でチェックしているわけです。「先程まで使っていたcontextだから、終了しているわけないじゃん」と思ったあなた。実はここに大きな落とし穴がありました。誰も気づかない大きな大きな落とし穴が・・・

このgraphql処理を呼んでいる関数doSomething()ですが、呼び出し元ではこんな呼び方をしていました。

func do(ctx context.Context) {
  go doSomething(ctx)
}

Goroutineで呼び出していたわけです。

つまり、先程までcontextが生きていたからといって、違うGoroutineでも生きている保証はないんです。なぜかというと、waitGroupも何も使ってなく、ただ別のGoroutineで呼び出しているだけなので、呼び出し元のGoroutineが待ってくれる保証がないからです。

呼び出し元の親Goroutineが先に終了した場合、子Goroutine内ではctx<-Done()になってしまい、エラーを返してくるわけです。

ここまでのまとめ

親子間で異なるcontextを使った場合は、何の心配も不要です。ご自由にご利用ください。

以下、親子間で同じcontextを使いまわした場合の挙動をまとめます。

同期処理の場合
- 親Goroutineが子Goroutine処理を待ってくれるので、contextを使いまわしても問題ない
非同期処理の場合
- 親Goroutineが子Goroutine処理を待ってくれる保証はないので、親Goroutineが先に終了すると、子Goroutine内のctx<-Done()になる。もちろんその次の挙動は実装によります。上記の例では、ctx.Err()を返してくれていたので、エラーになったわけです。

同じ`context`を使い回す意味

「じゃ、毎回違うcontextを使ったらいいじゃん」と思うかもしれませんが、一度原点に戻って考えてみます。そもそもcontextは、プロセス間を横断する必要のあるリクエストスコープな値を伝達させるためのものです。

Package context defines the Context type, which carries deadlines, cancellation signals, and other request-scoped values across API boundaries and between processes.

（訳）パッケージコンテキストはコンテキストタイプを定義します。コンテキストタイプは、APIの境界を越えて、プロセス間で、期限、キャンセルシグナル、およびその他のリクエストスコープの値を伝達します。
出典：pkg.go.dev - context pkg

今回はトレースログを収集するために、同じcontextを使い回すことが必要でした。親子間で異なるcontextを使ったら、途中でトレースが追えなくなってしまいます。

考えられる対応策

元のcontextの値を引き継いだ、新しいcontextを生成することが求められます。実際contextパッケージ内には、次のようなメソッドが提供されています。いずれも元のcontextの値を引き継ぎながら、オプション的な機能が追加された新しいcontextが生成されます。

context.go

func WithCancel(parent Context) (ctx Context, cancel CancelFunc)
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc)
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)
func WithValue(parent Context, key, val any) Context

しかし、今回は元のcontextとまったく同じcontextを生成するだけで物足ります。オプション的な機能は必要ありません。既存のメソッドを利用すると、使用しない機能が追加された（ゴミデータが混ざっている）contextが生成されるので、使用を躊躇しちゃいました。

ということでたどり着いたのが、次に示す方法です。

採用した対応策

元のcontextの値のみをコピーした、新しいcontextを生成して使用することにしました。公式のcontextパッケージを眺めてみると、Contextの定義は次のようになっています。

context.go

type Context interface {
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}
	Err() error
	Value(key any) any
}

公式のcontextを内包した構造体を作り、上記interfaceを実装すれば行けそうな気がします。そして、実際実装したcustomｃontextの中身がこちらです。

customcontext.go

package customcontext

import (
	"context"
	"time"
)

type customContext struct {	ctx context.Context	}

func (customContext) Deadline() (time.Time, bool) { return time.Time{}, false }
func (customContext) Done() <-chan struct{} { return nil }
func (customContext) Err() error { return nil }
func (x customContext) Value(key any) any { return x.ctx.Value(key) }

func Clone(ctx context.Context) context.Context { return customContext{ctx} }

使用する箇所ではcustomcontext.Clone()メソッドで、値がコピーされた新しいcontextを作って使用するだけです。これで、元のcontextの値を引き継いでおり、且つ元のcontextとは生存関係を持たない新しいcontextが作れました。
めでたし❣　めでたし❣

func do(ctx context.Context) {
  newCtx := customcontext.Clone(ctx)
  go doSomething(newCtx)
}

参考資料