Using DataLoader with AsyncLocalStorage over Fastify/GraphQL Yoga | Pure DI in TypeScript Node.js

In the previous tutorial, I've implemented GraphQL Server using GraphQL Yoga.
Today, I'm going to extend our service by using DataLoader.

I will use the repo as starting point.

DataLoader

GraphQL gives us the ability to combine multiple queries into one request like this:

query {
  order0: order(id: "order-0") {
    id
  }

  order1: order(id: "order-1") {
    id
  }

  order2: order(id: "order-1") {
    id
  }
}

If we invoke the query, we will see that our StubOrderStore.find was hit 3 times.
It's not an issue with our in-memory implementation but might become a problem when the OrderStore implementation is I/O related.

What if we could combine all the underlying requests into one batch request, this would definitely improve the performance.

This is one of the use cases where DataLoader can help.

Let's add the DataLoader to our StubOrderStore.

// stub-order-store.ts

import DataLoader from "dataloader";
...

export class StubOrderStore implements OrderStore {
  private readonly loader: DataLoader<ID, Order | undefined>;

  constructor(private readonly logger: Logger) {
    this.loader = new DataLoader(this.batchLoad);
  }

  find(id: ID): Promise<Order | undefined> {
    return this.loader.load(id);
  }

  private batchLoad: DataLoader.BatchLoadFn<ID, Order | undefined> = (ids) => {
    this.logger.info(`calling StubOrderStore.batchLoad with orderIds: ${ids}`);

    return Promise.resolve(
      ids.map((id) => ({
        id,
      }))
    );
  };
}

Now, when we call our API we have only one message logged:

calling StubOrderStore.batchLoad with orderIds: order-0,order-1

So, DataLoader combines all different find calls into a single call to batchLoad. Moreover, it has deduplicated order-1 calls, which is cool.

Scoped DataLoader State

There is a problem we start facing if we call our API with the same orderIds one more time: we don't get any StubOrderStore.batchLoad messages anymore, meaning the Orders are stale now. That happens because DataLoader caches the batchLoad results per key per DataLoader instance. This is a very useful feature for cases when we want to cache data in the scope of the HTTP request. Unfortunately, with the current implementation, our Orders will never be reloaded from the persistence, since we have one instance of DataLoader for the whole application lifetime.

To fix the issue, we can store DataLoaders in AsyncLocalStorage as we did for RequestId before.

Let's add a new dependency to our StubOrderStore class: GetDataLoader fn:

// stub-order-store.ts
import { BatchLoadFn, GetDataLoader } from "./dataloader";
...

export class StubOrderStore implements OrderStore {
  constructor(
    private readonly logger: Logger,
    private readonly getDataLoader: GetDataLoader<ID, Order | undefined>
  ) {}

  find(id: ID): Promise<Order | undefined> {
    const loader = this.getDataLoader(this.batchLoad);

    return loader.load(id);
  }

  private batchLoad: BatchLoadFn<ID, Order | undefined> = (ids) => {
    this.logger.info(`calling StubOrderStore.batchLoad with orderIds: ${ids}`);

    return Promise.resolve(
      ids.map((id) => ({
        id,
      }))
    );
  };
}

In the code above, we invoke getDataLoader for every find, this gives us the flexibility to work with scoped DataLoader.

Now we also need to update our registry code:

// registry/stub-order-store.ts
export function orderStore() {
  return (deps: {
    createLogger: CreateLogger;
    getDataLoader: GetDataLoader;
  }) => {
    const orderStore = new StubOrderStore(
      deps.createLogger(StubOrderStore.name),
      deps.getDataLoader
    );

    return {
      orderStore,
    };
  };
}

Let's define the GetDataLoader interface:

// dataloader.ts
import DataLoader, { BatchLoadFn } from "dataloader";

export { DataLoader, BatchLoadFn };

export type GetDataLoader<K = any, V = any> = (
  loadFn: BatchLoadFn<K, V>
) => DataLoader<K, V>;

export type RunWithDataLoaders = <R>(callback: (...args: any[]) => R) => R;

Now we need to add the implementations for the interfaces.

Calling getDataLoader should either create a new DataLoader for a given loadFn, or return the existing one if exists.

type DataLoadersMap = Map<
  BatchLoadFn<unknown, unknown>,
  DataLoader<unknown, unknown>
>;

export class DataLoaderStore {
  private readonly dataLoadersMapAls = new AsyncLocalStorage<DataLoadersMap>();

  getDataLoader: GetDataLoader = (loadFn) => {
    const dataLoadersMap = this.dataLoadersMapAls.getStore();
    if (!dataLoadersMap) {
      // throwing an exception here to simplify the cient's code
      throw new Error(
        "'getDataLoader' should not be called outside of 'runWithDataLoaders' callback"
      );
    }

    return getOrCreateDataLoader(loadFn, dataLoadersMap);
  };

  runWithDataLoaders: RunWithDataLoaders = (callback) => {
    return this.dataLoadersMapAls.run(new Map(), callback);
  };
}

function getOrCreateDataLoader(
  loadFn: BatchLoadFn<unknown, unknown>,
  dataLoadersMap: DataLoadersMap
) {
  const existingDataLoader = dataLoadersMap.get(loadFn);
  if (existingDataLoader) {
    return existingDataLoader;
  }

  const newDataLoader = new DataLoader(loadFn);
  dataLoadersMap.set(loadFn, newDataLoader);

  return newDataLoader;
}

// registry/dataloader-store.ts
export function dataLoaderStore() {
  return () => {
    const { getDataLoader, runWithDataLoaders } = new DataLoaderStore();

    return {
      getDataLoader,
      runWithDataLoaders,
    };
  };
}

// create-app-registry.ts
export function createAppRegistry() {
  return new RegistryComposer()
    ...
    .add(dataLoaderStore())
    ...
    .compose();
}

The next step is to run runWithDataLoaders on every Fastify request, similar to what we did for RequestId.

// fastify-plugins/run-with-dataloaders-plugin.ts
export function runWithDataLoadersPlugin(deps: {
  runWithDataLoaders: RunWithDataLoaders;
}): FastifyPluginCallback {
  const plugin: FastifyPluginCallback = (fastify, _, next) => {
    fastify.addHook("onRequest", (_request, _reply, callback) => {
      deps.runWithDataLoaders(callback);
    });

    next();
  };

  return fp(plugin);
}

Registering runWithDataLoadersPlugin on our Fastify server:

// registry/fastify-server.ts
export function fastifyServer() {
  return (
    deps: Parameters<typeof runWithRequestIdPlugin>[0] &
      Parameters<typeof runWithDataLoadersPlugin>[0] &
      ...
  ) => {
    const server = Fastify({});

    server.register(runWithRequestIdPlugin(deps));
    server.register(runWithDataLoadersPlugin(deps));
    ...

    return {
      fastifyServer: server,
    };
  };
}

Finally, we can run our server again and repeat the same request.
We can now see the message for every HTTP request:

calling StubOrderStore.batchLoad with orderIds: order-0,order-1

...

calling StubOrderStore.batchLoad with orderIds: order-0,order-1

Full example source code can be found here

Conclusion

In this tutorial, we covered the usage of the DataLoader library using a global approach and a request-based with AsyncLocalStorage.
In a subsequent article, I will switch the gears and focus on the DI frameworks benchmarking.