Using DataLoader with AsyncLocalStorage over Fastify/GraphQL Yoga | Pure DI in TypeScript Node.js
Published: 2023-02-10 14:24:26

---
title: Using DataLoader with AsyncLocalStorage over Fastify/GraphQL Yoga | Pure DI in TypeScript Node.js
published: true
date: 2023-02-10 02:07:29 UTC
tags: 
canonical_url: https://blog.vady.dev/using-dataloader-with-asynclocalstorage-over-fastifygraphql-yoga
---

In the previous tutorial, I've implemented GraphQL Server using GraphQL Yoga.
Today, I'm going to extend our service by using [DataLoader](https://github.com/graphql/dataloader).

> I will use [the repo](https://github.com/vad3x/typesafe-dependency-injection-in-typescript-samples/tree/main/samples/part4-graphql) as starting point.

## DataLoader

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

```gql
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](https://github.com/graphql/dataloader) can help.

Let's add the `DataLoader` to our `StubOrderStore`.

```ts
// 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 `Order`s 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 `Order`s 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 `DataLoader`s in `AsyncLocalStorage` as we did for `RequestId` [before](https://dev.to/vad3x/step-by-step-guide-of-implementing-scoped-like-dependencies-using-asynclocalstorage-within-fastify-3cg9).

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

```ts
// 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:

```ts
// 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:

```ts
// 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.

```ts
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`.

```ts
// 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:

```ts
// 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](https://github.com/vad3x/typesafe-dependency-injection-in-typescript-samples/blob/main/samples/part5-dataloader)

## 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.

<!-- In a subsequent article, I will cover another important topic: `Configuration`. -->