Go에서 context.WithValue() 안전하게 사용하기

Go에서 context.WithValue()를 사용해 값을 전달할 때, 습관적으로 string 타입을 key로 사용하곤 했습니다. 하지만 이 방식은 key 충돌과 같은 문제가 발생할 수 있습니다.

이 글에서는 안전한 context key 사용법과 주의사항에 대해서 알아보겠습니다.

string key 사용시 발생할 수 있는 문제

const UserIdKey string = "user_id"

func withUserId(ctx context.Context, user *User) context.Context {
	return context.WithValue(ctx, UserIdKey, user)
}

위의 코드는 go-staticcheck에서 경고를 발생시킨다.

should not use built-in type string as key for value; 
define your own type to avoid collisions (SA1029)

string 타입의 키를 사용했을때 문제는 다른 패키지나 라이브러리에서 동일한 문자열을 Key로 설정할 경우 충돌이 발생할 수 있다는 것이다. 원하지 않는 값 덮어쓰기가 발생해서 디버깅하기 어려운 버그가 생길 수 있다.

전용 타입을 정의해서 해결하자

type UserIdKeyType string

const UserIdKey UserIdKeyType = "user_id"

func withUserId(ctx context.Context, user *User) context.Context {
	return context.WithValue(ctx, UserIdKey, user)
}

이렇게 문자열 기반의 커스텀 타입을 정의해서 key로 사용할 경우 key 충돌이 발생하지 않는다.

그 이유는, Go의 context.WithValue의 key 비교는 값(value)과 타입(type)을 함께 보기 때문이다. 만약 두 개의 패키지 안에서 동일한 커스텀 타입을 정의해서 사용한다고 하더라도 패키지가 다르기 때문에 별개의 key로 등록된다.

// pkgA
package pkgA
type TransactionKeyType string
const TransactionKey TransactionKeyType = "tx"

// pkgB
package pkgB
type TransactionKeyType string
const TransactionKey TransactionKeyType = "tx"

ctx = context.WithValue(ctx, pkgA.TransactionKey, "txA")
ctx = context.WithValue(ctx, pkgB.TransactionKey, "txB")

// 각각 독립적으로 저장됨
valueA := ctx.Value(pkgA.TransactionKey) // txA
valueB := ctx.Value(pkgB.TransactionKey) // txB

정리

• string 타입의 키를 사용하면 충돌 위험이 있다.
• 각 서비스에서만 사용되는 key는 각 서비스의 특정 패키지에서 커스텀키를 정의해서 사용한다.
• 공통으로 사용되는 key는 공통라이브러리에서 커스텀키를 정의하고 사용한다.