JSC: evaluate dynamically imported modules under the importer's async context

Source/JavaScriptCore/runtime/CyclicModuleRecord.cpp

@@ -41,6 +41,7 @@
 #include "SourceProfiler.h"
 #include "SymbolTableInlines.h"
 #if USE(BUN_JSC_ADDITIONS)
+#include "AsyncContextSwapScope.h"
 #include "SyntheticModuleRecord.h"
 #endif
 #include "UnlinkedModuleProgramCodeBlock.h"
@@ -555,7 +556,16 @@ void CyclicModuleRecord::executeAsync(JSGlobalObject* globalObject)
     // 7. Let onRejected be CreateBuiltinFunction(rejectedClosure, 0, "", « »).
     // Also handled in JSMicrotask.cpp.
     // 8. Perform PerformPromiseThen(capability.[[Promise]], onFulfilled, onRejected).
+#if USE(BUN_JSC_ADDITIONS)
+    // AsyncModuleExecutionFulfilled runs the bodies of the modules waiting on this one, so
+    // it has to run under the async context this evaluation was started under (for a
+    // dynamic import(), the importer's context installed by dynamicImportLoadSettled).
+    // Capture it with the reaction, as resolveWithInternalMicrotaskForAsyncAwait does for
+    // AsyncModuleExecutionResume; the AsyncModuleExecutionDone microtask reinstalls it.
+    promise->performPromiseThenWithInternalMicrotask(vm, InternalMicrotask::AsyncModuleExecutionDone, nullptr, this, AsyncContextSwapScope::current(vm, globalObject));
+#else
     promise->performPromiseThenWithInternalMicrotask(vm, InternalMicrotask::AsyncModuleExecutionDone, nullptr, this);
+#endif
     // 9. Perform ! module.ExecuteModule(capability).
     execute(globalObject, promise);
     RETURN_IF_EXCEPTION(scope, void());

Source/JavaScriptCore/runtime/JSMicrotask.cpp

