DataLoaderを用いてN+1 問題対処

間違っている部分、エラーが発生する部分、わかりにくい部分などがございましたら、是非ご指摘をお願い致します。

今回はPart5になります。

基本的な事をメインに行っているので前回までの内容を知らなくても基本的には問題ありません。前回までの内容をgithubから取得して来ても良いかと思います。（importのgithub.comを調整する必要があるかも知れないですが...）

前回：一人のユーザから複数のTodoを取得する処理を実装しました。

この記事の内容

前回一人のユーザから複数のTodoを取得する処理で発覚したN+1問題を対処していきます。

前回のおさらい
DataLoaderの実装

Github Repo

この記事で紹介するコードはGithubに置いています。

Part5までの切り取り

前回のおさらい

query users  {
  users{
    id,
    todos{
      id,
      text
    }
  }
}

上記のようにusersQuery内でtodosという複数の値を取得する際に下記のような複数のクエリを発行しているという問題がありました。usersの取得という１回と追加でユーザ（N人）に紐づいたtodoの取得の処理でN+1回のクエリを取得している為N＋1問題と言われています。

2022/12/31 09:35:54 /app/graph/schema.resolvers.go:65
[6.096ms] [rows:0] SELECT * FROM "todos" WHERE user_id = 'U93'

2022/12/31 09:35:54 /app/graph/schema.resolvers.go:65
[1.829ms] [rows:0] SELECT * FROM "todos" WHERE user_id = 'U12'

2022/12/31 09:35:54 /app/graph/schema.resolvers.go:65
[49.961ms] [rows:0] SELECT * FROM "todos" WHERE user_id = 'U3'

2022/12/31 09:35:54 /app/graph/schema.resolvers.go:65
[45.902ms] [rows:0] SELECT * FROM "todos" WHERE user_id = 'U87'

2022/12/31 09:35:54 /app/graph/schema.resolvers.go:65
[52.091ms] [rows:0] SELECT * FROM "todos" WHERE user_id = 'U38'

2022/12/31 09:35:54 /app/graph/schema.resolvers.go:65
[58.972ms] [rows:0] SELECT * FROM "todos" WHERE user_id = 'U1'

2022/12/31 09:35:54 /app/graph/schema.resolvers.go:65
[61.057ms] [rows:2] SELECT * FROM "todos" WHERE user_id = 'U21'

2022/12/31 09:35:54 /app/graph/schema.resolvers.go:65
[63.121ms] [rows:0] SELECT * FROM "todos" WHERE user_id = 'U47'

2022/12/31 09:35:54 /app/graph/schema.resolvers.go:65
[65.210ms] [rows:0] SELECT * FROM "todos" WHERE user_id = 'U49'

DataLoaderの実装

他にもDataLoaderはありますが、今回はgqlgenの公式サイトで紹介されているgithub.com/graph-gophers/dataloaderを使用していきます。

DataLoaderは上記のN＋1問題を二つのSQL文にまとめることで解決します。
その仕組みとしては、「一定時間待って、その間に実行されたデータ取得リクエストをバッチ化する」というアプローチを取っています。

SELECT * FROM "users";
SELECT * FROM "todos" WHERE user_id in (U93,U49,U47,U21,U1,U38,....);

それでは、実装していきます。

最初にdockerを起動します。

$ docker-compose start
Starting db     ... done
Starting server ... done

DataLoaderのインストール

最初にデータローダーのインストールを行っていきます。

go get -u github.com/graph-gophers/dataloader

loaderの作成

公式サイトと参考に記載したサイトをもとに実装を行っていきます。

ここからは複数のページに分かれて記述します。また、記述内容が多いいのでコードの説明などは、ファイルを記述後にピックアップして説明をしていきます。

今回新たに追加するディレクトリ構成です。

gqlgen-todos
|
...
├── loader
│   ├── loaders.go
│   └── user.go
...

loader/loaders.goこのファイルはコンテキストに各データローダを注入する役割を持っています。

loader/loaders.go

package loader

import (
	"context"
	"github.com/graph-gophers/dataloader"
	"gorm.io/gorm"
	"net/http"
)

type ctxKey string

const (
	loadersKey = ctxKey("dataloaders")
)

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

// Loadersの初期化メソッド
func NewLoaders(db *gorm.DB) *Loaders {

	//ローダーの定義
	userLoader := &UserLoader{
		DB: db,
	}
	loaders := &Loaders{
		UserLoader: dataloader.NewBatchedLoader(userLoader.BatchGetUsers),
	}
	return loaders
}


