goでcontext.WithValue()のget/setを手軽に書く

愚直に定義

以下のように書くのがめんどくさい。

type ctxKey string

const ctxTraceIDkey ctxKey = "TraceID"

type TraceID string

func WithContext(ctx context.Context, traceID TraceID) context.Context {
	return context.WithValue(ctx, ctxTraceIDkey, traceID)
}
func FromContext(ctx context.Context) TraceID {
	v, _ := ctx.Value(ctxTraceIDkey).(TraceID)
	return v
}

ちょっと冗長な書き方をしている。個別に型さえ定義しておけば値は何でも良いのだけど、個人的にはstruct{}などを使うとfmt.Printlnとかで値を確認したときなどに何がなんだかわからなくなって嫌という感覚がある。そのせいでctxTraceIDKey周りがだるい。

必要なのは以下の２点

• keyの一意性
• 煩雑な定義をサボる

fmt.Sprintf("%T") は一意では？

元々の定義で毎回型を個別に定義する理由は同じリテラルから作られたとしても型が異なればユニークということが保証されるためだった。例えば、Type X string; type Y stringと定義された別々の型X,Yは同じ""から作られても別の値として扱われる。

ところで、普通に考えてfmt.Sprintf("%T")が返す文字列は「あるパッケージに所属する型名」となるのでこれもまたユニークになるはず。別に型を個別に定義することにこだわらなくても良い。

fmt.Sprintf("%T", TraceID(""))  // => main.TraceID

ちなみに実装を覗いてみるとreflect.TypeOf(v).String()が使われているらしい。こちらを使っても良い。

genericsでどうにかなるのでは？

さて、ユニークであることを保証することができたので、今度は定義自体を手軽にする方法を考えてみよう。

以下のような型を持つ関数を定義してあげれば良さそうだ。ちょっとreactのhookっぽい形で定義してみる。型変数を一つ取るgenerics関数になっている。型を指定して呼び出すとgetter,setterを返すような関数（別にこの辺はstructを作ってメソッドでアクセスしても良い。好み。）。

func UseContextAccessor[T any]() (get func(context.Context) T, set func(context.Context, T) context.Context) 

以下のように渡した型のゼロ値を用いてkeyを生成する。

func UseContextAccessor[T any]() (get func(context.Context) T, set func(context.Context, T) context.Context) {
	var z T
	k := ctxKey(reflect.TypeOf(z).String()) // <package>.<type>名なのでunique

	get = func(ctx context.Context) T {
		v, _ := ctx.Value(k).(T)
		return v
	}
	set = func(ctx context.Context, v T) context.Context {
		return context.WithValue(ctx, k, v)
	}
	return
}

冒頭のTraceIDのaccessorの定義は以下のようになる。

var GetTraceID, WithTraceID = UseContextAccessor[TraceID]()

作ったaccessorは冒頭のものと同じように使える。

ctx = WithTraceID(ctx, TraceID("xxx"))
ctx = WithTraceID(ctx, TraceID("yyy"))

fmt.Println(GetTraceID(ctx)) // yyy

文字列だと比較が遅いのでは？

たしかに、contextのwith valueの探索は原理的にはリストの線形探索なのでO (N)でequalのチェックが文字列だとコストが嵩むとかありそうだけど、アプリケーション開発程度なら誤差だよ誤差という気がしてる。

気になるならuintptrのようなもので定義しても良い。これもuniqueになるはず？
たぶん、reflect.TypeOf()から追って言ってunsafe.Pointerか何か経由で型ごとに別れたアドレスを見つけるはずなのでその値をkeyにしてあげれば良い。

generics無しに定義できないか？

呼び出し方を考えてみると、ゼロ値を使ってreflect.TypeOf()でkeyを作っているのだから型を指定するかわりに引数として値を渡せば良いような気もしてくる。

UseContextAccessor[TraceID]()

のかわりに

UseContextAccessor(TraceID(""))

としてやれば良いと思ってしまう。

しかしこれではgetter,setterの型がanyになってしまう。これではあまり嬉しくない。返ってくる型が固定されたから嬉しいかったのだった。そしてこれはgenerics以後にようやく記述できるようになったものだった。