【GraphQL × Go】 N+1問題を解決するgqlgen + dataloaderの実装方法とCacheの実装オプション

こんにちは。バクラク事業部でバクラク申請の開発を担当している@sh_komineです。この記事は、6月から始まっている #LXベッテク月間 9日目の記事です。前回はPrivacyTech事業部の@cipepserさんによる合成データとは - 統計的な有用性を維持する架空のパーソナルデータでした。ものすごくBet Technologyな合成データのお話で読んでいてワクワクする記事です。気になる方は是非読んでみてください！

本日は、一般的なWebアプリケーション開発の技術で、バクラク事業部の開発で実際に使っているgqlgenとdataloaderの実装について紹介したいと思います。 gqlgen + dataloaderの記事自体は巷にだんだんと出揃ってきていると思いますが、弊社が使っている技術として改めてご紹介できたらと思います。

前提の話

今回の記事は自分が前回のエンジニアブログで書いた【GraphQL × Go】gqlgenの基本構成とオーバーフェッチを防ぐmodel resolverの実装で紹介できなかったDataLoaderの実装について紹介をしていきます。 tech.layerx.co.jp

gqlgenの基本的な使い方についてはこちらの記事から読まれることをお勧めします。

GraphQLの利便性とN+1問題とDataLoaderの役割

前段の記事でも触れていますが、GraphQLの大きな利便性はクライアント側から必要な時に必要なリソースを要求することができ、これにより不要なデータフェッチ（オーバーフェッチ）を防ぐことができることです。

オーバーフェッチ = リクエスト元で必要ないのに余分にリソースをフェッチしてしまうこと

少しおさらいすると、gqlgenではqueryで指定したtypeとgoで定義したstruct差分がある時にはmodel resolverが作成され、 queryに指定された時のみmodel resolverのメソッドが実行されるところまでご紹介しました。（サンプルコードは前回に引き続き、一般的なTODOリストのアプリケーションです。）

graph/schema.graphqls

type Todo {
  id: ID!
  text: String!
  done: Boolean!
  user: User! # ⭐️ graphqlではUserを指定
}

type User {
  id: ID!
  name: String!
}

#  type Query: データフェッチ系のエンドポイント定義
type Query {
  todos: [Todo!]!
}

graph/model/models.go

package model

type Todo struct {
    ID   string `json:"id"`
    Text string `json:"text"`
    Done bool   `json:"done"`
    UserID *string `json:"user_id"` //  ⭐️ todoスキーマに持っているUserIDまで定義
}

todoの一覧取得でuserも取得するクエリ

query listTodos {
  todos {
    text
    done
    user {
      name
    }
  }
}

graph/schema.resolvers.go

func (r *todoResolver) User(ctx context.Context, obj *model.Todo) (*model.User, error) {
        // ⭐️ query指定があった時だけ、このメソッドが呼び出される
        r.userRepo.GetUserByTodoID(obj.ID)
}

// ...

// Todo returns generated.TodoResolver implementation.
func (r *Resolver) Todo() generated.TodoResolver { return &todoResolver{r} }

// ...

type todoResolver struct{ *Resolver }

前回きちんと触れることはできなかったのですが、上記の実装はN+1問題という有名な問題を抱えています。

N+1問題 = N件のレコードを取得した時、関連レコードの取得にN回別のフェッチを行い、合わせてN+1回のフェッチとなってしまうこと

今回の実装で具体的なシーケンス図を描くと以下のような流れになります。 TodoがN = 3件ある時、Userの取得フェッチは3回走ることとなります。 Todoの取得で1回、Userの取得で3回なので、N + 1 = 3 + 1ですね。これは非常に効率が悪い、、、。

（※ 今回はわかりやすさのため、データ取得をSQLクエリで書きましたが、現実ではマイクロサービスの別のAPIなどを呼び出すことも多いかと思います。）

