Don’t cancel in-progress diagnostic generation when calling DiagnosticReportManager.removeItemsFromCache

Sources/LanguageServerProtocol/CMakeLists.txt

@@ -19,6 +19,7 @@ add_library(LanguageServerProtocol STATIC
   Notifications/LogMessageNotification.swift
   Notifications/LogTraceNotification.swift
   Notifications/PublishDiagnosticsNotification.swift
+  Notifications/ReopenTextDocumentNotifications.swift
   Notifications/SetTraceNotification.swift
   Notifications/ShowMessageNotification.swift
   Notifications/TextSynchronizationNotifications.swift

Sources/LanguageServerProtocol/Notifications/ReopenTextDocumentNotifications.swift

@@ -0,0 +1,29 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+/// Re-open the given document, discarding any in-memory state.
+///
+/// This notification is designed to be used internally in SourceKit-LSP: When build setting have changed, we re-open
+/// the document in sourcekitd to re-build its AST. This needs to be handled via a notification to ensure that no other
+/// request for this document is executing at the same time.
+///
+/// **(LSP Extension)**
+public struct ReopenTextDocumentNotification: NotificationType, Hashable {
+  public static let method: String = "textDocument/reopen"
+
+  /// The document identifier and initial contents.
+  public var textDocument: TextDocumentIdentifier
+
+  public init(textDocument: TextDocumentIdentifier) {
+    self.textDocument = textDocument
+  }
+}

Sources/SourceKitLSP/Clang/ClangLanguageService.swift

@@ -478,6 +478,8 @@ extension ClangLanguageService {
     clangd.send(notification)
   }
 
+  func reopenDocument(_ notification: ReopenTextDocumentNotification) {}
+
   public func changeDocument(_ notification: DidChangeTextDocumentNotification) {
     clangd.send(notification)
   }

Sources/SourceKitLSP/LanguageService.swift

@@ -118,6 +118,14 @@ public protocol LanguageService: AnyObject, Sendable {
 
   /// Sent to close a document on the Language Server.
   func closeDocument(_ notification: DidCloseTextDocumentNotification) async
+
+  /// Re-open the given document, discarding any in-memory state and forcing an AST to be re-built after build settings
+  /// have been changed. This needs to be handled via a notification to ensure that no other request for this document
+  /// is executing at the same time.
+  ///
+  /// Only intended for `SwiftLanguageService`.
+  func reopenDocument(_ notification: ReopenTextDocumentNotification) async
+
   func changeDocument(_ notification: DidChangeTextDocumentNotification) async
   func willSaveDocument(_ notification: WillSaveTextDocumentNotification) async
   func didSaveDocument(_ notification: DidSaveTextDocumentNotification) async

Sources/SourceKitLSP/MessageHandlingDependencyTracker.swift

@@ -126,6 +126,8 @@ enum MessageHandlingDependencyTracker: DependencyTracker {
       self = .freestanding
     case is PublishDiagnosticsNotification:
       self = .freestanding
+    case let notification as ReopenTextDocumentNotification:
+      self = .documentUpdate(notification.textDocument.uri)
     case is SetTraceNotification:
       self = .globalConfigurationChange
     case is ShowMessageNotification:

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -569,6 +569,8 @@ extension SourceKitLSPServer: MessageHandler {
       await self.openDocument(notification)
     case let notification as DidCloseTextDocumentNotification:
       await self.closeDocument(notification)
+    case let notification as ReopenTextDocumentNotification:
+      await self.reopenDocument(notification)
     case let notification as DidChangeTextDocumentNotification:
       await self.changeDocument(notification)
     case let notification as DidChangeWorkspaceFoldersNotification:
@@ -1300,6 +1302,17 @@ extension SourceKitLSPServer {
     await self.closeDocument(notification, workspace: workspace)
   }
 
+  func reopenDocument(_ notification: ReopenTextDocumentNotification) async {
+    let uri = notification.textDocument.uri
+    guard let workspace = await workspaceForDocument(uri: uri) else {
+      logger.error(
+        "received reopen notification for file '\(uri.forLogging)' without a corresponding workspace, ignoring..."
+      )
+      return
+    }
+    await workspace.documentService.value[uri]?.reopenDocument(notification)
+  }
+
   func closeDocument(_ notification: DidCloseTextDocumentNotification, workspace: Workspace) async {
     // Immediately close the document. We need to be sure to clear our pending work queue in case
     // the build system still isn't ready.

Sources/SourceKitLSP/Swift/DiagnosticReportManager.swift

@@ -86,11 +86,7 @@ actor DiagnosticReportManager {
   }
 
   func removeItemsFromCache(with uri: DocumentURI) async {
-    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(

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -327,15 +327,24 @@ extension SwiftLanguageService {
 
   // MARK: - Build System Integration
 
-  private func reopenDocument(_ snapshot: DocumentSnapshot, _ compileCmd: SwiftCompileCommand?) async {
+  public func reopenDocument(_ notification: ReopenTextDocumentNotification) async {
+    let snapshot = orLog("Getting snapshot to re-open document") {
+      try documentManager.latestSnapshot(notification.textDocument.uri)
+    }
+    guard let snapshot else {
+      return
+    }
     cancelInFlightPublishDiagnosticsTask(for: snapshot.uri)
     await diagnosticReportManager.removeItemsFromCache(with: snapshot.uri)
 
     let closeReq = closeDocumentSourcekitdRequest(uri: snapshot.uri)
-    _ = try? await self.sourcekitd.send(closeReq, fileContents: nil)
+    _ = await orLog("Closing document to re-open it") { try await self.sourcekitd.send(closeReq, fileContents: nil) }
 
-    let openReq = openDocumentSourcekitdRequest(snapshot: snapshot, compileCommand: compileCmd)
-    _ = try? await self.sourcekitd.send(openReq, fileContents: snapshot.text)
+    let openReq = openDocumentSourcekitdRequest(
+      snapshot: snapshot,
+      compileCommand: await buildSettings(for: snapshot.uri)
+    )
+    _ = await orLog("Re-opening document") { try await self.sourcekitd.send(openReq, fileContents: snapshot.text) }
 
     if await capabilityRegistry.clientSupportsPullDiagnostics(for: .swift) {
       await self.refreshDiagnosticsDebouncer.scheduleCall()
@@ -345,14 +354,11 @@ extension SwiftLanguageService {
   }
 
   public func documentUpdatedBuildSettings(_ uri: DocumentURI) async {
-    // We may not have a snapshot if this is called just before `openDocument`.
-    guard let snapshot = try? self.documentManager.latestSnapshot(uri) else {
-      return
-    }
-
-    // Close and re-open the document internally to inform sourcekitd to update the compile
-    // command. At the moment there's no better way to do this.
-    await self.reopenDocument(snapshot, await self.buildSettings(for: uri))
+    // Close and re-open the document internally to inform sourcekitd to update the compile command. At the moment
+    // there's no better way to do this.
+    // Schedule the document re-open in the SourceKit-LSP server. This ensures that the re-open happens exclusively with
+    // no other request running at the same time.
+    sourceKitLSPServer?.handle(ReopenTextDocumentNotification(textDocument: TextDocumentIdentifier(uri)))
   }
 
   public func documentDependenciesUpdated(_ uri: DocumentURI) async {