Cancel the sourcekitd request to retrieve diagnostics if noone depends on it anymore

After a document has been modified, the corresponding publish diagnsotics request gets cancelled. But that cancellation didn’t propagate through `DiagnosticReportManager`.

rdar://121323969

Author: ahoppen

Sources/SKSupport/AsyncUtils.swift

@@ -12,6 +12,55 @@
 
 import Foundation
 
+/// Wrapper around a task that allows multiple clients to depend on the task's value.
+///
+/// If all of the dependents are cancelled, the underlying task is cancelled as well.
+public actor RefCountedCancellableTask<Success> {
+  public let task: Task<Success, Error>
+
+  /// The number of clients that depend on the task's result and that are not cancelled.
+  private var refCount: Int = 0
+
+  /// Whether the task has been cancelled.
+  public private(set) var isCancelled: Bool = false
+
+  public init(priority: TaskPriority? = nil, operation: @escaping @Sendable () async throws -> Success) {
+    self.task = Task(priority: priority, operation: operation)
+  }
+
+  private func decrementRefCount() {
+    refCount -= 1
+    if refCount == 0 {
+      self.cancel()
+    }
+  }
+
+  /// Get the task's value.
+  ///
+  /// If all callers of `value` are cancelled, the underlying task gets cancelled as well.
+  public var value: Success {
+    get async throws {
+      if isCancelled {
+        throw CancellationError()
+      }
+      refCount += 1
+      return try await withTaskCancellationHandler {
+        return try await task.value
+      } onCancel: {
+        Task {
+          await self.decrementRefCount()
+        }
+      }
+    }
+  }
+
+  /// Cancel the task and throw a `CancellationError` to all clients that are awaiting the value.
+  public func cancel() {
+    isCancelled = true
+    task.cancel()
+  }
+}
+
 public extension Task {
   /// Awaits the value of the result.
   ///

Sources/SourceKitLSP/Swift/DiagnosticReportManager.swift

@@ -12,12 +12,13 @@
 
 import LSPLogging
 import LanguageServerProtocol
+import SKSupport
 import SourceKitD
 import SwiftParserDiagnostics
 
 actor DiagnosticReportManager {
   /// A task to produce diagnostics, either from a diagnostics request to `sourcektid` or by using the built-in swift-syntax.
-  private typealias ReportTask = Task<RelatedFullDocumentDiagnosticReport, Error>
+  private typealias ReportTask = RefCountedCancellableTask<RelatedFullDocumentDiagnosticReport>
 
   private let sourcekitd: SourceKitD
   private let syntaxTreeManager: SyntaxTreeManager
@@ -60,13 +61,13 @@ actor DiagnosticReportManager {
     for snapshot: DocumentSnapshot,
     buildSettings: SwiftCompileCommand?
   ) async throws -> RelatedFullDocumentDiagnosticReport {
-    if let reportTask = reportTask(for: snapshot.id, buildSettings: buildSettings) {
+    if let reportTask = reportTask(for: snapshot.id, buildSettings: buildSettings), await !reportTask.isCancelled {
       return try await reportTask.value
     }
-    let reportTask: Task<RelatedFullDocumentDiagnosticReport, Error>
+    let reportTask: ReportTask
     if let buildSettings, !buildSettings.isFallback {
-      reportTask = Task {
-        return try await requestReport(with: snapshot, compilerArgs: buildSettings.compilerArgs)
+      reportTask = ReportTask {
+        return try await self.requestReport(with: snapshot, compilerArgs: buildSettings.compilerArgs)
       }
     } else {
       logger.log(
@@ -76,21 +77,20 @@ actor DiagnosticReportManager {
       // sourcekitd won't be able to give us accurate semantic diagnostics.
       // Fall back to providing syntactic diagnostics from the built-in
       // swift-syntax. That's the best we can do for now.
-      reportTask = Task {
-        return try await requestFallbackReport(with: snapshot)
+      reportTask = ReportTask {
+        return try await self.requestFallbackReport(with: snapshot)
       }
     }
     setReportTask(for: snapshot.id, buildSettings: buildSettings, reportTask: reportTask)
     return try await reportTask.value
   }
 
   func removeItemsFromCache(with uri: DocumentURI) async {
-    for item in reportTaskCache {
-      if item.snapshotID.uri == uri {
-        item.reportTask.cancel()
-      }
-    }
+    let tasksToCancel = reportTaskCache.filter { $0.snapshotID.uri == uri }
     reportTaskCache.removeAll(where: { $0.snapshotID.uri == uri })
+    for task in tasksToCancel {
+      await task.reportTask.cancel()
+    }
   }
 
   private func requestReport(

Tests/SourceKitLSPTests/LifecycleTests.swift

@@ -99,10 +99,16 @@ final class LifecycleTests: XCTestCase {
       )
     )
 
+    let cursorInfoStartDate = Date()
     // Check that semantic functionality based on the AST is working again.
     let symbolInfo = try await testClient.send(
       SymbolInfoRequest(textDocument: TextDocumentIdentifier(uri), position: positions["3️⃣"])
     )
+    XCTAssertLessThan(
+      Date().timeIntervalSince(cursorInfoStartDate),
+      2,
+      "Cursor info request wasn't fast. sourcekitd still blocked?"
+    )
     XCTAssertGreaterThan(symbolInfo.count, 0)
   }
 }