Merge pull request #823 from PassiveLogic/kr/job-queue-insert

JavaScriptEventLoop: Avoid O(N^2) job queue insert for bursts of jobs

Author: krodak

Sources/JavaScriptEventLoop/JobQueue.swift

@@ -1,6 +1,6 @@
 // This file contains the job queue implementation which re-order jobs based on their priority.
-// The current implementation is much simple to be easily debugged, but should be re-implemented
-// using priority queue ideally.
+// Jobs form runs, one per priority in descending order, FIFO within a run. The first job of each
+// run points to its last job, so an insert steps over whole runs.
 
 import _Concurrency
 import _CJavaScriptEventLoop
@@ -19,15 +19,18 @@ extension JavaScriptEventLoop {
     func insertJobQueue(job newJob: UnownedJob) {
         withUnsafeMutablePointer(to: &queueState.headJob) { headJobPtr in
             var position: UnsafeMutablePointer<UnownedJob?> = headJobPtr
-            while let cur = position.pointee {
-                if cur.rawPriority < newJob.rawPriority {
-                    newJob.nextInQueue().pointee = cur
-                    position.pointee = newJob
+            while let runHead = position.pointee, runHead.rawPriority >= newJob.rawPriority {
+                let runTail = runHead.lastInRun().pointee!
+                if runHead.rawPriority == newJob.rawPriority {
+                    newJob.nextInQueue().pointee = runTail.nextInQueue().pointee
+                    runTail.nextInQueue().pointee = newJob
+                    runHead.lastInRun().pointee = newJob
                     return
                 }
-                position = cur.nextInQueue()
+                position = runTail.nextInQueue()
             }
-            newJob.nextInQueue().pointee = nil
+            newJob.nextInQueue().pointee = position.pointee
+            newJob.lastInRun().pointee = newJob
             position.pointee = newJob
         }
 
@@ -81,7 +84,14 @@ extension JavaScriptEventLoop {
 
     func claimNextFromQueue() -> UnownedJob? {
         if let job = self.queueState.headJob {
-            self.queueState.headJob = job.nextInQueue().pointee
+            let nextJob = job.nextInQueue().pointee
+            if let nextJob {
+                let runTail = job.lastInRun().pointee!
+                if !runTail.isSameJob(as: job) {
+                    nextJob.lastInRun().pointee = runTail
+                }
+            }
+            self.queueState.headJob = nextJob
             return job
         }
         return nil
@@ -107,6 +117,16 @@ extension UnownedJob {
         }
     }
 
+    /// The last job of the run. Valid only on the first job of a run.
+    fileprivate func lastInRun() -> UnsafeMutablePointer<UnownedJob?> {
+        return withUnsafeMutablePointer(to: &asImpl().pointee.SchedulerPrivate.1) { rawLastJobPtr in
+            UnsafeMutableRawPointer(rawLastJobPtr).bindMemory(to: UnownedJob?.self, capacity: 1)
+        }
+    }
+
+    fileprivate func isSameJob(as other: UnownedJob) -> Bool {
+        asImpl() == other.asImpl()
+    }
 }
 
 private struct JobFlags {

Tests/JavaScriptEventLoopTests/JavaScriptEventLoopTests.swift

@@ -113,6 +113,113 @@ final class JavaScriptEventLoopTests: XCTestCase {
         XCTAssertEqual(context.completed, ["t4", "t3", "t5"])
     }
 
+    final class OrderContext: @unchecked Sendable {
+        var completed: [String] = []
+        var expectedCount = 0
+        var continuation: CheckedContinuation<Void, Never>?
+
+        func record(_ name: String) {
+            completed.append(name)
+            if completed.count == expectedCount {
+                continuation?.resume()
+                continuation = nil
+            }
+        }
+    }
+
+    /// Starts tasks without awaiting them, so no await raises the priority of a queued task.
+    func runUnawaited(expectedCount: Int, _ start: (OrderContext) -> Void) async -> [String] {
+        let context = OrderContext()
+        context.expectedCount = expectedCount
+        await withCheckedContinuation { continuation in
+            context.continuation = continuation
+            start(context)
+        }
+        return context.completed
+    }
+
+    func testTaskPriorityOrderWithMixedPriorities() async throws {
+        let burst: [(String, TaskPriority)] = [
+            ("low1", .low), ("medium1", .medium), ("low2", .low), ("high1", .high),
+            ("medium2", .medium), ("low3", .low), ("high2", .high), ("medium3", .medium),
+        ]
+        let burstOrder = await runUnawaited(expectedCount: burst.count) { context in
+            for (name, priority) in burst {
+                Task(priority: priority) { context.record(name) }
+            }
+        }
+        XCTAssertEqual(burstOrder, ["high1", "high2", "medium1", "medium2", "medium3", "low1", "low2", "low3"])
+
+        let nestedOrder = await runUnawaited(expectedCount: 7) { context in
+            Task(priority: .medium) {
+                context.record("outer")
+                let nested: [(String, TaskPriority)] = [
+                    ("nestedLow1", .low), ("nestedHigh", .high), ("nestedMedium", .medium), ("nestedLow2", .low),
+                ]
+                for (name, priority) in nested {
+                    Task(priority: priority) { context.record(name) }
+                }
+            }
+            Task(priority: .medium) { context.record("medium") }
+            Task(priority: .low) { context.record("low") }
+        }
+        XCTAssertEqual(
+            nestedOrder,
+            ["outer", "nestedHigh", "medium", "nestedMedium", "low", "nestedLow1", "nestedLow2"]
+        )
+
+        let priorities: [TaskPriority] = [
+            .low, .high, .background, .medium, TaskPriority(rawValue: 22), TaskPriority(rawValue: 5),
+        ]
+        let many = (0..<600).map { index in ("\(index)", priorities[index % priorities.count]) }
+        let manyOrder = await runUnawaited(expectedCount: many.count) { context in
+            for (name, priority) in many {
+                Task(priority: priority) { context.record(name) }
+            }
+        }
+        let expectedManyOrder = many.enumerated()
+            .sorted { lhs, rhs in
+                lhs.element.1.rawValue != rhs.element.1.rawValue
+                    ? lhs.element.1.rawValue > rhs.element.1.rawValue
+                    : lhs.offset < rhs.offset
+            }
+            .map(\.element.0)
+        XCTAssertEqual(manyOrder, expectedManyOrder)
+    }
+
+    func testTaskPriorityOrderAfterQueueEmpties() async throws {
+        let rounds: [[(String, TaskPriority)]] = [
+            [("low1", .low), ("low2", .low), ("high1", .high), ("low3", .low)],
+            [("high1", .high), ("medium1", .medium), ("high2", .high), ("low1", .low)],
+            [("medium1", .medium), ("medium2", .medium), ("medium3", .medium)],
+        ]
+        let expectedOrders = [
+            ["high1", "low1", "low2", "low3"],
+            ["high1", "high2", "medium1", "low1"],
+            ["medium1", "medium2", "medium3"],
+        ]
+        for (round, expectedOrder) in zip(rounds, expectedOrders) {
+            let order = await runUnawaited(expectedCount: round.count) { context in
+                for (name, priority) in round {
+                    Task(priority: priority) { context.record(name) }
+                }
+            }
+            XCTAssertEqual(order, expectedOrder)
+        }
+    }
+
+    func testTaskPriorityOrderAfterSingleJobRunIsClaimed() async throws {
+        let order = await runUnawaited(expectedCount: 4) { context in
+            Task(priority: .low) { context.record("low1") }
+            Task(priority: .low) { context.record("low2") }
+            Task(priority: .high) {
+                context.record("high")
+                Task(priority: .low) { context.record("low3") }
+            }
+        }
+        XCTAssertEqual(order, ["high", "low1", "low2", "low3"])
+    }
+
     // MARK: - Promise Tests
 
     func testPromiseResolution() async throws {

Tests/JavaScriptEventLoopTests/JavaScriptEventLoopYieldIntervalTests.swift

@@ -76,4 +76,31 @@ final class JavaScriptEventLoopYieldIntervalTests: XCTestCase {
         for task in tasks { await task.value }
         XCTAssertEqual(order.value, Array(0..<20))
     }
+
+    func testPriorityOrderAcrossYields() async {
+        JavaScriptEventLoop.shared.yieldInterval = .milliseconds(1)
+        let order = Box([String]())
+        let expectedOrder = ["h0", "h1", "h2", "h3", "h4", "m0", "m1", "m2", "m3", "l0", "l1", "l2", "l3", "l4"]
+        // Tasks are not awaited, so no await raises the priority of a queued task.
+        await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
+            @Sendable func record(_ name: String) {
+                Self.busyWait(milliseconds: 1)
+                order.value.append(name)
+                if order.value.count == expectedOrder.count { continuation.resume() }
+            }
+            let priorities: [(String, TaskPriority)] = [("l", .low), ("m", .medium), ("h", .high)]
+            for index in 0..<12 {
+                let (prefix, priority) = priorities[index % 3]
+                let name = "\(prefix)\(index / 3)"
+                Task(priority: priority) {
+                    record(name)
+                    if name == "h0" {
+                        Task(priority: .low) { record("l4") }
+                        Task(priority: .high) { record("h4") }
+                    }
+                }
+            }
+        }
+        XCTAssertEqual(order.value, expectedOrder)
+    }
 }