Using AsyncLocalStorage in Node.js: Real-World Use Cases

In modern web applications, many operations happen asynchronously. Whether it's fetching data from a remote server, performing a complex database query, or handling an incoming request, your code needs to work with asynchronous operations effectively. However, asynchronous code can be challenging to manage, especially when you need to share context across different parts of your application. AsyncLocalStorage is a powerful tool that helps you manage context across asynchronous operations. In this article, we will explore AsyncLocalStorage and learn how to use it in your Node.js applications.

Introduction

Asynchronous programming is a fundamental concept in modern web development. Asynchronous programming is all about performing tasks that don't block the main thread of execution. Asynchronous operations are crucial for building high-performance applications that can handle many concurrent requests. However, asynchronous programming can be tricky to manage, especially when you need to share context across different parts of your application. For example, you may need to pass user information, tracing information, or transaction objects across asynchronous operations.

AsyncLocalStorage is a new Node.js API that provides a way to store and retrieve data across asynchronous operations. AsyncLocalStorage allows you to create a context that is available to all the asynchronous operations executed within that context. You can use this context to store data that needs to be shared across different parts of your application.

Example 1: Tracing and Logging

One common use case for AsyncLocalStorage is tracing and logging. When an error occurs, it's essential to have all the relevant information available to debug the issue quickly. One way to achieve this is to pass a unique identifier (traceId) to all the logs generated during a request's lifecycle. With AsyncLocalStorage, you can store the traceId and retrieve it whenever you need it.

// logging.js
import { AsyncLocalStorage } from "async_hooks";

const asyncLocalStorage = new AsyncLocalStorage();

export function withTraceId(traceId, fn) {
  return asyncLocalStorage.run(traceId, fn);
}

export function log(message) {
  const traceId = asyncLocalStorage.getStore();
  console.log(`${message} traceId=${traceId}`);
}

// app.js
import express from "express";
import { withTraceId, log } from "./logging";

app.use((req, res, next) => {
  const traceId = req.header("x-trace-id") || Math.random().toString(36).substring(7);
  withTraceId(traceId, next);
});

app.use((req, res, next) => {
  log(`processing request ${req.url}`);
  next();
});

app.get("/post/:postId", async (req, res) => {
  const post = await fetchPost(req.params.postId);
  res.send(post);
});

async function fetchPost(postId: string) {
  log(`fetching data for post id: ${postId}`);
  // some logic to fetch post
  return `Hello world, ${postId}`;
}

This code defines an asynchronous middleware function authMiddleware that authenticates and stores the resulting user object using asyncLocalStorage. And getUserSecret uses getCurrentUser to retrieve the current user and then performs some asynchronous operation to fetch data for the user.

AsyncLocalStorage can be used to pass other tags or metadata through to logs as well, not just trace IDs. This can be helpful for tracking other information relevant to a particular request or context, such as user IDs, session IDs, or any other relevant data.

Example 2: Authentication

Another use case for AsyncLocalStorage is authentication. When a user logs in, you need to store the user's information somewhere so that you can access it across different parts of your application. With AsyncLocalStorage, you can store the user's information and retrieve it whenever you need it.

// auth.js
async function authenticate(authToken) {
  // fetch user data with authToken
  return {
    id: 1,
    username: "john",
  };
}

async function authMiddleware(req, res, next) {
  const authToken = req.header("x-auth-token");
  if (!authToken) {
    next();
  }

  authenticate(authToken!)
    .then((user) => {
      return asyncLocalStorage.run(user, next as any);
    })
    .catch(next);
}

function getCurrentUser() {
  return asyncLocalStorage.getStore() as any;
}

// secret.js
async function getUserSecret() {
  const user = getCurrentUser();
  // some async operation to fetch data for user.id
  const secret = `secret of ${user.username}`;
  return secret;
}

// app.js
app.use(authMiddleware);

app.get("/secret", async (req, res) => {
  const user = getCurrentUser();
  if (!user) {
    res.status(401).send("Unauthorized");
    return;
  }
  const secret = await getUserSecret();
  res.send(secret);
});

This approach eliminates the need to pass the user object explicitly through function parameters, making the code more concise and easier to read.

Example 3: Transactions

When you work with databases, you sometimes need to execute multiple queries within a single transaction. However, passing the transaction object explicitly can be non-trivial, especially when you have multiple queries executed in different code places. With AsyncLocalStorage, you can store the transaction object and retrieve it whenever you need it.

import { AsyncLocalStorage } from "async_hooks";
import knex, { Knex } from "knex";

// database.ts
const db = knex({
  // database configuration
});

const asyncLocalStorage = new AsyncLocalStorage();

export async function withTransaction(fn) {
  const trx = db.transaction();
  return asyncLocalStorage.run(trx, fn);
}

export function getCurrentTransaction() {
  const trx = asyncLocalStorage.getStore();
  return (trx as Knex.Transaction) || db;
}

// task.ts
export function getTask() {
  const trx = getCurrentTransaction();
  return trx("tasks").forUpdate().skipLocked().first();
}

export function completeTask(id) {
  const trx = getCurrentTransaction();
  return trx("tasks").update({ done: true }).where("id", id);
}

// worker.ts
async function work() {
  return withTransaction(async () => {
    const task = await getTask();

    // do some work

    await completeTask(task.id);
  });
}

The withTransaction function creates a transaction using the knex library and runs the provided function with the transaction context.

The getTask and completeTask functions retrieve and update a task from the database using the transaction from getCurrentTransaction.

Finally, the work function uses withTransaction to execute the getTask and completeTask functions within a transaction.

Conclusion

AsyncLocalStorage is a powerful tool that allows you to manage context across asynchronous operations. You can use AsyncLocalStorage to store and retrieve data that needs to be shared across different parts of your application. However, AsyncLocalStorage should be used with caution, and you should be aware of its limitations. In some cases, it's better to pass data explicitly, especially in simple applications where passing the context as a parameter is easy.

In conclusion, AsyncLocalStorage is a great tool that you should consider when you need to share context across asynchronous operations. Whether you're working with tracing information, user data, or transaction objects, AsyncLocalStorage can help you manage your context effectively.