Implicitly cancel text document requests when the document is edited or closed

As a user makes an edit to a file, these requests are most likely no longer relevant. It also makes sure that a long-running sourcekitd request can't block the entire language server if the client does not cancel all requests. For example, consider the following sequence of requests:
 - `textDocument/semanticTokens/full` for document A
 - `textDocument/didChange` for document A
 - `textDocument/formatting` for document A

If the editor is not cancelling the semantic tokens request on edit (like VS Code does), then the `didChange` notification is blocked on the semantic tokens request finishing. Hence, we also can't run the `textDocument/formatting` request. Cancelling the semantic tokens on the edit fixes the issue.

rdar://133987424

Author: ahoppen

Documentation/Configuration File.md

@@ -44,5 +44,6 @@ The structure of the file is currently not guaranteed to be stable. Options may
 - `generatedFilesPath: string`: Directory in which generated interfaces and macro expansions should be stored.
 - `backgroundIndexing: bool`: Explicitly enable or disable background indexing.
 - `backgroundPreparationMode: "build"|"noLazy"|"enabled"`: Determines how background indexing should prepare a target. Possible values are: `build`: Build a target to prepare it, `noLazy`: Prepare a target without generating object files but do not do lazy type checking and function body skipping, `enabled`: Prepare a target without generating object files and the like
+- `cancelTextDocumentRequestsOnEditAndClose: bool`: Whether sending a `textDocument/didChange` or `textDocument/didClose` notification for a document should cancel all pending requests for that document.
 - `experimentalFeatures: string[]`: Experimental features to enable
 - `swiftPublishDiagnosticsDebounce`: The time that `SwiftLanguageService` should wait after an edit before starting to compute diagnostics and sending a `PublishDiagnosticsNotification`.

Sources/SKOptions/SourceKitLSPOptions.swift

@@ -286,6 +286,14 @@ public struct SourceKitLSPOptions: Sendable, Codable, CustomLogStringConvertible
     return .build
   }
 
+  /// Whether sending a `textDocument/didChange` or `textDocument/didClose` notification for a document should cancel
+  /// all pending requests for that document.
+  public var cancelTextDocumentRequestsOnEditAndClose: Bool? = nil
+
+  public var cancelTextDocumentRequestsOnEditAndCloseOrDefault: Bool {
+    return cancelTextDocumentRequestsOnEditAndClose ?? true
+  }
+
   /// Experimental features that are enabled.
   public var experimentalFeatures: Set<ExperimentalFeature>? = nil
 
@@ -343,6 +351,7 @@ public struct SourceKitLSPOptions: Sendable, Codable, CustomLogStringConvertible
     generatedFilesPath: String? = nil,
     backgroundIndexing: Bool? = nil,
     backgroundPreparationMode: String? = nil,
+    cancelTextDocumentRequestsOnEditAndClose: Bool? = nil,
     experimentalFeatures: Set<ExperimentalFeature>? = nil,
     swiftPublishDiagnosticsDebounceDuration: Double? = nil,
     workDoneProgressDebounceDuration: Double? = nil,
@@ -358,6 +367,7 @@ public struct SourceKitLSPOptions: Sendable, Codable, CustomLogStringConvertible
     self.defaultWorkspaceType = defaultWorkspaceType
     self.backgroundIndexing = backgroundIndexing
     self.backgroundPreparationMode = backgroundPreparationMode
+    self.cancelTextDocumentRequestsOnEditAndClose = cancelTextDocumentRequestsOnEditAndClose
     self.experimentalFeatures = experimentalFeatures
     self.swiftPublishDiagnosticsDebounceDuration = swiftPublishDiagnosticsDebounceDuration
     self.workDoneProgressDebounceDuration = workDoneProgressDebounceDuration
@@ -412,6 +422,8 @@ public struct SourceKitLSPOptions: Sendable, Codable, CustomLogStringConvertible
       generatedFilesPath: override?.generatedFilesPath ?? base.generatedFilesPath,
       backgroundIndexing: override?.backgroundIndexing ?? base.backgroundIndexing,
       backgroundPreparationMode: override?.backgroundPreparationMode ?? base.backgroundPreparationMode,
