Feature request: Context management thru AsyncLocalStorage?

[everett1992]

Use case

Currently if you want to context in async workflows you have to pass child loggers to the methods you call

const tasks = batch.map(async (id) => {
  doWork(id, { logger: logger.createChild() })
})
await Promise.all(tasks)

Solution/User Experience

We could use node's AsyncLocalStorage to save the 'current' logger, and functions could fetch the logger from the store

This could be combined with explicit resource management to allow writing code like

const handler = (event, context) => {
  using logger = withContext(context)
  const tasks = event.batch.map(doWork)
 await Promise.all(tasks) 
}

async function doWork(task: Task) {
  using logger = withAdditionalKeys({ id: task.id })
  await step1(task)
  await step2(task)
}

async function step1(task: Task) {
  using logger = withAdditionalKeys({ step: "1" )
  logger.info()
}

Alternative solutions

This can be built ad-hoc but would be better to include in the library.

Acknowledgment

• [x] This feature request meets Powertools for AWS Lambda (TypeScript) Tenets
• [x] Should this be considered in other Powertools for AWS Lambda languages? i.e. Python, Java, and .NET

Future readers

Please react with 👍 and your use case to help us understand customer demand.

[dreamorosi]

Hi @everett1992 thank you for opening this issue.

I've been looking at explicit resource management for a while and did some prototyping around Logger, Tracer, and Metrics primarily around "clean up" mechanisms, for example doing using logger ... within the scope of a function would clear any keys that were added, for metrics it would flush the buffer, and so on.

I didn't however think of combining this with AsyncLocalStorage to avoid props drilling - it's very neat.

When using the AsyncLocalStorage, I'd definitely want to use the enterWith pattern, so we don't have to force customers to wrap their functions in a callback.

My only concern with introducing this module is around testing. I want to make sure that we provide a seamless way for customers to test in isolation code that uses for example withAdditonalKeys without having to do any complex setup.

To be honest I'm fairly new to AsyncLocalStorage, I was wondering if you'd be open to sharing/collaborating on a PoC for the feature to see how it would work and whether there are limitations.

It's definitely something I'd like to add to the library, but having some help in these initial phases - if not during the implementation - would be a great help.

[everett1992]

Node is adding AsyncLocalStorage#use() as a Disposable alternative to enterWith, tho you can also use enterWith in your own Dispose impl. https://github.com/nodejs/node/pull/58104

using and AsyncLocalStorage are both new to me, I'm not sure what the Gotchas are.

class LoggerContext {
  constructor(defaultLogger: Logger) {
    this.als = new AsyncLocalStorage({ default: defaultLogger })
  }

  withKeys(keys) {
    const current = this.get()
    const child =  current.child(keys)
    // XXX If 'use' is not supported, could be `enterWith` and return a Logger with [Symbol.dispose] that restores
    // `current` in this.als.
    return this.als.use(child)
  }

  get() {
      return = this.als.getStore()
  }
}

const LOG_CTX = new LoggerContext(new Logger())

async function main() {
  // Create a new child logger with given keys, `logger` is normal Logger, call withAdditonalKeys or w/e
  using logger = LOG_CTX.withKeys({ id: id++ })
  await Promise.all([
    task(1),
    task(2),
  ])
}

async function task(n) {
  using logger = LOG_CTX.withKeys({task: n})
  await step(n)
}

async function step(n) {
  const logger = LOG_CTX.get() // Get the current context logger, without creating a child
  logger.info() // Would have id, and task key based on async callstack.
}

[dreamorosi]

Hi @everett1992, sorry for the delay on this.

I think we could consider splitting the feature request for AsyncLocalStorage and explicit resource management into separate work streams.

While the pattern you described above definitely works well together, making the logger instance disposable can be useful also without shared context and at this stage I am still not clear on whether there's any performance overhead in spawning a new AsyncLocalStorage.

For example, following your example, today you can do the same by passing

logger.ts

export cost logger = new Logger({});

handler.ts

import { logger as mainLogger } from './logger.js';

const processFn = async (n) => {
  const logger = mainLogger.createChild();
  logger.appendKeys({ task: n });
}

Since with AsyncLocalContext you'd still need to pass around a reference to the store (or bound getter function) the code would change slightly:

logger.ts

export cost logger = new Logger({}); // this now has an internal `AsyncLocalContext` store

import { logger as mainLogger } from './logger.js';

const processFn = async (n) => {
  const logger = mainLogger.withKeys({ task: n });
}

Additionally, since the children is anyway created as part of the function scope if there's no other reference to it, it'd still be garbage collected once processFn completes, so even the explicit resource management use case is not as compelling when using child loggers.

Something that might be more useful would be if we could change the internal implementation of the logger so that the key store itself is disposable, rather than the entire logger instance.

This way, even without creating child loggers, one could do this:

const appLogger = new Logger();

const processFn = async (n) => {
  const logger = mainLogger.withKeys({ task: n });
}

export const handler = async (event, context) => {
  using logger = appLogger.withKeys({ userId: event.id });

  await processFn(event.taskId); 
}

---

We are currently working on some Logger improvements that we are not ready to share publicly yet, but as part of these we will definitely keep in mind this feature request.

I'll share more updates as soon as I can, but most likely not before the next couple months.