💽

NestJSでDataLoaderをどこで初期化するか問題

2024/06/11に公開

ユビーではGraphQLのバックエンドサービスをNestJSでモジュラモノリスな構成で作っていますが、その中で得られたDataLoaderの知見について紹介します。

背景

GraphQLでN+1を回避するにはDataLoaderを使うのが一般的なプラクティスですが、DataLoaderはインタフェースの都合上リクエスト毎に初期化をおこなう必要があります。この初期化をどこで行うかが悩みどころで、いくつかパターンがあるので今回はそれについて紹介します。

DataLoaderの基本

このようなスキーマとResolverがあったとします。

type Query {
  posts: [Post!]!
}

type Post {
  id: ID!
  title: String!
  author: User!
}

type User {
  id: ID!
  name: String!
}
@Resolver("Post")
class PostResolver {
  constructor(private readonly db: DbService) {}

  @Query()
  posts(): Promise<Post[]> {
    return this.db.findPosts();
  }

  @ResolveField()
  author(@Parent() post: Post): Promise<User> {
        // ここがpostsの要素の数だけ呼ばれる
    return this.db.findUser(post.userId);
  }
}

これに対して以下のようなクエリを発行すると db.findUser が posts の配列の要素の数だけ呼ばれてパフォーマンスが劣化するというのが一般的なN+1の問題です。

query {
  posts {
    id
    title
    author {
      id
      name
    }
  }
}

この問題を解決するためにDataLoaderを使ってまとめてuserを取ってこれるようにします。

一回だけ初期化する

まずは何も考えずに以下のように書きます。

@Injectable()
class UserDataLoader extends DataLoader<string, User> {
  constructor(private readonly db: DbService) {
    // 最後にまとめてここが呼ばれるのでuserの取得が一回にまとまる
    super((keys: readonly string[]) => {
      return this.db.findUserByIds(keys);
    });
  }
}

@Resolver("Post")
class PostResolver {
  constructor(
    private readonly db: DbService,
    private readonly userDataLoader: UserDataLoader
  ) {}

  @Query()
  async posts(): Promise<Post[]> {
    return this.db.findPosts();
  }

  @ResolveField()
  author(@Parent() post: Post): Promise<User> {
    // findUserを呼び出していたところをdataloaderにidを渡すだけにした
    return this.userDataLoader.load(post.userId);
  }
}

この例は一見普通に動きますが全然ダメな実装で、DataLoaderはキーを元にデータをキャッシュするので、リクエストをまたいでキャッシュが混ざって意図せぬデータを返してしまったり、データを更新してもキャッシュが更新されずに古いデータを返し続けるという羽目になります。

NestJSでは上記のように書くと UserDataLoader は起動時の最初の一回しか初期化されずリクエストをまたいでインスタンスが共有されるのが原因です。DataLoaderのドキュメントにもリクエストごとにDataLoaderを初期化するべきと書かれています。

In many applications, a web server using DataLoader serves requests to many different users with different access permissions. It may be dangerous to use one cache across many users, and is encouraged to create a new DataLoader per request.

https://github.com/graphql/dataloader?tab=readme-ov-file#creating-a-new-dataloader-per-request

NestJSのRequest Scopeを使う

これを解決するためにNestJSのRequest Scopeを使います。以下のようにscopeを指定するとリクエスト毎に初期化をおこなってくれるという動作になります。

@Injectable({ scope: Scope.REQUEST }) // ここのscopeを指定した
class UserDataLoader extends DataLoader<string, User> {
  constructor(private readonly db: DbService) {
    super((keys: readonly string[]) => {
      return this.db.findUserByIds(keys);
    });
  }
}

これはNestJSでDataLoaderを利用する例ではよく見る方法ですが、これも推奨はしません。先程の方法に比べると意図せぬキャッシュによって壊れるということはなくなるので圧倒的にマシなのですが、パフォーマンスの面で問題になります。NestJSの公式にも以下のような注意書きがあります。

Using request-scoped providers will have an impact on application performance. While Nest tries to cache as much metadata as possible, it will still have to create an instance of your class on each request. Hence, it will slow down your average response time and overall benchmarking result. Unless a provider must be request-scoped, it is strongly recommended that you use the default singleton scope.

https://docs.nestjs.com/fundamentals/injection-scopes

DataLoaderのclassだけがRequest Scopeで初期化されるならいいのですが、依存するclassも全てリクエストごとに初期化されるようになるというのが問題です。特にDataLoaderは多くの場合DBのサービスに依存しているので、実装によってはリクエストごとにコネクションが作られることになるかもしれません。そこまで考えて依存するサービスを作るか、パフォーマンスのことを考慮しなくてもいいケースなどではRequest Scopeでもよいと思いますが、そうでない場合は別の方法を検討すべきです。

contextにインスタンスをもたせる

ここからはNestJSはあまり関係なくGraphQLのサーバー実装であれば同じように使える話です。Node.jsのGraphQLのサーバー実装にはリクエストごとに context というオブジェクトが作成され、Resolver 関数で参照することができます。リクエストごとに初期化されるのでここでDataLoaderのインスタンスを作ってResolverに流します。

