[JSC] A job of the module loader runs under the async context of the import() that asked for it

JSTests/BUN-TEST-DIFFERENCES.md

@@ -184,7 +184,7 @@ Not upstream's: a sync neither brings nor changes them. Most are named `stress/b
 
 | file | what it holds the engine to |
 |---|---|
-| `stress/bun-async-context-propagation.js` | the embedder's async-context slot is captured when a job is scheduled and current while it runs |
+| `stress/bun-async-context-propagation.js` | the embedder's async-context slot is captured when a job is scheduled and current while it runs, the jobs of the module loader included (its modules are in `stress/resources/bun-async-context-propagation/`) |
 | `stress/bun-aot-builtins-table.js` | the ahead-of-time compiler's table of built-ins (`Source/JavaScriptCore/aot/builtins/*.builtins.ts`) agrees with the engine. **After a sync that removes, renames, moves or reimplements a built-in that the table knows this fails, and each line it prints names the entry to fix.** |
 | `stress/sort-scratch-double-pivot.js` | double arithmetic in the comparator of an inlined sort stores no double into an Int32 array |
 | `stress/array-sort-inline-int32-double-comparator.js` | an earlier, smaller test for the sort fix |

JSTests/stress/bun-async-context-propagation.js

@@ -174,3 +174,124 @@ let log = [];
     shouldBe(seen, undefined, "residue from a sibling microtask");
     shouldBe(get(), undefined, "residue after a microtask that captured nothing");
 }
