JSC: async-context propagation for Bun's span context (one slot per suspended generator, PromiseReactionJobWithAsyncContext, AsyncFromSync reject paths)

Source/JavaScriptCore/runtime/AsyncContextSwapScope.h

@@ -28,6 +28,8 @@
 #if USE(BUN_JSC_ADDITIONS)
 
 #include "InternalFieldTuple.h"
+#include "JSAsyncFunctionGenerator.h"
+#include "JSAsyncGenerator.h"
 #include "JSCast.h"
 #include "JSGlobalObject.h"
 #include <wtf/ForbidHeapAllocation.h>
@@ -38,8 +40,8 @@ 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
+// is empty, or undefined while async context tracking has never been enabled on
+// the global, so the common path (no async context in use) costs a 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.
 class AsyncContextSwapScope {
@@ -49,13 +51,21 @@ class AsyncContextSwapScope {
     ALWAYS_INLINE AsyncContextSwapScope(VM& vm, JSGlobalObject* globalObject, JSValue asyncContext)
         : m_vm(vm)
     {
-        if (asyncContext.isEmpty() || asyncContext.isUndefined())
+        if (asyncContext.isEmpty())
+            return;
+        // Once anything uses async context, every continuation installs the
+        // context it captured -- including "none" -- so a context entered inside
+        // one continuation (AsyncLocalStorage.enterWith, an activated span) is
+        // scoped to that continuation instead of leaking into whichever job
+        // happens to run next. Until then this stays a single branch.
+        if (asyncContext.isUndefined() && !globalObject->isAsyncContextTrackingEnabled())
             return;
         m_asyncContextData = globalObject->m_asyncContextData.get();
         if (!m_asyncContextData)
             return;
         m_restoreAsyncContext = m_asyncContextData->getInternalField(0);
-        m_asyncContextData->putInternalField(vm, 0, asyncContext);
+        if (m_restoreAsyncContext != asyncContext)
+            m_asyncContextData->putInternalField(vm, 0, asyncContext);
     }
 
     ALWAYS_INLINE ~AsyncContextSwapScope()
@@ -111,6 +121,54 @@ class AsyncContextSwapScope {
         return InternalFieldTuple::create(vm, globalObject->internalFieldTupleStructure(), userContext, asyncContext);
     }
 
+    // Await-side capture. When the suspending object is an async (generator)
+    // function, record the current async context in its AsyncContext slot — a
+    // suspended async function has exactly one outstanding await, so this is a
+    // single barriered store instead of an InternalFieldTuple allocation per
+    // await — and return the object itself. Other drivers (AsyncFromSyncIterator,
+    // top-level-await module records) are rare and keep the tuple path.
+    static ALWAYS_INLINE JSValue captureForAwait(VM& vm, JSGlobalObject* globalObject, JSValue driver)
+    {
+        auto* asyncContextData = globalObject->m_asyncContextData.get();
+        if (!asyncContextData)
+            return driver;
+        JSValue asyncContext = asyncContextData->getInternalField(0);
+        if (driver.isCell()) {
+            JSCell* cell = driver.asCell();
+            JSType type = cell->type();
+            if (type == JSAsyncFunctionGeneratorType) {
+                auto* generator = uncheckedDowncast<JSAsyncFunctionGenerator>(cell);
+                if (generator->asyncContext() != asyncContext)
+                    generator->setAsyncContext(vm, asyncContext);
+                return driver;
+            }
+            if (type == JSAsyncGeneratorType) {
+                auto* generator = uncheckedDowncast<JSAsyncGenerator>(cell);
+                if (generator->asyncContext() != asyncContext)
+                    generator->setAsyncContext(vm, asyncContext);
+                return driver;
+            }
+        }
+        if (asyncContext.isUndefined())
+            return driver;
+        return InternalFieldTuple::create(vm, globalObject->internalFieldTupleStructure(), driver, asyncContext);
+    }
+
+    // Resume-side counterpart of captureForAwait(): yields the async context to
+    // restore and leaves contextArg pointing at the unwrapped driver.
+    static ALWAYS_INLINE JSValue contextForResume(JSValue& contextArg)
+    {
+        if (contextArg.isCell()) {
+            JSCell* cell = contextArg.asCell();
+            JSType type = cell->type();
+            if (type == JSAsyncFunctionGeneratorType)
+                return uncheckedDowncast<JSAsyncFunctionGenerator>(cell)->asyncContext();
+            if (type == JSAsyncGeneratorType)
+                return uncheckedDowncast<JSAsyncGenerator>(cell)->asyncContext();
+        }
+        return unwrapContextTuple(contextArg);
+    }
+
 private:
     VM& m_vm;
     InternalFieldTuple* m_asyncContextData { nullptr };

Source/JavaScriptCore/runtime/AsyncFromSyncIteratorPrototype.cpp

@@ -26,6 +26,9 @@
 
 #include "config.h"
 #include "AsyncFromSyncIteratorPrototype.h"
+#if USE(BUN_JSC_ADDITIONS)
+#include "AsyncContextSwapScope.h"
+#endif
 
 #include "IteratorOperations.h"
 #include "JSArrayInlines.h"
@@ -270,7 +273,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, AsyncContextSwapScope::captureForAwait(vm, globalObject, driver));
+#else
             JSPromise::rejectWithInternalMicrotask(vm, globalObject, error, InternalMicrotask::AsyncGeneratorDriverResume, driver);
+#endif
             return;
         }
     }

