Carry the AsyncLocalStorage frame in promise reactions and microtask arguments instead of an InternalFieldTuple

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

@@ -0,0 +1,176 @@
+//@ requireOptions("--useDollarVM=1")
+// The embedder async-context slot ($vm.asyncContext) must be captured when a
+// promise reaction / async continuation is scheduled and be the current value
+// while it runs; a job that captured nothing runs with nothing, and whatever a
+// job installs does not leak into the next one.
+
+function shouldBe(actual, expected, msg) {
+    if (actual !== expected)
+        throw new Error((msg || "") + " expected " + String(expected) + " but got " + String(actual));
+}
+
+const get = () => $vm.asyncContext();
+const set = v => $vm.setAsyncContext(v);
+function inContext(ctx, fn) {
+    const prev = get();
+    set(ctx);
+    try {
+        return fn();
+    } finally {
+        set(prev);
+    }
+}
+
+const A = ["A"];
+const B = ["B"];
+let log = [];
+
+// .then on a pending promise, one and two handlers; .then on settled promises.
+{
+    let resolve;
+    const pending = new Promise(r => { resolve = r; });
+    const settled = Promise.resolve(1);
+    const rejected = Promise.reject(1);
+    inContext(A, () => {
+        pending.then(() => log.push(["pending.then", get()]));
+        pending.then(() => log.push(["pending.then#2", get()]), () => {});
+        settled.then(() => log.push(["settled.then", get()]));
+        rejected.then(() => {}, () => log.push(["rejected.then", get()]));
+        rejected.catch(() => log.push(["rejected.catch", get()]));
+        settled.finally(() => log.push(["settled.finally", get()]));
+    });
+    pending.then(() => log.push(["pending.then (no ctx)", get()]));
+    inContext(B, () => {
+        pending.then(() => { log.push(["pending.then B", get()]); set(A); /* residue */ });
+        pending.then(() => log.push(["pending.then B#2", get()]));
+    });
+    resolve(1);
+    drainMicrotasks();
+    shouldBe(get(), undefined, "residue after drain");
+    const expected = {
+        "pending.then": A, "pending.then#2": A, "settled.then": A, "rejected.then": A, "rejected.catch": A, "settled.finally": A,
+        "pending.then (no ctx)": undefined, "pending.then B": B, "pending.then B#2": B,
+    };
+    for (const [name, ctx] of log)
+        shouldBe(ctx, expected[name], name);
+    shouldBe(log.length, Object.keys(expected).length, "all handlers ran");
+}
+
+// await: pending native promise, settled promise, non-promise, thenable, and a
+// promise that already has reactions (spilled inline reaction).
+{
+    log = [];
+    let resolveP;
+    const p = new Promise(r => { resolveP = r; });
+    async function f(name, value) {
+        log.push([name + " before", get()]);
+        await value;
+        log.push([name + " after", get()]);
+        await null;
+        log.push([name + " after2", get()]);
+    }
+    const thenable = { then(r) { log.push(["thenable.then", get()]); r(1); } };
+    inContext(A, () => {
+        f("pendingA", p);
+        f("settledA", Promise.resolve(1));
+        f("valueA", 42);
+        f("thenableA", thenable);
+    });
+    inContext(B, () => {
+        f("pendingB", p); // second reaction on p: p's inline reaction spills to the list
+        f("pendingB2", p);
+    });
+    f("pendingNone", p);
+    // enterWith-style residue inside an async function persists across its own awaits only
+    (async () => {
+        set(B);
+        log.push(["self before", get()]);
+        await p;
+        log.push(["self after", get()]);
+        await 0;
+        log.push(["self after2", get()]);
+    })();
+    set(undefined);
+    resolveP(1);
+    drainMicrotasks();
+    shouldBe(get(), undefined, "residue after drain 2");
+    const want = name => {
+        const who = name.split(" ")[0];
+        if (who === "pendingNone")
+            return undefined;
+        if (who === "self" || who === "pendingB" || who === "pendingB2")
+            return B;
+        return A; // pendingA settledA valueA thenableA thenable.then
+    };
+    for (const [name, ctx] of log)
+        shouldBe(ctx, want(name), name);
+    shouldBe(log.filter(([n]) => n.endsWith("after2")).length, 8, "all continuations ran");
+}
+
+// async generators and for-await
+{
+    log = [];
+    async function* gen(name) {
+        log.push([name + " gen start", get()]);
+        yield 1;
+        log.push([name + " gen after yield", get()]);
+        await null;
+        yield 2;
+        log.push([name + " gen end", get()]);
+    }
+    async function consume(name) {
+        for await (const v of gen(name))
+            log.push([name + " body", get()]);
+        log.push([name + " done", get()]);
+    }
+    inContext(A, () => consume("A"));
+    consume("none");
+    drainMicrotasks();
+    let sawA = 0, sawNone = 0;
+    for (const [name, ctx] of log) {
+        if (name.startsWith("A ")) { shouldBe(ctx, A, name); sawA++; }
+        if (name.startsWith("none ")) { shouldBe(ctx, undefined, name); sawNone++; }
+    }
+    shouldBe(sawA, 6);
+    shouldBe(sawNone, 6);
+    shouldBe(get(), undefined, "residue after drain 3");
+}
+
+// Promise combinators keep the context of the awaiter.
+{
+    log = [];
+    let r1, r2;
+    const p1 = new Promise(r => { r1 = r; }), p2 = new Promise(r => { r2 = r; });
+    inContext(A, () => {
+        (async () => { await Promise.all([p1, p2]); log.push(["all", get()]); })();
+        (async () => { await Promise.race([p1, p2]); log.push(["race", get()]); })();
+        (async () => { await Promise.allSettled([p1, p2]); log.push(["allSettled", get()]); })();
+        (async () => { await Promise.any([p1, p2]); log.push(["any", get()]); })();
+    });
+    inContext(B, () => r1(1));
+    r2(2);
+    drainMicrotasks();
+    shouldBe(log.length, 4);
+    for (const [name, ctx] of log)
+        shouldBe(ctx, A, name);
+}
+
+// A handler that throws still restores the slot.
+{
+    inContext(A, () => {
+        Promise.resolve().then(() => { set(B); throw new Error("x"); }).catch(() => {});
+    });
+    drainMicrotasks();
+    shouldBe(get(), undefined, "restored after throwing handler");
+}
+
+// A job that captured no context and installs one itself does not leak it past its end.
+{
+    set(undefined);
+    Promise.resolve().then(() => { set(A); });
+    let seen = "unset";
+    Promise.resolve().then(() => { seen = get(); });
+    drainMicrotasks();
+    shouldBe(seen, undefined, "residue from a sibling microtask");
+    shouldBe(get(), undefined, "residue after a microtask that captured nothing");
+}