ここにDataLoaderを導入するとこうなります。（※ DataLoaderのクラスはもっと細かくありますが、ここでは簡略化して書いています。）

この図をみるとデータフェッチが激減していることがわかります。今回はN = 3のサンプルですが、N = 100などとなった時に、大きくパフォーマンスが変わることは想像に難くないです。 DataLoaderが各userResolverからのデータフェッチを一時受けして、一定時間待ち合わせた上でまとめて取得してきてくれるおかげで、N+1回のデータフェッチが2回のデータフェッチとなりました。物凄いパフォーマンス改善です。画期的です。

このDataLoaderの技術はGraphQLの仕組みに欠かせない技術となっています。ちなみに、DataLoaderの待ち合わせて読み込む仕組みを遅延読み込み（Lazy Loading）と言います。遅延読み込みの説明は以下の記事がわかりやすく勉強になりましたので、ぜひご参照ください。感謝。

GraphQL と N+1 SQL 問題と dataloader - Qiita

dataloaderの実装サンプル

さて、具体的な実装サンプルもご紹介します。 golangのDataLoaderの実装はいくつかありますが、私たちのチームではgqlgenのreference「Optimizing N+1 database queries using Dataloaders — gqlgen」でも紹介されている graph-gophers/dataloader を利用しています。

今回はこちらのreferenceの実装方法を参考にしつつ、自分なりに少しわかりやすくリファクタしたものをサンプルとしました。

github.com

まず、dataloaderのuserの取得処理を実装したloader/user.goです。

package loader

import (
    "context"
    "fmt"
    "github.com/graph-gophers/dataloader"
    "github.com/shkomine/gqlgen-todos/graph/model"
    "github.com/shkomine/gqlgen-todos/repository"
    "log"
    "strings"
)

type UserLoader struct {
    userRepo repository.User
}

// BatchGetUsers dataloader.BatchFuncを実装したメソッド
// ユーザーの一覧を取得する
func (u *UserLoader) BatchGetUsers(ctx context.Context, keys dataloader.Keys) []*dataloader.Result {
    // dataloader.Keysの型を[]stringに変換する
    userIDs := make([]string, len(keys))
    for ix, key := range keys {
        userIDs[ix] = key.String()
    }
    // 実処理
    log.Printf("BatchGetUsers(id = %s)\n", strings.Join(userIDs, ","))
    userByID, err := u.userRepo.GetMapInIDs(ctx, userIDs)
    if err != nil {
        err := fmt.Errorf("fail get users, %w", err)
        log.Printf("%v\n", err)
        return nil
    }
    // []*model.User[]*dataloader.Resultに変換する
    output := make([]*dataloader.Result, len(keys))
    for index, userKey := range keys {
        user, ok := userByID[userKey.String()]
        if ok {
            output[index] = &dataloader.Result{Data: user, Error: nil}
        } else {
            err := fmt.Errorf("user not found %s", userKey.String())
            output[index] = &dataloader.Result{Data: nil, Error: err}
        }
    }
    return output
}

// LoadUser dataloader.Loadをwrapして型づけした実装
func LoadUser(ctx context.Context, userID string) (*model.User, error) {
    log.Printf("LoadUser(id = %s)\n", userID)
    loaders := GetLoaders(ctx)
    thunk := loaders.UserLoader.Load(ctx, dataloader.StringKey(userID))
    result, err := thunk()
    if err != nil {
        return nil, err
    }
    user := result.(*model.User)
    log.Printf("return LoadUser(id = %s, name = %s)\n", user.ID, user.Name)
    return user, nil
}

一つめのBatchGetUsersは dataloader.BatchFunc を実装しています。

graph-gophers/dataloader/dataloader.go

// BatchFunc is a function, which when given a slice of keys (string), returns an slice of `results`.
// It's important that the length of the input keys matches the length of the output results.
//
// The keys passed to this function are guaranteed to be unique
type BatchFunc func(context.Context, Keys) []*Result

