From 9e1027e13a35c1e6503f4e1b72e9c1aa3bb1af69 Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Thu, 22 Oct 2026 11:46:44 -0500 Subject: [PATCH] Run async generator steps in the AsyncLocalStorage context of their resumer An async generator kept the context of its creation for its whole life, and a store its body entered carried into every later step. Node runs each step in the execution that resumes it: a step starts with the context of the next(), return() or throw() call, a store it enters before its first await is the caller's after the call, and after an await the step continues in its own copy. Async generators now start each new step from the resumer's context and hand the prefix's context back at the step's first suspension, as synchronous generator steps and async function prefixes already do; the continuation after an await keeps the generator's own. Two requests that start from the settlement job, which runs from the loop, take Node's contexts too: a request queued behind a running step starts in that step's context, and a return() at a yield, whose unwind waits a turn, runs in the return() call's context. Corpus 7381 and 7382. --- .../7381-als-async-generator-steps.ts.json | 6 + .../7382-als-async-generator-steps.mjs.json | 6 + packages/runtime/src/scr_async.c | 80 +++++++++-- .../corpus/7381-als-async-generator-steps.ts | 133 ++++++++++++++++++ .../corpus/7382-als-async-generator-steps.mjs | 43 ++++++ .../7381-als-async-generator-steps.ts.json | 1 + .../7382-als-async-generator-steps.mjs.json | 1 + 7 files changed, 257 insertions(+), 13 deletions(-) create mode 100644 packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7381-als-async-generator-steps.ts.json create mode 100644 packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7382-als-async-generator-steps.mjs.json create mode 100644 tests/corpus/7381-als-async-generator-steps.ts create mode 100644 tests/corpus/7382-als-async-generator-steps.mjs create mode 100644 tests/harness/corpus-costs/7381-als-async-generator-steps.ts.json create mode 100644 tests/harness/corpus-costs/7382-als-async-generator-steps.mjs.json diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7381-als-async-generator-steps.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7381-als-async-generator-steps.ts.json new file mode 100644 index 000000000..47f93314d --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7381-als-async-generator-steps.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/7381-als-async-generator-steps.ts" + ], + "diags": [] +} diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7382-als-async-generator-steps.mjs.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7382-als-async-generator-steps.mjs.json new file mode 100644 index 000000000..5f55ba7bd --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7382-als-async-generator-steps.mjs.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/7382-als-async-generator-steps.mjs" + ], + "diags": [] +} diff --git a/packages/runtime/src/scr_async.c b/packages/runtime/src/scr_async.c index 90e7dd04f..e6184ed9b 100644 --- a/packages/runtime/src/scr_async.c +++ b/packages/runtime/src/scr_async.c @@ -715,11 +715,13 @@ static void scr_als_prefix_end(ScrFiber *f) { } } -/* A synchronous generator has no context of its own: each step runs in its - * resumer's execution, as Node runs it. A fiber-backed one switches to its - * own slot, so its step starts from the resumer's store and a store it - * enters is the resumer's after the step (scr_als_prefix_end on the way - * back); a stackless one runs in the resumer's slot (scr_gen_coro_hop). */ +/* A generator step runs in its resumer's execution, as Node runs it. A + * fiber-backed generator switches to its own slot, so its step starts from + * the resumer's store and a store it enters is the resumer's after the + * step (scr_als_prefix_end on the way back); a stackless synchronous one + * runs in the resumer's slot (scr_gen_coro_hop). An async generator's step + * starts the same way, and after an await it continues in its own copy of + * the context, as an async function's body does after its prefix. */ static void scr_als_step_begin(ScrFiber *f) { ScrAlsCtx **slot = SCR_ALS_SLOT(); if (f->als != *slot) { @@ -5303,6 +5305,10 @@ struct ScrGen { * since that return takes its own turn to await x. */ bool return_queued; bool return_awaited; + /* The AsyncLocalStorage context of a return() request deferred at a + * yield (+1): the body's unwind runs in it, as Node resumes it from the + * return value's await. */ + ScrAlsCtx *return_ctx; #ifdef SCR_WORKERS bool context_abandoned; #endif @@ -5406,6 +5412,7 @@ void scr_gen_release(ScrGen *g) { scr_gen_slot_reset(&g->in); scr_gen_slot_reset(&g->ret); scr_gen_exc_reset(&g->async_error); + scr_als_ctx_release(g->return_ctx); if (g->fiber != NULL) { if (g->state == SCR_GEN_UNSTARTED) { /* Never ran: nothing on the stack owns anything — clean teardown. @@ -5760,6 +5767,12 @@ __attribute__((always_inline)) static inline void scr_gen_switched_back(ScrGen * if (release) scr_gen_release(g); } +/* Out of line, so the hop that synchronous generators share stays as it + * was: an async generator's step begins from its resumer's context and + * hands the prefix's context back at its first suspension. */ +static __attribute__((noinline)) void scr_async_gen_step_begin(ScrFiber *f) { scr_als_step_begin(f); } +static __attribute__((noinline)) void scr_async_gen_step_end(ScrFiber *f) { scr_als_prefix_end(f); } + #ifndef __wasi__ /* The stackless form of the hop: the coroutine frame runs on this stack * until its next yield or its final suspension returns here. `retain` @@ -5768,14 +5781,17 @@ __attribute__((always_inline)) static inline void scr_gen_switched_back(ScrGen * __attribute__((always_inline)) static inline void scr_gen_coro_hop(ScrGen *g, bool retain) { ScrFiber *f = g->fiber; if (retain) scr_gen_retain(g); + /* A synchronous generator has no AsyncLocalStorage context of its own: + * its step runs in the resumer's slot, as Node runs it in the execution + * that called next(). An async generator's step runs in its own slot, + * which a new step (not the continuation after an await) starts from the + * resumer's context (scr_als_step_begin). */ + if (g->is_async && g->state != SCR_GEN_AWAITING) scr_async_gen_step_begin(f); g->state = SCR_GEN_RUNNING; f->resumer = scr_current; ScrAlsCtx **resumer_slot = scr_als_active; scr_current = f; scr_exc_swap_cell(&f->exc); - /* A synchronous generator has no AsyncLocalStorage context of its own: - * its step runs in the resumer's slot, as Node runs it in the execution - * that called next(). An async generator's step runs in its own. */ if (g->is_async) scr_als_active = &f->als; if (f->coro == NULL) f->entry(f, f->argpack); else scr_coro_resume(f->coro); @@ -5786,6 +5802,7 @@ __attribute__((always_inline)) static inline void scr_gen_coro_hop(ScrGen *g, bo scr_current = me; scr_exc_swap_cell(me != NULL ? &me->exc : NULL); scr_als_active = resumer_slot; + if (g->is_async) scr_async_gen_step_end(f); scr_gen_switched_back(g, f, retain, true); } @@ -5806,8 +5823,10 @@ static void scr_gen_switch_in(ScrGen *g) { } #endif scr_gen_retain(g); + /* Each step starts from the resumer's context; an async generator's + * continuation after an await keeps its own. */ + if (!g->is_async || g->state != SCR_GEN_AWAITING) scr_als_step_begin(f); g->state = SCR_GEN_RUNNING; - if (!g->is_async) scr_als_step_begin(f); #ifdef _WIN32 ScrCtx here = scr_win_self(); #elif defined(__wasi__) @@ -5979,6 +5998,25 @@ void scr_gen_resume_throw(ScrGen *g) { static void scr_async_gen_start_queued(ScrGen *g); +/* Runs fn(g) with `ctx` (borrowed) as the active AsyncLocalStorage + * context. A settlement job runs from the loop: a request queued behind + * another starts in the context of the step before it (Node drains the + * queue in the generator's own execution), and a deferred return() in its + * call's. */ +static __attribute__((noinline)) void scr_async_gen_in_context(ScrGen *g, ScrAlsCtx *ctx, + void (*fn)(ScrGen *)) { + ScrAlsCtx **slot = SCR_ALS_SLOT(); + ScrAlsCtx *prev = *slot; + *slot = scr_als_ctx_retain(ctx); + fn(g); + scr_als_restore(prev); +} + +/* Starts the queued requests in the context of the step before them. */ +static __attribute__((noinline)) void scr_async_gen_start_queued_in_own(ScrGen *g) { + scr_async_gen_in_context(g, g->fiber != NULL ? g->fiber->als : NULL, scr_async_gen_start_queued); +} + static void scr_async_gen_schedule_settle(ScrGen *g) { if (g->settle_queued) return; ScrFiber *job = scr_mem_calloc(sizeof *job); @@ -6026,6 +6064,11 @@ static bool scr_async_gen_settles_now(ScrGen *g) { return false; } +static __attribute__((noinline)) void scr_async_gen_capture_return_ctx(ScrGen *g) { + scr_als_ctx_release(g->return_ctx); + g->return_ctx = scr_als_ctx_retain(*SCR_ALS_SLOT()); +} + static void scr_async_gen_run_active(ScrGen *g, ScrAsyncGenRequest *req) { const bool was_done = g->state == SCR_GEN_DONE; if (req->mode == SCR_AG_NEXT) { @@ -6041,6 +6084,7 @@ static void scr_async_gen_run_active(ScrGen *g, ScrAsyncGenRequest *req) { /* At a yield, AsyncGeneratorYield awaits the return value before * the body resumes, so finally blocks run one turn later. */ g->return_queued = true; + scr_async_gen_capture_return_ctx(g); scr_async_gen_schedule_settle(g); return; } @@ -6071,6 +6115,16 @@ static void scr_async_gen_run_active(ScrGen *g, ScrAsyncGenRequest *req) { } } +static void scr_async_gen_resume_returning(ScrGen *g); + +/* The deferred return() resumes in its call's context. */ +static __attribute__((noinline)) void scr_async_gen_resume_returning_in_context(ScrGen *g) { + ScrAlsCtx *ctx = g->return_ctx; + g->return_ctx = NULL; + scr_async_gen_in_context(g, ctx, scr_async_gen_resume_returning); + scr_als_ctx_release(ctx); +} + /* The return() request's await has taken its turn: resume the body with * the return completion. */ static void scr_async_gen_resume_returning(ScrGen *g) { @@ -6203,12 +6257,12 @@ static void scr_async_gen_job(ScrGen *g) { } if (g->return_queued) { scr_gen_retain(g); - scr_async_gen_resume_returning(g); + scr_async_gen_resume_returning_in_context(g); scr_gen_release(g); return; } scr_async_gen_finish_active(g, false); - scr_async_gen_start_queued(g); + scr_async_gen_start_queued_in_own(g); } static void scr_async_gen_resume_ready(ScrGen *g) { @@ -6222,11 +6276,11 @@ static void scr_async_gen_resume_ready(ScrGen *g) { scr_gen_switch_in(g); if (g->async_error.kind != SCR_EXC_NONE) { scr_async_gen_finish_active(g, true); - scr_async_gen_start_queued(g); + scr_async_gen_start_queued_in_own(g); } else if (g->state != SCR_GEN_AWAITING) { if (scr_async_gen_settles_now(g)) { scr_async_gen_finish_active(g, false); - scr_async_gen_start_queued(g); + scr_async_gen_start_queued_in_own(g); } else { scr_async_gen_schedule_settle(g); } diff --git a/tests/corpus/7381-als-async-generator-steps.ts b/tests/corpus/7381-als-async-generator-steps.ts new file mode 100644 index 000000000..4f6ace58d --- /dev/null +++ b/tests/corpus/7381-als-async-generator-steps.ts @@ -0,0 +1,133 @@ +// AsyncLocalStorage contexts of async generator steps. Each step starts in +// the context of the next(), return() or throw() call that resumes it, +// whatever context created the generator; a store the step enters before +// its first await is the caller's after the call, and after an await the +// step keeps its own copy. A request queued behind a running step starts +// in that step's context, and a return() at a yield unwinds in the +// return() call's context. Covers yield*, for await with break, and two +// generators resumed in turn. +import { AsyncLocalStorage } from "node:async_hooks"; + +const store = new AsyncLocalStorage(); +const lines: string[] = []; +const log = (what: string): void => { + lines.push(`${what}: ${store.getStore() ?? "none"}`); +}; + +async function* steps(): AsyncGenerator { + log("body start"); + yield 1; + log("after the first yield"); + await null; + log("after an await"); + store.enterWith("inside"); + yield 2; + log("after the second yield"); + yield 3; + log("body end"); +} +const it = store.run("create", () => steps()); +await store.run("next-1", () => it.next()); +log("caller after next-1"); +await store.run("next-2", () => it.next()); +log("caller after next-2"); +await store.run("next-3", () => it.next()); +await store.run("next-4", () => it.next()); + +async function* prefixes(): AsyncGenerator { + store.enterWith("prefix"); + yield 1; + store.enterWith("step-2"); + await null; + log("after an await in step 2"); + yield 2; +} +await store.run("outer", async () => { + const gen = prefixes(); + const first = gen.next(); + log("synchronously after next"); + await first; + log("after awaiting next"); + const second = gen.next(); + log("synchronously after the second next"); + await second; + log("after awaiting the second next"); +}); + +async function* queued(): AsyncGenerator { + log("queued start"); + yield 1; + log("queued resumed"); + yield 2; +} +const q = queued(); +await Promise.all([store.run("q-1", () => q.next()), store.run("q-2", () => q.next())]); + +async function* inner(): AsyncGenerator { + log("inner start"); + yield "a"; + log("inner resumed"); + yield "b"; +} +async function* outerGen(): AsyncGenerator { + log("outer start"); + yield* inner(); + log("outer after yield*"); +} +const delegating = outerGen(); +await store.run("o-1", () => delegating.next()); +await store.run("o-2", () => delegating.next()); +await store.run("o-3", () => delegating.next()); + +async function* catching(): AsyncGenerator { + try { + yield 1; + } catch (e) { + log(`caught ${(e as Error).message}`); + yield 2; + } +} +const c = catching(); +await store.run("t-1", () => c.next()); +await store.run("t-2", () => c.throw(new Error("boom"))); + +async function* finallyGen(): AsyncGenerator { + try { + yield 1; + } finally { + log("finally"); + } +} +const f = finallyGen(); +await store.run("n", () => f.next()); +await store.run("return-call", () => f.return()); + +async function* endless(): AsyncGenerator { + let i = 0; + try { + while (true) { + log(`endless ${i}`); + yield i++; + } + } finally { + log("endless finally"); + } +} +await store.run("loop", async () => { + for await (const v of endless()) { + if (v === 1) break; + } +}); + +async function* named(name: string): AsyncGenerator { + log(`${name} 1`); + yield 1; + log(`${name} 2`); + yield 2; +} +const a = named("A"); +const b = named("B"); +await Promise.all([store.run("a-1", () => a.next()), store.run("b-1", () => b.next())]); +await Promise.all([store.run("b-2", () => b.next()), store.run("a-2", () => a.next())]); + +for (const line of lines) console.log(line); diff --git a/tests/corpus/7382-als-async-generator-steps.mjs b/tests/corpus/7382-als-async-generator-steps.mjs new file mode 100644 index 000000000..bbabdb592 --- /dev/null +++ b/tests/corpus/7382-als-async-generator-steps.mjs @@ -0,0 +1,43 @@ +// AsyncLocalStorage contexts of async generator steps from JavaScript: +// steps start in the context of the call that resumes them, a store a step +// enters before its first await is the caller's, queued requests start in +// the running step's context, and return() at a yield unwinds in the +// return() call's context. +import { AsyncLocalStorage } from "node:async_hooks"; +const s = new AsyncLocalStorage(); +const lines = []; +const g = (t) => lines.push(`${t}: ${s.getStore() ?? "none"}`); +async function* gen() { + g("body start"); + yield 1; + g("after first yield"); + await null; + g("after await"); + s.enterWith("inside"); + yield 2; + g("after second yield"); + yield 3; + g("body end"); +} +const it = s.run("create", () => gen()); +await s.run("next1", () => it.next()); +g("caller after next1"); +await s.run("next2", () => it.next()); +g("caller after next2"); +await s.run("next3", () => it.next()); +await s.run("next4", () => it.next()); +// concurrent next calls queued +async function* gen2() { g("g2 start"); yield 1; g("g2 resumed"); yield 2; } +const it2 = gen2(); +const p1 = s.run("q1", () => it2.next()); +const p2 = s.run("q2", () => it2.next()); +await Promise.all([p1, p2]); +// for await +async function* gen3() { for (let i = 0; i < 2; i++) { g("g3 step " + i); yield i; } } +await s.run("for-await", async () => { for await (const v of gen3()) g("loop body " + v); }); +// return/throw +async function* gen4() { try { yield 1; } finally { g("g4 finally"); } } +const it4 = gen4(); +await s.run("n", () => it4.next()); +await s.run("ret", () => it4.return()); +for (const line of lines) console.log(line); diff --git a/tests/harness/corpus-costs/7381-als-async-generator-steps.ts.json b/tests/harness/corpus-costs/7381-als-async-generator-steps.ts.json new file mode 100644 index 000000000..3f7384f20 --- /dev/null +++ b/tests/harness/corpus-costs/7381-als-async-generator-steps.ts.json @@ -0,0 +1 @@ +{ "plain": 3500, "san": 6000 } diff --git a/tests/harness/corpus-costs/7382-als-async-generator-steps.mjs.json b/tests/harness/corpus-costs/7382-als-async-generator-steps.mjs.json new file mode 100644 index 000000000..3f7384f20 --- /dev/null +++ b/tests/harness/corpus-costs/7382-als-async-generator-steps.mjs.json @@ -0,0 +1 @@ +{ "plain": 3500, "san": 6000 }