APIリクエスト内で同一関数呼び出しをキャッシュする仕組みとGo実装

ソフトウェアエンジニアの福井です。

Rettyでは小さなデータを複数取得し、それらを組み合わせて1つの大きなデータにする処理がいくつかあります。そして小さなデータ取得の判定に使うためのデータを取得することがあります。例えば飲食店の画像、メニュー、電話番号を取得し、それらを組み合わせた飲食店データを返すAPIがあります。この際、会員種別によって画像、メニュー、電話番号の取得方法が異なります。いずれのデータもSQLで取得しています。
通常は同じSQLを何回も実行しないよう、会員種別を前段で取得したあとに、それを使ってそれぞれの小さなデータを取得することになると思います。ただ再利用性などのいくつかの観点からそれぞれの小さなデータの直前でデータの分、会員種別を取得したくなりました。そうすると同じSQLを何回も実行してしまうという問題に戻ります。

問題を解消する仕組み

この問題を解消するため、1回のAPIリクエスト内で同じ関数呼び出しをキャッシュする仕組みを作りました。 キャッシュは他のAPIリクエストと共有しません。 これはGraphQLのN+1解消で使われるDataLoaderとは異なります。DataLoaderは1回のAPIリクエスト内の同じ関数呼び出しをバッチにできます。

実際のコード

言語はGo(1.23)でプロトコルはgRPC(google.golang.org/grpc v1.68.0)です。

package cache

// import ...

type cacheKeyLock struct {
    sync.Map
}

func (c *cacheKeyLock) loadOrStoreLock(key uint32) *sync.Mutex {
    actual, _ := c.LoadOrStore(key, &sync.Mutex{})
    return actual.(*sync.Mutex)
}

var globalCacheKeyLock = new(cacheKeyLock)

type cacheStoreCtxKey struct{}

func newCacheStore(ctx context.Context) context.Context {
    return context.WithValue(ctx, cacheStoreCtxKey{}, &sync.Map{})
}

func getCacheStore(ctx context.Context) *sync.Map {
    return ctx.Value(cacheStoreCtxKey{}).(*sync.Map)
}

func Wrap[T any](f T) T {
    rfn := reflect.ValueOf(f)

    w := reflect.MakeFunc(
        rfn.Type(),
        func(in []reflect.Value) []reflect.Value {
            h := fnv.New32a()
            _, _ = h.Write([]byte(runtime.FuncForPC(rfn.Pointer()).Name()))

            for i := 1; i < len(in); i++ {
                setTypeValueForKey(h, in[i].Kind(), in[i])
                _, _ = h.Write([]byte{31})
            }

            cacheKey := h.Sum32()

            ctx := in[0].Interface().(context.Context)
            s := getCacheStore(ctx)
            if v, ok := s.Load(cacheKey); ok {
                return v.([]reflect.Value)
            }

            lock := globalCacheKeyLock.loadOrStoreLock(cacheKey)
            lock.Lock()
            defer lock.Unlock()

            if v, ok := s.Load(cacheKey); ok { // 同じ関数呼び出しがないか探す
                return v.([]reflect.Value)
            }

            rtn := rfn.Call(in)
            s.Store(cacheKey, rtn) // 同じ関数呼び出しの返り値をキャッシュする

            return rtn
        },
    )

    return w.Interface().(T)
}

func setTypeValueForKey(h hash.Hash32, kind reflect.Kind, v reflect.Value) {
    switch kind {
    case reflect.String:
        _, _ = h.Write([]byte(v.String()))
        return
    case reflect.Uint64:
        _ = binary.Write(h, binary.LittleEndian, v.Uint())
        return
    case reflect.Slice:
        for i := range v.Len() {
            v2 := v.Index(i)
            setTypeValueForKey(h, v2.Kind(), v2)
            _, _ = h.Write([]byte{31})
        }
        return
    }
    // other kinds...

    panic(fmt.Errorf("not implemented: %v", kind))
}

func Interceptor() grpc.UnaryServerInterceptor {
    return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
        return handler(newCacheStore(ctx), req)
    }
}

キャッシュはインメモリで持ちます。キャッシュ生存期間は1回のAPIリクエスト単位になります。
利用側はcache.Interceptor()でinterceptorを設定し、以下のように返り値をキャッシュしたい関数を引数にした関数を実行します。

r := cache.Wrap(<キャッシュしたい関数>)(ctx, <キャッシュしたい関数の引数>)

例えばresult, err := GetMembership(ctx, restaurantID)と書いていたところは、
result, err := cache.Wrap(GetMembership)(ctx, restaurantID)となります。

処理の流れ

1. APIリクエストが来たらキャッシュストアを生成しcontextに入れる
2. Wrapした関数を実行する際、contextからキャッシュストアを取得する
3. キャッシュストアから関数と引数のkeyを探す
   1. なければ関数を実行し、返り値をキャッシュストアに入れ、返り値を返す
   2. あればキャッシュをそのまま返す
4. APIリクエスト完了とともにキャッシュストア自体が解放される

より具体的にすると以下の図のような流れになります。

関数と引数をkeyにするため、引数によって返り値が一意にならない関数にはこの仕組みは使えません。

ポイント

keyの作り方
ハッシュ関数はFNVを使用しています。MD5やSHA-256より速いです。FNVは暗号学的強度はないですがkeyにするだけのため、強度は必要ありません。

関数名はruntime.FuncForPC(rfn.Pointer(<関数のuintptr>)).Name()で取得しています。
関数のimportパスがgithub.com/RettyInc/restaurant/internalで関数名GetMembershipの場合、
文字列github.com/RettyInc/restaurant/internal.GetMembershipが取得できます。
この関数がレシーバー名をRestaurantを持つ場合、
文字列github.com/RettyInc/restaurant/internal.(*Restaurant).GetMembershipになります。

引数の値はreflect.Kindを判定してbyteにします。sliceやmapなどの場合は再起的に判定します。
fmt.Printfなども似たような再起的な処理をしています。 fmt.Printfは内部でfmr.printValue関数を呼び、これはreflect.Kindを判定して出力値をbufに書き込みます。 fmt.Printf("%v",

keyごとのロック
並列で実行することでキャッシュが使われずにWrapした関数を複数叩くことになります。そのためキャッシュストアの操作をロックしていますが、異なるkeyの処理までロックする必要はありません。そのためkeyごとにロックしています。 下段の飲食店ID: 123のAPIリクエストは前の処理を待ちます。そしてロックが解除されると前の返り値がキャッシュされるのでそれを使います。飲食店ID: 456の処理は飲食店ID: 123の処理を待つ必要はありません。

まとめ

Goで1回のAPIリクエスト内の同一関数呼び出しをキャッシュする仕組みと実装を紹介しました。これにより小さなデータの直前でデータの分、会員種別を取得しても同じSQLを何回も実行してしまうことはないです。一方、コードで1つのWrap関数を見ただけでは他のどこでキャッシュされうるかがわかりにくいというデメリットもあります。乱用せずに必要な場面で使うことが重要です。

※この記事はRetty Advent Calendar 2024の14日目の記事でした adventar.org