コードはざっくりこんな感じです。

type UserDataLoader = DataLoader<string, User>;
type AppContext = {
  userDataLoader: UserDataLoader;
};

@Injectable()
class UserDataLoaderService {
  // 依存するDbServiceは一回だけしか初期化されない
  constructor(private readonly db: DbService) {}

  // ここはリクエストごとに呼ばれる
  getLoader(): UserDataLoader {
    return new DataLoader((ids) => this.db.findUserByIds(ids));
  }
}

@Resolver("Post")
class PostResolver {
  // ...

  @ResolveField()
  author(@Parent() post: Post, @Context() context: AppContext): Promise<User> {
    // contextを通してDataLoaderのインスタンスを使う
    return context.userDataLoader.load(post.userId);
  }
}

@Module({
  imports: [
    GraphQLModule.forRootAsync<ApolloDriverConfig>({
      driver: ApolloDriver,
      imports: [ResolverModule],
      inject: [UserDataLoaderService],
      useFactory: (userDataLoaderService: UserDataLoaderService) => {
        return {
          typePaths: ["./schema.graphqls"],
          // この関数はリクエスト毎に呼ばれてcontextを作成する
          context: (): AppContext => ({
            userDataLoader: userDataLoaderService.getLoader(),
          }),
        };
      },
    }),
  ],
})
class AppModule {}

contextの初期化は上記のようにルートモジュールの設定部分に記述します。このとき、依存するclassをinjectできますがinjectするclassの初期化は初回に一回だけおこなわれるのでDataLoaderのclassが依存しているDbServiceは一回だけしか初期化されません。一方でDataLoaderServiceに実装した getLoader() はリクエスト毎に呼ばれるのでリクエスト毎にDataLoaderのインスタンスを作ることができます。

このように、contextの初期時にDataLoaderのインスタンスを作りcontextに紐づける方法はApollo Serverのドキュメントにも記述があり、GraphQLのサーバー実装では一般的な方法と言えます。

https://www.apollographql.com/docs/apollo-server/data/fetching-data/#batching-and-caching

基本的にはこの方法がシンプルかつパフォーマンス的にもよいと思うのですが、ひとつ微妙な点があって、contextというひとつのオブジェクトにアプリケーション横断で様々な機能が詰め込まれる可能性があるという点です。特にモジュラモノリスのような関心事をモジュールごとに分離したい設計においては中央集権的なオブジェクトにモジュール横断の関心事を詰め込みたくありません。そこで考えたのが次に紹介する方法です。

WeakMapにDataLoaderのインスタンスを保存する

アイデアとしては、contextオブジェクトをkeyにしたWeakMapにDataloaderのインスタンスを保存するというものです。WeakMapは弱参照のMapで、keyがGCで消えるとWeakMapにsetしたデータも自動で消えるという性質を持っています。

https://developer.mozilla.org/ja/docs/Web/JavaScript/Reference/Global_Objects/WeakMap

contextオブジェクトはリクエスト単位で作成され、リクエスト完了とともに消えることが期待できるので、contextをWeakMapのkeyに使うことでリクエスト単位のライフサイクルを持ったデータストアを得ることができます。コードとしては以下のようになります。

type UserDataLoader = DataLoader<string, User>;
type AppContext = {};

@Injectable()
class UserDataLoaderService {
  private readonly map = new WeakMap<AppContext, UserDataLoader>();

    // 一回だけ初期化される
  constructor(private readonly db: DbService) {}

  // DataLoaderが必要なタイミングで呼ばれる
  getLoader(context: AppContext): UserDataLoader {
    // contextオブジェクトをkeyにして、すでにインスタンスが作らていれば返す
    if (this.map.has(context)) {
      return this.map.get(context)!;
    }
    // 作られてなければ初期化してcontextをkeyにして保存する
    const loader = new DataLoader<string, User>((ids) =>
      this.db.findUserByIds(ids)
    );
    this.map.set(context, loader);
    return loader;
  }
}

@Resolver("Post")
class PostResolver {
  constructor(private readonly userDataLoader: UserDataLoaderService) {}
  
  // ...

  @ResolveField()
  author(@Parent() post: Post, @Context() context: AppContext): Promise<User> {
    // contextを引数にしてDataLoaderのインスタンスを得る
    return this.userDataLoader.getLoader(context).load(post.userId);
  }
}

これでcontextにDataLoaderのインスタンスを詰めることなく、リクエストごとに初期化されたDataLoaderを得ることができます。DataLoaderがcontextの型を知る必要があるという点だけやや微妙ですが、contextオブジェクトが太るよりはマシかなと個人的には思います。

まとめ

NestJSでDataLoaderを利用するいくつかの方法について説明しました。基本的にはリクエストごとに初期化するということだけを守れば実装は好みや状況次第ですが、ユビーのバックエンドサーバーではモジュラモノリスを採用しており、可能な限りモジュールの中に関心を閉じ込めたいのでWeakMapを用いた方法を採用しています。

Ubie テックブログ

Discussion