Asynchronous context tracking in NodeJS, introduced with version 16, addresses a common challenge in node applications: maintaining context across asynchronous operations. In such environment, where non-blocking I/O operations are the norm, it can be difficult to preserve a "context" or "state" across callbacks, promises, or async/await operations. This is crucial for tasks like tracking user sessions, handling transactions, or implementing logging that depends on knowing the sequence of operations that led to a particular state.

It provides a way for developers to preserve this context without resorting to complex workarounds. It uses the AsyncLocalStorage API, part of the async_hooks module, allowing developers to store and access data that is specific to a particular sequence of asynchronous operations. This makes it possible to easily pass along context through the many layers of asynchronous calls, improving the ability to monitor, debug, and write more maintainable applications. Essentially, it gives developers the power to keep track of the execution flow, even in the inherently asynchronous environment of NodeJS, making it easier to manage application state across asynchronous boundaries.

AsyncLocalStorage within AWS Lambda Environments

In the dynamic world of AWS Lambda, where functions respond to events in isolated invocations, managing context across asynchronous operations can be a bit of a juggling act. AsyncLocalStorage is designed to elegantly handle this exact scenario. It offers a seamless way to maintain context without the cumbersome need to pass state around through function parameters. Each Lambda invocation is treated as a standalone execution, ensuring that context does not leak between invocations, even in warm containers that are reused for efficiency. This setup guarantees that every function run has its clean slate regarding asynchronous context, making code more readable, maintainable, and significantly reducing the likelihood of bugs related to improper context management.

Let's consider the following example, which tracks the user claims by introducing a middy middleware

import { AsyncLocalStorage } from 'node:async_hooks'
import middy from '@middy/core'
import zod from 'zod'

const Claims = zod.object({ ... })
type Claims = zod.infer<typeof Claims>

const claimsStorage = new AsyncLocalStorage<Claims>();

const useClaims = () => {
    const store = claimsStorage.getStore()
    if (!store) {
        throw new Error('invalid claims')
    }
    return store
}

const withUserStoredInContext = (): middy.MiddlewareObj<APIGatewayProxyEventV2WithJWTAuthorizer, APIGatewayProxyStructuredResultV2> => {
  return {
    before: (handler: middy.Request<APIGatewayProxyEventV2WithJWTAuthorizer>) => {
      const claims = Claims.parse(handler.event.requestContext.authorizer.jwt.claims)
      claimsStorage.enterWith(claims)
    },
  }
}

we implement a function useClaims, and can access claims throughout the whole request lifecycle now.

import { useClaims } from '...'
import * as DB from './db'
import { isEmpty } from 'remeda'

const list = async () => {
    const { sub } = useClaims()

    if(!sub) return []

    const allPromotions = await DB.listBySub(sub)

    if(!allPromotions || isEmpty(allPromotions)) return []

    return transformPromotions(allPromotions)
}

Pitfalls

Using AsyncLocalStorage with AWS Lambda offers many benefits for context management across asynchronous operations. However, there are some considerations and potential pitfalls to be aware of:

1. enterWith is still experimental. That only stable function to populate the store is to use the run method. See this stackoverflow discussion.

2. Understanding Context Propagation: Developers must have a clear understanding of how context is propagated across asynchronous calls to effectively use AsyncLocalStorage. Misunderstandings can lead to context loss or incorrect assumptions about the availability of context, resulting in bugs that are difficult to diagnose.

   💡

   personally I already faced this quite often e.g. calling the getStore and not knowing why it returns undefined even though I thought I populated the store

3. Potential for Context Leaks: Incorrect usage of AsyncLocalStorage, especially not properly entering and exiting the context, can lead to context information leaking across Lambda invocations in warm containers. Although AWS Lambda provides isolation between invocations, improper management of context can introduce subtle bugs related to context contamination.

4. Cold Start & Memory Consumption: I did not verify this yet, but there are quite some old discussions about performance degradations when using these hooks. I can imagine the use of async hooks do have some impact on the performance of an application, be it the time to start a container or the memory consumption.

How do you use the AsyncLocalStorage? Let's discuss!

Else, give it a try, cheers!