[JSC] Microtask drain scopes

Source/JavaScriptCore/runtime/Microtask.h

@@ -111,6 +111,16 @@ constexpr bool isModuleLoaderInternalMicrotask(InternalMicrotask task)
     return static_cast<uint8_t>(task) >= static_cast<uint8_t>(InternalMicrotask::AsyncModuleExecutionResume)
         && static_cast<uint8_t>(task) <= static_cast<uint8_t>(InternalMicrotask::ImportModuleNamespace);
 }
+
+// The module-loader pipeline steps a MicrotaskQueue::DrainScope keeps in the queue
+// when it admits loader jobs. AsyncModuleExecutionResume resumes user module code and
+// PromiseFulfillWithoutHandlerJob is a plain settlement; both wait like anything else.
+constexpr bool isDrainScopeLoaderJob(InternalMicrotask task)
+{
+    return task != InternalMicrotask::AsyncModuleExecutionResume
+        && task != InternalMicrotask::PromiseFulfillWithoutHandlerJob
+        && isModuleLoaderInternalMicrotask(task);
+}
 #else
 constexpr unsigned maxMicrotaskArguments = 3;
 #endif

Source/JavaScriptCore/runtime/MicrotaskQueue.cpp

@@ -122,9 +122,43 @@ void MicrotaskQueue::visitAggregateImpl(Visitor& visitor)
 {
     m_queue.visitAggregate(visitor);
     m_toKeep.visitAggregate(visitor);
+#if USE(BUN_JSC_ADDITIONS)
+    for (auto& scope : m_drainScopes)
+        scope->deferred.visitAggregate(visitor);
+#endif
 }
 DEFINE_VISIT_AGGREGATE(MicrotaskQueue);
 
