Fix some more quadratic performance issues in AsyncQueue

We would hit quadratic behavior in `AsyncQueue` when the build system floods us with `build/logMessage` or `build/task(Start|Progress|Finish)` notifications because we record a dependency on all of the pending log message handling tasks.

We can extend the improvement made in #1840 to fix this quadratic problem: If the current task depends on a task with metadata that depends on itself (ie. all tasks of metadata that needs to be executed in-order), we only need to depend on the last task with that metadata.

This covers many message types and we can now only get into quadratic behavior if we get flooded with two different kinds of messages: One that does not have a self-dependency and one that depends on the message without a self-dependency. In terms of LSP messages, this could be a document read followed by a document update, but I think this is a lot more unlikely than getting spammed with one type of message.

Author: ahoppen

Sources/BuildSystemIntegration/BuildSystemMessageDependencyTracker.swift

@@ -74,10 +74,6 @@ package enum BuildSystemMessageDependencyTracker: QueueBasedMessageHandlerDepend
     }
   }
 
-  package func dependencies(in pendingTasks: [PendingTask<Self>]) -> [PendingTask<Self>] {
-    return pendingTasks.filter { $0.metadata.isDependency(of: self) }
-  }
-
   package init(_ request: some RequestType) {
     switch request {
     case is BuildShutdownRequest:

Sources/SourceKitLSP/MessageHandlingDependencyTracker.swift

@@ -13,7 +13,7 @@
 package import LanguageServerProtocol
 import LanguageServerProtocolExtensions
 import SKLogging
-package import SwiftExtensions
+import SwiftExtensions
 
 /// A lightweight way of describing tasks that are created from handling LSP
 /// requests or notifications for the purpose of dependency tracking.
@@ -83,10 +83,6 @@ package enum MessageHandlingDependencyTracker: QueueBasedMessageHandlerDependenc
     }
   }
 
-  package func dependencies(in pendingTasks: [PendingTask<Self>]) -> [PendingTask<Self>] {
-    return pendingTasks.filter { $0.metadata.isDependency(of: self) }
-  }
-
   package init(_ notification: some NotificationType) {
     switch notification {
     case is CancelRequestNotification:

Sources/SourceKitLSP/Swift/SyntacticTestIndex.swift

@@ -61,10 +61,6 @@ fileprivate enum TaskMetadata: DependencyTracker, Equatable {
       return !lhsUris.intersection(rhsUris).isEmpty
     }
   }
-
-  package func dependencies(in pendingTasks: [PendingTask<Self>]) -> [PendingTask<Self>] {
-    return pendingTasks.filter { $0.metadata.isDependency(of: self) }
-  }
 }
 
 /// Data from a syntactic scan of a source file for tests.

Sources/SwiftExtensions/AsyncQueue.swift

@@ -27,29 +27,23 @@ extension Task: AnyTask {
 }
 
 /// A type that is able to track dependencies between tasks.