+
+// The module loader loads, links and evaluates an import() in jobs of its own. Where a context was current at the
+// import(), they run under it, and what a module puts in the slot ends with the job that ran the module.
+{
+    log = [];
+    const entered = ["entered"];
+    globalThis.entered = entered;
+    const settled = [];
+    const settle = name => [() => settled.push([name, get()]), error => settled.push([name, get(), error])];
+    const resources = "./resources/bun-async-context-propagation/";
+    inContext(A, () => {
+        import(resources + "graph.js").then(...settle("graph"));
+        import(resources + "tla-importer.js").then(...settle("tla importer"));
+        import(resources + "thenable.js").then(...settle("thenable"));
+        import(resources + "enters.js").then(...settle("enters"));
+    });
+    inContext(B, () => {
+        import(resources + "other-graph.js").then(...settle("other graph"));
+    });
+    import(resources + "plain.js").then(...settle("plain"));
+    drainMicrotasks();
+    shouldBe(get(), undefined, "residue after the jobs of the module loader");
+    const expected = {
+        // Module bodies, a static dependency first, and an import() that a body starts.
+        "graph dependency": A, "graph": A, "graph nested": A,
+        "other graph dependency": B, "other graph": B,
+        // A body behind its own await, and the module that waits for it.
+        "tla before await": A, "tla after await": A, "tla importer": A,
+        // The `then` of a module namespace.
+        "namespace then": A,
+        // A module that enters a context keeps it to the end of its job, and no longer.
+        "enters dependency": A, "enters": entered,
+        // An import() under no context.
+        "plain dependency": undefined, "plain": undefined,
+    };
+    for (const [name, ctx] of log)
+        shouldBe(ctx, expected[name], name);
+    shouldBe(log.length, Object.keys(expected).length, "all module bodies ran");
+    const expectedSettled = { "graph": A, "tla importer": A, "thenable": A, "enters": A, "other graph": B, "plain": undefined };
+    for (const [name, ctx, error] of settled) {
+        shouldBe(error, undefined, "import() of " + name + " rejected with " + error + ":");
+        shouldBe(ctx, expectedSettled[name], "import() of " + name);
+    }
+    shouldBe(settled.length, Object.keys(expectedSettled).length, "all imports settled");
+}
+
+// A module that waits for a dependency with a top-level await runs when the body of the dependency ends, in the job of
+// whoever began to evaluate the dependency. So a module that B imports runs under A where A's import() began the
+// dependency, as in Node.js. Each import() settles under its own context.
+{
+    log = [];
+    let open;
+    globalThis.gate = new Promise(resolve => { open = resolve; });
+    const settled = [];
+    const settle = name => [() => settled.push([name, get()]), error => settled.push([name, get(), error])];
+    const resources = "./resources/bun-async-context-propagation/";
+    inContext(A, () => {
+        import(resources + "waits-first.js").then(...settle("waits first"));
+    });
+    drainMicrotasks();
+    inContext(B, () => {
+        import(resources + "waits-second.js").then(...settle("waits second"));
+    });
+    drainMicrotasks();
+    shouldBe(log.length, 1, "the dependency waits");
+    open();
+    drainMicrotasks();
+    shouldBe(get(), undefined, "residue after the modules that waited");
+    const expected = { "waits before its await": A, "waits behind its await": A, "waits first": A, "waits second": A };
+    for (const [name, ctx] of log)
+        shouldBe(ctx, expected[name], name);
+    shouldBe(log.length, Object.keys(expected).length, "all modules that waited ran");
+    const expectedSettled = { "waits first": A, "waits second": B };
+    for (const [name, ctx, error] of settled) {
+        shouldBe(error, undefined, "import() of " + name + " rejected with " + error + ":");
+        shouldBe(ctx, expectedSettled[name], "import() of " + name);
+    }
+    shouldBe(settled.length, Object.keys(expectedSettled).length, "both imports settled");
+}
+
+// The job that runs the modules behind a top-level await gets the context of the load, not what the slot holds where
+// the load begins to evaluate. Here an import() under no context is loaded by a drain that something calls under A:
+// the awaiting module runs under A as far as its await, and behind it. The module that waits for it runs later, under
+// no context.
+{
+    log = [];
+    let open;
+    globalThis.laterGate = new Promise(resolve => { open = resolve; });
+    let settledUnder = "not settled";
+    const resources = "./resources/bun-async-context-propagation/";
+    import(resources + "later-importer.js").then(() => { settledUnder = get(); });
+    inContext(A, () => drainMicrotasks());
+    shouldBe(log.length, 1, "the dependency waits");
+    open();
+    drainMicrotasks();
+    shouldBe(get(), undefined, "residue after the module that waited");
+    const expected = { "later before its await": A, "later behind its await": A, "later importer": undefined };
+    for (const [name, ctx] of log)
+        shouldBe(ctx, expected[name], name);
+    shouldBe(log.length, Object.keys(expected).length, "the module that waited ran");
+    shouldBe(settledUnder, undefined, "import() of later importer");
+}
+
+// Where the job of a load entered the context of the load, the job behind an await gets what is current where the
+// awaiting module begins, as any continuation does. Here that is what a module before it has put in the slot.
+{
+    log = [];
+    let open;
+    globalThis.enteredGate = new Promise(resolve => { open = resolve; });
+    const resources = "./resources/bun-async-context-propagation/";
+    inContext(A, () => {
+        import(resources + "sets-then-awaits.js");
+    });
+    drainMicrotasks();
+    shouldBe(log.length, 0, "the module waits");
+    open();
+    drainMicrotasks();
+    shouldBe(get(), undefined, "residue after the module that waited");
+    shouldBe(log.length, 1, "the module that waited ran");
+    shouldBe(log[0][1], globalThis.entered, "behind sets and awaits");
+}

JSTests/stress/resources/bun-async-context-propagation/awaits-gate.js

@@ -0,0 +1 @@
+await enteredGate;

JSTests/stress/resources/bun-async-context-propagation/enters-dependency.js