+#if USE(BUN_JSC_ADDITIONS)
+void MicrotaskQueue::beginDrainScope(bool admitLoaderJobs)
+{
+    auto scope = makeUnique<DrainScope>();
+    // Everything queued so far predates the scope. (Order within `deferred` and within
+    // what stays is preserved; kept loader jobs simply run earlier than they would have.)
+    MarkedMicrotaskDeque kept;
+    while (!m_queue.isEmpty()) {
+        QueuedTask task = m_queue.dequeue();
+        if (admitLoaderJobs && isDrainScopeLoaderJob(task.job()))
+            kept.enqueue(WTF::move(task));
+        else
+            scope->deferred.enqueue(WTF::move(task));
+    }
+    m_queue.swap(kept);
+    m_drainScopes.append(WTF::move(scope));
+}
+
+void MicrotaskQueue::endDrainScope()
+{
+    ASSERT(hasOpenDrainScope());
+    // Prepend while the scope is still on the stack, so the deferred tasks are never
+    // unreachable from visitAggregate.
+    auto& deferred = m_drainScopes.last()->deferred;
+    while (!deferred.isEmpty())
+        m_queue.prepend(deferred.takeLast());
+    m_drainScopes.removeLast();
+}
+#endif
+
 void MicrotaskQueue::enqueueSlow(QueuedTask&& task)
 {
     auto* globalObject = task.globalObject();

Source/JavaScriptCore/runtime/MicrotaskQueue.h

@@ -35,6 +35,7 @@
 #include <wtf/RefCounted.h>
 #include <wtf/SentinelLinkedList.h>
 #include <wtf/TZoneMalloc.h>
+#include <wtf/Vector.h>
 #include <wtf/VectorTraits.h>
 
 namespace JSC {
@@ -190,6 +191,22 @@ class MarkedMicrotaskDeque {
 #if USE(BUN_JSC_ADDITIONS)
     // Defined in MicrotaskQueueInlines.h (requires globalObject() from QueuedTask).
     inline void clearForGlobalObject(JSGlobalObject* targetGlobalObject);
+
+    // Re-insert at the front. Everything already in the deque shifts by one, so the
+    // "already marked" prefix can no longer be trusted; rescan from the start.
+    void prepend(QueuedTask&& task)
+    {
+        m_queue.prepend(WTF::move(task));
+        m_markedBefore = 0;
+    }
+
+    QueuedTask takeLast()
+    {
+        QueuedTask task = m_queue.takeLast();
+        if (m_markedBefore > m_queue.size())
+            m_markedBefore = m_queue.size();
+        return task;
+    }
 #endif
 
     void beginMarking()
@@ -233,17 +250,49 @@ class MicrotaskQueue : public BasicRawSentinelNode<MicrotaskQueue>, public RefCo
     {
         m_queue.clear();
         m_toKeep.clear();
+#if USE(BUN_JSC_ADDITIONS)
+        for (auto& scope : m_drainScopes)
+            scope->deferred.clear();
+#endif
     }
 
 #if USE(BUN_JSC_ADDITIONS)
     // Defined in MicrotaskQueueInlines.h (requires MarkedMicrotaskDeque::clearForGlobalObject).
     inline void clearForGlobalObject(JSGlobalObject* targetGlobalObject);
+
+    // Drain scopes.
+    //
+    // The embedder can turn its event loop from inside a synchronous frame while
+    // admitting only work that frame causes. beginDrainScope() sets aside every task
+    // already queued — they belong to the code the frame interrupted — so that until
+    // endDrainScope() a checkpoint runs only what has been queued since (by the scope's
+    // own code, transitively). endDrainScope() puts the set-aside tasks back at the
+    // front of the queue in their original order. Scopes nest: an inner scope sets
+    // aside the outer scope's pending tasks the same way. While any scope is open,
+    // VM::drainMicrotasks() is not the end of the outer frame's job, so it skips
+    // unhandled-rejection notification and WeakRef finalization.
+    //
+    // `admitLoaderJobs`: module-loader pipeline jobs already queued are keyed by loader
+    // state the scope shares with the outer program (a scope awaiting an `import()` of
+    // a module whose fetch has already settled depends on them), so a scope that may
+    // import keeps them in the queue. See isDrainScopeLoaderJob.
+    struct DrainScope {
+        WTF_DEPRECATED_MAKE_STRUCT_FAST_ALLOCATED(DrainScope);
+        MarkedMicrotaskDeque deferred;
+    };
+    bool hasOpenDrainScope() const { return !m_drainScopes.isEmpty(); }
+    JS_EXPORT_PRIVATE void beginDrainScope(bool admitLoaderJobs);
+    JS_EXPORT_PRIVATE void endDrainScope();
 #endif
 
     void beginMarking()
     {
         m_queue.beginMarking();
         m_toKeep.beginMarking();
+#if USE(BUN_JSC_ADDITIONS)
+        for (auto& scope : m_drainScopes)
+            scope->deferred.beginMarking();
+#endif
     }
 
     DECLARE_VISIT_AGGREGATE;
@@ -288,6 +337,9 @@ class MicrotaskQueue : public BasicRawSentinelNode<MicrotaskQueue>, public RefCo
 
     MarkedMicrotaskDeque m_queue;
     MarkedMicrotaskDeque m_toKeep;
+#if USE(BUN_JSC_ADDITIONS)
+    Vector<std::unique_ptr<DrainScope>, 2> m_drainScopes; // innermost last
+#endif
 };
 
 JS_EXPORT_PRIVATE void runMicrotaskWithDebugger(JSGlobalObject*, VM&, QueuedTask&);

Source/JavaScriptCore/runtime/MicrotaskQueueInlines.h

@@ -95,6 +95,8 @@ inline void MicrotaskQueue::clearForGlobalObject(JSGlobalObject* targetGlobalObj
         return;
     m_queue.clearForGlobalObject(targetGlobalObject);
     m_toKeep.clearForGlobalObject(targetGlobalObject);
+    for (auto& scope : m_drainScopes)
+        scope->deferred.clearForGlobalObject(targetGlobalObject);
 }
 #endif
 

Source/JavaScriptCore/runtime/VM.cpp

@@ -1546,6 +1546,21 @@ void VM::drainMicrotasks()
 
     if (executionForbidden()) [[unlikely]]
         m_defaultMicrotaskQueue->clear();
+#if USE(BUN_JSC_ADDITIONS)
+    else if (m_defaultMicrotaskQueue->hasOpenDrainScope()) [[unlikely]] {
+        // See MicrotaskQueue::DrainScope: run what the scope has queued, but this is not
+        // the end of the outer frame's synchronous execution.
+        std::optional<VMEntryScope> entryScope;
+        if (!m_defaultMicrotaskQueue->isEmpty())
+            entryScope.emplace(*this, nullptr);
+        m_defaultMicrotaskQueue->performMicrotaskCheckpoint</* useCallOnEachMicrotask */ true>(*this,
+            [&](JSGlobalObject*, JSGlobalObject* nextGlobalObject) {
+                if (entryScope && nextGlobalObject)
+                    entryScope->setGlobalObject(nextGlobalObject);
+            });
+        return;
+    }
+#endif
     else {
         std::optional<VMEntryScope> entryScope;
         if (!m_defaultMicrotaskQueue->isEmpty())