Goのcontextのkeyについて

はじめに

Goの標準パッケージの context.Context は context.WithValue 関数を使って任意の値を格納することができます。

context.WithValue には以下のようなコメントが書かれています。

// WithValue returns a copy of parent in which the value associated with key is
// val.
//
// Use context Values only for request-scoped data that transits processes and
// APIs, not for passing optional parameters to functions.
//
// The provided key must be comparable and should not be of type
// string or any other built-in type to avoid collisions between
// packages using context. Users of WithValue should define their own
// types for keys. To avoid allocating when assigning to an
// interface{}, context keys often have concrete type
// struct{}. Alternatively, exported context key variables' static
// type should be a pointer or interface.

組み込み型よりも struct{} のように自分で定義した型を使ったほうが良い、というようなことが書いてあります。

この記述を勘違いして以下のようにキーを定義したところ、期待通りの値が取れたり取れなかったりするバグができてしまいました。

var hogeKey = struct{}{}

func SetHoge(ctx context.Context, i int) context.Context {
	return context.WithValue(ctx, &hogeKey, i)
}

func GetHoge(ctx context.Context) (int, bool) {
	i, ok := ctx.Value(&hogeKey).(int)
	return i, ok
}

どのように定義すると同じキーとみなされるのか、調査しました。

調査

以下の5パターンのキーを用意し、setterとgetterを作ります。

空の構造体struct{}の実体として定義する
空の構造体struct{}の実体として定義し、ポインタで使う
空の構造体struct{}のポインタとして定義する
空の構造体struct{}を型定義する
空の構造体struct{}を型定義し、実体として定義する

package a

import "context"

// 1. 空の構造体struct{}の実体として定義する
var emptyStruct = struct{}{}

func SetStruct(ctx context.Context, i int) context.Context {
	return context.WithValue(ctx, emptyStruct, i)
}

func GetStruct(ctx context.Context) (int, bool) {
	i, ok := ctx.Value(emptyStruct).(int)
	return i, ok
}

// 2. 空の構造体struct{}の実体として定義し、ポインタで使う
var emptyStructAsPointer = struct{}{}

func SetStructAsPointer(ctx context.Context, i int) context.Context {
	return context.WithValue(ctx, &emptyStructAsPointer, i)
}

func GetStructAsPointer(ctx context.Context) (int, bool) {
	i, ok := ctx.Value(&emptyStructAsPointer).(int)
	return i, ok
}

// 3. 空の構造体struct{}のポインタとして定義する
var emptyStructPointer = &struct{}{}

func SetStructPointer(ctx context.Context, i int) context.Context {
	return context.WithValue(ctx, emptyStructPointer, i)
}

func GetStructPointer(ctx context.Context) (int, bool) {
	i, ok := ctx.Value(emptyStructPointer).(int)
	return i, ok
}

// 4. 空の構造体struct{}を型定義する
type emptyStructType struct{}

func SetStructType(ctx context.Context, i int) context.Context {
	return context.WithValue(ctx, emptyStructType{}, i)
}

func GetStructType(ctx context.Context) (int, bool) {
	i, ok := ctx.Value(emptyStructType{}).(int)
	return i, ok
}

// 5. 空の構造体struct{}を型定義し、実体として定義する
var emptyStructTypeValue = emptyStructType{}

func SetStructTypeValue(ctx context.Context, i int) context.Context {
	return context.WithValue(ctx, emptyStructTypeValue, i)
}

func GetStructTypeValue(ctx context.Context) (int, bool) {
	i, ok := ctx.Value(emptyStructTypeValue).(int)
	return i, ok
}

これと全く同じキーや関数の定義を2つのpackageで用意します。

.
├── a
│   └── a.go
├── b
│   └── b.go
└── main.go

a.goとb.goはpackage名以外は全く同じです。

各キーでcontextに値をセットし、取得してみます。

package main

import (
	"context"
	"fmt"

	"ctxkey/a"
	"ctxkey/b"
)

func main() {
	for i, set := range []func(ctx context.Context, i int) context.Context{
		a.SetStruct,
		a.SetStructAsPointer,
		a.SetStructPointer,
		a.SetStructType,
		a.SetStructTypeValue,
	} {
		ctx := set(context.Background(), 1)
		for j, get := range []func(ctx context.Context) (int, bool){
			a.GetStruct,
			a.GetStructAsPointer,
			a.GetStructPointer,
			a.GetStructType,
			a.GetStructTypeValue,
			b.GetStruct,
			b.GetStructAsPointer,
			b.GetStructPointer,
			b.GetStructType,
			b.GetStructTypeValue,
		} {
			res, ok := get(ctx)
			fmt.Printf("set: %v get: %v result: %v %v\n", i, j, res, ok)
		}
	}
}