+      cancelTextDocumentRequestsOnEditAndClose: override?.cancelTextDocumentRequestsOnEditAndClose
+        ?? base.cancelTextDocumentRequestsOnEditAndClose,
       experimentalFeatures: override?.experimentalFeatures ?? base.experimentalFeatures,
       swiftPublishDiagnosticsDebounceDuration: override?.swiftPublishDiagnosticsDebounceDuration
         ?? base.swiftPublishDiagnosticsDebounceDuration,

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -177,7 +177,10 @@ package actor SourceKitLSPServer {
   /// The requests that we are currently handling.
   ///
   /// Used to cancel the tasks if the client requests cancellation.
-  private var inProgressRequests: [RequestID: Task<(), Never>] = [:]
+  private var inProgressRequestsByID: [RequestID: Task<(), Never>] = [:]
+
+  /// For all currently handled text document requests a mapping from the document to the corresponding request ID.
+  private var inProgressTextDocumentRequests: [DocumentURI: Set<RequestID>] = [:]
 
   /// Up to 10 request IDs that have recently finished.
   ///
@@ -187,14 +190,22 @@ package actor SourceKitLSPServer {
 
   /// - Note: Needed so we can set an in-progress request from a different
   ///   isolation context.
-  private func setInProgressRequest(for id: RequestID, task: Task<(), Never>?) {
-    self.inProgressRequests[id] = task
+  private func setInProgressRequest(for id: RequestID, _ request: some RequestType, task: Task<(), Never>?) {
+    self.inProgressRequestsByID[id] = task
     if task == nil {
       recentlyFinishedRequests.append(id)
       while recentlyFinishedRequests.count > 10 {
         recentlyFinishedRequests.removeFirst()
       }
     }
+
+    if let request = request as? any TextDocumentRequest {
+      if task == nil {
+        inProgressTextDocumentRequests[request.textDocument.uri, default: []].remove(id)
+      } else {
+        inProgressTextDocumentRequests[request.textDocument.uri, default: []].insert(id)
+      }
+    }
   }
 
   var onExit: () -> Void
@@ -547,12 +558,19 @@ extension SourceKitLSPServer: MessageHandler {
   package nonisolated func handle(_ params: some NotificationType) {
     let notificationID = notificationIDForLogging.fetchAndIncrement()
     withLoggingScope("notification-\(notificationID % 100)") {
-      if let params = params as? CancelRequestNotification {
-        // Request cancellation needs to be able to overtake any other message we
-        // are currently handling. Ordering is not important here. We thus don't
-        // need to execute it on `messageHandlingQueue`.
+      // Request cancellation needs to be able to overtake any other message we
+      // are currently handling. Ordering is not important here. We thus don't
+      // need to execute it on `messageHandlingQueue`.
+      switch params {
+      case let params as CancelRequestNotification:
         self.cancelRequest(params)
         return
+      case let params as DidChangeTextDocumentNotification:
+        self.cancelTextDocumentRequests(for: params.textDocument.uri)
+      case let params as DidCloseTextDocumentNotification:
+        self.cancelTextDocumentRequests(for: params.textDocument.uri)
+      default:
+        break
       }
 
       let signposter = Logger(subsystem: LoggingScope.subsystem, category: "message-handling")
@@ -617,6 +635,7 @@ extension SourceKitLSPServer: MessageHandler {
       // The last 2 digits should be sufficient to differentiate between multiple concurrently running requests.
       await withLoggingScope("request-\(id.numericValue % 100)") {
         await withTaskCancellationHandler {
+          await self.testHooks.handleRequest?(params)
           await self.handleImpl(params, id: id, reply: reply)
           signposter.endInterval("Request", state, "Done")
         } onCancel: {
@@ -626,14 +645,14 @@ extension SourceKitLSPServer: MessageHandler {
       // We have handled the request and can't cancel it anymore.
       // Stop keeping track of it to free the memory.
       self.cancellationMessageHandlingQueue.async(priority: .background) {
-        await self.setInProgressRequest(for: id, task: nil)
+        await self.setInProgressRequest(for: id, params, task: nil)
       }
     }
     // Keep track of the ID -> Task management with low priority. Once we cancel
     // a request, the cancellation task runs with a high priority and depends on
     // this task, which will elevate this task's priority.
     cancellationMessageHandlingQueue.async(priority: .background) {
-      await self.setInProgressRequest(for: id, task: task)
+      await self.setInProgressRequest(for: id, params, task: task)
     }
   }
 
@@ -1222,11 +1241,11 @@ extension SourceKitLSPServer {
     // Nothing to do.
   }
 
-  nonisolated func cancelRequest(_ notification: CancelRequestNotification) {
+  private nonisolated func cancelRequest(_ notification: CancelRequestNotification) {
     // Since the request is very cheap to execute and stops other requests
     // from performing more work, we execute it with a high priority.
     cancellationMessageHandlingQueue.async(priority: .high) {
-      if let task = await self.inProgressRequests[notification.id] {
+      if let task = await self.inProgressRequestsByID[notification.id] {
         task.cancel()
         return
       }
@@ -1238,6 +1257,38 @@ extension SourceKitLSPServer {
     }
   }
 
+  /// Cancel all in-progress text document requests for the given document.
+  ///
+  /// As a user makes an edit to a file, these requests are most likely no longer relevant. It also makes sure that a
+  /// long-running sourcekitd request can't block the entire language server if the client does not cancel all requests.
+  /// For example, consider the following sequence of requests:
+  ///  - `textDocument/semanticTokens/full` for document A
+  ///  - `textDocument/didChange` for document A
+  ///  - `textDocument/formatting` for document A
+  ///
+  /// If the editor is not cancelling the semantic tokens request on edit (like VS Code does), then the `didChange`
+  /// notification is blocked on the semantic tokens request finishing. Hence, we also can't run the
+  /// `textDocument/formatting` request. Cancelling the semantic tokens on the edit fixes the issue.
+  ///
+  /// This method is a no-op if `cancelTextDocumentRequestsOnEditAndClose` is disabled.
+  private nonisolated func cancelTextDocumentRequests(for uri: DocumentURI) {
+    // Since the request is very cheap to execute and stops other requests
+    // from performing more work, we execute it with a high priority.
+    cancellationMessageHandlingQueue.async(priority: .high) {
+      await self.cancelTextDocumentRequestsImpl(for: uri)
+    }
+  }
+
+  private func cancelTextDocumentRequestsImpl(for uri: DocumentURI) {
+    guard self.options.cancelTextDocumentRequestsOnEditAndCloseOrDefault else {
+      return
+    }
+    for requestID in self.inProgressTextDocumentRequests[uri, default: []] {
+      logger.info("Implicitly cancelling request \(requestID)")
+      self.inProgressRequestsByID[requestID]?.cancel()
+    }
+  }
+
   /// The server is about to exit, and the server should flush any buffered state.
   ///
   /// The server shall not be used to handle more requests (other than possibly

Sources/SourceKitLSP/Swift/SemanticTokens.swift

@@ -36,6 +36,8 @@ extension SwiftLanguageService {
       return nil
     }
 
+    try Task.checkCancellation()
+
     return SyntaxHighlightingTokenParser(sourcekitd: sourcekitd).parseTokens(skTokens, in: snapshot)
   }
 
@@ -51,6 +53,8 @@ extension SwiftLanguageService {
     for snapshot: DocumentSnapshot,
     in range: Range<Position>? = nil
   ) async throws -> SyntaxHighlightingTokens {
+    try Task.checkCancellation()
+
     async let tree = syntaxTreeManager.syntaxTree(for: snapshot)
     let semanticTokens = await orLog("Loading semantic tokens") { try await semanticHighlightingTokens(for: snapshot) }
 
@@ -61,12 +65,16 @@ extension SwiftLanguageService {
         await tree.range
       }
 
+    try Task.checkCancellation()
+
     let tokens =
       await tree
       .classifications(in: range)
       .map { $0.highlightingTokens(in: snapshot) }
       .reduce(into: SyntaxHighlightingTokens(tokens: [])) { $0.tokens += $1.tokens }
 
+    try Task.checkCancellation()
+
     return
       tokens
       .mergingTokens(with: semanticTokens ?? SyntaxHighlightingTokens(tokens: []))

Sources/SourceKitLSP/TestHooks.swift

@@ -25,15 +25,22 @@ public struct TestHooks: Sendable {
 
   package var swiftpmTestHooks: SwiftPMTestHooks
 
+  /// A hook that will be executed before a request is handled.
+  ///
+  /// This allows requests to be artificially delayed.
+  package var handleRequest: (@Sendable (any RequestType) async -> Void)?
+
   public init() {
-    self.init(indexTestHooks: IndexTestHooks(), swiftpmTestHooks: SwiftPMTestHooks())
+    self.init(indexTestHooks: IndexTestHooks(), swiftpmTestHooks: SwiftPMTestHooks(), handleRequest: nil)
   }
 
   package init(
     indexTestHooks: IndexTestHooks = IndexTestHooks(),
-    swiftpmTestHooks: SwiftPMTestHooks = SwiftPMTestHooks()
+    swiftpmTestHooks: SwiftPMTestHooks = SwiftPMTestHooks(),
+    handleRequest: (@Sendable (any RequestType) async -> Void)? = nil
   ) {
     self.indexTestHooks = indexTestHooks
     self.swiftpmTestHooks = swiftpmTestHooks
+    self.handleRequest = handleRequest
   }
 }

Tests/SourceKitLSPTests/SemanticTokensTests.swift

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 import LanguageServerProtocol
+import SKOptions
 import SKSupport
 import SKTestSupport
 import SourceKitD
@@ -943,6 +944,61 @@ final class SemanticTokensTests: XCTestCase {
       ]
     )
   }
+
+  func testImplicitCancellationOnEdit() async throws {
+    let testClient = try await TestSourceKitLSPClient(
+      testHooks: TestHooks(handleRequest: { request in
+        if request is DocumentSemanticTokensRequest {
+          // Sleep long enough for the edit to be handled
+          try? await Task.sleep(for: .seconds(10))
+        }
+      })
+    )
+    let uri = DocumentURI(for: .swift)
+    let positions = testClient.openDocument("1️⃣", uri: uri)
+
+    let receivedSemanticTokensResponse = self.expectation(description: "Received semantic tokens response")
+    testClient.send(DocumentSemanticTokensRequest(textDocument: TextDocumentIdentifier(uri))) { result in
+      XCTAssertEqual(result, .failure(ResponseError.cancelled))
+      receivedSemanticTokensResponse.fulfill()
+    }
+    testClient.send(
+      DidChangeTextDocumentNotification(
+        textDocument: VersionedTextDocumentIdentifier(uri, version: 2),
+        contentChanges: [TextDocumentContentChangeEvent(range: Range(positions["1️⃣"]), text: "let x = 1")]
+      )
+    )
+    try await fulfillmentOfOrThrow([receivedSemanticTokensResponse])
+  }
+
+  func testNoImplicitCancellationOnEditIfImplicitCancellationIsDisabled() async throws {
+    try SkipUnless.longTestsEnabled()
+
+    let testClient = try await TestSourceKitLSPClient(
+      options: SourceKitLSPOptions(cancelTextDocumentRequestsOnEditAndClose: false),
+      testHooks: TestHooks(handleRequest: { request in
+        if request is DocumentSemanticTokensRequest {
+          // Sleep long enough for the edit to be handled
+          try? await Task.sleep(for: .seconds(2))
+        }
+      })
+    )
+    let uri = DocumentURI(for: .swift)
+    let positions = testClient.openDocument("1️⃣", uri: uri)
+
+    let receivedSemanticTokensResponse = self.expectation(description: "Received semantic tokens response")
+    testClient.send(DocumentSemanticTokensRequest(textDocument: TextDocumentIdentifier(uri))) { result in
+      XCTAssertEqual(result, .success(DocumentSemanticTokensResponse(data: [])))
+      receivedSemanticTokensResponse.fulfill()
+    }
+    testClient.send(
+      DidChangeTextDocumentNotification(
+        textDocument: VersionedTextDocumentIdentifier(uri, version: 2),
+        contentChanges: [TextDocumentContentChangeEvent(range: Range(positions["1️⃣"]), text: "let x = 1")]
+      )
+    )
+    try await fulfillmentOfOrThrow([receivedSemanticTokensResponse])
+  }
 }
 
 fileprivate struct TokenSpec {