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 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.
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 DI benchmarking.