Sharing Resources in an Asynchronous Environment with AsyncLocalStorage

“In the tapestry of life, context weaves the threads that give meaning and depth to our experiences.” — ChatGPT

TL;DR :

- Having a shared resource in an asynchronous function can lead to inconsistency
- You can avoid those problems by using Async Local Storage (available in Node13+, stable in Node16+)

A bit of … context

Back in January 2023, I was working on an internal context logger for OVRSEA. It was inspired by the Mapped Diagnostic Context concept found in popular loggers like Log4j, where you can set a variable in the context at any point to have it in subsequent logs.

import {logger} from "@ovrsea/ovrlogger"
const sendMessage = async (userId: string) => {
   logger.setContext("user", userId)
   logger.info("Trying to send a message") 
   // This will log : 
   // INFO: Trying to send a message, userId: "achraf"
   
   // * some computations
   logger.info("Sent the message successfully") 
   // This will log :
   // INFO: Trying to send a message, userId: "achraf"
}

If we had deeply nested calls to the logger, it will still have the userId logged !
And it was working quite well for a lot of case….until it didn’t.

The fall

One time, while investigating the logs of a document service, I noticed that the document IDs shown did not correspond to the URL of the document.
But, there was no sign of a bug in the document service when examining the data, and no bugs were reported.

INFO: Uploading document, documentId: 222
INFO: Uploaded document to /111.pdf, documentId: 222
INFO: Uploading document, documentId: 222
INFO: Uploading document to /222, documentId: 333 // WTF ?
INFO: Uploading document, documentId: 333
INFO: Uploading document to /333, documentId: 111

The code associated with the logs was:

export const uploadDocuments = (documentIds: string[]) => {
   await Promise.all(documentsIds.map((documentId) => {
     logger.setContext("documentId", documentId)
     logger.info("Uploading document")
     const url = await determineUrl(documentId)
     // more work
     logger.info(`Uploaded document to ${url}`)
   }))
}

So even If I didn’t want to believe it, the logger was sending untrustworthy logs. That’s worst than no logs !

So I made a Minimal Reproducible Example to try to understand what was happening.

Minimal Reproducible Example

I didn’t want to make actual HTTP requests or Database queries, so I simulated the asynchronous behavior with a handy function calledrandomDelay which adds a random delay

const randomDelay = (delayInSeconds: number) => {
   return new Promise((resolve) => {
     setTimeout(resolve, delayInSeconds * Math.random() * 1000);
   });
};

I also needed a minimal, naive logger with context capabilities, as I knew this was the faulty part.

const buildLogger = () => {
   const context: Record<string, any> = {};
   const logger = {
     log: (…args: any[]) => console.log(…args, context),
     setContext: (key: string, value: any) => {
       context[key] = value;
     },
   };
   return logger
  }
const logger = buildLogger();

It’s a builder that creates a closure to utilize the context in logs, allowing for multiple loggers with uncorrelated contexts.
Additionally, I replicated the actual asynchronous code with a fake mutation for testing purposes.

const mutation = async ({ documentId }: { documentId: number }) => {
   logger.setContext("documentId", documentId);
   logger.log(`Before async, true documentId = ${documentId}`)
   await randomDelay(2);
   logger.log(`After async, true documentId = ${documentId}`)
};

const uploadMultipleDocuments = async () => {
   const documentsIds = [111, 222, 333]
   await Promise.all(
     documentsIds.map(async (documentId) => {
       await mutation({ documentId });
     })
   );
};

Putting all of this together, we have a similar function to the actual code that sent inconsistent logs.

const randomDelay = (delayInSeconds: number) => {
  return new Promise((resolve) => {
    setTimeout(resolve, delayInSeconds * Math.random() * 1000);
  });
};

const buildLogger = () => {
 const context: Record<string, any> = {};
 const logger = {
   log: (...args: any[]) => console.log(...args, context),
   setContext: (key: string, value: any) => {
     context[key] = value;
   },
 };
 return logger
}

const logger = buildLogger();

const mutation = async ({ documentId }: { documentId: number }) => {
  logger.setContext("documentId", documentId);
  logger.log(`Before async, true documentId = ${documentId}`)
  await randomDelay(2);
  logger.log(`After async, true documentId = ${documentId}`)
};

const uploadMultipleDocuments = async () => {
  const documentsIds = [111, 222, 333]
  await Promise.all(
    documentsIds.map(async (documentId) => {
      await mutation({ documentId });
    })
  );
};

Okay, we’re all set! Let’s call this function and hopefully, we can figure out what’s going on. Hopefully, it’s not a real bug in the document service.

Results of the MRE

Hmmm… well, it’s definitely the logger.
Before the asynchronous call, it seems correct, but once we start the asynchronous call, we never update the context, so the context will just stay at the last set value.
That’s no good.

A simplified illustration on what’s exactly going on

The issue with this was that I realized any context set at the HTTP Request level (such as adding the targeted resource by the request to the context) was also at risk of being corrupted by parallel requests.

Therefore, we couldn’t set a context at the top level and expect it to be present in every log within nested function calls.

The firefighting

I tried to solve the issue in a naive way, but unfortunately, I couldn’t find a solution that didn’t negatively impact the developer experience, which is not ideal.