@@ -1207,7 +1207,7 @@ static void moduleLoadTopSettled(JSGlobalObject* globalObject, VM& vm, ThrowScop
             innerLoadFlags.add(ModuleLoadFlag::UseImportMap);
         if (context->dynamic()) {
 #if USE(BUN_JSC_ADDITIONS)
-            combinedCell = ModuleLoaderPayload::create(vm, statePromise, context->deferred(), context->referrerAsyncOrder());
+            combinedCell = ModuleLoaderPayload::create(vm, statePromise, context->deferred(), context->referrerAsyncOrder(), context->importerAsyncContext());
 #else
             combinedCell = ModuleLoaderPayload::create(vm, statePromise, context->deferred());
 #endif
@@ -1500,7 +1500,16 @@ static void dynamicImportLoadSettled(JSGlobalObject* globalObject, VM& vm, Throw
     if (!deferred) {
         // 6.c. Let evaluatePromise be module.Evaluate().
 #if USE(BUN_JSC_ADDITIONS)
-        JSPromise* evaluatePromise = module->evaluate(globalObject, dynamicPayload->referrerAsyncOrder(), capabilityPromise);
+        JSPromise* evaluatePromise = nullptr;
+        {
+            // The module bodies in this graph are the continuation of the import() call, so
+            // they run under the async context captured there (JSModuleLoader::loadModule).
+            // Modules held back by a top-level-await dependency execute later, from
+            // AsyncModuleExecutionDone, which reinstalls the context executeAsync snapshots
+            // while running inside this scope.
+            AsyncContextSwapScope asyncContextScope(vm, globalObject, dynamicPayload->importerAsyncContext());
+            evaluatePromise = module->evaluate(globalObject, dynamicPayload->referrerAsyncOrder(), capabilityPromise);
+        }
 #else
         JSPromise* evaluatePromise = module->evaluate(globalObject);
 #endif
@@ -1531,6 +1540,11 @@ static void dynamicImportLoadSettled(JSGlobalObject* globalObject, VM& vm, Throw
 
     // For each Module Record dep of evaluationList, append dep.Evaluate() to asyncDepsEvaluationPromises.
     MarkedArgumentBuffer asyncDepsEvaluationPromises;
+#if USE(BUN_JSC_ADDITIONS)
+    // As above: the eagerly evaluated async dependencies are the continuation of the
+    // import.defer() call, so they run under the async context captured there.
+    AsyncContextSwapScope asyncContextScope(vm, globalObject, dynamicPayload->importerAsyncContext());
+#endif
     for (AbstractModuleRecord* dep : evaluationList) {
 #if USE(BUN_JSC_ADDITIONS)
         JSPromise* depPromise = dep->evaluate(globalObject, dynamicPayload->referrerAsyncOrder(), capabilityPromise);
@@ -2220,6 +2234,12 @@ void runInternalMicrotask(JSGlobalObject* globalObject, VM& vm, InternalMicrotas
 #endif
 
     case InternalMicrotask::AsyncModuleExecutionDone: {
+        // CyclicModuleRecord::executeAsync captures the async context the module
+        // started executing under, so that the waiting ancestors executed by
+        // asyncExecutionFulfilled run under that context too.
+#if USE(BUN_JSC_ADDITIONS)
+        AsyncContextSwapScope asyncContextScope(vm, globalObject, arguments[3]);
+#endif
         auto* module = uncheckedDowncast<JSModuleRecord>(arguments[2]);
         RELEASE_AND_RETURN(scope, asyncModuleExecutionDone(module->realm(), module, arguments[1], static_cast<JSPromise::Status>(payload)));
     }

Source/JavaScriptCore/runtime/JSModuleLoader.cpp

@@ -50,6 +50,9 @@
 #include "SyntheticModuleRecord.h"
 #include "TopExceptionScope.h"
 #include "VMTrapsInlines.h"
+#if USE(BUN_JSC_ADDITIONS)
+#include "AsyncContextSwapScope.h"
+#endif
 #include <wtf/text/MakeString.h>
 
 namespace JSC {
@@ -354,10 +357,19 @@ JSPromise* JSModuleLoader::loadModule(JSGlobalObject* globalObject, const Identi
     JSPromise* promise = nullptr;
 
     ScriptFetchParameters::Type type = parameters ? parameters->type() : ScriptFetchParameters::Type::JavaScript;
+    JSValue importerAsyncContext = jsUndefined();
 #if USE(BUN_JSC_ADDITIONS)
     // Preserve the caller's ScriptFetchParameters (HostDefined data) into the loading context
     // before `parameters` is moved into fetch() below.
     RefPtr<ScriptFetchParameters> contextParameters = parameters ? parameters : ScriptFetchParameters::create(type);
+
+    // A dynamic import() evaluates its module graph under the async context (AsyncLocalStorage
+    // store) that was active at the import() call site. We are still synchronous with that call
+    // here, so the slot holds the importer's context; the graph is linked and evaluated several
+    // internal microtasks later, by which point the slot has been reset. Capture it now, before
+    // fetch() below can run host code, and carry it through to dynamicImportLoadSettled.
+    if (flags.contains(ModuleLoadFlag::Dynamic))
+        importerAsyncContext = AsyncContextSwapScope::current(vm, globalObject);
 #endif
 
     if (ModuleRegistryEntry* entry = getRegisteredMayBeNull(specifier, type)) {
@@ -384,7 +396,7 @@ JSPromise* JSModuleLoader::loadModule(JSGlobalObject* globalObject, const Identi
 #else
     AbstractModuleRecord::ModuleRequest request { specifier, ScriptFetchParameters::create(type) };
 #endif
-    auto* context = ModuleLoadingContext::create(vm, request, WTF::move(scriptFetcher), flags, referrerAsyncOrder);
+    auto* context = ModuleLoadingContext::create(vm, request, WTF::move(scriptFetcher), flags, referrerAsyncOrder, importerAsyncContext);
 
     JSPromise* intermediatePromise = JSPromise::create(vm, globalObject->promiseStructure());
     intermediatePromise->markAsHandled();

Source/JavaScriptCore/runtime/ModuleLoaderPayload.cpp

@@ -33,16 +33,18 @@ namespace JSC {
 
 const ClassInfo ModuleLoaderPayload::s_info = { "ModuleLoaderPayload"_s, nullptr, nullptr, nullptr, CREATE_METHOD_TABLE(ModuleLoaderPayload) };
 
-ModuleLoaderPayload::ModuleLoaderPayload(VM& vm, Structure* structure, JSPromise* promise, bool deferred, int64_t referrerAsyncOrder)
+ModuleLoaderPayload::ModuleLoaderPayload(VM& vm, Structure* structure, JSPromise* promise, bool deferred, int64_t referrerAsyncOrder, JSValue importerAsyncContext)
     : Base(vm, structure)
     , m_promise(promise, WriteBarrierEarlyInit)
 #if USE(BUN_JSC_ADDITIONS)
+    , m_importerAsyncContext(importerAsyncContext, WriteBarrierEarlyInit)
     , m_referrerAsyncOrder(referrerAsyncOrder)
 #endif
     , m_deferred(deferred)
 {
 #if !USE(BUN_JSC_ADDITIONS)
     UNUSED_PARAM(referrerAsyncOrder);
+    UNUSED_PARAM(importerAsyncContext);
 #endif
 }
 
@@ -60,13 +62,16 @@ void ModuleLoaderPayload::visitChildrenImpl(JSCell* cell, Visitor& visitor)
     Base::visitChildren(thisObject, visitor);
     visitor.append(thisObject->m_promise);
     visitor.append(thisObject->m_fulfillment);
+#if USE(BUN_JSC_ADDITIONS)
+    visitor.append(thisObject->m_importerAsyncContext);
+#endif
 }
 
 DEFINE_VISIT_CHILDREN(ModuleLoaderPayload);
 
-ModuleLoaderPayload* ModuleLoaderPayload::create(VM& vm, JSPromise* promise, bool deferred, int64_t referrerAsyncOrder)
+ModuleLoaderPayload* ModuleLoaderPayload::create(VM& vm, JSPromise* promise, bool deferred, int64_t referrerAsyncOrder, JSValue importerAsyncContext)
 {
-    ModuleLoaderPayload* instance = new (NotNull, allocateCell<ModuleLoaderPayload>(vm)) ModuleLoaderPayload(vm, vm.moduleLoaderPayloadStructure.get(), promise, deferred, referrerAsyncOrder);
+    ModuleLoaderPayload* instance = new (NotNull, allocateCell<ModuleLoaderPayload>(vm)) ModuleLoaderPayload(vm, vm.moduleLoaderPayloadStructure.get(), promise, deferred, referrerAsyncOrder, importerAsyncContext);
     instance->finishCreation(vm);
     return instance;
 }

Source/JavaScriptCore/runtime/ModuleLoaderPayload.h

@@ -49,7 +49,7 @@ class ModuleLoaderPayload final : public JSCell {
     }
 
     inline static Structure* createStructure(VM&, JSGlobalObject*, JSValue);
-    static ModuleLoaderPayload* create(VM&, JSPromise*, bool deferred = false, int64_t referrerAsyncOrder = -1);
+    static ModuleLoaderPayload* create(VM&, JSPromise*, bool deferred = false, int64_t referrerAsyncOrder = -1, JSValue importerAsyncContext = jsUndefined());
 
     JSPromise* promise() const { return m_promise.get(); }
 
@@ -59,6 +59,8 @@ class ModuleLoaderPayload final : public JSCell {
     bool deferred() const { return m_deferred; }
 #if USE(BUN_JSC_ADDITIONS)
     int64_t referrerAsyncOrder() const { return m_referrerAsyncOrder; }
+    // The async context (AsyncLocalStorage store) active at the import() call site.
+    JSValue importerAsyncContext() const { return m_importerAsyncContext.get(); }
 #endif
 
     bool decrementRemaining()
@@ -68,13 +70,14 @@ class ModuleLoaderPayload final : public JSCell {
     }
 
 private:
-    ModuleLoaderPayload(VM&, Structure*, JSPromise*, bool deferred, int64_t referrerAsyncOrder);
+    ModuleLoaderPayload(VM&, Structure*, JSPromise*, bool deferred, int64_t referrerAsyncOrder, JSValue importerAsyncContext);
 
     void finishCreation(VM&);
 
     WriteBarrier<JSPromise> m_promise;
     WriteBarrier<Unknown> m_fulfillment;
 #if USE(BUN_JSC_ADDITIONS)
+    WriteBarrier<Unknown> m_importerAsyncContext;
     int64_t m_referrerAsyncOrder { -1 };
 #endif
     uint8_t m_remainingFulfillments { 2 };

Source/JavaScriptCore/runtime/ModuleLoadingContext.cpp

@@ -62,24 +62,26 @@ ModuleLoadingContext* ModuleLoadingContext::create(VM& vm, Step step, const JSMo
     return context;
 }
 
-ModuleLoadingContext::ModuleLoadingContext(VM& vm, Structure* structure, AbstractModuleRecord::ModuleRequest&& moduleRequest, RefPtr<ScriptFetcher> scriptFetcher, OptionSet<ModuleLoadFlag> flags, int64_t referrerAsyncOrder)
+ModuleLoadingContext::ModuleLoadingContext(VM& vm, Structure* structure, AbstractModuleRecord::ModuleRequest&& moduleRequest, RefPtr<ScriptFetcher> scriptFetcher, OptionSet<ModuleLoadFlag> flags, int64_t referrerAsyncOrder, JSValue importerAsyncContext)
     : Base(vm, structure)
     , m_moduleRequest(WTF::move(moduleRequest))
     , m_scriptFetcher(WTF::move(scriptFetcher))
 #if USE(BUN_JSC_ADDITIONS)
+    , m_importerAsyncContext(importerAsyncContext, WriteBarrierEarlyInit)
     , m_referrerAsyncOrder(referrerAsyncOrder)
 #endif
     , m_flags(flags)
 {
 #if !USE(BUN_JSC_ADDITIONS)
     UNUSED_PARAM(referrerAsyncOrder);
+    UNUSED_PARAM(importerAsyncContext);
 #endif
 }
 
-ModuleLoadingContext* ModuleLoadingContext::create(VM& vm, const AbstractModuleRecord::ModuleRequest& moduleRequest, RefPtr<ScriptFetcher> scriptFetcher, OptionSet<ModuleLoadFlag> flags, int64_t referrerAsyncOrder)
+ModuleLoadingContext* ModuleLoadingContext::create(VM& vm, const AbstractModuleRecord::ModuleRequest& moduleRequest, RefPtr<ScriptFetcher> scriptFetcher, OptionSet<ModuleLoadFlag> flags, int64_t referrerAsyncOrder, JSValue importerAsyncContext)
 {
     AbstractModuleRecord::ModuleRequest requestCopy { moduleRequest };
-    auto* context = new (NotNull, allocateCell<ModuleLoadingContext>(vm)) ModuleLoadingContext(vm, vm.moduleLoadingContextStructure.get(), WTF::move(requestCopy), WTF::move(scriptFetcher), flags, referrerAsyncOrder);
+    auto* context = new (NotNull, allocateCell<ModuleLoadingContext>(vm)) ModuleLoadingContext(vm, vm.moduleLoadingContextStructure.get(), WTF::move(requestCopy), WTF::move(scriptFetcher), flags, referrerAsyncOrder, importerAsyncContext);
     context->finishCreation(vm);
     return context;
 }
@@ -104,6 +106,9 @@ void ModuleLoadingContext::visitChildrenImpl(JSCell* cell, Visitor& visitor)
     visitor.append(thisObject->m_entry);
     visitor.append(thisObject->m_referrer);
     visitor.append(thisObject->m_module);
+#if USE(BUN_JSC_ADDITIONS)
+    visitor.append(thisObject->m_importerAsyncContext);
+#endif
 }
 
 DEFINE_VISIT_CHILDREN(ModuleLoadingContext);

Source/JavaScriptCore/runtime/ModuleLoadingContext.h

@@ -61,7 +61,7 @@ class ModuleLoadingContext final : public JSCell {
     };
 
     static ModuleLoadingContext* create(VM&, Step, const JSModuleLoader::ModuleReferrer&, const AbstractModuleRecord::ModuleRequest&, JSCell* payload, ModuleRegistryEntry*, RefPtr<ScriptFetcher>);
-    static ModuleLoadingContext* create(VM&, const AbstractModuleRecord::ModuleRequest&, RefPtr<ScriptFetcher>, OptionSet<ModuleLoadFlag>, int64_t referrerAsyncOrder = -1);
+    static ModuleLoadingContext* create(VM&, const AbstractModuleRecord::ModuleRequest&, RefPtr<ScriptFetcher>, OptionSet<ModuleLoadFlag>, int64_t referrerAsyncOrder = -1, JSValue importerAsyncContext = jsUndefined());
 
     Step step() const { return m_step; }
     void setStep(Step s) { m_step = s; }
@@ -80,11 +80,13 @@ class ModuleLoadingContext final : public JSCell {
     bool deferred() const { return m_flags.contains(ModuleLoadFlag::Deferred); }
 #if USE(BUN_JSC_ADDITIONS)
     int64_t referrerAsyncOrder() const { return m_referrerAsyncOrder; }
+    // The async context (AsyncLocalStorage store) active at the import() call site.
+    JSValue importerAsyncContext() const { return m_importerAsyncContext.get(); }
 #endif
 
 private:
     ModuleLoadingContext(VM&, Structure*, Step, const JSModuleLoader::ModuleReferrer&, AbstractModuleRecord::ModuleRequest&&, JSCell* payload, ModuleRegistryEntry*, RefPtr<ScriptFetcher>);
-    ModuleLoadingContext(VM&, Structure*, AbstractModuleRecord::ModuleRequest&&, RefPtr<ScriptFetcher>, OptionSet<ModuleLoadFlag>, int64_t referrerAsyncOrder);
+    ModuleLoadingContext(VM&, Structure*, AbstractModuleRecord::ModuleRequest&&, RefPtr<ScriptFetcher>, OptionSet<ModuleLoadFlag>, int64_t referrerAsyncOrder, JSValue importerAsyncContext);
 
     Step m_step { Step::Main };
     AbstractModuleRecord::ModuleRequest m_moduleRequest;
@@ -94,6 +96,7 @@ class ModuleLoadingContext final : public JSCell {
     WriteBarrier<Unknown> m_referrer;
     WriteBarrier<AbstractModuleRecord> m_module;
 #if USE(BUN_JSC_ADDITIONS)
+    WriteBarrier<Unknown> m_importerAsyncContext;
     int64_t m_referrerAsyncOrder { -1 };
 #endif
     OptionSet<ModuleLoadFlag> m_flags;