diff --git a/packages/runtime/src/scr_async.c b/packages/runtime/src/scr_async.c index 9d3a25469..9e3365fbe 100644 --- a/packages/runtime/src/scr_async.c +++ b/packages/runtime/src/scr_async.c @@ -674,6 +674,32 @@ void scr_als_ctx_release(ScrAlsCtx *c) { free(c); } +/* enterWith on the main stack lasts until that code returns to the loop, + * so the store does not reach callbacks scheduled without it: the loop + * clears the main slot when the main module's body returns, and an + * enterWith inside a callback the loop runs puts a tick that clears it at + * the front of the tick queue, so it runs as soon as that callback + * returns. A store that an async callee's prefix or a generator step + * enters becomes the caller's at scr_als_prefix_end and is cleared the + * same way. */ +static SCR_TL bool scr_loop_started = false; + +static void scr_als_main_clear(void) { + scr_als_ctx_release(scr_als_main_slot); + scr_als_main_slot = NULL; +} + +static void scr_next_tick_raw_first(void (*fn)(void)); + +/* A new store on the main slot: one a loop callback enters ends with it. */ +static void scr_als_main_entered(void) { + if (scr_loop_started) scr_next_tick_raw_first(scr_als_main_clear); +} + +void scr_als_entered_with(void) { + if (SCR_ALS_SLOT() == &scr_als_main_slot) scr_als_main_entered(); +} + /* The end of an async body's synchronous prefix: its first await, or its * end. Node runs that prefix inside the caller's execution, so a store it * entered (enterWith) is still the caller's store after the call; the body @@ -685,6 +711,7 @@ static void scr_als_prefix_end(ScrFiber *f) { if (*slot != f->als) { scr_als_ctx_release(*slot); *slot = scr_als_ctx_retain(f->als); + if (f->als != NULL && slot == &scr_als_main_slot) scr_als_main_entered(); } } @@ -758,26 +785,6 @@ ScrDyn *scr_als_peek(double id) { return NULL; } -/* enterWith on the main stack lasts until that code returns to the loop, - * so the store does not reach callbacks scheduled without it: the loop - * clears the main slot when the main module's body returns, and an - * enterWith inside a callback the loop runs puts a tick that clears it at - * the front of the tick queue, so it runs as soon as that callback - * returns. */ -static SCR_TL bool scr_loop_started = false; - -static void scr_als_main_clear(void) { - scr_als_ctx_release(scr_als_main_slot); - scr_als_main_slot = NULL; -} - -static void scr_next_tick_raw_first(void (*fn)(void)); - -void scr_als_entered_with(void) { - if (scr_loop_started && SCR_ALS_SLOT() == &scr_als_main_slot) - scr_next_tick_raw_first(scr_als_main_clear); -} - static SCR_TL ScrCtx scr_loop_ctx; /* the main stack (scheduler home) */ static SCR_TL ScrFiber *scr_current = NULL; /* A finished fiber whose last switch went to the loop (scr_resume_fiber). */ diff --git a/tests/corpus/7106-async-prefix-enter-with.ts b/tests/corpus/7106-async-prefix-enter-with.ts index 27dff250b..d70822171 100644 --- a/tests/corpus/7106-async-prefix-enter-with.ts +++ b/tests/corpus/7106-async-prefix-enter-with.ts @@ -3,7 +3,9 @@ // after the call (Node's semantics), and the callee keeps it after its own // await; after the first await it stays inside the callee. Covers callees // that never suspend, stackless and fiber-backed callees, nested callees, -// run() scopes inside the prefix, and calls without an await. +// run() scopes inside the prefix, and calls without an await. A store that +// an event-loop callback enters this way, or through a generator step, ends +// with that callback and does not reach the next one. import { AsyncLocalStorage } from "node:async_hooks"; const store = new AsyncLocalStorage(); @@ -45,6 +47,33 @@ async function fibered(tag: string): Promise { return n; } +function* steps(): Generator { + store.enterWith("generator-step"); + yield 1; +} +function callbacks(): Promise { + return new Promise((resolve) => { + setTimeout(() => { + void prefixThenAwait("9"); + log.push(`timer after the call: ${current()}`); + }, 0); + setTimeout(() => { + log.push(`second timer: ${current()}`); + steps().next(); + log.push(`second timer after a generator step: ${current()}`); + }, 0); + setTimeout(() => { + log.push(`third timer: ${current()}`); + void neverSuspends("10"); + log.push(`third timer after the call: ${current()}`); + }, 0); + setTimeout(() => { + log.push(`last timer: ${current()}`); + resolve(); + }, 0); + }); +} + async function main(): Promise { await store.run("A", async () => { await neverSuspends("1"); @@ -83,6 +112,8 @@ async function main(): Promise { await fibered("8"); log.push(`H after fibered: ${current()}`); }); + await callbacks(); + log.push(`main after the callbacks: ${current()}`); for (const line of log) console.log(line); } void main();