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

This was causing a non-deterministic test failure: When target preparation finishes while a diagnostic request is in progress, it will re-open the document, which calls `DiagnosticReportManager.removeItemsFromCache` for that document’s URI. With the old implementation, we would thus cancel the diagnostics sourcekitd request and return a cancelled error to the diagnostics LSP request.

While doing this, I also realized that there was a race condition: Document re-opening would happen outside of the SourceKit-LSP message handling queue and could thus run concurrently to any other request. This means that a sourcekitd request could run after `reopenDocument` had closed the document but before it was opened again. Introduce an internal reopen request that can be handled on the main message handling queue and thus doesn’t have this problem

Author: ahoppen

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 {