// ミドルウェアはデータ ローダーをコンテキストに挿入します
func Middleware(loaders *Loaders, next http.Handler) http.Handler {
	loaders.UserLoader.ClearAll()
	// ローダーをリクエストコンテキストに挿入するミドルウェアを返す
	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)
	})
}

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

loader/user.goデータローダのユーザの取得処理を実装しています。

loader/user.go

package loader

import (
	"context"
	"fmt"
	"github.com/graph-gophers/dataloader"
	"github.com/shion0625/gqlgen-todos/graph/model"
	"gorm.io/gorm"
	"log"
)

// UserLoader はデータベースからユーザーを読み取ります
type UserLoader struct {
	DB *gorm.DB
}

// BatchGetUsers は、ID によって多くのユーザーを取得できるバッチ関数を実装します。
func (u *UserLoader) BatchGetUsers(ctx context.Context, keys dataloader.Keys) []*dataloader.Result {
	// 単一のクエリで要求されたすべてのユーザーを読み取ります
	userIDs := make([]string, len(keys))
	for ix, key := range keys {
		userIDs[ix] = key.String()
	}

	usersTemp := []*model.User{}

	if err := u.DB.Debug().Where("id IN ?", userIDs).Find(&usersTemp).Error; err != nil {
		err := fmt.Errorf("fail get users, %w", err)
		log.Printf("%v\n", err)
		return nil
	}

	usersByUserId := map[string]*model.User{}
	for _, user := range usersTemp {
		usersByUserId[user.ID] = user
	}

	users := make([]*model.User, len(userIDs))

	for i, id := range userIDs {
		users[i] = usersByUserId[id]
	}

	output := make([]*dataloader.Result, len(keys))
	for index := range userIDs {
		user := users[index]
		output[index] = &dataloader.Result{Data: user, Error: nil}
	}
	return output
}


// dataloader.Loadをwrapして型づけした実装
func LoadUser(ctx context.Context, userID string) (*model.User, error) {
	loaders := GetLoaders(ctx)
	thunk := loaders.UserLoader.Load(ctx, dataloader.StringKey(userID))
	result, err := thunk()
	if err != nil {
		return nil, err
	}
	return result.(*model.User), nil
}

ここまで実装できたらコードの説明を行っていきます。

NewLoaders関数の説明

NewLoaders関数

func NewLoaders(db *gorm.DB) *Loaders {

	//ローダーの定義
	userLoader := &UserLoader{
		DB: db,
	}
	loaders := &Loaders{
		UserLoader: dataloader.NewBatchedLoader(userLoader.BatchGetUsers),
	}
	return loaders
}

ここでは、dbの情報を取得してそれを使用してローダーの定義をしています。

Middlewareの説明