この dataloader.BatchFuncは dataloader.NewBatchLoaderの第一引数に指定することで、dataloader.Loaderの初期化に利用します。そしてdataloader.BatchFuncはdataloader.Loadを一定待ち合わせた後に、dataloader.Loaderの中から呼び出されます。

このメソッドのメインの処理としては、userRepo.GetMapInIDsを呼び出していています。このrepository/user.goは一般的なRepositoryなので詳しい説明は省略しますが、実装ではdbやapiからデータをフェッチする処理が書かれているイメージです。

repository/user.go

type User interface {
    GetByID(ctx context.Context, id string) (*model.User, error)
    GetMapInIDs(ctx context.Context, ids []string) (map[string]*model.User, error)
}

// ...

二つめのLoadUserメソッドはdataloader.LoadをUser用にwrapした関数で、todoResolverから以下のように呼び出します。後述しますが、loaderの本体は context.Context にMiddlewareでInjectするので、呼び出し元からはシンプルにloader.LoadUserを呼び出すことができます。

schema.resolvers.go

func (r *todoResolver) User(ctx context.Context, obj *model.Todo) (*model.User, error) {
    if obj.UserID == nil {
        return nil, nil
    }
    user, err := loader.LoadUser(ctx, *obj.UserID)
    if err != nil {
        return nil, err
    }
    return user, nil
}

// ...

type todoResolver struct{ *Resolver }

次に、各ドメインのローダーを束ねて、contextにInjectするloader/loaders.goの実装です。

package loader

import (
    "context"
    "database/sql"
    "github.com/graph-gophers/dataloader"
    "github.com/shkomine/gqlgen-todos/repository"
    "net/http"
)

type ctxKey string

const (
    loadersKey = ctxKey("dataloaders")
)

// Loaders 各DataLoaderを取りまとめるstruct
type Loaders struct {
    UserLoader *dataloader.Loader
}

// NewLoaders Loadersの初期化メソッド
func NewLoaders(conn *sql.DB) *Loaders {
    // define the data loader
    userLoader := &UserLoader{
        userRepo: repository.NewUserRepo(conn),
    }
    loaders := &Loaders{
        UserLoader: dataloader.NewBatchedLoader(userLoader.BatchGetUsers),
    }
    return loaders
}

// Middleware LoadersをcontextにインジェクトするHTTPミドルウェア
func Middleware(loaders *Loaders, next http.Handler) http.Handler {
    loaders.UserLoader.ClearAll()
    // return a middleware that injects the loader to the request context
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        nextCtx := context.WithValue(r.Context(), loadersKey, loaders)
        r = r.WithContext(nextCtx)
        next.ServeHTTP(w, r)
    })
}

// GetLoaders ContextからLoadersを取得する
func GetLoaders(ctx context.Context) *Loaders {
    return ctx.Value(loadersKey).(*Loaders)
}

このサンプルではUserLoaderしかないですが、他のドメインの初期化もここに書かれます。上記のloaderをmain関数の中でGraphQLに使うエンドポイントにInjectします。

server.go

package main

import (
    "database/sql"
    "github.com/shkomine/gqlgen-todos/loader"
    "github.com/shkomine/gqlgen-todos/repository"
    "log"
    "net/http"
    "os"

    "github.com/99designs/gqlgen/graphql/handler"
    "github.com/99designs/gqlgen/graphql/playground"
    _ "github.com/go-sql-driver/mysql"
    "github.com/shkomine/gqlgen-todos/graph"
    "github.com/shkomine/gqlgen-todos/graph/generated"
)

const defaultPort = "8080"