-package protocol DependencyTracker: Sendable {
-  /// Which tasks need to finish before a task described by `self` may start executing.
-  /// `pendingTasks` is sorted in the order in which the tasks were enqueued to `AsyncQueue`.
-  func dependencies(in pendingTasks: [PendingTask<Self>]) -> [PendingTask<Self>]
+package protocol DependencyTracker: Sendable, Hashable {
+  /// Whether the task described by `self` needs to finish executing before `other` can start executing.
+  func isDependency(of other: Self) -> Bool
 }
 
 /// A dependency tracker where each task depends on every other, i.e. a serial
 /// queue.
 package struct Serial: DependencyTracker {
-  package func dependencies(in pendingTasks: [PendingTask<Self>]) -> [PendingTask<Self>] {
-    if let lastTask = pendingTasks.last {
-      return [lastTask]
-    }
-    return []
+  package func isDependency(of other: Serial) -> Bool {
+    return true
   }
 }
 
-package struct PendingTask<TaskMetadata: Sendable>: Sendable {
+package struct PendingTask<TaskMetadata: Sendable & Hashable>: Sendable {
   /// The task that is pending.
   fileprivate let task: any AnyTask
 
-  package let metadata: TaskMetadata
-
   /// A unique value used to identify the task. This allows tasks to get
   /// removed from `pendingTasks` again after they finished executing.
   fileprivate let id: UUID
@@ -58,23 +52,25 @@ package struct PendingTask<TaskMetadata: Sendable>: Sendable {
 /// A list of pending tasks that can be sent across actor boundaries and is guarded by a lock.
 ///
 /// - Note: Unchecked sendable because the tasks are being protected by a lock.
-private class PendingTasks<TaskMetadata: Sendable>: @unchecked Sendable {
+private final class PendingTasks<TaskMetadata: Sendable & Hashable>: Sendable {
   ///  Lock guarding `pendingTasks`.
   private let lock = NSLock()
 
   /// Pending tasks that have not finished execution yet.
   ///
   /// - Important: This must only be accessed while `lock` has been acquired.
-  private var tasks: [PendingTask<TaskMetadata>] = []
+  private nonisolated(unsafe) var tasksByMetadata: [TaskMetadata: [PendingTask<TaskMetadata>]] = [:]
 
   init() {
     self.lock.name = "AsyncQueue"
   }
 
   /// Capture a lock and execute the closure, which may modify the pending tasks.
-  func withLock<T>(_ body: (_ pendingTasks: inout [PendingTask<TaskMetadata>]) throws -> T) rethrows -> T {
+  func withLock<T>(
+    _ body: (_ tasksByMetadata: inout [TaskMetadata: [PendingTask<TaskMetadata>]]) throws -> T
+  ) rethrows -> T {
     try lock.withLock {
-      try body(&tasks)
+      try body(&tasksByMetadata)
     }
   }
 }
@@ -122,10 +118,29 @@ package final class AsyncQueue<TaskMetadata: DependencyTracker>: Sendable {
   ) -> Task<Success, any Error> {
     let id = UUID()
 
-    return pendingTasks.withLock { tasks in
+    return pendingTasks.withLock { tasksByMetadata in
       // Build the list of tasks that need to finished execution before this one
       // can be executed
-      let dependencies = metadata.dependencies(in: tasks)
+      var dependencies: [PendingTask<TaskMetadata>] = []
+      for (pendingMetadata, pendingTasks) in tasksByMetadata {
+        guard pendingMetadata.isDependency(of: metadata) else {
+          // No dependency
+          continue
+        }
+        if metadata.isDependency(of: metadata), let lastPendingTask = pendingTasks.last {
+          // This kind of task depends on all other tasks of the same kind finishing. It is sufficient to just wait on
+          // the last task with this metadata, it will have all the other tasks with the same metadata as transitive
+          // dependencies.
+          dependencies.append(lastPendingTask)
+        } else {
+          // We depend on tasks with this metadata, but they don't have any dependencies between them, eg.
+          // `documentUpdate` depends on all `documentRequest` but `documentRequest` don't have dependencies between
+          // them. We need to depend on all of them unless we knew that we depended on some other task that already
+          // depends on all of these. But determining that would also require knowledge about the entire dependency
+          // graph, which is likely as expensive as depending on all of these tasks.
+          dependencies += pendingTasks
+        }
+      }
 
       // Schedule the task.
       let task = Task(priority: priority) { [pendingTasks] in
@@ -139,14 +154,17 @@ package final class AsyncQueue<TaskMetadata: DependencyTracker>: Sendable {
 
         let result = try await operation()
 
-        pendingTasks.withLock { tasks in
-          tasks.removeAll(where: { $0.id == id })
+        pendingTasks.withLock { tasksByMetadata in
+          tasksByMetadata[metadata, default: []].removeAll(where: { $0.id == id })
+          if tasksByMetadata[metadata]?.isEmpty ?? false {
+            tasksByMetadata[metadata] = nil
+          }
         }
 
         return result
       }
 
-      tasks.append(PendingTask(task: task, metadata: metadata, id: id))
+      tasksByMetadata[metadata, default: []].append(PendingTask(task: task, id: id))
 
       return task
     }