Spectral Contexts in Go (ファントム型を使った Context)

まず初めに Go を書いてるみなさんは良くこのようなコードを見かけませんか？

定義 A:

type UserID string

type userIDKey struct{}

func GetUserID(ctx context.Context) (UserID, bool) {
  val, ok := ctx.Value(userIDKey{}).(UserID)
  return val, ok
}

func WithUserID(ctx context.Context, userID UserID) {
  return context.WithValue(ctx, userIDKey{}, userID)
}

Web アプリケーションの場合、各ハンドラで定義 A の関数を利用するとこう書けます。

ctx := req.Context() // リクエストの context.Context

// JWT などから取り出した User ID を context へセット
ctx = WithUserID(ctx, "user-id")

// ビジネスロジックで context.Context が渡ってきたときに取得できる
fmt.Println(GetUserID(ctx))

このパターンを用いると、例えば teamIDKey を新しく作成する必要が出てきた時に定義 A に似たコードを書く必要があります。やりたいことに対してちょっとコード量が多いですね。

これを Go のジェネリクスを使ってシンプルに書き上げようという話です。

ファントム型を使って書き直す

ファントム型^[1]を使用し値を context に渡すコードは次のようになります。

package ctxutil

import "context"

// key は、型 T に基づくファントム型のキーを定義します。
type key[T any] struct{}

// WithValue は、指定された context に型 T の値をアタッチします。
// 新しい作成される context は、アタッチされた値を持ちます。
func WithValue[T any](ctx context.Context, val T) context.Context {
  return context.WithValue(ctx, key[T]{}, val)
}

// Value は、指定された context から型 T の値を取得します。
// 値が存在しない場合、型 T のゼロ値が返されます。
func Value[T any](ctx context.Context) (T, bool) {
  val, ok := ctx.Value(key[T]{}).(T)
  return val, ok
}

ここでは ctxutil をパッケージ名にしました。良い名前を考えるのは難しいですね。
このパッケージの関数を利用すると定義 A を利用したコードと同じことができます。

type UserID string

ctx := req.Context() // リクエストの context.Context

// JWT などから取り出した User ID を context へセット
ctx = ctxutil.WithValue[UserID](ctx, "user-id")

// ビジネスロジックで context.Context が渡ってきたときに取得できる
fmt.Println(ctxutil.Value[UserID](ctx))

これを利用すると新しくキーとなる型を定義してあげるだけで良くなります。

https://go.dev/play/p/0s7fFSJPYte

package main

import (
	"context"
	"fmt"

	"play.ground/ctxutil"
)

type UserID string

type TeamID string

func main() {
	ctx := context.Background()

	ctx = ctxutil.WithValue[UserID](ctx, "user-id")
	ctx = ctxutil.WithValue[TeamID](ctx, "team-id")

	fmt.Println(ctxutil.Value[UserID](ctx))
	fmt.Println(ctxutil.Value[TeamID](ctx))
}

利用用途を明確にしたい方もいると思うので、その場合は関数に閉じましょう。

type UserID string

func GetUserID(ctx context.Context) (UserID, bool) {
  return ctxutil.Value[UserID](ctx)
}

func WithUserID(ctx context.Context, userID UserID) context.Context {
  return ctxutil.WithValue[UserID](ctx, userID)
}

参考

https://hypirion.com/musings/spectral-contexts-in-go

脚注

1. 幽霊型(Phantom Type)とは実行時には存在しないけれども、コンパイル時に静的に型チェックされるような型のことです。構造体などのデータ型は、ジェネリック型パラメータを一つ余分に持ち、それをマーカーとして使ったりコンパイル時の型検査に使ったりすることができます。https://doc.rust-jp.rs/rust-by-example-ja/generics/phantom.html#幽霊型パラメータ