Since the logger wasn’t used everywhere, there was still room for improvement and fixing it. The solution I opted for was advocating for building a different logger in each asynchronous context.

While it did the job, I didn’t want to invest more time into it as the workaround wasn’t too burdensome.

However, this approach greatly limits the potential of the logger, as an ideal logger would have context isolation.

Context Isolation — Ideal Logger

The Rise

While working on another issue, I realized that I was essentially facing the same problem but in a different form.

I needed a value to be associated with the user request, and simply passing down the value from the top level to every part of the code where it was needed was not a feasible option.

So, I decided to revisit the logger and search online for potential solutions to this problem. After all, I couldn’t be the only one facing this issue!

As it turned out, I wasn’t alone.

Some incredibly intelligent individuals had implemented async_hooks to track the lifecycle of various asynchronous resources. While exploring the Node.js documentation for async_hooks, I stumbled upon something that appeared to be the solution to my problem: Async Local Storage.

It’s a class that provides helpers to insure that our context is there when we need it and to insure it’s correct !

• Set something as a local context : asyncLocalStorage.enterWith(context)
• Retrieve the current context : asyncLocalStorage.getStore()
• Run an asynchronous function within a specific context: asyncLocalStorage.run(context, fn)

Using run is crucial. Attempting to call the other functions outside this scope will result in undefined stores

Now that we have all the necessary tools, we can modify the logger!

To transform the Minimal Reproducible Example into a Proof of Concept, we first need to do is initialize the Async Local Storage.

import { AsyncLocalStorage } from "async_hooks";
type Context = Record<string, any>;
const asyncLocalStorage = new AsyncLocalStorage<Context>();

Then we need to change our logger to use the store

const buildLogger = () => {
 // we no longer need a closure actually
  const logger = {
    log: (...args: any[]) => {
      console.log(...args, asyncLocalStorage.getStore());
    },
    setContext: (key: string, value: any) => {
      const oldContext = asyncLocalStorage.getStore() ?? {};
      oldContext[key] = value;
      asyncLocalStorage.enterWith(oldContext);
    },
  };

  return logger;
};

We also need to initialize the context at the top level. To accomplish this, I created a small wrapper helper function that allows us to wrap our mutation in the asyncLocalStorage.run() scope. This ensures that the context is properly set and available throughout the execution.

const withAsyncContext =
  <T, F extends (...args: any[]) => Promise<T>>(fn: F) =>
  (...args: Parameters<F>): Promise<T> => {
    return asyncLocalStorage.run({}, async () => {
      const result = await fn(...args);

      return result;
    });
  };

It’s simply a curried function that takes another function as an argument, runs it inside the scope of asyncLocalStorage.run(), and returns the result.
Let's put all of this together:

import { AsyncLocalStorage } from "async_hooks";

type Context = Record<string, any>;
const asyncLocalStorage = new AsyncLocalStorage<Context>();

const randomDelay = (delayInSeconds: number) => {
  return new Promise((resolve) => {
    setTimeout(resolve, delayInSeconds * Math.random() * 1000);
  });
};

const buildLogger = () => {
  const logger = {
    log: (...args: any[]) => {
      console.log(...args, asyncLocalStorage.getStore());
    },
    setContext: (key: string, value: any) => {
      const oldContext = asyncLocalStorage.getStore() ?? {};

      oldContext[key] = value;
      asyncLocalStorage.enterWith(oldContext);
    },
  };

  return logger;
};

const logger = buildLogger();

const withAsyncContext =
  <T, F extends (...args: any[]) => Promise<T>>(fn: F) =>
  (...args: Parameters<F>): Promise<T> => {
    return asyncLocalStorage.run({}, async () => {
      const result = await fn(...args);

      return result;
    });
  };

const mutation = async ({ documentId }: { documentId: number }) => {
  logger.setContext("documentId", documentId);

  logger.log(`Before async, true documentId = ${documentId}`);
  await randomDelay(2);
  logger.log(`After async, true documentId = ${documentId}`);
};

const uploadMultipleDocuments = async () => {
  const documentsIds = [111, 222, 333];

  await Promise.all(
    documentsIds.map(async (documentId) => {
      await withAsyncContext(mutation)({ documentId });
    })
  );
};

This example is primarily intended for demonstration purposes. In practical scenarios, the withAsyncContext function would likely serve as middleware within an HTTP server or as a wrapper for the top-level handler of a serverless function.

Okay, it’s now time test it !

Reliable logs due to AsyncLocalStorage

That’s exactly what we’d expect ! That’s perfect for us !

Conclusion

Logging can be a challenging task since it primarily serves as a message to our future selves when we encounter bugs during debugging. To be considerate to our future selves, it is crucial to provide sufficient context in our logs to aid in bug resolution. However, achieving this can be tricky in asynchronous functions. Thankfully, by leveraging AsyncLocalStorage, we can bind the context to our scope and generate reliable logs.

Thank you for joining me on this journey through the peculiar yet fascinating aspects of Node.js. I hope the insights shared here prove valuable for your projects.
If you do something cool using AsyncLocalStorage, I’d be curious to hear about it.
Feel free to reach out and share your experiences.