Source/JavaScriptCore/runtime/JSAsyncFunctionGenerator.h

@@ -30,9 +30,15 @@
 
 namespace JSC {
 
-class JSAsyncFunctionGenerator final : public JSInternalFieldObjectImpl<5> {
+#if USE(BUN_JSC_ADDITIONS)
+static constexpr unsigned jsAsyncFunctionGeneratorNumberOfInternalFields = 6;
+#else
+static constexpr unsigned jsAsyncFunctionGeneratorNumberOfInternalFields = 5;
+#endif
+
+class JSAsyncFunctionGenerator final : public JSInternalFieldObjectImpl<jsAsyncFunctionGeneratorNumberOfInternalFields> {
 public:
-    using Base = JSInternalFieldObjectImpl<5>;
+    using Base = JSInternalFieldObjectImpl<jsAsyncFunctionGeneratorNumberOfInternalFields>;
 
     template<typename CellType, SubspaceAccess mode>
     static GCClient::IsoSubspace* subspaceFor(VM& vm)
@@ -52,8 +58,14 @@ class JSAsyncFunctionGenerator final : public JSInternalFieldObjectImpl<5> {
         This,
         Frame,
         Context,
+#if USE(BUN_JSC_ADDITIONS)
+        // Bun async context (AsyncLocalStorage) captured at the most recent await and
+        // restored when the function resumes. A suspended async function has exactly one
+        // outstanding await, so one slot replaces a per-await InternalFieldTuple.
+        AsyncContext,
+#endif
     };
-    static_assert(numberOfInternalFields == 5);
+    static_assert(numberOfInternalFields == jsAsyncFunctionGeneratorNumberOfInternalFields);
     static_assert(static_cast<uint32_t>(Field::State) == static_cast<uint32_t>(JSGenerator::Field::State));
     static_assert(static_cast<uint32_t>(Field::Next) == static_cast<uint32_t>(JSGenerator::Field::Next));
     static_assert(static_cast<uint32_t>(Field::This) == static_cast<uint32_t>(JSGenerator::Field::This));
@@ -67,6 +79,9 @@ class JSAsyncFunctionGenerator final : public JSInternalFieldObjectImpl<5> {
             jsUndefined(),
             jsUndefined(),
             jsUndefined(),
+#if USE(BUN_JSC_ADDITIONS)
+            jsUndefined(),
+#endif
         } };
     }
 
@@ -108,6 +123,18 @@ class JSAsyncFunctionGenerator final : public JSInternalFieldObjectImpl<5> {
         return Base::internalField(static_cast<unsigned>(Field::Context)).get();
     }
 
+#if USE(BUN_JSC_ADDITIONS)
+    JSValue asyncContext() const
+    {
+        return Base::internalField(static_cast<unsigned>(Field::AsyncContext)).get();
+    }
+
+    void setAsyncContext(VM& vm, JSValue value)
+    {
+        Base::internalField(static_cast<unsigned>(Field::AsyncContext)).set(vm, this, value);
+    }
+#endif
+
     DECLARE_EXPORT_INFO;
 
     DECLARE_VISIT_CHILDREN;

Source/JavaScriptCore/runtime/JSAsyncGenerator.h

@@ -31,9 +31,15 @@
 
 namespace JSC {
 
-class JSAsyncGenerator final : public JSInternalFieldObjectImpl<10> {
+#if USE(BUN_JSC_ADDITIONS)
+static constexpr unsigned jsAsyncGeneratorNumberOfInternalFields = 11;
+#else
+static constexpr unsigned jsAsyncGeneratorNumberOfInternalFields = 10;
+#endif
+
+class JSAsyncGenerator final : public JSInternalFieldObjectImpl<jsAsyncGeneratorNumberOfInternalFields> {
 public:
-    using Base = JSInternalFieldObjectImpl<10>;
+    using Base = JSInternalFieldObjectImpl<jsAsyncGeneratorNumberOfInternalFields>;
 
     template<typename CellType, SubspaceAccess mode>
     static GCClient::IsoSubspace* subspaceFor(VM& vm)
@@ -86,8 +92,12 @@ class JSAsyncGenerator final : public JSInternalFieldObjectImpl<10> {
         ResumePromise,
         CachedDriverResult,
         CachedDriverResultTarget,
+#if USE(BUN_JSC_ADDITIONS)
+        // Bun async context captured at the most recent await; see JSAsyncFunctionGenerator.
+        AsyncContext,
+#endif
     };
-    static_assert(numberOfInternalFields == 10);
+    static_assert(numberOfInternalFields == jsAsyncGeneratorNumberOfInternalFields);
     static std::array<JSValue, numberOfInternalFields> initialValues()
     {
         return { {
@@ -101,6 +111,9 @@ class JSAsyncGenerator final : public JSInternalFieldObjectImpl<10> {
             jsUndefined(),
             jsUndefined(),
             jsUndefined(),
+#if USE(BUN_JSC_ADDITIONS)
+            jsUndefined(),
+#endif
         } };
     }
 
@@ -197,6 +210,18 @@ class JSAsyncGenerator final : public JSInternalFieldObjectImpl<10> {
         Base::internalField(static_cast<unsigned>(Field::CachedDriverResultTarget)).set(vm, this, value);
     }
 
+#if USE(BUN_JSC_ADDITIONS)
+    JSValue asyncContext() const
+    {
+        return Base::internalField(static_cast<unsigned>(Field::AsyncContext)).get();
+    }
+
+    void setAsyncContext(VM& vm, JSValue value)
+    {
+        Base::internalField(static_cast<unsigned>(Field::AsyncContext)).set(vm, this, value);
+    }
+#endif
+
     bool isQueueEmpty() const
     {
         return resumeMode() == static_cast<int32_t>(AsyncGeneratorResumeMode::Empty);

Source/JavaScriptCore/runtime/JSMicrotask.cpp

@@ -307,9 +307,9 @@ static ALWAYS_INLINE void settleDriverWithIteratorResult(JSGlobalObject* globalO
 {
     JSGlobalObject* realm = producer->realm();
 #if USE(BUN_JSC_ADDITIONS)
-    // Capture Bun's async context alongside the driver so AsyncGeneratorDriverResume can restore it.
+    // Capture Bun's async context on the driver so AsyncGeneratorDriverResume can restore it.
     // The unwrapped target is still used for the cached-result identity check below.
-    JSValue wrappedTarget = AsyncContextSwapScope::wrapWithCurrent(vm, globalObject, target);
+    JSValue wrappedTarget = AsyncContextSwapScope::captureForAwait(vm, globalObject, target);
 #else
     UNUSED_PARAM(globalObject);
     JSValue wrappedTarget = target;
@@ -372,8 +372,13 @@ static void asyncFromSyncIteratorContinueOrDone(JSGlobalObject* globalObject, VM
         scope.release();
         if (auto* promise = dynamicDowncast<JSPromise>(target))
             promise->reject(vm, result);
-        else
+        else {
+#if USE(BUN_JSC_ADDITIONS)
+            JSPromise::rejectWithInternalMicrotask(vm, globalObject, result, InternalMicrotask::AsyncGeneratorDriverResume, AsyncContextSwapScope::captureForAwait(vm, globalObject, target));
+#else
             JSPromise::rejectWithInternalMicrotask(vm, globalObject, result, InternalMicrotask::AsyncGeneratorDriverResume, target);
+#endif
+        }
         break;
     }
     case JSPromise::Status::Fulfilled: {
@@ -551,7 +556,7 @@ static void asyncGeneratorCompleteStep(JSGlobalObject* globalObject, JSAsyncGene
     // resolveWithInternalMicrotask keeps resolvePromise's thenable check, matching a real Promise settlement.
     if (isThrow) {
 #if USE(BUN_JSC_ADDITIONS)
-        JSValue wrappedTarget = AsyncContextSwapScope::wrapWithCurrent(vm, globalObject, target);
+        JSValue wrappedTarget = AsyncContextSwapScope::captureForAwait(vm, globalObject, target);
 #else
         JSValue wrappedTarget = target;
 #endif
@@ -1845,7 +1850,7 @@ void runInternalMicrotask(JSGlobalObject* globalObject, VM& vm, InternalMicrotas
         // context may be an InternalFieldTuple [userContext, asyncContext]; the resolving
         // functions keep the tuple as-is, so only peek at field 1 for the swap.
         JSValue peek = context;
-        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::unwrapContextTuple(peek));
+        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::contextForResume(peek));
 #endif
 
         auto [resolve, reject] = JSPromise::createResolvingFunctionsWithInternalMicrotask(vm, globalObject, task, context);
@@ -1907,16 +1912,24 @@ void runInternalMicrotask(JSGlobalObject* globalObject, VM& vm, InternalMicrotas
         RELEASE_AND_RETURN(scope, promiseAnyResolveJob(resultPromise->realm(), vm, globalContext, arguments[1], static_cast<uint64_t>(arguments[2].asAnyInt()), static_cast<JSPromise::Status>(payload)));
     }
 
+#if USE(BUN_JSC_ADDITIONS)
+    case InternalMicrotask::PromiseReactionJobWithAsyncContext:
+#endif
     case InternalMicrotask::PromiseReactionJob: {
         JSValue promiseOrCapability = arguments[0];
         JSValue handler = arguments[1];
 #if USE(BUN_JSC_ADDITIONS)
-        // arguments[3] is either an InternalFieldTuple [userContext, asyncContext]
-        // or userContext directly (legacy behavior). The scope stays active through
-        // resolvePromise/rejectPromise so thenables returned from the handler
-        // capture the correct async context, and restores on every return.
+        // arguments[3] is Bun's async context itself (PromiseReactionJobWithAsyncContext), an
+        // InternalFieldTuple [userContext, asyncContext] (performPromiseThenWithContext), or a
+        // bare userContext. The scope stays active through resolvePromise/rejectPromise so
+        // thenables returned from the handler capture the correct async context.
         JSValue userContext = arguments[3];
-        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::unwrapContextTuple(userContext));
+        JSValue asyncContext;
+        if (task == InternalMicrotask::PromiseReactionJobWithAsyncContext)
+            asyncContext = std::exchange(userContext, JSValue());
+        else
+            asyncContext = AsyncContextSwapScope::unwrapContextTuple(userContext);
+        AsyncContextSwapScope asyncContextScope(vm, globalObject, asyncContext);
 #endif
 
         JSValue result;
@@ -2001,8 +2014,7 @@ void runInternalMicrotask(JSGlobalObject* globalObject, VM& vm, InternalMicrotas
         JSValue contextArg = arguments[2];
 
 #if USE(BUN_JSC_ADDITIONS)
-        // contextArg may be an InternalFieldTuple [generator, asyncContext].
-        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::unwrapContextTuple(contextArg));
+        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::contextForResume(contextArg));
 #endif
         auto* generator = uncheckedDowncast<JSAsyncFunctionGenerator>(contextArg);
         JSGlobalObject* generatorGlobalObject = generator->realm();
@@ -2026,7 +2038,7 @@ void runInternalMicrotask(JSGlobalObject* globalObject, VM& vm, InternalMicrotas
     case InternalMicrotask::AsyncGeneratorYieldAwaited: {
         JSValue contextArg = arguments[2];
 #if USE(BUN_JSC_ADDITIONS)
-        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::unwrapContextTuple(contextArg));
+        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::contextForResume(contextArg));
 #endif
         auto* generator = uncheckedDowncast<JSAsyncGenerator>(contextArg);
         scope.release();
@@ -2037,7 +2049,7 @@ void runInternalMicrotask(JSGlobalObject* globalObject, VM& vm, InternalMicrotas
     case InternalMicrotask::AsyncGeneratorBodyCallNormal: {
         JSValue contextArg = arguments[2];
 #if USE(BUN_JSC_ADDITIONS)
-        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::unwrapContextTuple(contextArg));
+        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::contextForResume(contextArg));
 #endif
         auto* generator = uncheckedDowncast<JSAsyncGenerator>(contextArg);
         scope.release();
@@ -2048,7 +2060,7 @@ void runInternalMicrotask(JSGlobalObject* globalObject, VM& vm, InternalMicrotas
     case InternalMicrotask::AsyncGeneratorBodyCallReturn: {
         JSValue contextArg = arguments[2];
 #if USE(BUN_JSC_ADDITIONS)
-        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::unwrapContextTuple(contextArg));
+        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::contextForResume(contextArg));
 #endif
         auto* generator = uncheckedDowncast<JSAsyncGenerator>(contextArg);
         scope.release();
@@ -2059,7 +2071,7 @@ void runInternalMicrotask(JSGlobalObject* globalObject, VM& vm, InternalMicrotas
     case InternalMicrotask::AsyncGeneratorAwaitReturn: {
         JSValue contextArg = arguments[2];
 #if USE(BUN_JSC_ADDITIONS)
-        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::unwrapContextTuple(contextArg));
+        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::contextForResume(contextArg));
 #endif
         auto* generator = uncheckedDowncast<JSAsyncGenerator>(contextArg);
         scope.release();
@@ -2070,7 +2082,7 @@ void runInternalMicrotask(JSGlobalObject* globalObject, VM& vm, InternalMicrotas
     case InternalMicrotask::AsyncGeneratorDriverResume: {
         JSValue contextArg = arguments[2];
 #if USE(BUN_JSC_ADDITIONS)
-        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::unwrapContextTuple(contextArg));
+        AsyncContextSwapScope asyncContextScope(vm, globalObject, AsyncContextSwapScope::contextForResume(contextArg));
 #endif
         scope.release();
         asyncGeneratorDriverResume(vm, contextArg, arguments[1], static_cast<JSPromise::Status>(payload), microtaskCallCache);