func main() {
    port := os.Getenv("PORT")
    if port == "" {
        port = defaultPort
    }

    // dbの初期化
    db, err := sql.Open("mysql", "root:password@tcp(127.0.0.1:13306)/test_db")
    if err != nil {
        log.Fatalf("main sql.Open error err:%v", err)
    }
    defer db.Close()

    // loaderの初期化
    ldrs := loader.NewLoaders(db)

    rslvr := graph.Resolver{
        TodoRepo: repository.NewTodoRepo(db),
        UserRepo: repository.NewUserRepo(db),
    }
    srv := handler.NewDefaultServer(
        generated.NewExecutableSchema(generated.Config{Resolvers: &rslvr}),
    )

    http.Handle("/", playground.Handler("GraphQL playground", "/query"))
    http.Handle("/query", loader.Middleware(ldrs, srv))

    log.Printf("connect to http://localhost:%s/ for GraphQL playground", port)
    log.Fatal(http.ListenAndServe(":"+port, nil))
}

これを実際にGraphQL queryで呼び出してみると、、、。

query listTodos {
  todos {
    id
    text
    done
    user {
      id
      name
    }
  }
}

以下のように取得できました

{
  "data": {
    "todos": [
      {
        "id": "1",
        "text": "技術記事を書く",
        "done": false,
        "user": {
          "id": "1",
          "name": "田中 太郎"
        }
      },
      {
        "id": "2",
        "text": "実装をする",
        "done": false,
        "user": {
          "id": "2",
          "name": "佐藤 二郎"
        }
      },
      {
        "id": "3",
        "text": "日記を書く",
        "done": false,
        "user": {
          "id": "1",
          "name": "田中 太郎"
        }
      }
    ]
  }
}

なお、ソースコード上で仕込んだログは以下のように吐かれていて、きちんと待ち合わせて取得できていることがわかります。 DataLoaderめちゃくちゃ便利ですね。

2022/06/13 00:14:10 LoadUser(id = 1)
2022/06/13 00:14:10 LoadUser(id = 2)
2022/06/13 00:14:10 LoadUser(id = 1)
2022/06/13 00:14:10 BatchGetUsers(id = 1,2,1)
2022/06/13 00:14:10 return LoadUser(id = 1, name = 田中 太郎)
2022/06/13 00:14:10 return LoadUser(id = 1, name = 田中 太郎)
2022/06/13 00:14:10 return LoadUser(id = 2, name = 佐藤 二郎)

`dataloader`のキャッシュの実装オプション

なお、今回は極力gqlgenのreference「Optimizing N+1 database queries using Dataloaders — gqlgen」通りに実装したのですが、キャッシュの実装は改善の余地があります。

graph-gophers/dataloaderに以下のような一文があり、DataLoaderには、http request単位の短いライフサイクルを考慮した簡単なCacheが実装されていることが書かれています。

This implementation contains a very basic cache that is intended only to be used for short lived DataLoaders (i.e. DataLoaders that only exist for the life of an http request). You may use your own implementation if you want.

it also has a NoCache type that implements the cache interface but all methods are noop. If you do not wish to cache anything.

https://github.com/graph-gophers/dataloader#cache

確かに NewBatchedLoaderのNewCache実装をみると、InMemoryCacheが使われていて、 mapとsync.RWMutexを使ったシンプルな実装になっています。このCacheは dataloader.Load関数内の最初に呼び出されており、Cacheから取得できたリソースはbatch処理で取得しにいきません。サンプルの実装ではこの初期化時で生成したCacheをリクエストを跨いで使用しており、意図せぬ長い期間、リソースのキャッシュが残り続ける可能性があります。

graph-gophers/dataloader/dataloader.go

// NewCache constructs a new InMemoryCache
func NewCache() *InMemoryCache {
    return &InMemoryCache{
        items: &sync.Map{},
    }
}

// ...

// NewBatchedLoader constructs a new Loader with given options.
func NewBatchedLoader(batchFn BatchFunc, opts ...Option) *Loader {
    loader := &Loader{
        batchFn:  batchFn,
        inputCap: 1000,
        wait:     16 * time.Millisecond,
    }

    // Apply options
    for _, apply := range opts {
        apply(loader)
    }

    // Set defaults
    if loader.cache == nil {
        loader.cache = NewCache()
    }

    if loader.tracer == nil {
        loader.tracer = &NoopTracer{}
    }

    return loader
}

