async generators do not capture async context

[devsnek]

tested on v17 but i don't suspect that matters much.

'use strict';

const { AsyncLocalStorage } = require('async_hooks');

const s = new AsyncLocalStorage();

async function normal() {
  console.log('async fn', s.getStore());
}

async function* gen() {
  console.log('async gen', s.getStore());
}

s.run('hello', normal);

s.run('hello', gen).next();

output:

async fn hello
async gen undefined

/cc @nodejs/async_hooks

i'm assuming this probably will need some sort of effort on the v8 side?

[benjamingr]

@devsnek is this v8 not firing the right promise hooks or on our side (or you haven't had a chance to investigate yet)?

[devsnek]

I'm not sure.

[Flarna]

Should async generators capture the async context?

[devsnek]

i dunno... its technically an abstraction that's outside of promises, but in terms of the usage of apis like AsyncLocalStorage, i think i would expect hooks to apply to them.

[Flarna]

next() is a sync call which runs using the asyncId of the caller. Doesn't seem that wrong to me.
e.g. if you use s.run('hello', () => g.next()) the context is visible in this next() call. in your sample code next is called directly after returning from s.run() therefore the store not active anymore.

If the generator would capture the id where it was created everything after await next() would end up in this context. Hard to tell if this is good or bad.

[iamgabrielsoft]

Are you missing the yield keyword, I can't find it

Your using a generator

[benjamingr]

Should async generators capture the async context?

User expectation would probably be "yes"

[RMHonor]

Using inspiration from open-telemetry/opentelemetry-js#2951 (comment), I was able to build a generic wrapper to work around this issue in the mean time.

export function bindAsyncGenerator<T = unknown, TReturn = any, TNext = unknown>(
  store: AsyncLocalStorage<unknown>,
  generator: AsyncGenerator<T, TReturn, TNext>,
): AsyncGenerator<T, TReturn, TNext> {
  const ctx = store.getStore();
  return {
    next: () => store.run(ctx, generator.next.bind(generator)),
    return: (args) => store.run(ctx, generator.return.bind(generator), args),
    throw: (args) => store.run(ctx, generator.throw.bind(generator), args),

    [Symbol.asyncIterator]() {
      return this;
    },
  };
}

[anonrig]

cc @nodejs/tsc who would have the most information to resolve this bug?

[mcollina]

I think @bengl can coordinate this.

[Flarna]

FWIW there is an ongoing discussion in scope of TC39 AsyncContext proposal regarding the same topic.

[mrmar99]

Got this problem when I was writing logging middleware for nice-grpc server.

Middleware:

export default function loggingMiddleware(logger: ILogger) {
  return async function* <Request, Response>(
    call: ServerMiddlewareCall<Request, Response>,
    context: CallContext,
  ) {
    const someId = context.metadata.get("some-id");

    return yield* withContext({ someId }, async function* () {
      try {
        return yield* call.next(call.request, context);
      } catch (err) {
        // ...
        throw err;
      }
    });
  };
}

In withContext function I use enterWith to attach storeData to my store:

import store from "path/to/AsyncLocalStorage/instance";

type TStoreData = {
  someId: string;
}

export function withContext<T>(storeData: TStoreData, callback: () => T): T {
  store.enterWith(storeData);
  return store.run(storeData, callback);
}

Now, I have access to data in storage in my callback function by importing store and calling store.getStore().

Let me know if there is a better way.