Clear a store that an async prefix enters in a loop callback

When an event-loop callback calls an async function or steps a generator
that enters a store before its first await, the store becomes the
callback's at the end of the prefix. On the main slot it now ends with
that callback, as an enterWith there does: the copy schedules the tick
that clears the main slot, so the store does not reach later callbacks.

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). */

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<string>();
@@ -45,6 +47,33 @@ async function fibered(tag: string): Promise<number> {
   return n;
 }
 
+function* steps(): Generator<number, void, undefined> {
+  store.enterWith("generator-step");
+  yield 1;
+}
+function callbacks(): Promise<void> {
+  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<void> {
   await store.run("A", async () => {
     await neverSuspends("1");
@@ -83,6 +112,8 @@ async function main(): Promise<void> {
     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();