今回はキャッシュの実装オプションとして、以下の3通りのアプローチを紹介します。

そもそもCacheを使わない
Batch実行ごとにCacheをクリアする
http requestごとにLoaderを生成をすることでCacheのライフサイクルを制御する

1. そもそもCacheを使わない

dataloader.NewBatchedLoaderメソッドでは、第二引数以降にOptionを渡すことができるのですが、dataloaderにはCacheを外から設定できる dataloader.WithCache メソッドが事前に定義されています。キャッシュをしないdataloader.NoCacheも事前に定義されているので、これを設定することでキャッシュを無効化できます。常に最新のデータを取りたい場合、確実に待ち合わせができるような実装であれば、これで十分かと思います。

loader/loaders.go

// NewLoaders Loadersの初期化メソッド
func NewLoaders(conn *sql.DB) *Loaders {
    // define the data loader
    userLoader := &UserLoader{
        userRepo: repository.NewUserRepo(conn),
    }
    loaders := &Loaders{
        UserLoader: dataloader.NewBatchedLoader(
            userLoader.BatchGetUsers,
            dataloader.WithCache(&dataloader.NoCache{}), // ⭐️NoCacheを設定する
        ),
    }
    return loaders
}

2. Batch実行ごとにCacheをクリアする

またオプション設定なのですが、Batch実行の度にキャッシュをクリアするオプションも提供されています。この場合、リソースのキャッシュは同時に来たリクエストでは共有されます。同時に複数のrequestが呼び出された時にキャッシュを共有したいが、あまりキャッシュを長く持ちたくない場合などには向いているかと思います。

loader/loaders.go

// NewLoaders Loadersの初期化メソッド
func NewLoaders(conn *sql.DB) *Loaders {
    // define the data loader
    userLoader := &UserLoader{
        userRepo: repository.NewUserRepo(conn),
    }
    loaders := &Loaders{
        UserLoader: dataloader.NewBatchedLoader(
            userLoader.BatchGetUsers,
            dataloader.WithClearCacheOnBatch(), // ⭐ Batch実行ごとにCacheをクリア
        ),
    }
    return loaders
}

3. http requestごとにLoaderを生成をすることでCacheのライフサイクルを制御する

最後に、requestごとにLoaderを生成する方法です。

GraphQLのN+1問題を解決する DataLoaderの使い方 - 一休.com Developers Blog

こちらの記事でも触れられているのですが、Loaderの初期化をリクエストごとに生成することで、リクエストごとに確実にCacheを初期化することができます。弊社ではマルチテナントSaaSでrequestを跨いだキャッシュなどは持ちたくなかったことから、この方法で初期化をしています。

その場合はMiddleware内でLoadersの初期化を行うことで、確実にhttp request内に閉じた状態でCacheを持つことができます。

loader/loaders.go

// Middleware LoadersをcontextにインジェクトするHTTPミドルウェア
func Middleware(db *sql.DB, next http.Handler) http.Handler {
    // loaderの初期化
    ldrs := NewLoaders(db)

    // return a middleware that injects the loader to the request context
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        nextCtx := context.WithValue(r.Context(), loadersKey, ldrs)
        r = r.WithContext(nextCtx)

        next.ServeHTTP(w, r)
    })
}

キャッシュの実装オプションについて、簡単に3つ紹介してみました。「1. そもそもCacheを使わない」や「2. Batch実行ごとにCacheをクリアする」で紹介したように dataloader.NewBatchedLoaderメソッドでは、第二引数以降にいろいろなOptionを渡すことができます。今回紹介した以外にも、WithWait、WithInputCapacity、WithBatchCapacityなどを使えばバッチの待ち時間や同時実行数、インプットの上限などでパフォーマンスチューニングしたりもできますので、是非使い倒して快適なGraphQLライフをお楽しみください。