Source/JavaScriptCore/runtime/AbstractModuleRecord.cpp

@@ -41,6 +41,7 @@
 #include "JSAsyncFromSyncIterator.h"
 #include "JSAsyncFunctionGenerator.h"
 #include "JSPromiseCombinatorsGlobalContext.h"
+#include "JSPromiseReaction.h"
 #endif
 #include "ObjectConstructor.h"
 #include "SyntheticModuleRecord.h"
@@ -1283,6 +1284,14 @@ static bool importPromiseGatesAsyncDependency(JSPromise* importPromise, CyclicMo
                 followPromiseOrDriver(iterator->target());
             break;
         }
+        case InternalMicrotask::PromiseFinallyReactionJob:
+        case InternalMicrotask::PromiseFinallyAwaitJob: {
+            // The context record holds the promise that .finally() returned.
+            JSCell* contextCell = cellOf(context);
+            if (auto* record = contextCell ? dynamicDowncast<JSSlimPromiseReaction>(contextCell) : nullptr)
+                follow(record->promise());
+            break;
+        }
         case InternalMicrotask::PromiseAllResolveJob:
         case InternalMicrotask::PromiseAllSettledResolveJob: {
             JSCell* contextCell = cellOf(cell);
@@ -1297,8 +1306,6 @@ static bool importPromiseGatesAsyncDependency(JSPromise* importPromise, CyclicMo
         case InternalMicrotask::PromiseResolveThenableJobWithInternalMicrotask:
         case InternalMicrotask::PromiseResolveWithoutHandlerJob:
         case InternalMicrotask::PromiseFulfillWithoutHandlerJob:
-        case InternalMicrotask::PromiseFinallyReactionJob:
-        case InternalMicrotask::PromiseFinallyAwaitJob:
         case InternalMicrotask::PromiseReactionJob:
         case InternalMicrotask::ModuleLoadStep:
         case InternalMicrotask::ModuleLoadTopSettled:

Source/JavaScriptCore/runtime/AsyncContextSwapScope.h

@@ -35,27 +35,45 @@
 
 namespace JSC {
 
-// RAII helper for Bun's AsyncLocalStorage: swaps an async context value into
-// JSGlobalObject::m_asyncContextData field 0 for the lifetime of the scope and
-// restores the previous value on destruction. A no-op when the supplied context
-// is empty or undefined, so the common path (no async context active) costs a
-// single branch. Also provides helpers for the snapshot side (capturing the
-// current context and wrapping it into an InternalFieldTuple alongside a user
-// context) and for unwrapping such a tuple on the restore side.
+// RAII helper for Bun's AsyncLocalStorage. A job (microtask, timer, ...)
+// captures the value of JSGlobalObject::m_asyncContextData field 0 when it is
+// scheduled; constructing this scope with that captured value installs it for
+// the lifetime of the scope and restores the previous value on destruction.
+//
+// A job that captured "no context" (undefined, or an empty JSValue for callers
+// that never capture) runs with no context: whatever an earlier job left in the
+// slot via AsyncLocalStorage.enterWith() is not inherited, and whatever the job
+// itself leaves there does not outlive it. Until the VM has enabled tracking
+// (VM::isAsyncContextTrackingEnabled) nothing can have been captured, so every
+// entry point here reduces to that flag test.
 class AsyncContextSwapScope {
     WTF_MAKE_NONCOPYABLE(AsyncContextSwapScope);
     WTF_FORBID_HEAP_ALLOCATION;
 public:
     ALWAYS_INLINE AsyncContextSwapScope(VM& vm, JSGlobalObject* globalObject, JSValue asyncContext)
         : m_vm(vm)
     {
-        if (asyncContext.isEmpty() || asyncContext.isUndefined())
+        if (!vm.isAsyncContextTrackingEnabled())
             return;
-        m_asyncContextData = globalObject->m_asyncContextData.get();
-        if (!m_asyncContextData)
+        enter(globalObject, asyncContext);
+    }
+
+    // For internal microtasks: the captured context is in the dedicated argument
+    // when the fast paths filled it, otherwise in an InternalFieldTuple
+    // [context, asyncContext] that contextArg is unwrapped from (see wrap()).
+    ALWAYS_INLINE AsyncContextSwapScope(VM& vm, JSGlobalObject* globalObject, JSValue asyncContextArg, JSValue& contextArg)
+        : m_vm(vm)
+    {
+        if (!vm.isAsyncContextTrackingEnabled()) {
+            ASSERT(!isContextTuple(contextArg));
             return;
-        m_restoreAsyncContext = m_asyncContextData->getInternalField(0);
-        m_asyncContextData->putInternalField(vm, 0, asyncContext);
+        }
+        JSValue asyncContext = asyncContextArg;
+        if (asyncContext.isEmpty() || asyncContext.isUndefined()) {
+            if (isContextTuple(contextArg)) [[unlikely]]
+                asyncContext = unwrapContextTuple(contextArg);
+        }
+        enter(globalObject, asyncContext);
     }
 
     ALWAYS_INLINE ~AsyncContextSwapScope()
@@ -75,43 +93,68 @@ class AsyncContextSwapScope {
         }
     }
 
+    static ALWAYS_INLINE bool isContextTuple(JSValue contextArg)
+    {
+        // JSType test first: rejects the usual non-tuple cells (generators,
+        // iterators, module records) without the ClassInfo walk.
+        return !contextArg.isEmpty() && contextArg.isCell() && contextArg.asCell()->type() == InternalFieldTupleType && contextArg.asCell()->inherits<InternalFieldTuple>();
+    }
+
     // If contextArg is an InternalFieldTuple [userContext, asyncContext],
     // overwrite contextArg with field 0 and return field 1. Otherwise leave
     // contextArg untouched and return jsUndefined(). Empty contextArg is
-    // tolerated (dynamicDowncast<T>(JSValue) is not empty-safe on its own).
+    // tolerated.
     static ALWAYS_INLINE JSValue unwrapContextTuple(JSValue& contextArg)
     {
-        if (contextArg.isEmpty())
+        if (!isContextTuple(contextArg))
             return jsUndefined();
-        if (auto* tuple = dynamicDowncast<InternalFieldTuple>(contextArg)) {
-            contextArg = tuple->getInternalField(0);
-            return tuple->getInternalField(1);
-        }
-        return jsUndefined();
+        auto* tuple = uncheckedDowncast<InternalFieldTuple>(contextArg.asCell());
+        contextArg = tuple->getInternalField(0);
+        return tuple->getInternalField(1);
     }
 
-    // Read the current async context (field 0 of m_asyncContextData), or
-    // jsUndefined() when tracking has not been enabled on this global.
-    static ALWAYS_INLINE JSValue current(JSGlobalObject* globalObject)
+    // The async context to capture for a job being scheduled now: field 0 of
+    // m_asyncContextData, or jsUndefined() when there is none.
+    static ALWAYS_INLINE JSValue current(VM& vm, JSGlobalObject* globalObject)
     {
-        if (auto* asyncContextData = globalObject->m_asyncContextData.get())
-            return asyncContextData->getInternalField(0);
-        return jsUndefined();
+        if (!vm.isAsyncContextTrackingEnabled())
+            return jsUndefined();
+        ASSERT(globalObject->m_asyncContextData);
+        return globalObject->m_asyncContextData->getInternalField(0);
     }
 
-    // Snapshot the current async context alongside userContext in an
-    // InternalFieldTuple [userContext, asyncContext]. When no async context is
-    // active, returns userContext unchanged so the caller keeps using the
-    // allocation-free inline/slim reaction fast paths.
-    static ALWAYS_INLINE JSValue wrapWithCurrent(VM& vm, JSGlobalObject* globalObject, JSValue userContext)
+    // Pair userContext with asyncContext in an InternalFieldTuple
+    // [userContext, asyncContext] for the paths that only have one slot to
+    // carry both. When asyncContext is none, returns userContext unchanged.
+    static ALWAYS_INLINE JSValue wrap(VM& vm, JSGlobalObject* globalObject, JSValue userContext, JSValue asyncContext)
     {
-        JSValue asyncContext = current(globalObject);
-        if (asyncContext.isUndefined())
+        if (asyncContext.isEmpty() || asyncContext.isUndefined())
             return userContext;
+        ASSERT(vm.isAsyncContextTrackingEnabled());
         return InternalFieldTuple::create(vm, globalObject->internalFieldTupleStructure(), userContext, asyncContext);
     }
 
+    static ALWAYS_INLINE JSValue wrapWithCurrent(VM& vm, JSGlobalObject* globalObject, JSValue userContext)
+    {
+        return wrap(vm, globalObject, userContext, current(vm, globalObject));
+    }
+
 private:
+    // The previous value is put back on exit even when nothing had to be
+    // installed, so whatever the job itself leaves in the slot (enterWith())
+    // ends with the job.
+    ALWAYS_INLINE void enter(JSGlobalObject* globalObject, JSValue asyncContext)
+    {
+        ASSERT(m_vm.isAsyncContextTrackingEnabled());
+        ASSERT(globalObject->m_asyncContextData);
+        if (asyncContext.isEmpty())
+            asyncContext = jsUndefined();
+        m_asyncContextData = globalObject->m_asyncContextData.get();
+        m_restoreAsyncContext = m_asyncContextData->getInternalField(0);
+        if (m_restoreAsyncContext != asyncContext)
+            m_asyncContextData->putInternalField(m_vm, 0, asyncContext);
+    }
+
     VM& m_vm;
     InternalFieldTuple* m_asyncContextData { nullptr };
     JSValue m_restoreAsyncContext;

Source/JavaScriptCore/runtime/AsyncFromSyncIteratorPrototype.cpp

@@ -27,6 +27,8 @@
 #include "config.h"
 #include "AsyncFromSyncIteratorPrototype.h"
 
+#include "AsyncContextSwapScope.h"
+
 #include "IteratorOperations.h"
 #include "JSArrayInlines.h"
 #include "JSArrayIterator.h"
@@ -270,7 +272,11 @@ void driveAsyncFromSyncIteratorWithDriver(JSGlobalObject* globalObject, JSAsyncF
             JSValue error = catchScope.exception()->value();
             if (!catchScope.clearExceptionExceptTermination()) [[unlikely]]
                 return;
+#if USE(BUN_JSC_ADDITIONS)
+            JSPromise::rejectWithInternalMicrotask(vm, globalObject, error, InternalMicrotask::AsyncGeneratorDriverResume, driver, AsyncContextSwapScope::current(vm, globalObject));
+#else
             JSPromise::rejectWithInternalMicrotask(vm, globalObject, error, InternalMicrotask::AsyncGeneratorDriverResume, driver);
+#endif
             return;
         }
     }

Source/JavaScriptCore/runtime/JSGlobalObject.cpp

@@ -3830,11 +3830,6 @@ void JSGlobalObject::queueMicrotaskSlow(VM& vm, QueuedTask&& task)
 }
 
 #if USE(BUN_JSC_ADDITIONS)
-void JSGlobalObject::queueMicrotask(VM& vm, InternalMicrotask job, uint8_t payload, JSValue argument0, JSValue argument1, JSValue argument2, JSValue argument3)
-{
-    queueMicrotask(vm, QueuedTask { nullptr, job, payload, this, argument0, argument1, argument2, argument3 });
-}
-
 FFI::FFIContext& JSGlobalObject::ffiContext()
 {
     if (!m_ffiContext) [[unlikely]] {