func Middleware(loaders *Loaders, next http.Handler) http.Handler { loaders.UserLoader.ClearAll() // ローダーをリクエストコンテキストに挿入するミドルウェアを返す 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) }) } コンテキストの挿入を行うミドルウェアです。 `/query`にリクエストが到達する前にローダーをコンテキストに注入しています。

GetLoaders関数の説明

func GetLoaders(ctx context.Context) *Loaders { return ctx.Value(loadersKey).(*Loaders) } 　コンテキストからローダーの取得を行っています。

BatchGetUsers関数の説明

BatchGetUsers関数

	userIDs := make([]string, len(keys))
	users :=  []*model.User{}
	for ix, key := range keys {
		userIDs[ix] = key.String()
	}

ここでは、userIdsという配列の作成を行い値はここでのkeysはN+1の際に何回にも分けて呼ばれていた値になります。私の場合はuser_idになります。そちらを先程定義した配列に入れます。

BatchGetUsers関数

	usersTemp := []*model.User{}
	if err := u.DB.Debug().Where("id IN ?", userIDs).Find(&usersTemp).Error; err != nil {
		err := fmt.Errorf("fail get users, %w", err)
		log.Printf("%v\n", err)
		return nil
	}

そして、その配列を用いて一度のクエリで値を取得する処理を記述しています。私の場合はuserテーブルからWhere("id IN ?", userIDs).Find(&usersTemp)こちらでUserモデルと一致するテーブルから先程の配列に入れたIDのデータの取得を行っています。

BatchGetUsers関数

	usersByUserId := map[string]*model.User{}
	for _, user := range usersTemp {
		usersByUserId[user.ID] = user
	}

	users := make([]*model.User, len(userIDs))

	for i, id := range userIDs {
		users[i] = usersByUserId[id]
	}

keysの順番にusersを変換しています。そうすること呼ばれたユーザIDの順番でデータを返すことができ対応するデータを返すことが出来ます。

並び替えないとユーザIDに対応していないデータを返すことになってしまいます。
この処理がないとU01のデータが欲しいのにU02のデータを返してしまうような事になります。

BatchGetUsers関数

	output := make([]*dataloader.Result, len(keys))
	for index := range userIDs {
		user := users[index]
		output[index] = &dataloader.Result{Data: user, Error: nil}
	}

ユーザIDの分だけ(r *todoResolver) Userが呼ばれているのでその回数分取得したデータの配列を作成しています。

戻り値の型が[]*dataloader.Resultなのでその型に合うようにデータを移し替えています。

LoadUser関数の説明

LoadUser関数

	loaders := GetLoaders(ctx)
	thunk := loaders.UserLoader.Load(ctx, dataloader.StringKey(userID))
	result, err := thunk()
	if err != nil {
		return nil, err
	}
	return result.(*model.User), nil

loders.go内で設定したGetLoaders関数を用いてコンテキストからloaderの取得を行っています。ここで取得するのは先程作成したuserLoaderになります。
そして、取得したuserLoaderを呼び出し、コンテキストとユーザのIDを引数として渡します。
そして、先程の戻り値を取得できたらその値を返して終了です。

loaderの使用

ここまででloaderが作成できました。
それでは、作成したloaderを使用してみます。

graph/schema.resolvers.go

package graph

import (
	"context"
	"crypto/rand"
	"fmt"
	"math/big"

	"github.com/shion0625/gqlgen-todos/graph/model"
+	"github.com/shion0625/gqlgen-todos/loader"
)

...

// User is the resolver for the user field.
func (r *todoResolver) User(ctx context.Context, obj *model.Todo) (*model.User, error) {
-	user := model.User{ID: obj.UserId}
-	r.DB.First(&user)
-	return &user, nil
+	user, err := loader.LoadUser(ctx, obj.UserId)
+   if err != nil {
+     return nil, err
+   }
+	return user, nil
}
...

作成したloadUserを呼び出します。この関数内でデータの取得処理を行っているのでエラーが起こるかの確認をします。

server.go

package main

import (
	"log"
	"net/http"
	"os"
	"fmt"
	"github.com/99designs/gqlgen/graphql/handler"
	"github.com/99designs/gqlgen/graphql/playground"
	"github.com/shion0625/gqlgen-todos/graph"
	"github.com/shion0625/gqlgen-todos/db"
	"github.com/joho/godotenv"
+	"github.com/shion0625/gqlgen-todos/loader"
)

const defaultPort = "8080"

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

	//データベースへの接続処理
	db := db.ConnectGORM() //追加
+    // loaderの初期化
+    ldrs := loader.NewLoaders(db)

	// resolver内でデータベースを扱えるように設定
	srv := handler.NewDefaultServer(graph.NewExecutableSchema(graph.Config{Resolvers: &graph.Resolver{
		DB:         db, // 追加
	}}))

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

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

最後にデータローダーの初期化の処理を追加し、/queryへリクエストが投げられたら先程作成したmiddlewareが呼ばれるといった処理を追加します。以上でDataLoaderを使用したN＋1問題の解決は完了です。

実装後のSQLログ

2023/01/04 07:04:23 /app/graph/schema.resolvers.go:45
[2.099ms] [rows:4] SELECT * FROM "todos"

2023/01/04 07:04:23 /app/loader/user.go:26
[3.545ms] [rows:3] SELECT * FROM "users" WHERE id IN ('U21','U12','U93')

キャッシュに関して

今回は先程のでメインの実装に関しての処理は終了です。

本当にお世話になっているこちらの記事の下部にてキャッシュに関する説明実装方法が記されています。ご参考にして実装をお願いします。

私は、1. そもそもCacheを使わないを使わないを選択しました。キャッシュ化することで同じkeyならSQLクエリが何度も飛ばないようになっているみたい!?

まとめ

DataLoaderを用いたN＋1問題の対処を行いました。
複雑なコードだったりで大変だとは思いますが概要であったりそういった点をお伝えでき実装していただければ幸いです。

次回は、todoの方のDataLoaderの処理を行っていきます。
今回は多対１のDataLoaderの実装を行って来ました。
todoに紐づいている一つのuserの取得でした。
userに紐づいている複数のtodoの値をDataLoaderで取得していきます。

参考