@@ -0,0 +1 @@
+log.push(["enters dependency", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/enters.js

@@ -0,0 +1,4 @@
+import "./enters-dependency.js";
+
+$vm.setAsyncContext(entered);
+log.push(["enters", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/graph-dependency.js

@@ -0,0 +1 @@
+log.push(["graph dependency", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/graph-nested.js

@@ -0,0 +1 @@
+log.push(["graph nested", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/graph.js

@@ -0,0 +1,4 @@
+import "./graph-dependency.js";
+
+log.push(["graph", $vm.asyncContext()]);
+import("./graph-nested.js");

JSTests/stress/resources/bun-async-context-propagation/later-importer.js

@@ -0,0 +1,3 @@
+import "./later.js";
+
+log.push(["later importer", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/later.js

@@ -0,0 +1,3 @@
+log.push(["later before its await", $vm.asyncContext()]);
+await laterGate;
+log.push(["later behind its await", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/other-graph-dependency.js

@@ -0,0 +1 @@
+log.push(["other graph dependency", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/other-graph.js

@@ -0,0 +1,3 @@
+import "./other-graph-dependency.js";
+
+log.push(["other graph", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/plain-dependency.js

@@ -0,0 +1 @@
+log.push(["plain dependency", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/plain.js

@@ -0,0 +1,3 @@
+import "./plain-dependency.js";
+
+log.push(["plain", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/sets-then-awaits.js

@@ -0,0 +1,4 @@
+import "./sets.js";
+import "./awaits-gate.js";
+
+log.push(["behind sets and awaits", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/sets.js

@@ -0,0 +1 @@
+$vm.setAsyncContext(entered);

JSTests/stress/resources/bun-async-context-propagation/thenable.js

@@ -0,0 +1,4 @@
+export function then(resolve) {
+    log.push(["namespace then", $vm.asyncContext()]);
+    resolve("thenable");
+}

JSTests/stress/resources/bun-async-context-propagation/tla-importer.js

@@ -0,0 +1,3 @@
+import "./tla.js";
+
+log.push(["tla importer", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/tla.js

@@ -0,0 +1,3 @@
+log.push(["tla before await", $vm.asyncContext()]);
+await null;
+log.push(["tla after await", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/waits-first.js

@@ -0,0 +1,3 @@
+import "./waits.js";
+
+log.push(["waits first", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/waits-second.js

@@ -0,0 +1,3 @@
+import "./waits.js";
+
+log.push(["waits second", $vm.asyncContext()]);

JSTests/stress/resources/bun-async-context-propagation/waits.js

@@ -0,0 +1,3 @@
+log.push(["waits before its await", $vm.asyncContext()]);
+await gate;
+log.push(["waits behind its await", $vm.asyncContext()]);

Source/JavaScriptCore/aot/runtime/AOTModuleEvaluation.cpp

@@ -20,6 +20,7 @@
 #include "AOTState.h"
 #include "AOTStronglyConnectedComponents.h"
 #include "AOTStubEntrypoints.h"
+#include "AsyncContextSwapScope.h"
 #include "ExceptionHelpers.h"
 #include "GlobalObjectMethodTable.h"
 #include "JSCInlines.h"
@@ -856,7 +857,7 @@ void Modules::resumeDeferredLoadsOfFinishedFiles()
 }
 
 // `bun run` runs an imported CommonJS file while it loads the graph of the importer. No ES module of that graph has begun by then, so a require() in the file that leads back to one evaluates it there and then.
-bool Modules::evaluateImportedCommonJSFiles(ModuleID root, JobsOfImportedFiles jobs, ErrorOfImportedFile error, ModuleGraphLoadingState* load)
+bool Modules::evaluateImportedCommonJSFiles(ModuleID root, JobsOfImportedFiles jobs, ErrorOfImportedFile error, ModuleGraphLoadingState* load, JSValue asyncContextOfImport)
 {
     VM& vm = *m_instance.vm;
     auto scope = DECLARE_THROW_SCOPE(vm);
@@ -894,10 +895,14 @@ bool Modules::evaluateImportedCommonJSFiles(ModuleID root, JobsOfImportedFiles j
     bool didRunFile = false;
     MicrotaskQueue& queue = vm.defaultMicrotaskQueue();
     FilesBeingRun filesBeingRun { root, std::nullopt, m_innermostFilesBeingRun, jobs == JobsOfImportedFiles::WaitForLoader, load };
+    filesBeingRun.asyncContext = AsyncContextSwapScope::current(vm, m_instance.globalObject);
+    // What the jobs of the loader would run under: the async context of the load, not what is current here.
+    JSValue asyncContextOfLoad = load ? load->asyncContext() : asyncContextOfImport;
     SetForScope innermostFilesBeingRun(m_innermostFilesBeingRun, &filesBeingRun);
     SetForScope modulesWhoseFilesRunAheadOfLoader(vm.aotState().modulesWhoseFilesRunAheadOfLoader, filesBeingRun.isAheadOfJobsOfLoader ? this : vm.aotState().modulesWhoseFilesRunAheadOfLoader);
     // The start of the root calls the files one after the other, and leaves where one of them has queued a job or where the root would begin.
-    if (jobs == JobsOfImportedFiles::WaitForLoader && data(root).evaluatesItself() && data(root).isEntryPoint() && queue.isEmpty() && !m_rootWhoseFilesRun) {
+    // Not where the files run under the async context of a load: each of them has a scope of its own then, in the loop below, as each has a job of its own under `bun run`.
+    if (jobs == JobsOfImportedFiles::WaitForLoader && data(root).evaluatesItself() && data(root).isEntryPoint() && queue.isEmpty() && !m_rootWhoseFilesRun && !ModuleLoaderAsyncContextScope::wouldEnter(vm, asyncContextOfLoad)) {
         ensureRecord(root)->retain();
         if (storedStatus(root) == StoredModuleStatus::InRecord) {
             callBodyUnderOwnMark(root, false, true);
@@ -977,11 +982,17 @@ bool Modules::evaluateImportedCommonJSFiles(ModuleID root, JobsOfImportedFiles j
         size_t numberOfJobs = numberOfJobsThatWouldWait(jobs);
         filesBeingRun.fileOrFace = imported;
         filesBeingRun.numberOfJobsInFrontOfFile = numberOfJobs;
-        // As the start of the root would call it.
-        if (storedStatus(imported) < StoredModuleStatus::InRecord && data(imported).evaluatesItself())
-            evaluateByCallingBody(imported);
-        else
-            evaluateModuleInRecord(imported);
+        {
+            // The loader of `bun run` runs the file in a job of its own, under the async context of the load that asked for the file: what the file enters ends with that job where the load has a store.
+            std::optional<ModuleLoaderAsyncContextScope> asyncContextOfFile;
+            if (jobs == JobsOfImportedFiles::WaitForLoader)
+                asyncContextOfFile.emplace(vm, m_instance.globalObject, asyncContextOfLoad);
+            // As the start of the root would call it.
+            if (storedStatus(imported) < StoredModuleStatus::InRecord && data(imported).evaluatesItself())
+                evaluateByCallingBody(imported);
+            else
+                evaluateModuleInRecord(imported);
+        }
         Exception* exception = scope.exception();
         if (filesBeingRun.fileHandedToLoader) [[unlikely]] {
             JSValue errorOfFile;
@@ -1110,6 +1121,8 @@ void Modules::handGraphToLoader(FilesBeingRun& files)
         SetForScope fileBeingHandedToLoader(m_fileBeingHandedToLoader, fileOf(*running));
         // The steps were due before any require() of the file began.
         SetForScope noSynchronousLoad(vm.m_synchronousModuleQueueSetAside, vm.m_synchronousModuleQueue);
+        // And before the file entered an async context: a step of a load that keeps none runs under what was current then.
+        AsyncContextSwapScope asyncContextBeforeFiles(vm, m_instance.globalObject, files.asyncContext);
         if (files.load) {
             ModuleRecord* record = recordIfExists(files.root);
             files.load->forgetVisited(record);

Source/JavaScriptCore/aot/runtime/AOTModuleMicrotasks.h

@@ -16,9 +16,10 @@ class JSGlobalObject;
 
 namespace AOT {
 
-void runCompiledModuleImport(JSGlobalObject*, JSValue importPromise, JSValue module, JSValue referrerAsyncOrder);
-void runCompiledModuleImportSettled(JSGlobalObject*, JSValue importPromise, JSValue resolution, JSValue module, JSPromise::Status);
-void runCompiledModuleImportStep(JSValue importPromise, JSValue module, JSValue error, uint8_t stepsLeft);
+// asyncContext: what the job was queued with, which the job that it queues next is queued with too.
+void runCompiledModuleImport(JSGlobalObject*, JSValue importPromise, JSValue module, JSValue referrerAsyncOrder, JSValue asyncContext);
+void runCompiledModuleImportSettled(JSGlobalObject*, JSValue importPromise, JSValue resolution, JSValue module, JSPromise::Status, JSValue asyncContext);
+void runCompiledModuleImportStep(JSValue importPromise, JSValue module, JSValue error, uint8_t stepsLeft, JSValue asyncContext);
 // Behind the last step that the loader takes for a graph.
 void didSettleLoad(JSPromise* settled);