Add a withTaskPriorityChangedHandler function and use it in TaskScheduler

We need this function anyway to escalate process priorities when we set `nice`ness values for them. It also simplifies the task scheduler and I’m hoping that it fixes a non-deterministic failure that causes task priority elevation to not work properly.

Author: ahoppen

Sources/SKCore/TaskScheduler.swift

@@ -13,6 +13,7 @@
 import CAtomics
 import Foundation
 import LSPLogging
+import SKSupport
 
 /// See comment on ``TaskDescriptionProtocol/dependencies(to:taskPriority:)``
 public enum TaskDependencyAction<TaskDescription: TaskDescriptionProtocol> {
@@ -125,10 +126,6 @@ public actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
   /// Every time `execute` gets called, a new task is placed in this continuation. See comment on `executionTask`.
   private let executionTaskCreatedContinuation: AsyncStream<Task<ExecutionTaskFinishStatus, Never>>.Continuation
 
-  /// Placing a new value in this continuation will cause `resultTask` to query its priority and set
-  /// `QueuedTask.priority`.
-  private let updatePriorityContinuation: AsyncStream<Void>.Continuation
-
   nonisolated(unsafe) private var _priority: AtomicUInt8
 
   /// The latest known priority of the task.
@@ -189,16 +186,10 @@ public actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
     description: TaskDescription,
     executionStateChangedCallback: (@Sendable (QueuedTask, TaskExecutionState) async -> Void)?
   ) async {
-    self._priority = .init(initialValue: priority?.rawValue ?? Task.currentPriority.rawValue)
+    self._priority = AtomicUInt8(initialValue: priority?.rawValue ?? Task.currentPriority.rawValue)
     self.description = description
     self.executionStateChangedCallback = executionStateChangedCallback
 
-    var updatePriorityContinuation: AsyncStream<Void>.Continuation!
-    let updatePriorityStream = AsyncStream {
-      updatePriorityContinuation = $0
-    }
-    self.updatePriorityContinuation = updatePriorityContinuation
-
     var executionTaskCreatedContinuation: AsyncStream<Task<ExecutionTaskFinishStatus, Never>>.Continuation!
     let executionTaskCreatedStream = AsyncStream {
       executionTaskCreatedContinuation = $0
@@ -207,38 +198,24 @@ public actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
 
     self.resultTask = Task.detached(priority: priority) {
       await withTaskCancellationHandler {
-        await withTaskGroup(of: Void.self) { taskGroup in
-          taskGroup.addTask {
-            for await _ in updatePriorityStream {
-              if Task.currentPriority != self.priority {
-                withLoggingSubsystemAndScope(subsystem: taskSchedulerSubsystem, scope: nil) {
-                  logger.debug(
-                    "Updating priority of \(self.description.forLogging) from \(self.priority.rawValue) to \(Task.currentPriority.rawValue)"
-                  )
-                }
-                self.priority = Task.currentPriority
-              }
+        await withTaskPriorityChangedHandler(initialPriority: self.priority) {
+          for await task in executionTaskCreatedStream {
+            switch await task.valuePropagatingCancellation {
+            case .cancelledToBeRescheduled:
+              // Break the switch and wait for a new `executionTask` to be placed into `executionTaskCreatedStream`.
+              break
+            case .terminated:
+              // The task finished. We are done with this `QueuedTask`
+              return
             }
           }
-          taskGroup.addTask {
-            for await task in executionTaskCreatedStream {
-              switch await task.valuePropagatingCancellation {
-              case .cancelledToBeRescheduled:
-                // Break the switch and wait for a new `executionTask` to be placed into `executionTaskCreatedStream`.
-                break
-              case .terminated:
-                // The task finished. We are done with this `QueuedTask`
-                return
-              }
-            }
-          }
-          // The first (update priority) task never finishes, so this waits for the second (wait for execution) task
-          // to terminate.
-          // Afterwards we also cancel the update priority task.
-          for await _ in taskGroup {
-            taskGroup.cancelAll()
-            return
+        } taskPriorityChanged: {
+          withLoggingSubsystemAndScope(subsystem: taskSchedulerSubsystem, scope: nil) {
+            logger.debug(
+              "Updating priority of \(self.description.forLogging) from \(self.priority.rawValue) to \(Task.currentPriority.rawValue)"
+            )
           }
+          self.priority = Task.currentPriority
         }
       } onCancel: {
         self.resultTaskCancelled.value = true
@@ -291,16 +268,6 @@ public actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
     self.executionTask = nil
   }
 
-  /// Trigger `QueuedTask.priority` to be updated with the current priority of the underlying task.
-  ///
-  /// This is an asynchronous operation that makes no guarantees when the updated priority will be available.
-  ///
-  /// This is needed because tasks can't subscribe to priority updates (ie. there is no `withPriorityHandler` similar to
-  /// `withCancellationHandler`, https://github.com/apple/swift/issues/73367).
-  func triggerPriorityUpdate() {
-    updatePriorityContinuation.yield()
-  }
-
   /// If the priority of this task is less than `targetPriority`, elevate the priority to `targetPriority` by spawning
   /// a new task that depends on it. Otherwise a no-op.
   nonisolated func elevatePriority(to targetPriority: TaskPriority) {
@@ -399,16 +366,6 @@ public actor TaskScheduler<TaskDescription: TaskDescriptionProtocol> {
     return queuedTask
   }
 
-  /// Trigger all queued tasks to update their priority.
-  ///
-  /// Should be called occasionally to elevate tasks in the queue whose underlying `Swift.Task` had their priority
-  /// elevated because a higher-priority task started depending on them.
-  private func triggerPriorityUpdateOfQueuedTasks() async {
-    for task in pendingTasks {
-      await task.triggerPriorityUpdate()
-    }
-  }
-
   /// Returns the maximum number of concurrent tasks that are allowed to execute at the given priority.
   private func maxConcurrentTasks(at priority: TaskPriority) -> Int {
     for (atPriority, maxConcurrentTasks) in maxConcurrentTasksByPriority {
@@ -431,9 +388,8 @@ public actor TaskScheduler<TaskDescription: TaskDescriptionProtocol> {
       {
         // We don't have any execution slots left. Thus, this poker has nothing to do and is done.
         // When the next task finishes, it calls `poke` again.
-        // If the low priority task's priority gets elevated, that will be picked up when the next task in the
-        // `TaskScheduler` finishes, which causes  `triggerPriorityUpdateOfQueuedTasks` to be called, which transfers
-        // the new elevated priority to `QueuedTask.priority` and which can then be picked up by the next `poke` call.
+        // If the low priority task's priority gets elevated that task's priority will get elevated and it will be
+        // picked up on the next `poke` call.
         return
       }
       let dependencies = task.description.dependencies(to: currentlyExecutingTasks.map(\.description))
@@ -535,7 +491,6 @@ public actor TaskScheduler<TaskDescription: TaskDescriptionProtocol> {
     case .terminated: break
     case .cancelledToBeRescheduled: pendingTasks.append(task)
     }
-    await self.triggerPriorityUpdateOfQueuedTasks()
     self.poke()
   }
 }

Sources/SKSupport/CMakeLists.txt

@@ -17,6 +17,7 @@ add_library(SKSupport STATIC
   Result.swift
   Sequence+AsyncMap.swift
   SwitchableProcessResultExitStatus.swift
+  Task+WithPriorityChangedHandler.swift
   ThreadSafeBox.swift
   WorkspaceType.swift
 )

Sources/SKSupport/Task+WithPriorityChangedHandler.swift

@@ -0,0 +1,60 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+/// Runs `operation`. If the task's priority changes while the operation is running, calls `taskPriorityChanged`.
+///
+/// Since Swift Concurrency doesn't support direct observation of a task's priority, this polls the task's priority at
+/// `pollingInterval`.
+/// The function assumes that the original priority of the task is `initialPriority`. If the task priority changed
+/// compared to `initialPriority`, the `taskPriorityChanged` will be called.
+public func withTaskPriorityChangedHandler(
+  initialPriority: TaskPriority = Task.currentPriority,
+  pollingInterval: Duration = .seconds(0.1),
+  @_inheritActorContext operation: @escaping @Sendable () async -> Void,
+  taskPriorityChanged: @escaping @Sendable () -> Void
+) async {
+  let lastPriority = ThreadSafeBox(initialValue: initialPriority)
+  await withTaskGroup(of: Void.self) { taskGroup in
+    taskGroup.addTask {
+      while true {
+        if Task.isCancelled {
+          return
+        }
+        let newPriority = Task.currentPriority
+        let didChange = lastPriority.withLock { lastPriority in
+          if newPriority != lastPriority {
+            lastPriority = newPriority
+            return true
+          }
+          return false
+        }
+        if didChange {
+          taskPriorityChanged()
+        }
+        do {
+          try await Task.sleep(for: pollingInterval)
+        } catch {
+          break
+        }
+      }
+    }
+    taskGroup.addTask {
+      await operation()
+    }
+    // The first task that watches the priority never finishes, so we are effectively await the `operation` task here
+    // and cancelling the priority observation task once the operation task is done.
+    // We do need to await the observation task as well so that priority escalation also affects the observation task.
+    for await _ in taskGroup {
+      taskGroup.cancelAll()
+    }
+  }
+}