これを実行すると以下のようになります。

set: 0 get: 0 result: 1 true
set: 0 get: 1 result: 0 false
set: 0 get: 2 result: 0 false
set: 0 get: 3 result: 0 false
set: 0 get: 4 result: 0 false
set: 0 get: 5 result: 1 true
set: 0 get: 6 result: 0 false
set: 0 get: 7 result: 0 false
set: 0 get: 8 result: 0 false
set: 0 get: 9 result: 0 false
set: 1 get: 0 result: 0 false
set: 1 get: 1 result: 1 true
set: 1 get: 2 result: 1 true
set: 1 get: 3 result: 0 false
set: 1 get: 4 result: 0 false
set: 1 get: 5 result: 0 false
set: 1 get: 6 result: 1 true
set: 1 get: 7 result: 1 true
set: 1 get: 8 result: 0 false
set: 1 get: 9 result: 0 false
set: 2 get: 0 result: 0 false
set: 2 get: 1 result: 1 true
set: 2 get: 2 result: 1 true
set: 2 get: 3 result: 0 false
set: 2 get: 4 result: 0 false
set: 2 get: 5 result: 0 false
set: 2 get: 6 result: 1 true
set: 2 get: 7 result: 1 true
set: 2 get: 8 result: 0 false
set: 2 get: 9 result: 0 false
set: 3 get: 0 result: 0 false
set: 3 get: 1 result: 0 false
set: 3 get: 2 result: 0 false
set: 3 get: 3 result: 1 true
set: 3 get: 4 result: 1 true
set: 3 get: 5 result: 0 false
set: 3 get: 6 result: 0 false
set: 3 get: 7 result: 0 false
set: 3 get: 8 result: 0 false
set: 3 get: 9 result: 0 false
set: 4 get: 0 result: 0 false
set: 4 get: 1 result: 0 false
set: 4 get: 2 result: 0 false
set: 4 get: 3 result: 1 true
set: 4 get: 4 result: 1 true
set: 4 get: 5 result: 0 false
set: 4 get: 6 result: 0 false
set: 4 get: 7 result: 0 false
set: 4 get: 8 result: 0 false
set: 4 get: 9 result: 0 false

表にしました。

○をつけた箇所は、値を取得できた組み合わせです。

struct ,structAspointer ,structPointer から分かるように、struct{}{} をそのまま使うと、別のpackageであっても同じキーとみなされ、値が取れてしまいます。つまり、複数のpackageでstruct{}{} をキーに値をセットすると、後からセットした値で上書きされます。

	ctx := context.Background()
	ctx = a.SetStruct(ctx, 1)
	ctx = b.SetStruct(ctx, 2)
	fmt.Println(a.GetStruct(ctx)) // 2 true
	fmt.Println(b.GetStruct(ctx)) // 2 true

また、structType と structTypeValue から分かるように、同じ型の場合、実体としては別でも同じキーとして扱われてしまいます。そのため、キーごとに別々の型で定義する必要があります。

contextのキーについて調べてみると以下の記事を見つけました。

どのような条件で同じキーとみなされるか、詳しく書かれており、非常に参考になりました。

結論

contextのキーを定義する場合、キーごとに別々の型で定義する必要があります。

type hogeKey struct{}

func SetHoge(ctx context.Context, i int) context.Context {
	return context.WithValue(ctx, hogeKey{}, i)
}

func GetHoge(ctx context.Context) (int, bool) {
	i, ok := ctx.Value(hogeKey{}).(int)
	return i, ok
}

type fugaKey struct{}

func SetFuga(ctx context.Context, i int) context.Context {
	return context.WithValue(ctx, fugaKey{}, i)
}

func GetFuga(ctx context.Context) (int, bool) {
	i, ok := ctx.Value(fugaKey{}).(int)
	return i, ok
}

ちなみに、同じ型で複数のキーを作りたい場合は

type ctxKey struct{ string }

のようにstringを含めるように定義し、キーごとに別の文字列をセットすれば、別のpackageで同じ文字列をセットしていても、別のキーとして扱われます。

// package a

type ctxKey struct{ string }

var hogeKey = ctxKey{ "hoge" }
var fugaKey = ctxKey{ "fuga" }

// package b

type ctxKey struct{ string }

var hogeKey = ctxKey{ "hoge" }

// a.hogeKey, a.fugaKey, b.hogeKey はそれぞれ別のキーとして扱われる