【Go】 context の使い方

はじめに

業務で cotenxt を使う事があったので、整理する意味を込めて使い方を記事にまとめます。仕組みの部分はもう少し理解を深めたらまとめようかなと思っています。今回の記事は Go の基本を理解している方向けの内容になっています。

context の 3つの使い方

キャンセルの伝播

WithCancelを使うと、cancel関数を取得することができます。
ctx.Done() でキャンセル関数の実行をしているかどうかを検知することができる。

func doWork(ctx context.Context) {
	select {
	case <-ctx.Done():
		fmt.Println("Work completed")
	default:
		fmt.Println("default")
	}
}

func main() {
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	cancel()
	go doWork(ctx) 
	time.Sleep(1 * time.Second) // go doWork が非同期で走っているので、timeを入れている
}

最後の行にスリープを入れているのは go doWork(ctx) は非同期で実行されているからです。非同期なので sleepを入れないと doWork関数の実行が終わる前にmain関数が終了してしまい select処理 の処理が終わらず、プリントが出力されません。

出力

Work completed

仮にcancel() → defer cancel()に変更すると最後に context にキャンセルの伝達されるので defaultが出力される。

タイムアウトの伝播

WithTimeoutメソッド を利用すると、タイムアウトを設定したコンテキストを取得できます。

package main

import (
	"context"
	"fmt"
	"time"
)

func doWork(ctx context.Context) {
	select {
	case <-time.After(3 * time.Second):
		fmt.Println("Work completed")
	case <-ctx.Done():
		fmt.Println("Work cancelled due to timeout")
	}
}

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) // ここでタイムアウトを設定する
	defer cancel()
	go doWork(ctx)
	time.Sleep(4 * time.Second) // main関数が終了する前に dowork の実行し終わるために設定
}

出力

Work cancelled due to timeout

値の伝播

WithValueメソッドを使用することで、key-value の形式で含まれた context を新たに生成します。

type KeyType string

func doWork(ctx context.Context) {
	value := ctx.Value(KeyType("myKey"))  // context から値を取得している
	fmt.Printf("Received value: %v\n", value)
}

func main() {
	ctx := context.WithValue(context.Background(), KeyType("myKey"), "Hello from main!") // context に値を含めている
	go doWork(ctx)
	time.Sleep(1 * time.Second)
}

JWT を含める

業務で携わっているプロダクトは ミドルウェア層で context に JWT や TraceID などの情報をcontext に含めて、その context をアプリケーション層で使用しています。

以下がコードの例です。

func JWTMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		token := r.Header.Get("Authorization")

		if token == "" {
			next.ServeHTTP(w, r)
			return
		}

		jwt, err := validateJWT(token) //JWTの検証
		if err != nil {
			http.Error(w, "Invalid JWT", http.StatusForbidden)
			return
		}

		ctx := contextWithJWT(r.Context(), jwt) // JWTをcontext 煮含める
		next.ServeHTTP(w, r.WithContext(ctx)) // 次のミドルウェアに渡す
	})
}

小ネタ

パッケージの関数で引数に context を求められることがよくあり、中身をたどるとこの使われ方をされているのをよく見ます。いくつか例を上げてみます。

これはboiler パッケージで context が使われている例であり、デバッグモードが有効かどうかを判別する処理になっています。詳しくはパッケージの中身を見てみてください。

func IsDebug(ctx context.Context) bool {
	debug, ok := ctx.Value(ctxDebug).(bool)
	if ok {
		return debug
	}
	return DebugMode
}

firebase のパッケージの中身で http クライアントを取得する処理です。

func ContextClient(ctx context.Context) *http.Client {
	if ctx != nil {
		if hc, ok := ctx.Value(HTTPClient).(*http.Client); ok {
			return hc
		}
	}
	if appengineClientHook != nil {
		return appengineClientHook(ctx)
	}
	return http.DefaultClient
}

最後に

今までパッケージの中身を見ることはなかったが、ctx って辿ってみると意外とシンプル使われ方をしているなと思いました。もう少し理解が深まったら context の仕組みの部分も記事にできたらなと思っています。